1use std::collections::HashSet;
18use std::path::Path;
19
20use mif_ontology::ExpansionConfig;
21use serde::{Deserialize, Serialize};
22
23use crate::error::MifRhError;
24use crate::index::Miss;
25use crate::suggest::TypeSuggestion;
26
27pub const STATUS_PENDING: &str = "pending";
30
31#[derive(Debug, Clone, Serialize, Deserialize)]
34pub struct SuggestionEntry {
35 pub finding_id: String,
37 #[serde(default, skip_serializing_if = "Option::is_none")]
42 pub file: Option<String>,
43 pub basis: String,
46 pub run_id: String,
48 pub candidates: Vec<TypeSuggestion>,
50 #[serde(default = "default_status")]
54 pub status: String,
55}
56
57fn default_status() -> String {
58 STATUS_PENDING.to_string()
59}
60
61#[derive(Debug, Clone, Serialize, Deserialize)]
63pub struct SuggestionQueue {
64 pub topic: String,
66 pub entries: Vec<SuggestionEntry>,
68}
69
70pub fn upsert_suggestions(
101 path: &Path,
102 topic: &str,
103 fresh: Vec<SuggestionEntry>,
104) -> Result<SuggestionQueue, MifRhError> {
105 let mut queue = if path.exists() {
106 let contents = std::fs::read_to_string(path).map_err(|source| MifRhError::Io {
107 path: path.display().to_string(),
108 source,
109 })?;
110 let queue = serde_json::from_str::<SuggestionQueue>(&contents).map_err(|source| {
111 MifRhError::Json {
112 path: path.display().to_string(),
113 source,
114 }
115 })?;
116 if queue.topic != topic {
119 return Err(MifRhError::QueueTopicMismatch {
120 path: path.display().to_string(),
121 expected: topic.to_string(),
122 found: queue.topic,
123 });
124 }
125 queue
126 } else {
127 SuggestionQueue {
128 topic: topic.to_string(),
129 entries: Vec::new(),
130 }
131 };
132
133 for entry in fresh {
134 match queue
135 .entries
136 .iter_mut()
137 .find(|existing| existing.finding_id == entry.finding_id)
138 {
139 Some(existing) if existing.status == STATUS_PENDING => *existing = entry,
140 Some(_) => {}, None => queue.entries.push(entry),
142 }
143 }
144 queue
145 .entries
146 .sort_by(|a, b| a.finding_id.cmp(&b.finding_id));
147
148 if let Some(parent) = path.parent() {
149 std::fs::create_dir_all(parent).map_err(|source| MifRhError::Io {
150 path: parent.display().to_string(),
151 source,
152 })?;
153 }
154 crate::write_json_atomic(path, &queue)?;
155 Ok(queue)
156}
157
158#[derive(Debug, Clone, Serialize, Deserialize)]
160pub struct ClusterMember {
161 pub finding_id: String,
163 pub topic: String,
165 pub content: String,
167 pub run_id: String,
169}
170
171#[derive(Debug, Clone, Serialize, Deserialize)]
174pub struct ExpansionCandidate {
175 pub size: usize,
177 pub runs: usize,
179 pub members: Vec<ClusterMember>,
181}
182
183#[must_use]
198pub fn expansion_candidates(misses: &[Miss], cfg: &ExpansionConfig) -> Vec<ExpansionCandidate> {
199 let vectors: Vec<Vec<f32>> = misses.iter().map(|m| m.vector.clone()).collect();
200 let clusters = mif_ontology::cluster_by_mutual_similarity(&vectors, cfg.cluster_similarity);
201
202 clusters
203 .into_iter()
204 .filter_map(|indices| {
205 let members: Vec<ClusterMember> = indices
206 .iter()
207 .map(|&i| ClusterMember {
208 finding_id: misses[i].finding_id.clone(),
209 topic: misses[i].topic.clone(),
210 content: misses[i].content.clone(),
211 run_id: misses[i].run_id.clone(),
212 })
213 .collect();
214 let size = members
215 .iter()
216 .map(|m| m.finding_id.as_str())
217 .collect::<HashSet<_>>()
218 .len();
219 let runs = members
220 .iter()
221 .map(|m| m.run_id.as_str())
222 .collect::<HashSet<_>>()
223 .len();
224 (size >= cfg.min_cluster_size && runs >= cfg.min_distinct_runs).then_some(
225 ExpansionCandidate {
226 size,
227 runs,
228 members,
229 },
230 )
231 })
232 .collect()
233}
234
235#[cfg(test)]
236mod tests {
237 use mif_ontology::{CalibrationConfig, ConfidenceTier, ExpansionConfig};
238
239 use super::{STATUS_PENDING, SuggestionEntry, expansion_candidates, upsert_suggestions};
240 use crate::index::Miss;
241 use crate::suggest::TypeSuggestion;
242
243 fn entry(finding_id: &str, run_id: &str, status: &str) -> SuggestionEntry {
244 SuggestionEntry {
245 finding_id: finding_id.to_string(),
246 file: None,
247 basis: "untyped".to_string(),
248 run_id: run_id.to_string(),
249 candidates: vec![TypeSuggestion {
250 entity_type: "control".to_string(),
251 ontology_id: "sec".to_string(),
252 score: 0.7,
253 tier: ConfidenceTier::FlagForReview,
254 margin: None,
255 calibrated: false,
256 negative_demoted: false,
257 }],
258 status: status.to_string(),
259 }
260 }
261
262 fn miss(finding_id: &str, run_id: &str, vector: Vec<f32>) -> Miss {
263 Miss {
264 finding_id: finding_id.to_string(),
265 topic: "sec".to_string(),
266 content: format!("content of {finding_id}"),
267 vector,
268 run_id: run_id.to_string(),
269 model: "test-model".to_string(),
270 }
271 }
272
273 #[test]
274 fn upsert_creates_refreshes_pending_and_preserves_verdicts() {
275 let dir = tempfile::tempdir().unwrap();
276 let path = dir.path().join("suggestions/sec.json");
277
278 upsert_suggestions(
280 &path,
281 "sec",
282 vec![
283 entry("f-1", "run-1", STATUS_PENDING),
284 entry("f-2", "run-1", STATUS_PENDING),
285 ],
286 )
287 .unwrap();
288
289 let contents = std::fs::read_to_string(&path).unwrap();
291 let confirmed = contents.replacen("\"pending\"", "\"confirmed\"", 1);
292 std::fs::write(&path, confirmed).unwrap();
293
294 let queue = upsert_suggestions(
296 &path,
297 "sec",
298 vec![
299 entry("f-1", "run-2", STATUS_PENDING),
300 entry("f-2", "run-2", STATUS_PENDING),
301 entry("f-3", "run-2", STATUS_PENDING),
302 ],
303 )
304 .unwrap();
305
306 assert_eq!(queue.entries.len(), 3);
307 let f1 = queue
309 .entries
310 .iter()
311 .find(|e| e.finding_id == "f-1")
312 .unwrap();
313 assert_eq!(f1.status, "confirmed");
314 assert_eq!(f1.run_id, "run-1");
315 let f2 = queue
317 .entries
318 .iter()
319 .find(|e| e.finding_id == "f-2")
320 .unwrap();
321 assert_eq!(f2.run_id, "run-2");
322 let ids: Vec<&str> = queue
324 .entries
325 .iter()
326 .map(|e| e.finding_id.as_str())
327 .collect();
328 assert_eq!(ids, ["f-1", "f-2", "f-3"]);
329 }
330
331 #[test]
332 fn a_queue_recorded_for_another_topic_is_rejected_not_absorbed() {
333 let dir = tempfile::tempdir().unwrap();
334 let path = dir.path().join("sec.json");
335 upsert_suggestions(&path, "sec", vec![entry("f-1", "run-1", STATUS_PENDING)]).unwrap();
336
337 let error = upsert_suggestions(&path, "edu", vec![]).unwrap_err();
338 assert!(matches!(
339 error,
340 crate::MifRhError::QueueTopicMismatch { .. }
341 ));
342 }
343
344 #[test]
345 fn a_corrupt_queue_file_is_an_explicit_error_never_a_silent_reset() {
346 let dir = tempfile::tempdir().unwrap();
347 let path = dir.path().join("sec.json");
348 std::fs::write(&path, "{ not json").unwrap();
349
350 let error = upsert_suggestions(&path, "sec", vec![]).unwrap_err();
351 assert!(matches!(error, crate::MifRhError::Json { .. }));
352 }
353
354 #[test]
355 fn expansion_candidates_require_size_and_distinct_run_gates() {
356 let cfg = ExpansionConfig {
357 cluster_similarity: 0.95,
358 min_cluster_size: 2,
359 min_distinct_runs: 2,
360 };
361 let one_run = [
364 miss("f-1", "run-1", vec![1.0, 0.0]),
365 miss("f-2", "run-1", vec![1.0, 0.0]),
366 ];
367 assert!(expansion_candidates(&one_run, &cfg).is_empty());
368
369 let two_runs = [
371 miss("f-1", "run-1", vec![1.0, 0.0]),
372 miss("f-2", "run-2", vec![1.0, 0.0]),
373 miss("f-3", "run-2", vec![0.0, 1.0]), ];
375 let candidates = expansion_candidates(&two_runs, &cfg);
376 assert_eq!(candidates.len(), 1);
377 assert_eq!(candidates[0].size, 2);
378 assert_eq!(candidates[0].runs, 2);
379 }
380
381 #[test]
382 fn the_same_finding_recurring_across_runs_does_not_inflate_cluster_size() {
383 let cfg = ExpansionConfig {
384 cluster_similarity: 0.95,
385 min_cluster_size: 2,
386 min_distinct_runs: 2,
387 };
388 let same_finding = [
391 miss("f-1", "run-1", vec![1.0, 0.0]),
392 miss("f-1", "run-2", vec![1.0, 0.0]),
393 ];
394 assert!(expansion_candidates(&same_finding, &cfg).is_empty());
395 }
396
397 #[test]
398 fn calibration_expansion_defaults_flow_through() {
399 let cfg = CalibrationConfig::default().expansion;
402 assert!((cfg.cluster_similarity - 0.80).abs() < f32::EPSILON);
403 assert_eq!(cfg.min_cluster_size, 3);
404 assert_eq!(cfg.min_distinct_runs, 2);
405 }
406}