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 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 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}