From 9751e6110176a4bb7931264ede907bd0acc06535 Mon Sep 17 00:00:00 2001 From: Max Brunsfeld Date: Wed, 31 Jul 2024 12:05:19 -0700 Subject: [PATCH] Show rate limit notices (#15515) This UI change is behind a `ZedPro` feature flag so that it won't be visible until we're ready to launch that service. Release Notes: - N/A --------- Co-authored-by: Marshall Bowers Co-authored-by: Marshall --- Cargo.lock | 4 + crates/assistant/Cargo.toml | 4 + crates/assistant/src/inline_assistant.rs | 159 +++++++++++++++++++++-- crates/collab/src/rate_limiter.rs | 8 +- crates/proto/proto/zed.proto | 1 + crates/rpc/src/peer.rs | 6 +- 6 files changed, 160 insertions(+), 22 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 17b27bbf10..4ded0abd22 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -407,8 +407,10 @@ dependencies = [ "collections", "command_palette_hooks", "ctor", + "db", "editor", "env_logger", + "feature_flags", "fs", "futures 0.3.28", "fuzzy", @@ -430,6 +432,7 @@ dependencies = [ "paths", "picker", "project", + "proto", "rand 0.8.5", "regex", "rope", @@ -454,6 +457,7 @@ dependencies = [ "util", "uuid", "workspace", + "zed_actions", ] [[package]] diff --git a/crates/assistant/Cargo.toml b/crates/assistant/Cargo.toml index 0bdc4642c5..bcff377a45 100644 --- a/crates/assistant/Cargo.toml +++ b/crates/assistant/Cargo.toml @@ -32,7 +32,9 @@ client.workspace = true clock.workspace = true collections.workspace = true command_palette_hooks.workspace = true +db.workspace = true editor.workspace = true +feature_flags.workspace = true fs.workspace = true futures.workspace = true fuzzy.workspace = true @@ -53,6 +55,7 @@ ordered-float.workspace = true parking_lot.workspace = true paths.workspace = true project.workspace = true +proto.workspace = true regex.workspace = true rope.workspace = true schemars.workspace = true @@ -74,6 +77,7 @@ util.workspace = true uuid.workspace = true workspace.workspace = true picker.workspace = true +zed_actions.workspace = true [dev-dependencies] ctor.workspace = true diff --git a/crates/assistant/src/inline_assistant.rs b/crates/assistant/src/inline_assistant.rs index 7eed29a76d..4d2a1225c2 100644 --- a/crates/assistant/src/inline_assistant.rs +++ b/crates/assistant/src/inline_assistant.rs @@ -3,7 +3,7 @@ use crate::{ Hunk, ModelSelector, StreamingDiff, }; use anyhow::{anyhow, Context as _, Result}; -use client::telemetry::Telemetry; +use client::{telemetry::Telemetry, ErrorExt}; use collections::{hash_map, HashMap, HashSet, VecDeque}; use editor::{ actions::{MoveDown, MoveUp, SelectAll}, @@ -14,6 +14,7 @@ use editor::{ Anchor, AnchorRangeExt, Editor, EditorElement, EditorEvent, EditorMode, EditorStyle, ExcerptRange, GutterDimensions, MultiBuffer, MultiBufferSnapshot, ToOffset, ToPoint, }; +use feature_flags::{FeatureFlagAppExt as _, ZedPro}; use fs::Fs; use futures::{ channel::mpsc, @@ -22,9 +23,9 @@ use futures::{ SinkExt, Stream, StreamExt, }; use gpui::{ - point, AppContext, EventEmitter, FocusHandle, FocusableView, Global, HighlightStyle, Model, - ModelContext, Subscription, Task, TextStyle, UpdateGlobal, View, ViewContext, WeakView, - WindowContext, + anchored, deferred, point, AppContext, ClickEvent, EventEmitter, FocusHandle, FocusableView, + FontWeight, Global, HighlightStyle, Model, ModelContext, Subscription, Task, TextStyle, + UpdateGlobal, View, ViewContext, WeakView, WindowContext, }; use language::{Buffer, IndentKind, Point, Selection, TransactionId}; use language_model::{ @@ -47,7 +48,7 @@ use std::{ time::{Duration, Instant}, }; use theme::ThemeSettings; -use ui::{prelude::*, IconButtonShape, Tooltip}; +use ui::{prelude::*, CheckboxWithLabel, IconButtonShape, Popover, Tooltip}; use util::{RangeExt, ResultExt}; use workspace::{notifications::NotificationId, Toast, Workspace}; @@ -1187,6 +1188,7 @@ struct PromptEditor { token_count: Option, _token_count_subscriptions: Vec, workspace: Option>, + show_rate_limit_notice: bool, } impl EventEmitter for PromptEditor {} @@ -1319,10 +1321,36 @@ impl Render for PromptEditor { assistant panel tab.", ), ) - .children( - if let CodegenStatus::Error(error) = &self.codegen.read(cx).status { - let error_message = SharedString::from(error.to_string()); - Some( + .map(|el| { + let CodegenStatus::Error(error) = &self.codegen.read(cx).status else { + return el; + }; + + let error_message = SharedString::from(error.to_string()); + if error.error_code() == proto::ErrorCode::RateLimitExceeded + && cx.has_flag::() + { + el.child( + v_flex() + .child( + IconButton::new("rate-limit-error", IconName::XCircle) + .selected(self.show_rate_limit_notice) + .shape(IconButtonShape::Square) + .icon_size(IconSize::Small) + .on_click(cx.listener(Self::toggle_rate_limit_notice)), + ) + .children(self.show_rate_limit_notice.then(|| { + deferred( + anchored() + .position_mode(gpui::AnchoredPositionMode::Local) + .position(point(px(0.), px(24.))) + .anchor(gpui::AnchorCorner::TopLeft) + .child(self.render_rate_limit_notice(cx)), + ) + })), + ) + } else { + el.child( div() .id("error") .tooltip(move |cx| Tooltip::text(error_message.clone(), cx)) @@ -1332,10 +1360,8 @@ impl Render for PromptEditor { .color(Color::Error), ), ) - } else { - None - }, - ), + } + }), ) .child(div().flex_1().child(self.render_prompt_editor(cx))) .child( @@ -1413,6 +1439,7 @@ impl PromptEditor { token_count: None, _token_count_subscriptions: token_count_subscriptions, workspace, + show_rate_limit_notice: false, }; this.count_tokens(cx); this.subscribe_to_editor(cx); @@ -1455,6 +1482,14 @@ impl PromptEditor { self.editor.read(cx).text(cx) } + fn toggle_rate_limit_notice(&mut self, _: &ClickEvent, cx: &mut ViewContext) { + self.show_rate_limit_notice = !self.show_rate_limit_notice; + if self.show_rate_limit_notice { + cx.focus_view(&self.editor); + } + cx.notify(); + } + fn handle_parent_editor_event( &mut self, _: View, @@ -1520,6 +1555,12 @@ impl PromptEditor { EditorEvent::BufferEdited => { self.count_tokens(cx); } + EditorEvent::Blurred => { + if self.show_rate_limit_notice { + self.show_rate_limit_notice = false; + cx.notify(); + } + } _ => {} } } @@ -1534,7 +1575,20 @@ impl PromptEditor { self.editor .update(cx, |editor, _| editor.set_read_only(true)); } - CodegenStatus::Done | CodegenStatus::Error(_) => { + CodegenStatus::Done => { + self.edited_since_done = false; + self.editor + .update(cx, |editor, _| editor.set_read_only(false)); + } + CodegenStatus::Error(error) => { + if cx.has_flag::() + && error.error_code() == proto::ErrorCode::RateLimitExceeded + && !dismissed_rate_limit_notice() + { + self.show_rate_limit_notice = true; + cx.notify(); + } + self.edited_since_done = false; self.editor .update(cx, |editor, _| editor.set_read_only(false)); @@ -1694,6 +1748,83 @@ impl PromptEditor { }, ) } + + fn render_rate_limit_notice(&self, cx: &mut ViewContext) -> impl IntoElement { + Popover::new().child( + v_flex() + .occlude() + .p_2() + .child( + Label::new("Out of Tokens") + .size(LabelSize::Small) + .weight(FontWeight::BOLD), + ) + .child(Label::new( + "Try Zed Pro for higher limits, a wider range of models, and more.", + )) + .child( + h_flex() + .justify_between() + .child(CheckboxWithLabel::new( + "dont-show-again", + Label::new("Don't show again"), + if dismissed_rate_limit_notice() { + ui::Selection::Selected + } else { + ui::Selection::Unselected + }, + |selection, cx| { + let is_dismissed = match selection { + ui::Selection::Unselected => false, + ui::Selection::Indeterminate => return, + ui::Selection::Selected => true, + }; + + set_rate_limit_notice_dismissed(is_dismissed, cx) + }, + )) + .child( + h_flex() + .gap_2() + .child( + Button::new("dismiss", "Dismiss") + .style(ButtonStyle::Transparent) + .on_click(cx.listener(Self::toggle_rate_limit_notice)), + ) + .child(Button::new("more-info", "More Info").on_click( + |_event, cx| { + cx.dispatch_action(Box::new( + zed_actions::OpenAccountSettings, + )) + }, + )), + ), + ), + ) + } +} + +const DISMISSED_RATE_LIMIT_NOTICE_KEY: &str = "dismissed-rate-limit-notice"; + +fn dismissed_rate_limit_notice() -> bool { + db::kvp::KEY_VALUE_STORE + .read_kvp(DISMISSED_RATE_LIMIT_NOTICE_KEY) + .log_err() + .map_or(false, |s| s.is_some()) +} + +fn set_rate_limit_notice_dismissed(is_dismissed: bool, cx: &mut AppContext) { + db::write_and_log(cx, move || async move { + if is_dismissed { + db::kvp::KEY_VALUE_STORE + .write_kvp(DISMISSED_RATE_LIMIT_NOTICE_KEY.into(), "1".into()) + .await + } else { + db::kvp::KEY_VALUE_STORE + .delete_kvp(DISMISSED_RATE_LIMIT_NOTICE_KEY.into()) + .await + } + }) } struct InlineAssist { diff --git a/crates/collab/src/rate_limiter.rs b/crates/collab/src/rate_limiter.rs index 5e619f7db2..b6ed7dcc4e 100644 --- a/crates/collab/src/rate_limiter.rs +++ b/crates/collab/src/rate_limiter.rs @@ -1,7 +1,7 @@ use crate::{db::UserId, executor::Executor, Database, Error, Result}; -use anyhow::anyhow; use chrono::{DateTime, Duration, Utc}; use dashmap::{DashMap, DashSet}; +use rpc::ErrorCodeExt; use sea_orm::prelude::DateTimeUtc; use std::sync::Arc; use util::ResultExt; @@ -73,7 +73,9 @@ impl RateLimiter { self.dirty_buckets.insert(bucket_key); Ok(()) } else { - Err(anyhow!("rate limit exceeded"))? + Err(rpc::proto::ErrorCode::RateLimitExceeded + .message("rate limit exceeded".into()) + .anyhow())? } } @@ -122,7 +124,7 @@ impl RateLimiter { } } -#[derive(Clone)] +#[derive(Clone, Debug)] struct RateBucket { capacity: usize, token_count: usize, diff --git a/crates/proto/proto/zed.proto b/crates/proto/proto/zed.proto index e6cc7c74a2..b396abe8e3 100644 --- a/crates/proto/proto/zed.proto +++ b/crates/proto/proto/zed.proto @@ -311,6 +311,7 @@ enum ErrorCode { DevServerOffline = 15; DevServerProjectPathDoesNotExist = 16; RemoteUpgradeRequired = 17; + RateLimitExceeded = 18; reserved 6; } diff --git a/crates/rpc/src/peer.rs b/crates/rpc/src/peer.rs index e0d38f9f37..242b3b9116 100644 --- a/crates/rpc/src/peer.rs +++ b/crates/rpc/src/peer.rs @@ -516,11 +516,7 @@ impl Peer { future::ready(match response { Ok(response) => { if let Some(proto::envelope::Payload::Error(error)) = &response.payload { - Some(Err(anyhow!( - "RPC request {} failed - {}", - T::NAME, - error.message - ))) + Some(Err(RpcError::from_proto(&error, T::NAME))) } else if let Some(proto::envelope::Payload::EndStream(_)) = &response.payload {