durable-lambda-core 1.2.0

//! DurableContext — the main context struct passed to handler functions.
//!
//! Own the replay state machine, backend connection, and execution metadata.
//! Provide methods for all durable operations to interact with the replay engine.

use std::collections::HashMap;
use std::sync::Arc;

use aws_sdk_lambda::types::Operation;
use aws_sdk_lambda::types::OperationUpdate;

use crate::backend::DurableBackend;
use crate::error::DurableError;
use crate::replay::ReplayEngine;
use crate::types::{CompensationRecord, ExecutionMode};

/// Main context for a durable execution invocation.
///
/// `DurableContext` is created at the start of each Lambda invocation. It loads
/// the complete operation state from AWS (paginating if necessary), initializes
/// the replay engine, and provides the interface for durable operations.
///
/// # Construction
///
/// Use [`DurableContext::new`] to create a context from the invocation payload.
/// The constructor paginates through all remaining operations automatically.
///
/// # Examples
///
/// ```no_run
/// use durable_lambda_core::context::DurableContext;
/// use durable_lambda_core::backend::RealBackend;
/// use durable_lambda_core::types::ExecutionMode;
/// use std::sync::Arc;
/// use std::collections::HashMap;
///
/// # async fn example() -> Result<(), durable_lambda_core::error::DurableError> {
/// let config = aws_config::load_defaults(aws_config::BehaviorVersion::latest()).await;
/// let client = aws_sdk_lambda::Client::new(&config);
/// let backend = Arc::new(RealBackend::new(client));
///
/// let ctx = DurableContext::new(
///     backend,
///     "arn:aws:lambda:us-east-1:123456789:durable-execution/my-exec".to_string(),
///     "initial-token".to_string(),
///     vec![],       // initial operations from invocation payload
///     None,         // no more pages
/// ).await?;
///
/// match ctx.execution_mode() {
///     ExecutionMode::Replaying => println!("Replaying from history"),
///     ExecutionMode::Executing => println!("Executing new operations"),
/// }
/// # Ok(())
/// # }
/// ```
pub struct DurableContext {
    backend: Arc<dyn DurableBackend>,
    replay_engine: ReplayEngine,
    durable_execution_arn: String,
    checkpoint_token: String,
    parent_op_id: Option<String>,
    batch_mode: bool,
    pending_updates: Vec<OperationUpdate>,
    /// Registered compensation closures for the saga/compensation pattern.
    /// Populated by `step_with_compensation` on forward step success.
    /// Drained and executed in reverse order by `run_compensations`.
    compensations: Vec<CompensationRecord>,
}

/// Maximum items per page when paginating execution state.
const PAGE_SIZE: i32 = 1000;

impl DurableContext {
    /// Create a new `DurableContext` from invocation parameters.
    ///
    /// Loads the complete operation state by paginating through
    /// `get_execution_state` until all pages are fetched. Initializes the
    /// replay engine with the full operations map.
    ///
    /// # Arguments
    ///
    /// * `backend` — The durable execution backend (real or mock).
    /// * `arn` — The durable execution ARN.
    /// * `checkpoint_token` — The initial checkpoint token from the invocation payload.
    /// * `initial_operations` — First page of operations from the invocation payload.
    /// * `next_marker` — Pagination marker for additional pages (`None` if complete).
    ///
    /// # Errors
    ///
    /// Returns [`DurableError`] if paginating the execution state fails.
    pub async fn new(
        backend: Arc<dyn DurableBackend>,
        arn: String,
        checkpoint_token: String,
        initial_operations: Vec<Operation>,
        next_marker: Option<String>,
    ) -> Result<Self, DurableError> {
        let mut operations: HashMap<String, Operation> = initial_operations
            .into_iter()
            .map(|op| (op.id().to_string(), op))
            .collect();

        // Paginate remaining operations.
        let mut marker = next_marker;
        while let Some(ref m) = marker {
            if m.is_empty() {
                break;
            }
            let response = backend
                .get_execution_state(&arn, &checkpoint_token, m, PAGE_SIZE)
                .await?;

            for op in response.operations() {
                operations.insert(op.id().to_string(), op.clone());
            }

            marker = response.next_marker().map(|s| s.to_string());
        }

        let replay_engine = ReplayEngine::new(operations, None);

        Ok(Self {
            backend,
            replay_engine,
            durable_execution_arn: arn,
            checkpoint_token,
            parent_op_id: None,
            batch_mode: false,
            pending_updates: Vec::new(),
            compensations: Vec::new(),
        })
    }

