tenant.openagents/omega

1use anyhow::{Context, Result};
2use collections::{HashMap, TypeIdHashMap};
3use futures::{
4    Future, FutureExt as _, Stream, StreamExt as _,
5    channel::oneshot,
6    future::{BoxFuture, LocalBoxFuture},
7    stream::BoxStream,
8};
9use gpui::{AnyEntity, AnyWeakEntity, AsyncApp, BackgroundExecutor, Entity, FutureExt as _};
10use parking_lot::Mutex;
11use proto::{
12    AnyTypedEnvelope, EntityMessage, Envelope, EnvelopedMessage, LspRequestId, LspRequestMessage,
13    RequestMessage, TypedEnvelope, error::ErrorExt as _,
14};
15use std::{
16    any::{Any, TypeId},
17    sync::{
18        Arc, OnceLock,
19        atomic::{self, AtomicU64},
20    },
21    time::Duration,
22};
23
24#[derive(Debug, Clone)]
25pub struct AnyProtoClient(Arc<State>);
26
27type RequestIds = Arc<
28    Mutex<
29        HashMap<
30            LspRequestId,
31            oneshot::Sender<
32                Result<
33                    Option<TypedEnvelope<Vec<proto::ProtoLspResponse<Box<dyn AnyTypedEnvelope>>>>>,
34                >,
35            >,
36        >,
37    >,
38>;
39
40static NEXT_LSP_REQUEST_ID: OnceLock<Arc<AtomicU64>> = OnceLock::new();
41static REQUEST_IDS: OnceLock<RequestIds> = OnceLock::new();
42
43struct State {
44    client: Arc<dyn ProtoClient>,
45    next_lsp_request_id: Arc<AtomicU64>,
46    request_ids: RequestIds,
47}
48
49impl std::fmt::Debug for State {
50    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
51        f.debug_struct("State")
52            .field("next_lsp_request_id", &self.next_lsp_request_id)
53            .field("request_ids", &self.request_ids)
54            .finish_non_exhaustive()
55    }
56}
57
58pub trait ProtoClient: Send + Sync {
59    fn request(
60        &self,
61        envelope: Envelope,
62        request_type: &'static str,
63    ) -> BoxFuture<'static, Result<Envelope>>;
64
65    fn request_stream(
66        &self,
67        envelope: Envelope,
68        request_type: &'static str,
69    ) -> BoxFuture<'static, Result<BoxStream<'static, Result<Envelope>>>> {
70        async move {
71            anyhow::bail!(
72                "stream requests are not supported for {request_type}: {:?}",
73                envelope.payload
74            )
75        }
76        .boxed()
77    }
78
79    fn send(&self, envelope: Envelope, message_type: &'static str) -> Result<()>;
80
81    fn send_response(&self, envelope: Envelope, message_type: &'static str) -> Result<()>;
82
83    fn message_handler_set(&self) -> &parking_lot::Mutex<ProtoMessageHandlerSet>;
84
85    fn is_via_collab(&self) -> bool;
86    fn has_wsl_interop(&self) -> bool;
87}
88
89#[derive(Default)]
90pub struct ProtoMessageHandlerSet {
91    pub entity_types_by_message_type: TypeIdHashMap<TypeId>,
92    pub entities_by_type_and_remote_id: HashMap<(TypeId, u64), EntityMessageSubscriber>,
93    pub entity_id_extractors: TypeIdHashMap<fn(&dyn AnyTypedEnvelope) -> u64>,
94    pub entities_by_message_type: TypeIdHashMap<AnyWeakEntity>,
95    pub message_handlers: TypeIdHashMap<ProtoMessageHandler>,
96}
97
98pub type ProtoMessageHandler = Arc<
99    dyn Send
100        + Sync
101        + Fn(
102            AnyEntity,
103            Box<dyn AnyTypedEnvelope>,
104            AnyProtoClient,
105            AsyncApp,
106        ) -> LocalBoxFuture<'static, Result<()>>,
107>;
108
109impl ProtoMessageHandlerSet {
110    pub fn clear(&mut self) {
111        self.message_handlers.clear();
112        self.entities_by_message_type.clear();
113        self.entities_by_type_and_remote_id.clear();
114        self.entity_id_extractors.clear();
115    }
116
117    fn add_message_handler(
118        &mut self,
119        message_type_id: TypeId,
120        entity: gpui::AnyWeakEntity,
121        handler: ProtoMessageHandler,
122    ) {
123        self.entities_by_message_type
124            .insert(message_type_id, entity);
125        let prev_handler = self.message_handlers.insert(message_type_id, handler);
126        if prev_handler.is_some() {
127            panic!("registered handler for the same message twice");
128        }
129    }
130
131    fn add_entity_message_handler(
132        &mut self,
133        message_type_id: TypeId,
134        entity_type_id: TypeId,
135        entity_id_extractor: fn(&dyn AnyTypedEnvelope) -> u64,
136        handler: ProtoMessageHandler,
137    ) {
138        self.entity_id_extractors
139            .entry(message_type_id)
140            .or_insert(entity_id_extractor);
141        self.entity_types_by_message_type
142            .insert(message_type_id, entity_type_id);
143        let prev_handler = self.message_handlers.insert(message_type_id, handler);
144        if prev_handler.is_some() {
145            panic!("registered handler for the same message twice");
146        }
147    }
148
149    pub fn handle_message(
150        this: &parking_lot::Mutex<Self>,
151        message: Box<dyn AnyTypedEnvelope>,
152        client: AnyProtoClient,
153        cx: AsyncApp,
154    ) -> Option<LocalBoxFuture<'static, Result<()>>> {
155        let payload_type_id = message.payload_type_id();
156        let mut this = this.lock();
157        let handler = this.message_handlers.get(&payload_type_id)?.clone();
158        let entity = if let Some(entity) = this.entities_by_message_type.get(&payload_type_id) {
159            entity.upgrade()?
160        } else {
161            let extract_entity_id = *this.entity_id_extractors.get(&payload_type_id)?;
162            let entity_type_id = *this.entity_types_by_message_type.get(&payload_type_id)?;
163            let entity_id = (extract_entity_id)(message.as_ref());
164            match this
165                .entities_by_type_and_remote_id
166                .get_mut(&(entity_type_id, entity_id))?
167            {
168                EntityMessageSubscriber::Pending(pending) => {
169                    pending.push(message);
170                    return None;
171                }
172                EntityMessageSubscriber::Entity { handle } => handle.upgrade()?,
173            }
174        };
175        drop(this);
176        Some(handler(entity, message, client, cx))
177    }
178}
179
180pub enum EntityMessageSubscriber {
181    Entity { handle: AnyWeakEntity },
182    Pending(Vec<Box<dyn AnyTypedEnvelope>>),
183}
184
185impl std::fmt::Debug for EntityMessageSubscriber {
186    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
187        match self {
188            EntityMessageSubscriber::Entity { handle } => f
189                .debug_struct("EntityMessageSubscriber::Entity")
190                .field("handle", handle)
191                .finish(),
192            EntityMessageSubscriber::Pending(vec) => f
193                .debug_struct("EntityMessageSubscriber::Pending")
194                .field(
195                    "envelopes",
196                    &vec.iter()
197                        .map(|envelope| envelope.payload_type_name())
198                        .collect::<Vec<_>>(),
199                )
200                .finish(),
201        }
202    }
203}
204
205impl<T> From<Arc<T>> for AnyProtoClient
206where
207    T: ProtoClient + 'static,
208{
209    fn from(client: Arc<T>) -> Self {
210        Self::new(client)
211    }
212}
213
214impl AnyProtoClient {
215    pub fn new<T: ProtoClient + 'static>(client: Arc<T>) -> Self {
216        Self(Arc::new(State {
217            client,
218            next_lsp_request_id: NEXT_LSP_REQUEST_ID
219                .get_or_init(|| Arc::new(AtomicU64::new(0)))
220                .clone(),
221            request_ids: REQUEST_IDS.get_or_init(RequestIds::default).clone(),
222        }))
223    }
224
225    pub fn is_via_collab(&self) -> bool {
226        self.0.client.is_via_collab()
227    }
228
229    pub fn request<T: RequestMessage>(
230        &self,
231        request: T,
232    ) -> impl Future<Output = Result<T::Response>> + use<T> {
233        let envelope = request.into_envelope(0, None, None);
234        let response = self.0.client.request(envelope, T::NAME);
235        async move {
236            T::Response::from_envelope(response.await?)
237                .context("received response of the wrong type")
238        }
239    }
240
241    pub fn request_stream<T: RequestMessage>(
242        &self,
243        request: T,
244    ) -> impl Future<Output = Result<BoxStream<'static, Result<T::Response>>>> + use<T> {
245        let envelope = request.into_envelope(0, None, None);
246        let response_stream = self.0.client.request_stream(envelope, T::NAME);
247        async move {
248            Ok(response_stream
249                .await?
250                .map(|response| {
251                    T::Response::from_envelope(response?)
252                        .context("received response of the wrong type")
253                })
254                .boxed())
255        }
256    }
257
258    pub fn send<T: EnvelopedMessage>(&self, request: T) -> Result<()> {
259        let envelope = request.into_envelope(0, None, None);
260        self.0.client.send(envelope, T::NAME)
261    }
262
263    pub fn send_response<T: EnvelopedMessage>(&self, request_id: u32, request: T) -> Result<()> {
264        let envelope = request.into_envelope(0, Some(request_id), None);
265        self.0.client.send(envelope, T::NAME)
266    }
267
268    pub fn request_lsp<T>(
269        &self,
270        project_id: u64,
271        server_id: Option<u64>,
272        timeout: Duration,
273        executor: BackgroundExecutor,
274        request: T,
275    ) -> impl Future<
276        Output = Result<Option<TypedEnvelope<Vec<proto::ProtoLspResponse<T::Response>>>>>,
277    > + use<T>
278    where
279        T: LspRequestMessage,
280    {
281        let new_id = LspRequestId(
282            self.0
283                .next_lsp_request_id
284                .fetch_add(1, atomic::Ordering::Acquire),
285        );
286        let (tx, rx) = oneshot::channel();
287        {
288            self.0.request_ids.lock().insert(new_id, tx);
289        }
290
291        let query = proto::LspQuery {
292            project_id,
293            server_id,
294            lsp_request_id: new_id.0,
295            request: Some(request.to_proto_query()),
296        };
297        let request = self.request(query);
298        let request_ids = self.0.request_ids.clone();
299        async move {
300            match request.await {
301                Ok(_request_enqueued) => {}
302                Err(e) => {
303                    request_ids.lock().remove(&new_id);
304                    return Err(e).context("sending LSP proto request");
305                }
306            }
307
308            let response = rx.with_timeout(timeout, &executor).await;
309            {
310                request_ids.lock().remove(&new_id);
311            }
312            match response {
313                Ok(Ok(response)) => {
314                    let response = response
315                        .context("waiting for LSP proto response")?
316                        .map(|response| {
317                            anyhow::Ok(TypedEnvelope {
318                                payload: response
319                                    .payload
320                                    .into_iter()
321                                    .map(|lsp_response| lsp_response.into_response::<T>())
322                                    .collect::<Result<Vec<_>>>()?,
323                                sender_id: response.sender_id,
324                                original_sender_id: response.original_sender_id,
325                                message_id: response.message_id,
326                                received_at: response.received_at,
327                            })
328                        })
329                        .transpose()
330                        .context("converting LSP proto response")?;
331                    Ok(response)
332                }
333                Err(_cancelled_due_timeout) => Ok(None),
334                Ok(Err(_channel_dropped)) => Ok(None),
335            }
336        }
337    }
338
339    pub fn send_lsp_response<T: LspRequestMessage>(
340        &self,
341        project_id: u64,
342        lsp_request_id: LspRequestId,
343        server_responses: HashMap<u64, T::Response>,
344    ) -> Result<()> {
345        self.send(proto::LspQueryResponse {
346            project_id,
347            lsp_request_id: lsp_request_id.0,
348            responses: server_responses
349                .into_iter()
350                .map(|(server_id, response)| proto::LspResponse {
351                    server_id,
352                    response: Some(T::response_to_proto_query(response)),
353                })
354                .collect(),
355        })
356    }
357
358    pub fn handle_lsp_response(&self, mut envelope: TypedEnvelope<proto::LspQueryResponse>) {
359        let request_id = LspRequestId(envelope.payload.lsp_request_id);
360        let mut response_senders = self.0.request_ids.lock();
361        if let Some(tx) = response_senders.remove(&request_id) {
362            let responses = envelope.payload.responses.drain(..).collect::<Vec<_>>();
363            tx.send(Ok(Some(proto::TypedEnvelope {
364                sender_id: envelope.sender_id,
365                original_sender_id: envelope.original_sender_id,
366                message_id: envelope.message_id,
367                received_at: envelope.received_at,
368                payload: responses
369                    .into_iter()
370                    .filter_map(|response| {
371                        use proto::lsp_response::Response;
372
373                        let server_id = response.server_id;
374                        let response = match response.response? {
375                            Response::GetReferencesResponse(response) => {
376                                to_any_envelope(&envelope, response)
377                            }
378                            Response::GetDocumentColorResponse(response) => {
379                                to_any_envelope(&envelope, response)
380                            }
381                            Response::GetHoverResponse(response) => {
382                                to_any_envelope(&envelope, response)
383                            }
384                            Response::GetCodeActionsResponse(response) => {
385                                to_any_envelope(&envelope, response)
386                            }
387                            Response::GetSignatureHelpResponse(response) => {
388                                to_any_envelope(&envelope, response)
389                            }
390                            Response::GetCodeLensResponse(response) => {
391                                to_any_envelope(&envelope, response)
392                            }
393                            Response::GetDocumentDiagnosticsResponse(response) => {
394                                to_any_envelope(&envelope, response)
395                            }
396                            Response::GetDefinitionResponse(response) => {
397                                to_any_envelope(&envelope, response)
398                            }
399                            Response::GetEditPredictionDefinitionResponse(response) => {
400                                to_any_envelope(&envelope, response)
401                            }
402                            Response::GetDeclarationResponse(response) => {
403                                to_any_envelope(&envelope, response)
404                            }
405                            Response::GetTypeDefinitionResponse(response) => {
406                                to_any_envelope(&envelope, response)
407                            }
408                            Response::GetEditPredictionTypeDefinitionResponse(response) => {
409                                to_any_envelope(&envelope, response)
410                            }
411                            Response::GetImplementationResponse(response) => {
412                                to_any_envelope(&envelope, response)
413                            }
414                            Response::InlayHintsResponse(response) => {
415                                to_any_envelope(&envelope, response)
416                            }
417                            Response::SemanticTokensResponse(response) => {
418                                to_any_envelope(&envelope, response)
419                            }
420                            Response::GetFoldingRangesResponse(response) => {
421                                to_any_envelope(&envelope, response)
422                            }
423                            Response::GetDocumentSymbolsResponse(response) => {
424                                to_any_envelope(&envelope, response)
425                            }
426                            Response::GetDocumentLinksResponse(response) => {
427                                to_any_envelope(&envelope, response)
428                            }
429                        };
430                        Some(proto::ProtoLspResponse {
431                            server_id,
432                            response,
433                        })
434                    })
435                    .collect(),
436            })))
437            .ok();
438        }
439    }
440
441    pub fn add_request_handler<M, E, H, F>(&self, entity: gpui::WeakEntity<E>, handler: H)
442    where
443        M: RequestMessage,
444        E: 'static,
445        H: 'static + Sync + Fn(Entity<E>, TypedEnvelope<M>, AsyncApp) -> F + Send + Sync,
446        F: 'static + Future<Output = Result<M::Response>>,
447    {
448        self.0
449            .client
450            .message_handler_set()
451            .lock()
452            .add_message_handler(
453                TypeId::of::<M>(),
454                entity.into(),
455                Arc::new(move |entity, envelope, client, cx| {
456                    let entity = entity.downcast::<E>().unwrap();
457                    let envelope = envelope.into_any().downcast::<TypedEnvelope<M>>().unwrap();
458                    let request_id = envelope.message_id();
459                    handler(entity, *envelope, cx)
460                        .then(move |result| async move {
461                            match result {
462                                Ok(response) => {
463                                    client.send_response(request_id, response)?;
464                                    Ok(())
465                                }
466                                Err(error) => {
467                                    client.send_response(request_id, error.to_proto())?;
468                                    Err(error)
469                                }
470                            }
471                        })
472                        .boxed_local()
473                }),
474            )
475    }
476
477    pub fn add_entity_request_handler<M, E, H, F>(&self, handler: H)
478    where
479        M: EnvelopedMessage + RequestMessage + EntityMessage,
480        E: 'static,
481        H: 'static + Sync + Send + Fn(gpui::Entity<E>, TypedEnvelope<M>, AsyncApp) -> F,
482        F: 'static + Future<Output = Result<M::Response>>,
483    {
484        let message_type_id = TypeId::of::<M>();
485        let entity_type_id = TypeId::of::<E>();
486        let entity_id_extractor = |envelope: &dyn AnyTypedEnvelope| {
487            (envelope as &dyn Any)
488                .downcast_ref::<TypedEnvelope<M>>()
489                .unwrap()
490                .payload
491                .remote_entity_id()
492        };
493        self.0
494            .client
495            .message_handler_set()
496            .lock()
497            .add_entity_message_handler(
498                message_type_id,
499                entity_type_id,
500                entity_id_extractor,
501                Arc::new(move |entity, envelope, client, cx| {
502                    let entity = entity.downcast::<E>().unwrap();
503                    let envelope = envelope.into_any().downcast::<TypedEnvelope<M>>().unwrap();
504                    let request_id = envelope.message_id();
505                    handler(entity, *envelope, cx)
506                        .then(move |result| async move {
507                            match result {
508                                Ok(response) => {
509                                    client.send_response(request_id, response)?;
510                                    Ok(())
511                                }
512                                Err(error) => {
513                                    client.send_response(request_id, error.to_proto())?;
514                                    Err(error)
515                                }
516                            }
517                        })
518                        .boxed_local()
519                }),
520            );
521    }
522
523    pub fn add_entity_stream_request_handler<M, E, H, F, S>(&self, handler: H)
524    where
525        M: EnvelopedMessage + RequestMessage + EntityMessage,
526        E: 'static,
527        H: 'static + Sync + Send + Fn(gpui::Entity<E>, TypedEnvelope<M>, AsyncApp) -> F,
528        F: 'static + Future<Output = Result<S>>,
529        S: 'static + Stream<Item = Result<M::Response>>,
530    {
531        let message_type_id = TypeId::of::<M>();
532        let entity_type_id = TypeId::of::<E>();
533        let entity_id_extractor = |envelope: &dyn AnyTypedEnvelope| {
534            (envelope as &dyn Any)
535                .downcast_ref::<TypedEnvelope<M>>()
536                .unwrap()
537                .payload
538                .remote_entity_id()
539        };
540        self.0
541            .client
542            .message_handler_set()
543            .lock()
544            .add_entity_message_handler(
545                message_type_id,
546                entity_type_id,
547                entity_id_extractor,
548                Arc::new(move |entity, envelope, client, cx| {
549                    let entity = entity.downcast::<E>().unwrap();
550                    let envelope = envelope.into_any().downcast::<TypedEnvelope<M>>().unwrap();
551                    let request_id = envelope.message_id();
552                    let stream = handler(entity, *envelope, cx);
553                    async move {
554                        // An Error response is itself a terminal stream frame on
555                        // both transports (Peer and ChannelClient), so we don't
556                        // need to follow it with an EndStream.
557                        match stream.await {
558                            Ok(stream) => {
559                                futures::pin_mut!(stream);
560                                while let Some(result) = stream.next().await {
561                                    match result {
562                                        Ok(response) => {
563                                            client.send_response(request_id, response)?
564                                        }
565                                        Err(error) => {
566                                            client.send_response(request_id, error.to_proto())?;
567                                            return Err(error);
568                                        }
569                                    }
570                                }
571                                client.send_response(request_id, proto::EndStream {})?;
572                                Ok(())
573                            }
574                            Err(error) => {
575                                client.send_response(request_id, error.to_proto())?;
576                                Err(error)
577                            }
578                        }
579                    }
580                    .boxed_local()
581                }),
582            );
583    }
584
585    pub fn add_entity_message_handler<M, E, H, F>(&self, handler: H)
586    where
587        M: EnvelopedMessage + EntityMessage,
588        E: 'static,
589        H: 'static + Sync + Send + Fn(gpui::Entity<E>, TypedEnvelope<M>, AsyncApp) -> F,
590        F: 'static + Future<Output = Result<()>>,
591    {
592        let message_type_id = TypeId::of::<M>();
593        let entity_type_id = TypeId::of::<E>();
594        let entity_id_extractor = |envelope: &dyn AnyTypedEnvelope| {
595            (envelope as &dyn Any)
596                .downcast_ref::<TypedEnvelope<M>>()
597                .unwrap()
598                .payload
599                .remote_entity_id()
600        };
601        self.0
602            .client
603            .message_handler_set()
604            .lock()
605            .add_entity_message_handler(
606                message_type_id,
607                entity_type_id,
608                entity_id_extractor,
609                Arc::new(move |entity, envelope, _, cx| {
610                    let entity = entity.downcast::<E>().unwrap();
611                    let envelope = envelope.into_any().downcast::<TypedEnvelope<M>>().unwrap();
612                    handler(entity, *envelope, cx).boxed_local()
613                }),
614            );
615    }
616
617    pub fn subscribe_to_entity<E: 'static>(&self, remote_id: u64, entity: &Entity<E>) {
618        let id = (TypeId::of::<E>(), remote_id);
619
620        let mut message_handlers = self.0.client.message_handler_set().lock();
621        if message_handlers
622            .entities_by_type_and_remote_id
623            .contains_key(&id)
624        {
625            panic!("already subscribed to entity");
626        }
627
628        message_handlers.entities_by_type_and_remote_id.insert(
629            id,
630            EntityMessageSubscriber::Entity {
631                handle: entity.downgrade().into(),
632            },
633        );
634    }
635
636    pub fn has_wsl_interop(&self) -> bool {
637        self.0.client.has_wsl_interop()
638    }
639}
640
641fn to_any_envelope<T: EnvelopedMessage>(
642    envelope: &TypedEnvelope<proto::LspQueryResponse>,
643    response: T,
644) -> Box<dyn AnyTypedEnvelope> {
645    Box::new(proto::TypedEnvelope {
646        sender_id: envelope.sender_id,
647        original_sender_id: envelope.original_sender_id,
648        message_id: envelope.message_id,
649        received_at: envelope.received_at,
650        payload: response,
651    }) as Box<_>
652}
653
654#[cfg(any(test, feature = "test-support"))]
655pub struct NoopProtoClient {
656    handler_set: parking_lot::Mutex<ProtoMessageHandlerSet>,
657}
658
659#[cfg(any(test, feature = "test-support"))]
660impl NoopProtoClient {
661    pub fn new() -> Arc<Self> {
662        Arc::new(Self {
663            handler_set: parking_lot::Mutex::new(ProtoMessageHandlerSet::default()),
664        })
665    }
666}
667
668#[cfg(any(test, feature = "test-support"))]
669impl ProtoClient for NoopProtoClient {
670    fn request(
671        &self,
672        _: proto::Envelope,
673        _: &'static str,
674    ) -> futures::future::BoxFuture<'static, Result<proto::Envelope>> {
675        unimplemented!()
676    }
677    fn send(&self, _: proto::Envelope, _: &'static str) -> Result<()> {
678        Ok(())
679    }
680    fn send_response(&self, _: proto::Envelope, _: &'static str) -> Result<()> {
681        Ok(())
682    }
683    fn message_handler_set(&self) -> &parking_lot::Mutex<ProtoMessageHandlerSet> {
684        &self.handler_set
685    }
686    fn is_via_collab(&self) -> bool {
687        false
688    }
689    fn has_wsl_interop(&self) -> bool {
690        false
691    }
692}
693