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