ep: Batch settled edit prediction telemetry (#59857)

# Objective Closes EP-78 - Describe the objective or issue this PR addresses. - If you're fixing a specific issue, use "Fixes #X" for each issue as [described in the GitHub docs](https://docs.github.com/en/issues/tracking-your-work-with-issues/using-issues/linking-a-pull-request-to-an-issue#linking-a-pull-request-to-an-issue-using-a-keyword). ## Solution - Describe the solution used to achieve the objective above. ## Testing - Did you test these changes? If so, how? - Are there any parts that need more testing? - How can other people (reviewers) test your changes? Is there anything specific they need to know? - If relevant, what platforms did you test these changes on, and are there any important ones you can't test? ## Self-Review Checklist: - [x] I've reviewed my own diff for quality, security, and reliability - [x] Unsafe blocks (if any) have justifying comments - [x] The content adheres to Zed's UI standards ([UX/UI](https://github.com/zed-industries/zed/blob/main/CONTRIBUTING.md#uiux-checklist) and [icon](https://github.com/zed-industries/zed/blob/main/crates/icons/README.md) guidelines) - [x] Tests cover the new/changed behavior - [x] Performance impact has been considered and is acceptable ## Showcase > This section is optional. If this PR does not include a visual change or does not add a new user-facing feature, you can delete this section. - Help others understand the result of this PR by showcasing your awesome work! - If this PR includes a visual change, consider adding a screenshot, GIF, or video - A before/after comparison is very useful for changes to existing features! While a showcase should aim to be brief and digestible, you can use a toggleable section to save space on longer showcases: <details> <summary>Click to view showcase</summary> My super cool demos here </details> --- Release Notes: - N/A or Added/Fixed/Improved ...
2026-06-25 01:32:40 +00:00
1 parent 60c6fd3690
commit 0227e83849
@@ -1,95 +0,0 @@
-use crate::{
-    EditPredictionStore, StoredEvent,
-    data_collection::{compute_uncommitted_diff, uncommitted_diffs_for_events},
-    zeta,
-};
-use anyhow::Result;
-use gpui::{Context, Entity, Task};
-use language::{Buffer, Point, ToPoint as _};
-use project::Project;
-use text::OffsetRangeExt as _;
-
-const MAX_UNCOMMITTED_DIFF_SIZE: usize = 64 * 1024;
-
-pub(crate) struct CapturedPredictionContext {
-    pub(crate) repository_url: Option<String>,
-    pub(crate) revision: Option<String>,
-    pub(crate) uncommitted_diff: Option<String>,
-    pub(crate) buffer_diagnostics: Vec<zeta_prompt::ActiveBufferDiagnostic>,
-    pub(crate) editable_context: Vec<zeta_prompt::RelatedFile>,
-}
-
-pub(crate) fn capture_prediction_context(
-    project: Entity<Project>,
-    buffer: Entity<Buffer>,
-    cursor_anchor: language::Anchor,
-    stored_events: Vec<StoredEvent>,
-    repository_url: Option<String>,
-    revision: Option<String>,
-    editable_context_task: Task<Result<Vec<zeta_prompt::RelatedFile>>>,
-    cx: &mut Context<EditPredictionStore>,
-) -> Option<Task<Result<CapturedPredictionContext>>> {
-    let snapshot = buffer.read(cx).snapshot();
-    let worktree_id = snapshot.file()?.worktree_id(cx);
-    let uncommitted_diff_task =
-        uncommitted_diffs_for_events(project, worktree_id, stored_events, cx);
-
-    Some(cx.spawn(async move |_this, cx| {
-        let uncommitted_diff_snapshot = match uncommitted_diff_task.await {
-            Ok(snapshot) => Some(snapshot),
-            Err(error) => {
-                log::debug!("failed to capture uncommitted diff: {error:?}");
-                None
-            }
-        };
-
-        let uncommitted_diff = if let Some(uncommitted_diff_snapshot) = uncommitted_diff_snapshot {
-            let estimated_uncommitted_diff_size = uncommitted_diff_snapshot
-                .iter()
-                .map(|(_, buffer_snapshot, diff_snapshot)| {
-                    diff_snapshot
-                        .hunks(buffer_snapshot)
-                        .map(|hunk| {
-                            hunk.diff_base_byte_range.len()
-                                + hunk.range.to_offset(buffer_snapshot).len()
-                        })
-                        .sum::<usize>()
-                })
-                .sum::<usize>();
-
-            if estimated_uncommitted_diff_size <= MAX_UNCOMMITTED_DIFF_SIZE {
-                let uncommitted_diff = cx
-                    .background_executor()
-                    .spawn(async move { compute_uncommitted_diff(uncommitted_diff_snapshot) })
-                    .await;
-                (uncommitted_diff.len() <= MAX_UNCOMMITTED_DIFF_SIZE).then_some(uncommitted_diff)
-            } else {
-                None
-            }
-        } else {
-            None
-        };
-
-        let buffer_diagnostics = zeta::active_buffer_diagnostics(
-            &snapshot,
-            Point::new(0, 0)..snapshot.max_point(),
-            cursor_anchor.to_point(&snapshot).row,
-            100,
-        );
-        let editable_context = match editable_context_task.await {
-            Ok(editable_context) => editable_context,
-            Err(error) => {
-                log::debug!("failed to capture editable context: {error:?}");
-                Vec::new()
-            }
-        };
-
-        Ok(CapturedPredictionContext {
-            repository_url,
-            revision,
-            uncommitted_diff,
-            buffer_diagnostics,
-            editable_context,
-        })
-    }))
-}
@@ -1,20 +1,105 @@
-use crate::{EditPredictionStore, StoredEvent};
+use crate::{EditPredictionStore, StoredEvent, zeta};

-use anyhow::Context as _;
+use anyhow::{Context as _, Result};
 use buffer_diff::BufferDiffSnapshot;
 use collections::HashMap;
 use gpui::{Context, Entity, Task};
-use language::BufferSnapshot;
+use language::{Buffer, BufferSnapshot, ToPoint as _};
 use project::{Project, ProjectPath, WorktreeId};
 use std::{fmt::Write as _, ops::Range, path::Path, sync::Arc};
-use text::Point;
+use text::{OffsetRangeExt as _, Point};
 use util::rel_path::RelPath;

+const MAX_UNCOMMITTED_DIFF_SIZE: usize = 64 * 1024;
+
 pub type UncommittedDiffSnapshot = Vec<(Arc<Path>, BufferSnapshot, BufferDiffSnapshot)>;
 pub type UncommittedDiffResult = std::result::Result<UncommittedDiffSnapshot, Arc<anyhow::Error>>;

 pub use zeta_prompt::udiff::CURSOR_POSITION_MARKER;

+pub(crate) struct CapturedPredictionContext {
+    pub(crate) repository_url: Option<String>,
+    pub(crate) revision: Option<String>,
+    pub(crate) uncommitted_diff: Option<String>,
+    pub(crate) buffer_diagnostics: Vec<zeta_prompt::ActiveBufferDiagnostic>,
+    pub(crate) editable_context: Vec<zeta_prompt::RelatedFile>,
+}
+
+pub(crate) fn capture_prediction_context(
+    project: Entity<Project>,
+    buffer: Entity<Buffer>,
+    cursor_anchor: language::Anchor,
+    stored_events: Vec<StoredEvent>,
+    repository_url: Option<String>,
+    revision: Option<String>,
+    editable_context_task: Task<Result<Vec<zeta_prompt::RelatedFile>>>,
+    cx: &mut Context<EditPredictionStore>,
+) -> Option<Task<Result<CapturedPredictionContext>>> {
+    let snapshot = buffer.read(cx).snapshot();
+    let worktree_id = snapshot.file()?.worktree_id(cx);
+    let uncommitted_diff_task =
+        uncommitted_diffs_for_events(project, worktree_id, stored_events, cx);
+
+    Some(cx.spawn(async move |_this, cx| {
+        let uncommitted_diff_snapshot = match uncommitted_diff_task.await {
+            Ok(snapshot) => Some(snapshot),
+            Err(error) => {
+                log::debug!("failed to capture uncommitted diff: {error:?}");
+                None
+            }
+        };
+
+        let uncommitted_diff = if let Some(uncommitted_diff_snapshot) = uncommitted_diff_snapshot {
+            let estimated_uncommitted_diff_size = uncommitted_diff_snapshot
+                .iter()
+                .map(|(_, buffer_snapshot, diff_snapshot)| {
+                    diff_snapshot
+                        .hunks(buffer_snapshot)
+                        .map(|hunk| {
+                            hunk.diff_base_byte_range.len()
+                                + hunk.range.to_offset(buffer_snapshot).len()
+                        })
+                        .sum::<usize>()
+                })
+                .sum::<usize>();
+
+            if estimated_uncommitted_diff_size <= MAX_UNCOMMITTED_DIFF_SIZE {
+                let uncommitted_diff = cx
+                    .background_executor()
+                    .spawn(async move { compute_uncommitted_diff(uncommitted_diff_snapshot) })
+                    .await;
+                (uncommitted_diff.len() <= MAX_UNCOMMITTED_DIFF_SIZE).then_some(uncommitted_diff)
+            } else {
+                None
+            }
+        } else {
+            None
+        };
+
+        let buffer_diagnostics = zeta::active_buffer_diagnostics(
+            &snapshot,
+            Point::new(0, 0)..snapshot.max_point(),
+            cursor_anchor.to_point(&snapshot).row,
+            100,
+        );
+        let editable_context = match editable_context_task.await {
+            Ok(editable_context) => editable_context,
+            Err(error) => {
+                log::debug!("failed to capture editable context: {error:?}");
+                Vec::new()
+            }
+        };
+
+        Ok(CapturedPredictionContext {
+            repository_url,
+            revision,
+            uncommitted_diff,
+            buffer_diagnostics,
+            editable_context,
+        })
+    }))
+}
+
 pub fn uncommitted_diffs_for_events(
    project: Entity<Project>,
    worktree_id: WorktreeId,
@@ -3,8 +3,10 @@ use buffer_diff::BufferDiff;
 use client::{Client, EditPredictionUsage, UserStore, global_llm_token};
 use cloud_api_client::LlmApiToken;
 use cloud_api_types::{
-    EditPredictionRecentFile, EditPredictionSettledKeptChars, OrganizationId,
-    SettledEditPrediction, SettledEditPredictionSampleData, SubmitEditPredictionFeedbackBody,
+    EditPredictionRecentFile, EditPredictionSettledKeptChars,
+    MAX_EDIT_PREDICTION_SETTLED_PER_REQUEST, OrganizationId, SettledEditPrediction,
+    SettledEditPredictionSampleData, SubmitEditPredictionFeedbackBody,
+    SubmitEditPredictionSettledBatchBody, SubmitEditPredictionSettledResponse,
 };
 use cloud_llm_client::predict_edits_v3::{
    PREDICT_EDITS_MODE_HEADER_NAME, PREDICT_EDITS_REQUEST_ID_HEADER_NAME,
@@ -83,7 +85,6 @@ mod prediction;

 pub mod udiff;

-mod capture_prediction_context;
 pub mod open_ai_compatible;
 mod zed_edit_prediction_delegate;
 pub mod zeta;
@@ -91,8 +92,8 @@ pub mod zeta;
 #[cfg(test)]
 mod edit_prediction_tests;

-use crate::capture_prediction_context::{CapturedPredictionContext, capture_prediction_context};
 use crate::cursor_excerpt::expand_context_syntactically_then_linewise;
+use crate::data_collection::{CapturedPredictionContext, capture_prediction_context};
 use crate::example_spec::RecentFile;
 use crate::license_detection::LicenseDetectionWatcher;
 use crate::mercury::Mercury;
@@ -130,8 +131,6 @@ const REQUEST_TIMEOUT_BACKOFF: Duration = Duration::from_secs(10);
 const EDIT_PREDICTION_SETTLED_TTL: Duration = Duration::from_secs(60 * 5);
 const EDIT_PREDICTION_SETTLED_QUIESCENCE: Duration = Duration::from_secs(10);
 const EDIT_PREDICTION_CAPTURE_MAX_FUTURE_EVENTS: usize = 4;
-/// The server rejects settled bodies larger than 64 KiB (compressed).
-const EDIT_PREDICTION_SETTLED_MAX_BODY_BYTES: usize = 63 * 1024;
 const EDIT_PREDICTION_SETTLED_MAX_EDITABLE_REGION_BYTES: usize = 4 * 1024;

 pub struct EditPredictionJumpsFeatureFlag;
@@ -2026,25 +2025,11 @@ impl EditPredictionStore {
                }
            });

+            let mut ready_predictions_by_organization_id: HashMap<_, Vec<_>> = HashMap::default();
            for (pending_capture, settled_editable_region) in ready_predictions {
-                let PendingPredictionCapture {
-                    request_id,
-                    editable_region_before_prediction,
-                    predicted_editable_region,
-                    ts_error_count_before_prediction,
-                    ts_error_count_after_prediction,
-                    organization_id,
-                    can_collect_data,
-                    is_in_open_source_repo,
-                    sample_data,
-                    model_version,
-                    e2e_latency,
-                    ..
-                } = pending_capture;
-                let settled_editable_region_for_metrics = settled_editable_region.clone();
                #[cfg(test)]
                {
-                    let request_id = request_id.clone();
+                    let request_id = pending_capture.request_id.clone();
                    let settled_editable_region = settled_editable_region.clone();
                    this.update(cx, |this, _| {
                        if let Some(callback) = &this.settled_event_callback {
@@ -2052,111 +2037,28 @@ impl EditPredictionStore {
                        }
                    });
                }
-                cx.background_spawn({
-                    let client = client.clone();
-                    let llm_token = llm_token.clone();
-                    let app_version = app_version.clone();
-                    async move {
-                        let kept_rate_result = compute_kept_rate(
-                            &editable_region_before_prediction,
-                            &predicted_editable_region,
-                            &settled_editable_region_for_metrics,
-                        );
-
-                        let result: anyhow::Result<()> = async {
-                            let sample_data = if can_collect_data
-                                && let Some(sample_data) = sample_data
-                                && let Ok(context) = sample_data.context_task.await
-                            {
-                                Some(SettledEditPredictionSampleData {
-                                    repository_url: context.repository_url,
-                                    revision: context.revision,
-                                    uncommitted_diff: context.uncommitted_diff,
-                                    editable_path: sample_data.editable_path,
-                                    editable_offset_range: sample_data.editable_offset_range,
-                                    buffer_diagnostics: context.buffer_diagnostics,
-                                    editable_context: context.editable_context,
-                                    future_edit_history_events: sample_data
-                                        .future_edit_history_events,
-                                    navigation_history: sample_data
-                                        .navigation_history
-                                        .into_iter()
-                                        .map(|file| EditPredictionRecentFile {
-                                            path: file.path,
-                                            cursor_position: file.cursor_position,
-                                        })
-                                        .collect(),
-                                    edit_events_before_quiescence: sample_data
-                                        .edit_events_before_quiescence,
-                                    next_edit_cursor_offset: sample_data.next_edit_cursor_offset,
-                                })
-                            } else {
-                                None
-                            };
-                            let settled_editable_region =
-                                can_collect_data.then_some(settled_editable_region);
-
-                            let mut body = SettledEditPrediction {
-                                request_id: request_id.0.to_string(),
-                                settled_editable_region,
-                                ts_error_count_before_prediction,
-                                ts_error_count_after_prediction,
-                                can_collect_data,
-                                is_in_open_source_repo,
-                                sample_data,
-                                kept_chars: EditPredictionSettledKeptChars {
-                                    candidate_new: kept_rate_result.candidate_new_chars,
-                                    reference_new: kept_rate_result.reference_new_chars,
-                                    candidate_deleted: kept_rate_result.candidate_deleted_chars,
-                                    reference_deleted: kept_rate_result.reference_deleted_chars,
-                                    kept: kept_rate_result.kept_chars,
-                                    correctly_deleted: kept_rate_result.correctly_deleted_chars,
-                                    discarded: kept_rate_result.discarded_chars,
-                                    context: kept_rate_result.context_chars,
-                                    kept_rate: kept_rate_result.kept_rate,
-                                    recall_rate: kept_rate_result.recall_rate,
-                                },
-                                example: None,
-                                model_version,
-                                e2e_latency_ms: e2e_latency.as_millis().min(u128::from(u64::MAX))
-                                    as u64,
-                            };
-
-                            let mut compressed =
-                                zstd::encode_all(&serde_json::to_vec(&body)?[..], 3)?;
-                            if compressed.len() > EDIT_PREDICTION_SETTLED_MAX_BODY_BYTES {
-                                body.sample_data = None;
-                                compressed = zstd::encode_all(&serde_json::to_vec(&body)?[..], 3)?;
-                            }
-
-                            let url = client
-                                .http_client()
-                                .build_zed_llm_url("/predict_edits/settled", &[])?;
-                            Self::send_api_request::<serde_json::Value>(
-                                |builder| {
-                                    Ok(builder
-                                        .uri(url.as_ref())
-                                        .header("Content-Encoding", "zstd")
-                                        .body(compressed.clone().into())?)
-                                },
-                                client,
-                                llm_token,
-                                organization_id,
-                                app_version,
-                            )
-                            .await?;
-                            Ok(())
-                        }
-                        .await;
-
-                        if let Err(error) = result {
-                            log::error!("failed to submit edit prediction settled: {error:?}");
-                        }
-                    }
-                })
-                .detach();
+                ready_predictions_by_organization_id
+                    .entry(pending_capture.organization_id.clone())
+                    .or_default()
+                    .push((pending_capture, settled_editable_region));
            }

+            cx.background_spawn({
+                let client = client.clone();
+                let llm_token = llm_token.clone();
+                let app_version = app_version.clone();
+                async move {
+                    send_settled_batches(
+                        client,
+                        llm_token,
+                        app_version,
+                        ready_predictions_by_organization_id,
+                    )
+                    .await;
+                }
+            })
+            .detach();
+
            next_wake_time = oldest_edited_at.map(|time| time + EDIT_PREDICTION_SETTLED_QUIESCENCE);
        }
    }
