diff --git a/assets/keymaps/default-macos.json b/assets/keymaps/default-macos.json index 6ea3aa2182..f7b43cc2bf 100644 --- a/assets/keymaps/default-macos.json +++ b/assets/keymaps/default-macos.json @@ -66,6 +66,7 @@ "cmd-v": "editor::Paste", "cmd-z": "editor::Undo", "cmd-shift-z": "editor::Redo", + "ctrl-shift-z": "zeta::RateCompletions", "up": "editor::MoveUp", "ctrl-up": "editor::MoveToStartOfParagraph", "pageup": "editor::MovePageUp", @@ -788,5 +789,25 @@ "ctrl-k left": "pane::SplitLeft", "ctrl-k right": "pane::SplitRight" } + }, + { + "context": "RateCompletionModal", + "use_key_equivalents": true, + "bindings": { + "cmd-enter": "zeta::ThumbsUp", + "cmd-delete": "zeta::ThumbsDown", + "shift-down": "zeta::NextEdit", + "shift-up": "zeta::PreviousEdit", + "space": "zeta::PreviewCompletion" + } + }, + { + "context": "RateCompletionModal > Editor", + "use_key_equivalents": true, + "bindings": { + "escape": "zeta::FocusCompletions", + "cmd-shift-enter": "zeta::ThumbsUpActiveCompletion", + "cmd-shift-backspace": "zeta::ThumbsDownActiveCompletion" + } } ] diff --git a/crates/inline_completion_button/src/inline_completion_button.rs b/crates/inline_completion_button/src/inline_completion_button.rs index 06664f403c..f5a129de1b 100644 --- a/crates/inline_completion_button/src/inline_completion_button.rs +++ b/crates/inline_completion_button/src/inline_completion_button.rs @@ -204,7 +204,7 @@ impl Render for InlineCompletionButton { } div().child( - Button::new("zeta", "Zeta") + Button::new("zeta", "ζ") .label_size(LabelSize::Small) .on_click(cx.listener(|this, _, cx| { if let Some(workspace) = this.workspace.upgrade() { diff --git a/crates/ui/src/components/list/list_item.rs b/crates/ui/src/components/list/list_item.rs index c65832d3e9..bf216649e7 100644 --- a/crates/ui/src/components/list/list_item.rs +++ b/crates/ui/src/components/list/list_item.rs @@ -39,6 +39,7 @@ pub struct ListItem { children: SmallVec<[AnyElement; 2]>, selectable: bool, overflow_x: bool, + focused: Option, } impl ListItem { @@ -62,6 +63,7 @@ impl ListItem { children: SmallVec::new(), selectable: true, overflow_x: false, + focused: None, } } @@ -140,6 +142,11 @@ impl ListItem { self.overflow_x = true; self } + + pub fn focused(mut self, focused: bool) -> Self { + self.focused = Some(focused); + self + } } impl Disableable for ListItem { @@ -177,9 +184,14 @@ impl RenderOnce for ListItem { this // TODO: Add focus state // .when(self.state == InteractionState::Focused, |this| { - // this.border_1() - // .border_color(cx.theme().colors().border_focused) - // }) + .when_some(self.focused, |this, focused| { + if focused { + this.border_1() + .border_color(cx.theme().colors().border_focused) + } else { + this.border_1() + } + }) .when(self.selectable, |this| { this.hover(|style| style.bg(cx.theme().colors().ghost_element_hover)) .active(|style| style.bg(cx.theme().colors().ghost_element_active)) @@ -204,10 +216,15 @@ impl RenderOnce for ListItem { .when(self.inset && !self.disabled, |this| { this // TODO: Add focus state - // .when(self.state == InteractionState::Focused, |this| { - // this.border_1() - // .border_color(cx.theme().colors().border_focused) - // }) + //.when(self.state == InteractionState::Focused, |this| { + .when_some(self.focused, |this, focused| { + if focused { + this.border_1() + .border_color(cx.theme().colors().border_focused) + } else { + this.border_1() + } + }) .when(self.selectable, |this| { this.hover(|style| { style.bg(cx.theme().colors().ghost_element_hover) diff --git a/crates/zed/src/main.rs b/crates/zed/src/main.rs index 3e2ec18f1f..5251364de4 100644 --- a/crates/zed/src/main.rs +++ b/crates/zed/src/main.rs @@ -463,6 +463,7 @@ fn main() { welcome::init(cx); settings_ui::init(cx); extensions_ui::init(cx); + zeta::init(cx); cx.observe_global::({ let languages = app_state.languages.clone(); diff --git a/crates/zed/src/zed/inline_completion_registry.rs b/crates/zed/src/zed/inline_completion_registry.rs index d32f5516d4..a2a59dd45c 100644 --- a/crates/zed/src/zed/inline_completion_registry.rs +++ b/crates/zed/src/zed/inline_completion_registry.rs @@ -165,7 +165,7 @@ fn assign_inline_completion_provider( } } language::language_settings::InlineCompletionProvider::Zeta => { - if cx.has_flag::() { + if cx.has_flag::() || cfg!(debug_assertions) { let zeta = zeta::Zeta::register(client.clone(), cx); if let Some(buffer) = editor.buffer().read(cx).as_singleton() { if buffer.read(cx).file().is_some() { diff --git a/crates/zeta/Cargo.toml b/crates/zeta/Cargo.toml index 0b07703eff..5ac4e514f4 100644 --- a/crates/zeta/Cargo.toml +++ b/crates/zeta/Cargo.toml @@ -13,6 +13,9 @@ workspace = true path = "src/zeta.rs" doctest = false +[features] +test-support = [] + [dependencies] anyhow.workspace = true client.workspace = true @@ -21,6 +24,7 @@ editor.workspace = true futures.workspace = true gpui.workspace = true http_client.workspace = true +indoc.workspace = true inline_completion.workspace = true language.workspace = true language_models.workspace = true @@ -32,8 +36,8 @@ settings.workspace = true similar.workspace = true telemetry_events.workspace = true theme.workspace = true -util.workspace = true ui.workspace = true +util.workspace = true uuid.workspace = true workspace.workspace = true diff --git a/crates/zeta/src/rate_completion_modal.rs b/crates/zeta/src/rate_completion_modal.rs index c32959b459..d644923c13 100644 --- a/crates/zeta/src/rate_completion_modal.rs +++ b/crates/zeta/src/rate_completion_modal.rs @@ -1,18 +1,44 @@ use crate::{InlineCompletion, InlineCompletionRating, Zeta}; use editor::Editor; use gpui::{ - prelude::*, AppContext, DismissEvent, EventEmitter, FocusHandle, FocusableView, HighlightStyle, - Model, StyledText, TextStyle, View, ViewContext, + actions, prelude::*, AppContext, DismissEvent, EventEmitter, FocusHandle, FocusableView, + HighlightStyle, Model, StyledText, TextStyle, View, ViewContext, }; use language::{language_settings, OffsetRangeExt}; + use settings::Settings; use theme::ThemeSettings; use ui::{prelude::*, List, ListItem, ListItemSpacing, TintColor}; use workspace::{ModalView, Workspace}; +actions!( + zeta, + [ + RateCompletions, + ThumbsUp, + ThumbsDown, + ThumbsUpActiveCompletion, + ThumbsDownActiveCompletion, + NextEdit, + PreviousEdit, + FocusCompletions, + PreviewCompletion, + ] +); + +pub fn init(cx: &mut AppContext) { + cx.observe_new_views(move |workspace: &mut Workspace, _cx| { + workspace.register_action(|workspace, _: &RateCompletions, cx| { + RateCompletionModal::toggle(workspace, cx); + }); + }) + .detach(); +} + pub struct RateCompletionModal { zeta: Model, active_completion: Option, + selected_index: usize, focus_handle: FocusHandle, _subscription: gpui::Subscription, } @@ -33,6 +59,7 @@ impl RateCompletionModal { let subscription = cx.observe(&zeta, |_, _, cx| cx.notify()); Self { zeta, + selected_index: 0, focus_handle: cx.focus_handle(), active_completion: None, _subscription: subscription, @@ -43,15 +70,211 @@ impl RateCompletionModal { cx.emit(DismissEvent); } + fn select_next(&mut self, _: &menu::SelectNext, cx: &mut ViewContext) { + self.selected_index += 1; + self.selected_index = usize::min( + self.selected_index, + self.zeta.read(cx).recent_completions().count(), + ); + cx.notify(); + } + + fn select_prev(&mut self, _: &menu::SelectPrev, cx: &mut ViewContext) { + self.selected_index = self.selected_index.saturating_sub(1); + cx.notify(); + } + + fn select_next_edit(&mut self, _: &NextEdit, cx: &mut ViewContext) { + let next_index = self + .zeta + .read(cx) + .recent_completions() + .skip(self.selected_index) + .enumerate() + .skip(1) // Skip straight to the next item + .find(|(_, completion)| !completion.edits.is_empty()) + .map(|(ix, _)| ix + self.selected_index); + + if let Some(next_index) = next_index { + self.selected_index = next_index; + cx.notify(); + } + } + + fn select_prev_edit(&mut self, _: &PreviousEdit, cx: &mut ViewContext) { + let zeta = self.zeta.read(cx); + let completions_len = zeta.recent_completions_len(); + + let prev_index = self + .zeta + .read(cx) + .recent_completions() + .rev() + .skip((completions_len - 1) - self.selected_index) + .enumerate() + .skip(1) // Skip straight to the previous item + .find(|(_, completion)| !completion.edits.is_empty()) + .map(|(ix, _)| self.selected_index - ix); + + if let Some(prev_index) = prev_index { + self.selected_index = prev_index; + cx.notify(); + } + cx.notify(); + } + + fn select_first(&mut self, _: &menu::SelectFirst, cx: &mut ViewContext) { + self.selected_index = 0; + cx.notify(); + } + + fn select_last(&mut self, _: &menu::SelectLast, cx: &mut ViewContext) { + self.selected_index = self.zeta.read(cx).recent_completions_len() - 1; + cx.notify(); + } + + fn thumbs_up(&mut self, _: &ThumbsUp, cx: &mut ViewContext) { + self.zeta.update(cx, |zeta, cx| { + let completion = zeta + .recent_completions() + .skip(self.selected_index) + .next() + .cloned(); + + if let Some(completion) = completion { + zeta.rate_completion( + &completion, + InlineCompletionRating::Positive, + "".to_string(), + cx, + ); + } + }); + self.select_next_edit(&Default::default(), cx); + cx.notify(); + } + + fn thumbs_down(&mut self, _: &ThumbsDown, cx: &mut ViewContext) { + self.zeta.update(cx, |zeta, cx| { + let completion = zeta + .recent_completions() + .skip(self.selected_index) + .next() + .cloned(); + + if let Some(completion) = completion { + zeta.rate_completion( + &completion, + InlineCompletionRating::Negative, + "".to_string(), + cx, + ); + } + }); + self.select_next_edit(&Default::default(), cx); + cx.notify(); + } + + fn thumbs_up_active(&mut self, _: &ThumbsUpActiveCompletion, cx: &mut ViewContext) { + self.zeta.update(cx, |zeta, cx| { + if let Some(active) = &self.active_completion { + zeta.rate_completion( + &active.completion, + InlineCompletionRating::Positive, + active.feedback_editor.read(cx).text(cx), + cx, + ); + } + }); + + let current_completion = self + .active_completion + .as_ref() + .map(|completion| completion.completion.clone()); + self.select_completion(current_completion, false, cx); + self.select_next_edit(&Default::default(), cx); + self.confirm(&Default::default(), cx); + + cx.notify(); + } + + fn thumbs_down_active(&mut self, _: &ThumbsDownActiveCompletion, cx: &mut ViewContext) { + self.zeta.update(cx, |zeta, cx| { + if let Some(active) = &self.active_completion { + zeta.rate_completion( + &active.completion, + InlineCompletionRating::Negative, + active.feedback_editor.read(cx).text(cx), + cx, + ); + } + }); + + let current_completion = self + .active_completion + .as_ref() + .map(|completion| completion.completion.clone()); + self.select_completion(current_completion, false, cx); + self.select_next_edit(&Default::default(), cx); + self.confirm(&Default::default(), cx); + + cx.notify(); + } + + fn focus_completions(&mut self, _: &FocusCompletions, cx: &mut ViewContext) { + cx.focus_self(); + cx.notify(); + } + + fn preview_completion(&mut self, _: &PreviewCompletion, cx: &mut ViewContext) { + let completion = self + .zeta + .read(cx) + .recent_completions() + .skip(self.selected_index) + .take(1) + .next() + .cloned(); + + self.select_completion(completion, false, cx); + } + + fn confirm(&mut self, _: &menu::Confirm, cx: &mut ViewContext) { + let completion = self + .zeta + .read(cx) + .recent_completions() + .skip(self.selected_index) + .take(1) + .next() + .cloned(); + + self.select_completion(completion, true, cx); + } + pub fn select_completion( &mut self, completion: Option, + focus: bool, cx: &mut ViewContext, ) { // Avoid resetting completion rating if it's already selected. if let Some(completion) = completion.as_ref() { + self.selected_index = self + .zeta + .read(cx) + .recent_completions() + .enumerate() + .find(|(_, completion_b)| completion.id == completion_b.id) + .map(|(ix, _)| ix) + .unwrap_or(self.selected_index); + cx.notify(); + if let Some(prev_completion) = self.active_completion.as_ref() { if completion.id == prev_completion.completion.id { + if focus { + cx.focus_view(&prev_completion.feedback_editor); + } return; } } @@ -70,9 +293,13 @@ impl RateCompletionModal { editor.set_show_indent_guides(false, cx); editor.set_show_inline_completions(Some(false), cx); editor.set_placeholder_text("Add your feedback about this completion…", cx); + if focus { + cx.focus_self(); + } editor }), }); + cx.notify(); } fn render_active_completion(&mut self, cx: &mut ViewContext) -> Option { @@ -204,21 +431,12 @@ impl RateCompletionModal { .icon_position(IconPosition::Start) .icon_color(Color::Error) .disabled(rated) - .on_click({ - let completion = active_completion.completion.clone(); - let feedback_editor = - active_completion.feedback_editor.clone(); - cx.listener(move |this, _, cx| { - this.zeta.update(cx, |zeta, cx| { - zeta.rate_completion( - &completion, - InlineCompletionRating::Negative, - feedback_editor.read(cx).text(cx), - cx, - ) - }) - }) - }), + .on_click(cx.listener(move |this, _, cx| { + this.thumbs_down_active( + &ThumbsDownActiveCompletion, + cx, + ); + })), ) .child( Button::new("good", "Good Completion") @@ -228,21 +446,9 @@ impl RateCompletionModal { .icon_position(IconPosition::Start) .icon_color(Color::Success) .disabled(rated) - .on_click({ - let completion = active_completion.completion.clone(); - let feedback_editor = - active_completion.feedback_editor.clone(); - cx.listener(move |this, _, cx| { - this.zeta.update(cx, |zeta, cx| { - zeta.rate_completion( - &completion, - InlineCompletionRating::Positive, - feedback_editor.read(cx).text(cx), - cx, - ) - }) - }) - }), + .on_click(cx.listener(move |this, _, cx| { + this.thumbs_up_active(&ThumbsUpActiveCompletion, cx); + })), ), ), ), @@ -257,7 +463,23 @@ impl Render for RateCompletionModal { h_flex() .key_context("RateCompletionModal") .track_focus(&self.focus_handle) + .focus(|this| { + this.border_1().border_color(cx.theme().colors().border_focused) + }) .on_action(cx.listener(Self::dismiss)) + .on_action(cx.listener(Self::confirm)) + .on_action(cx.listener(Self::select_prev)) + .on_action(cx.listener(Self::select_prev_edit)) + .on_action(cx.listener(Self::select_next)) + .on_action(cx.listener(Self::select_next_edit)) + .on_action(cx.listener(Self::select_first)) + .on_action(cx.listener(Self::select_last)) + .on_action(cx.listener(Self::thumbs_up)) + .on_action(cx.listener(Self::thumbs_down)) + .on_action(cx.listener(Self::thumbs_up_active)) + .on_action(cx.listener(Self::thumbs_down_active)) + .on_action(cx.listener(Self::focus_completions)) + .on_action(cx.listener(Self::preview_completion)) .bg(cx.theme().colors().elevated_surface_background) .border_1() .border_color(border_color) @@ -285,8 +507,8 @@ impl Render for RateCompletionModal { ) .into_any_element(), ) - .children(self.zeta.read(cx).recent_completions().cloned().map( - |completion| { + .children(self.zeta.read(cx).recent_completions().cloned().enumerate().map( + |(index, completion)| { let selected = self.active_completion.as_ref().map_or(false, |selected| { selected.completion.id == completion.id @@ -296,6 +518,7 @@ impl Render for RateCompletionModal { ListItem::new(completion.id) .inset(true) .spacing(ListItemSpacing::Sparse) + .focused(index == self.selected_index) .selected(selected) .start_slot(if rated { Icon::new(IconName::Check).color(Color::Success) @@ -316,7 +539,7 @@ impl Render for RateCompletionModal { .size(LabelSize::XSmall)), ) .on_click(cx.listener(move |this, _, cx| { - this.select_completion(Some(completion.clone()), cx); + this.select_completion(Some(completion.clone()), true, cx); })) }, )), diff --git a/crates/zeta/src/zeta.rs b/crates/zeta/src/zeta.rs index ed1dcdbeb2..b958e75f68 100644 --- a/crates/zeta/src/zeta.rs +++ b/crates/zeta/src/zeta.rs @@ -18,6 +18,7 @@ use std::{ borrow::Cow, cmp, fmt::Write, + future::Future, mem, ops::Range, path::Path, @@ -253,12 +254,17 @@ impl Zeta { } } - pub fn request_completion( + pub fn request_completion_impl( &mut self, buffer: &Model, position: language::Anchor, cx: &mut ModelContext, - ) -> Task> { + perform_predict_edits: F, + ) -> Task> + where + F: FnOnce(Arc, LlmApiToken, PredictEditsParams) -> R + 'static, + R: Future> + Send + 'static, + { let snapshot = self.report_changes_for_buffer(buffer, cx); let point = position.to_point(&snapshot); let offset = point.to_offset(&snapshot); @@ -292,7 +298,7 @@ impl Zeta { input_excerpt: input_excerpt.clone(), }; - let response = Self::perform_predict_edits(&client, llm_token, body).await?; + let response = perform_predict_edits(client, llm_token, body).await?; let output_excerpt = response.output_excerpt; log::debug!("prediction took: {:?}", start.elapsed()); @@ -320,50 +326,210 @@ impl Zeta { }) } - async fn perform_predict_edits( - client: &Arc, + // Generates several example completions of various states to fill the Zeta completion modal + #[cfg(any(test, feature = "test-support"))] + pub fn fill_with_fake_completions(&mut self, cx: &mut ModelContext) -> Task<()> { + let test_buffer_text = indoc::indoc! {r#"a longggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg line + And maybe a short line + + Then a few lines + + and then another + "#}; + + let buffer = cx.new_model(|cx| Buffer::local(test_buffer_text, cx)); + let position = buffer.read(cx).anchor_before(Point::new(1, 0)); + + let completion_tasks = vec![ + self.fake_completion( + &buffer, + position, + PredictEditsResponse { + output_excerpt: format!("{EDITABLE_REGION_START_MARKER} +a longggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg line +[here's an edit] +And maybe a short line +Then a few lines +and then another +{EDITABLE_REGION_END_MARKER} + ", ), + }, + cx, + ), + self.fake_completion( + &buffer, + position, + PredictEditsResponse { + output_excerpt: format!(r#"{EDITABLE_REGION_START_MARKER} +a longggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg line +And maybe a short line +[and another edit] +Then a few lines +and then another +{EDITABLE_REGION_END_MARKER} + "#), + }, + cx, + ), + self.fake_completion( + &buffer, + position, + PredictEditsResponse { + output_excerpt: format!(r#"{EDITABLE_REGION_START_MARKER} +a longggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg line +And maybe a short line + +Then a few lines + +and then another +{EDITABLE_REGION_END_MARKER} + "#), + }, + cx, + ), + self.fake_completion( + &buffer, + position, + PredictEditsResponse { + output_excerpt: format!(r#"{EDITABLE_REGION_START_MARKER} +a longggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg line +And maybe a short line + +Then a few lines + +and then another +{EDITABLE_REGION_END_MARKER} + "#), + }, + cx, + ), + self.fake_completion( + &buffer, + position, + PredictEditsResponse { + output_excerpt: format!(r#"{EDITABLE_REGION_START_MARKER} +a longggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg line +And maybe a short line +Then a few lines +[a third completion] +and then another +{EDITABLE_REGION_END_MARKER} + "#), + }, + cx, + ), + self.fake_completion( + &buffer, + position, + PredictEditsResponse { + output_excerpt: format!(r#"{EDITABLE_REGION_START_MARKER} +a longggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg line +And maybe a short line +and then another +[fourth completion example] +{EDITABLE_REGION_END_MARKER} + "#), + }, + cx, + ), + self.fake_completion( + &buffer, + position, + PredictEditsResponse { + output_excerpt: format!(r#"{EDITABLE_REGION_START_MARKER} +a longggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg line +And maybe a short line +Then a few lines +and then another +[fifth and final completion] +{EDITABLE_REGION_END_MARKER} + "#), + }, + cx, + ), + ]; + + cx.spawn(|zeta, mut cx| async move { + for task in completion_tasks { + task.await.unwrap(); + } + + zeta.update(&mut cx, |zeta, _cx| { + zeta.recent_completions.get_mut(2).unwrap().edits = Arc::new([]); + zeta.recent_completions.get_mut(3).unwrap().edits = Arc::new([]); + }) + .ok(); + }) + } + + #[cfg(any(test, feature = "test-support"))] + pub fn fake_completion( + &mut self, + buffer: &Model, + position: language::Anchor, + response: PredictEditsResponse, + cx: &mut ModelContext, + ) -> Task> { + use std::future::ready; + + self.request_completion_impl(buffer, position, cx, |_, _, _| ready(Ok(response))) + } + + pub fn request_completion( + &mut self, + buffer: &Model, + position: language::Anchor, + cx: &mut ModelContext, + ) -> Task> { + self.request_completion_impl(buffer, position, cx, Self::perform_predict_edits) + } + + fn perform_predict_edits( + client: Arc, llm_token: LlmApiToken, body: PredictEditsParams, - ) -> Result { - let http_client = client.http_client(); - let mut token = llm_token.acquire(client).await?; - let mut did_retry = false; + ) -> impl Future> { + async move { + let http_client = client.http_client(); + let mut token = llm_token.acquire(&client).await?; + let mut did_retry = false; - loop { - let request_builder = http_client::Request::builder(); - let request = request_builder - .method(Method::POST) - .uri( - http_client - .build_zed_llm_url("/predict_edits", &[])? - .as_ref(), - ) - .header("Content-Type", "application/json") - .header("Authorization", format!("Bearer {}", token)) - .body(serde_json::to_string(&body)?.into())?; + loop { + let request_builder = http_client::Request::builder(); + let request = request_builder + .method(Method::POST) + .uri( + http_client + .build_zed_llm_url("/predict_edits", &[])? + .as_ref(), + ) + .header("Content-Type", "application/json") + .header("Authorization", format!("Bearer {}", token)) + .body(serde_json::to_string(&body)?.into())?; - let mut response = http_client.send(request).await?; + let mut response = http_client.send(request).await?; - if response.status().is_success() { - let mut body = String::new(); - response.body_mut().read_to_string(&mut body).await?; - return Ok(serde_json::from_str(&body)?); - } else if !did_retry - && response - .headers() - .get(EXPIRED_LLM_TOKEN_HEADER_NAME) - .is_some() - { - did_retry = true; - token = llm_token.refresh(client).await?; - } else { - let mut body = String::new(); - response.body_mut().read_to_string(&mut body).await?; - return Err(anyhow!( - "error predicting edits.\nStatus: {:?}\nBody: {}", - response.status(), - body - )); + if response.status().is_success() { + let mut body = String::new(); + response.body_mut().read_to_string(&mut body).await?; + return Ok(serde_json::from_str(&body)?); + } else if !did_retry + && response + .headers() + .get(EXPIRED_LLM_TOKEN_HEADER_NAME) + .is_some() + { + did_retry = true; + token = llm_token.refresh(&client).await?; + } else { + let mut body = String::new(); + response.body_mut().read_to_string(&mut body).await?; + return Err(anyhow!( + "error predicting edits.\nStatus: {:?}\nBody: {}", + response.status(), + body + )); + } } } } @@ -409,7 +575,7 @@ impl Zeta { }) } - fn compute_edits( + pub fn compute_edits( old_text: String, new_text: &str, offset: usize, @@ -500,10 +666,14 @@ impl Zeta { cx.notify(); } - pub fn recent_completions(&self) -> impl Iterator { + pub fn recent_completions(&self) -> impl DoubleEndedIterator { self.recent_completions.iter() } + pub fn recent_completions_len(&self) -> usize { + self.recent_completions.len() + } + fn report_changes_for_buffer( &mut self, buffer: &Model,