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 =
194            CompiledGraph::new(self.nodes.clone(), self.edges.clone(), entry, merge_reducer);
195
196        for (name, router) in &self.conditional_routers {
197            compiled.add_router(name.clone(), router.clone());
198        }
199
200        compiled.validate()?;
201        Ok(compiled)
202    }
203
204    fn find_first_node_after_start(&self) -> Option<String> {
205        for edge in &self.edges {
206            if edge.source() == START {
207                if let Some(target) = edge.fixed_target() {
208                    return Some(target.to_string());
209                }
210            }
211        }
212        None
213    }
214}
215
216impl<S: StateSchema + 'static> Default for StateGraph<S> {
217    fn default() -> Self {
218        Self::new()
219    }
220}
221
222/// GraphBuilder - Fluent builder pattern for StateGraph
223pub struct GraphBuilder<S: StateSchema> {
224    graph: StateGraph<S>,
225}
226
227impl<S: StateSchema + 'static> GraphBuilder<S> {
228    /// Create a new empty graph builder.
229    pub fn new() -> Self {
230        Self {
231            graph: StateGraph::new(),
232        }
233    }
234
235    /// Add a node to the graph.
236    pub fn add_node<N: GraphNode<S> + 'static>(mut self, node: N) -> Self {
237        self.graph.add_node(node);
238        self
239    }
240
241    /// Add a synchronous function node to the graph.
242    pub fn add_node_fn<F>(mut self, name: impl Into<String>, func: F) -> Self
243    where
244        F: Fn(&S) -> Result<StateUpdate<S>, GraphError> + Send + Sync + 'static,
245    {
246        self.graph.add_node_fn(name, func);
247        self
248    }
249
250    /// Add an async function node to the graph.
251    pub fn add_async_node<F>(mut self, name: impl Into<String>, func: F) -> Self
252    where
253        F: AsyncFn<S> + 'static,
254    {
255        self.graph.add_async_node(name, func);
256        self
257    }
258
259    /// Add a subgraph node with input/output mappers.
260    pub fn add_subgraph<SubS: StateSchema + 'static>(
261        mut self,
262        name: impl Into<String>,
263        subgraph: CompiledGraph<SubS>,
264        input_mapper: impl Fn(&S) -> SubS + Send + Sync + 'static,
265        output_mapper: impl Fn(&SubS, &mut S) + Send + Sync + 'static,
266    ) -> Self {
267        self.graph
268            .add_subgraph(name, subgraph, input_mapper, output_mapper);
269        self
270    }
271
272    /// Add a subgraph node that shares the same state type.
273    pub fn add_subgraph_same_state(
274        mut self,
275        name: impl Into<String>,
276        subgraph: CompiledGraph<S>,
277    ) -> Self {
278        self.graph.add_subgraph_same_state(name, subgraph);
279        self
280    }
281
282    /// Add a fixed edge between two nodes.
283    pub fn add_edge(mut self, source: impl Into<String>, target: impl Into<String>) -> Self {
284        self.graph.add_edge(source, target);
285        self
286    }
287
288    /// Add a conditional edge routed by a named router.
289    pub fn add_conditional_edges(
290        mut self,
291        source: impl Into<String>,
292        router_name: impl Into<String>,
293        targets: HashMap<String, String>,
294        default: Option<String>,
295    ) -> Self {
296        self.graph
297            .add_conditional_edges(source, router_name, targets, default);
298        self
299    }
300
301    /// Add a FanOut edge for parallel execution.
302    pub fn add_fan_out(mut self, source: impl Into<String>, targets: Vec<String>) -> Self {
303        self.graph.add_fan_out(source, targets);
304        self
305    }
306
307    /// Add a FanIn edge joining multiple sources into one target.
308    pub fn add_fan_in(mut self, sources: Vec<String>, target: impl Into<String>) -> Self {
309        self.graph.add_fan_in(sources, target);
310        self
311    }
312
313    /// Set the entry point node for the graph.
314    pub fn set_entry_point(mut self, node: impl Into<String>) -> Self {
315        self.graph.set_entry_point(node);
316        self
317    }
318
319    /// Compile the graph into an executable `CompiledGraph`.
320    pub fn compile(self) -> GraphResult<CompiledGraph<S>> {
321        self.graph.compile()
322    }
323
324    /// Build and return the underlying `StateGraph`.
325    pub fn build(self) -> StateGraph<S> {
326        self.graph
327    }
328}
329
330impl<S: StateSchema + 'static> Default for GraphBuilder<S> {
331    fn default() -> Self {
332        Self::new()
333    }
334}
335
336#[cfg(test)]
337mod tests {
338    use super::*;
339    use crate::state::AgentState;
340
341    #[test]
342    fn test_graph_creation() {
343        let graph: StateGraph<AgentState> = StateGraph::new();
344        assert!(graph.nodes.is_empty());
345        assert!(graph.edges.is_empty());
346    }
347
348    #[test]
349    fn test_add_node_fn() {
350        let mut graph: StateGraph<AgentState> = StateGraph::new();
351        graph.add_node_fn("test_node", |state: &AgentState| {
352            Ok(StateUpdate::full(state.clone()))
353        });
354        assert_eq!(graph.nodes.len(), 1);
355    }
356
357    #[test]
358    fn test_add_edge() {
359        let mut graph: StateGraph<AgentState> = StateGraph::new();
360        graph.add_edge(START, "node1");
361        graph.add_edge("node1", END);
362        assert_eq!(graph.edges.len(), 2);
363    }
364
365    #[test]
366    fn test_compile_empty_graph() {
367        let graph: StateGraph<AgentState> = StateGraph::new();
368        let result = graph.compile();
369        assert!(result.is_err());
370    }
371
372    #[test]
373    fn test_builder_pattern() {
374        let compiled = GraphBuilder::<AgentState>::new()
375            .add_node_fn("process", |state| Ok(StateUpdate::full(state.clone())))
376            .add_edge(START, "process")
377            .add_edge("process", END)
378            .compile();
379        assert!(compiled.is_ok());
380    }
381
382    // ------------------------------------------------------------------
383    // 0.22.0 C3 fixes: field reducers honored, merge on main-path base,
384    // stream fan-out runs all branches, FanIn topology compiles (H-A10).
385    // ------------------------------------------------------------------
386
387    use crate::compiled::types::StreamEvent;
388    use crate::state::{AppendMessagesReducer, MessageEntry, MessageRole};
389
390    /// C3-1: a field reducer registered via `set_reducer` merges its field
391    /// into the state instead of being dropped (previously the registered
392    /// reducer never reached `CompiledGraph` and every merge was replace).
393    #[tokio::test]
394    async fn test_field_reducer_accumulates_across_nodes() {
395        let mut graph: StateGraph<AgentState> = StateGraph::new();
396        graph.add_node_fn("a", |_state: &AgentState| {
397            // Partial update: only this node's message (drop the rest).
398            Ok(StateUpdate {
399                update: Some(AgentState {
400                    input: String::new(),
401                    messages: vec![MessageEntry {
402                        role: MessageRole::AI,
403                        content: "from-a".into(),
404                    }],
405                    steps: vec![],
406                    output: None,
407                }),
408                metadata: Default::default(),
409            })
410        });
411        graph.add_node_fn("b", |_state: &AgentState| {
412            Ok(StateUpdate {
413                update: Some(AgentState {
414                    input: String::new(),
415                    messages: vec![MessageEntry {
416                        role: MessageRole::AI,
417                        content: "from-b".into(),
418                    }],
419                    steps: vec![],
420                    output: None,
421                }),
422                metadata: Default::default(),
423            })
424        });
425        graph.set_entry_point("a");
426        graph.add_edge("a", "b");
427        graph.add_edge("b", END);
428        graph.set_reducer("messages", Arc::new(AppendMessagesReducer));
429
430        let compiled = graph.compile().unwrap();
431        let mut initial = AgentState::new("start"); // messages = [human "start"]
432        initial.messages.clear();
433        let result = compiled.invoke(initial).await.unwrap();
434        let contents: Vec<&str> = result
435            .final_state
436            .messages
437            .iter()
438            .map(|m| m.content.as_str())
439            .collect();
440        // Both node messages survived the merges (replace would keep only b's).
441        assert!(
442            contents.contains(&"from-a") && contents.contains(&"from-b"),
443            "field reducer should accumulate, got {contents:?}"
444        );
445    }
446
447    /// C3-2: `merge_parallel_states` folds on top of the pre-fan-out
448    /// main-path state — the fan-out source node's writes are not dropped.
449    #[tokio::test]
450    async fn test_parallel_merge_keeps_main_path_state() {
451        let mut graph: StateGraph<AgentState> = StateGraph::new();
452        graph.add_node_fn("source", |state: &AgentState| {
453            let mut s = state.clone();
454            s.messages.push(MessageEntry {
455                role: MessageRole::AI,
456                content: "main-path".into(),
457            });
458            Ok(StateUpdate::full(s))
459        });
460        graph.add_node_fn("branch1", |state: &AgentState| {
461            let mut s = state.clone();
462            s.messages.push(MessageEntry {
463                role: MessageRole::AI,
464                content: "b1".into(),
465            });
466            Ok(StateUpdate::full(s))
467        });
468        graph.add_node_fn("branch2", |state: &AgentState| {
469            let mut s = state.clone();
470            s.messages.push(MessageEntry {
471                role: MessageRole::AI,
472                content: "b2".into(),
473            });
474            Ok(StateUpdate::full(s))
475        });
476        graph.add_node_fn("merge", |state: &AgentState| {
477            let mut s = state.clone();
478            s.set_output("merged");
479            Ok(StateUpdate::full(s))
480        });
481        graph.set_entry_point("source");
482        graph.add_fan_out("source", vec!["branch1".into(), "branch2".into()]);
483        graph.add_fan_in(vec!["branch1".into(), "branch2".into()], "merge");
484        graph.add_edge("merge", END);
485        graph.set_reducer("messages", Arc::new(AppendMessagesReducer));
486
487        let compiled = graph.compile().unwrap();
488        let mut initial = AgentState::new("start");
489        initial.messages.clear();
490        let result = compiled.invoke(initial).await.unwrap();
491        let contents: Vec<&str> = result
492            .final_state
493            .messages
494            .iter()
495            .map(|m| m.content.as_str())
496            .collect();
497        // All three writers survive: source (main path) + both branches.
498        assert!(
499            contents.contains(&"main-path") && contents.contains(&"b1") && contents.contains(&"b2"),
500            "merge must keep main-path + all branch writes, got {contents:?}"
501        );
502        assert_eq!(result.final_state.output.as_deref(), Some("merged"));
503    }
504
505    /// C3-3: `CompiledGraph::stream` executes ALL fan-out branches (it used
506    /// to run only `targets[0]` and silently drop the rest).
507    #[tokio::test]
508    async fn test_stream_fan_out_runs_all_branches() {
509        let mut graph: StateGraph<AgentState> = StateGraph::new();
510        graph.add_node_fn("source", |state: &AgentState| {
511            Ok(StateUpdate::full(state.clone()))
512        });
513        graph.add_node_fn("branch1", |state: &AgentState| {
514            let mut s = state.clone();
515            s.messages.push(MessageEntry {
516                role: MessageRole::AI,
517                content: "s-b1".into(),
518            });
519            Ok(StateUpdate::full(s))
520        });
521        graph.add_node_fn("branch2", |state: &AgentState| {
522            let mut s = state.clone();
523            s.messages.push(MessageEntry {
524                role: MessageRole::AI,
525                content: "s-b2".into(),
526            });
527            Ok(StateUpdate::full(s))
528        });
529        graph.set_entry_point("source");
530        graph.add_fan_out("source", vec!["branch1".into(), "branch2".into()]);
531        graph.add_fan_in(vec!["branch1".into(), "branch2".into()], END);
532        graph.set_reducer("messages", Arc::new(AppendMessagesReducer));
533
534        let compiled = graph.compile().unwrap();
535        let mut initial = AgentState::new("start");
536        initial.messages.clear();
537        let events = compiled.stream_collected(initial).await.unwrap();
538        let end = events
539            .into_iter()
540            .find_map(|e| match e {
541                StreamEvent::End(s) => Some(s),
542                _ => None,
543            })
544            .expect("stream should end");
545        let contents: Vec<&str> = end.messages.iter().map(|m| m.content.as_str()).collect();
546        assert!(
547            contents.contains(&"s-b1") && contents.contains(&"s-b2"),
548            "stream must execute all fan-out branches, got {contents:?}"
549        );
550    }
551
552    /// H-A10: a standard fan-out + fan-in topology compiles — FanIn edges
553    /// were previously invisible to reachability (`source()` returned a
554    /// constant), so `compile()` rejected the graph with "Unreachable node".
555    #[test]
556    fn test_fan_in_topology_compiles() {
557        let mut graph: StateGraph<AgentState> = StateGraph::new();
558        graph.add_node_fn("source", |state: &AgentState| {
559            Ok(StateUpdate::full(state.clone()))
560        });
561        graph.add_node_fn("b1", |state: &AgentState| {
562            Ok(StateUpdate::full(state.clone()))
563        });
564        graph.add_node_fn("b2", |state: &AgentState| {
565            Ok(StateUpdate::full(state.clone()))
566        });
567        graph.add_node_fn("merge", |state: &AgentState| {
568            Ok(StateUpdate::full(state.clone()))
569        });
570        graph.set_entry_point("source");
571        graph.add_fan_out("source", vec!["b1".into(), "b2".into()]);
572        graph.add_fan_in(vec!["b1".into(), "b2".into()], "merge");
573        graph.add_edge("merge", END);
574        assert!(
575            graph.compile().is_ok(),
576            "fan-out + fan-in topology must compile"
577        );
578    }
579}
lc_langgraph/graph.rs

lc_langgraph/
graph.rs
