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