Skip to main content

vtcode_memory/
query.rs

1//! Read-only queries across sessions for analytics and long-term learning.
2
3use std::collections::HashMap;
4use std::path::Path;
5use std::sync::Mutex;
6use std::time::SystemTime;
7
8use chrono::{DateTime, Utc};
9use lru::LruCache;
10
11use crate::error::SessionStoreError;
12use crate::sessions_root;
13
14const MANIFEST_CACHE_CAPACITY: usize = 200;
15const MANIFEST_CACHE_NONZERO_CAPACITY: std::num::NonZeroUsize =
16    std::num::NonZeroUsize::new(MANIFEST_CACHE_CAPACITY).unwrap();
17static MANIFEST_CACHE: std::sync::OnceLock<Mutex<LruCache<String, CachedManifest>>> = std::sync::OnceLock::new();
18
19#[derive(Debug, Clone)]
20struct CachedManifest {
21    summary: SessionSummary,
22    signature: ManifestSignature,
23}
24
25#[derive(Debug, Clone, Copy, PartialEq, Eq)]
26struct ManifestSignature {
27    modified: Option<SystemTime>,
28    len: u64,
29}
30
31fn manifest_signature(path: &Path) -> Option<ManifestSignature> {
32    let metadata = std::fs::symlink_metadata(path).ok()?;
33    Some(ManifestSignature {
34        modified: metadata.modified().ok(),
35        len: metadata.len(),
36    })
37}
38
39/// Invalidate a cached manifest after an atomic replacement.
40pub(crate) fn invalidate_manifest_cache(path: &Path) {
41    if let Some(cache) = MANIFEST_CACHE.get()
42        && let Ok(mut cache) = cache.lock()
43    {
44        let key = path.to_string_lossy();
45        cache.pop(key.as_ref());
46    }
47}
48
49/// Lightweight summary of a single session, read from its `manifest.json`.
50#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, PartialEq, Eq)]
51pub struct SessionSummary {
52    /// Session identifier (directory name).
53    pub session_id: String,
54    /// Number of completed turns.
55    pub turn_count: u64,
56    /// Total events recorded.
57    pub event_count: u64,
58    /// Lifecycle status.
59    pub status: String,
60    /// RFC3339 last-update timestamp (used for ordering).
61    pub updated_at: String,
62}
63
64/// A single grounded fact drawn from a session's memory envelope.
65#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, PartialEq, Eq)]
66pub struct FactRecord {
67    /// The fact text.
68    pub fact: String,
69    /// Session the fact originated from.
70    pub session_id: String,
71}
72
73/// One result returned from a memory search.
74///
75/// Mirrors the shape used by the grok-build memory subsystem so that
76/// higher-level consumers (tool bridge, context injection) can share
77/// formatting logic once a richer backend is available.
78#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, PartialEq)]
79pub struct MemorySearchResult {
80    /// Stable identifier for this chunk (session_id + fact index).
81    chunk_id: String,
82    /// Source memory file path.
83    path: String,
84    /// 0-based start line in the source file (0 for derived facts).
85    start_line: usize,
86    /// 0-based end line in the source file (0 for derived facts).
87    end_line: usize,
88    /// Relevance score (higher = more relevant).
89    score: f64,
90    /// Text snippet from the chunk.
91    snippet: String,
92    /// Source scope: `"session"` for per-session memory files.
93    source: String,
94    /// Unix timestamp (seconds) when the source memory was created.
95    created_at: Option<i64>,
96}
97
98/// List up to `n` most-recently-updated sessions.
99#[must_use]
100pub fn recent_sessions(workspace: &Path, n: usize) -> Vec<SessionSummary> {
101    let root = sessions_root(workspace);
102    if !root.exists() {
103        return Vec::new();
104    }
105    let mut out = Vec::new();
106    let entries = match std::fs::read_dir(&root) {
107        Ok(e) => e,
108        Err(_) => return Vec::new(),
109    };
110    let mut cache = MANIFEST_CACHE
111        .get_or_init(|| Mutex::new(LruCache::new(MANIFEST_CACHE_NONZERO_CAPACITY)))
112        .lock()
113        .unwrap_or_else(std::sync::PoisonError::into_inner);
114    for entry in entries.filter_map(Result::ok) {
115        let manifest = entry.path().join("manifest.json");
116        let key = manifest.to_string_lossy().into_owned();
117        let signature = manifest_signature(&manifest);
118        if let Some(cached) = cache.get(&key)
119            && signature == Some(cached.signature)
120        {
121            out.push(cached.summary.clone());
122            continue;
123        }
124        if let Ok(bytes) = std::fs::read(&manifest)
125            && let Ok(s) = serde_json::from_slice::<SessionSummary>(&bytes)
126        {
127            if let Some(signature) = signature.or_else(|| manifest_signature(&manifest)) {
128                cache.put(key, CachedManifest { summary: s.clone(), signature });
129            }
130            out.push(s);
131        }
132    }
133    out.sort_by(|a, b| b.updated_at.cmp(&a.updated_at));
134    out.truncate(n);
135    out
136}
137
138/// Cross-session long-term-learning query: collect grounded facts from every
139/// session's derived memory envelope. This is how the agent learns across
140/// sessions without loading any history into context.
141pub fn query_facts(workspace: &Path, limit: usize) -> Result<Vec<FactRecord>, SessionStoreError> {
142    let root = sessions_root(workspace);
143    if !root.exists() {
144        return Ok(Vec::new());
145    }
146    let mut facts: Vec<FactRecord> = Vec::new();
147    let entries = std::fs::read_dir(&root).map_err(|e| SessionStoreError::io(root.clone(), e))?;
148    for entry in entries.filter_map(Result::ok) {
149        let memory = entry.path().join(crate::DERIVED_DIR).join("memory.json");
150        let Ok(bytes) = std::fs::read(&memory) else {
151            continue;
152        };
153        let Ok(value) = serde_json::from_slice::<serde_json::Value>(&bytes) else {
154            continue;
155        };
156        let session_id = entry.file_name().to_string_lossy().into_owned();
157        if let Some(arr) = value.get("grounded_facts").and_then(|v| v.as_array()) {
158            for item in arr {
159                if let Some(fact) = item.get("fact").and_then(|f| f.as_str()) {
160                    facts.push(FactRecord {
161                        fact: fact.to_string(),
162                        session_id: session_id.clone(),
163                    });
164                }
165            }
166        }
167    }
168    facts.truncate(limit);
169    Ok(facts)
170}
171
172/// Cross-session memory search: scan every session's derived memory envelope
173/// for facts matching `query`. Returns up to `max_results` results with
174/// score >= `min_score`, sorted by descending relevance.
175///
176/// Scoring uses BM25 (`k1=1.2`, `b=0.75`) over tokenized facts. Ties are
177/// deterministic by the stable `session_id:index` chunk identifier.
178pub fn search_memory(
179    workspace: &Path,
180    query: &str,
181    max_results: usize,
182    min_score: f64,
183) -> Result<Vec<MemorySearchResult>, SessionStoreError> {
184    if query.is_empty() {
185        return Ok(Vec::new());
186    }
187
188    let root = sessions_root(workspace);
189    if !root.exists() {
190        return Ok(Vec::new());
191    }
192
193    let query_terms = tokenize(query);
194    if query_terms.is_empty() {
195        return Ok(Vec::new());
196    }
197    let mut results: Vec<MemorySearchResult> = Vec::new();
198    let session_source = String::from("session");
199    let mut documents = Vec::new();
200
201    let entries = std::fs::read_dir(&root).map_err(|e| SessionStoreError::io(root.clone(), e))?;
202    for entry in entries.filter_map(Result::ok) {
203        let session_dir = entry.path();
204        let memory = session_dir.join(crate::DERIVED_DIR).join("memory.json");
205        let Ok(bytes) = std::fs::read(&memory) else {
206            continue;
207        };
208        let Ok(value) = serde_json::from_slice::<serde_json::Value>(&bytes) else {
209            continue;
210        };
211        let session_id = entry.file_name().to_string_lossy().into_owned();
212        let memory_path = memory.to_string_lossy().into_owned();
213        let created_at = value
214            .get("created_at")
215            .and_then(parse_timestamp)
216            .or_else(|| value.get("updated_at").and_then(parse_timestamp));
217
218        if let Some(arr) = value.get("grounded_facts").and_then(|v| v.as_array()) {
219            for (idx, item) in arr.iter().enumerate() {
220                let Some(fact) = item.get("fact").and_then(|f| f.as_str()) else {
221                    continue;
222                };
223                documents.push((
224                    format!("{session_id}:{idx}"),
225                    memory_path.clone(),
226                    fact.to_owned(),
227                    tokenize(fact),
228                    created_at,
229                ));
230            }
231        }
232    }
233
234    let document_count = documents.len();
235    if document_count == 0 {
236        return Ok(Vec::new());
237    }
238    let average_length =
239        documents.iter().map(|(_, _, _, terms, _)| terms.len() as f64).sum::<f64>() / document_count as f64;
240    let mut document_frequency: HashMap<String, usize> = HashMap::new();
241    for (_, _, _, terms, _) in &documents {
242        let mut seen = std::collections::HashSet::new();
243        for term in terms {
244            if seen.insert(term.as_str()) {
245                *document_frequency.entry(term.clone()).or_insert(0) += 1;
246            }
247        }
248    }
249    let k1 = 1.2;
250    let b = 0.75;
251    let now = Utc::now();
252    for (chunk_id, path, fact, terms, created_at) in documents {
253        let length = terms.len() as f64;
254        let mut term_frequency = HashMap::<&str, usize>::new();
255        for term in &terms {
256            *term_frequency.entry(term.as_str()).or_insert(0) += 1;
257        }
258        let mut score = 0.0;
259        for query_term in &query_terms {
260            let Some(&frequency) = term_frequency.get(query_term.as_str()) else {
261                continue;
262            };
263            let Some(&frequency_in_documents) = document_frequency.get(query_term) else {
264                continue;
265            };
266            let idf = (((document_count - frequency_in_documents) as f64 + 0.5)
267                / (frequency_in_documents as f64 + 0.5)
268                + 1.0)
269                .ln();
270            let denominator = frequency as f64 + k1 * (1.0 - b + b * length / average_length.max(1.0));
271            score += idf * (frequency as f64 * (k1 + 1.0)) / denominator;
272        }
273        if score <= 0.0 {
274            continue;
275        }
276        // A mild recency preference helps current session facts win close BM25
277        // matches without allowing recency to create a result that did not
278        // match the query. Missing or malformed timestamps deliberately keep
279        // the neutral multiplier of 1.0.
280        score *= recency_multiplier(created_at, now);
281        if score < min_score {
282            continue;
283        }
284        results.push(MemorySearchResult {
285            chunk_id,
286            path,
287            start_line: 0,
288            end_line: 0,
289            score,
290            snippet: fact,
291            source: session_source.clone(),
292            created_at,
293        });
294    }
295
296    results.sort_by(|a, b| {
297        b.score
298            .partial_cmp(&a.score)
299            .unwrap_or(std::cmp::Ordering::Equal)
300            .then_with(|| a.chunk_id.cmp(&b.chunk_id))
301    });
302    results.truncate(max_results);
303    Ok(results)
304}
305
306fn parse_timestamp(value: &serde_json::Value) -> Option<i64> {
307    if let Some(timestamp) = value.as_i64() {
308        return DateTime::<Utc>::from_timestamp(timestamp, 0).map(|_| timestamp);
309    }
310    let text = value.as_str()?.trim();
311    if let Ok(timestamp) = text.parse::<i64>() {
312        return DateTime::<Utc>::from_timestamp(timestamp, 0).map(|_| timestamp);
313    }
314    DateTime::parse_from_rfc3339(text).ok().map(|timestamp| timestamp.timestamp())
315}
316
317fn recency_multiplier(created_at: Option<i64>, now: DateTime<Utc>) -> f64 {
318    let Some(created_at) = created_at.and_then(|timestamp| DateTime::<Utc>::from_timestamp(timestamp, 0)) else {
319        return 1.0;
320    };
321    let age_days = (now - created_at).num_seconds().max(0) as f64 / 86_400.0;
322    1.0 + 0.15 * 0.5_f64.powf(age_days / 30.0)
323}
324
325/// Return the configured default for `max_results` in search queries.
326pub fn default_search_max_results() -> usize {
327    6
328}
329
330/// Return the configured default for `min_score` in search queries.
331pub fn default_search_min_score() -> f64 {
332    0.0
333}
334
335fn count_substring_matches(text: &str, lowered_query: &str) -> usize {
336    if lowered_query.is_empty() {
337        return 0;
338    }
339    let lowered = text.to_ascii_lowercase();
340    let mut count = 0;
341    let mut start = 0;
342    while let Some(pos) = lowered[start..].find(lowered_query) {
343        count += 1;
344        start += pos + lowered_query.len();
345    }
346    count
347}
348
349fn tokenize(text: &str) -> Vec<String> {
350    let mut tokens = Vec::new();
351    let mut current = String::new();
352    for character in text.chars() {
353        if character.is_alphanumeric() {
354            current.extend(character.to_lowercase());
355        } else if !current.is_empty() {
356            tokens.push(std::mem::take(&mut current));
357        }
358    }
359    if !current.is_empty() {
360        tokens.push(current);
361    }
362    tokens
363}
364
365#[cfg(test)]
366mod tests {
367    use super::*;
368    use tempfile::TempDir;
369
370    #[test]
371    fn count_substring_matches_counts_overlapping() {
372        assert_eq!(count_substring_matches("aaaa", "aa"), 2);
373        assert_eq!(count_substring_matches("ababa", "aba"), 1);
374        assert_eq!(count_substring_matches("hello world", "ll"), 1);
375        assert_eq!(count_substring_matches("", "x"), 0);
376    }
377
378    #[test]
379    fn search_memory_returns_matching_facts() {
380        let dir = TempDir::new().expect("tempdir");
381        let sess = crate::session_dir(dir.path(), "s1");
382        std::fs::create_dir_all(sess.join(crate::DERIVED_DIR)).expect("mkdir");
383        let memory = serde_json::json!({
384            "grounded_facts": [
385                {"fact": "the widget is blue"},
386                {"fact": "the server runs on port 8080"},
387                {"fact": "use PostgreSQL for persistence"},
388            ]
389        });
390        std::fs::write(sess.join(crate::DERIVED_DIR).join("memory.json"), serde_json::to_string(&memory).expect("ser"))
391            .expect("write");
392
393        let results = search_memory(dir.path(), "blue", 10, 0.0).expect("search");
394        assert_eq!(results.len(), 1);
395        assert_eq!(results[0].snippet, "the widget is blue");
396        assert_eq!(results[0].chunk_id, "s1:0");
397        assert!(results[0].score > 0.0);
398    }
399
400    #[test]
401    fn search_memory_scores_multiple_matches() {
402        let dir = TempDir::new().expect("tempdir");
403        let sess = crate::session_dir(dir.path(), "s2");
404        std::fs::create_dir_all(sess.join(crate::DERIVED_DIR)).expect("mkdir");
405        let memory = serde_json::json!({
406            "grounded_facts": [
407                {"fact": "rust uses rustc and cargo"},
408                {"fact": "cargo is the rust build tool"},
409            ]
410        });
411        std::fs::write(sess.join(crate::DERIVED_DIR).join("memory.json"), serde_json::to_string(&memory).expect("ser"))
412            .expect("write");
413
414        let results = search_memory(dir.path(), "cargo", 10, 0.0).expect("search");
415        assert_eq!(results.len(), 2);
416        assert!(results.iter().all(|r| r.score > 0.0));
417    }
418
419    #[test]
420    fn search_memory_respects_min_score() {
421        let dir = TempDir::new().expect("tempdir");
422        let sess = crate::session_dir(dir.path(), "s3");
423        std::fs::create_dir_all(sess.join(crate::DERIVED_DIR)).expect("mkdir");
424        let memory = serde_json::json!({
425            "grounded_facts": [
426                {"fact": "alpha beta gamma"},
427            ]
428        });
429        std::fs::write(sess.join(crate::DERIVED_DIR).join("memory.json"), serde_json::to_string(&memory).expect("ser"))
430            .expect("write");
431
432        let results = search_memory(dir.path(), "beta", 10, 2.0).expect("search");
433        assert!(results.is_empty());
434    }
435
436    #[test]
437    fn search_memory_empty_query_returns_empty() {
438        let dir = TempDir::new().expect("tempdir");
439        let results = search_memory(dir.path(), "", 10, 0.0).expect("search");
440        assert!(results.is_empty());
441    }
442
443    #[test]
444    fn search_memory_sorts_by_score_descending() {
445        let dir = TempDir::new().expect("tempdir");
446        for i in 0..3 {
447            let sess = crate::session_dir(dir.path(), &format!("s{i}"));
448            std::fs::create_dir_all(sess.join(crate::DERIVED_DIR)).expect("mkdir");
449            let memory = serde_json::json!({
450                "grounded_facts": [
451                    {"fact": format!("fact {i} appears twice twice")},
452                ]
453            });
454            std::fs::write(
455                sess.join(crate::DERIVED_DIR).join("memory.json"),
456                serde_json::to_string(&memory).expect("ser"),
457            )
458            .expect("write");
459        }
460
461        let results = search_memory(dir.path(), "twice", 10, 0.0).expect("search");
462        assert_eq!(results.len(), 3);
463        assert!(results.windows(2).all(|w| w[0].score >= w[1].score));
464    }
465
466    #[test]
467    fn search_memory_uses_bm25_term_coverage_and_deterministic_ties() {
468        let dir = TempDir::new().expect("tempdir");
469        for (session_id, facts) in [
470            ("s1", vec!["rust cargo tool", "unrelated note"]),
471            ("s2", vec!["cargo build tool"]),
472        ] {
473            let session = crate::session_dir(dir.path(), session_id);
474            std::fs::create_dir_all(session.join(crate::DERIVED_DIR)).expect("mkdir");
475            let memory = serde_json::json!({
476                "grounded_facts": facts.into_iter().map(|fact| serde_json::json!({"fact": fact})).collect::<Vec<_>>()
477            });
478            std::fs::write(
479                session.join(crate::DERIVED_DIR).join("memory.json"),
480                serde_json::to_string(&memory).expect("serialize"),
481            )
482            .expect("write");
483        }
484
485        let results = search_memory(dir.path(), "rust cargo", 10, 0.0).expect("search");
486        assert_eq!(results.first().map(|result| result.chunk_id.as_str()), Some("s1:0"));
487        assert!(results[0].score > results[1].score);
488    }
489
490    #[test]
491    fn search_memory_applies_recency_to_valid_rfc3339_timestamps() {
492        let dir = TempDir::new().expect("tempdir");
493        for (session_id, timestamp) in [("old", "2020-01-01T00:00:00Z"), ("new", "2099-01-01T00:00:00Z")] {
494            let session = crate::session_dir(dir.path(), session_id);
495            std::fs::create_dir_all(session.join(crate::DERIVED_DIR)).expect("mkdir");
496            let memory = serde_json::json!({
497                "created_at": timestamp,
498                "grounded_facts": [{"fact": "shared search fact"}]
499            });
500            std::fs::write(
501                session.join(crate::DERIVED_DIR).join("memory.json"),
502                serde_json::to_vec(&memory).expect("serialize"),
503            )
504            .expect("write");
505        }
506
507        let results = search_memory(dir.path(), "shared", 10, 0.0).expect("search");
508        assert_eq!(results.len(), 2);
509        assert_eq!(results[0].chunk_id, "new:0");
510        assert!(results[0].score > results[1].score);
511    }
512
513    #[test]
514    fn search_memory_does_not_boost_missing_or_invalid_timestamps() {
515        let dir = TempDir::new().expect("tempdir");
516        for (session_id, timestamp) in [
517            ("missing", serde_json::Value::Null),
518            ("invalid", serde_json::json!("not-a-date")),
519        ] {
520            let session = crate::session_dir(dir.path(), session_id);
521            std::fs::create_dir_all(session.join(crate::DERIVED_DIR)).expect("mkdir");
522            let mut memory = serde_json::json!({"grounded_facts": [{"fact": "neutral search fact"}]});
523            if !timestamp.is_null() {
524                memory["created_at"] = timestamp;
525            }
526            std::fs::write(
527                session.join(crate::DERIVED_DIR).join("memory.json"),
528                serde_json::to_vec(&memory).expect("serialize"),
529            )
530            .expect("write");
531        }
532
533        let results = search_memory(dir.path(), "neutral", 10, 0.0).expect("search");
534        assert_eq!(results.len(), 2);
535        assert!((results[0].score - results[1].score).abs() < f64::EPSILON);
536        assert_eq!(results[0].created_at, None);
537        assert_eq!(results[1].created_at, None);
538    }
539
540    #[test]
541    fn recency_multiplier_matches_thirty_day_half_life() {
542        let now = DateTime::parse_from_rfc3339("2026-01-31T00:00:00Z")
543            .expect("timestamp")
544            .with_timezone(&Utc);
545        let thirty_days_ago = now - chrono::Duration::days(30);
546        let multiplier = recency_multiplier(Some(thirty_days_ago.timestamp()), now);
547        assert!((multiplier - 1.075).abs() < 1e-12);
548        assert!((recency_multiplier(None, now) - 1.0).abs() < f64::EPSILON);
549    }
550
551    #[test]
552    fn recent_sessions_invalidates_manifest_cache_after_replacement() {
553        let dir = TempDir::new().expect("tempdir");
554        let session = crate::session_dir(dir.path(), "cache-session");
555        std::fs::create_dir_all(&session).expect("mkdir");
556        let manifest = serde_json::json!({
557            "session_id": "cache-session",
558            "schema_version": 1,
559            "created_at": "2026-01-01T00:00:00Z",
560            "updated_at": "2026-01-01T00:00:00Z",
561            "turn_count": 1,
562            "event_count": 1,
563            "status": "active"
564        });
565        let path = session.join("manifest.json");
566        std::fs::write(&path, serde_json::to_vec(&manifest).expect("serialize")).expect("write");
567        assert_eq!(recent_sessions(dir.path(), 1)[0].updated_at, "2026-01-01T00:00:00Z");
568
569        let mut replaced = manifest;
570        replaced["updated_at"] = serde_json::json!("2099-01-01T00:00:00Z");
571        std::fs::write(&path, serde_json::to_vec(&replaced).expect("serialize")).expect("replace");
572        assert_eq!(recent_sessions(dir.path(), 1)[0].updated_at, "2099-01-01T00:00:00Z");
573    }
574}