edit_prediction_button.rs

   1use anyhow::Result;
   2use client::{Client, UserStore, zed_urls};
   3use cloud_llm_client::UsageLimit;
   4use codestral::CodestralCompletionProvider;
   5use copilot::{Copilot, Status};
   6use editor::{Editor, SelectionEffects, actions::ShowEditPrediction, scroll::Autoscroll};
   7use feature_flags::{FeatureFlagAppExt, PredictEditsRateCompletionsFeatureFlag};
   8use fs::Fs;
   9use gpui::{
  10    Action, Animation, AnimationExt, App, AsyncWindowContext, Corner, Entity, FocusHandle,
  11    Focusable, IntoElement, ParentElement, Render, Subscription, WeakEntity, actions, div,
  12    pulsating_between,
  13};
  14use indoc::indoc;
  15use language::{
  16    EditPredictionsMode, File, Language,
  17    language_settings::{self, AllLanguageSettings, EditPredictionProvider, all_language_settings},
  18};
  19use project::DisableAiSettings;
  20use regex::Regex;
  21use settings::{Settings, SettingsStore, update_settings_file};
  22use std::{
  23    sync::{Arc, LazyLock},
  24    time::Duration,
  25};
  26use supermaven::{AccountStatus, Supermaven};
  27use ui::{
  28    Clickable, ContextMenu, ContextMenuEntry, DocumentationEdge, DocumentationSide, IconButton,
  29    IconButtonShape, Indicator, PopoverMenu, PopoverMenuHandle, ProgressBar, Tooltip, prelude::*,
  30};
  31use workspace::{
  32    StatusItemView, Toast, Workspace, create_and_open_local_file, item::ItemHandle,
  33    notifications::NotificationId,
  34};
  35use zed_actions::OpenBrowser;
  36use zeta::RateCompletions;
  37
  38actions!(
  39    edit_prediction,
  40    [
  41        /// Toggles the edit prediction menu.
  42        ToggleMenu
  43    ]
  44);
  45
  46const COPILOT_SETTINGS_URL: &str = "https://github.com/settings/copilot";
  47const PRIVACY_DOCS: &str = "https://zed.dev/docs/ai/privacy-and-security";
  48
  49struct CopilotErrorToast;
  50
  51pub struct EditPredictionButton {
  52    editor_subscription: Option<(Subscription, usize)>,
  53    editor_enabled: Option<bool>,
  54    editor_show_predictions: bool,
  55    editor_focus_handle: Option<FocusHandle>,
  56    language: Option<Arc<Language>>,
  57    file: Option<Arc<dyn File>>,
  58    edit_prediction_provider: Option<Arc<dyn edit_prediction::EditPredictionProviderHandle>>,
  59    fs: Arc<dyn Fs>,
  60    user_store: Entity<UserStore>,
  61    popover_menu_handle: PopoverMenuHandle<ContextMenu>,
  62}
  63
  64enum SupermavenButtonStatus {
  65    Ready,
  66    Errored(String),
  67    NeedsActivation(String),
  68    Initializing,
  69}
  70
  71impl Render for EditPredictionButton {
  72    fn render(&mut self, _: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
  73        // Return empty div if AI is disabled
  74        if DisableAiSettings::get_global(cx).disable_ai {
  75            return div().hidden();
  76        }
  77
  78        let all_language_settings = all_language_settings(None, cx);
  79
  80        match all_language_settings.edit_predictions.provider {
  81            EditPredictionProvider::None => div().hidden(),
  82
  83            EditPredictionProvider::Copilot => {
  84                let Some(copilot) = Copilot::global(cx) else {
  85                    return div().hidden();
  86                };
  87                let status = copilot.read(cx).status();
  88
  89                let enabled = self.editor_enabled.unwrap_or(false);
  90
  91                let icon = match status {
  92                    Status::Error(_) => IconName::CopilotError,
  93                    Status::Authorized => {
  94                        if enabled {
  95                            IconName::Copilot
  96                        } else {
  97                            IconName::CopilotDisabled
  98                        }
  99                    }
 100                    _ => IconName::CopilotInit,
 101                };
 102
 103                if let Status::Error(e) = status {
 104                    return div().child(
 105                        IconButton::new("copilot-error", icon)
 106                            .icon_size(IconSize::Small)
 107                            .on_click(cx.listener(move |_, _, window, cx| {
 108                                if let Some(workspace) = window.root::<Workspace>().flatten() {
 109                                    workspace.update(cx, |workspace, cx| {
 110                                        workspace.show_toast(
 111                                            Toast::new(
 112                                                NotificationId::unique::<CopilotErrorToast>(),
 113                                                format!("Copilot can't be started: {}", e),
 114                                            )
 115                                            .on_click(
 116                                                "Reinstall Copilot",
 117                                                |window, cx| {
 118                                                    copilot::reinstall_and_sign_in(window, cx)
 119                                                },
 120                                            ),
 121                                            cx,
 122                                        );
 123                                    });
 124                                }
 125                            }))
 126                            .tooltip(|_window, cx| {
 127                                Tooltip::for_action("GitHub Copilot", &ToggleMenu, cx)
 128                            }),
 129                    );
 130                }
 131                let this = cx.weak_entity();
 132
 133                div().child(
 134                    PopoverMenu::new("copilot")
 135                        .menu(move |window, cx| {
 136                            let current_status = Copilot::global(cx)?.read(cx).status();
 137                            match current_status {
 138                                Status::Authorized => this.update(cx, |this, cx| {
 139                                    this.build_copilot_context_menu(window, cx)
 140                                }),
 141                                _ => this.update(cx, |this, cx| {
 142                                    this.build_copilot_start_menu(window, cx)
 143                                }),
 144                            }
 145                            .ok()
 146                        })
 147                        .anchor(Corner::BottomRight)
 148                        .trigger_with_tooltip(
 149                            IconButton::new("copilot-icon", icon),
 150                            |_window, cx| Tooltip::for_action("GitHub Copilot", &ToggleMenu, cx),
 151                        )
 152                        .with_handle(self.popover_menu_handle.clone()),
 153                )
 154            }
 155
 156            EditPredictionProvider::Supermaven => {
 157                let Some(supermaven) = Supermaven::global(cx) else {
 158                    return div();
 159                };
 160
 161                let supermaven = supermaven.read(cx);
 162
 163                let status = match supermaven {
 164                    Supermaven::Starting => SupermavenButtonStatus::Initializing,
 165                    Supermaven::FailedDownload { error } => {
 166                        SupermavenButtonStatus::Errored(error.to_string())
 167                    }
 168                    Supermaven::Spawned(agent) => {
 169                        let account_status = agent.account_status.clone();
 170                        match account_status {
 171                            AccountStatus::NeedsActivation { activate_url } => {
 172                                SupermavenButtonStatus::NeedsActivation(activate_url)
 173                            }
 174                            AccountStatus::Unknown => SupermavenButtonStatus::Initializing,
 175                            AccountStatus::Ready => SupermavenButtonStatus::Ready,
 176                        }
 177                    }
 178                    Supermaven::Error { error } => {
 179                        SupermavenButtonStatus::Errored(error.to_string())
 180                    }
 181                };
 182
 183                let icon = status.to_icon();
 184                let tooltip_text = status.to_tooltip();
 185                let has_menu = status.has_menu();
 186                let this = cx.weak_entity();
 187                let fs = self.fs.clone();
 188
 189                div().child(
 190                    PopoverMenu::new("supermaven")
 191                        .menu(move |window, cx| match &status {
 192                            SupermavenButtonStatus::NeedsActivation(activate_url) => {
 193                                Some(ContextMenu::build(window, cx, |menu, _, _| {
 194                                    let fs = fs.clone();
 195                                    let activate_url = activate_url.clone();
 196
 197                                    menu.entry("Sign In", None, move |_, cx| {
 198                                        cx.open_url(activate_url.as_str())
 199                                    })
 200                                    .entry(
 201                                        "Use Zed AI",
 202                                        None,
 203                                        move |_, cx| {
 204                                            set_completion_provider(
 205                                                fs.clone(),
 206                                                cx,
 207                                                EditPredictionProvider::Zed,
 208                                            )
 209                                        },
 210                                    )
 211                                }))
 212                            }
 213                            SupermavenButtonStatus::Ready => this
 214                                .update(cx, |this, cx| {
 215                                    this.build_supermaven_context_menu(window, cx)
 216                                })
 217                                .ok(),
 218                            _ => None,
 219                        })
 220                        .anchor(Corner::BottomRight)
 221                        .trigger_with_tooltip(
 222                            IconButton::new("supermaven-icon", icon),
 223                            move |window, cx| {
 224                                if has_menu {
 225                                    Tooltip::for_action(tooltip_text.clone(), &ToggleMenu, cx)
 226                                } else {
 227                                    Tooltip::text(tooltip_text.clone())(window, cx)
 228                                }
 229                            },
 230                        )
 231                        .with_handle(self.popover_menu_handle.clone()),
 232                )
 233            }
 234
 235            EditPredictionProvider::Codestral => {
 236                let enabled = self.editor_enabled.unwrap_or(true);
 237                let has_api_key = CodestralCompletionProvider::has_api_key(cx);
 238                let fs = self.fs.clone();
 239                let this = cx.weak_entity();
 240
 241                div().child(
 242                    PopoverMenu::new("codestral")
 243                        .menu(move |window, cx| {
 244                            if has_api_key {
 245                                this.update(cx, |this, cx| {
 246                                    this.build_codestral_context_menu(window, cx)
 247                                })
 248                                .ok()
 249                            } else {
 250                                Some(ContextMenu::build(window, cx, |menu, _, _| {
 251                                    let fs = fs.clone();
 252
 253                                    menu.entry(
 254                                        "Configure Codestral API Key",
 255                                        None,
 256                                        move |window, cx| {
 257                                            window.dispatch_action(
 258                                                zed_actions::agent::OpenSettings.boxed_clone(),
 259                                                cx,
 260                                            );
 261                                        },
 262                                    )
 263                                    .separator()
 264                                    .entry(
 265                                        "Use Zed AI instead",
 266                                        None,
 267                                        move |_, cx| {
 268                                            set_completion_provider(
 269                                                fs.clone(),
 270                                                cx,
 271                                                EditPredictionProvider::Zed,
 272                                            )
 273                                        },
 274                                    )
 275                                }))
 276                            }
 277                        })
 278                        .anchor(Corner::BottomRight)
 279                        .trigger_with_tooltip(
 280                            IconButton::new("codestral-icon", IconName::AiMistral)
 281                                .shape(IconButtonShape::Square)
 282                                .when(!has_api_key, |this| {
 283                                    this.indicator(Indicator::dot().color(Color::Error))
 284                                        .indicator_border_color(Some(
 285                                            cx.theme().colors().status_bar_background,
 286                                        ))
 287                                })
 288                                .when(has_api_key && !enabled, |this| {
 289                                    this.indicator(Indicator::dot().color(Color::Ignored))
 290                                        .indicator_border_color(Some(
 291                                            cx.theme().colors().status_bar_background,
 292                                        ))
 293                                }),
 294                            move |_window, cx| Tooltip::for_action("Codestral", &ToggleMenu, cx),
 295                        )
 296                        .with_handle(self.popover_menu_handle.clone()),
 297                )
 298            }
 299
 300            EditPredictionProvider::Zed => {
 301                let enabled = self.editor_enabled.unwrap_or(true);
 302
 303                let zeta_icon = if enabled {
 304                    IconName::ZedPredict
 305                } else {
 306                    IconName::ZedPredictDisabled
 307                };
 308
 309                if zeta::should_show_upsell_modal() {
 310                    let tooltip_meta = if self.user_store.read(cx).current_user().is_some() {
 311                        "Choose a Plan"
 312                    } else {
 313                        "Sign In"
 314                    };
 315
 316                    return div().child(
 317                        IconButton::new("zed-predict-pending-button", zeta_icon)
 318                            .shape(IconButtonShape::Square)
 319                            .indicator(Indicator::dot().color(Color::Muted))
 320                            .indicator_border_color(Some(cx.theme().colors().status_bar_background))
 321                            .tooltip(move |_window, cx| {
 322                                Tooltip::with_meta("Edit Predictions", None, tooltip_meta, cx)
 323                            })
 324                            .on_click(cx.listener(move |_, _, window, cx| {
 325                                telemetry::event!(
 326                                    "Pending ToS Clicked",
 327                                    source = "Edit Prediction Status Button"
 328                                );
 329                                window.dispatch_action(
 330                                    zed_actions::OpenZedPredictOnboarding.boxed_clone(),
 331                                    cx,
 332                                );
 333                            })),
 334                    );
 335                }
 336
 337                let mut over_limit = false;
 338
 339                if let Some(usage) = self
 340                    .edit_prediction_provider
 341                    .as_ref()
 342                    .and_then(|provider| provider.usage(cx))
 343                {
 344                    over_limit = usage.over_limit()
 345                }
 346
 347                let show_editor_predictions = self.editor_show_predictions;
 348
 349                let icon_button = IconButton::new("zed-predict-pending-button", zeta_icon)
 350                    .shape(IconButtonShape::Square)
 351                    .when(
 352                        enabled && (!show_editor_predictions || over_limit),
 353                        |this| {
 354                            this.indicator(Indicator::dot().when_else(
 355                                over_limit,
 356                                |dot| dot.color(Color::Error),
 357                                |dot| dot.color(Color::Muted),
 358                            ))
 359                            .indicator_border_color(Some(cx.theme().colors().status_bar_background))
 360                        },
 361                    )
 362                    .when(!self.popover_menu_handle.is_deployed(), |element| {
 363                        element.tooltip(move |_window, cx| {
 364                            if enabled {
 365                                if show_editor_predictions {
 366                                    Tooltip::for_action("Edit Prediction", &ToggleMenu, cx)
 367                                } else {
 368                                    Tooltip::with_meta(
 369                                        "Edit Prediction",
 370                                        Some(&ToggleMenu),
 371                                        "Hidden For This File",
 372                                        cx,
 373                                    )
 374                                }
 375                            } else {
 376                                Tooltip::with_meta(
 377                                    "Edit Prediction",
 378                                    Some(&ToggleMenu),
 379                                    "Disabled For This File",
 380                                    cx,
 381                                )
 382                            }
 383                        })
 384                    });
 385
 386                let this = cx.weak_entity();
 387
 388                let mut popover_menu = PopoverMenu::new("zeta")
 389                    .menu(move |window, cx| {
 390                        this.update(cx, |this, cx| this.build_zeta_context_menu(window, cx))
 391                            .ok()
 392                    })
 393                    .anchor(Corner::BottomRight)
 394                    .with_handle(self.popover_menu_handle.clone());
 395
 396                let is_refreshing = self
 397                    .edit_prediction_provider
 398                    .as_ref()
 399                    .is_some_and(|provider| provider.is_refreshing(cx));
 400
 401                if is_refreshing {
 402                    popover_menu = popover_menu.trigger(
 403                        icon_button.with_animation(
 404                            "pulsating-label",
 405                            Animation::new(Duration::from_secs(2))
 406                                .repeat()
 407                                .with_easing(pulsating_between(0.2, 1.0)),
 408                            |icon_button, delta| icon_button.alpha(delta),
 409                        ),
 410                    );
 411                } else {
 412                    popover_menu = popover_menu.trigger(icon_button);
 413                }
 414
 415                div().child(popover_menu.into_any_element())
 416            }
 417        }
 418    }
 419}
 420
 421impl EditPredictionButton {
 422    pub fn new(
 423        fs: Arc<dyn Fs>,
 424        user_store: Entity<UserStore>,
 425        popover_menu_handle: PopoverMenuHandle<ContextMenu>,
 426        client: Arc<Client>,
 427        cx: &mut Context<Self>,
 428    ) -> Self {
 429        if let Some(copilot) = Copilot::global(cx) {
 430            cx.observe(&copilot, |_, _, cx| cx.notify()).detach()
 431        }
 432
 433        cx.observe_global::<SettingsStore>(move |_, cx| cx.notify())
 434            .detach();
 435
 436        CodestralCompletionProvider::ensure_api_key_loaded(client.http_client(), cx);
 437
 438        Self {
 439            editor_subscription: None,
 440            editor_enabled: None,
 441            editor_show_predictions: true,
 442            editor_focus_handle: None,
 443            language: None,
 444            file: None,
 445            edit_prediction_provider: None,
 446            user_store,
 447            popover_menu_handle,
 448            fs,
 449        }
 450    }
 451
 452    fn get_available_providers(&self, cx: &App) -> Vec<EditPredictionProvider> {
 453        let mut providers = Vec::new();
 454
 455        providers.push(EditPredictionProvider::Zed);
 456
 457        if let Some(copilot) = Copilot::global(cx) {
 458            if matches!(copilot.read(cx).status(), Status::Authorized) {
 459                providers.push(EditPredictionProvider::Copilot);
 460            }
 461        }
 462
 463        if let Some(supermaven) = Supermaven::global(cx) {
 464            if let Supermaven::Spawned(agent) = supermaven.read(cx) {
 465                if matches!(agent.account_status, AccountStatus::Ready) {
 466                    providers.push(EditPredictionProvider::Supermaven);
 467                }
 468            }
 469        }
 470
 471        if CodestralCompletionProvider::has_api_key(cx) {
 472            providers.push(EditPredictionProvider::Codestral);
 473        }
 474
 475        providers
 476    }
 477
 478    fn add_provider_switching_section(
 479        &self,
 480        mut menu: ContextMenu,
 481        current_provider: EditPredictionProvider,
 482        cx: &App,
 483    ) -> ContextMenu {
 484        let available_providers = self.get_available_providers(cx);
 485
 486        let other_providers: Vec<_> = available_providers
 487            .into_iter()
 488            .filter(|p| *p != current_provider && *p != EditPredictionProvider::None)
 489            .collect();
 490
 491        if !other_providers.is_empty() {
 492            menu = menu.separator().header("Switch Providers");
 493
 494            for provider in other_providers {
 495                let fs = self.fs.clone();
 496
 497                menu = match provider {
 498                    EditPredictionProvider::Zed => menu.item(
 499                        ContextMenuEntry::new("Zed AI")
 500                            .documentation_aside(
 501                                DocumentationSide::Left,
 502                                DocumentationEdge::Top,
 503                                |_| {
 504                                    Label::new("Zed's edit prediction is powered by Zeta, an open-source, dataset mode.")
 505                                        .into_any_element()
 506                                },
 507                            )
 508                            .handler(move |_, cx| {
 509                                set_completion_provider(fs.clone(), cx, provider);
 510                            }),
 511                    ),
 512                    EditPredictionProvider::Copilot => {
 513                        menu.entry("GitHub Copilot", None, move |_, cx| {
 514                            set_completion_provider(fs.clone(), cx, provider);
 515                        })
 516                    }
 517                    EditPredictionProvider::Supermaven => {
 518                        menu.entry("Supermaven", None, move |_, cx| {
 519                            set_completion_provider(fs.clone(), cx, provider);
 520                        })
 521                    }
 522                    EditPredictionProvider::Codestral => {
 523                        menu.entry("Codestral", None, move |_, cx| {
 524                            set_completion_provider(fs.clone(), cx, provider);
 525                        })
 526                    }
 527                    EditPredictionProvider::None => continue,
 528                };
 529            }
 530        }
 531
 532        menu
 533    }
 534
 535    pub fn build_copilot_start_menu(
 536        &mut self,
 537        window: &mut Window,
 538        cx: &mut Context<Self>,
 539    ) -> Entity<ContextMenu> {
 540        let fs = self.fs.clone();
 541        ContextMenu::build(window, cx, |menu, _, _| {
 542            menu.entry("Sign In to Copilot", None, copilot::initiate_sign_in)
 543                .entry("Disable Copilot", None, {
 544                    let fs = fs.clone();
 545                    move |_window, cx| hide_copilot(fs.clone(), cx)
 546                })
 547                .separator()
 548                .entry("Use Zed AI", None, {
 549                    let fs = fs.clone();
 550                    move |_window, cx| {
 551                        set_completion_provider(fs.clone(), cx, EditPredictionProvider::Zed)
 552                    }
 553                })
 554        })
 555    }
 556
 557    pub fn build_language_settings_menu(
 558        &self,
 559        mut menu: ContextMenu,
 560        window: &Window,
 561        cx: &mut App,
 562    ) -> ContextMenu {
 563        let fs = self.fs.clone();
 564        let line_height = window.line_height();
 565
 566        menu = menu.header("Show Edit Predictions For");
 567
 568        let language_state = self.language.as_ref().map(|language| {
 569            (
 570                language.clone(),
 571                language_settings::language_settings(Some(language.name()), None, cx)
 572                    .show_edit_predictions,
 573            )
 574        });
 575
 576        if let Some(editor_focus_handle) = self.editor_focus_handle.clone() {
 577            let entry = ContextMenuEntry::new("This Buffer")
 578                .toggleable(IconPosition::Start, self.editor_show_predictions)
 579                .action(Box::new(editor::actions::ToggleEditPrediction))
 580                .handler(move |window, cx| {
 581                    editor_focus_handle.dispatch_action(
 582                        &editor::actions::ToggleEditPrediction,
 583                        window,
 584                        cx,
 585                    );
 586                });
 587
 588            match language_state.clone() {
 589                Some((language, false)) => {
 590                    menu = menu.item(
 591                        entry
 592                            .disabled(true)
 593                            .documentation_aside(DocumentationSide::Left, DocumentationEdge::Top, move |_cx| {
 594                                Label::new(format!("Edit predictions cannot be toggled for this buffer because they are disabled for {}", language.name()))
 595                                    .into_any_element()
 596                            })
 597                    );
 598                }
 599                Some(_) | None => menu = menu.item(entry),
 600            }
 601        }
 602
 603        if let Some((language, language_enabled)) = language_state {
 604            let fs = fs.clone();
 605
 606            menu = menu.toggleable_entry(
 607                language.name(),
 608                language_enabled,
 609                IconPosition::Start,
 610                None,
 611                move |_, cx| {
 612                    toggle_show_edit_predictions_for_language(language.clone(), fs.clone(), cx)
 613                },
 614            );
 615        }
 616
 617        let settings = AllLanguageSettings::get_global(cx);
 618
 619        let globally_enabled = settings.show_edit_predictions(None, cx);
 620        let entry = ContextMenuEntry::new("All Files")
 621            .toggleable(IconPosition::Start, globally_enabled)
 622            .action(workspace::ToggleEditPrediction.boxed_clone())
 623            .handler(|window, cx| {
 624                window.dispatch_action(workspace::ToggleEditPrediction.boxed_clone(), cx)
 625            });
 626        menu = menu.item(entry);
 627
 628        let provider = settings.edit_predictions.provider;
 629        let current_mode = settings.edit_predictions_mode();
 630        let subtle_mode = matches!(current_mode, EditPredictionsMode::Subtle);
 631        let eager_mode = matches!(current_mode, EditPredictionsMode::Eager);
 632
 633        if matches!(
 634            provider,
 635            EditPredictionProvider::Zed
 636                | EditPredictionProvider::Copilot
 637                | EditPredictionProvider::Supermaven
 638                | EditPredictionProvider::Codestral
 639        ) {
 640            menu = menu
 641                .separator()
 642                .header("Display Modes")
 643                .item(
 644                    ContextMenuEntry::new("Eager")
 645                        .toggleable(IconPosition::Start, eager_mode)
 646                        .documentation_aside(DocumentationSide::Left, DocumentationEdge::Top, move |_| {
 647                            Label::new("Display predictions inline when there are no language server completions available.").into_any_element()
 648                        })
 649                        .handler({
 650                            let fs = fs.clone();
 651                            move |_, cx| {
 652                                toggle_edit_prediction_mode(fs.clone(), EditPredictionsMode::Eager, cx)
 653                            }
 654                        }),
 655                )
 656                .item(
 657                    ContextMenuEntry::new("Subtle")
 658                        .toggleable(IconPosition::Start, subtle_mode)
 659                        .documentation_aside(DocumentationSide::Left, DocumentationEdge::Top, move |_| {
 660                            Label::new("Display predictions inline only when holding a modifier key (alt by default).").into_any_element()
 661                        })
 662                        .handler({
 663                            let fs = fs.clone();
 664                            move |_, cx| {
 665                                toggle_edit_prediction_mode(fs.clone(), EditPredictionsMode::Subtle, cx)
 666                            }
 667                        }),
 668                );
 669        }
 670
 671        menu = menu.separator().header("Privacy");
 672
 673        if let Some(provider) = &self.edit_prediction_provider {
 674            let data_collection = provider.data_collection_state(cx);
 675
 676            if data_collection.is_supported() {
 677                let provider = provider.clone();
 678                let enabled = data_collection.is_enabled();
 679                let is_open_source = data_collection.is_project_open_source();
 680                let is_collecting = data_collection.is_enabled();
 681                let (icon_name, icon_color) = if is_open_source && is_collecting {
 682                    (IconName::Check, Color::Success)
 683                } else {
 684                    (IconName::Check, Color::Accent)
 685                };
 686
 687                menu = menu.item(
 688                    ContextMenuEntry::new("Training Data Collection")
 689                        .toggleable(IconPosition::Start, data_collection.is_enabled())
 690                        .icon(icon_name)
 691                        .icon_color(icon_color)
 692                        .documentation_aside(DocumentationSide::Left, DocumentationEdge::Top, move |cx| {
 693                            let (msg, label_color, icon_name, icon_color) = match (is_open_source, is_collecting) {
 694                                (true, true) => (
 695                                    "Project identified as open source, and you're sharing data.",
 696                                    Color::Default,
 697                                    IconName::Check,
 698                                    Color::Success,
 699                                ),
 700                                (true, false) => (
 701                                    "Project identified as open source, but you're not sharing data.",
 702                                    Color::Muted,
 703                                    IconName::Close,
 704                                    Color::Muted,
 705                                ),
 706                                (false, true) => (
 707                                    "Project not identified as open source. No data captured.",
 708                                    Color::Muted,
 709                                    IconName::Close,
 710                                    Color::Muted,
 711                                ),
 712                                (false, false) => (
 713                                    "Project not identified as open source, and setting turned off.",
 714                                    Color::Muted,
 715                                    IconName::Close,
 716                                    Color::Muted,
 717                                ),
 718                            };
 719                            v_flex()
 720                                .gap_2()
 721                                .child(
 722                                    Label::new(indoc!{
 723                                        "Help us improve our open dataset model by sharing data from open source repositories. \
 724                                        Zed must detect a license file in your repo for this setting to take effect. \
 725                                        Files with sensitive data and secrets are excluded by default."
 726                                    })
 727                                )
 728                                .child(
 729                                    h_flex()
 730                                        .items_start()
 731                                        .pt_2()
 732                                        .pr_1()
 733                                        .flex_1()
 734                                        .gap_1p5()
 735                                        .border_t_1()
 736                                        .border_color(cx.theme().colors().border_variant)
 737                                        .child(h_flex().flex_shrink_0().h(line_height).child(Icon::new(icon_name).size(IconSize::XSmall).color(icon_color)))
 738                                        .child(div().child(msg).w_full().text_sm().text_color(label_color.color(cx)))
 739                                )
 740                                .into_any_element()
 741                        })
 742                        .handler(move |_, cx| {
 743                            provider.toggle_data_collection(cx);
 744
 745                            if !enabled {
 746                                telemetry::event!(
 747                                    "Data Collection Enabled",
 748                                    source = "Edit Prediction Status Menu"
 749                                );
 750                            } else {
 751                                telemetry::event!(
 752                                    "Data Collection Disabled",
 753                                    source = "Edit Prediction Status Menu"
 754                                );
 755                            }
 756                        })
 757                );
 758
 759                if is_collecting && !is_open_source {
 760                    menu = menu.item(
 761                        ContextMenuEntry::new("No data captured.")
 762                            .disabled(true)
 763                            .icon(IconName::Close)
 764                            .icon_color(Color::Error)
 765                            .icon_size(IconSize::Small),
 766                    );
 767                }
 768            }
 769        }
 770
 771        menu = menu.item(
 772            ContextMenuEntry::new("Configure Excluded Files")
 773                .icon(IconName::LockOutlined)
 774                .icon_color(Color::Muted)
 775                .documentation_aside(DocumentationSide::Left, DocumentationEdge::Top, |_| {
 776                    Label::new(indoc!{"
 777                        Open your settings to add sensitive paths for which Zed will never predict edits."}).into_any_element()
 778                })
 779                .handler(move |window, cx| {
 780                    if let Some(workspace) = window.root().flatten() {
 781                        let workspace = workspace.downgrade();
 782                        window
 783                            .spawn(cx, async |cx| {
 784                                open_disabled_globs_setting_in_editor(
 785                                    workspace,
 786                                    cx,
 787                                ).await
 788                            })
 789                            .detach_and_log_err(cx);
 790                    }
 791                }),
 792        ).item(
 793            ContextMenuEntry::new("View Docs")
 794                .icon(IconName::FileGeneric)
 795                .icon_color(Color::Muted)
 796                .handler(move |_, cx| {
 797                    cx.open_url(PRIVACY_DOCS);
 798                })
 799        );
 800
 801        if !self.editor_enabled.unwrap_or(true) {
 802            menu = menu.item(
 803                ContextMenuEntry::new("This file is excluded.")
 804                    .disabled(true)
 805                    .icon(IconName::ZedPredictDisabled)
 806                    .icon_size(IconSize::Small),
 807            );
 808        }
 809
 810        if let Some(editor_focus_handle) = self.editor_focus_handle.clone() {
 811            menu = menu
 812                .separator()
 813                .header("Actions")
 814                .entry(
 815                    "Predict Edit at Cursor",
 816                    Some(Box::new(ShowEditPrediction)),
 817                    {
 818                        let editor_focus_handle = editor_focus_handle.clone();
 819                        move |window, cx| {
 820                            editor_focus_handle.dispatch_action(&ShowEditPrediction, window, cx);
 821                        }
 822                    },
 823                )
 824                .context(editor_focus_handle)
 825                .when(
 826                    cx.has_flag::<PredictEditsRateCompletionsFeatureFlag>(),
 827                    |this| this.action("Rate Completions", RateCompletions.boxed_clone()),
 828                );
 829        }
 830
 831        menu
 832    }
 833
 834    fn build_copilot_context_menu(
 835        &self,
 836        window: &mut Window,
 837        cx: &mut Context<Self>,
 838    ) -> Entity<ContextMenu> {
 839        ContextMenu::build(window, cx, |menu, window, cx| {
 840            let menu = self.build_language_settings_menu(menu, window, cx);
 841            let menu =
 842                self.add_provider_switching_section(menu, EditPredictionProvider::Copilot, cx);
 843
 844            menu.separator()
 845                .link(
 846                    "Go to Copilot Settings",
 847                    OpenBrowser {
 848                        url: COPILOT_SETTINGS_URL.to_string(),
 849                    }
 850                    .boxed_clone(),
 851                )
 852                .action("Sign Out", copilot::SignOut.boxed_clone())
 853        })
 854    }
 855
 856    fn build_supermaven_context_menu(
 857        &self,
 858        window: &mut Window,
 859        cx: &mut Context<Self>,
 860    ) -> Entity<ContextMenu> {
 861        ContextMenu::build(window, cx, |menu, window, cx| {
 862            let menu = self.build_language_settings_menu(menu, window, cx);
 863            let menu =
 864                self.add_provider_switching_section(menu, EditPredictionProvider::Supermaven, cx);
 865
 866            menu.separator()
 867                .action("Sign Out", supermaven::SignOut.boxed_clone())
 868        })
 869    }
 870
 871    fn build_codestral_context_menu(
 872        &self,
 873        window: &mut Window,
 874        cx: &mut Context<Self>,
 875    ) -> Entity<ContextMenu> {
 876        ContextMenu::build(window, cx, |menu, window, cx| {
 877            let menu = self.build_language_settings_menu(menu, window, cx);
 878            let menu =
 879                self.add_provider_switching_section(menu, EditPredictionProvider::Codestral, cx);
 880
 881            menu.separator()
 882                .entry("Configure Codestral API Key", None, move |window, cx| {
 883                    window.dispatch_action(zed_actions::agent::OpenSettings.boxed_clone(), cx);
 884                })
 885        })
 886    }
 887
 888    fn build_zeta_context_menu(
 889        &self,
 890        window: &mut Window,
 891        cx: &mut Context<Self>,
 892    ) -> Entity<ContextMenu> {
 893        ContextMenu::build(window, cx, |mut menu, window, cx| {
 894            if let Some(usage) = self
 895                .edit_prediction_provider
 896                .as_ref()
 897                .and_then(|provider| provider.usage(cx))
 898            {
 899                menu = menu.header("Usage");
 900                menu = menu
 901                    .custom_entry(
 902                        move |_window, cx| {
 903                            let used_percentage = match usage.limit {
 904                                UsageLimit::Limited(limit) => {
 905                                    Some((usage.amount as f32 / limit as f32) * 100.)
 906                                }
 907                                UsageLimit::Unlimited => None,
 908                            };
 909
 910                            h_flex()
 911                                .flex_1()
 912                                .gap_1p5()
 913                                .children(
 914                                    used_percentage.map(|percent| {
 915                                        ProgressBar::new("usage", percent, 100., cx)
 916                                    }),
 917                                )
 918                                .child(
 919                                    Label::new(match usage.limit {
 920                                        UsageLimit::Limited(limit) => {
 921                                            format!("{} / {limit}", usage.amount)
 922                                        }
 923                                        UsageLimit::Unlimited => format!("{} / ∞", usage.amount),
 924                                    })
 925                                    .size(LabelSize::Small)
 926                                    .color(Color::Muted),
 927                                )
 928                                .into_any_element()
 929                        },
 930                        move |_, cx| cx.open_url(&zed_urls::account_url(cx)),
 931                    )
 932                    .when(usage.over_limit(), |menu| -> ContextMenu {
 933                        menu.entry("Subscribe to increase your limit", None, |_window, cx| {
 934                            cx.open_url(&zed_urls::account_url(cx))
 935                        })
 936                    })
 937                    .separator();
 938            } else if self.user_store.read(cx).account_too_young() {
 939                menu = menu
 940                    .custom_entry(
 941                        |_window, _cx| {
 942                            Label::new("Your GitHub account is less than 30 days old.")
 943                                .size(LabelSize::Small)
 944                                .color(Color::Warning)
 945                                .into_any_element()
 946                        },
 947                        |_window, cx| cx.open_url(&zed_urls::account_url(cx)),
 948                    )
 949                    .entry("Upgrade to Zed Pro or contact us.", None, |_window, cx| {
 950                        cx.open_url(&zed_urls::account_url(cx))
 951                    })
 952                    .separator();
 953            } else if self.user_store.read(cx).has_overdue_invoices() {
 954                menu = menu
 955                    .custom_entry(
 956                        |_window, _cx| {
 957                            Label::new("You have an outstanding invoice")
 958                                .size(LabelSize::Small)
 959                                .color(Color::Warning)
 960                                .into_any_element()
 961                        },
 962                        |_window, cx| {
 963                            cx.open_url(&zed_urls::account_url(cx))
 964                        },
 965                    )
 966                    .entry(
 967                        "Check your payment status or contact us at billing-support@zed.dev to continue using this feature.",
 968                        None,
 969                        |_window, cx| {
 970                            cx.open_url(&zed_urls::account_url(cx))
 971                        },
 972                    )
 973                    .separator();
 974            }
 975
 976            let menu = self.build_language_settings_menu(menu, window, cx);
 977            let menu = self.add_provider_switching_section(menu, EditPredictionProvider::Zed, cx);
 978
 979            menu
 980        })
 981    }
 982
 983    pub fn update_enabled(&mut self, editor: Entity<Editor>, cx: &mut Context<Self>) {
 984        let editor = editor.read(cx);
 985        let snapshot = editor.buffer().read(cx).snapshot(cx);
 986        let suggestion_anchor = editor.selections.newest_anchor().start;
 987        let language = snapshot.language_at(suggestion_anchor);
 988        let file = snapshot.file_at(suggestion_anchor).cloned();
 989        self.editor_enabled = {
 990            let file = file.as_ref();
 991            Some(
 992                file.map(|file| {
 993                    all_language_settings(Some(file), cx)
 994                        .edit_predictions_enabled_for_file(file, cx)
 995                })
 996                .unwrap_or(true),
 997            )
 998        };
 999        self.editor_show_predictions = editor.edit_predictions_enabled();
1000        self.edit_prediction_provider = editor.edit_prediction_provider();
1001        self.language = language.cloned();
1002        self.file = file;
1003        self.editor_focus_handle = Some(editor.focus_handle(cx));
1004
1005        cx.notify();
1006    }
1007}
1008
1009impl StatusItemView for EditPredictionButton {
1010    fn set_active_pane_item(
1011        &mut self,
1012        item: Option<&dyn ItemHandle>,
1013        _: &mut Window,
1014        cx: &mut Context<Self>,
1015    ) {
1016        if let Some(editor) = item.and_then(|item| item.act_as::<Editor>(cx)) {
1017            self.editor_subscription = Some((
1018                cx.observe(&editor, Self::update_enabled),
1019                editor.entity_id().as_u64() as usize,
1020            ));
1021            self.update_enabled(editor, cx);
1022        } else {
1023            self.language = None;
1024            self.editor_subscription = None;
1025            self.editor_enabled = None;
1026        }
1027        cx.notify();
1028    }
1029}
1030
1031impl SupermavenButtonStatus {
1032    fn to_icon(&self) -> IconName {
1033        match self {
1034            SupermavenButtonStatus::Ready => IconName::Supermaven,
1035            SupermavenButtonStatus::Errored(_) => IconName::SupermavenError,
1036            SupermavenButtonStatus::NeedsActivation(_) => IconName::SupermavenInit,
1037            SupermavenButtonStatus::Initializing => IconName::SupermavenInit,
1038        }
1039    }
1040
1041    fn to_tooltip(&self) -> String {
1042        match self {
1043            SupermavenButtonStatus::Ready => "Supermaven is ready".to_string(),
1044            SupermavenButtonStatus::Errored(error) => format!("Supermaven error: {}", error),
1045            SupermavenButtonStatus::NeedsActivation(_) => "Supermaven needs activation".to_string(),
1046            SupermavenButtonStatus::Initializing => "Supermaven initializing".to_string(),
1047        }
1048    }
1049
1050    fn has_menu(&self) -> bool {
1051        match self {
1052            SupermavenButtonStatus::Ready | SupermavenButtonStatus::NeedsActivation(_) => true,
1053            SupermavenButtonStatus::Errored(_) | SupermavenButtonStatus::Initializing => false,
1054        }
1055    }
1056}
1057
1058async fn open_disabled_globs_setting_in_editor(
1059    workspace: WeakEntity<Workspace>,
1060    cx: &mut AsyncWindowContext,
1061) -> Result<()> {
1062    let settings_editor = workspace
1063        .update_in(cx, |_, window, cx| {
1064            create_and_open_local_file(paths::settings_file(), window, cx, || {
1065                settings::initial_user_settings_content().as_ref().into()
1066            })
1067        })?
1068        .await?
1069        .downcast::<Editor>()
1070        .unwrap();
1071
1072    settings_editor
1073        .downgrade()
1074        .update_in(cx, |item, window, cx| {
1075            let text = item.buffer().read(cx).snapshot(cx).text();
1076
1077            let settings = cx.global::<SettingsStore>();
1078
1079            // Ensure that we always have "edit_predictions { "disabled_globs": [] }"
1080            let edits = settings.edits_for_update(&text, |file| {
1081                file.project
1082                    .all_languages
1083                    .edit_predictions
1084                    .get_or_insert_with(Default::default)
1085                    .disabled_globs
1086                    .get_or_insert_with(Vec::new);
1087            });
1088
1089            if !edits.is_empty() {
1090                item.edit(edits, cx);
1091            }
1092
1093            let text = item.buffer().read(cx).snapshot(cx).text();
1094
1095            static DISABLED_GLOBS_REGEX: LazyLock<Regex> = LazyLock::new(|| {
1096                Regex::new(r#""disabled_globs":\s*\[\s*(?P<content>(?:.|\n)*?)\s*\]"#).unwrap()
1097            });
1098            // Only capture [...]
1099            let range = DISABLED_GLOBS_REGEX.captures(&text).and_then(|captures| {
1100                captures
1101                    .name("content")
1102                    .map(|inner_match| inner_match.start()..inner_match.end())
1103            });
1104            if let Some(range) = range {
1105                item.change_selections(
1106                    SelectionEffects::scroll(Autoscroll::newest()),
1107                    window,
1108                    cx,
1109                    |selections| {
1110                        selections.select_ranges(vec![range]);
1111                    },
1112                );
1113            }
1114        })?;
1115
1116    anyhow::Ok(())
1117}
1118
1119fn set_completion_provider(fs: Arc<dyn Fs>, cx: &mut App, provider: EditPredictionProvider) {
1120    update_settings_file(fs, cx, move |settings, _| {
1121        settings
1122            .project
1123            .all_languages
1124            .features
1125            .get_or_insert_default()
1126            .edit_prediction_provider = Some(provider);
1127    });
1128}
1129
1130fn toggle_show_edit_predictions_for_language(
1131    language: Arc<Language>,
1132    fs: Arc<dyn Fs>,
1133    cx: &mut App,
1134) {
1135    let show_edit_predictions =
1136        all_language_settings(None, cx).show_edit_predictions(Some(&language), cx);
1137    update_settings_file(fs, cx, move |settings, _| {
1138        settings
1139            .project
1140            .all_languages
1141            .languages
1142            .0
1143            .entry(language.name().0)
1144            .or_default()
1145            .show_edit_predictions = Some(!show_edit_predictions);
1146    });
1147}
1148
1149fn hide_copilot(fs: Arc<dyn Fs>, cx: &mut App) {
1150    update_settings_file(fs, cx, move |settings, _| {
1151        settings
1152            .project
1153            .all_languages
1154            .features
1155            .get_or_insert(Default::default())
1156            .edit_prediction_provider = Some(EditPredictionProvider::None);
1157    });
1158}
1159
1160fn toggle_edit_prediction_mode(fs: Arc<dyn Fs>, mode: EditPredictionsMode, cx: &mut App) {
1161    let settings = AllLanguageSettings::get_global(cx);
1162    let current_mode = settings.edit_predictions_mode();
1163
1164    if current_mode != mode {
1165        update_settings_file(fs, cx, move |settings, _cx| {
1166            if let Some(edit_predictions) = settings.project.all_languages.edit_predictions.as_mut()
1167            {
1168                edit_predictions.mode = Some(mode);
1169            } else {
1170                settings.project.all_languages.edit_predictions =
1171                    Some(settings::EditPredictionSettingsContent {
1172                        mode: Some(mode),
1173                        ..Default::default()
1174                    });
1175            }
1176        });
1177    }
1178}