Session-Aware Guards

/// Thread-safe, append-only, hash-chained session journal.
///
/// Create one per session and share it (via `Arc<SessionJournal>`) with all
/// guards that need session-aware context.
pub struct SessionJournal {
    inner: Mutex<JournalInner>,
    session_id: String,
}

/// Inner journal state (not thread-safe -- wrapped by `SessionJournal`).
#[derive(Debug)]
struct JournalInner {
    /// The capacity-bounded ring of retained entries.
    entries: chio_bounded::Ring<JournalEntry>,
    /// Running fold over the hashes of entries evicted from the ring, so the
    /// chain stays committed after the prefix is dropped. `None` until the first
    /// eviction.
    evicted_head_hash: Option<String>,
    /// Monotonic sequence counter, independent of the (bounded) ring length so
    /// sequence numbers never repeat after eviction.
    next_sequence: u64,
    /// Cumulative data flow stats.
    data_flow: CumulativeDataFlow,
    /// Tool invocation sequence (tool names in order), bounded to the entry cap.
    tool_sequence: chio_bounded::Ring<String>,
    /// Per-tool invocation counts (cumulative). The counts survive entry-ring
    /// eviction so the behavioral-sequence guard can answer "was this tool ever
    /// invoked", but the set of DISTINCT keys is bounded fail-closed by
    /// `tool_counts_cap`. See `record` for the overflow rule.
    tool_counts: HashMap<String, u64>,
    /// Maximum number of distinct tool names retained in `tool_counts`. Once
    /// reached, a previously-unseen tool name is dropped from the cumulative
    /// counts (fail-closed): a dependent required-predecessor check then treats
    /// it as never-invoked and denies. Already-seen tools keep
    /// counting so legitimate (registry-bounded) predecessor checks stay
    /// correct across ring eviction.
    tool_counts_cap: usize,
    /// The tool name of the current consecutive run, or `None` before the first
    /// record. Together with `current_streak_len` this is an O(1), bounded
    /// cumulative streak counter (last-tool plus consecutive-count) that survives
    /// entry-ring eviction, so a `max_consecutive` check can be enforced even when
    /// the streak is longer than `journal_entry_cap` and the older part of the
    /// streak has been evicted from `tool_sequence`.
    current_streak_tool: Option<String>,
    /// Length of the current consecutive run of `current_streak_tool`. Reset to 1
    /// when a different tool is recorded, so it never accumulates unboundedly as a
    /// collection: it is a single running scalar.
    current_streak_len: u64,
}

/// Cumulative data flow statistics for a session.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct CumulativeDataFlow {
    /// Total bytes read across all invocations in the session.
    pub total_bytes_read: u64,
    /// Total bytes written across all invocations in the session.
    pub total_bytes_written: u64,
    /// Total number of tool invocations recorded.
    pub total_invocations: u64,
    /// Maximum delegation depth seen in the session.
    pub max_delegation_depth: u32,
}

/// Configuration for cumulative data flow limits.
#[derive(Clone, Debug, Default)]
pub struct DataFlowConfig {
    /// Maximum cumulative bytes read per session. None means unlimited.
    pub max_bytes_read: Option<u64>,
    /// Maximum cumulative bytes written per session. None means unlimited.
    pub max_bytes_written: Option<u64>,
    /// Maximum cumulative bytes (read + written) per session. None means unlimited.
    pub max_bytes_total: Option<u64>,
}

/// Guard that enforces cumulative data flow limits using the session journal.
///
/// Reads the journal's cumulative data flow statistics and denies requests
/// if any configured limit has been reached.
pub struct DataFlowGuard {
    journal: Arc<SessionJournal>,
    config: DataFlowConfig,
}

