1use serde::{Deserialize, Serialize};
12
13use super::context_field::{ContextState, ViewKind};
14
15#[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#[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 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 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 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 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 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 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 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}