Skip to main content

lean_ctx/core/
context_policies.rs

1//! Context Policy Engine -- declarative rules for context governance.
2//!
3//! Extends the existing profile/role system with match-based policies
4//! that automatically include/exclude/pin/transform context items.
5//!
6//! Integrates with:
7//!   - io_boundary.rs (secret path detection)
8//!   - profiles.rs (compression/routing config)
9//!   - roles.rs (role-based access control)
10
11use serde::{Deserialize, Serialize};
12
13use super::context_field::{ContextState, ViewKind};
14
15/// A declarative context policy rule.
16#[derive(Debug, Clone, Serialize, Deserialize)]
17pub struct ContextPolicy {
18    pub name: String,
19    #[serde(rename = "match")]
20    pub match_pattern: String,
21    pub action: PolicyAction,
22    #[serde(default)]
23    pub condition: Option<PolicyCondition>,
24    #[serde(default)]
25    pub reason: Option<String>,
26}
27
28#[derive(Debug, Clone, Serialize, Deserialize)]
29#[serde(rename_all = "snake_case")]
30pub enum PolicyAction {
31    Exclude,
32    Include,
33    Pin,
34    SetView { view: String },
35    MaxTokens { limit: usize },
36    MarkOutdated,
37    Redact,
38    Audit,
39}
40
41#[derive(Debug, Clone, Serialize, Deserialize)]
42#[serde(rename_all = "snake_case")]
43pub enum PolicyCondition {
44    SourceSeenBefore,
45    SourceModifiedRecently,
46    TokensAbove { threshold: usize },
47    Always,
48    AgentIs { agent_id: String },
49    AgentRoleIs { role: String },
50    ContentContainsSecret,
51}
52
53/// A set of loaded policies.
54#[derive(Debug, Clone, Default, Serialize, Deserialize)]
55pub struct PolicySet {
56    pub policies: Vec<ContextPolicy>,
57}
58
59impl PolicySet {
60    pub fn new() -> Self {
61        Self::default()
62    }
63
64    /// Built-in default policies that align with existing LeanCTX behavior.
65    pub fn defaults() -> Self {
66        Self {
67            policies: vec![
68                ContextPolicy {
69                    name: "never_include_secrets".to_string(),
70                    match_pattern: "**/.env*".to_string(),
71                    action: PolicyAction::Exclude,
72                    condition: None,
73                    reason: Some("secrets".to_string()),
74                },
75                ContextPolicy {
76                    name: "exclude_private_keys".to_string(),
77                    match_pattern: "**/*private_key*".to_string(),
78                    action: PolicyAction::Exclude,
79                    condition: None,
80                    reason: Some("private key material".to_string()),
81                },
82                ContextPolicy {
83                    name: "exclude_credentials".to_string(),
84                    match_pattern: "**/credentials*".to_string(),
85                    action: PolicyAction::Exclude,
86                    condition: None,
87                    reason: Some("credentials".to_string()),
88                },
89                ContextPolicy {
90                    name: "delta_after_first_read".to_string(),
91                    match_pattern: "src/**".to_string(),
92                    action: PolicyAction::SetView {
93                        view: "diff".to_string(),
94                    },
95                    condition: Some(PolicyCondition::SourceSeenBefore),
96                    reason: Some("predictive coding: only send prediction errors".to_string()),
97                },
98                ContextPolicy {
99                    name: "compress_large_files".to_string(),
100                    match_pattern: "**/*".to_string(),
101                    action: PolicyAction::SetView {
102                        view: "signatures".to_string(),
103                    },
104                    condition: Some(PolicyCondition::TokensAbove { threshold: 8000 }),
105                    reason: Some("large file budget protection".to_string()),
106                },
107            ],
108        }
109    }
110
111    /// Evaluate all policies against a path, returning applicable actions.
112    pub fn evaluate(
113        &self,
114        path: &str,
115        seen_before: bool,
116        token_count: usize,
117    ) -> Vec<PolicyEvalResult> {
118        self.evaluate_full(path, seen_before, token_count, None, None, None)
119    }
120
121    /// Evaluate with full context including agent/role/content dimensions.
122    pub fn evaluate_full(
123        &self,
124        path: &str,
125        seen_before: bool,
126        token_count: usize,
127        agent_id: Option<&str>,
128        role: Option<&str>,
129        content: Option<&str>,
130    ) -> Vec<PolicyEvalResult> {
131        let mut results = Vec::new();
132        for policy in &self.policies {
133            if !path_matches(&policy.match_pattern, path) {
134                continue;
135            }
136            if let Some(ref condition) = policy.condition {
137                if !check_condition(
138                    condition,
139                    seen_before,
140                    token_count,
141                    path,
142                    agent_id,
143                    role,
144                    content,
145                ) {
146                    continue;
147                }
148            }
149            results.push(PolicyEvalResult {
150                policy_name: policy.name.clone(),
151                action: policy.action.clone(),
152                reason: policy.reason.clone().unwrap_or_else(|| policy.name.clone()),
153            });
154        }
155        results
156    }
157
158    /// Determine the effective state for an item after policy evaluation.
159    pub fn effective_state(
160        &self,
161        path: &str,
162        current: ContextState,
163        seen_before: bool,
164        token_count: usize,
165    ) -> ContextState {
166        let evals = self.evaluate(path, seen_before, token_count);
167        let mut state = current;
168        for eval in &evals {
169            match &eval.action {
170                PolicyAction::Exclude => state = ContextState::Excluded,
171                PolicyAction::Pin => state = ContextState::Pinned,
172                PolicyAction::Include => {
173                    if state == ContextState::Candidate {
174                        state = ContextState::Included;
175                    }
176                }
177                PolicyAction::MarkOutdated => state = ContextState::Stale,
178                PolicyAction::MaxTokens { limit } => {
179                    if token_count > *limit {
180                        state = ContextState::Excluded;
181                    }
182                }
183                PolicyAction::SetView { .. } | PolicyAction::Redact | PolicyAction::Audit => {}
184            }
185        }
186        state
187    }
188
189    /// Determine the recommended view for an item after policy evaluation.
190    pub fn recommended_view(
191        &self,
192        path: &str,
193        seen_before: bool,
194        token_count: usize,
195    ) -> Option<ViewKind> {
196        let evals = self.evaluate(path, seen_before, token_count);
197        for eval in evals.iter().rev() {
198            if let PolicyAction::SetView { view } = &eval.action {
199                return Some(ViewKind::parse(view));
200            }
201        }
202        None
203    }
204
205    /// Load policies from a project's .lean-ctx/policies.json file.
206    pub fn load_project(project_root: &std::path::Path) -> Self {
207        let path = project_root.join(".lean-ctx").join("policies.json");
208        std::fs::read_to_string(&path)
209            .ok()
210            .and_then(|s| serde_json::from_str(&s).ok())
211            .unwrap_or_else(Self::defaults)
212    }
213
214    /// Save policies to a project's .lean-ctx/policies.json file.
215    pub fn save_project(&self, project_root: &std::path::Path) -> Result<(), String> {
216        let dir = project_root.join(".lean-ctx");
217        std::fs::create_dir_all(&dir).map_err(|e| e.to_string())?;
218        let path = dir.join("policies.json");
219        let json = serde_json::to_string_pretty(self).map_err(|e| e.to_string())?;
220        crate::config_io::write_atomic(&path, &json)
221    }
222}
223
224#[derive(Debug, Clone)]
225pub struct PolicyEvalResult {
226    pub policy_name: String,
227    pub action: PolicyAction,
228    pub reason: String,
229}
230
231fn path_matches(pattern: &str, path: &str) -> bool {
232    if pattern == "**/*" {
233        return true;
234    }
235
236    if let Some(suffix) = pattern.strip_prefix("**/") {
237        if suffix.contains('*') {
238            let inner = suffix.replace('*', "");
239            return path.contains(&inner);
240        }
241        return path.contains(suffix) || path.ends_with(suffix);
242    }
243
244    if let Some(prefix) = pattern.strip_suffix("/**") {
245        return path.starts_with(prefix);
246    }
247
248    if pattern.contains("**") {
249        let parts: Vec<&str> = pattern.split("**").collect();
250        if parts.len() == 2 {
251            return path.starts_with(parts[0]) && path.ends_with(parts[1]);
252        }
253    }
254
255    if let Some(prefix) = pattern.strip_suffix('*') {
256        return path.starts_with(prefix);
257    }
258
259    path == pattern || path.ends_with(pattern)
260}
261
262fn check_condition(
263    condition: &PolicyCondition,
264    seen_before: bool,
265    token_count: usize,
266    path: &str,
267    agent_id: Option<&str>,
268    role: Option<&str>,
269    content: Option<&str>,
270) -> bool {
271    match condition {
272        PolicyCondition::SourceSeenBefore => seen_before,
273        PolicyCondition::TokensAbove { threshold } => token_count > *threshold,
274        PolicyCondition::SourceModifiedRecently => {
275            const RECENT_SECS: u64 = 3600;
276            std::fs::metadata(path)
277                .and_then(|m| m.modified())
278                .ok()
279                .and_then(|t| t.elapsed().ok())
280                .is_some_and(|elapsed| elapsed.as_secs() < RECENT_SECS)
281        }
282        PolicyCondition::Always => true,
283        PolicyCondition::AgentIs { agent_id: expected } => {
284            agent_id.is_some_and(|id| id == expected)
285        }
286        PolicyCondition::AgentRoleIs { role: expected } => role.is_some_and(|r| r == expected),
287        PolicyCondition::ContentContainsSecret => {
288            content.is_some_and(|c| !crate::core::secret_detection::detect_secrets(c).is_empty())
289        }
290    }
291}
292
293#[cfg(test)]
294mod tests {
295    use super::*;
296
297    #[test]
298    fn default_policies_exclude_env_files() {
299        let ps = PolicySet::defaults();
300        let results = ps.evaluate(".env", false, 100);
301        assert!(
302            results
303                .iter()
304                .any(|r| matches!(r.action, PolicyAction::Exclude)),
305            "should exclude .env files"
306        );
307    }
308
309    #[test]
310    fn default_policies_exclude_private_keys() {
311        let ps = PolicySet::defaults();
312        let results = ps.evaluate("secrets/private_key.pem", false, 100);
313        assert!(
314            results
315                .iter()
316                .any(|r| matches!(r.action, PolicyAction::Exclude)),
317            "should exclude private key files"
318        );
319    }
320
321    #[test]
322    fn delta_policy_only_when_seen_before() {
323        let ps = PolicySet::defaults();
324        let first = ps.evaluate("src/main.rs", false, 500);
325        let second = ps.evaluate("src/main.rs", true, 500);
326        assert!(
327            !first
328                .iter()
329                .any(|r| matches!(&r.action, PolicyAction::SetView { view } if view == "diff")),
330            "should NOT suggest diff on first read"
331        );
332        assert!(
333            second
334                .iter()
335                .any(|r| matches!(&r.action, PolicyAction::SetView { view } if view == "diff")),
336            "should suggest diff on subsequent read"
337        );
338    }
339
340    #[test]
341    fn large_file_policy_triggers_above_threshold() {
342        let ps = PolicySet::defaults();
343        let small = ps.evaluate("src/main.rs", false, 500);
344        let large = ps.evaluate("src/main.rs", false, 10000);
345        assert!(!small
346            .iter()
347            .any(|r| matches!(&r.action, PolicyAction::SetView { view } if view == "signatures")),);
348        assert!(large
349            .iter()
350            .any(|r| matches!(&r.action, PolicyAction::SetView { view } if view == "signatures")),);
351    }
352
353    #[test]
354    fn effective_state_excludes_secrets() {
355        let ps = PolicySet::defaults();
356        let state = ps.effective_state(".env.local", ContextState::Candidate, false, 100);
357        assert_eq!(state, ContextState::Excluded);
358    }
359
360    #[test]
361    fn recommended_view_for_seen_file() {
362        let ps = PolicySet::defaults();
363        let view = ps.recommended_view("src/main.rs", true, 500);
364        assert_eq!(view, Some(ViewKind::Diff));
365    }
366
367    #[test]
368    fn recommended_view_none_for_new_file() {
369        let ps = PolicySet::defaults();
370        let view = ps.recommended_view("src/main.rs", false, 500);
371        assert!(view.is_none() || view == Some(ViewKind::Diff),);
372    }
373
374    #[test]
375    fn path_matches_glob_patterns() {
376        assert!(path_matches("**/.env*", ".env"));
377        assert!(path_matches("**/.env*", ".env.local"));
378        assert!(path_matches("**/.env*", "config/.env.prod"));
379        assert!(path_matches("src/**", "src/main.rs"));
380        assert!(path_matches("src/**", "src/core/mod.rs"));
381        assert!(path_matches("**/*", "anything.txt"));
382        assert!(!path_matches("src/**", "tests/test.rs"));
383    }
384
385    #[test]
386    fn empty_policy_set_changes_nothing() {
387        let ps = PolicySet::new();
388        let state = ps.effective_state("src/main.rs", ContextState::Included, false, 100);
389        assert_eq!(state, ContextState::Included);
390    }
391
392    #[test]
393    fn custom_policy_works() {
394        let ps = PolicySet {
395            policies: vec![ContextPolicy {
396                name: "pin_readme".to_string(),
397                match_pattern: "README.md".to_string(),
398                action: PolicyAction::Pin,
399                condition: None,
400                reason: None,
401            }],
402        };
403        let state = ps.effective_state("README.md", ContextState::Candidate, false, 100);
404        assert_eq!(state, ContextState::Pinned);
405    }
406}