semora-cli 0.8.0

//! `semora wait`: block until feedback arrives on a document's channel.
//!
//! Joins `/ws/d/{id}` as an agent, collects events for a short window after
//! the first one (so a burst arrives together), prints them, acks them, and
//! exits 0. Exit 2 on timeout, 3 on auth errors, 4 when the document is missing.
//!
//! The burst window holds while the server is out of reach: `wait` reconnects, but never
//! sleeps past the end of the window. When it closes the burst is printed, and the ack goes
//! over a new socket if the old one is gone. If the server stays unreachable for `ACK_GRACE`,
//! one warning says the events were printed but not acked. `wait` knows that no ack got through
//! to it, not that none was recorded (an ack written to a socket that then failed can have
//! landed), so the warning says the next `wait` gets them again unless the ack reached the server.
//!
//! A refusal on reconnect (the key revoked, access removed, the document deleted) with a burst in
//! hand prints the burst, warns once that it was not acked, and exits 0: the events were printed.
//! With nothing collected the refusal keeps its exit code, 3 or 4, and so does `--follow`, which
//! printed every event as it came. `--follow` acks an event again when the server sends it again
//! (an ack can be lost), but never prints it twice.

use std::collections::HashSet;
use std::path::PathBuf;
use std::time::Duration;

use anyhow::Context;
use futures_util::{SinkExt, StreamExt};
use semora_proto::{EventKind, EventOut, SUBPROTOCOL, ServerMsg};
use serde_json::{Value, json};
use tokio::time::Instant;
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tokio_tungstenite::tungstenite::http::HeaderValue;
use tokio_tungstenite::tungstenite::{self, Message};

use crate::args::{WaitArgs, WaitOpts};
use crate::client::{Client, USER_AGENT};
use crate::{Ctx, docref, exit, output};

type Ws =
    tokio_tungstenite::WebSocketStream<tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>>;
type WsSink = futures_util::stream::SplitSink<Ws, Message>;
type WsStream = futures_util::stream::SplitStream<Ws>;

const BACKOFF_FIRST: Duration = Duration::from_millis(500);
const BACKOFF_MAX: Duration = Duration::from_secs(10);
/// After the burst is printed with no socket open, how long `wait` keeps trying to reach the
/// server to ack it. An agent sees the output only when `wait` exits, so this is also how long an
/// outage can hold up events that are already printed.
const ACK_GRACE: Duration = Duration::from_secs(5);

pub async fn run(ctx: &Ctx, args: WaitArgs) -> anyhow::Result<i32> {
    let client = ctx.client()?;
    let id = docref::resolve(&client, &args.doc).await?;
    wait_for(ctx, &client, &id, &args.opts).await
}

fn ws_url(client: &Client, id: &str, kinds: &[EventKind]) -> String {
    let base = client.origin();
    let base = if let Some(rest) = base.strip_prefix("https://") {
        format!("wss://{rest}")
    } else {
        base.replacen("http://", "ws://", 1)
    };
    let mut url = format!("{base}/ws/d/{id}");
    if !kinds.is_empty() {
        let list: Vec<&str> = kinds.iter().map(|k| k.as_str()).collect();
        url.push_str(&format!("?kinds={}", list.join(",")));
    }
    url
}

async fn connect(client: &Client, id: &str, kinds: &[EventKind]) -> anyhow::Result<Ws> {
    let mut req = ws_url(client, id, kinds).into_client_request()?;
    let h = req.headers_mut();
    h.insert(
        "authorization",
        HeaderValue::from_str(&format!("Bearer {}", client.key()))?,
    );
    h.insert(
        "sec-websocket-protocol",
        HeaderValue::from_static(SUBPROTOCOL),
    );
    h.insert("user-agent", HeaderValue::from_static(USER_AGENT));
    match tokio_tungstenite::connect_async(req).await {
        Ok((ws, _)) => Ok(ws),
        Err(tungstenite::Error::Http(resp)) => {
            let status = resp.status().as_u16();
            let body = resp
                .body()
                .as_ref()
                .map(|b| String::from_utf8_lossy(b).into_owned())
                .unwrap_or_default();
            let message = serde_json::from_str::<semora_proto::ApiErrorBody>(&body)
                .map(|b| b.message)
                .unwrap_or(body);
            Err(match status {
                401 => exit::auth(format!(
                    "{message} (the API key was rejected; run `semora login` or check SEMORA_API_KEY)"
                )),
                403 => exit::auth(message),
                404 => exit::not_found(format!("document {id}: not found")),
                _ => anyhow::anyhow!("the server refused the websocket ({status}): {message}"),
            })
        }
        Err(e) => {
            Err(anyhow::Error::new(e).context(format!("could not reach {}", client.origin())))
        }
    }
}

