weavegraph 0.7.0

//! Versioned state management for workflow execution.
//!
//! State is organized into three independent channels — messages, extras, and errors —
//! each carrying its own version number for change detection and optimistic concurrency.
//!
//! The main types are [`VersionedState`] (mutable runtime state) and [`StateSnapshot`]
//! (point-in-time read-only view passed to nodes during execution).

use rustc_hash::FxHashMap;
use serde::{Serialize, de::DeserializeOwned};
use serde_json::Value;
use std::hash::{Hash, Hasher};
use std::marker::PhantomData;
use thiserror::Error;

use crate::{
    channels::{Channel, ErrorsChannel, ExtrasChannel, MessagesChannel},
    message::{Message, Role},
};

/// Lifecycle classification for a state slot.
///
/// Lifecycle is **metadata** — it does not affect the storage key or identity comparisons.
/// Two `StateKey` values with the same `(namespace, name, schema_version)` but different
/// lifecycle annotations refer to the same slot and compare as equal.
///
/// Consumed by [`StateNormalizeProfile`](crate::runtimes::replay::StateNormalizeProfile)
/// and [`NodePartial::clear_typed_extra_key`](crate::node::NodePartial::clear_typed_extra_key)
/// to distinguish durable state from per-invocation scratch values.
///
/// # Registration-time conflict detection
///
/// When you register a key with a lifecycle annotation, the profile detects and panics on
/// conflicting annotations for the same storage key.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum StateLifecycle {
    /// Persists across invocations. This is the default.
    Durable,
    /// Per-invocation scratch data excluded from durable comparisons and resume normalization.
    InvocationScoped,
}

/// Schema-versioned key for typed values stored in [`VersionedState::extra`].
///
/// Domain crates define constants using `StateKey` so nodes, reducers, tests, and replay
/// code can reference typed slots without repeating string literals.
///
/// # Equality and hashing
///
/// `PartialEq`, `Eq`, and `Hash` are based solely on `(namespace, name, schema_version)`.
/// The `lifecycle` field is excluded so two keys for the same slot compare equal regardless
/// of lifecycle annotation.
///
/// ```rust
/// use serde::{Deserialize, Serialize};
/// use weavegraph::state::{StateKey, StateLifecycle};
///
/// #[derive(Serialize, Deserialize)]
/// struct PortfolioSnapshot { cash: i64 }
///
/// const PORTFOLIO: StateKey<PortfolioSnapshot> = StateKey::new("wq", "portfolio_snapshot", 1);
/// const CURRENT_EVENT: StateKey<u64> = StateKey::new("wq", "event", 1).invocation_scoped();
///
/// assert_eq!(PORTFOLIO.storage_key(), "wq:portfolio_snapshot:v1");
/// assert_eq!(CURRENT_EVENT.lifecycle(), StateLifecycle::InvocationScoped);
/// // Same slot identity regardless of lifecycle annotation:
/// assert_eq!(
///     StateKey::<u64>::new("wq", "event", 1),
///     StateKey::<u64>::new("wq", "event", 1).invocation_scoped(),
/// );
/// ```
#[derive(Debug)]
pub struct StateKey<T> {
    namespace: &'static str,
    name: &'static str,
    schema_version: u32,
    lifecycle: StateLifecycle,
    _marker: PhantomData<fn() -> T>,
}

impl<T> Clone for StateKey<T> {
    fn clone(&self) -> Self {
        *self
    }
}

impl<T> Copy for StateKey<T> {}

// Equality and Hash intentionally exclude `lifecycle` — it is metadata,
// not identity. Two keys for the same slot are the same key.
impl<T> PartialEq for StateKey<T> {
    fn eq(&self, other: &Self) -> bool {
        self.namespace == other.namespace
            && self.name == other.name
            && self.schema_version == other.schema_version
    }
}

impl<T> Eq for StateKey<T> {}

impl<T> Hash for StateKey<T> {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.namespace.hash(state);
        self.name.hash(state);
        self.schema_version.hash(state);
    }
}

