cellstate_core/

drift.rs

//! Composite divergence metric for multi-agent drift detection.
//!
//! The `DriftMeter` computes a weighted divergence score between two agents
//! by composing three orthogonal signals:
//!
//! 1. **State divergence**: How different are the agents' derived states?
//! 2. **Causal divergence**: How much have their Event DAG lanes diverged?
//! 3. **Context divergence**: How much overlap exists in their scored context?
//!
//! Each signal produces a value in [0, 1] (0 = identical, 1 = fully diverged).
//! A Weibull temporal decay weights recent divergence higher than stale divergence.
//!
//! The composite score is a weighted sum using the same `ScoringWeights::validate()`
//! and `normalize()` pattern, ensuring mathematical consistency with the rest of the
//! scoring system.
//!
//! Re-export path: `cellstate_core::drift::*`

use crate::{AgentId, AgentState, DecayParams, Event, Timestamp};
use serde::{Deserialize, Serialize};

/// Errors from validating [`DriftWeights`].
#[derive(Debug, Clone, PartialEq)]
pub enum DriftWeightsError {
    /// One or more weights are NaN or Infinity.
    NonFinite,
    /// One or more weights are negative.
    Negative,
    /// Weights do not sum to 1.0 (±epsilon). Contains the actual sum.
    SumNotOne(f32),
}

impl std::fmt::Display for DriftWeightsError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::NonFinite => write!(f, "one or more drift weights are NaN or Infinity"),
            Self::Negative => write!(f, "one or more drift weights are negative"),
            Self::SumNotOne(sum) => {
                write!(f, "drift weights sum to {} instead of 1.0", sum)
            }
        }
    }
}

impl std::error::Error for DriftWeightsError {}

/// Divergence weights for the three DriftMeter signals.
///
/// Must sum to 1.0 (same constraint as `ScoringWeights`).
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct DriftWeights {
    /// Weight for agent state divergence signal.
    pub state: f32,
    /// Weight for causal (DAG lane) divergence signal.
    pub causal: f32,
    /// Weight for context overlap divergence signal.
    pub context: f32,
}

impl Default for DriftWeights {
    fn default() -> Self {
        Self {
            state: 0.40,
            causal: 0.35,
            context: 0.25,
        }
    }
}

impl DriftWeights {
    /// Validate weights are finite, non-negative, and sum to 1.0 (±epsilon).
    pub fn validate(&self) -> Result<(), DriftWeightsError> {
        const EPSILON: f32 = 0.01;
        if !self.state.is_finite() || !self.causal.is_finite() || !self.context.is_finite() {
            return Err(DriftWeightsError::NonFinite);
        }
        if self.state < 0.0 || self.causal < 0.0 || self.context < 0.0 {
            return Err(DriftWeightsError::Negative);
        }
        let sum = self.state + self.causal + self.context;
        if (sum - 1.0).abs() > EPSILON {
            return Err(DriftWeightsError::SumNotOne(sum));
        }
        Ok(())
    }

    /// Normalize weights to sum to exactly 1.0.
    ///
    /// If the sum is NaN, Infinity, or zero the weights are left unchanged
    /// to avoid silently producing all-NaN values.
    pub fn normalize(&mut self) {
        let sum = self.state + self.causal + self.context;
        if sum.is_finite() && sum > 0.0 {
            self.state /= sum;
            self.causal /= sum;
            self.context /= sum;
        }
        // If sum is NaN/Inf/zero, leave weights unchanged
    }
}

/// Input parameters for computing drift between two agents.
#[derive(Debug, Clone)]
pub struct DriftInput<'a> {
    /// First agent ID.
    pub agent_a: AgentId,
    /// Second agent ID.
    pub agent_b: AgentId,
    /// Event stream for agent A.
    pub events_a: &'a [Event<serde_json::Value>],
    /// Event stream for agent B.
    pub events_b: &'a [Event<serde_json::Value>],
    /// Entity IDs scored into agent A's L1 context.
    ///
    /// Uses raw `uuid::Uuid` because context items are heterogeneous (notes,
    /// artifacts, trajectories, etc.). Type checking happens at insertion time
    /// via the `EntityIdType` trait on each concrete ID type. When Phase 2
    /// populates these from the database, consider migrating to a typed
    /// `ContextEntity` enum if cross-type confusion becomes a risk.
    pub context_items_a: &'a [uuid::Uuid],
    /// Entity IDs scored into agent B's L1 context. See `context_items_a` docs.
    pub context_items_b: &'a [uuid::Uuid],
    /// Temporal decay parameters (Weibull).
    pub decay: &'a DecayParams,
    /// Divergence signal weights.
    pub weights: &'a DriftWeights,
    /// Alignment threshold (from IntentDef).
    pub threshold: f64,
}

