crank/crates/crank-runtime/src/idempotency.rs

use std::time::{Duration, Instant};

use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
use crank_core::{
    CacheScope, CoordinationStateReservation, CoordinationStateStore, CoordinationStateValue,
    HttpMethod, IdempotencyMode, IdempotencyPolicy, Target,
};
use serde_json::{Map, Value, json};
use sha2::{Digest, Sha256};
use uuid::Uuid;

use crate::{
    AdapterResponse, PreparedRequest, RuntimeError, RuntimeOperation, RuntimeRequestContext,
};

const POLL_INTERVAL: Duration = Duration::from_millis(10);

pub(crate) enum IdempotencyAction {
    Disabled,
    Execute(IdempotencyReservation),
    Replay(AdapterResponse),
}

#[derive(Clone)]
pub(crate) struct IdempotencyReservation {
    key: String,
    initial: CoordinationStateValue,
    fingerprint: String,
    result_ttl: Duration,
}

pub fn prepare_idempotency(
    operation: &RuntimeOperation,
    input: &Value,
    prepared_request: &mut PreparedRequest,
) -> Result<Option<String>, RuntimeError> {
    let Some(policy) = policy(operation) else {
        return Ok(None);
    };

    let key = key_from_policy(policy, input, prepared_request);
    let Some(key) = key else {
        if policy.mode == IdempotencyMode::Required {
            return Err(RuntimeError::InvalidPreparedRequest {
                field: "execution_config.idempotency".to_owned(),
                reason: "required idempotency key was not provided".to_owned(),
            });
        }
        return Ok(None);
    };

    if let Some(header_name) = policy
        .header_name
        .as_deref()
        .filter(|value| !value.is_empty())
    {
        prepared_request
            .headers
            .entry(header_name.to_owned())
            .or_insert_with(|| key.clone());
    }

    Ok(Some(key))
}

pub(crate) async fn begin(
    store: Option<&dyn CoordinationStateStore>,
    operation: &RuntimeOperation,
    input: &Value,
    idempotency_key: Option<&str>,
    request_context: Option<&RuntimeRequestContext>,
) -> Result<IdempotencyAction, RuntimeError> {
    let Some(policy) = policy(operation) else {
        return Ok(IdempotencyAction::Disabled);
    };
    let Some(idempotency_key) = idempotency_key else {
        return Ok(IdempotencyAction::Disabled);
    };
    let Some(key) = cache_key(operation, idempotency_key, request_context) else {
        return Err(RuntimeError::IdempotencyStoreUnavailable {
            operation_id: operation.operation_id.as_str().to_owned(),
        });
    };
    let Some(store) = store else {
        return Err(RuntimeError::IdempotencyStoreUnavailable {
            operation_id: operation.operation_id.as_str().to_owned(),
        });
    };

    let fingerprint = request_fingerprint(input)?;
    let result_ttl = Duration::from_millis(policy.ttl_ms);
    let reservation_ttl = result_ttl.max(Duration::from_millis(
        operation.execution_config.timeout_ms.saturating_add(1_000),
    ));
    let initial = in_progress_value(&fingerprint);
    let reservation = store
        .reserve_value(
            CacheScope::Coordination,
            &key,
            initial.clone(),
            reservation_ttl,
        )
        .await
        .map_err(|_| RuntimeError::IdempotencyStoreUnavailable {
            operation_id: operation.operation_id.as_str().to_owned(),
        })?;

    match reservation {
        CoordinationStateReservation::Reserved => {
            Ok(IdempotencyAction::Execute(IdempotencyReservation {
                key,
                initial,
                fingerprint,
                result_ttl,
            }))
        }
        CoordinationStateReservation::Existing(existing) => {
            resolve_existing(store, operation, &key, &fingerprint, existing).await
        }
    }
}

pub(crate) async fn complete(
    store: &dyn CoordinationStateStore,
    operation: &RuntimeOperation,
    reservation: &IdempotencyReservation,
    response: &AdapterResponse,
) -> Result<(), RuntimeError> {
    let completed = CoordinationStateValue {
        payload: json!({
            "state": "completed",
            "fingerprint": reservation.fingerprint,
            "response": response,
        }),
    };
    let replaced = store
        .compare_and_set_value(
            CacheScope::Coordination,
            &reservation.key,
            &reservation.initial,
            completed,
            reservation.result_ttl,
        )
        .await
        .map_err(|_| RuntimeError::IdempotencyStoreUnavailable {
            operation_id: operation.operation_id.as_str().to_owned(),
        })?;
    if replaced {
        Ok(())
    } else {
        Err(RuntimeError::IdempotencyOutcomeUnknown {
            operation_id: operation.operation_id.as_str().to_owned(),
        })
    }
}

pub(crate) async fn mark_outcome_unknown(
    store: &dyn CoordinationStateStore,
    operation: &RuntimeOperation,
    reservation: &IdempotencyReservation,
) -> Result<(), RuntimeError> {
    let unknown = CoordinationStateValue {
        payload: json!({
            "state": "outcome_unknown",
            "fingerprint": reservation.fingerprint,
        }),
    };
    let replaced = store
        .compare_and_set_value(
            CacheScope::Coordination,
            &reservation.key,
            &reservation.initial,
            unknown,
            reservation.result_ttl,
        )
        .await
        .map_err(|_| RuntimeError::IdempotencyStoreUnavailable {
            operation_id: operation.operation_id.as_str().to_owned(),
        })?;
    if replaced {
        Ok(())
    } else {
        Err(RuntimeError::IdempotencyOutcomeUnknown {
            operation_id: operation.operation_id.as_str().to_owned(),
        })
    }
}

