tenant.openagents/omega

1use crate::{Envelope, PeerId};
2use anyhow::{Context as _, Result};
3use serde::Serialize;
4use std::{
5    any::{Any, TypeId},
6    cmp,
7    fmt::{self, Debug},
8};
9use std::{marker::PhantomData, time::Instant};
10
11pub trait EnvelopedMessage: Clone + Debug + Serialize + Sized + Send + Sync + 'static {
12    const NAME: &'static str;
13    const PRIORITY: MessagePriority;
14    fn into_envelope(
15        self,
16        id: u32,
17        responding_to: Option<u32>,
18        original_sender_id: Option<PeerId>,
19    ) -> Envelope;
20    fn from_envelope(envelope: Envelope) -> Option<Self>;
21}
22
23pub trait EntityMessage: EnvelopedMessage {
24    type Entity;
25    fn remote_entity_id(&self) -> u64;
26}
27
28pub trait RequestMessage: EnvelopedMessage {
29    type Response: EnvelopedMessage;
30}
31
32/// A trait to bind LSP request and responses for the proto layer.
33/// Should be used for every LSP request that has to traverse through the proto layer.
34///
35/// `lsp_messages` macro in the same crate provides a convenient way to implement this.
36pub trait LspRequestMessage: EnvelopedMessage {
37    type Response: EnvelopedMessage;
38
39    fn to_proto_query(self) -> crate::lsp_query::Request;
40
41    fn response_to_proto_query(response: Self::Response) -> crate::lsp_response::Response;
42
43    fn buffer_id(&self) -> u64;
44
45    fn buffer_version(&self) -> &[crate::VectorClockEntry];
46
47    /// Whether to deduplicate the requests, or keep the previous ones running when another
48    /// request of the same kind is processed.
49    fn stop_previous_requests() -> bool;
50}
51
52#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
53pub struct LspRequestId(pub u64);
54
55/// A response from a single language server.
56/// There could be multiple responses for a single LSP request,
57/// from different servers.
58pub struct ProtoLspResponse<R> {
59    pub server_id: u64,
60    pub response: R,
61}
62
63impl ProtoLspResponse<Box<dyn AnyTypedEnvelope>> {
64    pub fn into_response<T: LspRequestMessage>(self) -> Result<ProtoLspResponse<T::Response>> {
65        let envelope = self
66            .response
67            .into_any()
68            .downcast::<TypedEnvelope<T::Response>>()
69            .map_err(|_| {
70                anyhow::anyhow!(
71                    "cannot downcast LspResponse to {} for message {}",
72                    T::Response::NAME,
73                    T::NAME,
74                )
75            })?;
76
77        Ok(ProtoLspResponse {
78            server_id: self.server_id,
79            response: envelope.payload,
80        })
81    }
82}
83
84pub trait AnyTypedEnvelope: Any + Send + Sync {
85    fn payload_type_id(&self) -> TypeId;
86    fn payload_type_name(&self) -> &'static str;
87    fn into_any(self: Box<Self>) -> Box<dyn Any + Send + Sync>;
88    fn is_background(&self) -> bool;
89    fn original_sender_id(&self) -> Option<PeerId>;
90    fn sender_id(&self) -> PeerId;
91    fn message_id(&self) -> u32;
92}
93
94pub enum MessagePriority {
95    Foreground,
96    Background,
97}
98
99impl<T: EnvelopedMessage> AnyTypedEnvelope for TypedEnvelope<T> {
100    fn payload_type_id(&self) -> TypeId {
101        TypeId::of::<T>()
102    }
103
104    fn payload_type_name(&self) -> &'static str {
105        T::NAME
106    }
107
108    fn into_any(self: Box<Self>) -> Box<dyn Any + Send + Sync> {
109        self
110    }
111
112    fn is_background(&self) -> bool {
113        matches!(T::PRIORITY, MessagePriority::Background)
114    }
115
116    fn original_sender_id(&self) -> Option<PeerId> {
117        self.original_sender_id
118    }
119
120    fn sender_id(&self) -> PeerId {
121        self.sender_id
122    }
123
124    fn message_id(&self) -> u32 {
125        self.message_id
126    }
127}
128
129impl PeerId {
130    pub fn from_u64(peer_id: u64) -> Self {
131        let owner_id = (peer_id >> 32) as u32;
132        let id = peer_id as u32;
133        Self { owner_id, id }
134    }
135
136    pub fn as_u64(self) -> u64 {
137        ((self.owner_id as u64) << 32) | (self.id as u64)
138    }
139}
140
141impl Copy for PeerId {}
142
143impl Eq for PeerId {}
144
145impl Ord for PeerId {
146    fn cmp(&self, other: &Self) -> cmp::Ordering {
147        self.owner_id
148            .cmp(&other.owner_id)
149            .then_with(|| self.id.cmp(&other.id))
150    }
151}
152
153impl PartialOrd for PeerId {
154    fn partial_cmp(&self, other: &Self) -> Option<cmp::Ordering> {
155        Some(self.cmp(other))
156    }
157}
158
159impl std::hash::Hash for PeerId {
160    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
161        self.owner_id.hash(state);
162        self.id.hash(state);
163    }
164}
165
166impl fmt::Display for PeerId {
167    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
168        write!(f, "{}/{}", self.owner_id, self.id)
169    }
170}
171
172pub struct Receipt<T> {
173    pub sender_id: PeerId,
174    pub message_id: u32,
175    payload_type: PhantomData<T>,
176}
177
178impl<T> Clone for Receipt<T> {
179    fn clone(&self) -> Self {
180        *self
181    }
182}
183
184impl<T> Copy for Receipt<T> {}
185
186#[derive(Clone, Debug)]
187pub struct TypedEnvelope<T> {
188    pub sender_id: PeerId,
189    pub original_sender_id: Option<PeerId>,
190    pub message_id: u32,
191    pub payload: T,
192    pub received_at: Instant,
193}
194
195impl<T> TypedEnvelope<T> {
196    pub fn original_sender_id(&self) -> Result<PeerId> {
197        self.original_sender_id
198            .context("missing original_sender_id")
199    }
200}
201
202impl<T: RequestMessage> TypedEnvelope<T> {
203    pub fn receipt(&self) -> Receipt<T> {
204        Receipt {
205            sender_id: self.sender_id,
206            message_id: self.message_id,
207            payload_type: PhantomData,
208        }
209    }
210}
211