Skip to main content

lean_ctx/server/
context_gate.rs

1use crate::core::context_field::{ContextItemId, ContextState};
2use crate::core::context_ledger::{ContextLedger, PressureAction};
3use crate::core::context_overlay::{OverlayOp, OverlayStore};
4
5#[derive(Debug, Clone)]
6pub struct PreDispatchResult {
7    pub overridden_mode: Option<String>,
8    pub reason: Option<&'static str>,
9    pub pressure_downgraded: bool,
10}
11
12#[derive(Debug, Clone)]
13pub struct PostDispatchResult {
14    pub eviction_hint: Option<String>,
15    pub elicitation_hint: Option<String>,
16    pub resource_changed: bool,
17}
18
19pub fn pre_dispatch_read(
20    path: &str,
21    requested_mode: &str,
22    task: Option<&str>,
23    project_root: Option<&str>,
24    pressure: Option<&PressureAction>,
25) -> PreDispatchResult {
26    let no_change = PreDispatchResult {
27        overridden_mode: None,
28        reason: None,
29        pressure_downgraded: false,
30    };
31
32    if requested_mode == "diff" || requested_mode.starts_with("lines") {
33        return no_change;
34    }
35
36    if let Some(root) = project_root {
37        let overlay = OverlayStore::load_project(&std::path::PathBuf::from(root));
38        if let Some(result) = check_overlay_mode_override(path, requested_mode, &overlay) {
39            return result;
40        }
41    }
42
43    if let Some(action) = pressure {
44        if let Some(downgraded) = pressure_downgrade(requested_mode, action) {
45            return PreDispatchResult {
46                overridden_mode: Some(downgraded),
47                reason: Some("pressure-auto-downgrade"),
48                pressure_downgraded: true,
49            };
50        }
51    }
52
53    if requested_mode == "full" {
54        return no_change;
55    }
56
57    if let Ok(bt) = crate::core::bounce_tracker::global().lock() {
58        if bt.should_force_full(path) {
59            return PreDispatchResult {
60                overridden_mode: Some("full".to_string()),
61                reason: Some("bounce-prevention"),
62                pressure_downgraded: false,
63            };
64        }
65    }
66
67    if let Some(task_str) = task {
68        let intent = crate::core::intent_engine::StructuredIntent::from_query(task_str);
69        let norm = crate::core::pathutil::normalize_tool_path(path);
70        let is_target = intent
71            .targets
72            .iter()
73            .any(|t| norm.ends_with(t) || norm.contains(t));
74        if is_target {
75            return PreDispatchResult {
76                overridden_mode: Some("full".to_string()),
77                reason: Some("intent-target"),
78                pressure_downgraded: false,
79            };
80        }
81    }
82
83    if let Some(root) = project_root {
84        if let Some(index) = try_load_graph(root) {
85            let related = index.get_related(path, 1);
86            if let Some(task_str) = task {
87                let intent = crate::core::intent_engine::StructuredIntent::from_query(task_str);
88                for target in &intent.targets {
89                    let target_related = index.get_related(target, 1);
90                    let norm = crate::core::pathutil::normalize_tool_path(path);
91                    if target_related
92                        .iter()
93                        .any(|r| r.contains(&norm) || norm.contains(r))
94                    {
95                        return PreDispatchResult {
96                            overridden_mode: Some("map".to_string()),
97                            reason: Some("graph-direct-import"),
98                            pressure_downgraded: false,
99                        };
100                    }
101                }
102            }
103            if !related.is_empty() && requested_mode == "auto" {
104                let reverse_deps = index.get_reverse_deps(path, 1);
105                if reverse_deps.len() > 3 {
106                    return PreDispatchResult {
107                        overridden_mode: Some("map".to_string()),
108                        reason: Some("graph-hub-file"),
109                        pressure_downgraded: false,
110                    };
111                }
112            }
113        }
114    }
115
116    if let Some(root) = project_root {
117        if let Some(knowledge) = crate::core::knowledge::ProjectKnowledge::load(root) {
118            let norm = crate::core::pathutil::normalize_tool_path(path);
119            let mentions = knowledge
120                .facts
121                .iter()
122                .filter(|f| f.value.contains(&norm) || f.key.contains(&norm))
123                .count();
124            if mentions >= 3 {
125                return PreDispatchResult {
126                    overridden_mode: Some("map".to_string()),
127                    reason: Some("knowledge-high-relevance"),
128                    pressure_downgraded: false,
129                };
130            }
131        }
132    }
133
134    no_change
135}
136
137fn pressure_downgrade(requested_mode: &str, action: &PressureAction) -> Option<String> {
138    match action {
139        PressureAction::ForceCompression => match requested_mode {
140            "full" => Some("map".to_string()),
141            "map" => Some("signatures".to_string()),
142            _ => None,
143        },
144        PressureAction::EvictLeastRelevant => match requested_mode {
145            "full" => Some("map".to_string()),
146            "map" | "auto" => Some("signatures".to_string()),
147            _ => None,
148        },
149        PressureAction::NoAction | PressureAction::SuggestCompression => None,
150    }
151}
152
153fn check_overlay_mode_override(
154    path: &str,
155    requested_mode: &str,
156    overlay: &OverlayStore,
157) -> Option<PreDispatchResult> {
158    let item_id = ContextItemId::from_file(path);
159    let overlays = overlay.for_item(&item_id);
160
161    for ov in overlays.iter().rev() {
162        match &ov.operation {
163            OverlayOp::SetView(view) => {
164                let mode_str = view.as_str();
165                if mode_str != requested_mode {
166                    return Some(PreDispatchResult {
167                        overridden_mode: Some(mode_str.to_string()),
168                        reason: Some("overlay-set-view"),
169                        pressure_downgraded: false,
170                    });
171                }
172            }
173            OverlayOp::Pin { .. } if requested_mode != "full" => {
174                return Some(PreDispatchResult {
175                    overridden_mode: Some("full".to_string()),
176                    reason: Some("pinned"),
177                    pressure_downgraded: false,
178                });
179            }
180            OverlayOp::Exclude { .. } if requested_mode != "signatures" => {
181                return Some(PreDispatchResult {
182                    overridden_mode: Some("signatures".to_string()),
183                    reason: Some("excluded"),
184                    pressure_downgraded: false,
185                });
186            }
187            _ => {}
188        }
189    }
190    None
191}
192
193pub fn post_dispatch_record(
194    path: &str,
195    mode: &str,
196    original_tokens: usize,
197    sent_tokens: usize,
198    ledger: &mut ContextLedger,
199    overlay: &OverlayStore,
200) -> PostDispatchResult {
201    post_dispatch_record_with_task(
202        path,
203        mode,
204        original_tokens,
205        sent_tokens,
206        ledger,
207        overlay,
208        None,
209    )
210}
211
212pub fn post_dispatch_record_with_task(
213    path: &str,
214    mode: &str,
215    original_tokens: usize,
216    sent_tokens: usize,
217    ledger: &mut ContextLedger,
218    overlay: &OverlayStore,
219    task: Option<&str>,
220) -> PostDispatchResult {
221    let prev_count = ledger.entries.len();
222    let prev_pressure = ledger.pressure().recommendation;
223
224    ledger.record_with_task(path, mode, original_tokens, sent_tokens, task);
225
226    let item_id = ContextItemId::from_file(path);
227    let state = overlay.apply_to_state(&item_id, ContextState::Included);
228
229    if state == ContextState::Excluded {
230        return PostDispatchResult {
231            eviction_hint: Some(format!("File '{path}' is excluded by overlay.")),
232            elicitation_hint: None,
233            resource_changed: true,
234        };
235    }
236
237    let elicitation =
238        super::elicitation::check_elicitation_needed(ledger, Some(path), Some(sent_tokens))
239            .map(|s| s.format_fallback_hint());
240
241    let pressure = ledger.pressure();
242
243    apply_reinjection_plan(ledger, &pressure.recommendation);
244
245    let new_entry = ledger.entries.len() != prev_count;
246    let pressure_shifted = pressure.recommendation != prev_pressure;
247    let resource_changed = new_entry || pressure_shifted;
248
249    if pressure.utilization > 0.9 {
250        let candidates = ledger.eviction_candidates_by_phi(3);
251        if !candidates.is_empty() {
252            let names: Vec<_> = candidates
253                .iter()
254                .take(3)
255                .map(|p| crate::core::protocol::shorten_path(p))
256                .collect();
257            return PostDispatchResult {
258                eviction_hint: Some(format!(
259                    "Context pressure {:.0}%. Consider evicting: {}",
260                    pressure.utilization * 100.0,
261                    names.join(", ")
262                )),
263                elicitation_hint: elicitation,
264                resource_changed,
265            };
266        }
267    }
268
269    PostDispatchResult {
270        eviction_hint: None,
271        elicitation_hint: elicitation,
272        resource_changed,
273    }
274}
275
276fn apply_reinjection_plan(ledger: &mut ContextLedger, action: &PressureAction) {
277    if *action != PressureAction::ForceCompression && *action != PressureAction::EvictLeastRelevant
278    {
279        return;
280    }
281    for entry in &mut ledger.entries {
282        if entry.mode == "full" {
283            entry.mode = "map".to_string();
284        }
285    }
286}
287
288fn try_load_graph(project_root: &str) -> Option<crate::core::graph_index::ProjectIndex> {
289    crate::core::graph_index::ProjectIndex::load(project_root)
290}
291
292#[cfg(test)]
293mod tests {
294    use super::*;
295
296    #[test]
297    fn pre_dispatch_passthrough_for_full() {
298        let result = pre_dispatch_read("src/main.rs", "full", None, None, None);
299        assert!(result.overridden_mode.is_none());
300    }
301
302    #[test]
303    fn pre_dispatch_passthrough_for_diff() {
304        let result = pre_dispatch_read("src/main.rs", "diff", None, None, None);
305        assert!(result.overridden_mode.is_none());
306    }
307
308    #[test]
309    fn pre_dispatch_no_override_without_signals() {
310        let result = pre_dispatch_read("src/unknown.rs", "auto", None, None, None);
311        assert!(result.overridden_mode.is_none());
312    }
313
314    #[test]
315    fn pre_dispatch_bounce_prevention_forces_full() {
316        {
317            let mut bt = crate::core::bounce_tracker::global().lock().unwrap();
318            bt.set_seq(1);
319            bt.record_read("src/bouncy.yml", "map", 30, 400);
320            bt.set_seq(2);
321            bt.record_read("src/bouncy.yml", "full", 400, 400);
322            bt.set_seq(3);
323            bt.record_read("a2.yml", "map", 30, 400);
324            bt.set_seq(4);
325            bt.record_read("a2.yml", "full", 400, 400);
326            bt.set_seq(5);
327            bt.record_read("a3.yml", "map", 30, 400);
328            bt.set_seq(6);
329            bt.record_read("a3.yml", "full", 400, 400);
330        }
331        let result = pre_dispatch_read("new.yml", "auto", None, None, None);
332        assert_eq!(result.overridden_mode, Some("full".to_string()));
333        assert_eq!(result.reason, Some("bounce-prevention"));
334    }
335
336    #[test]
337    fn pressure_downgrade_full_to_map() {
338        let result = pre_dispatch_read(
339            "c.rs",
340            "full",
341            None,
342            None,
343            Some(&PressureAction::ForceCompression),
344        );
345        assert_eq!(result.overridden_mode, Some("map".to_string()));
346        assert_eq!(result.reason, Some("pressure-auto-downgrade"));
347        assert!(result.pressure_downgraded);
348    }
349
350    #[test]
351    fn pressure_downgrade_map_to_signatures_on_evict() {
352        let result = pre_dispatch_read(
353            "c.rs",
354            "map",
355            None,
356            None,
357            Some(&PressureAction::EvictLeastRelevant),
358        );
359        assert_eq!(result.overridden_mode, Some("signatures".to_string()));
360        assert!(result.pressure_downgraded);
361    }
362
363    #[test]
364    fn no_pressure_downgrade_when_low() {
365        let result = pre_dispatch_read("c.rs", "full", None, None, Some(&PressureAction::NoAction));
366        assert!(result.overridden_mode.is_none());
367        assert!(!result.pressure_downgraded);
368    }
369
370    #[test]
371    fn post_dispatch_reinjection_downgrades_entries() {
372        let mut ledger = ContextLedger::with_window_size(1000);
373        ledger.record("a.rs", "full", 400, 400);
374        ledger.record("b.rs", "full", 400, 400);
375        let overlay = OverlayStore::new();
376        let result = post_dispatch_record("c.rs", "full", 300, 300, &mut ledger, &overlay);
377        assert!(result.resource_changed);
378        let a_entry = ledger.entries.iter().find(|e| e.path == "a.rs").unwrap();
379        assert_eq!(a_entry.mode, "map");
380    }
381
382    #[test]
383    fn overlay_pin_forces_full_mode() {
384        let dir = tempfile::tempdir().expect("tmp dir");
385        let root = dir.path();
386        let mut store = OverlayStore::new();
387        let target = ContextItemId::from_file("src/important.rs");
388        store.add(crate::core::context_overlay::ContextOverlay::new(
389            target,
390            OverlayOp::Pin { verbatim: false },
391            crate::core::context_overlay::OverlayScope::Project,
392            String::new(),
393            crate::core::context_overlay::OverlayAuthor::User,
394        ));
395        store.save_project(root).unwrap();
396
397        let result = pre_dispatch_read(
398            "src/important.rs",
399            "auto",
400            None,
401            Some(root.to_str().unwrap()),
402            None,
403        );
404        assert_eq!(result.overridden_mode, Some("full".to_string()));
405        assert_eq!(result.reason, Some("pinned"));
406    }
407
408    #[test]
409    fn overlay_exclude_forces_signatures_mode() {
410        let dir = tempfile::tempdir().expect("tmp dir");
411        let root = dir.path();
412        let mut store = OverlayStore::new();
413        let target = ContextItemId::from_file("src/noisy.rs");
414        store.add(crate::core::context_overlay::ContextOverlay::new(
415            target,
416            OverlayOp::Exclude {
417                reason: "noise".to_string(),
418            },
419            crate::core::context_overlay::OverlayScope::Project,
420            String::new(),
421            crate::core::context_overlay::OverlayAuthor::User,
422        ));
423        store.save_project(root).unwrap();
424
425        let result = pre_dispatch_read(
426            "src/noisy.rs",
427            "auto",
428            None,
429            Some(root.to_str().unwrap()),
430            None,
431        );
432        assert_eq!(result.overridden_mode, Some("signatures".to_string()));
433        assert_eq!(result.reason, Some("excluded"));
434    }
435
436    #[test]
437    fn overlay_set_view_forces_specified_mode() {
438        let dir = tempfile::tempdir().expect("tmp dir");
439        let root = dir.path();
440        let mut store = OverlayStore::new();
441        let target = ContextItemId::from_file("src/big.rs");
442        store.add(crate::core::context_overlay::ContextOverlay::new(
443            target,
444            OverlayOp::SetView(crate::core::context_field::ViewKind::Map),
445            crate::core::context_overlay::OverlayScope::Project,
446            String::new(),
447            crate::core::context_overlay::OverlayAuthor::User,
448        ));
449        store.save_project(root).unwrap();
450
451        let result = pre_dispatch_read(
452            "src/big.rs",
453            "auto",
454            None,
455            Some(root.to_str().unwrap()),
456            None,
457        );
458        assert_eq!(result.overridden_mode, Some("map".to_string()));
459        assert_eq!(result.reason, Some("overlay-set-view"));
460    }
461}