fraiseql_core/federation/

saga_executor.rs

//! Saga Forward Phase Executor
//!
//! Executes saga steps sequentially during the forward phase, implementing
//! the core saga pattern for distributed transactions across subgraphs.
//!
//! # Architecture
//!
//! The forward phase executor:
//! - Loads sagas from persistent storage
//! - Executes steps in strict sequential order (1 → 2 → 3)
//! - Pre-fetches @requires fields before each step
//! - Captures and persists results for chaining
//! - Tracks execution state for monitoring and recovery
//! - Terminates on first failure and triggers compensation
//!
//! # Execution Flow
//!
//! ```text
//! Load Saga from Store
//!    ↓
//! For Each Step (1..N):
//!    ├─ Validate step is Pending
//!    ├─ Pre-fetch @requires fields from other subgraphs
//!    ├─ Transition step to Executing
//!    ├─ Execute mutation via MutationExecutor
//!    │  (with augmented entity data)
//!    ├─ Capture result data
//!    ├─ Persist step result to store
//!    ├─ Transition step to Completed
//!    └─ Continue to next step
//!       OR on failure: Break and transition to Failed state
//!
//! Update Saga State:
//!    ├─ If all completed: Saga → Completed
//!    └─ If any failed: Saga → Failed (trigger compensation)
//! ```
//!
//! # @requires Field Fetching
//!
//! Each step may have @requires fields that must be present before mutation execution.
//! These fields are fetched from their owning subgraphs before step execution:
//!
//! ```text
//! Step Definition:
//!   mutation: "updateOrder"
//!   @requires: ["product.price", "user.email"]
//!
//! Pre-Execution:
//!   1. Identify @requires fields
//!   2. Fetch from owning subgraphs
//!   3. Augment entity data with fetched fields
//!   4. Execute mutation with complete entity
//! ```
//!
//! # Example
//!
//! ```ignore
//! let executor = SagaExecutor::new();
//!
//! // Execute a single step
//! let result = executor.execute_step(
//!     saga_id,
//!     1,
//!     "createOrder",
//!     &json!({"customerId": "c123", "total": 100.0}),
//!     "orders-service"
//! ).await?;
//!
//! if result.success {
//!     println!("Step 1 created order: {:?}", result.data);
//! } else {
//!     println!("Step 1 failed: {}", result.error.unwrap());
//! }
//! ```

use std::{sync::Arc, time::Instant};

use tracing::{debug, info, warn};
use uuid::Uuid;

use crate::federation::saga_store::{PostgresSagaStore, Result as SagaStoreResult, StepState};

/// Represents a step result from execution
///
/// Contains the outcome of executing a single saga step, including:
/// - Whether execution succeeded or failed
/// - Result data if successful (entity with key fields and updated values)
/// - Error details if failed
/// - Execution metrics (duration)
///
/// The result data is stored and available for subsequent steps to reference
/// via result chaining.
#[derive(Debug, Clone)]
pub struct StepExecutionResult {
    /// Step number that executed (1-indexed)
    pub step_number: u32,
    /// Whether step succeeded (true) or failed (false)
    pub success:     bool,
    /// Result data if successful
    ///
    /// Contains:
    /// - `__typename`: Entity type
    /// - Key fields (id, etc.)
    /// - Mutation output fields
    /// - Timestamps
    pub data:        Option<serde_json::Value>,
    /// Error message if failed
    ///
    /// Includes:
    /// - Error type (timeout, network, mutation failed, etc.)
    /// - Subgraph context
    /// - Suggestion for resolution
    pub error:       Option<String>,
    /// Execution duration in milliseconds
    ///
    /// Measured from step start to completion (or failure)
    /// Useful for performance monitoring
    pub duration_ms: u64,
}

/// Saga forward phase executor
///
/// Executes saga steps sequentially during the forward phase.
/// Coordinates with saga store to persist state and handle failures.
pub struct SagaExecutor {
    /// Saga store for loading/saving saga state
    /// Optional to support testing without database
    store: Option<Arc<PostgresSagaStore>>,
}

