gemini_adk_rs/

state.rs

//! Typed key-value state container for agents.
//!
//! Supports optional delta tracking for transactional state management
//! and prefix-scoped accessors for namespace isolation.

use std::collections::{HashMap, VecDeque};
use std::marker::PhantomData;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::time::SystemTime;

use dashmap::DashMap;
use serde_json::Value;

const DEFAULT_MUTATION_JOURNAL_CAPACITY: usize = 1024;

/// A compile-time typed state key that eliminates typo bugs and type mismatches.
///
/// Create as a const and use with `State::get_key()` / `State::set_key()`:
///
/// ```rust,ignore
/// const TURN_COUNT: StateKey<u32> = StateKey::new("session:turn_count");
/// const SENTIMENT: StateKey<String> = StateKey::new("derived:sentiment");
///
/// state.set_key(&TURN_COUNT, 5);
/// let count: Option<u32> = state.get_key(&TURN_COUNT);
/// ```
pub struct StateKey<T> {
    key: &'static str,
    _phantom: PhantomData<fn() -> T>,
}

impl<T> StateKey<T> {
    /// Create a new typed state key.
    pub const fn new(key: &'static str) -> Self {
        Self {
            key,
            _phantom: PhantomData,
        }
    }

    /// The string key.
    pub const fn key(&self) -> &'static str {
        self.key
    }
}

/// Where a state mutation came from.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum StateMutationOrigin {
    /// Regular `State::set` or prefixed state write.
    Set,
    /// Direct committed-store write that bypasses delta tracking.
    SetCommitted,
    /// Removal of a single key.
    Remove,
    /// Removal caused by clearing a prefix.
    ClearPrefix,
    /// Delta changes committed into the base state.
    Commit,
}

/// A single state mutation recorded in the bounded mutation journal.
#[derive(Debug, Clone, PartialEq)]
pub struct StateMutation {
    /// Monotonic sequence number assigned when the mutation was recorded.
    pub sequence: u64,
    /// State key that changed.
    pub key: String,
    /// Value before the mutation, or `None` when the key did not exist.
    pub old: Option<Value>,
    /// Value after the mutation, or `None` when the key was removed.
    pub new: Option<Value>,
    /// Operation that recorded the mutation.
    pub origin: StateMutationOrigin,
    /// Wall-clock time at which the mutation was recorded.
    pub timestamp: SystemTime,
    /// Whether the mutation was written to a delta-tracked view.
    pub delta: bool,
}

/// A concurrent, type-safe state container that agents read from and write to.
///
/// By default, `set()` writes directly to the inner store. When delta tracking
/// is enabled via `with_delta_tracking()`, writes go to a separate delta map
/// that can be committed or rolled back.
#[derive(Debug, Clone)]
pub struct State {
    inner: Arc<DashMap<String, Value>>,
    delta: Arc<DashMap<String, Value>>,
    mutations: Arc<std::sync::Mutex<VecDeque<StateMutation>>>,
    next_mutation_sequence: Arc<AtomicU64>,
    mutation_capacity: usize,
    track_delta: bool,
}

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

impl State {
    /// Create a new empty state container.
    pub fn new() -> Self {
        Self {
            inner: Arc::new(DashMap::new()),
            delta: Arc::new(DashMap::new()),
            mutations: Arc::new(std::sync::Mutex::new(VecDeque::new())),
            next_mutation_sequence: Arc::new(AtomicU64::new(1)),
            mutation_capacity: DEFAULT_MUTATION_JOURNAL_CAPACITY,
            track_delta: false,
        }
    }

    /// Create a new State with delta tracking enabled.
    /// Writes go to the delta map; reads check delta first, then inner.
    pub fn with_delta_tracking(&self) -> State {
        State {
            inner: self.inner.clone(),
            delta: Arc::new(DashMap::new()),
            mutations: self.mutations.clone(),
            next_mutation_sequence: self.next_mutation_sequence.clone(),
            mutation_capacity: self.mutation_capacity,
            track_delta: true,
        }
    }

    /// Get a value by key, attempting to deserialize to the requested type.
    /// When delta tracking is enabled, checks delta first, then inner.
    pub fn get<T: serde::de::DeserializeOwned>(&self, key: &str) -> Option<T> {
        self.get_raw(key)
            .and_then(|v| serde_json::from_value(v).ok())
    }

    /// Borrow a value by key without cloning, applying `f` to the reference.
    ///
    /// This is the zero-copy alternative to `get_raw()`. The closure receives
    /// a `&Value` directly from the DashMap ref-guard, avoiding the
    /// `Value::clone()` + `serde_json::from_value()` overhead of `get()`.
    ///
    /// Lookup order: delta (if tracking) → inner → derived fallback.
    pub fn with<F, R>(&self, key: &str, f: F) -> Option<R>
    where
        F: FnOnce(&Value) -> R,
    {
        if self.track_delta {
            if let Some(ref_multi) = self.delta.get(key) {
                return Some(f(ref_multi.value()));
            }
        }
        if let Some(ref_multi) = self.inner.get(key) {
            return Some(f(ref_multi.value()));
        }
        if !key.contains(':') {
            let mut derived_key = String::with_capacity(8 + key.len());
            use std::fmt::Write;
            let _ = write!(derived_key, "derived:{}", key);
            if self.track_delta {
                if let Some(ref_multi) = self.delta.get(&derived_key) {
                    return Some(f(ref_multi.value()));
                }
            }
            if let Some(ref_multi) = self.inner.get(&derived_key) {
                return Some(f(ref_multi.value()));
            }
        }
        None
    }

