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(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 std::fs::read_to_string(&path)
235 .ok()
236 .and_then(|s| serde_json::from_str(&s).ok())
237 .unwrap_or_default()
238 }
239}
240
241#[cfg(test)]
246mod tests {
247 use super::*;
248
249 fn make_target() -> ContextItemId {
250 ContextItemId::from_file("src/main.rs")
251 }
252
253 fn make_overlay(op: OverlayOp) -> ContextOverlay {
254 ContextOverlay::new(
255 make_target(),
256 op,
257 OverlayScope::Session,
258 "abc123".into(),
259 OverlayAuthor::User,
260 )
261 }
262
263 #[test]
266 fn exclude_sets_excluded_state() {
267 let mut store = OverlayStore::new();
268 store.add(make_overlay(OverlayOp::Exclude {
269 reason: "too large".into(),
270 }));
271 let state = store.apply_to_state(&make_target(), ContextState::Candidate);
272 assert_eq!(state, ContextState::Excluded);
273 }
274
275 #[test]
276 fn include_sets_included_state() {
277 let mut store = OverlayStore::new();
278 store.add(make_overlay(OverlayOp::Include));
279 let state = store.apply_to_state(&make_target(), ContextState::Candidate);
280 assert_eq!(state, ContextState::Included);
281 }
282
283 #[test]
284 fn pin_sets_pinned_state() {
285 let mut store = OverlayStore::new();
286 store.add(make_overlay(OverlayOp::Pin { verbatim: true }));
287 let state = store.apply_to_state(&make_target(), ContextState::Candidate);
288 assert_eq!(state, ContextState::Pinned);
289 }
290
291 #[test]
292 fn unpin_resets_to_candidate() {
293 let mut store = OverlayStore::new();
294 store.add(make_overlay(OverlayOp::Unpin));
295 let state = store.apply_to_state(&make_target(), ContextState::Pinned);
296 assert_eq!(state, ContextState::Candidate);
297 }
298
299 #[test]
300 fn mark_outdated_sets_stale_state() {
301 let mut store = OverlayStore::new();
302 store.add(make_overlay(OverlayOp::MarkOutdated));
303 let state = store.apply_to_state(&make_target(), ContextState::Included);
304 assert_eq!(state, ContextState::Stale);
305 }
306
307 #[test]
308 fn non_state_ops_preserve_current_state() {
309 let mut store = OverlayStore::new();
310 store.add(make_overlay(OverlayOp::SetPriority(0.9)));
311 let state = store.apply_to_state(&make_target(), ContextState::Included);
312 assert_eq!(state, ContextState::Included);
313 }
314
315 #[test]
318 fn mark_stale_when_hash_changes() {
319 let mut store = OverlayStore::new();
320 store.add(make_overlay(OverlayOp::Include));
321 assert!(!store.overlays[0].stale);
322
323 store.mark_stale_by_hash(&make_target(), "different_hash");
324 assert!(store.overlays[0].stale);
325 }
326
327 #[test]
328 fn no_stale_when_hash_matches() {
329 let mut store = OverlayStore::new();
330 store.add(make_overlay(OverlayOp::Include));
331 store.mark_stale_by_hash(&make_target(), "abc123");
332 assert!(!store.overlays[0].stale);
333 }
334
335 #[test]
338 fn active_for_scope_filters_correctly() {
339 let mut store = OverlayStore::new();
340 store.add(make_overlay(OverlayOp::Include));
341 store.add(ContextOverlay::new(
342 ContextItemId::from_file("other.rs"),
343 OverlayOp::Include,
344 OverlayScope::Project,
345 "xyz".into(),
346 OverlayAuthor::User,
347 ));
348
349 let session = store.active_for_scope(&OverlayScope::Session);
350 assert_eq!(session.len(), 1);
351
352 let project = store.active_for_scope(&OverlayScope::Project);
353 assert_eq!(project.len(), 1);
354
355 let global = store.active_for_scope(&OverlayScope::Global);
356 assert!(global.is_empty());
357 }
358
359 #[test]
362 fn prune_removes_expired_overlays() {
363 let mut store = OverlayStore::new();
364 let mut expired = make_overlay(OverlayOp::Expire { after_secs: 0 });
365 expired.created_at = Utc::now() - chrono::Duration::seconds(10);
366 store.add(expired);
367 store.add(make_overlay(OverlayOp::Include));
368
369 assert_eq!(store.overlays.len(), 2);
370 store.prune_expired();
371 assert_eq!(store.overlays.len(), 1);
372 }
373
374 #[test]
375 fn prune_keeps_unexpired_overlays() {
376 let mut store = OverlayStore::new();
377 store.add(make_overlay(OverlayOp::Expire { after_secs: 99999 }));
378 store.prune_expired();
379 assert_eq!(store.overlays.len(), 1);
380 }
381
382 #[test]
385 fn save_and_load_roundtrip() {
386 let dir = tempfile::tempdir().expect("tmp dir");
387 let root = dir.path();
388
389 let mut store = OverlayStore::new();
390 store.add(make_overlay(OverlayOp::Include));
391 store.add(make_overlay(OverlayOp::Exclude {
392 reason: "noise".into(),
393 }));
394 store.add(make_overlay(OverlayOp::SetView(ViewKind::Signatures)));
395
396 store.save_project(root).expect("save");
397 let loaded = OverlayStore::load_project(root);
398 assert_eq!(loaded.overlays.len(), store.overlays.len());
399 }
400
401 #[test]
402 fn load_missing_file_returns_empty() {
403 let dir = tempfile::tempdir().expect("tmp dir");
404 let store = OverlayStore::load_project(dir.path());
405 assert!(store.overlays.is_empty());
406 }
407
408 #[test]
411 fn newer_overlay_replaces_same_target_and_op() {
412 let mut store = OverlayStore::new();
413 store.add(make_overlay(OverlayOp::Exclude {
414 reason: "first".into(),
415 }));
416 assert_eq!(store.overlays.len(), 1);
417 assert_eq!(
418 store.overlays[0].operation,
419 OverlayOp::Exclude {
420 reason: "first".into()
421 }
422 );
423
424 store.add(make_overlay(OverlayOp::Exclude {
425 reason: "second".into(),
426 }));
427 assert_eq!(store.overlays.len(), 1);
428 assert_eq!(
429 store.overlays[0].operation,
430 OverlayOp::Exclude {
431 reason: "second".into()
432 }
433 );
434 }
435
436 #[test]
437 fn different_ops_coexist_for_same_target() {
438 let mut store = OverlayStore::new();
439 store.add(make_overlay(OverlayOp::Include));
440 store.add(make_overlay(OverlayOp::SetPriority(0.8)));
441 assert_eq!(store.overlays.len(), 2);
442 }
443
444 #[test]
447 fn history_returns_chronological_order() {
448 let mut store = OverlayStore::new();
449 let mut older = make_overlay(OverlayOp::Include);
450 older.created_at = Utc::now() - chrono::Duration::seconds(60);
451 store.overlays.push(older);
452
453 let newer = make_overlay(OverlayOp::SetPriority(0.5));
454 store.overlays.push(newer);
455
456 let hist = store.history(&make_target());
457 assert_eq!(hist.len(), 2);
458 assert!(hist[0].created_at <= hist[1].created_at);
459 }
460
461 #[test]
464 fn remove_deletes_by_id() {
465 let mut store = OverlayStore::new();
466 let ov = make_overlay(OverlayOp::Include);
467 let id = ov.id.clone();
468 store.add(ov);
469 assert_eq!(store.overlays.len(), 1);
470
471 store.remove(&id);
472 assert!(store.overlays.is_empty());
473 }
474}