lean_ctx/tools/
search_kernel.rs

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

lean_ctx/tools/
search_kernel.rs
