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

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

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

#[tokio::test]
async fn replays_mutation_result_for_same_idempotency_key() {
    let call_count = Arc::new(AtomicUsize::new(0));
    let executor = RuntimeExecutorBuilder::new()
        .register_adapter(Arc::new(CountingAdapter {
            call_count: Arc::clone(&call_count),
        }))
        .with_response_cache(Arc::new(InMemoryResponseCacheStore::default()))
        .build();
    let operation = idempotent_post_operation().into();
    let context = crank_runtime::RuntimeRequestContext::from_request_id("req_1")
        .with_response_cache_scope("workspace_1", "agent_1");
    let input = json!({ "request_id": "order-123" });

    let first = executor
        .execute_with_context(&operation, &input, Some(&context))
        .await
        .unwrap();
    let second = executor
        .execute_with_context(&operation, &input, Some(&context))
        .await
        .unwrap();

    assert_eq!(first, json!({ "result": "created-1" }));
    assert_eq!(second, first);
    assert_eq!(call_count.load(Ordering::SeqCst), 1);
}

#[tokio::test]
async fn required_idempotency_rejects_missing_key_before_adapter_call() {
    let call_count = Arc::new(AtomicUsize::new(0));
    let executor = RuntimeExecutorBuilder::new()
        .register_adapter(Arc::new(CountingAdapter {
            call_count: Arc::clone(&call_count),
        }))
        .with_response_cache(Arc::new(InMemoryResponseCacheStore::default()))
        .build();
    let mut operation: crank_runtime::RuntimeOperation = idempotent_post_operation().into();
    let policy = operation.execution_config.idempotency.as_mut().unwrap();
    policy.input_field = Some("missing_idempotency_key".to_owned());
    policy.header_name = None;

    let error = executor
        .execute(&operation, &json!({ "request_id": "shape-ok" }))
        .await
        .expect_err("missing idempotency key must fail");

    assert!(matches!(
        error,
        crank_runtime::RuntimeError::InvalidPreparedRequest { .. }
    ));
    assert_eq!(call_count.load(Ordering::SeqCst), 0);
}

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: &RuntimeRequestContext,
    ) -> Result<AdapterResponse, ProtocolAdapterError> {
        assert_eq!(
            prepared.headers.get("Idempotency-Key").map(String::as_str),
            Some("order-123")
        );
        let call_number = self.call_count.fetch_add(1, Ordering::SeqCst) + 1;
        Ok(AdapterResponse {
            status_code: 201,
            headers: BTreeMap::new(),
            body: json!({ "result": format!("created-{call_number}") }),
            data: json!({ "result": format!("created-{call_number}") }),
        })
    }
}

fn idempotent_post_operation() -> Operation<Schema, MappingSet> {
    Operation {
        id: OperationId::new("op_create_order"),
        name: "create_order".to_owned(),
        display_name: "Create 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::Post,
            path_template: "/orders".to_owned(),
            static_headers: BTreeMap::new(),
        }),
        input_schema: object_schema(BTreeMap::from([("request_id".to_owned(), string_schema())])),
        output_schema: object_schema(BTreeMap::from([("result".to_owned(), string_schema())])),
        input_mapping: MappingSet { rules: Vec::new() },
        output_mapping: MappingSet {
            rules: vec![MappingRule {
                source: "$.response.body.result".to_owned(),
                target: "$.output.result".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: Some(IdempotencyPolicy {
                mode: IdempotencyMode::Required,
                ttl_ms: 60_000,
                input_field: Some("request_id".to_owned()),
                header_name: Some("Idempotency-Key".to_owned()),
            }),
            safety: None,
            approval_policy: None,
            auth_profile_ref: None,
            headers: BTreeMap::new(),
        },
        tool_description: ToolDescription {
            title: "Create order".to_owned(),
            description: "Creates an order exactly once for a stable request id.".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(),
    }
}