klieo-core 2.2.0

//! Persist + resume behaviour for suspended runs (ADR-045). The checkpoint
//! data type lives in [`crate::checkpoint`]; this module holds the runtime
//! logic that writes it to `KvStore` and replays it back into a run.

use chrono::{DateTime, Utc};
use std::sync::Arc;
use std::time::Duration;

use crate::agent::AgentContext;
use crate::bus::KvStore;
use crate::checkpoint::{ApprovalDecision, RunCheckpoint, CHECKPOINT_BUCKET};
use crate::error::Error;
use crate::ids::ThreadId;
use crate::llm::{FinishReason, Message, Role};

/// Single builder for the blocking and streaming suspend paths (ADR-045), so
/// both persist an identical checkpoint shape.
pub(crate) async fn build_suspend_checkpoint(
    ctx: &AgentContext,
    thread: &ThreadId,
    step: u32,
    message: &Message,
    finish_reason: FinishReason,
    max_history_tokens: usize,
) -> Result<RunCheckpoint, Error> {
    let pending_tool_calls =
        matches!(finish_reason, FinishReason::ToolCalls).then(|| message.tool_calls.clone());
    Ok(RunCheckpoint {
        run_id: ctx.run_id,
        step_index: step,
        thread_id: thread.clone(),
        messages: ctx
            .short_term
            .load(thread.clone(), max_history_tokens)
            .await?,
        pending_tool_calls,
        resume_attempted: false,
        created_at: Utc::now(),
    })
}

/// Persisting under a `KvStore` bucket (rather than holding state in memory) is
/// what lets a different process resume the run. Serialization or store failure
/// surfaces as `Error` — the run never proceeds past a checkpoint it could not
/// persist, so an unpersisted suspend cannot be silently lost.
pub(crate) async fn persist_checkpoint(
    ctx: &AgentContext,
    bucket: &str,
    checkpoint: &RunCheckpoint,
) -> Result<(), Error> {
    let bytes = serde_json::to_vec(checkpoint).map_err(|e| Error::Other {
        message: "checkpoint serialize".into(),
        source: Some(Box::new(e)),
    })?;
    ctx.kv
        .put(bucket, &checkpoint.run_id.to_string(), bytes.into())
        .await?;
    Ok(())
}

/// Reap abandoned suspended-run checkpoints, returning the count deleted. An
/// entry whose `created_at` is at or before `cutoff` is removed; the bound is
/// inclusive. Best-effort throughout: a checkpoint that fails to read or
/// deserialize, or whose delete fails, is logged and skipped so one bad entry
/// never stalls the sweep. The `cutoff` is typically `Utc::now() - ttl`.
///
/// Enumerates with `keys_paginated` one bounded page at a time and fetches one
/// checkpoint per key, so a large backlog of abandoned runs pulls neither the
/// whole key list nor every (potentially conversation-sized) payload into
/// memory at once — only one page of keys plus a single checkpoint are resident.
///
/// A host can drive this on its own schedule, or spawn the opt-in
/// [`spawn_checkpoint_gc`] task. Distinct from [`crate::spawn_kv_reaper`], which
/// evicts `{kind}.{stream_id}` keys by resume-buffer liveness: checkpoints key
/// on `run_id` and have no resume buffer, so age is the only signal an abandoned
/// one leaves behind.
pub async fn gc_checkpoints(kv: &dyn KvStore, cutoff: DateTime<Utc>) -> Result<u64, Error> {
    gc_checkpoints_paged(kv, cutoff, GC_KEY_PAGE).await
}

/// Handle for the spawned checkpoint-GC task; its `Drop` aborts the task.
pub struct CheckpointGcHandle {
    task: Option<tokio::task::JoinHandle<()>>,
}

impl Drop for CheckpointGcHandle {
    fn drop(&mut self) {
        if let Some(task) = self.task.take() {
            task.abort();
        }
    }
}

/// Spawn the opt-in background task that reaps abandoned suspended-run
/// checkpoints every `interval`, deleting any older than `ttl`. Key the `ttl` to
/// the review-item timeout SLA so a checkpoint never outlives the review it backs
/// (ADR-045 item 7) — without it, abandoned suspensions leak `KvStore` entries
/// forever. Memory-bounded via [`gc_checkpoints`] (page-walk, one checkpoint
/// resident at a time). A `ttl` too large to represent skips that sweep rather
/// than reaping everything. Returns a [`CheckpointGcHandle`] whose `Drop` aborts
/// the task.
pub fn spawn_checkpoint_gc(
    kv: Arc<dyn KvStore>,
    ttl: Duration,
    interval: Duration,
) -> CheckpointGcHandle {
    let task = tokio::spawn(async move {
        loop {
            tokio::time::sleep(interval).await;
            let Ok(ttl_chrono) = chrono::Duration::from_std(ttl) else {
                tracing::warn!(
                    target: "klieo.checkpoint.gc",
                    ttl = ?ttl,
                    "configured ttl is out of range; skipping this sweep"
                );
                continue;
            };
            match gc_checkpoints(kv.as_ref(), Utc::now() - ttl_chrono).await {
                Ok(reaped) if reaped > 0 => tracing::info!(
                    target: "klieo.checkpoint.gc",
                    reaped,
                    "reaped abandoned suspended-run checkpoints"
                ),
                Ok(_) => {}
                Err(err) => tracing::warn!(
                    target: "klieo.checkpoint.gc",
                    error = %err,
                    "checkpoint gc sweep failed; will retry next interval"
                ),
            }
        }
    });
    CheckpointGcHandle { task: Some(task) }
}

