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