fn evaluate(&self, _ctx: &GuardContext) -> Result<GuardDecision, KernelError> {
    let snapshot = self.journal.snapshot().map_err(|e| {
        KernelError::Internal(format!("data-flow guard journal error (fail-closed): {e}"))
    })?;
    let flow = snapshot.data_flow;

    // Check bytes read limit.
    if let Some(max_read) = self.config.max_bytes_read {
        if flow.total_bytes_read >= max_read {
            return Ok(GuardDecision::deny(Vec::new()));
        }
    }

    // Check bytes written limit.
    if let Some(max_written) = self.config.max_bytes_written {
        if flow.total_bytes_written >= max_written {
            return Ok(GuardDecision::deny(Vec::new()));
        }
    }

    // Check total I/O limit.
    if let Some(max_total) = self.config.max_bytes_total {
        let total = flow
            .total_bytes_read
            .saturating_add(flow.total_bytes_written);
        if total >= max_total {
            return Ok(GuardDecision::deny(Vec::new()));
        }
    }

    Ok(GuardDecision::allow())
}

/// Policy configuration for the behavioral sequence guard.
#[derive(Clone, Debug, Default)]
pub struct SequencePolicy {
    /// Tools that must have been invoked before a given tool can run.
    /// Map from tool_name to set of required predecessor tools.
    pub required_predecessors: HashMap<String, HashSet<String>>,
    /// Forbidden immediate transitions: (from_tool, to_tool) pairs.
    /// If the last invoked tool is `from_tool`, then `to_tool` is denied.
    pub forbidden_transitions: Vec<(String, String)>,
    /// Maximum consecutive invocations of the same tool.
    /// None means unlimited.
    pub max_consecutive: Option<u32>,
    /// If set, the first tool in the session must match this name.
    pub required_first_tool: Option<String>,
}

/// Guard that enforces tool ordering policies using the session journal.
pub struct BehavioralSequenceGuard {
    journal: Arc<SessionJournal>,
    policy: SequencePolicy,
}

fn evaluate(&self, ctx: &GuardContext) -> Result<GuardDecision, KernelError> {
    let tool_name = &ctx.request.tool_name;

    let snapshot = self.journal.snapshot().map_err(|e| {
        KernelError::Internal(format!(
            "behavioral-sequence guard journal error (fail-closed): {e}"
        ))
    })?;

    // Check required first tool. "Have any tools run yet" is a cumulative
    // property, so consult the journal's cumulative O(1) last-tool field
    // (`current_streak_tool`), which is `None` only before the first record,
    // NOT the bounded `tool_sequence` ring. When `journal_entry_cap` is 0 the
    // ring stores no tool names (capacity 0 = disabled), so `tool_sequence`
    // would report EVERY call as the first and mis-fire this check; the
    // cumulative field is correct at any entry cap.
    if snapshot.current_streak_tool.is_none() {
        if let Some(ref required_first) = self.policy.required_first_tool {
            if tool_name != required_first {
                return Ok(GuardDecision::deny(Vec::new()));
            }
        }
    }

    // Check required predecessors. "Has this tool ever been invoked" is a
    // cumulative property, so consult the journal's cumulative `tool_counts`
    // (which survives ring eviction) rather than the
    // bounded `tool_sequence` tail. A predecessor invoked once and then
    // pushed out of the retained window is still known to have run, so a
    // workflow that runs setup once and then more than `journal_entry_cap`
    // other calls is no longer falsely denied when a dependent tool needs
    // that evicted predecessor. `tool_counts` cannot grow without bound: the
    // journal caps its distinct-key set fail-closed
    // (`journal_tool_counts_cap`). Legitimate registry-bounded predecessors
    // stay recorded, but a predecessor that overflowed the cap is absent here
    // and therefore denies (fail-closed) rather than being falsely treated as
    // invoked.
    if let Some(required) = self.policy.required_predecessors.get(tool_name) {
        for req in required {
            if !snapshot.tool_counts.contains_key(req) {
                return Ok(GuardDecision::deny(Vec::new()));
            }
        }
    }

    // Check forbidden transitions. The "last invoked tool" comes from the
    // journal's cumulative O(1) last-tool field (`current_streak_tool`), NOT
    // the bounded `tool_sequence` tail: when `journal_entry_cap` is 0 the ring
    // stores no tool names (capacity 0 = disabled), so `tool_sequence.last()`
    // is always None and a forbidden transition would silently never fire
    // (fail-OPEN), letting a memory-budget setting disable a transition-deny
    // policy. The cumulative field tracks the most recent recorded tool at any
    // entry cap, so the check holds fail-closed.
    if let Some(last_tool) = snapshot.current_streak_tool.as_deref() {
        for (from, to) in &self.policy.forbidden_transitions {
            if last_tool == from && tool_name == to {
                return Ok(GuardDecision::deny(Vec::new()));
            }
        }
    }

    // Check max consecutive. The count of prior consecutive same-tool
    // invocations comes from the journal's cumulative O(1) streak counter
    // (`current_streak_tool` + `current_streak_len`), NOT the bounded
    // `tool_sequence` tail. When `journal_entry_cap` is smaller than
    // `max_consecutive`, the ring evicts the older part of a same-tool streak,
    // so counting the retained tail would undercount and ALLOW a call that
    // must be DENIED. The cumulative counter survives ring eviction, so the
    // streak limit holds regardless of the entry cap. If the
    // request tool differs from the current-streak tool, no prior consecutive
    // run exists for it (it would start a fresh streak).
    if let Some(max_consec) = self.policy.max_consecutive {
        let prior_streak =
            if snapshot.current_streak_tool.as_deref() == Some(tool_name.as_str()) {
                snapshot.current_streak_len
            } else {
                0
            };
        if prior_streak >= u64::from(max_consec) {
            return Ok(GuardDecision::deny(Vec::new()));
        }
    }

    Ok(GuardDecision::allow())
}

