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}