1use sqlx::{PgPool, postgres::PgListener};
10use tokio::sync::broadcast;
11use tokio_util::sync::CancellationToken;
12use uuid::Uuid;
13
14pub const PG_CHANNEL: &str = "bus_events";
15
16#[derive(Clone, Debug)]
19pub struct BusEvent(pub serde_json::Value);
20
21impl BusEvent {
22 pub fn kind(&self) -> &str {
23 self.0.get("kind").and_then(|v| v.as_str()).unwrap_or("")
24 }
25
26 fn uuid_field(&self, key: &str) -> Option<Uuid> {
27 self.0
28 .get(key)
29 .and_then(|v| v.as_str())
30 .and_then(|s| Uuid::parse_str(s).ok())
31 }
32
33 pub fn team_id(&self) -> Option<Uuid> {
34 self.uuid_field("team_id")
35 }
36
37 pub fn recipient_agent_id(&self) -> Option<Uuid> {
38 self.uuid_field("recipient_agent_id")
39 }
40
41 pub fn sender_agent_id(&self) -> Option<Uuid> {
42 self.uuid_field("sender_agent_id")
43 }
44
45 pub fn channel_id(&self) -> Option<Uuid> {
46 self.uuid_field("channel_id")
47 }
48
49 pub fn message_id(&self) -> Option<i64> {
50 self.0.get("id").and_then(|v| v.as_i64())
51 }
52
53 pub fn recipient_session(&self) -> Option<&str> {
56 self.0.get("recipient_session").and_then(|v| v.as_str())
57 }
58
59 pub fn sender_session(&self) -> Option<&str> {
61 self.0.get("sender_session").and_then(|v| v.as_str())
62 }
63
64 pub fn is_announcement(&self) -> bool {
67 self.0
68 .get("announce")
69 .and_then(|v| v.as_bool())
70 .unwrap_or(false)
71 }
72
73 pub fn is_direct_message(&self) -> bool {
74 self.kind() == "message" && self.recipient_agent_id().is_some()
75 }
76
77 pub fn visible_to(&self, team_id: Uuid, agent_id: Uuid, session: &str) -> bool {
86 if self.team_id() != Some(team_id) {
87 return false;
88 }
89 if self.is_direct_message() {
90 if self.sender_agent_id() == Some(agent_id) {
91 return true;
92 }
93 if self.recipient_agent_id() != Some(agent_id) {
94 return false;
95 }
96 return match self.recipient_session() {
97 Some(addressed) => addressed == session,
98 None => true,
100 };
101 }
102 true
103 }
104}
105
106#[derive(Clone)]
107pub struct EventHub {
108 tx: broadcast::Sender<BusEvent>,
109}
110
111impl Default for EventHub {
112 fn default() -> Self {
113 Self::new()
114 }
115}
116
117impl EventHub {
118 pub fn new() -> Self {
119 let (tx, _) = broadcast::channel(256);
122 Self { tx }
123 }
124
125 pub fn subscribe(&self) -> broadcast::Receiver<BusEvent> {
126 self.tx.subscribe()
127 }
128
129 pub fn publish(&self, event: BusEvent) {
130 let _ = self.tx.send(event);
132 }
133}
134
135pub async fn run_pg_listener(pool: PgPool, hub: EventHub, ct: CancellationToken) {
138 loop {
139 if ct.is_cancelled() {
140 return;
141 }
142 match PgListener::connect_with(&pool).await {
143 Ok(mut listener) => {
144 if let Err(e) = listener.listen(PG_CHANNEL).await {
145 tracing::warn!(error = %e, "LISTEN failed; retrying");
146 } else {
147 tracing::info!("event listener attached to '{PG_CHANNEL}'");
148 loop {
149 tokio::select! {
150 _ = ct.cancelled() => return,
151 recv = listener.try_recv() => match recv {
152 Ok(Some(notification)) => {
153 match serde_json::from_str(notification.payload()) {
154 Ok(value) => hub.publish(BusEvent(value)),
155 Err(e) => tracing::warn!(
156 error = %e,
157 payload = notification.payload(),
158 "unparseable bus event"
159 ),
160 }
161 }
162 Ok(None) => tracing::debug!("event listener reconnected"),
166 Err(e) => {
167 tracing::warn!(error = %e, "event listener error");
168 break;
169 }
170 }
171 }
172 }
173 }
174 }
175 Err(e) => {
176 tracing::warn!(error = %e, "could not attach event listener; retrying");
177 }
178 }
179 tokio::select! {
180 _ = ct.cancelled() => return,
181 _ = tokio::time::sleep(std::time::Duration::from_secs(2)) => {}
182 }
183 }
184}