autoagents_core/
utils.rs

1use futures_core::Stream;
2#[cfg(all(target_arch = "wasm32", target_os = "wasi"))]
3use std::future::Future;
4use std::pin::Pin;
5#[cfg(all(target_arch = "wasm32", target_os = "wasi"))]
6use std::task::{Context, Poll};
7
8// -----------------------------
9// Channel aliases
10// -----------------------------
11#[cfg(not(target_arch = "wasm32"))]
12pub use tokio::sync::mpsc::{Receiver, Sender, channel};
13
14#[cfg(target_arch = "wasm32")]
15pub use futures::channel::mpsc::{Receiver, Sender, channel};
16
17// -----------------------------
18// Unified boxed stream type
19// -----------------------------
20#[cfg(not(all(target_arch = "wasm32", target_os = "wasi")))]
21pub type BoxRuntimeStream<T> = Pin<Box<dyn Stream<Item = T> + Send>>;
22
23#[cfg(all(target_arch = "wasm32", target_os = "wasi"))]
24pub type BoxRuntimeStream<T> = Pin<Box<dyn Stream<Item = T>>>;
25
26#[cfg(not(target_arch = "wasm32"))]
27pub type BoxEventStream<T> = Pin<Box<dyn Stream<Item = T> + Send + Sync>>;
28
29#[cfg(target_arch = "wasm32")]
30pub type BoxEventStream<T> = Pin<Box<dyn Stream<Item = T> + Send>>;
31
32// -----------------------------
33// Conversion helpers
34// -----------------------------
35#[cfg(not(target_arch = "wasm32"))]
36pub(crate) fn receiver_into_stream<T: 'static + Send>(rx: Receiver<T>) -> BoxEventStream<T> {
37    use tokio_stream::wrappers::ReceiverStream;
38    Box::pin(ReceiverStream::new(rx))
39}
40
41#[cfg(target_arch = "wasm32")]
42pub(crate) fn receiver_into_stream<T: 'static + Send>(rx: Receiver<T>) -> BoxEventStream<T> {
43    Box::pin(rx)
44}
45
46#[cfg(all(target_arch = "wasm32", target_os = "wasi"))]
47struct WasiDrivenStream<T> {
48    producer: Pin<Box<dyn Future<Output = ()>>>,
49    receiver: BoxRuntimeStream<T>,
50    producer_done: bool,
51}
52
53#[cfg(all(target_arch = "wasm32", target_os = "wasi"))]
54impl<T> Stream for WasiDrivenStream<T> {
55    type Item = T;
56
57    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
58        if !self.producer_done && self.producer.as_mut().poll(cx).is_ready() {
59            self.producer_done = true;
60        }
61
62        match self.receiver.as_mut().poll_next(cx) {
63            Poll::Ready(item) => Poll::Ready(item),
64            Poll::Pending if self.producer_done => Poll::Ready(None),
65            Poll::Pending => Poll::Pending,
66        }
67    }
68}
69
70// Platform-specific spawn functions
71#[cfg(not(target_arch = "wasm32"))]
72pub(crate) fn spawn_future<F>(fut: F) -> tokio::task::JoinHandle<F::Output>
73where
74    F: std::future::Future + Send + 'static,
75    F::Output: Send + 'static,
76{
77    tokio::spawn(fut)
78}
79
80#[cfg(all(target_arch = "wasm32", not(target_os = "wasi")))]
81pub(crate) fn spawn_future<F>(fut: F)
82where
83    F: std::future::Future<Output = ()> + 'static,
84{
85    wasm_bindgen_futures::spawn_local(fut)
86}
87
88#[cfg(not(target_arch = "wasm32"))]
89pub(crate) fn stream_from_producer<T, F>(rx: Receiver<T>, producer: F) -> BoxRuntimeStream<T>
90where
91    T: 'static + Send,
92    F: std::future::Future<Output = ()> + Send + 'static,
93{
94    use tokio_stream::wrappers::ReceiverStream;
95
96    spawn_future(producer);
97    Box::pin(ReceiverStream::new(rx))
98}
99
100#[cfg(all(target_arch = "wasm32", not(target_os = "wasi")))]
101pub(crate) fn stream_from_producer<T, F>(rx: Receiver<T>, producer: F) -> BoxRuntimeStream<T>
102where
103    T: 'static + Send,
104    F: std::future::Future<Output = ()> + 'static,
105{
106    spawn_future(producer);
107    Box::pin(rx)
108}
109
110#[cfg(all(target_arch = "wasm32", target_os = "wasi"))]
111pub(crate) fn stream_from_producer<T, F>(rx: Receiver<T>, producer: F) -> BoxRuntimeStream<T>
112where
113    T: 'static + Send,
114    F: std::future::Future<Output = ()> + 'static,
115{
116    Box::pin(WasiDrivenStream {
117        producer: Box::pin(producer),
118        receiver: Box::pin(rx),
119        producer_done: false,
120    })
121}
122
123#[cfg(all(target_arch = "wasm32", target_os = "wasi"))]
124pub fn block_on_local_executor<F>(future: F) -> F::Output
125where
126    F: std::future::Future,
127{
128    // Keep the public WASI entrypoint stable for downstream crates while the
129    // streaming internals run cooperatively instead of on a detached executor.
130    futures::executor::block_on(future)
131}
132
133#[cfg(test)]
134mod tests {
135    use super::*;
136    use futures_util::StreamExt;
137
138    #[tokio::test]
139    async fn receiver_into_stream_yields_messages_in_order() {
140        let (tx, rx) = channel(2);
141        tx.send(1).await.expect("send first message");
142        tx.send(2).await.expect("send second message");
143        drop(tx);
144
145        let values: Vec<_> = receiver_into_stream(rx).collect().await;
146        assert_eq!(values, vec![1, 2]);
147    }
148
149    #[tokio::test]
150    async fn spawn_future_returns_joinhandle_output() {
151        let handle = spawn_future(async { 7usize });
152
153        assert_eq!(handle.await.expect("task joins"), 7);
154    }
155
156    #[tokio::test]
157    async fn stream_from_producer_forwards_background_values() {
158        let (tx, rx) = channel(3);
159        let stream = stream_from_producer(rx, async move {
160            for value in [1, 2, 3] {
161                tx.send(value).await.expect("send produced value");
162            }
163        });
164
165        let values: Vec<_> = stream.collect().await;
166        assert_eq!(values, vec![1, 2, 3]);
167    }
168}
autoagents_core/utils.rs

autoagents_core/
utils.rs
