yenru0/zed
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

ep: Combine settled and example capture (#58759)

Browse Source 
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 is consistent with the [UI/UX
checklist](https://github.com/zed-industries/zed/blob/main/CONTRIBUTING.md#uiux-checklist)
- [x] Tests cover the new/changed behavior
- [x] Performance impact has been considered and is acceptable

Closes #ISSUE

Release Notes:

- N/A or Added/Fixed/Improved ...
This commit is contained in:
Ben Kunkle authored and GitHub committed 2026-06-22 16:45:50 +00:00
1 parent 3e9408a9f4
commit 834b56642a
13 files changed
+973 -1265
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

Cargo.lock
Generated
-1
View File
@@ -3265,7 +3265,6 @@ dependencies = [
"serde",
"serde_json",
"strum 0.27.2",
"uuid",
"zeta_prompt",
]
crates/cloud_api_types/Cargo.toml
-1
View File
@@ -19,7 +19,6 @@ cloud_llm_client.workspace = true
serde.workspace = true
serde_json.workspace = true
strum.workspace = true
uuid.workspace = true
zeta_prompt.workspace = true
[dev-dependencies]
crates/cloud_api_types/src/cloud_api_types.rs
+35 -34
View File
@@ -6,6 +6,7 @@ mod timestamp;
pub mod websocket_protocol;
use std::collections::BTreeMap;
use std::ops::Range;
use std::path::Path;
use std::sync::Arc;
@@ -129,58 +130,58 @@ pub struct SubmitEditPredictionFeedbackBody {
}
#[derive(Debug, PartialEq, Serialize, Deserialize)]
pub struct SubmitEditPredictionSettledBody {
pub struct SettledEditPrediction {
pub request_id: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub settled_editable_region: Option<String>,
pub ts_error_count_before_prediction: usize,
pub ts_error_count_after_prediction: usize,
pub can_collect_data: bool,
pub is_in_open_source_repo: bool,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub sample_data: Option<SettledEditPredictionSampleData>,
#[serde(flatten)]
pub kept_chars: EditPredictionSettledKeptChars,
pub example: Option<serde_json::Value>,
pub model_version: Option<String>,
#[serde(rename = "e2e_latency")]
pub e2e_latency_ms: u128,
pub e2e_latency_ms: u64,
}
#[derive(Debug, PartialEq, Serialize, Deserialize)]
pub struct SettledEditPredictionSampleData {
pub repository_url: Option<String>,
pub revision: Option<String>,
/// Note: this is only the uncommitted diff for files in `edit_history`
/// This is done to avoid excessive memory usage
#[serde(default, skip_serializing_if = "Option::is_none")]
pub uncommitted_diff: Option<String>,
pub editable_path: Arc<Path>,
pub editable_offset_range: Range<usize>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub buffer_diagnostics: Vec<zeta_prompt::ActiveBufferDiagnostic>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub future_edit_history_events: Vec<Arc<zeta_prompt::Event>>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub navigation_history: Vec<EditPredictionRecentFile>,
pub edit_events_before_quiescence: u32,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub next_edit_cursor_offset: Option<usize>,
}
pub const MAX_EDIT_PREDICTION_SETTLED_PER_REQUEST: usize = 32;
#[derive(Debug, PartialEq, Serialize, Deserialize)]
pub struct SubmitEditPredictionSettledBatchBody {
pub predictions: Vec<SettledEditPrediction>,
}
#[derive(Debug, PartialEq, Serialize, Deserialize)]
pub struct SubmitEditPredictionSettledResponse {}
#[derive(Debug, PartialEq, Serialize, Deserialize)]
pub struct SubmitEditPredictionJumpExampleBody {
pub request_id: uuid::Uuid,
pub trigger: JumpExampleTrigger,
pub repository_url: Option<String>,
pub revision: Option<String>,
/// Note: this is only the uncommitted diff for files in `edit_history`
/// This is done to avoid excessive memory usage
pub uncommitted_diff: Option<String>,
pub recently_opened_files: Vec<JumpExampleRecentFile>,
pub recently_viewed_files: Vec<JumpExampleRecentFile>,
pub cursor_path: Arc<Path>,
pub cursor_position: String,
pub edit_history: Vec<Arc<zeta_prompt::Event>>,
pub diagnostics: Vec<zeta_prompt::ActiveBufferDiagnostic>,
pub future_edit_history: String,
pub navigation_history: Vec<JumpExampleRecentFile>,
pub is_in_open_source_repo: bool,
pub can_collect_data: bool,
}
#[derive(Debug, PartialEq, Serialize, Deserialize)]
pub struct SubmitEditPredictionJumpExampleResponse {}
#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum JumpExampleTrigger {
Prediction,
Diagnostic,
}
#[derive(Debug, PartialEq, Serialize, Deserialize)]
pub struct JumpExampleRecentFile {
pub struct EditPredictionRecentFile {
pub path: Arc<Path>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub cursor_position: Option<usize>,
crates/edit_prediction/src/capture_example.rs
-565
View File
@@ -1,565 +0,0 @@
use crate::{
StoredEvent,
data_collection::{
UncommittedDiffSnapshot, compute_cursor_excerpt, compute_uncommitted_diff,
format_cursor_excerpt,
},
example_spec::{ExampleSpec, RecentFile},
};
use anyhow::Result;
use gpui::{App, Entity, Task};
use language::Buffer;
use project::Project;
use std::{fmt::Write, path::Path, sync::Arc};
pub fn capture_example(
project: Entity<Project>,
buffer: Entity<Buffer>,
cursor_anchor: language::Anchor,
events: Vec<StoredEvent>,
recently_opened_files: Vec<RecentFile>,
recently_viewed_files: Vec<RecentFile>,
uncommitted_diff_snapshot: UncommittedDiffSnapshot,
populate_expected_patch: bool,
cx: &mut App,
) -> Option<Task<Result<ExampleSpec>>> {
let snapshot = buffer.read(cx).snapshot();
let file = snapshot.file()?;
let worktree_id = file.worktree_id(cx);
let repository = project.read(cx).active_repository(cx)?;
let repository_snapshot = repository.read(cx).snapshot();
let worktree = project.read(cx).worktree_for_id(worktree_id, cx)?;
let root_name = worktree.read(cx).root_name_str().to_owned();
let cursor_path: Arc<Path> = file.path().as_std_path().into();
if worktree.read(cx).abs_path() != repository_snapshot.work_directory_abs_path {
return None;
}
let repository_url = repository_snapshot
.remote_origin_url
.clone()
.or_else(|| repository_snapshot.remote_upstream_url.clone())?;
let revision = repository_snapshot.head_commit.as_ref()?.sha.to_string();
Some(cx.spawn(async move |cx| {
let uncommitted_diff = cx
.background_executor()
.spawn(async move { compute_uncommitted_diff(uncommitted_diff_snapshot) })
.await;
let line_comment_prefix = snapshot
.language()
.and_then(|lang| lang.config().line_comments.first())
.map(|s| s.to_string())
.unwrap_or_default();
let (cursor_excerpt, cursor_offset_in_excerpt, cursor_excerpt_range) = cx
.background_executor()
.spawn(async move { compute_cursor_excerpt(&snapshot, cursor_anchor) })
.await;
let mut edit_history = String::new();
for stored_event in &events {
write_event_with_relative_paths(&mut edit_history, &stored_event.event, &root_name);
if !edit_history.ends_with('\n') {
edit_history.push('\n');
}
}
let uncommitted_diff_contains_edit_history = !edit_history.is_empty();
// Initialize an empty patch with context lines, to make it easy
// to write the expected patch by hand.
let mut expected_patches = Vec::new();
let mut rejected_patch = None;
if populate_expected_patch {
let mut empty_patch = String::new();
let start_row = cursor_excerpt_range.start.row + 1;
let row_count = cursor_excerpt_range.end.row - cursor_excerpt_range.start.row + 1;
writeln!(&mut empty_patch, "--- a/{}", cursor_path.display()).ok();
writeln!(&mut empty_patch, "+++ b/{}", cursor_path.display()).ok();
writeln!(
&mut empty_patch,
"@@ -{},{} +{},{} @@",
start_row, row_count, start_row, row_count,
)
.ok();
for line in cursor_excerpt.lines() {
writeln!(&mut empty_patch, " {}", line).ok();
}
expected_patches.push(empty_patch.clone());
rejected_patch = Some(empty_patch);
}
let spec = ExampleSpec {
name: generate_timestamp_name(),
repository_url,
revision,
tags: Vec::new(),
reasoning: None,
uncommitted_diff,
recently_opened_files,
recently_viewed_files,
uncommitted_diff_contains_edit_history,
cursor_path,
cursor_position: format_cursor_excerpt(
&cursor_excerpt,
cursor_offset_in_excerpt,
&line_comment_prefix,
),
edit_history,
expected_patches,
rejected_patch,
telemetry: None,
human_feedback: Vec::new(),
rating: None,
};
Ok(spec)
}))
}
pub(crate) fn write_event_with_relative_paths(
output: &mut String,
event: &zeta_prompt::Event,
root_name: &str,
) {
fn write_relative_path(output: &mut String, path: &Path, root_name: &str) {
for component in path.strip_prefix(root_name).unwrap_or(path).components() {
output.push('/');
write!(output, "{}", component.as_os_str().to_string_lossy()).ok();
}
}
let zeta_prompt::Event::BufferChange {
path,
old_path,
diff,
..
} = event;
output.push_str("--- a");
write_relative_path(output, old_path.as_ref(), root_name);
output.push_str("\n+++ b");
write_relative_path(output, path.as_ref(), root_name);
output.push('\n');
output.push_str(diff);
}
fn generate_timestamp_name() -> String {
let format = time::format_description::parse("[year]-[month]-[day] [hour]:[minute]:[second]");
match format {
Ok(format) => {
let now = time::OffsetDateTime::now_local()
.unwrap_or_else(|_| time::OffsetDateTime::now_utc());
now.format(&format)
.unwrap_or_else(|_| "unknown-time".to_string())
}
Err(_) => "unknown-time".to_string(),
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::EditPredictionStore;
use crate::data_collection::uncommitted_diffs_for_events;
use client::RefreshLlmTokenListener;
use client::{Client, UserStore};
use clock::FakeSystemClock;
use gpui::{AppContext as _, TestAppContext, http_client::FakeHttpClient};
use indoc::indoc;
use language::{Anchor, Point};
use project::{FakeFs, Project};
use serde_json::json;
use settings::SettingsStore;
use std::path::Path;
#[gpui::test]
async fn test_capture_example(cx: &mut TestAppContext) {
init_test(cx);
let fs = FakeFs::new(cx.executor());
let committed_contents = indoc! {"
fn main() {
one();
two();
three();
four();
five();
six();
seven();
eight();
nine();
}
"};
let disk_contents = indoc! {"
fn main() {
// comment 1
one();
two();
three();
four();
five();
six();
seven();
eight();
// comment 2
nine();
}
"};
fs.insert_tree(
"/project",
json!({
".git": {},
"src": {
"deleted.rs": "pub fn deleted_file() {\n deleted();\n}\n",
"main.rs": disk_contents,
"new.rs": "pub fn new_file() {\n}\n",
}
}),
)
.await;
// Create an external file outside the main project
fs.insert_tree(
"/external",
json!({
"external.rs": "fn external() {}\n",
}),
)
.await;
fs.set_head_for_repo(
Path::new("/project/.git"),
&[
(
"src/deleted.rs",
"pub fn deleted_file() {\n deleted();\n}\n".to_string(),
),
("src/main.rs", committed_contents.to_string()),
],
"abc123def456",
);
fs.set_remote_for_repo(
Path::new("/project/.git"),
"origin",
"https://github.com/test/repo.git",
);
let project = Project::test(fs.clone(), ["/project".as_ref()], cx).await;
let buffer = project
.update(cx, |project, cx| {
project.open_local_buffer("/project/src/main.rs", cx)
})
.await
.unwrap();
let ep_store = cx.read(|cx| EditPredictionStore::try_global(cx).unwrap());
ep_store.update(cx, |ep_store, cx| {
ep_store.register_buffer(&buffer, &project, cx)
});
cx.run_until_parked();
let deleted_file_buffer = project
.update(cx, |project, cx| {
project.open_local_buffer("/project/src/deleted.rs", cx)
})
.await
.unwrap();
ep_store.update(cx, |ep_store, cx| {
ep_store.register_buffer(&deleted_file_buffer, &project, cx)
});
cx.run_until_parked();
deleted_file_buffer.update(cx, |buffer, cx| {
buffer.edit([(0..buffer.len(), "")], None, cx);
});
cx.run_until_parked();
buffer.update(cx, |buffer, cx| {
let point = Point::new(6, 0);
buffer.edit([(point..point, " // comment 3\n")], None, cx);
let point = Point::new(4, 0);
buffer.edit([(point..point, " // comment 4\n")], None, cx);
pretty_assertions::assert_eq!(
buffer.text(),
indoc! {"
fn main() {
// comment 1
one();
two();
// comment 4
three();
four();
// comment 3
five();
six();
seven();
eight();
// comment 2
nine();
}
"}
);
});
cx.run_until_parked();
let new_file_buffer = project
.update(cx, |project, cx| {
project.open_local_buffer("/project/src/new.rs", cx)
})
.await
.unwrap();
ep_store.update(cx, |ep_store, cx| {
ep_store.register_buffer(&new_file_buffer, &project, cx)
});
cx.run_until_parked();
new_file_buffer.update(cx, |buffer, cx| {
let point = Point::new(1, 0);
buffer.edit([(point..point, " created();\n")], None, cx);
});
cx.run_until_parked();
// Open and edit an external file (outside the main project's worktree)
let external_buffer = project
.update(cx, |project, cx| {
project.open_local_buffer("/external/external.rs", cx)
})
.await
.unwrap();
ep_store.update(cx, |ep_store, cx| {
ep_store.register_buffer(&external_buffer, &project, cx)
});
cx.run_until_parked();
external_buffer.update(cx, |buffer, cx| {
let point = Point::new(0, 0);
buffer.edit([(point..point, "// external edit\n")], None, cx);
});
cx.run_until_parked();
// Verify the external edit was recorded in events
let events = ep_store.update(cx, |store, cx| store.edit_history_for_project(&project, cx));
assert!(
matches!(
events
.last()
.unwrap()
.event
.as_ref(),
zeta_prompt::Event::BufferChange { path, .. } if path.as_ref() == "/external/external.rs"
),
"external file edit should be in events"
);
let worktree_id = buffer.read_with(cx, |buffer, cx| buffer.file().unwrap().worktree_id(cx));
let unfiltered_uncommitted_diffs = ep_store
.update(cx, |_store, cx| {
uncommitted_diffs_for_events(project.clone(), worktree_id, events.clone(), cx)
})
.await
.unwrap();
assert!(
unfiltered_uncommitted_diffs
.iter()
.all(|(path, _, _)| path.as_ref() != Path::new("external.rs"))
);
let project_events = events
.into_iter()
.filter(|event| {
let zeta_prompt::Event::BufferChange { path, .. } = event.event.as_ref();
path.as_ref() != "/external/external.rs"
})
.collect::<Vec<_>>();
let uncommitted_diffs_by_path = ep_store
.update(cx, |_store, cx| {
uncommitted_diffs_for_events(
project.clone(),
worktree_id,
project_events.clone(),
cx,
)
})
.await
.unwrap();
let mut example = cx
.update(|cx| {
capture_example(
project.clone(),
buffer.clone(),
Anchor::min_for_buffer(buffer.read(cx).remote_id()),
project_events,
Vec::new(),
Vec::new(),
uncommitted_diffs_by_path,
true,
cx,
)
.unwrap()
})
.await
.unwrap();
example.name = "test".to_string();
pretty_assertions::assert_eq!(
example,
ExampleSpec {
name: "test".to_string(),
repository_url: "https://github.com/test/repo.git".to_string(),
revision: "abc123def456".to_string(),
tags: Vec::new(),
reasoning: None,
uncommitted_diff: indoc! {"
--- a/src/deleted.rs
+++ b/src/deleted.rs
@@ -1,3 +1,0 @@
-pub fn deleted_file() {
- deleted();
-}
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,11 +1,15 @@
fn main() {
+ // comment 1
one();
two();
+ // comment 4
three();
four();
+ // comment 3
five();
six();
seven();
eight();
+ // comment 2
nine();
}
--- /dev/null
+++ b/src/new.rs
@@ -0,0 +1,3 @@
+pub fn new_file() {
+ created();
+}
"}
.to_string(),
recently_opened_files: Vec::new(),
recently_viewed_files: Vec::new(),
uncommitted_diff_contains_edit_history: true,
cursor_path: Path::new("src/main.rs").into(),
cursor_position: indoc! {"
fn main() {
^[CURSOR_POSITION]
// comment 1
one();
two();
// comment 4
three();
four();
// comment 3
five();
six();
seven();
eight();
// comment 2
nine();
}
"}
.to_string(),
edit_history: indoc! {"
--- a/src/deleted.rs
+++ b/src/deleted.rs
@@ -1,3 +1,0 @@
-pub fn deleted_file() {
- deleted();
-}
--- a/src/main.rs
+++ b/src/main.rs
@@ -2,8 +2,10 @@
// comment 1
one();
two();
+ // comment 4
three();
four();
+ // comment 3
five();
six();
seven();
--- a/src/new.rs
+++ b/src/new.rs
@@ -1,2 +1,3 @@
pub fn new_file() {
+ created();
}
"}
.to_string(),
expected_patches: vec![
indoc! {"
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,16 +1,16 @@
fn main() {
// comment 1
one();
two();
// comment 4
three();
four();
// comment 3
five();
six();
seven();
eight();
// comment 2
nine();
}
"}
.to_string()
],
rejected_patch: Some(
indoc! {"
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,16 +1,16 @@
fn main() {
// comment 1
one();
two();
// comment 4
three();
four();
// comment 3
five();
six();
seven();
eight();
// comment 2
nine();
}
"}
.to_string()
),
telemetry: None,
human_feedback: Vec::new(),
rating: None,
}
);
}
fn init_test(cx: &mut TestAppContext) {
cx.update(|cx| {
let settings_store = SettingsStore::test(cx);
cx.set_global(settings_store);
zlog::init_test();
let http_client = FakeHttpClient::with_404_response();
let client = Client::new(Arc::new(FakeSystemClock::new()), http_client, cx);
let user_store = cx.new(|cx| UserStore::new(client.clone(), cx));
language_model::init(cx);
RefreshLlmTokenListener::register(client.clone(), user_store.clone(), cx);
EditPredictionStore::global(&client, &user_store, cx);
})
}
}
crates/edit_prediction/src/capture_prediction_context.rs
+85
View File
@@ -0,0 +1,85 @@
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) 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>,
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,
);
Ok(CapturedPredictionContext {
repository_url,
revision,
uncommitted_diff,
buffer_diagnostics,
})
}))
}
crates/edit_prediction/src/data_collection.rs
+1 -12
View File
@@ -7,7 +7,7 @@ use gpui::{Context, Entity, Task};
use language::BufferSnapshot;
use project::{Project, ProjectPath, WorktreeId};
use std::{fmt::Write as _, ops::Range, path::Path, sync::Arc};
use text::{OffsetRangeExt, Point};
use text::Point;
use util::rel_path::RelPath;
pub type UncommittedDiffSnapshot = Vec<(Arc<Path>, BufferSnapshot, BufferDiffSnapshot)>;
@@ -234,17 +234,6 @@ pub(crate) fn compute_uncommitted_diff(snapshot: UncommittedDiffSnapshot) -> Str
uncommitted_diff
}
pub(crate) fn estimate_uncommitted_diff_byte_size(snapshot: &UncommittedDiffSnapshot) -> usize {
let mut size = 0;
for (_, buffer_snapshot, diff_snapshot) in snapshot {
for hunk in diff_snapshot.hunks(buffer_snapshot) {
size += hunk.diff_base_byte_range.len();
size += hunk.range.to_offset(buffer_snapshot).len();
}
}
size
}
fn row_start_or_max(snapshot: &language::BufferSnapshot, row: u32) -> Point {
if row >= snapshot.max_point().row {
snapshot.max_point()
crates/edit_prediction/src/edit_prediction.rs
+380 -177
View File
@@ -3,8 +3,8 @@ use buffer_diff::BufferDiff;
use client::{Client, EditPredictionUsage, UserStore, global_llm_token};
use cloud_api_client::LlmApiToken;
use cloud_api_types::{
EditPredictionSettledKeptChars, OrganizationId, SubmitEditPredictionFeedbackBody,
SubmitEditPredictionSettledBody,
EditPredictionRecentFile, EditPredictionSettledKeptChars, OrganizationId,
SettledEditPrediction, SettledEditPredictionSampleData, SubmitEditPredictionFeedbackBody,
};
use cloud_llm_client::predict_edits_v3::{
PREDICT_EDITS_MODE_HEADER_NAME, PREDICT_EDITS_REQUEST_ID_HEADER_NAME,
@@ -74,7 +74,6 @@ pub mod cursor_excerpt;
pub mod data_collection;
pub mod example_spec;
pub mod fim;
mod jump_example;
mod license_detection;
pub mod mercury;
pub mod metrics;
@@ -85,7 +84,7 @@ mod prediction;
pub mod udiff;
mod capture_example;
mod capture_prediction_context;
pub mod open_ai_compatible;
mod zed_edit_prediction_delegate;
pub mod zeta;
@@ -93,14 +92,9 @@ 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::uncommitted_diffs_for_events;
use crate::example_spec::ExampleSpec;
use crate::example_spec::RecentFile;
use crate::jump_example::{
JUMP_EXAMPLE_NAVIGATION_COUNT, JumpExampleTrigger, PendingJumpExampleCapture,
PendingJumpExampleCaptureKey,
};
use crate::license_detection::LicenseDetectionWatcher;
use crate::mercury::Mercury;
pub use crate::metrics::{KeptRateResult, compute_kept_rate};
@@ -136,6 +130,10 @@ 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;
@@ -206,7 +204,6 @@ pub struct EditPredictionModelInput {
snapshot: BufferSnapshot,
position: Anchor,
events: Vec<Arc<zeta_prompt::Event>>,
stored_events: Vec<StoredEvent>,
related_files: Vec<RelatedFile>,
mode: PredictEditsMode,
trigger: PredictEditsRequestTrigger,
@@ -358,6 +355,7 @@ fn push_recent_file(files: &mut VecDeque<RecentFile>, mut file: RecentFile) {
struct ProjectState {
events: VecDeque<StoredEvent>,
last_event: Option<LastEvent>,
next_last_event_seq: u64,
recently_viewed_files: VecDeque<RecentFile>,
recently_opened_files: VecDeque<RecentFile>,
registered_buffers: HashMap<gpui::EntityId, RegisteredBuffer>,
@@ -366,8 +364,7 @@ struct ProjectState {
last_edit_source: Option<BufferEditSource>,
next_pending_prediction_id: usize,
pending_predictions: ArrayVec<PendingPrediction, 2, u8>,
pending_jump_example_captures: Vec<PendingJumpExampleCapture>,
starting_jump_example_captures: Vec<PendingJumpExampleCaptureKey>,
pending_prediction_captures: Vec<PendingPredictionCapture>,
debug_tx: Option<mpsc::UnboundedSender<DebugEvent>>,
last_edit_prediction_refresh: Option<(EntityId, Instant)>,
last_jump_prediction_refresh: Option<(EntityId, Instant)>,
@@ -471,22 +468,36 @@ impl ProjectState {
}
fn finalize_last_event(&mut self, cx: &mut Context<EditPredictionStore>) {
let Some(event) = self.last_event.take() else {
return;
};
let Some(event) = event.finalize(&self.license_detection_watchers, cx) else {
let Some(last_event) = self.last_event.take() else {
return;
};
let event = last_event.finalize(&self.license_detection_watchers, cx);
for capture in &mut self.pending_jump_example_captures {
capture.future_events.push(event.event.clone());
for capture in &mut self.pending_prediction_captures {
capture.try_record_future_event(
&last_event,
event.as_ref(),
&self.license_detection_watchers,
cx,
);
}
jump_example::drain_completed_jump_example_captures(self, cx);
let Some(event) = event else {
return;
};
if self.events.len() + 1 >= EVENT_COUNT_MAX {
self.events.pop_front();
}
self.events.push_back(event);
}
fn clear_history(&mut self) {
self.events.clear();
self.last_event.take();
for capture in &mut self.pending_prediction_captures {
capture.sample_data = None;
}
}
}
#[derive(Debug, Clone)]
@@ -588,9 +599,9 @@ impl std::ops::Deref for BufferEditPrediction<'_> {
}
}
#[derive(Clone)]
struct PendingSettledPrediction {
struct PendingPredictionCapture {
request_id: EditPredictionId,
edited_buffer_id: EntityId,
editable_anchor_range: Range<Anchor>,
editable_region_before_prediction: String,
predicted_editable_region: String,
@@ -599,23 +610,104 @@ struct PendingSettledPrediction {
organization_id: Option<OrganizationId>,
can_collect_data: bool,
is_in_open_source_repo: bool,
example: Option<ExampleSpec>,
sample_data: Option<PendingPredictionCaptureSampleData>,
model_version: Option<String>,
enqueued_at: Instant,
last_edit_at: Instant,
e2e_latency: std::time::Duration,
}
struct PendingPredictionCaptureSampleData {
context_task: Task<Result<CapturedPredictionContext>>,
editable_path: Arc<Path>,
editable_offset_range: Range<usize>,
next_edit_cursor_offset: Option<usize>,
future_edit_history_events: Vec<Arc<zeta_prompt::Event>>,
navigation_history: VecDeque<RecentFile>,
edit_events_before_quiescence: u32,
prompt_history_boundary: Option<PromptHistoryBoundary>,
}
/// Marks where the prompt's edit history ended. Sample data may only include
/// content the user produced after this point.
struct PromptHistoryBoundary {
/// The seq of the first event this capture is expected to observe: the
/// event that was pending when the prediction was requested, or the next
/// event to be created if none was pending. Observing a later seq first
/// means events were lost while the prediction request was in flight.
first_event_seq: u64,
/// The prompt's end snapshot within the event that was pending when the
/// prediction was requested, if any. The first observed event is trimmed
/// to its suffix after this snapshot.
snapshot: Option<TextBufferSnapshot>,
}
impl PendingPredictionCapture {
/// Records the project's last event (pending or finalizing) into this
/// sample's future edit history. Returns false if the sample must be
/// dropped because its future history can't be captured accurately.
fn try_record_future_event(
&mut self,
last_event: &LastEvent,
finalized_event: Option<&StoredEvent>,
license_detection_watchers: &HashMap<WorktreeId, Rc<LicenseDetectionWatcher>>,
cx: &App,
) {
let Some(sample) = &mut self.sample_data else {
return;
};
let boundary = sample.prompt_history_boundary.take();
let suffix_snapshot = match &boundary {
Some(boundary) => {
if last_event.seq != boundary.first_event_seq {
// Events were finalized before this capture was enqueued,
// so events are missing from the future history.
self.sample_data.take();
return;
}
boundary.snapshot.as_ref()
}
None => None,
};
let event = match suffix_snapshot {
Some(snapshot) => {
let suffix = last_event
.suffix_after(snapshot)
.and_then(|suffix| suffix.finalize(license_detection_watchers, cx));
let Some(suffix) = suffix else {
return;
};
suffix.event
}
None => match finalized_event {
Some(event) => event.event.clone(),
None => return,
},
};
if !event.in_open_source_repo() {
self.sample_data.take();
return;
}
sample.edit_events_before_quiescence += 1;
if sample.future_edit_history_events.len() < EDIT_PREDICTION_CAPTURE_MAX_FUTURE_EVENTS {
sample.future_edit_history_events.push(event);
}
}
}
struct RegisteredBuffer {
file: Option<Arc<dyn File>>,
snapshot: TextBufferSnapshot,
pending_predictions: Vec<PendingSettledPrediction>,
last_position: Option<Anchor>,
_subscriptions: [gpui::Subscription; 2],
}
#[derive(Clone)]
struct LastEvent {
/// Project-wide monotonic sequence number identifying this event.
seq: u64,
old_snapshot: TextBufferSnapshot,
new_snapshot: TextBufferSnapshot,
old_file: Option<Arc<dyn File>>,
@@ -649,7 +741,7 @@ impl LastEvent {
})
});
let (diff, edit_range) = compute_diff_between_snapshots_in_range(
let (diff, old_range, new_range) = compute_diff_between_snapshots_in_range(
&self.old_snapshot,
&self.new_snapshot,
&self.total_edit_range,
@@ -663,13 +755,15 @@ impl LastEvent {
old_path,
path,
diff,
old_range,
new_range: new_range.clone(),
in_open_source_repo,
predicted: self.predicted,
}),
old_snapshot: self.old_snapshot.clone(),
new_snapshot_version: self.new_snapshot.version.clone(),
total_edit_range: self.new_snapshot.anchor_before(edit_range.start)
..self.new_snapshot.anchor_before(edit_range.end),
total_edit_range: self.new_snapshot.anchor_before(new_range.start)
..self.new_snapshot.anchor_before(new_range.end),
file_context: self.file_context.clone(),
})
}
@@ -680,50 +774,41 @@ impl LastEvent {
return (self.clone(), None);
};
let Some(after) = self.suffix_after(boundary_snapshot) else {
return (self.clone(), None);
};
let total_edit_range_before_pause = self
.total_edit_range_at_last_pause_boundary
.clone()
.unwrap_or_else(|| self.total_edit_range.clone());
let Some(total_edit_range_after_pause) =
compute_total_edit_range_between_snapshots(boundary_snapshot, &self.new_snapshot)
else {
return (self.clone(), None);
};
let latest_edit_range_before_pause = total_edit_range_before_pause.clone();
let latest_edit_range_after_pause = total_edit_range_after_pause.clone();
let before = LastEvent {
old_snapshot: self.old_snapshot.clone(),
new_snapshot: boundary_snapshot.clone(),
old_file: self.old_file.clone(),
new_file: self.new_file.clone(),
latest_edit_range: latest_edit_range_before_pause,
latest_edit_range: total_edit_range_before_pause.clone(),
total_edit_range: total_edit_range_before_pause,
total_edit_range_at_last_pause_boundary: None,
predicted: self.predicted,
snapshot_after_last_editing_pause: None,
last_edit_time: self.last_edit_time,
file_context: self.file_context.clone(),
};
let after = LastEvent {
old_snapshot: boundary_snapshot.clone(),
new_snapshot: self.new_snapshot.clone(),
old_file: self.old_file.clone(),
new_file: self.new_file.clone(),
latest_edit_range: latest_edit_range_after_pause,
total_edit_range: total_edit_range_after_pause,
total_edit_range_at_last_pause_boundary: None,
predicted: self.predicted,
snapshot_after_last_editing_pause: None,
last_edit_time: self.last_edit_time,
file_context: self.file_context.clone(),
..self.clone()
};
(before, Some(after))
}
/// The portion of this event that happened after `boundary_snapshot`, or
/// None if the buffer hasn't changed since.
pub fn suffix_after(&self, boundary_snapshot: &TextBufferSnapshot) -> Option<LastEvent> {
let total_edit_range =
compute_total_edit_range_between_snapshots(boundary_snapshot, &self.new_snapshot)?;
Some(LastEvent {
old_snapshot: boundary_snapshot.clone(),
latest_edit_range: total_edit_range.clone(),
total_edit_range,
total_edit_range_at_last_pause_boundary: None,
snapshot_after_last_editing_pause: None,
..self.clone()
})
}
}
fn compute_total_edit_range_between_snapshots(
@@ -791,12 +876,16 @@ fn compute_diff_between_snapshots_in_range(
old_snapshot: &TextBufferSnapshot,
new_snapshot: &TextBufferSnapshot,
total_edit_range: &Range<Anchor>,
) -> Option<(String, Range<Point>)> {
let new_start_point = total_edit_range.start.to_point(new_snapshot);
let new_end_point = total_edit_range.end.to_point(new_snapshot);
) -> Option<(String, Range<usize>, Range<usize>)> {
let new_start_offset = total_edit_range.start.to_offset(new_snapshot);
let new_end_offset = total_edit_range.end.to_offset(new_snapshot);
let new_start_point = new_snapshot.offset_to_point(new_start_offset);
let new_end_point = new_snapshot.offset_to_point(new_end_offset);
let old_range = compute_old_range_for_new_range(old_snapshot, new_snapshot, total_edit_range)?;
let old_start_point = old_range.start;
let old_end_point = old_range.end;
let old_start_offset = old_snapshot.point_to_offset(old_start_point);
let old_end_offset = old_snapshot.point_to_offset(old_end_point);
const CONTEXT_LINES: u32 = 3;
@@ -832,7 +921,11 @@ fn compute_diff_between_snapshots_in_range(
new_context_start_row,
);
Some((diff, new_start_point..new_end_point))
Some((
diff,
old_start_offset..old_end_offset,
new_start_offset..new_end_offset,
))
}
pub(crate) fn buffer_path_with_id_fallback(
@@ -1131,15 +1224,13 @@ impl EditPredictionStore {
pub fn clear_history(&mut self) {
for project_state in self.projects.values_mut() {
project_state.events.clear();
project_state.last_event.take();
project_state.clear_history();
}
}
pub fn clear_history_for_project(&mut self, project: &Entity<Project>) {
if let Some(project_state) = self.projects.get_mut(&project.entity_id()) {
project_state.events.clear();
project_state.last_event.take();
project_state.clear_history();
}
}
@@ -1337,6 +1428,7 @@ impl EditPredictionStore {
},
events: VecDeque::new(),
last_event: None,
next_last_event_seq: 0,
recently_viewed_files: VecDeque::new(),
recently_opened_files: VecDeque::new(),
debug_tx: None,
@@ -1346,8 +1438,7 @@ impl EditPredictionStore {
last_edit_source: None,
cancelled_predictions: HashSet::default(),
pending_predictions: ArrayVec::new(),
pending_jump_example_captures: Vec::new(),
starting_jump_example_captures: Vec::new(),
pending_prediction_captures: Vec::new(),
next_pending_prediction_id: 0,
last_edit_prediction_refresh: None,
last_jump_prediction_refresh: None,
@@ -1475,14 +1566,23 @@ impl EditPredictionStore {
path: path.path.as_std_path().into(),
cursor_position,
};
for capture in &mut project_state.pending_jump_example_captures {
capture.navigation_history.push(recent_file.clone());
if capture.navigation_history.len() > JUMP_EXAMPLE_NAVIGATION_COUNT {
capture.navigation_history.remove(0);
let can_collect_navigation = project_state
.license_detection_watchers
.get(&path.worktree_id)
.is_some_and(|watcher| watcher.is_project_open_source());
for capture in &mut project_state.pending_prediction_captures {
if let Some(sample_data) = capture.sample_data.as_mut() {
if can_collect_navigation {
push_recent_file(
&mut sample_data.navigation_history,
recent_file.clone(),
);
} else {
capture.sample_data = None;
}
}
}
push_recent_file(&mut project_state.recently_viewed_files, recent_file);
jump_example::drain_completed_jump_example_captures(project_state, cx);
}
}
project::Event::DiagnosticsUpdated { .. } => {
@@ -1549,7 +1649,6 @@ impl EditPredictionStore {
snapshot,
file,
last_position: None,
pending_predictions: Vec::new(),
_subscriptions: [
cx.subscribe(buffer, {
let project = project.downgrade();
@@ -1617,9 +1716,19 @@ impl EditPredictionStore {
return;
};
for pending_prediction in &mut registered_buffer.pending_predictions {
if edit_range.overlaps(&pending_prediction.editable_anchor_range, &new_snapshot) {
pending_prediction.last_edit_at = now;
for pending_capture in &mut project_state.pending_prediction_captures {
if pending_capture.edited_buffer_id == buffer.entity_id()
&& edit_range.overlaps(&pending_capture.editable_anchor_range, &new_snapshot)
{
pending_capture.last_edit_at = now;
if is_local
&& !is_predicted
&& let Some(sample_data) = pending_capture.sample_data.as_mut()
&& sample_data.next_edit_cursor_offset.is_none()
{
sample_data.next_edit_cursor_offset =
Some(edit_range.start.to_offset(&new_snapshot));
}
}
}
@@ -1697,7 +1806,10 @@ impl EditPredictionStore {
file_context
});
let seq = project_state.next_last_event_seq;
project_state.next_last_event_seq += 1;
project_state.last_event = Some(LastEvent {
seq,
old_file,
new_file,
old_snapshot,
@@ -1900,48 +2012,79 @@ impl EditPredictionStore {
let mut oldest_edited_at = None;
let mut ready_predictions = Vec::new();
this.update(cx, |this, _| {
for (_, project_state) in this.projects.iter_mut() {
for (_, registered_buffer) in project_state.registered_buffers.iter_mut() {
let mut pending_index = 0;
while pending_index < registered_buffer.pending_predictions.len() {
let pending_prediction =
&registered_buffer.pending_predictions[pending_index];
let age = now.saturating_duration_since(pending_prediction.enqueued_at);
if age >= EDIT_PREDICTION_SETTLED_TTL {
registered_buffer.pending_predictions.remove(pending_index);
continue;
}
let quiet_for =
now.saturating_duration_since(pending_prediction.last_edit_at);
if quiet_for >= EDIT_PREDICTION_SETTLED_QUIESCENCE {
let pending_prediction =
registered_buffer.pending_predictions.remove(pending_index);
let settled_editable_region = registered_buffer
.snapshot
.text_for_range(
pending_prediction.editable_anchor_range.clone(),
)
.collect::<String>();
ready_predictions
.push((pending_prediction, settled_editable_region));
continue;
}
if oldest_edited_at
.is_none_or(|time| pending_prediction.last_edit_at < time)
{
oldest_edited_at = Some(pending_prediction.last_edit_at);
}
pending_index += 1;
this.update(cx, |this, cx| {
for project_state in this.projects.values_mut() {
let ProjectState {
last_event,
registered_buffers,
license_detection_watchers,
pending_prediction_captures,
..
} = project_state;
let pending_last_event = last_event.as_ref().map(|last_event| {
(
last_event,
last_event.finalize(license_detection_watchers, cx),
)
});
let mut pending_index = 0;
while pending_index < pending_prediction_captures.len() {
let pending_capture = &pending_prediction_captures[pending_index];
let age = now.saturating_duration_since(pending_capture.enqueued_at);
if age >= EDIT_PREDICTION_SETTLED_TTL {
pending_prediction_captures.remove(pending_index);
continue;
}
let quiet_for = now.saturating_duration_since(pending_capture.last_edit_at);
if quiet_for >= EDIT_PREDICTION_SETTLED_QUIESCENCE {
let Some(registered_buffer) =
registered_buffers.get(&pending_capture.edited_buffer_id)
else {
pending_prediction_captures.remove(pending_index);
continue;
};
let editable_offset_range = pending_capture
.editable_anchor_range
.to_offset(&registered_buffer.snapshot);
if editable_offset_range.len()
> EDIT_PREDICTION_SETTLED_MAX_EDITABLE_REGION_BYTES
{
// The prediction was obliterated by a huge edit;
// kept-rate against it would be meaningless and the
// region would blow the body size cap.
pending_prediction_captures.remove(pending_index);
continue;
}
let settled_editable_region = registered_buffer
.snapshot
.text_for_range(editable_offset_range)
.collect::<String>();
let mut pending_capture =
pending_prediction_captures.remove(pending_index);
if let Some((last_event, finalized_event)) = pending_last_event.as_ref()
{
pending_capture.try_record_future_event(
last_event,
finalized_event.as_ref(),
license_detection_watchers,
cx,
);
}
ready_predictions.push((pending_capture, settled_editable_region));
continue;
}
if oldest_edited_at.is_none_or(|time| pending_capture.last_edit_at < time) {
oldest_edited_at = Some(pending_capture.last_edit_at);
}
pending_index += 1;
}
}
});
for (pending_prediction, settled_editable_region) in ready_predictions {
let PendingSettledPrediction {
for (pending_capture, settled_editable_region) in ready_predictions {
let PendingPredictionCapture {
request_id,
editable_region_before_prediction,
predicted_editable_region,
@@ -1950,11 +2093,11 @@ impl EditPredictionStore {
organization_id,
can_collect_data,
is_in_open_source_repo,
example,
sample_data,
model_version,
e2e_latency,
..
} = pending_prediction;
} = pending_capture;
let settled_editable_region_for_metrics = settled_editable_region.clone();
#[cfg(test)]
{
@@ -1978,21 +2121,46 @@ impl EditPredictionStore {
);
let result: anyhow::Result<()> = async {
let settled_editable_region =
can_collect_data.then_some(settled_editable_region);
let example = if can_collect_data {
example.map(serde_json::to_value).transpose()?
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,
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 body = SubmitEditPredictionSettledBody {
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,
@@ -2005,13 +2173,18 @@ impl EditPredictionStore {
kept_rate: kept_rate_result.kept_rate,
recall_rate: kept_rate_result.recall_rate,
},
example,
example: None,
model_version,
e2e_latency_ms: e2e_latency.as_millis(),
e2e_latency_ms: e2e_latency.as_millis().min(u128::from(u64::MAX))
as u64,
};
let json_bytes = serde_json::to_vec(&body)?;
let compressed = zstd::encode_all(&json_bytes[..], 3)?;
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()
@@ -2053,7 +2226,8 @@ impl EditPredictionStore {
edited_buffer_snapshot: &BufferSnapshot,
editable_offset_range: Range<usize>,
edit_preview: &EditPreview,
example: Option<ExampleSpec>,
context_task: Option<Task<Result<CapturedPredictionContext>>>,
prompt_history_boundary: Option<PromptHistoryBoundary>,
model_version: Option<String>,
e2e_latency: std::time::Duration,
cx: &mut Context<Self>,
@@ -2073,12 +2247,12 @@ impl EditPredictionStore {
.current_organization()
.map(|organization| organization.id.clone());
let project_state = this.get_or_init_project(project, cx);
let Some(registered_buffer) = project_state
if !project_state
.registered_buffers
.get_mut(&edited_buffer.entity_id())
else {
.contains_key(&edited_buffer.entity_id())
{
return;
};
}
let editable_region_before_prediction = edited_buffer_snapshot
.text_for_range(editable_offset_range.clone())
@@ -2101,12 +2275,30 @@ impl EditPredictionStore {
),
);
let editable_anchor_range =
edited_buffer_snapshot.anchor_range_inside(editable_offset_range);
edited_buffer_snapshot.anchor_range_inside(editable_offset_range.clone());
let now = cx.background_executor().now();
registered_buffer
.pending_predictions
.push(PendingSettledPrediction {
let sample_data = if can_collect_data
&& let Some(context_task) = context_task
&& let Some(file) = edited_buffer_snapshot.file()
{
Some(PendingPredictionCaptureSampleData {
context_task,
editable_path: file.path().as_std_path().into(),
editable_offset_range,
next_edit_cursor_offset: None,
future_edit_history_events: Vec::new(),
navigation_history: VecDeque::new(),
edit_events_before_quiescence: 0,
prompt_history_boundary,
})
} else {
None
};
project_state
.pending_prediction_captures
.push(PendingPredictionCapture {
request_id,
edited_buffer_id: edited_buffer.entity_id(),
editable_anchor_range,
editable_region_before_prediction,
predicted_editable_region,
@@ -2115,7 +2307,7 @@ impl EditPredictionStore {
organization_id,
can_collect_data,
is_in_open_source_repo,
example,
sample_data,
model_version,
e2e_latency,
enqueued_at: now,
@@ -2749,6 +2941,18 @@ impl EditPredictionStore {
self.get_or_init_project(&project, cx);
let project_state = self.projects.get(&project.entity_id()).unwrap();
let stored_events = project_state.events(cx);
let prompt_history_boundary = Some(PromptHistoryBoundary {
first_event_seq: project_state
.last_event
.as_ref()
.map_or(project_state.next_last_event_seq, |last_event| {
last_event.seq
}),
snapshot: project_state
.last_event
.as_ref()
.map(|last_event| last_event.new_snapshot.clone()),
});
let has_events = !stored_events.is_empty();
let events: Vec<Arc<zeta_prompt::Event>> =
stored_events.iter().map(|e| e.event.clone()).collect();
@@ -2768,15 +2972,28 @@ impl EditPredictionStore {
};
let buffer_id = active_buffer.read(cx).remote_id();
let repo_url = project
let (repository_url, revision) = project
.read(cx)
.git_store()
.read(cx)
.repository_and_path_for_buffer_id(buffer_id, cx)
.and_then(|(repo, _)| repo.read(cx).default_remote_url());
.map(|(repository, _)| {
let snapshot = repository.read(cx).snapshot();
(
snapshot
.remote_origin_url
.clone()
.or_else(|| snapshot.remote_upstream_url.clone()),
snapshot
.head_commit
.as_ref()
.map(|commit| commit.sha.to_string()),
)
})
.unwrap_or_default();
let is_staff_zed_repo = cx.is_staff()
&& repo_url
&& repository_url
.as_ref()
.is_some_and(|url| is_zed_industries_repo(url));
let is_open_source = is_staff_zed_repo
@@ -2790,14 +3007,12 @@ impl EditPredictionStore {
&& is_open_source
&& self.is_data_collection_enabled(cx)
&& matches!(self.edit_prediction_model, EditPredictionModel::Zeta);
let capture_worktree_id = snapshot.file().map(|file| file.worktree_id(cx));
let inputs = EditPredictionModelInput {
project: project.clone(),
buffer: active_buffer,
snapshot,
position,
events,
stored_events: stored_events.clone(),
related_files,
mode,
trigger,
@@ -2809,41 +3024,27 @@ impl EditPredictionStore {
let task = match self.edit_prediction_model {
EditPredictionModel::Zeta => {
let capture_data = if let Some(worktree_id) = capture_worktree_id
&& can_collect_data
{
let uncommitted_diff_snapshot = uncommitted_diffs_for_events(
project.clone(),
worktree_id,
stored_events.clone(),
cx,
)
.shared();
jump_example::try_start_jump_example_capture(
project_state,
uncommitted_diff_snapshot.clone(),
inputs.project.clone(),
inputs.snapshot.clone(),
inputs.position,
match trigger {
PredictEditsRequestTrigger::Diagnostics
| PredictEditsRequestTrigger::DiagnosticNavigation => {
JumpExampleTrigger::Diagnostic
}
_ => JumpExampleTrigger::Prediction,
},
stored_events,
inputs.diagnostic_search_range.clone(),
can_collect_data,
is_open_source,
cx,
);
rand::random_ratio(1, 10).then(|| uncommitted_diff_snapshot)
} else {
None
};
zeta::request_prediction_with_zeta(self, inputs, capture_data, repo_url, cx)
let context_task = can_collect_data
.then(|| {
capture_prediction_context(
inputs.project.clone(),
inputs.buffer.clone(),
inputs.position,
stored_events,
repository_url.clone(),
revision,
cx,
)
})
.flatten();
zeta::request_prediction_with_zeta(
self,
inputs,
context_task,
prompt_history_boundary,
repository_url,
cx,
)
}
EditPredictionModel::Fim { format } => fim::request_prediction(inputs, format, cx),
EditPredictionModel::Mercury => {
@@ -3440,7 +3641,7 @@ fn merge_trailing_events_if_needed(
merge_anchor_ranges(&merged_edit_range, &event.total_edit_range, latest_snapshot);
}
if let Some((diff, edit_range)) = compute_diff_between_snapshots_in_range(
if let Some((diff, old_range, new_range)) = compute_diff_between_snapshots_in_range(
&oldest_snapshot,
newest_snapshot,
&merged_edit_range,
@@ -3456,6 +3657,8 @@ fn merge_trailing_events_if_needed(
old_path: old_path.clone(),
path: path.clone(),
diff,
old_range,
new_range: new_range.clone(),
in_open_source_repo: *in_open_source_repo,
predicted: events.range(merge_start..).all(|event| {
matches!(
@@ -3469,8 +3672,8 @@ fn merge_trailing_events_if_needed(
}),
old_snapshot: oldest_snapshot.clone(),
new_snapshot_version: newest_snapshot.version.clone(),
total_edit_range: newest_snapshot.anchor_before(edit_range.start)
..newest_snapshot.anchor_before(edit_range.end),
total_edit_range: newest_snapshot.anchor_before(new_range.start)
..newest_snapshot.anchor_before(new_range.end),
file_context: oldest_event.file_context.clone(),
},
};
crates/edit_prediction/src/edit_prediction_tests.rs
+431 -48
View File
@@ -3,7 +3,7 @@ use clock::FakeSystemClock;
use clock::ReplicaId;
use cloud_api_types::{
CreateLlmTokenResponse, LlmToken, Organization, OrganizationConfiguration,
OrganizationEditPredictionConfiguration, OrganizationId, SubmitEditPredictionSettledBody,
OrganizationEditPredictionConfiguration, OrganizationId, SettledEditPrediction,
SubmitEditPredictionSettledResponse,
};
use cloud_llm_client::{
@@ -1455,6 +1455,7 @@ async fn test_empty_prediction(cx: &mut TestAppContext) {
let position = snapshot.anchor_before(language::Point::new(1, 3));
ep_store.update(cx, |ep_store, cx| {
ep_store.register_buffer(&buffer, &project, cx);
ep_store.refresh_prediction_from_buffer(
project.clone(),
buffer.clone(),
@@ -1495,7 +1496,7 @@ async fn test_empty_prediction(cx: &mut TestAppContext) {
assert_eq!(
&reject_request.rejections,
&[EditPredictionRejection {
request_id: id,
request_id: id.clone(),
reason: EditPredictionRejectReason::Empty,
was_shown: false,
model_version: Some("zeta2:test-empty".to_string()),
@@ -2686,6 +2687,7 @@ fn test_active_buffer_diagnostics_collection_limits(cx: &mut TestAppContext) {
let text = (0..300)
.map(|row| format!("line {row} has some diagnostic context\n"))
.collect::<String>();
let long_message = "diagnostic message ".repeat(1000);
let buffer = cx.new(|cx| Buffer::local(&text, cx));
buffer.update(cx, |buffer, cx| {
@@ -2695,7 +2697,7 @@ fn test_active_buffer_diagnostics_collection_limits(cx: &mut TestAppContext) {
range: text::PointUtf16::new(150, 0)..text::PointUtf16::new(150, 4),
diagnostic: Diagnostic {
severity: DiagnosticSeverity::ERROR,
message: "long snippet".to_string(),
message: long_message.clone(),
group_id: 1,
is_primary: true,
source_kind: language::DiagnosticSourceKind::Pushed,
@@ -2716,7 +2718,15 @@ fn test_active_buffer_diagnostics_collection_limits(cx: &mut TestAppContext) {
);
assert_eq!(active_buffer_diagnostics.len(), 1);
assert!(active_buffer_diagnostics[0].snippet.len() <= 512 * 3 + 2);
assert!(
active_buffer_diagnostics[0].message.len()
<= crate::zeta::MAX_ACTIVE_BUFFER_DIAGNOSTIC_MESSAGE_TOKENS_TO_COLLECT * 3 + 2
);
assert!(active_buffer_diagnostics[0].message.len() < long_message.len());
assert!(
active_buffer_diagnostics[0].snippet.len()
<= crate::zeta::MAX_ACTIVE_BUFFER_DIAGNOSTIC_SNIPPET_TOKENS_TO_COLLECT * 3 + 2
);
assert!(active_buffer_diagnostics[0].snippet.len() < text.len());
}
@@ -2819,7 +2829,7 @@ struct RequestChannels {
oneshot::Sender<PredictEditsV3Response>,
)>,
reject: mpsc::UnboundedReceiver<(RejectEditPredictionsBody, oneshot::Sender<()>)>,
settled: mpsc::UnboundedReceiver<SubmitEditPredictionSettledBody>,
settled: mpsc::UnboundedReceiver<SettledEditPrediction>,
}
fn init_test_with_fake_client(
@@ -3833,6 +3843,7 @@ async fn test_edit_prediction_settled(cx: &mut TestAppContext) {
&edit_preview_a,
None,
None,
None,
Duration::from_secs(0),
cx,
);
@@ -3900,6 +3911,7 @@ async fn test_edit_prediction_settled(cx: &mut TestAppContext) {
&edit_preview_b,
None,
None,
None,
Duration::from_secs(0),
cx,
);
@@ -3949,19 +3961,40 @@ async fn test_edit_prediction_settled(cx: &mut TestAppContext) {
}
}
#[gpui::test]
async fn test_edit_prediction_settled_omits_body_when_data_collection_is_disabled(
const MIT_LICENSE: &str = indoc! {r#"
MIT License
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
"#};
async fn init_sample_capture_test(
tree: serde_json::Value,
cx: &mut TestAppContext,
) -> (
Entity<EditPredictionStore>,
RequestChannels,
Entity<Project>,
Entity<Buffer>,
) {
let (ep_store, mut requests) = init_test_with_fake_client(cx);
let (ep_store, requests) = init_test_with_fake_client(cx);
let fs = FakeFs::new(cx.executor());
fs.insert_tree(
"/root",
json!({
"foo.md": "sensitive source\n"
}),
)
.await;
fs.insert_tree("/root", tree).await;
let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await;
let buffer = project
.update(cx, |project, cx| {
@@ -3970,52 +4003,404 @@ async fn test_edit_prediction_settled_omits_body_when_data_collection_is_disable
})
.await
.unwrap();
ep_store.update(cx, |ep_store, cx| {
ep_store.register_buffer(&buffer, &project, cx);
});
cx.run_until_parked();
(ep_store, requests, project, buffer)
}
/// Enqueues a settled prediction capture with injected sample data, as the
/// real request path would when data collection is enabled. Returns the
/// editable offset range.
async fn enqueue_sample_capture(
ep_store: &Entity<EditPredictionStore>,
project: &Entity<Project>,
buffer: &Entity<Buffer>,
id: &str,
editable_point_range: Range<Point>,
context: CapturedPredictionContext,
prompt_history_boundary: Option<PromptHistoryBoundary>,
navigation_history: VecDeque<RecentFile>,
cx: &mut TestAppContext,
) -> Range<usize> {
let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot());
let edits: Arc<[(Range<Anchor>, Arc<str>)]> =
cx.update(|cx| to_completion_edits([(0..9, "replacement".into())], &buffer, cx).into());
let editable_offset_range = snapshot.point_to_offset(editable_point_range.start)
..snapshot.point_to_offset(editable_point_range.end);
let empty_edits: Arc<[(Range<Anchor>, Arc<str>)]> = Vec::new().into();
let edit_preview = buffer
.read_with(cx, |buffer, cx| buffer.preview_edits(edits, cx))
.read_with(cx, |buffer, cx| buffer.preview_edits(empty_edits, cx))
.await;
ep_store.update(cx, |ep_store, cx| {
ep_store.enqueue_settled_prediction(
EditPredictionId("prediction-private".into()),
&project,
&buffer,
EditPredictionId(id.to_string().into()),
project,
buffer,
&snapshot,
0..snapshot.len(),
editable_offset_range.clone(),
&edit_preview,
Some(ExampleSpec {
name: "test example".to_string(),
repository_url: "https://example.com/repo".to_string(),
revision: "rev".to_string(),
tags: Vec::new(),
reasoning: None,
uncommitted_diff: String::new(),
recently_opened_files: Vec::new(),
recently_viewed_files: Vec::new(),
uncommitted_diff_contains_edit_history: false,
cursor_path: Path::new("foo.md").into(),
cursor_position: "0".to_string(),
edit_history: "sensitive edit history".to_string(),
expected_patches: vec!["sensitive patch".to_string()],
rejected_patch: None,
telemetry: None,
human_feedback: Vec::new(),
rating: None,
}),
Some("test-model".to_string()),
Duration::from_millis(42),
None,
None,
None,
Duration::from_secs(0),
cx,
);
let pending_capture = ep_store
.projects
.get_mut(&project.entity_id())
.unwrap()
.pending_prediction_captures
.last_mut()
.unwrap();
pending_capture.can_collect_data = true;
pending_capture.sample_data = Some(PendingPredictionCaptureSampleData {
context_task: Task::ready(Ok(context)),
editable_path: snapshot.file().unwrap().path().as_std_path().into(),
editable_offset_range: editable_offset_range.clone(),
next_edit_cursor_offset: None,
future_edit_history_events: Vec::new(),
navigation_history,
edit_events_before_quiescence: 0,
prompt_history_boundary,
});
});
editable_offset_range
}
#[gpui::test]
async fn test_edit_prediction_settled_sends_sample_data_after_quiescence(cx: &mut TestAppContext) {
let (ep_store, mut requests, project, buffer) = init_sample_capture_test(
json!({
"LICENSE": MIT_LICENSE,
"foo.md": (0..60).map(|ix| format!("line {ix}\n")).collect::<String>(),
}),
cx,
)
.await;
buffer.update(cx, |buffer, cx| {
let offset = Point::new(1, 0).to_offset(buffer);
buffer.edit(vec![(offset..offset, "prompted ")], None, cx);
});
cx.run_until_parked();
let boundary = ep_store.update(cx, |ep_store, _cx| {
let project_state = ep_store.projects.get(&project.entity_id()).unwrap();
PromptHistoryBoundary {
first_event_seq: project_state
.last_event
.as_ref()
.map_or(project_state.next_last_event_seq, |last_event| {
last_event.seq
}),
snapshot: project_state
.last_event
.as_ref()
.map(|last_event| last_event.new_snapshot.clone()),
}
});
let editable_offset_range = enqueue_sample_capture(
&ep_store,
&project,
&buffer,
"prediction-sample",
Point::new(2, 0)..Point::new(3, 0),
CapturedPredictionContext {
repository_url: Some("https://example.com/repo.git".to_string()),
revision: Some("abc123".to_string()),
uncommitted_diff: Some("--- a/foo.md\n+++ b/foo.md\n".to_string()),
buffer_diagnostics: vec![zeta_prompt::ActiveBufferDiagnostic {
severity: Some(1),
message: "sample diagnostic".to_string(),
snippet: String::new(),
snippet_buffer_row_range: 0..0,
diagnostic_range_in_snippet: 0..0,
}],
},
Some(boundary),
VecDeque::from([RecentFile {
path: Path::new("foo.md").into(),
cursor_position: Some(3),
}]),
cx,
)
.await;
let expected_next_edit_cursor_offset = editable_offset_range.start;
// A second capture whose user has data collection disabled; its sample
// and settled region must be redacted at send time.
let boundary = ep_store.update(cx, |ep_store, _cx| {
let project_state = ep_store.projects.get(&project.entity_id()).unwrap();
PromptHistoryBoundary {
first_event_seq: project_state
.last_event
.as_ref()
.map_or(project_state.next_last_event_seq, |last_event| {
last_event.seq
}),
snapshot: project_state
.last_event
.as_ref()
.map(|last_event| last_event.new_snapshot.clone()),
}
});
enqueue_sample_capture(
&ep_store,
&project,
&buffer,
"prediction-redacted",
Point::new(2, 0)..Point::new(3, 0),
CapturedPredictionContext {
repository_url: None,
revision: None,
uncommitted_diff: None,
buffer_diagnostics: Vec::new(),
},
Some(boundary),
VecDeque::new(),
cx,
)
.await;
ep_store.update(cx, |ep_store, _cx| {
ep_store
.projects
.get_mut(&project.entity_id())
.unwrap()
.pending_prediction_captures
.last_mut()
.unwrap()
.can_collect_data = false;
});
cx.executor().advance_clock(LAST_CHANGE_GROUPING_TIME * 2);
cx.run_until_parked();
for (ix, row) in [2, 20, 30, 40, 50].into_iter().enumerate() {
buffer.update(cx, |buffer, cx| {
let start = Point::new(row, 0).to_offset(buffer);
let end = Point::new(row, 4).to_offset(buffer);
buffer.edit(vec![(start..end, format!("future {ix}"))], None, cx);
});
cx.run_until_parked();
}
cx.executor()
.advance_clock(EDIT_PREDICTION_SETTLED_QUIESCENCE);
cx.run_until_parked();
let mut settled_by_id = std::collections::HashMap::new();
for _ in 0..2 {
let request = requests
.settled
.next()
.await
.expect("settled request should be sent");
settled_by_id.insert(request.request_id.clone(), request);
}
let redacted_request = settled_by_id.remove("prediction-redacted").unwrap();
assert!(!redacted_request.can_collect_data);
assert_eq!(redacted_request.settled_editable_region, None);
assert_eq!(redacted_request.sample_data, None);
let settled_request = settled_by_id.remove("prediction-sample").unwrap();
let sample_data = settled_request
.sample_data
.expect("sample data should be sent after quiescence");
assert_eq!(
sample_data.repository_url.as_deref(),
Some("https://example.com/repo.git")
);
assert_eq!(sample_data.revision.as_deref(), Some("abc123"));
assert_eq!(
sample_data.uncommitted_diff.as_deref(),
Some("--- a/foo.md\n+++ b/foo.md\n")
);
assert_eq!(sample_data.editable_path.as_ref(), Path::new("foo.md"));
assert_eq!(sample_data.editable_offset_range, editable_offset_range);
assert_eq!(sample_data.buffer_diagnostics.len(), 1);
assert_eq!(sample_data.future_edit_history_events.len(), 4);
assert_eq!(sample_data.navigation_history.len(), 1);
assert_eq!(sample_data.edit_events_before_quiescence, 5);
assert_eq!(
sample_data.next_edit_cursor_offset,
Some(expected_next_edit_cursor_offset)
);
let future_event_diffs = sample_data
.future_edit_history_events
.iter()
.map(|event| match event.as_ref() {
zeta_prompt::Event::BufferChange { diff, .. } => diff.as_str(),
})
.collect::<Vec<_>>();
assert!(future_event_diffs.iter().all(|diff| {
diff.lines()
.all(|line| !line.starts_with("+prompted ") && line != "-line 1")
}));
assert!(
future_event_diffs
.iter()
.any(|diff| diff.contains("future 0"))
);
assert!(
!future_event_diffs
.iter()
.any(|diff| diff.contains("future 4"))
);
}
#[gpui::test]
async fn test_edit_prediction_settled_sample_data_requires_observing_all_events_since_request(
cx: &mut TestAppContext,
) {
let (ep_store, mut requests, project, buffer) = init_sample_capture_test(
json!({
"LICENSE": MIT_LICENSE,
"foo.md": (0..30).map(|ix| format!("line {ix}\n")).collect::<String>(),
}),
cx,
)
.await;
// Two predictions are requested while no event is pending.
let (boundary_observed, boundary_missed) = ep_store.update(cx, |ep_store, _cx| {
let project_state = ep_store.projects.get(&project.entity_id()).unwrap();
let first_event_seq = project_state
.last_event
.as_ref()
.map_or(project_state.next_last_event_seq, |last_event| {
last_event.seq
});
let snapshot = project_state
.last_event
.as_ref()
.map(|last_event| last_event.new_snapshot.clone());
(
PromptHistoryBoundary {
first_event_seq,
snapshot: snapshot.clone(),
},
PromptHistoryBoundary {
first_event_seq,
snapshot,
},
)
});
assert!(boundary_observed.snapshot.is_none());
// The first capture is enqueued immediately, so it observes both
// subsequent events.
enqueue_sample_capture(
&ep_store,
&project,
&buffer,
"prediction-observed",
Point::new(1, 0)..Point::new(2, 0),
CapturedPredictionContext {
repository_url: None,
revision: None,
uncommitted_diff: None,
buffer_diagnostics: Vec::new(),
},
Some(boundary_observed),
VecDeque::new(),
cx,
)
.await;
buffer.update(cx, |buffer, cx| {
let offset = Point::new(1, 0).to_offset(buffer);
buffer.edit(vec![(offset..offset, "first ")], None, cx);
});
cx.run_until_parked();
buffer.update(cx, |buffer, cx| {
let offset = Point::new(20, 0).to_offset(buffer);
buffer.edit(vec![(offset..offset, "second ")], None, cx);
});
cx.run_until_parked();
// The second capture is enqueued only after the first event was
// finalized, so its future history has a gap and it must be dropped.
enqueue_sample_capture(
&ep_store,
&project,
&buffer,
"prediction-missed",
Point::new(1, 0)..Point::new(2, 0),
CapturedPredictionContext {
repository_url: None,
revision: None,
uncommitted_diff: None,
buffer_diagnostics: Vec::new(),
},
Some(boundary_missed),
VecDeque::new(),
cx,
)
.await;
cx.executor()
.advance_clock(EDIT_PREDICTION_SETTLED_QUIESCENCE);
cx.run_until_parked();
let mut settled_by_id = std::collections::HashMap::new();
for _ in 0..2 {
let request = requests
.settled
.next()
.await
.expect("settled request should be sent");
settled_by_id.insert(request.request_id.clone(), request);
}
let observed_request = settled_by_id.remove("prediction-observed").unwrap();
let sample_data = observed_request
.sample_data
.expect("sample data should be sent");
assert_eq!(sample_data.edit_events_before_quiescence, 2);
assert_eq!(sample_data.future_edit_history_events.len(), 2);
let missed_request = settled_by_id.remove("prediction-missed").unwrap();
assert_eq!(missed_request.sample_data, None);
}
#[gpui::test]
async fn test_edit_prediction_settled_drops_future_events_when_their_oss_status_is_unknown(
cx: &mut TestAppContext,
) {
let (ep_store, mut requests, project, buffer) = init_sample_capture_test(
json!({ "foo.md": (0..30).map(|ix| format!("line {ix}\n")).collect::<String>() }),
cx,
)
.await;
enqueue_sample_capture(
&ep_store,
&project,
&buffer,
"prediction-non-oss",
Point::new(1, 0)..Point::new(2, 0),
CapturedPredictionContext {
repository_url: None,
revision: None,
uncommitted_diff: None,
buffer_diagnostics: Vec::new(),
},
None,
VecDeque::new(),
cx,
)
.await;
// There is no LICENSE file, so this event's open-source status is unknown
// and the sample must be dropped.
buffer.update(cx, |buffer, cx| {
let offset = Point::new(20, 0).to_offset(buffer);
buffer.edit(vec![(offset..offset, "outside ")], None, cx);
});
cx.run_until_parked();
cx.executor()
.advance_clock(EDIT_PREDICTION_SETTLED_QUIESCENCE);
cx.run_until_parked();
@@ -4025,9 +4410,7 @@ async fn test_edit_prediction_settled_omits_body_when_data_collection_is_disable
.next()
.await
.expect("settled request should be sent");
assert!(!settled_request.can_collect_data);
assert_eq!(settled_request.settled_editable_region, None);
assert_eq!(settled_request.example, None);
assert_eq!(settled_request.sample_data, None);
}
#[gpui::test]
crates/edit_prediction/src/jump_example.rs
-372
View File
@@ -1,372 +0,0 @@
use std::{
ops::Range,
sync::Arc,
time::{Duration, Instant},
};
use anyhow::{Context as _, Result};
pub use cloud_api_types::JumpExampleTrigger;
use cloud_api_types::{
JumpExampleRecentFile, SubmitEditPredictionJumpExampleBody,
SubmitEditPredictionJumpExampleResponse,
};
use futures::future::Shared;
use gpui::{AppContext as _, AsyncApp, Context, Entity, Task, TaskExt as _, WeakEntity};
use language::{BufferSnapshot, File, Point};
use project::{Project, WorktreeId};
use release_channel::AppVersion;
use text::ToPoint as _;
use util::rel_path::RelPath;
use crate::{
EditPredictionStore, ProjectState, StoredEvent,
data_collection::{
UncommittedDiffResult, compute_cursor_excerpt, compute_uncommitted_diff,
estimate_uncommitted_diff_byte_size, format_cursor_excerpt,
},
example_spec::RecentFile,
zeta,
};
pub const JUMP_EXAMPLE_MAX_PENDING_CAPTURE_COUNT: usize = 10;
pub const JUMP_EXAMPLE_FUTURE_EVENT_COUNT: usize = 2;
pub const JUMP_EXAMPLE_TTL: Duration = Duration::from_secs(60 * 2);
pub const JUMP_EXAMPLE_NAVIGATION_COUNT: usize = 20;
pub const JUMP_EXAMPLE_MAX_UNCOMMITTED_DIFF_SIZE: usize = 64 * 1024;
pub struct PendingJumpExampleCapture {
key: PendingJumpExampleCaptureKey,
trigger: JumpExampleTrigger,
file: Arc<dyn File>,
edit_history: Vec<Arc<zeta_prompt::Event>>,
recently_opened_files: Vec<RecentFile>,
recently_viewed_files: Vec<RecentFile>,
worktree_root_name: String,
cursor_position: String,
started_at: Instant,
uncommitted_diff: Option<String>,
pub future_events: Vec<Arc<zeta_prompt::Event>>,
pub navigation_history: Vec<RecentFile>,
diagnostics: Vec<zeta_prompt::ActiveBufferDiagnostic>,
repository_url: Option<String>,
revision: Option<String>,
can_collect_data: bool,
is_in_open_source_repo: bool,
}
#[derive(Eq, PartialEq, Hash, Clone)]
pub struct PendingJumpExampleCaptureKey {
worktree_id: WorktreeId,
file_path: Arc<RelPath>,
row_bucket: u32,
}
pub fn try_start_jump_example_capture(
project_state: &ProjectState,
uncommitted_diffs: Shared<Task<UncommittedDiffResult>>,
project: Entity<Project>,
snapshot: BufferSnapshot,
position: language::Anchor,
trigger: JumpExampleTrigger,
stored_events: Vec<StoredEvent>,
diagnostic_search_range: Range<Point>,
can_collect_data: bool,
is_in_open_source_repo: bool,
cx: &mut Context<EditPredictionStore>,
) {
let Some(file) = snapshot.file().cloned() else {
return;
};
let example_key = PendingJumpExampleCaptureKey {
worktree_id: file.worktree_id(cx),
file_path: file.path().clone(),
row_bucket: position.to_point(&snapshot).row / 10,
};
let should_capture_example = project_state.pending_jump_example_captures.len()
< JUMP_EXAMPLE_MAX_PENDING_CAPTURE_COUNT
&& !project_state
.starting_jump_example_captures
.contains(&example_key)
&& !project_state
.pending_jump_example_captures
.iter()
.any(|capture| &capture.key == &example_key);
if !should_capture_example {
return;
}
let _project = project.clone();
let _example_key = example_key.clone();
let task = cx.spawn(async move |ep_store, cx| {
let project = _project;
let example_key = _example_key;
let Some(ep_store) = ep_store.upgrade() else {
return anyhow::Ok(());
};
ep_store.update(cx, |ep_store, cx| {
let project_state = ep_store.get_or_init_project(&project, cx);
project_state
.starting_jump_example_captures
.push(example_key.clone());
});
let (repository, worktree) = project.read_with(cx, |project, cx| {
let repository = project.active_repository(cx);
let worktree_id = file.worktree_id(cx);
let worktree = project.worktree_for_id(worktree_id, cx);
(repository, worktree)
});
let Some(worktree) = worktree else {
return Ok(());
};
let diagnostics = zeta::active_buffer_diagnostics(
&snapshot,
diagnostic_search_range.clone(),
position.to_point(&snapshot).row,
100,
);
let uncommitted_diff = 'uncommitted_diff: {
if repository.is_none() {
break 'uncommitted_diff None;
}
let uncommitted_diff_snapshot = uncommitted_diffs
.await
.map_err(|error| anyhow::anyhow!("{error:?}"))
.context("failed to capture uncommitted diff")?;
let estimated_byte_size =
estimate_uncommitted_diff_byte_size(&uncommitted_diff_snapshot);
if estimated_byte_size > JUMP_EXAMPLE_MAX_UNCOMMITTED_DIFF_SIZE {
break 'uncommitted_diff None;
}
let uncommitted_diff = cx
.background_executor()
.spawn(async move { compute_uncommitted_diff(uncommitted_diff_snapshot) })
.await;
if uncommitted_diff.len() > JUMP_EXAMPLE_MAX_UNCOMMITTED_DIFF_SIZE {
break 'uncommitted_diff None;
}
Some(uncommitted_diff)
};
let edit_history = stored_events
.iter()
.map(|e| e.event.clone())
.collect::<Vec<_>>();
let (repository_url, revision) = if let Some(repository) = &repository {
repository.read_with(cx, |repository, _| {
let snapshot = repository.snapshot();
(
snapshot
.remote_origin_url
.clone()
.or_else(|| snapshot.remote_upstream_url.clone()),
snapshot
.head_commit
.as_ref()
.map(|commit| commit.sha.to_string()),
)
})
} else {
(None, None)
};
let line_comment_prefix = snapshot
.language()
.and_then(|language| language.config().line_comments.first())
.map(|prefix| prefix.to_string())
.unwrap_or_default();
let (cursor_excerpt, cursor_offset_in_excerpt, _) = cx
.background_executor()
.spawn(async move { compute_cursor_excerpt(&snapshot, position) })
.await;
let cursor_position = format_cursor_excerpt(
&cursor_excerpt,
cursor_offset_in_excerpt,
&line_comment_prefix,
);
let now = cx.background_executor().now();
ep_store.update(cx, |ep_store, cx| {
let recently_opened_files = ep_store.recently_opened_files_for_project(&project);
let recently_viewed_files = ep_store.recently_viewed_files_for_project(&project);
let project_state = ep_store.get_or_init_project(&project, cx);
project_state
.pending_jump_example_captures
.push(PendingJumpExampleCapture {
key: example_key,
trigger,
file,
uncommitted_diff,
edit_history,
recently_opened_files,
recently_viewed_files,
repository_url,
revision,
diagnostics,
worktree_root_name: worktree.read(cx).root_name_str().to_owned(),
cursor_position,
started_at: now,
future_events: Vec::new(),
navigation_history: Vec::new(),
is_in_open_source_repo,
can_collect_data,
});
drain_completed_jump_example_captures(project_state, cx);
});
Ok(())
});
cx.spawn(async move |ep_store, cx| {
let result = task.await;
ep_store
.update(cx, |ep_store, cx| {
ep_store
.get_or_init_project(&project, cx)
.starting_jump_example_captures
.retain(|key| key != &example_key);
})
.ok();
result
})
.detach_and_log_err(cx);
}
pub fn drain_completed_jump_example_captures(
project_state: &mut ProjectState,
cx: &mut Context<EditPredictionStore>,
) {
let now = cx.background_executor().now();
let mut capture_index = 0;
while capture_index < project_state.pending_jump_example_captures.len() {
let capture = &project_state.pending_jump_example_captures[capture_index];
let finished = capture.future_events.len() >= JUMP_EXAMPLE_FUTURE_EVENT_COUNT
|| now.saturating_duration_since(capture.started_at) >= JUMP_EXAMPLE_TTL;
if !finished {
capture_index += 1;
continue;
}
let capture = project_state
.pending_jump_example_captures
.remove(capture_index);
cx.spawn(async move |this, cx| {
let result = submit_jump_example_capture_task(this, capture, cx).await;
if let Err(error) = result {
log::error!("failed to submit jump opportunity capture: {error:?}");
}
})
.detach();
}
}
fn submit_jump_example_capture_task(
this: WeakEntity<EditPredictionStore>,
capture: PendingJumpExampleCapture,
cx: &mut AsyncApp,
) -> Task<Result<()>> {
let Some((organization_id, client, llm_token, app_version)) = this
.update(cx, |this, cx| {
(
this.user_store
.read(cx)
.current_organization()
.map(|organization| organization.id.clone()),
this.client.clone(),
this.llm_token.clone(),
AppVersion::global(cx),
)
})
.ok()
else {
return Task::ready(Ok(()));
};
cx.background_spawn(async move {
let PendingJumpExampleCapture {
key: _,
trigger,
file,
edit_history,
recently_opened_files,
recently_viewed_files,
worktree_root_name,
cursor_position,
started_at: _,
uncommitted_diff,
future_events,
navigation_history,
diagnostics,
repository_url,
revision,
is_in_open_source_repo,
can_collect_data,
} = capture;
let future_edit_history = render_jump_example_events(&future_events, &worktree_root_name);
let cursor_path = file.path().as_std_path().into();
let example = SubmitEditPredictionJumpExampleBody {
request_id: uuid::Uuid::new_v4(),
trigger,
repository_url,
revision,
uncommitted_diff,
recently_opened_files: jump_example_recent_files(recently_opened_files),
recently_viewed_files: jump_example_recent_files(recently_viewed_files),
cursor_path,
cursor_position,
edit_history,
diagnostics,
future_edit_history,
navigation_history: jump_example_recent_files(navigation_history),
is_in_open_source_repo,
can_collect_data,
};
let json_bytes = serde_json::to_vec(&example)?;
let compressed = zstd::encode_all(&json_bytes[..], 3)?;
let url = client
.http_client()
.build_zed_llm_url("/predict_edits/jump_example", &[])?;
EditPredictionStore::send_api_request::<SubmitEditPredictionJumpExampleResponse>(
|builder| {
Ok(builder
.uri(url.as_ref())
.header("Content-Encoding", "zstd")
.body(compressed.clone().into())?)
},
client,
llm_token,
organization_id,
app_version,
)
.await?;
Ok(())
})
}
fn jump_example_recent_files(files: Vec<RecentFile>) -> Vec<JumpExampleRecentFile> {
files
.into_iter()
.map(|file| JumpExampleRecentFile {
path: file.path,
cursor_position: file.cursor_position,
})
.collect()
}
fn render_jump_example_events(events: &[Arc<zeta_prompt::Event>], root_name: &str) -> String {
let mut edit_history = String::new();
for event in events {
crate::capture_example::write_event_with_relative_paths(
&mut edit_history,
event,
root_name,
);
if !edit_history.ends_with('\n') {
edit_history.push('\n');
}
}
edit_history
}
crates/edit_prediction/src/zeta.rs
+11 -52
View File
@@ -1,15 +1,15 @@
use crate::{
CloudRequestTimeoutError, CurrentEditPrediction, DebugEvent, EditPredictionFinishedDebugEvent,
EditPredictionId, EditPredictionModelInput, EditPredictionStartedDebugEvent,
EditPredictionStore, ZedUpdateRequiredError, buffer_path_with_id_fallback,
EditPredictionStore, PromptHistoryBoundary, ZedUpdateRequiredError,
buffer_path_with_id_fallback,
capture_prediction_context::CapturedPredictionContext,
cursor_excerpt::{self, compute_cursor_excerpt, compute_syntax_ranges},
data_collection::UncommittedDiffResult,
prediction::EditPredictionResult,
};
use anyhow::{Context as _, Result};
use anyhow::Result;
use cloud_llm_client::{AcceptEditPredictionBody, predict_edits_v3::RawCompletionRequest};
use edit_prediction_types::PredictedCursorPosition;
use futures::future::Shared;
use gpui::{App, AppContext as _, Entity, Task, TaskExt, WeakEntity, prelude::*};
use language::{
Buffer, BufferSnapshot, DiagnosticSeverity, EditPredictionPromptFormat, OffsetRangeExt as _,
@@ -33,7 +33,7 @@ use crate::open_ai_compatible::{
load_open_ai_compatible_api_key_if_needed, send_custom_server_request,
};
pub fn request_prediction_with_zeta(
pub(crate) fn request_prediction_with_zeta(
store: &mut EditPredictionStore,
EditPredictionModelInput {
buffer,
@@ -41,7 +41,6 @@ pub fn request_prediction_with_zeta(
position,
related_files,
events,
stored_events,
debug_tx,
mode,
trigger,
@@ -51,7 +50,8 @@ pub fn request_prediction_with_zeta(
is_open_source,
..
}: EditPredictionModelInput,
capture_data: Option<Shared<Task<UncommittedDiffResult>>>,
context_task: Option<Task<Result<CapturedPredictionContext>>>,
prompt_history_boundary: Option<PromptHistoryBoundary>,
repo_url: Option<String>,
cx: &mut Context<EditPredictionStore>,
) -> Task<Result<Option<EditPredictionResult>>> {
@@ -413,49 +413,7 @@ pub fn request_prediction_with_zeta(
let editable_range_in_buffer = editable_range_in_buffer.clone();
let edit_preview = prediction.edit_preview.clone();
let model_version = prediction.model_version.clone();
let example_task = capture_data.and_then(|uncommitted_diffs| {
let (recently_opened_files, recently_viewed_files) = this
.read_with(cx, |this, _| {
(
this.recently_opened_files_for_project(&project),
this.recently_viewed_files_for_project(&project),
)
})
.ok()?;
Some(cx.spawn({
let project = project.clone();
let edited_buffer = edited_buffer.clone();
async move |cx| {
let uncommitted_diffs = uncommitted_diffs
.await
.map_err(|error| anyhow::anyhow!("{error:?}"))
.context("failed to capture uncommitted diff")?;
let Some(task) = cx.update(|cx| {
crate::capture_example::capture_example(
project.clone(),
edited_buffer.clone(),
position,
stored_events,
recently_opened_files,
recently_viewed_files,
uncommitted_diffs,
false,
cx,
)
}) else {
return Err(anyhow::anyhow!("failed to capture example"));
};
task.await
}
}))
});
cx.spawn(async move |cx| {
let example_spec = if let Some(task) = example_task {
task.await.ok()
} else {
None
};
weak_this
.update(cx, |this, cx| {
this.enqueue_settled_prediction(
@@ -465,7 +423,8 @@ pub fn request_prediction_with_zeta(
&edited_buffer_snapshot,
editable_range_in_buffer,
&edit_preview,
example_spec,
context_task,
prompt_history_boundary,
model_version,
request_duration,
cx,
@@ -549,8 +508,8 @@ fn handle_api_response<T>(
const ACTIVE_BUFFER_DIAGNOSTIC_ADDITIONAL_CONTEXT_TOKEN_COUNT: usize = 100;
const MAX_ACTIVE_BUFFER_DIAGNOSTICS_TO_COLLECT: usize = 20;
const MAX_ACTIVE_BUFFER_DIAGNOSTIC_MESSAGE_TOKENS_TO_COLLECT: usize = 512;
const MAX_ACTIVE_BUFFER_DIAGNOSTIC_SNIPPET_TOKENS_TO_COLLECT: usize = 512;
pub(crate) const MAX_ACTIVE_BUFFER_DIAGNOSTIC_MESSAGE_TOKENS_TO_COLLECT: usize = 512;
pub(crate) const MAX_ACTIVE_BUFFER_DIAGNOSTIC_SNIPPET_TOKENS_TO_COLLECT: usize = 512;
pub(crate) fn active_buffer_diagnostics(
snapshot: &language::BufferSnapshot,
crates/edit_prediction_cli/src/qa.rs
+1 -2
View File
@@ -87,8 +87,7 @@ pub fn build_prompt(example: &Example) -> Result<String> {
path,
old_path,
diff,
predicted: _,
in_open_source_repo: _,
..
} = event.as_ref();
edit_history.push_str(&format!("--- a{}\n", old_path.display()));
edit_history.push_str(&format!("+++ b{}\n", path.display()));
crates/edit_prediction_metrics/src/reversal.rs
+18
View File
@@ -1181,6 +1181,8 @@ mod tests {
-old
+new"}
.into(),
old_range: 0..1,
new_range: 0..1,
predicted: false,
in_open_source_repo: true,
}),
@@ -1192,6 +1194,8 @@ mod tests {
-a
+b"}
.into(),
old_range: 0..1,
new_range: 0..1,
predicted: false,
in_open_source_repo: true,
}),
@@ -1203,6 +1207,8 @@ mod tests {
-x
+y"}
.into(),
old_range: 0..1,
new_range: 0..1,
predicted: false,
in_open_source_repo: true,
}),
@@ -1931,6 +1937,8 @@ mod tests {
line2
"}
.into(),
old_range: 0..1,
new_range: 0..1,
predicted: false,
in_open_source_repo: false,
})],
@@ -1969,6 +1977,8 @@ mod tests {
line11
"}
.into(),
old_range: 0..1,
new_range: 0..1,
predicted: false,
in_open_source_repo: false,
})],
@@ -2029,6 +2039,8 @@ mod tests {
line2
"}
.into(),
old_range: 0..1,
new_range: 0..1,
predicted: false,
in_open_source_repo: false,
})],
@@ -2066,6 +2078,8 @@ mod tests {
more_wrong
"}
.into(),
old_range: 0..1,
new_range: 0..1,
predicted: false,
in_open_source_repo: false,
})],
@@ -2134,6 +2148,8 @@ mod tests {
line2
"}
.into(),
old_range: 0..1,
new_range: 0..1,
predicted: false,
in_open_source_repo: false,
}),
@@ -2147,6 +2163,8 @@ mod tests {
line2
"}
.into(),
old_range: 0..1,
new_range: 0..1,
predicted: false,
in_open_source_repo: false,
}),
crates/zeta_prompt/src/zeta_prompt.rs
+11 -1
View File
@@ -167,6 +167,10 @@ impl ZetaFormat {
}
}
fn empty_range() -> Range<usize> {
0..0
}
#[derive(Clone, Debug, PartialEq, Hash, Serialize, Deserialize)]
#[serde(tag = "event")]
pub enum Event {
@@ -174,6 +178,10 @@ pub enum Event {
path: Arc<Path>,
old_path: Arc<Path>,
diff: String,
#[serde(default = "empty_range")]
old_range: Range<usize>,
#[serde(default = "empty_range")]
new_range: Range<usize>,
predicted: bool,
in_open_source_repo: bool,
},
@@ -203,7 +211,7 @@ pub fn write_event(prompt: &mut String, event: &Event) {
old_path,
diff,
predicted,
in_open_source_repo: _,
..
} => {
if *predicted {
prompt.push_str("// User accepted prediction:\n");
@@ -5071,6 +5079,8 @@ mod tests {
path: Path::new(path).into(),
old_path: Path::new(path).into(),
diff: diff.to_string(),
old_range: 0..diff.len(),
new_range: 0..diff.len(),
predicted: false,
in_open_source_repo: false,
}