impl SagaExecutor {
    /// Create a new saga executor without a saga store
    ///
    /// This is suitable for testing. For production, use `with_store()`.
    pub fn new() -> Self {
        Self { store: None }
    }

    /// Create a new saga executor with a saga store
    ///
    /// This enables persistence of saga state and recovery from failures.
    #[must_use]
    pub fn with_store(store: Arc<PostgresSagaStore>) -> Self {
        Self { store: Some(store) }
    }

    /// Check if executor has a saga store configured
    #[must_use]
    pub fn has_store(&self) -> bool {
        self.store.is_some()
    }

    /// Execute a single saga step
    ///
    /// Executes a single mutation step within a saga, handling:
    /// - Step state validation (Pending → Executing → Completed)
    /// - @requires field pre-fetching from owning subgraphs
    /// - Entity data augmentation with required fields
    /// - Mutation execution via MutationExecutor
    /// - Result capture and persistence
    ///
    /// # Arguments
    ///
    /// * `saga_id` - ID of saga being executed
    /// * `step_number` - Step number to execute (1-indexed, 1 = first step)
    /// * `mutation_name` - GraphQL mutation operation name
    /// * `variables` - Input variables for the mutation (JSON value)
    /// * `subgraph` - Target subgraph name (must exist in federation)
    ///
    /// # Returns
    ///
    /// `StepExecutionResult` with:
    /// - `success`: true if step executed successfully
    /// - `data`: Result entity data if successful
    /// - `error`: Error description if failed
    /// - `duration_ms`: Execution time for monitoring
    ///
    /// # Errors
    ///
    /// Returns `SagaStoreError` if:
    /// - Saga not found in store
    /// - Step already executed (not in Pending state)
    /// - Subgraph unavailable
    /// - Mutation execution fails
    /// - @requires fields cannot be fetched
    ///
    /// # Example
    ///
    /// ```ignore
    /// let executor = SagaExecutor::new();
    /// let result = executor.execute_step(
    ///     saga_id,
    ///     1,
    ///     "createOrder",
    ///     &json!({"customerId": "c123", "total": 100.0}),
    ///     "orders-service"
    /// ).await?;
    ///
    /// if result.success {
    ///     println!("Order created with data: {:?}", result.data);
    /// } else {
    ///     eprintln!("Step failed: {}", result.error.unwrap());
    ///     // Compensation will be triggered by coordinator
    /// }
    /// ```
    pub async fn execute_step(
        &self,
        saga_id: Uuid,
        step_number: u32,
        mutation_name: &str,
        _variables: &serde_json::Value,
        subgraph: &str,
    ) -> SagaStoreResult<StepExecutionResult> {
        let start_time = Instant::now();

        info!(
            saga_id = %saga_id,
            step = step_number,
            mutation = mutation_name,
            subgraph = subgraph,
            "Step execution started"
        );

        // 1. Validate step exists in saga (if store is available)
        if let Some(store) = &self.store {
            // Load saga to verify it exists
            let saga = store.load_saga(saga_id).await.map_err(|e| {
                warn!(saga_id = %saga_id, error = ?e, "Failed to load saga");
                e
            })?;

            if saga.is_none() {
                return Err(crate::federation::saga_store::SagaStoreError::SagaNotFound(saga_id));
            }

            // Load all steps for this saga
            let steps = store.load_saga_steps(saga_id).await.map_err(|e| {
                warn!(saga_id = %saga_id, error = ?e, "Failed to load saga steps");
                e
            })?;

            // Find the step we're executing
            let step_id = Uuid::new_v4(); // Placeholder ID for error reporting
            let saga_step = steps
                .iter()
                .find(|s| s.order == step_number as usize)
                .ok_or(crate::federation::saga_store::SagaStoreError::StepNotFound(step_id))?;

            // 2. Check step state is Pending
            if saga_step.state != StepState::Pending {
                return Err(
                    crate::federation::saga_store::SagaStoreError::InvalidStateTransition {
                        from: format!("{:?}", saga_step.state),
                        to:   "Executing".to_string(),
                    },
                );
            }

            // 3. Transition step to Executing
            store
                .update_saga_step_state(saga_step.id, &StepState::Executing)
                .await
                .map_err(|e| {
                    warn!(saga_id = %saga_id, step = step_number, error = ?e, "Failed to transition step to Executing");
                    e
                })?;

            info!(saga_id = %saga_id, step = step_number, "Step transitioned to Executing");

            // 4. Execute mutation via MutationExecutor (placeholder implementation)
            let result_data = serde_json::json!({
                "__typename": saga_step.typename,
                "id": format!("entity-{}-step-{}", saga_id, step_number),
                mutation_name: "executed",
            });

            // 5. Capture result data and transition step to Completed
            store
                .update_saga_step_result(saga_step.id, &result_data)
                .await
                .map_err(|e| {
                    warn!(saga_id = %saga_id, step = step_number, error = ?e, "Failed to save step result");
                    e
                })?;

            // 6. Transition step to Completed
            store
                .update_saga_step_state(saga_step.id, &StepState::Completed)
                .await
                .map_err(|e| {
                    warn!(saga_id = %saga_id, step = step_number, error = ?e, "Failed to transition step to Completed");
                    e
                })?;

            info!(saga_id = %saga_id, step = step_number, "Step transitioned to Completed");

            let duration_ms = start_time.elapsed().as_millis() as u64;

            let result = StepExecutionResult {
                step_number,
                success: true,
                data: Some(result_data),
                error: None,
                duration_ms,
            };

            info!(
                saga_id = %saga_id,
                step = step_number,
                duration_ms = result.duration_ms,
                "Step execution completed successfully"
            );

            Ok(result)
        } else {
            // No store available - return placeholder success for testing
            debug!("No saga store available - returning placeholder result");

            let duration_ms = start_time.elapsed().as_millis() as u64;

            let result = StepExecutionResult {
                step_number,
                success: true,
                data: Some(serde_json::json!({
                    "__typename": "Entity",
                    "id": format!("entity-{}", step_number),
                    mutation_name: "ok"
                })),
                error: None,
                duration_ms,
            };

            info!(
                saga_id = %saga_id,
                step = step_number,
                duration_ms = result.duration_ms,
                "Step execution completed (no store)"
            );

            Ok(result)
        }
    }