/// Composite divergence metric between two agents.
///
/// All scores are in [0, 1]: 0.0 = identical, 1.0 = fully diverged.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
pub struct DriftMeter {
    /// First agent in the comparison.
    pub agent_a: AgentId,
    /// Second agent in the comparison.
    pub agent_b: AgentId,
    /// State-level divergence (0 = same state, 1 = fully different).
    pub state_divergence: f64,
    /// Causal divergence from DAG lane comparison.
    pub causal_divergence: f64,
    /// Context overlap divergence (1 - Jaccard similarity of scored items).
    pub context_divergence: f64,
    /// Weighted composite score.
    pub composite_score: f64,
    /// Threshold from IntentDef.drift_threshold (default 0.85).
    /// Drift is detected when `1.0 - composite_score < threshold`,
    /// i.e. when composite divergence is high enough.
    pub threshold: f64,
    /// Whether the agents are considered to be drifting apart.
    pub is_drifting: bool,
    /// When this measurement was taken.
    #[cfg_attr(feature = "openapi", schema(value_type = String, format = "date-time"))]
    pub computed_at: Timestamp,
}

impl DriftMeter {
    /// Compute the drift between two agents from their event streams.
    ///
    /// This is a pure function — no I/O.
    pub fn compute(input: DriftInput<'_>) -> Self {
        let now = chrono::Utc::now();

        // Signal 1: State divergence
        let state_a = AgentState::from_events(input.events_a);
        let state_b = AgentState::from_events(input.events_b);
        let state_divergence = Self::compute_state_divergence(&state_a, &state_b);

        // Signal 2: Causal divergence (DAG lane comparison)
        let causal_divergence =
            Self::compute_causal_divergence(input.events_a, input.events_b, input.decay);

        // Signal 3: Context divergence (Jaccard distance of scored items)
        let context_divergence =
            Self::compute_context_divergence(input.context_items_a, input.context_items_b);

        // Composite weighted score
        let composite_score = input.weights.state as f64 * state_divergence
            + input.weights.causal as f64 * causal_divergence
            + input.weights.context as f64 * context_divergence;

        // Guard NaN/Infinity: treat as no divergence rather than silently
        // poisoning downstream decisions.
        let composite_score = if composite_score.is_finite() {
            composite_score.clamp(0.0, 1.0)
        } else {
            0.0 // fail-safe: treat as no divergence
        };

        // Drift detected when alignment (1 - divergence) drops below threshold
        let alignment = 1.0 - composite_score;
        let is_drifting = alignment < input.threshold;

        Self {
            agent_a: input.agent_a,
            agent_b: input.agent_b,
            state_divergence,
            causal_divergence,
            context_divergence,
            composite_score,
            threshold: input.threshold,
            is_drifting,
            computed_at: now,
        }
    }

    /// Compare two AgentState values for divergence.
    ///
    /// Returns 0.0 if identical variant, 0.5 if same category but different
    /// details, 1.0 if completely different states.
    fn compute_state_divergence(a: &AgentState, b: &AgentState) -> f64 {
        use std::mem::discriminant;

        if a == b {
            return 0.0;
        }
        if discriminant(a) == discriminant(b) {
            // Same variant but different inner data (e.g. different scope_ids)
            return 0.5;
        }
        1.0
    }

    /// Compare DAG lane sequences between two agents.
    ///
    /// Uses the ratio of shared event kinds in recent history as a proxy
    /// for causal alignment. Apply Weibull decay to weight recent events more.
    fn compute_causal_divergence(
        events_a: &[Event<serde_json::Value>],
        events_b: &[Event<serde_json::Value>],
        decay: &DecayParams,
    ) -> f64 {
        if events_a.is_empty() && events_b.is_empty() {
            return 0.0;
        }
        if events_a.is_empty() || events_b.is_empty() {
            return 1.0;
        }

        // Extract recent event kinds with decay-weighted counts
        let weighted_kinds_a = Self::decay_weighted_kind_histogram(events_a, decay);
        let weighted_kinds_b = Self::decay_weighted_kind_histogram(events_b, decay);

        // Compute cosine-like similarity between the two histograms
        let all_kinds: std::collections::HashSet<u16> = weighted_kinds_a
            .keys()
            .chain(weighted_kinds_b.keys())
            .copied()
            .collect();

        if all_kinds.is_empty() {
            return 0.0;
        }

        let mut dot = 0.0f64;
        let mut norm_a = 0.0f64;
        let mut norm_b = 0.0f64;

        for kind in &all_kinds {
            let wa = *weighted_kinds_a.get(kind).unwrap_or(&0.0) as f64;
            let wb = *weighted_kinds_b.get(kind).unwrap_or(&0.0) as f64;
            dot += wa * wb;
            norm_a += wa * wa;
            norm_b += wb * wb;
        }

        let denominator = norm_a.sqrt() * norm_b.sqrt();
        if denominator == 0.0 {
            return 1.0;
        }

        let cosine_similarity = dot / denominator;
        (1.0 - cosine_similarity).clamp(0.0, 1.0)
    }

