use super::error::ExecutionError;
use super::execution_state::{ExecutionState, StepState};
use super::ids::{CallableType, ExecutionId, ParentLink, StepId, StepSource, StepType, TenantId};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::time::Instant;
#[derive(Debug)]
pub struct Execution {
pub id: ExecutionId,
pub tenant_id: Option<TenantId>,
pub state: ExecutionState,
pub parent: Option<ParentLink>,
pub steps: HashMap<StepId, Step>,
pub step_order: Vec<StepId>,
pub schema_version: Option<String>,
pub started_at: Option<Instant>,
pub ended_at: Option<Instant>,
pub output: Option<String>,
pub error: Option<ExecutionError>,
}
impl Execution {
pub fn new() -> Self {
Self {
id: ExecutionId::new(),
tenant_id: None,
state: ExecutionState::Created,
parent: None,
steps: HashMap::new(),
step_order: Vec::new(),
schema_version: None,
started_at: None,
ended_at: None,
output: None,
error: None,
}
}
pub fn with_id(id: ExecutionId) -> Self {
Self {
id,
tenant_id: None,
state: ExecutionState::Created,
parent: None,
steps: HashMap::new(),
step_order: Vec::new(),
schema_version: None,
started_at: None,
ended_at: None,
output: None,
error: None,
}
}
pub fn with_tenant(tenant_id: TenantId) -> Self {
Self {
id: ExecutionId::new(),
tenant_id: Some(tenant_id),
state: ExecutionState::Created,
parent: None,
steps: HashMap::new(),
step_order: Vec::new(),
schema_version: None,
started_at: None,
ended_at: None,
output: None,
error: None,
}
}
pub fn child(&self) -> Self {
let mut child = Self::new();
child.parent = Some(ParentLink::execution(self.id.clone()));
child.tenant_id = self.tenant_id.clone(); child
}
pub fn with_schema_version(mut self, version: impl Into<String>) -> Self {
self.schema_version = Some(version.into());
self
}
pub fn get_step(&self, id: &StepId) -> Option<&Step> {
self.steps.get(id)
}
pub fn get_step_mut(&mut self, id: &StepId) -> Option<&mut Step> {
self.steps.get_mut(id)
}
pub fn add_step(&mut self, step: Step) {
self.step_order.push(step.id.clone());
self.steps.insert(step.id.clone(), step);
}
pub fn duration_ms(&self) -> Option<u64> {
match (self.started_at, self.ended_at) {
(Some(start), Some(end)) => Some(end.duration_since(start).as_millis() as u64),
(Some(start), None) => Some(start.elapsed().as_millis() as u64),
_ => None,
}
}
pub fn is_terminal(&self) -> bool {
self.state.is_terminal()
}
}
impl Default for Execution {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct Step {
#[serde(rename = "stepId")]
pub id: StepId,
#[serde(rename = "parentStepId")]
pub parent_step_id: Option<StepId>,
pub step_type: StepType,
pub name: String,
pub state: StepState,
pub input: Option<String>,
pub output: Option<String>,
pub error: Option<ExecutionError>,
pub duration_ms: Option<u64>,
pub started_at: Option<i64>,
pub ended_at: Option<i64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub source: Option<StepSource>,
#[serde(skip_serializing_if = "Option::is_none")]
pub callable_id: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub callable_type: Option<CallableType>,
}
impl Step {
pub fn new(step_type: StepType, name: impl Into<String>) -> Self {
Self {
id: StepId::new(),
parent_step_id: None,
step_type,
name: name.into(),
state: StepState::Pending,
input: None,
output: None,
error: None,
duration_ms: None,
started_at: None,
ended_at: None,
source: None,
callable_id: None,
callable_type: None,
}
}
pub fn nested(parent_id: &StepId, step_type: StepType, name: impl Into<String>) -> Self {
let mut step = Self::new(step_type, name);
step.parent_step_id = Some(parent_id.clone());
step
}
pub fn with_input(mut self, input: impl Into<String>) -> Self {
self.input = Some(input.into());
self
}
pub fn with_source(mut self, source: StepSource) -> Self {
self.source = Some(source);
self
}
pub fn with_callable(
mut self,
callable_id: impl Into<String>,
callable_type: CallableType,
) -> Self {
self.callable_id = Some(callable_id.into());
self.callable_type = Some(callable_type);
self
}
}
#[cfg(test)]
mod tests {
use super::super::execution_state::WaitReason;
use super::super::ids::StepSourceType;
use super::*;
#[test]
fn test_execution_new() {
let exec = Execution::new();
assert!(exec.id.as_str().starts_with("exec_"));
assert_eq!(exec.state, ExecutionState::Created);
assert!(exec.parent.is_none());
assert!(exec.steps.is_empty());
assert!(exec.step_order.is_empty());
assert!(exec.schema_version.is_none());
assert!(exec.started_at.is_none());
assert!(exec.ended_at.is_none());
assert!(exec.output.is_none());
assert!(exec.error.is_none());
}
#[test]
fn test_execution_with_id() {
let id = ExecutionId::from_string("exec_custom_id");
let exec = Execution::with_id(id.clone());
assert_eq!(exec.id.as_str(), "exec_custom_id");
assert_eq!(exec.state, ExecutionState::Created);
}
#[test]
fn test_execution_child() {
let parent = Execution::new();
let child = parent.child();
assert!(child.parent.is_some());
let parent_link = child.parent.unwrap();
assert_eq!(parent_link.parent_id, parent.id.as_str());
}
#[test]
fn test_execution_with_schema_version() {
let exec = Execution::new().with_schema_version("v1.0.0");
assert_eq!(exec.schema_version, Some("v1.0.0".to_string()));
}
#[test]
fn test_execution_with_schema_version_owned_string() {
let exec = Execution::new().with_schema_version(String::from("v2.0.0"));
assert_eq!(exec.schema_version, Some("v2.0.0".to_string()));
}
#[test]
fn test_execution_add_step() {
let mut exec = Execution::new();
let step = Step::new(StepType::LlmNode, "test_step");
let step_id = step.id.clone();
exec.add_step(step);
assert_eq!(exec.steps.len(), 1);
assert_eq!(exec.step_order.len(), 1);
assert!(exec.steps.contains_key(&step_id));
}
#[test]
fn test_execution_get_step() {
let mut exec = Execution::new();
let step = Step::new(StepType::ToolNode, "get_step_test");
let step_id = step.id.clone();
exec.add_step(step);
let retrieved = exec.get_step(&step_id);
assert!(retrieved.is_some());
assert_eq!(retrieved.unwrap().name, "get_step_test");
}
#[test]
fn test_execution_get_step_not_found() {
let exec = Execution::new();
let nonexistent_id = StepId::from_string("step_nonexistent");
assert!(exec.get_step(&nonexistent_id).is_none());
}
#[test]
fn test_execution_get_step_mut() {
let mut exec = Execution::new();
let step = Step::new(StepType::FunctionNode, "mutable_step");
let step_id = step.id.clone();
exec.add_step(step);
let step_mut = exec.get_step_mut(&step_id);
assert!(step_mut.is_some());
step_mut.unwrap().output = Some("modified".to_string());
let step = exec.get_step(&step_id).unwrap();
assert_eq!(step.output, Some("modified".to_string()));
}
#[test]
fn test_execution_step_order_preserved() {
let mut exec = Execution::new();
let step1 = Step::new(StepType::LlmNode, "step1");
let step2 = Step::new(StepType::ToolNode, "step2");
let step3 = Step::new(StepType::FunctionNode, "step3");
let id1 = step1.id.clone();
let id2 = step2.id.clone();
let id3 = step3.id.clone();
exec.add_step(step1);
exec.add_step(step2);
exec.add_step(step3);
assert_eq!(exec.step_order[0], id1);
assert_eq!(exec.step_order[1], id2);
assert_eq!(exec.step_order[2], id3);
}
#[test]
fn test_execution_duration_ms_not_started() {
let exec = Execution::new();
assert!(exec.duration_ms().is_none());
}
#[test]
fn test_execution_duration_ms_started_not_ended() {
let mut exec = Execution::new();
exec.started_at = Some(Instant::now());
std::thread::sleep(std::time::Duration::from_millis(10));
let duration = exec.duration_ms();
assert!(duration.is_some());
assert!(duration.unwrap() >= 10);
}
#[test]
fn test_execution_duration_ms_completed() {
let mut exec = Execution::new();
let start = Instant::now();
std::thread::sleep(std::time::Duration::from_millis(20));
let end = Instant::now();
exec.started_at = Some(start);
exec.ended_at = Some(end);
let duration = exec.duration_ms();
assert!(duration.is_some());
assert!(duration.unwrap() >= 20);
}
#[test]
fn test_execution_is_terminal_created() {
let exec = Execution::new();
assert!(!exec.is_terminal());
}
#[test]
fn test_execution_is_terminal_running() {
let mut exec = Execution::new();
exec.state = ExecutionState::Running;
assert!(!exec.is_terminal());
}
#[test]
fn test_execution_is_terminal_completed() {
let mut exec = Execution::new();
exec.state = ExecutionState::Completed;
assert!(exec.is_terminal());
}
#[test]
fn test_execution_is_terminal_failed() {
let mut exec = Execution::new();
exec.state = ExecutionState::Failed;
assert!(exec.is_terminal());
}
#[test]
fn test_execution_is_terminal_cancelled() {
let mut exec = Execution::new();
exec.state = ExecutionState::Cancelled;
assert!(exec.is_terminal());
}
#[test]
fn test_execution_is_terminal_paused() {
let mut exec = Execution::new();
exec.state = ExecutionState::Paused;
assert!(!exec.is_terminal());
}
#[test]
fn test_execution_is_terminal_waiting() {
let mut exec = Execution::new();
exec.state = ExecutionState::Waiting(WaitReason::Approval);
assert!(!exec.is_terminal());
}
#[test]
fn test_execution_default() {
let exec: Execution = Default::default();
assert!(exec.id.as_str().starts_with("exec_"));
assert_eq!(exec.state, ExecutionState::Created);
}
#[test]
fn test_step_new() {
let step = Step::new(StepType::LlmNode, "test_step");
assert!(step.id.as_str().starts_with("step_"));
assert_eq!(step.step_type, StepType::LlmNode);
assert_eq!(step.name, "test_step");
assert_eq!(step.state, StepState::Pending);
assert!(step.parent_step_id.is_none());
assert!(step.input.is_none());
assert!(step.output.is_none());
assert!(step.error.is_none());
assert!(step.duration_ms.is_none());
}
#[test]
fn test_step_new_all_types() {
let types = vec![
StepType::LlmNode,
StepType::GraphNode,
StepType::ToolNode,
StepType::FunctionNode,
StepType::RouterNode,
StepType::BranchNode,
StepType::LoopNode,
];
for step_type in types {
let step = Step::new(step_type.clone(), "test");
assert_eq!(step.step_type, step_type);
}
}
#[test]
fn test_step_nested() {
let parent_id = StepId::from_string("step_parent");
let step = Step::nested(&parent_id, StepType::ToolNode, "nested_step");
assert!(step.parent_step_id.is_some());
assert_eq!(step.parent_step_id.unwrap().as_str(), "step_parent");
assert_eq!(step.step_type, StepType::ToolNode);
assert_eq!(step.name, "nested_step");
}
#[test]
fn test_step_with_input() {
let step = Step::new(StepType::LlmNode, "input_step").with_input("Hello, AI!");
assert!(step.input.is_some());
assert_eq!(step.input.unwrap(), "Hello, AI!");
}
#[test]
fn test_step_with_input_owned_string() {
let step =
Step::new(StepType::LlmNode, "input_step").with_input(String::from("Owned input"));
assert!(step.input.is_some());
assert_eq!(step.input.unwrap(), "Owned input");
}
#[test]
fn test_step_clone() {
let step = Step::new(StepType::FunctionNode, "cloneable").with_input("input data");
let cloned = step.clone();
assert_eq!(step.id.as_str(), cloned.id.as_str());
assert_eq!(step.name, cloned.name);
assert_eq!(step.input, cloned.input);
}
#[test]
fn test_step_serde() {
let step = Step::new(StepType::GraphNode, "serializable").with_input("input");
let json = serde_json::to_string(&step).unwrap();
let parsed: Step = serde_json::from_str(&json).unwrap();
assert_eq!(step.name, parsed.name);
assert_eq!(step.step_type, parsed.step_type);
assert_eq!(step.input, parsed.input);
}
#[test]
fn test_step_serde_field_names() {
let step = Step::new(StepType::LlmNode, "test_step").with_input("test input");
let json = serde_json::to_string(&step).unwrap();
assert!(json.contains("\"stepId\""), "Should have stepId field");
assert!(json.contains("\"stepType\""), "Should have stepType field");
assert!(json.contains("\"state\""), "Should have state field");
assert!(
json.contains("\"durationMs\""),
"Should have durationMs field"
);
assert!(
json.contains("\"startedAt\""),
"Should have startedAt field"
);
assert!(json.contains("\"endedAt\""), "Should have endedAt field");
assert!(
!json.contains("\"step_id\""),
"Should NOT have step_id field"
);
assert!(
!json.contains("\"step_type\""),
"Should NOT have step_type field"
);
assert!(
!json.contains("\"duration_ms\""),
"Should NOT have duration_ms field"
);
assert!(
!json.contains("\"started_at\""),
"Should NOT have started_at field"
);
assert!(
!json.contains("\"ended_at\""),
"Should NOT have ended_at field"
);
}
#[test]
fn test_step_timestamps() {
let mut step = Step::new(StepType::ToolNode, "timestamped");
let now = chrono::Utc::now().timestamp_millis();
step.started_at = Some(now);
step.ended_at = Some(now + 1000);
step.duration_ms = Some(1000);
assert_eq!(step.started_at, Some(now));
assert_eq!(step.ended_at, Some(now + 1000));
assert_eq!(step.duration_ms, Some(1000));
}
#[test]
fn test_step_state_modifications() {
let mut step = Step::new(StepType::LlmNode, "state_step");
assert_eq!(step.state, StepState::Pending);
step.state = StepState::Running;
assert_eq!(step.state, StepState::Running);
step.state = StepState::Completed;
step.output = Some("Result".to_string());
assert_eq!(step.state, StepState::Completed);
assert_eq!(step.output, Some("Result".to_string()));
}
#[test]
fn test_step_error_handling() {
let mut step = Step::new(StepType::ToolNode, "error_step");
step.state = StepState::Failed;
step.error = Some(ExecutionError::kernel_internal("Test error"));
assert!(step.error.is_some());
}
#[test]
fn test_execution_with_multiple_steps() {
let mut exec = Execution::new();
exec.state = ExecutionState::Running;
exec.started_at = Some(Instant::now());
for i in 0..5 {
let step = Step::new(StepType::FunctionNode, format!("step_{}", i))
.with_input(format!("input_{}", i));
exec.add_step(step);
}
assert_eq!(exec.steps.len(), 5);
assert_eq!(exec.step_order.len(), 5);
for step_id in &exec.step_order {
assert!(exec.get_step(step_id).is_some());
}
}
#[test]
fn test_nested_execution_structure() {
let root = Execution::new();
let child1 = root.child();
let child2 = root.child();
assert!(child1.parent.is_some());
assert!(child2.parent.is_some());
assert_eq!(
child1.parent.as_ref().unwrap().parent_id,
child2.parent.as_ref().unwrap().parent_id
);
assert_ne!(child1.id.as_str(), child2.id.as_str());
}
#[test]
fn test_nested_steps_structure() {
let parent_step = Step::new(StepType::GraphNode, "parent");
let parent_id = parent_step.id.clone();
let child1 = Step::nested(&parent_id, StepType::LlmNode, "child1");
let child2 = Step::nested(&parent_id, StepType::ToolNode, "child2");
assert_eq!(
child1.parent_step_id.as_ref().unwrap().as_str(),
parent_id.as_str()
);
assert_eq!(
child2.parent_step_id.as_ref().unwrap().as_str(),
parent_id.as_str()
);
}
#[test]
fn test_step_with_callable() {
let step = Step::new(StepType::GraphNode, "Research Agent")
.with_callable("research-agent-v2", CallableType::Agent);
assert!(step.callable_id.is_some());
assert_eq!(step.callable_id.unwrap(), "research-agent-v2");
assert!(step.callable_type.is_some());
assert_eq!(step.callable_type.unwrap(), CallableType::Agent);
}
#[test]
fn test_step_callable_serde() {
let step = Step::new(StepType::GraphNode, "Chat Handler")
.with_callable("chat-handler", CallableType::Chat);
let json = serde_json::to_string(&step).unwrap();
let parsed: Step = serde_json::from_str(&json).unwrap();
assert_eq!(parsed.callable_id, Some("chat-handler".to_string()));
assert_eq!(parsed.callable_type, Some(CallableType::Chat));
assert!(
json.contains("\"callableId\""),
"Should have callableId field"
);
assert!(
json.contains("\"callableType\""),
"Should have callableType field"
);
}
#[test]
fn test_step_callable_none_not_serialized() {
let step = Step::new(StepType::LlmNode, "No Callable Info");
let json = serde_json::to_string(&step).unwrap();
assert!(
!json.contains("callableId"),
"Should NOT serialize None callableId"
);
assert!(
!json.contains("callableType"),
"Should NOT serialize None callableType"
);
}
#[test]
fn test_callable_type_display() {
assert_eq!(format!("{}", CallableType::Completion), "completion");
assert_eq!(format!("{}", CallableType::Chat), "chat");
assert_eq!(format!("{}", CallableType::Agent), "agent");
assert_eq!(format!("{}", CallableType::Workflow), "workflow");
assert_eq!(format!("{}", CallableType::Background), "background");
assert_eq!(format!("{}", CallableType::Composite), "composite");
assert_eq!(format!("{}", CallableType::Tool), "tool");
assert_eq!(format!("{}", CallableType::Custom), "custom");
}
#[test]
fn test_callable_type_default() {
let default_type = CallableType::default();
assert_eq!(default_type, CallableType::Agent);
}
#[test]
fn test_callable_type_serde_all_variants() {
let variants = vec![
CallableType::Completion,
CallableType::Chat,
CallableType::Agent,
CallableType::Workflow,
CallableType::Background,
CallableType::Composite,
CallableType::Tool,
CallableType::Custom,
];
for variant in variants {
let json = serde_json::to_string(&variant).unwrap();
let parsed: CallableType = serde_json::from_str(&json).unwrap();
assert_eq!(parsed, variant);
}
}
#[test]
fn test_step_chained_builders() {
let _parent_id = StepId::from_string("step_parent");
let step = Step::new(StepType::GraphNode, "Full Step")
.with_input("User request")
.with_callable("research-agent", CallableType::Agent)
.with_source(StepSource {
source_type: StepSourceType::Discovered,
triggered_by: Some("step_123".to_string()),
reason: Some("LLM suggested sub-task".to_string()),
depth: Some(1),
spawn_mode: None,
});
assert_eq!(step.name, "Full Step");
assert_eq!(step.input, Some("User request".to_string()));
assert_eq!(step.callable_id, Some("research-agent".to_string()));
assert_eq!(step.callable_type, Some(CallableType::Agent));
assert!(step.source.is_some());
assert_eq!(
step.source.as_ref().unwrap().source_type,
StepSourceType::Discovered
);
}
}