/// Default EMA smoothing factor. Equivalent to a ~10-sample window.
pub const DEFAULT_EMA_ALPHA: f64 = 0.2;
/// Default sigma threshold above which a window is flagged.
pub const DEFAULT_SIGMA_THRESHOLD: f64 = 2.0;
/// Default rolling window length in seconds.
pub const DEFAULT_WINDOW_SECS: u64 = 60;
/// Default number of historical windows used to prime the baseline
/// before the guard starts emitting signals. Guarantees the z-score
/// has enough history to be meaningful.
pub const DEFAULT_BASELINE_MIN_WINDOWS: u64 = 3;

/// EMA (exponentially-weighted moving average) baseline state for a
/// single (agent, metric) pair. Used by behavioral profiling to detect
/// z-score anomalies without storing every historical sample.
///
/// The baseline uses Welford-style incremental tracking of mean and
/// variance so callers can compute a z-score for any new sample
/// without re-reading history.
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)]
#[serde(rename_all = "camelCase")]
pub struct EmaBaselineState {
    /// Number of samples folded into the baseline.
    pub sample_count: u64,
    /// Exponentially-weighted mean.
    pub ema_mean: f64,
    /// Exponentially-weighted variance.
    pub ema_variance: f64,
    /// Last update timestamp (unix seconds).
    pub last_update: u64,
}

impl EmaBaselineState {
    /// Fold a new sample into the baseline with the provided smoothing
    /// factor `alpha` (0.0..=1.0). Higher alpha weighs recent samples
    /// more heavily.
    ///
    /// `alpha` is clamped to `(0.0, 1.0]`. `now` is recorded as
    /// `last_update`.
    pub fn update(&mut self, sample: f64, alpha: f64, now: u64) {
        let alpha = alpha.clamp(f64::MIN_POSITIVE, 1.0);
        if self.sample_count == 0 {
            self.ema_mean = sample;
            self.ema_variance = 0.0;
        } else {
            let prev_mean = self.ema_mean;
            self.ema_mean = prev_mean + alpha * (sample - prev_mean);
            // Incremental EWMA variance, following West (1979) / Welford.
            let diff = sample - prev_mean;
            self.ema_variance = (1.0 - alpha) * (self.ema_variance + alpha * diff * diff);
        }
        self.sample_count = self.sample_count.saturating_add(1);
        self.last_update = now;
    }

