1mod edit_prediction_button;
2mod edit_prediction_context_view;
3mod rate_prediction_modal;
4
5use command_palette_hooks::CommandPaletteFilter;
6use edit_prediction::{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::{EditPredictionButton, ToggleMenu};
20
21use crate::rate_prediction_modal::PredictEditsRatePredictionsFeatureFlag;
22
23actions!(
24 dev,
25 [
26 /// Opens the edit prediction context view.
27 OpenEditPredictionContextView,
28 ]
29);
30
31actions!(
32 edit_prediction,
33 [
34 /// Opens the rate completions modal.
35 RatePredictions,
36 /// Captures an ExampleSpec from the current editing session and opens it as Markdown.
37 CaptureExample,
38 ]
39);
40
41pub fn init(cx: &mut App) {
42 feature_gate_predict_edits_actions(cx);
43
44 cx.observe_new(move |workspace: &mut Workspace, _, _cx| {
45 workspace.register_action(|workspace, _: &RatePredictions, window, cx| {
46 if cx.has_flag::<PredictEditsRatePredictionsFeatureFlag>() {
47 RatePredictionsModal::toggle(workspace, window, cx);
48 }
49 });
50
51 workspace.register_action(|workspace, _: &CaptureExample, window, cx| {
52 capture_example_as_markdown(workspace, window, cx);
53 });
54 workspace.register_action_renderer(|div, _, _, cx| {
55 let has_flag = cx.has_flag::<Zeta2FeatureFlag>();
56 div.when(has_flag, |div| {
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 })
79 .detach();
80}
81
82fn feature_gate_predict_edits_actions(cx: &mut App) {
83 let rate_completion_action_types = [TypeId::of::<RatePredictions>()];
84 let reset_onboarding_action_types = [TypeId::of::<ResetOnboarding>()];
85 let all_action_types = [
86 TypeId::of::<RatePredictions>(),
87 TypeId::of::<CaptureExample>(),
88 TypeId::of::<edit_prediction::ResetOnboarding>(),
89 zed_actions::OpenZedPredictOnboarding.type_id(),
90 TypeId::of::<edit_prediction::ClearHistory>(),
91 TypeId::of::<rate_prediction_modal::ThumbsUpActivePrediction>(),
92 TypeId::of::<rate_prediction_modal::ThumbsDownActivePrediction>(),
93 TypeId::of::<rate_prediction_modal::NextEdit>(),
94 TypeId::of::<rate_prediction_modal::PreviousEdit>(),
95 ];
96
97 CommandPaletteFilter::update_global(cx, |filter, _cx| {
98 filter.hide_action_types(&rate_completion_action_types);
99 filter.hide_action_types(&reset_onboarding_action_types);
100 filter.hide_action_types(&[zed_actions::OpenZedPredictOnboarding.type_id()]);
101 });
102
103 cx.observe_global::<SettingsStore>(move |cx| {
104 let is_ai_disabled = DisableAiSettings::get_global(cx).disable_ai;
105 let has_feature_flag = cx.has_flag::<PredictEditsRatePredictionsFeatureFlag>();
106
107 CommandPaletteFilter::update_global(cx, |filter, _cx| {
108 if is_ai_disabled {
109 filter.hide_action_types(&all_action_types);
110 } else if has_feature_flag {
111 filter.show_action_types(&rate_completion_action_types);
112 } else {
113 filter.hide_action_types(&rate_completion_action_types);
114 }
115 });
116 })
117 .detach();
118
119 cx.observe_flag::<PredictEditsRatePredictionsFeatureFlag, _>(move |is_enabled, cx| {
120 if !DisableAiSettings::get_global(cx).disable_ai {
121 if is_enabled {
122 CommandPaletteFilter::update_global(cx, |filter, _cx| {
123 filter.show_action_types(&rate_completion_action_types);
124 });
125 } else {
126 CommandPaletteFilter::update_global(cx, |filter, _cx| {
127 filter.hide_action_types(&rate_completion_action_types);
128 });
129 }
130 }
131 })
132 .detach();
133}
134
135fn capture_example_as_markdown(
136 workspace: &mut Workspace,
137 window: &mut Window,
138 cx: &mut Context<Workspace>,
139) -> Option<()> {
140 let markdown_language = workspace
141 .app_state()
142 .languages
143 .language_for_name("Markdown");
144
145 let fs = workspace.app_state().fs.clone();
146 let project = workspace.project().clone();
147 let editor = workspace.active_item_as::<Editor>(cx)?;
148 let editor = editor.read(cx);
149 let (buffer, cursor_anchor) = editor
150 .buffer()
151 .read(cx)
152 .text_anchor_for_position(editor.selections.newest_anchor().head(), cx)?;
153 let example = capture_example(project.clone(), buffer, cursor_anchor, true, cx)?;
154
155 let examples_dir = AllLanguageSettings::get_global(cx)
156 .edit_predictions
157 .examples_dir
158 .clone();
159
160 cx.spawn_in(window, async move |workspace_entity, cx| {
161 let markdown_language = markdown_language.await?;
162 let example_spec = example.await?;
163 let buffer = if let Some(dir) = examples_dir {
164 fs.create_dir(&dir).await.ok();
165 let mut path = dir.join(&example_spec.name.replace(' ', "--").replace(':', "-"));
166 path.set_extension("md");
167 project.update(cx, |project, cx| project.open_local_buffer(&path, cx))
168 } else {
169 project.update(cx, |project, cx| project.create_buffer(false, cx))
170 }?
171 .await?;
172
173 buffer.update(cx, |buffer, cx| {
174 buffer.set_text(example_spec.to_markdown(), cx);
175 buffer.set_language(Some(markdown_language), cx);
176 })?;
177 workspace_entity.update_in(cx, |workspace, window, cx| {
178 workspace.add_item_to_active_pane(
179 Box::new(
180 cx.new(|cx| Editor::for_buffer(buffer, Some(project.clone()), window, cx)),
181 ),
182 None,
183 true,
184 window,
185 cx,
186 );
187 })
188 })
189 .detach_and_log_err(cx);
190 None
191}