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