Skip to content

feat(observability): correlate agent relay delivery #4769

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
MajorTal wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat(observability): correlate agent relay delivery #4769

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions Cargo.lock
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

137 changes: 122 additions & 15 deletions crates/buzz-acp/src/relay.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,9 @@ use std::time::Duration;
const EVENT_CHANNEL_CAPACITY_DEFAULT: usize = 256;
/// Capacity of the command channel from harness to background task.
const CMD_CHANNEL_CAPACITY: usize = 64;
/// Upgrade-response header emitted by buzz-relay for joining client and server
/// diagnostics without exposing a tenant or identity.
const RELAY_CONNECTION_ID_HEADER: &str = "x-buzz-connection-id";

/// Read the event channel capacity from the environment, falling back to the
/// compiled-in default. Parsed once at call-site (connect time).
Expand Down Expand Up @@ -121,7 +124,7 @@ use buzz_core::kind::{
KIND_TYPING_INDICATOR,
};
use futures_util::{SinkExt, StreamExt};
use nostr::{Event, EventBuilder, Keys, Kind, RelayUrl, Tag};
use nostr::{Event, EventBuilder, Keys, Kind, RelayUrl, Tag, Timestamp};
use serde_json::{json, Value};
use tokio::sync::mpsc;
use tokio::time::timeout;
Expand Down Expand Up @@ -620,7 +623,7 @@ impl HarnessRelay {
// jittered backoff. A terminal error (bad URL, bad auth tag,
// rejected/invalid signing key) fails immediately — see
// `is_terminal_connect_error`.
let (ws, handshake_buffer) =
let (ws, handshake_buffer, relay_connection_id) =
retry_initial_connect(|| do_connect(relay_url, keys, auth_tag.as_ref())).await?;

let (event_tx, event_rx) = mpsc::channel::<Option<BuzzEvent>>(event_channel_capacity());
Expand All @@ -644,6 +647,7 @@ impl HarnessRelay {
bg_relay_url,
bg_agent_pubkey_hex,
bg_auth_tag,
relay_connection_id,
)
.await;
});
Expand Down Expand Up @@ -990,6 +994,15 @@ impl TwoGenDedup {

/// State maintained by the background WebSocket task.
struct BgState {
/// Opaque connection ID supplied by the relay in the WebSocket upgrade.
/// This joins harness logs to one exact server-side connection without
/// logging message content or client identity.
relay_connection_id: Option<String>,
/// Monotonic socket generation within this harness process.
connection_generation: u64,
/// Relay connection establishment time, used to identify events that
/// predate the current socket and therefore likely arrived via replay.
connected_at: u64,
/// Active subscriptions: channel_id → subscription_id string.
active_subscriptions: HashMap<Uuid, String>,
/// Most recent `created_at` timestamp seen per channel (for `since` filter).
Expand Down Expand Up @@ -1078,6 +1091,9 @@ struct BgState {
impl BgState {
fn new() -> Self {
Self {
relay_connection_id: None,
connection_generation: 0,
connected_at: 0,
active_subscriptions: HashMap::new(),
last_seen: HashMap::new(),
seen_ids: TwoGenDedup::new(SEEN_ID_LIMIT),
Expand All @@ -1102,6 +1118,12 @@ impl BgState {
}
}

fn note_connection(&mut self, relay_connection_id: Option<String>) {
self.relay_connection_id = relay_connection_id;
self.connection_generation = self.connection_generation.saturating_add(1);
self.connected_at = Timestamp::now().as_secs();
}

/// Record a received event for dedup and `since` tracking.
/// Returns `true` if the event is new (not a duplicate).
fn record_event(&mut self, channel_id: Uuid, event: &Event) -> bool {
Expand Down Expand Up @@ -1557,8 +1579,10 @@ async fn run_background_task(
relay_url: String,
agent_pubkey_hex: String,
auth_tag: Option<nostr::Tag>,
initial_relay_connection_id: Option<String>,
) {
let mut state = BgState::new();
state.note_connection(initial_relay_connection_id);

let handshake_ok = process_handshake_buffer(
&mut ws,
Expand Down Expand Up @@ -1822,11 +1846,20 @@ async fn run_background_task(
}
}
Some(Err(e)) => {
warn!("WebSocket error in background task: {e}");
warn!(
relay_connection_id = state.relay_connection_id.as_deref().unwrap_or("unavailable"),
connection_generation = state.connection_generation,
error = %e,
"WebSocket error in background task"
);
true
}
None => {
debug!("WebSocket stream ended");
warn!(
relay_connection_id = state.relay_connection_id.as_deref().unwrap_or("unavailable"),
connection_generation = state.connection_generation,
"WebSocket stream ended without a Close frame"
);
true
}
};
Expand Down Expand Up @@ -1953,7 +1986,12 @@ async fn run_background_task(
_ = ping_interval.tick() => {
if ping_sent && last_pong.elapsed() > PONG_TIMEOUT {
// No pong received after our last ping — connection is dead.
warn!("no pong received within {:?} — connection dead, reconnecting", PONG_TIMEOUT);
warn!(
relay_connection_id = state.relay_connection_id.as_deref().unwrap_or("unavailable"),
connection_generation = state.connection_generation,
timeout_secs = PONG_TIMEOUT.as_secs(),
"no pong received — connection dead, reconnecting"
);
// Use try_send to avoid blocking on backpressure during recovery.
let _ = event_tx.try_send(None);
match try_autonomous_reconnect(
Expand Down Expand Up @@ -2156,6 +2194,18 @@ async fn handle_ws_message(
let ts = event.created_at.as_secs();
let event_id_hex = event.id.to_hex();
if state.record_event(channel_id, &event) {
let received_at = Timestamp::now().as_secs();
debug!(
relay_connection_id = state.relay_connection_id.as_deref().unwrap_or("unavailable"),
connection_generation = state.connection_generation,
subscription_id = %subscription_id,
channel_id = %channel_id,
event_id = %event_id_hex,
event_created_at = ts,
delivery_age_secs = received_at.saturating_sub(ts),
event_predates_connection = ts < state.connected_at,
"relay channel event received"
);
let buzz_event = BuzzEvent {
channel_id,
event: *event,
Expand Down Expand Up @@ -2390,8 +2440,24 @@ async fn handle_ws_message(
}
true
}
Message::Close(_) => {
debug!("relay sent Close frame");
Message::Close(frame) => {
match frame {
Some(frame) => warn!(
relay_connection_id = state.relay_connection_id.as_deref().unwrap_or("unavailable"),
connection_generation = state.connection_generation,
close_code = ?frame.code,
close_reason = %frame.reason,
"relay sent Close frame"
),
None => warn!(
relay_connection_id = state
.relay_connection_id
.as_deref()
.unwrap_or("unavailable"),
connection_generation = state.connection_generation,
"relay sent Close frame without code or reason"
),
}
false
}
// Binary, Pong, Frame — ignore
Expand Down Expand Up @@ -2938,9 +3004,18 @@ async fn try_autonomous_reconnect(
backoffs.len()
);
match do_connect(relay_url, keys, auth_tag).await {
Ok((new_ws, handshake_buffer)) => {
Ok((new_ws, handshake_buffer, relay_connection_id)) => {
*ws = new_ws;
info!("autonomous reconnect succeeded (attempt {})", attempt + 1);
state.note_connection(relay_connection_id);
info!(
relay_connection_id = state
.relay_connection_id
.as_deref()
.unwrap_or("unavailable"),
connection_generation = state.connection_generation,
attempt = attempt + 1,
"autonomous reconnect succeeded"
);
let handshake_ok = process_handshake_buffer(
ws,
handshake_buffer,
Expand Down Expand Up @@ -3076,9 +3151,18 @@ async fn wait_for_reconnect(
loop {
info!("attempting relay reconnect to {relay_url}…");
match do_connect(relay_url, keys, auth_tag).await {
Ok((new_ws, handshake_buffer)) => {
Ok((new_ws, handshake_buffer, relay_connection_id)) => {
*ws = new_ws;
info!("relay reconnected to {relay_url}");
state.note_connection(relay_connection_id);
info!(
relay_connection_id = state
.relay_connection_id
.as_deref()
.unwrap_or("unavailable"),
connection_generation = state.connection_generation,
relay_url,
"relay reconnected"
);
let handshake_ok = process_handshake_buffer(
ws,
handshake_buffer,
Expand Down Expand Up @@ -3842,16 +3926,24 @@ async fn do_connect(
relay_url: &str,
keys: &Keys,
auth_tag: Option<&nostr::Tag>,
) -> Result<(WsStream, VecDeque<RelayMessage>), RelayError> {
) -> Result<(WsStream, VecDeque<RelayMessage>, Option<String>), RelayError> {
let parsed = relay_url
.parse::<url::Url>()
.map_err(|e| RelayError::Http(format!("invalid relay URL: {e}")))?;

let (ws, _response) = tokio::time::timeout(CONNECT_TIMEOUT, connect_async(parsed.as_str()))
let (ws, response) = tokio::time::timeout(CONNECT_TIMEOUT, connect_async(parsed.as_str()))
.await
.map_err(|_| RelayError::ConnectionClosed)? // timeout → treat as connection failure
.map_err(|e| RelayError::WebSocket(Box::new(e)))?;
debug!("connected to relay at {relay_url}");
let relay_connection_id = response
.headers()
.get(RELAY_CONNECTION_ID_HEADER)
.and_then(|value| value.to_str().ok())
.map(str::to_owned);
info!(
relay_connection_id = relay_connection_id.as_deref().unwrap_or("unavailable"),
relay_url, "connected to relay"
);

let mut ws = ws;
let mut buffer: VecDeque<RelayMessage> = VecDeque::new();
Expand All @@ -3874,7 +3966,7 @@ async fn do_connect(
};

debug!("NIP-42 authentication successful (event {event_id})");
Ok((ws, buffer))
Ok((ws, buffer, relay_connection_id))
}

/// Wait for an `AUTH` challenge from the relay, buffering any other messages.
Expand Down Expand Up @@ -4400,6 +4492,21 @@ mod tests {
);
}

#[test]
fn connection_diagnostics_advance_generation_and_replace_relay_id() {
let mut state = BgState::new();

state.note_connection(Some("first".to_string()));
let first_connected_at = state.connected_at;
assert_eq!(state.connection_generation, 1);
assert_eq!(state.relay_connection_id.as_deref(), Some("first"));

state.note_connection(Some("second".to_string()));
assert_eq!(state.connection_generation, 2);
assert_eq!(state.relay_connection_id.as_deref(), Some("second"));
assert!(state.connected_at >= first_connected_at);
}

#[tokio::test]
async fn fresh_reconnect_preserves_gate_until_pending_replay_resumes() {
let (mut client, mut server) = test_ws_pair().await;
Expand Down
1 change: 1 addition & 0 deletions crates/buzz-pubsub/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ tracing = { workspace = true }
thiserror = { workspace = true }
nostr = { workspace = true }
futures-util = { workspace = true }
metrics = { workspace = true }

[dev-dependencies]
tokio = { workspace = true }
23 changes: 23 additions & 0 deletions crates/buzz-pubsub/src/subscriber.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,16 +56,23 @@ pub(crate) async fn run_subscriber(
// established and ran successfully, so reset backoff to the initial
// value — a brief Redis restart should reconnect quickly.
backoff_secs = BACKOFF_INITIAL_SECS;
metrics::gauge!("buzz_pubsub_connected").set(0.0);
metrics::counter!("buzz_pubsub_disconnects_total", "reason" => "clean")
.increment(1);
tracing::warn!("Redis pub/sub stream ended (clean disconnect) — reconnecting in {backoff_secs}s");
}
Err(e) => {
metrics::gauge!("buzz_pubsub_connected").set(0.0);
metrics::counter!("buzz_pubsub_disconnects_total", "reason" => "error")
.increment(1);
tracing::error!("Redis pub/sub error: {e} — reconnecting in {backoff_secs}s");
}
}

tokio::time::sleep(Duration::from_secs(backoff_secs)).await;
backoff_secs = (backoff_secs * 2).min(BACKOFF_MAX_SECS);

metrics::counter!("buzz_pubsub_reconnect_attempts_total").increment(1);
tracing::info!("Attempting to reconnect to Redis pub/sub...");
}
}
Expand Down Expand Up @@ -101,6 +108,9 @@ async fn connect_and_subscribe(
topic_count = active_topics.len(),
"Redis pub/sub subscriber connected with dynamic scoped subscriptions"
);
metrics::gauge!("buzz_pubsub_connected").set(1.0);
metrics::counter!("buzz_pubsub_connections_total").increment(1);
record_topic_gauges(&desired_topics, &active_topics).await;

loop {
tokio::select! {
Expand All @@ -110,13 +120,15 @@ async fn connect_and_subscribe(
let channel = topic.redis_channel();
if active_topics.insert(channel.clone()) {
sink.subscribe(&channel).await?;
record_topic_gauges(&desired_topics, &active_topics).await;
}
}
SubscriptionCommand::UnsubscribeIfIdle(topic) => {
if desired_refcount(&desired_topics, topic).await == 0 {
let channel = topic.redis_channel();
if active_topics.remove(&channel) {
sink.unsubscribe(&channel).await?;
record_topic_gauges(&desired_topics, &active_topics).await;
}
}
}
Expand Down Expand Up @@ -171,6 +183,17 @@ async fn connect_and_subscribe(
}
}

async fn record_topic_gauges(desired_topics: &DesiredTopics, active_topics: &HashSet<String>) {
let desired_count = desired_topics
.lock()
.await
.values()
.filter(|count| **count > 0)
.count();
metrics::gauge!("buzz_pubsub_topics_desired").set(desired_count as f64);
metrics::gauge!("buzz_pubsub_topics_active").set(active_topics.len() as f64);
}

async fn desired_refcount(desired_topics: &DesiredTopics, topic: EventTopicKey) -> usize {
desired_topics
.lock()
Expand Down
2 changes: 1 addition & 1 deletion crates/buzz-relay/src/connection.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -125,8 +125,8 @@ pub async fn handle_connection(
state: Arc<AppState>,
addr: SocketAddr,
tenant: TenantContext,
conn_id: Uuid,
) {
let conn_id = Uuid::new_v4();
let cancel = CancellationToken::new();
let community_id = tenant.community();
let registry = Arc::clone(&state.community_connections);
Expand Down
Loading