aether_core/testing/
agent_trace.rs

1use crate::events::{ToolEvent, TurnEvent};
2use std::sync::{Arc, Mutex};
3
4use crate::events::{AgentEvent, LlmCallOutcome, LlmCallPurpose, TurnOutcome};
5
6/// A recorded agent event stream (as seen by an
7/// [`AgentObserver`](crate::events::AgentObserver)), with assertion helpers
8/// for trace-shaped tests. Build one with
9/// [`TestAgentBuilder::run_trace`](crate::testing::TestAgentBuilder::run_trace).
10pub struct AgentTrace {
11    events: Vec<AgentEvent>,
12}
13
14impl AgentTrace {
15    pub fn from_events(events: Vec<AgentEvent>) -> Self {
16        Self { events }
17    }
18
19    pub fn from_observer_events(events: &Arc<Mutex<Vec<AgentEvent>>>) -> Self {
20        Self::from_events(events.lock().unwrap().clone())
21    }
22
23    pub fn events(&self) -> &[AgentEvent] {
24        &self.events
25    }
26
27    pub fn assert_names(&self, expected: &[&str]) {
28        assert_eq!(map_event_names(&self.events), expected, "unexpected trace: {:?}", self.events);
29    }
30
31    /// Index into [`Self::events`] of the first event matching `predicate`.
32    pub fn position(&self, predicate: impl Fn(&AgentEvent) -> bool) -> usize {
33        self.events
34            .iter()
35            .position(predicate)
36            .unwrap_or_else(|| panic!("expected event not found in trace: {:?}", self.events))
37    }
38
39    /// Indexes into [`Self::events`] of every event matching `predicate`.
40    pub fn positions(&self, predicate: impl Fn(&AgentEvent) -> bool) -> Vec<usize> {
41        self.events.iter().enumerate().filter(|(_, event)| predicate(event)).map(|(index, _)| index).collect()
42    }
43
44    pub fn call_usage(&self, for_purpose: LlmCallPurpose) -> Option<llm::TokenUsage> {
45        self.events.iter().find_map(|event| match event {
46            AgentEvent::Turn(TurnEvent::LlmCallEnded { purpose, outcome: LlmCallOutcome::Completed { usage, .. } })
47                if *purpose == for_purpose =>
48            {
49                *usage
50            }
51            _ => None,
52        })
53    }
54}
55
56/// Maps trace-relevant events to compact names for order assertions,
57/// skipping content messages.
58pub fn map_event_names(events: &[AgentEvent]) -> Vec<String> {
59    events
60        .iter()
61        .filter_map(|event| match event {
62            AgentEvent::Turn(TurnEvent::Started { .. }) => Some("turn_started".to_string()),
63            AgentEvent::Turn(TurnEvent::Ended { outcome }) => {
64                let outcome = match outcome {
65                    TurnOutcome::Completed => "completed",
66                    TurnOutcome::Cancelled => "cancelled",
67                    TurnOutcome::Failed { .. } => "failed",
68                };
69                Some(format!("turn_ended:{outcome}"))
70            }
71            AgentEvent::Turn(TurnEvent::RetryScheduled { purpose, attempt, .. }) => {
72                Some(format!("retry_scheduled:{purpose:?}:{attempt}"))
73            }
74            AgentEvent::Turn(TurnEvent::LlmCallStarted { purpose, attempt, .. }) => {
75                Some(format!("call_started:{purpose:?}:{attempt}"))
76            }
77            AgentEvent::Turn(TurnEvent::LlmCallEnded { purpose, outcome }) => {
78                let outcome = match outcome {
79                    LlmCallOutcome::Completed { .. } => "completed",
80                    LlmCallOutcome::Failed { will_retry: true, .. } => "failed_will_retry",
81                    LlmCallOutcome::Failed { will_retry: false, .. } => "failed_terminal",
82                    LlmCallOutcome::Cancelled => "cancelled",
83                };
84                Some(format!("call_ended:{purpose:?}:{outcome}"))
85            }
86            AgentEvent::Tool(ToolEvent::ExecutionStarted { .. }) => Some("tool_execution_started".to_string()),
87            AgentEvent::Tool(ToolEvent::DefinitionsUpdated { .. }) => Some("tool_definitions".to_string()),
88            _ => None,
89        })
90        .collect()
91}
aether_core/testing/agent_trace.rs

aether_core/testing/
agent_trace.rs
