klieo-workflow-api 3.5.0

//! HTTP handlers for the run service.

use crate::auth::can_read_run;
use crate::error::WorkflowApiError;
use crate::executor::{spawn_run, RunTask};
use crate::provenance::RunProvenance;
use crate::run_store::{ClaimOutcome, RunRecordView};
use crate::state::WorkflowRunState;
use axum::{
    extract::{Path, State},
    http::StatusCode,
    response::{IntoResponse, Response},
    Extension, Json,
};
use klieo_auth_common::Identity;
use klieo_flows::Flow;
use klieo_workflow::{canonical_json_bytes, compile, WorkflowDef};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use sha2::{Digest, Sha256};
use std::sync::Arc;
use tokio::sync::OwnedSemaphorePermit;

/// Agent name stamped on the `AgentContext` minted per run.
const RUNNER_AGENT_NAME: &str = "workflow-run";

/// Body of `POST /runs`: the definition to run plus its input envelope.
#[derive(Debug, Deserialize)]
pub(crate) struct SubmitBody {
    /// The declarative workflow to compile and run.
    pub definition: WorkflowDef,
    /// The initial envelope; must be a JSON object.
    #[serde(default)]
    pub input: Value,
}

/// `202 Accepted` body: the id to poll.
#[derive(Debug, Serialize)]
struct SubmitAccepted {
    run_id: String,
}

/// `POST /runs` — compile the definition and start a supervised run.
///
/// Order: reject a non-object input, compile (sync `400` on failure), acquire
/// a concurrency permit (`429` when saturated), then claim the idempotency key
/// via CAS. Compile + permit precede the claim so an invalid def or a
/// saturated server never strands an idempotency claim on a run that will
/// never exist. On the winning path the run-start entry is appended to the
/// durable audit chain before the run is spawned.
pub(crate) async fn submit(
    State(state): State<Arc<WorkflowRunState>>,
    Extension(identity): Extension<Identity>,
    Json(body): Json<SubmitBody>,
) -> Result<Response, WorkflowApiError> {
    if !body.input.is_object() {
        return Err(WorkflowApiError::BadRequest(
            "input must be a JSON object".into(),
        ));
    }
    let flow = compile(&body.definition, &state.registry)?;
    let permit = Arc::clone(&state.permits)
        .try_acquire_owned()
        .map_err(|_| WorkflowApiError::TooManyRuns)?;

    let run_id = ulid::Ulid::new().to_string();
    let key = idempotency_key(&body.definition, &body.input);
    let provenance = run_provenance(&state, &run_id, &body.definition, identity.as_str());
    if let Some(existing) = claim_or_existing(&state, &provenance, &key).await? {
        return Ok(accepted(existing));
    }
    if let Err(e) = provenance.record_started().await {
        roll_back_claim(&state, &key, &run_id).await;
        return Err(WorkflowApiError::Internal(Box::new(e)));
    }

    spawn_task(&state, provenance, flow, body.input, permit);
    Ok(accepted(run_id))
}

/// Register the run's live-progress channel and spawn its supervised executor.
fn spawn_task(
    state: &WorkflowRunState,
    provenance: RunProvenance,
    flow: Arc<dyn Flow>,
    input: Value,
    permit: OwnedSemaphorePermit,
) {
    let task = RunTask {
        flow,
        ctx: state.app.context(RUNNER_AGENT_NAME),
        input,
        timeout: state.run_timeout,
        permit,
        progress_sender: state.progress.register(&provenance.run_id),
    };
    spawn_run(
        state.run_store.clone(),
        provenance,
        state.progress.clone(),
        task,
    );
}

/// Undo the idempotency claim + pending record after a run-start audit
/// failure, so a genuine retry re-runs. Best-effort — logged on failure, with
/// the startup sweep as the cross-restart backstop.
async fn roll_back_claim(state: &WorkflowRunState, idempotency_key: &str, run_id: &str) {
    if let Err(e) = state
        .run_store
        .rollback_claim(idempotency_key, run_id)
        .await
    {
        tracing::error!(run_id, error = %e, "failed to roll back claim after run-start audit failure");
    }
}