    /// Get a raw JSON value by key.
    /// When delta tracking is enabled, checks delta first, then inner.
    /// If the key is not found and doesn't contain a prefix, also checks `derived:{key}`
    /// as a transparent fallback for computed variables.
    pub fn get_raw(&self, key: &str) -> Option<Value> {
        if self.track_delta {
            if let Some(v) = self.delta.get(key) {
                return Some(v.value().clone());
            }
        }
        if let Some(v) = self.inner.get(key) {
            return Some(v.value().clone());
        }
        // Transparent derived fallback: if key has no prefix, check derived:{key}
        if !key.contains(':') {
            use std::fmt::Write;
            let mut derived_key = String::with_capacity(8 + key.len());
            let _ = write!(derived_key, "derived:{}", key);
            if self.track_delta {
                if let Some(v) = self.delta.get(&derived_key) {
                    return Some(v.value().clone());
                }
            }
            return self.inner.get(&derived_key).map(|v| v.value().clone());
        }
        None
    }

    /// Get a typed value using a `StateKey<T>`.
    pub fn get_key<T: serde::de::DeserializeOwned>(&self, key: &StateKey<T>) -> Option<T> {
        self.get(key.key())
    }

    /// Set a typed value using a `StateKey<T>`.
    pub fn set_key<T: serde::Serialize>(&self, key: &StateKey<T>, value: T) {
        self.set(key.key(), value);
    }

    /// Zero-copy borrow using a `StateKey<T>`.
    pub fn with_key<T, F, R>(&self, key: &StateKey<T>, f: F) -> Option<R>
    where
        F: FnOnce(&Value) -> R,
    {
        self.with(key.key(), f)
    }

    /// Set a value by key.
    /// When delta tracking is enabled, writes to delta instead of inner.
    pub fn set(&self, key: impl Into<String>, value: impl serde::Serialize) {
        let key = key.into();
        let v = serde_json::to_value(value).expect("value must be serializable");
        let old = self.get_raw(&key);
        if self.track_delta {
            self.delta.insert(key.clone(), v.clone());
            self.record_mutation(key, old, Some(v), StateMutationOrigin::Set);
        } else {
            self.inner.insert(key.clone(), v.clone());
            self.record_mutation(key, old, Some(v), StateMutationOrigin::Set);
        }
    }

    /// Set a value directly in the committed store, bypassing delta tracking.
    pub fn set_committed(&self, key: impl Into<String>, value: impl serde::Serialize) {
        let key = key.into();
        let v = serde_json::to_value(value).expect("value must be serializable");
        let old = self.inner.insert(key.clone(), v.clone());
        self.record_mutation(key, old, Some(v), StateMutationOrigin::SetCommitted);
    }

    /// Atomically read-modify-write a value.
    ///
    /// If the key doesn't exist, `default` is used as the initial value.
    /// The function `f` receives the current value and returns the new value.
    /// Returns the new value after modification.
    pub fn modify<T, F>(&self, key: &str, default: T, f: F) -> T
    where
        T: serde::Serialize + serde::de::DeserializeOwned,
        F: FnOnce(T) -> T,
    {
        // Read current value from whichever store has it
        let current: T = self.get(key).unwrap_or(default);
        let new_val = f(current);
        self.set(key, &new_val);
        new_val
    }

    /// Check if a key exists (in delta or inner).
    pub fn contains(&self, key: &str) -> bool {
        if self.track_delta && self.delta.contains_key(key) {
            return true;
        }
        self.inner.contains_key(key)
    }

    /// Remove a key.
    pub fn remove(&self, key: &str) -> Option<Value> {
        if self.track_delta {
            // Remove from delta if present, but also check inner
            let from_delta = self.delta.remove(key).map(|(_, v)| v);
            let from_inner = self.inner.remove(key).map(|(_, v)| v);
            let removed = from_delta.or(from_inner);
            if let Some(ref old) = removed {
                self.record_mutation(
                    key.to_string(),
                    Some(old.clone()),
                    None,
                    StateMutationOrigin::Remove,
                );
            }
            removed
        } else {
            let removed = self.inner.remove(key).map(|(_, v)| v);
            if let Some(ref old) = removed {
                self.record_mutation(
                    key.to_string(),
                    Some(old.clone()),
                    None,
                    StateMutationOrigin::Remove,
                );
            }
            removed
        }
    }

    /// Get all keys (from both inner and delta when tracking).
    pub fn keys(&self) -> Vec<String> {
        if !self.track_delta || self.delta.is_empty() {
            return self.inner.iter().map(|r| r.key().clone()).collect();
        }
        let mut seen =
            std::collections::HashSet::with_capacity(self.inner.len() + self.delta.len());
        let mut keys = Vec::with_capacity(self.inner.len() + self.delta.len());
        for entry in self.inner.iter() {
            let key = entry.key().clone();
            seen.insert(key.clone());
            keys.push(key);
        }
        for entry in self.delta.iter() {
            let key = entry.key().clone();
            if seen.insert(key.clone()) {
                keys.push(key);
            }
        }
        keys
    }

    /// Create a new State containing only the specified keys.
    pub fn pick(&self, keys: &[&str]) -> State {
        let new = State::new();
        for key in keys {
            if let Some(v) = self.get_raw(key) {
                new.set(*key, v);
            }
        }
        new
    }

    /// Merge another state into this one (other's values overwrite on conflict).
    pub fn merge(&self, other: &State) {
        for entry in other.inner.iter() {
            self.set(entry.key().clone(), entry.value().clone());
        }
    }