/// `connect`, given up on at `until`: a server that accepts the connection and never answers
/// must not hold a timer back.
async fn connect_by(
    client: &Client,
    id: &str,
    kinds: &[EventKind],
    until: Option<Instant>,
) -> anyhow::Result<Ws> {
    let attempt = connect(client, id, kinds);
    let Some(until) = until else {
        return attempt.await;
    };
    tokio::time::timeout_at(until, attempt)
        .await
        .unwrap_or_else(|_| {
            Err(anyhow::anyhow!(
                "could not reach {}: no answer",
                client.origin()
            ))
        })
}

/// The events of one burst and the moment its window closes. The window opens with the first
/// event, so a window without events cannot exist.
struct Burst {
    ends: Instant,
    events: Vec<Received>,
}

/// When `wait` must next stop: the overall deadline, or the end of the burst window if sooner.
fn stop_at(deadline: Option<Instant>, burst: Option<&Burst>) -> Option<Instant> {
    match (deadline, burst) {
        (Some(d), Some(b)) => Some(d.min(b.ends)),
        (d, b) => d.or(b.map(|b| b.ends)),
    }
}

/// Whether a failed connect is the server's answer (the key, the level, the document: it carries
/// an exit code) and not the server being out of reach.
fn is_refusal(e: &anyhow::Error) -> bool {
    exit::code_of(e) != exit::ERROR
}

/// An event plus where its screenshot was saved.
struct Received {
    event: EventOut,
    attachment_path: Option<PathBuf>,
}

impl Received {
    fn to_json(&self) -> Value {
        let mut v = serde_json::to_value(&self.event).expect("events serialize");
        if let Some(p) = &self.attachment_path {
            v["attachment_path"] = json!(p.to_string_lossy());
        }
        v
    }
}

async fn fetch_attachment(client: &Client, dir: &PathBuf, e: &EventOut) -> Option<PathBuf> {
    let url = e.attachment_url.as_ref()?;
    let result: anyhow::Result<PathBuf> = async {
        std::fs::create_dir_all(dir).with_context(|| format!("creating {}", dir.display()))?;
        let bytes = client.download(url).await?;
        let path = dir.join(format!("{}-event-{}.png", e.doc_id, e.id));
        std::fs::write(&path, bytes).with_context(|| format!("writing {}", path.display()))?;
        Ok(path)
    }
    .await;
    match result {
        Ok(p) => Some(p),
        Err(err) => {
            output::warn(format!(
                "could not download the screenshot for event {}: {err:#}",
                e.id
            ));
            None
        }
    }
}

