1mod edit_prediction_button;
2mod edit_prediction_context_view;
3mod rate_prediction_modal;
4
5use command_palette_hooks::CommandPaletteFilter;
6use edit_prediction::{EditPredictionStore, ResetOnboarding, Zeta2FeatureFlag, capture_example};
7use edit_prediction_context_view::EditPredictionContextView;
8use editor::Editor;
9use feature_flags::FeatureFlagAppExt as _;
10use gpui::actions;
11use language::language_settings::AllLanguageSettings;
12use project::DisableAiSettings;
13use rate_prediction_modal::RatePredictionsModal;
14use settings::{Settings as _, SettingsStore};
15use std::any::{Any as _, TypeId};
16use ui::{App, prelude::*};
17use workspace::{SplitDirection, Workspace};
18
19pub use edit_prediction_button::{
20 EditPredictionButton, ToggleMenu, get_available_providers, set_completion_provider,
21};
22
23use crate::rate_prediction_modal::PredictEditsRatePredictionsFeatureFlag;
24
25actions!(
26 dev,
27 [
28 /// Opens the edit prediction context view.
29 OpenEditPredictionContextView,
30 ]
31);
32
33actions!(
34 edit_prediction,
35 [
36 /// Opens the rate completions modal.
37 RatePredictions,
38 /// Captures an ExampleSpec from the current editing session and opens it as Markdown.
39 CaptureExample,
40 ]
41);
42
43pub fn init(cx: &mut App) {
44 feature_gate_predict_edits_actions(cx);
45
46 cx.observe_new(move |workspace: &mut Workspace, _, _cx| {
47 workspace.register_action(|workspace, _: &RatePredictions, window, cx| {
48 if cx.has_flag::<PredictEditsRatePredictionsFeatureFlag>() {
49 RatePredictionsModal::toggle(workspace, window, cx);
50 }
51 });
52
53 workspace.register_action(|workspace, _: &CaptureExample, window, cx| {
54 capture_example_as_markdown(workspace, window, cx);
55 });
56 workspace.register_action_renderer(|div, _, _, cx| {
57 let has_flag = cx.has_flag::<Zeta2FeatureFlag>();
58 div.when(has_flag, |div| {
59 div.on_action(cx.listener(
60 move |workspace, _: &OpenEditPredictionContextView, window, cx| {
61 let project = workspace.project();
62 workspace.split_item(
63 SplitDirection::Right,
64 Box::new(cx.new(|cx| {
65 EditPredictionContextView::new(
66 project.clone(),
67 workspace.client(),
68 workspace.user_store(),
69 window,
70 cx,
71 )
72 })),
73 window,
74 cx,
75 );
76 },
77 ))
78 })
79 });
80 })
81 .detach();
82}
83
84fn feature_gate_predict_edits_actions(cx: &mut App) {
85 let rate_completion_action_types = [TypeId::of::<RatePredictions>()];
86 let reset_onboarding_action_types = [TypeId::of::<ResetOnboarding>()];
87 let all_action_types = [
88 TypeId::of::<RatePredictions>(),
89 TypeId::of::<CaptureExample>(),
90 TypeId::of::<edit_prediction::ResetOnboarding>(),
91 zed_actions::OpenZedPredictOnboarding.type_id(),
92 TypeId::of::<edit_prediction::ClearHistory>(),
93 TypeId::of::<rate_prediction_modal::ThumbsUpActivePrediction>(),
94 TypeId::of::<rate_prediction_modal::ThumbsDownActivePrediction>(),
95 TypeId::of::<rate_prediction_modal::NextEdit>(),
96 TypeId::of::<rate_prediction_modal::PreviousEdit>(),
97 ];
98
99 CommandPaletteFilter::update_global(cx, |filter, _cx| {
100 filter.hide_action_types(&rate_completion_action_types);
101 filter.hide_action_types(&reset_onboarding_action_types);
102 filter.hide_action_types(&[zed_actions::OpenZedPredictOnboarding.type_id()]);
103 });
104
105 cx.observe_global::<SettingsStore>(move |cx| {
106 let is_ai_disabled = DisableAiSettings::get_global(cx).disable_ai;
107 let has_feature_flag = cx.has_flag::<PredictEditsRatePredictionsFeatureFlag>();
108
109 CommandPaletteFilter::update_global(cx, |filter, _cx| {
110 if is_ai_disabled {
111 filter.hide_action_types(&all_action_types);
112 } else if has_feature_flag {
113 filter.show_action_types(&rate_completion_action_types);
114 } else {
115 filter.hide_action_types(&rate_completion_action_types);
116 }
117 });
118 })
119 .detach();
120
121 cx.observe_flag::<PredictEditsRatePredictionsFeatureFlag, _>(move |is_enabled, cx| {
122 if !DisableAiSettings::get_global(cx).disable_ai {
123 if is_enabled {
124 CommandPaletteFilter::update_global(cx, |filter, _cx| {
125 filter.show_action_types(&rate_completion_action_types);
126 });
127 } else {
128 CommandPaletteFilter::update_global(cx, |filter, _cx| {
129 filter.hide_action_types(&rate_completion_action_types);
130 });
131 }
132 }
133 })
134 .detach();
135}
136
137fn capture_example_as_markdown(
138 workspace: &mut Workspace,
139 window: &mut Window,
140 cx: &mut Context<Workspace>,
141) -> Option<()> {
142 let markdown_language = workspace
143 .app_state()
144 .languages
145 .language_for_name("Markdown");
146
147 let fs = workspace.app_state().fs.clone();
148 let project = workspace.project().clone();
149 let editor = workspace.active_item_as::<Editor>(cx)?;
150 let editor = editor.read(cx);
151 let (buffer, cursor_anchor) = editor
152 .buffer()
153 .read(cx)
154 .text_anchor_for_position(editor.selections.newest_anchor().head(), cx)?;
155 let ep_store = EditPredictionStore::try_global(cx)?;
156 let events = ep_store.update(cx, |store, cx| {
157 store.edit_history_for_project_with_pause_split_last_event(&project, cx)
158 });
159 let example = capture_example(
160 project.clone(),
161 buffer,
162 cursor_anchor,
163 events,
164 Vec::new(),
165 true,
166 cx,
167 )?;
168
169 let examples_dir = AllLanguageSettings::get_global(cx)
170 .edit_predictions
171 .examples_dir
172 .clone();
173
174 cx.spawn_in(window, async move |workspace_entity, cx| {
175 let markdown_language = markdown_language.await?;
176 let example_spec = example.await?;
177 let buffer = if let Some(dir) = examples_dir {
178 fs.create_dir(&dir).await.ok();
179 let mut path = dir.join(&example_spec.name.replace(' ', "--").replace(':', "-"));
180 path.set_extension("md");
181 project
182 .update(cx, |project, cx| project.open_local_buffer(&path, cx))
183 .await?
184 } else {
185 project
186 .update(cx, |project, cx| {
187 project.create_buffer(Some(markdown_language.clone()), false, cx)
188 })
189 .await?
190 };
191
192 buffer.update(cx, |buffer, cx| {
193 buffer.set_text(example_spec.to_markdown(), cx);
194 buffer.set_language(Some(markdown_language), cx);
195 });
196 workspace_entity.update_in(cx, |workspace, window, cx| {
197 workspace.add_item_to_active_pane(
198 Box::new(
199 cx.new(|cx| Editor::for_buffer(buffer, Some(project.clone()), window, cx)),
200 ),
201 None,
202 true,
203 window,
204 cx,
205 );
206 })
207 })
208 .detach_and_log_err(cx);
209 None
210}