Skip to main content

lc_langgraph/compiled/
graph.rs

1// crates/lc-langgraph/src/compiled/graph.rs
2//! CompiledGraph struct definition, constructors, getters, and runtime injection
3
4use super::types::{DynamicTask, GraphExecution, GraphInvocation, PendingInterrupt};
5use crate::checkpointer::{CheckpointInfo, Checkpointer};
6use crate::edge::{ConditionalEdge, GraphEdge};
7use crate::errors::{GraphError, GraphResult};
8use crate::node::GraphNode;
9use crate::state::{Reducer, StateSchema};
10use crate::subgraph::SubgraphNode;
11use std::collections::HashMap;
12use std::sync::Arc;
13use tokio::sync::Mutex;
14use tokio::sync::RwLock;
15
16/// CompiledGraph - Ready-to-execute graph
17///
18/// Created from StateGraph.compile(). Handles:
19/// - State management and updates
20/// - Edge routing (fixed and conditional)
21/// - Execution with recursion limits
22/// - Checkpointing for persistence
23#[allow(clippy::type_complexity)]
24#[derive(Clone)]
25pub struct CompiledGraph<S: StateSchema> {
26    pub(super) nodes: HashMap<String, Arc<dyn GraphNode<S>>>,
27    pub(super) edges: Vec<GraphEdge>,
28    pub(super) entry_point: String,
29    pub(super) default_reducer: Arc<dyn Reducer<S>>,
30    pub(super) conditional_routers: HashMap<String, Arc<dyn ConditionalEdge<S>>>,
31    pub(super) checkpointer: Option<Arc<Mutex<dyn Checkpointer<S> + Send>>>,
32    pub(super) recursion_limit: usize,
33    pub(super) interrupt_before: Vec<String>,
34    pub(super) interrupt_after: Vec<String>,
35
36    pub(super) runtime_nodes: Arc<RwLock<HashMap<String, Arc<dyn GraphNode<S>>>>>,
37    pub(super) runtime_edges: Arc<RwLock<Vec<GraphEdge>>>,
38    pub(super) runtime_conditional_routers:
39        Arc<RwLock<HashMap<String, Arc<dyn ConditionalEdge<S>>>>>,
40    pub(super) task_inbox: Arc<Mutex<Vec<DynamicTask>>>,
41}
42
43impl<S: StateSchema> std::fmt::Debug for CompiledGraph<S> {
44    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
45        f.debug_struct("CompiledGraph")
46            .field("nodes", &self.nodes.keys().collect::<Vec<_>>())
47            .field("edges", &self.edges)
48            .field("entry_point", &self.entry_point)
49            .field("recursion_limit", &self.recursion_limit)
50            .field("interrupt_before", &self.interrupt_before)
51            .field("interrupt_after", &self.interrupt_after)
52            .finish()
53    }
54}
55
56impl<S: StateSchema> CompiledGraph<S> {
57    pub(crate) fn new(
58        nodes: HashMap<String, Arc<dyn GraphNode<S>>>,
59        edges: Vec<GraphEdge>,
60        entry_point: String,
61        default_reducer: Arc<dyn Reducer<S>>,
62    ) -> Self {
63        Self {
64            nodes,
65            edges,
66            entry_point,
67            default_reducer,
68            conditional_routers: HashMap::new(),
69            checkpointer: None,
70            recursion_limit: 25,
71            interrupt_before: Vec::new(),
72            interrupt_after: Vec::new(),
73            runtime_nodes: Arc::new(RwLock::new(HashMap::new())),
74            runtime_edges: Arc::new(RwLock::new(Vec::new())),
75            runtime_conditional_routers: Arc::new(RwLock::new(HashMap::new())),
76            task_inbox: Arc::new(Mutex::new(Vec::new())),
77        }
78    }
79
80    pub(crate) fn add_router(&mut self, name: String, router: Arc<dyn ConditionalEdge<S>>) {
81        self.conditional_routers.insert(name, router);
82    }
83
84    /// Attach a checkpointer to this graph for state persistence.
85    pub fn with_checkpointer<C: Checkpointer<S> + 'static>(mut self, checkpointer: C) -> Self {
86        self.checkpointer = Some(Arc::new(Mutex::new(checkpointer)));
87        self
88    }
89
90    /// Set the maximum recursion depth before execution aborts.
91    pub fn with_recursion_limit(mut self, limit: usize) -> Self {
92        self.recursion_limit = limit;
93        self
94    }
95
96    /// Set the node names before which execution should be interrupted.
97    pub fn with_interrupt_before(mut self, nodes: Vec<String>) -> Self {
98        self.interrupt_before = nodes;
99        self
100    }
101
102    /// Set the node names after which execution should be interrupted.
103    pub fn with_interrupt_after(mut self, nodes: Vec<String>) -> Self {
104        self.interrupt_after = nodes;
105        self
106    }
107
108    /// Return the names of all static nodes registered in the graph.
109    pub fn node_names(&self) -> Vec<String> {
110        self.nodes.keys().cloned().collect()
111    }
112
113    /// Return the graph's fixed edges.
114    pub fn get_edges(&self) -> &[GraphEdge] {
115        &self.edges
116    }
117
118    /// Return the entry point node name.
119    pub fn entry_point(&self) -> &str {
120        &self.entry_point
121    }
122
123    /// Return the configured recursion limit.
124    pub fn recursion_limit(&self) -> usize {
125        self.recursion_limit
126    }
127
128    /// Return the nodes that interrupt execution before running.
129    pub fn interrupt_before(&self) -> &[String] {
130        &self.interrupt_before
131    }
132
133    /// Return the nodes that interrupt execution after running.
134    pub fn interrupt_after(&self) -> &[String] {
135        &self.interrupt_after
136    }
137
138    /// Get the last checkpoint state (for interrupt recovery), together with the
139    /// recursion budget consumed up to that checkpoint.
140    ///
141    /// H5: no longer guesses "most recent" from `list()`'s HashMap order + `ids.last()`;
142    /// instead the checkpointer's `last()` returns the truly newest checkpoint by
143    /// (timestamp, seq).
144    pub async fn last_checkpoint_state(&self) -> Option<(S, usize)> {
145        if let Some(ref cp) = self.checkpointer {
146            let guard = cp.lock().await;
147            guard.last().await.ok()?
148        } else {
149            None
150        }
151    }
152
153    /// Return a locked handle to the attached checkpointer, or an error when none
154    /// is attached (state history / time-travel need one) or it is contended.
155    async fn checkpointer_locked(
156        &self,
157    ) -> GraphResult<tokio::sync::MutexGuard<'_, dyn Checkpointer<S> + Send>> {
158        let Some(cp) = &self.checkpointer else {
159            return Err(GraphError::CheckpointError(
160                "state history / time-travel require a checkpointer (attach one with with_checkpointer)"
161                    .to_string(),
162            ));
163        };
164        Ok(cp.lock().await)
165    }
166
167    /// State history: every step's checkpoint snapshot, oldest first. This is the
168    /// "agent 当时为什么这么决策" inspection primitive — each snapshot exposes the
169    /// state, ordering keys, and recursion budget at that point in the run.
170    pub async fn get_state_history(&self) -> GraphResult<Vec<CheckpointInfo<S>>> {
171        let guard = self.checkpointer_locked().await?;
172        guard.snapshots().await
173    }
174
175    /// Time-travel: fork a fresh execution timeline from an old checkpoint.
176    ///
177    /// - `checkpoint_id`: the snapshot to branch from (see [`get_state_history`](Self::get_state_history)).
178    /// - `continue_at_node`: the node the forked run resumes at — normally the
179    ///   `next_node` recorded in the history step immediately after that snapshot.
180    /// - `override_state`: optionally replace the snapshot state with new input
181    ///   ("改输入、从那条线分叉重跑") before running forward.
182    ///
183    /// The fork is seeded as a NEW checkpoint so the original history is untouched
184    /// (its own lineage), continues counting against the snapshot's recursion
185    /// budget, and — because it starts at `continue_at_node` and only runs
186    /// FORWARD via `invoke_from_node` — never replays side effects emitted before
187    /// the fork point.
188    pub async fn fork_from(
189        &self,
190        checkpoint_id: &str,
191        continue_at_node: impl Into<String>,
192        override_state: Option<S>,
193    ) -> GraphResult<GraphInvocation<S>> {
194        let snapshot = {
195            let guard = self.checkpointer_locked().await?;
196            guard
197                .snapshots()
198                .await?
199                .into_iter()
200                .find(|s| s.id == checkpoint_id)
201                .ok_or_else(|| {
202                    GraphError::CheckpointError(format!("Checkpoint '{checkpoint_id}' not found"))
203                })?
204        };
205        let base_state = override_state.unwrap_or(snapshot.state);
206
207        // Seed a fresh checkpoint so the fork lives in its own lineage.
208        let guard = self.checkpointer_locked().await?;
209        guard.save(&base_state, snapshot.recursion_count).await?;
210        drop(guard);
211
212        self.invoke_from_node_with_count(
213            continue_at_node.into(),
214            base_state,
215            snapshot.recursion_count,
216        )
217        .await
218    }
219
220    /// Create a resume execution context (from the last checkpoint)
221    /// interrupted_node may be "node_name" or "after_node_name"
222    pub async fn create_resume_execution(
223        &self,
224        interrupted_node: &str,
225    ) -> Option<GraphExecution<S>> {
226        let (state, recursion_count) = self.last_checkpoint_state().await?;
227        let current = interrupted_node
228            .strip_prefix("after_")
229            .unwrap_or(interrupted_node);
230        Some(GraphExecution {
231            state,
232            current_node: current.to_string(),
233            steps: Vec::new(),
234            // M6: resume carries over the recursion budget already consumed before the
235            // interrupt, instead of restarting from zero — otherwise repeated
236            // interrupt→resume could bypass recursion_limit indefinitely.
237            recursion_count,
238            interrupted_at: interrupted_node.to_string(),
239            pending_interrupt: None,
240        })
241    }
242
243    /// Resume from a runtime interrupt at `interrupted_node`, feeding the
244    /// human's `value` back into that node.
245    ///
246    /// The interrupted node is re-entered (LangGraph-style) and the value is
247    /// injected into its `NodeConfig.metadata` under
248    /// [`crate::node::INTERRUPT_RESUME_KEY`], so an [`crate::node::InterruptibleNode`]
249    /// closure can branch on `Some(decision)` to continue past the suspension
250    /// without re-running side effects. Requires a checkpointer with the
251    /// checkpoint that was persisted when the interrupt fired.
252    pub async fn resume_with_value(
253        &self,
254        interrupted_node: &str,
255        value: serde_json::Value,
256    ) -> GraphResult<GraphInvocation<S>> {
257        let mut execution = self
258            .create_resume_execution(interrupted_node)
259            .await
260            .ok_or_else(|| {
261                GraphError::ResumeError(format!(
262                    "cannot resume '{}': no checkpoint to resume from (attach a checkpointer)",
263                    interrupted_node
264                ))
265            })?;
266        execution.pending_interrupt = Some(PendingInterrupt {
267            id: interrupted_node.to_string(),
268            value,
269        });
270        self.invoke_with_execution(execution).await
271    }
272
273    pub(super) async fn get_node(&self, name: &str) -> GraphResult<Arc<dyn GraphNode<S>>> {
274        if let Some(node) = self.nodes.get(name) {
275            return Ok(node.clone());
276        }
277        if let Some(node) = self.runtime_nodes.read().await.get(name) {
278            return Ok(node.clone());
279        }
280        Err(GraphError::ExecutionError(format!(
281            "Node '{}' not found",
282            name
283        )))
284    }
285
286    /// Submit a new task for dynamic planning mid-execution.
287    ///
288    /// Returns an error instead of silently dropping the task when the inbox
289    /// lock is contended (previously a failed `try_lock` swallowed it, so a task
290    /// could vanish without any signal).
291    pub fn submit_task(&self, description: impl Into<String>) -> Result<(), GraphError> {
292        let mut inbox = self.task_inbox.try_lock().map_err(|_| {
293            GraphError::ExecutionError(
294                "dynamic task inbox is busy; task refused (not silently dropped)".to_string(),
295            )
296        })?;
297        inbox.push(DynamicTask {
298            id: uuid::Uuid::new_v4().to_string(),
299            description: description.into(),
300        });
301        Ok(())
302    }
303
304    /// Inject a runtime node directly
305    pub async fn inject_node(&self, name: &str, node: Arc<dyn GraphNode<S>>) -> GraphResult<()> {
306        self.runtime_nodes
307            .write()
308            .await
309            .insert(name.to_string(), node);
310        Ok(())
311    }
312
313    /// Inject a runtime edge directly
314    pub async fn inject_edge(&self, source: &str, target: &str) -> GraphResult<()> {
315        self.runtime_edges
316            .write()
317            .await
318            .push(GraphEdge::fixed(source, target));
319        Ok(())
320    }
321
322    /// Inject a subgraph as a runtime node
323    pub async fn inject_subgraph<SubS: StateSchema + 'static>(
324        &self,
325        name: &str,
326        subgraph: CompiledGraph<SubS>,
327        input_mapper: impl Fn(&S) -> SubS + Send + Sync + 'static,
328        output_mapper: impl Fn(&SubS, &mut S) + Send + Sync + 'static,
329    ) -> GraphResult<()> {
330        let node = SubgraphNode::new(name, subgraph, input_mapper, output_mapper);
331        self.runtime_nodes
332            .write()
333            .await
334            .insert(name.to_string(), Arc::new(node));
335        Ok(())
336    }
337}