3501 lines · 126.6 KB · Rust
RawDownload
use anyhow::{Context as _, Result};use buffer_diff::BufferDiff;use clock;use collections::{BTreeMap, HashMap};use fs::MTime;use futures::{FutureExt, StreamExt, channel::mpsc};use gpui::{    App, AppContext, AsyncApp, Context, Entity, SharedString, Subscription, Task, WeakEntity,};use language::{Anchor, Buffer, BufferEvent, Point, ToOffset, ToPoint};use project::{Project, ProjectItem, lsp_store::OpenLspBufferHandle};use std::{    cmp,    ops::Range,    path::{Path, PathBuf},    sync::Arc,};use text::{Edit, Patch, Rope};use util::{RangeExt, ResultExt as _}; /// Stores undo information for a single buffer's rejected edits#[derive(Clone)]pub struct PerBufferUndo {    pub buffer: WeakEntity<Buffer>,    pub edits_to_restore: Vec<(Range<Anchor>, String)>,    pub status: UndoBufferStatus,} /// Tracks the buffer status for undo purposes#[derive(Clone, Debug)]pub enum UndoBufferStatus {    Modified,    /// Buffer was created by the agent.    /// - `had_existing_content: true` - Agent overwrote an existing file. On reject, the    ///   original content was restored. Undo is supported: we restore the agent's content.    /// - `had_existing_content: false` - Agent created a new file that didn't exist before.    ///   On reject, the file was deleted. Undo is NOT currently supported (would require    ///   recreating the file). Future TODO.    Created {        had_existing_content: bool,    },} /// Stores undo information for the most recent reject operation#[derive(Clone)]pub struct LastRejectUndo {    /// Per-buffer undo information    pub buffers: Vec<PerBufferUndo>,} /// Tracks actions performed by tools in a threadpub struct ActionLog {    /// Buffers that we want to notify the model about when they change.    tracked_buffers: BTreeMap<Entity<Buffer>, TrackedBuffer>,    /// The project this action log is associated with    project: Entity<Project>,    /// An action log to forward all public methods to    /// Useful in cases like subagents, where we want to track individual diffs for this subagent,    /// but also want to associate the reads/writes with a parent review experience    linked_action_log: Option<Entity<ActionLog>>,    /// Stores undo information for the most recent reject operation    last_reject_undo: Option<LastRejectUndo>,    /// Tracks the last time files were read by the agent, to detect external modifications    file_read_times: HashMap<PathBuf, MTime>,} impl ActionLog {    /// Creates a new, empty action log associated with the given project.    pub fn new(project: Entity<Project>) -> Self {        Self {            tracked_buffers: BTreeMap::default(),            project,            linked_action_log: None,            last_reject_undo: None,            file_read_times: HashMap::default(),        }    }     pub fn with_linked_action_log(mut self, linked_action_log: Entity<ActionLog>) -> Self {        self.linked_action_log = Some(linked_action_log);        self    }     pub fn project(&self) -> &Entity<Project> {        &self.project    }     pub fn file_read_time(&self, path: &Path) -> Option<MTime> {        self.file_read_times.get(path).copied()    }     fn update_file_read_time(&mut self, buffer: &Entity<Buffer>, cx: &App) {        let buffer = buffer.read(cx);        if let Some(file) = buffer.file() {            if let Some(local_file) = file.as_local() {                if let Some(mtime) = file.disk_state().mtime() {                    let abs_path = local_file.abs_path(cx);                    self.file_read_times.insert(abs_path, mtime);                }            }        }    }     fn remove_file_read_time(&mut self, buffer: &Entity<Buffer>, cx: &App) {        let buffer = buffer.read(cx);        if let Some(file) = buffer.file() {            if let Some(local_file) = file.as_local() {                let abs_path = local_file.abs_path(cx);                self.file_read_times.remove(&abs_path);            }        }    }     fn track_buffer_internal(        &mut self,        buffer: Entity<Buffer>,        is_created: bool,        cx: &mut Context<Self>,    ) -> &mut TrackedBuffer {        let status = if is_created {            if let Some(tracked) = self.tracked_buffers.remove(&buffer) {                match tracked.status {                    TrackedBufferStatus::Created {                        existing_file_content,                    } => TrackedBufferStatus::Created {                        existing_file_content,                    },                    TrackedBufferStatus::Modified | TrackedBufferStatus::Deleted => {                        TrackedBufferStatus::Created {                            existing_file_content: Some(tracked.diff_base),                        }                    }                }            } else if buffer                .read(cx)                .file()                .is_some_and(|file| file.disk_state().exists())            {                TrackedBufferStatus::Created {                    existing_file_content: Some(buffer.read(cx).as_rope().clone()),                }            } else {                TrackedBufferStatus::Created {                    existing_file_content: None,                }            }        } else {            TrackedBufferStatus::Modified        };         let tracked_buffer = self            .tracked_buffers            .entry(buffer.clone())            .or_insert_with(|| {                let open_lsp_handle = self.project.update(cx, |project, cx| {                    project.register_buffer_with_language_servers(&buffer, cx)                });                 let text_snapshot = buffer.read(cx).text_snapshot();                let language = buffer.read(cx).language().cloned();                let language_registry = buffer.read(cx).language_registry();                let diff = cx.new(|cx| {                    BufferDiff::new(                        &text_snapshot,                        language,                        language_registry,                        buffer_diff::DiffBaseKind::Custom,                        cx,                    )                });                let (diff_update_tx, diff_update_rx) = mpsc::unbounded();                let diff_base;                let unreviewed_edits;                if is_created {                    diff_base = Rope::default();                    unreviewed_edits = Patch::new(vec![Edit {                        old: 0..1,                        new: 0..text_snapshot.max_point().row + 1,                    }])                } else {                    diff_base = buffer.read(cx).as_rope().clone();                    unreviewed_edits = Patch::default();                }                TrackedBuffer {                    buffer: buffer.clone(),                    diff_base,                    unreviewed_edits,                    snapshot: text_snapshot,                    status,                    version: buffer.read(cx).version(),                    diff,                    diff_update: diff_update_tx,                    _open_lsp_handle: open_lsp_handle,                    _maintain_diff: cx.spawn({                        let buffer = buffer.clone();                        async move |this, cx| {                            Self::maintain_diff(this, buffer, diff_update_rx, cx)                                .await                                .ok();                        }                    }),                    _subscription: cx.subscribe(&buffer, Self::handle_buffer_event),                }            });        tracked_buffer.version = buffer.read(cx).version();        tracked_buffer    }     fn handle_buffer_event(        &mut self,        buffer: Entity<Buffer>,        event: &BufferEvent,        cx: &mut Context<Self>,    ) {        match event {            BufferEvent::Edited { .. } => {                let Some(tracked_buffer) = self.tracked_buffers.get_mut(&buffer) else {                    return;                };                let buffer_version = buffer.read(cx).version();                if !buffer_version.changed_since(&tracked_buffer.version) {                    return;                }                self.handle_buffer_edited(buffer, cx);            }            BufferEvent::FileHandleChanged => {                self.handle_buffer_file_changed(buffer, cx);            }            _ => {}        };    }     fn handle_buffer_edited(&mut self, buffer: Entity<Buffer>, cx: &mut Context<Self>) {        let Some(tracked_buffer) = self.tracked_buffers.get_mut(&buffer) else {            return;        };        tracked_buffer.schedule_diff_update(ChangeAuthor::User, cx);    }     fn handle_buffer_file_changed(&mut self, buffer: Entity<Buffer>, cx: &mut Context<Self>) {        let Some(tracked_buffer) = self.tracked_buffers.get_mut(&buffer) else {            return;        };         match tracked_buffer.status {            TrackedBufferStatus::Created { .. } | TrackedBufferStatus::Modified => {                if buffer                    .read(cx)                    .file()                    .is_some_and(|file| file.disk_state().is_deleted())                {                    // If the buffer had been edited by a tool, but it got                    // deleted externally, we want to stop tracking it.                    self.tracked_buffers.remove(&buffer);                }                cx.notify();            }            TrackedBufferStatus::Deleted => {                if buffer                    .read(cx)                    .file()                    .is_some_and(|file| !file.disk_state().is_deleted())                {                    // If the buffer had been deleted by a tool, but it got                    // resurrected externally, we want to clear the edits we                    // were tracking and reset the buffer's state.                    self.tracked_buffers.remove(&buffer);                    self.track_buffer_internal(buffer, false, cx);                }                cx.notify();            }        }    }     async fn maintain_diff(        this: WeakEntity<Self>,        buffer: Entity<Buffer>,        mut buffer_updates: mpsc::UnboundedReceiver<(ChangeAuthor, text::BufferSnapshot)>,        cx: &mut AsyncApp,    ) -> Result<()> {        let git_diff = this            .update(cx, |this, cx| {                this.project.update(cx, |project, cx| {                    project.open_uncommitted_diff(buffer.clone(), cx)                })            })?            .await            .ok();        let (mut git_diff_updates_tx, mut git_diff_updates_rx) = watch::channel(());        let _diff_subscription = if let Some(git_diff) = git_diff.as_ref() {            cx.update(|cx| {                Some(cx.subscribe(git_diff, move |_, event, _cx| {                    if matches!(event, buffer_diff::BufferDiffEvent::BaseTextChanged) {                        git_diff_updates_tx.send(()).ok();                    }                }))            })        } else {            None        };         loop {            futures::select_biased! {                buffer_update = buffer_updates.next() => {                    if let Some((author, buffer_snapshot)) = buffer_update {                        Self::track_edits(&this, &buffer, author, buffer_snapshot, cx).await?;                    } else {                        break;                    }                }                _ = git_diff_updates_rx.changed().fuse() => {                    if let Some(git_diff) = git_diff.as_ref() {                        Self::keep_committed_edits(&this, &buffer, git_diff, cx).await?;                    }                }            }        }         Ok(())    }     async fn track_edits(        this: &WeakEntity<ActionLog>,        buffer: &Entity<Buffer>,        author: ChangeAuthor,        buffer_snapshot: text::BufferSnapshot,        cx: &mut AsyncApp,    ) -> Result<()> {        let rebase = this.update(cx, |this, cx| {            let tracked_buffer = this                .tracked_buffers                .get_mut(buffer)                .context("buffer not tracked")?;             let rebase = cx.background_spawn({                let mut base_text = tracked_buffer.diff_base.clone();                let old_snapshot = tracked_buffer.snapshot.clone();                let new_snapshot = buffer_snapshot.clone();                let unreviewed_edits = tracked_buffer.unreviewed_edits.clone();                let edits = diff_snapshots(&old_snapshot, &new_snapshot);                async move {                    if let ChangeAuthor::User = author {                        apply_non_conflicting_edits(                            &unreviewed_edits,                            edits,                            &mut base_text,                            new_snapshot.as_rope(),                        );                    }                     (Arc::from(base_text.to_string().as_str()), base_text)                }            });             anyhow::Ok(rebase)        })??;        let (new_base_text, new_diff_base) = rebase.await;         Self::update_diff(            this,            buffer,            buffer_snapshot,            new_base_text,            new_diff_base,            cx,        )        .await    }     async fn keep_committed_edits(        this: &WeakEntity<ActionLog>,        buffer: &Entity<Buffer>,        git_diff: &Entity<BufferDiff>,        cx: &mut AsyncApp,    ) -> Result<()> {        let buffer_snapshot = this.read_with(cx, |this, _cx| {            let tracked_buffer = this                .tracked_buffers                .get(buffer)                .context("buffer not tracked")?;            anyhow::Ok(tracked_buffer.snapshot.clone())        })??;        let (new_base_text, new_diff_base) = this            .read_with(cx, |this, cx| {                let tracked_buffer = this                    .tracked_buffers                    .get(buffer)                    .context("buffer not tracked")?;                let old_unreviewed_edits = tracked_buffer.unreviewed_edits.clone();                let agent_diff_base = tracked_buffer.diff_base.clone();                let git_diff_base = git_diff.read(cx).base_text(cx).as_rope().clone();                let buffer_text = tracked_buffer.snapshot.as_rope().clone();                anyhow::Ok(cx.background_spawn(async move {                    if buffer_text.len() == git_diff_base.len()                        && buffer_text.chars_at(0).eq(git_diff_base.chars_at(0))                    {                        return (Arc::<str>::from(git_diff_base.to_string()), git_diff_base);                    }                    let mut old_unreviewed_edits = old_unreviewed_edits.into_iter().peekable();                    let committed_edits = language::line_diff(                        &agent_diff_base.to_string(),                        &git_diff_base.to_string(),                    )                    .into_iter()                    .map(|(old, new)| Edit { old, new });                     let mut new_agent_diff_base = agent_diff_base.clone();                    let mut row_delta = 0i32;                    for committed in committed_edits {                        while let Some(unreviewed) = old_unreviewed_edits.peek() {                            // If the committed edit matches the unreviewed                            // edit, assume the user wants to keep it.                            if committed.old == unreviewed.old {                                let unreviewed_new =                                    buffer_text.slice_rows(unreviewed.new.clone()).to_string();                                let committed_new =                                    git_diff_base.slice_rows(committed.new.clone()).to_string();                                if unreviewed_new == committed_new {                                    let old_byte_start =                                        new_agent_diff_base.point_to_offset(Point::new(                                            (unreviewed.old.start as i32 + row_delta) as u32,                                            0,                                        ));                                    let old_byte_end =                                        new_agent_diff_base.point_to_offset(cmp::min(                                            Point::new(                                                (unreviewed.old.end as i32 + row_delta) as u32,                                                0,                                            ),                                            new_agent_diff_base.max_point(),                                        ));                                    new_agent_diff_base                                        .replace(old_byte_start..old_byte_end, &unreviewed_new);                                    row_delta +=                                        unreviewed.new_len() as i32 - unreviewed.old_len() as i32;                                }                            } else if unreviewed.old.start >= committed.old.end {                                break;                            }                             old_unreviewed_edits.next().unwrap();                        }                    }                     (                        Arc::from(new_agent_diff_base.to_string().as_str()),                        new_agent_diff_base,                    )                }))            })??            .await;         Self::update_diff(            this,            buffer,            buffer_snapshot,            new_base_text,            new_diff_base,            cx,        )        .await    }     async fn update_diff(        this: &WeakEntity<ActionLog>,        buffer: &Entity<Buffer>,        buffer_snapshot: text::BufferSnapshot,        new_base_text: Arc<str>,        new_diff_base: Rope,        cx: &mut AsyncApp,    ) -> Result<()> {        let diff = this.read_with(cx, |this, _cx| {            let tracked_buffer = this                .tracked_buffers                .get(buffer)                .context("buffer not tracked")?;            anyhow::Ok(tracked_buffer.diff.clone())        })??;        diff.update(cx, |diff, cx| {            diff.set_base_text(Some(new_base_text), buffer_snapshot.clone(), cx)        })        .await;        let diff_snapshot = diff.update(cx, |diff, cx| diff.snapshot(cx));         let unreviewed_edits = cx            .background_spawn({                let buffer_snapshot = buffer_snapshot.clone();                let new_diff_base = new_diff_base.clone();                async move {                    let mut unreviewed_edits = Patch::default();                    for hunk in diff_snapshot.hunks_intersecting_range(                        Anchor::min_for_buffer(buffer_snapshot.remote_id())                            ..Anchor::max_for_buffer(buffer_snapshot.remote_id()),                        &buffer_snapshot,                    ) {                        let old_range = new_diff_base                            .offset_to_point(hunk.diff_base_byte_range.start)                            ..new_diff_base.offset_to_point(hunk.diff_base_byte_range.end);                        let new_range = hunk.range.start..hunk.range.end;                        unreviewed_edits.push(point_to_row_edit(

Lines 1–500 of 3501.