cloud_api_client.rs

  1use std::sync::Arc;
  2
  3use anyhow::{Result, anyhow};
  4pub use cloud_api_types::*;
  5use futures::AsyncReadExt as _;
  6use http_client::http::request;
  7use http_client::{AsyncBody, HttpClientWithUrl, Method, Request};
  8use parking_lot::RwLock;
  9
 10struct Credentials {
 11    user_id: u32,
 12    access_token: String,
 13}
 14
 15pub struct CloudApiClient {
 16    credentials: RwLock<Option<Credentials>>,
 17    http_client: Arc<HttpClientWithUrl>,
 18}
 19
 20impl CloudApiClient {
 21    pub fn new(http_client: Arc<HttpClientWithUrl>) -> Self {
 22        Self {
 23            credentials: RwLock::new(None),
 24            http_client,
 25        }
 26    }
 27
 28    pub fn has_credentials(&self) -> bool {
 29        self.credentials.read().is_some()
 30    }
 31
 32    pub fn set_credentials(&self, user_id: u32, access_token: String) {
 33        *self.credentials.write() = Some(Credentials {
 34            user_id,
 35            access_token,
 36        });
 37    }
 38
 39    pub fn clear_credentials(&self) {
 40        *self.credentials.write() = None;
 41    }
 42
 43    fn authorization_header(&self) -> Result<String> {
 44        let guard = self.credentials.read();
 45        let credentials = guard
 46            .as_ref()
 47            .ok_or_else(|| anyhow!("No credentials provided"))?;
 48
 49        Ok(format!(
 50            "{} {}",
 51            credentials.user_id, credentials.access_token
 52        ))
 53    }
 54
 55    fn build_request(
 56        &self,
 57        req: request::Builder,
 58        body: impl Into<AsyncBody>,
 59    ) -> Result<Request<AsyncBody>> {
 60        Ok(req
 61            .header("Content-Type", "application/json")
 62            .header("Authorization", self.authorization_header()?)
 63            .body(body.into())?)
 64    }
 65
 66    pub async fn get_authenticated_user(&self) -> Result<GetAuthenticatedUserResponse> {
 67        let request = self.build_request(
 68            Request::builder().method(Method::GET).uri(
 69                self.http_client
 70                    .build_zed_cloud_url("/client/users/me", &[])?
 71                    .as_ref(),
 72            ),
 73            AsyncBody::default(),
 74        )?;
 75
 76        let mut response = self.http_client.send(request).await?;
 77
 78        if !response.status().is_success() {
 79            let mut body = String::new();
 80            response.body_mut().read_to_string(&mut body).await?;
 81
 82            anyhow::bail!(
 83                "Failed to get authenticated user.\nStatus: {:?}\nBody: {body}",
 84                response.status()
 85            )
 86        }
 87
 88        let mut body = String::new();
 89        response.body_mut().read_to_string(&mut body).await?;
 90
 91        Ok(serde_json::from_str(&body)?)
 92    }
 93
 94    pub async fn accept_terms_of_service(&self) -> Result<AcceptTermsOfServiceResponse> {
 95        let request = self.build_request(
 96            Request::builder().method(Method::POST).uri(
 97                self.http_client
 98                    .build_zed_cloud_url("/client/terms_of_service/accept", &[])?
 99                    .as_ref(),
100            ),
101            AsyncBody::default(),
102        )?;
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 accept terms of service.\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
122    pub async fn create_llm_token(
123        &self,
124        system_id: Option<String>,
125    ) -> Result<CreateLlmTokenResponse> {
126        let mut request_builder = Request::builder().method(Method::POST).uri(
127            self.http_client
128                .build_zed_cloud_url("/client/llm_tokens", &[])?
129                .as_ref(),
130        );
131
132        if let Some(system_id) = system_id {
133            request_builder = request_builder.header(ZED_SYSTEM_ID_HEADER_NAME, system_id);
134        }
135
136        let request = self.build_request(request_builder, AsyncBody::default())?;
137
138        let mut response = self.http_client.send(request).await?;
139
140        if !response.status().is_success() {
141            let mut body = String::new();
142            response.body_mut().read_to_string(&mut body).await?;
143
144            anyhow::bail!(
145                "Failed to create LLM token.\nStatus: {:?}\nBody: {body}",
146                response.status()
147            )
148        }
149
150        let mut body = String::new();
151        response.body_mut().read_to_string(&mut body).await?;
152
153        Ok(serde_json::from_str(&body)?)
154    }
155}