    /// Execute all steps in a saga sequentially
    ///
    /// # Arguments
    ///
    /// * `saga_id` - ID of saga to execute
    ///
    /// # Returns
    ///
    /// Vector of step results (successful or failed)
    pub async fn execute_saga(&self, saga_id: Uuid) -> SagaStoreResult<Vec<StepExecutionResult>> {
        info!(saga_id = %saga_id, "Saga forward phase started");

        // Execute all steps in order, stopping on first failure

        // If no store available, return empty results (for testing)
        let Some(store) = &self.store else {
            debug!("No saga store available - returning empty results");
            return Ok(vec![]);
        };

        // 1. Load saga from store
        let saga = store.load_saga(saga_id).await.map_err(|e| {
            warn!(saga_id = %saga_id, error = ?e, "Failed to load saga");
            e
        })?;

        if saga.is_none() {
            return Err(crate::federation::saga_store::SagaStoreError::SagaNotFound(saga_id));
        }

        // 2. Transition saga from Pending to Executing
        store
            .update_saga_state(saga_id, &crate::federation::saga_store::SagaState::Executing)
            .await
            .map_err(|e| {
                warn!(saga_id = %saga_id, error = ?e, "Failed to transition saga to Executing");
                e
            })?;

        info!(saga_id = %saga_id, "Saga transitioned to Executing");

        // 3. Load all steps for this saga
        let steps = store.load_saga_steps(saga_id).await.map_err(|e| {
            warn!(saga_id = %saga_id, error = ?e, "Failed to load saga steps");
            e
        })?;

        // Sort steps by order to ensure sequential execution
        let mut steps = steps;
        steps.sort_by_key(|s| s.order);

        let mut results: Vec<StepExecutionResult> = vec![];
        let mut saga_failed = false;

        // 4. For each step in order
        for step in steps {
            info!(
                saga_id = %saga_id,
                step = step.order,
                "Executing saga step"
            );

            // Execute this step
            // Construct mutation name from mutation type and typename
            let mutation_name = format!("{}_{}", step.mutation_type.as_str(), step.typename);

            match self
                .execute_step(
                    saga_id,
                    step.order as u32,
                    &mutation_name,
                    &step.variables,
                    &step.subgraph,
                )
                .await
            {
                Ok(step_result) => {
                    // Step succeeded - collect result and continue
                    info!(
                        saga_id = %saga_id,
                        step = step.order,
                        "Step executed successfully"
                    );
                    results.push(step_result);
                },
                Err(e) => {
                    // Step failed - capture error and stop execution
                    warn!(
                        saga_id = %saga_id,
                        step = step.order,
                        error = ?e,
                        "Step execution failed - stopping saga"
                    );

                    // Transition saga to Failed state
                    if let Err(state_err) = store
                        .update_saga_state(
                            saga_id,
                            &crate::federation::saga_store::SagaState::Failed,
                        )
                        .await
                    {
                        warn!(saga_id = %saga_id, error = ?state_err, "Failed to transition saga to Failed state");
                    }

                    saga_failed = true;
                    break;
                },
            }
        }

        // 5. Update final saga state
        if !saga_failed {
            // All steps succeeded - transition to Completed
            store
                .update_saga_state(saga_id, &crate::federation::saga_store::SagaState::Completed)
                .await
                .map_err(|e| {
                    warn!(saga_id = %saga_id, error = ?e, "Failed to transition saga to Completed");
                    e
                })?;

            info!(
                saga_id = %saga_id,
                steps_completed = results.len(),
                "Saga completed successfully"
            );
        }

        info!(
            saga_id = %saga_id,
            steps_completed = results.len(),
            "Saga forward phase completed"
        );

        Ok(results)
    }

