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