loop_agent_sdk/

state.rs

//! Loop Agent SDK - State Interface
//! 
//! Context persistence for stateless Lambda execution.
//! Handles "memory" across invocations in <100ms reload time.

use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// User context - reloaded at start of each Lambda invocation
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct UserContext {
    /// User's Solana pubkey
    pub pubkey: String,
    
    /// Current vault state (cached, refresh on demand)
    pub vault: VaultState,
    
    /// User preferences
    pub preferences: UserPreferences,
    
    /// Recent activity (last 10 events)
    pub recent_events: Vec<ContextEvent>,
    
    /// Agent-specific memory (key-value)
    pub agent_memory: HashMap<String, serde_json::Value>,
    
    /// Last context update timestamp
    pub updated_at: i64,
    
    /// Context version for optimistic locking
    pub version: u64,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct VaultState {
    pub cred_balance: u64,
    pub oxo_balance: u64,
    pub total_staked: u64,
    pub pending_yield: u64,
    pub active_positions: u16,
    /// Cached at
    pub snapshot_at: i64,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct UserPreferences {
    /// Auto-stake captured Cred
    pub auto_stake: bool,
    /// Preferred stake duration
    pub default_duration_days: u16,
    /// Auto-compound yield
    pub auto_compound: bool,
    /// Yield threshold for auto-compound (lamports)
    pub compound_threshold: u64,
    /// Risk tolerance for recommendations
    pub risk_tolerance: String,
    /// Notification preferences
    pub notifications: NotificationPrefs,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NotificationPrefs {
    pub on_capture: bool,
    pub on_stake: bool,
    pub on_yield: bool,
    pub on_unlock: bool,
    /// Minimum amount to notify (in Cred lamports)
    pub min_notify_amount: u64,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ContextEvent {
    pub event_type: String,
    pub timestamp: i64,
    pub summary: String,
    pub amount: Option<u64>,
}

impl Default for UserPreferences {
    fn default() -> Self {
        Self {
            auto_stake: true,
            default_duration_days: 90,
            auto_compound: true,
            compound_threshold: 10_000_000_000, // 10 Cred
            risk_tolerance: "balanced".to_string(),
            notifications: NotificationPrefs {
                on_capture: true,
                on_stake: false,
                on_yield: true,
                on_unlock: true,
                min_notify_amount: 1_000_000_000, // 1 Cred
            },
        }
    }
}

/// State storage backend trait
/// Can be backed by DynamoDB, Redis, or Solana account data
pub trait StateStore {
    /// Fetch user context (<100ms target)
    fn context_fetch(&self, user_pubkey: &str) -> Result<UserContext, StateError>;
    
    /// Push updated context (with optimistic locking)
    fn context_push(&self, context: UserContext) -> Result<(), StateError>;
    
    /// Fetch multiple users (for batch operations)
    fn context_batch_fetch(&self, user_pubkeys: &[String]) -> Result<Vec<UserContext>, StateError>;
    
    /// Update specific field without full fetch/push
    fn context_patch(
        &self,
        user_pubkey: &str,
        path: &str,
        value: serde_json::Value,
    ) -> Result<(), StateError>;
    
    /// Add event to recent events (ring buffer)
    fn push_event(&self, user_pubkey: &str, event: ContextEvent) -> Result<(), StateError>;
    
    /// Get or initialize context for new user
    fn get_or_init(&self, user_pubkey: &str) -> Result<UserContext, StateError>;
}

#[derive(Debug, Clone)]
pub enum StateError {
    /// User context not found
    NotFound,
    /// Version conflict (optimistic lock failed)
    VersionConflict { expected: u64, actual: u64 },
    /// Storage backend error
    StorageError(String),
    /// Serialization error
    SerializationError(String),
    /// Timeout fetching state
    Timeout,
}

/// DynamoDB-backed state store configuration
#[derive(Debug, Clone)]
pub struct DynamoStateConfig {
    pub table_name: String,
    pub region: String,
    /// TTL for cached vault state (seconds)
    pub vault_cache_ttl: u32,
    /// Max recent events to keep
    pub max_events: usize,
}

impl Default for DynamoStateConfig {
    fn default() -> Self {
        Self {
            table_name: "loop-agent-state".to_string(),
            region: "us-east-1".to_string(),
            vault_cache_ttl: 60,  // Refresh vault every 60s
            max_events: 10,
        }
    }
}

/// DynamoDB table schema
/// 
/// Primary Key: pk = "USER#{pubkey}"
/// Sort Key: sk = "CONTEXT"
/// 
/// GSI1: For batch queries by last update
/// GSI1PK: "ACTIVE"
/// GSI1SK: updated_at
/// 
/// Example item:
/// ```json
/// {
///   "pk": "USER#5A7E...YKTZ",
///   "sk": "CONTEXT",
///   "pubkey": "5A7E...YKTZ",
///   "vault": { ... },
///   "preferences": { ... },
///   "recent_events": [ ... ],
///   "agent_memory": { ... },
///   "updated_at": 1711468800,
///   "version": 42
/// }
/// ```

/// Helper to create initial context for new user
pub fn init_user_context(pubkey: &str) -> UserContext {
    UserContext {
        pubkey: pubkey.to_string(),
        vault: VaultState {
            cred_balance: 0,
            oxo_balance: 0,
            total_staked: 0,
            pending_yield: 0,
            active_positions: 0,
            snapshot_at: 0,
        },
        preferences: UserPreferences::default(),
        recent_events: Vec::new(),
        agent_memory: HashMap::new(),
        updated_at: chrono::Utc::now().timestamp(),
        version: 1,
    }
}

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

    #[test]
    fn default_preferences_sensible() {
        let prefs = UserPreferences::default();
        assert!(prefs.auto_stake);
        assert!(prefs.auto_compound);
        assert_eq!(prefs.default_duration_days, 90);
    }

    #[test]
    fn context_serializes() {
        let ctx = init_user_context("5A7ETEST");
        let json = serde_json::to_string(&ctx).unwrap();
        let parsed: UserContext = serde_json::from_str(&json).unwrap();
        assert_eq!(parsed.pubkey, "5A7ETEST");
    }
}