From aa0cfcd190124c99bea2e5609f0557e93fff97af Mon Sep 17 00:00:00 2001 From: Ben Kunkle Date: Fri, 26 Jun 2026 11:47:51 -0400 Subject: [PATCH] ep: Send V4 prediction requests (#59930) # Objective - 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:
Click to view showcase My super cool demos here
--- Release Notes: - N/A or Added/Fixed/Improved ... --- crates/edit_prediction/src/edit_prediction.rs | 118 +++++-- .../src/edit_prediction_tests.rs | 172 +++++++++- crates/edit_prediction/src/mercury.rs | 8 +- crates/edit_prediction/src/udiff.rs | 250 ++++++++++++-- crates/edit_prediction/src/zeta.rs | 314 ++++++++++++++---- 5 files changed, 737 insertions(+), 125 deletions(-) diff --git a/crates/edit_prediction/src/edit_prediction.rs b/crates/edit_prediction/src/edit_prediction.rs index f0520bdf15..cef4a63e33 100644 --- a/crates/edit_prediction/src/edit_prediction.rs +++ b/crates/edit_prediction/src/edit_prediction.rs @@ -13,6 +13,7 @@ use cloud_llm_client::predict_edits_v3::{ PREDICT_EDITS_TRIGGER_HEADER_NAME, PredictEditsMode, PredictEditsV3Request, PredictEditsV3Response, RawCompletionRequest, RawCompletionResponse, }; +use cloud_llm_client::predict_edits_v4::{PredictEditsV4Request, PredictEditsV4Response}; use cloud_llm_client::{ EditPredictionRejectReason, EditPredictionRejection, MAX_EDIT_PREDICTION_REJECTIONS_PER_REQUEST, MINIMUM_REQUIRED_VERSION_HEADER_NAME, @@ -59,7 +60,7 @@ use std::rc::Rc; use text::{AnchorRangeExt, Edit}; use workspace::{AppState, Workspace}; use zeta_prompt::ContextSource; -use zeta_prompt::{Zeta2PromptInput, ZetaFormat}; +use zeta_prompt::{Zeta2PromptInput, Zeta3PromptInput, ZetaFormat}; use std::mem; use std::ops::Range; @@ -190,7 +191,6 @@ pub enum EditPredictionModel { Mercury, } -#[derive(Clone)] pub struct EditPredictionModelInput { project: Entity, buffer: Entity, @@ -198,12 +198,14 @@ pub struct EditPredictionModelInput { position: Anchor, events: Vec>, related_files: Vec, + editable_context: Option>>>, mode: PredictEditsMode, trigger: PredictEditsRequestTrigger, diagnostic_search_range: Range, debug_tx: Option>, can_collect_data: bool, is_open_source: bool, + allow_jump: bool, } #[derive(Debug)] @@ -1107,6 +1109,7 @@ impl EditPredictionStore { let client = self.client.clone(); let llm_token = self.llm_token.clone(); let app_version = AppVersion::global(cx); + let is_jumps_api = cx.has_flag::(); let organization_id = self .user_store .read(cx) @@ -1118,9 +1121,10 @@ impl EditPredictionStore { .background_spawn(async move { let organization_id = organization_id.ok_or_else(|| anyhow!("No organization selected."))?; - let url = client - .http_client() - .build_zed_llm_url("/edit_prediction_experiments", &[])?; + let url = client.http_client().build_zed_llm_url( + "/edit_prediction_experiments", + &[("is_jumps_api", if is_jumps_api { "true" } else { "false" })], + )?; let mut response = client .authenticated_llm_request(&llm_token, organization_id, |token| { Ok(http_client::Request::builder() @@ -2740,23 +2744,27 @@ 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 + let (stored_events, prompt_history_boundary, debug_tx) = { + let project_state = self.projects.get(&project.entity_id()).unwrap(); + ( + project_state.events(cx), + 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()), }), - snapshot: project_state - .last_event - .as_ref() - .map(|last_event| last_event.new_snapshot.clone()), - }); + project_state.debug_tx.clone(), + ) + }; let events: Vec> = stored_events.iter().map(|e| e.event.clone()).collect(); - let debug_tx = project_state.debug_tx.clone(); let snapshot = active_buffer.read(cx).snapshot(); let cursor_point = position.to_point(&snapshot); @@ -2766,6 +2774,7 @@ impl EditPredictionStore { Point::new(diagnostic_search_start, 0)..Point::new(diagnostic_search_end, 0); let related_files = self.context_for_project(&project, cx); + let allow_jump = is_cloud_zeta && cx.has_flag::(); let mode = match all_language_settings(snapshot.file(), cx).edit_predictions_mode() { EditPredictionsMode::Eager => PredictEditsMode::Eager, EditPredictionsMode::Subtle => PredictEditsMode::Subtle, @@ -2807,6 +2816,16 @@ impl EditPredictionStore { && is_open_source && self.is_data_collection_enabled(cx) && matches!(self.edit_prediction_model, EditPredictionModel::Zeta); + let editable_context = allow_jump.then(|| { + self.collect_editable_context( + project.clone(), + active_buffer.clone(), + position, + Vec::new(), + vec![ContextSource::CurrentFile, ContextSource::EditHistory], + cx, + ) + }); let inputs = EditPredictionModelInput { project: project.clone(), buffer: active_buffer, @@ -2814,12 +2833,14 @@ impl EditPredictionStore { position, events, related_files, + editable_context, mode, trigger, diagnostic_search_range, debug_tx, can_collect_data, is_open_source, + allow_jump, }; let task = match self.edit_prediction_model { @@ -2906,11 +2927,62 @@ impl EditPredictionStore { trigger: PredictEditsRequestTrigger, mode: PredictEditsMode, ) -> Result<(PredictEditsV3Response, Option)> { - let url = client - .http_client() - .build_zed_llm_url("/predict_edits/v3", &[])?; - let request = PredictEditsV3Request { input }; + Self::send_predict_edits_request( + "/predict_edits/v3", + request, + preferred_experiment, + client, + llm_token, + organization_id, + app_version, + trigger, + mode, + ) + .await + } + + pub(crate) async fn send_v4_request( + input: Zeta3PromptInput, + preferred_experiment: Option, + client: Arc, + llm_token: LlmApiToken, + organization_id: Option, + app_version: Version, + trigger: PredictEditsRequestTrigger, + mode: PredictEditsMode, + ) -> Result<(PredictEditsV4Response, Option)> { + let request = PredictEditsV4Request { input }; + Self::send_predict_edits_request( + "/predict_edits/v4", + request, + preferred_experiment, + client, + llm_token, + organization_id, + app_version, + trigger, + mode, + ) + .await + } + + async fn send_predict_edits_request( + path: &str, + request: Req, + preferred_experiment: Option, + client: Arc, + llm_token: LlmApiToken, + organization_id: Option, + app_version: Version, + trigger: PredictEditsRequestTrigger, + mode: PredictEditsMode, + ) -> Result<(Res, Option)> + where + Req: serde::Serialize, + Res: serde::de::DeserializeOwned, + { + let url = client.http_client().build_zed_llm_url(path, &[])?; let request_id = uuid::Uuid::new_v4().to_string(); let json_bytes = serde_json::to_vec(&request)?; diff --git a/crates/edit_prediction/src/edit_prediction_tests.rs b/crates/edit_prediction/src/edit_prediction_tests.rs index 9a6c47ffa2..a2a196d135 100644 --- a/crates/edit_prediction/src/edit_prediction_tests.rs +++ b/crates/edit_prediction/src/edit_prediction_tests.rs @@ -10,22 +10,24 @@ use cloud_llm_client::{ EditPredictionRejectReason, EditPredictionRejection, PredictEditsRequestTrigger, RejectEditPredictionsBody, predict_edits_v3::{PredictEditsV3Request, PredictEditsV3Response}, + predict_edits_v4::{PredictEditsV4Request, PredictEditsV4Response}, }; use db::AppDatabase; use edit_prediction_types::EditPredictionRequestTrigger; +use feature_flags::{FeatureFlag as _, FeatureFlagAppExt as _, FeatureFlagsSettings}; use futures::{ AsyncReadExt, FutureExt, StreamExt, channel::{mpsc, oneshot}, }; use gpui::App; use gpui::{ - Entity, TestAppContext, + Entity, TestAppContext, UpdateGlobal, http_client::{FakeHttpClient, Response}, }; use indoc::indoc; use language::{ Anchor, Buffer, Capability, Diagnostic, DiagnosticEntry, DiagnosticSet, DiagnosticSeverity, - Point, + Point, unified_diff_with_offsets, }; use lsp::LanguageServerId; use parking_lot::Mutex; @@ -43,6 +45,7 @@ use uuid::Uuid; use workspace::{AppState, CollaboratorId, MultiWorkspace}; use zeta_prompt::Zeta2PromptInput; +use crate::prediction::EditPredictionInputs; use crate::udiff::apply_diff_to_string; use crate::{ BufferEditPrediction, EDIT_PREDICTION_SETTLED_QUIESCENCE, EditPredictionId, @@ -54,6 +57,7 @@ use super::*; #[gpui::test] async fn test_current_state(cx: &mut TestAppContext) { let (ep_store, mut requests) = init_test_with_fake_client(cx); + cx.update(|cx| set_jumps_feature_flag_override(cx, "on")); let fs = FakeFs::new(cx.executor()); fs.insert_tree( "/root", @@ -92,12 +96,12 @@ async fn test_current_state(cx: &mut TestAppContext) { cx, ) }); - let (request, respond_tx) = requests.predict.next().await.unwrap(); + let (_request, respond_tx) = requests.predict_v4.next().await.unwrap(); respond_tx - .send(model_response( - &request, - indoc! {r" + .send(PredictEditsV4Response { + request_id: Uuid::new_v4().to_string(), + patch: indoc! {r" --- a/root/1.txt +++ b/root/1.txt @@ ... @@ @@ -105,8 +109,10 @@ async fn test_current_state(cx: &mut TestAppContext) { -How +How are you? Bye - "}, - )) + "} + .to_string(), + model_version: None, + }) .unwrap(); cx.run_until_parked(); @@ -2501,11 +2507,8 @@ fn prompt_from_request(request: &PredictEditsV3Request) -> String { .expect("default zeta prompt formatting should succeed in edit prediction tests") } -fn assert_no_predict_request_ready( - requests: &mut mpsc::UnboundedReceiver<( - PredictEditsV3Request, - oneshot::Sender, - )>, +fn assert_no_predict_request_ready( + requests: &mut mpsc::UnboundedReceiver<(Request, oneshot::Sender)>, ) { if requests.next().now_or_never().flatten().is_some() { panic!("Unexpected prediction request while throttled."); @@ -2517,6 +2520,10 @@ struct RequestChannels { PredictEditsV3Request, oneshot::Sender, )>, + predict_v4: mpsc::UnboundedReceiver<( + PredictEditsV4Request, + oneshot::Sender, + )>, reject: mpsc::UnboundedReceiver<(RejectEditPredictionsBody, oneshot::Sender<()>)>, settled: mpsc::UnboundedReceiver, } @@ -2535,6 +2542,7 @@ fn init_test_with_fake_client_and_legacy_data_collection( cx.set_global(AppDatabase::test_new()); let settings_store = SettingsStore::test(cx); cx.set_global(settings_store); + disable_jumps_feature_flag(cx); zlog::init_test(); if let Some(legacy_data_collection_choice) = legacy_data_collection_choice { @@ -2549,6 +2557,7 @@ fn init_test_with_fake_client_and_legacy_data_collection( } let (predict_req_tx, predict_req_rx) = mpsc::unbounded(); + let (predict_v4_req_tx, predict_v4_req_rx) = mpsc::unbounded(); let (reject_req_tx, reject_req_rx) = mpsc::unbounded(); let (settled_req_tx, settled_req_rx) = mpsc::unbounded(); @@ -2562,6 +2571,7 @@ fn init_test_with_fake_client_and_legacy_data_collection( .map(str::to_owned); let mut body = req.into_body(); let predict_req_tx = predict_req_tx.clone(); + let predict_v4_req_tx = predict_v4_req_tx.clone(); let reject_req_tx = reject_req_tx.clone(); let settled_req_tx = settled_req_tx.clone(); async move { @@ -2591,6 +2601,27 @@ fn init_test_with_fake_client_and_legacy_data_collection( } serde_json::to_string(&response).unwrap() } + "/predict_edits/v4" => { + let mut buf = Vec::new(); + body.read_to_end(&mut buf).await.ok(); + let decompressed = zstd::decode_all(&buf[..]).unwrap(); + let req = serde_json::from_slice(&decompressed).unwrap(); + + let (res_tx, res_rx) = oneshot::channel::(); + predict_v4_req_tx.unbounded_send((req, res_tx)).unwrap(); + let response = res_rx.await?; + if response.request_id == REQUEST_TIMEOUT_RESPONSE_ID { + return Ok(Response::builder() + .status(http_client::http::StatusCode::REQUEST_TIMEOUT) + .body( + http_client::http::StatusCode::REQUEST_TIMEOUT + .as_str() + .into(), + ) + .unwrap()); + } + serde_json::to_string(&response).unwrap() + } "/predict_edits/reject" => { let mut buf = Vec::new(); body.read_to_end(&mut buf).await.ok(); @@ -2638,6 +2669,7 @@ fn init_test_with_fake_client_and_legacy_data_collection( user_store, RequestChannels { predict: predict_req_rx, + predict_v4: predict_v4_req_rx, reject: reject_req_rx, settled: settled_req_rx, }, @@ -2867,6 +2899,63 @@ async fn test_edit_prediction_end_of_buffer(cx: &mut TestAppContext) { ); } +#[gpui::test] +async fn test_edit_prediction_v4_end_of_buffer(cx: &mut TestAppContext) { + init_test(cx); + + let fs = FakeFs::new(cx.executor()); + fs.insert_tree( + "/project", + json!({ + "file.txt": "lorem\n" + }), + ) + .await; + let project = Project::test(fs, [path!("/project").as_ref()], cx).await; + let buffer = project + .update(cx, |project, cx| { + let path = project + .find_project_path(path!("/project/file.txt"), cx) + .unwrap(); + project.open_buffer(path, cx) + }) + .await + .unwrap(); + let (ep_store, response) = make_test_ep_store(&project, cx).await; + *response.lock() = "lorem\nipsum\n".to_string(); + + let position = buffer.read_with(cx, |buffer, _| buffer.anchor_before(Point::new(1, 0))); + ep_store.update(cx, |ep_store, cx| { + ep_store.register_project(&project, cx); + ep_store.register_buffer(&buffer, &project, cx); + ep_store.refresh_prediction_from_buffer( + project.clone(), + buffer.clone(), + position, + EditPredictionRequestTrigger::Other, + cx, + ); + }); + cx.run_until_parked(); + + let edits = ep_store.update(cx, |ep_store, cx| { + let prediction = ep_store + .prediction_at(&buffer, None, &project, cx) + .expect("should have prediction"); + let prediction = match prediction { + BufferEditPrediction::Local { prediction } + | BufferEditPrediction::Jump { prediction } => prediction, + }; + assert!(prediction.editable_range.is_some()); + prediction.edits.iter().cloned().collect::>() + }); + buffer.update(cx, |buffer, cx| buffer.edit(edits, None, cx)); + + buffer.read_with(cx, |buffer, _| { + assert_eq!(buffer.text(), "lorem\nipsum\n"); + }); +} + #[gpui::test] async fn test_edit_prediction_no_spurious_trailing_newline(cx: &mut TestAppContext) { // Test that zeta2's newline normalization logic doesn't insert spurious newlines. @@ -3013,6 +3102,26 @@ fn init_test(cx: &mut TestAppContext) { cx.set_global(AppDatabase::test_new()); let settings_store = SettingsStore::test(cx); cx.set_global(settings_store); + disable_jumps_feature_flag(cx); + }); +} + +fn disable_jumps_feature_flag(cx: &mut App) { + SettingsStore::update_global(cx, |store, _| { + store.register_setting::(); + }); + set_jumps_feature_flag_override(cx, "off"); + cx.update_flags(false, vec![]); +} + +fn set_jumps_feature_flag_override(cx: &mut App, value: &str) { + SettingsStore::update_global(cx, |store, cx| { + store.update_user_settings(cx, |content| { + content.feature_flags.get_or_insert_default().insert( + EditPredictionJumpsFeatureFlag::NAME.to_string(), + value.to_string(), + ); + }); }); } @@ -3105,6 +3214,43 @@ async fn make_test_ep_store( ) .unwrap()) } + (Method::POST, "/predict_edits/v4") => { + let mut buf = Vec::new(); + body.read_to_end(&mut buf).await.ok(); + let decompressed = zstd::decode_all(&buf[..]).unwrap(); + let req: PredictEditsV4Request = + serde_json::from_slice(&decompressed).unwrap(); + + next_request_id += 1; + let output = completion_response.lock().clone(); + let file = req + .input + .editable_context + .iter() + .find(|file| file.path == req.input.cursor_path) + .or_else(|| req.input.editable_context.first()) + .expect("V4 requests should include editable context"); + let excerpt = file + .excerpts + .first() + .expect("V4 editable context should include an excerpt"); + let diff = unified_diff_with_offsets( + &excerpt.text, + &output, + excerpt.row_range.start, + excerpt.row_range.start, + ); + let path = file.path.to_string_lossy(); + let response = PredictEditsV4Response { + request_id: format!("request-{next_request_id}"), + patch: format!("--- a/{path}\n+++ b/{path}\n{diff}"), + model_version: None, + }; + Ok(http_client::Response::builder() + .status(200) + .body(serde_json::to_string(&response).unwrap().into()) + .unwrap()) + } _ => Ok(http_client::Response::builder() .status(404) .body("Not Found".to_string().into()) diff --git a/crates/edit_prediction/src/mercury.rs b/crates/edit_prediction/src/mercury.rs index 986d708b11..26e0b5ef24 100644 --- a/crates/edit_prediction/src/mercury.rs +++ b/crates/edit_prediction/src/mercury.rs @@ -1,7 +1,9 @@ use crate::{ - DebugEvent, EditPredictionFinishedDebugEvent, EditPredictionId, EditPredictionInputs, - EditPredictionModelInput, EditPredictionStartedDebugEvent, EditPredictionStore, - open_ai_response::text_from_response, prediction::EditPredictionResult, zeta::compute_edits, + DebugEvent, EditPredictionFinishedDebugEvent, EditPredictionId, EditPredictionModelInput, + EditPredictionStartedDebugEvent, EditPredictionStore, + open_ai_response::text_from_response, + prediction::{EditPredictionInputs, EditPredictionResult}, + zeta::compute_edits, }; use anyhow::{Context as _, Result}; use cloud_llm_client::EditPredictionRejectReason; diff --git a/crates/edit_prediction/src/udiff.rs b/crates/edit_prediction/src/udiff.rs index b2468755a8..8dfefa391e 100644 --- a/crates/edit_prediction/src/udiff.rs +++ b/crates/edit_prediction/src/udiff.rs @@ -2,14 +2,18 @@ use std::{mem, ops::Range, path::Path, path::PathBuf, sync::Arc}; use anyhow::{Context as _, Result, anyhow}; use collections::{HashMap, hash_map::Entry}; +use edit_prediction_types::PredictedCursorPosition; use gpui::{AsyncApp, Entity}; -use language::{Anchor, Buffer, OffsetRangeExt as _, TextBufferSnapshot, text_diff}; +use language::{ + Anchor, Buffer, BufferSnapshot, OffsetRangeExt as _, TextBufferSnapshot, text_diff, +}; use postage::stream::Stream as _; use project::Project; use util::{paths::PathStyle, rel_path::RelPath}; use worktree::Worktree; use zeta_prompt::udiff::{ - DiffEvent, DiffParser, FileStatus, Hunk, disambiguate_by_line_number, find_context_candidates, + DiffEvent, DiffParser, FileStatus, Hunk, INLINE_CURSOR_MARKER, disambiguate_by_line_number, + find_context_candidates, }; pub use zeta_prompt::udiff::{ @@ -30,6 +34,83 @@ impl OpenedBuffers { } } +pub async fn prediction_edits_for_single_file_diff( + diff_str: &str, + project: &Entity, + cx: &mut AsyncApp, +) -> Result< + Option<( + Entity, + BufferSnapshot, + Vec<(Range, Arc)>, + Option, + )>, +> { + let mut diff = DiffParser::new(diff_str); + let mut target_file = None; + let mut edits = Vec::new(); + let mut cursor_position = None; + + while let Some(event) = diff.next()? { + match event { + DiffEvent::Hunk { path, hunk, status } => { + anyhow::ensure!( + status == FileStatus::Modified, + "V4 edit predictions only support modifying existing files" + ); + + let path = path.to_string(); + if let Some((target_path, _, _)) = &target_file { + anyhow::ensure!( + target_path == &path, + "V4 edit predictions only support one file" + ); + } else { + let project_path = project + .update(cx, |project, cx| { + project.find_project_path(Path::new(&path), cx) + }) + .with_context(|| format!("no such path: {path}"))?; + let buffer = project + .update(cx, |project, cx| project.open_buffer(project_path, cx)) + .await?; + let snapshot = buffer.read_with(cx, |buffer, _| buffer.snapshot()); + + target_file = Some((path, buffer, snapshot)); + } + + let (_, _, snapshot) = target_file.as_ref().context("missing target file")?; + for (range, text) in resolve_hunk_edits_in_buffer( + hunk, + snapshot, + &[Anchor::min_max_range_for_buffer(snapshot.remote_id())], + status, + )? { + if let Some(marker_offset) = text.find(INLINE_CURSOR_MARKER) { + cursor_position.get_or_insert_with(|| { + PredictedCursorPosition::new(range.start, marker_offset) + }); + let text = text.replace(INLINE_CURSOR_MARKER, ""); + if range.start != range.end || !text.is_empty() { + edits.push((range, text.into())); + } + } else { + edits.push((range, text)); + } + } + } + DiffEvent::FileEnd { renamed_to } => { + anyhow::ensure!( + renamed_to.is_none(), + "V4 edit predictions do not support renames" + ); + } + } + } + + Ok(target_file.map(|(_, buffer, snapshot)| (buffer, snapshot, edits, cursor_position))) +} + #[must_use] pub async fn apply_diff( diff_str: &str, @@ -243,7 +324,7 @@ fn resolve_hunk_edits_in_buffer( buffer: &TextBufferSnapshot, ranges: &[Range], status: FileStatus, -) -> Result, Arc)>, anyhow::Error> { +) -> Result, Arc)>, anyhow::Error> { let context_offset = if status == FileStatus::Created || hunk.context.is_empty() { 0 } else { @@ -272,22 +353,26 @@ fn resolve_hunk_edits_in_buffer( return Err(anyhow!("Edit range {:?} exceeds buffer length", edit.range)); } - let iter = hunk.edits.into_iter().flat_map(move |edit| { - let old_text = buffer - .text_for_range(context_offset + edit.range.start..context_offset + edit.range.end) - .collect::(); - let edits_within_hunk = language::text_diff(&old_text, &edit.text); - edits_within_hunk - .into_iter() - .map(move |(inner_range, inner_text)| { - ( - buffer.anchor_after(context_offset + edit.range.start + inner_range.start) - ..buffer.anchor_before(context_offset + edit.range.start + inner_range.end), - inner_text, - ) - }) - }); - Ok(iter) + Ok(hunk + .edits + .into_iter() + .flat_map(move |edit| { + let old_text = buffer + .text_for_range(context_offset + edit.range.start..context_offset + edit.range.end) + .collect::(); + let edits_within_hunk = language::text_diff(&old_text, &edit.text); + edits_within_hunk + .into_iter() + .map(move |(inner_range, inner_text)| { + ( + buffer.anchor_after(context_offset + edit.range.start + inner_range.start) + ..buffer + .anchor_before(context_offset + edit.range.start + inner_range.end), + inner_text, + ) + }) + }) + .collect()) } #[cfg(test)] @@ -295,10 +380,12 @@ mod tests { use super::*; use gpui::TestAppContext; use indoc::indoc; + use language::ToOffset as _; use pretty_assertions::assert_eq; use project::{FakeFs, Project}; use serde_json::json; use settings::SettingsStore; + use std::path::Path; use util::path; #[test] @@ -363,6 +450,131 @@ mod tests { assert_eq!(text, "REPLACED"); } + #[gpui::test] + async fn test_prediction_edits_for_single_file_diff_can_target_project_file( + cx: &mut TestAppContext, + ) { + let fs = init_test(cx); + fs.insert_tree( + path!("/root"), + json!({ + "file1": "Hello!\nHow\nBye\n", + "file2": "Hola!\nComo\nAdios\n", + }), + ) + .await; + let project = Project::test(fs, [path!("/root").as_ref()], cx).await; + + let diff = indoc! {r#" + --- a/file2 + +++ b/file2 + @@ ... @@ + Hola! + -Como + +Como estas? + Adios + "#}; + + let (buffer, snapshot, edits, cursor_position) = + prediction_edits_for_single_file_diff(diff, &project, &mut cx.to_async()) + .await + .unwrap() + .unwrap(); + + assert!(cursor_position.is_none()); + buffer.update(cx, |buffer, cx| buffer.edit(edits, None, cx)); + assert_eq!( + snapshot.file().unwrap().path().as_std_path(), + Path::new("file2") + ); + buffer.read_with(cx, |buffer, _cx| { + assert_eq!(buffer.text(), "Hola!\nComo estas?\nAdios\n"); + }); + } + + #[gpui::test] + async fn test_prediction_edits_for_single_file_diff_strips_inline_cursor_marker( + cx: &mut TestAppContext, + ) { + let fs = init_test(cx); + fs.insert_tree( + path!("/root"), + json!({ + "file": "Hello!\nHow\nBye\n", + }), + ) + .await; + let project = Project::test(fs, [path!("/root").as_ref()], cx).await; + + let diff = indoc! {r#" + --- a/file + +++ b/file + @@ ... @@ + Hello! + -How + +How are <|user_cursor|>you? + Bye + "#}; + + let (buffer, snapshot, edits, cursor_position) = + prediction_edits_for_single_file_diff(diff, &project, &mut cx.to_async()) + .await + .unwrap() + .unwrap(); + + assert!( + edits + .iter() + .all(|(_, text)| !text.contains(INLINE_CURSOR_MARKER)) + ); + let cursor_position = cursor_position.unwrap(); + assert_eq!( + cursor_position.anchor.to_offset(&snapshot), + "Hello!\nHow".len() + ); + assert_eq!(cursor_position.offset, " are ".len()); + + buffer.update(cx, |buffer, cx| buffer.edit(edits, None, cx)); + buffer.read_with(cx, |buffer, _cx| { + assert_eq!(buffer.text(), "Hello!\nHow are you?\nBye\n"); + }); + } + + #[gpui::test] + async fn test_prediction_edits_for_single_file_diff_does_not_treat_completed_literal_marker_as_cursor( + cx: &mut TestAppContext, + ) { + let fs = init_test(cx); + fs.insert_tree( + path!("/root"), + json!({ + "file": "text <|user_cursor\n", + }), + ) + .await; + let project = Project::test(fs, [path!("/root").as_ref()], cx).await; + + let diff = indoc! {r#" + --- a/file + +++ b/file + @@ ... @@ + -text <|user_cursor + +text <|user_cursor|> + "#}; + + let (buffer, _, edits, cursor_position) = + prediction_edits_for_single_file_diff(diff, &project, &mut cx.to_async()) + .await + .unwrap() + .unwrap(); + + assert!(cursor_position.is_none()); + buffer.update(cx, |buffer, cx| buffer.edit(edits, None, cx)); + buffer.read_with(cx, |buffer, _cx| { + assert_eq!(buffer.text(), "text <|user_cursor|>\n"); + }); + } + #[gpui::test] async fn test_apply_diff_successful(cx: &mut TestAppContext) { let fs = init_test(cx); diff --git a/crates/edit_prediction/src/zeta.rs b/crates/edit_prediction/src/zeta.rs index cff9bb4d85..009d91c5b5 100644 --- a/crates/edit_prediction/src/zeta.rs +++ b/crates/edit_prediction/src/zeta.rs @@ -6,8 +6,9 @@ use crate::{ cursor_excerpt::{self, compute_cursor_excerpt, compute_syntax_ranges}, data_collection::CapturedPredictionContext, prediction::EditPredictionResult, + udiff::prediction_edits_for_single_file_diff, }; -use anyhow::Result; +use anyhow::{Context as _, Result}; use cloud_llm_client::{AcceptEditPredictionBody, predict_edits_v3::RawCompletionRequest}; use edit_prediction_types::PredictedCursorPosition; use gpui::{App, AppContext as _, Entity, Task, TaskExt, WeakEntity, prelude::*}; @@ -15,6 +16,7 @@ use language::{ Buffer, BufferSnapshot, DiagnosticSeverity, EditPredictionPromptFormat, OffsetRangeExt as _, ToOffset as _, ZetaVersion, language_settings::all_language_settings, text_diff, }; +use project::Project; use release_channel::AppVersion; use text::{Anchor, Bias, Point}; use ui::SharedString; @@ -24,8 +26,9 @@ use workspace::workspace_error::{ErrorAction, ErrorSeverity, WorkspaceError}; use std::{ops::Range, path::Path, sync::Arc}; use zeta_prompt::{ - ParsedOutput, Zeta2PromptInput, ZetaFormat, excerpt_ranges_for_format, format_zeta_prompt, - get_prefill, parse_zeta2_model_output, stop_tokens_for_format, + FilePosition, ParsedOutput, Zeta2PromptInput, Zeta3PromptInput, ZetaFormat, + excerpt_ranges_for_format, format_zeta_prompt, get_prefill, parse_zeta2_model_output, + stop_tokens_for_format, zeta1::{self, EDITABLE_REGION_END_MARKER}, }; @@ -40,6 +43,7 @@ pub(crate) fn request_prediction_with_zeta( snapshot, position, related_files, + editable_context, events, debug_tx, mode, @@ -48,6 +52,7 @@ pub(crate) fn request_prediction_with_zeta( diagnostic_search_range, can_collect_data, is_open_source, + allow_jump, .. }: EditPredictionModelInput, context_task: Option>>, @@ -83,15 +88,25 @@ pub(crate) fn request_prediction_with_zeta( .map(|organization| organization.id.clone()); let app_version = AppVersion::global(cx); - struct Prediction { - prompt_input: Zeta2PromptInput, - buffer: Entity, - snapshot: BufferSnapshot, - edits: Vec<(Range, Arc)>, - cursor_position: Option, - editable_range_in_buffer: Range, + enum PendingPrediction { + Resolved { + inputs: EditPredictionInputs, + buffer: Entity, + snapshot: BufferSnapshot, + edits: Vec<(Range, Arc)>, + cursor_position: Option, + editable_range_in_buffer: Option>, + }, + Patch { + inputs: EditPredictionInputs, + project: Entity, + buffer: Entity, + snapshot: BufferSnapshot, + patch: String, + }, } + let project_for_request = project.clone(); let request_task = cx.background_spawn({ async move { let local_zeta_version = custom_server_settings @@ -119,27 +134,78 @@ pub(crate) fn request_prediction_with_zeta( }) .unwrap_or_default(); - let cursor_offset = position.to_offset(&snapshot); - let (full_context_offset_range, prompt_input) = zeta2_prompt_input( - &snapshot, - related_files, - events, - diagnostic_search_range, - excerpt_path, - cursor_offset, - is_open_source, - can_collect_data, - repo_url, - ); + enum RequestInput { + V3 { + full_context_offset_range: Range, + prompt_input: Zeta2PromptInput, + formatted_prompt: Option, + }, + V4(Zeta3PromptInput), + } - let formatted_prompt = format_zeta_prompt(&prompt_input, zeta_format); + let cursor_offset = position.to_offset(&snapshot); + let request_input = + if allow_jump && custom_server_settings.is_none() && raw_config.is_none() { + let cursor_point = snapshot.offset_to_point(cursor_offset); + let editable_context = editable_context + .context("missing editable context task")? + .await?; + let active_buffer_diagnostics = active_buffer_diagnostics( + &snapshot, + diagnostic_search_range, + cursor_point.row, + ACTIVE_BUFFER_DIAGNOSTIC_ADDITIONAL_CONTEXT_TOKEN_COUNT, + ); + let syntax_ranges = + compute_syntax_ranges(&snapshot, cursor_offset, &(0..snapshot.len())); + + RequestInput::V4(Zeta3PromptInput { + cursor_path: excerpt_path, + cursor_position: FilePosition { + row: cursor_point.row, + column: cursor_point.column, + }, + events, + editable_context, + syntax_ranges, + active_buffer_diagnostics, + in_open_source_repo: is_open_source, + can_collect_data, + repo_url, + }) + } else { + let (full_context_offset_range, prompt_input) = zeta2_prompt_input( + &snapshot, + related_files, + events, + diagnostic_search_range, + excerpt_path, + cursor_offset, + is_open_source, + can_collect_data, + repo_url, + ); + let formatted_prompt = format_zeta_prompt(&prompt_input, zeta_format); + + RequestInput::V3 { + full_context_offset_range, + prompt_input, + formatted_prompt, + } + }; if let Some(debug_tx) = &debug_tx { + let prompt = match &request_input { + RequestInput::V3 { + formatted_prompt, .. + } => formatted_prompt.clone(), + RequestInput::V4(_) => None, + }; debug_tx .unbounded_send(DebugEvent::EditPredictionStarted( EditPredictionStartedDebugEvent { buffer: buffer.downgrade(), - prompt: formatted_prompt.clone(), + prompt, position, }, )) @@ -148,6 +214,45 @@ pub(crate) fn request_prediction_with_zeta( log::trace!("Sending edit prediction request"); + let (full_context_offset_range, prompt_input, formatted_prompt) = match request_input { + RequestInput::V3 { + full_context_offset_range, + prompt_input, + formatted_prompt, + } => (full_context_offset_range, prompt_input, formatted_prompt), + RequestInput::V4(v4_prompt_input) => { + let inputs = EditPredictionInputs::V3(v4_prompt_input.clone()); + let (response, usage) = EditPredictionStore::send_v4_request( + v4_prompt_input, + preferred_experiment.clone(), + client, + llm_token, + organization_id, + app_version, + trigger, + mode, + ) + .await?; + + let request_id = EditPredictionId(response.request_id.into()); + let model_version = response.model_version; + return anyhow::Ok(( + Some(( + request_id, + Some(PendingPrediction::Patch { + inputs, + project: project_for_request, + buffer, + snapshot: snapshot.clone(), + patch: response.patch, + }), + model_version, + )), + usage, + )); + } + }; + let Some((request_id, output, model_version, usage)) = (if let Some(custom_settings) = &custom_server_settings { let max_tokens = custom_settings.max_output_tokens * 4; @@ -272,7 +377,6 @@ pub(crate) fn request_prediction_with_zeta( Some((request_id, output, None, usage)) } else { - // Use V3 endpoint - server handles model/version selection and suffix stripping let (response, usage) = EditPredictionStore::send_v3_request( prompt_input.clone(), preferred_experiment.clone(), @@ -357,14 +461,14 @@ pub(crate) fn request_prediction_with_zeta( }) .unwrap_or_default(); - let prediction = Prediction { - prompt_input, + let prediction = Some(PendingPrediction::Resolved { + inputs: EditPredictionInputs::V2(prompt_input), buffer, snapshot: snapshot.clone(), edits, cursor_position, - editable_range_in_buffer, - }; + editable_range_in_buffer: Some(editable_range_in_buffer), + }); anyhow::Ok((Some((request_id, prediction, model_version)), usage)) } @@ -378,25 +482,95 @@ pub(crate) fn request_prediction_with_zeta( }; let request_duration = cx.background_executor().now() - request_start; - let Prediction { - prompt_input: inputs, - buffer: edited_buffer, - snapshot: edited_buffer_snapshot, + let Some(prediction) = prediction else { + return Ok(None); + }; + + let ( + inputs, + edited_buffer, + edited_buffer_snapshot, edits, cursor_position, editable_range_in_buffer, - } = prediction; + ) = match prediction { + PendingPrediction::Resolved { + inputs, + buffer, + snapshot, + edits, + cursor_position, + editable_range_in_buffer, + } => ( + inputs, + buffer, + snapshot, + edits, + cursor_position, + editable_range_in_buffer, + ), + PendingPrediction::Patch { + inputs, + project, + buffer: fallback_buffer, + snapshot: fallback_snapshot, + patch, + } => { + let Some((buffer, snapshot, edits, cursor_position)) = + prediction_edits_for_single_file_diff(&patch, &project, cx).await? + else { + return Ok(Some( + EditPredictionResult::new( + id, + &fallback_buffer, + &fallback_snapshot, + Arc::new([]), + None, + None, + inputs, + model_version, + trigger, + request_duration, + cx, + ) + .await, + )); + }; + let editable_range_in_buffer = + edits + .iter() + .fold(None, |range: Option>, (edit_range, _)| { + let edit_range = edit_range.start.to_offset(&snapshot) + ..edit_range.end.to_offset(&snapshot); + Some(match range { + Some(range) => { + range.start.min(edit_range.start)..range.end.max(edit_range.end) + } + None => edit_range, + }) + }); + + ( + inputs, + buffer, + snapshot, + edits, + cursor_position, + editable_range_in_buffer, + ) + } + }; - let editable_anchor_range = - edited_buffer_snapshot.anchor_range_inside(editable_range_in_buffer.clone()); let result = EditPredictionResult::new( id, &edited_buffer, &edited_buffer_snapshot, edits.into(), cursor_position, - Some(editable_anchor_range), - EditPredictionInputs::V2(inputs), + editable_range_in_buffer + .clone() + .map(|range| edited_buffer_snapshot.anchor_range_inside(range)), + inputs, model_version, trigger, request_duration, @@ -404,34 +578,35 @@ pub(crate) fn request_prediction_with_zeta( ) .await; - 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(); + if let Some(editable_range_in_buffer) = editable_range_in_buffer.clone() { + 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 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)) }) @@ -594,7 +769,7 @@ pub(crate) fn active_buffer_diagnostics( pub fn zeta2_prompt_input( snapshot: &language::BufferSnapshot, - related_files: Vec, + mut related_files: Vec, events: Vec>, diagnostic_search_range: Range, excerpt_path: Arc, @@ -623,6 +798,11 @@ pub fn zeta2_prompt_input( snapshot.offset_to_point(cursor_offset).row, ACTIVE_BUFFER_DIAGNOSTIC_ADDITIONAL_CONTEXT_TOKEN_COUNT, ); + for file in &mut related_files { + for excerpt in &mut file.excerpts { + excerpt.context_source = zeta_prompt::ContextSource::Lsp; + } + } let prompt_input = zeta_prompt::Zeta2PromptInput { cursor_path: excerpt_path,