Restore Checkpoint now appears if you press Cancel (#31310)

Richard Feldman created 6 months ago

## Before


https://github.com/user-attachments/assets/0da54afd-78bb-4fee-ab0c-f6ff96f89550

## After


https://github.com/user-attachments/assets/e840e642-714b-4ed7-99cf-a972f50361ba

Release Notes:

- In the Agent Panel, Restore Checkpoint now appears if you press Cancel
during generation.

Change summary

crates/agent/src/thread.rs | 19 +++++++++++++++++--
1 file changed, 17 insertions(+), 2 deletions(-)

Detailed changes

crates/agent/src/thread.rs

@@ -757,6 +757,14 @@ impl Thread {
             return;
         };
 
+        self.finalize_checkpoint(pending_checkpoint, cx);
+    }
+
+    fn finalize_checkpoint(
+        &mut self,
+        pending_checkpoint: ThreadCheckpoint,
+        cx: &mut Context<Self>,
+    ) {
         let git_store = self.project.read(cx).git_store().clone();
         let final_checkpoint = git_store.update(cx, |git_store, cx| git_store.checkpoint(cx));
         cx.spawn(async move |this, cx| match final_checkpoint.await {
@@ -2248,10 +2256,17 @@ impl Thread {
             );
         }
 
-        self.finalize_pending_checkpoint(cx);
-
         if canceled {
             cx.emit(ThreadEvent::CompletionCanceled);
+
+            // When canceled, we always want to insert the checkpoint.
+            // (We skip over finalize_pending_checkpoint, because it
+            // would conclude we didn't have anything to insert here.)
+            if let Some(checkpoint) = self.pending_checkpoint.take() {
+                self.insert_checkpoint(checkpoint, cx);
+            }
+        } else {
+            self.finalize_pending_checkpoint(cx);
         }
 
         canceled