    /// Return the current execution mode (Replaying or Executing).
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # async fn example(ctx: durable_lambda_core::context::DurableContext) {
    /// use durable_lambda_core::types::ExecutionMode;
    /// match ctx.execution_mode() {
    ///     ExecutionMode::Replaying => { /* returning cached results */ }
    ///     ExecutionMode::Executing => { /* running new operations */ }
    /// }
    /// # }
    /// ```
    pub fn execution_mode(&self) -> ExecutionMode {
        self.replay_engine.execution_mode()
    }

    /// Return whether the context is currently replaying from history.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # async fn example(ctx: durable_lambda_core::context::DurableContext) {
    /// if ctx.is_replaying() {
    ///     println!("Replaying cached operations");
    /// }
    /// # }
    /// ```
    pub fn is_replaying(&self) -> bool {
        self.replay_engine.is_replaying()
    }

    /// Return a reference to the durable execution ARN.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # async fn example(ctx: durable_lambda_core::context::DurableContext) {
    /// println!("Execution ARN: {}", ctx.arn());
    /// # }
    /// ```
    pub fn arn(&self) -> &str {
        &self.durable_execution_arn
    }

    /// Return the current checkpoint token.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # async fn example(ctx: durable_lambda_core::context::DurableContext) {
    /// let token = ctx.checkpoint_token();
    /// # }
    /// ```
    pub fn checkpoint_token(&self) -> &str {
        &self.checkpoint_token
    }

    /// Update the checkpoint token (called after a successful checkpoint).
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # async fn example(mut ctx: durable_lambda_core::context::DurableContext) {
    /// ctx.set_checkpoint_token("new-token-from-aws".to_string());
    /// # }
    /// ```
    pub fn set_checkpoint_token(&mut self, token: String) {
        self.checkpoint_token = token;
    }

    /// Return a reference to the backend.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # async fn example(ctx: durable_lambda_core::context::DurableContext) {
    /// let _backend = ctx.backend();
    /// # }
    /// ```
    pub fn backend(&self) -> &Arc<dyn DurableBackend> {
        &self.backend
    }

    /// Return a mutable reference to the replay engine.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # async fn example(mut ctx: durable_lambda_core::context::DurableContext) {
    /// let engine = ctx.replay_engine_mut();
    /// # }
    /// ```
    pub fn replay_engine_mut(&mut self) -> &mut ReplayEngine {
        &mut self.replay_engine
    }

    /// Create a child context for isolated operation ID namespacing.
    ///
    /// The child context shares the same backend and ARN but gets its own
    /// `ReplayEngine` with a parent-scoped `OperationIdGenerator`. Operations
    /// within the child context produce deterministic IDs scoped under
    /// `parent_op_id`, preventing collisions with the parent or sibling contexts.
    ///
    /// Used internally by parallel and child_context operations.
    ///
    /// # Arguments
    ///
    /// * `parent_op_id` — The operation ID that scopes this child context
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # async fn example(ctx: &durable_lambda_core::context::DurableContext) {
    /// let child = ctx.create_child_context("branch-op-id");
    /// // child operations will have IDs scoped under "branch-op-id"
    /// # }
    /// ```
    pub fn create_child_context(&self, parent_op_id: &str) -> DurableContext {
        let operations = self.replay_engine.operations().clone();
        let replay_engine = ReplayEngine::new(operations, Some(parent_op_id.to_string()));

        DurableContext {
            backend: self.backend.clone(),
            replay_engine,
            durable_execution_arn: self.durable_execution_arn.clone(),
            checkpoint_token: self.checkpoint_token.clone(),
            parent_op_id: Some(parent_op_id.to_string()),
            batch_mode: false,
            pending_updates: Vec::new(),
            compensations: Vec::new(), // NOT inherited from parent (isolated per context)
        }
    }

    /// Return a reference to the replay engine.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # async fn example(ctx: durable_lambda_core::context::DurableContext) {
    /// let engine = ctx.replay_engine();
    /// assert!(!engine.operations().is_empty() || true);
    /// # }
    /// ```
    pub fn replay_engine(&self) -> &ReplayEngine {
        &self.replay_engine
    }

    /// Return the parent operation ID, if this is a child context.
    ///
    /// Returns `None` for the root context. Returns the parent's operation ID
    /// for child contexts created via [`create_child_context`](Self::create_child_context).
    /// Used by replay-safe logging for hierarchical tracing.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # async fn example(ctx: &durable_lambda_core::context::DurableContext) {
    /// if let Some(parent_id) = ctx.parent_op_id() {
    ///     println!("Child context under parent: {parent_id}");
    /// }
    /// # }
    /// ```
    pub fn parent_op_id(&self) -> Option<&str> {
        self.parent_op_id.as_deref()
    }

