1use 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#[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 {
86 self.checkpointer = Some(Arc::new(Mutex::new(checkpointer)));
87 self
88 }
89
90 pub fn with_recursion_limit(mut self, limit: usize) -> Self {
92 self.recursion_limit = limit;
93 self
94 }
95
96 pub fn with_interrupt_before(mut self, nodes: Vec<String>) -> Self {
98 self.interrupt_before = nodes;
99 self
100 }
101
102 pub fn with_interrupt_after(mut self, nodes: Vec<String>) -> Self {
104 self.interrupt_after = nodes;
105 self
106 }
107
108 pub fn node_names(&self) -> Vec<String> {
110 self.nodes.keys().cloned().collect()
111 }
112
113 pub fn get_edges(&self) -> &[GraphEdge] {
115 &self.edges
116 }
117
118 pub fn entry_point(&self) -> &str {
120 &self.entry_point
121 }
122
123 pub fn recursion_limit(&self) -> usize {
125 self.recursion_limit
126 }
127
128 pub fn interrupt_before(&self) -> &[String] {
130 &self.interrupt_before
131 }
132
133 pub fn interrupt_after(&self) -> &[String] {
135 &self.interrupt_after
136 }
137
138 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 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 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 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 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 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 recursion_count,
238 interrupted_at: interrupted_node.to_string(),
239 pending_interrupt: None,
240 })
241 }
242
243 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 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 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 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 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}