cellstate_core/

session.rs

//! Session typestate for compile-time safety of stateful LLM session lifecycle.
//!
//! Uses the typestate pattern to make invalid state transitions uncompilable.
//! Sessions track stateful LLM connections where only deltas (new messages)
//! are sent after the initial context, avoiding the "Groundhog Day" problem
//! of re-sending full context every tool call.
//!
//! # State Transition Diagram
//!
//! ```text
//! Session::new() → Created ── activate() ──→ Active ── close() ──→ Closed (terminal)
//!                                              │
//!                                         record_delta() ↺
//!                                              │
//!                                         expire() ──→ Expired (terminal)
//! ```

use crate::{EnumParseError, SessionId, TenantId, Timestamp};
use serde::{Deserialize, Serialize};
use std::fmt;
use std::marker::PhantomData;
use std::str::FromStr;

// ============================================================================
// SESSION STATUS ENUM
// ============================================================================

/// Status of a stateful LLM session.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
pub enum SessionStatus {
    /// Session has been created but not yet activated
    Created,
    /// Session is active and accepting deltas
    Active,
    /// Session was closed normally
    Closed,
    /// Session expired due to TTL
    Expired,
}

impl SessionStatus {
    /// Convert to database string representation.
    pub fn as_db_str(&self) -> &'static str {
        match self {
            SessionStatus::Created => "created",
            SessionStatus::Active => "active",
            SessionStatus::Closed => "closed",
            SessionStatus::Expired => "expired",
        }
    }

    /// Parse from database string representation.
    pub fn from_db_str(s: &str) -> Result<Self, EnumParseError> {
        match s.to_lowercase().as_str() {
            "created" => Ok(SessionStatus::Created),
            "active" => Ok(SessionStatus::Active),
            "closed" => Ok(SessionStatus::Closed),
            "expired" => Ok(SessionStatus::Expired),
            _ => Err(EnumParseError {
                enum_name: "SessionStatus",
                input: s.to_string(),
            }),
        }
    }

    /// Check if this is a terminal state (no further transitions possible).
    pub fn is_terminal(&self) -> bool {
        matches!(self, SessionStatus::Closed | SessionStatus::Expired)
    }
}

impl fmt::Display for SessionStatus {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.as_db_str())
    }
}

impl FromStr for SessionStatus {
    type Err = EnumParseError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        Self::from_db_str(s)
    }
}

// ============================================================================
// SESSION DATA (internal storage, state-independent)
// ============================================================================

/// Internal data storage for a session, independent of typestate.
/// This is what gets persisted to the database.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
pub struct SessionRecord {
    #[cfg_attr(feature = "openapi", schema(value_type = String, format = "uuid"))]
    pub session_id: SessionId,
    #[cfg_attr(feature = "openapi", schema(value_type = String, format = "uuid"))]
    pub tenant_id: TenantId,
    /// Provider that owns this session
    pub provider_id: String,
    /// Model locked to this session
    pub model: String,
    #[cfg_attr(feature = "openapi", schema(value_type = String, format = "date-time"))]
    pub created_at: Timestamp,
    #[cfg_attr(feature = "openapi", schema(value_type = Option<String>, format = "date-time"))]
    pub activated_at: Option<Timestamp>,
    #[cfg_attr(feature = "openapi", schema(value_type = Option<String>, format = "date-time"))]
    pub closed_at: Option<Timestamp>,
    /// Time-to-live in seconds for this session
    pub ttl_secs: u64,
    /// Number of tool loop rounds completed
    pub round_count: u64,
    /// Token count of the initial context (first request)
    pub initial_token_count: u64,
    /// Cumulative token count of all deltas sent
    pub delta_token_count: u64,
    /// Arbitrary metadata
    #[serde(default)]
    pub metadata: serde_json::Value,
}

// ============================================================================
// TYPESTATE MARKERS
// ============================================================================

/// Marker trait for session states.
pub trait SessionState: private::Sealed + Send + Sync {}

/// Session has been created but not yet activated.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Created;
impl SessionState for Created {}

/// Session is active and accepting deltas.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Active;
impl SessionState for Active {}

