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;
8use std::time::Duration;
9use tokio::sync::{Notify, mpsc};
10
11use crate::context::CompactionConfig;
12use crate::core::{AgentBuilder, 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                    | AgentEvent::Context(ContextEvent::ReplyRecorded { .. })
51            )
52        })
53        .collect()
54}
55
56pub async fn ready_catalog(builder: McpBuilder) -> McpCatalog {
57    let catalog = builder.spawn().expect("the MCP catalog spawns");
58    catalog.ready().await;
59    catalog
60}
61
62pub fn mcp_instructions(entries: &[(&str, &str)]) -> McpInstructions {
63    let servers = entries.iter().map(|(k, v)| ((*k).to_string(), (*v).to_string())).collect();
64    McpInstructions { servers, progressive_discovery: false }
65}
66
67/// Millisecond-scale retry delays so retry tests run fast under virtual time.
68pub fn fast_retry(max_attempts: u32) -> RetryConfig {
69    RetryConfig { max_attempts, base_delay: Duration::from_millis(1), max_delay: Duration::from_millis(5) }
70}
71
72pub fn test_agent() -> TestAgentBuilder {
73    TestAgentBuilder::default()
74}
75
76/// A fluent sequence of commands and synchronization points for a test agent.
77#[derive(Default)]
78pub struct TestScenario {
79    steps: Vec<Step>,
80}
81
82impl TestScenario {
83    pub fn new() -> Self {
84        Self::default()
85    }
86
87    pub fn send(self, command: Command) -> Self {
88        self.step(Step::Send(command))
89    }
90
91    pub fn user_text(self, text: impl Into<String>) -> Self {
92        self.send(Command::text(text))
93    }
94
95    pub fn cancel(self) -> Self {
96        self.send(Command::UserCommand(UserCommand::Cancel))
97    }
98
99    pub fn switch_model(self, provider: impl StreamingModelProvider + 'static) -> Self {
100        self.send(Command::AgentCommand(AgentCommand::SwitchModel(Box::new(provider))))
101    }
102
103    pub fn replace_conversation(self, messages: Vec<ChatMessage>) -> Self {
104        self.send(Command::AgentCommand(AgentCommand::ReplaceConversation(messages)))
105    }
106
107    pub fn wait_for(self, predicate: impl Fn(&AgentEvent) -> bool + Send + 'static) -> Self {
108        self.step(Step::WaitFor(Box::new(predicate)))
109    }
110
111    pub fn wait_for_turn_end(self) -> Self {
112        self.wait_for(|event| matches!(event, AgentEvent::Turn(TurnEvent::Ended { .. })))
113    }
114
115    pub fn wait_for_compaction_start(self) -> Self {
116        self.wait_for(|event| matches!(event, AgentEvent::Context(ContextEvent::CompactionStarted { .. })))
117    }
118
119    pub fn wait_for_retry(self, attempt: u32) -> Self {
120        self.wait_for(
121            move |event| matches!(event, AgentEvent::Turn(TurnEvent::RetryScheduled { attempt: actual, .. }) if *actual == attempt),
122        )
123    }
124
125    /// Run an arbitrary side effect between scenario steps, e.g. releasing a
126    /// paused LLM stream so a queued message can be injected mid-turn.
127    pub fn perform(self, action: impl FnOnce() + Send + 'static) -> Self {
128        self.step(Step::Perform(Box::new(action)))
129    }
130
131    fn step(mut self, step: Step) -> Self {
132        self.steps.push(step);
133        self
134    }
135}
136
137/// Result of running a test agent, including messages, captured contexts, and the MCP elicitations it answered.
138pub struct TestAgentResult {
139    pub messages: Vec<AgentEvent>,
140    pub contexts: Vec<Context>,
141    pub elicitations: Vec<CapturedElicitation>,
142}
143
144/// Error returned by [`TestAgentBuilder::run`], [`TestAgentBuilder::run_trace`], and
145/// [`TestAgentBuilder::run_with_context`].
146#[derive(Debug, thiserror::Error)]
147pub enum TestAgentError {
148    #[error(transparent)]
149    Agent(#[from] AgentError),
150    #[error("failed to send command to the test agent: {0}")]
151    SendCommand(#[from] mpsc::error::SendError<Command>),
152}
153
154pub type TestResult<T> = std::result::Result<T, TestAgentError>;
155
156#[derive(Default)]
157pub struct TestAgentBuilder {
158    responses: Vec<Vec<Result<LlmResponse, LlmError>>>,
159    llm_config: Vec<Configure<FakeLlmProvider>>,
160    agent_config: Vec<Configure<AgentBuilder>>,
161    /// `None` connects a single fake server named `test`.
162    mcp_servers: Option<Vec<ServerSpec>>,
163    elicitations: Option<ElicitationScriptBuilder>,
164    scenario: Option<TestScenario>,
165}
166
167impl TestAgentBuilder {
168    pub fn commands(self, commands: Vec<Command>) -> Self {
169        assert!(!commands.is_empty(), "commands() requires at least one command");
170        let scenario = commands.into_iter().fold(TestScenario::new(), TestScenario::send);
171        self.scenario(scenario.wait_for_turn_end())
172    }
173
174    pub fn scenario(mut self, scenario: TestScenario) -> Self {
175        assert!(!scenario.steps.is_empty(), "scenario() requires at least one step");
176        assert!(self.scenario.is_none(), "commands(), user_text(), and scenario() are mutually exclusive");
177        self.scenario = Some(scenario);
178        self
179    }
180
181    pub fn user_text(self, text: &str) -> Self {
182        self.commands(vec![Command::text(text)])
183    }
184
185    pub fn llm_responses(mut self, llm_responses: &[Vec<LlmResponse>]) -> Self {
186        self.responses = llm_responses.iter().map(|turn| turn.iter().cloned().map(Ok).collect()).collect();
187        self
188    }
189
190    pub fn llm_result_responses(mut self, llm_responses: &[Vec<Result<LlmResponse, LlmError>>]) -> Self {
191        self.responses = Vec::from(llm_responses);
192        self
193    }
194
195    pub fn model(self, model: LlmModel) -> Self {
196        self.configure_llm(|llm| llm.with_model(model))
197    }
198
199    pub fn provider_context_window(self, window: Option<u32>) -> Self {
200        self.configure_llm(move |llm| llm.with_context_window(window))
201    }
202
203    pub fn context_window_override(self, window: u32) -> Self {
204        self.configure(move |agent| agent.context_window(Some(window)))
205    }
206
207    pub fn tool_timeout(self, timeout: Duration) -> Self {
208        self.configure(move |agent| agent.tool_timeout(timeout))
209    }
210
211    pub fn max_auto_continues(self, max: u32) -> Self {
212        self.configure(move |agent| agent.max_auto_continues(max))
213    }
214
215    pub fn retry_config(self, config: RetryConfig) -> Self {
216        self.configure(move |agent| agent.retry(config))
217    }
218
219    /// Run without the default fake MCP server when the scenario does not exercise tools.
220    pub fn without_mcp(mut self) -> Self {
221        self.mcp_servers = Some(Vec::new());
222        self
223    }
224
225    pub fn mcp_server(mut self, name: &str, server: impl Into<McpServer>) -> Self {
226        self.mcp_servers.get_or_insert_with(Vec::new).push(ServerSpec::in_process(name, server));
227        self
228    }
229
230    pub fn elicitation_response(mut self, response: ElicitResult) -> Self {
231        self.elicitations = Some(self.elicitations.take().unwrap_or_default().response(response));
232        self
233    }
234
235    pub fn on_url_elicitation<F, Fut>(mut self, handler: F) -> Self
236    where
237        F: Fn(String, String) -> Fut + Send + Sync + 'static,
238        Fut: Future<Output = ()> + Send + 'static,
239    {
240        self.elicitations = Some(self.elicitations.take().unwrap_or_default().on_url(handler));
241        self
242    }
243
244    /// Pre-populate the context with conversation history.
245    pub fn messages(self, messages: Vec<ChatMessage>) -> Self {
246        self.configure(|agent| agent.messages(messages))
247    }
248
249    /// Set the system prompt. Multiple prompts are concatenated with double newlines.
250    pub fn system_prompt(self, prompt: Prompt) -> Self {
251        self.configure(|agent| agent.system_prompt(prompt))
252    }
253
254    /// Set the session affinity key attached to every LLM call; defaults to a
255    /// fresh random key per agent.
256    pub fn session_affinity_key(self, key: impl Into<String>) -> Self {
257        let key = key.into();
258        self.configure(|agent| agent.session_affinity_key(key))
259    }
260
261    /// Configure context compaction settings.
262    pub fn compaction_config(self, config: CompactionConfig) -> Self {
263        self.configure(|agent| agent.compaction(config))
264    }
265
266    /// Set the model settings applied to every LLM call.
267    pub fn model_settings(self, settings: ModelSettings) -> Self {
268        self.configure(|agent| agent.model_settings(settings))
269    }
270
271    /// Pause the fake LLM stream at `turn_index` / `chunk_index` until
272    /// `release.notify_one()` is called. Used for deterministic timing tests.
273    pub fn pause_turn_after(self, turn_index: usize, chunk_index: usize, release: Arc<Notify>) -> Self {
274        self.configure_llm(move |llm| llm.pause_turn_after(turn_index, chunk_index, release))
275    }
276
277    /// Attach an observer of the test agent's event stream.
278    pub fn observer(self, observer: Box<dyn AgentObserver>) -> Self {
279        self.configure(|agent| agent.observer(observer))
280    }
281
282    pub async fn run(self) -> TestResult<Vec<AgentEvent>> {
283        let result = self.run_with_context().await?;
284        Ok(result.messages)
285    }
286
287    /// Runs the test agent with a recording observer attached and returns the
288    /// full event trace, including internal events.
289    pub async fn run_trace(self) -> TestResult<AgentTrace> {
290        let observer = FakeAgentObserver::new();
291        let events = observer.events();
292        self.observer(Box::new(observer)).run().await?;
293        Ok(AgentTrace::from_observer_events(&events))
294    }
295
296    /// Runs the test agent and returns both messages and captured contexts.
297    ///
298    /// Use this when you need to verify what context was passed to the LLM,
299    /// for example when testing that file attachments are properly formatted.
300    pub async fn run_with_context(self) -> TestResult<TestAgentResult> {
301        let Self { responses, llm_config, agent_config, mcp_servers, elicitations, scenario } = self;
302        let scenario = scenario.expect("test agent requires commands(), user_text(), or scenario()");
303        let llm =
304            llm_config.into_iter().fold(FakeLlmProvider::from_results(responses), |llm, configure| configure(llm));
305        let captured_contexts = llm.captured_contexts();
306
307        let servers = mcp_servers.unwrap_or_else(|| vec![ServerSpec::in_process("test", FakeMcpServer::math())]);
308        let mut mcp = mcp("/workspace").with_servers(servers);
309        let mut script = None;
310        if let Some(elicitations) = elicitations {
311            let (sink, events) = mpsc::channel(32);
312            mcp = mcp.with_elicitation(ElicitationSink::form_and_url(sink));
313            script = Some(elicitations.spawn(events));
314        }
315
316        let builder = agent(llm).mcp(ready_catalog(mcp).await).retry(RetryConfig::disabled());
317        let builder = agent_config.into_iter().fold(builder, |builder, configure| configure(builder));
318        let (tx, mut rx, handle) = builder.spawn().await?;
319        let mut messages = Vec::new();
320
321        for step in scenario.steps {
322            match step {
323                Step::Send(command) => tx.send(command).await?,
324                Step::WaitFor(predicate) => loop {
325                    let message = rx.recv().await.expect("agent event channel closed before scenario step matched");
326                    let matched = predicate(&message);
327                    messages.push(message);
328                    if matched {
329                        break;
330                    }
331                },
332                Step::Perform(action) => action(),
333            }
334        }
335        drop(tx);
336
337        handle.await_completion().await;
338
339        let contexts = captured_contexts.lock().expect("captured contexts lock poisoned").clone();
340        let elicitations = script.map(|script| script.captured()).unwrap_or_default();
341        Ok(TestAgentResult { messages, contexts, elicitations })
342    }
343
344    fn configure(mut self, configure: impl FnOnce(AgentBuilder) -> AgentBuilder + Send + 'static) -> Self {
345        self.agent_config.push(Box::new(configure));
346        self
347    }
348
349    fn configure_llm(mut self, configure: impl FnOnce(FakeLlmProvider) -> FakeLlmProvider + Send + 'static) -> Self {
350        self.llm_config.push(Box::new(configure));
351        self
352    }
353}
354
355type Configure<T> = Box<dyn FnOnce(T) -> T + Send>;
356
357enum Step {
358    Send(Command),
359    WaitFor(Box<dyn Fn(&AgentEvent) -> bool + Send>),
360    Perform(Box<dyn FnOnce() + Send>),
361}
aether_core/testing/utils.rs

aether_core/testing/
utils.rs
