lc_langgraph/
graph.rs

1// crates/lc-langgraph/src/graph.rs
2//! StateGraph - Main graph class for LangGraph
3
4use crate::compiled::CompiledGraph;
5use crate::edge::{ConditionalEdge, GraphEdge};
6use crate::errors::{GraphError, GraphResult};
7use crate::node::{AsyncFn, AsyncNode, GraphNode, SyncNode};
8use crate::state::{MergeReducer, Reducer, ReplaceReducer, StateSchema, StateUpdate};
9use crate::subgraph::SubgraphNode;
10use std::collections::HashMap;
11use std::sync::Arc;
12
13/// START sentinel node identifier
14pub const START: &str = "__start__";
15/// END sentinel node identifier
16pub const END: &str = "__end__";
17
18/// StateGraph - Main graph builder
19///
20/// Manages nodes, edges, and state schema for graph execution.
21/// After building, compile to get an executable CompiledGraph.
22pub struct StateGraph<S: StateSchema> {
23    nodes: HashMap<String, Arc<dyn GraphNode<S>>>,
24    edges: Vec<GraphEdge>,
25    entry_point: Option<String>,
26    reducers: HashMap<String, Arc<dyn Reducer<S>>>,
27    default_reducer: Arc<dyn Reducer<S>>,
28    conditional_routers: HashMap<String, Arc<dyn ConditionalEdge<S>>>,
29}
30
31impl<S: StateSchema + 'static> StateGraph<S> {
32    /// Create a new empty state graph.
33    pub fn new() -> Self {
34        Self {
35            nodes: HashMap::new(),
36            edges: Vec::new(),
37            entry_point: None,
38            reducers: HashMap::new(),
39            default_reducer: Arc::new(ReplaceReducer),
40            conditional_routers: HashMap::new(),
41        }
42    }
43
44    /// Add a node to the graph.
45    pub fn add_node<N: GraphNode<S> + 'static>(&mut self, node: N) -> &mut Self {
46        let name = node.name().to_string();
47        self.nodes.insert(name, Arc::new(node));
48        self
49    }
50
51    /// Add a synchronous function node to the graph.
52    pub fn add_node_fn<F>(&mut self, name: impl Into<String>, func: F) -> &mut Self
53    where
54        F: Fn(&S) -> Result<StateUpdate<S>, GraphError> + Send + Sync + 'static,
55    {
56        let node = SyncNode::new(name, func);
57        let node_name = node.name().to_string();
58        self.nodes.insert(node_name, Arc::new(node));
59        self
60    }
61
62    /// Add an async function node to the graph.
63    pub fn add_async_node<F>(&mut self, name: impl Into<String>, func: F) -> &mut Self
64    where
65        F: AsyncFn<S> + 'static,
66    {
67        let node = AsyncNode::new(name, func);
68        let node_name = node.name().to_string();
69        self.nodes.insert(node_name, Arc::new(node));
70        self
71    }
72
73    /// Add a subgraph node with input/output mappers.
74    pub fn add_subgraph<SubS: StateSchema + 'static>(
75        &mut self,
76        name: impl Into<String>,
77        subgraph: CompiledGraph<SubS>,
78        input_mapper: impl Fn(&S) -> SubS + Send + Sync + 'static,
79        output_mapper: impl Fn(&SubS, &mut S) + Send + Sync + 'static,
80    ) -> &mut Self {
81        let node = SubgraphNode::new(name, subgraph, input_mapper, output_mapper);
82        let node_name = node.name().to_string();
83        self.nodes.insert(node_name, Arc::new(node));
84        self
85    }
86
87    /// Add a subgraph node that shares the same state type.
88    pub fn add_subgraph_same_state(
89        &mut self,
90        name: impl Into<String>,
91        subgraph: CompiledGraph<S>,
92    ) -> &mut Self {
93        let node: SubgraphNode<S, S> = SubgraphNode::same_state(name, subgraph);
94        let node_name = node.name().to_string();
95        self.nodes.insert(node_name, Arc::new(node));
96        self
97    }
98
99    /// Add a fixed edge between two nodes.
100    pub fn add_edge(&mut self, source: impl Into<String>, target: impl Into<String>) -> &mut Self {
101        let edge = GraphEdge::fixed(source, target);
102        self.edges.push(edge);
103        self
104    }
105
106    /// Add a conditional edge routed by a named router.
107    pub fn add_conditional_edges(
108        &mut self,
109        source: impl Into<String>,
110        router_name: impl Into<String>,
111        targets: HashMap<String, String>,
112        default: Option<String>,
113    ) -> &mut Self {
114        let edge = GraphEdge::conditional(source, router_name, targets, default);
115        self.edges.push(edge);
116        self
117    }
118
119    /// Add a FanOut edge for parallel execution.
120    pub fn add_fan_out(&mut self, source: impl Into<String>, targets: Vec<String>) -> &mut Self {
121        let edge = GraphEdge::fan_out(source, targets);
122        self.edges.push(edge);
123        self
124    }
125
126    /// Add a FanIn edge joining multiple sources into one target.
127    pub fn add_fan_in(&mut self, sources: Vec<String>, target: impl Into<String>) -> &mut Self {
128        let edge = GraphEdge::fan_in(sources, target);
129        self.edges.push(edge);
130        self
131    }
132
133    /// Register a conditional routing function by name.
134    pub fn set_conditional_router<R: ConditionalEdge<S> + 'static>(
135        &mut self,
136        name: impl Into<String>,
137        router: R,
138    ) -> &mut Self {
139        self.conditional_routers
140            .insert(name.into(), Arc::new(router));
141        self
142    }
143
144    /// Set the entry point node for the graph.
145    pub fn set_entry_point(&mut self, node: impl Into<String>) -> &mut Self {
146        self.entry_point = Some(node.into());
147        self
148    }
149
150    /// Register a reducer for a specific state field.
151    pub fn set_reducer(
152        &mut self,
153        field: impl Into<String>,
154        reducer: Arc<dyn Reducer<S>>,
155    ) -> &mut Self {
156        self.reducers.insert(field.into(), reducer);
157        self
158    }
159
160    /// Compile the graph into an executable `CompiledGraph`.
161    pub fn compile(&self) -> GraphResult<CompiledGraph<S>> {
162        if self.nodes.is_empty() {
163            return Err(GraphError::ValidationError(
164                "Graph has no nodes".to_string(),
165            ));
166        }
167
168        let entry = self
169            .entry_point
170            .clone()
171            .or_else(|| self.find_first_node_after_start())
172            .ok_or_else(|| GraphError::ValidationError("No entry point defined".to_string()))?;
173
174        if !self.nodes.contains_key(&entry) && entry != START {
175            return Err(GraphError::ValidationError(format!(
176                "Entry point '{}' not found",
177                entry
178            )));
179        }
180
181        // 0.22.0 C3 fix: honor the per-field reducers registered via
182        // `set_reducer`. Previously they were dropped here and every merge
183        // point used plain replace, silently overwriting accumulating fields.
184        let merge_reducer: Arc<dyn Reducer<S>> = if self.reducers.is_empty() {
185            self.default_reducer.clone()
186        } else {
187            Arc::new(MergeReducer::new(
188                self.default_reducer.clone(),
189                self.reducers.values().cloned().collect(),
190            ))
191        };
192
193        let mut compiled = CompiledGraph::new(
194            self.nodes.clone(),
195            self.edges.clone(),
196            entry,
197            merge_reducer,
198        );
199
200        for (name, router) in &self.conditional_routers {
201            compiled.add_router(name.clone(), router.clone());
202        }
203
204        compiled.validate()?;
205        Ok(compiled)
206    }
207
208    fn find_first_node_after_start(&self) -> Option<String> {
209        for edge in &self.edges {
210            if edge.source() == START {
211                if let Some(target) = edge.fixed_target() {
212                    return Some(target.to_string());
213                }
214            }
215        }
216        None
217    }
218}
219
220impl<S: StateSchema + 'static> Default for StateGraph<S> {
221    fn default() -> Self {
222        Self::new()
223    }
224}
225
226/// GraphBuilder - Fluent builder pattern for StateGraph
227pub struct GraphBuilder<S: StateSchema> {
228    graph: StateGraph<S>,
229}
230
231impl<S: StateSchema + 'static> GraphBuilder<S> {
232    /// Create a new empty graph builder.
233    pub fn new() -> Self {
234        Self {
235            graph: StateGraph::new(),
236        }
237    }
238
239    /// Add a node to the graph.
240    pub fn add_node<N: GraphNode<S> + 'static>(mut self, node: N) -> Self {
241        self.graph.add_node(node);
242        self
243    }
244
245    /// Add a synchronous function node to the graph.
246    pub fn add_node_fn<F>(mut self, name: impl Into<String>, func: F) -> Self
247    where
248        F: Fn(&S) -> Result<StateUpdate<S>, GraphError> + Send + Sync + 'static,
249    {
250        self.graph.add_node_fn(name, func);
251        self
252    }
253
254    /// Add an async function node to the graph.
255    pub fn add_async_node<F>(mut self, name: impl Into<String>, func: F) -> Self
256    where
257        F: AsyncFn<S> + 'static,
258    {
259        self.graph.add_async_node(name, func);
260        self
261    }
262
263    /// Add a subgraph node with input/output mappers.
264    pub fn add_subgraph<SubS: StateSchema + 'static>(
265        mut self,
266        name: impl Into<String>,
267        subgraph: CompiledGraph<SubS>,
268        input_mapper: impl Fn(&S) -> SubS + Send + Sync + 'static,
269        output_mapper: impl Fn(&SubS, &mut S) + Send + Sync + 'static,
270    ) -> Self {
271        self.graph
272            .add_subgraph(name, subgraph, input_mapper, output_mapper);
273        self
274    }
275
276    /// Add a subgraph node that shares the same state type.
277    pub fn add_subgraph_same_state(
278        mut self,
279        name: impl Into<String>,
280        subgraph: CompiledGraph<S>,
281    ) -> Self {
282        self.graph.add_subgraph_same_state(name, subgraph);
283        self
284    }
285
286    /// Add a fixed edge between two nodes.
287    pub fn add_edge(mut self, source: impl Into<String>, target: impl Into<String>) -> Self {
288        self.graph.add_edge(source, target);
289        self
290    }
291
292    /// Add a conditional edge routed by a named router.
293    pub fn add_conditional_edges(
294        mut self,
295        source: impl Into<String>,
296        router_name: impl Into<String>,
297        targets: HashMap<String, String>,
298        default: Option<String>,
299    ) -> Self {
300        self.graph
301            .add_conditional_edges(source, router_name, targets, default);
302        self
303    }
304
305    /// Add a FanOut edge for parallel execution.
306    pub fn add_fan_out(mut self, source: impl Into<String>, targets: Vec<String>) -> Self {
307        self.graph.add_fan_out(source, targets);
308        self
309    }
310
311    /// Add a FanIn edge joining multiple sources into one target.
312    pub fn add_fan_in(mut self, sources: Vec<String>, target: impl Into<String>) -> Self {
313        self.graph.add_fan_in(sources, target);
314        self
315    }
316
317    /// Set the entry point node for the graph.
318    pub fn set_entry_point(mut self, node: impl Into<String>) -> Self {
319        self.graph.set_entry_point(node);
320        self
321    }
322
323    /// Compile the graph into an executable `CompiledGraph`.
324    pub fn compile(self) -> GraphResult<CompiledGraph<S>> {
325        self.graph.compile()
326    }
327
328    /// Build and return the underlying `StateGraph`.
329    pub fn build(self) -> StateGraph<S> {
330        self.graph
331    }
332}
333
334impl<S: StateSchema + 'static> Default for GraphBuilder<S> {
335    fn default() -> Self {
336        Self::new()
337    }
338}
339
340#[cfg(test)]
341mod tests {
342    use super::*;
343    use crate::state::AgentState;
344
345    #[test]
346    fn test_graph_creation() {
347        let graph: StateGraph<AgentState> = StateGraph::new();
348        assert!(graph.nodes.is_empty());
349        assert!(graph.edges.is_empty());
350    }
351
352    #[test]
353    fn test_add_node_fn() {
354        let mut graph: StateGraph<AgentState> = StateGraph::new();
355        graph.add_node_fn("test_node", |state: &AgentState| {
356            Ok(StateUpdate::full(state.clone()))
357        });
358        assert_eq!(graph.nodes.len(), 1);
359    }
360
361    #[test]
362    fn test_add_edge() {
363        let mut graph: StateGraph<AgentState> = StateGraph::new();
364        graph.add_edge(START, "node1");
365        graph.add_edge("node1", END);
366        assert_eq!(graph.edges.len(), 2);
367    }
368
369    #[test]
370    fn test_compile_empty_graph() {
371        let graph: StateGraph<AgentState> = StateGraph::new();
372        let result = graph.compile();
373        assert!(result.is_err());
374    }
375
376    #[test]
377    fn test_builder_pattern() {
378        let compiled = GraphBuilder::<AgentState>::new()
379            .add_node_fn("process", |state| Ok(StateUpdate::full(state.clone())))
380            .add_edge(START, "process")
381            .add_edge("process", END)
382            .compile();
383        assert!(compiled.is_ok());
384    }
385
386    // ------------------------------------------------------------------
387    // 0.22.0 C3 fixes: field reducers honored, merge on main-path base,
388    // stream fan-out runs all branches, FanIn topology compiles (H-A10).
389    // ------------------------------------------------------------------
390
391    use crate::compiled::types::StreamEvent;
392    use crate::state::{AppendMessagesReducer, MessageEntry, MessageRole};
393    use futures_util::StreamExt;
394
395    /// C3-1: a field reducer registered via `set_reducer` merges its field
396    /// into the state instead of being dropped (previously the registered
397    /// reducer never reached `CompiledGraph` and every merge was replace).
398    #[tokio::test]
399    async fn test_field_reducer_accumulates_across_nodes() {
400        let mut graph: StateGraph<AgentState> = StateGraph::new();
401        graph.add_node_fn("a", |_state: &AgentState| {
402            // Partial update: only this node's message (drop the rest).
403            Ok(StateUpdate {
404                update: Some(AgentState {
405                    input: String::new(),
406                    messages: vec![MessageEntry {
407                        role: MessageRole::AI,
408                        content: "from-a".into(),
409                    }],
410                    steps: vec![],
411                    output: None,
412                }),
413                metadata: Default::default(),
414            })
415        });
416        graph.add_node_fn("b", |_state: &AgentState| {
417            Ok(StateUpdate {
418                update: Some(AgentState {
419                    input: String::new(),
420                    messages: vec![MessageEntry {
421                        role: MessageRole::AI,
422                        content: "from-b".into(),
423                    }],
424                    steps: vec![],
425                    output: None,
426                }),
427                metadata: Default::default(),
428            })
429        });
430        graph.set_entry_point("a");
431        graph.add_edge("a", "b");
432        graph.add_edge("b", END);
433        graph.set_reducer("messages", Arc::new(AppendMessagesReducer));
434
435        let compiled = graph.compile().unwrap();
436        let mut initial = AgentState::new("start"); // messages = [human "start"]
437        initial.messages.clear();
438        let result = compiled.invoke(initial).await.unwrap();
439        let contents: Vec<&str> = result
440            .final_state
441            .messages
442            .iter()
443            .map(|m| m.content.as_str())
444            .collect();
445        // Both node messages survived the merges (replace would keep only b's).
446        assert!(
447            contents.contains(&"from-a") && contents.contains(&"from-b"),
448            "field reducer should accumulate, got {contents:?}"
449        );
450    }
451
452    /// C3-2: `merge_parallel_states` folds on top of the pre-fan-out
453    /// main-path state — the fan-out source node's writes are not dropped.
454    #[tokio::test]
455    async fn test_parallel_merge_keeps_main_path_state() {
456        let mut graph: StateGraph<AgentState> = StateGraph::new();
457        graph.add_node_fn("source", |state: &AgentState| {
458            let mut s = state.clone();
459            s.messages.push(MessageEntry {
460                role: MessageRole::AI,
461                content: "main-path".into(),
462            });
463            Ok(StateUpdate::full(s))
464        });
465        graph.add_node_fn("branch1", |state: &AgentState| {
466            let mut s = state.clone();
467            s.messages.push(MessageEntry {
468                role: MessageRole::AI,
469                content: "b1".into(),
470            });
471            Ok(StateUpdate::full(s))
472        });
473        graph.add_node_fn("branch2", |state: &AgentState| {
474            let mut s = state.clone();
475            s.messages.push(MessageEntry {
476                role: MessageRole::AI,
477                content: "b2".into(),
478            });
479            Ok(StateUpdate::full(s))
480        });
481        graph.add_node_fn("merge", |state: &AgentState| {
482            let mut s = state.clone();
483            s.set_output("merged");
484            Ok(StateUpdate::full(s))
485        });
486        graph.set_entry_point("source");
487        graph.add_fan_out("source", vec!["branch1".into(), "branch2".into()]);
488        graph.add_fan_in(vec!["branch1".into(), "branch2".into()], "merge");
489        graph.add_edge("merge", END);
490        graph.set_reducer("messages", Arc::new(AppendMessagesReducer));
491
492        let compiled = graph.compile().unwrap();
493        let mut initial = AgentState::new("start");
494        initial.messages.clear();
495        let result = compiled.invoke(initial).await.unwrap();
496        let contents: Vec<&str> = result
497            .final_state
498            .messages
499            .iter()
500            .map(|m| m.content.as_str())
501            .collect();
502        // All three writers survive: source (main path) + both branches.
503        assert!(
504            contents.contains(&"main-path")
505                && contents.contains(&"b1")
506                && contents.contains(&"b2"),
507            "merge must keep main-path + all branch writes, got {contents:?}"
508        );
509        assert_eq!(result.final_state.output.as_deref(), Some("merged"));
510    }
511
512    /// C3-3: `CompiledGraph::stream` executes ALL fan-out branches (it used
513    /// to run only `targets[0]` and silently drop the rest).
514    #[tokio::test]
515    async fn test_stream_fan_out_runs_all_branches() {
516        let mut graph: StateGraph<AgentState> = StateGraph::new();
517        graph.add_node_fn("source", |state: &AgentState| {
518            Ok(StateUpdate::full(state.clone()))
519        });
520        graph.add_node_fn("branch1", |state: &AgentState| {
521            let mut s = state.clone();
522            s.messages.push(MessageEntry {
523                role: MessageRole::AI,
524                content: "s-b1".into(),
525            });
526            Ok(StateUpdate::full(s))
527        });
528        graph.add_node_fn("branch2", |state: &AgentState| {
529            let mut s = state.clone();
530            s.messages.push(MessageEntry {
531                role: MessageRole::AI,
532                content: "s-b2".into(),
533            });
534            Ok(StateUpdate::full(s))
535        });
536        graph.set_entry_point("source");
537        graph.add_fan_out("source", vec!["branch1".into(), "branch2".into()]);
538        graph.add_fan_in(vec!["branch1".into(), "branch2".into()], END);
539        graph.set_reducer("messages", Arc::new(AppendMessagesReducer));
540
541        let compiled = graph.compile().unwrap();
542        let mut initial = AgentState::new("start");
543        initial.messages.clear();
544        let events = compiled.stream_collected(initial).await.unwrap();
545        let end = events
546            .into_iter()
547            .find_map(|e| match e {
548                StreamEvent::End(s) => Some(s),
549                _ => None,
550            })
551            .expect("stream should end");
552        let contents: Vec<&str> = end
553            .messages
554            .iter()
555            .map(|m| m.content.as_str())
556            .collect();
557        assert!(
558            contents.contains(&"s-b1") && contents.contains(&"s-b2"),
559            "stream must execute all fan-out branches, got {contents:?}"
560        );
561    }
562
563    /// H-A10: a standard fan-out + fan-in topology compiles — FanIn edges
564    /// were previously invisible to reachability (`source()` returned a
565    /// constant), so `compile()` rejected the graph with "Unreachable node".
566    #[test]
567    fn test_fan_in_topology_compiles() {
568        let mut graph: StateGraph<AgentState> = StateGraph::new();
569        graph.add_node_fn("source", |state: &AgentState| {
570            Ok(StateUpdate::full(state.clone()))
571        });
572        graph.add_node_fn("b1", |state: &AgentState| Ok(StateUpdate::full(state.clone())));
573        graph.add_node_fn("b2", |state: &AgentState| Ok(StateUpdate::full(state.clone())));
574        graph.add_node_fn("merge", |state: &AgentState| Ok(StateUpdate::full(state.clone())));
575        graph.set_entry_point("source");
576        graph.add_fan_out("source", vec!["b1".into(), "b2".into()]);
577        graph.add_fan_in(vec!["b1".into(), "b2".into()], "merge");
578        graph.add_edge("merge", END);
579        assert!(
580            graph.compile().is_ok(),
581            "fan-out + fan-in topology must compile"
582        );
583    }
584}
lc_langgraph/graph.rs

lc_langgraph/
graph.rs
