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        messages: Vec<LlmMessage>,
678        _config: &LlmCallConfig,
679    ) -> Result<LlmResponseStream> {
680        // Log the call
681        self.call_log.write().await.push(messages);
682
683        // Get next response
684        let mut index = self.call_index.write().await;
685        let responses = self.responses.read().await;
686
687        let response = responses.get(*index).cloned().unwrap_or_else(|| {
688            MockLlmResponse::text("Mock response (no more responses configured)")
689        });
690
691        *index += 1;
692        drop(index);
693        drop(responses);
694
695        // Create a stream that emits the response
696        let events = vec![
697            Ok(LlmStreamEvent::TextDelta(response.text.clone())),
698            if let Some(tool_calls) = response.tool_calls {
699                Ok(LlmStreamEvent::ToolCalls(tool_calls))
700            } else {
701                Ok(LlmStreamEvent::Done(Box::default()))
702            },
703            Ok(LlmStreamEvent::Done(Box::default())),
704        ];
705
706        Ok(Box::pin(stream::iter(events)))
707    }
708}
709
710// ============================================================================
711// InMemoryEventEmitter - Stores events in memory for testing
712// ============================================================================
713
714/// In-memory event emitter for testing
715///
716/// Stores emitted events in memory for inspection.
717/// Useful for testing and examples where you want to verify events without a database.
718///
719/// # Example
720///
721/// ```ignore
722/// use everruns_core::in_memory::InMemoryEventEmitter;
723///
724/// let emitter = InMemoryEventEmitter::new();
725///
726/// // Emit events...
727///
728/// // Check emitted events
729/// let events = emitter.events().await;
730/// assert_eq!(events.len(), 2);
731/// ```
732#[derive(Debug, Default, Clone)]
733pub struct InMemoryEventEmitter {
734    events: Arc<RwLock<Vec<Event>>>,
735    sequence: Arc<RwLock<i32>>,
736}
737
738impl InMemoryEventEmitter {
739    /// Create a new in-memory event emitter
740    pub fn new() -> Self {
741        Self {
742            events: Arc::new(RwLock::new(Vec::new())),
743            sequence: Arc::new(RwLock::new(0)),
744        }
745    }
746
747    /// Get all emitted events
748    pub async fn events(&self) -> Vec<Event> {
749        self.events.read().await.clone()
750    }
751
752    /// Get the count of emitted events
753    pub async fn event_count(&self) -> usize {
754        self.events.read().await.len()
755    }
756
757    /// Clear all events
758    pub async fn clear(&self) {
759        self.events.write().await.clear();
760        *self.sequence.write().await = 0;
761    }
762
763    /// Get events by type
764    pub async fn events_by_type(&self, event_type: &str) -> Vec<Event> {
765        self.events
766            .read()
767            .await
768            .iter()
769            .filter(|e| e.event_type == event_type)
770            .cloned()
771            .collect()
772    }
773
774    /// Get events for a specific session
775    pub async fn events_for_session(&self, session_id: Uuid) -> Vec<Event> {
776        self.events
777            .read()
778            .await
779            .iter()
780            .filter(|e| e.session_uuid() == session_id)
781            .cloned()
782            .collect()
783    }
784}
785
786#[async_trait]
787impl EventEmitter for InMemoryEventEmitter {
788    async fn emit(&self, request: EventRequest) -> Result<Event> {
789        let mut sequence = self.sequence.write().await;
790        *sequence += 1;
791        let seq = *sequence;
792        drop(sequence);
793
794        // Convert EventRequest to Event with generated id and sequence
795        let event = request.into_event(EventId::new(), seq);
796        self.events.write().await.push(event.clone());
797        Ok(event)
798    }
799}
800
801#[cfg(test)]
802mod tests {
803    use super::*;
804    use uuid::Uuid;
805
806    #[tokio::test]
807    async fn test_in_memory_message_retriever() {
808        let store = InMemoryMessageRetriever::new();
809        let session_id: SessionId = Uuid::now_v7().into();
810
811        store
812            .store(session_id, Message::user("Hello"))
813            .await
814            .unwrap();
815
816        let messages = store.load(session_id).await.unwrap();
817        assert_eq!(messages.len(), 1);
818        assert_eq!(messages[0].text(), Some("Hello"));
819    }
820
821    #[tokio::test]
822    async fn test_in_memory_message_retriever_add_and_get() {
823        let store = InMemoryMessageRetriever::new();
824        let session_id: SessionId = Uuid::now_v7().into();
825
826        // Add a message using the add method
827        let message = store
828            .add(session_id, InputMessage::user("Hello via add"))
829            .await
830            .unwrap();
831
832        // Get the message by ID
833        let retrieved = store.get(session_id, message.id).await.unwrap();
834        assert!(retrieved.is_some());
835        assert_eq!(retrieved.unwrap().text(), Some("Hello via add"));
836
837        // Get non-existent message
838        let missing = store.get(session_id, MessageId::new()).await.unwrap();
839        assert!(missing.is_none());
840    }
841
842    /// Regression test: add() must return message with ID usable for get()
843    ///
844    /// This test documents a critical invariant: the ID in the message returned by
845    /// add() must match the ID stored internally, so that get(returned_id) succeeds.
846    #[tokio::test]
847    async fn test_message_retriever_add_returns_consistent_id() {
848        let store = InMemoryMessageRetriever::new();
849        let session_id: SessionId = Uuid::now_v7().into();
850
851        // Add a message
852        let added = store
853            .add(session_id, InputMessage::user("Test consistency"))
854            .await
855            .unwrap();
856
857        // The returned message ID must be retrievable
858        let retrieved = store.get(session_id, added.id).await.unwrap();
859        assert!(
860            retrieved.is_some(),
861            "Message must be retrievable by the ID returned from add()"
862        );
863
864        // The retrieved message must have the same ID
865        let retrieved = retrieved.unwrap();
866        assert_eq!(
867            retrieved.id, added.id,
868            "Retrieved message ID must match the ID returned from add()"
869        );
870
871        // The message must also appear in load() with the same ID
872        let all_messages = store.load(session_id).await.unwrap();
873        let found = all_messages.iter().find(|m| m.id == added.id);
874        assert!(
875            found.is_some(),
876            "Message with returned ID must appear in load() results"
877        );
878    }
879
880    #[tokio::test]
881    async fn test_mock_tool_executor() {
882        let executor = MockToolExecutor::new();
883        executor
884            .set_result("get_weather", serde_json::json!({"temp": 72}))
885            .await;
886
887        let tool_call = ToolCall {
888            id: "call_1".to_string(),
889            name: "get_weather".to_string(),
890            arguments: serde_json::json!({"city": "NYC"}),
891        };
892
893        let tool_def = ToolDefinition::Builtin(crate::tool_types::BuiltinTool {
894            name: "get_weather".to_string(),
895            display_name: None,
896            description: "Get weather".to_string(),
897            parameters: serde_json::json!({}),
898            policy: crate::tool_types::ToolPolicy::Auto,
899            category: None,
900            deferrable: crate::tool_types::DeferrablePolicy::default(),
901            hints: crate::tool_types::ToolHints::default(),
902            full_parameters: None,
903        });
904
905        let result = executor.execute(&tool_call, &tool_def).await.unwrap();
906
907        assert!(result.error.is_none());
908        assert_eq!(result.result, Some(serde_json::json!({"temp": 72})));
909    }
910
911    #[tokio::test]
912    async fn test_in_memory_event_emitter() {
913        use crate::events::{EventContext, EventRequest, InputMessageData};
914
915        let emitter = InMemoryEventEmitter::new();
916        let session_id: SessionId = Uuid::now_v7().into();
917        let event_context = EventContext::empty();
918
919        // Emit an event
920        let event1 = emitter
921            .emit(EventRequest::new(
922                session_id,
923                event_context.clone(),
924                InputMessageData::new(Message::user("test1")),
925            ))
926            .await
927            .unwrap();
928        assert_eq!(event1.sequence, Some(1));
929
930        // Emit another event
931        let event2 = emitter
932            .emit(EventRequest::new(
933                session_id,
934                event_context,
935                InputMessageData::new(Message::user("test2")),
936            ))
937            .await
938            .unwrap();
939        assert_eq!(event2.sequence, Some(2));
940
941        // Check events
942        let events = emitter.events().await;
943        assert_eq!(events.len(), 2);
944        assert_eq!(emitter.event_count().await, 2);
945    }
946
947    #[tokio::test]
948    async fn test_in_memory_event_emitter_filter_by_type() {
949        use crate::events::{
950            EventContext, EventRequest, INPUT_MESSAGE, InputMessageData, REASON_STARTED,
951            ReasonStartedData,
952        };
953
954        let emitter = InMemoryEventEmitter::new();
955        let session_id: SessionId = Uuid::now_v7().into();
956        let event_context = EventContext::empty();
957
958        // Emit different event types
959        emitter
960            .emit(EventRequest::new(
961                session_id,
962                event_context.clone(),
963                InputMessageData::new(Message::user("test")),
964            ))
965            .await
966            .unwrap();
967
968        emitter
969            .emit(EventRequest::new(
970                session_id,
971                event_context,
972                ReasonStartedData {
973                    harness_id: HarnessId::from_seed(1),
974                    agent_id: Some(AgentId::new()),
975                    metadata: None,
976                },
977            ))
978            .await
979            .unwrap();
980
981        // Filter by type
982        let received_events = emitter.events_by_type(INPUT_MESSAGE).await;
983        assert_eq!(received_events.len(), 1);
984
985        let started_events = emitter.events_by_type(REASON_STARTED).await;
986        assert_eq!(started_events.len(), 1);
987    }
988
989    #[tokio::test]
990    async fn test_in_memory_event_emitter_filter_by_session() {
991        use crate::events::{EventContext, EventRequest, InputMessageData};
992
993        let emitter = InMemoryEventEmitter::new();
994        let session1: SessionId = Uuid::now_v7().into();
995        let session2: SessionId = Uuid::now_v7().into();
996
997        // Emit events for different sessions
998        let context = EventContext::empty();
999
1000        emitter
1001            .emit(EventRequest::new(
1002                session1,
1003                context.clone(),
1004                InputMessageData::new(Message::user("session1")),
1005            ))
1006            .await
1007            .unwrap();
1008        emitter
1009            .emit(EventRequest::new(
1010                session2,
1011                context,
1012                InputMessageData::new(Message::user("session2")),
1013            ))
1014            .await
1015            .unwrap();
1016
1017        // Filter by session
1018        let session1_events = emitter.events_for_session(session1.uuid()).await;
1019        assert_eq!(session1_events.len(), 1);
1020
1021        let session2_events = emitter.events_for_session(session2.uuid()).await;
1022        assert_eq!(session2_events.len(), 1);
1023    }
1024
1025    #[tokio::test]
1026    async fn test_in_memory_event_emitter_clear() {
1027        use crate::events::{EventContext, EventRequest, InputMessageData};
1028
1029        let emitter = InMemoryEventEmitter::new();
1030        let session_id: SessionId = Uuid::now_v7().into();
1031        let event_context = EventContext::empty();
1032
1033        emitter
1034            .emit(EventRequest::new(
1035                session_id,
1036                event_context,
1037                InputMessageData::new(Message::user("test")),
1038            ))
1039            .await
1040            .unwrap();
1041
1042        assert_eq!(emitter.event_count().await, 1);
1043
1044        emitter.clear().await;
1045
1046        assert_eq!(emitter.event_count().await, 0);
1047    }
1048}
everruns_core/in_memory.rs

everruns_core/
in_memory.rs
