Skip to main content

zagens_tools/
policy_engine.rs

1//! Tool policy engine (kernel-v2 M3).
2//!
3//! Single decision point for approval, sandbox class, and parallel scheduling
4//! keys. Replaces scattered string heuristics once `tools.policy = "engine"`.
5
6use std::sync::atomic::{AtomicU64, Ordering};
7
8use serde::{Deserialize, Serialize};
9
10use crate::ApprovalRequirement;
11use crate::tool_manifest::{Footprint, FootprintProvenance, ResourceSet, SpawnClass, ToolManifest};
12
13static POLICY_SHADOW_COMPARISONS: AtomicU64 = AtomicU64::new(0);
14static POLICY_SHADOW_DIFFS: AtomicU64 = AtomicU64::new(0);
15
16/// Session mode slice consumed by policy rules.
17#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
18#[serde(rename_all = "snake_case")]
19pub enum PolicySessionMode {
20    Agent,
21    Plan,
22    Yolo,
23}
24
25/// Whether a tool call needs user approval before execution.
26#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
27#[serde(rename_all = "snake_case")]
28pub enum ApprovalNeed {
29    Auto,
30    Ask,
31    Deny,
32}
33
34/// Sandbox enforcement tier derived from footprint (proposal §8.1).
35#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
36#[serde(rename_all = "snake_case")]
37pub enum SandboxClass {
38    None,
39    Sandboxed,
40    Strict,
41}
42
43/// Resource bucket for DAG / batch parallel scheduling (M4 consumer).
44#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
45#[serde(rename_all = "snake_case")]
46pub enum ParallelResourceKey {
47    ReadOnlyBatch,
48}
49
50/// Policy output consumed by the turn loop and future DAG scheduler.
51#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
52pub struct PolicyDecision {
53    pub approval: ApprovalNeed,
54    pub sandbox: SandboxClass,
55    pub parallel_key: Option<ParallelResourceKey>,
56    /// Scheduling / UI read-only flag (distinct from parallel eligibility).
57    pub read_only: bool,
58}
59
60/// Legacy-compatible plan flags derived from a policy decision.
61#[derive(Debug, Clone, Copy, PartialEq, Eq)]
62pub struct PolicyPlanMeta {
63    pub approval_required: bool,
64    pub read_only: bool,
65    pub supports_parallel: bool,
66}
67
68/// Inputs for a single tool policy evaluation.
69#[derive(Debug, Clone)]
70pub struct PolicyInput {
71    pub session_mode: PolicySessionMode,
72    pub manifest: ToolManifest,
73    /// Registry hint until footprints fully replace capability metadata.
74    pub legacy_approval: ApprovalRequirement,
75    /// Built-in registry hint for parallel eligibility.
76    pub supports_parallel_hint: bool,
77    /// Session trust / YOLO bypasses approval prompts.
78    pub trust_mode: bool,
79}
80
81/// Shadow-mode counters (`tools.policy = "shadow"`).
82#[derive(Debug, Clone, Copy, PartialEq, Eq)]
83pub struct PolicyShadowStats {
84    pub comparisons: u64,
85    pub diffs: u64,
86}
87
88impl PolicyDecision {
89    #[must_use]
90    pub fn plan_meta(&self) -> PolicyPlanMeta {
91        PolicyPlanMeta {
92            approval_required: !matches!(self.approval, ApprovalNeed::Auto),
93            read_only: self.read_only,
94            supports_parallel: self.parallel_key.is_some(),
95        }
96    }
97}
98
99impl PolicyPlanMeta {
100    #[must_use]
101    pub fn differs_from(&self, other: &Self) -> bool {
102        self != other
103    }
104}
105
106/// Policy engine — stateless rule set (proposal §8.1, §8.1.2).
107#[derive(Debug, Clone, Copy, Default)]
108pub struct PolicyEngine;
109
110impl PolicyEngine {
111    #[must_use]
112    pub fn decide(input: &PolicyInput) -> PolicyDecision {
113        let footprint = effective_footprint(&input.manifest);
114        let approval = decide_approval(input, &footprint);
115        let sandbox = decide_sandbox(&footprint, input.manifest.provenance);
116        let parallel_key = decide_parallel_key(input, &footprint, approval);
117        let read_only = decide_read_only(&footprint, input.manifest.provenance);
118        PolicyDecision {
119            approval,
120            sandbox,
121            parallel_key,
122            read_only,
123        }
124    }
125}
126
127/// Record a shadow comparison between legacy and engine plan flags.
128pub fn record_policy_shadow_diff(
129    tool_name: &str,
130    legacy: &PolicyPlanMeta,
131    engine: &PolicyPlanMeta,
132) {
133    POLICY_SHADOW_COMPARISONS.fetch_add(1, Ordering::Relaxed);
134    if legacy.differs_from(engine) {
135        POLICY_SHADOW_DIFFS.fetch_add(1, Ordering::Relaxed);
136        tracing::debug!(
137            tool = tool_name,
138            legacy_approval = legacy.approval_required,
139            engine_approval = engine.approval_required,
140            legacy_read_only = legacy.read_only,
141            engine_read_only = engine.read_only,
142            legacy_parallel = legacy.supports_parallel,
143            engine_parallel = engine.supports_parallel,
144            "policy_engine shadow diff"
145        );
146    }
147}
148
149#[must_use]
150pub fn policy_shadow_stats() -> PolicyShadowStats {
151    PolicyShadowStats {
152        comparisons: POLICY_SHADOW_COMPARISONS.load(Ordering::Relaxed),
153        diffs: POLICY_SHADOW_DIFFS.load(Ordering::Relaxed),
154    }
155}
156
157fn effective_footprint(manifest: &ToolManifest) -> Footprint {
158    match manifest.provenance {
159        FootprintProvenance::McpSelfDeclared => {
160            let mut footprint = manifest.footprint.clone();
161            if footprint.writes.is_empty() {
162                footprint.writes = ResourceSet {
163                    workspace_write: true,
164                    network_write: true,
165                    ..ResourceSet::default()
166                };
167            }
168            footprint
169        }
170        _ => manifest.footprint.clone(),
171    }
172}
173
174fn decide_approval(input: &PolicyInput, footprint: &Footprint) -> ApprovalNeed {
175    if input.trust_mode || matches!(input.session_mode, PolicySessionMode::Yolo) {
176        return ApprovalNeed::Auto;
177    }
178
179    match input.manifest.provenance {
180        FootprintProvenance::McpSelfDeclared => ApprovalNeed::Ask,
181        _ => match input.legacy_approval {
182            ApprovalRequirement::Auto => {
183                if footprint.writes.is_empty() && footprint.spawns == SpawnClass::None {
184                    ApprovalNeed::Auto
185                } else {
186                    ApprovalNeed::Ask
187                }
188            }
189            ApprovalRequirement::Suggest | ApprovalRequirement::Required => ApprovalNeed::Ask,
190        },
191    }
192}
193
194fn decide_sandbox(footprint: &Footprint, provenance: FootprintProvenance) -> SandboxClass {
195    if provenance == FootprintProvenance::McpSelfDeclared {
196        return SandboxClass::Strict;
197    }
198    match footprint.spawns {
199        SpawnClass::Privileged => SandboxClass::Strict,
200        SpawnClass::Sandboxed => SandboxClass::Sandboxed,
201        SpawnClass::None if !footprint.writes.is_empty() => SandboxClass::Strict,
202        SpawnClass::None => SandboxClass::None,
203    }
204}
205
206fn decide_read_only(footprint: &Footprint, provenance: FootprintProvenance) -> bool {
207    if provenance == FootprintProvenance::McpSelfDeclared {
208        return false;
209    }
210    footprint.writes.is_empty() && footprint.spawns == SpawnClass::None
211}
212
213fn decide_parallel_key(
214    input: &PolicyInput,
215    footprint: &Footprint,
216    approval: ApprovalNeed,
217) -> Option<ParallelResourceKey> {
218    if input.manifest.provenance == FootprintProvenance::McpSelfDeclared {
219        return None;
220    }
221    if approval != ApprovalNeed::Auto {
222        return None;
223    }
224    if !footprint.writes.is_empty() || footprint.spawns != SpawnClass::None {
225        return None;
226    }
227    if !input.supports_parallel_hint {
228        return None;
229    }
230    Some(ParallelResourceKey::ReadOnlyBatch)
231}
232
233#[cfg(test)]
234mod tests {
235    use super::*;
236    use crate::ToolCapability;
237
238    fn read_only_manifest(name: &str, provenance: FootprintProvenance) -> ToolManifest {
239        ToolManifest::derive_conservative(name, &[ToolCapability::ReadOnly], false, provenance)
240    }
241
242    fn default_input(manifest: ToolManifest) -> PolicyInput {
243        PolicyInput {
244            session_mode: PolicySessionMode::Agent,
245            manifest,
246            legacy_approval: ApprovalRequirement::Auto,
247            supports_parallel_hint: true,
248            trust_mode: false,
249        }
250    }
251
252    #[test]
253    fn builtin_read_only_tool_can_parallelize() {
254        let decision = PolicyEngine::decide(&default_input(read_only_manifest(
255            "read_file",
256            FootprintProvenance::BuiltIn,
257        )));
258        assert_eq!(decision.approval, ApprovalNeed::Auto);
259        assert_eq!(
260            decision.parallel_key,
261            Some(ParallelResourceKey::ReadOnlyBatch)
262        );
263        let meta = decision.plan_meta();
264        assert!(meta.read_only);
265        assert!(meta.supports_parallel);
266    }
267
268    #[test]
269    fn mcp_self_declared_read_only_still_requires_approval() {
270        let manifest = read_only_manifest("mcp_evil_read", FootprintProvenance::McpSelfDeclared);
271        let decision = PolicyEngine::decide(&PolicyInput {
272            legacy_approval: ApprovalRequirement::Auto,
273            supports_parallel_hint: true,
274            ..default_input(manifest)
275        });
276        assert_eq!(decision.approval, ApprovalNeed::Ask);
277        assert_eq!(decision.sandbox, SandboxClass::Strict);
278        assert!(decision.parallel_key.is_none());
279        let meta = decision.plan_meta();
280        assert!(meta.approval_required);
281        assert!(!meta.read_only);
282        assert!(!meta.supports_parallel);
283    }
284
285    #[test]
286    fn mcp_self_declared_write_is_not_relaxed() {
287        let manifest = ToolManifest::derive_conservative(
288            "mcp_writer",
289            &[ToolCapability::Network, ToolCapability::RequiresApproval],
290            false,
291            FootprintProvenance::McpSelfDeclared,
292        );
293        let decision = PolicyEngine::decide(&default_input(manifest));
294        assert_eq!(decision.approval, ApprovalNeed::Ask);
295        assert!(decision.parallel_key.is_none());
296    }
297
298    #[test]
299    fn trust_mode_skips_approval() {
300        let manifest = read_only_manifest("mcp_evil", FootprintProvenance::McpSelfDeclared);
301        let decision = PolicyEngine::decide(&PolicyInput {
302            trust_mode: true,
303            ..default_input(manifest)
304        });
305        assert_eq!(decision.approval, ApprovalNeed::Auto);
306    }
307
308    #[test]
309    fn write_footprint_never_parallelizes() {
310        let manifest = ToolManifest::derive_conservative(
311            "write_file",
312            &[ToolCapability::WritesFiles],
313            true,
314            FootprintProvenance::BuiltIn,
315        );
316        let decision = PolicyEngine::decide(&PolicyInput {
317            legacy_approval: ApprovalRequirement::Suggest,
318            supports_parallel_hint: false,
319            ..default_input(manifest)
320        });
321        assert_eq!(decision.approval, ApprovalNeed::Ask);
322        assert!(decision.parallel_key.is_none());
323    }
324
325    #[test]
326    fn shadow_stats_increment_on_diff() {
327        let before = policy_shadow_stats();
328        record_policy_shadow_diff(
329            "read_file",
330            &PolicyPlanMeta {
331                approval_required: false,
332                read_only: true,
333                supports_parallel: true,
334            },
335            &PolicyPlanMeta {
336                approval_required: true,
337                read_only: false,
338                supports_parallel: false,
339            },
340        );
341        let after = policy_shadow_stats();
342        assert_eq!(after.comparisons, before.comparisons + 1);
343        assert_eq!(after.diffs, before.diffs + 1);
344    }
345}