mockforge_world_state/
engine.rs

1//! World State Engine - Central coordinator for unified world state
2//!
3//! This module provides the core engine that aggregates state from all
4//! subsystems and maintains unified state snapshots.
5
6use crate::aggregators::StateAggregator;
7use crate::model::{StateLayer, WorldStateSnapshot};
8use crate::query::WorldStateQuery;
9use anyhow::Result;
10use std::collections::HashMap;
11use std::sync::Arc;
12use tokio::sync::RwLock;
13use tracing::{debug, info, warn};
14
15/// World State Engine
16///
17/// Central coordinator that aggregates state from all MockForge subsystems
18/// and provides a unified view of the entire system state.
19pub struct WorldStateEngine {
20    /// Registered aggregators for each layer
21    aggregators: HashMap<StateLayer, Arc<dyn StateAggregator>>,
22    /// Historical snapshots (limited to recent ones)
23    snapshots: Arc<RwLock<Vec<WorldStateSnapshot>>>,
24    /// Maximum number of snapshots to keep
25    max_snapshots: usize,
26}
27
28impl WorldStateEngine {
29    /// Create a new world state engine
30    pub fn new() -> Self {
31        Self {
32            aggregators: HashMap::new(),
33            snapshots: Arc::new(RwLock::new(Vec::new())),
34            max_snapshots: 100,
35        }
36    }
37
38    /// Register an aggregator for a layer
39    pub fn register_aggregator(&mut self, aggregator: Arc<dyn StateAggregator>) {
40        let layer = aggregator.layer();
41        self.aggregators.insert(layer, aggregator);
42        info!("Registered aggregator for layer: {:?}", layer);
43    }
44
45    /// Create a snapshot of the current world state
46    pub async fn create_snapshot(&self) -> Result<WorldStateSnapshot> {
47        debug!("Creating world state snapshot");
48
49        let mut snapshot = WorldStateSnapshot::new();
50        let mut all_nodes = Vec::new();
51        let mut all_edges = Vec::new();
52
53        // Aggregate state from all registered aggregators
54        for (layer, aggregator) in &self.aggregators {
55            match aggregator.aggregate().await {
56                Ok((nodes, edges)) => {
57                    debug!(
58                        "Aggregated {} nodes and {} edges from layer: {:?}",
59                        nodes.len(),
60                        edges.len(),
61                        layer
62                    );
63                    all_nodes.extend(nodes);
64                    all_edges.extend(edges);
65                    snapshot.layers.insert(*layer, true);
66                }
67                Err(e) => {
68                    warn!("Failed to aggregate state from layer {:?}: {}", layer, e);
69                    snapshot.layers.insert(*layer, false);
70                }
71            }
72        }
73
74        snapshot.nodes = all_nodes;
75        snapshot.edges = all_edges;
76
77        // Store snapshot
78        let mut snapshots = self.snapshots.write().await;
79        snapshots.push(snapshot.clone());
80
81        // Limit snapshot history
82        if snapshots.len() > self.max_snapshots {
83            snapshots.remove(0);
84        }
85
86        info!(
87            "Created world state snapshot with {} nodes and {} edges",
88            snapshot.nodes.len(),
89            snapshot.edges.len()
90        );
91
92        Ok(snapshot)
93    }
94
95    /// Get the current world state snapshot
96    pub async fn get_current_snapshot(&self) -> Result<WorldStateSnapshot> {
97        self.create_snapshot().await
98    }
99
100    /// Get a snapshot by ID
101    pub async fn get_snapshot(&self, snapshot_id: &str) -> Option<WorldStateSnapshot> {
102        let snapshots = self.snapshots.read().await;
103        snapshots.iter().find(|s| s.id == snapshot_id).cloned()
104    }
105
106    /// Get all available snapshots
107    pub async fn get_all_snapshots(&self) -> Vec<WorldStateSnapshot> {
108        let snapshots = self.snapshots.read().await;
109        snapshots.clone()
110    }
111
112    /// Query the current world state with filters
113    pub async fn query(&self, query: &WorldStateQuery) -> Result<WorldStateSnapshot> {
114        let snapshot = self.create_snapshot().await?;
115
116        // Filter nodes
117        let filtered_nodes: Vec<_> =
118            snapshot.nodes.iter().filter(|node| query.matches_node(node)).cloned().collect();
119
120        // Filter edges
121        let filtered_edges: Vec<_> = if query.include_edges {
122            snapshot.edges.iter().filter(|edge| query.matches_edge(edge)).cloned().collect()
123        } else {
124            Vec::new()
125        };
126
127        // Create filtered snapshot
128        let mut filtered_snapshot = snapshot;
129        filtered_snapshot.nodes = filtered_nodes;
130        filtered_snapshot.edges = filtered_edges;
131
132        Ok(filtered_snapshot)
133    }
134
135    /// Get available layers
136    pub fn get_layers(&self) -> Vec<StateLayer> {
137        self.aggregators.keys().copied().collect()
138    }
139
140    /// Set maximum number of snapshots to keep
141    pub fn set_max_snapshots(&mut self, max: usize) {
142        self.max_snapshots = max;
143    }
144}
145
146impl Default for WorldStateEngine {
147    fn default() -> Self {
148        Self::new()
149    }
150}
151
152#[cfg(test)]
153mod tests {
154    use super::*;
155    use crate::model::{NodeType, StateNode};
156    use std::collections::HashSet;
157
158    // Mock aggregator for testing
159    struct MockAggregator {
160        layer: StateLayer,
161        nodes: Vec<StateNode>,
162        edges: Vec<crate::model::StateEdge>,
163        should_fail: bool,
164    }
165
166    impl MockAggregator {
167        fn new(layer: StateLayer) -> Self {
168            Self {
169                layer,
170                nodes: Vec::new(),
171                edges: Vec::new(),
172                should_fail: false,
173            }
174        }
175
176        fn with_nodes(mut self, nodes: Vec<StateNode>) -> Self {
177            self.nodes = nodes;
178            self
179        }
180
181        fn with_edges(mut self, edges: Vec<crate::model::StateEdge>) -> Self {
182            self.edges = edges;
183            self
184        }
185
186        fn with_failure(mut self) -> Self {
187            self.should_fail = true;
188            self
189        }
190    }
191
192    #[async_trait::async_trait]
193    impl StateAggregator for MockAggregator {
194        async fn aggregate(&self) -> Result<(Vec<StateNode>, Vec<crate::model::StateEdge>)> {
195            if self.should_fail {
196                anyhow::bail!("Mock aggregator failure");
197            }
198            Ok((self.nodes.clone(), self.edges.clone()))
199        }
200
201        fn layer(&self) -> StateLayer {
202            self.layer
203        }
204    }
205
206    #[test]
207    fn test_world_state_engine_new() {
208        let engine = WorldStateEngine::new();
209        assert_eq!(engine.aggregators.len(), 0);
210        assert_eq!(engine.max_snapshots, 100);
211    }
212
213    #[test]
214    fn test_world_state_engine_default() {
215        let engine = WorldStateEngine::default();
216        assert_eq!(engine.aggregators.len(), 0);
217        assert_eq!(engine.max_snapshots, 100);
218    }
219
220    #[test]
221    fn test_register_aggregator() {
222        let mut engine = WorldStateEngine::new();
223        let aggregator = Arc::new(MockAggregator::new(StateLayer::Personas));
224
225        engine.register_aggregator(aggregator);
226        assert_eq!(engine.aggregators.len(), 1);
227        assert!(engine.aggregators.contains_key(&StateLayer::Personas));
228    }
229
230    #[test]
231    fn test_register_multiple_aggregators() {
232        let mut engine = WorldStateEngine::new();
233
234        engine.register_aggregator(Arc::new(MockAggregator::new(StateLayer::Personas)));
235        engine.register_aggregator(Arc::new(MockAggregator::new(StateLayer::Lifecycle)));
236        engine.register_aggregator(Arc::new(MockAggregator::new(StateLayer::Reality)));
237
238        assert_eq!(engine.aggregators.len(), 3);
239        assert!(engine.aggregators.contains_key(&StateLayer::Personas));
240        assert!(engine.aggregators.contains_key(&StateLayer::Lifecycle));
241        assert!(engine.aggregators.contains_key(&StateLayer::Reality));
242    }
243
244    #[test]
245    fn test_register_aggregator_replacement() {
246        let mut engine = WorldStateEngine::new();
247
248        // Register first aggregator for Personas layer
249        engine.register_aggregator(Arc::new(MockAggregator::new(StateLayer::Personas)));
250        assert_eq!(engine.aggregators.len(), 1);
251
252        // Register second aggregator for same layer - should replace
253        engine.register_aggregator(Arc::new(MockAggregator::new(StateLayer::Personas)));
254        assert_eq!(engine.aggregators.len(), 1);
255    }
256
257    #[tokio::test]
258    async fn test_create_snapshot_empty() {
259        let engine = WorldStateEngine::new();
260        let snapshot = engine.create_snapshot().await.unwrap();
261
262        assert!(snapshot.nodes.is_empty());
263        assert!(snapshot.edges.is_empty());
264        assert!(snapshot.layers.is_empty());
265    }
266
267    #[tokio::test]
268    async fn test_create_snapshot_with_aggregator() {
269        let mut engine = WorldStateEngine::new();
270
271        let node = StateNode::new(
272            "test-node".to_string(),
273            "Test Node".to_string(),
274            NodeType::Persona,
275            StateLayer::Personas,
276        );
277
278        let aggregator =
279            Arc::new(MockAggregator::new(StateLayer::Personas).with_nodes(vec![node.clone()]));
280
281        engine.register_aggregator(aggregator);
282
283        let snapshot = engine.create_snapshot().await.unwrap();
284
285        assert_eq!(snapshot.nodes.len(), 1);
286        assert_eq!(snapshot.nodes[0].id, "test-node");
287        assert_eq!(snapshot.edges.len(), 0);
288        assert!(snapshot.layers.contains_key(&StateLayer::Personas));
289        assert_eq!(snapshot.layers[&StateLayer::Personas], true);
290    }
291
292    #[tokio::test]
293    async fn test_create_snapshot_with_edges() {
294        let mut engine = WorldStateEngine::new();
295
296        let node1 = StateNode::new(
297            "node1".to_string(),
298            "Node 1".to_string(),
299            NodeType::Persona,
300            StateLayer::Personas,
301        );
302        let node2 = StateNode::new(
303            "node2".to_string(),
304            "Node 2".to_string(),
305            NodeType::Persona,
306            StateLayer::Personas,
307        );
308        let edge = crate::model::StateEdge::new(
309            "node1".to_string(),
310            "node2".to_string(),
311            "relates_to".to_string(),
312        );
313
314        let aggregator = Arc::new(
315            MockAggregator::new(StateLayer::Personas)
316                .with_nodes(vec![node1, node2])
317                .with_edges(vec![edge]),
318        );
319
320        engine.register_aggregator(aggregator);
321
322        let snapshot = engine.create_snapshot().await.unwrap();
323
324        assert_eq!(snapshot.nodes.len(), 2);
325        assert_eq!(snapshot.edges.len(), 1);
326        assert_eq!(snapshot.edges[0].from, "node1");
327        assert_eq!(snapshot.edges[0].to, "node2");
328    }
329
330    #[tokio::test]
331    async fn test_create_snapshot_multiple_aggregators() {
332        let mut engine = WorldStateEngine::new();
333
334        let persona_node = StateNode::new(
335            "persona1".to_string(),
336            "Persona 1".to_string(),
337            NodeType::Persona,
338            StateLayer::Personas,
339        );
340        let lifecycle_node = StateNode::new(
341            "lifecycle1".to_string(),
342            "Lifecycle 1".to_string(),
343            NodeType::Entity,
344            StateLayer::Lifecycle,
345        );
346
347        engine.register_aggregator(Arc::new(
348            MockAggregator::new(StateLayer::Personas).with_nodes(vec![persona_node]),
349        ));
350        engine.register_aggregator(Arc::new(
351            MockAggregator::new(StateLayer::Lifecycle).with_nodes(vec![lifecycle_node]),
352        ));
353
354        let snapshot = engine.create_snapshot().await.unwrap();
355
356        assert_eq!(snapshot.nodes.len(), 2);
357        assert_eq!(snapshot.layers.len(), 2);
358        assert!(snapshot.layers.contains_key(&StateLayer::Personas));
359        assert!(snapshot.layers.contains_key(&StateLayer::Lifecycle));
360    }
361
362    #[tokio::test]
363    async fn test_create_snapshot_aggregator_failure() {
364        let mut engine = WorldStateEngine::new();
365
366        let success_node = StateNode::new(
367            "success".to_string(),
368            "Success Node".to_string(),
369            NodeType::Persona,
370            StateLayer::Personas,
371        );
372
373        // One aggregator succeeds
374        engine.register_aggregator(Arc::new(
375            MockAggregator::new(StateLayer::Personas).with_nodes(vec![success_node]),
376        ));
377
378        // One aggregator fails
379        engine.register_aggregator(Arc::new(
380            MockAggregator::new(StateLayer::Lifecycle).with_failure(),
381        ));
382
383        let snapshot = engine.create_snapshot().await.unwrap();
384
385        // Should have successful nodes but mark failed layer
386        assert_eq!(snapshot.nodes.len(), 1);
387        assert_eq!(snapshot.layers[&StateLayer::Personas], true);
388        assert_eq!(snapshot.layers[&StateLayer::Lifecycle], false);
389    }
390
391    #[tokio::test]
392    async fn test_get_current_snapshot() {
393        let mut engine = WorldStateEngine::new();
394
395        let node = StateNode::new(
396            "test".to_string(),
397            "Test".to_string(),
398            NodeType::Entity,
399            StateLayer::System,
400        );
401
402        engine.register_aggregator(Arc::new(
403            MockAggregator::new(StateLayer::System).with_nodes(vec![node]),
404        ));
405
406        let snapshot = engine.get_current_snapshot().await.unwrap();
407
408        assert_eq!(snapshot.nodes.len(), 1);
409        assert_eq!(snapshot.nodes[0].id, "test");
410    }
411
412    #[tokio::test]
413    async fn test_snapshot_storage() {
414        let mut engine = WorldStateEngine::new();
415
416        let node = StateNode::new(
417            "test".to_string(),
418            "Test".to_string(),
419            NodeType::Entity,
420            StateLayer::System,
421        );
422
423        engine.register_aggregator(Arc::new(
424            MockAggregator::new(StateLayer::System).with_nodes(vec![node]),
425        ));
426
427        // Create first snapshot
428        let snapshot1 = engine.create_snapshot().await.unwrap();
429        let snapshot1_id = snapshot1.id.clone();
430
431        // Verify snapshot is stored
432        let snapshots = engine.snapshots.read().await;
433        assert_eq!(snapshots.len(), 1);
434        assert_eq!(snapshots[0].id, snapshot1_id);
435    }
436
437    #[tokio::test]
438    async fn test_get_snapshot() {
439        let mut engine = WorldStateEngine::new();
440
441        let node = StateNode::new(
442            "test".to_string(),
443            "Test".to_string(),
444            NodeType::Entity,
445            StateLayer::System,
446        );
447
448        engine.register_aggregator(Arc::new(
449            MockAggregator::new(StateLayer::System).with_nodes(vec![node]),
450        ));
451
452        // Create snapshot
453        let snapshot = engine.create_snapshot().await.unwrap();
454        let snapshot_id = snapshot.id.clone();
455
456        // Retrieve by ID
457        let retrieved = engine.get_snapshot(&snapshot_id).await;
458        assert!(retrieved.is_some());
459        assert_eq!(retrieved.unwrap().id, snapshot_id);
460
461        // Try to retrieve non-existent snapshot
462        let not_found = engine.get_snapshot("nonexistent").await;
463        assert!(not_found.is_none());
464    }
465
466    #[tokio::test]
467    async fn test_get_all_snapshots() {
468        let mut engine = WorldStateEngine::new();
469
470        let node = StateNode::new(
471            "test".to_string(),
472            "Test".to_string(),
473            NodeType::Entity,
474            StateLayer::System,
475        );
476
477        engine.register_aggregator(Arc::new(
478            MockAggregator::new(StateLayer::System).with_nodes(vec![node]),
479        ));
480
481        // Create multiple snapshots
482        engine.create_snapshot().await.unwrap();
483        engine.create_snapshot().await.unwrap();
484        engine.create_snapshot().await.unwrap();
485
486        let all_snapshots = engine.get_all_snapshots().await;
487        assert_eq!(all_snapshots.len(), 3);
488    }
489
490    #[tokio::test]
491    async fn test_max_snapshots_limit() {
492        let mut engine = WorldStateEngine::new();
493        engine.set_max_snapshots(2);
494
495        let node = StateNode::new(
496            "test".to_string(),
497            "Test".to_string(),
498            NodeType::Entity,
499            StateLayer::System,
500        );
501
502        engine.register_aggregator(Arc::new(
503            MockAggregator::new(StateLayer::System).with_nodes(vec![node]),
504        ));
505
506        // Create 3 snapshots (limit is 2)
507        engine.create_snapshot().await.unwrap();
508        engine.create_snapshot().await.unwrap();
509        engine.create_snapshot().await.unwrap();
510
511        let snapshots = engine.get_all_snapshots().await;
512        assert_eq!(snapshots.len(), 2);
513    }
514
515    #[tokio::test]
516    async fn test_query_no_filters() {
517        let mut engine = WorldStateEngine::new();
518
519        let node1 = StateNode::new(
520            "node1".to_string(),
521            "Node 1".to_string(),
522            NodeType::Persona,
523            StateLayer::Personas,
524        );
525        let node2 = StateNode::new(
526            "node2".to_string(),
527            "Node 2".to_string(),
528            NodeType::Entity,
529            StateLayer::Lifecycle,
530        );
531
532        engine.register_aggregator(Arc::new(
533            MockAggregator::new(StateLayer::Personas).with_nodes(vec![node1]),
534        ));
535        engine.register_aggregator(Arc::new(
536            MockAggregator::new(StateLayer::Lifecycle).with_nodes(vec![node2]),
537        ));
538
539        let query = WorldStateQuery::new();
540        let result = engine.query(&query).await.unwrap();
541
542        assert_eq!(result.nodes.len(), 2);
543    }
544
545    #[tokio::test]
546    async fn test_query_filter_by_node_type() {
547        let mut engine = WorldStateEngine::new();
548
549        let persona_node = StateNode::new(
550            "persona".to_string(),
551            "Persona".to_string(),
552            NodeType::Persona,
553            StateLayer::Personas,
554        );
555        let entity_node = StateNode::new(
556            "entity".to_string(),
557            "Entity".to_string(),
558            NodeType::Entity,
559            StateLayer::Lifecycle,
560        );
561
562        engine.register_aggregator(Arc::new(
563            MockAggregator::new(StateLayer::Personas).with_nodes(vec![persona_node]),
564        ));
565        engine.register_aggregator(Arc::new(
566            MockAggregator::new(StateLayer::Lifecycle).with_nodes(vec![entity_node]),
567        ));
568
569        let mut node_types = HashSet::new();
570        node_types.insert(NodeType::Persona);
571
572        let query = WorldStateQuery::new().with_node_types(node_types);
573        let result = engine.query(&query).await.unwrap();
574
575        assert_eq!(result.nodes.len(), 1);
576        assert_eq!(result.nodes[0].node_type, NodeType::Persona);
577    }
578
579    #[tokio::test]
580    async fn test_query_filter_by_layer() {
581        let mut engine = WorldStateEngine::new();
582
583        let node1 = StateNode::new(
584            "node1".to_string(),
585            "Node 1".to_string(),
586            NodeType::Persona,
587            StateLayer::Personas,
588        );
589        let node2 = StateNode::new(
590            "node2".to_string(),
591            "Node 2".to_string(),
592            NodeType::Entity,
593            StateLayer::Lifecycle,
594        );
595
596        engine.register_aggregator(Arc::new(
597            MockAggregator::new(StateLayer::Personas).with_nodes(vec![node1]),
598        ));
599        engine.register_aggregator(Arc::new(
600            MockAggregator::new(StateLayer::Lifecycle).with_nodes(vec![node2]),
601        ));
602
603        let mut layers = HashSet::new();
604        layers.insert(StateLayer::Personas);
605
606        let query = WorldStateQuery::new().with_layers(layers);
607        let result = engine.query(&query).await.unwrap();
608
609        assert_eq!(result.nodes.len(), 1);
610        assert_eq!(result.nodes[0].layer, StateLayer::Personas);
611    }
612
613    #[tokio::test]
614    async fn test_query_include_edges_false() {
615        let mut engine = WorldStateEngine::new();
616
617        let node = StateNode::new(
618            "node".to_string(),
619            "Node".to_string(),
620            NodeType::Persona,
621            StateLayer::Personas,
622        );
623        let edge = crate::model::StateEdge::new(
624            "node".to_string(),
625            "node".to_string(),
626            "self".to_string(),
627        );
628
629        engine.register_aggregator(Arc::new(
630            MockAggregator::new(StateLayer::Personas)
631                .with_nodes(vec![node])
632                .with_edges(vec![edge]),
633        ));
634
635        let query = WorldStateQuery::new().include_edges(false);
636        let result = engine.query(&query).await.unwrap();
637
638        assert_eq!(result.nodes.len(), 1);
639        assert_eq!(result.edges.len(), 0);
640    }
641
642    #[tokio::test]
643    async fn test_query_filter_edges() {
644        let mut engine = WorldStateEngine::new();
645
646        let edge1 =
647            crate::model::StateEdge::new("a".to_string(), "b".to_string(), "owns".to_string());
648        let edge2 = crate::model::StateEdge::new(
649            "b".to_string(),
650            "c".to_string(),
651            "references".to_string(),
652        );
653
654        engine.register_aggregator(Arc::new(
655            MockAggregator::new(StateLayer::Personas).with_edges(vec![edge1, edge2]),
656        ));
657
658        let mut relationship_types = HashSet::new();
659        relationship_types.insert("owns".to_string());
660
661        let query = WorldStateQuery::new().with_relationship_types(relationship_types);
662        let result = engine.query(&query).await.unwrap();
663
664        assert_eq!(result.edges.len(), 1);
665        assert_eq!(result.edges[0].relationship_type, "owns");
666    }
667
668    #[test]
669    fn test_get_layers_empty() {
670        let engine = WorldStateEngine::new();
671        let layers = engine.get_layers();
672        assert!(layers.is_empty());
673    }
674
675    #[test]
676    fn test_get_layers() {
677        let mut engine = WorldStateEngine::new();
678
679        engine.register_aggregator(Arc::new(MockAggregator::new(StateLayer::Personas)));
680        engine.register_aggregator(Arc::new(MockAggregator::new(StateLayer::Lifecycle)));
681        engine.register_aggregator(Arc::new(MockAggregator::new(StateLayer::Reality)));
682
683        let layers = engine.get_layers();
684        assert_eq!(layers.len(), 3);
685        assert!(layers.contains(&StateLayer::Personas));
686        assert!(layers.contains(&StateLayer::Lifecycle));
687        assert!(layers.contains(&StateLayer::Reality));
688    }
689
690    #[test]
691    fn test_set_max_snapshots() {
692        let mut engine = WorldStateEngine::new();
693        assert_eq!(engine.max_snapshots, 100);
694
695        engine.set_max_snapshots(50);
696        assert_eq!(engine.max_snapshots, 50);
697
698        engine.set_max_snapshots(200);
699        assert_eq!(engine.max_snapshots, 200);
700    }
701
702    #[tokio::test]
703    async fn test_snapshot_pruning() {
704        let mut engine = WorldStateEngine::new();
705        engine.set_max_snapshots(3);
706
707        let node = StateNode::new(
708            "test".to_string(),
709            "Test".to_string(),
710            NodeType::Entity,
711            StateLayer::System,
712        );
713
714        engine.register_aggregator(Arc::new(
715            MockAggregator::new(StateLayer::System).with_nodes(vec![node]),
716        ));
717
718        // Create 5 snapshots
719        let snapshot1 = engine.create_snapshot().await.unwrap();
720        let snapshot2 = engine.create_snapshot().await.unwrap();
721        let snapshot3 = engine.create_snapshot().await.unwrap();
722        let snapshot4 = engine.create_snapshot().await.unwrap();
723        let snapshot5 = engine.create_snapshot().await.unwrap();
724
725        let all_snapshots = engine.get_all_snapshots().await;
726        assert_eq!(all_snapshots.len(), 3);
727
728        // First two snapshots should be pruned
729        assert!(engine.get_snapshot(&snapshot1.id).await.is_none());
730        assert!(engine.get_snapshot(&snapshot2.id).await.is_none());
731
732        // Last three should exist
733        assert!(engine.get_snapshot(&snapshot3.id).await.is_some());
734        assert!(engine.get_snapshot(&snapshot4.id).await.is_some());
735        assert!(engine.get_snapshot(&snapshot5.id).await.is_some());
736    }
737
738    #[tokio::test]
739    async fn test_concurrent_snapshot_access() {
740        let mut engine = WorldStateEngine::new();
741
742        let node = StateNode::new(
743            "test".to_string(),
744            "Test".to_string(),
745            NodeType::Entity,
746            StateLayer::System,
747        );
748
749        engine.register_aggregator(Arc::new(
750            MockAggregator::new(StateLayer::System).with_nodes(vec![node]),
751        ));
752
753        let engine = Arc::new(engine);
754
755        // Create snapshot in background
756        let engine_clone = Arc::clone(&engine);
757        let handle1 = tokio::spawn(async move { engine_clone.create_snapshot().await });
758
759        // Read snapshots concurrently
760        let engine_clone = Arc::clone(&engine);
761        let handle2 = tokio::spawn(async move { engine_clone.get_all_snapshots().await });
762
763        // Both should complete successfully
764        let snapshot = handle1.await.unwrap();
765        let snapshots = handle2.await.unwrap();
766
767        assert!(snapshot.is_ok());
768        assert!(!snapshots.is_empty());
769    }
770
771    #[tokio::test]
772    async fn test_query_with_multiple_filters() {
773        let mut engine = WorldStateEngine::new();
774
775        let node1 = StateNode::new(
776            "node1".to_string(),
777            "Node 1".to_string(),
778            NodeType::Persona,
779            StateLayer::Personas,
780        );
781        let node2 = StateNode::new(
782            "node2".to_string(),
783            "Node 2".to_string(),
784            NodeType::Entity,
785            StateLayer::Lifecycle,
786        );
787        let node3 = StateNode::new(
788            "node3".to_string(),
789            "Node 3".to_string(),
790            NodeType::Persona,
791            StateLayer::Lifecycle,
792        );
793
794        engine.register_aggregator(Arc::new(
795            MockAggregator::new(StateLayer::Personas).with_nodes(vec![node1]),
796        ));
797        engine.register_aggregator(Arc::new(
798            MockAggregator::new(StateLayer::Lifecycle).with_nodes(vec![node2, node3]),
799        ));
800
801        let mut node_types = HashSet::new();
802        node_types.insert(NodeType::Persona);
803
804        let mut layers = HashSet::new();
805        layers.insert(StateLayer::Personas);
806
807        let query = WorldStateQuery::new().with_node_types(node_types).with_layers(layers);
808
809        let result = engine.query(&query).await.unwrap();
810
811        // Should only match node1 (Persona type AND Personas layer)
812        assert_eq!(result.nodes.len(), 1);
813        assert_eq!(result.nodes[0].id, "node1");
814    }
815
816    #[tokio::test]
817    async fn test_empty_query_result() {
818        let mut engine = WorldStateEngine::new();
819
820        let node = StateNode::new(
821            "node".to_string(),
822            "Node".to_string(),
823            NodeType::Persona,
824            StateLayer::Personas,
825        );
826
827        engine.register_aggregator(Arc::new(
828            MockAggregator::new(StateLayer::Personas).with_nodes(vec![node]),
829        ));
830
831        // Query for non-existent node type
832        let mut node_types = HashSet::new();
833        node_types.insert(NodeType::System);
834
835        let query = WorldStateQuery::new().with_node_types(node_types);
836        let result = engine.query(&query).await.unwrap();
837
838        assert_eq!(result.nodes.len(), 0);
839        assert_eq!(result.edges.len(), 0);
840    }
841
842    #[tokio::test]
843    async fn test_aggregator_metadata() {
844        struct MetadataAggregator {
845            layer: StateLayer,
846        }
847
848        #[async_trait::async_trait]
849        impl StateAggregator for MetadataAggregator {
850            async fn aggregate(&self) -> Result<(Vec<StateNode>, Vec<crate::model::StateEdge>)> {
851                Ok((Vec::new(), Vec::new()))
852            }
853
854            fn layer(&self) -> StateLayer {
855                self.layer
856            }
857
858            fn metadata(&self) -> HashMap<String, serde_json::Value> {
859                let mut map = HashMap::new();
860                map.insert("version".to_string(), serde_json::json!("1.0"));
861                map.insert("enabled".to_string(), serde_json::json!(true));
862                map
863            }
864        }
865
866        let aggregator = MetadataAggregator {
867            layer: StateLayer::Personas,
868        };
869
870        let metadata = aggregator.metadata();
871        assert_eq!(metadata.len(), 2);
872        assert_eq!(metadata.get("version"), Some(&serde_json::json!("1.0")));
873        assert_eq!(metadata.get("enabled"), Some(&serde_json::json!(true)));
874    }
875}
mockforge_world_state/engine.rs

mockforge_world_state/
engine.rs
