1mod llm_token;
2mod websocket;
3
4use std::sync::Arc;
5
6use anyhow::{Context, Result, anyhow};
7use cloud_api_types::websocket_protocol::{PROTOCOL_VERSION, PROTOCOL_VERSION_HEADER_NAME};
8pub use cloud_api_types::*;
9use futures::AsyncReadExt as _;
10use gpui::{App, Task};
11use gpui_tokio::Tokio;
12use http_client::http::request;
13use http_client::{
14 AsyncBody, HttpClientWithUrl, HttpRequestExt, Json, Method, Request, StatusCode,
15};
16use parking_lot::RwLock;
17use thiserror::Error;
18use yawc::WebSocket;
19
20use crate::websocket::Connection;
21
22pub use llm_token::LlmApiToken;
23
24struct Credentials {
25 user_id: u32,
26 access_token: String,
27}
28
29#[derive(Debug, Error)]
30pub enum ClientApiError {
31 #[error("Unauthorized")]
32 Unauthorized,
33 #[error(transparent)]
34 Other(#[from] anyhow::Error),
35}
36
37pub struct CloudApiClient {
38 credentials: RwLock<Option<Credentials>>,
39 http_client: Arc<HttpClientWithUrl>,
40}
41
42impl CloudApiClient {
43 pub fn new(http_client: Arc<HttpClientWithUrl>) -> Self {
44 Self {
45 credentials: RwLock::new(None),
46 http_client,
47 }
48 }
49
50 pub fn has_credentials(&self) -> bool {
51 self.credentials.read().is_some()
52 }
53
54 pub fn set_credentials(&self, user_id: u32, access_token: String) {
55 *self.credentials.write() = Some(Credentials {
56 user_id,
57 access_token,
58 });
59 }
60
61 pub fn clear_credentials(&self) {
62 *self.credentials.write() = None;
63 }
64
65 fn build_request(
66 &self,
67 req: request::Builder,
68 body: impl Into<AsyncBody>,
69 ) -> Result<Request<AsyncBody>> {
70 let credentials = self.credentials.read();
71 let credentials = credentials.as_ref().context("no credentials provided")?;
72 build_request(req, body, credentials)
73 }
74
75 pub async fn get_authenticated_user(
76 &self,
77 ) -> Result<GetAuthenticatedUserResponse, ClientApiError> {
78 let request = self.build_request(
79 Request::builder().method(Method::GET).uri(
80 self.http_client
81 .build_zed_cloud_url("/client/users/me")?
82 .as_ref(),
83 ),
84 AsyncBody::default(),
85 )?;
86
87 let mut response = self.http_client.send(request).await?;
88
89 if !response.status().is_success() {
90 if response.status() == StatusCode::UNAUTHORIZED {
91 return Err(ClientApiError::Unauthorized);
92 }
93
94 let mut body = String::new();
95 response
96 .body_mut()
97 .read_to_string(&mut body)
98 .await
99 .context("failed to read response body")?;
100
101 return Err(ClientApiError::Other(anyhow::anyhow!(
102 "Failed to get authenticated user.\nStatus: {:?}\nBody: {body}",
103 response.status()
104 )));
105 }
106
107 let mut body = String::new();
108 response
109 .body_mut()
110 .read_to_string(&mut body)
111 .await
112 .context("failed to read response body")?;
113
114 Ok(serde_json::from_str(&body).context("failed to parse response body")?)
115 }
116
117 pub fn connect(&self, cx: &App) -> Result<Task<Result<Connection>>> {
118 let mut connect_url = self
119 .http_client
120 .build_zed_cloud_url("/client/users/connect")?;
121 connect_url
122 .set_scheme(match connect_url.scheme() {
123 "https" => "wss",
124 "http" => "ws",
125 scheme => Err(anyhow!("invalid URL scheme: {scheme}"))?,
126 })
127 .map_err(|_| anyhow!("failed to set URL scheme"))?;
128
129 let credentials = self.credentials.read();
130 let credentials = credentials.as_ref().context("no credentials provided")?;
131 let authorization_header = format!("{} {}", credentials.user_id, credentials.access_token);
132
133 Ok(Tokio::spawn_result(cx, async move {
134 let ws = WebSocket::connect(connect_url)
135 .with_request(
136 request::Builder::new()
137 .header("Authorization", authorization_header)
138 .header(PROTOCOL_VERSION_HEADER_NAME, PROTOCOL_VERSION.to_string()),
139 )
140 .await?;
141
142 Ok(Connection::new(ws))
143 }))
144 }
145
146 pub async fn create_llm_token(
147 &self,
148 system_id: Option<String>,
149 organization_id: Option<OrganizationId>,
150 ) -> Result<CreateLlmTokenResponse, ClientApiError> {
151 let request_builder = Request::builder()
152 .method(Method::POST)
153 .uri(
154 self.http_client
155 .build_zed_cloud_url("/client/llm_tokens")?
156 .as_ref(),
157 )
158 .when_some(system_id, |builder, system_id| {
159 builder.header(ZED_SYSTEM_ID_HEADER_NAME, system_id)
160 });
161
162 let request = self.build_request(
163 request_builder,
164 Json(CreateLlmTokenBody { organization_id }),
165 )?;
166
167 let mut response = self.http_client.send(request).await?;
168
169 if !response.status().is_success() {
170 if response.status() == StatusCode::UNAUTHORIZED {
171 return Err(ClientApiError::Unauthorized);
172 }
173
174 let mut body = String::new();
175 response
176 .body_mut()
177 .read_to_string(&mut body)
178 .await
179 .context("failed to read response body")?;
180
181 return Err(ClientApiError::Other(anyhow::anyhow!(
182 "Failed to create LLM token.\nStatus: {:?}\nBody: {body}",
183 response.status()
184 )));
185 }
186
187 let mut body = String::new();
188 response
189 .body_mut()
190 .read_to_string(&mut body)
191 .await
192 .context("failed to read response body")?;
193
194 Ok(serde_json::from_str(&body).context("failed to parse response body")?)
195 }
196
197 pub async fn validate_credentials(&self, user_id: u32, access_token: &str) -> Result<bool> {
198 let request = build_request(
199 Request::builder().method(Method::GET).uri(
200 self.http_client
201 .build_zed_cloud_url("/client/users/me")?
202 .as_ref(),
203 ),
204 AsyncBody::default(),
205 &Credentials {
206 user_id,
207 access_token: access_token.into(),
208 },
209 )?;
210
211 let mut response = self.http_client.send(request).await?;
212
213 if response.status().is_success() {
214 Ok(true)
215 } else {
216 let mut body = String::new();
217 response.body_mut().read_to_string(&mut body).await?;
218 if response.status() == StatusCode::UNAUTHORIZED {
219 Ok(false)
220 } else {
221 Err(anyhow!(
222 "Failed to get authenticated user.\nStatus: {:?}\nBody: {body}",
223 response.status()
224 ))
225 }
226 }
227 }
228
229 pub async fn submit_agent_feedback(&self, body: SubmitAgentThreadFeedbackBody) -> Result<()> {
230 let request = self.build_request(
231 Request::builder().method(Method::POST).uri(
232 self.http_client
233 .build_zed_cloud_url("/client/feedback/agent_thread")?
234 .as_ref(),
235 ),
236 AsyncBody::from(serde_json::to_string(&body)?),
237 )?;
238
239 let mut response = self.http_client.send(request).await?;
240
241 if !response.status().is_success() {
242 let mut body = String::new();
243 response.body_mut().read_to_string(&mut body).await?;
244
245 anyhow::bail!(
246 "Failed to submit agent feedback.\nStatus: {:?}\nBody: {body}",
247 response.status()
248 )
249 }
250
251 Ok(())
252 }
253
254 pub async fn submit_agent_feedback_comments(
255 &self,
256 body: SubmitAgentThreadFeedbackCommentsBody,
257 ) -> Result<()> {
258 let request = self.build_request(
259 Request::builder().method(Method::POST).uri(
260 self.http_client
261 .build_zed_cloud_url("/client/feedback/agent_thread_comments")?
262 .as_ref(),
263 ),
264 AsyncBody::from(serde_json::to_string(&body)?),
265 )?;
266
267 let mut response = self.http_client.send(request).await?;
268
269 if !response.status().is_success() {
270 let mut body = String::new();
271 response.body_mut().read_to_string(&mut body).await?;
272
273 anyhow::bail!(
274 "Failed to submit agent feedback comments.\nStatus: {:?}\nBody: {body}",
275 response.status()
276 )
277 }
278
279 Ok(())
280 }
281
282 pub async fn submit_edit_prediction_feedback(
283 &self,
284 body: SubmitEditPredictionFeedbackBody,
285 ) -> Result<()> {
286 let request = self.build_request(
287 Request::builder().method(Method::POST).uri(
288 self.http_client
289 .build_zed_cloud_url("/client/feedback/edit_prediction")?
290 .as_ref(),
291 ),
292 AsyncBody::from(serde_json::to_string(&body)?),
293 )?;
294
295 let mut response = self.http_client.send(request).await?;
296
297 if !response.status().is_success() {
298 let mut body = String::new();
299 response.body_mut().read_to_string(&mut body).await?;
300
301 anyhow::bail!(
302 "Failed to submit edit prediction feedback.\nStatus: {:?}\nBody: {body}",
303 response.status()
304 )
305 }
306
307 Ok(())
308 }
309}
310
311fn build_request(
312 req: request::Builder,
313 body: impl Into<AsyncBody>,
314 credentials: &Credentials,
315) -> Result<Request<AsyncBody>> {
316 Ok(req
317 .header("Content-Type", "application/json")
318 .header(
319 "Authorization",
320 format!("{} {}", credentials.user_id, credentials.access_token),
321 )
322 .body(body.into())?)
323}