lc_langgraph/compiled/
stream.rs

1// crates/lc-langgraph/src/compiled/stream.rs
2//! CompiledGraph stream, find_next_node, find_fan_out_targets,
3//! find_fan_in_target, and execute_parallel_branches methods
4
5use super::graph::CompiledGraph;
6use super::types::{GraphInvocation, StreamEvent};
7use crate::edge::GraphEdge;
8use crate::errors::{GraphError, GraphResult};
9use crate::graph::END;
10use crate::node::NodeConfig;
11use crate::state::StateSchema;
12use futures_util::future::join_all;
13use futures_util::Stream;
14use std::collections::HashMap;
15use std::pin::Pin;
16use tokio_stream::wrappers::ReceiverStream;
17
18impl<S: StateSchema + Send + Sync + 'static> CompiledGraph<S> {
19    /// Stream graph execution as a true async stream.
20    ///
21    /// Each `StreamEvent` is emitted as soon as it occurs (node entry,
22    /// node completion, state update), enabling real-time consumption.
23    ///
24    /// This is the preferred streaming API. For the old all-at-once
25    /// behavior, use `stream_collected()`.
26    pub fn stream(&self, input: S) -> Pin<Box<dyn Stream<Item = Result<StreamEvent<S>, GraphError>> + Send>> {
27        let (tx, rx) = tokio::sync::mpsc::channel(64);
28
29        let graph = self.clone();
30        tokio::spawn(async move {
31            let mut state = input;
32            let mut current_node = graph.entry_point.clone();
33            let mut recursion_count = 0;
34
35            if tx.send(Ok(StreamEvent::start(state.clone()))).await.is_err() {
36                return;
37            }
38
39            if let Some(ref checkpointer) = graph.checkpointer {
40                match checkpointer.lock().await.save(&state).await {
41                    Ok(_) => {}
42                    Err(e) => {
43                        let _ = tx.send(Err(e)).await;
44                        return;
45                    }
46                }
47            }
48
49            while current_node != END && recursion_count < graph.recursion_limit {
50                if graph.interrupt_before.contains(&current_node) {
51                    let _ = tx.send(Err(GraphError::ExecutionInterrupted(current_node.clone()))).await;
52                    return;
53                }
54
55                recursion_count += 1;
56
57                if tx.send(Ok(StreamEvent::enter_node(current_node.clone(), state.clone()))).await.is_err() {
58                    return;
59                }
60
61                let node = match graph.get_node(&current_node).await {
62                    Ok(n) => n,
63                    Err(e) => {
64                        let _ = tx.send(Err(e)).await;
65                        return;
66                    }
67                };
68
69                let config = NodeConfig {
70                    recursion_limit: graph.recursion_limit,
71                    debug: false,
72                    metadata: HashMap::new(),
73                };
74
75                let update = match node.execute(&state, Some(config)).await {
76                    Ok(u) => u,
77                    Err(e) => {
78                        let _ = tx.send(Err(e)).await;
79                        return;
80                    }
81                };
82
83                if tx.send(Ok(StreamEvent::node_complete(current_node.clone(), update.clone()))).await.is_err() {
84                    return;
85                }
86
87                if let Some(new_state) = update.update {
88                    state = graph.default_reducer.reduce(&state, &new_state);
89                    if tx.send(Ok(StreamEvent::state_update(state.clone()))).await.is_err() {
90                        return;
91                    }
92                }
93
94                if graph.interrupt_after.contains(&current_node) {
95                    if let Some(ref checkpointer) = graph.checkpointer {
96                        let _ = checkpointer.lock().await.save(&state).await;
97                    }
98                    let _ = tx.send(Err(GraphError::ExecutionInterrupted(format!(
99                        "after_{}",
100                        current_node
101                    )))).await;
102                    return;
103                }
104
105                let next_node = match graph.find_next_node(&current_node, &state).await {
106                    Ok(n) => n,
107                    Err(e) => {
108                        let _ = tx.send(Err(e)).await;
109                        return;
110                    }
111                };
112
113                if let Some(ref checkpointer) = graph.checkpointer {
114                    let _ = checkpointer.lock().await.save(&state).await;
115                }
116
117                current_node = next_node;
118            }
119
120            let _ = tx.send(Ok(StreamEvent::end(state))).await;
121        });
122
123        Box::pin(ReceiverStream::new(rx))
124    }
125
126    /// Collect all stream events into a Vec (backward-compatible API).
127    ///
128    /// This is the old `stream()` behavior. Prefer `stream()` for
129    /// real-time consumption.
130    pub async fn stream_collected(&self, input: S) -> GraphResult<Vec<StreamEvent<S>>> {
131        let mut events = Vec::new();
132        let mut stream = self.stream(input);
133        use futures_util::StreamExt;
134        while let Some(event) = stream.next().await {
135            events.push(event?);
136        }
137        Ok(events)
138    }
139
140    pub(super) async fn find_next_node(&self, current: &str, state: &S) -> GraphResult<String> {
141        // 1. Check runtime edges — edge is cloned out of the guard before any .await
142        'rt: {
143            let edge = {
144                let re = self.runtime_edges.read().await;
145                match re.iter().find(|e| e.source() == current) {
146                    Some(e) => e.clone(),
147                    None => break 'rt,
148                }
149            }; // re dropped here
150            match edge {
151                GraphEdge::Fixed { target, .. } => return Ok(target),
152                GraphEdge::Conditional {
153                    router_name,
154                    targets,
155                    default_target,
156                    ..
157                } => {
158                    let router = self
159                        .conditional_routers
160                        .get(&router_name)
161                        .cloned()
162                        .or_else(|| {
163                            self.runtime_conditional_routers
164                                .try_read()
165                                .ok()
166                                .and_then(|guard| guard.get(&router_name).cloned())
167                        })
168                        .ok_or_else(|| {
169                            GraphError::ExecutionError(format!(
170                                "Router '{}' not found (runtime)",
171                                router_name
172                            ))
173                        })?;
174                    let route_key = router.route(state).await?;
175                    let target = targets
176                        .get(&route_key)
177                        .or(default_target.as_ref())
178                        .ok_or_else(|| {
179                            GraphError::RoutingError(format!(
180                                "No target for route '{}' (runtime)",
181                                route_key
182                            ))
183                        })?;
184                    return Ok(target.clone());
185                }
186                GraphEdge::FanOut { targets, .. } => {
187                    if targets.is_empty() {
188                        return Err(GraphError::RoutingError(
189                            "FanOut has no targets (runtime)".to_string(),
190                        ));
191                    }
192                    return Ok(targets[0].clone());
193                }
194                GraphEdge::FanIn { .. } => {}
195            }
196        }
197
198        // 2. Check static edges
199        for edge in &self.edges {
200            if edge.source() == current {
201                match edge {
202                    GraphEdge::Fixed { target, .. } => {
203                        return Ok(target.clone());
204                    }
205                    GraphEdge::Conditional {
206                        router_name,
207                        targets,
208                        default_target,
209                        ..
210                    } => {
211                        let router = self
212                            .conditional_routers
213                            .get(router_name)
214                            .cloned()
215                            .or_else(|| {
216                                self.runtime_conditional_routers
217                                    .try_read()
218                                    .ok()
219                                    .and_then(|guard| guard.get(router_name).cloned())
220                            })
221                            .ok_or_else(|| {
222                                GraphError::ExecutionError(format!(
223                                    "Router '{}' not found",
224                                    router_name
225                                ))
226                            })?;
227
228                        let route_key = router.route(state).await?;
229
230                        let target = targets
231                            .get(&route_key)
232                            .or(default_target.as_ref())
233                            .ok_or_else(|| {
234                                GraphError::RoutingError(format!(
235                                    "No target for route '{}'",
236                                    route_key
237                                ))
238                            })?;
239
240                        return Ok(target.clone());
241                    }
242                    GraphEdge::FanOut { targets, .. } => {
243                        if targets.is_empty() {
244                            return Err(GraphError::RoutingError(
245                                "FanOut has no targets".to_string(),
246                            ));
247                        }
248                        return Ok(targets[0].clone());
249                    }
250                    GraphEdge::FanIn { .. } => {
251                        continue;
252                    }
253                }
254            }
255        }
256
257        if current == self.entry_point && self.nodes.len() == 1 {
258            return Ok(END.to_string());
259        }
260
261        Err(GraphError::RoutingError(format!(
262            "No outgoing edge from node '{}'",
263            current
264        )))
265    }
266
267    pub(super) async fn find_fan_out_targets(&self, current: &str) -> Option<Vec<String>> {
268        {
269            let re = self.runtime_edges.read().await;
270            if let Some(GraphEdge::FanOut { targets, .. }) =
271                re.iter().find(|e| e.source() == current)
272            {
273                return Some(targets.clone());
274            }
275        }
276        for edge in &self.edges {
277            if edge.source() == current {
278                if let GraphEdge::FanOut { targets, .. } = edge {
279                    return Some(targets.clone());
280                }
281            }
282        }
283        None
284    }
285
286    pub(super) async fn find_fan_in_target(&self, sources: &[String]) -> Option<String> {
287        {
288            let re = self.runtime_edges.read().await;
289            if let Some(GraphEdge::FanIn {
290                sources: edge_sources,
291                target,
292            }) = re.iter().find(|e| matches!(e, GraphEdge::FanIn { .. }))
293            {
294                if edge_sources.iter().all(|s| sources.contains(s)) {
295                    return Some(target.clone());
296                }
297            }
298        }
299        for edge in &self.edges {
300            if let GraphEdge::FanIn {
301                sources: edge_sources,
302                target,
303            } = edge
304            {
305                if edge_sources.iter().all(|s| sources.contains(s)) {
306                    return Some(target.clone());
307                }
308            }
309        }
310        None
311    }
312
313    pub(super) async fn execute_parallel_branches(
314        &self,
315        targets: &[String],
316        state: &S,
317    ) -> GraphResult<Vec<(String, GraphInvocation<S>)>> {
318        let futures: Vec<_> = targets
319            .iter()
320            .filter(|t| *t != END)
321            .map(|target| {
322                let target = target.clone();
323                let state_clone = state.clone();
324                async move {
325                    let result = self.invoke_from_node(target.clone(), state_clone).await;
326                    result.map(|inv| (target, inv))
327                }
328            })
329            .collect();
330
331        let results = join_all(futures).await;
332
333        let mut successful = Vec::new();
334        for result in results {
335            match result {
336                Ok((name, inv)) => successful.push((name, inv)),
337                Err(e) => return Err(e),
338            }
339        }
340
341        Ok(successful)
342    }
343}
lc_langgraph/compiled/stream.rs

lc_langgraph/compiled/
stream.rs
