theway_core/agent/runtime_extensions/
state.rs

1use std::sync::Arc;
2
3use async_trait::async_trait;
4use theway_contract::extension::ExtensionDurableEntry;
5use theway_contract::session::{SessionError, SessionStore, StoredSessionEntry};
6use thiserror::Error;
7
8#[async_trait]
9pub trait SessionExtensionStatePort: Send + Sync {
10    async fn append_durable_entries(
11        &self,
12        extension_id: &str,
13        entries: Vec<ExtensionDurableEntry>,
14    ) -> Result<Vec<String>, SessionExtensionStateError>;
15
16    async fn replay_durable_entries(
17        &self,
18        extension_id: &str,
19        leaf_id: Option<&str>,
20    ) -> Result<Vec<ExtensionDurableEntry>, SessionExtensionStateError>;
21}
22
23pub struct PersistentSessionExtensionStatePort {
24    store: Arc<dyn SessionStore>,
25}
26
27impl PersistentSessionExtensionStatePort {
28    pub fn new(store: Arc<dyn SessionStore>) -> Self {
29        Self { store }
30    }
31
32    pub fn store(&self) -> &Arc<dyn SessionStore> {
33        &self.store
34    }
35}
36
37#[async_trait]
38impl SessionExtensionStatePort for PersistentSessionExtensionStatePort {
39    async fn append_durable_entries(
40        &self,
41        extension_id: &str,
42        entries: Vec<ExtensionDurableEntry>,
43    ) -> Result<Vec<String>, SessionExtensionStateError> {
44        if entries.is_empty() {
45            return Ok(Vec::new());
46        }
47        for entry in &entries {
48            entry
49                .validate()
50                .map_err(|error| SessionExtensionStateError::InvalidEntry(error.to_string()))?;
51            if entry.extension_id != extension_id {
52                return Err(SessionExtensionStateError::OwnerMismatch {
53                    expected: extension_id.to_string(),
54                    actual: entry.extension_id.clone(),
55                });
56            }
57        }
58
59        let mut parent_id = self.store.get_leaf_id().await?;
60        let timestamp = chrono::Utc::now().to_rfc3339();
61        let mut ids = Vec::with_capacity(entries.len());
62        let mut stored = Vec::with_capacity(entries.len());
63        for entry in entries {
64            let id = self.store.create_entry_id().await?;
65            stored.push(StoredSessionEntry::extension(
66                id.clone(),
67                parent_id,
68                timestamp.clone(),
69                entry,
70            )?);
71            parent_id = Some(id.clone());
72            ids.push(id);
73        }
74        self.store.append_entries(stored).await?;
75        Ok(ids)
76    }
77
78    async fn replay_durable_entries(
79        &self,
80        extension_id: &str,
81        leaf_id: Option<&str>,
82    ) -> Result<Vec<ExtensionDurableEntry>, SessionExtensionStateError> {
83        self.store
84            .get_extension_entries(extension_id, leaf_id)
85            .await?
86            .into_iter()
87            .map(|entry| {
88                entry.extension_payload()?.ok_or_else(|| {
89                    SessionExtensionStateError::InvalidEntry(
90                        "extension query returned a non-extension entry".into(),
91                    )
92                })
93            })
94            .collect()
95    }
96}
97
98#[derive(Clone, Copy, Debug, Default)]
99pub struct NoopSessionExtensionStatePort;
100
101#[async_trait]
102impl SessionExtensionStatePort for NoopSessionExtensionStatePort {
103    async fn append_durable_entries(
104        &self,
105        _extension_id: &str,
106        entries: Vec<ExtensionDurableEntry>,
107    ) -> Result<Vec<String>, SessionExtensionStateError> {
108        if entries.is_empty() {
109            Ok(Vec::new())
110        } else {
111            Err(SessionExtensionStateError::Unavailable)
112        }
113    }
114
115    async fn replay_durable_entries(
116        &self,
117        _extension_id: &str,
118        _leaf_id: Option<&str>,
119    ) -> Result<Vec<ExtensionDurableEntry>, SessionExtensionStateError> {
120        Ok(Vec::new())
121    }
122}
123
124#[derive(Debug, Error)]
125pub enum SessionExtensionStateError {
126    #[error("session extension state persistence is unavailable")]
127    Unavailable,
128    #[error("extension durable entry is invalid: {0}")]
129    InvalidEntry(String),
130    #[error("extension durable entry owner mismatch: expected {expected}, got {actual}")]
131    OwnerMismatch { expected: String, actual: String },
132    #[error(transparent)]
133    Session(#[from] SessionError),
134}
theway_core/agent/runtime_extensions/state.rs

theway_core/agent/runtime_extensions/
state.rs
