1201 lines · 44.0 KB · Rust
RawDownload
pub mod request;mod sign_in; use anyhow::{anyhow, Context, Result};use async_compression::futures::bufread::GzipDecoder;use async_tar::Archive;use collections::{HashMap, HashSet};use futures::{channel::oneshot, future::Shared, Future, FutureExt, TryFutureExt};use gpui::{    actions, AppContext, AsyncAppContext, Entity, ModelContext, ModelHandle, Task, WeakModelHandle,};use language::{    language_settings::{all_language_settings, language_settings},    point_from_lsp, point_to_lsp, Anchor, Bias, Buffer, BufferSnapshot, Language, PointUtf16,    ToPointUtf16,};use log::{debug, error};use lsp::{LanguageServer, LanguageServerBinary, LanguageServerId};use node_runtime::NodeRuntime;use request::{LogMessage, StatusNotification};use settings::SettingsStore;use smol::{fs, io::BufReader, stream::StreamExt};use std::{    ffi::OsString,    mem,    ops::Range,    path::{Path, PathBuf},    pin::Pin,    sync::Arc,};use util::{    fs::remove_matching, github::latest_github_release, http::HttpClient, paths, ResultExt,}; const COPILOT_AUTH_NAMESPACE: &'static str = "copilot_auth";actions!(copilot_auth, [SignIn, SignOut]); const COPILOT_NAMESPACE: &'static str = "copilot";actions!(    copilot,    [Suggest, NextSuggestion, PreviousSuggestion, Reinstall]); pub fn init(http: Arc<dyn HttpClient>, node_runtime: Arc<dyn NodeRuntime>, cx: &mut AppContext) {    let copilot = cx.add_model({        let node_runtime = node_runtime.clone();        move |cx| Copilot::start(http, node_runtime, cx)    });    cx.set_global(copilot.clone());     cx.observe(&copilot, |handle, cx| {        let status = handle.read(cx).status();        cx.update_default_global::<collections::CommandPaletteFilter, _, _>(move |filter, _cx| {            match status {                Status::Disabled => {                    filter.filtered_namespaces.insert(COPILOT_NAMESPACE);                    filter.filtered_namespaces.insert(COPILOT_AUTH_NAMESPACE);                }                Status::Authorized => {                    filter.filtered_namespaces.remove(COPILOT_NAMESPACE);                    filter.filtered_namespaces.remove(COPILOT_AUTH_NAMESPACE);                }                _ => {                    filter.filtered_namespaces.insert(COPILOT_NAMESPACE);                    filter.filtered_namespaces.remove(COPILOT_AUTH_NAMESPACE);                }            }        });    })    .detach();     sign_in::init(cx);    cx.add_global_action(|_: &SignIn, cx| {        if let Some(copilot) = Copilot::global(cx) {            copilot                .update(cx, |copilot, cx| copilot.sign_in(cx))                .detach_and_log_err(cx);        }    });    cx.add_global_action(|_: &SignOut, cx| {        if let Some(copilot) = Copilot::global(cx) {            copilot                .update(cx, |copilot, cx| copilot.sign_out(cx))                .detach_and_log_err(cx);        }    });     cx.add_global_action(|_: &Reinstall, cx| {        if let Some(copilot) = Copilot::global(cx) {            copilot                .update(cx, |copilot, cx| copilot.reinstall(cx))                .detach();        }    });} enum CopilotServer {    Disabled,    Starting { task: Shared<Task<()>> },    Error(Arc<str>),    Running(RunningCopilotServer),} impl CopilotServer {    fn as_authenticated(&mut self) -> Result<&mut RunningCopilotServer> {        let server = self.as_running()?;        if matches!(server.sign_in_status, SignInStatus::Authorized { .. }) {            Ok(server)        } else {            Err(anyhow!("must sign in before using copilot"))        }    }     fn as_running(&mut self) -> Result<&mut RunningCopilotServer> {        match self {            CopilotServer::Starting { .. } => Err(anyhow!("copilot is still starting")),            CopilotServer::Disabled => Err(anyhow!("copilot is disabled")),            CopilotServer::Error(error) => Err(anyhow!(                "copilot was not started because of an error: {}",                error            )),            CopilotServer::Running(server) => Ok(server),        }    }} struct RunningCopilotServer {    lsp: Arc<LanguageServer>,    sign_in_status: SignInStatus,    registered_buffers: HashMap<usize, RegisteredBuffer>,} #[derive(Clone, Debug)]enum SignInStatus {    Authorized,    Unauthorized,    SigningIn {        prompt: Option<request::PromptUserDeviceFlow>,        task: Shared<Task<Result<(), Arc<anyhow::Error>>>>,    },    SignedOut,} #[derive(Debug, Clone)]pub enum Status {    Starting {        task: Shared<Task<()>>,    },    Error(Arc<str>),    Disabled,    SignedOut,    SigningIn {        prompt: Option<request::PromptUserDeviceFlow>,    },    Unauthorized,    Authorized,} impl Status {    pub fn is_authorized(&self) -> bool {        matches!(self, Status::Authorized)    }} struct RegisteredBuffer {    uri: lsp::Url,    language_id: String,    snapshot: BufferSnapshot,    snapshot_version: i32,    _subscriptions: [gpui::Subscription; 2],    pending_buffer_change: Task<Option<()>>,} impl RegisteredBuffer {    fn report_changes(        &mut self,        buffer: &ModelHandle<Buffer>,        cx: &mut ModelContext<Copilot>,    ) -> oneshot::Receiver<(i32, BufferSnapshot)> {        let (done_tx, done_rx) = oneshot::channel();         if buffer.read(cx).version() == self.snapshot.version {            let _ = done_tx.send((self.snapshot_version, self.snapshot.clone()));        } else {            let buffer = buffer.downgrade();            let id = buffer.id();            let prev_pending_change =                mem::replace(&mut self.pending_buffer_change, Task::ready(None));            self.pending_buffer_change = cx.spawn_weak(|copilot, mut cx| async move {                prev_pending_change.await;                 let old_version = copilot.upgrade(&cx)?.update(&mut cx, |copilot, _| {                    let server = copilot.server.as_authenticated().log_err()?;                    let buffer = server.registered_buffers.get_mut(&id)?;                    Some(buffer.snapshot.version.clone())                })?;                let new_snapshot = buffer                    .upgrade(&cx)?                    .read_with(&cx, |buffer, _| buffer.snapshot());                 let content_changes = cx                    .background()                    .spawn({                        let new_snapshot = new_snapshot.clone();                        async move {                            new_snapshot                                .edits_since::<(PointUtf16, usize)>(&old_version)                                .map(|edit| {                                    let edit_start = edit.new.start.0;                                    let edit_end = edit_start + (edit.old.end.0 - edit.old.start.0);                                    let new_text = new_snapshot                                        .text_for_range(edit.new.start.1..edit.new.end.1)                                        .collect();                                    lsp::TextDocumentContentChangeEvent {                                        range: Some(lsp::Range::new(                                            point_to_lsp(edit_start),                                            point_to_lsp(edit_end),                                        )),                                        range_length: None,                                        text: new_text,                                    }                                })                                .collect::<Vec<_>>()                        }                    })                    .await;                 copilot.upgrade(&cx)?.update(&mut cx, |copilot, _| {                    let server = copilot.server.as_authenticated().log_err()?;                    let buffer = server.registered_buffers.get_mut(&id)?;                    if !content_changes.is_empty() {                        buffer.snapshot_version += 1;                        buffer.snapshot = new_snapshot;                        server                            .lsp                            .notify::<lsp::notification::DidChangeTextDocument>(                                lsp::DidChangeTextDocumentParams {                                    text_document: lsp::VersionedTextDocumentIdentifier::new(                                        buffer.uri.clone(),                                        buffer.snapshot_version,                                    ),                                    content_changes,                                },                            )                            .log_err();                    }                    let _ = done_tx.send((buffer.snapshot_version, buffer.snapshot.clone()));                    Some(())                })?;                 Some(())            });        }         done_rx    }} #[derive(Debug)]pub struct Completion {    pub uuid: String,    pub range: Range<Anchor>,    pub text: String,} pub struct Copilot {    http: Arc<dyn HttpClient>,    node_runtime: Arc<dyn NodeRuntime>,    server: CopilotServer,    buffers: HashSet<WeakModelHandle<Buffer>>,} impl Entity for Copilot {    type Event = ();     fn app_will_quit(        &mut self,        _: &mut AppContext,    ) -> Option<Pin<Box<dyn 'static + Future<Output = ()>>>> {        match mem::replace(&mut self.server, CopilotServer::Disabled) {            CopilotServer::Running(server) => Some(Box::pin(async move {                if let Some(shutdown) = server.lsp.shutdown() {                    shutdown.await;                }            })),            _ => None,        }    }} impl Copilot {    pub fn global(cx: &AppContext) -> Option<ModelHandle<Self>> {        if cx.has_global::<ModelHandle<Self>>() {            Some(cx.global::<ModelHandle<Self>>().clone())        } else {            None        }    }     fn start(        http: Arc<dyn HttpClient>,        node_runtime: Arc<dyn NodeRuntime>,        cx: &mut ModelContext<Self>,    ) -> Self {        let mut this = Self {            http,            node_runtime,            server: CopilotServer::Disabled,            buffers: Default::default(),        };        this.enable_or_disable_copilot(cx);        cx.observe_global::<SettingsStore, _>(move |this, cx| this.enable_or_disable_copilot(cx))            .detach();        this    }     fn enable_or_disable_copilot(&mut self, cx: &mut ModelContext<Copilot>) {        let http = self.http.clone();        let node_runtime = self.node_runtime.clone();        if all_language_settings(None, cx).copilot_enabled(None, None) {            if matches!(self.server, CopilotServer::Disabled) {                let start_task = cx                    .spawn({                        move |this, cx| Self::start_language_server(http, node_runtime, this, cx)                    })                    .shared();                self.server = CopilotServer::Starting { task: start_task };                cx.notify();            }        } else {            self.server = CopilotServer::Disabled;            cx.notify();        }    }     #[cfg(any(test, feature = "test-support"))]    pub fn fake(cx: &mut gpui::TestAppContext) -> (ModelHandle<Self>, lsp::FakeLanguageServer) {        use node_runtime::FakeNodeRuntime;         let (server, fake_server) =            LanguageServer::fake("copilot".into(), Default::default(), cx.to_async());        let http = util::http::FakeHttpClient::create(|_| async { unreachable!() });        let node_runtime = FakeNodeRuntime::new();        let this = cx.add_model(|_| Self {            http: http.clone(),            node_runtime,            server: CopilotServer::Running(RunningCopilotServer {                lsp: Arc::new(server),                sign_in_status: SignInStatus::Authorized,                registered_buffers: Default::default(),            }),            buffers: Default::default(),        });        (this, fake_server)    }     fn start_language_server(        http: Arc<dyn HttpClient>,        node_runtime: Arc<dyn NodeRuntime>,        this: ModelHandle<Self>,        mut cx: AsyncAppContext,    ) -> impl Future<Output = ()> {        async move {            let start_language_server = async {                let server_path = get_copilot_lsp(http).await?;                let node_path = node_runtime.binary_path().await?;                let arguments: Vec<OsString> = vec![server_path.into(), "--stdio".into()];                let binary = LanguageServerBinary {                    path: node_path,                    arguments,                };                let server = LanguageServer::new(                    LanguageServerId(0),                    binary,                    Path::new("/"),                    None,                    cx.clone(),                )?;                 server                    .on_notification::<LogMessage, _>(|params, _cx| {                        match params.level {                            // Copilot is pretty aggressive about logging                            0 => debug!("copilot: {}", params.message),                            1 => debug!("copilot: {}", params.message),                            _ => error!("copilot: {}", params.message),                        }                         debug!("copilot metadata: {}", params.metadata_str);                        debug!("copilot extra: {:?}", params.extra);                    })                    .detach();                 server                    .on_notification::<StatusNotification, _>(                        |_, _| { /* Silence the notification */ },                    )                    .detach();                 let server = server.initialize(Default::default()).await?;                 let status = server                    .request::<request::CheckStatus>(request::CheckStatusParams {                        local_checks_only: false,                    })                    .await?;                 server                    .request::<request::SetEditorInfo>(request::SetEditorInfoParams {                        editor_info: request::EditorInfo {                            name: "zed".into(),                            version: env!("CARGO_PKG_VERSION").into(),                        },                        editor_plugin_info: request::EditorPluginInfo {                            name: "zed-copilot".into(),                            version: "0.0.1".into(),                        },                    })                    .await?;                 anyhow::Ok((server, status))            };             let server = start_language_server.await;            this.update(&mut cx, |this, cx| {                cx.notify();                match server {                    Ok((server, status)) => {                        this.server = CopilotServer::Running(RunningCopilotServer {                            lsp: server,                            sign_in_status: SignInStatus::SignedOut,                            registered_buffers: Default::default(),                        });                        this.update_sign_in_status(status, cx);                    }                    Err(error) => {                        this.server = CopilotServer::Error(error.to_string().into());                        cx.notify()                    }                }            })        }    }     pub fn sign_in(&mut self, cx: &mut ModelContext<Self>) -> Task<Result<()>> {        if let CopilotServer::Running(server) = &mut self.server {            let task = match &server.sign_in_status {                SignInStatus::Authorized { .. } => Task::ready(Ok(())).shared(),                SignInStatus::SigningIn { task, .. } => {                    cx.notify();                    task.clone()                }                SignInStatus::SignedOut | SignInStatus::Unauthorized { .. } => {                    let lsp = server.lsp.clone();                    let task = cx                        .spawn(|this, mut cx| async move {                            let sign_in = async {                                let sign_in = lsp                                    .request::<request::SignInInitiate>(                                        request::SignInInitiateParams {},                                    )                                    .await?;                                match sign_in {                                    request::SignInInitiateResult::AlreadySignedIn { user } => {                                        Ok(request::SignInStatus::Ok { user })                                    }                                    request::SignInInitiateResult::PromptUserDeviceFlow(flow) => {                                        this.update(&mut cx, |this, cx| {                                            if let CopilotServer::Running(RunningCopilotServer {                                                sign_in_status: status,                                                ..                                            }) = &mut this.server                                            {                                                if let SignInStatus::SigningIn {                                                    prompt: prompt_flow,                                                    ..                                                } = status                                                {                                                    *prompt_flow = Some(flow.clone());                                                    cx.notify();                                                }                                            }                                        });                                        let response = lsp                                            .request::<request::SignInConfirm>(                                                request::SignInConfirmParams {                                                    user_code: flow.user_code,                                                },                                            )                                            .await?;                                        Ok(response)                                    }                                }                            };                             let sign_in = sign_in.await;                            this.update(&mut cx, |this, cx| match sign_in {                                Ok(status) => {                                    this.update_sign_in_status(status, cx);                                    Ok(())

Lines 1–500 of 1201.

Keyboard shortcuts

General

Show keyboard shortcuts
?
Focus the current search
/

Issues and pull requests

Open filters
f

Issues, pull requests and Inbox

Select next item
j↓
Select previous item
k↑
Open selected item
Enter
Dismiss filters or clear selection
Esc

Repository code

Go to file
t