flashmind_types/
stream.rs

1//! [`AgentStream`] — a stream that yields events and terminates with a result.
2//!
3//! Replaces the `(impl Stream<Item = I>, oneshot::Receiver<O>)` pattern used
4//! throughout the agent runtime. Producers yield [`Outcome::Item`] during
5//! processing and [`Outcome::Done`] as the final value. Consumers iterate
6//! the stream normally and call [`take_result`](AgentStream::take_result)
7//! after exhaustion.
8//!
9//! # Why not just `Stream`?
10//!
11//! The agent loop needs both incremental events (text deltas, tool progress)
12//! and a final assembled result (full response text, token counts, finish reason).
13//! This type packages both into a single handle so callers don't need to juggle
14//! two separate handles or risk deadlocks from dropping the receiver.
15
16use std::pin::Pin;
17use std::task::{Context, Poll};
18
19use futures_core::Stream;
20
21/// A single element in an [`AgentStream`]: either an intermediate event or the
22/// final result.
23pub enum Outcome<I, O> {
24    /// An intermediate item yielded during processing.
25    Item(I),
26    /// The final result — must be yielded exactly once, as the last element.
27    Done(O),
28}
29
30/// A stream that yields events of type `I` and produces a final result `O`.
31///
32/// Implements [`Stream<Item = I>`] so it can be used with `.next().await` and
33/// `StreamExt` combinators. When the inner stream yields [`Outcome::Done`],
34/// the result is stored internally and the stream terminates. Retrieve it
35/// with [`take_result`](Self::take_result).
36///
37/// # Thread safety
38///
39/// The inner stream must satisfy `Send + 'a`, which means `AgentStream` can be
40/// safely moved across thread boundaries (e.g., spawned on a Tokio task).
41///
42/// # Example
43///
44/// ```ignore
45/// fn work() -> AgentStream<'static, String, u32> {
46///     AgentStream::new(async_stream::stream! {
47///         yield Outcome::Item("hello".into());
48///         yield Outcome::Item("world".into());
49///         yield Outcome::Done(42);
50///     })
51/// }
52///
53/// let mut s = work();
54/// while let Some(item) = s.next().await {
55///     println!("{item}");
56/// }
57/// let answer = s.take_result().unwrap();
58/// ```
59pub struct AgentStream<'a, I, O> {
60    inner: Pin<Box<dyn Stream<Item = Outcome<I, O>> + Send + 'a>>,
61    result: Option<O>,
62}
63
64impl<'a, I, O> AgentStream<'a, I, O> {
65    /// Wrap a stream of [`Outcome`] values into an `AgentStream`.
66    pub fn new(inner: impl Stream<Item = Outcome<I, O>> + Send + 'a) -> Self {
67        Self {
68            inner: Box::pin(inner),
69            result: None,
70        }
71    }
72
73    /// Take the final result after the stream has been fully consumed.
74    ///
75    /// Returns `None` if the stream hasn't ended yet or if the producer
76    /// never yielded [`Outcome::Done`].
77    ///
78    /// # Panics
79    ///
80    /// This method does not panic. It always returns an [`Option`], yielding
81    /// `None` when no result is available.
82    pub fn take_result(&mut self) -> Option<O> {
83        self.result.take()
84    }
85}
86
87impl<I, O: Unpin> AgentStream<'_, I, O> {
88    /// Poll for the next item without requiring manual pinning or `StreamExt`.
89    pub async fn next(&mut self) -> Option<I> {
90        use futures_util::StreamExt;
91        StreamExt::next(self).await
92    }
93}
94
95impl<I, O: Unpin> Stream for AgentStream<'_, I, O> {
96    type Item = I;
97
98    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<I>> {
99        let this = self.get_mut();
100        match this.inner.as_mut().poll_next(cx) {
101            Poll::Ready(Some(Outcome::Item(item))) => Poll::Ready(Some(item)),
102            Poll::Ready(Some(Outcome::Done(result))) => {
103                this.result = Some(result);
104                Poll::Ready(None)
105            }
106            Poll::Ready(None) => Poll::Ready(None),
107            Poll::Pending => Poll::Pending,
108        }
109    }
110}
111
112#[cfg(test)]
113mod tests {
114    use super::*;
115    use futures_util::StreamExt;
116
117    #[tokio::test]
118    async fn yields_items_then_result() {
119        let mut s = AgentStream::new(async_stream::stream! {
120            yield Outcome::Item("a");
121            yield Outcome::Item("b");
122            yield Outcome::Done(42);
123        });
124
125        let items: Vec<&str> = s.by_ref().collect().await;
126        assert_eq!(items, vec!["a", "b"]);
127        assert_eq!(s.take_result(), Some(42));
128    }
129
130    #[tokio::test]
131    async fn empty_stream_no_result() {
132        let mut s: AgentStream<&str, i32> = AgentStream::new(futures_util::stream::empty());
133
134        let items: Vec<&str> = s.by_ref().collect().await;
135        assert!(items.is_empty());
136        assert_eq!(s.take_result(), None);
137    }
138
139    #[tokio::test]
140    async fn result_only() {
141        let mut s: AgentStream<&str, i32> = AgentStream::new(async_stream::stream! {
142            yield Outcome::Done(99);
143        });
144
145        let items: Vec<&str> = s.by_ref().collect().await;
146        assert!(items.is_empty());
147        assert_eq!(s.take_result(), Some(99));
148    }
149
150    #[tokio::test]
151    async fn next_convenience() {
152        let mut s = AgentStream::new(async_stream::stream! {
153            yield Outcome::Item("x");
154            yield Outcome::Item("y");
155            yield Outcome::Done(7);
156        });
157
158        assert_eq!(s.next().await, Some("x"));
159        assert_eq!(s.next().await, Some("y"));
160        assert_eq!(s.next().await, None);
161        assert_eq!(s.take_result(), Some(7));
162    }
163}
flashmind_types/stream.rs

flashmind_types/
stream.rs