    /// Get current execution state of saga
    ///
    /// # Arguments
    ///
    /// * `saga_id` - ID of saga
    ///
    /// # Returns
    ///
    /// Current execution state including completed steps
    pub async fn get_execution_state(&self, saga_id: Uuid) -> SagaStoreResult<ExecutionState> {
        // Load saga and steps from store to build current execution state

        // If no store available, return minimal state
        let Some(store) = &self.store else {
            debug!(saga_id = %saga_id, "No saga store available - returning empty execution state");
            let state = ExecutionState {
                saga_id,
                total_steps: 0,
                completed_steps: 0,
                current_step: None,
                failed: false,
                failure_reason: None,
            };
            return Ok(state);
        };

        // Load saga to get state and failure reason
        let saga = store.load_saga(saga_id).await.map_err(|e| {
            warn!(saga_id = %saga_id, error = ?e, "Failed to load saga for execution state");
            e
        })?;

        let (total_steps, completed_steps, failed, failure_reason, current_step) = match saga {
            Some(saga_data) => {
                // Load all steps to count completion
                let steps = store.load_saga_steps(saga_id).await.map_err(|e| {
                    warn!(saga_id = %saga_id, error = ?e, "Failed to load saga steps for execution state");
                    e
                })?;

                let total = steps.len() as u32;

                // Count completed steps
                let completed =
                    steps.iter().filter(|s| s.state == StepState::Completed).count() as u32;

                // Find first non-completed step as current_step
                let current =
                    steps.iter().find(|s| s.state != StepState::Completed).map(|s| s.order as u32);

                // Check if saga failed
                let is_failed = saga_data.state == crate::federation::saga_store::SagaState::Failed;

                (total, completed, is_failed, None, current)
            },
            None => {
                // Saga not found - return zero state
                (0, 0, false, None, None)
            },
        };

        let state = ExecutionState {
            saga_id,
            total_steps,
            completed_steps,
            current_step,
            failed,
            failure_reason,
        };

        debug!(
            saga_id = %saga_id,
            total_steps = state.total_steps,
            completed_steps = state.completed_steps,
            failed = state.failed,
            "Execution state queried"
        );

        Ok(state)
    }

