mockforge-core 0.3.116

Shared logic for MockForge - routing, validation, latency, proxy
Documentation
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
//! Consistency engine implementation
//!
//! The consistency engine coordinates state across all protocols, ensuring
//! that persona, scenario, reality level, and entity state are synchronized.

use crate::consistency::adapters::ProtocolAdapter;
use crate::consistency::types::{EntityState, ProtocolState, StateChangeEvent, UnifiedState};
use crate::protocol_abstraction::Protocol;
use crate::reality::RealityLevel;
use crate::Result;
// ChaosScenario is defined in mockforge-chaos, but we use serde_json::Value to avoid circular dependency
type ChaosScenario = serde_json::Value;
use mockforge_data::PersonaProfile;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::{broadcast, RwLock};
use tracing::{debug, error, info, warn};

/// Default broadcast channel capacity for state change events
const DEFAULT_BROADCAST_CAPACITY: usize = 1000;

/// Get the broadcast channel capacity from environment or use default
fn get_broadcast_capacity() -> usize {
    std::env::var("MOCKFORGE_BROADCAST_CAPACITY")
        .ok()
        .and_then(|s| s.parse().ok())
        .unwrap_or(DEFAULT_BROADCAST_CAPACITY)
}

/// Consistency engine for coordinating state across all protocols
///
/// The engine maintains unified state per workspace and broadcasts state
/// changes to all registered protocol adapters. This ensures that all
/// protocols reflect the same persona, scenario, reality level, and entity state.
pub struct ConsistencyEngine {
    /// Workspace state storage (workspace_id -> UnifiedState)
    states: Arc<RwLock<HashMap<String, UnifiedState>>>,
    /// Event broadcaster for state change notifications
    event_tx: broadcast::Sender<StateChangeEvent>,
    /// Registered protocol adapters
    adapters: Arc<RwLock<Vec<Arc<dyn ProtocolAdapter + Send + Sync>>>>,
}

impl ConsistencyEngine {
    /// Create a new consistency engine
    ///
    /// The broadcast channel capacity can be configured via the
    /// `MOCKFORGE_BROADCAST_CAPACITY` environment variable.
    pub fn new() -> Self {
        let capacity = get_broadcast_capacity();
        let (event_tx, _) = broadcast::channel(capacity);
        Self {
            states: Arc::new(RwLock::new(HashMap::new())),
            event_tx,
            adapters: Arc::new(RwLock::new(Vec::new())),
        }
    }

    /// Register a protocol adapter
    ///
    /// Adapters are notified of all state changes for their protocol.
    pub async fn register_adapter(&self, adapter: Arc<dyn ProtocolAdapter + Send + Sync>) {
        let mut adapters = self.adapters.write().await;
        adapters.push(adapter);
        info!("Registered protocol adapter");
    }

    /// Get or create unified state for a workspace
    pub async fn get_or_create_state(&self, workspace_id: &str) -> UnifiedState {
        let mut states = self.states.write().await;
        states
            .entry(workspace_id.to_string())
            .or_insert_with(|| UnifiedState::new(workspace_id.to_string()))
            .clone()
    }

    /// Get unified state for a workspace (returns None if not found)
    pub async fn get_state(&self, workspace_id: &str) -> Option<UnifiedState> {
        let states = self.states.read().await;
        states.get(workspace_id).cloned()
    }

    /// Set active persona for a workspace
    ///
    /// This updates the unified state and broadcasts the change to all
    /// protocol adapters, ensuring all protocols use the new persona.
    pub async fn set_active_persona(
        &self,
        workspace_id: &str,
        persona: PersonaProfile,
    ) -> Result<()> {
        let mut states = self.states.write().await;
        let state = states
            .entry(workspace_id.to_string())
            .or_insert_with(|| UnifiedState::new(workspace_id.to_string()));

        state.active_persona = Some(persona.clone());
        state.increment_version();

        let event = StateChangeEvent::PersonaChanged {
            workspace_id: workspace_id.to_string(),
            persona,
        };

        drop(states); // Release lock before broadcasting

        self.broadcast_event(&event).await;
        info!("Set active persona for workspace {}", workspace_id);
        Ok(())
    }

