1use std::collections::{BTreeSet, HashMap};
29use std::sync::Arc;
30
31use futures::future::join_all;
32use oneai_core::error::{OneAIError, Result};
33use oneai_core::traits::{InteractionGate, LlmProvider, PermissionResolver, Tool};
34use oneai_core::{InferenceRequest, InferenceResponse, Message, PermissionAction, Role};
35
36use crate::state_graph::{
37 EdgeCondition, GraphEdge, GraphExecutionResult, GraphState, NodeAction, StateGraph,
38};
39
40pub type CustomConditionFn = Arc<dyn Fn(&GraphState) -> bool + Send + Sync>;
46
47pub(crate) fn evaluate_custom_condition(
54 registry: &HashMap<String, CustomConditionFn>,
55 name: &str,
56 description: &str,
57 state: &GraphState,
58) -> bool {
59 match registry.get(name) {
60 Some(evaluator) => evaluator(state),
61 None => {
62 tracing::warn!(
63 "Custom condition '{}' ('{}') not registered. Defaulting to false.",
64 name,
65 description
66 );
67 false
68 }
69 }
70}
71
72#[async_trait::async_trait]
83pub trait DelegateFactory: Send + Sync {
84 async fn delegate(&self, agent_kind: &str, task: &str) -> Result<String>;
86}
87
88pub struct NoopDelegateFactory;
91
92#[async_trait::async_trait]
93impl DelegateFactory for NoopDelegateFactory {
94 async fn delegate(&self, agent_kind: &str, _task: &str) -> Result<String> {
95 Err(OneAIError::Workflow(format!(
96 "Delegate action '{}' not supported — no DelegateFactory configured",
97 agent_kind
98 )))
99 }
100}
101
102#[derive(Debug, Clone)]
106pub struct ActionResult {
107 pub output: String,
109 pub error: Option<String>,
111}
112
113struct BranchResult {
120 node_id: String,
122 action_result: ActionResult,
124 branch_state: GraphState,
126 next_targets: Vec<String>,
129}
130
131#[async_trait::async_trait]
152pub trait GraphActionExecutor: Send + Sync {
153 async fn execute_llm_infer(
162 &self,
163 action: &NodeAction,
164 state: &mut GraphState,
165 ) -> Result<ActionResult>;
166
167 async fn execute_tool_call(
174 &self,
175 tool_name: &str,
176 args: &serde_json::Value,
177 state: &mut GraphState,
178 ) -> Result<ActionResult>;
179
180 async fn execute_paradigm_switch(
186 &self,
187 paradigm: &str,
188 state: &mut GraphState,
189 ) -> Result<ActionResult>;
190
191 async fn parse_decision(
196 &self,
197 response: &InferenceResponse,
198 state: &mut GraphState,
199 ) -> Result<oneai_core::GraphDecision>;
200}
201
202pub struct DirectProviderActionExecutor {
216 provider: Arc<dyn LlmProvider>,
217 tools: Arc<tokio::sync::RwLock<HashMap<String, Arc<dyn Tool>>>>,
218 permission_resolver: Option<Arc<dyn PermissionResolver>>,
223}
224
225impl DirectProviderActionExecutor {
226 pub fn new(
228 provider: Arc<dyn LlmProvider>,
229 tools: Arc<tokio::sync::RwLock<HashMap<String, Arc<dyn Tool>>>>,
230 ) -> Self {
231 Self {
232 provider,
233 tools,
234 permission_resolver: None,
235 }
236 }
237
238 pub fn with_permission_resolver(mut self, resolver: Arc<dyn PermissionResolver>) -> Self {
241 self.permission_resolver = Some(resolver);
242 self
243 }
244}
245
246#[async_trait::async_trait]
247impl GraphActionExecutor for DirectProviderActionExecutor {
248 async fn execute_llm_infer(
249 &self,
250 action: &NodeAction,
251 state: &mut GraphState,
252 ) -> Result<ActionResult> {
253 let (
255 system_prompt_override,
256 _use_streaming,
257 include_tool_definitions,
258 tool_filter_override,
259 thinking_budget,
260 temperature,
261 max_tokens,
262 ) = match action {
263 NodeAction::LlmInfer {
264 system_prompt_override,
265 use_streaming,
266 include_tool_definitions,
267 tool_filter_override,
268 thinking_budget,
269 temperature,
270 max_tokens,
271 } => (
272 system_prompt_override.clone(),
273 *use_streaming,
274 *include_tool_definitions,
275 tool_filter_override.clone(),
276 *thinking_budget,
277 *temperature,
278 *max_tokens,
279 ),
280 _ => return Err(OneAIError::Workflow("Expected LlmInfer action".to_string())),
281 };
282
283 let system_prompt = system_prompt_override
285 .unwrap_or_else(|| "You are an intelligent agent. Respond to the task.".to_string());
286
287 let mut conversation = state.conversation.clone();
288 if !conversation.messages.iter().any(|m| m.role == Role::System) {
289 conversation.add_message(Message::system(&system_prompt));
290 }
291
292 let tool_defs = if include_tool_definitions {
294 let tools_map = self.tools.read().await;
295 if let Some(filter) = &tool_filter_override {
296 tools_map
298 .values()
299 .filter(|t| filter.contains(&t.name().to_string()))
300 .map(|t| oneai_core::ToolDefinition {
301 name: t.name().to_string(),
302 description: t.description().to_string(),
303 parameters_schema: t.parameters_schema(),
304 })
305 .collect()
306 } else {
307 tools_map
309 .values()
310 .map(|t| oneai_core::ToolDefinition {
311 name: t.name().to_string(),
312 description: t.description().to_string(),
313 parameters_schema: t.parameters_schema(),
314 })
315 .collect()
316 }
317 } else {
318 vec![] };
320
321 let request = InferenceRequest {
322 conversation,
323 tools: tool_defs,
324 max_tokens: max_tokens.or(Some(4096)),
325 temperature: temperature.or(Some(0.3)),
326 top_p: None,
327 stop_sequences: vec![],
328 constrained_output: None,
329 thinking_budget,
330 metadata: HashMap::new(),
331 };
332
333 let response = self.provider.infer(request).await?;
334 let output = response.message.text_content();
335
336 state.conversation.add_message(response.message.clone());
338
339 let _decision = self.parse_decision(&response, state).await?;
341
342 Ok(ActionResult {
343 output,
344 error: None,
345 })
346 }
347
348 async fn execute_tool_call(
349 &self,
350 tool_name: &str,
351 args: &serde_json::Value,
352 state: &mut GraphState,
353 ) -> Result<ActionResult> {
354 if let Some(resolver) = self.permission_resolver.as_ref() {
359 match resolver.resolve(tool_name, args) {
360 PermissionAction::Deny { reason } => {
361 tracing::warn!(
362 "ToolCall node '{}' denied by domain policy: {}",
363 tool_name,
364 reason
365 );
366 return Ok(ActionResult {
367 output: format!("Denied by domain policy: {}", reason),
368 error: Some(format!("Denied by domain policy: {}", reason)),
369 });
370 }
371 PermissionAction::AutoApprove => {
372 tracing::info!(
373 "ToolCall node '{}' auto-approved by domain policy",
374 tool_name
375 );
376 }
377 _ => {}
382 }
383 }
384
385 let tools_map = self.tools.read().await;
386 let tool = tools_map.get(tool_name).ok_or_else(|| {
387 OneAIError::Workflow(format!("Tool '{}' not found for ToolCall node", tool_name))
388 })?;
389
390 let output = tool.execute(args.clone()).await?;
391
392 state.conversation.add_message(Message::tool_result(
393 format!("graph_tool_{}", tool_name),
394 output.content.clone(),
395 ));
396
397 Ok(ActionResult {
398 output: output.content,
399 error: output.error,
400 })
401 }
402
403 async fn execute_paradigm_switch(
404 &self,
405 paradigm: &str,
406 state: &mut GraphState,
407 ) -> Result<ActionResult> {
408 state.active_paradigm = Some(paradigm.to_string());
409 state.parsed_decision = None; Ok(ActionResult {
412 output: format!("Paradigm switched to: {}", paradigm),
413 error: None,
414 })
415 }
416
417 async fn parse_decision(
418 &self,
419 response: &InferenceResponse,
420 state: &mut GraphState,
421 ) -> Result<oneai_core::GraphDecision> {
422 let mut tool_calls = Vec::new();
425 let mut text_parts = Vec::new();
426
427 for block in &response.message.content {
428 match block {
429 oneai_core::ContentBlock::ToolCall { id: _, name, args } => {
430 if name == "delegate" {
432 let args_value: serde_json::Value =
434 serde_json::from_str(args).unwrap_or_else(|_| serde_json::json!({}));
435 let agent_kind = args_value
436 .get("agent_type")
437 .and_then(|v| v.as_str())
438 .unwrap_or("Explore")
439 .to_string();
440 let task = args_value
441 .get("task")
442 .and_then(|v| v.as_str())
443 .unwrap_or("")
444 .to_string();
445 let decision = oneai_core::GraphDecision::Delegate { agent_kind, task };
446 state.parsed_decision = Some(decision.clone());
447 return Ok(decision);
448 }
449 if name == "switch_paradigm" {
450 let args_value: serde_json::Value =
451 serde_json::from_str(args).unwrap_or_else(|_| serde_json::json!({}));
452 let paradigm = args_value
453 .get("paradigm")
454 .and_then(|v| v.as_str())
455 .unwrap_or("react")
456 .to_string();
457 let decision = oneai_core::GraphDecision::SwitchParadigm { paradigm };
458 state.parsed_decision = Some(decision.clone());
459 return Ok(decision);
460 }
461 tool_calls.push(name.clone());
462 }
463 oneai_core::ContentBlock::Text { text } => {
464 text_parts.push(text.clone());
465 }
466 _ => {}
467 }
468 }
469
470 let decision = if !tool_calls.is_empty() {
471 oneai_core::GraphDecision::ToolCalls {
472 count: tool_calls.len(),
473 }
474 } else {
475 oneai_core::GraphDecision::DirectAnswer {
476 text: text_parts.join("\n"),
477 }
478 };
479
480 state.parsed_decision = Some(decision.clone());
481 Ok(decision)
482 }
483}
484
485pub struct StateGraphExecutor {
504 action_executor: Arc<dyn GraphActionExecutor>,
508 delegate_factory: Arc<dyn DelegateFactory>,
510 interaction_gate: Arc<dyn InteractionGate>,
512 max_iterations: usize,
515 custom_conditions: HashMap<String, CustomConditionFn>,
518}
519
520impl StateGraphExecutor {
521 pub fn new(
526 action_executor: Arc<dyn GraphActionExecutor>,
527 delegate_factory: Arc<dyn DelegateFactory>,
528 interaction_gate: Arc<dyn InteractionGate>,
529 max_iterations: usize,
530 ) -> Self {
531 Self {
532 action_executor,
533 delegate_factory,
534 interaction_gate,
535 max_iterations,
536 custom_conditions: HashMap::new(),
537 }
538 }
539
540 pub fn with_defaults(
542 action_executor: Arc<dyn GraphActionExecutor>,
543 delegate_factory: Arc<dyn DelegateFactory>,
544 interaction_gate: Arc<dyn InteractionGate>,
545 ) -> Self {
546 Self::new(action_executor, delegate_factory, interaction_gate, 50)
547 }
548
549 pub fn with_custom_condition(
557 mut self,
558 name: impl Into<String>,
559 evaluator: CustomConditionFn,
560 ) -> Self {
561 self.custom_conditions.insert(name.into(), evaluator);
562 self
563 }
564
565 pub fn with_direct_provider(
572 provider: Arc<dyn LlmProvider>,
573 tools: Arc<tokio::sync::RwLock<HashMap<String, Arc<dyn Tool>>>>,
574 delegate_factory: Arc<dyn DelegateFactory>,
575 interaction_gate: Arc<dyn InteractionGate>,
576 max_iterations: usize,
577 ) -> Self {
578 let action_executor = Arc::new(DirectProviderActionExecutor::new(provider, tools));
579 Self::new(
580 action_executor,
581 delegate_factory,
582 interaction_gate,
583 max_iterations,
584 )
585 }
586
587 pub fn with_direct_provider_defaults(
590 provider: Arc<dyn LlmProvider>,
591 tools: Arc<tokio::sync::RwLock<HashMap<String, Arc<dyn Tool>>>>,
592 delegate_factory: Arc<dyn DelegateFactory>,
593 interaction_gate: Arc<dyn InteractionGate>,
594 ) -> Self {
595 Self::with_direct_provider(provider, tools, delegate_factory, interaction_gate, 50)
596 }
597
598 pub async fn execute(
624 &self,
625 graph: &StateGraph,
626 initial_state: GraphState,
627 ) -> Result<GraphExecutionResult> {
628 let mut state = initial_state;
629 let mut frontier: BTreeSet<String> = std::iter::once(graph.entry_point.clone()).collect();
633 let mut iterations = 0;
634 let mut interrupt_checkpoints = Vec::new();
635
636 while iterations < self.max_iterations {
637 iterations += 1;
638 state.iteration_count = iterations;
639
640 if frontier.is_empty() {
641 break;
642 }
643
644 if let Some(term_id) = frontier
649 .iter()
650 .find(|id| graph.terminal_nodes.contains(id))
651 .cloned()
652 {
653 let node = graph.get_node(&term_id).ok_or_else(|| {
654 OneAIError::Workflow(format!(
655 "Terminal node '{}' not found in graph '{}'",
656 term_id, graph.name
657 ))
658 })?;
659 let action_result = self.execute_node_action(&node.action, &mut state).await?;
660 state.last_result = Some(action_result.output.clone());
661 state.last_error = action_result.error.clone();
662
663 return Ok(GraphExecutionResult {
664 name: graph.name.clone(),
665 final_state: state,
666 completed: true,
667 terminal_node: Some(term_id),
668 iterations,
669 interrupt_checkpoints,
670 });
671 }
672
673 let has_interrupt = frontier
676 .iter()
677 .any(|id| graph.get_node(id).map(|n| n.interrupt).unwrap_or(false));
678
679 let next_targets: Vec<String> = if frontier.len() == 1 || has_interrupt {
680 self.execute_frontier_sequential(
681 graph,
682 &frontier,
683 &mut state,
684 &mut interrupt_checkpoints,
685 )
686 .await?
687 } else {
688 self.execute_frontier_parallel(graph, &frontier, &mut state)
689 .await?
690 };
691
692 if state.should_terminate {
693 break;
694 }
695
696 if next_targets.is_empty() {
697 tracing::warn!(
698 "No matching edge condition from frontier {:?} in graph '{}'. Terminating.",
699 frontier,
700 graph.name
701 );
702 break;
703 }
704
705 frontier = next_targets.into_iter().collect();
706 }
707
708 let should_terminate = state.should_terminate;
710 if iterations >= self.max_iterations {
711 tracing::warn!(
712 "StateGraph '{}' exceeded max iterations ({}). Terminating.",
713 graph.name,
714 self.max_iterations
715 );
716 }
717
718 Ok(GraphExecutionResult {
719 name: graph.name.clone(),
720 final_state: state,
721 completed: !should_terminate,
722 terminal_node: None,
723 iterations,
724 interrupt_checkpoints,
725 })
726 }
727
728 async fn execute_frontier_sequential(
738 &self,
739 graph: &StateGraph,
740 frontier: &BTreeSet<String>,
741 state: &mut GraphState,
742 interrupt_checkpoints: &mut Vec<String>,
743 ) -> Result<Vec<String>> {
744 let mut next = Vec::new();
745 for node_id in frontier {
746 let node = graph.get_node(node_id).ok_or_else(|| {
747 OneAIError::Workflow(format!(
748 "Node '{}' not found in graph '{}'",
749 node_id, graph.name
750 ))
751 })?;
752
753 if node.interrupt {
756 let checkpoint_id = format!("interrupt_{}_{}", graph.name, node_id);
757 interrupt_checkpoints.push(checkpoint_id.clone());
758 let approval_request = oneai_core::ApprovalRequest {
759 tool_name: "state_graph_interrupt".into(),
760 args: serde_json::json!({
761 "node": node_id,
762 "description": match &node.action {
763 NodeAction::HumanApproval { description } => description.clone(),
764 _ => "Interrupt point reached".to_string(),
765 },
766 "state": state.variables,
767 }),
768 risk_level: oneai_core::RiskLevel::Medium,
769 permission_level: Some(oneai_core::PermissionLevel::Standard),
770 justification: format!(
771 "StateGraph interrupt at node '{}' in graph '{}'",
772 node_id, graph.name
773 ),
774 };
775 let approval = self
776 .interaction_gate
777 .request(oneai_core::InteractionRequest::ToolApproval {
778 approval: approval_request,
779 })
780 .await?;
781 match approval {
782 oneai_core::InteractionResponse::Abort { reason } => {
783 state.should_terminate = true;
784 state.last_error = Some(format!("Interrupt denied: {}", reason));
785 return Ok(next);
786 }
787 oneai_core::InteractionResponse::Revise { feedback } => {
788 state.should_terminate = true;
789 state.last_error = Some(format!("Interrupt denied: {}", feedback));
790 return Ok(next);
791 }
792 _ => { }
793 }
794 }
795
796 let action_result = self.execute_node_action(&node.action, state).await?;
797 state.last_result = Some(action_result.output.clone());
798 state.last_error = action_result.error.clone();
799 if state.should_terminate {
800 return Ok(next);
801 }
802
803 let edges = graph.get_edges_from(node_id);
804 next.extend(self.route_next_nodes(&edges, state)?);
805 }
806 Ok(next)
807 }
808
809 async fn execute_frontier_parallel(
819 &self,
820 graph: &StateGraph,
821 frontier: &BTreeSet<String>,
822 state: &mut GraphState,
823 ) -> Result<Vec<String>> {
824 let baseline_msgs = state.conversation.messages.len();
827 let incoming = state.clone();
831
832 let branch_futs: Vec<_> = frontier
840 .iter()
841 .map(|node_id| {
842 let seed = incoming.clone();
847 async move {
848 let node = match graph.get_node(node_id) {
849 Some(n) => n,
850 None => {
851 return BranchResult {
852 node_id: node_id.clone(),
853 action_result: ActionResult {
854 output: String::new(),
855 error: Some(format!(
856 "Node '{}' not found in graph '{}'",
857 node_id, graph.name
858 )),
859 },
860 branch_state: seed,
861 next_targets: Vec::new(),
862 };
863 }
864 };
865 let mut branch_state = seed;
866 let action_result = match self
867 .execute_node_action(&node.action, &mut branch_state)
868 .await
869 {
870 Ok(ar) => ar,
871 Err(e) => ActionResult {
872 output: String::new(),
873 error: Some(e.to_string()),
874 },
875 };
876 let edges = graph.get_edges_from(node_id);
877 let next_targets = self
880 .route_next_nodes(&edges, &branch_state)
881 .unwrap_or_default();
882 BranchResult {
883 node_id: node_id.clone(),
884 action_result,
885 branch_state,
886 next_targets,
887 }
888 }
889 })
890 .collect();
891
892 let mut branches: Vec<BranchResult> = join_all(branch_futs).await;
893 branches.sort_by(|a, b| a.node_id.cmp(&b.node_id));
898
899 let mut joined_output = Vec::new();
902 let mut first_error: Option<String> = None;
903 let mut any_terminate = false;
904 for b in &branches {
905 joined_output.push(b.action_result.output.clone());
906 if first_error.is_none() {
907 first_error = b.action_result.error.clone();
908 }
909 any_terminate |= b.branch_state.should_terminate;
910 }
911 state.last_result = Some(joined_output.join("\n---\n"));
912 state.last_error = first_error;
913 state.should_terminate = any_terminate;
914
915 for b in &branches {
919 let extra_from = baseline_msgs.min(b.branch_state.conversation.messages.len());
920 for m in &b.branch_state.conversation.messages[extra_from..] {
921 state.conversation.add_message(m.clone());
922 }
923 }
924 for b in &branches {
928 for (k, v) in &b.branch_state.variables {
929 state.variables.insert(k.clone(), v.clone());
930 }
931 }
932 if let Some(last) = branches.last() {
937 state.parsed_decision = last.branch_state.parsed_decision.clone();
938 }
939
940 let mut next = Vec::new();
941 for b in &branches {
942 next.extend(b.next_targets.iter().cloned());
943 }
944 Ok(next)
945 }
946
947 async fn execute_node_action(
953 &self,
954 action: &NodeAction,
955 state: &mut GraphState,
956 ) -> Result<ActionResult> {
957 match action {
958 NodeAction::LlmInfer { .. } => {
959 self.action_executor.execute_llm_infer(action, state).await
962 }
963
964 NodeAction::ToolCall {
965 tool_name,
966 args_template,
967 } => {
968 let resolved_name = interpolate_graph_template(tool_name, &state.variables);
970 let resolved_args = if let Some(template) = args_template {
971 let json_str = interpolate_graph_template(template, &state.variables);
972 serde_json::from_str(&json_str).unwrap_or(serde_json::json!({}))
973 } else {
974 serde_json::json!({})
975 };
976
977 self.action_executor
980 .execute_tool_call(&resolved_name, &resolved_args, state)
981 .await
982 }
983
984 NodeAction::Delegate {
985 agent_kind,
986 task_template,
987 } => {
988 let task = interpolate_graph_template(task_template, &state.variables);
989
990 let result = self.delegate_factory.delegate(agent_kind, &task).await?;
991
992 state.conversation.add_message(Message::assistant(format!(
994 "[Delegate {}]: {}",
995 agent_kind, result
996 )));
997
998 state.parsed_decision = Some(oneai_core::GraphDecision::Delegate {
1000 agent_kind: agent_kind.clone(),
1001 task,
1002 });
1003
1004 Ok(ActionResult {
1005 output: result,
1006 error: None,
1007 })
1008 }
1009
1010 NodeAction::HumanApproval { description } => {
1011 Ok(ActionResult {
1014 output: description.clone(),
1015 error: None,
1016 })
1017 }
1018
1019 NodeAction::ConditionCheck { condition } => {
1020 let result = self.evaluate_condition_expression(condition, state)?;
1021 state
1022 .variables
1023 .insert("_condition_result".to_string(), result.to_string());
1024 Ok(ActionResult {
1025 output: result.to_string(),
1026 error: None,
1027 })
1028 }
1029
1030 NodeAction::SwitchParadigm { paradigm } => {
1031 self.action_executor
1033 .execute_paradigm_switch(paradigm, state)
1034 .await
1035 }
1036 }
1037 }
1038
1039 fn route_next_nodes(&self, edges: &[&GraphEdge], state: &GraphState) -> Result<Vec<String>> {
1051 let mut next = Vec::with_capacity(edges.len());
1052 for edge in edges {
1053 let satisfied = match &edge.condition {
1054 Some(condition) => self.evaluate_edge_condition(condition, state)?,
1055 None => true, };
1057 if satisfied {
1058 next.push(edge.to.clone());
1059 }
1060 }
1061 Ok(next)
1062 }
1063
1064 fn evaluate_edge_condition(
1071 &self,
1072 condition: &EdgeCondition,
1073 state: &GraphState,
1074 ) -> Result<bool> {
1075 match condition {
1076 EdgeCondition::HasToolCalls => {
1077 Ok(state
1079 .parsed_decision
1080 .as_ref()
1081 .map(|d| d.has_tool_calls())
1082 .unwrap_or(false))
1083 }
1084
1085 EdgeCondition::IsFinalAnswer => {
1086 Ok(state
1088 .parsed_decision
1089 .as_ref()
1090 .map(|d| d.is_final())
1091 .unwrap_or(false))
1092 }
1093
1094 EdgeCondition::RequestsDelegation => {
1095 Ok(state
1097 .parsed_decision
1098 .as_ref()
1099 .map(|d| d.is_delegation())
1100 .unwrap_or(false))
1101 }
1102
1103 EdgeCondition::ErrorOccurred => Ok(state.last_error.is_some()),
1104
1105 EdgeCondition::StateEquals { variable, value } => {
1106 Ok(state.variables.get(variable) == Some(value))
1107 }
1108
1109 EdgeCondition::Always => Ok(true),
1110
1111 EdgeCondition::Custom { name, description } => Ok(evaluate_custom_condition(
1112 &self.custom_conditions,
1113 name,
1114 description,
1115 state,
1116 )),
1117
1118 EdgeCondition::ParadigmEquals { paradigm } => {
1119 Ok(state.active_paradigm.as_ref() == Some(paradigm))
1120 }
1121
1122 EdgeCondition::IterationExceeds { count } => Ok(state.iteration_count > *count),
1123 }
1124 }
1125
1126 fn evaluate_condition_expression(&self, condition: &str, state: &GraphState) -> Result<bool> {
1134 if condition == "has_tool_calls" {
1135 return self.evaluate_edge_condition(&EdgeCondition::HasToolCalls, state);
1136 }
1137 if condition == "is_final_answer" {
1138 return self.evaluate_edge_condition(&EdgeCondition::IsFinalAnswer, state);
1139 }
1140 if condition == "error_occurred" {
1141 return self.evaluate_edge_condition(&EdgeCondition::ErrorOccurred, state);
1142 }
1143 if let Some((var, val)) = condition.split_once("==") {
1145 return Ok(state.variables.get(var.trim()) == Some(&val.trim().to_string()));
1146 }
1147 Ok(state
1149 .variables
1150 .get(condition)
1151 .map(|v| v == "true")
1152 .unwrap_or(false))
1153 }
1154}
1155
1156pub fn interpolate_graph_template(template: &str, variables: &HashMap<String, String>) -> String {
1163 let mut result = template.to_string();
1164 for (key, value) in variables {
1165 result = result.replace(&format!("{{{{{}}}}}", key), value);
1166 }
1167 result
1168}
1169
1170#[cfg(test)]
1173mod tests {
1174 use super::*;
1175 use crate::state_graph::{GraphNode, StateGraph};
1176
1177 #[allow(dead_code)] fn make_simple_graph() -> StateGraph {
1179 let mut graph = StateGraph::new("test", "start");
1181
1182 graph.add_node(GraphNode {
1183 id: "start".to_string(),
1184 action: NodeAction::ConditionCheck {
1185 condition: "ready==true".to_string(),
1186 },
1187 interrupt: false,
1188 metadata: HashMap::new(),
1189 });
1190
1191 graph.add_node(GraphNode {
1192 id: "end".to_string(),
1193 action: NodeAction::LlmInfer {
1194 system_prompt_override: Some("Final answer".to_string()),
1195 use_streaming: false,
1196 include_tool_definitions: false,
1197 tool_filter_override: None,
1198 thinking_budget: None,
1199 temperature: None,
1200 max_tokens: None,
1201 },
1202 interrupt: false,
1203 metadata: HashMap::new(),
1204 });
1205
1206 graph.add_edge(GraphEdge {
1207 from: "start".to_string(),
1208 to: "end".to_string(),
1209 condition: Some(EdgeCondition::Always),
1210 metadata: HashMap::new(),
1211 });
1212
1213 graph.add_terminal("end".to_string());
1214
1215 graph
1216 }
1217
1218 #[test]
1219 fn test_evaluate_edge_condition_always() {
1220 let state = GraphState::new();
1221 let executor = make_test_executor();
1222
1223 let result = executor
1224 .evaluate_edge_condition(&EdgeCondition::Always, &state)
1225 .unwrap();
1226 assert!(result);
1227 }
1228
1229 #[test]
1230 fn test_evaluate_edge_condition_error_occurred() {
1231 let mut state = GraphState::new();
1232 let executor = make_test_executor();
1233
1234 assert!(!executor
1236 .evaluate_edge_condition(&EdgeCondition::ErrorOccurred, &state)
1237 .unwrap());
1238
1239 state.last_error = Some("test error".to_string());
1241 assert!(executor
1242 .evaluate_edge_condition(&EdgeCondition::ErrorOccurred, &state)
1243 .unwrap());
1244 }
1245
1246 #[test]
1247 fn test_evaluate_edge_condition_state_equals() {
1248 let mut state = GraphState::new();
1249 let executor = make_test_executor();
1250
1251 let cond = EdgeCondition::StateEquals {
1253 variable: "mode".to_string(),
1254 value: "react".to_string(),
1255 };
1256 assert!(!executor.evaluate_edge_condition(&cond, &state).unwrap());
1257
1258 state
1260 .variables
1261 .insert("mode".to_string(), "plan".to_string());
1262 assert!(!executor.evaluate_edge_condition(&cond, &state).unwrap());
1263
1264 state
1266 .variables
1267 .insert("mode".to_string(), "react".to_string());
1268 assert!(executor.evaluate_edge_condition(&cond, &state).unwrap());
1269 }
1270
1271 #[test]
1272 fn test_evaluate_edge_condition_has_tool_calls() {
1273 let mut state = GraphState::new();
1274 let executor = make_test_executor();
1275
1276 assert!(!executor
1278 .evaluate_edge_condition(&EdgeCondition::HasToolCalls, &state)
1279 .unwrap());
1280
1281 state.parsed_decision = Some(oneai_core::GraphDecision::DirectAnswer {
1283 text: "Just a text answer".to_string(),
1284 });
1285 assert!(!executor
1286 .evaluate_edge_condition(&EdgeCondition::HasToolCalls, &state)
1287 .unwrap());
1288
1289 state.parsed_decision = Some(oneai_core::GraphDecision::ToolCalls { count: 1 });
1291 assert!(executor
1292 .evaluate_edge_condition(&EdgeCondition::HasToolCalls, &state)
1293 .unwrap());
1294 }
1295
1296 #[test]
1297 fn test_evaluate_edge_condition_is_final_answer() {
1298 let mut state = GraphState::new();
1299 let executor = make_test_executor();
1300
1301 assert!(!executor
1303 .evaluate_edge_condition(&EdgeCondition::IsFinalAnswer, &state)
1304 .unwrap());
1305
1306 state.parsed_decision = Some(oneai_core::GraphDecision::DirectAnswer {
1308 text: "The answer is 42".to_string(),
1309 });
1310 assert!(executor
1311 .evaluate_edge_condition(&EdgeCondition::IsFinalAnswer, &state)
1312 .unwrap());
1313
1314 state.parsed_decision = Some(oneai_core::GraphDecision::ToolCalls { count: 2 });
1316 assert!(!executor
1317 .evaluate_edge_condition(&EdgeCondition::IsFinalAnswer, &state)
1318 .unwrap());
1319 }
1320
1321 #[test]
1322 fn test_evaluate_edge_condition_requests_delegation() {
1323 let mut state = GraphState::new();
1324 let executor = make_test_executor();
1325
1326 state.parsed_decision = Some(oneai_core::GraphDecision::Delegate {
1328 agent_kind: "Explore".to_string(),
1329 task: "Search the codebase".to_string(),
1330 });
1331 assert!(executor
1332 .evaluate_edge_condition(&EdgeCondition::RequestsDelegation, &state)
1333 .unwrap());
1334
1335 state.parsed_decision = Some(oneai_core::GraphDecision::ToolCalls { count: 1 });
1337 assert!(!executor
1338 .evaluate_edge_condition(&EdgeCondition::RequestsDelegation, &state)
1339 .unwrap());
1340 }
1341
1342 #[test]
1343 fn test_evaluate_edge_condition_paradigm_equals() {
1344 let mut state = GraphState::new();
1345 let executor = make_test_executor();
1346
1347 let cond = EdgeCondition::ParadigmEquals {
1349 paradigm: "react".to_string(),
1350 };
1351 assert!(!executor.evaluate_edge_condition(&cond, &state).unwrap());
1352
1353 state.active_paradigm = Some("plan".to_string());
1355 assert!(!executor.evaluate_edge_condition(&cond, &state).unwrap());
1356
1357 state.active_paradigm = Some("react".to_string());
1359 assert!(executor.evaluate_edge_condition(&cond, &state).unwrap());
1360 }
1361
1362 #[test]
1363 fn test_evaluate_edge_condition_iteration_exceeds() {
1364 let mut state = GraphState::new();
1365 let executor = make_test_executor();
1366
1367 let cond = EdgeCondition::IterationExceeds { count: 5 };
1369 assert!(!executor.evaluate_edge_condition(&cond, &state).unwrap());
1370
1371 state.iteration_count = 10;
1373 assert!(executor.evaluate_edge_condition(&cond, &state).unwrap());
1374 }
1375
1376 #[test]
1377 fn test_interpolate_graph_template() {
1378 let vars = HashMap::from([
1379 ("selected_tool".to_string(), "shell".to_string()),
1380 ("command".to_string(), "ls -la".to_string()),
1381 ]);
1382
1383 let result = interpolate_graph_template("{{selected_tool}} with {{command}}", &vars);
1384 assert_eq!(result, "shell with ls -la");
1385 }
1386
1387 #[test]
1388 fn test_evaluate_condition_expression() {
1389 let mut state = GraphState::new();
1390 let executor = make_test_executor();
1391
1392 state
1394 .variables
1395 .insert("mode".to_string(), "react".to_string());
1396 assert!(executor
1397 .evaluate_condition_expression("mode==react", &state)
1398 .unwrap());
1399 assert!(!executor
1400 .evaluate_condition_expression("mode==plan", &state)
1401 .unwrap());
1402
1403 state.last_error = Some("error".to_string());
1405 assert!(executor
1406 .evaluate_condition_expression("error_occurred", &state)
1407 .unwrap());
1408 }
1409
1410 #[allow(dead_code)] struct TestStateGraphExecutor {
1415 #[allow(dead_code)]
1416 max_iterations: usize,
1417 }
1418
1419 impl TestStateGraphExecutor {
1420 fn evaluate_edge_condition(
1421 &self,
1422 condition: &EdgeCondition,
1423 state: &GraphState,
1424 ) -> Result<bool> {
1425 match condition {
1426 EdgeCondition::HasToolCalls => Ok(state
1427 .parsed_decision
1428 .as_ref()
1429 .map(|d| d.has_tool_calls())
1430 .unwrap_or(false)),
1431 EdgeCondition::IsFinalAnswer => Ok(state
1432 .parsed_decision
1433 .as_ref()
1434 .map(|d| d.is_final())
1435 .unwrap_or(false)),
1436 EdgeCondition::RequestsDelegation => Ok(state
1437 .parsed_decision
1438 .as_ref()
1439 .map(|d| d.is_delegation())
1440 .unwrap_or(false)),
1441 EdgeCondition::ErrorOccurred => Ok(state.last_error.is_some()),
1442 EdgeCondition::StateEquals { variable, value } => {
1443 Ok(state.variables.get(variable) == Some(value))
1444 }
1445 EdgeCondition::Always => Ok(true),
1446 EdgeCondition::Custom { name, .. } => {
1447 tracing::warn!(
1448 "Custom condition '{}' not registered, defaulting to false",
1449 name
1450 );
1451 Ok(false)
1452 }
1453 EdgeCondition::ParadigmEquals { paradigm } => {
1454 Ok(state.active_paradigm.as_ref() == Some(paradigm))
1455 }
1456 EdgeCondition::IterationExceeds { count } => Ok(state.iteration_count > *count),
1457 }
1458 }
1459
1460 fn evaluate_condition_expression(
1461 &self,
1462 condition: &str,
1463 state: &GraphState,
1464 ) -> Result<bool> {
1465 if condition == "has_tool_calls" {
1466 return self.evaluate_edge_condition(&EdgeCondition::HasToolCalls, state);
1467 }
1468 if condition == "is_final_answer" {
1469 return self.evaluate_edge_condition(&EdgeCondition::IsFinalAnswer, state);
1470 }
1471 if condition == "error_occurred" {
1472 return self.evaluate_edge_condition(&EdgeCondition::ErrorOccurred, state);
1473 }
1474 if let Some((var, val)) = condition.split_once("==") {
1475 return Ok(state.variables.get(var.trim()) == Some(&val.trim().to_string()));
1476 }
1477 Ok(state
1478 .variables
1479 .get(condition)
1480 .map(|v| v == "true")
1481 .unwrap_or(false))
1482 }
1483 }
1484
1485 fn make_test_executor() -> TestStateGraphExecutor {
1486 TestStateGraphExecutor { max_iterations: 50 }
1487 }
1488
1489 #[test]
1490 fn test_graph_decision_enum() {
1491 let decision = oneai_core::GraphDecision::DirectAnswer {
1492 text: "42".to_string(),
1493 };
1494 assert!(decision.is_final());
1495 assert!(!decision.has_tool_calls());
1496
1497 let tool_calls = oneai_core::GraphDecision::ToolCalls { count: 2 };
1498 assert!(tool_calls.has_tool_calls());
1499 assert!(!tool_calls.is_final());
1500
1501 let delegate = oneai_core::GraphDecision::Delegate {
1502 agent_kind: "Explore".to_string(),
1503 task: "Search".to_string(),
1504 };
1505 assert!(delegate.is_delegation());
1506 assert!(!delegate.has_tool_calls());
1507
1508 let switch = oneai_core::GraphDecision::SwitchParadigm {
1509 paradigm: "plan".to_string(),
1510 };
1511 assert!(!switch.is_final());
1512 assert!(!switch.has_tool_calls());
1513 }
1514
1515 #[test]
1516 fn test_custom_condition_registry_lookup() {
1517 use std::sync::Arc;
1518 let mut registry: HashMap<String, CustomConditionFn> = HashMap::new();
1519 let mut state = GraphState::new();
1520 state
1521 .variables
1522 .insert("attempts".to_string(), "3".to_string());
1523
1524 assert!(!evaluate_custom_condition(
1526 ®istry,
1527 "not_registered",
1528 "missing",
1529 &state
1530 ));
1531
1532 registry.insert(
1534 "attempts_exhausted".to_string(),
1535 Arc::new(|s: &GraphState| {
1536 s.variables
1537 .get("attempts")
1538 .map(|v| v == "3")
1539 .unwrap_or(false)
1540 }),
1541 );
1542 assert!(evaluate_custom_condition(
1543 ®istry,
1544 "attempts_exhausted",
1545 "tries >= 3",
1546 &state
1547 ));
1548
1549 let other_state = GraphState::new();
1551 assert!(!evaluate_custom_condition(
1552 ®istry,
1553 "attempts_exhausted",
1554 "tries >= 3",
1555 &other_state
1556 ));
1557 }
1558
1559 use async_trait::async_trait;
1568 use oneai_core::InferenceResponse;
1569
1570 struct NullGraphActionExecutor;
1575
1576 #[async_trait]
1577 impl GraphActionExecutor for NullGraphActionExecutor {
1578 async fn execute_llm_infer(
1579 &self,
1580 _action: &NodeAction,
1581 _state: &mut GraphState,
1582 ) -> Result<ActionResult> {
1583 Err(OneAIError::Workflow(
1584 "NullGraphActionExecutor::execute_llm_infer \
1585 called — this test fixture only supports Delegate/ConditionCheck nodes"
1586 .into(),
1587 ))
1588 }
1589 async fn execute_tool_call(
1590 &self,
1591 _tool_name: &str,
1592 _args: &serde_json::Value,
1593 _state: &mut GraphState,
1594 ) -> Result<ActionResult> {
1595 Err(OneAIError::Workflow(
1596 "NullGraphActionExecutor::execute_tool_call called".into(),
1597 ))
1598 }
1599 async fn execute_paradigm_switch(
1600 &self,
1601 _paradigm: &str,
1602 state: &mut GraphState,
1603 ) -> Result<ActionResult> {
1604 state.active_paradigm = Some(_paradigm.to_string());
1607 state.parsed_decision = None;
1608 Ok(ActionResult {
1609 output: format!("Paradigm switched to: {}", _paradigm),
1610 error: None,
1611 })
1612 }
1613 async fn parse_decision(
1614 &self,
1615 _response: &InferenceResponse,
1616 _state: &mut GraphState,
1617 ) -> Result<oneai_core::GraphDecision> {
1618 Ok(oneai_core::GraphDecision::DirectAnswer {
1619 text: String::new(),
1620 })
1621 }
1622 }
1623
1624 struct RecordingDelegateFactory;
1629
1630 #[async_trait]
1631 impl DelegateFactory for RecordingDelegateFactory {
1632 async fn delegate(&self, agent_kind: &str, _task: &str) -> Result<String> {
1633 Ok(format!("result:{agent_kind}"))
1634 }
1635 }
1636
1637 fn make_frontier_executor() -> StateGraphExecutor {
1641 use oneai_core::platform::StubPlatformInteractionGate;
1642 StateGraphExecutor::new(
1643 Arc::new(NullGraphActionExecutor),
1644 Arc::new(RecordingDelegateFactory),
1645 Arc::new(StubPlatformInteractionGate::macos()) as Arc<dyn InteractionGate>,
1646 50,
1647 )
1648 }
1649
1650 fn make_fork_graph(
1653 entry_kind: &str,
1654 a_kind: &str,
1655 b_kind: &str,
1656 join_kind: &str,
1657 ) -> StateGraph {
1658 let mut graph = StateGraph::new("fork", "entry");
1659 graph.add_node(GraphNode {
1660 id: "entry".to_string(),
1661 action: NodeAction::Delegate {
1662 agent_kind: entry_kind.to_string(),
1663 task_template: "t".to_string(),
1664 },
1665 interrupt: false,
1666 metadata: HashMap::new(),
1667 });
1668 graph.add_node(GraphNode {
1669 id: "a".to_string(),
1670 action: NodeAction::Delegate {
1671 agent_kind: a_kind.to_string(),
1672 task_template: "t".to_string(),
1673 },
1674 interrupt: false,
1675 metadata: HashMap::new(),
1676 });
1677 graph.add_node(GraphNode {
1678 id: "b".to_string(),
1679 action: NodeAction::Delegate {
1680 agent_kind: b_kind.to_string(),
1681 task_template: "t".to_string(),
1682 },
1683 interrupt: false,
1684 metadata: HashMap::new(),
1685 });
1686 graph.add_node(GraphNode {
1687 id: "join".to_string(),
1688 action: NodeAction::Delegate {
1689 agent_kind: join_kind.to_string(),
1690 task_template: "t".to_string(),
1691 },
1692 interrupt: false,
1693 metadata: HashMap::new(),
1694 });
1695 graph.add_edge(GraphEdge {
1698 from: "entry".to_string(),
1699 to: "a".to_string(),
1700 condition: Some(EdgeCondition::Always),
1701 metadata: HashMap::new(),
1702 });
1703 graph.add_edge(GraphEdge {
1704 from: "entry".to_string(),
1705 to: "b".to_string(),
1706 condition: Some(EdgeCondition::Always),
1707 metadata: HashMap::new(),
1708 });
1709 graph.add_edge(GraphEdge {
1711 from: "a".to_string(),
1712 to: "join".to_string(),
1713 condition: Some(EdgeCondition::Always),
1714 metadata: HashMap::new(),
1715 });
1716 graph.add_edge(GraphEdge {
1717 from: "b".to_string(),
1718 to: "join".to_string(),
1719 condition: Some(EdgeCondition::Always),
1720 metadata: HashMap::new(),
1721 });
1722 graph.add_terminal("join");
1723 graph
1724 }
1725
1726 fn delegate_msg(kind: &str) -> String {
1729 format!("[Delegate {}]: result:{}", kind, kind)
1730 }
1731
1732 #[tokio::test]
1733 async fn frontier_fork_runs_both_branches_and_merges() {
1734 let graph = make_fork_graph("S", "A", "B", "J");
1739 let executor = make_frontier_executor();
1740 let result = executor.execute(&graph, GraphState::new()).await.unwrap();
1741
1742 assert!(result.completed, "should reach terminal join");
1743 assert_eq!(result.terminal_node.as_deref(), Some("join"));
1744 let texts: Vec<String> = result
1746 .final_state
1747 .conversation
1748 .messages
1749 .iter()
1750 .map(|m| m.text_content())
1751 .collect();
1752 assert_eq!(
1753 texts.len(),
1754 4,
1755 "expected 4 delegate messages, got {texts:?}"
1756 );
1757 assert_eq!(texts[0], delegate_msg("S"));
1758 assert_eq!(texts[1], delegate_msg("A"));
1759 assert_eq!(texts[2], delegate_msg("B"));
1760 assert_eq!(texts[3], delegate_msg("J"));
1761 }
1762
1763 #[tokio::test]
1764 async fn frontier_join_runs_only_after_both_branches_complete() {
1765 let graph = make_fork_graph("S", "A", "B", "J");
1770 let executor = make_frontier_executor();
1771 let result = executor.execute(&graph, GraphState::new()).await.unwrap();
1772
1773 let join_count = result
1775 .final_state
1776 .conversation
1777 .messages
1778 .iter()
1779 .filter(|m| m.text_content() == delegate_msg("J"))
1780 .count();
1781 assert_eq!(
1782 join_count, 1,
1783 "join must run exactly once (natural-join dedup)"
1784 );
1785 assert_eq!(result.iterations, 3);
1787 }
1788
1789 #[tokio::test]
1790 async fn frontier_single_branch_unchanged_from_legacy_behaviour() {
1791 let mut graph = StateGraph::new("linear", "entry");
1796 graph.add_node(GraphNode {
1797 id: "entry".to_string(),
1798 action: NodeAction::Delegate {
1799 agent_kind: "S".to_string(),
1800 task_template: "t".to_string(),
1801 },
1802 interrupt: false,
1803 metadata: HashMap::new(),
1804 });
1805 graph.add_node(GraphNode {
1806 id: "join".to_string(),
1807 action: NodeAction::Delegate {
1808 agent_kind: "J".to_string(),
1809 task_template: "t".to_string(),
1810 },
1811 interrupt: false,
1812 metadata: HashMap::new(),
1813 });
1814 graph.add_edge(GraphEdge {
1815 from: "entry".to_string(),
1816 to: "join".to_string(),
1817 condition: Some(EdgeCondition::Always),
1818 metadata: HashMap::new(),
1819 });
1820 graph.add_terminal("join");
1821
1822 let executor = make_frontier_executor();
1823 let result = executor.execute(&graph, GraphState::new()).await.unwrap();
1824 assert!(result.completed);
1825 assert_eq!(result.terminal_node.as_deref(), Some("join"));
1826 assert_eq!(result.iterations, 2);
1827 let texts: Vec<String> = result
1828 .final_state
1829 .conversation
1830 .messages
1831 .iter()
1832 .map(|m| m.text_content())
1833 .collect();
1834 assert_eq!(texts, vec![delegate_msg("S"), delegate_msg("J")]);
1835 }
1836}