rustvello 0.9.1

use super::*;
use rustvello_core::broker::Broker;
use rustvello_core::error::RustvelloError;
use rustvello_core::observability::{
    EventEmitter, EventLevel, TaskLifecycleEvent, TaskLifecycleKind,
};
use rustvello_core::orchestrator::InvocationControlBackend;
use rustvello_core::runner::Runner;
use rustvello_core::state_backend::StateBackend;
use rustvello_core::task::{TaskDefinition, TaskRegistry};
use rustvello_proto::call::{CallDTO, SerializedArguments};
use rustvello_proto::config::{AppConfig, TaskConfig};
use rustvello_proto::identifiers::{InvocationId, RunnerId, TaskId};
use rustvello_proto::invocation::{InvocationDTO, TraceContextCarrier};
use rustvello_proto::status::{ConcurrencyControlType, InvocationStatus};
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::Arc;
use std::sync::Mutex;
use std::time::Duration;

pub(super) fn make_runner() -> (
    PersistentTokioRunner,
    Arc<dyn InvocationControlBackend>,
    Arc<dyn StateBackend>,
) {
    let broker: Arc<dyn Broker> = Arc::new(rustvello_mem::broker::MemBroker::new());
    let orchestrator: Arc<dyn InvocationControlBackend> =
        Arc::new(rustvello_mem::orchestrator::MemOrchestrator::new());
    let state_backend: Arc<dyn StateBackend> =
        Arc::new(rustvello_mem::state_backend::MemStateBackend::new());

    let mut registry = TaskRegistry::new();
    registry
        .register(TaskDefinition::new(
            TaskId::new("test", "double"),
            TaskConfig::default(),
            Arc::new(|args_json: String| {
                let args: std::collections::BTreeMap<String, String> =
                    serde_json::from_str(&args_json).map_err(|e| {
                        RustvelloError::Serialization {
                            message: e.to_string(),
                        }
                    })?;
                let x: i64 = args.get("x").and_then(|v| v.parse().ok()).unwrap_or(0);
                serde_json::to_string(&(x * 2)).map_err(|e| RustvelloError::Serialization {
                    message: e.to_string(),
                })
            }),
        ))
        .unwrap();

    let runner = PersistentTokioRunner::new(
        "test-app".to_string(),
        AppConfig::default(),
        Arc::clone(&broker),
        Arc::clone(&orchestrator),
        Arc::clone(&state_backend),
        Arc::new(registry),
        None,
    );

    (runner, orchestrator, state_backend)
}

#[derive(Clone, Default)]
struct RecordingEmitter {
    events: Arc<Mutex<Vec<TaskLifecycleEvent>>>,
}

impl EventEmitter for RecordingEmitter {
    fn on_task_lifecycle(&self, event: &TaskLifecycleEvent) {
        self.events.lock().unwrap().push(event.clone());
    }
}

#[tokio::test]
async fn test_run_one_no_work() {
    let (runner, _, _) = make_runner();
    let did_work = runner.run_one().await.unwrap();
    assert!(!did_work);
}

#[tokio::test]
async fn test_full_invocation_cycle() {
    let (runner, orchestrator, state_backend) = make_runner();

    let task_id = TaskId::new("test", "double");
    let mut args = SerializedArguments::new();
    args.insert("x", "21");
    let call = CallDTO::new(task_id.clone(), args);

    let inv_id = orchestrator.register_invocation(&call).await.unwrap();
    let inv_dto = InvocationDTO::new(inv_id.clone(), task_id, call.call_id.clone());
    state_backend
        .upsert_invocation(&inv_dto, &call)
        .await
        .unwrap();
    runner
        .control_plane
        .broker
        .route_invocation(&inv_id)
        .await
        .unwrap();

    let did_work = runner.run_one().await.unwrap();
    assert!(did_work);

    let status = orchestrator.get_invocation_status(&inv_id).await.unwrap();
    assert_eq!(status.status, InvocationStatus::Success);

    let result = state_backend.get_result(&inv_id).await.unwrap();
    assert_eq!(result, Some("42".to_string()));
}

#[tokio::test]
async fn lifecycle_hooks_carry_attempt_worker_and_task_context() {
    let (runner, orchestrator, state_backend) = make_runner();
    let recording = RecordingEmitter::default();
    let events = Arc::clone(&recording.events);
    let runner = runner.with_event_emitter(EventLevel::TaskLifecycle, recording);

    let task_id = TaskId::new("test", "double");
    let call = CallDTO::new(task_id.clone(), SerializedArguments::new());
    let invocation_id = orchestrator.register_invocation(&call).await.unwrap();
    state_backend
        .upsert_invocation(
            &InvocationDTO::new(invocation_id.clone(), task_id.clone(), call.call_id.clone()),
            &call,
        )
        .await
        .unwrap();
    runner
        .control_plane
        .broker
        .route_invocation(&invocation_id)
        .await
        .unwrap();

    assert!(runner.run_one().await.unwrap());

    let events = events.lock().unwrap();
    assert_eq!(events.len(), 2);
    assert!(matches!(events[0].kind, TaskLifecycleKind::Started));
    assert!(matches!(
        events[1].kind,
        TaskLifecycleKind::Succeeded { .. }
    ));
    for event in events.iter() {
        assert_eq!(event.context.app_id.as_ref(), "test-app");
        assert_eq!(event.context.task_id, task_id);
        assert_eq!(event.context.invocation_id, invocation_id);
        assert_eq!(event.context.attempt, 0);
        assert_eq!(event.context.queue.as_ref(), "default");
        assert_eq!(
            event.context.worker.runner_cls.as_ref(),
            "PersistentTokioWorker"
        );
        assert!(event.context.worker.parent_runner_id.is_some());
        assert!(event.context.trace_context.traceparent.is_none());
    }
    assert!(events[0].event_time <= events[1].event_time);
}

