use super::{DiffPage, PreparedSnapshot};
use crate::{
cancellation::{AgentCancellation, AgentCancellationHandle},
diff_review::{self, CommentChange},
};
use std::sync::Arc;
use std::{
path::{Path, PathBuf},
thread::{self, JoinHandle},
time::{Duration, Instant},
};
enum Outcome {
Review {
snapshot: Result<Arc<PreparedSnapshot>, String>,
mutation: Option<Result<(), String>>,
fingerprint: Result<(usize, u64), String>,
},
Notification(Result<(usize, u64), String>),
}
#[derive(Default)]
pub(crate) struct DiffWorker {
task: Option<JoinHandle<Outcome>>,
cancel: Option<AgentCancellationHandle>,
refresh_invalidated: bool,
full_refresh_pending: bool,
mutation_in_flight: bool,
cwd: PathBuf,
next_refresh: Option<Instant>,
notification_only: bool,
active: bool,
cache: Arc<std::sync::Mutex<diff_review::SnapshotCache>>,
projections: Arc<std::sync::Mutex<super::projection::ProjectionCache>>,
}
impl DiffWorker {
pub(crate) fn is_pending(&self) -> bool {
self.task.is_some()
}
pub(crate) fn poll(&mut self, page: &mut DiffPage, cwd: &Path, active: bool) -> bool {
let mut changed = false;
if active && page.has_unseen_changes {
page.has_unseen_changes = false;
changed = true;
}
if active != self.active {
self.active = active;
self.next_refresh = None;
}
if std::mem::take(&mut page.full_refresh_requested) {
self.full_refresh_pending = true;
self.next_refresh = None;
if self.task.is_some() {
self.refresh_invalidated = true;
if !self.mutation_in_flight
&& let Some(cancel) = &self.cancel
{
cancel.cancel();
}
}
changed = true;
}
if !active
&& !self.notification_only
&& !page.saving
&& let Some(cancel) = &self.cancel
{
cancel.cancel();
self.refresh_invalidated = true;
}
if self.cwd != cwd {
if let Some(cancel) = &self.cancel {
cancel.cancel();
self.refresh_invalidated = true;
}
self.next_refresh = None;
}
if self.task.as_ref().is_some_and(|task| task.is_finished()) {
let result = self.task.take().expect("finished worker").join();
self.cancel = None;
self.mutation_in_flight = false;
let refresh_invalidated = std::mem::take(&mut self.refresh_invalidated);
if self.cwd == cwd {
match result {
Ok(Outcome::Notification(result)) => {
if !refresh_invalidated {
apply_change_fingerprint(page, result, active);
}
}
Ok(Outcome::Review {
snapshot,
mutation,
fingerprint,
}) => {
if mutation.is_none()
&& !refresh_invalidated
&& snapshot.is_ok()
&& page.dialog.is_none()
{
page.error = None;
}
if let Some(result) = mutation {
page.saving = false;
match result {
Ok(()) => {
page.dialog = None;
page.error = None;
}
Err(error) => page.error = Some(error),
}
}
if !refresh_invalidated {
apply_change_fingerprint(page, fingerprint, active);
match snapshot {
Ok(snapshot) => page.apply_snapshot(snapshot),
Err(error) if active => page.error = Some(error),
Err(_) => {}
}
}
}
Err(_) => {
page.saving = false;
page.error = Some("Diff worker stopped unexpectedly".into());
}
}
changed = true;
}
self.next_refresh =
(!(refresh_invalidated || active && self.notification_only)).then(|| {
Instant::now()
+ if active
&& page
.snapshot
.as_ref()
.is_some_and(|s| s.highlighting_pending)
{
Duration::from_millis(16)
} else {
Duration::from_secs(1)
}
});
}
if self.task.is_some() {
return changed;
}
if self.cwd != cwd {
self.cwd = cwd.to_path_buf();
*page = DiffPage::default();
self.next_refresh = None;
changed = true;
}
if page.pending.is_none() && self.next_refresh.is_some_and(|next| Instant::now() < next) {
return changed;
}
let change: Option<CommentChange> = page.pending.take();
let mutation_requested = change.is_some();
self.notification_only = !active && !mutation_requested;
let notification_only = self.notification_only;
let full_refresh = !notification_only && std::mem::take(&mut self.full_refresh_pending);
if full_refresh {
self.cache = Arc::default();
self.projections = Arc::default();
}
let path = cwd.to_path_buf();
let (token, cancel) = AgentCancellation::default().child_token();
let previous = page.snapshot.clone();
let selected_path = previous
.as_ref()
.and_then(|snapshot| snapshot.files.get(page.selected))
.map(|file| file.path.clone());
let cache = Arc::clone(&self.cache);
let projections = Arc::clone(&self.projections);
match thread::Builder::new()
.name("magi-diff-review".into())
.spawn(move || {
if notification_only {
return Outcome::Notification(diff_review::change_fingerprint(&path, &token));
}
let mutation = change.map(|change| {
if token.is_canceled() {
Err("Comment update canceled".into())
} else {
diff_review::save_comment(&path, change, &token)
}
});
let snapshot = if token.is_canceled() {
Err("Diff refresh canceled".into())
} else {
cache
.lock()
.map_err(|_| "Diff cache lock poisoned".to_owned())
.and_then(|mut cache| {
load_and_prepare_snapshot(&path, &token, &mut cache, |source| {
let mut projections = projections.lock().map_err(|_| {
"Diff projection cache lock poisoned".to_owned()
})?;
PreparedSnapshot::prepare(
source,
if full_refresh {
None
} else {
previous.as_deref()
},
selected_path.as_deref(),
&token,
&mut projections,
)
.map_err(|error| error.to_string())
})
})
};
let (snapshot, fingerprint) = match snapshot {
Ok((snapshot, fingerprint)) => (Ok(snapshot), Ok(fingerprint)),
Err(error) => (Err(error.clone()), Err(error)),
};
Outcome::Review {
snapshot,
mutation,
fingerprint,
}
}) {
Ok(task) => {
self.task = Some(task);
self.cancel = Some(cancel);
self.mutation_in_flight = mutation_requested;
}
Err(error) => {
self.full_refresh_pending |= full_refresh;
page.saving = false;
page.error = Some(format!("Cannot start diff worker: {error}"));
self.next_refresh = Some(Instant::now() + Duration::from_secs(1));
changed = true;
}
}
changed || mutation_requested
}
}
fn load_and_prepare_snapshot(
path: &Path,
token: &AgentCancellation,
cache: &mut diff_review::SnapshotCache,
prepare: impl FnOnce(diff_review::ReviewSnapshot) -> Result<PreparedSnapshot, String>,
) -> Result<(Arc<PreparedSnapshot>, (usize, u64)), String> {
let source = diff_review::load_snapshot(path, token, cache)?;
let fingerprint = cache.fingerprint;
token.check().map_err(|error| error.to_string())?;
let snapshot = prepare(source)?;
token.check().map_err(|error| error.to_string())?;
Ok((Arc::new(snapshot), fingerprint))
}
fn apply_change_fingerprint(
page: &mut DiffPage,
result: Result<(usize, u64), String>,
active: bool,
) {
if let Ok((count, fingerprint)) = result {
let changed = page
.last_change_fingerprint
.map_or(count > 0, |previous| previous != fingerprint);
page.has_unseen_changes = !active && (page.has_unseen_changes || changed);
page.last_change_fingerprint = Some(fingerprint);
}
}
impl Drop for DiffWorker {
fn drop(&mut self) {
if let Some(cancel) = &self.cancel {
cancel.cancel();
}
if self.task.as_ref().is_some_and(|task| task.is_finished()) {
let _ = self.task.take().expect("finished worker").join();
}
}
}