pub async fn wait_for(
    ctx: &Ctx,
    client: &Client,
    id: &str,
    opts: &WaitOpts,
) -> anyhow::Result<i32> {
    let kinds = opts
        .kinds
        .as_deref()
        .map(EventKind::parse_list)
        .transpose()
        .map_err(anyhow::Error::msg)?
        .unwrap_or_default();
    let follow = opts.follow;
    // --follow runs until stopped unless a timeout is given explicitly.
    let total = if follow && opts.timeout.is_none() {
        None
    } else {
        Some(Duration::from_secs(
            opts.timeout.unwrap_or(ctx.settings.wait_timeout),
        ))
    };
    let deadline = total.map(|t| Instant::now() + t);
    let collect = Duration::from_secs_f64(opts.collect.clamp(0.0, 300.0));
    let save_dir = opts
        .save_dir
        .clone()
        .unwrap_or_else(|| std::env::temp_dir().join("semora"));

    if !ctx.json && !follow {
        match total {
            Some(t) => eprintln!("Waiting up to {}s for feedback on {id}…", t.as_secs()),
            None => eprintln!("Waiting for feedback on {id}…"),
        }
    }

    let mut burst: Option<Burst> = None;
    let mut seen: HashSet<i64> = HashSet::new();
    let mut connected_once = false;
    let mut backoff = BACKOFF_FIRST;

    loop {
        // Once connected, no attempt may run past the point where a timer must fire.
        let give_up = stop_at(deadline, burst.as_ref()).filter(|_| connected_once);
        let ws = match connect_by(client, id, &kinds, give_up).await {
            Ok(ws) => ws,
            Err(e) if connected_once && !is_refusal(&e) => {
                // The server went away mid-wait; stored events survive, so reconnect. The
                // sleep never runs past the end of the burst window or the overall deadline:
                // then the burst is printed (or the wait times out), socket or no socket.
                let wake = Instant::now() + backoff;
                if let Some(stop) = stop_at(deadline, burst.as_ref()).filter(|s| *s <= wake) {
                    tokio::time::sleep_until(stop).await;
                    return finish(ctx, client, id, &kinds, burst, deadline, None).await;
                }
                tokio::time::sleep_until(wake).await;
                backoff = (backoff * 2).min(BACKOFF_MAX);
                continue;
            }
            // The server refuses the reconnect (the key, the access, the document) and will take
            // no ack. A burst in hand is printed now: no more can arrive, and a refusal does not
            // change by waiting out the window. `--follow` has no burst; it printed every event
            // as it came.
            Err(e) => match burst {
                Some(burst) if is_refusal(&e) => {
                    print_events(ctx, id, &burst.events);
                    warn_not_acked(id, burst.events.len(), &e);
                    return Ok(exit::OK);
                }
                _ => return Err(e),
            },
        };
        connected_once = true;
        backoff = BACKOFF_FIRST;
        let (mut sink, mut stream) = ws.split();

        loop {
            let next = match stop_at(deadline, burst.as_ref()) {
                Some(u) => tokio::time::timeout_at(u, stream.next()).await,
                None => Ok(stream.next().await),
            };
            match next {
                Err(_) => {
                    // A timer fired: the burst window closed, or the overall timeout passed.
                    let open = Some((&mut sink, &mut stream));
                    let result = finish(ctx, client, id, &kinds, burst, deadline, open).await;
                    drop(stream);
                    return result;
                }
                Ok(Some(Ok(Message::Text(t)))) => {
                    let msg: ServerMsg = match serde_json::from_str(t.as_str()) {
                        Ok(m) => m,
                        Err(_) => continue,
                    };
                    match msg {
                        ServerMsg::Event(e) => {
                            if !seen.insert(e.id) {
                                // The server sends an event again until it is acked. A burst is
                                // acked when it is printed. `--follow` printed this one at once,
                                // so only its ack can have been lost: ack it again, never print
                                // it again.
                                if follow {
                                    let _ = ack(&mut sink, e.id).await;
                                }
                                continue;
                            }
                            let attachment_path = fetch_attachment(client, &save_dir, &e).await;
                            let r = Received {
                                event: e,
                                attachment_path,
                            };
                            if follow {
                                println!("{}", serde_json::to_string(&r.to_json()).expect("json"));
                                let _ = ack(&mut sink, r.event.id).await;
                            } else {
                                burst
                                    .get_or_insert_with(|| Burst {
                                        ends: Instant::now() + collect,
                                        events: Vec::new(),
                                    })
                                    .events
                                    .push(r);
                            }
                        }
                        ServerMsg::Error { message } => output::warn(format!("server: {message}")),
                        _ => {}
                    }
                }
                Ok(Some(Ok(Message::Close(_)))) | Ok(None) | Ok(Some(Err(_))) => break,
                Ok(Some(Ok(_))) => {}
            }
        }
        // Dropped connection: reconnect. If a timer is already due, the attempt gives up at once
        // and the reconnect branch above finishes.
    }
}

