tenant.openagents/omega

1mod llm_token;
2mod websocket;
3
4use std::sync::Arc;
5
6use anyhow::{Result, anyhow};
7pub use cloud_api_types::*;
8use futures::AsyncReadExt as _;
9use http_client::http::request;
10use http_client::{
11    AsyncBody, HttpClientWithUrl, HttpRequestExt, Json, Method, Request, Response, StatusCode,
12};
13use parking_lot::RwLock;
14use serde::de::DeserializeOwned;
15use thiserror::Error;
16
17pub use llm_token::LlmApiToken;
18
19struct Credentials {
20    user_id: u32,
21    access_token: String,
22}
23
24#[derive(Clone, Copy)]
25enum Authentication {
26    Credentials,
27    Session,
28}
29
30#[derive(Debug, Error)]
31pub enum ClientApiError {
32    /// 401 — credentials are invalid or expired.
33    #[error("Unauthorized")]
34    Unauthorized,
35    /// No credentials have been set on the client.
36    #[error("not signed in")]
37    NotSignedIn,
38    /// Connection-level failure: DNS, TCP, TLS, timeout, etc.
39    /// The HTTP request never received a response.
40    #[error("connection to {host} failed")]
41    ConnectionFailed {
42        host: String,
43        #[source]
44        source: anyhow::Error,
45    },
46    /// Server returned a non-success HTTP status (other than 401).
47    #[error("{host} returned {status}")]
48    ServerError {
49        host: String,
50        status: StatusCode,
51        body: String,
52    },
53    /// Failed to read or parse the response body after a successful HTTP status.
54    #[error("invalid response")]
55    InvalidResponse(#[source] anyhow::Error),
56    /// Failed to build the HTTP request (URL construction, serialization, etc.).
57    /// This typically indicates a programming error.
58    #[error("failed to build request")]
59    RequestBuildFailed(#[source] anyhow::Error),
60}
61
62pub struct CloudApiClient {
63    credentials: RwLock<Option<Credentials>>,
64    http_client: Arc<HttpClientWithUrl>,
65    authentication: Authentication,
66}
67
68impl CloudApiClient {
69    pub fn new(http_client: Arc<HttpClientWithUrl>) -> Self {
70        Self {
71            credentials: RwLock::new(None),
72            http_client,
73            authentication: Authentication::Credentials,
74        }
75    }
76
77    pub fn new_with_session_authentication(http_client: Arc<HttpClientWithUrl>) -> Self {
78        Self {
79            credentials: RwLock::new(None),
80            http_client,
81            authentication: Authentication::Session,
82        }
83    }
84
85    pub fn has_credentials(&self) -> bool {
86        self.credentials.read().is_some()
87    }
88
89    pub fn set_credentials(&self, user_id: u32, access_token: String) {
90        *self.credentials.write() = Some(Credentials {
91            user_id,
92            access_token,
93        });
94    }
95
96    pub fn clear_credentials(&self) {
97        *self.credentials.write() = None;
98    }
99
100    pub fn cloud_host(&self) -> String {
101        self.http_client
102            .build_zed_cloud_url("/")
103            .ok()
104            .and_then(|url| url.host_str().map(String::from))
105            .unwrap_or_else(|| "cloud.zed.dev".into())
106    }
107
108    pub async fn get_authenticated_user(
109        &self,
110        system_id: Option<String>,
111    ) -> Result<GetAuthenticatedUserResponse, ClientApiError> {
112        let request_builder = Request::builder()
113            .method(Method::GET)
114            .uri(
115                self.http_client
116                    .build_zed_cloud_url("/client/users/me")
117                    .map_err(ClientApiError::RequestBuildFailed)?
118                    .as_ref(),
119            )
120            .when_some(system_id, |builder, system_id| {
121                builder.header(ZED_SYSTEM_ID_HEADER_NAME, system_id)
122            });
123
124        self.send_authenticated_json_request(request_builder, AsyncBody::default())
125            .await
126    }
127
128    async fn create_llm_token(
129        &self,
130        system_id: Option<String>,
131        organization_id: OrganizationId,
132    ) -> Result<CreateLlmTokenResponse, ClientApiError> {
133        let request_builder = Request::builder()
134            .method(Method::POST)
135            .uri(
136                self.http_client
137                    .build_zed_cloud_url("/client/llm_tokens")
138                    .map_err(ClientApiError::RequestBuildFailed)?
139                    .as_ref(),
140            )
141            .when_some(system_id, |builder, system_id| {
142                builder.header(ZED_SYSTEM_ID_HEADER_NAME, system_id)
143            });
144
145        self.send_authenticated_json_request(
146            request_builder,
147            Json(CreateLlmTokenBody { organization_id }),
148        )
149        .await
150    }
151
152    pub async fn update_system_settings(
153        &self,
154        system_id: String,
155        body: UpdateSystemSettingsBody,
156    ) -> Result<SystemSettings, ClientApiError> {
157        let request_builder = Request::builder()
158            .method(Method::PATCH)
159            .uri(
160                self.http_client
161                    .build_zed_cloud_url("/client/system_settings")
162                    .map_err(ClientApiError::RequestBuildFailed)?
163                    .as_ref(),
164            )
165            .header(ZED_SYSTEM_ID_HEADER_NAME, system_id);
166
167        self.send_authenticated_json_request(request_builder, Json(body))
168            .await
169    }
170
171    pub async fn send_authenticated_json_request<T: DeserializeOwned>(
172        &self,
173        request_builder: request::Builder,
174        body: impl Into<AsyncBody>,
175    ) -> Result<T, ClientApiError> {
176        let mut response = self
177            .send_authenticated_request(request_builder, body)
178            .await?;
179        Self::read_response_json(&mut response).await
180    }
181
182    async fn send_authenticated_request(
183        &self,
184        request_builder: request::Builder,
185        body: impl Into<AsyncBody>,
186    ) -> Result<Response<AsyncBody>, ClientApiError> {
187        let request = if matches!(self.authentication, Authentication::Session) {
188            build_request(request_builder, body, None)
189                .map_err(ClientApiError::RequestBuildFailed)?
190        } else {
191            let credentials = self.credentials.read();
192            let credentials = credentials.as_ref().ok_or(ClientApiError::NotSignedIn)?;
193            build_request(request_builder, body, Some(credentials))
194                .map_err(ClientApiError::RequestBuildFailed)?
195        };
196
197        let host = self.cloud_host();
198        let mut response = self.http_client.send(request).await.map_err(|source| {
199            ClientApiError::ConnectionFailed {
200                host: host.clone(),
201                source,
202            }
203        })?;
204
205        let status = response.status();
206        if status.is_success() {
207            return Ok(response);
208        }
209
210        if status == StatusCode::UNAUTHORIZED {
211            return Err(ClientApiError::Unauthorized);
212        }
213
214        let body = match Self::read_response_body(&mut response).await {
215            Ok(body) => body,
216            Err(error) => format!("failed to read response body: {error}"),
217        };
218        Err(ClientApiError::ServerError { host, status, body })
219    }
220
221    async fn read_response_json<T: DeserializeOwned>(
222        response: &mut Response<AsyncBody>,
223    ) -> Result<T, ClientApiError> {
224        let body = Self::read_response_body(response).await?;
225        serde_json::from_str(&body).map_err(|error| ClientApiError::InvalidResponse(error.into()))
226    }
227
228    async fn read_response_body(
229        response: &mut Response<AsyncBody>,
230    ) -> Result<String, ClientApiError> {
231        let mut body = String::new();
232        response
233            .body_mut()
234            .read_to_string(&mut body)
235            .await
236            .map_err(|error| ClientApiError::InvalidResponse(error.into()))?;
237        Ok(body)
238    }
239
240    pub async fn validate_credentials(&self, user_id: u32, access_token: &str) -> Result<bool> {
241        let request = build_request(
242            Request::builder().method(Method::GET).uri(
243                self.http_client
244                    .build_zed_cloud_url("/client/users/me")?
245                    .as_ref(),
246            ),
247            AsyncBody::default(),
248            Some(&Credentials {
249                user_id,
250                access_token: access_token.into(),
251            }),
252        )?;
253
254        let mut response = self.http_client.send(request).await?;
255
256        if response.status().is_success() {
257            Ok(true)
258        } else {
259            let mut body = String::new();
260            response.body_mut().read_to_string(&mut body).await?;
261            if response.status() == StatusCode::UNAUTHORIZED {
262                Ok(false)
263            } else {
264                Err(anyhow!(
265                    "Failed to get authenticated user.\nStatus: {:?}\nBody: {body}",
266                    response.status()
267                ))
268            }
269        }
270    }
271
272    pub async fn submit_agent_feedback(&self, body: SubmitAgentThreadFeedbackBody) -> Result<()> {
273        let request = Request::builder().method(Method::POST).uri(
274            self.http_client
275                .build_zed_cloud_url("/client/feedback/agent_thread")?
276                .as_ref(),
277        );
278
279        self.send_authenticated_request(request, AsyncBody::from(serde_json::to_string(&body)?))
280            .await?;
281        Ok(())
282    }
283
284    pub async fn submit_agent_feedback_comments(
285        &self,
286        body: SubmitAgentThreadFeedbackCommentsBody,
287    ) -> Result<()> {
288        let request = Request::builder().method(Method::POST).uri(
289            self.http_client
290                .build_zed_cloud_url("/client/feedback/agent_thread_comments")?
291                .as_ref(),
292        );
293
294        self.send_authenticated_request(request, AsyncBody::from(serde_json::to_string(&body)?))
295            .await?;
296        Ok(())
297    }
298
299    pub async fn submit_edit_prediction_feedback(
300        &self,
301        body: SubmitEditPredictionFeedbackBody,
302    ) -> Result<()> {
303        let request = Request::builder().method(Method::POST).uri(
304            self.http_client
305                .build_zed_cloud_url("/client/feedback/edit_prediction")?
306                .as_ref(),
307        );
308
309        self.send_authenticated_request(request, AsyncBody::from(serde_json::to_string(&body)?))
310            .await?;
311        Ok(())
312    }
313}
314
315fn build_request(
316    req: request::Builder,
317    body: impl Into<AsyncBody>,
318    credentials: Option<&Credentials>,
319) -> Result<Request<AsyncBody>> {
320    Ok(req
321        .header("Content-Type", "application/json")
322        .when_some(credentials, |request, credentials| {
323            request.header(
324                "Authorization",
325                format!("{} {}", credentials.user_id, credentials.access_token),
326            )
327        })
328        .body(body.into())?)
329}
330
331#[cfg(test)]
332mod tests {
333    use super::*;
334
335    #[test]
336    fn build_session_authenticated_request_without_authorization_header() -> Result<()> {
337        let request = build_request(
338            Request::builder().uri("https://cloud.zed.dev/client/users/me"),
339            AsyncBody::default(),
340            None,
341        )?;
342
343        assert_eq!(
344            request
345                .headers()
346                .get("Content-Type")
347                .and_then(|value| value.to_str().ok()),
348            Some("application/json")
349        );
350        assert!(!request.headers().contains_key("Authorization"));
351        Ok(())
352    }
353
354    #[test]
355    fn build_credentials_authenticated_request_with_authorization_header() -> Result<()> {
356        let request = build_request(
357            Request::builder().uri("https://cloud.zed.dev/client/users/me"),
358            AsyncBody::default(),
359            Some(&Credentials {
360                user_id: 123,
361                access_token: "token".into(),
362            }),
363        )?;
364
365        assert_eq!(
366            request
367                .headers()
368                .get("Authorization")
369                .and_then(|value| value.to_str().ok()),
370            Some("123 token")
371        );
372        Ok(())
373    }
374}
375