Skip to main content

lean_ctx/core/
search_kernel.rs

1// SPDX-License-Identifier: Apache-2.0
2//! Evidence and repeated-query tracking for the search tool.
3
4use 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/// Cumulative search activity for the current session.
16#[derive(Debug, Clone, Default, serde::Serialize)]
17pub struct SearchSummary {
18    /// Number of recorded searches.
19    pub total_searches: usize,
20    /// Number of distinct query hashes.
21    pub unique_queries: usize,
22    /// Number of searches whose query was previously recorded.
23    pub repeated_queries: usize,
24    /// Number of result tokens across all recorded searches.
25    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
34/// Records search evidence and remembers the query for repetition detection.
35pub 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/// Returns whether the query has already been recorded in this session.
51#[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/// Returns cumulative search activity for the current session.
61#[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
76/// Clears all recorded search activity.
77pub 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}