tenant.openagents/omega

1use std::time::Duration;
2
3use anyhow::{Context as _, Result, anyhow};
4use cloud_api_types::websocket_protocol::{PROTOCOL_VERSION, PROTOCOL_VERSION_HEADER_NAME};
5use futures::channel::mpsc::unbounded;
6use futures::stream::{SplitSink, SplitStream};
7use futures::{FutureExt as _, SinkExt as _, StreamExt as _, TryStreamExt as _};
8use gpui::{App, Task};
9use http_client::http::request;
10use yawc::frame::Frame;
11use yawc::{TcpWebSocket, WebSocket};
12
13use super::{MessageStream, forward_frame};
14use crate::CloudApiClient;
15
16const KEEPALIVE_INTERVAL: Duration = Duration::from_secs(1);
17
18pub struct Connection {
19    tx: SplitSink<TcpWebSocket, Frame>,
20    rx: SplitStream<TcpWebSocket>,
21}
22
23impl Connection {
24    fn new(websocket: TcpWebSocket) -> Self {
25        let (tx, rx) = websocket.split();
26        Self { tx, rx }
27    }
28
29    pub fn spawn(self, cx: &App) -> (MessageStream, Task<()>) {
30        let mut tx = self.tx;
31        let rx = self.rx.fuse();
32        let (message_tx, message_rx) = unbounded();
33        let executor = cx.background_executor().clone();
34        let task = cx.spawn(async move |_cx| {
35            let keepalive_timer = executor.timer(KEEPALIVE_INTERVAL).fuse();
36            futures::pin_mut!(keepalive_timer, rx);
37
38            loop {
39                futures::select_biased! {
40                    _ = keepalive_timer => {
41                        if tx.send(Frame::ping(Vec::new())).await.is_err() {
42                            break;
43                        }
44                        keepalive_timer.set(executor.timer(KEEPALIVE_INTERVAL).fuse());
45                    }
46                    frame = rx.next() => {
47                        let Some(frame) = frame else {
48                            break;
49                        };
50                        if !forward_frame(frame, &message_tx) {
51                            break;
52                        }
53                    }
54                }
55            }
56        });
57
58        (message_rx.into_stream().boxed(), task)
59    }
60}
61
62impl CloudApiClient {
63    pub fn connect(self: &std::sync::Arc<Self>, cx: &App) -> Result<Task<Result<Connection>>> {
64        let mut connect_url = self
65            .http_client
66            .build_zed_cloud_url("/client/users/connect")?;
67        connect_url
68            .set_scheme(match connect_url.scheme() {
69                "https" => "wss",
70                "http" => "ws",
71                scheme => Err(anyhow!("invalid URL scheme: {scheme}"))?,
72            })
73            .map_err(|_| anyhow!("failed to set URL scheme"))?;
74
75        let credentials = self.credentials.read();
76        let credentials = credentials.as_ref().context("no credentials provided")?;
77        let authorization_header = format!("{} {}", credentials.user_id, credentials.access_token);
78
79        Ok(gpui_tokio::Tokio::spawn_result(cx, async move {
80            let websocket = WebSocket::connect(connect_url)
81                .with_request(
82                    request::Builder::new()
83                        .header("Authorization", authorization_header)
84                        .header(PROTOCOL_VERSION_HEADER_NAME, PROTOCOL_VERSION.to_string()),
85                )
86                .await?;
87
88            Ok(Connection::new(websocket))
89        }))
90    }
91}
92