newton-task-submission 0.7.2

use crate::{
    contract_response_hash, contract_task_hash, submission_ids_field, task_ids_field, BatchIntentItem,
    PendingTaskRecord, SubmissionState, TaskChainPolicy, TaskExecutionIntent, TaskOperation, TaskPlanMemberRecord,
    TaskPlanRecord, TaskPlanState, TaskPlannerConfig, TaskPlannerObserver, TaskPlanningCommit, TaskPlanningStore,
    TaskPlanningWrite, TaskProjection,
};
use alloy::primitives::B256;
use newton_submission_protocol::{
    CompletedExecution, EffectOutcome, ExecutionChannel, ExecutionId, ExecutionOutcome, ExecutionPhase, ExecutionStatus,
};
use std::{collections::HashMap, sync::Arc, time::SystemTime};
use tokio_util::sync::CancellationToken;
use tracing::{error, info, warn};

const ACTIVE_PLAN_LIMIT: usize = 128;

/// Durable task planner, statically composed from one store implementing its
/// domain persistence and execution-channel APIs, plus runtime observation.
#[derive(Debug)]
pub struct TaskPlanner<S, O> {
    store: Arc<S>,
    observer: Arc<O>,
    config: TaskPlannerConfig,
    chains: HashMap<u64, TaskChainPolicy>,
}

