tenant.openagents/omega

1mod edit_prediction_button;
2mod edit_prediction_context_view;
3mod rate_prediction_modal;
4
5use command_palette_hooks::CommandPaletteFilter;
6use edit_prediction::ResetOnboarding;
7use edit_prediction_context_view::EditPredictionContextView;
8use feature_flags::FeatureFlagAppExt as _;
9use gpui::actions;
10use project::DisableAiSettings;
11use rate_prediction_modal::RatePredictionsModal;
12use settings::{Settings as _, SettingsStore};
13use std::any::{Any as _, TypeId};
14use ui::{App, prelude::*};
15use workspace::{SplitDirection, Workspace};
16
17pub use edit_prediction_button::{
18    EditPredictionButton, ToggleMenu, get_available_providers, set_completion_provider,
19};
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    ]
37);
38
39pub fn init(cx: &mut App) {
40    feature_gate_predict_edits_actions(cx);
41
42    cx.observe_new(move |workspace: &mut Workspace, _, _cx| {
43        workspace.register_action(|workspace, _: &RatePredictions, window, cx| {
44            if cx.has_flag::<PredictEditsRatePredictionsFeatureFlag>() {
45                RatePredictionsModal::toggle(workspace, window, cx);
46            }
47        });
48
49        workspace.register_action_renderer(|div, _, _, cx| {
50            div.on_action(cx.listener(
51                move |workspace, _: &OpenEditPredictionContextView, window, cx| {
52                    let project = workspace.project();
53                    workspace.split_item(
54                        SplitDirection::Right,
55                        Box::new(cx.new(|cx| {
56                            EditPredictionContextView::new(
57                                project.clone(),
58                                workspace.client(),
59                                workspace.user_store(),
60                                window,
61                                cx,
62                            )
63                        })),
64                        window,
65                        cx,
66                    );
67                },
68            ))
69        });
70    })
71    .detach();
72}
73
74fn feature_gate_predict_edits_actions(cx: &mut App) {
75    let rate_completion_action_types = [TypeId::of::<RatePredictions>()];
76    let reset_onboarding_action_types = [TypeId::of::<ResetOnboarding>()];
77    let all_action_types = [
78        TypeId::of::<RatePredictions>(),
79        TypeId::of::<edit_prediction::ResetOnboarding>(),
80        zed_actions::OpenOmegaPredictOnboarding.type_id(),
81        TypeId::of::<edit_prediction::ClearHistory>(),
82        TypeId::of::<rate_prediction_modal::ThumbsUpActivePrediction>(),
83        TypeId::of::<rate_prediction_modal::ThumbsDownActivePrediction>(),
84        TypeId::of::<rate_prediction_modal::NextEdit>(),
85        TypeId::of::<rate_prediction_modal::PreviousEdit>(),
86    ];
87
88    CommandPaletteFilter::update_global(cx, |filter, _cx| {
89        filter.hide_action_types(&rate_completion_action_types);
90        filter.hide_action_types(&reset_onboarding_action_types);
91        filter.hide_action_types(&[zed_actions::OpenOmegaPredictOnboarding.type_id()]);
92    });
93
94    cx.observe_global::<SettingsStore>(move |cx| {
95        let is_ai_disabled = DisableAiSettings::get_global(cx).disable_ai;
96        let has_feature_flag = cx.has_flag::<PredictEditsRatePredictionsFeatureFlag>();
97
98        CommandPaletteFilter::update_global(cx, |filter, _cx| {
99            if is_ai_disabled {
100                filter.hide_action_types(&all_action_types);
101            } else if has_feature_flag {
102                filter.show_action_types(&rate_completion_action_types);
103            } else {
104                filter.hide_action_types(&rate_completion_action_types);
105            }
106        });
107    })
108    .detach();
109
110    cx.observe_flag::<PredictEditsRatePredictionsFeatureFlag, _>(move |value, cx| {
111        if !DisableAiSettings::get_global(cx).disable_ai {
112            if *value {
113                CommandPaletteFilter::update_global(cx, |filter, _cx| {
114                    filter.show_action_types(&rate_completion_action_types);
115                });
116            } else {
117                CommandPaletteFilter::update_global(cx, |filter, _cx| {
118                    filter.hide_action_types(&rate_completion_action_types);
119                });
120            }
121        }
122    })
123    .detach();
124}
125