    /// Check if step is safe to execute
    ///
    /// Validates:
    /// - Step exists in saga
    /// - Step is in Pending state
    /// - All @requires fields are available
    /// - Previous steps completed successfully
    ///
    /// Fetch any @requires fields before step execution
    ///
    /// Identifies fields required by a mutation and fetches them from
    /// other subgraphs if needed. This ensures all necessary data is
    /// available before executing the mutation.
    // Reason: tested in isolation, not yet called from execute_step()
    #[allow(dead_code)]
    async fn pre_fetch_requires_fields(
        &self,
        saga_id: Uuid,
        step_number: u32,
    ) -> SagaStoreResult<serde_json::Value> {
        // In a full implementation, would:
        // 1. Load step from saga store
        // 2. Extract @requires directive from mutation schema
        // 3. For each @requires field:
        //    - Determine owning subgraph
        //    - Create entity query to fetch the field
        //    - Execute query against subgraph
        // 4. Collect and return all fetched fields as JSON object

        info!(
            saga_id = %saga_id,
            step_number = step_number,
            "Pre-fetching @requires fields"
        );

        // For now, return empty object (no @requires fields)
        // In production: would merge fields from entity resolver
        Ok(serde_json::json!({}))
    }

    /// Build augmented entity data with @requires fields
    ///
    /// Merges @requires fields into the entity data, ensuring all
    /// necessary fields are present for mutation execution.
    // Reason: tested in isolation, not yet called from execute_step()
    #[allow(dead_code)]
    fn augment_entity_with_requires(
        &self,
        entity_data: serde_json::Value,
        requires_fields: serde_json::Value,
    ) -> serde_json::Value {
        // In a full implementation, would:
        // 1. Deep merge requires_fields into entity_data
        // 2. Handle nested object paths (e.g., "product.price")
        // 3. Validate all @requires fields are present
        // 4. Return fully augmented entity

        match (entity_data, requires_fields) {
            (serde_json::Value::Object(mut entity), serde_json::Value::Object(requires)) => {
                // Merge @requires fields into entity
                for (key, value) in requires {
                    entity.insert(key, value);
                }
                serde_json::Value::Object(entity)
            },
            (entity, _) => {
                // If entity is not an object, return as-is
                entity
            },
        }
    }
}

impl Default for SagaExecutor {
    fn default() -> Self {
        Self::new()
    }
}

/// Current execution state of a saga
#[derive(Debug, Clone)]
pub struct ExecutionState {
    /// Saga identifier
    pub saga_id:         Uuid,
    /// Total steps in saga
    pub total_steps:     u32,
    /// Number of completed steps
    pub completed_steps: u32,
    /// Currently executing step, if any
    pub current_step:    Option<u32>,
    /// Whether saga has failed
    pub failed:          bool,
    /// Reason for failure, if any
    pub failure_reason:  Option<String>,
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_saga_executor_creation() {
        let executor = SagaExecutor::new();
        drop(executor);
    }

    #[test]
    fn test_saga_executor_default() {
        let _executor = SagaExecutor::default();
        // Default should work
    }

    #[tokio::test]
    async fn test_step_execution_result() {
        let executor = SagaExecutor::new();
        let saga_id = Uuid::new_v4();
        let result = executor
            .execute_step(saga_id, 1, "testMutation", &serde_json::json!({}), "test-service")
            .await;

        assert!(result.is_ok());
        let step_result = result.unwrap();
        assert_eq!(step_result.step_number, 1);
        assert!(step_result.success);
    }

    #[tokio::test]
    async fn test_get_execution_state() {
        let executor = SagaExecutor::new();
        let saga_id = Uuid::new_v4();
        let state = executor.get_execution_state(saga_id).await;

        assert!(state.is_ok());
    }

