aether_core/testing/
utils.rs

1use std::collections::BTreeMap;
2use std::sync::{Arc, Mutex};
3use std::time::Duration;
4use tokio::sync::{Notify, mpsc};
5
6use crate::context::CompactionConfig;
7use crate::core::{AgentError, Prompt, RetryConfig, agent};
8use crate::events::{
9    AgentCommand, AgentEvent, AgentObserver, Command, ContextEvent, ToolEvent, TurnEvent, UserCommand,
10};
11use crate::mcp::mcp;
12use crate::testing::fake_mcp::fake_mcp;
13use crate::testing::{AgentTrace, FakeAgentObserver, FakeMcpServer};
14use llm::{ChatMessage, Context, LlmError, LlmModel, LlmResponse, ModelSettings, StreamingModelProvider};
15
16use llm::testing::FakeLlmProvider;
17
18pub async fn drain_until(
19    receiver: &mut mpsc::Receiver<AgentEvent>,
20    predicate: impl Fn(&AgentEvent) -> bool,
21) -> Vec<AgentEvent> {
22    let mut events = Vec::new();
23    while let Some(event) = receiver.recv().await {
24        let matched = predicate(&event);
25        events.push(event);
26        if matched {
27            return events;
28        }
29    }
30    panic!("agent event channel closed before predicate matched");
31}
32
33pub fn content_events(events: Vec<AgentEvent>) -> Vec<AgentEvent> {
34    events
35        .into_iter()
36        .filter(|event| {
37            !matches!(
38                event,
39                AgentEvent::Turn(
40                    TurnEvent::Started { .. } | TurnEvent::LlmCallStarted { .. } | TurnEvent::LlmCallEnded { .. }
41                ) | AgentEvent::Tool(ToolEvent::ExecutionStarted { .. } | ToolEvent::DefinitionsUpdated { .. })
42            )
43        })
44        .collect()
45}
46
47pub fn mcp_instructions(entries: &[(&str, &str)]) -> BTreeMap<String, String> {
48    entries.iter().map(|(k, v)| ((*k).to_string(), (*v).to_string())).collect()
49}
50
51pub fn test_agent() -> TestAgentBuilder {
52    TestAgentBuilder::new()
53}
54
55/// An ordered interaction with a test agent.
56pub enum TestAgentStep {
57    Send(Command),
58    WaitFor(Box<dyn Fn(&AgentEvent) -> bool + Send>),
59    /// Run an arbitrary side effect (e.g. releasing a paused LLM stream) between steps.
60    Perform(Box<dyn FnOnce() + Send>),
61}
62
63impl TestAgentStep {
64    pub fn send(command: Command) -> Self {
65        Self::Send(command)
66    }
67
68    pub fn user_text(text: impl Into<String>) -> Self {
69        Self::send(Command::UserCommand(UserCommand::Text { content: vec![llm::ContentBlock::text(text.into())] }))
70    }
71
72    pub fn cancel() -> Self {
73        Self::send(Command::UserCommand(UserCommand::Cancel))
74    }
75
76    pub fn switch_model(provider: impl StreamingModelProvider + 'static) -> Self {
77        Self::send(Command::AgentCommand(AgentCommand::SwitchModel(Box::new(provider))))
78    }
79
80    pub fn replace_conversation(messages: Vec<ChatMessage>) -> Self {
81        Self::send(Command::AgentCommand(AgentCommand::ReplaceConversation(messages)))
82    }
83
84    pub fn perform(action: impl FnOnce() + Send + 'static) -> Self {
85        Self::Perform(Box::new(action))
86    }
87
88    pub fn wait_for(predicate: impl Fn(&AgentEvent) -> bool + Send + 'static) -> Self {
89        Self::WaitFor(Box::new(predicate))
90    }
91
92    pub fn wait_for_turn_end() -> Self {
93        Self::wait_for(|event| matches!(event, AgentEvent::Turn(TurnEvent::Ended { .. })))
94    }
95
96    pub fn wait_for_compaction_start() -> Self {
97        Self::wait_for(|event| matches!(event, AgentEvent::Context(ContextEvent::CompactionStarted { .. })))
98    }
99
100    pub fn wait_for_retry(attempt: u32) -> Self {
101        Self::wait_for(
102            move |event| matches!(event, AgentEvent::Turn(TurnEvent::RetryScheduled { attempt: actual, .. }) if *actual == attempt),
103        )
104    }
105}
106
107/// A fluent sequence of commands and synchronization points for a test agent.
108#[derive(Default)]
109pub struct TestScenario {
110    steps: Vec<TestAgentStep>,
111}
112
113impl TestScenario {
114    pub fn new() -> Self {
115        Self::default()
116    }
117
118    pub fn send(mut self, command: Command) -> Self {
119        self.steps.push(TestAgentStep::send(command));
120        self
121    }
122
123    pub fn user_text(mut self, text: impl Into<String>) -> Self {
124        self.steps.push(TestAgentStep::user_text(text));
125        self
126    }
127
128    pub fn cancel(mut self) -> Self {
129        self.steps.push(TestAgentStep::cancel());
130        self
131    }
132
133    pub fn switch_model(mut self, provider: impl StreamingModelProvider + 'static) -> Self {
134        self.steps.push(TestAgentStep::switch_model(provider));
135        self
136    }
137
138    pub fn replace_conversation(mut self, messages: Vec<ChatMessage>) -> Self {
139        self.steps.push(TestAgentStep::replace_conversation(messages));
140        self
141    }
142
143    pub fn wait_for(mut self, predicate: impl Fn(&AgentEvent) -> bool + Send + 'static) -> Self {
144        self.steps.push(TestAgentStep::wait_for(predicate));
145        self
146    }
147
148    pub fn wait_for_turn_end(mut self) -> Self {
149        self.steps.push(TestAgentStep::wait_for_turn_end());
150        self
151    }
152
153    pub fn wait_for_compaction_start(mut self) -> Self {
154        self.steps.push(TestAgentStep::wait_for_compaction_start());
155        self
156    }
157
158    pub fn wait_for_retry(mut self, attempt: u32) -> Self {
159        self.steps.push(TestAgentStep::wait_for_retry(attempt));
160        self
161    }
162
163    /// Run an arbitrary side effect between scenario steps, e.g. releasing a
164    /// paused LLM stream so a queued message can be injected mid-turn.
165    pub fn perform(mut self, action: impl FnOnce() + Send + 'static) -> Self {
166        self.steps.push(TestAgentStep::perform(action));
167        self
168    }
169}
170
171impl From<Vec<TestAgentStep>> for TestScenario {
172    fn from(steps: Vec<TestAgentStep>) -> Self {
173        Self { steps }
174    }
175}
176
177/// Result of running a test agent, including messages and captured contexts.
178pub struct TestAgentResult {
179    pub messages: Vec<AgentEvent>,
180    pub captured_contexts: Arc<Mutex<Vec<Context>>>,
181}
182
183/// Error returned by [`TestAgentBuilder::run`], [`TestAgentBuilder::run_trace`], and
184/// [`TestAgentBuilder::run_with_context`].
185#[derive(Debug, thiserror::Error)]
186pub enum TestAgentError {
187    #[error(transparent)]
188    Agent(#[from] AgentError),
189    #[error("failed to send command to the test agent: {0}")]
190    SendCommand(#[from] mpsc::error::SendError<Command>),
191}
192
193pub type TestResult<T> = std::result::Result<T, TestAgentError>;
194
195struct ProviderTestConfig {
196    responses: Vec<Vec<Result<LlmResponse, LlmError>>>,
197    model: Option<LlmModel>,
198    context_window: Option<u32>,
199    pause: Option<(usize, usize, Arc<Notify>)>,
200}
201
202struct AgentTestConfig {
203    context_window_override: Option<u32>,
204    timeout: Option<Duration>,
205    max_auto_continues: Option<u32>,
206    retry_config: Option<RetryConfig>,
207    observers: Vec<Box<dyn AgentObserver>>,
208    include_fake_mcp: bool,
209    initial_messages: Vec<ChatMessage>,
210    system_prompt: Option<Prompt>,
211    compaction: Option<CompactionConfig>,
212    model_settings: Option<ModelSettings>,
213}
214
215enum TestExecution {
216    CommandsUntilTurnEnd(Vec<Command>),
217    Scenario(TestScenario),
218}
219
220pub struct TestAgentBuilder {
221    provider: ProviderTestConfig,
222    agent: AgentTestConfig,
223    execution: Option<TestExecution>,
224}
225
226impl Default for TestAgentBuilder {
227    fn default() -> Self {
228        Self::new()
229    }
230}
231
232impl TestAgentBuilder {
233    pub fn new() -> Self {
234        Self {
235            provider: ProviderTestConfig { responses: Vec::new(), model: None, context_window: None, pause: None },
236            agent: AgentTestConfig {
237                context_window_override: None,
238                timeout: None,
239                max_auto_continues: None,
240                retry_config: None,
241                observers: Vec::new(),
242                include_fake_mcp: true,
243                initial_messages: Vec::new(),
244                system_prompt: None,
245                compaction: None,
246                model_settings: None,
247            },
248            execution: None,
249        }
250    }
251
252    pub fn commands(self, commands: Vec<Command>) -> Self {
253        self.with_execution(TestExecution::CommandsUntilTurnEnd(commands))
254    }
255
256    pub fn scenario(self, scenario: impl Into<TestScenario>) -> Self {
257        self.with_execution(TestExecution::Scenario(scenario.into()))
258    }
259
260    pub fn user_text(self, text: &str) -> Self {
261        self.commands(vec![Command::UserCommand(UserCommand::Text { content: vec![llm::ContentBlock::text(text)] })])
262    }
263
264    pub fn llm_responses(mut self, llm_responses: &[Vec<LlmResponse>]) -> Self {
265        self.provider.responses = llm_responses.iter().map(|turn| turn.iter().cloned().map(Ok).collect()).collect();
266        self
267    }
268
269    pub fn llm_result_responses(mut self, llm_responses: &[Vec<Result<LlmResponse, LlmError>>]) -> Self {
270        self.provider.responses = Vec::from(llm_responses);
271        self
272    }
273
274    pub fn model(mut self, model: LlmModel) -> Self {
275        self.provider.model = Some(model);
276        self
277    }
278
279    pub fn provider_context_window(mut self, window: Option<u32>) -> Self {
280        self.provider.context_window = window;
281        self
282    }
283
284    pub fn context_window_override(mut self, window: u32) -> Self {
285        self.agent.context_window_override = Some(window);
286        self
287    }
288
289    pub fn tool_timeout(mut self, timeout: Duration) -> Self {
290        self.agent.timeout = Some(timeout);
291        self
292    }
293
294    pub fn max_auto_continues(mut self, max: u32) -> Self {
295        self.agent.max_auto_continues = Some(max);
296        self
297    }
298
299    pub fn retry_config(mut self, config: RetryConfig) -> Self {
300        self.agent.retry_config = Some(config);
301        self
302    }
303
304    /// Run without the default fake MCP server when the scenario does not exercise tools.
305    pub fn without_mcp(mut self) -> Self {
306        self.agent.include_fake_mcp = false;
307        self
308    }
309
310    /// Pre-populate the context with conversation history.
311    pub fn messages(mut self, messages: Vec<ChatMessage>) -> Self {
312        self.agent.initial_messages = messages;
313        self
314    }
315
316    /// Set the system prompt.
317    pub fn system_prompt(mut self, prompt: Prompt) -> Self {
318        self.agent.system_prompt = Some(prompt);
319        self
320    }
321
322    /// Configure context compaction settings.
323    pub fn compaction_config(mut self, config: CompactionConfig) -> Self {
324        self.agent.compaction = Some(config);
325        self
326    }
327
328    /// Set the model settings applied to every LLM call.
329    pub fn model_settings(mut self, settings: ModelSettings) -> Self {
330        self.agent.model_settings = Some(settings);
331        self
332    }
333
334    /// Pause the fake LLM stream at `turn_index` / `chunk_index` until
335    /// `release.notify_one()` is called. Used for deterministic timing tests.
336    pub fn pause_turn_after(mut self, turn_index: usize, chunk_index: usize, release: Arc<Notify>) -> Self {
337        self.provider.pause = Some((turn_index, chunk_index, release));
338        self
339    }
340
341    /// Attach an observer of the test agent's event stream.
342    pub fn observer(mut self, observer: Box<dyn AgentObserver>) -> Self {
343        self.agent.observers.push(observer);
344        self
345    }
346
347    pub async fn run(self) -> TestResult<Vec<AgentEvent>> {
348        let result = self.run_with_context().await?;
349        Ok(result.messages)
350    }
351
352    /// Runs the test agent with a recording observer attached and returns the
353    /// full event trace, including internal events.
354    pub async fn run_trace(self) -> TestResult<AgentTrace> {
355        let observer = FakeAgentObserver::new();
356        let events = observer.events();
357        self.observer(Box::new(observer)).run().await?;
358        Ok(AgentTrace::from_observer_events(&events))
359    }
360
361    /// Runs the test agent and returns both messages and captured contexts.
362    ///
363    /// Use this when you need to verify what context was passed to the LLM,
364    /// for example when testing that file attachments are properly formatted.
365    pub async fn run_with_context(self) -> TestResult<TestAgentResult> {
366        let Self { provider, agent: config, execution } = self;
367        let mut llm = FakeLlmProvider::from_results(provider.responses).with_context_window(provider.context_window);
368        if let Some(model) = provider.model {
369            llm = llm.with_model(model);
370        }
371        if let Some((turn_index, chunk_index, release)) = provider.pause {
372            llm = llm.pause_turn_after(turn_index, chunk_index, release);
373        }
374        let captured_contexts = llm.captured_contexts();
375
376        let mut mcp_spawn = if config.include_fake_mcp {
377            Some(
378                mcp("/workspace")
379                    .with_servers(vec![fake_mcp("test", FakeMcpServer::new())])
380                    .spawn()
381                    .await
382                    .map_err(AgentError::from)?,
383            )
384        } else {
385            None
386        };
387
388        let mut builder = agent(llm);
389        if let Some(spawn) = &mut mcp_spawn {
390            let snapshot = spawn.block_until_ready().await.expect("bootstrap completes");
391            builder = builder.tools(spawn.command_tx.clone(), snapshot.tool_definitions);
392        }
393        if let Some(timeout) = config.timeout {
394            builder = builder.tool_timeout(timeout);
395        }
396        if let Some(max) = config.max_auto_continues {
397            builder = builder.max_auto_continues(max);
398        }
399        if let Some(retry) = config.retry_config {
400            builder = builder.retry(retry);
401        } else {
402            builder = builder.retry(RetryConfig::disabled());
403        }
404        if let Some(prompt) = config.system_prompt {
405            builder = builder.system_prompt(prompt);
406        }
407        if let Some(compaction) = config.compaction {
408            builder = builder.compaction(compaction);
409        }
410        if let Some(settings) = config.model_settings {
411            builder = builder.model_settings(settings);
412        }
413        builder = builder.context_window(config.context_window_override);
414        if !config.initial_messages.is_empty() {
415            builder = builder.messages(config.initial_messages);
416        }
417        for observer in config.observers {
418            builder = builder.observer(observer);
419        }
420
421        let steps = match execution.expect("test agent requires commands(), user_text(), or scenario()") {
422            TestExecution::CommandsUntilTurnEnd(commands) => {
423                assert!(!commands.is_empty(), "commands() requires at least one command");
424                let mut steps = commands.into_iter().map(TestAgentStep::send).collect::<Vec<_>>();
425                steps.push(TestAgentStep::wait_for_turn_end());
426                steps
427            }
428            TestExecution::Scenario(scenario) => {
429                assert!(!scenario.steps.is_empty(), "scenario() requires at least one step");
430                scenario.steps
431            }
432        };
433        let (tx, mut rx, handle) = builder.spawn().await?;
434        let mut messages = Vec::new();
435
436        for step in steps {
437            match step {
438                TestAgentStep::Send(command) => tx.send(command).await?,
439                TestAgentStep::WaitFor(predicate) => loop {
440                    let message = rx.recv().await.expect("agent event channel closed before scenario step matched");
441                    let matched = predicate(&message);
442                    messages.push(message);
443                    if matched {
444                        break;
445                    }
446                },
447                TestAgentStep::Perform(action) => action(),
448            }
449        }
450        drop(tx);
451
452        handle.await_completion().await;
453
454        Ok(TestAgentResult { messages, captured_contexts })
455    }
456
457    fn with_execution(mut self, execution: TestExecution) -> Self {
458        assert!(self.execution.is_none(), "commands(), user_text(), and scenario() are mutually exclusive");
459        self.execution = Some(execution);
460        self
461    }
462}
aether_core/testing/utils.rs

aether_core/testing/
utils.rs