fn run_provenance(
    state: &WorkflowRunState,
    run_id: &str,
    definition: &WorkflowDef,
    author: &str,
) -> RunProvenance {
    RunProvenance {
        repo: Arc::clone(&state.provenance),
        run_id: run_id.to_string(),
        author: author.to_string(),
        workflow_id: definition.id.clone(),
        content_hash: definition.canonical_hash(),
    }
}

/// Claim the idempotency key; on the CAS-loser path return the winning run's
/// id (the caller replays it), otherwise create the pending record and
/// return `None`.
async fn claim_or_existing(
    state: &WorkflowRunState,
    provenance: &RunProvenance,
    idempotency_key: &str,
) -> Result<Option<String>, WorkflowApiError> {
    match state
        .run_store
        .claim_idempotent(idempotency_key, &provenance.run_id)
        .await
    {
        Ok(ClaimOutcome::AlreadyExists(existing)) => Ok(Some(existing)),
        Ok(ClaimOutcome::Created) => {
            let view = RunRecordView::pending(
                &provenance.run_id,
                &provenance.workflow_id,
                &provenance.author,
                now_rfc3339(),
            );
            state
                .run_store
                .create(&view)
                .await
                .map_err(|e| WorkflowApiError::Internal(Box::new(e)))?;
            Ok(None)
        }
        Err(e) => Err(WorkflowApiError::Internal(Box::new(e))),
    }
}

/// `GET /runs/{id}` — return the run's current status snapshot.
///
/// `404` (not `403`) for a run authored by someone else: the caller's
/// [`WORKFLOW_READ_SCOPE`](crate::WORKFLOW_READ_SCOPE) alone only proves it
/// may read runs, not which ones — a run's mere existence is not disclosed
/// to a non-owner. [`WORKFLOW_READ_ALL_SCOPE`](crate::WORKFLOW_READ_ALL_SCOPE)
/// lifts that restriction.
pub(crate) async fn get_run(
    State(state): State<Arc<WorkflowRunState>>,
    Extension(identity): Extension<Identity>,
    Path(run_id): Path<String>,
) -> Result<Json<RunRecordView>, WorkflowApiError> {
    let view = match state.run_store.get(&run_id).await {
        Ok(Some(view)) => view,
        Ok(None) => return Err(WorkflowApiError::RunNotFound),
        Err(e) => return Err(WorkflowApiError::Internal(Box::new(e))),
    };
    if !can_read_run(&identity, &view.author) {
        return Err(WorkflowApiError::RunNotFound);
    }
    Ok(Json(view))
}

/// The palette a caller may compose a workflow against.
#[derive(Debug, Serialize)]
pub(crate) struct RegistryView {
    models: Vec<String>,
    tools: Vec<String>,
    subflows: Vec<String>,
}

/// `GET /health` — unauthenticated liveness probe. A caller that gets a
/// response at all (regardless of body) knows the process is up and the
/// router is serving; readiness (dependency connectivity) is not checked.
pub(crate) async fn health() -> &'static str {
    "ok"
}

/// `GET /registry` — enumerate the allowed model / tool / subflow ids.
pub(crate) async fn get_registry(State(state): State<Arc<WorkflowRunState>>) -> Json<RegistryView> {
    Json(RegistryView {
        models: owned(state.registry.model_ids()),
        tools: owned(state.registry.tool_ids()),
        subflows: owned(state.registry.subflow_ids()),
    })
}

fn owned(ids: Vec<&str>) -> Vec<String> {
    ids.into_iter().map(String::from).collect()
}

fn accepted(run_id: String) -> Response {
    (StatusCode::ACCEPTED, Json(SubmitAccepted { run_id })).into_response()
}

/// Idempotency key = SHA-256 over the definition's canonical hash and the
/// input's canonical bytes, so resubmits that differ only in JSON key order
/// (of the definition OR the input) collapse to one run.
fn idempotency_key(definition: &WorkflowDef, input: &Value) -> String {
    let mut hasher = Sha256::new();
    hasher.update(definition.canonical_hash().as_bytes());
    hasher.update([0u8]);
    hasher.update(canonical_json_bytes(input));
    hex::encode(hasher.finalize())
}

fn now_rfc3339() -> String {
    chrono::Utc::now().to_rfc3339()
}