1use std::sync::Arc;
2
3use anyhow::{Context, Result, anyhow};
4pub use cloud_api_types::*;
5use futures::AsyncReadExt as _;
6use http_client::http::request;
7use http_client::{AsyncBody, HttpClientWithUrl, Method, Request, StatusCode};
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 build_request(
44 &self,
45 req: request::Builder,
46 body: impl Into<AsyncBody>,
47 ) -> Result<Request<AsyncBody>> {
48 let credentials = self.credentials.read();
49 let credentials = credentials.as_ref().context("no credentials provided")?;
50 build_request(req, body, credentials)
51 }
52
53 pub async fn get_authenticated_user(&self) -> Result<GetAuthenticatedUserResponse> {
54 let request = self.build_request(
55 Request::builder().method(Method::GET).uri(
56 self.http_client
57 .build_zed_cloud_url("/client/users/me", &[])?
58 .as_ref(),
59 ),
60 AsyncBody::default(),
61 )?;
62
63 let mut response = self.http_client.send(request).await?;
64
65 if !response.status().is_success() {
66 let mut body = String::new();
67 response.body_mut().read_to_string(&mut body).await?;
68
69 anyhow::bail!(
70 "Failed to get authenticated user.\nStatus: {:?}\nBody: {body}",
71 response.status()
72 )
73 }
74
75 let mut body = String::new();
76 response.body_mut().read_to_string(&mut body).await?;
77
78 Ok(serde_json::from_str(&body)?)
79 }
80
81 pub async fn accept_terms_of_service(&self) -> Result<AcceptTermsOfServiceResponse> {
82 let request = self.build_request(
83 Request::builder().method(Method::POST).uri(
84 self.http_client
85 .build_zed_cloud_url("/client/terms_of_service/accept", &[])?
86 .as_ref(),
87 ),
88 AsyncBody::default(),
89 )?;
90
91 let mut response = self.http_client.send(request).await?;
92
93 if !response.status().is_success() {
94 let mut body = String::new();
95 response.body_mut().read_to_string(&mut body).await?;
96
97 anyhow::bail!(
98 "Failed to accept terms of service.\nStatus: {:?}\nBody: {body}",
99 response.status()
100 )
101 }
102
103 let mut body = String::new();
104 response.body_mut().read_to_string(&mut body).await?;
105
106 Ok(serde_json::from_str(&body)?)
107 }
108
109 pub async fn create_llm_token(
110 &self,
111 system_id: Option<String>,
112 ) -> Result<CreateLlmTokenResponse> {
113 let mut request_builder = Request::builder().method(Method::POST).uri(
114 self.http_client
115 .build_zed_cloud_url("/client/llm_tokens", &[])?
116 .as_ref(),
117 );
118
119 if let Some(system_id) = system_id {
120 request_builder = request_builder.header(ZED_SYSTEM_ID_HEADER_NAME, system_id);
121 }
122
123 let request = self.build_request(request_builder, AsyncBody::default())?;
124
125 let mut response = self.http_client.send(request).await?;
126
127 if !response.status().is_success() {
128 let mut body = String::new();
129 response.body_mut().read_to_string(&mut body).await?;
130
131 anyhow::bail!(
132 "Failed to create LLM token.\nStatus: {:?}\nBody: {body}",
133 response.status()
134 )
135 }
136
137 let mut body = String::new();
138 response.body_mut().read_to_string(&mut body).await?;
139
140 Ok(serde_json::from_str(&body)?)
141 }
142
143 pub async fn validate_credentials(&self, user_id: u32, access_token: &str) -> Result<bool> {
144 let request = build_request(
145 Request::builder().method(Method::GET).uri(
146 self.http_client
147 .build_zed_cloud_url("/client/users/me", &[])?
148 .as_ref(),
149 ),
150 AsyncBody::default(),
151 &Credentials {
152 user_id,
153 access_token: access_token.into(),
154 },
155 )?;
156
157 let mut response = self.http_client.send(request).await?;
158
159 if response.status().is_success() {
160 Ok(true)
161 } else {
162 let mut body = String::new();
163 response.body_mut().read_to_string(&mut body).await?;
164 if response.status() == StatusCode::UNAUTHORIZED {
165 return Ok(false);
166 } else {
167 return Err(anyhow!(
168 "Failed to get authenticated user.\nStatus: {:?}\nBody: {body}",
169 response.status()
170 ));
171 }
172 }
173 }
174}
175
176fn build_request(
177 req: request::Builder,
178 body: impl Into<AsyncBody>,
179 credentials: &Credentials,
180) -> Result<Request<AsyncBody>> {
181 Ok(req
182 .header("Content-Type", "application/json")
183 .header(
184 "Authorization",
185 format!("{} {}", credentials.user_id, credentials.access_token),
186 )
187 .body(body.into())?)
188}