#[tokio::test]
async fn test_retry_on_failure() {
    let broker: Arc<dyn Broker> = Arc::new(rustvello_mem::broker::MemBroker::new());
    let orchestrator: Arc<dyn InvocationControlBackend> =
        Arc::new(rustvello_mem::orchestrator::MemOrchestrator::new());
    let state_backend: Arc<dyn StateBackend> =
        Arc::new(rustvello_mem::state_backend::MemStateBackend::new());

    let mut registry = TaskRegistry::new();
    registry
        .register({
            let mut config = TaskConfig::default();
            config.max_retries = 2;
            TaskDefinition::new(
                TaskId::new("test", "failing"),
                config,
                Arc::new(|_args: String| {
                    Err(RustvelloError::runner_err("always fails".to_string()))
                }),
            )
        })
        .unwrap();

    let recording = RecordingEmitter::default();
    let events = Arc::clone(&recording.events);
    let runner = PersistentTokioRunner::new(
        "test-app".to_string(),
        AppConfig::default(),
        Arc::clone(&broker),
        Arc::clone(&orchestrator),
        Arc::clone(&state_backend),
        Arc::new(registry),
        None,
    )
    .with_event_emitter(EventLevel::TaskLifecycle, recording);

    let task_id = TaskId::new("test", "failing");
    let args = SerializedArguments::new();
    let call = CallDTO::new(task_id.clone(), args);

    let inv_id = orchestrator.register_invocation(&call).await.unwrap();
    let trace_context = TraceContextCarrier {
        traceparent: Some("00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01".to_string()),
        tracestate: Some("example=retry".to_string()),
    };
    let inv_dto = InvocationDTO::new(inv_id.clone(), task_id, call.call_id.clone())
        .with_trace_context(trace_context.clone());
    state_backend
        .upsert_invocation(&inv_dto, &call)
        .await
        .unwrap();
    broker.route_invocation(&inv_id).await.unwrap();

    runner.run_one().await.unwrap();
    let status = orchestrator.get_invocation_status(&inv_id).await.unwrap();
    assert_eq!(status.status, InvocationStatus::Retry);

    runner.run_one().await.unwrap();
    let status = orchestrator.get_invocation_status(&inv_id).await.unwrap();
    assert_eq!(status.status, InvocationStatus::Retry);

    runner.run_one().await.unwrap();
    let status = orchestrator.get_invocation_status(&inv_id).await.unwrap();
    assert_eq!(status.status, InvocationStatus::Failed);

    let error = state_backend.get_error(&inv_id).await.unwrap();
    assert!(error.is_some());

    let events = events.lock().unwrap();
    let attempts: Vec<_> = events
        .iter()
        .filter(|event| matches!(event.kind, TaskLifecycleKind::Started))
        .collect();
    assert_eq!(attempts.len(), 3);
    assert_eq!(attempts[0].context.attempt, 0);
    assert_eq!(attempts[1].context.attempt, 1);
    assert_eq!(attempts[2].context.attempt, 2);
    assert!(attempts
        .windows(2)
        .all(|pair| pair[0].context.worker.runner_id != pair[1].context.worker.runner_id));
    let lifecycle: Vec<_> = events
        .iter()
        .map(|event| (&event.kind, event.context.attempt))
        .collect();
    assert!(matches!(lifecycle[0], (&TaskLifecycleKind::Started, 0)));
    assert!(matches!(
        lifecycle[1],
        (&TaskLifecycleKind::Failed { .. }, 0)
    ));
    assert!(matches!(
        lifecycle[2],
        (&TaskLifecycleKind::RetryScheduled { .. }, 0)
    ));
    assert!(matches!(lifecycle[3], (&TaskLifecycleKind::Started, 1)));
    assert!(matches!(
        lifecycle[4],
        (&TaskLifecycleKind::Failed { .. }, 1)
    ));
    assert!(matches!(
        lifecycle[5],
        (&TaskLifecycleKind::RetryScheduled { .. }, 1)
    ));
    assert!(matches!(lifecycle[6], (&TaskLifecycleKind::Started, 2)));
    assert!(matches!(
        lifecycle[7],
        (&TaskLifecycleKind::Failed { .. }, 2)
    ));
    assert!(events
        .iter()
        .all(|event| event.context.trace_context == trace_context));
}

#[tokio::test]
async fn test_heartbeat_registers_with_orchestrator() {
    let (runner, orchestrator, _) = make_runner();
    runner.heartbeat().await.unwrap();
    let stale = orchestrator
        .get_stale_running_invocations(60)
        .await
        .unwrap();
    assert!(stale.is_empty());
}

#[tokio::test]
async fn test_recover_stale_pending_invocation() {
    let broker: Arc<dyn Broker> = Arc::new(rustvello_mem::broker::MemBroker::new());
    let mem_orchestrator = Arc::new(rustvello_mem::orchestrator::MemOrchestrator::new());
    let orchestrator: Arc<dyn InvocationControlBackend> = Arc::clone(&mem_orchestrator) as _;
    let state_backend: Arc<dyn StateBackend> =
        Arc::new(rustvello_mem::state_backend::MemStateBackend::new());

    let mut registry = TaskRegistry::new();
    registry
        .register(TaskDefinition::new(
            TaskId::new("test", "double"),
            TaskConfig::default(),
            Arc::new(|args_json: String| {
                let args: std::collections::BTreeMap<String, String> =
                    serde_json::from_str(&args_json).map_err(|e| {
                        RustvelloError::Serialization {
                            message: e.to_string(),
                        }
                    })?;
                let x: i64 = args.get("x").and_then(|v| v.parse().ok()).unwrap_or(0);
                serde_json::to_string(&(x * 2)).map_err(|e| RustvelloError::Serialization {
                    message: e.to_string(),
                })
            }),
        ))
        .unwrap();

    let runner = PersistentTokioRunner::new(
        "test-app".to_string(),
        AppConfig::default(),
        Arc::clone(&broker),
        Arc::clone(&orchestrator),
        Arc::clone(&state_backend),
        Arc::new(registry),
        None,
    );

    let task_id = TaskId::new("test", "double");
    let mut args = SerializedArguments::new();
    args.insert("x", "5");
    let call = CallDTO::new(task_id.clone(), args);

    let dead_runner_id = RunnerId::from_string("dead-runner-pending");
    let inv_id = orchestrator.register_invocation(&call).await.unwrap();
    let inv_dto = InvocationDTO::new(inv_id.clone(), task_id, call.call_id.clone());
    state_backend
        .upsert_invocation(&inv_dto, &call)
        .await
        .unwrap();
    orchestrator
        .set_invocation_status(&inv_id, InvocationStatus::Pending, Some(&dead_runner_id))
        .await
        .unwrap();

    {
        mem_orchestrator
            .backdate_status_for_testing(&inv_id, chrono::Duration::seconds(600))
            .await;
    }

    let recovered = runner.recover_stale_invocations().await.unwrap();
    assert_eq!(recovered, 1);

    let status = orchestrator.get_invocation_status(&inv_id).await.unwrap();
    assert_eq!(status.status, InvocationStatus::Rerouted);

    let did_work = runner.run_one().await.unwrap();
    assert!(did_work);

    let status = orchestrator.get_invocation_status(&inv_id).await.unwrap();
    assert_eq!(status.status, InvocationStatus::Success);
}

#[tokio::test]
async fn test_recover_stale_running_invocation() {
    let broker: Arc<dyn Broker> = Arc::new(rustvello_mem::broker::MemBroker::new());
    let mem_orchestrator = Arc::new(rustvello_mem::orchestrator::MemOrchestrator::new());
    let orchestrator: Arc<dyn InvocationControlBackend> = Arc::clone(&mem_orchestrator) as _;
    let state_backend: Arc<dyn StateBackend> =
        Arc::new(rustvello_mem::state_backend::MemStateBackend::new());

    let mut registry = TaskRegistry::new();
    registry
        .register(TaskDefinition::new(
            TaskId::new("test", "double"),
            TaskConfig::default(),
            Arc::new(|args_json: String| {
                let args: std::collections::BTreeMap<String, String> =
                    serde_json::from_str(&args_json).map_err(|e| {
                        RustvelloError::Serialization {
                            message: e.to_string(),
                        }
                    })?;
                let x: i64 = args.get("x").and_then(|v| v.parse().ok()).unwrap_or(0);
                serde_json::to_string(&(x * 2)).map_err(|e| RustvelloError::Serialization {
                    message: e.to_string(),
                })
            }),
        ))
        .unwrap();

    let runner = PersistentTokioRunner::new(
        "test-app".to_string(),
        AppConfig::default(),
        Arc::clone(&broker),
        Arc::clone(&orchestrator),
        Arc::clone(&state_backend),
        Arc::new(registry),
        None,
    );

    let dead_runner_id = RunnerId::from_string("dead-runner");

    let task_id = TaskId::new("test", "double");
    let mut args = SerializedArguments::new();
    args.insert("x", "7");
    let call = CallDTO::new(task_id.clone(), args);

    let inv_id = orchestrator.register_invocation(&call).await.unwrap();
    let inv_dto = InvocationDTO::new(inv_id.clone(), task_id, call.call_id.clone());
    state_backend
        .upsert_invocation(&inv_dto, &call)
        .await
        .unwrap();
    orchestrator
        .set_invocation_status(&inv_id, InvocationStatus::Pending, Some(&dead_runner_id))
        .await
        .unwrap();
    orchestrator
        .set_invocation_status(&inv_id, InvocationStatus::Running, Some(&dead_runner_id))
        .await
        .unwrap();
    // The dead runner never heartbeated and its RUNNING status is old.
    mem_orchestrator
        .backdate_status_for_testing(&inv_id, chrono::Duration::seconds(600))
        .await;

    let recovered = runner.recover_stale_invocations().await.unwrap();
    assert_eq!(recovered, 1);

    let status = orchestrator.get_invocation_status(&inv_id).await.unwrap();
    assert_eq!(status.status, InvocationStatus::Rerouted);

    let did_work = runner.run_one().await.unwrap();
    assert!(did_work);

    let status = orchestrator.get_invocation_status(&inv_id).await.unwrap();
    assert_eq!(status.status, InvocationStatus::Success);
}

