use std::collections::{BTreeSet, HashMap};
use std::sync::Arc;
use futures::future::join_all;
use oneai_core::error::{OneAIError, Result};
use oneai_core::traits::{InteractionGate, LlmProvider, PermissionResolver, Tool};
use oneai_core::{InferenceRequest, InferenceResponse, Message, PermissionAction, Role};
use crate::state_graph::{
EdgeCondition, GraphEdge, GraphExecutionResult, GraphState, NodeAction, StateGraph,
};
pub type CustomConditionFn = Arc<dyn Fn(&GraphState) -> bool + Send + Sync>;
pub(crate) fn evaluate_custom_condition(
registry: &HashMap<String, CustomConditionFn>,
name: &str,
description: &str,
state: &GraphState,
) -> bool {
match registry.get(name) {
Some(evaluator) => evaluator(state),
None => {
tracing::warn!(
"Custom condition '{}' ('{}') not registered. Defaulting to false.",
name,
description
);
false
}
}
}
#[async_trait::async_trait]
pub trait DelegateFactory: Send + Sync {
async fn delegate(&self, agent_kind: &str, task: &str) -> Result<String>;
}
pub struct NoopDelegateFactory;
#[async_trait::async_trait]
impl DelegateFactory for NoopDelegateFactory {
async fn delegate(&self, agent_kind: &str, _task: &str) -> Result<String> {
Err(OneAIError::Workflow(format!(
"Delegate action '{}' not supported — no DelegateFactory configured",
agent_kind
)))
}
}
#[derive(Debug, Clone)]
pub struct ActionResult {
pub output: String,
pub error: Option<String>,
}
struct BranchResult {
node_id: String,
action_result: ActionResult,
branch_state: GraphState,
next_targets: Vec<String>,
}
#[async_trait::async_trait]
pub trait GraphActionExecutor: Send + Sync {
async fn execute_llm_infer(
&self,
action: &NodeAction,
state: &mut GraphState,
) -> Result<ActionResult>;
async fn execute_tool_call(
&self,
tool_name: &str,
args: &serde_json::Value,
state: &mut GraphState,
) -> Result<ActionResult>;
async fn execute_paradigm_switch(
&self,
paradigm: &str,
state: &mut GraphState,
) -> Result<ActionResult>;
async fn parse_decision(
&self,
response: &InferenceResponse,
state: &mut GraphState,
) -> Result<oneai_core::GraphDecision>;
}
pub struct DirectProviderActionExecutor {
provider: Arc<dyn LlmProvider>,
tools: Arc<tokio::sync::RwLock<HashMap<String, Arc<dyn Tool>>>>,
permission_resolver: Option<Arc<dyn PermissionResolver>>,
}
impl DirectProviderActionExecutor {
pub fn new(
provider: Arc<dyn LlmProvider>,
tools: Arc<tokio::sync::RwLock<HashMap<String, Arc<dyn Tool>>>>,
) -> Self {
Self {
provider,
tools,
permission_resolver: None,
}
}
pub fn with_permission_resolver(mut self, resolver: Arc<dyn PermissionResolver>) -> Self {
self.permission_resolver = Some(resolver);
self
}
}
#[async_trait::async_trait]
impl GraphActionExecutor for DirectProviderActionExecutor {
async fn execute_llm_infer(
&self,
action: &NodeAction,
state: &mut GraphState,
) -> Result<ActionResult> {
let (
system_prompt_override,
_use_streaming,
include_tool_definitions,
tool_filter_override,
thinking_budget,
temperature,
max_tokens,
) = match action {
NodeAction::LlmInfer {
system_prompt_override,
use_streaming,
include_tool_definitions,
tool_filter_override,
thinking_budget,
temperature,
max_tokens,
} => (
system_prompt_override.clone(),
*use_streaming,
*include_tool_definitions,
tool_filter_override.clone(),
*thinking_budget,
*temperature,
*max_tokens,
),
_ => return Err(OneAIError::Workflow("Expected LlmInfer action".to_string())),
};
let system_prompt = system_prompt_override
.unwrap_or_else(|| "You are an intelligent agent. Respond to the task.".to_string());
let mut conversation = state.conversation.clone();
if !conversation.messages.iter().any(|m| m.role == Role::System) {
conversation.add_message(Message::system(&system_prompt));
}
let tool_defs = if include_tool_definitions {
let tools_map = self.tools.read().await;
if let Some(filter) = &tool_filter_override {
tools_map
.values()
.filter(|t| filter.contains(&t.name().to_string()))
.map(|t| oneai_core::ToolDefinition {
name: t.name().to_string(),
description: t.description().to_string(),
parameters_schema: t.parameters_schema(),
})
.collect()
} else {
tools_map
.values()
.map(|t| oneai_core::ToolDefinition {
name: t.name().to_string(),
description: t.description().to_string(),
parameters_schema: t.parameters_schema(),
})
.collect()
}
} else {
vec![] };
let request = InferenceRequest {
conversation,
tools: tool_defs,
max_tokens: max_tokens.or(Some(4096)),
temperature: temperature.or(Some(0.3)),
top_p: None,
stop_sequences: vec![],
constrained_output: None,
thinking_budget,
metadata: HashMap::new(),
};
let response = self.provider.infer(request).await?;
let output = response.message.text_content();
state.conversation.add_message(response.message.clone());
let _decision = self.parse_decision(&response, state).await?;
Ok(ActionResult {
output,
error: None,
})
}
async fn execute_tool_call(
&self,
tool_name: &str,
args: &serde_json::Value,
state: &mut GraphState,
) -> Result<ActionResult> {
if let Some(resolver) = self.permission_resolver.as_ref() {
match resolver.resolve(tool_name, args) {
PermissionAction::Deny { reason } => {
tracing::warn!(
"ToolCall node '{}' denied by domain policy: {}",
tool_name,
reason
);
return Ok(ActionResult {
output: format!("Denied by domain policy: {}", reason),
error: Some(format!("Denied by domain policy: {}", reason)),
});
}
PermissionAction::AutoApprove => {
tracing::info!(
"ToolCall node '{}' auto-approved by domain policy",
tool_name
);
}
_ => {}
}
}
let tools_map = self.tools.read().await;
let tool = tools_map.get(tool_name).ok_or_else(|| {
OneAIError::Workflow(format!("Tool '{}' not found for ToolCall node", tool_name))
})?;
let output = tool.execute(args.clone()).await?;
state.conversation.add_message(Message::tool_result(
format!("graph_tool_{}", tool_name),
output.content.clone(),
));
Ok(ActionResult {
output: output.content,
error: output.error,
})
}
async fn execute_paradigm_switch(
&self,
paradigm: &str,
state: &mut GraphState,
) -> Result<ActionResult> {
state.active_paradigm = Some(paradigm.to_string());
state.parsed_decision = None;
Ok(ActionResult {
output: format!("Paradigm switched to: {}", paradigm),
error: None,
})
}
async fn parse_decision(
&self,
response: &InferenceResponse,
state: &mut GraphState,
) -> Result<oneai_core::GraphDecision> {
let mut tool_calls = Vec::new();
let mut text_parts = Vec::new();
for block in &response.message.content {
match block {
oneai_core::ContentBlock::ToolCall { id: _, name, args } => {
if name == "delegate" {
let args_value: serde_json::Value =
serde_json::from_str(args).unwrap_or_else(|_| serde_json::json!({}));
let agent_kind = args_value
.get("agent_type")
.and_then(|v| v.as_str())
.unwrap_or("Explore")
.to_string();
let task = args_value
.get("task")
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string();
let decision = oneai_core::GraphDecision::Delegate { agent_kind, task };
state.parsed_decision = Some(decision.clone());
return Ok(decision);
}
if name == "switch_paradigm" {
let args_value: serde_json::Value =
serde_json::from_str(args).unwrap_or_else(|_| serde_json::json!({}));
let paradigm = args_value
.get("paradigm")
.and_then(|v| v.as_str())
.unwrap_or("react")
.to_string();
let decision = oneai_core::GraphDecision::SwitchParadigm { paradigm };
state.parsed_decision = Some(decision.clone());
return Ok(decision);
}
tool_calls.push(name.clone());
}
oneai_core::ContentBlock::Text { text } => {
text_parts.push(text.clone());
}
_ => {}
}
}
let decision = if !tool_calls.is_empty() {
oneai_core::GraphDecision::ToolCalls {
count: tool_calls.len(),
}
} else {
oneai_core::GraphDecision::DirectAnswer {
text: text_parts.join("\n"),
}
};
state.parsed_decision = Some(decision.clone());
Ok(decision)
}
}
pub struct StateGraphExecutor {
action_executor: Arc<dyn GraphActionExecutor>,
delegate_factory: Arc<dyn DelegateFactory>,
interaction_gate: Arc<dyn InteractionGate>,
max_iterations: usize,
custom_conditions: HashMap<String, CustomConditionFn>,
}
impl StateGraphExecutor {
pub fn new(
action_executor: Arc<dyn GraphActionExecutor>,
delegate_factory: Arc<dyn DelegateFactory>,
interaction_gate: Arc<dyn InteractionGate>,
max_iterations: usize,
) -> Self {
Self {
action_executor,
delegate_factory,
interaction_gate,
max_iterations,
custom_conditions: HashMap::new(),
}
}
pub fn with_defaults(
action_executor: Arc<dyn GraphActionExecutor>,
delegate_factory: Arc<dyn DelegateFactory>,
interaction_gate: Arc<dyn InteractionGate>,
) -> Self {
Self::new(action_executor, delegate_factory, interaction_gate, 50)
}
pub fn with_custom_condition(
mut self,
name: impl Into<String>,
evaluator: CustomConditionFn,
) -> Self {
self.custom_conditions.insert(name.into(), evaluator);
self
}
pub fn with_direct_provider(
provider: Arc<dyn LlmProvider>,
tools: Arc<tokio::sync::RwLock<HashMap<String, Arc<dyn Tool>>>>,
delegate_factory: Arc<dyn DelegateFactory>,
interaction_gate: Arc<dyn InteractionGate>,
max_iterations: usize,
) -> Self {
let action_executor = Arc::new(DirectProviderActionExecutor::new(provider, tools));
Self::new(
action_executor,
delegate_factory,
interaction_gate,
max_iterations,
)
}
pub fn with_direct_provider_defaults(
provider: Arc<dyn LlmProvider>,
tools: Arc<tokio::sync::RwLock<HashMap<String, Arc<dyn Tool>>>>,
delegate_factory: Arc<dyn DelegateFactory>,
interaction_gate: Arc<dyn InteractionGate>,
) -> Self {
Self::with_direct_provider(provider, tools, delegate_factory, interaction_gate, 50)
}
pub async fn execute(
&self,
graph: &StateGraph,
initial_state: GraphState,
) -> Result<GraphExecutionResult> {
let mut state = initial_state;
let mut frontier: BTreeSet<String> = std::iter::once(graph.entry_point.clone()).collect();
let mut iterations = 0;
let mut interrupt_checkpoints = Vec::new();
while iterations < self.max_iterations {
iterations += 1;
state.iteration_count = iterations;
if frontier.is_empty() {
break;
}
if let Some(term_id) = frontier
.iter()
.find(|id| graph.terminal_nodes.contains(id))
.cloned()
{
let node = graph.get_node(&term_id).ok_or_else(|| {
OneAIError::Workflow(format!(
"Terminal node '{}' not found in graph '{}'",
term_id, graph.name
))
})?;
let action_result = self.execute_node_action(&node.action, &mut state).await?;
state.last_result = Some(action_result.output.clone());
state.last_error = action_result.error.clone();
return Ok(GraphExecutionResult {
name: graph.name.clone(),
final_state: state,
completed: true,
terminal_node: Some(term_id),
iterations,
interrupt_checkpoints,
});
}
let has_interrupt = frontier
.iter()
.any(|id| graph.get_node(id).map(|n| n.interrupt).unwrap_or(false));
let next_targets: Vec<String> = if frontier.len() == 1 || has_interrupt {
self.execute_frontier_sequential(
graph,
&frontier,
&mut state,
&mut interrupt_checkpoints,
)
.await?
} else {
self.execute_frontier_parallel(graph, &frontier, &mut state)
.await?
};
if state.should_terminate {
break;
}
if next_targets.is_empty() {
tracing::warn!(
"No matching edge condition from frontier {:?} in graph '{}'. Terminating.",
frontier,
graph.name
);
break;
}
frontier = next_targets.into_iter().collect();
}
let should_terminate = state.should_terminate;
if iterations >= self.max_iterations {
tracing::warn!(
"StateGraph '{}' exceeded max iterations ({}). Terminating.",
graph.name,
self.max_iterations
);
}
Ok(GraphExecutionResult {
name: graph.name.clone(),
final_state: state,
completed: !should_terminate,
terminal_node: None,
iterations,
interrupt_checkpoints,
})
}
async fn execute_frontier_sequential(
&self,
graph: &StateGraph,
frontier: &BTreeSet<String>,
state: &mut GraphState,
interrupt_checkpoints: &mut Vec<String>,
) -> Result<Vec<String>> {
let mut next = Vec::new();
for node_id in frontier {
let node = graph.get_node(node_id).ok_or_else(|| {
OneAIError::Workflow(format!(
"Node '{}' not found in graph '{}'",
node_id, graph.name
))
})?;
if node.interrupt {
let checkpoint_id = format!("interrupt_{}_{}", graph.name, node_id);
interrupt_checkpoints.push(checkpoint_id.clone());
let approval_request = oneai_core::ApprovalRequest {
tool_name: "state_graph_interrupt".into(),
args: serde_json::json!({
"node": node_id,
"description": match &node.action {
NodeAction::HumanApproval { description } => description.clone(),
_ => "Interrupt point reached".to_string(),
},
"state": state.variables,
}),
risk_level: oneai_core::RiskLevel::Medium,
permission_level: Some(oneai_core::PermissionLevel::Standard),
justification: format!(
"StateGraph interrupt at node '{}' in graph '{}'",
node_id, graph.name
),
};
let approval = self
.interaction_gate
.request(oneai_core::InteractionRequest::ToolApproval {
approval: approval_request,
})
.await?;
match approval {
oneai_core::InteractionResponse::Abort { reason } => {
state.should_terminate = true;
state.last_error = Some(format!("Interrupt denied: {}", reason));
return Ok(next);
}
oneai_core::InteractionResponse::Revise { feedback } => {
state.should_terminate = true;
state.last_error = Some(format!("Interrupt denied: {}", feedback));
return Ok(next);
}
_ => { }
}
}
let action_result = self.execute_node_action(&node.action, state).await?;
state.last_result = Some(action_result.output.clone());
state.last_error = action_result.error.clone();
if state.should_terminate {
return Ok(next);
}
let edges = graph.get_edges_from(node_id);
next.extend(self.route_next_nodes(&edges, state)?);
}
Ok(next)
}
async fn execute_frontier_parallel(
&self,
graph: &StateGraph,
frontier: &BTreeSet<String>,
state: &mut GraphState,
) -> Result<Vec<String>> {
let baseline_msgs = state.conversation.messages.len();
let incoming = state.clone();
let branch_futs: Vec<_> = frontier
.iter()
.map(|node_id| {
let seed = incoming.clone();
async move {
let node = match graph.get_node(node_id) {
Some(n) => n,
None => {
return BranchResult {
node_id: node_id.clone(),
action_result: ActionResult {
output: String::new(),
error: Some(format!(
"Node '{}' not found in graph '{}'",
node_id, graph.name
)),
},
branch_state: seed,
next_targets: Vec::new(),
};
}
};
let mut branch_state = seed;
let action_result = match self
.execute_node_action(&node.action, &mut branch_state)
.await
{
Ok(ar) => ar,
Err(e) => ActionResult {
output: String::new(),
error: Some(e.to_string()),
},
};
let edges = graph.get_edges_from(node_id);
let next_targets = self
.route_next_nodes(&edges, &branch_state)
.unwrap_or_default();
BranchResult {
node_id: node_id.clone(),
action_result,
branch_state,
next_targets,
}
}
})
.collect();
let mut branches: Vec<BranchResult> = join_all(branch_futs).await;
branches.sort_by(|a, b| a.node_id.cmp(&b.node_id));
let mut joined_output = Vec::new();
let mut first_error: Option<String> = None;
let mut any_terminate = false;
for b in &branches {
joined_output.push(b.action_result.output.clone());
if first_error.is_none() {
first_error = b.action_result.error.clone();
}
any_terminate |= b.branch_state.should_terminate;
}
state.last_result = Some(joined_output.join("\n---\n"));
state.last_error = first_error;
state.should_terminate = any_terminate;
for b in &branches {
let extra_from = baseline_msgs.min(b.branch_state.conversation.messages.len());
for m in &b.branch_state.conversation.messages[extra_from..] {
state.conversation.add_message(m.clone());
}
}
for b in &branches {
for (k, v) in &b.branch_state.variables {
state.variables.insert(k.clone(), v.clone());
}
}
if let Some(last) = branches.last() {
state.parsed_decision = last.branch_state.parsed_decision.clone();
}
let mut next = Vec::new();
for b in &branches {
next.extend(b.next_targets.iter().cloned());
}
Ok(next)
}
async fn execute_node_action(
&self,
action: &NodeAction,
state: &mut GraphState,
) -> Result<ActionResult> {
match action {
NodeAction::LlmInfer { .. } => {
self.action_executor.execute_llm_infer(action, state).await
}
NodeAction::ToolCall {
tool_name,
args_template,
} => {
let resolved_name = interpolate_graph_template(tool_name, &state.variables);
let resolved_args = if let Some(template) = args_template {
let json_str = interpolate_graph_template(template, &state.variables);
serde_json::from_str(&json_str).unwrap_or(serde_json::json!({}))
} else {
serde_json::json!({})
};
self.action_executor
.execute_tool_call(&resolved_name, &resolved_args, state)
.await
}
NodeAction::Delegate {
agent_kind,
task_template,
} => {
let task = interpolate_graph_template(task_template, &state.variables);
let result = self.delegate_factory.delegate(agent_kind, &task).await?;
state.conversation.add_message(Message::assistant(format!(
"[Delegate {}]: {}",
agent_kind, result
)));
state.parsed_decision = Some(oneai_core::GraphDecision::Delegate {
agent_kind: agent_kind.clone(),
task,
});
Ok(ActionResult {
output: result,
error: None,
})
}
NodeAction::HumanApproval { description } => {
Ok(ActionResult {
output: description.clone(),
error: None,
})
}
NodeAction::ConditionCheck { condition } => {
let result = self.evaluate_condition_expression(condition, state)?;
state
.variables
.insert("_condition_result".to_string(), result.to_string());
Ok(ActionResult {
output: result.to_string(),
error: None,
})
}
NodeAction::SwitchParadigm { paradigm } => {
self.action_executor
.execute_paradigm_switch(paradigm, state)
.await
}
}
}
fn route_next_nodes(&self, edges: &[&GraphEdge], state: &GraphState) -> Result<Vec<String>> {
let mut next = Vec::with_capacity(edges.len());
for edge in edges {
let satisfied = match &edge.condition {
Some(condition) => self.evaluate_edge_condition(condition, state)?,
None => true, };
if satisfied {
next.push(edge.to.clone());
}
}
Ok(next)
}
fn evaluate_edge_condition(
&self,
condition: &EdgeCondition,
state: &GraphState,
) -> Result<bool> {
match condition {
EdgeCondition::HasToolCalls => {
Ok(state
.parsed_decision
.as_ref()
.map(|d| d.has_tool_calls())
.unwrap_or(false))
}
EdgeCondition::IsFinalAnswer => {
Ok(state
.parsed_decision
.as_ref()
.map(|d| d.is_final())
.unwrap_or(false))
}
EdgeCondition::RequestsDelegation => {
Ok(state
.parsed_decision
.as_ref()
.map(|d| d.is_delegation())
.unwrap_or(false))
}
EdgeCondition::ErrorOccurred => Ok(state.last_error.is_some()),
EdgeCondition::StateEquals { variable, value } => {
Ok(state.variables.get(variable) == Some(value))
}
EdgeCondition::Always => Ok(true),
EdgeCondition::Custom { name, description } => Ok(evaluate_custom_condition(
&self.custom_conditions,
name,
description,
state,
)),
EdgeCondition::ParadigmEquals { paradigm } => {
Ok(state.active_paradigm.as_ref() == Some(paradigm))
}
EdgeCondition::IterationExceeds { count } => Ok(state.iteration_count > *count),
}
}
fn evaluate_condition_expression(&self, condition: &str, state: &GraphState) -> Result<bool> {
if condition == "has_tool_calls" {
return self.evaluate_edge_condition(&EdgeCondition::HasToolCalls, state);
}
if condition == "is_final_answer" {
return self.evaluate_edge_condition(&EdgeCondition::IsFinalAnswer, state);
}
if condition == "error_occurred" {
return self.evaluate_edge_condition(&EdgeCondition::ErrorOccurred, state);
}
if let Some((var, val)) = condition.split_once("==") {
return Ok(state.variables.get(var.trim()) == Some(&val.trim().to_string()));
}
Ok(state
.variables
.get(condition)
.map(|v| v == "true")
.unwrap_or(false))
}
}
pub fn interpolate_graph_template(template: &str, variables: &HashMap<String, String>) -> String {
let mut result = template.to_string();
for (key, value) in variables {
result = result.replace(&format!("{{{{{}}}}}", key), value);
}
result
}
#[cfg(test)]
mod tests {
use super::*;
use crate::state_graph::{GraphNode, StateGraph};
#[allow(dead_code)] fn make_simple_graph() -> StateGraph {
let mut graph = StateGraph::new("test", "start");
graph.add_node(GraphNode {
id: "start".to_string(),
action: NodeAction::ConditionCheck {
condition: "ready==true".to_string(),
},
interrupt: false,
metadata: HashMap::new(),
});
graph.add_node(GraphNode {
id: "end".to_string(),
action: NodeAction::LlmInfer {
system_prompt_override: Some("Final answer".to_string()),
use_streaming: false,
include_tool_definitions: false,
tool_filter_override: None,
thinking_budget: None,
temperature: None,
max_tokens: None,
},
interrupt: false,
metadata: HashMap::new(),
});
graph.add_edge(GraphEdge {
from: "start".to_string(),
to: "end".to_string(),
condition: Some(EdgeCondition::Always),
metadata: HashMap::new(),
});
graph.add_terminal("end".to_string());
graph
}
#[test]
fn test_evaluate_edge_condition_always() {
let state = GraphState::new();
let executor = make_test_executor();
let result = executor
.evaluate_edge_condition(&EdgeCondition::Always, &state)
.unwrap();
assert!(result);
}
#[test]
fn test_evaluate_edge_condition_error_occurred() {
let mut state = GraphState::new();
let executor = make_test_executor();
assert!(!executor
.evaluate_edge_condition(&EdgeCondition::ErrorOccurred, &state)
.unwrap());
state.last_error = Some("test error".to_string());
assert!(executor
.evaluate_edge_condition(&EdgeCondition::ErrorOccurred, &state)
.unwrap());
}
#[test]
fn test_evaluate_edge_condition_state_equals() {
let mut state = GraphState::new();
let executor = make_test_executor();
let cond = EdgeCondition::StateEquals {
variable: "mode".to_string(),
value: "react".to_string(),
};
assert!(!executor.evaluate_edge_condition(&cond, &state).unwrap());
state
.variables
.insert("mode".to_string(), "plan".to_string());
assert!(!executor.evaluate_edge_condition(&cond, &state).unwrap());
state
.variables
.insert("mode".to_string(), "react".to_string());
assert!(executor.evaluate_edge_condition(&cond, &state).unwrap());
}
#[test]
fn test_evaluate_edge_condition_has_tool_calls() {
let mut state = GraphState::new();
let executor = make_test_executor();
assert!(!executor
.evaluate_edge_condition(&EdgeCondition::HasToolCalls, &state)
.unwrap());
state.parsed_decision = Some(oneai_core::GraphDecision::DirectAnswer {
text: "Just a text answer".to_string(),
});
assert!(!executor
.evaluate_edge_condition(&EdgeCondition::HasToolCalls, &state)
.unwrap());
state.parsed_decision = Some(oneai_core::GraphDecision::ToolCalls { count: 1 });
assert!(executor
.evaluate_edge_condition(&EdgeCondition::HasToolCalls, &state)
.unwrap());
}
#[test]
fn test_evaluate_edge_condition_is_final_answer() {
let mut state = GraphState::new();
let executor = make_test_executor();
assert!(!executor
.evaluate_edge_condition(&EdgeCondition::IsFinalAnswer, &state)
.unwrap());
state.parsed_decision = Some(oneai_core::GraphDecision::DirectAnswer {
text: "The answer is 42".to_string(),
});
assert!(executor
.evaluate_edge_condition(&EdgeCondition::IsFinalAnswer, &state)
.unwrap());
state.parsed_decision = Some(oneai_core::GraphDecision::ToolCalls { count: 2 });
assert!(!executor
.evaluate_edge_condition(&EdgeCondition::IsFinalAnswer, &state)
.unwrap());
}
#[test]
fn test_evaluate_edge_condition_requests_delegation() {
let mut state = GraphState::new();
let executor = make_test_executor();
state.parsed_decision = Some(oneai_core::GraphDecision::Delegate {
agent_kind: "Explore".to_string(),
task: "Search the codebase".to_string(),
});
assert!(executor
.evaluate_edge_condition(&EdgeCondition::RequestsDelegation, &state)
.unwrap());
state.parsed_decision = Some(oneai_core::GraphDecision::ToolCalls { count: 1 });
assert!(!executor
.evaluate_edge_condition(&EdgeCondition::RequestsDelegation, &state)
.unwrap());
}
#[test]
fn test_evaluate_edge_condition_paradigm_equals() {
let mut state = GraphState::new();
let executor = make_test_executor();
let cond = EdgeCondition::ParadigmEquals {
paradigm: "react".to_string(),
};
assert!(!executor.evaluate_edge_condition(&cond, &state).unwrap());
state.active_paradigm = Some("plan".to_string());
assert!(!executor.evaluate_edge_condition(&cond, &state).unwrap());
state.active_paradigm = Some("react".to_string());
assert!(executor.evaluate_edge_condition(&cond, &state).unwrap());
}
#[test]
fn test_evaluate_edge_condition_iteration_exceeds() {
let mut state = GraphState::new();
let executor = make_test_executor();
let cond = EdgeCondition::IterationExceeds { count: 5 };
assert!(!executor.evaluate_edge_condition(&cond, &state).unwrap());
state.iteration_count = 10;
assert!(executor.evaluate_edge_condition(&cond, &state).unwrap());
}
#[test]
fn test_interpolate_graph_template() {
let vars = HashMap::from([
("selected_tool".to_string(), "shell".to_string()),
("command".to_string(), "ls -la".to_string()),
]);
let result = interpolate_graph_template("{{selected_tool}} with {{command}}", &vars);
assert_eq!(result, "shell with ls -la");
}
#[test]
fn test_evaluate_condition_expression() {
let mut state = GraphState::new();
let executor = make_test_executor();
state
.variables
.insert("mode".to_string(), "react".to_string());
assert!(executor
.evaluate_condition_expression("mode==react", &state)
.unwrap());
assert!(!executor
.evaluate_condition_expression("mode==plan", &state)
.unwrap());
state.last_error = Some("error".to_string());
assert!(executor
.evaluate_condition_expression("error_occurred", &state)
.unwrap());
}
#[allow(dead_code)] struct TestStateGraphExecutor {
#[allow(dead_code)]
max_iterations: usize,
}
impl TestStateGraphExecutor {
fn evaluate_edge_condition(
&self,
condition: &EdgeCondition,
state: &GraphState,
) -> Result<bool> {
match condition {
EdgeCondition::HasToolCalls => Ok(state
.parsed_decision
.as_ref()
.map(|d| d.has_tool_calls())
.unwrap_or(false)),
EdgeCondition::IsFinalAnswer => Ok(state
.parsed_decision
.as_ref()
.map(|d| d.is_final())
.unwrap_or(false)),
EdgeCondition::RequestsDelegation => Ok(state
.parsed_decision
.as_ref()
.map(|d| d.is_delegation())
.unwrap_or(false)),
EdgeCondition::ErrorOccurred => Ok(state.last_error.is_some()),
EdgeCondition::StateEquals { variable, value } => {
Ok(state.variables.get(variable) == Some(value))
}
EdgeCondition::Always => Ok(true),
EdgeCondition::Custom { name, .. } => {
tracing::warn!(
"Custom condition '{}' not registered, defaulting to false",
name
);
Ok(false)
}
EdgeCondition::ParadigmEquals { paradigm } => {
Ok(state.active_paradigm.as_ref() == Some(paradigm))
}
EdgeCondition::IterationExceeds { count } => Ok(state.iteration_count > *count),
}
}
fn evaluate_condition_expression(
&self,
condition: &str,
state: &GraphState,
) -> Result<bool> {
if condition == "has_tool_calls" {
return self.evaluate_edge_condition(&EdgeCondition::HasToolCalls, state);
}
if condition == "is_final_answer" {
return self.evaluate_edge_condition(&EdgeCondition::IsFinalAnswer, state);
}
if condition == "error_occurred" {
return self.evaluate_edge_condition(&EdgeCondition::ErrorOccurred, state);
}
if let Some((var, val)) = condition.split_once("==") {
return Ok(state.variables.get(var.trim()) == Some(&val.trim().to_string()));
}
Ok(state
.variables
.get(condition)
.map(|v| v == "true")
.unwrap_or(false))
}
}
fn make_test_executor() -> TestStateGraphExecutor {
TestStateGraphExecutor { max_iterations: 50 }
}
#[test]
fn test_graph_decision_enum() {
let decision = oneai_core::GraphDecision::DirectAnswer {
text: "42".to_string(),
};
assert!(decision.is_final());
assert!(!decision.has_tool_calls());
let tool_calls = oneai_core::GraphDecision::ToolCalls { count: 2 };
assert!(tool_calls.has_tool_calls());
assert!(!tool_calls.is_final());
let delegate = oneai_core::GraphDecision::Delegate {
agent_kind: "Explore".to_string(),
task: "Search".to_string(),
};
assert!(delegate.is_delegation());
assert!(!delegate.has_tool_calls());
let switch = oneai_core::GraphDecision::SwitchParadigm {
paradigm: "plan".to_string(),
};
assert!(!switch.is_final());
assert!(!switch.has_tool_calls());
}
#[test]
fn test_custom_condition_registry_lookup() {
use std::sync::Arc;
let mut registry: HashMap<String, CustomConditionFn> = HashMap::new();
let mut state = GraphState::new();
state
.variables
.insert("attempts".to_string(), "3".to_string());
assert!(!evaluate_custom_condition(
®istry,
"not_registered",
"missing",
&state
));
registry.insert(
"attempts_exhausted".to_string(),
Arc::new(|s: &GraphState| {
s.variables
.get("attempts")
.map(|v| v == "3")
.unwrap_or(false)
}),
);
assert!(evaluate_custom_condition(
®istry,
"attempts_exhausted",
"tries >= 3",
&state
));
let other_state = GraphState::new();
assert!(!evaluate_custom_condition(
®istry,
"attempts_exhausted",
"tries >= 3",
&other_state
));
}
use async_trait::async_trait;
use oneai_core::InferenceResponse;
struct NullGraphActionExecutor;
#[async_trait]
impl GraphActionExecutor for NullGraphActionExecutor {
async fn execute_llm_infer(
&self,
_action: &NodeAction,
_state: &mut GraphState,
) -> Result<ActionResult> {
Err(OneAIError::Workflow(
"NullGraphActionExecutor::execute_llm_infer \
called — this test fixture only supports Delegate/ConditionCheck nodes"
.into(),
))
}
async fn execute_tool_call(
&self,
_tool_name: &str,
_args: &serde_json::Value,
_state: &mut GraphState,
) -> Result<ActionResult> {
Err(OneAIError::Workflow(
"NullGraphActionExecutor::execute_tool_call called".into(),
))
}
async fn execute_paradigm_switch(
&self,
_paradigm: &str,
state: &mut GraphState,
) -> Result<ActionResult> {
state.active_paradigm = Some(_paradigm.to_string());
state.parsed_decision = None;
Ok(ActionResult {
output: format!("Paradigm switched to: {}", _paradigm),
error: None,
})
}
async fn parse_decision(
&self,
_response: &InferenceResponse,
_state: &mut GraphState,
) -> Result<oneai_core::GraphDecision> {
Ok(oneai_core::GraphDecision::DirectAnswer {
text: String::new(),
})
}
}
struct RecordingDelegateFactory;
#[async_trait]
impl DelegateFactory for RecordingDelegateFactory {
async fn delegate(&self, agent_kind: &str, _task: &str) -> Result<String> {
Ok(format!("result:{agent_kind}"))
}
}
fn make_frontier_executor() -> StateGraphExecutor {
use oneai_core::platform::StubPlatformInteractionGate;
StateGraphExecutor::new(
Arc::new(NullGraphActionExecutor),
Arc::new(RecordingDelegateFactory),
Arc::new(StubPlatformInteractionGate::macos()) as Arc<dyn InteractionGate>,
50,
)
}
fn make_fork_graph(
entry_kind: &str,
a_kind: &str,
b_kind: &str,
join_kind: &str,
) -> StateGraph {
let mut graph = StateGraph::new("fork", "entry");
graph.add_node(GraphNode {
id: "entry".to_string(),
action: NodeAction::Delegate {
agent_kind: entry_kind.to_string(),
task_template: "t".to_string(),
},
interrupt: false,
metadata: HashMap::new(),
});
graph.add_node(GraphNode {
id: "a".to_string(),
action: NodeAction::Delegate {
agent_kind: a_kind.to_string(),
task_template: "t".to_string(),
},
interrupt: false,
metadata: HashMap::new(),
});
graph.add_node(GraphNode {
id: "b".to_string(),
action: NodeAction::Delegate {
agent_kind: b_kind.to_string(),
task_template: "t".to_string(),
},
interrupt: false,
metadata: HashMap::new(),
});
graph.add_node(GraphNode {
id: "join".to_string(),
action: NodeAction::Delegate {
agent_kind: join_kind.to_string(),
task_template: "t".to_string(),
},
interrupt: false,
metadata: HashMap::new(),
});
graph.add_edge(GraphEdge {
from: "entry".to_string(),
to: "a".to_string(),
condition: Some(EdgeCondition::Always),
metadata: HashMap::new(),
});
graph.add_edge(GraphEdge {
from: "entry".to_string(),
to: "b".to_string(),
condition: Some(EdgeCondition::Always),
metadata: HashMap::new(),
});
graph.add_edge(GraphEdge {
from: "a".to_string(),
to: "join".to_string(),
condition: Some(EdgeCondition::Always),
metadata: HashMap::new(),
});
graph.add_edge(GraphEdge {
from: "b".to_string(),
to: "join".to_string(),
condition: Some(EdgeCondition::Always),
metadata: HashMap::new(),
});
graph.add_terminal("join");
graph
}
fn delegate_msg(kind: &str) -> String {
format!("[Delegate {}]: result:{}", kind, kind)
}
#[tokio::test]
async fn frontier_fork_runs_both_branches_and_merges() {
let graph = make_fork_graph("S", "A", "B", "J");
let executor = make_frontier_executor();
let result = executor.execute(&graph, GraphState::new()).await.unwrap();
assert!(result.completed, "should reach terminal join");
assert_eq!(result.terminal_node.as_deref(), Some("join"));
let texts: Vec<String> = result
.final_state
.conversation
.messages
.iter()
.map(|m| m.text_content())
.collect();
assert_eq!(
texts.len(),
4,
"expected 4 delegate messages, got {texts:?}"
);
assert_eq!(texts[0], delegate_msg("S"));
assert_eq!(texts[1], delegate_msg("A"));
assert_eq!(texts[2], delegate_msg("B"));
assert_eq!(texts[3], delegate_msg("J"));
}
#[tokio::test]
async fn frontier_join_runs_only_after_both_branches_complete() {
let graph = make_fork_graph("S", "A", "B", "J");
let executor = make_frontier_executor();
let result = executor.execute(&graph, GraphState::new()).await.unwrap();
let join_count = result
.final_state
.conversation
.messages
.iter()
.filter(|m| m.text_content() == delegate_msg("J"))
.count();
assert_eq!(
join_count, 1,
"join must run exactly once (natural-join dedup)"
);
assert_eq!(result.iterations, 3);
}
#[tokio::test]
async fn frontier_single_branch_unchanged_from_legacy_behaviour() {
let mut graph = StateGraph::new("linear", "entry");
graph.add_node(GraphNode {
id: "entry".to_string(),
action: NodeAction::Delegate {
agent_kind: "S".to_string(),
task_template: "t".to_string(),
},
interrupt: false,
metadata: HashMap::new(),
});
graph.add_node(GraphNode {
id: "join".to_string(),
action: NodeAction::Delegate {
agent_kind: "J".to_string(),
task_template: "t".to_string(),
},
interrupt: false,
metadata: HashMap::new(),
});
graph.add_edge(GraphEdge {
from: "entry".to_string(),
to: "join".to_string(),
condition: Some(EdgeCondition::Always),
metadata: HashMap::new(),
});
graph.add_terminal("join");
let executor = make_frontier_executor();
let result = executor.execute(&graph, GraphState::new()).await.unwrap();
assert!(result.completed);
assert_eq!(result.terminal_node.as_deref(), Some("join"));
assert_eq!(result.iterations, 2);
let texts: Vec<String> = result
.final_state
.conversation
.messages
.iter()
.map(|m| m.text_content())
.collect();
assert_eq!(texts, vec![delegate_msg("S"), delegate_msg("J")]);
}
}