async fn resolve_existing(
    store: &dyn CoordinationStateStore,
    operation: &RuntimeOperation,
    key: &str,
    fingerprint: &str,
    mut existing: CoordinationStateValue,
) -> Result<IdempotencyAction, RuntimeError> {
    let operation_id = operation.operation_id.as_str().to_owned();
    let deadline =
        Instant::now() + Duration::from_millis(operation.execution_config.timeout_ms.max(1));
    loop {
        let existing_fingerprint = existing.payload.get("fingerprint").and_then(Value::as_str);
        if existing_fingerprint != Some(fingerprint) {
            return Err(RuntimeError::IdempotencyConflict { operation_id });
        }
        match existing.payload.get("state").and_then(Value::as_str) {
            Some("completed") => {
                let response = existing
                    .payload
                    .get("response")
                    .cloned()
                    .and_then(|value| serde_json::from_value(value).ok())
                    .ok_or_else(|| RuntimeError::InvalidPreparedRequest {
                        field: "idempotency_state".to_owned(),
                        reason: "completed idempotency state has no valid response".to_owned(),
                    })?;
                return Ok(IdempotencyAction::Replay(response));
            }
            Some("outcome_unknown") => {
                return Err(RuntimeError::IdempotencyOutcomeUnknown { operation_id });
            }
            Some("in_progress") if Instant::now() < deadline => {
                tokio::time::sleep(POLL_INTERVAL).await;
                existing = store
                    .get_value(CacheScope::Coordination, key)
                    .await
                    .map_err(|_| RuntimeError::IdempotencyStoreUnavailable {
                        operation_id: operation_id.clone(),
                    })?
                    .ok_or_else(|| RuntimeError::IdempotencyOutcomeUnknown {
                        operation_id: operation_id.clone(),
                    })?;
            }
            Some("in_progress") => {
                return Err(RuntimeError::IdempotencyInProgress { operation_id });
            }
            _ => {
                return Err(RuntimeError::InvalidPreparedRequest {
                    field: "idempotency_state".to_owned(),
                    reason: "unknown idempotency state".to_owned(),
                });
            }
        }
    }
}

pub fn policy(operation: &RuntimeOperation) -> Option<&IdempotencyPolicy> {
    let is_mutating_rest =
        matches!(&operation.target, Target::Rest(target) if target.method != HttpMethod::Get);
    if !is_mutating_rest {
        return None;
    }

    let policy = operation.execution_config.idempotency.as_ref()?;
    if policy.mode == IdempotencyMode::Disabled || policy.ttl_ms == 0 {
        return None;
    }

    Some(policy)
}

fn cache_key(
    operation: &RuntimeOperation,
    idempotency_key: &str,
    request_context: Option<&RuntimeRequestContext>,
) -> Option<String> {
    if idempotency_key.is_empty() {
        return None;
    }

    let scope = request_context?.response_cache_scope()?;
    let key_hash = URL_SAFE_NO_PAD.encode(Sha256::digest(idempotency_key.as_bytes()));

    Some(format!(
        "crank:idempotency:workspace:{}:agent:{}:operation:{}:version:{}:key:{}",
        scope.workspace_key,
        scope.agent_key,
        operation.operation_id.as_str(),
        operation.operation_version,
        key_hash
    ))
}

fn request_fingerprint(input: &Value) -> Result<String, RuntimeError> {
    let canonical = canonical_json(input);
    let encoded =
        serde_json::to_vec(&canonical).map_err(|error| RuntimeError::InvalidPreparedRequest {
            field: "idempotency_fingerprint".to_owned(),
            reason: error.to_string(),
        })?;
    Ok(URL_SAFE_NO_PAD.encode(Sha256::digest(encoded)))
}

fn canonical_json(value: &Value) -> Value {
    match value {
        Value::Object(object) => {
            let mut entries = object.iter().collect::<Vec<_>>();
            entries.sort_unstable_by_key(|(key, _)| *key);
            let mut canonical = Map::new();
            for (key, value) in entries {
                canonical.insert(key.clone(), canonical_json(value));
            }
            Value::Object(canonical)
        }
        Value::Array(values) => Value::Array(values.iter().map(canonical_json).collect()),
        _ => value.clone(),
    }
}

fn in_progress_value(fingerprint: &str) -> CoordinationStateValue {
    CoordinationStateValue {
        payload: json!({
            "state": "in_progress",
            "fingerprint": fingerprint,
            "owner": Uuid::now_v7().simple().to_string(),
        }),
    }
}

fn key_from_policy(
    policy: &IdempotencyPolicy,
    input: &Value,
    prepared_request: &PreparedRequest,
) -> Option<String> {
    if let Some(header_name) = policy
        .header_name
        .as_deref()
        .filter(|value| !value.is_empty())
        && let Some(value) = prepared_request.headers.get(header_name)
    {
        return Some(value.clone());
    }

    let field_name = policy.input_field.as_deref()?.trim();
    if field_name.is_empty() {
        return None;
    }

    input.get(field_name).and_then(|value| match value {
        Value::String(value) if !value.trim().is_empty() => Some(value.clone()),
        Value::Number(value) => Some(value.to_string()),
        _ => None,
    })
}