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