1use async_trait::async_trait;
12use serde_json::{json, Value};
13use std::path::PathBuf;
14use std::sync::{Arc, Mutex};
15
16use super::Tool;
17use crate::checkpoint::{CheckpointId, ShadowRepo};
18use crate::checkpoint_log::{read_log, CheckpointLogWriter, CheckpointRecord, TouchedVia};
19use crate::error::{Error, Result};
20use crate::llm::ToolSpec;
21use crate::tools::TouchedFiles;
22
23#[derive(Clone)]
28pub struct CheckpointToolCtx {
29 pub repo: Arc<Mutex<ShadowRepo>>,
30 pub session_id: String,
31}
32
33pub struct CheckpointList {
36 ctx: CheckpointToolCtx,
37}
38
39impl CheckpointList {
40 pub fn new(ctx: CheckpointToolCtx) -> Self {
41 Self { ctx }
42 }
43}
44
45#[async_trait]
46impl Tool for CheckpointList {
47 fn is_deferred(&self) -> bool {
48 true
49 }
50
51 fn spec(&self) -> ToolSpec {
52 ToolSpec {
53 name: "checkpoint_list".into(),
54 description: "List checkpoints for the current session, newest first. \
55 Checkpoints are saved by the agent via `checkpoint_save` or \
56 by the runtime at session boundaries. Use `checkpoint_diff` \
57 to inspect the changes captured in a checkpoint."
58 .into(),
59 parameters: json!({"type": "object", "properties": {}}),
60 }
61 }
62
63 fn side_effect_class(&self) -> crate::tools::ToolSideEffect {
64 crate::tools::ToolSideEffect::ReadOnly
65 }
66
67 async fn execute(&self, _args: Value) -> Result<String> {
68 let infos = {
69 let repo = self.ctx.repo.lock().map_err(|_| Error::Tool {
70 name: "checkpoint_list".into(),
71 call_id: None,
72 message: "lock poisoned".into(),
73 })?;
74 repo.list_for_session(&self.ctx.session_id)?
75 };
76 if infos.is_empty() {
77 return Ok("no checkpoints in this session yet".to_string());
78 }
79 let mut out = String::new();
80 for i in &infos {
81 out.push_str(&format!(
82 "[{}] {} ({} file(s) changed) ts={}\n",
83 i.id, i.message, i.files_changed, i.timestamp
84 ));
85 }
86 Ok(out.trim_end().to_string())
87 }
88}
89
90pub struct CheckpointDiff {
93 ctx: CheckpointToolCtx,
94}
95
96impl CheckpointDiff {
97 pub fn new(ctx: CheckpointToolCtx) -> Self {
98 Self { ctx }
99 }
100}
101
102#[async_trait]
103impl Tool for CheckpointDiff {
104 fn is_deferred(&self) -> bool {
105 true
106 }
107
108 fn spec(&self) -> ToolSpec {
109 ToolSpec {
110 name: "checkpoint_diff".into(),
111 description: "Show a unified diff between two checkpoints in this session. \
112 If `b` is omitted, diffs `a` against the current workspace state."
113 .into(),
114 parameters: json!({
115 "type": "object",
116 "properties": {
117 "a": {"type": "string", "description": "Older checkpoint id."},
118 "b": {"type": "string", "description": "Optional newer checkpoint id; omit for current workspace."}
119 },
120 "required": ["a"]
121 }),
122 }
123 }
124
125 fn side_effect_class(&self) -> crate::tools::ToolSideEffect {
126 crate::tools::ToolSideEffect::ReadOnly
127 }
128
129 async fn execute(&self, args: Value) -> Result<String> {
130 let a = args["a"].as_str().ok_or_else(|| Error::BadToolArgs {
131 name: "checkpoint_diff".into(),
132 message: "missing `a`".into(),
133 })?;
134 let b = args["b"].as_str();
135 let aid = CheckpointId(a.to_string());
136 let bid = b.map(|s| CheckpointId(s.to_string()));
137 let diff = {
138 let repo = self.ctx.repo.lock().map_err(|_| Error::Tool {
139 name: "checkpoint_diff".into(),
140 call_id: None,
141 message: "lock poisoned".into(),
142 })?;
143 repo.diff(&aid, bid.as_ref(), &[])?
144 };
145 if diff.trim().is_empty() {
146 Ok("no differences".to_string())
147 } else {
148 Ok(diff)
149 }
150 }
151}
152
153#[derive(Clone)]
161pub struct CheckpointSaveCtx {
162 pub repo: Arc<Mutex<ShadowRepo>>,
163 pub session_id: String,
164 pub touched_files: Option<Arc<Mutex<TouchedFiles>>>,
165 pub writer: Arc<Mutex<CheckpointLogWriter>>,
166 pub turn_index: Arc<std::sync::atomic::AtomicUsize>,
167 pub log_path: PathBuf,
168}
169
170pub struct CheckpointSave {
171 ctx: CheckpointSaveCtx,
172}
173
174impl CheckpointSave {
175 pub fn new(ctx: CheckpointSaveCtx) -> Self {
176 Self { ctx }
177 }
178}
179
180#[async_trait]
181impl Tool for CheckpointSave {
182 fn is_deferred(&self) -> bool {
183 true
184 }
185
186 fn spec(&self) -> ToolSpec {
187 ToolSpec {
188 name: "checkpoint_save".into(),
189 description: "Save an explicit restore point for the current session. \
190 Call this before making a risky batch of changes, or after \
191 completing a logical unit of work you might want to revert \
192 to. Unlike automatic checkpoints (which no longer exist), \
193 this runs only when you call it."
194 .into(),
195 parameters: json!({
196 "type": "object",
197 "properties": {
198 "message": {
199 "type": "string",
200 "description": "A short label for this checkpoint, e.g. 'before refactor' or 'after adding tests'. Defaults to the current turn number if omitted."
201 }
202 }
203 }),
204 }
205 }
206
207 fn side_effect_class(&self) -> crate::tools::ToolSideEffect {
208 crate::tools::ToolSideEffect::ReadOnly
209 }
210
211 async fn execute(&self, args: Value) -> Result<String> {
212 let turn = self
213 .ctx
214 .turn_index
215 .load(std::sync::atomic::Ordering::Relaxed);
216 let message = args
217 .get("message")
218 .and_then(|v| v.as_str())
219 .map(|s| s.to_string())
220 .unwrap_or_else(|| format!("turn {turn}"));
221
222 let (mut paths, saw_shell) = match &self.ctx.touched_files {
224 Some(slot) => match slot.lock() {
225 Ok(t) => (t.paths_sorted(), t.saw_shell),
226 Err(_) => (vec![], false),
227 },
228 None => (vec![], true),
229 };
230
231 let last_id: Option<CheckpointId> = {
233 match read_log(&self.ctx.log_path) {
234 Ok(recs) => recs.last().map(|r| r.id.clone()),
235 Err(_) => None,
236 }
237 };
238
239 let id = {
241 let repo = self.ctx.repo.lock().map_err(|_| Error::Tool {
242 name: "checkpoint_save".into(),
243 call_id: None,
244 message: "lock poisoned".into(),
245 })?;
246 repo.snapshot_for_session(&self.ctx.session_id, &message)?
247 };
248
249 let mut via = TouchedVia::Structured;
252 if saw_shell {
253 via = TouchedVia::ShellDiff;
254 if let Some(ref last) = last_id {
255 let repo = self.ctx.repo.lock().map_err(|_| Error::Tool {
256 name: "checkpoint_save".into(),
257 call_id: None,
258 message: "lock poisoned".into(),
259 })?;
260 if let Ok(diff_paths) = repo.changed_paths(last, &id) {
261 let mut set: std::collections::HashSet<String> = paths.drain(..).collect();
262 for p in diff_paths {
263 set.insert(p);
264 }
265 paths = {
266 let mut v: Vec<String> = set.into_iter().collect();
267 v.sort();
268 v
269 };
270 }
271 }
272 }
273
274 {
276 let writer = self.ctx.writer.lock().map_err(|_| Error::Tool {
277 name: "checkpoint_save".into(),
278 call_id: None,
279 message: "writer lock poisoned".into(),
280 })?;
281 let rec = CheckpointRecord {
282 turn,
283 pre: last_id,
284 id: id.clone(),
285 message: Some(message),
286 touched_files: paths,
287 touched_via: via,
288 started_at: 0,
289 finished_at: 0,
290 saved_at: unix_now(),
291 };
292 writer.append(&rec).map_err(|e| Error::Tool {
293 name: "checkpoint_save".into(),
294 call_id: None,
295 message: format!("failed to append checkpoint record: {e}"),
296 })?;
297 }
298
299 Ok(id.0)
300 }
301}
302
303pub fn build_checkpoint_save_tool(
305 repo: Arc<Mutex<ShadowRepo>>,
306 session_id: String,
307 touched_files: Option<Arc<Mutex<TouchedFiles>>>,
308 writer: Arc<Mutex<CheckpointLogWriter>>,
309 turn_index: Arc<std::sync::atomic::AtomicUsize>,
310 log_path: PathBuf,
311) -> CheckpointSave {
312 CheckpointSave::new(CheckpointSaveCtx {
313 repo,
314 session_id,
315 touched_files,
316 writer,
317 turn_index,
318 log_path,
319 })
320}
321
322fn unix_now() -> i64 {
324 std::time::SystemTime::now()
325 .duration_since(std::time::UNIX_EPOCH)
326 .map(|d| d.as_secs() as i64)
327 .unwrap_or(0)
328}
329
330pub fn build_checkpoint_tools(
335 workspace: impl Into<std::path::PathBuf>,
336 session_id: impl Into<String>,
337) -> Option<(CheckpointList, CheckpointDiff, Arc<Mutex<ShadowRepo>>)> {
338 let repo = match ShadowRepo::open(workspace) {
339 Ok(r) => Arc::new(Mutex::new(r)),
340 Err(e) => {
341 tracing::warn!("checkpoint tools unavailable: {e}");
342 return None;
343 }
344 };
345 build_checkpoint_tools_inner(repo, session_id.into())
346}
347
348pub fn build_checkpoint_tools_at(
351 workspace: impl Into<std::path::PathBuf>,
352 shadow_dir: impl Into<std::path::PathBuf>,
353 session_id: impl Into<String>,
354) -> Option<(CheckpointList, CheckpointDiff, Arc<Mutex<ShadowRepo>>)> {
355 let repo = match ShadowRepo::open_at(workspace, shadow_dir) {
356 Ok(r) => Arc::new(Mutex::new(r)),
357 Err(e) => {
358 tracing::warn!("checkpoint tools unavailable: {e}");
359 return None;
360 }
361 };
362 build_checkpoint_tools_inner(repo, session_id.into())
363}
364
365fn build_checkpoint_tools_inner(
366 repo: Arc<Mutex<ShadowRepo>>,
367 session_id: String,
368) -> Option<(CheckpointList, CheckpointDiff, Arc<Mutex<ShadowRepo>>)> {
369 let ctx = CheckpointToolCtx {
370 repo: Arc::clone(&repo),
371 session_id,
372 };
373 Some((
374 CheckpointList::new(ctx.clone()),
375 CheckpointDiff::new(ctx),
376 repo,
377 ))
378}
379
380#[cfg(test)]
383mod tests {
384 use super::*;
385 use std::fs;
386 use std::process::Command;
387 use tempfile::TempDir;
388
389 fn has_git() -> bool {
390 Command::new("git").arg("--version").output().is_ok()
391 }
392
393 struct TestWs {
398 workspace: TempDir,
399 shadow: TempDir,
400 }
401
402 impl TestWs {
403 fn path(&self) -> &std::path::Path {
404 self.workspace.path()
405 }
406 fn shadow_dir(&self) -> std::path::PathBuf {
407 self.shadow.path().join("shadow-git")
408 }
409 }
410
411 fn ws() -> TestWs {
412 TestWs {
413 workspace: tempfile::tempdir().expect("workspace tempdir"),
414 shadow: tempfile::tempdir().expect("shadow tempdir"),
415 }
416 }
417
418 #[test]
419 fn unix_now_is_positive() {
420 let ts = unix_now();
422 assert!(
424 ts > 1_000_000_000,
425 "unix_now() returned suspicious value: {ts}"
426 );
427 }
428
429 #[tokio::test]
430 async fn list_tool_is_deferred() {
431 if !has_git() {
433 return;
434 }
435 let w = ws();
436 let (list, _, _) = build_checkpoint_tools_at(w.path(), w.shadow_dir(), "s").unwrap();
437 assert!(
438 list.is_deferred(),
439 "checkpoint_list must be a deferred tool"
440 );
441 }
442
443 #[tokio::test]
444 async fn diff_tool_is_deferred() {
445 if !has_git() {
447 return;
448 }
449 let w = ws();
450 let (_, diff, _) = build_checkpoint_tools_at(w.path(), w.shadow_dir(), "s").unwrap();
451 assert!(
452 diff.is_deferred(),
453 "checkpoint_diff must be a deferred tool"
454 );
455 }
456
457 #[tokio::test]
458 async fn list_tool_empty_session_returns_sentinel() {
459 if !has_git() {
461 return;
462 }
463 let w = ws();
464 let (list, _, _) = build_checkpoint_tools_at(w.path(), w.shadow_dir(), "s").unwrap();
465 let out = list.execute(json!({})).await.unwrap();
466 assert_eq!(out, "no checkpoints in this session yet");
467 }
468
469 #[tokio::test]
470 async fn list_tool_shows_session_checkpoints() {
471 if !has_git() {
472 return;
473 }
474 let w = ws();
475 fs::write(w.path().join("a.txt"), "hi").unwrap();
476 let (list, _, repo) = build_checkpoint_tools_at(w.path(), w.shadow_dir(), "alpha").unwrap();
477 repo.lock()
478 .unwrap()
479 .snapshot_for_session("alpha", "t0")
480 .unwrap();
481 let out = list.execute(json!({})).await.unwrap();
482 assert!(out.contains("t0"));
483 }
484
485 #[tokio::test]
486 async fn list_tool_only_shows_own_session() {
487 if !has_git() {
488 return;
489 }
490 let w = ws();
491 fs::write(w.path().join("a.txt"), "hi").unwrap();
492 let (list_a, _, repo) =
493 build_checkpoint_tools_at(w.path(), w.shadow_dir(), "alpha").unwrap();
494 repo.lock()
495 .unwrap()
496 .snapshot_for_session("alpha", "tA")
497 .unwrap();
498 repo.lock()
499 .unwrap()
500 .snapshot_for_session("beta", "tB")
501 .unwrap();
502 let out_a = list_a.execute(json!({})).await.unwrap();
503 assert!(out_a.contains("tA"));
504 assert!(
505 !out_a.contains("tB"),
506 "alpha must not see beta's checkpoints"
507 );
508 }
509
510 #[tokio::test]
511 async fn diff_tool_returns_empty_for_no_change() {
512 if !has_git() {
513 return;
514 }
515 let w = ws();
516 fs::write(w.path().join("a.txt"), "x").unwrap();
517 let (_, diff, repo) = build_checkpoint_tools_at(w.path(), w.shadow_dir(), "s").unwrap();
518 let id = repo
519 .lock()
520 .unwrap()
521 .snapshot_for_session("s", "snap")
522 .unwrap();
523 let out = diff.execute(json!({"a": id.0})).await.unwrap();
524 assert_eq!(out, "no differences");
525 }
526
527 #[tokio::test]
530 async fn checkpoint_save_tool_creates_entry() {
531 if !has_git() {
532 return;
533 }
534 let w = ws();
535 let log_path = w.path().join("checkpoints.jsonl");
536 fs::write(w.path().join("a.txt"), "hello").unwrap();
537
538 let repo = Arc::new(Mutex::new(
539 ShadowRepo::open_at(w.path(), w.shadow_dir()).unwrap(),
540 ));
541 let touched = Arc::new(Mutex::new(TouchedFiles::new()));
542 {
543 let mut t = touched.lock().unwrap();
544 t.paths.insert("a.txt".to_string());
545 }
546 let writer = Arc::new(Mutex::new(CheckpointLogWriter::open(&log_path).unwrap()));
547 let turn_index = Arc::new(std::sync::atomic::AtomicUsize::new(3));
548
549 let tool = build_checkpoint_save_tool(
550 repo.clone(),
551 "sess".into(),
552 Some(touched.clone()),
553 writer.clone(),
554 turn_index.clone(),
555 log_path.clone(),
556 );
557
558 let id_str = tool
559 .execute(json!({"message": "my save point"}))
560 .await
561 .unwrap();
562 assert!(!id_str.is_empty(), "should return a checkpoint id");
563
564 let cid = CheckpointId(id_str.clone());
566 let data = repo.lock().unwrap().read_file_at(&cid, "a.txt").unwrap();
567 assert_eq!(data, Some(b"hello".to_vec()));
568
569 let recs = read_log(&log_path).unwrap();
571 assert_eq!(recs.len(), 1);
572 assert_eq!(recs[0].turn, 3);
573 assert_eq!(recs[0].id.0, id_str);
574 assert_eq!(recs[0].message.as_deref(), Some("my save point"));
575 assert_eq!(recs[0].touched_files, vec!["a.txt".to_string()]);
576 assert_eq!(recs[0].touched_via, TouchedVia::Structured);
577 }
578
579 #[tokio::test]
580 async fn diff_tool_missing_a_arg_errors() {
581 if !has_git() {
583 return;
584 }
585 let w = ws();
586 let (_, diff, _) = build_checkpoint_tools_at(w.path(), w.shadow_dir(), "s").unwrap();
587 let res = diff.execute(json!({})).await;
588 assert!(res.is_err(), "diff without 'a' argument must fail");
589 }
590
591 #[tokio::test]
592 async fn checkpoint_save_defaults_message_to_turn_index() {
593 if !has_git() {
595 return;
596 }
597 let w = ws();
598 let log_path = w.path().join("checkpoints.jsonl");
599 fs::write(w.path().join("b.txt"), "world").unwrap();
600
601 let repo = Arc::new(Mutex::new(
602 ShadowRepo::open_at(w.path(), w.shadow_dir()).unwrap(),
603 ));
604 let writer = Arc::new(Mutex::new(CheckpointLogWriter::open(&log_path).unwrap()));
605 let turn_index = Arc::new(std::sync::atomic::AtomicUsize::new(7));
606
607 let tool = build_checkpoint_save_tool(
608 repo.clone(),
609 "sess2".into(),
610 None,
611 writer.clone(),
612 turn_index.clone(),
613 log_path.clone(),
614 );
615
616 tool.execute(json!({})).await.unwrap();
618
619 let recs = read_log(&log_path).unwrap();
620 assert_eq!(recs.len(), 1);
621 assert_eq!(
622 recs[0].message.as_deref(),
623 Some("turn 7"),
624 "default message must be 'turn N' where N is the turn index"
625 );
626 }
627}