    /// Rename a key.
    pub fn rename(&self, from: &str, to: &str) {
        if let Some(v) = self.remove(from) {
            if self.track_delta {
                self.set(to.to_string(), v);
            } else {
                self.set(to.to_string(), v);
            }
        }
    }

    // ── Delta methods ──────────────────────────────────────────────────────

    /// Whether delta tracking is enabled.
    pub fn is_tracking_delta(&self) -> bool {
        self.track_delta
    }

    /// Whether there are uncommitted delta changes.
    pub fn has_delta(&self) -> bool {
        self.track_delta && !self.delta.is_empty()
    }

    /// Get a snapshot of the current delta.
    pub fn delta(&self) -> HashMap<String, Value> {
        self.delta
            .iter()
            .map(|entry| (entry.key().clone(), entry.value().clone()))
            .collect()
    }

    /// Commit delta changes into the inner store, then clear the delta.
    pub fn commit(&self) {
        for entry in self.delta.iter() {
            let key = entry.key().clone();
            let value = entry.value().clone();
            let old = self.inner.insert(key.clone(), value.clone());
            self.record_mutation_with_delta(
                key,
                old,
                Some(value),
                StateMutationOrigin::Commit,
                false,
            );
        }
        self.delta.clear();
    }

    /// Discard all uncommitted delta changes.
    pub fn rollback(&self) {
        self.delta.clear();
    }

    // ── Prefix accessors ───────────────────────────────────────────────────

    /// Access state with the `app:` prefix scope.
    pub fn app(&self) -> PrefixedState<'_> {
        PrefixedState {
            state: self,
            prefix: "app:",
        }
    }

    /// Access state with the `user:` prefix scope.
    pub fn user(&self) -> PrefixedState<'_> {
        PrefixedState {
            state: self,
            prefix: "user:",
        }
    }

    /// Access state with the `temp:` prefix scope.
    pub fn temp(&self) -> PrefixedState<'_> {
        PrefixedState {
            state: self,
            prefix: "temp:",
        }
    }

    /// Access state with the `session:` prefix scope (auto-tracked signals).
    pub fn session(&self) -> PrefixedState<'_> {
        PrefixedState {
            state: self,
            prefix: "session:",
        }
    }

    /// Access state with the `turn:` prefix scope (reset each turn).
    pub fn turn(&self) -> PrefixedState<'_> {
        PrefixedState {
            state: self,
            prefix: "turn:",
        }
    }

    /// Access state with the `bg:` prefix scope (background tasks).
    pub fn bg(&self) -> PrefixedState<'_> {
        PrefixedState {
            state: self,
            prefix: "bg:",
        }
    }

    /// Access read-only state with the `derived:` prefix scope (computed vars only).
    pub fn derived(&self) -> ReadOnlyPrefixedState<'_> {
        ReadOnlyPrefixedState {
            state: self,
            prefix: "derived:",
        }
    }

    // ── Utility methods ───────────────────────────────────────────────────

    /// Snapshot the values of specific keys. Returns HashMap of key -> current value.
    /// Used by watchers to capture state before mutations.
    pub fn snapshot_values(&self, keys: &[&str]) -> HashMap<String, Value> {
        keys.iter()
            .filter_map(|&k| self.get_raw(k).map(|v| (k.to_string(), v)))
            .collect()
    }

    /// Diff current state against a previous snapshot.
    /// Returns Vec of (key, old_value, new_value) for keys that changed.
    pub fn diff_values(
        &self,
        prev: &HashMap<String, Value>,
        keys: &[&str],
    ) -> Vec<(String, Value, Value)> {
        keys.iter()
            .filter_map(|&k| {
                let old = prev.get(k);
                let new = self.get_raw(k);
                match (old, new) {
                    (Some(o), Some(n)) if o != &n => Some((k.to_string(), o.clone(), n)),
                    (None, Some(n)) => Some((k.to_string(), Value::Null, n)),
                    (Some(o), None) => Some((k.to_string(), o.clone(), Value::Null)),
                    _ => None,
                }
            })
            .collect()
    }

    /// Export all state as a HashMap (for persistence/serialization).
    pub fn to_hashmap(&self) -> std::collections::HashMap<String, serde_json::Value> {
        self.inner
            .iter()
            .map(|entry| (entry.key().clone(), entry.value().clone()))
            .collect()
    }

    /// Restore state from a HashMap (for persistence/deserialization).
    pub fn from_hashmap(&self, map: std::collections::HashMap<String, serde_json::Value>) {
        for (key, value) in map {
            self.set_committed(key, value);
        }
    }

    /// Remove all keys with the given prefix.
    pub fn clear_prefix(&self, prefix: &str) {
        let keys_to_remove: Vec<String> = self
            .inner
            .iter()
            .filter(|entry| entry.key().starts_with(prefix))
            .map(|entry| entry.key().clone())
            .collect();
        for key in keys_to_remove {
            if let Some((_, old)) = self.inner.remove(&key) {
                self.record_mutation(key, Some(old), None, StateMutationOrigin::ClearPrefix);
            }
        }
        if self.track_delta {
            let delta_keys: Vec<String> = self
                .delta
                .iter()
                .filter(|entry| entry.key().starts_with(prefix))
                .map(|entry| entry.key().clone())
                .collect();
            for key in delta_keys {
                if let Some((_, old)) = self.delta.remove(&key) {
                    self.record_mutation(key, Some(old), None, StateMutationOrigin::ClearPrefix);
                }
            }
        }
    }

    /// Return a snapshot of recent state mutations.
    pub fn recent_mutations(&self) -> Vec<StateMutation> {
        self.mutations
            .lock()
            .expect("state mutation journal poisoned")
            .iter()
            .cloned()
            .collect()
    }

    /// Return the current monotonic cursor for the mutation journal.
    pub fn mutation_cursor(&self) -> u64 {
        self.next_mutation_sequence.load(Ordering::Relaxed) - 1
    }

    /// Return mutations appended after a previously captured cursor.
    pub fn mutations_since(&self, cursor: u64) -> Vec<StateMutation> {
        let mutations = self
            .mutations
            .lock()
            .expect("state mutation journal poisoned");
        mutations
            .iter()
            .filter(|mutation| mutation.sequence > cursor)
            .cloned()
            .collect()
    }

    /// Drain and return all recorded state mutations.
    pub fn drain_mutations(&self) -> Vec<StateMutation> {
        self.mutations
            .lock()
            .expect("state mutation journal poisoned")
            .drain(..)
            .collect()
    }

    fn record_mutation(
        &self,
        key: String,
        old: Option<Value>,
        new: Option<Value>,
        origin: StateMutationOrigin,
    ) {
        self.record_mutation_with_delta(key, old, new, origin, self.track_delta);
    }

    fn record_mutation_with_delta(
        &self,
        key: String,
        old: Option<Value>,
        new: Option<Value>,
        origin: StateMutationOrigin,
        delta: bool,
    ) {
        let mut mutations = self
            .mutations
            .lock()
            .expect("state mutation journal poisoned");
        if mutations.len() >= self.mutation_capacity {
            mutations.pop_front();
        }
        let sequence = self.next_mutation_sequence.fetch_add(1, Ordering::Relaxed);
        mutations.push_back(StateMutation {
            sequence,
            key,
            old,
            new,
            origin,
            timestamp: SystemTime::now(),
            delta,
        });
    }
}

