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};
5use crate::checkpointer::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    pub fn with_checkpointer<C: Checkpointer<S> + 'static>(mut self, checkpointer: C) -> Self {
85        self.checkpointer = Some(Arc::new(Mutex::new(checkpointer)));
86        self
87    }
88
89    pub fn with_recursion_limit(mut self, limit: usize) -> Self {
90        self.recursion_limit = limit;
91        self
92    }
93
94    pub fn with_interrupt_before(mut self, nodes: Vec<String>) -> Self {
95        self.interrupt_before = nodes;
96        self
97    }
98
99    pub fn with_interrupt_after(mut self, nodes: Vec<String>) -> Self {
100        self.interrupt_after = nodes;
101        self
102    }
103
104    pub fn node_names(&self) -> Vec<String> {
105        self.nodes.keys().cloned().collect()
106    }
107
108    pub fn get_edges(&self) -> &[GraphEdge] {
109        &self.edges
110    }
111
112    pub fn entry_point(&self) -> &str {
113        &self.entry_point
114    }
115
116    pub fn recursion_limit(&self) -> usize {
117        self.recursion_limit
118    }
119
120    pub fn interrupt_before(&self) -> &[String] {
121        &self.interrupt_before
122    }
123
124    pub fn interrupt_after(&self) -> &[String] {
125        &self.interrupt_after
126    }
127
128    /// Get the last checkpoint state (for interrupt recovery), together with the
129    /// recursion budget consumed up to that checkpoint.
130    ///
131    /// H5: 不再用 `list()` 的 HashMap 乱序 + `ids.last()` 猜"最近";直接由
132    /// checkpointer 的 `last()` 按 (timestamp, seq) 返回真正最近的 checkpoint。
133    pub async fn last_checkpoint_state(&self) -> Option<(S, usize)> {
134        if let Some(ref cp) = self.checkpointer {
135            let guard = cp.lock().await;
136            guard.last().await.ok()?
137        } else {
138            None
139        }
140    }
141
142    /// Create a resume execution context (from the last checkpoint)
143    /// interrupted_node may be "node_name" or "after_node_name"
144    pub async fn create_resume_execution(
145        &self,
146        interrupted_node: &str,
147    ) -> Option<GraphExecution<S>> {
148        let (state, recursion_count) = self.last_checkpoint_state().await?;
149        let current = interrupted_node
150            .strip_prefix("after_")
151            .unwrap_or(interrupted_node);
152        Some(GraphExecution {
153            state,
154            current_node: current.to_string(),
155            steps: Vec::new(),
156            // M6: 续跑沿用中断前已消耗的递归预算,而不是清零重来——
157            // 否则反复 interrupt→resume 可以无限绕过 recursion_limit。
158            recursion_count,
159            interrupted_at: interrupted_node.to_string(),
160        })
161    }
162
163    pub(super) async fn get_node(&self, name: &str) -> GraphResult<Arc<dyn GraphNode<S>>> {
164        if let Some(node) = self.nodes.get(name) {
165            return Ok(node.clone());
166        }
167        if let Some(node) = self.runtime_nodes.read().await.get(name) {
168            return Ok(node.clone());
169        }
170        Err(GraphError::ExecutionError(format!(
171            "Node '{}' not found",
172            name
173        )))
174    }
175
176    /// Submit a new task for dynamic planning mid-execution
177    pub fn submit_task(&self, description: String) {
178        if let Ok(mut inbox) = self.task_inbox.try_lock() {
179            inbox.push(DynamicTask {
180                id: uuid::Uuid::new_v4().to_string(),
181                description,
182            });
183        }
184    }
185
186    /// Inject a runtime node directly
187    pub async fn inject_node(&self, name: &str, node: Arc<dyn GraphNode<S>>) -> GraphResult<()> {
188        self.runtime_nodes
189            .write()
190            .await
191            .insert(name.to_string(), node);
192        Ok(())
193    }
194
195    /// Inject a runtime edge directly
196    pub async fn inject_edge(&self, source: &str, target: &str) -> GraphResult<()> {
197        self.runtime_edges
198            .write()
199            .await
200            .push(GraphEdge::fixed(source, target));
201        Ok(())
202    }
203
204    /// Inject a subgraph as a runtime node
205    pub async fn inject_subgraph<SubS: StateSchema + 'static>(
206        &self,
207        name: &str,
208        subgraph: CompiledGraph<SubS>,
209        input_mapper: impl Fn(&S) -> SubS + Send + Sync + 'static,
210        output_mapper: impl Fn(&SubS, &mut S) + Send + Sync + 'static,
211    ) -> GraphResult<()> {
212        let node = SubgraphNode::new(name, subgraph, input_mapper, output_mapper);
213        self.runtime_nodes
214            .write()
215            .await
216            .insert(name.to_string(), Arc::new(node));
217        Ok(())
218    }
219}