From ed34600e5b4ab478f9de3a135973281b68e42c18 Mon Sep 17 00:00:00 2001 From: tlongwell-block <109685178+tlongwell-block@users.noreply.github.com> Date: Sun, 15 Mar 2026 14:44:26 -0400 Subject: [PATCH] fix: WebSocket hardening + agent reliability (8 WS bugs, member-only discovery, mention tag fix) (#67) --- crates/sprout-acp/README.md | 2 +- crates/sprout-acp/src/config.rs | 4 +- crates/sprout-acp/src/relay.rs | 392 ++++- crates/sprout-db/src/channel.rs | 19 +- crates/sprout-db/src/lib.rs | 4 +- crates/sprout-mcp/Cargo.toml | 4 + crates/sprout-mcp/src/lib.rs | 24 +- crates/sprout-mcp/src/relay_client.rs | 1486 +++++++++++++---- crates/sprout-relay/src/api/agents.rs | 2 +- crates/sprout-relay/src/api/channels.rs | 4 +- crates/sprout-relay/src/connection.rs | 111 +- crates/sprout-relay/src/state.rs | 182 +- .../features/messages/ui/MessageComposer.tsx | 48 +- 13 files changed, 1798 insertions(+), 484 deletions(-) diff --git a/crates/sprout-acp/README.md b/crates/sprout-acp/README.md index 4d28c9bd7d..9d501bf60e 100644 --- a/crates/sprout-acp/README.md +++ b/crates/sprout-acp/README.md @@ -40,7 +40,7 @@ This prints an `nsec1...` private key and an API token. **Save both immediately The harness discovers channels by querying the relay with the agent's authenticated identity. -**Open channels** (the default for local dev) are accessible to any authenticated pubkey — no extra setup needed. Just start the harness and it will find and subscribe to all open channels. +By default, the harness discovers only channels the agent is a **member** of (`GET /api/channels?member=true`). When the agent is added to a new channel, the membership notification subscription auto-subscribes to it. **Private channels** require explicit membership. The relay doesn't yet have a REST/event API for managing channel members — this is a known gap. For now, use `create_channel` via the Sprout MCP tools to create new channels (the creator is automatically a member). diff --git a/crates/sprout-acp/src/config.rs b/crates/sprout-acp/src/config.rs index 7b984f03a2..c359848951 100644 --- a/crates/sprout-acp/src/config.rs +++ b/crates/sprout-acp/src/config.rs @@ -142,7 +142,7 @@ pub struct CliArgs { #[arg(long, env = "SPROUT_ACP_CONFIG", default_value = "./sprout-acp.toml")] pub config: PathBuf, - #[arg(long, env = "SPROUT_ACP_DEDUP", default_value = "drop", value_enum)] + #[arg(long, env = "SPROUT_ACP_DEDUP", default_value = "queue", value_enum)] pub dedup: DedupMode, #[arg(long, env = "SPROUT_ACP_NO_IGNORE_SELF")] @@ -566,7 +566,7 @@ mod tests { system_prompt: None, initial_message: None, subscribe_mode: mode, - dedup_mode: DedupMode::Drop, + dedup_mode: DedupMode::Queue, ignore_self: true, kinds_override: None, channels_override: None, diff --git a/crates/sprout-acp/src/relay.rs b/crates/sprout-acp/src/relay.rs index 67be64276b..ce3eb6d1cb 100644 --- a/crates/sprout-acp/src/relay.rs +++ b/crates/sprout-acp/src/relay.rs @@ -24,16 +24,28 @@ use std::time::Duration; // ─── Named constants ────────────────────────────────────────────────────────── -/// Capacity of the event channel from background task to harness. -const EVENT_CHANNEL_CAPACITY: usize = 256; +/// Default capacity of the event channel from background task to harness. +/// Override with `SPROUT_ACP_EVENT_BUFFER` env var at startup. +const EVENT_CHANNEL_CAPACITY_DEFAULT: usize = 256; /// Capacity of the command channel from harness to background task. const CMD_CHANNEL_CAPACITY: usize = 64; + +/// Read the event channel capacity from the environment, falling back to the +/// compiled-in default. Parsed once at call-site (connect time). +fn event_channel_capacity() -> usize { + std::env::var("SPROUT_ACP_EVENT_BUFFER") + .ok() + .and_then(|v| v.parse::().ok()) + .unwrap_or(EVENT_CHANNEL_CAPACITY_DEFAULT) +} /// Maximum number of seen event IDs before the dedup set is cleared. const SEEN_ID_LIMIT: usize = 12_000; /// Seconds subtracted from `since` on resubscribe to tolerate clock skew. const SINCE_SKEW_SECS: u64 = 5; /// Timeout for the NIP-42 auth handshake steps. const AUTH_TIMEOUT: Duration = Duration::from_secs(5); +/// Timeout for the TCP + WebSocket handshake in `do_connect`. +const CONNECT_TIMEOUT: Duration = Duration::from_secs(10); use futures_util::{SinkExt, StreamExt}; use nostr::{Event, EventBuilder, Keys, Kind, Tag, Url as NostrUrl}; @@ -190,7 +202,7 @@ impl HarnessRelay { // Perform the initial connection and auth handshake. let (ws, _buffer) = do_connect(relay_url, keys, api_token).await?; - let (event_tx, event_rx) = mpsc::channel::>(EVENT_CHANNEL_CAPACITY); + let (event_tx, event_rx) = mpsc::channel::>(event_channel_capacity()); let (cmd_tx, cmd_rx) = mpsc::channel::(CMD_CHANNEL_CAPACITY); let bg_keys = keys.clone(); @@ -227,10 +239,10 @@ impl HarnessRelay { }) } - /// Discover channels the harness is a member of via `GET /api/channels`. + /// Discover channels the agent is a member of via `GET /api/channels?member=true`. pub async fn discover_channels(&self) -> Result, RelayError> { let http_url = relay_ws_to_http(&self.relay_url); - let url = format!("{http_url}/api/channels"); + let url = format!("{http_url}/api/channels?member=true"); let builder = self.http.get(&url); let builder = apply_auth(builder, &self.api_token, &self.keys); @@ -359,6 +371,11 @@ struct BgState { membership_last_seen: Option, /// Whether the membership notification subscription is active. membership_sub_active: bool, + /// Oldest dropped channel-event timestamp per channel, keyed by channel_id. + /// Mirrors `membership_dropped_since` but for ordinary channel events. + /// On reconnect resubscribe, `since` = min(last_seen, channel_dropped_since). + /// Cleared per-channel after a successful resubscribe. + channel_dropped_since: HashMap, } impl BgState { @@ -371,6 +388,7 @@ impl BgState { membership_dropped_since: None, membership_last_seen: None, membership_sub_active: false, + channel_dropped_since: HashMap::new(), } } @@ -423,9 +441,10 @@ async fn run_background_task( tokio::select! { // ── Incoming WebSocket message ──────────────────────────────────── raw = ws.next() => { - match raw { + // Determine if the socket is lost. + let socket_lost = match raw { Some(Ok(msg)) => { - if !handle_ws_message( + !handle_ws_message( msg, &mut ws, &event_tx, @@ -435,41 +454,62 @@ async fn run_background_task( api_token.as_deref(), ) .await - { - // handle_ws_message returns false on connection loss. - // Signal the caller and wait for a Reconnect command. - let _ = event_tx.send(None).await; - wait_for_reconnect( - &mut ws, - &mut cmd_rx, - &mut state, - &keys, - &relay_url, - api_token.as_deref(), - &agent_pubkey_hex, - false, // skip_drain: wait for Reconnect command - ) - .await; - } } Some(Err(e)) => { warn!("WebSocket error in background task: {e}"); - let _ = event_tx.send(None).await; - wait_for_reconnect( - &mut ws, - &mut cmd_rx, - &mut state, - &keys, - &relay_url, - api_token.as_deref(), - &agent_pubkey_hex, - false, - ) - .await; + true } None => { debug!("WebSocket stream ended"); - let _ = event_tx.send(None).await; + true + } + }; + + if socket_lost { + // Signal the caller, then attempt autonomous reconnect. + let _ = event_tx.send(None).await; + let reconnected = try_autonomous_reconnect( + &mut ws, + &mut state, + &keys, + &relay_url, + api_token.as_deref(), + &agent_pubkey_hex, + ) + .await; + if reconnected { + // Drain stale Reconnect commands the caller queued after + // seeing the None sentinel. Process other commands normally + // so Subscribe/Unsubscribe/SubscribeMembership aren't lost. + while let Ok(cmd) = cmd_rx.try_recv() { + match cmd { + RelayCommand::Reconnect => { + debug!("drained stale Reconnect command after autonomous reconnect"); + } + RelayCommand::Shutdown => { + debug!("shutdown received during reconnect drain"); + return; + } + RelayCommand::Subscribe { channel_id, filter } => { + send_subscribe(&mut ws, &state, channel_id, &agent_pubkey_hex, None, &filter).await; + state.active_subscriptions.insert(channel_id, channel_sub_id(channel_id)); + state.active_filters.insert(channel_id, filter); + } + RelayCommand::Unsubscribe { channel_id } => { + if let Some(sub_id) = state.active_subscriptions.remove(&channel_id) { + let msg = json!(["CLOSE", sub_id]); + if let Ok(text) = serde_json::to_string(&msg) { + let _ = ws.send(Message::Text(text.into())).await; + } + } + } + RelayCommand::SubscribeMembership => { + let _ = send_membership_subscribe(&mut ws, &agent_pubkey_hex, None).await; + state.membership_sub_active = true; + } + } + } + } else { wait_for_reconnect( &mut ws, &mut cmd_rx, @@ -603,6 +643,8 @@ async fn handle_ws_message( Err(mpsc::error::TrySendError::Closed(_)) => return false, } } else if let Some(channel_id) = channel_id_from_sub_id(&subscription_id) { + let ts = event.created_at.as_u64(); + let event_id_hex = event.id.to_hex(); if state.record_event(channel_id, &event) { let sprout_event = SproutEvent { channel_id, @@ -611,7 +653,21 @@ async fn handle_ws_message( match event_tx.try_send(Some(sprout_event)) { Ok(()) => {} Err(mpsc::error::TrySendError::Full(_)) => { - warn!("event channel full — dropping event for channel {channel_id}"); + // Remove from dedup set so the replayed event + // won't be rejected as a duplicate after reconnect. + state.seen_ids.remove(&event_id_hex); + // Track the oldest dropped timestamp so reconnect + // replay starts early enough to re-deliver it. + state + .channel_dropped_since + .entry(channel_id) + .and_modify(|d| *d = (*d).min(ts)) + .or_insert(ts); + warn!( + channel_id = %channel_id, + ts, + "event channel full — dropping event for channel {channel_id} — will replay from {ts} on reconnect" + ); } Err(mpsc::error::TrySendError::Closed(_)) => { // Receiver dropped — shut down. @@ -673,6 +729,115 @@ async fn handle_ws_message( } } +/// Resubscribe all active channels and membership notifications after a +/// successful reconnect. Computes `since = min(last_seen, channel_dropped_since)` +/// per channel, and only clears the drop tracker when the REQ is confirmed sent. +async fn resubscribe_after_reconnect( + ws: &mut WsStream, + state: &mut BgState, + agent_pubkey_hex: &str, +) { + let channels: Vec = state.active_subscriptions.keys().copied().collect(); + if !channels.is_empty() { + info!( + "resubscribing to {} channel(s) after reconnect", + channels.len() + ); + for channel_id in channels { + let last_seen = state.last_seen.get(&channel_id).copied(); + let dropped = state.channel_dropped_since.get(&channel_id).copied(); + let since = match (last_seen, dropped) { + (Some(l), Some(d)) => Some(l.min(d)), + (Some(l), None) => Some(l), + (None, Some(d)) => Some(d), + (None, None) => None, + }; + let filter = state + .active_filters + .get(&channel_id) + .cloned() + .unwrap_or(ChannelFilter { + kinds: None, + require_mention: false, + }); + let sent = + send_subscribe(ws, state, channel_id, agent_pubkey_hex, since, &filter).await; + if sent { + state.channel_dropped_since.remove(&channel_id); + } + } + } + + if state.membership_sub_active { + let replay_since = match (state.membership_dropped_since, state.membership_last_seen) { + (Some(d), Some(l)) => Some(d.min(l)), + (Some(d), None) => Some(d), + (None, Some(l)) => Some(l), + (None, None) => None, + }; + let sent = send_membership_subscribe(ws, agent_pubkey_hex, replay_since).await; + if sent { + state.membership_dropped_since = None; + } + } +} + +/// Attempt autonomous reconnect on socket loss — 3 attempts with 1s→2s→4s backoff. +/// +/// If any attempt succeeds, resubscribes all active channels and membership +/// notifications via [`resubscribe_after_reconnect`] and returns `true`. +/// If all 3 attempts fail, returns `false` so the caller can fall back to +/// `wait_for_reconnect` (which blocks until the caller sends a `Reconnect`). +#[allow(clippy::too_many_arguments)] +async fn try_autonomous_reconnect( + ws: &mut WsStream, + state: &mut BgState, + keys: &Keys, + relay_url: &str, + api_token: Option<&str>, + agent_pubkey_hex: &str, +) -> bool { + let backoffs = [ + Duration::from_secs(1), + Duration::from_secs(2), + Duration::from_secs(4), + ]; + + for (attempt, delay) in backoffs.iter().enumerate() { + info!( + "autonomous reconnect attempt {}/{} to {relay_url}…", + attempt + 1, + backoffs.len() + ); + match do_connect(relay_url, keys, api_token).await { + Ok((new_ws, _buffer)) => { + *ws = new_ws; + info!("autonomous reconnect succeeded (attempt {})", attempt + 1); + resubscribe_after_reconnect(ws, state, agent_pubkey_hex).await; + return true; + } + Err(e) => { + warn!( + "autonomous reconnect attempt {} failed: {e} — {}", + attempt + 1, + if attempt + 1 < backoffs.len() { + format!("retrying in {}s", delay.as_secs()) + } else { + "falling back to caller-driven reconnect".to_string() + } + ); + // Skip sleep on the final attempt — no point waiting before + // falling back to caller-driven reconnect. + if attempt + 1 < backoffs.len() { + tokio::time::sleep(*delay).await; + } + } + } + } + + false +} + /// Attempt reconnection with exponential backoff. Resubscribes all active /// channels with `since` filters on success. /// @@ -725,42 +890,7 @@ async fn wait_for_reconnect( Ok((new_ws, _buffer)) => { *ws = new_ws; info!("relay reconnected to {relay_url}"); - - // Resubscribe all active channels with `since` filter. - let channels: Vec = state.active_subscriptions.keys().copied().collect(); - if !channels.is_empty() { - info!("resubscribing to {} channel(s)", channels.len()); - for channel_id in channels { - let since = state.last_seen.get(&channel_id).copied(); - let filter = state.active_filters.get(&channel_id).cloned().unwrap_or( - ChannelFilter { - kinds: None, - require_mention: false, - }, - ); - send_subscribe(ws, state, channel_id, agent_pubkey_hex, since, &filter) - .await; - } - } - - // Resubscribe to membership notifications if active. - // Use the oldest of (dropped, last_seen) so dropped events are replayed. - if state.membership_sub_active { - let replay_since = - match (state.membership_dropped_since, state.membership_last_seen) { - (Some(d), Some(l)) => Some(d.min(l)), - (Some(d), None) => Some(d), - (None, Some(l)) => Some(l), - (None, None) => None, - }; - let sent = send_membership_subscribe(ws, agent_pubkey_hex, replay_since).await; - if sent { - // Only clear drop tracker if the REQ was actually sent — - // if send failed, retain it so the next reconnect retries. - state.membership_dropped_since = None; - } - } - + resubscribe_after_reconnect(ws, state, agent_pubkey_hex).await; return; } Err(e) => { @@ -782,6 +912,8 @@ async fn wait_for_reconnect( /// - `#h` is always included (channel-scoped subscription). /// - On first subscribe (`since` is `None`) adds `since=now` to avoid replaying /// history. On reconnect (`since` is `Some`) subtracts [`SINCE_SKEW_SECS`]. +/// +/// Returns `true` if the REQ was successfully written to the WebSocket. async fn send_subscribe( ws: &mut WsStream, _state: &BgState, @@ -789,7 +921,7 @@ async fn send_subscribe( agent_pubkey_hex: &str, since: Option, filter: &ChannelFilter, -) { +) -> bool { let sub_id = channel_sub_id(channel_id); let mut req_filter = serde_json::Map::new(); @@ -821,10 +953,8 @@ async fn send_subscribe( let req = json!(["REQ", sub_id, Value::Object(req_filter)]); match serde_json::to_string(&req) { - Ok(text) => { - if let Err(e) = ws.send(Message::Text(text.into())).await { - warn!("failed to send REQ for channel {channel_id}: {e}"); - } else { + Ok(text) => match ws.send(Message::Text(text.into())).await { + Ok(()) => { debug!( "subscribed to channel {channel_id}{}", if since.is_some() { @@ -833,10 +963,16 @@ async fn send_subscribe( " (since=now)" } ); + true } - } + Err(e) => { + warn!("failed to send REQ for channel {channel_id}: {e}"); + false + } + }, Err(e) => { warn!("failed to serialize REQ for channel {channel_id}: {e}"); + false } } } @@ -1078,8 +1214,9 @@ async fn do_connect( .parse::() .map_err(|e| RelayError::Http(format!("invalid relay URL: {e}")))?; - let (ws, _response) = 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}"); @@ -1648,4 +1785,99 @@ mod tests { state.seen_ids.len() ); } + + // ── Bug 5: channel_dropped_since tracking ───────────────────────────────── + + /// Test 8: channel_dropped_since records the OLDEST dropped timestamp. + /// + /// Simulates the backpressure path directly on BgState: + /// - First drop at ts=1000 → entry is 1000 + /// - Second drop at ts=2000 (later) → entry stays 1000 (min) + /// - Third drop at ts=500 (earlier) → entry updates to 500 (min) + #[test] + fn acp_records_channel_dropped_since_on_backpressure() { + let mut state = BgState::new(); + let channel_id = Uuid::new_v4(); + + // Simulate the backpressure path: record ts=1000. + let ts1: u64 = 1_000; + state + .channel_dropped_since + .entry(channel_id) + .and_modify(|d| *d = (*d).min(ts1)) + .or_insert(ts1); + assert_eq!( + state.channel_dropped_since.get(&channel_id).copied(), + Some(1_000), + "first drop should record ts=1000" + ); + + // Later timestamp (2000) — entry should stay at 1000. + let ts2: u64 = 2_000; + state + .channel_dropped_since + .entry(channel_id) + .and_modify(|d| *d = (*d).min(ts2)) + .or_insert(ts2); + assert_eq!( + state.channel_dropped_since.get(&channel_id).copied(), + Some(1_000), + "later drop should not overwrite earlier timestamp" + ); + + // Earlier timestamp (500) — entry should update to 500. + let ts3: u64 = 500; + state + .channel_dropped_since + .entry(channel_id) + .and_modify(|d| *d = (*d).min(ts3)) + .or_insert(ts3); + assert_eq!( + state.channel_dropped_since.get(&channel_id).copied(), + Some(500), + "earlier drop should update entry to 500" + ); + } + + /// Test 9: reconnect since filter = min(last_seen, channel_dropped_since) - SINCE_SKEW_SECS. + /// + /// With last_seen=1000 and channel_dropped_since=900, the effective since + /// passed to send_subscribe should be min(1000, 900) - SINCE_SKEW_SECS = 895. + #[test] + fn acp_reconnect_uses_dropped_since_for_replay() { + let mut state = BgState::new(); + let channel_id = Uuid::new_v4(); + + // Set up state: last_seen=1000, channel_dropped_since=900. + state.last_seen.insert(channel_id, 1_000); + state.channel_dropped_since.insert(channel_id, 900); + + // Compute the since value the reconnect path would use. + let last_seen = state.last_seen.get(&channel_id).copied(); + let dropped = state.channel_dropped_since.get(&channel_id).copied(); + let since = match (last_seen, dropped) { + (Some(l), Some(d)) => Some(l.min(d)), + (Some(l), None) => Some(l), + (None, Some(d)) => Some(d), + (None, None) => None, + }; + + // The since passed to send_subscribe (which subtracts SINCE_SKEW_SECS internally). + assert_eq!(since, Some(900), "since should be min(1000, 900) = 900"); + + // After subtracting skew (as send_subscribe does), the REQ filter value is: + let req_since = since.unwrap().saturating_sub(SINCE_SKEW_SECS); + assert_eq!( + req_since, 895, + "REQ since filter should be 900 - {} = 895", + SINCE_SKEW_SECS + ); + + // Simulate clearing after resubscribe. + state.channel_dropped_since.remove(&channel_id); + assert!( + !state.channel_dropped_since.contains_key(&channel_id), + "channel_dropped_since should be cleared after resubscribe" + ); + } } diff --git a/crates/sprout-db/src/channel.rs b/crates/sprout-db/src/channel.rs index 8965d3fb66..ebe8b5817d 100644 --- a/crates/sprout-db/src/channel.rs +++ b/crates/sprout-db/src/channel.rs @@ -707,8 +707,20 @@ pub async fn get_accessible_channels( pool: &MySqlPool, pubkey: &[u8], visibility_filter: Option<&str>, + member_only: Option, ) -> Result> { - let base = r#" + // When `member_only` is `Some(true)`, restrict to channels where the user + // has an active membership (cm.channel_id IS NOT NULL). This is a strict + // subset of the default result set and is pushed into SQL so the LIMIT 1000 + // applies to the filtered set, not the pre-filter set. + let membership_clause = if member_only == Some(true) { + "AND cm.channel_id IS NOT NULL" + } else { + "AND (c.visibility = 'open' OR cm.channel_id IS NOT NULL)" + }; + + let base = format!( + r#" SELECT DISTINCT c.id, c.name, c.channel_type, c.visibility, c.description, c.canvas, c.created_by, c.created_at, c.updated_at, c.archived_at, c.deleted_at, c.nip29_group_id, c.topic_required, c.max_members, @@ -719,8 +731,9 @@ pub async fn get_accessible_channels( LEFT JOIN channel_members cm ON c.id = cm.channel_id AND cm.pubkey = ? AND cm.removed_at IS NULL WHERE c.deleted_at IS NULL - AND (c.visibility = 'open' OR cm.channel_id IS NOT NULL) - "#; + {membership_clause} + "# + ); let sql = if visibility_filter.is_some() { format!("{base} AND c.visibility = ?\n ORDER BY FIELD(c.channel_type, 'stream', 'forum', 'dm'), c.name\n LIMIT 1000") diff --git a/crates/sprout-db/src/lib.rs b/crates/sprout-db/src/lib.rs index 365fc1aafe..740c7d52f5 100644 --- a/crates/sprout-db/src/lib.rs +++ b/crates/sprout-db/src/lib.rs @@ -403,12 +403,14 @@ impl Db { /// /// If `visibility_filter` is `Some("open")` or `Some("private")`, only channels /// with that visibility are returned. + /// If `member_only` is `Some(true)`, only channels the user is a member of are returned. pub async fn get_accessible_channels( &self, pubkey: &[u8], visibility_filter: Option<&str>, + member_only: Option, ) -> Result> { - channel::get_accessible_channels(&self.pool, pubkey, visibility_filter).await + channel::get_accessible_channels(&self.pool, pubkey, visibility_filter, member_only).await } /// Returns all bot-role members with aggregated channel names. diff --git a/crates/sprout-mcp/Cargo.toml b/crates/sprout-mcp/Cargo.toml index 21b38923d6..6d8482eecf 100644 --- a/crates/sprout-mcp/Cargo.toml +++ b/crates/sprout-mcp/Cargo.toml @@ -41,3 +41,7 @@ tracing-subscriber = { workspace = true } thiserror = { workspace = true } anyhow = { workspace = true } url = { workspace = true } + +[dev-dependencies] +tokio = { workspace = true, features = ["test-util"] } +tokio-tungstenite = { workspace = true } diff --git a/crates/sprout-mcp/src/lib.rs b/crates/sprout-mcp/src/lib.rs index 5e0106a78e..2fc09292a4 100644 --- a/crates/sprout-mcp/src/lib.rs +++ b/crates/sprout-mcp/src/lib.rs @@ -37,11 +37,27 @@ //! `SPROUT_API_TOKEN` is set the token is embedded in the auth event tags so the relay can //! verify the agent's API permissions. //! -//! ## WebSocket Reconnection +//! ## WebSocket Connection Management //! -//! [`relay_client::RelayClient`] supports automatic reconnection with exponential backoff -//! (1 s → 2 s → 4 s → … → 30 s cap). After reconnecting it re-authenticates via NIP-42 and -//! resubmits all subscriptions that were active at the time of the disconnect. +//! [`relay_client::RelayClient`] uses a background tokio task that owns the WebSocket +//! connection. The background task: +//! +//! - Responds to Ping frames immediately — preventing relay disconnects during long LLM turns +//! - Handles mid-session NIP-42 AUTH challenges automatically +//! - Reconnects with exponential backoff (1 s → 2 s → 4 s → … → 30 s cap) on any +//! connection loss, without any action required from the caller +//! - Re-authenticates via NIP-42 after each reconnect +//! - Replays all active subscriptions after reconnect +//! +//! ```text +//! RelayClient (Clone) +//! ├── cmd_tx: mpsc::Sender ← send_event / subscribe / close +//! └── bg_handle: JoinHandle<()> +//! └── run_background_task() +//! ├── ws.next() → handle_ws_message() // Ping→Pong, AUTH→respond, OK→resolve +//! ├── cmd_rx → handle_command() // SendEvent, Subscribe, Close +//! └── tick → expire_timed_out() // 10s timeouts +//! ``` //! //! ## Available Tools //! diff --git a/crates/sprout-mcp/src/relay_client.rs b/crates/sprout-mcp/src/relay_client.rs index ed1fed5dfe..943f5f75ec 100644 --- a/crates/sprout-mcp/src/relay_client.rs +++ b/crates/sprout-mcp/src/relay_client.rs @@ -1,15 +1,27 @@ -use std::collections::{HashMap, VecDeque}; -use std::sync::Arc; +use std::collections::HashMap; use std::time::Duration; use futures_util::{SinkExt, StreamExt}; use nostr::{Event, EventBuilder, Filter, Keys, Kind, Tag, Url}; use serde_json::{json, Value}; use thiserror::Error; -use tokio::sync::Mutex; -use tokio::time::timeout; +use tokio::sync::{mpsc, oneshot}; +use tokio::task::JoinHandle; use tokio_tungstenite::{connect_async, tungstenite::Message, MaybeTlsStream, WebSocketStream}; -use tracing::debug; +use tracing::{debug, warn}; + +// ── Timeouts ────────────────────────────────────────────────────────────────── + +/// How long to wait for an OK acknowledgement after sending an event. +const SEND_EVENT_TIMEOUT: Duration = Duration::from_secs(10); +/// How long to wait for EOSE after sending a REQ. +const SUBSCRIBE_TIMEOUT: Duration = Duration::from_secs(10); +/// Timeout for the TCP + WebSocket handshake in `do_connect`. +const CONNECT_TIMEOUT: Duration = Duration::from_secs(10); +/// Capacity of the command channel. +const CMD_CHANNEL_CAPACITY: usize = 64; + +// ── Public error type ───────────────────────────────────────────────────────── /// Errors that can occur when communicating with a Sprout relay. #[derive(Debug, Error)] @@ -57,6 +69,8 @@ impl From for RelayClientError { } } +// ── Public relay message type ───────────────────────────────────────────────── + /// A message received from a Nostr relay. #[derive(Debug, Clone)] pub enum RelayMessage { @@ -104,370 +118,625 @@ pub struct OkResponse { pub message: String, } +// ── Internal types ──────────────────────────────────────────────────────────── + type WsStream = WebSocketStream>; -struct Inner { - ws: WsStream, - buffer: VecDeque, - pending_challenge: Option, +/// Commands sent from `RelayClient` to the background WebSocket task. +#[allow(clippy::large_enum_variant)] +enum RelayCommand { + SendEvent { + event: Event, + reply: oneshot::Sender>, + }, + Subscribe { + sub_id: String, + filters: Vec, + reply: oneshot::Sender, RelayClientError>>, + }, + CloseSubscription { + sub_id: String, + reply: oneshot::Sender>, + }, + Shutdown, } -impl Inner { - async fn send_raw(&mut self, value: &Value) -> Result<(), RelayClientError> { - let text = serde_json::to_string(value)?; - self.ws.send(Message::Text(text.into())).await?; +/// A subscription waiting for EOSE. +struct PendingSubscription { + events: Vec, + reply: oneshot::Sender, RelayClientError>>, + deadline: tokio::time::Instant, +} + +/// State owned exclusively by the background task. +struct BgState { + /// Active subscriptions: sub_id → filters (for reconnect replay). + active_subscriptions: HashMap>, + /// Pending OK waiters: event_id → (reply, deadline). + pending_ok: HashMap< + String, + ( + oneshot::Sender>, + tokio::time::Instant, + ), + >, + /// Pending EOSE collectors: sub_id → collector. + pending_eose: HashMap, +} + +impl BgState { + fn new() -> Self { + Self { + active_subscriptions: HashMap::new(), + pending_ok: HashMap::new(), + pending_eose: HashMap::new(), + } + } + + /// Resolve all pending operations with `ConnectionClosed` (called on reconnect). + fn cancel_pending(&mut self) { + for (_, (reply, _)) in self.pending_ok.drain() { + let _ = reply.send(Err(RelayClientError::ConnectionClosed)); + } + for (_, sub) in self.pending_eose.drain() { + let _ = sub.reply.send(Err(RelayClientError::ConnectionClosed)); + } + } + + /// Expire any pending operations whose deadline has passed. + fn expire_timed_out(&mut self) { + let now = tokio::time::Instant::now(); + + let expired_ok: Vec = self + .pending_ok + .iter() + .filter(|(_, (_, dl))| now >= *dl) + .map(|(k, _)| k.clone()) + .collect(); + for k in expired_ok { + if let Some((reply, _)) = self.pending_ok.remove(&k) { + let _ = reply.send(Err(RelayClientError::Timeout)); + } + } + + let expired_eose: Vec = self + .pending_eose + .iter() + .filter(|(_, sub)| now >= sub.deadline) + .map(|(k, _)| k.clone()) + .collect(); + for k in expired_eose { + if let Some(sub) = self.pending_eose.remove(&k) { + let _ = sub.reply.send(Err(RelayClientError::Timeout)); + } + } + } +} + +// ── Background task ─────────────────────────────────────────────────────────── + +/// Perform a single NIP-42 connection + auth handshake. +/// Returns the authenticated WebSocket stream on success. +async fn do_connect( + relay_url: &str, + keys: &Keys, + api_token: Option<&str>, +) -> Result { + let parsed = relay_url + .parse::() + .map_err(|e| RelayClientError::Url(e.to_string()))?; + + let (mut ws, _) = tokio::time::timeout(CONNECT_TIMEOUT, connect_async(parsed.as_str())) + .await + .map_err(|_| RelayClientError::ConnectionClosed)? // timeout → treat as connection failure + .map_err(RelayClientError::WebSocket)?; + + debug!("connected to relay at {relay_url}"); + + // Wait for AUTH challenge (5s timeout). + let challenge = wait_for_auth_challenge(&mut ws, Duration::from_secs(5)).await?; + + let auth_event = build_auth_event(&challenge, relay_url, keys, api_token)?; + let event_id = auth_event.id.to_hex(); + debug!("sending AUTH event {event_id}"); + let auth_msg = serde_json::to_string(&json!(["AUTH", auth_event]))?; + ws.send(Message::Text(auth_msg.into())).await?; + + let ok = wait_for_ok(&mut ws, &event_id, Duration::from_secs(5)).await?; + if !ok.accepted { + return Err(RelayClientError::AuthFailed(ok.message)); + } + + debug!("NIP-42 authentication successful"); + Ok(ws) +} + +/// Wait for an AUTH challenge frame, responding to Pings along the way. +async fn wait_for_auth_challenge( + ws: &mut WsStream, + timeout_dur: Duration, +) -> Result { + let deadline = tokio::time::Instant::now() + timeout_dur; + loop { + let remaining = deadline + .checked_duration_since(tokio::time::Instant::now()) + .unwrap_or(Duration::ZERO); + if remaining.is_zero() { + return Err(RelayClientError::NoAuthChallenge); + } + let raw = tokio::time::timeout(remaining, ws.next()) + .await + .map_err(|_| RelayClientError::NoAuthChallenge)? + .ok_or(RelayClientError::ConnectionClosed)? + .map_err(RelayClientError::WebSocket)?; + match raw { + Message::Text(text) => { + if let RelayMessage::Auth { challenge } = parse_relay_message(&text)? { + return Ok(challenge); + } + } + Message::Ping(data) => { + ws.send(Message::Pong(data)).await?; + } + Message::Close(_) => return Err(RelayClientError::ConnectionClosed), + _ => {} + } + } +} + +/// Wait for an OK frame matching `event_id`, responding to Pings along the way. +async fn wait_for_ok( + ws: &mut WsStream, + event_id: &str, + timeout_dur: Duration, +) -> Result { + let deadline = tokio::time::Instant::now() + timeout_dur; + loop { + let remaining = deadline + .checked_duration_since(tokio::time::Instant::now()) + .unwrap_or(Duration::ZERO); + if remaining.is_zero() { + return Err(RelayClientError::Timeout); + } + let raw = tokio::time::timeout(remaining, ws.next()) + .await + .map_err(|_| RelayClientError::Timeout)? + .ok_or(RelayClientError::ConnectionClosed)? + .map_err(RelayClientError::WebSocket)?; + match raw { + Message::Text(text) => match parse_relay_message(&text)? { + RelayMessage::Ok(ok) if ok.event_id == event_id => return Ok(ok), + _ => {} // discard other messages during handshake + }, + Message::Ping(data) => { + ws.send(Message::Pong(data)).await?; + } + Message::Close(_) => return Err(RelayClientError::ConnectionClosed), + _ => {} + } + } +} + +/// Build a NIP-42 AUTH event for the given challenge. +#[allow(clippy::result_large_err)] +fn build_auth_event( + challenge: &str, + relay_url: &str, + keys: &Keys, + api_token: Option<&str>, +) -> Result { + let relay_nostr_url: Url = relay_url + .parse() + .map_err(|e: url::ParseError| RelayClientError::Url(e.to_string()))?; + if let Some(token) = api_token { + let tags = vec![ + Tag::parse(&["relay", relay_url]) + .map_err(|e| RelayClientError::EventBuilder(e.to_string()))?, + Tag::parse(&["challenge", challenge]) + .map_err(|e| RelayClientError::EventBuilder(e.to_string()))?, + Tag::parse(&["auth_token", token]) + .map_err(|e| RelayClientError::EventBuilder(e.to_string()))?, + ]; + Ok(EventBuilder::new(Kind::Authentication, "", tags).sign_with_keys(keys)?) + } else { + Ok(EventBuilder::auth(challenge, relay_nostr_url).sign_with_keys(keys)?) + } +} + +/// Send a NIP-42 AUTH response for a mid-session challenge. +/// +/// Fire-and-forget: we don't wait for the relay's OK. If the relay rejects +/// the re-auth it will close the connection, which triggers our reconnect logic. +async fn send_auth_response( + ws: &mut WsStream, + challenge: &str, + relay_url: &str, + keys: &Keys, + api_token: Option<&str>, +) { + let result: Result<(), RelayClientError> = async { + let auth_event = build_auth_event(challenge, relay_url, keys, api_token)?; + let msg = serde_json::to_string(&json!(["AUTH", auth_event]))?; + ws.send(Message::Text(msg.into())).await?; + debug!("sent AUTH response for mid-session challenge"); Ok(()) } - - // wait_for_auth_challenge, wait_for_ok, and collect_until_eose share a similar - // deadline-loop structure but differ in termination condition and what they do - // with interleaved messages, so they cannot be collapsed into a single helper. - async fn wait_for_auth_challenge( - &mut self, - timeout_dur: Duration, - ) -> Result { - if let Some(challenge) = self.pending_challenge.take() { - return Ok(challenge); - } - - if let Some(idx) = self - .buffer - .iter() - .position(|m| matches!(m, RelayMessage::Auth { .. })) - { - if let Some(RelayMessage::Auth { challenge }) = self.buffer.remove(idx) { - return Ok(challenge); - } - } - - let deadline = tokio::time::Instant::now() + timeout_dur; - - loop { - let remaining = deadline - .checked_duration_since(tokio::time::Instant::now()) - .unwrap_or(Duration::ZERO); - - if remaining.is_zero() { - return Err(RelayClientError::NoAuthChallenge); - } - - let raw = timeout(remaining, self.ws.next()) - .await - .map_err(|_| RelayClientError::NoAuthChallenge)? - .ok_or(RelayClientError::ConnectionClosed)? - .map_err(RelayClientError::WebSocket)?; - - match raw { - Message::Text(text) => { - let msg = parse_relay_message(&text)?; - match msg { - RelayMessage::Auth { challenge } => return Ok(challenge), - other => self.buffer.push_back(other), - } - } - Message::Ping(data) => { - self.ws.send(Message::Pong(data)).await?; - } - Message::Close(_) => return Err(RelayClientError::ConnectionClosed), - _ => {} - } - } + .await; + if let Err(e) = result { + warn!("failed to respond to mid-session AUTH challenge: {e}"); } +} - async fn wait_for_ok( - &mut self, - event_id: &str, - timeout_dur: Duration, - ) -> Result { - let deadline = tokio::time::Instant::now() + timeout_dur; - - if let Some(idx) = self - .buffer - .iter() - .position(|m| matches!(m, RelayMessage::Ok(ok) if ok.event_id == event_id)) - { - if let Some(RelayMessage::Ok(ok)) = self.buffer.remove(idx) { - return Ok(ok); - } - } - - loop { - let remaining = deadline - .checked_duration_since(tokio::time::Instant::now()) - .unwrap_or(Duration::ZERO); - - if remaining.is_zero() { - return Err(RelayClientError::Timeout); - } - - let raw = timeout(remaining, self.ws.next()) - .await - .map_err(|_| RelayClientError::Timeout)? - .ok_or(RelayClientError::ConnectionClosed)? - .map_err(RelayClientError::WebSocket)?; - - match raw { - Message::Text(text) => { - let msg = parse_relay_message(&text)?; - match msg { - RelayMessage::Ok(ok) if ok.event_id == event_id => return Ok(ok), - RelayMessage::Auth { ref challenge } => { - self.pending_challenge = Some(challenge.clone()); - self.buffer.push_back(msg); - } - other => self.buffer.push_back(other), - } +/// Handle a single WebSocket message in the background task. +/// +/// Returns `false` if the connection has been lost (Close frame or error). +async fn handle_ws_message( + msg: Message, + ws: &mut WsStream, + state: &mut BgState, + keys: &Keys, + relay_url: &str, + api_token: Option<&str>, +) -> bool { + match msg { + Message::Text(text) => { + let relay_msg = match parse_relay_message(&text) { + Ok(m) => m, + Err(e) => { + warn!("failed to parse relay message: {e}"); + return true; } - Message::Ping(data) => { - self.ws.send(Message::Pong(data)).await?; - } - Message::Close(_) => return Err(RelayClientError::ConnectionClosed), - _ => {} - } - } - } - - async fn collect_until_eose( - &mut self, - sub_id: &str, - timeout_dur: Duration, - ) -> Result, RelayClientError> { - let deadline = tokio::time::Instant::now() + timeout_dur; - let mut events = Vec::new(); - - let old_buffer = std::mem::take(&mut self.buffer); - let mut found_eose = false; - for msg in old_buffer { - if found_eose { - self.buffer.push_back(msg); - continue; - } - match msg { + }; + match relay_msg { RelayMessage::Event { subscription_id, event, - } if subscription_id == sub_id => { - events.push(*event); - } - RelayMessage::Eose { subscription_id } if subscription_id == sub_id => { - found_eose = true; - } - other => self.buffer.push_back(other), - } - } - if found_eose { - return Ok(events); - } - - loop { - let remaining = deadline - .checked_duration_since(tokio::time::Instant::now()) - .unwrap_or(Duration::ZERO); - - if remaining.is_zero() { - return Err(RelayClientError::Timeout); - } - - let raw = timeout(remaining, self.ws.next()) - .await - .map_err(|_| RelayClientError::Timeout)? - .ok_or(RelayClientError::ConnectionClosed)? - .map_err(RelayClientError::WebSocket)?; - - match raw { - Message::Text(text) => { - let msg = parse_relay_message(&text)?; - match msg { - RelayMessage::Event { - subscription_id, - event, - } if subscription_id == sub_id => { - events.push(*event); - } - RelayMessage::Eose { subscription_id } if subscription_id == sub_id => { - return Ok(events); - } - RelayMessage::Auth { ref challenge } => { - self.pending_challenge = Some(challenge.clone()); - self.buffer.push_back(msg); - } - other => self.buffer.push_back(other), + } => { + if let Some(sub) = state.pending_eose.get_mut(&subscription_id) { + sub.events.push(*event); + } else { + debug!("EVENT for unknown/completed subscription {subscription_id}"); } } - Message::Ping(data) => { - self.ws.send(Message::Pong(data)).await?; + RelayMessage::Ok(ok) => { + if let Some((reply, _)) = state.pending_ok.remove(&ok.event_id) { + let _ = reply.send(Ok(ok)); + } else { + debug!("OK for unknown event {}", ok.event_id); + } } - Message::Close(_) => return Err(RelayClientError::ConnectionClosed), - _ => {} + RelayMessage::Eose { subscription_id } => { + if let Some(sub) = state.pending_eose.remove(&subscription_id) { + let _ = sub.reply.send(Ok(sub.events)); + // One-shot subscription fulfilled — don't replay on reconnect. + state.active_subscriptions.remove(&subscription_id); + } else { + debug!("EOSE for unknown subscription {subscription_id}"); + } + } + RelayMessage::Closed { + subscription_id, + message, + } => { + warn!("subscription {subscription_id} closed by relay: {message}"); + state.active_subscriptions.remove(&subscription_id); + if let Some(sub) = state.pending_eose.remove(&subscription_id) { + let _ = sub.reply.send(Err(RelayClientError::ConnectionClosed)); + } + } + RelayMessage::Notice { message } => { + debug!("relay NOTICE: {message}"); + } + RelayMessage::Auth { challenge } => { + debug!("received mid-session AUTH challenge — re-authenticating"); + send_auth_response(ws, &challenge, relay_url, keys, api_token).await; + } + } + true + } + Message::Ping(data) => { + if let Err(e) = ws.send(Message::Pong(data)).await { + warn!("failed to send Pong: {e}"); + return false; + } + true + } + Message::Close(_) => { + debug!("relay sent Close frame"); + false + } + _ => true, + } +} + +/// Reconnect with backoff, cancel pending ops, then replay subscriptions. +/// +/// Returns `true` on successful reconnect, `false` if the task should exit +/// (Shutdown received or command channel closed during backoff). +/// +/// Processes commands during backoff sleeps so that Shutdown is honoured +/// promptly and new operations fail fast with `ConnectionClosed`. +async fn do_reconnect( + ws: &mut WsStream, + state: &mut BgState, + cmd_rx: &mut mpsc::Receiver, + keys: &Keys, + relay_url: &str, + api_token: Option<&str>, +) -> bool { + warn!("relay connection lost — reconnecting…"); + state.cancel_pending(); + + let mut delay = Duration::from_secs(1); + loop { + match do_connect(relay_url, keys, api_token).await { + Ok(new_ws) => { + tracing::info!("reconnected to relay at {relay_url}"); + *ws = new_ws; + + // Replay active subscriptions. + let subs: Vec<(String, Vec)> = state + .active_subscriptions + .iter() + .map(|(k, v)| (k.clone(), v.clone())) + .collect(); + for (sub_id, filters) in subs { + let mut msg: Vec = Vec::with_capacity(2 + filters.len()); + msg.push(json!("REQ")); + msg.push(json!(sub_id)); + for f in &filters { + match serde_json::to_value(f) { + Ok(v) => msg.push(v), + Err(e) => warn!("failed to serialize filter for {sub_id}: {e}"), + } + } + let text = match serde_json::to_string(&Value::Array(msg)) { + Ok(t) => t, + Err(e) => { + warn!("failed to serialize REQ for {sub_id}: {e}"); + continue; + } + }; + if let Err(e) = ws.send(Message::Text(text.into())).await { + warn!("failed to resubscribe to {sub_id}: {e}"); + } + } + return true; + } + Err(e) => { + warn!("reconnect failed: {e}, retrying in {delay:?}"); + // Wait for backoff delay while still processing commands. + tokio::select! { + _ = tokio::time::sleep(delay) => {} + cmd = cmd_rx.recv() => { + match cmd { + Some(RelayCommand::Shutdown) | None => { + debug!("shutdown during reconnect"); + state.cancel_pending(); + return false; + } + // Fail new operations immediately — we're disconnected. + Some(RelayCommand::SendEvent { reply, .. }) => { + let _ = reply.send(Err(RelayClientError::ConnectionClosed)); + } + Some(RelayCommand::Subscribe { reply, .. }) => { + let _ = reply.send(Err(RelayClientError::ConnectionClosed)); + } + Some(RelayCommand::CloseSubscription { reply, .. }) => { + let _ = reply.send(Err(RelayClientError::ConnectionClosed)); + } + } + } + } + delay = (delay * 2).min(Duration::from_secs(30)); } } } } +/// The main background task loop. +/// +/// Owns the WebSocket, responds to Pings, routes relay messages to pending +/// waiters, and handles reconnection transparently. +async fn run_background_task( + mut ws: WsStream, + mut cmd_rx: mpsc::Receiver, + keys: Keys, + relay_url: String, + api_token: Option, +) { + let mut state = BgState::new(); + // Ticker for expiring timed-out pending operations (~1s granularity). + let mut tick = tokio::time::interval(Duration::from_secs(1)); + tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + + loop { + tokio::select! { + // ── Incoming WebSocket message ──────────────────────────────────── + raw = ws.next() => { + let needs_reconnect = match raw { + Some(Ok(msg)) => { + !handle_ws_message( + msg, &mut ws, &mut state, &keys, &relay_url, api_token.as_deref(), + ).await + } + Some(Err(e)) => { warn!("WebSocket error: {e}"); true } + None => { debug!("WebSocket stream ended"); true } + }; + if needs_reconnect + && !do_reconnect(&mut ws, &mut state, &mut cmd_rx, &keys, &relay_url, api_token.as_deref()).await + { + return; // Shutdown received during reconnect + } + } + + // ── Command from RelayClient ────────────────────────────────────── + cmd = cmd_rx.recv() => { + match cmd { + Some(RelayCommand::SendEvent { event, reply }) => { + let event_id = event.id.to_hex(); + let msg = match serde_json::to_string(&json!(["EVENT", event])) { + Ok(t) => t, + Err(e) => { let _ = reply.send(Err(e.into())); continue; } + }; + if let Err(e) = ws.send(Message::Text(msg.into())).await { + let _ = reply.send(Err(RelayClientError::WebSocket(e))); + if !do_reconnect(&mut ws, &mut state, &mut cmd_rx, &keys, &relay_url, api_token.as_deref()).await { + return; + } + continue; + } + let deadline = tokio::time::Instant::now() + SEND_EVENT_TIMEOUT; + state.pending_ok.insert(event_id, (reply, deadline)); + } + + Some(RelayCommand::Subscribe { sub_id, filters, reply }) => { + let mut msg: Vec = Vec::with_capacity(2 + filters.len()); + msg.push(json!("REQ")); + msg.push(json!(sub_id)); + let mut ser_err: Option = None; + for f in &filters { + match serde_json::to_value(f) { + Ok(v) => msg.push(v), + Err(e) => { ser_err = Some(e); break; } + } + } + if let Some(e) = ser_err { + let _ = reply.send(Err(e.into())); + continue; + } + let text = match serde_json::to_string(&Value::Array(msg)) { + Ok(t) => t, + Err(e) => { let _ = reply.send(Err(e.into())); continue; } + }; + if let Err(e) = ws.send(Message::Text(text.into())).await { + let _ = reply.send(Err(RelayClientError::WebSocket(e))); + if !do_reconnect(&mut ws, &mut state, &mut cmd_rx, &keys, &relay_url, api_token.as_deref()).await { + return; + } + continue; + } + state.active_subscriptions.insert(sub_id.clone(), filters); + let deadline = tokio::time::Instant::now() + SUBSCRIBE_TIMEOUT; + state.pending_eose.insert(sub_id, PendingSubscription { + events: Vec::new(), + reply, + deadline, + }); + } + + Some(RelayCommand::CloseSubscription { sub_id, reply }) => { + state.active_subscriptions.remove(&sub_id); + if let Some(sub) = state.pending_eose.remove(&sub_id) { + let _ = sub.reply.send(Err(RelayClientError::ConnectionClosed)); + } + let msg = match serde_json::to_string(&json!(["CLOSE", sub_id])) { + Ok(t) => t, + Err(e) => { let _ = reply.send(Err(e.into())); continue; } + }; + if let Err(e) = ws.send(Message::Text(msg.into())).await { + let _ = reply.send(Err(RelayClientError::WebSocket(e))); + if !do_reconnect(&mut ws, &mut state, &mut cmd_rx, &keys, &relay_url, api_token.as_deref()).await { + return; + } + continue; + } + let _ = reply.send(Ok(())); + } + + Some(RelayCommand::Shutdown) | None => { + debug!("background task shutting down"); + state.cancel_pending(); + return; + } + } + } + + // ── Timeout ticker ──────────────────────────────────────────────── + _ = tick.tick() => { + state.expire_timed_out(); + } + } + } +} + +// ── Public client ───────────────────────────────────────────────────────────── + +/// Shared handle to the background task. When the last `Arc` clone drops, +/// the task is signalled to shut down and then aborted as a safety net. +struct BgTaskHandle { + cmd_tx: mpsc::Sender, + handle: JoinHandle<()>, +} + +impl Drop for BgTaskHandle { + fn drop(&mut self) { + let _ = self.cmd_tx.try_send(RelayCommand::Shutdown); + self.handle.abort(); + } +} + /// Clone-able WebSocket client for the Sprout relay. /// -/// All clones share the same underlying connection via `Arc>`. -/// Active subscriptions are tracked so they can be resubmitted after a reconnect. +/// Internally, a background tokio task owns the WebSocket connection. All +/// clones share the same command channel to that task. The background task: +/// - Responds to Ping frames immediately (prevents relay disconnect) +/// - Handles mid-session AUTH challenges automatically +/// - Reconnects with exponential backoff on connection loss +/// - Processes Shutdown commands even during reconnect backoff +/// - Replays active subscriptions after reconnect +/// +/// When the last clone is dropped, the background task is automatically +/// shut down via [`BgTaskHandle`]'s `Drop` implementation. #[derive(Clone)] pub struct RelayClient { - inner: Arc>, + /// Shared background task handle — Drop sends Shutdown + abort. + bg: std::sync::Arc, keys: Keys, /// WebSocket URL of the relay (e.g. "ws://localhost:3000"). relay_url: String, /// Shared reqwest client for REST API calls. http: reqwest::Client, /// Optional API token for Bearer auth on REST endpoints. - /// When present, REST calls send `Authorization: Bearer ` instead of `X-Pubkey`. api_token: Option, - /// Active subscriptions: sub_id → filters. Used to resubscribe after reconnect. - active_subscriptions: Arc>>>, } impl RelayClient { - /// Perform a single connection + NIP-42 authentication attempt. - /// Returns the authenticated `Inner` on success. - async fn try_connect( - relay_url: &str, - keys: &Keys, - api_token: Option<&str>, - ) -> Result { - let parsed = relay_url - .parse::() - .map_err(|e| RelayClientError::Url(e.to_string()))?; - - let (ws, _response) = connect_async(parsed.as_str()) - .await - .map_err(RelayClientError::WebSocket)?; - - debug!("connected to relay at {relay_url}"); - - let mut inner = Inner { - ws, - buffer: VecDeque::new(), - pending_challenge: None, - }; - - let challenge = inner - .wait_for_auth_challenge(Duration::from_secs(5)) - .await?; - - let relay_nostr_url: Url = relay_url - .parse() - .map_err(|e: url::ParseError| RelayClientError::Url(e.to_string()))?; - - let auth_event = if let Some(token) = api_token { - let tags = vec![ - Tag::parse(&["relay", relay_url]) - .map_err(|e| RelayClientError::EventBuilder(e.to_string()))?, - Tag::parse(&["challenge", &challenge]) - .map_err(|e| RelayClientError::EventBuilder(e.to_string()))?, - Tag::parse(&["auth_token", token]) - .map_err(|e| RelayClientError::EventBuilder(e.to_string()))?, - ]; - EventBuilder::new(Kind::Authentication, "", tags).sign_with_keys(keys)? - } else { - EventBuilder::auth(&challenge, relay_nostr_url).sign_with_keys(keys)? - }; - - let event_id = auth_event.id.to_hex(); - // Log only the event ID, never the full AUTH payload which may contain tokens. - debug!("sending AUTH event {event_id}"); - let msg = json!(["AUTH", auth_event]); - inner.send_raw(&msg).await?; - - let ok = inner.wait_for_ok(&event_id, Duration::from_secs(5)).await?; - - if !ok.accepted { - return Err(RelayClientError::AuthFailed(ok.message)); - } - - debug!("NIP-42 authentication successful"); - Ok(inner) - } - - /// Connect to the relay with exponential-backoff retry. + /// Connect to the relay and start the background task. /// - /// Attempts `try_connect` in a loop, doubling the delay on each failure - /// (1 s → 2 s → 4 s → … → 30 s max). Returns only when a connection - /// and NIP-42 auth handshake succeed. - async fn connect_with_retry(relay_url: &str, keys: &Keys, api_token: Option<&str>) -> Inner { - let mut delay = Duration::from_secs(1); - let max_delay = Duration::from_secs(30); - loop { - match Self::try_connect(relay_url, keys, api_token).await { - Ok(inner) => { - tracing::info!("connected to relay at {relay_url}"); - return inner; - } - Err(e) => { - tracing::warn!("connection failed: {e}, retrying in {delay:?}"); - tokio::time::sleep(delay).await; - delay = (delay * 2).min(max_delay); - } - } - } - } - - /// Connect to the relay (first connection; returns an error rather than retrying - /// so the caller can surface a startup failure immediately). + /// Performs the initial NIP-42 handshake synchronously so startup failures + /// are surfaced immediately. After that, reconnection is automatic. pub async fn connect( relay_url: &str, keys: &Keys, api_token: Option<&str>, ) -> Result { - let inner = Self::try_connect(relay_url, keys, api_token).await?; + let ws = do_connect(relay_url, keys, api_token).await?; + + let (cmd_tx, cmd_rx) = mpsc::channel(CMD_CHANNEL_CAPACITY); + + let bg_keys = keys.clone(); + let bg_relay_url = relay_url.to_string(); + let bg_api_token = api_token.map(|t| t.to_string()); + + let handle = tokio::spawn(async move { + run_background_task(ws, cmd_rx, bg_keys, bg_relay_url, bg_api_token).await; + }); Ok(Self { + bg: std::sync::Arc::new(BgTaskHandle { cmd_tx, handle }), keys: keys.clone(), relay_url: relay_url.to_string(), - // SAFETY: default builder with only timeout config cannot fail + // Default builder with only timeout config — infallible in practice. http: reqwest::Client::builder() .timeout(std::time::Duration::from_secs(10)) .connect_timeout(std::time::Duration::from_secs(5)) .build() - .expect("SAFETY: default builder with only timeout config cannot fail"), - inner: Arc::new(Mutex::new(inner)), + .map_err(|e| RelayClientError::Url(format!("HTTP client build failed: {e}")))?, api_token: api_token.map(|t| t.to_string()), - active_subscriptions: Arc::new(Mutex::new(HashMap::new())), }) } - /// Reconnect after a connection loss: replace the inner WebSocket with a fresh - /// authenticated connection (using exponential backoff), then resubscribe to all - /// subscriptions that were active at the time of the disconnect. - pub async fn reconnect(&self) { - tracing::warn!("relay connection lost — reconnecting…"); - let new_inner = - Self::connect_with_retry(&self.relay_url, &self.keys, self.api_token.as_deref()).await; - - { - let mut inner = self.inner.lock().await; - *inner = new_inner; - } - - let subs = self.active_subscriptions.lock().await.clone(); - if !subs.is_empty() { - tracing::info!("resubscribing to {} active subscription(s)", subs.len()); - for (sub_id, filters) in &subs { - let mut inner = self.inner.lock().await; - let mut msg: Vec = Vec::with_capacity(2 + filters.len()); - msg.push(json!("REQ")); - msg.push(json!(sub_id)); - for f in filters { - match serde_json::to_value(f) { - Ok(v) => msg.push(v), - Err(e) => { - tracing::warn!("failed to serialize filter for {sub_id}: {e}"); - } - } - } - if let Err(e) = inner.send_raw(&Value::Array(msg)).await { - tracing::warn!("failed to resubscribe to {sub_id}: {e}"); - } - } - } - } - /// Returns the Nostr keypair used for signing and authentication. pub fn keys(&self) -> &Keys { &self.keys @@ -582,11 +851,18 @@ impl RelayClient { /// Publish a signed Nostr event to the relay and wait for the `OK` acknowledgement. pub async fn send_event(&self, event: Event) -> Result { - let mut inner = self.inner.lock().await; - let event_id = event.id.to_hex(); - let msg = json!(["EVENT", event]); - inner.send_raw(&msg).await?; - inner.wait_for_ok(&event_id, Duration::from_secs(10)).await + let (reply_tx, reply_rx) = oneshot::channel(); + self.bg + .cmd_tx + .send(RelayCommand::SendEvent { + event, + reply: reply_tx, + }) + .await + .map_err(|_| RelayClientError::ConnectionClosed)?; + reply_rx + .await + .map_err(|_| RelayClientError::ConnectionClosed)? } /// Open a subscription with the given filters and collect all stored events until `EOSE`. @@ -595,45 +871,49 @@ impl RelayClient { sub_id: &str, filters: Vec, ) -> Result, RelayClientError> { - // Track this subscription so it can be resubmitted after a reconnect. - self.active_subscriptions - .lock() + let (reply_tx, reply_rx) = oneshot::channel(); + self.bg + .cmd_tx + .send(RelayCommand::Subscribe { + sub_id: sub_id.to_string(), + filters, + reply: reply_tx, + }) .await - .insert(sub_id.to_string(), filters.clone()); - - let mut inner = self.inner.lock().await; - - let mut msg: Vec = Vec::with_capacity(2 + filters.len()); - msg.push(json!("REQ")); - msg.push(json!(sub_id)); - for f in &filters { - msg.push(serde_json::to_value(f)?); - } - inner.send_raw(&Value::Array(msg)).await?; - - inner - .collect_until_eose(sub_id, Duration::from_secs(10)) + .map_err(|_| RelayClientError::ConnectionClosed)?; + reply_rx .await + .map_err(|_| RelayClientError::ConnectionClosed)? } /// Send a `CLOSE` message to the relay and remove the subscription from the active set. pub async fn close_subscription(&self, sub_id: &str) -> Result<(), RelayClientError> { - // Remove from active subscriptions — no longer needs to be resubscribed. - self.active_subscriptions.lock().await.remove(sub_id); - - let mut inner = self.inner.lock().await; - let msg = json!(["CLOSE", sub_id]); - inner.send_raw(&msg).await + let (reply_tx, reply_rx) = oneshot::channel(); + self.bg + .cmd_tx + .send(RelayCommand::CloseSubscription { + sub_id: sub_id.to_string(), + reply: reply_tx, + }) + .await + .map_err(|_| RelayClientError::ConnectionClosed)?; + reply_rx + .await + .map_err(|_| RelayClientError::ConnectionClosed)? } - /// Perform a clean WebSocket close handshake. + /// Signal the background task to shut down. + /// + /// The task will also be aborted when the last `RelayClient` clone is + /// dropped, so calling this explicitly is optional but allows a prompt stop. pub async fn close(&self) -> Result<(), RelayClientError> { - let mut inner = self.inner.lock().await; - inner.ws.close(None).await?; + let _ = self.bg.cmd_tx.send(RelayCommand::Shutdown).await; Ok(()) } } +// ── Free functions ──────────────────────────────────────────────────────────── + /// Convert a WebSocket URL to its HTTP equivalent. /// Converts `ws://` → `http://` and `wss://` → `https://`, strips trailing slash. /// @@ -932,4 +1212,514 @@ mod tests { let result = parse_relay_message(text); assert!(result.is_err()); } + + // ── Integration tests: mini relay ───────────────────────────────────────── + // + // Each test spins up a lightweight in-process WebSocket server that performs + // the NIP-42 handshake, then runs a caller-supplied scenario closure. + // The closure receives the split sink+stream so it can drive the test. + + #[cfg(test)] + mod integration { + use super::*; + use futures_util::stream::{SplitSink, SplitStream}; + use futures_util::{SinkExt, StreamExt}; + use std::future::Future; + use tokio::net::TcpListener; + use tokio_tungstenite::{accept_async, tungstenite::Message, WebSocketStream}; + + type PlainWs = WebSocketStream; + + /// Spawn a mini relay that performs NIP-42 auth handshake then runs `scenario`. + /// Returns the `ws://127.0.0.1:{port}` URL. + async fn spawn_mini_relay(scenario: F) -> String + where + F: FnOnce(SplitSink, SplitStream) -> Fut + Send + 'static, + Fut: Future + Send, + { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let port = listener.local_addr().unwrap().port(); + + tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + let ws = accept_async(stream).await.unwrap(); + let (mut sink, mut stream) = ws.split(); + + // Send AUTH challenge. + sink.send(Message::Text( + r#"["AUTH","test-challenge"]"#.to_string().into(), + )) + .await + .unwrap(); + + // Wait for AUTH response, send OK for the event. + while let Some(Ok(msg)) = stream.next().await { + if let Message::Text(text) = msg { + let arr: Vec = serde_json::from_str(&text).unwrap(); + if arr.first().and_then(|v| v.as_str()) == Some("AUTH") { + let event_id = arr[1]["id"].as_str().unwrap().to_string(); + sink.send(Message::Text( + format!(r#"["OK","{}",true,""]"#, event_id).into(), + )) + .await + .unwrap(); + break; + } + } + } + + scenario(sink, stream).await; + }); + + format!("ws://127.0.0.1:{}", port) + } + + // ── Test 1: background task responds to Ping when idle ──────────────── + + #[tokio::test] + async fn bg_responds_to_ping_without_caller_activity() { + let url = spawn_mini_relay(|mut sink, mut stream| async move { + // Send a Ping — background task should Pong immediately. + sink.send(Message::Ping(b"abc".to_vec().into())) + .await + .unwrap(); + + // Drain until we see the Pong. + while let Some(Ok(msg)) = stream.next().await { + if let Message::Pong(data) = msg { + assert_eq!(data.as_ref(), b"abc"); + return; + } + } + panic!("never received Pong"); + }) + .await; + + let keys = Keys::generate(); + let client = RelayClient::connect(&url, &keys, None).await.unwrap(); + // Give the background task a moment to process the Ping. + tokio::time::sleep(Duration::from_millis(200)).await; + let _ = client.close().await; + } + + // ── Test 2: background task handles mid-session AUTH challenge ──────── + + #[tokio::test] + async fn bg_handles_mid_session_auth_challenge() { + let url = spawn_mini_relay(|mut sink, mut stream| async move { + // Send a fresh AUTH challenge after initial handshake. + sink.send(Message::Text( + r#"["AUTH","challenge-2"]"#.to_string().into(), + )) + .await + .unwrap(); + + // Expect a new AUTH event with kind 22242. + while let Some(Ok(msg)) = stream.next().await { + if let Message::Text(text) = msg { + let arr: Vec = serde_json::from_str(&text).unwrap(); + if arr.first().and_then(|v| v.as_str()) == Some("AUTH") { + let kind = arr[1]["kind"].as_u64().unwrap(); + assert_eq!(kind, 22242, "expected kind 22242 (Authentication)"); + // Verify the challenge tag is present. + let tags = arr[1]["tags"].as_array().unwrap(); + let has_challenge = tags.iter().any(|t| { + t.as_array().and_then(|a| a.get(1)).and_then(|v| v.as_str()) + == Some("challenge-2") + }); + assert!(has_challenge, "AUTH event missing challenge tag"); + return; + } + } + } + panic!("never received AUTH response"); + }) + .await; + + let keys = Keys::generate(); + let client = RelayClient::connect(&url, &keys, None).await.unwrap(); + tokio::time::sleep(Duration::from_millis(200)).await; + let _ = client.close().await; + } + + // ── Test 3: send_event receives OK response ─────────────────────────── + + #[tokio::test] + async fn send_event_receives_ok_response() { + let url = spawn_mini_relay(|mut sink, mut stream| async move { + // Wait for EVENT, send matching OK. + while let Some(Ok(msg)) = stream.next().await { + if let Message::Text(text) = msg { + let arr: Vec = serde_json::from_str(&text).unwrap(); + if arr.first().and_then(|v| v.as_str()) == Some("EVENT") { + let event_id = arr[1]["id"].as_str().unwrap().to_string(); + sink.send(Message::Text( + format!(r#"["OK","{}",true,""]"#, event_id).into(), + )) + .await + .unwrap(); + return; + } + } + } + }) + .await; + + let keys = Keys::generate(); + let client = RelayClient::connect(&url, &keys, None).await.unwrap(); + + let event = EventBuilder::new(Kind::Custom(9), "test", []) + .sign_with_keys(&keys) + .unwrap(); + let expected_id = event.id.to_hex(); + + let ok = client.send_event(event).await.unwrap(); + assert_eq!(ok.event_id, expected_id); + assert!(ok.accepted); + assert_eq!(ok.message, ""); + + let _ = client.close().await; + } + + // ── Test 4: subscribe collects events until EOSE ────────────────────── + + #[tokio::test] + async fn subscribe_collects_events_until_eose() { + let url = spawn_mini_relay(|mut sink, mut stream| async move { + // Wait for REQ. + while let Some(Ok(msg)) = stream.next().await { + if let Message::Text(text) = msg { + let arr: Vec = serde_json::from_str(&text).unwrap(); + if arr.first().and_then(|v| v.as_str()) == Some("REQ") { + let sub_id = arr[1].as_str().unwrap().to_string(); + + // Build 3 minimal valid events. + let relay_keys = Keys::generate(); + for i in 0u8..3 { + let ev = EventBuilder::new(Kind::TextNote, format!("msg {i}"), []) + .sign_with_keys(&relay_keys) + .unwrap(); + let frame = serde_json::to_string(&serde_json::json!([ + "EVENT", sub_id, ev + ])) + .unwrap(); + sink.send(Message::Text(frame.into())).await.unwrap(); + } + + // Send EOSE. + sink.send(Message::Text(format!(r#"["EOSE","{}"]"#, sub_id).into())) + .await + .unwrap(); + return; + } + } + } + }) + .await; + + let keys = Keys::generate(); + let client = RelayClient::connect(&url, &keys, None).await.unwrap(); + + let events = client + .subscribe("sub-1", vec![Filter::new()]) + .await + .unwrap(); + assert_eq!(events.len(), 3, "expected 3 events before EOSE"); + + let _ = client.close().await; + } + + // ── Test 5: send_event times out when relay never sends OK ──────────── + // + // We connect first (real time), then pause time and advance past the + // 10s SEND_EVENT_TIMEOUT to avoid a real wait. + + #[tokio::test] + async fn send_event_times_out_when_no_ok() { + let url = spawn_mini_relay(|_sink, mut stream| async move { + // Consume the EVENT frame but never respond with OK. + // Hold connection open by draining until the client drops. + while let Some(Ok(_)) = stream.next().await {} + }) + .await; + + let keys = Keys::generate(); + let client = RelayClient::connect(&url, &keys, None).await.unwrap(); + + let event = EventBuilder::new(Kind::Custom(9), "timeout-test", []) + .sign_with_keys(&keys) + .unwrap(); + + // Pause time AFTER connecting so the auth handshake completes normally. + tokio::time::pause(); + + // Start the send — enqueues EVENT to background task. + let send_fut = client.send_event(event); + tokio::pin!(send_fut); + + // Yield to let the background task send the EVENT frame. + tokio::task::yield_now().await; + tokio::task::yield_now().await; + + // Advance time past the 10s SEND_EVENT_TIMEOUT + 1s tick granularity. + tokio::time::advance(Duration::from_secs(12)).await; + // Let the background task's tick fire and expire the pending_ok entry. + tokio::task::yield_now().await; + tokio::task::yield_now().await; + + let result = send_fut.await; + assert!( + matches!(result, Err(RelayClientError::Timeout)), + "expected Timeout, got: {:?}", + result + ); + + let _ = client.close().await; + } + + // ── Test 6: close_subscription sends CLOSE message to relay ────────── + + #[tokio::test] + async fn close_subscription_sends_close_message() { + let (close_tx, close_rx) = tokio::sync::oneshot::channel::(); + let close_tx = std::sync::Arc::new(tokio::sync::Mutex::new(Some(close_tx))); + + let url = spawn_mini_relay({ + let close_tx = close_tx.clone(); + move |mut sink, mut stream| async move { + // Handle REQ, send EOSE immediately so subscribe() returns. + while let Some(Ok(msg)) = stream.next().await { + if let Message::Text(text) = msg { + let arr: Vec = + serde_json::from_str(&text).unwrap_or_default(); + match arr.first().and_then(|v| v.as_str()) { + Some("REQ") => { + let sub_id = arr[1].as_str().unwrap().to_string(); + sink.send(Message::Text( + format!(r#"["EOSE","{}"]"#, sub_id).into(), + )) + .await + .unwrap(); + } + Some("CLOSE") => { + let sub_id = arr[1].as_str().unwrap().to_string(); + if let Some(tx) = close_tx.lock().await.take() { + let _ = tx.send(sub_id); + } + return; + } + _ => {} + } + } + } + } + }) + .await; + + let keys = Keys::generate(); + let client = RelayClient::connect(&url, &keys, None).await.unwrap(); + + // Subscribe (EOSE comes back immediately). + client + .subscribe("sub-close", vec![Filter::new()]) + .await + .unwrap(); + + // Close the subscription — relay should receive CLOSE. + client.close_subscription("sub-close").await.unwrap(); + + let closed_id = tokio::time::timeout(Duration::from_secs(2), close_rx) + .await + .expect("timed out waiting for CLOSE") + .expect("channel dropped"); + + assert_eq!(closed_id, "sub-close"); + + let _ = client.close().await; + } + + // ── Test 7: reconnect on transport close ────────────────────────────── + // + // Strategy: use a shared TcpListener that accepts two connections. + // First connection: close immediately after auth. + // Second connection: full relay that handles send_event. + + #[tokio::test] + async fn bg_reconnects_on_transport_close() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let port = listener.local_addr().unwrap().port(); + let url = format!("ws://127.0.0.1:{}", port); + + // Channel so the test can wait until the second relay is ready. + let (ok_tx, ok_rx) = tokio::sync::oneshot::channel::<()>(); + let ok_tx = std::sync::Arc::new(tokio::sync::Mutex::new(Some(ok_tx))); + + tokio::spawn({ + let ok_tx = ok_tx.clone(); + async move { + // ── Connection 1: close right after auth ────────────────── + { + let (stream, _) = listener.accept().await.unwrap(); + let ws = accept_async(stream).await.unwrap(); + let (mut sink, mut stream) = ws.split(); + + sink.send(Message::Text( + r#"["AUTH","test-challenge"]"#.to_string().into(), + )) + .await + .unwrap(); + + while let Some(Ok(msg)) = stream.next().await { + if let Message::Text(text) = msg { + let arr: Vec = + serde_json::from_str(&text).unwrap_or_default(); + if arr.first().and_then(|v| v.as_str()) == Some("AUTH") { + let event_id = arr[1]["id"].as_str().unwrap().to_string(); + sink.send(Message::Text( + format!(r#"["OK","{}",true,""]"#, event_id).into(), + )) + .await + .unwrap(); + break; + } + } + } + // Drop sink+stream — closes the WS connection. + } + + // ── Connection 2: full relay that handles EVENT ──────────── + { + let (stream, _) = listener.accept().await.unwrap(); + let ws = accept_async(stream).await.unwrap(); + let (mut sink, mut stream) = ws.split(); + + sink.send(Message::Text( + r#"["AUTH","test-challenge"]"#.to_string().into(), + )) + .await + .unwrap(); + + while let Some(Ok(msg)) = stream.next().await { + if let Message::Text(text) = msg { + let arr: Vec = + serde_json::from_str(&text).unwrap_or_default(); + match arr.first().and_then(|v| v.as_str()) { + Some("AUTH") => { + let event_id = arr[1]["id"].as_str().unwrap().to_string(); + sink.send(Message::Text( + format!(r#"["OK","{}",true,""]"#, event_id).into(), + )) + .await + .unwrap(); + // Signal AFTER auth handshake completes. + if let Some(tx) = ok_tx.lock().await.take() { + let _ = tx.send(()); + } + } + Some("EVENT") => { + let event_id = arr[1]["id"].as_str().unwrap().to_string(); + sink.send(Message::Text( + format!(r#"["OK","{}",true,""]"#, event_id).into(), + )) + .await + .unwrap(); + return; + } + _ => {} + } + } + } + } + } + }); + + let keys = Keys::generate(); + let client = RelayClient::connect(&url, &keys, None).await.unwrap(); + + // Wait for the second relay to complete auth (background task reconnected). + tokio::time::timeout(Duration::from_secs(5), ok_rx) + .await + .expect("timed out waiting for reconnect") + .unwrap(); + + // send_event should succeed on the new connection. + let event = EventBuilder::new(Kind::Custom(9), "after-reconnect", []) + .sign_with_keys(&keys) + .unwrap(); + let ok = client.send_event(event).await.unwrap(); + assert!(ok.accepted); + + let _ = client.close().await; + } + + // ── Test 8: shutdown during reconnect ───────────────────────────── + + #[tokio::test] + async fn shutdown_during_reconnect_exits_promptly() { + // Connect to a relay that closes immediately after auth, + // then never accepts again. The background task enters the + // reconnect loop. Verify that Shutdown is processed during + // reconnect backoff — the task exits gracefully, NOT via abort. + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let port = listener.local_addr().unwrap().port(); + let url = format!("ws://127.0.0.1:{}", port); + + tokio::spawn(async move { + // Accept one connection, auth, then close. + let (stream, _) = listener.accept().await.unwrap(); + let ws = accept_async(stream).await.unwrap(); + let (mut sink, mut stream) = ws.split(); + + sink.send(Message::Text( + r#"["AUTH","test-challenge"]"#.to_string().into(), + )) + .await + .unwrap(); + + while let Some(Ok(msg)) = stream.next().await { + if let Message::Text(text) = msg { + let arr: Vec = + serde_json::from_str(&text).unwrap_or_default(); + if arr.first().and_then(|v| v.as_str()) == Some("AUTH") { + let event_id = arr[1]["id"].as_str().unwrap().to_string(); + sink.send(Message::Text( + format!(r#"["OK","{}",true,""]"#, event_id).into(), + )) + .await + .unwrap(); + break; + } + } + } + // Drop — closes connection, triggering reconnect. + // Don't accept any more connections — reconnect will fail forever. + drop(listener); + }); + + let keys = Keys::generate(); + let client = RelayClient::connect(&url, &keys, None).await.unwrap(); + + // Give the background task time to notice the close and enter + // the reconnect backoff loop. + tokio::time::sleep(Duration::from_millis(200)).await; + + // Send Shutdown via close() — the reconnect loop processes this + // during its backoff sleep and exits the task gracefully. + let _ = client.close().await; + + // Wait for the task to process Shutdown. Do NOT drop the client + // yet — we want to prove the task exits via Shutdown processing, + // not via BgTaskHandle::drop calling abort(). + tokio::time::sleep(Duration::from_millis(500)).await; + + // After graceful shutdown, the background task dropped its cmd_rx. + // Sending another command should fail with a closed-channel error, + // proving the task exited on its own. + let result = client.bg.cmd_tx.send(RelayCommand::Shutdown).await; + assert!( + result.is_err(), + "cmd channel should be closed after graceful shutdown — task exited via Shutdown, not abort" + ); + } + } } diff --git a/crates/sprout-relay/src/api/agents.rs b/crates/sprout-relay/src/api/agents.rs index e54f4ffe03..3221d3dcb9 100644 --- a/crates/sprout-relay/src/api/agents.rs +++ b/crates/sprout-relay/src/api/agents.rs @@ -30,7 +30,7 @@ pub async fn agents_handler( // Get requester's accessible channels to filter bot channel visibility. let accessible_channels = state .db - .get_accessible_channels(&pubkey_bytes, None) + .get_accessible_channels(&pubkey_bytes, None, None) .await .map_err(|e| { tracing::error!("agents: failed to load accessible channels: {e}"); diff --git a/crates/sprout-relay/src/api/channels.rs b/crates/sprout-relay/src/api/channels.rs index bb4a3c589d..ecffc26bee 100644 --- a/crates/sprout-relay/src/api/channels.rs +++ b/crates/sprout-relay/src/api/channels.rs @@ -26,6 +26,8 @@ use super::{api_error, extract_auth_context, internal_error}; pub struct ListChannelsParams { /// Optional visibility filter: `"open"` or `"private"`. pub visibility: Option, + /// When `true`, return only channels the user is a member of. + pub member: Option, } /// Returns all channels accessible to the authenticated user. @@ -43,7 +45,7 @@ pub async fn channels_handler( let channels = state .db - .get_accessible_channels(&pubkey_bytes, params.visibility.as_deref()) + .get_accessible_channels(&pubkey_bytes, params.visibility.as_deref(), params.member) .await .map_err(|e| internal_error(&format!("db error: {e}")))?; diff --git a/crates/sprout-relay/src/connection.rs b/crates/sprout-relay/src/connection.rs index 7fdf66b6f3..bc991a6370 100644 --- a/crates/sprout-relay/src/connection.rs +++ b/crates/sprout-relay/src/connection.rs @@ -20,6 +20,10 @@ use crate::handlers; use crate::protocol::{ClientMessage, RelayMessage}; use crate::state::AppState; +/// Number of buffer-full events tolerated before cancelling a slow client. +/// Prevents transient read stalls from hard-disconnecting agents mid-inference. +pub(crate) const SLOW_CLIENT_GRACE_LIMIT: u8 = 3; + /// NIP-42 authentication state for a single connection. #[derive(Debug, Clone)] pub enum AuthState { @@ -37,7 +41,7 @@ pub enum AuthState { /// Per-connection state split by access pattern: /// - `auth_state`: RwLock (read-heavy after initial auth) /// - `subscriptions`: Mutex (write-heavy during REQ/CLOSE) -/// - `send_tx`, `cancel`: outside any lock (Clone+Send, no coordination needed) +/// - `send_tx`, `ctrl_tx`, `cancel`: outside any lock (Clone+Send, no coordination needed) pub struct ConnectionState { /// Unique identifier for this connection. pub conn_id: Uuid, @@ -47,23 +51,41 @@ pub struct ConnectionState { pub auth_state: RwLock, /// Active subscriptions keyed by subscription ID. pub subscriptions: Mutex>>, - /// Sender for outbound WebSocket messages. + /// Sender for outbound data messages (EVENT, NOTICE, OK, etc.). pub send_tx: mpsc::Sender, + /// Sender for outbound control frames (Pong, Close). + /// Separate channel with priority drain — if this channel fills too, + /// the connection is closed (writer is completely stalled). + pub ctrl_tx: mpsc::Sender, /// Token used to signal graceful shutdown of this connection's tasks. pub cancel: CancellationToken, + /// Consecutive buffer-full events. Cancel only after [`SLOW_CLIENT_GRACE_LIMIT`]. + /// Shared with `ConnectionManager::ConnEntry` so both direct sends and + /// fan-out broadcasts track the same counter. + pub backpressure_count: Arc, } impl ConnectionState { - /// Sends a message to this connection's outbound channel. + /// Sends a data message to this connection's outbound channel. /// - /// If the send buffer is full (slow client), cancels the connection - /// via the `CancellationToken` to prevent unbounded memory growth. + /// On a full buffer, increments the backpressure counter. The first + /// [`SLOW_CLIENT_GRACE_LIMIT`] occurrences log a warning; sustained + /// backpressure cancels the connection to prevent unbounded memory growth. pub fn send(&self, msg: String) -> bool { match self.send_tx.try_send(WsMessage::Text(msg.into())) { - Ok(_) => true, + Ok(_) => { + // Successful send resets the grace counter. + self.backpressure_count.store(0, Ordering::Relaxed); + true + } Err(tokio::sync::mpsc::error::TrySendError::Full(_)) => { - warn!(conn_id = %self.conn_id, "send buffer full — closing slow client"); - self.cancel.cancel(); + let count = self.backpressure_count.fetch_add(1, Ordering::Relaxed) + 1; + if count >= SLOW_CLIENT_GRACE_LIMIT { + warn!(conn_id = %self.conn_id, count, "sustained backpressure — closing slow client"); + self.cancel.cancel(); + } else { + warn!(conn_id = %self.conn_id, count, grace = SLOW_CLIENT_GRACE_LIMIT, "send buffer full — grace {count}/{SLOW_CLIENT_GRACE_LIMIT}"); + } false } Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => { @@ -92,6 +114,11 @@ pub async fn handle_connection(socket: WebSocket, state: Arc, addr: So let cancel = CancellationToken::new(); let (tx, rx) = mpsc::channel::(state.config.send_buffer_size); + // Control channel for Pong/Close — small capacity, guaranteed delivery + // even when the data buffer is full. + let (ctrl_tx, ctrl_rx) = mpsc::channel::(8); + + let backpressure_count = Arc::new(AtomicU8::new(0)); let conn = Arc::new(ConnectionState { conn_id, @@ -101,7 +128,9 @@ pub async fn handle_connection(socket: WebSocket, state: Arc, addr: So }), subscriptions: Mutex::new(HashMap::new()), send_tx: tx.clone(), + ctrl_tx: ctrl_tx.clone(), cancel: cancel.clone(), + backpressure_count: Arc::clone(&backpressure_count), }); info!(conn_id = %conn_id, addr = %addr, "WebSocket connection established"); @@ -117,19 +146,22 @@ pub async fn handle_connection(socket: WebSocket, state: Arc, addr: So } // Register after challenge succeeds — avoids leaked entries on early disconnect. - state - .conn_manager - .register(conn_id, tx.clone(), cancel.clone()); + state.conn_manager.register( + conn_id, + tx.clone(), + cancel.clone(), + Arc::clone(&backpressure_count), + ); let (ws_send, ws_recv) = socket.split(); let send_cancel = cancel.child_token(); - let send_task = tokio::spawn(send_loop(ws_send, rx, send_cancel)); + let send_task = tokio::spawn(send_loop(ws_send, rx, ctrl_rx, send_cancel)); let missed_pongs = Arc::new(AtomicU8::new(0)); let heartbeat_cancel = cancel.clone(); let heartbeat_task = tokio::spawn(heartbeat_loop( - tx.clone(), + ctrl_tx, Arc::clone(&missed_pongs), heartbeat_cancel, )); @@ -154,29 +186,55 @@ pub async fn handle_connection(socket: WebSocket, state: Arc, addr: So drop(permit); } +/// Outbound send loop with control-frame priority. +/// +/// Control frames (Pong, Close) are drained first on every iteration, +/// giving them priority over data frames. If the underlying socket writer +/// is stalled, control frames queue in the small ctrl_rx buffer; callers +/// treat a full control channel as terminal (Bug 7 fix). async fn send_loop( mut ws_send: futures_util::stream::SplitSink, - mut rx: mpsc::Receiver, + mut data_rx: mpsc::Receiver, + mut ctrl_rx: mpsc::Receiver, cancel: CancellationToken, ) { loop { - tokio::select! { - Some(msg) = rx.recv() => { - if ws_send.send(msg).await.is_err() { - break; - } + // Priority: drain all pending control frames before data. + while let Ok(ctrl_msg) = ctrl_rx.try_recv() { + if ws_send.send(ctrl_msg).await.is_err() { + return; } + } + + tokio::select! { + // Biased: cancel > control > data. Cancel must win immediately + // so backpressure-triggered shutdown isn't starved by queued data. + biased; _ = cancel.cancelled() => { let _ = ws_send.send(WsMessage::Close(None)).await; break; } + Some(ctrl_msg) = ctrl_rx.recv() => { + if ws_send.send(ctrl_msg).await.is_err() { + break; + } + } + Some(msg) = data_rx.recv() => { + if ws_send.send(msg).await.is_err() { + break; + } + } } } } /// 3 missed pongs → disconnect. +/// +/// Sends Ping through the control channel so it isn't blocked by a full +/// data buffer. Uses `try_send` to keep the select loop responsive to +/// cancellation — a full control channel means the writer is stalled. async fn heartbeat_loop( - tx: mpsc::Sender, + ctrl_tx: mpsc::Sender, missed_pongs: Arc, cancel: CancellationToken, ) { @@ -194,7 +252,9 @@ async fn heartbeat_loop( cancel.cancel(); break; } - if tx.send(WsMessage::Ping(axum::body::Bytes::new())).await.is_err() { + if ctrl_tx.try_send(WsMessage::Ping(axum::body::Bytes::new())).is_err() { + warn!("control channel full — cannot send Ping, closing"); + cancel.cancel(); break; } } @@ -241,7 +301,14 @@ async fn recv_loop( missed_pongs.store(0, Ordering::Relaxed); } Some(Ok(WsMessage::Ping(data))) => { - let _ = conn.send_tx.try_send(WsMessage::Pong(data)); + // Send Pong through the control channel — priority + // delivery even when the data buffer is full (Bug 7 fix). + if conn.ctrl_tx.try_send(WsMessage::Pong(data)).is_err() { + // Control channel full means the socket writer is + // completely stalled — treat as terminal. + warn!(conn_id = %conn.conn_id, "control channel full — cannot send Pong, closing"); + break; + } } Some(Ok(WsMessage::Close(_))) | None => { debug!("WebSocket closed by client"); diff --git a/crates/sprout-relay/src/state.rs b/crates/sprout-relay/src/state.rs index ea11682725..dc333acce6 100644 --- a/crates/sprout-relay/src/state.rs +++ b/crates/sprout-relay/src/state.rs @@ -1,5 +1,6 @@ //! Shared application state — Arc-wrapped, shared across all connections. +use std::sync::atomic::{AtomicU8, Ordering}; use std::sync::Arc; use std::time::Instant; @@ -19,13 +20,21 @@ use sprout_workflow::WorkflowEngine; use crate::api::tokens::MintRateLimiter; use crate::config::Config; +use crate::connection::SLOW_CLIENT_GRACE_LIMIT; use crate::subscription::SubscriptionRegistry; +/// Per-connection entry in the connection manager. +struct ConnEntry { + tx: mpsc::Sender, + cancel: CancellationToken, + /// Shared with `ConnectionState` — both direct sends and fan-out + /// broadcasts track the same consecutive-full counter. + backpressure_count: Arc, +} + /// Tracks active WebSocket connections and provides message routing by connection ID. pub struct ConnectionManager { - /// Map from connection ID to the sender half of the connection's outbound channel - /// and the cancellation token for the connection. - connections: DashMap, CancellationToken)>, + connections: DashMap, } impl ConnectionManager { @@ -36,9 +45,23 @@ impl ConnectionManager { } } - /// Registers a connection with its outbound sender and cancellation token. - pub fn register(&self, conn_id: Uuid, tx: mpsc::Sender, cancel: CancellationToken) { - self.connections.insert(conn_id, (tx, cancel)); + /// Registers a connection with its outbound sender, cancellation token, + /// and shared backpressure counter (same `Arc` as `ConnectionState`). + pub fn register( + &self, + conn_id: Uuid, + tx: mpsc::Sender, + cancel: CancellationToken, + backpressure_count: Arc, + ) { + self.connections.insert( + conn_id, + ConnEntry { + tx, + cancel, + backpressure_count, + }, + ); } /// Removes a connection from the registry. @@ -49,15 +72,24 @@ impl ConnectionManager { /// Sends a text message to the given connection. /// /// Returns `false` if the connection is gone or the buffer is full. - /// On a full buffer, cancels the slow client's connection to prevent silent drops. + /// On sustained backpressure (>[`SLOW_CLIENT_GRACE_LIMIT`] consecutive full + /// buffers), cancels the connection. Transient stalls get a warning only. pub fn send_to(&self, conn_id: Uuid, msg: String) -> bool { if let Some(entry) = self.connections.get(&conn_id) { - let (tx, cancel) = entry.value(); - match tx.try_send(WsMessage::Text(msg.into())) { - Ok(_) => true, + let conn = entry.value(); + match conn.tx.try_send(WsMessage::Text(msg.into())) { + Ok(_) => { + conn.backpressure_count.store(0, Ordering::Relaxed); + true + } Err(tokio::sync::mpsc::error::TrySendError::Full(_)) => { - tracing::warn!(conn_id = %conn_id, "fan-out: send buffer full — cancelling slow client"); - cancel.cancel(); + let count = conn.backpressure_count.fetch_add(1, Ordering::Relaxed) + 1; + if count >= SLOW_CLIENT_GRACE_LIMIT { + tracing::warn!(conn_id = %conn_id, count, "fan-out: sustained backpressure — cancelling slow client"); + conn.cancel.cancel(); + } else { + tracing::warn!(conn_id = %conn_id, count, grace = SLOW_CLIENT_GRACE_LIMIT, "fan-out: send buffer full — grace {count}/{SLOW_CLIENT_GRACE_LIMIT}"); + } false } Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => { @@ -175,3 +207,129 @@ impl std::fmt::Debug for AppState { .finish() } } + +// ── Unit tests ──────────────────────────────────────────────────────────────── + +#[cfg(test)] +mod tests { + use super::*; + use crate::connection::{AuthState, ConnectionState}; + use std::collections::HashMap; + use tokio::sync::{Mutex, RwLock}; + + /// Helper: create a ConnectionManager with one registered connection. + /// Returns (manager, conn_id, receiver, cancel, shared_backpressure_count). + fn setup_conn( + buffer_size: usize, + ) -> ( + ConnectionManager, + Uuid, + mpsc::Receiver, + CancellationToken, + Arc, + ) { + let mgr = ConnectionManager::new(); + let conn_id = Uuid::new_v4(); + let (tx, rx) = mpsc::channel(buffer_size); + let cancel = CancellationToken::new(); + let bp = Arc::new(AtomicU8::new(0)); + mgr.register(conn_id, tx, cancel.clone(), Arc::clone(&bp)); + (mgr, conn_id, rx, cancel, bp) + } + + #[test] + fn send_to_resets_grace_counter_on_success() { + let (mgr, id, _rx, _cancel, bp) = setup_conn(16); + // Simulate prior backpressure. + bp.store(2, Ordering::Relaxed); + assert!(mgr.send_to(id, "hello".into())); + assert_eq!( + bp.load(Ordering::Relaxed), + 0, + "successful send should reset counter" + ); + } + + #[test] + fn send_to_increments_grace_counter_on_full() { + // Buffer size 1 — fill it, then the next send is Full. + let (mgr, id, _rx, cancel, bp) = setup_conn(1); + assert!(mgr.send_to(id, "fill".into())); + // Buffer is now full. + assert!(!mgr.send_to(id, "overflow-1".into())); + assert_eq!(bp.load(Ordering::Relaxed), 1, "first overflow → count=1"); + assert!( + !cancel.is_cancelled(), + "should not cancel on first overflow" + ); + + assert!(!mgr.send_to(id, "overflow-2".into())); + assert_eq!(bp.load(Ordering::Relaxed), 2); + assert!( + !cancel.is_cancelled(), + "should not cancel on second overflow" + ); + } + + #[test] + fn send_to_cancels_after_grace_limit() { + let (mgr, id, _rx, cancel, _bp) = setup_conn(1); + assert!(mgr.send_to(id, "fill".into())); + // Exhaust grace: 3 consecutive Full events. + for _ in 0..SLOW_CLIENT_GRACE_LIMIT { + mgr.send_to(id, "overflow".into()); + } + assert!( + cancel.is_cancelled(), + "should cancel after SLOW_CLIENT_GRACE_LIMIT overflows" + ); + } + + #[test] + fn shared_counter_between_direct_and_fanout() { + // Verify that ConnectionState::send() and ConnectionManager::send_to() + // share the same backpressure counter via Arc. + let conn_id = Uuid::new_v4(); + let (tx, _rx) = mpsc::channel(1); + let (ctrl_tx, _ctrl_rx) = mpsc::channel(8); + let cancel = CancellationToken::new(); + let bp = Arc::new(AtomicU8::new(0)); + + let conn = ConnectionState { + conn_id, + remote_addr: "127.0.0.1:1234".parse().unwrap(), + auth_state: RwLock::new(AuthState::Failed), + subscriptions: Mutex::new(HashMap::new()), + send_tx: tx.clone(), + ctrl_tx, + cancel: cancel.clone(), + backpressure_count: Arc::clone(&bp), + }; + + let mgr = ConnectionManager::new(); + mgr.register(conn_id, tx, cancel.clone(), Arc::clone(&bp)); + + // Fill the buffer via direct send. + assert!(conn.send("fill".into())); + // Overflow via fan-out. + assert!(!mgr.send_to(conn_id, "overflow-fanout".into())); + assert_eq!( + bp.load(Ordering::Relaxed), + 1, + "fan-out overflow increments shared counter" + ); + // Overflow via direct send. + assert!(!conn.send("overflow-direct".into())); + assert_eq!( + bp.load(Ordering::Relaxed), + 2, + "direct overflow increments same counter" + ); + // One more fan-out overflow → should cancel (3 consecutive). + mgr.send_to(conn_id, "overflow-final".into()); + assert!( + cancel.is_cancelled(), + "shared counter reached limit via mixed path" + ); + } +} diff --git a/desktop/src/features/messages/ui/MessageComposer.tsx b/desktop/src/features/messages/ui/MessageComposer.tsx index d1ef1ac16c..974b8a7449 100644 --- a/desktop/src/features/messages/ui/MessageComposer.tsx +++ b/desktop/src/features/messages/ui/MessageComposer.tsx @@ -136,18 +136,48 @@ export function MessageComposer({ [content, mentionStartIndex], ); - const extractMentionPubkeys = React.useCallback((text: string): string[] => { - const pubkeys: string[] = []; - const mentionMap = mentionMapRef.current; + const extractMentionPubkeys = React.useCallback( + (text: string): string[] => { + const pubkeys: string[] = []; - for (const [displayName, pubkey] of mentionMap) { - if (text.includes(`@${displayName}`)) { - pubkeys.push(pubkey); + /** Test whether `@name` appears as a mention token in `text`. + * Left boundary: start-of-string or whitespace. + * Right boundary: end-of-string, whitespace, or punctuation. + * Case-insensitive. Prevents @Ann from matching @Anna or foo@Ann. */ + const hasMention = (name: string): boolean => { + const escaped = name.replace(/[.*+?^${}()|[\]\\]/g, "\\$&"); + const pattern = new RegExp( + `(?:^|\\s)@${escaped}(?=[\\s,;.!?:)\\]}]|$)`, + "i", + ); + return pattern.test(text); + }; + + // 1. Check the autocomplete-inserted mention map. + for (const [displayName, pubkey] of mentionMapRef.current) { + if (hasMention(displayName)) { + pubkeys.push(pubkey); + } } - } - return [...new Set(pubkeys)]; - }, []); + // 2. Fall back to scanning channel members for manually-typed @mentions + // that bypassed autocomplete (e.g. user typed "@Daphne" without selecting). + for (const member of members) { + if (pubkeys.includes(member.pubkey)) { + continue; // already matched via autocomplete map + } + const name = + member.displayName ?? + managedAgentNamesByPubkey.get(member.pubkey.toLowerCase()); + if (name && hasMention(name)) { + pubkeys.push(member.pubkey); + } + } + + return [...new Set(pubkeys)]; + }, + [members, managedAgentNamesByPubkey], + ); const submitMessage = React.useCallback(async () => { const trimmed = content.trim();