mod websocket; use std::sync::Arc; use anyhow::{Context, Result, anyhow};use cloud_api_types::websocket_protocol::{PROTOCOL_VERSION, PROTOCOL_VERSION_HEADER_NAME};pub use cloud_api_types::*;use futures::AsyncReadExt as _;use gpui::{App, Task};use gpui_tokio::Tokio;use http_client::http::request;use http_client::{ AsyncBody, HttpClientWithUrl, HttpRequestExt, Json, Method, Request, StatusCode,};use parking_lot::RwLock;use thiserror::Error;use yawc::WebSocket; use crate::websocket::Connection; struct Credentials { user_id: u32, access_token: String,} #[derive(Debug, Error)]pub enum ClientApiError { #[error("Unauthorized")] Unauthorized, #[error(transparent)] Other(#[from] anyhow::Error),} pub struct CloudApiClient { credentials: RwLock<Option<Credentials>>, http_client: Arc<HttpClientWithUrl>,} impl CloudApiClient { pub fn new(http_client: Arc<HttpClientWithUrl>) -> Self { Self { credentials: RwLock::new(None), http_client, } } pub fn has_credentials(&self) -> bool { self.credentials.read().is_some() } pub fn set_credentials(&self, user_id: u32, access_token: String) { *self.credentials.write() = Some(Credentials { user_id, access_token, }); } pub fn clear_credentials(&self) { *self.credentials.write() = None; } fn build_request( &self, req: request::Builder, body: impl Into<AsyncBody>, ) -> Result<Request<AsyncBody>> { let credentials = self.credentials.read(); let credentials = credentials.as_ref().context("no credentials provided")?; build_request(req, body, credentials) } pub async fn get_authenticated_user( &self, ) -> Result<GetAuthenticatedUserResponse, ClientApiError> { let request = self.build_request( Request::builder().method(Method::GET).uri( self.http_client .build_zed_cloud_url("/client/users/me")? .as_ref(), ), AsyncBody::default(), )?; let mut response = self.http_client.send(request).await?; if !response.status().is_success() { if response.status() == StatusCode::UNAUTHORIZED { return Err(ClientApiError::Unauthorized); } let mut body = String::new(); response .body_mut() .read_to_string(&mut body) .await .context("failed to read response body")?; return Err(ClientApiError::Other(anyhow::anyhow!( "Failed to get authenticated user.\nStatus: {:?}\nBody: {body}", response.status() ))); } let mut body = String::new(); response .body_mut() .read_to_string(&mut body) .await .context("failed to read response body")?; Ok(serde_json::from_str(&body).context("failed to parse response body")?) } pub fn connect(&self, cx: &App) -> Result<Task<Result<Connection>>> { let mut connect_url = self .http_client .build_zed_cloud_url("/client/users/connect")?; connect_url .set_scheme(match connect_url.scheme() { "https" => "wss", "http" => "ws", scheme => Err(anyhow!("invalid URL scheme: {scheme}"))?, }) .map_err(|_| anyhow!("failed to set URL scheme"))?; let credentials = self.credentials.read(); let credentials = credentials.as_ref().context("no credentials provided")?; let authorization_header = format!("{} {}", credentials.user_id, credentials.access_token); Ok(Tokio::spawn_result(cx, async move { let ws = WebSocket::connect(connect_url) .with_request( request::Builder::new() .header("Authorization", authorization_header) .header(PROTOCOL_VERSION_HEADER_NAME, PROTOCOL_VERSION.to_string()), ) .await?; Ok(Connection::new(ws)) })) } pub async fn create_llm_token( &self, system_id: Option<String>, organization_id: Option<OrganizationId>, ) -> Result<CreateLlmTokenResponse, ClientApiError> { let request_builder = Request::builder() .method(Method::POST) .uri( self.http_client .build_zed_cloud_url("/client/llm_tokens")? .as_ref(), ) .when_some(system_id, |builder, system_id| { builder.header(ZED_SYSTEM_ID_HEADER_NAME, system_id) }); let request = self.build_request( request_builder, Json(CreateLlmTokenBody { organization_id }), )?; let mut response = self.http_client.send(request).await?; if !response.status().is_success() { if response.status() == StatusCode::UNAUTHORIZED { return Err(ClientApiError::Unauthorized); } let mut body = String::new(); response .body_mut() .read_to_string(&mut body) .await .context("failed to read response body")?; return Err(ClientApiError::Other(anyhow::anyhow!( "Failed to create LLM token.\nStatus: {:?}\nBody: {body}", response.status() ))); } let mut body = String::new(); response .body_mut() .read_to_string(&mut body) .await .context("failed to read response body")?; Ok(serde_json::from_str(&body).context("failed to parse response body")?) } pub async fn validate_credentials(&self, user_id: u32, access_token: &str) -> Result<bool> { let request = build_request( Request::builder().method(Method::GET).uri( self.http_client .build_zed_cloud_url("/client/users/me")? .as_ref(), ), AsyncBody::default(), &Credentials { user_id, access_token: access_token.into(), }, )?; let mut response = self.http_client.send(request).await?; if response.status().is_success() { Ok(true) } else { let mut body = String::new(); response.body_mut().read_to_string(&mut body).await?; if response.status() == StatusCode::UNAUTHORIZED { Ok(false) } else { Err(anyhow!( "Failed to get authenticated user.\nStatus: {:?}\nBody: {body}", response.status() )) } } } pub async fn submit_agent_feedback(&self, body: SubmitAgentThreadFeedbackBody) -> Result<()> { let request = self.build_request( Request::builder().method(Method::POST).uri( self.http_client .build_zed_cloud_url("/client/feedback/agent_thread")? .as_ref(), ), AsyncBody::from(serde_json::to_string(&body)?), )?; let mut response = self.http_client.send(request).await?; if !response.status().is_success() { let mut body = String::new(); response.body_mut().read_to_string(&mut body).await?; anyhow::bail!( "Failed to submit agent feedback.\nStatus: {:?}\nBody: {body}", response.status() ) } Ok(()) } pub async fn submit_agent_feedback_comments( &self, body: SubmitAgentThreadFeedbackCommentsBody, ) -> Result<()> { let request = self.build_request( Request::builder().method(Method::POST).uri( self.http_client .build_zed_cloud_url("/client/feedback/agent_thread_comments")? .as_ref(), ), AsyncBody::from(serde_json::to_string(&body)?), )?; let mut response = self.http_client.send(request).await?; if !response.status().is_success() { let mut body = String::new(); response.body_mut().read_to_string(&mut body).await?; anyhow::bail!( "Failed to submit agent feedback comments.\nStatus: {:?}\nBody: {body}", response.status() ) } Ok(()) } pub async fn submit_edit_prediction_feedback( &self, body: SubmitEditPredictionFeedbackBody, ) -> Result<()> { let request = self.build_request( Request::builder().method(Method::POST).uri( self.http_client .build_zed_cloud_url("/client/feedback/edit_prediction")? .as_ref(), ), AsyncBody::from(serde_json::to_string(&body)?), )?; let mut response = self.http_client.send(request).await?; if !response.status().is_success() { let mut body = String::new(); response.body_mut().read_to_string(&mut body).await?; anyhow::bail!( "Failed to submit edit prediction feedback.\nStatus: {:?}\nBody: {body}", response.status() ) } Ok(()) }} fn build_request( req: request::Builder, body: impl Into<AsyncBody>, credentials: &Credentials,) -> Result<Request<AsyncBody>> { Ok(req .header("Content-Type", "application/json") .header( "Authorization", format!("{} {}", credentials.user_id, credentials.access_token), ) .body(body.into())?)}
mod websocket;
use std::sync::Arc;
use anyhow::{Context, Result, anyhow};
use cloud_api_types::websocket_protocol::{PROTOCOL_VERSION, PROTOCOL_VERSION_HEADER_NAME};
pub use cloud_api_types::*;
use futures::AsyncReadExt as _;
use gpui::{App, Task};
use gpui_tokio::Tokio;
use http_client::http::request;
use http_client::{
AsyncBody, HttpClientWithUrl, HttpRequestExt, Json, Method, Request, StatusCode,
};
use parking_lot::RwLock;
use thiserror::Error;
use yawc::WebSocket;
use crate::websocket::Connection;
struct Credentials {
user_id: u32,
access_token: String,
}
#[derive(Debug, Error)]
pub enum ClientApiError {
#[error("Unauthorized")]
Unauthorized,
#[error(transparent)]
Other(#[from] anyhow::Error),
}
pub struct CloudApiClient {
credentials: RwLock<Option<Credentials>>,
http_client: Arc<HttpClientWithUrl>,
}
impl CloudApiClient {
pub fn new(http_client: Arc<HttpClientWithUrl>) -> Self {
Self {
credentials: RwLock::new(None),
http_client,
}
}
pub fn has_credentials(&self) -> bool {
self.credentials.read().is_some()
}
pub fn set_credentials(&self, user_id: u32, access_token: String) {
*self.credentials.write() = Some(Credentials {
user_id,
access_token,
});
}
pub fn clear_credentials(&self) {
*self.credentials.write() = None;
}
fn build_request(
&self,
req: request::Builder,
body: impl Into<AsyncBody>,
) -> Result<Request<AsyncBody>> {
let credentials = self.credentials.read();
let credentials = credentials.as_ref().context("no credentials provided")?;
build_request(req, body, credentials)
}
pub async fn get_authenticated_user(
&self,
) -> Result<GetAuthenticatedUserResponse, ClientApiError> {
let request = self.build_request(
Request::builder().method(Method::GET).uri(
self.http_client
.build_zed_cloud_url("/client/users/me")?
.as_ref(),
),
AsyncBody::default(),
)?;
let mut response = self.http_client.send(request).await?;
if !response.status().is_success() {
if response.status() == StatusCode::UNAUTHORIZED {
return Err(ClientApiError::Unauthorized);
}
let mut body = String::new();
response
.body_mut()
.read_to_string(&mut body)
.await
.context("failed to read response body")?;
return Err(ClientApiError::Other(anyhow::anyhow!(
"Failed to get authenticated user.\nStatus: {:?}\nBody: {body}",
response.status()
)));
}
let mut body = String::new();
response
.body_mut()
.read_to_string(&mut body)
.await
.context("failed to read response body")?;
Ok(serde_json::from_str(&body).context("failed to parse response body")?)
}
pub fn connect(&self, cx: &App) -> Result<Task<Result<Connection>>> {
let mut connect_url = self
.http_client
.build_zed_cloud_url("/client/users/connect")?;
connect_url
.set_scheme(match connect_url.scheme() {
"https" => "wss",
"http" => "ws",
scheme => Err(anyhow!("invalid URL scheme: {scheme}"))?,
})
.map_err(|_| anyhow!("failed to set URL scheme"))?;
let credentials = self.credentials.read();
let credentials = credentials.as_ref().context("no credentials provided")?;
let authorization_header = format!("{} {}", credentials.user_id, credentials.access_token);
Ok(Tokio::spawn_result(cx, async move {
let ws = WebSocket::connect(connect_url)
.with_request(
request::Builder::new()
.header("Authorization", authorization_header)
.header(PROTOCOL_VERSION_HEADER_NAME, PROTOCOL_VERSION.to_string()),
)
.await?;
Ok(Connection::new(ws))
}))
}
pub async fn create_llm_token(
&self,
system_id: Option<String>,
organization_id: Option<OrganizationId>,
) -> Result<CreateLlmTokenResponse, ClientApiError> {
let request_builder = Request::builder()
.method(Method::POST)
.uri(
self.http_client
.build_zed_cloud_url("/client/llm_tokens")?
.as_ref(),
)
.when_some(system_id, |builder, system_id| {
builder.header(ZED_SYSTEM_ID_HEADER_NAME, system_id)
});
let request = self.build_request(
request_builder,
Json(CreateLlmTokenBody { organization_id }),
)?;
let mut response = self.http_client.send(request).await?;
if !response.status().is_success() {
if response.status() == StatusCode::UNAUTHORIZED {
return Err(ClientApiError::Unauthorized);
}
let mut body = String::new();
response
.body_mut()
.read_to_string(&mut body)
.await
.context("failed to read response body")?;
return Err(ClientApiError::Other(anyhow::anyhow!(
"Failed to create LLM token.\nStatus: {:?}\nBody: {body}",
response.status()
)));
}
let mut body = String::new();
response
.body_mut()
.read_to_string(&mut body)
.await
.context("failed to read response body")?;
Ok(serde_json::from_str(&body).context("failed to parse response body")?)
}
pub async fn validate_credentials(&self, user_id: u32, access_token: &str) -> Result<bool> {
let request = build_request(
Request::builder().method(Method::GET).uri(
self.http_client
.build_zed_cloud_url("/client/users/me")?
.as_ref(),
),
AsyncBody::default(),
&Credentials {
user_id,
access_token: access_token.into(),
},
)?;
let mut response = self.http_client.send(request).await?;
if response.status().is_success() {
Ok(true)
} else {
let mut body = String::new();
response.body_mut().read_to_string(&mut body).await?;
if response.status() == StatusCode::UNAUTHORIZED {
Ok(false)
} else {
Err(anyhow!(
"Failed to get authenticated user.\nStatus: {:?}\nBody: {body}",
response.status()
))
}
}
}
pub async fn submit_agent_feedback(&self, body: SubmitAgentThreadFeedbackBody) -> Result<()> {
let request = self.build_request(
Request::builder().method(Method::POST).uri(
self.http_client
.build_zed_cloud_url("/client/feedback/agent_thread")?
.as_ref(),
),
AsyncBody::from(serde_json::to_string(&body)?),
)?;
let mut response = self.http_client.send(request).await?;
if !response.status().is_success() {
let mut body = String::new();
response.body_mut().read_to_string(&mut body).await?;
anyhow::bail!(
"Failed to submit agent feedback.\nStatus: {:?}\nBody: {body}",
response.status()
)
}
Ok(())
}
pub async fn submit_agent_feedback_comments(
&self,
body: SubmitAgentThreadFeedbackCommentsBody,
) -> Result<()> {
let request = self.build_request(
Request::builder().method(Method::POST).uri(
self.http_client
.build_zed_cloud_url("/client/feedback/agent_thread_comments")?
.as_ref(),
),
AsyncBody::from(serde_json::to_string(&body)?),
)?;
let mut response = self.http_client.send(request).await?;
if !response.status().is_success() {
let mut body = String::new();
response.body_mut().read_to_string(&mut body).await?;
anyhow::bail!(
"Failed to submit agent feedback comments.\nStatus: {:?}\nBody: {body}",
response.status()
)
}
Ok(())
}
pub async fn submit_edit_prediction_feedback(
&self,
body: SubmitEditPredictionFeedbackBody,
) -> Result<()> {
let request = self.build_request(
Request::builder().method(Method::POST).uri(
self.http_client
.build_zed_cloud_url("/client/feedback/edit_prediction")?
.as_ref(),
),
AsyncBody::from(serde_json::to_string(&body)?),
)?;
let mut response = self.http_client.send(request).await?;
if !response.status().is_success() {
let mut body = String::new();
response.body_mut().read_to_string(&mut body).await?;
anyhow::bail!(
"Failed to submit edit prediction feedback.\nStatus: {:?}\nBody: {body}",
response.status()
)
}
Ok(())
}
}
fn build_request(
req: request::Builder,
body: impl Into<AsyncBody>,
credentials: &Credentials,
) -> Result<Request<AsyncBody>> {
Ok(req
.header("Content-Type", "application/json")
.header(
"Authorization",
format!("{} {}", credentials.user_id, credentials.access_token),
)
.body(body.into())?)
}