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, 0).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 =
78                        checkpointer.lock().await.save(&state, recursion_count).await?;
79                    steps.push(ExecutionStep::checkpoint(
80                        checkpoint_id,
81                        current_node.clone(),
82                    ));
83                }
84                continue;
85            }
86
87            recursion_count += 1;
88
89            let node = self.get_node(&current_node).await?;
90
91            let config = NodeConfig {
92                recursion_limit: self.recursion_limit,
93                debug: false,
94                metadata: HashMap::new(),
95            };
96
97            let update = node.execute(&state, Some(config)).await?;
98
99            if let Some(new_state) = update.update {
100                state = self.default_reducer.reduce(&state, &new_state);
101            }
102
103            steps.push(ExecutionStep::node(
104                current_node.clone(),
105                update.metadata.clone(),
106            ));
107
108            if self.interrupt_after.contains(&current_node) {
109                return Err(GraphError::ExecutionInterrupted(format!(
110                    "after_{}",
111                    current_node
112                )));
113            }
114
115            let next_node = self.find_next_node(&current_node, &state).await?;
116
117            if let Some(ref checkpointer) = self.checkpointer {
118                let checkpoint_id =
119                    checkpointer.lock().await.save(&state, recursion_count).await?;
120                steps.push(ExecutionStep::checkpoint(checkpoint_id, next_node.clone()));
121            }
122
123            current_node = next_node;
124        }
125
126        Ok(GraphInvocation {
127            final_state: state,
128            steps,
129            recursion_count,
130        })
131    }
132
133    pub async fn invoke_with_execution(
134        &self,
135        execution: GraphExecution<S>,
136    ) -> GraphResult<GraphInvocation<S>> {
137        let mut state = execution.state;
138        let mut current_node = if execution.interrupted_at.starts_with("after_") {
139            self.find_next_node(&execution.current_node, &state).await?
140        } else {
141            execution.current_node
142        };
143        let mut steps = execution.steps;
144        let mut recursion_count = execution.recursion_count;
145        let first_node = current_node.clone();
146
147        loop {
148            if current_node == END {
149                break;
150            }
151
152            // Q3: check the limit before executing a step (same guard as `invoke`).
153            if recursion_count >= self.recursion_limit {
154                return Err(GraphError::RecursionLimitReached(self.recursion_limit));
155            }
156
157            if current_node != first_node && self.interrupt_before.contains(&current_node) {
158                return Err(GraphError::ExecutionInterrupted(current_node.clone()));
159            }
160
161            recursion_count += 1;
162
163            let node = self.get_node(&current_node).await?;
164
165            let config = NodeConfig {
166                recursion_limit: self.recursion_limit,
167                debug: false,
168                metadata: HashMap::new(),
169            };
170
171            let update = node.execute(&state, Some(config)).await?;
172
173            if let Some(new_state) = update.update {
174                state = self.default_reducer.reduce(&state, &new_state);
175            }
176
177            steps.push(ExecutionStep::node(
178                current_node.clone(),
179                update.metadata.clone(),
180            ));
181
182            if self.interrupt_after.contains(&current_node) {
183                return Err(GraphError::ExecutionInterrupted(format!(
184                    "after_{}",
185                    current_node
186                )));
187            }
188
189            let next_node = self.find_next_node(&current_node, &state).await?;
190
191            if let Some(ref checkpointer) = self.checkpointer {
192                let checkpoint_id =
193                    checkpointer.lock().await.save(&state, recursion_count).await?;
194                steps.push(ExecutionStep::checkpoint(checkpoint_id, next_node.clone()));
195            }
196
197            current_node = next_node;
198        }
199
200        Ok(GraphInvocation {
201            final_state: state,
202            steps,
203            recursion_count,
204        })
205    }
206
207    pub async fn resume(&self, execution: GraphExecution<S>) -> GraphResult<GraphInvocation<S>> {
208        self.invoke_with_execution(execution).await
209    }
210
211    pub async fn invoke_from_node(
212        &self,
213        start_node: String,
214        input: S,
215    ) -> GraphResult<GraphInvocation<S>> {
216        self.invoke_from_node_with_count(start_node, input, 0).await
217    }
218
219    /// Like [`invoke_from_node`](Self::invoke_from_node), but continues from an
220    /// existing recursion count.
221    ///
222    /// Q3/Q6: this is the single enforcement point for the recursion limit on
223    /// the "start from an arbitrary node" path. Parallel FanOut branches call
224    /// this with the main path's current count so their depth stays visible to
225    /// the shared `recursion_limit` budget instead of restarting from zero.
226    pub(super) async fn invoke_from_node_with_count(
227        &self,
228        start_node: String,
229        input: S,
230        mut recursion_count: usize,
231    ) -> GraphResult<GraphInvocation<S>> {
232        let mut state = input;
233        let mut current_node = start_node;
234        let mut steps: Vec<ExecutionStep> = Vec::new();
235
236        loop {
237            if current_node == END {
238                break;
239            }
240
241            if recursion_count >= self.recursion_limit {
242                return Err(GraphError::RecursionLimitReached(self.recursion_limit));
243            }
244
245            if self.interrupt_before.contains(&current_node) {
246                return Err(GraphError::ExecutionInterrupted(current_node.clone()));
247            }
248
249            recursion_count += 1;
250
251            let node = self.get_node(&current_node).await?;
252
253            let config = NodeConfig {
254                recursion_limit: self.recursion_limit,
255                debug: false,
256                metadata: HashMap::new(),
257            };
258
259            let update = node.execute(&state, Some(config)).await?;
260
261            if let Some(new_state) = update.update {
262                state = self.default_reducer.reduce(&state, &new_state);
263            }
264
265            steps.push(ExecutionStep::node(
266                current_node.clone(),
267                update.metadata.clone(),
268            ));
269
270            if self.interrupt_after.contains(&current_node) {
271                return Err(GraphError::ExecutionInterrupted(format!(
272                    "after_{}",
273                    current_node
274                )));
275            }
276
277            current_node = self.find_next_node(&current_node, &state).await?;
278        }
279
280        Ok(GraphInvocation {
281            final_state: state,
282            steps,
283            recursion_count,
284        })
285    }
286}