rig_core/test_utils/
memory.rs

1//! Conversation memory helpers for deterministic agent tests.
2
3use std::sync::{
4    Arc,
5    atomic::{AtomicUsize, Ordering},
6};
7
8use crate::{
9    completion::Message,
10    memory::{ConversationMemory, InMemoryConversationMemory, MemoryError},
11    wasm_compat::WasmBoxedFuture,
12};
13
14/// Memory backend that records load and append calls while delegating storage to
15/// [`InMemoryConversationMemory`].
16#[derive(Clone, Default)]
17pub struct CountingMemory {
18    inner: InMemoryConversationMemory,
19    loads: Arc<AtomicUsize>,
20    appends: Arc<AtomicUsize>,
21}
22
23impl CountingMemory {
24    /// Return the backing in-memory store.
25    pub fn inner(&self) -> &InMemoryConversationMemory {
26        &self.inner
27    }
28
29    /// Return the number of calls to [`ConversationMemory::load`].
30    pub fn load_count(&self) -> usize {
31        self.loads.load(Ordering::SeqCst)
32    }
33
34    /// Return the number of calls to [`ConversationMemory::append`].
35    pub fn append_count(&self) -> usize {
36        self.appends.load(Ordering::SeqCst)
37    }
38}
39
40impl ConversationMemory for CountingMemory {
41    fn load<'a>(
42        &'a self,
43        conversation_id: &'a str,
44    ) -> WasmBoxedFuture<'a, Result<Vec<Message>, MemoryError>> {
45        self.loads.fetch_add(1, Ordering::SeqCst);
46        self.inner.load(conversation_id)
47    }
48
49    fn append<'a>(
50        &'a self,
51        conversation_id: &'a str,
52        messages: Vec<Message>,
53    ) -> WasmBoxedFuture<'a, Result<(), MemoryError>> {
54        self.appends.fetch_add(1, Ordering::SeqCst);
55        self.inner.append(conversation_id, messages)
56    }
57
58    fn clear<'a>(
59        &'a self,
60        conversation_id: &'a str,
61    ) -> WasmBoxedFuture<'a, Result<(), MemoryError>> {
62        self.inner.clear(conversation_id)
63    }
64}
65
66/// Memory backend that always fails on load and no-ops append and clear.
67#[derive(Clone)]
68pub struct FailingMemory {
69    message: String,
70}
71
72impl FailingMemory {
73    /// Create a load-failing memory backend.
74    pub fn new(message: impl Into<String>) -> Self {
75        Self {
76            message: message.into(),
77        }
78    }
79}
80
81impl Default for FailingMemory {
82    fn default() -> Self {
83        Self::new("load boom")
84    }
85}
86
87impl ConversationMemory for FailingMemory {
88    fn load<'a>(
89        &'a self,
90        _conversation_id: &'a str,
91    ) -> WasmBoxedFuture<'a, Result<Vec<Message>, MemoryError>> {
92        let message = self.message.clone();
93        Box::pin(async move { Err(MemoryError::backend(std::io::Error::other(message))) })
94    }
95
96    fn append<'a>(
97        &'a self,
98        _conversation_id: &'a str,
99        _messages: Vec<Message>,
100    ) -> WasmBoxedFuture<'a, Result<(), MemoryError>> {
101        Box::pin(async { Ok(()) })
102    }
103
104    fn clear<'a>(
105        &'a self,
106        _conversation_id: &'a str,
107    ) -> WasmBoxedFuture<'a, Result<(), MemoryError>> {
108        Box::pin(async { Ok(()) })
109    }
110}
111
112/// Memory backend that loads empty history and always fails on append.
113#[derive(Clone)]
114pub struct AppendFailingMemory {
115    message: String,
116}
117
118impl AppendFailingMemory {
119    /// Create an append-failing memory backend.
120    pub fn new(message: impl Into<String>) -> Self {
121        Self {
122            message: message.into(),
123        }
124    }
125}
126
127impl Default for AppendFailingMemory {
128    fn default() -> Self {
129        Self::new("append boom")
130    }
131}
132
133impl ConversationMemory for AppendFailingMemory {
134    fn load<'a>(
135        &'a self,
136        _conversation_id: &'a str,
137    ) -> WasmBoxedFuture<'a, Result<Vec<Message>, MemoryError>> {
138        Box::pin(async { Ok(Vec::new()) })
139    }
140
141    fn append<'a>(
142        &'a self,
143        _conversation_id: &'a str,
144        _messages: Vec<Message>,
145    ) -> WasmBoxedFuture<'a, Result<(), MemoryError>> {
146        let message = self.message.clone();
147        Box::pin(async move { Err(MemoryError::backend(std::io::Error::other(message))) })
148    }
149
150    fn clear<'a>(
151        &'a self,
152        _conversation_id: &'a str,
153    ) -> WasmBoxedFuture<'a, Result<(), MemoryError>> {
154        Box::pin(async { Ok(()) })
155    }
156}
rig_core/test_utils/memory.rs

rig_core/test_utils/
memory.rs
