okbot_testkit/
memory_store.rs

1//! An in-memory [`ConversationStore`].
2
3use std::collections::HashMap;
4use std::sync::{Arc, Mutex};
5use std::time::Duration;
6
7use jiff::{SignedDuration, Timestamp};
8use okbot_core::{
9    AgentRunRecord, ChannelId, Clock, ConversationKey, ConversationStore, Fire, FireState,
10    NewMessage, PAIRING_MAX_FAILURES, StoreError, StoredMessage, UserId, async_trait,
11};
12
13#[derive(Default)]
14struct State {
15    seen: HashMap<(ChannelId, String), Timestamp>,
16    messages: Vec<(u64, NewMessage)>,
17    reset_after: HashMap<ConversationKey, u64>,
18    sessions: HashMap<ConversationKey, (String, String)>,
19    runs: Vec<AgentRunRecord>,
20    next: u64,
21    owners: Vec<(ChannelId, UserId)>,
22    /// code → (expires, failures)
23    codes: HashMap<String, (Timestamp, u32)>,
24    fires: HashMap<(String, Timestamp), Fire>,
25    direct: HashMap<(ChannelId, UserId), ConversationKey>,
26}
27
28/// A [`ConversationStore`] kept in memory, for runtime tests.
29pub struct MemoryStore {
30    state: Mutex<State>,
31    clock: Arc<dyn Clock>,
32}
33
34impl std::fmt::Debug for MemoryStore {
35    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
36        f.debug_struct("MemoryStore").finish_non_exhaustive()
37    }
38}
39
40impl MemoryStore {
41    /// An empty store reading time from `clock`.
42    pub fn new(clock: Arc<dyn Clock>) -> Self {
43        Self {
44            state: Mutex::new(State::default()),
45            clock,
46        }
47    }
48
49    fn lock(&self) -> std::sync::MutexGuard<'_, State> {
50        self.state
51            .lock()
52            .unwrap_or_else(std::sync::PoisonError::into_inner)
53    }
54
55    /// Recorded agent runs, oldest first.
56    pub fn runs(&self) -> Vec<AgentRunRecord> {
57        self.lock().runs.clone()
58    }
59}
60
61fn signed(d: Duration) -> SignedDuration {
62    SignedDuration::try_from(d).unwrap_or(SignedDuration::MAX)
63}
64
65#[async_trait]
66impl ConversationStore for MemoryStore {
67    async fn mark_seen(
68        &self,
69        channel: &ChannelId,
70        event_id: &str,
71        ttl: Duration,
72    ) -> Result<bool, StoreError> {
73        let now = self.clock.now();
74        let mut state = self.lock();
75        let key = (channel.clone(), event_id.to_string());
76        if state.seen.get(&key).is_some_and(|expires| *expires > now) {
77            return Ok(false);
78        }
79        state
80            .seen
81            .insert(key, now.saturating_add(signed(ttl)).unwrap_or(now));
82        Ok(true)
83    }
84
85    async fn append(&self, message: NewMessage) -> Result<(), StoreError> {
86        let mut state = self.lock();
87        state.next += 1;
88        let id = state.next;
89        state.messages.push((id, message));
90        Ok(())
91    }
92
93    async fn recent(
94        &self,
95        conversation: &ConversationKey,
96        limit: usize,
97    ) -> Result<Vec<StoredMessage>, StoreError> {
98        let state = self.lock();
99        let after = state.reset_after.get(conversation).copied().unwrap_or(0);
100        let mut out: Vec<StoredMessage> = state
101            .messages
102            .iter()
103            .rev()
104            .filter(|(id, m)| *id > after && m.conversation == *conversation)
105            .take(limit)
106            .map(|(_, m)| StoredMessage {
107                role: m.role,
108                author_name: m.author_name.clone(),
109                text: m.text.clone(),
110                at: m.at,
111            })
112            .collect();
113        out.reverse();
114        Ok(out)
115    }
116
117    async fn reset(&self, conversation: &ConversationKey) -> Result<(), StoreError> {
118        let mut state = self.lock();
119        let last = state.next;
120        state.reset_after.insert(conversation.clone(), last);
121        state.sessions.remove(conversation);
122        Ok(())
123    }
124
125    async fn session(
126        &self,
127        conversation: &ConversationKey,
128    ) -> Result<Option<(String, String)>, StoreError> {
129        Ok(self.lock().sessions.get(conversation).cloned())
130    }
131
132    async fn set_session(
133        &self,
134        conversation: &ConversationKey,
135        backend: &str,
136        session: &str,
137    ) -> Result<(), StoreError> {
138        self.lock().sessions.insert(
139            conversation.clone(),
140            (backend.to_string(), session.to_string()),
141        );
142        Ok(())
143    }
144
145    async fn owners(&self) -> Result<Vec<(ChannelId, UserId)>, StoreError> {
146        let mut owners = self.lock().owners.clone();
147        owners.sort();
148        Ok(owners)
149    }
150
151    async fn remove_owner(&self, channel: &ChannelId, user: &UserId) -> Result<bool, StoreError> {
152        let mut state = self.lock();
153        let before = state.owners.len();
154        state.owners.retain(|(c, u)| c != channel || u != user);
155        Ok(state.owners.len() != before)
156    }
157
158    async fn create_pairing_code(&self, code: &str, ttl: Duration) -> Result<(), StoreError> {
159        let expires = self.clock.now() + signed(ttl);
160        self.lock().codes.insert(code.to_string(), (expires, 0));
161        Ok(())
162    }
163
164    async fn redeem_pairing_code(
165        &self,
166        code: &str,
167        channel: &ChannelId,
168        user: &UserId,
169    ) -> Result<bool, StoreError> {
170        let now = self.clock.now();
171        let mut state = self.lock();
172        state.codes.retain(|_, (expires, _)| *expires > now);
173        if state.codes.remove(code).is_some() {
174            let account = (channel.clone(), user.clone());
175            if !state.owners.contains(&account) {
176                state.owners.push(account);
177            }
178            return Ok(true);
179        }
180        for (_, failures) in state.codes.values_mut() {
181            *failures += 1;
182        }
183        state
184            .codes
185            .retain(|_, (_, failures)| *failures < PAIRING_MAX_FAILURES);
186        Ok(false)
187    }
188
189    async fn record_run(&self, run: AgentRunRecord) -> Result<(), StoreError> {
190        self.lock().runs.push(run);
191        Ok(())
192    }
193
194    async fn claim_fire(&self, task_id: &str, fire_at: Timestamp) -> Result<Fire, StoreError> {
195        let now = self.clock.now();
196        Ok(self
197            .lock()
198            .fires
199            .entry((task_id.to_string(), fire_at))
200            .or_insert_with(|| Fire {
201                task_id: task_id.to_string(),
202                fire_at,
203                state: FireState::Sending,
204                attempts: 0,
205                updated_at: now,
206            })
207            .clone())
208    }
209
210    async fn set_fire(
211        &self,
212        task_id: &str,
213        fire_at: Timestamp,
214        state: FireState,
215        attempts: u32,
216    ) -> Result<(), StoreError> {
217        let fire = Fire {
218            task_id: task_id.to_string(),
219            fire_at,
220            state,
221            attempts,
222            updated_at: self.clock.now(),
223        };
224        self.lock()
225            .fires
226            .insert((task_id.to_string(), fire_at), fire);
227        Ok(())
228    }
229
230    async fn recent_fires(&self, limit: usize) -> Result<Vec<Fire>, StoreError> {
231        let mut fires: Vec<Fire> = self.lock().fires.values().cloned().collect();
232        fires.sort_by_key(|f| std::cmp::Reverse((f.updated_at, f.fire_at)));
233        fires.truncate(limit);
234        Ok(fires)
235    }
236
237    async fn set_direct_chat(
238        &self,
239        channel: &ChannelId,
240        user: &UserId,
241        conversation: &ConversationKey,
242    ) -> Result<(), StoreError> {
243        self.lock()
244            .direct
245            .insert((channel.clone(), user.clone()), conversation.clone());
246        Ok(())
247    }
248
249    async fn direct_chat(
250        &self,
251        channel: &ChannelId,
252        user: &UserId,
253    ) -> Result<Option<ConversationKey>, StoreError> {
254        Ok(self
255            .lock()
256            .direct
257            .get(&(channel.clone(), user.clone()))
258            .cloned())
259    }
260
261    async fn purge_before(&self, before: Timestamp) -> Result<u64, StoreError> {
262        let now = self.clock.now();
263        let mut state = self.lock();
264        let count = |s: &State| s.messages.len() + s.runs.len() + s.seen.len() + s.fires.len();
265        let start = count(&state);
266        state.messages.retain(|(_, m)| m.at >= before);
267        state.runs.retain(|r| r.started_at >= before);
268        state.seen.retain(|_, expires| *expires > now);
269        state
270            .fires
271            .retain(|_, f| f.updated_at >= before || f.state == FireState::Sending);
272        let end = count(&state);
273        Ok((start - end) as u64)
274    }
275}
276
277#[cfg(test)]
278mod tests {
279    use super::*;
280    use crate::ManualClock;
281
282    #[tokio::test]
283    async fn conformance() {
284        let clock = Arc::new(ManualClock::default());
285        let store = MemoryStore::new(clock.clone());
286        crate::conformance::conversation_store(&store, &clock).await;
287    }
288}
okbot_testkit/memory_store.rs

okbot_testkit/
memory_store.rs