impl<T> StateKey<T> {
    /// Create a typed state key with [`StateLifecycle::Durable`] (default).
    pub const fn new(namespace: &'static str, name: &'static str, schema_version: u32) -> Self {
        Self {
            namespace,
            name,
            schema_version,
            lifecycle: StateLifecycle::Durable,
            _marker: PhantomData,
        }
    }

    /// Return a copy annotated as [`StateLifecycle::InvocationScoped`].
    ///
    /// The returned key compares equal to the original. Use this when defining constants
    /// for per-invocation scratch slots so normalization and cleanup helpers can distinguish
    /// them from durable state.
    ///
    /// ```rust
    /// use weavegraph::state::{StateKey, StateLifecycle};
    ///
    /// const TICK: StateKey<u64> = StateKey::new("wq", "tick", 1).invocation_scoped();
    /// assert_eq!(TICK.lifecycle(), StateLifecycle::InvocationScoped);
    /// ```
    #[must_use]
    pub const fn invocation_scoped(mut self) -> Self {
        self.lifecycle = StateLifecycle::InvocationScoped;
        self
    }

    /// Return the lifecycle classification of this key.
    #[must_use]
    pub fn lifecycle(&self) -> StateLifecycle {
        self.lifecycle
    }

    /// Return the namespace component.
    #[must_use]
    pub fn namespace(&self) -> &'static str {
        self.namespace
    }

    /// Return the key name component.
    #[must_use]
    pub fn name(&self) -> &'static str {
        self.name
    }

    /// Return the schema version component.
    #[must_use]
    pub fn schema_version(&self) -> u32 {
        self.schema_version
    }

    /// Return the `extra` map key used for storage: `namespace:name:v{schema_version}`.
    ///
    /// Bumping the schema version writes to a new slot, preventing silent collisions
    /// between incompatible payload shapes.
    #[must_use]
    pub fn storage_key(&self) -> String {
        format!("{}:{}:v{}", self.namespace, self.name, self.schema_version)
    }
}

/// Errors produced by typed state-slot helpers.
#[derive(Debug, Error)]
#[cfg_attr(feature = "diagnostics", derive(miette::Diagnostic))]
#[non_exhaustive]
pub enum StateSlotError {
    /// The requested slot was absent.
    #[error("state slot not found: {key}")]
    #[cfg_attr(
        feature = "diagnostics",
        diagnostic(code(weavegraph::state::slot_missing))
    )]
    Missing {
        /// Concrete storage key that was not found.
        key: String,
    },

    /// A slot value could not be serialized to JSON.
    #[error("failed to serialize state slot {key}: {source}")]
    #[cfg_attr(
        feature = "diagnostics",
        diagnostic(code(weavegraph::state::slot_serialize))
    )]
    Serialize {
        /// Concrete storage key being written.
        key: String,
        /// Underlying serde serialization error.
        #[source]
        source: serde_json::Error,
    },

    /// A slot value could not be deserialized from JSON.
    #[error("failed to deserialize state slot {key}: {source}")]
    #[cfg_attr(
        feature = "diagnostics",
        diagnostic(code(weavegraph::state::slot_deserialize))
    )]
    Deserialize {
        /// Concrete storage key being read.
        key: String,
        /// Underlying serde deserialization error.
        #[source]
        source: serde_json::Error,
    },
}

/// Runtime state container for workflow execution.
///
/// Manages three independent channels of versioned data: messages, custom extras, and
/// error events. Each channel tracks its own version number for change detection.
///
/// ```rust
/// use weavegraph::state::VersionedState;
/// use weavegraph::channels::Channel;
/// use serde_json::json;
///
/// let mut state = VersionedState::new_with_user_message("Process this");
/// state.add_extra("session_id", json!("sess_123"));
///
/// let snapshot = state.snapshot();
/// assert_eq!(snapshot.messages.len(), 1);
/// assert_eq!(snapshot.extra.get("session_id"), Some(&json!("sess_123")));
/// ```
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct VersionedState {
    /// Conversation messages.
    pub messages: MessagesChannel,
    /// Custom metadata and intermediate results.
    pub extra: ExtrasChannel,
    /// Error events and diagnostics.
    pub errors: ErrorsChannel,
}

