1use std::collections::{BTreeMap, BTreeSet, HashMap};
5
6use super::StalenessStatus;
7
8pub const SUGGESTION_WINDOW: usize = 24;
9pub const MEDIUM_SUGGESTION_THRESHOLD: u32 = 45;
10pub const HIGH_SUGGESTION_THRESHOLD: u32 = 70;
11pub const MAJOR_REWRITE_THRESHOLD_PCT: u32 = 50;
12
13const CHANGE_WEIGHT: u32 = 16;
14const DISTINCT_STATE_WEIGHT: u32 = 8;
15const DISTINCT_AGENT_WEIGHT: u32 = 10;
16const RECENCY_WEIGHT: u32 = 12;
17const STALE_WEIGHT: u32 = 18;
18const HAS_CONTEXT_PENALTY: u32 = 35;
19
20#[derive(Clone, Debug, PartialEq, Eq)]
21pub enum ContextSuggestionTier {
22 Medium,
23 High,
24}
25
26#[derive(Clone, Debug, PartialEq, Eq)]
27pub struct ContextSuggestion {
28 pub path: String,
29 pub score: u32,
30 pub tier: ContextSuggestionTier,
31 pub reasons: Vec<String>,
32 pub recent_changes: u32,
33 pub distinct_states: u32,
34 pub distinct_agents: u32,
35 pub has_context: bool,
36 pub stale_annotations: u32,
37}
38
39#[derive(Default)]
40pub struct SuggestionSignal {
41 pub recent_changes: u32,
42 pub distinct_states: BTreeSet<String>,
43 pub distinct_agents: BTreeSet<String>,
44 pub latest_seen_index: Option<usize>,
45}
46
47pub struct SuggestionInputs {
48 pub signals: BTreeMap<String, SuggestionSignal>,
49 pub stale_map: HashMap<String, StalenessStatus>,
50 pub active_paths: BTreeSet<String>,
51 pub history_len: usize,
52}
53
54pub fn score_suggestions(inputs: SuggestionInputs, limit: usize) -> Vec<ContextSuggestion> {
55 let SuggestionInputs {
56 signals,
57 stale_map,
58 active_paths,
59 history_len,
60 } = inputs;
61 let mut suggestions = Vec::new();
62
63 for (path, signal) in signals {
64 let has_context = active_paths.contains(&path);
65 let stale_annotations = stale_map
66 .iter()
67 .filter(|(key, status)| {
68 key.starts_with(&format!("{path}:")) && !matches!(status, StalenessStatus::Fresh)
69 })
70 .count() as u32;
71
72 let mut score = signal.recent_changes.saturating_mul(CHANGE_WEIGHT);
73 score += (signal.distinct_states.len() as u32).saturating_mul(DISTINCT_STATE_WEIGHT);
74 score += (signal.distinct_agents.len() as u32).saturating_mul(DISTINCT_AGENT_WEIGHT);
75 if signal.latest_seen_index.unwrap_or(usize::MAX) <= 3 {
76 score += RECENCY_WEIGHT;
77 }
78 if stale_annotations > 0 {
79 score += stale_annotations.saturating_mul(STALE_WEIGHT);
80 }
81 if has_context && stale_annotations == 0 {
82 score = score.saturating_sub(HAS_CONTEXT_PENALTY);
83 }
84
85 let tier = if score >= HIGH_SUGGESTION_THRESHOLD {
86 Some(ContextSuggestionTier::High)
87 } else if score >= MEDIUM_SUGGESTION_THRESHOLD {
88 Some(ContextSuggestionTier::Medium)
89 } else {
90 None
91 };
92
93 let Some(tier) = tier else {
94 continue;
95 };
96
97 let mut reasons = Vec::new();
98 if signal.recent_changes >= 3 {
99 reasons.push(format!(
100 "{} recent changes across the last {} states",
101 signal.recent_changes, history_len
102 ));
103 }
104 if signal.distinct_agents.len() >= 2 {
105 reasons.push(format!(
106 "{} distinct agents touched this file",
107 signal.distinct_agents.len()
108 ));
109 }
110 if stale_annotations > 0 {
111 reasons.push(format!("{stale_annotations} annotation(s) may be stale"));
112 }
113 if !has_context {
114 reasons.push("no active file guidance exists yet".to_string());
115 }
116
117 suggestions.push(ContextSuggestion {
118 path,
119 score,
120 tier,
121 reasons,
122 recent_changes: signal.recent_changes,
123 distinct_states: signal.distinct_states.len() as u32,
124 distinct_agents: signal.distinct_agents.len() as u32,
125 has_context,
126 stale_annotations,
127 });
128 }
129
130 suggestions.sort_by(|a, b| b.score.cmp(&a.score).then_with(|| a.path.cmp(&b.path)));
131 suggestions.truncate(limit);
132 suggestions
133}
134
135#[cfg(test)]
136mod tests {
137 use super::*;
138
139 #[test]
140 fn score_suggestions_hits_tier_boundaries_and_context_penalty() {
141 let mut signals = BTreeMap::new();
142 signals.insert(
143 "high.rs".to_string(),
144 signal(2, [], ["anthropic/claude", "openai/codex"], None),
145 );
146 signals.insert("penalty.rs".to_string(), signal(5, [], [], None));
147 signals.insert("low.rs".to_string(), signal(1, ["s1"], [], None));
148
149 let mut stale_map = HashMap::new();
150 stale_map.insert("high.rs:File".to_string(), StalenessStatus::Unknown);
151 stale_map.insert("penalty.rs:File".to_string(), StalenessStatus::Fresh);
152
153 let inputs = SuggestionInputs {
154 signals,
155 stale_map,
156 active_paths: BTreeSet::from(["penalty.rs".to_string()]),
157 history_len: 9,
158 };
159
160 let suggestions = score_suggestions(inputs, 10);
161
162 assert_eq!(
163 suggestions,
164 vec![
165 ContextSuggestion {
166 path: "high.rs".to_string(),
167 score: HIGH_SUGGESTION_THRESHOLD,
168 tier: ContextSuggestionTier::High,
169 reasons: vec![
170 "2 distinct agents touched this file".to_string(),
171 "1 annotation(s) may be stale".to_string(),
172 "no active file guidance exists yet".to_string(),
173 ],
174 recent_changes: 2,
175 distinct_states: 0,
176 distinct_agents: 2,
177 has_context: false,
178 stale_annotations: 1,
179 },
180 ContextSuggestion {
181 path: "penalty.rs".to_string(),
182 score: MEDIUM_SUGGESTION_THRESHOLD,
183 tier: ContextSuggestionTier::Medium,
184 reasons: vec!["5 recent changes across the last 9 states".to_string()],
185 recent_changes: 5,
186 distinct_states: 0,
187 distinct_agents: 0,
188 has_context: true,
189 stale_annotations: 0,
190 },
191 ]
192 );
193 }
194
195 fn signal<const S: usize, const A: usize>(
196 recent_changes: u32,
197 states: [&str; S],
198 agents: [&str; A],
199 latest_seen_index: Option<usize>,
200 ) -> SuggestionSignal {
201 SuggestionSignal {
202 recent_changes,
203 distinct_states: states.into_iter().map(str::to_string).collect(),
204 distinct_agents: agents.into_iter().map(str::to_string).collect(),
205 latest_seen_index,
206 }
207 }
208}