#[test]
fn test_default_workers_matches_cpu_count() {
    let (runner, _, _) = make_runner();
    let expected = std::thread::available_parallelism().map_or(1, std::num::NonZero::get);
    assert_eq!(runner.num_workers(), expected);
    assert!(runner.num_workers() >= 1);
}

#[test]
fn test_with_num_workers_override() {
    let (runner, _, _) = make_runner();
    let runner = runner.with_num_workers(4);
    assert_eq!(runner.num_workers(), 4);
}

#[test]
fn test_with_num_workers_clamps_to_one() {
    let (runner, _, _) = make_runner();
    let runner = runner.with_num_workers(0);
    assert_eq!(runner.num_workers(), 1);
}

#[tokio::test]
async fn test_concurrent_workers_process_invocations() {
    let broker: Arc<dyn Broker> = Arc::new(rustvello_mem::broker::MemBroker::new());
    let orchestrator: Arc<dyn InvocationControlBackend> =
        Arc::new(rustvello_mem::orchestrator::MemOrchestrator::new());
    let state_backend: Arc<dyn StateBackend> =
        Arc::new(rustvello_mem::state_backend::MemStateBackend::new());

    let mut registry = TaskRegistry::new();
    registry
        .register(TaskDefinition::new(
            TaskId::new("test", "double"),
            TaskConfig::default(),
            Arc::new(|args_json: String| {
                let args: std::collections::BTreeMap<String, String> =
                    serde_json::from_str(&args_json).map_err(|e| {
                        RustvelloError::Serialization {
                            message: e.to_string(),
                        }
                    })?;
                let x: i64 = args.get("x").and_then(|v| v.parse().ok()).unwrap_or(0);
                serde_json::to_string(&(x * 2)).map_err(|e| RustvelloError::Serialization {
                    message: e.to_string(),
                })
            }),
        ))
        .unwrap();

    let runner = PersistentTokioRunner::new(
        "test-app".to_string(),
        AppConfig::default(),
        Arc::clone(&broker),
        Arc::clone(&orchestrator),
        Arc::clone(&state_backend),
        Arc::new(registry),
        None,
    )
    .with_num_workers(4)
    .with_idle_sleep(10);

    let task_id = TaskId::new("test", "double");
    let mut inv_ids = Vec::new();
    for i in 0..8 {
        let mut args = SerializedArguments::new();
        args.insert("x", &i.to_string());
        let call = CallDTO::new(task_id.clone(), args);
        let inv_id = orchestrator.register_invocation(&call).await.unwrap();
        let inv_dto = InvocationDTO::new(inv_id.clone(), task_id.clone(), call.call_id.clone());
        state_backend
            .upsert_invocation(&inv_dto, &call)
            .await
            .unwrap();
        broker.route_invocation(&inv_id).await.unwrap();
        inv_ids.push(inv_id);
    }

    let _ = runner
        .with_graceful_shutdown(async {
            tokio::time::sleep(Duration::from_millis(500)).await;
        })
        .await;

    for inv_id in &inv_ids {
        let status = orchestrator.get_invocation_status(inv_id).await.unwrap();
        assert_eq!(
            status.status,
            InvocationStatus::Success,
            "Invocation {} should be Success",
            inv_id
        );
    }
}

#[test]
fn test_runner_cls() {
    let (runner, _, _) = make_runner();
    assert_eq!(runner.runner_cls(), "PersistentTokioRunner");
}

#[test]
fn test_max_parallel_slots() {
    let (runner, _, _) = make_runner();
    let runner = runner.with_num_workers(8);
    assert_eq!(runner.max_parallel_slots(), 8);
}

// ===================================================================
// Blocking priority tests
// ===================================================================

#[tokio::test]
async fn test_blocking_priority_over_fifo() {
    let broker: Arc<dyn Broker> = Arc::new(rustvello_mem::broker::MemBroker::new());
    let orchestrator: Arc<dyn InvocationControlBackend> =
        Arc::new(rustvello_mem::orchestrator::MemOrchestrator::new());
    let state_backend: Arc<dyn StateBackend> =
        Arc::new(rustvello_mem::state_backend::MemStateBackend::new());

    let task_id = TaskId::new("test", "double");
    let mut registry = TaskRegistry::new();
    registry
        .register(TaskDefinition::new(
            task_id.clone(),
            TaskConfig::default(),
            Arc::new(|args_json: String| {
                let args: std::collections::BTreeMap<String, String> =
                    serde_json::from_str(&args_json).map_err(|e| {
                        RustvelloError::Serialization {
                            message: e.to_string(),
                        }
                    })?;
                let x: i64 = args.get("x").and_then(|v| v.parse().ok()).unwrap_or(0);
                serde_json::to_string(&(x * 2)).map_err(|e| RustvelloError::Serialization {
                    message: e.to_string(),
                })
            }),
        ))
        .unwrap();

    let runner = PersistentTokioRunner::new(
        "test-app".to_string(),
        AppConfig::default(),
        Arc::clone(&broker),
        Arc::clone(&orchestrator),
        Arc::clone(&state_backend),
        Arc::new(registry),
        None,
    );

    let mut args1 = SerializedArguments::new();
    args1.insert("x", "10");
    let call1 = CallDTO::new(task_id.clone(), args1);
    let regular_inv = orchestrator.register_invocation(&call1).await.unwrap();
    let dto1 = InvocationDTO::new(regular_inv.clone(), task_id.clone(), call1.call_id.clone());
    state_backend
        .upsert_invocation(&dto1, &call1)
        .await
        .unwrap();
    broker.route_invocation(&regular_inv).await.unwrap();

    let mut args2 = SerializedArguments::new();
    args2.insert("x", "21");
    let call2 = CallDTO::new(task_id.clone(), args2);
    let blocking_inv = orchestrator.register_invocation(&call2).await.unwrap();
    let dto2 = InvocationDTO::new(blocking_inv.clone(), task_id.clone(), call2.call_id.clone());
    state_backend
        .upsert_invocation(&dto2, &call2)
        .await
        .unwrap();
    broker.route_invocation(&blocking_inv).await.unwrap();

    let call3 = CallDTO::new(task_id.clone(), SerializedArguments::new());
    let waiter_inv = orchestrator.register_invocation(&call3).await.unwrap();
    orchestrator
        .set_waiting_for(&waiter_inv, &blocking_inv)
        .await
        .unwrap();

    runner.run_one().await.unwrap();

    let blocking_status = orchestrator
        .get_invocation_status(&blocking_inv)
        .await
        .unwrap();
    assert_eq!(
        blocking_status.status,
        InvocationStatus::Success,
        "Blocking invocation should be prioritized and executed first"
    );

    let regular_status = orchestrator
        .get_invocation_status(&regular_inv)
        .await
        .unwrap();
    assert_eq!(
        regular_status.status,
        InvocationStatus::Registered,
        "Regular invocation should not have been executed yet"
    );

    let result = state_backend.get_result(&blocking_inv).await.unwrap();
    assert_eq!(result, Some("42".to_string()));
}