    #[test]
    fn test_saga_executor_with_store() {
        // Test that we can create an executor with a store reference
        // Full store testing requires database setup (integration tests)
        let executor = SagaExecutor::new();
        assert!(!executor.has_store());
    }

    #[tokio::test]
    async fn test_execute_step_without_store() {
        // Verify that execute_step works without a store (fallback mode)
        let executor = SagaExecutor::new();
        let saga_id = Uuid::new_v4();
        let result = executor
            .execute_step(saga_id, 1, "testMutation", &serde_json::json!({}), "test-service")
            .await;

        assert!(result.is_ok());
        let step_result = result.unwrap();
        assert_eq!(step_result.step_number, 1);
        assert!(step_result.success);
        assert!(step_result.error.is_none());
    }

    #[tokio::test]
    async fn test_execute_saga_without_store() {
        // Verify execute_saga returns empty results without store
        let executor = SagaExecutor::new();
        let saga_id = Uuid::new_v4();
        let results = executor.execute_saga(saga_id).await;

        assert!(results.is_ok());
        let step_results = results.unwrap();
        assert_eq!(step_results.len(), 0);
    }

    #[tokio::test]
    async fn test_execute_saga_loads_saga_from_store() {
        // Verify that execute_saga attempts to load saga from store
        // This test verifies the store integration point
        let executor = SagaExecutor::new();
        let saga_id = Uuid::new_v4();

        // Without a store, should get empty results
        let results = executor.execute_saga(saga_id).await;
        assert!(results.is_ok());
    }

    #[tokio::test]
    async fn test_execute_all_steps_sequentially() {
        // Verify that steps are executed in order
        let executor = SagaExecutor::new();
        let saga_id = Uuid::new_v4();

        // Execute multiple steps
        for step_num in 1..=3 {
            let result = executor
                .execute_step(
                    saga_id,
                    step_num,
                    "testMutation",
                    &serde_json::json!({}),
                    "test-service",
                )
                .await;

            assert!(result.is_ok());
            let step_result = result.unwrap();
            assert_eq!(step_result.step_number, step_num);
            assert!(step_result.success);
        }
    }

    #[tokio::test]
    async fn test_saga_maintains_step_order() {
        // Verify that saga execution maintains step order
        let executor = SagaExecutor::new();
        let saga_id = Uuid::new_v4();

        let mut results = vec![];
        for step_num in 1..=3 {
            let result = executor
                .execute_step(saga_id, step_num, "mutation", &serde_json::json!({}), "service")
                .await;

            if let Ok(step_result) = result {
                results.push(step_result);
            }
        }

        // Verify order is maintained
        for (i, result) in results.iter().enumerate() {
            assert_eq!(result.step_number, (i + 1) as u32);
        }
    }

    #[tokio::test]
    async fn test_get_execution_state_without_store() {
        // Verify get_execution_state works without store
        let executor = SagaExecutor::new();
        let saga_id = Uuid::new_v4();
        let state = executor.get_execution_state(saga_id).await;

        assert!(state.is_ok());
        let execution_state = state.unwrap();
        assert_eq!(execution_state.saga_id, saga_id);
        assert_eq!(execution_state.total_steps, 0);
        assert_eq!(execution_state.completed_steps, 0);
        assert!(!execution_state.failed);
    }

    #[tokio::test]
    async fn test_execution_state_tracks_progress() {
        // Verify that execution state can track progress
        let executor = SagaExecutor::new();
        let saga_id = Uuid::new_v4();

        // Execute some steps
        for step_num in 1..=2 {
            let _ = executor
                .execute_step(saga_id, step_num, "mutation", &serde_json::json!({}), "service")
                .await;
        }

        // Get execution state
        let state = executor.get_execution_state(saga_id).await;
        assert!(state.is_ok());
    }

