Replay idempotent mutation results

crates/crank-runtime/src/executor.rs

@@ -4,8 +4,9 @@ use std::time::{Duration, Instant};
 
 use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
 use crank_core::{
-    AdapterRegistry, CachedHeader, CachedResponse, ExecutionMode, HttpMethod, InvocationStatus,
+    AdapterRegistry, CachedHeader, CachedResponse, ExecutionMode, HttpMethod, IdempotencyMode,
-    MeteringEvent, ResponseCacheStore, SharedMeteringSink, SharedProtocolAdapter, Target,
+    IdempotencyPolicy, InvocationStatus, MeteringEvent, ResponseCacheStore, SharedMeteringSink,
+    SharedProtocolAdapter, Target,
 };
 use serde_json::{Map, Value, json};
 use sha2::{Digest, Sha256};
@@ -109,7 +110,7 @@ impl RuntimeExecutor {
         let prepared_request = self.prepare_request(operation, input)?;
         let prepared_request = apply_resolved_auth(prepared_request, resolved_auth);
         let result = self
-            .execute_prepared(operation, prepared_request, request_context)
+            .execute_prepared(operation, input, prepared_request, request_context)
             .await;
         self.record_metering(operation, request_context, &result, started_at)
             .await;
@@ -134,9 +135,25 @@ impl RuntimeExecutor {
     async fn execute_prepared(
         &self,
         operation: &RuntimeOperation,
+        input: &Value,
         prepared_request: PreparedRequest,
         request_context: Option<&RuntimeRequestContext>,
     ) -> Result<Value, RuntimeError> {
+        let mut prepared_request = prepared_request;
+        let idempotency_key = self.prepare_idempotency(operation, input, &mut prepared_request)?;
+        if let Some(response) = self
+            .load_idempotent_adapter_response(
+                operation,
+                idempotency_key.as_deref(),
+                request_context,
+            )
+            .await
+        {
+            let finalized_output = finalize_output(operation, &response)?;
+            operation.output_schema.validate_shape(&finalized_output)?;
+            return Ok(finalized_output);
+        }
+
         let adapter_response = match self
             .load_cached_adapter_response(operation, &prepared_request, request_context)
             .await
@@ -153,6 +170,13 @@ impl RuntimeExecutor {
                     &adapter_response,
                 )
                 .await;
+                self.store_idempotent_adapter_response(
+                    operation,
+                    idempotency_key.as_deref(),
+                    request_context,
+                    &adapter_response,
+                )
+                .await;
                 adapter_response
             }
         };
@@ -163,6 +187,41 @@ impl RuntimeExecutor {
         Ok(finalized_output)
     }
 
+    fn prepare_idempotency(
+        &self,
+        operation: &RuntimeOperation,
+        input: &Value,
+        prepared_request: &mut PreparedRequest,
+    ) -> Result<Option<String>, RuntimeError> {
+        let Some(policy) = idempotency_policy(operation) else {
+            return Ok(None);
+        };
+
+        let key = idempotency_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))
+    }
+
     async fn record_metering<T>(
         &self,
         operation: &RuntimeOperation,
@@ -311,6 +370,74 @@ impl RuntimeExecutor {
             );
         }
     }
+
+    async fn load_idempotent_adapter_response(
+        &self,
+        operation: &RuntimeOperation,
+        idempotency_key: Option<&str>,
+        request_context: Option<&RuntimeRequestContext>,
+    ) -> Option<AdapterResponse> {
+        let response_cache = self.response_cache.as_ref()?;
+        let cache_key = idempotency_cache_key(operation, idempotency_key?, request_context)?;
+        let cached = match response_cache.get(&cache_key).await {
+            Ok(cached) => cached?,
+            Err(error) => {
+                debug!(
+                    operation_id = %operation.operation_id,
+                    cache_key,
+                    error = %error,
+                    "idempotency cache lookup skipped"
+                );
+                return None;
+            }
+        };
+
+        adapter_response_from_cached(cached).ok()
+    }
+
+    async fn store_idempotent_adapter_response(
+        &self,
+        operation: &RuntimeOperation,
+        idempotency_key: Option<&str>,
+        request_context: Option<&RuntimeRequestContext>,
+        adapter_response: &AdapterResponse,
+    ) {
+        let Some(response_cache) = self.response_cache.as_ref() else {
+            return;
+        };
+        if !(200..=299).contains(&adapter_response.status_code) {
+            return;
+        }
+        let Some(policy) = idempotency_policy(operation) else {
+            return;
+        };
+        let Some(cache_key) = idempotency_cache_key(
+            operation,
+            idempotency_key.unwrap_or_default(),
+            request_context,
+        ) else {
+            return;
+        };
+        let Some(cached_response) = cached_response_from_adapter(adapter_response) else {
+            return;
+        };
+
+        if let Err(error) = response_cache
+            .put(
+                &cache_key,
+                cached_response,
+                Duration::from_millis(policy.ttl_ms),
+            )
+            .await
+        {
+            debug!(
+                operation_id = %operation.operation_id,
+                cache_key,
+                error = %error,
+                "idempotency cache write skipped"
+            );
+        }
+    }
 }
 
 impl PreparedRequest {
@@ -490,6 +617,70 @@ fn response_cache_ttl(operation: &RuntimeOperation) -> Option<Duration> {
     Some(Duration::from_millis(policy.ttl_ms))
 }
 
+fn idempotency_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 idempotency_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())
+    {
+        if 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,
+    })
+}
+
+fn idempotency_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 response_cache_key(
     operation: &RuntimeOperation,
     prepared_request: &PreparedRequest,

crates/crank-runtime/tests/integration.rs

@@ -1,3 +1,4 @@
 mod integration {
+    mod idempotency;
     mod no_input_get;
 }

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

@@ -0,0 +1,193 @@
+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()),
+            }),
+            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(),
+    }
+}