tinyagents/graph/testkit/
types.rs

1//! Data types for the graph testkit.
2//!
3//! See the module [`mod`](super) docs for the building blocks these back: the
4//! event recorder ([`GraphEventRecorder`]), the projection collector
5//! ([`StreamCollector`]), the bundled run-under-test ([`GraphRun`]), the fluent
6//! assertion builder ([`GraphAssertions`]), and the stateful retry-counting node
7//! ([`RetryCountingNode`]). This file holds the plain definitions; the impls and
8//! the node-helper free functions live in `mod.rs`.
9
10use std::sync::Arc;
11use std::sync::atomic::AtomicUsize;
12
13use crate::graph::compiled::{GraphExecution, StateSnapshot};
14use crate::graph::stream::{CollectingSink, GraphEvent};
15
16/// Captures the durable executor's [`GraphEvent`] stream for inspection.
17///
18/// Wraps a [`CollectingSink`]; hand its [`sink`](GraphEventRecorder::sink) to
19/// [`CompiledGraph::with_event_sink`](crate::graph::CompiledGraph::with_event_sink)
20/// (or use [`run_recorded`](super::run_recorded), which wires it for you) so a
21/// run's events are recorded, then read them back as raw events, `kind` strings,
22/// or a [`StreamCollector`] of higher-level projections.
23#[derive(Clone, Default)]
24pub struct GraphEventRecorder {
25    pub(crate) sink: CollectingSink,
26}
27
28/// A higher-level projection over a recorded [`GraphEvent`] stream.
29///
30/// Where [`GraphEventRecorder`] hands back raw events, `StreamCollector` slices
31/// them into the views a test usually wants: the executed-node order, the nodes
32/// that wrote state, the `(from, to)` routes the executor selected, the emitted
33/// interrupts, the persisted-checkpoint count, and any custom node writes.
34#[derive(Clone, Default)]
35pub struct StreamCollector {
36    pub(crate) events: Vec<GraphEvent>,
37}
38
39/// A completed (or interrupted) run bundled with everything the assertions need.
40///
41/// [`assert_graph`](super::assert_graph) reads its truth from one `GraphRun`:
42/// the [`GraphExecution`] (visited nodes, steps, status, interrupts), the
43/// recorded [`GraphEvent`] stream (real route selections and checkpoint saves),
44/// and the thread's checkpoint `history` (newest-first, as
45/// [`get_state_history`](crate::graph::CompiledGraph::get_state_history) returns
46/// it). Build one with [`run_recorded`](super::run_recorded), or assemble it by
47/// hand via [`GraphRun::new`] and the `with_*` builders.
48pub struct GraphRun<State> {
49    /// The execution result (final state, visited nodes, steps, status).
50    pub execution: GraphExecution<State>,
51    /// The recorded low-level graph events, in emission order.
52    pub events: Vec<GraphEvent>,
53    /// The thread's checkpoint history, newest-first (empty for a non-durable
54    /// run, or one executed without a thread id).
55    pub history: Vec<StateSnapshot<State>>,
56}
57
58/// A fluent, panic-on-failure assertion builder over a [`GraphRun`].
59///
60/// Returned by [`assert_graph`](super::assert_graph). Every method asserts and
61/// returns `&Self` so checks chain:
62///
63/// ```ignore
64/// assert_graph(&run)
65///     .visited(["agent", "tools", "agent"])
66///     .routed("agent", "tools")
67///     .checkpoint_count(3)
68///     .completed();
69/// ```
70pub struct GraphAssertions<'a, State> {
71    pub(crate) run: &'a GraphRun<State>,
72}
73
74/// A stateful node double that counts its activations and fails the first
75/// `fail_times` of them.
76///
77/// Each [`handler`](RetryCountingNode::handler) it produces shares the same
78/// activation counter, so a test can run the graph and then read
79/// [`attempts`](RetryCountingNode::attempts) to assert how many times the node
80/// ran. The first `fail_times` activations return
81/// [`TinyAgentsError::Graph`](crate::TinyAgentsError::Graph); subsequent ones
82/// succeed with the configured update.
83#[derive(Clone)]
84pub struct RetryCountingNode {
85    pub(crate) attempts: Arc<AtomicUsize>,
86    pub(crate) fail_times: usize,
87}
tinyagents/graph/testkit/types.rs

tinyagents/graph/testkit/
types.rs