    /// Standard deviation (sqrt of EWMA variance).
    #[must_use]
    pub fn stddev(&self) -> f64 {
        self.ema_variance.max(0.0).sqrt()
    }

    /// Z-score for a new sample. Returns `None` when the baseline has
    /// fewer than two samples or zero variance (no meaningful signal).
    #[must_use]
    pub fn z_score(&self, sample: f64) -> Option<f64> {
        if self.sample_count < 2 {
            return None;
        }
        let stddev = self.stddev();
        if stddev <= f64::EPSILON {
            return None;
        }
        Some((sample - self.ema_mean) / stddev)
    }
}

/// Z-score with a Poisson-style stddev floor.
///
/// For count metrics (call rate, deny count, unique tools) a zero
/// measured variance is an artifact of a short baseline rather than a
/// true zero-noise process. We floor the effective stddev at
/// `sqrt(max(mean, 1))` so that a 50x spike over a steady 10/window
/// baseline is detected as an anomaly even when the EWMA variance
/// happens to be numerically zero.
fn robust_z_score(state: &EmaBaselineState, sample: f64) -> Option<f64> {
    if state.sample_count < 2 {
        return None;
    }
    let measured = state.stddev();
    let floor = state.ema_mean.max(1.0).sqrt();
    let effective = measured.max(floor);
    if effective <= f64::EPSILON {
        return None;
    }
    Some((sample - state.ema_mean) / effective)
}

fn current_window_start(&self, now: u64) -> u64 {
    let window = self.config.window_secs.max(1);
    (now / window) * window
}

/// Summary of behavioral-anomaly signals derived from receipts over a
/// window. Used by `BehavioralProfileGuard` and surfaced in operator
/// UIs.
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)]
#[serde(rename_all = "camelCase")]
pub struct BehavioralAnomalyScore {
    /// Agent subject this anomaly score applies to.
    pub agent_id: String,
    /// Baseline statistic the z-score is computed against.
    pub baseline: EmaBaselineState,
    /// Current-window sample value (e.g. call count per window).
    pub current_sample: f64,
    /// Computed z-score, or `None` when baseline is too small.
    pub z_score: Option<f64>,
    /// Threshold above which an advisory signal is raised.
    pub sigma_threshold: f64,
    /// Whether the current sample crossed the threshold.
    pub anomaly: bool,
    /// Unix timestamp (seconds) at which the score was computed.
    pub generated_at: u64,
}

use std::sync::Arc;
use chio_guards::{
    DataFlowGuard, DataFlowConfig,
    BehavioralSequenceGuard, SequencePolicy,
    BehavioralProfileGuard, BehavioralProfileConfig,
    InMemoryReceiptFeed,
};
use chio_http_session::SessionJournal;

let journal = Arc::new(SessionJournal::new("sess-1".to_string()));
let feed = InMemoryReceiptFeed::new();

let mut pipeline = chio_guards::GuardPipeline::new();

pipeline.add(Box::new(DataFlowGuard::new(
    journal.clone(),
    DataFlowConfig {
        max_bytes_read: Some(50 * 1024 * 1024),
        max_bytes_written: Some(10 * 1024 * 1024),
        max_bytes_total: None,
    },
)));

let mut policy = SequencePolicy::default();
policy.required_first_tool = Some("init".to_string());
pipeline.add(Box::new(BehavioralSequenceGuard::new(journal.clone(), policy)));

pipeline.add(Box::new(BehavioralProfileGuard::with_config(
    Box::new(feed),
    BehavioralProfileConfig::default(),
)));