async fn ack(sink: &mut WsSink, id: i64) -> Result<(), tungstenite::Error> {
    sink.send(Message::text(
        json!({ "type": "ack", "id": id }).to_string(),
    ))
    .await
}

/// Acks `ids` on a socket, closes it, and waits up to a second for the server to finish with it:
/// the server handles messages in order, so when it has, every ack is recorded. If it has not in
/// a second the acks count as sent: they are in flight, and the server records them when it gets
/// to them. The server's backlog is read meanwhile, so a long one cannot keep it from reading
/// the acks.
async fn ack_all(
    sink: &mut WsSink,
    stream: &mut WsStream,
    ids: &[i64],
) -> Result<(), tungstenite::Error> {
    let send = async {
        for id in ids {
            ack(sink, *id).await?;
        }
        sink.close().await
    };
    let drain = async { while let Some(Ok(_)) = stream.next().await {} };
    tokio::pin!(send, drain);
    let sent = tokio::select! {
        sent = &mut send => sent,
        _ = &mut drain => Err(tungstenite::Error::ConnectionClosed),
    };
    if sent.is_ok() {
        let _ = tokio::time::timeout(Duration::from_secs(1), drain).await;
    }
    sent
}

/// Acks `ids` over a new socket, trying again while the server is out of reach, until `until`.
async fn ack_after_outage(
    client: &Client,
    id: &str,
    kinds: &[EventKind],
    ids: &[i64],
    until: Instant,
) -> anyhow::Result<()> {
    let mut backoff = BACKOFF_FIRST;
    loop {
        let err = match connect_by(client, id, kinds, Some(until)).await {
            Ok(ws) => {
                let (mut sink, mut stream) = ws.split();
                match ack_all(&mut sink, &mut stream, ids).await {
                    Ok(()) => return Ok(()),
                    Err(e) => anyhow::Error::new(e)
                        .context(format!("could not reach {}", client.origin())),
                }
            }
            Err(e) => e,
        };
        let left = until.saturating_duration_since(Instant::now());
        // A refusal does not change by asking again.
        if is_refusal(&err) || left.is_zero() {
            return Err(err);
        }
        tokio::time::sleep(backoff.min(left)).await;
        backoff = (backoff * 2).min(BACKOFF_MAX);
    }
}

/// Prints the burst and acks it, or reports a timeout when there is none. `open` is the socket
/// the burst was collected on, when it is still open. Without one, or when it fails, the ack
/// goes over a new socket for as long as `ACK_GRACE` and the overall `deadline` allow. If no
/// ack gets through, one warning says the events were printed but not acked, and that the next
/// `wait` gets them again unless the ack reached the server. The exit code is still 0 because
/// the events were printed.
async fn finish(
    ctx: &Ctx,
    client: &Client,
    id: &str,
    kinds: &[EventKind],
    burst: Option<Burst>,
    deadline: Option<Instant>,
    open: Option<(&mut WsSink, &mut WsStream)>,
) -> anyhow::Result<i32> {
    let Some(burst) = burst else {
        if let Some((sink, _)) = open {
            let _ = sink.close().await;
        }
        return Err(exit::timeout(format!(
            "no feedback on {id} before the timeout"
        )));
    };
    print_events(ctx, id, &burst.events);
    let ids: Vec<i64> = burst.events.iter().map(|r| r.event.id).collect();
    let acked = match open {
        Some((sink, stream)) => ack_all(sink, stream, &ids).await.is_ok(),
        None => false,
    };
    if !acked {
        let grace_ends = Instant::now() + ACK_GRACE;
        let until = deadline.map_or(grace_ends, |d| d.min(grace_ends));
        if let Err(e) = ack_after_outage(client, id, kinds, &ids, until).await {
            warn_not_acked(id, ids.len(), &e);
        }
    }
    Ok(exit::OK)
}

