graphflow_stream/
stream.rs

1#![doc = include_str!("../README.md")]
2
3use graph_flow::{
4    Context, ExecutionResult, ExecutionStatus, FlowRunner, Graph, InMemorySessionStorage,
5    NextAction, Session, SessionStorage, Task, TaskResult,
6    error::{GraphError, Result},
7};
8use serde::{Deserialize, Serialize};
9use std::future::Future;
10use std::sync::Arc;
11use std::sync::atomic::{AtomicU64, Ordering};
12use std::time::Duration;
13use tokio::sync::mpsc;
14use tokio::task::JoinHandle;
15
16/// Sending half of a stream channel; a running task's `emit_*` calls write here.
17pub type StreamSender = mpsc::Sender<StreamEvent>;
18/// Receiving half of a stream channel; drain this to observe a run in progress.
19pub type StreamReceiver = mpsc::Receiver<StreamEvent>;
20
21/// An incremental event emitted while a task or graph run is in flight.
22#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
23pub enum StreamEvent {
24    /// A task started running.
25    TaskStarted { task_id: String },
26    /// A partial chunk of output (e.g. one LLM token).
27    Token { task_id: String, delta: String },
28    /// A task finished successfully.
29    TaskFinished { task_id: String },
30    /// A task returned an error instead of finishing normally.
31    TaskFailed { task_id: String, error: String },
32}
33
34tokio::task_local! {
35    static STREAM_TX: StreamSender;
36}
37
38/// Send an event on the ambient stream, if one is active; a no-op otherwise.
39pub async fn emit(event: StreamEvent) {
40    if let Ok(tx) = STREAM_TX.try_with(|tx| tx.clone()) {
41        let _ = tx.send(event).await;
42    }
43}
44
45/// Emit a [`StreamEvent::TaskStarted`] on the ambient stream.
46pub async fn emit_started(task_id: impl Into<String>) {
47    emit(StreamEvent::TaskStarted {
48        task_id: task_id.into(),
49    })
50    .await;
51}
52
53/// Emit a [`StreamEvent::Token`] on the ambient stream.
54pub async fn emit_token(task_id: impl Into<String>, delta: impl Into<String>) {
55    emit(StreamEvent::Token {
56        task_id: task_id.into(),
57        delta: delta.into(),
58    })
59    .await;
60}
61
62/// Emit a [`StreamEvent::TaskFinished`] on the ambient stream.
63pub async fn emit_finished(task_id: impl Into<String>) {
64    emit(StreamEvent::TaskFinished {
65        task_id: task_id.into(),
66    })
67    .await;
68}
69
70/// Emit a [`StreamEvent::TaskFailed`] on the ambient stream.
71pub async fn emit_failed(task_id: impl Into<String>, error: impl Into<String>) {
72    emit(StreamEvent::TaskFailed {
73        task_id: task_id.into(),
74        error: error.into(),
75    })
76    .await;
77}
78
79/// Drain any `Stream<Item = String>` (e.g. a mapped LLM completion) into `emit_token` calls.
80pub async fn forward_text_stream<S>(task_id: impl Into<String>, mut stream: S)
81where
82    S: futures_core::Stream<Item = String> + Unpin,
83{
84    use tokio_stream::StreamExt;
85
86    let task_id = task_id.into();
87    while let Some(delta) = stream.next().await {
88        emit_token(task_id.clone(), delta).await;
89    }
90    emit_finished(task_id).await;
91}
92
93/// Run `fut` with an ambient stream scope active; the primitive `spawn_task`/`spawn_graph` use.
94pub fn run_streaming<F>(buffer: usize, fut: F) -> (StreamReceiver, JoinHandle<F::Output>)
95where
96    F: Future + Send + 'static,
97    F::Output: Send + 'static,
98{
99    let (tx, rx) = mpsc::channel(buffer);
100    let handle = tokio::spawn(STREAM_TX.scope(tx, fut));
101    (rx, handle)
102}
103
104/// Run a single [`Task`], streaming its `emit_*` calls out through the returned receiver.
105pub fn spawn_task<T>(
106    task: Arc<T>,
107    context: Context,
108    buffer: usize,
109) -> (StreamReceiver, JoinHandle<Result<TaskResult>>)
110where
111    T: Task + Send + Sync + 'static,
112{
113    run_streaming(buffer, async move { task.run(context).await })
114}
115
116/// Run a [`Task`] and join its streamed [`StreamEvent::Token`] deltas into one `String`.
117pub async fn collect_text<T>(task: Arc<T>, context: Context, buffer: usize) -> Result<String>
118where
119    T: Task + Send + Sync + 'static,
120{
121    let (mut rx, handle) = spawn_task(task, context, buffer);
122    let mut text = String::new();
123    while let Some(event) = rx.recv().await {
124        if let StreamEvent::Token { delta, .. } = event {
125            text.push_str(&delta);
126        }
127    }
128    handle.await.map_err(|e| GraphError::Other(e.into()))??;
129    Ok(text)
130}
131
132/// Drive a [`FlowRunner`] to completion, streaming every task's `emit_*` calls out.
133pub fn spawn_graph(
134    flow_runner: FlowRunner,
135    session_id: impl Into<String>,
136    buffer: usize,
137) -> (StreamReceiver, JoinHandle<Result<ExecutionResult>>) {
138    let session_id = session_id.into();
139    run_streaming(buffer, async move {
140        run_to_completion(&flow_runner, &session_id).await
141    })
142}
143
144async fn run_to_completion(flow_runner: &FlowRunner, session_id: &str) -> Result<ExecutionResult> {
145    loop {
146        let result = flow_runner.run(session_id).await?;
147        if !matches!(result.status, ExecutionStatus::Paused { .. }) {
148            return Ok(result);
149        }
150    }
151}
152
153static SUBGRAPH_SESSION_COUNTER: AtomicU64 = AtomicU64::new(0);
154
155/// A [`Task`] that drives a whole inner [`Graph`], so a graph can be a node in another graph.
156pub struct SubgraphTask {
157    id: String,
158    graph: Arc<Graph>,
159    storage: Arc<dyn SessionStorage>,
160}
161
162impl SubgraphTask {
163    /// Wrap `graph` as a task with the given id, using in-memory session storage.
164    pub fn new(id: impl Into<String>, graph: Arc<Graph>) -> Self {
165        Self {
166            id: id.into(),
167            graph,
168            storage: Arc::new(InMemorySessionStorage::new()),
169        }
170    }
171
172    /// Use a custom [`SessionStorage`] backend for the inner graph's sessions.
173    pub fn with_storage(mut self, storage: Arc<dyn SessionStorage>) -> Self {
174        self.storage = storage;
175        self
176    }
177}
178
179#[async_trait::async_trait]
180impl Task for SubgraphTask {
181    fn id(&self) -> &str {
182        &self.id
183    }
184
185    async fn run(&self, context: Context) -> Result<TaskResult> {
186        let start_task_id = self.graph.start_task_id().ok_or_else(|| {
187            GraphError::TaskNotFound(format!("subgraph '{}' has no start task", self.graph.id))
188        })?;
189
190        let suffix = SUBGRAPH_SESSION_COUNTER.fetch_add(1, Ordering::Relaxed);
191        let session_id = format!("{}:{}", self.id, suffix);
192
193        let mut session = Session::new_from_task(session_id.clone(), start_task_id)
194            .with_graph_id(self.graph.id.clone());
195        session.context = context;
196        self.storage.save(session).await?;
197
198        let runner = FlowRunner::new(self.graph.clone(), self.storage.clone());
199        let result = run_to_completion(&runner, &session_id).await?;
200
201        let next_action = match result.status {
202            ExecutionStatus::WaitingForInput => NextAction::WaitForInput,
203            _ => NextAction::Continue,
204        };
205        Ok(TaskResult::new(result.response, next_action))
206    }
207}
208
209/// A [`StreamEvent`] plus its offset from the start of the recording.
210#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
211pub struct RecordedEvent {
212    pub event: StreamEvent,
213    pub at: Duration,
214}
215
216/// A recorded run: every [`StreamEvent`] a `StreamReceiver` produced, with timing.
217///
218/// Serializable, so a recording can be saved and replayed later — e.g. to debug a
219/// past run, or demo a UI without hitting an LLM again.
220#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
221pub struct Recording {
222    pub events: Vec<RecordedEvent>,
223}
224
225impl Recording {
226    /// Replay this recording on a fresh channel, preserving the original timing between events.
227    pub fn replay(self, buffer: usize) -> StreamReceiver {
228        let (tx, rx) = mpsc::channel(buffer);
229        tokio::spawn(async move {
230            let mut last = Duration::ZERO;
231            for recorded in self.events {
232                let wait = recorded.at.saturating_sub(last);
233                if !wait.is_zero() {
234                    tokio::time::sleep(wait).await;
235                }
236                last = recorded.at;
237                if tx.send(recorded.event).await.is_err() {
238                    return;
239                }
240            }
241        });
242        rx
243    }
244}
245
246/// Drain a `StreamReceiver` into a [`Recording`], timestamping each event from the start.
247pub async fn record(mut rx: StreamReceiver) -> Recording {
248    let start = std::time::Instant::now();
249    let mut events = Vec::new();
250    while let Some(event) = rx.recv().await {
251        events.push(RecordedEvent {
252            event,
253            at: start.elapsed(),
254        });
255    }
256    Recording { events }
257}
graphflow_stream/stream.rs

graphflow_stream/
stream.rs
