diff --git a/Cargo.lock b/Cargo.lock index d94903d229..d7c2c59c87 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3303,6 +3303,7 @@ dependencies = [ "thiserror", "tokio", "tokio-tungstenite 0.26.2", + "tokio-util", "toml", "tracing", "tracing-subscriber", diff --git a/Cargo.toml b/Cargo.toml index b2e18462fc..bdff97b2dc 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -31,7 +31,7 @@ repository = "https://github.com/sprout-rs/sprout" [workspace.dependencies] # Runtime tokio = { version = "1", features = ["rt-multi-thread", "macros", "net", "time", "sync", "io-util", "signal", "process"] } -tokio-util = { version = "0.7", features = ["rt"] } +tokio-util = { version = "0.7", features = ["rt", "codec"] } # HTTP + WebSocket axum = { version = "0.8", features = ["ws", "macros"] } diff --git a/crates/sprout-acp/Cargo.toml b/crates/sprout-acp/Cargo.toml index 66c593e816..8c78dcd4cd 100644 --- a/crates/sprout-acp/Cargo.toml +++ b/crates/sprout-acp/Cargo.toml @@ -24,6 +24,9 @@ tokio = { workspace = true } # WebSocket tokio-tungstenite = { workspace = true } + +# Codec (bounded line reads) +tokio-util = { workspace = true } rustls = { version = "0.23", default-features = false, features = ["ring", "std"] } futures-util = { workspace = true } diff --git a/crates/sprout-acp/src/acp.rs b/crates/sprout-acp/src/acp.rs index 679882fc5c..f6958b8c38 100644 --- a/crates/sprout-acp/src/acp.rs +++ b/crates/sprout-acp/src/acp.rs @@ -8,8 +8,14 @@ //! 4. [`AcpClient::session_prompt_with_idle_timeout`] — send prompt with idle/hard deadline, return stop reason //! 5. [`AcpClient::session_cancel`] / [`AcpClient::cancel_with_cleanup`] — cancel in-flight turn -use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader}; +use futures_util::StreamExt; +use tokio::io::AsyncWriteExt; use tokio::process::{Child, ChildStdin, ChildStdout}; +use tokio_util::codec::{FramedRead, LinesCodec, LinesCodecError}; + +/// Maximum allowed size of a single NDJSON line from the agent's stdout. +/// Lines exceeding this limit are rejected to prevent OOM from rogue agents. +const MAX_LINE_SIZE: usize = 10_000_000; // 10 MB // ─── Public types ──────────────────────────────────────────────────────────── @@ -52,8 +58,11 @@ pub enum StopReason { impl StopReason { /// Parse a `stopReason` string from the ACP wire format. + /// + /// Matching is case-insensitive so agents that send `"END_TURN"` or + /// `"Cancelled"` are handled correctly without a protocol error. pub fn from_str(s: &str) -> Option { - match s { + match s.to_ascii_lowercase().as_str() { "end_turn" => Some(Self::EndTurn), "cancelled" => Some(Self::Cancelled), "max_tokens" => Some(Self::MaxTokens), @@ -82,6 +91,12 @@ pub enum AcpError { #[error("Hard turn timeout exceeded")] HardTimeout, + #[error("Request timeout — agent did not respond within {0:?}")] + Timeout(std::time::Duration), + + #[error("Write timeout — agent stopped reading stdin (blocked for {0:?})")] + WriteTimeout(std::time::Duration), + #[error("Protocol error: {0}")] Protocol(String), } @@ -97,8 +112,10 @@ pub struct AcpClient { child: Child, /// Write end of the agent's stdin pipe. stdin: ChildStdin, - /// Buffered reader over the agent's stdout pipe (line-oriented). - reader: BufReader, + /// Framed reader over the agent's stdout pipe (line-oriented, bounded). + /// Uses `LinesCodec::new_with_max_length` to enforce MAX_LINE_SIZE at the + /// read level — prevents OOM from rogue agents writing infinite non-newline bytes. + reader: FramedRead, /// Monotonically increasing JSON-RPC request id counter. /// Harness-generated IDs are always numeric. next_id: u64, @@ -131,8 +148,21 @@ impl AcpClient { /// Call this when you need guaranteed cleanup — e.g., in `run_models` /// before process exit. pub async fn shutdown(&mut self) { + // Send SIGKILL to the direct child. On Unix, process_group(0) was set + // at spawn time, so the child runs in its own process group. Ideally + // we'd kill the entire group via kill(-pgid, SIGKILL), but that + // requires `unsafe` (libc::kill) which is denied by #![deny(unsafe_code)]. + // When the direct child dies, its subprocesses become orphans adopted + // by init — acceptable for a harness that respawns or exits. let _ = self.child.start_kill(); - let _ = self.child.wait().await; + // Bounded wait: if the child doesn't exit within 5s after SIGKILL, + // give up and let Drop/OS handle it. An unbounded wait here would + // wedge the harness during respawn or shutdown if a child is stuck. + match tokio::time::timeout(std::time::Duration::from_secs(5), self.child.wait()).await { + Ok(Ok(_)) => {} + Ok(Err(e)) => tracing::debug!("child wait error after kill: {e}"), + Err(_) => tracing::warn!("child did not exit within 5s after SIGKILL — abandoning"), + } } /// Spawn the agent binary as a subprocess and connect to its stdio pipes. @@ -146,7 +176,16 @@ impl AcpClient { .stdin(Stdio::piped()) .stdout(Stdio::piped()) // Inherit stderr so agent logs are visible in the harness terminal. - .stderr(Stdio::inherit()); + .stderr(Stdio::inherit()) + // Ensure the child is killed when the AcpClient is dropped (best-effort). + // Callers MUST still call shutdown().await for guaranteed cleanup. + .kill_on_drop(true); + + // Spawn the agent in its own process group so SIGKILL doesn't propagate + // to the harness's own process group on Unix. + // tokio::process::Command::process_group is a stable tokio API (no extra imports needed). + #[cfg(unix)] + cmd.process_group(0); let mut child = cmd.spawn()?; @@ -162,7 +201,7 @@ impl AcpClient { Ok(Self { child, stdin, - reader: BufReader::new(stdout), + reader: FramedRead::new(stdout, LinesCodec::new_with_max_length(MAX_LINE_SIZE)), next_id: 0, pending_permission_id: None, permission_responded: false, @@ -343,7 +382,7 @@ impl AcpClient { pub async fn cancel_with_cleanup( &mut self, session_id: &str, - idle_timeout: std::time::Duration, + _idle_timeout: std::time::Duration, ) -> Result { // Inherit the hard deadline from the timed-out turn so the drain loop // doesn't start a fresh timer (prevents double-jeopardy). If the original @@ -392,9 +431,13 @@ impl AcpClient { // Step 2: send session/cancel notification (no id) self.session_cancel(session_id).await?; tracing::info!(target: "acp::cancel", "sent session/cancel for {session_id}"); - // Clamp idle timeout to at least 30s during cleanup — the cancel notification + // Use a fixed 30s idle timeout during cleanup — the cancel notification // needs time to propagate and the agent may go silent while winding down. - let cleanup_idle = idle_timeout.max(std::time::Duration::from_secs(30)); + // We do NOT use the caller's idle_timeout here: that value was tuned for + // normal prompt activity and may be very short (e.g. 5s). Using it during + // cancel would cause premature IdleTimeout before the cancelled response + // arrives, leaving the session in an inconsistent state. + let cleanup_idle = std::time::Duration::from_secs(30); let result = self .read_until_response_with_idle_timeout(prompt_id, cleanup_idle, hard_deadline) .await?; @@ -404,18 +447,35 @@ impl AcpClient { // ── Internal helpers ────────────────────────────────────────────────── /// Serialize `value` as a single NDJSON line and flush to the agent's stdin. + /// + /// Bounded by a 30-second write timeout. If the agent stops reading stdin + /// (e.g., it's stuck or dead), the write would otherwise block forever. async fn write_ndjson(&mut self, value: &serde_json::Value) -> Result<(), AcpError> { + const WRITE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30); let line = serde_json::to_string(value)?; - self.stdin.write_all(line.as_bytes()).await?; - self.stdin.write_all(b"\n").await?; - self.stdin.flush().await?; - Ok(()) + tokio::time::timeout(WRITE_TIMEOUT, async { + self.stdin.write_all(line.as_bytes()).await?; + self.stdin.write_all(b"\n").await?; + self.stdin.flush().await?; + Ok::<(), std::io::Error>(()) + }) + .await + .map_err(|_| AcpError::WriteTimeout(WRITE_TIMEOUT))? + .map_err(AcpError::Io) } + /// Default timeout for non-prompt RPCs (initialize, session/new, etc.). + const REQUEST_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(60); + /// Send a JSON-RPC request and wait for the matching response. /// /// Assigns the next available id, writes the NDJSON line to stdin, /// then calls [`read_until_response`](Self::read_until_response). + /// + /// The write phase is bounded by `WRITE_TIMEOUT` (30s) and the read phase + /// by `REQUEST_TIMEOUT` (60s), so worst-case wall clock is ~90s. Non-prompt + /// RPCs like `initialize` and `session/new` should complete in seconds; + /// if they don't, the agent is likely stuck and we must not block forever. async fn send_request( &mut self, method: &str, @@ -432,9 +492,52 @@ impl AcpClient { }); tracing::debug!(target: "acp::wire", "→ {}", &serde_json::to_string(&msg).unwrap_or_default()); - self.write_ndjson(&msg).await?; - self.read_until_response(id).await + // Wrap write + read in a single timeout so a hung agent can't block forever. + // We cannot use an async block that borrows `self` mutably across two awaits + // inside timeout(), so we sequence them with early-return on timeout. + let timeout = Self::REQUEST_TIMEOUT; + match tokio::time::timeout(timeout, self.write_ndjson(&msg)).await { + Ok(result) => result?, + Err(_) => return Err(AcpError::Timeout(timeout)), + } + + match tokio::time::timeout(timeout, self.read_until_response(id)).await { + Ok(result) => result, + Err(_) => Err(AcpError::Timeout(timeout)), + } + } + + /// Drain any buffered lines from the agent's stdout without blocking. + /// + /// After a [`AcpError::Timeout`] from [`send_request`], the agent may + /// eventually send the late response. That stale message will sit in the + /// `BufReader` buffer and be silently skipped by the next `read_until_response` + /// call (ID mismatch). However, if the caller wants a clean slate — e.g. + /// before retrying the same method — they can call this to consume any + /// buffered data with a short deadline. + /// + /// This is a best-effort drain: it reads until the buffer is empty or + /// `drain_timeout` elapses, whichever comes first. Errors are ignored. + #[allow(dead_code)] // Scaffolding for future model-switch timeout cleanup; not yet wired. + pub async fn drain_stale_responses(&mut self, drain_timeout: std::time::Duration) { + let deadline = tokio::time::Instant::now() + drain_timeout; + loop { + let remaining = deadline.saturating_duration_since(tokio::time::Instant::now()); + if remaining.is_zero() { + break; + } + let read_result = tokio::time::timeout(remaining, self.reader.next()).await; + match read_result { + // Timeout or stream ended — buffer is empty or agent exited. + Err(_) | Ok(None) => break, + Ok(Some(Ok(_))) => { + // Consumed one buffered line; loop to drain more. + tracing::debug!(target: "acp::wire", "drained stale buffered line"); + } + Ok(Some(Err(_))) => break, + } + } } /// Send a JSON-RPC **notification** — no `id` field, no response expected. @@ -463,7 +566,8 @@ impl AcpClient { /// While waiting, handles: /// - `session/update` notifications → logged via tracing /// - `session/request_permission` requests → auto-approved with `allow_once` - /// - Any other messages → debug-logged and ignored + /// - Any other messages → debug-logged and ignored; if they carry an `id` + /// (i.e. they are requests, not notifications), a JSON-RPC -32601 error is sent. /// /// Compares the incoming `id` field as a `serde_json::Value` against /// `json!(expected_id)` so that both numeric and string IDs work correctly. @@ -472,17 +576,28 @@ impl AcpClient { expected_id: u64, ) -> Result { loop { - let mut line = String::new(); - let n = self.reader.read_line(&mut line).await?; - if n == 0 { - return Err(AcpError::AgentExited); - } + // LinesCodec::new_with_max_length enforces MAX_LINE_SIZE at the + // read level — the buffer never grows beyond the limit, preventing + // OOM from rogue agents writing infinite non-newline bytes. + let line = match self.reader.next().await { + None => return Err(AcpError::AgentExited), + Some(Err(LinesCodecError::MaxLineLengthExceeded)) => { + return Err(AcpError::Protocol( + "agent stdout line exceeded 10MB limit".into(), + )); + } + Some(Err(e)) => { + return Err(AcpError::Io(std::io::Error::other(e))); + } + Some(Ok(line)) => line, + }; let trimmed = line.trim(); if trimmed.is_empty() { continue; } + // Only log and reset idle after we have a valid non-empty line. tracing::debug!(target: "acp::wire", "← {trimmed}"); let msg: serde_json::Value = match serde_json::from_str(trimmed) { @@ -519,6 +634,19 @@ impl AcpClient { self.handle_permission_request(&msg).await?; } other => { + // If the unknown message has an id, it's a request expecting a reply. + // Silence would cause the agent to hang waiting for a response. + // Send a JSON-RPC -32601 "Method not found" error. + if msg.get("id").is_some() { + let err_resp = serde_json::json!({ + "jsonrpc": "2.0", + "id": msg["id"], + "error": {"code": -32601, "message": format!("Method not found: {other}")} + }); + // Surface write failures — a broken pipe means the + // agent process is dead and continuing would hang. + self.write_ndjson(&err_resp).await?; + } tracing::debug!(target: "acp::wire", "ignoring unknown method: {other}"); } } @@ -554,19 +682,21 @@ impl AcpClient { }; let remaining = next_deadline.saturating_duration_since(Instant::now()); - let read_result = tokio::time::timeout(remaining, async { - let mut line = String::new(); - let n = self.reader.read_line(&mut line).await?; - Ok::<(usize, String), std::io::Error>((n, line)) - }) - .await; + // LinesCodec::new_with_max_length enforces MAX_LINE_SIZE at the + // read level — the buffer never grows beyond the limit. + let read_result = tokio::time::timeout(remaining, self.reader.next()).await; match read_result { - Ok(Ok((0, _))) => return Err(AcpError::AgentExited), - Ok(Ok((_, line))) => { - // Any stdout activity resets the idle clock. - idle_deadline = Instant::now() + idle_timeout; - + Ok(None) => return Err(AcpError::AgentExited), + Ok(Some(Err(LinesCodecError::MaxLineLengthExceeded))) => { + return Err(AcpError::Protocol( + "agent stdout line exceeded 10MB limit".into(), + )); + } + Ok(Some(Err(e))) => { + return Err(AcpError::Io(std::io::Error::other(e))); + } + Ok(Some(Ok(line))) => { let trimmed = line.trim(); if trimmed.is_empty() { continue; @@ -585,6 +715,10 @@ impl AcpClient { } }; + // Only reset the idle clock on lines that parse as valid JSON. + // Malformed lines (skipped above) don't count as real agent activity. + idle_deadline = Instant::now() + idle_timeout; + // Check for matching response. if let Some(id) = msg.get("id") { if *id == serde_json::json!(expected_id) { @@ -595,7 +729,7 @@ impl AcpClient { } } - // Dispatch notifications. + // Dispatch notifications and agent-initiated requests. if let Some(method) = msg.get("method").and_then(|v| v.as_str()) { match method { "session/update" => self.handle_session_update(&msg), @@ -603,12 +737,24 @@ impl AcpClient { self.handle_permission_request(&msg).await?; } other => { + // If the unknown message has an id, it's a request expecting a reply. + // Silence would cause the agent to hang waiting for a response. + // Send a JSON-RPC -32601 "Method not found" error. + if msg.get("id").is_some() { + let err_resp = serde_json::json!({ + "jsonrpc": "2.0", + "id": msg["id"], + "error": {"code": -32601, "message": format!("Method not found: {other}")} + }); + // Surface write failures — a broken pipe means the + // agent process is dead and continuing would hang. + self.write_ndjson(&err_resp).await?; + } tracing::debug!(target: "acp::wire", "ignoring unknown method: {other}"); } } } } - Ok(Err(e)) => return Err(AcpError::Io(e)), Err(_elapsed) => { // Classification was determined before sleeping — not // affected by scheduler jitter between deadline and wakeup. @@ -739,11 +885,22 @@ impl AcpClient { } }; - // Mark as responded BEFORE the write — if the future is dropped by a - // timeout between write and flag-set, cancel_with_cleanup would otherwise - // see permission_responded=false and send a second response (protocol violation). - self.permission_responded = true; + // Write the response first, then mark as responded. + // + // Previous ordering (flag-before-write) was intended to guard against a + // double-response if a timeout fires between write and flag-set. However, + // the deadlock risk is worse: if write_ndjson fails (e.g. WriteTimeout), + // the flag would be true but no response was actually sent. Then + // cancel_with_cleanup would see permission_responded=true, skip sending + // the cancelled outcome, and the agent would hang waiting for a reply + // that never arrives — a guaranteed deadlock. + // + // The correct fix: set the flag AFTER a successful write. The double- + // response window (between write completion and flag-set) is negligibly + // small and bounded by a single memory store; the deadlock window was + // unbounded. self.write_ndjson(&response).await?; + self.permission_responded = true; self.pending_permission_id = None; Ok(()) } @@ -875,8 +1032,18 @@ pub fn resolve_model_switch_method( impl Drop for AcpClient { fn drop(&mut self) { - // Best-effort kill — ignore errors (process may have already exited). + // Best-effort SIGKILL + reap. We cannot `await` in Drop (sync context), + // so we use `start_kill()` followed by `try_wait()`. The `try_wait()` + // reaps the child if it has already exited (likely after SIGKILL on most + // platforms). If the process hasn't exited yet, it will be reaped by the + // OS when the harness process itself exits (init adopts orphans). + // + // Callers SHOULD still call `shutdown().await` whenever possible for + // guaranteed reaping. Drop is the safety net for panics and early exits. let _ = self.child.start_kill(); + // Non-blocking reap attempt — prevents zombie accumulation in the + // common case where SIGKILL takes effect before Drop returns. + let _ = self.child.try_wait(); } } @@ -910,8 +1077,26 @@ mod tests { fn stop_reason_returns_none_for_unknown() { assert_eq!(StopReason::from_str("unknown_value"), None); assert_eq!(StopReason::from_str(""), None); - assert_eq!(StopReason::from_str("END_TURN"), None); // case-sensitive - assert_eq!(StopReason::from_str("endturn"), None); // no camelCase + assert_eq!(StopReason::from_str("endturn"), None); // no camelCase — still unknown + } + + #[test] + fn stop_reason_is_case_insensitive() { + // Agents may send uppercase or mixed-case variants — all should parse correctly. + assert_eq!(StopReason::from_str("END_TURN"), Some(StopReason::EndTurn)); + assert_eq!( + StopReason::from_str("CANCELLED"), + Some(StopReason::Cancelled) + ); + assert_eq!( + StopReason::from_str("Max_Tokens"), + Some(StopReason::MaxTokens) + ); + assert_eq!( + StopReason::from_str("MAX_TURN_REQUESTS"), + Some(StopReason::MaxTurnRequests) + ); + assert_eq!(StopReason::from_str("Refusal"), Some(StopReason::Refusal)); } // ── Permission option finding ───────────────────────────────────────── @@ -1436,9 +1621,12 @@ mod tests { #[tokio::test] async fn idle_resets_on_stdout_activity() { - let mut client = - spawn_script("for i in $(seq 1 10); do echo 'keepalive'; sleep 0.05; done; sleep 10") - .await; + // Send valid JSON (session/update notifications) to reset the idle timer. + // Non-JSON lines no longer reset idle (Finding #6 hardening). + let mut client = spawn_script( + r#"for i in $(seq 1 10); do echo '{"jsonrpc":"2.0","method":"session/update","params":{"update":{"sessionUpdate":"agent_thought_chunk","content":{"text":"thinking"}}}}'; sleep 0.05; done; sleep 10"#, + ) + .await; let hard_deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(10); let start = std::time::Instant::now(); let result = client @@ -1449,6 +1637,7 @@ mod tests { ) .await; let elapsed = start.elapsed(); + // 10 messages × 50ms = ~500ms of activity, then idle timeout fires after 200ms more assert!(elapsed >= std::time::Duration::from_millis(400)); assert!(elapsed < std::time::Duration::from_secs(3)); assert!(matches!(result, Err(AcpError::IdleTimeout(_)))); diff --git a/crates/sprout-acp/src/config.rs b/crates/sprout-acp/src/config.rs index 0e567d6c3a..0450f273f9 100644 --- a/crates/sprout-acp/src/config.rs +++ b/crates/sprout-acp/src/config.rs @@ -380,27 +380,42 @@ pub fn normalize_agent_args(command: &str, agent_args: Vec) -> Vec Result { - // Propagate legacy env vars so clap sees them under the canonical names. - // Only set if the canonical var is absent — clap's `env` reads these. - // NOTE: std::env::set_var is safe in edition 2021 but will require - // `unsafe` in edition 2024. This runs inside #[tokio::main] so worker - // threads are already alive. When migrating to edition 2024, move this - // to a non-async fn main() wrapper that runs before the runtime starts. - for (legacy, canonical) in [ - ("SPROUT_ACP_PRIVATE_KEY", "SPROUT_PRIVATE_KEY"), - ("SPROUT_ACP_API_TOKEN", "SPROUT_API_TOKEN"), - ] { - if std::env::var(canonical).is_err() { - if let Ok(val) = std::env::var(legacy) { - std::env::set_var(canonical, &val); - } +/// Propagate legacy env-var aliases to their canonical names. +/// +/// Must be called **before** the tokio runtime starts — i.e. from the sync +/// `fn main()` wrapper, not from inside `#[tokio::main]`. +/// +/// `std::env::set_var` is safe in Rust 2021 only when no other threads are +/// running. In Rust 2024 it requires `unsafe`. Calling this before +/// `#[tokio::main]` ensures worker threads are not yet alive. +/// +/// // Must be called before tokio runtime starts — see Rust 2024 edition safety. +pub fn propagate_legacy_env_vars() { + for (legacy, canonical) in [ + ("SPROUT_ACP_PRIVATE_KEY", "SPROUT_PRIVATE_KEY"), + ("SPROUT_ACP_API_TOKEN", "SPROUT_API_TOKEN"), + ] { + if std::env::var(canonical).is_err() { + if let Ok(val) = std::env::var(legacy) { + std::env::set_var(canonical, &val); } } + } +} - let args = CliArgs::parse(); +impl Config { + pub fn from_cli() -> Result { + // Legacy env-var propagation is intentionally NOT done here. + // Call `propagate_legacy_env_vars()` before the tokio runtime starts + // (in the sync `fn main()` wrapper) — see Rust 2024 edition safety. + let mut args = CliArgs::parse(); let keys = Keys::parse(&args.private_key)?; + // Best-effort zeroize: overwrite the raw private key string to reduce + // exposure via core dumps or heap inspection (#41). Without the `zeroize` + // crate we can only clear the String — the allocator may retain copies. + args.private_key + .replace_range(.., &"0".repeat(args.private_key.len())); + args.private_key.clear(); let system_prompt = if let Some(text) = args.system_prompt { Some(text) @@ -437,55 +452,99 @@ impl Config { } let agent_command = args.agent_command; + + // Finding #49a — agent_command must not be empty. + if agent_command.trim().is_empty() { + return Err(ConfigError::ConfigFile( + "agent_command must not be empty".into(), + )); + } + let agent_args = normalize_agent_args(&agent_command, args.agent_args); - Ok(Config { + // Finding #49b — warn on invalid UUIDs in --channels. + if let Some(ref channels) = args.channels { + for ch in channels { + if ch.parse::().is_err() { + tracing::warn!( + channel = %ch, + "--channels entry is not a valid UUID and will be ignored" + ); + } + } + } + + // Finding #49c — cap heartbeat interval at 86400s (24h). + let heartbeat_interval = if args.heartbeat_interval > 86400 { + tracing::warn!( + interval = args.heartbeat_interval, + "heartbeat interval exceeds 24h — capping at 86400s" + ); + 86400u64 + } else { + args.heartbeat_interval + }; + + // Resolve idle_timeout_secs with deprecation handling. + // Precedence: explicit --idle-timeout > --turn-timeout (deprecated) > default 300. + let idle_timeout_secs = { + let raw = match (args.idle_timeout, args.turn_timeout) { + (Some(idle), Some(_turn)) => { + tracing::warn!( + "--turn-timeout / SPROUT_ACP_TURN_TIMEOUT is deprecated and ignored \ + when --idle-timeout / SPROUT_ACP_IDLE_TIMEOUT is also set" + ); + idle + } + (Some(idle), None) => idle, + (None, Some(turn)) => { + tracing::warn!( + "--turn-timeout / SPROUT_ACP_TURN_TIMEOUT is deprecated; \ + use --idle-timeout / SPROUT_ACP_IDLE_TIMEOUT instead" + ); + turn + } + (None, None) => 300, // default + }; + if raw == 0 { + tracing::warn!("idle timeout of 0 is invalid — using 1s minimum"); + 1 + } else { + raw + } + }; + + let max_turn_duration_secs = { + let raw = args.max_turn_duration; + if raw == 0 { + tracing::warn!("max turn duration of 0 is invalid — using 60s minimum"); + 60 + } else { + raw + } + }; + + // Finding #20 — idle_timeout must be strictly less than max_turn_duration. + // If idle_timeout >= max_turn_duration, the absolute wall-clock cap would + // fire before the idle timeout ever could, making idle_timeout a dead letter. + if idle_timeout_secs >= max_turn_duration_secs { + return Err(ConfigError::ConfigFile(format!( + "idle_timeout ({}s) must be less than max_turn_duration ({}s)", + idle_timeout_secs, max_turn_duration_secs + ))); + } + + let config = Config { keys, api_token: args.api_token, relay_url: args.relay_url, agent_command, agent_args, mcp_command: args.mcp_command, - // Deprecated --turn-timeout is a fallback for backward compat. - // New deployments should use --idle-timeout exclusively. - // Precedence: explicit --idle-timeout > --turn-timeout (deprecated) > default 300. - idle_timeout_secs: { - let raw = match (args.idle_timeout, args.turn_timeout) { - (Some(idle), Some(_turn)) => { - tracing::warn!( - "--turn-timeout / SPROUT_ACP_TURN_TIMEOUT is deprecated and ignored \ - when --idle-timeout / SPROUT_ACP_IDLE_TIMEOUT is also set" - ); - idle - } - (Some(idle), None) => idle, - (None, Some(turn)) => { - tracing::warn!( - "--turn-timeout / SPROUT_ACP_TURN_TIMEOUT is deprecated; \ - use --idle-timeout / SPROUT_ACP_IDLE_TIMEOUT instead" - ); - turn - } - (None, None) => 300, // default - }; - if raw == 0 { - tracing::warn!("idle timeout of 0 is invalid — using 1s minimum"); - 1 - } else { - raw - } - }, - max_turn_duration_secs: { - let raw = args.max_turn_duration; - if raw == 0 { - tracing::warn!("max turn duration of 0 is invalid — using 60s minimum"); - 60 - } else { - raw - } - }, + idle_timeout_secs, + max_turn_duration_secs, agents: args.agents, - heartbeat_interval_secs: args.heartbeat_interval, + heartbeat_interval_secs: heartbeat_interval, heartbeat_prompt, system_prompt, initial_message: args.initial_message, @@ -502,7 +561,9 @@ impl Config { typing_enabled: !args.no_typing, model: args.model, permission_mode: args.permission_mode, - }) + }; + + Ok(config) } /// Human-readable summary (no secrets). @@ -540,8 +601,11 @@ struct TomlConfig { } pub fn load_rules(path: &std::path::Path) -> Result, ConfigError> { + use std::sync::atomic::AtomicU32; + use std::sync::Arc; + let content = std::fs::read_to_string(path)?; - let config: TomlConfig = + let mut config: TomlConfig = toml::from_str(&content).map_err(|e| ConfigError::ConfigFile(e.to_string()))?; if config.rules.len() > 100 { @@ -551,14 +615,22 @@ pub fn load_rules(path: &std::path::Path) -> Result, Confi ))); } + // Finding #49d — warn when Config mode has no rules; agent will receive nothing. + if config.rules.is_empty() { + tracing::warn!( + path = %path.display(), + "config file contains zero rules — agent will receive no events in Config mode" + ); + } + let mut seen_names = std::collections::HashSet::new(); - for rule in &config.rules { + for rule in &mut config.rules { if rule.name.trim().is_empty() { return Err(ConfigError::ConfigFile( "rule name must not be empty".into(), )); } - if !seen_names.insert(&rule.name) { + if !seen_names.insert(rule.name.clone()) { return Err(ConfigError::ConfigFile(format!( "duplicate rule name: {}", rule.name @@ -574,11 +646,17 @@ pub fn load_rules(path: &std::path::Path) -> Result, Confi } // Fail fast: parse the expression at load time so typos don't // silently produce dead rules at runtime. - if let Err(e) = evalexpr::build_operator_tree(expr) { - return Err(ConfigError::ConfigFile(format!( - "rule '{}': invalid filter expression: {e}", - rule.name, - ))); + // Finding #34 — store the compiled AST so match_event never re-parses. + match evalexpr::build_operator_tree(expr) { + Ok(node) => { + rule.compiled_filter = Some(Arc::new(node)); + } + Err(e) => { + return Err(ConfigError::ConfigFile(format!( + "rule '{}': invalid filter expression: {e}", + rule.name, + ))); + } } } // Validate channel scope — catch typos like "ALL" or "All" early. @@ -590,6 +668,10 @@ pub fn load_rules(path: &std::path::Path) -> Result, Confi ))); } } + // Initialise the consecutive-timeout counter (finding #25). + // Deserialization leaves it at default (new Arc) + // but we set it explicitly here for clarity. + rule.consecutive_timeouts = Arc::new(AtomicU32::new(0)); } Ok(config.rules) @@ -836,6 +918,8 @@ mod tests { kinds: Vec, mention: bool, ) -> SubscriptionRule { + use std::sync::atomic::AtomicU32; + use std::sync::Arc; SubscriptionRule { name: name.into(), channels, @@ -843,6 +927,8 @@ mod tests { require_mention: mention, filter: None, prompt_tag: None, + compiled_filter: None, + consecutive_timeouts: Arc::new(AtomicU32::new(0)), } } diff --git a/crates/sprout-acp/src/filter.rs b/crates/sprout-acp/src/filter.rs index 446552771f..b99120c7be 100644 --- a/crates/sprout-acp/src/filter.rs +++ b/crates/sprout-acp/src/filter.rs @@ -5,9 +5,11 @@ //! - Evaluating boolean filter expressions with a hard timeout //! - Matching events against ordered subscription rules (first match wins) +use std::sync::atomic::{AtomicU32, Ordering}; +use std::sync::Arc; use std::time::Duration; -use tracing::warn; +use tracing::{error, warn}; // ── Errors ──────────────────────────────────────────────────────────────────── @@ -77,7 +79,13 @@ impl ChannelScope { } /// A single subscription rule from the agent config. -#[derive(Debug, Clone, serde::Deserialize)] +/// +/// # Thread safety +/// +/// `consecutive_timeouts` is an `AtomicU32` so `match_event` can update it +/// without requiring `&mut self` — rules are shared via `Arc<[SubscriptionRule]>` +/// across the event-dispatch loop. +#[derive(Debug, serde::Deserialize)] pub struct SubscriptionRule { /// Human-readable rule name; used as fallback `prompt_tag`. pub name: String, @@ -95,6 +103,52 @@ pub struct SubscriptionRule { /// Tag passed to the prompt template. Falls back to `name` if absent. #[serde(default)] pub prompt_tag: Option, + /// Pre-compiled evalexpr AST for the `filter` expression (finding #34). + /// + /// Populated by `load_rules()` at startup so `match_event` never re-parses + /// the expression string on the hot path. `None` when `filter` is `None` + /// or the rule was constructed without calling `load_rules()` (e.g. tests). + #[serde(skip)] + pub compiled_filter: Option>, + /// Consecutive filter-evaluation timeout counter (finding #25). + /// + /// Incremented on each timeout; reset on any successful evaluation. + /// When this reaches `MAX_CONSECUTIVE_TIMEOUTS`, the rule is treated as + /// disabled and `match_event` returns `None` (fail-closed). + #[serde(skip)] + pub consecutive_timeouts: Arc, +} + +impl Default for SubscriptionRule { + fn default() -> Self { + Self { + name: String::new(), + channels: ChannelScope::All("all".into()), + kinds: Vec::new(), + require_mention: false, + filter: None, + prompt_tag: None, + compiled_filter: None, + consecutive_timeouts: Arc::new(AtomicU32::new(0)), + } + } +} + +impl Clone for SubscriptionRule { + fn clone(&self) -> Self { + Self { + name: self.name.clone(), + channels: self.channels.clone(), + kinds: self.kinds.clone(), + require_mention: self.require_mention, + filter: self.filter.clone(), + prompt_tag: self.prompt_tag.clone(), + compiled_filter: self.compiled_filter.clone(), + // Share the same counter across clones so all copies of a rule + // agree on the timeout state. + consecutive_timeouts: self.consecutive_timeouts.clone(), + } + } } // ── MatchedRule ─────────────────────────────────────────────────────────────── @@ -120,13 +174,41 @@ const MAX_EXPR_LEN: usize = 4096; /// Maximum wall-clock time allowed for a single evalexpr evaluation. const EVAL_TIMEOUT: Duration = Duration::from_millis(100); +/// Maximum concurrent blocking filter evaluations (finding #13 / Issue 3). +/// +/// The semaphore permit is moved into each `spawn_blocking` closure so it is +/// held until the blocking thread finishes — not just until the caller's timeout +/// fires. This truly bounds the number of live blocking evals even under repeated +/// slow expressions. +const MAX_CONCURRENT_FILTER_EVALS: usize = 4; + +/// Semaphore that bounds concurrent `spawn_blocking` filter evaluations. +/// +/// Wrapped in `Arc` so `acquire_owned()` can be used, which returns an +/// `OwnedSemaphorePermit` that can be moved into the `spawn_blocking` closure. +/// This ensures the permit is held until the blocking task actually finishes — +/// not just until the caller's timeout fires — so the semaphore truly bounds +/// the number of live blocking threads (finding #13 / Issue 3). +static FILTER_EVAL_SEMAPHORE: std::sync::LazyLock> = + std::sync::LazyLock::new(|| Arc::new(tokio::sync::Semaphore::new(MAX_CONCURRENT_FILTER_EVALS))); + /// Evaluate a boolean filter expression against a `FilterContext`. /// /// - Caps expression length at [`MAX_EXPR_LEN`] bytes. +/// - Acquires an owned permit from [`FILTER_EVAL_SEMAPHORE`] and moves it into +/// the blocking closure so it is held until the task finishes, not just until +/// the caller's timeout fires (finding #13 / Issue 3). /// - Runs evaluation on a blocking thread with a [`EVAL_TIMEOUT`] hard timeout. +/// - When a pre-compiled `node` is provided (via `Arc`), uses +/// `node.eval_boolean_with_context()` instead of re-parsing the expression +/// string on every call (finding #34). /// - Registers custom string helpers: `str_contains`, `str_starts_with`, /// `str_ends_with`, `str_len` (duplicated intentionally from sprout-workflow). -pub async fn evaluate_filter(expr: &str, ctx: &FilterContext) -> Result { +pub async fn evaluate_filter( + expr: &str, + ctx: &FilterContext, + node: Option>, +) -> Result { if expr.len() > MAX_EXPR_LEN { return Err(FilterError::ExpressionTooLong { len: expr.len(), @@ -137,10 +219,33 @@ pub async fn evaluate_filter(expr: &str, ctx: &FilterContext) -> Result Result Result {} - Ok(false) => continue, - Err(e) => { - warn!(rule = %rule.name, error = %e, "filter expression error; skipping rule"); + // Skip rules that have timed out too many times — treat as disabled. + let prior_timeouts = rule.consecutive_timeouts.load(Ordering::Relaxed); + if prior_timeouts >= MAX_CONSECUTIVE_TIMEOUTS { + error!( + rule = %rule.name, + rule_index = index, + timeouts = prior_timeouts, + "filter rule disabled after too many consecutive timeouts; \ + failing closed (no match for any rule)" + ); + // Fail-closed: disabled rule → no match for this event. + return None; + } + + match evaluate_filter(expr, &filter_ctx, rule.compiled_filter.clone()).await { + Ok(true) => { + // Successful match — reset timeout counter. + rule.consecutive_timeouts.store(0, Ordering::Relaxed); + } + Ok(false) => { + rule.consecutive_timeouts.store(0, Ordering::Relaxed); continue; } + Err(FilterError::Timeout) => { + let n = rule.consecutive_timeouts.fetch_add(1, Ordering::Relaxed) + 1; + warn!( + rule = %rule.name, + rule_index = index, + consecutive_timeouts = n, + "filter expression timed out; failing closed (no match for any rule)" + ); + // Fail-closed: timeout → no match, not next rule. + return None; + } + Err(e) => { + warn!( + rule = %rule.name, + rule_index = index, + error = %e, + "filter expression error; failing closed (no match for any rule)" + ); + // Fail-closed: any error → no match, not next rule. + return None; + } } } @@ -337,6 +503,26 @@ mod tests { Uuid::new_v4() } + fn make_rule( + name: &str, + channels: ChannelScope, + kinds: Vec, + mention: bool, + filter: Option<&str>, + prompt_tag: Option<&str>, + ) -> SubscriptionRule { + SubscriptionRule { + name: name.into(), + channels, + kinds, + require_mention: mention, + filter: filter.map(|s| s.into()), + prompt_tag: prompt_tag.map(|s| s.into()), + compiled_filter: None, + consecutive_timeouts: Arc::new(AtomicU32::new(0)), + } + } + // ── FilterContext ───────────────────────────────────────────────────────── #[test] @@ -359,12 +545,12 @@ mod tests { let event = make_event(9, "P1 incident in production"); let ctx = FilterContext::from_event(&event, any_channel()); - let result = evaluate_filter(r#"str_contains(content, "P1")"#, &ctx) + let result = evaluate_filter(r#"str_contains(content, "P1")"#, &ctx, None) .await .unwrap(); assert!(result); - let result = evaluate_filter(r#"str_contains(content, "P2")"#, &ctx) + let result = evaluate_filter(r#"str_contains(content, "P2")"#, &ctx, None) .await .unwrap(); assert!(!result); @@ -375,10 +561,10 @@ mod tests { let event = make_event(9, "some content"); let ctx = FilterContext::from_event(&event, any_channel()); - let result = evaluate_filter("kind == 9", &ctx).await.unwrap(); + let result = evaluate_filter("kind == 9", &ctx, None).await.unwrap(); assert!(result); - let result = evaluate_filter("kind == 1", &ctx).await.unwrap(); + let result = evaluate_filter("kind == 1", &ctx, None).await.unwrap(); assert!(!result); } @@ -388,7 +574,7 @@ mod tests { let ctx = FilterContext::from_event(&event, any_channel()); let long_expr = "a".repeat(MAX_EXPR_LEN + 1); - let err = evaluate_filter(&long_expr, &ctx).await.unwrap_err(); + let err = evaluate_filter(&long_expr, &ctx, None).await.unwrap_err(); assert!(matches!( err, @@ -397,6 +583,19 @@ mod tests { )); } + #[tokio::test] + async fn test_evaluate_filter_precompiled_node() { + let event = make_event(9, "hello world"); + let ctx = FilterContext::from_event(&event, any_channel()); + + let node = + Arc::new(evalexpr::build_operator_tree(r#"str_contains(content, "hello")"#).unwrap()); + let result = evaluate_filter(r#"str_contains(content, "hello")"#, &ctx, Some(node)) + .await + .unwrap(); + assert!(result); + } + // ── match_event ─────────────────────────────────────────────────────────── #[tokio::test] @@ -405,22 +604,22 @@ mod tests { let channel_id = any_channel(); let rules = vec![ - SubscriptionRule { - name: "first".into(), - channels: ChannelScope::All("all".into()), - kinds: vec![], - require_mention: false, - filter: None, - prompt_tag: Some("tag-first".into()), - }, - SubscriptionRule { - name: "second".into(), - channels: ChannelScope::All("all".into()), - kinds: vec![], - require_mention: false, - filter: None, - prompt_tag: Some("tag-second".into()), - }, + make_rule( + "first", + ChannelScope::All("all".into()), + vec![], + false, + None, + Some("tag-first"), + ), + make_rule( + "second", + ChannelScope::All("all".into()), + vec![], + false, + None, + Some("tag-second"), + ), ]; let matched = match_event(&event, channel_id, &rules, "").await.unwrap(); @@ -434,22 +633,22 @@ mod tests { let channel_id = any_channel(); let rules = vec![ - SubscriptionRule { - name: "wrong-kind".into(), - channels: ChannelScope::All("all".into()), - kinds: vec![1], - require_mention: false, - filter: None, - prompt_tag: None, - }, - SubscriptionRule { - name: "right-kind".into(), - channels: ChannelScope::All("all".into()), - kinds: vec![9], - require_mention: false, - filter: None, - prompt_tag: Some("matched".into()), - }, + make_rule( + "wrong-kind", + ChannelScope::All("all".into()), + vec![1], + false, + None, + None, + ), + make_rule( + "right-kind", + ChannelScope::All("all".into()), + vec![9], + false, + None, + Some("matched"), + ), ]; let matched = match_event(&event, channel_id, &rules, "").await.unwrap(); @@ -465,14 +664,14 @@ mod tests { let event_with_mention = make_event_with_p_tag(9, "hello", agent_pubkey); let channel_id = any_channel(); - let rules = vec![SubscriptionRule { - name: "mention-only".into(), - channels: ChannelScope::All("all".into()), - kinds: vec![], - require_mention: true, - filter: None, - prompt_tag: Some("mentioned".into()), - }]; + let rules = vec![make_rule( + "mention-only", + ChannelScope::All("all".into()), + vec![], + true, + None, + Some("mentioned"), + )]; // Without mention — no match. let result = match_event(&event_no_mention, channel_id, &rules, agent_pubkey).await; @@ -490,14 +689,14 @@ mod tests { let event = make_event(1, "hello"); let channel_id = any_channel(); - let rules = vec![SubscriptionRule { - name: "kind-9-only".into(), - channels: ChannelScope::All("all".into()), - kinds: vec![9], - require_mention: false, - filter: None, - prompt_tag: None, - }]; + let rules = vec![make_rule( + "kind-9-only", + ChannelScope::All("all".into()), + vec![9], + false, + None, + None, + )]; let result = match_event(&event, channel_id, &rules, "").await; assert!(result.is_none()); @@ -542,16 +741,74 @@ mod tests { let event = make_event(9, "hello"); let channel_id = any_channel(); - let rules = vec![SubscriptionRule { - name: "my-rule".into(), - channels: ChannelScope::All("all".into()), - kinds: vec![], - require_mention: false, - filter: None, - prompt_tag: None, // no explicit tag - }]; + let rules = vec![make_rule( + "my-rule", + ChannelScope::All("all".into()), + vec![], + false, + None, + None, // no explicit tag + )]; let matched = match_event(&event, channel_id, &rules, "").await.unwrap(); assert_eq!(matched.prompt_tag, "my-rule"); } + + // ── Fail-closed filter error handling (finding #25) ─────────────────────── + + #[tokio::test] + async fn test_filter_error_fails_closed_no_fallthrough() { + // A broken filter on rule[0] must NOT fall through to rule[1]. + let event = make_event(9, "hello"); + let channel_id = any_channel(); + + let rules = vec![ + make_rule( + "broken-filter", + ChannelScope::All("all".into()), + vec![], + false, + Some("this is not valid evalexpr syntax !!!"), + Some("should-not-match"), + ), + make_rule( + "catch-all", + ChannelScope::All("all".into()), + vec![], + false, + None, + Some("catch-all"), + ), + ]; + + // Must return None — not "catch-all". + let result = match_event(&event, channel_id, &rules, "").await; + assert!( + result.is_none(), + "filter error must fail closed, not fall through to next rule" + ); + } + + #[tokio::test] + async fn test_consecutive_timeouts_disables_rule() { + // After MAX_CONSECUTIVE_TIMEOUTS, the rule is skipped and None returned. + let event = make_event(9, "hello"); + let channel_id = any_channel(); + + let rule = make_rule( + "timed-out-rule", + ChannelScope::All("all".into()), + vec![], + false, + Some("kind == 9"), + Some("should-not-match"), + ); + // Pre-seed the counter at the threshold. + rule.consecutive_timeouts + .store(MAX_CONSECUTIVE_TIMEOUTS, Ordering::Relaxed); + + let rules = vec![rule]; + let result = match_event(&event, channel_id, &rules, "").await; + assert!(result.is_none(), "disabled rule must return None"); + } } diff --git a/crates/sprout-acp/src/main.rs b/crates/sprout-acp/src/main.rs index 4046cca097..56e35d34bc 100644 --- a/crates/sprout-acp/src/main.rs +++ b/crates/sprout-acp/src/main.rs @@ -27,7 +27,7 @@ use sprout_core::kind::{ KIND_MEMBER_ADDED_NOTIFICATION, KIND_MEMBER_REMOVED_NOTIFICATION, KIND_STREAM_MESSAGE, KIND_STREAM_REMINDER, KIND_WORKFLOW_APPROVAL_REQUESTED, }; -use tokio::sync::watch; +use tokio::sync::{mpsc, watch}; use tracing_subscriber::EnvFilter; use uuid::Uuid; @@ -48,8 +48,217 @@ fn is_subcommand(name: &str) -> bool { /// Timeout for the `sprout-acp models` subcommand (spawn + init + session/new). const MODELS_TIMEOUT: Duration = Duration::from_secs(10); +/// Maximum crashes in a 60-second window before a slot's circuit opens. +const CIRCUIT_BREAKER_THRESHOLD: usize = 3; +/// Window for circuit-breaker crash counting. +const CIRCUIT_BREAKER_WINDOW: Duration = Duration::from_secs(60); +/// Cooldown before a tripped circuit breaker allows a probe respawn. +const CIRCUIT_BREAKER_COOLDOWN: Duration = Duration::from_secs(300); // 5 minutes +/// Base backoff delay for respawn (doubles per recent crash, capped at 30s). +const RESPAWN_BASE_DELAY: Duration = Duration::from_secs(1); +/// Maximum respawn backoff delay. +const RESPAWN_MAX_DELAY: Duration = Duration::from_secs(30); + +/// Per-slot circuit breaker state. +/// +/// `crash_times` holds timestamps of recent crashes within `CIRCUIT_BREAKER_WINDOW`. +/// `open_until` is set when the threshold is hit; the circuit stays open until that +/// instant, then allows one probe respawn (half-open). If the probe crashes, the +/// circuit re-opens for another `CIRCUIT_BREAKER_COOLDOWN` period. +/// +/// All state transitions go through methods on this struct — callers never +/// manipulate `crash_times` or `open_until` directly. +struct SlotCircuit { + crash_times: Vec, + open_until: Option, + /// True while a background respawn/refill task is in flight for this slot. + /// Prevents duplicate spawns from maintenance ticks that fire before the + /// previous spawn_and_init completes. + respawn_in_flight: bool, +} + +/// Result of [`SlotCircuit::record_crash`]. +enum CrashVerdict { + /// Respawn is allowed after sleeping for this duration (jittered backoff). + Respawn(Duration), + /// Circuit is open — do not respawn. + CircuitOpen, + /// Circuit was open but cooldown has elapsed — one probe respawn is allowed + /// (no backoff sleep). If the probe crashes, the next `record_crash` will + /// immediately re-open the circuit. + HalfOpenProbe, +} + +impl SlotCircuit { + /// Record a crash and decide whether to respawn. + /// + /// This is the **single canonical path** for all crash → respawn decisions. + /// Called by `respawn_agent_into`, `recover_panicked_agent`, and slot refill. + fn record_crash(&mut self) -> CrashVerdict { + let now = std::time::Instant::now(); + + // Half-open: cooldown elapsed → allow one probe. + if let Some(open_until) = self.open_until { + if now >= open_until { + // Pre-seed crash_times to threshold-1 so that if the probe + // itself crashes on the *next* call, the threshold is hit + // immediately and the circuit re-opens. This implements a + // "prove stability for one full window" policy. + self.crash_times.clear(); + for _ in 0..(CIRCUIT_BREAKER_THRESHOLD - 1) { + self.crash_times.push(now); + } + self.open_until = None; + return CrashVerdict::HalfOpenProbe; + } else { + return CrashVerdict::CircuitOpen; + } + } + + // Record this crash and prune old entries. + self.crash_times.push(now); + self.crash_times + .retain(|&t| now.duration_since(t) < CIRCUIT_BREAKER_WINDOW); + + let recent = self.crash_times.len(); + + if recent >= CIRCUIT_BREAKER_THRESHOLD { + self.open_until = Some(now + CIRCUIT_BREAKER_COOLDOWN); + return CrashVerdict::CircuitOpen; + } + + // Exponential backoff: 1s * 2^(recent-1), capped at 30s, with ±20% jitter. + let base = RESPAWN_BASE_DELAY.saturating_mul(1u32 << (recent - 1).min(5)); + let capped = base.min(RESPAWN_MAX_DELAY); + let jitter = (std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .subsec_nanos() as f64) + / 1_000_000_000.0; // 0.0..1.0 + let factor = 0.8 + jitter * 0.4; // 0.8..1.2 + CrashVerdict::Respawn(capped.mul_f64(factor)) + } + + /// Mark a spawn failure — opens the circuit so the slot isn't retried + /// on every heartbeat tick. Uses fresh `Instant::now()` so spawn latency + /// doesn't shorten the effective cooldown. + fn mark_spawn_failed(&mut self) { + self.open_until = Some(std::time::Instant::now() + CIRCUIT_BREAKER_COOLDOWN); + } + + /// Check if an empty slot can be refilled. Unlike `record_crash`, this + /// does NOT record a new crash — it only checks whether the circuit + /// allows a respawn attempt. + /// + /// Returns `true` if respawn is allowed. For half-open probes, pre-seeds + /// crash_times so the next crash re-opens immediately. For normal refills + /// (no circuit was ever opened), crash history is preserved so the breaker + /// can still trip if the refilled agent crashes quickly. + fn can_refill(&mut self) -> bool { + let now = std::time::Instant::now(); + match self.open_until { + Some(open_until) => { + if now >= open_until { + // Half-open probe: pre-seed crash_times. + self.crash_times.clear(); + for _ in 0..(CIRCUIT_BREAKER_THRESHOLD - 1) { + self.crash_times.push(now); + } + self.open_until = None; + true + } else { + false // cooldown not elapsed + } + } + None => true, // no circuit open — normal refill, preserve crash history + } + } +} + +/// True if any slot has a respawn task in flight. Used to prevent premature +/// "all agents dead" exits — a respawning agent may succeed in seconds. +fn any_respawn_in_flight(crash_history: &[SlotCircuit]) -> bool { + crash_history.iter().any(|s| s.respawn_in_flight) +} + +/// Result of a background respawn task. +struct RespawnResult { + index: usize, + result: Result, +} + +/// RAII guard that ensures a `RespawnResult` is sent even if the task panics. +/// Without this, a panicked respawn task would leave `respawn_in_flight = true` +/// permanently, silently losing the slot forever. +struct RespawnGuard { + index: usize, + tx: mpsc::Sender, + sent: bool, +} + +impl RespawnGuard { + fn new(index: usize, tx: mpsc::Sender) -> Self { + Self { + index, + tx, + sent: false, + } + } + + /// Send the result and disarm the guard. Uses `try_send` (sync) so there + /// is no await boundary between marking `sent` and actually enqueueing — + /// cancellation cannot slip between the two. + fn send(mut self, result: Result) { + // Invariant: try_send succeeds because the channel capacity equals the + // slot count, and respawn_in_flight guarantees at most one outstanding + // result per slot. If this ever fails, the channel sizing or the + // respawn_in_flight guard has drifted — that's a bug, not a transient. + match self.tx.try_send(RespawnResult { + index: self.index, + result, + }) { + Ok(()) => self.sent = true, + Err(e) => { + tracing::error!( + agent = self.index, + "respawn result channel full or closed: {e}" + ); + // Drop will fire and send a failure result as fallback. + } + } + } +} + +impl Drop for RespawnGuard { + fn drop(&mut self) { + if !self.sent { + tracing::error!( + agent = self.index, + "respawn task exited without sending result — sending failure" + ); + // Best-effort: try_send in Drop (can't await). + let _ = self.tx.try_send(RespawnResult { + index: self.index, + result: Err(anyhow::anyhow!("respawn task panicked or was cancelled")), + }); + } + } +} + +// ── Finding #16: propagate_legacy_env_vars before tokio runtime ─────────────── +// +// Sync env-var propagation must run before the tokio runtime starts so that +// any child processes inherit the correct environment. This must happen in the +// sync entry point — `std::env::set_var` is only safe before tokio spawns +// worker threads (Rust 2024 edition safety requirement). + +fn main() -> Result<()> { + config::propagate_legacy_env_vars(); + tokio_main() +} + #[tokio::main] -async fn main() -> Result<()> { +async fn tokio_main() -> Result<()> { // Install the ring crypto provider for rustls (required for wss:// connections). rustls::crypto::ring::default_provider() .install_default() @@ -78,21 +287,79 @@ async fn main() -> Result<()> { tracing::info!("sprout-acp starting: {}", config.summary()); // ── Step 1: Spawn N ACP agent subprocesses and initialize ───────────────── - let mut agents = Vec::with_capacity(config.agents as usize); + // + // Finding #10: one agent failing to start must not kill the whole pool. + // We attempt each spawn under a 60-second timeout; failures are logged and + // skipped. If ALL agents fail we return an error. A partial pool is valid — + // the harness continues with reduced capacity and logs a warning. + let mut agent_slots: Vec> = Vec::with_capacity(config.agents as usize); for i in 0..config.agents as usize { - let acp = spawn_and_init(&config).await?; - agents.push(OwnedAgent { - index: i, - acp, - state: SessionState::default(), - model_capabilities: None, - desired_model: config.model.clone(), - }); + // Spawn OUTSIDE the timeout so we always own the child for cleanup. + // This matches the run_models pattern and prevents zombie leaks on + // init timeout (the cancelled future would drop the AcpClient via + // Drop which is best-effort only). + let spawn_result = AcpClient::spawn(&config.agent_command, &config.agent_args).await; + match spawn_result { + Ok(mut acp) => { + match tokio::time::timeout(Duration::from_secs(60), acp.initialize()).await { + Ok(Ok(init_result)) => { + tracing::info!(agent = i, "agent initialized: {init_result}"); + agent_slots.push(Some(OwnedAgent { + index: i, + acp, + state: SessionState::default(), + model_capabilities: None, + desired_model: config.model.clone(), + })); + } + Ok(Err(e)) => { + tracing::error!(agent = i, "agent initialize failed: {e}"); + acp.shutdown().await; + agent_slots.push(None); + } + Err(_) => { + tracing::error!(agent = i, "agent timed out during init (60s)"); + acp.shutdown().await; + agent_slots.push(None); + } + } + } + Err(e) => { + tracing::error!(agent = i, "agent failed to spawn: {e}"); + agent_slots.push(None); + } + } } - tracing::info!("agent_pool_ready agents={}", agents.len()); - let mut pool = AgentPool::new(agents); + let live_count = agent_slots.iter().filter(|s| s.is_some()).count(); + if live_count == 0 { + return Err(anyhow::anyhow!( + "all {} agents failed to start — cannot continue", + config.agents + )); + } + if live_count < config.agents as usize { + tracing::warn!( + "started {}/{} agents — continuing with reduced pool", + live_count, + config.agents + ); + } + tracing::info!("agent_pool_ready agents={}", live_count); + let mut pool = AgentPool::from_slots(agent_slots); // ── Step 2: Connect to Sprout relay ────────────────────────────────────── + // + // Finding #22: capture a startup watermark BEFORE connecting to the relay. + // This timestamp is used for membership notification replay (via + // startup_watermark) and as the initial subscribe_since for channels + // discovered at startup. The Subscribe handler falls back to + // subscribe_since when last_seen is None, closing the blind spot + // between "agents ready" and "first REQ sent". + let startup_watermark: u64 = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_secs(); + let pubkey_hex = config.keys.public_key().to_hex(); let mut relay = HarnessRelay::connect( &config.relay_url, @@ -103,6 +370,14 @@ async fn main() -> Result<()> { .await .map_err(|e| anyhow::anyhow!("relay connect error: {e}"))?; + // Finding #22: tell the relay background task the watermark so it can use + // `since = watermark - 5s` on the first REQ instead of `since=now`. + // Best-effort: a failure here is non-fatal (we just lose the startup window + // protection, which is the same as the pre-fix behaviour). + if let Err(e) = relay.set_startup_watermark(startup_watermark).await { + tracing::warn!("failed to set startup watermark: {e}"); + } + tracing::info!("connected to relay at {}", config.relay_url); // ── Step 2b: Subscribe to membership notifications ──────────────────────── @@ -173,6 +448,8 @@ async fn main() -> Result<()> { }), require_mention: !config.no_mention_filter, filter: None, + compiled_filter: None, + consecutive_timeouts: std::sync::Arc::new(std::sync::atomic::AtomicU32::new(0)), prompt_tag: Some("@mention".into()), }] } @@ -183,18 +460,14 @@ async fn main() -> Result<()> { kinds: config.kinds_override.clone().unwrap_or_default(), require_mention: false, filter: None, + compiled_filter: None, + consecutive_timeouts: std::sync::Arc::new(std::sync::atomic::AtomicU32::new(0)), prompt_tag: Some("all".into()), }] } SubscribeMode::Config => { - let loaded = config::load_rules(&config.config_path)?; - if loaded.is_empty() { - tracing::warn!( - "config file {} contains zero rules — agent will receive no events", - config.config_path.display() - ); - } - loaded + // load_rules() already warns if the config file has zero rules. + config::load_rules(&config.config_path)? } }; @@ -270,6 +543,19 @@ async fn main() -> Result<()> { let mut typing_channels: HashSet = HashSet::new(); let mut presence_task: Option> = None; + // ── Step 6d: Maintenance (slot refill + queue compaction) ──────────────── + // Runs at the TOP of every loop iteration via Instant check — cannot be + // starved by the biased select. Slot refill spawns background tasks so + // spawn_and_init never blocks the main loop. + let maintenance_interval = Duration::from_secs(30); + let mut last_maintenance = std::time::Instant::now(); + + // Channel for background respawn tasks to return completed agents. + // Bounded to agent count — at most one respawn per slot in flight. + let (respawn_tx, mut respawn_rx) = mpsc::channel::(config.agents as usize); + // JoinSet for respawn tasks so shutdown can abort them. + let mut respawn_tasks: tokio::task::JoinSet<()> = tokio::task::JoinSet::new(); + // ── Step 7: Shutdown signal ─────────────────────────────────────────────── let (shutdown_tx, mut shutdown_rx) = watch::channel(()); @@ -299,8 +585,10 @@ async fn main() -> Result<()> { // second are both processed. The seen_membership_ids set handles exact // replays that share the same timestamp. let mut membership_newest_ts: HashMap = HashMap::new(); - // Dedup set for exact membership event replays (bounded, cleared at 2000). - let mut seen_membership_ids: HashSet = HashSet::new(); + // Two-generation dedup for membership event replays (bounded, no amnesia). + // Rotates at 1000 entries instead of clearing the entire set at 2000. + let mut seen_membership_current: HashSet = HashSet::new(); + let mut seen_membership_previous: HashSet = HashSet::new(); // Channels the agent has been removed from. When a checked-out agent is // returned to the pool, its sessions for these channels are stripped, and @@ -320,6 +608,20 @@ async fn main() -> Result<()> { // and capture it in TaskMeta at dispatch time. let mut removed_channels: HashSet = HashSet::new(); + // ── Finding #14: Per-slot crash history for circuit breaker ─────────────── + // + // One SlotCircuit per agent slot. crash_times entries are pruned to the last + // CIRCUIT_BREAKER_WINDOW on each respawn attempt. The Vec is indexed by + // agent slot index, so it must be sized to the configured pool capacity + // (not the live count, which may be smaller after partial startup). + let mut crash_history: Vec = (0..config.agents as usize) + .map(|_| SlotCircuit { + crash_times: Vec::new(), + open_until: None, + respawn_in_flight: false, + }) + .collect(); + // ── Step 8: Main orchestration loop ────────────────────────────────────── // // Branches 1 & 2 both need to borrow `pool`, but they access different @@ -331,12 +633,86 @@ async fn main() -> Result<()> { } loop { + // ── Maintenance (runs at loop top — cannot be starved by biased select) ── + if last_maintenance.elapsed() >= maintenance_interval { + last_maintenance = std::time::Instant::now(); + queue.compact_expired_state(); + + // Slot refill: spawn background tasks for empty slots whose + // circuit breaker allows it. spawn_and_init runs off the main + // loop so it never blocks event processing. + for (idx, slot) in crash_history.iter_mut().enumerate() { + if pool.slot_alive(idx) || slot.respawn_in_flight { + continue; + } + if !slot.can_refill() { + continue; + } + slot.respawn_in_flight = true; + tracing::info!(agent = idx, "slot refill: spawning background respawn"); + let cmd = config.agent_command.clone(); + let args = config.agent_args.clone(); + let guard = RespawnGuard::new(idx, respawn_tx.clone()); + respawn_tasks.spawn(async move { + let result = spawn_and_init(&cmd, &args).await; + guard.send(result); + }); + } + + // Flush requeued batches whose retry_after has expired. Without + // this, a batch requeued during crash recovery can sit idle + // indefinitely on quiet channels — dispatch_pending is only + // called on relay events or pool results, neither of which + // arrive when the channel is silent. + if queue.has_flushable_work() { + typing_channels.extend(dispatch_pending(&mut pool, &mut queue, &ctx)); + } + } + + // ── Collect completed background respawns (non-blocking) ───────────── + let mut respawn_collected = false; + while let Ok(rr) = respawn_rx.try_recv() { + crash_history[rr.index].respawn_in_flight = false; + match rr.result { + Ok(acp) => { + let agent = OwnedAgent { + index: rr.index, + acp, + state: SessionState::default(), + model_capabilities: None, + desired_model: config.model.clone(), + }; + pool.return_agent(agent); + tracing::info!(agent = rr.index, "respawn complete"); + respawn_collected = true; + } + Err(e) => { + crash_history[rr.index].mark_spawn_failed(); + tracing::warn!(agent = rr.index, "respawn failed: {e} — circuit re-opened"); + } + } + } + // Flush requeued events that were waiting for a live agent. Without + // this, batches requeued during crash recovery sit idle until the + // next relay event arrives — which can be minutes on quiet channels. + if respawn_collected { + typing_channels.extend(dispatch_pending(&mut pool, &mut queue, &ctx)); + } + // Borrow result_rx and join_set simultaneously via split-borrow helper. let pool_event: Option = { let (result_rx, join_set) = pool.rx_and_join_set(); tokio::select! { biased; - r = result_rx.recv() => Some(PoolEvent::Result(Box::new(r.expect("result channel closed")))), + // Finding #24: recv() returning None means all senders dropped + // (pool was torn down). Break cleanly instead of panicking. + r = result_rx.recv() => match r { + Some(result) => Some(PoolEvent::Result(Box::new(result))), + None => { + tracing::info!("result channel closed — exiting main loop"); + break; + } + }, // Guard: join_next() returns None immediately when JoinSet is // empty, which would cause a tight spin. Only poll when there // are in-flight tasks. @@ -372,7 +748,10 @@ async fn main() -> Result<()> { // Why not `<=`? That would suppress legitimate live // add→remove (or remove→add) sequences in the same second, // leaving the harness in the wrong membership state. - if !seen_membership_ids.insert(eid.clone()) { + // Two-generation dedup: check both sets before inserting. + if seen_membership_current.contains(&eid) + || seen_membership_previous.contains(&eid) + { tracing::debug!( channel_id = %ch, kind = kind_u32, @@ -380,13 +759,11 @@ async fn main() -> Result<()> { ); continue; } - // Bound the dedup set to prevent unbounded growth. - // Re-insert the current ID after clearing so it stays - // protected against immediate replay (same pattern as - // relay.rs BgState::record_event). - if seen_membership_ids.len() > 2000 { - seen_membership_ids.clear(); - seen_membership_ids.insert(eid); + seen_membership_current.insert(eid); + // Rotate at 1000: current → previous, no amnesia window. + if seen_membership_current.len() >= 1000 { + seen_membership_previous = + std::mem::take(&mut seen_membership_current); } if let Some(&newest) = membership_newest_ts.get(&ch) { if ts < newest { @@ -592,10 +969,13 @@ async fn main() -> Result<()> { } } => { let _ = result_rx; + // Use try_publish (non-blocking) for typing indicators — + // they're ephemeral and must not block the main loop during + // relay reconnection (#35). for &ch in &typing_channels { if let Ok(event) = relay.build_typing_event(ch) { - if let Err(e) = relay.publish_event(event).await { - tracing::debug!("typing indicator failed for {ch}: {e}"); + if let Err(e) = relay.try_publish_event(event) { + tracing::debug!("typing indicator dropped for {ch}: {e}"); } } } @@ -621,9 +1001,10 @@ async fn main() -> Result<()> { *result, &mut heartbeat_in_flight, &removed_channels, - ) - .await - == LoopAction::Exit + &mut crash_history, + &respawn_tx, + &mut respawn_tasks, + ) == LoopAction::Exit { break; } @@ -634,9 +1015,10 @@ async fn main() -> Result<()> { &mut heartbeat_in_flight, &removed_channels, &mut typing_channels, - ) - .await - == LoopAction::Exit + &mut crash_history, + &respawn_tx, + &mut respawn_tasks, + ) == LoopAction::Exit { break; } @@ -652,9 +1034,11 @@ async fn main() -> Result<()> { &mut heartbeat_in_flight, &removed_channels, &mut typing_channels, - ) - .await; - if pool.live_count() == 0 { + &mut crash_history, + &respawn_tx, + &mut respawn_tasks, + ); + if pool.live_count() == 0 && !any_respawn_in_flight(&crash_history) { tracing::error!("all agents dead — exiting"); break; } @@ -669,10 +1053,30 @@ async fn main() -> Result<()> { // 30 s is generous for in-flight prompts to be cancelled; using // max_turn_duration here would cause Ctrl+C to hang for up to an hour. let grace = Duration::from_secs(30); + // Best-effort drain of both join_set and result_rx during the grace period. + // Tasks that finish normally send their OwnedAgent through result_rx — we + // explicitly shut them down here to reap child processes. If the grace + // period expires, remaining tasks are aborted and fall back to + // AcpClient::Drop (start_kill + try_wait — best-effort, not guaranteed). + let (rx_ref, js_ref) = pool.rx_and_join_set(); let shutdown_result = tokio::time::timeout(grace, async { - while let Some(result) = pool.join_set.join_next().await { - if let Err(e) = result { - tracing::warn!("task finished with error during shutdown: {e}"); + loop { + tokio::select! { + result = js_ref.join_next() => { + match result { + Some(Err(e)) => tracing::warn!("task error during shutdown: {e}"), + Some(Ok(())) => {} + None => break, // join_set empty + } + } + maybe_result = rx_ref.recv() => { + if let Some(mut pr) = maybe_result { + let idx = pr.agent.index; + pr.agent.acp.shutdown().await; + tracing::debug!(agent = idx, "reaped checked-out agent on shutdown"); + } + // If None, channel closed — tasks are done. + } } } }) @@ -681,8 +1085,39 @@ async fn main() -> Result<()> { tracing::warn!("grace period expired, aborting remaining tasks"); pool.join_set.shutdown().await; } + // Drain any remaining results that arrived after join_set drained but + // before tasks were aborted. + while let Ok(mut pr) = pool.result_rx_try_recv() { + let idx = pr.agent.index; + pr.agent.acp.shutdown().await; + tracing::debug!(agent = idx, "reaped late-arriving agent on shutdown"); + } + // Explicitly shut down idle agents still sitting in their slots. + for slot in pool.agents_mut().iter_mut() { + if let Some(agent) = slot.take() { + let idx = agent.index; + let mut acp = agent.acp; + acp.shutdown().await; + tracing::debug!(agent = idx, "reaped idle agent on shutdown"); + } + } drop(pool); + // Abort any in-flight respawn tasks. They may be sleeping in backoff or + // running spawn_and_init — either way, we don't want them spawning new + // children after the main loop has exited. RespawnGuard::Drop sends a + // failure result for aborted tasks, so respawn_in_flight is cleared. + respawn_tasks.shutdown().await; + + // Drain any respawn results that completed before the abort. Explicitly + // shut down returned agents instead of relying on AcpClient::Drop. + while let Ok(rr) = respawn_rx.try_recv() { + if let Ok(mut acp) = rr.result { + acp.shutdown().await; + tracing::debug!(agent = rr.index, "reaped respawned agent on shutdown"); + } + } + // Cancel any in-flight presence heartbeat before sending offline. if let Some(h) = presence_task.take() { h.abort(); @@ -703,6 +1138,10 @@ async fn main() -> Result<()> { } } + // Graceful relay shutdown — sends WebSocket close frame and waits up to 5s + // for the background task to finish, rather than aborting immediately (#40). + relay.shutdown().await; + tracing::info!("sprout-acp stopped"); Ok(()) } @@ -778,13 +1217,17 @@ fn dispatch_pending( // ── handle_prompt_result ────────────────────────────────────────────────────── -async fn handle_prompt_result( +#[allow(clippy::too_many_arguments)] +fn handle_prompt_result( pool: &mut AgentPool, queue: &mut EventQueue, config: &Config, mut result: PromptResult, heartbeat_in_flight: &mut bool, removed_channels: &HashSet, + crash_history: &mut [SlotCircuit], + respawn_tx: &mpsc::Sender, + respawn_tasks: &mut tokio::task::JoinSet<()>, ) -> LoopAction { let before = pool.task_map().len(); let agent_index = result.agent.index; @@ -792,11 +1235,11 @@ async fn handle_prompt_result( .retain(|_, meta| meta.agent_index != agent_index); debug_assert_eq!(before, pool.task_map().len() + 1); - match &result.source { - PromptSource::Channel(ch) => queue.mark_complete(*ch), - PromptSource::Heartbeat => *heartbeat_in_flight = false, - } - + // Requeue BEFORE mark_complete: requeue() sets retry_after with a future + // deadline, and mark_complete() checks for it to decide whether to preserve + // retry_counts. If mark_complete runs first, retry_counts is cleared and + // every retry starts at attempt 1 — defeating exponential backoff and + // dead-letter protection. if let Some(batch) = result.batch { // Don't requeue batches for channels the agent was removed from — // those events are stale and should be silently dropped. @@ -811,6 +1254,11 @@ async fn handle_prompt_result( } } + match &result.source { + PromptSource::Channel(ch) => queue.mark_complete(*ch), + PromptSource::Heartbeat => *heartbeat_in_flight = false, + } + // Strip sessions for channels the agent was removed from while this // agent was checked out. This covers the gap where invalidate_channel_sessions // only touches idle agents. @@ -827,6 +1275,15 @@ async fn handle_prompt_result( let agent_index = result.agent.index; match result.outcome { + // Successful prompt — return agent to pool. + PromptOutcome::Ok(_) => { + tracing::debug!( + agent = agent_index, + outcome = outcome_label, + "agent_returned" + ); + pool.return_agent(result.agent); + } // Fatal outcomes: the agent subprocess is dead or poisoned — respawn it. PromptOutcome::AgentExited | PromptOutcome::Timeout => { tracing::debug!( @@ -835,24 +1292,67 @@ async fn handle_prompt_result( "agent_returned — respawning" ); let index = result.agent.index; - match respawn_agent_into(result.agent, config).await { - Ok(agent) => pool.return_agent(agent), - Err(e) => { - tracing::error!("failed to respawn agent {index}: {e}"); - if pool.live_count() == 0 { - tracing::error!("all agents dead — exiting"); - return LoopAction::Exit; - } + let slot_history = &mut crash_history[index]; + if !spawn_respawn_task( + result.agent, + config, + slot_history, + respawn_tx, + respawn_tasks, + ) { + // Circuit open — slot stays empty until maintenance refill. + if pool.live_count() == 0 && !any_respawn_in_flight(crash_history) { + tracing::error!("all agents dead — exiting"); + return LoopAction::Exit; } } } - _ => { - tracing::debug!( - agent = agent_index, - outcome = outcome_label, - "agent_returned" + // Errors fall into two categories: + // + // 1. Transport-class (Io, WriteTimeout, Timeout, Protocol): the stdio + // pipe may be corrupted or the agent desynchronized. These are fatal + // to the agent regardless of whether they occurred during session + // creation or an active prompt — respawn unconditionally. + // + // 2. Application-class (IdleTimeout, HardTimeout, Json): the pipe is + // intact but the prompt failed. Return the agent to the pool so it + // can be reused for the next event. + PromptOutcome::Error(ref e) => { + let is_transport_error = matches!( + e, + acp::AcpError::Io(_) + | acp::AcpError::WriteTimeout(_) + | acp::AcpError::Timeout(_) + | acp::AcpError::Protocol(_) ); - pool.return_agent(result.agent); + if is_transport_error { + tracing::warn!( + agent = agent_index, + outcome = outcome_label, + "transport/protocol error — respawning agent" + ); + let index = result.agent.index; + let slot_history = &mut crash_history[index]; + if !spawn_respawn_task( + result.agent, + config, + slot_history, + respawn_tx, + respawn_tasks, + ) && pool.live_count() == 0 + && !any_respawn_in_flight(crash_history) + { + tracing::error!("all agents dead — exiting"); + return LoopAction::Exit; + } + } else { + tracing::debug!( + agent = agent_index, + outcome = outcome_label, + "agent_returned (application error — pipe intact)" + ); + pool.return_agent(result.agent); + } } } LoopAction::Continue @@ -860,7 +1360,8 @@ async fn handle_prompt_result( // ── recover_panicked_agent ──────────────────────────────────────────────────── -async fn recover_panicked_agent( +#[allow(clippy::too_many_arguments)] +fn recover_panicked_agent( pool: &mut AgentPool, queue: &mut EventQueue, config: &Config, @@ -868,6 +1369,9 @@ async fn recover_panicked_agent( heartbeat_in_flight: &mut bool, removed_channels: &HashSet, typing_channels: &mut HashSet, + crash_history: &mut [SlotCircuit], + respawn_tx: &mpsc::Sender, + respawn_tasks: &mut tokio::task::JoinSet<()>, ) { let task_id = join_error.id(); let Some(meta) = pool.task_map_mut().remove(&task_id) else { @@ -876,6 +1380,21 @@ async fn recover_panicked_agent( }; let i = meta.agent_index; + // Requeue BEFORE mark_complete (same rationale as handle_prompt_result). + if let Some(batch) = meta.recoverable_batch { + if let Some(ch) = meta.channel_id { + if !removed_channels.contains(&ch) { + queue.requeue(batch); + tracing::warn!("requeued batch for panicked agent {i}"); + } else { + tracing::debug!( + channel_id = %ch, + "dropping panicked batch for removed channel" + ); + } + } + } + if let Some(ch) = meta.channel_id { queue.mark_complete(ch); typing_channels.remove(&ch); @@ -885,45 +1404,57 @@ async fn recover_panicked_agent( tracing::warn!("cleared wedged heartbeat_in_flight from panicked agent {i}"); } - if let Some(batch) = meta.recoverable_batch { - // Don't requeue batches for removed channels. - if !removed_channels.contains(&batch.channel_id) { - queue.requeue(batch); - tracing::warn!("requeued batch for panicked agent {i}"); - } else { - tracing::debug!( - channel_id = %batch.channel_id, - "dropping panicked batch for removed channel" - ); - } - } + // Panics count as crashes for the circuit breaker. + // The panicked task already dropped the AcpClient, so we just need to + // check the circuit and spawn a fresh agent in the background. + let slot = &mut crash_history[i]; - match spawn_and_init(config).await { - Ok(acp) => { - pool.agents_mut()[i] = Some(OwnedAgent { - index: i, - acp, - state: SessionState::default(), - model_capabilities: None, - desired_model: config.model.clone(), - }); - tracing::info!("respawned agent {i} after panic"); + let delay = match slot.record_crash() { + CrashVerdict::CircuitOpen => { + tracing::error!(agent = i, "circuit open after panic — not respawning"); + return; } - Err(e) => { - tracing::error!("failed to respawn agent {i} after panic: {e}"); + CrashVerdict::HalfOpenProbe => { + tracing::info!(agent = i, "circuit half-open — probe respawn after panic"); + Duration::ZERO } - } + CrashVerdict::Respawn(d) => { + tracing::info!( + agent = i, + delay_ms = d.as_millis(), + "respawn backoff after panic" + ); + d + } + }; + + // Spawn respawn work off the main loop. + slot.respawn_in_flight = true; + let cmd = config.agent_command.clone(); + let args = config.agent_args.clone(); + let guard = RespawnGuard::new(i, respawn_tx.clone()); + respawn_tasks.spawn(async move { + if !delay.is_zero() { + tokio::time::sleep(delay).await; + } + let result = spawn_and_init(&cmd, &args).await; + guard.send(result); + }); } // ── drain_ready_join_results ────────────────────────────────────────────────── -async fn drain_ready_join_results( +#[allow(clippy::too_many_arguments)] +fn drain_ready_join_results( pool: &mut AgentPool, queue: &mut EventQueue, config: &Config, heartbeat_in_flight: &mut bool, removed_channels: &HashSet, typing_channels: &mut HashSet, + crash_history: &mut [SlotCircuit], + respawn_tx: &mpsc::Sender, + respawn_tasks: &mut tokio::task::JoinSet<()>, ) -> LoopAction { while let Some(Some(join_result)) = pool.join_set.join_next().now_or_never() { if let Err(join_error) = join_result { @@ -936,9 +1467,11 @@ async fn drain_ready_join_results( heartbeat_in_flight, removed_channels, typing_channels, - ) - .await; - if pool.live_count() == 0 { + crash_history, + respawn_tx, + respawn_tasks, + ); + if pool.live_count() == 0 && !any_respawn_in_flight(crash_history) { return LoopAction::Exit; } } @@ -1007,33 +1540,85 @@ fn default_heartbeat_prompt() -> String { // ── respawn_agent_into ──────────────────────────────────────────────────────── -async fn respawn_agent_into(old_agent: OwnedAgent, config: &Config) -> Result { +/// Spawn a background respawn task for a crashed agent slot. +/// +/// Does the circuit breaker check synchronously (non-blocking), then spawns +/// the actual shutdown + backoff + spawn_and_init work into a background task. +/// The result comes back through `respawn_tx` so the main loop stays responsive. +/// +/// Returns `true` if a respawn task was spawned, `false` if the circuit is open. +fn spawn_respawn_task( + old_agent: OwnedAgent, + config: &Config, + slot: &mut SlotCircuit, + respawn_tx: &mpsc::Sender, + respawn_tasks: &mut tokio::task::JoinSet<()>, +) -> bool { let index = old_agent.index; - drop(old_agent); // kill the old process via AcpClient Drop - let acp = spawn_and_init(config).await?; - Ok(OwnedAgent { - index, - acp, - state: SessionState::default(), - model_capabilities: None, - desired_model: config.model.clone(), - }) + + // Circuit breaker: record crash, decide whether to respawn. + let delay = match slot.record_crash() { + CrashVerdict::CircuitOpen => { + tracing::error!(agent = index, "circuit open — not respawning"); + return false; + } + CrashVerdict::HalfOpenProbe => { + tracing::info!(agent = index, "circuit half-open — probe respawn"); + Duration::ZERO + } + CrashVerdict::Respawn(d) => { + tracing::info!(agent = index, delay_ms = d.as_millis(), "respawn backoff"); + d + } + }; + + slot.respawn_in_flight = true; + + // Spawn the actual work (shutdown + sleep + spawn + init) off the main loop. + let cmd = config.agent_command.clone(); + let args = config.agent_args.clone(); + let guard = RespawnGuard::new(index, respawn_tx.clone()); + respawn_tasks.spawn(async move { + // Shutdown old agent (reap child, prevent zombie). + let mut agent = old_agent; + agent.acp.shutdown().await; + drop(agent); + + if !delay.is_zero() { + tokio::time::sleep(delay).await; + } + + let result = spawn_and_init(&cmd, &args).await; + guard.send(result); + }); + + true } // ── spawn_and_init ──────────────────────────────────────────────────────────── -async fn spawn_and_init(config: &Config) -> Result { - let mut acp = AcpClient::spawn(&config.agent_command, &config.agent_args) +/// Spawn an agent subprocess and run the MCP `initialize` handshake. +/// +/// Takes owned args so it can run in a background `tokio::spawn` task without +/// borrowing `Config`. All respawn/refill paths use this. +async fn spawn_and_init(command: &str, args: &[String]) -> Result { + let mut acp = AcpClient::spawn(command, args) .await .map_err(|e| anyhow::anyhow!("failed to spawn agent: {e}"))?; - let init_result = acp - .initialize() - .await - .map_err(|e| anyhow::anyhow!("agent initialize failed: {e}"))?; - - tracing::info!("agent initialized: {init_result}"); - Ok(acp) + match acp.initialize().await { + Ok(init_result) => { + tracing::info!("agent initialized: {init_result}"); + Ok(acp) + } + Err(e) => { + // Explicitly shut down the spawned child to prevent zombie/leak. + // Drop only does start_kill + try_wait (best-effort); shutdown() + // does start_kill + bounded wait (guaranteed reap). + acp.shutdown().await; + Err(anyhow::anyhow!("agent initialize failed: {e}")) + } + } } // ── build_mcp_servers ───────────────────────────────────────────────────────── @@ -1191,6 +1776,9 @@ fn build_mcp_servers(config: &Config) -> Vec { }, EnvVar { name: "SPROUT_PRIVATE_KEY".into(), + // bech32 encoding of a valid secret key is infallible. + // Panic here is correct: injecting a bogus secret would cause + // delayed, hard-to-diagnose agent failures downstream. value: config .keys .secret_key() diff --git a/crates/sprout-acp/src/pool.rs b/crates/sprout-acp/src/pool.rs index 4755d0fb04..be45dc4499 100644 --- a/crates/sprout-acp/src/pool.rs +++ b/crates/sprout-acp/src/pool.rs @@ -191,6 +191,11 @@ impl AgentPool { /// /// Agents are placed into indexed slots. The unbounded channel is created /// here; tasks send results back through `result_tx`. + /// + /// Prefer [`AgentPool::from_slots`] for startup paths where some agents may + /// have failed — `new()` packs agents densely and will break the + /// `agent.index` invariant if any slot was skipped. + #[allow(dead_code)] pub fn new(agents: Vec) -> Self { let (result_tx, result_rx) = mpsc::unbounded_channel(); let slots = agents.into_iter().map(Some).collect(); @@ -203,6 +208,23 @@ impl AgentPool { } } + /// Create a pool from pre-indexed slots (may contain None for failed startups). + /// + /// Slot positions are preserved so that `agent.index` always matches the + /// index into `self.agents`. Use this instead of `new()` when the startup + /// loop skips failed agents — `new()` would pack agents densely and break + /// the index invariant. + pub fn from_slots(slots: Vec>) -> Self { + let (result_tx, result_rx) = mpsc::unbounded_channel(); + Self { + agents: slots, + result_tx, + result_rx, + join_set: JoinSet::new(), + task_map: HashMap::new(), + } + } + /// Try to claim an idle agent for the given channel (or heartbeat if `None`). /// /// Pass 1: prefer an agent that already has a session for `channel_id`. @@ -230,10 +252,16 @@ impl AgentPool { /// Return an agent to its slot after a task completes. pub fn return_agent(&mut self, agent: OwnedAgent) { let idx = agent.index; - debug_assert!( - self.agents[idx].is_none(), - "return_agent: slot {idx} already occupied" - ); + if self.agents[idx].is_some() { + // This is a bug: two tasks returned the same agent index. Log it + // loudly so it shows up in production logs, then overwrite — the + // alternative (dropping the incoming agent) would permanently leak + // the slot. + tracing::error!( + idx, + "BUG: return_agent called for slot {idx} which is already occupied — overwriting" + ); + } self.agents[idx] = Some(agent); } @@ -284,6 +312,24 @@ impl AgentPool { (&mut self.result_rx, &mut self.join_set) } + /// Non-blocking drain of the result channel. Used during shutdown to + /// collect agents that completed while join_set was being drained. + pub fn result_rx_try_recv(&mut self) -> Result { + self.result_rx.try_recv() + } + + /// Check whether a slot is alive: either idle in the pool or checked out + /// for an in-flight task. Returns `false` only when the slot is truly + /// empty and available for refill. + pub fn slot_alive(&self, index: usize) -> bool { + let idle = self.agents.get(index).is_some_and(|s| s.is_some()); + if idle { + return true; + } + // Check if the agent is checked out (in-flight on a task). + self.task_map.values().any(|m| m.agent_index == index) + } + pub fn agents_mut(&mut self) -> &mut Vec> { &mut self.agents } @@ -311,8 +357,13 @@ impl AgentPool { // ── run_prompt_task ─────────────────────────────────────────────────────────── -/// Timeout for pre-prompt context fetches (thread/DM history). -const CONTEXT_FETCH_TIMEOUT: Duration = Duration::from_millis(500); +/// Timeout for a single pre-prompt context fetch attempt (thread/DM history). +/// Each call gets this budget; with one retry the total worst-case is +/// 2 × CONTEXT_FETCH_TIMEOUT + CONTEXT_FETCH_RETRY_DELAY ≈ 6.5 s. +const CONTEXT_FETCH_TIMEOUT: Duration = Duration::from_millis(3_000); + +/// Delay between the first failed context fetch and the single retry. +const CONTEXT_FETCH_RETRY_DELAY: Duration = Duration::from_millis(500); /// Timeout for model-switch requests (`session/set_config_option`, `session/set_model`). const MODEL_SWITCH_TIMEOUT: Duration = Duration::from_secs(5); @@ -347,7 +398,7 @@ async fn create_session_and_apply_model( if let Some(ref desired) = agent.desired_model { match resolve_model_switch_method(&resp.raw, desired) { Some(method) => { - apply_model_switch(&mut agent.acp, &resp.session_id, desired, &method).await; + apply_model_switch(&mut agent.acp, &resp.session_id, desired, &method).await?; } None => { tracing::warn!( @@ -382,7 +433,7 @@ async fn apply_model_switch( session_id: &str, desired: &str, method: &ModelSwitchMethod, -) { +) -> Result<(), AcpError> { let method_label = match method { ModelSwitchMethod::ConfigOption { config_id, .. } => { format!("configOption (configId={config_id})") @@ -413,6 +464,20 @@ async fn apply_model_switch( "applied model {desired} via {method_label} on session {session_id}" ); } + // Transport-class errors may have corrupted the stdio stream — propagate + // so the caller can respawn the agent instead of reusing a poisoned one. + Ok(Err(e @ AcpError::Io(_))) + | Ok(Err(e @ AcpError::WriteTimeout(_))) + | Ok(Err(e @ AcpError::Timeout(_))) + | Ok(Err(e @ AcpError::Protocol(_))) + | Ok(Err(e @ AcpError::AgentExited)) => { + tracing::error!( + target: "pool::model", + "fatal error setting model {desired} via {method_label}: {e}" + ); + return Err(e); + } + // Application-level errors (Json, etc.) — agent is fine, just uses default model. Ok(Err(e)) => { tracing::warn!( target: "pool::model", @@ -420,12 +485,16 @@ async fn apply_model_switch( ); } Err(_) => { - tracing::warn!( + // Outer timeout fired — the inner send_request may have left the + // stream in an unknown state. Treat as transport error. + tracing::error!( target: "pool::model", - "model set via {method_label} timed out ({MODEL_SWITCH_TIMEOUT:?}) — proceeding with agent default" + "model set via {method_label} timed out ({MODEL_SWITCH_TIMEOUT:?}) — treating as fatal" ); + return Err(AcpError::Timeout(MODEL_SWITCH_TIMEOUT)); } } + Ok(()) } /// Set the session permission mode via `session/set_config_option`. @@ -471,16 +540,20 @@ async fn apply_permission_mode( "applied permission mode {wire:?} on session {session_id}" ); } - Ok(Err(AcpError::AgentExited)) => { - // Fatal: the agent process crashed (e.g., goose doesn't support - // session/set_config_option and exits instead of returning an error). - // Propagate so the caller can respawn. + // Transport-class errors may have corrupted the stdio stream — propagate + // so the caller can respawn the agent. + Ok(Err(e @ AcpError::Io(_))) + | Ok(Err(e @ AcpError::WriteTimeout(_))) + | Ok(Err(e @ AcpError::Timeout(_))) + | Ok(Err(e @ AcpError::Protocol(_))) + | Ok(Err(e @ AcpError::AgentExited)) => { tracing::error!( target: "pool::permission", - "agent exited while setting permission mode {wire:?} — process crashed" + "fatal error setting permission mode {wire:?}: {e}" ); - return Err(AcpError::AgentExited); + return Err(e); } + // Application-level errors — agent is fine, just uses default permission mode. Ok(Err(e)) => { tracing::warn!( target: "pool::permission", @@ -488,10 +561,12 @@ async fn apply_permission_mode( ); } Err(_) => { - tracing::warn!( + // Outer timeout fired — stream may be in unknown state. + tracing::error!( target: "pool::permission", - "permission mode set timed out ({PERMISSION_MODE_TIMEOUT:?}) — falling back to per-tool auto-approval" + "permission mode set timed out ({PERMISSION_MODE_TIMEOUT:?}) — treating as fatal" ); + return Err(AcpError::Timeout(PERMISSION_MODE_TIMEOUT)); } } Ok(()) @@ -923,44 +998,63 @@ pub async fn run_prompt_task( // ── Context fetching ────────────────────────────────────────────────────────── +/// Retry wrapper for context fetches: one retry with `CONTEXT_FETCH_RETRY_DELAY` +/// on any `None` result. The closure is called twice at most. +/// +/// Using a closure (not a `Future`) so the retry can construct a fresh `Future` +/// each attempt without requiring `Clone` or re-boxing. +async fn fetch_with_retry(f: F) -> Option +where + F: Fn() -> Fut, + Fut: std::future::Future>, +{ + if let Some(result) = f().await { + return Some(result); + } + tokio::time::sleep(CONTEXT_FETCH_RETRY_DELAY).await; + f().await +} + /// Lazy-fetch channel metadata for a channel not in the startup discovery cache. /// /// Handles channels added dynamically via membership notifications after startup. -/// Uses the same 500ms timeout as context fetches. Returns `None` on failure -/// (graceful degradation — prompt will lack channel name and DM detection). +/// Uses `CONTEXT_FETCH_TIMEOUT` with one retry on failure. Returns `None` on +/// persistent failure (graceful degradation — prompt will lack channel name and +/// DM detection). async fn fetch_channel_info(channel_id: Uuid, rest: &RestClient) -> Option { let path = format!("/api/channels/{}", channel_id); - let result = timeout(CONTEXT_FETCH_TIMEOUT, rest.get_json(&path)).await; - - match result { - Ok(Ok(json)) => { - let name = json - .get("name") - .and_then(|v| v.as_str()) - .unwrap_or("unknown") - .to_string(); - let channel_type = json - .get("channel_type") - .and_then(|v| v.as_str()) - .unwrap_or("stream") - .to_string(); - Some(PromptChannelInfo { name, channel_type }) + fetch_with_retry(|| async { + match timeout(CONTEXT_FETCH_TIMEOUT, rest.get_json(&path)).await { + Ok(Ok(json)) => { + let name = json + .get("name") + .and_then(|v| v.as_str()) + .unwrap_or("unknown") + .to_string(); + let channel_type = json + .get("channel_type") + .and_then(|v| v.as_str()) + .unwrap_or("stream") + .to_string(); + Some(PromptChannelInfo { name, channel_type }) + } + Ok(Err(e)) => { + tracing::debug!( + channel_id = %channel_id, + "channel info fetch failed: {e} — will retry" + ); + None + } + Err(_) => { + tracing::debug!( + channel_id = %channel_id, + "channel info fetch timed out — will retry" + ); + None + } } - Ok(Err(e)) => { - tracing::debug!( - channel_id = %channel_id, - "channel info fetch failed: {e} — using defaults" - ); - None - } - Err(_) => { - tracing::debug!( - channel_id = %channel_id, - "channel info fetch timed out — using defaults" - ); - None - } - } + }) + .await } /// Fetch conversation context (thread or DM) for a batch before prompting. @@ -1085,23 +1179,25 @@ async fn fetch_prompt_profile_lookup( } let body = serde_json::json!({ "pubkeys": pubkeys }); - let result = timeout( - CONTEXT_FETCH_TIMEOUT, - rest.post_json("/api/users/batch", &body), - ) - .await; - - match result { - Ok(Ok(json)) => parse_profile_lookup_response(json), - Ok(Err(e)) => { - tracing::debug!("prompt profile lookup failed: {e} — using raw pubkeys"); - None + fetch_with_retry(|| async { + match timeout( + CONTEXT_FETCH_TIMEOUT, + rest.post_json("/api/users/batch", &body), + ) + .await + { + Ok(Ok(json)) => parse_profile_lookup_response(json), + Ok(Err(e)) => { + tracing::debug!("prompt profile lookup failed: {e} — will retry"); + None + } + Err(_) => { + tracing::debug!("prompt profile lookup timed out — will retry"); + None + } } - Err(_) => { - tracing::debug!("prompt profile lookup timed out — using raw pubkeys"); - None - } - } + }) + .await } /// Fetch thread context via REST: `GET /api/channels/{id}/threads/{event_id}?limit=N` @@ -1129,27 +1225,28 @@ async fn fetch_thread_context( channel_id, root_event_id, limit ); - let result = timeout(CONTEXT_FETCH_TIMEOUT, rest.get_json(&path)).await; - - match result { - Ok(Ok(json)) => parse_thread_response(json), - Ok(Err(e)) => { - tracing::warn!( - channel_id = %channel_id, - root = root_event_id, - "thread context fetch failed: {e} — falling back to hints-only" - ); - None + fetch_with_retry(|| async { + match timeout(CONTEXT_FETCH_TIMEOUT, rest.get_json(&path)).await { + Ok(Ok(json)) => parse_thread_response(json), + Ok(Err(e)) => { + tracing::warn!( + channel_id = %channel_id, + root = root_event_id, + "thread context fetch failed: {e} — will retry" + ); + None + } + Err(_) => { + tracing::warn!( + channel_id = %channel_id, + root = root_event_id, + "thread context fetch timed out — will retry" + ); + None + } } - Err(_) => { - tracing::warn!( - channel_id = %channel_id, - root = root_event_id, - "thread context fetch timed out — falling back to hints-only" - ); - None - } - } + }) + .await } /// Fetch DM context via REST: `GET /api/channels/{id}/messages?limit=N` @@ -1160,25 +1257,26 @@ async fn fetch_dm_context( ) -> Option { let path = format!("/api/channels/{}/messages?limit={}", channel_id, limit); - let result = timeout(CONTEXT_FETCH_TIMEOUT, rest.get_json(&path)).await; - - match result { - Ok(Ok(json)) => parse_dm_response(json, limit), - Ok(Err(e)) => { - tracing::warn!( - channel_id = %channel_id, - "DM context fetch failed: {e} — falling back to hints-only" - ); - None + fetch_with_retry(|| async { + match timeout(CONTEXT_FETCH_TIMEOUT, rest.get_json(&path)).await { + Ok(Ok(json)) => parse_dm_response(json, limit), + Ok(Err(e)) => { + tracing::warn!( + channel_id = %channel_id, + "DM context fetch failed: {e} — will retry" + ); + None + } + Err(_) => { + tracing::warn!( + channel_id = %channel_id, + "DM context fetch timed out — will retry" + ); + None + } } - Err(_) => { - tracing::warn!( - channel_id = %channel_id, - "DM context fetch timed out — falling back to hints-only" - ); - None - } - } + }) + .await } /// Parse the thread REST response into a `ConversationContext::Thread`. @@ -1367,13 +1465,18 @@ impl ReactionGuard { impl Drop for ReactionGuard { fn drop(&mut self) { - // Safety: always called from within a tokio task (run_prompt_task is - // spawned via JoinSet::spawn), so a runtime context is guaranteed. - // During shutdown the spawned cleanup may be dropped before completing - // — acceptable for a cosmetic indicator. + // Guard against drop outside a tokio runtime (e.g., in unit tests or + // during process teardown before the runtime is fully initialized). + // `run_prompt_task` is always spawned via `JoinSet::spawn`, so a + // runtime handle is normally available; `try_current` is the safe + // fallback for the rare cases it isn't. if let Some(rest) = self.rest.take() { let ids = std::mem::take(&mut self.ids); - tokio::spawn(clear_reactions(rest, ids)); + if let Ok(handle) = tokio::runtime::Handle::try_current() { + handle.spawn(clear_reactions(rest, ids)); + } + // If no runtime is available, reactions are left as-is — they are + // cosmetic indicators and the stale state is harmless. } } } @@ -1537,28 +1640,38 @@ pub(crate) async fn reaction_remove(rest: &crate::relay::RestClient, event_id: & } } -/// Add 💬 to all events concurrently. Awaited inline before the prompt fires. -/// Bounded by `REACTION_TIMEOUT` per call — worst case is a single 500ms wait -/// regardless of batch size. +/// Maximum concurrent reaction HTTP requests per fan-out call. +/// Prevents unbounded parallelism when a large batch of events arrives. +const REACTION_CONCURRENCY: usize = 10; + +/// Add 💬 to all events, capped at `REACTION_CONCURRENCY` concurrent requests. +/// Awaited inline before the prompt fires. async fn react_working(rest: &crate::relay::RestClient, event_ids: &[String]) { - futures_util::future::join_all( - event_ids - .iter() - .map(|eid| reaction_add(rest, eid, REACTION_WORKING)), - ) - .await; + for chunk in event_ids.chunks(REACTION_CONCURRENCY) { + futures_util::future::join_all( + chunk + .iter() + .map(|eid| reaction_add(rest, eid, REACTION_WORKING)), + ) + .await; + } } /// Fire-and-forget: remove both 👀 and 💬 from all events. Spawned on turn complete. -/// All removals run concurrently — bounded by `REACTION_TIMEOUT` per call. +/// Capped at `REACTION_CONCURRENCY` concurrent requests per chunk to avoid +/// unbounded HTTP fan-out on large batches. async fn clear_reactions(rest: crate::relay::RestClient, event_ids: Vec) { - futures_util::future::join_all(event_ids.iter().flat_map(|eid| { - [ - reaction_remove(&rest, eid, REACTION_SEEN), - reaction_remove(&rest, eid, REACTION_WORKING), - ] - })) - .await; + // Each event needs two removals (👀 and 💬); pair them and chunk by + // REACTION_CONCURRENCY pairs so the total concurrent requests stay bounded. + for chunk in event_ids.chunks(REACTION_CONCURRENCY) { + futures_util::future::join_all(chunk.iter().flat_map(|eid| { + [ + reaction_remove(&rest, eid, REACTION_SEEN), + reaction_remove(&rest, eid, REACTION_WORKING), + ] + })) + .await; + } } // ─── Unit Tests ────────────────────────────────────────────────────────────── diff --git a/crates/sprout-acp/src/queue.rs b/crates/sprout-acp/src/queue.rs index 6f227315f9..0805278339 100644 --- a/crates/sprout-acp/src/queue.rs +++ b/crates/sprout-acp/src/queue.rs @@ -20,6 +20,26 @@ use uuid::Uuid; use crate::config::DedupMode; +// ── Reliability constants ───────────────────────────────────────────────────── + +/// Maximum events queued per channel before oldest events are dropped. +const MAX_PENDING_PER_CHANNEL: usize = 500; + +/// Maximum events drained into a single batch. +const MAX_BATCH_EVENTS: usize = 50; + +/// Maximum retry attempts before a batch is dead-lettered. +const MAX_RETRIES: u32 = 10; + +/// Base retry delay in seconds (doubled each attempt). +const BASE_RETRY_DELAY_SECS: u64 = 5; + +/// Cap on retry delay in seconds. +const MAX_RETRY_DELAY_SECS: u64 = 300; + +/// In-flight deadline: max_turn (3600s) + 100s buffer. +const IN_FLIGHT_DEADLINE_SECS: u64 = 3700; + // ── Types ───────────────────────────────────────────────────────────────────── /// An event waiting in the queue. @@ -55,36 +75,53 @@ pub struct FlushBatch { /// /// ```text /// State: -/// queues: Map> -/// in_flight_channels: HashSet -/// retry_after: Map -/// dedup_mode: DedupMode +/// queues: Map> (capped at MAX_PENDING_PER_CHANNEL) +/// in_flight_channels: HashSet +/// in_flight_deadlines: Map (auto-expire after IN_FLIGHT_DEADLINE_SECS) +/// retry_after: Map +/// retry_counts: Map (dead-letter after MAX_RETRIES) +/// dedup_mode: DedupMode /// /// Transitions: /// push(event): /// if dedup_mode == Drop AND in_flight_channels.contains(event.channel_id): /// debug log + discard +/// else if queues[channel].len() >= MAX_PENDING_PER_CHANNEL: +/// drop oldest (pop_front), warn, push_back new event /// else: /// queues[event.channel_id].push_back(event) /// /// flush_next() → Option: +/// expire any stuck in-flight entries past their deadline /// candidates = channels where queue non-empty /// AND NOT in in_flight_channels /// AND (no retry_after OR retry_after[c] <= now) /// if candidates empty: return None /// channel = pick candidate with oldest head event (min received_at) -/// events = drain queues[channel] +/// events = drain up to MAX_BATCH_EVENTS from queues[channel] /// in_flight_channels.insert(channel) +/// in_flight_deadlines.insert(channel, now + IN_FLIGHT_DEADLINE_SECS) /// return Some(FlushBatch { channel, events }) /// /// mark_complete(channel_id): /// in_flight_channels.remove(channel_id) -/// (retry_after entries expire naturally via Instant check) +/// in_flight_deadlines.remove(channel_id) +/// retry_counts.remove(channel_id) +/// clean up expired retry_after entry if present +/// +/// requeue(batch): +/// increment retry_counts[channel] +/// if retry_counts[channel] > MAX_RETRIES: dead-letter (log ERROR, discard) +/// else: push_front with original received_at, set exponential backoff retry_after with jitter /// ``` pub struct EventQueue { queues: HashMap>, in_flight_channels: HashSet, + /// Per-channel deadline for auto-expiring stuck in-flight entries. + in_flight_deadlines: HashMap, retry_after: HashMap, + /// Per-channel retry attempt counter for exponential backoff / dead-lettering. + retry_counts: HashMap, dedup_mode: DedupMode, } @@ -94,7 +131,9 @@ impl EventQueue { Self { queues: HashMap::new(), in_flight_channels: HashSet::new(), + in_flight_deadlines: HashMap::new(), retry_after: HashMap::new(), + retry_counts: HashMap::new(), dedup_mode, } } @@ -115,10 +154,17 @@ impl EventQueue { ); return false; } - self.queues - .entry(event.channel_id) - .or_default() - .push_back(event); + let queue = self.queues.entry(event.channel_id).or_default(); + // Enforce per-channel depth cap: drop oldest to make room. + if queue.len() >= MAX_PENDING_PER_CHANNEL { + queue.pop_front(); + tracing::warn!( + channel_id = %event.channel_id, + limit = MAX_PENDING_PER_CHANNEL, + "queue depth cap reached — dropped oldest event" + ); + } + queue.push_back(event); true } @@ -131,6 +177,22 @@ impl EventQueue { pub fn flush_next(&mut self) -> Option { let now = Instant::now(); + // Auto-expire any stuck in-flight entries that missed mark_complete. + let expired: Vec = self + .in_flight_deadlines + .iter() + .filter(|(_, deadline)| now >= **deadline) + .map(|(id, _)| *id) + .collect(); + for id in expired { + tracing::error!( + channel_id = %id, + "BUG: in-flight channel expired without mark_complete — auto-releasing" + ); + self.in_flight_channels.remove(&id); + self.in_flight_deadlines.remove(&id); + } + // Find the channel whose head event has the oldest received_at, // excluding in-flight channels and throttled channels. let channel_id = self @@ -144,10 +206,11 @@ impl EventQueue { .min_by_key(|(_, q)| q.front().unwrap().received_at) .map(|(id, _)| *id)?; - // Drain ALL events for that channel. - let queue = self.queues.remove(&channel_id)?; + // Drain up to MAX_BATCH_EVENTS; leave any remainder in the queue. + let queue = self.queues.entry(channel_id).or_default(); + let drain_count = MAX_BATCH_EVENTS.min(queue.len()); let events: Vec = queue - .into_iter() + .drain(..drain_count) .map(|qe| BatchEvent { event: qe.event, prompt_tag: qe.prompt_tag, @@ -155,17 +218,47 @@ impl EventQueue { }) .collect(); + // Remove the queue entry if now empty (keeps pending_channels() accurate). + if self.queues.get(&channel_id).is_some_and(|q| q.is_empty()) { + self.queues.remove(&channel_id); + } + self.in_flight_channels.insert(channel_id); + self.in_flight_deadlines.insert( + channel_id, + now + Duration::from_secs(IN_FLIGHT_DEADLINE_SECS), + ); Some(FlushBatch { channel_id, events }) } /// Mark the prompt for `channel_id` as complete. /// - /// Removes the channel from `in_flight_channels`. Does NOT clear - /// `retry_after` — those entries expire naturally via Instant check. + /// Removes the channel from `in_flight_channels` and `in_flight_deadlines`. + /// + /// If the channel was NOT requeued (no active `retry_after` throttle), the + /// retry counter is reset — the channel is healthy and the next failure + /// starts fresh. If the channel WAS requeued, `retry_counts` is left intact + /// so the backoff sequence continues on the next attempt. + /// + /// Also cleans up any already-expired `retry_after` entry. pub fn mark_complete(&mut self, channel_id: Uuid) { self.in_flight_channels.remove(&channel_id); + self.in_flight_deadlines.remove(&channel_id); + let now = Instant::now(); + match self.retry_after.get(&channel_id) { + // Active throttle → channel was requeued; keep retry_counts intact. + Some(&deadline) if deadline > now => {} + // Expired or absent throttle → successful completion; reset counter + // and clean up the stale retry_after entry. + Some(_) => { + self.retry_after.remove(&channel_id); + self.retry_counts.remove(&channel_id); + } + None => { + self.retry_counts.remove(&channel_id); + } + } } /// Re-queue a batch of events that failed to process. @@ -174,25 +267,84 @@ impl EventQueue { /// are processed first on the next flush cycle. This prevents event loss /// when session creation or `session/prompt` fails transiently. /// - /// `received_at` is reset to `Instant::now()` for re-queued events. - /// A 5-second `retry_after` throttle is set so the channel is not - /// immediately re-flushed. + /// Original `received_at` timestamps are preserved so the channel retains + /// its fairness position. The retry delay comes from exponential backoff, + /// not from resetting received_at. + /// + /// After [`MAX_RETRIES`] attempts the batch is dead-lettered: logged at + /// ERROR and discarded rather than requeued. This prevents poison batches + /// from looping forever. /// /// Note: does NOT remove from `in_flight_channels` — caller must call /// `mark_complete` separately. pub fn requeue(&mut self, batch: FlushBatch) { - let queue = self.queues.entry(batch.channel_id).or_default(); + let channel_id = batch.channel_id; + let attempt = { + let count = self.retry_counts.entry(channel_id).or_insert(0); + *count += 1; + *count + }; + + if attempt > MAX_RETRIES { + tracing::error!( + channel_id = %channel_id, + attempt, + events = batch.events.len(), + "dead-lettering batch after {} retries — discarding {} events", + MAX_RETRIES, + batch.events.len(), + ); + self.retry_counts.remove(&channel_id); + // Also clear retry_after so fresh traffic on this channel isn't + // throttled by stale backoff from the discarded poison batch. + self.retry_after.remove(&channel_id); + return; + } + + // Exponential backoff: BASE * 2^(attempt-1), capped at MAX, with ±20% jitter. + let base_secs = BASE_RETRY_DELAY_SECS.saturating_mul(1u64 << (attempt - 1).min(6)); + let capped_secs = base_secs.min(MAX_RETRY_DELAY_SECS); + // Jitter: multiply by 0.8..1.2 using subsecond nanos as entropy source. + let jitter = { + let nanos = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .subsec_nanos(); + 0.8 + (nanos as f64 / u32::MAX as f64) * 0.4 + }; + let delay = Duration::from_secs_f64(capped_secs as f64 * jitter); + + tracing::warn!( + channel_id = %channel_id, + attempt, + max = MAX_RETRIES, + delay_secs = delay.as_secs_f64(), + events = batch.events.len(), + "requeueing failed batch with backoff" + ); + + let queue = self.queues.entry(channel_id).or_default(); // Push to front in reverse order so original order is preserved. for be in batch.events.into_iter().rev() { queue.push_front(QueuedEvent { - channel_id: batch.channel_id, + channel_id, event: be.event, prompt_tag: be.prompt_tag, - received_at: Instant::now(), + received_at: be.received_at, // preserve original timestamp (#46) }); } - self.retry_after - .insert(batch.channel_id, Instant::now() + Duration::from_secs(5)); + // Enforce per-channel cap: trim oldest (back) events if requeue pushed + // the queue over the limit. Without this, repeated requeue+push cycles + // can grow the queue unboundedly. + while queue.len() > MAX_PENDING_PER_CHANNEL { + queue.pop_back(); + tracing::warn!( + channel_id = %channel_id, + limit = MAX_PENDING_PER_CHANNEL, + "requeue overflow — dropped oldest event to enforce cap" + ); + } + self.retry_after.insert(channel_id, Instant::now() + delay); } /// Re-queue a batch preserving original `received_at` timestamps. @@ -218,8 +370,29 @@ impl EventQueue { /// Returns `true` if any channel has pending events that are not in-flight /// and not throttled by `retry_after`. - pub fn has_flushable_work(&self) -> bool { + /// + /// Also auto-expires any stuck in-flight entries whose deadline has passed. + /// This is a `&mut self` method so expiry can happen without requiring a + /// full `flush_next` call. + pub fn has_flushable_work(&mut self) -> bool { let now = Instant::now(); + + // Auto-expire stuck in-flight entries (same logic as flush_next). + let expired: Vec = self + .in_flight_deadlines + .iter() + .filter(|(_, deadline)| now >= **deadline) + .map(|(id, _)| *id) + .collect(); + for id in expired { + tracing::error!( + channel_id = %id, + "BUG: in-flight channel expired without mark_complete — auto-releasing" + ); + self.in_flight_channels.remove(&id); + self.in_flight_deadlines.remove(&id); + } + self.queues.iter().any(|(id, q)| { !q.is_empty() && !self.in_flight_channels.contains(id) @@ -266,6 +439,12 @@ impl EventQueue { .map(|q| q.into_iter().map(|e| e.event.id.to_hex()).collect()) .unwrap_or_default(); self.retry_after.remove(&channel_id); + self.retry_counts.remove(&channel_id); + // Preserve in_flight_channels AND in_flight_deadlines: the in-flight + // task will eventually complete (calling mark_complete) or the deadline + // will expire (auto-cleaning the channel). Removing deadlines without + // removing in_flight_channels would disable auto-expiry and leave a + // wedged task permanently blocking the channel. ids } @@ -274,6 +453,35 @@ impl EventQueue { pub fn is_channel_in_flight(&self, channel_id: Uuid) -> bool { self.in_flight_channels.contains(&channel_id) } + + /// Compact expired metadata entries to prevent unbounded map growth. + /// + /// Removes `retry_after` entries whose deadline has already passed, and + /// cleans up orphaned `retry_counts` entries for channels that have no + /// queued events, no active throttle, and no in-flight prompt. Without + /// this, channels that completed their retry cycle but never received + /// fresh traffic would leak a `u32` entry in `retry_counts` indefinitely. + /// + /// The in-flight guard is critical: a channel whose throttle expired and + /// whose queue is empty because it was flushed may still have a retry + /// attempt in flight. Removing its `retry_counts` would reset the + /// backoff sequence if that attempt fails and requeues. + /// + /// Should be called periodically from the main event loop (e.g., every + /// 30 seconds). `flush_next` and `has_flushable_work` handle in-flight + /// expiry inline; this covers the `retry_after` and `retry_counts` maps. + pub fn compact_expired_state(&mut self) { + let now = Instant::now(); + self.retry_after.retain(|_, deadline| *deadline > now); + // Remove retry_counts for channels with no active throttle, no + // queued events, AND no in-flight prompt — they completed their + // retry cycle and are truly idle. + self.retry_counts.retain(|ch, _| { + self.retry_after.contains_key(ch) + || self.queues.get(ch).is_some_and(|q| !q.is_empty()) + || self.in_flight_channels.contains(ch) + }); + } } impl Default for EventQueue { @@ -641,7 +849,13 @@ pub fn format_prompt( // one the agent is responding to. Thread/DM context is supplementary info // included alongside, not a scope override. This prevents mixed batches // (thread reply + later plain message) from being mislabeled as "thread". - let last_event = batch.events.last().expect("batch must have ≥1 event"); + let last_event = match batch.events.last() { + Some(e) => e, + None => { + tracing::error!("format_prompt called with empty batch — returning empty prompt"); + return String::new(); + } + }; let thread_tags = parse_thread_tags(&last_event.event); let is_dm = channel_info .map(|ci| ci.channel_type == "dm") @@ -1903,4 +2117,84 @@ mod tests { assert_eq!(drained.len(), 1); assert!(q.is_in_flight()); // in-flight unaffected } + + // ── compact_expired_state ───────────────────────────────────────────── + + #[test] + fn test_compact_cleans_orphaned_retry_counts() { + let mut q = EventQueue::new(DedupMode::Queue); + let ch = Uuid::new_v4(); + + // Simulate: push, flush, requeue (sets retry_after + retry_counts), + // then mark_complete (preserves retry_counts because throttle is active). + q.push(make_queued(ch, "msg1")); + let batch = q.flush_next().unwrap(); + q.requeue(batch); + q.mark_complete(ch); + assert!(q.retry_after.contains_key(&ch)); + assert!(q.retry_counts.contains_key(&ch)); + + // The requeued event is back in the queue. Flush it again so the + // queue is empty (simulating a successful retry dispatch). + // We need to wait for retry_after to expire first. + q.retry_after + .insert(ch, Instant::now() - Duration::from_secs(1)); + let _batch2 = q.flush_next().unwrap(); + // Now mark_complete with no active throttle — clears retry_counts. + q.mark_complete(ch); + assert!(!q.retry_counts.contains_key(&ch)); + + // Re-create the orphan scenario: manually insert stale retry_counts + // with no queue, no throttle, and no in-flight. + q.retry_counts.insert(ch, 3); + q.compact_expired_state(); + assert!( + !q.retry_counts.contains_key(&ch), + "orphaned retry_counts should be removed" + ); + } + + #[test] + fn test_compact_preserves_retry_counts_when_in_flight() { + let mut q = EventQueue::new(DedupMode::Queue); + let ch = Uuid::new_v4(); + + // Push, flush, requeue, mark_complete — sets up retry state. + q.push(make_queued(ch, "msg1")); + let batch = q.flush_next().unwrap(); + q.requeue(batch); + q.mark_complete(ch); + + // Expire the throttle so the requeued event can be flushed. + q.retry_after + .insert(ch, Instant::now() - Duration::from_secs(1)); + let _batch2 = q.flush_next().unwrap(); + // Channel is now in-flight with empty queue and expired throttle. + assert!(q.in_flight_channels.contains(&ch)); + assert!(q.queues.get(&ch).is_none_or(|q| q.is_empty())); + + // compact must NOT remove retry_counts — the in-flight attempt + // may fail and requeue, which needs the existing count. + q.compact_expired_state(); + assert!( + q.retry_counts.contains_key(&ch), + "retry_counts must survive while channel is in-flight" + ); + } + + #[test] + fn test_compact_preserves_retry_counts_with_queued_events() { + let mut q = EventQueue::new(DedupMode::Queue); + let ch = Uuid::new_v4(); + + // Manually set up: retry_counts exists, queue is non-empty, no throttle. + q.push(make_queued(ch, "msg1")); + q.retry_counts.insert(ch, 2); + + q.compact_expired_state(); + assert!( + q.retry_counts.contains_key(&ch), + "retry_counts should survive when queue is non-empty" + ); + } } diff --git a/crates/sprout-acp/src/relay.rs b/crates/sprout-acp/src/relay.rs index 61d547f5f4..e368641eb0 100644 --- a/crates/sprout-acp/src/relay.rs +++ b/crates/sprout-acp/src/relay.rs @@ -37,8 +37,21 @@ fn event_channel_capacity() -> usize { .map(|v| v.max(1)) // mpsc::channel panics on capacity 0 .unwrap_or(EVENT_CHANNEL_CAPACITY_DEFAULT) } -/// Maximum number of seen event IDs before the dedup set is cleared. +/// Maximum number of seen event IDs before the dedup set is rotated. +/// Two-generation dedup: each generation holds up to SEEN_ID_LIMIT/2 entries. const SEEN_ID_LIMIT: usize = 12_000; + +/// Interval between client-initiated WebSocket pings. +const PING_INTERVAL: Duration = Duration::from_secs(30); +/// If no pong is received within this duration after a ping, the connection is +/// considered dead and the background task triggers a reconnect. +const PONG_TIMEOUT: Duration = Duration::from_secs(10); +/// Timeout for individual ws.send() calls. Prevents a stalled socket from +/// wedging the background task indefinitely. +const WS_SEND_TIMEOUT_SECS: u64 = 10; +/// Diagnostic threshold: log when a connection has been stable for this long. +/// No backoff reset is implemented yet — this is a hook for future improvement. +const STABLE_CONNECTION_SECS: u64 = 60; /// Seconds subtracted from `since` on resubscribe to tolerate clock skew. const SINCE_SKEW_SECS: u64 = 5; /// Timeout for the NIP-42 auth handshake steps. @@ -46,6 +59,8 @@ 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 std::time::Instant; + use futures_util::{SinkExt, StreamExt}; use nostr::{Event, EventBuilder, Keys, Kind, Tag, Url as NostrUrl}; use serde_json::{json, Value}; @@ -81,26 +96,92 @@ pub struct RestClient { pub keys: Keys, } +/// Whether an HTTP status code is retriable (transient server/rate-limit errors). +fn is_retriable_status(status: reqwest::StatusCode) -> bool { + matches!(status.as_u16(), 429 | 502 | 503 | 504) +} + +/// Base retry delays for transient HTTP failures: 500ms, 1s, 2s. +/// Jitter (±20%) is applied at call time via `jittered_duration`. +const REST_RETRY_BASE_DELAYS: [Duration; 3] = [ + Duration::from_millis(500), + Duration::from_millis(1000), + Duration::from_millis(2000), +]; + impl RestClient { - /// GET a JSON endpoint, returning the parsed value on success. - pub async fn get_json(&self, path: &str) -> Result { - let url = format!("{}{}", self.base_url, path); - let builder = self.http.get(&url); - let builder = apply_auth(builder, &self.api_token, &self.keys); + /// Retry helper: executes `build_request` up to 4 times (1 attempt + 3 retries) + /// on transient failures (429, 502, 503, 504, timeout, connect errors). + /// Retry delays are jittered to prevent thundering-herd. + /// + /// Safety: all Sprout REST endpoints used by the harness are idempotent or + /// deduplicated server-side. GET/PUT/DELETE are inherently safe to retry. + /// POST /api/events publishes signed Nostr events whose IDs are deterministic + /// hashes — the relay deduplicates by event ID per NIP-01, so retries cannot + /// produce duplicate side effects. + async fn request_with_retry( + &self, + method: &str, + path: &str, + build_request: F, + ) -> Result + where + F: Fn() -> Fut, + Fut: std::future::Future>, + { + let mut last_err = None; - let resp = builder - .send() - .await - .map_err(|e| RelayError::Http(e.to_string()))?; + for (attempt, delay) in std::iter::once(None) + .chain(REST_RETRY_BASE_DELAYS.iter().map(|d| Some(*d))) + .enumerate() + { + if let Some(base) = delay { + let jittered = jittered_duration(base); + tracing::debug!( + "retrying {method} {path} (attempt {attempt}) in {:.1}s", + jittered.as_secs_f64() + ); + tokio::time::sleep(jittered).await; + } - if !resp.status().is_success() { - return Err(RelayError::Http(format!( - "GET {} returned HTTP {}", - path, - resp.status() - ))); + match build_request().await { + Ok(resp) if resp.status().is_success() => return Ok(resp), + Ok(resp) if is_retriable_status(resp.status()) => { + let status = resp.status(); + tracing::warn!("{method} {path} returned retriable HTTP {status}"); + last_err = Some(RelayError::Http(format!( + "{method} {path} returned HTTP {status}" + ))); + } + Ok(resp) => { + // Non-retriable error (401, 403, 404, etc.) — fail immediately. + return Err(RelayError::Http(format!( + "{method} {} returned HTTP {}", + path, + resp.status() + ))); + } + Err(e) if e.is_timeout() || e.is_connect() => { + tracing::warn!("{method} {path} network error: {e}"); + last_err = Some(RelayError::Http(e.to_string())); + } + Err(e) => return Err(RelayError::Http(e.to_string())), + } } + Err(last_err + .unwrap_or_else(|| RelayError::Http(format!("{method} {path} failed after retries")))) + } + + /// GET a JSON endpoint with retry on transient failures (429, 502, 503, 504). + pub async fn get_json(&self, path: &str) -> Result { + let url = format!("{}{}", self.base_url, path); + let resp = self + .request_with_retry("GET", path, || { + let builder = apply_auth(self.http.get(&url), &self.api_token, &self.keys); + builder.send() + }) + .await?; resp.json() .await .map_err(|e| RelayError::Http(e.to_string())) @@ -111,22 +192,14 @@ impl RestClient { /// Returns `Value::Null` for empty response bodies (e.g. 204 No Content). pub async fn put_json(&self, path: &str, body: &Value) -> Result { let url = format!("{}{}", self.base_url, path); - let builder = self.http.put(&url).json(body); - let builder = apply_auth(builder, &self.api_token, &self.keys); - - let resp = builder - .send() - .await - .map_err(|e| RelayError::Http(e.to_string()))?; - - if !resp.status().is_success() { - return Err(RelayError::Http(format!( - "PUT {} returned HTTP {}", - path, - resp.status() - ))); - } - + let body = body.clone(); + let resp = self + .request_with_retry("PUT", path, || { + let builder = + apply_auth(self.http.put(&url).json(&body), &self.api_token, &self.keys); + builder.send() + }) + .await?; let text = resp .text() .await @@ -142,22 +215,17 @@ impl RestClient { /// Returns `Value::Null` for empty response bodies (e.g. 204 No Content). pub async fn post_json(&self, path: &str, body: &Value) -> Result { let url = format!("{}{}", self.base_url, path); - let builder = self.http.post(&url).json(body); - let builder = apply_auth(builder, &self.api_token, &self.keys); - - let resp = builder - .send() - .await - .map_err(|e| RelayError::Http(e.to_string()))?; - - if !resp.status().is_success() { - return Err(RelayError::Http(format!( - "POST {} returned HTTP {}", - path, - resp.status() - ))); - } - + let body = body.clone(); + let resp = self + .request_with_retry("POST", path, || { + let builder = apply_auth( + self.http.post(&url).json(&body), + &self.api_token, + &self.keys, + ); + builder.send() + }) + .await?; let text = resp .text() .await @@ -172,21 +240,11 @@ impl RestClient { #[allow(dead_code)] pub async fn delete(&self, path: &str) -> Result<(), RelayError> { let url = format!("{}{}", self.base_url, path); - let builder = self.http.delete(&url); - let builder = apply_auth(builder, &self.api_token, &self.keys); - - let resp = builder - .send() - .await - .map_err(|e| RelayError::Http(e.to_string()))?; - - if !resp.status().is_success() { - return Err(RelayError::Http(format!( - "DELETE {} returned HTTP {}", - path, - resp.status() - ))); - } + self.request_with_retry("DELETE", path, || { + let builder = apply_auth(self.http.delete(&url), &self.api_token, &self.keys); + builder.send() + }) + .await?; Ok(()) } } @@ -286,6 +344,10 @@ enum RelayCommand { SubscribeMembership, /// Publish a signed event to the relay (for typing indicators, etc.). PublishEvent { event: Box }, + /// Set the startup watermark timestamp for Finding #22. + /// The background task uses this as the floor `since` for membership + /// notification replay so events before startup are never re-delivered. + SetStartupWatermark { ts: u64 }, } // ── WebSocket stream type alias ─────────────────────────────────────────────── @@ -318,7 +380,9 @@ pub struct HarnessRelay { #[allow(dead_code)] agent_pubkey_hex: String, /// Handle to the background task (for clean shutdown). - bg_handle: tokio::task::JoinHandle<()>, + /// Wrapped in `Option` so `shutdown()` can take ownership without conflicting + /// with `Drop` (which only has `&mut self`). + bg_handle: Option>, } impl HarnessRelay { @@ -332,7 +396,9 @@ impl HarnessRelay { agent_pubkey_hex: &str, ) -> Result { // Perform the initial connection and auth handshake. - let (ws, _buffer) = do_connect(relay_url, keys, api_token).await?; + // Finding #8: capture the handshake buffer and pass it to the background + // task so buffered messages aren't silently discarded. + let (ws, handshake_buffer) = do_connect(relay_url, keys, api_token).await?; let (event_tx, event_rx) = mpsc::channel::>(event_channel_capacity()); let (cmd_tx, cmd_rx) = mpsc::channel::(CMD_CHANNEL_CAPACITY); @@ -345,6 +411,7 @@ impl HarnessRelay { let bg_handle = tokio::spawn(async move { run_background_task( ws, + handshake_buffer, event_tx, cmd_rx, bg_keys, @@ -367,34 +434,17 @@ impl HarnessRelay { api_token: api_token.map(|t| t.to_string()), keys: keys.clone(), agent_pubkey_hex: agent_pubkey_hex.to_string(), - bg_handle, + bg_handle: Some(bg_handle), }) } /// Discover channels the agent is a member of via `GET /api/channels?member=true`. + /// + /// Uses the retry-enabled `RestClient::get_json` so transient 502/503/429 + /// errors during startup don't abort the harness. 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?member=true"); - - let builder = self.http.get(&url); - let builder = apply_auth(builder, &self.api_token, &self.keys); - - let resp = builder - .send() - .await - .map_err(|e| RelayError::Http(e.to_string()))?; - - if !resp.status().is_success() { - return Err(RelayError::Http(format!( - "GET /api/channels returned HTTP {}", - resp.status() - ))); - } - - let body: Value = resp - .json() - .await - .map_err(|e| RelayError::Http(e.to_string()))?; + let rest = self.rest_client(); + let body = rest.get_json("/api/channels?member=true").await?; let channels = body .as_array() @@ -488,6 +538,10 @@ impl HarnessRelay { } /// Publish a signed event to the relay via the background WebSocket task. + /// + /// Blocks until the command channel has capacity. For ephemeral events + /// (typing indicators) prefer [`try_publish_event`] which never blocks. + #[allow(dead_code)] // Public API — callers outside the harness may use this pub async fn publish_event(&self, event: Event) -> Result<(), RelayError> { self.cmd_tx .send(RelayCommand::PublishEvent { @@ -497,6 +551,18 @@ impl HarnessRelay { .map_err(|_| RelayError::ConnectionClosed) } + /// Fire-and-forget publish — uses `try_send` so it never blocks the caller. + /// + /// Suitable for ephemeral commands like typing indicators where dropping + /// the event on a full command channel is acceptable. + pub fn try_publish_event(&self, event: Event) -> Result<(), RelayError> { + self.cmd_tx + .try_send(RelayCommand::PublishEvent { + event: Box::new(event), + }) + .map_err(|_| RelayError::ConnectionClosed) + } + /// Build a typing indicator event (kind:20002) for a channel. pub fn build_typing_event(&self, channel_id: Uuid) -> Result { let h_tag = Tag::parse(&["h", &channel_id.to_string()]) @@ -506,6 +572,19 @@ impl HarnessRelay { Ok(event) } + /// Set the startup watermark timestamp (Finding #22). + /// + /// Call this once after `connect()` with the Unix timestamp captured just + /// before the relay connection was established. The background task uses + /// this as the floor `since` for membership notification replay so events + /// predating this session are never re-delivered after reconnect. + pub async fn set_startup_watermark(&self, ts: u64) -> Result<(), RelayError> { + self.cmd_tx + .send(RelayCommand::SetStartupWatermark { ts }) + .await + .map_err(|_| RelayError::ConnectionClosed) + } + /// Reconnect after connection loss. Instructs the background task to /// re-authenticate and resubscribe to all previously active channels. pub async fn reconnect(&mut self) -> Result<(), RelayError> { @@ -518,24 +597,95 @@ impl HarnessRelay { } } +impl HarnessRelay { + /// Graceful async shutdown — sends Shutdown command and waits up to 5s for + /// the background task to finish. Use this from async contexts instead of + /// relying on `Drop` (which aborts immediately). + pub async fn shutdown(mut self) { + let _ = self.cmd_tx.send(RelayCommand::Shutdown).await; + if let Some(handle) = self.bg_handle.take() { + let abort_handle = handle.abort_handle(); + if tokio::time::timeout(Duration::from_secs(5), handle) + .await + .is_err() + { + tracing::warn!("relay background task did not finish in 5s — aborting"); + abort_handle.abort(); + } + } + } +} + impl Drop for HarnessRelay { fn drop(&mut self) { // Best-effort shutdown signal; ignore errors (task may already be done). let _ = self.cmd_tx.try_send(RelayCommand::Shutdown); - self.bg_handle.abort(); + if let Some(handle) = self.bg_handle.take() { + handle.abort(); + } } } // ── Background task ─────────────────────────────────────────────────────────── +/// Two-generation dedup set with bounded memory. +/// +/// Mitigates the "amnesia window" caused by clearing the entire set at once. +/// When `current` reaches `limit/2` entries it is rotated into `previous`. +/// At any point we remember between `limit/2` and `limit` recent IDs. +/// The oldest `limit/2` IDs are forgotten on each rotation — this is the +/// inherent tradeoff of bounded-memory dedup. For the default limit of +/// 12,000, the worst case is that an ID seen 6,001+ inserts ago may be +/// replayed as new. This is acceptable for Nostr event dedup where the +/// `since` filter provides the primary replay protection. +struct TwoGenDedup { + current: HashSet, + previous: HashSet, + limit: usize, +} + +impl TwoGenDedup { + fn new(limit: usize) -> Self { + Self { + current: HashSet::new(), + previous: HashSet::new(), + limit, + } + } + + fn contains(&self, id: &str) -> bool { + self.current.contains(id) || self.previous.contains(id) + } + + /// Insert `id`. Returns `true` if it was new (not a duplicate). + fn insert(&mut self, id: String) -> bool { + if self.contains(&id) { + return false; + } + self.current.insert(id); + if self.current.len() >= self.limit / 2 { + // Rotate: current → previous, start fresh current. + self.previous = std::mem::take(&mut self.current); + } + true + } + + /// Remove an ID (used to un-deduplicate a dropped event so it can be + /// replayed after reconnect). + fn remove(&mut self, id: &str) { + self.current.remove(id); + self.previous.remove(id); + } +} + /// State maintained by the background WebSocket task. struct BgState { /// Active subscriptions: channel_id → subscription_id string. active_subscriptions: HashMap, /// Most recent `created_at` timestamp seen per channel (for `since` filter). last_seen: HashMap, - /// Set of event IDs seen, for deduplication. - seen_ids: HashSet, + /// Two-generation dedup set of event IDs seen. + seen_ids: TwoGenDedup, /// Per-channel filter used on subscribe (for resubscribe after reconnect). active_filters: HashMap, /// Oldest timestamp of a membership notification that was dropped due to @@ -552,6 +702,21 @@ struct BgState { /// On reconnect resubscribe, `since` = min(last_seen, channel_dropped_since). /// Cleared per-channel after a successful resubscribe. channel_dropped_since: HashMap, + /// Set by the backpressure handler when the event channel is full. + /// The main loop checks this flag and triggers a proactive resubscribe + /// (without waiting for a disconnect) so dropped events are replayed. + proactive_resubscribe_needed: bool, + /// Unix timestamp captured just before the relay connection was established + /// (Finding #22). Used as the floor `since` for membership notification + /// replay so events predating this session are never re-delivered. + startup_watermark: Option, + /// Wall-clock timestamp when each channel was first subscribed. + /// Used as the `since` fallback on reconnect for channels that have no + /// `last_seen` or `channel_dropped_since`. This prevents channels joined + /// after startup from replaying from `startup_watermark` (which could be + /// hours old), while still allowing startup-era channels to use the + /// startup watermark via their `subscribe_since ≈ startup_watermark`. + subscribe_since: HashMap, } impl BgState { @@ -559,12 +724,15 @@ impl BgState { Self { active_subscriptions: HashMap::new(), last_seen: HashMap::new(), - seen_ids: HashSet::new(), + seen_ids: TwoGenDedup::new(SEEN_ID_LIMIT), active_filters: HashMap::new(), membership_dropped_since: None, membership_last_seen: None, membership_sub_active: false, channel_dropped_since: HashMap::new(), + proactive_resubscribe_needed: false, + startup_watermark: None, + subscribe_since: HashMap::new(), } } @@ -573,20 +741,11 @@ impl BgState { fn record_event(&mut self, channel_id: Uuid, event: &Event) -> bool { let id_hex = event.id.to_hex(); - // Deduplicate. + // Two-generation dedup: no amnesia window on rotation. if !self.seen_ids.insert(id_hex) { return false; } - // Bound seen_ids to prevent unbounded memory growth. - if self.seen_ids.len() > SEEN_ID_LIMIT { - // HashSet has no ordering, so we clear and re-insert the current - // event to avoid a false-negative dedup gap for this event. - let current_id = event.id.to_hex(); - self.seen_ids.clear(); - self.seen_ids.insert(current_id); - } - // Update last_seen timestamp. let ts = event.created_at.as_u64(); self.last_seen @@ -596,14 +755,193 @@ impl BgState { true } + + /// Clear all per-channel state for a channel that is being unsubscribed. + /// Prevents stale replay on re-subscribe and avoids unbounded state growth + /// for channels that are removed and never re-added. + fn clear_channel_state(&mut self, channel_id: &Uuid) { + self.last_seen.remove(channel_id); + self.subscribe_since.remove(channel_id); + self.channel_dropped_since.remove(channel_id); + self.active_filters.remove(channel_id); + } +} + +/// Record a command's intent in state while disconnected (no WebSocket). +/// +/// Subscribe/Unsubscribe/SubscribeMembership record intent so reconnect +/// restores the right subscriptions. SetStartupWatermark floors the replay +/// window. PublishEvent and Reconnect are no-ops while disconnected. +/// +/// Callers MUST handle `Shutdown` before calling — reaching the Shutdown +/// arm here is a logic error. +fn apply_command_to_state(state: &mut BgState, cmd: RelayCommand) { + match cmd { + RelayCommand::Subscribe { channel_id, filter } => { + state + .active_subscriptions + .insert(channel_id, channel_sub_id(channel_id)); + state.active_filters.insert(channel_id, filter); + state.subscribe_since.entry(channel_id).or_insert_with(|| { + // Use startup_watermark as floor when available — closes the + // blind spot between watermark capture and first REQ. + state.startup_watermark.unwrap_or_else(|| { + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_secs() + }) + }); + } + RelayCommand::Unsubscribe { channel_id } => { + state.active_subscriptions.remove(&channel_id); + state.clear_channel_state(&channel_id); + } + RelayCommand::SubscribeMembership => { + state.membership_sub_active = true; + } + RelayCommand::SetStartupWatermark { ts } => { + state.startup_watermark = Some(ts); + if state.membership_last_seen.is_none() { + state.membership_last_seen = Some(ts); + } + } + // Ephemeral events are meaningless while disconnected. + RelayCommand::PublishEvent { .. } => {} + // Already reconnecting — redundant. + RelayCommand::Reconnect => {} + // Callers MUST handle Shutdown before calling this function. + RelayCommand::Shutdown => { + debug_assert!( + false, + "Shutdown must be handled by caller, not apply_command_to_state" + ); + } + } +} + +/// Execute a command on a live WebSocket connection. +/// +/// Handles the five data commands: Subscribe, Unsubscribe, +/// SubscribeMembership, PublishEvent, SetStartupWatermark. Callers handle +/// Shutdown and Reconnect for control flow before dispatching here. +/// +/// Returns `true` if the command succeeded (or was a no-op). Returns `false` +/// if a WebSocket send failed — the caller should treat this as a dead socket +/// and trigger reconnect. On failure, subscription intent is preserved in +/// state via [`apply_command_to_state`] so reconnect will restore it. +async fn execute_connected_command( + ws: &mut WsStream, + state: &mut BgState, + agent_pubkey_hex: &str, + cmd: RelayCommand, +) -> bool { + match cmd { + RelayCommand::Subscribe { channel_id, filter } => { + // Seed subscribe_since BEFORE computing since — on first + // subscribe, this provides the fallback timestamp that + // closes the startup blind spot. Use startup_watermark as + // floor when available so events between watermark capture + // and this REQ are not missed. + state.subscribe_since.entry(channel_id).or_insert_with(|| { + state.startup_watermark.unwrap_or_else(|| { + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_secs() + }) + }); + let since = state + .last_seen + .get(&channel_id) + .copied() + .or_else(|| state.subscribe_since.get(&channel_id).copied()); + let sent = + send_subscribe(ws, state, channel_id, agent_pubkey_hex, since, &filter).await; + if sent { + state + .active_subscriptions + .insert(channel_id, channel_sub_id(channel_id)); + state.active_filters.insert(channel_id, filter); + true + } else { + // Send failed — record intent so reconnect restores it. + warn!("subscribe REQ failed for channel {channel_id} — recording intent for reconnect"); + apply_command_to_state(state, RelayCommand::Subscribe { channel_id, filter }); + false + } + } + 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) { + // Best-effort CLOSE — don't fail the command if send fails, + // because the intent (unsubscribe) is already applied to state. + let _ = + ws_send_timeout(ws, Message::Text(text.into()), WS_SEND_TIMEOUT_SECS).await; + } + debug!("unsubscribed from channel {channel_id}"); + } + state.clear_channel_state(&channel_id); + true + } + RelayCommand::SubscribeMembership => { + let since = state.membership_last_seen.or(state.startup_watermark); + let sent = send_membership_subscribe(ws, agent_pubkey_hex, since).await; + if sent { + state.membership_sub_active = true; + if state.membership_last_seen.is_none() { + state.membership_last_seen = since; + } + true + } else { + // Send failed — record intent so reconnect restores it. + warn!("membership subscribe REQ failed — recording intent for reconnect"); + state.membership_sub_active = true; + false + } + } + RelayCommand::PublishEvent { event } => { + let msg = json!(["EVENT", event]); + if let Ok(text) = serde_json::to_string(&msg) { + if let Err(e) = + ws_send_timeout(ws, Message::Text(text.into()), WS_SEND_TIMEOUT_SECS).await + { + // Ephemeral events (typing indicators) are best-effort. + // Log the failure but don't trigger reconnect — the next + // ping or read will detect the dead socket. + warn!("failed to publish event: {e}"); + } + } + true + } + RelayCommand::SetStartupWatermark { ts } => { + state.startup_watermark = Some(ts); + if state.membership_last_seen.is_none() { + state.membership_last_seen = Some(ts); + } + debug!("startup watermark set to {ts}"); + true + } + // Control-flow commands — callers handle these before dispatching. + RelayCommand::Shutdown | RelayCommand::Reconnect => { + debug_assert!( + false, + "Shutdown/Reconnect must be handled by caller, not execute_connected_command" + ); + true + } + } } /// The main background task loop. /// /// Owns the WebSocket stream, responds to Pings, forwards events, and handles /// reconnection. +#[allow(clippy::too_many_arguments)] async fn run_background_task( mut ws: WsStream, + initial_handshake_buffer: std::collections::VecDeque, event_tx: mpsc::Sender>, mut cmd_rx: mpsc::Receiver, keys: Keys, @@ -613,23 +951,165 @@ async fn run_background_task( ) { let mut state = BgState::new(); + // Finding #8: process any messages buffered during the initial auth handshake. + // If a buffered message signals connection drop, trigger reconnect immediately. + let handshake_ok = process_handshake_buffer( + &mut ws, + initial_handshake_buffer, + &event_tx, + &mut state, + &keys, + &relay_url, + api_token.as_deref(), + &agent_pubkey_hex, + ) + .await; + if !handshake_ok { + warn!("handshake buffer contained a drop signal — attempting autonomous reconnect"); + // Don't wait for a caller-driven Reconnect command — the caller was + // never notified (no sentinel sent). Go straight to reconnect loop. + let _ = event_tx.try_send(None); + match try_autonomous_reconnect( + &mut ws, + &mut cmd_rx, + &mut state, + &keys, + &relay_url, + api_token.as_deref(), + &agent_pubkey_hex, + &event_tx, + ) + .await + { + ReconnectOutcome::Ok => { + if matches!( + drain_post_reconnect(&mut ws, &mut cmd_rx, &mut state, &agent_pubkey_hex).await, + ReconnectOutcome::Shutdown + ) { + return; + } + } + ReconnectOutcome::Shutdown => return, + ReconnectOutcome::Failed => { + if matches!( + wait_for_reconnect( + &mut ws, + &mut cmd_rx, + &mut state, + &keys, + &relay_url, + api_token.as_deref(), + &agent_pubkey_hex, + &event_tx, + true, + ) + .await, + ReconnectOutcome::Shutdown + ) { + return; + } + } + } + // ping_sent, last_pong, connected_since are initialized below — + // no reset needed here since they haven't been declared yet. + } + + // Finding #31: client-initiated ping to detect silent connection death. + let mut ping_interval = tokio::time::interval(PING_INTERVAL); + ping_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + let mut last_pong = Instant::now(); + let mut ping_sent = false; + + // Finding #42: track connection stability for backoff reset. + let mut connected_since = Instant::now(); + let mut stable_logged = false; + loop { + // Finding #3: check proactive resubscribe flag before blocking on select! + if state.proactive_resubscribe_needed { + state.proactive_resubscribe_needed = false; + info!("proactive resubscribe triggered by backpressure event loss"); + if !resubscribe_after_reconnect(&mut ws, &mut state, &agent_pubkey_hex).await { + warn!("proactive resubscribe had failures — triggering reconnect"); + let _ = event_tx.try_send(None); + match try_autonomous_reconnect( + &mut ws, + &mut cmd_rx, + &mut state, + &keys, + &relay_url, + api_token.as_deref(), + &agent_pubkey_hex, + &event_tx, + ) + .await + { + ReconnectOutcome::Ok => { + if matches!( + drain_post_reconnect( + &mut ws, + &mut cmd_rx, + &mut state, + &agent_pubkey_hex + ) + .await, + ReconnectOutcome::Shutdown + ) { + return; + } + } + ReconnectOutcome::Shutdown => return, + ReconnectOutcome::Failed => { + if matches!( + wait_for_reconnect( + &mut ws, + &mut cmd_rx, + &mut state, + &keys, + &relay_url, + api_token.as_deref(), + &agent_pubkey_hex, + &event_tx, + true, + ) + .await, + ReconnectOutcome::Shutdown + ) { + return; + } + } + } + ping_sent = false; + last_pong = Instant::now(); + connected_since = Instant::now(); + stable_logged = false; + } + } + tokio::select! { // ── Incoming WebSocket message ──────────────────────────────────── raw = ws.next() => { // Determine if the socket is lost. let socket_lost = match raw { Some(Ok(msg)) => { - !handle_ws_message( - msg, - &mut ws, - &event_tx, - &mut state, - &keys, - &relay_url, - api_token.as_deref(), - ) - .await + // Finding #31: track pong replies directly, before dispatch. + if matches!(msg, Message::Pong(_)) { + last_pong = Instant::now(); + ping_sent = false; + false // pong is healthy — not a socket loss + } else { + !handle_ws_message( + msg, + &mut ws, + &event_tx, + &mut state, + &keys, + &relay_url, + api_token.as_deref(), + &agent_pubkey_hex, + ) + .await + } } Some(Err(e)) => { warn!("WebSocket error in background task: {e}"); @@ -643,131 +1123,194 @@ async fn run_background_task( if socket_lost { // Signal the caller, then attempt autonomous reconnect. - let _ = event_tx.send(None).await; - let reconnected = try_autonomous_reconnect( + // Use try_send to avoid blocking on backpressure — recovery + // must not stall when the event channel is full. + let _ = event_tx.try_send(None); + let outcome = try_autonomous_reconnect( &mut ws, + &mut cmd_rx, &mut state, &keys, &relay_url, api_token.as_deref(), &agent_pubkey_hex, + &event_tx, ) .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; - } - RelayCommand::PublishEvent { event } => { - let msg = json!(["EVENT", event]); - if let Ok(text) = serde_json::to_string(&msg) { - let _ = ws.send(Message::Text(text.into())).await; - } - } - } - } - } else { - wait_for_reconnect( - &mut ws, - &mut cmd_rx, - &mut state, - &keys, - &relay_url, - api_token.as_deref(), - &agent_pubkey_hex, - false, - ) - .await; + match outcome { + ReconnectOutcome::Shutdown => return, + ReconnectOutcome::Ok => { + if matches!( + drain_post_reconnect(&mut ws, &mut cmd_rx, &mut state, &agent_pubkey_hex).await, + ReconnectOutcome::Shutdown + ) { return; } + // Reset ping state after reconnect. + ping_sent = false; + last_pong = Instant::now(); + connected_since = Instant::now(); + stable_logged = false; } + ReconnectOutcome::Failed => { + if matches!( + wait_for_reconnect( + &mut ws, &mut cmd_rx, &mut state, &keys, &relay_url, + api_token.as_deref(), &agent_pubkey_hex, &event_tx, true, + ).await, + ReconnectOutcome::Shutdown + ) { return; } + ping_sent = false; + last_pong = Instant::now(); + connected_since = Instant::now(); + stable_logged = false; + } + } // end match outcome } } // ── Command from HarnessRelay ───────────────────────────────────── cmd = cmd_rx.recv() => { match cmd { - Some(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); - } - Some(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; - } - debug!("unsubscribed from channel {channel_id}"); - } - } - Some(RelayCommand::SubscribeMembership) => { - let _ = - send_membership_subscribe(&mut ws, &agent_pubkey_hex, None).await; - state.membership_sub_active = true; - // Seed the watermark so reconnect replays from this point - // rather than falling back to since=now (which could miss - // notifications during the reconnect gap). - if state.membership_last_seen.is_none() { - state.membership_last_seen = Some( - std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap_or_default() - .as_secs(), - ); - } - } - Some(RelayCommand::PublishEvent { event }) => { - let msg = json!(["EVENT", event]); - if let Ok(text) = serde_json::to_string(&msg) { - if let Err(e) = ws.send(Message::Text(text.into())).await { - warn!("failed to publish event: {e}"); - } - } - } Some(RelayCommand::Reconnect) => { - // Reconnect command already consumed — skip the drain loop. - wait_for_reconnect( - &mut ws, - &mut cmd_rx, - &mut state, - &keys, - &relay_url, - api_token.as_deref(), - &agent_pubkey_hex, - true, // skip_drain: command already consumed - ) - .await; + if matches!( + wait_for_reconnect( + &mut ws, &mut cmd_rx, &mut state, &keys, &relay_url, + api_token.as_deref(), &agent_pubkey_hex, &event_tx, true, + ).await, + ReconnectOutcome::Shutdown + ) { return; } + ping_sent = false; + last_pong = Instant::now(); + connected_since = Instant::now(); + stable_logged = false; } Some(RelayCommand::Shutdown) | None => { - debug!("background task shutting down"); + debug!("background task shutting down — sending close frame"); + let _ = ws_send_timeout( + &mut ws, + Message::Close(None), + WS_SEND_TIMEOUT_SECS, + ) + .await; return; } + Some(cmd) => { + let ok = execute_connected_command( + &mut ws, + &mut state, + &agent_pubkey_hex, + cmd, + ) + .await; + if !ok { + // Send failed — socket is likely dead. Trigger reconnect. + warn!("command send failed — triggering reconnect"); + let _ = event_tx.try_send(None); + match try_autonomous_reconnect( + &mut ws, &mut cmd_rx, &mut state, &keys, &relay_url, + api_token.as_deref(), &agent_pubkey_hex, &event_tx, + ).await { + ReconnectOutcome::Shutdown => return, + ReconnectOutcome::Ok => { + if matches!( + drain_post_reconnect(&mut ws, &mut cmd_rx, &mut state, &agent_pubkey_hex).await, + ReconnectOutcome::Shutdown + ) { return; } + } + ReconnectOutcome::Failed => { + if matches!( + wait_for_reconnect( + &mut ws, &mut cmd_rx, &mut state, &keys, &relay_url, + api_token.as_deref(), &agent_pubkey_hex, &event_tx, true, + ).await, + ReconnectOutcome::Shutdown + ) { return; } + } + } + ping_sent = false; + last_pong = Instant::now(); + connected_since = Instant::now(); + stable_logged = false; + } + } } } + + // ── Finding #31: client-initiated ping ──────────────────────────── + _ = 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); + // Use try_send to avoid blocking on backpressure during recovery. + let _ = event_tx.try_send(None); + match try_autonomous_reconnect( + &mut ws, &mut cmd_rx, &mut state, &keys, &relay_url, + api_token.as_deref(), &agent_pubkey_hex, &event_tx, + ).await { + ReconnectOutcome::Shutdown => return, + ReconnectOutcome::Ok => { + if matches!( + drain_post_reconnect(&mut ws, &mut cmd_rx, &mut state, &agent_pubkey_hex).await, + ReconnectOutcome::Shutdown + ) { return; } + } + ReconnectOutcome::Failed => { + if matches!( + wait_for_reconnect( + &mut ws, &mut cmd_rx, &mut state, &keys, &relay_url, + api_token.as_deref(), &agent_pubkey_hex, &event_tx, true, + ).await, + ReconnectOutcome::Shutdown + ) { return; } + } + } + ping_sent = false; + last_pong = Instant::now(); + connected_since = Instant::now(); + stable_logged = false; + } else if !ping_sent { + if let Err(e) = ws_send_timeout(&mut ws, Message::Ping(vec![].into()), WS_SEND_TIMEOUT_SECS).await { + warn!("failed to send ping: {e} — triggering reconnect"); + // Use try_send to avoid blocking on backpressure during recovery. + let _ = event_tx.try_send(None); + match try_autonomous_reconnect( + &mut ws, &mut cmd_rx, &mut state, &keys, &relay_url, + api_token.as_deref(), &agent_pubkey_hex, &event_tx, + ).await { + ReconnectOutcome::Shutdown => return, + ReconnectOutcome::Ok => { + if matches!( + drain_post_reconnect(&mut ws, &mut cmd_rx, &mut state, &agent_pubkey_hex).await, + ReconnectOutcome::Shutdown + ) { return; } + } + ReconnectOutcome::Failed => { + if matches!( + wait_for_reconnect( + &mut ws, &mut cmd_rx, &mut state, &keys, &relay_url, + api_token.as_deref(), &agent_pubkey_hex, &event_tx, true, + ).await, + ReconnectOutcome::Shutdown + ) { return; } + } + } + ping_sent = false; + last_pong = Instant::now(); + connected_since = Instant::now(); + stable_logged = false; + } else { + ping_sent = true; + debug!("sent ping to relay"); + } + } + } + } + + // Finding #42: log when connection has been stable for STABLE_CONNECTION_SECS. + // Log once when the connection has been stable. Diagnostic only. + if !stable_logged && connected_since.elapsed() > Duration::from_secs(STABLE_CONNECTION_SECS) + { + stable_logged = true; + debug!("connection stable for >{}s", STABLE_CONNECTION_SECS); } } } @@ -776,6 +1319,7 @@ async fn run_background_task( /// /// Returns `false` if the connection has been lost (Close frame or unrecoverable /// error), `true` otherwise. +#[allow(clippy::too_many_arguments)] async fn handle_ws_message( msg: Message, ws: &mut WsStream, @@ -784,6 +1328,7 @@ async fn handle_ws_message( keys: &Keys, relay_url: &str, api_token: Option<&str>, + agent_pubkey_hex: &str, ) -> bool { match msg { Message::Text(text) => { @@ -809,25 +1354,57 @@ async fn handle_ws_message( return true; } }; + // Dedup membership notifications through TwoGenDedup. + // We use seen_ids directly instead of record_event() + // because record_event() also updates last_seen, which + // would contaminate per-channel replay watermarks with + // membership-event timestamps and cause channel event + // loss on reconnect. + let event_id_hex = event.id.to_hex(); + if !state.seen_ids.insert(event_id_hex.clone()) { + debug!( + channel_id = %channel_uuid, + event_id = %event_id_hex, + "duplicate membership notification — skipping" + ); + return true; + } let ts = event.created_at.as_u64(); let sprout_event = SproutEvent { channel_id: channel_uuid, event: *event, }; + // Finding #3: warn at 80% capacity. + let cap = event_tx.max_capacity(); + let used = cap - event_tx.capacity(); + if used >= (cap * 4 / 5) { + warn!( + used, + capacity = cap, + "event channel at ≥80% capacity — backpressure imminent" + ); + } match event_tx.try_send(Some(sprout_event)) { Ok(()) => { state.membership_last_seen = Some(state.membership_last_seen.unwrap_or(0).max(ts)); } Err(mpsc::error::TrySendError::Full(_)) => { + // Remove from dedup so reconnect replay can + // re-deliver this event (it was never forwarded + // to the harness). + state.seen_ids.remove(&event_id_hex); // Track the oldest dropped timestamp so reconnect // replay starts early enough to re-deliver it. state.membership_dropped_since = Some(state.membership_dropped_since.map_or(ts, |d| d.min(ts))); + // Finding #3: proactively trigger resubscribe without + // waiting for a disconnect. + state.proactive_resubscribe_needed = true; warn!( channel_id = %channel_uuid, ts, - "membership notification dropped (backpressure) — will replay from {ts} on reconnect" + "membership notification dropped (backpressure) — proactive resubscribe queued" ); } Err(mpsc::error::TrySendError::Closed(_)) => return false, @@ -840,6 +1417,16 @@ async fn handle_ws_message( channel_id, event: *event, }; + // Finding #3: warn at 80% capacity. + let cap = event_tx.max_capacity(); + let used = cap - event_tx.capacity(); + if used >= (cap * 4 / 5) { + warn!( + used, + capacity = cap, + "event channel at ≥80% capacity — backpressure imminent" + ); + } match event_tx.try_send(Some(sprout_event)) { Ok(()) => {} Err(mpsc::error::TrySendError::Full(_)) => { @@ -853,10 +1440,12 @@ async fn handle_ws_message( .entry(channel_id) .and_modify(|d| *d = (*d).min(ts)) .or_insert(ts); + // Finding #3: proactively trigger resubscribe. + state.proactive_resubscribe_needed = true; warn!( channel_id = %channel_id, ts, - "event channel full — dropping event for channel {channel_id} — will replay from {ts} on reconnect" + "event channel full — dropping event for channel {channel_id} — proactive resubscribe queued" ); } Err(mpsc::error::TrySendError::Closed(_)) => { @@ -882,15 +1471,110 @@ async fn handle_ws_message( subscription_id, message, } => { - warn!("subscription {subscription_id} closed by relay: {message}"); + // Finding #15: CLOSED needs cleanup and resubscribe, not just logging. + // Classify the error to decide how to respond. + let is_auth_error = message.starts_with("auth-required") + || message.starts_with("restricted") + || message.contains("auth"); + warn!( + "subscription {subscription_id} closed by relay: {message}{}", + if is_auth_error { + " [auth error — reconnect required]" + } else { + "" + } + ); + + if is_auth_error { + // Auth errors require a full reconnect (re-handshake). + return false; + } + + // Attempt targeted resubscribe. State is NOT cleared before + // the attempt — if the send fails and triggers reconnect, + // resubscribe_after_reconnect() needs the subscription to + // still be in state so it can restore it. + if subscription_id == MEMBERSHIP_NOTIF_SUB_ID { + let 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) => state.startup_watermark, + }; + let sent = send_membership_subscribe(ws, agent_pubkey_hex, since).await; + if sent { + // Success — subscription is live again. + state.membership_dropped_since = None; + } else { + // Resubscribe failed — likely half-dead socket. + // Keep membership_sub_active = true so reconnect restores it. + warn!( + "membership resubscribe failed after CLOSED — triggering reconnect" + ); + return false; + } + } else if let Some(channel_id) = channel_id_from_sub_id(&subscription_id) { + // Guard: only resubscribe if the channel is still active. + // A delayed CLOSED for an already-unsubscribed channel must + // NOT resurrect the subscription (especially with a default + // permissive filter, which would be a fail-open regression). + if !state.active_subscriptions.contains_key(&channel_id) { + debug!("ignoring CLOSED for already-unsubscribed channel {channel_id}"); + } else { + let last_seen = state.last_seen.get(&channel_id).copied(); + let dropped = state.channel_dropped_since.get(&channel_id).copied(); + let subscribe_ts = state.subscribe_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) => subscribe_ts.or(state.startup_watermark), + }; + let filter = match state.active_filters.get(&channel_id).cloned() { + Some(f) => f, + None => { + // Fail closed: missing filter state means the subscription + // intent is inconsistent. Trigger reconnect rather than + // resubscribing with a permissive wildcard. + warn!("missing filter for channel {channel_id} after CLOSED — triggering reconnect (fail-closed)"); + return false; + } + }; + let sent = send_subscribe( + ws, + state, + channel_id, + agent_pubkey_hex, + since, + &filter, + ) + .await; + if sent { + // Success — update subscription ID (relay may assign new one). + state + .active_subscriptions + .insert(channel_id, channel_sub_id(channel_id)); + state.channel_dropped_since.remove(&channel_id); + } else { + // Resubscribe failed — likely half-dead socket. + // Keep channel in active_subscriptions so reconnect restores it. + warn!("channel {channel_id} resubscribe failed after CLOSED — triggering reconnect"); + return false; + } + } // end: channel is still active + } else { + warn!("CLOSED for unknown subscription {subscription_id} — ignoring"); + } } RelayMessage::Auth { challenge } => { - // Fix 5: Handle mid-session AUTH challenge by re-authenticating. + // Finding #18: AUTH send failure must trigger reconnect. debug!("received mid-session AUTH challenge — re-authenticating"); if let Err(e) = send_auth_response(ws, &challenge, relay_url, keys, api_token).await { - warn!("failed to respond to mid-session AUTH challenge: {e}"); + warn!("failed to respond to mid-session AUTH challenge: {e} — triggering reconnect"); + return false; } } RelayMessage::Ok { @@ -898,13 +1582,18 @@ async fn handle_ws_message( accepted, message, } => { + if !accepted && message.starts_with("auth") { + // Finding #18: AUTH OK with accepted=false means auth was rejected. + warn!("mid-session AUTH rejected (event {event_id}): {message} — triggering reconnect"); + return false; + } debug!("OK for event {event_id}: accepted={accepted} message={message}"); } } true } Message::Ping(data) => { - if let Err(e) = ws.send(Message::Pong(data)).await { + if let Err(e) = ws_send_timeout(ws, Message::Pong(data), WS_SEND_TIMEOUT_SECS).await { warn!("failed to send pong: {e}"); return false; } @@ -919,14 +1608,88 @@ async fn handle_ws_message( } } +/// Process messages buffered during the NIP-42 auth handshake (Finding #8). +/// +/// `do_connect` buffers any non-AUTH/non-OK messages it receives while waiting +/// for the challenge and OK. Those messages would otherwise be silently +/// discarded. We replay them through the normal handler here. +#[allow(clippy::too_many_arguments)] +/// Returns `false` if any buffered message signals the connection should be dropped. +async fn process_handshake_buffer( + ws: &mut WsStream, + buffer: std::collections::VecDeque, + event_tx: &mpsc::Sender>, + state: &mut BgState, + keys: &Keys, + relay_url: &str, + api_token: Option<&str>, + agent_pubkey_hex: &str, +) -> bool { + if buffer.is_empty() { + return true; + } + debug!("processing {} buffered handshake message(s)", buffer.len()); + for relay_msg in buffer { + // Re-encode to text so we can reuse handle_ws_message. + // This is slightly wasteful but keeps the handler as the single + // source of truth for message dispatch. + let text = match &relay_msg { + RelayMessage::Event { + subscription_id, + event, + } => serde_json::to_string(&json!(["EVENT", subscription_id, event])).ok(), + RelayMessage::Eose { subscription_id } => { + serde_json::to_string(&json!(["EOSE", subscription_id])).ok() + } + RelayMessage::Notice { message } => { + serde_json::to_string(&json!(["NOTICE", message])).ok() + } + RelayMessage::Closed { + subscription_id, + message, + } => serde_json::to_string(&json!(["CLOSED", subscription_id, message])).ok(), + RelayMessage::Ok { + event_id, + accepted, + message, + } => serde_json::to_string(&json!(["OK", event_id, accepted, message])).ok(), + // AUTH in the buffer is stale — skip it. + RelayMessage::Auth { .. } => None, + }; + if let Some(text) = text { + let should_continue = handle_ws_message( + Message::Text(text.into()), + ws, + event_tx, + state, + keys, + relay_url, + api_token, + agent_pubkey_hex, + ) + .await; + if !should_continue { + return false; + } + } + } + true +} + /// 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. +/// +/// Returns `true` if ALL subscriptions were sent successfully. Returns `false` +/// if any send failed — the caller should treat this as a failed reconnect +/// and retry, because a "connected" socket with missing subscriptions causes +/// silent event loss. async fn resubscribe_after_reconnect( ws: &mut WsStream, state: &mut BgState, agent_pubkey_hex: &str, -) { +) -> bool { + let mut all_ok = true; let channels: Vec = state.active_subscriptions.keys().copied().collect(); if !channels.is_empty() { info!( @@ -936,24 +1699,33 @@ async fn resubscribe_after_reconnect( 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(); + // Fall back to per-channel subscribe_since (not startup_watermark) + // so channels joined after startup don't replay stale history. + let subscribe_ts = state.subscribe_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, + (None, None) => subscribe_ts.or(state.startup_watermark), + }; + let filter = match state.active_filters.get(&channel_id).cloned() { + Some(f) => f, + None => { + // Fail closed: missing filter state means the subscription + // intent is inconsistent. Skip rather than resubscribe with + // a permissive wildcard that would widen the subscription. + warn!("missing filter for channel {channel_id} — skipping resubscribe (fail-closed)"); + all_ok = false; + continue; + } }; - 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); + } else { + warn!("failed to resubscribe channel {channel_id} after reconnect"); + all_ok = false; } } } @@ -963,34 +1735,111 @@ async fn resubscribe_after_reconnect( (Some(d), Some(l)) => Some(d.min(l)), (Some(d), None) => Some(d), (None, Some(l)) => Some(l), - (None, None) => None, + (None, None) => state.startup_watermark, }; let sent = send_membership_subscribe(ws, agent_pubkey_hex, replay_since).await; if sent { state.membership_dropped_since = None; + } else { + warn!("failed to resubscribe membership after reconnect"); + all_ok = false; } } + + all_ok } -/// Attempt autonomous reconnect on socket loss — 3 attempts with 1s→2s→4s backoff. +/// Attempt autonomous reconnect on socket loss. /// -/// 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`). +/// Finding #42: 5 attempts with 1s→2s→4s→8s→16s backoff (was 3 attempts). +/// Finding #27: ±20% jitter on each sleep. +/// Finding #8: process handshake buffer on success. +/// +/// Outcome of an autonomous reconnect attempt. +enum ReconnectOutcome { + /// Reconnected and resubscribed successfully. + Ok, + /// All attempts exhausted — caller should fall back to wait_for_reconnect. + Failed, + /// A Shutdown command was received during backoff — caller must return immediately. + Shutdown, +} + +/// Drain all pending commands after a successful reconnect. +/// +/// Processes queued commands that arrived while reconnecting. Reconnect +/// commands are silently dropped (already reconnected). Shutdown causes an +/// immediate close-frame + return of `ReconnectOutcome::Shutdown`. All other +/// commands are executed on the live socket via [`execute_connected_command`]. +/// If any send fails, remaining commands are recorded as intent via +/// [`apply_command_to_state`] and the drain continues (the caller's next +/// read/ping will detect the dead socket). +async fn drain_post_reconnect( + ws: &mut WsStream, + cmd_rx: &mut mpsc::Receiver, + state: &mut BgState, + agent_pubkey_hex: &str, +) -> ReconnectOutcome { + let mut send_failed = false; + while let Ok(cmd) = cmd_rx.try_recv() { + if send_failed { + match cmd { + RelayCommand::Shutdown => { + let _ = ws_send_timeout(ws, Message::Close(None), WS_SEND_TIMEOUT_SECS).await; + return ReconnectOutcome::Shutdown; + } + RelayCommand::Reconnect => {} + cmd => apply_command_to_state(state, cmd), + } + continue; + } + match cmd { + RelayCommand::Reconnect => { + debug!("drained stale Reconnect after reconnect"); + } + RelayCommand::Shutdown => { + debug!("shutdown received during post-reconnect drain"); + let _ = ws_send_timeout(ws, Message::Close(None), WS_SEND_TIMEOUT_SECS).await; + return ReconnectOutcome::Shutdown; + } + cmd => { + let ok = execute_connected_command(ws, state, agent_pubkey_hex, cmd).await; + if !ok { + warn!("send failed during post-reconnect drain — recording remaining commands as intent"); + send_failed = true; + } + } + } + } + ReconnectOutcome::Ok +} + +/// Attempt autonomous reconnect on socket loss. +/// +/// Returns [`ReconnectOutcome::Ok`] on success, [`ReconnectOutcome::Failed`] +/// if all attempts are exhausted, or [`ReconnectOutcome::Shutdown`] if a +/// Shutdown command was received during backoff sleep. Callers MUST check +/// for `Shutdown` and return immediately — do NOT fall through to +/// `wait_for_reconnect`, which would loop forever since the Shutdown command +/// was already consumed. #[allow(clippy::too_many_arguments)] async fn try_autonomous_reconnect( ws: &mut WsStream, + cmd_rx: &mut mpsc::Receiver, state: &mut BgState, keys: &Keys, relay_url: &str, api_token: Option<&str>, agent_pubkey_hex: &str, -) -> bool { + event_tx: &mpsc::Sender>, +) -> ReconnectOutcome { + // Finding #42: 5 attempts, up to 16s base backoff. let backoffs = [ Duration::from_secs(1), Duration::from_secs(2), Duration::from_secs(4), + Duration::from_secs(8), + Duration::from_secs(16), ]; for (attempt, delay) in backoffs.iter().enumerate() { @@ -1000,32 +1849,70 @@ async fn try_autonomous_reconnect( backoffs.len() ); match do_connect(relay_url, keys, api_token).await { - Ok((new_ws, _buffer)) => { + Ok((new_ws, handshake_buffer)) => { *ws = new_ws; info!("autonomous reconnect succeeded (attempt {})", attempt + 1); - resubscribe_after_reconnect(ws, state, agent_pubkey_hex).await; - return true; + // Finding #8: process buffered messages from the handshake. + let handshake_ok = process_handshake_buffer( + ws, + handshake_buffer, + event_tx, + state, + keys, + relay_url, + api_token, + agent_pubkey_hex, + ) + .await; + if !handshake_ok { + warn!( + "handshake buffer drop signal after autonomous reconnect (attempt {})", + attempt + 1 + ); + // Fall through to backoff sleep instead of returning immediately. + // Returning false here would skip remaining attempts; continuing + // without sleep would drive a tight reconnect storm. + } else if resubscribe_after_reconnect(ws, state, agent_pubkey_hex).await { + return ReconnectOutcome::Ok; + } else { + warn!("resubscribe failed after autonomous reconnect — treating as failed attempt"); + // Fall through to backoff sleep and retry. + } } 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() + warn!("autonomous reconnect attempt {} failed: {e}", attempt + 1); + } + } + + // Backoff sleep between attempts (shared by handshake-drop and connect-error). + // Skip sleep on the final attempt — we'll fall through to the caller. + // Use select! so Shutdown commands are honoured during sleep. + if attempt + 1 < backoffs.len() { + let jittered = jittered_duration(*delay); + tracing::info!( + "retrying autonomous reconnect in {:.1}s", + jittered.as_secs_f64() + ); + // Deadline-based sleep: commands processed during the wait don't + // reset the timer (prevents PublishEvent traffic from collapsing backoff). + let deadline = tokio::time::Instant::now() + jittered; + let sleep = tokio::time::sleep_until(deadline); + tokio::pin!(sleep); + loop { + tokio::select! { + _ = &mut sleep => break, + cmd = cmd_rx.recv() => { + match cmd { + Some(RelayCommand::Shutdown) | None => return ReconnectOutcome::Shutdown, + Some(cmd) => apply_command_to_state(state, cmd), + } } - ); - // 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 + ReconnectOutcome::Failed } /// Attempt reconnection with exponential backoff. Resubscribes all active @@ -1045,55 +1932,97 @@ async fn wait_for_reconnect( relay_url: &str, api_token: Option<&str>, agent_pubkey_hex: &str, + event_tx: &mpsc::Sender>, skip_drain: bool, -) { +) -> ReconnectOutcome { if !skip_drain { // Drain commands until we get Reconnect (or Shutdown). + // Other commands update state so reconnect reflects latest intent. loop { match cmd_rx.recv().await { Some(RelayCommand::Reconnect) => break, - Some(RelayCommand::Shutdown) | None => return, - // Apply Subscribe/Unsubscribe to state so reconnect reflects - // the latest caller intent (not just pre-disconnect state). - Some(RelayCommand::Subscribe { channel_id, filter }) => { - state - .active_subscriptions - .insert(channel_id, channel_sub_id(channel_id)); - state.active_filters.insert(channel_id, filter); - } - Some(RelayCommand::Unsubscribe { channel_id }) => { - state.active_subscriptions.remove(&channel_id); - state.active_filters.remove(&channel_id); - } - Some(RelayCommand::SubscribeMembership) => { - state.membership_sub_active = true; - } - // Ephemeral events are meaningless while disconnected — drop silently. - Some(RelayCommand::PublishEvent { .. }) => {} + Some(RelayCommand::Shutdown) | None => return ReconnectOutcome::Shutdown, + Some(cmd) => apply_command_to_state(state, cmd), } } } - // Attempt reconnection with backoff. + // Finding #42: 6 attempts with backoff up to 32s + jitter (Finding #27). + // Finding #27: use tokio::select! so shutdown is honoured during sleep. + let backoffs = [ + Duration::from_secs(1), + Duration::from_secs(2), + Duration::from_secs(4), + Duration::from_secs(8), + Duration::from_secs(16), + Duration::from_secs(32), + ]; + let mut attempt = 0usize; let mut delay = Duration::from_secs(1); loop { info!("attempting relay reconnect to {relay_url}…"); match do_connect(relay_url, keys, api_token).await { - Ok((new_ws, _buffer)) => { + Ok((new_ws, handshake_buffer)) => { *ws = new_ws; info!("relay reconnected to {relay_url}"); - resubscribe_after_reconnect(ws, state, agent_pubkey_hex).await; - return; + // Finding #8: process buffered messages from the handshake. + let handshake_ok = process_handshake_buffer( + ws, + handshake_buffer, + event_tx, + state, + keys, + relay_url, + api_token, + agent_pubkey_hex, + ) + .await; + if !handshake_ok { + warn!("handshake buffer contained a drop signal after reconnect — will retry with backoff"); + // Fall through to the backoff sleep below instead of + // tight-looping. A relay that consistently fails the + // handshake would otherwise drive a reconnect storm. + } else if resubscribe_after_reconnect(ws, state, agent_pubkey_hex).await { + // Drain any commands that arrived during the final + // do_connect() + resubscribe (which don't poll cmd_rx). + return drain_post_reconnect(ws, cmd_rx, state, agent_pubkey_hex).await; + } else { + warn!("resubscribe failed after reconnect — will retry with backoff"); + // Fall through to backoff sleep. + } } Err(e) => { - warn!( - "relay reconnect failed: {e} — retrying in {}s", - delay.as_secs() - ); - tokio::time::sleep(delay).await; - delay = (delay * 2).min(Duration::from_secs(60)); + warn!("relay reconnect failed: {e}"); } } + + // Backoff sleep — shared by both handshake-drop and connect-error paths. + // Uses a deadline so commands processed during the wait don't reset + // the timer. Without this, periodic PublishEvent traffic (typing + // refresh every 3s) would collapse the jittered backoff into a + // reconnect storm. + let jittered = jittered_duration(delay); + warn!("retrying reconnect in {:.1}s", jittered.as_secs_f64()); + let deadline = tokio::time::Instant::now() + jittered; + let sleep = tokio::time::sleep_until(deadline); + tokio::pin!(sleep); + loop { + tokio::select! { + _ = &mut sleep => break, + cmd = cmd_rx.recv() => { + match cmd { + Some(RelayCommand::Shutdown) | None => return ReconnectOutcome::Shutdown, + Some(cmd) => apply_command_to_state(state, cmd), + } + } + } + } + attempt += 1; + delay = if attempt < backoffs.len() { + backoffs[attempt] + } else { + Duration::from_secs(60) + }; } } @@ -1145,23 +2074,25 @@ async fn send_subscribe( let req = json!(["REQ", sub_id, Value::Object(req_filter)]); match serde_json::to_string(&req) { - Ok(text) => match ws.send(Message::Text(text.into())).await { - Ok(()) => { - debug!( - "subscribed to channel {channel_id}{}", - if since.is_some() { - " (with since filter)" - } else { - " (since=now)" - } - ); - true + Ok(text) => { + match ws_send_timeout(ws, Message::Text(text.into()), WS_SEND_TIMEOUT_SECS).await { + Ok(()) => { + debug!( + "subscribed to channel {channel_id}{}", + if since.is_some() { + " (with since filter)" + } else { + " (since=now)" + } + ); + true + } + Err(e) => { + warn!("failed to send REQ for channel {channel_id}: {e}"); + false + } } - 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 @@ -1197,16 +2128,18 @@ async fn send_membership_subscribe( let req = json!(["REQ", MEMBERSHIP_NOTIF_SUB_ID, Value::Object(req_filter)]); match serde_json::to_string(&req) { - Ok(text) => match ws.send(Message::Text(text.into())).await { - Ok(()) => { - debug!("subscribed to membership notifications (since={since_ts})"); - true + Ok(text) => { + match ws_send_timeout(ws, Message::Text(text.into()), WS_SEND_TIMEOUT_SECS).await { + Ok(()) => { + debug!("subscribed to membership notifications (since={since_ts})"); + true + } + Err(e) => { + warn!("failed to send membership notification REQ: {e}"); + false + } } - Err(e) => { - warn!("failed to send membership notification REQ: {e}"); - false - } - }, + } Err(e) => { warn!("failed to serialize membership notification REQ: {e}"); false @@ -1214,6 +2147,34 @@ async fn send_membership_subscribe( } } +/// Send a WebSocket message with a hard timeout. +/// +/// All `ws.send()` calls go through here so a stalled TCP socket can't wedge +/// the background task. On timeout the caller should break out of the loop to +/// trigger reconnect. +async fn ws_send_timeout( + ws: &mut WsStream, + msg: Message, + timeout_secs: u64, +) -> Result<(), RelayError> { + tokio::time::timeout(Duration::from_secs(timeout_secs), ws.send(msg)) + .await + .map_err(|_| RelayError::Timeout)? + .map_err(|e| RelayError::WebSocket(Box::new(e))) +} + +/// Add ±20% jitter to a backoff duration using the nanosecond sub-second +/// component of the system clock as a cheap entropy source (no `rand` dep). +fn jittered_duration(base: Duration) -> Duration { + let nanos = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .subsec_nanos(); + // factor ∈ [0.8, 1.2) + let factor = 0.8 + (nanos as f64 / u32::MAX as f64) * 0.4; + base.mul_f64(factor) +} + /// Extract a channel UUID from the h tag of a Nostr event. fn extract_h_tag_uuid(event: &nostr::Event) -> Option { event.tags.iter().find_map(|tag| { @@ -1252,9 +2213,7 @@ async fn send_auth_response( }; let auth_msg = serde_json::to_string(&json!(["AUTH", auth_event]))?; - ws.send(Message::Text(auth_msg.into())) - .await - .map_err(|e| RelayError::WebSocket(Box::new(e)))?; + ws_send_timeout(ws, Message::Text(auth_msg.into()), WS_SEND_TIMEOUT_SECS).await?; debug!("sent AUTH response for challenge"); Ok(()) } @@ -1481,9 +2440,9 @@ async fn wait_for_auth_challenge( } } Message::Ping(data) => { - ws.send(Message::Pong(data)) + ws_send_timeout(ws, Message::Pong(data), WS_SEND_TIMEOUT_SECS) .await - .map_err(|e| RelayError::WebSocket(Box::new(e)))?; + .map_err(|_| RelayError::Timeout)?; } Message::Close(_) => return Err(RelayError::ConnectionClosed), _ => {} @@ -1559,9 +2518,9 @@ async fn wait_for_any_ok( } } Message::Ping(data) => { - ws.send(Message::Pong(data)) + ws_send_timeout(ws, Message::Pong(data), WS_SEND_TIMEOUT_SECS) .await - .map_err(|e| RelayError::WebSocket(Box::new(e)))?; + .map_err(|_| RelayError::Timeout)?; } Message::Close(_) => return Err(RelayError::ConnectionClosed), _ => {} @@ -1941,40 +2900,93 @@ mod tests { assert_eq!(state.last_seen.get(&ch2).copied(), Some(2_000_000)); } + /// Two-generation dedup: no amnesia window on rotation. + /// + /// The old implementation cleared the entire set at 12_001, creating a gap + /// where all previously-seen IDs became eligible again. The new TwoGenDedup + /// rotates at SEEN_ID_LIMIT/2 = 6_000, keeping the previous generation so + /// IDs from both generations are still recognised as duplicates. + #[test] + fn bg_state_two_gen_dedup_no_amnesia_on_rotation() { + let mut dedup = TwoGenDedup::new(SEEN_ID_LIMIT); + + // Fill current generation to the rotation threshold (limit/2 = 6_000). + // After inserting the 6_000th item, current rotates into previous. + let mut ids: Vec = Vec::new(); + for i in 0u64..6_000 { + let id = format!("{:0>64x}", i); + ids.push(id.clone()); + dedup.insert(id); + } + + // All 6_000 IDs were rotated into `previous`. `current` is now empty. + // They must still be recognised as duplicates. + for id in &ids { + assert!( + dedup.contains(id), + "rotated ID {id} should still be a duplicate" + ); + } + + // New IDs after rotation must be accepted. + let new_id = format!("{:0>64x}", 99_999u64); + assert!( + dedup.insert(new_id.clone()), + "new ID after rotation should be accepted" + ); + assert!( + dedup.contains(&new_id), + "new ID should be found after insert" + ); + } + + #[test] + fn bg_state_two_gen_dedup_duplicate_rejected_across_generations() { + let mut dedup = TwoGenDedup::new(12); + // limit/2 = 6, so rotation happens at 6 inserts. + for i in 0u64..6 { + dedup.insert(format!("id-{i}")); + } + // id-0 is now in `previous` (rotated). Inserting it again must return false. + assert!( + !dedup.insert("id-0".to_string()), + "cross-generation duplicate must be rejected" + ); + } + #[test] fn bg_state_seen_ids_cleared_at_limit() { + // Compatibility test: BgState.record_event still deduplicates correctly + // after the TwoGenDedup rotation threshold is crossed. let mut state = BgState::new(); let channel_id = Uuid::new_v4(); - // Pre-populate seen_ids to just below the threshold (12_000 - 1). - // Use synthetic hex strings — we're testing the clear logic, not signing. - for i in 0u64..11_999 { + // Insert SEEN_ID_LIMIT/2 synthetic IDs to trigger the first rotation. + for i in 0u64..(SEEN_ID_LIMIT as u64 / 2) { state.seen_ids.insert(format!("{:0>64x}", i)); } - assert_eq!(state.seen_ids.len(), 11_999); - // Now insert two real events. The first will bring us to 12_000 (no - // clear yet), the second will push us to 12_001 and trigger the clear. - let keys = nostr::Keys::generate(); - let event1 = make_test_event(&keys, 1_000_000); - let keys2 = nostr::Keys::generate(); - let event2 = make_test_event(&keys2, 1_000_001); - - // First insert: 12_000 entries — no clear triggered yet. - state.record_event(channel_id, &event1); - assert_eq!( - state.seen_ids.len(), - 12_000, - "should be at 12_000 before clear" + // The first generation has been rotated into `previous`. All IDs are + // still present across the two generations — no amnesia window. + assert!( + state + .seen_ids + .contains("0000000000000000000000000000000000000000000000000000000000000000"), + "first ID should still be recognised after rotation" ); - // Second insert: 12_001 entries — triggers clear, then re-inserts. - state.record_event(channel_id, &event2); - // After clear + re-insert of event2, seen_ids should be very small. + // A new real event should be accepted (not a duplicate). + let keys = nostr::Keys::generate(); + let event = make_test_event(&keys, 1_000_000); assert!( - state.seen_ids.len() < 12_000, - "seen_ids should have been cleared, got {}", - state.seen_ids.len() + state.record_event(channel_id, &event), + "new event after rotation should be accepted" + ); + + // The same event must be rejected as a duplicate. + assert!( + !state.record_event(channel_id, &event), + "duplicate event after rotation should be rejected" ); } @@ -2072,4 +3084,64 @@ mod tests { "channel_dropped_since should be cleared after resubscribe" ); } + + // ── Membership dedup regression tests (M4) ─────────────────────────── + + /// Membership dedup must NOT contaminate per-channel `last_seen`. + /// Using `record_event()` for membership notifications would update + /// `last_seen[channel_uuid]`, causing channel resubscribe to use a + /// membership timestamp as the `since` filter — skipping channel events. + /// The fix uses `seen_ids.insert()` directly. + #[test] + fn membership_dedup_does_not_touch_last_seen() { + let mut state = BgState::new(); + let channel_id = Uuid::new_v4(); + let keys = nostr::Keys::generate(); + + // Simulate: a channel event sets last_seen to 1000. + let event1 = make_test_event(&keys, 1_000); + assert!(state.record_event(channel_id, &event1)); + assert_eq!(state.last_seen.get(&channel_id).copied(), Some(1_000)); + + // Simulate: a membership notification for the same channel at ts=2000. + // This should go through seen_ids only, NOT update last_seen. + let membership_event = make_test_event(&keys, 2_000); + let membership_id = membership_event.id.to_hex(); + assert!( + state.seen_ids.insert(membership_id), + "membership event should be accepted by dedup" + ); + // last_seen must still be 1000, not 2000. + assert_eq!( + state.last_seen.get(&channel_id).copied(), + Some(1_000), + "membership dedup must not contaminate last_seen" + ); + } + + /// On membership backpressure (TrySendError::Full), the dedup ID must + /// be removed from seen_ids so reconnect replay can re-deliver the event. + /// Without this, a dropped membership notification would be permanently + /// rejected as a duplicate on replay. + #[test] + fn membership_backpressure_removes_dedup_id() { + let mut state = BgState::new(); + let keys = nostr::Keys::generate(); + + let event = make_test_event(&keys, 1_000); + let event_id_hex = event.id.to_hex(); + + // Insert into dedup (simulating the pre-try_send path). + assert!(state.seen_ids.insert(event_id_hex.clone())); + assert!(state.seen_ids.contains(&event_id_hex)); + + // Simulate backpressure: remove the ID (matching the production code). + state.seen_ids.remove(&event_id_hex); + + // The ID should now be accepted again on replay. + assert!( + state.seen_ids.insert(event_id_hex), + "after backpressure removal, replay must be accepted" + ); + } }