/// Session was closed normally (terminal).
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Closed;
impl SessionState for Closed {}

/// Session expired due to TTL (terminal).
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct SessionExpired;
impl SessionState for SessionExpired {}

mod private {
    pub trait Sealed {}
    impl Sealed for super::Created {}
    impl Sealed for super::Active {}
    impl Sealed for super::Closed {}
    impl Sealed for super::SessionExpired {}
}

// ============================================================================
// SESSION TYPESTATE WRAPPER
// ============================================================================

/// A stateful LLM session with compile-time state tracking.
///
/// The type parameter `S` indicates the current state of the session.
/// Methods are only available in appropriate states:
/// - `Session<Created>`: Can be activated
/// - `Session<Active>`: Can record deltas, close, or expire
/// - `Session<Closed>`: Terminal, no further transitions
/// - `Session<SessionExpired>`: Terminal, no further transitions
#[derive(Debug, Clone)]
pub struct Session<S: SessionState> {
    data: SessionRecord,
    _state: PhantomData<S>,
}

impl<S: SessionState> Session<S> {
    /// Access the underlying session data (read-only).
    pub fn data(&self) -> &SessionRecord {
        &self.data
    }

    /// Get the session ID.
    pub fn session_id(&self) -> SessionId {
        self.data.session_id
    }

    /// Get the tenant ID.
    pub fn tenant_id(&self) -> TenantId {
        self.data.tenant_id
    }

    /// Get the provider ID.
    pub fn provider_id(&self) -> &str {
        &self.data.provider_id
    }

    /// Get the model.
    pub fn model(&self) -> &str {
        &self.data.model
    }

    /// Get when the session was created.
    pub fn created_at(&self) -> Timestamp {
        self.data.created_at
    }

    /// Get the round count.
    pub fn round_count(&self) -> u64 {
        self.data.round_count
    }

    /// Consume and return the underlying data (for serialization).
    pub fn into_data(self) -> SessionRecord {
        self.data
    }
}

impl Session<Created> {
    /// Create a new session in the Created state.
    pub fn new(data: SessionRecord) -> Self {
        Session {
            data,
            _state: PhantomData,
        }
    }

    /// Activate the session.
    ///
    /// Transitions to `Session<Active>`.
    /// Consumes the current session.
    pub fn activate(mut self, activated_at: Timestamp) -> Session<Active> {
        self.data.activated_at = Some(activated_at);
        Session {
            data: self.data,
            _state: PhantomData,
        }
    }
}

impl Session<Active> {
    /// Get when the session was activated.
    pub fn activated_at(&self) -> Timestamp {
        self.data
            .activated_at
            .expect("Active session must have activated_at")
    }

    /// Record a delta (new messages sent to the provider).
    ///
    /// This is a looping state — it does not consume self.
    /// Increments round_count and accumulates delta_tokens.
    pub fn record_delta(&mut self, delta_tokens: u64) {
        self.data.round_count += 1;
        self.data.delta_token_count += delta_tokens;
    }

    /// Calculate total tokens saved by using stateful sessions.
    ///
    /// Without sessions, each round would re-send the full initial context.
    /// Savings = (round_count * initial_tokens) - delta_tokens
    pub fn tokens_saved(&self) -> u64 {
        let would_have_sent = self.data.round_count * self.data.initial_token_count;
        would_have_sent.saturating_sub(self.data.delta_token_count)
    }

    /// Close the session normally.
    ///
    /// Transitions to `Session<Closed>` (terminal state).
    /// Consumes the current session.
    pub fn close(mut self, closed_at: Timestamp) -> Session<Closed> {
        self.data.closed_at = Some(closed_at);
        Session {
            data: self.data,
            _state: PhantomData,
        }
    }

    /// Expire the session due to TTL.
    ///
    /// Transitions to `Session<SessionExpired>` (terminal state).
    /// Consumes the current session.
    pub fn expire(mut self, expired_at: Timestamp) -> Session<SessionExpired> {
        self.data.closed_at = Some(expired_at);
        Session {
            data: self.data,
            _state: PhantomData,
        }
    }
}

