tenant.openagents/omega

1use std::collections::BTreeMap;
2use std::sync::Arc;
3
4use anyhow::{Context as _, Result, anyhow};
5use cloud_api_client::{
6    AuthenticatedUser, GetAuthenticatedUserResponse, KnownOrUnknown, Plan, PlanInfo,
7};
8use cloud_llm_client::{CurrentUsage, UsageData, UsageLimit};
9use futures::{StreamExt, stream::BoxStream};
10use gpui::{AppContext as _, TestAppContext};
11use http_client::{AsyncBody, Method, Request, http};
12use parking_lot::Mutex;
13use rpc::{ConnectionId, Peer, Receipt, TypedEnvelope, proto};
14
15use crate::{Client, Connection, Credentials, EstablishConnectionError};
16
17pub struct FakeServer {
18    peer: Arc<Peer>,
19    state: Arc<Mutex<FakeServerState>>,
20}
21
22#[derive(Default)]
23struct FakeServerState {
24    incoming: Option<BoxStream<'static, Box<dyn proto::AnyTypedEnvelope>>>,
25    connection_id: Option<ConnectionId>,
26    forbid_connections: bool,
27    auth_count: usize,
28    access_token: usize,
29}
30
31impl FakeServer {
32    pub async fn for_client(
33        client_user_id: u64,
34        client: &Arc<Client>,
35        cx: &TestAppContext,
36    ) -> Self {
37        let server = Self {
38            peer: Peer::new(0),
39            state: Default::default(),
40        };
41
42        client.http_client().as_fake().replace_handler({
43            let state = server.state.clone();
44            move |old_handler, req| {
45                let state = state.clone();
46                let old_handler = old_handler.clone();
47                async move {
48                    match (req.method(), req.uri().path()) {
49                        (&Method::GET, "/client/users/me") => {
50                            let credentials = parse_authorization_header(&req);
51                            if credentials
52                                != Some(Credentials {
53                                    user_id: client_user_id,
54                                    access_token: state.lock().access_token.to_string(),
55                                })
56                            {
57                                return Ok(http_client::Response::builder()
58                                    .status(401)
59                                    .body("Unauthorized".into())
60                                    .unwrap());
61                            }
62
63                            Ok(http_client::Response::builder()
64                                .status(200)
65                                .body(
66                                    serde_json::to_string(&make_get_authenticated_user_response(
67                                        client_user_id as i32,
68                                        format!("user-{client_user_id}"),
69                                    ))
70                                    .unwrap()
71                                    .into(),
72                                )
73                                .unwrap())
74                        }
75                        _ => old_handler(req).await,
76                    }
77                }
78            }
79        });
80        client
81            .override_authenticate({
82                let state = Arc::downgrade(&server.state);
83                move |cx| {
84                    let state = state.clone();
85                    cx.spawn(async move |_| {
86                        let state = state.upgrade().context("server dropped")?;
87                        let mut state = state.lock();
88                        state.auth_count += 1;
89                        let access_token = state.access_token.to_string();
90                        Ok(Credentials {
91                            user_id: client_user_id,
92                            access_token,
93                        })
94                    })
95                }
96            })
97            .override_establish_connection({
98                let peer = Arc::downgrade(&server.peer);
99                let state = Arc::downgrade(&server.state);
100                move |credentials, cx| {
101                    let peer = peer.clone();
102                    let state = state.clone();
103                    let credentials = credentials.clone();
104                    cx.spawn(async move |cx| {
105                        let state = state.upgrade().context("server dropped")?;
106                        let peer = peer.upgrade().context("server dropped")?;
107                        if state.lock().forbid_connections {
108                            Err(EstablishConnectionError::Other(anyhow!(
109                                "server is forbidding connections"
110                            )))?
111                        }
112
113                        if credentials
114                            != (Credentials {
115                                user_id: client_user_id,
116                                access_token: state.lock().access_token.to_string(),
117                            })
118                        {
119                            Err(EstablishConnectionError::Unauthorized)?
120                        }
121
122                        let (client_conn, server_conn, _) =
123                            Connection::in_memory(cx.background_executor().clone());
124                        let (connection_id, io, incoming) =
125                            peer.add_test_connection(server_conn, cx.background_executor().clone());
126                        cx.background_spawn(io).detach();
127                        {
128                            let mut state = state.lock();
129                            state.connection_id = Some(connection_id);
130                            state.incoming = Some(incoming);
131                        }
132                        peer.send(
133                            connection_id,
134                            proto::Hello {
135                                peer_id: Some(connection_id.into()),
136                            },
137                        )
138                        .unwrap();
139
140                        Ok(client_conn)
141                    })
142                }
143            });
144
145        client
146            .connect(false, &cx.to_async())
147            .await
148            .into_response()
149            .unwrap();
150
151        server
152    }
153
154    pub fn disconnect(&self) {
155        if self.state.lock().connection_id.is_some() {
156            self.peer.disconnect(self.connection_id());
157            let mut state = self.state.lock();
158            state.connection_id.take();
159            state.incoming.take();
160        }
161    }
162
163    pub fn auth_count(&self) -> usize {
164        self.state.lock().auth_count
165    }
166
167    pub fn roll_access_token(&self) {
168        self.state.lock().access_token += 1;
169    }
170
171    pub fn forbid_connections(&self) {
172        self.state.lock().forbid_connections = true;
173    }
174
175    pub fn allow_connections(&self) {
176        self.state.lock().forbid_connections = false;
177    }
178
179    pub fn send<T: proto::EnvelopedMessage>(&self, message: T) {
180        self.peer.send(self.connection_id(), message).unwrap();
181    }
182
183    #[allow(clippy::await_holding_lock)]
184    pub async fn receive<M: proto::EnvelopedMessage>(&self) -> Result<TypedEnvelope<M>> {
185        let message = self
186            .state
187            .lock()
188            .incoming
189            .as_mut()
190            .expect("not connected")
191            .next()
192            .await
193            .context("other half hung up")?;
194        let type_name = message.payload_type_name();
195        let message = message.into_any();
196
197        if message.is::<TypedEnvelope<M>>() {
198            return Ok(*message.downcast().unwrap());
199        }
200
201        panic!(
202            "fake server received unexpected message type: {:?}",
203            type_name
204        );
205    }
206
207    pub fn respond<T: proto::RequestMessage>(&self, receipt: Receipt<T>, response: T::Response) {
208        self.peer.respond(receipt, response).unwrap()
209    }
210
211    fn connection_id(&self) -> ConnectionId {
212        self.state.lock().connection_id.expect("not connected")
213    }
214}
215
216impl Drop for FakeServer {
217    fn drop(&mut self) {
218        self.disconnect();
219    }
220}
221
222pub fn parse_authorization_header(req: &Request<AsyncBody>) -> Option<Credentials> {
223    let mut auth_header = req
224        .headers()
225        .get(http::header::AUTHORIZATION)?
226        .to_str()
227        .ok()?
228        .split_whitespace();
229    let user_id = auth_header.next()?.parse().ok()?;
230    let access_token = auth_header.next()?;
231    Some(Credentials {
232        user_id,
233        access_token: access_token.to_string(),
234    })
235}
236
237pub fn make_get_authenticated_user_response(
238    user_id: i32,
239    username: String,
240) -> GetAuthenticatedUserResponse {
241    GetAuthenticatedUserResponse {
242        user: AuthenticatedUser {
243            id_v2: format!("user_{user_id}"),
244            legacy_user_id: user_id,
245            metrics_id: format!("metrics-id-{user_id}"),
246            username: username.clone(),
247            avatar_url: "".to_string(),
248            github_login: username,
249            name: None,
250            is_staff: false,
251            accepted_tos_at: None,
252            has_connected_to_collab_once: false,
253        },
254        feature_flags: vec![],
255        organizations: vec![],
256        default_organization_id: None,
257        plans_by_organization: BTreeMap::new(),
258        configuration_by_organization: BTreeMap::new(),
259        plan: PlanInfo {
260            plan: KnownOrUnknown::Known(Plan::ZedPro),
261            subscription_period: None,
262            usage: CurrentUsage {
263                edit_predictions: UsageData {
264                    used: 250,
265                    limit: UsageLimit::Unlimited,
266                },
267            },
268            trial_started_at: None,
269            is_account_too_young: false,
270            has_overdue_invoices: false,
271        },
272    }
273}
274