use std::collections::HashMap;
use std::sync::Arc;
use serde::{Deserialize, Serialize};
use oneai_core::error::{OneAIError, Result};
use oneai_core::traits::{InteractionGate, LlmProvider, PermissionResolver, Tool};
use oneai_core::{Conversation, InferenceRequest, Message};
use crate::config::{RetryPolicy, WorkflowConfig};
use crate::dag::WorkflowDag;
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[non_exhaustive]
pub enum StepStatus {
Pending,
Running,
Completed,
Failed,
Skipped,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StepResult {
pub step_id: String,
pub status: StepStatus,
pub output: Option<String>,
pub error: Option<String>,
pub retries_used: usize,
pub execution_time_ms: Option<u64>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WorkflowResult {
pub name: String,
pub step_results: HashMap<String, StepResult>,
pub variables: HashMap<String, String>,
pub success: bool,
pub total_time_ms: u64,
pub completed_levels: usize,
}
impl WorkflowResult {
pub fn get_step_result(&self, step_id: &str) -> Option<&StepResult> {
self.step_results.get(step_id)
}
pub fn completed_steps(&self) -> Vec<String> {
self.step_results
.iter()
.filter(|(_, r)| r.status == StepStatus::Completed)
.map(|(id, _)| id.clone())
.collect()
}
pub fn failed_steps(&self) -> Vec<String> {
self.step_results
.iter()
.filter(|(_, r)| r.status == StepStatus::Failed)
.map(|(id, _)| id.clone())
.collect()
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WorkflowContext {
pub variables: HashMap<String, String>,
pub step_outputs: HashMap<String, String>,
}
impl WorkflowContext {
pub fn from_config(config: &WorkflowConfig) -> Self {
Self {
variables: config.variables.clone(),
step_outputs: HashMap::new(),
}
}
pub fn set_variable(&mut self, key: impl Into<String>, value: impl Into<String>) {
self.variables.insert(key.into(), value.into());
}
pub fn get_variable(&self, key: &str) -> Option<&String> {
self.variables.get(key)
}
pub fn set_step_output(&mut self, step_id: impl Into<String>, output: impl Into<String>) {
self.step_outputs.insert(step_id.into(), output.into());
}
}
pub struct WorkflowExecutor {
tools: Arc<tokio::sync::RwLock<HashMap<String, Arc<dyn Tool>>>>,
interaction_gate: Arc<dyn InteractionGate>,
provider: Option<Arc<dyn LlmProvider>>,
permission_resolver: Option<Arc<dyn PermissionResolver>>,
}
impl WorkflowExecutor {
pub fn new(
tools: Arc<HashMap<String, Arc<dyn Tool>>>,
interaction_gate: Arc<dyn InteractionGate>,
) -> Self {
Self {
tools: Arc::new(tokio::sync::RwLock::new((*tools).clone())),
interaction_gate,
provider: None,
permission_resolver: None,
}
}
pub fn new_empty(interaction_gate: Arc<dyn InteractionGate>) -> Self {
Self {
tools: Arc::new(tokio::sync::RwLock::new(HashMap::new())),
interaction_gate,
provider: None,
permission_resolver: None,
}
}
pub fn with_provider(
tools: Arc<HashMap<String, Arc<dyn Tool>>>,
interaction_gate: Arc<dyn InteractionGate>,
provider: Arc<dyn LlmProvider>,
) -> Self {
Self {
tools: Arc::new(tokio::sync::RwLock::new((*tools).clone())),
interaction_gate,
provider: Some(provider),
permission_resolver: None,
}
}
pub fn with_permission_resolver(mut self, resolver: Arc<dyn PermissionResolver>) -> Self {
self.permission_resolver = Some(resolver);
self
}
pub fn set_provider(&mut self, provider: Arc<dyn LlmProvider>) {
self.provider = Some(provider);
}
pub async fn register_tool(&self, tool: Arc<dyn Tool>) {
self.tools
.write()
.await
.insert(tool.name().to_string(), tool);
}
pub fn tools_handle(&self) -> Arc<tokio::sync::RwLock<HashMap<String, Arc<dyn Tool>>>> {
self.tools.clone()
}
pub async fn execute(
&self,
dag: &WorkflowDag,
config: &WorkflowConfig,
) -> Result<WorkflowResult> {
let start_time = std::time::Instant::now();
let mut step_results: HashMap<String, StepResult> = HashMap::new();
let mut context = WorkflowContext::from_config(config);
let mut completed_levels = 0;
for id in dag.nodes.keys() {
step_results.insert(
id.clone(),
StepResult {
step_id: id.clone(),
status: StepStatus::Pending,
output: None,
error: None,
retries_used: 0,
execution_time_ms: None,
},
);
}
for level_ids in &dag.levels {
let _level_start = std::time::Instant::now();
let mut skip_level = false;
if !config.continue_on_failure {
for step_id in level_ids {
if let Some(node) = dag.get_node(step_id) {
for dep_id in &node.depends_on {
if let Some(dep_result) = step_results.get(dep_id) {
if dep_result.status == StepStatus::Failed
|| dep_result.status == StepStatus::Skipped
{
skip_level = true;
break;
}
}
}
}
if skip_level {
break;
}
}
}
if skip_level {
for step_id in level_ids {
if let Some(result) = step_results.get_mut(step_id) {
result.status = StepStatus::Skipped;
result.error = Some("Dependency failed — step skipped".to_string());
}
}
continue;
}
let mut level_futures = Vec::new();
for step_id in level_ids {
let node = dag.get_node(step_id).unwrap();
let step_config = node.step.clone();
let retry_policy = config.effective_retry_policy(step_id);
let timeout_secs = config.effective_timeout(step_id);
let tools = self.tools.clone();
let interaction_gate = self.interaction_gate.clone();
let context_snapshot = context.clone();
let provider = self.provider.clone();
let permission_resolver = self.permission_resolver.clone();
level_futures.push(tokio::spawn(async move {
execute_step(
step_config,
retry_policy,
timeout_secs,
tools,
interaction_gate,
context_snapshot,
provider,
permission_resolver,
)
.await
}));
}
for future in level_futures {
let result = future
.await
.map_err(|e| OneAIError::Workflow(format!("Task join error: {}", e)))?;
match result {
Ok(step_result) => {
step_results.insert(step_result.step_id.clone(), step_result.clone());
if step_result.status == StepStatus::Completed {
if let Some(output) = &step_result.output {
context.set_step_output(&step_result.step_id, output);
}
}
}
Err(e) => {
tracing::error!("Step execution error: {}", e);
}
}
}
completed_levels += 1;
}
let total_time_ms = start_time.elapsed().as_millis() as u64;
let success = step_results
.values()
.all(|r| r.status == StepStatus::Completed || r.status == StepStatus::Skipped);
Ok(WorkflowResult {
name: dag.name.clone(),
step_results,
variables: context.variables,
success,
total_time_ms,
completed_levels,
})
}
}
#[allow(clippy::too_many_arguments)]
async fn execute_step(
mut step: crate::config::StepConfig,
retry_policy: RetryPolicy,
timeout_secs: Option<u64>,
tools: Arc<tokio::sync::RwLock<HashMap<String, Arc<dyn Tool>>>>,
interaction_gate: Arc<dyn InteractionGate>,
context: WorkflowContext,
provider: Option<Arc<dyn LlmProvider>>,
permission_resolver: Option<Arc<dyn PermissionResolver>>,
) -> Result<StepResult> {
let start_time = std::time::Instant::now();
let step_id = step.id.clone();
let mut retries_used = 0;
let mut last_error: Option<String> = None;
let interpolated_prompt = step
.prompt
.as_ref()
.map(|p| interpolate_template(p, &context));
let interpolated_tool_args = step.tool_args.as_ref().map(|a| {
let json_str = serde_json::to_string(a).unwrap_or_default();
let interpolated = interpolate_template(&json_str, &context);
serde_json::from_str(&interpolated).unwrap_or_else(|_| a.clone())
});
if let (Some(resolver), Some(tool_name)) = (permission_resolver.as_ref(), step.tool.as_ref()) {
let check_args = interpolated_tool_args
.clone()
.or_else(|| step.tool_args.clone())
.unwrap_or_else(|| serde_json::json!({}));
match resolver.resolve(tool_name, &check_args) {
oneai_core::PermissionAction::Deny { reason } => {
tracing::warn!(
"Workflow step '{}' (tool '{}') denied by domain policy: {}",
step_id,
tool_name,
reason
);
return Ok(StepResult {
step_id,
status: StepStatus::Failed,
output: None,
error: Some(format!("Denied by domain policy: {}", reason)),
retries_used: 0,
execution_time_ms: Some(start_time.elapsed().as_millis() as u64),
});
}
oneai_core::PermissionAction::AutoApprove => {
tracing::info!(
"Workflow step '{}' (tool '{}') auto-approved by domain policy",
step_id,
tool_name
);
}
oneai_core::PermissionAction::RequireConfirmation => {
tracing::info!(
"Workflow step '{}' (tool '{}') requires confirmation by domain policy",
step_id,
tool_name
);
step.requires_approval = true;
}
oneai_core::PermissionAction::UseDefaultPermission { .. } => {
}
}
}
if step.requires_approval {
let approval_request = oneai_core::ApprovalRequest {
tool_name: step.id.clone(),
args: step.tool_args.clone().unwrap_or(serde_json::json!({})),
risk_level: oneai_core::RiskLevel::High,
permission_level: Some(oneai_core::PermissionLevel::Full),
justification: format!("Workflow step '{}' requires human approval", step.id),
};
let response = interaction_gate
.request(oneai_core::InteractionRequest::ToolApproval {
approval: approval_request,
})
.await?;
match response {
oneai_core::InteractionResponse::Abort { reason } => {
return Ok(StepResult {
step_id,
status: StepStatus::Failed,
output: None,
error: Some(format!("Approval denied: {}", reason)),
retries_used: 0,
execution_time_ms: Some(start_time.elapsed().as_millis() as u64),
});
}
oneai_core::InteractionResponse::Revise { feedback } => {
return Ok(StepResult {
step_id,
status: StepStatus::Failed,
output: None,
error: Some(format!("Approval denied: {}", feedback)),
retries_used: 0,
execution_time_ms: Some(start_time.elapsed().as_millis() as u64),
});
}
oneai_core::InteractionResponse::ProceedWith {
modification: oneai_core::InteractionModification::ReplaceToolArgs(_args),
} => {
}
oneai_core::InteractionResponse::ProceedWith { .. }
| oneai_core::InteractionResponse::Proceed
| oneai_core::InteractionResponse::Choose { .. } => {
}
_ => {}
}
}
for attempt in 0..=retry_policy.max_retries {
let exec_result: std::result::Result<oneai_core::ToolOutput, OneAIError> =
if let Some(tool_name) = &step.tool {
let tools_map = tools.read().await;
let tool = tools_map.get(tool_name);
if let Some(tool) = tool {
let args = match &interpolated_tool_args {
Some(a) => a.clone(),
None => step.tool_args.clone().unwrap_or(serde_json::json!({})),
};
let timeout = timeout_secs.unwrap_or(60);
let timeout_result = tokio::time::timeout(
std::time::Duration::from_secs(timeout),
tool.execute(args.clone()),
)
.await;
match timeout_result {
Ok(inner_result) => inner_result, Err(_) => Err(OneAIError::Timeout(format!(
"Step '{}' timed out after {} seconds",
step_id, timeout
))),
}
} else {
Err(OneAIError::Workflow(format!(
"Tool '{}' not found for step '{}'",
tool_name, step_id
)))
}
} else if let Some(provider) = &provider {
let prompt_text = match &interpolated_prompt {
Some(p) => p.clone(),
None => step.prompt.clone().unwrap_or_default(),
};
if prompt_text.is_empty() {
Ok(oneai_core::ToolOutput {
success: true,
content: String::new(),
error: None,
..Default::default()
})
} else {
let mut conversation = Conversation::new();
conversation.add_message(Message::system(
"You are executing a step in a deterministic workflow. Respond concisely.",
));
conversation.add_message(Message::user(prompt_text));
let request = InferenceRequest {
conversation,
tools: vec![], max_tokens: Some(2048),
temperature: Some(0.3),
top_p: None,
stop_sequences: vec![],
constrained_output: None,
thinking_budget: None,
metadata: HashMap::new(),
};
let response = provider.infer(request).await;
match response {
Ok(inference_response) => Ok(oneai_core::ToolOutput {
success: true,
content: inference_response.message.text_content(),
error: None,
..Default::default()
}),
Err(e) => Err(OneAIError::Provider(format!(
"LLM inference failed for step '{}': {}",
step_id, e
))),
}
}
} else {
let prompt_text = match &interpolated_prompt {
Some(p) => p.clone(),
None => step.prompt.clone().unwrap_or_default(),
};
Ok(oneai_core::ToolOutput {
success: true,
content: prompt_text,
error: None,
..Default::default()
})
};
match exec_result {
Ok(output) => {
if output.success {
return Ok(StepResult {
step_id,
status: StepStatus::Completed,
output: Some(output.content),
error: None,
retries_used,
execution_time_ms: Some(start_time.elapsed().as_millis() as u64),
});
} else {
last_error = Some(
output
.error
.unwrap_or_else(|| "Tool execution failed".to_string()),
);
}
}
Err(e) => {
last_error = Some(format!("Step error: {}", e));
}
}
retries_used = attempt;
if attempt < retry_policy.max_retries {
tokio::time::sleep(std::time::Duration::from_secs(
retry_policy.retry_delay_secs,
))
.await;
}
}
Ok(StepResult {
step_id,
status: StepStatus::Failed,
output: None,
error: last_error,
retries_used,
execution_time_ms: Some(start_time.elapsed().as_millis() as u64),
})
}
pub fn interpolate_template(template: &str, context: &WorkflowContext) -> String {
let mut result = template.to_string();
for (step_id, output) in &context.step_outputs {
let long_pattern = "{{".to_string() + step_id + "_output}}";
result = result.replace(&long_pattern, output);
}
for (step_id, output) in &context.step_outputs {
let short_pattern = "{{".to_string() + step_id + "}}";
result = result.replace(&short_pattern, output);
}
for (key, value) in &context.variables {
result = result.replace(&format!("{{{{{}}}}}", key), value);
}
result
}
#[cfg(test)]
mod interpolation_tests {
use super::*;
#[test]
fn test_interpolate_step_outputs() {
let mut context = WorkflowContext::from_config(&WorkflowConfig::new("test", vec![]));
context.set_step_output("step1", "result of step1");
context.set_step_output("step2", "result of step2");
let result = interpolate_template("{{step1}} then {{step2}}", &context);
assert_eq!(result, "result of step1 then result of step2");
let result = interpolate_template("{{step1_output}} and {{step2_output}}", &context);
assert_eq!(result, "result of step1 and result of step2");
}
#[test]
fn test_interpolate_variables() {
let mut context = WorkflowContext::from_config(&WorkflowConfig::new("test", vec![]));
context.set_variable("focus_area", "security");
context.set_variable("language", "Rust");
let result = interpolate_template("Focus on {{focus_area}} in {{language}}", &context);
assert_eq!(result, "Focus on security in Rust");
}
#[test]
fn test_interpolate_mixed() {
let mut context = WorkflowContext::from_config(&WorkflowConfig::new("test", vec![]));
context.set_step_output("read_diff", "auth.rs changes");
context.set_variable("focus", "security");
let result = interpolate_template(
"Review {{read_diff_output}} with focus on {{focus}}",
&context,
);
assert_eq!(result, "Review auth.rs changes with focus on security");
}
#[test]
fn test_interpolate_no_match() {
let context = WorkflowContext::from_config(&WorkflowConfig::new("test", vec![]));
let result = interpolate_template("No variables here {{unknown}}", &context);
assert_eq!(result, "No variables here {{unknown}}");
}
}
#[cfg(test)]
mod permission_resolver_tests {
use super::*;
use crate::compiler::compile;
use crate::config::StepConfig;
use oneai_core::error::Result;
use oneai_core::traits::{InteractionGate, PermissionResolver, Tool};
use oneai_core::{
InteractionRequest, InteractionResponse, PermissionAction, RiskLevel, ToolOutput,
};
struct EchoTool;
#[async_trait::async_trait]
impl Tool for EchoTool {
fn name(&self) -> &str {
"echo"
}
fn description(&self) -> &str {
"echo the value"
}
fn parameters_schema(&self) -> serde_json::Value {
serde_json::json!({"type": "object"})
}
fn risk_level(&self) -> RiskLevel {
RiskLevel::Low
}
async fn execute(&self, args: serde_json::Value) -> Result<ToolOutput> {
Ok(ToolOutput {
success: true,
content: args
.get("value")
.and_then(|v| v.as_str())
.unwrap_or("ok")
.to_string(),
error: None,
..Default::default()
})
}
}
struct NoopGate;
#[async_trait::async_trait]
impl InteractionGate for NoopGate {
async fn request(&self, _req: InteractionRequest) -> Result<InteractionResponse> {
Ok(InteractionResponse::Proceed)
}
fn enabled(&self, _point: oneai_core::InteractionPoint) -> bool {
false
}
}
struct DenyResolver {
denied: String,
}
impl PermissionResolver for DenyResolver {
fn resolve(&self, tool_name: &str, _args: &serde_json::Value) -> PermissionAction {
if tool_name == self.denied {
PermissionAction::Deny {
reason: "forbidden by domain policy".to_string(),
}
} else {
PermissionAction::AutoApprove
}
}
}
fn step(id: &str, tool: &str) -> StepConfig {
StepConfig {
id: id.to_string(),
description: String::new(),
depends_on: vec![],
tool: Some(tool.to_string()),
tool_args: Some(serde_json::json!({"value": "ran"})),
prompt: None,
requires_approval: false,
timeout_secs: None,
retry_policy: None,
metadata: HashMap::new(),
}
}
#[tokio::test]
async fn deny_by_default_short_circuits_a_tool_step() {
let config = WorkflowConfig::new("deny-test", vec![step("s1", "echo")]);
let dag = compile(&config);
let executor = WorkflowExecutor::new_empty(Arc::new(NoopGate)).with_permission_resolver(
Arc::new(DenyResolver {
denied: "echo".to_string(),
}),
);
executor.register_tool(Arc::new(EchoTool)).await;
let result = executor.execute(&dag, &config).await.unwrap();
let step_result = result.step_results.get("s1").expect("step result present");
assert_eq!(step_result.status, StepStatus::Failed);
let err = step_result.error.as_ref().expect("error present");
assert!(err.contains("Denied by domain policy"));
assert!(err.contains("forbidden by domain policy"));
}
#[tokio::test]
async fn auto_approve_allows_a_tool_step_to_run() {
let config = WorkflowConfig::new("allow-test", vec![step("s1", "echo")]);
let dag = compile(&config);
let executor = WorkflowExecutor::new_empty(Arc::new(NoopGate)).with_permission_resolver(
Arc::new(DenyResolver {
denied: "not_echo".to_string(),
}),
);
executor.register_tool(Arc::new(EchoTool)).await;
let result = executor.execute(&dag, &config).await.unwrap();
let step_result = result.step_results.get("s1").expect("step result present");
assert_eq!(step_result.status, StepStatus::Completed);
assert_eq!(step_result.output.as_deref(), Some("ran"));
}
#[tokio::test]
async fn no_resolver_preserves_legacy_behaviour() {
let config = WorkflowConfig::new("legacy-test", vec![step("s1", "echo")]);
let dag = compile(&config);
let executor = WorkflowExecutor::new_empty(Arc::new(NoopGate));
executor.register_tool(Arc::new(EchoTool)).await;
let result = executor.execute(&dag, &config).await.unwrap();
let step_result = result.step_results.get("s1").expect("step result present");
assert_eq!(step_result.status, StepStatus::Completed);
assert_eq!(step_result.output.as_deref(), Some("ran"));
}
}