    /// Set active scenario for a workspace
    pub async fn set_active_scenario(&self, workspace_id: &str, scenario_id: String) -> Result<()> {
        let mut states = self.states.write().await;
        let state = states
            .entry(workspace_id.to_string())
            .or_insert_with(|| UnifiedState::new(workspace_id.to_string()));

        state.active_scenario = Some(scenario_id.clone());
        state.increment_version();

        let event = StateChangeEvent::ScenarioChanged {
            workspace_id: workspace_id.to_string(),
            scenario_id,
        };

        drop(states);

        self.broadcast_event(&event).await;
        info!("Set active scenario for workspace {}", workspace_id);
        Ok(())
    }

    /// Set reality level for a workspace
    pub async fn set_reality_level(&self, workspace_id: &str, level: RealityLevel) -> Result<()> {
        let mut states = self.states.write().await;
        let state = states
            .entry(workspace_id.to_string())
            .or_insert_with(|| UnifiedState::new(workspace_id.to_string()));

        state.reality_level = level;
        state.increment_version();

        let event = StateChangeEvent::RealityLevelChanged {
            workspace_id: workspace_id.to_string(),
            level,
        };

        drop(states);

        self.broadcast_event(&event).await;
        debug!("Set reality level {:?} for workspace {}", level, workspace_id);
        Ok(())
    }

    /// Set reality continuum ratio for a workspace
    pub async fn set_reality_ratio(&self, workspace_id: &str, ratio: f64) -> Result<()> {
        let ratio = ratio.clamp(0.0, 1.0);
        let mut states = self.states.write().await;
        let state = states
            .entry(workspace_id.to_string())
            .or_insert_with(|| UnifiedState::new(workspace_id.to_string()));

        state.reality_continuum_ratio = ratio;
        state.increment_version();

        let event = StateChangeEvent::RealityRatioChanged {
            workspace_id: workspace_id.to_string(),
            ratio,
        };

        drop(states);

        self.broadcast_event(&event).await;
        debug!("Set reality ratio {} for workspace {}", ratio, workspace_id);
        Ok(())
    }

    /// Register or update an entity
    ///
    /// Entities are tracked across all protocols. When an entity is created
    /// via HTTP, it becomes immediately available in GraphQL, gRPC, etc.
    /// Also automatically adds the entity to the persona graph if a persona_id is present.
    pub async fn register_entity(&self, workspace_id: &str, entity: EntityState) -> Result<()> {
        let mut states = self.states.write().await;
        let state = states
            .entry(workspace_id.to_string())
            .or_insert_with(|| UnifiedState::new(workspace_id.to_string()));

        let is_new = !state
            .entity_state
            .contains_key(&UnifiedState::entity_key(&entity.entity_type, &entity.entity_id));

        // Add entity to persona graph if persona_id is present
        #[cfg(feature = "persona-graph")]
        if let Some(ref persona_id) = entity.persona_id {
            let graph = state.get_or_create_persona_graph();
            graph.get_or_create_node_with_links(persona_id, &entity.entity_type, None, None);

            // If entity data contains related entity IDs, link them in the graph
            #[cfg(feature = "persona-graph")]
            if let Some(user_id) = entity.data.get("user_id").or_else(|| entity.data.get("userId"))
            {
                if let Some(user_id_str) = user_id.as_str() {
                    let user_persona_id = format!("user:{}", user_id_str);
                    graph.link_entity_types(
                        &user_persona_id,
                        "user",
                        persona_id,
                        &entity.entity_type,
                    );
                }
            }

            // Link orders to payments
            #[cfg(feature = "persona-graph")]
            if entity.entity_type == "payment" {
                if let Some(order_id) =
                    entity.data.get("order_id").or_else(|| entity.data.get("orderId"))
                {
                    if let Some(order_id_str) = order_id.as_str() {
                        let order_persona_id = format!("order:{}", order_id_str);
                        graph.link_entity_types(&order_persona_id, "order", persona_id, "payment");
                    }
                }
            }
        }

        let entity_clone = entity.clone();
        state.register_entity(entity_clone.clone());

        let event = if is_new {
            StateChangeEvent::EntityCreated {
                workspace_id: workspace_id.to_string(),
                entity: entity_clone,
            }
        } else {
            StateChangeEvent::EntityUpdated {
                workspace_id: workspace_id.to_string(),
                entity: entity_clone,
            }
        };

        drop(states);

        self.broadcast_event(&event).await;
        debug!(
            "Registered entity {}:{} for workspace {}",
            entity.entity_type, entity.entity_id, workspace_id
        );
        Ok(())
    }