#[tokio::test]
async fn test_no_blocking_falls_back_to_fifo() {
    let (runner, orchestrator, state_backend) = make_runner();

    let task_id = TaskId::new("test", "double");
    let mut args = SerializedArguments::new();
    args.insert("x", "5");
    let call = CallDTO::new(task_id.clone(), args);

    let inv_id = orchestrator.register_invocation(&call).await.unwrap();
    let inv_dto = InvocationDTO::new(inv_id.clone(), task_id, call.call_id.clone());
    state_backend
        .upsert_invocation(&inv_dto, &call)
        .await
        .unwrap();
    runner
        .control_plane
        .broker
        .route_invocation(&inv_id)
        .await
        .unwrap();

    let did_work = runner.run_one().await.unwrap();
    assert!(did_work);

    let status = orchestrator.get_invocation_status(&inv_id).await.unwrap();
    assert_eq!(status.status, InvocationStatus::Success);
    let result = state_backend.get_result(&inv_id).await.unwrap();
    assert_eq!(result, Some("10".to_string()));
}

#[tokio::test]
async fn test_blocking_race_handled_gracefully() {
    let broker: Arc<dyn Broker> = Arc::new(rustvello_mem::broker::MemBroker::new());
    let orchestrator: Arc<dyn InvocationControlBackend> =
        Arc::new(rustvello_mem::orchestrator::MemOrchestrator::new());
    let state_backend: Arc<dyn StateBackend> =
        Arc::new(rustvello_mem::state_backend::MemStateBackend::new());

    let task_id = TaskId::new("test", "double");
    let mut registry = TaskRegistry::new();
    registry
        .register(TaskDefinition::new(
            task_id.clone(),
            TaskConfig::default(),
            Arc::new(|args_json: String| {
                let args: std::collections::BTreeMap<String, String> =
                    serde_json::from_str(&args_json).map_err(|e| {
                        RustvelloError::Serialization {
                            message: e.to_string(),
                        }
                    })?;
                let x: i64 = args.get("x").and_then(|v| v.parse().ok()).unwrap_or(0);
                serde_json::to_string(&(x * 2)).map_err(|e| RustvelloError::Serialization {
                    message: e.to_string(),
                })
            }),
        ))
        .unwrap();

    let runner = PersistentTokioRunner::new(
        "test-app".to_string(),
        AppConfig::default(),
        Arc::clone(&broker),
        Arc::clone(&orchestrator),
        Arc::clone(&state_backend),
        Arc::new(registry),
        None,
    );

    let mut args = SerializedArguments::new();
    args.insert("x", "7");
    let call = CallDTO::new(task_id.clone(), args);
    let blocking_inv = orchestrator.register_invocation(&call).await.unwrap();
    let dto = InvocationDTO::new(blocking_inv.clone(), task_id.clone(), call.call_id.clone());
    state_backend.upsert_invocation(&dto, &call).await.unwrap();
    broker.route_invocation(&blocking_inv).await.unwrap();

    let call2 = CallDTO::new(task_id.clone(), SerializedArguments::new());
    let waiter = orchestrator.register_invocation(&call2).await.unwrap();
    orchestrator
        .set_waiting_for(&waiter, &blocking_inv)
        .await
        .unwrap();

    runner.run_one().await.unwrap();
    let status = orchestrator
        .get_invocation_status(&blocking_inv)
        .await
        .unwrap();
    assert_eq!(status.status, InvocationStatus::Success);

    runner.run_one().await.unwrap();
}

/// Helper: create a runner with a custom task config for the "double" task.
fn make_runner_with_config(
    task_config: TaskConfig,
) -> (
    PersistentTokioRunner,
    Arc<dyn InvocationControlBackend>,
    Arc<dyn StateBackend>,
    Arc<dyn Broker>,
) {
    let broker: Arc<dyn Broker> = Arc::new(rustvello_mem::broker::MemBroker::new());
    let orchestrator: Arc<dyn InvocationControlBackend> =
        Arc::new(rustvello_mem::orchestrator::MemOrchestrator::new());
    let state_backend: Arc<dyn StateBackend> =
        Arc::new(rustvello_mem::state_backend::MemStateBackend::new());

    let task_id = TaskId::new("test", "double");
    let mut registry = TaskRegistry::new();
    registry
        .register(TaskDefinition::new(
            task_id.clone(),
            task_config,
            Arc::new(|args_json: String| {
                let args: std::collections::BTreeMap<String, String> =
                    serde_json::from_str(&args_json).map_err(|e| {
                        RustvelloError::Serialization {
                            message: e.to_string(),
                        }
                    })?;
                let x: i64 = args.get("x").and_then(|v| v.parse().ok()).unwrap_or(0);
                serde_json::to_string(&(x * 2)).map_err(|e| RustvelloError::Serialization {
                    message: e.to_string(),
                })
            }),
        ))
        .unwrap();

    let runner = PersistentTokioRunner::new(
        "test-app".to_string(),
        AppConfig::default(),
        Arc::clone(&broker),
        Arc::clone(&orchestrator),
        Arc::clone(&state_backend),
        Arc::new(registry),
        None,
    );

    (runner, orchestrator, state_backend, broker)
}

/// An invocation registered and stored but not queued: it already runs on
/// another runner, so this runner must never dequeue it.
async fn register_running_elsewhere(
    orchestrator: &dyn InvocationControlBackend,
    state_backend: &dyn StateBackend,
    task_id: &TaskId,
    args: SerializedArguments,
) -> InvocationId {
    let call = CallDTO::new(task_id.clone(), args);
    let inv_id = orchestrator.register_invocation(&call).await.unwrap();
    let dto = InvocationDTO::new(inv_id.clone(), task_id.clone(), call.call_id.clone());
    state_backend.upsert_invocation(&dto, &call).await.unwrap();
    inv_id
}

async fn submit_invocation(
    orchestrator: &dyn InvocationControlBackend,
    state_backend: &dyn StateBackend,
    broker: &dyn Broker,
    task_id: &TaskId,
    args: SerializedArguments,
) -> InvocationId {
    let call = CallDTO::new(task_id.clone(), args);
    let inv_id = orchestrator.register_invocation(&call).await.unwrap();
    let dto = InvocationDTO::new(inv_id.clone(), task_id.clone(), call.call_id.clone());
    state_backend.upsert_invocation(&dto, &call).await.unwrap();
    broker
        .route_invocation_for_task(&inv_id, task_id)
        .await
        .unwrap();
    inv_id
}

