1use std::collections::HashMap;
17use std::fmt;
18
19use serde::{Deserialize, Serialize};
20
21#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
29pub struct ContextItemId(pub String);
30
31impl ContextItemId {
32 pub fn from_file(path: &str) -> Self {
33 Self(format!("file:{path}"))
34 }
35 pub fn from_shell(command: &str) -> Self {
36 let hash = crate::core::project_hash::hash_project_root(command);
37 Self(format!("shell:{hash}"))
38 }
39 pub fn from_knowledge(category: &str, key: &str) -> Self {
40 Self(format!("knowledge:{category}:{key}"))
41 }
42 pub fn from_memory(key: &str) -> Self {
43 Self(format!("memory:{key}"))
44 }
45 pub fn from_provider(provider: &str, key: &str) -> Self {
46 Self(format!("provider:{provider}:{key}"))
47 }
48 pub fn as_str(&self) -> &str {
49 &self.0
50 }
51}
52
53impl fmt::Display for ContextItemId {
54 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
55 f.write_str(&self.0)
56 }
57}
58
59#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
60#[serde(rename_all = "snake_case")]
61pub enum ContextKind {
62 File,
63 Shell,
64 Knowledge,
65 Memory,
66 Provider,
67}
68
69#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
70#[serde(rename_all = "snake_case")]
71#[derive(Default)]
72pub enum ContextState {
73 #[default]
74 Candidate,
75 Included,
76 Excluded,
77 Pinned,
78 Stale,
79 Shadowed,
80}
81
82#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
83#[serde(rename_all = "snake_case")]
84pub enum ViewKind {
85 Full,
86 Signatures,
87 Map,
88 Diff,
89 Aggressive,
90 Entropy,
91 Lines,
92 Reference,
93 Handle,
94}
95
96impl ViewKind {
97 pub fn as_str(&self) -> &'static str {
98 match self {
99 Self::Full => "full",
100 Self::Signatures => "signatures",
101 Self::Map => "map",
102 Self::Diff => "diff",
103 Self::Aggressive => "aggressive",
104 Self::Entropy => "entropy",
105 Self::Lines => "lines",
106 Self::Reference => "reference",
107 Self::Handle => "handle",
108 }
109 }
110
111 pub fn parse(s: &str) -> Self {
112 match s.trim().to_lowercase().as_str() {
113 "signatures" => Self::Signatures,
114 "map" => Self::Map,
115 "diff" => Self::Diff,
116 "aggressive" => Self::Aggressive,
117 "entropy" => Self::Entropy,
118 "lines" => Self::Lines,
119 "reference" => Self::Reference,
120 "handle" => Self::Handle,
121 _ => Self::Full,
122 }
123 }
124
125 pub fn density_rank(&self) -> u8 {
127 match self {
128 Self::Full => 0,
129 Self::Aggressive => 1,
130 Self::Diff => 2,
131 Self::Lines => 3,
132 Self::Entropy => 4,
133 Self::Signatures => 5,
134 Self::Map => 6,
135 Self::Reference => 7,
136 Self::Handle => 8,
137 }
138 }
139}
140
141#[derive(Debug, Clone, Default, Serialize, Deserialize)]
143pub struct ViewCosts {
144 pub estimates: HashMap<ViewKind, usize>,
145}
146
147impl ViewCosts {
148 pub fn new() -> Self {
149 Self::default()
150 }
151
152 pub fn set(&mut self, view: ViewKind, tokens: usize) {
153 self.estimates.insert(view, tokens);
154 }
155
156 pub fn get(&self, view: &ViewKind) -> usize {
157 self.estimates.get(view).copied().unwrap_or(0)
158 }
159
160 pub fn cheapest_content_view(&self) -> Option<(ViewKind, usize)> {
162 self.estimates
163 .iter()
164 .filter(|(v, _)| **v != ViewKind::Handle)
165 .min_by_key(|&(_, &tokens)| tokens)
166 .map(|(&v, &t)| (v, t))
167 }
168
169 pub fn from_full_tokens(full_tokens: usize) -> Self {
170 let mut vc = Self::new();
171 vc.set(ViewKind::Full, full_tokens);
172 vc.set(ViewKind::Signatures, full_tokens / 5);
173 vc.set(ViewKind::Map, full_tokens / 8);
174 vc.set(ViewKind::Reference, full_tokens / 20);
175 vc.set(ViewKind::Handle, 25);
176 vc
177 }
178}
179
180#[derive(Debug, Clone, Default, Serialize, Deserialize)]
181pub struct Provenance {
182 pub tool: Option<String>,
183 pub agent_id: Option<String>,
184 pub client_name: Option<String>,
185 pub timestamp: Option<String>,
186}
187
188#[derive(Debug, Clone, Serialize, Deserialize)]
195pub struct FieldWeights {
196 pub w_relevance: f64,
197 pub w_surprise: f64,
198 pub w_graph: f64,
199 pub w_history: f64,
200 pub w_cost: f64,
201 pub w_redundancy: f64,
202}
203
204impl Default for FieldWeights {
205 fn default() -> Self {
206 Self {
207 w_relevance: 0.35,
208 w_surprise: 0.15,
209 w_graph: 0.20,
210 w_history: 0.10,
211 w_cost: 0.10,
212 w_redundancy: 0.10,
213 }
214 }
215}
216
217impl FieldWeights {
218 pub fn conservative() -> Self {
221 Self {
222 w_relevance: 0.45,
223 w_surprise: 0.10,
224 w_graph: 0.20,
225 w_history: 0.15,
226 w_cost: 0.05,
227 w_redundancy: 0.05,
228 }
229 }
230
231 pub fn balanced() -> Self {
233 Self::default()
234 }
235
236 pub fn aggressive() -> Self {
239 Self {
240 w_relevance: 0.30,
241 w_surprise: 0.20,
242 w_graph: 0.15,
243 w_history: 0.05,
244 w_cost: 0.20,
245 w_redundancy: 0.10,
246 }
247 }
248
249 pub fn from_arm(arm: &crate::core::bandit::BanditArm) -> Self {
254 match arm.name.as_str() {
255 "conservative" => Self::conservative(),
256 "aggressive" => Self::aggressive(),
257 _ => Self::balanced(),
258 }
259 }
260}
261
262static ACTIVE_WEIGHTS: std::sync::RwLock<Option<FieldWeights>> = std::sync::RwLock::new(None);
267
268pub fn set_active_weights(weights: FieldWeights) {
271 if let Ok(mut w) = ACTIVE_WEIGHTS.write() {
272 *w = Some(weights);
273 }
274}
275
276pub fn active_weights() -> FieldWeights {
279 ACTIVE_WEIGHTS
280 .read()
281 .ok()
282 .and_then(|w| w.clone())
283 .unwrap_or_default()
284}
285
286#[derive(Debug, Clone, Default, Serialize, Deserialize)]
288pub struct FieldSignals {
289 pub relevance: f64,
290 pub surprise: f64,
291 pub graph_proximity: f64,
292 pub history_signal: f64,
293 pub token_cost_norm: f64,
294 pub redundancy: f64,
295}
296
297#[derive(Debug, Clone, Serialize, Deserialize)]
299pub struct FieldPotential {
300 pub signals: FieldSignals,
301 pub phi: f64,
302 pub view_costs: ViewCosts,
303 pub best_view: ViewKind,
304}
305
306#[derive(Debug, Clone, Copy)]
308pub struct TokenBudget {
309 pub total: usize,
310 pub used: usize,
311}
312
313impl TokenBudget {
314 pub fn remaining(&self) -> usize {
315 self.total.saturating_sub(self.used)
316 }
317 pub fn utilization(&self) -> f64 {
318 if self.total == 0 {
319 return 1.0;
320 }
321 self.used as f64 / self.total as f64
322 }
323 pub fn temperature(&self) -> f64 {
327 let u = self.utilization();
328 (0.1 + u * 1.9).clamp(0.1, 2.0)
329 }
330}
331
332pub struct ContextField {
334 weights: FieldWeights,
335}
336
337impl Default for ContextField {
338 fn default() -> Self {
339 Self::new()
340 }
341}
342
343impl ContextField {
344 pub fn new() -> Self {
345 Self {
346 weights: FieldWeights::default(),
347 }
348 }
349
350 pub fn with_weights(weights: FieldWeights) -> Self {
351 Self { weights }
352 }
353
354 pub fn active() -> Self {
358 Self {
359 weights: active_weights(),
360 }
361 }
362
363 pub fn compute_phi(&self, signals: &FieldSignals) -> f64 {
368 let w = &self.weights;
369 let phi = w.w_relevance * signals.relevance
370 + w.w_surprise * signals.surprise
371 + w.w_graph * signals.graph_proximity
372 + w.w_history * signals.history_signal
373 - w.w_cost * signals.token_cost_norm
374 - w.w_redundancy * signals.redundancy;
375 phi.clamp(0.0, 1.0)
376 }
377
378 pub fn select_view(&self, costs: &ViewCosts, temperature: f64) -> ViewKind {
386 if costs.estimates.is_empty() {
387 return ViewKind::Full;
388 }
389
390 let t = temperature.max(0.01);
391 let max_cost = costs.estimates.values().copied().max().unwrap_or(1).max(1) as f64;
392
393 let mut best_view = ViewKind::Full;
394 let mut best_score = f64::NEG_INFINITY;
395
396 for (&view, &tokens) in &costs.estimates {
397 let normalized_cost = tokens as f64 / max_cost;
398 let density_bonus = 1.0 - (view.density_rank() as f64 / 8.0);
399 let score = density_bonus * (2.0 - t) - normalized_cost * t;
402 if score > best_score {
403 best_score = score;
404 best_view = view;
405 }
406 }
407
408 best_view
409 }
410
411 pub fn compute_batch(
413 &self,
414 items: &[(ContextItemId, FieldSignals, ViewCosts)],
415 budget: TokenBudget,
416 ) -> HashMap<ContextItemId, FieldPotential> {
417 let temperature = budget.temperature();
418 let mut result = HashMap::new();
419
420 for (id, signals, costs) in items {
421 let phi = self.compute_phi(signals);
422 let best_view = self.select_view(costs, temperature);
423 result.insert(
424 id.clone(),
425 FieldPotential {
426 signals: signals.clone(),
427 phi,
428 view_costs: costs.clone(),
429 best_view,
430 },
431 );
432 }
433
434 result
435 }
436}
437
438pub fn normalize_relevance(score: f64, max_score: f64) -> f64 {
444 if max_score <= 0.0 {
445 return 0.0;
446 }
447 (score / max_score).clamp(0.0, 1.0)
448}
449
450pub fn normalize_surprise(surprise: f64) -> f64 {
453 ((surprise - 5.0) / 12.0).clamp(0.0, 1.0)
454}
455
456pub fn normalize_graph_proximity(distance: usize) -> f64 {
459 1.0 / (1.0 + distance as f64)
460}
461
462pub fn normalize_token_cost(tokens: usize, budget_total: usize) -> f64 {
464 if budget_total == 0 {
465 return 1.0;
466 }
467 (tokens as f64 / budget_total as f64).clamp(0.0, 1.0)
468}
469
470pub fn efficiency(phi: f64, tokens: usize) -> f64 {
473 if tokens == 0 {
474 return phi;
475 }
476 phi / tokens as f64
477}
478
479pub const MMR_LAMBDA: f64 = 0.7;
483
484pub fn mmr_score(phi: f64, max_similarity: f64, lambda: f64) -> f64 {
490 let l = lambda.clamp(0.0, 1.0);
491 l * phi - (1.0 - l) * max_similarity.clamp(0.0, 1.0)
492}
493
494pub fn compute_signals_for_path(
498 path: &str,
499 task: Option<&str>,
500 file_content: Option<&str>,
501 budget_total: usize,
502 full_tokens: usize,
503) -> (FieldSignals, ViewCosts) {
504 let mut signals = FieldSignals::default();
505
506 let heatmap = super::heatmap::HeatMap::load();
507 let heat_entry = heatmap.entries.get(path);
508
509 if let Some(task_desc) = task {
511 let (_, keywords) = super::task_relevance::parse_task_hints(task_desc);
512 let path_lower = path.to_lowercase();
513 let keyword_hits = keywords
514 .iter()
515 .filter(|kw| path_lower.contains(&kw.to_lowercase()))
516 .count();
517 let keyword_score = (keyword_hits as f64 * 0.3).min(1.0);
518 let freq_score = heat_entry.map_or(0.0, |e| (e.access_count as f64 / 10.0).min(1.0));
519 signals.relevance = normalize_relevance(keyword_score + freq_score, 2.0);
520 } else {
521 let freq = heat_entry.map_or(0.0, |e| e.access_count as f64);
522 signals.relevance = normalize_relevance(freq, 10.0);
523 }
524
525 if let Some(content) = file_content {
527 let surprise_val = super::surprise::line_surprise(content);
528 signals.surprise = normalize_surprise(surprise_val);
529 }
530
531 let depth = path.matches('/').count();
534 signals.graph_proximity = normalize_graph_proximity(depth);
535
536 let access_count = heat_entry.map_or(0, |e| e.access_count);
538 signals.history_signal = (access_count as f64 / 20.0).min(1.0);
539
540 signals.token_cost_norm = normalize_token_cost(full_tokens, budget_total);
542
543 signals.redundancy = 0.0;
545
546 let view_costs = ViewCosts::from_full_tokens(full_tokens);
547 (signals, view_costs)
548}
549
550#[cfg(test)]
555mod tests {
556 use super::*;
557
558 #[test]
559 fn phi_increases_with_relevance() {
560 let field = ContextField::new();
561 let low = field.compute_phi(&FieldSignals {
562 relevance: 0.2,
563 ..Default::default()
564 });
565 let high = field.compute_phi(&FieldSignals {
566 relevance: 0.9,
567 ..Default::default()
568 });
569 assert!(high > low, "higher relevance should yield higher phi");
570 }
571
572 #[test]
573 fn phi_decreases_with_cost() {
574 let field = ContextField::new();
575 let cheap = field.compute_phi(&FieldSignals {
576 relevance: 0.5,
577 token_cost_norm: 0.1,
578 ..Default::default()
579 });
580 let expensive = field.compute_phi(&FieldSignals {
581 relevance: 0.5,
582 token_cost_norm: 0.9,
583 ..Default::default()
584 });
585 assert!(cheap > expensive, "higher cost should reduce phi");
586 }
587
588 #[test]
589 fn phi_decreases_with_redundancy() {
590 let field = ContextField::new();
591 let unique = field.compute_phi(&FieldSignals {
592 relevance: 0.5,
593 redundancy: 0.0,
594 ..Default::default()
595 });
596 let redundant = field.compute_phi(&FieldSignals {
597 relevance: 0.5,
598 redundancy: 0.9,
599 ..Default::default()
600 });
601 assert!(unique > redundant, "redundancy should reduce phi");
602 }
603
604 #[test]
605 fn phi_is_clamped_to_unit_interval() {
606 let field = ContextField::new();
607 let phi = field.compute_phi(&FieldSignals {
608 relevance: 1.0,
609 surprise: 1.0,
610 graph_proximity: 1.0,
611 history_signal: 1.0,
612 token_cost_norm: 0.0,
613 redundancy: 0.0,
614 });
615 assert!(phi <= 1.0);
616 assert!(phi >= 0.0);
617 }
618
619 #[test]
620 fn view_selection_prefers_dense_at_low_temperature() {
621 let field = ContextField::new();
622 let costs = ViewCosts::from_full_tokens(5000);
623 let view = field.select_view(&costs, 0.1);
624 assert_eq!(
625 view,
626 ViewKind::Full,
627 "low temperature (relaxed budget) should prefer full view"
628 );
629 }
630
631 #[test]
632 fn view_selection_prefers_sparse_at_high_temperature() {
633 let field = ContextField::new();
634 let costs = ViewCosts::from_full_tokens(5000);
635 let view = field.select_view(&costs, 2.0);
636 assert_ne!(
637 view,
638 ViewKind::Full,
639 "high temperature (tight budget) should prefer sparser view"
640 );
641 }
642
643 #[test]
644 fn budget_temperature_scales_with_utilization() {
645 let low = TokenBudget {
646 total: 10000,
647 used: 1000,
648 };
649 let high = TokenBudget {
650 total: 10000,
651 used: 9000,
652 };
653 assert!(
654 high.temperature() > low.temperature(),
655 "higher utilization should increase temperature"
656 );
657 }
658
659 #[test]
660 fn normalize_surprise_maps_range() {
661 assert!((normalize_surprise(5.0) - 0.0).abs() < 0.01);
662 assert!((normalize_surprise(17.0) - 1.0).abs() < 0.01);
663 assert!((normalize_surprise(11.0) - 0.5).abs() < 0.01);
664 }
665
666 #[test]
667 fn normalize_graph_proximity_inverse_distance() {
668 assert!((normalize_graph_proximity(0) - 1.0).abs() < f64::EPSILON);
669 assert!((normalize_graph_proximity(1) - 0.5).abs() < f64::EPSILON);
670 assert!(normalize_graph_proximity(10) < 0.15);
671 }
672
673 #[test]
674 fn efficiency_ratio_is_phi_per_token() {
675 let e = efficiency(0.8, 400);
676 assert!((e - 0.002).abs() < 0.0001);
677 }
678
679 #[test]
680 fn field_weights_from_arm_maps_presets() {
681 let mut bandit = crate::core::bandit::ThresholdBandit::default();
683 let con = FieldWeights::from_arm(&bandit.arms[0]); let agg = FieldWeights::from_arm(&bandit.arms[2]); assert!(
686 con.w_relevance > agg.w_relevance,
687 "conservative trusts relevance more"
688 );
689 assert!(
690 agg.w_cost > con.w_cost,
691 "aggressive penalizes cost more (denser)"
692 );
693 let _ = bandit.choose_arm();
694 }
695
696 #[test]
697 fn learned_weights_shift_after_feedback() {
698 let mut bandit = crate::core::bandit::ThresholdBandit::default();
701 for _ in 0..20 {
702 bandit.update("aggressive", true);
703 }
704 for _ in 0..20 {
705 bandit.update("conservative", false);
706 }
707 let arm = bandit.arms[bandit.best_arm_idx_by_mean()].clone();
708 let learned = FieldWeights::from_arm(&arm);
709 let expected = FieldWeights::aggressive();
710 assert!((learned.w_cost - expected.w_cost).abs() < f64::EPSILON);
711 }
712
713 #[test]
714 fn active_weights_default_is_balanced() {
715 set_active_weights(FieldWeights::default());
718 let w = active_weights();
719 let d = FieldWeights::default();
720 assert!((w.w_relevance - d.w_relevance).abs() < f64::EPSILON);
721 }
722
723 #[test]
724 fn mmr_penalizes_redundancy() {
725 let unique = mmr_score(0.8, 0.0, MMR_LAMBDA);
727 let redundant = mmr_score(0.8, 0.9, MMR_LAMBDA);
728 assert!(
729 unique > redundant,
730 "redundant item ({redundant}) must score below unique ({unique})"
731 );
732 }
733
734 #[test]
735 fn mmr_lambda_one_is_pure_relevance() {
736 assert!((mmr_score(0.6, 0.9, 1.0) - 0.6).abs() < f64::EPSILON);
738 }
739
740 #[test]
741 fn context_item_id_stable() {
742 let a = ContextItemId::from_file("src/main.rs");
743 let b = ContextItemId::from_file("src/main.rs");
744 assert_eq!(a, b);
745 }
746
747 #[test]
748 fn view_costs_from_full() {
749 let vc = ViewCosts::from_full_tokens(5000);
750 assert_eq!(vc.get(&ViewKind::Full), 5000);
751 assert_eq!(vc.get(&ViewKind::Signatures), 1000);
752 assert_eq!(vc.get(&ViewKind::Map), 625);
753 assert_eq!(vc.get(&ViewKind::Handle), 25);
754 }
755
756 #[test]
757 fn batch_compute_produces_results_for_all_items() {
758 let field = ContextField::new();
759 let items = vec![
760 (
761 ContextItemId::from_file("a.rs"),
762 FieldSignals {
763 relevance: 0.8,
764 ..Default::default()
765 },
766 ViewCosts::from_full_tokens(2000),
767 ),
768 (
769 ContextItemId::from_file("b.rs"),
770 FieldSignals {
771 relevance: 0.3,
772 ..Default::default()
773 },
774 ViewCosts::from_full_tokens(500),
775 ),
776 ];
777 let budget = TokenBudget {
778 total: 10000,
779 used: 2000,
780 };
781 let results = field.compute_batch(&items, budget);
782 assert_eq!(results.len(), 2);
783 assert!(results.contains_key(&ContextItemId::from_file("a.rs")));
784 assert!(results.contains_key(&ContextItemId::from_file("b.rs")));
785 }
786}