/// Immutable point-in-time view of workflow state.
///
/// Created by [`VersionedState::snapshot`] and passed to nodes during execution.
/// Contains cloned channel data and version numbers at the moment of the snapshot.
///
/// ```rust
/// use weavegraph::state::VersionedState;
/// use weavegraph::channels::Channel;
/// use serde_json::json;
///
/// let mut state = VersionedState::new_with_user_message("Hello");
/// state.extra.get_mut().insert("key".to_string(), json!("value"));
///
/// let snapshot = state.snapshot();
/// state.extra.get_mut().clear();
///
/// assert_eq!(snapshot.extra.get("key"), Some(&json!("value")));
/// assert!(state.extra.snapshot().is_empty());
/// ```
#[derive(Clone, Debug)]
pub struct StateSnapshot {
    /// Messages at snapshot time.
    pub messages: Vec<Message>,
    /// Version of the messages channel when snapshotted.
    pub messages_version: u32,
    /// Extra data at snapshot time.
    pub extra: FxHashMap<String, Value>,
    /// Version of the extra channel when snapshotted.
    pub extra_version: u32,
    /// Error events at snapshot time.
    pub errors: Vec<crate::channels::errors::ErrorEvent>,
    /// Version of the errors channel when snapshotted.
    pub errors_version: u32,
}

impl VersionedState {
    /// Construct state initialized with a single user message.
    ///
    /// ```rust
    /// use weavegraph::state::VersionedState;
    ///
    /// let state = VersionedState::new_with_user_message("Analyze this");
    /// let snap = state.snapshot();
    /// assert_eq!(snap.messages.len(), 1);
    /// assert_eq!(snap.messages[0].role, weavegraph::message::Role::User);
    /// ```
    pub fn new_with_user_message(user_text: &str) -> Self {
        Self {
            messages: MessagesChannel::new(vec![Message::with_role(Role::User, user_text)], 1),
            extra: ExtrasChannel::default(),
            errors: ErrorsChannel::default(),
        }
    }

    /// Construct state initialized with an existing message list.
    ///
    /// ```rust
    /// use weavegraph::state::VersionedState;
    /// use weavegraph::message::{Message, Role};
    ///
    /// let messages = vec![
    ///     Message::with_role(Role::User, "Explain error handling in Rust"),
    ///     Message::with_role(Role::Assistant, "Use Result and the ? operator."),
    /// ];
    /// let state = VersionedState::new_with_messages(messages);
    /// assert_eq!(state.snapshot().messages.len(), 2);
    /// ```
    pub fn new_with_messages(messages: Vec<Message>) -> Self {
        Self {
            messages: MessagesChannel::new(messages, 1),
            extra: ExtrasChannel::default(),
            errors: ErrorsChannel::default(),
        }
    }

    /// Return a builder for fluent state construction.
    pub fn builder() -> VersionedStateBuilder {
        VersionedStateBuilder::default()
    }

    /// Append a message to the messages channel.
    pub fn add_message(&mut self, role: &str, content: &str) -> &mut Self {
        self.messages
            .get_mut()
            .push(Message::with_role(Role::from(role), content));
        self
    }

    /// Insert a key-value pair into the extra channel.
    pub fn add_extra(&mut self, key: &str, value: Value) -> &mut Self {
        self.extra.get_mut().insert(key.to_owned(), value);
        self
    }

    /// Serialize `value` and insert it under `key.storage_key()` in the extra channel.
    pub fn add_typed_extra<T: Serialize>(
        &mut self,
        key: StateKey<T>,
        value: T,
    ) -> Result<&mut Self, StateSlotError> {
        let storage_key = key.storage_key();
        let json_value =
            serde_json::to_value(value).map_err(|source| StateSlotError::Serialize {
                key: storage_key.clone(),
                source,
            })?;
        self.extra.get_mut().insert(storage_key, json_value);
        Ok(self)
    }

