zeta2: Fix double lease panic in event capture (#46096)

Ben Kunkle created 6 days ago

Closes #ISSUE

Release Notes:

- N/A *or* Added/Fixed/Improved ...

Change summary

crates/edit_prediction/src/capture_example.rs       | 16 ++++++++------
crates/edit_prediction/src/edit_prediction.rs       |  3 ++
crates/edit_prediction_ui/src/edit_prediction_ui.rs |  8 +++++-
3 files changed, 18 insertions(+), 9 deletions(-)

Detailed changes

crates/edit_prediction/src/capture_example.rs

@@ -1,5 +1,5 @@
 use crate::{
-    EditPredictionExampleCaptureFeatureFlag, EditPredictionStore, StoredEvent,
+    EditPredictionExampleCaptureFeatureFlag, StoredEvent,
     cursor_excerpt::editable_and_context_ranges_for_cursor_position, example_spec::ExampleSpec,
 };
 use anyhow::Result;
@@ -18,9 +18,9 @@ pub fn capture_example(
     project: Entity<Project>,
     buffer: Entity<Buffer>,
     cursor_anchor: language::Anchor,
+    mut events: Vec<StoredEvent>,
     cx: &mut App,
 ) -> Option<Task<Result<ExampleSpec>>> {
-    let ep_store = EditPredictionStore::try_global(cx)?;
     let snapshot = buffer.read(cx).snapshot();
     let file = snapshot.file()?;
     let worktree_id = file.worktree_id(cx);
@@ -38,10 +38,6 @@ pub fn capture_example(
         .or_else(|| repository_snapshot.remote_upstream_url.clone())?;
     let revision = repository_snapshot.head_commit.as_ref()?.sha.to_string();
 
-    let mut events = ep_store.update(cx, |store, cx| {
-        store.edit_history_for_project_with_pause_split_last_event(&project, cx)
-    });
-
     let git_store = project.read(cx).git_store().clone();
 
     Some(cx.spawn(async move |mut cx| {
@@ -204,6 +200,7 @@ pub(crate) fn should_sample_edit_prediction_example_capture(cx: &App) -> bool {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::EditPredictionStore;
     use client::{Client, UserStore};
     use clock::FakeSystemClock;
     use gpui::{AppContext as _, TestAppContext, http_client::FakeHttpClient};
@@ -315,8 +312,13 @@ mod tests {
         });
         cx.run_until_parked();
 
+        let events = ep_store.update(cx, |store, cx| {
+            store.edit_history_for_project_with_pause_split_last_event(&project, cx)
+        });
         let mut example = cx
-            .update(|cx| capture_example(project.clone(), buffer.clone(), Anchor::MIN, cx).unwrap())
+            .update(|cx| {
+                capture_example(project.clone(), buffer.clone(), Anchor::MIN, events, cx).unwrap()
+            })
             .await
             .unwrap();
         example.name = "test".to_string();

crates/edit_prediction/src/edit_prediction.rs

@@ -1645,10 +1645,13 @@ impl EditPredictionStore {
             && self.can_collect_events(&inputs.events);
 
         if can_collect_example && should_sample_edit_prediction_example_capture(cx) {
+            let events_for_capture =
+                self.edit_history_for_project_with_pause_split_last_event(&project, cx);
             if let Some(example_task) = capture_example::capture_example(
                 project.clone(),
                 active_buffer.clone(),
                 position,
+                events_for_capture,
                 cx,
             ) {
                 cx.spawn(async move |_this, _cx| {

crates/edit_prediction_ui/src/edit_prediction_ui.rs

@@ -3,7 +3,7 @@ mod edit_prediction_context_view;
 mod rate_prediction_modal;
 
 use command_palette_hooks::CommandPaletteFilter;
-use edit_prediction::{ResetOnboarding, Zeta2FeatureFlag, capture_example};
+use edit_prediction::{EditPredictionStore, ResetOnboarding, Zeta2FeatureFlag, capture_example};
 use edit_prediction_context_view::EditPredictionContextView;
 use editor::Editor;
 use feature_flags::FeatureFlagAppExt as _;
@@ -150,7 +150,11 @@ fn capture_example_as_markdown(
         .buffer()
         .read(cx)
         .text_anchor_for_position(editor.selections.newest_anchor().head(), cx)?;
-    let example = capture_example(project.clone(), buffer, cursor_anchor, cx)?;
+    let ep_store = EditPredictionStore::try_global(cx)?;
+    let events = ep_store.update(cx, |store, cx| {
+        store.edit_history_for_project_with_pause_split_last_event(&project, cx)
+    });
+    let example = capture_example(project.clone(), buffer, cursor_anchor, events, cx)?;
 
     let examples_dir = AllLanguageSettings::get_global(cx)
         .edit_predictions