crank/crates/crank-runtime/tests/integration/confirmation.rs

use std::collections::BTreeMap;
use std::sync::{
    Arc,
    atomic::{AtomicUsize, Ordering},
};

use async_trait::async_trait;
use crank_core::{
    AdapterResponse, ConfirmationPolicy, ExecutionConfig, ExecutionMode, HttpMethod, Operation,
    OperationId, OperationSafetyClass, OperationSafetyPolicy, OperationSecurityLevel,
    OperationStatus, Protocol, ProtocolAdapter, ProtocolAdapterError, RestTarget, Target,
    ToolDescription,
};
use crank_mapping::{MappingRule, MappingSet};
use crank_runtime::{
    InMemoryCoordinationStateStore, RuntimeError, RuntimeExecutorBuilder, RuntimeRequestContext,
};
use crank_schema::{Schema, SchemaKind};
use serde_json::json;
use time::OffsetDateTime;

#[tokio::test]
async fn destructive_operation_requires_single_use_confirmation() {
    let call_count = Arc::new(AtomicUsize::new(0));
    let executor = RuntimeExecutorBuilder::new()
        .register_adapter(Arc::new(CountingAdapter {
            call_count: Arc::clone(&call_count),
        }))
        .with_coordination_store(Arc::new(InMemoryCoordinationStateStore::default()))
        .build();
    let operation = destructive_delete_operation().into();
    let context = RuntimeRequestContext::from_request_id("req_confirm")
        .with_response_cache_scope("workspace_1", "agent_1");

    let first = executor
        .execute_with_context(
            &operation,
            &json!({ "order_id": "ord_123" }),
            Some(&context),
        )
        .await
        .expect_err("first destructive call must only stage confirmation");
    let confirmation_token = match first {
        RuntimeError::ConfirmationRequired {
            confirmation_token, ..
        } => confirmation_token,
        error => panic!("expected confirmation required, got {error:?}"),
    };
    assert_eq!(call_count.load(Ordering::SeqCst), 0);

    let confirmed_context = context
        .clone()
        .with_confirmation_token(confirmation_token.clone());
    let confirmed = executor
        .execute_with_context(
            &operation,
            &json!({ "order_id": "ord_123" }),
            Some(&confirmed_context),
        )
        .await
        .unwrap();
    assert_eq!(confirmed, json!({ "deleted": true }));
    assert_eq!(call_count.load(Ordering::SeqCst), 1);

    let replay = executor
        .execute_with_context(
            &operation,
            &json!({ "order_id": "ord_123" }),
            Some(&confirmed_context),
        )
        .await
        .expect_err("confirmation token must be single-use");
    assert!(matches!(
        replay,
        RuntimeError::InvalidConfirmationToken { .. }
    ));
    assert_eq!(call_count.load(Ordering::SeqCst), 1);

    let approved_context = context.with_approval_granted();
    let approved = executor
        .execute_with_context(
            &operation,
            &json!({ "order_id": "ord_123" }),
            Some(&approved_context),
        )
        .await
        .unwrap();
    assert_eq!(approved, json!({ "deleted": true }));
    assert_eq!(call_count.load(Ordering::SeqCst), 2);
}

struct CountingAdapter {
    call_count: Arc<AtomicUsize>,
}

#[async_trait]
impl ProtocolAdapter for CountingAdapter {
    fn protocol(&self) -> Protocol {
        Protocol::Rest
    }

    fn supports_mode(&self, mode: ExecutionMode) -> bool {
        mode == ExecutionMode::Unary
    }

    async fn invoke_unary(
        &self,
        _target: &Target,
        prepared: &crank_core::PreparedRequest,
        _context: &crank_core::RuntimeRequestContext,
    ) -> Result<AdapterResponse, ProtocolAdapterError> {
        assert_eq!(
            prepared.path_params.get("order_id"),
            Some(&"ord_123".to_owned())
        );
        self.call_count.fetch_add(1, Ordering::SeqCst);
        Ok(AdapterResponse {
            status_code: 200,
            headers: BTreeMap::new(),
            body: json!({ "deleted": true }),
            data: json!({ "deleted": true }),
        })
    }
}

fn destructive_delete_operation() -> Operation<Schema, MappingSet> {
    Operation {
        id: OperationId::new("op_delete_order"),
        name: "delete_order".to_owned(),
        display_name: "Delete order".to_owned(),
        category: "orders".to_owned(),
        protocol: Protocol::Rest,
        security_level: OperationSecurityLevel::Standard,
        status: OperationStatus::Published,
        version: 1,
        target: Target::Rest(RestTarget {
            base_url: "https://api.example.invalid".to_owned(),
            method: HttpMethod::Delete,
            path_template: "/orders/{order_id}".to_owned(),
            static_headers: BTreeMap::new(),
        }),
        input_schema: object_schema(BTreeMap::from([("order_id".to_owned(), string_schema())])),
        output_schema: object_schema(BTreeMap::from([("deleted".to_owned(), bool_schema())])),
        input_mapping: MappingSet {
            rules: vec![MappingRule {
                source: "$.mcp.order_id".to_owned(),
                target: "$.request.path.order_id".to_owned(),
                required: true,
                default_value: None,
                transform: None,
                condition: None,
                notes: None,
            }],
        },
        output_mapping: MappingSet {
            rules: vec![MappingRule {
                source: "$.response.body.deleted".to_owned(),
                target: "$.output.deleted".to_owned(),
                required: true,
                default_value: None,
                transform: None,
                condition: None,
                notes: None,
            }],
        },
        execution_config: ExecutionConfig {
            timeout_ms: 1_000,
            retry_policy: None,
            response_cache: None,
            idempotency: None,
            safety: Some(OperationSafetyPolicy {
                class: OperationSafetyClass::Destructive,
                confirmation: Some(ConfirmationPolicy { ttl_ms: 60_000 }),
            }),
            approval_policy: None,
            auth_profile_ref: None,
            headers: BTreeMap::new(),
        },
        tool_description: ToolDescription {
            title: "Delete order".to_owned(),
            description: "Deletes an order after explicit confirmation.".to_owned(),
            tags: Vec::new(),
            examples: Vec::new(),
        },
        samples: None,
        generated_draft: None,
        config_export: None,
        wizard_state: None,
        created_at: OffsetDateTime::UNIX_EPOCH,
        updated_at: OffsetDateTime::UNIX_EPOCH,
        published_at: None,
    }
}

fn object_schema(fields: BTreeMap<String, Schema>) -> Schema {
    Schema {
        kind: SchemaKind::Object,
        description: None,
        required: true,
        nullable: false,
        default_value: None,
        fields,
        items: None,
        enum_values: Vec::new(),
        variants: Vec::new(),
    }
}

fn string_schema() -> Schema {
    Schema {
        kind: SchemaKind::String,
        description: None,
        required: true,
        nullable: false,
        default_value: None,
        fields: BTreeMap::new(),
        items: None,
        enum_values: Vec::new(),
        variants: Vec::new(),
    }
}

fn bool_schema() -> Schema {
    Schema {
        kind: SchemaKind::Boolean,
        description: None,
        required: true,
        nullable: false,
        default_value: None,
        fields: BTreeMap::new(),
        items: None,
        enum_values: Vec::new(),
        variants: Vec::new(),
    }
}