Skip to main content

fakecloud_stepfunctions/
interpreter.rs

1use std::collections::HashMap;
2use std::sync::Arc;
3
4use chrono::Utc;
5use serde_json::{json, Value};
6use tracing::{debug, warn};
7
8use fakecloud_aws::arn::Arn;
9use fakecloud_core::delivery::DeliveryBus;
10use fakecloud_dynamodb::SharedDynamoDbState;
11
12use crate::choice::evaluate_choice;
13use crate::error_handling::{find_catcher, should_retry};
14use crate::io_processing::{apply_input_path, apply_output_path, apply_result_path};
15use crate::service::SharedServiceRegistry;
16use crate::state::{ExecutionStatus, HistoryEvent, SharedStepFunctionsState};
17
18/// Execute a state machine definition with the given input.
19/// Updates the execution record in shared state as it progresses.
20#[allow(clippy::too_many_arguments)]
21pub async fn execute_state_machine(
22    state: SharedStepFunctionsState,
23    execution_arn: String,
24    definition: String,
25    input: Option<String>,
26    delivery: Option<Arc<DeliveryBus>>,
27    dynamodb_state: Option<SharedDynamoDbState>,
28    registry: Option<SharedServiceRegistry>,
29    logging_configuration: Option<Value>,
30) {
31    let def: Value = match serde_json::from_str(&definition) {
32        Ok(v) => v,
33        Err(e) => {
34            fail_execution(
35                &state,
36                &execution_arn,
37                "States.Runtime",
38                &format!("Failed to parse definition: {e}"),
39            );
40            return;
41        }
42    };
43
44    let raw_input: Value = input
45        .as_deref()
46        .and_then(|s| serde_json::from_str(s).ok())
47        .unwrap_or(json!({}));
48
49    // Record ExecutionStarted event
50    add_event(
51        &state,
52        &execution_arn,
53        "ExecutionStarted",
54        0,
55        json!({
56            "input": serde_json::to_string(&raw_input).expect("serde_json::Value serialization is infallible"),
57            "roleArn": "arn:aws:iam::123456789012:role/execution-role"
58        }),
59    );
60
61    // Run the state machine inside an inner tokio::spawn so that any panic
62    // bubbles up as a JoinError instead of tearing down the caller. Without
63    // this the panic propagates through the outer spawn in `start_execution`
64    // which leaves the execution stuck in Running and leaks the panic to
65    // tokio's default hook.
66    let def_owned = def;
67    let state_clone = state.clone();
68    let execution_arn_clone = execution_arn.clone();
69    let delivery_clone = delivery.clone();
70    let dynamodb_state_clone = dynamodb_state.clone();
71    let registry_clone = registry.clone();
72    let handle = tokio::spawn(async move {
73        run_states(
74            &def_owned,
75            raw_input,
76            &delivery_clone,
77            &dynamodb_state_clone,
78            &registry_clone,
79            &state_clone,
80            &execution_arn_clone,
81        )
82        .await
83    });
84
85    match handle.await {
86        Ok(Ok(output)) => {
87            succeed_execution(&state, &execution_arn, &output);
88        }
89        Ok(Err((error, cause))) => {
90            fail_execution(&state, &execution_arn, &error, &cause);
91        }
92        Err(join_err) => {
93            let msg = if join_err.is_panic() {
94                let payload = join_err.into_panic();
95                if let Some(s) = payload.downcast_ref::<String>() {
96                    s.clone()
97                } else if let Some(s) = payload.downcast_ref::<&'static str>() {
98                    (*s).to_string()
99                } else {
100                    "execution task panicked".to_string()
101                }
102            } else {
103                format!("execution task cancelled: {join_err}")
104            };
105            tracing::error!(
106                execution_arn = %execution_arn,
107                panic = %msg,
108                "Step Functions execution panicked"
109            );
110            fail_execution(&state, &execution_arn, "States.Runtime", &msg);
111        }
112    }
113
114    // Deliver execution history to CloudWatch Logs when logging is configured.
115    deliver_execution_logs(
116        &state,
117        &execution_arn,
118        delivery.as_ref(),
119        logging_configuration.as_ref(),
120    );
121}
122
123type StatesResult<'a> = std::pin::Pin<
124    Box<dyn std::future::Future<Output = Result<Value, (String, String)>> + Send + 'a>,
125>;
126
127/// Result of executing a single state in the state machine.
128pub(crate) enum Advance {
129    /// Continue to the given state with the given input.
130    Next(String, Value),
131    /// Terminate the state machine with the given output.
132    End(Value),
133    /// Fail the state machine with the given error and cause.
134    Fail(String, String),
135}
136
137async fn run_wait_state(
138    name: &str,
139    state_def: &Value,
140    input: Value,
141    shared_state: &SharedStepFunctionsState,
142    execution_arn: &str,
143) -> Advance {
144    let entered_event_id = add_event(
145        shared_state,
146        execution_arn,
147        "WaitStateEntered",
148        0,
149        json!({
150            "name": name,
151            "input": serde_json::to_string(&input).expect("serde_json::Value serialization is infallible"),
152        }),
153    );
154
155    execute_wait_state(state_def, &input).await;
156
157    add_event(
158        shared_state,
159        execution_arn,
160        "WaitStateExited",
161        entered_event_id,
162        json!({
163            "name": name,
164            "output": serde_json::to_string(&input).expect("serde_json::Value serialization is infallible"),
165        }),
166    );
167
168    advance_from_next(state_def, input)
169}
170
171#[allow(clippy::too_many_arguments)]
172async fn run_task_state(
173    name: &str,
174    state_def: &Value,
175    input: Value,
176    delivery: &Option<Arc<DeliveryBus>>,
177    dynamodb_state: &Option<SharedDynamoDbState>,
178    registry: &Option<SharedServiceRegistry>,
179    shared_state: &SharedStepFunctionsState,
180    execution_arn: &str,
181) -> Advance {
182    let entered_event_id = add_event(
183        shared_state,
184        execution_arn,
185        "TaskStateEntered",
186        0,
187        json!({
188            "name": name,
189            "input": serde_json::to_string(&input).expect("serde_json::Value serialization is infallible"),
190        }),
191    );
192
193    let result = execute_task_state(
194        name,
195        state_def,
196        &input,
197        delivery,
198        dynamodb_state,
199        registry,
200        shared_state,
201        execution_arn,
202        entered_event_id,
203    )
204    .await;
205
206    match result {
207        Ok(output) => {
208            add_event(
209                shared_state,
210                execution_arn,
211                "TaskStateExited",
212                entered_event_id,
213                json!({
214                    "name": name,
215                    "output": serde_json::to_string(&output).expect("serde_json::Value serialization is infallible"),
216                }),
217            );
218            advance_from_next(state_def, output)
219        }
220        // Task-originated error → the States.TaskFailed wildcard applies.
221        Err((error, cause)) => advance_from_error(state_def, &input, error, cause, true),
222    }
223}
224
225#[allow(clippy::too_many_arguments)]
226async fn run_parallel_state(
227    name: &str,
228    state_def: &Value,
229    input: Value,
230    delivery: &Option<Arc<DeliveryBus>>,
231    dynamodb_state: &Option<SharedDynamoDbState>,
232    registry: &Option<SharedServiceRegistry>,
233    shared_state: &SharedStepFunctionsState,
234    execution_arn: &str,
235) -> Advance {
236    let entered_event_id = add_event(
237        shared_state,
238        execution_arn,
239        "ParallelStateEntered",
240        0,
241        json!({
242            "name": name,
243            "input": serde_json::to_string(&input).expect("serde_json::Value serialization is infallible"),
244        }),
245    );
246
247    let result = execute_parallel_state(
248        state_def,
249        &input,
250        delivery,
251        dynamodb_state,
252        registry,
253        shared_state,
254        execution_arn,
255    )
256    .await;
257
258    match result {
259        Ok(output) => {
260            add_event(
261                shared_state,
262                execution_arn,
263                "ParallelStateExited",
264                entered_event_id,
265                json!({
266                    "name": name,
267                    "output": serde_json::to_string(&output).expect("serde_json::Value serialization is infallible"),
268                }),
269            );
270            advance_from_next(state_def, output)
271        }
272        // A Parallel branch failure is not itself a task error; the
273        // States.TaskFailed wildcard must not swallow branch-level errors.
274        Err((error, cause)) => advance_from_error(state_def, &input, error, cause, false),
275    }
276}
277
278#[allow(clippy::too_many_arguments)]
279async fn run_map_state(
280    name: &str,
281    state_def: &Value,
282    input: Value,
283    delivery: &Option<Arc<DeliveryBus>>,
284    dynamodb_state: &Option<SharedDynamoDbState>,
285    registry: &Option<SharedServiceRegistry>,
286    shared_state: &SharedStepFunctionsState,
287    execution_arn: &str,
288) -> Advance {
289    let entered_event_id = add_event(
290        shared_state,
291        execution_arn,
292        "MapStateEntered",
293        0,
294        json!({
295            "name": name,
296            "input": serde_json::to_string(&input).expect("serde_json::Value serialization is infallible"),
297        }),
298    );
299
300    let result = execute_map_state(
301        state_def,
302        &input,
303        delivery,
304        dynamodb_state,
305        registry,
306        shared_state,
307        execution_arn,
308    )
309    .await;
310
311    match result {
312        Ok(output) => {
313            add_event(
314                shared_state,
315                execution_arn,
316                "MapStateExited",
317                entered_event_id,
318                json!({
319                    "name": name,
320                    "output": serde_json::to_string(&output).expect("serde_json::Value serialization is infallible"),
321                }),
322            );
323            advance_from_next(state_def, output)
324        }
325        // A Map iteration failure is not itself a task error; the
326        // States.TaskFailed wildcard must not swallow iteration-level errors.
327        Err((error, cause)) => advance_from_error(state_def, &input, error, cause, false),
328    }
329}
330
331/// Execute a Wait state: pause execution for a specified duration or until a timestamp.
332async fn execute_wait_state(state_def: &Value, input: &Value) {
333    if let Some(seconds) = state_def["Seconds"].as_u64() {
334        tokio::time::sleep(tokio::time::Duration::from_secs(seconds)).await;
335        return;
336    }
337
338    if let Some(path) = state_def["SecondsPath"].as_str() {
339        let val = crate::io_processing::resolve_path(input, path);
340        if let Some(seconds) = val.as_u64() {
341            tokio::time::sleep(tokio::time::Duration::from_secs(seconds)).await;
342        }
343        return;
344    }
345
346    if let Some(ts_str) = state_def["Timestamp"].as_str() {
347        if let Ok(target) = chrono::DateTime::parse_from_rfc3339(ts_str) {
348            let now = Utc::now();
349            let target_utc = target.with_timezone(&chrono::Utc);
350            if target_utc > now {
351                let duration = (target_utc - now).to_std().unwrap_or_default();
352                tokio::time::sleep(duration).await;
353            }
354        }
355        return;
356    }
357
358    if let Some(path) = state_def["TimestampPath"].as_str() {
359        let val = crate::io_processing::resolve_path(input, path);
360        if let Some(ts_str) = val.as_str() {
361            if let Ok(target) = chrono::DateTime::parse_from_rfc3339(ts_str) {
362                let now = Utc::now();
363                let target_utc = target.with_timezone(&chrono::Utc);
364                if target_utc > now {
365                    let duration = (target_utc - now).to_std().unwrap_or_default();
366                    tokio::time::sleep(duration).await;
367                }
368            }
369        }
370        return;
371    }
372
373    warn!(
374        "Wait state has no valid Seconds, SecondsPath, Timestamp, or TimestampPath — skipping wait"
375    );
376}
377
378/// Execute a Task state: invoke the resource (Lambda, SQS, SNS, EventBridge, DynamoDB),
379/// apply I/O processing, handle Retry.
380#[allow(clippy::too_many_arguments)]
381async fn execute_task_state(
382    name: &str,
383    state_def: &Value,
384    input: &Value,
385    delivery: &Option<Arc<DeliveryBus>>,
386    dynamodb_state: &Option<SharedDynamoDbState>,
387    registry: &Option<SharedServiceRegistry>,
388    shared_state: &SharedStepFunctionsState,
389    execution_arn: &str,
390    entered_event_id: i64,
391) -> Result<Value, (String, String)> {
392    let resource = state_def["Resource"].as_str().unwrap_or("").to_string();
393
394    let input_path = state_def["InputPath"].as_str();
395    let result_path = state_def["ResultPath"].as_str();
396    let output_path = state_def["OutputPath"].as_str();
397
398    let effective_input = if input_path == Some("null") {
399        json!({})
400    } else {
401        apply_input_path(input, input_path)
402    };
403
404    let retriers = state_def["Retry"].as_array().cloned().unwrap_or_default();
405    let timeout_seconds = state_def["TimeoutSeconds"].as_u64();
406    let heartbeat_seconds = state_def["HeartbeatSeconds"].as_u64();
407    let mut attempt = 0u32;
408
409    let is_wait_for_task_token = resource.contains(".waitForTaskToken");
410    let task_token = if is_wait_for_task_token {
411        let token = format!(
412            "FCToken-{}-{}",
413            chrono::Utc::now().timestamp_nanos_opt().unwrap_or(0),
414            uuid::Uuid::new_v4().simple(),
415        );
416        let account_id = account_id_from_arn(execution_arn);
417        let context = json!({
418            "Task": { "Token": token.clone() },
419            "Execution": { "Id": execution_arn },
420            "State": { "Name": name },
421        });
422        {
423            let mut accounts = shared_state.write();
424            let state = accounts.get_or_create(account_id);
425            state.task_tokens.insert(
426                token.clone(),
427                crate::state::TaskTokenState {
428                    activity_arn: String::new(),
429                    status: "PENDING".to_string(),
430                    output: None,
431                    error: None,
432                    cause: None,
433                    input: None,
434                    created_at: chrono::Utc::now(),
435                    last_heartbeat_at: None,
436                    heartbeat_seconds: heartbeat_seconds.map(|s| s as i64),
437                    timeout_seconds: timeout_seconds.map(|s| s as i64),
438                },
439            );
440        }
441        Some((token, context))
442    } else {
443        None
444    };
445
446    let task_input = if let Some(params) = state_def.get("Parameters") {
447        if let Some((_, ctx)) = &task_token {
448            apply_parameters(params, &effective_input, Some(ctx))
449        } else {
450            apply_parameters(params, &effective_input, None)
451        }
452    } else {
453        effective_input
454    };
455
456    loop {
457        add_event(
458            shared_state,
459            execution_arn,
460            "TaskScheduled",
461            entered_event_id,
462            json!({
463                "resource": resource,
464                "region": "us-east-1",
465                "parameters": serde_json::to_string(&task_input).expect("serde_json::Value serialization is infallible"),
466            }),
467        );
468
469        add_event(
470            shared_state,
471            execution_arn,
472            "TaskStarted",
473            entered_event_id,
474            json!({ "resource": resource }),
475        );
476
477        let invoke_result = invoke_resource(
478            &resource,
479            &task_input,
480            delivery,
481            dynamodb_state,
482            registry,
483            execution_arn,
484            timeout_seconds,
485            heartbeat_seconds,
486            shared_state,
487        )
488        .await;
489
490        match invoke_result {
491            Ok(result) => {
492                if let Some((token, _)) = &task_token {
493                    let account_id = account_id_from_arn(execution_arn);
494                    match poll_task_token(
495                        shared_state,
496                        account_id,
497                        token,
498                        timeout_seconds,
499                        heartbeat_seconds,
500                    )
501                    .await
502                    {
503                        Ok(output) => {
504                            add_event(
505                                shared_state,
506                                execution_arn,
507                                "TaskSucceeded",
508                                entered_event_id,
509                                json!({
510                                    "resource": resource,
511                                    "output": serde_json::to_string(&output).expect("serde_json::Value serialization is infallible"),
512                                }),
513                            );
514
515                            let selected = if let Some(selector) = state_def.get("ResultSelector") {
516                                apply_parameters(selector, &output, None)
517                            } else {
518                                output
519                            };
520
521                            let after_result = if result_path == Some("null") {
522                                input.clone()
523                            } else {
524                                apply_result_path(input, &selected, result_path)
525                            };
526
527                            let output = if output_path == Some("null") {
528                                json!({})
529                            } else {
530                                apply_output_path(&after_result, output_path)
531                            };
532
533                            return Ok(output);
534                        }
535                        Err((error, cause)) => {
536                            add_event(
537                                shared_state,
538                                execution_arn,
539                                "TaskFailed",
540                                entered_event_id,
541                                json!({ "error": error, "cause": cause }),
542                            );
543
544                            if let Some(delay_ms) = should_retry(&retriers, &error, attempt, true) {
545                                attempt += 1;
546                                // Honour the ASL-computed backoff (already bounded
547                                // by MaxDelaySeconds inside `should_retry`).
548                                tokio::time::sleep(tokio::time::Duration::from_millis(delay_ms))
549                                    .await;
550                                continue;
551                            }
552
553                            return Err((error, cause));
554                        }
555                    }
556                }
557
558                add_event(
559                    shared_state,
560                    execution_arn,
561                    "TaskSucceeded",
562                    entered_event_id,
563                    json!({
564                        "resource": resource,
565                        "output": serde_json::to_string(&result).expect("serde_json::Value serialization is infallible"),
566                    }),
567                );
568
569                let selected = if let Some(selector) = state_def.get("ResultSelector") {
570                    apply_parameters(selector, &result, None)
571                } else {
572                    result
573                };
574
575                let after_result = if result_path == Some("null") {
576                    input.clone()
577                } else {
578                    apply_result_path(input, &selected, result_path)
579                };
580
581                let output = if output_path == Some("null") {
582                    json!({})
583                } else {
584                    apply_output_path(&after_result, output_path)
585                };
586
587                return Ok(output);
588            }
589            Err((error, cause)) => {
590                add_event(
591                    shared_state,
592                    execution_arn,
593                    "TaskFailed",
594                    entered_event_id,
595                    json!({ "error": error, "cause": cause }),
596                );
597
598                if let Some(delay_ms) = should_retry(&retriers, &error, attempt, true) {
599                    attempt += 1;
600                    // Honour the ASL-computed backoff (already bounded by
601                    // MaxDelaySeconds inside `should_retry`).
602                    tokio::time::sleep(tokio::time::Duration::from_millis(delay_ms)).await;
603                    continue;
604                }
605
606                return Err((error, cause));
607            }
608        }
609    }
610}
611
612/// Execute a Parallel state: run all branches concurrently, collect results into an array.
613async fn execute_parallel_state(
614    state_def: &Value,
615    input: &Value,
616    delivery: &Option<Arc<DeliveryBus>>,
617    dynamodb_state: &Option<SharedDynamoDbState>,
618    registry: &Option<SharedServiceRegistry>,
619    shared_state: &SharedStepFunctionsState,
620    execution_arn: &str,
621) -> Result<Value, (String, String)> {
622    let input_path = state_def["InputPath"].as_str();
623    let result_path = state_def["ResultPath"].as_str();
624    let output_path = state_def["OutputPath"].as_str();
625
626    let effective_input = if input_path == Some("null") {
627        json!({})
628    } else {
629        apply_input_path(input, input_path)
630    };
631
632    let branches = state_def["Branches"]
633        .as_array()
634        .cloned()
635        .unwrap_or_default();
636
637    if branches.is_empty() {
638        return Err((
639            "States.Runtime".to_string(),
640            "Parallel state has no Branches".to_string(),
641        ));
642    }
643
644    // Spawn all branches concurrently
645    let mut handles = Vec::new();
646    for branch_def in &branches {
647        let branch = branch_def.clone();
648        let branch_input = effective_input.clone();
649        let delivery = delivery.clone();
650        let ddb = dynamodb_state.clone();
651        let reg = registry.clone();
652        let state = shared_state.clone();
653        let arn = execution_arn.to_string();
654
655        handles.push(tokio::spawn(async move {
656            run_states(&branch, branch_input, &delivery, &ddb, &reg, &state, &arn).await
657        }));
658    }
659
660    // Collect results in order
661    let mut results = Vec::with_capacity(handles.len());
662    for handle in handles {
663        let result = handle.await.map_err(|e| {
664            (
665                "States.Runtime".to_string(),
666                format!("Branch execution panicked: {e}"),
667            )
668        })??;
669        results.push(result);
670    }
671
672    let branch_output = Value::Array(results);
673
674    // Apply ResultSelector if present
675    let selected = if let Some(selector) = state_def.get("ResultSelector") {
676        apply_parameters(selector, &branch_output, None)
677    } else {
678        branch_output
679    };
680
681    // Apply ResultPath
682    let after_result = if result_path == Some("null") {
683        input.clone()
684    } else {
685        apply_result_path(input, &selected, result_path)
686    };
687
688    // Apply OutputPath
689    let output = if output_path == Some("null") {
690        json!({})
691    } else {
692        apply_output_path(&after_result, output_path)
693    };
694
695    Ok(output)
696}
697
698/// Execute a Map state: iterate over an array and run a sub-state machine per item.
699/// Supports distributed mode: ItemReader (S3), ItemBatcher, ResultWriter (S3),
700/// ToleratedFailurePercentage, and MaxConcurrencyPath.
701#[allow(clippy::too_many_arguments)]
702async fn execute_map_state(
703    state_def: &Value,
704    input: &Value,
705    delivery: &Option<Arc<DeliveryBus>>,
706    dynamodb_state: &Option<SharedDynamoDbState>,
707    registry: &Option<SharedServiceRegistry>,
708    shared_state: &SharedStepFunctionsState,
709    execution_arn: &str,
710) -> Result<Value, (String, String)> {
711    let input_path = state_def["InputPath"].as_str();
712    let result_path = state_def["ResultPath"].as_str();
713    let output_path = state_def["OutputPath"].as_str();
714
715    let effective_input = if input_path == Some("null") {
716        json!({})
717    } else {
718        apply_input_path(input, input_path)
719    };
720
721    // Resolve MaxConcurrencyPath if present
722    let max_concurrency = if let Some(path) = state_def["MaxConcurrencyPath"].as_str() {
723        crate::io_processing::resolve_path(&effective_input, path)
724            .as_u64()
725            .unwrap_or(0)
726    } else {
727        state_def["MaxConcurrency"].as_u64().unwrap_or(0)
728    };
729    let effective_concurrency = if max_concurrency == 0 {
730        40
731    } else {
732        max_concurrency as usize
733    };
734
735    // Read items from ItemReader (S3) or ItemsPath
736    let items = if let Some(item_reader) = state_def.get("ItemReader") {
737        read_items_from_s3(item_reader, registry, execution_arn).await?
738    } else {
739        let items_path = state_def["ItemsPath"].as_str().unwrap_or("$");
740        let items_value = crate::io_processing::resolve_path(&effective_input, items_path);
741        items_value.as_array().cloned().unwrap_or_default()
742    };
743
744    // Apply ItemBatcher if present
745    let batch_config = state_def.get("ItemBatcher").cloned();
746    let batched_items = if let Some(ref batcher) = batch_config {
747        apply_item_batcher(&items, batcher, &effective_input)
748    } else {
749        items
750    };
751
752    // Get the iterator definition (ItemProcessor or Iterator for backwards compat)
753    let iterator_def = state_def
754        .get("ItemProcessor")
755        .or_else(|| state_def.get("Iterator"))
756        .cloned()
757        .ok_or_else(|| {
758            (
759                "States.Runtime".to_string(),
760                "Map state has no ItemProcessor or Iterator".to_string(),
761            )
762        })?;
763
764    let tolerated_failure_percentage = state_def["ToleratedFailurePercentage"]
765        .as_f64()
766        .unwrap_or(0.0);
767    let total_items = batched_items.len() as f64;
768    let mut failure_count = 0usize;
769
770    let semaphore = Arc::new(tokio::sync::Semaphore::new(effective_concurrency));
771
772    // Process all items
773    let mut handles = Vec::new();
774    for (index, batch_item) in batched_items.into_iter().enumerate() {
775        let iter_def = iterator_def.clone();
776        let delivery = delivery.clone();
777        let ddb = dynamodb_state.clone();
778        let reg = registry.clone();
779        let state = shared_state.clone();
780        let arn = execution_arn.to_string();
781        let sem = semaphore.clone();
782
783        // Apply ItemSelector if present
784        let item_input = if let Some(selector) = state_def.get("ItemSelector") {
785            let mut ctx = serde_json::Map::new();
786            ctx.insert("value".to_string(), batch_item.clone());
787            ctx.insert("index".to_string(), json!(index));
788            apply_parameters(selector, &Value::Object(ctx), None)
789        } else {
790            batch_item
791        };
792
793        add_event(
794            shared_state,
795            execution_arn,
796            "MapIterationStarted",
797            0,
798            json!({ "index": index }),
799        );
800
801        handles.push(tokio::spawn(async move {
802            let _permit = sem
803                .acquire()
804                .await
805                .map_err(|_| ("States.Runtime".to_string(), "Semaphore closed".to_string()))?;
806            let result =
807                run_states(&iter_def, item_input, &delivery, &ddb, &reg, &state, &arn).await;
808            Ok::<(usize, Result<Value, (String, String)>), (String, String)>((index, result))
809        }));
810    }
811
812    // Collect results in order
813    let mut results: Vec<(usize, Value)> = Vec::with_capacity(handles.len());
814    for handle in handles {
815        let (index, result) = handle.await.map_err(|e| {
816            (
817                "States.Runtime".to_string(),
818                format!("Map iteration panicked: {e}"),
819            )
820        })??;
821
822        match result {
823            Ok(output) => {
824                add_event(
825                    shared_state,
826                    execution_arn,
827                    "MapIterationSucceeded",
828                    0,
829                    json!({ "index": index }),
830                );
831                results.push((index, output));
832            }
833            Err((error, cause)) => {
834                add_event(
835                    shared_state,
836                    execution_arn,
837                    "MapIterationFailed",
838                    0,
839                    json!({ "index": index, "error": error }),
840                );
841                failure_count += 1;
842                let failure_percentage = (failure_count as f64 / total_items) * 100.0;
843                if failure_percentage > tolerated_failure_percentage {
844                    return Err((error, cause));
845                }
846            }
847        }
848    }
849
850    // Sort by index to maintain order
851    results.sort_by_key(|(i, _)| *i);
852    let map_output = Value::Array(results.into_iter().map(|(_, v)| v).collect());
853
854    // Write results to S3 if ResultWriter is configured
855    if let Some(result_writer) = state_def.get("ResultWriter") {
856        write_map_results_to_s3(result_writer, registry, execution_arn, &map_output).await?;
857    }
858
859    // Apply ResultSelector if present
860    let selected = if let Some(selector) = state_def.get("ResultSelector") {
861        apply_parameters(selector, &map_output, None)
862    } else {
863        map_output
864    };
865
866    // Apply ResultPath
867    let after_result = if result_path == Some("null") {
868        input.clone()
869    } else {
870        apply_result_path(input, &selected, result_path)
871    };
872
873    // Apply OutputPath
874    let output = if output_path == Some("null") {
875        json!({})
876    } else {
877        apply_output_path(&after_result, output_path)
878    };
879
880    Ok(output)
881}
882
883/// Read items from S3 for distributed Map mode ItemReader.
884async fn read_items_from_s3(
885    item_reader: &Value,
886    registry: &Option<SharedServiceRegistry>,
887    execution_arn: &str,
888) -> Result<Vec<Value>, (String, String)> {
889    let resource = item_reader["Resource"]
890        .as_str()
891        .unwrap_or("arn:aws:states:::s3:getObject");
892    if !resource.contains("s3:getObject") {
893        return Err((
894            "States.Runtime".to_string(),
895            format!("ItemReader unsupported resource: {resource}"),
896        ));
897    }
898
899    let params = item_reader
900        .get("Parameters")
901        .cloned()
902        .unwrap_or_else(|| json!({}));
903    let bucket = params["Bucket"].as_str().ok_or_else(|| {
904        (
905            "States.Runtime".to_string(),
906            "ItemReader missing Bucket".to_string(),
907        )
908    })?;
909    let key = params["Key"].as_str().ok_or_else(|| {
910        (
911            "States.Runtime".to_string(),
912            "ItemReader missing Key".to_string(),
913        )
914    })?;
915
916    let registry_arc = resolve_registry(registry)?;
917    let account_id = account_from_execution_arn(execution_arn);
918
919    let body = call_sdk_action_raw_bytes(
920        &registry_arc,
921        "s3",
922        "GetObject",
923        &json!({ "Bucket": bucket, "Key": key }),
924        &account_id,
925    )
926    .await?;
927
928    // Parse JSON array from the S3 object body
929    let parsed: Value = serde_json::from_slice(&body).map_err(|e| {
930        (
931            "States.Runtime".to_string(),
932            format!("ItemReader failed to parse S3 object as JSON: {e}"),
933        )
934    })?;
935
936    parsed.as_array().cloned().ok_or_else(|| {
937        (
938            "States.Runtime".to_string(),
939            "ItemReader S3 object is not a JSON array".to_string(),
940        )
941    })
942}
943
944/// Write Map results to S3 for distributed Map mode ResultWriter.
945async fn write_map_results_to_s3(
946    result_writer: &Value,
947    registry: &Option<SharedServiceRegistry>,
948    execution_arn: &str,
949    results: &Value,
950) -> Result<(), (String, String)> {
951    let resource = result_writer["Resource"]
952        .as_str()
953        .unwrap_or("arn:aws:states:::s3:putObject");
954    if !resource.contains("s3:putObject") {
955        return Err((
956            "States.Runtime".to_string(),
957            format!("ResultWriter unsupported resource: {resource}"),
958        ));
959    }
960
961    let params = result_writer
962        .get("Parameters")
963        .cloned()
964        .unwrap_or_else(|| json!({}));
965    let bucket = params["Bucket"].as_str().ok_or_else(|| {
966        (
967            "States.Runtime".to_string(),
968            "ResultWriter missing Bucket".to_string(),
969        )
970    })?;
971    let prefix = params["Prefix"].as_str().unwrap_or("map-results/");
972
973    let registry_arc = resolve_registry(registry)?;
974    let account_id = account_from_execution_arn(execution_arn);
975
976    use bytes::Bytes;
977    let body = Bytes::from(
978        serde_json::to_vec(results).expect("serde_json::Value serialization is infallible"),
979    );
980
981    // Build a raw S3 PutObject request with the JSON body
982    use fakecloud_core::service::AwsRequest;
983    use http::{HeaderMap, Method};
984    let service = registry_arc.get("s3").ok_or_else(|| {
985        (
986            "States.TaskFailed".to_string(),
987            "S3 service not available for ResultWriter".to_string(),
988        )
989    })?;
990
991    let req = AwsRequest {
992        service: "s3".to_string(),
993        action: "PutObject".to_string(),
994        region: "us-east-1".to_string(),
995        account_id: account_id.to_string(),
996        request_id: uuid::Uuid::new_v4().to_string(),
997        headers: HeaderMap::new(),
998        query_params: std::collections::HashMap::new(),
999        body,
1000        body_stream: parking_lot::Mutex::new(None),
1001        path_segments: vec![bucket.to_string(), format!("{prefix}result.json")],
1002        raw_path: format!("/{bucket}/{prefix}result.json"),
1003        raw_query: String::new(),
1004        method: Method::PUT,
1005        is_query_protocol: false,
1006        access_key_id: None,
1007        principal: None,
1008    };
1009
1010    service.handle(req).await.map_err(|err| {
1011        let code = err.code().to_string();
1012        let msg = err.message();
1013        (
1014            format!("S3.{code}"),
1015            format!("ResultWriter PutObject failed: {msg}"),
1016        )
1017    })?;
1018
1019    Ok(())
1020}
1021
1022/// Apply ItemBatcher configuration to group items into batches.
1023fn apply_item_batcher(items: &[Value], batcher: &Value, _effective_input: &Value) -> Vec<Value> {
1024    let max_per_batch = batcher["MaxItemsPerBatch"].as_u64().unwrap_or(u64::MAX) as usize;
1025    let max_bytes = batcher["MaxInputBytesPerBatch"].as_u64().unwrap_or(0) as usize;
1026    let batch_input = batcher.get("BatchInput").cloned();
1027
1028    let mut batches: Vec<Vec<Value>> = Vec::new();
1029    let mut current_batch: Vec<Value> = Vec::new();
1030    let mut current_bytes = 0usize;
1031
1032    for item in items.iter().cloned() {
1033        let item_bytes = serde_json::to_vec(&item).unwrap_or_default().len();
1034        if !current_batch.is_empty()
1035            && (current_batch.len() >= max_per_batch
1036                || (max_bytes > 0 && current_bytes + item_bytes > max_bytes))
1037        {
1038            batches.push(current_batch);
1039            current_batch = Vec::new();
1040            current_bytes = 0;
1041        }
1042        current_bytes += item_bytes;
1043        current_batch.push(item);
1044    }
1045    if !current_batch.is_empty() {
1046        batches.push(current_batch);
1047    }
1048
1049    batches
1050        .into_iter()
1051        .enumerate()
1052        .map(|(index, batch)| {
1053            let mut map = serde_json::Map::new();
1054            map.insert("index".to_string(), json!(index));
1055            map.insert("items".to_string(), Value::Array(batch));
1056            if let Some(Value::Object(ref obj)) = batch_input {
1057                for (k, v) in obj {
1058                    map.insert(k.clone(), v.clone());
1059                }
1060            }
1061            Value::Object(map)
1062        })
1063        .collect()
1064}
1065
1066/// Invoke a resource (Lambda function or SDK integration).
1067#[allow(clippy::too_many_arguments)]
1068async fn invoke_resource(
1069    resource: &str,
1070    input: &Value,
1071    delivery: &Option<Arc<DeliveryBus>>,
1072    dynamodb_state: &Option<SharedDynamoDbState>,
1073    registry: &Option<SharedServiceRegistry>,
1074    execution_arn: &str,
1075    timeout_seconds: Option<u64>,
1076    heartbeat_seconds: Option<u64>,
1077    shared_state: &SharedStepFunctionsState,
1078) -> Result<Value, (String, String)> {
1079    // Direct activity ARN: arn:aws:states:<region>:<account>:activity:<name>
1080    if resource.contains(":states:") && resource.contains(":activity:") {
1081        return invoke_activity(
1082            resource,
1083            input,
1084            shared_state,
1085            timeout_seconds,
1086            heartbeat_seconds,
1087        )
1088        .await;
1089    }
1090
1091    // Direct Lambda ARN: arn:aws:lambda:<region>:<account>:function:<name>
1092    if resource.contains(":lambda:") && resource.contains(":function:") {
1093        return invoke_lambda_direct(resource, input, delivery, timeout_seconds).await;
1094    }
1095
1096    // SDK integration patterns: arn:aws:states:::<service>:<action>
1097    if resource.starts_with("arn:aws:states:::lambda:invoke") {
1098        let function_name = input["FunctionName"].as_str().unwrap_or("");
1099        let payload = if let Some(p) = input.get("Payload") {
1100            p.clone()
1101        } else {
1102            input.clone()
1103        };
1104        // The optimized `lambda:invoke` integration returns the AWS Invoke API
1105        // response envelope, not the bare function output — `ResultSelector`
1106        // expressions like `{"value.$": "$.Payload"}` depend on it. Wrap only
1107        // the successful result; error results (function errors and transport
1108        // failures) pass through unchanged, since a failed lambda:invoke task
1109        // has no result envelope on AWS. The direct-ARN path above
1110        // intentionally stays unwrapped — there real AWS returns the bare
1111        // payload. `SdkHttpMetadata`/`SdkResponseMetadata` are omitted; real
1112        // templates select `Payload`/`StatusCode`.
1113        return invoke_lambda_direct(function_name, &payload, delivery, timeout_seconds)
1114            .await
1115            .map(|payload| {
1116                json!({
1117                    "ExecutedVersion": "$LATEST",
1118                    "Payload": payload,
1119                    "StatusCode": 200,
1120                })
1121            });
1122    }
1123
1124    if resource.starts_with("arn:aws:states:::sqs:sendMessage") {
1125        return invoke_sqs_send_message(input, delivery);
1126    }
1127
1128    if resource.starts_with("arn:aws:states:::sns:publish") {
1129        return invoke_sns_publish(input, delivery);
1130    }
1131
1132    if resource.starts_with("arn:aws:states:::events:putEvents") {
1133        return invoke_eventbridge_put_events(input, delivery);
1134    }
1135
1136    if resource.starts_with("arn:aws:states:::dynamodb:getItem") {
1137        return invoke_dynamodb_get_item(input, dynamodb_state);
1138    }
1139
1140    if resource.starts_with("arn:aws:states:::dynamodb:putItem") {
1141        return invoke_dynamodb_put_item(input, dynamodb_state);
1142    }
1143
1144    if resource.starts_with("arn:aws:states:::dynamodb:deleteItem") {
1145        return invoke_dynamodb_delete_item(input, dynamodb_state);
1146    }
1147
1148    if resource.starts_with("arn:aws:states:::dynamodb:updateItem") {
1149        return invoke_dynamodb_update_item(input, dynamodb_state);
1150    }
1151
1152    // Nested Step Functions execution:
1153    //   arn:aws:states:::states:startExecution[.sync|.waitForTaskToken]
1154    // Routes through the same registry path as the generic aws-sdk
1155    // integration so the inner execution sees a real `StartExecution`
1156    // call. `.sync` polls `DescribeExecution` until the inner execution
1157    // reaches a terminal state, mirroring AWS' behavior of bubbling the
1158    // inner Output back to the parent.
1159    if let Some(tail) = resource.strip_prefix("arn:aws:states:::") {
1160        if tail.starts_with("states:startExecution") {
1161            let account_id = account_from_execution_arn(execution_arn);
1162            let result =
1163                invoke_aws_sdk_integration(tail, input, registry, &account_id, timeout_seconds)
1164                    .await;
1165            // Patch the inner execution's parent_execution_arn so the
1166            // `/execution-tree` introspection can walk back-references.
1167            // For `.sync`, the result is the DescribeExecution shape
1168            // (has `executionArn`); for fire-and-forget StartExecution the
1169            // initial StartExecution response also carries `executionArn`.
1170            if let Ok(ref value) = result {
1171                if let Some(inner_arn) = value
1172                    .get("executionArn")
1173                    .or_else(|| value.get("ExecutionArn"))
1174                    .and_then(Value::as_str)
1175                {
1176                    let mut accounts = shared_state.write();
1177                    if let Some(state) = accounts.get_mut(&account_id) {
1178                        if let Some(exec) = state.executions.get_mut(inner_arn) {
1179                            exec.parent_execution_arn = Some(execution_arn.to_string());
1180                        }
1181                    }
1182                }
1183            }
1184            return result;
1185        }
1186    }
1187
1188    // Generic AWS SDK integration: arn:aws:states:::aws-sdk:<service>:<action>[.<wait>]
1189    // Routes the call to the registered service via the central
1190    // ServiceRegistry, passing the Task's `Parameters` block as the
1191    // request body. Mirrors the AWS SDK service integration pattern in
1192    // real Step Functions.
1193    if let Some(rest) = resource.strip_prefix("arn:aws:states:::aws-sdk:") {
1194        let account_id = account_from_execution_arn(execution_arn);
1195        return invoke_aws_sdk_integration(rest, input, registry, &account_id, timeout_seconds)
1196            .await;
1197    }
1198
1199    // Optimized service integrations expose `.sync` variants for ECS,
1200    // Athena, and Glue. Route them through the same waiter machinery as
1201    // `aws-sdk:` so callers can write the AWS-blessed ARN forms.
1202    if let Some(tail) = resource.strip_prefix("arn:aws:states:::") {
1203        if tail.contains(".sync") {
1204            let account_id = account_from_execution_arn(execution_arn);
1205            return invoke_aws_sdk_integration(tail, input, registry, &account_id, timeout_seconds)
1206                .await;
1207        }
1208    }
1209
1210    Err((
1211        "States.TaskFailed".to_string(),
1212        format!("Unsupported resource: {resource}"),
1213    ))
1214}
1215
1216/// Convert an SDK integration action (camelCase, e.g. `getItem`) to its
1217/// PascalCase wire form (`GetItem`). Step Functions Task ARNs use
1218/// camelCase, but the underlying AWS service handlers expect PascalCase
1219/// `X-Amz-Target`-style action names.
1220fn camel_to_pascal(action: &str) -> String {
1221    let mut chars = action.chars();
1222    match chars.next() {
1223        None => String::new(),
1224        Some(first) => first.to_ascii_uppercase().to_string() + chars.as_str(),
1225    }
1226}
1227
1228/// Map a Step Functions SDK integration service id to the corresponding
1229/// fakecloud `service_name()`. Most match 1:1, but a handful of AWS SDK
1230/// service ids differ from the SigV4 service identifier we register
1231/// services under (e.g. `sfn` -> `states`, `cloudwatchlogs` -> `logs`).
1232fn map_sdk_service_id(service_id: &str) -> &str {
1233    match service_id {
1234        "sfn" => "states",
1235        "cloudwatchlogs" => "logs",
1236        // Default: pass through unchanged.
1237        other => other,
1238    }
1239}
1240
1241/// Extract the AWS account id from a Step Functions execution ARN
1242/// (`arn:aws:states:<region>:<account>:execution:...`). Falls back to
1243/// the AWS-conventional fixture id if the ARN is malformed.
1244fn account_from_execution_arn(execution_arn: &str) -> String {
1245    execution_arn
1246        .split(':')
1247        .nth(4)
1248        .filter(|s| !s.is_empty())
1249        .unwrap_or("123456789012")
1250        .to_string()
1251}
1252
1253/// Dispatch a Step Functions `aws-sdk:<service>:<action>` Task to the
1254/// registered service via the central [`fakecloud_core::registry::ServiceRegistry`].
1255/// `tail` is the trailing portion of the resource ARN after the
1256/// `aws-sdk:` prefix (e.g. `dynamodb:getItem` or `sqs:sendMessage.waitForTaskToken`).
1257async fn invoke_aws_sdk_integration(
1258    tail: &str,
1259    input: &Value,
1260    registry: &Option<SharedServiceRegistry>,
1261    account_id: &str,
1262    timeout_seconds: Option<u64>,
1263) -> Result<Value, (String, String)> {
1264    let registry_arc = resolve_registry(registry)?;
1265
1266    // Split `<service>:<action>[.<modifier>]`. The `.waitForTaskToken`
1267    // modifier is accepted but ignored — the integration runs synchronously
1268    // regardless. The `.sync` modifier triggers a polling loop after the
1269    // initial call to wait for the downstream operation to reach a terminal
1270    // state.
1271    let mut parts = tail.splitn(2, ':');
1272    let service_id = parts.next().filter(|s| !s.is_empty()).ok_or_else(|| {
1273        (
1274            "States.TaskFailed".to_string(),
1275            format!("Invalid aws-sdk Resource ARN: missing service in '{tail}'"),
1276        )
1277    })?;
1278    let action_with_mod = parts.next().filter(|s| !s.is_empty()).ok_or_else(|| {
1279        (
1280            "States.TaskFailed".to_string(),
1281            format!("Invalid aws-sdk Resource ARN: missing action in '{tail}'"),
1282        )
1283    })?;
1284    let action_camel = action_with_mod
1285        .split('.')
1286        .next()
1287        .filter(|s| !s.is_empty())
1288        .ok_or_else(|| {
1289            (
1290                "States.TaskFailed".to_string(),
1291                format!("Invalid aws-sdk Resource ARN: empty action in '{tail}'"),
1292            )
1293        })?;
1294    let modifiers: Vec<&str> = action_with_mod.split('.').skip(1).collect();
1295    let is_sync = modifiers.iter().any(|m| *m == "sync" || *m == "sync:2");
1296
1297    let action_pascal = camel_to_pascal(action_camel);
1298    let service_name = map_sdk_service_id(service_id).to_string();
1299
1300    // ECS's wire format is camelCase, but Step Functions's optimized
1301    // integration takes PascalCase parameter names in the state machine
1302    // definition. Translate top-level keys for the ECS service so AWS
1303    // Step Functions definitions Just Work against the fake.
1304    let translated_input = match service_name.as_str() {
1305        "ecs" => translate_ecs_keys_to_camel(input),
1306        _ => input.clone(),
1307    };
1308
1309    let initial = call_sdk_action(
1310        &registry_arc,
1311        &service_name,
1312        &action_pascal,
1313        &translated_input,
1314        account_id,
1315    )
1316    .await?;
1317
1318    if !is_sync {
1319        return Ok(initial);
1320    }
1321
1322    // `.sync` pattern: dispatch to the per-service waiter that polls until
1323    // the downstream operation reaches a terminal state. Returns the full
1324    // describe-shape result, or surfaces a terminal failure as
1325    // `States.TaskFailed`.
1326    sync_wait(
1327        &registry_arc,
1328        &service_name,
1329        &action_pascal,
1330        &initial,
1331        &translated_input,
1332        account_id,
1333        timeout_seconds,
1334    )
1335    .await
1336}
1337
1338/// Translate the top-level PascalCase keys that AWS Step Functions
1339/// state machines use for `ecs:runTask` Parameters into the camelCase
1340/// shape that the AWS ECS API (and our handler) expects. Unknown keys
1341/// pass through unchanged so callers can still send raw camelCase. The
1342/// translation is shallow on purpose — nested overrides keep their
1343/// existing camelCase shape on real AWS too.
1344fn translate_ecs_keys_to_camel(input: &Value) -> Value {
1345    let Some(obj) = input.as_object() else {
1346        return input.clone();
1347    };
1348    let mut out = serde_json::Map::with_capacity(obj.len());
1349    for (k, v) in obj.iter() {
1350        let camel = match k.as_str() {
1351            "Cluster" => "cluster",
1352            "TaskDefinition" => "taskDefinition",
1353            "LaunchType" => "launchType",
1354            "Group" => "group",
1355            "Overrides" => "overrides",
1356            "PlatformVersion" => "platformVersion",
1357            "NetworkConfiguration" => "networkConfiguration",
1358            "Tags" => "tags",
1359            "EnableExecuteCommand" => "enableExecuteCommand",
1360            "PropagateTags" => "propagateTags",
1361            "ReferenceId" => "referenceId",
1362            "StartedBy" => "startedBy",
1363            "Count" => "count",
1364            "CapacityProviderStrategy" => "capacityProviderStrategy",
1365            "PlacementConstraints" => "placementConstraints",
1366            "PlacementStrategy" => "placementStrategy",
1367            other => other,
1368        };
1369        out.insert(camel.to_string(), v.clone());
1370    }
1371    Value::Object(out)
1372}
1373
1374fn resolve_registry(
1375    registry: &Option<SharedServiceRegistry>,
1376) -> Result<Arc<fakecloud_core::registry::ServiceRegistry>, (String, String)> {
1377    let registry_handle = registry.as_ref().ok_or_else(|| {
1378        (
1379            "States.TaskFailed".to_string(),
1380            "No service registry configured for aws-sdk integration".to_string(),
1381        )
1382    })?;
1383    registry_handle.get().cloned().ok_or_else(|| {
1384        (
1385            "States.TaskFailed".to_string(),
1386            "Service registry not yet initialised; aws-sdk integration unavailable".to_string(),
1387        )
1388    })
1389}
1390
1391/// Call a single AWS SDK action against the registered service handler.
1392async fn call_sdk_action(
1393    registry: &Arc<fakecloud_core::registry::ServiceRegistry>,
1394    service_name: &str,
1395    action_pascal: &str,
1396    input: &Value,
1397    account_id: &str,
1398) -> Result<Value, (String, String)> {
1399    use bytes::Bytes;
1400    use fakecloud_core::service::AwsRequest;
1401    use http::{HeaderMap, Method};
1402
1403    let service = registry.get(service_name).ok_or_else(|| {
1404        (
1405            "States.TaskFailed".to_string(),
1406            format!("Unknown aws-sdk service '{service_name}'"),
1407        )
1408    })?;
1409
1410    let body_bytes = Bytes::from(
1411        serde_json::to_vec(input).expect("serde_json::Value serialization is infallible"),
1412    );
1413
1414    let req = AwsRequest {
1415        service: service_name.to_string(),
1416        action: action_pascal.to_string(),
1417        region: "us-east-1".to_string(),
1418        account_id: account_id.to_string(),
1419        request_id: uuid::Uuid::new_v4().to_string(),
1420        headers: HeaderMap::new(),
1421        query_params: std::collections::HashMap::new(),
1422        body: body_bytes,
1423        body_stream: parking_lot::Mutex::new(None),
1424        path_segments: vec![],
1425        raw_path: "/".to_string(),
1426        raw_query: String::new(),
1427        method: Method::POST,
1428        is_query_protocol: false,
1429        access_key_id: None,
1430        principal: None,
1431    };
1432
1433    let response = service.handle(req).await.map_err(|err| {
1434        let code = err.code().to_string();
1435        let msg = err.message();
1436        let prefix_service = match service_name {
1437            "dynamodb" => "DynamoDb".to_string(),
1438            "states" => "Sfn".to_string(),
1439            other => camel_to_pascal(other),
1440        };
1441        (
1442            format!("{prefix_service}.{code}"),
1443            format!("{action_pascal} failed: {msg}"),
1444        )
1445    })?;
1446
1447    let response_bytes = match response.body {
1448        fakecloud_core::service::ResponseBody::Bytes(b) => b,
1449        fakecloud_core::service::ResponseBody::File { .. } => {
1450            return Err((
1451                "States.TaskFailed".to_string(),
1452                "aws-sdk integration: file-backed response not supported".to_string(),
1453            ));
1454        }
1455    };
1456
1457    if response_bytes.is_empty() {
1458        return Ok(json!({}));
1459    }
1460    serde_json::from_slice(&response_bytes).map_err(|e| {
1461        (
1462            "States.TaskFailed".to_string(),
1463            format!("aws-sdk integration: failed to parse response JSON: {e}"),
1464        )
1465    })
1466}
1467
1468/// Call an SDK action and return raw bytes without JSON parsing.
1469/// Used by distributed Map mode to read/write S3 objects.
1470async fn call_sdk_action_raw_bytes(
1471    registry: &Arc<fakecloud_core::registry::ServiceRegistry>,
1472    service_name: &str,
1473    action_pascal: &str,
1474    input: &Value,
1475    account_id: &str,
1476) -> Result<bytes::Bytes, (String, String)> {
1477    use bytes::Bytes;
1478    use fakecloud_core::service::AwsRequest;
1479    use http::{HeaderMap, Method};
1480
1481    let service = registry.get(service_name).ok_or_else(|| {
1482        (
1483            "States.TaskFailed".to_string(),
1484            format!("Unknown aws-sdk service '{service_name}'"),
1485        )
1486    })?;
1487
1488    let body_bytes = Bytes::from(
1489        serde_json::to_vec(input).expect("serde_json::Value serialization is infallible"),
1490    );
1491
1492    let req = AwsRequest {
1493        service: service_name.to_string(),
1494        action: action_pascal.to_string(),
1495        region: "us-east-1".to_string(),
1496        account_id: account_id.to_string(),
1497        request_id: uuid::Uuid::new_v4().to_string(),
1498        headers: HeaderMap::new(),
1499        query_params: std::collections::HashMap::new(),
1500        body: body_bytes,
1501        body_stream: parking_lot::Mutex::new(None),
1502        path_segments: vec![],
1503        raw_path: "/".to_string(),
1504        raw_query: String::new(),
1505        method: Method::POST,
1506        is_query_protocol: false,
1507        access_key_id: None,
1508        principal: None,
1509    };
1510
1511    let response = service.handle(req).await.map_err(|err| {
1512        let code = err.code().to_string();
1513        let msg = err.message();
1514        let prefix_service = match service_name {
1515            "dynamodb" => "DynamoDb".to_string(),
1516            "states" => "Sfn".to_string(),
1517            other => camel_to_pascal(other),
1518        };
1519        (
1520            format!("{prefix_service}.{code}"),
1521            format!("{action_pascal} failed: {msg}"),
1522        )
1523    })?;
1524
1525    match response.body {
1526        fakecloud_core::service::ResponseBody::Bytes(b) => Ok(b),
1527        fakecloud_core::service::ResponseBody::File { .. } => Err((
1528            "States.TaskFailed".to_string(),
1529            "aws-sdk integration: file-backed response not supported".to_string(),
1530        )),
1531    }
1532}
1533
1534/// Cap on `.sync` polling so a stuck downstream task can't hang an
1535/// execution forever. Mirrors the `TimeoutSeconds` knob on the Task state
1536/// when one is set; otherwise defaults to 5 minutes which is enough for
1537/// the in-process services that ship today (Athena returns synchronously,
1538/// ECS tasks finish within seconds even when docker-less).
1539const SYNC_DEFAULT_TIMEOUT_SECS: u64 = 300;
1540const SYNC_POLL_INTERVAL_MS: u64 = 200;
1541
1542/// Dispatch `.sync` waiters by service+action. Each waiter polls the
1543/// matching describe-style API until the downstream operation reaches a
1544/// terminal state, then returns the full describe response.
1545async fn sync_wait(
1546    registry: &Arc<fakecloud_core::registry::ServiceRegistry>,
1547    service_name: &str,
1548    action_pascal: &str,
1549    initial: &Value,
1550    input: &Value,
1551    account_id: &str,
1552    timeout_seconds: Option<u64>,
1553) -> Result<Value, (String, String)> {
1554    match (service_name, action_pascal) {
1555        ("ecs", "RunTask") => {
1556            sync_wait_ecs_run_task(registry, initial, input, account_id, timeout_seconds).await
1557        }
1558        ("athena", "StartQueryExecution") => {
1559            sync_wait_athena_query(registry, initial, account_id, timeout_seconds).await
1560        }
1561        ("states", "StartExecution") => {
1562            sync_wait_states_start_execution(registry, initial, account_id, timeout_seconds).await
1563        }
1564        ("glue", "StartJobRun") => {
1565            // Poll the real Glue `GetJobRun` (the job run was created by the
1566            // StartJobRun integration that ran just before this waiter), so the
1567            // `.sync` result is the actual run — full JobRun shape, and a
1568            // FAILED/STOPPED/TIMEOUT state surfaces as a task failure rather
1569            // than a hardcoded SUCCEEDED.
1570            let job_run_id = initial
1571                .get("JobRunId")
1572                .and_then(Value::as_str)
1573                .unwrap_or_default()
1574                .to_string();
1575            let job_name = input
1576                .get("JobName")
1577                .and_then(Value::as_str)
1578                .unwrap_or_default()
1579                .to_string();
1580            let deadline = sync_deadline(timeout_seconds);
1581            loop {
1582                let described = call_sdk_action(
1583                    registry,
1584                    "glue",
1585                    "GetJobRun",
1586                    &json!({ "JobName": job_name, "RunId": job_run_id }),
1587                    account_id,
1588                )
1589                .await?;
1590                let state = described
1591                    .get("JobRun")
1592                    .and_then(|r| r.get("JobRunState"))
1593                    .and_then(Value::as_str)
1594                    .unwrap_or("");
1595                match state {
1596                    "SUCCEEDED" => return Ok(described),
1597                    "FAILED" | "STOPPED" | "TIMEOUT" | "ERROR" => {
1598                        return Err((
1599                            "States.TaskFailed".to_string(),
1600                            format!("Glue job run {job_run_id} ended in state {state}"),
1601                        ));
1602                    }
1603                    _ => {}
1604                }
1605                if std::time::Instant::now() >= deadline {
1606                    return Err((
1607                        "States.Timeout".to_string(),
1608                        format!(
1609                            "glue:startJobRun.sync timed out after {}s for run {job_run_id}",
1610                            sync_timeout_secs(timeout_seconds)
1611                        ),
1612                    ));
1613                }
1614                tokio::time::sleep(std::time::Duration::from_millis(SYNC_POLL_INTERVAL_MS)).await;
1615            }
1616        }
1617        _ => Err((
1618            "States.TaskFailed".to_string(),
1619            format!(
1620                "`.sync` is not supported for {service_name}:{action_pascal} yet — \
1621                 supported: ecs:RunTask, athena:StartQueryExecution, glue:StartJobRun, states:StartExecution"
1622            ),
1623        )),
1624    }
1625}
1626
1627async fn sync_wait_ecs_run_task(
1628    registry: &Arc<fakecloud_core::registry::ServiceRegistry>,
1629    initial: &Value,
1630    input: &Value,
1631    account_id: &str,
1632    timeout_seconds: Option<u64>,
1633) -> Result<Value, (String, String)> {
1634    let tasks = initial
1635        .get("tasks")
1636        .and_then(Value::as_array)
1637        .ok_or_else(|| {
1638            (
1639                "States.TaskFailed".to_string(),
1640                "ecs:RunTask.sync: response missing 'tasks' array".to_string(),
1641            )
1642        })?;
1643    if tasks.is_empty() {
1644        return Err((
1645            "States.TaskFailed".to_string(),
1646            "ecs:RunTask.sync: no tasks were started".to_string(),
1647        ));
1648    }
1649    let task_arns: Vec<String> = tasks
1650        .iter()
1651        .filter_map(|t| t.get("taskArn").and_then(Value::as_str).map(String::from))
1652        .collect();
1653    let cluster = input
1654        .get("cluster")
1655        .or_else(|| input.get("Cluster"))
1656        .and_then(Value::as_str)
1657        .map(String::from);
1658
1659    let deadline = sync_deadline(timeout_seconds);
1660    loop {
1661        let mut describe_input = json!({ "tasks": task_arns });
1662        if let Some(c) = &cluster {
1663            describe_input["cluster"] = json!(c);
1664        }
1665        let described = call_sdk_action(
1666            registry,
1667            "ecs",
1668            "DescribeTasks",
1669            &describe_input,
1670            account_id,
1671        )
1672        .await?;
1673        let described_tasks = described
1674            .get("tasks")
1675            .and_then(Value::as_array)
1676            .cloned()
1677            .unwrap_or_default();
1678        let all_stopped = !described_tasks.is_empty()
1679            && described_tasks
1680                .iter()
1681                .all(|t| t.get("lastStatus").and_then(Value::as_str) == Some("STOPPED"));
1682        if all_stopped {
1683            // Surface non-zero container exit codes or stop codes that map
1684            // to AWS-style failures. Per AWS docs, ECS RunTask.sync raises
1685            // States.TaskFailed when any container exits non-zero or the
1686            // task stops with a failure code.
1687            let any_failed = described_tasks.iter().any(|t| {
1688                let stop_code = t.get("stopCode").and_then(Value::as_str);
1689                let bad_stop = matches!(
1690                    stop_code,
1691                    Some(
1692                        "TaskFailedToStart"
1693                            | "EssentialContainerExited"
1694                            | "ServiceSchedulerInitiated"
1695                    )
1696                );
1697                let bad_exit = t
1698                    .get("containers")
1699                    .and_then(Value::as_array)
1700                    .map(|cs| {
1701                        cs.iter().any(|c| {
1702                            c.get("exitCode")
1703                                .and_then(Value::as_i64)
1704                                .map(|n| n != 0)
1705                                .unwrap_or(false)
1706                        })
1707                    })
1708                    .unwrap_or(false);
1709                bad_stop || bad_exit
1710            });
1711            if any_failed {
1712                let cause = described_tasks
1713                    .iter()
1714                    .find_map(|t| {
1715                        t.get("stoppedReason")
1716                            .and_then(Value::as_str)
1717                            .map(String::from)
1718                    })
1719                    .unwrap_or_else(|| "ECS task failed".to_string());
1720                return Err(("States.TaskFailed".to_string(), cause));
1721            }
1722            return Ok(described);
1723        }
1724        if std::time::Instant::now() >= deadline {
1725            return Err((
1726                "States.Timeout".to_string(),
1727                format!(
1728                    "ecs:RunTask.sync timed out after {}s waiting for {} task(s) to STOP",
1729                    sync_timeout_secs(timeout_seconds),
1730                    task_arns.len()
1731                ),
1732            ));
1733        }
1734        tokio::time::sleep(std::time::Duration::from_millis(SYNC_POLL_INTERVAL_MS)).await;
1735    }
1736}
1737
1738async fn sync_wait_athena_query(
1739    registry: &Arc<fakecloud_core::registry::ServiceRegistry>,
1740    initial: &Value,
1741    account_id: &str,
1742    timeout_seconds: Option<u64>,
1743) -> Result<Value, (String, String)> {
1744    let qid = initial
1745        .get("QueryExecutionId")
1746        .and_then(Value::as_str)
1747        .ok_or_else(|| {
1748            (
1749                "States.TaskFailed".to_string(),
1750                "athena:StartQueryExecution.sync: response missing QueryExecutionId".to_string(),
1751            )
1752        })?
1753        .to_string();
1754
1755    let deadline = sync_deadline(timeout_seconds);
1756    loop {
1757        let described = call_sdk_action(
1758            registry,
1759            "athena",
1760            "GetQueryExecution",
1761            &json!({ "QueryExecutionId": qid }),
1762            account_id,
1763        )
1764        .await?;
1765        let state = described
1766            .get("QueryExecution")
1767            .and_then(|qe| qe.get("Status"))
1768            .and_then(|s| s.get("State"))
1769            .and_then(Value::as_str)
1770            .unwrap_or("");
1771        match state {
1772            "SUCCEEDED" => return Ok(described),
1773            "FAILED" | "CANCELLED" => {
1774                let cause = described
1775                    .get("QueryExecution")
1776                    .and_then(|qe| qe.get("Status"))
1777                    .and_then(|s| s.get("StateChangeReason"))
1778                    .and_then(Value::as_str)
1779                    .unwrap_or("Athena query reached terminal failure state")
1780                    .to_string();
1781                return Err(("States.TaskFailed".to_string(), cause));
1782            }
1783            _ => {}
1784        }
1785        if std::time::Instant::now() >= deadline {
1786            return Err((
1787                "States.Timeout".to_string(),
1788                format!(
1789                    "athena:StartQueryExecution.sync timed out after {}s for query {qid}",
1790                    sync_timeout_secs(timeout_seconds)
1791                ),
1792            ));
1793        }
1794        tokio::time::sleep(std::time::Duration::from_millis(SYNC_POLL_INTERVAL_MS)).await;
1795    }
1796}
1797
1798/// Wait for a nested Step Functions execution to reach a terminal
1799/// state, then return the bubbled `Output` like AWS does for
1800/// `arn:aws:states:::states:startExecution.sync`. The initial
1801/// `StartExecution` response carries `executionArn`; we poll
1802/// `DescribeExecution` until status leaves `RUNNING`.
1803async fn sync_wait_states_start_execution(
1804    registry: &Arc<fakecloud_core::registry::ServiceRegistry>,
1805    initial: &Value,
1806    account_id: &str,
1807    timeout_seconds: Option<u64>,
1808) -> Result<Value, (String, String)> {
1809    let exec_arn = initial
1810        .get("executionArn")
1811        .or_else(|| initial.get("ExecutionArn"))
1812        .and_then(Value::as_str)
1813        .ok_or_else(|| {
1814            (
1815                "States.TaskFailed".to_string(),
1816                "states:startExecution.sync: response missing executionArn".to_string(),
1817            )
1818        })?
1819        .to_string();
1820
1821    let deadline = sync_deadline(timeout_seconds);
1822    loop {
1823        let described = call_sdk_action(
1824            registry,
1825            "states",
1826            "DescribeExecution",
1827            &json!({ "executionArn": exec_arn }),
1828            account_id,
1829        )
1830        .await?;
1831        let status = described
1832            .get("status")
1833            .or_else(|| described.get("Status"))
1834            .and_then(Value::as_str)
1835            .unwrap_or("");
1836        match status {
1837            "SUCCEEDED" => return Ok(described),
1838            "FAILED" | "TIMED_OUT" | "ABORTED" => {
1839                let cause = described
1840                    .get("cause")
1841                    .or_else(|| described.get("Cause"))
1842                    .and_then(Value::as_str)
1843                    .unwrap_or("Nested execution reached terminal failure state")
1844                    .to_string();
1845                return Err(("States.TaskFailed".to_string(), cause));
1846            }
1847            _ => {}
1848        }
1849        if std::time::Instant::now() >= deadline {
1850            return Err((
1851                "States.Timeout".to_string(),
1852                format!(
1853                    "states:startExecution.sync timed out after {}s for {exec_arn}",
1854                    sync_timeout_secs(timeout_seconds)
1855                ),
1856            ));
1857        }
1858        tokio::time::sleep(std::time::Duration::from_millis(SYNC_POLL_INTERVAL_MS)).await;
1859    }
1860}
1861
1862fn sync_timeout_secs(timeout_seconds: Option<u64>) -> u64 {
1863    timeout_seconds.unwrap_or(SYNC_DEFAULT_TIMEOUT_SECS)
1864}
1865
1866fn sync_deadline(timeout_seconds: Option<u64>) -> std::time::Instant {
1867    std::time::Instant::now() + std::time::Duration::from_secs(sync_timeout_secs(timeout_seconds))
1868}
1869
1870#[derive(Clone, Copy)]
1871pub(crate) enum UpdateClause {
1872    Set,
1873    Remove,
1874    Add,
1875    Delete,
1876}
1877
1878/// Invoke a Lambda function directly via DeliveryBus.
1879async fn invoke_lambda_direct(
1880    function_arn: &str,
1881    input: &Value,
1882    delivery: &Option<Arc<DeliveryBus>>,
1883    timeout_seconds: Option<u64>,
1884) -> Result<Value, (String, String)> {
1885    let delivery = delivery.as_ref().ok_or_else(|| {
1886        (
1887            "States.TaskFailed".to_string(),
1888            "No delivery bus configured for Lambda invocation".to_string(),
1889        )
1890    })?;
1891
1892    let payload =
1893        serde_json::to_string(input).expect("serde_json::Value serialization is infallible");
1894
1895    let invoke_future = delivery.invoke_lambda(function_arn, &payload);
1896
1897    let result = if let Some(timeout) = timeout_seconds {
1898        match tokio::time::timeout(tokio::time::Duration::from_secs(timeout), invoke_future).await {
1899            Ok(r) => r,
1900            Err(_) => {
1901                return Err((
1902                    "States.Timeout".to_string(),
1903                    format!("Task timed out after {timeout} seconds"),
1904                ));
1905            }
1906        }
1907    } else {
1908        invoke_future.await
1909    };
1910
1911    match result {
1912        Some(Ok(bytes)) => {
1913            let response_str = String::from_utf8_lossy(&bytes);
1914            let value: Value =
1915                serde_json::from_str(&response_str).unwrap_or(json!(response_str.to_string()));
1916
1917            // The Lambda runtime returns HTTP 200 even for unhandled function
1918            // errors, encoding the failure in the payload as
1919            // `{"errorMessage", "errorType", "stackTrace"}`. The delivery trait
1920            // only conveys bytes, so the payload is the only available signal —
1921            // the same heuristic the Lambda service uses for async destination
1922            // routing. Mirror real AWS: a function error fails the task with the
1923            // function's `errorType` as the Error and the serialized payload as
1924            // the Cause, which is exactly what makes Retry/Catch on custom
1925            // exception names work.
1926            //
1927            // Require BOTH `errorType` and `errorMessage` (not either) to avoid
1928            // misreading a legitimate success payload that merely happens to
1929            // contain one of those field names as a function error.
1930            if let Some(obj) = value.as_object() {
1931                if obj.contains_key("errorType") && obj.contains_key("errorMessage") {
1932                    let error_type = obj
1933                        .get("errorType")
1934                        .and_then(Value::as_str)
1935                        .unwrap_or("Exception")
1936                        .to_string();
1937                    let cause = serde_json::to_string(&value)
1938                        .expect("serde_json::Value serialization is infallible");
1939                    return Err((error_type, cause));
1940                }
1941            }
1942
1943            Ok(value)
1944        }
1945        Some(Err(e)) => {
1946            // Failures of the Lambda Invoke API call itself surface on real AWS
1947            // with `Lambda.`-prefixed error names (`Lambda.ServiceException`,
1948            // `Lambda.SdkClientException`, `Lambda.Unknown`, ...) — never
1949            // `States.TaskFailed` — and those are the names AWS documents for
1950            // Retry rules. The delivery error is an unstructured string with
1951            // the outcome flattened, so the honest default is `Lambda.Unknown`
1952            // ("outcome unknown"). The one case we can cheaply distinguish is a
1953            // missing function, which AWS reports as
1954            // `Lambda.ResourceNotFoundException`.
1955            if e.starts_with("Function not found") {
1956                Err(("Lambda.ResourceNotFoundException".to_string(), e))
1957            } else {
1958                Err(("Lambda.Unknown".to_string(), e))
1959            }
1960        }
1961        None => {
1962            // No runtime available — return empty result
1963            Ok(json!({}))
1964        }
1965    }
1966}
1967
1968/// Invoke an activity worker. Inserts a `PENDING` token into shared state
1969/// so a worker can claim it via `GetActivityTask`, then polls until the
1970/// worker calls `SendTaskSuccess` / `SendTaskFailure` or the heartbeat /
1971/// timeout windows expire.
1972async fn invoke_activity(
1973    activity_arn: &str,
1974    input: &Value,
1975    shared_state: &SharedStepFunctionsState,
1976    timeout_seconds: Option<u64>,
1977    heartbeat_seconds: Option<u64>,
1978) -> Result<Value, (String, String)> {
1979    use crate::state::TaskTokenState;
1980
1981    // Activity must exist (look up across accounts via ARN segment).
1982    let activity_account = activity_arn.split(':').nth(4).unwrap_or("").to_string();
1983    {
1984        let accounts = shared_state.read();
1985        let exists = accounts
1986            .get(&activity_account)
1987            .map(|s| s.activities.contains_key(activity_arn))
1988            .unwrap_or(false);
1989        if !exists {
1990            return Err((
1991                "States.TaskFailed".to_string(),
1992                format!("Activity does not exist: {activity_arn}"),
1993            ));
1994        }
1995    }
1996
1997    let token = format!(
1998        "FCToken-{}-{}",
1999        chrono::Utc::now().timestamp_nanos_opt().unwrap_or(0),
2000        uuid::Uuid::new_v4().simple(),
2001    );
2002    let now = chrono::Utc::now();
2003    let input_str =
2004        serde_json::to_string(input).expect("serde_json::Value serialization is infallible");
2005    {
2006        let mut accounts = shared_state.write();
2007        let state = accounts.get_or_create(&activity_account);
2008        state.task_tokens.insert(
2009            token.clone(),
2010            TaskTokenState {
2011                activity_arn: activity_arn.to_string(),
2012                status: "PENDING".to_string(),
2013                output: None,
2014                error: None,
2015                cause: None,
2016                input: Some(input_str),
2017                created_at: now,
2018                last_heartbeat_at: None,
2019                heartbeat_seconds: heartbeat_seconds.map(|s| s as i64),
2020                timeout_seconds: timeout_seconds.map(|s| s as i64),
2021            },
2022        );
2023    }
2024
2025    poll_task_token(
2026        shared_state,
2027        &activity_account,
2028        &token,
2029        timeout_seconds,
2030        heartbeat_seconds,
2031    )
2032    .await
2033}
2034
2035pub(crate) enum NextState {
2036    Name(String),
2037    End,
2038    Error(String),
2039}
2040
2041#[path = "interpreter_helpers.rs"]
2042mod interpreter_helpers;
2043pub(crate) use interpreter_helpers::*;
2044
2045#[cfg(test)]
2046#[path = "interpreter_tests.rs"]
2047mod tests;