impl Session<Closed> {
    /// Get when the session was closed.
    pub fn closed_at(&self) -> Timestamp {
        self.data
            .closed_at
            .expect("Closed session must have closed_at")
    }
}

impl Session<SessionExpired> {
    /// Get when the session expired.
    pub fn expired_at(&self) -> Timestamp {
        self.data
            .closed_at
            .expect("Expired session must have closed_at")
    }
}

// ============================================================================
// DATABASE BOUNDARY: STORED SESSION
// ============================================================================

/// A session as stored in the database (status-agnostic).
///
/// When loading from the database, we don't know the state at compile time.
/// Use the `into_typed` method to validate and convert to a typed session.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct StoredSession {
    pub data: SessionRecord,
    pub status: SessionStatus,
}

/// Enum representing all possible runtime states of a session.
/// Use this when you need to handle sessions loaded from the database.
#[derive(Debug, Clone)]
pub enum LoadedSession {
    Created(Session<Created>),
    Active(Session<Active>),
    Closed(Session<Closed>),
    Expired(Session<SessionExpired>),
}

impl StoredSession {
    /// Convert to a typed session based on the stored status.
    pub fn into_typed(self) -> LoadedSession {
        match self.status {
            SessionStatus::Created => LoadedSession::Created(Session {
                data: self.data,
                _state: PhantomData,
            }),
            SessionStatus::Active => LoadedSession::Active(Session {
                data: self.data,
                _state: PhantomData,
            }),
            SessionStatus::Closed => LoadedSession::Closed(Session {
                data: self.data,
                _state: PhantomData,
            }),
            SessionStatus::Expired => LoadedSession::Expired(Session {
                data: self.data,
                _state: PhantomData,
            }),
        }
    }

    /// Try to convert to a created session.
    pub fn into_created(self) -> Result<Session<Created>, SessionStateError> {
        if self.status != SessionStatus::Created {
            return Err(SessionStateError::WrongState {
                session_id: self.data.session_id,
                expected: SessionStatus::Created,
                actual: self.status,
            });
        }
        Ok(Session {
            data: self.data,
            _state: PhantomData,
        })
    }

    /// Try to convert to an active session.
    pub fn into_active(self) -> Result<Session<Active>, SessionStateError> {
        if self.status != SessionStatus::Active {
            return Err(SessionStateError::WrongState {
                session_id: self.data.session_id,
                expected: SessionStatus::Active,
                actual: self.status,
            });
        }
        Ok(Session {
            data: self.data,
            _state: PhantomData,
        })
    }

    /// Get the underlying data without state validation.
    pub fn data(&self) -> &SessionRecord {
        &self.data
    }

    /// Get the current status.
    pub fn status(&self) -> SessionStatus {
        self.status
    }
}

/// Errors when transitioning session states.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum SessionStateError {
    /// Session is not in the expected state.
    WrongState {
        session_id: SessionId,
        expected: SessionStatus,
        actual: SessionStatus,
    },
}

impl fmt::Display for SessionStateError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            SessionStateError::WrongState {
                session_id,
                expected,
                actual,
            } => {
                write!(
                    f,
                    "Session {} is in state {} but expected {}",
                    session_id, actual, expected
                )
            }
        }
    }
}

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

#[cfg(test)]
mod tests {
    use super::*;
    use crate::EntityIdType;
    use chrono::Utc;

    fn make_session_data() -> SessionRecord {
        let now = Utc::now();
        SessionRecord {
            session_id: SessionId::now_v7(),
            tenant_id: TenantId::now_v7(),
            provider_id: "openai".to_string(),
            model: "gpt-4".to_string(),
            created_at: now,
            activated_at: None,
            closed_at: None,
            ttl_secs: 300,
            round_count: 0,
            initial_token_count: 5000,
            delta_token_count: 0,
            metadata: serde_json::Value::Object(serde_json::Map::new()),
        }
    }

    #[test]
    fn test_session_status_roundtrip() {
        for status in [
            SessionStatus::Created,
            SessionStatus::Active,
            SessionStatus::Closed,
            SessionStatus::Expired,
        ] {
            let db_str = status.as_db_str();
            let parsed =
                SessionStatus::from_db_str(db_str).expect("SessionStatus roundtrip should succeed");
            assert_eq!(status, parsed);
        }
    }

