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::{
    InMemoryCoordinationStateStore, InMemoryResponseCacheStore, RuntimeError,
    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),
            release: None,
        }))
        .with_response_cache(Arc::new(InMemoryResponseCacheStore::default()))
        .with_coordination_store(Arc::new(InMemoryCoordinationStateStore::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),
            release: None,
        }))
        .with_response_cache(Arc::new(InMemoryResponseCacheStore::default()))
        .with_coordination_store(Arc::new(InMemoryCoordinationStateStore::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);
}

#[tokio::test]
async fn required_idempotency_fails_closed_without_coordination_store() {
    let call_count = Arc::new(AtomicUsize::new(0));
    let executor = RuntimeExecutorBuilder::new()
        .register_adapter(Arc::new(CountingAdapter {
            call_count: Arc::clone(&call_count),
            release: None,
        }))
        .build();
    let operation = idempotent_post_operation().into();
    let context = crank_runtime::RuntimeRequestContext::from_request_id("req_no_store")
        .with_response_cache_scope("workspace_1", "agent_1");

    let error = executor
        .execute_with_context(
            &operation,
            &json!({ "request_id": "must-not-run" }),
            Some(&context),
        )
        .await
        .expect_err("required idempotency must not execute without an atomic store");

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

#[tokio::test]
async fn concurrent_calls_with_same_key_execute_adapter_once() {
    let call_count = Arc::new(AtomicUsize::new(0));
    let release = Arc::new(tokio::sync::Notify::new());
    let executor = RuntimeExecutorBuilder::new()
        .register_adapter(Arc::new(CountingAdapter {
            call_count: Arc::clone(&call_count),
            release: Some(Arc::clone(&release)),
        }))
        .with_coordination_store(Arc::new(InMemoryCoordinationStateStore::default()))
        .build();
    let operation: crank_runtime::RuntimeOperation = idempotent_post_operation().into();
    let context = crank_runtime::RuntimeRequestContext::from_request_id("req_concurrent")
        .with_response_cache_scope("workspace_1", "agent_1");
    let input = json!({ "request_id": "order-concurrent" });

    let first_executor = executor.clone();
    let first_operation = operation.clone();
    let first_context = context.clone();
    let first_input = input.clone();
    let first = tokio::spawn(async move {
        first_executor
            .execute_with_context(&first_operation, &first_input, Some(&first_context))
            .await
    });
    wait_for_call_count(&call_count, 1).await;

    let second_executor = executor.clone();
    let second_operation = operation.clone();
    let second_context = context.clone();
    let second_input = input.clone();
    let second = tokio::spawn(async move {
        second_executor
            .execute_with_context(&second_operation, &second_input, Some(&second_context))
            .await
    });
    for _ in 0..100 {
        tokio::task::yield_now().await;
    }
    assert_eq!(call_count.load(Ordering::SeqCst), 1);

    release.notify_waiters();
    let first_result = first.await.unwrap().unwrap();
    let second_result = second.await.unwrap().unwrap();
    assert_eq!(first_result, second_result);
    assert_eq!(call_count.load(Ordering::SeqCst), 1);
}

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

    executor
        .execute_with_context(
            &operation,
            &json!({ "request_id": "stable-key", "amount": 10 }),
            Some(&context),
        )
        .await
        .unwrap();
    let error = executor
        .execute_with_context(
            &operation,
            &json!({ "request_id": "stable-key", "amount": 20 }),
            Some(&context),
        )
        .await
        .expect_err("same key must not accept a different request fingerprint");

    assert!(matches!(error, RuntimeError::IdempotencyConflict { .. }));
    assert_eq!(call_count.load(Ordering::SeqCst), 1);
}

#[tokio::test]
async fn uncertain_adapter_failure_blocks_automatic_retry() {
    let call_count = Arc::new(AtomicUsize::new(0));
    let executor = RuntimeExecutorBuilder::new()
        .register_adapter(Arc::new(FailingAdapter {
            call_count: Arc::clone(&call_count),
        }))
        .with_coordination_store(Arc::new(InMemoryCoordinationStateStore::default()))
        .build();
    let operation = idempotent_post_operation().into();
    let context = crank_runtime::RuntimeRequestContext::from_request_id("req_unknown")
        .with_response_cache_scope("workspace_1", "agent_1");
    let input = json!({ "request_id": "uncertain-outcome" });

    let first = executor
        .execute_with_context(&operation, &input, Some(&context))
        .await
        .expect_err("adapter failure must be returned");
    assert!(matches!(first, RuntimeError::ProtocolAdapter(_)));

    let retry = executor
        .execute_with_context(&operation, &input, Some(&context))
        .await
        .expect_err("an uncertain external outcome must not be retried automatically");
    assert!(matches!(
        retry,
        RuntimeError::IdempotencyOutcomeUnknown { .. }
    ));
    assert_eq!(call_count.load(Ordering::SeqCst), 1);
}

async fn wait_for_call_count(call_count: &AtomicUsize, expected: usize) {
    for _ in 0..1_000 {
        if call_count.load(Ordering::SeqCst) >= expected {
            return;
        }
        tokio::task::yield_now().await;
    }
    panic!("adapter did not receive {expected} call(s)");
}

struct CountingAdapter {
    call_count: Arc<AtomicUsize>,
    release: Option<Arc<tokio::sync::Notify>>,
}

struct FailingAdapter {
    call_count: Arc<AtomicUsize>,
}

#[async_trait]
impl ProtocolAdapter for FailingAdapter {
    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> {
        self.call_count.fetch_add(1, Ordering::SeqCst);
        Err(ProtocolAdapterError::Message(
            "upstream outcome is unknown".to_owned(),
        ))
    }
}

#[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!(prepared.headers.contains_key("Idempotency-Key"));
        let call_number = self.call_count.fetch_add(1, Ordering::SeqCst) + 1;
        if let Some(release) = &self.release {
            release.notified().await;
        }
        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(),
    }
}