aether_core/testing/
agent_trace.rs

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

aether_core/testing/
agent_trace.rs
