1use chrono::{DateTime, Utc};
5use serde::{Deserialize, Serialize};
6use std::fmt;
7use std::path::Path;
8
9use super::context_field::{ContextItemId, ContextState, ViewKind};
10
11#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
16pub struct OverlayId(pub String);
17
18impl OverlayId {
19 pub fn generate(target: &ContextItemId) -> Self {
20 Self(format!(
21 "ov_{}_{}",
22 target.as_str(),
23 Utc::now().timestamp_millis()
24 ))
25 }
26
27 pub fn as_str(&self) -> &str {
28 &self.0
29 }
30}
31
32impl fmt::Display for OverlayId {
33 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
34 f.write_str(&self.0)
35 }
36}
37
38#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
39#[serde(tag = "type", rename_all = "snake_case")]
40pub enum OverlayOp {
41 Include,
42 Exclude { reason: String },
43 Pin { verbatim: bool },
44 Unpin,
45 Rewrite { content: String },
46 SetView(ViewKind),
47 SetPriority { set_priority: f64 },
48 MarkOutdated,
49 Expire { after_secs: u64 },
50}
51
52impl OverlayOp {
53 fn discriminant(&self) -> &'static str {
54 match self {
55 Self::Include => "include",
56 Self::Exclude { .. } => "exclude",
57 Self::Pin { .. } => "pin",
58 Self::Unpin => "unpin",
59 Self::Rewrite { .. } => "rewrite",
60 Self::SetView(_) => "set_view",
61 Self::SetPriority { .. } => "set_priority",
62 Self::MarkOutdated => "mark_outdated",
63 Self::Expire { .. } => "expire",
64 }
65 }
66}
67
68#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
69#[serde(rename_all = "snake_case")]
70pub enum OverlayScope {
71 Call,
72 Session,
73 Project,
74 Agent(String),
75 Global,
76}
77
78#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
79#[serde(rename_all = "snake_case")]
80pub enum OverlayAuthor {
81 User,
82 Policy(String),
83 Agent(String),
84}
85
86#[derive(Debug, Clone, Serialize, Deserialize)]
91pub struct ContextOverlay {
92 pub id: OverlayId,
93 pub target: ContextItemId,
94 pub operation: OverlayOp,
95 pub scope: OverlayScope,
96 pub before_hash: String,
97 pub author: OverlayAuthor,
98 pub created_at: DateTime<Utc>,
99 pub stale: bool,
100}
101
102impl ContextOverlay {
103 pub fn new(
104 target: ContextItemId,
105 operation: OverlayOp,
106 scope: OverlayScope,
107 before_hash: String,
108 author: OverlayAuthor,
109 ) -> Self {
110 Self {
111 id: OverlayId::generate(&target),
112 target,
113 operation,
114 scope,
115 before_hash,
116 author,
117 created_at: Utc::now(),
118 stale: false,
119 }
120 }
121
122 fn is_expired(&self) -> bool {
123 if let OverlayOp::Expire { after_secs } = &self.operation {
124 let elapsed = Utc::now()
125 .signed_duration_since(self.created_at)
126 .num_seconds();
127 elapsed >= *after_secs as i64
128 } else {
129 false
130 }
131 }
132}
133
134const OVERLAY_FILE: &str = ".lean-ctx/overlays.json";
139
140#[derive(Debug, Clone, Default, Serialize, Deserialize)]
141pub struct OverlayStore {
142 overlays: Vec<ContextOverlay>,
143}
144
145impl OverlayStore {
146 pub fn new() -> Self {
147 Self::default()
148 }
149
150 pub fn add(&mut self, overlay: ContextOverlay) {
153 let disc = overlay.operation.discriminant();
154 self.overlays.retain(|existing| {
155 !(existing.target == overlay.target && existing.operation.discriminant() == disc)
156 });
157 self.overlays.push(overlay);
158 }
159
160 pub fn remove(&mut self, id: &OverlayId) {
161 self.overlays.retain(|o| o.id != *id);
162 }
163
164 pub fn for_item(&self, target: &ContextItemId) -> Vec<&ContextOverlay> {
165 self.overlays
166 .iter()
167 .filter(|o| o.target == *target)
168 .collect()
169 }
170
171 pub fn active_for_scope(&self, scope: &OverlayScope) -> Vec<&ContextOverlay> {
172 self.overlays.iter().filter(|o| o.scope == *scope).collect()
173 }
174
175 pub fn apply_to_state(
178 &self,
179 target: &ContextItemId,
180 current_state: ContextState,
181 ) -> ContextState {
182 let mut state = current_state;
183 for overlay in self.overlays.iter().filter(|o| o.target == *target) {
184 state = match &overlay.operation {
185 OverlayOp::Include => ContextState::Included,
186 OverlayOp::Exclude { .. } => ContextState::Excluded,
187 OverlayOp::Pin { .. } => ContextState::Pinned,
188 OverlayOp::Unpin => ContextState::Candidate,
189 OverlayOp::MarkOutdated => ContextState::Stale,
190 _ => state,
191 };
192 }
193 state
194 }
195
196 pub fn mark_stale_by_hash(&mut self, target: &ContextItemId, new_hash: &str) {
198 for overlay in self.overlays.iter_mut().filter(|o| o.target == *target) {
199 if overlay.before_hash != new_hash {
200 overlay.stale = true;
201 }
202 }
203 }
204
205 pub fn prune_expired(&mut self) {
207 self.overlays.retain(|o| !o.is_expired());
208 }
209
210 pub fn history(&self, target: &ContextItemId) -> Vec<&ContextOverlay> {
212 let mut items: Vec<&ContextOverlay> = self.for_item(target);
213 items.sort_by_key(|o| o.created_at);
214 items
215 }
216
217 pub fn remove_for_item(&mut self, target: &ContextItemId) {
218 self.overlays.retain(|o| o.target != *target);
219 }
220
221 pub fn all(&self) -> &[ContextOverlay] {
222 &self.overlays
223 }
224
225 pub fn save_project(&self, project_root: &Path) -> Result<(), String> {
226 let path = project_root.join(OVERLAY_FILE);
227 let json =
228 serde_json::to_string_pretty(self).map_err(|e| format!("serialize overlays: {e}"))?;
229 crate::config_io::write_atomic(&path, &json)
230 }
231
232 pub fn load_project(project_root: &Path) -> Self {
233 let path = project_root.join(OVERLAY_FILE);
234 let mut store: Self = std::fs::read_to_string(&path)
235 .ok()
236 .and_then(|s| serde_json::from_str(&s).ok())
237 .unwrap_or_default();
238 let now = chrono::Utc::now();
239 let session_ttl = chrono::Duration::hours(24);
240 store.overlays.retain(|o| match &o.scope {
241 OverlayScope::Session | OverlayScope::Call => {
242 now.signed_duration_since(o.created_at) < session_ttl
243 }
244 _ => true,
245 });
246 store
247 }
248}
249
250#[cfg(test)]
255mod tests {
256 use super::*;
257
258 fn make_target() -> ContextItemId {
259 ContextItemId::from_file("src/main.rs")
260 }
261
262 fn make_overlay(op: OverlayOp) -> ContextOverlay {
263 ContextOverlay::new(
264 make_target(),
265 op,
266 OverlayScope::Session,
267 "abc123".into(),
268 OverlayAuthor::User,
269 )
270 }
271
272 #[test]
275 fn exclude_sets_excluded_state() {
276 let mut store = OverlayStore::new();
277 store.add(make_overlay(OverlayOp::Exclude {
278 reason: "too large".into(),
279 }));
280 let state = store.apply_to_state(&make_target(), ContextState::Candidate);
281 assert_eq!(state, ContextState::Excluded);
282 }
283
284 #[test]
285 fn include_sets_included_state() {
286 let mut store = OverlayStore::new();
287 store.add(make_overlay(OverlayOp::Include));
288 let state = store.apply_to_state(&make_target(), ContextState::Candidate);
289 assert_eq!(state, ContextState::Included);
290 }
291
292 #[test]
293 fn pin_sets_pinned_state() {
294 let mut store = OverlayStore::new();
295 store.add(make_overlay(OverlayOp::Pin { verbatim: true }));
296 let state = store.apply_to_state(&make_target(), ContextState::Candidate);
297 assert_eq!(state, ContextState::Pinned);
298 }
299
300 #[test]
301 fn unpin_resets_to_candidate() {
302 let mut store = OverlayStore::new();
303 store.add(make_overlay(OverlayOp::Unpin));
304 let state = store.apply_to_state(&make_target(), ContextState::Pinned);
305 assert_eq!(state, ContextState::Candidate);
306 }
307
308 #[test]
309 fn mark_outdated_sets_stale_state() {
310 let mut store = OverlayStore::new();
311 store.add(make_overlay(OverlayOp::MarkOutdated));
312 let state = store.apply_to_state(&make_target(), ContextState::Included);
313 assert_eq!(state, ContextState::Stale);
314 }
315
316 #[test]
317 fn non_state_ops_preserve_current_state() {
318 let mut store = OverlayStore::new();
319 store.add(make_overlay(OverlayOp::SetPriority { set_priority: 0.9 }));
320 let state = store.apply_to_state(&make_target(), ContextState::Included);
321 assert_eq!(state, ContextState::Included);
322 }
323
324 #[test]
327 fn mark_stale_when_hash_changes() {
328 let mut store = OverlayStore::new();
329 store.add(make_overlay(OverlayOp::Include));
330 assert!(!store.overlays[0].stale);
331
332 store.mark_stale_by_hash(&make_target(), "different_hash");
333 assert!(store.overlays[0].stale);
334 }
335
336 #[test]
337 fn no_stale_when_hash_matches() {
338 let mut store = OverlayStore::new();
339 store.add(make_overlay(OverlayOp::Include));
340 store.mark_stale_by_hash(&make_target(), "abc123");
341 assert!(!store.overlays[0].stale);
342 }
343
344 #[test]
347 fn active_for_scope_filters_correctly() {
348 let mut store = OverlayStore::new();
349 store.add(make_overlay(OverlayOp::Include));
350 store.add(ContextOverlay::new(
351 ContextItemId::from_file("other.rs"),
352 OverlayOp::Include,
353 OverlayScope::Project,
354 "xyz".into(),
355 OverlayAuthor::User,
356 ));
357
358 let session = store.active_for_scope(&OverlayScope::Session);
359 assert_eq!(session.len(), 1);
360
361 let project = store.active_for_scope(&OverlayScope::Project);
362 assert_eq!(project.len(), 1);
363
364 let global = store.active_for_scope(&OverlayScope::Global);
365 assert!(global.is_empty());
366 }
367
368 #[test]
371 fn prune_removes_expired_overlays() {
372 let mut store = OverlayStore::new();
373 let mut expired = make_overlay(OverlayOp::Expire { after_secs: 0 });
374 expired.created_at = Utc::now() - chrono::Duration::seconds(10);
375 store.add(expired);
376 store.add(make_overlay(OverlayOp::Include));
377
378 assert_eq!(store.overlays.len(), 2);
379 store.prune_expired();
380 assert_eq!(store.overlays.len(), 1);
381 }
382
383 #[test]
384 fn prune_keeps_unexpired_overlays() {
385 let mut store = OverlayStore::new();
386 store.add(make_overlay(OverlayOp::Expire { after_secs: 99999 }));
387 store.prune_expired();
388 assert_eq!(store.overlays.len(), 1);
389 }
390
391 #[test]
394 fn save_and_load_roundtrip() {
395 let dir = tempfile::tempdir().expect("tmp dir");
396 let root = dir.path();
397
398 let mut store = OverlayStore::new();
399 store.add(make_overlay(OverlayOp::Include));
400 store.add(make_overlay(OverlayOp::Exclude {
401 reason: "noise".into(),
402 }));
403 store.add(make_overlay(OverlayOp::SetView(ViewKind::Signatures)));
404
405 store.save_project(root).expect("save");
406 let loaded = OverlayStore::load_project(root);
407 assert_eq!(loaded.overlays.len(), store.overlays.len());
408 }
409
410 #[test]
411 fn load_missing_file_returns_empty() {
412 let dir = tempfile::tempdir().expect("tmp dir");
413 let store = OverlayStore::load_project(dir.path());
414 assert!(store.overlays.is_empty());
415 }
416
417 #[test]
420 fn newer_overlay_replaces_same_target_and_op() {
421 let mut store = OverlayStore::new();
422 store.add(make_overlay(OverlayOp::Exclude {
423 reason: "first".into(),
424 }));
425 assert_eq!(store.overlays.len(), 1);
426 assert_eq!(
427 store.overlays[0].operation,
428 OverlayOp::Exclude {
429 reason: "first".into()
430 }
431 );
432
433 store.add(make_overlay(OverlayOp::Exclude {
434 reason: "second".into(),
435 }));
436 assert_eq!(store.overlays.len(), 1);
437 assert_eq!(
438 store.overlays[0].operation,
439 OverlayOp::Exclude {
440 reason: "second".into()
441 }
442 );
443 }
444
445 #[test]
446 fn different_ops_coexist_for_same_target() {
447 let mut store = OverlayStore::new();
448 store.add(make_overlay(OverlayOp::Include));
449 store.add(make_overlay(OverlayOp::SetPriority { set_priority: 0.8 }));
450 assert_eq!(store.overlays.len(), 2);
451 }
452
453 #[test]
456 fn history_returns_chronological_order() {
457 let mut store = OverlayStore::new();
458 let mut older = make_overlay(OverlayOp::Include);
459 older.created_at = Utc::now() - chrono::Duration::seconds(60);
460 store.overlays.push(older);
461
462 let newer = make_overlay(OverlayOp::SetPriority { set_priority: 0.5 });
463 store.overlays.push(newer);
464
465 let hist = store.history(&make_target());
466 assert_eq!(hist.len(), 2);
467 assert!(hist[0].created_at <= hist[1].created_at);
468 }
469
470 #[test]
473 fn remove_deletes_by_id() {
474 let mut store = OverlayStore::new();
475 let ov = make_overlay(OverlayOp::Include);
476 let id = ov.id.clone();
477 store.add(ov);
478 assert_eq!(store.overlays.len(), 1);
479
480 store.remove(&id);
481 assert!(store.overlays.is_empty());
482 }
483}