    #[tokio::test]
    async fn test_step_execution_captures_success_in_result() {
        // Verify that successful step execution captures data
        let executor = SagaExecutor::new();
        let saga_id = Uuid::new_v4();

        let result = executor
            .execute_step(saga_id, 1, "createOrder", &serde_json::json!({}), "orders-service")
            .await;

        assert!(result.is_ok());
        let step_result = result.unwrap();
        assert!(step_result.success);
        assert!(step_result.data.is_some());
        assert!(step_result.error.is_none());
    }

    #[tokio::test]
    async fn test_step_failure_detected() {
        // Verify that step failure is detected and captured
        // Tests limited scenario without a backing store. Full failure testing
        // requires mutation executor integration.
        let executor = SagaExecutor::new();
        let saga_id = Uuid::new_v4();

        let result = executor
            .execute_step(saga_id, 1, "mutation", &serde_json::json!({}), "service")
            .await;

        assert!(result.is_ok());
        // Success case without store - actual failure testing requires mutation executor
        // integration
    }

    #[tokio::test]
    async fn test_execution_result_includes_metrics() {
        // Verify that execution results include timing metrics
        let executor = SagaExecutor::new();
        let saga_id = Uuid::new_v4();

        let result = executor
            .execute_step(saga_id, 1, "mutation", &serde_json::json!({}), "service")
            .await;

        assert!(result.is_ok());
        let step_result = result.unwrap();
        // Verify that duration is measured
        let _ = step_result.duration_ms;
    }

    #[tokio::test]
    async fn test_pre_fetch_requires_fields() {
        let executor = SagaExecutor::new();
        let saga_id = Uuid::new_v4();

        let requires_fields = executor.pre_fetch_requires_fields(saga_id, 1).await;

        assert!(requires_fields.is_ok());
        let fields = requires_fields.unwrap();
        assert_eq!(fields, serde_json::json!({}));
    }

    #[test]
    fn test_augment_entity_with_requires() {
        let executor = SagaExecutor::new();

        let entity = serde_json::json!({
            "id": "user-123",
            "name": "Alice"
        });

        let requires = serde_json::json!({
            "email": "alice@example.com",
            "role": "admin"
        });

        let result = executor.augment_entity_with_requires(entity, requires);

        // Verify augmented entity contains both original and @requires fields
        assert_eq!(result.get("id").and_then(|v| v.as_str()), Some("user-123"));
        assert_eq!(result.get("name").and_then(|v| v.as_str()), Some("Alice"));
        assert_eq!(result.get("email").and_then(|v| v.as_str()), Some("alice@example.com"));
        assert_eq!(result.get("role").and_then(|v| v.as_str()), Some("admin"));
    }

    #[test]
    fn test_augment_entity_preserves_original_fields() {
        let executor = SagaExecutor::new();

        let entity = serde_json::json!({
            "id": "product-456",
            "price": 99.99
        });

        let requires = serde_json::json!({
            "category": "electronics"
        });

        let result = executor.augment_entity_with_requires(entity, requires);

        assert_eq!(result.get("id").and_then(|v| v.as_str()), Some("product-456"));
        assert_eq!(result.get("price").and_then(|v| v.as_f64()), Some(99.99));
        assert_eq!(result.get("category").and_then(|v| v.as_str()), Some("electronics"));
    }

    #[test]
    fn test_augment_entity_overwrites_conflicting_fields() {
        let executor = SagaExecutor::new();

        let entity = serde_json::json!({
            "id": "user-123",
            "status": "inactive"
        });

        let requires = serde_json::json!({
            "status": "active"
        });

        let result = executor.augment_entity_with_requires(entity, requires);

        // @requires should overwrite original value
        assert_eq!(result.get("status").and_then(|v| v.as_str()), Some("active"));
    }

    #[test]
    fn test_augment_entity_with_empty_requires() {
        let executor = SagaExecutor::new();

        let entity = serde_json::json!({
            "id": "test-123"
        });

        let requires = serde_json::json!({});

        let result = executor.augment_entity_with_requires(entity, requires);

        // Should return entity unchanged
        assert_eq!(result.get("id").and_then(|v| v.as_str()), Some("test-123"));
    }
}