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 create_llm_token(
119 &self,
120 system_id: Option<String>,
121 ) -> Result<CreateLlmTokenResponse> {
122 let mut request_builder = Request::builder().method(Method::POST).uri(
123 self.http_client
124 .build_zed_cloud_url("/client/llm_tokens", &[])?
125 .as_ref(),
126 );
127
128 if let Some(system_id) = system_id {
129 request_builder = request_builder.header(ZED_SYSTEM_ID_HEADER_NAME, system_id);
130 }
131
132 let request = self.build_request(request_builder, AsyncBody::default())?;
133
134 let mut response = self.http_client.send(request).await?;
135
136 if !response.status().is_success() {
137 let mut body = String::new();
138 response.body_mut().read_to_string(&mut body).await?;
139
140 anyhow::bail!(
141 "Failed to create LLM token.\nStatus: {:?}\nBody: {body}",
142 response.status()
143 )
144 }
145
146 let mut body = String::new();
147 response.body_mut().read_to_string(&mut body).await?;
148
149 Ok(serde_json::from_str(&body)?)
150 }
151
152 pub async fn validate_credentials(&self, user_id: u32, access_token: &str) -> Result<bool> {
153 let request = build_request(
154 Request::builder().method(Method::GET).uri(
155 self.http_client
156 .build_zed_cloud_url("/client/users/me", &[])?
157 .as_ref(),
158 ),
159 AsyncBody::default(),
160 &Credentials {
161 user_id,
162 access_token: access_token.into(),
163 },
164 )?;
165
166 let mut response = self.http_client.send(request).await?;
167
168 if response.status().is_success() {
169 Ok(true)
170 } else {
171 let mut body = String::new();
172 response.body_mut().read_to_string(&mut body).await?;
173 if response.status() == StatusCode::UNAUTHORIZED {
174 Ok(false)
175 } else {
176 Err(anyhow!(
177 "Failed to get authenticated user.\nStatus: {:?}\nBody: {body}",
178 response.status()
179 ))
180 }
181 }
182 }
183}
184
185fn build_request(
186 req: request::Builder,
187 body: impl Into<AsyncBody>,
188 credentials: &Credentials,
189) -> Result<Request<AsyncBody>> {
190 Ok(req
191 .header("Content-Type", "application/json")
192 .header(
193 "Authorization",
194 format!("{} {}", credentials.user_id, credentials.access_token),
195 )
196 .body(body.into())?)
197}