everruns_core/
in_memory.rs

1// In-memory implementations for examples and testing
2//
3// These implementations keep all data in memory, making them perfect for:
4// - Standalone examples that don't need a database
5// - Unit tests
6// - Quick prototyping
7
8use crate::agent::Agent;
9use crate::credential_provider::CredentialProvider;
10use crate::harness::Harness;
11use crate::provider::DriverId;
12use crate::session::Session;
13use crate::tool_types::{ToolCall, ToolDefinition, ToolResult};
14use crate::traits::ResolvedModel;
15use crate::typed_id::{AgentId, EventId, HarnessId, MessageId, ModelId, SessionId};
16use async_trait::async_trait;
17use std::collections::HashMap;
18use std::sync::Arc;
19use tokio::sync::RwLock;
20use uuid::Uuid;
21
22use crate::error::Result;
23use crate::message::Message;
24use crate::message_filter::MessageQuery;
25use crate::message_retriever::{InputMessage, MessageHistory, MessageRetriever};
26use crate::traits::{AgentStore, HarnessStore, ProviderStore, SessionStore, ToolExecutor};
27use chrono::Utc;
28
29// ============================================================================
30// InMemoryMessageRetriever - In-memory message storage for testing
31// ============================================================================
32
33/// In-memory message retriever
34///
35/// Stores messages in a HashMap keyed by session ID.
36/// Implements the `MessageRetriever` trait for retrieval operations.
37///
38/// Note: Write operations (add, store) are provided as inherent methods
39/// for testing purposes. In production, messages are stored via EventEmitter.
40#[derive(Debug, Default, Clone)]
41pub struct InMemoryMessageRetriever {
42    messages: Arc<RwLock<HashMap<SessionId, Vec<Message>>>>,
43}
44
45impl InMemoryMessageRetriever {
46    /// Create a new in-memory message retriever
47    pub fn new() -> Self {
48        Self {
49            messages: Arc::new(RwLock::new(HashMap::new())),
50        }
51    }
52
53    /// Get all sessions
54    pub async fn sessions(&self) -> Vec<SessionId> {
55        self.messages.read().await.keys().copied().collect()
56    }
57
58    /// Clear all messages
59    pub async fn clear(&self) {
60        self.messages.write().await.clear();
61    }
62
63    /// Clear messages for a specific session
64    pub async fn clear_session(&self, session_id: SessionId) {
65        self.messages.write().await.remove(&session_id);
66    }
67
68    /// Pre-populate with messages (useful for testing)
69    pub async fn seed(&self, session_id: SessionId, messages: Vec<Message>) {
70        self.messages.write().await.insert(session_id, messages);
71    }
72
73    /// Add a new message and return it with generated ID (for testing)
74    ///
75    /// Note: In production, messages are stored via EventService.
76    /// This method is provided for test setup and in-memory usage.
77    pub async fn add(&self, session_id: SessionId, input: InputMessage) -> Result<Message> {
78        let message = Message {
79            id: MessageId::new(),
80            role: input.role,
81            content: input.content,
82            phase: None,
83            thinking: None, // InputMessage doesn't include thinking (user messages don't have thinking)
84            thinking_signature: None,
85            controls: input.controls,
86            metadata: input.metadata,
87            external_actor: None,
88            created_at: Utc::now(),
89        };
90
91        self.messages
92            .write()
93            .await
94            .entry(session_id)
95            .or_default()
96            .push(message.clone());
97
98        Ok(message)
99    }
100
101    /// Store an existing message (for testing)
102    ///
103    /// Note: In production, messages are stored via EventEmitter.
104    /// This method is provided for test setup and in-memory usage.
105    pub async fn store(&self, session_id: SessionId, message: Message) -> Result<()> {
106        self.messages
107            .write()
108            .await
109            .entry(session_id)
110            .or_default()
111            .push(message);
112        Ok(())
113    }
114}
115
116#[async_trait]
117impl MessageRetriever for InMemoryMessageRetriever {
118    async fn get(&self, session_id: SessionId, message_id: MessageId) -> Result<Option<Message>> {
119        Ok(self
120            .messages
121            .read()
122            .await
123            .get(&session_id)
124            .and_then(|messages| messages.iter().find(|m| m.id == message_id).cloned()))
125    }
126
127    async fn load(&self, session_id: SessionId) -> Result<Vec<Message>> {
128        Ok(self
129            .messages
130            .read()
131            .await
132            .get(&session_id)
133            .cloned()
134            .unwrap_or_default())
135    }
136
137    async fn load_filtered(&self, query: MessageQuery) -> Result<Vec<Message>> {
138        use crate::message_filter::MessageFilter;
139
140        let mut messages = self.load(query.session_id).await?;
141        if let Some(after) = query.after_sequence {
142            messages = messages.into_iter().skip(after.max(0) as usize).collect();
143        }
144
145        // Apply filters
146        for filter in &query.filters {
147            match filter {
148                MessageFilter::TimeRange { from, to } => {
149                    messages.retain(|m| {
150                        let after_from = from.is_none_or(|t| m.created_at >= t);
151                        let before_to = to.is_none_or(|t| m.created_at <= t);
152                        after_from && before_to
153                    });
154                }
155                MessageFilter::Search(q) => {
156                    let q_lower = q.to_lowercase();
157                    messages.retain(|m| {
158                        m.text()
159                            .is_some_and(|t| t.to_lowercase().contains(&q_lower))
160                    });
161                }
162                MessageFilter::Custom(predicate) => {
163                    messages.retain(|m| predicate(m));
164                }
165                // Other filters not commonly used in-memory
166                _ => {}
167            }
168        }
169
170        query.apply_windowing(&mut messages);
171
172        // Apply injections
173        if query.has_injections() {
174            query.apply_injections(&mut messages);
175        }
176
177        Ok(messages)
178    }
179
180    async fn load_filtered_history(&self, query: MessageQuery) -> Result<MessageHistory> {
181        let source_sequence = self
182            .messages
183            .read()
184            .await
185            .get(&query.session_id)
186            .map(|messages| messages.len() as i64)
187            .unwrap_or(0);
188        Ok(MessageHistory {
189            messages: self.load_filtered(query).await?,
190            source_sequence: Some(source_sequence),
191        })
192    }
193
194    async fn count(&self, session_id: SessionId) -> Result<usize> {
195        Ok(self
196            .messages
197            .read()
198            .await
199            .get(&session_id)
200            .map(|m| m.len())
201            .unwrap_or(0))
202    }
203}
204
205// ============================================================================
206// InMemoryAgentStore - Stores agents in memory
207// ============================================================================
208
209/// In-memory agent store
210///
211/// Stores agents in a HashMap keyed by agent ID.
212/// Useful for testing and examples where you want to configure agents without a database.
213#[derive(Debug, Default, Clone)]
214pub struct InMemoryAgentStore {
215    agents: Arc<RwLock<HashMap<AgentId, Agent>>>,
216}
217
218impl InMemoryAgentStore {
219    /// Create a new in-memory agent store
220    pub fn new() -> Self {
221        Self {
222            agents: Arc::new(RwLock::new(HashMap::new())),
223        }
224    }
225
226    /// Add an agent to the store
227    pub async fn add_agent(&self, agent: Agent) {
228        self.agents.write().await.insert(agent.public_id, agent);
229    }
230
231    /// Get all agent IDs
232    pub async fn agent_ids(&self) -> Vec<AgentId> {
233        self.agents.read().await.keys().copied().collect()
234    }
235
236    /// Clear all agents
237    pub async fn clear(&self) {
238        self.agents.write().await.clear();
239    }
240}
241
242#[async_trait]
243impl AgentStore for InMemoryAgentStore {
244    async fn get_agent(&self, agent_id: AgentId) -> Result<Option<Agent>> {
245        Ok(self.agents.read().await.get(&agent_id).cloned())
246    }
247}
248
249// ============================================================================
250// InMemoryHarnessStore - Stores harnesses in memory
251// ============================================================================
252
253/// In-memory harness store
254///
255/// Stores harnesses in a HashMap keyed by harness ID.
256/// Useful for testing and examples where you want to configure harnesses without a database.
257#[derive(Debug, Default, Clone)]
258pub struct InMemoryHarnessStore {
259    harnesses: Arc<RwLock<HashMap<HarnessId, Harness>>>,
260}
261
262impl InMemoryHarnessStore {
263    /// Create a new in-memory harness store
264    pub fn new() -> Self {
265        Self {
266            harnesses: Arc::new(RwLock::new(HashMap::new())),
267        }
268    }
269
270    /// Add a harness to the store
271    pub async fn add_harness(&self, harness: Harness) {
272        self.harnesses.write().await.insert(harness.id, harness);
273    }
274}
275
276#[async_trait]
277impl HarnessStore for InMemoryHarnessStore {
278    async fn get_harness_chain(&self, harness_id: HarnessId) -> Result<Vec<Harness>> {
279        Ok(self
280            .harnesses
281            .read()
282            .await
283            .get(&harness_id)
284            .cloned()
285            .into_iter()
286            .collect())
287    }
288}
289
290// ============================================================================
291// InMemorySessionStore - Stores sessions in memory
292// ============================================================================
293
294/// In-memory session store
295///
296/// Stores sessions in a HashMap keyed by session ID.
297/// Useful for testing and examples where you want to configure sessions without a database.
298#[derive(Debug, Default, Clone)]
299pub struct InMemorySessionStore {
300    sessions: Arc<RwLock<HashMap<SessionId, Session>>>,
301}
302
303impl InMemorySessionStore {
304    /// Create a new in-memory session store
305    pub fn new() -> Self {
306        Self {
307            sessions: Arc::new(RwLock::new(HashMap::new())),
308        }
309    }
310
311    /// Add a session to the store
312    pub async fn add_session(&self, session: Session) {
313        self.sessions.write().await.insert(session.id, session);
314    }
315
316    /// Get all session IDs
317    pub async fn session_ids(&self) -> Vec<SessionId> {
318        self.sessions.read().await.keys().copied().collect()
319    }
320
321    /// Clear all sessions
322    pub async fn clear(&self) {
323        self.sessions.write().await.clear();
324    }
325}
326
327#[async_trait]
328impl SessionStore for InMemorySessionStore {
329    async fn get_session(&self, session_id: SessionId) -> Result<Option<Session>> {
330        Ok(self.sessions.read().await.get(&session_id).cloned())
331    }
332}
333
334// ============================================================================
335// InMemoryProviderStore - Stores LLM provider configurations in memory
336// ============================================================================
337
338/// In-memory LLM provider store
339///
340/// Stores model configurations in a HashMap keyed by model UUID.
341/// Useful for testing and examples where you want to configure providers without a database.
342///
343/// # Example
344///
345/// ```ignore
346/// use everruns_core::in_memory::InMemoryProviderStore;
347/// use everruns_core::EnvCredentialProvider;
348///
349/// let store = InMemoryProviderStore::from_credential_provider(&EnvCredentialProvider).await;
350/// // Uses OPENAI_API_KEY or ANTHROPIC_API_KEY via the injected provider
351/// ```
352#[derive(Debug, Default, Clone)]
353pub struct InMemoryProviderStore {
354    models: Arc<RwLock<HashMap<ModelId, ResolvedModel>>>,
355    default_model: Arc<RwLock<Option<ResolvedModel>>>,
356}
357
358impl InMemoryProviderStore {
359    /// Create a new empty in-memory provider store
360    pub fn new() -> Self {
361        Self {
362            models: Arc::new(RwLock::new(HashMap::new())),
363            default_model: Arc::new(RwLock::new(None)),
364        }
365    }
366
367    /// Create a provider store from an injected [`CredentialProvider`].
368    ///
369    /// Checks OpenAI first, then Anthropic, and configures a default model for
370    /// whichever the provider resolves credentials for. The store never reads
371    /// the process environment itself; standalone/dev callers pass
372    /// [`EnvCredentialProvider`](crate::credential_provider::EnvCredentialProvider)
373    /// to opt into env-based credentials.
374    pub async fn from_credential_provider(provider: &dyn CredentialProvider) -> Self {
375        let store = Self::new();
376
377        // Check for OpenAI first, then Anthropic.
378        if let Some(creds) = provider
379            .resolve(&DriverId::OpenAI)
380            .filter(|c| c.api_key.is_some())
381        {
382            store
383                .set_default_model(ResolvedModel {
384                    model: "gpt-5.4".to_string(),
385                    provider_type: DriverId::OpenAI,
386                    api_key: creds.api_key,
387                    base_url: creds.base_url,
388                    provider_metadata: None,
389                })
390                .await;
391        } else if let Some(creds) = provider
392            .resolve(&DriverId::Anthropic)
393            .filter(|c| c.api_key.is_some())
394        {
395            store
396                .set_default_model(ResolvedModel {
397                    model: "claude-sonnet-4-20250514".to_string(),
398                    provider_type: DriverId::Anthropic,
399                    api_key: creds.api_key,
400                    base_url: creds.base_url,
401                    provider_metadata: None,
402                })
403                .await;
404        }
405
406        store
407    }
408
409    /// Create a provider store with a specific default model
410    pub async fn with_default(model: ResolvedModel) -> Self {
411        let store = Self::new();
412        store.set_default_model(model).await;
413        store
414    }
415
416    /// Add a model to the store
417    pub async fn add_model(&self, model_id: ModelId, model: ResolvedModel) {
418        self.models.write().await.insert(model_id, model);
419    }
420
421    /// Set the default model
422    pub async fn set_default_model(&self, model: ResolvedModel) {
423        *self.default_model.write().await = Some(model);
424    }
425
426    /// Clear all models
427    pub async fn clear(&self) {
428        self.models.write().await.clear();
429        *self.default_model.write().await = None;
430    }
431}
432
433#[async_trait]
434impl ProviderStore for InMemoryProviderStore {
435    async fn get_resolved_model(&self, model_id: ModelId) -> Result<Option<ResolvedModel>> {
436        Ok(self.models.read().await.get(&model_id).cloned())
437    }
438
439    async fn get_default_model(&self) -> Result<Option<ResolvedModel>> {
440        Ok(self.default_model.read().await.clone())
441    }
442}
443
444// ============================================================================
445// MockToolExecutor - Returns predefined results
446// ============================================================================
447
448/// Mock tool executor for testing
449///
450/// Returns predefined results based on tool name.
451#[derive(Debug, Default)]
452pub struct MockToolExecutor {
453    results: Arc<RwLock<HashMap<String, serde_json::Value>>>,
454    call_log: Arc<RwLock<Vec<ToolCall>>>,
455}
456
457impl MockToolExecutor {
458    /// Create a new mock tool executor
459    pub fn new() -> Self {
460        Self {
461            results: Arc::new(RwLock::new(HashMap::new())),
462            call_log: Arc::new(RwLock::new(Vec::new())),
463        }
464    }
465
466    /// Set the result for a specific tool
467    pub async fn set_result(&self, tool_name: impl Into<String>, result: serde_json::Value) {
468        self.results.write().await.insert(tool_name.into(), result);
469    }
470
471    /// Get the call log
472    pub async fn calls(&self) -> Vec<ToolCall> {
473        self.call_log.read().await.clone()
474    }
475
476    /// Clear the call log
477    pub async fn clear_calls(&self) {
478        self.call_log.write().await.clear();
479    }
480}
481
482#[async_trait]
483impl ToolExecutor for MockToolExecutor {
484    async fn execute(
485        &self,
486        tool_call: &ToolCall,
487        _tool_def: &ToolDefinition,
488    ) -> Result<ToolResult> {
489        // Log the call
490        self.call_log.write().await.push(tool_call.clone());
491
492        // Return predefined result or default
493        let result = self
494            .results
495            .read()
496            .await
497            .get(&tool_call.name)
498            .cloned()
499            .unwrap_or_else(|| serde_json::json!({"status": "ok"}));
500
501        Ok(ToolResult {
502            tool_call_id: tool_call.id.clone(),
503            result: Some(result),
504            images: None,
505            error: None,
506            connection_required: None,
507            raw_output: None,
508        })
509    }
510}
511
512// ============================================================================
513// EchoToolExecutor - Echoes back the arguments
514// ============================================================================
515
516/// Tool executor that echoes back the arguments
517///
518/// Useful for simple testing without setting up mock results.
519#[derive(Debug, Default, Clone, Copy)]
520pub struct EchoToolExecutor;
521
522impl EchoToolExecutor {
523    pub fn new() -> Self {
524        Self
525    }
526}
527
528#[async_trait]
529impl ToolExecutor for EchoToolExecutor {
530    async fn execute(
531        &self,
532        tool_call: &ToolCall,
533        _tool_def: &ToolDefinition,
534    ) -> Result<ToolResult> {
535        Ok(ToolResult {
536            tool_call_id: tool_call.id.clone(),
537            result: Some(serde_json::json!({
538                "echoed_tool": tool_call.name,
539                "echoed_arguments": tool_call.arguments
540            })),
541            images: None,
542            error: None,
543            connection_required: None,
544            raw_output: None,
545        })
546    }
547}
548
549// ============================================================================
550// FailingToolExecutor - Always returns an error
551// ============================================================================
552
553/// Tool executor that always fails
554///
555/// Useful for testing error handling.
556#[derive(Debug, Clone)]
557pub struct FailingToolExecutor {
558    error_message: String,
559}
560
561impl FailingToolExecutor {
562    pub fn new(error_message: impl Into<String>) -> Self {
563        Self {
564            error_message: error_message.into(),
565        }
566    }
567}
568
569impl Default for FailingToolExecutor {
570    fn default() -> Self {
571        Self::new("Tool execution failed")
572    }
573}
574
575#[async_trait]
576impl ToolExecutor for FailingToolExecutor {
577    async fn execute(
578        &self,
579        tool_call: &ToolCall,
580        _tool_def: &ToolDefinition,
581    ) -> Result<ToolResult> {
582        Ok(ToolResult {
583            tool_call_id: tool_call.id.clone(),
584            result: None,
585            images: None,
586            error: Some(self.error_message.clone()),
587            connection_required: None,
588            raw_output: None,
589        })
590    }
591}
592
593// ============================================================================
594// MockProvider - Returns predefined responses
595// ============================================================================
596
597use crate::driver_registry::{
598    ChatDriver, LlmCallConfig, LlmMessage, LlmResponseStream, LlmStreamEvent,
599};
600use crate::events::{Event, EventRequest};
601use crate::traits::EventEmitter;
602use futures::stream;
603
604/// Mock LLM provider for testing
605///
606/// Returns predefined responses in sequence.
607#[derive(Debug, Default)]
608pub struct MockProvider {
609    responses: Arc<RwLock<Vec<MockLlmResponse>>>,
610    call_index: Arc<RwLock<usize>>,
611    call_log: Arc<RwLock<Vec<Vec<LlmMessage>>>>,
612}
613
614/// A mock LLM response
615#[derive(Debug, Clone)]
616pub struct MockLlmResponse {
617    pub text: String,
618    pub tool_calls: Option<Vec<ToolCall>>,
619}
620
621impl MockLlmResponse {
622    /// Create a text-only response
623    pub fn text(text: impl Into<String>) -> Self {
624        Self {
625            text: text.into(),
626            tool_calls: None,
627        }
628    }
629
630    /// Create a response with tool calls
631    pub fn with_tools(text: impl Into<String>, tool_calls: Vec<ToolCall>) -> Self {
632        Self {
633            text: text.into(),
634            tool_calls: Some(tool_calls),
635        }
636    }
637}
638
639impl MockProvider {
640    /// Create a new mock LLM provider
641    pub fn new() -> Self {
642        Self {
643            responses: Arc::new(RwLock::new(Vec::new())),
644            call_index: Arc::new(RwLock::new(0)),
645            call_log: Arc::new(RwLock::new(Vec::new())),
646        }
647    }
648
649    /// Add a response to the queue
650    pub async fn add_response(&self, response: MockLlmResponse) {
651        self.responses.write().await.push(response);
652    }
653
654    /// Set all responses at once
655    pub async fn set_responses(&self, responses: Vec<MockLlmResponse>) {
656        *self.responses.write().await = responses;
657        *self.call_index.write().await = 0;
658    }
659
660    /// Get the call log
661    pub async fn calls(&self) -> Vec<Vec<LlmMessage>> {
662        self.call_log.read().await.clone()
663    }
664
665    /// Reset the provider
666    pub async fn reset(&self) {
667        self.responses.write().await.clear();
668        *self.call_index.write().await = 0;
669        self.call_log.write().await.clear();
670    }
671}
672
673#[async_trait]
674impl ChatDriver for MockProvider {
675    async fn chat_completion_stream(
676        &self,
677        _endpoint: &crate::ProviderEndpoint,
678        messages: Vec<LlmMessage>,
679        _config: &LlmCallConfig,
680    ) -> Result<LlmResponseStream> {
681        // Log the call
682        self.call_log.write().await.push(messages);
683
684        // Get next response
685        let mut index = self.call_index.write().await;
686        let responses = self.responses.read().await;
687
688        let response = responses.get(*index).cloned().unwrap_or_else(|| {
689            MockLlmResponse::text("Mock response (no more responses configured)")
690        });
691
692        *index += 1;
693        drop(index);
694        drop(responses);
695
696        // Create a stream that emits the response
697        let events = vec![
698            Ok(LlmStreamEvent::TextDelta(response.text.clone())),
699            if let Some(tool_calls) = response.tool_calls {
700                Ok(LlmStreamEvent::ToolCalls(tool_calls))
701            } else {
702                Ok(LlmStreamEvent::Done(Box::default()))
703            },
704            Ok(LlmStreamEvent::Done(Box::default())),
705        ];
706
707        Ok(Box::pin(stream::iter(events)))
708    }
709}
710
711// ============================================================================
712// InMemoryEventEmitter - Stores events in memory for testing
713// ============================================================================
714
715/// In-memory event emitter for testing
716///
717/// Stores emitted events in memory for inspection.
718/// Useful for testing and examples where you want to verify events without a database.
719///
720/// # Example
721///
722/// ```ignore
723/// use everruns_core::in_memory::InMemoryEventEmitter;
724///
725/// let emitter = InMemoryEventEmitter::new();
726///
727/// // Emit events...
728///
729/// // Check emitted events
730/// let events = emitter.events().await;
731/// assert_eq!(events.len(), 2);
732/// ```
733#[derive(Debug, Default, Clone)]
734pub struct InMemoryEventEmitter {
735    events: Arc<RwLock<Vec<Event>>>,
736    sequence: Arc<RwLock<i32>>,
737}
738
739impl InMemoryEventEmitter {
740    /// Create a new in-memory event emitter
741    pub fn new() -> Self {
742        Self {
743            events: Arc::new(RwLock::new(Vec::new())),
744            sequence: Arc::new(RwLock::new(0)),
745        }
746    }
747
748    /// Get all emitted events
749    pub async fn events(&self) -> Vec<Event> {
750        self.events.read().await.clone()
751    }
752
753    /// Get the count of emitted events
754    pub async fn event_count(&self) -> usize {
755        self.events.read().await.len()
756    }
757
758    /// Clear all events
759    pub async fn clear(&self) {
760        self.events.write().await.clear();
761        *self.sequence.write().await = 0;
762    }
763
764    /// Get events by type
765    pub async fn events_by_type(&self, event_type: &str) -> Vec<Event> {
766        self.events
767            .read()
768            .await
769            .iter()
770            .filter(|e| e.event_type == event_type)
771            .cloned()
772            .collect()
773    }
774
775    /// Get events for a specific session
776    pub async fn events_for_session(&self, session_id: Uuid) -> Vec<Event> {
777        self.events
778            .read()
779            .await
780            .iter()
781            .filter(|e| e.session_uuid() == session_id)
782            .cloned()
783            .collect()
784    }
785}
786
787#[async_trait]
788impl EventEmitter for InMemoryEventEmitter {
789    async fn emit(&self, request: EventRequest) -> Result<Event> {
790        let mut sequence = self.sequence.write().await;
791        *sequence += 1;
792        let seq = *sequence;
793        drop(sequence);
794
795        // Convert EventRequest to Event with generated id and sequence
796        let event = request.into_event(EventId::new(), seq);
797        self.events.write().await.push(event.clone());
798        Ok(event)
799    }
800}
801
802#[cfg(test)]
803mod tests {
804    use super::*;
805    use uuid::Uuid;
806
807    #[tokio::test]
808    async fn test_in_memory_message_retriever() {
809        let store = InMemoryMessageRetriever::new();
810        let session_id: SessionId = Uuid::now_v7().into();
811
812        store
813            .store(session_id, Message::user("Hello"))
814            .await
815            .unwrap();
816
817        let messages = store.load(session_id).await.unwrap();
818        assert_eq!(messages.len(), 1);
819        assert_eq!(messages[0].text(), Some("Hello"));
820    }
821
822    #[tokio::test]
823    async fn test_in_memory_message_retriever_add_and_get() {
824        let store = InMemoryMessageRetriever::new();
825        let session_id: SessionId = Uuid::now_v7().into();
826
827        // Add a message using the add method
828        let message = store
829            .add(session_id, InputMessage::user("Hello via add"))
830            .await
831            .unwrap();
832
833        // Get the message by ID
834        let retrieved = store.get(session_id, message.id).await.unwrap();
835        assert!(retrieved.is_some());
836        assert_eq!(retrieved.unwrap().text(), Some("Hello via add"));
837
838        // Get non-existent message
839        let missing = store.get(session_id, MessageId::new()).await.unwrap();
840        assert!(missing.is_none());
841    }
842
843    /// Regression test: add() must return message with ID usable for get()
844    ///
845    /// This test documents a critical invariant: the ID in the message returned by
846    /// add() must match the ID stored internally, so that get(returned_id) succeeds.
847    #[tokio::test]
848    async fn test_message_retriever_add_returns_consistent_id() {
849        let store = InMemoryMessageRetriever::new();
850        let session_id: SessionId = Uuid::now_v7().into();
851
852        // Add a message
853        let added = store
854            .add(session_id, InputMessage::user("Test consistency"))
855            .await
856            .unwrap();
857
858        // The returned message ID must be retrievable
859        let retrieved = store.get(session_id, added.id).await.unwrap();
860        assert!(
861            retrieved.is_some(),
862            "Message must be retrievable by the ID returned from add()"
863        );
864
865        // The retrieved message must have the same ID
866        let retrieved = retrieved.unwrap();
867        assert_eq!(
868            retrieved.id, added.id,
869            "Retrieved message ID must match the ID returned from add()"
870        );
871
872        // The message must also appear in load() with the same ID
873        let all_messages = store.load(session_id).await.unwrap();
874        let found = all_messages.iter().find(|m| m.id == added.id);
875        assert!(
876            found.is_some(),
877            "Message with returned ID must appear in load() results"
878        );
879    }
880
881    #[tokio::test]
882    async fn test_mock_tool_executor() {
883        let executor = MockToolExecutor::new();
884        executor
885            .set_result("get_weather", serde_json::json!({"temp": 72}))
886            .await;
887
888        let tool_call = ToolCall {
889            id: "call_1".to_string(),
890            name: "get_weather".to_string(),
891            arguments: serde_json::json!({"city": "NYC"}),
892        };
893
894        let tool_def = ToolDefinition::Builtin(crate::tool_types::BuiltinTool {
895            name: "get_weather".to_string(),
896            display_name: None,
897            description: "Get weather".to_string(),
898            parameters: serde_json::json!({}),
899            policy: crate::tool_types::ToolPolicy::Auto,
900            category: None,
901            deferrable: crate::tool_types::DeferrablePolicy::default(),
902            hints: crate::tool_types::ToolHints::default(),
903            full_parameters: None,
904        });
905
906        let result = executor.execute(&tool_call, &tool_def).await.unwrap();
907
908        assert!(result.error.is_none());
909        assert_eq!(result.result, Some(serde_json::json!({"temp": 72})));
910    }
911
912    #[tokio::test]
913    async fn test_in_memory_event_emitter() {
914        use crate::events::{EventContext, EventRequest, InputMessageData};
915
916        let emitter = InMemoryEventEmitter::new();
917        let session_id: SessionId = Uuid::now_v7().into();
918        let event_context = EventContext::empty();
919
920        // Emit an event
921        let event1 = emitter
922            .emit(EventRequest::new(
923                session_id,
924                event_context.clone(),
925                InputMessageData::new(Message::user("test1")),
926            ))
927            .await
928            .unwrap();
929        assert_eq!(event1.sequence, Some(1));
930
931        // Emit another event
932        let event2 = emitter
933            .emit(EventRequest::new(
934                session_id,
935                event_context,
936                InputMessageData::new(Message::user("test2")),
937            ))
938            .await
939            .unwrap();
940        assert_eq!(event2.sequence, Some(2));
941
942        // Check events
943        let events = emitter.events().await;
944        assert_eq!(events.len(), 2);
945        assert_eq!(emitter.event_count().await, 2);
946    }
947
948    #[tokio::test]
949    async fn test_in_memory_event_emitter_filter_by_type() {
950        use crate::events::{
951            EventContext, EventRequest, INPUT_MESSAGE, InputMessageData, REASON_STARTED,
952            ReasonStartedData,
953        };
954
955        let emitter = InMemoryEventEmitter::new();
956        let session_id: SessionId = Uuid::now_v7().into();
957        let event_context = EventContext::empty();
958
959        // Emit different event types
960        emitter
961            .emit(EventRequest::new(
962                session_id,
963                event_context.clone(),
964                InputMessageData::new(Message::user("test")),
965            ))
966            .await
967            .unwrap();
968
969        emitter
970            .emit(EventRequest::new(
971                session_id,
972                event_context,
973                ReasonStartedData {
974                    harness_id: HarnessId::from_seed(1),
975                    agent_id: Some(AgentId::new()),
976                    metadata: None,
977                },
978            ))
979            .await
980            .unwrap();
981
982        // Filter by type
983        let received_events = emitter.events_by_type(INPUT_MESSAGE).await;
984        assert_eq!(received_events.len(), 1);
985
986        let started_events = emitter.events_by_type(REASON_STARTED).await;
987        assert_eq!(started_events.len(), 1);
988    }
989
990    #[tokio::test]
991    async fn test_in_memory_event_emitter_filter_by_session() {
992        use crate::events::{EventContext, EventRequest, InputMessageData};
993
994        let emitter = InMemoryEventEmitter::new();
995        let session1: SessionId = Uuid::now_v7().into();
996        let session2: SessionId = Uuid::now_v7().into();
997
998        // Emit events for different sessions
999        let context = EventContext::empty();
1000
1001        emitter
1002            .emit(EventRequest::new(
1003                session1,
1004                context.clone(),
1005                InputMessageData::new(Message::user("session1")),
1006            ))
1007            .await
1008            .unwrap();
1009        emitter
1010            .emit(EventRequest::new(
1011                session2,
1012                context,
1013                InputMessageData::new(Message::user("session2")),
1014            ))
1015            .await
1016            .unwrap();
1017
1018        // Filter by session
1019        let session1_events = emitter.events_for_session(session1.uuid()).await;
1020        assert_eq!(session1_events.len(), 1);
1021
1022        let session2_events = emitter.events_for_session(session2.uuid()).await;
1023        assert_eq!(session2_events.len(), 1);
1024    }
1025
1026    #[tokio::test]
1027    async fn test_in_memory_event_emitter_clear() {
1028        use crate::events::{EventContext, EventRequest, InputMessageData};
1029
1030        let emitter = InMemoryEventEmitter::new();
1031        let session_id: SessionId = Uuid::now_v7().into();
1032        let event_context = EventContext::empty();
1033
1034        emitter
1035            .emit(EventRequest::new(
1036                session_id,
1037                event_context,
1038                InputMessageData::new(Message::user("test")),
1039            ))
1040            .await
1041            .unwrap();
1042
1043        assert_eq!(emitter.event_count().await, 1);
1044
1045        emitter.clear().await;
1046
1047        assert_eq!(emitter.event_count().await, 0);
1048    }
1049}
everruns_core/in_memory.rs

everruns_core/
in_memory.rs
