aether_core/testing/
utils.rs

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

aether_core/testing/
utils.rs