/// Sweep page size: large enough to amortize the per-page round trip, small
/// enough that one page of keys stays a bounded slice rather than the whole
/// bucket. The exact value is not load-bearing — any backend that cares
/// overrides `keys_paginated`.
const GC_KEY_PAGE: usize = 256;

/// Page-walking core of [`gc_checkpoints`], parameterized on page size so tests
/// can drive the multi-page cursor loop without seeding a full page of keys.
async fn gc_checkpoints_paged(
    kv: &dyn KvStore,
    cutoff: DateTime<Utc>,
    page_size: usize,
) -> Result<u64, Error> {
    let mut reaped = 0u64;
    let mut cursor = None;
    loop {
        let page = kv
            .keys_paginated(CHECKPOINT_BUCKET, cursor, page_size)
            .await?;
        for key in &page.keys {
            let Some(checkpoint) = load_checkpoint_for_gc(kv, key).await else {
                continue;
            };
            if checkpoint.created_at > cutoff {
                continue;
            }
            match kv.delete(CHECKPOINT_BUCKET, key).await {
                Ok(()) => reaped += 1,
                Err(e) => tracing::warn!(
                    target: "klieo.checkpoint.gc",
                    operation = "delete",
                    key = %key,
                    error = %e,
                    "checkpoint delete failed; continuing sweep"
                ),
            }
        }
        match page.next {
            Some(c) => cursor = Some(c),
            None => break,
        }
    }
    Ok(reaped)
}

/// Read and deserialize one checkpoint for the GC sweep. Returns `None` — after
/// logging the cause — when the key has already been deleted (a concurrent
/// reaper raced us), the read fails, or the stored value is not a checkpoint, so
/// the caller can skip it without aborting the whole sweep.
async fn load_checkpoint_for_gc(kv: &dyn KvStore, key: &str) -> Option<RunCheckpoint> {
    let entry = match kv.get(CHECKPOINT_BUCKET, key).await {
        Ok(Some(entry)) => entry,
        Ok(None) => return None,
        Err(e) => {
            tracing::warn!(
                target: "klieo.checkpoint.gc",
                operation = "read",
                key = %key,
                error = %e,
                "checkpoint read failed; skipping"
            );
            return None;
        }
    };
    match serde_json::from_slice(&entry.value) {
        Ok(checkpoint) => Some(checkpoint),
        Err(e) => {
            tracing::warn!(
                target: "klieo.checkpoint.gc",
                operation = "deserialize",
                key = %key,
                error = %e,
                "skipping undeserializable checkpoint entry"
            );
            None
        }
    }
}

/// Resume a run suspended via the review gate (ADR-045).
///
/// On `Approved` with pending tool calls, those calls are dispatched here. Each
/// resume first claims a one-shot latch on the persisted checkpoint by
/// compare-and-set ([`claim_resume_latch`]): a sequential retry (the store is
/// already latched) and a concurrent resume (it loses the CAS) are both refused
/// for a non-idempotent pending call — returning [`Error::ResumeReplayBlocked`]
/// and leaving the checkpoint for operator reconciliation — rather than risk a
/// duplicate side effect (e.g. a double payout). Idempotent tools
/// (`ToolInvoker::is_tool_idempotent`) re-dispatch freely. A within-process
/// resume with no persisted checkpoint cannot be retried or raced and is
/// pre-dispatch-safe, so it skips the latch.
pub async fn resume_from_checkpoint(
    ctx: &AgentContext,
    system_prompt: &str,
    checkpoint: RunCheckpoint,
    decision: ApprovalDecision,
    opts: super::RunOptions,
) -> Result<String, Error> {
    let thread = checkpoint.thread_id.clone();
    let latch = RunCheckpoint {
        resume_attempted: true,
        ..checkpoint.clone()
    };
    ctx.short_term.clear(thread.clone()).await?;
    ctx.short_term
        .append_batch(thread.clone(), checkpoint.messages)
        .await?;
    match (decision, checkpoint.pending_tool_calls) {
        (ApprovalDecision::Approved, Some(calls)) => {
            dispatch_pending_on_resume(ctx, &thread, &calls, &latch, &opts).await?;
        }
        (ApprovalDecision::Approved, None) => {}
        (ApprovalDecision::Rejected { reason }, _) => {
            ctx.short_term
                .append(
                    thread.clone(),
                    Message {
                        role: Role::Tool,
                        content: format!("Human reviewer rejected this step: {reason}"),
                        tool_calls: vec![],
                        tool_call_id: None,
                    },
                )
                .await?;
        }
    }
    super::run_loop(ctx, system_prompt, &thread, &opts, checkpoint.step_index).await
}

