Skip to main content

aether_core/testing/
fake_agent_observer.rs

1use crate::events::{AgentEvent, AgentObserver};
2use std::sync::{Arc, Mutex};
3
4/// In-memory [`AgentObserver`] that records every event it receives, for
5/// asserting on the stream an agent emits.
6#[derive(Default)]
7pub struct FakeAgentObserver {
8    events: Arc<Mutex<Vec<AgentEvent>>>,
9    system_prompts: Arc<Mutex<Vec<String>>>,
10}
11
12impl FakeAgentObserver {
13    pub fn new() -> Self {
14        Self::default()
15    }
16
17    /// Shared handle to the recorded events; clones observe future events too.
18    pub fn events(&self) -> Arc<Mutex<Vec<AgentEvent>>> {
19        Arc::clone(&self.events)
20    }
21
22    /// Shared handle to the system prompts reported for each LLM request.
23    pub fn system_prompts(&self) -> Arc<Mutex<Vec<String>>> {
24        Arc::clone(&self.system_prompts)
25    }
26}
27
28impl AgentObserver for FakeAgentObserver {
29    fn on_event(&mut self, message: &AgentEvent) {
30        self.events.lock().unwrap().push(message.clone());
31    }
32
33    fn on_system_prompt(&mut self, prompt: &str) {
34        self.system_prompts.lock().unwrap().push(prompt.to_string());
35    }
36}