diff --git a/crates/aether-core/src/testing/utils.rs b/crates/aether-core/src/testing/utils.rs index 8e364cd5..f1c546e3 100644 --- a/crates/aether-core/src/testing/utils.rs +++ b/crates/aether-core/src/testing/utils.rs @@ -1,11 +1,10 @@ use std::collections::BTreeMap; -use std::error::Error; use std::sync::{Arc, Mutex}; use std::time::Duration; use tokio::sync::{Notify, mpsc}; use crate::context::CompactionConfig; -use crate::core::{Prompt, RetryConfig, agent}; +use crate::core::{AgentError, Prompt, RetryConfig, agent}; use crate::events::{ AgentCommand, AgentEvent, AgentObserver, Command, ContextEvent, ToolEvent, TurnEvent, UserCommand, }; @@ -181,6 +180,18 @@ pub struct TestAgentResult { pub captured_contexts: Arc>>, } +/// Error returned by [`TestAgentBuilder::run`], [`TestAgentBuilder::run_trace`], and +/// [`TestAgentBuilder::run_with_context`]. +#[derive(Debug, thiserror::Error)] +pub enum TestAgentError { + #[error(transparent)] + Agent(#[from] AgentError), + #[error("failed to send command to the test agent: {0}")] + SendCommand(#[from] mpsc::error::SendError), +} + +pub type TestResult = std::result::Result; + struct ProviderTestConfig { responses: Vec>>, model: Option, @@ -333,14 +344,14 @@ impl TestAgentBuilder { self } - pub async fn run(self) -> Result, Box> { + pub async fn run(self) -> TestResult> { let result = self.run_with_context().await?; Ok(result.messages) } /// Runs the test agent with a recording observer attached and returns the /// full event trace, including internal events. - pub async fn run_trace(self) -> Result> { + pub async fn run_trace(self) -> TestResult { let observer = FakeAgentObserver::new(); let events = observer.events(); self.observer(Box::new(observer)).run().await?; @@ -351,7 +362,7 @@ impl TestAgentBuilder { /// /// Use this when you need to verify what context was passed to the LLM, /// for example when testing that file attachments are properly formatted. - pub async fn run_with_context(self) -> Result> { + pub async fn run_with_context(self) -> TestResult { let Self { provider, agent: config, execution } = self; let mut llm = FakeLlmProvider::from_results(provider.responses).with_context_window(provider.context_window); if let Some(model) = provider.model { @@ -363,7 +374,13 @@ impl TestAgentBuilder { let captured_contexts = llm.captured_contexts(); let mut mcp_spawn = if config.include_fake_mcp { - Some(mcp("/workspace").with_servers(vec![fake_mcp("test", FakeMcpServer::new())]).spawn().await?) + Some( + mcp("/workspace") + .with_servers(vec![fake_mcp("test", FakeMcpServer::new())]) + .spawn() + .await + .map_err(AgentError::from)?, + ) } else { None }; diff --git a/crates/aether-telemetry/tests/observer_tests.rs b/crates/aether-telemetry/tests/observer_tests.rs index 1dcdb91f..1cb5236f 100644 --- a/crates/aether-telemetry/tests/observer_tests.rs +++ b/crates/aether-telemetry/tests/observer_tests.rs @@ -428,12 +428,12 @@ async fn happy_tool_trace() -> Result> { .build(), llm_response("m2").text(&["The sum is 8"]).usage(30, 7).build(), ]; - test_agent() + Ok(test_agent() .model("anthropic:claude-sonnet-4-5".parse()?) .llm_responses(&responses) .user_text("3+5 = ?") .run_trace() - .await + .await?) } fn chat_call(provider: &str, model: &str, outcome: LlmCallOutcome) -> [AgentEvent; 2] {