impl<S, O> TaskPlanner<S, O>
where
    S: TaskPlanningStore + ExecutionChannel<TaskExecutionIntent> + 'static,
    O: TaskPlannerObserver + 'static,
{
    /// Statically composes the task store, observation, and policy.
    pub fn new(
        store: Arc<S>,
        observer: Arc<O>,
        config: TaskPlannerConfig,
        chains: HashMap<u64, TaskChainPolicy>,
    ) -> Self {
        Self {
            store,
            observer,
            config,
            chains,
        }
    }

    /// Reconciles and plans tasks until cancellation.
    pub async fn run(&self, cancellation: CancellationToken) {
        self.observer.heartbeat();
        let mut planning = tokio::time::interval(self.config.poll_interval());
        planning.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
        let mut chains = self.chains.values().collect::<Vec<_>>();
        chains.sort_unstable_by_key(|chain| chain.chain_id);
        let mut first_chain = 0_usize;
        loop {
            tokio::select! {
                _ = cancellation.cancelled() => break,
                _ = planning.tick() => {
                    // Heartbeat AFTER the tick, and only if every chain planned
                    // cleanly. Beating first marked the planner healthy before
                    // doing any work, so a planner failing every chain flipped
                    // healthy -> failed -> healthy once per poll interval
                    // (100ms by default) and whichever value a 5-30s scrape
                    // happened to catch was effectively a coin flip.
                    let mut clean = true;
                    for offset in 0..chains.len() {
                        let chain = chains[(first_chain + offset) % chains.len()];
                        match self.tick_chain(chain).await {
                            Ok(_) => {}
                            Err(error) => {
                                clean = false;
                                self.observer.failed(error.to_string());
                                error!(chain_id = chain.chain_id, %error, "task planning failed");
                            }
                        }
                    }
                    if clean {
                        self.observer.heartbeat();
                    }
                    if !chains.is_empty() {
                        first_chain = (first_chain + 1) % chains.len();
                    }
                }
            }
        }
    }

    /// Reconciles active plans and freezes at most one new batch for a chain.
    pub async fn tick_chain(&self, chain: &TaskChainPolicy) -> eyre::Result<Option<ExecutionId>> {
        let pending_limit = chain.max_batch_size.max(1).saturating_mul(2);
        let snapshot = self
            .store
            .load(chain.chain_id, pending_limit, ACTIVE_PLAN_LIMIT)
            .await
            .map_err(eyre::Report::new)?;
        for plan in &snapshot.plans {
            self.reconcile_plan(plan).await?;
        }
        if snapshot.plans.len() == ACTIVE_PLAN_LIMIT {
            return Ok(None);
        }

        let Some(plan) = select_plan(chain, self.config.batch_interval_ms, now_ms()?, &snapshot.pending)? else {
            return Ok(None);
        };
        let committed = self
            .store
            .commit(&TaskPlanningWrite::InsertPlan(plan.clone()))
            .await
            .map_err(eyre::Report::new)?;
        if committed == TaskPlanningCommit::Stale {
            return Ok(None);
        }
        self.observer.planned(chain.chain_id, plan.operation);
        info!(
            plan_id = %plan.plan_id,
            chain_id = chain.chain_id,
            operation = plan.operation.as_str(),
            member_count = plan.members.len(),
            submission_ids = %submission_ids_field(plan.intent.items()),
            task_ids = %task_ids_field(plan.intent.items()),
            "planned durable task batch"
        );
        self.emit_plan(&plan).await?;
        Ok(Some(plan.plan_id))
    }

    /// Advances one immutable plan by observing only durable common-channel
    /// state. Outcome projection covers all members atomically. For a projected
    /// plan, common acknowledgement is written before producer
    /// acknowledgement; retention preserves the common row across that crash
    /// boundary, so the sequence can resume on the next tick.
    async fn reconcile_plan(&self, plan: &TaskPlanRecord) -> eyre::Result<()> {
        match plan.state {
            TaskPlanState::Planned => self.emit_plan(plan).await,
            TaskPlanState::Submitted => {
                match self.store.status(plan.plan_id).await.map_err(eyre::Report::new)? {
                    ExecutionStatus::Pending(progress) => {
                        self.store
                            .commit(&TaskPlanningWrite::ProjectProgress {
                                plan_id: plan.plan_id,
                                state: task_state(progress.phase),
                                progress,
                            })
                            .await
                            .map_err(eyre::Report::new)?;
                    }
                    ExecutionStatus::Completed(completed) => {
                        let projections = project_outcome(plan, &completed)?;
                        let committed = self
                            .store
                            .commit(&TaskPlanningWrite::ProjectOutcome {
                                plan_id: plan.plan_id,
                                progress: completed.progress,
                                projections: projections.clone(),
                            })
                            .await
                            .map_err(eyre::Report::new)?;
                        // Only the write that actually applied may log. A plan
                        // already projected by an earlier tick would otherwise
                        // repeat every member's terminal line.
                        if committed == TaskPlanningCommit::Applied {
                            log_projections(plan, &projections);
                        }
                    }
                }
                Ok(())
            }
            TaskPlanState::Projected => {
                self.store.acknowledge(plan.plan_id).await.map_err(eyre::Report::new)?;
                self.store
                    .commit(&TaskPlanningWrite::MarkAcknowledged(plan.plan_id))
                    .await
                    .map_err(eyre::Report::new)?;
                Ok(())
            }
            TaskPlanState::Acknowledged => Ok(()),
        }
    }

    /// Hands a durable plan to the common execution channel before marking it
    /// submitted. Repeating this after a crash is safe because `plan_id` binds
    /// the same immutable request bytes in both domains.
    async fn emit_plan(&self, plan: &TaskPlanRecord) -> eyre::Result<()> {
        self.store.submit(&plan.execution()).await.map_err(eyre::Report::new)?;
        self.store
            .commit(&TaskPlanningWrite::MarkSubmitted(plan.plan_id))
            .await
            .map_err(eyre::Report::new)?;
        Ok(())
    }
}

