agent: Create checkpoints when editing a past message (#30831)

Max Brunsfeld created 6 months ago

Release Notes:

- N/A

Change summary

crates/agent/src/active_thread.rs |  7 ++++++-
crates/agent/src/thread.rs        | 12 +++++++++++-
crates/project/src/git_store.rs   |  2 +-
3 files changed, 18 insertions(+), 3 deletions(-)

Detailed changes

crates/agent/src/active_thread.rs

@@ -1542,11 +1542,15 @@ impl ActiveThread {
         let project = self.thread.read(cx).project().clone();
         let prompt_store = self.thread_store.read(cx).prompt_store().clone();
 
+        let git_store = project.read(cx).git_store().clone();
+        let checkpoint = git_store.update(cx, |git_store, cx| git_store.checkpoint(cx));
+
         let load_context_task =
             crate::context::load_context(new_context, &project, &prompt_store, cx);
         self._load_edited_message_context_task =
             Some(cx.spawn_in(window, async move |this, cx| {
-                let context = load_context_task.await;
+                let (context, checkpoint) =
+                    futures::future::join(load_context_task, checkpoint).await;
                 let _ = this
                     .update_in(cx, |this, window, cx| {
                         this.thread.update(cx, |thread, cx| {
@@ -1555,6 +1559,7 @@ impl ActiveThread {
                                 Role::User,
                                 vec![MessageSegment::Text(edited_text)],
                                 Some(context.loaded_context),
+                                checkpoint.ok(),
                                 cx,
                             );
                             for message_id in this.messages_after(message_id) {

crates/agent/src/thread.rs

@@ -214,7 +214,7 @@ pub struct GitState {
     pub diff: Option<String>,
 }
 
-#[derive(Clone)]
+#[derive(Clone, Debug)]
 pub struct ThreadCheckpoint {
     message_id: MessageId,
     git_checkpoint: GitStoreCheckpoint,
@@ -996,6 +996,7 @@ impl Thread {
         new_role: Role,
         new_segments: Vec<MessageSegment>,
         loaded_context: Option<LoadedContext>,
+        checkpoint: Option<GitStoreCheckpoint>,
         cx: &mut Context<Self>,
     ) -> bool {
         let Some(message) = self.messages.iter_mut().find(|message| message.id == id) else {
@@ -1006,6 +1007,15 @@ impl Thread {
         if let Some(context) = loaded_context {
             message.loaded_context = context;
         }
+        if let Some(git_checkpoint) = checkpoint {
+            self.checkpoints_by_message.insert(
+                id,
+                ThreadCheckpoint {
+                    message_id: id,
+                    git_checkpoint,
+                },
+            );
+        }
         self.touch_updated_at();
         cx.emit(ThreadEvent::MessageEdited(id));
         true

crates/project/src/git_store.rs

@@ -163,7 +163,7 @@ struct LocalDownstreamState {
     _task: Task<Result<()>>,
 }
 
-#[derive(Clone)]
+#[derive(Clone, Debug)]
 pub struct GitStoreCheckpoint {
     checkpoints_by_work_dir_abs_path: HashMap<Arc<Path>, GitRepositoryCheckpoint>,
 }