1use std::ops::Range;
2use std::str::FromStr as _;
3
4use anyhow::{Context as _, Result};
5use gpui::http_client::http::{HeaderMap, HeaderValue};
6use gpui::{App, Context, Entity, SharedString};
7use language::Buffer;
8use project::Project;
9use zed_llm_client::{
10 EDIT_PREDICTIONS_USAGE_AMOUNT_HEADER_NAME, EDIT_PREDICTIONS_USAGE_LIMIT_HEADER_NAME, UsageLimit,
11};
12
13// TODO: Find a better home for `Direction`.
14//
15// This should live in an ancestor crate of `editor` and `inline_completion`,
16// but at time of writing there isn't an obvious spot.
17#[derive(Copy, Clone, PartialEq, Eq)]
18pub enum Direction {
19 Prev,
20 Next,
21}
22
23#[derive(Clone)]
24pub struct InlineCompletion {
25 /// The ID of the completion, if it has one.
26 pub id: Option<SharedString>,
27 pub edits: Vec<(Range<language::Anchor>, String)>,
28 pub edit_preview: Option<language::EditPreview>,
29}
30
31pub enum DataCollectionState {
32 /// The provider doesn't support data collection.
33 Unsupported,
34 /// Data collection is enabled.
35 Enabled { is_project_open_source: bool },
36 /// Data collection is disabled or unanswered.
37 Disabled { is_project_open_source: bool },
38}
39
40impl DataCollectionState {
41 pub fn is_supported(&self) -> bool {
42 !matches!(self, DataCollectionState::Unsupported { .. })
43 }
44
45 pub fn is_enabled(&self) -> bool {
46 matches!(self, DataCollectionState::Enabled { .. })
47 }
48
49 pub fn is_project_open_source(&self) -> bool {
50 match self {
51 Self::Enabled {
52 is_project_open_source,
53 }
54 | Self::Disabled {
55 is_project_open_source,
56 } => *is_project_open_source,
57 _ => false,
58 }
59 }
60}
61
62#[derive(Debug, Clone, Copy)]
63pub struct EditPredictionUsage {
64 pub limit: UsageLimit,
65 pub amount: i32,
66}
67
68impl EditPredictionUsage {
69 pub fn from_headers(headers: &HeaderMap<HeaderValue>) -> Result<Self> {
70 let limit = headers
71 .get(EDIT_PREDICTIONS_USAGE_LIMIT_HEADER_NAME)
72 .with_context(|| {
73 format!("missing {EDIT_PREDICTIONS_USAGE_LIMIT_HEADER_NAME:?} header")
74 })?;
75 let limit = UsageLimit::from_str(limit.to_str()?)?;
76
77 let amount = headers
78 .get(EDIT_PREDICTIONS_USAGE_AMOUNT_HEADER_NAME)
79 .with_context(|| {
80 format!("missing {EDIT_PREDICTIONS_USAGE_AMOUNT_HEADER_NAME:?} header")
81 })?;
82 let amount = amount.to_str()?.parse::<i32>()?;
83
84 Ok(Self { limit, amount })
85 }
86
87 pub fn over_limit(&self) -> bool {
88 match self.limit {
89 UsageLimit::Limited(limit) => self.amount >= limit,
90 UsageLimit::Unlimited => false,
91 }
92 }
93}
94
95pub trait EditPredictionProvider: 'static + Sized {
96 fn name() -> &'static str;
97 fn display_name() -> &'static str;
98 fn show_completions_in_menu() -> bool;
99 fn show_tab_accept_marker() -> bool {
100 false
101 }
102 fn data_collection_state(&self, _cx: &App) -> DataCollectionState {
103 DataCollectionState::Unsupported
104 }
105
106 fn usage(&self, _cx: &App) -> Option<EditPredictionUsage> {
107 None
108 }
109
110 fn toggle_data_collection(&mut self, _cx: &mut App) {}
111 fn is_enabled(
112 &self,
113 buffer: &Entity<Buffer>,
114 cursor_position: language::Anchor,
115 cx: &App,
116 ) -> bool;
117 fn is_refreshing(&self) -> bool;
118 fn refresh(
119 &mut self,
120 project: Option<Entity<Project>>,
121 buffer: Entity<Buffer>,
122 cursor_position: language::Anchor,
123 debounce: bool,
124 cx: &mut Context<Self>,
125 );
126 fn needs_terms_acceptance(&self, _cx: &App) -> bool {
127 false
128 }
129 fn cycle(
130 &mut self,
131 buffer: Entity<Buffer>,
132 cursor_position: language::Anchor,
133 direction: Direction,
134 cx: &mut Context<Self>,
135 );
136 fn accept(&mut self, cx: &mut Context<Self>);
137 fn discard(&mut self, cx: &mut Context<Self>);
138 fn suggest(
139 &mut self,
140 buffer: &Entity<Buffer>,
141 cursor_position: language::Anchor,
142 cx: &mut Context<Self>,
143 ) -> Option<InlineCompletion>;
144}
145
146pub trait InlineCompletionProviderHandle {
147 fn name(&self) -> &'static str;
148 fn display_name(&self) -> &'static str;
149 fn is_enabled(
150 &self,
151 buffer: &Entity<Buffer>,
152 cursor_position: language::Anchor,
153 cx: &App,
154 ) -> bool;
155 fn show_completions_in_menu(&self) -> bool;
156 fn show_tab_accept_marker(&self) -> bool;
157 fn data_collection_state(&self, cx: &App) -> DataCollectionState;
158 fn usage(&self, cx: &App) -> Option<EditPredictionUsage>;
159 fn toggle_data_collection(&self, cx: &mut App);
160 fn needs_terms_acceptance(&self, cx: &App) -> bool;
161 fn is_refreshing(&self, cx: &App) -> bool;
162 fn refresh(
163 &self,
164 project: Option<Entity<Project>>,
165 buffer: Entity<Buffer>,
166 cursor_position: language::Anchor,
167 debounce: bool,
168 cx: &mut App,
169 );
170 fn cycle(
171 &self,
172 buffer: Entity<Buffer>,
173 cursor_position: language::Anchor,
174 direction: Direction,
175 cx: &mut App,
176 );
177 fn accept(&self, cx: &mut App);
178 fn discard(&self, cx: &mut App);
179 fn suggest(
180 &self,
181 buffer: &Entity<Buffer>,
182 cursor_position: language::Anchor,
183 cx: &mut App,
184 ) -> Option<InlineCompletion>;
185}
186
187impl<T> InlineCompletionProviderHandle for Entity<T>
188where
189 T: EditPredictionProvider,
190{
191 fn name(&self) -> &'static str {
192 T::name()
193 }
194
195 fn display_name(&self) -> &'static str {
196 T::display_name()
197 }
198
199 fn show_completions_in_menu(&self) -> bool {
200 T::show_completions_in_menu()
201 }
202
203 fn show_tab_accept_marker(&self) -> bool {
204 T::show_tab_accept_marker()
205 }
206
207 fn data_collection_state(&self, cx: &App) -> DataCollectionState {
208 self.read(cx).data_collection_state(cx)
209 }
210
211 fn usage(&self, cx: &App) -> Option<EditPredictionUsage> {
212 self.read(cx).usage(cx)
213 }
214
215 fn toggle_data_collection(&self, cx: &mut App) {
216 self.update(cx, |this, cx| this.toggle_data_collection(cx))
217 }
218
219 fn is_enabled(
220 &self,
221 buffer: &Entity<Buffer>,
222 cursor_position: language::Anchor,
223 cx: &App,
224 ) -> bool {
225 self.read(cx).is_enabled(buffer, cursor_position, cx)
226 }
227
228 fn needs_terms_acceptance(&self, cx: &App) -> bool {
229 self.read(cx).needs_terms_acceptance(cx)
230 }
231
232 fn is_refreshing(&self, cx: &App) -> bool {
233 self.read(cx).is_refreshing()
234 }
235
236 fn refresh(
237 &self,
238 project: Option<Entity<Project>>,
239 buffer: Entity<Buffer>,
240 cursor_position: language::Anchor,
241 debounce: bool,
242 cx: &mut App,
243 ) {
244 self.update(cx, |this, cx| {
245 this.refresh(project, buffer, cursor_position, debounce, cx)
246 })
247 }
248
249 fn cycle(
250 &self,
251 buffer: Entity<Buffer>,
252 cursor_position: language::Anchor,
253 direction: Direction,
254 cx: &mut App,
255 ) {
256 self.update(cx, |this, cx| {
257 this.cycle(buffer, cursor_position, direction, cx)
258 })
259 }
260
261 fn accept(&self, cx: &mut App) {
262 self.update(cx, |this, cx| this.accept(cx))
263 }
264
265 fn discard(&self, cx: &mut App) {
266 self.update(cx, |this, cx| this.discard(cx))
267 }
268
269 fn suggest(
270 &self,
271 buffer: &Entity<Buffer>,
272 cursor_position: language::Anchor,
273 cx: &mut App,
274 ) -> Option<InlineCompletion> {
275 self.update(cx, |this, cx| this.suggest(buffer, cursor_position, cx))
276 }
277}