@@ -2416,6 +2318,137 @@ impl EditPredictionStore {
    pub const THROTTLE_TIMEOUT: Duration = Duration::from_millis(300);
 }

+async fn send_settled_batches(
+    client: Arc<Client>,
+    llm_token: LlmApiToken,
+    app_version: Version,
+    ready_predictions_by_organization_id: hash_map::HashMap<
+        Option<OrganizationId>,
+        Vec<(PendingPredictionCapture, String)>,
+        collections::FxBuildHasher,
+    >,
+) {
+    let Some(url) = client
+        .http_client()
+        .build_zed_llm_url("/predict_edits/settled", &[])
+        .context("failed to build edit predictions settled url")
+        .log_err()
+    else {
+        return;
+    };
+
+    for (organization_id, ready_predictions) in ready_predictions_by_organization_id {
+        let mut ready_predictions = ready_predictions.into_iter();
+        loop {
+            let done_batch = ready_predictions
+                .by_ref()
+                .take(MAX_EDIT_PREDICTION_SETTLED_PER_REQUEST);
+            let mut batch = Vec::with_capacity(MAX_EDIT_PREDICTION_SETTLED_PER_REQUEST);
+            for (pending_capture, settled_editable_region) in done_batch {
+                let PendingPredictionCapture {
+                    request_id,
+                    editable_region_before_prediction,
+                    predicted_editable_region,
+                    ts_error_count_before_prediction,
+                    ts_error_count_after_prediction,
+                    can_collect_data,
+                    is_in_open_source_repo,
+                    sample_data,
+                    model_version,
+                    e2e_latency,
+                    ..
+                } = pending_capture;
+                let kept_rate_result = compute_kept_rate(
+                    &editable_region_before_prediction,
+                    &predicted_editable_region,
+                    &settled_editable_region,
+                );
+
+                let sample_data = if can_collect_data
+                    && let Some(sample_data) = sample_data
+                    && let Ok(context) = sample_data.context_task.await
+                {
+                    Some(SettledEditPredictionSampleData {
+                        repository_url: context.repository_url,
+                        revision: context.revision,
+                        uncommitted_diff: context.uncommitted_diff,
+                        editable_path: sample_data.editable_path,
+                        editable_offset_range: sample_data.editable_offset_range,
+                        buffer_diagnostics: context.buffer_diagnostics,
+                        editable_context: context.editable_context,
+                        future_edit_history_events: sample_data.future_edit_history_events,
+                        navigation_history: sample_data
+                            .navigation_history
+                            .into_iter()
+                            .map(|file| EditPredictionRecentFile {
+                                path: file.path,
+                                cursor_position: file.cursor_position,
+                            })
+                            .collect(),
+                        edit_events_before_quiescence: sample_data.edit_events_before_quiescence,
+                        next_edit_cursor_offset: sample_data.next_edit_cursor_offset,
+                    })
+                } else {
+                    None
+                };
+
+                batch.push(SettledEditPrediction {
+                    request_id: request_id.0.to_string(),
+                    settled_editable_region: can_collect_data.then_some(settled_editable_region),
+                    ts_error_count_before_prediction,
+                    ts_error_count_after_prediction,
+                    can_collect_data,
+                    is_in_open_source_repo,
+                    sample_data,
+                    kept_chars: EditPredictionSettledKeptChars {
+                        candidate_new: kept_rate_result.candidate_new_chars,
+                        reference_new: kept_rate_result.reference_new_chars,
+                        candidate_deleted: kept_rate_result.candidate_deleted_chars,
+                        reference_deleted: kept_rate_result.reference_deleted_chars,
+                        kept: kept_rate_result.kept_chars,
+                        correctly_deleted: kept_rate_result.correctly_deleted_chars,
+                        discarded: kept_rate_result.discarded_chars,
+                        context: kept_rate_result.context_chars,
+                        kept_rate: kept_rate_result.kept_rate,
+                        recall_rate: kept_rate_result.recall_rate,
+                    },
+                    example: None,
+                    model_version,
+                    e2e_latency_ms: e2e_latency.as_millis().min(u128::from(u64::MAX)) as u64,
+                });
+            }
+
+            if batch.is_empty() {
+                break;
+            }
+
+            let result = async {
+                let body = SubmitEditPredictionSettledBatchBody { predictions: batch };
+                let compressed = zstd::encode_all(&serde_json::to_vec(&body)?[..], 3)?;
+                EditPredictionStore::send_api_request::<SubmitEditPredictionSettledResponse>(
+                    |builder| {
+                        Ok(builder
+                            .uri(url.as_ref())
+                            .header("Content-Encoding", "zstd")
+                            .body(compressed.clone().into())?)
+                    },
+                    client.clone(),
+                    llm_token.clone(),
+                    organization_id.clone(),
+                    app_version.clone(),
+                )
+                .await?;
+                anyhow::Ok(())
+            }
+            .await;
+
+            if let Err(error) = result {
+                log::error!("failed to submit edit predictions settled: {error:?}");
+            }
+        }
+    }
+}
+
 fn currently_following(project: &Entity<Project>, cx: &App) -> bool {
    let Some(app_state) = AppState::try_global(cx) else {
        return false;
@@ -4,7 +4,7 @@ use clock::ReplicaId;
 use cloud_api_types::{
    CreateLlmTokenResponse, LlmToken, Organization, OrganizationConfiguration,
    OrganizationEditPredictionConfiguration, OrganizationId, SettledEditPrediction,
-    SubmitEditPredictionSettledResponse,
+    SubmitEditPredictionSettledBatchBody, SubmitEditPredictionSettledResponse,
 };
 use cloud_llm_client::{
    EditPredictionRejectReason, EditPredictionRejection, PredictEditsRequestTrigger,
@@ -259,6 +259,68 @@ async fn test_simple_request(cx: &mut TestAppContext) {
    assert_eq!(prediction.edits[0].1.as_ref(), " are you?");
 }

+#[gpui::test]
+async fn test_zeta_request_sends_settled_body_when_data_collection_is_disabled(
+    cx: &mut TestAppContext,
+) {
+    let (ep_store, mut requests) = init_test_with_fake_client(cx);
+    let fs = FakeFs::new(cx.executor());
+    fs.insert_tree(
+        "/root",
+        json!({
+            "foo.md":  "Hello!\nHow\nBye\n"
+        }),
+    )
+    .await;
+    let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await;
+
+    let buffer = project
+        .update(cx, |project, cx| {
+            let path = project.find_project_path(path!("root/foo.md"), cx).unwrap();
+            project.open_buffer(path, cx)
+        })
+        .await
+        .unwrap();
+    let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot());
+    let position = snapshot.anchor_before(language::Point::new(1, 3));
+
+    ep_store.update(cx, |ep_store, cx| {
+        ep_store.register_buffer(&buffer, &project, cx);
+    });
+
+    let prediction_task = ep_store.update(cx, |ep_store, cx| {
+        ep_store.request_prediction(
+            &project,
+            &buffer,
+            position,
+            PredictEditsRequestTrigger::Other,
+            cx,
+        )
+    });
+
+    let (request, respond_tx) = requests.predict.next().await.unwrap();
+    assert!(!request.input.can_collect_data);
+    respond_tx
+        .send(model_response(&request, SIMPLE_DIFF))
+        .unwrap();
+
+    prediction_task.await.unwrap().unwrap();
+    cx.run_until_parked();
+    cx.executor()
+        .advance_clock(EDIT_PREDICTION_SETTLED_QUIESCENCE);
+    cx.run_until_parked();
+
+    let settled_request = requests
+        .settled
+        .next()
+        .now_or_never()
+        .flatten()
+        .expect("settled request should be sent");
+    assert!(!settled_request.can_collect_data);
+    assert_eq!(settled_request.settled_editable_region, None);
+    assert_eq!(settled_request.sample_data, None);
+}
+
 #[gpui::test]
 async fn test_request_events(cx: &mut TestAppContext) {
    let (ep_store, mut requests) = init_test_with_fake_client(cx);
@@ -1270,6 +1332,19 @@ async fn test_empty_prediction(cx: &mut TestAppContext) {
            e2e_latency_ms: Some(0),
        }]
    );
+    cx.executor()
+        .advance_clock(EDIT_PREDICTION_SETTLED_QUIESCENCE);
+    cx.run_until_parked();
+
+    let settled_request = requests
+        .settled
+        .next()
+        .now_or_never()
+        .flatten()
+        .expect("empty prediction should still send settled request");
+    assert_eq!(settled_request.request_id, id);
+    assert_eq!(settled_request.settled_editable_region, None);
+    assert_eq!(settled_request.sample_data, None);
 }

 #[gpui::test]
