tenant.openagents/omega

1use std::time::Duration;
2
3use anyhow::{Result, anyhow};
4
5use futures::channel::mpsc::unbounded;
6use futures::stream::SplitStream;
7use futures::{FutureExt as _, StreamExt as _, TryStreamExt as _};
8use gpui::{App, Task};
9use yawc::WebSocket;
10
11use super::{MessageStream, forward_frame};
12use crate::CloudApiClient;
13
14const CONNECT_TIMEOUT: Duration = Duration::from_secs(30);
15
16pub struct Connection {
17    rx: SplitStream<WebSocket>,
18}
19
20impl Connection {
21    fn new(websocket: WebSocket) -> Self {
22        let (_, rx) = websocket.split();
23        Self { rx }
24    }
25
26    pub fn spawn(self, cx: &App) -> (MessageStream, Task<()>) {
27        let mut rx = self.rx;
28        let (message_tx, message_rx) = unbounded();
29        let task = cx.spawn(async move |_cx| {
30            while let Some(frame) = rx.next().await {
31                let frame = match frame {
32                    Ok(frame) => frame,
33                    Err(error) => {
34                        let error = anyhow!("Cloud WebSocket error: {error}");
35                        if message_tx.unbounded_send(Err(error)).is_err() {
36                            break;
37                        }
38                        continue;
39                    }
40                };
41                if !forward_frame(frame, &message_tx) {
42                    break;
43                }
44            }
45        });
46
47        (message_rx.into_stream().boxed(), task)
48    }
49}
50
51impl CloudApiClient {
52    pub fn connect(self: &std::sync::Arc<Self>, cx: &App) -> Result<Task<Result<Connection>>> {
53        let client = self.clone();
54        let executor = cx.background_executor().clone();
55        Ok(cx.spawn(async move |_cx| {
56            let mut connect_url = client
57                .http_client
58                .build_zed_cloud_url("/client/users/connect")?;
59            connect_url
60                .set_scheme(match connect_url.scheme() {
61                    "https" => "wss",
62                    "http" => "ws",
63                    scheme => return Err(anyhow!("invalid URL scheme: {scheme}")),
64                })
65                .map_err(|_| anyhow!("failed to set URL scheme"))?;
66
67
68            let connect = WebSocket::connect(connect_url).fuse();
69            let timeout = executor.timer(CONNECT_TIMEOUT).fuse();
70            futures::pin_mut!(connect, timeout);
71            let websocket = futures::select_biased! {
72                result = connect => result.map_err(|error| anyhow!("failed to connect to Cloud WebSocket: {error}"))?,
73                _ = timeout => return Err(anyhow!("timed out connecting to Cloud WebSocket")),
74            };
75
76            Ok(Connection::new(websocket))
77        }))
78    }
79}
80