    /// Clone all channel data and version numbers into an immutable [`StateSnapshot`].
    ///
    /// ```rust
    /// use weavegraph::state::VersionedState;
    /// use weavegraph::channels::Channel;
    /// use serde_json::json;
    ///
    /// let mut state = VersionedState::new_with_user_message("Test");
    /// state.extra.get_mut().insert("status".to_string(), json!("processing"));
    /// let snapshot = state.snapshot();
    /// state.extra.get_mut().insert("status".to_string(), json!("complete"));
    ///
    /// assert_eq!(snapshot.extra.get("status"), Some(&json!("processing")));
    /// ```
    pub fn snapshot(&self) -> StateSnapshot {
        StateSnapshot {
            messages: self.messages.snapshot(),
            messages_version: self.messages.version(),
            extra: self.extra.snapshot(),
            extra_version: self.extra.version(),
            errors: self.errors.snapshot(),
            errors_version: self.errors.version(),
        }
    }
}

impl StateSnapshot {
    /// Deserialize an optional typed value from the extra channel.
    ///
    /// Returns `Ok(None)` when the slot is absent.
    pub fn get_typed<T: DeserializeOwned>(
        &self,
        key: StateKey<T>,
    ) -> Result<Option<T>, StateSlotError> {
        let storage_key = key.storage_key();
        self.extra
            .get(&storage_key)
            .cloned()
            .map(|v| {
                serde_json::from_value(v).map_err(|source| StateSlotError::Deserialize {
                    key: storage_key,
                    source,
                })
            })
            .transpose()
    }

    /// Deserialize a required typed value from the extra channel.
    ///
    /// Returns [`StateSlotError::Missing`] when the slot is absent.
    pub fn require_typed<T: DeserializeOwned>(
        &self,
        key: StateKey<T>,
    ) -> Result<T, StateSlotError> {
        let storage_key = key.storage_key();
        self.get_typed(key)?
            .ok_or(StateSlotError::Missing { key: storage_key })
    }
}

/// Fluent builder for [`VersionedState`].
///
/// ```rust
/// use weavegraph::state::VersionedState;
/// use weavegraph::channels::Channel;
/// use serde_json::json;
///
/// let state = VersionedState::builder()
///     .with_user_message("What's the weather like?")
///     .with_assistant_message("I'll help you check the weather.")
///     .with_extra("location", json!("New York"))
///     .build();
///
/// assert_eq!(state.snapshot().messages.len(), 2);
/// ```
#[derive(Debug, Default)]
pub struct VersionedStateBuilder {
    messages: Vec<Message>,
    extra: FxHashMap<String, Value>,
}

impl VersionedStateBuilder {
    /// Append a user message.
    pub fn with_user_message(mut self, content: &str) -> Self {
        self.messages.push(Message::with_role(Role::User, content));
        self
    }

    /// Append an assistant message.
    pub fn with_assistant_message(mut self, content: &str) -> Self {
        self.messages
            .push(Message::with_role(Role::Assistant, content));
        self
    }

    /// Append a system message.
    pub fn with_system_message(mut self, content: &str) -> Self {
        self.messages
            .push(Message::with_role(Role::System, content));
        self
    }

    /// Append a message with a custom role.
    pub fn with_message(mut self, role: &str, content: &str) -> Self {
        self.messages
            .push(Message::with_role(Role::from(role), content));
        self
    }

    /// Insert a key-value pair into the extra channel.
    pub fn with_extra(mut self, key: &str, value: Value) -> Self {
        self.extra.insert(key.to_owned(), value);
        self
    }

    /// Serialize `value` and insert it under `key.storage_key()`.
    pub fn with_typed_extra<T: Serialize>(
        mut self,
        key: StateKey<T>,
        value: T,
    ) -> Result<Self, StateSlotError> {
        let storage_key = key.storage_key();
        let json_value =
            serde_json::to_value(value).map_err(|source| StateSlotError::Serialize {
                key: storage_key.clone(),
                source,
            })?;
        self.extra.insert(storage_key, json_value);
        Ok(self)
    }

    /// Construct the [`VersionedState`]. All channels start at version 1.
    pub fn build(self) -> VersionedState {
        VersionedState {
            messages: MessagesChannel::new(self.messages, 1),
            extra: ExtrasChannel::new(self.extra, 1),
            errors: ErrorsChannel::default(),
        }
    }
}