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

okbot_testkit/
memory_store.rs