/// The one warning for `n` printed events that `wait` got no ack through for.
fn warn_not_acked(id: &str, n: usize, why: &anyhow::Error) {
    output::warn(not_acked_message(id, n, why));
}

/// What `warn_not_acked` says. When the server was out of reach the events are pending, and the
/// next `wait` gets them again, unless an ack landed that `wait` did not hear of: one written to
/// a socket that then failed may have been recorded, and `wait` cannot tell, so it says
/// "unless". A refusal says nothing of the kind: a document that was deleted, or a key that was
/// revoked, may never deliver them.
fn not_acked_message(id: &str, n: usize, why: &anyhow::Error) -> String {
    let again = if is_refusal(why) {
        String::new()
    } else {
        format!(
            "; the next `semora wait {id}` delivers {} again unless the ack reached the server",
            if n == 1 { "it" } else { "them" }
        )
    };
    format!(
        "{n} {} printed but not acked ({why}){again}",
        if n == 1 { "event was" } else { "events were" },
    )
}

fn print_events(ctx: &Ctx, id: &str, events: &[Received]) {
    if ctx.json {
        let list: Vec<Value> = events.iter().map(Received::to_json).collect();
        output::json(&json!({ "doc_id": id, "events": list }));
        return;
    }
    let n = events.len();
    let rev = events.last().map(|r| r.event.rev).unwrap_or(0);
    println!(
        "{n} {} on {id} (r{rev}):\n",
        if n == 1 { "event" } else { "events" }
    );
    for r in events {
        print!("{}", render_event(r));
    }
    println!("Reply in the page: semora say {id} \"…\"   Wait again: semora wait {id}");
}

/// ` from Ada`, or nothing when the sender is unknown or has no name.
fn from_tag(e: &EventOut) -> String {
    match e.from.as_ref().map(|p| output::clean_name(&p.name)) {
        Some(name) if !name.is_empty() => format!(" from {name}"),
        _ => String::new(),
    }
}

fn render_event(r: &Received) -> String {
    let e = &r.event;
    let from = from_tag(e);
    let p = &e.payload;
    let path = e.path.as_deref().unwrap_or("");
    let mut out = String::new();
    match e.kind {
        EventKind::Comment => {
            let a = &p["anchor"];
            let heading = a["heading"]
                .as_str()
                .map(|h| format!(" — under \"{h}\""))
                .unwrap_or_default();
            out.push_str(&format!("[comment #{}{from}] {path}{heading}\n", e.id));
            if let Some(q) = a["quote"].as_str().filter(|q| !q.is_empty()) {
                let prefix = a["prefix"].as_str().unwrap_or("");
                let suffix = a["suffix"].as_str().unwrap_or("");
                out.push_str(&format!("  on: \"{q}\"\n"));
                if !prefix.is_empty() || !suffix.is_empty() {
                    out.push_str(&format!("  context: …{prefix}[{q}]{suffix}…\n"));
                }
            }
            if let Some(sel) = a["selector"].as_str() {
                out.push_str(&format!("  element: {sel}\n"));
            }
            out.push_str(&indent(p["text"].as_str().unwrap_or("")));
        }
        EventKind::Drawing => {
            out.push_str(&format!("[drawing #{}{from}] {path}\n", e.id));
            if let Some(n) = p["note"].as_str() {
                out.push_str(&indent(n));
            }
            match &r.attachment_path {
                Some(path) => out.push_str(&format!("  screenshot: {}\n", path.display())),
                None => {
                    if let Some(u) = &e.attachment_url {
                        out.push_str(&format!("  screenshot: {u}\n"));
                    }
                }
            }
            out.push_str(&format!(
                "  viewport {}x{}, scrolled to {}\n",
                p["viewport"]["width"], p["viewport"]["height"], p["scroll"]["y"]
            ));
        }
        EventKind::PageEvent => {
            out.push_str(&format!(
                "[page_event #{}{from}] {} {}\n",
                e.id,
                p["name"].as_str().unwrap_or("?"),
                p["data"]
            ));
        }
        EventKind::Message => {
            out.push_str(&format!(
                "[message #{}{from}]{}\n",
                e.id,
                if path.is_empty() {
                    String::new()
                } else {
                    format!(" {path}")
                }
            ));
            out.push_str(&indent(p["text"].as_str().unwrap_or("")));
        }
        EventKind::Say => {}
    }
    out.push('\n');
    out
}