    /// Get entity by type and ID
    pub async fn get_entity(
        &self,
        workspace_id: &str,
        entity_type: &str,
        entity_id: &str,
    ) -> Option<EntityState> {
        let states = self.states.read().await;
        states.get(workspace_id)?.get_entity(entity_type, entity_id).cloned()
    }

    /// Find related entities using the persona graph
    ///
    /// Given a persona ID and entity type, finds all related entities of the target type
    /// by traversing the persona graph.
    ///
    /// # Arguments
    /// * `workspace_id` - Workspace identifier
    /// * `persona_id` - Starting persona ID
    /// * `target_entity_type` - Entity type to find (e.g., "order", "payment")
    /// * `relationship_type` - Optional relationship type filter (e.g., "has_orders")
    ///
    /// # Returns
    /// Vector of entity states matching the criteria
    pub async fn find_related_entities(
        &self,
        workspace_id: &str,
        persona_id: &str,
        target_entity_type: &str,
        relationship_type: Option<&str>,
    ) -> Vec<EntityState> {
        #[cfg(not(feature = "persona-graph"))]
        let _ = (persona_id, relationship_type);

        let states = self.states.read().await;
        let state = match states.get(workspace_id) {
            Some(s) => s,
            None => return Vec::new(),
        };

        #[cfg(feature = "persona-graph")]
        let graph = match state.persona_graph() {
            Some(g) => g,
            None => return Vec::new(),
        };

        // Find related persona IDs in the graph
        #[cfg(feature = "persona-graph")]
        let related_persona_ids =
            graph.find_related_by_entity_type(persona_id, target_entity_type, relationship_type);

        #[cfg(not(feature = "persona-graph"))]
        let related_persona_ids: Vec<String> = Vec::new();

        // Convert persona IDs to entity states
        let mut related_entities = Vec::new();
        for related_persona_id in related_persona_ids {
            // Extract entity ID from persona ID (format: "entity_type:entity_id")
            if let Some((_, entity_id)) = related_persona_id.split_once(':') {
                if let Some(entity) = state.get_entity(target_entity_type, entity_id) {
                    related_entities.push(entity.clone());
                }
            }
        }

        related_entities
    }

    /// Activate a chaos rule
    pub async fn activate_chaos_rule(&self, workspace_id: &str, rule: ChaosScenario) -> Result<()> {
        let mut states = self.states.write().await;
        let state = states
            .entry(workspace_id.to_string())
            .or_insert_with(|| UnifiedState::new(workspace_id.to_string()));

        // Remove existing rule with same name if present
        // Note: ChaosScenario is serde_json::Value, so we compare by serializing
        if let Some(rule_name) = rule.get("name").and_then(|v| v.as_str()) {
            state
                .active_chaos_rules
                .retain(|r| r.get("name").and_then(|v| v.as_str()) != Some(rule_name));
        }
        state.active_chaos_rules.push(rule.clone());
        state.increment_version();

        let event = StateChangeEvent::ChaosRuleActivated {
            workspace_id: workspace_id.to_string(),
            rule,
        };

        drop(states);

        self.broadcast_event(&event).await;
        info!("Activated chaos rule for workspace {}", workspace_id);
        Ok(())
    }

    /// Deactivate a chaos rule
    pub async fn deactivate_chaos_rule(&self, workspace_id: &str, rule_name: &str) -> Result<()> {
        let mut states = self.states.write().await;
        let state = states
            .entry(workspace_id.to_string())
            .or_insert_with(|| UnifiedState::new(workspace_id.to_string()));

        // Note: ChaosScenario is serde_json::Value, so we compare by serializing
        let removed = state
            .active_chaos_rules
            .iter()
            .any(|r| r.get("name").and_then(|v| v.as_str()) == Some(rule_name));

        if removed {
            state
                .active_chaos_rules
                .retain(|r| r.get("name").and_then(|v| v.as_str()) != Some(rule_name));
            state.increment_version();

            let event = StateChangeEvent::ChaosRuleDeactivated {
                workspace_id: workspace_id.to_string(),
                rule_name: rule_name.to_string(),
            };

            drop(states);

            self.broadcast_event(&event).await;
            info!("Deactivated chaos rule {} for workspace {}", rule_name, workspace_id);
        }

        Ok(())
    }

