Skip to main content

vtcode_memory/
query.rs

1//! Read-only queries across sessions for analytics and long-term learning.
2
3use std::path::Path;
4use std::sync::Mutex;
5
6use lru::LruCache;
7
8use crate::error::SessionStoreError;
9use crate::sessions_root;
10
11const MANIFEST_CACHE_CAPACITY: usize = 200;
12const MANIFEST_CACHE_NONZERO_CAPACITY: std::num::NonZeroUsize =
13    std::num::NonZeroUsize::new(MANIFEST_CACHE_CAPACITY).unwrap();
14static MANIFEST_CACHE: std::sync::OnceLock<Mutex<LruCache<String, SessionSummary>>> = std::sync::OnceLock::new();
15
16/// Lightweight summary of a single session, read from its `manifest.json`.
17#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, PartialEq, Eq)]
18pub struct SessionSummary {
19    /// Session identifier (directory name).
20    pub session_id: String,
21    /// Number of completed turns.
22    pub turn_count: u64,
23    /// Total events recorded.
24    pub event_count: u64,
25    /// Lifecycle status.
26    pub status: String,
27    /// RFC3339 last-update timestamp (used for ordering).
28    pub updated_at: String,
29}
30
31/// A single grounded fact drawn from a session's memory envelope.
32#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, PartialEq, Eq)]
33pub struct FactRecord {
34    /// The fact text.
35    pub fact: String,
36    /// Session the fact originated from.
37    pub session_id: String,
38}
39
40/// One result returned from a memory search.
41///
42/// Mirrors the shape used by the grok-build memory subsystem so that
43/// higher-level consumers (tool bridge, context injection) can share
44/// formatting logic once a richer backend is available.
45#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, PartialEq)]
46pub struct MemorySearchResult {
47    /// Stable identifier for this chunk (session_id + fact index).
48    chunk_id: String,
49    /// Source memory file path.
50    path: String,
51    /// 0-based start line in the source file (0 for derived facts).
52    start_line: usize,
53    /// 0-based end line in the source file (0 for derived facts).
54    end_line: usize,
55    /// Relevance score (higher = more relevant).
56    score: f64,
57    /// Text snippet from the chunk.
58    snippet: String,
59    /// Source scope: `"session"` for per-session memory files.
60    source: String,
61    /// Unix timestamp (seconds) when the source memory was created.
62    created_at: Option<i64>,
63}
64
65/// List up to `n` most-recently-updated sessions.
66#[must_use]
67pub fn recent_sessions(workspace: &Path, n: usize) -> Vec<SessionSummary> {
68    let root = sessions_root(workspace);
69    if !root.exists() {
70        return Vec::new();
71    }
72    let mut out = Vec::new();
73    let entries = match std::fs::read_dir(&root) {
74        Ok(e) => e,
75        Err(_) => return Vec::new(),
76    };
77    let mut cache = MANIFEST_CACHE
78        .get_or_init(|| Mutex::new(LruCache::new(MANIFEST_CACHE_NONZERO_CAPACITY)))
79        .lock()
80        .unwrap_or_else(std::sync::PoisonError::into_inner);
81    for entry in entries.filter_map(Result::ok) {
82        let manifest = entry.path().join("manifest.json");
83        let key = manifest.to_string_lossy().into_owned();
84        if let Some(s) = cache.get(&key) {
85            out.push(s.clone());
86            continue;
87        }
88        if let Ok(bytes) = std::fs::read(&manifest)
89            && let Ok(s) = serde_json::from_slice::<SessionSummary>(&bytes)
90        {
91            cache.put(key, s.clone());
92            out.push(s);
93        }
94    }
95    out.sort_by(|a, b| b.updated_at.cmp(&a.updated_at));
96    out.truncate(n);
97    out
98}
99
100/// Cross-session long-term-learning query: collect grounded facts from every
101/// session's derived memory envelope. This is how the agent learns across
102/// sessions without loading any history into context.
103pub fn query_facts(workspace: &Path, limit: usize) -> Result<Vec<FactRecord>, SessionStoreError> {
104    let root = sessions_root(workspace);
105    if !root.exists() {
106        return Ok(Vec::new());
107    }
108    let mut facts: Vec<FactRecord> = Vec::new();
109    let entries = std::fs::read_dir(&root).map_err(|e| SessionStoreError::io(root.clone(), e))?;
110    for entry in entries.filter_map(Result::ok) {
111        let memory = entry.path().join(crate::DERIVED_DIR).join("memory.json");
112        let Ok(bytes) = std::fs::read(&memory) else {
113            continue;
114        };
115        let Ok(value) = serde_json::from_slice::<serde_json::Value>(&bytes) else {
116            continue;
117        };
118        let session_id = entry.file_name().to_string_lossy().into_owned();
119        if let Some(arr) = value.get("grounded_facts").and_then(|v| v.as_array()) {
120            for item in arr {
121                if let Some(fact) = item.get("fact").and_then(|f| f.as_str()) {
122                    facts.push(FactRecord {
123                        fact: fact.to_string(),
124                        session_id: session_id.clone(),
125                    });
126                }
127            }
128        }
129    }
130    facts.truncate(limit);
131    Ok(facts)
132}
133
134/// Cross-session memory search: scan every session's derived memory envelope
135/// for facts matching `query`. Returns up to `max_results` results with
136/// score >= `min_score`, sorted by descending relevance.
137///
138/// Scoring is based on the number of case-insensitive query matches found
139/// in each fact. A simple BMH-style substring count keeps this zero-dependency
140/// while still surfacing the most-relevant chunks first.
141pub fn search_memory(
142    workspace: &Path,
143    query: &str,
144    max_results: usize,
145    min_score: f64,
146) -> Result<Vec<MemorySearchResult>, SessionStoreError> {
147    if query.is_empty() {
148        return Ok(Vec::new());
149    }
150
151    let root = sessions_root(workspace);
152    if !root.exists() {
153        return Ok(Vec::new());
154    }
155
156    let lowered = query.to_ascii_lowercase();
157    let mut results: Vec<MemorySearchResult> = Vec::new();
158    let session_source = String::from("session");
159
160    let entries = std::fs::read_dir(&root).map_err(|e| SessionStoreError::io(root.clone(), e))?;
161    for entry in entries.filter_map(Result::ok) {
162        let session_dir = entry.path();
163        let memory = session_dir.join(crate::DERIVED_DIR).join("memory.json");
164        let Ok(bytes) = std::fs::read(&memory) else {
165            continue;
166        };
167        let Ok(value) = serde_json::from_slice::<serde_json::Value>(&bytes) else {
168            continue;
169        };
170        let session_id = entry.file_name().to_string_lossy().into_owned();
171        let memory_path = memory.to_string_lossy().into_owned();
172        let created_at = value
173            .get("created_at")
174            .and_then(|v| v.as_i64())
175            .or_else(|| value.get("updated_at").and_then(|v| v.as_i64()));
176
177        if let Some(arr) = value.get("grounded_facts").and_then(|v| v.as_array()) {
178            for (idx, item) in arr.iter().enumerate() {
179                let Some(fact) = item.get("fact").and_then(|f| f.as_str()) else {
180                    continue;
181                };
182                let score = count_substring_matches(fact, &lowered) as f64;
183                if score <= 0.0 || score < min_score {
184                    continue;
185                }
186                results.push(MemorySearchResult {
187                    chunk_id: format!("{session_id}:{idx}"),
188                    path: memory_path.clone(),
189                    start_line: 0,
190                    end_line: 0,
191                    score,
192                    snippet: fact.to_string(),
193                    source: session_source.clone(),
194                    created_at,
195                });
196            }
197        }
198    }
199
200    results.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap_or(std::cmp::Ordering::Equal));
201    results.truncate(max_results);
202    Ok(results)
203}
204
205/// Return the configured default for `max_results` in search queries.
206pub fn default_search_max_results() -> usize {
207    6
208}
209
210/// Return the configured default for `min_score` in search queries.
211pub fn default_search_min_score() -> f64 {
212    0.0
213}
214
215fn count_substring_matches(text: &str, lowered_query: &str) -> usize {
216    if lowered_query.is_empty() {
217        return 0;
218    }
219    let lowered = text.to_ascii_lowercase();
220    let mut count = 0;
221    let mut start = 0;
222    while let Some(pos) = lowered[start..].find(lowered_query) {
223        count += 1;
224        start += pos + lowered_query.len();
225    }
226    count
227}
228
229#[cfg(test)]
230mod tests {
231    use super::*;
232    use tempfile::TempDir;
233
234    #[test]
235    fn count_substring_matches_counts_overlapping() {
236        assert_eq!(count_substring_matches("aaaa", "aa"), 2);
237        assert_eq!(count_substring_matches("ababa", "aba"), 1);
238        assert_eq!(count_substring_matches("hello world", "ll"), 1);
239        assert_eq!(count_substring_matches("", "x"), 0);
240    }
241
242    #[test]
243    fn search_memory_returns_matching_facts() {
244        let dir = TempDir::new().expect("tempdir");
245        let sess = crate::session_dir(dir.path(), "s1");
246        std::fs::create_dir_all(sess.join(crate::DERIVED_DIR)).expect("mkdir");
247        let memory = serde_json::json!({
248            "grounded_facts": [
249                {"fact": "the widget is blue"},
250                {"fact": "the server runs on port 8080"},
251                {"fact": "use PostgreSQL for persistence"},
252            ]
253        });
254        std::fs::write(sess.join(crate::DERIVED_DIR).join("memory.json"), serde_json::to_string(&memory).expect("ser"))
255            .expect("write");
256
257        let results = search_memory(dir.path(), "blue", 10, 0.0).expect("search");
258        assert_eq!(results.len(), 1);
259        assert_eq!(results[0].snippet, "the widget is blue");
260        assert_eq!(results[0].chunk_id, "s1:0");
261        assert!((results[0].score - 1.0).abs() < f64::EPSILON);
262    }
263
264    #[test]
265    fn search_memory_scores_multiple_matches() {
266        let dir = TempDir::new().expect("tempdir");
267        let sess = crate::session_dir(dir.path(), "s2");
268        std::fs::create_dir_all(sess.join(crate::DERIVED_DIR)).expect("mkdir");
269        let memory = serde_json::json!({
270            "grounded_facts": [
271                {"fact": "rust uses rustc and cargo"},
272                {"fact": "cargo is the rust build tool"},
273            ]
274        });
275        std::fs::write(sess.join(crate::DERIVED_DIR).join("memory.json"), serde_json::to_string(&memory).expect("ser"))
276            .expect("write");
277
278        let results = search_memory(dir.path(), "cargo", 10, 0.0).expect("search");
279        assert_eq!(results.len(), 2);
280        assert!(results.iter().all(|r| (r.score - 1.0).abs() < f64::EPSILON));
281    }
282
283    #[test]
284    fn search_memory_respects_min_score() {
285        let dir = TempDir::new().expect("tempdir");
286        let sess = crate::session_dir(dir.path(), "s3");
287        std::fs::create_dir_all(sess.join(crate::DERIVED_DIR)).expect("mkdir");
288        let memory = serde_json::json!({
289            "grounded_facts": [
290                {"fact": "alpha beta gamma"},
291            ]
292        });
293        std::fs::write(sess.join(crate::DERIVED_DIR).join("memory.json"), serde_json::to_string(&memory).expect("ser"))
294            .expect("write");
295
296        let results = search_memory(dir.path(), "beta", 10, 2.0).expect("search");
297        assert!(results.is_empty());
298    }
299
300    #[test]
301    fn search_memory_empty_query_returns_empty() {
302        let dir = TempDir::new().expect("tempdir");
303        let results = search_memory(dir.path(), "", 10, 0.0).expect("search");
304        assert!(results.is_empty());
305    }
306
307    #[test]
308    fn search_memory_sorts_by_score_descending() {
309        let dir = TempDir::new().expect("tempdir");
310        for i in 0..3 {
311            let sess = crate::session_dir(dir.path(), &format!("s{i}"));
312            std::fs::create_dir_all(sess.join(crate::DERIVED_DIR)).expect("mkdir");
313            let memory = serde_json::json!({
314                "grounded_facts": [
315                    {"fact": format!("fact {i} appears twice twice")},
316                ]
317            });
318            std::fs::write(
319                sess.join(crate::DERIVED_DIR).join("memory.json"),
320                serde_json::to_string(&memory).expect("ser"),
321            )
322            .expect("write");
323        }
324
325        let results = search_memory(dir.path(), "twice", 10, 0.0).expect("search");
326        assert_eq!(results.len(), 3);
327        assert!(results.windows(2).all(|w| w[0].score >= w[1].score));
328    }
329}