use serde_json::Value;
use std::collections::HashMap;
use std::io::{ErrorKind, Write};
use std::net::TcpStream;
use std::time::{Duration, Instant};
use tungstenite::error::UrlError;
use tungstenite::stream::MaybeTlsStream;
use tungstenite::{Message, WebSocket};
pub const NODE_TYPES: [&str; 7] = [
"goal",
"option",
"decision",
"action",
"outcome",
"observation",
"revisit",
];
pub const READ_TIMEOUT: Duration = Duration::from_secs(75);
#[derive(Debug, Default, Clone)]
pub struct Filter {
pub types: Vec<String>,
pub branches: Vec<String>,
pub edges: bool,
pub json: bool,
}
impl Filter {
pub fn unknown_types(&self) -> Vec<&str> {
self.types
.iter()
.map(String::as_str)
.filter(|t| !NODE_TYPES.contains(t))
.collect()
}
pub fn admits(&self, event: &Value) -> bool {
let branch = event.get("branch").and_then(Value::as_str);
if !self.branches.is_empty()
&& !branch.is_some_and(|b| self.branches.iter().any(|f| f == b))
{
return false;
}
match event.get("table").and_then(Value::as_str) {
Some("decision_nodes") => {
let node_type = event.get("node_type").and_then(Value::as_str).unwrap_or("");
self.types.is_empty() || self.types.iter().any(|t| t == node_type)
}
Some("decision_edges") => self.edges,
_ => true,
}
}
}
pub fn quote(title: &str) -> String {
let mut out = String::with_capacity(title.len() + 2);
out.push('"');
for c in title.chars() {
match c {
'"' => out.push_str("\\\""),
'\\' => out.push_str("\\\\"),
'\n' => out.push_str("\\n"),
'\t' => out.push_str("\\t"),
'\r' => out.push_str("\\r"),
c if c.is_control() => out.push_str(&format!("\\u{{{:x}}}", c as u32)),
c => out.push(c),
}
}
out.push('"');
out
}
pub fn format_event(event: &Value, at: &str) -> Option<String> {
let table = event.get("table").and_then(Value::as_str)?;
let op = event.get("op").and_then(Value::as_str).unwrap_or("INSERT");
let branch = event
.get("branch")
.and_then(Value::as_str)
.filter(|b| !b.is_empty())
.unwrap_or("-");
let short = |key: &str| -> String {
event
.get(key)
.and_then(Value::as_str)
.map(|s| s.chars().take(8).collect())
.unwrap_or_else(|| "????????".to_string())
};
let line = match table {
"decision_nodes" => {
let node_type = event
.get("node_type")
.and_then(Value::as_str)
.unwrap_or("node");
let kind = match op {
"INSERT" => node_type.to_string(),
"UPDATE" => format!("{node_type} (updated)"),
"DELETE" => format!("{node_type} (deleted)"),
other => format!("{node_type} ({})", other.to_lowercase()),
};
let what = match event.get("title").and_then(Value::as_str) {
Some(title) => quote(title),
None => short("change_id"),
};
format!("{at} {branch:<8} {kind:<11} {what}")
}
"decision_edges" => {
let edge_type = event
.get("edge_type")
.and_then(Value::as_str)
.unwrap_or("leads_to");
let verb = match op {
"INSERT" => format!("edge {edge_type}"),
"DELETE" => format!("edge {edge_type} (deleted)"),
other => format!("edge {edge_type} ({})", other.to_lowercase()),
};
format!(
"{at} {branch:<8} {verb} {} -> {}",
short("from_change_id"),
short("to_change_id")
)
}
other => format!("{at} {branch:<8} {other} {op}"),
};
Some(line)
}
pub struct Dedup {
seen: HashMap<String, Instant>,
window: Duration,
}
impl Dedup {
pub fn new(window: Duration) -> Self {
Self {
seen: HashMap::new(),
window,
}
}
pub fn first_sighting(&mut self, frame: &str) -> bool {
self.first_sighting_at(frame, Instant::now())
}
fn first_sighting_at(&mut self, frame: &str, now: Instant) -> bool {
self.seen
.retain(|_, at| now.duration_since(*at) < self.window);
match self.seen.get(frame) {
Some(_) => false,
None => {
self.seen.insert(frame.to_string(), now);
true
}
}
}
}
pub fn connect_failure(e: &tungstenite::Error) -> String {
match e {
tungstenite::Error::Url(UrlError::UnableToConnect(_)) => {
"connect failed: unable to connect".to_string()
}
tungstenite::Error::Io(io) => format!("connect failed: {}", io.kind()),
tungstenite::Error::Http(resp) => format!("connect failed: HTTP {}", resp.status()),
other => format!("connect failed: {other}"),
}
}
fn is_permanent(e: &tungstenite::Error) -> Option<String> {
match e {
tungstenite::Error::Http(resp) if resp.status().is_client_error() => Some(format!(
"the server refused the connection with HTTP {}; check the token and the workspace",
resp.status()
)),
_ => None,
}
}
pub fn run(url: &str, filter: &Filter, out: &mut dyn Write) -> Result<(), String> {
let mut backoff = Duration::from_secs(1);
let max_backoff = Duration::from_secs(30);
let mut dedup = Dedup::new(Duration::from_secs(60));
loop {
match tungstenite::connect(url) {
Ok((socket, _response)) => {
set_read_timeout(&socket);
let (why, delivered) = read_until_closed(socket, filter, &mut dedup, out)?;
if delivered {
backoff = Duration::from_secs(1);
}
eprintln!("reconnecting ({why})");
}
Err(e) => {
if let Some(reason) = is_permanent(&e) {
return Err(reason);
}
eprintln!("reconnecting ({})", connect_failure(&e));
}
}
std::thread::sleep(backoff);
backoff = (backoff * 2).min(max_backoff);
}
}
fn set_read_timeout(socket: &WebSocket<MaybeTlsStream<TcpStream>>) {
let tcp = match socket.get_ref() {
MaybeTlsStream::Plain(s) => s,
MaybeTlsStream::Rustls(s) => s.get_ref(),
_ => return,
};
let _ = tcp.set_read_timeout(Some(READ_TIMEOUT));
}
fn read_until_closed(
mut socket: WebSocket<MaybeTlsStream<TcpStream>>,
filter: &Filter,
dedup: &mut Dedup,
out: &mut dyn Write,
) -> Result<(String, bool), String> {
let mut delivered = false;
loop {
match socket.read() {
Ok(Message::Text(text)) => {
delivered = true;
let text = text.as_str();
let event: Value = match serde_json::from_str(text) {
Ok(v) => v,
Err(_) => {
emit(out, text)?;
continue;
}
};
if !filter.admits(&event) || !dedup.first_sighting(text) {
continue;
}
if filter.json {
emit(out, text)?;
} else if let Some(line) = format_event(&event, &now_hms()) {
emit(out, &line)?;
}
}
Ok(Message::Close(frame)) => {
let why = match frame {
Some(f) => format!("close {}", u16::from(f.code)),
None => "close".to_string(),
};
return Ok((why, delivered));
}
Ok(_) => {}
Err(tungstenite::Error::ConnectionClosed) => {
return Ok(("closed".to_string(), delivered))
}
Err(tungstenite::Error::Io(io))
if matches!(io.kind(), ErrorKind::WouldBlock | ErrorKind::TimedOut) =>
{
return Ok((
format!("read timed out after {}s", READ_TIMEOUT.as_secs()),
delivered,
));
}
Err(e) => return Ok((format!("error: {e}"), delivered)),
}
}
}
fn emit(out: &mut dyn Write, line: &str) -> Result<(), String> {
match writeln!(out, "{line}").and_then(|_| out.flush()) {
Ok(()) => Ok(()),
Err(e) if e.kind() == ErrorKind::BrokenPipe => std::process::exit(0),
Err(e) => Err(e.to_string()),
}
}
fn now_hms() -> String {
chrono::Local::now().format("%H:%M:%S").to_string()
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
fn node(op: &str) -> Value {
json!({
"table": "decision_nodes", "op": op, "workspace": "tetris-arena",
"id": "d43f55ba-7c6e-4fdb-b528-9850fabf4c20",
"change_id": "83c0f4bb-1870-4eb2-86d3-eedc62c6e25b",
"node_type": "outcome", "title": "Round 4: clamp DAS through the clear flash",
"branch": "agent-3"
})
}
#[test]
fn an_insert_quotes_the_title() {
assert_eq!(
format_event(&node("INSERT"), "05:01:12").unwrap(),
"05:01:12 agent-3 outcome \"Round 4: clamp DAS through the clear flash\""
);
}
#[test]
fn branches_line_up_across_one_and_two_digit_agents() {
let mut ten = node("INSERT");
ten["branch"] = json!("agent-10");
let a = format_event(&node("INSERT"), "t").unwrap();
let b = format_event(&ten, "t").unwrap();
assert_eq!(a.find("outcome"), b.find("outcome"), "{a}\n{b}");
}
#[test]
fn a_title_keeps_its_emoji_and_escapes_its_newlines() {
let mut e = node("INSERT");
e["title"] = json!("I \u{2764}\u{fe0f} kicks\nline two \"quoted\"");
let line = format_event(&e, "t").unwrap();
assert!(line.contains("I \u{2764}\u{fe0f} kicks"), "{line}");
assert!(line.contains("\\n"), "{line}");
assert!(line.contains("\\\"quoted\\\""), "{line}");
assert!(!line.contains("u{fe0f}"), "{line}");
}
#[test]
fn a_refused_connection_never_prints_the_url() {
let url = "wss://example.test/events?workspace=w&token=SECRET".to_string();
let e = tungstenite::Error::Url(UrlError::UnableToConnect(url));
let msg = connect_failure(&e);
assert!(!msg.contains("token="), "{msg}");
assert!(!msg.contains("SECRET"), "{msg}");
assert_eq!(msg, "connect failed: unable to connect");
}
#[test]
fn a_typo_in_types_is_named() {
let f = Filter {
types: vec!["outcome".into(), "outcomes".into()],
..Default::default()
};
assert_eq!(f.unknown_types(), vec!["outcomes"]);
assert!(Filter::default().unknown_types().is_empty());
}
#[test]
fn an_update_is_not_a_second_node() {
let line = format_event(&node("UPDATE"), "05:01:12").unwrap();
assert!(line.contains("outcome (updated)"), "{line}");
}
#[test]
fn a_missing_title_falls_back_to_the_change_id_prefix() {
let mut e = node("INSERT");
e.as_object_mut().unwrap().remove("title");
assert_eq!(
format_event(&e, "05:01:12").unwrap(),
"05:01:12 agent-3 outcome 83c0f4bb"
);
}
#[test]
fn a_missing_branch_prints_a_dash() {
let mut e = node("INSERT");
e.as_object_mut().unwrap().remove("branch");
assert!(format_event(&e, "t")
.unwrap()
.starts_with("t - outcome"));
}
#[test]
fn an_edge_shows_both_endpoints() {
let e = json!({
"table": "decision_edges", "op": "INSERT", "workspace": "tetris-arena",
"id": "x", "edge_type": "chosen", "branch": "agent-3",
"from_change_id": "a1b2c3d4-0000", "to_change_id": "e5f6a7b8-0000"
});
assert_eq!(
format_event(&e, "t").unwrap(),
"t agent-3 edge chosen a1b2c3d4 -> e5f6a7b8"
);
}
#[test]
fn a_frame_without_a_table_is_not_an_event() {
assert!(format_event(&json!({"hello": 1}), "t").is_none());
}
#[test]
fn edges_are_hidden_unless_asked_for() {
let e = json!({"table": "decision_edges", "op": "INSERT", "id": "x"});
assert!(!Filter::default().admits(&e));
let f = Filter {
edges: true,
..Default::default()
};
assert!(f.admits(&e));
}
#[test]
fn type_and_branch_filters_narrow_nodes() {
let f = Filter {
types: vec!["outcome".into()],
branches: vec!["agent-3".into()],
..Default::default()
};
assert!(f.admits(&node("INSERT")));
let mut other_type = node("INSERT");
other_type["node_type"] = json!("option");
assert!(!f.admits(&other_type));
let mut other_branch = node("INSERT");
other_branch["branch"] = json!("agent-4");
assert!(!f.admits(&other_branch));
let mut no_branch = node("INSERT");
no_branch.as_object_mut().unwrap().remove("branch");
assert!(
!f.admits(&no_branch),
"no branch cannot be the asked-for branch"
);
}
#[test]
fn a_repeated_frame_prints_once_inside_the_window() {
let mut d = Dedup::new(Duration::from_secs(60));
let t0 = Instant::now();
let frame = node("INSERT").to_string();
assert!(d.first_sighting_at(&frame, t0));
assert!(!d.first_sighting_at(&frame, t0 + Duration::from_secs(1)));
assert!(d.first_sighting_at(&frame, t0 + Duration::from_secs(61)));
}
#[test]
fn two_different_updates_to_one_node_both_print() {
let mut d = Dedup::new(Duration::from_secs(60));
let t0 = Instant::now();
let mut first = node("UPDATE");
first["status"] = json!("active");
let mut second = node("UPDATE");
second["status"] = json!("completed");
assert!(d.first_sighting_at(&first.to_string(), t0));
assert!(
d.first_sighting_at(&second.to_string(), t0 + Duration::from_secs(1)),
"a second update that says something different is a second event"
);
}
}