1pub mod request;
   2mod sign_in;
   3
   4use anyhow::{anyhow, Context, Result};
   5use async_compression::futures::bufread::GzipDecoder;
   6use async_tar::Archive;
   7use collections::{HashMap, HashSet};
   8use futures::{channel::oneshot, future::Shared, Future, FutureExt, TryFutureExt};
   9use gpui::{
  10    actions, AppContext, AsyncAppContext, Entity, ModelContext, ModelHandle, Task, WeakModelHandle,
  11};
  12use language::{
  13    language_settings::{all_language_settings, language_settings},
  14    point_from_lsp, point_to_lsp, Anchor, Bias, Buffer, BufferSnapshot, Language,
  15    LanguageServerName, PointUtf16, ToPointUtf16,
  16};
  17use lsp::{LanguageServer, LanguageServerBinary, LanguageServerId};
  18use node_runtime::NodeRuntime;
  19use parking_lot::Mutex;
  20use request::StatusNotification;
  21use settings::SettingsStore;
  22use smol::{fs, io::BufReader, stream::StreamExt};
  23use std::{
  24    ffi::OsString,
  25    mem,
  26    ops::Range,
  27    path::{Path, PathBuf},
  28    pin::Pin,
  29    sync::Arc,
  30};
  31use util::{
  32    fs::remove_matching, github::latest_github_release, http::HttpClient, paths, ResultExt,
  33};
  34
  35const COPILOT_AUTH_NAMESPACE: &'static str = "copilot_auth";
  36actions!(copilot_auth, [SignIn, SignOut]);
  37
  38const COPILOT_NAMESPACE: &'static str = "copilot";
  39actions!(
  40    copilot,
  41    [Suggest, NextSuggestion, PreviousSuggestion, Reinstall]
  42);
  43
  44pub fn init(
  45    new_server_id: LanguageServerId,
  46    http: Arc<dyn HttpClient>,
  47    node_runtime: Arc<dyn NodeRuntime>,
  48    cx: &mut AppContext,
  49) {
  50    let copilot = cx.add_model({
  51        let node_runtime = node_runtime.clone();
  52        move |cx| Copilot::start(new_server_id, http, node_runtime, cx)
  53    });
  54    cx.set_global(copilot.clone());
  55
  56    cx.observe(&copilot, |handle, cx| {
  57        let status = handle.read(cx).status();
  58        cx.update_default_global::<collections::CommandPaletteFilter, _, _>(move |filter, _cx| {
  59            match status {
  60                Status::Disabled => {
  61                    filter.hidden_namespaces.insert(COPILOT_NAMESPACE);
  62                    filter.hidden_namespaces.insert(COPILOT_AUTH_NAMESPACE);
  63                }
  64                Status::Authorized => {
  65                    filter.hidden_namespaces.remove(COPILOT_NAMESPACE);
  66                    filter.hidden_namespaces.remove(COPILOT_AUTH_NAMESPACE);
  67                }
  68                _ => {
  69                    filter.hidden_namespaces.insert(COPILOT_NAMESPACE);
  70                    filter.hidden_namespaces.remove(COPILOT_AUTH_NAMESPACE);
  71                }
  72            }
  73        });
  74    })
  75    .detach();
  76
  77    sign_in::init(cx);
  78    cx.add_global_action(|_: &SignIn, cx| {
  79        if let Some(copilot) = Copilot::global(cx) {
  80            copilot
  81                .update(cx, |copilot, cx| copilot.sign_in(cx))
  82                .detach_and_log_err(cx);
  83        }
  84    });
  85    cx.add_global_action(|_: &SignOut, cx| {
  86        if let Some(copilot) = Copilot::global(cx) {
  87            copilot
  88                .update(cx, |copilot, cx| copilot.sign_out(cx))
  89                .detach_and_log_err(cx);
  90        }
  91    });
  92
  93    cx.add_global_action(|_: &Reinstall, cx| {
  94        if let Some(copilot) = Copilot::global(cx) {
  95            copilot
  96                .update(cx, |copilot, cx| copilot.reinstall(cx))
  97                .detach();
  98        }
  99    });
 100}
 101
 102enum CopilotServer {
 103    Disabled,
 104    Starting { task: Shared<Task<()>> },
 105    Error(Arc<str>),
 106    Running(RunningCopilotServer),
 107}
 108
 109impl CopilotServer {
 110    fn as_authenticated(&mut self) -> Result<&mut RunningCopilotServer> {
 111        let server = self.as_running()?;
 112        if matches!(server.sign_in_status, SignInStatus::Authorized { .. }) {
 113            Ok(server)
 114        } else {
 115            Err(anyhow!("must sign in before using copilot"))
 116        }
 117    }
 118
 119    fn as_running(&mut self) -> Result<&mut RunningCopilotServer> {
 120        match self {
 121            CopilotServer::Starting { .. } => Err(anyhow!("copilot is still starting")),
 122            CopilotServer::Disabled => Err(anyhow!("copilot is disabled")),
 123            CopilotServer::Error(error) => Err(anyhow!(
 124                "copilot was not started because of an error: {}",
 125                error
 126            )),
 127            CopilotServer::Running(server) => Ok(server),
 128        }
 129    }
 130}
 131
 132struct RunningCopilotServer {
 133    name: LanguageServerName,
 134    lsp: Arc<LanguageServer>,
 135    sign_in_status: SignInStatus,
 136    registered_buffers: HashMap<usize, RegisteredBuffer>,
 137}
 138
 139#[derive(Clone, Debug)]
 140enum SignInStatus {
 141    Authorized,
 142    Unauthorized,
 143    SigningIn {
 144        prompt: Option<request::PromptUserDeviceFlow>,
 145        task: Shared<Task<Result<(), Arc<anyhow::Error>>>>,
 146    },
 147    SignedOut,
 148}
 149
 150#[derive(Debug, Clone)]
 151pub enum Status {
 152    Starting {
 153        task: Shared<Task<()>>,
 154    },
 155    Error(Arc<str>),
 156    Disabled,
 157    SignedOut,
 158    SigningIn {
 159        prompt: Option<request::PromptUserDeviceFlow>,
 160    },
 161    Unauthorized,
 162    Authorized,
 163}
 164
 165impl Status {
 166    pub fn is_authorized(&self) -> bool {
 167        matches!(self, Status::Authorized)
 168    }
 169}
 170
 171struct RegisteredBuffer {
 172    uri: lsp::Url,
 173    language_id: String,
 174    snapshot: BufferSnapshot,
 175    snapshot_version: i32,
 176    _subscriptions: [gpui::Subscription; 2],
 177    pending_buffer_change: Task<Option<()>>,
 178}
 179
 180impl RegisteredBuffer {
 181    fn report_changes(
 182        &mut self,
 183        buffer: &ModelHandle<Buffer>,
 184        cx: &mut ModelContext<Copilot>,
 185    ) -> oneshot::Receiver<(i32, BufferSnapshot)> {
 186        let (done_tx, done_rx) = oneshot::channel();
 187
 188        if buffer.read(cx).version() == self.snapshot.version {
 189            let _ = done_tx.send((self.snapshot_version, self.snapshot.clone()));
 190        } else {
 191            let buffer = buffer.downgrade();
 192            let id = buffer.id();
 193            let prev_pending_change =
 194                mem::replace(&mut self.pending_buffer_change, Task::ready(None));
 195            self.pending_buffer_change = cx.spawn_weak(|copilot, mut cx| async move {
 196                prev_pending_change.await;
 197
 198                let old_version = copilot.upgrade(&cx)?.update(&mut cx, |copilot, _| {
 199                    let server = copilot.server.as_authenticated().log_err()?;
 200                    let buffer = server.registered_buffers.get_mut(&id)?;
 201                    Some(buffer.snapshot.version.clone())
 202                })?;
 203                let new_snapshot = buffer
 204                    .upgrade(&cx)?
 205                    .read_with(&cx, |buffer, _| buffer.snapshot());
 206
 207                let content_changes = cx
 208                    .background()
 209                    .spawn({
 210                        let new_snapshot = new_snapshot.clone();
 211                        async move {
 212                            new_snapshot
 213                                .edits_since::<(PointUtf16, usize)>(&old_version)
 214                                .map(|edit| {
 215                                    let edit_start = edit.new.start.0;
 216                                    let edit_end = edit_start + (edit.old.end.0 - edit.old.start.0);
 217                                    let new_text = new_snapshot
 218                                        .text_for_range(edit.new.start.1..edit.new.end.1)
 219                                        .collect();
 220                                    lsp::TextDocumentContentChangeEvent {
 221                                        range: Some(lsp::Range::new(
 222                                            point_to_lsp(edit_start),
 223                                            point_to_lsp(edit_end),
 224                                        )),
 225                                        range_length: None,
 226                                        text: new_text,
 227                                    }
 228                                })
 229                                .collect::<Vec<_>>()
 230                        }
 231                    })
 232                    .await;
 233
 234                copilot.upgrade(&cx)?.update(&mut cx, |copilot, _| {
 235                    let server = copilot.server.as_authenticated().log_err()?;
 236                    let buffer = server.registered_buffers.get_mut(&id)?;
 237                    if !content_changes.is_empty() {
 238                        buffer.snapshot_version += 1;
 239                        buffer.snapshot = new_snapshot;
 240                        server
 241                            .lsp
 242                            .notify::<lsp::notification::DidChangeTextDocument>(
 243                                lsp::DidChangeTextDocumentParams {
 244                                    text_document: lsp::VersionedTextDocumentIdentifier::new(
 245                                        buffer.uri.clone(),
 246                                        buffer.snapshot_version,
 247                                    ),
 248                                    content_changes,
 249                                },
 250                            )
 251                            .log_err();
 252                    }
 253                    let _ = done_tx.send((buffer.snapshot_version, buffer.snapshot.clone()));
 254                    Some(())
 255                })?;
 256
 257                Some(())
 258            });
 259        }
 260
 261        done_rx
 262    }
 263}
 264
 265#[derive(Debug)]
 266pub struct Completion {
 267    pub uuid: String,
 268    pub range: Range<Anchor>,
 269    pub text: String,
 270}
 271
 272pub struct Copilot {
 273    http: Arc<dyn HttpClient>,
 274    node_runtime: Arc<dyn NodeRuntime>,
 275    server: CopilotServer,
 276    buffers: HashSet<WeakModelHandle<Buffer>>,
 277    server_id: LanguageServerId,
 278}
 279
 280pub enum Event {
 281    CopilotLanguageServerStarted,
 282}
 283
 284impl Entity for Copilot {
 285    type Event = Event;
 286
 287    fn app_will_quit(
 288        &mut self,
 289        _: &mut AppContext,
 290    ) -> Option<Pin<Box<dyn 'static + Future<Output = ()>>>> {
 291        match mem::replace(&mut self.server, CopilotServer::Disabled) {
 292            CopilotServer::Running(server) => Some(Box::pin(async move {
 293                if let Some(shutdown) = server.lsp.shutdown() {
 294                    shutdown.await;
 295                }
 296            })),
 297            _ => None,
 298        }
 299    }
 300}
 301
 302impl Copilot {
 303    pub fn global(cx: &AppContext) -> Option<ModelHandle<Self>> {
 304        if cx.has_global::<ModelHandle<Self>>() {
 305            Some(cx.global::<ModelHandle<Self>>().clone())
 306        } else {
 307            None
 308        }
 309    }
 310
 311    fn start(
 312        new_server_id: LanguageServerId,
 313        http: Arc<dyn HttpClient>,
 314        node_runtime: Arc<dyn NodeRuntime>,
 315        cx: &mut ModelContext<Self>,
 316    ) -> Self {
 317        let mut this = Self {
 318            server_id: new_server_id,
 319            http,
 320            node_runtime,
 321            server: CopilotServer::Disabled,
 322            buffers: Default::default(),
 323        };
 324        this.enable_or_disable_copilot(cx);
 325        cx.observe_global::<SettingsStore, _>(move |this, cx| this.enable_or_disable_copilot(cx))
 326            .detach();
 327        this
 328    }
 329
 330    fn enable_or_disable_copilot(&mut self, cx: &mut ModelContext<Copilot>) {
 331        let server_id = self.server_id;
 332        let http = self.http.clone();
 333        let node_runtime = self.node_runtime.clone();
 334        if all_language_settings(None, cx).copilot_enabled(None, None) {
 335            if matches!(self.server, CopilotServer::Disabled) {
 336                let start_task = cx
 337                    .spawn({
 338                        move |this, cx| {
 339                            Self::start_language_server(server_id, http, node_runtime, this, cx)
 340                        }
 341                    })
 342                    .shared();
 343                self.server = CopilotServer::Starting { task: start_task };
 344                cx.notify();
 345            }
 346        } else {
 347            self.server = CopilotServer::Disabled;
 348            cx.notify();
 349        }
 350    }
 351
 352    #[cfg(any(test, feature = "test-support"))]
 353    pub fn fake(cx: &mut gpui::TestAppContext) -> (ModelHandle<Self>, lsp::FakeLanguageServer) {
 354        use node_runtime::FakeNodeRuntime;
 355
 356        let (server, fake_server) =
 357            LanguageServer::fake("copilot".into(), Default::default(), cx.to_async());
 358        let http = util::http::FakeHttpClient::create(|_| async { unreachable!() });
 359        let node_runtime = FakeNodeRuntime::new();
 360        let this = cx.add_model(|_| Self {
 361            server_id: LanguageServerId(0),
 362            http: http.clone(),
 363            node_runtime,
 364            server: CopilotServer::Running(RunningCopilotServer {
 365                name: LanguageServerName(Arc::from("copilot")),
 366                lsp: Arc::new(server),
 367                sign_in_status: SignInStatus::Authorized,
 368                registered_buffers: Default::default(),
 369            }),
 370            buffers: Default::default(),
 371        });
 372        (this, fake_server)
 373    }
 374
 375    fn start_language_server(
 376        new_server_id: LanguageServerId,
 377        http: Arc<dyn HttpClient>,
 378        node_runtime: Arc<dyn NodeRuntime>,
 379        this: ModelHandle<Self>,
 380        mut cx: AsyncAppContext,
 381    ) -> impl Future<Output = ()> {
 382        async move {
 383            let start_language_server = async {
 384                let server_path = get_copilot_lsp(http).await?;
 385                let node_path = node_runtime.binary_path().await?;
 386                let arguments: Vec<OsString> = vec![server_path.into(), "--stdio".into()];
 387                let binary = LanguageServerBinary {
 388                    path: node_path,
 389                    arguments,
 390                };
 391
 392                let server = LanguageServer::new(
 393                    Arc::new(Mutex::new(None)),
 394                    new_server_id,
 395                    binary,
 396                    Path::new("/"),
 397                    None,
 398                    cx.clone(),
 399                )?;
 400
 401                server
 402                    .on_notification::<StatusNotification, _>(
 403                        |_, _| { /* Silence the notification */ },
 404                    )
 405                    .detach();
 406
 407                let server = server.initialize(Default::default()).await?;
 408
 409                let status = server
 410                    .request::<request::CheckStatus>(request::CheckStatusParams {
 411                        local_checks_only: false,
 412                    })
 413                    .await?;
 414
 415                server
 416                    .request::<request::SetEditorInfo>(request::SetEditorInfoParams {
 417                        editor_info: request::EditorInfo {
 418                            name: "zed".into(),
 419                            version: env!("CARGO_PKG_VERSION").into(),
 420                        },
 421                        editor_plugin_info: request::EditorPluginInfo {
 422                            name: "zed-copilot".into(),
 423                            version: "0.0.1".into(),
 424                        },
 425                    })
 426                    .await?;
 427
 428                anyhow::Ok((server, status))
 429            };
 430
 431            let server = start_language_server.await;
 432            this.update(&mut cx, |this, cx| {
 433                cx.notify();
 434                match server {
 435                    Ok((server, status)) => {
 436                        this.server = CopilotServer::Running(RunningCopilotServer {
 437                            name: LanguageServerName(Arc::from("copilot")),
 438                            lsp: server,
 439                            sign_in_status: SignInStatus::SignedOut,
 440                            registered_buffers: Default::default(),
 441                        });
 442                        cx.emit(Event::CopilotLanguageServerStarted);
 443                        this.update_sign_in_status(status, cx);
 444                    }
 445                    Err(error) => {
 446                        this.server = CopilotServer::Error(error.to_string().into());
 447                        cx.notify()
 448                    }
 449                }
 450            })
 451        }
 452    }
 453
 454    pub fn sign_in(&mut self, cx: &mut ModelContext<Self>) -> Task<Result<()>> {
 455        if let CopilotServer::Running(server) = &mut self.server {
 456            let task = match &server.sign_in_status {
 457                SignInStatus::Authorized { .. } => Task::ready(Ok(())).shared(),
 458                SignInStatus::SigningIn { task, .. } => {
 459                    cx.notify();
 460                    task.clone()
 461                }
 462                SignInStatus::SignedOut | SignInStatus::Unauthorized { .. } => {
 463                    let lsp = server.lsp.clone();
 464                    let task = cx
 465                        .spawn(|this, mut cx| async move {
 466                            let sign_in = async {
 467                                let sign_in = lsp
 468                                    .request::<request::SignInInitiate>(
 469                                        request::SignInInitiateParams {},
 470                                    )
 471                                    .await?;
 472                                match sign_in {
 473                                    request::SignInInitiateResult::AlreadySignedIn { user } => {
 474                                        Ok(request::SignInStatus::Ok { user })
 475                                    }
 476                                    request::SignInInitiateResult::PromptUserDeviceFlow(flow) => {
 477                                        this.update(&mut cx, |this, cx| {
 478                                            if let CopilotServer::Running(RunningCopilotServer {
 479                                                sign_in_status: status,
 480                                                ..
 481                                            }) = &mut this.server
 482                                            {
 483                                                if let SignInStatus::SigningIn {
 484                                                    prompt: prompt_flow,
 485                                                    ..
 486                                                } = status
 487                                                {
 488                                                    *prompt_flow = Some(flow.clone());
 489                                                    cx.notify();
 490                                                }
 491                                            }
 492                                        });
 493                                        let response = lsp
 494                                            .request::<request::SignInConfirm>(
 495                                                request::SignInConfirmParams {
 496                                                    user_code: flow.user_code,
 497                                                },
 498                                            )
 499                                            .await?;
 500                                        Ok(response)
 501                                    }
 502                                }
 503                            };
 504
 505                            let sign_in = sign_in.await;
 506                            this.update(&mut cx, |this, cx| match sign_in {
 507                                Ok(status) => {
 508                                    this.update_sign_in_status(status, cx);
 509                                    Ok(())
 510                                }
 511                                Err(error) => {
 512                                    this.update_sign_in_status(
 513                                        request::SignInStatus::NotSignedIn,
 514                                        cx,
 515                                    );
 516                                    Err(Arc::new(error))
 517                                }
 518                            })
 519                        })
 520                        .shared();
 521                    server.sign_in_status = SignInStatus::SigningIn {
 522                        prompt: None,
 523                        task: task.clone(),
 524                    };
 525                    cx.notify();
 526                    task
 527                }
 528            };
 529
 530            cx.foreground()
 531                .spawn(task.map_err(|err| anyhow!("{:?}", err)))
 532        } else {
 533            // If we're downloading, wait until download is finished
 534            // If we're in a stuck state, display to the user
 535            Task::ready(Err(anyhow!("copilot hasn't started yet")))
 536        }
 537    }
 538
 539    fn sign_out(&mut self, cx: &mut ModelContext<Self>) -> Task<Result<()>> {
 540        self.update_sign_in_status(request::SignInStatus::NotSignedIn, cx);
 541        if let CopilotServer::Running(RunningCopilotServer { lsp: server, .. }) = &self.server {
 542            let server = server.clone();
 543            cx.background().spawn(async move {
 544                server
 545                    .request::<request::SignOut>(request::SignOutParams {})
 546                    .await?;
 547                anyhow::Ok(())
 548            })
 549        } else {
 550            Task::ready(Err(anyhow!("copilot hasn't started yet")))
 551        }
 552    }
 553
 554    pub fn reinstall(&mut self, cx: &mut ModelContext<Self>) -> Task<()> {
 555        let start_task = cx
 556            .spawn({
 557                let http = self.http.clone();
 558                let node_runtime = self.node_runtime.clone();
 559                let server_id = self.server_id;
 560                move |this, cx| async move {
 561                    clear_copilot_dir().await;
 562                    Self::start_language_server(server_id, http, node_runtime, this, cx).await
 563                }
 564            })
 565            .shared();
 566
 567        self.server = CopilotServer::Starting {
 568            task: start_task.clone(),
 569        };
 570
 571        cx.notify();
 572
 573        cx.foreground().spawn(start_task)
 574    }
 575
 576    pub fn language_server(&self) -> Option<(&LanguageServerName, &Arc<LanguageServer>)> {
 577        if let CopilotServer::Running(server) = &self.server {
 578            Some((&server.name, &server.lsp))
 579        } else {
 580            None
 581        }
 582    }
 583
 584    pub fn register_buffer(&mut self, buffer: &ModelHandle<Buffer>, cx: &mut ModelContext<Self>) {
 585        let weak_buffer = buffer.downgrade();
 586        self.buffers.insert(weak_buffer.clone());
 587
 588        if let CopilotServer::Running(RunningCopilotServer {
 589            lsp: server,
 590            sign_in_status: status,
 591            registered_buffers,
 592            ..
 593        }) = &mut self.server
 594        {
 595            if !matches!(status, SignInStatus::Authorized { .. }) {
 596                return;
 597            }
 598
 599            registered_buffers.entry(buffer.id()).or_insert_with(|| {
 600                let uri: lsp::Url = uri_for_buffer(buffer, cx);
 601                let language_id = id_for_language(buffer.read(cx).language());
 602                let snapshot = buffer.read(cx).snapshot();
 603                server
 604                    .notify::<lsp::notification::DidOpenTextDocument>(
 605                        lsp::DidOpenTextDocumentParams {
 606                            text_document: lsp::TextDocumentItem {
 607                                uri: uri.clone(),
 608                                language_id: language_id.clone(),
 609                                version: 0,
 610                                text: snapshot.text(),
 611                            },
 612                        },
 613                    )
 614                    .log_err();
 615
 616                RegisteredBuffer {
 617                    uri,
 618                    language_id,
 619                    snapshot,
 620                    snapshot_version: 0,
 621                    pending_buffer_change: Task::ready(Some(())),
 622                    _subscriptions: [
 623                        cx.subscribe(buffer, |this, buffer, event, cx| {
 624                            this.handle_buffer_event(buffer, event, cx).log_err();
 625                        }),
 626                        cx.observe_release(buffer, move |this, _buffer, _cx| {
 627                            this.buffers.remove(&weak_buffer);
 628                            this.unregister_buffer(&weak_buffer);
 629                        }),
 630                    ],
 631                }
 632            });
 633        }
 634    }
 635
 636    fn handle_buffer_event(
 637        &mut self,
 638        buffer: ModelHandle<Buffer>,
 639        event: &language::Event,
 640        cx: &mut ModelContext<Self>,
 641    ) -> Result<()> {
 642        if let Ok(server) = self.server.as_running() {
 643            if let Some(registered_buffer) = server.registered_buffers.get_mut(&buffer.id()) {
 644                match event {
 645                    language::Event::Edited => {
 646                        let _ = registered_buffer.report_changes(&buffer, cx);
 647                    }
 648                    language::Event::Saved => {
 649                        server
 650                            .lsp
 651                            .notify::<lsp::notification::DidSaveTextDocument>(
 652                                lsp::DidSaveTextDocumentParams {
 653                                    text_document: lsp::TextDocumentIdentifier::new(
 654                                        registered_buffer.uri.clone(),
 655                                    ),
 656                                    text: None,
 657                                },
 658                            )?;
 659                    }
 660                    language::Event::FileHandleChanged | language::Event::LanguageChanged => {
 661                        let new_language_id = id_for_language(buffer.read(cx).language());
 662                        let new_uri = uri_for_buffer(&buffer, cx);
 663                        if new_uri != registered_buffer.uri
 664                            || new_language_id != registered_buffer.language_id
 665                        {
 666                            let old_uri = mem::replace(&mut registered_buffer.uri, new_uri);
 667                            registered_buffer.language_id = new_language_id;
 668                            server
 669                                .lsp
 670                                .notify::<lsp::notification::DidCloseTextDocument>(
 671                                    lsp::DidCloseTextDocumentParams {
 672                                        text_document: lsp::TextDocumentIdentifier::new(old_uri),
 673                                    },
 674                                )?;
 675                            server
 676                                .lsp
 677                                .notify::<lsp::notification::DidOpenTextDocument>(
 678                                    lsp::DidOpenTextDocumentParams {
 679                                        text_document: lsp::TextDocumentItem::new(
 680                                            registered_buffer.uri.clone(),
 681                                            registered_buffer.language_id.clone(),
 682                                            registered_buffer.snapshot_version,
 683                                            registered_buffer.snapshot.text(),
 684                                        ),
 685                                    },
 686                                )?;
 687                        }
 688                    }
 689                    _ => {}
 690                }
 691            }
 692        }
 693
 694        Ok(())
 695    }
 696
 697    fn unregister_buffer(&mut self, buffer: &WeakModelHandle<Buffer>) {
 698        if let Ok(server) = self.server.as_running() {
 699            if let Some(buffer) = server.registered_buffers.remove(&buffer.id()) {
 700                server
 701                    .lsp
 702                    .notify::<lsp::notification::DidCloseTextDocument>(
 703                        lsp::DidCloseTextDocumentParams {
 704                            text_document: lsp::TextDocumentIdentifier::new(buffer.uri),
 705                        },
 706                    )
 707                    .log_err();
 708            }
 709        }
 710    }
 711
 712    pub fn completions<T>(
 713        &mut self,
 714        buffer: &ModelHandle<Buffer>,
 715        position: T,
 716        cx: &mut ModelContext<Self>,
 717    ) -> Task<Result<Vec<Completion>>>
 718    where
 719        T: ToPointUtf16,
 720    {
 721        self.request_completions::<request::GetCompletions, _>(buffer, position, cx)
 722    }
 723
 724    pub fn completions_cycling<T>(
 725        &mut self,
 726        buffer: &ModelHandle<Buffer>,
 727        position: T,
 728        cx: &mut ModelContext<Self>,
 729    ) -> Task<Result<Vec<Completion>>>
 730    where
 731        T: ToPointUtf16,
 732    {
 733        self.request_completions::<request::GetCompletionsCycling, _>(buffer, position, cx)
 734    }
 735
 736    pub fn accept_completion(
 737        &mut self,
 738        completion: &Completion,
 739        cx: &mut ModelContext<Self>,
 740    ) -> Task<Result<()>> {
 741        let server = match self.server.as_authenticated() {
 742            Ok(server) => server,
 743            Err(error) => return Task::ready(Err(error)),
 744        };
 745        let request =
 746            server
 747                .lsp
 748                .request::<request::NotifyAccepted>(request::NotifyAcceptedParams {
 749                    uuid: completion.uuid.clone(),
 750                });
 751        cx.background().spawn(async move {
 752            request.await?;
 753            Ok(())
 754        })
 755    }
 756
 757    pub fn discard_completions(
 758        &mut self,
 759        completions: &[Completion],
 760        cx: &mut ModelContext<Self>,
 761    ) -> Task<Result<()>> {
 762        let server = match self.server.as_authenticated() {
 763            Ok(server) => server,
 764            Err(error) => return Task::ready(Err(error)),
 765        };
 766        let request =
 767            server
 768                .lsp
 769                .request::<request::NotifyRejected>(request::NotifyRejectedParams {
 770                    uuids: completions
 771                        .iter()
 772                        .map(|completion| completion.uuid.clone())
 773                        .collect(),
 774                });
 775        cx.background().spawn(async move {
 776            request.await?;
 777            Ok(())
 778        })
 779    }
 780
 781    fn request_completions<R, T>(
 782        &mut self,
 783        buffer: &ModelHandle<Buffer>,
 784        position: T,
 785        cx: &mut ModelContext<Self>,
 786    ) -> Task<Result<Vec<Completion>>>
 787    where
 788        R: 'static
 789            + lsp::request::Request<
 790                Params = request::GetCompletionsParams,
 791                Result = request::GetCompletionsResult,
 792            >,
 793        T: ToPointUtf16,
 794    {
 795        self.register_buffer(buffer, cx);
 796
 797        let server = match self.server.as_authenticated() {
 798            Ok(server) => server,
 799            Err(error) => return Task::ready(Err(error)),
 800        };
 801        let lsp = server.lsp.clone();
 802        let registered_buffer = server.registered_buffers.get_mut(&buffer.id()).unwrap();
 803        let snapshot = registered_buffer.report_changes(buffer, cx);
 804        let buffer = buffer.read(cx);
 805        let uri = registered_buffer.uri.clone();
 806        let position = position.to_point_utf16(buffer);
 807        let settings = language_settings(buffer.language_at(position).as_ref(), buffer.file(), cx);
 808        let tab_size = settings.tab_size;
 809        let hard_tabs = settings.hard_tabs;
 810        let relative_path = buffer
 811            .file()
 812            .map(|file| file.path().to_path_buf())
 813            .unwrap_or_default();
 814
 815        cx.foreground().spawn(async move {
 816            let (version, snapshot) = snapshot.await?;
 817            let result = lsp
 818                .request::<R>(request::GetCompletionsParams {
 819                    doc: request::GetCompletionsDocument {
 820                        uri,
 821                        tab_size: tab_size.into(),
 822                        indent_size: 1,
 823                        insert_spaces: !hard_tabs,
 824                        relative_path: relative_path.to_string_lossy().into(),
 825                        position: point_to_lsp(position),
 826                        version: version.try_into().unwrap(),
 827                    },
 828                })
 829                .await?;
 830            let completions = result
 831                .completions
 832                .into_iter()
 833                .map(|completion| {
 834                    let start = snapshot
 835                        .clip_point_utf16(point_from_lsp(completion.range.start), Bias::Left);
 836                    let end =
 837                        snapshot.clip_point_utf16(point_from_lsp(completion.range.end), Bias::Left);
 838                    Completion {
 839                        uuid: completion.uuid,
 840                        range: snapshot.anchor_before(start)..snapshot.anchor_after(end),
 841                        text: completion.text,
 842                    }
 843                })
 844                .collect();
 845            anyhow::Ok(completions)
 846        })
 847    }
 848
 849    pub fn status(&self) -> Status {
 850        match &self.server {
 851            CopilotServer::Starting { task } => Status::Starting { task: task.clone() },
 852            CopilotServer::Disabled => Status::Disabled,
 853            CopilotServer::Error(error) => Status::Error(error.clone()),
 854            CopilotServer::Running(RunningCopilotServer { sign_in_status, .. }) => {
 855                match sign_in_status {
 856                    SignInStatus::Authorized { .. } => Status::Authorized,
 857                    SignInStatus::Unauthorized { .. } => Status::Unauthorized,
 858                    SignInStatus::SigningIn { prompt, .. } => Status::SigningIn {
 859                        prompt: prompt.clone(),
 860                    },
 861                    SignInStatus::SignedOut => Status::SignedOut,
 862                }
 863            }
 864        }
 865    }
 866
 867    fn update_sign_in_status(
 868        &mut self,
 869        lsp_status: request::SignInStatus,
 870        cx: &mut ModelContext<Self>,
 871    ) {
 872        self.buffers.retain(|buffer| buffer.is_upgradable(cx));
 873
 874        if let Ok(server) = self.server.as_running() {
 875            match lsp_status {
 876                request::SignInStatus::Ok { .. }
 877                | request::SignInStatus::MaybeOk { .. }
 878                | request::SignInStatus::AlreadySignedIn { .. } => {
 879                    server.sign_in_status = SignInStatus::Authorized;
 880                    for buffer in self.buffers.iter().cloned().collect::<Vec<_>>() {
 881                        if let Some(buffer) = buffer.upgrade(cx) {
 882                            self.register_buffer(&buffer, cx);
 883                        }
 884                    }
 885                }
 886                request::SignInStatus::NotAuthorized { .. } => {
 887                    server.sign_in_status = SignInStatus::Unauthorized;
 888                    for buffer in self.buffers.iter().copied().collect::<Vec<_>>() {
 889                        self.unregister_buffer(&buffer);
 890                    }
 891                }
 892                request::SignInStatus::NotSignedIn => {
 893                    server.sign_in_status = SignInStatus::SignedOut;
 894                    for buffer in self.buffers.iter().copied().collect::<Vec<_>>() {
 895                        self.unregister_buffer(&buffer);
 896                    }
 897                }
 898            }
 899
 900            cx.notify();
 901        }
 902    }
 903}
 904
 905fn id_for_language(language: Option<&Arc<Language>>) -> String {
 906    let language_name = language.map(|language| language.name());
 907    match language_name.as_deref() {
 908        Some("Plain Text") => "plaintext".to_string(),
 909        Some(language_name) => language_name.to_lowercase(),
 910        None => "plaintext".to_string(),
 911    }
 912}
 913
 914fn uri_for_buffer(buffer: &ModelHandle<Buffer>, cx: &AppContext) -> lsp::Url {
 915    if let Some(file) = buffer.read(cx).file().and_then(|file| file.as_local()) {
 916        lsp::Url::from_file_path(file.abs_path(cx)).unwrap()
 917    } else {
 918        format!("buffer://{}", buffer.id()).parse().unwrap()
 919    }
 920}
 921
 922async fn clear_copilot_dir() {
 923    remove_matching(&paths::COPILOT_DIR, |_| true).await
 924}
 925
 926async fn get_copilot_lsp(http: Arc<dyn HttpClient>) -> anyhow::Result<PathBuf> {
 927    const SERVER_PATH: &'static str = "dist/agent.js";
 928
 929    ///Check for the latest copilot language server and download it if we haven't already
 930    async fn fetch_latest(http: Arc<dyn HttpClient>) -> anyhow::Result<PathBuf> {
 931        let release = latest_github_release("zed-industries/copilot", false, http.clone()).await?;
 932
 933        let version_dir = &*paths::COPILOT_DIR.join(format!("copilot-{}", release.name));
 934
 935        fs::create_dir_all(version_dir).await?;
 936        let server_path = version_dir.join(SERVER_PATH);
 937
 938        if fs::metadata(&server_path).await.is_err() {
 939            // Copilot LSP looks for this dist dir specifcially, so lets add it in.
 940            let dist_dir = version_dir.join("dist");
 941            fs::create_dir_all(dist_dir.as_path()).await?;
 942
 943            let url = &release
 944                .assets
 945                .get(0)
 946                .context("Github release for copilot contained no assets")?
 947                .browser_download_url;
 948
 949            let mut response = http
 950                .get(&url, Default::default(), true)
 951                .await
 952                .map_err(|err| anyhow!("error downloading copilot release: {}", err))?;
 953            let decompressed_bytes = GzipDecoder::new(BufReader::new(response.body_mut()));
 954            let archive = Archive::new(decompressed_bytes);
 955            archive.unpack(dist_dir).await?;
 956
 957            remove_matching(&paths::COPILOT_DIR, |entry| entry != version_dir).await;
 958        }
 959
 960        Ok(server_path)
 961    }
 962
 963    match fetch_latest(http).await {
 964        ok @ Result::Ok(..) => ok,
 965        e @ Err(..) => {
 966            e.log_err();
 967            // Fetch a cached binary, if it exists
 968            (|| async move {
 969                let mut last_version_dir = None;
 970                let mut entries = fs::read_dir(paths::COPILOT_DIR.as_path()).await?;
 971                while let Some(entry) = entries.next().await {
 972                    let entry = entry?;
 973                    if entry.file_type().await?.is_dir() {
 974                        last_version_dir = Some(entry.path());
 975                    }
 976                }
 977                let last_version_dir =
 978                    last_version_dir.ok_or_else(|| anyhow!("no cached binary"))?;
 979                let server_path = last_version_dir.join(SERVER_PATH);
 980                if server_path.exists() {
 981                    Ok(server_path)
 982                } else {
 983                    Err(anyhow!(
 984                        "missing executable in directory {:?}",
 985                        last_version_dir
 986                    ))
 987                }
 988            })()
 989            .await
 990        }
 991    }
 992}
 993
 994#[cfg(test)]
 995mod tests {
 996    use super::*;
 997    use gpui::{executor::Deterministic, TestAppContext};
 998
 999    #[gpui::test(iterations = 10)]
1000    async fn test_buffer_management(deterministic: Arc<Deterministic>, cx: &mut TestAppContext) {
1001        deterministic.forbid_parking();
1002        let (copilot, mut lsp) = Copilot::fake(cx);
1003
1004        let buffer_1 = cx.add_model(|cx| Buffer::new(0, cx.model_id() as u64, "Hello"));
1005        let buffer_1_uri: lsp::Url = format!("buffer://{}", buffer_1.id()).parse().unwrap();
1006        copilot.update(cx, |copilot, cx| copilot.register_buffer(&buffer_1, cx));
1007        assert_eq!(
1008            lsp.receive_notification::<lsp::notification::DidOpenTextDocument>()
1009                .await,
1010            lsp::DidOpenTextDocumentParams {
1011                text_document: lsp::TextDocumentItem::new(
1012                    buffer_1_uri.clone(),
1013                    "plaintext".into(),
1014                    0,
1015                    "Hello".into()
1016                ),
1017            }
1018        );
1019
1020        let buffer_2 = cx.add_model(|cx| Buffer::new(0, cx.model_id() as u64, "Goodbye"));
1021        let buffer_2_uri: lsp::Url = format!("buffer://{}", buffer_2.id()).parse().unwrap();
1022        copilot.update(cx, |copilot, cx| copilot.register_buffer(&buffer_2, cx));
1023        assert_eq!(
1024            lsp.receive_notification::<lsp::notification::DidOpenTextDocument>()
1025                .await,
1026            lsp::DidOpenTextDocumentParams {
1027                text_document: lsp::TextDocumentItem::new(
1028                    buffer_2_uri.clone(),
1029                    "plaintext".into(),
1030                    0,
1031                    "Goodbye".into()
1032                ),
1033            }
1034        );
1035
1036        buffer_1.update(cx, |buffer, cx| buffer.edit([(5..5, " world")], None, cx));
1037        assert_eq!(
1038            lsp.receive_notification::<lsp::notification::DidChangeTextDocument>()
1039                .await,
1040            lsp::DidChangeTextDocumentParams {
1041                text_document: lsp::VersionedTextDocumentIdentifier::new(buffer_1_uri.clone(), 1),
1042                content_changes: vec![lsp::TextDocumentContentChangeEvent {
1043                    range: Some(lsp::Range::new(
1044                        lsp::Position::new(0, 5),
1045                        lsp::Position::new(0, 5)
1046                    )),
1047                    range_length: None,
1048                    text: " world".into(),
1049                }],
1050            }
1051        );
1052
1053        // Ensure updates to the file are reflected in the LSP.
1054        buffer_1.update(cx, |buffer, cx| {
1055            buffer.file_updated(
1056                Arc::new(File {
1057                    abs_path: "/root/child/buffer-1".into(),
1058                    path: Path::new("child/buffer-1").into(),
1059                }),
1060                cx,
1061            )
1062        });
1063        assert_eq!(
1064            lsp.receive_notification::<lsp::notification::DidCloseTextDocument>()
1065                .await,
1066            lsp::DidCloseTextDocumentParams {
1067                text_document: lsp::TextDocumentIdentifier::new(buffer_1_uri),
1068            }
1069        );
1070        let buffer_1_uri = lsp::Url::from_file_path("/root/child/buffer-1").unwrap();
1071        assert_eq!(
1072            lsp.receive_notification::<lsp::notification::DidOpenTextDocument>()
1073                .await,
1074            lsp::DidOpenTextDocumentParams {
1075                text_document: lsp::TextDocumentItem::new(
1076                    buffer_1_uri.clone(),
1077                    "plaintext".into(),
1078                    1,
1079                    "Hello world".into()
1080                ),
1081            }
1082        );
1083
1084        // Ensure all previously-registered buffers are closed when signing out.
1085        lsp.handle_request::<request::SignOut, _, _>(|_, _| async {
1086            Ok(request::SignOutResult {})
1087        });
1088        copilot
1089            .update(cx, |copilot, cx| copilot.sign_out(cx))
1090            .await
1091            .unwrap();
1092        assert_eq!(
1093            lsp.receive_notification::<lsp::notification::DidCloseTextDocument>()
1094                .await,
1095            lsp::DidCloseTextDocumentParams {
1096                text_document: lsp::TextDocumentIdentifier::new(buffer_1_uri.clone()),
1097            }
1098        );
1099        assert_eq!(
1100            lsp.receive_notification::<lsp::notification::DidCloseTextDocument>()
1101                .await,
1102            lsp::DidCloseTextDocumentParams {
1103                text_document: lsp::TextDocumentIdentifier::new(buffer_2_uri.clone()),
1104            }
1105        );
1106
1107        // Ensure all previously-registered buffers are re-opened when signing in.
1108        lsp.handle_request::<request::SignInInitiate, _, _>(|_, _| async {
1109            Ok(request::SignInInitiateResult::AlreadySignedIn {
1110                user: "user-1".into(),
1111            })
1112        });
1113        copilot
1114            .update(cx, |copilot, cx| copilot.sign_in(cx))
1115            .await
1116            .unwrap();
1117        assert_eq!(
1118            lsp.receive_notification::<lsp::notification::DidOpenTextDocument>()
1119                .await,
1120            lsp::DidOpenTextDocumentParams {
1121                text_document: lsp::TextDocumentItem::new(
1122                    buffer_1_uri.clone(),
1123                    "plaintext".into(),
1124                    0,
1125                    "Hello world".into()
1126                ),
1127            }
1128        );
1129        assert_eq!(
1130            lsp.receive_notification::<lsp::notification::DidOpenTextDocument>()
1131                .await,
1132            lsp::DidOpenTextDocumentParams {
1133                text_document: lsp::TextDocumentItem::new(
1134                    buffer_2_uri.clone(),
1135                    "plaintext".into(),
1136                    0,
1137                    "Goodbye".into()
1138                ),
1139            }
1140        );
1141
1142        // Dropping a buffer causes it to be closed on the LSP side as well.
1143        cx.update(|_| drop(buffer_2));
1144        assert_eq!(
1145            lsp.receive_notification::<lsp::notification::DidCloseTextDocument>()
1146                .await,
1147            lsp::DidCloseTextDocumentParams {
1148                text_document: lsp::TextDocumentIdentifier::new(buffer_2_uri),
1149            }
1150        );
1151    }
1152
1153    struct File {
1154        abs_path: PathBuf,
1155        path: Arc<Path>,
1156    }
1157
1158    impl language::File for File {
1159        fn as_local(&self) -> Option<&dyn language::LocalFile> {
1160            Some(self)
1161        }
1162
1163        fn mtime(&self) -> std::time::SystemTime {
1164            unimplemented!()
1165        }
1166
1167        fn path(&self) -> &Arc<Path> {
1168            &self.path
1169        }
1170
1171        fn full_path(&self, _: &AppContext) -> PathBuf {
1172            unimplemented!()
1173        }
1174
1175        fn file_name<'a>(&'a self, _: &'a AppContext) -> &'a std::ffi::OsStr {
1176            unimplemented!()
1177        }
1178
1179        fn is_deleted(&self) -> bool {
1180            unimplemented!()
1181        }
1182
1183        fn as_any(&self) -> &dyn std::any::Any {
1184            unimplemented!()
1185        }
1186
1187        fn to_proto(&self) -> rpc::proto::File {
1188            unimplemented!()
1189        }
1190
1191        fn worktree_id(&self) -> usize {
1192            0
1193        }
1194    }
1195
1196    impl language::LocalFile for File {
1197        fn abs_path(&self, _: &AppContext) -> PathBuf {
1198            self.abs_path.clone()
1199        }
1200
1201        fn load(&self, _: &AppContext) -> Task<Result<String>> {
1202            unimplemented!()
1203        }
1204
1205        fn buffer_reloaded(
1206            &self,
1207            _: u64,
1208            _: &clock::Global,
1209            _: language::RopeFingerprint,
1210            _: language::LineEnding,
1211            _: std::time::SystemTime,
1212            _: &mut AppContext,
1213        ) {
1214            unimplemented!()
1215        }
1216    }
1217}