Skip to main content

vtcode_core/compaction/
prefire.rs

1use std::sync::Mutex;
2use std::sync::atomic::AtomicBool;
3use std::sync::atomic::Ordering;
4
5/// Cached result of an async (background / prefire) pass-1 sample for
6/// two-pass compaction. Held between the background pass-1 and the
7/// synchronous pass-2 apply at compaction time.
8#[derive(Debug, Clone)]
9pub struct AsyncCompactionCache {
10    /// The successor-usable NOTE₁ text (extracted `<summary>` or full pass-1 output).
11    pub note1: String,
12    /// Number of leading conversation items pass-1 summarized.
13    pub prefix_len: usize,
14    /// Fingerprint of `conversation[..prefix_len]` at pass-1 time. Pass-2 only
15    /// applies NOTE₁ when the current conversation still has this exact prefix.
16    pub fingerprint: u64,
17    /// Model slug pass-1 ran under; invalidated on model switch.
18    pub model_slug: String,
19    /// Wall time pass-1 took (ms).
20    pub pass1_latency_ms: u64,
21}
22
23/// Prefire two-pass state. Manages the in-flight guard and cached async
24/// pass-1 result between turns.
25#[derive(Default)]
26pub struct PrefireState {
27    /// Set while a background pass-1 sample is running.
28    in_flight: AtomicBool,
29    /// Cached async pass-1 result, ready for pass-2 apply.
30    cache: Mutex<Option<AsyncCompactionCache>>,
31}
32
33impl PrefireState {
34    /// Try to claim the single in-flight slot. Returns `true` iff this caller
35    /// won the race and should start pass-1.
36    pub fn try_begin(&self) -> bool {
37        self.in_flight
38            .compare_exchange(false, true, Ordering::AcqRel, Ordering::Relaxed)
39            .is_ok()
40    }
41
42    /// Release the in-flight slot.
43    pub fn finish(&self) {
44        self.in_flight.store(false, Ordering::Release);
45    }
46
47    /// Whether a pass-1 is currently running.
48    pub fn is_in_flight(&self) -> bool {
49        self.in_flight.load(Ordering::Acquire)
50    }
51
52    /// Stash a completed pass-1 cache for later pass-2 use.
53    pub fn store(&self, cache: AsyncCompactionCache) {
54        *self.cache.lock().unwrap_or_else(std::sync::PoisonError::into_inner) = Some(cache);
55    }
56
57    /// Take the cache, leaving `None`.
58    pub fn take(&self) -> Option<AsyncCompactionCache> {
59        self.cache.lock().unwrap_or_else(std::sync::PoisonError::into_inner).take()
60    }
61
62    /// Drop any cached async pass-1 result.
63    pub fn clear(&self) {
64        *self.cache.lock().unwrap_or_else(std::sync::PoisonError::into_inner) = None;
65    }
66
67    /// Whether a valid cache is available.
68    pub fn has_cache(&self) -> bool {
69        self.cache.lock().unwrap_or_else(std::sync::PoisonError::into_inner).is_some()
70    }
71}
72
73#[cfg(test)]
74mod tests {
75    use super::*;
76
77    #[test]
78    fn try_begin_blocks_concurrent_wins() {
79        let state = PrefireState::default();
80        assert!(state.try_begin());
81        assert!(!state.try_begin());
82        state.finish();
83        assert!(state.try_begin());
84    }
85
86    #[test]
87    fn store_take_roundtrip() {
88        let state = PrefireState::default();
89        state.store(AsyncCompactionCache {
90            note1: "note".to_string(),
91            prefix_len: 3,
92            fingerprint: 42,
93            model_slug: "model".to_string(),
94            pass1_latency_ms: 5,
95        });
96        assert!(state.has_cache());
97        let cache = state.take().unwrap();
98        assert_eq!(cache.note1, "note");
99        assert_eq!(cache.prefix_len, 3);
100        assert!(!state.has_cache());
101    }
102
103    #[test]
104    fn clear_drops_cache() {
105        let state = PrefireState::default();
106        state.store(AsyncCompactionCache {
107            note1: "note".to_string(),
108            prefix_len: 3,
109            fingerprint: 42,
110            model_slug: "model".to_string(),
111            pass1_latency_ms: 5,
112        });
113        assert!(state.has_cache());
114        state.clear();
115        assert!(!state.has_cache());
116    }
117
118    #[test]
119    fn is_in_flight_reflects_try_begin() {
120        let state = PrefireState::default();
121        assert!(!state.is_in_flight());
122        state.try_begin();
123        assert!(state.is_in_flight());
124        state.finish();
125        assert!(!state.is_in_flight());
126    }
127}