/// A borrowed view of state that automatically prepends a prefix to all keys.
pub struct PrefixedState<'a> {
    state: &'a State,
    prefix: &'static str,
}

impl<'a> PrefixedState<'a> {
    fn prefixed_key(&self, key: &str) -> String {
        format!("{}{}", self.prefix, key)
    }

    /// Get a value by key (with prefix applied).
    pub fn get<T: serde::de::DeserializeOwned>(&self, key: &str) -> Option<T> {
        self.state.get(&self.prefixed_key(key))
    }

    /// Get a raw JSON value by key (with prefix applied).
    pub fn get_raw(&self, key: &str) -> Option<Value> {
        self.state.get_raw(&self.prefixed_key(key))
    }

    /// Zero-copy borrow a value by key (with prefix applied).
    pub fn with<F, R>(&self, key: &str, f: F) -> Option<R>
    where
        F: FnOnce(&Value) -> R,
    {
        self.state.with(&self.prefixed_key(key), f)
    }

    /// Set a value by key (with prefix applied).
    pub fn set(&self, key: impl AsRef<str>, value: impl serde::Serialize) {
        self.state.set(self.prefixed_key(key.as_ref()), value);
    }

    /// Check if a key exists (with prefix applied).
    pub fn contains(&self, key: &str) -> bool {
        self.state.contains(&self.prefixed_key(key))
    }

    /// Remove a key (with prefix applied).
    pub fn remove(&self, key: &str) -> Option<Value> {
        self.state.remove(&self.prefixed_key(key))
    }

    /// Get all keys within this prefix scope (prefix stripped from results).
    pub fn keys(&self) -> Vec<String> {
        self.state
            .keys()
            .into_iter()
            .filter_map(|k| k.strip_prefix(self.prefix).map(|s| s.to_string()))
            .collect()
    }
}

/// A borrowed, read-only view of state that automatically prepends a prefix to all keys.
///
/// Unlike `PrefixedState`, this does not expose `set()` or `remove()` methods,
/// making it suitable for computed/derived state that user code should not mutate.
pub struct ReadOnlyPrefixedState<'a> {
    state: &'a State,
    prefix: &'static str,
}

impl<'a> ReadOnlyPrefixedState<'a> {
    fn prefixed_key(&self, key: &str) -> String {
        format!("{}{}", self.prefix, key)
    }

    /// Get a value by key (with prefix applied).
    pub fn get<T: serde::de::DeserializeOwned>(&self, key: &str) -> Option<T> {
        self.state.get(&self.prefixed_key(key))
    }

    /// Get a raw JSON value by key (with prefix applied).
    pub fn get_raw(&self, key: &str) -> Option<Value> {
        self.state.get_raw(&self.prefixed_key(key))
    }

    /// Zero-copy borrow a value by key (with prefix applied).
    pub fn with<F, R>(&self, key: &str, f: F) -> Option<R>
    where
        F: FnOnce(&Value) -> R,
    {
        self.state.with(&self.prefixed_key(key), f)
    }

    /// Check if a key exists (with prefix applied).
    pub fn contains(&self, key: &str) -> bool {
        self.state.contains(&self.prefixed_key(key))
    }

