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        if crate::core::pathutil::is_data_dir_collision(project_root) {
208            return Self::defaults();
209        }
210        let path = project_root.join(".lean-ctx").join("policies.json");
211        std::fs::read_to_string(&path)
212            .ok()
213            .and_then(|s| serde_json::from_str(&s).ok())
214            .unwrap_or_else(Self::defaults)
215    }
216
217    /// Save policies to a project's .lean-ctx/policies.json file.
218    pub fn save_project(&self, project_root: &std::path::Path) -> Result<(), String> {
219        let dir = crate::core::pathutil::safe_project_data_dir(project_root)?;
220        std::fs::create_dir_all(&dir).map_err(|e| e.to_string())?;
221        let path = dir.join("policies.json");
222        let json = serde_json::to_string_pretty(self).map_err(|e| e.to_string())?;
223        crate::config_io::write_atomic(&path, &json)
224    }
225}
226
227#[derive(Debug, Clone)]
228pub struct PolicyEvalResult {
229    pub policy_name: String,
230    pub action: PolicyAction,
231    pub reason: String,
232}
233
234fn path_matches(pattern: &str, path: &str) -> bool {
235    if pattern == "**/*" {
236        return true;
237    }
238
239    if let Some(suffix) = pattern.strip_prefix("**/") {
240        if suffix.contains('*') {
241            let inner = suffix.replace('*', "");
242            return path.contains(&inner);
243        }
244        return path.contains(suffix) || path.ends_with(suffix);
245    }
246
247    if let Some(prefix) = pattern.strip_suffix("/**") {
248        return path.starts_with(prefix);
249    }
250
251    if pattern.contains("**") {
252        let parts: Vec<&str> = pattern.split("**").collect();
253        if parts.len() == 2 {
254            return path.starts_with(parts[0]) && path.ends_with(parts[1]);
255        }
256    }
257
258    if let Some(prefix) = pattern.strip_suffix('*') {
259        return path.starts_with(prefix);
260    }
261
262    path == pattern || path.ends_with(pattern)
263}
264
265fn check_condition(
266    condition: &PolicyCondition,
267    seen_before: bool,
268    token_count: usize,
269    path: &str,
270    agent_id: Option<&str>,
271    role: Option<&str>,
272    content: Option<&str>,
273) -> bool {
274    match condition {
275        PolicyCondition::SourceSeenBefore => seen_before,
276        PolicyCondition::TokensAbove { threshold } => token_count > *threshold,
277        PolicyCondition::SourceModifiedRecently => {
278            const RECENT_SECS: u64 = 3600;
279            std::fs::metadata(path)
280                .and_then(|m| m.modified())
281                .ok()
282                .and_then(|t| t.elapsed().ok())
283                .is_some_and(|elapsed| elapsed.as_secs() < RECENT_SECS)
284        }
285        PolicyCondition::Always => true,
286        PolicyCondition::AgentIs { agent_id: expected } => {
287            agent_id.is_some_and(|id| id == expected)
288        }
289        PolicyCondition::AgentRoleIs { role: expected } => role.is_some_and(|r| r == expected),
290        PolicyCondition::ContentContainsSecret => {
291            content.is_some_and(|c| !crate::core::secret_detection::detect_secrets(c).is_empty())
292        }
293    }
294}
295
296#[cfg(test)]
297mod tests {
298    use super::*;
299
300    #[test]
301    fn default_policies_exclude_env_files() {
302        let ps = PolicySet::defaults();
303        let results = ps.evaluate(".env", false, 100);
304        assert!(
305            results
306                .iter()
307                .any(|r| matches!(r.action, PolicyAction::Exclude)),
308            "should exclude .env files"
309        );
310    }
311
312    #[test]
313    fn default_policies_exclude_private_keys() {
314        let ps = PolicySet::defaults();
315        let results = ps.evaluate("secrets/private_key.pem", false, 100);
316        assert!(
317            results
318                .iter()
319                .any(|r| matches!(r.action, PolicyAction::Exclude)),
320            "should exclude private key files"
321        );
322    }
323
324    #[test]
325    fn delta_policy_only_when_seen_before() {
326        let ps = PolicySet::defaults();
327        let first = ps.evaluate("src/main.rs", false, 500);
328        let second = ps.evaluate("src/main.rs", true, 500);
329        assert!(
330            !first
331                .iter()
332                .any(|r| matches!(&r.action, PolicyAction::SetView { view } if view == "diff")),
333            "should NOT suggest diff on first read"
334        );
335        assert!(
336            second
337                .iter()
338                .any(|r| matches!(&r.action, PolicyAction::SetView { view } if view == "diff")),
339            "should suggest diff on subsequent read"
340        );
341    }
342
343    #[test]
344    fn large_file_policy_triggers_above_threshold() {
345        let ps = PolicySet::defaults();
346        let small = ps.evaluate("src/main.rs", false, 500);
347        let large = ps.evaluate("src/main.rs", false, 10000);
348        assert!(!small
349            .iter()
350            .any(|r| matches!(&r.action, PolicyAction::SetView { view } if view == "signatures")),);
351        assert!(large
352            .iter()
353            .any(|r| matches!(&r.action, PolicyAction::SetView { view } if view == "signatures")),);
354    }
355
356    #[test]
357    fn effective_state_excludes_secrets() {
358        let ps = PolicySet::defaults();
359        let state = ps.effective_state(".env.local", ContextState::Candidate, false, 100);
360        assert_eq!(state, ContextState::Excluded);
361    }
362
363    #[test]
364    fn recommended_view_for_seen_file() {
365        let ps = PolicySet::defaults();
366        let view = ps.recommended_view("src/main.rs", true, 500);
367        assert_eq!(view, Some(ViewKind::Diff));
368    }
369
370    #[test]
371    fn recommended_view_none_for_new_file() {
372        let ps = PolicySet::defaults();
373        let view = ps.recommended_view("src/main.rs", false, 500);
374        assert!(view.is_none() || view == Some(ViewKind::Diff),);
375    }
376
377    #[test]
378    fn path_matches_glob_patterns() {
379        assert!(path_matches("**/.env*", ".env"));
380        assert!(path_matches("**/.env*", ".env.local"));
381        assert!(path_matches("**/.env*", "config/.env.prod"));
382        assert!(path_matches("src/**", "src/main.rs"));
383        assert!(path_matches("src/**", "src/core/mod.rs"));
384        assert!(path_matches("**/*", "anything.txt"));
385        assert!(!path_matches("src/**", "tests/test.rs"));
386    }
387
388    #[test]
389    fn empty_policy_set_changes_nothing() {
390        let ps = PolicySet::new();
391        let state = ps.effective_state("src/main.rs", ContextState::Included, false, 100);
392        assert_eq!(state, ContextState::Included);
393    }
394
395    #[test]
396    fn custom_policy_works() {
397        let ps = PolicySet {
398            policies: vec![ContextPolicy {
399                name: "pin_readme".to_string(),
400                match_pattern: "README.md".to_string(),
401                action: PolicyAction::Pin,
402                condition: None,
403                reason: None,
404            }],
405        };
406        let state = ps.effective_state("README.md", ContextState::Candidate, false, 100);
407        assert_eq!(state, ContextState::Pinned);
408    }
409}