    /// Enable batch checkpoint mode.
    ///
    /// When enabled, step operation checkpoints (START and SUCCEED/FAIL)
    /// are accumulated in memory instead of being sent immediately.
    /// Call [`flush_batch`](Self::flush_batch) to send all accumulated
    /// updates in a single AWS API call.
    ///
    /// Batch mode applies only to `step` operations. `wait`, `invoke`,
    /// and `callback` always send checkpoints immediately because they
    /// produce suspension errors that require the checkpoint to be
    /// persisted before the function exits.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # async fn example(mut ctx: durable_lambda_core::context::DurableContext) -> Result<(), durable_lambda_core::error::DurableError> {
    /// ctx.enable_batch_mode();
    /// let _: Result<i32, String> = ctx.step("step1", || async { Ok(1) }).await?;
    /// let _: Result<i32, String> = ctx.step("step2", || async { Ok(2) }).await?;
    /// ctx.flush_batch().await?;  // sends all updates in one call
    /// # Ok(())
    /// # }
    /// ```
    pub fn enable_batch_mode(&mut self) {
        self.batch_mode = true;
    }

    /// Return whether batch checkpoint mode is active.
    pub fn is_batch_mode(&self) -> bool {
        self.batch_mode
    }

    /// Accumulate an operation update for later batch flush.
    ///
    /// Called internally by step operations when batch mode is active.
    pub(crate) fn push_pending_update(&mut self, update: OperationUpdate) {
        self.pending_updates.push(update);
    }

    /// Return the number of pending (unflushed) updates.
    pub fn pending_update_count(&self) -> usize {
        self.pending_updates.len()
    }

    /// Return the number of registered compensation closures.
    ///
    /// Useful for asserting compensation registration in tests.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # async fn example(ctx: durable_lambda_core::context::DurableContext) {
    /// assert_eq!(ctx.compensation_count(), 0);
    /// # }
    /// ```
    pub fn compensation_count(&self) -> usize {
        self.compensations.len()
    }

    /// Register a compensation closure after a successful forward step.
    ///
    /// Called by `step_with_compensation` when the forward step succeeds.
    pub(crate) fn push_compensation(&mut self, record: CompensationRecord) {
        self.compensations.push(record);
    }

    /// Drain all registered compensations for execution.
    ///
    /// Returns the compensations vec (emptying the field) so `run_compensations`
    /// can execute them. Reversing the returned vec gives LIFO order.
    pub(crate) fn take_compensations(&mut self) -> Vec<CompensationRecord> {
        std::mem::take(&mut self.compensations)
    }

    /// Flush all accumulated checkpoint updates in a single AWS API call.
    ///
    /// No-op if no updates are pending. After flushing, the checkpoint
    /// token is updated from the response.
    ///
    /// # Errors
    ///
    /// Returns [`DurableError`] if the batch checkpoint call fails.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # async fn example(mut ctx: durable_lambda_core::context::DurableContext) -> Result<(), durable_lambda_core::error::DurableError> {
    /// ctx.enable_batch_mode();
    /// // ... run several steps ...
    /// ctx.flush_batch().await?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn flush_batch(&mut self) -> Result<(), DurableError> {
        if self.pending_updates.is_empty() {
            return Ok(());
        }
        let updates = std::mem::take(&mut self.pending_updates);
        let response = self
            .backend()
            .batch_checkpoint(self.arn(), self.checkpoint_token(), updates, None)
            .await?;
        let new_token = response.checkpoint_token().ok_or_else(|| {
            DurableError::checkpoint_failed(
                "batch",
                std::io::Error::new(
                    std::io::ErrorKind::InvalidData,
                    "batch checkpoint response missing checkpoint_token",
                ),
            )
        })?;
        self.set_checkpoint_token(new_token.to_string());
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use aws_sdk_lambda::operation::checkpoint_durable_execution::CheckpointDurableExecutionOutput;
    use aws_sdk_lambda::operation::get_durable_execution_state::GetDurableExecutionStateOutput;
    use aws_sdk_lambda::types::{OperationStatus, OperationType, OperationUpdate};
    /// A simple mock backend for testing context construction.
    struct TestBackend {
        pages: Vec<(Vec<Operation>, Option<String>)>,
    }

    #[async_trait::async_trait]
    impl DurableBackend for TestBackend {
        async fn checkpoint(
            &self,
            _arn: &str,
            _checkpoint_token: &str,
            _updates: Vec<OperationUpdate>,
            _client_token: Option<&str>,
        ) -> Result<CheckpointDurableExecutionOutput, DurableError> {
            unimplemented!("not needed for context tests")
        }

