Skip to main content

lc_langgraph/compiled/
invoke.rs

1// crates/lc-langgraph/src/compiled/invoke.rs
2//! CompiledGraph invoke, invoke_with_execution, resume, and invoke_from_node methods
3
4use super::graph::CompiledGraph;
5use super::types::{ExecutionStep, GraphExecution, GraphInvocation, ParallelBranch};
6use crate::errors::{GraphError, GraphResult};
7use crate::node::NodeConfig;
8use crate::state::StateSchema;
9use crate::END;
10use std::collections::HashMap;
11
12impl<S: StateSchema> CompiledGraph<S> {
13    pub async fn invoke(&self, input: S) -> GraphResult<GraphInvocation<S>> {
14        let mut state = input;
15        let mut current_node = self.entry_point.clone();
16        let mut steps: Vec<ExecutionStep> = Vec::new();
17        let mut recursion_count = 0;
18
19        if let Some(ref checkpointer) = self.checkpointer {
20            let checkpoint_id = checkpointer.lock().await.save(&state).await?;
21            steps.push(ExecutionStep::checkpoint(
22                checkpoint_id,
23                current_node.clone(),
24            ));
25        }
26
27        loop {
28            if current_node == END {
29                break;
30            }
31
32            // Q3: check the limit BEFORE executing a step. A `loop` with this
33            // guard means a graph that legitimately uses exactly `limit` steps
34            // and then reaches END is NOT misreported as exceeding the limit
35            // (the old `count >= limit` post-check fired at `count == limit`).
36            if recursion_count >= self.recursion_limit {
37                return Err(GraphError::RecursionLimitReached(self.recursion_limit));
38            }
39
40            if self.interrupt_before.contains(&current_node) {
41                return Err(GraphError::ExecutionInterrupted(current_node.clone()));
42            }
43
44            // Check for FanOut edge — execute all branches in parallel and merge
45            let fan_out_targets = self.find_fan_out_targets(&current_node).await;
46            if let Some(targets) = fan_out_targets {
47                recursion_count += 1;
48                let mut parallel_branches: Vec<ParallelBranch<S>> = Vec::new();
49
50                // Q6: branches share the main-path recursion budget so the limit
51                // cannot be bypassed by fanning out into deep sub-executions.
52                let branch_results = self
53                    .execute_parallel_branches(&targets, &state, recursion_count)
54                    .await?;
55                for (name, inv) in branch_results {
56                    parallel_branches.push(ParallelBranch {
57                        name: name.clone(),
58                        final_state: inv.final_state.clone(),
59                        steps: inv.steps.clone(),
60                    });
61                    steps.push(ExecutionStep::ParallelNode {
62                        branch: name,
63                        metadata: HashMap::new(),
64                    });
65                }
66
67                let merge_target = self.find_fan_in_target(&targets).await;
68                if let Some(merge_node) = merge_target {
69                    state = self.merge_parallel_states(&parallel_branches)?;
70                    current_node = merge_node;
71                } else {
72                    state = self.merge_parallel_states(&parallel_branches)?;
73                    current_node = END.to_string();
74                }
75
76                if let Some(ref checkpointer) = self.checkpointer {
77                    let checkpoint_id = checkpointer.lock().await.save(&state).await?;
78                    steps.push(ExecutionStep::checkpoint(
79                        checkpoint_id,
80                        current_node.clone(),
81                    ));
82                }
83                continue;
84            }
85
86            recursion_count += 1;
87
88            let node = self.get_node(&current_node).await?;
89
90            let config = NodeConfig {
91                recursion_limit: self.recursion_limit,
92                debug: false,
93                metadata: HashMap::new(),
94            };
95
96            let update = node.execute(&state, Some(config)).await?;
97
98            if let Some(new_state) = update.update {
99                state = self.default_reducer.reduce(&state, &new_state);
100            }
101
102            steps.push(ExecutionStep::node(
103                current_node.clone(),
104                update.metadata.clone(),
105            ));
106
107            if self.interrupt_after.contains(&current_node) {
108                return Err(GraphError::ExecutionInterrupted(format!(
109                    "after_{}",
110                    current_node
111                )));
112            }
113
114            let next_node = self.find_next_node(&current_node, &state).await?;
115
116            if let Some(ref checkpointer) = self.checkpointer {
117                let checkpoint_id = checkpointer.lock().await.save(&state).await?;
118                steps.push(ExecutionStep::checkpoint(checkpoint_id, next_node.clone()));
119            }
120
121            current_node = next_node;
122        }
123
124        Ok(GraphInvocation {
125            final_state: state,
126            steps,
127            recursion_count,
128        })
129    }
130
131    pub async fn invoke_with_execution(
132        &self,
133        execution: GraphExecution<S>,
134    ) -> GraphResult<GraphInvocation<S>> {
135        let mut state = execution.state;
136        let mut current_node = if execution.interrupted_at.starts_with("after_") {
137            self.find_next_node(&execution.current_node, &state).await?
138        } else {
139            execution.current_node
140        };
141        let mut steps = execution.steps;
142        let mut recursion_count = execution.recursion_count;
143        let first_node = current_node.clone();
144
145        loop {
146            if current_node == END {
147                break;
148            }
149
150            // Q3: check the limit before executing a step (same guard as `invoke`).
151            if recursion_count >= self.recursion_limit {
152                return Err(GraphError::RecursionLimitReached(self.recursion_limit));
153            }
154
155            if current_node != first_node && self.interrupt_before.contains(&current_node) {
156                return Err(GraphError::ExecutionInterrupted(current_node.clone()));
157            }
158
159            recursion_count += 1;
160
161            let node = self.get_node(&current_node).await?;
162
163            let config = NodeConfig {
164                recursion_limit: self.recursion_limit,
165                debug: false,
166                metadata: HashMap::new(),
167            };
168
169            let update = node.execute(&state, Some(config)).await?;
170
171            if let Some(new_state) = update.update {
172                state = self.default_reducer.reduce(&state, &new_state);
173            }
174
175            steps.push(ExecutionStep::node(
176                current_node.clone(),
177                update.metadata.clone(),
178            ));
179
180            if self.interrupt_after.contains(&current_node) {
181                return Err(GraphError::ExecutionInterrupted(format!(
182                    "after_{}",
183                    current_node
184                )));
185            }
186
187            let next_node = self.find_next_node(&current_node, &state).await?;
188
189            if let Some(ref checkpointer) = self.checkpointer {
190                let checkpoint_id = checkpointer.lock().await.save(&state).await?;
191                steps.push(ExecutionStep::checkpoint(checkpoint_id, next_node.clone()));
192            }
193
194            current_node = next_node;
195        }
196
197        Ok(GraphInvocation {
198            final_state: state,
199            steps,
200            recursion_count,
201        })
202    }
203
204    pub async fn resume(&self, execution: GraphExecution<S>) -> GraphResult<GraphInvocation<S>> {
205        self.invoke_with_execution(execution).await
206    }
207
208    pub async fn invoke_from_node(
209        &self,
210        start_node: String,
211        input: S,
212    ) -> GraphResult<GraphInvocation<S>> {
213        self.invoke_from_node_with_count(start_node, input, 0).await
214    }
215
216    /// Like [`invoke_from_node`](Self::invoke_from_node), but continues from an
217    /// existing recursion count.
218    ///
219    /// Q3/Q6: this is the single enforcement point for the recursion limit on
220    /// the "start from an arbitrary node" path. Parallel FanOut branches call
221    /// this with the main path's current count so their depth stays visible to
222    /// the shared `recursion_limit` budget instead of restarting from zero.
223    pub(super) async fn invoke_from_node_with_count(
224        &self,
225        start_node: String,
226        input: S,
227        mut recursion_count: usize,
228    ) -> GraphResult<GraphInvocation<S>> {
229        let mut state = input;
230        let mut current_node = start_node;
231        let mut steps: Vec<ExecutionStep> = Vec::new();
232
233        loop {
234            if current_node == END {
235                break;
236            }
237
238            if recursion_count >= self.recursion_limit {
239                return Err(GraphError::RecursionLimitReached(self.recursion_limit));
240            }
241
242            if self.interrupt_before.contains(&current_node) {
243                return Err(GraphError::ExecutionInterrupted(current_node.clone()));
244            }
245
246            recursion_count += 1;
247
248            let node = self.get_node(&current_node).await?;
249
250            let config = NodeConfig {
251                recursion_limit: self.recursion_limit,
252                debug: false,
253                metadata: HashMap::new(),
254            };
255
256            let update = node.execute(&state, Some(config)).await?;
257
258            if let Some(new_state) = update.update {
259                state = self.default_reducer.reduce(&state, &new_state);
260            }
261
262            steps.push(ExecutionStep::node(
263                current_node.clone(),
264                update.metadata.clone(),
265            ));
266
267            if self.interrupt_after.contains(&current_node) {
268                return Err(GraphError::ExecutionInterrupted(format!(
269                    "after_{}",
270                    current_node
271                )));
272            }
273
274            current_node = self.find_next_node(&current_node, &state).await?;
275        }
276
277        Ok(GraphInvocation {
278            final_state: state,
279            steps,
280            recursion_count,
281        })
282    }
283}