    /// Build a decay-weighted histogram of event kinds.
    fn decay_weighted_kind_histogram(
        events: &[Event<serde_json::Value>],
        decay: &DecayParams,
    ) -> std::collections::HashMap<u16, f32> {
        let mut histogram = std::collections::HashMap::new();
        let len = events.len();

        for (i, event) in events.iter().enumerate() {
            // Use position from end as the "age" for decay
            let age = (len - 1 - i) as f32;
            let weight = crate::parametric_decay(age, decay);
            *histogram.entry(event.header.event_kind.0).or_insert(0.0f32) += weight;
        }

        histogram
    }

    /// Compute context divergence as 1 - Jaccard similarity.
    fn compute_context_divergence(items_a: &[uuid::Uuid], items_b: &[uuid::Uuid]) -> f64 {
        if items_a.is_empty() && items_b.is_empty() {
            return 0.0;
        }

        let set_a: std::collections::HashSet<uuid::Uuid> = items_a.iter().copied().collect();
        let set_b: std::collections::HashSet<uuid::Uuid> = items_b.iter().copied().collect();

        let intersection = set_a.intersection(&set_b).count();
        let union = set_a.union(&set_b).count();

        if union == 0 {
            return 0.0;
        }

        let jaccard = intersection as f64 / union as f64;
        1.0 - jaccard
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::identity::EntityIdType;
    use crate::{DagPosition, Event, EventHeader, EventId, EventKind};
    use serde_json::json;

    fn make_scope_event(kind: EventKind, scope_id: crate::ScopeId) -> Event<serde_json::Value> {
        Event::new(
            EventHeader::new(
                EventId::now_v7(),
                EventId::now_v7(),
                chrono::Utc::now().timestamp_micros(),
                DagPosition::root(),
                0,
                kind,
            ),
            json!({ "scope_id": scope_id }),
        )
    }

    #[test]
    fn test_identical_agents_zero_divergence() {
        let agent_a = AgentId::now_v7();
        let agent_b = AgentId::now_v7();
        let scope_id = crate::ScopeId::now_v7();

        let events = vec![make_scope_event(EventKind::SCOPE_CREATED, scope_id)];
        let context: Vec<uuid::Uuid> = vec![uuid::Uuid::now_v7()];
        let decay = DecayParams::exponential(10.0);
        let weights = DriftWeights::default();

        let meter = DriftMeter::compute(DriftInput {
            agent_a,
            agent_b,
            events_a: &events,
            events_b: &events,
            context_items_a: &context,
            context_items_b: &context,
            decay: &decay,
            weights: &weights,
            threshold: 0.85,
        });

        assert_eq!(meter.state_divergence, 0.0);
        assert_eq!(meter.causal_divergence, 0.0);
        assert_eq!(meter.context_divergence, 0.0);
        assert_eq!(meter.composite_score, 0.0);
        assert!(!meter.is_drifting);
    }

    #[test]
    fn test_fully_diverged_agents() {
        let agent_a = AgentId::now_v7();
        let agent_b = AgentId::now_v7();
        let scope_a = crate::ScopeId::now_v7();
        let scope_b = crate::ScopeId::now_v7();

        let events_a = vec![make_scope_event(EventKind::SCOPE_CREATED, scope_a)];
        let events_b = vec![
            make_scope_event(EventKind::SCOPE_CREATED, scope_b),
            make_scope_event(EventKind::SCOPE_CLOSED, scope_b),
        ];
        let context_a: Vec<uuid::Uuid> = vec![uuid::Uuid::now_v7()];
        let context_b: Vec<uuid::Uuid> = vec![uuid::Uuid::now_v7()];
        let decay = DecayParams::exponential(10.0);
        let weights = DriftWeights::default();

        let meter = DriftMeter::compute(DriftInput {
            agent_a,
            agent_b,
            events_a: &events_a,
            events_b: &events_b,
            context_items_a: &context_a,
            context_items_b: &context_b,
            decay: &decay,
            weights: &weights,
            threshold: 0.85,
        });

        // State: Gathering vs Complete = 1.0 divergence
        assert_eq!(meter.state_divergence, 1.0);
        // Context: disjoint sets = 1.0 divergence
        assert_eq!(meter.context_divergence, 1.0);
        // Composite should be high
        assert!(meter.composite_score > 0.5);
        assert!(meter.is_drifting);
    }

    #[test]
    fn test_empty_events_no_divergence() {
        let meter = DriftMeter::compute(DriftInput {
            agent_a: AgentId::now_v7(),
            agent_b: AgentId::now_v7(),
            events_a: &[],
            events_b: &[],
            context_items_a: &[],
            context_items_b: &[],
            decay: &DecayParams::exponential(10.0),
            weights: &DriftWeights::default(),
            threshold: 0.85,
        });
        assert_eq!(meter.composite_score, 0.0);
        assert!(!meter.is_drifting);
    }

    #[test]
    fn test_one_empty_stream_fully_diverged() {
        let scope_id = crate::ScopeId::now_v7();
        let events = vec![make_scope_event(EventKind::SCOPE_CREATED, scope_id)];
        let meter = DriftMeter::compute(DriftInput {
            agent_a: AgentId::now_v7(),
            agent_b: AgentId::now_v7(),
            events_a: &events,
            events_b: &[],
            context_items_a: &[],
            context_items_b: &[],
            decay: &DecayParams::exponential(10.0),
            weights: &DriftWeights::default(),
            threshold: 0.85,
        });
        // State: Gathering (has scope_id) vs Idle = 1.0
        assert_eq!(meter.state_divergence, 1.0);
        // Causal: one empty = 1.0
        assert_eq!(meter.causal_divergence, 1.0);
        assert!(meter.is_drifting);
    }

    #[test]
    fn test_drift_weights_validate() {
        assert!(DriftWeights::default().validate().is_ok());
        assert!(DriftWeights {
            state: 0.5,
            causal: 0.5,
            context: 0.5,
        }
        .validate()
        .is_err());
    }

    #[test]
    fn test_drift_weights_normalize() {
        let mut weights = DriftWeights {
            state: 2.0,
            causal: 2.0,
            context: 1.0,
        };
        weights.normalize();
        assert!(weights.validate().is_ok());
        assert!((weights.state - 0.4).abs() < 0.01);
    }

    #[test]
    fn test_context_divergence_partial_overlap() {
        let shared = uuid::Uuid::now_v7();
        let only_a = uuid::Uuid::now_v7();
        let only_b = uuid::Uuid::now_v7();

        let items_a = vec![shared, only_a];
        let items_b = vec![shared, only_b];

        let divergence = DriftMeter::compute_context_divergence(&items_a, &items_b);
        // Jaccard: 1 intersection / 3 union = 0.333, divergence = 0.667
        assert!((divergence - 0.667).abs() < 0.01);
    }

    // ── NaN / Infinity production guard tests ────────────────────────

    #[test]
    fn test_nan_weight_detected() {
        let w = DriftWeights {
            state: f32::NAN,
            causal: 0.5,
            context: 0.5,
        };
        assert_eq!(w.validate(), Err(DriftWeightsError::NonFinite));
    }

    #[test]
    fn test_infinity_weight_detected() {
        let w = DriftWeights {
            state: f32::INFINITY,
            causal: 0.0,
            context: 0.0,
        };
        assert_eq!(w.validate(), Err(DriftWeightsError::NonFinite));
    }

    #[test]
    fn test_neg_infinity_weight_detected() {
        let w = DriftWeights {
            state: f32::NEG_INFINITY,
            causal: 0.5,
            context: 0.5,
        };
        assert_eq!(w.validate(), Err(DriftWeightsError::NonFinite));
    }

    #[test]
    fn test_negative_weight_detected() {
        let w = DriftWeights {
            state: -0.1,
            causal: 0.6,
            context: 0.5,
        };
        assert_eq!(w.validate(), Err(DriftWeightsError::Negative));
    }

    #[test]
    fn test_sum_not_one_detected() {
        let w = DriftWeights {
            state: 0.5,
            causal: 0.5,
            context: 0.5,
        };
        match w.validate() {
            Err(DriftWeightsError::SumNotOne(sum)) => {
                assert!((sum - 1.5).abs() < 0.001);
            }
            other => panic!("Expected SumNotOne, got {:?}", other),
        }
    }

    #[test]
    fn test_all_zero_weights_normalize_unchanged() {
        let mut w = DriftWeights {
            state: 0.0,
            causal: 0.0,
            context: 0.0,
        };
        let before = w.clone();
        w.normalize();
        assert_eq!(w, before, "all-zero weights should be left unchanged");
    }

    #[test]
    fn test_nan_weights_normalize_unchanged() {
        let mut w = DriftWeights {
            state: f32::NAN,
            causal: 0.5,
            context: 0.5,
        };
        // NaN sum means normalize should leave weights unchanged
        w.normalize();
        // state stays NaN (NaN != NaN, so check with is_nan)
        assert!(w.state.is_nan());
        assert_eq!(w.causal, 0.5);
        assert_eq!(w.context, 0.5);
    }

    #[test]
    fn test_infinity_weights_normalize_unchanged() {
        let mut w = DriftWeights {
            state: f32::INFINITY,
            causal: 0.5,
            context: 0.5,
        };
        let causal_before = w.causal;
        let context_before = w.context;
        w.normalize();
        // Infinity sum is not finite, so weights should be unchanged
        assert_eq!(w.state, f32::INFINITY);
        assert_eq!(w.causal, causal_before);
        assert_eq!(w.context, context_before);
    }

    #[test]
    fn test_nan_composite_score_guarded() {
        // NaN weights should produce composite_score = 0.0, not NaN
        let nan_weights = DriftWeights {
            state: f32::NAN,
            causal: 0.5,
            context: 0.5,
        };
        let meter = DriftMeter::compute(DriftInput {
            agent_a: AgentId::now_v7(),
            agent_b: AgentId::now_v7(),
            events_a: &[],
            events_b: &[],
            context_items_a: &[],
            context_items_b: &[],
            decay: &DecayParams::exponential(10.0),
            weights: &nan_weights,
            threshold: 0.85,
        });
        assert!(
            meter.composite_score.is_finite(),
            "composite_score must be finite, got {}",
            meter.composite_score
        );
        assert_eq!(meter.composite_score, 0.0);
        assert!(!meter.is_drifting);
    }

    #[test]
    fn test_infinity_in_composite_score_guarded() {
        // Infinity weights should produce composite_score = 0.0, not Infinity
        let inf_weights = DriftWeights {
            state: f32::INFINITY,
            causal: 0.0,
            context: 0.0,
        };
        let meter = DriftMeter::compute(DriftInput {
            agent_a: AgentId::now_v7(),
            agent_b: AgentId::now_v7(),
            events_a: &[],
            events_b: &[],
            context_items_a: &[],
            context_items_b: &[],
            decay: &DecayParams::exponential(10.0),
            weights: &inf_weights,
            threshold: 0.85,
        });
        assert!(
            meter.composite_score.is_finite(),
            "composite_score must be finite, got {}",
            meter.composite_score
        );
        assert_eq!(meter.composite_score, 0.0);
    }

    #[test]
    fn test_validate_returns_specific_error_variants() {
        // NonFinite takes priority over other checks
        let w = DriftWeights {
            state: f32::NAN,
            causal: -1.0,
            context: 0.5,
        };
        assert!(matches!(w.validate(), Err(DriftWeightsError::NonFinite)));

        // Negative takes priority over SumNotOne
        let w = DriftWeights {
            state: -0.5,
            causal: 1.0,
            context: 1.0,
        };
        assert!(matches!(w.validate(), Err(DriftWeightsError::Negative)));

        // SumNotOne when all are valid but don't sum to 1
        let w = DriftWeights {
            state: 0.1,
            causal: 0.1,
            context: 0.1,
        };
        assert!(matches!(w.validate(), Err(DriftWeightsError::SumNotOne(_))));
    }

    #[test]
    fn test_mixed_nan_valid_weights() {
        // Only one NaN among otherwise valid values
        let w = DriftWeights {
            state: 0.4,
            causal: f32::NAN,
            context: 0.25,
        };
        assert_eq!(w.validate(), Err(DriftWeightsError::NonFinite));
    }

    #[test]
    fn test_drift_weights_error_display() {
        let e = DriftWeightsError::NonFinite;
        assert!(e.to_string().contains("NaN or Infinity"));
        let e = DriftWeightsError::Negative;
        assert!(e.to_string().contains("negative"));
        let e = DriftWeightsError::SumNotOne(1.5);
        assert!(e.to_string().contains("1.5"));
    }
}