behest_store/memory/
execution.rs

1//! In-memory execution store for tool calls and token usage tracking.
2
3use std::collections::HashMap;
4
5use async_trait::async_trait;
6use tokio::sync::RwLock;
7use uuid::Uuid;
8
9use crate::{
10    ExecutionStore, SessionStats, StoreResult, ToolExecution, ToolExecutionStatus, UsageRecord,
11};
12
13/// In-memory execution store for testing, development, and prototyping.
14///
15/// Tracks tool executions and token usage records in memory-backed
16/// `HashMap`s protected by `RwLock`. Data is lost when the process exits.
17/// Implements [`ExecutionStore`].
18#[derive(Default)]
19pub struct MemoryExecutionStore {
20    executions: RwLock<HashMap<Uuid, ToolExecution>>,
21    usage_records: RwLock<HashMap<Uuid, UsageRecord>>,
22    message_counts: RwLock<HashMap<Uuid, u64>>,
23}
24
25impl MemoryExecutionStore {
26    /// Creates an empty in-memory execution store.
27    #[must_use]
28    pub fn new() -> Self {
29        Self::default()
30    }
31
32    /// Sets the message count for a session, used by [`session_stats`](ExecutionStore::session_stats).
33    ///
34    /// In a real backend this would be computed from a messages table;
35    /// the in-memory store requires the caller to explicitly provide it
36    /// since messages are maintained in the session store, not here.
37    pub async fn set_message_count(&self, session_id: Uuid, count: u64) {
38        self.message_counts.write().await.insert(session_id, count);
39    }
40}
41
42#[async_trait]
43impl ExecutionStore for MemoryExecutionStore {
44    async fn record_execution(&self, execution: ToolExecution) -> StoreResult<ToolExecution> {
45        let id = execution.id;
46        self.executions.write().await.insert(id, execution.clone());
47        Ok(execution)
48    }
49
50    async fn list_executions(&self, session_id: &Uuid) -> StoreResult<Vec<ToolExecution>> {
51        let executions = self.executions.read().await;
52        let mut result: Vec<ToolExecution> = executions
53            .values()
54            .filter(|e| e.session_id == *session_id)
55            .cloned()
56            .collect();
57        result.sort_by_key(|e| e.created_at);
58        Ok(result)
59    }
60
61    async fn list_executions_by_message(
62        &self,
63        message_id: &Uuid,
64    ) -> StoreResult<Vec<ToolExecution>> {
65        let executions = self.executions.read().await;
66        let mut result: Vec<ToolExecution> = executions
67            .values()
68            .filter(|e| e.message_id == *message_id)
69            .cloned()
70            .collect();
71        result.sort_by_key(|e| e.created_at);
72        Ok(result)
73    }
74
75    async fn record_usage(&self, record: UsageRecord) -> StoreResult<UsageRecord> {
76        let id = record.id;
77        self.usage_records.write().await.insert(id, record.clone());
78        Ok(record)
79    }
80
81    async fn list_usage(&self, session_id: &Uuid) -> StoreResult<Vec<UsageRecord>> {
82        let records = self.usage_records.read().await;
83        let mut result: Vec<UsageRecord> = records
84            .values()
85            .filter(|r| r.session_id == *session_id)
86            .cloned()
87            .collect();
88        result.sort_by_key(|r| r.created_at);
89        Ok(result)
90    }
91
92    async fn session_stats(&self, session_id: &Uuid) -> StoreResult<SessionStats> {
93        let executions = self.executions.read().await;
94        let usage_records = self.usage_records.read().await;
95        let message_counts = self.message_counts.read().await;
96
97        let session_execs: Vec<&ToolExecution> = executions
98            .values()
99            .filter(|e| e.session_id == *session_id)
100            .collect();
101
102        let session_usage: Vec<&UsageRecord> = usage_records
103            .values()
104            .filter(|r| r.session_id == *session_id)
105            .collect();
106
107        let tool_call_count = session_execs.len() as u64;
108        let tool_success_count = session_execs
109            .iter()
110            .filter(|e| e.status == ToolExecutionStatus::Success)
111            .count() as u64;
112        let tool_failure_count = session_execs
113            .iter()
114            .filter(|e| e.status == ToolExecutionStatus::Failed)
115            .count() as u64;
116
117        let total_duration_ms: u64 = session_execs
118            .iter()
119            .map(|e| u64::try_from(e.duration.as_millis()).unwrap_or(u64::MAX))
120            .sum();
121
122        let avg_tool_duration_ms = total_duration_ms.checked_div(tool_call_count).unwrap_or(0);
123
124        let total_input_tokens: u64 = session_usage.iter().map(|r| r.input_tokens).sum();
125        let total_output_tokens: u64 = session_usage.iter().map(|r| r.output_tokens).sum();
126        let total_tokens: u64 = session_usage.iter().map(|r| r.total_tokens).sum();
127
128        let message_count = message_counts.get(session_id).copied().unwrap_or(0);
129
130        Ok(SessionStats {
131            session_id: *session_id,
132            message_count,
133            tool_call_count,
134            tool_success_count,
135            tool_failure_count,
136            total_input_tokens,
137            total_output_tokens,
138            total_tokens,
139            avg_tool_duration_ms,
140        })
141    }
142}
143
144#[cfg(test)]
145#[allow(clippy::unwrap_used)]
146mod tests {
147    use super::*;
148    use behest_provider::TokenUsage;
149    use serde_json::json;
150    use std::time::Duration;
151
152    fn test_session_id() -> Uuid {
153        Uuid::now_v7()
154    }
155
156    fn test_message_id() -> Uuid {
157        Uuid::now_v7()
158    }
159
160    #[tokio::test]
161    async fn memory_execution_store_should_record_and_list_executions() {
162        let store = MemoryExecutionStore::new();
163        let session_id = test_session_id();
164        let message_id = test_message_id();
165
166        let exec = ToolExecution::new(
167            session_id,
168            message_id,
169            "call_1",
170            "get_weather",
171            json!({"city": "London"}),
172        )
173        .with_success(json!({"temp": 22}), Duration::from_millis(150));
174
175        store.record_execution(exec).await.unwrap();
176
177        let execs = store.list_executions(&session_id).await.unwrap();
178        assert_eq!(execs.len(), 1);
179        assert_eq!(execs[0].tool_name, "get_weather");
180        assert_eq!(execs[0].status, ToolExecutionStatus::Success);
181        assert_eq!(execs[0].duration, Duration::from_millis(150));
182    }
183
184    #[tokio::test]
185    async fn memory_execution_store_should_list_by_message() {
186        let store = MemoryExecutionStore::new();
187        let session_id = test_session_id();
188        let msg1 = test_message_id();
189        let msg2 = test_message_id();
190
191        store
192            .record_execution(ToolExecution::new(
193                session_id,
194                msg1,
195                "call_1",
196                "tool_a",
197                json!({}),
198            ))
199            .await
200            .unwrap();
201        store
202            .record_execution(ToolExecution::new(
203                session_id,
204                msg2,
205                "call_2",
206                "tool_b",
207                json!({}),
208            ))
209            .await
210            .unwrap();
211
212        let by_msg1 = store.list_executions_by_message(&msg1).await.unwrap();
213        assert_eq!(by_msg1.len(), 1);
214        assert_eq!(by_msg1[0].tool_name, "tool_a");
215    }
216
217    #[tokio::test]
218    async fn memory_execution_store_should_record_failed_execution() {
219        let store = MemoryExecutionStore::new();
220        let session_id = test_session_id();
221        let message_id = test_message_id();
222
223        let exec = ToolExecution::new(session_id, message_id, "call_1", "broken_tool", json!({}))
224            .with_failure("tool crashed", Duration::from_millis(50));
225
226        store.record_execution(exec).await.unwrap();
227
228        let execs = store.list_executions(&session_id).await.unwrap();
229        assert_eq!(execs[0].status, ToolExecutionStatus::Failed);
230        assert_eq!(execs[0].error.as_deref(), Some("tool crashed"));
231    }
232
233    #[tokio::test]
234    async fn memory_execution_store_should_record_and_list_usage() {
235        let store = MemoryExecutionStore::new();
236        let session_id = test_session_id();
237        let message_id = test_message_id();
238
239        let record = UsageRecord::new(
240            session_id,
241            message_id,
242            "openai",
243            "gpt-4",
244            TokenUsage::new(100, 50),
245        );
246
247        store.record_usage(record).await.unwrap();
248
249        let usage = store.list_usage(&session_id).await.unwrap();
250        assert_eq!(usage.len(), 1);
251        assert_eq!(usage[0].provider, "openai");
252        assert_eq!(usage[0].model, "gpt-4");
253        assert_eq!(usage[0].input_tokens, 100);
254        assert_eq!(usage[0].output_tokens, 50);
255        assert_eq!(usage[0].total_tokens, 150);
256    }
257
258    #[tokio::test]
259    async fn memory_execution_store_should_compute_session_stats() {
260        let store = MemoryExecutionStore::new();
261        let session_id = test_session_id();
262        let msg = test_message_id();
263
264        store.set_message_count(session_id, 10).await;
265
266        // 3 tool executions: 2 success, 1 failure
267        store
268            .record_execution(
269                ToolExecution::new(session_id, msg, "c1", "t1", json!({}))
270                    .with_success(json!(null), Duration::from_millis(100)),
271            )
272            .await
273            .unwrap();
274        store
275            .record_execution(
276                ToolExecution::new(session_id, msg, "c2", "t2", json!({}))
277                    .with_success(json!(null), Duration::from_millis(200)),
278            )
279            .await
280            .unwrap();
281        store
282            .record_execution(
283                ToolExecution::new(session_id, msg, "c3", "t3", json!({}))
284                    .with_failure("err", Duration::from_millis(300)),
285            )
286            .await
287            .unwrap();
288
289        // 2 usage records
290        store
291            .record_usage(UsageRecord::new(
292                session_id,
293                msg,
294                "openai",
295                "gpt-4",
296                TokenUsage::new(100, 50),
297            ))
298            .await
299            .unwrap();
300        store
301            .record_usage(UsageRecord::new(
302                session_id,
303                msg,
304                "openai",
305                "gpt-4",
306                TokenUsage::new(200, 80),
307            ))
308            .await
309            .unwrap();
310
311        let stats = store.session_stats(&session_id).await.unwrap();
312
313        assert_eq!(stats.message_count, 10);
314        assert_eq!(stats.tool_call_count, 3);
315        assert_eq!(stats.tool_success_count, 2);
316        assert_eq!(stats.tool_failure_count, 1);
317        assert_eq!(stats.total_input_tokens, 300);
318        assert_eq!(stats.total_output_tokens, 130);
319        assert_eq!(stats.total_tokens, 430);
320        assert_eq!(stats.avg_tool_duration_ms, 200); // (100+200+300)/3
321    }
322
323    #[tokio::test]
324    async fn memory_execution_store_should_return_empty_stats() {
325        let store = MemoryExecutionStore::new();
326        let stats = store.session_stats(&test_session_id()).await.unwrap();
327
328        assert_eq!(stats.message_count, 0);
329        assert_eq!(stats.tool_call_count, 0);
330        assert_eq!(stats.total_tokens, 0);
331    }
332}
behest_store/memory/execution.rs

behest_store/memory/
execution.rs