/// Claim the resume latch (fail-closed, ADR-045) then dispatch the approved
/// pending calls. The latch only applies to a persisted checkpoint — the only
/// thing a retry or a concurrent resume can race on; a bucketless within-process
/// resume is pre-dispatch-safe and dispatches directly.
async fn dispatch_pending_on_resume(
    ctx: &AgentContext,
    thread: &ThreadId,
    calls: &[crate::llm::ToolCall],
    latch: &RunCheckpoint,
    opts: &super::RunOptions,
) -> Result<(), Error> {
    let non_idempotent: Vec<String> = calls
        .iter()
        .filter(|c| !ctx.tools.is_tool_idempotent(&c.name))
        .map(|c| c.name.clone())
        .collect();
    if let Some(bucket) = &opts.checkpoint_kv_bucket {
        claim_resume_latch(ctx, bucket, latch, &non_idempotent).await?;
    }
    super::dispatch_tool_calls(ctx, thread, calls, "resume").await
}

/// Atomically claim the one-shot resume latch by compare-and-set, so neither a
/// sequential retry nor a concurrent resume can re-fire a non-idempotent tool
/// (ADR-045). A pending non-idempotent call is refused when the persisted
/// checkpoint is already latched or the CAS loses to a concurrent claimer;
/// idempotent batches proceed in either case.
async fn claim_resume_latch(
    ctx: &AgentContext,
    bucket: &str,
    latch: &RunCheckpoint,
    non_idempotent: &[String],
) -> Result<(), Error> {
    let key = latch.run_id.to_string();
    let expected = match ctx.kv.get(bucket, &key).await? {
        Some(entry) => {
            if checkpoint_is_latched(&entry.value) && !non_idempotent.is_empty() {
                return Err(resume_blocked(latch.run_id, non_idempotent));
            }
            Some(entry.revision)
        }
        None => None,
    };
    let bytes = serde_json::to_vec(latch).map_err(|e| Error::wrap("checkpoint serialize", e))?;
    match ctx.kv.cas(bucket, &key, bytes.into(), expected).await {
        Ok(_) => Ok(()),
        Err(crate::error::BusError::CasConflict { .. }) if !non_idempotent.is_empty() => {
            Err(resume_blocked(latch.run_id, non_idempotent))
        }
        Err(crate::error::BusError::CasConflict { .. }) => Ok(()),
        Err(e) => Err(Error::Bus(e)),
    }
}

/// A persisted checkpoint counts as latched when it parses and `resume_attempted`
/// is set. An unparseable value is treated as latched — the fail-closed
/// direction: refuse rather than risk re-firing on a checkpoint we cannot read.
fn checkpoint_is_latched(value: &[u8]) -> bool {
    serde_json::from_slice::<RunCheckpoint>(value)
        .map(|c| c.resume_attempted)
        .unwrap_or(true)
}

