zed/crates/cloud_api_client/src/cloud_api_client.rs

mod llm_token;
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;

pub use llm_token::LlmApiToken;

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())?)
}