1use std::collections::HashSet;
25use std::time::Duration;
26
27use async_stream::try_stream;
28use futures_util::{SinkExt, Stream, StreamExt};
29use tokio_tungstenite::tungstenite::Message;
30
31pub use cloacina_api_types::delivery::DELIVERY_PROTOCOL_VERSION;
32use cloacina_api_types::delivery::{ClientMessage, ServerMessage};
33
34use crate::error::ClientError;
35use crate::Client;
36
37#[derive(Debug, Clone)]
39pub struct DeliveryPush {
40 pub id: i64,
42 pub kind: String,
44 pub recipient: String,
45 pub tenant_id: Option<String>,
46 pub payload: Vec<u8>,
48}
49
50#[derive(Debug, Clone)]
52pub struct SubscribeOptions {
53 pub reconnect: bool,
55 pub reconnect_initial: Duration,
57 pub reconnect_max: Duration,
59}
60
61impl Default for SubscribeOptions {
62 fn default() -> Self {
63 Self {
64 reconnect: true,
65 reconnect_initial: Duration::from_millis(100),
66 reconnect_max: Duration::from_secs(30),
67 }
68 }
69}
70
71fn ws_base(server: &str) -> Result<String, ClientError> {
72 if let Some(rest) = server.strip_prefix("https://") {
73 Ok(format!("wss://{rest}"))
74 } else if let Some(rest) = server.strip_prefix("http://") {
75 Ok(format!("ws://{rest}"))
76 } else {
77 Err(ClientError::Config(format!(
78 "server must start with http:// or https:// (got {server})"
79 )))
80 }
81}
82
83pub(crate) fn subscribe_delivery(
84 client: Client,
85 recipient: String,
86 options: SubscribeOptions,
87) -> impl Stream<Item = Result<DeliveryPush, ClientError>> {
88 try_stream! {
89 let base = ws_base(client.server())?;
90 let mut seen: HashSet<i64> = HashSet::new();
91 let mut backoff = options.reconnect_initial;
92
93 loop {
94 let ticket = client.create_ws_ticket().await?.ticket;
96 let url = format!(
97 "{base}/v1/ws/delivery/{}?token={}",
98 urlencoding::encode(&recipient),
99 urlencoding::encode(&ticket),
100 );
101
102 let (mut socket, _resp) = tokio_tungstenite::connect_async(&url)
103 .await
104 .map_err(|e| ClientError::Ws(format!("connect failed for {url}: {e}")))?;
105
106 let hello = serde_json::to_string(&ClientMessage::Hello {
109 protocol_version: DELIVERY_PROTOCOL_VERSION,
110 since_id: None,
111 })
112 .expect("hello serializes");
113 socket
114 .send(Message::Text(hello))
115 .await
116 .map_err(|e| ClientError::Ws(format!("hello send failed: {e}")))?;
117
118 let mut close_code: Option<u16> = None;
119
120 while let Some(msg) = socket.next().await {
121 let msg = match msg {
122 Ok(m) => m,
123 Err(e) => {
124 if !options.reconnect {
125 Err(ClientError::Ws(format!("recv error: {e}")))?;
126 }
127 break;
128 }
129 };
130 match msg {
131 Message::Text(text) => {
132 let frame: ServerMessage = match serde_json::from_str(&text) {
133 Ok(f) => f,
134 Err(_) => continue, };
136 if let ServerMessage::Push { id, kind, recipient: r, tenant_id, .. } = &frame {
137 let payload = frame
138 .decode_push_payload()
139 .map_err(|e| ClientError::Ws(format!("bad push payload: {e}")))?;
140 let push = DeliveryPush {
141 id: *id,
142 kind: kind.clone(),
143 recipient: r.clone(),
144 tenant_id: tenant_id.clone(),
145 payload,
146 };
147 let fresh = seen.insert(push.id);
148 let ack_id = push.id;
149 if fresh {
150 yield push;
151 }
152 let ack = serde_json::to_string(&ClientMessage::Ack {
155 protocol_version: DELIVERY_PROTOCOL_VERSION,
156 id: ack_id,
157 })
158 .expect("ack serializes");
159 if socket.send(Message::Text(ack)).await.is_err() {
160 break;
161 }
162 }
163 backoff = options.reconnect_initial;
164 }
165 Message::Close(frame) => {
166 close_code = frame.map(|f| f.code.into());
167 break;
168 }
169 _ => {}
170 }
171 }
172
173 if close_code == Some(4426) {
175 Err(ClientError::ProtocolVersion {
176 client_version: DELIVERY_PROTOCOL_VERSION,
177 })?;
178 }
179 if !options.reconnect {
180 break;
181 }
182 tokio::time::sleep(backoff).await;
183 backoff = (backoff * 2).min(options.reconnect_max);
184 }
185 }
186}
187
188pub(crate) fn follow_execution_events(
189 client: Client,
190 execution_id: String,
191 options: SubscribeOptions,
192) -> impl Stream<Item = Result<serde_json::Value, ClientError>> {
193 try_stream! {
194 let recipient = format!("exec_events:{execution_id}");
195 let stream = subscribe_delivery(client, recipient, options);
196 let mut stream = std::pin::pin!(stream);
197 while let Some(push) = stream.next().await {
198 let push = push?;
199 let event: serde_json::Value = serde_json::from_slice(&push.payload)
200 .map_err(|e| ClientError::Ws(format!("push payload is not JSON: {e}")))?;
201 yield event;
202 }
203 }
204}