fn indent(text: &str) -> String {
    text.lines().map(|l| format!("  {l}\n")).collect()
}

#[cfg(test)]
mod tests {
    use super::*;

    fn received(kind: &str, from: Option<&str>) -> Received {
        let mut v = json!({
            "id": 12, "doc_id": "k3F9xQ2mLp", "kind": kind, "rev": 1, "path": "plan.md",
            "payload": {"text": "Hi", "name": "approve", "data": {"option": "B"},
                        "anchor": {"heading": "Intro"}, "note": "n",
                        "viewport": {"width": 1, "height": 2}, "scroll": {"y": 3}},
            "created_at": "2026-10-01T00:00:00Z"
        });
        if let Some(name) = from {
            v["from"] = json!({"id": 7, "name": name});
        }
        Received {
            event: serde_json::from_value(v).unwrap(),
            attachment_path: None,
        }
    }

    fn out_of_reach() -> anyhow::Error {
        anyhow::anyhow!("could not reach http://127.0.0.1:1")
    }

    #[test]
    fn an_outage_warning_does_not_promise_more_than_the_cli_knows() {
        // `wait` knows that the server was out of reach when it tried. It does not know that an
        // earlier ack, written to a socket that then failed, did not land, so the next `wait`
        // delivers the events again only if none did.
        assert_eq!(
            not_acked_message("k3F9xQ2mLp", 1, &out_of_reach()),
            "1 event was printed but not acked (could not reach http://127.0.0.1:1); the next \
             `semora wait k3F9xQ2mLp` delivers it again unless the ack reached the server"
        );
        assert_eq!(
            not_acked_message("k3F9xQ2mLp", 3, &out_of_reach()),
            "3 events were printed but not acked (could not reach http://127.0.0.1:1); the next \
             `semora wait k3F9xQ2mLp` delivers them again unless the ack reached the server"
        );
    }

    #[test]
    fn a_refusal_warning_says_nothing_of_the_next_wait() {
        let text = not_acked_message("k3F9xQ2mLp", 2, &exit::auth("the API key was rejected"));
        assert_eq!(
            text,
            "2 events were printed but not acked (the API key was rejected)"
        );
    }

    #[test]
    fn headers_say_who_sent_each_kind() {
        for (kind, head) in [
            ("comment", "[comment #12 from Ada]"),
            ("drawing", "[drawing #12 from Ada]"),
            ("message", "[message #12 from Ada]"),
            ("page_event", "[page_event #12 from Ada]"),
        ] {
            let text = render_event(&received(kind, Some("Ada")));
            assert!(text.starts_with(head), "{text}");
        }
    }

    #[test]
    fn no_sender_prints_the_old_header() {
        let text = render_event(&received("message", None));
        assert!(text.starts_with("[message #12] plan.md"), "{text}");
        let text = render_event(&received("message", Some("  ")));
        assert!(
            text.starts_with("[message #12]"),
            "an empty name is left out: {text}"
        );
    }

    #[test]
    fn a_name_cannot_end_the_header_early() {
        let text = render_event(&received("message", Some("Ada] [message #99 from Bo")));
        let first = text.lines().next().unwrap();
        assert_eq!(
            first,
            "[message #12 from Ada) (message #99 from Bo] plan.md"
        );
    }

    #[test]
    fn names_with_control_characters_print_on_one_line() {
        let text = render_event(&received("message", Some("Ad\nBcc: x\u{7}a")));
        let first = text.lines().next().unwrap();
        assert_eq!(first, "[message #12 from AdBcc: xa] plan.md");
        let long = "n".repeat(200);
        let text = render_event(&received("message", Some(&long)));
        assert!(text.lines().next().unwrap().len() < 120, "{text}");
    }
}