        async fn get_execution_state(
            &self,
            _arn: &str,
            _checkpoint_token: &str,
            next_marker: &str,
            _max_items: i32,
        ) -> Result<GetDurableExecutionStateOutput, DurableError> {
            let page_idx: usize = next_marker.parse().unwrap_or(0);
            if page_idx >= self.pages.len() {
                return Ok(GetDurableExecutionStateOutput::builder().build().unwrap());
            }
            let (ops, marker) = &self.pages[page_idx];
            let mut builder = GetDurableExecutionStateOutput::builder();
            for op in ops {
                builder = builder.operations(op.clone());
            }
            if let Some(m) = marker {
                builder = builder.next_marker(m);
            }
            Ok(builder.build().unwrap())
        }
    }

    fn make_op(id: &str, status: OperationStatus) -> Operation {
        Operation::builder()
            .id(id)
            .r#type(OperationType::Step)
            .status(status)
            .start_timestamp(aws_smithy_types::DateTime::from_secs(0))
            .build()
            .unwrap()
    }

    #[tokio::test]
    async fn empty_history_creates_executing_context() {
        let backend = Arc::new(TestBackend { pages: vec![] });
        let ctx = DurableContext::new(backend, "arn:test".into(), "tok".into(), vec![], None)
            .await
            .unwrap();

        assert_eq!(ctx.execution_mode(), ExecutionMode::Executing);
        assert!(!ctx.is_replaying());
        assert_eq!(ctx.arn(), "arn:test");
        assert_eq!(ctx.checkpoint_token(), "tok");
    }

    #[tokio::test]
    async fn initial_operations_loaded() {
        let backend = Arc::new(TestBackend { pages: vec![] });
        let ops = vec![make_op("op1", OperationStatus::Succeeded)];
        let ctx = DurableContext::new(backend, "arn:test".into(), "tok".into(), ops, None)
            .await
            .unwrap();

        assert!(ctx.is_replaying());
        assert!(ctx.replay_engine().check_result("op1").is_some());
    }

    #[tokio::test]
    async fn pagination_loads_all_pages() {
        let backend = Arc::new(TestBackend {
            pages: vec![
                (
                    vec![make_op("op2", OperationStatus::Succeeded)],
                    Some("1".to_string()),
                ),
                (vec![make_op("op3", OperationStatus::Succeeded)], None),
            ],
        });

        let initial = vec![make_op("op1", OperationStatus::Succeeded)];
        let ctx = DurableContext::new(
            backend,
            "arn:test".into(),
            "tok".into(),
            initial,
            Some("0".to_string()),
        )
        .await
        .unwrap();

        assert!(ctx.replay_engine().check_result("op1").is_some());
        assert!(ctx.replay_engine().check_result("op2").is_some());
        assert!(ctx.replay_engine().check_result("op3").is_some());
    }

    #[tokio::test]
    async fn set_checkpoint_token_updates() {
        let backend = Arc::new(TestBackend { pages: vec![] });
        let mut ctx = DurableContext::new(backend, "arn:test".into(), "tok1".into(), vec![], None)
            .await
            .unwrap();

        assert_eq!(ctx.checkpoint_token(), "tok1");
        ctx.set_checkpoint_token("tok2".to_string());
        assert_eq!(ctx.checkpoint_token(), "tok2");
    }

    // --- compensation field tests ---

    #[tokio::test]
    async fn new_context_has_empty_compensations() {
        let backend = Arc::new(TestBackend { pages: vec![] });
        let ctx = DurableContext::new(backend, "arn:test".into(), "tok".into(), vec![], None)
            .await
            .unwrap();

        assert_eq!(ctx.compensation_count(), 0);
    }

    #[tokio::test]
    async fn create_child_context_has_empty_compensations_not_inherited() {
        use crate::types::CompensationRecord;

        let backend = Arc::new(TestBackend { pages: vec![] });
        let mut ctx = DurableContext::new(backend, "arn:test".into(), "tok".into(), vec![], None)
            .await
            .unwrap();

        // Register a compensation on the parent
        let record = CompensationRecord {
            name: "parent_comp".to_string(),
            forward_result_json: serde_json::Value::Null,
            compensate_fn: Box::new(|_| Box::pin(async { Ok(()) })),
        };
        ctx.push_compensation(record);
        assert_eq!(ctx.compensation_count(), 1);

        // Child context should start with 0 compensations
        let child = ctx.create_child_context("some-op-id");
        assert_eq!(
            child.compensation_count(),
            0,
            "child context must NOT inherit parent compensations"
        );
    }
}