#[tokio::test]
async fn test_cc_task_level_blocks_second_invocation() {
    let mut config = TaskConfig::default();
    config.concurrency_control = ConcurrencyControlType::Task;
    config.running_concurrency = Some(1);
    let (runner, orchestrator, state_backend, broker) = make_runner_with_config(config);

    let task_id = TaskId::new("test", "double");
    let mut args1 = SerializedArguments::new();
    args1.insert("x", "1");
    let mut args2 = SerializedArguments::new();
    args2.insert("x", "2");

    let inv1 = submit_invocation(&*orchestrator, &*state_backend, &*broker, &task_id, args1).await;
    let inv2 = submit_invocation(&*orchestrator, &*state_backend, &*broker, &task_id, args2).await;

    let did_work = runner.run_one().await.unwrap();
    assert!(did_work);

    let s1 = orchestrator.get_invocation_status(&inv1).await.unwrap();
    assert_eq!(s1.status, InvocationStatus::Success);

    let did_work = runner.run_one().await.unwrap();
    assert!(did_work);

    let s2 = orchestrator.get_invocation_status(&inv2).await.unwrap();
    assert_eq!(s2.status, InvocationStatus::Success);
}

#[tokio::test]
async fn test_cc_reroute_on_cc_re_enqueues() {
    let mut config = TaskConfig::default();
    config.concurrency_control = ConcurrencyControlType::Task;
    config.running_concurrency = Some(1);
    config.reroute_on_cc = true;
    let (runner, orchestrator, state_backend, broker) = make_runner_with_config(config);

    let task_id = TaskId::new("test", "double");
    let mut args1 = SerializedArguments::new();
    args1.insert("x", "5");
    let mut args2 = SerializedArguments::new();
    args2.insert("x", "6");

    let inv1 = register_running_elsewhere(&*orchestrator, &*state_backend, &task_id, args1).await;
    let _inv2 = submit_invocation(&*orchestrator, &*state_backend, &*broker, &task_id, args2).await;

    let fake_runner = RunnerId::from_string("fake-runner");
    orchestrator
        .index_for_concurrency_control(&inv1, &task_id, Some(&SerializedArguments::new()))
        .await
        .unwrap();
    orchestrator
        .set_invocation_status(&inv1, InvocationStatus::Pending, Some(&fake_runner))
        .await
        .unwrap();
    orchestrator
        .set_invocation_status(&inv1, InvocationStatus::Running, Some(&fake_runner))
        .await
        .unwrap();

    let did_work = runner.run_one().await.unwrap();
    assert!(!did_work);

    let s2 = orchestrator.get_invocation_status(&_inv2).await.unwrap();
    assert_eq!(s2.status, InvocationStatus::Rerouted);
}

#[tokio::test]
async fn test_cc_final_rejection_without_reroute() {
    let mut config = TaskConfig::default();
    config.concurrency_control = ConcurrencyControlType::Task;
    config.running_concurrency = Some(1);
    config.reroute_on_cc = false;
    let (runner, orchestrator, state_backend, broker) = make_runner_with_config(config);

    let task_id = TaskId::new("test", "double");
    let mut args1 = SerializedArguments::new();
    args1.insert("x", "10");
    let mut args2 = SerializedArguments::new();
    args2.insert("x", "20");

    let inv1 = register_running_elsewhere(&*orchestrator, &*state_backend, &task_id, args1).await;
    let inv2 = submit_invocation(&*orchestrator, &*state_backend, &*broker, &task_id, args2).await;

    let fake_runner = RunnerId::from_string("fake-runner");
    orchestrator
        .index_for_concurrency_control(&inv1, &task_id, Some(&SerializedArguments::new()))
        .await
        .unwrap();
    orchestrator
        .set_invocation_status(&inv1, InvocationStatus::Pending, Some(&fake_runner))
        .await
        .unwrap();
    orchestrator
        .set_invocation_status(&inv1, InvocationStatus::Running, Some(&fake_runner))
        .await
        .unwrap();

    let did_work = runner.run_one().await.unwrap();
    assert!(!did_work);

    let s2 = orchestrator.get_invocation_status(&inv2).await.unwrap();
    assert_eq!(s2.status, InvocationStatus::ConcurrencyControlledFinal);
}

#[tokio::test]
async fn test_cc_unlimited_allows_all() {
    let config = TaskConfig::default();
    let (runner, orchestrator, state_backend, broker) = make_runner_with_config(config);

    let task_id = TaskId::new("test", "double");
    let mut args1 = SerializedArguments::new();
    args1.insert("x", "1");
    let mut args2 = SerializedArguments::new();
    args2.insert("x", "2");

    submit_invocation(&*orchestrator, &*state_backend, &*broker, &task_id, args1).await;
    submit_invocation(&*orchestrator, &*state_backend, &*broker, &task_id, args2).await;

    assert!(runner.run_one().await.unwrap());
    assert!(runner.run_one().await.unwrap());
}

#[tokio::test]
async fn test_cc_index_cleanup_on_success() {
    let mut config = TaskConfig::default();
    config.concurrency_control = ConcurrencyControlType::Task;
    config.running_concurrency = Some(1);
    let (runner, orchestrator, state_backend, broker) = make_runner_with_config(config);

    let task_id = TaskId::new("test", "double");

    for i in 0..3 {
        let mut args = SerializedArguments::new();
        args.insert("x", &i.to_string());
        let inv =
            submit_invocation(&*orchestrator, &*state_backend, &*broker, &task_id, args).await;

        let did_work = runner.run_one().await.unwrap();
        assert!(did_work, "Invocation {} should have executed", i);

        let status = orchestrator.get_invocation_status(&inv).await.unwrap();
        assert_eq!(status.status, InvocationStatus::Success);
    }
}

