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::SuggestCompression => match requested_mode {
249            "auto" => Some("map".to_string()),
250            _ => None,
251        },
252        PressureAction::ForceCompression => match requested_mode {
253            "full" => Some("map".to_string()),
254            "auto" | "map" => Some("signatures".to_string()),
255            _ => None,
256        },
257        PressureAction::EvictLeastRelevant => match requested_mode {
258            "full" => Some("map".to_string()),
259            "auto" | "map" => Some("signatures".to_string()),
260            _ => None,
261        },
262        PressureAction::NoAction => None,
263    }
264}
265
266fn check_overlay_mode_override(
267    path: &str,
268    requested_mode: &str,
269    overlay: &OverlayStore,
270) -> Option<PreDispatchResult> {
271    let item_id = ContextItemId::from_file(path);
272    let overlays = overlay.for_item(&item_id);
273
274    for ov in overlays.iter().rev() {
275        match &ov.operation {
276            OverlayOp::SetView(view) => {
277                let mode_str = view.as_str();
278                if mode_str != requested_mode {
279                    return Some(PreDispatchResult {
280                        overridden_mode: Some(mode_str.to_string()),
281                        reason: Some("overlay-set-view"),
282                        pressure_downgraded: false,
283                        budget_blocked: false,
284                        budget_warning: None,
285                    });
286                }
287            }
288            OverlayOp::Pin { .. } if requested_mode != "full" => {
289                return Some(PreDispatchResult {
290                    overridden_mode: Some("full".to_string()),
291                    reason: Some("pinned"),
292                    pressure_downgraded: false,
293                    budget_blocked: false,
294                    budget_warning: None,
295                });
296            }
297            OverlayOp::Exclude { .. } if requested_mode != "signatures" => {
298                return Some(PreDispatchResult {
299                    overridden_mode: Some("signatures".to_string()),
300                    reason: Some("excluded"),
301                    pressure_downgraded: false,
302                    budget_blocked: false,
303                    budget_warning: None,
304                });
305            }
306            _ => {}
307        }
308    }
309    None
310}
311
312pub fn post_dispatch_record(
313    path: &str,
314    mode: &str,
315    original_tokens: usize,
316    sent_tokens: usize,
317    ledger: &mut ContextLedger,
318    overlay: &OverlayStore,
319) -> PostDispatchResult {
320    post_dispatch_record_with_task(
321        path,
322        mode,
323        original_tokens,
324        sent_tokens,
325        ledger,
326        overlay,
327        None,
328    )
329}
330
331pub fn post_dispatch_record_with_task(
332    path: &str,
333    mode: &str,
334    original_tokens: usize,
335    sent_tokens: usize,
336    ledger: &mut ContextLedger,
337    overlay: &OverlayStore,
338    task: Option<&str>,
339) -> PostDispatchResult {
340    let prev_count = ledger.entries.len();
341    let prev_pressure = ledger.pressure().recommendation;
342
343    ledger.record_with_task(path, mode, original_tokens, sent_tokens, task);
344
345    let item_id = ContextItemId::from_file(path);
346    let state = overlay.apply_to_state(&item_id, ContextState::Included);
347
348    if state == ContextState::Excluded {
349        return PostDispatchResult {
350            eviction_hint: Some(format!("File '{path}' is excluded by overlay.")),
351            elicitation_hint: None,
352            resource_changed: true,
353        };
354    }
355
356    let elicitation =
357        super::elicitation::check_elicitation_needed(ledger, Some(path), Some(sent_tokens))
358            .map(|s| s.format_fallback_hint());
359
360    let pressure = ledger.pressure();
361
362    apply_reinjection_plan(ledger, &pressure.recommendation);
363
364    let new_entry = ledger.entries.len() != prev_count;
365    let pressure_shifted = pressure.recommendation != prev_pressure;
366    let resource_changed = new_entry || pressure_shifted;
367
368    if pressure.utilization > 0.9 {
369        let candidates = ledger.eviction_candidates_by_phi(3);
370        if !candidates.is_empty() {
371            let names: Vec<_> = candidates
372                .iter()
373                .take(3)
374                .map(|p| crate::core::protocol::shorten_path(p))
375                .collect();
376            return PostDispatchResult {
377                eviction_hint: Some(format!(
378                    "Context pressure {:.0}%. Evict: ctx_ledger(action=\"evict\", targets=\"{}\")",
379                    pressure.utilization * 100.0,
380                    names.join(", ")
381                )),
382                elicitation_hint: elicitation,
383                resource_changed,
384            };
385        }
386    }
387
388    PostDispatchResult {
389        eviction_hint: None,
390        elicitation_hint: elicitation,
391        resource_changed,
392    }
393}
394
395fn apply_reinjection_plan(ledger: &mut ContextLedger, action: &PressureAction) {
396    if *action != PressureAction::ForceCompression && *action != PressureAction::EvictLeastRelevant
397    {
398        return;
399    }
400    for entry in &mut ledger.entries {
401        if entry.mode == "full" {
402            entry.mode = "map".to_string();
403        }
404    }
405}
406
407fn try_load_graph(project_root: &str) -> Option<crate::core::graph_index::ProjectIndex> {
408    crate::core::graph_index::ProjectIndex::load(project_root)
409}
410
411#[cfg(test)]
412mod tests {
413    use super::*;
414
415    #[test]
416    fn pre_dispatch_passthrough_for_full() {
417        let result = pre_dispatch_read("src/main.rs", "full", None, None, None);
418        assert!(result.overridden_mode.is_none());
419    }
420
421    #[test]
422    fn pre_dispatch_passthrough_for_diff() {
423        let result = pre_dispatch_read("src/main.rs", "diff", None, None, None);
424        assert!(result.overridden_mode.is_none());
425    }
426
427    #[test]
428    fn pre_dispatch_no_override_without_signals() {
429        let result = pre_dispatch_read("src/unknown.rs", "auto", None, None, None);
430        assert!(result.overridden_mode.is_none());
431    }
432
433    #[test]
434    fn pre_dispatch_bounce_prevention_forces_full() {
435        {
436            let mut bt = crate::core::bounce_tracker::global().lock().unwrap();
437            bt.set_seq(1);
438            bt.record_read("src/bouncy.yml", "map", 30, 400);
439            bt.set_seq(2);
440            bt.record_read("src/bouncy.yml", "full", 400, 400);
441            bt.set_seq(3);
442            bt.record_read("a2.yml", "map", 30, 400);
443            bt.set_seq(4);
444            bt.record_read("a2.yml", "full", 400, 400);
445            bt.set_seq(5);
446            bt.record_read("a3.yml", "map", 30, 400);
447            bt.set_seq(6);
448            bt.record_read("a3.yml", "full", 400, 400);
449        }
450        let result = pre_dispatch_read("new.yml", "auto", None, None, None);
451        assert_eq!(result.overridden_mode, Some("full".to_string()));
452        assert_eq!(result.reason, Some("bounce-prevention"));
453    }
454
455    #[test]
456    fn pressure_downgrade_full_to_map() {
457        let result = pre_dispatch_read(
458            "c.rs",
459            "full",
460            None,
461            None,
462            Some(&PressureAction::ForceCompression),
463        );
464        assert_eq!(result.overridden_mode, Some("map".to_string()));
465        assert_eq!(result.reason, Some("pressure-auto-downgrade"));
466        assert!(result.pressure_downgraded);
467    }
468
469    #[test]
470    fn pressure_downgrade_map_to_signatures_on_evict() {
471        let result = pre_dispatch_read(
472            "c.rs",
473            "map",
474            None,
475            None,
476            Some(&PressureAction::EvictLeastRelevant),
477        );
478        assert_eq!(result.overridden_mode, Some("signatures".to_string()));
479        assert!(result.pressure_downgraded);
480    }
481
482    #[test]
483    fn no_pressure_downgrade_when_low() {
484        let result = pre_dispatch_read("c.rs", "full", None, None, Some(&PressureAction::NoAction));
485        assert!(result.overridden_mode.is_none());
486        assert!(!result.pressure_downgraded);
487    }
488
489    #[test]
490    fn suggest_compression_downgrades_auto_to_map() {
491        let result = pre_dispatch_read(
492            "c.rs",
493            "auto",
494            None,
495            None,
496            Some(&PressureAction::SuggestCompression),
497        );
498        assert_eq!(result.overridden_mode, Some("map".to_string()));
499        assert!(result.pressure_downgraded);
500    }
501
502    #[test]
503    fn suggest_compression_does_not_touch_explicit_full() {
504        let result = pre_dispatch_read(
505            "c.rs",
506            "full",
507            None,
508            None,
509            Some(&PressureAction::SuggestCompression),
510        );
511        assert!(result.overridden_mode.is_none());
512        assert!(!result.pressure_downgraded);
513    }
514
515    #[test]
516    fn post_dispatch_reinjection_downgrades_entries() {
517        let mut ledger = ContextLedger::with_window_size(1000);
518        ledger.record("a.rs", "full", 400, 400);
519        ledger.record("b.rs", "full", 400, 400);
520        let overlay = OverlayStore::new();
521        let result = post_dispatch_record("c.rs", "full", 300, 300, &mut ledger, &overlay);
522        assert!(result.resource_changed);
523        let a_entry = ledger.entries.iter().find(|e| e.path == "a.rs").unwrap();
524        assert_eq!(a_entry.mode, "map");
525    }
526
527    #[test]
528    fn overlay_pin_forces_full_mode() {
529        let dir = tempfile::tempdir().expect("tmp dir");
530        let root = dir.path();
531        let mut store = OverlayStore::new();
532        let target = ContextItemId::from_file("src/important.rs");
533        store.add(crate::core::context_overlay::ContextOverlay::new(
534            target,
535            OverlayOp::Pin { verbatim: false },
536            crate::core::context_overlay::OverlayScope::Project,
537            String::new(),
538            crate::core::context_overlay::OverlayAuthor::User,
539        ));
540        store.save_project(root).unwrap();
541
542        let result = pre_dispatch_read(
543            "src/important.rs",
544            "auto",
545            None,
546            Some(root.to_str().unwrap()),
547            None,
548        );
549        assert_eq!(result.overridden_mode, Some("full".to_string()));
550        assert_eq!(result.reason, Some("pinned"));
551    }
552
553    #[test]
554    fn overlay_exclude_forces_signatures_mode() {
555        let dir = tempfile::tempdir().expect("tmp dir");
556        let root = dir.path();
557        let mut store = OverlayStore::new();
558        let target = ContextItemId::from_file("src/noisy.rs");
559        store.add(crate::core::context_overlay::ContextOverlay::new(
560            target,
561            OverlayOp::Exclude {
562                reason: "noise".to_string(),
563            },
564            crate::core::context_overlay::OverlayScope::Project,
565            String::new(),
566            crate::core::context_overlay::OverlayAuthor::User,
567        ));
568        store.save_project(root).unwrap();
569
570        let result = pre_dispatch_read(
571            "src/noisy.rs",
572            "auto",
573            None,
574            Some(root.to_str().unwrap()),
575            None,
576        );
577        assert_eq!(result.overridden_mode, Some("signatures".to_string()));
578        assert_eq!(result.reason, Some("excluded"));
579    }
580
581    #[test]
582    fn overlay_set_view_forces_specified_mode() {
583        let dir = tempfile::tempdir().expect("tmp dir");
584        let root = dir.path();
585        let mut store = OverlayStore::new();
586        let target = ContextItemId::from_file("src/big.rs");
587        store.add(crate::core::context_overlay::ContextOverlay::new(
588            target,
589            OverlayOp::SetView(crate::core::context_field::ViewKind::Map),
590            crate::core::context_overlay::OverlayScope::Project,
591            String::new(),
592            crate::core::context_overlay::OverlayAuthor::User,
593        ));
594        store.save_project(root).unwrap();
595
596        let result = pre_dispatch_read(
597            "src/big.rs",
598            "auto",
599            None,
600            Some(root.to_str().unwrap()),
601            None,
602        );
603        assert_eq!(result.overridden_mode, Some("map".to_string()));
604        assert_eq!(result.reason, Some("overlay-set-view"));
605    }
606}