use std::collections::{BTreeMap, BTreeSet};
use std::ffi::OsStr;
use std::io::{self, ErrorKind, Write as _};
use std::path::{Component, Path, PathBuf};
use std::process::{Command, Stdio};
use std::sync::{Arc, Mutex, PoisonError};
use std::time::{Duration, SystemTime};
use rig_core::completion::UnsupportedOption;
use rig_core::effect::{EffectKind, HandlerDescriptor, HandlerKey};
use rig_core::serve::{Dispatch, ErasedHandler, Reply, Serve};
use rig_core::tool::{
ContentHash, FileChange, FileState, FilesChanged, PublishedContext, Touches, content_hash,
};
use crate::fs::Previous;
use crate::process::Launcher;
use crate::workspace::{Listing, Stat, Workspace};
use crate::{MAX_FILE_BYTES, blocking};
#[derive(Clone, Debug)]
pub struct Snapshot(Taken);
#[derive(Clone, Debug)]
enum Taken {
Git { git: Git, tree: String },
Scan {
files: Arc<BTreeMap<PathBuf, Seen>>,
began: SystemTime,
blobs: Option<PathBuf>,
},
}
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct Taking {
pub keep: Option<PathBuf>,
pub also: Vec<PathBuf>,
pub launcher: Launcher,
}
#[derive(Clone, Debug, PartialEq, Eq)]
struct Seen {
size: u64,
modified: Option<SystemTime>,
hash: Option<ContentHash>,
lines: usize,
}
impl Seen {
fn state(&self) -> FileState {
FileState {
bytes: self.size,
lines: self.hash.is_some().then_some(self.lines),
hash: self.hash.clone(),
}
}
}
pub fn snapshot(
workspace: &dyn Workspace,
root: &Path,
taking: &Taking,
earlier: Option<&Snapshot>,
) -> io::Result<Snapshot> {
let keep = taking.keep.as_deref();
let git = match earlier.map(|it| &it.0) {
Some(Taken::Git { git, .. }) => Some(git.clone()),
Some(Taken::Scan { .. }) => None,
None => keep
.filter(|_| workspace.host().is_none())
.and_then(|keep| Git::find(root, keep, &taking.launcher)),
};
let also: Vec<PathBuf> = taking
.also
.iter()
.map(|path| root.join(path))
.filter(|path| {
path.starts_with(root) && !path.components().any(|part| part == Component::ParentDir)
})
.collect();
match git {
Some(git) => {
let tree = git.write_tree(&also)?;
Ok(Snapshot(Taken::Git { git, tree }))
}
None => {
let before = match earlier.map(|it| &it.0) {
Some(Taken::Scan { files, began, .. }) => Some((&**files, *began)),
_ => None,
};
let blobs = keep.map(|keep| private(keep, root).join("blobs"));
let began = SystemTime::now();
let files = scan(workspace, root, &also, before, blobs.as_deref())?;
let files = Arc::new(files);
Ok(Snapshot(Taken::Scan {
files,
began,
blobs,
}))
}
}
}
fn private(keep: &Path, dir: &Path) -> PathBuf {
let name = content_hash(dir.as_os_str().as_encoded_bytes()).0;
keep.join("checkpoints")
.join(name.get(..16).unwrap_or(&name))
}
pub fn forget_kept(keep: &Path) -> io::Result<()> {
let Ok(dirs) = std::fs::read_dir(keep.join("checkpoints")) else {
return Ok(());
};
for dir in dirs.flatten() {
match std::fs::remove_dir_all(dir.path().join("blobs")) {
Err(error) if error.kind() != ErrorKind::NotFound => return Err(error),
_ => {}
}
}
Ok(())
}
impl Snapshot {
pub fn changes_since(&self, earlier: &Snapshot) -> io::Result<Vec<FileChange>> {
match (&earlier.0, &self.0) {
(Taken::Git { tree: before, .. }, Taken::Git { git, tree }) => {
git.changes(before, tree)
}
(Taken::Scan { files: before, .. }, Taken::Scan { files, blobs, .. }) => {
Ok(scan_changes(before, files, blobs.as_deref()))
}
_ => Err(io::Error::new(
ErrorKind::InvalidInput,
"the two snapshots were taken different ways",
)),
}
}
pub fn content(&self, path: &Path) -> Option<io::Result<Vec<u8>>> {
match &self.0 {
Taken::Git { git, tree } => Some(git.content(tree, path)),
Taken::Scan { files, blobs, .. } => {
let blobs = blobs.as_deref()?;
Some(match files.get(path).map(|seen| &seen.hash) {
Some(Some(hash)) => std::fs::read(blobs.join(&hash.0)),
Some(None) => Err(io::Error::new(
ErrorKind::FileTooLarge,
"the file was too large to keep",
)),
None => Err(ErrorKind::NotFound.into()),
})
}
}
}
}
fn scan(
workspace: &dyn Workspace,
root: &Path,
also: &[PathBuf],
earlier: Option<(&BTreeMap<PathBuf, Seen>, SystemTime)>,
blobs: Option<&Path>,
) -> io::Result<BTreeMap<PathBuf, Seen>> {
if let Some(blobs) = blobs {
std::fs::create_dir_all(blobs)?;
}
let mut listed = workspace.files(root, None, Listing::Visible)?;
for path in also {
listed.extend(
workspace
.files(path, None, Listing::All)
.unwrap_or_default(),
);
}
let mut seen = BTreeMap::new();
for path in listed {
let Ok(Stat::File { size, modified, .. }) = workspace.stat(&path) else {
continue;
};
let same = earlier.and_then(|(files, began)| {
let was = files.get(&path)?;
let settled = modified.is_some_and(|modified| settled_before(modified, began));
(settled && was.size == size && was.modified == modified).then_some(was)
});
let file = match same {
Some(was) => was.clone(),
None => {
let bytes = (size <= MAX_FILE_BYTES)
.then(|| workspace.read(&path, MAX_FILE_BYTES).ok())
.flatten()
.map(|contents| contents.bytes);
let hash = bytes.as_deref().map(content_hash);
if let (Some(blobs), Some(bytes), Some(hash)) = (blobs, &bytes, &hash) {
let blob = blobs.join(&hash.0);
if !blob.exists() {
crate::fs::write_scratch(&blob, bytes).ok();
}
}
Seen {
size,
modified,
hash,
lines: bytes.map_or(0, |bytes| String::from_utf8_lossy(&bytes).lines().count()),
}
}
};
seen.insert(path, file);
}
Ok(seen)
}
fn settled_before(modified: SystemTime, began: SystemTime) -> bool {
let whole_seconds = modified
.duration_since(SystemTime::UNIX_EPOCH)
.is_ok_and(|since| since.subsec_nanos() == 0);
let precision = if whole_seconds {
Duration::from_secs(2)
} else {
Duration::from_millis(100)
};
began
.checked_sub(precision)
.is_some_and(|settled| modified < settled)
}
fn scan_changes(
before: &BTreeMap<PathBuf, Seen>,
after: &BTreeMap<PathBuf, Seen>,
blobs: Option<&Path>,
) -> Vec<FileChange> {
let paths: BTreeSet<&PathBuf> = before.keys().chain(after.keys()).collect();
paths
.into_iter()
.filter_map(|path| {
let shown = path.to_string_lossy();
let (was, now) = (before.get(path), after.get(path));
match (was, now) {
(Some(was), Some(now)) if was.hash.is_some() && was.hash == now.hash => {
return None;
}
(Some(was), Some(now)) if was.hash.is_none() && was == now => return None,
(None, None) => return None,
_ => {}
}
let kept = |seen: Option<&Seen>| match (blobs, seen) {
(_, None) => Some(None),
(
Some(blobs),
Some(Seen {
hash: Some(hash), ..
}),
) => std::fs::read(blobs.join(&hash.0)).ok().map(Some),
_ => None,
};
if let (Some(was_bytes), Some(now_bytes)) = (kept(was), kept(now)) {
return Some(change(&shown, was_bytes, now_bytes.as_deref()));
}
let mut change = match (was, now) {
(None, Some(now)) => FileChange::created(shown, now.lines),
(Some(was), None) => FileChange::deleted(shown, was.lines),
_ => FileChange::modified(shown, 0, 0),
};
change.before = was.map(Seen::state);
change.after = now.map(Seen::state);
Some(change)
})
.collect()
}
#[derive(Clone, Debug)]
struct Git {
top: PathBuf,
dir: PathBuf,
objects: PathBuf,
root: PathBuf,
below: PathBuf,
keep: PathBuf,
launcher: Launcher,
}
impl Git {
fn find(root: &Path, keep: &Path, launcher: &Launcher) -> Option<Self> {
let git = |args: &[&OsStr]| {
launcher
.command("git")
.arg("-C")
.arg(root)
.args(args)
.env_remove("GIT_DIR")
.env_remove("GIT_WORK_TREE")
.env_remove("GIT_INDEX_FILE")
.env_remove("GIT_OBJECT_DIRECTORY")
.stdin(Stdio::null())
.stderr(Stdio::null())
.output()
.ok()
};
let output = git(&[
"rev-parse",
"--show-toplevel",
"--absolute-git-dir",
"--git-path",
"objects",
]
.map(OsStr::new))
.filter(|output| output.status.success())?;
let text = String::from_utf8(output.stdout).ok()?;
let mut lines = text.lines();
let (top, dir, objects) = (lines.next()?, lines.next()?, lines.next()?);
let top = PathBuf::from(top);
let canonical = std::fs::canonicalize(root).ok()?;
let below = canonical.strip_prefix(&top).ok()?.to_path_buf();
let check = ["check-ignore", "-q", "--"].map(OsStr::new);
if !below.as_os_str().is_empty()
&& git(&[check.as_slice(), &[canonical.as_os_str()]].concat())
.is_some_and(|output| output.status.success())
{
return None;
}
let keep = private(keep, &top);
Some(Self {
objects: canonical.join(objects),
root: root.to_path_buf(),
top,
dir: PathBuf::from(dir),
below,
keep,
launcher: launcher.clone(),
})
}
fn command(&self, args: &[&OsStr]) -> Command {
let mut command = self.launcher.command("git");
command
.arg("--git-dir")
.arg(&self.dir)
.arg("--work-tree")
.arg(&self.top)
.args([
"-c",
"core.fsmonitor=false",
"-c",
"core.untrackedCache=false",
])
.args(args)
.current_dir(&self.top)
.env("GIT_INDEX_FILE", self.keep.join("index"))
.env("GIT_OBJECT_DIRECTORY", self.keep.join("objects"))
.env("GIT_ALTERNATE_OBJECT_DIRECTORIES", &self.objects)
.env("GIT_OPTIONAL_LOCKS", "0")
.env_remove("GIT_DIR")
.env_remove("GIT_WORK_TREE")
.stdin(Stdio::null());
command
}
fn run(&self, args: &[&OsStr]) -> io::Result<Vec<u8>> {
let output = self.command(args).output()?;
match output.status.success() {
true => Ok(output.stdout),
false => Err(io::Error::other(format!(
"git {} failed: {}",
args.first()
.map_or(std::borrow::Cow::Borrowed(""), |it| it.to_string_lossy()),
String::from_utf8_lossy(&output.stderr).trim()
))),
}
}
fn write_tree(&self, also: &[PathBuf]) -> io::Result<String> {
static INDEX: Mutex<()> = Mutex::new(());
let _index = INDEX.lock().unwrap_or_else(PoisonError::into_inner);
std::fs::create_dir_all(self.keep.join("objects"))?;
let below = match self.below.as_os_str().is_empty() {
true => OsStr::new("."),
false => self.below.as_os_str(),
};
let add = ["add", "--all", "--ignore-errors", "--"].map(OsStr::new);
self.run(&[add.as_slice(), &[below]].concat())?;
let also: Vec<PathBuf> = also
.iter()
.filter(|path| path.exists())
.filter_map(|path| Some(self.below.join(path.strip_prefix(&self.root).ok()?)))
.collect();
if !also.is_empty() {
let force = ["add", "--all", "--force", "--ignore-errors", "--"].map(OsStr::new);
let also = also.iter().map(|path| path.as_os_str());
self.run(&force.into_iter().chain(also).collect::<Vec<_>>())?;
}
let tree = self.run(&[OsStr::new("write-tree")])?;
Ok(String::from_utf8_lossy(&tree).trim().to_owned())
}
fn blobs(&self, ids: &[&str]) -> io::Result<Vec<Vec<u8>>> {
if ids.is_empty() {
return Ok(Vec::new());
}
let mut child = self
.command(&["cat-file", "--batch=%(objectsize)"].map(OsStr::new))
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::null())
.spawn()?;
let input: String = ids.iter().map(|id| format!("{id}\n")).collect();
let stdin = child.stdin.take();
let output = std::thread::scope(|scope| {
scope.spawn(|| stdin.map(|mut stdin| stdin.write_all(input.as_bytes())));
child.wait_with_output()
})?;
let mut rest = output.stdout.as_slice();
let mut blobs = Vec::with_capacity(ids.len());
for id in ids {
let missing = || io::Error::new(ErrorKind::NotFound, format!("git has no {id}"));
let end = rest.iter().position(|byte| *byte == b'\n');
let header = rest.get(..end.ok_or_else(missing)?).unwrap_or_default();
let size: usize = String::from_utf8_lossy(header)
.parse()
.map_err(|_| missing())?;
let body = rest.get(header.len() + 1..).unwrap_or_default();
blobs.push(body.get(..size).ok_or_else(missing)?.to_vec());
rest = body.get(size + 1..).unwrap_or_default();
}
Ok(blobs)
}
fn content(&self, tree: &str, path: &Path) -> io::Result<Vec<u8>> {
let below = match path.strip_prefix(&self.root) {
Ok(rest) => self.below.join(rest),
Err(_) => path
.strip_prefix(&self.top)
.map_err(|_| {
io::Error::new(ErrorKind::InvalidInput, "the path is outside the work tree")
})?
.to_path_buf(),
};
let spec = format!("{tree}:{}", below.to_string_lossy());
self.run(&["cat-file", "blob", &spec].map(OsStr::new))
}
fn shown(&self, path: &str) -> String {
match Path::new(path).strip_prefix(&self.below) {
Ok(rest) => self.root.join(rest),
Err(_) => self.top.join(path),
}
.to_string_lossy()
.into_owned()
}
fn changes(&self, before: &str, after: &str) -> io::Result<Vec<FileChange>> {
let raw = self.run(
&[
"diff-tree",
"-r",
"-z",
"--no-renames",
"--raw",
before,
after,
]
.map(OsStr::new),
)?;
let text = String::from_utf8_lossy(&raw);
let mut fields = text.split('\0');
let mut changed = Vec::new();
while let (Some(header), Some(path)) = (fields.next(), fields.next()) {
let mut parts = header.split(' ').skip(2);
if let (Some(old), Some(new)) = (parts.next(), parts.next())
&& old != new
{
changed.push((path, old, new));
}
}
let present = |id: &&str| id.bytes().any(|byte| byte != b'0');
let sides = changed.iter().flat_map(|(_, old, new)| [*old, *new]);
let mut blobs = self
.blobs(&sides.filter(present).collect::<Vec<_>>())?
.into_iter();
let mut side = |id: &str| present(&id).then(|| blobs.next()).flatten();
let changes = changed.into_iter().map(|(path, old, new)| {
let (was, now) = (side(old), side(new));
change(&self.shown(path), was, now.as_deref())
});
Ok(changes.collect())
}
}
fn change(path: &str, before: Option<Vec<u8>>, after: Option<&[u8]>) -> FileChange {
let small =
|bytes: Option<&[u8]>| bytes.is_none_or(|bytes| bytes.len() as u64 <= MAX_FILE_BYTES);
if small(before.as_deref()) && small(after) {
return Previous::change(before.map(Previous::of), path, after);
}
let was = before.as_deref().map(FileState::of);
let now = after.map(FileState::of);
let lines = |state: &Option<FileState>| state.as_ref().and_then(|it| it.lines).unwrap_or(0);
let mut change = match (&was, &now) {
(_, None) => FileChange::deleted(path, lines(&was)),
(None, Some(_)) => FileChange::created(path, lines(&now)),
(Some(_), Some(_)) => FileChange::modified(path, 0, 0),
};
change.before = was;
change.after = now;
change
}
#[derive(Clone, Debug)]
pub struct Checkpointed {
pub inner: ErasedHandler,
pub workspace: Arc<dyn Workspace>,
pub root: PathBuf,
pub taking: Taking,
pub newest: Arc<Mutex<Option<Snapshot>>>,
pub first: Arc<Mutex<Option<Snapshot>>>,
pub before: Arc<Mutex<Option<Snapshot>>>,
}
impl Checkpointed {
async fn take(&self, earlier: Option<Snapshot>) -> Option<(Option<Snapshot>, Snapshot)> {
let (workspace, root, taking) = (
self.workspace.clone(),
self.root.clone(),
self.taking.clone(),
);
let taken = blocking(move || {
let now = snapshot(&*workspace, &root, &taking, earlier.as_ref())
.map_err(|error| rig_core::tool::ToolExecutionError::other(error.to_string()))?;
Ok((earlier, now))
})
.await;
taken
.inspect_err(|error| {
tracing::warn!(root = %self.root.display(), "no checkpoint of the call: {error}");
})
.ok()
}
}
impl Serve for Checkpointed {
type Family = rig_core::effect::family::Dynamic;
fn descriptor(&self) -> HandlerDescriptor {
self.inner.descriptor()
}
fn key(&self) -> HandlerKey {
self.inner.key().clone()
}
async fn serve(&self, kind: EffectKind, dispatch: Dispatch) -> Reply {
let published = dispatch.scope::<PublishedContext>();
let newest = self
.newest
.lock()
.unwrap_or_else(PoisonError::into_inner)
.clone();
let before = self.take(newest).await.map(|(_, now)| now);
if let Some(before) = &before {
self.first
.lock()
.unwrap_or_else(PoisonError::into_inner)
.get_or_insert_with(|| before.clone());
*self.newest.lock().unwrap_or_else(PoisonError::into_inner) = Some(before.clone());
*self.before.lock().unwrap_or_else(PoisonError::into_inner) = Some(before.clone());
}
let reply = self.inner.handle(kind, dispatch).await;
if !matches!(reply, Reply::Outcome(_)) {
return reply;
}
let Some((Some(before), after)) = self.take(before).await else {
return reply;
};
let changes = after
.changes_since(&before)
.inspect_err(|error| {
tracing::warn!(root = %self.root.display(), "no changes of the call: {error}");
})
.unwrap_or_default();
*self.newest.lock().unwrap_or_else(PoisonError::into_inner) = Some(after);
if let Some(published) = published {
let own = published
.result_context()
.and_then(|context| context.get::<FilesChanged>().ok().flatten());
if own.is_none() && !changes.is_empty() {
published.insert(FilesChanged(changes)).ok();
}
}
reply
}
fn footprint(&self, kind: &EffectKind) -> Option<Touches> {
self.inner.footprint(kind)
}
fn check(&self, kind: &EffectKind) -> Vec<UnsupportedOption> {
self.inner.check(kind)
}
}
#[cfg(test)]
mod tests;