    /// Update protocol state
    pub async fn update_protocol_state(
        &self,
        workspace_id: &str,
        protocol: Protocol,
        protocol_state: ProtocolState,
    ) -> Result<()> {
        let mut states = self.states.write().await;
        let state = states
            .entry(workspace_id.to_string())
            .or_insert_with(|| UnifiedState::new(workspace_id.to_string()));

        state.protocol_states.insert(protocol, protocol_state);
        state.increment_version();

        Ok(())
    }

    /// Get protocol state
    pub async fn get_protocol_state(
        &self,
        workspace_id: &str,
        protocol: Protocol,
    ) -> Option<ProtocolState> {
        let states = self.states.read().await;
        states.get(workspace_id)?.protocol_states.get(&protocol).cloned()
    }

    /// Subscribe to state change events
    ///
    /// Returns a receiver that will receive all state change events for
    /// the specified workspace (or all workspaces if None).
    pub fn subscribe_to_events(
        &self,
        _workspace_id: Option<&str>,
    ) -> broadcast::Receiver<StateChangeEvent> {
        self.event_tx.subscribe()
    }

    /// Restore unified state from a snapshot
    ///
    /// This replaces the current state for a workspace with the provided state.
    /// All protocol adapters will be notified of the state changes.
    pub async fn restore_state(&self, state: UnifiedState) -> Result<()> {
        let workspace_id = state.workspace_id.clone();
        let mut states = self.states.write().await;
        states.insert(workspace_id.clone(), state.clone());
        drop(states);

        // Broadcast state change events for all components
        if let Some(ref persona) = state.active_persona {
            self.broadcast_event(&StateChangeEvent::PersonaChanged {
                workspace_id: workspace_id.clone(),
                persona: persona.clone(),
            })
            .await;
        }

        if let Some(ref scenario_id) = state.active_scenario {
            self.broadcast_event(&StateChangeEvent::ScenarioChanged {
                workspace_id: workspace_id.clone(),
                scenario_id: scenario_id.clone(),
            })
            .await;
        }

        self.broadcast_event(&StateChangeEvent::RealityLevelChanged {
            workspace_id: workspace_id.clone(),
            level: state.reality_level,
        })
        .await;

        self.broadcast_event(&StateChangeEvent::RealityRatioChanged {
            workspace_id: workspace_id.clone(),
            ratio: state.reality_continuum_ratio,
        })
        .await;

        // Broadcast entity events
        for entity in state.entity_state.values() {
            self.broadcast_event(&StateChangeEvent::EntityCreated {
                workspace_id: workspace_id.clone(),
                entity: entity.clone(),
            })
            .await;
        }

        // Broadcast chaos rule events
        for rule in &state.active_chaos_rules {
            self.broadcast_event(&StateChangeEvent::ChaosRuleActivated {
                workspace_id: workspace_id.clone(),
                rule: rule.clone(),
            })
            .await;
        }

        info!("Restored state for workspace {}", workspace_id);
        Ok(())
    }

    /// Broadcast a state change event to all adapters
    async fn broadcast_event(&self, event: &StateChangeEvent) {
        // Send to event subscribers
        if let Err(e) = self.event_tx.send(event.clone()) {
            warn!("Failed to broadcast state change event: {}", e);
        }

        // Notify all protocol adapters
        let adapters = self.adapters.read().await;
        for adapter in adapters.iter() {
            if let Err(e) = adapter.on_state_change(event).await {
                error!(
                    "Protocol adapter {:?} failed to handle state change: {}",
                    adapter.protocol(),
                    e
                );
            }
        }
    }
}

impl Default for ConsistencyEngine {
    fn default() -> Self {
        Self::new()
    }
}