1mod edit_prediction_button;
2mod edit_prediction_context_view;
3mod rate_prediction_modal;
4
5use command_palette_hooks::CommandPaletteFilter;
6use edit_prediction::{EditPredictionStore, ResetOnboarding, 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 div.on_action(cx.listener(
58 move |workspace, _: &OpenEditPredictionContextView, window, cx| {
59 let project = workspace.project();
60 workspace.split_item(
61 SplitDirection::Right,
62 Box::new(cx.new(|cx| {
63 EditPredictionContextView::new(
64 project.clone(),
65 workspace.client(),
66 workspace.user_store(),
67 window,
68 cx,
69 )
70 })),
71 window,
72 cx,
73 );
74 },
75 ))
76 });
77 })
78 .detach();
79}
80
81fn feature_gate_predict_edits_actions(cx: &mut App) {
82 let rate_completion_action_types = [TypeId::of::<RatePredictions>()];
83 let reset_onboarding_action_types = [TypeId::of::<ResetOnboarding>()];
84 let all_action_types = [
85 TypeId::of::<RatePredictions>(),
86 TypeId::of::<CaptureExample>(),
87 TypeId::of::<edit_prediction::ResetOnboarding>(),
88 zed_actions::OpenZedPredictOnboarding.type_id(),
89 TypeId::of::<edit_prediction::ClearHistory>(),
90 TypeId::of::<rate_prediction_modal::ThumbsUpActivePrediction>(),
91 TypeId::of::<rate_prediction_modal::ThumbsDownActivePrediction>(),
92 TypeId::of::<rate_prediction_modal::NextEdit>(),
93 TypeId::of::<rate_prediction_modal::PreviousEdit>(),
94 ];
95
96 CommandPaletteFilter::update_global(cx, |filter, _cx| {
97 filter.hide_action_types(&rate_completion_action_types);
98 filter.hide_action_types(&reset_onboarding_action_types);
99 filter.hide_action_types(&[zed_actions::OpenZedPredictOnboarding.type_id()]);
100 });
101
102 cx.observe_global::<SettingsStore>(move |cx| {
103 let is_ai_disabled = DisableAiSettings::get_global(cx).disable_ai;
104 let has_feature_flag = cx.has_flag::<PredictEditsRatePredictionsFeatureFlag>();
105
106 CommandPaletteFilter::update_global(cx, |filter, _cx| {
107 if is_ai_disabled {
108 filter.hide_action_types(&all_action_types);
109 } else if has_feature_flag {
110 filter.show_action_types(&rate_completion_action_types);
111 } else {
112 filter.hide_action_types(&rate_completion_action_types);
113 }
114 });
115 })
116 .detach();
117
118 cx.observe_flag::<PredictEditsRatePredictionsFeatureFlag, _>(move |is_enabled, cx| {
119 if !DisableAiSettings::get_global(cx).disable_ai {
120 if is_enabled {
121 CommandPaletteFilter::update_global(cx, |filter, _cx| {
122 filter.show_action_types(&rate_completion_action_types);
123 });
124 } else {
125 CommandPaletteFilter::update_global(cx, |filter, _cx| {
126 filter.hide_action_types(&rate_completion_action_types);
127 });
128 }
129 }
130 })
131 .detach();
132}
133
134fn capture_example_as_markdown(
135 workspace: &mut Workspace,
136 window: &mut Window,
137 cx: &mut Context<Workspace>,
138) -> Option<()> {
139 let markdown_language = workspace
140 .app_state()
141 .languages
142 .language_for_name("Markdown");
143
144 let fs = workspace.app_state().fs.clone();
145 let project = workspace.project().clone();
146 let editor = workspace.active_item_as::<Editor>(cx)?;
147 let editor = editor.read(cx);
148 let (buffer, cursor_anchor) = editor
149 .buffer()
150 .read(cx)
151 .text_anchor_for_position(editor.selections.newest_anchor().head(), cx)?;
152 let ep_store = EditPredictionStore::try_global(cx)?;
153 let events = ep_store.update(cx, |store, cx| store.edit_history_for_project(&project, cx));
154 let example = capture_example(project.clone(), buffer, cursor_anchor, events, true, cx)?;
155
156 let examples_dir = AllLanguageSettings::get_global(cx)
157 .edit_predictions
158 .examples_dir
159 .clone();
160
161 cx.spawn_in(window, async move |workspace_entity, cx| {
162 let markdown_language = markdown_language.await?;
163 let example_spec = example.await?;
164 let buffer = if let Some(dir) = examples_dir {
165 fs.create_dir(&dir).await.ok();
166 let mut path = dir.join(&example_spec.name.replace(' ', "--").replace(':', "-"));
167 path.set_extension("md");
168 project
169 .update(cx, |project, cx| project.open_local_buffer(&path, cx))
170 .await?
171 } else {
172 project
173 .update(cx, |project, cx| {
174 project.create_buffer(Some(markdown_language.clone()), false, cx)
175 })
176 .await?
177 };
178
179 buffer.update(cx, |buffer, cx| {
180 buffer.set_text(example_spec.to_markdown(), cx);
181 buffer.set_language(Some(markdown_language), cx);
182 });
183 workspace_entity.update_in(cx, |workspace, window, cx| {
184 workspace.add_item_to_active_pane(
185 Box::new(
186 cx.new(|cx| Editor::for_buffer(buffer, Some(project.clone()), window, cx)),
187 ),
188 None,
189 true,
190 window,
191 cx,
192 );
193 })
194 })
195 .detach_and_log_err(cx);
196 None
197}