From ef2912855b54a0fa8015eff6bfae2022927bd957 Mon Sep 17 00:00:00 2001 From: github-ops Date: Sat, 20 Jun 2026 21:19:25 +0000 Subject: [PATCH] Replay idempotent mutation results --- crates/crank-runtime/src/executor.rs | 197 +++++++++++++++++- crates/crank-runtime/tests/integration.rs | 1 + .../tests/integration/idempotency.rs | 193 +++++++++++++++++ 3 files changed, 388 insertions(+), 3 deletions(-) create mode 100644 crates/crank-runtime/tests/integration/idempotency.rs diff --git a/crates/crank-runtime/src/executor.rs b/crates/crank-runtime/src/executor.rs index 4400eda..eca9af4 100644 --- a/crates/crank-runtime/src/executor.rs +++ b/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, - 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 { + 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, 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( &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 { + 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 { 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 { + 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 { + 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, diff --git a/crates/crank-runtime/tests/integration.rs b/crates/crank-runtime/tests/integration.rs index 9270287..ac5b022 100644 --- a/crates/crank-runtime/tests/integration.rs +++ b/crates/crank-runtime/tests/integration.rs @@ -1,3 +1,4 @@ mod integration { + mod idempotency; mod no_input_get; } diff --git a/crates/crank-runtime/tests/integration/idempotency.rs b/crates/crank-runtime/tests/integration/idempotency.rs new file mode 100644 index 0000000..1c88137 --- /dev/null +++ b/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, +} + +#[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 { + 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 { + 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) -> 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(), + } +}