    /// Get all keys within this prefix scope (prefix stripped from results).
    pub fn keys(&self) -> Vec<String> {
        self.state
            .keys()
            .into_iter()
            .filter_map(|k| k.strip_prefix(self.prefix).map(|s| s.to_string()))
            .collect()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn set_and_get_string() {
        let state = State::new();
        state.set("name", "Alice");
        assert_eq!(state.get::<String>("name"), Some("Alice".to_string()));
    }

    #[test]
    fn set_and_get_json() {
        let state = State::new();
        state.set("data", serde_json::json!({"temp": 22}));
        let v: Value = state.get("data").unwrap();
        assert_eq!(v["temp"], 22);
    }

    #[test]
    fn pick_subset() {
        let state = State::new();
        state.set("a", 1);
        state.set("b", 2);
        state.set("c", 3);
        let picked = state.pick(&["a", "c"]);
        assert!(picked.contains("a"));
        assert!(!picked.contains("b"));
        assert!(picked.contains("c"));
    }

    #[test]
    fn merge_states() {
        let s1 = State::new();
        s1.set("a", 1);
        let s2 = State::new();
        s2.set("b", 2);
        s1.merge(&s2);
        assert!(s1.contains("a"));
        assert!(s1.contains("b"));
    }

    #[test]
    fn rename_key() {
        let state = State::new();
        state.set("old", "value");
        state.rename("old", "new");
        assert!(!state.contains("old"));
        assert_eq!(state.get::<String>("new"), Some("value".to_string()));
    }

    #[test]
    fn remove_returns_value() {
        let state = State::new();
        state.set("key", 42);
        let removed = state.remove("key");
        assert!(removed.is_some());
        assert!(!state.contains("key"));
    }

    #[test]
    fn get_missing_returns_none() {
        let state = State::new();
        assert_eq!(state.get::<String>("nope"), None);
    }

    // ── Delta tracking tests ──────────────────────────────────────────────

    #[test]
    fn delta_tracking_writes_to_delta() {
        let state = State::new();
        state.set("committed", "yes");

        let tracked = state.with_delta_tracking();
        tracked.set("new_key", "new_value");

        // New key visible through tracked state
        assert_eq!(
            tracked.get::<String>("new_key"),
            Some("new_value".to_string())
        );
        // But NOT visible in original (non-delta) state's inner
        assert!(!state.contains("new_key"));
        // Committed key still visible through tracked state
        assert_eq!(tracked.get::<String>("committed"), Some("yes".to_string()));
    }

    #[test]
    fn delta_has_delta_reports_correctly() {
        let state = State::new();
        let tracked = state.with_delta_tracking();
        assert!(!tracked.has_delta());

        tracked.set("key", "val");
        assert!(tracked.has_delta());
    }

    #[test]
    fn delta_commit_merges_to_inner() {
        let state = State::new();
        let tracked = state.with_delta_tracking();
        tracked.set("key", "val");
        assert!(!state.contains("key"));

        tracked.commit();
        // Now visible in original state
        assert_eq!(state.get::<String>("key"), Some("val".to_string()));
        assert!(!tracked.has_delta());
    }

    #[test]
    fn delta_rollback_discards_changes() {
        let state = State::new();
        let tracked = state.with_delta_tracking();
        tracked.set("key", "val");
        assert!(tracked.has_delta());

        tracked.rollback();
        assert!(!tracked.has_delta());
        assert!(!state.contains("key"));
        assert!(!tracked.contains("key"));
    }

    #[test]
    fn delta_snapshot() {
        let state = State::new();
        let tracked = state.with_delta_tracking();
        tracked.set("a", 1);
        tracked.set("b", 2);

        let snapshot = tracked.delta();
        assert_eq!(snapshot.len(), 2);
        assert!(snapshot.contains_key("a"));
        assert!(snapshot.contains_key("b"));
    }

    #[test]
    fn set_committed_bypasses_delta() {
        let state = State::new();
        let tracked = state.with_delta_tracking();
        tracked.set_committed("direct", "value");

        // Visible immediately in inner
        assert_eq!(state.get::<String>("direct"), Some("value".to_string()));
        // Not in delta
        assert!(!tracked.has_delta());
        // Still visible through tracked (reads inner too)
        assert_eq!(tracked.get::<String>("direct"), Some("value".to_string()));
    }

    #[test]
    fn mutation_journal_records_set_and_remove() {
        let state = State::new();
        state.set("key", "first");
        state.set("key", "second");
        state.remove("key");

        let mutations = state.recent_mutations();
        assert_eq!(mutations.len(), 3);
        assert_eq!(mutations[0].key, "key");
        assert_eq!(mutations[0].old, None);
        assert_eq!(mutations[0].new, Some(serde_json::json!("first")));
        assert_eq!(mutations[0].origin, StateMutationOrigin::Set);

        assert_eq!(mutations[1].old, Some(serde_json::json!("first")));
        assert_eq!(mutations[1].new, Some(serde_json::json!("second")));

        assert_eq!(mutations[2].old, Some(serde_json::json!("second")));
        assert_eq!(mutations[2].new, None);
        assert_eq!(mutations[2].origin, StateMutationOrigin::Remove);
    }

    #[test]
    fn mutation_journal_is_shared_with_delta_tracking() {
        let state = State::new();
        state.set("committed", "yes");

        let tracked = state.with_delta_tracking();
        tracked.set("committed", "maybe");
        tracked.commit();

        let mutations = state.recent_mutations();
        assert_eq!(mutations.len(), 3);
        assert_eq!(mutations[1].key, "committed");
        assert_eq!(mutations[1].old, Some(serde_json::json!("yes")));
        assert_eq!(mutations[1].new, Some(serde_json::json!("maybe")));
        assert_eq!(mutations[1].origin, StateMutationOrigin::Set);
        assert!(mutations[1].delta);

        assert_eq!(mutations[2].origin, StateMutationOrigin::Commit);
        assert!(!mutations[2].delta);
    }

    #[test]
    fn drain_mutations_clears_journal() {
        let state = State::new();
        state.set("a", 1);
        state.set("b", 2);

        let drained = state.drain_mutations();
        assert_eq!(drained.len(), 2);
        assert!(state.recent_mutations().is_empty());
    }

    #[test]
    fn mutation_cursor_reads_only_later_changes() {
        let state = State::new();
        state.set("before", 1);
        let cursor = state.mutation_cursor();

        state.set("after", 2);
        state.remove("before");

        let mutations = state.mutations_since(cursor);
        assert_eq!(mutations.len(), 2);
        assert_eq!(mutations[0].key, "after");
        assert_eq!(mutations[1].key, "before");
    }

    #[test]
    fn no_delta_tracking_preserves_existing_behavior() {
        let state = State::new();
        assert!(!state.is_tracking_delta());
        state.set("key", "val");
        assert_eq!(state.get::<String>("key"), Some("val".to_string()));
        assert!(!state.has_delta());
    }

    // ── Prefix tests ──────────────────────────────────────────────────────

    #[test]
    fn prefix_app_set_and_get() {
        let state = State::new();
        state.app().set("flag", true);

        // Accessible via prefix accessor
        assert_eq!(state.app().get::<bool>("flag"), Some(true));
        // Also accessible via raw key
        assert_eq!(state.get::<bool>("app:flag"), Some(true));
    }

    #[test]
    fn prefix_user_set_and_get() {
        let state = State::new();
        state.user().set("name", "Alice");
        assert_eq!(
            state.user().get::<String>("name"),
            Some("Alice".to_string())
        );
        assert_eq!(state.get::<String>("user:name"), Some("Alice".to_string()));
    }

    #[test]
    fn prefix_temp_set_and_get() {
        let state = State::new();
        state.temp().set("scratch", 42);
        assert_eq!(state.temp().get::<i32>("scratch"), Some(42));
    }

    #[test]
    fn prefix_contains_and_remove() {
        let state = State::new();
        state.app().set("x", 1);
        assert!(state.app().contains("x"));
        state.app().remove("x");
        assert!(!state.app().contains("x"));
    }

    #[test]
    fn prefix_keys() {
        let state = State::new();
        state.app().set("a", 1);
        state.app().set("b", 2);
        state.user().set("c", 3);

        let app_keys = state.app().keys();
        assert_eq!(app_keys.len(), 2);
        assert!(app_keys.contains(&"a".to_string()));
        assert!(app_keys.contains(&"b".to_string()));

        let user_keys = state.user().keys();
        assert_eq!(user_keys.len(), 1);
        assert!(user_keys.contains(&"c".to_string()));
    }

    #[test]
    fn prefix_with_delta_tracking() {
        let state = State::new();
        let tracked = state.with_delta_tracking();
        tracked.app().set("flag", true);

        // Visible in tracked state via prefix
        assert_eq!(tracked.app().get::<bool>("flag"), Some(true));
        // In delta, not committed
        assert!(tracked.has_delta());
        assert!(!state.contains("app:flag"));

        tracked.commit();
        assert_eq!(state.get::<bool>("app:flag"), Some(true));
    }

    // ── New prefix accessor tests ────────────────────────────────────────

    #[test]
    fn prefix_session_set_and_get() {
        let state = State::new();
        state.session().set("turn_count", 5);
        assert_eq!(state.session().get::<i32>("turn_count"), Some(5));
        assert_eq!(state.get::<i32>("session:turn_count"), Some(5));
    }

    #[test]
    fn prefix_turn_set_and_get() {
        let state = State::new();
        state.turn().set("transcript", "hello");
        assert_eq!(
            state.turn().get::<String>("transcript"),
            Some("hello".to_string())
        );
        assert_eq!(
            state.get::<String>("turn:transcript"),
            Some("hello".to_string())
        );
    }

    #[test]
    fn prefix_bg_set_and_get() {
        let state = State::new();
        state.bg().set("task_id", "abc-123");
        assert_eq!(
            state.bg().get::<String>("task_id"),
            Some("abc-123".to_string())
        );
        assert_eq!(
            state.get::<String>("bg:task_id"),
            Some("abc-123".to_string())
        );
    }

    #[test]
    fn prefix_session_contains_and_remove() {
        let state = State::new();
        state.session().set("x", 1);
        assert!(state.session().contains("x"));
        state.session().remove("x");
        assert!(!state.session().contains("x"));
    }

    #[test]
    fn prefix_turn_keys() {
        let state = State::new();
        state.turn().set("a", 1);
        state.turn().set("b", 2);
        state.session().set("c", 3);

        let turn_keys = state.turn().keys();
        assert_eq!(turn_keys.len(), 2);
        assert!(turn_keys.contains(&"a".to_string()));
        assert!(turn_keys.contains(&"b".to_string()));
    }

    // ── ReadOnlyPrefixedState (derived) tests ────────────────────────────

    #[test]
    fn derived_read_only_get() {
        let state = State::new();
        // Write via raw key (simulating ComputedRegistry)
        state.set("derived:sentiment", "positive");
        assert_eq!(
            state.derived().get::<String>("sentiment"),
            Some("positive".to_string())
        );
    }

    #[test]
    fn derived_read_only_get_raw() {
        let state = State::new();
        state.set("derived:score", serde_json::json!(0.95));
        let raw = state.derived().get_raw("score");
        assert!(raw.is_some());
        assert_eq!(raw.unwrap(), serde_json::json!(0.95));
    }

    #[test]
    fn derived_read_only_contains() {
        let state = State::new();
        state.set("derived:exists", true);
        assert!(state.derived().contains("exists"));
        assert!(!state.derived().contains("missing"));
    }

    #[test]
    fn derived_read_only_keys() {
        let state = State::new();
        state.set("derived:a", 1);
        state.set("derived:b", 2);
        state.set("app:c", 3);

        let derived_keys = state.derived().keys();
        assert_eq!(derived_keys.len(), 2);
        assert!(derived_keys.contains(&"a".to_string()));
        assert!(derived_keys.contains(&"b".to_string()));
    }

    #[test]
    fn derived_missing_key_returns_none() {
        let state = State::new();
        assert_eq!(state.derived().get::<String>("nope"), None);
        assert_eq!(state.derived().get_raw("nope"), None);
    }

    // ── snapshot_values tests ────────────────────────────────────────────

    #[test]
    fn snapshot_values_captures_existing_keys() {
        let state = State::new();
        state.set("a", 1);
        state.set("b", "hello");
        state.set("c", true);

        let snap = state.snapshot_values(&["a", "b", "missing"]);
        assert_eq!(snap.len(), 2);
        assert_eq!(snap.get("a"), Some(&serde_json::json!(1)));
        assert_eq!(snap.get("b"), Some(&serde_json::json!("hello")));
        assert!(!snap.contains_key("missing"));
    }

    #[test]
    fn snapshot_values_empty_keys() {
        let state = State::new();
        state.set("a", 1);
        let snap = state.snapshot_values(&[]);
        assert!(snap.is_empty());
    }

    // ── diff_values tests ────────────────────────────────────────────────

    #[test]
    fn diff_values_detects_changed_value() {
        let state = State::new();
        state.set("x", 1);
        let snap = state.snapshot_values(&["x"]);

        state.set("x", 2);
        let diffs = state.diff_values(&snap, &["x"]);
        assert_eq!(diffs.len(), 1);
        assert_eq!(diffs[0].0, "x");
        assert_eq!(diffs[0].1, serde_json::json!(1));
        assert_eq!(diffs[0].2, serde_json::json!(2));
    }

    #[test]
    fn diff_values_detects_new_key() {
        let state = State::new();
        let snap = state.snapshot_values(&["y"]);

        state.set("y", "new");
        let diffs = state.diff_values(&snap, &["y"]);
        assert_eq!(diffs.len(), 1);
        assert_eq!(diffs[0].0, "y");
        assert_eq!(diffs[0].1, Value::Null);
        assert_eq!(diffs[0].2, serde_json::json!("new"));
    }

    #[test]
    fn diff_values_detects_removed_key() {
        let state = State::new();
        state.set("z", 42);
        let snap = state.snapshot_values(&["z"]);

        state.remove("z");
        let diffs = state.diff_values(&snap, &["z"]);
        assert_eq!(diffs.len(), 1);
        assert_eq!(diffs[0].0, "z");
        assert_eq!(diffs[0].1, serde_json::json!(42));
        assert_eq!(diffs[0].2, Value::Null);
    }

    #[test]
    fn diff_values_no_change() {
        let state = State::new();
        state.set("stable", 10);
        let snap = state.snapshot_values(&["stable"]);

        // No mutation
        let diffs = state.diff_values(&snap, &["stable"]);
        assert!(diffs.is_empty());
    }

    #[test]
    fn diff_values_multiple_keys_mixed_changes() {
        let state = State::new();
        state.set("a", 1);
        state.set("b", 2);
        let snap = state.snapshot_values(&["a", "b", "c"]);

        state.set("a", 10); // changed
                            // b unchanged
        state.set("c", 3); // new

        let diffs = state.diff_values(&snap, &["a", "b", "c"]);
        assert_eq!(diffs.len(), 2); // a changed, c new; b unchanged
        let diff_keys: Vec<&str> = diffs.iter().map(|(k, _, _)| k.as_str()).collect();
        assert!(diff_keys.contains(&"a"));
        assert!(diff_keys.contains(&"c"));
    }

    // ── clear_prefix tests ───────────────────────────────────────────────

    #[test]
    fn clear_prefix_removes_matching_keys() {
        let state = State::new();
        state.set("turn:a", 1);
        state.set("turn:b", 2);
        state.set("app:c", 3);
        state.set("session:d", 4);

        state.clear_prefix("turn:");
        assert!(!state.contains("turn:a"));
        assert!(!state.contains("turn:b"));
        assert!(state.contains("app:c"));
        assert!(state.contains("session:d"));
    }

    #[test]
    fn clear_prefix_no_matching_keys_is_noop() {
        let state = State::new();
        state.set("app:x", 1);
        state.clear_prefix("turn:");
        assert!(state.contains("app:x"));
    }

    #[test]
    fn clear_prefix_also_clears_delta() {
        let state = State::new();
        state.set("turn:committed", 1);
        let tracked = state.with_delta_tracking();
        tracked.set("turn:delta_val", 2);

        // Both committed and delta have turn: keys
        assert!(tracked.contains("turn:committed"));
        assert!(tracked.contains("turn:delta_val"));

        tracked.clear_prefix("turn:");
        assert!(!tracked.contains("turn:committed"));
        assert!(!tracked.contains("turn:delta_val"));
    }

    #[test]
    fn clear_prefix_via_turn_accessor() {
        let state = State::new();
        state.turn().set("x", 1);
        state.turn().set("y", 2);
        state.app().set("z", 3);

        state.clear_prefix("turn:");
        assert!(state.turn().keys().is_empty());
        assert!(state.app().contains("z"));
    }

    // ── modify() tests ──────────────────────────────────────────────────

    #[test]
    fn modify_increment_existing() {
        let state = State::new();
        state.set("count", 5u32);
        let result = state.modify("count", 0u32, |n| n + 1);
        assert_eq!(result, 6);
        assert_eq!(state.get::<u32>("count"), Some(6));
    }

    #[test]
    fn modify_uses_default_when_missing() {
        let state = State::new();
        let result = state.modify("new_count", 0u32, |n| n + 1);
        assert_eq!(result, 1);
        assert_eq!(state.get::<u32>("new_count"), Some(1));
    }

    #[test]
    fn modify_with_delta_tracking() {
        let state = State::new();
        state.set("x", 10u32);
        let tracked = state.with_delta_tracking();
        let result = tracked.modify("x", 0u32, |n| n * 2);
        assert_eq!(result, 20);
        // Written to delta, not committed
        assert_eq!(tracked.get::<u32>("x"), Some(20));
        assert_eq!(state.get::<u32>("x"), Some(10)); // original unchanged
    }

    // ── derived fallback tests ──────────────────────────────────────────

    #[test]
    fn get_falls_back_to_derived_prefix() {
        let state = State::new();
        state.set("derived:risk", 0.85);
        // Access without prefix — should find derived:risk
        assert_eq!(state.get::<f64>("risk"), Some(0.85));
    }

    #[test]
    fn get_prefers_direct_key_over_derived() {
        let state = State::new();
        state.set("score", 1.0);
        state.set("derived:score", 0.5);
        // Direct key should win
        assert_eq!(state.get::<f64>("score"), Some(1.0));
    }

    #[test]
    fn get_derived_fallback_skipped_for_prefixed_keys() {
        let state = State::new();
        state.set("derived:risk", 0.85);
        // Prefixed key should NOT trigger fallback
        assert_eq!(state.get::<f64>("app:risk"), None);
    }

    #[test]
    fn get_derived_fallback_with_delta_tracking() {
        let state = State::new();
        let tracked = state.with_delta_tracking();
        tracked.set("derived:computed_val", 42);
        assert_eq!(tracked.get::<i32>("computed_val"), Some(42));
    }

    // ── with() zero-copy borrow tests ──────────────────────────────────

    #[test]
    fn with_reads_from_inner() {
        let state = State::new();
        state.set("name", "Alice");
        let len = state.with("name", |v| v.as_str().unwrap().len());
        assert_eq!(len, Some(5));
    }

    #[test]
    fn with_reads_from_delta_first() {
        let state = State::new();
        state.set("x", 1);
        let tracked = state.with_delta_tracking();
        tracked.set("x", 99);
        let val = tracked.with("x", |v| v.as_i64().unwrap());
        assert_eq!(val, Some(99));
    }

    #[test]
    fn with_falls_back_to_inner_when_not_in_delta() {
        let state = State::new();
        state.set("committed", "yes");
        let tracked = state.with_delta_tracking();
        let val = tracked.with("committed", |v| v.as_str().unwrap().to_string());
        assert_eq!(val, Some("yes".to_string()));
    }

    #[test]
    fn with_falls_back_to_derived() {
        let state = State::new();
        state.set("derived:risk", 0.85);
        let val = state.with("risk", |v| v.as_f64().unwrap());
        assert_eq!(val, Some(0.85));
    }

    #[test]
    fn with_derived_fallback_skipped_for_prefixed() {
        let state = State::new();
        state.set("derived:risk", 0.85);
        let val = state.with("app:risk", |v| v.as_f64().unwrap());
        assert_eq!(val, None);
    }

    #[test]
    fn with_returns_none_for_missing() {
        let state = State::new();
        let val = state.with("missing", |v| v.clone());
        assert_eq!(val, None);
    }

    #[test]
    fn with_on_prefixed_state() {
        let state = State::new();
        state.app().set("flag", true);
        let val = state.app().with("flag", |v| v.as_bool().unwrap());
        assert_eq!(val, Some(true));
    }

    #[test]
    fn with_on_read_only_prefixed_state() {
        let state = State::new();
        state.set("derived:score", serde_json::json!(0.95));
        let val = state.derived().with("score", |v| v.as_f64().unwrap());
        assert_eq!(val, Some(0.95));
    }

    // ── StateKey typed key tests ───────────────────────────────────────

    const TURN_COUNT: StateKey<u32> = StateKey::new("session:turn_count");
    const NAME: StateKey<String> = StateKey::new("user:name");

    #[test]
    fn state_key_get_and_set() {
        let state = State::new();
        state.set_key(&TURN_COUNT, 5);
        assert_eq!(state.get_key(&TURN_COUNT), Some(5));
    }

    #[test]
    fn state_key_get_missing() {
        let state = State::new();
        assert_eq!(state.get_key(&TURN_COUNT), None);
    }

    #[test]
    fn state_key_string_type() {
        let state = State::new();
        state.set_key(&NAME, "Alice".to_string());
        assert_eq!(state.get_key(&NAME), Some("Alice".to_string()));
    }

    #[test]
    fn state_key_with() {
        let state = State::new();
        state.set_key(&TURN_COUNT, 42);
        let val = state.with_key(&TURN_COUNT, |v| v.as_u64().unwrap());
        assert_eq!(val, Some(42));
    }

    #[test]
    fn state_key_interop_with_raw() {
        let state = State::new();
        state.set_key(&TURN_COUNT, 10);
        // Can also read via raw key
        assert_eq!(state.get::<u32>("session:turn_count"), Some(10));
    }
}