#[tokio::test]
async fn test_trigger_static_string_args_are_json_serialized_for_tasks() {
    let broker: Arc<dyn Broker> = Arc::new(rustvello_mem::broker::MemBroker::new());
    let orchestrator: Arc<dyn InvocationControlBackend> =
        Arc::new(rustvello_mem::orchestrator::MemOrchestrator::new());
    let state_backend: Arc<dyn StateBackend> =
        Arc::new(rustvello_mem::state_backend::MemStateBackend::new());
    let trigger_store: Arc<dyn rustvello_core::trigger::TriggerStore> =
        Arc::new(rustvello_mem::trigger::MemTriggerStore::new());
    let trigger_manager = rustvello_core::trigger::TriggerManager::new(Arc::clone(&trigger_store));

    let task_id = TaskId::new("test", "echo_message");
    let mut registry = TaskRegistry::new();
    registry
        .register(TaskDefinition::new(
            task_id.clone(),
            TaskConfig::default(),
            Arc::new(|args_json: String| {
                let args: std::collections::BTreeMap<String, String> =
                    serde_json::from_str(&args_json).map_err(|e| {
                        RustvelloError::Serialization {
                            message: e.to_string(),
                        }
                    })?;
                let raw = args
                    .get("message")
                    .ok_or_else(|| RustvelloError::Serialization {
                        message: "missing message argument".to_owned(),
                    })?;
                let message: String =
                    serde_json::from_str(raw).map_err(|e| RustvelloError::Serialization {
                        message: e.to_string(),
                    })?;
                serde_json::to_string(&message).map_err(|e| RustvelloError::Serialization {
                    message: e.to_string(),
                })
            }),
        ))
        .unwrap();

    crate::trigger_builder::TriggerBuilder::new()
        .on_event("demo.event")
        .with_static_args(serde_json::json!({"message": "hello"}))
        .build_and_register(&task_id, &trigger_store)
        .await
        .unwrap();

    trigger_manager
        .emit_event("demo.event", serde_json::json!({}))
        .await
        .unwrap();

    let runner = PersistentTokioRunner::new(
        "test-app".to_string(),
        AppConfig::default(),
        Arc::clone(&broker),
        Arc::clone(&orchestrator),
        Arc::clone(&state_backend),
        Arc::new(registry),
        Some(trigger_manager.clone()),
    );

    runner.evaluate_triggers(&trigger_manager).await.unwrap();

    let invocations = orchestrator
        .get_invocations_by_task(&task_id)
        .await
        .unwrap();
    assert_eq!(invocations.len(), 1);
    assert!(runner.run_one().await.unwrap());

    let result = state_backend.get_result(&invocations[0]).await.unwrap();
    assert_eq!(result, Some("\"hello\"".to_owned()));
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn test_cc_slot_is_atomic_across_runner_instances() {
    let broker: Arc<dyn Broker> = Arc::new(rustvello_mem::broker::MemBroker::new());
    let orchestrator: Arc<dyn InvocationControlBackend> =
        Arc::new(rustvello_mem::orchestrator::MemOrchestrator::new());
    let state_backend: Arc<dyn StateBackend> =
        Arc::new(rustvello_mem::state_backend::MemStateBackend::new());
    let active = Arc::new(AtomicUsize::new(0));
    let peak = Arc::new(AtomicUsize::new(0));
    let first_started = Arc::new(AtomicBool::new(false));
    let release_first = Arc::new(AtomicBool::new(false));

    let mut config = TaskConfig::default();
    config.concurrency_control = ConcurrencyControlType::Task;
    config.running_concurrency = Some(1);

    let mut registry = TaskRegistry::new();
    let active_for_task = Arc::clone(&active);
    let peak_for_task = Arc::clone(&peak);
    let started_for_task = Arc::clone(&first_started);
    let release_for_task = Arc::clone(&release_first);
    registry
        .register(TaskDefinition::new(
            TaskId::new("test", "contended"),
            config,
            Arc::new(move |_args: String| {
                let now = active_for_task.fetch_add(1, Ordering::SeqCst) + 1;
                peak_for_task.fetch_max(now, Ordering::SeqCst);
                if now == 1 {
                    started_for_task.store(true, Ordering::SeqCst);
                    while !release_for_task.load(Ordering::SeqCst) {
                        std::thread::sleep(Duration::from_millis(1));
                    }
                }
                active_for_task.fetch_sub(1, Ordering::SeqCst);
                Ok("null".to_string())
            }),
        ))
        .unwrap();

    let registry = Arc::new(registry);
    let runner_a = PersistentTokioRunner::new(
        "atomic-cc".to_string(),
        AppConfig::default(),
        Arc::clone(&broker),
        Arc::clone(&orchestrator),
        Arc::clone(&state_backend),
        Arc::clone(&registry),
        None,
    );
    let runner_b = PersistentTokioRunner::new(
        "atomic-cc".to_string(),
        AppConfig::default(),
        Arc::clone(&broker),
        Arc::clone(&orchestrator),
        Arc::clone(&state_backend),
        registry,
        None,
    );

    let task_id = TaskId::new("test", "contended");
    let first = submit_invocation(
        &*orchestrator,
        &*state_backend,
        &*broker,
        &task_id,
        SerializedArguments::new(),
    )
    .await;
    let second = submit_invocation(
        &*orchestrator,
        &*state_backend,
        &*broker,
        &task_id,
        SerializedArguments::new(),
    )
    .await;

    let a = tokio::spawn(async move { runner_a.run_one().await });
    tokio::time::timeout(Duration::from_secs(2), async {
        while !first_started.load(Ordering::SeqCst) {
            tokio::task::yield_now().await;
        }
    })
    .await
    .expect("first runner did not start its invocation");
    let b = tokio::spawn(async move { runner_b.run_one().await });
    let b_result = tokio::time::timeout(Duration::from_secs(2), b).await;
    release_first.store(true, Ordering::SeqCst);
    a.await.unwrap().unwrap();
    b_result
        .expect("second runner did not complete its CC check")
        .unwrap()
        .unwrap();

    let statuses = [
        orchestrator
            .get_invocation_status(&first)
            .await
            .unwrap()
            .status,
        orchestrator
            .get_invocation_status(&second)
            .await
            .unwrap()
            .status,
    ];
    assert_eq!(peak.load(Ordering::SeqCst), 1);
    assert_eq!(
        statuses
            .iter()
            .filter(|status| **status == InvocationStatus::Success)
            .count(),
        1
    );
    assert_eq!(
        statuses
            .iter()
            .filter(|status| **status == InvocationStatus::ConcurrencyControlledFinal)
            .count(),
        1
    );
}

/// A broker whose `wait_for_work` only returns on cancellation, like a
/// database broker with a very slow poll interval.
struct NeverWakesBroker(rustvello_mem::broker::MemBroker);

#[async_trait::async_trait]
impl Broker for NeverWakesBroker {
    async fn route_invocation_with_options(
        &self,
        invocation_id: &InvocationId,
        task_id: Option<&TaskId>,
        queue_name: &str,
        priority: f64,
    ) -> RustvelloResult<()> {
        self.0
            .route_invocation_with_options(invocation_id, task_id, queue_name, priority)
            .await
    }

    async fn retrieve_invocation_from_queue(
        &self,
        queue_name: &str,
        task_id: Option<&TaskId>,
    ) -> RustvelloResult<Option<InvocationId>> {
        self.0
            .retrieve_invocation_from_queue(queue_name, task_id)
            .await
    }

    async fn retrieve_invocation_for_language_from_queue(
        &self,
        language: rustvello_proto::identifiers::TaskLanguage,
        queue_name: &str,
    ) -> RustvelloResult<Option<InvocationId>> {
        self.0
            .retrieve_invocation_for_language_from_queue(language, queue_name)
            .await
    }

    async fn count_invocations_in_queues(
        &self,
        queue_names: &[String],
        task_id: Option<&TaskId>,
    ) -> RustvelloResult<usize> {
        self.0
            .count_invocations_in_queues(queue_names, task_id)
            .await
    }

    async fn retrieve_invocation_for_language(
        &self,
        language: rustvello_proto::identifiers::TaskLanguage,
    ) -> RustvelloResult<Option<InvocationId>> {
        self.0.retrieve_invocation_for_language(language).await
    }

    async fn route_invocation(&self, invocation_id: &InvocationId) -> RustvelloResult<()> {
        self.0.route_invocation(invocation_id).await
    }

    async fn route_invocation_for_task(
        &self,
        invocation_id: &InvocationId,
        task_id: &TaskId,
    ) -> RustvelloResult<()> {
        self.0
            .route_invocation_for_task(invocation_id, task_id)
            .await
    }

    async fn retrieve_invocation(
        &self,
        task_id: Option<&TaskId>,
    ) -> RustvelloResult<Option<InvocationId>> {
        self.0.retrieve_invocation(task_id).await
    }

    async fn wait_for_work(&self, cancel: &tokio_util::sync::CancellationToken) -> bool {
        cancel.cancelled().await;
        false
    }

    async fn count_invocations(&self, task_id: Option<&TaskId>) -> RustvelloResult<usize> {
        self.0.count_invocations(task_id).await
    }

    async fn purge(&self, task_id: Option<&TaskId>) -> RustvelloResult<()> {
        self.0.purge(task_id).await
    }
}

#[tokio::test]
async fn idle_wait_is_capped_by_idle_sleep_ms() {
    let runner = PersistentTokioRunner::new(
        "test-app".to_string(),
        AppConfig::default(),
        Arc::new(NeverWakesBroker(rustvello_mem::broker::MemBroker::new())),
        Arc::new(rustvello_mem::orchestrator::MemOrchestrator::new()),
        Arc::new(rustvello_mem::state_backend::MemStateBackend::new()),
        Arc::new(TaskRegistry::new()),
        None,
    )
    .with_idle_sleep(20);
    let cancel = tokio_util::sync::CancellationToken::new();
    let started = std::time::Instant::now();
    assert!(
        runner.wait_idle(&cancel).await,
        "keep polling after the cap"
    );
    assert!(
        started.elapsed() < Duration::from_secs(2),
        "idle wait ignored idle_sleep_ms: {:?}",
        started.elapsed()
    );
    cancel.cancel();
    assert!(!runner.wait_idle(&cancel).await, "stop once cancelled");
}

// --- Recovery must never orphan an invocation (unknown task, interrupted
// recovery, a task whose worker dies every time). ---

struct RecoveryFixture {
    runner: PersistentTokioRunner,
    mem_orchestrator: Arc<rustvello_mem::orchestrator::MemOrchestrator>,
    orchestrator: Arc<dyn InvocationControlBackend>,
    state_backend: Arc<dyn StateBackend>,
}

fn recovery_fixture(config: AppConfig) -> RecoveryFixture {
    let broker: Arc<dyn Broker> = Arc::new(rustvello_mem::broker::MemBroker::new());
    let mem_orchestrator = Arc::new(rustvello_mem::orchestrator::MemOrchestrator::new());
    let orchestrator: Arc<dyn InvocationControlBackend> = Arc::clone(&mem_orchestrator) as _;
    let state_backend: Arc<dyn StateBackend> =
        Arc::new(rustvello_mem::state_backend::MemStateBackend::new());
    recovery_fixture_with(
        config,
        broker,
        mem_orchestrator,
        orchestrator,
        state_backend,
    )
}

fn recovery_fixture_with(
    config: AppConfig,
    broker: Arc<dyn Broker>,
    mem_orchestrator: Arc<rustvello_mem::orchestrator::MemOrchestrator>,
    orchestrator: Arc<dyn InvocationControlBackend>,
    state_backend: Arc<dyn StateBackend>,
) -> RecoveryFixture {
    let mut registry = TaskRegistry::new();
    registry
        .register(TaskDefinition::new(
            TaskId::new("test", "known"),
            TaskConfig::default(),
            Arc::new(|_args: String| Ok("1".to_owned())),
        ))
        .unwrap();
    let runner = PersistentTokioRunner::new(
        "test-app".to_string(),
        config,
        broker,
        Arc::clone(&orchestrator),
        Arc::clone(&state_backend),
        Arc::new(registry),
        None,
    );
    RecoveryFixture {
        runner,
        mem_orchestrator,
        orchestrator,
        state_backend,
    }
}

/// An invocation of `task` left Running by a runner that never heartbeats.
async fn orphaned_running(fixture: &RecoveryFixture, task: &str) -> InvocationId {
    let call = CallDTO::new(TaskId::new("test", task), SerializedArguments::new());
    let inv_id = fixture
        .orchestrator
        .register_invocation(&call)
        .await
        .unwrap();
    let dto = InvocationDTO::new(
        inv_id.clone(),
        TaskId::new("test", task),
        call.call_id.clone(),
    );
    fixture
        .state_backend
        .upsert_invocation(&dto, &call)
        .await
        .unwrap();
    let dead = RunnerId::from_string("dead-runner");
    for status in [InvocationStatus::Pending, InvocationStatus::Running] {
        fixture
            .orchestrator
            .set_invocation_status(&inv_id, status, Some(&dead))
            .await
            .unwrap();
    }
    // Old enough for recovery: the dead runner never heartbeated.
    fixture
        .mem_orchestrator
        .backdate_status_for_testing(&inv_id, chrono::Duration::seconds(600))
        .await;
    inv_id
}

#[tokio::test]
async fn recovery_leaves_unknown_tasks_untouched_and_continues_the_pass() {
    let fixture = recovery_fixture(AppConfig::default());
    let unknown = orphaned_running(&fixture, "unknown").await;
    let known = orphaned_running(&fixture, "known").await;

    // The unknown task used to abort the pass after writing RUNNING_RECOVERY.
    let recovered = fixture.runner.recover_stale_invocations().await.unwrap();
    assert_eq!(recovered, 1, "the known task is still recovered");
    let status = |id: InvocationId| {
        let orchestrator = Arc::clone(&fixture.orchestrator);
        async move {
            orchestrator
                .get_invocation_status(&id)
                .await
                .unwrap()
                .status
        }
    };
    assert_eq!(status(known).await, InvocationStatus::Rerouted);
    assert_eq!(
        status(unknown).await,
        InvocationStatus::Running,
        "nothing written for a task this runner cannot route"
    );
}

#[tokio::test]
async fn interrupted_recovery_is_resumed_once_stale() {
    let fixture = recovery_fixture(AppConfig::default());
    let inv_id = orphaned_running(&fixture, "known").await;
    // A recovery that wrote its first status and then died.
    fixture
        .orchestrator
        .set_invocation_status(
            &inv_id,
            InvocationStatus::RunningRecovery,
            Some(&RunnerId::from_string("crashed-recoverer")),
        )
        .await
        .unwrap();
    fixture
        .mem_orchestrator
        .backdate_status_for_testing(&inv_id, chrono::Duration::seconds(600))
        .await;

    let recovered = fixture.runner.recover_stale_invocations().await.unwrap();
    assert_eq!(recovered, 1);
    let status = fixture
        .orchestrator
        .get_invocation_status(&inv_id)
        .await
        .unwrap();
    assert_eq!(status.status, InvocationStatus::Rerouted);
}

#[tokio::test]
async fn recovery_limit_fails_the_invocation_with_a_stored_error() {
    let mut config = AppConfig::default();
    config.max_invocation_recoveries = 2;
    let fixture = recovery_fixture(config);
    // Works for a task no runner knows, too: the limit needs no route.
    let inv_id = orphaned_running(&fixture, "unknown").await;
    for _ in 0..2 {
        fixture
            .state_backend
            .add_history(&rustvello_proto::invocation::InvocationHistory::new(
                inv_id.clone(),
                rustvello_proto::status::InvocationStatusRecord::new(
                    InvocationStatus::RunningRecovery,
                    None,
                ),
                None,
            ))
            .await
            .unwrap();
    }

    let recovered = fixture.runner.recover_stale_invocations().await.unwrap();
    assert_eq!(recovered, 0);
    let status = fixture
        .orchestrator
        .get_invocation_status(&inv_id)
        .await
        .unwrap();
    assert_eq!(status.status, InvocationStatus::Failed);
    let error = fixture
        .state_backend
        .get_error(&inv_id)
        .await
        .unwrap()
        .expect("stored error");
    assert_eq!(error.error_type, "RecoveryLimitExceeded");
}

/// The same three cases on SQLite, whose recovery takes the atomic
/// publication path instead of two status writes.
#[cfg(feature = "sqlite")]
mod sqlite_recovery {
    use super::*;

    fn sqlite_fixture(dir: &std::path::Path, config: AppConfig) -> RecoveryFixture {
        let db = Arc::new(
            rustvello_sqlite::db::Database::open(dir.join("recovery.db"), "test-app").unwrap(),
        );
        let broker: Arc<dyn Broker> =
            Arc::new(rustvello_sqlite::broker::SqliteBroker::new(Arc::clone(&db)));
        let orchestrator: Arc<dyn InvocationControlBackend> = Arc::new(
            rustvello_sqlite::orchestrator::SqliteOrchestrator::new(Arc::clone(&db)),
        );
        let state_backend: Arc<dyn StateBackend> =
            Arc::new(rustvello_sqlite::state_backend::SqliteStateBackend::new(db));
        assert!(
            orchestrator.runtime_publication().is_some(),
            "SQLite recovery must exercise the publication path"
        );
        recovery_fixture_with(
            config,
            broker,
            Arc::new(rustvello_mem::orchestrator::MemOrchestrator::new()),
            orchestrator,
            state_backend,
        )
    }

    async fn status(fixture: &RecoveryFixture, id: &InvocationId) -> InvocationStatus {
        fixture
            .orchestrator
            .get_invocation_status(id)
            .await
            .unwrap()
            .status
    }

    #[tokio::test]
    async fn unknown_task_is_left_untouched_and_the_pass_continues() {
        let dir = tempfile::tempdir().unwrap();
        let mut config = AppConfig::default();
        config.runner_dead_after_seconds = 0;
        let fixture = sqlite_fixture(dir.path(), config);
        let unknown = orphaned_running(&fixture, "unknown").await;
        let known = orphaned_running(&fixture, "known").await;
        assert_eq!(fixture.runner.recover_stale_invocations().await.unwrap(), 1);
        assert_eq!(status(&fixture, &known).await, InvocationStatus::Rerouted);
        assert_eq!(status(&fixture, &unknown).await, InvocationStatus::Running);
    }

    #[tokio::test]
    async fn interrupted_recovery_is_resumed() {
        let dir = tempfile::tempdir().unwrap();
        let mut config = AppConfig::default();
        // Every recovery status is already "stale": no clock to backdate here.
        config.runner_dead_after_seconds = 0;
        let fixture = sqlite_fixture(dir.path(), config);
        let inv_id = orphaned_running(&fixture, "known").await;
        fixture
            .orchestrator
            .set_invocation_status(
                &inv_id,
                InvocationStatus::RunningRecovery,
                Some(&RunnerId::from_string("crashed-recoverer")),
            )
            .await
            .unwrap();
        tokio::time::sleep(Duration::from_millis(20)).await;
        assert_eq!(fixture.runner.recover_stale_invocations().await.unwrap(), 1);
        assert_eq!(status(&fixture, &inv_id).await, InvocationStatus::Rerouted);
    }

    #[tokio::test]
    async fn recovery_limit_fails_with_a_stored_error() {
        let dir = tempfile::tempdir().unwrap();
        let mut config = AppConfig::default();
        config.max_invocation_recoveries = 1;
        config.runner_dead_after_seconds = 0;
        let fixture = sqlite_fixture(dir.path(), config);
        let inv_id = orphaned_running(&fixture, "known").await;
        fixture
            .state_backend
            .add_history(&rustvello_proto::invocation::InvocationHistory::new(
                inv_id.clone(),
                rustvello_proto::status::InvocationStatusRecord::new(
                    InvocationStatus::RunningRecovery,
                    None,
                ),
                None,
            ))
            .await
            .unwrap();
        assert_eq!(fixture.runner.recover_stale_invocations().await.unwrap(), 0);
        assert_eq!(status(&fixture, &inv_id).await, InvocationStatus::Failed);
        let error = fixture
            .state_backend
            .get_error(&inv_id)
            .await
            .unwrap()
            .unwrap();
        assert_eq!(error.error_type, "RecoveryLimitExceeded");
    }
}

/// Client report: a fresh worker's running invocation was rerouted by another
/// pod's recovery; when the task finished, the completion was rejected
/// (`REROUTED -> SUCCESS`) and the concurrency slot leaked for good, so the
/// task could never run again. The attempt must end cleanly and free the slot.
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn rejected_completion_after_a_reroute_releases_the_slot() {
    let mut config = TaskConfig::default();
    config.concurrency_control = ConcurrencyControlType::Task;
    config.running_concurrency = Some(1);

    let broker: Arc<dyn Broker> = Arc::new(rustvello_mem::broker::MemBroker::new());
    let orchestrator: Arc<dyn InvocationControlBackend> =
        Arc::new(rustvello_mem::orchestrator::MemOrchestrator::new());
    let state_backend: Arc<dyn StateBackend> =
        Arc::new(rustvello_mem::state_backend::MemStateBackend::new());
    let task_id = TaskId::new("test", "slow");
    let mut registry = TaskRegistry::new();
    registry
        .register(TaskDefinition::new(
            task_id.clone(),
            config.clone(),
            Arc::new(|_args: String| {
                std::thread::sleep(Duration::from_millis(400));
                Ok("1".to_owned())
            }),
        ))
        .unwrap();
    let runner = Arc::new(PersistentTokioRunner::new(
        "test-app".to_string(),
        AppConfig::default(),
        Arc::clone(&broker),
        Arc::clone(&orchestrator),
        Arc::clone(&state_backend),
        Arc::new(registry),
        None,
    ));

    let inv_id = submit_invocation(
        &*orchestrator,
        &*state_backend,
        &*broker,
        &task_id,
        SerializedArguments::new(),
    )
    .await;
    let attempt = {
        let runner = Arc::clone(&runner);
        tokio::spawn(async move { runner.run_one().await })
    };
    for _ in 0..200 {
        if orchestrator
            .get_invocation_status(&inv_id)
            .await
            .unwrap()
            .status
            == InvocationStatus::Running
        {
            break;
        }
        tokio::time::sleep(Duration::from_millis(5)).await;
    }

    // Another pod's recovery takes the invocation over while it runs.
    let recoverer = RunnerId::from_string("other-pod-recovery");
    for status in [
        InvocationStatus::RunningRecovery,
        InvocationStatus::Rerouted,
    ] {
        orchestrator
            .set_invocation_status(&inv_id, status, Some(&recoverer))
            .await
            .unwrap();
    }

    let outcome = attempt.await.unwrap();
    assert!(
        outcome.is_ok(),
        "a rejected completion is a discarded outcome, not a worker error: {outcome:?}"
    );
    assert_eq!(
        orchestrator
            .get_invocation_status(&inv_id)
            .await
            .unwrap()
            .status,
        InvocationStatus::Rerouted
    );

    // The slot is free: another invocation of the task can take it.
    let other = submit_invocation(
        &*orchestrator,
        &*state_backend,
        &*broker,
        &task_id,
        SerializedArguments::new(),
    )
    .await;
    orchestrator
        .set_invocation_status(&other, InvocationStatus::Pending, Some(&recoverer))
        .await
        .unwrap();
    assert!(
        orchestrator
            .try_acquire_concurrency_slot(
                &other,
                &task_id,
                &config,
                Some(&SerializedArguments::new())
            )
            .await
            .unwrap(),
        "the taken-over attempt released its concurrency slot"
    );
}