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, Method, Request, StatusCode};
13use parking_lot::RwLock;
14use yawc::WebSocket;
15
16use crate::websocket::Connection;
17
18struct Credentials {
19 user_id: u32,
20 access_token: String,
21}
22
23pub struct CloudApiClient {
24 credentials: RwLock<Option<Credentials>>,
25 http_client: Arc<HttpClientWithUrl>,
26}
27
28impl CloudApiClient {
29 pub fn new(http_client: Arc<HttpClientWithUrl>) -> Self {
30 Self {
31 credentials: RwLock::new(None),
32 http_client,
33 }
34 }
35
36 pub fn has_credentials(&self) -> bool {
37 self.credentials.read().is_some()
38 }
39
40 pub fn set_credentials(&self, user_id: u32, access_token: String) {
41 *self.credentials.write() = Some(Credentials {
42 user_id,
43 access_token,
44 });
45 }
46
47 pub fn clear_credentials(&self) {
48 *self.credentials.write() = None;
49 }
50
51 fn build_request(
52 &self,
53 req: request::Builder,
54 body: impl Into<AsyncBody>,
55 ) -> Result<Request<AsyncBody>> {
56 let credentials = self.credentials.read();
57 let credentials = credentials.as_ref().context("no credentials provided")?;
58 build_request(req, body, credentials)
59 }
60
61 pub async fn get_authenticated_user(&self) -> Result<GetAuthenticatedUserResponse> {
62 let request = self.build_request(
63 Request::builder().method(Method::GET).uri(
64 self.http_client
65 .build_zed_cloud_url("/client/users/me", &[])?
66 .as_ref(),
67 ),
68 AsyncBody::default(),
69 )?;
70
71 let mut response = self.http_client.send(request).await?;
72
73 if !response.status().is_success() {
74 let mut body = String::new();
75 response.body_mut().read_to_string(&mut body).await?;
76
77 anyhow::bail!(
78 "Failed to get authenticated user.\nStatus: {:?}\nBody: {body}",
79 response.status()
80 )
81 }
82
83 let mut body = String::new();
84 response.body_mut().read_to_string(&mut body).await?;
85
86 Ok(serde_json::from_str(&body)?)
87 }
88
89 pub fn connect(&self, cx: &App) -> Result<Task<Result<Connection>>> {
90 let mut connect_url = self
91 .http_client
92 .build_zed_cloud_url("/client/users/connect", &[])?;
93 connect_url
94 .set_scheme(match connect_url.scheme() {
95 "https" => "wss",
96 "http" => "ws",
97 scheme => Err(anyhow!("invalid URL scheme: {scheme}"))?,
98 })
99 .map_err(|_| anyhow!("failed to set URL scheme"))?;
100
101 let credentials = self.credentials.read();
102 let credentials = credentials.as_ref().context("no credentials provided")?;
103 let authorization_header = format!("{} {}", credentials.user_id, credentials.access_token);
104
105 Ok(Tokio::spawn_result(cx, async move {
106 let ws = WebSocket::connect(connect_url)
107 .with_request(
108 request::Builder::new()
109 .header("Authorization", authorization_header)
110 .header(PROTOCOL_VERSION_HEADER_NAME, PROTOCOL_VERSION.to_string()),
111 )
112 .await?;
113
114 Ok(Connection::new(ws))
115 }))
116 }
117
118 pub async fn accept_terms_of_service(&self) -> Result<AcceptTermsOfServiceResponse> {
119 let request = self.build_request(
120 Request::builder().method(Method::POST).uri(
121 self.http_client
122 .build_zed_cloud_url("/client/terms_of_service/accept", &[])?
123 .as_ref(),
124 ),
125 AsyncBody::default(),
126 )?;
127
128 let mut response = self.http_client.send(request).await?;
129
130 if !response.status().is_success() {
131 let mut body = String::new();
132 response.body_mut().read_to_string(&mut body).await?;
133
134 anyhow::bail!(
135 "Failed to accept terms of service.\nStatus: {:?}\nBody: {body}",
136 response.status()
137 )
138 }
139
140 let mut body = String::new();
141 response.body_mut().read_to_string(&mut body).await?;
142
143 Ok(serde_json::from_str(&body)?)
144 }
145
146 pub async fn create_llm_token(
147 &self,
148 system_id: Option<String>,
149 ) -> Result<CreateLlmTokenResponse> {
150 let mut request_builder = Request::builder().method(Method::POST).uri(
151 self.http_client
152 .build_zed_cloud_url("/client/llm_tokens", &[])?
153 .as_ref(),
154 );
155
156 if let Some(system_id) = system_id {
157 request_builder = request_builder.header(ZED_SYSTEM_ID_HEADER_NAME, system_id);
158 }
159
160 let request = self.build_request(request_builder, AsyncBody::default())?;
161
162 let mut response = self.http_client.send(request).await?;
163
164 if !response.status().is_success() {
165 let mut body = String::new();
166 response.body_mut().read_to_string(&mut body).await?;
167
168 anyhow::bail!(
169 "Failed to create LLM token.\nStatus: {:?}\nBody: {body}",
170 response.status()
171 )
172 }
173
174 let mut body = String::new();
175 response.body_mut().read_to_string(&mut body).await?;
176
177 Ok(serde_json::from_str(&body)?)
178 }
179
180 pub async fn validate_credentials(&self, user_id: u32, access_token: &str) -> Result<bool> {
181 let request = build_request(
182 Request::builder().method(Method::GET).uri(
183 self.http_client
184 .build_zed_cloud_url("/client/users/me", &[])?
185 .as_ref(),
186 ),
187 AsyncBody::default(),
188 &Credentials {
189 user_id,
190 access_token: access_token.into(),
191 },
192 )?;
193
194 let mut response = self.http_client.send(request).await?;
195
196 if response.status().is_success() {
197 Ok(true)
198 } else {
199 let mut body = String::new();
200 response.body_mut().read_to_string(&mut body).await?;
201 if response.status() == StatusCode::UNAUTHORIZED {
202 Ok(false)
203 } else {
204 Err(anyhow!(
205 "Failed to get authenticated user.\nStatus: {:?}\nBody: {body}",
206 response.status()
207 ))
208 }
209 }
210 }
211}
212
213fn build_request(
214 req: request::Builder,
215 body: impl Into<AsyncBody>,
216 credentials: &Credentials,
217) -> Result<Request<AsyncBody>> {
218 Ok(req
219 .header("Content-Type", "application/json")
220 .header(
221 "Authorization",
222 format!("{} {}", credentials.user_id, credentials.access_token),
223 )
224 .body(body.into())?)
225}