ZIm/crates/copilot/src/copilot.rs

497 lines
17 KiB
Rust

mod request;
mod sign_in;
use anyhow::{anyhow, Result};
use async_compression::futures::bufread::GzipDecoder;
use client::Client;
use futures::{future::Shared, FutureExt, TryFutureExt};
use gpui::{actions, AppContext, Entity, ModelContext, ModelHandle, MutableAppContext, Task};
use language::{point_from_lsp, point_to_lsp, Anchor, Bias, Buffer, BufferSnapshot, ToPointUtf16};
use lsp::LanguageServer;
use node_runtime::NodeRuntime;
use settings::Settings;
use smol::{fs, io::BufReader, stream::StreamExt};
use std::{
env::consts,
path::{Path, PathBuf},
sync::Arc,
};
use util::{
fs::remove_matching, github::latest_github_release, http::HttpClient, paths, ResultExt,
};
actions!(copilot, [SignIn, SignOut, NextSuggestion]);
pub fn init(client: Arc<Client>, node_runtime: Arc<NodeRuntime>, cx: &mut MutableAppContext) {
let copilot = cx.add_model(|cx| Copilot::start(client.http_client(), node_runtime, cx));
cx.set_global(copilot.clone());
cx.add_global_action(|_: &SignIn, cx| {
let copilot = Copilot::global(cx).unwrap();
copilot
.update(cx, |copilot, cx| copilot.sign_in(cx))
.detach_and_log_err(cx);
});
cx.add_global_action(|_: &SignOut, cx| {
let copilot = Copilot::global(cx).unwrap();
copilot
.update(cx, |copilot, cx| copilot.sign_out(cx))
.detach_and_log_err(cx);
});
sign_in::init(cx);
}
enum CopilotServer {
Downloading,
Error(Arc<str>),
Started {
server: Arc<LanguageServer>,
status: SignInStatus,
},
}
#[derive(Clone, Debug)]
enum SignInStatus {
Authorized {
_user: String,
},
Unauthorized {
_user: String,
},
SigningIn {
prompt: Option<request::PromptUserDeviceFlow>,
task: Shared<Task<Result<(), Arc<anyhow::Error>>>>,
},
SignedOut,
}
#[derive(Debug, PartialEq, Eq)]
pub enum Status {
Downloading,
Error(Arc<str>),
SignedOut,
SigningIn {
prompt: Option<request::PromptUserDeviceFlow>,
},
Unauthorized,
Authorized,
}
impl Status {
pub fn is_authorized(&self) -> bool {
matches!(self, Status::Authorized)
}
}
#[derive(Debug, PartialEq, Eq)]
pub struct Completion {
pub position: Anchor,
pub text: String,
}
pub struct Copilot {
server: CopilotServer,
}
impl Entity for Copilot {
type Event = ();
}
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<NodeRuntime>,
cx: &mut ModelContext<Self>,
) -> Self {
// TODO: Don't eagerly download the LSP
cx.spawn(|this, mut cx| async move {
let start_language_server = async {
let server_path = get_lsp_binary(http).await?;
let server =
LanguageServer::new(0, &server_path, &["--stdio"], Path::new("/"), cx.clone())?;
let server = server.initialize(Default::default()).await?;
let status = server
.request::<request::CheckStatus>(request::CheckStatusParams {
local_checks_only: false,
})
.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::Started {
server,
status: SignInStatus::SignedOut,
};
this.update_sign_in_status(status, cx);
}
Err(error) => {
this.server = CopilotServer::Error(error.to_string().into());
}
}
})
})
.detach();
Self {
server: CopilotServer::Downloading,
}
}
fn sign_in(&mut self, cx: &mut ModelContext<Self>) -> Task<Result<()>> {
if let CopilotServer::Started { server, status } = &mut self.server {
let task = match status {
SignInStatus::Authorized { .. } | SignInStatus::Unauthorized { .. } => {
Task::ready(Ok(())).shared()
}
SignInStatus::SigningIn { task, .. } => task.clone(),
SignInStatus::SignedOut => {
let server = server.clone();
let task = cx
.spawn(|this, mut cx| async move {
let sign_in = async {
let sign_in = server
.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::Started { status, .. } =
&mut this.server
{
if let SignInStatus::SigningIn {
prompt: prompt_flow,
..
} = status
{
*prompt_flow = Some(flow.clone());
cx.notify();
}
}
});
let response = server
.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(())
}
Err(error) => {
this.update_sign_in_status(
request::SignInStatus::NotSignedIn,
cx,
);
Err(Arc::new(error))
}
})
})
.shared();
*status = SignInStatus::SigningIn {
prompt: None,
task: task.clone(),
};
cx.notify();
task
}
};
cx.foreground()
.spawn(task.map_err(|err| anyhow!("{:?}", err)))
} else {
Task::ready(Err(anyhow!("copilot hasn't started yet")))
}
}
fn sign_out(&mut self, cx: &mut ModelContext<Self>) -> Task<Result<()>> {
if let CopilotServer::Started { server, status } = &mut self.server {
*status = SignInStatus::SignedOut;
cx.notify();
let server = server.clone();
cx.background().spawn(async move {
server
.request::<request::SignOut>(request::SignOutParams {})
.await?;
anyhow::Ok(())
})
} else {
Task::ready(Err(anyhow!("copilot hasn't started yet")))
}
}
pub fn completion<T>(
&self,
buffer: &ModelHandle<Buffer>,
position: T,
cx: &mut ModelContext<Self>,
) -> Task<Result<Option<Completion>>>
where
T: ToPointUtf16,
{
let server = match self.authorized_server() {
Ok(server) => server,
Err(error) => return Task::ready(Err(error)),
};
let buffer = buffer.read(cx).snapshot();
let request = server
.request::<request::GetCompletions>(build_completion_params(&buffer, position, cx));
cx.background().spawn(async move {
let result = request.await?;
let completion = result
.completions
.into_iter()
.next()
.map(|completion| completion_from_lsp(completion, &buffer));
anyhow::Ok(completion)
})
}
pub fn completions_cycling<T>(
&self,
buffer: &ModelHandle<Buffer>,
position: T,
cx: &mut ModelContext<Self>,
) -> Task<Result<Vec<Completion>>>
where
T: ToPointUtf16,
{
let server = match self.authorized_server() {
Ok(server) => server,
Err(error) => return Task::ready(Err(error)),
};
let buffer = buffer.read(cx).snapshot();
let request = server.request::<request::GetCompletionsCycling>(build_completion_params(
&buffer, position, cx,
));
cx.background().spawn(async move {
let result = request.await?;
let completions = result
.completions
.into_iter()
.map(|completion| completion_from_lsp(completion, &buffer))
.collect();
anyhow::Ok(completions)
})
}
pub fn status(&self) -> Status {
match &self.server {
CopilotServer::Downloading => Status::Downloading,
CopilotServer::Error(error) => Status::Error(error.clone()),
CopilotServer::Started { status, .. } => match status {
SignInStatus::Authorized { .. } => Status::Authorized,
SignInStatus::Unauthorized { .. } => Status::Unauthorized,
SignInStatus::SigningIn { prompt, .. } => Status::SigningIn {
prompt: prompt.clone(),
},
SignInStatus::SignedOut => Status::SignedOut,
},
}
}
fn update_sign_in_status(
&mut self,
lsp_status: request::SignInStatus,
cx: &mut ModelContext<Self>,
) {
if let CopilotServer::Started { status, .. } = &mut self.server {
*status = match lsp_status {
request::SignInStatus::Ok { user } | request::SignInStatus::MaybeOk { user } => {
SignInStatus::Authorized { _user: user }
}
request::SignInStatus::NotAuthorized { user } => {
SignInStatus::Unauthorized { _user: user }
}
_ => SignInStatus::SignedOut,
};
cx.notify();
}
}
fn authorized_server(&self) -> Result<Arc<LanguageServer>> {
match &self.server {
CopilotServer::Downloading => Err(anyhow!("copilot is still downloading")),
CopilotServer::Error(error) => Err(anyhow!(
"copilot was not started because of an error: {}",
error
)),
CopilotServer::Started { server, status } => {
if matches!(status, SignInStatus::Authorized { .. }) {
Ok(server.clone())
} else {
Err(anyhow!("must sign in before using copilot"))
}
}
}
}
}
fn build_completion_params<T>(
buffer: &BufferSnapshot,
position: T,
cx: &AppContext,
) -> request::GetCompletionsParams
where
T: ToPointUtf16,
{
let position = position.to_point_utf16(&buffer);
let language_name = buffer.language_at(position).map(|language| language.name());
let language_name = language_name.as_deref();
let path;
let relative_path;
if let Some(file) = buffer.file() {
if let Some(file) = file.as_local() {
path = file.abs_path(cx);
} else {
path = file.full_path(cx);
}
relative_path = file.path().to_path_buf();
} else {
path = PathBuf::from("/untitled");
relative_path = PathBuf::from("untitled");
}
let settings = cx.global::<Settings>();
let language_id = match language_name {
Some("Plain Text") => "plaintext".to_string(),
Some(language_name) => language_name.to_lowercase(),
None => "plaintext".to_string(),
};
request::GetCompletionsParams {
doc: request::GetCompletionsDocument {
source: buffer.text(),
tab_size: settings.tab_size(language_name).into(),
indent_size: 1,
insert_spaces: !settings.hard_tabs(language_name),
uri: lsp::Url::from_file_path(&path).unwrap(),
path: path.to_string_lossy().into(),
relative_path: relative_path.to_string_lossy().into(),
language_id,
position: point_to_lsp(position),
version: 0,
},
}
}
fn completion_from_lsp(completion: request::Completion, buffer: &BufferSnapshot) -> Completion {
let position = buffer.clip_point_utf16(point_from_lsp(completion.position), Bias::Left);
Completion {
position: buffer.anchor_before(position),
text: completion.display_text,
}
}
async fn get_lsp_binary(http: Arc<dyn HttpClient>) -> anyhow::Result<PathBuf> {
///Check for the latest copilot language server and download it if we haven't already
async fn fetch_latest(http: Arc<dyn HttpClient>) -> anyhow::Result<PathBuf> {
let release = latest_github_release("zed-industries/copilot", http.clone()).await?;
let asset_name = format!("copilot-darwin-{}.gz", consts::ARCH);
let asset = release
.assets
.iter()
.find(|asset| asset.name == asset_name)
.ok_or_else(|| anyhow!("no asset found matching {:?}", asset_name))?;
fs::create_dir_all(&*paths::COPILOT_DIR).await?;
let destination_path =
paths::COPILOT_DIR.join(format!("copilot-{}-{}", release.name, consts::ARCH));
if fs::metadata(&destination_path).await.is_err() {
let mut response = http
.get(&asset.browser_download_url, Default::default(), true)
.await
.map_err(|err| anyhow!("error downloading release: {}", err))?;
let decompressed_bytes = GzipDecoder::new(BufReader::new(response.body_mut()));
let mut file = fs::File::create(&destination_path).await?;
futures::io::copy(decompressed_bytes, &mut file).await?;
fs::set_permissions(
&destination_path,
<fs::Permissions as fs::unix::PermissionsExt>::from_mode(0o755),
)
.await?;
remove_matching(&paths::COPILOT_DIR, |entry| entry != destination_path).await;
}
Ok(destination_path)
}
match fetch_latest(http).await {
ok @ Result::Ok(..) => ok,
e @ Err(..) => {
e.log_err();
// Fetch a cached binary, if it exists
(|| async move {
let mut last = None;
let mut entries = fs::read_dir(paths::COPILOT_DIR.as_path()).await?;
while let Some(entry) = entries.next().await {
last = Some(entry?.path());
}
last.ok_or_else(|| anyhow!("no cached binary"))
})()
.await
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use gpui::TestAppContext;
use util::http;
#[gpui::test]
async fn test_smoke(cx: &mut TestAppContext) {
Settings::test_async(cx);
let http = http::client();
let node_runtime = NodeRuntime::new(http.clone(), cx.background());
let copilot = cx.add_model(|cx| Copilot::start(http, node_runtime, cx));
smol::Timer::after(std::time::Duration::from_secs(2)).await;
copilot
.update(cx, |copilot, cx| copilot.sign_in(cx))
.await
.unwrap();
copilot.read_with(cx, |copilot, _| copilot.status());
let buffer = cx.add_model(|cx| language::Buffer::new(0, "fn foo() -> ", cx));
dbg!(copilot
.update(cx, |copilot, cx| copilot.completion(&buffer, 12, cx))
.await
.unwrap());
dbg!(copilot
.update(cx, |copilot, cx| copilot
.completions_cycling(&buffer, 12, cx))
.await
.unwrap());
}
}