af_workflow/

state.rs

//! Branch-scoped run state.
//!
//! Port of the `ctx.state` surface used by nodes (`append` / `set` / `get`).
//! The trait is async because production backs it with Redis; this crate ships
//! an in-memory implementation good for tests, single-process runs, and as the
//! reference semantics (sliding-window `append`, TTL expiry).
//!
//! R6: nodes scope keys by `"<branch_id>."` prefix before calling these — the
//! store itself is prefix-agnostic.

use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};

use async_trait::async_trait;
use serde_json::Value;

/// Branch-scoped durable key/value state with TTL.
#[async_trait]
pub trait State: Send + Sync {
    /// Append `item` to the list at `key`, keeping at most `max_len` entries
    /// (oldest dropped first). Creates the list if absent. `ttl` bounds the
    /// key's lifetime.
    async fn append(&self, key: &str, item: Value, max_len: Option<usize>, ttl: Option<Duration>);

    /// Set `key` to `value` (overwrite), with optional `ttl`.
    async fn set(&self, key: &str, value: Value, ttl: Option<Duration>);

    /// Read `key`, or `None` if absent/expired.
    async fn get(&self, key: &str) -> Option<Value>;
}

/// A namespaced view over a shared state backend. Use one namespace per
/// workflow instance so identical branch/key names cannot collide.
pub struct NamespacedState {
    inner: Arc<dyn State>,
    prefix: String,
}

impl NamespacedState {
    /// View of `inner` with every key prefixed by `namespace`.
    pub fn new(inner: Arc<dyn State>, namespace: impl std::fmt::Display) -> Self {
        Self {
            inner,
            prefix: format!("{namespace}."),
        }
    }

    /// Key prefix.
    pub fn prefix(&self) -> &str {
        &self.prefix
    }

    fn key(&self, key: &str) -> String {
        format!("{}{key}", self.prefix)
    }
}

#[async_trait]
impl State for NamespacedState {
    async fn append(&self, key: &str, item: Value, max_len: Option<usize>, ttl: Option<Duration>) {
        self.inner.append(&self.key(key), item, max_len, ttl).await;
    }

    async fn set(&self, key: &str, value: Value, ttl: Option<Duration>) {
        self.inner.set(&self.key(key), value, ttl).await;
    }

    async fn get(&self, key: &str) -> Option<Value> {
        self.inner.get(&self.key(key)).await
    }
}

struct Entry {
    value: Value,
    expires_at: Option<Instant>,
}

impl Entry {
    fn is_expired(&self, now: Instant) -> bool {
        self.expires_at.map(|e| now >= e).unwrap_or(false)
    }
}

/// In-memory [`State`]. Cheap to share behind an `Arc`.
#[derive(Default)]
pub struct MemoryState {
    map: Mutex<HashMap<String, Entry>>,
}

impl MemoryState {
    /// Empty in-memory state.
    pub fn new() -> Self {
        Self::default()
    }

    /// Rebuild from a persisted snapshot. TTLs are not persisted: a durable
    /// instance's expiry is owned by its lifecycle, not by per-key timers.
    pub fn from_snapshot(entries: serde_json::Map<String, Value>) -> Self {
        Self {
            map: Mutex::new(
                entries
                    .into_iter()
                    .map(|(key, value)| {
                        (
                            key,
                            Entry {
                                value,
                                expires_at: None,
                            },
                        )
                    })
                    .collect(),
            ),
        }
    }

    /// Every live (non-expired) key, ready to be persisted as JSON.
    pub fn snapshot(&self) -> serde_json::Map<String, Value> {
        let now = Instant::now();
        self.map
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
            .iter()
            .filter(|(_, entry)| !entry.is_expired(now))
            .map(|(key, entry)| (key.clone(), entry.value.clone()))
            .collect()
    }

    fn expiry(ttl: Option<Duration>) -> Option<Instant> {
        ttl.map(|d| Instant::now() + d)
    }
}

#[async_trait]
impl State for MemoryState {
    async fn append(&self, key: &str, item: Value, max_len: Option<usize>, ttl: Option<Duration>) {
        let now = Instant::now();
        let mut map = self
            .map
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);

        // Drop the entry if it expired, so we start a fresh window.
        if let Some(e) = map.get(key) {
            if e.is_expired(now) {
                map.remove(key);
            }
        }

        let entry = map.entry(key.to_string()).or_insert_with(|| Entry {
            value: Value::Array(Vec::new()),
            expires_at: Self::expiry(ttl),
        });

        if !entry.value.is_array() {
            entry.value = Value::Array(Vec::new());
        }
        if let Value::Array(list) = &mut entry.value {
            list.push(item);
            if let Some(cap) = max_len {
                while list.len() > cap {
                    list.remove(0);
                }
            }
        }
        if ttl.is_some() {
            entry.expires_at = Self::expiry(ttl);
        }
    }

    async fn set(&self, key: &str, value: Value, ttl: Option<Duration>) {
        let mut map = self
            .map
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        map.insert(
            key.to_string(),
            Entry {
                value,
                expires_at: Self::expiry(ttl),
            },
        );
    }

    async fn get(&self, key: &str) -> Option<Value> {
        let now = Instant::now();
        let mut map = self
            .map
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        match map.get(key) {
            Some(e) if e.is_expired(now) => {
                map.remove(key);
                None
            }
            Some(e) => Some(e.value.clone()),
            None => None,
        }
    }
}

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

    #[tokio::test]
    async fn append_respects_max_len() {
        let s = MemoryState::new();
        for i in 0..5 {
            s.append("k", json!(i), Some(3), None).await;
        }
        let v = s.get("k").await.unwrap();
        assert_eq!(v, json!([2, 3, 4]));
    }

    #[tokio::test]
    async fn set_and_get() {
        let s = MemoryState::new();
        s.set("k", json!({"a": 1}), None).await;
        assert_eq!(s.get("k").await.unwrap(), json!({"a": 1}));
        assert!(s.get("missing").await.is_none());
    }

    #[tokio::test]
    async fn namespaces_isolate_instances() {
        let shared: Arc<dyn State> = Arc::new(MemoryState::new());
        let first = NamespacedState::new(shared.clone(), "first");
        let second = NamespacedState::new(shared, "second");
        first.set("branch.value", json!(1), None).await;
        second.set("branch.value", json!(2), None).await;
        assert_eq!(first.get("branch.value").await, Some(json!(1)));
        assert_eq!(second.get("branch.value").await, Some(json!(2)));
    }
}