Add ExecPolicyManager (#8349)

Move exec policy management into services to keep turn context
immutable.

codex-rs/core/src/exec_policy.rs

@@ -3,6 +3,8 @@ use std::path::Path;
 use std::path::PathBuf;
 use std::sync::Arc;
 
+use arc_swap::ArcSwap;
+
 use crate::command_safety::is_dangerous_command::requires_initial_appoval;
 use codex_execpolicy::AmendError;
 use codex_execpolicy::Decision;
@@ -17,7 +19,6 @@ use codex_protocol::protocol::AskForApproval;
 use codex_protocol::protocol::SandboxPolicy;
 use thiserror::Error;
 use tokio::fs;
-use tokio::sync::RwLock;
 use tokio::task::spawn_blocking;
 
 use crate::bash::parse_shell_lc_plain_commands;
@@ -80,7 +81,118 @@ pub enum ExecPolicyUpdateError {
     FeatureDisabled,
 }
 
-pub(crate) async fn load_exec_policy_for_features(
+pub(crate) struct ExecPolicyManager {
+    policy: ArcSwap<Policy>,
+}
+
+impl ExecPolicyManager {
+    pub(crate) fn new(policy: Arc<Policy>) -> Self {
+        Self {
+            policy: ArcSwap::from(policy),
+        }
+    }
+
+    pub(crate) async fn load(
+        features: &Features,
+        codex_home: &Path,
+    ) -> Result<Self, ExecPolicyError> {
+        let policy = load_exec_policy_for_features(features, codex_home).await?;
+        Ok(Self::new(Arc::new(policy)))
+    }
+
+    pub(crate) fn current(&self) -> Arc<Policy> {
+        self.policy.load_full()
+    }
+
+    pub(crate) async fn create_exec_approval_requirement_for_command(
+        &self,
+        features: &Features,
+        command: &[String],
+        approval_policy: AskForApproval,
+        sandbox_policy: &SandboxPolicy,
+        sandbox_permissions: SandboxPermissions,
+    ) -> ExecApprovalRequirement {
+        let exec_policy = self.current();
+        let commands =
+            parse_shell_lc_plain_commands(command).unwrap_or_else(|| vec![command.to_vec()]);
+        let heuristics_fallback = |cmd: &[String]| {
+            if requires_initial_appoval(approval_policy, sandbox_policy, cmd, sandbox_permissions) {
+                Decision::Prompt
+            } else {
+                Decision::Allow
+            }
+        };
+        let evaluation = exec_policy.check_multiple(commands.iter(), &heuristics_fallback);
+
+        match evaluation.decision {
+            Decision::Forbidden => ExecApprovalRequirement::Forbidden {
+                reason: FORBIDDEN_REASON.to_string(),
+            },
+            Decision::Prompt => {
+                if matches!(approval_policy, AskForApproval::Never) {
+                    ExecApprovalRequirement::Forbidden {
+                        reason: PROMPT_CONFLICT_REASON.to_string(),
+                    }
+                } else {
+                    ExecApprovalRequirement::NeedsApproval {
+                        reason: derive_prompt_reason(&evaluation),
+                        proposed_execpolicy_amendment: if features.enabled(Feature::ExecPolicy) {
+                            try_derive_execpolicy_amendment_for_prompt_rules(
+                                &evaluation.matched_rules,
+                            )
+                        } else {
+                            None
+                        },
+                    }
+                }
+            }
+            Decision::Allow => ExecApprovalRequirement::Skip {
+                // Bypass sandbox if execpolicy allows the command
+                bypass_sandbox: evaluation.matched_rules.iter().any(|rule_match| {
+                    is_policy_match(rule_match) && rule_match.decision() == Decision::Allow
+                }),
+                proposed_execpolicy_amendment: if features.enabled(Feature::ExecPolicy) {
+                    try_derive_execpolicy_amendment_for_allow_rules(&evaluation.matched_rules)
+                } else {
+                    None
+                },
+            },
+        }
+    }
+
+    pub(crate) async fn append_amendment_and_update(
+        &self,
+        codex_home: &Path,
+        amendment: &ExecPolicyAmendment,
+    ) -> Result<(), ExecPolicyUpdateError> {
+        let policy_path = default_policy_path(codex_home);
+        let prefix = amendment.command.clone();
+        spawn_blocking({
+            let policy_path = policy_path.clone();
+            let prefix = prefix.clone();
+            move || blocking_append_allow_prefix_rule(&policy_path, &prefix)
+        })
+        .await
+        .map_err(|source| ExecPolicyUpdateError::JoinBlockingTask { source })?
+        .map_err(|source| ExecPolicyUpdateError::AppendRule {
+            path: policy_path,
+            source,
+        })?;
+
+        let mut updated_policy = self.current().as_ref().clone();
+        updated_policy.add_prefix_rule(&prefix, Decision::Allow)?;
+        self.policy.store(Arc::new(updated_policy));
+        Ok(())
+    }
+}
+
+impl Default for ExecPolicyManager {
+    fn default() -> Self {
+        Self::new(Arc::new(Policy::empty()))
+    }
+}
+
+async fn load_exec_policy_for_features(
     features: &Features,
     codex_home: &Path,
 ) -> Result<Policy, ExecPolicyError> {
@@ -123,37 +235,10 @@ pub async fn load_exec_policy(codex_home: &Path) -> Result<Policy, ExecPolicyErr
     Ok(policy)
 }
 
-pub(crate) fn default_policy_path(codex_home: &Path) -> PathBuf {
+fn default_policy_path(codex_home: &Path) -> PathBuf {
     codex_home.join(RULES_DIR_NAME).join(DEFAULT_POLICY_FILE)
 }
 
-pub(crate) async fn append_execpolicy_amendment_and_update(
-    codex_home: &Path,
-    current_policy: &Arc<RwLock<Policy>>,
-    prefix: &[String],
-) -> Result<(), ExecPolicyUpdateError> {
-    let policy_path = default_policy_path(codex_home);
-    let prefix = prefix.to_vec();
-    spawn_blocking({
-        let policy_path = policy_path.clone();
-        let prefix = prefix.clone();
-        move || blocking_append_allow_prefix_rule(&policy_path, &prefix)
-    })
-    .await
-    .map_err(|source| ExecPolicyUpdateError::JoinBlockingTask { source })?
-    .map_err(|source| ExecPolicyUpdateError::AppendRule {
-        path: policy_path,
-        source,
-    })?;
-
-    current_policy
-        .write()
-        .await
-        .add_prefix_rule(&prefix, Decision::Allow)?;
-
-    Ok(())
-}
-
 /// Derive a proposed execpolicy amendment when a command requires user approval
 /// - If any execpolicy rule prompts, return None, because an amendment would not skip that policy requirement.
 /// - Otherwise return the first heuristics Prompt.
@@ -217,59 +302,6 @@ fn derive_prompt_reason(evaluation: &Evaluation) -> Option<String> {
     })
 }
 
-pub(crate) async fn create_exec_approval_requirement_for_command(
-    exec_policy: &Arc<RwLock<Policy>>,
-    features: &Features,
-    command: &[String],
-    approval_policy: AskForApproval,
-    sandbox_policy: &SandboxPolicy,
-    sandbox_permissions: SandboxPermissions,
-) -> ExecApprovalRequirement {
-    let commands = parse_shell_lc_plain_commands(command).unwrap_or_else(|| vec![command.to_vec()]);
-    let heuristics_fallback = |cmd: &[String]| {
-        if requires_initial_appoval(approval_policy, sandbox_policy, cmd, sandbox_permissions) {
-            Decision::Prompt
-        } else {
-            Decision::Allow
-        }
-    };
-    let policy = exec_policy.read().await;
-    let evaluation = policy.check_multiple(commands.iter(), &heuristics_fallback);
-
-    match evaluation.decision {
-        Decision::Forbidden => ExecApprovalRequirement::Forbidden {
-            reason: FORBIDDEN_REASON.to_string(),
-        },
-        Decision::Prompt => {
-            if matches!(approval_policy, AskForApproval::Never) {
-                ExecApprovalRequirement::Forbidden {
-                    reason: PROMPT_CONFLICT_REASON.to_string(),
-                }
-            } else {
-                ExecApprovalRequirement::NeedsApproval {
-                    reason: derive_prompt_reason(&evaluation),
-                    proposed_execpolicy_amendment: if features.enabled(Feature::ExecPolicy) {
-                        try_derive_execpolicy_amendment_for_prompt_rules(&evaluation.matched_rules)
-                    } else {
-                        None
-                    },
-                }
-            }
-        }
-        Decision::Allow => ExecApprovalRequirement::Skip {
-            // Bypass sandbox if execpolicy allows the command
-            bypass_sandbox: evaluation.matched_rules.iter().any(|rule_match| {
-                is_policy_match(rule_match) && rule_match.decision() == Decision::Allow
-            }),
-            proposed_execpolicy_amendment: if features.enabled(Feature::ExecPolicy) {
-                try_derive_execpolicy_amendment_for_allow_rules(&evaluation.matched_rules)
-            } else {
-                None
-            },
-        },
-    }
-}
-
 async fn collect_policy_files(dir: &Path) -> Result<Vec<PathBuf>, ExecPolicyError> {
     let mut read_dir = match fs::read_dir(dir).await {
         Ok(read_dir) => read_dir,
@@ -334,9 +366,10 @@ mod tests {
         features.disable(Feature::ExecPolicy);
         let temp_dir = tempdir().expect("create temp dir");
 
-        let policy = load_exec_policy_for_features(&features, temp_dir.path())
+        let manager = ExecPolicyManager::load(&features, temp_dir.path())
             .await
-            .expect("policy result");
+            .expect("manager result");
+        let policy = manager.current();
 
         let commands = [vec!["rm".to_string()]];
         assert_eq!(
@@ -425,7 +458,7 @@ prefix_rule(pattern=["rm"], decision="forbidden")
         parser
             .parse("test.rules", policy_src)
             .expect("parse policy");
-        let policy = Arc::new(RwLock::new(parser.build()));
+        let policy = Arc::new(parser.build());
 
         let forbidden_script = vec![
             "bash".to_string(),
@@ -433,15 +466,16 @@ prefix_rule(pattern=["rm"], decision="forbidden")
             "rm -rf /tmp".to_string(),
         ];
 
-        let requirement = create_exec_approval_requirement_for_command(
-            &policy,
-            &Features::with_defaults(),
-            &forbidden_script,
-            AskForApproval::OnRequest,
-            &SandboxPolicy::DangerFullAccess,
-            SandboxPermissions::UseDefault,
-        )
-        .await;
+        let manager = ExecPolicyManager::new(policy);
+        let requirement = manager
+            .create_exec_approval_requirement_for_command(
+                &Features::with_defaults(),
+                &forbidden_script,
+                AskForApproval::OnRequest,
+                &SandboxPolicy::DangerFullAccess,
+                SandboxPermissions::UseDefault,
+            )
+            .await;
 
         assert_eq!(
             requirement,
@@ -458,18 +492,19 @@ prefix_rule(pattern=["rm"], decision="forbidden")
         parser
             .parse("test.rules", policy_src)
             .expect("parse policy");
-        let policy = Arc::new(RwLock::new(parser.build()));
+        let policy = Arc::new(parser.build());
         let command = vec!["rm".to_string()];
 
-        let requirement = create_exec_approval_requirement_for_command(
-            &policy,
-            &Features::with_defaults(),
-            &command,
-            AskForApproval::OnRequest,
-            &SandboxPolicy::DangerFullAccess,
-            SandboxPermissions::UseDefault,
-        )
-        .await;
+        let manager = ExecPolicyManager::new(policy);
+        let requirement = manager
+            .create_exec_approval_requirement_for_command(
+                &Features::with_defaults(),
+                &command,
+                AskForApproval::OnRequest,
+                &SandboxPolicy::DangerFullAccess,
+                SandboxPermissions::UseDefault,
+            )
+            .await;
 
         assert_eq!(
             requirement,
@@ -487,18 +522,19 @@ prefix_rule(pattern=["rm"], decision="forbidden")
         parser
             .parse("test.rules", policy_src)
             .expect("parse policy");
-        let policy = Arc::new(RwLock::new(parser.build()));
+        let policy = Arc::new(parser.build());
         let command = vec!["rm".to_string()];
 
-        let requirement = create_exec_approval_requirement_for_command(
-            &policy,
-            &Features::with_defaults(),
-            &command,
-            AskForApproval::Never,
-            &SandboxPolicy::DangerFullAccess,
-            SandboxPermissions::UseDefault,
-        )
-        .await;
+        let manager = ExecPolicyManager::new(policy);
+        let requirement = manager
+            .create_exec_approval_requirement_for_command(
+                &Features::with_defaults(),
+                &command,
+                AskForApproval::Never,
+                &SandboxPolicy::DangerFullAccess,
+                SandboxPermissions::UseDefault,
+            )
+            .await;
 
         assert_eq!(
             requirement,
@@ -512,16 +548,16 @@ prefix_rule(pattern=["rm"], decision="forbidden")
     async fn exec_approval_requirement_falls_back_to_heuristics() {
         let command = vec!["cargo".to_string(), "build".to_string()];
 
-        let empty_policy = Arc::new(RwLock::new(Policy::empty()));
-        let requirement = create_exec_approval_requirement_for_command(
-            &empty_policy,
-            &Features::with_defaults(),
-            &command,
-            AskForApproval::UnlessTrusted,
-            &SandboxPolicy::ReadOnly,
-            SandboxPermissions::UseDefault,
-        )
-        .await;
+        let manager = ExecPolicyManager::default();
+        let requirement = manager
+            .create_exec_approval_requirement_for_command(
+                &Features::with_defaults(),
+                &command,
+                AskForApproval::UnlessTrusted,
+                &SandboxPolicy::ReadOnly,
+                SandboxPermissions::UseDefault,
+            )
+            .await;
 
         assert_eq!(
             requirement,
@@ -539,7 +575,7 @@ prefix_rule(pattern=["rm"], decision="forbidden")
         parser
             .parse("test.rules", policy_src)
             .expect("parse policy");
-        let policy = Arc::new(RwLock::new(parser.build()));
+        let policy = Arc::new(parser.build());
         let command = vec![
             "bash".to_string(),
             "-lc".to_string(),
@@ -547,15 +583,15 @@ prefix_rule(pattern=["rm"], decision="forbidden")
         ];
 
         assert_eq!(
-            create_exec_approval_requirement_for_command(
-                &policy,
-                &Features::with_defaults(),
-                &command,
-                AskForApproval::UnlessTrusted,
-                &SandboxPolicy::DangerFullAccess,
-                SandboxPermissions::UseDefault,
-            )
-            .await,
+            ExecPolicyManager::new(policy)
+                .create_exec_approval_requirement_for_command(
+                    &Features::with_defaults(),
+                    &command,
+                    AskForApproval::UnlessTrusted,
+                    &SandboxPolicy::DangerFullAccess,
+                    SandboxPermissions::UseDefault,
+                )
+                .await,
             ExecApprovalRequirement::NeedsApproval {
                 reason: None,
                 proposed_execpolicy_amendment: Some(ExecPolicyAmendment::new(vec![
@@ -568,14 +604,16 @@ prefix_rule(pattern=["rm"], decision="forbidden")
     #[tokio::test]
     async fn append_execpolicy_amendment_updates_policy_and_file() {
         let codex_home = tempdir().expect("create temp dir");
-        let current_policy = Arc::new(RwLock::new(Policy::empty()));
         let prefix = vec!["echo".to_string(), "hello".to_string()];
+        let manager = ExecPolicyManager::default();
 
-        append_execpolicy_amendment_and_update(codex_home.path(), &current_policy, &prefix)
+        manager
+            .append_amendment_and_update(codex_home.path(), &ExecPolicyAmendment::from(prefix))
             .await
             .expect("update policy");
+        let updated_policy = manager.current();
 
-        let evaluation = current_policy.read().await.check(
+        let evaluation = updated_policy.check(
             &["echo".to_string(), "hello".to_string(), "world".to_string()],
             &|_| Decision::Allow,
         );
@@ -599,10 +637,11 @@ prefix_rule(pattern=["rm"], decision="forbidden")
     #[tokio::test]
     async fn append_execpolicy_amendment_rejects_empty_prefix() {
         let codex_home = tempdir().expect("create temp dir");
-        let current_policy = Arc::new(RwLock::new(Policy::empty()));
+        let manager = ExecPolicyManager::default();
 
-        let result =
-            append_execpolicy_amendment_and_update(codex_home.path(), &current_policy, &[]).await;
+        let result = manager
+            .append_amendment_and_update(codex_home.path(), &ExecPolicyAmendment::from(vec![]))
+            .await;
 
         assert!(matches!(
             result,
@@ -617,16 +656,16 @@ prefix_rule(pattern=["rm"], decision="forbidden")
     async fn proposed_execpolicy_amendment_is_present_for_single_command_without_policy_match() {
         let command = vec!["cargo".to_string(), "build".to_string()];
 
-        let empty_policy = Arc::new(RwLock::new(Policy::empty()));
-        let requirement = create_exec_approval_requirement_for_command(
-            &empty_policy,
-            &Features::with_defaults(),
-            &command,
-            AskForApproval::UnlessTrusted,
-            &SandboxPolicy::ReadOnly,
-            SandboxPermissions::UseDefault,
-        )
-        .await;
+        let manager = ExecPolicyManager::default();
+        let requirement = manager
+            .create_exec_approval_requirement_for_command(
+                &Features::with_defaults(),
+                &command,
+                AskForApproval::UnlessTrusted,
+                &SandboxPolicy::ReadOnly,
+                SandboxPermissions::UseDefault,
+            )
+            .await;
 
         assert_eq!(
             requirement,
@@ -644,15 +683,16 @@ prefix_rule(pattern=["rm"], decision="forbidden")
         let mut features = Features::with_defaults();
         features.disable(Feature::ExecPolicy);
 
-        let requirement = create_exec_approval_requirement_for_command(
-            &Arc::new(RwLock::new(Policy::empty())),
-            &features,
-            &command,
-            AskForApproval::UnlessTrusted,
-            &SandboxPolicy::ReadOnly,
-            SandboxPermissions::UseDefault,
-        )
-        .await;
+        let manager = ExecPolicyManager::default();
+        let requirement = manager
+            .create_exec_approval_requirement_for_command(
+                &features,
+                &command,
+                AskForApproval::UnlessTrusted,
+                &SandboxPolicy::ReadOnly,
+                SandboxPermissions::UseDefault,
+            )
+            .await;
 
         assert_eq!(
             requirement,
@@ -670,18 +710,19 @@ prefix_rule(pattern=["rm"], decision="forbidden")
         parser
             .parse("test.rules", policy_src)
             .expect("parse policy");
-        let policy = Arc::new(RwLock::new(parser.build()));
+        let policy = Arc::new(parser.build());
         let command = vec!["rm".to_string()];
 
-        let requirement = create_exec_approval_requirement_for_command(
-            &policy,
-            &Features::with_defaults(),
-            &command,
-            AskForApproval::OnRequest,
-            &SandboxPolicy::DangerFullAccess,
-            SandboxPermissions::UseDefault,
-        )
-        .await;
+        let manager = ExecPolicyManager::new(policy);
+        let requirement = manager
+            .create_exec_approval_requirement_for_command(
+                &Features::with_defaults(),
+                &command,
+                AskForApproval::OnRequest,
+                &SandboxPolicy::DangerFullAccess,
+                SandboxPermissions::UseDefault,
+            )
+            .await;
 
         assert_eq!(
             requirement,
@@ -699,15 +740,16 @@ prefix_rule(pattern=["rm"], decision="forbidden")
             "-lc".to_string(),
             "cargo build && echo ok".to_string(),
         ];
-        let requirement = create_exec_approval_requirement_for_command(
-            &Arc::new(RwLock::new(Policy::empty())),
-            &Features::with_defaults(),
-            &command,
-            AskForApproval::UnlessTrusted,
-            &SandboxPolicy::ReadOnly,
-            SandboxPermissions::UseDefault,
-        )
-        .await;
+        let manager = ExecPolicyManager::default();
+        let requirement = manager
+            .create_exec_approval_requirement_for_command(
+                &Features::with_defaults(),
+                &command,
+                AskForApproval::UnlessTrusted,
+                &SandboxPolicy::ReadOnly,
+                SandboxPermissions::UseDefault,
+            )
+            .await;
 
         assert_eq!(
             requirement,
@@ -728,7 +770,7 @@ prefix_rule(pattern=["rm"], decision="forbidden")
         parser
             .parse("test.rules", policy_src)
             .expect("parse policy");
-        let policy = Arc::new(RwLock::new(parser.build()));
+        let policy = Arc::new(parser.build());
 
         let command = vec![
             "bash".to_string(),
@@ -737,15 +779,15 @@ prefix_rule(pattern=["rm"], decision="forbidden")
         ];
 
         assert_eq!(
-            create_exec_approval_requirement_for_command(
-                &policy,
-                &Features::with_defaults(),
-                &command,
-                AskForApproval::UnlessTrusted,
-                &SandboxPolicy::ReadOnly,
-                SandboxPermissions::UseDefault,
-            )
-            .await,
+            ExecPolicyManager::new(policy)
+                .create_exec_approval_requirement_for_command(
+                    &Features::with_defaults(),
+                    &command,
+                    AskForApproval::UnlessTrusted,
+                    &SandboxPolicy::ReadOnly,
+                    SandboxPermissions::UseDefault,
+                )
+                .await,
             ExecApprovalRequirement::NeedsApproval {
                 reason: None,
                 proposed_execpolicy_amendment: Some(ExecPolicyAmendment::new(vec![
@@ -759,15 +801,16 @@ prefix_rule(pattern=["rm"], decision="forbidden")
     async fn proposed_execpolicy_amendment_is_present_when_heuristics_allow() {
         let command = vec!["echo".to_string(), "safe".to_string()];
 
-        let requirement = create_exec_approval_requirement_for_command(
-            &Arc::new(RwLock::new(Policy::empty())),
-            &Features::with_defaults(),
-            &command,
-            AskForApproval::OnRequest,
-            &SandboxPolicy::ReadOnly,
-            SandboxPermissions::UseDefault,
-        )
-        .await;
+        let manager = ExecPolicyManager::default();
+        let requirement = manager
+            .create_exec_approval_requirement_for_command(
+                &Features::with_defaults(),
+                &command,
+                AskForApproval::OnRequest,
+                &SandboxPolicy::ReadOnly,
+                SandboxPermissions::UseDefault,
+            )
+            .await;
 
         assert_eq!(
             requirement,
@@ -785,18 +828,19 @@ prefix_rule(pattern=["rm"], decision="forbidden")
         parser
             .parse("test.rules", policy_src)
             .expect("parse policy");
-        let policy = Arc::new(RwLock::new(parser.build()));
+        let policy = Arc::new(parser.build());
         let command = vec!["echo".to_string(), "safe".to_string()];
 
-        let requirement = create_exec_approval_requirement_for_command(
-            &policy,
-            &Features::with_defaults(),
-            &command,
-            AskForApproval::OnRequest,
-            &SandboxPolicy::ReadOnly,
-            SandboxPermissions::UseDefault,
-        )
-        .await;
+        let manager = ExecPolicyManager::new(policy);
+        let requirement = manager
+            .create_exec_approval_requirement_for_command(
+                &Features::with_defaults(),
+                &command,
+                AskForApproval::OnRequest,
+                &SandboxPolicy::ReadOnly,
+                SandboxPermissions::UseDefault,
+            )
+            .await;
 
         assert_eq!(
             requirement,