tenant.openagents/omega

1mod copilot_edit_prediction_delegate;
2pub mod request;
3
4use crate::request::{
5    DidFocus, DidFocusParams, FormattingOptions, InlineCompletionContext,
6    InlineCompletionTriggerKind, InlineCompletions, NextEditSuggestions,
7};
8use ::fs::Fs;
9use anyhow::{Context as _, Result, anyhow};
10use collections::{HashMap, HashSet};
11use command_palette_hooks::CommandPaletteFilter;
12use futures::future;
13use futures::{Future, FutureExt, TryFutureExt, channel::oneshot, future::Shared, select_biased};
14use gpui::{
15    App, AppContext as _, AsyncApp, Context, Entity, EntityId, EventEmitter, Global, Subscription,
16    Task, WeakEntity, actions,
17};
18use language::language_settings::{AllLanguageSettings, CopilotSettings};
19use language::{
20    Anchor, Bias, Buffer, BufferSnapshot, Language, PointUtf16, ToPointUtf16,
21    language_settings::{EditPredictionProvider, all_language_settings},
22    point_from_lsp, point_to_lsp,
23};
24use lsp::{LanguageServer, LanguageServerBinary, LanguageServerId, LanguageServerName};
25use node_runtime::{NodeRuntime, VersionStrategy};
26use parking_lot::Mutex;
27use project::project_settings::ProjectSettings;
28use project::{DisableAiSettings, Project};
29use request::DidChangeStatus;
30use serde_json::json;
31use settings::{Settings, SettingsStore};
32use std::{
33    any::TypeId,
34    collections::hash_map::Entry,
35    env,
36    ffi::OsString,
37    mem,
38    ops::Range,
39    path::{Path, PathBuf},
40    sync::Arc,
41};
42use sum_tree::Dimensions;
43use util::{ResultExt, fs::remove_matching};
44use workspace::AppState;
45
46pub use crate::copilot_edit_prediction_delegate::CopilotEditPredictionDelegate;
47
48actions!(
49    copilot,
50    [
51        /// Requests a code completion suggestion from Copilot.
52        Suggest,
53        /// Cycles to the next Copilot suggestion.
54        NextSuggestion,
55        /// Cycles to the previous Copilot suggestion.
56        PreviousSuggestion,
57        /// Reinstalls the Copilot language server.
58        Reinstall,
59        /// Signs in to GitHub Copilot.
60        SignIn,
61        /// Signs out of GitHub Copilot.
62        SignOut
63    ]
64);
65
66enum CopilotServer {
67    Disabled,
68    Starting { task: Shared<Task<()>> },
69    Error(Arc<str>),
70    Running(RunningCopilotServer),
71}
72
73impl CopilotServer {
74    fn as_authenticated(&mut self) -> Result<&mut RunningCopilotServer> {
75        let server = self.as_running()?;
76        anyhow::ensure!(
77            matches!(server.sign_in_status, SignInStatus::Authorized),
78            "must sign in before using copilot"
79        );
80        Ok(server)
81    }
82
83    fn as_running(&mut self) -> Result<&mut RunningCopilotServer> {
84        match self {
85            CopilotServer::Starting { .. } => anyhow::bail!("copilot is still starting"),
86            CopilotServer::Disabled => anyhow::bail!("copilot is disabled"),
87            CopilotServer::Error(error) => {
88                anyhow::bail!("copilot was not started because of an error: {error}")
89            }
90            CopilotServer::Running(server) => Ok(server),
91        }
92    }
93}
94
95struct RunningCopilotServer {
96    lsp: Arc<LanguageServer>,
97    sign_in_status: SignInStatus,
98    registered_buffers: HashMap<EntityId, RegisteredBuffer>,
99}
100
101#[derive(Clone, Debug)]
102enum SignInStatus {
103    Authorized,
104    Unauthorized,
105    SigningIn {
106        prompt: Option<request::PromptUserDeviceFlow>,
107        task: Shared<Task<Result<(), Arc<anyhow::Error>>>>,
108    },
109    SignedOut {
110        awaiting_signing_in: bool,
111    },
112}
113
114#[derive(Debug, Clone)]
115pub enum Status {
116    Starting {
117        task: Shared<Task<()>>,
118    },
119    Error(Arc<str>),
120    Disabled,
121    SignedOut {
122        awaiting_signing_in: bool,
123    },
124    SigningIn {
125        prompt: Option<request::PromptUserDeviceFlow>,
126    },
127    Unauthorized,
128    Authorized,
129}
130
131impl Status {
132    pub fn is_authorized(&self) -> bool {
133        matches!(self, Status::Authorized)
134    }
135
136    pub fn is_configured(&self) -> bool {
137        matches!(
138            self,
139            Status::Starting { .. }
140                | Status::Error(_)
141                | Status::SigningIn { .. }
142                | Status::Authorized
143        )
144    }
145}
146
147struct RegisteredBuffer {
148    uri: lsp::Uri,
149    language_id: String,
150    snapshot: BufferSnapshot,
151    snapshot_version: i32,
152    _subscriptions: [gpui::Subscription; 2],
153    pending_buffer_change: Task<Option<()>>,
154}
155
156impl RegisteredBuffer {
157    fn report_changes(
158        &mut self,
159        buffer: &Entity<Buffer>,
160        cx: &mut Context<Copilot>,
161    ) -> oneshot::Receiver<(i32, BufferSnapshot)> {
162        let (done_tx, done_rx) = oneshot::channel();
163
164        if buffer.read(cx).version() == self.snapshot.version {
165            let _ = done_tx.send((self.snapshot_version, self.snapshot.clone()));
166        } else {
167            let buffer = buffer.downgrade();
168            let id = buffer.entity_id();
169            let prev_pending_change =
170                mem::replace(&mut self.pending_buffer_change, Task::ready(None));
171            self.pending_buffer_change = cx.spawn(async move |copilot, cx| {
172                prev_pending_change.await;
173
174                let old_version = copilot
175                    .update(cx, |copilot, _| {
176                        let server = copilot.server.as_authenticated().log_err()?;
177                        let buffer = server.registered_buffers.get_mut(&id)?;
178                        Some(buffer.snapshot.version.clone())
179                    })
180                    .ok()??;
181                let new_snapshot = buffer.read_with(cx, |buffer, _| buffer.snapshot()).ok()?;
182
183                let content_changes = cx
184                    .background_spawn({
185                        let new_snapshot = new_snapshot.clone();
186                        async move {
187                            new_snapshot
188                                .edits_since::<Dimensions<PointUtf16, usize>>(&old_version)
189                                .map(|edit| {
190                                    let edit_start = edit.new.start.0;
191                                    let edit_end = edit_start + (edit.old.end.0 - edit.old.start.0);
192                                    let new_text = new_snapshot
193                                        .text_for_range(edit.new.start.1..edit.new.end.1)
194                                        .collect();
195                                    lsp::TextDocumentContentChangeEvent {
196                                        range: Some(lsp::Range::new(
197                                            point_to_lsp(edit_start),
198                                            point_to_lsp(edit_end),
199                                        )),
200                                        range_length: None,
201                                        text: new_text,
202                                    }
203                                })
204                                .collect::<Vec<_>>()
205                        }
206                    })
207                    .await;
208
209                copilot
210                    .update(cx, |copilot, _| {
211                        let server = copilot.server.as_authenticated().log_err()?;
212                        let buffer = server.registered_buffers.get_mut(&id)?;
213                        if !content_changes.is_empty() {
214                            buffer.snapshot_version += 1;
215                            buffer.snapshot = new_snapshot;
216                            server
217                                .lsp
218                                .notify::<lsp::notification::DidChangeTextDocument>(
219                                    lsp::DidChangeTextDocumentParams {
220                                        text_document: lsp::VersionedTextDocumentIdentifier::new(
221                                            buffer.uri.clone(),
222                                            buffer.snapshot_version,
223                                        ),
224                                        content_changes,
225                                    },
226                                )
227                                .ok();
228                        }
229                        let _ = done_tx.send((buffer.snapshot_version, buffer.snapshot.clone()));
230                        Some(())
231                    })
232                    .ok()?;
233
234                Some(())
235            });
236        }
237
238        done_rx
239    }
240}
241
242#[derive(Debug)]
243pub struct Completion {
244    pub uuid: String,
245    pub range: Range<Anchor>,
246    pub text: String,
247}
248
249pub struct Copilot {
250    fs: Arc<dyn Fs>,
251    node_runtime: NodeRuntime,
252    server: CopilotServer,
253    buffers: HashSet<WeakEntity<Buffer>>,
254    server_id: LanguageServerId,
255    _subscriptions: Vec<Subscription>,
256}
257
258pub enum Event {
259    CopilotAuthSignedIn,
260    CopilotAuthSignedOut,
261}
262
263impl EventEmitter<Event> for Copilot {}
264
265#[derive(Clone)]
266pub struct GlobalCopilotAuth(pub Entity<Copilot>);
267
268impl GlobalCopilotAuth {
269    pub fn set_global(
270        server_id: LanguageServerId,
271        fs: Arc<dyn Fs>,
272        node_runtime: NodeRuntime,
273        cx: &mut App,
274    ) -> GlobalCopilotAuth {
275        let auth =
276            GlobalCopilotAuth(cx.new(|cx| Copilot::new(None, server_id, fs, node_runtime, cx)));
277        cx.set_global(auth.clone());
278        auth
279    }
280    pub fn try_global(cx: &mut App) -> Option<&GlobalCopilotAuth> {
281        cx.try_global()
282    }
283
284    pub fn try_get_or_init(app_state: Arc<AppState>, cx: &mut App) -> Option<GlobalCopilotAuth> {
285        let ai_enabled = !DisableAiSettings::get(None, cx).disable_ai;
286
287        if let Some(copilot) = cx.try_global::<Self>().cloned() {
288            if ai_enabled {
289                Some(copilot)
290            } else {
291                cx.remove_global::<Self>();
292                None
293            }
294        } else if ai_enabled {
295            Some(Self::set_global(
296                app_state.languages.next_language_server_id(),
297                app_state.fs.clone(),
298                app_state.node_runtime.clone(),
299                cx,
300            ))
301        } else {
302            None
303        }
304    }
305}
306impl Global for GlobalCopilotAuth {}
307
308#[derive(Clone, Copy, Debug, PartialEq, Eq)]
309pub(crate) enum CompletionSource {
310    NextEditSuggestion,
311    InlineCompletion,
312}
313
314/// Copilot's NextEditSuggestion response, with coordinates converted to Anchors.
315#[derive(Clone)]
316pub(crate) struct CopilotEditPrediction {
317    pub(crate) buffer: Entity<Buffer>,
318    pub(crate) range: Range<Anchor>,
319    pub(crate) text: String,
320    pub(crate) command: Option<lsp::Command>,
321    pub(crate) snapshot: BufferSnapshot,
322    pub(crate) source: CompletionSource,
323}
324
325impl Copilot {
326    pub fn new(
327        project: Option<Entity<Project>>,
328        new_server_id: LanguageServerId,
329        fs: Arc<dyn Fs>,
330        node_runtime: NodeRuntime,
331        cx: &mut Context<Self>,
332    ) -> Self {
333        let send_focus_notification = project.map(|project| {
334            cx.subscribe(&project, |this, project, e: &project::Event, cx| {
335                if let project::Event::ActiveEntryChanged(new_entry) = e
336                    && let Ok(running) = this.server.as_authenticated()
337                {
338                    let uri = new_entry
339                        .and_then(|id| project.read(cx).path_for_entry(id, cx))
340                        .and_then(|entry| project.read(cx).absolute_path(&entry, cx))
341                        .and_then(|abs_path| lsp::Uri::from_file_path(abs_path).ok());
342
343                    _ = running.lsp.notify::<DidFocus>(DidFocusParams { uri });
344                }
345            })
346        });
347        let global_authentication_events =
348            cx.try_global::<GlobalCopilotAuth>().cloned().map(|auth| {
349                cx.subscribe(&auth.0, |_, _, _: &Event, cx| {
350                    let request_timeout = ProjectSettings::get_global(cx)
351                        .global_lsp_settings
352                        .get_request_timeout();
353                    cx.spawn(async move |this, cx| {
354                        let Some(server) = this
355                            .update(cx, |this, _| this.language_server().cloned())
356                            .ok()
357                            .flatten()
358                        else {
359                            return;
360                        };
361                        let status = server
362                            .request::<request::CheckStatus>(
363                                request::CheckStatusParams {
364                                    local_checks_only: false,
365                                },
366                                request_timeout,
367                            )
368                            .await
369                            .into_response()
370                            .ok();
371                        if let Some(status) = status {
372                            this.update(cx, |copilot, cx| {
373                                copilot.update_sign_in_status(status, cx);
374                            })
375                            .ok();
376                        }
377                    })
378                    .detach()
379                })
380            });
381        let _subscriptions = std::iter::once(cx.on_app_quit(Self::shutdown_language_server))
382            .chain(send_focus_notification)
383            .chain(global_authentication_events)
384            .collect();
385        let mut this = Self {
386            server_id: new_server_id,
387            fs,
388            node_runtime,
389            server: CopilotServer::Disabled,
390            buffers: Default::default(),
391            _subscriptions,
392        };
393        this.start_copilot(true, false, cx);
394        cx.observe_global::<SettingsStore>(move |this, cx| {
395            let ai_disabled = DisableAiSettings::get_global(cx).disable_ai;
396
397            if ai_disabled {
398                // Stop the server if AI is disabled
399                if !matches!(this.server, CopilotServer::Disabled) {
400                    let shutdown = match mem::replace(&mut this.server, CopilotServer::Disabled) {
401                        CopilotServer::Running(server) => {
402                            let shutdown_future = server.lsp.shutdown();
403                            Some(cx.background_spawn(async move {
404                                if let Some(fut) = shutdown_future {
405                                    fut.await;
406                                }
407                            }))
408                        }
409                        _ => None,
410                    };
411                    if let Some(task) = shutdown {
412                        task.detach();
413                    }
414                    cx.notify();
415                }
416            } else {
417                // Only start if AI is enabled
418                this.start_copilot(true, false, cx);
419                if let Ok(server) = this.server.as_running() {
420                    notify_did_change_config_to_server(&server.lsp, cx)
421                        .context("copilot setting change: did change configuration")
422                        .log_err();
423                }
424            }
425            this.update_action_visibilities(cx);
426        })
427        .detach();
428        cx.observe_self(|copilot, cx| {
429            copilot.update_action_visibilities(cx);
430        })
431        .detach();
432        this
433    }
434
435    fn shutdown_language_server(
436        &mut self,
437        _cx: &mut Context<Self>,
438    ) -> impl Future<Output = ()> + use<> {
439        let shutdown = match mem::replace(&mut self.server, CopilotServer::Disabled) {
440            CopilotServer::Running(server) => Some(Box::pin(async move { server.lsp.shutdown() })),
441            _ => None,
442        };
443
444        async move {
445            if let Some(shutdown) = shutdown {
446                shutdown.await;
447            }
448        }
449    }
450
451    pub fn start_copilot(
452        &mut self,
453        check_edit_prediction_provider: bool,
454        awaiting_sign_in_after_start: bool,
455        cx: &mut Context<Self>,
456    ) {
457        if DisableAiSettings::get_global(cx).disable_ai {
458            return;
459        }
460        if !matches!(self.server, CopilotServer::Disabled) {
461            return;
462        }
463        let language_settings = all_language_settings(None, cx);
464        if check_edit_prediction_provider
465            && language_settings.edit_predictions.provider != EditPredictionProvider::Copilot
466        {
467            return;
468        }
469        let server_id = self.server_id;
470        let fs = self.fs.clone();
471        let node_runtime = self.node_runtime.clone();
472        let env = self.build_env(&language_settings.edit_predictions.copilot);
473        let start_task = cx
474            .spawn(async move |this, cx| {
475                Self::start_language_server(
476                    server_id,
477                    fs,
478                    node_runtime,
479                    env,
480                    this,
481                    awaiting_sign_in_after_start,
482                    cx,
483                )
484                .await
485            })
486            .shared();
487        self.server = CopilotServer::Starting { task: start_task };
488        cx.notify();
489    }
490
491    fn build_env(&self, copilot_settings: &CopilotSettings) -> Option<HashMap<String, String>> {
492        let proxy_url = copilot_settings.proxy.clone()?;
493        let no_verify = copilot_settings.proxy_no_verify;
494        let http_or_https_proxy = if proxy_url.starts_with("http:") {
495            Some("HTTP_PROXY")
496        } else if proxy_url.starts_with("https:") {
497            Some("HTTPS_PROXY")
498        } else {
499            log::error!(
500                "Unsupported protocol scheme for language server proxy (must be http or https)"
501            );
502            None
503        };
504
505        let mut env = HashMap::default();
506
507        if let Some(proxy_type) = http_or_https_proxy {
508            env.insert(proxy_type.to_string(), proxy_url);
509            if let Some(true) = no_verify {
510                env.insert("NODE_TLS_REJECT_UNAUTHORIZED".to_string(), "0".to_string());
511            };
512        }
513
514        for env_var in [
515            copilot_chat::COPILOT_OAUTH_ENV_VAR,
516            copilot_chat::GITHUB_COPILOT_OAUTH_ENV_VAR,
517        ] {
518            if let Ok(oauth_token) = env::var(env_var) {
519                env.insert(env_var.to_string(), oauth_token);
520                break;
521            }
522        }
523
524        if env.is_empty() { None } else { Some(env) }
525    }
526
527    #[cfg(any(test, feature = "test-support"))]
528    pub fn fake(cx: &mut gpui::TestAppContext) -> (Entity<Self>, lsp::FakeLanguageServer) {
529        use fs::FakeFs;
530        use gpui::Subscription;
531        use lsp::FakeLanguageServer;
532        use node_runtime::NodeRuntime;
533
534        let (server, fake_server) = FakeLanguageServer::new(
535            LanguageServerId(0),
536            LanguageServerBinary {
537                path: "path/to/copilot".into(),
538                arguments: vec![],
539                env: None,
540            },
541            "copilot".into(),
542            Default::default(),
543            &mut cx.to_async(),
544        );
545        let node_runtime = NodeRuntime::unavailable();
546        let send_focus_notification = Subscription::new(|| {});
547        let this = cx.new(|cx| Self {
548            server_id: LanguageServerId(0),
549            fs: FakeFs::new(cx.background_executor().clone()),
550            node_runtime,
551            server: CopilotServer::Running(RunningCopilotServer {
552                lsp: Arc::new(server),
553                sign_in_status: SignInStatus::Authorized,
554                registered_buffers: Default::default(),
555            }),
556            _subscriptions: vec![
557                send_focus_notification,
558                cx.on_app_quit(Self::shutdown_language_server),
559            ],
560            buffers: Default::default(),
561        });
562        (this, fake_server)
563    }
564
565    async fn start_language_server(
566        new_server_id: LanguageServerId,
567        fs: Arc<dyn Fs>,
568        node_runtime: NodeRuntime,
569        env: Option<HashMap<String, String>>,
570        this: WeakEntity<Self>,
571        awaiting_sign_in_after_start: bool,
572        cx: &mut AsyncApp,
573    ) {
574        let start_language_server = async {
575            let server_path = get_copilot_lsp(fs, node_runtime).await?;
576
577            let arguments: Vec<OsString> = vec!["--stdio".into()];
578            let binary = LanguageServerBinary {
579                path: server_path,
580                arguments,
581                env,
582            };
583
584            let root_path = if cfg!(target_os = "windows") {
585                Path::new("C:/")
586            } else {
587                Path::new("/")
588            };
589
590            let server_name = LanguageServerName("copilot".into());
591            let server = LanguageServer::new(
592                Arc::new(Mutex::new(None)),
593                new_server_id,
594                server_name,
595                binary,
596                root_path,
597                None,
598                Default::default(),
599                cx,
600            )?;
601
602            server
603                .on_notification::<DidChangeStatus, _>({
604                    let this = this.clone();
605                    move |params, cx| {
606                        if params.kind == request::StatusKind::Normal {
607                            let this = this.clone();
608                            cx.spawn(async move |cx| {
609                                let lsp = this
610                                    .read_with(cx, |copilot, _| {
611                                        if let CopilotServer::Running(server) = &copilot.server {
612                                            Some(server.lsp.clone())
613                                        } else {
614                                            None
615                                        }
616                                    })
617                                    .ok()
618                                    .flatten();
619                                let Some(lsp) = lsp else { return };
620                                let request_timeout = cx.update(|cx| {
621                                    ProjectSettings::get_global(cx)
622                                        .global_lsp_settings
623                                        .get_request_timeout()
624                                });
625                                let status = lsp
626                                    .request::<request::CheckStatus>(
627                                        request::CheckStatusParams {
628                                            local_checks_only: false,
629                                        },
630                                        request_timeout,
631                                    )
632                                    .await
633                                    .into_response()
634                                    .ok();
635                                if let Some(status) = status {
636                                    this.update(cx, |copilot, cx| {
637                                        copilot.update_sign_in_status(status, cx);
638                                    })
639                                    .ok();
640                                }
641                            })
642                            .detach();
643                        }
644                    }
645                })
646                .detach();
647
648            server
649                .on_request::<lsp::request::ShowDocument, _, _>(move |params, cx| {
650                    if params.external.unwrap_or(false) {
651                        let url = params.uri.to_string();
652                        cx.update(|cx| cx.open_url(&url));
653                    }
654                    async move { Ok(lsp::ShowDocumentResult { success: true }) }
655                })
656                .detach();
657
658            let configuration = lsp::DidChangeConfigurationParams {
659                settings: Default::default(),
660            };
661
662            let editor_info = request::SetEditorInfoParams {
663                editor_info: request::EditorInfo {
664                    name: "zed".into(),
665                    version: env!("CARGO_PKG_VERSION").into(),
666                },
667                editor_plugin_info: request::EditorPluginInfo {
668                    name: "zed-copilot".into(),
669                    version: "0.0.1".into(),
670                },
671            };
672            let editor_info_json = serde_json::to_value(&editor_info)?;
673
674            let request_timeout = cx.update(|app| {
675                ProjectSettings::get_global(app)
676                    .global_lsp_settings
677                    .get_request_timeout()
678            });
679
680            let server = cx
681                .update(|cx| {
682                    let mut params = server.default_initialize_params(false, false, cx);
683                    params.initialization_options = Some(editor_info_json);
684                    params
685                        .capabilities
686                        .window
687                        .get_or_insert_with(Default::default)
688                        .show_document =
689                        Some(lsp::ShowDocumentClientCapabilities { support: true });
690                    server.initialize(params, configuration.into(), request_timeout, cx)
691                })
692                .await?;
693
694            this.update(cx, |_, cx| notify_did_change_config_to_server(&server, cx))?
695                .context("copilot: did change configuration")?;
696
697            let status = server
698                .request::<request::CheckStatus>(
699                    request::CheckStatusParams {
700                        local_checks_only: false,
701                    },
702                    request_timeout,
703                )
704                .await
705                .into_response()
706                .context("copilot: check status")?;
707
708            anyhow::Ok((server, status))
709        };
710
711        let server = start_language_server.await;
712        this.update(cx, |this, cx| {
713            cx.notify();
714
715            if env::var("ZED_FORCE_COPILOT_ERROR").is_ok() {
716                this.server = CopilotServer::Error(
717                    "Forced error for testing (ZED_FORCE_COPILOT_ERROR)".into(),
718                );
719                return;
720            }
721
722            match server {
723                Ok((server, status)) => {
724                    this.server = CopilotServer::Running(RunningCopilotServer {
725                        lsp: server,
726                        sign_in_status: SignInStatus::SignedOut {
727                            awaiting_signing_in: awaiting_sign_in_after_start,
728                        },
729                        registered_buffers: Default::default(),
730                    });
731                    this.update_sign_in_status(status, cx);
732                }
733                Err(error) => {
734                    this.server = CopilotServer::Error(error.to_string().into());
735                    cx.notify()
736                }
737            }
738        })
739        .ok();
740    }
741
742    pub fn is_authenticated(&self) -> bool {
743        return matches!(
744            self.server,
745            CopilotServer::Running(RunningCopilotServer {
746                sign_in_status: SignInStatus::Authorized,
747                ..
748            })
749        );
750    }
751
752    pub fn sign_in(&mut self, cx: &mut Context<Self>) -> Task<Result<()>> {
753        if let CopilotServer::Running(server) = &mut self.server {
754            let task = match &server.sign_in_status {
755                SignInStatus::Authorized => Task::ready(Ok(())).shared(),
756                SignInStatus::SigningIn { task, .. } => {
757                    cx.notify();
758                    task.clone()
759                }
760                SignInStatus::SignedOut { .. } | SignInStatus::Unauthorized => {
761                    let lsp = server.lsp.clone();
762
763                    let request_timeout = ProjectSettings::get_global(cx)
764                        .global_lsp_settings
765                        .get_request_timeout();
766
767                    let task = cx
768                        .spawn(async move |this, cx| {
769                            let sign_in = async {
770                                let flow = lsp
771                                    .request::<request::SignIn>(
772                                        request::SignInParams {},
773                                        request_timeout,
774                                    )
775                                    .await
776                                    .into_response()
777                                    .context("copilot sign-in")?;
778
779                                this.update(cx, |this, cx| {
780                                    if let CopilotServer::Running(RunningCopilotServer {
781                                        sign_in_status: status,
782                                        ..
783                                    }) = &mut this.server
784                                        && let SignInStatus::SigningIn {
785                                            prompt: prompt_flow,
786                                            ..
787                                        } = status
788                                    {
789                                        *prompt_flow = Some(flow.clone());
790                                        cx.notify();
791                                    }
792                                })?;
793
794                                anyhow::Ok(())
795                            };
796
797                            let sign_in = sign_in.await;
798                            this.update(cx, |this, cx| match sign_in {
799                                Ok(()) => Ok(()),
800                                Err(error) => {
801                                    this.update_sign_in_status(
802                                        request::SignInStatus::NotSignedIn,
803                                        cx,
804                                    );
805                                    Err(Arc::new(error))
806                                }
807                            })?
808                        })
809                        .shared();
810                    server.sign_in_status = SignInStatus::SigningIn {
811                        prompt: None,
812                        task: task.clone(),
813                    };
814                    cx.notify();
815                    task
816                }
817            };
818
819            cx.background_spawn(task.map_err(|err| anyhow!("{err:?}")))
820        } else {
821            // If we're downloading, wait until download is finished
822            // If we're in a stuck state, display to the user
823            Task::ready(Err(anyhow!("copilot hasn't started yet")))
824        }
825    }
826
827    pub fn sign_out(&mut self, cx: &mut Context<Self>) -> Task<Result<()>> {
828        self.update_sign_in_status(request::SignInStatus::NotSignedIn, cx);
829        match &self.server {
830            CopilotServer::Running(RunningCopilotServer { lsp: server, .. }) => {
831                let request_timeout = ProjectSettings::get_global(cx)
832                    .global_lsp_settings
833                    .get_request_timeout();
834
835                let server = server.clone();
836                cx.background_spawn(async move {
837                    server
838                        .request::<request::SignOut>(request::SignOutParams {}, request_timeout)
839                        .await
840                        .into_response()
841                        .context("copilot: sign in confirm")?;
842                    anyhow::Ok(())
843                })
844            }
845            CopilotServer::Disabled => cx.background_spawn(async {
846                clear_copilot_config_dir().await;
847                anyhow::Ok(())
848            }),
849            _ => Task::ready(Err(anyhow!("copilot hasn't started yet"))),
850        }
851    }
852
853    pub fn reinstall(&mut self, cx: &mut Context<Self>) -> Shared<Task<()>> {
854        let language_settings = all_language_settings(None, cx);
855        let env = self.build_env(&language_settings.edit_predictions.copilot);
856        let start_task = cx
857            .spawn({
858                let fs = self.fs.clone();
859                let node_runtime = self.node_runtime.clone();
860                let server_id = self.server_id;
861                async move |this, cx| {
862                    clear_copilot_dir().await;
863                    Self::start_language_server(server_id, fs, node_runtime, env, this, false, cx)
864                        .await
865                }
866            })
867            .shared();
868
869        self.server = CopilotServer::Starting {
870            task: start_task.clone(),
871        };
872
873        cx.notify();
874
875        start_task
876    }
877
878    pub fn language_server(&self) -> Option<&Arc<LanguageServer>> {
879        if let CopilotServer::Running(server) = &self.server {
880            Some(&server.lsp)
881        } else {
882            None
883        }
884    }
885
886    pub fn register_buffer(&mut self, buffer: &Entity<Buffer>, cx: &mut Context<Self>) {
887        let weak_buffer = buffer.downgrade();
888        self.buffers.insert(weak_buffer.clone());
889
890        if let CopilotServer::Running(RunningCopilotServer {
891            lsp: server,
892            sign_in_status: status,
893            registered_buffers,
894            ..
895        }) = &mut self.server
896        {
897            if !matches!(status, SignInStatus::Authorized) {
898                return;
899            }
900
901            let entry = registered_buffers.entry(buffer.entity_id());
902            if let Entry::Vacant(e) = entry {
903                let Ok(uri) = uri_for_buffer(buffer, cx) else {
904                    return;
905                };
906                let language_id = id_for_language(buffer.read(cx).language());
907                let snapshot = buffer.read(cx).snapshot();
908                server
909                    .notify::<lsp::notification::DidOpenTextDocument>(
910                        lsp::DidOpenTextDocumentParams {
911                            text_document: lsp::TextDocumentItem {
912                                uri: uri.clone(),
913                                language_id: language_id.clone(),
914                                version: 0,
915                                text: snapshot.text(),
916                            },
917                        },
918                    )
919                    .ok();
920
921                e.insert(RegisteredBuffer {
922                    uri,
923                    language_id,
924                    snapshot,
925                    snapshot_version: 0,
926                    pending_buffer_change: Task::ready(Some(())),
927                    _subscriptions: [
928                        cx.subscribe(buffer, |this, buffer, event, cx| {
929                            this.handle_buffer_event(buffer, event, cx).log_err();
930                        }),
931                        cx.observe_release(buffer, move |this, _buffer, _cx| {
932                            this.buffers.remove(&weak_buffer);
933                            this.unregister_buffer(&weak_buffer);
934                        }),
935                    ],
936                });
937            }
938        }
939    }
940
941    fn handle_buffer_event(
942        &mut self,
943        buffer: Entity<Buffer>,
944        event: &language::BufferEvent,
945        cx: &mut Context<Self>,
946    ) -> Result<()> {
947        if let Ok(server) = self.server.as_running()
948            && let Some(registered_buffer) = server.registered_buffers.get_mut(&buffer.entity_id())
949        {
950            match event {
951                language::BufferEvent::Edited { .. } => {
952                    drop(registered_buffer.report_changes(&buffer, cx));
953                }
954                language::BufferEvent::Saved => {
955                    server
956                        .lsp
957                        .notify::<lsp::notification::DidSaveTextDocument>(
958                            lsp::DidSaveTextDocumentParams {
959                                text_document: lsp::TextDocumentIdentifier::new(
960                                    registered_buffer.uri.clone(),
961                                ),
962                                text: None,
963                            },
964                        )
965                        .ok();
966                }
967                language::BufferEvent::FileHandleChanged
968                | language::BufferEvent::LanguageChanged(_) => {
969                    let new_language_id = id_for_language(buffer.read(cx).language());
970                    let Ok(new_uri) = uri_for_buffer(&buffer, cx) else {
971                        return Ok(());
972                    };
973                    if new_uri != registered_buffer.uri
974                        || new_language_id != registered_buffer.language_id
975                    {
976                        let old_uri = mem::replace(&mut registered_buffer.uri, new_uri);
977                        registered_buffer.language_id = new_language_id;
978                        server
979                            .lsp
980                            .notify::<lsp::notification::DidCloseTextDocument>(
981                                lsp::DidCloseTextDocumentParams {
982                                    text_document: lsp::TextDocumentIdentifier::new(old_uri),
983                                },
984                            )
985                            .ok();
986                        server
987                            .lsp
988                            .notify::<lsp::notification::DidOpenTextDocument>(
989                                lsp::DidOpenTextDocumentParams {
990                                    text_document: lsp::TextDocumentItem::new(
991                                        registered_buffer.uri.clone(),
992                                        registered_buffer.language_id.clone(),
993                                        registered_buffer.snapshot_version,
994                                        registered_buffer.snapshot.text(),
995                                    ),
996                                },
997                            )
998                            .ok();
999                    }
1000                }
1001                _ => {}
1002            }
1003        }
1004
1005        Ok(())
1006    }
1007
1008    fn unregister_buffer(&mut self, buffer: &WeakEntity<Buffer>) {
1009        if let Ok(server) = self.server.as_running()
1010            && let Some(buffer) = server.registered_buffers.remove(&buffer.entity_id())
1011        {
1012            server
1013                .lsp
1014                .notify::<lsp::notification::DidCloseTextDocument>(
1015                    lsp::DidCloseTextDocumentParams {
1016                        text_document: lsp::TextDocumentIdentifier::new(buffer.uri),
1017                    },
1018                )
1019                .ok();
1020        }
1021    }
1022
1023    pub(crate) fn completions(
1024        &mut self,
1025        buffer: &Entity<Buffer>,
1026        position: Anchor,
1027        cx: &mut Context<Self>,
1028    ) -> Task<Result<Vec<CopilotEditPrediction>>> {
1029        self.register_buffer(buffer, cx);
1030
1031        let server = match self.server.as_authenticated() {
1032            Ok(server) => server,
1033            Err(error) => return Task::ready(Err(error)),
1034        };
1035        let buffer_entity = buffer.clone();
1036        let lsp = server.lsp.clone();
1037        let Some(registered_buffer) = server.registered_buffers.get_mut(&buffer.entity_id()) else {
1038            return Task::ready(Err(anyhow::anyhow!("buffer not registered")));
1039        };
1040        let pending_snapshot = registered_buffer.report_changes(buffer, cx);
1041        let buffer = buffer.read(cx);
1042        let uri = registered_buffer.uri.clone();
1043        let position = position.to_point_utf16(buffer);
1044        let snapshot = buffer.snapshot();
1045        let settings = snapshot.settings_at(0, cx);
1046        let tab_size = settings.tab_size.get();
1047        let hard_tabs = settings.hard_tabs;
1048        drop(settings);
1049
1050        let request_timeout = ProjectSettings::get_global(cx)
1051            .global_lsp_settings
1052            .get_request_timeout();
1053
1054        let nes_enabled = AllLanguageSettings::get_global(cx)
1055            .edit_predictions
1056            .copilot
1057            .enable_next_edit_suggestions
1058            .unwrap_or(true);
1059
1060        cx.background_spawn(async move {
1061            let (version, snapshot) = pending_snapshot.await?;
1062            let lsp_position = point_to_lsp(position);
1063
1064            let nes_fut = if nes_enabled {
1065                lsp.request::<NextEditSuggestions>(
1066                    request::NextEditSuggestionsParams {
1067                        text_document: lsp::VersionedTextDocumentIdentifier {
1068                            uri: uri.clone(),
1069                            version,
1070                        },
1071                        position: lsp_position,
1072                    },
1073                    request_timeout,
1074                )
1075                .map(|resp| {
1076                    resp.into_response()
1077                        .ok()
1078                        .map(|result| {
1079                            result
1080                                .edits
1081                                .into_iter()
1082                                .map(|completion| {
1083                                    let start = snapshot.clip_point_utf16(
1084                                        point_from_lsp(completion.range.start),
1085                                        Bias::Left,
1086                                    );
1087                                    let end = snapshot.clip_point_utf16(
1088                                        point_from_lsp(completion.range.end),
1089                                        Bias::Left,
1090                                    );
1091                                    CopilotEditPrediction {
1092                                        buffer: buffer_entity.clone(),
1093                                        range: snapshot.anchor_before(start)
1094                                            ..snapshot.anchor_after(end),
1095                                        text: completion.text,
1096                                        command: completion.command,
1097                                        snapshot: snapshot.clone(),
1098                                        source: CompletionSource::NextEditSuggestion,
1099                                    }
1100                                })
1101                                .collect::<Vec<_>>()
1102                        })
1103                        .unwrap_or_default()
1104                })
1105                .left_future()
1106                .fuse()
1107            } else {
1108                future::ready(Vec::<CopilotEditPrediction>::new())
1109                    .right_future()
1110                    .fuse()
1111            };
1112
1113            let inline_fut = lsp
1114                .request::<InlineCompletions>(
1115                    request::InlineCompletionsParams {
1116                        text_document: lsp::VersionedTextDocumentIdentifier {
1117                            uri: uri.clone(),
1118                            version,
1119                        },
1120                        position: lsp_position,
1121                        context: InlineCompletionContext {
1122                            trigger_kind: InlineCompletionTriggerKind::Automatic,
1123                        },
1124                        formatting_options: Some(FormattingOptions {
1125                            tab_size,
1126                            insert_spaces: !hard_tabs,
1127                        }),
1128                    },
1129                    request_timeout,
1130                )
1131                .map(|resp| {
1132                    resp.into_response()
1133                        .ok()
1134                        .map(|result| {
1135                            result
1136                                .items
1137                                .into_iter()
1138                                .map(|item| {
1139                                    let start = snapshot.clip_point_utf16(
1140                                        point_from_lsp(item.range.start),
1141                                        Bias::Left,
1142                                    );
1143                                    let end = snapshot.clip_point_utf16(
1144                                        point_from_lsp(item.range.end),
1145                                        Bias::Left,
1146                                    );
1147                                    CopilotEditPrediction {
1148                                        buffer: buffer_entity.clone(),
1149                                        range: snapshot.anchor_before(start)
1150                                            ..snapshot.anchor_after(end),
1151                                        text: item.insert_text,
1152                                        command: item.command,
1153                                        snapshot: snapshot.clone(),
1154                                        source: CompletionSource::InlineCompletion,
1155                                    }
1156                                })
1157                                .collect::<Vec<_>>()
1158                        })
1159                        .unwrap_or_default()
1160                })
1161                .fuse();
1162
1163            futures::pin_mut!(nes_fut, inline_fut);
1164
1165            let mut nes_result: Option<Vec<CopilotEditPrediction>> = None;
1166            let mut inline_result: Option<Vec<CopilotEditPrediction>> = None;
1167
1168            loop {
1169                select_biased! {
1170                    nes = nes_fut => {
1171                        if !nes.is_empty() {
1172                            return Ok(nes);
1173                        }
1174                        nes_result = Some(nes);
1175                    }
1176                    inline = inline_fut => {
1177                        if !inline.is_empty() {
1178                            return Ok(inline);
1179                        }
1180                        inline_result = Some(inline);
1181                    }
1182                    complete => break,
1183                }
1184
1185                if let (Some(nes), Some(inline)) = (&nes_result, &inline_result) {
1186                    return if !nes.is_empty() {
1187                        Ok(nes.clone())
1188                    } else {
1189                        Ok(inline.clone())
1190                    };
1191                }
1192            }
1193
1194            Ok(nes_result.or(inline_result).unwrap_or_default())
1195        })
1196    }
1197
1198    pub(crate) fn accept_completion(
1199        &mut self,
1200        completion: &CopilotEditPrediction,
1201        cx: &mut Context<Self>,
1202    ) -> Task<Result<()>> {
1203        let server = match self.server.as_authenticated() {
1204            Ok(server) => server,
1205            Err(error) => return Task::ready(Err(error)),
1206        };
1207        if let Some(command) = &completion.command {
1208            let request_timeout = ProjectSettings::get_global(cx)
1209                .global_lsp_settings
1210                .get_request_timeout();
1211
1212            let request = server.lsp.request::<lsp::ExecuteCommand>(
1213                lsp::ExecuteCommandParams {
1214                    command: command.command.clone(),
1215                    arguments: command.arguments.clone().unwrap_or_default(),
1216                    ..Default::default()
1217                },
1218                request_timeout,
1219            );
1220            cx.background_spawn(async move {
1221                request
1222                    .await
1223                    .into_response()
1224                    .context("copilot: notify accepted")?;
1225                Ok(())
1226            })
1227        } else {
1228            Task::ready(Ok(()))
1229        }
1230    }
1231
1232    pub fn status(&self) -> Status {
1233        match &self.server {
1234            CopilotServer::Starting { task } => Status::Starting { task: task.clone() },
1235            CopilotServer::Disabled => Status::Disabled,
1236            CopilotServer::Error(error) => Status::Error(error.clone()),
1237            CopilotServer::Running(RunningCopilotServer { sign_in_status, .. }) => {
1238                match sign_in_status {
1239                    SignInStatus::Authorized => Status::Authorized,
1240                    SignInStatus::Unauthorized => Status::Unauthorized,
1241                    SignInStatus::SigningIn { prompt, .. } => Status::SigningIn {
1242                        prompt: prompt.clone(),
1243                    },
1244                    SignInStatus::SignedOut {
1245                        awaiting_signing_in,
1246                    } => Status::SignedOut {
1247                        awaiting_signing_in: *awaiting_signing_in,
1248                    },
1249                }
1250            }
1251        }
1252    }
1253
1254    pub fn update_sign_in_status(
1255        &mut self,
1256        lsp_status: request::SignInStatus,
1257        cx: &mut Context<Self>,
1258    ) {
1259        self.buffers.retain(|buffer| buffer.is_upgradable());
1260
1261        if let Ok(server) = self.server.as_running() {
1262            match lsp_status {
1263                request::SignInStatus::Ok { user: Some(_) }
1264                | request::SignInStatus::MaybeOk { .. }
1265                | request::SignInStatus::AlreadySignedIn { .. } => {
1266                    server.sign_in_status = SignInStatus::Authorized;
1267                    cx.emit(Event::CopilotAuthSignedIn);
1268                    notify_copilot_chat_auth_changed(cx);
1269                    for buffer in self.buffers.iter().cloned().collect::<Vec<_>>() {
1270                        if let Some(buffer) = buffer.upgrade() {
1271                            self.register_buffer(&buffer, cx);
1272                        }
1273                    }
1274                }
1275                request::SignInStatus::NotAuthorized { .. } => {
1276                    server.sign_in_status = SignInStatus::Unauthorized;
1277                    for buffer in self.buffers.iter().cloned().collect::<Vec<_>>() {
1278                        self.unregister_buffer(&buffer);
1279                    }
1280                }
1281                request::SignInStatus::Ok { user: None } | request::SignInStatus::NotSignedIn => {
1282                    if !matches!(server.sign_in_status, SignInStatus::SignedOut { .. }) {
1283                        server.sign_in_status = SignInStatus::SignedOut {
1284                            awaiting_signing_in: false,
1285                        };
1286                    }
1287                    cx.emit(Event::CopilotAuthSignedOut);
1288                    notify_copilot_chat_auth_changed(cx);
1289                    for buffer in self.buffers.iter().cloned().collect::<Vec<_>>() {
1290                        self.unregister_buffer(&buffer);
1291                    }
1292                }
1293            }
1294
1295            cx.notify();
1296        }
1297    }
1298
1299    fn update_action_visibilities(&self, cx: &mut App) {
1300        let signed_in_actions = [
1301            TypeId::of::<Suggest>(),
1302            TypeId::of::<NextSuggestion>(),
1303            TypeId::of::<PreviousSuggestion>(),
1304            TypeId::of::<Reinstall>(),
1305        ];
1306        let auth_actions = [TypeId::of::<SignOut>()];
1307        let no_auth_actions = [TypeId::of::<SignIn>()];
1308        let status = self.status();
1309
1310        let is_ai_disabled = DisableAiSettings::get_global(cx).disable_ai;
1311        let filter = CommandPaletteFilter::global_mut(cx);
1312
1313        if is_ai_disabled {
1314            filter.hide_action_types(&signed_in_actions);
1315            filter.hide_action_types(&auth_actions);
1316            filter.hide_action_types(&no_auth_actions);
1317        } else {
1318            match status {
1319                Status::Disabled => {
1320                    filter.hide_action_types(&signed_in_actions);
1321                    filter.hide_action_types(&auth_actions);
1322                    filter.hide_action_types(&no_auth_actions);
1323                }
1324                Status::Authorized => {
1325                    filter.hide_action_types(&no_auth_actions);
1326                    filter.show_action_types(signed_in_actions.iter().chain(&auth_actions));
1327                }
1328                _ => {
1329                    filter.hide_action_types(&signed_in_actions);
1330                    filter.hide_action_types(&auth_actions);
1331                    filter.show_action_types(&no_auth_actions);
1332                }
1333            }
1334        }
1335    }
1336}
1337
1338fn id_for_language(language: Option<&Arc<Language>>) -> String {
1339    language
1340        .map(|language| language.lsp_id())
1341        .unwrap_or_else(|| "plaintext".to_string())
1342}
1343
1344fn uri_for_buffer(buffer: &Entity<Buffer>, cx: &App) -> Result<lsp::Uri, ()> {
1345    if let Some(file) = buffer.read(cx).file().and_then(|file| file.as_local()) {
1346        lsp::Uri::from_file_path(file.abs_path(cx))
1347    } else {
1348        format!("buffer://{}", buffer.entity_id())
1349            .parse()
1350            .map_err(|_| ())
1351    }
1352}
1353
1354fn notify_did_change_config_to_server(
1355    server: &Arc<LanguageServer>,
1356    cx: &mut Context<Copilot>,
1357) -> std::result::Result<(), anyhow::Error> {
1358    let copilot_settings = all_language_settings(None, cx)
1359        .edit_predictions
1360        .copilot
1361        .clone();
1362
1363    if let Some(copilot_chat) = copilot_chat::CopilotChat::global(cx) {
1364        copilot_chat.update(cx, |chat, cx| {
1365            chat.set_configuration(
1366                copilot_chat::CopilotChatConfiguration {
1367                    enterprise_uri: copilot_settings.enterprise_uri.clone(),
1368                },
1369                cx,
1370            );
1371        });
1372    }
1373
1374    let settings = json!({
1375        "http": {
1376            "proxy": copilot_settings.proxy,
1377            "proxyStrictSSL": !copilot_settings.proxy_no_verify.unwrap_or(false)
1378        },
1379        "github-enterprise": {
1380            "uri": copilot_settings.enterprise_uri
1381        }
1382    });
1383
1384    server
1385        .notify::<lsp::notification::DidChangeConfiguration>(lsp::DidChangeConfigurationParams {
1386            settings,
1387        })
1388        .ok();
1389    Ok(())
1390}
1391
1392/// Notify Copilot Chat after the Copilot LSP reports an auth state change.
1393/// This replaces watching the SDK's token files, which is unreliable for
1394/// SQLite backed auth because writes may go through WAL files.
1395fn notify_copilot_chat_auth_changed(cx: &mut Context<Copilot>) {
1396    if let Some(copilot_chat) = copilot_chat::CopilotChat::global(cx) {
1397        copilot_chat.update(cx, |chat, cx| chat.reload_auth(cx));
1398    }
1399}
1400
1401async fn clear_copilot_dir() {
1402    remove_matching(paths::copilot_dir(), |_| true).await
1403}
1404
1405async fn clear_copilot_config_dir() {
1406    remove_matching(copilot_chat::copilot_chat_config_dir(), |_| true).await
1407}
1408
1409async fn get_copilot_lsp(fs: Arc<dyn Fs>, node_runtime: NodeRuntime) -> anyhow::Result<PathBuf> {
1410    const PACKAGE_NAME: &str = "@github/copilot-language-server";
1411    const SERVER_PATH: &str =
1412        "node_modules/@github/copilot-language-server/dist/language-server.js";
1413
1414    let latest_version = node_runtime
1415        .npm_package_latest_version(PACKAGE_NAME)
1416        .await?;
1417    let server_path = paths::copilot_dir().join(SERVER_PATH);
1418    let binary_path = copilot_lsp_native_binary_path()?;
1419
1420    fs.create_dir(paths::copilot_dir()).await?;
1421
1422    let should_install = !fs.is_file(&binary_path).await
1423        || node_runtime
1424            .should_install_npm_package(
1425                PACKAGE_NAME,
1426                &server_path,
1427                paths::copilot_dir(),
1428                VersionStrategy::Latest(&latest_version),
1429            )
1430            .await;
1431    if should_install {
1432        node_runtime
1433            .npm_install_latest_packages(paths::copilot_dir(), &[PACKAGE_NAME])
1434            .await?;
1435    }
1436
1437    if fs.is_file(&binary_path).await {
1438        return Ok(binary_path);
1439    }
1440
1441    anyhow::bail!("GitHub Copilot native language server binary was not installed")
1442}
1443
1444fn copilot_lsp_native_binary_path() -> anyhow::Result<PathBuf> {
1445    let platform = match env::consts::OS {
1446        "linux" => "linux",
1447        "macos" => "darwin",
1448        "windows" => "win32",
1449        platform => anyhow::bail!("unsupported Copilot language server platform: {platform}"),
1450    };
1451    let architecture = match env::consts::ARCH {
1452        "aarch64" => "arm64",
1453        "x86_64" => "x64",
1454        architecture => {
1455            anyhow::bail!("unsupported Copilot language server architecture: {architecture}")
1456        }
1457    };
1458
1459    let package_name = format!("copilot-language-server-{platform}-{architecture}");
1460
1461    let executable_name = if cfg!(target_os = "windows") {
1462        "copilot-language-server.exe"
1463    } else {
1464        "copilot-language-server"
1465    };
1466    Ok(paths::copilot_dir()
1467        .join("node_modules")
1468        .join("@github")
1469        .join(package_name)
1470        .join(executable_name))
1471}
1472
1473#[cfg(test)]
1474mod tests {
1475    use super::*;
1476    use fs::FakeFs;
1477    use gpui::TestAppContext;
1478    use language::language_settings::AllLanguageSettings;
1479    use node_runtime::NodeRuntime;
1480    use settings::{Settings, SettingsStore};
1481    use util::{
1482        path,
1483        paths::PathStyle,
1484        rel_path::{RelPath, rel_path},
1485    };
1486
1487    #[gpui::test]
1488    async fn test_copilot_does_not_start_when_ai_disabled(cx: &mut TestAppContext) {
1489        cx.update(|cx| {
1490            let store = SettingsStore::test(cx);
1491            cx.set_global(store);
1492            DisableAiSettings::register(cx);
1493            AllLanguageSettings::register(cx);
1494
1495            // Set disable_ai to true before creating Copilot
1496            DisableAiSettings::override_global(DisableAiSettings { disable_ai: true }, cx);
1497        });
1498
1499        let copilot = cx.new(|cx| Copilot {
1500            server_id: LanguageServerId(0),
1501            fs: FakeFs::new(cx.background_executor().clone()),
1502            node_runtime: NodeRuntime::unavailable(),
1503            server: CopilotServer::Disabled,
1504            buffers: Default::default(),
1505            _subscriptions: vec![],
1506        });
1507
1508        // Try to start copilot - it should remain disabled
1509        copilot.update(cx, |copilot, cx| {
1510            copilot.start_copilot(false, false, cx);
1511        });
1512
1513        // Verify the server is still disabled
1514        copilot.read_with(cx, |copilot, _| {
1515            assert!(
1516                matches!(copilot.server, CopilotServer::Disabled),
1517                "Copilot should not start when disable_ai is true"
1518            );
1519        });
1520    }
1521
1522    #[gpui::test]
1523    async fn test_copilot_stops_when_ai_becomes_disabled(cx: &mut TestAppContext) {
1524        cx.update(|cx| {
1525            let store = SettingsStore::test(cx);
1526            cx.set_global(store);
1527            DisableAiSettings::register(cx);
1528            AllLanguageSettings::register(cx);
1529
1530            // AI is initially enabled
1531            DisableAiSettings::override_global(DisableAiSettings { disable_ai: false }, cx);
1532        });
1533
1534        // Create a fake Copilot that's already running, with the settings observer
1535        let (copilot, _lsp) = Copilot::fake(cx);
1536
1537        // Add the settings observer that handles disable_ai changes
1538        copilot.update(cx, |_, cx| {
1539            cx.observe_global::<SettingsStore>(move |this, cx| {
1540                let ai_disabled = DisableAiSettings::get_global(cx).disable_ai;
1541
1542                if ai_disabled {
1543                    if !matches!(this.server, CopilotServer::Disabled) {
1544                        let shutdown = match mem::replace(&mut this.server, CopilotServer::Disabled)
1545                        {
1546                            CopilotServer::Running(server) => {
1547                                let shutdown_future = server.lsp.shutdown();
1548                                Some(cx.background_spawn(async move {
1549                                    if let Some(fut) = shutdown_future {
1550                                        fut.await;
1551                                    }
1552                                }))
1553                            }
1554                            _ => None,
1555                        };
1556                        if let Some(task) = shutdown {
1557                            task.detach();
1558                        }
1559                        cx.notify();
1560                    }
1561                }
1562            })
1563            .detach();
1564        });
1565
1566        // Verify copilot is running
1567        copilot.read_with(cx, |copilot, _| {
1568            assert!(
1569                matches!(copilot.server, CopilotServer::Running(_)),
1570                "Copilot should be running initially"
1571            );
1572        });
1573
1574        // Now disable AI
1575        cx.update(|cx| {
1576            DisableAiSettings::override_global(DisableAiSettings { disable_ai: true }, cx);
1577        });
1578
1579        // The settings observer should have stopped the server
1580        cx.run_until_parked();
1581
1582        copilot.read_with(cx, |copilot, _| {
1583            assert!(
1584                matches!(copilot.server, CopilotServer::Disabled),
1585                "Copilot should be disabled after disable_ai is set to true"
1586            );
1587        });
1588    }
1589
1590    #[gpui::test(iterations = 10)]
1591    async fn test_buffer_management(cx: &mut TestAppContext) {
1592        init_test(cx);
1593        let (copilot, mut lsp) = Copilot::fake(cx);
1594
1595        let buffer_1 = cx.new(|cx| Buffer::local("Hello", cx));
1596        let buffer_1_uri: lsp::Uri = format!("buffer://{}", buffer_1.entity_id().as_u64())
1597            .parse()
1598            .unwrap();
1599        copilot.update(cx, |copilot, cx| copilot.register_buffer(&buffer_1, cx));
1600        assert_eq!(
1601            lsp.receive_notification::<lsp::notification::DidOpenTextDocument>()
1602                .await,
1603            lsp::DidOpenTextDocumentParams {
1604                text_document: lsp::TextDocumentItem::new(
1605                    buffer_1_uri.clone(),
1606                    "plaintext".into(),
1607                    0,
1608                    "Hello".into()
1609                ),
1610            }
1611        );
1612
1613        let buffer_2 = cx.new(|cx| Buffer::local("Goodbye", cx));
1614        let buffer_2_uri: lsp::Uri = format!("buffer://{}", buffer_2.entity_id().as_u64())
1615            .parse()
1616            .unwrap();
1617        copilot.update(cx, |copilot, cx| copilot.register_buffer(&buffer_2, cx));
1618        assert_eq!(
1619            lsp.receive_notification::<lsp::notification::DidOpenTextDocument>()
1620                .await,
1621            lsp::DidOpenTextDocumentParams {
1622                text_document: lsp::TextDocumentItem::new(
1623                    buffer_2_uri.clone(),
1624                    "plaintext".into(),
1625                    0,
1626                    "Goodbye".into()
1627                ),
1628            }
1629        );
1630
1631        buffer_1.update(cx, |buffer, cx| buffer.edit([(5..5, " world")], None, cx));
1632        assert_eq!(
1633            lsp.receive_notification::<lsp::notification::DidChangeTextDocument>()
1634                .await,
1635            lsp::DidChangeTextDocumentParams {
1636                text_document: lsp::VersionedTextDocumentIdentifier::new(buffer_1_uri.clone(), 1),
1637                content_changes: vec![lsp::TextDocumentContentChangeEvent {
1638                    range: Some(lsp::Range::new(
1639                        lsp::Position::new(0, 5),
1640                        lsp::Position::new(0, 5)
1641                    )),
1642                    range_length: None,
1643                    text: " world".into(),
1644                }],
1645            }
1646        );
1647
1648        // Ensure updates to the file are reflected in the LSP.
1649        buffer_1.update(cx, |buffer, cx| {
1650            buffer.file_updated(
1651                Arc::new(File {
1652                    abs_path: path!("/root/child/buffer-1").into(),
1653                    path: rel_path("child/buffer-1").into(),
1654                }),
1655                cx,
1656            )
1657        });
1658        assert_eq!(
1659            lsp.receive_notification::<lsp::notification::DidCloseTextDocument>()
1660                .await,
1661            lsp::DidCloseTextDocumentParams {
1662                text_document: lsp::TextDocumentIdentifier::new(buffer_1_uri),
1663            }
1664        );
1665        let buffer_1_uri = lsp::Uri::from_file_path(path!("/root/child/buffer-1")).unwrap();
1666        assert_eq!(
1667            lsp.receive_notification::<lsp::notification::DidOpenTextDocument>()
1668                .await,
1669            lsp::DidOpenTextDocumentParams {
1670                text_document: lsp::TextDocumentItem::new(
1671                    buffer_1_uri.clone(),
1672                    "plaintext".into(),
1673                    1,
1674                    "Hello world".into()
1675                ),
1676            }
1677        );
1678
1679        // Ensure all previously-registered buffers are closed when signing out.
1680        lsp.set_request_handler::<request::SignOut, _, _>(|_, _| async {
1681            Ok(request::SignOutResult {})
1682        });
1683        copilot
1684            .update(cx, |copilot, cx| copilot.sign_out(cx))
1685            .await
1686            .unwrap();
1687        let mut received_close_notifications = vec![
1688            lsp.receive_notification::<lsp::notification::DidCloseTextDocument>()
1689                .await,
1690            lsp.receive_notification::<lsp::notification::DidCloseTextDocument>()
1691                .await,
1692        ];
1693        received_close_notifications
1694            .sort_by_key(|notification| notification.text_document.uri.clone());
1695        assert_eq!(
1696            received_close_notifications,
1697            vec![
1698                lsp::DidCloseTextDocumentParams {
1699                    text_document: lsp::TextDocumentIdentifier::new(buffer_2_uri.clone()),
1700                },
1701                lsp::DidCloseTextDocumentParams {
1702                    text_document: lsp::TextDocumentIdentifier::new(buffer_1_uri.clone()),
1703                },
1704            ],
1705        );
1706
1707        // Ensure all previously-registered buffers are re-opened when signing in.
1708        lsp.set_request_handler::<request::SignIn, _, _>(|_, _| async {
1709            Ok(request::PromptUserDeviceFlow {
1710                user_code: "test-code".into(),
1711                command: lsp::Command {
1712                    title: "Sign in".into(),
1713                    command: "github.copilot.finishDeviceFlow".into(),
1714                    arguments: None,
1715                },
1716            })
1717        });
1718        copilot
1719            .update(cx, |copilot, cx| copilot.sign_in(cx))
1720            .await
1721            .unwrap();
1722
1723        // Simulate auth completion by directly updating sign-in status
1724        copilot.update(cx, |copilot, cx| {
1725            copilot.update_sign_in_status(
1726                request::SignInStatus::Ok {
1727                    user: Some("user-1".into()),
1728                },
1729                cx,
1730            );
1731        });
1732
1733        let mut received_open_notifications = vec![
1734            lsp.receive_notification::<lsp::notification::DidOpenTextDocument>()
1735                .await,
1736            lsp.receive_notification::<lsp::notification::DidOpenTextDocument>()
1737                .await,
1738        ];
1739        received_open_notifications
1740            .sort_by_key(|notification| notification.text_document.uri.clone());
1741        assert_eq!(
1742            received_open_notifications,
1743            vec![
1744                lsp::DidOpenTextDocumentParams {
1745                    text_document: lsp::TextDocumentItem::new(
1746                        buffer_2_uri.clone(),
1747                        "plaintext".into(),
1748                        0,
1749                        "Goodbye".into()
1750                    ),
1751                },
1752                lsp::DidOpenTextDocumentParams {
1753                    text_document: lsp::TextDocumentItem::new(
1754                        buffer_1_uri.clone(),
1755                        "plaintext".into(),
1756                        0,
1757                        "Hello world".into()
1758                    ),
1759                }
1760            ]
1761        );
1762        // Dropping a buffer causes it to be closed on the LSP side as well.
1763        cx.update(|_| drop(buffer_2));
1764        assert_eq!(
1765            lsp.receive_notification::<lsp::notification::DidCloseTextDocument>()
1766                .await,
1767            lsp::DidCloseTextDocumentParams {
1768                text_document: lsp::TextDocumentIdentifier::new(buffer_2_uri),
1769            }
1770        );
1771    }
1772
1773    struct File {
1774        abs_path: PathBuf,
1775        path: Arc<RelPath>,
1776    }
1777
1778    impl language::File for File {
1779        fn as_local(&self) -> Option<&dyn language::LocalFile> {
1780            Some(self)
1781        }
1782
1783        fn disk_state(&self) -> language::DiskState {
1784            language::DiskState::Present {
1785                mtime: ::fs::MTime::from_seconds_and_nanos(100, 42),
1786                size: 0,
1787            }
1788        }
1789
1790        fn path(&self) -> &Arc<RelPath> {
1791            &self.path
1792        }
1793
1794        fn path_style(&self, _: &App) -> PathStyle {
1795            PathStyle::local()
1796        }
1797
1798        fn full_path(&self, _: &App) -> PathBuf {
1799            unimplemented!()
1800        }
1801
1802        fn file_name<'a>(&'a self, _: &'a App) -> &'a str {
1803            unimplemented!()
1804        }
1805
1806        fn to_proto(&self, _: &App) -> rpc::proto::File {
1807            unimplemented!()
1808        }
1809
1810        fn worktree_id(&self, _: &App) -> settings::WorktreeId {
1811            settings::WorktreeId::from_usize(0)
1812        }
1813
1814        fn is_private(&self) -> bool {
1815            false
1816        }
1817    }
1818
1819    impl language::LocalFile for File {
1820        fn abs_path(&self, _: &App) -> PathBuf {
1821            self.abs_path.clone()
1822        }
1823
1824        fn load(&self, _: &App) -> Task<Result<String>> {
1825            unimplemented!()
1826        }
1827
1828        fn load_bytes(&self, _cx: &App) -> Task<Result<Vec<u8>>> {
1829            unimplemented!()
1830        }
1831    }
1832
1833    #[gpui::test]
1834    async fn test_copilot_starts_when_ai_becomes_enabled(cx: &mut TestAppContext) {
1835        cx.update(|cx| {
1836            let store = SettingsStore::test(cx);
1837            cx.set_global(store);
1838            DisableAiSettings::register(cx);
1839            AllLanguageSettings::register(cx);
1840
1841            // AI is initially disabled
1842            DisableAiSettings::override_global(DisableAiSettings { disable_ai: true }, cx);
1843        });
1844
1845        let copilot = cx.new(|cx| Copilot {
1846            server_id: LanguageServerId(0),
1847            fs: FakeFs::new(cx.background_executor().clone()),
1848            node_runtime: NodeRuntime::unavailable(),
1849            server: CopilotServer::Disabled,
1850            buffers: Default::default(),
1851            _subscriptions: vec![],
1852        });
1853
1854        // Verify copilot is disabled initially
1855        copilot.read_with(cx, |copilot, _| {
1856            assert!(
1857                matches!(copilot.server, CopilotServer::Disabled),
1858                "Copilot should be disabled initially"
1859            );
1860        });
1861
1862        // Try to start - should fail because AI is disabled
1863        // Use check_edit_prediction_provider=false to skip provider check
1864        copilot.update(cx, |copilot, cx| {
1865            copilot.start_copilot(false, false, cx);
1866        });
1867
1868        copilot.read_with(cx, |copilot, _| {
1869            assert!(
1870                matches!(copilot.server, CopilotServer::Disabled),
1871                "Copilot should remain disabled when disable_ai is true"
1872            );
1873        });
1874
1875        // Now enable AI
1876        cx.update(|cx| {
1877            DisableAiSettings::override_global(DisableAiSettings { disable_ai: false }, cx);
1878        });
1879
1880        // Try to start again - should work now
1881        copilot.update(cx, |copilot, cx| {
1882            copilot.start_copilot(false, false, cx);
1883        });
1884
1885        copilot.read_with(cx, |copilot, _| {
1886            assert!(
1887                matches!(copilot.server, CopilotServer::Starting { .. }),
1888                "Copilot should be starting after disable_ai is set to false"
1889            );
1890        });
1891    }
1892
1893    fn init_test(cx: &mut TestAppContext) {
1894        zlog::init_test();
1895
1896        cx.update(|cx| {
1897            let settings_store = SettingsStore::test(cx);
1898            cx.set_global(settings_store);
1899        });
1900    }
1901}
1902