/// Emits one terminal or requeue line per plan member.
///
/// The executor and common channel are batch-oriented, so without this a task
/// that was admitted individually has no individual outcome anywhere in the
/// logs. `task_id` is recovered from the immutable intent, which is the only
/// place the planner still holds the contract identity at projection time.
fn log_projections(plan: &TaskPlanRecord, projections: &[TaskProjection]) {
    let task_ids = plan
        .intent
        .items()
        .iter()
        .map(|item| (item.submission_id, item.task.taskId))
        .collect::<HashMap<_, _>>();
    for projection in projections {
        let task_id = task_ids.get(&projection.submission_id).copied().unwrap_or_default();
        match projection.state {
            SubmissionState::Failed => warn!(
                submission_id = %projection.submission_id,
                %task_id,
                plan_id = %plan.plan_id,
                chain_id = plan.chain_id,
                operation = plan.operation.as_str(),
                terminal_error = projection.terminal_error.as_deref().unwrap_or("unknown"),
                "task submission failed permanently"
            ),
            SubmissionState::BatchPending => info!(
                submission_id = %projection.submission_id,
                %task_id,
                plan_id = %plan.plan_id,
                chain_id = plan.chain_id,
                retry_operation = projection.operation.map_or("unchanged", TaskOperation::as_str),
                "task submission requeued for another batch"
            ),
            state => info!(
                submission_id = %projection.submission_id,
                %task_id,
                plan_id = %plan.plan_id,
                chain_id = plan.chain_id,
                operation = plan.operation.as_str(),
                state = state.as_str(),
                "task submission completed on-chain"
            ),
        }
    }
}

fn select_plan(
    chain: &TaskChainPolicy,
    interval_ms: u64,
    now: i64,
    pending: &[PendingTaskRecord],
) -> eyre::Result<Option<TaskPlanRecord>> {
    let max_batch = chain.max_batch_size.max(1);
    let mut groups = [
        (TaskOperation::CombinedCreateAndRespond, Vec::new()),
        (TaskOperation::RespondOnly, Vec::new()),
    ];
    for candidate in pending {
        let group = if candidate.operation == TaskOperation::CombinedCreateAndRespond {
            &mut groups[0].1
        } else {
            &mut groups[1].1
        };
        if group.len() < max_batch {
            group.push(candidate);
        }
    }
    let selected = groups
        .into_iter()
        .filter(|(_, group)| {
            group.len() == max_batch
                || group.first().is_some_and(|oldest| {
                    now.saturating_sub(oldest.accepted_at_ms) >= i64::try_from(interval_ms).unwrap_or(i64::MAX)
                })
        })
        .min_by_key(|(_, group)| {
            group
                .first()
                .map(|candidate| (candidate.accepted_at_ms, candidate.submission_id.to_string()))
        });
    let Some((operation, candidates)) = selected else {
        return Ok(None);
    };

    let mut items = Vec::with_capacity(candidates.len());
    let mut members = Vec::with_capacity(candidates.len());
    let mut deadline = i64::MAX;
    for candidate in candidates {
        items.push(BatchIntentItem {
            submission_id: candidate.submission_id,
            expected_task_hash: contract_task_hash(&candidate.payload.task),
            expected_response_hash: contract_response_hash(&candidate.payload.task_response),
            task: candidate.payload.task.clone(),
            response: candidate.payload.task_response.clone(),
            signature_data: candidate.payload.signature_data.clone(),
            attestation_data: candidate.payload.attestation_data.clone(),
        });
        members.push(TaskPlanMemberRecord {
            submission_id: candidate.submission_id,
            effect_retry_count: candidate.effect_retry_count,
            max_effect_retries: candidate.max_effect_retries,
        });
        deadline = deadline.min(candidate.deadline_at_ms);
    }
    let intent = match operation {
        TaskOperation::CombinedCreateAndRespond => TaskExecutionIntent::CreateAndRespond {
            contract_role: "batch_task_manager".to_string(),
            items,
        },
        TaskOperation::RespondOnly => TaskExecutionIntent::Respond {
            contract_role: "batch_task_manager".to_string(),
            items,
        },
    };
    Ok(Some(TaskPlanRecord {
        plan_id: ExecutionId::new(),
        chain_id: chain.chain_id,
        operation,
        intent,
        deadline_at_ms: Some(deadline),
        state: TaskPlanState::Planned,
        members,
    }))
}

