klieo_core/runtime/
checkpoint.rs

1//! Persist + resume behaviour for suspended runs (ADR-045). The checkpoint
2//! data type lives in [`crate::checkpoint`]; this module holds the runtime
3//! logic that writes it to `KvStore` and replays it back into a run.
4
5use chrono::{DateTime, Utc};
6use std::sync::Arc;
7use std::time::Duration;
8
9use crate::agent::AgentContext;
10use crate::bus::KvStore;
11use crate::checkpoint::{ApprovalDecision, RunCheckpoint, CHECKPOINT_BUCKET};
12use crate::error::Error;
13use crate::ids::ThreadId;
14use crate::llm::{FinishReason, Message, Role};
15
16/// Single builder for the blocking and streaming suspend paths (ADR-045), so
17/// both persist an identical checkpoint shape.
18pub(crate) async fn build_suspend_checkpoint(
19    ctx: &AgentContext,
20    thread: &ThreadId,
21    step: u32,
22    message: &Message,
23    finish_reason: FinishReason,
24    max_history_tokens: usize,
25) -> Result<RunCheckpoint, Error> {
26    let pending_tool_calls =
27        matches!(finish_reason, FinishReason::ToolCalls).then(|| message.tool_calls.clone());
28    Ok(RunCheckpoint {
29        run_id: ctx.run_id,
30        step_index: step,
31        thread_id: thread.clone(),
32        messages: ctx
33            .short_term
34            .load(thread.clone(), max_history_tokens)
35            .await?,
36        pending_tool_calls,
37        resume_attempted: false,
38        created_at: Utc::now(),
39    })
40}
41
42/// Persisting under a `KvStore` bucket (rather than holding state in memory) is
43/// what lets a different process resume the run. Serialization or store failure
44/// surfaces as `Error` — the run never proceeds past a checkpoint it could not
45/// persist, so an unpersisted suspend cannot be silently lost.
46pub(crate) async fn persist_checkpoint(
47    ctx: &AgentContext,
48    bucket: &str,
49    checkpoint: &RunCheckpoint,
50) -> Result<(), Error> {
51    let bytes = serde_json::to_vec(checkpoint).map_err(|e| Error::Other {
52        message: "checkpoint serialize".into(),
53        source: Some(Box::new(e)),
54    })?;
55    ctx.kv
56        .put(bucket, &checkpoint.run_id.to_string(), bytes.into())
57        .await?;
58    Ok(())
59}
60
61/// Reap abandoned suspended-run checkpoints, returning the count deleted. An
62/// entry whose `created_at` is at or before `cutoff` is removed; the bound is
63/// inclusive. Best-effort throughout: a checkpoint that fails to read or
64/// deserialize, or whose delete fails, is logged and skipped so one bad entry
65/// never stalls the sweep. The `cutoff` is typically `Utc::now() - ttl`.
66///
67/// Enumerates with `keys_paginated` one bounded page at a time and fetches one
68/// checkpoint per key, so a large backlog of abandoned runs pulls neither the
69/// whole key list nor every (potentially conversation-sized) payload into
70/// memory at once — only one page of keys plus a single checkpoint are resident.
71///
72/// A host can drive this on its own schedule, or spawn the opt-in
73/// [`spawn_checkpoint_gc`] task. Distinct from [`crate::spawn_kv_reaper`], which
74/// evicts `{kind}.{stream_id}` keys by resume-buffer liveness: checkpoints key
75/// on `run_id` and have no resume buffer, so age is the only signal an abandoned
76/// one leaves behind.
77pub async fn gc_checkpoints(kv: &dyn KvStore, cutoff: DateTime<Utc>) -> Result<u64, Error> {
78    gc_checkpoints_paged(kv, cutoff, GC_KEY_PAGE).await
79}
80
81/// Handle for the spawned checkpoint-GC task; its `Drop` aborts the task.
82pub struct CheckpointGcHandle {
83    task: Option<tokio::task::JoinHandle<()>>,
84}
85
86impl Drop for CheckpointGcHandle {
87    fn drop(&mut self) {
88        if let Some(task) = self.task.take() {
89            task.abort();
90        }
91    }
92}
93
94/// Spawn the opt-in background task that reaps abandoned suspended-run
95/// checkpoints every `interval`, deleting any older than `ttl`. Key the `ttl` to
96/// the review-item timeout SLA so a checkpoint never outlives the review it backs
97/// (ADR-045 item 7) — without it, abandoned suspensions leak `KvStore` entries
98/// forever. Memory-bounded via [`gc_checkpoints`] (page-walk, one checkpoint
99/// resident at a time). A `ttl` too large to represent skips that sweep rather
100/// than reaping everything. Returns a [`CheckpointGcHandle`] whose `Drop` aborts
101/// the task.
102pub fn spawn_checkpoint_gc(
103    kv: Arc<dyn KvStore>,
104    ttl: Duration,
105    interval: Duration,
106) -> CheckpointGcHandle {
107    let task = tokio::spawn(async move {
108        loop {
109            tokio::time::sleep(interval).await;
110            let Ok(ttl_chrono) = chrono::Duration::from_std(ttl) else {
111                tracing::warn!(
112                    target: "klieo.checkpoint.gc",
113                    ttl = ?ttl,
114                    "configured ttl is out of range; skipping this sweep"
115                );
116                continue;
117            };
118            match gc_checkpoints(kv.as_ref(), Utc::now() - ttl_chrono).await {
119                Ok(reaped) if reaped > 0 => tracing::info!(
120                    target: "klieo.checkpoint.gc",
121                    reaped,
122                    "reaped abandoned suspended-run checkpoints"
123                ),
124                Ok(_) => {}
125                Err(err) => tracing::warn!(
126                    target: "klieo.checkpoint.gc",
127                    error = %err,
128                    "checkpoint gc sweep failed; will retry next interval"
129                ),
130            }
131        }
132    });
133    CheckpointGcHandle { task: Some(task) }
134}
135
136/// Sweep page size: large enough to amortize the per-page round trip, small
137/// enough that one page of keys stays a bounded slice rather than the whole
138/// bucket. The exact value is not load-bearing — any backend that cares
139/// overrides `keys_paginated`.
140const GC_KEY_PAGE: usize = 256;
141
142/// Page-walking core of [`gc_checkpoints`], parameterized on page size so tests
143/// can drive the multi-page cursor loop without seeding a full page of keys.
144async fn gc_checkpoints_paged(
145    kv: &dyn KvStore,
146    cutoff: DateTime<Utc>,
147    page_size: usize,
148) -> Result<u64, Error> {
149    let mut reaped = 0u64;
150    let mut cursor = None;
151    loop {
152        let page = kv
153            .keys_paginated(CHECKPOINT_BUCKET, cursor, page_size)
154            .await?;
155        for key in &page.keys {
156            let Some(checkpoint) = load_checkpoint_for_gc(kv, key).await else {
157                continue;
158            };
159            if checkpoint.created_at > cutoff {
160                continue;
161            }
162            match kv.delete(CHECKPOINT_BUCKET, key).await {
163                Ok(()) => reaped += 1,
164                Err(e) => tracing::warn!(
165                    target: "klieo.checkpoint.gc",
166                    operation = "delete",
167                    key = %key,
168                    error = %e,
169                    "checkpoint delete failed; continuing sweep"
170                ),
171            }
172        }
173        match page.next {
174            Some(c) => cursor = Some(c),
175            None => break,
176        }
177    }
178    Ok(reaped)
179}
180
181/// Read and deserialize one checkpoint for the GC sweep. Returns `None` — after
182/// logging the cause — when the key has already been deleted (a concurrent
183/// reaper raced us), the read fails, or the stored value is not a checkpoint, so
184/// the caller can skip it without aborting the whole sweep.
185async fn load_checkpoint_for_gc(kv: &dyn KvStore, key: &str) -> Option<RunCheckpoint> {
186    let entry = match kv.get(CHECKPOINT_BUCKET, key).await {
187        Ok(Some(entry)) => entry,
188        Ok(None) => return None,
189        Err(e) => {
190            tracing::warn!(
191                target: "klieo.checkpoint.gc",
192                operation = "read",
193                key = %key,
194                error = %e,
195                "checkpoint read failed; skipping"
196            );
197            return None;
198        }
199    };
200    match serde_json::from_slice(&entry.value) {
201        Ok(checkpoint) => Some(checkpoint),
202        Err(e) => {
203            tracing::warn!(
204                target: "klieo.checkpoint.gc",
205                operation = "deserialize",
206                key = %key,
207                error = %e,
208                "skipping undeserializable checkpoint entry"
209            );
210            None
211        }
212    }
213}
214
215/// Resume a run suspended via the review gate (ADR-045).
216///
217/// On `Approved` with pending tool calls, those calls are dispatched here. Each
218/// resume first claims a one-shot latch on the persisted checkpoint by
219/// compare-and-set ([`claim_resume_latch`]): a sequential retry (the store is
220/// already latched) and a concurrent resume (it loses the CAS) are both refused
221/// for a non-idempotent pending call — returning [`Error::ResumeReplayBlocked`]
222/// and leaving the checkpoint for operator reconciliation — rather than risk a
223/// duplicate side effect (e.g. a double payout). Idempotent tools
224/// (`ToolInvoker::is_tool_idempotent`) re-dispatch freely. A within-process
225/// resume with no persisted checkpoint cannot be retried or raced and is
226/// pre-dispatch-safe, so it skips the latch.
227pub async fn resume_from_checkpoint(
228    ctx: &AgentContext,
229    system_prompt: &str,
230    checkpoint: RunCheckpoint,
231    decision: ApprovalDecision,
232    opts: super::RunOptions,
233) -> Result<String, Error> {
234    let thread = checkpoint.thread_id.clone();
235    let latch = RunCheckpoint {
236        resume_attempted: true,
237        ..checkpoint.clone()
238    };
239    ctx.short_term.clear(thread.clone()).await?;
240    ctx.short_term
241        .append_batch(thread.clone(), checkpoint.messages)
242        .await?;
243    match (decision, checkpoint.pending_tool_calls) {
244        (ApprovalDecision::Approved, Some(calls)) => {
245            dispatch_pending_on_resume(ctx, &thread, &calls, &latch, &opts).await?;
246        }
247        (ApprovalDecision::Approved, None) => {}
248        (ApprovalDecision::Rejected { reason }, _) => {
249            ctx.short_term
250                .append(
251                    thread.clone(),
252                    Message {
253                        role: Role::Tool,
254                        content: format!("Human reviewer rejected this step: {reason}"),
255                        tool_calls: vec![],
256                        tool_call_id: None,
257                    },
258                )
259                .await?;
260        }
261    }
262    super::run_loop(ctx, system_prompt, &thread, &opts, checkpoint.step_index).await
263}
264
265/// Claim the resume latch (fail-closed, ADR-045) then dispatch the approved
266/// pending calls. The latch only applies to a persisted checkpoint — the only
267/// thing a retry or a concurrent resume can race on; a bucketless within-process
268/// resume is pre-dispatch-safe and dispatches directly.
269async fn dispatch_pending_on_resume(
270    ctx: &AgentContext,
271    thread: &ThreadId,
272    calls: &[crate::llm::ToolCall],
273    latch: &RunCheckpoint,
274    opts: &super::RunOptions,
275) -> Result<(), Error> {
276    let non_idempotent: Vec<String> = calls
277        .iter()
278        .filter(|c| !ctx.tools.is_tool_idempotent(&c.name))
279        .map(|c| c.name.clone())
280        .collect();
281    if let Some(bucket) = &opts.checkpoint_kv_bucket {
282        claim_resume_latch(ctx, bucket, latch, &non_idempotent).await?;
283    }
284    super::dispatch_tool_calls(ctx, thread, calls, "resume").await
285}
286
287/// Atomically claim the one-shot resume latch by compare-and-set, so neither a
288/// sequential retry nor a concurrent resume can re-fire a non-idempotent tool
289/// (ADR-045). A pending non-idempotent call is refused when the persisted
290/// checkpoint is already latched or the CAS loses to a concurrent claimer;
291/// idempotent batches proceed in either case.
292async fn claim_resume_latch(
293    ctx: &AgentContext,
294    bucket: &str,
295    latch: &RunCheckpoint,
296    non_idempotent: &[String],
297) -> Result<(), Error> {
298    let key = latch.run_id.to_string();
299    let expected = match ctx.kv.get(bucket, &key).await? {
300        Some(entry) => {
301            if checkpoint_is_latched(&entry.value) && !non_idempotent.is_empty() {
302                return Err(resume_blocked(latch.run_id, non_idempotent));
303            }
304            Some(entry.revision)
305        }
306        None => None,
307    };
308    let bytes = serde_json::to_vec(latch).map_err(|e| Error::wrap("checkpoint serialize", e))?;
309    match ctx.kv.cas(bucket, &key, bytes.into(), expected).await {
310        Ok(_) => Ok(()),
311        Err(crate::error::BusError::CasConflict { .. }) if !non_idempotent.is_empty() => {
312            Err(resume_blocked(latch.run_id, non_idempotent))
313        }
314        Err(crate::error::BusError::CasConflict { .. }) => Ok(()),
315        Err(e) => Err(Error::Bus(e)),
316    }
317}
318
319/// A persisted checkpoint counts as latched when it parses and `resume_attempted`
320/// is set. An unparseable value is treated as latched — the fail-closed
321/// direction: refuse rather than risk re-firing on a checkpoint we cannot read.
322fn checkpoint_is_latched(value: &[u8]) -> bool {
323    serde_json::from_slice::<RunCheckpoint>(value)
324        .map(|c| c.resume_attempted)
325        .unwrap_or(true)
326}
327
328fn resume_blocked(run_id: crate::ids::RunId, tools: &[String]) -> Error {
329    tracing::error!(
330        target: "klieo.checkpoint.resume",
331        run_id = %run_id,
332        tools = ?tools,
333        "resume blocked: non-idempotent pending tool calls cannot be proven un-fired (fail-closed, ADR-045) — operator reconciliation required"
334    );
335    Error::ResumeReplayBlocked {
336        run_id,
337        tools: tools.to_vec(),
338    }
339}
340
341#[cfg(test)]
342mod tests {
343    use super::*;
344    use crate::bus::{KvEntry, Lease, Revision};
345    use crate::error::BusError;
346    use crate::ids::{RunId, ThreadId};
347    use crate::llm::ToolCall;
348    use crate::runtime::RunOptions;
349    use crate::test_utils::{fake_context, fake_kv, FakeLlmClient, FakeLlmStep, FakeToolInvoker};
350    use async_trait::async_trait;
351    use bytes::Bytes;
352    use chrono::Utc;
353    use std::time::Duration;
354
355    struct KeysFailKv;
356
357    #[async_trait]
358    impl crate::bus::KvStore for KeysFailKv {
359        async fn get(&self, _: &str, _: &str) -> Result<Option<KvEntry>, BusError> {
360            Err(BusError::Unsupported("get".into()))
361        }
362        async fn put(&self, _: &str, _: &str, _: Bytes) -> Result<Revision, BusError> {
363            Err(BusError::Unsupported("put".into()))
364        }
365        async fn cas(
366            &self,
367            _: &str,
368            _: &str,
369            _: Bytes,
370            _: Option<Revision>,
371        ) -> Result<Revision, BusError> {
372            Err(BusError::Unsupported("cas".into()))
373        }
374        async fn delete(&self, _: &str, _: &str) -> Result<(), BusError> {
375            Err(BusError::Unsupported("delete".into()))
376        }
377        async fn lease(&self, _: &str, _: &str, _: Duration) -> Result<Lease, BusError> {
378            Err(BusError::Unsupported("lease".into()))
379        }
380        async fn keys(&self, _: &str) -> Result<Vec<String>, BusError> {
381            Err(BusError::Unsupported("enumerate unavailable".into()))
382        }
383    }
384
385    struct DeleteFailKv {
386        value: Bytes,
387    }
388
389    #[async_trait]
390    impl crate::bus::KvStore for DeleteFailKv {
391        async fn get(&self, _: &str, _: &str) -> Result<Option<KvEntry>, BusError> {
392            Ok(Some(KvEntry {
393                value: self.value.clone(),
394                revision: 1,
395            }))
396        }
397        async fn put(&self, _: &str, _: &str, _: Bytes) -> Result<Revision, BusError> {
398            Err(BusError::Unsupported("put".into()))
399        }
400        async fn cas(
401            &self,
402            _: &str,
403            _: &str,
404            _: Bytes,
405            _: Option<Revision>,
406        ) -> Result<Revision, BusError> {
407            Err(BusError::Unsupported("cas".into()))
408        }
409        async fn delete(&self, _: &str, _: &str) -> Result<(), BusError> {
410            Err(BusError::Unsupported("delete boom".into()))
411        }
412        async fn lease(&self, _: &str, _: &str, _: Duration) -> Result<Lease, BusError> {
413            Err(BusError::Unsupported("lease".into()))
414        }
415        async fn keys(&self, _: &str) -> Result<Vec<String>, BusError> {
416            Ok(vec!["stale".to_string()])
417        }
418    }
419
420    /// `get` returns a stored (un-latched) value, but every `cas` loses with a
421    /// `CasConflict` — models a concurrent resume that claimed the latch between
422    /// our read and our compare-and-set.
423    struct ConflictKv {
424        value: Bytes,
425    }
426
427    #[async_trait]
428    impl crate::bus::KvStore for ConflictKv {
429        async fn get(&self, _: &str, _: &str) -> Result<Option<KvEntry>, BusError> {
430            Ok(Some(KvEntry {
431                value: self.value.clone(),
432                revision: 1,
433            }))
434        }
435        async fn put(&self, _: &str, _: &str, _: Bytes) -> Result<Revision, BusError> {
436            Err(BusError::Unsupported("put".into()))
437        }
438        async fn cas(
439            &self,
440            _: &str,
441            _: &str,
442            _: Bytes,
443            _: Option<Revision>,
444        ) -> Result<Revision, BusError> {
445            Err(BusError::CasConflict {
446                expected: 1,
447                actual: 2,
448            })
449        }
450        async fn delete(&self, _: &str, _: &str) -> Result<(), BusError> {
451            Err(BusError::Unsupported("delete".into()))
452        }
453        async fn lease(&self, _: &str, _: &str, _: Duration) -> Result<Lease, BusError> {
454            Err(BusError::Unsupported("lease".into()))
455        }
456        async fn keys(&self, _: &str) -> Result<Vec<String>, BusError> {
457            Err(BusError::Unsupported("keys".into()))
458        }
459    }
460    use std::sync::Arc;
461
462    #[test]
463    fn checkpoint_round_trips_through_json() {
464        let cp = checkpoint_with_pending_tool(ThreadId::new("t-rt"), RunId::new());
465        let bytes = serde_json::to_vec(&cp).unwrap();
466        let back: RunCheckpoint = serde_json::from_slice(&bytes).unwrap();
467        assert_eq!(back.step_index, cp.step_index);
468        assert_eq!(back.run_id, cp.run_id);
469        assert_eq!(back.thread_id, cp.thread_id);
470
471        assert_eq!(back.messages.len(), cp.messages.len(), "history length");
472        assert_eq!(back.messages[0].role, cp.messages[0].role);
473        assert_eq!(back.messages[0].content, cp.messages[0].content);
474
475        let restored = back
476            .pending_tool_calls
477            .expect("pending tool calls must survive the round-trip");
478        let original = cp.pending_tool_calls.unwrap();
479        assert_eq!(restored.len(), original.len());
480        assert_eq!(restored[0].id, original[0].id);
481        assert_eq!(restored[0].name, original[0].name);
482        assert_eq!(restored[0].args, original[0].args);
483    }
484
485    fn echo_tool_invoker() -> Arc<FakeToolInvoker> {
486        Arc::new(FakeToolInvoker::new().with_tool("echo", "echo back", Ok))
487    }
488
489    fn checkpoint_with_pending_tool(thread: ThreadId, run_id: RunId) -> RunCheckpoint {
490        RunCheckpoint {
491            run_id,
492            step_index: 1,
493            thread_id: thread,
494            messages: vec![Message {
495                role: Role::User,
496                content: "go".into(),
497                tool_calls: vec![],
498                tool_call_id: None,
499            }],
500            pending_tool_calls: Some(vec![ToolCall {
501                id: "tc-1".into(),
502                name: "echo".into(),
503                args: serde_json::json!({"x": 1}),
504            }]),
505            resume_attempted: false,
506            created_at: Utc::now(),
507        }
508    }
509
510    fn checkpoint_no_pending(thread: ThreadId, run_id: RunId) -> RunCheckpoint {
511        RunCheckpoint {
512            run_id,
513            step_index: 1,
514            thread_id: thread,
515            messages: vec![Message {
516                role: Role::User,
517                content: "go".into(),
518                tool_calls: vec![],
519                tool_call_id: None,
520            }],
521            pending_tool_calls: None,
522            resume_attempted: false,
523            created_at: Utc::now(),
524        }
525    }
526
527    #[tokio::test]
528    async fn resume_approved_dispatches_pending_then_completes() {
529        let mut ctx = fake_context("resume-test");
530        ctx.llm =
531            Arc::new(FakeLlmClient::new("fake").with_steps(vec![FakeLlmStep::Text("done".into())]));
532        ctx.kv = fake_kv();
533        ctx.tools = echo_tool_invoker();
534
535        let thread = ThreadId::new("t-resume-approved");
536        let cp = checkpoint_with_pending_tool(thread.clone(), ctx.run_id);
537
538        let out = resume_from_checkpoint(
539            &ctx,
540            "sys",
541            cp,
542            ApprovalDecision::Approved,
543            RunOptions::default(),
544        )
545        .await
546        .unwrap();
547
548        assert_eq!(out, "done");
549
550        let history = ctx.short_term.load(thread, 8192).await.unwrap();
551        let tool_msgs: Vec<_> = history.iter().filter(|m| m.role == Role::Tool).collect();
552        assert_eq!(
553            tool_msgs.len(),
554            1,
555            "dispatched tool must leave a Role::Tool message"
556        );
557        assert_eq!(tool_msgs[0].tool_call_id.as_deref(), Some("tc-1"));
558    }
559
560    #[tokio::test]
561    async fn resume_approved_without_pending_calls_completes_and_injects_no_tool_message() {
562        let mut ctx = fake_context("resume-approved-none");
563        ctx.llm =
564            Arc::new(FakeLlmClient::new("fake").with_steps(vec![FakeLlmStep::Text("done".into())]));
565        ctx.kv = fake_kv();
566
567        let thread = ThreadId::new("t-resume-approved-none");
568        let cp = checkpoint_no_pending(thread.clone(), ctx.run_id);
569
570        let out = resume_from_checkpoint(
571            &ctx,
572            "sys",
573            cp,
574            ApprovalDecision::Approved,
575            RunOptions::default(),
576        )
577        .await
578        .unwrap();
579
580        assert_eq!(out, "done");
581
582        let history = ctx.short_term.load(thread, 8192).await.unwrap();
583        let tool_msgs = history.iter().filter(|m| m.role == Role::Tool).count();
584        assert_eq!(
585            tool_msgs, 0,
586            "approving a step with no pending tool calls must inject nothing"
587        );
588    }
589
590    #[tokio::test]
591    async fn resume_rejected_injects_tool_message_then_completes() {
592        let mut ctx = fake_context("resume-reject");
593        ctx.llm =
594            Arc::new(FakeLlmClient::new("fake").with_steps(vec![FakeLlmStep::Text("done".into())]));
595        ctx.kv = fake_kv();
596
597        let thread = ThreadId::new("t-resume-rejected");
598        let cp = checkpoint_no_pending(thread.clone(), ctx.run_id);
599
600        let out = resume_from_checkpoint(
601            &ctx,
602            "sys",
603            cp,
604            ApprovalDecision::Rejected {
605                reason: "no".into(),
606            },
607            RunOptions::default(),
608        )
609        .await
610        .unwrap();
611
612        assert_eq!(out, "done");
613
614        let history = ctx.short_term.load(thread, 8192).await.unwrap();
615        let rejection: Vec<_> = history
616            .iter()
617            .filter(|m| m.role == Role::Tool && m.content.contains("no"))
618            .collect();
619        assert_eq!(rejection.len(), 1, "rejected message must be appended");
620        assert!(rejection[0]
621            .content
622            .contains("Human reviewer rejected this step: no"));
623    }
624
625    #[tokio::test]
626    async fn resume_rejected_with_pending_calls_does_not_dispatch_them() {
627        let mut ctx = fake_context("resume-reject-pending");
628        ctx.llm =
629            Arc::new(FakeLlmClient::new("fake").with_steps(vec![FakeLlmStep::Text("done".into())]));
630        ctx.kv = fake_kv();
631        ctx.tools = echo_tool_invoker();
632
633        let thread = ThreadId::new("t-resume-reject-pending");
634        let cp = checkpoint_with_pending_tool(thread.clone(), ctx.run_id);
635
636        let out = resume_from_checkpoint(
637            &ctx,
638            "sys",
639            cp,
640            ApprovalDecision::Rejected {
641                reason: "blocked".into(),
642            },
643            RunOptions::default(),
644        )
645        .await
646        .unwrap();
647
648        assert_eq!(out, "done");
649
650        let history = ctx.short_term.load(thread, 8192).await.unwrap();
651        assert!(
652            !history
653                .iter()
654                .any(|m| m.tool_call_id.as_deref() == Some("tc-1")),
655            "a rejected step must NOT dispatch its pending tool calls"
656        );
657        assert!(
658            history
659                .iter()
660                .any(|m| m.role == Role::Tool && m.content.contains("blocked")),
661            "the rejection reason must be fed back to the model"
662        );
663    }
664
665    async fn seed_latched(ctx: &AgentContext, thread: &ThreadId) {
666        let mut latched = checkpoint_with_pending_tool(thread.clone(), ctx.run_id);
667        latched.resume_attempted = true;
668        ctx.kv
669            .put(
670                CHECKPOINT_BUCKET,
671                &ctx.run_id.to_string(),
672                serde_json::to_vec(&latched).unwrap().into(),
673            )
674            .await
675            .unwrap();
676    }
677
678    #[tokio::test]
679    async fn resume_retry_with_non_idempotent_pending_fails_closed() {
680        let mut ctx = fake_context("resume-retry-nonidem");
681        ctx.llm =
682            Arc::new(FakeLlmClient::new("fake").with_steps(vec![FakeLlmStep::Text("done".into())]));
683        ctx.kv = fake_kv();
684        ctx.tools = echo_tool_invoker(); // `echo` uses the default: non-idempotent
685
686        let thread = ThreadId::new("t-retry-nonidem");
687        seed_latched(&ctx, &thread).await; // a prior resume already latched the store
688        let cp = checkpoint_with_pending_tool(thread.clone(), ctx.run_id);
689
690        let err = resume_from_checkpoint(
691            &ctx,
692            "sys",
693            cp,
694            ApprovalDecision::Approved,
695            RunOptions::default().with_checkpoint_bucket(CHECKPOINT_BUCKET),
696        )
697        .await
698        .unwrap_err();
699
700        match err {
701            Error::ResumeReplayBlocked { tools, .. } => assert_eq!(
702                tools,
703                vec!["echo".to_string()],
704                "the refused non-idempotent tool must be named"
705            ),
706            other => panic!("expected ResumeReplayBlocked, got {other:?}"),
707        }
708
709        let history = ctx.short_term.load(thread, 8192).await.unwrap();
710        assert!(
711            !history
712                .iter()
713                .any(|m| m.tool_call_id.as_deref() == Some("tc-1")),
714            "a fail-closed resume must NOT re-dispatch the non-idempotent pending call"
715        );
716    }
717
718    #[tokio::test]
719    async fn resume_retry_with_idempotent_pending_redispatches() {
720        let mut ctx = fake_context("resume-retry-idem");
721        ctx.llm =
722            Arc::new(FakeLlmClient::new("fake").with_steps(vec![FakeLlmStep::Text("done".into())]));
723        ctx.kv = fake_kv();
724        ctx.tools = Arc::new(FakeToolInvoker::new().with_idempotent_tool("echo", "echo back", Ok));
725
726        let thread = ThreadId::new("t-retry-idem");
727        seed_latched(&ctx, &thread).await;
728        let cp = checkpoint_with_pending_tool(thread.clone(), ctx.run_id);
729
730        let out = resume_from_checkpoint(
731            &ctx,
732            "sys",
733            cp,
734            ApprovalDecision::Approved,
735            RunOptions::default().with_checkpoint_bucket(CHECKPOINT_BUCKET),
736        )
737        .await
738        .unwrap();
739        assert_eq!(out, "done");
740
741        let history = ctx.short_term.load(thread, 8192).await.unwrap();
742        assert!(
743            history
744                .iter()
745                .any(|m| m.tool_call_id.as_deref() == Some("tc-1")),
746            "an idempotent tool is safe to re-dispatch even after a prior resume latch"
747        );
748    }
749
750    #[tokio::test]
751    async fn concurrent_resume_losing_the_cas_fails_closed() {
752        // The store still reads as un-latched, but the compare-and-set loses —
753        // a concurrent resume claimed the latch between our read and write. A
754        // non-idempotent call must be refused, not re-fired (CWE-367 TOCTOU).
755        let mut ctx = fake_context("resume-cas-conflict");
756        ctx.llm =
757            Arc::new(FakeLlmClient::new("fake").with_steps(vec![FakeLlmStep::Text("done".into())]));
758        ctx.tools = echo_tool_invoker();
759        let thread = ThreadId::new("t-cas-conflict");
760        let unlatched =
761            serde_json::to_vec(&checkpoint_with_pending_tool(thread.clone(), ctx.run_id))
762                .unwrap()
763                .into();
764        ctx.kv = Arc::new(ConflictKv { value: unlatched });
765
766        let err = resume_from_checkpoint(
767            &ctx,
768            "sys",
769            checkpoint_with_pending_tool(thread, ctx.run_id),
770            ApprovalDecision::Approved,
771            RunOptions::default().with_checkpoint_bucket(CHECKPOINT_BUCKET),
772        )
773        .await
774        .unwrap_err();
775
776        assert!(
777            matches!(err, Error::ResumeReplayBlocked { .. }),
778            "losing the latch CAS to a concurrent resume must fail closed, got {err:?}"
779        );
780    }
781
782    #[tokio::test]
783    async fn first_resume_latches_attempt_before_dispatch() {
784        let mut ctx = fake_context("resume-latch");
785        ctx.llm =
786            Arc::new(FakeLlmClient::new("fake").with_steps(vec![FakeLlmStep::Text("done".into())]));
787        ctx.kv = fake_kv();
788        ctx.tools = echo_tool_invoker();
789
790        let thread = ThreadId::new("t-latch");
791        let run_id = ctx.run_id;
792        let cp = checkpoint_with_pending_tool(thread, run_id);
793
794        resume_from_checkpoint(
795            &ctx,
796            "sys",
797            cp,
798            ApprovalDecision::Approved,
799            RunOptions::default().with_checkpoint_bucket(CHECKPOINT_BUCKET),
800        )
801        .await
802        .unwrap();
803
804        let entry = ctx
805            .kv
806            .get(CHECKPOINT_BUCKET, &run_id.to_string())
807            .await
808            .unwrap()
809            .expect("first resume must persist the latched checkpoint");
810        let latched: RunCheckpoint = serde_json::from_slice(&entry.value).unwrap();
811        assert!(
812            latched.resume_attempted,
813            "first resume must latch resume_attempted=true before dispatch, so a retry is recognised"
814        );
815    }
816
817    #[test]
818    fn corrupt_checkpoint_bytes_treated_as_latched() {
819        assert!(
820            checkpoint_is_latched(b"not-json"),
821            "an unparseable stored checkpoint must read as latched — refuse rather than risk a re-fire"
822        );
823        let unlatched = serde_json::to_vec(&checkpoint_no_pending(
824            ThreadId::new("t-parse"),
825            RunId::new(),
826        ))
827        .unwrap();
828        assert!(
829            !checkpoint_is_latched(&unlatched),
830            "a well-formed un-latched checkpoint must read as not latched"
831        );
832    }
833
834    #[tokio::test]
835    async fn bucketless_resume_skips_latch_and_dispatches() {
836        let mut ctx = fake_context("resume-no-bucket");
837        ctx.llm =
838            Arc::new(FakeLlmClient::new("fake").with_steps(vec![FakeLlmStep::Text("done".into())]));
839        ctx.kv = fake_kv();
840        ctx.tools = echo_tool_invoker(); // non-idempotent
841
842        let thread = ThreadId::new("t-no-bucket");
843        let run_id = ctx.run_id;
844        let cp = checkpoint_with_pending_tool(thread.clone(), run_id);
845
846        let out = resume_from_checkpoint(
847            &ctx,
848            "sys",
849            cp,
850            ApprovalDecision::Approved,
851            RunOptions::default(),
852        )
853        .await
854        .unwrap();
855        assert_eq!(out, "done");
856
857        let history = ctx.short_term.load(thread, 8192).await.unwrap();
858        assert!(
859            history
860                .iter()
861                .any(|m| m.tool_call_id.as_deref() == Some("tc-1")),
862            "a within-process resume with no bucket is pre-dispatch-safe and dispatches directly"
863        );
864        assert!(
865            ctx.kv
866                .get(CHECKPOINT_BUCKET, &run_id.to_string())
867                .await
868                .unwrap()
869                .is_none(),
870            "no bucket means no latch is persisted"
871        );
872    }
873
874    #[tokio::test]
875    async fn concurrent_cas_conflict_with_idempotent_batch_proceeds() {
876        let mut ctx = fake_context("resume-conflict-idem");
877        ctx.llm =
878            Arc::new(FakeLlmClient::new("fake").with_steps(vec![FakeLlmStep::Text("done".into())]));
879        ctx.tools = Arc::new(FakeToolInvoker::new().with_idempotent_tool("echo", "echo back", Ok));
880        let thread = ThreadId::new("t-conflict-idem");
881        let unlatched =
882            serde_json::to_vec(&checkpoint_with_pending_tool(thread.clone(), ctx.run_id))
883                .unwrap()
884                .into();
885        ctx.kv = Arc::new(ConflictKv { value: unlatched });
886
887        let out = resume_from_checkpoint(
888            &ctx,
889            "sys",
890            checkpoint_with_pending_tool(thread.clone(), ctx.run_id),
891            ApprovalDecision::Approved,
892            RunOptions::default().with_checkpoint_bucket(CHECKPOINT_BUCKET),
893        )
894        .await
895        .unwrap();
896        assert_eq!(out, "done");
897
898        let history = ctx.short_term.load(thread, 8192).await.unwrap();
899        assert!(
900            history
901                .iter()
902                .any(|m| m.tool_call_id.as_deref() == Some("tc-1")),
903            "an idempotent tool proceeds even when it loses the latch CAS"
904        );
905    }
906
907    #[test]
908    fn legacy_checkpoint_without_resume_attempted_defaults_false() {
909        // A checkpoint persisted before the latch field existed must stay
910        // readable and default to the first-attempt state — never silently
911        // treat an old checkpoint as already-attempted (which would block a
912        // legitimate first resume).
913        let run_id = RunId::new();
914        let legacy = format!(
915            r#"{{"run_id":"{run_id}","step_index":1,"thread_id":"t-legacy","messages":[],"pending_tool_calls":null,"created_at":"2026-06-12T00:00:00Z"}}"#
916        );
917        let back: RunCheckpoint = serde_json::from_str(&legacy).unwrap();
918        assert!(!back.resume_attempted);
919    }
920
921    async fn seed_checkpoint(
922        kv: &Arc<dyn crate::bus::KvStore>,
923        run_id: RunId,
924        created_at: chrono::DateTime<Utc>,
925    ) {
926        let key = run_id.to_string();
927        let mut cp = checkpoint_with_pending_tool(ThreadId::new("t-gc"), run_id);
928        cp.created_at = created_at;
929        let bytes = serde_json::to_vec(&cp).unwrap();
930        kv.put(CHECKPOINT_BUCKET, &key, bytes.into()).await.unwrap();
931    }
932
933    #[tokio::test]
934    async fn gc_checkpoints_reaps_only_entries_at_or_before_cutoff() {
935        let kv = fake_kv();
936        let now = Utc::now();
937        let cutoff = now - chrono::Duration::hours(1);
938        let stale = RunId::new();
939        let at_cutoff = RunId::new();
940        let fresh = RunId::new();
941        let stale_key = stale.to_string();
942        let at_cutoff_key = at_cutoff.to_string();
943        let fresh_key = fresh.to_string();
944        seed_checkpoint(&kv, stale, now - chrono::Duration::hours(2)).await;
945        seed_checkpoint(&kv, at_cutoff, cutoff).await;
946        seed_checkpoint(&kv, fresh, now).await;
947
948        let reaped = gc_checkpoints(kv.as_ref(), cutoff).await.unwrap();
949
950        assert_eq!(
951            reaped, 2,
952            "the stale checkpoint and the one exactly at the cutoff are both reaped (bound is <=)"
953        );
954        assert!(
955            kv.get(CHECKPOINT_BUCKET, &stale_key)
956                .await
957                .unwrap()
958                .is_none(),
959            "the stale checkpoint is deleted"
960        );
961        assert!(
962            kv.get(CHECKPOINT_BUCKET, &at_cutoff_key)
963                .await
964                .unwrap()
965                .is_none(),
966            "a checkpoint whose created_at equals the cutoff is reaped — the bound is inclusive"
967        );
968        assert!(
969            kv.get(CHECKPOINT_BUCKET, &fresh_key)
970                .await
971                .unwrap()
972                .is_some(),
973            "a checkpoint newer than the cutoff must survive the sweep"
974        );
975    }
976
977    #[tokio::test]
978    async fn gc_checkpoints_paged_reaps_eligible_across_pages() {
979        let kv = fake_kv();
980        let now = Utc::now();
981        let cutoff = now;
982        let stale = now - chrono::Duration::hours(1);
983        let ids: Vec<RunId> = (0..5).map(|_| RunId::new()).collect();
984        for id in &ids {
985            seed_checkpoint(&kv, *id, stale).await;
986        }
987
988        // page_size 2 over 5 keys drives three pages + two cursor advances.
989        let reaped = gc_checkpoints_paged(kv.as_ref(), cutoff, 2).await.unwrap();
990
991        assert_eq!(reaped, 5);
992        for id in &ids {
993            assert!(
994                kv.get(CHECKPOINT_BUCKET, &id.to_string())
995                    .await
996                    .unwrap()
997                    .is_none(),
998                "every eligible checkpoint across all pages is reaped"
999            );
1000        }
1001    }
1002
1003    #[tokio::test]
1004    async fn gc_checkpoints_skips_undeserializable_entry_without_failing() {
1005        let kv = fake_kv();
1006        let now = Utc::now();
1007        let stale = RunId::new();
1008        let stale_key = stale.to_string();
1009        seed_checkpoint(&kv, stale, now - chrono::Duration::hours(2)).await;
1010        kv.put(
1011            CHECKPOINT_BUCKET,
1012            "not-a-checkpoint",
1013            b"junk".to_vec().into(),
1014        )
1015        .await
1016        .unwrap();
1017
1018        let reaped = gc_checkpoints(kv.as_ref(), now).await.unwrap();
1019
1020        assert_eq!(reaped, 1, "the valid stale checkpoint is still reaped");
1021        assert!(
1022            kv.get(CHECKPOINT_BUCKET, &stale_key)
1023                .await
1024                .unwrap()
1025                .is_none(),
1026            "the stale checkpoint is gone"
1027        );
1028        assert!(
1029            kv.get(CHECKPOINT_BUCKET, "not-a-checkpoint")
1030                .await
1031                .unwrap()
1032                .is_some(),
1033            "an undeserializable entry is skipped, not deleted — a foreign value cannot stall the sweep"
1034        );
1035    }
1036
1037    #[tokio::test]
1038    async fn gc_checkpoints_empty_bucket_is_zero() {
1039        let kv = fake_kv();
1040        let reaped = gc_checkpoints(kv.as_ref(), Utc::now()).await.unwrap();
1041        assert_eq!(reaped, 0, "an empty bucket reaps nothing");
1042    }
1043
1044    #[tokio::test]
1045    async fn gc_checkpoints_propagates_enumerate_failure() {
1046        let kv = KeysFailKv;
1047        let err = gc_checkpoints(&kv, Utc::now()).await.unwrap_err();
1048        assert!(
1049            matches!(err, Error::Bus(BusError::Unsupported(_))),
1050            "a keys() failure must propagate — the sweep cannot enumerate, so it must surface the error, not silently report zero reaped"
1051        );
1052    }
1053
1054    #[tokio::test]
1055    async fn gc_checkpoints_skips_entry_whose_delete_fails() {
1056        let mut cp = checkpoint_with_pending_tool(ThreadId::new("t-del"), RunId::new());
1057        cp.created_at = Utc::now() - chrono::Duration::hours(1);
1058        let kv = DeleteFailKv {
1059            value: Bytes::from(serde_json::to_vec(&cp).unwrap()),
1060        };
1061
1062        let reaped = gc_checkpoints(&kv, Utc::now()).await.unwrap();
1063
1064        assert_eq!(
1065            reaped, 0,
1066            "a stale checkpoint whose delete fails is logged and skipped, not counted; the sweep still returns Ok"
1067        );
1068    }
1069
1070    #[tokio::test]
1071    async fn spawn_checkpoint_gc_reaps_stale_and_keeps_fresh() {
1072        let kv = fake_kv();
1073        let stale = RunId::new();
1074        let fresh = RunId::new();
1075        seed_checkpoint(&kv, stale, Utc::now() - chrono::Duration::hours(2)).await;
1076        seed_checkpoint(&kv, fresh, Utc::now()).await;
1077
1078        // ttl 1h: the 2h-old checkpoint is past it, the just-created one is not.
1079        let handle = spawn_checkpoint_gc(
1080            Arc::clone(&kv),
1081            Duration::from_secs(3600),
1082            Duration::from_millis(5),
1083        );
1084        // Let several sweep intervals elapse.
1085        tokio::time::sleep(Duration::from_millis(80)).await;
1086
1087        assert!(
1088            kv.get(CHECKPOINT_BUCKET, &stale.to_string())
1089                .await
1090                .unwrap()
1091                .is_none(),
1092            "the spawned gc must reap a checkpoint older than the ttl"
1093        );
1094        assert!(
1095            kv.get(CHECKPOINT_BUCKET, &fresh.to_string())
1096                .await
1097                .unwrap()
1098                .is_some(),
1099            "a checkpoint within the ttl must survive"
1100        );
1101        drop(handle);
1102    }
1103
1104    #[tokio::test]
1105    async fn checkpoint_gc_handle_drop_stops_reaping() {
1106        let kv = fake_kv();
1107        let stale = RunId::new();
1108        seed_checkpoint(&kv, stale, Utc::now() - chrono::Duration::hours(2)).await;
1109
1110        // Drop before the first interval elapses, then wait well past several
1111        // intervals: the abort must prevent any sweep (contrast with
1112        // `spawn_checkpoint_gc_reaps_stale_and_keeps_fresh`, identical timing
1113        // but no drop, where the same stale checkpoint IS reaped).
1114        let handle = spawn_checkpoint_gc(
1115            Arc::clone(&kv),
1116            Duration::from_secs(3600),
1117            Duration::from_millis(5),
1118        );
1119        drop(handle);
1120        tokio::time::sleep(Duration::from_millis(80)).await;
1121
1122        assert!(
1123            kv.get(CHECKPOINT_BUCKET, &stale.to_string())
1124                .await
1125                .unwrap()
1126                .is_some(),
1127            "dropping the handle aborts the task, so the stale checkpoint is never reaped"
1128        );
1129    }
1130
1131    #[tokio::test]
1132    async fn spawn_checkpoint_gc_with_out_of_range_ttl_skips_sweep() {
1133        let kv = fake_kv();
1134        let stale = RunId::new();
1135        seed_checkpoint(&kv, stale, Utc::now() - chrono::Duration::hours(2)).await;
1136
1137        // `Duration::MAX` overflows `chrono::Duration::from_std`, so each sweep
1138        // skips rather than reaping everything against a degenerate cutoff.
1139        let handle = spawn_checkpoint_gc(Arc::clone(&kv), Duration::MAX, Duration::from_millis(5));
1140        tokio::time::sleep(Duration::from_millis(80)).await;
1141
1142        assert!(
1143            kv.get(CHECKPOINT_BUCKET, &stale.to_string())
1144                .await
1145                .unwrap()
1146                .is_some(),
1147            "an out-of-range ttl must skip the sweep, never reap against a degenerate cutoff"
1148        );
1149        drop(handle);
1150    }
1151}
klieo_core/runtime/checkpoint.rs

klieo_core/runtime/
checkpoint.rs
