ag_ui/client/transport/
replay.rs

1//! A transport that replays a scripted list of events.
2//!
3//! Testing a client against a live agent is slow, flaky, and needs a model. It
4//! is also unnecessary: the agent's half of the conversation is just a list of
5//! events. [`ReplayTransport`] serves one — and records the
6//! [`RunAgentInput`](https://kimsoungryoul.github.io/ag-ui-rust/api/ag_ui/input/struct.RunAgentInput.html)s it was handed, which is how a test asserts that a resume
7//! carried the right answers.
8//!
9//! ```
10//! use ag_ui::client::transport::ReplayTransport;
11//! use ag_ui::Event;
12//!
13//! let transport = ReplayTransport::new([
14//!     Event::run_started("thread-1", "run-1"),
15//!     Event::run_finished_success("thread-1", "run-1"),
16//! ]);
17//! ```
18
19use std::collections::VecDeque;
20use std::sync::{Arc, Mutex, MutexGuard};
21
22use crate::{Event, RunAgentInput};
23
24use crate::client::error::Error;
25use crate::client::transport::{EventStream, Transport, TransportFuture};
26
27/// A [`Transport`] that answers each run from a script.
28///
29/// Cloning shares the script and the recording, so a test can keep a handle
30/// after handing one to a [`Session`](crate::client::Session).
31#[derive(Clone, Debug, Default)]
32pub struct ReplayTransport {
33    inner: Arc<Mutex<Script>>,
34}
35
36#[derive(Debug, Default)]
37struct Script {
38    runs: VecDeque<Vec<Event>>,
39    requests: Vec<RunAgentInput>,
40}
41
42impl ReplayTransport {
43    /// A transport that answers the first run with these events, and every
44    /// later run with an error.
45    pub fn new(events: impl IntoIterator<Item = Event>) -> Self {
46        Self::with_runs([events.into_iter().collect::<Vec<_>>()])
47    }
48
49    /// A transport that answers each run with the next list in the script.
50    ///
51    /// This is what a human-in-the-loop round trip needs: the first run pauses
52    /// on an interrupt, the second — the resume — carries on.
53    pub fn with_runs(runs: impl IntoIterator<Item = Vec<Event>>) -> Self {
54        Self {
55            inner: Arc::new(Mutex::new(Script {
56                runs: runs.into_iter().collect(),
57                requests: Vec::new(),
58            })),
59        }
60    }
61
62    /// Every request this transport has been handed, in order.
63    pub fn requests(&self) -> Vec<RunAgentInput> {
64        self.lock().requests.clone()
65    }
66
67    /// The most recent request, if there has been one.
68    pub fn last_request(&self) -> Option<RunAgentInput> {
69        self.lock().requests.last().cloned()
70    }
71
72    /// How many runs are left in the script.
73    pub fn remaining(&self) -> usize {
74        self.lock().runs.len()
75    }
76
77    /// A poisoned lock still holds a perfectly good script — a test that
78    /// panicked mid-assert should fail on that panic, not on this mutex.
79    fn lock(&self) -> MutexGuard<'_, Script> {
80        self.inner.lock().unwrap_or_else(|error| error.into_inner())
81    }
82}
83
84impl Transport for ReplayTransport {
85    fn run(&self, input: RunAgentInput) -> TransportFuture {
86        let mut script = self.lock();
87        script.requests.push(input);
88        let next = script.runs.pop_front();
89        drop(script);
90
91        Box::pin(async move {
92            let Some(events) = next else {
93                return Err(Error::Transport(
94                    "the replay script has no runs left".into(),
95                ));
96            };
97            let stream = futures_util::stream::iter(events.into_iter().map(Ok));
98            Ok(Box::pin(stream) as EventStream)
99        })
100    }
101}
ag_ui/client/transport/replay.rs

ag_ui/client/transport/
replay.rs
