diff --git a/crates/edit_prediction/src/edit_prediction.rs b/crates/edit_prediction/src/edit_prediction.rs index ab88dbf2fa6a66..0f2a67c6a8a3aa 100644 --- a/crates/edit_prediction/src/edit_prediction.rs +++ b/crates/edit_prediction/src/edit_prediction.rs @@ -129,6 +129,7 @@ const GIT_CHANGED_FILE_SETS_COMMIT_LIMIT: usize = 100; const LAST_CHANGE_GROUPING_TIME: Duration = Duration::from_secs(1); const ZED_PREDICT_DATA_COLLECTION_CHOICE: &str = "zed_predict_data_collection_choice"; const REJECT_REQUEST_DEBOUNCE: Duration = Duration::from_secs(15); +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); @@ -165,6 +166,7 @@ pub struct EditPredictionStore { update_required: bool, edit_prediction_model: EditPredictionModel, zeta2_raw_config: Option, + request_backoff_until: Option, preferred_experiment: Option, available_experiments: Vec, pub mercury: Mercury, @@ -925,6 +927,7 @@ impl EditPredictionStore { update_required: false, edit_prediction_model: EditPredictionModel::Zeta, zeta2_raw_config: Self::zeta2_raw_config_from_env(), + request_backoff_until: None, preferred_experiment: None, available_experiments: Vec::new(), mercury: Mercury::new(cx), @@ -967,6 +970,27 @@ impl EditPredictionStore { self.zeta2_raw_config.as_ref() } + pub(crate) fn back_off_requests_after_timeout(&mut self, cx: &mut Context) { + self.request_backoff_until = Some(cx.background_executor().now() + REQUEST_TIMEOUT_BACKOFF); + log::info!( + "Backing off edit prediction requests for {:?} after Cloud timeout", + REQUEST_TIMEOUT_BACKOFF + ); + } + + fn request_backoff_active(&mut self, cx: &App) -> bool { + let Some(backoff_until) = self.request_backoff_until else { + return false; + }; + + if cx.background_executor().now() < backoff_until { + true + } else { + self.request_backoff_until = None; + false + } + } + pub fn preferred_experiment(&self) -> Option<&str> { self.preferred_experiment.as_deref() } @@ -2655,6 +2679,13 @@ impl EditPredictionStore { return Task::ready(Ok(None)); } + if is_cloud_zeta && self.request_backoff_active(cx) { + log::debug!( + "Skipping Zeta edit prediction request while backing off after Cloud timeout" + ); + return Task::ready(Ok(None)); + } + self.get_or_init_project(&project, cx); let project_state = self.projects.get(&project.entity_id()).unwrap(); let stored_events = project_state.events(cx); @@ -3029,6 +3060,9 @@ impl EditPredictionStore { let status = response.status(); let mut body = String::new(); response.body_mut().read_to_string(&mut body).await?; + if status == http_client::http::StatusCode::REQUEST_TIMEOUT { + return Err(anyhow::Error::new(CloudRequestTimeoutError)); + } anyhow::bail!("Request failed with status: {status:?}\nBody: {body}"); } } @@ -3358,6 +3392,10 @@ pub struct ZedUpdateRequiredError { minimum_version: Version, } +#[derive(Error, Debug)] +#[error("Cloud request timed out")] +pub(crate) struct CloudRequestTimeoutError; + struct ZedPredictUpsell; fn is_upsell_dismissed(cx: &App) -> bool { diff --git a/crates/edit_prediction/src/edit_prediction_tests.rs b/crates/edit_prediction/src/edit_prediction_tests.rs index 52c6948cdc5a21..7d4a5f54c38df1 100644 --- a/crates/edit_prediction/src/edit_prediction_tests.rs +++ b/crates/edit_prediction/src/edit_prediction_tests.rs @@ -48,6 +48,7 @@ use zeta_prompt::ZetaPromptInput; use crate::{ BufferEditPrediction, EDIT_PREDICTION_SETTLED_QUIESCENCE, EditPredictionId, EditPredictionJumpsFeatureFlag, EditPredictionStore, REJECT_REQUEST_DEBOUNCE, + REQUEST_TIMEOUT_BACKOFF, }; #[gpui::test] @@ -2066,6 +2067,63 @@ async fn test_jump_and_edit_throttles_are_independent(cx: &mut TestAppContext) { cx.run_until_parked(); } +#[gpui::test] +async fn test_cloud_timeout_backs_off_zeta_requests(cx: &mut TestAppContext) { + let (ep_store, mut requests) = init_test_with_fake_client(cx); + let fs = FakeFs::new(cx.executor()); + fs.insert_tree( + "/root", + json!({ + "foo.md": "Hello!\nHow\nBye\n" + }), + ) + .await; + let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await; + + let buffer = project + .update(cx, |project, cx| { + let path = project.find_project_path(path!("root/foo.md"), cx).unwrap(); + project.open_buffer(path, cx) + }) + .await + .unwrap(); + let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot()); + let position = snapshot.anchor_before(language::Point::new(1, 3)); + + ep_store.update(cx, |ep_store, cx| { + ep_store.register_project(&project, cx); + ep_store.register_buffer(&buffer, &project, cx); + }); + + ep_store.update(cx, |ep_store, cx| { + ep_store.refresh_prediction_from_buffer(project.clone(), buffer.clone(), position, cx); + }); + let (_request, respond_tx) = requests.predict.next().await.unwrap(); + respond_tx.send(request_timeout_response()).unwrap(); + cx.run_until_parked(); + + ep_store.update(cx, |ep_store, cx| { + ep_store.refresh_prediction_from_buffer(project.clone(), buffer.clone(), position, cx); + }); + cx.background_executor + .advance_clock(EditPredictionStore::THROTTLE_TIMEOUT); + cx.background_executor.run_until_parked(); + cx.run_until_parked(); + assert_no_predict_request_ready(&mut requests.predict); + + cx.background_executor + .advance_clock(REQUEST_TIMEOUT_BACKOFF); + cx.background_executor.run_until_parked(); + cx.run_until_parked(); + + ep_store.update(cx, |ep_store, cx| { + ep_store.refresh_prediction_from_buffer(project.clone(), buffer.clone(), position, cx); + }); + let (_request, respond_tx) = requests.predict.next().await.unwrap(); + respond_tx.send(empty_response()).unwrap(); + cx.run_until_parked(); +} + #[gpui::test] async fn test_same_frame_duplicate_requests_deduplicated(cx: &mut TestAppContext) { let (ep_store, mut requests) = init_test_with_fake_client(cx); @@ -2546,6 +2604,18 @@ fn empty_response() -> PredictEditsV3Response { } } +const REQUEST_TIMEOUT_RESPONSE_ID: &str = "__request_timeout__"; + +fn request_timeout_response() -> PredictEditsV3Response { + PredictEditsV3Response { + request_id: REQUEST_TIMEOUT_RESPONSE_ID.to_string(), + editable_range: 0..0, + output: String::new(), + cursor_offset: None, + model_version: None, + } +} + fn prompt_from_request(request: &PredictEditsV3Request) -> String { zeta_prompt::format_zeta_prompt(&request.input, zeta_prompt::ZetaFormat::default()) .expect("default zeta prompt formatting should succeed in edit prediction tests") @@ -2659,9 +2729,20 @@ fn init_test_with_fake_client_and_legacy_data_collection( let decompressed = zstd::decode_all(&buf[..]).unwrap(); let req = serde_json::from_slice(&decompressed).unwrap(); - let (res_tx, res_rx) = oneshot::channel(); + let (res_tx, res_rx) = oneshot::channel::(); predict_req_tx.unbounded_send((req, res_tx)).unwrap(); - serde_json::to_string(&res_rx.await?).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(); diff --git a/crates/edit_prediction/src/zeta.rs b/crates/edit_prediction/src/zeta.rs index ee2bcd62f04aa6..72463392dad977 100644 --- a/crates/edit_prediction/src/zeta.rs +++ b/crates/edit_prediction/src/zeta.rs @@ -1,7 +1,7 @@ use crate::{ - CurrentEditPrediction, DebugEvent, EditPredictionFinishedDebugEvent, EditPredictionId, - EditPredictionModelInput, EditPredictionStartedDebugEvent, EditPredictionStore, - ZedUpdateRequiredError, buffer_path_with_id_fallback, + CloudRequestTimeoutError, CurrentEditPrediction, DebugEvent, EditPredictionFinishedDebugEvent, + EditPredictionId, EditPredictionModelInput, EditPredictionStartedDebugEvent, + EditPredictionStore, ZedUpdateRequiredError, buffer_path_with_id_fallback, cursor_excerpt::{self, compute_cursor_excerpt, compute_syntax_ranges}, data_collection::UncommittedDiffResult, prediction::EditPredictionResult, @@ -491,6 +491,11 @@ fn handle_api_response( Ok(data) } Err(err) => { + if err.is::() { + this.update(cx, |this, cx| this.back_off_requests_after_timeout(cx)) + .ok(); + } + if err.is::() { cx.update(|cx| { this.update(cx, |this, _cx| {