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

rig_core/test_utils/
memory.rs
