Genarrative/rust/crates/agent-runtime-sqlite/tests/control.rs

use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};

use agent_runtime_core::{ApprovalDecision, Message, RunStatus};
use agent_runtime_sqlite::{RuntimeService, RuntimeServiceError, WorkerLease};

fn user(text: &str) -> Message {
    Message::user(text).expect("valid user message")
}

fn wait_until_epoch_ms(target: i64) {
    let deadline = Instant::now() + Duration::from_secs(2);
    loop {
        let now = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .expect("system clock")
            .as_millis() as i64;
        if now >= target {
            return;
        }
        assert!(Instant::now() < deadline, "lease did not expire");
        std::thread::sleep(Duration::from_millis(1));
    }
}

#[test]
fn queued_cancel_is_atomic_and_runtime_visible() {
    let service = RuntimeService::in_memory().expect("runtime");
    let handle = service.prepare_run("queued cancel").expect("prepare");

    let cancelled = service.cancel_run(&handle.run_id).expect("cancel");
    assert_eq!(cancelled.status, "cancelled");
    assert_eq!(
        service
            .get_run(&handle.run_id)
            .expect("run")
            .unwrap()
            .status,
        "cancelled"
    );
    let snapshot = service
        .load_runtime_snapshot(&handle.runtime_id)
        .expect("snapshot")
        .expect("runtime exists");
    assert_eq!(
        snapshot.run(&handle.run_id).expect("run snapshot").status(),
        RunStatus::Cancelled
    );
}

#[test]
fn active_and_expired_cancel_stay_on_the_reconciliation_gate() {
    let service = RuntimeService::in_memory().expect("runtime");
    let active = service.prepare_run("active cancel").expect("prepare");
    let active_lease = WorkerLease::new(&active.run_id);
    service
        .claim_run_with_lease(&active.run_id, &active_lease, Duration::from_secs(30))
        .expect("claim active");
    let requested = service.cancel_run(&active.run_id).expect("request cancel");
    assert_eq!(requested.status, "cancel_requested");
    assert_eq!(
        service
            .get_run(&active.run_id)
            .expect("run")
            .unwrap()
            .status,
        "cancel_requested"
    );

    // A short lease represents a worker that disappeared before its first
    // heartbeat; cancellation must reconcile it instead of guessing safe.
    let stale = service.prepare_run("stale cancel").expect("prepare");
    let stale_lease = WorkerLease::new(&stale.run_id);
    let (_, stale_record) = service
        .claim_run_with_lease(&stale.run_id, &stale_lease, Duration::from_millis(1))
        .expect("claim stale");
    wait_until_epoch_ms(stale_record.lease_expires_at);
    let reconciled = service.cancel_run(&stale.run_id).expect("reconcile cancel");
    assert_eq!(reconciled.status, "reconciling");
    let snapshot = service
        .load_runtime_snapshot(&stale.runtime_id)
        .expect("snapshot")
        .expect("runtime exists");
    assert_eq!(
        snapshot.run(&stale.run_id).expect("run snapshot").status(),
        RunStatus::Reconciling
    );
}

#[test]
fn fail_unclaimed_run_closes_run_runtime_and_session() {
    let service = RuntimeService::in_memory().expect("runtime");
    let handle = service.prepare_run("setup failure").expect("prepare");

    let failed = service
        .fail_unclaimed_run(&handle.run_id, "provider 配置失败")
        .expect("fail run");
    assert_eq!(failed.status, "failed");
    assert_eq!(
        service
            .get_session(&handle.session_id)
            .expect("session")
            .unwrap()
            .status,
        "failed"
    );
    let snapshot = service
        .load_runtime_snapshot(&handle.runtime_id)
        .expect("snapshot")
        .expect("runtime exists");
    assert_eq!(
        snapshot.run(&handle.run_id).expect("run snapshot").status(),
        RunStatus::Failed
    );
}

#[test]
fn invalid_approval_and_reconciliation_leave_runtime_unchanged() {
    let service = RuntimeService::in_memory().expect("runtime");
    let handle = service
        .prepare_run_with_messages("invalid control", vec![user("invalid control")])
        .expect("prepare");
    let before = service
        .load_runtime_snapshot(&handle.runtime_id)
        .expect("snapshot")
        .expect("runtime exists");

    let approval_error = service
        .resolve_approval_decision("missing-approval", ApprovalDecision::Allow)
        .expect_err("missing approval must fail");
    assert!(!matches!(approval_error, RuntimeServiceError::Core(_)));
    let reconcile_error = service
        .reconcile_provider_result(&handle.run_id, "missing-provider-request", vec![user("x")])
        .expect_err("missing checkpoint must fail");
    assert!(matches!(
        reconcile_error,
        RuntimeServiceError::InvalidInput(_)
    ));

    let after = service
        .load_runtime_snapshot(&handle.runtime_id)
        .expect("snapshot")
        .expect("runtime exists");
    assert_eq!(after.revision(), before.revision());
    assert_eq!(
        after.run(&handle.run_id).expect("run snapshot").status(),
        before.run(&handle.run_id).expect("run snapshot").status()
    );
    assert!(
        service
            .get_approval("missing-approval")
            .expect("approval lookup")
            .is_none()
    );
}