enya-agent 0.1.0

package/crates/enact-core/src/kernel/execution_model.rs

@@ -0,0 +1,831 @@
+//! Execution Model - Runtime instance types for graph/agent execution
+//!
+//! This module defines the core model types:
+//! - `Execution`: One run of a blueprint (graph, agent, workflow)
+//! - `Step`: A distinct action within an execution
+//!
+//! These are the **model types** that hold execution state.
+//! The **engine** (ExecutionKernel) uses these models to manage execution.
+//!
+//! ## ⚠️ CRITICAL INVARIANT: Data + Bookkeeping Only
+//!
+//! **This module MUST NOT contain execution logic or side effects.**
+//!
+//! This module is strictly for:
+//! - Data structures (Execution, Step)
+//! - Bookkeeping methods (getters, setters, state queries)
+//!
+//! This module MUST NEVER:
+//! - Execute logic (no business logic, no decision-making)
+//! - Call providers (no external service calls)
+//! - Embed policy checks (policy belongs in kernel::enforcement)
+//! - Perform side effects (no I/O, no mutations beyond self)
+//!
+//! If you find yourself adding execution logic here, it belongs in:
+//! - `kernel::reducer` (for state transitions)
+//! - `kernel::enforcement` (for policy checks)
+//! - `kernel::execution_kernel` (for orchestration)
+//!
+//! ## Error Handling (feat-02)
+//!
+//! All errors are represented using `ExecutionError` which provides:
+//! - Deterministic retry policies
+//! - Structured error categories
+//! - HTTP status code mapping
+//! - Idempotency tracking
+
+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;
+
+/// Execution - one run of a blueprint
+#[derive(Debug)]
+pub struct Execution {
+    /// Unique execution ID
+    pub id: ExecutionId,
+    /// Tenant ID (REQUIRED for audit trail and multi-tenant isolation)
+    /// This is set from TenantContext when the kernel is created.
+    pub tenant_id: Option<TenantId>,
+    /// Current state in the lifecycle
+    pub state: ExecutionState,
+    /// Parent execution (if this is a sub-execution)
+    pub parent: Option<ParentLink>,
+    /// All steps in this execution
+    pub steps: HashMap<StepId, Step>,
+    /// Ordered list of step IDs (for replay)
+    pub step_order: Vec<StepId>,
+    /// Schema version hash (for replay validation)
+    pub schema_version: Option<String>,
+    /// Start time
+    pub started_at: Option<Instant>,
+    /// End time
+    pub ended_at: Option<Instant>,
+    /// Final output (if completed)
+    pub output: Option<String>,
+    /// Structured error (if failed) - see feat-02 Error Taxonomy
+    pub error: Option<ExecutionError>,
+}
+
+impl Execution {
+    /// Create a new 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,
+        }
+    }
+
+    /// Create a new execution with a specific ID
+    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,
+        }
+    }
+
+    /// Create a new execution with tenant ID
+    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,
+        }
+    }
+
+    /// Create a child execution
+    /// Inherits tenant_id from parent for multi-tenant isolation
+    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(); // Inherit tenant from parent
+        child
+    }
+
+    /// Set schema version for replay validation
+    pub fn with_schema_version(mut self, version: impl Into<String>) -> Self {
+        self.schema_version = Some(version.into());
+        self
+    }
+
+    /// Get a step by ID
+    pub fn get_step(&self, id: &StepId) -> Option<&Step> {
+        self.steps.get(id)
+    }
+
+    /// Get a mutable step by ID
+    pub fn get_step_mut(&mut self, id: &StepId) -> Option<&mut Step> {
+        self.steps.get_mut(id)
+    }
+
+    /// Add a new step
+    pub fn add_step(&mut self, step: Step) {
+        self.step_order.push(step.id.clone());
+        self.steps.insert(step.id.clone(), step);
+    }
+
+    /// Get execution duration in milliseconds
+    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,
+        }
+    }
+
+    /// Check if execution is in a terminal state
+    pub fn is_terminal(&self) -> bool {
+        self.state.is_terminal()
+    }
+}
+
+impl Default for Execution {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+/// Step - a distinct action within an execution
+///
+/// @see packages/enact-schemas/src/streaming.schemas.ts - stepEventDataSchema
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct Step {
+    /// Unique step ID
+    #[serde(rename = "stepId")]
+    pub id: StepId,
+    /// Parent step (for nested operations)
+    #[serde(rename = "parentStepId")]
+    pub parent_step_id: Option<StepId>,
+    /// Step type
+    pub step_type: StepType,
+    /// Step name/label
+    pub name: String,
+    /// Current state (matches `state` in TypeScript schema)
+    pub state: StepState,
+    /// Input to this step
+    pub input: Option<String>,
+    /// Output from this step
+    pub output: Option<String>,
+    /// Structured error (if failed) - see feat-02 Error Taxonomy
+    pub error: Option<ExecutionError>,
+    /// Duration in milliseconds
+    pub duration_ms: Option<u64>,
+    /// Timestamp when step started (unix millis)
+    pub started_at: Option<i64>,
+    /// Timestamp when step ended (unix millis)
+    pub ended_at: Option<i64>,
+    /// Source/origin of this step (how/why it was created)
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub source: Option<StepSource>,
+    /// Callable ID (stable identifier for billing/traceability)
+    /// Unlike `name` which can change, this provides a stable reference.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub callable_id: Option<String>,
+    /// Callable type (for billing differentiation and audit trails)
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub callable_type: Option<CallableType>,
+}
+
+impl Step {
+    /// Create a new 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,
+        }
+    }
+
+    /// Create a nested step under a parent
+    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
+    }
+
+    /// Set input
+    pub fn with_input(mut self, input: impl Into<String>) -> Self {
+        self.input = Some(input.into());
+        self
+    }
+
+    /// Set source/origin of this step
+    pub fn with_source(mut self, source: StepSource) -> Self {
+        self.source = Some(source);
+        self
+    }
+
+    /// Set callable info (for billing/traceability)
+    ///
+    /// Unlike `name` which can change, callable_id provides a stable reference
+    /// for cost attribution and audit trails.
+    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::*;
+
+    // =========================================================================
+    // Execution Tests
+    // =========================================================================
+
+    #[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);
+    }
+
+    // =========================================================================
+    // Step Tests
+    // =========================================================================
+
+    #[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() {
+        // Verify JSON field names match TypeScript schema (camelCase)
+        let step = Step::new(StepType::LlmNode, "test_step").with_input("test input");
+
+        let json = serde_json::to_string(&step).unwrap();
+
+        // Check that camelCase field names are used
+        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");
+
+        // Verify no snake_case field names
+        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());
+    }
+
+    // =========================================================================
+    // Integration Tests
+    // =========================================================================
+
+    #[test]
+    fn test_execution_with_multiple_steps() {
+        let mut exec = Execution::new();
+        exec.state = ExecutionState::Running;
+        exec.started_at = Some(Instant::now());
+
+        // Add multiple steps
+        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);
+
+        // Verify all steps are accessible
+        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();
+
+        // Both children should have the same parent
+        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
+        );
+
+        // But different IDs
+        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()
+        );
+    }
+
+    // =========================================================================
+    // Callable Tracking Tests (for billing/traceability)
+    // =========================================================================
+
+    #[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));
+
+        // Verify camelCase field names
+        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();
+
+        // Optional None fields should not appear in JSON (skip_serializing_if)
+        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() {
+        // Test that all builders can be chained
+        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
+        );
+    }
+}