fn project_outcome(plan: &TaskPlanRecord, completed: &CompletedExecution) -> eyre::Result<Vec<TaskProjection>> {
    let dispositions = match &completed.outcome {
        ExecutionOutcome::Effects(effects) => {
            if effects.len() != plan.members.len() {
                return Err(eyre::eyre!(
                    "task plan {} expected {} effects but received {}",
                    plan.plan_id,
                    plan.members.len(),
                    effects.len()
                ));
            }
            effects
                .iter()
                .map(|effect| match effect {
                    EffectOutcome::Succeeded => TaskDisposition::Succeeded,
                    EffectOutcome::Missing => TaskDisposition::Retry(plan.operation),
                    EffectOutcome::PartiallySucceeded => TaskDisposition::Retry(TaskOperation::RespondOnly),
                    EffectOutcome::Conflict => TaskDisposition::Failed("onchain_effect_conflict".to_string()),
                    EffectOutcome::Failed(error) => TaskDisposition::Failed(error.clone()),
                })
                .collect()
        }
        ExecutionOutcome::RetryableFailure(_) => vec![TaskDisposition::Retry(plan.operation); plan.members.len()],
        ExecutionOutcome::PermanentFailure(error) => {
            vec![TaskDisposition::Failed(error.clone()); plan.members.len()]
        }
    };
    Ok(plan
        .members
        .iter()
        .zip(dispositions)
        .map(|(member, disposition)| project_member(member, disposition))
        .collect())
}

#[derive(Debug, Clone)]
enum TaskDisposition {
    Succeeded,
    Retry(TaskOperation),
    Failed(String),
}

fn project_member(member: &TaskPlanMemberRecord, disposition: TaskDisposition) -> TaskProjection {
    match disposition {
        TaskDisposition::Succeeded => TaskProjection {
            submission_id: member.submission_id,
            state: SubmissionState::Succeeded,
            operation: None,
            increment_effect_retry: false,
            terminal_error: None,
        },
        TaskDisposition::Retry(_) if member.effect_retry_count >= member.max_effect_retries => TaskProjection {
            submission_id: member.submission_id,
            state: SubmissionState::Failed,
            operation: None,
            increment_effect_retry: false,
            terminal_error: Some("onchain_effect_retry_exhausted".to_string()),
        },
        TaskDisposition::Retry(operation) => TaskProjection {
            submission_id: member.submission_id,
            state: SubmissionState::BatchPending,
            operation: Some(operation),
            increment_effect_retry: true,
            terminal_error: None,
        },
        TaskDisposition::Failed(error) => TaskProjection {
            submission_id: member.submission_id,
            state: SubmissionState::Failed,
            operation: None,
            increment_effect_retry: false,
            terminal_error: Some(error),
        },
    }
}

const fn task_state(phase: ExecutionPhase) -> SubmissionState {
    match phase {
        ExecutionPhase::Ready => SubmissionState::ReadyForSubmission,
        ExecutionPhase::Assigned => SubmissionState::Assigned,
        ExecutionPhase::Prepared => SubmissionState::Prepared,
        ExecutionPhase::Broadcast => SubmissionState::Broadcast,
        ExecutionPhase::Mined => SubmissionState::Mined,
    }
}

fn now_ms() -> eyre::Result<i64> {
    let elapsed = SystemTime::now().duration_since(SystemTime::UNIX_EPOCH)?;
    i64::try_from(elapsed.as_millis()).map_err(Into::into)
}

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

    #[test]
    fn partial_combined_effect_retries_response_only() {
        let member = TaskPlanMemberRecord {
            submission_id: crate::SubmissionId::new(),
            effect_retry_count: 0,
            max_effect_retries: 3,
        };
        let projected = project_member(&member, TaskDisposition::Retry(TaskOperation::RespondOnly));
        assert_eq!(projected.state, SubmissionState::BatchPending);
        assert_eq!(projected.operation, Some(TaskOperation::RespondOnly));
        assert!(projected.increment_effect_retry);
    }

    #[test]
    fn exhausted_retry_becomes_a_planner_selected_failure() {
        let member = TaskPlanMemberRecord {
            submission_id: crate::SubmissionId::new(),
            effect_retry_count: 3,
            max_effect_retries: 3,
        };
        let projected = project_member(&member, TaskDisposition::Retry(TaskOperation::CombinedCreateAndRespond));
        assert_eq!(projected.state, SubmissionState::Failed);
        assert_eq!(
            projected.terminal_error.as_deref(),
            Some("onchain_effect_retry_exhausted")
        );
    }
}