1use std::sync::Arc;
  2
  3use anyhow::{Result, anyhow};
  4pub use cloud_api_types::*;
  5use futures::AsyncReadExt as _;
  6use http_client::{AsyncBody, HttpClientWithUrl, Method, Request};
  7use parking_lot::RwLock;
  8
  9struct Credentials {
 10    user_id: u32,
 11    access_token: String,
 12}
 13
 14pub struct CloudApiClient {
 15    credentials: RwLock<Option<Credentials>>,
 16    http_client: Arc<HttpClientWithUrl>,
 17}
 18
 19impl CloudApiClient {
 20    pub fn new(http_client: Arc<HttpClientWithUrl>) -> Self {
 21        Self {
 22            credentials: RwLock::new(None),
 23            http_client,
 24        }
 25    }
 26
 27    pub fn has_credentials(&self) -> bool {
 28        self.credentials.read().is_some()
 29    }
 30
 31    pub fn set_credentials(&self, user_id: u32, access_token: String) {
 32        *self.credentials.write() = Some(Credentials {
 33            user_id,
 34            access_token,
 35        });
 36    }
 37
 38    pub fn clear_credentials(&self) {
 39        *self.credentials.write() = None;
 40    }
 41
 42    fn authorization_header(&self) -> Result<String> {
 43        let guard = self.credentials.read();
 44        let credentials = guard
 45            .as_ref()
 46            .ok_or_else(|| anyhow!("No credentials provided"))?;
 47
 48        Ok(format!(
 49            "{} {}",
 50            credentials.user_id, credentials.access_token
 51        ))
 52    }
 53
 54    pub async fn get_authenticated_user(&self) -> Result<GetAuthenticatedUserResponse> {
 55        let request = Request::builder()
 56            .method(Method::GET)
 57            .uri(
 58                self.http_client
 59                    .build_zed_cloud_url("/client/users/me", &[])?
 60                    .as_ref(),
 61            )
 62            .header("Content-Type", "application/json")
 63            .header("Authorization", self.authorization_header()?)
 64            .body(AsyncBody::default())?;
 65
 66        let mut response = self.http_client.send(request).await?;
 67
 68        if !response.status().is_success() {
 69            let mut body = String::new();
 70            response.body_mut().read_to_string(&mut body).await?;
 71
 72            anyhow::bail!(
 73                "Failed to get authenticated user.\nStatus: {:?}\nBody: {body}",
 74                response.status()
 75            )
 76        }
 77
 78        let mut body = String::new();
 79        response.body_mut().read_to_string(&mut body).await?;
 80
 81        Ok(serde_json::from_str(&body)?)
 82    }
 83
 84    pub async fn create_llm_token(
 85        &self,
 86        system_id: Option<String>,
 87    ) -> Result<CreateLlmTokenResponse> {
 88        let mut request_builder = Request::builder()
 89            .method(Method::POST)
 90            .uri(
 91                self.http_client
 92                    .build_zed_cloud_url("/client/llm_tokens", &[])?
 93                    .as_ref(),
 94            )
 95            .header("Content-Type", "application/json")
 96            .header("Authorization", self.authorization_header()?);
 97
 98        if let Some(system_id) = system_id {
 99            request_builder = request_builder.header(ZED_SYSTEM_ID_HEADER_NAME, system_id);
100        }
101
102        let request = request_builder.body(AsyncBody::default())?;
103
104        let mut response = self.http_client.send(request).await?;
105
106        if !response.status().is_success() {
107            let mut body = String::new();
108            response.body_mut().read_to_string(&mut body).await?;
109
110            anyhow::bail!(
111                "Failed to create LLM token.\nStatus: {:?}\nBody: {body}",
112                response.status()
113            )
114        }
115
116        let mut body = String::new();
117        response.body_mut().read_to_string(&mut body).await?;
118
119        Ok(serde_json::from_str(&body)?)
120    }
121}