fn resume_blocked(run_id: crate::ids::RunId, tools: &[String]) -> Error {
    tracing::error!(
        target: "klieo.checkpoint.resume",
        run_id = %run_id,
        tools = ?tools,
        "resume blocked: non-idempotent pending tool calls cannot be proven un-fired (fail-closed, ADR-045) — operator reconciliation required"
    );
    Error::ResumeReplayBlocked {
        run_id,
        tools: tools.to_vec(),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::bus::{KvEntry, Lease, Revision};
    use crate::error::BusError;
    use crate::ids::{RunId, ThreadId};
    use crate::llm::ToolCall;
    use crate::runtime::RunOptions;
    use crate::test_utils::{fake_context, fake_kv, FakeLlmClient, FakeLlmStep, FakeToolInvoker};
    use async_trait::async_trait;
    use bytes::Bytes;
    use chrono::Utc;
    use std::time::Duration;

    struct KeysFailKv;

    #[async_trait]
    impl crate::bus::KvStore for KeysFailKv {
        async fn get(&self, _: &str, _: &str) -> Result<Option<KvEntry>, BusError> {
            Err(BusError::Unsupported("get".into()))
        }
        async fn put(&self, _: &str, _: &str, _: Bytes) -> Result<Revision, BusError> {
            Err(BusError::Unsupported("put".into()))
        }
        async fn cas(
            &self,
            _: &str,
            _: &str,
            _: Bytes,
            _: Option<Revision>,
        ) -> Result<Revision, BusError> {
            Err(BusError::Unsupported("cas".into()))
        }
        async fn delete(&self, _: &str, _: &str) -> Result<(), BusError> {
            Err(BusError::Unsupported("delete".into()))
        }
        async fn lease(&self, _: &str, _: &str, _: Duration) -> Result<Lease, BusError> {
            Err(BusError::Unsupported("lease".into()))
        }
        async fn keys(&self, _: &str) -> Result<Vec<String>, BusError> {
            Err(BusError::Unsupported("enumerate unavailable".into()))
        }
    }

    struct DeleteFailKv {
        value: Bytes,
    }

    #[async_trait]
    impl crate::bus::KvStore for DeleteFailKv {
        async fn get(&self, _: &str, _: &str) -> Result<Option<KvEntry>, BusError> {
            Ok(Some(KvEntry {
                value: self.value.clone(),
                revision: 1,
            }))
        }
        async fn put(&self, _: &str, _: &str, _: Bytes) -> Result<Revision, BusError> {
            Err(BusError::Unsupported("put".into()))
        }
        async fn cas(
            &self,
            _: &str,
            _: &str,
            _: Bytes,
            _: Option<Revision>,
        ) -> Result<Revision, BusError> {
            Err(BusError::Unsupported("cas".into()))
        }
        async fn delete(&self, _: &str, _: &str) -> Result<(), BusError> {
            Err(BusError::Unsupported("delete boom".into()))
        }
        async fn lease(&self, _: &str, _: &str, _: Duration) -> Result<Lease, BusError> {
            Err(BusError::Unsupported("lease".into()))
        }
        async fn keys(&self, _: &str) -> Result<Vec<String>, BusError> {
            Ok(vec!["stale".to_string()])
        }
    }

    /// `get` returns a stored (un-latched) value, but every `cas` loses with a
    /// `CasConflict` — models a concurrent resume that claimed the latch between
    /// our read and our compare-and-set.
    struct ConflictKv {
        value: Bytes,
    }

    #[async_trait]
    impl crate::bus::KvStore for ConflictKv {
        async fn get(&self, _: &str, _: &str) -> Result<Option<KvEntry>, BusError> {
            Ok(Some(KvEntry {
                value: self.value.clone(),
                revision: 1,
            }))
        }
        async fn put(&self, _: &str, _: &str, _: Bytes) -> Result<Revision, BusError> {
            Err(BusError::Unsupported("put".into()))
        }
        async fn cas(
            &self,
            _: &str,
            _: &str,
            _: Bytes,
            _: Option<Revision>,
        ) -> Result<Revision, BusError> {
            Err(BusError::CasConflict {
                expected: 1,
                actual: 2,
            })
        }
        async fn delete(&self, _: &str, _: &str) -> Result<(), BusError> {
            Err(BusError::Unsupported("delete".into()))
        }
        async fn lease(&self, _: &str, _: &str, _: Duration) -> Result<Lease, BusError> {
            Err(BusError::Unsupported("lease".into()))
        }
        async fn keys(&self, _: &str) -> Result<Vec<String>, BusError> {
            Err(BusError::Unsupported("keys".into()))
        }
    }
    use std::sync::Arc;

    #[test]
    fn checkpoint_round_trips_through_json() {
        let cp = checkpoint_with_pending_tool(ThreadId::new("t-rt"), RunId::new());
        let bytes = serde_json::to_vec(&cp).unwrap();
        let back: RunCheckpoint = serde_json::from_slice(&bytes).unwrap();
        assert_eq!(back.step_index, cp.step_index);
        assert_eq!(back.run_id, cp.run_id);
        assert_eq!(back.thread_id, cp.thread_id);

        assert_eq!(back.messages.len(), cp.messages.len(), "history length");
        assert_eq!(back.messages[0].role, cp.messages[0].role);
        assert_eq!(back.messages[0].content, cp.messages[0].content);

        let restored = back
            .pending_tool_calls
            .expect("pending tool calls must survive the round-trip");
        let original = cp.pending_tool_calls.unwrap();
        assert_eq!(restored.len(), original.len());
        assert_eq!(restored[0].id, original[0].id);
        assert_eq!(restored[0].name, original[0].name);
        assert_eq!(restored[0].args, original[0].args);
    }

    fn echo_tool_invoker() -> Arc<FakeToolInvoker> {
        Arc::new(FakeToolInvoker::new().with_tool("echo", "echo back", Ok))
    }

    fn checkpoint_with_pending_tool(thread: ThreadId, run_id: RunId) -> RunCheckpoint {
        RunCheckpoint {
            run_id,
            step_index: 1,
            thread_id: thread,
            messages: vec![Message {
                role: Role::User,
                content: "go".into(),
                tool_calls: vec![],
                tool_call_id: None,
            }],
            pending_tool_calls: Some(vec![ToolCall {
                id: "tc-1".into(),
                name: "echo".into(),
                args: serde_json::json!({"x": 1}),
            }]),
            resume_attempted: false,
            created_at: Utc::now(),
        }
    }

    fn checkpoint_no_pending(thread: ThreadId, run_id: RunId) -> RunCheckpoint {
        RunCheckpoint {
            run_id,
            step_index: 1,
            thread_id: thread,
            messages: vec![Message {
                role: Role::User,
                content: "go".into(),
                tool_calls: vec![],
                tool_call_id: None,
            }],
            pending_tool_calls: None,
            resume_attempted: false,
            created_at: Utc::now(),
        }
    }

    #[tokio::test]
    async fn resume_approved_dispatches_pending_then_completes() {
        let mut ctx = fake_context("resume-test");
        ctx.llm =
            Arc::new(FakeLlmClient::new("fake").with_steps(vec![FakeLlmStep::Text("done".into())]));
        ctx.kv = fake_kv();
        ctx.tools = echo_tool_invoker();

        let thread = ThreadId::new("t-resume-approved");
        let cp = checkpoint_with_pending_tool(thread.clone(), ctx.run_id);

        let out = resume_from_checkpoint(
            &ctx,
            "sys",
            cp,
            ApprovalDecision::Approved,
            RunOptions::default(),
        )
        .await
        .unwrap();

        assert_eq!(out, "done");

        let history = ctx.short_term.load(thread, 8192).await.unwrap();
        let tool_msgs: Vec<_> = history.iter().filter(|m| m.role == Role::Tool).collect();
        assert_eq!(
            tool_msgs.len(),
            1,
            "dispatched tool must leave a Role::Tool message"
        );
        assert_eq!(tool_msgs[0].tool_call_id.as_deref(), Some("tc-1"));
    }

    #[tokio::test]
    async fn resume_approved_without_pending_calls_completes_and_injects_no_tool_message() {
        let mut ctx = fake_context("resume-approved-none");
        ctx.llm =
            Arc::new(FakeLlmClient::new("fake").with_steps(vec![FakeLlmStep::Text("done".into())]));
        ctx.kv = fake_kv();

        let thread = ThreadId::new("t-resume-approved-none");
        let cp = checkpoint_no_pending(thread.clone(), ctx.run_id);

        let out = resume_from_checkpoint(
            &ctx,
            "sys",
            cp,
            ApprovalDecision::Approved,
            RunOptions::default(),
        )
        .await
        .unwrap();

        assert_eq!(out, "done");

        let history = ctx.short_term.load(thread, 8192).await.unwrap();
        let tool_msgs = history.iter().filter(|m| m.role == Role::Tool).count();
        assert_eq!(
            tool_msgs, 0,
            "approving a step with no pending tool calls must inject nothing"
        );
    }

    #[tokio::test]
    async fn resume_rejected_injects_tool_message_then_completes() {
        let mut ctx = fake_context("resume-reject");
        ctx.llm =
            Arc::new(FakeLlmClient::new("fake").with_steps(vec![FakeLlmStep::Text("done".into())]));
        ctx.kv = fake_kv();

        let thread = ThreadId::new("t-resume-rejected");
        let cp = checkpoint_no_pending(thread.clone(), ctx.run_id);

        let out = resume_from_checkpoint(
            &ctx,
            "sys",
            cp,
            ApprovalDecision::Rejected {
                reason: "no".into(),
            },
            RunOptions::default(),
        )
        .await
        .unwrap();

        assert_eq!(out, "done");

        let history = ctx.short_term.load(thread, 8192).await.unwrap();
        let rejection: Vec<_> = history
            .iter()
            .filter(|m| m.role == Role::Tool && m.content.contains("no"))
            .collect();
        assert_eq!(rejection.len(), 1, "rejected message must be appended");
        assert!(rejection[0]
            .content
            .contains("Human reviewer rejected this step: no"));
    }

    #[tokio::test]
    async fn resume_rejected_with_pending_calls_does_not_dispatch_them() {
        let mut ctx = fake_context("resume-reject-pending");
        ctx.llm =
            Arc::new(FakeLlmClient::new("fake").with_steps(vec![FakeLlmStep::Text("done".into())]));
        ctx.kv = fake_kv();
        ctx.tools = echo_tool_invoker();

        let thread = ThreadId::new("t-resume-reject-pending");
        let cp = checkpoint_with_pending_tool(thread.clone(), ctx.run_id);

        let out = resume_from_checkpoint(
            &ctx,
            "sys",
            cp,
            ApprovalDecision::Rejected {
                reason: "blocked".into(),
            },
            RunOptions::default(),
        )
        .await
        .unwrap();

        assert_eq!(out, "done");

        let history = ctx.short_term.load(thread, 8192).await.unwrap();
        assert!(
            !history
                .iter()
                .any(|m| m.tool_call_id.as_deref() == Some("tc-1")),
            "a rejected step must NOT dispatch its pending tool calls"
        );
        assert!(
            history
                .iter()
                .any(|m| m.role == Role::Tool && m.content.contains("blocked")),
            "the rejection reason must be fed back to the model"
        );
    }

    async fn seed_latched(ctx: &AgentContext, thread: &ThreadId) {
        let mut latched = checkpoint_with_pending_tool(thread.clone(), ctx.run_id);
        latched.resume_attempted = true;
        ctx.kv
            .put(
                CHECKPOINT_BUCKET,
                &ctx.run_id.to_string(),
                serde_json::to_vec(&latched).unwrap().into(),
            )
            .await
            .unwrap();
    }

    #[tokio::test]
    async fn resume_retry_with_non_idempotent_pending_fails_closed() {
        let mut ctx = fake_context("resume-retry-nonidem");
        ctx.llm =
            Arc::new(FakeLlmClient::new("fake").with_steps(vec![FakeLlmStep::Text("done".into())]));
        ctx.kv = fake_kv();
        ctx.tools = echo_tool_invoker(); // `echo` uses the default: non-idempotent

        let thread = ThreadId::new("t-retry-nonidem");
        seed_latched(&ctx, &thread).await; // a prior resume already latched the store
        let cp = checkpoint_with_pending_tool(thread.clone(), ctx.run_id);

        let err = resume_from_checkpoint(
            &ctx,
            "sys",
            cp,
            ApprovalDecision::Approved,
            RunOptions::default().with_checkpoint_bucket(CHECKPOINT_BUCKET),
        )
        .await
        .unwrap_err();

        match err {
            Error::ResumeReplayBlocked { tools, .. } => assert_eq!(
                tools,
                vec!["echo".to_string()],
                "the refused non-idempotent tool must be named"
            ),
            other => panic!("expected ResumeReplayBlocked, got {other:?}"),
        }

        let history = ctx.short_term.load(thread, 8192).await.unwrap();
        assert!(
            !history
                .iter()
                .any(|m| m.tool_call_id.as_deref() == Some("tc-1")),
            "a fail-closed resume must NOT re-dispatch the non-idempotent pending call"
        );
    }

    #[tokio::test]
    async fn resume_retry_with_idempotent_pending_redispatches() {
        let mut ctx = fake_context("resume-retry-idem");
        ctx.llm =
            Arc::new(FakeLlmClient::new("fake").with_steps(vec![FakeLlmStep::Text("done".into())]));
        ctx.kv = fake_kv();
        ctx.tools = Arc::new(FakeToolInvoker::new().with_idempotent_tool("echo", "echo back", Ok));

        let thread = ThreadId::new("t-retry-idem");
        seed_latched(&ctx, &thread).await;
        let cp = checkpoint_with_pending_tool(thread.clone(), ctx.run_id);

        let out = resume_from_checkpoint(
            &ctx,
            "sys",
            cp,
            ApprovalDecision::Approved,
            RunOptions::default().with_checkpoint_bucket(CHECKPOINT_BUCKET),
        )
        .await
        .unwrap();
        assert_eq!(out, "done");

        let history = ctx.short_term.load(thread, 8192).await.unwrap();
        assert!(
            history
                .iter()
                .any(|m| m.tool_call_id.as_deref() == Some("tc-1")),
            "an idempotent tool is safe to re-dispatch even after a prior resume latch"
        );
    }

    #[tokio::test]
    async fn concurrent_resume_losing_the_cas_fails_closed() {
        // The store still reads as un-latched, but the compare-and-set loses —
        // a concurrent resume claimed the latch between our read and write. A
        // non-idempotent call must be refused, not re-fired (CWE-367 TOCTOU).
        let mut ctx = fake_context("resume-cas-conflict");
        ctx.llm =
            Arc::new(FakeLlmClient::new("fake").with_steps(vec![FakeLlmStep::Text("done".into())]));
        ctx.tools = echo_tool_invoker();
        let thread = ThreadId::new("t-cas-conflict");
        let unlatched =
            serde_json::to_vec(&checkpoint_with_pending_tool(thread.clone(), ctx.run_id))
                .unwrap()
                .into();
        ctx.kv = Arc::new(ConflictKv { value: unlatched });

        let err = resume_from_checkpoint(
            &ctx,
            "sys",
            checkpoint_with_pending_tool(thread, ctx.run_id),
            ApprovalDecision::Approved,
            RunOptions::default().with_checkpoint_bucket(CHECKPOINT_BUCKET),
        )
        .await
        .unwrap_err();

        assert!(
            matches!(err, Error::ResumeReplayBlocked { .. }),
            "losing the latch CAS to a concurrent resume must fail closed, got {err:?}"
        );
    }

    #[tokio::test]
    async fn first_resume_latches_attempt_before_dispatch() {
        let mut ctx = fake_context("resume-latch");
        ctx.llm =
            Arc::new(FakeLlmClient::new("fake").with_steps(vec![FakeLlmStep::Text("done".into())]));
        ctx.kv = fake_kv();
        ctx.tools = echo_tool_invoker();

        let thread = ThreadId::new("t-latch");
        let run_id = ctx.run_id;
        let cp = checkpoint_with_pending_tool(thread, run_id);

        resume_from_checkpoint(
            &ctx,
            "sys",
            cp,
            ApprovalDecision::Approved,
            RunOptions::default().with_checkpoint_bucket(CHECKPOINT_BUCKET),
        )
        .await
        .unwrap();

        let entry = ctx
            .kv
            .get(CHECKPOINT_BUCKET, &run_id.to_string())
            .await
            .unwrap()
            .expect("first resume must persist the latched checkpoint");
        let latched: RunCheckpoint = serde_json::from_slice(&entry.value).unwrap();
        assert!(
            latched.resume_attempted,
            "first resume must latch resume_attempted=true before dispatch, so a retry is recognised"
        );
    }

    #[test]
    fn corrupt_checkpoint_bytes_treated_as_latched() {
        assert!(
            checkpoint_is_latched(b"not-json"),
            "an unparseable stored checkpoint must read as latched — refuse rather than risk a re-fire"
        );
        let unlatched = serde_json::to_vec(&checkpoint_no_pending(
            ThreadId::new("t-parse"),
            RunId::new(),
        ))
        .unwrap();
        assert!(
            !checkpoint_is_latched(&unlatched),
            "a well-formed un-latched checkpoint must read as not latched"
        );
    }

    #[tokio::test]
    async fn bucketless_resume_skips_latch_and_dispatches() {
        let mut ctx = fake_context("resume-no-bucket");
        ctx.llm =
            Arc::new(FakeLlmClient::new("fake").with_steps(vec![FakeLlmStep::Text("done".into())]));
        ctx.kv = fake_kv();
        ctx.tools = echo_tool_invoker(); // non-idempotent

        let thread = ThreadId::new("t-no-bucket");
        let run_id = ctx.run_id;
        let cp = checkpoint_with_pending_tool(thread.clone(), run_id);

        let out = resume_from_checkpoint(
            &ctx,
            "sys",
            cp,
            ApprovalDecision::Approved,
            RunOptions::default(),
        )
        .await
        .unwrap();
        assert_eq!(out, "done");

        let history = ctx.short_term.load(thread, 8192).await.unwrap();
        assert!(
            history
                .iter()
                .any(|m| m.tool_call_id.as_deref() == Some("tc-1")),
            "a within-process resume with no bucket is pre-dispatch-safe and dispatches directly"
        );
        assert!(
            ctx.kv
                .get(CHECKPOINT_BUCKET, &run_id.to_string())
                .await
                .unwrap()
                .is_none(),
            "no bucket means no latch is persisted"
        );
    }

    #[tokio::test]
    async fn concurrent_cas_conflict_with_idempotent_batch_proceeds() {
        let mut ctx = fake_context("resume-conflict-idem");
        ctx.llm =
            Arc::new(FakeLlmClient::new("fake").with_steps(vec![FakeLlmStep::Text("done".into())]));
        ctx.tools = Arc::new(FakeToolInvoker::new().with_idempotent_tool("echo", "echo back", Ok));
        let thread = ThreadId::new("t-conflict-idem");
        let unlatched =
            serde_json::to_vec(&checkpoint_with_pending_tool(thread.clone(), ctx.run_id))
                .unwrap()
                .into();
        ctx.kv = Arc::new(ConflictKv { value: unlatched });

        let out = resume_from_checkpoint(
            &ctx,
            "sys",
            checkpoint_with_pending_tool(thread.clone(), ctx.run_id),
            ApprovalDecision::Approved,
            RunOptions::default().with_checkpoint_bucket(CHECKPOINT_BUCKET),
        )
        .await
        .unwrap();
        assert_eq!(out, "done");

        let history = ctx.short_term.load(thread, 8192).await.unwrap();
        assert!(
            history
                .iter()
                .any(|m| m.tool_call_id.as_deref() == Some("tc-1")),
            "an idempotent tool proceeds even when it loses the latch CAS"
        );
    }

    #[test]
    fn legacy_checkpoint_without_resume_attempted_defaults_false() {
        // A checkpoint persisted before the latch field existed must stay
        // readable and default to the first-attempt state — never silently
        // treat an old checkpoint as already-attempted (which would block a
        // legitimate first resume).
        let run_id = RunId::new();
        let legacy = format!(
            r#"{{"run_id":"{run_id}","step_index":1,"thread_id":"t-legacy","messages":[],"pending_tool_calls":null,"created_at":"2026-06-12T00:00:00Z"}}"#
        );
        let back: RunCheckpoint = serde_json::from_str(&legacy).unwrap();
        assert!(!back.resume_attempted);
    }

    async fn seed_checkpoint(
        kv: &Arc<dyn crate::bus::KvStore>,
        run_id: RunId,
        created_at: chrono::DateTime<Utc>,
    ) {
        let key = run_id.to_string();
        let mut cp = checkpoint_with_pending_tool(ThreadId::new("t-gc"), run_id);
        cp.created_at = created_at;
        let bytes = serde_json::to_vec(&cp).unwrap();
        kv.put(CHECKPOINT_BUCKET, &key, bytes.into()).await.unwrap();
    }

    #[tokio::test]
    async fn gc_checkpoints_reaps_only_entries_at_or_before_cutoff() {
        let kv = fake_kv();
        let now = Utc::now();
        let cutoff = now - chrono::Duration::hours(1);
        let stale = RunId::new();
        let at_cutoff = RunId::new();
        let fresh = RunId::new();
        let stale_key = stale.to_string();
        let at_cutoff_key = at_cutoff.to_string();
        let fresh_key = fresh.to_string();
        seed_checkpoint(&kv, stale, now - chrono::Duration::hours(2)).await;
        seed_checkpoint(&kv, at_cutoff, cutoff).await;
        seed_checkpoint(&kv, fresh, now).await;

        let reaped = gc_checkpoints(kv.as_ref(), cutoff).await.unwrap();

        assert_eq!(
            reaped, 2,
            "the stale checkpoint and the one exactly at the cutoff are both reaped (bound is <=)"
        );
        assert!(
            kv.get(CHECKPOINT_BUCKET, &stale_key)
                .await
                .unwrap()
                .is_none(),
            "the stale checkpoint is deleted"
        );
        assert!(
            kv.get(CHECKPOINT_BUCKET, &at_cutoff_key)
                .await
                .unwrap()
                .is_none(),
            "a checkpoint whose created_at equals the cutoff is reaped — the bound is inclusive"
        );
        assert!(
            kv.get(CHECKPOINT_BUCKET, &fresh_key)
                .await
                .unwrap()
                .is_some(),
            "a checkpoint newer than the cutoff must survive the sweep"
        );
    }

    #[tokio::test]
    async fn gc_checkpoints_paged_reaps_eligible_across_pages() {
        let kv = fake_kv();
        let now = Utc::now();
        let cutoff = now;
        let stale = now - chrono::Duration::hours(1);
        let ids: Vec<RunId> = (0..5).map(|_| RunId::new()).collect();
        for id in &ids {
            seed_checkpoint(&kv, *id, stale).await;
        }

        // page_size 2 over 5 keys drives three pages + two cursor advances.
        let reaped = gc_checkpoints_paged(kv.as_ref(), cutoff, 2).await.unwrap();

        assert_eq!(reaped, 5);
        for id in &ids {
            assert!(
                kv.get(CHECKPOINT_BUCKET, &id.to_string())
                    .await
                    .unwrap()
                    .is_none(),
                "every eligible checkpoint across all pages is reaped"
            );
        }
    }

    #[tokio::test]
    async fn gc_checkpoints_skips_undeserializable_entry_without_failing() {
        let kv = fake_kv();
        let now = Utc::now();
        let stale = RunId::new();
        let stale_key = stale.to_string();
        seed_checkpoint(&kv, stale, now - chrono::Duration::hours(2)).await;
        kv.put(
            CHECKPOINT_BUCKET,
            "not-a-checkpoint",
            b"junk".to_vec().into(),
        )
        .await
        .unwrap();

        let reaped = gc_checkpoints(kv.as_ref(), now).await.unwrap();

        assert_eq!(reaped, 1, "the valid stale checkpoint is still reaped");
        assert!(
            kv.get(CHECKPOINT_BUCKET, &stale_key)
                .await
                .unwrap()
                .is_none(),
            "the stale checkpoint is gone"
        );
        assert!(
            kv.get(CHECKPOINT_BUCKET, "not-a-checkpoint")
                .await
                .unwrap()
                .is_some(),
            "an undeserializable entry is skipped, not deleted — a foreign value cannot stall the sweep"
        );
    }

    #[tokio::test]
    async fn gc_checkpoints_empty_bucket_is_zero() {
        let kv = fake_kv();
        let reaped = gc_checkpoints(kv.as_ref(), Utc::now()).await.unwrap();
        assert_eq!(reaped, 0, "an empty bucket reaps nothing");
    }

    #[tokio::test]
    async fn gc_checkpoints_propagates_enumerate_failure() {
        let kv = KeysFailKv;
        let err = gc_checkpoints(&kv, Utc::now()).await.unwrap_err();
        assert!(
            matches!(err, Error::Bus(BusError::Unsupported(_))),
            "a keys() failure must propagate — the sweep cannot enumerate, so it must surface the error, not silently report zero reaped"
        );
    }

    #[tokio::test]
    async fn gc_checkpoints_skips_entry_whose_delete_fails() {
        let mut cp = checkpoint_with_pending_tool(ThreadId::new("t-del"), RunId::new());
        cp.created_at = Utc::now() - chrono::Duration::hours(1);
        let kv = DeleteFailKv {
            value: Bytes::from(serde_json::to_vec(&cp).unwrap()),
        };

        let reaped = gc_checkpoints(&kv, Utc::now()).await.unwrap();

        assert_eq!(
            reaped, 0,
            "a stale checkpoint whose delete fails is logged and skipped, not counted; the sweep still returns Ok"
        );
    }

    #[tokio::test]
    async fn spawn_checkpoint_gc_reaps_stale_and_keeps_fresh() {
        let kv = fake_kv();
        let stale = RunId::new();
        let fresh = RunId::new();
        seed_checkpoint(&kv, stale, Utc::now() - chrono::Duration::hours(2)).await;
        seed_checkpoint(&kv, fresh, Utc::now()).await;

        // ttl 1h: the 2h-old checkpoint is past it, the just-created one is not.
        let handle = spawn_checkpoint_gc(
            Arc::clone(&kv),
            Duration::from_secs(3600),
            Duration::from_millis(5),
        );
        // Let several sweep intervals elapse.
        tokio::time::sleep(Duration::from_millis(80)).await;

        assert!(
            kv.get(CHECKPOINT_BUCKET, &stale.to_string())
                .await
                .unwrap()
                .is_none(),
            "the spawned gc must reap a checkpoint older than the ttl"
        );
        assert!(
            kv.get(CHECKPOINT_BUCKET, &fresh.to_string())
                .await
                .unwrap()
                .is_some(),
            "a checkpoint within the ttl must survive"
        );
        drop(handle);
    }

    #[tokio::test]
    async fn checkpoint_gc_handle_drop_stops_reaping() {
        let kv = fake_kv();
        let stale = RunId::new();
        seed_checkpoint(&kv, stale, Utc::now() - chrono::Duration::hours(2)).await;

        // Drop before the first interval elapses, then wait well past several
        // intervals: the abort must prevent any sweep (contrast with
        // `spawn_checkpoint_gc_reaps_stale_and_keeps_fresh`, identical timing
        // but no drop, where the same stale checkpoint IS reaped).
        let handle = spawn_checkpoint_gc(
            Arc::clone(&kv),
            Duration::from_secs(3600),
            Duration::from_millis(5),
        );
        drop(handle);
        tokio::time::sleep(Duration::from_millis(80)).await;

        assert!(
            kv.get(CHECKPOINT_BUCKET, &stale.to_string())
                .await
                .unwrap()
                .is_some(),
            "dropping the handle aborts the task, so the stale checkpoint is never reaped"
        );
    }

    #[tokio::test]
    async fn spawn_checkpoint_gc_with_out_of_range_ttl_skips_sweep() {
        let kv = fake_kv();
        let stale = RunId::new();
        seed_checkpoint(&kv, stale, Utc::now() - chrono::Duration::hours(2)).await;

        // `Duration::MAX` overflows `chrono::Duration::from_std`, so each sweep
        // skips rather than reaping everything against a degenerate cutoff.
        let handle = spawn_checkpoint_gc(Arc::clone(&kv), Duration::MAX, Duration::from_millis(5));
        tokio::time::sleep(Duration::from_millis(80)).await;

        assert!(
            kv.get(CHECKPOINT_BUCKET, &stale.to_string())
                .await
                .unwrap()
                .is_some(),
            "an out-of-range ttl must skip the sweep, never reap against a degenerate cutoff"
        );
        drop(handle);
    }
}