stasis/infrastructure/runtime/
in_memory_thread_store.rs

1use std::collections::HashMap;
2use std::sync::{Arc, RwLock};
3
4use async_trait::async_trait;
5use chrono::{DateTime, Utc};
6
7use crate::domain::errors::{Result, StasisError};
8use crate::domain::runtime::thread::{
9    NewThread, NewThreadEvent, ThreadEvent, ThreadSnapshot,
10};
11use crate::ports::outbound::runtime::thread_store::ThreadStore;
12
13#[derive(Clone, Default)]
14pub struct InMemoryThreadStore {
15    threads: Arc<RwLock<HashMap<String, ThreadSnapshot>>>,
16    events: Arc<RwLock<HashMap<String, Vec<ThreadEvent>>>>,
17}
18
19#[async_trait]
20impl ThreadStore for InMemoryThreadStore {
21    async fn create_thread(&self, thread: NewThread) -> Result<ThreadSnapshot> {
22        let mut threads = self
23            .threads
24            .write()
25            .map_err(|_| StasisError::PortFailure("thread store lock poisoned".to_string()))?;
26
27        if threads.contains_key(&thread.thread_id) {
28            return Err(StasisError::PortFailure(format!(
29                "thread already exists: {}",
30                thread.thread_id
31            )));
32        }
33
34        if let Some(parent_thread_id) = &thread.parent_thread_id
35            && !threads.contains_key(parent_thread_id)
36        {
37            return Err(StasisError::PortFailure(format!(
38                "parent thread not found: {}",
39                parent_thread_id
40            )));
41        }
42
43        let record = ThreadSnapshot {
44            thread_id: thread.thread_id,
45            parent_thread_id: thread.parent_thread_id,
46            branch_label: thread.branch_label,
47            created_at: thread.created_at,
48            updated_at: thread.created_at,
49        };
50        threads.insert(record.thread_id.clone(), record.clone());
51        Ok(record)
52    }
53
54    async fn get_thread(&self, thread_id: &str) -> Result<Option<ThreadSnapshot>> {
55        let threads = self
56            .threads
57            .read()
58            .map_err(|_| StasisError::PortFailure("thread store lock poisoned".to_string()))?;
59
60        Ok(threads.get(thread_id).cloned())
61    }
62
63    async fn append_event(&self, event: NewThreadEvent) -> Result<ThreadEvent> {
64        {
65            let mut threads = self
66                .threads
67                .write()
68                .map_err(|_| StasisError::PortFailure("thread store lock poisoned".to_string()))?;
69            let Some(thread) = threads.get_mut(&event.thread_id) else {
70                return Err(StasisError::PortFailure(format!(
71                    "thread not found: {}",
72                    event.thread_id
73                )));
74            };
75            thread.updated_at = event.occurred_at;
76        }
77
78        let mut events = self.events.write().map_err(|_| {
79            StasisError::PortFailure("thread event store lock poisoned".to_string())
80        })?;
81
82        let record = ThreadEvent {
83            event_id: event.event_id,
84            thread_id: event.thread_id,
85            event_kind: event.event_kind,
86            payload_ref: event.payload_ref,
87            occurred_at: event.occurred_at,
88        };
89        events
90            .entry(record.thread_id.clone())
91            .or_insert_with(Vec::new)
92            .push(record.clone());
93
94        Ok(record)
95    }
96
97    async fn list_events(&self, thread_id: &str) -> Result<Vec<ThreadEvent>> {
98        let events = self.events.read().map_err(|_| {
99            StasisError::PortFailure("thread event store lock poisoned".to_string())
100        })?;
101
102        let mut result = events.get(thread_id).cloned().unwrap_or_default();
103        result.sort_by(|a, b| a.occurred_at.cmp(&b.occurred_at));
104        Ok(result)
105    }
106
107    async fn fork_thread(
108        &self,
109        parent_thread_id: &str,
110        child_thread_id: &str,
111        branch_label: Option<String>,
112        created_at: DateTime<Utc>,
113    ) -> Result<ThreadSnapshot> {
114        self.create_thread(NewThread {
115            thread_id: child_thread_id.to_string(),
116            parent_thread_id: Some(parent_thread_id.to_string()),
117            branch_label,
118            created_at,
119        })
120        .await
121    }
122
123    async fn list_lineage(&self, thread_id: &str) -> Result<Vec<ThreadSnapshot>> {
124        let threads = self
125            .threads
126            .read()
127            .map_err(|_| StasisError::PortFailure("thread store lock poisoned".to_string()))?;
128
129        let mut lineage = Vec::new();
130        let mut cursor = threads.get(thread_id).cloned();
131        while let Some(node) = cursor {
132            cursor = node
133                .parent_thread_id
134                .as_ref()
135                .and_then(|parent| threads.get(parent).cloned());
136            lineage.push(node);
137        }
138
139        lineage.reverse();
140        Ok(lineage)
141    }
142}
stasis/infrastructure/runtime/in_memory_thread_store.rs

stasis/infrastructure/runtime/
in_memory_thread_store.rs
