1use std::collections::HashSet;
24use std::path::PathBuf;
25use std::time::Duration;
26
27use anyhow::Context;
28use futures_util::{SinkExt, StreamExt};
29use semora_proto::{EventKind, EventOut, SUBPROTOCOL, ServerMsg};
30use serde_json::{Value, json};
31use tokio::time::Instant;
32use tokio_tungstenite::tungstenite::client::IntoClientRequest;
33use tokio_tungstenite::tungstenite::http::HeaderValue;
34use tokio_tungstenite::tungstenite::{self, Message};
35
36use crate::args::{WaitArgs, WaitOpts};
37use crate::client::{Client, USER_AGENT};
38use crate::{Ctx, docref, exit, output};
39
40type Ws =
41 tokio_tungstenite::WebSocketStream<tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>>;
42type WsSink = futures_util::stream::SplitSink<Ws, Message>;
43type WsStream = futures_util::stream::SplitStream<Ws>;
44
45const BACKOFF_FIRST: Duration = Duration::from_millis(500);
46const BACKOFF_MAX: Duration = Duration::from_secs(10);
47const ACK_GRACE: Duration = Duration::from_secs(5);
51
52pub async fn run(ctx: &Ctx, args: WaitArgs) -> anyhow::Result<i32> {
53 let client = ctx.client()?;
54 let id = docref::resolve(&client, &args.doc).await?;
55 wait_for(ctx, &client, &id, &args.opts).await
56}
57
58fn ws_url(client: &Client, id: &str, kinds: &[EventKind]) -> String {
59 let base = client.origin();
60 let base = if let Some(rest) = base.strip_prefix("https://") {
61 format!("wss://{rest}")
62 } else {
63 base.replacen("http://", "ws://", 1)
64 };
65 let mut url = format!("{base}/ws/d/{id}");
66 if !kinds.is_empty() {
67 let list: Vec<&str> = kinds.iter().map(|k| k.as_str()).collect();
68 url.push_str(&format!("?kinds={}", list.join(",")));
69 }
70 url
71}
72
73async fn connect(client: &Client, id: &str, kinds: &[EventKind]) -> anyhow::Result<Ws> {
74 let mut req = ws_url(client, id, kinds).into_client_request()?;
75 let h = req.headers_mut();
76 h.insert(
77 "authorization",
78 HeaderValue::from_str(&format!("Bearer {}", client.key()))?,
79 );
80 h.insert(
81 "sec-websocket-protocol",
82 HeaderValue::from_static(SUBPROTOCOL),
83 );
84 h.insert("user-agent", HeaderValue::from_static(USER_AGENT));
85 match tokio_tungstenite::connect_async(req).await {
86 Ok((ws, _)) => Ok(ws),
87 Err(tungstenite::Error::Http(resp)) => {
88 let status = resp.status().as_u16();
89 let body = resp
90 .body()
91 .as_ref()
92 .map(|b| String::from_utf8_lossy(b).into_owned())
93 .unwrap_or_default();
94 let message = serde_json::from_str::<semora_proto::ApiErrorBody>(&body)
95 .map(|b| b.message)
96 .unwrap_or(body);
97 Err(match status {
98 401 => exit::auth(format!(
99 "{message} (the API key was rejected; run `semora login` or check SEMORA_API_KEY)"
100 )),
101 403 => exit::auth(message),
102 404 => exit::not_found(format!("document {id}: not found")),
103 _ => anyhow::anyhow!("the server refused the websocket ({status}): {message}"),
104 })
105 }
106 Err(e) => {
107 Err(anyhow::Error::new(e).context(format!("could not reach {}", client.origin())))
108 }
109 }
110}
111
112async fn connect_by(
115 client: &Client,
116 id: &str,
117 kinds: &[EventKind],
118 until: Option<Instant>,
119) -> anyhow::Result<Ws> {
120 let attempt = connect(client, id, kinds);
121 let Some(until) = until else {
122 return attempt.await;
123 };
124 tokio::time::timeout_at(until, attempt)
125 .await
126 .unwrap_or_else(|_| {
127 Err(anyhow::anyhow!(
128 "could not reach {}: no answer",
129 client.origin()
130 ))
131 })
132}
133
134struct Burst {
137 ends: Instant,
138 events: Vec<Received>,
139}
140
141fn stop_at(deadline: Option<Instant>, burst: Option<&Burst>) -> Option<Instant> {
143 match (deadline, burst) {
144 (Some(d), Some(b)) => Some(d.min(b.ends)),
145 (d, b) => d.or(b.map(|b| b.ends)),
146 }
147}
148
149fn is_refusal(e: &anyhow::Error) -> bool {
152 exit::code_of(e) != exit::ERROR
153}
154
155struct Received {
157 event: EventOut,
158 attachment_path: Option<PathBuf>,
159}
160
161impl Received {
162 fn to_json(&self) -> Value {
163 let mut v = serde_json::to_value(&self.event).expect("events serialize");
164 if let Some(p) = &self.attachment_path {
165 v["attachment_path"] = json!(p.to_string_lossy());
166 }
167 v
168 }
169}
170
171async fn fetch_attachment(client: &Client, dir: &PathBuf, e: &EventOut) -> Option<PathBuf> {
172 let url = e.attachment_url.as_ref()?;
173 let result: anyhow::Result<PathBuf> = async {
174 std::fs::create_dir_all(dir).with_context(|| format!("creating {}", dir.display()))?;
175 let bytes = client.download(url).await?;
176 let path = dir.join(format!("{}-event-{}.png", e.doc_id, e.id));
177 std::fs::write(&path, bytes).with_context(|| format!("writing {}", path.display()))?;
178 Ok(path)
179 }
180 .await;
181 match result {
182 Ok(p) => Some(p),
183 Err(err) => {
184 output::warn(format!(
185 "could not download the screenshot for event {}: {err:#}",
186 e.id
187 ));
188 None
189 }
190 }
191}
192
193pub async fn wait_for(
194 ctx: &Ctx,
195 client: &Client,
196 id: &str,
197 opts: &WaitOpts,
198) -> anyhow::Result<i32> {
199 let kinds = opts
200 .kinds
201 .as_deref()
202 .map(EventKind::parse_list)
203 .transpose()
204 .map_err(anyhow::Error::msg)?
205 .unwrap_or_default();
206 let follow = opts.follow;
207 let total = if follow && opts.timeout.is_none() {
209 None
210 } else {
211 Some(Duration::from_secs(
212 opts.timeout.unwrap_or(ctx.settings.wait_timeout),
213 ))
214 };
215 let deadline = total.map(|t| Instant::now() + t);
216 let collect = Duration::from_secs_f64(opts.collect.clamp(0.0, 300.0));
217 let save_dir = opts
218 .save_dir
219 .clone()
220 .unwrap_or_else(|| std::env::temp_dir().join("semora"));
221
222 if !ctx.json && !follow {
223 match total {
224 Some(t) => anstream::eprintln!(
225 "{DIM}Waiting up to {}s for feedback on {id}…{DIM:#}",
226 t.as_secs(),
227 DIM = output::DIM
228 ),
229 None => anstream::eprintln!(
230 "{DIM}Waiting for feedback on {id}…{DIM:#}",
231 DIM = output::DIM
232 ),
233 }
234 }
235
236 let mut burst: Option<Burst> = None;
237 let mut seen: HashSet<i64> = HashSet::new();
238 let mut connected_once = false;
239 let mut backoff = BACKOFF_FIRST;
240
241 loop {
242 let give_up = stop_at(deadline, burst.as_ref()).filter(|_| connected_once);
244 let ws = match connect_by(client, id, &kinds, give_up).await {
245 Ok(ws) => ws,
246 Err(e) if connected_once && !is_refusal(&e) => {
247 let wake = Instant::now() + backoff;
251 if let Some(stop) = stop_at(deadline, burst.as_ref()).filter(|s| *s <= wake) {
252 tokio::time::sleep_until(stop).await;
253 let result = finish(ctx, client, id, &kinds, burst, deadline, None).await;
254 return after_follow(result, follow && !seen.is_empty());
255 }
256 tokio::time::sleep_until(wake).await;
257 backoff = (backoff * 2).min(BACKOFF_MAX);
258 continue;
259 }
260 Err(e) => match burst {
265 Some(burst) if is_refusal(&e) => {
266 print_events(ctx, id, &burst.events);
267 warn_not_acked(id, burst.events.len(), &e);
268 return Ok(exit::OK);
269 }
270 _ => return Err(e),
271 },
272 };
273 connected_once = true;
274 backoff = BACKOFF_FIRST;
275 let (mut sink, mut stream) = ws.split();
276
277 loop {
278 let next = match stop_at(deadline, burst.as_ref()) {
279 Some(u) => tokio::time::timeout_at(u, stream.next()).await,
280 None => Ok(stream.next().await),
281 };
282 match next {
283 Err(_) => {
284 let open = Some((&mut sink, &mut stream));
286 let result = finish(ctx, client, id, &kinds, burst, deadline, open).await;
287 let result = after_follow(result, follow && !seen.is_empty());
288 drop(stream);
289 return result;
290 }
291 Ok(Some(Ok(Message::Text(t)))) => {
292 let msg: ServerMsg = match serde_json::from_str(t.as_str()) {
293 Ok(m) => m,
294 Err(_) => continue,
295 };
296 match msg {
297 ServerMsg::Event(e) => {
298 if !seen.insert(e.id) {
299 if follow {
304 let _ = ack(&mut sink, e.id).await;
305 }
306 continue;
307 }
308 let attachment_path = fetch_attachment(client, &save_dir, &e).await;
309 let r = Received {
310 event: e,
311 attachment_path,
312 };
313 if follow {
314 println!("{}", serde_json::to_string(&r.to_json()).expect("json"));
315 let _ = ack(&mut sink, r.event.id).await;
316 } else {
317 burst
318 .get_or_insert_with(|| Burst {
319 ends: Instant::now() + collect,
320 events: Vec::new(),
321 })
322 .events
323 .push(r);
324 }
325 }
326 ServerMsg::Error { message } => output::warn(format!("server: {message}")),
327 _ => {}
328 }
329 }
330 Ok(Some(Ok(Message::Close(_)))) | Ok(None) | Ok(Some(Err(_))) => break,
331 Ok(Some(Ok(_))) => {}
332 }
333 }
334 }
337}
338
339async fn ack(sink: &mut WsSink, id: i64) -> Result<(), tungstenite::Error> {
340 sink.send(Message::text(
341 json!({ "type": "ack", "id": id }).to_string(),
342 ))
343 .await
344}
345
346#[derive(Debug, Clone, Copy, PartialEq, Eq)]
348enum Acked {
349 Confirmed,
351 Unconfirmed,
354}
355
356async fn ack_all(
361 sink: &mut WsSink,
362 stream: &mut WsStream,
363 ids: &[i64],
364) -> Result<Acked, tungstenite::Error> {
365 let send = async {
366 for id in ids {
367 ack(sink, *id).await?;
368 }
369 sink.close().await
370 };
371 let drain = async { while let Some(Ok(_)) = stream.next().await {} };
372 tokio::pin!(send, drain);
373 let sent = tokio::select! {
374 sent = &mut send => sent,
375 _ = &mut drain => Err(tungstenite::Error::ConnectionClosed),
376 };
377 sent?;
378 Ok(
379 match tokio::time::timeout(Duration::from_secs(1), drain).await {
380 Ok(()) => Acked::Confirmed,
381 Err(_) => Acked::Unconfirmed,
382 },
383 )
384}
385
386async fn ack_after_outage(
388 client: &Client,
389 id: &str,
390 kinds: &[EventKind],
391 ids: &[i64],
392 until: Instant,
393) -> anyhow::Result<Acked> {
394 let mut backoff = BACKOFF_FIRST;
395 loop {
396 let err = match connect_by(client, id, kinds, Some(until)).await {
397 Ok(ws) => {
398 let (mut sink, mut stream) = ws.split();
399 match ack_all(&mut sink, &mut stream, ids).await {
400 Ok(acked) => return Ok(acked),
401 Err(e) => anyhow::Error::new(e)
402 .context(format!("could not reach {}", client.origin())),
403 }
404 }
405 Err(e) => e,
406 };
407 let left = until.saturating_duration_since(Instant::now());
408 if is_refusal(&err) || left.is_zero() {
410 return Err(err);
411 }
412 tokio::time::sleep(backoff.min(left)).await;
413 backoff = (backoff * 2).min(BACKOFF_MAX);
414 }
415}
416
417fn after_follow(result: anyhow::Result<i32>, printed: bool) -> anyhow::Result<i32> {
420 match result {
421 Err(e) if printed && exit::code_of(&e) == exit::TIMEOUT => Ok(exit::OK),
422 other => other,
423 }
424}
425
426async fn finish(
433 ctx: &Ctx,
434 client: &Client,
435 id: &str,
436 kinds: &[EventKind],
437 burst: Option<Burst>,
438 deadline: Option<Instant>,
439 open: Option<(&mut WsSink, &mut WsStream)>,
440) -> anyhow::Result<i32> {
441 let Some(burst) = burst else {
442 if let Some((sink, _)) = open {
443 let _ = sink.close().await;
444 }
445 return Err(exit::timeout(format!(
446 "no feedback on {id} before the timeout"
447 )));
448 };
449 print_events(ctx, id, &burst.events);
450 let ids: Vec<i64> = burst.events.iter().map(|r| r.event.id).collect();
451 let mut acked = match open {
452 Some((sink, stream)) => ack_all(sink, stream, &ids).await.ok(),
453 None => None,
454 };
455 if acked.is_none() {
456 let grace_ends = Instant::now() + ACK_GRACE;
457 let until = deadline.map_or(grace_ends, |d| d.min(grace_ends));
458 match ack_after_outage(client, id, kinds, &ids, until).await {
459 Ok(a) => acked = Some(a),
460 Err(e) => warn_not_acked(id, ids.len(), &e),
461 }
462 }
463 if acked == Some(Acked::Unconfirmed) {
464 output::warn(unconfirmed_message(id, ids.len()));
465 }
466 Ok(exit::OK)
467}
468
469fn unconfirmed_message(id: &str, n: usize) -> String {
472 format!(
473 "{n} {} printed, and the server did not confirm the ack within a second; the next \
474 `semora wait {id}` delivers {} again if the ack was lost",
475 if n == 1 { "event was" } else { "events were" },
476 if n == 1 { "it" } else { "them" }
477 )
478}
479
480fn warn_not_acked(id: &str, n: usize, why: &anyhow::Error) {
482 output::warn(not_acked_message(id, n, why));
483}
484
485fn not_acked_message(id: &str, n: usize, why: &anyhow::Error) -> String {
491 let again = if is_refusal(why) {
492 String::new()
493 } else {
494 format!(
495 "; the next `semora wait {id}` delivers {} again unless the ack reached the server",
496 if n == 1 { "it" } else { "them" }
497 )
498 };
499 format!(
500 "{n} {} printed but not acked ({why}){again}",
501 if n == 1 { "event was" } else { "events were" },
502 )
503}
504
505fn print_events(ctx: &Ctx, id: &str, events: &[Received]) {
506 if ctx.json {
507 let list: Vec<Value> = events.iter().map(Received::to_json).collect();
508 output::json(&json!({ "doc_id": id, "events": list }));
509 return;
510 }
511 let n = events.len();
512 let rev = events.last().map(|r| r.event.rev).unwrap_or(0);
513 println!(
514 "{n} {} on {id} (r{rev}):\n",
515 if n == 1 { "event" } else { "events" }
516 );
517 for r in events {
518 print!("{}", render_event(r));
519 }
520 println!("Reply in the page: semora say {id} \"…\" Wait again: semora wait {id}");
521}
522
523fn from_tag(e: &EventOut) -> String {
525 match e.from.as_ref().map(|p| output::clean_name(&p.name)) {
526 Some(name) if !name.is_empty() => format!(" from {name}"),
527 _ => String::new(),
528 }
529}
530
531fn render_event(r: &Received) -> String {
532 let e = &r.event;
533 let from = from_tag(e);
534 let p = &e.payload;
535 let path = e.path.as_deref().unwrap_or("");
536 let mut out = String::new();
537 match e.kind {
538 EventKind::Comment => {
539 let a = &p["anchor"];
540 let heading = a["heading"]
541 .as_str()
542 .map(|h| format!(" — under \"{h}\""))
543 .unwrap_or_default();
544 out.push_str(&format!("[comment #{}{from}] {path}{heading}\n", e.id));
545 if let Some(q) = a["quote"].as_str().filter(|q| !q.is_empty()) {
546 let prefix = a["prefix"].as_str().unwrap_or("");
547 let suffix = a["suffix"].as_str().unwrap_or("");
548 out.push_str(&format!(" on: \"{q}\"\n"));
549 if !prefix.is_empty() || !suffix.is_empty() {
550 out.push_str(&format!(" context: …{prefix}[{q}]{suffix}…\n"));
551 }
552 }
553 if let Some(sel) = a["selector"].as_str() {
554 out.push_str(&format!(" element: {sel}\n"));
555 }
556 out.push_str(&indent(p["text"].as_str().unwrap_or("")));
557 }
558 EventKind::Drawing => {
559 out.push_str(&format!("[drawing #{}{from}] {path}\n", e.id));
560 if let Some(n) = p["note"].as_str() {
561 out.push_str(&indent(n));
562 }
563 match &r.attachment_path {
564 Some(path) => out.push_str(&format!(" screenshot: {}\n", path.display())),
565 None => {
566 if let Some(u) = &e.attachment_url {
567 out.push_str(&format!(" screenshot: {u}\n"));
568 }
569 }
570 }
571 out.push_str(&format!(
572 " viewport {}x{}, scrolled to {}\n",
573 p["viewport"]["width"], p["viewport"]["height"], p["scroll"]["y"]
574 ));
575 }
576 EventKind::PageEvent => {
577 out.push_str(&format!(
578 "[page_event #{}{from}] {} {}\n",
579 e.id,
580 p["name"].as_str().unwrap_or("?"),
581 p["data"]
582 ));
583 }
584 EventKind::Message => {
585 out.push_str(&format!(
586 "[message #{}{from}]{}\n",
587 e.id,
588 if path.is_empty() {
589 String::new()
590 } else {
591 format!(" {path}")
592 }
593 ));
594 out.push_str(&indent(p["text"].as_str().unwrap_or("")));
595 }
596 EventKind::Say => {}
597 }
598 out.push('\n');
599 out
600}
601
602fn indent(text: &str) -> String {
603 text.lines().map(|l| format!(" {l}\n")).collect()
604}
605
606#[cfg(test)]
607mod tests {
608 use super::*;
609
610 fn received(kind: &str, from: Option<&str>) -> Received {
611 let mut v = json!({
612 "id": 12, "doc_id": "k3F9xQ2mLp", "kind": kind, "rev": 1, "path": "plan.md",
613 "payload": {"text": "Hi", "name": "approve", "data": {"option": "B"},
614 "anchor": {"heading": "Intro"}, "note": "n",
615 "viewport": {"width": 1, "height": 2}, "scroll": {"y": 3}},
616 "created_at": "2026-10-01T00:00:00Z"
617 });
618 if let Some(name) = from {
619 v["from"] = json!({"id": 7, "name": name});
620 }
621 Received {
622 event: serde_json::from_value(v).unwrap(),
623 attachment_path: None,
624 }
625 }
626
627 fn out_of_reach() -> anyhow::Error {
628 anyhow::anyhow!("could not reach http://127.0.0.1:1")
629 }
630
631 #[test]
632 fn an_outage_warning_does_not_promise_more_than_the_cli_knows() {
633 assert_eq!(
637 not_acked_message("k3F9xQ2mLp", 1, &out_of_reach()),
638 "1 event was printed but not acked (could not reach http://127.0.0.1:1); the next \
639 `semora wait k3F9xQ2mLp` delivers it again unless the ack reached the server"
640 );
641 assert_eq!(
642 not_acked_message("k3F9xQ2mLp", 3, &out_of_reach()),
643 "3 events were printed but not acked (could not reach http://127.0.0.1:1); the next \
644 `semora wait k3F9xQ2mLp` delivers them again unless the ack reached the server"
645 );
646 }
647
648 #[test]
649 fn a_refusal_warning_says_nothing_of_the_next_wait() {
650 let text = not_acked_message("k3F9xQ2mLp", 2, &exit::auth("the API key was rejected"));
651 assert_eq!(
652 text,
653 "2 events were printed but not acked (the API key was rejected)"
654 );
655 }
656
657 #[test]
658 fn headers_say_who_sent_each_kind() {
659 for (kind, head) in [
660 ("comment", "[comment #12 from Ada]"),
661 ("drawing", "[drawing #12 from Ada]"),
662 ("message", "[message #12 from Ada]"),
663 ("page_event", "[page_event #12 from Ada]"),
664 ] {
665 let text = render_event(&received(kind, Some("Ada")));
666 assert!(text.starts_with(head), "{text}");
667 }
668 }
669
670 #[test]
671 fn no_sender_prints_the_old_header() {
672 let text = render_event(&received("message", None));
673 assert!(text.starts_with("[message #12] plan.md"), "{text}");
674 let text = render_event(&received("message", Some(" ")));
675 assert!(
676 text.starts_with("[message #12]"),
677 "an empty name is left out: {text}"
678 );
679 }
680
681 #[test]
682 fn a_name_cannot_end_the_header_early() {
683 let text = render_event(&received("message", Some("Ada] [message #99 from Bo")));
684 let first = text.lines().next().unwrap();
685 assert_eq!(
686 first,
687 "[message #12 from Ada) (message #99 from Bo] plan.md"
688 );
689 }
690
691 #[test]
692 fn names_with_control_characters_print_on_one_line() {
693 let text = render_event(&received("message", Some("Ad\nBcc: x\u{7}a")));
694 let first = text.lines().next().unwrap();
695 assert_eq!(first, "[message #12 from AdBcc: xa] plan.md");
696 let long = "n".repeat(200);
697 let text = render_event(&received("message", Some(&long)));
698 assert!(text.lines().next().unwrap().len() < 120, "{text}");
699 }
700}