lean_ctx/core/
search_kernel.rs1use std::collections::{HashMap, hash_map::DefaultHasher};
5use std::hash::{Hash, Hasher};
6use std::sync::atomic::{AtomicUsize, Ordering};
7use std::sync::{LazyLock, Mutex};
8
9use crate::core::context_kernel::evidence_hook;
10
11static QUERY_TOKENS: LazyLock<Mutex<HashMap<u64, usize>>> = LazyLock::new(Default::default);
12static TOTAL_SEARCHES: AtomicUsize = AtomicUsize::new(0);
13static TOTAL_TOKENS: AtomicUsize = AtomicUsize::new(0);
14
15#[derive(Debug, Clone, Default, serde::Serialize)]
17pub struct SearchSummary {
18 pub total_searches: usize,
20 pub unique_queries: usize,
22 pub repeated_queries: usize,
24 pub total_tokens: usize,
26}
27
28fn query_hash(query: &str) -> u64 {
29 let mut hasher = DefaultHasher::new();
30 query.hash(&mut hasher);
31 hasher.finish()
32}
33
34pub fn record_search(query: &str, result_count: usize, tokens: usize) {
36 if !crate::core::context_kernel::kernel_config::is_enabled() {
37 return;
38 }
39 evidence_hook::record_tool_call("ctx_search", result_count, tokens);
40 QUERY_TOKENS
41 .lock()
42 .unwrap_or_else(std::sync::PoisonError::into_inner)
43 .entry(query_hash(query))
44 .and_modify(|c| *c += 1)
45 .or_insert(1);
46 TOTAL_SEARCHES.fetch_add(1, Ordering::Relaxed);
47 TOTAL_TOKENS.fetch_add(tokens, Ordering::Relaxed);
48}
49
50#[must_use]
52pub fn is_repeated_query(query: &str) -> bool {
53 QUERY_TOKENS
54 .lock()
55 .unwrap_or_else(std::sync::PoisonError::into_inner)
56 .get(&query_hash(query))
57 .is_some_and(|&c| c > 1)
58}
59
60#[must_use]
62pub fn search_summary() -> SearchSummary {
63 let total_searches = TOTAL_SEARCHES.load(Ordering::Relaxed);
64 let unique_queries = QUERY_TOKENS
65 .lock()
66 .unwrap_or_else(std::sync::PoisonError::into_inner)
67 .len();
68 SearchSummary {
69 total_searches,
70 unique_queries,
71 repeated_queries: total_searches.saturating_sub(unique_queries),
72 total_tokens: TOTAL_TOKENS.load(Ordering::Relaxed),
73 }
74}
75
76pub fn reset() {
78 QUERY_TOKENS
79 .lock()
80 .unwrap_or_else(std::sync::PoisonError::into_inner)
81 .clear();
82 TOTAL_SEARCHES.store(0, Ordering::Relaxed);
83 TOTAL_TOKENS.store(0, Ordering::Relaxed);
84}
85
86#[cfg(test)]
87mod tests {
88 use super::{is_repeated_query, record_search, reset, search_summary};
89 use crate::core::context_kernel::kernel_config::KERNEL_TEST_LOCK;
90
91 fn isolated() -> std::sync::MutexGuard<'static, ()> {
92 let guard = KERNEL_TEST_LOCK
93 .lock()
94 .unwrap_or_else(std::sync::PoisonError::into_inner);
95 reset();
96 guard
97 }
98
99 #[test]
100 fn records_search_evidence() {
101 let _guard = isolated();
102 record_search("one", 1, 10);
103 record_search("two", 2, 20);
104 record_search("three", 3, 30);
105 assert_eq!(search_summary().total_searches, 3);
106 }
107
108 #[test]
109 fn detects_repeated_query() {
110 let _guard = isolated();
111 record_search("same", 1, 10);
112 record_search("same", 1, 10);
113 assert!(is_repeated_query("same"));
114 assert_eq!(search_summary().repeated_queries, 1);
115 }
116
117 #[test]
118 fn unique_queries_tracked() {
119 let _guard = isolated();
120 record_search("one", 1, 10);
121 record_search("two", 1, 20);
122 record_search("three", 1, 30);
123 let summary = search_summary();
124 assert_eq!(summary.unique_queries, 3);
125 assert_eq!(summary.repeated_queries, 0);
126 }
127}