@@ -2533,8 +2608,11 @@ fn init_test_with_fake_client_and_legacy_data_collection(
                            } else {
                                buf
                            };
-                            let req = serde_json::from_slice(&body).unwrap();
-                            settled_req_tx.unbounded_send(req).unwrap();
+                            let req: SubmitEditPredictionSettledBatchBody =
+                                serde_json::from_slice(&body).unwrap();
+                            for prediction in req.predictions {
+                                settled_req_tx.unbounded_send(prediction).unwrap();
+                            }
                            serde_json::to_string(&SubmitEditPredictionSettledResponse {}).unwrap()
                        }
                        _ => {
@@ -3,8 +3,8 @@ use crate::{
    EditPredictionId, EditPredictionModelInput, EditPredictionStartedDebugEvent,
    EditPredictionStore, PromptHistoryBoundary, ZedUpdateRequiredError,
    buffer_path_with_id_fallback,
-    capture_prediction_context::CapturedPredictionContext,
    cursor_excerpt::{self, compute_cursor_excerpt, compute_syntax_ranges},
+    data_collection::CapturedPredictionContext,
    prediction::EditPredictionResult,
 };
 use anyhow::Result;
@@ -404,36 +404,34 @@ pub(crate) fn request_prediction_with_zeta(
        )
        .await;

-        if can_collect_data {
-            let prediction = &result.prediction;
-            let weak_this = this.clone();
-            let request_id = prediction.id.clone();
-            let edited_buffer = edited_buffer.clone();
-            let edited_buffer_snapshot = edited_buffer_snapshot.clone();
-            let editable_range_in_buffer = editable_range_in_buffer.clone();
-            let edit_preview = prediction.edit_preview.clone();
-            let model_version = prediction.model_version.clone();
-            cx.spawn(async move |cx| {
-                weak_this
-                    .update(cx, |this, cx| {
-                        this.enqueue_settled_prediction(
-                            request_id.clone(),
-                            &project,
-                            &edited_buffer,
-                            &edited_buffer_snapshot,
-                            editable_range_in_buffer,
-                            &edit_preview,
-                            context_task,
-                            prompt_history_boundary,
-                            model_version,
-                            request_duration,
-                            cx,
-                        );
-                    })
-                    .ok();
-            })
-            .detach();
-        }
+        let prediction = &result.prediction;
+        let weak_this = this.clone();
+        let request_id = prediction.id.clone();
+        let edited_buffer = edited_buffer.clone();
+        let edited_buffer_snapshot = edited_buffer_snapshot.clone();
+        let editable_range_in_buffer = editable_range_in_buffer.clone();
+        let edit_preview = prediction.edit_preview.clone();
+        let model_version = prediction.model_version.clone();
+        cx.spawn(async move |cx| {
+            weak_this
+                .update(cx, |this, cx| {
+                    this.enqueue_settled_prediction(
+                        request_id.clone(),
+                        &project,
+                        &edited_buffer,
+                        &edited_buffer_snapshot,
+                        editable_range_in_buffer,
+                        &edit_preview,
+                        context_task,
+                        prompt_history_boundary,
+                        model_version,
+                        request_duration,
+                        cx,
+                    );
+                })
+                .ok();
+        })
+        .detach();

        Ok(Some(result))
    })