Replay idempotent mutation results

2026-06-20 21:19:25 +00:00
parent 79afb7df70
commit ef2912855b
3 changed files with 388 additions and 3 deletions
@@ -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,
-    MeteringEvent, ResponseCacheStore, SharedMeteringSink, SharedProtocolAdapter, Target,
+    AdapterRegistry, CachedHeader, CachedResponse, ExecutionMode, HttpMethod, IdempotencyMode,
+    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,