1use std::ops::Range;
2use std::str::FromStr as _;
3
4use anyhow::{Result, anyhow};
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 .ok_or_else(|| {
73 anyhow!("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 .ok_or_else(|| {
80 anyhow!("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
88pub trait EditPredictionProvider: 'static + Sized {
89 fn name() -> &'static str;
90 fn display_name() -> &'static str;
91 fn show_completions_in_menu() -> bool;
92 fn show_tab_accept_marker() -> bool {
93 false
94 }
95 fn data_collection_state(&self, _cx: &App) -> DataCollectionState {
96 DataCollectionState::Unsupported
97 }
98
99 fn usage(&self, _cx: &App) -> Option<EditPredictionUsage> {
100 None
101 }
102
103 fn toggle_data_collection(&mut self, _cx: &mut App) {}
104 fn is_enabled(
105 &self,
106 buffer: &Entity<Buffer>,
107 cursor_position: language::Anchor,
108 cx: &App,
109 ) -> bool;
110 fn is_refreshing(&self) -> bool;
111 fn refresh(
112 &mut self,
113 project: Option<Entity<Project>>,
114 buffer: Entity<Buffer>,
115 cursor_position: language::Anchor,
116 debounce: bool,
117 cx: &mut Context<Self>,
118 );
119 fn needs_terms_acceptance(&self, _cx: &App) -> bool {
120 false
121 }
122 fn cycle(
123 &mut self,
124 buffer: Entity<Buffer>,
125 cursor_position: language::Anchor,
126 direction: Direction,
127 cx: &mut Context<Self>,
128 );
129 fn accept(&mut self, cx: &mut Context<Self>);
130 fn discard(&mut self, cx: &mut Context<Self>);
131 fn suggest(
132 &mut self,
133 buffer: &Entity<Buffer>,
134 cursor_position: language::Anchor,
135 cx: &mut Context<Self>,
136 ) -> Option<InlineCompletion>;
137}
138
139pub trait InlineCompletionProviderHandle {
140 fn name(&self) -> &'static str;
141 fn display_name(&self) -> &'static str;
142 fn is_enabled(
143 &self,
144 buffer: &Entity<Buffer>,
145 cursor_position: language::Anchor,
146 cx: &App,
147 ) -> bool;
148 fn show_completions_in_menu(&self) -> bool;
149 fn show_tab_accept_marker(&self) -> bool;
150 fn data_collection_state(&self, cx: &App) -> DataCollectionState;
151 fn usage(&self, cx: &App) -> Option<EditPredictionUsage>;
152 fn toggle_data_collection(&self, cx: &mut App);
153 fn needs_terms_acceptance(&self, cx: &App) -> bool;
154 fn is_refreshing(&self, cx: &App) -> bool;
155 fn refresh(
156 &self,
157 project: Option<Entity<Project>>,
158 buffer: Entity<Buffer>,
159 cursor_position: language::Anchor,
160 debounce: bool,
161 cx: &mut App,
162 );
163 fn cycle(
164 &self,
165 buffer: Entity<Buffer>,
166 cursor_position: language::Anchor,
167 direction: Direction,
168 cx: &mut App,
169 );
170 fn accept(&self, cx: &mut App);
171 fn discard(&self, cx: &mut App);
172 fn suggest(
173 &self,
174 buffer: &Entity<Buffer>,
175 cursor_position: language::Anchor,
176 cx: &mut App,
177 ) -> Option<InlineCompletion>;
178}
179
180impl<T> InlineCompletionProviderHandle for Entity<T>
181where
182 T: EditPredictionProvider,
183{
184 fn name(&self) -> &'static str {
185 T::name()
186 }
187
188 fn display_name(&self) -> &'static str {
189 T::display_name()
190 }
191
192 fn show_completions_in_menu(&self) -> bool {
193 T::show_completions_in_menu()
194 }
195
196 fn show_tab_accept_marker(&self) -> bool {
197 T::show_tab_accept_marker()
198 }
199
200 fn data_collection_state(&self, cx: &App) -> DataCollectionState {
201 self.read(cx).data_collection_state(cx)
202 }
203
204 fn usage(&self, cx: &App) -> Option<EditPredictionUsage> {
205 self.read(cx).usage(cx)
206 }
207
208 fn toggle_data_collection(&self, cx: &mut App) {
209 self.update(cx, |this, cx| this.toggle_data_collection(cx))
210 }
211
212 fn is_enabled(
213 &self,
214 buffer: &Entity<Buffer>,
215 cursor_position: language::Anchor,
216 cx: &App,
217 ) -> bool {
218 self.read(cx).is_enabled(buffer, cursor_position, cx)
219 }
220
221 fn needs_terms_acceptance(&self, cx: &App) -> bool {
222 self.read(cx).needs_terms_acceptance(cx)
223 }
224
225 fn is_refreshing(&self, cx: &App) -> bool {
226 self.read(cx).is_refreshing()
227 }
228
229 fn refresh(
230 &self,
231 project: Option<Entity<Project>>,
232 buffer: Entity<Buffer>,
233 cursor_position: language::Anchor,
234 debounce: bool,
235 cx: &mut App,
236 ) {
237 self.update(cx, |this, cx| {
238 this.refresh(project, buffer, cursor_position, debounce, cx)
239 })
240 }
241
242 fn cycle(
243 &self,
244 buffer: Entity<Buffer>,
245 cursor_position: language::Anchor,
246 direction: Direction,
247 cx: &mut App,
248 ) {
249 self.update(cx, |this, cx| {
250 this.cycle(buffer, cursor_position, direction, cx)
251 })
252 }
253
254 fn accept(&self, cx: &mut App) {
255 self.update(cx, |this, cx| this.accept(cx))
256 }
257
258 fn discard(&self, cx: &mut App) {
259 self.update(cx, |this, cx| this.discard(cx))
260 }
261
262 fn suggest(
263 &self,
264 buffer: &Entity<Buffer>,
265 cursor_position: language::Anchor,
266 cx: &mut App,
267 ) -> Option<InlineCompletion> {
268 self.update(cx, |this, cx| this.suggest(buffer, cursor_position, cx))
269 }
270}