klieo-workflow-api 3.3.0

//! Supervised run executor.
//!
//! The executor upholds the terminal-write guarantee: every run writes
//! exactly one terminal status — `Succeeded`, `Failed` (error OR panic), or
//! `Aborted` (timeout) — so a run never sticks `Running`. The flow runs in a
//! child task whose `JoinError` (a panic) is caught and mapped to `Failed`,
//! and a per-run timeout aborts a runaway flow.
//!
//! The timeout `abort()` is cooperative: it cancels the child task at its
//! next `.await` point. An await-heavy flow (LLM calls, bus, tool dispatch)
//! stops promptly; a hypothetical CPU-bound body with no awaits would run to
//! completion before the abort takes effect. Workflow flows are await-heavy,
//! so this bound holds in practice.

use crate::progress::ProgressHub;
use crate::provenance::{RunProvenance, STATUS_ABORTED, STATUS_FAILED, STATUS_SUCCEEDED};
use crate::run_store::RunStore;
use klieo_core::agent::{AgentContext, AgentEvent};
use klieo_flows::Flow;
use serde_json::Value;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::{broadcast, OwnedSemaphorePermit};

const TIMEOUT_REASON: &str = "run exceeded the maximum duration";
const PANIC_REASON: &str = "run panicked";

/// Terminal outcome of a supervised run.
enum Outcome {
    Succeeded(Value),
    Failed(String),
    Aborted(String),
}

/// The flow + runtime inputs for one supervised run.
pub(crate) struct RunTask {
    pub flow: Arc<dyn Flow>,
    pub ctx: AgentContext,
    pub input: Value,
    pub timeout: Duration,
    pub permit: OwnedSemaphorePermit,
    /// Live-progress sender installed as `ctx.progress` for this run.
    pub progress_sender: broadcast::Sender<AgentEvent>,
}

/// Spawn a supervised run. The permit is held for the run's lifetime and
/// released when the task ends. Installs the run's live-progress sender on the
/// context, writes a terminal status on every exit path (appending the
/// terminal provenance entry FIRST, authoritative, before the non-authoritative
/// run-store update), and removes the run's progress channel on finalize so it
/// cannot leak.
pub(crate) fn spawn_run(
    run_store: RunStore,
    provenance: RunProvenance,
    progress: ProgressHub,
    task: RunTask,
) {
    tokio::spawn(async move {
        let _permit = task.permit;
        let mut ctx = task.ctx;
        ctx.progress = Some(task.progress_sender);
        if let Err(e) = run_store.set_running(&provenance.run_id).await {
            tracing::error!(run_id = provenance.run_id, error = %e, "failed to mark run running");
        }
        let outcome = execute(task.flow, ctx, task.input, task.timeout).await;
        finalize(&run_store, &provenance, outcome).await;
        progress.deregister(&provenance.run_id);
    });
}

/// Run the flow in a child task so a panic becomes a catchable `JoinError`,
/// bounded by `timeout`.
async fn execute(
    flow: Arc<dyn Flow>,
    ctx: AgentContext,
    input: Value,
    timeout: Duration,
) -> Outcome {
    let mut child = tokio::spawn(async move { flow.run(ctx, input).await });
    match tokio::time::timeout(timeout, &mut child).await {
        Ok(Ok(Ok(value))) => Outcome::Succeeded(value),
        Ok(Ok(Err(flow_err))) => Outcome::Failed(flow_err.to_string()),
        Ok(Err(join_err)) => {
            tracing::error!(error = %join_err, "workflow run task panicked");
            Outcome::Failed(PANIC_REASON.to_string())
        }
        Err(_elapsed) => {
            child.abort();
            Outcome::Aborted(TIMEOUT_REASON.to_string())
        }
    }
}

/// Record the terminal state: append to the audit chain FIRST (authoritative),
/// then update the volatile run-store cache. If the store write fails after
/// the chain append, the audit is still complete.
async fn finalize(run_store: &RunStore, provenance: &RunProvenance, outcome: Outcome) {
    let (status, reason) = outcome_labels(&outcome);
    if let Err(e) = provenance.record_terminal(status, reason).await {
        tracing::error!(run_id = provenance.run_id, error = %e, "failed to append terminal provenance entry");
    }
    let run_id = &provenance.run_id;
    let write = match outcome {
        Outcome::Succeeded(value) => run_store.set_succeeded(run_id, value).await,
        Outcome::Failed(msg) => run_store.set_failed(run_id, msg).await,
        Outcome::Aborted(msg) => run_store.set_aborted(run_id, msg).await,
    };
    if let Err(e) = write {
        tracing::error!(run_id, error = %e, "failed to write terminal run status");
    }
}

/// Map an outcome to its provenance `(status, reason)` metadata.
fn outcome_labels(outcome: &Outcome) -> (&'static str, Option<String>) {
    match outcome {
        Outcome::Succeeded(_) => (STATUS_SUCCEEDED, None),
        Outcome::Failed(msg) => (STATUS_FAILED, Some(msg.clone())),
        Outcome::Aborted(msg) => (STATUS_ABORTED, Some(msg.clone())),
    }
}