ruvector_sona/loops/
coordinator.rs

1//! Loop Coordinator - Orchestrates all learning loops
2
3use crate::ewc::{EwcConfig, EwcPlusPlus};
4use crate::loops::background::{BackgroundLoop, BackgroundLoopConfig, BackgroundResult};
5use crate::loops::instant::InstantLoop;
6use crate::lora::{BaseLoRA, MicroLoRA};
7use crate::reasoning_bank::{PatternConfig, ReasoningBank};
8use crate::types::{QueryTrajectory, SonaConfig};
9use parking_lot::RwLock;
10use std::sync::Arc;
11
12/// Loop coordinator managing all learning loops
13pub struct LoopCoordinator {
14    /// Configuration
15    _config: SonaConfig,
16    /// Instant loop (Loop A)
17    instant: InstantLoop,
18    /// Background loop (Loop B)
19    background: BackgroundLoop,
20    /// Shared components
21    reasoning_bank: Arc<RwLock<ReasoningBank>>,
22    ewc: Arc<RwLock<EwcPlusPlus>>,
23    base_lora: Arc<RwLock<BaseLoRA>>,
24    /// Enabled flags
25    instant_enabled: bool,
26    background_enabled: bool,
27}
28
29impl LoopCoordinator {
30    /// Create new coordinator with default config
31    pub fn new(hidden_dim: usize) -> Self {
32        Self::with_config(SonaConfig {
33            hidden_dim,
34            embedding_dim: hidden_dim,
35            ..Default::default()
36        })
37    }
38
39    /// Create with custom config
40    pub fn with_config(config: SonaConfig) -> Self {
41        let reasoning_bank = Arc::new(RwLock::new(ReasoningBank::new(PatternConfig {
42            embedding_dim: config.embedding_dim,
43            k_clusters: config.pattern_clusters,
44            ..Default::default()
45        })));
46
47        let ewc = Arc::new(RwLock::new(EwcPlusPlus::new(EwcConfig {
48            param_count: config.hidden_dim * config.base_lora_rank * 2,
49            initial_lambda: config.ewc_lambda,
50            ..Default::default()
51        })));
52
53        let base_lora = Arc::new(RwLock::new(BaseLoRA::new(
54            config.hidden_dim,
55            config.base_lora_rank,
56            12, // Default number of layers
57        )));
58
59        let instant = InstantLoop::from_sona_config(&config);
60        let background = BackgroundLoop::new(
61            BackgroundLoopConfig::from(&config),
62            reasoning_bank.clone(),
63            ewc.clone(),
64            base_lora.clone(),
65        );
66
67        Self {
68            _config: config,
69            instant,
70            background,
71            reasoning_bank,
72            ewc,
73            base_lora,
74            instant_enabled: true,
75            background_enabled: true,
76        }
77    }
78
79    /// Process inference trajectory (Loop A)
80    pub fn on_inference(&self, trajectory: QueryTrajectory) {
81        if self.instant_enabled {
82            self.instant.on_trajectory(trajectory);
83        }
84    }
85
86    /// Generate next trajectory ID
87    pub fn next_trajectory_id(&self) -> u64 {
88        self.instant.next_id()
89    }
90
91    /// Run background cycle if needed (Loop B)
92    pub fn maybe_run_background(&self) -> Option<BackgroundResult> {
93        if !self.background_enabled {
94            return None;
95        }
96
97        if self.background.should_run() {
98            let trajectories = self.instant.drain_trajectories();
99            if !trajectories.is_empty() {
100                return Some(self.background.run_cycle(trajectories, false));
101            }
102        }
103
104        None
105    }
106
107    /// Force background cycle (bypasses minimum trajectory check)
108    pub fn force_background(&self) -> BackgroundResult {
109        let trajectories = self.instant.drain_trajectories();
110        self.background.run_cycle(trajectories, true)
111    }
112
113    /// Flush instant loop updates
114    pub fn flush_instant(&self) {
115        self.instant.flush();
116    }
117
118    /// Get micro-LoRA for inference
119    pub fn micro_lora(&self) -> &Arc<RwLock<MicroLoRA>> {
120        self.instant.micro_lora()
121    }
122
123    /// Get micro-LoRA weights for external persistence (snapshot save)
124    ///
125    /// Returns cloned owned vectors so caller can serialize without lifetime issues.
126    pub fn get_micro_lora_weights(&self) -> (Vec<f32>, Vec<f32>) {
127        let guard = self.instant.micro_lora().read();
128        let (down, up) = guard.get_weights();
129        (down.clone(), up.clone())
130    }
131
132    /// Restore micro-LoRA weights from external persistence (snapshot load)
133    pub fn restore_micro_lora_weights(&self, down: Vec<f32>, up: Vec<f32>) -> Result<(), String> {
134        self.instant.micro_lora().write().set_weights(down, up)
135    }
136
137    /// Get base-LoRA for inference
138    pub fn base_lora(&self) -> &Arc<RwLock<BaseLoRA>> {
139        &self.base_lora
140    }
141
142    /// Get reasoning bank
143    pub fn reasoning_bank(&self) -> &Arc<RwLock<ReasoningBank>> {
144        &self.reasoning_bank
145    }
146
147    /// Get EWC++
148    pub fn ewc(&self) -> &Arc<RwLock<EwcPlusPlus>> {
149        &self.ewc
150    }
151
152    /// Enable/disable instant loop
153    pub fn set_instant_enabled(&mut self, enabled: bool) {
154        self.instant_enabled = enabled;
155    }
156
157    /// Enable/disable background loop
158    pub fn set_background_enabled(&mut self, enabled: bool) {
159        self.background_enabled = enabled;
160    }
161
162    /// Get statistics
163    pub fn stats(&self) -> CoordinatorStats {
164        let (buffer_len, dropped, success_rate) = self.instant.buffer_stats();
165
166        CoordinatorStats {
167            trajectories_recorded: buffer_len as u64 + dropped,
168            trajectories_buffered: buffer_len,
169            trajectories_dropped: dropped,
170            buffer_success_rate: success_rate,
171            patterns_stored: self.reasoning_bank.read().pattern_count(),
172            patterns_learned: self.reasoning_bank.read().pattern_count(),
173            ewc_tasks: self.ewc.read().task_count(),
174            instant_enabled: self.instant_enabled,
175            background_enabled: self.background_enabled,
176        }
177    }
178
179    /// Serialize full state to JSON for persistence (fixes #274)
180    pub fn serialize_state(&self) -> String {
181        let rb = self.reasoning_bank.read();
182        let patterns = rb.get_all_patterns();
183        let ewc = self.ewc.read();
184        serde_json::json!({
185            "version": 1,
186            "patterns": patterns,
187            "ewc_task_count": ewc.task_count(),
188            "instant_enabled": self.instant_enabled,
189            "background_enabled": self.background_enabled,
190        })
191        .to_string()
192    }
193
194    /// Restore state from JSON (fixes #274)
195    /// Call after construction to restore learned patterns from a previous session.
196    pub fn load_state(&self, json: &str) -> Result<usize, String> {
197        let state: serde_json::Value =
198            serde_json::from_str(json).map_err(|e| format!("Invalid state JSON: {}", e))?;
199
200        let mut loaded = 0;
201
202        // Restore patterns into reasoning bank
203        if let Some(patterns) = state.get("patterns").and_then(|p| p.as_array()) {
204            let mut rb = self.reasoning_bank.write();
205            for p in patterns {
206                if let Ok(pattern) = serde_json::from_value::<crate::LearnedPattern>(p.clone()) {
207                    rb.insert_pattern(pattern);
208                    loaded += 1;
209                }
210            }
211        }
212
213        Ok(loaded)
214    }
215}
216
217/// Coordinator statistics
218#[derive(Debug, Clone)]
219#[cfg_attr(
220    feature = "serde-support",
221    derive(serde::Serialize, serde::Deserialize)
222)]
223pub struct CoordinatorStats {
224    /// Total trajectories ever recorded (buffered + dropped) — fixes #273
225    #[cfg_attr(feature = "serde-support", serde(alias = "trajectoriesRecorded"))]
226    pub trajectories_recorded: u64,
227    pub trajectories_buffered: usize,
228    pub trajectories_dropped: u64,
229    pub buffer_success_rate: f64,
230    pub patterns_stored: usize,
231    /// Alias for patterns_stored — matches JS interface
232    #[cfg_attr(feature = "serde-support", serde(alias = "patternsLearned"))]
233    pub patterns_learned: usize,
234    pub ewc_tasks: usize,
235    pub instant_enabled: bool,
236    pub background_enabled: bool,
237}
238
239#[cfg(test)]
240mod tests {
241    use super::*;
242    use crate::types::TrajectoryStep;
243
244    fn make_trajectory(id: u64) -> QueryTrajectory {
245        let mut t = QueryTrajectory::new(id, vec![0.1; 256]);
246        t.add_step(TrajectoryStep::new(vec![0.5; 256], vec![], 0.8, 0));
247        t.finalize(0.8, 1000);
248        t
249    }
250
251    #[test]
252    fn test_coordinator_creation() {
253        let coord = LoopCoordinator::new(256);
254        let stats = coord.stats();
255        assert_eq!(stats.trajectories_buffered, 0);
256    }
257
258    #[test]
259    fn test_inference_processing() {
260        let coord = LoopCoordinator::new(256);
261
262        for i in 0..10 {
263            let t = make_trajectory(coord.next_trajectory_id());
264            coord.on_inference(t);
265        }
266
267        let stats = coord.stats();
268        assert_eq!(stats.trajectories_buffered, 10);
269    }
270
271    #[test]
272    fn test_force_background() {
273        let coord = LoopCoordinator::new(256);
274
275        for i in 0..150 {
276            let t = make_trajectory(coord.next_trajectory_id());
277            coord.on_inference(t);
278        }
279
280        let result = coord.force_background();
281        assert_eq!(result.trajectories_processed, 150);
282        assert!(result.patterns_extracted > 0);
283    }
284
285    #[test]
286    fn test_get_micro_lora_weights() {
287        let coord = LoopCoordinator::new(256);
288        let (down, up) = coord.get_micro_lora_weights();
289
290        assert_eq!(down.len(), 256 * 2);
291        assert_eq!(up.len(), 2 * 256);
292    }
293
294    #[test]
295    fn test_restore_micro_lora_weights() {
296        let coord = LoopCoordinator::new(256);
297
298        let custom_down = vec![0.5f32; 256 * 2];
299        let custom_up = vec![0.3f32; 2 * 256];
300
301        let result = coord.restore_micro_lora_weights(custom_down.clone(), custom_up.clone());
302        assert!(result.is_ok());
303
304        let (got_down, got_up) = coord.get_micro_lora_weights();
305        assert_eq!(got_down, custom_down);
306        assert_eq!(got_up, custom_up);
307    }
308
309    #[test]
310    fn test_restore_micro_lora_weights_wrong_dim() {
311        let coord = LoopCoordinator::new(256);
312
313        let wrong_down = vec![0.5f32; 256 * 3];
314        let custom_up = vec![0.3f32; 2 * 256];
315
316        let result = coord.restore_micro_lora_weights(wrong_down, custom_up);
317        assert!(result.is_err());
318        assert!(result.unwrap_err().contains("down_proj dimension mismatch"));
319    }
320}
ruvector_sona/loops/coordinator.rs

ruvector_sona/loops/
coordinator.rs