    #[test]
    fn test_session_status_terminal() {
        assert!(!SessionStatus::Created.is_terminal());
        assert!(!SessionStatus::Active.is_terminal());
        assert!(SessionStatus::Closed.is_terminal());
        assert!(SessionStatus::Expired.is_terminal());
    }

    #[test]
    fn test_session_create_activate_delta_close() {
        let now = Utc::now();
        let data = make_session_data();
        let session = Session::<Created>::new(data);

        // Activate
        let mut active = session.activate(now);
        assert_eq!(active.activated_at(), now);
        assert_eq!(active.round_count(), 0);

        // Record 3 deltas
        active.record_delta(100);
        active.record_delta(120);
        active.record_delta(150);

        assert_eq!(active.round_count(), 3);
        assert_eq!(active.data().delta_token_count, 370);

        // Verify tokens saved: 3 rounds * 5000 initial = 15000; deltas = 370; saved = 14630
        assert_eq!(active.tokens_saved(), 14630);

        // Close
        let closed = active.close(now);
        assert_eq!(closed.closed_at(), now);
    }

    #[test]
    fn test_session_create_activate_expire() {
        let now = Utc::now();
        let data = make_session_data();
        let session = Session::<Created>::new(data);

        let active = session.activate(now);
        let expired = active.expire(now);

        // Terminal — verify expired_at is set
        assert_eq!(expired.expired_at(), now);
    }

    #[test]
    fn test_tokens_saved_calculation() {
        let now = Utc::now();
        let data = make_session_data();
        let session = Session::<Created>::new(data);
        let mut active = session.activate(now);

        // 3 rounds × 5000 initial = 15000 would-have-sent
        // deltas = 370 → saved = 14630
        active.record_delta(100);
        active.record_delta(120);
        active.record_delta(150);

        assert_eq!(active.tokens_saved(), 14630);
    }

    #[test]
    fn test_tokens_saved_zero_rounds() {
        let now = Utc::now();
        let data = make_session_data();
        let session = Session::<Created>::new(data);
        let active = session.activate(now);

        // 0 rounds → 0 saved
        assert_eq!(active.tokens_saved(), 0);
    }

    #[test]
    fn test_stored_session_into_typed() {
        let data = make_session_data();

        // Created
        let stored = StoredSession {
            data: data.clone(),
            status: SessionStatus::Created,
        };
        assert!(matches!(stored.into_typed(), LoadedSession::Created(_)));

        // Active
        let stored = StoredSession {
            data: data.clone(),
            status: SessionStatus::Active,
        };
        assert!(matches!(stored.into_typed(), LoadedSession::Active(_)));

        // Closed
        let stored = StoredSession {
            data: data.clone(),
            status: SessionStatus::Closed,
        };
        assert!(matches!(stored.into_typed(), LoadedSession::Closed(_)));

        // Expired
        let stored = StoredSession {
            data,
            status: SessionStatus::Expired,
        };
        assert!(matches!(stored.into_typed(), LoadedSession::Expired(_)));
    }

    #[test]
    fn test_stored_session_wrong_state() {
        let data = make_session_data();
        let stored = StoredSession {
            data,
            status: SessionStatus::Active,
        };

        assert!(matches!(
            stored.into_created(),
            Err(SessionStateError::WrongState { .. })
        ));
    }

    #[test]
    fn test_stored_session_into_active() {
        let data = make_session_data();
        let stored = StoredSession {
            data: data.clone(),
            status: SessionStatus::Active,
        };

        let active = stored.into_active().expect("should convert to active");
        assert_eq!(active.session_id(), data.session_id);
    }

    #[test]
    fn test_session_status_display_fromstr() {
        for status in [
            SessionStatus::Created,
            SessionStatus::Active,
            SessionStatus::Closed,
            SessionStatus::Expired,
        ] {
            let s = status.to_string();
            let parsed: SessionStatus = s.parse().expect("should parse from display string");
            assert_eq!(status, parsed);
        }
    }
}