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

use std::time::Duration;

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

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

pub async fn confirm_operation(
    store: Option<&dyn CoordinationStateStore>,
    operation: &RuntimeOperation,
    input: &Value,
    request_context: Option<&RuntimeRequestContext>,
) -> Result<(), RuntimeError> {
    let safety = effective_safety_policy(operation);
    if !safety.class.requires_confirmation() {
        return Ok(());
    }

    let Some(store) = store else {
        return Err(RuntimeError::ConfirmationStoreUnavailable {
            operation_id: operation.operation_id.as_str().to_owned(),
        });
    };

    let scope = confirmation_scope(operation, request_context)?;
    let input_hash = hash_json(input)?;
    let provided_token = request_context
        .and_then(RuntimeRequestContext::confirmation_token)
        .or_else(|| confirmation_token_from_input(input));

    let Some(provided_token) = provided_token else {
        let token = issue_confirmation_token(store, &scope, &input_hash, &safety).await?;
        return Err(RuntimeError::ConfirmationRequired {
            operation_id: operation.operation_id.as_str().to_owned(),
            safety_class: safety.class,
            confirmation_token: token,
            expires_in_ms: confirmation_ttl_ms(&safety),
        });
    };

    consume_confirmation_token(store, &scope, provided_token, &input_hash).await
}

fn effective_safety_policy(operation: &RuntimeOperation) -> OperationSafetyPolicy {
    operation
        .execution_config
        .safety
        .clone()
        .unwrap_or_else(|| OperationSafetyPolicy {
            class: infer_safety_class(operation),
            confirmation: None,
        })
}

fn infer_safety_class(operation: &RuntimeOperation) -> OperationSafetyClass {
    match &operation.target {
        Target::Rest(target) => match target.method {
            HttpMethod::Get => OperationSafetyClass::Read,
            HttpMethod::Delete => OperationSafetyClass::Destructive,
            HttpMethod::Post | HttpMethod::Put | HttpMethod::Patch => OperationSafetyClass::Write,
        },
    }
}

fn confirmation_ttl_ms(policy: &OperationSafetyPolicy) -> u64 {
    policy
        .confirmation
        .as_ref()
        .map(|confirmation| confirmation.ttl_ms)
        .filter(|ttl_ms| *ttl_ms > 0)
        .unwrap_or(300_000)
}

fn confirmation_scope(
    operation: &RuntimeOperation,
    request_context: Option<&RuntimeRequestContext>,
) -> Result<String, RuntimeError> {
    let Some(scope) = request_context.and_then(RuntimeRequestContext::response_cache_scope) else {
        return Err(RuntimeError::InvalidPreparedRequest {
            field: "runtime_context.response_cache_scope".to_owned(),
            reason: "confirmation requires workspace and agent scope".to_owned(),
        });
    };

    Ok(format!(
        "workspace:{}:agent:{}:operation:{}:version:{}",
        scope.workspace_key,
        scope.agent_key,
        operation.operation_id.as_str(),
        operation.operation_version
    ))
}

async fn issue_confirmation_token(
    store: &dyn CoordinationStateStore,
    scope: &str,
    input_hash: &str,
    safety: &OperationSafetyPolicy,
) -> Result<String, RuntimeError> {
    let token = format!("ct_{}", Uuid::now_v7().simple());
    let key = confirmation_cache_key(scope, &token);
    let ttl_ms = confirmation_ttl_ms(safety);
    store
        .put_value(
            CacheScope::Coordination,
            &key,
            CoordinationStateValue {
                payload: json!({
                    "scope": scope,
                    "input_hash": input_hash,
                }),
            },
            Duration::from_millis(ttl_ms),
        )
        .await
        .map_err(|error| RuntimeError::InvalidPreparedRequest {
            field: "confirmation_token".to_owned(),
            reason: error.to_string(),
        })?;
    Ok(token)
}

async fn consume_confirmation_token(
    store: &dyn CoordinationStateStore,
    operation_scope: &str,
    token: &str,
    input_hash: &str,
) -> Result<(), RuntimeError> {
    let key = confirmation_cache_key(operation_scope, token);
    let stored = store
        .get_value(CacheScope::Coordination, &key)
        .await
        .map_err(|error| RuntimeError::InvalidPreparedRequest {
            field: "confirmation_token".to_owned(),
            reason: error.to_string(),
        })?;
    let _ = store.delete_value(CacheScope::Coordination, &key).await;

    let Some(stored) = stored else {
        return Err(RuntimeError::InvalidConfirmationToken {
            operation_id: operation_id_from_scope(operation_scope),
        });
    };

    let matches_scope =
        stored.payload.get("scope").and_then(Value::as_str) == Some(operation_scope);
    let matches_input =
        stored.payload.get("input_hash").and_then(Value::as_str) == Some(input_hash);

    if matches_scope && matches_input {
        Ok(())
    } else {
        Err(RuntimeError::InvalidConfirmationToken {
            operation_id: operation_id_from_scope(operation_scope),
        })
    }
}

fn confirmation_cache_key(scope: &str, token: &str) -> String {
    let token_hash = URL_SAFE_NO_PAD.encode(Sha256::digest(token.as_bytes()));
    format!("crank:confirmation:{scope}:token:{token_hash}")
}

fn confirmation_token_from_input(input: &Value) -> Option<&str> {
    input
        .get("_crank_confirmation_token")
        .and_then(Value::as_str)
        .filter(|value| !value.trim().is_empty())
}

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

fn input_without_confirmation_token(input: &Value) -> Value {
    let Value::Object(object) = input else {
        return input.clone();
    };
    let mut sanitized = object.clone();
    sanitized.remove("_crank_confirmation_token");
    Value::Object(sanitized)
}

fn operation_id_from_scope(scope: &str) -> String {
    scope
        .split(":operation:")
        .nth(1)
        .and_then(|tail| tail.split(":version:").next())
        .unwrap_or("<unknown>")
        .to_owned()
}