diff --git a/crates/buzz-acp/src/acp.rs b/crates/buzz-acp/src/acp.rs index 9eb668cbc2..faf17d6da0 100644 --- a/crates/buzz-acp/src/acp.rs +++ b/crates/buzz-acp/src/acp.rs @@ -112,6 +112,112 @@ pub enum AcpError { AgentError { code: i64, message: String }, } +/// Verification state for the directly spawned ACP adapter process. +#[derive(Clone, Debug, Eq, PartialEq)] +pub enum DirectChildShutdownStatus { + /// The direct child exited and its status was reaped. + Reaped, + /// Waiting for the direct child returned an operating-system error. + WaitFailed { + /// Portable error category reported by the operating system. + kind: std::io::ErrorKind, + /// Operating-system error text. This never includes ACP wire content. + message: String, + }, + /// The direct child did not produce an exit status within the bound. + TimedOut { + /// Maximum time spent waiting for this verification attempt. + timeout: std::time::Duration, + }, +} + +/// Verification state for the process group created for an ACP adapter. +#[derive(Clone, Debug, Eq, PartialEq)] +pub enum ProcessGroupShutdownStatus { + /// The original process group no longer exists (`killpg(..., 0)` returned `ESRCH`). + Absent, + /// The original process group still exists after the bounded probe window. + Present, + /// The operating system could not determine whether the group exists. + ProbeFailed { + /// Platform error number returned by the process-group probe. + errno: i32, + }, + /// Windows direct-child containment completed and the child was reaped. + /// + /// Windows does not yet place adapters in a Job Object, so this is an + /// explicit bounded platform boundary rather than descendant-absence proof. + #[cfg_attr(not(windows), allow(dead_code))] + WindowsDirectChildBoundary, + /// Process-group containment is unavailable on this platform. + /// + /// Direct-child termination alone is insufficient ownership proof. + #[cfg_attr(unix, allow(dead_code))] + Unavailable, +} + +/// Typed evidence that an ACP adapter shutdown could not be fully verified. +#[derive(Debug, thiserror::Error)] +#[error( + "agent shutdown unverified for pid {process_id}: direct_child={direct_child:?}, \ + process_group={process_group:?}" +)] +pub struct ShutdownError { + /// Process ID assigned to the adapter at spawn time. + pub process_id: u32, + /// Final direct-child termination evidence. + pub direct_child: DirectChildShutdownStatus, + /// Final process-group termination evidence. + pub process_group: ProcessGroupShutdownStatus, +} + +fn classify_shutdown( + process_id: u32, + direct_child: DirectChildShutdownStatus, + process_group: ProcessGroupShutdownStatus, +) -> Result<(), ShutdownError> { + classify_shutdown_with_windows_boundary(process_id, direct_child, process_group, cfg!(windows)) +} + +fn classify_shutdown_with_windows_boundary( + process_id: u32, + direct_child: DirectChildShutdownStatus, + process_group: ProcessGroupShutdownStatus, + windows_direct_child_boundary_supported: bool, +) -> Result<(), ShutdownError> { + let child_verified = direct_child == DirectChildShutdownStatus::Reaped; + let group_verified = match process_group { + ProcessGroupShutdownStatus::Absent => true, + ProcessGroupShutdownStatus::WindowsDirectChildBoundary => { + windows_direct_child_boundary_supported + } + ProcessGroupShutdownStatus::Unavailable + | ProcessGroupShutdownStatus::Present + | ProcessGroupShutdownStatus::ProbeFailed { .. } => false, + }; + if child_verified && group_verified { + Ok(()) + } else { + Err(ShutdownError { + process_id, + direct_child, + process_group, + }) + } +} + +impl AcpError { + /// Whether the adapter process must be replaced before another request. + /// + /// Every non-agent error means process ownership, stdio framing, or request + /// synchronization is uncertain. A JSON-RPC `AgentError` is the sole + /// application-class error: the adapter positively framed a response and + /// remains reusable unless a higher-level lifecycle boundary says otherwise. + pub fn requires_process_replacement(&self) -> bool { + !matches!(self, Self::AgentError { .. }) + } +} + /// Build an [`AcpError::AgentError`] from a JSON-RPC error object, /// preserving the numeric code. When the `message` field is missing or /// non-string, fall back to the full JSON object so provider-specific @@ -143,6 +249,23 @@ fn build_initialize_params() -> serde_json::Value { pub struct AcpClient { /// The agent child process (kept alive to prevent zombie). child: Child, + /// Process ID assigned to the direct child at spawn time. + /// + /// Stored independently because [`Child::id`] becomes `None` after the + /// process has been polled to completion, but shutdown evidence must still + /// identify the original process. + spawn_process_id: u32, + /// Original process-group ID created at spawn time on Unix. + /// + /// `None` on platforms without Unix process-group containment. Descendants + /// that deliberately escape via `setsid(2)` create a new session/group and + /// are outside what this stored group ID—and therefore [`Self::shutdown`]— + /// can prove. + process_group_id: Option, + /// Whether a prior [`Self::shutdown`] call verified every platform-supported + /// termination boundary. Prevents repeated shutdown or `Drop` from signaling + /// a process-group ID that the operating system may later reuse. + shutdown_verified: bool, /// Write end of the agent's stdin pipe. stdin: ChildStdin, /// Framed reader over the agent's stdout pipe (line-oriented, bounded). @@ -215,6 +338,13 @@ pub struct AcpClient { /// deltas. Both goose and buzz-agent emit this notification; goose gates /// on client capability advertisement, buzz-agent emits unconditionally. goose_usage: UsageTracker, + /// Whether the initialize response advertised ACP `session/close`. + /// + /// Session close is optional. Callers use this bit to choose between an + /// explicit close and a deliberate process recycle; guessing support and + /// treating method-not-found as a crash would trip the circuit breaker for + /// otherwise healthy legacy adapters. + supports_session_close: bool, } /// Recursively merge `overlay` into `base`, with `overlay` winning on scalar/shape @@ -415,32 +545,84 @@ fn build_client_capabilities() -> serde_json::Value { } impl AcpClient { - /// Kill the agent subprocess and wait for it to exit (no zombies). + /// Kill the agent process group and verify bounded cleanup. /// - /// `Drop` only calls `start_kill()` (sends SIGKILL but doesn't reap). - /// Call this when you need guaranteed cleanup — e.g., in `run_models` - /// before process exit. - pub async fn shutdown(&mut self) { - // Kill the entire process group when possible. The child was spawned - // with process_group(0), so its PID == its PGID. Killing the group - // ensures subprocesses (MCP servers, tool processes) are cleaned up - // rather than orphaned to init. - // - // Falls back to start_kill() (direct child only) on non-Unix or if - // the child has been polled to completion (id() returns None). - match self.child.id() { - Some(pid) if kill_process_group(pid) => {} - _ => { + /// On Unix, success proves both that the direct child was reaped and that + /// the original spawn-time process group is absent. A direct-child kill is + /// attempted when group signaling fails and again when the first bounded + /// wait cannot verify exit. The group ID is retained independently of + /// [`Child::id`], so verification still runs after Tokio clears the live ID. + /// + /// On Windows, a reaped direct child is the current explicit containment + /// boundary until adapter processes are placed in Job Objects. Other + /// non-Unix platforms fail closed because descendant containment remains + /// unavailable. + /// + /// Descendants that deliberately escape the original group with `setsid(2)` + /// are outside this API's containment and verification boundary. + pub async fn shutdown(&mut self) -> Result<(), ShutdownError> { + const SHUTDOWN_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5); + if self.shutdown_verified { + return Ok(()); + } + + #[cfg(unix)] + { + let process_group_id = self.process_group_id; + + // Signal the group only while Tokio still proves that the original + // direct child owns the spawn PID. Once Child::id() is cleared, the + // numeric PID/PGID may be reused; it remains safe to probe, never to + // signal. Direct-child fallback stays bounded by the Child handle. + if let Some(id) = process_group_signal_target( + self.spawn_process_id, + process_group_id, + self.child.id(), + ) { + if signal_process_group(id).is_err() { + let _ = self.child.start_kill(); + } + } else if self.child.id().is_some() { let _ = self.child.start_kill(); } + + let mut direct_child = wait_for_direct_child(&mut self.child, SHUTDOWN_TIMEOUT).await; + if direct_child != DirectChildShutdownStatus::Reaped { + // A group signal that did not yield a reaped direct child is + // not enough. Retry through Tokio's direct-child handle. + let _ = self.child.start_kill(); + direct_child = wait_for_direct_child(&mut self.child, SHUTDOWN_TIMEOUT).await; + } + + let process_group = match process_group_id { + Some(id) => wait_for_process_group_absence(id, SHUTDOWN_TIMEOUT).await, + None => ProcessGroupShutdownStatus::Absent, + }; + if process_group == ProcessGroupShutdownStatus::Absent { + // An absent ID can be reused by the OS. Never signal it again. + self.process_group_id = None; + } + let result = classify_shutdown(self.spawn_process_id, direct_child, process_group); + if result.is_ok() { + self.shutdown_verified = true; + } + result } - // 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"), + + #[cfg(not(unix))] + { + // No portable descendant process-group primitive is available. + let _ = self.child.start_kill(); + let direct_child = wait_for_direct_child(&mut self.child, SHUTDOWN_TIMEOUT).await; + #[cfg(windows)] + let containment = ProcessGroupShutdownStatus::WindowsDirectChildBoundary; + #[cfg(all(not(unix), not(windows)))] + let containment = ProcessGroupShutdownStatus::Unavailable; + let result = classify_shutdown(self.spawn_process_id, direct_child, containment); + if result.is_ok() { + self.shutdown_verified = true; + } + result } } @@ -538,6 +720,13 @@ impl AcpClient { configure_no_window(&mut cmd); let mut child = cmd.spawn()?; + let spawn_process_id = child + .id() + .ok_or_else(|| AcpError::Protocol("spawned agent has no process id".into()))?; + #[cfg(unix)] + let process_group_id = Some(spawn_process_id); + #[cfg(not(unix))] + let process_group_id = None; let stdin = child .stdin @@ -550,6 +739,9 @@ impl AcpClient { Ok(Self { child, + spawn_process_id, + process_group_id, + shutdown_verified: false, stdin, reader: FramedRead::new(stdout, LinesCodec::new_with_max_length(MAX_LINE_SIZE)), next_id: 0, @@ -564,6 +756,7 @@ impl AcpClient { steering_supported: false, steer_rx: None, goose_usage: UsageTracker::default(), + supports_session_close: false, }) } @@ -600,6 +793,37 @@ impl AcpClient { } } + /// Parse one nonempty ACP NDJSON frame without retaining malformed content. + /// + /// A malformed frame makes stdout synchronization uncertain, so callers + /// must fail closed instead of skipping ahead to a later response. Error + /// evidence intentionally includes only the byte length and parser error; + /// agent output may contain prompt text, credentials, or tool results. + fn parse_ndjson_frame(&self, line: &str) -> Result { + match serde_json::from_str(line) { + Ok(msg) => { + tracing::debug!(target: "acp::wire", "← {line}"); + Ok(msg) + } + Err(error) => { + self.observe( + "acp_parse_error", + serde_json::json!({ + "lineLength": line.len(), + "error": error.to_string(), + }), + ); + tracing::warn!( + target: "acp::wire", + line_length = line.len(), + error = %error, + "failed to parse agent stdout as JSON" + ); + Err(AcpError::Json(error)) + } + } + } + /// Send the `initialize` request and return the agent's response result value. /// /// Must be called exactly once, before any other ACP method. @@ -618,10 +842,18 @@ impl AcpClient { .pointer("/_meta/steering/supported") .and_then(|v| v.as_bool()) .unwrap_or(false); + self.supports_session_close = result + .pointer("/agentCapabilities/sessionCapabilities/close") + .is_some_and(serde_json::Value::is_object); tracing::debug!(target: "acp::init", "initialize response: {result}"); Ok(result) } + /// Whether the adapter explicitly advertised ACP `session/close`. + pub fn supports_session_close(&self) -> bool { + self.supports_session_close + } + /// Send the ACP `authenticate` request for an adapter-advertised method. pub async fn authenticate(&mut self, method_id: &str) -> Result { let params = serde_json::json!({ @@ -829,6 +1061,7 @@ impl AcpClient { /// that response. /// /// Note: async because writing to stdin requires async I/O. + #[allow(dead_code)] // Public ACP API; bounded cleanup uses its shared-deadline variant. pub async fn session_cancel(&mut self, session_id: &str) -> Result<(), AcpError> { let params = serde_json::json!({ "sessionId": session_id, @@ -836,6 +1069,28 @@ impl AcpClient { self.send_notification("session/cancel", params).await } + /// Close a session and release its adapter-owned resources. + /// + /// Session retirement is intentionally bounded more tightly than ordinary + /// non-prompt RPCs. A close that cannot complete promptly leaves ownership + /// uncertain, so callers must treat any error as poisoning the adapter + /// rather than deleting the local session ID and reusing the process. + pub async fn session_close(&mut self, session_id: &str) -> Result<(), AcpError> { + const SESSION_CLOSE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5); + let params = serde_json::json!({ + "sessionId": session_id, + }); + match tokio::time::timeout( + SESSION_CLOSE_TIMEOUT, + self.send_request("session/close", params), + ) + .await + { + Ok(result) => result.map(|_| ()), + Err(_) => Err(AcpError::Timeout(SESSION_CLOSE_TIMEOUT)), + } + } + /// Returns `true` if a `session/prompt` request is currently in flight. pub fn has_in_flight_prompt(&self) -> bool { self.last_prompt_id.is_some() @@ -1005,7 +1260,7 @@ impl AcpClient { if let Some(perm_id) = self.pending_permission_id.clone() { if !self.permission_responded { let response = permission_response_cancelled(&perm_id); - self.write_ndjson(&response).await?; + self.write_ndjson_until(&response, hard_deadline).await?; tracing::debug!( target: "acp::cancel", "responded cancelled to pending permission id={perm_id}" @@ -1016,7 +1271,12 @@ impl AcpClient { } // Step 2: send session/cancel notification (no id) - self.session_cancel(session_id).await?; + let cancel = serde_json::json!({ + "jsonrpc": "2.0", + "method": "session/cancel", + "params": {"sessionId": session_id}, + }); + self.write_ndjson_until(&cancel, hard_deadline).await?; tracing::info!(target: "acp::cancel", "sent session/cancel for {session_id}"); // Use a fixed 30s idle timeout during cleanup — the cancel notification // needs time to propagate and the agent may go silent while winding down. @@ -1058,6 +1318,20 @@ impl AcpClient { Ok(()) } + /// Write within an existing operation deadline instead of starting a new + /// 30-second write budget that can outlive its caller's cleanup grace. + async fn write_ndjson_until( + &mut self, + value: &serde_json::Value, + deadline: tokio::time::Instant, + ) -> Result<(), AcpError> { + tokio::time::timeout_at(deadline, self.write_ndjson(value)) + .await + .map_err(|_| AcpError::HardTimeout { + silence: std::time::Duration::ZERO, + })? + } + /// Default timeout for non-prompt RPCs (initialize, session/new, etc.). const REQUEST_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(60); @@ -1138,6 +1412,7 @@ impl AcpClient { /// /// Used for `session/cancel`. The absence of `id` is the JSON-RPC 2.0 /// distinguisher between requests and notifications. + #[allow(dead_code)] // Used by the public notification API above. async fn send_notification( &mut self, method: &str, @@ -1191,26 +1466,7 @@ impl AcpClient { 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) { - Ok(v) => v, - Err(e) => { - self.observe( - "acp_parse_error", - serde_json::json!({ - "line": trimmed, - "error": e.to_string(), - }), - ); - tracing::warn!( - target: "acp::wire", - "failed to parse line as JSON: {e} — skipping" - ); - continue; - } - }; + let msg = self.parse_ndjson_frame(trimmed)?; self.observe("acp_read", msg.clone()); // Check if this is a response to our expected request (has matching id @@ -1221,7 +1477,12 @@ impl AcpClient { if let Some(error) = msg.get("error") { return Err(agent_error_from_json(error)); } - return Ok(msg["result"].clone()); + if let Some(result) = msg.get("result") { + return Ok(result.clone()); + } + return Err(AcpError::Protocol( + "matching JSON-RPC response omitted both result and error".into(), + )); } } @@ -1515,23 +1776,13 @@ impl AcpClient { continue; } - tracing::debug!(target: "acp::wire", "← {trimmed}"); - - let msg: serde_json::Value = match serde_json::from_str(trimmed) { - Ok(v) => v, - Err(e) => { - self.observe( - "acp_parse_error", - serde_json::json!({ - "line": trimmed, - "error": e.to_string(), - }), - ); - tracing::warn!( - target: "acp::wire", - "failed to parse line as JSON: {e} — skipping" - ); - continue; + let msg = match self.parse_ndjson_frame(trimmed) { + Ok(msg) => msg, + Err(error) => { + if let Some((_, _, ack_tx)) = pending_steer.take() { + let _ = ack_tx.send(crate::pool::SteerAck::PromptCompletedNeutral); + } + return Err(error); } }; self.observe("acp_read", msg.clone()); @@ -1564,6 +1815,12 @@ impl AcpClient { crate::pool::SteerAck::Err( crate::pool::SteerError::AgentError { code, message }, ) + } else if msg.get("result").is_none() { + crate::pool::SteerAck::Err( + crate::pool::SteerError::Transport( + "steer response missing result or error".into(), + ), + ) } else { // Success result. Whether it counts as // a delivered steer — and whether the @@ -2174,24 +2431,66 @@ pub fn model_in_catalog( // ─── Drop: kill child process ───────────────────────────────────────────────── +async fn wait_for_direct_child( + child: &mut Child, + timeout: std::time::Duration, +) -> DirectChildShutdownStatus { + match tokio::time::timeout(timeout, child.wait()).await { + Ok(Ok(_)) => DirectChildShutdownStatus::Reaped, + Ok(Err(error)) => DirectChildShutdownStatus::WaitFailed { + kind: error.kind(), + message: error.to_string(), + }, + Err(_) => DirectChildShutdownStatus::TimedOut { timeout }, + } +} + impl Drop for AcpClient { fn drop(&mut self) { // Best-effort SIGKILL + reap. We cannot `await` in Drop (sync context). - // Kill the process group when possible so subprocesses don't leak. + // Signal the stored process group when possible so cleanup does not + // depend on Child::id() remaining available. // Callers SHOULD still call `shutdown().await` for guaranteed reaping. - match self.child.id() { - Some(pid) if kill_process_group(pid) => {} - _ => { + #[cfg(unix)] + if !self.shutdown_verified { + if let Some(id) = process_group_signal_target( + self.spawn_process_id, + self.process_group_id, + self.child.id(), + ) { + if signal_process_group(id).is_err() { + let _ = self.child.start_kill(); + } + } else if self.child.id().is_some() { let _ = self.child.start_kill(); } } + #[cfg(not(unix))] + if !self.shutdown_verified { + 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(); } } -/// Send SIGKILL to an entire process group. Returns `true` if the signal was sent. +/// Return the only process-group ID that is safe to signal. +/// +/// A stored numeric PGID is evidence for later absence probes, but after the +/// live child ID disappears it may name an unrelated, reused process group. +/// Signal only while both spawn-time and live identities still agree. +#[cfg(unix)] +fn process_group_signal_target( + spawn_process_id: u32, + process_group_id: Option, + live_child_id: Option, +) -> Option { + (process_group_id == Some(spawn_process_id) && live_child_id == Some(spawn_process_id)) + .then_some(spawn_process_id) +} + +/// Send SIGKILL to the original process group. /// /// The child is spawned with `process_group(0)`, so its PID equals its PGID. /// Killing the group ensures subprocesses (MCP servers, tool processes) are @@ -2200,19 +2499,58 @@ impl Drop for AcpClient { /// Uses `nix::sys::signal::killpg` — a safe wrapper around the POSIX `killpg` /// syscall — so the crate's `#![deny(unsafe_code)]` policy is preserved. #[cfg(unix)] -fn kill_process_group(pid: u32) -> bool { +fn signal_process_group(process_group_id: u32) -> Result<(), nix::errno::Errno> { + use nix::sys::signal::{killpg, Signal}; + use nix::unistd::Pid; + + killpg(Pid::from_raw(process_group_id as i32), Signal::SIGKILL) +} + +#[cfg(unix)] +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum ProcessGroupProbe { + Absent, + Present, + Failed { errno: i32 }, +} + +/// Probe the original group without sending a signal. +#[cfg(unix)] +fn probe_process_group(process_group_id: u32) -> ProcessGroupProbe { + use nix::errno::Errno; use nix::sys::signal::{killpg, Signal}; use nix::unistd::Pid; - // pid == pgid because the child was spawned with process_group(0). - killpg(Pid::from_raw(pid as i32), Signal::SIGKILL).is_ok() + match killpg(Pid::from_raw(process_group_id as i32), None::) { + Ok(()) => ProcessGroupProbe::Present, + Err(Errno::ESRCH) => ProcessGroupProbe::Absent, + Err(error) => ProcessGroupProbe::Failed { + errno: error as i32, + }, + } } -/// Fallback for non-Unix: process-group kill not available. -/// Returns `false` so the caller falls back to `child.start_kill()`. -#[cfg(not(unix))] -fn kill_process_group(_pid: u32) -> bool { - false +#[cfg(unix)] +async fn wait_for_process_group_absence( + process_group_id: u32, + timeout: std::time::Duration, +) -> ProcessGroupShutdownStatus { + const PROBE_INTERVAL: std::time::Duration = std::time::Duration::from_millis(10); + let deadline = tokio::time::Instant::now() + timeout; + loop { + match probe_process_group(process_group_id) { + ProcessGroupProbe::Absent => return ProcessGroupShutdownStatus::Absent, + ProcessGroupProbe::Failed { errno } => { + return ProcessGroupShutdownStatus::ProbeFailed { errno }; + } + ProcessGroupProbe::Present if tokio::time::Instant::now() >= deadline => { + return ProcessGroupShutdownStatus::Present; + } + ProcessGroupProbe::Present => { + tokio::time::sleep(PROBE_INTERVAL).await; + } + } + } } /// Suppress the console window that Windows otherwise allocates for every @@ -2938,7 +3276,10 @@ mod tests { .await .unwrap_or_else(|| panic!("child produced no output for {var}")) .expect("child stdout was not readable"); - client.shutdown().await; + client + .shutdown() + .await + .expect("env probe shutdown must be verified"); std::fs::remove_dir_all(&dir).expect("remove env probe dir"); observed } @@ -3004,6 +3345,369 @@ mod tests { ); } + #[test] + fn shutdown_classification_accepts_reaped_child_and_absent_group() { + assert!(classify_shutdown( + 41, + DirectChildShutdownStatus::Reaped, + ProcessGroupShutdownStatus::Absent, + ) + .is_ok()); + } + + #[test] + fn shutdown_classification_accepts_reaped_child_at_windows_containment_boundary() { + assert!(classify_shutdown_with_windows_boundary( + 41, + DirectChildShutdownStatus::Reaped, + ProcessGroupShutdownStatus::WindowsDirectChildBoundary, + true, + ) + .is_ok()); + } + + #[test] + fn windows_containment_boundary_still_rejects_unreaped_direct_child() { + let timeout = std::time::Duration::from_secs(5); + let error = classify_shutdown_with_windows_boundary( + 41, + DirectChildShutdownStatus::TimedOut { timeout }, + ProcessGroupShutdownStatus::WindowsDirectChildBoundary, + true, + ) + .expect_err("the Windows boundary requires verified direct-child reaping"); + + assert_eq!( + error.direct_child, + DirectChildShutdownStatus::TimedOut { timeout } + ); + assert_eq!( + error.process_group, + ProcessGroupShutdownStatus::WindowsDirectChildBoundary + ); + } + + #[cfg(not(windows))] + #[test] + fn windows_containment_boundary_is_rejected_off_windows() { + let error = classify_shutdown( + 41, + DirectChildShutdownStatus::Reaped, + ProcessGroupShutdownStatus::WindowsDirectChildBoundary, + ) + .expect_err("the Windows direct-child boundary must not weaken other platforms"); + + assert_eq!( + error.process_group, + ProcessGroupShutdownStatus::WindowsDirectChildBoundary + ); + } + + #[cfg(windows)] + #[test] + fn production_classifier_accepts_reaped_child_at_windows_boundary() { + assert!(classify_shutdown( + 41, + DirectChildShutdownStatus::Reaped, + ProcessGroupShutdownStatus::WindowsDirectChildBoundary, + ) + .is_ok()); + } + + #[cfg(unix)] + #[test] + fn process_group_signal_requires_live_spawn_identity_match() { + assert_eq!( + process_group_signal_target(41, Some(41), Some(41)), + Some(41) + ); + assert_eq!(process_group_signal_target(41, Some(41), None), None); + assert_eq!(process_group_signal_target(41, Some(41), Some(42)), None); + assert_eq!(process_group_signal_target(41, Some(42), Some(41)), None); + } + + #[test] + fn shutdown_classification_rejects_present_process_group() { + let error = classify_shutdown( + 41, + DirectChildShutdownStatus::Reaped, + ProcessGroupShutdownStatus::Present, + ) + .expect_err("a present original process group is not verified shutdown"); + + assert_eq!(error.process_id, 41); + assert_eq!(error.direct_child, DirectChildShutdownStatus::Reaped); + assert_eq!(error.process_group, ProcessGroupShutdownStatus::Present); + } + + #[test] + fn shutdown_classification_rejects_unreaped_child_even_when_group_is_absent() { + let timeout = std::time::Duration::from_secs(5); + let error = classify_shutdown( + 41, + DirectChildShutdownStatus::TimedOut { timeout }, + ProcessGroupShutdownStatus::Absent, + ) + .expect_err("group absence does not substitute for reaping the direct child"); + + assert_eq!( + error.direct_child, + DirectChildShutdownStatus::TimedOut { timeout } + ); + assert_eq!(error.process_group, ProcessGroupShutdownStatus::Absent); + } + + #[test] + fn shutdown_classification_preserves_process_group_probe_failure() { + let error = classify_shutdown( + 41, + DirectChildShutdownStatus::Reaped, + ProcessGroupShutdownStatus::ProbeFailed { errno: 1 }, + ) + .expect_err("an indeterminate process-group probe must fail closed"); + + assert_eq!( + error.process_group, + ProcessGroupShutdownStatus::ProbeFailed { errno: 1 } + ); + } + + #[cfg(all(not(unix), not(windows)))] + #[test] + fn shutdown_classification_quarantines_when_descendant_containment_is_unavailable() { + let error = classify_shutdown( + 41, + DirectChildShutdownStatus::Reaped, + ProcessGroupShutdownStatus::Unavailable, + ) + .expect_err("direct-child reaping cannot prove descendant absence"); + + assert_eq!(error.process_group, ProcessGroupShutdownStatus::Unavailable); + } + + #[test] + fn shutdown_classification_rejects_unavailable_containment() { + let error = classify_shutdown( + 41, + DirectChildShutdownStatus::Reaped, + ProcessGroupShutdownStatus::Unavailable, + ) + .expect_err("Unavailable containment must remain fail-closed on every platform"); + + assert_eq!(error.process_group, ProcessGroupShutdownStatus::Unavailable); + } + + #[cfg(unix)] + #[tokio::test] + async fn shutdown_kills_and_verifies_original_process_group() { + // The background sleep inherits bash's process group. A verified + // shutdown therefore covers both the direct child and its descendant. + let mut client = spawn_script("sleep 60 & wait").await; + let process_group_id = client + .process_group_id + .expect("Unix spawns must retain their original process-group id"); + assert_eq!( + probe_process_group(process_group_id), + ProcessGroupProbe::Present, + "the spawned process group must exist before shutdown" + ); + + client + .shutdown() + .await + .expect("shutdown must verify direct child and process group termination"); + + assert_eq!( + probe_process_group(process_group_id), + ProcessGroupProbe::Absent, + "shutdown must not return while the original process group exists" + ); + client + .shutdown() + .await + .expect("verified shutdown must be safely idempotent"); + } + + #[cfg(unix)] + #[tokio::test] + async fn shutdown_probes_but_never_signals_stored_group_after_child_id_is_cleared() { + let mut client = spawn_script("exit 0").await; + let process_group_id = client + .process_group_id + .expect("Unix spawns must retain their original process-group id"); + client + .child + .wait() + .await + .expect("test process must be reaped"); + assert!( + client.child.id().is_none(), + "Tokio must have cleared the live child id for this regression" + ); + assert_eq!( + process_group_signal_target( + client.spawn_process_id, + client.process_group_id, + client.child.id(), + ), + None, + "a numeric PGID is not a safe signal target after the child identity is gone" + ); + + client + .shutdown() + .await + .expect("an already-absent stored group may still be verified by a read-only probe"); + assert_eq!( + probe_process_group(process_group_id), + ProcessGroupProbe::Absent + ); + } + + #[tokio::test] + async fn send_request_rejects_malformed_ndjson_before_later_matching_response() { + const SECRET: &str = "TOP_SECRET_ORDINARY_PATH"; + let script = r#" + read -r _request + echo '' + echo ' ' + echo '{"jsonrpc":"2.0","id":0,"result":{"secret":"TOP_SECRET_ORDINARY_PATH"}' + echo '{"jsonrpc":"2.0","id":0,"result":{"accepted":true}}' + sleep 1 + "#; + let mut client = spawn_script(script).await; + let observer = ObserverHandle::in_process(); + client.set_observer(Some(observer.clone()), 0); + + let result = client + .send_request("test/malformed", serde_json::json!({})) + .await; + + assert!( + matches!(result, Err(AcpError::Json(_))), + "malformed nonempty NDJSON must fail before a later matching response: {result:?}" + ); + let events = observer.snapshot(); + let parse_events: Vec<_> = events + .iter() + .filter(|event| event.kind == "acp_parse_error") + .collect(); + assert_eq!(parse_events.len(), 1, "one parse failure must be observed"); + assert_eq!( + parse_events[0].payload["lineLength"], + serde_json::json!( + r#"{"jsonrpc":"2.0","id":0,"result":{"secret":"TOP_SECRET_ORDINARY_PATH"}"#.len() + ) + ); + assert!( + events + .iter() + .all(|event| !event.payload.to_string().contains(SECRET)), + "observer evidence must not retain malformed wire content" + ); + assert!( + parse_events[0].payload.get("line").is_none(), + "observer evidence must not expose the malformed line" + ); + } + + #[tokio::test] + async fn session_prompt_rejects_malformed_ndjson_before_later_matching_response() { + const SECRET: &str = "TOP_SECRET_PROMPT_PATH"; + let script = r#" + read -r _prompt + echo '' + echo ' ' + echo '{"jsonrpc":"2.0","id":0,"result":{"secret":"TOP_SECRET_PROMPT_PATH"}' + echo '{"jsonrpc":"2.0","id":0,"result":{"stopReason":"end_turn"}}' + sleep 1 + "#; + let mut client = spawn_script(script).await; + let observer = ObserverHandle::in_process(); + client.set_observer(Some(observer.clone()), 0); + + let result = client + .session_prompt_with_idle_timeout( + "session-test", + "hello", + std::time::Duration::from_secs(2), + std::time::Duration::from_secs(5), + ) + .await; + + assert!( + matches!(result, Err(AcpError::Json(_))), + "malformed nonempty NDJSON must fail before a later prompt response: {result:?}" + ); + let events = observer.snapshot(); + let parse_events: Vec<_> = events + .iter() + .filter(|event| event.kind == "acp_parse_error") + .collect(); + assert_eq!(parse_events.len(), 1, "one parse failure must be observed"); + assert_eq!( + parse_events[0].payload["lineLength"], + serde_json::json!( + r#"{"jsonrpc":"2.0","id":0,"result":{"secret":"TOP_SECRET_PROMPT_PATH"}"#.len() + ) + ); + assert!( + events + .iter() + .all(|event| !event.payload.to_string().contains(SECRET)), + "observer evidence must not retain malformed wire content" + ); + assert!( + parse_events[0].payload.get("line").is_none(), + "observer evidence must not expose the malformed line" + ); + } + + #[tokio::test] + async fn initialize_records_advertised_session_close_capability() { + let script = "read -r _initialize; \ + echo '{\"jsonrpc\":\"2.0\",\"id\":0,\"result\":{\"protocolVersion\":1,\"agentCapabilities\":{\"sessionCapabilities\":{\"close\":{}}}}}'; \ + sleep 10"; + let mut client = spawn_script(script).await; + + client.initialize().await.expect("initialize must succeed"); + + assert!( + client.supports_session_close(), + "an advertised ACP session close capability must be retained" + ); + } + + #[tokio::test] + async fn initialize_leaves_unadvertised_session_close_unsupported() { + let script = "read -r _initialize; \ + echo '{\"jsonrpc\":\"2.0\",\"id\":0,\"result\":{\"protocolVersion\":1,\"agentCapabilities\":{}}}'; \ + sleep 10"; + let mut client = spawn_script(script).await; + + client.initialize().await.expect("initialize must succeed"); + + assert!( + !client.supports_session_close(), + "absence of the optional capability must remain explicit" + ); + } + + #[tokio::test] + async fn session_close_rejects_matching_response_without_result_or_error() { + let script = "read -r _close; \ + echo '{\"jsonrpc\":\"2.0\",\"id\":0}'; \ + sleep 10"; + let mut client = spawn_script(script).await; + + let result = client.session_close("session-old").await; + + assert!( + matches!(result, Err(AcpError::Protocol(_))), + "a close response without result or error is not a positive acknowledgement: {result:?}" + ); + } + #[tokio::test] async fn idle_timeout_fires_on_silent_process() { let mut client = spawn_script("sleep 10").await; @@ -3052,9 +3756,9 @@ mod tests { /// must not dead-letter a drain that simply ran past its grace window. #[tokio::test] async fn cancel_with_cleanup_grace_maps_expiry_to_cancel_drain_timeout() { - // Agent ignores `session/cancel` on stdin and keeps producing noise - // forever — never drains within the grace window. - let mut client = spawn_script("while true; do echo 'noise'; sleep 0.01; done").await; + // Agent ignores `session/cancel` on stdin and keeps producing valid + // JSON activity forever — never drains within the grace window. + let mut client = spawn_script("while true; do echo '{}'; sleep 0.01; done").await; client.last_prompt_id = Some(999); let grace = std::time::Duration::from_millis(200); let result = client @@ -3066,6 +3770,30 @@ mod tests { ); } + #[tokio::test] + async fn cancel_cleanup_writes_cannot_outlive_the_shared_grace_deadline() { + let mut client = spawn_script("sleep 10").await; + client.last_prompt_id = Some(999); + // A response this large fills the child stdin pipe when the child never + // reads. Before the shared write deadline, each cleanup write carried + // its own 30s timeout and could outlive the caller's 5s grace. + client.pending_permission_id = Some(serde_json::json!("x".repeat(2_000_000))); + client.permission_responded = false; + + let grace = std::time::Duration::from_millis(50); + let result = tokio::time::timeout( + std::time::Duration::from_secs(1), + client.cancel_with_cleanup_grace("test-session", grace), + ) + .await + .expect("cleanup must obey the shared grace even under stdin backpressure"); + + assert!( + matches!(result, Err(AcpError::CancelDrainTimeout(g)) if g == grace), + "permission and cancel writes must share the caller's grace deadline: {result:?}" + ); + } + #[tokio::test] async fn idle_resets_on_stdout_activity() { // Send valid JSON (session/update notifications) to reset the idle timer. @@ -3755,6 +4483,69 @@ mod tests { } } + /// A matching JSON-RPC frame without either `result` or `error` is not + /// evidence that the steer was applied. It must release the withheld event + /// through the transport-failure path instead of falsely acknowledging + /// delivery. + #[tokio::test] + async fn native_steer_matching_response_without_result_or_error_fails_closed() { + let script = "sleep 0.5; \ + echo '{\"jsonrpc\":\"2.0\",\"id\":0}'; \ + sleep 10"; + let mut client = spawn_script(script).await; + + let update = session_info_update_msg(Some(serde_json::json!("run-malformed"))); + let _ = client.handle_session_update(&update); + + let (steer_tx, steer_rx) = tokio::sync::mpsc::channel::(1); + client.install_steer_rx(steer_rx); + + let (ack_tx, ack_rx) = tokio::sync::oneshot::channel::(); + let send_task = tokio::spawn(async move { + steer_tx + .send(crate::pool::SteerRequest { + prompt_blocks: vec!["steer body".into()], + ack_tx, + }) + .await + .expect("steer_tx send should succeed"); + }); + + let idle = std::time::Duration::from_secs(2); + let max_dur = std::time::Duration::from_secs(10); + let hard_deadline = tokio::time::Instant::now() + max_dur; + let read_result = client + .read_until_response_with_idle_timeout( + "sess-malformed", + 999, + idle, + hard_deadline, + max_dur, + ) + .await; + send_task.await.expect("send_task should complete"); + assert!( + matches!( + read_result, + Err(AcpError::IdleTimeout(_)) | Err(AcpError::AgentExited) + ), + "read loop should remain bounded after malformed steer response, got {read_result:?}" + ); + + let ack = ack_rx + .await + .expect("ack oneshot must receive a fail-closed outcome"); + match ack { + crate::pool::SteerAck::Err(crate::pool::SteerError::Transport(message)) => { + assert!( + message.contains("missing result or error"), + "malformed response should be classified without raw payload: {message}" + ); + } + other => panic!("expected transport failure, got {other:?}"), + } + } + /// Steer-success renewal keeps the turn alive past the original hard /// deadline. This is the red-on-old/green-on-new test for the core bug /// fix (acp.rs:1440-1444): without renewal, the read loop returns diff --git a/crates/buzz-acp/src/lib.rs b/crates/buzz-acp/src/lib.rs index d63f720c65..4435dc15d7 100644 --- a/crates/buzz-acp/src/lib.rs +++ b/crates/buzz-acp/src/lib.rs @@ -37,11 +37,12 @@ use filter::SubscriptionRule; use futures_util::FutureExt; use nostr::{PublicKey, ToBech32}; use pool::{ - AgentPool, ControlSignal, IdleSwitchResult, OwnedAgent, PromptContext, PromptOutcome, - PromptResult, PromptSource, SessionState, TimeoutKind, + AcceptedDropControl, AgentPool, ControlSignal, IdleSwitchResult, ModelSwitchRequest, + ModelSwitchRollback, OwnedAgent, PromptContext, PromptOutcome, PromptResult, PromptSource, + SessionState, TimeoutKind, }; use pool_lifecycle::PoolLifecycle; -use queue::{CancelReason, EventQueue, FlushBatch, QueuedEvent, ThreadTags}; +use queue::{CancelReason, EnqueueOccurrenceId, EventQueue, FlushBatch, QueuedEvent, ThreadTags}; use relay::{HarnessRelay, RelayEventPublisher}; use tokio::sync::{mpsc, watch}; use tracing_subscriber::EnvFilter; @@ -406,6 +407,16 @@ impl ObserverPublishPacer { return; } } + + async fn wait_priority(&mut self) { + let now = tokio::time::Instant::now(); + if self.next_publish > now { + tokio::time::sleep_until(self.next_publish).await; + } + let published_at = tokio::time::Instant::now(); + self.published.push_back(published_at); + self.next_publish = published_at + OBSERVER_PUBLISH_INTERVAL; + } } fn spawn_relay_observer_publisher( @@ -415,14 +426,19 @@ fn spawn_relay_observer_publisher( agent_pubkey_hex: String, owner_pubkey_hex: String, owner_pubkey: PublicKey, -) -> tokio::task::JoinHandle<()> { +) -> tokio::task::JoinHandle { tokio::spawn(async move { // Subscribe BEFORE snapshotting so an event emitted between the two // calls is never lost: it lands in the snapshot, the live receiver, or - // both. The overlap is deduped in the run loop via the snapshot's - // high-water `seq` (monotonic, assigned at emit). + // both. The overlap is deduped in the run loop via exact snapshot + // `seq` membership (monotonic, assigned at emit). let rx = observer.subscribe(); let snapshot = observer.snapshot(); + // The publisher owns only the receiver. Retaining this sender across + // the run-loop await would keep both broadcast lanes open forever, + // making graceful shutdown time out even after the process-level + // observer handle is dropped. + drop(observer); run_relay_observer_publisher( snapshot, rx, @@ -432,31 +448,33 @@ fn spawn_relay_observer_publisher( owner_pubkey_hex, owner_pubkey, ) - .await; + .await }) } async fn run_relay_observer_publisher( snapshot: Vec, - mut rx: tokio::sync::broadcast::Receiver, + mut rx: observer::ObserverReceiver, publisher: RelayEventPublisher, keys: nostr::Keys, agent_pubkey_hex: String, owner_pubkey_hex: String, owner_pubkey: PublicKey, -) { +) -> bool { + let mut terminal_delivery_ok = !rx.control_result_admission_failed(); let mut coalescer = ObserverChunkCoalescer::default(); let mut pacer = ObserverPublishPacer::new(); - let max_snapshot_seq = snapshot.iter().map(|event| event.seq).max().unwrap_or(0); + let snapshot_seqs: HashSet<_> = snapshot.iter().map(|event| event.seq).collect(); for event in snapshot { for event in coalescer.ingest(event) { - publish_relay_observer_event( + terminal_delivery_ok &= publish_relay_observer_event_preemptible( &publisher, &keys, &agent_pubkey_hex, &owner_pubkey_hex, &owner_pubkey, &mut pacer, + &mut rx, event, ) .await; @@ -470,32 +488,59 @@ async fn run_relay_observer_publisher( result = rx.recv() => { match result { Ok(event) => { - // Skip live events already delivered via the snapshot - // (the subscribe-before-snapshot overlap). - if event.seq <= max_snapshot_seq { + if observer::is_terminal_control_result(&event) { + // A terminal result returned by ObserverReceiver is + // still retained and unacknowledged. Publish it even + // when its sequence was captured in the snapshot: + // the snapshot attempt may have failed, and static + // membership is not delivery proof. + terminal_delivery_ok &= publish_relay_terminal_observer_event( + &publisher, + &keys, + &agent_pubkey_hex, + &owner_pubkey_hex, + &owner_pubkey, + &mut pacer, + &mut rx, + event, + ) + .await; + continue; + } + // Ordinary live telemetry already delivered via the + // snapshot is the subscribe-before-snapshot overlap. + if snapshot_seqs.contains(&event.seq) { continue; } for event in coalescer.ingest(event) { - publish_relay_observer_event( + terminal_delivery_ok &= publish_relay_observer_event_preemptible( &publisher, &keys, &agent_pubkey_hex, - &owner_pubkey_hex, &owner_pubkey, &mut pacer, event, + &owner_pubkey_hex, &owner_pubkey, &mut pacer, + &mut rx, event, ).await; } } Err(tokio::sync::broadcast::error::RecvError::Lagged(count)) => { + // Terminal-lane lag is recovered from the acknowledged + // sequence ledger inside ObserverReceiver, so only + // telemetry overflow reaches this branch. Log that + // observability loss without tainting the independent + // terminal-delivery proof used by graceful shutdown. for event in coalescer.flush() { - publish_relay_observer_event( + terminal_delivery_ok &= publish_relay_observer_event_preemptible( &publisher, &keys, &agent_pubkey_hex, - &owner_pubkey_hex, &owner_pubkey, &mut pacer, event, + &owner_pubkey_hex, &owner_pubkey, &mut pacer, + &mut rx, event, ).await; } tracing::warn!(dropped = count, "relay observer publisher lagged"); } Err(tokio::sync::broadcast::error::RecvError::Closed) => { for event in coalescer.flush() { - publish_relay_observer_event( + terminal_delivery_ok &= publish_relay_observer_event_preemptible( &publisher, &keys, &agent_pubkey_hex, - &owner_pubkey_hex, &owner_pubkey, &mut pacer, event, + &owner_pubkey_hex, &owner_pubkey, &mut pacer, + &mut rx, event, ).await; } break; @@ -505,14 +550,16 @@ async fn run_relay_observer_publisher( _ = flush_interval.tick() => { // Periodic flush ensures live streaming even during continuous chunk delivery. for event in coalescer.flush() { - publish_relay_observer_event( + terminal_delivery_ok &= publish_relay_observer_event_preemptible( &publisher, &keys, &agent_pubkey_hex, - &owner_pubkey_hex, &owner_pubkey, &mut pacer, event, + &owner_pubkey_hex, &owner_pubkey, &mut pacer, + &mut rx, event, ).await; } } } } + terminal_delivery_ok && !rx.control_result_admission_failed() } #[derive(Default)] @@ -794,9 +841,157 @@ async fn publish_relay_observer_event( owner_pubkey_hex: &str, owner_pubkey: &PublicKey, pacer: &mut ObserverPublishPacer, - mut event: observer::ObserverEvent, -) { + event: observer::ObserverEvent, +) -> bool { + let is_terminal_control_result = observer::is_terminal_control_result(&event); + if is_terminal_control_result { + pacer.wait_priority().await; + } else { + pacer.wait().await; + } + publish_relay_observer_event_now( + publisher, + keys, + agent_pubkey_hex, + owner_pubkey_hex, + owner_pubkey, + event, + ) + .await +} + +#[allow(clippy::too_many_arguments)] +async fn publish_relay_terminal_observer_event( + publisher: &RelayEventPublisher, + keys: &nostr::Keys, + agent_pubkey_hex: &str, + owner_pubkey_hex: &str, + owner_pubkey: &PublicKey, + pacer: &mut ObserverPublishPacer, + rx: &mut observer::ObserverReceiver, + event: observer::ObserverEvent, +) -> bool { + let seq = event.seq; + let delivered = publish_relay_observer_event( + publisher, + keys, + agent_pubkey_hex, + owner_pubkey_hex, + owner_pubkey, + pacer, + event, + ) + .await; + if delivered { + rx.acknowledge_control_result(seq); + true + } else { + // The relay publisher already performed its bounded reconnect and + // rate-gate wait. Release the ledger claim for one end-to-end retry; + // a second failure stays retained/in-flight and makes shutdown proof + // fail without an unbounded retry loop. + rx.retry_control_result(seq) + } +} + +#[allow(clippy::too_many_arguments)] +async fn publish_relay_observer_event_preemptible( + publisher: &RelayEventPublisher, + keys: &nostr::Keys, + agent_pubkey_hex: &str, + owner_pubkey_hex: &str, + owner_pubkey: &PublicKey, + pacer: &mut ObserverPublishPacer, + rx: &mut observer::ObserverReceiver, + event: observer::ObserverEvent, +) -> bool { + if observer::is_terminal_control_result(&event) { + return publish_relay_terminal_observer_event( + publisher, + keys, + agent_pubkey_hex, + owner_pubkey_hex, + owner_pubkey, + pacer, + rx, + event, + ) + .await; + } + + let mut priority_lane_open = true; + while priority_lane_open { + tokio::select! { + biased; + result = rx.recv_control_result() => { + match result { + Ok(priority) => { + // ObserverReceiver returns only retained terminal + // results. Snapshot membership alone is never an ACK: + // a captured publication may have failed before this + // overlapping wake-up was consumed. + if !publish_relay_terminal_observer_event( + publisher, + keys, + agent_pubkey_hex, + owner_pubkey_hex, + owner_pubkey, + pacer, + rx, + priority, + ) + .await + { + return false; + } + } + Err(tokio::sync::broadcast::error::RecvError::Lagged(count)) => { + tracing::warn!( + dropped = count, + "relay observer priority publisher lagged" + ); + return false; + } + Err(tokio::sync::broadcast::error::RecvError::Closed) => { + priority_lane_open = false; + } + } + } + _ = pacer.wait() => { + return publish_relay_observer_event_now( + publisher, + keys, + agent_pubkey_hex, + owner_pubkey_hex, + owner_pubkey, + event, + ) + .await; + } + } + } + pacer.wait().await; + publish_relay_observer_event_now( + publisher, + keys, + agent_pubkey_hex, + owner_pubkey_hex, + owner_pubkey, + event, + ) + .await +} + +async fn publish_relay_observer_event_now( + publisher: &RelayEventPublisher, + keys: &nostr::Keys, + agent_pubkey_hex: &str, + owner_pubkey_hex: &str, + owner_pubkey: &PublicKey, + mut event: observer::ObserverEvent, +) -> bool { + let is_terminal_control_result = observer::is_terminal_control_result(&event); // Trim oversized frames to fit the plaintext cap rather than letting // encrypt_observer_payload reject and drop them whole (silent telemetry loss). fit_observer_event_to_budget(&mut event); @@ -804,10 +999,10 @@ async fn publish_relay_observer_event( Ok(encrypted) => encrypted, Err(error) => { tracing::warn!("failed to encrypt relay observer event: {error}"); - return; + return !is_terminal_control_result; } }; - let builder = match buzz_sdk::build_agent_observer_frame( + let mut builder = match buzz_sdk::build_agent_observer_frame( owner_pubkey_hex, agent_pubkey_hex, OBSERVER_FRAME_TELEMETRY, @@ -816,28 +1011,109 @@ async fn publish_relay_observer_event( Ok(builder) => builder, Err(error) => { tracing::warn!("failed to build relay observer event: {error}"); - return; + return !is_terminal_control_result; } }; + if is_terminal_control_result { + let priority_tag = match nostr::Tag::parse(["priority", "control-result"]) { + Ok(tag) => tag, + Err(error) => { + tracing::warn!("failed to build observer control-result priority tag: {error}"); + return false; + } + }; + builder = builder.tag(priority_tag); + } let signed = match builder.sign_with_keys(keys) { Ok(event) => event, Err(error) => { tracing::warn!("failed to sign relay observer event: {error}"); - return; + return !is_terminal_control_result; } }; - if let Err(error) = publisher.publish_event(signed).await { + let publish_result = if is_terminal_control_result { + publisher.publish_terminal_event(signed).await + } else { + publisher.publish_event(signed).await + }; + if let Err(error) = publish_result { tracing::warn!("relay observer event dropped: {error}"); + return !is_terminal_control_result; } + true } /// Maximum age (seconds) for an observer control frame to be considered fresh. const OBSERVER_CONTROL_FRESHNESS_SECS: i64 = 300; +const MODEL_SWITCH_REQUEST_ID_LEN: usize = 32; +/// Bound replay memory while covering substantially more than the maximum +/// number of legitimate controls expected inside the freshness window. +const OBSERVER_CONTROL_DEDUP_CAPACITY: usize = 1_024; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum ObserverControlAdmission { + Admitted, + Replay, + CapacityExceeded, +} + +struct ObserverControlDedup { + capacity: usize, + order: VecDeque<(String, i64)>, + seen: HashSet, +} + +impl ObserverControlDedup { + fn new(capacity: usize) -> Self { + assert!( + capacity > 0, + "observer control dedup capacity must be positive" + ); + Self { + capacity, + order: VecDeque::with_capacity(capacity), + seen: HashSet::with_capacity(capacity), + } + } + + fn admit( + &mut self, + event_id: String, + event_timestamp: i64, + now: i64, + ) -> ObserverControlAdmission { + let seen = &mut self.seen; + self.order.retain(|(retained_id, retained_timestamp)| { + let fresh = retained_timestamp.abs_diff(now) <= OBSERVER_CONTROL_FRESHNESS_SECS as u64; + if !fresh { + seen.remove(retained_id); + } + fresh + }); + if self.seen.contains(&event_id) { + return ObserverControlAdmission::Replay; + } + if self.order.len() >= self.capacity { + return ObserverControlAdmission::CapacityExceeded; + } + self.seen.insert(event_id.clone()); + self.order.push_back((event_id, event_timestamp)); + ObserverControlAdmission::Admitted + } +} + +#[allow(clippy::too_many_arguments)] fn handle_relay_observer_control_event( keys: &nostr::Keys, event: nostr::Event, + dedup: &mut ObserverControlDedup, pool: &mut AgentPool, + queue: &mut EventQueue, + config: &Config, + crash_history: &mut [SlotCircuit], + respawn_tx: &mpsc::Sender, + respawn_tasks: &mut tokio::task::JoinSet<()>, observer: Option<&observer::ObserverHandle>, owner_pubkey_hex: &str, ) { @@ -860,7 +1136,7 @@ fn handle_relay_observer_control_event( // Freshness: reject stale/replayed frames outside ±5 minute window. let now = chrono::Utc::now().timestamp(); let event_ts = event.created_at.as_secs() as i64; - if (event_ts - now).unsigned_abs() > OBSERVER_CONTROL_FRESHNESS_SECS as u64 { + if event_ts.abs_diff(now) > OBSERVER_CONTROL_FRESHNESS_SECS as u64 { tracing::warn!( event_ts, now, @@ -876,6 +1152,21 @@ fn handle_relay_observer_control_event( return; } }; + match dedup.admit(event.id.to_hex(), event_ts, now) { + ObserverControlAdmission::Admitted => {} + ObserverControlAdmission::Replay => { + tracing::warn!(event_id = %event.id, "replayed observer control frame — dropping"); + return; + } + ObserverControlAdmission::CapacityExceeded => { + tracing::error!( + event_id = %event.id, + capacity = dedup.capacity, + "observer control replay window is full — failing closed" + ); + return; + } + } let command_type = payload.get("type").and_then(|value| value.as_str()); match command_type { @@ -883,7 +1174,16 @@ fn handle_relay_observer_control_event( handle_cancel_turn_control(&payload, pool, observer); } Some("switch_model") => { - handle_switch_model_control(&payload, pool, observer); + handle_switch_model_control( + &payload, + pool, + queue, + config, + crash_history, + respawn_tx, + respawn_tasks, + observer, + ); } _ => { tracing::debug!(payload = %payload, "ignoring unknown observer control frame"); @@ -906,8 +1206,11 @@ fn handle_cancel_turn_control( return; }; - let fired = signal_in_flight_task(pool, channel_id, ControlSignal::Cancel); - let status = if fired { "sent" } else { "no_active_turn" }; + let status = match signal_in_flight_task(pool, channel_id, ControlSignal::Cancel) { + ControlSignalResult::Delivered => "sent", + ControlSignalResult::DropRecorded => "drop_recorded", + ControlSignalResult::NotAccepted => "no_active_turn", + }; if let Some(observer) = observer { observer.emit( "control_result", @@ -934,12 +1237,18 @@ fn handle_cancel_turn_control( /// post-cancel via `create_session_and_apply_model` (the turn restarts on the /// unchanged model + an `unsupported_model` result). /// -/// Idle path: validate against the cached catalog *before* invalidating -/// (pre-cancel guard), then set `desired_model` + invalidate. The override -/// takes visible effect on the agent's next turn. +/// Idle path: validate against the cached catalog, claim the owning adapter, +/// and recycle its process. The override takes effect on the fresh process's +/// next session without locally abandoning the prior remote session. +#[allow(clippy::too_many_arguments)] fn handle_switch_model_control( payload: &serde_json::Value, pool: &mut AgentPool, + queue: &mut EventQueue, + config: &Config, + crash_history: &mut [SlotCircuit], + respawn_tx: &mpsc::Sender, + respawn_tasks: &mut tokio::task::JoinSet<()>, observer: Option<&observer::ObserverHandle>, ) { let Some(channel_id) = payload @@ -954,32 +1263,62 @@ fn handle_switch_model_control( tracing::warn!("observer switch_model control frame missing modelId"); return; }; + let Some(request_id) = parse_model_switch_request_id(payload) else { + tracing::warn!("observer switch_model control frame missing or invalid requestId"); + return; + }; + let request = ModelSwitchRequest::new(model_id, request_id); // A turn is in flight for this channel iff a task_map entry exists. The // agent is moved out of the pool during a turn, so the control oneshot is // the only reachable lever; an idle channel has no such entry. - let turn_in_flight = pool + let in_flight_agent = pool .task_map() .values() - .any(|m| m.channel_id == Some(channel_id)); + .find(|m| m.channel_id == Some(channel_id)) + .map(|m| m.agent_index); - let status = if turn_in_flight { + let status = if let Some(agent_index) = in_flight_agent { // Busy path: deliver over the oneshot. `false` means the oneshot was // already consumed this turn (a prior cancel/interrupt) — the turn is // already ending, so the switch cannot land on it. - if signal_in_flight_task( - pool, - channel_id, - ControlSignal::SwitchModel(model_id.to_string()), - ) { + if signal_in_flight_task(pool, channel_id, ControlSignal::SwitchModel(request)) + == ControlSignalResult::Delivered + { + queue.require_agent(channel_id, agent_index); "sent" } else { "turn_ending" } } else { - // Idle path: validate against the cached catalog before invalidating. - match pool.switch_idle_agent_model(channel_id, model_id) { - IdleSwitchResult::Switched => "switched", + // Idle path: validate before taking the adapter out of service. + match pool.switch_idle_agent_model(channel_id, model_id, request_id) { + IdleSwitchResult::Recycle(agent) => { + let agent = *agent; + let agent_index = agent.index; + queue.require_agent(channel_id, agent_index); + let scheduled = spawn_model_switch_recycle_task( + agent, + config, + &mut crash_history[agent_index], + respawn_tx, + respawn_tasks, + observer.cloned(), + PendingSwitchControl { + channel_id, + model_id: model_id.to_string(), + request_id: request_id.to_string(), + }, + ); + if scheduled { + // Scheduling a recycle is acceptance, not proof that the + // replacement initialized or applied this model. + "recycling" + } else { + queue.clear_required_agent(channel_id); + "switch_failed" + } + } IdleSwitchResult::UnsupportedModel => "unsupported_model", IdleSwitchResult::NoIdleAgent => "no_active_turn", } @@ -995,15 +1334,59 @@ fn handle_switch_model_control( turn_id: None, started_at: None, }, - serde_json::json!({ - "type": "switch_model", - "status": status, - "modelId": model_id, - }), + switch_model_control_result_payload(status, model_id, request_id), ); } } +fn switch_model_control_result_payload( + status: &str, + model_id: &str, + request_id: &str, +) -> serde_json::Value { + serde_json::json!({ + "type": "switch_model", + "status": status, + "modelId": model_id, + "requestId": request_id, + }) +} + +fn parse_model_switch_request_id(payload: &serde_json::Value) -> Option<&str> { + let request_id = payload.get("requestId")?.as_str()?; + (request_id.len() == MODEL_SWITCH_REQUEST_ID_LEN + && request_id + .bytes() + .all(|byte| byte.is_ascii_hexdigit() && !byte.is_ascii_uppercase())) + .then_some(request_id) +} + +#[allow(clippy::too_many_arguments)] +fn recycle_idle_channel_sessions( + pool: &mut AgentPool, + channel_id: Uuid, + config: &Config, + crash_history: &mut [SlotCircuit], + respawn_tx: &mpsc::Sender, + respawn_tasks: &mut tokio::task::JoinSet<()>, + observer: Option, +) -> usize { + let agents = pool.take_idle_agents_with_session(channel_id); + let count = agents.len(); + for agent in agents { + let index = agent.index; + let _ = spawn_recycle_task( + agent, + config, + &mut crash_history[index], + respawn_tx, + respawn_tasks, + observer.clone(), + ); + } + count +} + /// 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. @@ -1014,6 +1397,27 @@ const CIRCUIT_BREAKER_COOLDOWN: Duration = Duration::from_secs(300); // 5 minute const RESPAWN_BASE_DELAY: Duration = Duration::from_secs(1); /// Maximum respawn backoff delay. const RESPAWN_MAX_DELAY: Duration = Duration::from_secs(30); +/// Minimum spacing between negotiated compatibility recycles for one slot. +/// +/// Optional-close adapters require process replacement to retire sessions, but +/// normal cancellation or max-token responses must not become an unbounded +/// kill/spawn loop. +const RECYCLE_MIN_INTERVAL: Duration = Duration::from_secs(5); + +#[derive(Clone, Debug)] +struct ReplacementModelIntent { + desired_model: Option, + model_overridden: bool, + model_switch_request_id: Option, + model_switch_rollback: Option>, +} + +#[derive(Clone, Debug)] +struct PendingSwitchControl { + channel_id: Uuid, + model_id: String, + request_id: String, +} /// Per-slot circuit breaker state. /// @@ -1027,10 +1431,19 @@ const RESPAWN_MAX_DELAY: Duration = Duration::from_secs(30); struct SlotCircuit { crash_times: Vec, open_until: Option, + /// Permanent, process-lifetime quarantine after an adapter's termination + /// could not be proved. A timed crash cooldown must never authorize a new + /// owner while the prior process group may still exist. + cleanup_unverified: bool, /// 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, + /// Earliest time another negotiated compatibility recycle may begin. + next_recycle_not_before: Option, + /// Live model intent that must survive an intentional or cleanup-triggered + /// replacement until a replacement agent initializes successfully. + pending_model_intent: Option, } /// Result of [`SlotCircuit::record_crash`]. @@ -1046,11 +1459,36 @@ enum CrashVerdict { } impl SlotCircuit { + fn new() -> Self { + Self { + crash_times: Vec::new(), + open_until: None, + cleanup_unverified: false, + respawn_in_flight: false, + next_recycle_not_before: None, + pending_model_intent: None, + } + } + + /// Reserve the next compatibility-recycle slot and return its delay. + /// + /// This budget is separate from crash accounting: negotiated recycling is + /// healthy maintenance, but it still needs a hard rate bound. + fn schedule_recycle(&mut self) -> Duration { + let now = std::time::Instant::now(); + let scheduled = self.next_recycle_not_before.unwrap_or(now).max(now); + self.next_recycle_not_before = Some(scheduled + RECYCLE_MIN_INTERVAL); + scheduled.saturating_duration_since(now) + } + /// 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 { + if self.cleanup_unverified { + return CrashVerdict::CircuitOpen; + } let now = std::time::Instant::now(); // Half-open: cooldown elapsed → allow one probe. @@ -1102,6 +1540,38 @@ impl SlotCircuit { self.open_until = Some(std::time::Instant::now() + CIRCUIT_BREAKER_COOLDOWN); } + /// Quarantine this slot for the remainder of the supervisor process. + /// + /// Recovery requires an operator-visible process restart after external + /// process/cgroup verification; an ordinary cooldown is not cleanup proof. + fn mark_cleanup_unverified(&mut self) { + self.cleanup_unverified = true; + self.open_until = None; + } + + /// Restore the exact intent that preceded a terminally rejected switch. + /// + /// The observer result makes the current requested model final from the + /// caller's perspective, so a later cooldown refill must never silently + /// retry it. Nested rollback history represents an older pending switch + /// and remains eligible to resume. + fn rollback_terminal_switch_failure(&mut self) { + let rollback = self + .pending_model_intent + .take() + .and_then(|intent| intent.model_switch_rollback); + self.pending_model_intent = rollback.map(|rollback| ReplacementModelIntent { + desired_model: rollback.desired_model, + model_overridden: rollback.model_overridden, + model_switch_request_id: rollback.request_id, + model_switch_rollback: rollback.previous, + }); + } + + fn blocks_supervisor_exit(&self) -> bool { + self.cleanup_unverified + } + /// 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. @@ -1111,6 +1581,9 @@ impl SlotCircuit { /// (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 { + if self.cleanup_unverified { + return false; + } let now = std::time::Instant::now(); match self.open_until { Some(open_until) => { @@ -1131,17 +1604,74 @@ impl SlotCircuit { } } -/// 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. +/// True if any slot has recovery work in flight or is deliberately quarantined. +/// Used to prevent premature "all agents dead" exits: respawns may succeed in +/// seconds, while quarantined slots must keep the current process alive and +/// degraded so a service manager cannot restart into possible process overlap. fn any_respawn_in_flight(crash_history: &[SlotCircuit]) -> bool { - crash_history.iter().any(|s| s.respawn_in_flight) + crash_history + .iter() + .any(|s| s.respawn_in_flight || s.blocks_supervisor_exit()) } /// Result of a background respawn task. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum RespawnFailureClass { + /// The prior owner was retired, but replacement spawn/init failed. + Spawn, + /// The prior owner's direct child/process group could not be proved absent. + CleanupUnverified, +} + +/// Typed failure from spawning and initializing a replacement adapter. +/// +/// Cleanup proof is part of the outcome: a failed handshake is retryable only +/// when the newly spawned client was also proved shut down. +#[derive(Debug)] +enum SpawnInitFailure { + Spawn(anyhow::Error), + CleanupUnverified(anyhow::Error), +} + +impl SpawnInitFailure { + fn into_respawn_failure(self) -> (anyhow::Error, RespawnFailureClass) { + match self { + Self::Spawn(error) => (error, RespawnFailureClass::Spawn), + Self::CleanupUnverified(error) => (error, RespawnFailureClass::CleanupUnverified), + } + } +} + +type SpawnInitResult = std::result::Result<(AcpClient, u32, String), SpawnInitFailure>; + +fn missing_replacement_ownership( + guard: &mut RespawnGuard, + phase: &'static str, +) -> SpawnInitFailure { + guard.phase = RespawnPhase::OldOwnerUnverified; + SpawnInitFailure::CleanupUnverified(anyhow::anyhow!("replacement ownership missing {phase}")) +} + struct RespawnResult { index: usize, - /// Tuple: (initialized client, protocol version, agent name). - result: Result<(AcpClient, u32, String)>, + /// `Some`: initialized replacement; `None`: old adapter was fully shut + /// down while an open circuit intentionally suppressed replacement. + result: Result>, + failure_class: Option, + /// Model intent to apply to the replacement agent. + /// + /// Crash recovery supplies configured defaults. Intentional compatibility + /// recycling preserves a live `SwitchModel` override. + desired_model: Option, + model_overridden: bool, + /// Preserve the slot's pending model intent if replacement did not + /// initialize. Cleanup-triggered replacement and negotiated recycling set + /// this; ordinary crash recovery intentionally resets to config defaults. + retain_model_intent: bool, + /// True only when this result already emitted a terminal `switch_failed` + /// control result. The supervisor must atomically restore the preceding + /// intent before it permits any cooldown refill. + terminal_switch_failure: bool, } /// Outcome of a non-cancelling steer attempt, forwarded from a per-attempt @@ -1155,6 +1685,7 @@ struct RespawnResult { /// `event_id` is the hex id of the single event the steer carried. struct SteerAckEvent { channel_id: Uuid, + occurrence_id: EnqueueOccurrenceId, event_id: String, /// `Ok` if the read loop sent any of the locked `SteerAck` variants. /// `Err` if the oneshot was dropped without a send — should not happen @@ -1164,67 +1695,625 @@ struct SteerAckEvent { ack: std::result::Result, } +type InitializingReplacement = Arc>>; + +#[derive(Clone)] +enum RespawnPhase { + /// The task may still own the previous adapter process. + OldOwnerUnverified, + /// Prior-owner cleanup is verified and no replacement process exists. + CleanNoOwner, + /// A replacement process exists and must be explicitly shut down if the + /// surrounding respawn task is cancelled. + Initializing(InitializingReplacement), +} + /// 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, + desired_model: Option, + model_overridden: bool, + retain_model_intent: bool, + switch_failure: Option<(observer::ObserverHandle, PendingSwitchControl)>, + terminal_switch_failure_emitted: bool, + phase: RespawnPhase, sent: bool, } impl RespawnGuard { - fn new(index: usize, tx: mpsc::Sender) -> Self { + fn new( + index: usize, + tx: mpsc::Sender, + desired_model: Option, + model_overridden: bool, + retain_model_intent: bool, + ) -> Self { Self { index, tx, + desired_model, + model_overridden, + retain_model_intent, + switch_failure: None, + terminal_switch_failure_emitted: false, + phase: RespawnPhase::OldOwnerUnverified, 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<(AcpClient, u32, String)>) { - // 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. - } - } + fn mark_owner_cleanup_verified(&mut self) { + self.phase = RespawnPhase::CleanNoOwner; } -} -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")), - }); - } + fn mark_replacement_initializing(&mut self, client: InitializingReplacement) { + self.phase = RespawnPhase::Initializing(client); } -} -// -// 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 + fn with_switch_failure( + mut self, + observer: Option, + control: Option, + ) -> Self { + self.switch_failure = observer.zip(control); + self + } + + fn emit_switch_failure(&mut self) -> bool { + if self.terminal_switch_failure_emitted { + return true; + } + let Some((observer, control)) = self.switch_failure.take() else { + return false; + }; + observer.emit( + "control_result", + Some(self.index), + &observer::ObserverContext { + channel_id: Some(control.channel_id.to_string()), + session_id: None, + turn_id: None, + started_at: None, + }, + serde_json::json!({ + "type": "switch_model", + "status": "switch_failed", + "modelId": control.model_id, + "requestId": control.request_id, + }), + ); + self.terminal_switch_failure_emitted = true; + true + } + + /// 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: SpawnInitResult) { + match result { + Ok(initialized) => self.send_inner(Ok(Some(initialized)), None), + Err(failure) => { + let (error, failure_class) = failure.into_respawn_failure(); + self.send_inner(Err(error), Some(failure_class)); + } + } + } + + /// Report that bounded shutdown completed and an open circuit deliberately + /// suppressed replacement. + fn send_cleanup_complete(mut self) { + self.send_inner(Ok(None), None); + } + + /// Report that replacement is forbidden because prior-owner termination + /// could not be proved. + fn send_cleanup_unverified(mut self, error: anyhow::Error) { + self.send_inner(Err(error), Some(RespawnFailureClass::CleanupUnverified)); + } + + fn send_inner( + &mut self, + result: Result>, + failure_class: Option, + ) { + let terminal_switch_failure = result.is_err() && self.emit_switch_failure(); + // 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, + failure_class, + desired_model: self.desired_model.clone(), + model_overridden: self.model_overridden, + retain_model_intent: self.retain_model_intent, + terminal_switch_failure, + }) { + 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 { + return; + } + + match self.phase.clone() { + RespawnPhase::OldOwnerUnverified => { + tracing::error!( + agent = self.index, + "respawn task exited while process ownership was unverified" + ); + self.send_inner( + Err(anyhow::anyhow!("respawn task panicked or was cancelled")), + Some(RespawnFailureClass::CleanupUnverified), + ); + } + RespawnPhase::CleanNoOwner => { + tracing::debug!( + agent = self.index, + "respawn task cancelled after verified cleanup and before replacement spawn" + ); + self.send_inner(Ok(None), None); + } + RespawnPhase::Initializing(client) => { + let Some(runtime) = tokio::runtime::Handle::try_current().ok() else { + self.send_inner( + Err(anyhow::anyhow!( + "replacement initialization cancelled without an async cleanup runtime" + )), + Some(RespawnFailureClass::CleanupUnverified), + ); + return; + }; + + // Transfer result reporting and any terminal switch context to + // a detached bounded-cleanup task. The cancelled initializer + // releases the mutex guard while unwinding; this task then + // takes the replacement client and proves its process group is + // absent before classifying cancellation as retryable. + let mut cleanup_guard = RespawnGuard::new( + self.index, + self.tx.clone(), + self.desired_model.clone(), + self.model_overridden, + self.retain_model_intent, + ) + .with_switch_failure(None, None); + cleanup_guard.switch_failure = self.switch_failure.take(); + cleanup_guard.terminal_switch_failure_emitted = + self.terminal_switch_failure_emitted; + self.sent = true; + drop(runtime.spawn(async move { + let mut slot = client.lock().await; + let replacement = slot.take(); + drop(slot); + match replacement { + Some(mut acp) => match acp.shutdown().await { + Ok(()) => { + cleanup_guard.send(Err(SpawnInitFailure::Spawn(anyhow::anyhow!( + "replacement initialization cancelled after verified cleanup" + )))) + } + Err(error) => cleanup_guard.send_cleanup_unverified(anyhow::anyhow!( + "replacement initialization cancelled; cleanup unverified: {error}" + )), + }, + None => cleanup_guard.send_cleanup_complete(), + } + })); + } + } + } +} + +#[cfg(test)] +mod respawn_guard_tests { + use super::*; + + #[cfg(unix)] + struct ProcessGroupCleanup(Option); + + #[cfg(unix)] + impl Drop for ProcessGroupCleanup { + fn drop(&mut self) { + if let Some(process_group) = self.0 { + let _ = nix::sys::signal::killpg(process_group, nix::sys::signal::Signal::SIGKILL); + } + } + } + + #[tokio::test] + async fn cancelled_respawn_task_quarantines_slot_when_cleanup_is_unverified() { + let (tx, mut rx) = mpsc::channel(1); + let (armed_tx, armed_rx) = tokio::sync::oneshot::channel(); + + let task = tokio::spawn(async move { + let _guard = RespawnGuard::new(1, tx, None, false, false); + armed_tx.send(()).expect("test task should arm guard"); + std::future::pending::<()>().await; + }); + + armed_rx.await.expect("test task should start"); + task.abort(); + assert!(task + .await + .expect_err("task should be cancelled") + .is_cancelled()); + + let result = rx + .recv() + .await + .expect("guard drop must report the cancelled respawn"); + assert!(result.result.is_err()); + assert_eq!( + result.failure_class, + Some(RespawnFailureClass::CleanupUnverified) + ); + } + + #[tokio::test] + async fn cancellation_after_verified_old_owner_cleanup_is_clean() { + let (tx, mut rx) = mpsc::channel(1); + let (armed_tx, armed_rx) = tokio::sync::oneshot::channel(); + + let task = tokio::spawn(async move { + let mut guard = RespawnGuard::new(1, tx, None, false, false); + guard.mark_owner_cleanup_verified(); + armed_tx.send(()).expect("test task should arm guard"); + std::future::pending::<()>().await; + }); + + armed_rx.await.expect("test task should start"); + task.abort(); + assert!(task + .await + .expect_err("task should be cancelled") + .is_cancelled()); + + let result = rx + .recv() + .await + .expect("guard drop must report clean cancellation"); + assert!(matches!(result.result, Ok(None))); + assert_eq!(result.failure_class, None); + } + + #[cfg(unix)] + #[tokio::test] + async fn cancellation_during_replacement_initialize_verifies_process_group_cleanup() { + use nix::errno::Errno; + use nix::sys::signal::{killpg, Signal}; + use nix::unistd::Pid; + + let pid_file = std::env::temp_dir().join(format!( + "buzz-acp-respawn-init-{}-{}.pid", + std::process::id(), + Uuid::new_v4() + )); + let script = r#"printf '%s\n' "$$" > "$1"; exec /bin/sleep 300"#; + let args = vec![ + "-c".to_string(), + script.to_string(), + "buzz-acp-respawn-init-test".to_string(), + pid_file.to_string_lossy().into_owned(), + ]; + let (tx, mut rx) = mpsc::channel(1); + let task = tokio::spawn(async move { + let mut guard = RespawnGuard::new(1, tx, None, false, false); + guard.mark_owner_cleanup_verified(); + let result = spawn_and_init("bash", &args, &[], false, 1, None, &mut guard).await; + guard.send(result); + }); + + tokio::time::timeout(Duration::from_secs(10), async { + while !pid_file.exists() { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("replacement adapter must publish its PID"); + let process_group = Pid::from_raw( + std::fs::read_to_string(&pid_file) + .expect("replacement PID file") + .trim() + .parse() + .expect("replacement PID"), + ); + let mut cleanup = ProcessGroupCleanup(Some(process_group)); + assert_eq!(killpg(process_group, None::), Ok(())); + + task.abort(); + assert!(task + .await + .expect_err("initializer task should be cancelled") + .is_cancelled()); + let result = tokio::time::timeout(Duration::from_secs(10), rx.recv()) + .await + .expect("bounded cleanup must report") + .expect("cleanup result channel"); + assert!(result.result.is_err()); + assert_eq!( + result.failure_class, + Some(RespawnFailureClass::Spawn), + "verified replacement cleanup remains retryable" + ); + assert_eq!( + killpg(process_group, None::), + Err(Errno::ESRCH), + "cancelled replacement process group must be absent before reporting" + ); + cleanup.0 = None; + let _ = std::fs::remove_file(pid_file); + } + + #[test] + fn partial_init_cleanup_failure_is_not_downgraded_to_spawn_failure() { + let (tx, mut rx) = mpsc::channel(1); + let guard = RespawnGuard::new(1, tx, None, false, false); + + guard.send(Err(SpawnInitFailure::CleanupUnverified(anyhow::anyhow!( + "spawned adapter cleanup unverified" + )))); + + let result = rx + .try_recv() + .expect("typed spawn/init failure must reach the supervisor"); + assert!(result.result.is_err()); + assert_eq!( + result.failure_class, + Some(RespawnFailureClass::CleanupUnverified) + ); + } + + #[test] + fn missing_replacement_ownership_is_typed_cleanup_uncertainty() { + let (tx, _rx) = mpsc::channel(1); + let mut guard = RespawnGuard::new(1, tx, None, false, false); + guard.mark_owner_cleanup_verified(); + + let failure = missing_replacement_ownership(&mut guard, "before adapter initialization"); + + assert!(matches!(guard.phase, RespawnPhase::OldOwnerUnverified)); + assert!(matches!( + failure, + SpawnInitFailure::CleanupUnverified(error) + if error.to_string() + == "replacement ownership missing before adapter initialization" + )); + } + + #[test] + fn failed_idle_switch_recycle_emits_channel_scoped_terminal_failure() { + let (tx, mut rx) = mpsc::channel(1); + let observer = observer::ObserverHandle::in_process(); + let channel_id = Uuid::new_v4(); + let guard = RespawnGuard::new(2, tx, Some("model-b".into()), true, true) + .with_switch_failure( + Some(observer.clone()), + Some(PendingSwitchControl { + channel_id, + model_id: "model-b".into(), + request_id: "0123456789abcdef0123456789abcdef".into(), + }), + ); + + guard.send(Err(SpawnInitFailure::Spawn(anyhow::anyhow!( + "replacement failed" + )))); + + let result = rx.try_recv().expect("failure must reach the supervisor"); + assert!(result.result.is_err()); + assert_eq!( + result.failure_class, + Some(RespawnFailureClass::Spawn), + "a normally returned spawn/init failure remains eligible for cooldown retry" + ); + assert!( + result.terminal_switch_failure, + "the supervisor must know that the request received a terminal result" + ); + let failures: Vec<_> = observer + .snapshot() + .into_iter() + .filter(|event| event.kind == "control_result") + .collect(); + let channel_id_string = channel_id.to_string(); + assert_eq!(failures.len(), 1); + assert_eq!(failures[0].agent_index, Some(2)); + assert_eq!( + failures[0].channel_id.as_deref(), + Some(channel_id_string.as_str()) + ); + assert_eq!(failures[0].payload["status"], "switch_failed"); + assert_eq!(failures[0].payload["modelId"], "model-b"); + assert_eq!( + failures[0].payload["requestId"], + "0123456789abcdef0123456789abcdef" + ); + } + + #[test] + fn terminal_switch_failure_restores_exact_prior_intent_for_retry() { + let mut slot = SlotCircuit::new(); + slot.pending_model_intent = Some(ReplacementModelIntent { + desired_model: Some("model-b".into()), + model_overridden: true, + model_switch_request_id: Some("request-b".into()), + model_switch_rollback: Some(Box::new(ModelSwitchRollback { + desired_model: Some("model-a".into()), + model_overridden: true, + request_id: Some("request-a".into()), + previous: Some(Box::new(ModelSwitchRollback { + desired_model: Some("configured-default".into()), + model_overridden: false, + request_id: None, + previous: None, + })), + })), + }); + + slot.rollback_terminal_switch_failure(); + + let restored = slot + .pending_model_intent + .as_ref() + .expect("the exact prior pending intent must resume"); + assert_eq!(restored.desired_model.as_deref(), Some("model-a")); + assert!(restored.model_overridden); + assert_eq!( + restored.model_switch_request_id.as_deref(), + Some("request-a") + ); + let prior = restored + .model_switch_rollback + .as_ref() + .expect("nested rollback history must remain intact"); + assert_eq!(prior.desired_model.as_deref(), Some("configured-default")); + assert!(!prior.model_overridden); + assert_eq!(prior.request_id, None); + assert!(prior.previous.is_none()); + } +} + +#[cfg(test)] +mod pool_startup_tests { + use super::*; + + #[test] + fn cleanup_unverified_lazy_startup_is_not_retried() { + assert!( + PoolStartupFailure::CleanupUnverified("cleanup unverified".into()).cleanup_unverified() + ); + assert!(!PoolStartupFailure::Retryable("provider unavailable".into()).cleanup_unverified()); + assert!(lazy_pool_can_wake(true, false, false)); + assert!(!lazy_pool_can_wake(true, false, true)); + assert!(!lazy_pool_can_wake(false, false, false)); + assert!(!lazy_pool_can_wake(true, true, false)); + } + + #[test] + fn cleanup_unverified_eager_startup_enters_non_spawning_quarantine() { + let startup = resolve_initial_pool_startup( + false, + Some(Err(PoolStartupFailure::CleanupUnverified( + "adapter ownership is unverified".into(), + ))), + ) + .expect("cleanup uncertainty must not terminate the supervisor"); + + assert!( + matches!(startup, InitialPoolStartup::CleanupQuarantine(_)), + "eager cleanup uncertainty must keep the harness alive without a ready pool" + ); + } + + #[test] + fn missing_eager_startup_result_returns_typed_invariant_failure() { + let error = match resolve_initial_pool_startup(false, None) { + Ok(_) => panic!("missing eager startup evidence must fail closed"), + Err(error) => error, + }; + + assert!(matches!( + error, + PoolStartupFailure::Invariant(message) + if message == "eager pool startup did not provide an initialization result" + )); + } +} + +#[cfg(test)] +mod shutdown_verification_tests { + use super::*; + + #[test] + fn checked_out_respawn_and_idle_uncertainty_each_prevent_clean_shutdown() { + for owner in [ + ShutdownOwner::CheckedOut, + ShutdownOwner::Respawn, + ShutdownOwner::Idle, + ] { + let mut verification = ShutdownVerification::default(); + verification.record(owner, false); + assert!( + verification.into_result().is_err(), + "{owner:?} ownership uncertainty must make shutdown fail" + ); + } + + let mut verification = ShutdownVerification::default(); + verification.record(ShutdownOwner::CheckedOut, true); + verification.record(ShutdownOwner::Respawn, true); + verification.record(ShutdownOwner::Idle, true); + assert!( + verification.into_result().is_ok(), + "fully verified ownership may return clean" + ); + } + + #[tokio::test] + async fn pre_loop_failure_cannot_escape_existing_cleanup_quarantine() { + let (shutdown_tx, mut shutdown_rx) = watch::channel(()); + let task = tokio::spawn(async move { + let mut pool = AgentPool::from_slots(Vec::new()); + await_pre_loop_operation( + async { Err::<(), _>("relay unavailable") }, + &mut shutdown_rx, + &mut pool, + "relay connect", + true, + ) + .await + }); + + tokio::task::yield_now().await; + assert!( + !task.is_finished(), + "an automatic pre-loop error must remain inside cleanup quarantine" + ); + + shutdown_tx + .send(()) + .expect("explicit shutdown must release quarantine"); + let error = tokio::time::timeout(Duration::from_secs(1), task) + .await + .expect("quarantined failure must respond to explicit shutdown") + .expect("test task") + .expect_err("pre-loop operation must still fail"); + assert!( + error.to_string().contains("cleanup remains unverified"), + "the returned error must preserve cleanup uncertainty: {error}" + ); + } + + #[test] + fn cli_failure_preserves_cleanup_uncertainty_as_primary_error() { + let result: Result<()> = cli_failure_after_cleanup( + anyhow::anyhow!("adapter protocol failed"), + Err(anyhow::anyhow!("adapter cleanup unverified")), + ); + + let error = result.expect_err("cleanup uncertainty must fail the command"); + assert_eq!(error.to_string(), "adapter cleanup unverified"); + } +} + +// +// 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). @@ -1233,6 +2322,53 @@ pub fn run() -> Result<()> { tokio_main() } +async fn install_shutdown_channel() -> Result<(watch::Sender<()>, watch::Receiver<()>)> { + let (shutdown_tx, shutdown_rx) = watch::channel(()); + + #[cfg(unix)] + { + use tokio::signal::unix::{signal, SignalKind}; + + // Constructing the streams installs Tokio's process-level handlers + // synchronously. Do this before any adapter spawn so SIGINT/SIGTERM can + // never take the default exit path while a detached adapter PGID lives. + let mut sigint = signal(SignalKind::interrupt()).map_err(|error| anyhow::anyhow!(error))?; + let mut sigterm = + signal(SignalKind::terminate()).map_err(|error| anyhow::anyhow!(error))?; + let signal_tx = shutdown_tx.clone(); + tokio::spawn(async move { + loop { + let received = tokio::select! { + signal = sigint.recv() => signal, + signal = sigterm.recv() => signal, + }; + if received.is_none() { + break; + } + let _ = signal_tx.send(()); + } + }); + } + + #[cfg(not(unix))] + { + let signal_tx = shutdown_tx.clone(); + tokio::spawn(async move { + loop { + if tokio::signal::ctrl_c().await.is_err() { + break; + } + let _ = signal_tx.send(()); + } + }); + // Give the spawned ctrl_c future one poll so its platform handler is + // registered before the caller is allowed to spawn an adapter. + tokio::task::yield_now().await; + } + + Ok((shutdown_tx, shutdown_rx)) +} + #[tokio::main] async fn tokio_main() -> Result<()> { // Install the ring crypto provider for rustls (required for wss:// connections). @@ -1248,7 +2384,8 @@ async fn tokio_main() -> Result<()> { .map(|(_, a)| a) .collect(); let args = ModelsArgs::parse_from(&filtered); - return run_models(args).await; + let (_shutdown_tx, shutdown_rx) = install_shutdown_channel().await?; + return run_models(args, shutdown_rx).await; } if is_subcommand("auth-methods") { @@ -1258,7 +2395,8 @@ async fn tokio_main() -> Result<()> { .map(|(_, a)| a) .collect(); let args = AuthMethodsArgs::parse_from(&filtered); - return run_auth_methods(args).await; + let (_shutdown_tx, shutdown_rx) = install_shutdown_channel().await?; + return run_auth_methods(args, shutdown_rx).await; } if is_subcommand("authenticate") { @@ -1268,7 +2406,8 @@ async fn tokio_main() -> Result<()> { .map(|(_, a)| a) .collect(); let args = AuthenticateArgs::parse_from(&filtered); - return run_authenticate(args).await; + let (_shutdown_tx, shutdown_rx) = install_shutdown_channel().await?; + return run_authenticate(args, shutdown_rx).await; } tracing_subscriber::fmt() @@ -1292,6 +2431,8 @@ async fn tokio_main() -> Result<()> { return setup_mode::run_setup_listener(config, payload).await; } + let (shutdown_tx, mut shutdown_rx) = install_shutdown_channel().await?; + tracing::info!("buzz-acp starting: {}", config.summary()); let observer = config @@ -1312,12 +2453,39 @@ async fn tokio_main() -> Result<()> { ); } - let mut pool = if config.lazy_pool { - AgentPool::from_slots((0..config.agents).map(|_| None).collect()) + let eager_pool_result = if config.lazy_pool { + None } else { - initialize_agent_pool(&PoolStartup::from_config(&config, observer.clone()), None).await? + Some( + initialize_agent_pool( + &PoolStartup::from_config(&config, observer.clone()), + Some(shutdown_rx.clone()), + ) + .await, + ) }; - let mut pool_ready = !config.lazy_pool; + let (mut pool, mut pool_ready, mut pool_cleanup_unverified, initial_cleanup_error) = + match resolve_initial_pool_startup(config.lazy_pool, eager_pool_result)? { + InitialPoolStartup::Ready(pool) => (pool, true, false, None), + InitialPoolStartup::Dormant => ( + AgentPool::from_slots((0..config.agents).map(|_| None).collect()), + false, + false, + None, + ), + InitialPoolStartup::CleanupQuarantine(error) => { + tracing::error!( + error, + "eager pool cleanup could not be verified — entering non-spawning quarantine" + ); + ( + AgentPool::from_slots((0..config.agents).map(|_| None).collect()), + false, + true, + Some(error), + ) + } + }; let mut pool_lifecycle: PoolLifecycle = PoolLifecycle::listening(); // Capture a startup watermark BEFORE connecting to the relay. This timestamp @@ -1332,16 +2500,37 @@ async fn tokio_main() -> Result<()> { let pubkey_hex = config.keys.public_key().to_hex(); + // Resolve one canonical privileged owner before the relay background task + // starts. This closes the handshake-buffer window where a configured + // fallback owner could otherwise bypass relay freshness/replay admission. + // Priority: BUZZ_AUTH_TAG (NIP-OA attestation) → --agent-owner flag. + let startup_owner: Option = resolve_agent_owner(&config); + if let Some(ref owner) = startup_owner { + tracing::info!("agent owner: {owner}"); + } else { + tracing::info!("no agent owner configured"); + } + // Parse BUZZ_AUTH_TAG into a nostr::Tag for NIP-OA relay membership delegation. let relay_auth_tag: Option = std::env::var("BUZZ_AUTH_TAG") .ok() .filter(|s| !s.is_empty()) .and_then(|s| buzz_sdk::nip_oa::parse_auth_tag(&s).ok()); - let mut relay = - HarnessRelay::connect(&config.relay_url, &config.keys, &pubkey_hex, relay_auth_tag) - .await - .map_err(|e| anyhow::anyhow!("relay connect error: {e}"))?; + let mut relay = await_pre_loop_operation( + HarnessRelay::connect_with_owner( + &config.relay_url, + &config.keys, + &pubkey_hex, + relay_auth_tag, + startup_owner.clone(), + ), + &mut shutdown_rx, + &mut pool, + "relay connect", + pool_cleanup_unverified, + ) + .await?; // Tell the relay background task the watermark so it can use // `since = watermark - 5s` on the first REQ instead of `since=now`. @@ -1353,22 +2542,19 @@ async fn tokio_main() -> Result<()> { tracing::info!("connected to relay at {}", config.relay_url); - relay - .subscribe_membership_notifications() - .await - .map_err(|e| anyhow::anyhow!("membership notification subscribe error: {e}"))?; + await_pre_loop_operation( + relay.subscribe_membership_notifications(), + &mut shutdown_rx, + &mut pool, + "membership notification subscribe", + pool_cleanup_unverified, + ) + .await?; tracing::info!("subscribed to membership notifications"); let presence_publisher = relay.event_publisher(); let presence_keys = config.keys.clone(); - // Priority: BUZZ_AUTH_TAG (NIP-OA attestation) → --agent-owner flag. - let startup_owner: Option = resolve_agent_owner(&config); - if let Some(ref owner) = startup_owner { - tracing::info!("agent owner: {owner}"); - } else { - tracing::info!("no agent owner configured"); - } // Warn if owner-dependent mode but no owner resolved yet. if startup_owner.is_none() { match &config.respond_to { @@ -1391,6 +2577,7 @@ async fn tokio_main() -> Result<()> { let owner_cache = OwnerCache::new(startup_owner.clone()); let mut relay_observer_control_rx = None; + let mut observer_control_dedup = ObserverControlDedup::new(OBSERVER_CONTROL_DEDUP_CAPACITY); let mut relay_observer_publisher_task = None; let mut relay_observer_publisher = None; if config.relay_observer { @@ -1407,10 +2594,14 @@ async fn tokio_main() -> Result<()> { owner_pubkey_hex, owner_pubkey, )); - relay - .subscribe_observer_controls() - .await - .map_err(|e| anyhow::anyhow!("observer control subscribe error: {e}"))?; + await_pre_loop_operation( + relay.subscribe_observer_controls(), + &mut shutdown_rx, + &mut pool, + "observer control subscribe", + pool_cleanup_unverified, + ) + .await?; relay_observer_control_rx = relay.take_observer_control_rx(); tracing::info!("relay observer enabled"); } @@ -1426,10 +2617,14 @@ async fn tokio_main() -> Result<()> { } } - let channel_info_map = relay - .discover_channels() - .await - .map_err(|e| anyhow::anyhow!("channel discovery error: {e}"))?; + let channel_info_map = await_pre_loop_operation( + relay.discover_channels(), + &mut shutdown_rx, + &mut pool, + "channel discovery", + pool_cleanup_unverified, + ) + .await?; tracing::info!("discovered {} channel(s)", channel_info_map.len()); let channel_ids: Vec = channel_info_map.keys().copied().collect(); @@ -1467,7 +2662,20 @@ async fn tokio_main() -> Result<()> { } SubscribeMode::Config => { // load_rules() already warns if the config file has zero rules. - config::load_rules(&config.config_path)? + match config::load_rules(&config.config_path) { + Ok(rules) => rules, + Err(error) => { + return fail_pre_loop_with_pool_cleanup( + &mut pool, + &mut shutdown_rx, + "subscription rule loading", + anyhow::anyhow!(error), + pool_cleanup_unverified, + false, + ) + .await; + } + } } }; @@ -1522,6 +2730,15 @@ async fn tokio_main() -> Result<()> { "listening", None, ); + } else if let Some(error) = initial_cleanup_error.as_deref() { + emit_runtime_lifecycle( + observer.as_ref(), + &runtime_start_nonce, + &pubkey_hex, + &config.relay_url, + "failed", + Some(error), + ); } let base_prompt_content = config.base_prompt_content.take(); @@ -1612,7 +2829,8 @@ async fn tokio_main() -> Result<()> { 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(); - let (wake_tx, mut wake_rx) = mpsc::channel::<(u32, Result)>(1); + let (wake_tx, mut wake_rx) = + mpsc::channel::<(u32, std::result::Result)>(1); let mut wake_tasks: tokio::task::JoinSet<()> = tokio::task::JoinSet::new(); // Channel for non-cancelling steer ack watchers to forward outcomes back @@ -1627,26 +2845,6 @@ async fn tokio_main() -> Result<()> { // `IN_FLIGHT_DEADLINE_SECS` expires. let (steer_ack_tx, mut steer_ack_rx) = mpsc::unbounded_channel::(); - // ── Step 7: Shutdown signal ─────────────────────────────────────────────── - let (shutdown_tx, mut shutdown_rx) = watch::channel(()); - - let tx = shutdown_tx.clone(); - tokio::spawn(async move { - tokio::signal::ctrl_c().await.ok(); - let _ = tx.send(()); - }); - - #[cfg(unix)] - { - let tx = shutdown_tx.clone(); - tokio::spawn(async move { - use tokio::signal::unix::{signal, SignalKind}; - let mut sigterm = signal(SignalKind::terminate()).expect("SIGTERM handler"); - sigterm.recv().await; - let _ = tx.send(()); - }); - } - // Track the newest membership notification timestamp per channel. // On reconnect the relay replays events newest-first, so the first event // per channel is authoritative. Any later event with ts < newest is stale. @@ -1685,11 +2883,7 @@ async fn tokio_main() -> Result<()> { // 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, - }) + .map(|_| SlotCircuit::new()) .collect(); // @@ -1700,7 +2894,7 @@ async fn tokio_main() -> Result<()> { Result(Box), Panic(tokio::task::JoinError), SteerAck(SteerAckEvent), - Wake(u32, Result), + Wake(u32, std::result::Result), } loop { @@ -1710,7 +2904,7 @@ async fn tokio_main() -> Result<()> { // unconditionally would complete instantly on every iteration — a // busy spin — whenever the queued work drained after a failed wake. let mut lazy_wake_work_pending = false; - if config.lazy_pool && !pool_ready { + if lazy_pool_can_wake(config.lazy_pool, pool_ready, pool_cleanup_unverified) { lazy_wake_work_pending = queue.has_flushable_work(); if let Some(attempt) = pool_lifecycle .start_wake_if_due(lazy_wake_work_pending, tokio::time::Instant::now()) @@ -1727,9 +2921,7 @@ async fn tokio_main() -> Result<()> { let wake_tx = wake_tx.clone(); let wake_shutdown = shutdown_rx.clone(); wake_tasks.spawn(async move { - let result = initialize_agent_pool(&startup, Some(wake_shutdown)) - .await - .map_err(|error| error.to_string()); + let result = initialize_agent_pool(&startup, Some(wake_shutdown)).await; if let Err(error) = wake_tx.send((attempt, result)).await { let (_attempt, result) = error.0; if let Ok(mut abandoned_pool) = result { @@ -1761,9 +2953,24 @@ async fn tokio_main() -> Result<()> { let env = config.persona_env_vars.clone(); let has_codex = config.has_generated_codex_config; let observer = observer.clone(); - let guard = RespawnGuard::new(idx, respawn_tx.clone()); + let pending_intent = slot.pending_model_intent.clone(); + let (desired_model, model_overridden, retain_model_intent) = match pending_intent { + Some(intent) => (intent.desired_model, intent.model_overridden, true), + None => (config.model.clone(), false, false), + }; + let mut guard = RespawnGuard::new( + idx, + respawn_tx.clone(), + desired_model, + model_overridden, + retain_model_intent, + ); respawn_tasks.spawn(async move { - let result = spawn_and_init(&cmd, &args, &env, has_codex, idx, observer).await; + // This refill starts from an already-empty slot. + guard.mark_owner_cleanup_verified(); + let result = + spawn_and_init(&cmd, &args, &env, has_codex, idx, observer, &mut guard) + .await; guard.send(result); }); } @@ -1780,29 +2987,74 @@ async fn tokio_main() -> Result<()> { } } + // Reap every completed replacement task before entering the biased + // select. The select arm below is the quiet-loop wake path; this eager + // drain prevents sustained ready relay/result traffic from starving + // that lower-priority arm and retaining completed JoinSet entries. + while let Some(completed) = respawn_tasks.try_join_next() { + if let Err(error) = completed { + // RespawnGuard::Drop already enqueues the matching failure. + tracing::error!("respawn task failed: {error}"); + } + } + let mut respawn_collected = false; while let Ok(rr) = respawn_rx.try_recv() { - crash_history[rr.index].respawn_in_flight = false; + let slot = &mut crash_history[rr.index]; + slot.respawn_in_flight = false; match rr.result { - Ok((acp, protocol_version, agent_name)) => { + Ok(Some((acp, protocol_version, agent_name))) => { + let pending_model_intent = slot.pending_model_intent.take(); + let model_switch_request_id = pending_model_intent + .as_ref() + .and_then(|intent| intent.model_switch_request_id.clone()); + let model_switch_rollback = + pending_model_intent.and_then(|intent| intent.model_switch_rollback); let agent = OwnedAgent { index: rr.index, acp, state: SessionState::default(), model_capabilities: None, - desired_model: config.model.clone(), - model_overridden: false, + desired_model: rr.desired_model, + model_overridden: rr.model_overridden, + model_switch_request_id, + model_switch_rollback, agent_name, goose_system_prompt_supported: None, protocol_version, }; pool.return_agent(agent); + queue.release_required_agent(rr.index); tracing::info!(agent = rr.index, "respawn complete"); respawn_collected = true; } + Ok(None) => { + if !rr.retain_model_intent { + slot.pending_model_intent = None; + } + tracing::warn!( + agent = rr.index, + "old adapter shutdown complete; circuit remains open" + ); + } Err(e) => { - crash_history[rr.index].mark_spawn_failed(); - tracing::warn!(agent = rr.index, "respawn failed: {e} — circuit re-opened"); + if rr.terminal_switch_failure { + slot.rollback_terminal_switch_failure(); + } else if !rr.retain_model_intent { + slot.pending_model_intent = None; + } + if rr.failure_class == Some(RespawnFailureClass::Spawn) { + slot.mark_spawn_failed(); + tracing::warn!(agent = rr.index, "respawn failed: {e} — circuit re-opened"); + } else { + // Missing classification also fails closed. Internal + // bookkeeping ambiguity is not cleanup proof. + slot.mark_cleanup_unverified(); + tracing::error!( + agent = rr.index, + "adapter cleanup unverified: {e} — slot quarantined until process restart" + ); + } } } } @@ -1826,6 +3078,15 @@ async fn tokio_main() -> Result<()> { Some(result) => Some(PoolEvent::Result(Box::new(result))), None => { tracing::info!("result channel closed — exiting main loop"); + await_automatic_exit_permission( + &mut shutdown_rx, + cleanup_ownership_unverified( + pool_cleanup_unverified, + &crash_history, + ), + "result channel closed", + ) + .await; break; } }, @@ -1859,6 +3120,11 @@ async fn tokio_main() -> Result<()> { } } => None, Some(Err(error)) = wake_tasks.join_next(), if !wake_tasks.is_empty() => { + // A panicked/cancelled startup task may have dropped a + // partially spawned client without verified shutdown. + // Keep this supervisor process from attempting another + // lazy wake. + pool_cleanup_unverified = true; if let Some(attempt) = pool_lifecycle.waking_attempt() { let message = format!("pool wake task failed: {error}"); if pool_lifecycle.cancel_wake( @@ -1878,6 +3144,22 @@ async fn tokio_main() -> Result<()> { } None } + completed = respawn_tasks.join_next(), if !respawn_tasks.is_empty() => { + if let Some(Err(error)) = completed { + // RespawnGuard::Drop already sends the matching failure + // result, so this arm only reports and reaps the task. + tracing::error!("respawn task failed: {error}"); + } + None + } + _ = tokio::time::sleep_until(tokio::time::Instant::from_std( + last_maintenance + maintenance_interval + )), if pool_ready => { + // Wake the loop so cooldown probes, retry throttles, and + // empty-slot refill progress even when every optional + // heartbeat and the relay are quiet. + None + } control_event = async { match relay_observer_control_rx.as_mut() { Some(rx) => rx.recv().await, @@ -1888,7 +3170,19 @@ async fn tokio_main() -> Result<()> { match control_event { Some(event) => { if let Some(ref owner_hex) = owner_cache.pubkey { - handle_relay_observer_control_event(&config.keys, event, &mut pool, observer.as_ref(), owner_hex); + handle_relay_observer_control_event( + &config.keys, + event, + &mut observer_control_dedup, + &mut pool, + &mut queue, + &config, + &mut crash_history, + &respawn_tx, + &mut respawn_tasks, + observer.as_ref(), + owner_hex, + ); } else { tracing::warn!("observer control frame received but no owner resolved — dropping"); } @@ -1982,18 +3276,25 @@ async fn tokio_main() -> Result<()> { if let Err(e) = relay.unsubscribe_channel(ch).await { tracing::warn!("failed to unsubscribe from channel {ch}: {e}"); } - // Drain queued events and invalidate sessions for the - // removed channel. Events already in-flight will - // complete normally (the relay may reject actions if - // the agent lost access). + // Drain queued events and recycle idle adapters that + // still own the removed channel. Checked-out adapters + // are recycled when they return. let drained_ids = queue.drain_channel(ch); - let invalidated = if pool_ready { - pool.invalidate_channel_sessions(ch) + let recycled = if pool_ready { + recycle_idle_channel_sessions( + &mut pool, + ch, + &config, + &mut crash_history, + &respawn_tx, + &mut respawn_tasks, + observer.clone(), + ) } else { 0 }; - // Track removed channels so checked-out agents get - // their sessions stripped when they return to the pool. + // Track removed channels so checked-out adapters are + // recycled instead of locally forgetting ownership. removed_channels.insert(ch); typing_channels.remove(&ch); // Best-effort: clean up 👀 on drained events. @@ -2011,11 +3312,11 @@ async fn tokio_main() -> Result<()> { } }); } - if !drained_ids.is_empty() || invalidated > 0 { + if !drained_ids.is_empty() || recycled > 0 { tracing::info!( channel_id = %ch, drained = drained_ids.len(), - invalidated, + recycled, "cleaned up after membership removal" ); } @@ -2029,11 +3330,12 @@ async fn tokio_main() -> Result<()> { } // Check: kind:9, content "!shutdown", from owner, mentions THIS agent. - let is_shutdown = is_owner_control_command( + let is_shutdown = is_admitted_owner_control_command( &buzz_event.event, kind_u32, "!shutdown", &pubkey_hex, + buzz_event.privileged_control_admitted, ); if is_shutdown { let owner = owner_cache.get(); @@ -2060,25 +3362,31 @@ async fn tokio_main() -> Result<()> { // Mode-independent: !cancel fires regardless of // --multiple-event-handling. It is explicit user // intent, not an automatic policy decision. - let is_cancel = is_owner_control_command( + let is_cancel = is_admitted_owner_control_command( &buzz_event.event, kind_u32, "!cancel", &pubkey_hex, + buzz_event.privileged_control_admitted, ); if is_cancel { if let Some(owner) = owner_cache.get() { if buzz_event.event.pubkey.to_hex() == *owner { - let fired = signal_in_flight_task( + let control_result = signal_in_flight_task( &mut pool, buzz_event.channel_id, ControlSignal::Cancel, ); - if !fired { - tracing::warn!( + match control_result { + ControlSignalResult::Delivered => {} + ControlSignalResult::DropRecorded => tracing::info!( + channel_id = %buzz_event.channel_id, + "!cancel received after a prior control — batch drop recorded" + ), + ControlSignalResult::NotAccepted => tracing::warn!( channel_id = %buzz_event.channel_id, "!cancel received but no in-flight task — no-op" - ); + ), } continue; // consume event — do NOT push to queue } @@ -2098,32 +3406,46 @@ async fn tokio_main() -> Result<()> { // returns. If idle, invalidate the cached channel // session immediately. Queued future events remain // queued and will create a fresh session on dispatch. - let is_rotate = is_owner_control_command( + let is_rotate = is_admitted_owner_control_command( &buzz_event.event, kind_u32, "!rotate", &pubkey_hex, + buzz_event.privileged_control_admitted, ); if is_rotate { if let Some(owner) = owner_cache.get() { if buzz_event.event.pubkey.to_hex() == *owner { - let fired = signal_in_flight_task( + let control_result = signal_in_flight_task( &mut pool, buzz_event.channel_id, ControlSignal::Rotate, ); - if fired { - tracing::info!( + match control_result { + ControlSignalResult::Delivered => tracing::info!( channel_id = %buzz_event.channel_id, "!rotate received — cancelling in-flight turn and rotating session" - ); - } else { - let invalidated = pool.invalidate_channel_sessions(buzz_event.channel_id); - tracing::info!( + ), + ControlSignalResult::DropRecorded => tracing::info!( channel_id = %buzz_event.channel_id, - invalidated, - "!rotate received — invalidated idle channel session(s)" - ); + "!rotate received after a prior control — batch drop and rotation recorded" + ), + ControlSignalResult::NotAccepted => { + let recycled = recycle_idle_channel_sessions( + &mut pool, + buzz_event.channel_id, + &config, + &mut crash_history, + &respawn_tx, + &mut respawn_tasks, + observer.clone(), + ); + tracing::info!( + channel_id = %buzz_event.channel_id, + recycled, + "!rotate received — recycling idle channel session owner(s)" + ); + } } continue; // consume event — do NOT push to queue } @@ -2194,7 +3516,7 @@ async fn tokio_main() -> Result<()> { // backed payload) so the cost is negligible. let event_for_steer = buzz_event.event.clone(); let prompt_tag_for_steer = prompt_tag.clone(); - let accepted = queue.push(QueuedEvent { + let occurrence_id = queue.push(QueuedEvent { channel_id: buzz_event.channel_id, event: buzz_event.event, received_at: std::time::Instant::now(), @@ -2205,7 +3527,7 @@ async fn tokio_main() -> Result<()> { // Fire-and-forget: on rare fast-failure paths the // guard's cleanup may race with this add, leaving a // cosmetic stale 👀. Acceptable — see ReactionGuard docs. - if accepted { + if occurrence_id.is_some() { let rc = ctx.rest_client.clone(); let eid = event_id_hex.clone(); tokio::spawn(async move { @@ -2215,7 +3537,8 @@ async fn tokio_main() -> Result<()> { // Event is already queued. If mode requires it AND // the channel has an in-flight task, fire cancel — // OR take the non-cancelling (ACP steer) fork for Steer signals. - if accepted && queue.is_channel_in_flight(buzz_event.channel_id) { + if let Some(occurrence_id) = occurrence_id { + if queue.is_channel_in_flight(buzz_event.channel_id) { // Author eligibility (owner ∪ allowlist ∪ siblings) // is already enforced by the inbound author gate // above, so the mid-turn signal fires for every @@ -2243,6 +3566,7 @@ async fn tokio_main() -> Result<()> { &mut pool, &mut queue, buzz_event.channel_id, + occurrence_id, event_for_steer, prompt_tag_for_steer, &steer_ack_tx, @@ -2255,6 +3579,7 @@ async fn tokio_main() -> Result<()> { ); } } + } } if pool_ready { for (channel_id, thread_tags) in @@ -2269,6 +3594,15 @@ async fn tokio_main() -> Result<()> { if let Err(e) = relay.reconnect().await { tracing::error!("relay background task is gone: {e} — exiting"); tokio::time::sleep(Duration::from_secs(1)).await; + await_automatic_exit_permission( + &mut shutdown_rx, + cleanup_ownership_unverified( + pool_cleanup_unverified, + &crash_history, + ), + "relay reconnect failed", + ) + .await; break; } } @@ -2368,6 +3702,12 @@ async fn tokio_main() -> Result<()> { Some(&ctx.rest_client), ) == LoopAction::Exit { + await_automatic_exit_permission( + &mut shutdown_rx, + cleanup_ownership_unverified(pool_cleanup_unverified, &crash_history), + "prompt result requested exit", + ) + .await; break; } if drain_ready_join_results( @@ -2381,8 +3721,15 @@ async fn tokio_main() -> Result<()> { &respawn_tx, &mut respawn_tasks, observer.clone(), + Some(&ctx.rest_client), ) == LoopAction::Exit { + await_automatic_exit_permission( + &mut shutdown_rx, + cleanup_ownership_unverified(pool_cleanup_unverified, &crash_history), + "drained prompt result requested exit", + ) + .await; break; } for (channel_id, thread_tags) in dispatch_pending(&mut pool, &mut queue, &ctx) { @@ -2403,9 +3750,16 @@ async fn tokio_main() -> Result<()> { &respawn_tx, &mut respawn_tasks, observer.clone(), + Some(&ctx.rest_client), ); if pool.live_count() == 0 && !any_respawn_in_flight(&crash_history) { tracing::error!("all agents dead — exiting"); + await_automatic_exit_permission( + &mut shutdown_rx, + cleanup_ownership_unverified(pool_cleanup_unverified, &crash_history), + "all agents dead", + ) + .await; break; } for (channel_id, thread_tags) in dispatch_pending(&mut pool, &mut queue, &ctx) { @@ -2414,6 +3768,7 @@ async fn tokio_main() -> Result<()> { } Some(PoolEvent::SteerAck(SteerAckEvent { channel_id, + occurrence_id, event_id, ack, })) => { @@ -2530,10 +3885,10 @@ async fn tokio_main() -> Result<()> { queue.extend_in_flight_deadline(channel_id, config.max_turn_duration_secs); } if drop_withheld { - queue.remove_event(channel_id, &event_id); + queue.remove_native_steer(channel_id, occurrence_id); } if release_withheld { - queue.release_native_steer(channel_id, &event_id); + queue.release_native_steer(channel_id, occurrence_id); } if signal_fallback { // Universal cancel+merge fallback. Note: the @@ -2555,6 +3910,14 @@ async fn tokio_main() -> Result<()> { } } Some(PoolEvent::Wake(attempt, result)) => { + let cleanup_unverified = result + .as_ref() + .err() + .is_some_and(PoolStartupFailure::cleanup_unverified); + if cleanup_unverified { + pool_cleanup_unverified = true; + } + let result = result.map_err(|error| error.to_string()); let completion = result.as_ref().map(|_| ()).map_err(|error| error.clone()); if let Err(error) = pool_lifecycle.complete_wake(attempt, result, tokio::time::Instant::now()) @@ -2608,6 +3971,12 @@ async fn tokio_main() -> Result<()> { // signal handlers (result channel closed, LoopAction::Exit) cancel the wake // just as promptly. Timeout is a backstop for a slot stuck outside the // select (e.g. in spawn); only then do we fall back to aborting. + let mut shutdown_verification = ShutdownVerification::default(); + shutdown_verification.record(ShutdownOwner::Wake, !pool_cleanup_unverified); + shutdown_verification.record( + ShutdownOwner::Respawn, + !crash_history.iter().any(|slot| slot.cleanup_unverified), + ); let _ = shutdown_tx.send(()); let wake_drain = tokio::time::timeout(Duration::from_secs(30), async { while wake_tasks.join_next().await.is_some() {} @@ -2615,80 +3984,169 @@ async fn tokio_main() -> Result<()> { .await; if wake_drain.is_err() { tracing::warn!("wake task did not drain within grace period — aborting"); + shutdown_verification.record(ShutdownOwner::Wake, false); wake_tasks.shutdown().await; } while let Ok((_attempt, result)) = wake_rx.try_recv() { - if let Ok(mut awakened_pool) = result { - shutdown_agent_pool(&mut awakened_pool).await; + match result { + Ok(mut awakened_pool) => { + let verified = shutdown_agent_pool(&mut awakened_pool).await; + shutdown_verification.record(ShutdownOwner::Wake, verified); + } + Err(PoolStartupFailure::CleanupUnverified(_)) => { + shutdown_verification.record(ShutdownOwner::Wake, false); + } + Err(PoolStartupFailure::Invariant(error)) => { + tracing::error!(error, "wake pool startup invariant failed"); + shutdown_verification.record(ShutdownOwner::Wake, false); + } + Err(PoolStartupFailure::Retryable(_)) => {} } } - tracing::info!("shutdown: waiting for in-flight prompts"); - // 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 signalled = pool.cancel_checked_out_for_shutdown(); + tracing::info!(signalled, "shutdown: cancelling in-flight prompts"); + // Preserve the typed process owner until every task returns it. Aborting a + // task drops the only AcpClient handle and makes reaping/process-group + // absence impossible to prove. let (rx_ref, js_ref) = pool.rx_and_join_set(); - let shutdown_result = tokio::time::timeout(grace, async { - 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 + let mut checked_out_cleanup_verified = true; + loop { + tokio::select! { + result = js_ref.join_next() => { + match result { + Some(Err(e)) => { + tracing::warn!("task error during shutdown: {e}"); + checked_out_cleanup_verified = false; } + Some(Ok(())) => {} + None => break, } - maybe_result = rx_ref.recv() => { - if let Some(mut pr) = maybe_result { - let idx = pr.agent.index; - pr.agent.acp.shutdown().await; + } + maybe_result = rx_ref.recv() => { + if let Some(mut pr) = maybe_result { + let idx = pr.agent.index; + if shutdown_acp_with_log( + &mut pr.agent.acp, + Some(idx), + "checked-out agent shutdown", + ) + .await + { tracing::debug!(agent = idx, "reaped checked-out agent on shutdown"); + } else { + checked_out_cleanup_verified = false; } - // If None, channel closed — tasks are done. } } + _ = shutdown_rx.changed() => { + // A repeated SIGINT/SIGTERM does not abort the task and lose + // its sole typed process owner. Grace has already been skipped: + // every returned owner is immediately hard-shut down, reaped, + // and process-group-probed in the result arm above. + tracing::warn!( + "repeated shutdown signal received — hard cleanup active; \ + preserving typed ownership until verification completes" + ); + } } - }) - .await; - if shutdown_result.is_err() { - 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. + // Drain results that raced with the final join completion. 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"); + if shutdown_acp_with_log( + &mut pr.agent.acp, + Some(idx), + "late checked-out agent shutdown", + ) + .await + { + tracing::debug!(agent = idx, "reaped late-arriving agent on shutdown"); + } else { + checked_out_cleanup_verified = false; + } } + shutdown_verification.record(ShutdownOwner::CheckedOut, checked_out_cleanup_verified); // Explicitly shut down idle agents still sitting in their slots. + let mut idle_cleanup_verified = true; 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"); + if shutdown_acp_with_log(&mut acp, Some(idx), "idle agent shutdown").await { + tracing::debug!(agent = idx, "reaped idle agent on shutdown"); + } else { + idle_cleanup_verified = false; + } } } + shutdown_verification.record(ShutdownOwner::Idle, idle_cleanup_verified); 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. + // Track the exact slots whose ownership result has not reached the + // supervisor. JoinSet length is not a safe proxy: a task can enqueue its + // result and remain briefly joinable, which would make shutdown wait for a + // second result that will never exist. + let mut pending_respawn_results: HashSet = crash_history + .iter() + .enumerate() + .filter_map(|(index, slot)| slot.respawn_in_flight.then_some(index)) + .collect(); + let mut shutdown_respawn_results = Vec::new(); + while let Ok(result) = respawn_rx.try_recv() { + pending_respawn_results.remove(&result.index); + shutdown_respawn_results.push(result); + } + + // Abort any in-flight respawn tasks. Backoff-phase guards report a clean + // no-owner cancellation synchronously. Initializing guards transfer their + // AcpClient to a bounded cleanup task and report only after shutdown proof. 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"); + // AcpClient::shutdown may use two five-second child waits plus one + // five-second process-group absence probe. All replacement cleanups run + // concurrently, so one shared 16-second deadline covers their bounded + // verification path. + let cleanup_deadline = tokio::time::Instant::now() + Duration::from_secs(16); + while !pending_respawn_results.is_empty() { + match tokio::time::timeout_at(cleanup_deadline, respawn_rx.recv()).await { + Ok(Some(result)) => { + pending_respawn_results.remove(&result.index); + shutdown_respawn_results.push(result); + } + Ok(None) | Err(_) => { + shutdown_verification.record(ShutdownOwner::Respawn, false); + tracing::error!( + pending_slots = ?pending_respawn_results, + "timed out waiting for cancelled respawn ownership verification" + ); + break; + } + } + } + while let Ok(result) = respawn_rx.try_recv() { + shutdown_respawn_results.push(result); + } + + // Explicitly shut down successfully returned agents instead of relying on + // AcpClient::Drop. + for rr in shutdown_respawn_results { + match rr.result { + Ok(Some((mut acp, _, _))) => { + let verified = + shutdown_acp_with_log(&mut acp, Some(rr.index), "respawned agent shutdown") + .await; + shutdown_verification.record(ShutdownOwner::Respawn, verified); + if verified { + tracing::debug!(agent = rr.index, "reaped respawned agent on shutdown"); + } + } + Ok(None) => {} + Err(_) => { + if rr.failure_class == Some(RespawnFailureClass::CleanupUnverified) { + shutdown_verification.record(ShutdownOwner::Respawn, false); + } + } } } @@ -2711,14 +4169,36 @@ async fn tokio_main() -> Result<()> { } } + // Close the last process-level observer sender, then give the publisher a + // bounded chance to drain every terminal result into the relay's + // non-evicting priority state. A timeout or rejected admission is a + // shutdown verification failure, never a silent proof loss. + drop(observer); if let Some(handle) = relay_observer_publisher_task.take() { - handle.abort(); + let abort_handle = handle.abort_handle(); + match tokio::time::timeout(Duration::from_secs(5), handle).await { + Ok(Ok(true)) => {} + Ok(Ok(false)) => { + tracing::error!("terminal observer result delivery failed before shutdown"); + shutdown_verification.record(ShutdownOwner::ObserverPublisher, false); + } + Ok(Err(error)) => { + tracing::error!("observer publisher task failed during shutdown: {error}"); + shutdown_verification.record(ShutdownOwner::ObserverPublisher, false); + } + Err(_) => { + tracing::error!("observer publisher did not drain within 5s"); + abort_handle.abort(); + shutdown_verification.record(ShutdownOwner::ObserverPublisher, false); + } + } } // 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; + shutdown_verification.into_result()?; tracing::info!("buzz-acp stopped"); Ok(()) } @@ -2747,6 +4227,17 @@ fn is_owner_control_command( && event_mentions_agent(event, agent_pubkey_hex) } +fn is_admitted_owner_control_command( + event: &nostr::Event, + kind_u32: u32, + command: &str, + agent_pubkey_hex: &str, + privileged_control_admitted: bool, +) -> bool { + privileged_control_admitted + && is_owner_control_command(event, kind_u32, command, agent_pubkey_hex) +} + // ── signal_in_flight_task ───────────────────────────────────────────────────── /// Decide which [`ControlSignal`] (if any) to send to an in-flight turn when a @@ -2775,26 +4266,90 @@ fn mode_gate_signal( } } -/// Send a control signal to the in-flight task for `channel_id`. -/// Returns `true` if a signal was sent, `false` if no in-flight task was found. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum ControlSignalResult { + /// The task's control receiver accepted the signal. + Delivered, + /// A prior control already consumed (or closed) the receiver, but an + /// explicit owner Cancel/Rotate disposition was recorded at the supervisor + /// boundary and will dominate result/panic batch fate. + DropRecorded, + /// No matching task exists, or this non-drop control could not be sent. + NotAccepted, +} + +fn record_drop_control(meta: &mut pool::TaskMeta, control: AcceptedDropControl) { + if meta.accepted_drop_control != Some(AcceptedDropControl::Rotate) { + meta.accepted_drop_control = Some(control); + } +} + +/// Deliver a control signal or, for explicit Cancel/Rotate, persist its +/// stronger supervisor-owned drop disposition while the task still exists. fn signal_in_flight_task( pool: &mut AgentPool, channel_id: uuid::Uuid, mode: ControlSignal, -) -> bool { +) -> ControlSignalResult { let entry = pool .task_map_mut() .values_mut() .find(|m| m.channel_id == Some(channel_id)); - if let Some(meta) = entry { - if let Some(tx) = meta.control_tx.take() { - tracing::info!(channel = %channel_id, ?mode, "control signal sent to in-flight task"); - let _ = tx.send(mode); - return true; + let Some(meta) = entry else { + return ControlSignalResult::NotAccepted; + }; + let accepted_drop_control = match &mode { + ControlSignal::Cancel => Some(AcceptedDropControl::Cancel), + ControlSignal::Rotate => Some(AcceptedDropControl::Rotate), + ControlSignal::Steer | ControlSignal::Interrupt | ControlSignal::SwitchModel(_) => None, + }; + let accepted_model_switch = match &mode { + ControlSignal::SwitchModel(request) => Some(request.clone()), + ControlSignal::Steer + | ControlSignal::Interrupt + | ControlSignal::Cancel + | ControlSignal::Rotate => None, + }; + + let Some(tx) = meta.control_tx.take() else { + if let Some(control) = accepted_drop_control { + record_drop_control(meta, control); + return ControlSignalResult::DropRecorded; + } + return ControlSignalResult::NotAccepted; + }; + + match tx.send(mode) { + Ok(()) => { + if let Some(control) = accepted_drop_control { + record_drop_control(meta, control); + } + if let Some(request) = accepted_model_switch { + meta.desired_model = Some(request.model_id.clone()); + meta.model_overridden = true; + meta.accepted_model_switch = Some(request); + } + tracing::info!( + channel = %channel_id, + "control signal sent to in-flight task" + ); + ControlSignalResult::Delivered + } + Err(mode) => { + tracing::debug!( + channel = %channel_id, + ?mode, + "in-flight task stopped accepting control signals" + ); + if let Some(control) = accepted_drop_control { + record_drop_control(meta, control); + ControlSignalResult::DropRecorded + } else { + ControlSignalResult::NotAccepted + } } } - false } /// Attempt the non-cancelling (ACP) steer for a freshly-queued event. @@ -2825,6 +4380,7 @@ fn try_native_steer( pool: &mut AgentPool, queue: &mut EventQueue, channel_id: uuid::Uuid, + occurrence_id: EnqueueOccurrenceId, event: nostr::Event, prompt_tag: String, steer_ack_tx: &mpsc::UnboundedSender, @@ -2865,7 +4421,7 @@ fn try_native_steer( // clears `in_flight_channels` and a stray `flush_next` could // re-deliver the event via normal dispatch. See // `EventQueue::mark_native_steer_pending` docs at queue.rs:606. - let withheld = queue.mark_native_steer_pending(channel_id, &event_id_hex); + let withheld = queue.mark_native_steer_pending(channel_id, occurrence_id); if !withheld { // Race: the event was already drained out of the queue // before we got here (e.g. a concurrent flush picked it @@ -2887,6 +4443,7 @@ fn try_native_steer( let ack = ack_rx.await; let _ = ack_tx_clone.send(SteerAckEvent { channel_id, + occurrence_id, event_id: event_id_for_watcher, ack, }); @@ -2906,6 +4463,12 @@ fn try_native_steer( // ── dispatch_pending ────────────────────────────────────────────────────────── +fn defer_unpinned_batch_for_capacity(queue: &mut EventQueue, batch: FlushBatch) { + let channel_id = batch.channel_id; + queue.requeue_preserve_timestamps(batch); + queue.mark_deferred(channel_id); +} + /// Flush queued work to available agents. fn dispatch_pending( pool: &mut AgentPool, @@ -2924,14 +4487,33 @@ fn dispatch_pending( .last() .map(|event| queue::parse_thread_tags(&event.event)) .unwrap_or_default(); - let affinity_hit = pool.has_session_for(channel_id); - let mut agent = match pool.try_claim(Some(channel_id)) { + let required_agent = queue.required_agent(channel_id); + let affinity_hit = required_agent.is_some() || pool.has_session_for(channel_id); + let mut agent = match required_agent + .and_then(|index| pool.try_claim_index(index)) + .or_else(|| { + if required_agent.is_none() { + pool.try_claim(Some(channel_id)) + } else { + None + } + }) { Some(a) => a, None => { + if let Some(required_agent) = required_agent { + tracing::debug!( + channel = %channel_id, + agent = required_agent, + "required agent slot unavailable — deferring pinned batch" + ); + queue.requeue_preserve_timestamps(batch); + queue.block_required_agent(channel_id); + queue.mark_deferred(channel_id); + continue; + } let pending = queue.pending_channels(); tracing::debug!(pending_channels = pending, "pool_exhausted"); - queue.requeue_preserve_timestamps(batch); - queue.mark_complete(channel_id); + defer_unpinned_batch_for_capacity(queue, batch); break; } }; @@ -2945,6 +4527,8 @@ fn dispatch_pending( let result_tx = pool.result_tx(); let ctx_clone = Arc::clone(ctx); let agent_index = agent.index; + let desired_model = agent.desired_model.clone(); + let model_overridden = agent.model_overridden; // Mid-turn non-cancelling steer seam: install the per-turn steer // receiver on the read loop so the main loop's mode-gate fork @@ -2985,6 +4569,10 @@ fn dispatch_pending( channel_id: Some(channel_id), turn_id, recoverable_batch, + desired_model, + model_overridden, + accepted_model_switch: None, + accepted_drop_control: None, control_tx: Some(control_tx), steer_tx, }, @@ -3051,6 +4639,47 @@ fn spawn_failure_notice( } } +fn acp_error_class(error: &acp::AcpError) -> &'static str { + match error { + acp::AcpError::Io(_) => "io", + acp::AcpError::Json(_) => "json", + acp::AcpError::AgentExited => "agent_exited", + acp::AcpError::IdleTimeout(_) => "idle_timeout", + acp::AcpError::HardTimeout { .. } => "hard_timeout", + acp::AcpError::CancelDrainTimeout(_) => "cancel_drain_timeout", + acp::AcpError::Timeout(_) => "timeout", + acp::AcpError::WriteTimeout(_) => "write_timeout", + acp::AcpError::Protocol(_) => "protocol", + acp::AcpError::AgentError { .. } => "agent_error", + } +} + +fn outcome_replaces_adapter(outcome: &PromptOutcome) -> bool { + matches!( + outcome, + PromptOutcome::AgentExited + | PromptOutcome::Timeout(_) + | PromptOutcome::CancelDrainTimeout(_) + | PromptOutcome::CancelCleanupFailed(_) + | PromptOutcome::SessionCloseFailed(_) + | PromptOutcome::SessionRecycleRequired + | PromptOutcome::SessionSetupCloseFailed(_) + | PromptOutcome::SessionSetupRecycleRequired + | PromptOutcome::ControlPreemptedSetup + ) || matches!( + outcome, + PromptOutcome::Error(error) if error.requires_process_replacement() + ) +} + +fn outcome_releases_source_session(outcome: &PromptOutcome) -> bool { + outcome_replaces_adapter(outcome) + || matches!( + outcome, + PromptOutcome::SessionRetired(_) | PromptOutcome::Cancelled + ) +} + #[allow(clippy::too_many_arguments)] fn handle_prompt_result( pool: &mut AgentPool, @@ -3067,9 +4696,86 @@ fn handle_prompt_result( ) -> LoopAction { let before = pool.task_map().len(); let agent_index = result.agent.index; - pool.task_map_mut() - .retain(|_, meta| meta.agent_index != agent_index); + let completed_task_id = pool + .task_map() + .iter() + .find_map(|(task_id, meta)| (meta.agent_index == agent_index).then_some(*task_id)); + let completed_meta = completed_task_id.and_then(|task_id| pool.task_map_mut().remove(&task_id)); debug_assert_eq!(before, pool.task_map().len() + 1); + let accepted_drop_control = completed_meta + .as_ref() + .and_then(|meta| meta.accepted_drop_control); + let mut recycle_for_unconsumed_control = completed_meta + .as_ref() + .is_some_and(|meta| meta.accepted_drop_control == Some(AcceptedDropControl::Rotate)) + && !outcome_releases_source_session(&result.outcome); + if result.retry_agent_index.is_none() { + if let Some(request) = completed_meta + .as_ref() + .and_then(|meta| meta.accepted_model_switch.as_ref()) + { + // The main loop successfully delivered SwitchModel, but the prompt + // future may have won the task's biased simultaneous-ready race. + // Reuse the task-side catalog gate here: this returned agent still + // owns both the prior intent and the capabilities that accepted or + // reject the request. Only an accepted model may pin the replay and + // recycle the still-owning adapter. + if pool::stage_busy_model_switch_intent(&mut result.agent, request) { + result.retry_agent_index = Some(agent_index); + recycle_for_unconsumed_control |= !outcome_releases_source_session(&result.outcome); + tracing::info!( + agent = agent_index, + model = %request.model_id, + "reconciled acknowledged model switch after prompt/control ready race" + ); + } else { + tracing::warn!( + agent = agent_index, + model = %request.model_id, + "rejected unsupported model switch after prompt/control ready race" + ); + } + } + } + let removed_session_requires_recycle = removed_channels + .iter() + .any(|channel_id| result.agent.state.sessions.contains_key(channel_id)) + && !outcome_replaces_adapter(&result.outcome); + let lifecycle_recycle_required = + recycle_for_unconsumed_control || removed_session_requires_recycle; + if lifecycle_recycle_required && result.agent.model_overridden { + result.retry_agent_index.get_or_insert(agent_index); + } + let result_channel = match &result.source { + PromptSource::Channel(channel_id) => Some(*channel_id), + PromptSource::Heartbeat => None, + }; + let existing_required_agent = + result_channel.and_then(|channel_id| queue.required_agent(channel_id)); + let replacement_preserves_override = result.agent.model_switch_request_id.is_some() + || (result.agent.model_overridden + && matches!( + result.outcome, + PromptOutcome::SessionRetired(_) + | PromptOutcome::Cancelled + | PromptOutcome::CancelDrainTimeout(_) + | PromptOutcome::CancelCleanupFailed(_) + | PromptOutcome::SessionCloseFailed(_) + | PromptOutcome::SessionRecycleRequired + | PromptOutcome::SessionSetupCloseFailed(_) + | PromptOutcome::SessionSetupRecycleRequired + | PromptOutcome::ControlPreemptedSetup + )); + let established_required_agent = result + .retry_agent_index + .or(replacement_preserves_override.then_some(agent_index)); + let preserve_model_on_fatal_replacement = + established_required_agent.is_some() || existing_required_agent == Some(agent_index); + if let (Some(channel_id), Some(required_agent)) = (result_channel, established_required_agent) { + if !removed_channels.contains(&channel_id) { + queue.require_agent(channel_id, required_agent); + } + } // The hard-timeout death_message (below) must describe the batch's // *actual* fate, not just the `recently_active` eligibility flag — a @@ -3079,6 +4785,8 @@ fn handle_prompt_result( // branch below records what actually happened; only the hard-timeout // match arm in the death_message construction reads it. let mut hard_timeout_fate_suffix: Option<&'static str> = None; + let mut batch_requeued = false; + let mut batch_terminally_disposed = false; // Requeue BEFORE mark_complete: requeue() sets retry_after with a future // deadline, and mark_complete() checks for it to decide whether to preserve @@ -3086,12 +4794,26 @@ fn handle_prompt_result( // every retry starts at attempt 1 — defeating exponential backoff and // dead-letter protection. if let Some(batch) = result.batch.take() { + if let Some(control) = accepted_drop_control { + tracing::info!( + channel_id = %batch.channel_id, + events = batch.events.len(), + ?control, + "dropping batch after acknowledged owner control" + ); + hard_timeout_fate_suffix = Some(" — batch dropped (owner control)"); + batch_terminally_disposed = true; // Don't requeue batches for channels the agent was removed from — // those events are stale and should be silently dropped. - if !removed_channels.contains(&batch.channel_id) { + } else if !removed_channels.contains(&batch.channel_id) { if matches!( result.outcome, - PromptOutcome::Cancelled | PromptOutcome::CancelDrainTimeout(_) + PromptOutcome::Cancelled + | PromptOutcome::CancelDrainTimeout(_) + | PromptOutcome::CancelCleanupFailed(_) + | PromptOutcome::SessionCloseFailed(_) + | PromptOutcome::SessionRecycleRequired + | PromptOutcome::ControlPreemptedSetup ) { // Cancel re-prompt: store as cancelled events so flush_next() // merges them into the next FlushBatch.cancelled_events, @@ -3102,13 +4824,14 @@ fn handle_prompt_result( // — consistent with MergeFraming::for_reason(None) and the // system default — rather than telling the agent to supersede. // - // CancelDrainTimeout shares this path with Cancelled: a failed - // 5s drain after a control-signal cancel is a cleanup-deadline - // problem, not the deterministic hard-cap death below — the - // original batch must survive with no retry/dead-letter - // accounting, same as a clean cancel. + // CancelDrainTimeout and SessionCloseFailed share this path + // with Cancelled: cleanup did not complete, but that is not + // the deterministic hard-cap death below — the original batch + // must survive with no retry/dead-letter accounting, same as + // a clean cancel. let reason = batch.cancel_reason.unwrap_or(CancelReason::Steer); queue.requeue_as_cancelled(batch, reason); + batch_requeued = true; } else if matches!( result.outcome, PromptOutcome::Timeout(TimeoutKind::Hard { @@ -3127,6 +4850,7 @@ fn handle_prompt_result( ); spawn_failure_notice(rest_client, &batch, content); hard_timeout_fate_suffix = Some(" — dead-lettered (no recent activity)"); + batch_terminally_disposed = true; } else if matches!( result.outcome, PromptOutcome::Timeout(TimeoutKind::Hard { @@ -3145,8 +4869,10 @@ fn handle_prompt_result( ); spawn_failure_notice(rest_client, &dead, content); hard_timeout_fate_suffix = Some(" — dead-lettered (retry budget exhausted)"); + batch_terminally_disposed = true; } else { hard_timeout_fate_suffix = Some(" — requeued for retry (recently active)"); + batch_requeued = true; } } else if matches!(&result.outcome, PromptOutcome::Error(e) if is_auth_error(e)) { // Auth errors are non-retryable: the token won't self-repair @@ -3163,20 +4889,29 @@ fn handle_prompt_result( and then re-send." .to_string(); spawn_failure_notice(rest_client, &batch, content); - } else if let Some(dead) = queue.requeue(batch) { - let reason = match &result.outcome { - PromptOutcome::Timeout(TimeoutKind::Idle) => "the turn timed out".to_string(), - PromptOutcome::Timeout(TimeoutKind::Hard { .. }) => { - "the turn exceeded the maximum duration".to_string() + batch_terminally_disposed = true; + } else { + match queue.requeue(batch) { + Some(dead) => { + let reason = match &result.outcome { + PromptOutcome::Timeout(TimeoutKind::Idle) => { + "the turn timed out".to_string() + } + PromptOutcome::Timeout(TimeoutKind::Hard { .. }) => { + "the turn exceeded the maximum duration".to_string() + } + PromptOutcome::AgentExited => "the agent process exited".to_string(), + PromptOutcome::Error(e) => format!("{e}"), + _ => "repeated failures".to_string(), + }; + let content = format!( + "⚠️ I couldn't process the last request after multiple retries ({reason}). Please re-send if it's still needed." + ); + spawn_failure_notice(rest_client, &dead, content); + batch_terminally_disposed = true; } - PromptOutcome::AgentExited => "the agent process exited".to_string(), - PromptOutcome::Error(e) => format!("{e}"), - _ => "repeated failures".to_string(), - }; - let content = format!( - "⚠️ I couldn't process the last request after multiple retries ({reason}). Please re-send if it's still needed." - ); - spawn_failure_notice(rest_client, &dead, content); + None => batch_requeued = true, + } } } else { tracing::debug!( @@ -3185,31 +4920,57 @@ fn handle_prompt_result( "dropping failed batch for removed channel" ); hard_timeout_fate_suffix = Some(" — batch dropped (channel removed)"); + batch_terminally_disposed = true; } } match &result.source { - PromptSource::Channel(ch) => queue.mark_complete(*ch), + PromptSource::Channel(ch) => { + queue.mark_complete(*ch); + let replacement_resets_model = match &result.outcome { + PromptOutcome::AgentExited | PromptOutcome::Timeout(_) => true, + PromptOutcome::Error(error) => error.requires_process_replacement(), + _ => false, + }; + let keep_required_agent = !batch_terminally_disposed + && (established_required_agent.is_some() + || (replacement_resets_model && preserve_model_on_fatal_replacement) + || (batch_requeued && !replacement_resets_model)); + if !keep_required_agent { + if batch_terminally_disposed { + queue.clear_required_agent_if_drained(*ch); + } else { + queue.clear_required_agent(*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. - for ch in removed_channels { - result.agent.state.invalidate_channel(ch); + if lifecycle_recycle_required { + // Batch fate and exact-slot bookkeeping above intentionally used the + // original prompt outcome. Only now replace the healthy adapter, so an + // acknowledged rotate/model switch or membership removal can never + // abandon a remote session via local map deletion. + result.outcome = PromptOutcome::SessionRecycleRequired; } let outcome_label = match &result.outcome { PromptOutcome::Ok(_) => "ok", + PromptOutcome::SessionRetired(_) => "session_retired", PromptOutcome::Error(_) => "error", PromptOutcome::Timeout(TimeoutKind::Idle) => "idle_timeout", PromptOutcome::Timeout(TimeoutKind::Hard { .. }) => "hard_timeout", PromptOutcome::AgentExited => "exited", PromptOutcome::Cancelled => "cancelled", PromptOutcome::CancelDrainTimeout(_) => "cancel_drain_timeout", + PromptOutcome::CancelCleanupFailed(_) => "cancel_cleanup_failed", + PromptOutcome::SessionCloseFailed(_) => "session_close_failed", + PromptOutcome::SessionRecycleRequired => "session_recycle_required", + PromptOutcome::SessionSetupCloseFailed(_) => "session_setup_close_failed", + PromptOutcome::SessionSetupRecycleRequired => "session_setup_recycle_required", + PromptOutcome::ControlPreemptedSetup => "control_preempted_setup", }; - let agent_index = result.agent.index; // Capture the spawn-time configured model and our PID before the agent is // moved into match arms below. `desired_model` reflects the config/persona // model at spawn time — it does NOT reflect `session/set_model` overrides, @@ -3249,7 +5010,7 @@ fn handle_prompt_result( match result.outcome { // Successful prompt — return agent to pool. - PromptOutcome::Ok(_) => { + PromptOutcome::Ok(_) | PromptOutcome::SessionRetired(_) => { tracing::debug!( agent = agent_index, outcome = outcome_label, @@ -3284,14 +5045,26 @@ fn handle_prompt_result( 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, - observer.clone(), - ) { + let spawned = if preserve_model_on_fatal_replacement { + spawn_respawn_task_preserving_model_intent( + result.agent, + config, + slot_history, + respawn_tx, + respawn_tasks, + observer.clone(), + ) + } else { + spawn_respawn_task( + result.agent, + config, + slot_history, + respawn_tx, + respawn_tasks, + observer.clone(), + ) + }; + if !spawned { // 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"); @@ -3324,7 +5097,7 @@ fn handle_prompt_result( let index = result.agent.index; let slot_history = &mut crash_history[index]; - if !spawn_respawn_task( + if !spawn_respawn_task_preserving_model_intent( result.agent, config, slot_history, @@ -3339,14 +5112,138 @@ fn handle_prompt_result( } } } - // 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 + PromptOutcome::CancelCleanupFailed(error) => { + let error_class = acp_error_class(&error); + let error_code = match &error { + acp::AcpError::AgentError { code, .. } => Some(*code), + _ => None, + }; + tracing::warn!( + agent = agent_index, + outcome = outcome_label, + configured_model = %harness_configured_model, + pid = harness_pid, + error_class, + error_code = ?error_code, + "agent_returned — respawning (cancel cleanup failed; adapter-controlled details redacted)" + ); + let death_message = format!( + "Agent cancellation cleanup failed ({error_class}); the agent process is being replaced." + ); + emit_turn_error(&death_message, error_code); + + let index = result.agent.index; + let slot_history = &mut crash_history[index]; + if !spawn_respawn_task_preserving_model_intent( + result.agent, + config, + slot_history, + respawn_tx, + respawn_tasks, + observer.clone(), + ) && pool.live_count() == 0 + && !any_respawn_in_flight(crash_history) + { + tracing::error!("all agents dead — exiting"); + return LoopAction::Exit; + } + } + // `session/close` is optional ACP. A healthy adapter that does not + // advertise it cannot safely be reused after local retirement because + // only process teardown can release its session tree. Replace it + // immediately without charging crash budget or losing a live model + // override. + PromptOutcome::SessionRecycleRequired + | PromptOutcome::SessionSetupRecycleRequired + | PromptOutcome::ControlPreemptedSetup => { + tracing::info!( + agent = agent_index, + outcome = outcome_label, + configured_model = %harness_configured_model, + pid = harness_pid, + "agent_returned — recycling (session close unsupported)" + ); + let index = result.agent.index; + let slot_history = &mut crash_history[index]; + let switch_control = match ( + &result.source, + result.agent.desired_model.as_ref(), + result.agent.model_switch_request_id.as_ref(), + ) { + (PromptSource::Channel(channel_id), Some(model_id), Some(request_id)) => { + Some(PendingSwitchControl { + channel_id: *channel_id, + model_id: model_id.clone(), + request_id: request_id.clone(), + }) + } + _ => None, + }; + if !spawn_recycle_task_with_switch( + result.agent, + config, + slot_history, + respawn_tx, + respawn_tasks, + observer.clone(), + switch_control, + ) && pool.live_count() == 0 + && !any_respawn_in_flight(crash_history) + { + tracing::error!("all agents dead — exiting"); + return LoopAction::Exit; + } + } + // Session retirement is fail-closed: once `session/close` fails, Buzz + // no longer knows whether the adapter still owns the session's query + // tree. Error class is irrelevant — even method-not-found is unsafe to + // reuse because local invalidation would abandon those resources. + PromptOutcome::SessionCloseFailed(error) + | PromptOutcome::SessionSetupCloseFailed(error) => { + let error_class = acp_error_class(&error); + let error_code = match &error { + acp::AcpError::AgentError { code, .. } => Some(*code), + _ => None, + }; + tracing::warn!( + agent = agent_index, + outcome = outcome_label, + configured_model = %harness_configured_model, + pid = harness_pid, + error_class, + error_code = ?error_code, + "agent_returned — respawning (session close failed; adapter-controlled details redacted)" + ); + let death_message = format!( + "Agent session cleanup failed ({error_class}); the agent process is being replaced." + ); + emit_turn_error(&death_message, error_code); + + let index = result.agent.index; + let slot_history = &mut crash_history[index]; + if !spawn_respawn_task_preserving_model_intent( + result.agent, + config, + slot_history, + respawn_tx, + respawn_tasks, + observer.clone(), + ) && pool.live_count() == 0 + && !any_respawn_in_flight(crash_history) + { + tracing::error!("all agents dead — exiting"); + return LoopAction::Exit; + } + } + // Errors fall into two categories: + // + // 1. Transport/framing-class (Io, WriteTimeout, Timeout, Protocol, + // Json): 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 (agent JSON-RPC errors): the pipe is // intact but the prompt failed. Return the agent to the pool so it // can be reused for the next event. @@ -3365,13 +5262,7 @@ fn handle_prompt_result( pool.return_agent(result.agent); } PromptOutcome::Error(ref e) => { - let is_transport_error = matches!( - e, - acp::AcpError::Io(_) - | acp::AcpError::WriteTimeout(_) - | acp::AcpError::Timeout(_) - | acp::AcpError::Protocol(_) - ); + let is_transport_error = e.requires_process_replacement(); let error_code = match &e { acp::AcpError::AgentError { code, .. } => Some(*code), _ => None, @@ -3389,16 +5280,26 @@ fn handle_prompt_result( 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, - observer, - ) && pool.live_count() == 0 - && !any_respawn_in_flight(crash_history) - { + let spawned = if preserve_model_on_fatal_replacement { + spawn_respawn_task_preserving_model_intent( + result.agent, + config, + slot_history, + respawn_tx, + respawn_tasks, + observer, + ) + } else { + spawn_respawn_task( + result.agent, + config, + slot_history, + respawn_tx, + respawn_tasks, + observer, + ) + }; + if !spawned && pool.live_count() == 0 && !any_respawn_in_flight(crash_history) { tracing::error!("all agents dead — exiting"); return LoopAction::Exit; } @@ -3416,6 +5317,9 @@ fn handle_prompt_result( } } } + if pool.slot_alive(agent_index) { + queue.release_required_agent(agent_index); + } LoopAction::Continue } @@ -3429,9 +5333,10 @@ fn recover_panicked_agent( removed_channels: &HashSet, typing_channels: &mut HashMap, crash_history: &mut [SlotCircuit], - respawn_tx: &mpsc::Sender, - respawn_tasks: &mut tokio::task::JoinSet<()>, + _respawn_tx: &mpsc::Sender, + _respawn_tasks: &mut tokio::task::JoinSet<()>, observer: Option, + rest_client: Option<&relay::RestClient>, ) { let task_id = join_error.id(); let Some(meta) = pool.task_map_mut().remove(&task_id) else { @@ -3439,15 +5344,49 @@ fn recover_panicked_agent( return; }; let i = meta.agent_index; + let accepted_drop_control = meta.accepted_drop_control; + let preserve_model_intent = meta.model_overridden + || meta + .channel_id + .is_some_and(|channel_id| queue.required_agent(channel_id) == Some(i)); + if preserve_model_intent { + if let Some(channel_id) = meta.channel_id { + if !removed_channels.contains(&channel_id) { + queue.require_agent(channel_id, i); + } + } + } // 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) { - // Dead-letter on exhaustion is logged inside requeue(); a - // panic path has no outcome to report, so no notice here. - let _ = queue.requeue(batch); - tracing::warn!("requeued batch for panicked agent {i}"); + if let Some(control) = accepted_drop_control { + tracing::info!( + agent = i, + channel = %ch, + events = batch.events.len(), + ?control, + "dropping panicked batch after acknowledged owner control" + ); + } else if !removed_channels.contains(&ch) { + match queue.requeue(batch) { + Some(dead) => { + queue.clear_required_agent_if_drained(ch); + tracing::error!( + agent = i, + channel = %ch, + events = dead.events.len(), + "dead-lettered batch after repeated panics" + ); + spawn_failure_notice( + rest_client, + &dead, + "⚠️ I couldn't process the last request after repeated internal task failures. Please re-send if it's still needed." + .to_string(), + ); + } + None => tracing::warn!("requeued batch for panicked agent {i}"), + } } else { tracing::debug!( channel_id = %ch, @@ -3474,48 +5413,36 @@ fn recover_panicked_agent( serde_json::json!({ "outcome": "panic", "error": format!("Agent task panicked: {join_error}"), + "slotQuarantined": true, }), ); } - // 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. + // Unwinding destroyed the only AcpClient/Child handle. Drop issued a + // best-effort group signal, but it cannot await direct-child reaping or + // prove process-group absence. Never translate that ambiguity into a timed + // respawn: quarantine the slot for this supervisor process so maintenance + // and crash cooldowns cannot create an overlapping adapter. let slot = &mut crash_history[i]; - - let delay = match slot.record_crash() { - CrashVerdict::CircuitOpen => { - tracing::error!(agent = i, "circuit open after panic — not respawning"); - return; - } - 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 - } + let (desired_model, model_overridden) = if preserve_model_intent { + (meta.desired_model.clone(), meta.model_overridden) + } else { + (config.model.clone(), false) }; - - // 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 env = config.persona_env_vars.clone(); - let has_codex = config.has_generated_codex_config; - 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, &env, has_codex, i, observer).await; - guard.send(result); + slot.pending_model_intent = preserve_model_intent.then_some(ReplacementModelIntent { + desired_model, + model_overridden, + model_switch_request_id: meta + .accepted_model_switch + .as_ref() + .map(|request| request.request_id.clone()), + model_switch_rollback: None, }); + slot.mark_cleanup_unverified(); + tracing::error!( + agent = i, + "panicked adapter cleanup cannot be verified — slot quarantined until process restart" + ); } #[allow(clippy::too_many_arguments)] @@ -3530,6 +5457,7 @@ fn drain_ready_join_results( respawn_tx: &mpsc::Sender, respawn_tasks: &mut tokio::task::JoinSet<()>, observer: Option, + rest_client: Option<&relay::RestClient>, ) -> LoopAction { while let Some(Some(join_result)) = pool.join_set.join_next().now_or_never() { if let Err(join_error) = join_result { @@ -3546,6 +5474,7 @@ fn drain_ready_join_results( respawn_tx, respawn_tasks, observer.clone(), + rest_client, ); if pool.live_count() == 0 && !any_respawn_in_flight(crash_history) { return LoopAction::Exit; @@ -3575,8 +5504,11 @@ fn dispatch_heartbeat( let result_tx = pool.result_tx(); let ctx_clone = Arc::clone(ctx); let agent_index = agent.index; + let desired_model = agent.desired_model.clone(); + let model_overridden = agent.model_overridden; let turn_id = Uuid::new_v4().to_string(); let task_turn_id = turn_id.clone(); + let (control_tx, control_rx) = tokio::sync::oneshot::channel::(); let abort_handle = pool.join_set.spawn(async move { pool::run_prompt_task( @@ -3585,7 +5517,7 @@ fn dispatch_heartbeat( Some(prompt_text), ctx_clone, result_tx, - None, + Some(control_rx), task_turn_id, ) .await; @@ -3598,7 +5530,11 @@ fn dispatch_heartbeat( channel_id: None, turn_id, recoverable_batch: None, - control_tx: None, + desired_model, + model_overridden, + accepted_model_switch: None, + accepted_drop_control: None, + control_tx: Some(control_tx), steer_tx: None, }, ); @@ -3660,14 +5596,73 @@ fn spawn_respawn_task( respawn_tx: &mpsc::Sender, respawn_tasks: &mut tokio::task::JoinSet<()>, observer: Option, +) -> bool { + spawn_respawn_task_with_policy( + old_agent, + config, + slot, + respawn_tx, + respawn_tasks, + observer, + false, + ) +} + +/// Cleanup-triggered crash replacement preserves the live model intent until +/// a replacement initializes. This includes failed session retirement and +/// bounded cancel cleanup; neither is an operator-requested model reset. +fn spawn_respawn_task_preserving_model_intent( + old_agent: OwnedAgent, + config: &Config, + slot: &mut SlotCircuit, + respawn_tx: &mpsc::Sender, + respawn_tasks: &mut tokio::task::JoinSet<()>, + observer: Option, +) -> bool { + spawn_respawn_task_with_policy( + old_agent, + config, + slot, + respawn_tx, + respawn_tasks, + observer, + true, + ) +} + +#[allow(clippy::too_many_arguments)] +fn spawn_respawn_task_with_policy( + old_agent: OwnedAgent, + config: &Config, + slot: &mut SlotCircuit, + respawn_tx: &mpsc::Sender, + respawn_tasks: &mut tokio::task::JoinSet<()>, + observer: Option, + retain_model_intent: bool, ) -> bool { let index = old_agent.index; + let (desired_model, model_overridden) = if retain_model_intent { + (old_agent.desired_model.clone(), old_agent.model_overridden) + } else { + (config.model.clone(), false) + }; // 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; + tracing::error!( + agent = index, + "circuit open — suppressing replacement after bounded shutdown" + ); + return spawn_cleanup_only_task( + old_agent, + slot, + respawn_tx, + respawn_tasks, + desired_model, + model_overridden, + retain_model_intent, + ); } CrashVerdict::HalfOpenProbe => { tracing::info!(agent = index, "circuit half-open — probe respawn"); @@ -3679,25 +5674,219 @@ fn spawn_respawn_task( } }; + spawn_replacement_task( + old_agent, + config, + slot, + respawn_tx, + respawn_tasks, + observer, + delay, + desired_model, + model_overridden, + retain_model_intent, + None, + ) +} + +/// Replace a healthy adapter whose negotiated capabilities cannot explicitly +/// retire sessions. This is lifecycle maintenance, not crash recovery: it +/// bypasses crash accounting and preserves any live model override. +fn spawn_recycle_task( + old_agent: OwnedAgent, + config: &Config, + slot: &mut SlotCircuit, + respawn_tx: &mpsc::Sender, + respawn_tasks: &mut tokio::task::JoinSet<()>, + observer: Option, +) -> bool { + spawn_recycle_task_with_switch( + old_agent, + config, + slot, + respawn_tx, + respawn_tasks, + observer, + None, + ) +} + +#[allow(clippy::too_many_arguments)] +fn spawn_model_switch_recycle_task( + old_agent: OwnedAgent, + config: &Config, + slot: &mut SlotCircuit, + respawn_tx: &mpsc::Sender, + respawn_tasks: &mut tokio::task::JoinSet<()>, + observer: Option, + control: PendingSwitchControl, +) -> bool { + spawn_recycle_task_with_switch( + old_agent, + config, + slot, + respawn_tx, + respawn_tasks, + observer, + Some(control), + ) +} + +#[allow(clippy::too_many_arguments)] +fn spawn_recycle_task_with_switch( + old_agent: OwnedAgent, + config: &Config, + slot: &mut SlotCircuit, + respawn_tx: &mpsc::Sender, + respawn_tasks: &mut tokio::task::JoinSet<()>, + observer: Option, + switch_control: Option, +) -> bool { + let desired_model = old_agent.desired_model.clone(); + let model_overridden = old_agent.model_overridden; + let delay = slot.schedule_recycle(); + spawn_replacement_task( + old_agent, + config, + slot, + respawn_tx, + respawn_tasks, + observer, + delay, + desired_model, + model_overridden, + true, + switch_control, + ) +} + +#[allow(clippy::too_many_arguments)] +fn spawn_cleanup_only_task( + old_agent: OwnedAgent, + slot: &mut SlotCircuit, + respawn_tx: &mpsc::Sender, + respawn_tasks: &mut tokio::task::JoinSet<()>, + desired_model: Option, + model_overridden: bool, + retain_model_intent: bool, +) -> bool { + let index = old_agent.index; + debug_assert!( + !slot.respawn_in_flight, + "an owned completing agent cannot already have cleanup in flight" + ); + slot.respawn_in_flight = true; + slot.pending_model_intent = retain_model_intent.then_some(ReplacementModelIntent { + desired_model: desired_model.clone(), + model_overridden, + model_switch_request_id: if retain_model_intent { + old_agent.model_switch_request_id.clone() + } else { + None + }, + model_switch_rollback: if retain_model_intent { + old_agent.model_switch_rollback.clone() + } else { + None + }, + }); + let guard = RespawnGuard::new( + index, + respawn_tx.clone(), + desired_model, + model_overridden, + retain_model_intent, + ); + respawn_tasks.spawn(async move { + let mut agent = old_agent; + if let Err(error) = agent.acp.shutdown().await { + tracing::error!( + agent = index, + error = %error, + "old adapter shutdown could not be verified; circuit remains open" + ); + guard.send_cleanup_unverified(anyhow::anyhow!( + "old adapter cleanup unverified: {error}" + )); + return; + } + drop(agent); + guard.send_cleanup_complete(); + }); + true +} + +#[allow(clippy::too_many_arguments)] +fn spawn_replacement_task( + old_agent: OwnedAgent, + config: &Config, + slot: &mut SlotCircuit, + respawn_tx: &mpsc::Sender, + respawn_tasks: &mut tokio::task::JoinSet<()>, + observer: Option, + delay: Duration, + desired_model: Option, + model_overridden: bool, + retain_model_intent: bool, + switch_control: Option, +) -> bool { + let index = old_agent.index; + debug_assert!( + !slot.respawn_in_flight, + "an owned completing agent cannot already have a replacement in flight" + ); slot.respawn_in_flight = true; + slot.pending_model_intent = retain_model_intent.then_some(ReplacementModelIntent { + desired_model: desired_model.clone(), + model_overridden, + model_switch_request_id: if retain_model_intent { + old_agent.model_switch_request_id.clone() + } else { + None + }, + model_switch_rollback: if retain_model_intent { + old_agent.model_switch_rollback.clone() + } else { + None + }, + }); // 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 env = config.persona_env_vars.clone(); let has_codex = config.has_generated_codex_config; - let guard = RespawnGuard::new(index, respawn_tx.clone()); + let mut guard = RespawnGuard::new( + index, + respawn_tx.clone(), + desired_model, + model_overridden, + retain_model_intent, + ) + .with_switch_failure(observer.clone(), switch_control); respawn_tasks.spawn(async move { // Shutdown old agent (reap child, prevent zombie). let mut agent = old_agent; - agent.acp.shutdown().await; + if let Err(error) = agent.acp.shutdown().await { + tracing::error!( + agent = index, + error = %error, + "old adapter shutdown could not be verified; refusing replacement spawn" + ); + guard.send_cleanup_unverified(anyhow::anyhow!( + "old adapter cleanup unverified: {error}" + )); + return; + } drop(agent); + guard.mark_owner_cleanup_verified(); if !delay.is_zero() { tokio::time::sleep(delay).await; } - let result = spawn_and_init(&cmd, &args, &env, has_codex, index, observer).await; + let result = + spawn_and_init(&cmd, &args, &env, has_codex, index, observer, &mut guard).await; guard.send(result); }); @@ -3715,22 +5904,218 @@ fn normalized_agent_name(init_result: &serde_json::Value) -> String { .to_ascii_lowercase() } -async fn shutdown_agent_slots(slots: &mut [Option]) { - for slot in slots { - if let Some(mut agent) = slot.take() { - agent.acp.shutdown().await; +async fn shutdown_acp_with_log( + acp: &mut AcpClient, + agent_index: Option, + phase: &'static str, +) -> bool { + match acp.shutdown().await { + Ok(()) => true, + Err(error) => { + tracing::error!( + agent = agent_index, + phase, + error = %error, + "ACP adapter shutdown could not be verified" + ); + false } } } -async fn shutdown_agent_pool(pool: &mut AgentPool) { - pool.join_set.shutdown().await; - while let Ok(mut result) = pool.result_rx_try_recv() { - result.agent.acp.shutdown().await; +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum ShutdownOwner { + Wake, + CheckedOut, + Respawn, + Idle, + ObserverPublisher, +} + +impl ShutdownOwner { + fn label(self) -> &'static str { + match self { + Self::Wake => "pool initialization", + Self::CheckedOut => "checked-out agent", + Self::Respawn => "respawn agent", + Self::Idle => "idle agent", + Self::ObserverPublisher => "observer terminal-result publisher", + } + } +} + +#[derive(Default)] +struct ShutdownVerification { + unverified: Vec, +} + +impl ShutdownVerification { + fn record(&mut self, owner: ShutdownOwner, verified: bool) { + if !verified && !self.unverified.contains(&owner) { + self.unverified.push(owner); + } + } + + fn into_result(self) -> Result<()> { + if self.unverified.is_empty() { + return Ok(()); + } + let owners = self + .unverified + .into_iter() + .map(ShutdownOwner::label) + .collect::>() + .join(", "); + Err(anyhow::anyhow!( + "ACP shutdown ownership could not be verified for: {owners}" + )) + } +} + +async fn shutdown_agent_slots(slots: &mut [Option]) -> bool { + let mut cleanup_verified = true; + for slot in slots { + if let Some(mut agent) = slot.take() { + if !shutdown_acp_with_log( + &mut agent.acp, + Some(agent.index), + "pool initialization cleanup", + ) + .await + { + cleanup_verified = false; + } + } + } + cleanup_verified +} + +fn cleanup_ownership_unverified( + pool_cleanup_unverified: bool, + crash_history: &[SlotCircuit], +) -> bool { + pool_cleanup_unverified + || crash_history + .iter() + .any(SlotCircuit::blocks_supervisor_exit) +} + +async fn await_automatic_exit_permission( + shutdown: &mut watch::Receiver<()>, + cleanup_unverified: bool, + reason: &'static str, +) { + if !cleanup_unverified { + return; + } + tracing::error!( + reason, + "automatic loop exit blocked by unverified adapter ownership; \ + holding non-spawning quarantine until explicit shutdown" + ); + let _ = shutdown.changed().await; +} + +async fn shutdown_agent_pool(pool: &mut AgentPool) -> bool { + let mut cleanup_verified = true; + let signalled = pool.cancel_checked_out_for_shutdown(); + tracing::debug!(signalled, "cooperatively cancelling checked-out agents"); + while let Some(result) = pool.join_set.join_next().await { + if let Err(error) = result { + tracing::error!(error = %error, "pool task failed during shutdown cleanup"); + cleanup_verified = false; + } + } + while let Ok(mut result) = pool.result_rx_try_recv() { + if !shutdown_acp_with_log( + &mut result.agent.acp, + Some(result.agent.index), + "pool result cleanup", + ) + .await + { + cleanup_verified = false; + } } for slot in pool.agents_mut() { if let Some(mut agent) = slot.take() { - agent.acp.shutdown().await; + if !shutdown_acp_with_log(&mut agent.acp, Some(agent.index), "pool slot cleanup").await + { + cleanup_verified = false; + } + } + } + cleanup_verified +} + +async fn fail_pre_loop_with_pool_cleanup( + pool: &mut AgentPool, + shutdown: &mut watch::Receiver<()>, + phase: &'static str, + error: anyhow::Error, + cleanup_already_unverified: bool, + shutdown_requested: bool, +) -> Result { + let cleanup_unverified = cleanup_already_unverified || !shutdown_agent_pool(pool).await; + if cleanup_unverified && !shutdown_requested { + // An automatic pre-loop failure must not let a service manager restart + // this process while any previously spawned adapter may still exist. + // Remain in the non-spawning quarantine until an explicit signal + // requests process shutdown. + tracing::error!( + phase, + error = %error, + "pre-loop failure with unverified adapter cleanup — holding quarantine until shutdown" + ); + let _ = shutdown.changed().await; + } + + if cleanup_unverified { + Err(anyhow::anyhow!( + "{phase}: {error}; adapter cleanup remains unverified" + )) + } else { + Err(anyhow::anyhow!("{phase}: {error}")) + } +} + +async fn await_pre_loop_operation( + operation: F, + shutdown: &mut watch::Receiver<()>, + pool: &mut AgentPool, + phase: &'static str, + cleanup_already_unverified: bool, +) -> Result +where + E: std::fmt::Display, + F: std::future::Future>, +{ + tokio::select! { + biased; + _ = shutdown.changed() => { + fail_pre_loop_with_pool_cleanup( + pool, + shutdown, + phase, + anyhow::anyhow!("shutdown signal received before the main loop"), + cleanup_already_unverified, + true, + ) + .await + } + result = operation => match result { + Ok(value) => Ok(value), + Err(error) => { + fail_pre_loop_with_pool_cleanup( + pool, + shutdown, + phase, + anyhow::anyhow!("{error}"), + cleanup_already_unverified, + false, + ) + .await + } } } } @@ -3745,6 +6130,55 @@ struct PoolStartup { observer: Option, } +#[derive(Debug, thiserror::Error)] +enum PoolStartupFailure { + #[error("{0}")] + Retryable(String), + #[error("{0}")] + CleanupUnverified(String), + #[error("{0}")] + Invariant(String), +} + +enum InitialPoolStartup { + Ready(AgentPool), + Dormant, + CleanupQuarantine(String), +} + +fn resolve_initial_pool_startup( + lazy_pool: bool, + eager_result: Option>, +) -> std::result::Result { + if lazy_pool { + debug_assert!(eager_result.is_none()); + return Ok(InitialPoolStartup::Dormant); + } + + let eager_result = eager_result.ok_or_else(|| { + PoolStartupFailure::Invariant( + "eager pool startup did not provide an initialization result".to_string(), + ) + })?; + match eager_result { + Ok(pool) => Ok(InitialPoolStartup::Ready(pool)), + Err(PoolStartupFailure::CleanupUnverified(error)) => { + Ok(InitialPoolStartup::CleanupQuarantine(error)) + } + Err(error) => Err(error), + } +} + +fn lazy_pool_can_wake(lazy_pool: bool, pool_ready: bool, cleanup_unverified: bool) -> bool { + lazy_pool && !pool_ready && !cleanup_unverified +} + +impl PoolStartupFailure { + fn cleanup_unverified(&self) -> bool { + matches!(self, Self::CleanupUnverified(_)) + } +} + impl PoolStartup { fn from_config(config: &Config, observer: Option) -> Self { Self { @@ -3762,9 +6196,10 @@ impl PoolStartup { async fn initialize_agent_pool( startup: &PoolStartup, mut shutdown: Option>, -) -> Result { - // One agent failing to start must not kill the whole pool. - // Attempt each spawn under a 60-second timeout; a partial pool is valid. +) -> std::result::Result { + // One agent failing to start need not kill the whole pool after its child + // is proved absent. Attempt each initialization under a 60-second timeout; + // a partial pool is valid only across verified cleanup boundaries. let mut agent_slots: Vec> = Vec::with_capacity(startup.agents as usize); for i in 0..startup.agents as usize { let spawn_result = AcpClient::spawn( @@ -3782,9 +6217,23 @@ async fn initialize_agent_pool( Some(shutdown) => tokio::select! { biased; _ = shutdown.changed() => { - acp.shutdown().await; - shutdown_agent_slots(&mut agent_slots).await; - return Err(anyhow::anyhow!("pool initialization cancelled by shutdown")); + let current_cleanup_verified = shutdown_acp_with_log( + &mut acp, + Some(i), + "cancelled pool initialization", + ) + .await; + let owned_cleanup_verified = + shutdown_agent_slots(&mut agent_slots).await; + if !current_cleanup_verified || !owned_cleanup_verified { + return Err(PoolStartupFailure::CleanupUnverified(format!( + "pool initialization cancelled by shutdown; adapter cleanup \ + unverified at or before slot {i}" + ))); + } + return Err(PoolStartupFailure::Retryable( + "pool initialization cancelled by shutdown".into(), + )); } result = initialize => result, }, @@ -3821,6 +6270,8 @@ async fn initialize_agent_pool( model_capabilities: None, desired_model: startup.model.clone(), model_overridden: false, + model_switch_request_id: None, + model_switch_rollback: None, agent_name, goose_system_prompt_supported: None, protocol_version, @@ -3828,12 +6279,44 @@ async fn initialize_agent_pool( } Ok(Err(e)) => { tracing::error!(agent = i, "agent initialize failed: {e}"); - acp.shutdown().await; + if !shutdown_acp_with_log(&mut acp, Some(i), "failed pool initialization") + .await + { + let owned_cleanup_verified = + shutdown_agent_slots(&mut agent_slots).await; + return Err(PoolStartupFailure::CleanupUnverified(format!( + "agent {i} initialize failed: {e}; spawned adapter cleanup \ + unverified{}", + if owned_cleanup_verified { + "" + } else { + "; cleanup of an already-initialized adapter was also unverified" + } + ))); + } agent_slots.push(None); } Err(_) => { tracing::error!(agent = i, "agent timed out during init (60s)"); - acp.shutdown().await; + if !shutdown_acp_with_log( + &mut acp, + Some(i), + "timed-out pool initialization", + ) + .await + { + let owned_cleanup_verified = + shutdown_agent_slots(&mut agent_slots).await; + return Err(PoolStartupFailure::CleanupUnverified(format!( + "agent {i} timed out during initialization; spawned adapter \ + cleanup unverified{}", + if owned_cleanup_verified { + "" + } else { + "; cleanup of an already-initialized adapter was also unverified" + } + ))); + } agent_slots.push(None); } } @@ -3846,10 +6329,10 @@ async fn initialize_agent_pool( } let live_count = agent_slots.iter().filter(|slot| slot.is_some()).count(); if live_count == 0 { - return Err(anyhow::anyhow!( + return Err(PoolStartupFailure::Retryable(format!( "all {} agents failed to start — cannot continue", startup.agents - )); + ))); } if live_count < startup.agents as usize { tracing::warn!( @@ -3866,7 +6349,9 @@ async fn initialize_agent_pool( /// 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. +/// borrowing `Config`. All respawn/refill paths use this. Its failure type +/// preserves whether a partially initialized child was proved shut down, so +/// the supervisor cannot turn cleanup ambiguity into a cooldown retry. async fn spawn_and_init( command: &str, args: &[String], @@ -3874,16 +6359,39 @@ async fn spawn_and_init( has_generated_codex_config: bool, agent_index: usize, observer: Option, -) -> Result<(AcpClient, u32, String)> { + guard: &mut RespawnGuard, +) -> SpawnInitResult { let mut acp = AcpClient::spawn(command, args, extra_env, has_generated_codex_config) .await - .map_err(|e| anyhow::anyhow!("failed to spawn agent: {e}"))?; + .map_err(|e| SpawnInitFailure::Spawn(anyhow::anyhow!("failed to spawn agent: {e}")))?; acp.set_observer(observer, agent_index); + let initializing = Arc::new(tokio::sync::Mutex::new(Some(acp))); + guard.mark_replacement_initializing(Arc::clone(&initializing)); + let mut slot = initializing.lock().await; - match acp.initialize().await { + let initialize_result = match slot.as_mut() { + Some(acp) => acp.initialize().await, + None => { + drop(slot); + return Err(missing_replacement_ownership( + guard, + "before adapter initialization", + )); + } + }; + match initialize_result { Ok(init_result) => { tracing::info!("agent initialized: {init_result}"); let protocol_version = init_result["protocolVersion"].as_u64().unwrap_or(1) as u32; + let Some(acp) = slot.take() else { + drop(slot); + return Err(missing_replacement_ownership( + guard, + "after successful adapter initialization", + )); + }; + drop(slot); + guard.mark_owner_cleanup_verified(); acp.observe( "agent_initialized", serde_json::json!({ @@ -3896,10 +6404,35 @@ async fn spawn_and_init( } 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}")) + // Drop is best-effort; replacement must know whether bounded + // shutdown was actually verified. + let cleanup = match slot.as_mut() { + Some(acp) => acp.shutdown().await, + None => { + drop(slot); + return Err(missing_replacement_ownership( + guard, + "after failed adapter initialization", + )); + } + }; + let _ = slot.take(); + drop(slot); + match cleanup { + Ok(()) => { + guard.mark_owner_cleanup_verified(); + Err(SpawnInitFailure::Spawn(anyhow::anyhow!( + "agent initialize failed: {e}" + ))) + } + Err(cleanup_error) => { + guard.phase = RespawnPhase::OldOwnerUnverified; + Err(SpawnInitFailure::CleanupUnverified(anyhow::anyhow!( + "agent initialize failed: {e}; spawned adapter cleanup unverified: \ + {cleanup_error}" + ))) + } + } } } } @@ -3909,6 +6442,30 @@ async fn spawn_auth_client(agent: &AuthAgentArgs) -> Result Result<()> { + client + .shutdown() + .await + .map_err(|error| anyhow::anyhow!("{phase}: adapter cleanup unverified: {error}")) +} + +fn cli_failure_after_cleanup( + operation_error: anyhow::Error, + cleanup_result: Result<()>, +) -> Result { + cleanup_result?; + Err(operation_error) +} + +async fn fail_cli_client( + client: &mut AcpClient, + cleanup_phase: &'static str, + operation_error: anyhow::Error, +) -> Result { + let cleanup_result = shutdown_cli_client(client, cleanup_phase).await; + cli_failure_after_cleanup(operation_error, cleanup_result) +} + fn extract_auth_methods(init_result: &serde_json::Value) -> Vec { init_result .get("authMethods") @@ -3918,31 +6475,41 @@ fn extract_auth_methods(init_result: &serde_json::Value) -> Vec Result<()> { - let mut client = match spawn_auth_client(&args.agent).await { - Ok(c) => c, - Err(e) => { - eprintln!("error: failed to spawn agent: {e}"); - std::process::exit(1); +async fn run_auth_methods(args: AuthMethodsArgs, mut shutdown: watch::Receiver<()>) -> Result<()> { + let mut client = spawn_auth_client(&args.agent) + .await + .map_err(|error| anyhow::anyhow!("failed to spawn agent: {error}"))?; + + let initialize = tokio::select! { + biased; + _ = shutdown.changed() => { + shutdown_cli_client(&mut client, "auth-methods signal cleanup").await?; + return Err(anyhow::anyhow!("auth-methods cancelled by shutdown signal")); } + result = tokio::time::timeout(MODELS_TIMEOUT, client.initialize()) => result, }; - - let init_result = match tokio::time::timeout(MODELS_TIMEOUT, client.initialize()).await { + let init_result = match initialize { Ok(Ok(result)) => result, Ok(Err(e)) => { - client.shutdown().await; - eprintln!("error: agent initialize failed: {e}"); - std::process::exit(1); + return fail_cli_client( + &mut client, + "auth-methods initialization failure", + anyhow::anyhow!("agent initialize failed: {e}"), + ) + .await; } Err(_) => { - client.shutdown().await; - eprintln!("error: agent timed out ({MODELS_TIMEOUT:?})"); - std::process::exit(1); + return fail_cli_client( + &mut client, + "auth-methods initialization timeout", + anyhow::anyhow!("agent timed out ({MODELS_TIMEOUT:?})"), + ) + .await; } }; let methods = extract_auth_methods(&init_result); - client.shutdown().await; + shutdown_cli_client(&mut client, "auth-methods completion").await?; if args.json { let output = serde_json::json!({ "methods": methods }); @@ -3966,26 +6533,36 @@ async fn run_auth_methods(args: AuthMethodsArgs) -> Result<()> { } /// `buzz-acp authenticate` — invoke one adapter-owned auth method. -async fn run_authenticate(args: AuthenticateArgs) -> Result<()> { - let mut client = match spawn_auth_client(&args.agent).await { - Ok(c) => c, - Err(e) => { - eprintln!("error: failed to spawn agent: {e}"); - std::process::exit(1); +async fn run_authenticate(args: AuthenticateArgs, mut shutdown: watch::Receiver<()>) -> Result<()> { + let mut client = spawn_auth_client(&args.agent) + .await + .map_err(|error| anyhow::anyhow!("failed to spawn agent: {error}"))?; + + let initialize = tokio::select! { + biased; + _ = shutdown.changed() => { + shutdown_cli_client(&mut client, "authenticate signal cleanup").await?; + return Err(anyhow::anyhow!("authenticate cancelled by shutdown signal")); } + result = tokio::time::timeout(MODELS_TIMEOUT, client.initialize()) => result, }; - - let init_result = match tokio::time::timeout(MODELS_TIMEOUT, client.initialize()).await { + let init_result = match initialize { Ok(Ok(result)) => result, Ok(Err(e)) => { - client.shutdown().await; - eprintln!("error: agent initialize failed: {e}"); - std::process::exit(1); + return fail_cli_client( + &mut client, + "authenticate initialization failure", + anyhow::anyhow!("agent initialize failed: {e}"), + ) + .await; } Err(_) => { - client.shutdown().await; - eprintln!("error: agent initialize timed out ({MODELS_TIMEOUT:?})"); - std::process::exit(1); + return fail_cli_client( + &mut client, + "authenticate initialization timeout", + anyhow::anyhow!("agent initialize timed out ({MODELS_TIMEOUT:?})"), + ) + .await; } }; @@ -3993,38 +6570,56 @@ async fn run_authenticate(args: AuthenticateArgs) -> Result<()> { .iter() .any(|method| method.get("id").and_then(|id| id.as_str()) == Some(args.method_id.as_str())); if !supports_method { - client.shutdown().await; - eprintln!( - "error: auth method '{}' is not advertised by this adapter", - args.method_id - ); - std::process::exit(1); + return fail_cli_client( + &mut client, + "unsupported authenticate method", + anyhow::anyhow!( + "auth method '{}' is not advertised by this adapter", + args.method_id + ), + ) + .await; } - let result = - tokio::time::timeout(AUTHENTICATE_TIMEOUT, client.authenticate(&args.method_id)).await; + let result = tokio::select! { + biased; + _ = shutdown.changed() => { + shutdown_cli_client(&mut client, "authenticate signal cleanup").await?; + return Err(anyhow::anyhow!("authenticate cancelled by shutdown signal")); + } + result = tokio::time::timeout( + AUTHENTICATE_TIMEOUT, + client.authenticate(&args.method_id), + ) => result, + }; match result { Ok(Ok(_)) => { - client.shutdown().await; + shutdown_cli_client(&mut client, "authenticate completion").await?; Ok(()) } Ok(Err(e)) => { - client.shutdown().await; - eprintln!("error: authenticate failed: {e}"); - std::process::exit(1); + fail_cli_client( + &mut client, + "authenticate failure", + anyhow::anyhow!("authenticate failed: {e}"), + ) + .await } Err(_) => { - client.shutdown().await; - eprintln!("error: authenticate timed out ({AUTHENTICATE_TIMEOUT:?})"); - std::process::exit(1); + fail_cli_client( + &mut client, + "authenticate timeout", + anyhow::anyhow!("authenticate timed out ({AUTHENTICATE_TIMEOUT:?})"), + ) + .await } } } /// Flow: spawn → initialize → session/new → print models → shutdown. /// No relay connection, no MCP servers, no subscriptions. ~2-5s total. -async fn run_models(args: ModelsArgs) -> Result<()> { +async fn run_models(args: ModelsArgs, mut shutdown: watch::Receiver<()>) -> Result<()> { use acp::{extract_model_config_options, extract_model_state}; let agent_args = config::normalize_agent_args(&args.agent.agent_command, args.agent.agent_args); @@ -4035,35 +6630,42 @@ async fn run_models(args: ModelsArgs) -> Result<()> { // Spawn outside the timeout so we always own the child for cleanup. // `models` subcommand doesn't use persona packs — no extra env, no codex config. - let mut client = - match AcpClient::spawn(&args.agent.agent_command, &agent_args, &[], false).await { - Ok(c) => c, - Err(e) => { - eprintln!("error: failed to spawn agent: {e}"); - std::process::exit(1); - } - }; + let mut client = AcpClient::spawn(&args.agent.agent_command, &agent_args, &[], false) + .await + .map_err(|error| anyhow::anyhow!("failed to spawn agent: {error}"))?; // Initialize + session/new under a timeout. Client is owned above, // so shutdown() runs on all paths (success, error, timeout). - let protocol_result = tokio::time::timeout(MODELS_TIMEOUT, async { - let init = client.initialize().await?; - let session = client.session_new_full(&cwd, vec![], None, None).await?; - Ok::<_, acp::AcpError>((init, session)) - }) - .await; + let protocol_result = tokio::select! { + biased; + _ = shutdown.changed() => { + shutdown_cli_client(&mut client, "models signal cleanup").await?; + return Err(anyhow::anyhow!("models cancelled by shutdown signal")); + } + result = tokio::time::timeout(MODELS_TIMEOUT, async { + let init = client.initialize().await?; + let session = client.session_new_full(&cwd, vec![], None, None).await?; + Ok::<_, acp::AcpError>((init, session)) + }) => result, + }; let (init_result, session_resp) = match protocol_result { Ok(Ok(tuple)) => tuple, Ok(Err(e)) => { - client.shutdown().await; - eprintln!("error: agent communication failed: {e}"); - std::process::exit(1); + return fail_cli_client( + &mut client, + "models protocol failure", + anyhow::anyhow!("agent communication failed: {e}"), + ) + .await; } Err(_) => { - client.shutdown().await; - eprintln!("error: agent timed out ({MODELS_TIMEOUT:?})"); - std::process::exit(1); + return fail_cli_client( + &mut client, + "models protocol timeout", + anyhow::anyhow!("agent timed out ({MODELS_TIMEOUT:?})"), + ) + .await; } }; @@ -4084,6 +6686,7 @@ async fn run_models(args: ModelsArgs) -> Result<()> { // Extract model info from session/new response. let config_options = extract_model_config_options(&session_resp.raw); let model_state = extract_model_state(&session_resp.raw); + shutdown_cli_client(&mut client, "models completion").await?; if args.json { // Structured JSON output — consumed by Phase 3 `get_agent_models`. @@ -4156,7 +6759,6 @@ async fn run_models(args: ModelsArgs) -> Result<()> { } } - client.shutdown().await; Ok(()) } @@ -4298,6 +6900,33 @@ mod owner_control_command_tests { )); } + #[test] + fn owner_control_execution_requires_relay_admission() { + let agent = "ab".repeat(32); + let event = make_event(KIND_STREAM_MESSAGE, "!shutdown", Some(&agent)); + + assert!( + !is_admitted_owner_control_command( + &event, + KIND_STREAM_MESSAGE, + "!shutdown", + &agent, + false, + ), + "command shape alone must never authorize harness execution" + ); + assert!( + is_admitted_owner_control_command( + &event, + KIND_STREAM_MESSAGE, + "!shutdown", + &agent, + true, + ), + "relay-admitted exact command shape remains executable" + ); + } + #[test] fn mode_gate_signal_maps_handling_to_control_signal() { let owner = "a".repeat(64); @@ -4353,72 +6982,355 @@ mod owner_control_command_tests { channel_id: Some(channel_id), turn_id: "test-turn-id".to_string(), recoverable_batch: None, + desired_model: None, + model_overridden: false, + accepted_model_switch: None, + accepted_drop_control: None, control_tx: Some(control_tx), steer_tx: None, }, ); - assert!(!signal_in_flight_task( - &mut pool, - other_channel_id, - ControlSignal::Rotate - )); - assert!(signal_in_flight_task( - &mut pool, - channel_id, - ControlSignal::Rotate - )); + assert_eq!( + signal_in_flight_task(&mut pool, other_channel_id, ControlSignal::Rotate), + ControlSignalResult::NotAccepted + ); + assert_eq!( + signal_in_flight_task(&mut pool, channel_id, ControlSignal::Rotate), + ControlSignalResult::Delivered + ); + assert_eq!( + pool.task_map() + .values() + .next() + .and_then(|meta| meta.accepted_drop_control), + Some(AcceptedDropControl::Rotate), + "a successfully delivered rotate must persist its drop disposition" + ); assert_eq!(control_rx.await.unwrap(), ControlSignal::Rotate); - assert!(!signal_in_flight_task( - &mut pool, - channel_id, - ControlSignal::Rotate - )); + assert_eq!( + signal_in_flight_task(&mut pool, channel_id, ControlSignal::Rotate), + ControlSignalResult::DropRecorded + ); } -} -#[cfg(test)] -mod owner_cache_tests { - use super::*; + #[tokio::test] + async fn signal_in_flight_task_records_cancel_drop_disposition() { + let mut pool = AgentPool::from_slots(vec![]); + let channel_id = Uuid::new_v4(); + let (control_tx, control_rx) = tokio::sync::oneshot::channel(); - #[test] - fn new_with_some_caches_immediately() { - let cache = OwnerCache::new(Some("abcd".into())); - assert_eq!(cache.get(), Some("abcd")); - } + let abort_handle = pool.join_set.spawn(async {}); + pool.task_map_mut().insert( + abort_handle.id(), + pool::TaskMeta { + agent_index: 0, + channel_id: Some(channel_id), + turn_id: "cancel-drop".to_string(), + recoverable_batch: None, + desired_model: None, + model_overridden: false, + accepted_model_switch: None, + accepted_drop_control: None, + control_tx: Some(control_tx), + steer_tx: None, + }, + ); - #[test] - fn new_with_none_returns_none() { - let cache = OwnerCache::new(None); - assert!(cache.get().is_none()); + assert_eq!( + signal_in_flight_task(&mut pool, channel_id, ControlSignal::Cancel), + ControlSignalResult::Delivered + ); + assert_eq!( + pool.task_map() + .values() + .next() + .and_then(|meta| meta.accepted_drop_control), + Some(AcceptedDropControl::Cancel) + ); + assert_eq!(control_rx.await.unwrap(), ControlSignal::Cancel); } - #[test] - fn get_returns_cached_value() { - let cache = OwnerCache::new(Some("ab".repeat(32))); - assert_eq!(cache.get(), Some("ab".repeat(32)).as_deref()); - } -} + #[tokio::test] + async fn later_cancel_upgrades_batch_fate_after_interrupt_consumed_control_sender() { + let mut pool = AgentPool::from_slots(vec![]); + let channel_id = Uuid::new_v4(); + let (control_tx, control_rx) = tokio::sync::oneshot::channel(); -#[cfg(test)] -mod author_gate_tests { - use super::*; + let abort_handle = pool.join_set.spawn(async {}); + pool.task_map_mut().insert( + abort_handle.id(), + pool::TaskMeta { + agent_index: 0, + channel_id: Some(channel_id), + turn_id: "interrupt-then-cancel".to_string(), + recoverable_batch: None, + desired_model: None, + model_overridden: false, + accepted_model_switch: None, + accepted_drop_control: None, + control_tx: Some(control_tx), + steer_tx: None, + }, + ); - /// A `RestClient` for tests. The author-gate decisions exercised here all - /// resolve from the owner pubkey or sibling cache before any HTTP call, so - /// this client is never actually used to make a request. - fn dummy_rest_client() -> relay::RestClient { - relay::RestClient { - http: reqwest::Client::new(), - base_url: "http://localhost:0".into(), - keys: nostr::Keys::generate(), - auth_tag_json: None, - } + assert_eq!( + signal_in_flight_task(&mut pool, channel_id, ControlSignal::Interrupt), + ControlSignalResult::Delivered + ); + assert_eq!(control_rx.await.unwrap(), ControlSignal::Interrupt); + assert_eq!( + signal_in_flight_task(&mut pool, channel_id, ControlSignal::Cancel), + ControlSignalResult::DropRecorded, + "the owner drop disposition must still be recorded after the sender is consumed" + ); + assert_eq!( + pool.task_map() + .values() + .next() + .and_then(|meta| meta.accepted_drop_control), + Some(AcceptedDropControl::Cancel) + ); } - const OWNER: &str = "00"; - const SIBLING: &str = "11"; - const EXTERNAL: &str = "22"; + #[tokio::test] + async fn later_rotate_upgrades_batch_and_lifecycle_after_switch_consumed_control_sender() { + let mut pool = AgentPool::from_slots(vec![]); + let channel_id = Uuid::new_v4(); + let (control_tx, control_rx) = tokio::sync::oneshot::channel(); + + let abort_handle = pool.join_set.spawn(async {}); + pool.task_map_mut().insert( + abort_handle.id(), + pool::TaskMeta { + agent_index: 0, + channel_id: Some(channel_id), + turn_id: "switch-then-rotate".to_string(), + recoverable_batch: None, + desired_model: None, + model_overridden: false, + accepted_model_switch: None, + accepted_drop_control: None, + control_tx: Some(control_tx), + steer_tx: None, + }, + ); + + assert_eq!( + signal_in_flight_task( + &mut pool, + channel_id, + ControlSignal::SwitchModel(ModelSwitchRequest::new( + "runtime-model", + "0123456789abcdef0123456789abcdef", + )) + ), + ControlSignalResult::Delivered + ); + assert_eq!( + control_rx.await.unwrap(), + ControlSignal::SwitchModel(ModelSwitchRequest::new( + "runtime-model", + "0123456789abcdef0123456789abcdef", + )) + ); + assert_eq!( + signal_in_flight_task(&mut pool, channel_id, ControlSignal::Rotate), + ControlSignalResult::DropRecorded, + "rotate must record upgraded batch fate and lifecycle after a prior control" + ); + let meta = pool.task_map().values().next().unwrap(); + assert_eq!( + meta.accepted_drop_control, + Some(AcceptedDropControl::Rotate) + ); + assert_eq!( + meta.accepted_model_switch + .as_ref() + .map(|request| request.model_id.as_str()), + Some("runtime-model"), + "the earlier accepted model intent must survive the rotate upgrade" + ); + assert_eq!( + meta.accepted_model_switch + .as_ref() + .map(|request| request.request_id.as_str()), + Some("0123456789abcdef0123456789abcdef"), + "the exact request correlation must survive the rotate upgrade" + ); + } + + #[tokio::test] + async fn signal_in_flight_task_reports_closed_receiver_as_not_sent() { + let mut pool = AgentPool::from_slots(vec![]); + let channel_id = Uuid::new_v4(); + let (control_tx, control_rx) = tokio::sync::oneshot::channel(); + drop(control_rx); + + let abort_handle = pool.join_set.spawn(async {}); + pool.task_map_mut().insert( + abort_handle.id(), + pool::TaskMeta { + agent_index: 0, + channel_id: Some(channel_id), + turn_id: "post-prompt-cleanup".to_string(), + recoverable_batch: None, + desired_model: None, + model_overridden: false, + accepted_model_switch: None, + accepted_drop_control: None, + control_tx: Some(control_tx), + steer_tx: None, + }, + ); + + assert_eq!( + signal_in_flight_task( + &mut pool, + channel_id, + ControlSignal::SwitchModel(ModelSwitchRequest::new( + "gpt-5", + "fedcba9876543210fedcba9876543210", + )) + ), + ControlSignalResult::NotAccepted, + "a dropped task receiver must not be reported as a delivered control" + ); + } +} + +#[cfg(test)] +mod model_switch_request_id_tests { + use super::*; + + #[test] + fn accepts_only_exact_lowercase_hex_request_ids() { + assert_eq!( + parse_model_switch_request_id(&serde_json::json!({ + "requestId": "0123456789abcdef0123456789abcdef" + })), + Some("0123456789abcdef0123456789abcdef") + ); + + for invalid in [ + serde_json::json!({}), + serde_json::json!({"requestId": null}), + serde_json::json!({"requestId": 7}), + serde_json::json!({"requestId": "0123456789abcdef0123456789abcde"}), + serde_json::json!({"requestId": "0123456789abcdef0123456789abcdef0"}), + serde_json::json!({"requestId": "0123456789ABCDEF0123456789ABCDEF"}), + serde_json::json!({"requestId": "0123456789abcdef0123456789abcdeg"}), + ] { + assert_eq!( + parse_model_switch_request_id(&invalid), + None, + "invalid requestId must be rejected: {invalid}" + ); + } + } + + #[test] + fn every_immediate_status_echoes_the_exact_request_id() { + let request_id = "0123456789abcdef0123456789abcdef"; + for status in [ + "sent", + "recycling", + "switch_failed", + "unsupported_model", + "turn_ending", + "no_active_turn", + ] { + let payload = switch_model_control_result_payload(status, "runtime-model", request_id); + assert_eq!(payload["status"], status); + assert_eq!(payload["modelId"], "runtime-model"); + assert_eq!(payload["requestId"], request_id); + } + } +} + +#[cfg(test)] +mod observer_control_dedup_tests { + use super::*; + + #[test] + fn still_fresh_ids_are_retained_and_capacity_fails_closed() { + let mut dedup = ObserverControlDedup::new(2); + let now = 1_000; + + assert_eq!( + dedup.admit("event-1".into(), now - 2, now), + ObserverControlAdmission::Admitted + ); + assert_eq!( + dedup.admit("event-2".into(), now - 1, now), + ObserverControlAdmission::Admitted + ); + assert_eq!( + dedup.admit("event-3".into(), now, now), + ObserverControlAdmission::CapacityExceeded, + "a full fresh window must reject new controls instead of forgetting replay IDs" + ); + assert_eq!( + dedup.admit("event-1".into(), now - 2, now), + ObserverControlAdmission::Replay, + "capacity pressure must not make a still-fresh replay admissible" + ); + assert_eq!(dedup.order.len(), 2); + assert_eq!(dedup.seen.len(), 2); + assert_eq!( + dedup.admit( + "event-3".into(), + now + OBSERVER_CONTROL_FRESHNESS_SECS + 1, + now + OBSERVER_CONTROL_FRESHNESS_SECS + 1, + ), + ObserverControlAdmission::Admitted, + "expired IDs should free capacity for fresh controls" + ); + } +} + +#[cfg(test)] +mod owner_cache_tests { + use super::*; + + #[test] + fn new_with_some_caches_immediately() { + let cache = OwnerCache::new(Some("abcd".into())); + assert_eq!(cache.get(), Some("abcd")); + } + + #[test] + fn new_with_none_returns_none() { + let cache = OwnerCache::new(None); + assert!(cache.get().is_none()); + } + + #[test] + fn get_returns_cached_value() { + let cache = OwnerCache::new(Some("ab".repeat(32))); + assert_eq!(cache.get(), Some("ab".repeat(32)).as_deref()); + } +} + +#[cfg(test)] +mod author_gate_tests { + use super::*; + + /// A `RestClient` for tests. The author-gate decisions exercised here all + /// resolve from the owner pubkey or sibling cache before any HTTP call, so + /// this client is never actually used to make a request. + fn dummy_rest_client() -> relay::RestClient { + relay::RestClient { + http: reqwest::Client::new(), + base_url: "http://localhost:0".into(), + keys: nostr::Keys::generate(), + auth_tag_json: None, + } + } + + const OWNER: &str = "00"; + const SIBLING: &str = "11"; + const EXTERNAL: &str = "22"; const STRANGER: &str = "33"; /// Owner + a known sibling, none of them on the explicit allowlist. @@ -4786,8 +7698,89 @@ mod observer_snapshot_race_tests { ); } + #[tokio::test(start_paused = true)] + async fn non_terminal_control_result_uses_the_ordinary_observer_publisher() { + let observer = observer::ObserverHandle::in_process(); + let agent_keys = Keys::generate(); + let owner_keys = Keys::generate(); + let (publisher, mut published_rx) = RelayEventPublisher::test_pair(); + let rx = observer.subscribe(); + observer.emit( + "control_result", + None, + &observer::context_for(None, None, None), + serde_json::json!({ "status": "sent" }), + ); + let snapshot = observer.snapshot(); + drop(observer); + + let task = tokio::spawn(run_relay_observer_publisher( + snapshot, + rx, + publisher, + agent_keys.clone(), + agent_keys.public_key().to_hex(), + owner_keys.public_key().to_hex(), + owner_keys.public_key(), + )); + tokio::time::advance(Duration::from_secs(2)).await; + + let event = published_rx + .recv() + .await + .expect("non-terminal control progress must remain observable"); + assert!( + !event.tags.iter().any(|tag| { + tag.as_slice().first().map(String::as_str) == Some("priority") + && tag.as_slice().get(1).map(String::as_str) == Some("control-result") + }), + "non-terminal progress must not enter the protected terminal publisher" + ); + + assert!( + task.await.expect("publisher must drain cleanly"), + "ordinary control progress must not invalidate terminal delivery proof" + ); + assert!( + published_rx.try_recv().is_err(), + "subscribe-before-snapshot overlap must publish progress exactly once" + ); + } + + #[tokio::test] + async fn publisher_wrapper_releases_its_sender_and_drains_after_outer_drop() { + let observer = observer::ObserverHandle::in_process(); + let agent_keys = Keys::generate(); + let owner_keys = Keys::generate(); + let (publisher, _published_rx) = RelayEventPublisher::test_pair(); + observer.emit( + "control_result", + None, + &observer::context_for(None, None, None), + serde_json::json!({ "status": "shutdown" }), + ); + + let task = spawn_relay_observer_publisher( + observer.clone(), + publisher, + agent_keys.clone(), + agent_keys.public_key().to_hex(), + owner_keys.public_key().to_hex(), + owner_keys.public_key(), + ); + drop(observer); + + assert!( + tokio::time::timeout(Duration::from_secs(1), task) + .await + .expect("publisher wrapper must observe closed lanes and drain") + .expect("publisher wrapper must not panic"), + "confirmed terminal delivery must preserve successful shutdown verification" + ); + } + /// An event emitted between `subscribe()` and `snapshot()` lands in BOTH - /// the snapshot and the live receiver; the seq high-water dedupe must + /// the snapshot and the live receiver; exact snapshot membership must /// deliver it exactly once — and never lose events on either side of it. #[tokio::test(start_paused = true)] async fn overlap_between_subscribe_and_snapshot_publishes_exactly_once() { @@ -4835,37 +7828,335 @@ mod observer_snapshot_race_tests { "each event must be published exactly once, in order" ); } -} - -#[cfg(test)] -mod observer_publish_pacer_tests { - use super::*; #[tokio::test(start_paused = true)] - async fn starts_without_a_burst_and_spaces_frames() { - let started = tokio::time::Instant::now(); - let mut pacer = ObserverPublishPacer::new(); + async fn snapshot_membership_does_not_drop_an_earlier_uncaptured_control_result() { + let observer = observer::ObserverHandle::in_process(); + let agent_keys = Keys::generate(); + let owner_keys = Keys::generate(); + let (publisher, mut published_rx) = RelayEventPublisher::test_pair(); + let rx = observer.subscribe(); - pacer.wait().await; - let first = tokio::time::Instant::now(); - pacer.wait().await; - let second = tokio::time::Instant::now(); + observer.emit( + "control_result", + None, + &observer::context_for(None, None, None), + serde_json::json!({ "status": "shutdown" }), + ); + emit_marker(&observer, "captured-telemetry"); + let snapshot: Vec<_> = observer + .snapshot() + .into_iter() + .filter(|event| event.kind == "test_event") + .collect(); + assert_eq!(snapshot.len(), 1, "the synthetic snapshot contains only T"); + drop(observer); - assert_eq!(first.duration_since(started), OBSERVER_PUBLISH_INTERVAL); - assert_eq!(second.duration_since(first), OBSERVER_PUBLISH_INTERVAL); - } + let terminal_delivery_ok = run_relay_observer_publisher( + snapshot, + rx, + publisher, + agent_keys.clone(), + agent_keys.public_key().to_hex(), + owner_keys.public_key().to_hex(), + owner_keys.public_key(), + ) + .await; - #[tokio::test(start_paused = true)] - async fn limits_frames_in_each_rolling_minute() { - let mut pacer = ObserverPublishPacer::new(); - pacer.wait().await; - let first = tokio::time::Instant::now(); - for _ in 1..OBSERVER_PUBLISH_LIMIT_PER_MINUTE { - pacer.wait().await; + let mut published_kinds = Vec::new(); + while let Some(event) = published_rx.recv().await { + let payload: serde_json::Value = + decrypt_observer_payload(&owner_keys, &event).expect("decrypt published frame"); + published_kinds.push(payload["kind"].as_str().unwrap().to_string()); } + assert_eq!( + published_kinds + .iter() + .filter(|kind| kind.as_str() == "control_result") + .count(), + 1, + "C must publish exactly once even though the later T was the only snapshot member" + ); + assert!( + terminal_delivery_ok, + "confirmed terminal delivery must preserve successful shutdown verification" + ); + } - pacer.wait().await; - let ninety_first = tokio::time::Instant::now(); + #[tokio::test] + async fn failed_snapshot_terminal_publish_retries_live_overlap_before_acknowledgement() { + let observer = observer::ObserverHandle::in_process(); + let agent_keys = Keys::generate(); + let owner_keys = Keys::generate(); + let (publisher, mut published_rx) = + RelayEventPublisher::test_pair_dropping_first_terminal_confirmation(); + let rx = observer.subscribe(); + + observer.emit( + "control_result", + None, + &observer::context_for(None, None, None), + serde_json::json!({ "status": "retry-me" }), + ); + let snapshot = observer.snapshot(); + drop(observer); + + let terminal_delivery_ok = run_relay_observer_publisher( + snapshot, + rx, + publisher, + agent_keys.clone(), + agent_keys.public_key().to_hex(), + owner_keys.public_key().to_hex(), + owner_keys.public_key(), + ) + .await; + + let mut delivered_seqs = Vec::new(); + while let Some(event) = published_rx.recv().await { + let payload: serde_json::Value = + decrypt_observer_payload(&owner_keys, &event).expect("decrypt published frame"); + delivered_seqs.push(payload["seq"].as_u64().expect("observer sequence")); + } + assert_eq!( + delivered_seqs, + [1, 1], + "the unconfirmed snapshot attempt must release the same retained result for one retry" + ); + assert!( + terminal_delivery_ok, + "a later relay-confirmed retry must restore verified terminal delivery" + ); + } + + #[tokio::test] + async fn control_result_broadcast_lag_replays_every_terminal_result() { + let observer = observer::ObserverHandle::in_process(); + let agent_keys = Keys::generate(); + let owner_keys = Keys::generate(); + let (publisher, mut published_rx) = RelayEventPublisher::test_pair(); + let rx = observer.subscribe(); + + let burst = 64 + 17; + for marker in 0..burst { + observer.emit( + "control_result", + None, + &observer::context_for(None, None, None), + serde_json::json!({ "marker": marker }), + ); + } + drop(observer); + + let publisher_task = tokio::spawn(run_relay_observer_publisher( + Vec::new(), + rx, + publisher, + agent_keys.clone(), + agent_keys.public_key().to_hex(), + owner_keys.public_key().to_hex(), + owner_keys.public_key(), + )); + + let mut markers = Vec::new(); + while let Some(event) = published_rx.recv().await { + let payload: serde_json::Value = + decrypt_observer_payload(&owner_keys, &event).expect("decrypt published frame"); + markers.push(payload["payload"]["marker"].as_u64().unwrap() as usize); + } + let terminal_delivery_ok = publisher_task.await.expect("publisher must not panic"); + assert_eq!( + markers, + (0..burst).collect::>(), + "lag recovery must publish every accepted terminal result exactly once" + ); + assert!( + terminal_delivery_ok, + "replayed and acknowledged terminal results preserve verified delivery" + ); + } + + #[tokio::test(start_paused = true)] + async fn telemetry_broadcast_lag_does_not_taint_terminal_delivery_proof() { + let observer = observer::ObserverHandle::in_process(); + let agent_keys = Keys::generate(); + let owner_keys = Keys::generate(); + let (publisher, mut published_rx) = RelayEventPublisher::test_pair(); + let rx = observer.subscribe(); + + observer.emit( + "control_result", + None, + &observer::context_for(None, None, None), + serde_json::json!({ "status": "shutdown" }), + ); + for marker in 0..=2_000 { + observer.emit( + "acp_read", + None, + &observer::context_for(None, None, None), + serde_json::json!({ + "params": { + "update": { + "sessionUpdate": "agent_message_chunk", + "messageId": "lagged-telemetry", + "content": {"type": "text", "text": marker.to_string()}, + }, + }, + }), + ); + } + drop(observer); + + let publisher_task = tokio::spawn(run_relay_observer_publisher( + Vec::new(), + rx, + publisher, + agent_keys.clone(), + agent_keys.public_key().to_hex(), + owner_keys.public_key().to_hex(), + owner_keys.public_key(), + )); + tokio::time::advance(Duration::from_secs(2)).await; + + while published_rx.recv().await.is_some() {} + assert!( + publisher_task.await.expect("publisher must not panic"), + "ordinary telemetry loss must not invalidate acknowledged terminal delivery" + ); + } + + #[tokio::test(start_paused = true)] + async fn live_control_result_preempts_captured_telemetry_pacing() { + let observer = observer::ObserverHandle::in_process(); + let agent_keys = Keys::generate(); + let owner_keys = Keys::generate(); + let (publisher, mut published_rx) = RelayEventPublisher::test_pair(); + + for marker in 0..48 { + emit_marker(&observer, &marker.to_string()); + } + let rx = observer.subscribe(); + let snapshot = observer.snapshot(); + observer.emit( + "control_result", + None, + &observer::context_for(None, None, None), + serde_json::json!({ "status": "shutdown" }), + ); + + let task = tokio::spawn(run_relay_observer_publisher( + snapshot, + rx, + publisher, + agent_keys.clone(), + agent_keys.public_key().to_hex(), + owner_keys.public_key().to_hex(), + owner_keys.public_key(), + )); + + tokio::time::advance(Duration::from_secs(1)).await; + let event = published_rx + .recv() + .await + .expect("priority control result must publish before the snapshot drains"); + let payload: serde_json::Value = + decrypt_observer_payload(&owner_keys, &event).expect("decrypt published frame"); + assert_eq!(payload["kind"], "control_result"); + assert_eq!(payload["payload"]["status"], "shutdown"); + + task.abort(); + let _ = task.await; + } + + #[tokio::test(start_paused = true)] + async fn live_control_result_preempts_live_chunk_flush_pacing() { + let observer = observer::ObserverHandle::in_process(); + let agent_keys = Keys::generate(); + let owner_keys = Keys::generate(); + let (publisher, mut published_rx) = RelayEventPublisher::test_pair(); + let rx = observer.subscribe(); + let snapshot = observer.snapshot(); + observer.emit( + "acp_read", + None, + &observer::context_for(None, None, None), + serde_json::json!({ + "params": { + "update": { + "sessionUpdate": "agent_message_chunk", + "messageId": "message-1", + "content": {"type": "text", "text": "chunk"}, + }, + }, + }), + ); + + let task = tokio::spawn(run_relay_observer_publisher( + snapshot, + rx, + publisher, + agent_keys.clone(), + agent_keys.public_key().to_hex(), + owner_keys.public_key().to_hex(), + owner_keys.public_key(), + )); + tokio::time::advance(Duration::from_millis(100)).await; + assert!( + published_rx.try_recv().is_err(), + "the live chunk flush must still be in ordinary pacing" + ); + + observer.emit( + "control_result", + None, + &observer::context_for(None, None, None), + serde_json::json!({ "status": "shutdown" }), + ); + tokio::time::advance(Duration::from_millis(100)).await; + let event = published_rx + .recv() + .await + .expect("priority control result must interrupt live chunk pacing"); + let payload: serde_json::Value = + decrypt_observer_payload(&owner_keys, &event).expect("decrypt published frame"); + assert_eq!(payload["kind"], "control_result"); + assert_eq!(payload["payload"]["status"], "shutdown"); + + task.abort(); + let _ = task.await; + } +} + +#[cfg(test)] +mod observer_publish_pacer_tests { + use super::*; + + #[tokio::test(start_paused = true)] + async fn starts_without_a_burst_and_spaces_frames() { + let started = tokio::time::Instant::now(); + let mut pacer = ObserverPublishPacer::new(); + + pacer.wait().await; + let first = tokio::time::Instant::now(); + pacer.wait().await; + let second = tokio::time::Instant::now(); + + assert_eq!(first.duration_since(started), OBSERVER_PUBLISH_INTERVAL); + assert_eq!(second.duration_since(first), OBSERVER_PUBLISH_INTERVAL); + } + + #[tokio::test(start_paused = true)] + async fn limits_frames_in_each_rolling_minute() { + let mut pacer = ObserverPublishPacer::new(); + pacer.wait().await; + let first = tokio::time::Instant::now(); + for _ in 1..OBSERVER_PUBLISH_LIMIT_PER_MINUTE { + pacer.wait().await; + } + + pacer.wait().await; + let ninety_first = tokio::time::Instant::now(); assert_eq!(ninety_first.duration_since(first), Duration::from_secs(60)); } @@ -5186,12 +8477,14 @@ mod error_outcome_emission_tests { //! red. use super::*; - use crate::acp::{AcpClient, AcpError}; + use crate::acp::{AcpClient, AcpError, StopReason}; use crate::observer::ObserverHandle; use crate::pool::{ - AgentPool, OwnedAgent, PromptOutcome, PromptResult, PromptSource, TimeoutKind, + AgentPool, ChannelInfoResolver, ControlSignal, OwnedAgent, PromptContext, PromptOutcome, + PromptResult, PromptSource, SessionState, TimeoutKind, }; - use crate::queue::{BatchEvent, FlushBatch}; + use crate::queue::{BatchEvent, BatchOccurrenceIds, FlushBatch}; + use crate::relay::{ChannelInfo, RestClient}; use nostr::{EventBuilder, Keys, Kind}; use std::collections::HashSet; @@ -5243,124 +8536,1658 @@ mod error_outcome_emission_tests { } } - #[test] - fn normalizes_agent_name_from_initialize_result() { - assert_eq!( - normalized_agent_name(&serde_json::json!({ - "agentInfo": { "name": " Goose ", "version": "1.43.0" } - })), - "goose" + #[test] + fn normalizes_agent_name_from_initialize_result() { + assert_eq!( + normalized_agent_name(&serde_json::json!({ + "agentInfo": { "name": " Goose ", "version": "1.43.0" } + })), + "goose" + ); + assert_eq!( + normalized_agent_name(&serde_json::json!({ + "serverInfo": { "name": "buzz-agent" } + })), + "buzz-agent" + ); + } + + /// Spawn a real but inert agent subprocess (`cat`) so the error paths have + /// an `OwnedAgent` to move into respawn or return to the pool. The error + /// branches never talk to the subprocess. + async fn dummy_agent(index: usize) -> OwnedAgent { + OwnedAgent { + index, + acp: AcpClient::spawn("cat", &[], &[], false) + .await + .expect("spawn cat as inert agent"), + state: Default::default(), + model_capabilities: None, + desired_model: None, + model_overridden: false, + model_switch_request_id: None, + model_switch_rollback: None, + agent_name: "unknown".into(), + goose_system_prompt_supported: None, + // Error branches under test never read this; 1 is the legacy + // non-systemPrompt path, the simplest valid value. + protocol_version: 1, + } + } + + #[tokio::test] + async fn shutdown_cooperatively_recovers_and_reaps_active_prompt_and_heartbeat_agents() { + let mut pool = AgentPool::from_slots(vec![None, None]); + for (index, source) in [ + (0, PromptSource::Channel(Uuid::new_v4())), + (1, PromptSource::Heartbeat), + ] { + let agent = dummy_agent(index).await; + let result_tx = pool.result_tx(); + let (control_tx, control_rx) = tokio::sync::oneshot::channel(); + let channel_id = match &source { + PromptSource::Channel(channel_id) => Some(*channel_id), + PromptSource::Heartbeat => None, + }; + let task_source = source; + let abort_handle = pool.join_set.spawn(async move { + let signal = control_rx.await.expect("shutdown control"); + assert!(matches!(signal, ControlSignal::Cancel)); + let _ = result_tx.send(PromptResult { + agent, + source: task_source, + turn_id: format!("shutdown-turn-{index}"), + outcome: PromptOutcome::Cancelled, + batch: None, + retry_agent_index: None, + }); + }); + pool.task_map_mut().insert( + abort_handle.id(), + crate::pool::TaskMeta { + agent_index: index, + channel_id, + turn_id: format!("shutdown-turn-{index}"), + recoverable_batch: None, + desired_model: None, + model_overridden: false, + accepted_model_switch: None, + accepted_drop_control: None, + control_tx: Some(control_tx), + steer_tx: None, + }, + ); + } + + assert!( + tokio::time::timeout(Duration::from_secs(5), shutdown_agent_pool(&mut pool)) + .await + .expect("cooperative shutdown must be bounded"), + "every typed agent owner must be explicitly reaped with verified process-group absence" + ); + assert!(pool.join_set.is_empty()); + } + + fn register_completing_task(pool: &mut AgentPool, agent_index: usize, channel_id: Uuid) { + let task_id = pool.join_set.spawn(async {}).id(); + pool.task_map_mut().insert( + task_id, + crate::pool::TaskMeta { + agent_index, + channel_id: Some(channel_id), + turn_id: format!("turn-{agent_index}"), + recoverable_batch: None, + desired_model: None, + model_overridden: false, + accepted_model_switch: None, + accepted_drop_control: None, + control_tx: None, + steer_tx: None, + }, + ); + } + + fn one_event_batch(channel_id: Uuid, content: &str) -> FlushBatch { + FlushBatch { + channel_id, + events: vec![BatchEvent { + event: EventBuilder::new(Kind::Custom(9), content) + .sign_with_keys(&Keys::generate()) + .unwrap(), + prompt_tag: "test".into(), + received_at: std::time::Instant::now(), + }], + cancelled_events: vec![], + cancel_reason: Some(CancelReason::Interrupt), + occurrence_ids: BatchOccurrenceIds::for_test(1, 0), + } + } + + #[test] + fn ordinary_pool_capacity_deferral_preserves_retry_budget() { + let channel_id = Uuid::new_v4(); + let mut queue = EventQueue::new(config::DedupMode::Queue); + queue.push(crate::queue::QueuedEvent { + channel_id, + event: EventBuilder::new(Kind::Custom(9), "capacity deferred") + .sign_with_keys(&Keys::generate()) + .unwrap(), + received_at: std::time::Instant::now(), + prompt_tag: "test".into(), + }); + + let deferred = queue.flush_next().expect("queued batch"); + queue.set_retry_count_for_test(channel_id, crate::queue::MAX_RETRIES); + defer_unpinned_batch_for_capacity(&mut queue, deferred); + + let final_attempt = queue.flush_next().expect("deferred batch"); + assert!( + queue.requeue(final_attempt).is_some(), + "capacity deferral must not replenish an exhausted poison-batch retry budget" + ); + } + + #[tokio::test] + async fn control_preempted_setup_requeues_as_cancelled_without_charging_retry_budget() { + let channel_id = Uuid::new_v4(); + let original = EventBuilder::new(Kind::Custom(9), "setup-preempted") + .sign_with_keys(&Keys::generate()) + .unwrap(); + let agent = dummy_agent(0).await; + let mut pool = AgentPool::from_slots(vec![None]); + register_completing_task(&mut pool, 0, channel_id); + let mut queue = EventQueue::new(config::DedupMode::Queue); + queue.push(QueuedEvent { + channel_id, + event: original.clone(), + received_at: std::time::Instant::now(), + prompt_tag: "test".into(), + }); + let mut batch = queue.flush_next().expect("active setup batch"); + batch.cancel_reason = Some(CancelReason::Interrupt); + queue.set_retry_count_for_test(channel_id, crate::queue::MAX_RETRIES); + let config = test_config(); + let mut heartbeat_in_flight = false; + let removed_channels = HashSet::new(); + let mut crash_history = vec![SlotCircuit::new()]; + let (respawn_tx, _respawn_rx) = mpsc::channel(2); + let mut respawn_tasks = tokio::task::JoinSet::new(); + + let action = handle_prompt_result( + &mut pool, + &mut queue, + &config, + PromptResult { + agent, + source: PromptSource::Channel(channel_id), + turn_id: "turn-setup-preempted".into(), + outcome: PromptOutcome::ControlPreemptedSetup, + batch: Some(batch), + retry_agent_index: None, + }, + &mut heartbeat_in_flight, + &removed_channels, + &mut crash_history, + &respawn_tx, + &mut respawn_tasks, + None, + None, + ); + + assert!(matches!(action, LoopAction::Continue)); + let requeued = queue + .flush_next() + .expect("control-preempted setup must preserve the batch"); + let preserved_ids = requeued + .events + .iter() + .chain(requeued.cancelled_events.iter()) + .map(|event| event.event.id) + .collect::>(); + assert_eq!( + preserved_ids, + [original.id], + "even an exhausted retry budget must preserve the setup-preempted event exactly once" + ); + assert_eq!(requeued.cancel_reason, Some(CancelReason::Interrupt)); + assert_eq!( + crash_history[0].crash_times.len(), + 0, + "cooperative setup preemption must not consume crash budget" + ); + + respawn_tasks.shutdown().await; + } + + #[tokio::test] + async fn busy_switch_model_retry_pins_original_slot_not_other_idle_agent() { + let channel_id = Uuid::new_v4(); + let mut agent = dummy_agent(0).await; + agent.desired_model = Some("runtime-model".into()); + agent.model_overridden = true; + let other = dummy_agent(1).await; + let mut pool = AgentPool::from_slots(vec![None, Some(other)]); + register_completing_task(&mut pool, 0, channel_id); + let mut queue = EventQueue::new(config::DedupMode::Queue); + let config = test_config(); + let mut heartbeat_in_flight = false; + let mut crash_history = vec![SlotCircuit::new(), SlotCircuit::new()]; + let (respawn_tx, _respawn_rx) = mpsc::channel(2); + let mut respawn_tasks = tokio::task::JoinSet::new(); + + let action = handle_prompt_result( + &mut pool, + &mut queue, + &config, + PromptResult { + agent, + source: PromptSource::Channel(channel_id), + turn_id: "turn-0".into(), + outcome: PromptOutcome::Cancelled, + batch: Some(one_event_batch(channel_id, "switch-model-retry")), + retry_agent_index: Some(0), + }, + &mut heartbeat_in_flight, + &HashSet::new(), + &mut crash_history, + &respawn_tx, + &mut respawn_tasks, + None, + None, + ); + + assert!(action == LoopAction::Continue); + assert_eq!(queue.required_agent(channel_id), Some(0)); + assert!( + pool.slot_alive(1), + "unrelated idle slot must stay unclaimed" + ); + let exact = pool + .try_claim_index(0) + .expect("the replay must claim its original slot"); + assert_eq!(exact.desired_model.as_deref(), Some("runtime-model")); + pool.return_agent(exact); + shutdown_agent_pool(&mut pool).await; + } + + #[tokio::test] + async fn post_model_setup_exit_preserves_pending_switch_for_replacement_retry() { + let channel_id = Uuid::new_v4(); + let mut agent = dummy_agent(0).await; + agent.desired_model = Some("model-b".into()); + agent.model_overridden = true; + agent.model_switch_request_id = Some("0123456789abcdef0123456789abcdef".into()); + agent.model_switch_rollback = Some(Box::new(ModelSwitchRollback { + desired_model: Some("model-a".into()), + model_overridden: false, + request_id: None, + previous: None, + })); + let mut pool = AgentPool::from_slots(vec![None]); + register_completing_task(&mut pool, 0, channel_id); + let mut queue = EventQueue::new(config::DedupMode::Queue); + let config = test_config(); + let mut heartbeat_in_flight = false; + let mut crash_history = vec![SlotCircuit::new()]; + let (respawn_tx, _respawn_rx) = mpsc::channel(1); + let mut respawn_tasks = tokio::task::JoinSet::new(); + + let action = handle_prompt_result( + &mut pool, + &mut queue, + &config, + PromptResult { + agent, + source: PromptSource::Channel(channel_id), + turn_id: "turn-post-model-setup-exit".into(), + outcome: PromptOutcome::AgentExited, + batch: Some(one_event_batch(channel_id, "retry pending switch")), + retry_agent_index: None, + }, + &mut heartbeat_in_flight, + &HashSet::new(), + &mut crash_history, + &respawn_tx, + &mut respawn_tasks, + None, + None, + ); + + assert!(action == LoopAction::Continue); + assert_eq!( + queue.required_agent(channel_id), + Some(0), + "the requeued turn must stay pinned to the slot carrying its pending switch" + ); + let pending = crash_history[0] + .pending_model_intent + .as_ref() + .expect("fatal setup failure must retain pending switch state"); + assert_eq!(pending.desired_model.as_deref(), Some("model-b")); + assert!(pending.model_overridden); + assert_eq!( + pending.model_switch_request_id.as_deref(), + Some("0123456789abcdef0123456789abcdef") + ); + let rollback = pending + .model_switch_rollback + .as_ref() + .expect("replacement retry must retain prior-model rollback"); + assert_eq!(rollback.desired_model.as_deref(), Some("model-a")); + assert!(!rollback.model_overridden); + + respawn_tasks.shutdown().await; + } + + #[tokio::test] + async fn successful_session_retirement_keeps_override_pin_until_next_turn_succeeds() { + let channel_id = Uuid::new_v4(); + let mut agent = dummy_agent(0).await; + agent.desired_model = Some("runtime-model".into()); + agent.model_overridden = true; + let mut pool = AgentPool::from_slots(vec![None]); + register_completing_task(&mut pool, 0, channel_id); + let mut queue = EventQueue::new(config::DedupMode::Queue); + let config = test_config(); + let mut heartbeat_in_flight = false; + let mut crash_history = vec![SlotCircuit::new()]; + let (respawn_tx, _respawn_rx) = mpsc::channel(1); + let mut respawn_tasks = tokio::task::JoinSet::new(); + + handle_prompt_result( + &mut pool, + &mut queue, + &config, + PromptResult { + agent, + source: PromptSource::Channel(channel_id), + turn_id: "turn-0".into(), + outcome: PromptOutcome::SessionRetired(StopReason::EndTurn), + batch: None, + retry_agent_index: None, + }, + &mut heartbeat_in_flight, + &HashSet::new(), + &mut crash_history, + &respawn_tx, + &mut respawn_tasks, + None, + None, + ); + assert_eq!(queue.required_agent(channel_id), Some(0)); + + let replay_agent = pool.try_claim_index(0).expect("pinned replay agent"); + register_completing_task(&mut pool, 0, channel_id); + handle_prompt_result( + &mut pool, + &mut queue, + &config, + PromptResult { + agent: replay_agent, + source: PromptSource::Channel(channel_id), + turn_id: "turn-0-replay".into(), + outcome: PromptOutcome::Ok(StopReason::EndTurn), + batch: None, + retry_agent_index: None, + }, + &mut heartbeat_in_flight, + &HashSet::new(), + &mut crash_history, + &respawn_tx, + &mut respawn_tasks, + None, + None, + ); + assert_eq!( + queue.required_agent(channel_id), + None, + "a successful fresh-session turn can fall back to session affinity" + ); + shutdown_agent_pool(&mut pool).await; + } + + #[tokio::test] + async fn completed_before_switch_model_keeps_no_batch_pin_and_healthy_return_wakes_it() { + let source_channel = Uuid::new_v4(); + let waiting_channel = Uuid::new_v4(); + let mut agent = dummy_agent(0).await; + agent.desired_model = Some("runtime-model".into()); + agent.model_overridden = true; + let mut pool = AgentPool::from_slots(vec![None]); + register_completing_task(&mut pool, 0, source_channel); + let mut queue = EventQueue::new(config::DedupMode::Queue); + queue.push(QueuedEvent { + channel_id: waiting_channel, + event: EventBuilder::new(Kind::Custom(9), "waiting") + .sign_with_keys(&Keys::generate()) + .unwrap(), + received_at: std::time::Instant::now(), + prompt_tag: "test".into(), + }); + queue.require_agent(waiting_channel, 0); + queue.block_required_agent(waiting_channel); + let config = test_config(); + let mut heartbeat_in_flight = false; + let mut crash_history = vec![SlotCircuit::new()]; + let (respawn_tx, _respawn_rx) = mpsc::channel(1); + let mut respawn_tasks = tokio::task::JoinSet::new(); + + handle_prompt_result( + &mut pool, + &mut queue, + &config, + PromptResult { + agent, + source: PromptSource::Channel(source_channel), + turn_id: "turn-0".into(), + outcome: PromptOutcome::Ok(StopReason::EndTurn), + batch: None, + retry_agent_index: Some(0), + }, + &mut heartbeat_in_flight, + &HashSet::new(), + &mut crash_history, + &respawn_tx, + &mut respawn_tasks, + None, + None, + ); + + assert_eq!(queue.required_agent(source_channel), Some(0)); + assert_eq!( + queue + .flush_next() + .expect("healthy slot return must wake other pinned work") + .channel_id, + waiting_channel + ); + shutdown_agent_pool(&mut pool).await; + } + + #[tokio::test] + async fn acknowledged_switch_model_is_reconciled_when_prompt_result_wins_ready_race() { + let channel_id = Uuid::new_v4(); + let mut agent = dummy_agent(0).await; + agent + .state + .sessions + .insert(channel_id, "session-before-switch".into()); + let mut pool = AgentPool::from_slots(vec![None]); + let task_id = pool.join_set.spawn(async {}).id(); + pool.task_map_mut().insert( + task_id, + crate::pool::TaskMeta { + agent_index: 0, + channel_id: Some(channel_id), + turn_id: "turn-ready-race".into(), + recoverable_batch: None, + desired_model: Some("runtime-model".into()), + model_overridden: true, + accepted_model_switch: Some(ModelSwitchRequest::new( + "runtime-model", + "0123456789abcdef0123456789abcdef", + )), + accepted_drop_control: None, + control_tx: None, + steer_tx: None, + }, + ); + let mut queue = EventQueue::new(config::DedupMode::Queue); + let config = test_config(); + let mut heartbeat_in_flight = false; + let mut crash_history = vec![SlotCircuit::new()]; + let (respawn_tx, _respawn_rx) = mpsc::channel(1); + let mut respawn_tasks = tokio::task::JoinSet::new(); + + handle_prompt_result( + &mut pool, + &mut queue, + &config, + PromptResult { + agent, + source: PromptSource::Channel(channel_id), + turn_id: "turn-ready-race".into(), + outcome: PromptOutcome::Ok(StopReason::EndTurn), + batch: None, + retry_agent_index: None, + }, + &mut heartbeat_in_flight, + &HashSet::new(), + &mut crash_history, + &respawn_tx, + &mut respawn_tasks, + None, + None, + ); + + assert_eq!( + queue.required_agent(channel_id), + Some(0), + "an acknowledged switch must remain pinned to the accepting slot" + ); + assert_eq!( + pool.live_count(), + 0, + "the adapter still owns the old session and must not return idle" + ); + assert!( + crash_history[0].crash_times.is_empty(), + "an acknowledged lifecycle recycle must not consume crash budget" + ); + assert_eq!(respawn_tasks.len(), 1); + assert_eq!( + crash_history[0] + .pending_model_intent + .as_ref() + .and_then(|intent| intent.desired_model.as_deref()), + Some("runtime-model") + ); + assert_eq!( + crash_history[0] + .pending_model_intent + .as_ref() + .and_then(|intent| intent.model_switch_request_id.as_deref()), + Some("0123456789abcdef0123456789abcdef") + ); + respawn_tasks.shutdown().await; + } + + #[tokio::test] + async fn unsupported_switch_model_ready_race_restores_prior_intent_without_recycle() { + let channel_id = Uuid::new_v4(); + let mut agent = dummy_agent(0).await; + let observer = observer::ObserverHandle::in_process(); + agent.acp.set_observer(Some(observer.clone()), 0); + agent + .state + .sessions + .insert(channel_id, "session-before-switch".into()); + agent.model_capabilities = Some(pool::AgentModelCapabilities { + config_options_raw: vec![serde_json::json!({ + "configId": "model", + "category": "model", + "options": [{"value": "model-a"}], + })], + available_models_raw: None, + }); + agent.desired_model = Some("model-a".into()); + agent.model_overridden = true; + + let mut pool = AgentPool::from_slots(vec![None]); + let task_id = pool.join_set.spawn(async {}).id(); + pool.task_map_mut().insert( + task_id, + crate::pool::TaskMeta { + agent_index: 0, + channel_id: Some(channel_id), + turn_id: "turn-unsupported-ready-race".into(), + recoverable_batch: None, + // The sender currently records the accepted request here for + // panic recovery. The returned agent remains the authority for + // the prior, already-applied model intent. + desired_model: Some("model-b".into()), + model_overridden: true, + accepted_model_switch: Some(ModelSwitchRequest::new( + "model-b", + "abcdef0123456789abcdef0123456789", + )), + accepted_drop_control: None, + control_tx: None, + steer_tx: None, + }, + ); + let mut queue = EventQueue::new(config::DedupMode::Queue); + let config = test_config(); + let mut heartbeat_in_flight = false; + let mut crash_history = vec![SlotCircuit::new()]; + let (respawn_tx, _respawn_rx) = mpsc::channel(1); + let mut respawn_tasks = tokio::task::JoinSet::new(); + + handle_prompt_result( + &mut pool, + &mut queue, + &config, + PromptResult { + agent, + source: PromptSource::Channel(channel_id), + turn_id: "turn-unsupported-ready-race".into(), + outcome: PromptOutcome::Ok(StopReason::EndTurn), + batch: None, + retry_agent_index: None, + }, + &mut heartbeat_in_flight, + &HashSet::new(), + &mut crash_history, + &respawn_tx, + &mut respawn_tasks, + None, + None, + ); + + let live_count = pool.live_count(); + let required_agent = queue.required_agent(channel_id); + let respawn_count = respawn_tasks.len(); + let mut returned_agent = pool.try_claim_index(0); + let returned_intent = returned_agent.as_ref().map(|agent| { + ( + agent.desired_model.clone(), + agent.model_overridden, + agent.state.sessions.get(&channel_id).cloned(), + ) + }); + if let Some(agent) = returned_agent.as_mut() { + agent + .acp + .shutdown() + .await + .expect("returned adapter must shut down cleanly"); + } + respawn_tasks.shutdown().await; + + assert_eq!( + returned_intent, + Some(( + Some("model-a".to_string()), + true, + Some("session-before-switch".to_string()), + )), + "unsupported model B must preserve the prior applied model A and session" + ); + assert_eq!(live_count, 1, "the healthy prior-model adapter stays idle"); + assert_eq!( + required_agent, None, + "unsupported model B must not create an exact-slot retry" + ); + assert_eq!( + respawn_count, 0, + "unsupported model B must not recycle the healthy adapter" + ); + let unsupported = observer + .snapshot() + .into_iter() + .find(|event| event.kind == "control_result") + .expect("ready-race rejection must emit a terminal control result"); + assert_eq!(unsupported.payload["status"], "unsupported_model"); + assert_eq!(unsupported.payload["modelId"], "model-b"); + assert_eq!( + unsupported.payload["requestId"], + "abcdef0123456789abcdef0123456789" + ); + } + + #[tokio::test] + async fn acknowledged_rotate_recycles_when_prompt_result_wins_ready_race() { + let channel_id = Uuid::new_v4(); + let mut agent = dummy_agent(0).await; + agent + .state + .sessions + .insert(channel_id, "session-before-rotate".into()); + let mut pool = AgentPool::from_slots(vec![None]); + let task_id = pool.join_set.spawn(async {}).id(); + pool.task_map_mut().insert( + task_id, + crate::pool::TaskMeta { + agent_index: 0, + channel_id: Some(channel_id), + turn_id: "turn-rotate-ready-race".into(), + recoverable_batch: None, + desired_model: None, + model_overridden: false, + accepted_model_switch: None, + accepted_drop_control: Some(AcceptedDropControl::Rotate), + control_tx: None, + steer_tx: None, + }, + ); + let mut queue = EventQueue::new(config::DedupMode::Queue); + let config = test_config(); + let mut heartbeat_in_flight = false; + let mut crash_history = vec![SlotCircuit::new()]; + let (respawn_tx, _respawn_rx) = mpsc::channel(1); + let mut respawn_tasks = tokio::task::JoinSet::new(); + + handle_prompt_result( + &mut pool, + &mut queue, + &config, + PromptResult { + agent, + source: PromptSource::Channel(channel_id), + turn_id: "turn-rotate-ready-race".into(), + outcome: PromptOutcome::Ok(StopReason::EndTurn), + batch: None, + retry_agent_index: None, + }, + &mut heartbeat_in_flight, + &HashSet::new(), + &mut crash_history, + &respawn_tx, + &mut respawn_tasks, + None, + None, + ); + + assert_eq!(pool.live_count(), 0); + assert_eq!(respawn_tasks.len(), 1); + assert!(crash_history[0].crash_times.is_empty()); + respawn_tasks.shutdown().await; + } + + #[tokio::test] + async fn acknowledged_rotate_drops_batch_when_error_wins_ready_race() { + let channel_id = Uuid::new_v4(); + let mut agent = dummy_agent(0).await; + agent + .state + .sessions + .insert(channel_id, "session-before-rotate".into()); + let mut batch = one_event_batch(channel_id, "must-not-replay"); + batch.cancel_reason = None; + let mut pool = AgentPool::from_slots(vec![None]); + let task_id = pool.join_set.spawn(async {}).id(); + pool.task_map_mut().insert( + task_id, + crate::pool::TaskMeta { + agent_index: 0, + channel_id: Some(channel_id), + turn_id: "turn-rotate-error-race".into(), + recoverable_batch: None, + desired_model: None, + model_overridden: false, + accepted_model_switch: None, + accepted_drop_control: Some(AcceptedDropControl::Rotate), + control_tx: None, + steer_tx: None, + }, + ); + let mut queue = EventQueue::new(config::DedupMode::Queue); + let config = test_config(); + let mut heartbeat_in_flight = false; + let mut crash_history = vec![SlotCircuit::new()]; + let (respawn_tx, _respawn_rx) = mpsc::channel(1); + let mut respawn_tasks = tokio::task::JoinSet::new(); + + handle_prompt_result( + &mut pool, + &mut queue, + &config, + PromptResult { + agent, + source: PromptSource::Channel(channel_id), + turn_id: "turn-rotate-error-race".into(), + outcome: PromptOutcome::Error(AcpError::AgentError { + code: -32000, + message: "application failure won ready race".into(), + }), + batch: Some(batch), + retry_agent_index: None, + }, + &mut heartbeat_in_flight, + &HashSet::new(), + &mut crash_history, + &respawn_tx, + &mut respawn_tasks, + None, + None, + ); + + assert_eq!( + queue.pending_channels(), + 0, + "an accepted rotate must drop the losing prompt outcome's batch" + ); + assert_eq!(pool.live_count(), 0); + assert_eq!(respawn_tasks.len(), 1); + assert!(crash_history[0].crash_times.is_empty()); + respawn_tasks.shutdown().await; + } + + #[tokio::test] + async fn acknowledged_cancel_drops_batch_when_error_wins_ready_race() { + let channel_id = Uuid::new_v4(); + let agent = dummy_agent(0).await; + let mut batch = one_event_batch(channel_id, "must-not-replay"); + batch.cancel_reason = None; + let mut pool = AgentPool::from_slots(vec![None]); + let task_id = pool.join_set.spawn(async {}).id(); + pool.task_map_mut().insert( + task_id, + crate::pool::TaskMeta { + agent_index: 0, + channel_id: Some(channel_id), + turn_id: "turn-cancel-ready-race".into(), + recoverable_batch: None, + desired_model: None, + model_overridden: false, + accepted_model_switch: None, + accepted_drop_control: Some(AcceptedDropControl::Cancel), + control_tx: None, + steer_tx: None, + }, + ); + let mut queue = EventQueue::new(config::DedupMode::Queue); + let config = test_config(); + let mut heartbeat_in_flight = false; + let mut crash_history = vec![SlotCircuit::new()]; + let (respawn_tx, _respawn_rx) = mpsc::channel(1); + let mut respawn_tasks = tokio::task::JoinSet::new(); + + handle_prompt_result( + &mut pool, + &mut queue, + &config, + PromptResult { + agent, + source: PromptSource::Channel(channel_id), + turn_id: "turn-cancel-ready-race".into(), + outcome: PromptOutcome::Error(AcpError::AgentError { + code: -32000, + message: "application failure won ready race".into(), + }), + batch: Some(batch), + retry_agent_index: None, + }, + &mut heartbeat_in_flight, + &HashSet::new(), + &mut crash_history, + &respawn_tx, + &mut respawn_tasks, + None, + None, + ); + + assert_eq!( + queue.pending_channels(), + 0, + "an accepted cancel must dominate the losing prompt outcome's retry policy" + ); + assert_eq!(pool.live_count(), 1); + shutdown_agent_pool(&mut pool).await; + } + + #[tokio::test] + async fn membership_removed_while_checked_out_recycles_remote_session_owner() { + let channel_id = Uuid::new_v4(); + let mut agent = dummy_agent(0).await; + agent + .state + .sessions + .insert(channel_id, "session-before-removal".into()); + let mut pool = AgentPool::from_slots(vec![None]); + register_completing_task(&mut pool, 0, channel_id); + let mut queue = EventQueue::new(config::DedupMode::Queue); + let config = test_config(); + let mut heartbeat_in_flight = false; + let mut crash_history = vec![SlotCircuit::new()]; + let (respawn_tx, _respawn_rx) = mpsc::channel(1); + let mut respawn_tasks = tokio::task::JoinSet::new(); + + handle_prompt_result( + &mut pool, + &mut queue, + &config, + PromptResult { + agent, + source: PromptSource::Channel(channel_id), + turn_id: "turn-membership-removed".into(), + outcome: PromptOutcome::Ok(StopReason::EndTurn), + batch: None, + retry_agent_index: None, + }, + &mut heartbeat_in_flight, + &HashSet::from([channel_id]), + &mut crash_history, + &respawn_tx, + &mut respawn_tasks, + None, + None, + ); + + assert_eq!(pool.live_count(), 0); + assert_eq!(respawn_tasks.len(), 1); + assert!(crash_history[0].crash_times.is_empty()); + assert_eq!(queue.required_agent(channel_id), None); + respawn_tasks.shutdown().await; + } + + #[tokio::test] + async fn rejected_session_close_preserves_batch_and_respawns_poisoned_adapter() { + let capture = + std::env::temp_dir().join(format!("buzz-acp-close-rejected-{}.ndjson", Uuid::new_v4())); + let script = r#" + capture="$1" + IFS= read -r initialize + printf '%s\n' "$initialize" >> "$capture" + printf '%s\n' '{"jsonrpc":"2.0","id":0,"result":{"protocolVersion":2,"agentCapabilities":{"sessionCapabilities":{"close":{}}}}}' + IFS= read -r prompt + printf '%s\n' "$prompt" >> "$capture" + IFS= read -r cancel + printf '%s\n' "$cancel" >> "$capture" + printf '%s\n' '{"jsonrpc":"2.0","id":1,"result":{"stopReason":"cancelled"}}' + IFS= read -r close + printf '%s\n' "$close" >> "$capture" + printf '%s\n' '{"jsonrpc":"2.0","id":2,"error":{"code":-32601,"message":"Method not found"}}' + sleep 10 + "#; + let mut acp = AcpClient::spawn( + "bash", + &[ + "-c".to_string(), + script.to_string(), + "buzz-acp-close-rejected-test".to_string(), + capture.to_string_lossy().into_owned(), + ], + &[], + false, + ) + .await + .expect("failed to spawn fake ACP adapter"); + acp.initialize() + .await + .expect("fake adapter must advertise session/close"); + assert!(acp.supports_session_close()); + + let channel_id = Uuid::new_v4(); + let mut state = SessionState::default(); + state.sessions.insert(channel_id, "session-old".into()); + let agent = OwnedAgent { + index: 0, + acp, + state, + model_capabilities: None, + desired_model: None, + model_overridden: false, + model_switch_request_id: None, + model_switch_rollback: None, + agent_name: "claude-code-acp".into(), + goose_system_prompt_supported: None, + protocol_version: 2, + }; + let keys = Keys::generate(); + let rest_client = RestClient { + http: reqwest::Client::new(), + base_url: "http://127.0.0.1:0".into(), + keys: keys.clone(), + auth_tag_json: None, + }; + let ctx = Arc::new(PromptContext { + mcp_servers: vec![], + initial_message: None, + idle_timeout: Duration::from_secs(60), + max_turn_duration: Duration::from_secs(120), + turn_liveness_interval: Duration::ZERO, + dedup_mode: config::DedupMode::Queue, + system_prompt: None, + session_title: None, + team_instructions: None, + heartbeat_prompt: None, + base_prompt: None, + cwd: ".".into(), + rest_client: rest_client.clone(), + channel_info: ChannelInfoResolver::new( + std::collections::HashMap::from([( + channel_id, + ChannelInfo { + name: "test".into(), + channel_type: "public".into(), + }, + )]), + rest_client, + ), + context_message_limit: 0, + max_turns_per_session: 0, + permission_mode: config::PermissionMode::Default, + agent_keys: keys.clone(), + agent_owner_pubkey: None, + memory_enabled: false, + harness_name: "claude-code-acp".into(), + relay_url: "ws://127.0.0.1:3000".into(), + }); + let event = EventBuilder::new(Kind::Custom(9), "original") + .sign_with_keys(&keys) + .unwrap(); + let event_id = event.id; + let batch = FlushBatch { + channel_id, + events: vec![BatchEvent { + event, + prompt_tag: "test".into(), + received_at: std::time::Instant::now(), + }], + cancelled_events: vec![], + cancel_reason: None, + occurrence_ids: BatchOccurrenceIds::for_test(1, 0), + }; + let (result_tx, mut result_rx) = mpsc::unbounded_channel(); + let (control_tx, control_rx) = tokio::sync::oneshot::channel(); + let task = tokio::spawn(crate::pool::run_prompt_task( + agent, + Some(batch), + Some("prompt".into()), + ctx, + result_tx, + Some(control_rx), + "turn-close-rejected".into(), + )); + + tokio::time::timeout(Duration::from_secs(5), async { + loop { + if std::fs::read_to_string(&capture) + .is_ok_and(|contents| contents.lines().count() >= 2) + { + break; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("prompt request must reach the fake adapter"); + control_tx + .send(ControlSignal::Steer) + .expect("control receiver must still be live"); + let result = tokio::time::timeout(Duration::from_secs(5), result_rx.recv()) + .await + .expect("prompt task must return after close rejection") + .expect("prompt result channel must stay open"); + task.await.expect("prompt task must not panic"); + assert!(matches!( + result.outcome, + PromptOutcome::SessionCloseFailed(AcpError::AgentError { code: -32601, .. }) + )); + let captured = std::fs::read_to_string(&capture).expect("capture must remain readable"); + let messages: Vec = captured + .lines() + .map(|line| serde_json::from_str(line).expect("captured line must be JSON")) + .collect(); + assert_eq!(messages[0]["method"], "initialize"); + assert_eq!(messages[1]["method"], "session/prompt"); + assert_eq!(messages[2]["method"], "session/cancel"); + assert_eq!(messages[3]["method"], "session/close"); + + assert_eq!( + result + .agent + .state + .sessions + .get(&channel_id) + .map(String::as_str), + Some("session-old"), + "failed close must not silently discard local session ownership" + ); + + let mut pool = AgentPool::from_slots(vec![None]); + let task_id = pool.join_set.spawn(async {}).id(); + pool.task_map_mut().insert( + task_id, + crate::pool::TaskMeta { + agent_index: 0, + channel_id: Some(channel_id), + turn_id: "turn-close-rejected".into(), + recoverable_batch: None, + desired_model: None, + model_overridden: false, + accepted_model_switch: None, + accepted_drop_control: None, + control_tx: None, + steer_tx: None, + }, + ); + let mut queue = EventQueue::new(config::DedupMode::Queue); + let config = test_config(); + let mut heartbeat_in_flight = false; + let removed_channels = HashSet::new(); + let mut crash_history = vec![SlotCircuit::new()]; + let (respawn_tx, _respawn_rx) = mpsc::channel(8); + let mut respawn_tasks = tokio::task::JoinSet::new(); + + handle_prompt_result( + &mut pool, + &mut queue, + &config, + result, + &mut heartbeat_in_flight, + &removed_channels, + &mut crash_history, + &respawn_tx, + &mut respawn_tasks, + None, + None, + ); + + assert_eq!( + pool.live_count(), + 0, + "an adapter with uncertain remote ownership must never return idle" + ); + assert_eq!( + respawn_tasks.len(), + 1, + "close rejection must enter process-group shutdown and respawn" + ); + let requeued = queue + .flush_next() + .expect("steer-cancelled work must remain queued"); + assert_eq!( + requeued.events.len(), + 1, + "with no newer event to merge, the queue re-dispatches the cancelled batch" + ); + assert!(requeued.cancelled_events.is_empty()); + assert_eq!(requeued.events[0].event.id, event_id); + assert_eq!(requeued.cancel_reason, Some(CancelReason::Steer)); + + let _ = std::fs::remove_file(capture); + } + + struct OutcomeDisposition { + turn_errors: usize, + live_agents: usize, + respawn_tasks: usize, + recent_crashes: usize, + circuit_open: bool, + turn_error_texts: Vec, + } + + /// Drive one outcome through `handle_prompt_result` and report its + /// observable supervisor disposition. + async fn outcome_disposition_for(outcome: PromptOutcome) -> OutcomeDisposition { + let agent = dummy_agent(0).await; + let mut pool = AgentPool::from_slots(vec![None]); + + // `handle_prompt_result` asserts it removes exactly one in-flight task + // for the completing agent (the slot was checked out, not idle). Mirror + // the real dispatch path by registering a TaskMeta keyed on a genuine + // `task::Id` — only obtainable from inside a spawned task. + let task_id = pool.join_set.spawn(async {}).id(); + pool.task_map_mut().insert( + task_id, + crate::pool::TaskMeta { + agent_index: 0, + channel_id: None, + turn_id: "test-turn-id".to_string(), + recoverable_batch: None, + desired_model: None, + model_overridden: false, + accepted_model_switch: None, + accepted_drop_control: None, + control_tx: None, + steer_tx: None, + }, + ); + + let mut queue = EventQueue::new(config::DedupMode::Queue); + let config = test_config(); + let mut heartbeat_in_flight = false; + let removed_channels = HashSet::new(); + let mut crash_history = vec![SlotCircuit::new()]; + let (respawn_tx, _respawn_rx) = mpsc::channel(8); + let mut respawn_tasks = tokio::task::JoinSet::new(); + let observer = ObserverHandle::in_process(); + + let result = PromptResult { + agent, + source: PromptSource::Channel(Uuid::new_v4()), + turn_id: "test-turn-id".to_string(), + outcome, + batch: None, + retry_agent_index: None, + }; + + handle_prompt_result( + &mut pool, + &mut queue, + &config, + result, + &mut heartbeat_in_flight, + &removed_channels, + &mut crash_history, + &respawn_tx, + &mut respawn_tasks, + Some(observer.clone()), + None, + ); + + let turn_errors: Vec<_> = observer + .snapshot() + .into_iter() + .filter(|e| e.kind == "turn_error") + .collect(); + assert!( + turn_errors + .iter() + .all(|event| event.turn_id.as_deref() == Some("test-turn-id")), + "turn_error must retain the completed turn id" + ); + OutcomeDisposition { + turn_errors: turn_errors.len(), + live_agents: pool.live_count(), + respawn_tasks: respawn_tasks.len(), + recent_crashes: crash_history[0].crash_times.len(), + circuit_open: crash_history[0].open_until.is_some(), + turn_error_texts: turn_errors + .iter() + .filter_map(|event| event.payload["error"].as_str().map(str::to_owned)) + .collect(), + } + } + + /// Drive one error outcome through `handle_prompt_result` and return how + /// many `turn_error` events it emitted to the observer feed. + async fn turn_errors_emitted_for(outcome: PromptOutcome) -> usize { + outcome_disposition_for(outcome).await.turn_errors + } + + #[tokio::test] + async fn agent_exited_emits_exactly_one_feed_event() { + assert_eq!(turn_errors_emitted_for(PromptOutcome::AgentExited).await, 1); + } + + #[tokio::test] + async fn panic_event_retains_task_turn_id() { + let mut pool = AgentPool::from_slots(vec![]); + let channel_id = Uuid::new_v4(); + let (started_tx, started_rx) = tokio::sync::oneshot::channel(); + let abort_handle = pool.join_set.spawn(async move { + let _ = started_tx.send(()); + std::future::pending::<()>().await; + }); + let task_id = abort_handle.id(); + pool.task_map_mut().insert( + task_id, + crate::pool::TaskMeta { + agent_index: 0, + channel_id: Some(channel_id), + turn_id: "panic-turn-id".to_string(), + recoverable_batch: None, + desired_model: None, + model_overridden: false, + accepted_model_switch: None, + accepted_drop_control: None, + control_tx: None, + steer_tx: None, + }, + ); + started_rx.await.unwrap(); + abort_handle.abort(); + let join_error = pool.join_set.join_next().await.unwrap().unwrap_err(); + + let mut queue = EventQueue::new(config::DedupMode::Queue); + let config = test_config(); + let mut heartbeat_in_flight = false; + let removed_channels = HashSet::new(); + let mut typing_channels = HashMap::new(); + let mut crash_history = vec![SlotCircuit::new()]; + let (respawn_tx, _respawn_rx) = mpsc::channel(8); + let mut respawn_tasks = tokio::task::JoinSet::new(); + let observer = ObserverHandle::in_process(); + + recover_panicked_agent( + &mut pool, + &mut queue, + &config, + join_error, + &mut heartbeat_in_flight, + &removed_channels, + &mut typing_channels, + &mut crash_history, + &respawn_tx, + &mut respawn_tasks, + Some(observer.clone()), + None, + ); + + let panic = observer + .snapshot() + .into_iter() + .find(|event| event.kind == "agent_panic") + .expect("panic recovery emits an observer event"); + assert_eq!( + panic.channel_id.as_deref(), + Some(channel_id.to_string().as_str()) + ); + assert_eq!(panic.turn_id.as_deref(), Some("panic-turn-id")); + } + + #[test] + fn cleanup_unverified_quarantine_cannot_be_bypassed_by_refill_cooldown() { + let mut slot = SlotCircuit::new(); + slot.mark_cleanup_unverified(); + slot.open_until = Some(std::time::Instant::now() - Duration::from_secs(1)); + + assert!( + !slot.can_refill(), + "elapsed crash cooldown must not authorize overlap after cleanup became unverified" + ); + assert!( + slot.blocks_supervisor_exit(), + "the process must remain alive and degraded instead of restarting into possible overlap" + ); + } + + #[tokio::test] + async fn automatic_exit_waits_in_quarantine_until_explicit_shutdown() { + let (shutdown_tx, mut shutdown_rx) = watch::channel(()); + let mut wait = Box::pin(await_automatic_exit_permission( + &mut shutdown_rx, + true, + "deterministic test exit", + )); + + assert!( + tokio::time::timeout(Duration::from_millis(25), &mut wait) + .await + .is_err(), + "automatic failure must not let the supervisor exit while cleanup is unverified" + ); + + shutdown_tx.send(()).expect("explicit shutdown"); + tokio::time::timeout(Duration::from_secs(1), wait) + .await + .expect("explicit shutdown releases quarantine"); + } + + #[tokio::test] + async fn accepted_cancel_drops_recoverable_batch_when_task_panics() { + let mut pool = AgentPool::from_slots(vec![None]); + let channel_id = Uuid::new_v4(); + let mut batch = one_event_batch(channel_id, "must-not-replay"); + batch.cancel_reason = None; + let (started_tx, started_rx) = tokio::sync::oneshot::channel(); + let abort_handle = pool.join_set.spawn(async move { + let _ = started_tx.send(()); + std::future::pending::<()>().await; + }); + let task_id = abort_handle.id(); + pool.task_map_mut().insert( + task_id, + crate::pool::TaskMeta { + agent_index: 0, + channel_id: Some(channel_id), + turn_id: "panic-after-cancel".into(), + recoverable_batch: Some(batch), + desired_model: None, + model_overridden: false, + accepted_model_switch: None, + accepted_drop_control: Some(AcceptedDropControl::Cancel), + control_tx: None, + steer_tx: None, + }, + ); + started_rx.await.unwrap(); + abort_handle.abort(); + let join_error = pool.join_set.join_next().await.unwrap().unwrap_err(); + + let mut queue = EventQueue::new(config::DedupMode::Queue); + let config = test_config(); + let mut heartbeat_in_flight = false; + let removed_channels = HashSet::new(); + let mut typing_channels = HashMap::new(); + let mut crash_history = vec![SlotCircuit::new()]; + let (respawn_tx, _respawn_rx) = mpsc::channel(1); + let mut respawn_tasks = tokio::task::JoinSet::new(); + + recover_panicked_agent( + &mut pool, + &mut queue, + &config, + join_error, + &mut heartbeat_in_flight, + &removed_channels, + &mut typing_channels, + &mut crash_history, + &respawn_tx, + &mut respawn_tasks, + None, + None, + ); + + assert_eq!( + queue.pending_channels(), + 0, + "panic recovery must not resurrect work after an accepted cancel" + ); + respawn_tasks.shutdown().await; + } + + #[tokio::test] + async fn pinned_model_replay_panic_preserves_intent_in_quarantined_slot() { + let mut pool = AgentPool::from_slots(vec![None]); + let channel_id = Uuid::new_v4(); + let (started_tx, started_rx) = tokio::sync::oneshot::channel(); + let abort_handle = pool.join_set.spawn(async move { + let _ = started_tx.send(()); + std::future::pending::<()>().await; + }); + let task_id = abort_handle.id(); + pool.task_map_mut().insert( + task_id, + crate::pool::TaskMeta { + agent_index: 0, + channel_id: Some(channel_id), + turn_id: "panic-model-turn".into(), + recoverable_batch: Some(one_event_batch(channel_id, "model-retry")), + desired_model: Some("runtime-model".into()), + model_overridden: true, + accepted_model_switch: None, + accepted_drop_control: None, + control_tx: None, + steer_tx: None, + }, + ); + started_rx.await.unwrap(); + abort_handle.abort(); + let join_error = pool.join_set.join_next().await.unwrap().unwrap_err(); + + let mut queue = EventQueue::new(config::DedupMode::Queue); + queue.require_agent(channel_id, 0); + let config = test_config(); + let mut heartbeat_in_flight = false; + let removed_channels = HashSet::new(); + let mut typing_channels = HashMap::new(); + let mut crash_history = vec![SlotCircuit::new()]; + let (respawn_tx, _respawn_rx) = mpsc::channel(1); + let mut respawn_tasks = tokio::task::JoinSet::new(); + + recover_panicked_agent( + &mut pool, + &mut queue, + &config, + join_error, + &mut heartbeat_in_flight, + &removed_channels, + &mut typing_channels, + &mut crash_history, + &respawn_tx, + &mut respawn_tasks, + None, + None, + ); + + let pending = crash_history[0] + .pending_model_intent + .as_ref() + .expect("pinned panic must retain model intent through replacement"); + assert_eq!(pending.desired_model.as_deref(), Some("runtime-model")); + assert!(pending.model_overridden); + assert_eq!(queue.required_agent(channel_id), Some(0)); + assert!( + crash_history[0].cleanup_unverified, + "panic loses the typed cleanup owner and must quarantine the slot" + ); + assert_eq!( + respawn_tasks.len(), + 0, + "panic recovery must not spawn a replacement without cleanup proof" + ); + respawn_tasks.shutdown().await; + } + + #[tokio::test] + async fn circuit_open_panic_quarantines_instead_of_scheduling_cooldown_refill() { + let mut pool = AgentPool::from_slots(vec![None]); + let channel_id = Uuid::new_v4(); + let (started_tx, started_rx) = tokio::sync::oneshot::channel(); + let abort_handle = pool.join_set.spawn(async move { + let _ = started_tx.send(()); + std::future::pending::<()>().await; + }); + let task_id = abort_handle.id(); + pool.task_map_mut().insert( + task_id, + crate::pool::TaskMeta { + agent_index: 0, + channel_id: Some(channel_id), + turn_id: "panic-open-circuit".into(), + recoverable_batch: Some(one_event_batch(channel_id, "model-retry")), + desired_model: Some("runtime-model".into()), + model_overridden: true, + accepted_model_switch: Some(ModelSwitchRequest::new( + "runtime-model", + "fedcba9876543210fedcba9876543210", + )), + accepted_drop_control: None, + control_tx: None, + steer_tx: None, + }, + ); + started_rx.await.unwrap(); + abort_handle.abort(); + let join_error = pool.join_set.join_next().await.unwrap().unwrap_err(); + + let mut queue = EventQueue::new(config::DedupMode::Queue); + queue.require_agent(channel_id, 0); + let config = test_config(); + let mut heartbeat_in_flight = false; + let removed_channels = HashSet::new(); + let mut typing_channels = HashMap::new(); + let mut crash_history = vec![SlotCircuit::new()]; + let now = std::time::Instant::now(); + crash_history[0].crash_times = vec![now; CIRCUIT_BREAKER_THRESHOLD.saturating_sub(1)]; + let (respawn_tx, mut respawn_rx) = mpsc::channel(1); + let mut respawn_tasks = tokio::task::JoinSet::new(); + + recover_panicked_agent( + &mut pool, + &mut queue, + &config, + join_error, + &mut heartbeat_in_flight, + &removed_channels, + &mut typing_channels, + &mut crash_history, + &respawn_tx, + &mut respawn_tasks, + None, + None, + ); + + assert!(crash_history[0].cleanup_unverified); + assert!( + crash_history[0].blocks_supervisor_exit(), + "quarantine must keep the current process alive and degraded" + ); + assert!( + respawn_rx.try_recv().is_err(), + "panic recovery must not publish a synthetic cleanup-complete marker" ); assert_eq!( - normalized_agent_name(&serde_json::json!({ - "serverInfo": { "name": "buzz-agent" } - })), - "buzz-agent" + crash_history[0] + .pending_model_intent + .as_ref() + .and_then(|intent| intent.desired_model.as_deref()), + Some("runtime-model") ); + assert!(respawn_tasks.is_empty()); } - /// Spawn a real but inert agent subprocess (`cat`) so the error paths have - /// an `OwnedAgent` to move into respawn or return to the pool. The error - /// branches never talk to the subprocess. - async fn dummy_agent(index: usize) -> OwnedAgent { - OwnedAgent { - index, - acp: AcpClient::spawn("cat", &[], &[], false) - .await - .expect("spawn cat as inert agent"), - state: Default::default(), - model_capabilities: None, - desired_model: None, - model_overridden: false, - agent_name: "unknown".into(), - goose_system_prompt_supported: None, - // Error branches under test never read this; 1 is the legacy - // non-systemPrompt path, the simplest valid value. - protocol_version: 1, - } - } - - /// Drive one error outcome through `handle_prompt_result` and return how - /// many `turn_error` events it emitted to the observer feed. - async fn turn_errors_emitted_for(outcome: PromptOutcome) -> usize { - let agent = dummy_agent(0).await; + #[tokio::test] + async fn exhausted_panic_retry_dead_letters_and_clears_exact_slot_pin() { let mut pool = AgentPool::from_slots(vec![None]); - - // `handle_prompt_result` asserts it removes exactly one in-flight task - // for the completing agent (the slot was checked out, not idle). Mirror - // the real dispatch path by registering a TaskMeta keyed on a genuine - // `task::Id` — only obtainable from inside a spawned task. - let task_id = pool.join_set.spawn(async {}).id(); + let channel_id = Uuid::new_v4(); + let (started_tx, started_rx) = tokio::sync::oneshot::channel(); + let abort_handle = pool.join_set.spawn(async move { + let _ = started_tx.send(()); + std::future::pending::<()>().await; + }); + let task_id = abort_handle.id(); + let mut batch = one_event_batch(channel_id, "poison-batch"); + batch.cancel_reason = None; pool.task_map_mut().insert( task_id, crate::pool::TaskMeta { agent_index: 0, - channel_id: None, - turn_id: "test-turn-id".to_string(), - recoverable_batch: None, + channel_id: Some(channel_id), + turn_id: "panic-dead-letter".into(), + recoverable_batch: Some(batch), + desired_model: Some("runtime-model".into()), + model_overridden: true, + accepted_model_switch: None, + accepted_drop_control: None, control_tx: None, steer_tx: None, }, ); + started_rx.await.unwrap(); + abort_handle.abort(); + let join_error = pool.join_set.join_next().await.unwrap().unwrap_err(); let mut queue = EventQueue::new(config::DedupMode::Queue); + queue.require_agent(channel_id, 0); + queue.set_retry_count_for_test(channel_id, crate::queue::MAX_RETRIES); let config = test_config(); let mut heartbeat_in_flight = false; let removed_channels = HashSet::new(); - let mut crash_history = vec![SlotCircuit { - crash_times: Vec::new(), - open_until: None, - respawn_in_flight: false, - }]; - let (respawn_tx, _respawn_rx) = mpsc::channel(8); + let mut typing_channels = HashMap::new(); + let mut crash_history = vec![SlotCircuit::new()]; + let (respawn_tx, _respawn_rx) = mpsc::channel(1); let mut respawn_tasks = tokio::task::JoinSet::new(); - let observer = ObserverHandle::in_process(); - - let result = PromptResult { - agent, - source: PromptSource::Channel(Uuid::new_v4()), - turn_id: "test-turn-id".to_string(), - outcome, - batch: None, - }; - handle_prompt_result( + recover_panicked_agent( &mut pool, &mut queue, &config, - result, + join_error, &mut heartbeat_in_flight, &removed_channels, + &mut typing_channels, &mut crash_history, &respawn_tx, &mut respawn_tasks, - Some(observer.clone()), + None, None, ); - let turn_errors: Vec<_> = observer - .snapshot() - .into_iter() - .filter(|e| e.kind == "turn_error") - .collect(); - assert!( - turn_errors - .iter() - .all(|event| event.turn_id.as_deref() == Some("test-turn-id")), - "turn_error must retain the completed turn id" + assert_eq!( + queue.required_agent(channel_id), + None, + "dead-lettered work must not leave an orphan exact-slot requirement" ); - turn_errors.len() - } - - #[tokio::test] - async fn agent_exited_emits_exactly_one_feed_event() { - assert_eq!(turn_errors_emitted_for(PromptOutcome::AgentExited).await, 1); + assert!(!queue.has_flushable_work()); + respawn_tasks.shutdown().await; } #[tokio::test] - async fn panic_event_retains_task_turn_id() { - let mut pool = AgentPool::from_slots(vec![]); + async fn exhausted_panic_retry_preserves_exact_slot_pin_for_queued_residue() { let channel_id = Uuid::new_v4(); + let mut queue = EventQueue::new(config::DedupMode::Queue); + queue.require_agent(channel_id, 0); + for index in 0..=crate::queue::MAX_BATCH_EVENTS { + queue.push(crate::queue::QueuedEvent { + channel_id, + event: EventBuilder::new(Kind::Custom(9), format!("dependent-{index}")) + .sign_with_keys(&Keys::generate()) + .unwrap(), + received_at: std::time::Instant::now(), + prompt_tag: "test".into(), + }); + } + let poison_batch = queue.flush_next().expect("first bounded batch"); + assert_eq!(poison_batch.events.len(), crate::queue::MAX_BATCH_EVENTS); + assert_eq!(queue.queued_event_count(&channel_id), 1); + queue.set_retry_count_for_test(channel_id, crate::queue::MAX_RETRIES); + + let mut pool = AgentPool::from_slots(vec![None]); let (started_tx, started_rx) = tokio::sync::oneshot::channel(); let abort_handle = pool.join_set.spawn(async move { let _ = started_tx.send(()); @@ -5372,8 +10199,12 @@ mod error_outcome_emission_tests { crate::pool::TaskMeta { agent_index: 0, channel_id: Some(channel_id), - turn_id: "panic-turn-id".to_string(), - recoverable_batch: None, + turn_id: "panic-dead-letter-with-residue".into(), + recoverable_batch: Some(poison_batch), + desired_model: Some("runtime-model".into()), + model_overridden: true, + accepted_model_switch: None, + accepted_drop_control: None, control_tx: None, steer_tx: None, }, @@ -5382,19 +10213,13 @@ mod error_outcome_emission_tests { abort_handle.abort(); let join_error = pool.join_set.join_next().await.unwrap().unwrap_err(); - let mut queue = EventQueue::new(config::DedupMode::Queue); let config = test_config(); let mut heartbeat_in_flight = false; let removed_channels = HashSet::new(); let mut typing_channels = HashMap::new(); - let mut crash_history = vec![SlotCircuit { - crash_times: Vec::new(), - open_until: None, - respawn_in_flight: false, - }]; - let (respawn_tx, _respawn_rx) = mpsc::channel(8); + let mut crash_history = vec![SlotCircuit::new()]; + let (respawn_tx, _respawn_rx) = mpsc::channel(1); let mut respawn_tasks = tokio::task::JoinSet::new(); - let observer = ObserverHandle::in_process(); recover_panicked_agent( &mut pool, @@ -5407,19 +10232,21 @@ mod error_outcome_emission_tests { &mut crash_history, &respawn_tx, &mut respawn_tasks, - Some(observer.clone()), + None, + None, ); - let panic = observer - .snapshot() - .into_iter() - .find(|event| event.kind == "agent_panic") - .expect("panic recovery emits an observer event"); assert_eq!( - panic.channel_id.as_deref(), - Some(channel_id.to_string().as_str()) + queue.required_agent(channel_id), + Some(0), + "terminal disposal must preserve exact-slot affinity for dependent residue" ); - assert_eq!(panic.turn_id.as_deref(), Some("panic-turn-id")); + assert_eq!( + queue.queued_event_count(&channel_id), + 1, + "the batch boundary must leave dependent residue queued" + ); + respawn_tasks.shutdown().await; } #[tokio::test] @@ -5452,6 +10279,191 @@ mod error_outcome_emission_tests { ); } + #[tokio::test] + async fn optional_close_recycle_respawns_without_crash_penalty() { + let disposition = outcome_disposition_for(PromptOutcome::SessionRecycleRequired).await; + + assert_eq!( + disposition.live_agents, 0, + "the adapter must leave service until its process group is replaced" + ); + assert_eq!( + disposition.respawn_tasks, 1, + "optional-close fallback must start one intentional replacement" + ); + assert_eq!( + disposition.turn_errors, 0, + "a negotiated compatibility fallback is not an operator error" + ); + assert_eq!( + disposition.recent_crashes, 0, + "intentional replacement must not consume crash budget" + ); + assert!( + !disposition.circuit_open, + "normal retirements must never open the crash circuit" + ); + } + + #[tokio::test] + async fn optional_close_recycle_preserves_live_model_override() { + let mut agent = dummy_agent(0).await; + agent.desired_model = Some("runtime-model".into()); + agent.model_overridden = true; + let config = test_config(); + let mut slot = SlotCircuit::new(); + let (respawn_tx, mut respawn_rx) = mpsc::channel(1); + let mut respawn_tasks = tokio::task::JoinSet::new(); + + assert!(spawn_recycle_task( + agent, + &config, + &mut slot, + &respawn_tx, + &mut respawn_tasks, + None, + )); + let result = tokio::time::timeout(Duration::from_secs(5), respawn_rx.recv()) + .await + .expect("intentional recycle must report its spawn result") + .expect("respawn result channel must stay open"); + + assert_eq!(result.desired_model.as_deref(), Some("runtime-model")); + assert!( + result.model_overridden, + "busy SwitchModel intent must survive the compatibility recycle" + ); + assert!( + slot.crash_times.is_empty() && slot.open_until.is_none(), + "intentional recycle must not mutate the crash circuit" + ); + } + + #[tokio::test] + async fn optional_close_recycle_reserves_per_slot_rate_budget() { + let agent = dummy_agent(0).await; + let config = test_config(); + let mut slot = SlotCircuit::new(); + let (respawn_tx, _respawn_rx) = mpsc::channel(1); + let mut respawn_tasks = tokio::task::JoinSet::new(); + let started = std::time::Instant::now(); + + assert!(spawn_recycle_task( + agent, + &config, + &mut slot, + &respawn_tx, + &mut respawn_tasks, + None, + )); + let next = slot + .next_recycle_not_before + .expect("recycle must reserve the next per-slot budget"); + assert!( + next >= started + RECYCLE_MIN_INTERVAL, + "the next negotiated recycle must be delayed by the minimum interval" + ); + assert!(slot.crash_times.is_empty() && slot.open_until.is_none()); + respawn_tasks.shutdown().await; + } + + #[tokio::test] + async fn circuit_open_still_schedules_guaranteed_old_adapter_cleanup() { + let agent = dummy_agent(0).await; + let config = test_config(); + let mut slot = SlotCircuit::new(); + slot.open_until = Some(std::time::Instant::now() + CIRCUIT_BREAKER_COOLDOWN); + let (respawn_tx, mut respawn_rx) = mpsc::channel(1); + let mut respawn_tasks = tokio::task::JoinSet::new(); + + assert!( + spawn_respawn_task( + agent, + &config, + &mut slot, + &respawn_tx, + &mut respawn_tasks, + None, + ), + "an open circuit may suppress replacement, never bounded shutdown" + ); + assert!(slot.respawn_in_flight); + let result = tokio::time::timeout(Duration::from_secs(2), respawn_rx.recv()) + .await + .expect("cleanup completion must wake the supervisor") + .expect("cleanup result channel must stay open"); + assert!( + result.result.is_ok(), + "cleanup-only completion is not a replacement spawn failure" + ); + respawn_tasks + .join_next() + .await + .expect("cleanup task must be tracked") + .expect("cleanup task must not panic"); + } + + #[tokio::test] + async fn every_session_close_failure_poison_class_respawns_the_adapter() { + let cases = [ + ( + "unsupported", + AcpError::AgentError { + code: -32601, + message: "method not found".into(), + }, + ), + ( + "adapter error", + AcpError::AgentError { + code: -32000, + message: "close rejected".into(), + }, + ), + ( + "timeout", + AcpError::Timeout(std::time::Duration::from_secs(5)), + ), + ]; + + for (name, error) in cases { + let disposition = + outcome_disposition_for(PromptOutcome::SessionCloseFailed(error)).await; + assert_eq!( + disposition.live_agents, 0, + "{name}: poisoned adapter must not return to the idle pool" + ); + assert_eq!( + disposition.respawn_tasks, 1, + "{name}: poisoned adapter must enter process-group replacement" + ); + assert_eq!( + disposition.turn_errors, 1, + "{name}: operator feed must record the cleanup failure" + ); + } + } + + #[tokio::test] + async fn session_close_failure_redacts_adapter_controlled_message() { + let secret = "adapter-secret-that-must-not-enter-observer"; + let disposition = + outcome_disposition_for(PromptOutcome::SessionCloseFailed(AcpError::AgentError { + code: -32000, + message: secret.into(), + })) + .await; + + assert_eq!(disposition.turn_error_texts.len(), 1); + assert!( + disposition + .turn_error_texts + .iter() + .all(|text| !text.contains(secret)), + "adapter-controlled error text must not enter the observer feed" + ); + } + /// idle_timeout outcome_label is "idle_timeout"; hard_timeout is "hard_timeout". #[tokio::test] async fn timeout_outcome_labels_differ() { @@ -5466,6 +10478,10 @@ mod error_outcome_emission_tests { channel_id: None, turn_id: "test-turn-id".to_string(), recoverable_batch: None, + desired_model: None, + model_overridden: false, + accepted_model_switch: None, + accepted_drop_control: None, control_tx: None, steer_tx: None, }, @@ -5474,11 +10490,7 @@ mod error_outcome_emission_tests { let config = test_config(); let mut heartbeat_in_flight = false; let removed_channels = HashSet::new(); - let mut crash_history = vec![SlotCircuit { - crash_times: Vec::new(), - open_until: None, - respawn_in_flight: false, - }]; + let mut crash_history = vec![SlotCircuit::new()]; let (respawn_tx, _respawn_rx) = mpsc::channel(8); let mut respawn_tasks = tokio::task::JoinSet::new(); let observer = ObserverHandle::in_process(); @@ -5488,6 +10500,7 @@ mod error_outcome_emission_tests { turn_id: "test-turn-id".to_string(), outcome, batch: None, + retry_agent_index: None, }; handle_prompt_result( &mut pool, @@ -5541,6 +10554,7 @@ mod error_outcome_emission_tests { }], cancelled_events: vec![], cancel_reason: None, + occurrence_ids: BatchOccurrenceIds::for_test(1, 0), } }; @@ -5557,6 +10571,10 @@ mod error_outcome_emission_tests { channel_id: None, turn_id: "test-turn-id".to_string(), recoverable_batch: None, + desired_model: None, + model_overridden: false, + accepted_model_switch: None, + accepted_drop_control: None, control_tx: None, steer_tx: None, }, @@ -5565,11 +10583,7 @@ mod error_outcome_emission_tests { let config = test_config(); let mut heartbeat_in_flight = false; let removed_channels = HashSet::new(); - let mut crash_history = vec![SlotCircuit { - crash_times: Vec::new(), - open_until: None, - respawn_in_flight: false, - }]; + let mut crash_history = vec![SlotCircuit::new()]; let (respawn_tx, _respawn_rx) = mpsc::channel(8); let mut respawn_tasks = tokio::task::JoinSet::new(); let result = PromptResult { @@ -5578,6 +10592,7 @@ mod error_outcome_emission_tests { turn_id: "test-turn-id".to_string(), outcome, batch: Some(batch), + retry_agent_index: None, }; handle_prompt_result( &mut pool, @@ -5647,6 +10662,7 @@ mod error_outcome_emission_tests { }], cancelled_events: vec![], cancel_reason: None, + occurrence_ids: BatchOccurrenceIds::for_test(1, 0), } }; @@ -5662,6 +10678,10 @@ mod error_outcome_emission_tests { channel_id: None, turn_id: "test-turn-id".to_string(), recoverable_batch: None, + desired_model: None, + model_overridden: false, + accepted_model_switch: None, + accepted_drop_control: None, control_tx: None, steer_tx: None, }, @@ -5670,11 +10690,7 @@ mod error_outcome_emission_tests { let config = test_config(); let mut heartbeat_in_flight = false; let removed_channels = HashSet::new(); - let mut crash_history = vec![SlotCircuit { - crash_times: Vec::new(), - open_until: None, - respawn_in_flight: false, - }]; + let mut crash_history = vec![SlotCircuit::new()]; let (respawn_tx, _respawn_rx) = mpsc::channel(8); let mut respawn_tasks = tokio::task::JoinSet::new(); let result = PromptResult { @@ -5683,6 +10699,7 @@ mod error_outcome_emission_tests { turn_id: "test-turn-id".to_string(), outcome, batch: Some(batch), + retry_agent_index: None, }; handle_prompt_result( &mut pool, @@ -5738,6 +10755,10 @@ mod error_outcome_emission_tests { channel_id: None, turn_id: "test-turn-id".to_string(), recoverable_batch: None, + desired_model: None, + model_overridden: false, + accepted_model_switch: None, + accepted_drop_control: None, control_tx: None, steer_tx: None, }, @@ -5746,11 +10767,7 @@ mod error_outcome_emission_tests { let config = test_config(); let mut heartbeat_in_flight = false; let removed_channels = HashSet::new(); - let mut crash_history = vec![SlotCircuit { - crash_times: Vec::new(), - open_until: None, - respawn_in_flight: false, - }]; + let mut crash_history = vec![SlotCircuit::new()]; let (respawn_tx, _respawn_rx) = mpsc::channel(8); let mut respawn_tasks = tokio::task::JoinSet::new(); let observer = ObserverHandle::in_process(); @@ -5765,6 +10782,7 @@ mod error_outcome_emission_tests { }], cancelled_events: vec![], cancel_reason: None, + occurrence_ids: BatchOccurrenceIds::for_test(1, 0), }; let result = PromptResult { agent, @@ -5774,6 +10792,7 @@ mod error_outcome_emission_tests { recently_active: true, }), batch: Some(batch), + retry_agent_index: None, }; handle_prompt_result( &mut pool, @@ -5821,6 +10840,7 @@ mod error_outcome_emission_tests { // upcoming requeue() call in handle_prompt_result crosses the // dead-letter threshold. queue.set_retry_count_for_test(channel_id, crate::queue::MAX_RETRIES); + queue.require_agent(channel_id, 0); let agent = dummy_agent(0).await; let mut pool = AgentPool::from_slots(vec![None]); @@ -5832,6 +10852,10 @@ mod error_outcome_emission_tests { channel_id: None, turn_id: "test-turn-id".to_string(), recoverable_batch: None, + desired_model: None, + model_overridden: false, + accepted_model_switch: None, + accepted_drop_control: None, control_tx: None, steer_tx: None, }, @@ -5839,11 +10863,7 @@ mod error_outcome_emission_tests { let config = test_config(); let mut heartbeat_in_flight = false; let removed_channels = HashSet::new(); - let mut crash_history = vec![SlotCircuit { - crash_times: Vec::new(), - open_until: None, - respawn_in_flight: false, - }]; + let mut crash_history = vec![SlotCircuit::new()]; let (respawn_tx, _respawn_rx) = mpsc::channel(8); let mut respawn_tasks = tokio::task::JoinSet::new(); let observer = ObserverHandle::in_process(); @@ -5858,6 +10878,7 @@ mod error_outcome_emission_tests { }], cancelled_events: vec![], cancel_reason: None, + occurrence_ids: BatchOccurrenceIds::for_test(1, 0), }; let result = PromptResult { agent, @@ -5867,6 +10888,7 @@ mod error_outcome_emission_tests { recently_active: true, }), batch: Some(batch), + retry_agent_index: None, }; handle_prompt_result( &mut pool, @@ -5899,6 +10921,11 @@ mod error_outcome_emission_tests { 0, "batch with an exhausted retry budget must be dead-lettered, not requeued" ); + assert_eq!( + queue.required_agent(channel_id), + None, + "terminal dead-letter must release its exact-slot requirement" + ); } /// Cancel-drain-timeout batches are requeued as cancelled (merge into the @@ -5936,6 +10963,7 @@ mod error_outcome_emission_tests { }], cancelled_events: vec![], cancel_reason: Some(CancelReason::Steer), + occurrence_ids: BatchOccurrenceIds::for_test(1, 0), }; let agent = dummy_agent(0).await; @@ -5948,6 +10976,10 @@ mod error_outcome_emission_tests { channel_id: None, turn_id: "test-turn-id".to_string(), recoverable_batch: None, + desired_model: None, + model_overridden: false, + accepted_model_switch: None, + accepted_drop_control: None, control_tx: None, steer_tx: None, }, @@ -5966,11 +10998,7 @@ mod error_outcome_emission_tests { let config = test_config(); let mut heartbeat_in_flight = false; let removed_channels = HashSet::new(); - let mut crash_history = vec![SlotCircuit { - crash_times: Vec::new(), - open_until: None, - respawn_in_flight: false, - }]; + let mut crash_history = vec![SlotCircuit::new()]; let (respawn_tx, _respawn_rx) = mpsc::channel(8); let mut respawn_tasks = tokio::task::JoinSet::new(); let observer = ObserverHandle::in_process(); @@ -5981,6 +11009,7 @@ mod error_outcome_emission_tests { turn_id: "test-turn-id".to_string(), outcome: PromptOutcome::CancelDrainTimeout(grace), batch: Some(batch), + retry_agent_index: None, }; handle_prompt_result( @@ -6087,6 +11116,10 @@ mod error_outcome_emission_tests { channel_id: None, turn_id: "test-turn-id".to_string(), recoverable_batch: None, + desired_model: None, + model_overridden: false, + accepted_model_switch: None, + accepted_drop_control: None, control_tx: None, steer_tx: None, }, @@ -6095,11 +11128,7 @@ mod error_outcome_emission_tests { let config = test_config(); let mut heartbeat_in_flight = false; let removed_channels = HashSet::new(); - let mut crash_history = vec![SlotCircuit { - crash_times: Vec::new(), - open_until: None, - respawn_in_flight: false, - }]; + let mut crash_history = vec![SlotCircuit::new()]; let (respawn_tx, _respawn_rx) = mpsc::channel(8); let mut respawn_tasks = tokio::task::JoinSet::new(); let observer = ObserverHandle::in_process(); @@ -6113,6 +11142,7 @@ mod error_outcome_emission_tests { // `classify_control_cancel_failure` — `handle_prompt_result` // never sees one to requeue. batch: None, + retry_agent_index: None, }; handle_prompt_result( @@ -6177,6 +11207,19 @@ mod error_outcome_emission_tests { assert_eq!(turn_errors_emitted_for(PromptOutcome::Error(io)).await, 1); } + #[tokio::test] + async fn malformed_json_response_poisons_adapter_instead_of_reusing_it() { + let json_error = serde_json::from_str::("{") + .expect_err("fixture must be malformed JSON"); + let disposition = + outcome_disposition_for(PromptOutcome::Error(AcpError::Json(json_error))).await; + assert_eq!(disposition.live_agents, 0); + assert_eq!( + disposition.respawn_tasks, 1, + "JSON framing corruption must replace the adapter" + ); + } + #[tokio::test] async fn application_error_emits_exactly_one_feed_event() { let app = AcpError::IdleTimeout(std::time::Duration::from_secs(1)); @@ -6257,6 +11300,7 @@ mod error_outcome_emission_tests { }], cancelled_events: vec![], cancel_reason: None, + occurrence_ids: BatchOccurrenceIds::for_test(1, 0), }; let auth_error = acp::AcpError::AgentError { @@ -6275,6 +11319,10 @@ mod error_outcome_emission_tests { channel_id: None, turn_id: "test-turn-id".to_string(), recoverable_batch: None, + desired_model: None, + model_overridden: false, + accepted_model_switch: None, + accepted_drop_control: None, control_tx: None, steer_tx: None, }, @@ -6283,11 +11331,7 @@ mod error_outcome_emission_tests { let config = test_config(); let mut heartbeat_in_flight = false; let removed_channels = std::collections::HashSet::new(); - let mut crash_history = vec![SlotCircuit { - crash_times: Vec::new(), - open_until: None, - respawn_in_flight: false, - }]; + let mut crash_history = vec![SlotCircuit::new()]; let (respawn_tx, _respawn_rx) = mpsc::channel(8); let mut respawn_tasks = tokio::task::JoinSet::new(); let result = PromptResult { @@ -6296,6 +11340,7 @@ mod error_outcome_emission_tests { turn_id: "test-turn-id".to_string(), outcome: PromptOutcome::Error(auth_error), batch: Some(batch), + retry_agent_index: None, }; handle_prompt_result( &mut pool, @@ -6342,6 +11387,7 @@ mod error_outcome_emission_tests { }], cancelled_events: vec![], cancel_reason: None, + occurrence_ids: BatchOccurrenceIds::for_test(1, 0), }; // Usage-credits error — AgentError but NOT an auth error. @@ -6360,6 +11406,10 @@ mod error_outcome_emission_tests { channel_id: None, turn_id: "test-turn-id".to_string(), recoverable_batch: None, + desired_model: None, + model_overridden: false, + accepted_model_switch: None, + accepted_drop_control: None, control_tx: None, steer_tx: None, }, @@ -6368,11 +11418,7 @@ mod error_outcome_emission_tests { let config = test_config(); let mut heartbeat_in_flight = false; let removed_channels = std::collections::HashSet::new(); - let mut crash_history = vec![SlotCircuit { - crash_times: Vec::new(), - open_until: None, - respawn_in_flight: false, - }]; + let mut crash_history = vec![SlotCircuit::new()]; let (respawn_tx, _respawn_rx) = mpsc::channel(8); let mut respawn_tasks = tokio::task::JoinSet::new(); let result = PromptResult { @@ -6381,6 +11427,7 @@ mod error_outcome_emission_tests { turn_id: "test-turn-id".to_string(), outcome: PromptOutcome::Error(usage_error), batch: Some(batch), + retry_agent_index: None, }; handle_prompt_result( &mut pool, diff --git a/crates/buzz-acp/src/observer.rs b/crates/buzz-acp/src/observer.rs index 7029e5af6d..b2dcc2726b 100644 --- a/crates/buzz-acp/src/observer.rs +++ b/crates/buzz-acp/src/observer.rs @@ -5,7 +5,7 @@ //! frames without exposing a local HTTP port. use std::{ - collections::VecDeque, + collections::{BTreeMap, HashMap, HashSet, VecDeque}, sync::{ atomic::{AtomicU64, Ordering}, Arc, Mutex, @@ -16,6 +16,29 @@ use serde::Serialize; use tokio::sync::broadcast; const OBSERVER_BUFFER_CAP: usize = 1_000; +/// Reserved broadcast capacity for terminal control results. This is only the +/// live wake-up lane: terminal results themselves remain in the acknowledged +/// replay ledger until the relay publisher confirms delivery. +const OBSERVER_CONTROL_RESULT_CAP: usize = 64; +/// Hard process-local retention bound for terminal results awaiting relay ACK. +/// +/// The live publisher normally owns at most one result at a time. This larger +/// reserve absorbs an extended relay outage without allowing an authenticated +/// control stream to grow process memory indefinitely. Overflow rejects the +/// newest result, logs loudly, and permanently marks shutdown proof unverified +/// for this observer lifecycle; it never evicts an older accepted result. +const OBSERVER_CONTROL_RESULT_LEDGER_CAP: usize = 1_024; + +pub(crate) fn is_terminal_control_result(event: &ObserverEvent) -> bool { + event.kind == "control_result" + && !matches!( + event + .payload + .get("status") + .and_then(serde_json::Value::as_str), + Some("sent" | "recycling") + ) +} /// Best-effort metadata attached to observer events. #[derive(Clone, Debug, Default)] @@ -38,21 +61,278 @@ pub struct ObserverHandle { struct ObserverInner { tx: broadcast::Sender, + control_result_tx: broadcast::Sender, buffer: Mutex>, + control_results: Arc>>, + control_result_rejected: Arc, seq: AtomicU64, } fn new_observer_handle() -> ObserverHandle { let (tx, _) = broadcast::channel(OBSERVER_BUFFER_CAP); + let (control_result_tx, _) = broadcast::channel(OBSERVER_CONTROL_RESULT_CAP); ObserverHandle { inner: Arc::new(ObserverInner { tx, + control_result_tx, buffer: Mutex::new(VecDeque::with_capacity(OBSERVER_BUFFER_CAP)), + control_results: Arc::new(Mutex::new(BTreeMap::new())), + control_result_rejected: Arc::new(AtomicU64::new(0)), seq: AtomicU64::new(1), }), } } +/// Priority-aware receiver for the observer publisher. +/// +/// Terminal control results have reserved broadcast capacity and win when both +/// lanes are ready. This receiver is the single acknowledging consumer for +/// terminal results: the ledger retains each result until this publisher +/// confirms relay handoff, and any remainder is released when the process-local +/// observer lifecycle is dropped. +pub struct ObserverReceiver { + control_result_rx: broadcast::Receiver, + telemetry_rx: broadcast::Receiver, + control_results: Arc>>, + control_result_rejected: Arc, + control_result_replay: VecDeque, + in_flight_control_results: HashSet, + control_result_failures: HashMap, + control_result_closed: bool, + telemetry_closed: bool, +} + +impl ObserverReceiver { + /// Whether this observer lifecycle rejected any terminal result because its + /// bounded acknowledgement ledger was full. + pub(crate) fn control_result_admission_failed(&self) -> bool { + self.control_result_rejected.load(Ordering::Acquire) > 0 + } + + fn queue_retained_control_results(&mut self) { + let retained = match self.control_results.lock() { + Ok(results) => results + .values() + .filter(|event| !self.in_flight_control_results.contains(&event.seq)) + .cloned() + .collect::>(), + Err(error) => { + tracing::warn!( + target: "observer", + "observer control-result ledger lock poisoned; recovering retained results" + ); + error + .into_inner() + .values() + .filter(|event| !self.in_flight_control_results.contains(&event.seq)) + .cloned() + .collect() + } + }; + for event in retained { + self.control_result_replay.push_back(event); + } + } + + fn pop_retained_control_result(&mut self) -> Option { + while let Some(event) = self.control_result_replay.pop_front() { + let retained = match self.control_results.lock() { + Ok(results) => results.contains_key(&event.seq), + Err(error) => { + tracing::warn!( + target: "observer", + "observer control-result ledger lock poisoned; recovering delivery state" + ); + error.into_inner().contains_key(&event.seq) + } + }; + if retained && self.in_flight_control_results.insert(event.seq) { + return Some(event); + } + } + None + } + + fn accept_live_control_result(&mut self, event: ObserverEvent) -> Option { + let retained = match self.control_results.lock() { + Ok(results) => results.contains_key(&event.seq), + Err(error) => { + tracing::warn!( + target: "observer", + "observer control-result ledger lock poisoned; recovering live delivery state" + ); + error.into_inner().contains_key(&event.seq) + } + }; + (retained && self.in_flight_control_results.insert(event.seq)).then_some(event) + } + + /// Confirm that a terminal result reached the relay publisher. + /// + /// This is the sole normal-lifecycle removal point for retained terminal + /// outcomes. A lagging receiver resnapshots everything not acknowledged; + /// dropping the in-process observer lifecycle drops any remaining ledger. + pub(crate) fn acknowledge_control_result(&mut self, seq: u64) { + self.in_flight_control_results.remove(&seq); + self.control_result_failures.remove(&seq); + self.control_result_replay.retain(|event| event.seq != seq); + match self.control_results.lock() { + Ok(mut results) => { + results.remove(&seq); + } + Err(error) => { + tracing::warn!( + target: "observer", + "observer control-result ledger lock poisoned; recovering acknowledgement" + ); + error.into_inner().remove(&seq); + } + } + } + + /// Release one failed terminal delivery for a single end-to-end retry. + /// + /// Relay publication already performs reconnect and rate-gate recovery + /// within its own bounded acknowledgement window. This receiver grants one + /// additional replay for failures that escape that window, then leaves a + /// second failure retained and in-flight so shutdown reports unverified + /// delivery without spinning forever. + pub(crate) fn retry_control_result(&mut self, seq: u64) -> bool { + if self + .control_result_failures + .get(&seq) + .is_some_and(|failures| *failures >= 1) + { + return false; + } + let retained = match self.control_results.lock() { + Ok(results) => results.get(&seq).cloned(), + Err(error) => { + tracing::warn!( + target: "observer", + "observer control-result ledger lock poisoned; recovering retry state" + ); + error.into_inner().get(&seq).cloned() + } + }; + let Some(event) = retained else { + return false; + }; + + self.control_result_failures.insert(seq, 1); + self.in_flight_control_results.remove(&seq); + if !self + .control_result_replay + .iter() + .any(|queued| queued.seq == seq) + { + self.control_result_replay.push_back(event); + } + true + } + + /// Receive only from the reserved terminal-control lane. + /// + /// The relay publisher uses this while ordinary telemetry is being paced + /// so a newly emitted terminal result can interrupt that wait. + pub(crate) async fn recv_control_result( + &mut self, + ) -> Result { + loop { + if let Some(event) = self.pop_retained_control_result() { + return Ok(event); + } + if self.control_result_closed { + self.queue_retained_control_results(); + return self + .pop_retained_control_result() + .ok_or(broadcast::error::RecvError::Closed); + } + match self.control_result_rx.recv().await { + Ok(event) => { + if let Some(event) = self.accept_live_control_result(event) { + return Ok(event); + } + } + Err(broadcast::error::RecvError::Lagged(count)) => { + tracing::warn!( + skipped_wakeups = count, + "observer terminal-result receiver lagged; replaying retained sequences" + ); + self.queue_retained_control_results(); + } + Err(broadcast::error::RecvError::Closed) => { + self.control_result_closed = true; + self.queue_retained_control_results(); + } + } + } + } + + /// Receive the next observer event, draining terminal control results first. + pub async fn recv(&mut self) -> Result { + loop { + if let Some(event) = self.pop_retained_control_result() { + return Ok(event); + } + match (self.control_result_closed, self.telemetry_closed) { + (true, true) => { + self.queue_retained_control_results(); + return self + .pop_retained_control_result() + .ok_or(broadcast::error::RecvError::Closed); + } + (false, true) => return self.recv_control_result().await, + (true, false) => match self.telemetry_rx.recv().await { + Err(broadcast::error::RecvError::Closed) => { + self.telemetry_closed = true; + } + result => return result, + }, + (false, false) => { + tokio::select! { + biased; + result = self.control_result_rx.recv() => { + match result { + Ok(event) => { + if let Some(event) = self.accept_live_control_result(event) { + return Ok(event); + } + } + Err(broadcast::error::RecvError::Lagged(count)) => { + tracing::warn!( + skipped_wakeups = count, + "observer terminal-result receiver lagged; replaying retained sequences" + ); + self.queue_retained_control_results(); + if let Some(event) = self.pop_retained_control_result() { + return Ok(event); + } + } + Err(broadcast::error::RecvError::Closed) => { + self.control_result_closed = true; + self.queue_retained_control_results(); + if let Some(event) = self.pop_retained_control_result() { + return Ok(event); + } + } + } + } + result = self.telemetry_rx.recv() => { + match result { + Err(broadcast::error::RecvError::Closed) => { + self.telemetry_closed = true; + } + result => return result, + } + } + } + } + } + } + } +} + /// Event delivered through the in-process observer bus. #[derive(Clone, Serialize)] #[serde(rename_all = "camelCase")] @@ -85,19 +365,42 @@ impl ObserverHandle { } /// Subscribe to live observer events. - pub fn subscribe(&self) -> broadcast::Receiver { - self.inner.tx.subscribe() + pub fn subscribe(&self) -> ObserverReceiver { + ObserverReceiver { + control_result_rx: self.inner.control_result_tx.subscribe(), + telemetry_rx: self.inner.tx.subscribe(), + control_results: Arc::clone(&self.inner.control_results), + control_result_rejected: Arc::clone(&self.inner.control_result_rejected), + control_result_replay: VecDeque::new(), + in_flight_control_results: HashSet::new(), + control_result_failures: HashMap::new(), + control_result_closed: false, + telemetry_closed: false, + } } - /// Return the current replay buffer. + /// Return the current replay buffers with terminal control results first. + /// + /// Separate storage prevents a telemetry flood from evicting a result + /// before a newly started publisher snapshots the observer. pub fn snapshot(&self) -> Vec { + let mut snapshot = match self.inner.control_results.lock() { + Ok(results) => results.values().cloned().collect::>(), + Err(error) => { + tracing::warn!( + target: "observer", + "observer control-result ledger lock poisoned; recovering snapshot" + ); + error.into_inner().values().cloned().collect() + } + }; match self.inner.buffer.lock() { - Ok(buffer) => buffer.iter().cloned().collect(), + Ok(buffer) => snapshot.extend(buffer.iter().cloned()), Err(error) => { tracing::warn!(target: "observer", "observer replay buffer lock poisoned: {error}"); - Vec::new() } } + snapshot } /// Emit a local observer event. @@ -120,19 +423,63 @@ impl ObserverHandle { payload, }; - match self.inner.buffer.lock() { - Ok(mut buffer) => { - if buffer.len() >= OBSERVER_BUFFER_CAP { - buffer.pop_front(); + if event.kind == "control_result" && is_terminal_control_result(&event) { + let admitted = match self.inner.control_results.lock() { + Ok(mut results) => { + if results.len() >= OBSERVER_CONTROL_RESULT_LEDGER_CAP { + false + } else { + results.insert(event.seq, event.clone()); + true + } + } + Err(error) => { + tracing::warn!( + target: "observer", + "observer control-result ledger lock poisoned; recovering emission" + ); + let mut results = error.into_inner(); + if results.len() >= OBSERVER_CONTROL_RESULT_LEDGER_CAP { + false + } else { + results.insert(event.seq, event.clone()); + true + } } - buffer.push_back(event.clone()); + }; + if !admitted { + let rejected_total = self + .inner + .control_result_rejected + .fetch_add(1, Ordering::AcqRel) + .saturating_add(1); + tracing::error!( + target: "observer", + seq = event.seq, + rejected_total, + capacity = OBSERVER_CONTROL_RESULT_LEDGER_CAP, + "observer terminal-result ledger full; rejected newest result and invalidated delivery proof" + ); + return; } - Err(error) => { - tracing::warn!(target: "observer", "observer replay buffer lock poisoned: {error}"); + let _ = self.inner.control_result_tx.send(event); + } else { + match self.inner.buffer.lock() { + Ok(mut buffer) => { + if buffer.len() >= OBSERVER_BUFFER_CAP { + buffer.pop_front(); + } + buffer.push_back(event.clone()); + } + Err(error) => { + tracing::warn!( + target: "observer", + "observer replay buffer lock poisoned: {error}" + ); + } } + let _ = self.inner.tx.send(event); } - - let _ = self.inner.tx.send(event); } } @@ -164,3 +511,213 @@ pub fn context_for_turn( started_at: Some(started_at), } } + +#[cfg(test)] +mod tests { + use super::*; + + fn emit_kind(observer: &ObserverHandle, kind: &str, marker: usize) { + observer.emit( + kind, + None, + &ObserverContext::default(), + serde_json::json!({"marker": marker}), + ); + } + + fn emit_control_status(observer: &ObserverHandle, status: &str, marker: usize) { + observer.emit( + "control_result", + None, + &ObserverContext::default(), + serde_json::json!({"status": status, "marker": marker}), + ); + } + + #[tokio::test] + async fn non_terminal_control_statuses_remain_live_without_using_terminal_retention() { + let observer = ObserverHandle::in_process(); + let mut rx = observer.subscribe(); + + emit_control_status(&observer, "sent", 1); + let sent = rx + .recv() + .await + .expect("sent status remains observable live"); + assert_eq!(sent.payload["status"], "sent"); + + emit_control_status(&observer, "recycling", 2); + let recycling = rx + .recv() + .await + .expect("recycling status remains observable live"); + assert_eq!(recycling.payload["status"], "recycling"); + + for marker in 0..=OBSERVER_CONTROL_RESULT_LEDGER_CAP { + emit_control_status(&observer, "sent", marker); + } + emit_control_status(&observer, "switched", 3); + + assert!( + !rx.control_result_admission_failed(), + "non-terminal progress must not exhaust terminal-result admission" + ); + assert!( + observer.snapshot().iter().any(|event| { + event.payload["status"] == "switched" && event.payload["marker"] == 3 + }), + "the exact terminal result must remain retained for replay" + ); + } + + #[tokio::test] + async fn control_result_drains_ahead_of_lagged_telemetry_broadcast() { + let observer = ObserverHandle::in_process(); + let mut rx = observer.subscribe(); + for marker in 0..=OBSERVER_BUFFER_CAP { + emit_kind(&observer, "acp_read", marker); + } + emit_kind(&observer, "control_result", 7); + + let delivered = rx.recv().await.expect("protected result remains readable"); + assert_eq!(delivered.kind, "control_result"); + assert_eq!(delivered.payload["marker"], 7); + } + + #[tokio::test] + async fn control_result_lag_replays_every_unacknowledged_sequence() { + let observer = ObserverHandle::in_process(); + let mut rx = observer.subscribe(); + let burst = OBSERVER_CONTROL_RESULT_CAP + 17; + for marker in 0..burst { + emit_kind(&observer, "control_result", marker); + } + + let mut delivered = Vec::new(); + for _ in 0..burst { + let event = rx + .recv_control_result() + .await + .expect("lag recovery must replay every retained result"); + delivered.push(event.payload["marker"].as_u64().unwrap() as usize); + } + + assert_eq!( + delivered, + (0..burst).collect::>(), + "lag recovery must preserve every accepted terminal result exactly once" + ); + } + + #[tokio::test] + async fn control_result_remains_retained_until_explicit_acknowledgement() { + let observer = ObserverHandle::in_process(); + let mut rx = observer.subscribe(); + emit_kind(&observer, "control_result", 23); + + let delivered = rx + .recv_control_result() + .await + .expect("terminal result must reach its consumer"); + assert!( + observer + .snapshot() + .iter() + .any(|event| event.seq == delivered.seq), + "consumer receipt alone must not evict an unacknowledged result" + ); + + rx.acknowledge_control_result(delivered.seq); + assert!( + observer + .snapshot() + .iter() + .all(|event| event.seq != delivered.seq), + "explicit acknowledgement is the terminal ledger removal point" + ); + } + + #[tokio::test] + async fn failed_control_result_delivery_is_released_for_one_bounded_retry() { + let observer = ObserverHandle::in_process(); + let mut rx = observer.subscribe(); + emit_kind(&observer, "control_result", 29); + + let first = rx + .recv_control_result() + .await + .expect("terminal result must reach its consumer"); + assert!( + rx.retry_control_result(first.seq), + "the first failed delivery must be retained and scheduled once" + ); + + let retry = rx + .recv_control_result() + .await + .expect("released terminal result must replay"); + assert_eq!(retry.seq, first.seq); + assert_eq!(retry.payload["marker"], 29); + assert!( + !rx.retry_control_result(retry.seq), + "a second failed attempt must exhaust the bounded retry contract" + ); + + rx.acknowledge_control_result(retry.seq); + assert!( + observer + .snapshot() + .iter() + .all(|event| event.seq != retry.seq), + "a later confirmed delivery must clear retained and retry state" + ); + } + + #[test] + fn control_result_retention_is_bounded_and_rejection_is_observable() { + const EXPECTED_TERMINAL_LEDGER_CAP: usize = 1_024; + let observer = ObserverHandle::in_process(); + let rx = observer.subscribe(); + for marker in 0..=EXPECTED_TERMINAL_LEDGER_CAP { + emit_kind(&observer, "control_result", marker); + } + + assert_eq!( + observer + .snapshot() + .iter() + .filter(|event| event.kind == "control_result") + .count(), + EXPECTED_TERMINAL_LEDGER_CAP, + "terminal retention must reject newest work instead of growing without bound" + ); + assert!( + rx.control_result_admission_failed(), + "a rejected terminal result must make delivery verification fail closed" + ); + } + + #[test] + fn control_result_snapshot_capacity_is_reserved_from_telemetry_overflow() { + let observer = ObserverHandle::in_process(); + emit_kind(&observer, "control_result", 11); + for marker in 0..=OBSERVER_BUFFER_CAP { + emit_kind(&observer, "acp_read", marker); + } + + let snapshot = observer.snapshot(); + assert_eq!( + snapshot.first().map(|event| event.kind.as_str()), + Some("control_result"), + "protected results drain before the saturated telemetry snapshot" + ); + assert_eq!(snapshot[0].payload["marker"], 11); + assert_eq!( + snapshot + .iter() + .filter(|event| event.kind == "acp_read") + .count(), + OBSERVER_BUFFER_CAP + ); + } +} diff --git a/crates/buzz-acp/src/pool.rs b/crates/buzz-acp/src/pool.rs index b1fd68d044..726998696c 100644 --- a/crates/buzz-acp/src/pool.rs +++ b/crates/buzz-acp/src/pool.rs @@ -47,6 +47,13 @@ const RECENT_ACTIVITY_WINDOW: Duration = Duration::from_secs(60); // FlushBatch and BatchEvent derive Clone (added in queue.rs) so we can store // a recoverable copy in TaskMeta for panic recovery in Queue mode. +/// An acknowledged owner control whose batch fate is "drop, never replay." +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum AcceptedDropControl { + Cancel, + Rotate, +} + /// Metadata stored per in-flight task for panic recovery. pub struct TaskMeta { pub agent_index: usize, @@ -55,8 +62,23 @@ pub struct TaskMeta { pub turn_id: String, /// Clone of batch for Queue mode panic recovery. pub recoverable_batch: Option, - /// Control signal for the in-flight prompt task. - /// `None` for heartbeat tasks (not controllable) and after signal is consumed. + /// Model intent owned by the checked-out slot at dispatch time. Updated + /// when a busy `SwitchModel` control lands so panic recovery cannot pin a + /// replay to a freshly spawned default-model process. + pub desired_model: Option, + pub model_overridden: bool, + /// A `SwitchModel` command accepted by this task's oneshot but not yet + /// proven consumed by the task. The supervisor reconciles this intent + /// against the returned agent so a prompt/control simultaneous-ready race + /// cannot silently discard an acknowledged model change. + pub accepted_model_switch: Option, + /// An owner-requested drop control accepted by this task's oneshot but not + /// yet proven consumed. The supervisor owns this disposition so neither a + /// simultaneous-ready result nor a panic can replay explicitly dropped + /// work. `Rotate` also requires recycling an unretired source session. + pub accepted_drop_control: Option, + /// Control signal for the in-flight prompt or heartbeat task. + /// `None` only after a signal is consumed. pub control_tx: Option>, /// Steer request channel for non-cancelling mid-turn delivery. /// Capacity-1; `try_send` from the main loop fails on `Full`/`Closed`, @@ -153,13 +175,21 @@ pub struct OwnedAgent { pub state: SessionState, /// Model catalog from first session/new. None until first session created. pub model_capabilities: Option, - /// Desired model ID (from `Config.model`). Applied after every `session_new_full()`. + /// Desired model ID from configuration or a live switch. Applied after + /// every `session_new_full()`. pub desired_model: Option, /// Whether `desired_model` was set by a live `SwitchModel` control signal /// (as opposed to being derived from config/persona at spawn). Used by the /// desktop reader to distinguish a genuine runtime override from a stale - /// session whose persona model was edited. Reset on spawn/restart. + /// session whose persona model was edited. Preserved across an in-process + /// replacement while a live switch is pending. pub model_overridden: bool, + /// Correlation ID for a live model-switch control awaiting its terminal + /// `control_result`. Configured model defaults never carry one. + pub model_switch_request_id: Option, + /// Exact prior model intent to restore if the live switch is rejected by + /// the adapter. Nested state preserves a superseded pending switch too. + pub(crate) model_switch_rollback: Option>, /// Normalized agent name from initialize (`agentInfo.name`/`serverInfo.name`). pub agent_name: String, /// Whether Goose accepted its custom system-prompt method. `None` probes on @@ -226,6 +256,12 @@ pub struct PromptResult { pub outcome: PromptOutcome, /// Present on failure in Queue mode, for requeue. pub batch: Option, + /// Exact agent slot required when a busy `SwitchModel` turn is replayed. + /// + /// `None` for every other result. The supervisor persists this requirement + /// in `EventQueue` until the replay completes, including across process + /// recycling or cleanup-triggered replacement. + pub retry_agent_index: Option, } /// Whether the prompt came from a channel event or a heartbeat. @@ -235,30 +271,35 @@ pub enum PromptSource { Heartbeat, } -/// Apply state effects for Race 1, where a control signal arrives just after the -/// prompt completed naturally. The prompt result has already been consumed by -/// `select!`, so the harness must synthesize a successful result while still -/// honoring any load-bearing control signal semantics. -fn apply_completed_before_control_signal( - state: &mut SessionState, - source: &PromptSource, - control_signal: &ControlSignal, -) { - // Rotate and SwitchModel both invalidate so the next turn creates a fresh - // session. For SwitchModel the caller has already set `desired_model`, so - // the fresh session applies the new model on its next creation. - if matches!( - control_signal, - ControlSignal::Rotate | ControlSignal::SwitchModel(_) - ) { - state.invalidate(source); +/// Control signal for an in-flight channel turn. +/// +/// A validated live model-switch request. +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct ModelSwitchRequest { + pub model_id: String, + pub request_id: String, +} + +/// Transactional rollback state for a live model switch. +#[derive(Clone, Debug)] +pub(crate) struct ModelSwitchRollback { + pub desired_model: Option, + pub model_overridden: bool, + pub request_id: Option, + pub previous: Option>, +} + +impl ModelSwitchRequest { + pub fn new(model_id: impl Into, request_id: impl Into) -> Self { + Self { + model_id: model_id.into(), + request_id: request_id.into(), + } } } -/// Control signal for an in-flight channel turn. -/// -/// Not `Copy`: `SwitchModel` carries an owned `String`. Callers must clone when -/// a value is needed after a move, or match by reference. +/// Not `Copy`: `SwitchModel` carries owned request data. Callers must clone +/// when a value is needed after a move, or match by reference. #[derive(Clone, Debug, Eq, PartialEq)] pub enum ControlSignal { /// Stop the current turn and drop its triggering batch. @@ -279,8 +320,8 @@ pub enum ControlSignal { /// re-runs on a fresh session under the new model. The model lands by /// setting `OwnedAgent::desired_model` before invalidation; the requeued /// turn re-creates the session and re-applies `desired_model`. Runtime-only - /// — never persisted, gone on restart/respawn. - SwitchModel(String), + /// — never persisted, gone on supervisor restart. + SwitchModel(ModelSwitchRequest), } /// Goose-native non-cancelling steer request, sent from the main loop to an @@ -423,6 +464,11 @@ pub enum TimeoutKind { #[allow(dead_code)] pub enum PromptOutcome { Ok(StopReason), + /// The turn completed successfully and Buzz positively closed the session + /// before returning the healthy adapter. Distinct from `Ok` so the + /// supervisor can preserve exact-slot model affinity after local session + /// invalidation, including max-turn-count retirement with `EndTurn`. + SessionRetired(StopReason), Error(AcpError), AgentExited, Timeout(TimeoutKind), @@ -439,6 +485,37 @@ pub enum PromptOutcome { /// `CancelReason` on the batch (steer/interrupt requeue, explicit cancel /// drops) rather than the hard-cap's unconditional dead-letter. CancelDrainTimeout(Duration), + /// A post-`session/cancel` response failed in a way that does not prove + /// remote prompt/session cleanup completed. Even application-class + /// JSON-RPC errors poison the adapter at this boundary: local invalidation + /// plus reuse would abandon remote ownership. + CancelCleanupFailed(AcpError), + /// An adapter-owned session could not be closed after Buzz decided to + /// retire it. Remote resource ownership is now uncertain, so the adapter + /// must never be reused regardless of whether the underlying error looks + /// transport- or application-class. The supervisor replaces the process + /// group, which is the only fail-closed cleanup available when close did + /// not positively acknowledge. + SessionCloseFailed(AcpError), + /// A remote session was created, later setup failed, and close did not + /// positively acknowledge. The process is poisoned, while the untouched + /// batch follows ordinary bounded retry policy. + SessionSetupCloseFailed(AcpError), + /// The adapter did not advertise optional ACP `session/close`. + /// + /// Buzz must replace the process group to release adapter-owned sessions, + /// but this negotiated compatibility path is not a crash and must not + /// consume circuit-breaker budget. + SessionRecycleRequired, + /// A remote session was created, later setup failed, and optional close was + /// not advertised. Process recycling is the only safe cleanup; the untouched + /// batch follows ordinary bounded retry policy. + SessionSetupRecycleRequired, + /// An owner control preempted an in-flight session setup request. The + /// adapter must be recycled because the dropped request can still complete + /// remotely, while the control-framed batch is preserved without consuming + /// retry or dead-letter budget. + ControlPreemptedSetup, } /// Immutable config subset shared (via `Arc`) by all spawned prompt tasks. @@ -594,6 +671,14 @@ impl AgentPool { idx.map(|i| self.agents[i].take().unwrap()) } + /// Claim one exact idle slot. + /// + /// Used for busy `SwitchModel` replay: falling back to a different slot + /// would run the preserved batch under that slot's default model. + pub fn try_claim_index(&mut self, index: usize) -> Option { + self.agents.get_mut(index)?.take() + } + /// Return an agent to its slot after a task completes. pub fn return_agent(&mut self, agent: OwnedAgent) { let idx = agent.index; @@ -642,6 +727,24 @@ impl AgentPool { &mut self.task_map } + /// Cooperatively cancel every checked-out task while retaining its typed + /// [`OwnedAgent`] return path through `result_rx`. + /// + /// Shutdown must use this before joining tasks. Aborting a task would drop + /// the only typed process owner and make direct-child reaping and + /// process-group absence impossible to prove. + pub fn cancel_checked_out_for_shutdown(&mut self) -> usize { + let mut signalled = 0; + for meta in self.task_map.values_mut() { + if let Some(control_tx) = meta.control_tx.take() { + if control_tx.send(ControlSignal::Cancel).is_ok() { + signalled += 1; + } + } + } + signalled + } + /// Try to send a goose-native steer request to the in-flight task for /// `channel_id`. /// @@ -717,29 +820,30 @@ impl AgentPool { &mut self.agents } - /// Remove the session for `channel_id` from all idle agents. - /// - /// Called when the agent is removed from a channel — stale sessions - /// should not be reused. Checked-out agents (in-flight) are not - /// modified; their sessions will fail naturally on the next prompt - /// if the relay rejects the request. + /// Remove every idle adapter that still owns a session for `channel_id`. /// - /// Returns the number of sessions invalidated. - pub fn invalidate_channel_sessions(&mut self, channel_id: Uuid) -> usize { - let mut count = 0; - for slot in &mut self.agents { - if let Some(agent) = slot.as_mut() { - if agent.state.invalidate_channel(&channel_id) { - count += 1; + /// Callers must shut down/recycle each returned process before allowing its + /// slot back into service. Local-only invalidation is forbidden because it + /// would abandon remote subprocess ownership. + pub fn take_idle_agents_with_session(&mut self, channel_id: Uuid) -> Vec { + self.agents + .iter_mut() + .filter_map(|slot| { + if slot + .as_ref() + .is_some_and(|agent| agent.state.sessions.contains_key(&channel_id)) + { + slot.take() + } else { + None } - } - } - count + }) + .collect() } - /// Idle-path model switch: set `desired_model` on the idle agent for - /// `channel_id` and invalidate its session so the next turn re-creates the - /// session under the new model. + /// Idle-path model switch: validate the requested model, claim the adapter + /// that owns `channel_id`, and return it for process recycling. The next + /// process creates a fresh session under the preserved model intent. /// /// Pre-cancel guard: the desired model is validated against the agent's /// cached catalog *before* the session is invalidated, so an unsupported @@ -754,18 +858,20 @@ impl AgentPool { &mut self, channel_id: Uuid, model_id: &str, + request_id: &str, ) -> IdleSwitchResult { - let Some(agent) = self - .agents - .iter_mut() - .flatten() - .find(|a| a.state.sessions.contains_key(&channel_id)) - else { + let Some(position) = self.agents.iter().position(|slot| { + slot.as_ref() + .is_some_and(|agent| agent.state.sessions.contains_key(&channel_id)) + }) else { return IdleSwitchResult::NoIdleAgent; }; // Pre-cancel guard against the cached catalog. None = catalog not yet // populated (no session ever created); defer validation to apply time. + let Some(agent) = self.agents.get(position).and_then(Option::as_ref) else { + return IdleSwitchResult::NoIdleAgent; + }; if let Some(caps) = agent.model_capabilities.as_ref() { if !model_in_catalog( &caps.config_options_raw, @@ -776,18 +882,18 @@ impl AgentPool { } } - agent.desired_model = Some(model_id.to_string()); - agent.model_overridden = true; - agent.state.invalidate_channel(&channel_id); - IdleSwitchResult::Switched + let Some(mut agent) = self.agents.get_mut(position).and_then(Option::take) else { + return IdleSwitchResult::NoIdleAgent; + }; + begin_model_switch(&mut agent, &ModelSwitchRequest::new(model_id, request_id)); + IdleSwitchResult::Recycle(Box::new(agent)) } } /// Outcome of [`AgentPool::switch_idle_agent_model`]. -#[derive(Debug, PartialEq, Eq)] pub enum IdleSwitchResult { - /// `desired_model` set and the channel session invalidated. - Switched, + /// `desired_model` set; caller must recycle this still-owning adapter. + Recycle(Box), /// Desired model is not in the agent's cached catalog — pick rejected, /// session untouched. UnsupportedModel, @@ -795,6 +901,68 @@ pub enum IdleSwitchResult { NoIdleAgent, } +/// Stage a busy model switch without losing the slot's last accepted intent. +/// +/// A busy control is handled before its current session is retired, so this is +/// the last ownership boundary that still has both the prior intent and the +/// catalog that accepted it. Validate here before respawn preservation copies +/// `desired_model`: a rejected request must leave the exact prior model and +/// override provenance in place for the requeued turn. +/// +/// Returns `true` when the new intent was accepted, or `false` after restoring +/// the prior intent and emitting `unsupported_model`. +fn begin_model_switch(agent: &mut OwnedAgent, request: &ModelSwitchRequest) { + agent.model_switch_rollback = Some(Box::new(ModelSwitchRollback { + desired_model: agent.desired_model.take(), + model_overridden: agent.model_overridden, + request_id: agent.model_switch_request_id.take(), + previous: agent.model_switch_rollback.take(), + })); + agent.desired_model = Some(request.model_id.clone()); + agent.model_overridden = true; + agent.model_switch_request_id = Some(request.request_id.clone()); +} + +fn restore_previous_model_intent(agent: &mut OwnedAgent) { + let Some(rollback) = agent.model_switch_rollback.take() else { + agent.model_switch_request_id = None; + return; + }; + agent.desired_model = rollback.desired_model; + agent.model_overridden = rollback.model_overridden; + agent.model_switch_request_id = rollback.request_id; + agent.model_switch_rollback = rollback.previous; +} + +pub(crate) fn stage_busy_model_switch_intent( + agent: &mut OwnedAgent, + request: &ModelSwitchRequest, +) -> bool { + begin_model_switch(agent, request); + let rejected = agent.model_capabilities.as_ref().is_some_and(|caps| { + !model_in_catalog( + &caps.config_options_raw, + caps.available_models_raw.as_ref(), + &request.model_id, + ) + }); + if !rejected { + return true; + } + + restore_previous_model_intent(agent); + agent.acp.observe( + "control_result", + serde_json::json!({ + "type": "switch_model", + "status": "unsupported_model", + "modelId": request.model_id, + "requestId": request.request_id, + }), + ); + false +} + /// 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. @@ -859,19 +1027,63 @@ async fn resolve_new_session_channel_context( (is_dm, title_channel) } -/// Create a new ACP session via `session_new_full()`, populate model capabilities -/// on the agent (first session only), and apply `desired_model` if set. -/// -/// On error from `session_new_full()`, returns the `AcpError` — caller handles -/// error reporting. Model-switch failures are logged and gracefully ignored -/// (the agent proceeds with its default model). +enum SessionSetupError { + Agent(AcpError), + CloseFailed(AcpError), + RecycleRequired, +} + +struct PreparedSession { + session_id: String, + session_config: serde_json::Value, + terminal_model_switch: Option, +} + +impl From for SessionSetupError { + fn from(error: AcpError) -> Self { + Self::Agent(error) + } +} + +async fn cleanup_failed_session_setup( + agent: &mut OwnedAgent, + session_id: &str, + setup_error: AcpError, +) -> SessionSetupError { + // A transport/framing failure already poisons the stream. Sending another + // request cannot prove cleanup; the normal fatal-error path replaces it. + if setup_error.requires_process_replacement() { + return SessionSetupError::Agent(setup_error); + } + if !agent.acp.supports_session_close() { + return SessionSetupError::RecycleRequired; + } + match agent.acp.session_close(session_id).await { + Ok(_) => SessionSetupError::Agent(setup_error), + Err(close_error) => SessionSetupError::CloseFailed(close_error), + } +} + +fn session_setup_outcome(error: SessionSetupError) -> PromptOutcome { + match error { + SessionSetupError::Agent(AcpError::AgentExited) => PromptOutcome::AgentExited, + SessionSetupError::Agent(error) => PromptOutcome::Error(error), + SessionSetupError::CloseFailed(error) => PromptOutcome::SessionSetupCloseFailed(error), + SessionSetupError::RecycleRequired => PromptOutcome::SessionSetupRecycleRequired, + } +} + +/// Create a new ACP session, complete every required setup RPC, and return +/// ownership only after setup succeeds. If an application-class setup error +/// occurs after `session/new`, close the partial remote session before allowing +/// adapter reuse; otherwise require process replacement. async fn create_session_and_apply_model( agent: &mut OwnedAgent, ctx: &PromptContext, agent_core: Option<&str>, agent_canvas: Option<&str>, channel_name: Option<&str>, -) -> Result { +) -> Result { // Build base_prompt + system_prompt + agent core + canvas metadata into a // single prompt. Standard protocol-v2 agents receive it in `session/new`; // Goose receives it through the custom request below. Legacy agents receive @@ -907,7 +1119,8 @@ async fn create_session_and_apply_model( ), session_title.as_deref(), ) - .await?; + .await + .map_err(SessionSetupError::Agent)?; if is_goose && agent.goose_system_prompt_supported != Some(false) { if let Some(prompt) = combined_system_prompt.as_deref() { @@ -924,7 +1137,9 @@ async fn create_session_and_apply_model( "Goose does not support its system-prompt extension; using user-message framing" ); } - Err(error) => return Err(error), + Err(error) => { + return Err(cleanup_failed_session_setup(agent, &resp.session_id, error).await) + } } } } @@ -940,44 +1155,69 @@ async fn create_session_and_apply_model( // Apply desired_model if set, matching against the fresh session/new response. // Track whether the switch succeeded so session_config_captured reflects // the post-switch state (not the pre-switch desired state). - let switch_succeeded = if let Some(ref desired) = agent.desired_model { - match resolve_model_switch_method(&resp.raw, desired) { + let mut terminal_model_switch = None; + let mut switch_succeeded = false; + while let Some(desired) = agent.desired_model.clone() { + let switch_request_id = agent.model_switch_request_id.clone(); + match resolve_model_switch_method(&resp.raw, &desired) { Some(method) => { - apply_model_switch(&mut agent.acp, &resp.session_id, desired, &method).await?; - true + if let Err(error) = + apply_model_switch(&mut agent.acp, &resp.session_id, &desired, &method).await + { + if let Some(request_id) = switch_request_id.as_deref() { + agent.acp.observe( + "control_result", + serde_json::json!({ + "type": "switch_model", + "status": "switch_failed", + "modelId": desired, + "requestId": request_id, + }), + ); + restore_previous_model_intent(agent); + } + return Err(cleanup_failed_session_setup(agent, &resp.session_id, error).await); + } + if let Some(request_id) = switch_request_id.as_deref() { + terminal_model_switch = + Some(ModelSwitchRequest::new(desired.clone(), request_id)); + } + switch_succeeded = true; + break; } None => { + // Surface the miss so the desktop ModelPicker can reject a live + // pick rather than silently no-op. Restore and resolve the prior + // intent against this same fresh catalog before publishing the + // session: merely restoring its fields would otherwise commit an + // adapter-default session and strand any older request ID. + if let Some(request_id) = switch_request_id.as_deref() { + agent.acp.observe( + "control_result", + serde_json::json!({ + "type": "switch_model", + "status": "unsupported_model", + "modelId": desired, + "requestId": request_id, + }), + ); + restore_previous_model_intent(agent); + continue; + } tracing::warn!( target: "pool::model", - "desired model {desired} not found in agent's available models — proceeding with agent default" + "configured model {desired} not found in agent's available models — proceeding with agent default" ); - // Surface the miss so the desktop ModelPicker can reject a live - // pick rather than silently no-op. On the busy path the turn has - // already been cancelled+requeued by the time we get here, so the - // turn restarts on the unchanged model and the user is told no. - agent.acp.observe( - "control_result", - serde_json::json!({ - "type": "switch_model", - "status": "unsupported_model", - "modelId": desired, - }), - ); - false + break; } } - } else { - false - }; + } - // Emit session config for desktop consumption (config bridge tier 1b). - // Emitted AFTER desired_model resolution so the desktop caches the - // post-switch state. modelOverridden reflects whether the switch actually - // applied — false on the unsupported arm so the panel doesn't show a - // stale override badge. - agent.acp.observe( - "session_config_captured", - serde_json::json!({ + // Prepare session config for desktop consumption (config bridge tier 1b). + // Published AFTER desired_model resolution so the desktop caches the + // post-switch state. modelOverridden reflects the intent that was actually + // applied, including a prior pending override restored after a catalog miss. + let session_config = serde_json::json!({ "configOptions": resp.raw.get("configOptions").cloned().unwrap_or(serde_json::Value::Null), "modes": resp.raw.get("modes").cloned().unwrap_or(serde_json::Value::Null), "models": resp.raw.get("models").cloned().unwrap_or(serde_json::Value::Null), @@ -985,8 +1225,7 @@ async fn create_session_and_apply_model( // Pair identity for the desktop session-config cache, which is // keyed by (agent, relay) like the lifecycle frames. "relayUrl": ctx.relay_url, - }), - ); + }); // Apply permission mode if not the agent's built-in default AND the agent // advertises the requested mode in session/new. Agents that don't support @@ -995,18 +1234,114 @@ async fn create_session_and_apply_model( if !ctx.permission_mode.is_default() && agent_supports_mode(&resp.raw, ctx.permission_mode.as_wire_str()) { - apply_permission_mode(&mut agent.acp, &resp.session_id, &ctx.permission_mode).await?; + if let Err(error) = + apply_permission_mode(&mut agent.acp, &resp.session_id, &ctx.permission_mode).await + { + return Err(cleanup_failed_session_setup(agent, &resp.session_id, error).await); + } + } + + Ok(PreparedSession { + session_id: resp.session_id, + session_config, + terminal_model_switch, + }) +} + +/// Publish session-derived state only after the prepared remote session is +/// durably addressable through the local routing state. +/// +/// A failed invariant check leaves any live model-switch correlation pending so +/// replacement/retry can either complete it or report a terminal failure. +fn commit_prepared_session( + agent: &mut OwnedAgent, + source: &PromptSource, + prepared: &PreparedSession, +) -> bool { + let session_is_inserted = match source { + PromptSource::Channel(channel_id) => agent + .state + .sessions + .get(channel_id) + .is_some_and(|session_id| session_id == &prepared.session_id), + PromptSource::Heartbeat => { + agent.state.heartbeat_session.as_deref() == Some(prepared.session_id.as_str()) + } + }; + if !session_is_inserted { + tracing::error!( + target: "pool::session", + session_id = %prepared.session_id, + "prepared session was not inserted before commit" + ); + return false; + } + + if let Some(request) = prepared.terminal_model_switch.as_ref() { + let switch_still_pending = agent.desired_model.as_deref() + == Some(request.model_id.as_str()) + && agent.model_switch_request_id.as_deref() == Some(request.request_id.as_str()); + if !switch_still_pending { + tracing::error!( + target: "pool::model", + session_id = %prepared.session_id, + request_id = %request.request_id, + "prepared model switch no longer matches pending agent state" + ); + return false; + } } - Ok(resp.session_id) + agent + .acp + .observe("session_config_captured", prepared.session_config.clone()); + + if let Some(request) = prepared.terminal_model_switch.as_ref() { + agent.acp.observe( + "control_result", + serde_json::json!({ + "type": "switch_model", + "status": "switched", + "modelId": request.model_id, + "requestId": request.request_id, + }), + ); + agent.model_switch_request_id = None; + agent.model_switch_rollback = None; + } + + true +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum SessionRetirement { + Closed, + RecycleRequired, +} + +/// Close an adapter-owned session before forgetting its local routing state. +/// +/// If close fails, the local session remains present so callers cannot +/// accidentally reuse the adapter as though retirement had succeeded. +async fn retire_session( + agent: &mut OwnedAgent, + source: &PromptSource, + session_id: &str, +) -> Result { + if !agent.acp.supports_session_close() { + return Ok(SessionRetirement::RecycleRequired); + } + agent.acp.session_close(session_id).await?; + agent.state.invalidate(source); + Ok(SessionRetirement::Closed) } /// Send the appropriate ACP model-switch request with a timeout. /// -/// On timeout or error, logs a warning and returns — the caller proceeds -/// with the agent's default model. This is intentionally non-fatal: a stale -/// response from a timed-out request is safely ignored by `read_until_response` -/// (non-matching JSON-RPC IDs are skipped). +/// Every rejected or unverified request is returned to the caller. A framed +/// application error leaves the transport reusable, but it is still proof that +/// the requested model was not applied; the caller must report terminal switch +/// failure and retire the partially-created session before adapter reuse. async fn apply_model_switch( acp: &mut AcpClient, session_id: &str, @@ -1045,23 +1380,21 @@ async fn apply_model_switch( } // 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)) => { + Ok(Err(e)) if e.requires_process_replacement() => { 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. + // A framed JSON-RPC application error leaves the transport reusable, + // but is still a terminal failure for this model-switch attempt. Ok(Err(e)) => { tracing::warn!( target: "pool::model", - "failed to set model {desired} via {method_label}: {e} — proceeding with agent default" + "failed to set model {desired} via {method_label}: {e}" ); + return Err(e); } Err(_) => { // Outer timeout fired — the inner send_request may have left the @@ -1121,11 +1454,7 @@ async fn apply_permission_mode( } // 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)) => { + Ok(Err(e)) if e.requires_process_replacement() => { tracing::error!( target: "pool::permission", "fatal error setting permission mode {wire:?}: {e}" @@ -1304,12 +1633,24 @@ fn with_canvas(prompt: Option, canvas: Option<&str>) -> Option { /// /// On the happy path the read loop has already called `take()`, so this is a no-op. fn send_prompt_result( + result_tx: &mpsc::UnboundedSender, + turn_id: &str, + agent: OwnedAgent, + source: PromptSource, + outcome: PromptOutcome, + batch: Option, +) { + send_prompt_result_with_retry_agent(result_tx, turn_id, agent, source, outcome, batch, None); +} + +fn send_prompt_result_with_retry_agent( result_tx: &mpsc::UnboundedSender, turn_id: &str, mut agent: OwnedAgent, source: PromptSource, outcome: PromptOutcome, batch: Option, + retry_agent_index: Option, ) { agent.acp.clear_steer_rx(); let _ = result_tx.send(PromptResult { @@ -1318,6 +1659,7 @@ fn send_prompt_result( turn_id: turn_id.to_owned(), outcome, batch, + retry_agent_index, }); } @@ -1333,6 +1675,123 @@ fn send_prompt_result( /// /// The agent is ALWAYS returned — even on panic the `JoinSet` detects the /// abort and the caller uses `task_map` to recover the agent index. +enum Preemptible { + Completed(T), + Controlled(ControlSignal), +} + +async fn await_with_control( + control_rx: &mut Option>, + future: F, +) -> Preemptible +where + F: std::future::Future, +{ + let Some(rx) = control_rx.as_mut() else { + return Preemptible::Completed(future.await); + }; + tokio::pin!(future); + tokio::select! { + biased; + signal = rx => { + *control_rx = None; + Preemptible::Controlled(signal.unwrap_or(ControlSignal::Cancel)) + } + result = &mut future => Preemptible::Completed(result), + } +} + +#[allow(clippy::too_many_arguments)] +async fn send_pre_prompt_control_result( + result_tx: &mpsc::UnboundedSender, + turn_id: &str, + mut agent: OwnedAgent, + source: PromptSource, + signal: ControlSignal, + ctx: &PromptContext, + batch: Option, + session_id: Option, + setup_request_in_flight: bool, + retire_idle_session: bool, +) { + let retry_agent_index = matches!(&signal, ControlSignal::SwitchModel(_)).then_some(agent.index); + if let ControlSignal::SwitchModel(ref request) = signal { + stage_busy_model_switch_intent(&mut agent, request); + } + let retry_batch = requeue_cancelled_batch(ctx, signal.clone(), batch); + + let outcome = if setup_request_in_flight { + // Dropping session/new or model-application leaves an ACP request + // whose remote completion can no longer be correlated safely. Keep + // ownership typed, but force bounded process recycling. + agent.state.invalidate_all(); + PromptOutcome::ControlPreemptedSetup + } else if let Some(session_id) = session_id { + if agent.acp.has_in_flight_prompt() { + match agent + .acp + .cancel_with_cleanup_grace(&session_id, CONTROL_CANCEL_GRACE) + .await + { + Ok(stop_reason) => { + log_stop_reason(&source, &stop_reason); + match retire_session(&mut agent, &source, &session_id).await { + Ok(SessionRetirement::Closed) => PromptOutcome::Cancelled, + Ok(SessionRetirement::RecycleRequired) => { + PromptOutcome::SessionRecycleRequired + } + Err(error) => PromptOutcome::SessionCloseFailed(error), + } + } + Err(error) => { + let failure = + classify_control_cancel_failure(ctx, error, signal, retry_batch.clone()); + if failure.invalidate_all { + agent.state.invalidate_all(); + } else { + agent.state.invalidate(&source); + } + send_prompt_result_with_retry_agent( + result_tx, + turn_id, + agent, + source, + failure.outcome, + failure.retry_batch, + retry_agent_index, + ); + return; + } + } + } else if retire_idle_session + || matches!( + signal, + ControlSignal::Rotate | ControlSignal::SwitchModel(_) + ) + { + match retire_session(&mut agent, &source, &session_id).await { + Ok(SessionRetirement::Closed) => PromptOutcome::Cancelled, + Ok(SessionRetirement::RecycleRequired) => PromptOutcome::SessionRecycleRequired, + Err(error) => PromptOutcome::SessionCloseFailed(error), + } + } else { + PromptOutcome::Cancelled + } + } else { + PromptOutcome::Cancelled + }; + + send_prompt_result_with_retry_agent( + result_tx, + turn_id, + agent, + source, + outcome, + retry_batch, + retry_agent_index, + ); +} + pub async fn run_prompt_task( mut agent: OwnedAgent, batch: Option, @@ -1342,6 +1801,7 @@ pub async fn run_prompt_task( control_rx: Option>, turn_id: String, ) { + let mut control_rx = control_rx; // Is this a channel prompt or a heartbeat? let source = match &batch { Some(b) => PromptSource::Channel(b.channel_id), @@ -1460,16 +1920,31 @@ pub async fn run_prompt_task( &ctx.agent_keys, owner_pk, ); - let section = match tokio::time::timeout(CORE_FETCH_TIMEOUT, fetch).await { - Ok(s) => s, - Err(_) => { - tracing::warn!( - target: "engram::core", - channel = %cid, - timeout_ms = CORE_FETCH_TIMEOUT.as_millis() as u64, - "core fetch timed out — emitting no section" - ); - None + let section = match await_with_control( + &mut control_rx, + tokio::time::timeout(CORE_FETCH_TIMEOUT, fetch), + ) + .await + { + Preemptible::Completed(result) => match result { + Ok(s) => s, + Err(_) => { + tracing::warn!( + target: "engram::core", + channel = %cid, + timeout_ms = CORE_FETCH_TIMEOUT.as_millis() as u64, + "core fetch timed out — emitting no section" + ); + None + } + }, + Preemptible::Controlled(signal) => { + send_pre_prompt_control_result( + &result_tx, &turn_id, agent, source, signal, &ctx, batch, None, false, + false, + ) + .await; + return; } }; if let Some(rendered) = section { @@ -1506,13 +1981,43 @@ pub async fn run_prompt_task( let needs_canvas = is_new_channel_session && !agent.state.canvas_sections.contains_key(cid); let needs_title = is_new_channel_session && ctx.session_title.is_some(); if needs_canvas || needs_title { - let (is_dm, resolved_channel) = - resolve_new_session_channel_context(&ctx.channel_info, *cid).await; + let (is_dm, resolved_channel) = match await_with_control( + &mut control_rx, + resolve_new_session_channel_context(&ctx.channel_info, *cid), + ) + .await + { + Preemptible::Completed(result) => result, + Preemptible::Controlled(signal) => { + send_pre_prompt_control_result( + &result_tx, &turn_id, agent, source, signal, &ctx, batch, None, false, + false, + ) + .await; + return; + } + }; title_channel = resolved_channel; // A confirmed DM never receives a canvas section; an undeterminable // channel type fails closed as a DM for the same reason. if needs_canvas && !is_dm { - if let Some(section) = fetch_canvas_section(*cid, &ctx.rest_client).await { + let section = match await_with_control( + &mut control_rx, + fetch_canvas_section(*cid, &ctx.rest_client), + ) + .await + { + Preemptible::Completed(section) => section, + Preemptible::Controlled(signal) => { + send_pre_prompt_control_result( + &result_tx, &turn_id, agent, source, signal, &ctx, batch, None, false, + false, + ) + .await; + return; + } + }; + if let Some(section) = section { pending_canvas = Some((*cid, section)); } } @@ -1547,92 +2052,101 @@ pub async fn run_prompt_task( // agent in several channels doesn't produce identical session // rows; `title_channel` comes from the single resolve above and // is `None` for DM, unresolved, and unnamed channels. - match create_session_and_apply_model( - &mut agent, - &ctx, - agent_core.as_deref(), - agent_canvas.as_deref(), - title_channel.as_deref(), + match await_with_control( + &mut control_rx, + create_session_and_apply_model( + &mut agent, + &ctx, + agent_core.as_deref(), + agent_canvas.as_deref(), + title_channel.as_deref(), + ), ) .await { - Ok(sid) => { - tracing::info!( - target: "pool::session", - "created session {sid} for channel {cid}" - ); - agent.state.sessions.insert(*cid, sid.clone()); - // Commit canvas only after session creation succeeds (I3). - if let Some((pending_cid, section)) = pending_canvas.take() { - agent.state.canvas_sections.insert(pending_cid, section); - } - (sid, true) - } - Err(AcpError::AgentExited) => { - agent.state.invalidate_all(); - send_prompt_result( - &result_tx, - &turn_id, - agent, - source, - PromptOutcome::AgentExited, - requeue_batch_if_queue(&ctx, batch), - ); - return; - } - Err(e) => { - // Session creation failed; pending canvas was never committed, - // so the next retry will re-fetch a fresh revision. - send_prompt_result( - &result_tx, - &turn_id, - agent, - source, - PromptOutcome::Error(e), - requeue_batch_if_queue(&ctx, batch), - ); + Preemptible::Controlled(signal) => { + send_pre_prompt_control_result( + &result_tx, &turn_id, agent, source, signal, &ctx, batch, None, true, + false, + ) + .await; return; } - } - } - } - PromptSource::Heartbeat => { - if let Some(sid) = &agent.state.heartbeat_session { - (sid.clone(), false) - } else { - match create_session_and_apply_model(&mut agent, &ctx, None, None, None).await { - Ok(sid) => { - tracing::info!( - target: "pool::session", - "created heartbeat session {sid} for agent {}", - agent.index - ); - agent.state.heartbeat_session = Some(sid.clone()); - (sid, true) - } - Err(AcpError::AgentExited) => { - agent.state.invalidate_all(); - send_prompt_result( - &result_tx, - &turn_id, - agent, - source, - PromptOutcome::AgentExited, - None, - ); - return; - } - Err(e) => { - send_prompt_result( - &result_tx, - &turn_id, - agent, - source, - PromptOutcome::Error(e), - None, - ); + Preemptible::Completed(result) => match result { + Ok(prepared) => { + let sid = prepared.session_id.clone(); + tracing::info!( + target: "pool::session", + "created session {sid} for channel {cid}" + ); + agent.state.sessions.insert(*cid, sid.clone()); + // Commit canvas only after session creation succeeds (I3). + if let Some((pending_cid, section)) = pending_canvas.take() { + agent.state.canvas_sections.insert(pending_cid, section); + } + commit_prepared_session(&mut agent, &source, &prepared); + (sid, true) + } + Err(error) => { + // Session creation failed; pending canvas was never committed, + // so the next retry will re-fetch a fresh revision. + let outcome = session_setup_outcome(error); + if matches!(&outcome, PromptOutcome::AgentExited) { + agent.state.invalidate_all(); + } + send_prompt_result( + &result_tx, + &turn_id, + agent, + source, + outcome, + requeue_batch_if_queue(&ctx, batch), + ); + return; + } + }, + } + } + } + PromptSource::Heartbeat => { + if let Some(sid) = &agent.state.heartbeat_session { + (sid.clone(), false) + } else { + match await_with_control( + &mut control_rx, + create_session_and_apply_model(&mut agent, &ctx, None, None, None), + ) + .await + { + Preemptible::Controlled(signal) => { + send_pre_prompt_control_result( + &result_tx, &turn_id, agent, source, signal, &ctx, batch, None, true, + false, + ) + .await; return; } + Preemptible::Completed(result) => match result { + Ok(prepared) => { + let sid = prepared.session_id.clone(); + tracing::info!( + target: "pool::session", + "created heartbeat session {sid} for agent {}", + agent.index + ); + agent.state.heartbeat_session = Some(sid.clone()); + commit_prepared_session(&mut agent, &source, &prepared); + (sid, true) + } + Err(error) => { + let outcome = session_setup_outcome(error); + if matches!(&outcome, PromptOutcome::AgentExited) { + agent.state.invalidate_all(); + } + send_prompt_result(&result_tx, &turn_id, agent, source, outcome, None); + return; + } + }, } } } @@ -1685,15 +2199,35 @@ pub async fn run_prompt_task( agent_canvas.as_deref(), &init_msg, ); - let init_result = agent - .acp - .session_prompt_with_idle_timeout( + let init_result = match await_with_control( + &mut control_rx, + agent.acp.session_prompt_with_idle_timeout( &session_id, &init_msg, ctx.idle_timeout, ctx.max_turn_duration, - ) - .await; + ), + ) + .await + { + Preemptible::Completed(result) => result, + Preemptible::Controlled(signal) => { + send_pre_prompt_control_result( + &result_tx, + &turn_id, + agent, + source, + signal, + &ctx, + batch, + Some(session_id), + false, + true, + ) + .await; + return; + } + }; match init_result { Ok(stop_reason) => { @@ -1779,15 +2313,25 @@ pub async fn run_prompt_task( Err(e) => { tracing::error!( target: "pool::session", - "initial_message failed for channel {cid}: {e} — invalidating session" + "initial_message failed for channel {cid}: {e} — retiring partial session" ); - agent.state.invalidate(&source); + let outcome = if e.requires_process_replacement() { + PromptOutcome::Error(e) + } else { + match retire_session(&mut agent, &source, &session_id).await { + Ok(SessionRetirement::Closed) => PromptOutcome::Error(e), + Ok(SessionRetirement::RecycleRequired) => { + PromptOutcome::SessionSetupRecycleRequired + } + Err(close_error) => PromptOutcome::SessionSetupCloseFailed(close_error), + } + }; send_prompt_result( &result_tx, &turn_id, agent, source, - PromptOutcome::Error(e), + outcome, requeue_batch_if_queue(&ctx, batch), ); return; @@ -1818,16 +2362,81 @@ pub async fn run_prompt_task( } else if let Some(ref b) = batch { // Build prompt from batch with context enrichment. // Try startup cache first; lazy-fetch via REST for dynamic channels. - let channel_info = ctx.channel_info.resolve(b.channel_id).await; + let channel_info = + match await_with_control(&mut control_rx, ctx.channel_info.resolve(b.channel_id)).await + { + Preemptible::Completed(result) => result, + Preemptible::Controlled(signal) => { + send_pre_prompt_control_result( + &result_tx, + &turn_id, + agent, + source, + signal, + &ctx, + batch, + Some(session_id), + false, + false, + ) + .await; + return; + } + }; let conversation_context = if ctx.context_message_limit > 0 { - fetch_conversation_context(b, &channel_info, &ctx).await + match await_with_control( + &mut control_rx, + fetch_conversation_context(b, &channel_info, &ctx), + ) + .await + { + Preemptible::Completed(result) => result, + Preemptible::Controlled(signal) => { + send_pre_prompt_control_result( + &result_tx, + &turn_id, + agent, + source, + signal, + &ctx, + batch, + Some(session_id), + false, + false, + ) + .await; + return; + } + } } else { None }; - let profile_lookup = - fetch_prompt_profile_lookup(b, conversation_context.as_ref(), &ctx.rest_client).await; + let profile_lookup = match await_with_control( + &mut control_rx, + fetch_prompt_profile_lookup(b, conversation_context.as_ref(), &ctx.rest_client), + ) + .await + { + Preemptible::Completed(result) => result, + Preemptible::Controlled(signal) => { + send_pre_prompt_control_result( + &result_tx, + &turn_id, + agent, + source, + signal, + &ctx, + batch, + Some(session_id), + false, + false, + ) + .await; + return; + } + }; let known_names: Vec<&str> = profile_lookup .iter() @@ -1897,9 +2506,10 @@ pub async fn run_prompt_task( None => prompt_sections.iter().map(String::as_str).collect(), }; - // When control_rx is Some (channel tasks), wrap the prompt in select! so - // the main loop can cancel, interrupt, or rotate it. Heartbeats - // (control_rx=None) take the simple await path — they are not controllable. + // Every supervisor-owned prompt, including heartbeats, supplies + // `control_rx` so graceful shutdown can recover the typed agent owner. + // `None` remains supported for direct callers that intentionally do not + // need cooperative cancellation. // let prompt_result = match control_rx { None => { @@ -1925,13 +2535,17 @@ pub async fn run_prompt_task( ) => result, mode = rx => { let control_signal = mode.unwrap_or(ControlSignal::Cancel); - // Land the model switch before any cancel/requeue work: setting - // `desired_model` here means the fresh session created by the - // requeued turn (busy) or the next turn (already-completed) - // applies the new model. Runtime-only — never persisted. - if let ControlSignal::SwitchModel(ref model_id) = control_signal { - agent.desired_model = Some(model_id.clone()); - agent.model_overridden = true; + let retry_agent_index = matches!( + &control_signal, + ControlSignal::SwitchModel(_) + ) + .then_some(agent.index); + // Land an accepted model switch before any cancel/requeue + // work. A catalog rejection restores the exact prior intent + // before recycling can preserve it for the requeued turn. + // Runtime-only — never persisted. + if let ControlSignal::SwitchModel(ref request) = control_signal { + stage_busy_model_switch_intent(&mut agent, request); } // Control signal received. Guard against Race 1: the turn may // have completed naturally just as cancel fired. @@ -1944,27 +2558,50 @@ pub async fn run_prompt_task( { Ok(stop_reason) => { log_stop_reason(&source, &stop_reason); - agent.state.invalidate(&source); let retry_batch = requeue_cancelled_batch(&ctx, control_signal, batch); + let outcome = + match retire_session(&mut agent, &source, &session_id).await { + Ok(SessionRetirement::Closed) => PromptOutcome::Cancelled, + Ok(SessionRetirement::RecycleRequired) => { + PromptOutcome::SessionRecycleRequired + } + Err(error) => { + tracing::error!( + target: "pool::session", + "failed to close retired session {session_id}" + ); + PromptOutcome::SessionCloseFailed(error) + } + }; let usage = agent.acp.take_turn_usage(); + let metric_stop = if matches!( + &outcome, + PromptOutcome::Cancelled + | PromptOutcome::SessionRecycleRequired + ) { + buzz_core::agent_turn_metric::StopReason::Cancelled + } else { + buzz_core::agent_turn_metric::StopReason::Error + }; publish_agent_turn_metric( &ctx, usage, observer_channel_id, &session_id, &turn_id, - Some(buzz_core::agent_turn_metric::StopReason::Cancelled), + Some(metric_stop), ) .await; - send_prompt_result( + send_prompt_result_with_retry_agent( &result_tx, &turn_id, agent, source, - PromptOutcome::Cancelled, + outcome, retry_batch, + retry_agent_index, ); return; } @@ -1994,13 +2631,14 @@ pub async fn run_prompt_task( Some(buzz_core::agent_turn_metric::StopReason::Error), ) .await; - send_prompt_result( + send_prompt_result_with_retry_agent( &result_tx, &turn_id, agent, source, failure.outcome, failure.retry_batch, + retry_agent_index, ); return; } @@ -2020,25 +2658,37 @@ pub async fn run_prompt_task( // and last_prompt_id was cleared by the success path. // // MUST send a PromptResult or the main loop deadlocks. - if matches!( + let retires_session = matches!( control_signal, ControlSignal::Rotate | ControlSignal::SwitchModel(_) - ) { + ); + let outcome = if retires_session { tracing::debug!( target: "pool::prompt", - "rotate/switch signal arrived but turn already completed — invalidating session" + "rotate/switch signal arrived after turn completion — retiring session" ); + match retire_session(&mut agent, &source, &session_id).await { + Ok(SessionRetirement::Closed) => { + PromptOutcome::SessionRetired(StopReason::EndTurn) + } + Ok(SessionRetirement::RecycleRequired) => { + PromptOutcome::SessionRecycleRequired + } + Err(error) => PromptOutcome::SessionCloseFailed(error), + } } else { tracing::debug!( target: "pool::prompt", "control signal arrived but turn already completed — treating as success" ); - } - apply_completed_before_control_signal( - &mut agent.state, - &source, - &control_signal, - ); + PromptOutcome::Ok(StopReason::EndTurn) + }; + let metric_stop = + if matches!(&outcome, PromptOutcome::SessionCloseFailed(_)) { + buzz_core::agent_turn_metric::StopReason::Error + } else { + buzz_core::agent_turn_metric::StopReason::EndTurn + }; let usage = agent.acp.take_turn_usage(); publish_agent_turn_metric( &ctx, @@ -2046,16 +2696,17 @@ pub async fn run_prompt_task( observer_channel_id, &session_id, &turn_id, - Some(buzz_core::agent_turn_metric::StopReason::EndTurn), + Some(metric_stop), ) .await; - send_prompt_result( + send_prompt_result_with_retry_agent( &result_tx, &turn_id, agent, source, - PromptOutcome::Ok(StopReason::EndTurn), + outcome, None, // turn succeeded — batch was processed, no requeue + retry_agent_index, ); return; } @@ -2092,15 +2743,18 @@ pub async fn run_prompt_task( } }; - if should_rotate { + let retirement = if should_rotate { tracing::info!( target: "pool::session", "rotating session for {source:?} after {stop_reason:?}", ); - agent.state.invalidate(&source); - } + Some(retire_session(&mut agent, &source, &session_id).await) + } else { + None + }; - let core_stop = acp_stop_to_core(&stop_reason); + let retirement_failed = matches!(retirement, Some(Err(_))); + let core_stop = completed_turn_metric_stop(&stop_reason, retirement_failed); let usage = agent.acp.take_turn_usage(); publish_agent_turn_metric( &ctx, @@ -2112,14 +2766,21 @@ pub async fn run_prompt_task( ) .await; - send_prompt_result( - &result_tx, - &turn_id, - agent, - source, - PromptOutcome::Ok(stop_reason), - None, - ); + let outcome = match retirement { + Some(Err(error)) => { + tracing::error!( + target: "pool::session", + "failed to close rotated session {session_id}" + ); + PromptOutcome::SessionCloseFailed(error) + } + Some(Ok(SessionRetirement::RecycleRequired)) => { + PromptOutcome::SessionRecycleRequired + } + Some(Ok(SessionRetirement::Closed)) => PromptOutcome::SessionRetired(stop_reason), + None => PromptOutcome::Ok(stop_reason), + }; + send_prompt_result(&result_tx, &turn_id, agent, source, outcome, None); } Err(AcpError::AgentExited) => { tracing::error!(target: "pool::prompt", "agent {} exited during prompt", agent.index); @@ -3122,7 +3783,7 @@ fn classify_control_cancel_failure( PromptOutcome::CancelDrainTimeout(CONTROL_CANCEL_GRACE), false, ), - other => (PromptOutcome::Error(other), false), + other => (PromptOutcome::CancelCleanupFailed(other), false), }; ControlCancelFailure { outcome, @@ -3390,6 +4051,17 @@ fn acp_stop_to_core(r: &StopReason) -> buzz_core::agent_turn_metric::StopReason } } +fn completed_turn_metric_stop( + stop_reason: &StopReason, + retirement_failed: bool, +) -> buzz_core::agent_turn_metric::StopReason { + if retirement_failed { + buzz_core::agent_turn_metric::StopReason::Error + } else { + acp_stop_to_core(stop_reason) + } +} + /// Best-effort: build and publish a `kind:44200` NIP-AM agent turn metric event. /// /// Does nothing when `usage` is `None` (goose emitted no usage notification @@ -3730,6 +4402,12 @@ mod tests { use nostr::{EventBuilder, Keys, Kind, Tag, Timestamp}; use serde_json::json; + const TEST_MODEL_SWITCH_REQUEST_ID: &str = "0123456789abcdef0123456789abcdef"; + + fn model_switch_request(model_id: &str) -> ModelSwitchRequest { + ModelSwitchRequest::new(model_id, TEST_MODEL_SWITCH_REQUEST_ID) + } + // These pin the initial_message dispatch path (run_prompt_task, ~line 855): // a legacy agent WITH a base_prompt must get [Base] prepended to the user // message. This is the exact regression that shipped in the round-2 bug. @@ -4206,6 +4884,7 @@ mod tests { }], cancelled_events: vec![], cancel_reason: None, + occurrence_ids: Default::default(), }; let context = ConversationContext::Thread { messages: vec![ContextMessage { @@ -4339,43 +5018,6 @@ mod tests { (s, ch_a, ch_b) } - #[test] - fn test_rotate_after_natural_completion_invalidates_channel_state() { - let (mut s, ch_a, ch_b) = make_state(); - - apply_completed_before_control_signal( - &mut s, - &PromptSource::Channel(ch_a), - &ControlSignal::Rotate, - ); - - assert!(!s.sessions.contains_key(&ch_a)); - assert!(!s.turn_counts.contains_key(&ch_a)); - assert!(!s.core_sections.contains_key(&ch_a)); - assert!(!s.has_channel_state(&ch_a)); - assert_eq!(s.sessions.get(&ch_b).unwrap(), "sess-b"); - assert_eq!(*s.turn_counts.get(&ch_b).unwrap(), 3); - assert_eq!(s.core_sections.get(&ch_b).unwrap(), "core-b"); - assert_eq!(s.heartbeat_session.as_deref(), Some("sess-hb")); - assert_eq!(s.heartbeat_turn_count, 7); - } - - #[test] - fn test_cancel_after_natural_completion_preserves_channel_state() { - let (mut s, ch_a, ch_b) = make_state(); - - apply_completed_before_control_signal( - &mut s, - &PromptSource::Channel(ch_a), - &ControlSignal::Cancel, - ); - - assert_eq!(s.sessions.get(&ch_a).unwrap(), "sess-a"); - assert_eq!(*s.turn_counts.get(&ch_a).unwrap(), 5); - assert_eq!(s.core_sections.get(&ch_a).unwrap(), "core-a"); - assert_eq!(s.sessions.get(&ch_b).unwrap(), "sess-b"); - } - #[test] fn test_invalidate_channel_clears_session_and_turn_count() { let (mut s, ch_a, ch_b) = make_state(); @@ -4490,24 +5132,101 @@ mod tests { assert_eq!(s.core_sections.get(&ch_b).unwrap(), "core-b"); } - // ── ControlSignal::SwitchModel (Phase 3a, Option ii) ───────────────────── + async fn inert_owned_agent_with_session(index: usize, channel_id: Uuid) -> OwnedAgent { + let mut state = SessionState::default(); + state + .sessions + .insert(channel_id, format!("session-{index}")); + OwnedAgent { + index, + acp: AcpClient::spawn("cat", &[], &[], false) + .await + .expect("spawn inert adapter"), + state, + model_capabilities: None, + desired_model: None, + model_overridden: false, + model_switch_request_id: None, + model_switch_rollback: None, + agent_name: "inert".into(), + goose_system_prompt_supported: None, + protocol_version: 1, + } + } + + #[tokio::test] + async fn taking_idle_channel_owner_preserves_session_until_process_recycle() { + let channel_id = Uuid::new_v4(); + let agent = inert_owned_agent_with_session(1, channel_id).await; + let mut pool = AgentPool::from_slots(vec![None, Some(agent), None]); - #[test] - fn test_switch_model_after_natural_completion_invalidates_channel_state() { - let (mut s, ch_a, ch_b) = make_state(); + let mut owners = pool.take_idle_agents_with_session(channel_id); - // SwitchModel must invalidate just like Rotate so the requeued turn - // re-creates a fresh session that re-applies the new desired_model. - apply_completed_before_control_signal( - &mut s, - &PromptSource::Channel(ch_a), - &ControlSignal::SwitchModel("gpt-5".into()), + assert_eq!(pool.live_count(), 0); + assert_eq!(owners.len(), 1); + assert_eq!( + owners[0] + .state + .sessions + .get(&channel_id) + .map(String::as_str), + Some("session-1"), + "claiming for recycle must not locally forget remote ownership" ); + owners[0] + .acp + .shutdown() + .await + .expect("idle session owner must shut down cleanly"); + } - assert!(!s.has_channel_state(&ch_a)); - // ch_b untouched — the switch is channel-scoped. - assert_eq!(s.sessions.get(&ch_b).unwrap(), "sess-b"); - assert_eq!(*s.turn_counts.get(&ch_b).unwrap(), 3); + #[tokio::test] + async fn idle_model_switch_claim_preserves_session_and_model_intent_for_recycle() { + let channel_id = Uuid::new_v4(); + let agent = inert_owned_agent_with_session(1, channel_id).await; + let mut pool = AgentPool::from_slots(vec![None, Some(agent), None]); + + let mut claimed = match pool.switch_idle_agent_model( + channel_id, + "runtime-model", + TEST_MODEL_SWITCH_REQUEST_ID, + ) { + IdleSwitchResult::Recycle(agent) => *agent, + IdleSwitchResult::UnsupportedModel => panic!("model must be accepted without catalog"), + IdleSwitchResult::NoIdleAgent => panic!("idle session owner must be found"), + }; + + assert_eq!(pool.live_count(), 0); + assert_eq!( + claimed.state.sessions.get(&channel_id).map(String::as_str), + Some("session-1") + ); + assert_eq!(claimed.desired_model.as_deref(), Some("runtime-model")); + assert!(claimed.model_overridden); + assert_eq!( + claimed.model_switch_request_id.as_deref(), + Some(TEST_MODEL_SWITCH_REQUEST_ID) + ); + claimed + .acp + .shutdown() + .await + .expect("claimed agent must shut down cleanly"); + } + + #[test] + fn idle_model_switch_without_an_owned_slot_is_a_typed_no_op() { + let mut pool = AgentPool::from_slots(vec![None]); + + assert!(matches!( + pool.switch_idle_agent_model( + Uuid::new_v4(), + "runtime-model", + TEST_MODEL_SWITCH_REQUEST_ID, + ), + IdleSwitchResult::NoIdleAgent + )); + assert_eq!(pool.live_count(), 0); } // ── requeue_cancelled_batch ──────────────────────────────────────────── @@ -4532,67 +5251,1630 @@ mod tests { }], cancelled_events: vec![], cancel_reason: None, + occurrence_ids: Default::default(), } } - #[test] - fn test_requeue_cancelled_batch_maps_control_signal_to_cancel_reason() { - let cases = [ - (ControlSignal::Steer, Some(CancelReason::Steer)), - (ControlSignal::Interrupt, Some(CancelReason::Interrupt)), - ( - ControlSignal::SwitchModel("gpt-5".into()), - Some(CancelReason::Interrupt), - ), - (ControlSignal::Cancel, None), - (ControlSignal::Rotate, None), - ]; - let mut ctx = make_prompt_context_no_owner(); - ctx.dedup_mode = DedupMode::Queue; + fn capture_path(label: &str) -> std::path::PathBuf { + std::env::temp_dir().join(format!("buzz-acp-{label}-{}.ndjson", Uuid::new_v4())) + } - for (signal, expected_reason) in cases { - let channel_id = Uuid::new_v4(); - let batch = one_event_batch(channel_id); - let result = requeue_cancelled_batch(&ctx, signal.clone(), Some(batch)); - match expected_reason { - Some(reason) => { - let batch = result - .unwrap_or_else(|| panic!("{signal:?} must preserve the batch, got None")); - assert_eq!( - batch.cancel_reason, - Some(reason), - "{signal:?} must stamp {reason:?}" - ); + async fn wait_for_capture_lines(path: &std::path::Path, minimum: usize) -> String { + tokio::time::timeout(Duration::from_secs(5), async { + loop { + if let Ok(contents) = std::fs::read_to_string(path) { + if contents.lines().count() >= minimum { + return contents; + } } - None => assert!( - result.is_none(), - "{signal:?} must drop the batch, got {result:?}" - ), + tokio::time::sleep(Duration::from_millis(10)).await; } - } + }) + .await + .unwrap_or_else(|_| { + panic!( + "timed out waiting for {minimum} captured ACP messages at {}", + path.display() + ) + }) } - // ── classify_control_cancel_failure ───────────────────────────────────── - // Table-driven pin of the single production seam used by the - // `Err(error)` arm in `run_prompt_task`'s control-cancel branch. Crosses - // the exact error→outcome AND outcome→batch-fate boundary in one call, - // so a regression to the old per-arm duplication (or to routing an - // unexpected HardTimeout back through the real hard-cap path) fails - // here rather than only in independently-manufactured unit tests. + #[tokio::test] + async fn clean_control_cancel_closes_session_before_local_invalidation() { + let capture = capture_path("cancel-close"); + let script = r#" + capture="$1" + IFS= read -r initialize + printf '%s\n' "$initialize" >> "$capture" + printf '%s\n' '{"jsonrpc":"2.0","id":0,"result":{"protocolVersion":2,"agentCapabilities":{"sessionCapabilities":{"close":{}}}}}' + IFS= read -r prompt + printf '%s\n' "$prompt" >> "$capture" + IFS= read -r cancel + printf '%s\n' "$cancel" >> "$capture" + printf '%s\n' '{"jsonrpc":"2.0","id":1,"result":{"stopReason":"cancelled"}}' + if IFS= read -r close; then + printf '%s\n' "$close" >> "$capture" + printf '%s\n' '{"jsonrpc":"2.0","id":2,"result":{}}' + fi + sleep 10 + "#; + let mut acp = AcpClient::spawn( + "bash", + &[ + "-c".to_string(), + script.to_string(), + "buzz-acp-close-test".to_string(), + capture.to_string_lossy().into_owned(), + ], + &[], + false, + ) + .await + .expect("failed to spawn fake ACP adapter"); + acp.initialize() + .await + .expect("fake adapter must advertise session/close"); + assert!(acp.supports_session_close()); - /// Assert `outcome` is the expected `PromptOutcome` variant. `PromptOutcome` - /// has no `PartialEq` (it wraps `AcpError`, which isn't `PartialEq`), so - /// this matches by shape instead of deriving equality onto the whole enum. - fn assert_outcome_matches(outcome: &PromptOutcome, expected: &str) { - let label = match outcome { - PromptOutcome::AgentExited => "AgentExited", - PromptOutcome::Timeout(TimeoutKind::Idle) => "Timeout(Idle)", - PromptOutcome::Timeout(TimeoutKind::Hard { .. }) => "Timeout(Hard)", - PromptOutcome::CancelDrainTimeout(_) => "CancelDrainTimeout", - PromptOutcome::Error(_) => "Error", - PromptOutcome::Cancelled => "Cancelled", - PromptOutcome::Ok(_) => "Ok", - }; + let channel_id = Uuid::new_v4(); + let mut state = SessionState::default(); + state.sessions.insert(channel_id, "session-old".into()); + let agent = OwnedAgent { + index: 0, + acp, + state, + model_capabilities: None, + desired_model: None, + model_overridden: false, + model_switch_request_id: None, + model_switch_rollback: None, + agent_name: "claude-code-acp".into(), + goose_system_prompt_supported: None, + protocol_version: 2, + }; + let mut ctx = make_prompt_context_no_owner(); + ctx.dedup_mode = DedupMode::Queue; + let (result_tx, mut result_rx) = mpsc::unbounded_channel(); + let (control_tx, control_rx) = tokio::sync::oneshot::channel(); + let task = tokio::spawn(run_prompt_task( + agent, + Some(one_event_batch(channel_id)), + Some("prompt".into()), + Arc::new(ctx), + result_tx, + Some(control_rx), + "turn-cancel-close".into(), + )); + + let _ = wait_for_capture_lines(&capture, 2).await; + control_tx + .send(ControlSignal::Steer) + .expect("control receiver must still be live"); + let result = tokio::time::timeout(Duration::from_secs(5), result_rx.recv()) + .await + .expect("prompt task must return") + .expect("prompt result channel must stay open"); + task.await.expect("prompt task must not panic"); + + assert!( + matches!(result.outcome, PromptOutcome::Cancelled), + "clean cancel must remain a non-error outcome" + ); + assert!( + !result.agent.state.sessions.contains_key(&channel_id), + "local session state must be invalidated after the remote close succeeds" + ); + assert_eq!( + result.batch.as_ref().and_then(|batch| batch.cancel_reason), + Some(CancelReason::Steer), + "queue-mode steer must preserve the cancelled batch" + ); + + let captured = wait_for_capture_lines(&capture, 4).await; + let messages: Vec = captured + .lines() + .map(|line| serde_json::from_str(line).expect("captured line must be JSON")) + .collect(); + assert_eq!(messages[0]["method"], "initialize"); + assert_eq!(messages[1]["method"], "session/prompt"); + assert_eq!(messages[2]["method"], "session/cancel"); + assert_eq!( + messages[3]["method"], "session/close", + "the retired ACP session must be closed before the adapter is reused" + ); + assert_eq!(messages[3]["params"]["sessionId"], "session-old"); + + let _ = std::fs::remove_file(capture); + } + + #[tokio::test] + async fn shutdown_preempts_blocked_session_setup_and_returns_agent_owner() { + let capture = capture_path("shutdown-session-setup"); + let script = r#" + capture="$1" + IFS= read -r initialize + printf '%s\n' "$initialize" >> "$capture" + printf '%s\n' '{"jsonrpc":"2.0","id":0,"result":{"protocolVersion":2,"agentCapabilities":{"sessionCapabilities":{}}}}' + IFS= read -r new_session + printf '%s\n' "$new_session" >> "$capture" + sleep 10 + "#; + let mut acp = AcpClient::spawn( + "bash", + &[ + "-c".to_string(), + script.to_string(), + "buzz-acp-shutdown-session-setup-test".to_string(), + capture.to_string_lossy().into_owned(), + ], + &[], + false, + ) + .await + .expect("failed to spawn fake ACP adapter"); + acp.initialize() + .await + .expect("fake adapter must initialize"); + + let agent = OwnedAgent { + index: 0, + acp, + state: SessionState::default(), + model_capabilities: None, + desired_model: None, + model_overridden: false, + model_switch_request_id: None, + model_switch_rollback: None, + agent_name: "claude-code-acp".into(), + goose_system_prompt_supported: None, + protocol_version: 2, + }; + let mut ctx = make_prompt_context_no_owner(); + ctx.dedup_mode = DedupMode::Queue; + let (result_tx, mut result_rx) = mpsc::unbounded_channel(); + let (control_tx, control_rx) = tokio::sync::oneshot::channel(); + let task = tokio::spawn(run_prompt_task( + agent, + None, + Some("prompt".into()), + Arc::new(ctx), + result_tx, + Some(control_rx), + "turn-shutdown-session-setup".into(), + )); + + let _ = wait_for_capture_lines(&capture, 1).await; + control_tx + .send(ControlSignal::Cancel) + .expect("setup task must retain its shutdown receiver"); + let result = tokio::time::timeout(Duration::from_secs(1), result_rx.recv()) + .await + .expect("shutdown must not wait for blocked session setup") + .expect("typed agent owner must be returned"); + task.await.expect("setup task must not panic"); + + assert!(matches!( + result.outcome, + PromptOutcome::ControlPreemptedSetup + )); + assert!(result.batch.is_none(), "shutdown discards the active batch"); + assert_eq!(result.agent.index, 0, "typed agent owner must be recovered"); + + let _ = std::fs::remove_file(capture); + } + + #[tokio::test] + async fn shutdown_preempts_initial_message_and_closes_partial_session() { + let capture = capture_path("shutdown-initial-message"); + let script = r#" + capture="$1" + IFS= read -r initialize + printf '%s\n' "$initialize" >> "$capture" + printf '%s\n' '{"jsonrpc":"2.0","id":0,"result":{"protocolVersion":2,"agentCapabilities":{"sessionCapabilities":{"close":{}}}}}' + IFS= read -r new_session + printf '%s\n' "$new_session" >> "$capture" + printf '%s\n' '{"jsonrpc":"2.0","id":1,"result":{"sessionId":"session-created"}}' + IFS= read -r initial_prompt + printf '%s\n' "$initial_prompt" >> "$capture" + IFS= read -r cancel + printf '%s\n' "$cancel" >> "$capture" + printf '%s\n' '{"jsonrpc":"2.0","id":2,"result":{"stopReason":"cancelled"}}' + IFS= read -r close + printf '%s\n' "$close" >> "$capture" + printf '%s\n' '{"jsonrpc":"2.0","id":3,"result":{}}' + sleep 10 + "#; + let mut acp = AcpClient::spawn( + "bash", + &[ + "-c".to_string(), + script.to_string(), + "buzz-acp-shutdown-initial-message-test".to_string(), + capture.to_string_lossy().into_owned(), + ], + &[], + false, + ) + .await + .expect("failed to spawn fake ACP adapter"); + acp.initialize() + .await + .expect("fake adapter must initialize"); + + let channel_id = Uuid::new_v4(); + let agent = OwnedAgent { + index: 0, + acp, + state: SessionState::default(), + model_capabilities: None, + desired_model: None, + model_overridden: false, + model_switch_request_id: None, + model_switch_rollback: None, + agent_name: "claude-code-acp".into(), + goose_system_prompt_supported: None, + protocol_version: 2, + }; + let mut ctx = make_prompt_context_no_owner(); + ctx.initial_message = Some("initialize".into()); + ctx.channel_info = ChannelInfoResolver::new( + std::collections::HashMap::from([( + channel_id, + crate::relay::ChannelInfo { + name: "private".into(), + channel_type: "dm".into(), + }, + )]), + ctx.rest_client.clone(), + ); + let (result_tx, mut result_rx) = mpsc::unbounded_channel(); + let (control_tx, control_rx) = tokio::sync::oneshot::channel(); + let task = tokio::spawn(run_prompt_task( + agent, + Some(one_event_batch(channel_id)), + Some("prompt".into()), + Arc::new(ctx), + result_tx, + Some(control_rx), + "turn-shutdown-initial-message".into(), + )); + + let _ = wait_for_capture_lines(&capture, 3).await; + control_tx + .send(ControlSignal::Cancel) + .expect("initial-message task must retain its shutdown receiver"); + let result = tokio::time::timeout(Duration::from_secs(1), result_rx.recv()) + .await + .expect("shutdown must preempt the initial message") + .expect("typed agent owner must be returned"); + task.await.expect("initial-message task must not panic"); + + assert!(matches!(result.outcome, PromptOutcome::Cancelled)); + assert!(!result.agent.state.sessions.contains_key(&channel_id)); + let captured = wait_for_capture_lines(&capture, 5).await; + let messages: Vec = captured + .lines() + .map(|line| serde_json::from_str(line).expect("captured line must be JSON")) + .collect(); + assert_eq!(messages[3]["method"], "session/cancel"); + assert_eq!(messages[4]["method"], "session/close"); + + let _ = std::fs::remove_file(capture); + } + + #[tokio::test] + async fn control_before_initial_prompt_poll_closes_newly_created_idle_session() { + let capture = capture_path("pre-initial-poll-control"); + let script = r#" + capture="$1" + IFS= read -r initialize + printf '%s\n' "$initialize" >> "$capture" + printf '%s\n' '{"jsonrpc":"2.0","id":0,"result":{"protocolVersion":2,"agentCapabilities":{"sessionCapabilities":{"close":{}}}}}' + IFS= read -r close + printf '%s\n' "$close" >> "$capture" + printf '%s\n' '{"jsonrpc":"2.0","id":1,"result":{}}' + sleep 10 + "#; + let mut acp = AcpClient::spawn( + "bash", + &[ + "-c".to_string(), + script.to_string(), + "buzz-acp-pre-initial-poll-control-test".to_string(), + capture.to_string_lossy().into_owned(), + ], + &[], + false, + ) + .await + .expect("failed to spawn fake ACP adapter"); + acp.initialize() + .await + .expect("fake adapter must initialize"); + + let channel_id = Uuid::new_v4(); + let mut state = SessionState::default(); + state.sessions.insert(channel_id, "session-created".into()); + let agent = OwnedAgent { + index: 0, + acp, + state, + model_capabilities: None, + desired_model: None, + model_overridden: false, + model_switch_request_id: None, + model_switch_rollback: None, + agent_name: "claude-code-acp".into(), + goose_system_prompt_supported: None, + protocol_version: 2, + }; + assert!( + !agent.acp.has_in_flight_prompt(), + "fixture must reproduce the ready-control race before initial prompt polling" + ); + + let (result_tx, mut result_rx) = mpsc::unbounded_channel(); + send_pre_prompt_control_result( + &result_tx, + "turn-pre-initial-poll", + agent, + PromptSource::Channel(channel_id), + ControlSignal::Steer, + &make_prompt_context_no_owner(), + None, + Some("session-created".into()), + false, + true, + ) + .await; + + let result = result_rx.recv().await.expect("typed owner result"); + assert!(matches!(result.outcome, PromptOutcome::Cancelled)); + assert!( + !result.agent.state.sessions.contains_key(&channel_id), + "a control-ready partial session must not be reused without its initial message" + ); + let captured = wait_for_capture_lines(&capture, 2).await; + let messages: Vec = captured + .lines() + .map(|line| serde_json::from_str(line).expect("captured line must be JSON")) + .collect(); + assert_eq!(messages[1]["method"], "session/close"); + + let _ = std::fs::remove_file(capture); + } + + #[tokio::test] + async fn rejected_busy_model_switch_restores_prior_intent_and_never_retries_rejected_model() { + let capture = capture_path("busy-model-rejection"); + let script = r#" + capture="$1" + IFS= read -r initialize + printf '%s\n' "$initialize" >> "$capture" + printf '%s\n' '{"jsonrpc":"2.0","id":0,"result":{"protocolVersion":2,"agentCapabilities":{"sessionCapabilities":{"close":{}}}}}' + IFS= read -r first_prompt + printf '%s\n' "$first_prompt" >> "$capture" + IFS= read -r cancel + printf '%s\n' "$cancel" >> "$capture" + printf '%s\n' '{"jsonrpc":"2.0","id":1,"result":{"stopReason":"cancelled"}}' + IFS= read -r close + printf '%s\n' "$close" >> "$capture" + printf '%s\n' '{"jsonrpc":"2.0","id":2,"result":{}}' + IFS= read -r new_session + printf '%s\n' "$new_session" >> "$capture" + printf '%s\n' '{"jsonrpc":"2.0","id":3,"result":{"sessionId":"session-new","configOptions":[{"configId":"model","category":"model","options":[{"value":"model-a"}]}]}}' + IFS= read -r set_model + printf '%s\n' "$set_model" >> "$capture" + printf '%s\n' '{"jsonrpc":"2.0","id":4,"result":{}}' + IFS= read -r second_prompt + printf '%s\n' "$second_prompt" >> "$capture" + printf '%s\n' '{"jsonrpc":"2.0","id":5,"result":{"stopReason":"end_turn"}}' + sleep 10 + "#; + let mut acp = AcpClient::spawn( + "bash", + &[ + "-c".to_string(), + script.to_string(), + "buzz-acp-busy-model-rejection-test".to_string(), + capture.to_string_lossy().into_owned(), + ], + &[], + false, + ) + .await + .expect("failed to spawn fake ACP adapter"); + let observer = observer::ObserverHandle::in_process(); + acp.set_observer(Some(observer.clone()), 0); + acp.initialize() + .await + .expect("fake adapter must advertise session/close"); + + let channel_id = Uuid::new_v4(); + let mut state = SessionState::default(); + state.sessions.insert(channel_id, "session-old".into()); + let agent = OwnedAgent { + index: 0, + acp, + state, + model_capabilities: Some(AgentModelCapabilities { + config_options_raw: vec![serde_json::json!({ + "configId": "model", + "category": "model", + "options": [{"value": "model-a"}], + })], + available_models_raw: None, + }), + desired_model: Some("model-a".into()), + model_overridden: true, + model_switch_request_id: None, + model_switch_rollback: None, + agent_name: "claude-code-acp".into(), + goose_system_prompt_supported: None, + protocol_version: 2, + }; + let mut ctx = make_prompt_context_no_owner(); + ctx.dedup_mode = DedupMode::Queue; + let ctx = Arc::new(ctx); + let (first_tx, mut first_rx) = mpsc::unbounded_channel(); + let (control_tx, control_rx) = tokio::sync::oneshot::channel(); + let first_task = tokio::spawn(run_prompt_task( + agent, + Some(one_event_batch(channel_id)), + Some("first prompt".into()), + Arc::clone(&ctx), + first_tx, + Some(control_rx), + "turn-model-rejection-first".into(), + )); + + let _ = wait_for_capture_lines(&capture, 1).await; + control_tx + .send(ControlSignal::SwitchModel(model_switch_request("model-b"))) + .expect("busy task must accept model control"); + let first = tokio::time::timeout(Duration::from_secs(5), first_rx.recv()) + .await + .expect("busy task must return") + .expect("result channel must stay open"); + first_task.await.expect("busy task must not panic"); + + assert_eq!( + first.agent.desired_model.as_deref(), + Some("model-a"), + "rejecting model B must restore the exact prior desired model" + ); + assert!( + first.agent.model_overridden, + "rejecting model B must restore the prior override provenance" + ); + assert_eq!( + first.agent.model_switch_request_id, None, + "rejecting model B must restore the prior correlation state" + ); + + let retry_batch = first + .batch + .expect("busy switch must requeue the interrupted turn"); + let (second_tx, mut second_rx) = mpsc::unbounded_channel(); + run_prompt_task( + first.agent, + Some(retry_batch), + Some("second prompt".into()), + Arc::clone(&ctx), + second_tx, + None, + "turn-model-rejection-second".into(), + ) + .await; + let second = second_rx.recv().await.expect("retry result must be sent"); + assert!(matches!( + second.outcome, + PromptOutcome::Ok(StopReason::EndTurn) + )); + assert_eq!(second.agent.desired_model.as_deref(), Some("model-a")); + assert!(second.agent.model_overridden); + + let captured = wait_for_capture_lines(&capture, 7).await; + let messages: Vec = captured + .lines() + .map(|line| serde_json::from_str(line).expect("captured line must be JSON")) + .collect(); + let model_requests: Vec<&serde_json::Value> = messages + .iter() + .filter(|message| message["method"] == "session/set_config_option") + .collect(); + assert_eq!(model_requests.len(), 1); + assert_eq!(model_requests[0]["params"]["value"], "model-a"); + assert!( + !captured.contains("model-b"), + "the rejected model must never reach a future ACP session" + ); + let control_results: Vec<_> = observer + .snapshot() + .into_iter() + .filter(|event| event.kind == "control_result") + .collect(); + let channel_id_string = channel_id.to_string(); + assert!(control_results.iter().any(|event| { + event.channel_id.as_deref() == Some(channel_id_string.as_str()) + && event.payload["status"] == "unsupported_model" + && event.payload["modelId"] == "model-b" + && event.payload["requestId"] == TEST_MODEL_SWITCH_REQUEST_ID + })); + assert_eq!( + control_results.len(), + 1, + "a configured or legacy desired model without requestId must not emit an ambiguous result" + ); + + let _ = std::fs::remove_file(capture); + } + + #[tokio::test] + async fn unadvertised_session_close_recycles_without_sending_optional_rpc() { + let capture = capture_path("rotation-recycle"); + let script = r#" + capture="$1" + IFS= read -r initialize + printf '%s\n' "$initialize" >> "$capture" + printf '%s\n' '{"jsonrpc":"2.0","id":0,"result":{"protocolVersion":1,"agentCapabilities":{"sessionCapabilities":{}}}}' + IFS= read -r first_prompt + printf '%s\n' "$first_prompt" >> "$capture" + printf '%s\n' '{"jsonrpc":"2.0","id":1,"result":{"stopReason":"max_tokens"}}' + if IFS= read -r -t 2 unexpected; then + printf '%s\n' "$unexpected" >> "$capture" + printf '%s\n' '{"jsonrpc":"2.0","id":2,"error":{"code":-32601,"message":"Method not found"}}' + fi + sleep 10 + "#; + let mut acp = AcpClient::spawn( + "bash", + &[ + "-c".to_string(), + script.to_string(), + "buzz-acp-recycle-test".to_string(), + capture.to_string_lossy().into_owned(), + ], + &[], + false, + ) + .await + .expect("failed to spawn fake ACP adapter"); + acp.initialize() + .await + .expect("fake adapter must complete capability negotiation"); + assert!(!acp.supports_session_close()); + + let channel_id = Uuid::new_v4(); + let mut state = SessionState::default(); + state.sessions.insert(channel_id, "session-old".into()); + let agent = OwnedAgent { + index: 0, + acp, + state, + model_capabilities: None, + desired_model: None, + model_overridden: false, + model_switch_request_id: None, + model_switch_rollback: None, + agent_name: "legacy-acp".into(), + goose_system_prompt_supported: None, + protocol_version: 1, + }; + let (result_tx, mut result_rx) = mpsc::unbounded_channel(); + + run_prompt_task( + agent, + Some(one_event_batch(channel_id)), + Some("first prompt".into()), + Arc::new(make_prompt_context_no_owner()), + result_tx, + None, + "turn-rotation-recycle".into(), + ) + .await; + let result = result_rx + .recv() + .await + .expect("rotation result must be sent"); + + assert!( + matches!(result.outcome, PromptOutcome::SessionRecycleRequired), + "an adapter without advertised close support must be deliberately recycled" + ); + assert_eq!( + result + .agent + .state + .sessions + .get(&channel_id) + .map(String::as_str), + Some("session-old"), + "local ownership must remain until process-group replacement" + ); + let captured = wait_for_capture_lines(&capture, 2).await; + let messages: Vec = captured + .lines() + .map(|line| serde_json::from_str(line).expect("captured line must be JSON")) + .collect(); + assert_eq!(messages[0]["method"], "initialize"); + assert_eq!(messages[1]["method"], "session/prompt"); + assert_eq!( + captured.lines().count(), + 2, + "Buzz must not probe an optional method the adapter did not advertise" + ); + + let _ = std::fs::remove_file(capture); + } + + #[tokio::test] + async fn token_rotation_closes_old_session_before_next_session_new() { + let capture = capture_path("rotation-close"); + let script = r#" + capture="$1" + IFS= read -r initialize + printf '%s\n' "$initialize" >> "$capture" + printf '%s\n' '{"jsonrpc":"2.0","id":0,"result":{"protocolVersion":2,"agentCapabilities":{"sessionCapabilities":{"close":{}}}}}' + IFS= read -r first_prompt + printf '%s\n' "$first_prompt" >> "$capture" + printf '%s\n' '{"jsonrpc":"2.0","id":1,"result":{"stopReason":"max_tokens"}}' + IFS= read -r close + printf '%s\n' "$close" >> "$capture" + printf '%s\n' '{"jsonrpc":"2.0","id":2,"result":{}}' + IFS= read -r new_session + printf '%s\n' "$new_session" >> "$capture" + printf '%s\n' '{"jsonrpc":"2.0","id":3,"result":{"sessionId":"session-new"}}' + IFS= read -r second_prompt + printf '%s\n' "$second_prompt" >> "$capture" + printf '%s\n' '{"jsonrpc":"2.0","id":4,"result":{"stopReason":"end_turn"}}' + sleep 10 + "#; + let mut acp = AcpClient::spawn( + "bash", + &[ + "-c".to_string(), + script.to_string(), + "buzz-acp-rotation-test".to_string(), + capture.to_string_lossy().into_owned(), + ], + &[], + false, + ) + .await + .expect("failed to spawn fake ACP adapter"); + acp.initialize() + .await + .expect("fake adapter must advertise session/close"); + assert!(acp.supports_session_close()); + + let channel_id = Uuid::new_v4(); + let mut state = SessionState::default(); + state.sessions.insert(channel_id, "session-old".into()); + let agent = OwnedAgent { + index: 0, + acp, + state, + model_capabilities: None, + desired_model: None, + model_overridden: false, + model_switch_request_id: None, + model_switch_rollback: None, + agent_name: "claude-code-acp".into(), + goose_system_prompt_supported: None, + protocol_version: 2, + }; + let mut ctx = make_prompt_context_no_owner(); + ctx.channel_info = ChannelInfoResolver::new( + std::collections::HashMap::from([( + channel_id, + crate::relay::ChannelInfo { + name: "test".into(), + channel_type: "public".into(), + }, + )]), + ctx.rest_client.clone(), + ); + let ctx = Arc::new(ctx); + + let (first_tx, mut first_rx) = mpsc::unbounded_channel(); + run_prompt_task( + agent, + Some(one_event_batch(channel_id)), + Some("first prompt".into()), + Arc::clone(&ctx), + first_tx, + None, + "turn-rotation-first".into(), + ) + .await; + let first = first_rx.recv().await.expect("first result must be sent"); + assert!( + matches!( + first.outcome, + PromptOutcome::SessionRetired(StopReason::MaxTokens) + ), + "the completed turn must record positive remote session retirement" + ); + assert!( + !first.agent.state.sessions.contains_key(&channel_id), + "the rotated session must be absent after close succeeds" + ); + + let (second_tx, mut second_rx) = mpsc::unbounded_channel(); + run_prompt_task( + first.agent, + Some(one_event_batch(channel_id)), + Some("second prompt".into()), + Arc::clone(&ctx), + second_tx, + None, + "turn-rotation-second".into(), + ) + .await; + let second = second_rx.recv().await.expect("second result must be sent"); + assert!( + matches!(second.outcome, PromptOutcome::Ok(StopReason::EndTurn)), + "the replacement session must run the next turn successfully" + ); + assert_eq!( + second + .agent + .state + .sessions + .get(&channel_id) + .map(String::as_str), + Some("session-new") + ); + + let captured = wait_for_capture_lines(&capture, 5).await; + let messages: Vec = captured + .lines() + .map(|line| serde_json::from_str(line).expect("captured line must be JSON")) + .collect(); + assert_eq!(messages[0]["method"], "initialize"); + assert_eq!(messages[1]["method"], "session/prompt"); + assert_eq!( + messages[2]["method"], "session/close", + "the old session must close as part of rotation" + ); + assert_eq!(messages[2]["params"]["sessionId"], "session-old"); + assert_eq!( + messages[3]["method"], "session/new", + "a replacement session may be created only after close succeeds" + ); + assert_eq!(messages[4]["method"], "session/prompt"); + + let _ = std::fs::remove_file(capture); + } + + #[tokio::test] + async fn failed_goose_session_setup_closes_created_remote_session_before_reuse() { + let capture = capture_path("goose-setup-close"); + let script = r#" + capture="$1" + IFS= read -r initialize + printf '%s\n' "$initialize" >> "$capture" + printf '%s\n' '{"jsonrpc":"2.0","id":0,"result":{"protocolVersion":2,"agentCapabilities":{"sessionCapabilities":{"close":{}}}}}' + IFS= read -r new_session + printf '%s\n' "$new_session" >> "$capture" + printf '%s\n' '{"jsonrpc":"2.0","id":1,"result":{"sessionId":"session-created"}}' + IFS= read -r system_prompt + printf '%s\n' "$system_prompt" >> "$capture" + printf '%s\n' '{"jsonrpc":"2.0","id":2,"error":{"code":-32042,"message":"setup rejected"}}' + if IFS= read -r close; then + printf '%s\n' "$close" >> "$capture" + printf '%s\n' '{"jsonrpc":"2.0","id":3,"result":{}}' + fi + sleep 10 + "#; + let mut acp = AcpClient::spawn( + "bash", + &[ + "-c".to_string(), + script.to_string(), + "buzz-acp-goose-setup-close-test".to_string(), + capture.to_string_lossy().into_owned(), + ], + &[], + false, + ) + .await + .expect("failed to spawn fake ACP adapter"); + acp.initialize() + .await + .expect("fake adapter must advertise session/close"); + + let channel_id = Uuid::new_v4(); + let agent = OwnedAgent { + index: 0, + acp, + state: SessionState::default(), + model_capabilities: None, + desired_model: None, + model_overridden: false, + model_switch_request_id: None, + model_switch_rollback: None, + agent_name: "goose".into(), + goose_system_prompt_supported: None, + protocol_version: 2, + }; + let mut ctx = make_prompt_context_no_owner(); + ctx.base_prompt = Some("setup prompt"); + ctx.dedup_mode = DedupMode::Queue; + let (result_tx, mut result_rx) = mpsc::unbounded_channel(); + + run_prompt_task( + agent, + Some(one_event_batch(channel_id)), + Some("prompt".into()), + Arc::new(ctx), + result_tx, + None, + "turn-goose-setup-close".into(), + ) + .await; + let result = result_rx + .recv() + .await + .expect("setup failure result must be sent"); + + assert!(matches!( + result.outcome, + PromptOutcome::Error(AcpError::AgentError { code: -32042, .. }) + )); + assert!( + !result.agent.state.sessions.contains_key(&channel_id), + "a failed setup must never publish local ownership for the partial session" + ); + let captured = wait_for_capture_lines(&capture, 4).await; + let messages: Vec = captured + .lines() + .map(|line| serde_json::from_str(line).expect("captured line must be JSON")) + .collect(); + assert_eq!(messages[0]["method"], "initialize"); + assert_eq!(messages[1]["method"], "session/new"); + assert_eq!( + messages[2]["method"], + "_goose/unstable/session/system-prompt/set" + ); + assert_eq!(messages[3]["method"], "session/close"); + assert_eq!(messages[3]["params"]["sessionId"], "session-created"); + + let _ = std::fs::remove_file(capture); + } + + #[tokio::test] + async fn fresh_session_switch_success_echoes_request_id_once() { + let capture = capture_path("model-switch-request-id-success"); + let script = r#" + capture="$1" + IFS= read -r initialize + printf '%s\n' "$initialize" >> "$capture" + printf '%s\n' '{"jsonrpc":"2.0","id":0,"result":{"protocolVersion":2,"agentCapabilities":{"sessionCapabilities":{"close":{}}}}}' + IFS= read -r new_session + printf '%s\n' "$new_session" >> "$capture" + printf '%s\n' '{"jsonrpc":"2.0","id":1,"result":{"sessionId":"session-created","configOptions":[{"configId":"model","category":"model","options":[{"value":"model-b"}]}]}}' + IFS= read -r set_model + printf '%s\n' "$set_model" >> "$capture" + printf '%s\n' '{"jsonrpc":"2.0","id":2,"result":{}}' + sleep 10 + "#; + let mut acp = AcpClient::spawn( + "bash", + &[ + "-c".to_string(), + script.to_string(), + "buzz-acp-model-switch-request-id-success-test".to_string(), + capture.to_string_lossy().into_owned(), + ], + &[], + false, + ) + .await + .expect("failed to spawn fake ACP adapter"); + let observer = observer::ObserverHandle::in_process(); + acp.set_observer(Some(observer.clone()), 0); + acp.initialize() + .await + .expect("fake adapter must initialize"); + + let mut agent = OwnedAgent { + index: 0, + acp, + state: SessionState::default(), + model_capabilities: None, + desired_model: Some("model-b".into()), + model_overridden: true, + model_switch_request_id: Some(TEST_MODEL_SWITCH_REQUEST_ID.into()), + model_switch_rollback: None, + agent_name: "claude-code-acp".into(), + goose_system_prompt_supported: None, + protocol_version: 2, + }; + let ctx = make_prompt_context_no_owner(); + + let prepared = + match create_session_and_apply_model(&mut agent, &ctx, None, None, None).await { + Ok(prepared) => prepared, + Err(_) => panic!("fresh session model switch must succeed"), + }; + assert_eq!(prepared.session_id, "session-created"); + assert_eq!( + agent.model_switch_request_id.as_deref(), + Some(TEST_MODEL_SWITCH_REQUEST_ID), + "session setup alone must not consume the correlation ID" + ); + + assert!( + observer + .snapshot() + .iter() + .all(|event| event.kind != "control_result"), + "session setup alone must not emit terminal switch success" + ); + + let channel_id = Uuid::new_v4(); + let source = PromptSource::Channel(channel_id); + assert!( + !commit_prepared_session(&mut agent, &source, &prepared), + "terminal success must fail closed before local session insertion" + ); + assert_eq!( + agent.model_switch_request_id.as_deref(), + Some(TEST_MODEL_SWITCH_REQUEST_ID) + ); + agent + .state + .sessions + .insert(channel_id, prepared.session_id.clone()); + assert!(commit_prepared_session(&mut agent, &source, &prepared)); + + let results: Vec<_> = observer + .snapshot() + .into_iter() + .filter(|event| event.kind == "control_result") + .collect(); + assert_eq!(results.len(), 1); + assert_eq!(results[0].payload["status"], "switched"); + assert_eq!(results[0].payload["modelId"], "model-b"); + assert_eq!( + results[0].payload["requestId"], + TEST_MODEL_SWITCH_REQUEST_ID + ); + assert_eq!( + agent.model_switch_request_id, None, + "committed terminal success must consume the correlation ID" + ); + + agent + .acp + .shutdown() + .await + .expect("fake adapter must shut down cleanly"); + let _ = std::fs::remove_file(capture); + } + + #[tokio::test] + async fn permission_setup_failure_keeps_switch_pending_without_terminal_success() { + let capture = capture_path("model-switch-permission-failure"); + let script = r#" + capture="$1" + IFS= read -r initialize + printf '%s\n' "$initialize" >> "$capture" + printf '%s\n' '{"jsonrpc":"2.0","id":0,"result":{"protocolVersion":2,"agentCapabilities":{"sessionCapabilities":{"close":{}}}}}' + IFS= read -r new_session + printf '%s\n' "$new_session" >> "$capture" + printf '%s\n' '{"jsonrpc":"2.0","id":1,"result":{"sessionId":"session-created","configOptions":[{"configId":"model","category":"model","options":[{"value":"model-b"}]}],"modes":{"availableModes":[{"id":"bypassPermissions"}]}}}' + IFS= read -r set_model + printf '%s\n' "$set_model" >> "$capture" + printf '%s\n' '{"jsonrpc":"2.0","id":2,"result":{}}' + IFS= read -r set_mode + printf '%s\n' "$set_mode" >> "$capture" + exit 0 + "#; + let mut acp = AcpClient::spawn( + "bash", + &[ + "-c".to_string(), + script.to_string(), + "buzz-acp-model-switch-permission-failure-test".to_string(), + capture.to_string_lossy().into_owned(), + ], + &[], + false, + ) + .await + .expect("failed to spawn fake ACP adapter"); + let observer = observer::ObserverHandle::in_process(); + acp.set_observer(Some(observer.clone()), 0); + acp.initialize() + .await + .expect("fake adapter must initialize"); + + let mut agent = OwnedAgent { + index: 0, + acp, + state: SessionState::default(), + model_capabilities: None, + desired_model: Some("model-a".into()), + model_overridden: false, + model_switch_request_id: None, + model_switch_rollback: None, + agent_name: "claude-code-acp".into(), + goose_system_prompt_supported: None, + protocol_version: 2, + }; + assert!(stage_busy_model_switch_intent( + &mut agent, + &model_switch_request("model-b") + )); + let mut ctx = make_prompt_context_no_owner(); + ctx.permission_mode = PermissionMode::BypassPermissions; + + let result = create_session_and_apply_model(&mut agent, &ctx, None, None, None).await; + assert!(matches!( + result, + Err(SessionSetupError::Agent(AcpError::AgentExited)) + )); + assert!( + observer + .snapshot() + .iter() + .all(|event| event.kind != "control_result"), + "later setup failure must not leave a false terminal switched result" + ); + assert_eq!(agent.desired_model.as_deref(), Some("model-b")); + assert!(agent.model_overridden); + assert_eq!( + agent.model_switch_request_id.as_deref(), + Some(TEST_MODEL_SWITCH_REQUEST_ID), + "later setup failure must preserve correlation for replacement/retry" + ); + let rollback = agent + .model_switch_rollback + .as_ref() + .expect("pending switch must preserve prior intent"); + assert_eq!(rollback.desired_model.as_deref(), Some("model-a")); + assert!(!rollback.model_overridden); + + let captured = wait_for_capture_lines(&capture, 4).await; + let messages: Vec = captured + .lines() + .map(|line| serde_json::from_str(line).expect("captured line must be JSON")) + .collect(); + assert_eq!(messages[2]["method"], "session/set_config_option"); + assert_eq!(messages[2]["params"]["configId"], "model"); + assert_eq!(messages[3]["method"], "session/set_config_option"); + assert_eq!(messages[3]["params"]["configId"], "mode"); + let _ = std::fs::remove_file(capture); + } + + #[tokio::test] + async fn fresh_catalog_rejection_reapplies_restored_pending_model_before_session_commit() { + const PRIOR_REQUEST_ID: &str = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"; + + let capture = capture_path("model-switch-fresh-catalog-rollback"); + let script = r#" + capture="$1" + IFS= read -r initialize + printf '%s\n' "$initialize" >> "$capture" + printf '%s\n' '{"jsonrpc":"2.0","id":0,"result":{"protocolVersion":2,"agentCapabilities":{"sessionCapabilities":{"close":{}}}}}' + IFS= read -r new_session + printf '%s\n' "$new_session" >> "$capture" + printf '%s\n' '{"jsonrpc":"2.0","id":1,"result":{"sessionId":"session-created","configOptions":[{"configId":"model","category":"model","options":[{"value":"model-a"}]}]}}' + IFS= read -r set_model + printf '%s\n' "$set_model" >> "$capture" + printf '%s\n' '{"jsonrpc":"2.0","id":2,"result":{}}' + sleep 10 + "#; + let mut acp = AcpClient::spawn( + "bash", + &[ + "-c".to_string(), + script.to_string(), + "buzz-acp-model-switch-fresh-catalog-rollback-test".to_string(), + capture.to_string_lossy().into_owned(), + ], + &[], + false, + ) + .await + .expect("failed to spawn fake ACP adapter"); + let observer = observer::ObserverHandle::in_process(); + acp.set_observer(Some(observer.clone()), 0); + acp.initialize() + .await + .expect("fake adapter must initialize"); + + let mut agent = OwnedAgent { + index: 0, + acp, + state: SessionState::default(), + model_capabilities: None, + desired_model: Some("model-a".into()), + model_overridden: true, + model_switch_request_id: Some(PRIOR_REQUEST_ID.into()), + model_switch_rollback: Some(Box::new(ModelSwitchRollback { + desired_model: Some("configured-default".into()), + model_overridden: false, + request_id: None, + previous: None, + })), + agent_name: "claude-code-acp".into(), + goose_system_prompt_supported: None, + protocol_version: 2, + }; + assert!(stage_busy_model_switch_intent( + &mut agent, + &model_switch_request("model-b") + )); + + let ctx = make_prompt_context_no_owner(); + let prepared = + match create_session_and_apply_model(&mut agent, &ctx, None, None, None).await { + Ok(prepared) => prepared, + Err(_) => panic!("restored pending model must apply on the fresh session"), + }; + + let terminal = prepared + .terminal_model_switch + .as_ref() + .expect("the restored pending request must remain terminally correlated"); + assert_eq!(terminal.model_id, "model-a"); + assert_eq!(terminal.request_id, PRIOR_REQUEST_ID); + assert_eq!( + prepared.session_config["modelOverridden"], + serde_json::Value::Bool(true), + "session config must describe the model actually restored on this session" + ); + assert_eq!(agent.desired_model.as_deref(), Some("model-a")); + assert_eq!( + agent.model_switch_request_id.as_deref(), + Some(PRIOR_REQUEST_ID) + ); + + let captured = wait_for_capture_lines(&capture, 3).await; + let messages: Vec = captured + .lines() + .map(|line| serde_json::from_str(line).expect("captured line must be JSON")) + .collect(); + assert_eq!(messages[2]["method"], "session/set_config_option"); + assert_eq!(messages[2]["params"]["value"], "model-a"); + + let before_commit: Vec<_> = observer + .snapshot() + .into_iter() + .filter(|event| event.kind == "control_result") + .collect(); + assert!(before_commit.iter().any(|event| { + event.payload["status"] == "unsupported_model" + && event.payload["modelId"] == "model-b" + && event.payload["requestId"] == TEST_MODEL_SWITCH_REQUEST_ID + })); + assert!( + !before_commit + .iter() + .any(|event| event.payload["status"] == "switched"), + "restored request success must remain deferred until local commit" + ); + + let channel_id = Uuid::new_v4(); + agent + .state + .sessions + .insert(channel_id, prepared.session_id.clone()); + assert!(commit_prepared_session( + &mut agent, + &PromptSource::Channel(channel_id), + &prepared + )); + let after_commit: Vec<_> = observer + .snapshot() + .into_iter() + .filter(|event| event.kind == "control_result") + .collect(); + assert!(after_commit.iter().any(|event| { + event.payload["status"] == "switched" + && event.payload["modelId"] == "model-a" + && event.payload["requestId"] == PRIOR_REQUEST_ID + })); + assert!(agent.model_switch_request_id.is_none()); + assert!(agent.model_switch_rollback.is_none()); + + agent + .acp + .shutdown() + .await + .expect("fake adapter must shut down cleanly"); + let _ = std::fs::remove_file(capture); + } + + #[tokio::test] + async fn rejected_model_config_option_emits_failure_and_closes_partial_session() { + let capture = capture_path("model-config-rejection"); + let script = r#" + capture="$1" + IFS= read -r initialize + printf '%s\n' "$initialize" >> "$capture" + printf '%s\n' '{"jsonrpc":"2.0","id":0,"result":{"protocolVersion":2,"agentCapabilities":{"sessionCapabilities":{"close":{}}}}}' + IFS= read -r new_session + printf '%s\n' "$new_session" >> "$capture" + printf '%s\n' '{"jsonrpc":"2.0","id":1,"result":{"sessionId":"session-created","configOptions":[{"configId":"model","category":"model","options":[{"value":"model-b"}]}]}}' + IFS= read -r set_model + printf '%s\n' "$set_model" >> "$capture" + printf '%s\n' '{"jsonrpc":"2.0","id":2,"error":{"code":-32042,"message":"model rejected"}}' + IFS= read -r close + printf '%s\n' "$close" >> "$capture" + printf '%s\n' '{"jsonrpc":"2.0","id":3,"result":{}}' + IFS= read -r retry_session + printf '%s\n' "$retry_session" >> "$capture" + printf '%s\n' '{"jsonrpc":"2.0","id":4,"result":{"sessionId":"session-retry","configOptions":[{"configId":"model","category":"model","options":[{"value":"model-a"},{"value":"model-b"}]}]}}' + IFS= read -r retry_model + printf '%s\n' "$retry_model" >> "$capture" + printf '%s\n' '{"jsonrpc":"2.0","id":5,"result":{}}' + sleep 10 + "#; + let mut acp = AcpClient::spawn( + "bash", + &[ + "-c".to_string(), + script.to_string(), + "buzz-acp-model-config-rejection-test".to_string(), + capture.to_string_lossy().into_owned(), + ], + &[], + false, + ) + .await + .expect("failed to spawn fake ACP adapter"); + let observer = observer::ObserverHandle::in_process(); + acp.set_observer(Some(observer.clone()), 0); + acp.initialize() + .await + .expect("fake adapter must advertise session/close"); + + let channel_id = Uuid::new_v4(); + let mut agent = OwnedAgent { + index: 0, + acp, + state: SessionState::default(), + model_capabilities: None, + desired_model: Some("model-a".into()), + model_overridden: false, + model_switch_request_id: None, + model_switch_rollback: None, + agent_name: "claude-code-acp".into(), + goose_system_prompt_supported: None, + protocol_version: 2, + }; + assert!(stage_busy_model_switch_intent( + &mut agent, + &model_switch_request("model-b") + )); + let mut ctx = make_prompt_context_no_owner(); + ctx.dedup_mode = DedupMode::Queue; + let (result_tx, mut result_rx) = mpsc::unbounded_channel(); + + run_prompt_task( + agent, + Some(one_event_batch(channel_id)), + Some("prompt".into()), + Arc::new(ctx), + result_tx, + None, + "turn-model-config-rejection".into(), + ) + .await; + let mut result = result_rx + .recv() + .await + .expect("model rejection result must be sent"); + + assert!(matches!( + &result.outcome, + PromptOutcome::Error(AcpError::AgentError { code: -32042, .. }) + )); + assert!( + !result.agent.state.sessions.contains_key(&channel_id), + "a rejected model must never publish the partial session locally" + ); + assert_eq!( + result.agent.desired_model.as_deref(), + Some("model-a"), + "adapter rejection must restore the prior desired model" + ); + assert!( + !result.agent.model_overridden, + "adapter rejection must restore configured-model provenance" + ); + assert_eq!( + result.agent.model_switch_request_id, None, + "the rejected request must not remain eligible for a silent retry" + ); + assert!( + result.agent.model_switch_rollback.is_none(), + "restoring the prior intent must consume the rejected request's rollback frame" + ); + + let captured = wait_for_capture_lines(&capture, 4).await; + let messages: Vec = captured + .lines() + .map(|line| serde_json::from_str(line).expect("captured line must be JSON")) + .collect(); + assert_eq!(messages[0]["method"], "initialize"); + assert_eq!(messages[1]["method"], "session/new"); + assert_eq!(messages[2]["method"], "session/set_config_option"); + assert_eq!( + messages[3]["method"], "session/close", + "application rejection must close the partially-created session before reuse" + ); + assert_eq!(messages[3]["params"]["sessionId"], "session-created"); + + let control_results: Vec<_> = observer + .snapshot() + .into_iter() + .filter(|event| event.kind == "control_result") + .collect(); + assert!(control_results.iter().any(|event| { + event.payload["status"] == "switch_failed" + && event.payload["modelId"] == "model-b" + && event.payload["requestId"] == TEST_MODEL_SWITCH_REQUEST_ID + })); + assert!( + !control_results + .iter() + .any(|event| event.payload["status"] == "switched"), + "a framed model rejection must never become application proof" + ); + assert!( + !observer.snapshot().iter().any(|event| { + event.kind == "session_config_captured" + && event.payload["modelOverridden"] == serde_json::Value::Bool(true) + }), + "a rejected model must never be reported as an applied override" + ); + + let retry_ctx = make_prompt_context_no_owner(); + let prepared = + match create_session_and_apply_model(&mut result.agent, &retry_ctx, None, None, None) + .await + { + Ok(prepared) => prepared, + Err(_) => panic!("retry with the restored model must succeed"), + }; + result + .agent + .state + .sessions + .insert(channel_id, prepared.session_id.clone()); + assert!(commit_prepared_session( + &mut result.agent, + &PromptSource::Channel(channel_id), + &prepared + )); + + let captured = wait_for_capture_lines(&capture, 6).await; + let messages: Vec = captured + .lines() + .map(|line| serde_json::from_str(line).expect("captured line must be JSON")) + .collect(); + assert_eq!(messages[4]["method"], "session/new"); + assert_eq!(messages[5]["method"], "session/set_config_option"); + assert_eq!( + messages[5]["params"]["value"], "model-a", + "the next session must apply the restored model, never retry the rejected model" + ); + + result + .agent + .acp + .shutdown() + .await + .expect("fake adapter must shut down cleanly"); + let _ = std::fs::remove_file(capture); + } + + #[tokio::test] + async fn rejected_session_set_model_is_returned_as_switch_failure() { + let capture = capture_path("set-model-rejection"); + let script = r#" + capture="$1" + IFS= read -r set_model + printf '%s\n' "$set_model" >> "$capture" + printf '%s\n' '{"jsonrpc":"2.0","id":0,"error":{"code":-32043,"message":"set_model rejected"}}' + sleep 10 + "#; + let mut acp = AcpClient::spawn( + "bash", + &[ + "-c".to_string(), + script.to_string(), + "buzz-acp-set-model-rejection-test".to_string(), + capture.to_string_lossy().into_owned(), + ], + &[], + false, + ) + .await + .expect("failed to spawn fake ACP adapter"); + + let result = apply_model_switch( + &mut acp, + "session-created", + "model-b", + &ModelSwitchMethod::SetModel { + model_id: "model-b".into(), + }, + ) + .await; + + assert!(matches!( + result, + Err(AcpError::AgentError { code: -32043, .. }) + )); + let captured = wait_for_capture_lines(&capture, 1).await; + let request: serde_json::Value = + serde_json::from_str(captured.trim()).expect("captured request must be JSON"); + assert_eq!(request["method"], "session/set_model"); + assert_eq!(request["params"]["modelId"], "model-b"); + + acp.shutdown() + .await + .expect("fake adapter must shut down cleanly"); + let _ = std::fs::remove_file(capture); + } + + #[tokio::test] + async fn failed_goose_session_setup_without_close_support_requires_process_recycle() { + let capture = capture_path("goose-setup-recycle"); + let script = r#" + capture="$1" + IFS= read -r initialize + printf '%s\n' "$initialize" >> "$capture" + printf '%s\n' '{"jsonrpc":"2.0","id":0,"result":{"protocolVersion":2,"agentCapabilities":{"sessionCapabilities":{}}}}' + IFS= read -r new_session + printf '%s\n' "$new_session" >> "$capture" + printf '%s\n' '{"jsonrpc":"2.0","id":1,"result":{"sessionId":"session-created"}}' + IFS= read -r system_prompt + printf '%s\n' "$system_prompt" >> "$capture" + printf '%s\n' '{"jsonrpc":"2.0","id":2,"error":{"code":-32042,"message":"setup rejected"}}' + if IFS= read -r -t 1 unexpected; then + printf '%s\n' "$unexpected" >> "$capture" + fi + sleep 10 + "#; + let mut acp = AcpClient::spawn( + "bash", + &[ + "-c".to_string(), + script.to_string(), + "buzz-acp-goose-setup-recycle-test".to_string(), + capture.to_string_lossy().into_owned(), + ], + &[], + false, + ) + .await + .expect("failed to spawn fake ACP adapter"); + acp.initialize() + .await + .expect("fake adapter must complete capability negotiation"); + assert!(!acp.supports_session_close()); + + let channel_id = Uuid::new_v4(); + let agent = OwnedAgent { + index: 0, + acp, + state: SessionState::default(), + model_capabilities: None, + desired_model: None, + model_overridden: false, + model_switch_request_id: None, + model_switch_rollback: None, + agent_name: "goose".into(), + goose_system_prompt_supported: None, + protocol_version: 2, + }; + let mut ctx = make_prompt_context_no_owner(); + ctx.base_prompt = Some("setup prompt"); + ctx.dedup_mode = DedupMode::Queue; + let (result_tx, mut result_rx) = mpsc::unbounded_channel(); + + run_prompt_task( + agent, + Some(one_event_batch(channel_id)), + Some("prompt".into()), + Arc::new(ctx), + result_tx, + None, + "turn-goose-setup-recycle".into(), + ) + .await; + let result = result_rx + .recv() + .await + .expect("setup failure result must be sent"); + + assert!(matches!( + result.outcome, + PromptOutcome::SessionSetupRecycleRequired + )); + let captured = wait_for_capture_lines(&capture, 3).await; + assert_eq!( + captured.lines().count(), + 3, + "an unadvertised optional close method must never be probed" + ); + + let _ = std::fs::remove_file(capture); + } + + #[test] + fn test_requeue_cancelled_batch_maps_control_signal_to_cancel_reason() { + let cases = [ + (ControlSignal::Steer, Some(CancelReason::Steer)), + (ControlSignal::Interrupt, Some(CancelReason::Interrupt)), + ( + ControlSignal::SwitchModel(model_switch_request("gpt-5")), + Some(CancelReason::Interrupt), + ), + (ControlSignal::Cancel, None), + (ControlSignal::Rotate, None), + ]; + let mut ctx = make_prompt_context_no_owner(); + ctx.dedup_mode = DedupMode::Queue; + + for (signal, expected_reason) in cases { + let channel_id = Uuid::new_v4(); + let batch = one_event_batch(channel_id); + let result = requeue_cancelled_batch(&ctx, signal.clone(), Some(batch)); + match expected_reason { + Some(reason) => { + let batch = result + .unwrap_or_else(|| panic!("{signal:?} must preserve the batch, got None")); + assert_eq!( + batch.cancel_reason, + Some(reason), + "{signal:?} must stamp {reason:?}" + ); + } + None => assert!( + result.is_none(), + "{signal:?} must drop the batch, got {result:?}" + ), + } + } + } + + // ── classify_control_cancel_failure ───────────────────────────────────── + // Table-driven pin of the single production seam used by the + // `Err(error)` arm in `run_prompt_task`'s control-cancel branch. Crosses + // the exact error→outcome AND outcome→batch-fate boundary in one call, + // so a regression to the old per-arm duplication (or to routing an + // unexpected HardTimeout back through the real hard-cap path) fails + // here rather than only in independently-manufactured unit tests. + + /// Assert `outcome` is the expected `PromptOutcome` variant. `PromptOutcome` + /// has no `PartialEq` (it wraps `AcpError`, which isn't `PartialEq`), so + /// this matches by shape instead of deriving equality onto the whole enum. + fn assert_outcome_matches(outcome: &PromptOutcome, expected: &str) { + let label = match outcome { + PromptOutcome::AgentExited => "AgentExited", + PromptOutcome::Timeout(TimeoutKind::Idle) => "Timeout(Idle)", + PromptOutcome::Timeout(TimeoutKind::Hard { .. }) => "Timeout(Hard)", + PromptOutcome::CancelDrainTimeout(_) => "CancelDrainTimeout", + PromptOutcome::CancelCleanupFailed(_) => "CancelCleanupFailed", + PromptOutcome::SessionCloseFailed(_) => "SessionCloseFailed", + PromptOutcome::SessionRecycleRequired => "SessionRecycleRequired", + PromptOutcome::SessionSetupCloseFailed(_) => "SessionSetupCloseFailed", + PromptOutcome::SessionSetupRecycleRequired => "SessionSetupRecycleRequired", + PromptOutcome::ControlPreemptedSetup => "ControlPreemptedSetup", + PromptOutcome::Error(_) => "Error", + PromptOutcome::Cancelled => "Cancelled", + PromptOutcome::Ok(_) => "Ok", + PromptOutcome::SessionRetired(_) => "SessionRetired", + }; assert_eq!( label, expected, "got outcome shape {label}, want {expected}" @@ -4657,7 +6939,7 @@ mod tests { Case { name: "CancelDrainTimeout + SwitchModel preserves batch with Interrupt reason", error: || AcpError::CancelDrainTimeout(CONTROL_CANCEL_GRACE), - signal: ControlSignal::SwitchModel("gpt-5".to_string()), + signal: ControlSignal::SwitchModel(model_switch_request("gpt-5")), expected_outcome: "CancelDrainTimeout", batch_preserved: true, expected_reason: Some(CancelReason::Interrupt), @@ -4701,6 +6983,18 @@ mod tests { expected_reason: Some(CancelReason::Steer), invalidate_all: false, }, + Case { + name: "post-cancel AgentError poisons adapter and preserves Steer batch", + error: || AcpError::AgentError { + code: -32000, + message: "prompt rejected after cancel".into(), + }, + signal: ControlSignal::Steer, + expected_outcome: "CancelCleanupFailed", + batch_preserved: true, + expected_reason: Some(CancelReason::Steer), + invalidate_all: false, + }, ]; for case in cases { @@ -5075,6 +7369,8 @@ mod tests { model_capabilities: None, desired_model: None, model_overridden: false, + model_switch_request_id: None, + model_switch_rollback: None, agent_name: "unknown".into(), goose_system_prompt_supported: None, protocol_version: 2, @@ -5133,6 +7429,8 @@ mod tests { model_capabilities: None, desired_model: None, model_overridden: false, + model_switch_request_id: None, + model_switch_rollback: None, agent_name: "unknown".into(), goose_system_prompt_supported: None, protocol_version: 2, @@ -5193,6 +7491,20 @@ mod tests { assert_eq!(acp_stop_to_core(&StopReason::Refusal), CoreStop::Unknown); } + #[test] + fn failed_session_retirement_marks_completed_turn_metric_as_error() { + use buzz_core::agent_turn_metric::StopReason as CoreStop; + + assert_eq!( + completed_turn_metric_stop(&StopReason::MaxTokens, true), + CoreStop::Error + ); + assert_eq!( + completed_turn_metric_stop(&StopReason::MaxTurnRequests, true), + CoreStop::Error + ); + } + /// `publish_agent_turn_metric` is a no-op when `usage` is `None`. #[tokio::test] async fn test_publish_agent_turn_metric_noop_on_no_usage() { diff --git a/crates/buzz-acp/src/queue.rs b/crates/buzz-acp/src/queue.rs index 029bf86dbf..9ad66bb230 100644 --- a/crates/buzz-acp/src/queue.rs +++ b/crates/buzz-acp/src/queue.rs @@ -15,16 +15,17 @@ use nostr::{Event, ToBech32}; use std::collections::{HashMap, HashSet, VecDeque}; +use std::sync::atomic::{AtomicU64, Ordering}; use std::time::{Duration, Instant}; use uuid::Uuid; use crate::config::DedupMode; -/// Maximum events queued per channel before oldest events are dropped. +/// Maximum retained ordinary plus cancelled events per channel. const MAX_PENDING_PER_CHANNEL: usize = 500; /// Maximum events drained into a single batch. -const MAX_BATCH_EVENTS: usize = 50; +pub(crate) const MAX_BATCH_EVENTS: usize = 50; /// Maximum retry attempts before a batch is dead-lettered. pub(crate) const MAX_RETRIES: u32 = 10; @@ -41,6 +42,22 @@ const IN_FLIGHT_DEADLINE_BUFFER_SECS: u64 = 100; /// Default in-flight deadline: default max_turn (7200s) + 100s buffer. const DEFAULT_IN_FLIGHT_DEADLINE_SECS: u64 = 7300; +/// Process-local total order for accepted queue occurrences. +/// +/// `Instant` is monotonic but not unique. The ordinal is carried through +/// every queue/retry transition so equal-timestamp events retain enqueue FIFO. +static NEXT_OCCURRENCE_ORDER: AtomicU64 = AtomicU64::new(1); + +fn next_occurrence_order() -> u64 { + let order = NEXT_OCCURRENCE_ORDER.fetch_add(1, Ordering::Relaxed); + assert_ne!( + order, + u64::MAX, + "queue occurrence ordinal exhausted before process restart" + ); + order +} + /// An event waiting in the queue. #[derive(Debug, Clone)] pub struct QueuedEvent { @@ -51,6 +68,14 @@ pub struct QueuedEvent { pub prompt_tag: String, } +/// Opaque identity for one accepted [`EventQueue::push`] occurrence. +/// +/// This is intentionally unrelated to the signed Nostr event ID: relay replay +/// may enqueue the same signed event more than once, and an asynchronous steer +/// acknowledgement must affect only the occurrence that initiated it. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub(crate) struct EnqueueOccurrenceId(Uuid); + /// A single event inside a [`FlushBatch`]. #[derive(Debug, Clone)] pub struct BatchEvent { @@ -59,6 +84,84 @@ pub struct BatchEvent { pub received_at: Instant, } +/// Queue-private identity carried alongside a value for one enqueue +/// occurrence. Two deliveries of the same signed Nostr event intentionally +/// receive different IDs. +#[derive(Debug, Clone)] +struct Occurrence { + id: Uuid, + order: u64, + value: T, +} + +impl Occurrence { + fn new(value: T) -> Self { + Self { + id: Uuid::new_v4(), + order: next_occurrence_order(), + value, + } + } + + fn with_identity(id: Uuid, order: u64, value: T) -> Self { + Self { id, order, value } + } + + fn map(self, map: impl FnOnce(T) -> U) -> Occurrence { + Occurrence { + id: self.id, + order: self.order, + value: map(self.value), + } + } +} + +impl std::ops::Deref for Occurrence { + type Target = T; + + fn deref(&self) -> &Self::Target { + &self.value + } +} + +impl std::ops::DerefMut for Occurrence { + fn deref_mut(&mut self) -> &mut Self::Target { + &mut self.value + } +} + +/// Exact occurrence identity aligned position-for-position with a +/// [`FlushBatch`]'s two event vectors. +/// +/// This is crate-private so callers cannot manufacture or reinterpret retry +/// membership. Every queue transition validates exact lengths and uniqueness; +/// missing or malformed identity fails closed. +#[derive(Debug, Clone, Default)] +pub(crate) struct BatchOccurrenceIds { + events: Vec, + cancelled_events: Vec, + event_orders: Vec, + cancelled_event_orders: Vec, +} + +#[cfg(test)] +impl BatchOccurrenceIds { + /// Explicit identities for test-built batches that enter queue transitions. + /// + /// Formatting-only test batches can continue to use `Default`; production + /// transitions never infer missing occurrence identity. + pub(crate) fn for_test(event_count: usize, cancelled_event_count: usize) -> Self { + Self { + events: (0..event_count).map(|_| Uuid::new_v4()).collect(), + cancelled_events: (0..cancelled_event_count).map(|_| Uuid::new_v4()).collect(), + event_orders: (0..event_count).map(|_| next_occurrence_order()).collect(), + cancelled_event_orders: (0..cancelled_event_count) + .map(|_| next_occurrence_order()) + .collect(), + } + } +} + /// Why a batch's prior turn was cancelled — controls how `format_prompt` /// frames the merged re-prompt. #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -87,6 +190,21 @@ pub struct FlushBatch { /// [`Steer`](CancelReason::Steer) framing if a merge somehow lacks a reason /// (see [`MergeFraming::for_reason`]). pub cancel_reason: Option, + /// Queue-private per-enqueue identity, aligned with `events` and + /// `cancelled_events`. External behavior continues to use the signed Nostr + /// event; retry bookkeeping uses these occurrence IDs. + pub(crate) occurrence_ids: BatchOccurrenceIds, +} + +/// Immutable identity for one dispatched cohort. +/// +/// The occurrence membership is captured when the cohort first leaves the +/// queue. Same-channel arrivals after that point are intentionally excluded, +/// so they can never inherit this cohort's retry count or dead-letter fate. +#[derive(Debug)] +struct ActiveBatchIdentity { + id: Uuid, + occurrence_ids: HashSet, } /// Per-channel event queue with per-channel in-flight enforcement. @@ -95,30 +213,38 @@ pub struct FlushBatch { /// /// ```text /// State: -/// queues: Map> (capped at MAX_PENDING_PER_CHANNEL) +/// queues: Map> +/// cancelled_batches: Map> +/// (combined per-channel count capped at MAX_PENDING_PER_CHANNEL) /// in_flight_channels: HashSet /// in_flight_deadlines: Map (auto-expire after in_flight_deadline) -/// retry_after: Map -/// retry_counts: Map (dead-letter after MAX_RETRIES) +/// active_batches: Map +/// 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) +/// while queues[channel].len() + cancelled_batches[channel].len() +/// > MAX_PENDING_PER_CHANNEL: +/// drop the event with the oldest received_at across both stores /// /// 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) +/// AND (no active-batch retry_after OR deadline <= now) /// if candidates empty: return None /// channel = pick candidate with oldest head event (min received_at) -/// events = drain up to MAX_BATCH_EVENTS from queues[channel] +/// if channel has an active batch: +/// events = remove only that batch's immutable occurrence membership +/// else: +/// events = drain up to MAX_BATCH_EVENTS from queues[channel] +/// capture a fresh immutable batch identity /// in_flight_channels.insert(channel) /// in_flight_deadlines.insert(channel, now + in_flight_deadline) /// return Some(FlushBatch { channel, events }) @@ -126,44 +252,64 @@ pub struct FlushBatch { /// mark_complete(channel_id): /// in_flight_channels.remove(channel_id) /// in_flight_deadlines.remove(channel_id) -/// retry_counts.remove(channel_id) -/// clean up expired retry_after entry if present +/// if no active backoff remains: +/// remove active batch identity and its retry metadata /// /// requeue(batch): -/// increment retry_counts[channel] -/// if retry_counts[channel] > MAX_RETRIES: dead-letter (log ERROR, return batch to caller) +/// increment retry_counts[active batch ID] +/// if count > MAX_RETRIES: dead-letter only that immutable cohort /// else: push_front with original received_at, set exponential backoff retry_after with jitter /// ``` pub struct EventQueue { - queues: HashMap>, + queues: HashMap>>, in_flight_channels: HashSet, /// Per-channel deadline for auto-expiring stuck in-flight entries. in_flight_deadlines: HashMap, /// Number of events in each in-flight batch (for expiry logging). in_flight_batch_sizes: HashMap, + /// Retry deadline keyed by immutable batch identity, not channel. retry_after: HashMap, - /// Per-channel retry attempt counter for exponential backoff / dead-lettering. + /// Retry attempt counter keyed by immutable batch identity. retry_counts: HashMap, + /// The one dispatched/requeued cohort currently owning each channel. + /// + /// A channel remains single-flight, so it can have at most one active + /// cohort. New arrivals are absent from `occurrence_ids` and remain queued + /// for a later, fresh identity after this cohort completes or dead-letters. + active_batches: HashMap, dedup_mode: DedupMode, /// Events from cancelled batches, keyed by channel. Merged into the next /// `FlushBatch` for that channel as `cancelled_events` so `format_prompt()` /// can produce annotated "[Previous request — interrupted]" sections. - cancelled_batches: HashMap>, + cancelled_batches: HashMap>>, /// Why each channel's cancelled batch was cancelled (steer vs interrupt). /// Set by `requeue_as_cancelled`, consumed by `flush_next` to set /// `FlushBatch::cancel_reason`. Keyed by channel, cleared on flush. cancel_reasons: HashMap, + /// Exact agent slot required for a model-switch replay. + /// + /// The requirement survives flush/requeue cycles until the pinned batch + /// completes or the channel is removed. This prevents a requeued + /// `SwitchModel` turn from being claimed by a different idle slot whose + /// process still carries the default model. + required_agents: HashMap, + /// Required-agent channels whose exact slot is temporarily unavailable. + /// + /// They are excluded from `flush_next` so one recycling slot cannot + /// head-of-line block unrelated channels. `release_required_agent` wakes + /// them when that slot returns. + blocked_required_agents: HashSet, /// Events withheld from `queues` while a goose-native steer is in flight /// for that event. Invisible to `flush_next` / `has_flushable_work` / /// `drain` (the events have been moved out of `queues`), so the queue's /// no-double-deliver invariant holds without any change to the hot drain /// path. Populated by [`mark_native_steer_pending`]; drained back to the - /// queue front by [`release_native_steer`] (preserving original - /// `received_at` fairness, same discipline as `requeue_preserve_timestamps` - /// at line 453). Bulk recovery on in-flight deadline expiry is performed + /// queue by [`release_native_steer`] (preserving original `received_at` + /// plus enqueue-order fairness). Bulk recovery on in-flight deadline + /// expiry is performed /// by `flush_next` / `has_flushable_work` (recover, not log-and-drop — /// the events were never delivered to the agent). - withheld_native_steer: HashMap>, + withheld_native_steer: HashMap>>, /// Duration after which an in-flight channel is auto-expired as orphaned. /// Must be strictly greater than `max_turn_duration` so a turn running to /// the hard cap returns via `mark_complete` before the backstop fires. @@ -184,9 +330,12 @@ impl EventQueue { in_flight_batch_sizes: HashMap::new(), retry_after: HashMap::new(), retry_counts: HashMap::new(), + active_batches: HashMap::new(), dedup_mode, cancelled_batches: HashMap::new(), cancel_reasons: HashMap::new(), + required_agents: HashMap::new(), + blocked_required_agents: HashSet::new(), withheld_native_steer: HashMap::new(), in_flight_deadline: Duration::from_secs(DEFAULT_IN_FLIGHT_DEADLINE_SECS), } @@ -226,8 +375,9 @@ impl EventQueue { /// In [`DedupMode::Drop`], events for any currently in-flight channel are /// silently discarded (debug-logged). /// - /// Returns `true` if the event was accepted, `false` if dropped. - pub fn push(&mut self, event: QueuedEvent) -> bool { + /// Returns the queue-private occurrence identity when accepted, or `None` + /// when dropped. + pub(crate) fn push(&mut self, event: QueuedEvent) -> Option { if matches!(self.dedup_mode, DedupMode::Drop) && self.in_flight_channels.contains(&event.channel_id) { @@ -235,28 +385,295 @@ impl EventQueue { channel_id = %event.channel_id, "dropping event for in-flight channel (drop mode)" ); - return false; + return None; + } + let channel_id = event.channel_id; + let occurrence = Occurrence::new(event); + let occurrence_id = EnqueueOccurrenceId(occurrence.id); + self.queues + .entry(channel_id) + .or_default() + .push_back(occurrence); + self.enforce_pending_limit(channel_id, "push"); + Some(occurrence_id) + } + + /// Enforce the single per-channel retention budget shared by ordinary + /// queued events and cancelled-batch provenance. + /// + /// Eviction is chronological by the preserved `received_at` timestamp, + /// irrespective of which store owns the event. Ties evict cancelled + /// history before ordinary work. This keeps the newest retained context + /// while preserving FIFO/fairness timestamps for everything that remains. + /// If all cancelled history is evicted, its framing reason is removed too. + fn enforce_pending_limit(&mut self, channel_id: Uuid, mutation: &'static str) { + let mut dropped = 0usize; + loop { + let queued_len = self.queues.get(&channel_id).map_or(0, VecDeque::len); + let cancelled_len = self.cancelled_batches.get(&channel_id).map_or(0, Vec::len); + if queued_len + cancelled_len <= MAX_PENDING_PER_CHANNEL { + break; + } + + let oldest_queued = self.queues.get(&channel_id).and_then(|events| { + events + .iter() + .enumerate() + .min_by_key(|(_, event)| event.received_at) + .map(|(index, event)| (index, event.received_at)) + }); + let oldest_cancelled = self.cancelled_batches.get(&channel_id).and_then(|events| { + events + .iter() + .enumerate() + .min_by_key(|(_, event)| event.received_at) + .map(|(index, event)| (index, event.received_at)) + }); + + let dropped_occurrence_id = match (oldest_queued, oldest_cancelled) { + (Some((queued_index, queued_at)), Some((cancelled_index, cancelled_at))) => { + if cancelled_at <= queued_at { + self.cancelled_batches + .get_mut(&channel_id) + .map(|events| events.remove(cancelled_index).id) + } else if let Some(events) = self.queues.get_mut(&channel_id) { + events.remove(queued_index).map(|event| event.id) + } else { + None + } + } + (Some((queued_index, _)), None) => self + .queues + .get_mut(&channel_id) + .and_then(|events| events.remove(queued_index)) + .map(|event| event.id), + (None, Some((cancelled_index, _))) => self + .cancelled_batches + .get_mut(&channel_id) + .map(|events| events.remove(cancelled_index).id), + (None, None) => None, + }; + let Some(dropped_occurrence_id) = dropped_occurrence_id else { + break; + }; + if let Some(active) = self.active_batches.get_mut(&channel_id) { + active.occurrence_ids.remove(&dropped_occurrence_id); + } + dropped += 1; + } + + if self.queues.get(&channel_id).is_some_and(VecDeque::is_empty) { + self.queues.remove(&channel_id); + } + if self + .cancelled_batches + .get(&channel_id) + .is_some_and(Vec::is_empty) + { + self.cancelled_batches.remove(&channel_id); + } + if !self.cancelled_batches.contains_key(&channel_id) { + self.cancel_reasons.remove(&channel_id); } - 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(); + if dropped > 0 { tracing::warn!( - channel_id = %event.channel_id, + channel_id = %channel_id, + mutation, + dropped, limit = MAX_PENDING_PER_CHANNEL, - "queue depth cap reached — dropped oldest event" + "combined pending cap reached — dropped oldest retained event(s)" ); } - queue.push_back(event); - true + } + + fn validate_batch_occurrence_ids(batch: &FlushBatch) -> bool { + let event_count_matches = batch.occurrence_ids.events.len() == batch.events.len(); + let cancelled_count_matches = + batch.occurrence_ids.cancelled_events.len() == batch.cancelled_events.len(); + let event_order_count_matches = + batch.occurrence_ids.event_orders.len() == batch.events.len(); + let cancelled_order_count_matches = + batch.occurrence_ids.cancelled_event_orders.len() == batch.cancelled_events.len(); + let expected_count = batch.events.len() + batch.cancelled_events.len(); + let distinct_count = Self::batch_occurrence_membership(batch).len(); + let distinct_order_count = batch + .occurrence_ids + .event_orders + .iter() + .chain(batch.occurrence_ids.cancelled_event_orders.iter()) + .copied() + .collect::>() + .len(); + let identities_are_unique = distinct_count == expected_count; + let orders_are_unique = distinct_order_count == expected_count; + if event_count_matches + && cancelled_count_matches + && event_order_count_matches + && cancelled_order_count_matches + && identities_are_unique + && orders_are_unique + { + return true; + } + + tracing::error!( + channel_id = %batch.channel_id, + events = batch.events.len(), + event_occurrence_ids = batch.occurrence_ids.events.len(), + cancelled_events = batch.cancelled_events.len(), + cancelled_occurrence_ids = batch.occurrence_ids.cancelled_events.len(), + distinct_occurrence_ids = distinct_count, + event_occurrence_orders = batch.occurrence_ids.event_orders.len(), + cancelled_occurrence_orders = batch.occurrence_ids.cancelled_event_orders.len(), + distinct_occurrence_orders = distinct_order_count, + "batch occurrence identity invariant violated — refusing inferred retry membership" + ); + false + } + + fn batch_occurrence_membership(batch: &FlushBatch) -> HashSet { + batch + .occurrence_ids + .events + .iter() + .chain(batch.occurrence_ids.cancelled_events.iter()) + .copied() + .collect() + } + + fn activate_batch(&mut self, batch: &FlushBatch) -> Option { + if !Self::validate_batch_occurrence_ids(batch) { + return None; + } + let occurrence_ids = Self::batch_occurrence_membership(batch); + if let Some(active) = self.active_batches.get_mut(&batch.channel_id) { + // Test-built recovery batches can seed retry accounting before an + // identity has event members. Production batches always arrive + // here with the immutable membership captured by flush_next(). + if active.occurrence_ids.is_empty() { + active.occurrence_ids = occurrence_ids; + } else if active.occurrence_ids != occurrence_ids { + tracing::error!( + channel_id = %batch.channel_id, + active_occurrences = active.occurrence_ids.len(), + batch_occurrences = occurrence_ids.len(), + "batch occurrence identity differs from active cohort — refusing inferred retry membership" + ); + return None; + } + return Some(active.id); + } + + let id = Uuid::new_v4(); + self.active_batches + .insert(batch.channel_id, ActiveBatchIdentity { id, occurrence_ids }); + Some(id) + } + + fn clear_active_batch(&mut self, channel_id: Uuid) { + if let Some(active) = self.active_batches.remove(&channel_id) { + self.retry_after.remove(&active.id); + self.retry_counts.remove(&active.id); + } + } + + fn active_batch_oldest_pending(&self, channel_id: Uuid) -> Option { + let active = self.active_batches.get(&channel_id)?; + let ordinary = self.queues.get(&channel_id).and_then(|events| { + events + .iter() + .filter(|event| active.occurrence_ids.contains(&event.id)) + .map(|event| event.received_at) + .min() + }); + let cancelled = self.cancelled_batches.get(&channel_id).and_then(|events| { + events + .iter() + .filter(|event| active.occurrence_ids.contains(&event.id)) + .map(|event| event.received_at) + .min() + }); + match (ordinary, cancelled) { + (Some(left), Some(right)) => Some(left.min(right)), + (Some(value), None) | (None, Some(value)) => Some(value), + (None, None) => None, + } + } + + fn channel_oldest_pending(&self, channel_id: Uuid) -> Option { + if self.active_batches.contains_key(&channel_id) { + return self.active_batch_oldest_pending(channel_id); + } + + let ordinary = self + .queues + .get(&channel_id) + .and_then(|events| events.front()) + .map(|event| event.received_at); + let cancelled = self + .cancelled_batches + .get(&channel_id) + .and_then(|events| events.iter().map(|event| event.received_at).min()); + match (ordinary, cancelled) { + (Some(left), Some(right)) => Some(left.min(right)), + (Some(value), None) | (None, Some(value)) => Some(value), + (None, None) => None, + } + } + + fn prune_orphaned_active_batches(&mut self) { + let orphaned: Vec = self + .active_batches + .keys() + .copied() + .filter(|channel_id| { + !self.in_flight_channels.contains(channel_id) + && self.active_batch_oldest_pending(*channel_id).is_none() + }) + .collect(); + for channel_id in orphaned { + self.clear_active_batch(channel_id); + } + } + + fn oldest_ready_channel(&self, now: Instant) -> Option { + let channels: HashSet = self + .queues + .keys() + .chain(self.cancelled_batches.keys()) + .copied() + .collect(); + + channels + .into_iter() + .filter(|channel_id| { + !self.in_flight_channels.contains(channel_id) + && !self.blocked_required_agents.contains(channel_id) + && self + .active_batches + .get(channel_id) + .and_then(|active| self.retry_after.get(&active.id)) + .is_none_or(|deadline| *deadline <= now) + }) + .filter_map(|channel_id| { + self.channel_oldest_pending(channel_id) + .map(|received_at| (channel_id, received_at)) + }) + .min_by(|left, right| { + left.1 + .cmp(&right.1) + .then_with(|| left.0.as_u128().cmp(&right.0.as_u128())) + }) + .map(|(channel_id, _)| channel_id) } /// Try to flush the next batch. /// /// Returns `None` if all non-in-flight, non-throttled queues are empty. /// Otherwise picks the channel with the oldest pending event (FIFO fairness - /// across channels), drains ALL events for that channel into a single batch, - /// inserts into `in_flight_channels`, and returns the batch. + /// across channels), reconstructs an existing immutable retry/deferred + /// cohort or drains at most [`MAX_BATCH_EVENTS`] ordinary events into a + /// fresh cohort, inserts into `in_flight_channels`, and returns the batch. pub fn flush_next(&mut self) -> Option { let now = Instant::now(); @@ -278,6 +695,7 @@ impl EventQueue { ); self.in_flight_channels.remove(&id); self.in_flight_deadlines.remove(&id); + self.clear_active_batch(id); // Recover any withheld goose-native steer events for the expired // channel back to the queue front so normal dispatch delivers // them. Unlike the in-flight batch above (already delivered to a @@ -286,97 +704,139 @@ impl EventQueue { self.recover_withheld_for_expired_channel(id); } - // Find the channel whose head event has the oldest received_at, - // excluding in-flight channels and throttled channels. - let channel_id = self - .queues - .iter() - .filter(|(id, q)| { - !q.is_empty() - && !self.in_flight_channels.contains(id) - && self.retry_after.get(id).is_none_or(|&t| t <= now) - }) - .min_by_key(|(_, q)| q.front().unwrap().received_at) - .map(|(id, _)| *id); - - // Fallback: if no queued events are ready but a channel has cancelled - // events waiting (e.g., explicit !cancel with no new @mention), flush - // those as a regular batch (re-dispatch unchanged). - let channel_id = match channel_id { - Some(id) => id, - None => { - let cancelled_id = self - .cancelled_batches - .keys() - .find(|id| !self.in_flight_channels.contains(id)) - .copied(); - match cancelled_id { - Some(id) => { - // Move cancelled events into the regular events slot. - // No new events to merge — re-dispatch the original batch. - let cancelled = self.cancelled_batches.remove(&id).unwrap_or_default(); - let cancel_reason = self.cancel_reasons.remove(&id); - self.in_flight_channels.insert(id); - self.in_flight_deadlines - .insert(id, now + self.in_flight_deadline); - self.in_flight_batch_sizes.insert(id, cancelled.len()); - return Some(FlushBatch { - channel_id: id, - events: cancelled, - cancelled_events: vec![], - cancel_reason, - }); + self.prune_orphaned_active_batches(); + let channel_id = self.oldest_ready_channel(now)?; + let active_occurrence_ids = self + .active_batches + .get(&channel_id) + .map(|active| active.occurrence_ids.clone()); + + let mut event_occurrences = Vec::new(); + let mut cancelled_occurrences = Vec::new(); + if let Some(active_occurrence_ids) = active_occurrence_ids { + // A retry/deferred cohort is reconstructed exclusively from its + // immutable occurrence membership. Fresh arrivals remain in their + // original stores for a later batch and a fresh retry budget. + if let Some(queue) = self.queues.get_mut(&channel_id) { + let mut retained = VecDeque::with_capacity(queue.len()); + while let Some(event) = queue.pop_front() { + if active_occurrence_ids.contains(&event.id) { + event_occurrences.push(event.map(|event| BatchEvent { + event: event.event, + prompt_tag: event.prompt_tag, + received_at: event.received_at, + })); + } else { + retained.push_back(event); } - None => return None, } + *queue = retained; } - }; + if let Some(cancelled) = self.cancelled_batches.get_mut(&channel_id) { + let mut retained = Vec::with_capacity(cancelled.len()); + for event in cancelled.drain(..) { + if active_occurrence_ids.contains(&event.id) { + cancelled_occurrences.push(event); + } else { + retained.push(event); + } + } + *cancelled = retained; + } + } else { + // A fresh cohort preserves the historical merge behavior: up to + // 50 ordinary events plus all interrupted provenance currently + // waiting for this channel. + let queue = self.queues.entry(channel_id).or_default(); + let drain_count = MAX_BATCH_EVENTS.min(queue.len()); + event_occurrences = queue + .drain(..drain_count) + .map(|event| { + event.map(|event| BatchEvent { + event: event.event, + prompt_tag: event.prompt_tag, + received_at: event.received_at, + }) + }) + .collect(); + cancelled_occurrences = self + .cancelled_batches + .remove(&channel_id) + .unwrap_or_default(); + } - // 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 mut events: Vec = queue - .drain(..drain_count) - .map(|qe| BatchEvent { - event: qe.event, - prompt_tag: qe.prompt_tag, - received_at: qe.received_at, - }) - .collect(); // Relay replay delivers stored events newest-first (`ORDER BY // created_at DESC`), but batch consumers — `format_prompt` scope and // reply-anchor selection — require the LAST event to be the newest. // Stable sort: same-second events keep delivery order. - events.sort_by_key(|be| be.event.created_at); + event_occurrences.sort_by_key(|event| event.event.created_at); // Remove the queue entry if now empty. 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 + self.in_flight_deadline); - self.in_flight_batch_sizes.insert(channel_id, events.len()); - - // Merge any cancelled events stored by requeue_as_cancelled(). - let cancelled_events = self + if self .cancelled_batches - .remove(&channel_id) - .unwrap_or_default(); - let cancel_reason = if cancelled_events.is_empty() { + .get(&channel_id) + .is_some_and(Vec::is_empty) + { + self.cancelled_batches.remove(&channel_id); + } + + let cancel_reason = if cancelled_occurrences.is_empty() { self.cancel_reasons.remove(&channel_id); None + } else if self.cancelled_batches.contains_key(&channel_id) { + self.cancel_reasons.get(&channel_id).copied() } else { self.cancel_reasons.remove(&channel_id) }; - Some(FlushBatch { + // A cancelled-only fallback keeps the historical representation: + // provenance moves through `events` plus a cancel reason until a new + // ordinary request arrives and is merged. + if event_occurrences.is_empty() { + event_occurrences = std::mem::take(&mut cancelled_occurrences); + } + + let occurrence_ids = BatchOccurrenceIds { + events: event_occurrences.iter().map(|event| event.id).collect(), + cancelled_events: cancelled_occurrences.iter().map(|event| event.id).collect(), + event_orders: event_occurrences.iter().map(|event| event.order).collect(), + cancelled_event_orders: cancelled_occurrences + .iter() + .map(|event| event.order) + .collect(), + }; + let batch = FlushBatch { channel_id, - events, - cancelled_events, + events: event_occurrences + .into_iter() + .map(|event| event.value) + .collect(), + cancelled_events: cancelled_occurrences + .into_iter() + .map(|event| event.value) + .collect(), cancel_reason, - }) + occurrence_ids, + }; + if self.activate_batch(&batch).is_none() { + tracing::error!( + channel_id = %channel_id, + "dropping batch whose occurrence identity violated the active-cohort invariant" + ); + self.clear_active_batch(channel_id); + return None; + } + self.in_flight_channels.insert(channel_id); + self.in_flight_deadlines + .insert(channel_id, now + self.in_flight_deadline); + self.in_flight_batch_sizes.insert( + channel_id, + batch.events.len() + batch.cancelled_events.len(), + ); + Some(batch) } /// Mark the prompt for `channel_id` as complete. @@ -384,9 +844,9 @@ impl EventQueue { /// 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. + /// active batch identity and retry counter are reset — the next cohort + /// starts fresh. If the batch WAS requeued, its `retry_counts` entry 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) { @@ -394,21 +854,33 @@ impl EventQueue { self.in_flight_deadlines.remove(&channel_id); self.in_flight_batch_sizes.remove(&channel_id); let now = Instant::now(); - match self.retry_after.get(&channel_id) { + let active_batch_id = self.active_batches.get(&channel_id).map(|active| active.id); + match active_batch_id.and_then(|batch_id| { + self.retry_after + .get(&batch_id) + .copied() + .map(|deadline| (batch_id, deadline)) + }) { // Active throttle → channel was requeued; keep retry_counts intact. - Some(&deadline) if deadline > now => {} + 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); - } + // and release the active identity so later arrivals start fresh. + Some(_) | None => self.clear_active_batch(channel_id), } } + /// Release an in-flight batch that was deferred before delivery. + /// + /// Unlike [`Self::mark_complete`], this preserves retry bookkeeping. A + /// temporary resource constraint such as an unavailable exact agent slot + /// is not a successful delivery and must not replenish a poison batch's + /// retry budget. + pub fn mark_deferred(&mut self, channel_id: Uuid) { + self.in_flight_channels.remove(&channel_id); + self.in_flight_deadlines.remove(&channel_id); + self.in_flight_batch_sizes.remove(&channel_id); + } + /// Re-queue a batch of events that failed to process. /// /// Events are pushed back to the **front** of the channel's queue so they @@ -426,10 +898,14 @@ impl EventQueue { /// /// Note: does NOT remove from `in_flight_channels` — caller must call /// `mark_complete` separately. - pub fn requeue(&mut self, batch: FlushBatch) -> Option { + pub fn requeue(&mut self, mut batch: FlushBatch) -> Option { let channel_id = batch.channel_id; + let Some(batch_id) = self.activate_batch(&batch) else { + self.clear_active_batch(channel_id); + return Some(batch); + }; let attempt = { - let count = self.retry_counts.entry(channel_id).or_insert(0); + let count = self.retry_counts.entry(batch_id).or_insert(0); *count += 1; *count }; @@ -443,10 +919,7 @@ impl EventQueue { 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); + self.clear_active_batch(channel_id); return Some(batch); } @@ -472,28 +945,62 @@ impl EventQueue { "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, - event: be.event, - prompt_tag: be.prompt_tag, - received_at: be.received_at, // preserve original timestamp (#46) - }); + let BatchOccurrenceIds { + events: event_occurrence_ids, + cancelled_events: cancelled_occurrence_ids, + event_orders, + cancelled_event_orders, + } = std::mem::take(&mut batch.occurrence_ids); + let FlushBatch { + events, + cancelled_events, + cancel_reason, + .. + } = batch; + let mut event_occurrences: Vec> = events + .into_iter() + .zip(event_occurrence_ids) + .zip(event_orders) + .map(|((event, id), order)| Occurrence::with_identity(id, order, event)) + .collect(); + let cancelled_occurrences_empty = cancelled_occurrence_ids.is_empty(); + let cancelled_occurrences = cancelled_events + .into_iter() + .zip(cancelled_occurrence_ids) + .zip(cancelled_event_orders) + .map(|((event, id), order)| Occurrence::with_identity(id, order, event)); + if cancel_reason.is_some() && cancelled_occurrences_empty { + // A cancelled-only fallback is represented by events + reason. + // Keep that provenance separate so a newer event that arrives + // while the exact agent is unavailable still gets merged with + // interrupted-work framing. + self.cancelled_batches + .entry(channel_id) + .or_default() + .append(&mut event_occurrences); + } else { + let queue = self.queues.entry(channel_id).or_default(); + // Push to front in reverse order so original order is preserved. + for event in event_occurrences.into_iter().rev() { + queue.push_front(event.map(|event| QueuedEvent { + channel_id, + event: event.event, + prompt_tag: event.prompt_tag, + received_at: event.received_at, // preserve original timestamp (#46) + })); + } } - // 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" - ); + if !cancelled_occurrences_empty { + self.cancelled_batches + .entry(channel_id) + .or_default() + .extend(cancelled_occurrences); } - self.retry_after.insert(channel_id, Instant::now() + delay); + if let Some(reason) = cancel_reason { + self.cancel_reasons.insert(channel_id, reason); + } + self.enforce_pending_limit(channel_id, "requeue"); + self.retry_after.insert(batch_id, Instant::now() + delay); None } @@ -505,27 +1012,62 @@ impl EventQueue { /// /// Does NOT set `retry_after`. Does NOT remove from `in_flight_channels` — /// caller must call `mark_complete` separately. - pub fn requeue_preserve_timestamps(&mut self, batch: FlushBatch) { + pub fn requeue_preserve_timestamps(&mut self, mut batch: FlushBatch) { let channel_id = batch.channel_id; - 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, - event: be.event, - prompt_tag: be.prompt_tag, - received_at: be.received_at, - }); + if self.activate_batch(&batch).is_none() { + self.clear_active_batch(channel_id); + return; } - // Enforce per-channel cap: trim newest (back) events if over limit. - while queue.len() > MAX_PENDING_PER_CHANNEL { - queue.pop_back(); - tracing::warn!( - channel_id = %channel_id, - limit = MAX_PENDING_PER_CHANNEL, - "requeue_preserve overflow — dropped newest event to enforce cap" - ); + let BatchOccurrenceIds { + events: event_occurrence_ids, + cancelled_events: cancelled_occurrence_ids, + event_orders, + cancelled_event_orders, + } = std::mem::take(&mut batch.occurrence_ids); + let FlushBatch { + events, + cancelled_events, + cancel_reason, + .. + } = batch; + let mut event_occurrences: Vec> = events + .into_iter() + .zip(event_occurrence_ids) + .zip(event_orders) + .map(|((event, id), order)| Occurrence::with_identity(id, order, event)) + .collect(); + let cancelled_occurrences = cancelled_events + .into_iter() + .zip(cancelled_occurrence_ids.iter().copied()) + .zip(cancelled_event_orders) + .map(|((event, id), order)| Occurrence::with_identity(id, order, event)); + if cancel_reason.is_some() && cancelled_occurrence_ids.is_empty() { + self.cancelled_batches + .entry(channel_id) + .or_default() + .append(&mut event_occurrences); + } else { + let queue = self.queues.entry(channel_id).or_default(); + // Push to front in reverse order so original order is preserved. + for event in event_occurrences.into_iter().rev() { + queue.push_front(event.map(|event| QueuedEvent { + channel_id, + event: event.event, + prompt_tag: event.prompt_tag, + received_at: event.received_at, + })); + } + } + if !cancelled_occurrence_ids.is_empty() { + self.cancelled_batches + .entry(channel_id) + .or_default() + .extend(cancelled_occurrences); + } + if let Some(reason) = cancel_reason { + self.cancel_reasons.insert(channel_id, reason); } + self.enforce_pending_limit(channel_id, "requeue_preserve_timestamps"); } /// Requeue a cancelled batch so its events appear as `cancelled_events` @@ -539,12 +1081,38 @@ impl EventQueue { /// Unlike `requeue_preserve_timestamps`, events are NOT pushed back into /// the generic queue — they are stored separately and merged by /// `flush_next()`. No retry throttle, no backoff. - pub fn requeue_as_cancelled(&mut self, batch: FlushBatch, reason: CancelReason) { - let entry = self.cancelled_batches.entry(batch.channel_id).or_default(); + pub fn requeue_as_cancelled(&mut self, mut batch: FlushBatch, reason: CancelReason) { + let channel_id = batch.channel_id; + if self.activate_batch(&batch).is_none() { + self.clear_active_batch(channel_id); + return; + } + let BatchOccurrenceIds { + events: event_occurrence_ids, + cancelled_events: cancelled_occurrence_ids, + event_orders, + cancelled_event_orders, + } = std::mem::take(&mut batch.occurrence_ids); + let entry = self.cancelled_batches.entry(channel_id).or_default(); // Preserve any already-cancelled events from a prior cancel (double-cancel). - entry.extend(batch.cancelled_events); - entry.extend(batch.events); - self.cancel_reasons.insert(batch.channel_id, reason); + entry.extend( + batch + .cancelled_events + .into_iter() + .zip(cancelled_occurrence_ids) + .zip(cancelled_event_orders) + .map(|((event, id), order)| Occurrence::with_identity(id, order, event)), + ); + entry.extend( + batch + .events + .into_iter() + .zip(event_occurrence_ids) + .zip(event_orders) + .map(|((event, id), order)| Occurrence::with_identity(id, order, event)), + ); + self.cancel_reasons.insert(channel_id, reason); + self.enforce_pending_limit(channel_id, "requeue_as_cancelled"); } /// Returns `true` if any channel has pending events that are not in-flight @@ -574,20 +1142,15 @@ impl EventQueue { ); self.in_flight_channels.remove(&id); self.in_flight_deadlines.remove(&id); + self.clear_active_batch(id); // Symmetric with the flush_next expiry block: recover withheld // goose-native steer events for the expired channel so they are // not permanently orphaned in the side table. self.recover_withheld_for_expired_channel(id); } - self.queues.iter().any(|(id, q)| { - !q.is_empty() - && !self.in_flight_channels.contains(id) - && self.retry_after.get(id).is_none_or(|&t| t <= now) - }) || self - .cancelled_batches - .keys() - .any(|id| !self.in_flight_channels.contains(id)) + self.prune_orphaned_active_batches(); + self.oldest_ready_channel(now).is_some() } /// Number of channels with pending events. @@ -608,7 +1171,37 @@ impl EventQueue { /// `requeue()`'s dead-letter threshold directly. #[cfg(test)] pub fn set_retry_count_for_test(&mut self, channel_id: Uuid, count: u32) { - self.retry_counts.insert(channel_id, count); + let batch_id = self + .active_batches + .entry(channel_id) + .or_insert_with(|| ActiveBatchIdentity { + id: Uuid::new_v4(), + occurrence_ids: HashSet::new(), + }) + .id; + self.retry_counts.insert(batch_id, count); + } + + #[cfg(test)] + fn set_retry_deadline_for_test(&mut self, channel_id: Uuid, deadline: Instant) { + if let Some(active) = self.active_batches.get(&channel_id) { + self.retry_after.insert(active.id, deadline); + } + } + + #[cfg(test)] + fn retry_count_for_test(&self, channel_id: Uuid) -> Option { + self.active_batches + .get(&channel_id) + .and_then(|active| self.retry_counts.get(&active.id)) + .copied() + } + + #[cfg(test)] + fn has_retry_deadline_for_test(&self, channel_id: Uuid) -> bool { + self.active_batches + .get(&channel_id) + .is_some_and(|active| self.retry_after.contains_key(&active.id)) } /// Drop all queued (non-in-flight) events for a channel. @@ -628,10 +1221,11 @@ impl EventQueue { .remove(&channel_id) .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); + self.clear_active_batch(channel_id); self.cancelled_batches.remove(&channel_id); self.cancel_reasons.remove(&channel_id); + self.required_agents.remove(&channel_id); + self.blocked_required_agents.remove(&channel_id); self.withheld_native_steer.remove(&channel_id); // Preserve in_flight_channels AND in_flight_deadlines: the in-flight // task will eventually complete (calling mark_complete) or the deadline @@ -646,6 +1240,69 @@ impl EventQueue { self.in_flight_channels.contains(&channel_id) } + /// Pin a channel's retry to one exact agent slot. + pub fn require_agent(&mut self, channel_id: Uuid, agent_index: usize) { + if let Some(previous) = self.required_agents.insert(channel_id, agent_index) { + if previous != agent_index { + tracing::warn!( + channel_id = %channel_id, + previous_agent = previous, + agent = agent_index, + "replacing conflicting required-agent affinity" + ); + } + } + self.blocked_required_agents.remove(&channel_id); + } + + /// Return the exact slot required by a channel, if any. + pub fn required_agent(&self, channel_id: Uuid) -> Option { + self.required_agents.get(&channel_id).copied() + } + + /// Temporarily exclude a required-agent channel from fair queue selection. + pub fn block_required_agent(&mut self, channel_id: Uuid) { + if self.required_agents.contains_key(&channel_id) { + self.blocked_required_agents.insert(channel_id); + } + } + + /// Wake every channel pinned to `agent_index`. + pub fn release_required_agent(&mut self, agent_index: usize) { + self.blocked_required_agents.retain(|channel_id| { + self.required_agents.get(channel_id).copied() != Some(agent_index) + }); + } + + /// Clear a completed or abandoned exact-slot replay requirement. + pub fn clear_required_agent(&mut self, channel_id: Uuid) { + self.required_agents.remove(&channel_id); + self.blocked_required_agents.remove(&channel_id); + } + + /// Clear exact-slot affinity only after every dependent residue store is + /// drained. The currently completing/dead-lettered in-flight batch is not + /// counted; ordinary, cancelled, and native-steer-withheld residue is. + pub fn clear_required_agent_if_drained(&mut self, channel_id: Uuid) -> bool { + let has_residue = self + .queues + .get(&channel_id) + .is_some_and(|events| !events.is_empty()) + || self + .cancelled_batches + .get(&channel_id) + .is_some_and(|events| !events.is_empty()) + || self + .withheld_native_steer + .get(&channel_id) + .is_some_and(|events| !events.is_empty()); + if has_residue { + return false; + } + self.clear_required_agent(channel_id); + true + } + // ── Goose-native steer withhold (side table) ────────────────────────── // // While a goose-native `_goose/unstable/session/steer` write is in flight @@ -655,8 +1312,8 @@ impl EventQueue { // between `mark_complete` (which clears `in_flight_channels`) and the // ack arriving on the main loop. On `Success` the event is consumed // (`remove_event`); on `Err` / `PromptCompletedNeutral` it is released - // back to the queue front (`release_native_steer`), preserving its - // original `received_at` for FIFO fairness. + // back into the queue (`release_native_steer`), preserving its original + // `received_at` plus enqueue ordinal for FIFO fairness. /// Move a queued event out of `queues[channel_id]` into the side table /// to withhold it from `flush_next` while a goose-native steer is in @@ -670,16 +1327,20 @@ impl EventQueue { /// after `pool.send_steer` returns `Ok(())` and before any watcher task /// is spawned, so the withhold is established before `mark_complete` / /// any subsequent `flush_next` tick can run. - pub fn mark_native_steer_pending(&mut self, channel_id: Uuid, event_id: &str) -> bool { + pub fn mark_native_steer_pending( + &mut self, + channel_id: Uuid, + occurrence_id: EnqueueOccurrenceId, + ) -> bool { let Some(q) = self.queues.get_mut(&channel_id) else { return false; }; - let Some(pos) = q.iter().position(|qe| qe.event.id.to_hex() == event_id) else { + let Some(pos) = q.iter().position(|qe| qe.id == occurrence_id.0) else { + return false; + }; + let Some(qe) = q.remove(pos) else { return false; }; - let qe = q - .remove(pos) - .expect("position came from iter so remove must succeed"); if q.is_empty() { self.queues.remove(&channel_id); } @@ -690,8 +1351,20 @@ impl EventQueue { true } - /// Release a single withheld event back to the front of - /// `queues[channel_id]`, preserving its original `received_at`. + fn insert_queued_occurrence_by_arrival( + queue: &mut VecDeque>, + occurrence: Occurrence, + ) { + let key = (occurrence.received_at, occurrence.order); + let insert_at = queue + .iter() + .position(|queued| (queued.received_at, queued.order) > key) + .unwrap_or(queue.len()); + queue.insert(insert_at, occurrence); + } + + /// Release a single withheld event back into `queues[channel_id]`, + /// preserving its original `received_at` and enqueue ordinal. /// /// Called on `SteerAck::Err(_)` and `SteerAck::PromptCompletedNeutral` /// (delivery unknown after prompt completion; restoring queued event @@ -700,33 +1373,24 @@ impl EventQueue { /// /// Push-to-front matches the discipline of `requeue_preserve_timestamps` /// at line 453, preserving fairness across channels. - pub fn release_native_steer(&mut self, channel_id: Uuid, event_id: &str) { + pub fn release_native_steer(&mut self, channel_id: Uuid, occurrence_id: EnqueueOccurrenceId) { let Some(entries) = self.withheld_native_steer.get_mut(&channel_id) else { return; }; - let Some(pos) = entries - .iter() - .position(|qe| qe.event.id.to_hex() == event_id) - else { + let Some(pos) = entries.iter().position(|qe| qe.id == occurrence_id.0) else { return; }; let qe = entries.remove(pos); if entries.is_empty() { self.withheld_native_steer.remove(&channel_id); } - // Push to FRONT so original `received_at` keeps the event at the head - // of the channel's queue. Per-channel cap is enforced below in case - // a flood of events arrived during the ack window. + // Merge by the preserved arrival timestamp. Multiple steer + // acknowledgements can resolve independently, so repeated push_front + // would reverse their FIFO order and hide an older occurrence behind a + // newer one from cross-channel fairness selection. let queue = self.queues.entry(channel_id).or_default(); - queue.push_front(qe); - while queue.len() > MAX_PENDING_PER_CHANNEL { - queue.pop_back(); - tracing::warn!( - channel_id = %channel_id, - limit = MAX_PENDING_PER_CHANNEL, - "release_native_steer overflow — dropped newest event to enforce cap" - ); - } + Self::insert_queued_occurrence_by_arrival(queue, qe); + self.enforce_pending_limit(channel_id, "release_native_steer"); } /// Drop a specific event by id from both the side table and the main @@ -735,19 +1399,39 @@ impl EventQueue { /// Called on `SteerAck::Success` — the agent received the steer, so the /// event has been "delivered" via the non-cancelling path and must not /// be redelivered via normal dispatch. Idempotent across both stores. - pub fn remove_event(&mut self, channel_id: Uuid, event_id: &str) { + pub fn remove_native_steer(&mut self, channel_id: Uuid, occurrence_id: EnqueueOccurrenceId) { + let mut removed = false; if let Some(entries) = self.withheld_native_steer.get_mut(&channel_id) { - entries.retain(|qe| qe.event.id.to_hex() != event_id); + entries.retain(|event| { + if event.id == occurrence_id.0 { + removed = true; + false + } else { + true + } + }); if entries.is_empty() { self.withheld_native_steer.remove(&channel_id); } } if let Some(q) = self.queues.get_mut(&channel_id) { - q.retain(|qe| qe.event.id.to_hex() != event_id); + q.retain(|event| { + if event.id == occurrence_id.0 { + removed = true; + false + } else { + true + } + }); if q.is_empty() { self.queues.remove(&channel_id); } } + if removed { + if let Some(active) = self.active_batches.get_mut(&channel_id) { + active.occurrence_ids.remove(&occurrence_id.0); + } + } } /// Bulk-release every withheld event for `channel_id` back to the queue @@ -760,26 +1444,21 @@ impl EventQueue { /// events were never delivered to the agent, so normal dispatch must /// have a chance to deliver them. /// - /// Iterates the stored entries in reverse so per-entry `push_front` - /// composes to original-FIFO order at the queue front (same discipline - /// as `requeue_preserve_timestamps` at line 453). + /// Each occurrence is merged by its `(received_at, enqueue ordinal)` key, + /// so arbitrary acknowledgement/recovery order cannot reverse equal-time + /// events or hide an older event behind a newer queue head. fn recover_withheld_for_expired_channel(&mut self, channel_id: Uuid) { let Some(entries) = self.withheld_native_steer.remove(&channel_id) else { return; }; let n = entries.len(); - let queue = self.queues.entry(channel_id).or_default(); - for qe in entries.into_iter().rev() { - queue.push_front(qe); - } - while queue.len() > MAX_PENDING_PER_CHANNEL { - queue.pop_back(); - tracing::warn!( - channel_id = %channel_id, - limit = MAX_PENDING_PER_CHANNEL, - "withheld-steer recovery overflow — dropped newest event to enforce cap" - ); + { + let queue = self.queues.entry(channel_id).or_default(); + for occurrence in entries { + Self::insert_queued_occurrence_by_arrival(queue, occurrence); + } } + self.enforce_pending_limit(channel_id, "recover_withheld_for_expired_channel"); tracing::warn!( channel_id = %channel_id, recovered = n, @@ -790,15 +1469,14 @@ impl EventQueue { /// 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. + /// Removes `retry_after` entries whose deadline has already passed, prunes + /// active identities with no remaining members, and cleans up orphaned + /// `retry_counts` entries whose immutable batch no longer exists. Without + /// this, completed retry cycles could leak metadata 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 + /// The in-flight guard is critical: a batch whose throttle expired and + /// whose queued members were 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 @@ -807,13 +1485,14 @@ impl EventQueue { 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) + self.prune_orphaned_active_batches(); + let active_batch_ids: HashSet = self + .active_batches + .values() + .map(|active| active.id) + .collect(); + self.retry_counts.retain(|batch_id, _| { + self.retry_after.contains_key(batch_id) || active_batch_ids.contains(batch_id) }); } } @@ -1875,6 +2554,7 @@ mod tests { }], cancelled_events: vec![], cancel_reason: None, + occurrence_ids: Default::default(), }; let prompt = format_prompt(&batch, &FormatPromptArgs::default()).join("\n\n"); @@ -1909,6 +2589,7 @@ mod tests { received_at: Instant::now(), }], cancel_reason: reason, + occurrence_ids: Default::default(), } } @@ -2047,6 +2728,7 @@ mod tests { received_at: Instant::now(), }], cancel_reason: Some(CancelReason::Steer), + occurrence_ids: Default::default(), }; let prompt = format_prompt(&batch, &FormatPromptArgs::default()).join("\n\n"); assert!(prompt.contains("New messages — arrived while you were working — 2 events]")); @@ -2097,6 +2779,7 @@ mod tests { received_at: Instant::now(), }], cancel_reason: Some(CancelReason::Steer), + occurrence_ids: Default::default(), }; let prompt = format_prompt(&batch, &FormatPromptArgs::default()).join("\n\n"); @@ -2136,8 +2819,8 @@ mod tests { queue.requeue(batch); queue.mark_complete(ch); - // retry_after is set, so manually clear it for this test. - queue.retry_after.remove(&ch); + // Expire the immutable batch throttle for this test. + queue.set_retry_deadline_for_test(ch, Instant::now()); // Should be able to flush again and get the same events in order. let batch2 = queue.flush_next().unwrap(); @@ -2197,6 +2880,7 @@ mod tests { ], cancelled_events: vec![], cancel_reason: None, + occurrence_ids: Default::default(), }; let prompt = format_prompt(&batch, &FormatPromptArgs::default()).join("\n\n"); @@ -2225,6 +2909,7 @@ mod tests { }], cancelled_events: vec![], cancel_reason: None, + occurrence_ids: Default::default(), }; let prompt = format_prompt(&batch, &FormatPromptArgs::default()).join("\n\n"); @@ -2248,6 +2933,7 @@ mod tests { }], cancelled_events: vec![], cancel_reason: None, + occurrence_ids: Default::default(), }; let core = "[Agent Memory — core]\nbe helpful"; let prompt = format_prompt( @@ -2280,6 +2966,7 @@ mod tests { }], cancelled_events: vec![], cancel_reason: None, + occurrence_ids: Default::default(), }; let prompt = format_prompt( &batch, @@ -2310,6 +2997,7 @@ mod tests { }], cancelled_events: vec![], cancel_reason: None, + occurrence_ids: Default::default(), }; let core = "[Agent Memory — core]\nbe helpful"; let prompt = format_prompt( @@ -2337,6 +3025,7 @@ mod tests { }], cancelled_events: vec![], cancel_reason: None, + occurrence_ids: Default::default(), }; // format_prompt no longer accepts or emits base_prompt/system_prompt. @@ -2361,6 +3050,7 @@ mod tests { }], cancelled_events: vec![], cancel_reason: None, + occurrence_ids: Default::default(), }; let core = "[Agent Memory — core]\nremember this"; @@ -2417,6 +3107,7 @@ mod tests { }], cancelled_events: vec![], cancel_reason: None, + occurrence_ids: Default::default(), }; let prompt = format_prompt( @@ -2455,6 +3146,7 @@ mod tests { }], cancelled_events: vec![], cancel_reason: None, + occurrence_ids: Default::default(), }; let ctx = ConversationContext::Thread { @@ -2706,11 +3398,225 @@ mod tests { q.requeue_preserve_timestamps(batch); q.mark_complete(ch); - // No retry_after — channel should be immediately flushable. - assert!(!q.retry_after.contains_key(&ch)); + // No retry deadline — channel should be immediately flushable. + assert!(!q.has_retry_deadline_for_test(ch)); assert!(q.flush_next().is_some()); } + #[test] + fn exact_slot_deferral_preserves_retry_budget() { + let mut q = EventQueue::new(DedupMode::Queue); + let channel_id = Uuid::new_v4(); + q.push(make_queued(channel_id, "retry-on-exact-slot")); + let batch = q.flush_next().expect("initial batch"); + + q.set_retry_count_for_test(channel_id, MAX_RETRIES); + q.set_retry_deadline_for_test(channel_id, Instant::now() - Duration::from_secs(1)); + q.requeue_preserve_timestamps(batch); + q.mark_deferred(channel_id); + + assert_eq!( + q.retry_count_for_test(channel_id), + Some(MAX_RETRIES), + "resource deferral must not reset the delivery retry budget" + ); + let retried = q.flush_next().expect("deferred batch remains ready"); + assert!( + q.requeue(retried).is_some(), + "the next delivery failure must exhaust the preserved budget" + ); + } + + #[test] + fn fresh_arrival_does_not_inherit_poison_batch_dead_letter() { + let mut q = EventQueue::new(DedupMode::Queue); + let channel_id = Uuid::new_v4(); + q.push(make_queued(channel_id, "poison")); + let poison = q.flush_next().expect("initial poison batch"); + assert!(q.requeue(poison).is_none()); + q.mark_complete(channel_id); + + q.push(make_queued(channel_id, "fresh")); + q.set_retry_count_for_test(channel_id, MAX_RETRIES); + q.set_retry_deadline_for_test(channel_id, Instant::now()); + + let retry = q.flush_next().expect("poison retry after backoff"); + assert_eq!( + retry.events.len(), + 1, + "fresh same-channel traffic must not join the immutable retry batch" + ); + assert_eq!(retry.events[0].event.content, "poison"); + let dead = q + .requeue(retry) + .expect("poison batch should exhaust its retry budget"); + assert_eq!(dead.events.len(), 1); + q.mark_complete(channel_id); + + let fresh = q + .flush_next() + .expect("fresh arrival must survive poison disposal"); + assert_eq!(fresh.events.len(), 1); + assert_eq!(fresh.events[0].event.content, "fresh"); + assert!( + q.requeue(fresh).is_none(), + "fresh work must begin with a new retry budget" + ); + } + + #[test] + fn identical_event_id_fresh_occurrence_does_not_join_active_retry_cohort() { + let mut q = EventQueue::new(DedupMode::Queue); + let channel_id = Uuid::new_v4(); + let mut original = make_queued(channel_id, "same signed event"); + original.prompt_tag = "original-occurrence-1".into(); + let original_event_id = original.event.id; + let mut second_original = original.clone(); + second_original.prompt_tag = "original-occurrence-2".into(); + second_original.received_at = Instant::now(); + let mut fresh_duplicate = original.clone(); + fresh_duplicate.prompt_tag = "fresh-occurrence".into(); + fresh_duplicate.received_at = Instant::now(); + + q.push(original); + q.push(second_original); + let poison = q.flush_next().expect("initial occurrences"); + assert_eq!(poison.events.len(), 2); + assert_eq!(poison.events[0].event.id, original_event_id); + assert_eq!(poison.events[0].prompt_tag, "original-occurrence-1"); + assert_eq!(poison.events[1].prompt_tag, "original-occurrence-2"); + assert!(q.requeue(poison).is_none()); + q.mark_complete(channel_id); + + q.push(fresh_duplicate); + q.set_retry_count_for_test(channel_id, MAX_RETRIES); + q.set_retry_deadline_for_test(channel_id, Instant::now()); + + let retry = q.flush_next().expect("original occurrence retry"); + assert_eq!( + retry.events.len(), + 2, + "a fresh enqueue occurrence with the same signed event ID must not join the active cohort" + ); + assert_eq!(retry.events[0].event.id, original_event_id); + assert_eq!(retry.events[1].event.id, original_event_id); + assert_eq!(retry.events[0].prompt_tag, "original-occurrence-1"); + assert_eq!(retry.events[1].prompt_tag, "original-occurrence-2"); + let dead = q + .requeue(retry) + .expect("the original occurrences should exhaust their retry budget"); + assert_eq!(dead.events.len(), 2); + assert_eq!(dead.events[0].prompt_tag, "original-occurrence-1"); + assert_eq!(dead.events[1].prompt_tag, "original-occurrence-2"); + q.mark_complete(channel_id); + + let fresh = q + .flush_next() + .expect("the fresh duplicate occurrence must survive"); + assert_eq!(fresh.events.len(), 1); + assert_eq!(fresh.events[0].event.id, original_event_id); + assert_eq!(fresh.events[0].prompt_tag, "fresh-occurrence"); + assert!( + q.requeue(fresh).is_none(), + "the fresh occurrence must start with a new retry budget" + ); + } + + #[test] + fn identical_event_id_fresh_occurrence_does_not_join_deferred_cancelled_provenance() { + let mut q = EventQueue::new(DedupMode::Queue); + let channel_id = Uuid::new_v4(); + let mut original = make_queued(channel_id, "same signed cancellation event"); + original.prompt_tag = "original-cancelled-occurrence".into(); + let original_event_id = original.event.id; + let mut fresh_duplicate = original.clone(); + fresh_duplicate.prompt_tag = "fresh-after-cancel".into(); + fresh_duplicate.received_at = Instant::now(); + + q.push(original); + let interrupted = q.flush_next().expect("initial interrupted occurrence"); + q.requeue_as_cancelled(interrupted, CancelReason::Interrupt); + q.mark_complete(channel_id); + + let fallback = q + .flush_next() + .expect("cancelled-only fallback must become its own cohort"); + assert_eq!(fallback.events.len(), 1); + assert!(fallback.cancelled_events.is_empty()); + assert_eq!(fallback.cancel_reason, Some(CancelReason::Interrupt)); + q.requeue_preserve_timestamps(fallback); + q.mark_deferred(channel_id); + + q.push(fresh_duplicate); + let resumed = q + .flush_next() + .expect("deferred cancelled provenance must resume"); + assert_eq!( + resumed.events.len() + resumed.cancelled_events.len(), + 1, + "a fresh duplicate occurrence must not join deferred cancellation provenance" + ); + assert_eq!(resumed.events[0].event.id, original_event_id); + assert_eq!( + resumed.events[0].prompt_tag, + "original-cancelled-occurrence" + ); + assert_eq!(resumed.cancel_reason, Some(CancelReason::Interrupt)); + q.mark_complete(channel_id); + + let fresh = q + .flush_next() + .expect("fresh duplicate must remain for a later cohort"); + assert_eq!(fresh.events.len(), 1); + assert_eq!(fresh.events[0].event.id, original_event_id); + assert_eq!(fresh.events[0].prompt_tag, "fresh-after-cancel"); + } + + #[test] + fn missing_occurrence_identity_fails_closed_without_inference() { + let mut q = EventQueue::new(DedupMode::Queue); + let channel_id = Uuid::new_v4(); + q.push(make_queued(channel_id, "must-retain-original-identity")); + let mut batch = q.flush_next().expect("initial batch"); + + batch.occurrence_ids.events.clear(); + let rejected = q + .requeue(batch) + .expect("identity-free batch must be returned instead of reconstructed"); + assert_eq!(rejected.events.len(), 1); + assert_eq!(pending_count(&q), 0, "rejected work must not be replayed"); + assert!( + !q.active_batches.contains_key(&channel_id), + "fail-closed rejection must clear stale retry identity" + ); + } + + #[test] + fn duplicate_occurrence_identity_fails_closed() { + let mut q = EventQueue::new(DedupMode::Queue); + let channel_id = Uuid::new_v4(); + q.push(make_queued(channel_id, "first")); + let mut batch = q.flush_next().expect("initial batch"); + + batch.events.push(BatchEvent { + event: make_event("second"), + prompt_tag: "synthetic-corruption".into(), + received_at: Instant::now(), + }); + let duplicated_id = batch.occurrence_ids.events[0]; + batch.occurrence_ids.events.push(duplicated_id); + + let rejected = q + .requeue(batch) + .expect("duplicate occurrence identity must be returned, not replayed"); + assert_eq!(rejected.events.len(), 2); + assert_eq!(pending_count(&q), 0, "rejected work must not be replayed"); + assert!( + !q.active_batches.contains_key(&channel_id), + "fail-closed rejection must clear stale retry identity" + ); + } + #[test] fn test_requeue_preserve_timestamps_enforces_cap() { let mut q = EventQueue::new(DedupMode::Queue); @@ -2748,7 +3654,7 @@ mod tests { } #[test] - fn test_requeue_preserve_timestamps_overflow_keeps_requeued_events() { + fn test_requeue_preserve_timestamps_overflow_drops_oldest_events() { let mut q = EventQueue::new(DedupMode::Queue); let ch = Uuid::new_v4(); @@ -2757,7 +3663,7 @@ mod tests { q.push(make_queued(ch, &format!("original-{i}"))); } - // Flush a batch — these are the "requeued" events we want to survive. + // Flush the oldest batch. let batch = q.flush_next().expect("should flush"); let batch_size = batch.events.len(); @@ -2766,19 +3672,23 @@ mod tests { q.push(make_queued(ch, &format!("new-{i}"))); } - // Capture the content of the first requeued event for verification. - let requeued_first_content = batch.events[0].event.content.to_string(); - - // Requeue — older events go to front, overflow trims from back (newest). + // Requeue with preserved timestamps. The shared retention policy drops + // the chronologically oldest entries, even when they came from the + // just-requeued batch, so newer work is not sacrificed. q.requeue_preserve_timestamps(batch); q.mark_complete(ch); - // The requeued events should be at the front of the queue. + assert!( + q.queues + .get(&ch) + .is_some_and(|events| events.iter().any(|event| event.event.content == "new-49")), + "newest queued work must survive oldest-first overflow" + ); let batch2 = q.flush_next().expect("should flush after requeue"); assert_eq!( batch2.events[0].event.content.to_string(), - requeued_first_content, - "requeued events should be at the front (oldest), not trimmed" + "original-50", + "the oldest requeued batch must be evicted before newer retained work" ); } @@ -2812,8 +3722,7 @@ mod tests { ); // Manually expire the retry_after to simulate time passing. - q.retry_after - .insert(ch, Instant::now() - Duration::from_secs(1)); + q.set_retry_deadline_for_test(ch, Instant::now() - Duration::from_secs(1)); assert!( q.has_flushable_work(), "expired throttle should be flushable" @@ -2827,8 +3736,7 @@ mod tests { q.push(make_queued(ch, "poison")); for attempt in 1..=MAX_RETRIES { - q.retry_after - .insert(ch, Instant::now() - Duration::from_secs(1)); + q.set_retry_deadline_for_test(ch, Instant::now() - Duration::from_secs(1)); let batch = q.flush_next().expect("flush"); assert!( q.requeue(batch).is_none(), @@ -2838,16 +3746,15 @@ mod tests { } // The MAX_RETRIES+1'th failure dead-letters: batch is returned. - q.retry_after - .insert(ch, Instant::now() - Duration::from_secs(1)); + q.set_retry_deadline_for_test(ch, Instant::now() - Duration::from_secs(1)); let batch = q.flush_next().expect("flush"); let dead = q.requeue(batch).expect("should dead-letter"); assert_eq!(dead.channel_id, ch); assert_eq!(dead.events.len(), 1); q.mark_complete(ch); // Retry state is cleared so fresh traffic isn't throttled. - assert!(!q.retry_counts.contains_key(&ch)); - assert!(!q.retry_after.contains_key(&ch)); + assert_eq!(q.retry_count_for_test(ch), None); + assert!(!q.has_retry_deadline_for_test(ch)); } #[test] @@ -2872,8 +3779,7 @@ mod tests { assert_eq!(batch2.channel_id, ch2); // After retry_after expires, ch should be flushable again. - q.retry_after - .insert(ch, Instant::now() - Duration::from_secs(1)); + q.set_retry_deadline_for_test(ch, Instant::now() - Duration::from_secs(1)); q.mark_complete(ch2); let batch3 = q .flush_next() @@ -2972,6 +3878,7 @@ mod tests { }], cancelled_events: vec![], cancel_reason: None, + occurrence_ids: Default::default(), }; let ci = PromptChannelInfo { name: "engineering".into(), @@ -3003,6 +3910,7 @@ mod tests { }], cancelled_events: vec![], cancel_reason: None, + occurrence_ids: Default::default(), }; let ci = PromptChannelInfo { name: "DM".into(), @@ -3041,6 +3949,7 @@ mod tests { }], cancelled_events: vec![], cancel_reason: None, + occurrence_ids: Default::default(), }; let prompt = format_prompt(&batch, &FormatPromptArgs::default()).join("\n\n"); @@ -3069,6 +3978,7 @@ mod tests { }], cancelled_events: vec![], cancel_reason: None, + occurrence_ids: Default::default(), }; let ctx = ConversationContext::Thread { messages: vec![ @@ -3113,6 +4023,7 @@ mod tests { }], cancelled_events: vec![], cancel_reason: None, + occurrence_ids: Default::default(), }; let ci = PromptChannelInfo { name: "DM".into(), @@ -3162,6 +4073,7 @@ mod tests { }], cancelled_events: vec![], cancel_reason: None, + occurrence_ids: Default::default(), }; let ctx = ConversationContext::Thread { messages: vec![ContextMessage { @@ -3369,6 +4281,7 @@ mod tests { }], cancelled_events: vec![], cancel_reason: None, + occurrence_ids: Default::default(), }; let ci = PromptChannelInfo { name: "DM".into(), @@ -3426,6 +4339,7 @@ mod tests { }], cancelled_events: vec![], cancel_reason: None, + occurrence_ids: Default::default(), }; let ci = PromptChannelInfo { name: "DM".into(), @@ -3466,6 +4380,7 @@ mod tests { }], cancelled_events: vec![], cancel_reason: None, + occurrence_ids: Default::default(), }; let prompt = format_prompt(&batch, &FormatPromptArgs::default()).join("\n\n"); @@ -3490,6 +4405,7 @@ mod tests { }], cancelled_events: vec![], cancel_reason: None, + occurrence_ids: Default::default(), }; let prompt = format_prompt(&batch, &FormatPromptArgs::default()).join("\n\n"); @@ -3513,6 +4429,7 @@ mod tests { }], cancelled_events: vec![], cancel_reason: None, + occurrence_ids: Default::default(), }; let prompt = format_prompt(&batch, &FormatPromptArgs::default()).join("\n\n"); @@ -3603,25 +4520,25 @@ mod tests { 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)); + assert!(q.has_retry_deadline_for_test(ch)); + assert_eq!(q.retry_count_for_test(ch), Some(1)); // 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)); + q.set_retry_deadline_for_test(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)); + assert_eq!(q.retry_count_for_test(ch), None); // 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); + // with no active identity, throttle, queue, or in-flight owner. + let stale_batch_id = Uuid::new_v4(); + q.retry_counts.insert(stale_batch_id, 3); q.compact_expired_state(); assert!( - !q.retry_counts.contains_key(&ch), + !q.retry_counts.contains_key(&stale_batch_id), "orphaned retry_counts should be removed" ); } @@ -3638,8 +4555,7 @@ mod tests { 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)); + q.set_retry_deadline_for_test(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)); @@ -3649,7 +4565,7 @@ mod tests { // may fail and requeue, which needs the existing count. q.compact_expired_state(); assert!( - q.retry_counts.contains_key(&ch), + q.retry_count_for_test(ch).is_some(), "retry_counts must survive while channel is in-flight" ); } @@ -3659,13 +4575,17 @@ mod tests { let mut q = EventQueue::new(DedupMode::Queue); let ch = Uuid::new_v4(); - // Manually set up: retry_counts exists, queue is non-empty, no throttle. + // Set up an immutable deferred cohort with queued members and no + // throttle; maintenance must retain its retry budget. q.push(make_queued(ch, "msg1")); - q.retry_counts.insert(ch, 2); + let batch = q.flush_next().expect("active batch"); + q.set_retry_count_for_test(ch, 2); + q.requeue_preserve_timestamps(batch); + q.mark_deferred(ch); q.compact_expired_state(); assert!( - q.retry_counts.contains_key(&ch), + q.retry_count_for_test(ch).is_some(), "retry_counts should survive when queue is non-empty" ); } @@ -3800,6 +4720,64 @@ mod tests { ); } + #[test] + fn cancelled_only_and_ordinary_work_share_oldest_fifo() { + let mut q = EventQueue::new(DedupMode::Queue); + let cancelled_channel = Uuid::from_u128(1); + let ordinary_channel = Uuid::from_u128(2); + + q.push(make_queued_at( + cancelled_channel, + "older-interrupted-work", + Duration::from_secs(20), + )); + let interrupted = q.flush_next().expect("interrupted batch"); + q.requeue_as_cancelled(interrupted, CancelReason::Interrupt); + q.mark_complete(cancelled_channel); + + q.push(make_queued_at( + ordinary_channel, + "newer-ordinary-work", + Duration::from_secs(1), + )); + + let first = q.flush_next().expect("oldest ready work"); + assert_eq!( + first.channel_id, cancelled_channel, + "cancelled-only work must compete in the same FIFO as ordinary work" + ); + } + + #[test] + fn mixed_store_fifo_breaks_timestamp_ties_by_channel_id() { + let mut q = EventQueue::new(DedupMode::Queue); + let cancelled_channel = Uuid::from_u128(1); + let ordinary_channel = Uuid::from_u128(2); + let received_at = Instant::now() - Duration::from_secs(5); + + q.push(QueuedEvent { + channel_id: cancelled_channel, + event: make_event("interrupted-at-tie"), + received_at, + prompt_tag: "test".into(), + }); + let interrupted = q.flush_next().expect("interrupted batch"); + q.requeue_as_cancelled(interrupted, CancelReason::Interrupt); + q.mark_complete(cancelled_channel); + q.push(QueuedEvent { + channel_id: ordinary_channel, + event: make_event("ordinary-at-tie"), + received_at, + prompt_tag: "test".into(), + }); + + assert_eq!( + q.flush_next().expect("tied work").channel_id, + cancelled_channel, + "equal timestamps must have a deterministic channel-id tie break" + ); + } + #[test] fn test_drain_channel_clears_cancelled_batches() { let mut q = EventQueue::new(DedupMode::Queue); @@ -3862,6 +4840,161 @@ mod tests { ); } + #[test] + fn repeated_double_cancel_history_respects_combined_channel_cap() { + let mut q = EventQueue::new(DedupMode::Queue); + let ch = Uuid::new_v4(); + + q.push(make_queued(ch, "original")); + let original = q.flush_next().expect("original batch"); + q.requeue_as_cancelled(original, CancelReason::Interrupt); + q.mark_complete(ch); + + for i in 0..(MAX_PENDING_PER_CHANNEL + 25) { + q.push(make_queued(ch, &format!("new-{i}"))); + let merged = q.flush_next().expect("merged cancellation replay"); + q.requeue_as_cancelled(merged, CancelReason::Steer); + q.mark_complete(ch); + } + + let ordinary = q.queues.get(&ch).map_or(0, VecDeque::len); + let cancelled = q.cancelled_batches.get(&ch).map_or(0, Vec::len); + assert_eq!( + ordinary + cancelled, + MAX_PENDING_PER_CHANNEL, + "ordinary and cancelled work must share one channel capacity" + ); + assert!( + q.cancelled_batches + .get(&ch) + .is_some_and(|events| events.iter().any(|event| { + event.event.content == format!("new-{}", MAX_PENDING_PER_CHANNEL + 24) + })), + "the newest cancelled work must survive oldest-first eviction" + ); + assert_eq!(q.cancel_reasons.get(&ch), Some(&CancelReason::Steer)); + } + + #[test] + fn new_pushes_replace_old_cancelled_history_within_combined_cap() { + let mut q = EventQueue::new(DedupMode::Queue); + let ch = Uuid::new_v4(); + let cancelled_events = (0..MAX_PENDING_PER_CHANNEL) + .map(|i| { + let queued = make_queued(ch, &format!("cancelled-{i}")); + BatchEvent { + event: queued.event, + prompt_tag: queued.prompt_tag, + received_at: queued.received_at, + } + }) + .collect(); + q.requeue_as_cancelled( + FlushBatch { + channel_id: ch, + events: cancelled_events, + cancelled_events: vec![], + cancel_reason: None, + occurrence_ids: BatchOccurrenceIds::for_test(MAX_PENDING_PER_CHANNEL, 0), + }, + CancelReason::Interrupt, + ); + + q.push(make_queued(ch, "newest")); + let ordinary = q.queues.get(&ch).map_or(0, VecDeque::len); + let cancelled = q.cancelled_batches.get(&ch).map_or(0, Vec::len); + assert_eq!( + ordinary + cancelled, + MAX_PENDING_PER_CHANNEL, + "a new push must not add capacity beyond retained cancellation history" + ); + assert!( + q.queues + .get(&ch) + .is_some_and(|events| events.iter().any(|event| event.event.content == "newest")), + "new work must survive oldest-first eviction" + ); + assert!( + q.cancelled_batches.get(&ch).is_some_and(|events| events + .iter() + .all(|event| event.event.content != "cancelled-0")), + "the oldest cancelled history must be evicted first" + ); + + for i in 0..(MAX_PENDING_PER_CHANNEL - 1) { + q.push(make_queued(ch, &format!("replacement-{i}"))); + } + let ordinary = q.queues.get(&ch).map_or(0, VecDeque::len); + let cancelled = q.cancelled_batches.get(&ch).map_or(0, Vec::len); + assert_eq!(ordinary + cancelled, MAX_PENDING_PER_CHANNEL); + assert!( + !q.cancelled_batches.contains_key(&ch), + "fully evicted cancellation history must remove its empty side-table entry" + ); + assert!( + !q.cancel_reasons.contains_key(&ch), + "fully evicted cancellation history must remove its stale reason" + ); + } + + #[test] + fn retry_paths_bound_cancelled_provenance_with_ordinary_work() { + for preserve_timestamps in [false, true] { + let mut q = EventQueue::new(DedupMode::Queue); + let ch = Uuid::new_v4(); + let old = make_queued_at(ch, "old-cancelled", Duration::from_secs(10)); + let old_batch_event = BatchEvent { + event: old.event, + prompt_tag: old.prompt_tag, + received_at: old.received_at, + }; + q.cancelled_batches.insert( + ch, + (0..MAX_PENDING_PER_CHANNEL - 1) + .map(|_| Occurrence::new(old_batch_event.clone())) + .collect(), + ); + q.cancel_reasons.insert(ch, CancelReason::Interrupt); + + let newest = make_queued(ch, "newest-retry"); + let batch = FlushBatch { + channel_id: ch, + events: vec![BatchEvent { + event: newest.event, + prompt_tag: newest.prompt_tag, + received_at: newest.received_at, + }], + cancelled_events: vec![old_batch_event], + cancel_reason: Some(CancelReason::Steer), + occurrence_ids: BatchOccurrenceIds::for_test(1, 1), + }; + if preserve_timestamps { + q.requeue_preserve_timestamps(batch); + } else { + assert!(q.requeue(batch).is_none()); + } + + let ordinary = q.queues.get(&ch).map_or(0, VecDeque::len); + let cancelled = q.cancelled_batches.get(&ch).map_or(0, Vec::len); + assert_eq!( + ordinary + cancelled, + MAX_PENDING_PER_CHANNEL, + "retry path preserve_timestamps={preserve_timestamps} exceeded the shared cap" + ); + assert!( + q.queues.get(&ch).is_some_and(|events| events + .iter() + .any(|event| event.event.content == "newest-retry")), + "retry path preserve_timestamps={preserve_timestamps} evicted newest work" + ); + assert_eq!( + q.cancel_reasons.get(&ch), + Some(&CancelReason::Steer), + "retained cancellation history must retain the latest framing reason" + ); + } + } + #[test] fn test_reply_instruction_present_for_channel_thread_reply() { let ch = Uuid::new_v4(); @@ -3879,6 +5012,7 @@ mod tests { }], cancelled_events: vec![], cancel_reason: None, + occurrence_ids: Default::default(), }; // No profile lookup → sender treated as human → human-facing thread @@ -3921,6 +5055,7 @@ mod tests { }], cancelled_events: vec![], cancel_reason: None, + occurrence_ids: Default::default(), }; let ci = PromptChannelInfo { name: "DM".into(), @@ -3955,6 +5090,7 @@ mod tests { }], cancelled_events: vec![], cancel_reason: None, + occurrence_ids: Default::default(), }; // Top-level human message (no lookup → human): the reply opens a new @@ -3984,6 +5120,7 @@ mod tests { }], cancelled_events: vec![], cancel_reason: None, + occurrence_ids: Default::default(), }; let ci = PromptChannelInfo { name: "DM".into(), @@ -4026,6 +5163,7 @@ mod tests { }], cancelled_events: vec![], cancel_reason: None, + occurrence_ids: Default::default(), }; // Human-facing (no lookup) deep reply: anchor to the thread ROOT to @@ -4062,6 +5200,7 @@ mod tests { }], cancelled_events: vec![], cancel_reason: None, + occurrence_ids: Default::default(), }; let prompt = format_prompt(&batch, &FormatPromptArgs::default()).join("\n\n"); @@ -4104,6 +5243,7 @@ mod tests { ], cancelled_events: vec![], cancel_reason: None, + occurrence_ids: Default::default(), }; // Scope derives from the last (threaded) event; human-facing → anchor @@ -4141,6 +5281,7 @@ mod tests { ], cancelled_events: vec![], cancel_reason: None, + occurrence_ids: Default::default(), }; // Last event is top-level and human-facing → opens a new thread @@ -4167,6 +5308,7 @@ mod tests { }], cancelled_events: vec![], cancel_reason: None, + occurrence_ids: Default::default(), } } @@ -4280,10 +5422,9 @@ mod tests { let ch = Uuid::new_v4(); let qe = make_queued(ch, "hello"); - let event_id = qe.event.id.to_hex(); - q.push(qe); + let occurrence_id = q.push(qe).expect("accepted occurrence"); - assert!(q.mark_native_steer_pending(ch, &event_id)); + assert!(q.mark_native_steer_pending(ch, occurrence_id)); assert!( q.flush_next().is_none(), @@ -4297,6 +5438,126 @@ mod tests { assert_eq!(q.withheld_native_steer.get(&ch).map(|v| v.len()), Some(1)); } + #[test] + fn native_steer_withhold_missing_occurrence_is_a_lossless_no_op() { + let mut q = EventQueue::new(DedupMode::Queue); + let ch = Uuid::new_v4(); + let queued = q.push(make_queued(ch, "hello")).expect("queued occurrence"); + + assert!(!q.mark_native_steer_pending(ch, EnqueueOccurrenceId(Uuid::new_v4()))); + + let batch = q.flush_next().expect("original occurrence remains queued"); + assert_eq!(batch.occurrence_ids.events, vec![queued.0]); + assert_eq!(batch.events[0].event.content, "hello"); + assert!(!q.withheld_native_steer.contains_key(&ch)); + } + + #[test] + fn native_steer_success_removes_only_the_acknowledged_occurrence() { + let mut q = EventQueue::new(DedupMode::Queue); + let ch = Uuid::new_v4(); + let duplicate = make_queued(ch, "same signed delivery"); + + let first = q.push(duplicate.clone()).expect("first occurrence"); + let second = q.push(duplicate).expect("second occurrence"); + assert_ne!(first, second); + + assert!(q.mark_native_steer_pending(ch, first)); + q.remove_native_steer(ch, first); + + let remaining = q.flush_next().expect("second occurrence remains queued"); + assert_eq!(remaining.events.len(), 1); + assert_eq!(remaining.events[0].event.content, "same signed delivery"); + assert_eq!(remaining.occurrence_ids.events, vec![second.0]); + } + + #[test] + fn native_steer_failure_releases_only_the_acknowledged_occurrence() { + let mut q = EventQueue::new(DedupMode::Queue); + let ch = Uuid::new_v4(); + let duplicate = make_queued(ch, "same signed delivery"); + + let first = q.push(duplicate.clone()).expect("first occurrence"); + let second = q.push(duplicate).expect("second occurrence"); + + assert!(q.mark_native_steer_pending(ch, first)); + q.release_native_steer(ch, first); + + let batch = q.flush_next().expect("both occurrences remain deliverable"); + assert_eq!(batch.events.len(), 2); + let ids: HashSet<_> = batch.occurrence_ids.events.into_iter().collect(); + assert_eq!(ids, HashSet::from([first.0, second.0])); + } + + #[test] + fn native_steer_individual_releases_preserve_fifo_and_cross_channel_fairness() { + let mut q = EventQueue::new(DedupMode::Queue); + let steer_channel = Uuid::new_v4(); + let competing_channel = Uuid::new_v4(); + + let oldest = make_queued_at(steer_channel, "oldest", Duration::from_millis(30)); + let oldest_id = oldest.event.id.to_hex(); + let newest = make_queued_at(steer_channel, "newest", Duration::from_millis(10)); + let newest_id = newest.event.id.to_hex(); + let competitor = make_queued_at(competing_channel, "competitor", Duration::from_millis(20)); + + let oldest_occurrence = q.push(oldest).expect("accepted oldest steer"); + let newest_occurrence = q.push(newest).expect("accepted newest steer"); + q.push(competitor).expect("accepted competing event"); + assert!(q.mark_native_steer_pending(steer_channel, oldest_occurrence)); + assert!(q.mark_native_steer_pending(steer_channel, newest_occurrence)); + + q.release_native_steer(steer_channel, oldest_occurrence); + q.release_native_steer(steer_channel, newest_occurrence); + + let batch = q + .flush_next() + .expect("the channel holding the oldest occurrence must dispatch first"); + assert_eq!(batch.channel_id, steer_channel); + assert_eq!( + batch + .events + .iter() + .map(|event| event.event.id.to_hex()) + .collect::>(), + vec![oldest_id, newest_id], + "independent acknowledgements must not reverse withheld FIFO order" + ); + } + + #[test] + fn native_steer_reverse_ack_preserves_enqueue_order_when_instants_are_equal() { + let mut q = EventQueue::new(DedupMode::Queue); + let channel = Uuid::new_v4(); + let shared_received_at = Instant::now(); + let mut first = make_queued(channel, "first"); + first.received_at = shared_received_at; + let first_id = first.event.id.to_hex(); + let mut second = make_queued(channel, "second"); + second.received_at = shared_received_at; + let second_id = second.event.id.to_hex(); + + let first_occurrence = q.push(first).expect("accepted first occurrence"); + let second_occurrence = q.push(second).expect("accepted second occurrence"); + assert!(q.mark_native_steer_pending(channel, first_occurrence)); + assert!(q.mark_native_steer_pending(channel, second_occurrence)); + + // Acknowledgements may arrive in the opposite order from enqueue. + q.release_native_steer(channel, second_occurrence); + q.release_native_steer(channel, first_occurrence); + + let batch = q.flush_next().expect("released occurrences must dispatch"); + assert_eq!( + batch + .events + .iter() + .map(|event| event.event.id.to_hex()) + .collect::>(), + vec![first_id, second_id], + "the enqueue ordinal must break equal-Instant ties" + ); + } + /// Earlier events on the same channel must flush normally during the /// steer ack window. Only the specific withheld event is invisible. /// After `release_native_steer`, the released event sits at the queue @@ -4317,10 +5578,10 @@ mod tests { let e3_id = e3.event.id.to_hex(); q.push(e1); q.push(e2); - q.push(e3); + let e3_occurrence = q.push(e3).expect("accepted e3"); // Steer in flight for e3 — withhold it from normal dispatch. - assert!(q.mark_native_steer_pending(ch, &e3_id)); + assert!(q.mark_native_steer_pending(ch, e3_occurrence)); // Earlier events flush as a normal batch; e3 is invisible. let batch = q @@ -4335,7 +5596,7 @@ mod tests { q.mark_complete(ch); // Ack arrives as Err or PromptCompletedNeutral → release e3. - q.release_native_steer(ch, &e3_id); + q.release_native_steer(ch, e3_occurrence); let next = q.flush_next().expect("released e3 should now flush"); assert_eq!(next.channel_id, ch); @@ -4358,14 +5619,14 @@ mod tests { let qe = make_queued(ch, "withheld event"); let event_id = qe.event.id.to_hex(); - q.push(qe); + let occurrence_id = q.push(qe).expect("accepted occurrence"); // Simulate a prompt in flight for `ch`, then withhold the queued // event for an in-flight goose-native steer. q.in_flight_channels.insert(ch); q.in_flight_deadlines.insert(ch, Instant::now()); q.in_flight_batch_sizes.insert(ch, 1); - assert!(q.mark_native_steer_pending(ch, &event_id)); + assert!(q.mark_native_steer_pending(ch, occurrence_id)); // Force the in-flight deadline to be in the past, simulating the // steer ack never arriving and the read loop hanging long enough @@ -4395,8 +5656,8 @@ mod tests { } /// Bulk-release on expiry must preserve original FIFO. The - /// implementation iterates the side-table entries in reverse and - /// `push_front`s each — composing to original-FIFO at the queue front. + /// implementation merges every entry by its preserved receive timestamp + /// and enqueue ordinal. /// Test ≥2 withheld entries (3 here) with staggered `received_at`. #[test] fn test_native_steer_bulk_release_preserves_fifo() { @@ -4410,19 +5671,18 @@ mod tests { let e1_id = e1.event.id.to_hex(); let e2_id = e2.event.id.to_hex(); let e3_id = e3.event.id.to_hex(); - q.push(e1); - q.push(e2); - q.push(e3); + let e1_occurrence = q.push(e1).expect("accepted e1"); + let e2_occurrence = q.push(e2).expect("accepted e2"); + let e3_occurrence = q.push(e3).expect("accepted e3"); // Withhold all three in FIFO arrival order (e1, e2, e3 → side table). // This simulates a pathological repeated-steer flow; the more // realistic case (one withhold at a time) is covered by the other // tests. What matters here is that the bulk-recovery path - // (reverse iter + push_front) composes to original FIFO at the - // queue front. - assert!(q.mark_native_steer_pending(ch, &e1_id)); - assert!(q.mark_native_steer_pending(ch, &e2_id)); - assert!(q.mark_native_steer_pending(ch, &e3_id)); + // The arrival-key merge composes to original FIFO at the queue front. + assert!(q.mark_native_steer_pending(ch, e1_occurrence)); + assert!(q.mark_native_steer_pending(ch, e2_occurrence)); + assert!(q.mark_native_steer_pending(ch, e3_occurrence)); assert_eq!(pending_count(&q), 0); assert_eq!(q.withheld_native_steer.get(&ch).map(|v| v.len()), Some(3)); @@ -4461,6 +5721,7 @@ mod tests { }], cancelled_events: vec![], cancel_reason: None, + occurrence_ids: Default::default(), }; let prompt = format_prompt( &batch, @@ -4490,6 +5751,7 @@ mod tests { }], cancelled_events: vec![], cancel_reason: None, + occurrence_ids: Default::default(), }; let prompt = format_prompt( &batch, @@ -4518,6 +5780,7 @@ mod tests { }], cancelled_events: vec![], cancel_reason: None, + occurrence_ids: Default::default(), }; let prompt = format_prompt(&batch, &FormatPromptArgs::default()).join("\n\n"); assert!( @@ -4756,4 +6019,163 @@ mod tests { "second extend must not move deadline backward (monotonic)" ); } + + #[test] + fn required_agent_wait_does_not_block_other_channels() { + let mut q = EventQueue::new(DedupMode::Queue); + let pinned = Uuid::new_v4(); + let other = Uuid::new_v4(); + + q.push(make_queued(pinned, "switch-model-retry")); + q.push(make_queued(other, "independent-work")); + q.require_agent(pinned, 0); + q.block_required_agent(pinned); + + let first = q + .flush_next() + .expect("unrelated work must remain dispatchable"); + assert_eq!(first.channel_id, other); + q.mark_complete(other); + assert!( + q.flush_next().is_none(), + "the pinned channel must wait while its exact slot is unavailable" + ); + + q.release_required_agent(0); + let pinned_batch = q + .flush_next() + .expect("the pinned batch must wake when its slot returns"); + assert_eq!(pinned_batch.channel_id, pinned); + assert_eq!(q.required_agent(pinned), Some(0)); + + q.clear_required_agent(pinned); + assert_eq!(q.required_agent(pinned), None); + } + + #[test] + fn draining_removed_channel_clears_required_agent_wait() { + let mut q = EventQueue::new(DedupMode::Queue); + let channel_id = Uuid::new_v4(); + q.push(make_queued(channel_id, "stale-switch-model-retry")); + q.require_agent(channel_id, 1); + q.block_required_agent(channel_id); + + let drained = q.drain_channel(channel_id); + assert_eq!(drained.len(), 1); + assert_eq!(q.required_agent(channel_id), None); + + q.release_required_agent(1); + assert!( + q.flush_next().is_none(), + "membership removal must not resurrect affinity-held work" + ); + } + + #[test] + fn pool_exhaustion_requeue_preserves_cancelled_merge() { + let mut q = EventQueue::new(DedupMode::Queue); + let channel_id = Uuid::new_v4(); + q.push(make_queued(channel_id, "original")); + let original = q.flush_next().expect("original batch"); + q.requeue_as_cancelled(original, CancelReason::Interrupt); + q.mark_complete(channel_id); + q.push(make_queued(channel_id, "new")); + + let merged = q.flush_next().expect("merged batch"); + assert_eq!(merged.events.len(), 1); + assert_eq!(merged.cancelled_events.len(), 1); + q.requeue_preserve_timestamps(merged); + q.mark_complete(channel_id); + + let retried = q.flush_next().expect("preserved merged batch"); + assert_eq!(retried.events.len(), 1); + assert_eq!(retried.cancelled_events.len(), 1); + assert_eq!(retried.cancel_reason, Some(CancelReason::Interrupt)); + } + + #[test] + fn blocked_cancelled_only_fallback_keeps_interrupt_framing() { + let mut q = EventQueue::new(DedupMode::Queue); + let channel_id = Uuid::new_v4(); + q.push(make_queued(channel_id, "interrupted")); + let original = q.flush_next().expect("original batch"); + q.requeue_as_cancelled(original, CancelReason::Interrupt); + q.mark_complete(channel_id); + + let fallback = q.flush_next().expect("cancelled-only fallback"); + assert!(fallback.cancelled_events.is_empty()); + assert_eq!(fallback.cancel_reason, Some(CancelReason::Interrupt)); + q.requeue_preserve_timestamps(fallback); + q.mark_complete(channel_id); + q.push(make_queued(channel_id, "new-request")); + + let merged = q.flush_next().expect("merged retry"); + assert_eq!(merged.events.len(), 1); + assert_eq!(merged.cancelled_events.len(), 1); + assert_eq!(merged.cancel_reason, Some(CancelReason::Interrupt)); + } + + #[test] + fn failed_merged_batch_respects_retry_backoff_as_one_unit() { + let mut q = EventQueue::new(DedupMode::Queue); + let channel_id = Uuid::new_v4(); + q.push(make_queued(channel_id, "interrupted")); + let original = q.flush_next().expect("original batch"); + q.requeue_as_cancelled(original, CancelReason::Interrupt); + q.mark_complete(channel_id); + q.push(make_queued(channel_id, "new-request")); + let merged = q.flush_next().expect("merged batch"); + + assert!(q.requeue(merged).is_none()); + q.mark_complete(channel_id); + assert!( + q.flush_next().is_none(), + "neither half of a merged batch may bypass its retry throttle" + ); + + q.set_retry_deadline_for_test(channel_id, Instant::now()); + let retried = q.flush_next().expect("retry after backoff"); + assert_eq!(retried.events.len(), 1); + assert_eq!(retried.cancelled_events.len(), 1); + assert_eq!(retried.cancel_reason, Some(CancelReason::Interrupt)); + } + + #[test] + fn ordinary_failed_batch_does_not_create_empty_cancelled_fallback() { + let mut q = EventQueue::new(DedupMode::Queue); + let channel_id = Uuid::new_v4(); + q.push(make_queued(channel_id, "ordinary")); + let batch = q.flush_next().expect("ordinary batch"); + + assert!(q.requeue(batch).is_none()); + q.mark_complete(channel_id); + assert!(!q.cancelled_batches.contains_key(&channel_id)); + assert!(q.flush_next().is_none()); + } + + #[test] + fn compact_keeps_retry_budget_for_cancelled_only_fallback() { + let mut q = EventQueue::new(DedupMode::Queue); + let channel_id = Uuid::new_v4(); + q.push(make_queued(channel_id, "interrupted")); + let original = q.flush_next().expect("original"); + q.requeue_as_cancelled(original, CancelReason::Interrupt); + q.mark_complete(channel_id); + let fallback = q.flush_next().expect("cancelled-only fallback"); + + assert!(q.requeue(fallback).is_none()); + q.mark_complete(channel_id); + assert_eq!(q.retry_count_for_test(channel_id), Some(1)); + q.set_retry_deadline_for_test(channel_id, Instant::now()); + q.compact_expired_state(); + assert_eq!( + q.retry_count_for_test(channel_id), + Some(1), + "maintenance must not reset the retry budget while cancelled work remains" + ); + + let retry = q.flush_next().expect("retry after expiry"); + assert!(q.requeue(retry).is_none()); + assert_eq!(q.retry_count_for_test(channel_id), Some(2)); + } } diff --git a/crates/buzz-acp/src/relay.rs b/crates/buzz-acp/src/relay.rs index c8312cc61e..b26ad15a5e 100644 --- a/crates/buzz-acp/src/relay.rs +++ b/crates/buzz-acp/src/relay.rs @@ -22,6 +22,7 @@ //! channel. `next_event()` reads from the event receiver. use std::collections::{HashMap, HashSet, VecDeque}; +use std::sync::atomic::{AtomicBool, Ordering}; use std::time::Duration; /// Default capacity of the event channel from background task to harness. @@ -42,6 +43,23 @@ fn event_channel_capacity() -> usize { /// 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; +/// Maximum accepted age for ordinary owner-control candidates. +const PRIVILEGED_CONTROL_FRESHNESS_SECS: u64 = 300; +/// Bound privileged-control replay memory without evicting still-fresh IDs. +/// +/// On saturation, new controls fail closed until an entry expires. Legitimate +/// control volume is expected to be several orders of magnitude below this. +const PRIVILEGED_CONTROL_REPLAY_CAPACITY: usize = 1_024; +/// Maximum accepted age for authenticated owner-to-agent observer controls. +const OBSERVER_CONTROL_FRESHNESS_SECS: u64 = 300; +/// Bound observer-control replay memory without evicting still-fresh IDs. +const OBSERVER_CONTROL_REPLAY_CAPACITY: usize = 1_024; +/// Bound authenticated observer controls waiting for downstream capacity. +/// +/// This matches the non-evicting replay capacity: every still-fresh pending +/// control retains a replay entry, so a newly admitted control always has a +/// pending slot after expired entries are removed. +const OBSERVER_CONTROL_PENDING_CAPACITY: usize = OBSERVER_CONTROL_REPLAY_CAPACITY; /// Interval between client-initiated WebSocket pings. const PING_INTERVAL: Duration = Duration::from_secs(30); @@ -57,6 +75,12 @@ const WS_SEND_TIMEOUT_SECS: u64 = 10; const STABLE_CONNECTION_SECS: u64 = 60; /// Seconds subtracted from `since` on resubscribe to tolerate clock skew. const SINCE_SKEW_SECS: u64 = 5; +/// Maximum client clock lead accepted for ordinary channel messages. +/// +/// Privileged controls still require a non-future timestamp. Accepted skew is +/// clamped to the local clock before replay-watermark bookkeeping so a client +/// cannot push reconnect state into the future. +const CHANNEL_EVENT_FUTURE_SKEW_SECS: u64 = SINCE_SKEW_SECS; /// Timeout for the NIP-42 auth handshake steps. /// /// Raised from 5s to 20s (≈2 RTTs at the observed 10s max round-trip on degraded @@ -110,17 +134,40 @@ const DRAIN_BUDGET_PER_ITER: usize = 1; /// this covers ~40 s of gating; beyond that the oldest frames are dropped with /// visible accounting (`gated_observer_dropped`). const GATED_OBSERVER_QUEUE_CAP: usize = 256; +/// Reserved relay-gate capacity for terminal observer control results. +const GATED_OBSERVER_PRIORITY_CAP: usize = 64; +/// Largest relay retry hint accepted into the local admission gate. +/// +/// The production relay's fixed message window is sixty seconds, so the client +/// must honor that full reset horizon. Larger, malformed hints are bounded here +/// so they cannot park all local relay traffic indefinitely. +const MAX_RATE_LIMIT_RETRY_SECS: u64 = 60; +/// Maximum time a terminal observer publisher waits for the relay's exact +/// `OK accepted=true` admission proof. +/// +/// Ninety seconds is longer than the full sixty-second relay window plus its +/// positive-only twenty-percent desynchronization (72 s), with room for the +/// paced resend and relay OK. A later independent rate-limit response can still +/// exhaust this bounded attempt; the observer ledger then performs one bounded +/// end-to-end retry rather than silently discarding the terminal result. +const TERMINAL_OBSERVER_ACK_TIMEOUT: Duration = Duration::from_secs(90); +/// Cleartext marker added only to encrypted terminal control-result frames. +pub(crate) const OBSERVER_PRIORITY_TAG: &str = "priority"; +pub(crate) const OBSERVER_CONTROL_RESULT_PRIORITY: &str = "control-result"; use std::time::Instant; use buzz_core::kind::{ KIND_AGENT_OBSERVER_FRAME, KIND_MEMBER_ADDED_NOTIFICATION, KIND_MEMBER_REMOVED_NOTIFICATION, - KIND_TYPING_INDICATOR, + KIND_STREAM_MESSAGE, KIND_TYPING_INDICATOR, +}; +use buzz_core::observer::{ + content_looks_like_nip44, OBSERVER_AGENT_TAG, OBSERVER_FRAME_CONTROL, OBSERVER_FRAME_TAG, }; -use futures_util::{SinkExt, StreamExt}; +use futures_util::{Sink, SinkExt, StreamExt}; use nostr::{Event, EventBuilder, Keys, Kind, RelayUrl, Tag}; use serde_json::{json, Value}; -use tokio::sync::mpsc; +use tokio::sync::{mpsc, oneshot}; use tokio::time::timeout; use tokio_tungstenite::{connect_async, tungstenite::Message, MaybeTlsStream, WebSocketStream}; use tracing::{debug, info, warn}; @@ -258,6 +305,52 @@ fn unix_now_secs() -> u64 { .as_secs() } +/// Resolve the only identity allowed to exercise privileged relay controls. +/// +/// The relay-facing `auth` tag is parsed before `HarnessRelay::connect`, but +/// parsing alone does not authenticate its embedded owner. Verify the NIP-OA +/// signature against this agent key before retaining the owner identity. When +/// the attestation is absent or invalid, revalidate the trusted configured +/// owner fallback; if neither source is valid, privileged control is disabled. +fn resolve_authorized_owner_pubkey( + auth_tag: Option<&nostr::Tag>, + agent_keys: &Keys, + configured_owner: Option<&str>, +) -> Option { + if let Some(auth_tag) = auth_tag { + match serde_json::to_string(auth_tag) { + Ok(auth_tag_json) => { + match buzz_sdk::nip_oa::verify_auth_tag(&auth_tag_json, &agent_keys.public_key()) { + Ok(owner) => return Some(owner.to_hex()), + Err(error) => { + warn!( + "NIP-OA owner attestation verification failed; using configured fallback when valid: {error}" + ); + } + } + } + Err(error) => { + warn!( + "could not serialize NIP-OA owner attestation; using configured fallback when valid: {error}" + ); + } + } + } + + let configured_owner = configured_owner?; + match nostr::PublicKey::from_hex(configured_owner.trim()) { + Ok(owner) if owner != agent_keys.public_key() => Some(owner.to_hex()), + Ok(_) => { + warn!("configured privileged owner matches the agent key; failing closed"); + None + } + Err(error) => { + warn!("configured privileged owner is invalid; failing closed: {error}"); + None + } + } +} + impl RestClient { /// Sign a NIP-98 HTTP Auth event (kind:27235) for the given method/URL/body. /// @@ -430,6 +523,11 @@ pub struct BuzzEvent { pub channel_id: Uuid, /// The underlying Nostr event. pub event: Event, + /// Relay-issued authority for harness-level owner control execution. + /// + /// Command-shaped events without this marker remain ordinary messages, + /// even when their signer matches a configured owner. + pub privileged_control_admitted: bool, } /// Errors from relay operations. @@ -517,7 +615,10 @@ enum RelayCommand { /// Subscribe to encrypted observer control frames addressed to this agent. SubscribeObserverControls, /// Publish a signed event to the relay (for typing indicators, etc.). - PublishEvent { event: Box }, + PublishEvent { + event: Box, + admission: Option>>, + }, /// Floor `since` for membership notification replay; events before startup are never re-delivered. SetStartupWatermark { ts: u64 }, } @@ -564,11 +665,60 @@ impl RelayEventPublisher { self.cmd_tx .send(RelayCommand::PublishEvent { event: Box::new(event), + admission: None, }) .await .map_err(|_| RelayError::ConnectionClosed) } + /// Publish a terminal observer result only after the relay explicitly + /// acknowledges it with `OK accepted=true`. + pub async fn publish_terminal_event(&self, event: Event) -> Result<(), RelayError> { + self.publish_terminal_event_with_ack_timeout(event, TERMINAL_OBSERVER_ACK_TIMEOUT) + .await + } + + async fn publish_terminal_event_with_ack_timeout( + &self, + event: Event, + ack_timeout: Duration, + ) -> Result<(), RelayError> { + if !is_exact_terminal_observer_frame(&event) { + return Err(RelayError::Http( + "terminal observer event must be kind 24200 with exactly one control-result priority marker" + .to_string(), + )); + } + let deadline = tokio::time::Instant::now() + ack_timeout; + let (admission_tx, admission_rx) = oneshot::channel(); + match tokio::time::timeout_at( + deadline, + self.cmd_tx.send(RelayCommand::PublishEvent { + event: Box::new(event), + admission: Some(admission_tx), + }), + ) + .await + { + Ok(result) => result.map_err(|_| RelayError::ConnectionClosed)?, + Err(_) => return Err(RelayError::Timeout), + } + match tokio::time::timeout_at(deadline, admission_rx).await { + Ok(result) => result + .map_err(|_| RelayError::ConnectionClosed)? + .map_err(RelayError::Http), + Err(_) => { + // Dropping `admission_rx` above closes only the sender owned by + // this exact publication attempt. That per-attempt sender is the + // authoritative cancellation token checked before every park, + // socket write, and reconnect drain. An event-ID-only cancellation + // would be ambiguous when a delayed duplicate shares the original + // signed event ID, so timeout cleanup never enqueues one. + Err(RelayError::Timeout) + } + } + } + /// Test-only publisher pair: published events are forwarded to the /// returned receiver instead of a live relay socket. #[cfg(test)] @@ -577,10 +727,41 @@ impl RelayEventPublisher { let (event_tx, event_rx) = mpsc::channel(64); tokio::spawn(async move { while let Some(cmd) = cmd_rx.recv().await { - if let RelayCommand::PublishEvent { event } = cmd { + if let RelayCommand::PublishEvent { event, admission } = cmd { + if event_tx.send(*event).await.is_err() { + break; + } + if let Some(admission) = admission { + let _ = admission.send(Ok(())); + } + } + } + }); + (Self { cmd_tx }, event_rx) + } + + /// Test-only publisher that drops the first terminal admission sender, + /// then confirms every later publication. + #[cfg(test)] + pub(crate) fn test_pair_dropping_first_terminal_confirmation() -> (Self, mpsc::Receiver) + { + let (cmd_tx, mut cmd_rx) = mpsc::channel::(64); + let (event_tx, event_rx) = mpsc::channel(64); + tokio::spawn(async move { + let mut drop_first = true; + while let Some(cmd) = cmd_rx.recv().await { + if let RelayCommand::PublishEvent { event, admission } = cmd { if event_tx.send(*event).await.is_err() { break; } + if let Some(admission) = admission { + if drop_first { + drop_first = false; + drop(admission); + } else { + let _ = admission.send(Ok(())); + } + } } } }); @@ -598,6 +779,21 @@ impl HarnessRelay { keys: &Keys, agent_pubkey_hex: &str, auth_tag: Option, + ) -> Result { + Self::connect_with_owner(relay_url, keys, agent_pubkey_hex, auth_tag, None).await + } + + /// Connect with the canonical owner resolved by the trusted harness config. + /// + /// A valid NIP-OA attestation remains authoritative. `configured_owner` is + /// used only when the attestation is absent or invalid, and is normalized + /// and revalidated before the background task classifies any event. + pub async fn connect_with_owner( + relay_url: &str, + keys: &Keys, + agent_pubkey_hex: &str, + auth_tag: Option, + configured_owner: Option, ) -> Result { // Perform the initial connection and auth handshake, retrying // transient failures (dropped handshake, timeout) with bounded @@ -616,6 +812,7 @@ impl HarnessRelay { let bg_relay_url = relay_url.to_string(); let bg_agent_pubkey_hex = agent_pubkey_hex.to_string(); let bg_auth_tag = auth_tag.clone(); + let bg_configured_owner = configured_owner; let bg_handle = tokio::spawn(async move { run_background_task( @@ -628,6 +825,7 @@ impl HarnessRelay { bg_relay_url, bg_agent_pubkey_hex, bg_auth_tag, + bg_configured_owner, ) .await; }); @@ -822,6 +1020,7 @@ impl HarnessRelay { self.cmd_tx .send(RelayCommand::PublishEvent { event: Box::new(event), + admission: None, }) .await .map_err(|_| RelayError::ConnectionClosed) @@ -835,6 +1034,7 @@ impl HarnessRelay { self.cmd_tx .try_send(RelayCommand::PublishEvent { event: Box::new(event), + admission: None, }) .map_err(|_| RelayError::ConnectionClosed) } @@ -972,14 +1172,77 @@ impl TwoGenDedup { } } +#[derive(Debug, PartialEq, Eq)] +enum PrivilegedControlAdmission { + Admitted, + Replay, + Full, +} + +/// Non-evicting replay memory for ordinary privileged control candidates. +/// +/// Entries remain present for the event's entire freshness lifetime. Unlike +/// [`TwoGenDedup`], capacity pressure never evicts a still-admissible control +/// ID because doing so would reopen a replay path. Saturation fails closed. +struct PrivilegedControlReplay { + expires_at: HashMap, + capacity: usize, +} + +impl PrivilegedControlReplay { + fn new(capacity: usize) -> Self { + assert!( + capacity > 0, + "privileged control replay capacity must be positive" + ); + Self { + expires_at: HashMap::with_capacity(capacity), + capacity, + } + } + + fn admit(&mut self, event_id: String, expires_at: u64, now: u64) -> PrivilegedControlAdmission { + self.expires_at.retain(|_, expiry| *expiry >= now); + if self.expires_at.contains_key(&event_id) { + return PrivilegedControlAdmission::Replay; + } + if self.expires_at.len() >= self.capacity { + return PrivilegedControlAdmission::Full; + } + self.expires_at.insert(event_id, expires_at); + PrivilegedControlAdmission::Admitted + } + + fn remove(&mut self, event_id: &str) { + self.expires_at.remove(event_id); + } +} + /// State maintained by the background WebSocket task. struct BgState { /// Active subscriptions: channel_id → subscription_id string. active_subscriptions: HashMap, + /// Successful REQ completion second for each currently live channel. + /// + /// Owner controls must strictly post-date this boundary. Intent recorded + /// while disconnected never populates it. + channel_subscribed_at: HashMap, /// Most recent `created_at` timestamp seen per channel (for `since` filter). last_seen: HashMap, /// Two-generation dedup set of event IDs seen. seen_ids: TwoGenDedup, + /// Supervisor-process start second for privileged control admission. + /// + /// Nostr timestamps have one-second precision, so controls at exactly this + /// boundary are rejected: they cannot be proven to post-date startup. + privileged_control_start_boundary: u64, + /// Owner identity authenticated by the agent's NIP-OA attestation. + /// + /// `None` disables privileged classification. Command-looking events from + /// any other signer remain ordinary channel messages. + authorized_owner_pubkey: Option, + /// Replay IDs retained without eviction for their full freshness lifetime. + privileged_control_replay: PrivilegedControlReplay, /// Per-channel filter used on subscribe (for resubscribe after reconnect). active_filters: HashMap, /// Oldest timestamp of a membership notification that was dropped due to @@ -993,6 +1256,19 @@ struct BgState { membership_sub_active: bool, /// Whether the observer control subscription is active. observer_control_sub_active: bool, + /// Start second of the exact observer-control REQ currently live on the socket. + /// + /// `None` means subscription intent exists but no REQ is currently proven + /// live. Equality is rejected because Nostr timestamps have one-second + /// precision and cannot prove that a control followed the REQ. + observer_control_subscribed_at: Option, + /// Authenticated observer-control IDs retained for their freshness lifetime. + observer_control_replay: PrivilegedControlReplay, + /// Authenticated controls waiting for capacity in the downstream channel. + /// + /// This FIFO is drained by its own `select!` arm so relay reads, pings, + /// reconnects, and shutdown commands never wait for downstream capacity. + observer_control_pending: VecDeque, /// Oldest dropped channel-event timestamp per channel, keyed by channel_id. /// Mirrors `membership_dropped_since` but for ordinary channel events. /// On reconnect resubscribe, `since` = min(last_seen, channel_dropped_since). @@ -1039,13 +1315,25 @@ struct BgState { /// Bounded at `GATED_OBSERVER_QUEUE_CAP` (drop-oldest); drained by the /// main loop one frame per pacing tick once the gate clears. gated_observer_pending: VecDeque>, + /// Protected terminal control-result frames. This FIFO is never displaced + /// by ordinary observer telemetry and drains first after a relay gate. + gated_observer_priority_pending: VecDeque>, /// Observer frames written to the socket but not yet acknowledged. The /// relay's rate-limit NOTICE does not carry an event ID, so all unresolved /// observer writes are moved back ahead of the parked FIFO when one arrives. observer_in_flight: VecDeque>, + /// Protected acknowledgement window for terminal control-result frames. + observer_priority_in_flight: VecDeque>, + /// Terminal publisher completions keyed by relay event ID. A completion is + /// released only by an explicit relay `OK`, so local queue admission can + /// never be mistaken for externally durable proof. + observer_priority_confirmations: HashMap>>, /// Frames evicted from the bounded pending/in-flight observer buffers since /// summary log. Makes overflow loss visible instead of silent. gated_observer_dropped: u64, + /// Observer publishes rejected by a terminal relay denial. These are not + /// retried, but the count and exact denial are logged for auditability. + observer_terminal_denied: u64, /// Channels whose REQ failed during `resubscribe_after_reconnect`. /// /// A single failed channel REQ is parked here instead of aborting the whole @@ -1063,13 +1351,22 @@ impl BgState { fn new() -> Self { Self { active_subscriptions: HashMap::new(), + channel_subscribed_at: HashMap::new(), last_seen: HashMap::new(), seen_ids: TwoGenDedup::new(SEEN_ID_LIMIT), + privileged_control_start_boundary: unix_now_secs(), + authorized_owner_pubkey: None, + privileged_control_replay: PrivilegedControlReplay::new( + PRIVILEGED_CONTROL_REPLAY_CAPACITY, + ), active_filters: HashMap::new(), membership_dropped_since: None, membership_last_seen: None, membership_sub_active: false, observer_control_sub_active: false, + observer_control_subscribed_at: None, + observer_control_replay: PrivilegedControlReplay::new(OBSERVER_CONTROL_REPLAY_CAPACITY), + observer_control_pending: VecDeque::new(), channel_dropped_since: HashMap::new(), proactive_resubscribe_needed: false, startup_watermark: None, @@ -1079,16 +1376,25 @@ impl BgState { membership_resub_needed: false, observer_resub_needed: false, gated_observer_pending: VecDeque::new(), + gated_observer_priority_pending: VecDeque::new(), observer_in_flight: VecDeque::new(), + observer_priority_in_flight: VecDeque::new(), + observer_priority_confirmations: HashMap::new(), gated_observer_dropped: 0, + observer_terminal_denied: 0, resubscribe_retry: HashSet::new(), backoff_step: 0, } } - /// Record a received event for dedup and `since` tracking. + /// Record a received event for dedup and bounded `since` tracking. /// Returns `true` if the event is new (not a duplicate). - fn record_event(&mut self, channel_id: Uuid, event: &Event) -> bool { + fn record_event_with_watermark( + &mut self, + channel_id: Uuid, + event: &Event, + watermark_timestamp: u64, + ) -> bool { let id_hex = event.id.to_hex(); // Two-generation dedup: no amnesia window on rotation. @@ -1097,15 +1403,19 @@ impl BgState { } // Update last_seen timestamp. - let ts = event.created_at.as_secs(); self.last_seen .entry(channel_id) - .and_modify(|t| *t = (*t).max(ts)) - .or_insert(ts); + .and_modify(|t| *t = (*t).max(watermark_timestamp)) + .or_insert(watermark_timestamp); true } + #[cfg(test)] + fn record_event(&mut self, channel_id: Uuid, event: &Event) -> bool { + self.record_event_with_watermark(channel_id, event, event.created_at.as_secs()) + } + /// Compute the `since` timestamp for a channel (re)subscribe. /// /// Picks the earliest of `last_seen` and `channel_dropped_since` so @@ -1131,6 +1441,7 @@ impl BgState { /// 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.channel_subscribed_at.remove(channel_id); self.last_seen.remove(channel_id); self.subscribe_since.remove(channel_id); self.channel_dropped_since.remove(channel_id); @@ -1146,17 +1457,25 @@ impl BgState { /// low-quality hints from dropping the gate so short that re-queued REQs /// immediately re-trigger rate limiting. Note the deliberate asymmetry with /// the desktop TypeScript client, which uses a 10s no-hint default — both - /// values are conservative enough; the relay hint wins when present. + /// values are conservative enough. Relay hints win when present, up to the + /// bounded [`MAX_RATE_LIMIT_RETRY_SECS`] contract that keeps terminal + /// confirmation alive through the retry window. /// /// The gate takes the **maximum** of any existing deadline and the newly /// computed one so overlapping CLOSED/NOTICE messages can't shorten a gate /// that is already set further out. /// - /// Returns the gate deadline that was set. + /// Explicit reset hints use positive-only jitter: retrying before the + /// advertised reset can create a self-sustaining rate-limit loop. Returns + /// the gate deadline that was set. fn set_rate_limit_gate(&mut self, retry_secs: u64) -> tokio::time::Instant { - let secs = if retry_secs < 2 { 5 } else { retry_secs }; + let secs = if retry_secs < 2 { + 5 + } else { + retry_secs.min(MAX_RATE_LIMIT_RETRY_SECS) + }; let base = Duration::from_secs(secs); - let deadline = tokio::time::Instant::now() + jittered_duration(base); + let deadline = tokio::time::Instant::now() + rate_limit_retry_duration(base); let gate = match self.rate_limit_gate { Some(existing) if existing > deadline => existing, _ => deadline, @@ -1184,7 +1503,21 @@ impl BgState { /// /// Bounded drop-oldest queue: overflow evicts the oldest frame and counts /// it in `gated_observer_dropped` so the loss is visible, never silent. - fn park_gated_observer_frame(&mut self, event: Box) { + fn park_gated_observer_frame(&mut self, event: Box) -> bool { + if is_priority_observer_frame(&event) { + if self.gated_observer_priority_pending.len() + self.observer_priority_in_flight.len() + >= GATED_OBSERVER_PRIORITY_CAP + { + self.gated_observer_dropped += 1; + tracing::error!( + dropped_total = self.gated_observer_dropped, + "protected observer delivery state full — rejected newest terminal result" + ); + return false; + } + self.gated_observer_priority_pending.push_back(event); + return true; + } if self.gated_observer_pending.len() >= GATED_OBSERVER_QUEUE_CAP { self.gated_observer_pending.pop_front(); self.gated_observer_dropped += 1; @@ -1194,12 +1527,17 @@ impl BgState { ); } self.gated_observer_pending.push_back(event); + true } /// Restore unresolved observer writes ahead of frames parked after the /// gate armed. NOTICE has no event ID, so conservatively retry every frame /// without an OK; duplicate IDs are harmless at the relay. fn requeue_observer_in_flight(&mut self) { + self.prune_cancelled_terminal_admissions(); + while let Some(event) = self.observer_priority_in_flight.pop_back() { + self.gated_observer_priority_pending.push_front(event); + } while let Some(event) = self.observer_in_flight.pop_back() { self.gated_observer_pending.push_front(event); } @@ -1209,7 +1547,29 @@ impl BgState { } } - fn track_observer_in_flight(&mut self, event: Box) { + /// Retire authority belonging to one failed replacement socket while + /// preserving all unresolved observer delivery intent for the next one. + fn abandon_failed_reconnect_candidate(&mut self) { + self.requeue_observer_in_flight(); + self.observer_control_subscribed_at = None; + self.channel_subscribed_at.clear(); + } + + fn track_observer_in_flight(&mut self, event: Box) -> bool { + if is_priority_observer_frame(&event) { + if self.gated_observer_priority_pending.len() + self.observer_priority_in_flight.len() + >= GATED_OBSERVER_PRIORITY_CAP + { + self.gated_observer_dropped += 1; + tracing::error!( + dropped_total = self.gated_observer_dropped, + "protected observer acknowledgment state full — rejected newest terminal result" + ); + return false; + } + self.observer_priority_in_flight.push_back(event); + return true; + } if self.observer_in_flight.len() >= GATED_OBSERVER_QUEUE_CAP { self.observer_in_flight.pop_front(); self.gated_observer_dropped += 1; @@ -1219,15 +1579,193 @@ impl BgState { ); } self.observer_in_flight.push_back(event); + true + } + + fn has_observer_priority_capacity(&self) -> bool { + self.gated_observer_priority_pending.len() + self.observer_priority_in_flight.len() + < GATED_OBSERVER_PRIORITY_CAP + } + + fn contains_observer_priority_event(&self, event_id: &str) -> bool { + self.gated_observer_priority_pending + .iter() + .chain(self.observer_priority_in_flight.iter()) + .any(|event| event.id.to_hex() == event_id) + || self.observer_priority_confirmations.contains_key(event_id) + } + + fn reject_duplicate_observer_priority( + &self, + event_id: &str, + is_priority: bool, + admission: &mut Option>>, + ) -> bool { + if !is_priority || !self.contains_observer_priority_event(event_id) { + return false; + } + tracing::error!( + event_id = %event_id, + "rejected duplicate terminal observer result before queue or socket write" + ); + if let Some(admission) = admission.take() { + let _ = admission.send(Err( + "terminal observer event is already awaiting relay confirmation".to_string(), + )); + } + true + } + + /// A timed-out publisher drops its oneshot receiver synchronously. Purge + /// every exact terminal event whose confirmation sender observes that close, + /// independent of command-channel capacity or reconnect command ordering. + fn prune_cancelled_terminal_admissions(&mut self) { + let cancelled: Vec<_> = self + .observer_priority_confirmations + .iter() + .filter(|(_, confirmation)| confirmation.is_closed()) + .map(|(event_id, _)| event_id.clone()) + .collect(); + for event_id in cancelled { + self.cancel_terminal_admission(&event_id); + } + } + + /// Remove one timed-out terminal publication from every relay-owned state + /// lane. Event IDs are signed-event identities, so matching the exact ID + /// cannot cancel a different queued terminal result. + fn cancel_terminal_admission(&mut self, event_id: &str) -> bool { + let pending_len = self.gated_observer_priority_pending.len(); + self.gated_observer_priority_pending + .retain(|event| event.id.to_hex() != event_id); + let in_flight_len = self.observer_priority_in_flight.len(); + self.observer_priority_in_flight + .retain(|event| event.id.to_hex() != event_id); + let confirmation_removed = self + .observer_priority_confirmations + .remove(event_id) + .is_some(); + let removed = self.gated_observer_priority_pending.len() != pending_len + || self.observer_priority_in_flight.len() != in_flight_len + || confirmation_removed; + if removed { + warn!( + event_id, + "terminal observer relay acknowledgement timed out — cancelled exact delivery state" + ); + } else { + debug!( + event_id, + "terminal observer timeout cancellation found no unresolved delivery state" + ); + } + removed } fn acknowledge_observer_frame(&mut self, event_id: &str) { + if self.take_unresolved_observer_frame(event_id).is_some() { + if let Some(confirmation) = self.observer_priority_confirmations.remove(event_id) { + let _ = confirmation.send(Ok(())); + } + } + } + + /// Remove and return an exact unresolved observer write from either its + /// socket-bound acknowledgement window or its rate-gated retry FIFO. + fn take_unresolved_observer_frame(&mut self, event_id: &str) -> Option> { + if let Some(index) = self + .observer_priority_in_flight + .iter() + .position(|event| event.id.to_hex() == event_id) + { + return self.observer_priority_in_flight.remove(index); + } + if let Some(index) = self + .gated_observer_priority_pending + .iter() + .position(|event| event.id.to_hex() == event_id) + { + return self.gated_observer_priority_pending.remove(index); + } if let Some(index) = self .observer_in_flight .iter() .position(|event| event.id.to_hex() == event_id) { - self.observer_in_flight.remove(index); + return self.observer_in_flight.remove(index); + } + let index = self + .gated_observer_pending + .iter() + .position(|event| event.id.to_hex() == event_id)?; + self.gated_observer_pending.remove(index) + } + + /// Retry the exact rate-limited observer frame ahead of telemetry that was + /// parked after it was first written. + fn requeue_rate_limited_observer_frame(&mut self, event_id: &str) -> bool { + if self + .gated_observer_priority_pending + .iter() + .chain(self.gated_observer_pending.iter()) + .any(|event| event.id.to_hex() == event_id) + { + return true; + } + let Some(event) = self.take_unresolved_observer_frame(event_id) else { + return false; + }; + if is_priority_observer_frame(&event) { + self.gated_observer_priority_pending.push_front(event); + } else { + self.gated_observer_pending.push_front(event); + if self.gated_observer_pending.len() > GATED_OBSERVER_QUEUE_CAP { + self.gated_observer_pending.pop_back(); + self.gated_observer_dropped += 1; + } + } + true + } + + /// Retire a terminally denied observer frame without retrying it forever. + /// The caller logs the exact relay denial; the counter preserves a bounded + /// summary signal even after the frame leaves the acknowledgement window. + fn record_terminal_observer_denial(&mut self, event_id: &str, denial: &str) -> bool { + if self.take_unresolved_observer_frame(event_id).is_none() { + return false; + } + if let Some(confirmation) = self.observer_priority_confirmations.remove(event_id) { + let _ = confirmation.send(Err(format!( + "relay denied terminal observer frame: {denial}" + ))); + } + self.observer_terminal_denied = self.observer_terminal_denied.saturating_add(1); + true + } + + fn record_observer_confirmation( + &mut self, + event_id: String, + is_priority: bool, + admission: Option>>, + admitted: bool, + ) { + let Some(admission) = admission else { + return; + }; + if admission.is_closed() { + if admitted { + self.cancel_terminal_admission(&event_id); + } + } else if !admitted { + let _ = admission.send(Err("terminal observer priority state is full".to_string())); + } else if !is_priority { + let _ = admission.send(Err( + "confirmed delivery is only valid for terminal observer frames".to_string(), + )); + } else { + self.observer_priority_confirmations + .insert(event_id, admission); } } } @@ -1242,12 +1780,16 @@ impl BgState { /// 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) { + state.prune_cancelled_terminal_admissions(); match cmd { RelayCommand::Subscribe { channel_id, filter, replay_since, } => { + // Intent is not a live subscription. Any prior socket-bound + // authority stays disabled until a new REQ is written successfully. + state.channel_subscribed_at.remove(&channel_id); state .active_subscriptions .insert(channel_id, channel_sub_id(channel_id)); @@ -1270,6 +1812,7 @@ fn apply_command_to_state(state: &mut BgState, cmd: RelayCommand) { } RelayCommand::SubscribeObserverControls => { state.observer_control_sub_active = true; + state.observer_control_subscribed_at = None; } RelayCommand::SetStartupWatermark { ts } => { state.startup_watermark = Some(ts); @@ -1281,9 +1824,30 @@ fn apply_command_to_state(state: &mut BgState, cmd: RelayCommand) { // overflow) so they are delivered by the post-reconnect drain. Other // ephemeral publishes (typing indicators) are meaningless while // disconnected and are dropped. - RelayCommand::PublishEvent { event } => { + RelayCommand::PublishEvent { + event, + mut admission, + } => { + if admission.as_ref().is_some_and(oneshot::Sender::is_closed) { + debug!( + event_id = %event.id, + "discarding terminal observer publish cancelled before disconnected admission" + ); + return; + } if event.kind.as_u16() as u32 == KIND_AGENT_OBSERVER_FRAME { - state.park_gated_observer_frame(event); + let event_id = event.id.to_hex(); + let is_priority = is_priority_observer_frame(&event); + if state.reject_duplicate_observer_priority(&event_id, is_priority, &mut admission) + { + return; + } + let admitted = state.park_gated_observer_frame(event); + state.record_observer_confirmation(event_id, is_priority, admission, admitted); + } else if let Some(admission) = admission { + let _ = admission.send(Err( + "confirmed delivery is only valid for observer frames".to_string() + )); } } // Already reconnecting — redundant. @@ -1305,13 +1869,32 @@ fn apply_command_to_state(state: &mut BgState, cmd: RelayCommand) { /// discarded because replaying a typing indicator after reconnect is meaningless. /// `Shutdown` and `Reconnect` are handled by the caller. fn retain_failed_command_intent(state: &mut BgState, cmd: RelayCommand) { + state.prune_cancelled_terminal_admissions(); match cmd { - RelayCommand::PublishEvent { event } - if event.kind.as_u16() as u32 == KIND_AGENT_OBSERVER_FRAME => - { - state.park_gated_observer_frame(event); + RelayCommand::PublishEvent { + event, + mut admission, + } if event.kind.as_u16() as u32 == KIND_AGENT_OBSERVER_FRAME => { + if admission.as_ref().is_some_and(oneshot::Sender::is_closed) { + debug!( + event_id = %event.id, + "discarding terminal observer publish cancelled before failed-send retention" + ); + return; + } + let event_id = event.id.to_hex(); + let is_priority = is_priority_observer_frame(&event); + if state.reject_duplicate_observer_priority(&event_id, is_priority, &mut admission) { + return; + } + let admitted = state.park_gated_observer_frame(event); + state.record_observer_confirmation(event_id, is_priority, admission, admitted); + } + RelayCommand::PublishEvent { admission, .. } => { + if let Some(admission) = admission { + let _ = admission.send(Err("observer frame was not retained".to_string())); + } } - RelayCommand::PublishEvent { .. } => {} cmd => apply_command_to_state(state, cmd), } } @@ -1349,6 +1932,7 @@ async fn execute_connected_command( agent_pubkey_hex: &str, cmd: RelayCommand, ) -> bool { + state.prune_cancelled_terminal_admissions(); match cmd { RelayCommand::Subscribe { channel_id, @@ -1450,13 +2034,15 @@ async fn execute_connected_command( } RelayCommand::SubscribeObserverControls => { state.observer_control_sub_active = true; + state.observer_control_subscribed_at = None; if state.check_rate_gate().is_some() { debug!("rate-gated: deferring observer control subscription"); state.observer_resub_needed = true; return true; } - let sent = send_observer_control_subscribe(ws, agent_pubkey_hex).await; - if sent { + if let Some(subscribed_at) = send_observer_control_subscribe(ws, agent_pubkey_hex).await + { + state.observer_control_subscribed_at = Some(subscribed_at); state.observer_resub_needed = false; true } else { @@ -1465,20 +2051,39 @@ async fn execute_connected_command( false } } - RelayCommand::PublishEvent { event } => { + RelayCommand::PublishEvent { + event, + mut admission, + } => { + if admission.as_ref().is_some_and(oneshot::Sender::is_closed) { + debug!( + event_id = %event.id, + "discarding terminal observer publish cancelled before socket admission" + ); + return true; + } + let event_id = event.id.to_hex(); + let is_priority = is_priority_observer_frame(&event); + if state.reject_duplicate_observer_priority(&event_id, is_priority, &mut admission) { + return true; + } // Observer telemetry frames (kind 24200) are durable telemetry, not // droppable ephemera: park them while the rate-limit gate is armed — // and while earlier parked frames are still draining, so relative // order is preserved — then let the main-loop drain deliver them // one per pacing tick once the gate clears. if event.kind.as_u16() as u32 == KIND_AGENT_OBSERVER_FRAME - && (state.check_rate_gate().is_some() || !state.gated_observer_pending.is_empty()) + && (state.check_rate_gate().is_some() + || !state.gated_observer_priority_pending.is_empty() + || !state.gated_observer_pending.is_empty()) { debug!( + priority_pending = state.gated_observer_priority_pending.len(), pending = state.gated_observer_pending.len(), "rate-gated: parking observer frame for paced drain" ); - state.park_gated_observer_frame(event); + let admitted = state.park_gated_observer_frame(event); + state.record_observer_confirmation(event_id, is_priority, admission, admitted); return true; } // Drop remaining ephemeral publishes while rate-gated. Stale typing @@ -1494,18 +2099,66 @@ async fn execute_connected_command( debug!("rate-gated: dropping ephemeral PublishEvent (typing indicator)"); return true; } - // Best-effort: log a send failure but don't trigger reconnect — the - // next ping or read will detect the dead socket. A failed observer - // frame is parked so the post-reconnect drain redelivers it. + // A failed observer frame is parked so the post-reconnect drain + // redelivers it. If `start_send` already accepted bytes, any + // cancellation, timeout, or error taints the socket and forces + // reconnect before later traffic can flush those buffered bytes. let is_observer = event.kind.as_u16() as u32 == KIND_AGENT_OBSERVER_FRAME; - if send_publish_event_frame(ws, &event).await { - if is_observer { - state.track_observer_in_flight(event); + if is_priority && !state.has_observer_priority_capacity() { + state.gated_observer_dropped += 1; + tracing::error!( + dropped_total = state.gated_observer_dropped, + "protected observer acknowledgment state full — rejected newest terminal result before socket write" + ); + state.record_observer_confirmation(event_id, true, admission, false); + return true; + } + let outcome = + send_publish_event_frame_until_cancelled(ws, &event, admission.as_mut()).await; + match outcome { + PublishEventFrameOutcome::CancelledBeforeSend => { + debug!( + event_id, + "terminal observer publish cancelled before socket write started" + ); + } + PublishEventFrameOutcome::SocketTainted => { + warn!( + event_id, + "terminal observer publish cancelled after socket write started — retiring relay socket" + ); + } + PublishEventFrameOutcome::Sent => { + if is_observer { + let admitted = state.track_observer_in_flight(event); + state.record_observer_confirmation( + event_id, + is_priority, + admission, + admitted, + ); + } else if let Some(admission) = admission { + let _ = + admission + .send(Err("confirmed delivery is only valid for observer frames" + .to_string())); + } + } + PublishEventFrameOutcome::Failed + | PublishEventFrameOutcome::FailedSocketTainted + if is_observer => + { + let admitted = state.park_gated_observer_frame(event); + state.record_observer_confirmation(event_id, is_priority, admission, admitted); + } + PublishEventFrameOutcome::Failed + | PublishEventFrameOutcome::FailedSocketTainted => { + if let Some(admission) = admission { + let _ = admission.send(Err("relay publish failed".to_string())); + } } - } else if is_observer { - state.park_gated_observer_frame(event); } - true + !outcome.socket_tainted() } RelayCommand::SetStartupWatermark { ts } => { state.startup_watermark = Some(ts); @@ -1541,8 +2194,11 @@ async fn run_background_task( relay_url: String, agent_pubkey_hex: String, auth_tag: Option, + configured_owner: Option, ) { let mut state = BgState::new(); + state.authorized_owner_pubkey = + resolve_authorized_owner_pubkey(auth_tag.as_ref(), &keys, configured_owner.as_deref()); let handshake_ok = process_handshake_buffer( &mut ws, @@ -1574,14 +2230,7 @@ async fn run_background_task( ) .await { - ReconnectOutcome::Ok => { - if matches!( - drain_post_reconnect(&mut ws, &mut cmd_rx, &mut state, &agent_pubkey_hex).await, - ReconnectOutcome::Shutdown - ) { - return; - } - } + ReconnectOutcome::Ok => {} ReconnectOutcome::Shutdown => return, ReconnectOutcome::Failed => { if matches!( @@ -1657,20 +2306,7 @@ async fn run_background_task( ) .await { - ReconnectOutcome::Ok => { - if matches!( - drain_post_reconnect( - &mut ws, - &mut cmd_rx, - &mut state, - &agent_pubkey_hex - ) - .await, - ReconnectOutcome::Shutdown - ) { - return; - } - } + ReconnectOutcome::Ok => {} ReconnectOutcome::Shutdown => return, ReconnectOutcome::Failed => { if matches!( @@ -1731,7 +2367,11 @@ async fn run_background_task( } } if state.observer_resub_needed && budget > 0 { - if send_observer_control_subscribe(&mut ws, &agent_pubkey_hex).await { + state.observer_control_subscribed_at = None; + if let Some(subscribed_at) = + send_observer_control_subscribe(&mut ws, &agent_pubkey_hex).await + { + state.observer_control_subscribed_at = Some(subscribed_at); state.observer_resub_needed = false; budget = budget.saturating_sub(1); any_sent = true; @@ -1762,16 +2402,54 @@ async fn run_background_task( } } - if budget > 0 && !state.gated_observer_pending.is_empty() { - let sent = drain_gated_observer_pending(&mut ws, &mut state, budget).await; - if sent > 0 { - any_sent = true; - } - } - - if any_sent { + if budget > 0 + && (!state.gated_observer_priority_pending.is_empty() + || !state.gated_observer_pending.is_empty()) + { + match drain_gated_observer_pending(&mut ws, &mut state, budget).await { + GatedObserverDrainOutcome::Sent(sent) => { + if sent > 0 { + any_sent = true; + } + } + GatedObserverDrainOutcome::SocketTainted => { + warn!( + "terminal observer send was cancelled after start_send — reconnecting before later traffic" + ); + let _ = event_tx.try_send(None); + if matches!( + wait_for_reconnect( + &mut ws, + &mut cmd_rx, + &mut state, + &keys, + &relay_url, + &agent_pubkey_hex, + &event_tx, + &observer_control_tx, + true, + auth_tag.as_ref(), + ) + .await, + ReconnectOutcome::Shutdown + ) { + return; + } + ping_sent = false; + last_pong = Instant::now(); + connected_since = Instant::now(); + stable_logged = false; + drain_pacing_next = None; + continue; + } + } + } + + if any_sent { drain_pacing_next = Some(tokio::time::Instant::now() + REQ_PACING_INTERVAL); - } else if !state.gated_observer_pending.is_empty() { + } else if !state.gated_observer_priority_pending.is_empty() + || !state.gated_observer_pending.is_empty() + { // Nothing sent because the gate is still armed. Arm the pacing // timer to the gate deadline so parked observer frames drain // promptly even when no other traffic wakes the select loop. @@ -1782,6 +2460,23 @@ async fn run_background_task( } tokio::select! { + permit = observer_control_tx.reserve(), + if !state.observer_control_pending.is_empty() => + { + match permit { + Ok(permit) => { + deliver_next_pending_observer_control(&mut state, permit); + } + Err(_) => { + warn!( + pending = state.observer_control_pending.len(), + "observer control receiver closed with pending controls" + ); + return; + } + } + } + raw = ws.next() => { // Determine if the socket is lost. let socket_lost = match raw { @@ -1835,10 +2530,6 @@ async fn run_background_task( 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(); @@ -1908,12 +2599,7 @@ async fn run_background_task( auth_tag.as_ref(), ).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::Ok => {} ReconnectOutcome::Failed => { if matches!( wait_for_reconnect( @@ -1947,12 +2633,7 @@ async fn run_background_task( auth_tag.as_ref(), ).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::Ok => {} ReconnectOutcome::Failed => { if matches!( wait_for_reconnect( @@ -1980,12 +2661,7 @@ async fn run_background_task( auth_tag.as_ref(), ).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::Ok => {} ReconnectOutcome::Failed => { if matches!( wait_for_reconnect( @@ -2035,6 +2711,90 @@ async fn run_background_task( } } +/// Queue one fully authenticated observer control without awaiting downstream +/// capacity. Returns `false` only when the downstream receiver is closed. +fn queue_admitted_observer_control( + state: &mut BgState, + observer_control_tx: &mpsc::Sender, + event: Event, +) -> bool { + if state.observer_control_pending.is_empty() { + match observer_control_tx.try_send(event) { + Ok(()) => return true, + Err(mpsc::error::TrySendError::Full(event)) => { + state.observer_control_pending.push_back(event); + return true; + } + Err(mpsc::error::TrySendError::Closed(event)) => { + state.observer_control_replay.remove(&event.id.to_hex()); + warn!( + event_id = %event.id, + "observer control receiver closed before delivery" + ); + return false; + } + } + } + + if state.observer_control_pending.len() >= OBSERVER_CONTROL_PENDING_CAPACITY { + let now = unix_now_secs(); + let mut expired_ids = Vec::new(); + state.observer_control_pending.retain(|pending| { + let ts = pending.created_at.as_secs(); + let expired = now.saturating_sub(ts) > OBSERVER_CONTROL_FRESHNESS_SECS; + if expired { + expired_ids.push(pending.id.to_hex()); + warn!( + event_id = %pending.id, + event_ts = ts, + delivery_now = now, + "dropping observer control that expired while awaiting downstream capacity" + ); + } + !expired + }); + for event_id in expired_ids { + state.observer_control_replay.remove(&event_id); + } + } + + if state.observer_control_pending.len() >= OBSERVER_CONTROL_PENDING_CAPACITY { + state.observer_control_replay.remove(&event.id.to_hex()); + warn!( + event_id = %event.id, + capacity = OBSERVER_CONTROL_PENDING_CAPACITY, + "observer control pending queue full — failing closed without eviction" + ); + return true; + } + + state.observer_control_pending.push_back(event); + true +} + +/// Deliver the oldest pending observer control through a reserved downstream +/// slot, discarding only controls whose bounded freshness window elapsed while +/// they waited. +fn deliver_next_pending_observer_control(state: &mut BgState, permit: mpsc::Permit<'_, Event>) { + let delivery_now = unix_now_secs(); + while let Some(event) = state.observer_control_pending.pop_front() { + let event_id = event.id.to_hex(); + let ts = event.created_at.as_secs(); + if delivery_now.saturating_sub(ts) > OBSERVER_CONTROL_FRESHNESS_SECS { + state.observer_control_replay.remove(&event_id); + warn!( + event_id, + event_ts = ts, + delivery_now, + "dropping observer control that expired while awaiting downstream capacity" + ); + continue; + } + permit.send(event); + return; + } +} + /// Handle a single WebSocket message in the background task. /// /// Returns `false` if the connection has been lost (Close frame or unrecoverable @@ -2067,22 +2827,192 @@ async fn handle_ws_message( event, } => { if subscription_id == OBSERVER_CONTROL_SUB_ID { - match observer_control_tx.try_send(*event) { - Ok(()) => {} - Err(mpsc::error::TrySendError::Full(_)) => { - warn!("observer control event dropped because control channel is full"); + let Some(subscription_started_at) = state.observer_control_subscribed_at + else { + warn!( + event_id = %event.id, + "dropping observer control for inactive subscription" + ); + return true; + }; + if !state.observer_control_sub_active { + warn!( + event_id = %event.id, + "dropping observer control without subscription intent" + ); + return true; + } + if !is_authorized_observer_control( + &event, + agent_pubkey_hex, + state.authorized_owner_pubkey.as_deref(), + ) { + warn!( + event_id = %event.id, + "dropping observer frame outside the exact owner-control shape" + ); + return true; + } + + let ts = event.created_at.as_secs(); + let now = unix_now_secs(); + let admission_boundary = state + .privileged_control_start_boundary + .max(subscription_started_at); + if ts <= admission_boundary { + warn!( + event_id = %event.id, + event_ts = ts, + admission_boundary, + "dropping observer control that does not provably post-date startup and subscription" + ); + return true; + } + if ts > now || now - ts > OBSERVER_CONTROL_FRESHNESS_SECS { + warn!( + event_id = %event.id, + event_ts = ts, + now, + "dropping observer control outside the freshness window" + ); + return true; + } + + let event_for_verify = event.clone(); + match tokio::task::spawn_blocking(move || { + buzz_core::verify_event(&event_for_verify) + }) + .await + { + Ok(Ok(())) => {} + Ok(Err(error)) => { + warn!( + event_id = %event.id, + "dropping invalid observer control: {error}" + ); + return true; } - Err(mpsc::error::TrySendError::Closed(_)) => return false, + Err(error) => { + warn!( + event_id = %event.id, + "observer control verification task failed: {error}" + ); + return true; + } + } + + let event_id = event.id.to_hex(); + let expires_at = ts.saturating_add(OBSERVER_CONTROL_FRESHNESS_SECS); + match state + .observer_control_replay + .admit(event_id, expires_at, now) + { + PrivilegedControlAdmission::Admitted => {} + PrivilegedControlAdmission::Replay => { + warn!( + event_id = %event.id, + "dropping replayed observer control" + ); + return true; + } + PrivilegedControlAdmission::Full => { + warn!( + event_id = %event.id, + capacity = OBSERVER_CONTROL_REPLAY_CAPACITY, + "observer control replay guard full — failing closed" + ); + return true; + } + } + + if !queue_admitted_observer_control(state, observer_control_tx, *event) { + return false; } } else if subscription_id == MEMBERSHIP_NOTIF_SUB_ID { - // Membership notification — extract channel UUID from h tag. + if !state.membership_sub_active { + warn!( + event_id = %event.id, + "dropping membership notification for inactive subscription" + ); + return true; + } + + let kind = event.kind.as_u16() as u32; + if !matches!( + kind, + KIND_MEMBER_ADDED_NOTIFICATION | KIND_MEMBER_REMOVED_NOTIFICATION + ) { + warn!( + kind, + event_id = %event.id, + "dropping unexpected event kind on membership subscription" + ); + return true; + } + + if !event_has_tag_value(&event, "p", agent_pubkey_hex) { + warn!( + event_id = %event.id, + "dropping membership notification targeting another agent" + ); + return true; + } + + // Membership notification — extract the one unambiguous + // channel UUID from its h tag. let channel_uuid = match extract_h_tag_uuid(&event) { Some(uuid) => uuid, None => { - warn!("membership notification missing h tag — dropping"); + warn!( + event_id = %event.id, + "membership notification missing or ambiguous h tag — dropping" + ); return true; } }; + + // Relays can replay stored events without authenticating + // them. Verify before any dedup, watermark, or forwarding + // mutation so a forged frame cannot reserve its claimed + // ID and suppress a later authentic notification. + let event_for_verify = event.clone(); + match tokio::task::spawn_blocking(move || { + buzz_core::verify_event(&event_for_verify) + }) + .await + { + Ok(Ok(())) => {} + Ok(Err(error)) => { + warn!( + channel_id = %channel_uuid, + event_id = %event.id, + "dropping invalid membership notification: {error}" + ); + return true; + } + Err(error) => { + warn!( + channel_id = %channel_uuid, + event_id = %event.id, + "membership verification task failed: {error}" + ); + return true; + } + } + + let ts = event.created_at.as_secs(); + let now = unix_now_secs(); + if ts > now { + warn!( + channel_id = %channel_uuid, + event_id = %event.id, + event_ts = ts, + now, + "dropping future-dated membership notification" + ); + return true; + } + // Dedup membership notifications through TwoGenDedup. // We use seen_ids directly instead of record_event() // because record_event() also updates last_seen, which @@ -2098,10 +3028,10 @@ async fn handle_ws_message( ); return true; } - let ts = event.created_at.as_secs(); let buzz_event = BuzzEvent { channel_id: channel_uuid, event: *event, + privileged_control_admitted: false, }; let cap = event_tx.max_capacity(); let used = cap - event_tx.capacity(); @@ -2137,12 +3067,186 @@ async fn handle_ws_message( Err(mpsc::error::TrySendError::Closed(_)) => return false, } } else if let Some(channel_id) = channel_id_from_sub_id(&subscription_id) { + if state + .active_subscriptions + .get(&channel_id) + .map(String::as_str) + != Some(subscription_id.as_str()) + { + warn!( + channel_id = %channel_id, + event_id = %event.id, + "dropping event for inactive channel subscription" + ); + return true; + } + + if extract_h_tag_uuid(&event) != Some(channel_id) { + warn!( + channel_id = %channel_id, + event_id = %event.id, + "dropping event whose h tag does not match its subscription" + ); + return true; + } + + let Some(filter) = state.active_filters.get(&channel_id) else { + warn!( + channel_id = %channel_id, + event_id = %event.id, + "dropping event without a local subscription filter" + ); + return true; + }; + let kind = event.kind.as_u16() as u32; + if filter + .kinds + .as_ref() + .is_some_and(|kinds| !kinds.contains(&kind)) + { + warn!( + channel_id = %channel_id, + event_id = %event.id, + kind, + "dropping event outside the local kind filter" + ); + return true; + } + if filter.require_mention + && !event_has_tag_value(&event, "p", agent_pubkey_hex) + { + warn!( + channel_id = %channel_id, + event_id = %event.id, + "dropping event outside the local mention filter" + ); + return true; + } + + // Relays are not trusted to have verified stored or + // replayed events. Verify before touching replay + // watermarks or dedup state so a forged frame cannot + // reserve the claimed ID and suppress a later authentic + // event. Schnorr verification is CPU-bound. + let event_for_verify = event.clone(); + match tokio::task::spawn_blocking(move || { + buzz_core::verify_event(&event_for_verify) + }) + .await + { + Ok(Ok(())) => {} + Ok(Err(error)) => { + warn!( + channel_id = %channel_id, + event_id = %event.id, + "dropping invalid channel event: {error}" + ); + return true; + } + Err(error) => { + warn!( + channel_id = %channel_id, + event_id = %event.id, + "channel event verification task failed: {error}" + ); + return true; + } + } let ts = event.created_at.as_secs(); let event_id_hex = event.id.to_hex(); - if state.record_event(channel_id, &event) { + let now = unix_now_secs(); + let future_skew = ts.saturating_sub(now); + if future_skew > CHANNEL_EVENT_FUTURE_SKEW_SECS { + warn!( + channel_id = %channel_id, + event_id = %event.id, + event_ts = ts, + now, + "dropping future-dated channel event" + ); + return true; + } + let is_privileged_control_candidate = is_privileged_channel_control( + &event, + agent_pubkey_hex, + state.authorized_owner_pubkey.as_deref(), + ); + let privileged_control_boundary = + state.channel_subscribed_at.get(&channel_id).copied().map( + |subscribed_at| { + subscribed_at.max(state.privileged_control_start_boundary) + }, + ); + let privileged_control_admitted = if !is_privileged_control_candidate { + false + } else if ts > now { + warn!( + channel_id = %channel_id, + event_id = %event.id, + event_ts = ts, + now, + "future-skewed owner control — forwarding without execution authority" + ); + false + } else if privileged_control_boundary.is_none() { + warn!( + channel_id = %channel_id, + event_id = %event.id, + "owner control arrived without a successful live channel REQ — forwarding without execution authority" + ); + false + } else if privileged_control_boundary.is_some_and(|boundary| ts <= boundary) + { + warn!( + channel_id = %channel_id, + event_id = %event.id, + event_ts = ts, + admission_boundary = privileged_control_boundary, + "owner control does not provably post-date supervisor/channel subscription startup — forwarding without execution authority; retry with a newly signed event after the boundary" + ); + false + } else if now - ts > PRIVILEGED_CONTROL_FRESHNESS_SECS { + warn!( + channel_id = %channel_id, + event_id = %event.id, + event_ts = ts, + now, + "stale owner control — forwarding without execution authority" + ); + false + } else { + let expires_at = ts.saturating_add(PRIVILEGED_CONTROL_FRESHNESS_SECS); + match state.privileged_control_replay.admit( + event_id_hex.clone(), + expires_at, + now, + ) { + PrivilegedControlAdmission::Admitted => true, + PrivilegedControlAdmission::Replay => { + warn!( + channel_id = %channel_id, + event_id = %event.id, + "replayed owner control — withholding execution authority" + ); + false + } + PrivilegedControlAdmission::Full => { + warn!( + channel_id = %channel_id, + event_id = %event.id, + capacity = PRIVILEGED_CONTROL_REPLAY_CAPACITY, + "privileged control replay guard full — failing closed" + ); + false + } + } + }; + let replay_timestamp = ts.min(now); + if state.record_event_with_watermark(channel_id, &event, replay_timestamp) { let buzz_event = BuzzEvent { channel_id, event: *event, + privileged_control_admitted, }; // Warn at 80% capacity. let cap = event_tx.max_capacity(); @@ -2160,18 +3264,24 @@ async fn handle_ws_message( // Remove from dedup set so the replayed event // won't be rejected as a duplicate after reconnect. state.seen_ids.remove(&event_id_hex); + if privileged_control_admitted { + // The harness never received this control, + // so a reconnect replay must remain eligible. + state.privileged_control_replay.remove(&event_id_hex); + } // Track the oldest dropped timestamp so reconnect // replay starts early enough to re-deliver it. state .channel_dropped_since .entry(channel_id) - .and_modify(|d| *d = (*d).min(ts)) - .or_insert(ts); + .and_modify(|d| *d = (*d).min(replay_timestamp)) + .or_insert(replay_timestamp); // Proactively trigger resubscribe without waiting for a disconnect. state.proactive_resubscribe_needed = true; warn!( channel_id = %channel_id, - ts, + event_ts = ts, + replay_timestamp, "event channel full — dropping event for channel {channel_id} — proactive resubscribe queued" ); } @@ -2211,6 +3321,12 @@ async fn handle_ws_message( subscription_id, message, } => { + if subscription_id == OBSERVER_CONTROL_SUB_ID { + state.observer_control_subscribed_at = None; + } + if let Some(channel_id) = channel_id_from_sub_id(&subscription_id) { + state.channel_subscribed_at.remove(&channel_id); + } // A per-channel membership denial means THIS channel is // forbidden, not the whole connection. Drop just this // channel's subscription and keep the socket — otherwise the @@ -2269,9 +3385,11 @@ async fn handle_ws_message( // resubscribe_after_reconnect() needs the subscription to // still be in state so it can restore it. if subscription_id == OBSERVER_CONTROL_SUB_ID { - let sent = send_observer_control_subscribe(ws, agent_pubkey_hex).await; - if sent { + if let Some(subscribed_at) = + send_observer_control_subscribe(ws, agent_pubkey_hex).await + { state.observer_control_sub_active = true; + state.observer_control_subscribed_at = Some(subscribed_at); } else { warn!("observer control resubscribe failed after CLOSED — triggering reconnect"); return false; @@ -2361,7 +3479,29 @@ async fn handle_ws_message( warn!("mid-session AUTH rejected (event {event_id}): {message} — triggering reconnect"); return false; } - state.acknowledge_observer_frame(&event_id); + if accepted { + state.acknowledge_observer_frame(&event_id); + } else if message.starts_with("rate-limited:") { + let secs = parse_rate_limit_retry_secs(&message).unwrap_or(0); + let deadline = state.set_rate_limit_gate(secs); + if state.requeue_rate_limited_observer_frame(&event_id) { + warn!( + event_id, + "observer publish rate-limited via OK false — exact frame requeued; gate armed until ~{:.1}s", + deadline + .checked_duration_since(tokio::time::Instant::now()) + .unwrap_or_default() + .as_secs_f64() + ); + } + } else if state.record_terminal_observer_denial(&event_id, &message) { + warn!( + event_id, + denial = message, + denied_total = state.observer_terminal_denied, + "observer publish terminally denied — frame retired without retry" + ); + } debug!("OK for event {event_id}: accepted={accepted} message={message}"); } } @@ -2581,6 +3721,7 @@ async fn resubscribe_after_reconnect( } if state.observer_control_sub_active { + state.observer_control_subscribed_at = None; if state.check_rate_gate().is_some() { debug!("rate-gated: parking observer control resubscribe after reconnect"); state.observer_resub_needed = true; @@ -2588,10 +3729,15 @@ async fn resubscribe_after_reconnect( if !pacing_sleep(cmd_rx, &mut deferred_commands, REQ_PACING_INTERVAL).await { return ResubscribeResult::Shutdown; } - if !send_observer_control_subscribe(ws, agent_pubkey_hex).await { - warn!("failed to resubscribe observer controls after reconnect"); - retain_deferred_command_intent(state, &mut deferred_commands); - return ResubscribeResult::RetryConnection; + match send_observer_control_subscribe(ws, agent_pubkey_hex).await { + Some(subscribed_at) => { + state.observer_control_subscribed_at = Some(subscribed_at); + } + None => { + warn!("failed to resubscribe observer controls after reconnect"); + retain_deferred_command_intent(state, &mut deferred_commands); + return ResubscribeResult::RetryConnection; + } } state.observer_resub_needed = false; } @@ -2604,58 +3750,243 @@ async fn resubscribe_after_reconnect( } } -/// Send a signed EVENT frame on the live socket. Returns `false` on send failure. +/// Sink adapter that records whether `SinkExt::send` reached `start_send`. /// -/// Best-effort at the socket level: a failure is logged but does not trigger -/// reconnect — the next ping or read will detect the dead socket. -async fn send_publish_event_frame(ws: &mut WsStream, event: &Event) -> bool { - 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 - { - warn!("failed to publish event: {e}"); - return false; - } +/// Before that boundary cancellation is clean. After it, the sink may retain +/// buffered bytes even when the send future is dropped, so the connection must +/// be retired before any later frame is written. +struct StartSendTrackedSink<'a, S> { + inner: &'a mut S, + started: &'a AtomicBool, +} + +impl Sink for StartSendTrackedSink<'_, S> +where + S: Sink + Unpin, +{ + type Error = S::Error; + + fn poll_ready( + self: std::pin::Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + let this = self.get_mut(); + std::pin::Pin::new(&mut *this.inner).poll_ready(cx) + } + + fn start_send(self: std::pin::Pin<&mut Self>, item: Message) -> Result<(), Self::Error> { + let this = self.get_mut(); + // Conservatively taint as soon as `start_send` is invoked. Some sinks, + // including tungstenite write-failure paths, may retain the frame even + // when this call returns an error. + this.started.store(true, Ordering::SeqCst); + std::pin::Pin::new(&mut *this.inner).start_send(item) + } + + fn poll_flush( + self: std::pin::Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + let this = self.get_mut(); + std::pin::Pin::new(&mut *this.inner).poll_flush(cx) + } + + fn poll_close( + self: std::pin::Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + let this = self.get_mut(); + std::pin::Pin::new(&mut *this.inner).poll_close(cx) } - true } -/// Drain parked observer telemetry frames once the rate-limit gate clears. -/// -/// Called by the main loop pacing timer. Sends at most `budget` frames without -/// sleeping — pacing is enforced by the caller via `drain_pacing_next`. Stops -/// immediately if the gate re-arms mid-drain. When the queue empties, any +/// Send a signed EVENT frame on a sink. Returns `false` on send failure. +async fn send_publish_event_frame(ws: &mut S, event: &Event) -> bool +where + S: Sink + Unpin, + S::Error: std::fmt::Display, +{ + let msg = json!(["EVENT", event]); + if let Ok(text) = serde_json::to_string(&msg) { + match tokio::time::timeout( + Duration::from_secs(WS_SEND_TIMEOUT_SECS), + ws.send(Message::Text(text.into())), + ) + .await + { + Ok(Ok(())) => {} + Ok(Err(e)) => { + warn!("failed to publish event: {e}"); + return false; + } + Err(_) => { + warn!("failed to publish event: socket send timed out"); + return false; + } + } + } + true +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum PublishEventFrameOutcome { + Sent, + Failed, + CancelledBeforeSend, + /// Cancellation won after `start_send`; buffered bytes may remain. + SocketTainted, + /// The send failed or timed out after `start_send`; buffered bytes may remain. + FailedSocketTainted, +} + +impl PublishEventFrameOutcome { + fn socket_tainted(self) -> bool { + matches!(self, Self::SocketTainted | Self::FailedSocketTainted) + } +} + +/// Write an EVENT only while its terminal confirmation receiver remains open. +/// +/// The publisher drops that receiver at its end-to-end deadline. Selecting on +/// `closed()` prevents an expired command from reaching a reconnected socket +/// even when its FIFO cancellation command has not yet been processed. If +/// cancellation wins after `start_send`, the caller must retire the socket: +/// dropping `SinkExt::send` does not retract bytes already buffered by the sink. +async fn send_publish_event_frame_until_cancelled( + ws: &mut S, + event: &Event, + admission: Option<&mut oneshot::Sender>>, +) -> PublishEventFrameOutcome +where + S: Sink + Unpin, + S::Error: std::fmt::Display, +{ + let send_started = AtomicBool::new(false); + let sent = if let Some(admission) = admission { + if admission.is_closed() { + return PublishEventFrameOutcome::CancelledBeforeSend; + } + let selected = { + let mut tracked = StartSendTrackedSink { + inner: ws, + started: &send_started, + }; + tokio::select! { + biased; + _ = admission.closed() => None, + sent = send_publish_event_frame(&mut tracked, event) => Some(sent), + } + }; + match selected { + Some(sent) => sent, + None if send_started.load(Ordering::SeqCst) => { + return PublishEventFrameOutcome::SocketTainted; + } + None => return PublishEventFrameOutcome::CancelledBeforeSend, + } + } else { + let mut tracked = StartSendTrackedSink { + inner: ws, + started: &send_started, + }; + send_publish_event_frame(&mut tracked, event).await + }; + if sent { + PublishEventFrameOutcome::Sent + } else if send_started.load(Ordering::SeqCst) { + PublishEventFrameOutcome::FailedSocketTainted + } else { + PublishEventFrameOutcome::Failed + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum GatedObserverDrainOutcome { + Sent(usize), + SocketTainted, +} + +/// Drain parked observer telemetry frames once the rate-limit gate clears. +/// +/// Called by the main loop pacing timer. Sends at most `budget` frames without +/// sleeping — pacing is enforced by the caller via `drain_pacing_next`. Stops +/// immediately if the gate re-arms mid-drain. When the queue empties, any /// overflow loss is summarized in one warning. Returns the number of frames sent. async fn drain_gated_observer_pending( ws: &mut WsStream, state: &mut BgState, budget: usize, -) -> usize { +) -> GatedObserverDrainOutcome { let mut sent = 0; while sent < budget { + state.prune_cancelled_terminal_admissions(); if state.check_rate_gate().is_some() { break; } - let Some(event) = state.gated_observer_pending.pop_front() else { + let event = state + .gated_observer_priority_pending + .pop_front() + .or_else(|| state.gated_observer_pending.pop_front()); + let Some(event) = event else { break; }; - if !send_publish_event_frame(ws, &event).await { - // Socket may be dead — re-park at the front so the frame survives - // reconnect (the post-reconnect drain will retry it in order). - state.gated_observer_pending.push_front(event); - break; + let event_id = event.id.to_hex(); + let is_priority = is_priority_observer_frame(&event); + let mut confirmation = is_priority + .then(|| state.observer_priority_confirmations.remove(&event_id)) + .flatten(); + let outcome = + send_publish_event_frame_until_cancelled(ws, &event, confirmation.as_mut()).await; + match outcome { + PublishEventFrameOutcome::CancelledBeforeSend => continue, + PublishEventFrameOutcome::SocketTainted => { + return GatedObserverDrainOutcome::SocketTainted; + } + PublishEventFrameOutcome::Failed | PublishEventFrameOutcome::FailedSocketTainted => { + if confirmation + .as_ref() + .is_some_and(oneshot::Sender::is_closed) + { + if outcome.socket_tainted() { + return GatedObserverDrainOutcome::SocketTainted; + } + continue; + } + // Socket may be dead — re-park at the front so the frame survives + // reconnect (the post-reconnect drain will retry it in order). + if is_priority { + state.gated_observer_priority_pending.push_front(event); + if let Some(confirmation) = confirmation { + state + .observer_priority_confirmations + .insert(event_id, confirmation); + } + } else { + state.gated_observer_pending.push_front(event); + } + if outcome.socket_tainted() { + return GatedObserverDrainOutcome::SocketTainted; + } + break; + } + PublishEventFrameOutcome::Sent => { + let admitted = state.track_observer_in_flight(event); + state.record_observer_confirmation(event_id, is_priority, confirmation, admitted); + sent += 1; + } } - state.track_observer_in_flight(event); - sent += 1; } - if state.gated_observer_pending.is_empty() && state.gated_observer_dropped > 0 { + if state.gated_observer_priority_pending.is_empty() + && state.gated_observer_pending.is_empty() + && state.gated_observer_dropped > 0 + { warn!( observer_frames_dropped = state.gated_observer_dropped, "observer frames lost to gated-queue overflow" ); state.gated_observer_dropped = 0; } - sent + GatedObserverDrainOutcome::Sent(sent) } /// Drain `rate_limited_pending` channels whose retry deadline has passed. @@ -2813,7 +4144,10 @@ async fn drain_commands( if send_failed { match cmd { RelayCommand::Shutdown => { - let _ = ws_send_timeout(ws, Message::Close(None), WS_SEND_TIMEOUT_SECS).await; + // The failed send may have left bytes buffered after + // `start_send`. Do not write even a Close frame through + // this socket; returning drops or replaces it without a + // later flush. return ReconnectOutcome::Shutdown; } RelayCommand::Reconnect => {} @@ -2902,6 +4236,8 @@ async fn try_autonomous_reconnect( auth_tag: Option<&nostr::Tag>, ) -> ReconnectOutcome { state.requeue_observer_in_flight(); + state.observer_control_subscribed_at = None; + state.channel_subscribed_at.clear(); // 5 attempts, up to 16s base backoff. Shares delay values with the // initial-connect retry in `HarnessRelay::connect()` (STARTUP_CONNECT_BACKOFFS) — // see its doc comment for how the two loops consume the array differently. @@ -2949,7 +4285,23 @@ async fn try_autonomous_reconnect( match resubscribe_after_reconnect(ws, cmd_rx, state, agent_pubkey_hex, true) .await { - ResubscribeResult::Ok => return ReconnectOutcome::Ok, + ResubscribeResult::Ok => { + // Commands may arrive after the resubscribe helper's + // own drain. Finish that window here, and never return + // a socket whose final send failed or was tainted. + match drain_post_reconnect(ws, cmd_rx, state, agent_pubkey_hex).await { + ReconnectOutcome::Ok => return ReconnectOutcome::Ok, + ReconnectOutcome::Shutdown => { + return ReconnectOutcome::Shutdown; + } + ReconnectOutcome::Failed => { + warn!( + "post-reconnect command drain tainted or lost the autonomous socket — treating as failed attempt" + ); + // Fall through to backoff sleep and retry. + } + } + } ResubscribeResult::Shutdown => return ReconnectOutcome::Shutdown, ResubscribeResult::RetryConnection => { warn!("resubscribe failed after autonomous reconnect — treating as failed attempt"); @@ -2978,6 +4330,11 @@ async fn try_autonomous_reconnect( } } + // Any writes and subscription timestamps created on this failed + // candidate belong to that socket. Preserve unresolved publications, + // but never let its socket-bound authority reach the next candidate. + state.abandon_failed_reconnect_candidate(); + // Backoff sleep between ladder 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. @@ -3032,6 +4389,8 @@ async fn wait_for_reconnect( auth_tag: Option<&nostr::Tag>, ) -> ReconnectOutcome { state.requeue_observer_in_flight(); + state.observer_control_subscribed_at = None; + state.channel_subscribed_at.clear(); if !skip_drain { // Drain commands until we get Reconnect (or Shutdown). // Other commands update state so reconnect reflects latest intent. @@ -3086,8 +4445,21 @@ async fn wait_for_reconnect( { ResubscribeResult::Ok => { // Drain any commands that arrived during do_connect() + - // resubscribe (which don't poll cmd_rx). - return drain_post_reconnect(ws, cmd_rx, state, agent_pubkey_hex).await; + // resubscribe (which don't poll cmd_rx). A failed + // drain may have tainted this replacement socket, so + // retry instead of returning it to the main loop. + match drain_post_reconnect(ws, cmd_rx, state, agent_pubkey_hex).await { + ReconnectOutcome::Ok => return ReconnectOutcome::Ok, + ReconnectOutcome::Shutdown => { + return ReconnectOutcome::Shutdown; + } + ReconnectOutcome::Failed => { + warn!( + "post-reconnect command drain tainted or lost the replacement socket — retrying" + ); + // Fall through to the backoff sleep below. + } + } } ResubscribeResult::Shutdown => return ReconnectOutcome::Shutdown, ResubscribeResult::RetryConnection => { @@ -3114,6 +4486,11 @@ async fn wait_for_reconnect( } } + // A failed candidate may have written observer frames during replay or + // command drain. Requeue every unresolved frame and revoke timestamps + // that were authoritative only for that discarded socket. + state.abandon_failed_reconnect_candidate(); + // Persist ladder position before sleeping — if shutdown arrives mid-sleep, // the next session resumes from here rather than restarting at 0. state.backoff_step = attempt; @@ -3159,12 +4536,15 @@ async fn wait_for_reconnect( /// Returns `true` if the REQ was successfully written to the WebSocket. async fn send_subscribe( ws: &mut WsStream, - _state: &BgState, + state: &mut BgState, channel_id: Uuid, agent_pubkey_hex: &str, since: Option, filter: &ChannelFilter, ) -> bool { + // Clear before attempting the write so a failed or blocked replacement REQ + // can never inherit authority from the prior socket/subscription instance. + state.channel_subscribed_at.remove(&channel_id); let sub_id = channel_sub_id(channel_id); let mut req_filter = serde_json::Map::new(); @@ -3199,6 +4579,9 @@ async fn send_subscribe( Ok(text) => { match ws_send_timeout(ws, Message::Text(text.into()), WS_SEND_TIMEOUT_SECS).await { Ok(()) => { + state + .channel_subscribed_at + .insert(channel_id, unix_now_secs()); debug!( "subscribed to channel {channel_id}{}", if since.is_some() { @@ -3270,17 +4653,18 @@ async fn send_membership_subscribe( } /// Send a NIP-01 REQ for owner-to-agent observer control frames. -async fn send_observer_control_subscribe(ws: &mut WsStream, agent_pubkey_hex: &str) -> bool { +/// +/// Returns the second after the REQ is written successfully. Callers retain it +/// as the strict admission boundary for frames attributed to this subscription. +async fn send_observer_control_subscribe(ws: &mut WsStream, agent_pubkey_hex: &str) -> Option { + let replay_since = unix_now_secs(); let req = json!([ "REQ", OBSERVER_CONTROL_SUB_ID, { "kinds": [KIND_AGENT_OBSERVER_FRAME], "#p": [agent_pubkey_hex], - "since": std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap_or_default() - .as_secs(), + "since": replay_since, } ]); @@ -3289,17 +4673,17 @@ async fn send_observer_control_subscribe(ws: &mut WsStream, agent_pubkey_hex: &s match ws_send_timeout(ws, Message::Text(text.into()), WS_SEND_TIMEOUT_SECS).await { Ok(()) => { debug!("subscribed to observer control frames"); - true + Some(unix_now_secs()) } Err(e) => { warn!("failed to send observer control REQ: {e}"); - false + None } } } Err(e) => { warn!("failed to serialize observer control REQ: {e}"); - false + None } } } @@ -3344,6 +4728,21 @@ fn jittered_duration(base: Duration) -> Duration { base.mul_f64(factor) } +/// Add 0–20% positive jitter to an explicit relay reset horizon. +/// +/// Backoff attempts may safely jitter earlier, but a server-provided +/// `retry in Ns` value is a lower bound: retrying before it expires only +/// consumes another denied attempt and can perpetuate the shared gate. +fn rate_limit_retry_duration(base: Duration) -> Duration { + let nanos = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .subsec_nanos(); + // factor ∈ [1.0, 1.2) + let factor = 1.0 + (nanos as f64 / u32::MAX as f64) * 0.2; + base.mul_f64(factor) +} + /// Classify a `RelayError` as a DNS resolution failure. /// /// Matches the OS-level "name not found" strings surfaced by the platform's @@ -3414,18 +4813,98 @@ async fn dns_flat_sleep( } } -/// Extract a channel UUID from the h tag of a Nostr event. +/// Extract the one unambiguous channel UUID from the h tag of a Nostr event. +/// +/// Channel-scoped events must carry exactly one `h` tag. Accepting the first of +/// multiple values would let a validly-signed event match one relay filter while +/// being attributed to another channel locally. fn extract_h_tag_uuid(event: &nostr::Event) -> Option { - event.tags.iter().find_map(|tag| { + let mut h_tags = event.tags.iter().filter_map(|tag| { let tag_vec = tag.as_slice(); if tag_vec.len() >= 2 && tag_vec[0] == "h" { - tag_vec[1].parse::().ok() + Some(tag_vec[1].as_str()) } else { None } + }); + let channel_id = h_tags.next()?.parse::().ok()?; + if h_tags.next().is_some() { + return None; + } + Some(channel_id) +} + +fn event_has_tag_value(event: &nostr::Event, tag_name: &str, expected: &str) -> bool { + event.tags.iter().any(|tag| { + let tag_vec = tag.as_slice(); + tag_vec.first().map(String::as_str) == Some(tag_name) + && tag_vec.get(1).map(String::as_str) == Some(expected) }) } +fn is_priority_observer_frame(event: &Event) -> bool { + event.kind.as_u16() as u32 == KIND_AGENT_OBSERVER_FRAME + && event_has_tag_value( + event, + OBSERVER_PRIORITY_TAG, + OBSERVER_CONTROL_RESULT_PRIORITY, + ) +} + +fn is_exact_terminal_observer_frame(event: &Event) -> bool { + if event.kind.as_u16() as u32 != KIND_AGENT_OBSERVER_FRAME { + return false; + } + let mut priority_tags = event + .tags + .iter() + .filter(|tag| tag.as_slice().first().map(String::as_str) == Some(OBSERVER_PRIORITY_TAG)); + priority_tags.next().is_some_and(|tag| { + tag.as_slice().len() == 2 + && tag.as_slice().get(1).map(String::as_str) == Some(OBSERVER_CONTROL_RESULT_PRIORITY) + }) && priority_tags.next().is_none() +} + +fn event_has_single_tag_value(event: &nostr::Event, tag_name: &str, expected: &str) -> bool { + let mut values = event.tags.iter().filter_map(|tag| { + let tag_vec = tag.as_slice(); + (tag_vec.first().map(String::as_str) == Some(tag_name)) + .then(|| tag_vec.get(1).map(String::as_str)) + .flatten() + }); + values.next() == Some(expected) && values.next().is_none() +} + +fn is_authorized_observer_control( + event: &nostr::Event, + agent_pubkey_hex: &str, + authorized_owner_pubkey: Option<&str>, +) -> bool { + let Some(authorized_owner_pubkey) = authorized_owner_pubkey else { + return false; + }; + event.kind.as_u16() as u32 == KIND_AGENT_OBSERVER_FRAME + && event.pubkey.to_hex() == authorized_owner_pubkey + && event_has_single_tag_value(event, "p", agent_pubkey_hex) + && event_has_single_tag_value(event, OBSERVER_AGENT_TAG, agent_pubkey_hex) + && event_has_single_tag_value(event, OBSERVER_FRAME_TAG, OBSERVER_FRAME_CONTROL) + && content_looks_like_nip44(&event.content) +} + +fn is_privileged_channel_control( + event: &nostr::Event, + agent_pubkey_hex: &str, + authorized_owner_pubkey: Option<&str>, +) -> bool { + let Some(authorized_owner_pubkey) = authorized_owner_pubkey else { + return false; + }; + event.pubkey.to_hex() == authorized_owner_pubkey + && event.kind.as_u16() as u32 == KIND_STREAM_MESSAGE + && matches!(event.content.trim(), "!shutdown" | "!cancel" | "!rotate") + && event_has_tag_value(event, "p", agent_pubkey_hex) +} + /// Build and send a NIP-42 AUTH response event. /// /// If `auth_tag` is provided (NIP-OA owner attestation), it is included in the @@ -4337,762 +5816,2017 @@ mod tests { .expect("signing should succeed") } - async fn test_ws_pair() -> (WsStream, WebSocketStream) { - let listener = tokio::net::TcpListener::bind("127.0.0.1:0") - .await - .expect("bind test websocket"); - let address = listener.local_addr().expect("read test address"); - let server = tokio::spawn(async move { - let (stream, _) = listener.accept().await.expect("accept test websocket"); - tokio_tungstenite::accept_async(stream) - .await - .expect("complete server websocket handshake") - }); - let (client, _) = connect_async(format!("ws://{address}")) - .await - .expect("connect test websocket"); - (client, server.await.expect("join test websocket server")) - } - - async fn next_test_frame( - server: &mut WebSocketStream, - ) -> serde_json::Value { - let message = timeout(Duration::from_secs(1), server.next()) - .await - .expect("timed out waiting for websocket frame") - .expect("test websocket closed") - .expect("read test websocket frame"); - serde_json::from_str(message.to_text().expect("expected text frame")) - .expect("parse test websocket frame") + fn make_test_membership_event( + keys: &nostr::Keys, + channel_id: Uuid, + agent_pubkey_hex: &str, + created_at_secs: u64, + ) -> Event { + let channel_id = channel_id.to_string(); + let tags = [ + Tag::parse(["h", channel_id.as_str()]).expect("valid h tag"), + Tag::parse(["p", agent_pubkey_hex]).expect("valid p tag"), + ]; + EventBuilder::new( + Kind::Custom(KIND_MEMBER_ADDED_NOTIFICATION as u16), + "membership added", + ) + .tags(tags) + .custom_created_at(nostr::Timestamp::from(created_at_secs)) + .sign_with_keys(keys) + .expect("membership event signing should succeed") } - fn test_channel_filter() -> ChannelFilter { - ChannelFilter { - kinds: Some(vec![9]), - require_mention: false, + fn make_test_channel_event( + keys: &nostr::Keys, + channel_id: Uuid, + kind: u32, + content: &str, + mentioned_pubkey: Option<&str>, + created_at_secs: u64, + ) -> Event { + let channel_id = channel_id.to_string(); + let mut tags = vec![Tag::parse(["h", channel_id.as_str()]).expect("valid h tag")]; + if let Some(pubkey) = mentioned_pubkey { + tags.push(Tag::parse(["p", pubkey]).expect("valid p tag")); } - } - - fn seed_test_subscription(state: &mut BgState, channel_id: Uuid) { - apply_command_to_state( - state, - RelayCommand::Subscribe { - channel_id, - filter: test_channel_filter(), - replay_since: Some(1_000), - }, - ); + EventBuilder::new(Kind::Custom(kind as u16), content) + .tags(tags) + .custom_created_at(nostr::Timestamp::from(created_at_secs)) + .sign_with_keys(keys) + .expect("channel event signing should succeed") } #[tokio::test] - async fn fresh_reconnect_preserves_gate_until_pending_replay_resumes() { - let (mut client, mut server) = test_ws_pair().await; - let (_cmd_tx, mut cmd_rx) = mpsc::channel(1); + async fn membership_authentication_and_targeting_precede_dedup_and_forwarding() { + let (mut ws, _server) = test_ws_pair().await; + let (event_tx, mut event_rx) = mpsc::channel(4); + let (observer_control_tx, _observer_control_rx) = mpsc::channel(1); let mut state = BgState::new(); let channel_id = Uuid::new_v4(); - seed_test_subscription(&mut state, channel_id); - state.rate_limit_gate = Some(tokio::time::Instant::now() + Duration::from_millis(150)); + let agent_keys = Keys::generate(); + let relay_keys = Keys::generate(); + let agent_pubkey_hex = agent_keys.public_key().to_hex(); + let valid = + make_test_membership_event(&relay_keys, channel_id, &agent_pubkey_hex, 1_700_000); + let claimed_id = valid.id.to_hex(); + + let unsolicited_frame = Message::Text( + serde_json::to_string(&json!(["EVENT", MEMBERSHIP_NOTIF_SUB_ID, valid.clone()])) + .expect("serialize unsolicited membership frame") + .into(), + ); + assert!( + handle_ws_message( + unsolicited_frame, + &mut ws, + &event_tx, + &observer_control_tx, + &mut state, + &agent_keys, + "ws://localhost:3000", + &agent_pubkey_hex, + None, + ) + .await + ); + assert!( + !state.seen_ids.contains(&claimed_id), + "an inactive membership subscription must not admit or dedup events" + ); + assert!(event_rx.try_recv().is_err()); - let result = - resubscribe_after_reconnect(&mut client, &mut cmd_rx, &mut state, "agent-pubkey", true) - .await; + state.membership_sub_active = true; - assert!(matches!(result, ResubscribeResult::Ok)); - assert!(state.rate_limit_gate.is_some()); - assert!(state.rate_limited_pending.contains_key(&channel_id)); + let mut tampered_json = serde_json::to_value(&valid).expect("serialize membership event"); + tampered_json["content"] = Value::String("forged membership".to_string()); + let tampered: Event = + serde_json::from_value(tampered_json).expect("parse forged membership event"); + assert_eq!(tampered.id.to_hex(), claimed_id); + assert!(buzz_core::verify_event(&tampered).is_err()); + let forged_frame = Message::Text( + serde_json::to_string(&json!(["EVENT", MEMBERSHIP_NOTIF_SUB_ID, tampered])) + .expect("serialize forged membership frame") + .into(), + ); assert!( - timeout(Duration::from_millis(50), server.next()) - .await - .is_err(), - "fresh reconnect must not send REQ while the shared quota gate is active" + handle_ws_message( + forged_frame, + &mut ws, + &event_tx, + &observer_control_tx, + &mut state, + &agent_keys, + "ws://localhost:3000", + &agent_pubkey_hex, + None, + ) + .await + ); + assert!( + !state.seen_ids.contains(&claimed_id), + "a forged membership event must not poison dedup" + ); + assert!(event_rx.try_recv().is_err()); + + let wrong_target = + make_test_membership_event(&relay_keys, channel_id, "wrong-agent", 1_700_001); + let wrong_target_id = wrong_target.id.to_hex(); + let wrong_target_frame = Message::Text( + serde_json::to_string(&json!(["EVENT", MEMBERSHIP_NOTIF_SUB_ID, wrong_target])) + .expect("serialize wrong-target membership frame") + .into(), + ); + assert!( + handle_ws_message( + wrong_target_frame, + &mut ws, + &event_tx, + &observer_control_tx, + &mut state, + &agent_keys, + "ws://localhost:3000", + &agent_pubkey_hex, + None, + ) + .await + ); + assert!( + !state.seen_ids.contains(&wrong_target_id), + "a membership event targeting another agent must not poison dedup" ); + assert!(event_rx.try_recv().is_err()); - tokio::time::sleep(Duration::from_millis(125)).await; - assert_eq!( - drain_rate_limited_pending(&mut client, &mut state, "agent-pubkey", 1).await, - 1 + let valid_frame = Message::Text( + serde_json::to_string(&json!(["EVENT", MEMBERSHIP_NOTIF_SUB_ID, valid])) + .expect("serialize valid membership frame") + .into(), ); - let frame = next_test_frame(&mut server).await; - assert_eq!(frame[0], "REQ"); - assert_eq!(frame[1], channel_sub_id(channel_id)); + assert!( + handle_ws_message( + valid_frame, + &mut ws, + &event_tx, + &observer_control_tx, + &mut state, + &agent_keys, + "ws://localhost:3000", + &agent_pubkey_hex, + None, + ) + .await + ); + let forwarded = event_rx + .try_recv() + .expect("authentic same-ID membership event remains admissible") + .expect("membership event payload"); + assert_eq!(forwarded.channel_id, channel_id); + assert_eq!(forwarded.event.id.to_hex(), claimed_id); + assert!(state.seen_ids.contains(&claimed_id)); } #[tokio::test] - async fn subscribe_during_replay_pacing_is_sent_on_live_socket() { - let (client, mut server) = test_ws_pair().await; - let (cmd_tx, mut cmd_rx) = mpsc::channel(4); + async fn regular_event_verification_precedes_dedup_and_forwarding() { + let (mut ws, _server) = test_ws_pair().await; + let (event_tx, mut event_rx) = mpsc::channel(2); + let (observer_control_tx, _observer_control_rx) = mpsc::channel(1); let mut state = BgState::new(); - let replayed_channel = Uuid::new_v4(); - let deferred_channel = Uuid::new_v4(); - seed_test_subscription(&mut state, replayed_channel); - - let task = tokio::spawn(async move { - let mut client = client; - let result = resubscribe_after_reconnect( - &mut client, - &mut cmd_rx, + let channel_id = Uuid::new_v4(); + let keys = Keys::generate(); + seed_test_subscription(&mut state, channel_id); + let valid = make_test_channel_event(&keys, channel_id, 9, "test", None, 1_700_000); + let claimed_id = valid.id.to_hex(); + + // Retain the authentic event's claimed ID and signature while changing + // its content, making both the ID binding and signature invalid. + let mut tampered_json = serde_json::to_value(&valid).expect("serialize event"); + tampered_json["content"] = Value::String("tampered".to_string()); + let tampered: Event = serde_json::from_value(tampered_json).expect("parse tampered event"); + assert_eq!(tampered.id.to_hex(), claimed_id); + assert!(buzz_core::verify_event(&tampered).is_err()); + + let invalid_frame = Message::Text( + serde_json::to_string(&json!(["EVENT", channel_sub_id(channel_id), tampered])) + .expect("serialize invalid frame") + .into(), + ); + assert!( + handle_ws_message( + invalid_frame, + &mut ws, + &event_tx, + &observer_control_tx, &mut state, - "agent-pubkey", - true, + &keys, + "ws://localhost:3000", + &keys.public_key().to_hex(), + None, ) - .await; - (result, state) - }); + .await + ); + assert!( + !state.seen_ids.contains(&claimed_id), + "invalid event must not poison dedup" + ); + assert!( + event_rx.try_recv().is_err(), + "invalid event must not be forwarded" + ); - let replay = next_test_frame(&mut server).await; - assert_eq!(replay[1], channel_sub_id(replayed_channel)); - cmd_tx - .send(RelayCommand::Subscribe { - channel_id: deferred_channel, - filter: test_channel_filter(), - replay_since: Some(2_000), - }) + let valid_frame = Message::Text( + serde_json::to_string(&json!(["EVENT", channel_sub_id(channel_id), valid])) + .expect("serialize valid frame") + .into(), + ); + assert!( + handle_ws_message( + valid_frame, + &mut ws, + &event_tx, + &observer_control_tx, + &mut state, + &keys, + "ws://localhost:3000", + &keys.public_key().to_hex(), + None, + ) .await - .expect("queue subscribe during pacing"); - - let deferred = next_test_frame(&mut server).await; - assert_eq!(deferred[0], "REQ"); - assert_eq!(deferred[1], channel_sub_id(deferred_channel)); - let (result, state) = task.await.expect("join resubscribe task"); - assert!(matches!(result, ResubscribeResult::Ok)); - assert!(state.active_subscriptions.contains_key(&deferred_channel)); + ); + let forwarded = event_rx + .try_recv() + .expect("valid same-ID event remains admissible") + .expect("regular event payload"); + assert_eq!(forwarded.event.id.to_hex(), claimed_id); + assert!(state.seen_ids.contains(&claimed_id)); } #[tokio::test] - async fn unsubscribe_during_replay_pacing_sends_close_on_live_socket() { - let (client, mut server) = test_ws_pair().await; - let (cmd_tx, mut cmd_rx) = mpsc::channel(4); + async fn regular_events_must_match_active_subscription_and_local_filter() { + let (mut ws, _server) = test_ws_pair().await; + let (event_tx, mut event_rx) = mpsc::channel(8); + let (observer_control_tx, _observer_control_rx) = mpsc::channel(1); let mut state = BgState::new(); - let channel_id = Uuid::new_v4(); - seed_test_subscription(&mut state, channel_id); + let channel_a = Uuid::new_v4(); + let channel_b = Uuid::new_v4(); + let channel_inactive = Uuid::new_v4(); + let agent_keys = Keys::generate(); + let sender_keys = Keys::generate(); + let agent_pubkey_hex = agent_keys.public_key().to_hex(); + let mention_filter = ChannelFilter { + kinds: Some(vec![9]), + require_mention: true, + }; + seed_test_subscription_with_filter(&mut state, channel_a, mention_filter.clone()); + seed_test_subscription_with_filter(&mut state, channel_b, mention_filter); + + let routed = make_test_channel_event( + &sender_keys, + channel_a, + 9, + "routed", + Some(&agent_pubkey_hex), + 1_700_000, + ); + let routed_id = routed.id.to_hex(); + let misrouted_frame = Message::Text( + serde_json::to_string(&json!(["EVENT", channel_sub_id(channel_b), routed.clone()])) + .expect("serialize misrouted channel frame") + .into(), + ); + assert!( + handle_ws_message( + misrouted_frame, + &mut ws, + &event_tx, + &observer_control_tx, + &mut state, + &agent_keys, + "ws://localhost:3000", + &agent_pubkey_hex, + None, + ) + .await + ); + assert!( + !state.seen_ids.contains(&routed_id), + "an event delivered under the wrong active channel must not poison dedup" + ); + assert!(!state.last_seen.contains_key(&channel_b)); + assert!(event_rx.try_recv().is_err()); - let task = tokio::spawn(async move { - let mut client = client; - let result = resubscribe_after_reconnect( - &mut client, - &mut cmd_rx, + let correct_frame = Message::Text( + serde_json::to_string(&json!(["EVENT", channel_sub_id(channel_a), routed])) + .expect("serialize correctly routed channel frame") + .into(), + ); + assert!( + handle_ws_message( + correct_frame, + &mut ws, + &event_tx, + &observer_control_tx, &mut state, - "agent-pubkey", - true, + &agent_keys, + "ws://localhost:3000", + &agent_pubkey_hex, + None, ) - .await; - (result, state) - }); + .await + ); + let forwarded = event_rx + .try_recv() + .expect("correctly routed same-ID event remains admissible") + .expect("channel event payload"); + assert_eq!(forwarded.channel_id, channel_a); + assert_eq!(forwarded.event.id.to_hex(), routed_id); + + let inactive = make_test_channel_event( + &sender_keys, + channel_inactive, + 9, + "inactive", + Some(&agent_pubkey_hex), + 1_700_001, + ); + let inactive_id = inactive.id.to_hex(); + let inactive_frame = Message::Text( + serde_json::to_string(&json!([ + "EVENT", + channel_sub_id(channel_inactive), + inactive.clone() + ])) + .expect("serialize inactive channel frame") + .into(), + ); + assert!( + handle_ws_message( + inactive_frame, + &mut ws, + &event_tx, + &observer_control_tx, + &mut state, + &agent_keys, + "ws://localhost:3000", + &agent_pubkey_hex, + None, + ) + .await + ); + assert!( + !state.seen_ids.contains(&inactive_id), + "an inactive subscription must not poison dedup" + ); + assert!(event_rx.try_recv().is_err()); - let replay = next_test_frame(&mut server).await; - assert_eq!(replay[1], channel_sub_id(channel_id)); - cmd_tx - .send(RelayCommand::Unsubscribe { channel_id }) + seed_test_subscription_with_filter( + &mut state, + channel_inactive, + ChannelFilter { + kinds: Some(vec![9]), + require_mention: true, + }, + ); + let active_frame = Message::Text( + serde_json::to_string(&json!([ + "EVENT", + channel_sub_id(channel_inactive), + inactive + ])) + .expect("serialize newly active channel frame") + .into(), + ); + assert!( + handle_ws_message( + active_frame, + &mut ws, + &event_tx, + &observer_control_tx, + &mut state, + &agent_keys, + "ws://localhost:3000", + &agent_pubkey_hex, + None, + ) .await - .expect("queue unsubscribe during pacing"); + ); + assert_eq!( + event_rx + .try_recv() + .expect("event is admissible after local subscription activates") + .expect("channel event payload") + .event + .id + .to_hex(), + inactive_id + ); - let close = next_test_frame(&mut server).await; - assert_eq!(close, json!(["CLOSE", channel_sub_id(channel_id)])); - let (result, state) = task.await.expect("join resubscribe task"); - assert!(matches!(result, ResubscribeResult::Ok)); - assert!(!state.active_subscriptions.contains_key(&channel_id)); + let wrong_kind = make_test_channel_event( + &sender_keys, + channel_a, + 1, + "wrong kind", + Some(&agent_pubkey_hex), + 1_700_002, + ); + let wrong_kind_id = wrong_kind.id.to_hex(); + let wrong_kind_frame = Message::Text( + serde_json::to_string(&json!(["EVENT", channel_sub_id(channel_a), wrong_kind])) + .expect("serialize wrong-kind channel frame") + .into(), + ); + assert!( + handle_ws_message( + wrong_kind_frame, + &mut ws, + &event_tx, + &observer_control_tx, + &mut state, + &agent_keys, + "ws://localhost:3000", + &agent_pubkey_hex, + None, + ) + .await + ); + assert!(!state.seen_ids.contains(&wrong_kind_id)); + assert!(event_rx.try_recv().is_err()); + + let missing_mention = + make_test_channel_event(&sender_keys, channel_a, 9, "no mention", None, 1_700_003); + let missing_mention_id = missing_mention.id.to_hex(); + let missing_mention_frame = Message::Text( + serde_json::to_string(&json!([ + "EVENT", + channel_sub_id(channel_a), + missing_mention + ])) + .expect("serialize missing-mention channel frame") + .into(), + ); + assert!( + handle_ws_message( + missing_mention_frame, + &mut ws, + &event_tx, + &observer_control_tx, + &mut state, + &agent_keys, + "ws://localhost:3000", + &agent_pubkey_hex, + None, + ) + .await + ); + assert!(!state.seen_ids.contains(&missing_mention_id)); + assert!(event_rx.try_recv().is_err()); } #[tokio::test] - async fn publish_during_replay_pacing_is_sent_on_live_socket() { - let (client, mut server) = test_ws_pair().await; - let (cmd_tx, mut cmd_rx) = mpsc::channel(4); + async fn future_ordinary_event_cannot_poison_reconnect_watermark() { + let (mut ws, _server) = test_ws_pair().await; + let (event_tx, mut event_rx) = mpsc::channel(2); + let (observer_control_tx, _observer_control_rx) = mpsc::channel(1); let mut state = BgState::new(); let channel_id = Uuid::new_v4(); + let agent_keys = Keys::generate(); + let sender_keys = Keys::generate(); + let agent_pubkey_hex = agent_keys.public_key().to_hex(); seed_test_subscription(&mut state, channel_id); - let event = make_test_event(&nostr::Keys::generate(), 2_000); - let event_id = event.id.to_hex(); + let now = unix_now_secs(); + let future = + make_test_channel_event(&sender_keys, channel_id, 9, "future", None, now + 3_600); + let future_id = future.id.to_hex(); + let future_frame = Message::Text( + serde_json::to_string(&json!(["EVENT", channel_sub_id(channel_id), future])) + .expect("serialize future ordinary event") + .into(), + ); - let task = tokio::spawn(async move { - let mut client = client; - let result = resubscribe_after_reconnect( - &mut client, - &mut cmd_rx, + assert!( + handle_ws_message( + future_frame, + &mut ws, + &event_tx, + &observer_control_tx, &mut state, - "agent-pubkey", - true, + &agent_keys, + "ws://localhost:3000", + &agent_pubkey_hex, + None, ) - .await; - result - }); - - let replay = next_test_frame(&mut server).await; - assert_eq!(replay[1], channel_sub_id(channel_id)); - cmd_tx - .send(RelayCommand::PublishEvent { - event: Box::new(event), - }) .await - .expect("queue publish during pacing"); - - let publish = next_test_frame(&mut server).await; - assert_eq!(publish[0], "EVENT"); - assert_eq!(publish[1]["id"], event_id); - assert!(matches!( - task.await.expect("join resubscribe task"), - ResubscribeResult::Ok - )); + ); + assert!( + !state.seen_ids.contains(&future_id), + "future event must not reserve its ID" + ); + assert!( + !state.last_seen.contains_key(&channel_id), + "future event must not advance reconnect watermark" + ); + assert!(event_rx.try_recv().is_err()); + + let current = make_test_channel_event(&sender_keys, channel_id, 9, "current", None, now); + let current_id = current.id.to_hex(); + let current_frame = Message::Text( + serde_json::to_string(&json!(["EVENT", channel_sub_id(channel_id), current])) + .expect("serialize current ordinary event") + .into(), + ); + assert!( + handle_ws_message( + current_frame, + &mut ws, + &event_tx, + &observer_control_tx, + &mut state, + &agent_keys, + "ws://localhost:3000", + &agent_pubkey_hex, + None, + ) + .await + ); + let forwarded = event_rx + .try_recv() + .expect("current event should remain deliverable") + .expect("ordinary event payload"); + assert_eq!(forwarded.event.id.to_hex(), current_id); + assert_eq!(state.last_seen.get(&channel_id), Some(&now)); } - #[test] - fn failed_replay_retains_deferred_subscription_intent_in_fifo_order() { + #[tokio::test] + async fn bounded_future_skew_forwards_ordinary_event_without_future_watermark() { + let (mut ws, _server) = test_ws_pair().await; + let (event_tx, mut event_rx) = mpsc::channel(1); + let (observer_control_tx, _observer_control_rx) = mpsc::channel(1); let mut state = BgState::new(); - let kept_channel = Uuid::new_v4(); - let removed_channel = Uuid::new_v4(); - seed_test_subscription(&mut state, removed_channel); - let event = make_test_event(&nostr::Keys::generate(), 2_000); - let mut deferred = VecDeque::from([ - RelayCommand::Subscribe { - channel_id: kept_channel, - filter: test_channel_filter(), - replay_since: Some(2_000), - }, - RelayCommand::Unsubscribe { - channel_id: removed_channel, - }, - RelayCommand::PublishEvent { - event: Box::new(event), - }, - ]); - - retain_deferred_command_intent(&mut state, &mut deferred); + let channel_id = Uuid::new_v4(); + let agent_keys = Keys::generate(); + let sender_keys = Keys::generate(); + let agent_pubkey_hex = agent_keys.public_key().to_hex(); + seed_test_subscription(&mut state, channel_id); + let now = unix_now_secs(); + let event = make_test_channel_event( + &sender_keys, + channel_id, + 9, + "slightly ahead", + None, + now + CHANNEL_EVENT_FUTURE_SKEW_SECS, + ); + let event_id = event.id.to_hex(); + let frame = Message::Text( + serde_json::to_string(&json!(["EVENT", channel_sub_id(channel_id), event])) + .expect("serialize slightly future ordinary event") + .into(), + ); - assert!(deferred.is_empty()); - assert!(state.active_subscriptions.contains_key(&kept_channel)); - assert!(!state.active_subscriptions.contains_key(&removed_channel)); + assert!( + handle_ws_message( + frame, + &mut ws, + &event_tx, + &observer_control_tx, + &mut state, + &agent_keys, + "ws://localhost:3000", + &agent_pubkey_hex, + None, + ) + .await + ); + let forwarded = event_rx + .try_recv() + .expect("bounded client clock skew should remain deliverable") + .expect("ordinary event payload"); + assert_eq!(forwarded.event.id.to_hex(), event_id); + assert!(!forwarded.privileged_control_admitted); + assert!( + state + .last_seen + .get(&channel_id) + .is_some_and(|watermark| *watermark <= unix_now_secs()), + "accepted future skew must not move replay watermarks into the future" + ); } - #[test] - fn bg_state_dedup_first_event_accepted() { + #[tokio::test] + async fn future_membership_event_cannot_poison_replay_state() { + let (mut ws, _server) = test_ws_pair().await; + let (event_tx, mut event_rx) = mpsc::channel(2); + let (observer_control_tx, _observer_control_rx) = mpsc::channel(1); let mut state = BgState::new(); + state.membership_sub_active = true; let channel_id = Uuid::new_v4(); - let keys = nostr::Keys::generate(); - let event = make_test_event(&keys, 1_000_000); + let agent_keys = Keys::generate(); + let relay_keys = Keys::generate(); + let agent_pubkey_hex = agent_keys.public_key().to_hex(); + let now = unix_now_secs(); + let future = + make_test_membership_event(&relay_keys, channel_id, &agent_pubkey_hex, now + 3_600); + let future_id = future.id.to_hex(); + let future_frame = Message::Text( + serde_json::to_string(&json!(["EVENT", MEMBERSHIP_NOTIF_SUB_ID, future])) + .expect("serialize future membership event") + .into(), + ); + assert!( - state.record_event(channel_id, &event), - "first event should be accepted" + handle_ws_message( + future_frame, + &mut ws, + &event_tx, + &observer_control_tx, + &mut state, + &agent_keys, + "ws://localhost:3000", + &agent_pubkey_hex, + None, + ) + .await + ); + assert!( + !state.seen_ids.contains(&future_id), + "future membership event must not reserve its ID" + ); + assert_eq!( + state.membership_last_seen, None, + "future membership event must not advance replay state" + ); + assert!(event_rx.try_recv().is_err()); + + let current = make_test_membership_event(&relay_keys, channel_id, &agent_pubkey_hex, now); + let current_id = current.id.to_hex(); + let current_frame = Message::Text( + serde_json::to_string(&json!(["EVENT", MEMBERSHIP_NOTIF_SUB_ID, current])) + .expect("serialize current membership event") + .into(), + ); + assert!( + handle_ws_message( + current_frame, + &mut ws, + &event_tx, + &observer_control_tx, + &mut state, + &agent_keys, + "ws://localhost:3000", + &agent_pubkey_hex, + None, + ) + .await ); + let forwarded = event_rx + .try_recv() + .expect("current membership event should remain deliverable") + .expect("membership event payload"); + assert_eq!(forwarded.event.id.to_hex(), current_id); + assert_eq!(state.membership_last_seen, Some(now)); } - #[test] - fn bg_state_dedup_duplicate_rejected() { + #[tokio::test] + async fn stale_privileged_channel_control_is_rejected_before_dedup() { + let (mut ws, _server) = test_ws_pair().await; + let (event_tx, mut event_rx) = mpsc::channel(2); + let (observer_control_tx, _observer_control_rx) = mpsc::channel(1); let mut state = BgState::new(); let channel_id = Uuid::new_v4(); - let keys = nostr::Keys::generate(); - let event = make_test_event(&keys, 1_000_000); + let agent_keys = Keys::generate(); + let owner_keys = Keys::generate(); + let agent_pubkey_hex = agent_keys.public_key().to_hex(); + state.authorized_owner_pubkey = Some(owner_keys.public_key().to_hex()); + seed_test_subscription_with_filter( + &mut state, + channel_id, + ChannelFilter { + kinds: Some(vec![9]), + require_mention: true, + }, + ); + let stale = make_test_channel_event( + &owner_keys, + channel_id, + 9, + "!rotate", + Some(&agent_pubkey_hex), + unix_now_secs().saturating_sub(301), + ); + let stale_id = stale.id.to_hex(); + let stale_frame = Message::Text( + serde_json::to_string(&json!(["EVENT", channel_sub_id(channel_id), stale])) + .expect("serialize stale privileged control") + .into(), + ); + assert!( - state.record_event(channel_id, &event), - "first should be accepted" + handle_ws_message( + stale_frame, + &mut ws, + &event_tx, + &observer_control_tx, + &mut state, + &agent_keys, + "ws://localhost:3000", + &agent_pubkey_hex, + None, + ) + .await ); assert!( - !state.record_event(channel_id, &event), - "duplicate should be rejected" + state.seen_ids.contains(&stale_id), + "stale owner control should follow ordinary-message dedup" + ); + let forwarded = event_rx + .try_recv() + .expect("stale owner control should remain an ordinary message") + .expect("ordinary event payload"); + assert_eq!(forwarded.event.id.to_hex(), stale_id); + assert!( + !forwarded.privileged_control_admitted, + "stale owner control must never carry relay execution authority" ); } - #[test] - fn bg_state_dedup_different_ids_both_accepted() { + #[tokio::test] + async fn non_owner_command_shape_does_not_consume_privileged_replay_capacity() { + let (mut ws, _server) = test_ws_pair().await; + let (event_tx, mut event_rx) = mpsc::channel(2); + let (observer_control_tx, _observer_control_rx) = mpsc::channel(1); let mut state = BgState::new(); let channel_id = Uuid::new_v4(); - // Two different keys → two different event IDs. - let keys1 = nostr::Keys::generate(); - let keys2 = nostr::Keys::generate(); - let event1 = make_test_event(&keys1, 1_000_000); - let event2 = make_test_event(&keys2, 1_000_001); - assert!(state.record_event(channel_id, &event1)); - assert!(state.record_event(channel_id, &event2)); - } + let agent_keys = Keys::generate(); + let owner_keys = Keys::generate(); + let non_owner_keys = Keys::generate(); + let agent_pubkey_hex = agent_keys.public_key().to_hex(); + let now = unix_now_secs(); + state.authorized_owner_pubkey = Some(owner_keys.public_key().to_hex()); + state.privileged_control_start_boundary = now.saturating_sub(1); + state.privileged_control_replay = PrivilegedControlReplay::new(1); + assert_eq!( + state.privileged_control_replay.admit( + "owner-capacity".into(), + now.saturating_add(300), + now + ), + PrivilegedControlAdmission::Admitted + ); + seed_test_subscription_with_filter( + &mut state, + channel_id, + ChannelFilter { + kinds: Some(vec![9]), + require_mention: true, + }, + ); - #[test] - fn bg_state_last_seen_set_on_first_event() { - let mut state = BgState::new(); - let channel_id = Uuid::new_v4(); - let keys = nostr::Keys::generate(); - let event = make_test_event(&keys, 1_700_000); - state.record_event(channel_id, &event); - assert_eq!(state.last_seen.get(&channel_id).copied(), Some(1_700_000)); - } + let ordinary = make_test_channel_event( + &non_owner_keys, + channel_id, + 9, + "!shutdown", + Some(&agent_pubkey_hex), + now, + ); + let ordinary_id = ordinary.id.to_hex(); + let frame = Message::Text( + serde_json::to_string(&json!(["EVENT", channel_sub_id(channel_id), ordinary])) + .expect("serialize non-owner command-shaped event") + .into(), + ); - #[test] - fn bg_state_last_seen_advances_on_newer_event() { - let mut state = BgState::new(); - let channel_id = Uuid::new_v4(); - let keys1 = nostr::Keys::generate(); - let keys2 = nostr::Keys::generate(); - let event1 = make_test_event(&keys1, 1_700_000); - let event2 = make_test_event(&keys2, 1_800_000); - state.record_event(channel_id, &event1); - state.record_event(channel_id, &event2); - assert_eq!(state.last_seen.get(&channel_id).copied(), Some(1_800_000)); + assert!( + handle_ws_message( + frame, + &mut ws, + &event_tx, + &observer_control_tx, + &mut state, + &agent_keys, + "ws://localhost:3000", + &agent_pubkey_hex, + None, + ) + .await + ); + let forwarded = event_rx + .try_recv() + .expect("non-owner command shape follows ordinary-message admission") + .expect("ordinary event payload"); + assert_eq!(forwarded.event.id.to_hex(), ordinary_id); + assert!( + !forwarded.privileged_control_admitted, + "non-owner command shape must not carry relay execution authority" + ); + assert_eq!( + state.privileged_control_replay.expires_at.len(), + 1, + "non-owner traffic must not reserve privileged replay capacity" + ); } - #[test] - fn bg_state_last_seen_does_not_regress_on_older_event() { + #[tokio::test] + async fn owner_controls_require_live_channel_req_and_post_boundary_retry() { + let (mut ws, _server) = test_ws_pair().await; + let (event_tx, mut event_rx) = mpsc::channel(4); + let (observer_control_tx, _observer_control_rx) = mpsc::channel(1); let mut state = BgState::new(); let channel_id = Uuid::new_v4(); - let keys1 = nostr::Keys::generate(); - let keys2 = nostr::Keys::generate(); - let event_new = make_test_event(&keys1, 1_800_000); - let event_old = make_test_event(&keys2, 1_700_000); - state.record_event(channel_id, &event_new); - state.record_event(channel_id, &event_old); - // last_seen should remain at the higher timestamp - assert_eq!(state.last_seen.get(&channel_id).copied(), Some(1_800_000)); - } - - #[test] - fn bg_state_last_seen_independent_per_channel() { - let mut state = BgState::new(); - let ch1 = Uuid::new_v4(); - let ch2 = Uuid::new_v4(); - let keys1 = nostr::Keys::generate(); - let keys2 = nostr::Keys::generate(); - let event1 = make_test_event(&keys1, 1_000_000); - let event2 = make_test_event(&keys2, 2_000_000); - state.record_event(ch1, &event1); - state.record_event(ch2, &event2); - assert_eq!(state.last_seen.get(&ch1).copied(), Some(1_000_000)); - 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); + let agent_keys = Keys::generate(); + let owner_keys = Keys::generate(); + let agent_pubkey_hex = agent_keys.public_key().to_hex(); + let now = unix_now_secs(); + state.authorized_owner_pubkey = Some(owner_keys.public_key().to_hex()); + state.privileged_control_start_boundary = now.saturating_sub(10); + seed_test_subscription_with_filter( + &mut state, + channel_id, + ChannelFilter { + kinds: Some(vec![9]), + require_mention: true, + }, + ); - // 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); - } + for (command, boundary, expected_admitted) in [ + ("!shutdown", None, false), + ("!cancel", Some(now), false), + ("!rotate", Some(now.saturating_sub(1)), true), + ] { + match boundary { + Some(boundary) => { + state.channel_subscribed_at.insert(channel_id, boundary); + } + None => { + state.channel_subscribed_at.remove(&channel_id); + } + } + let event = make_test_channel_event( + &owner_keys, + channel_id, + KIND_STREAM_MESSAGE, + command, + Some(&agent_pubkey_hex), + now, + ); + let frame = Message::Text( + serde_json::to_string(&json!(["EVENT", channel_sub_id(channel_id), event])) + .expect("serialize owner control") + .into(), + ); - // 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" + handle_ws_message( + frame, + &mut ws, + &event_tx, + &observer_control_tx, + &mut state, + &agent_keys, + "ws://localhost:3000", + &agent_pubkey_hex, + None, + ) + .await + ); + let forwarded = event_rx + .try_recv() + .expect("owner command-shaped text remains deliverable") + .expect("ordinary event payload"); + assert_eq!( + forwarded.privileged_control_admitted, expected_admitted, + "{command} admission mismatch for boundary {boundary:?}" ); } + } - // 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" + #[tokio::test] + async fn channel_req_boundary_tracks_only_a_successful_live_send() { + let (mut client, mut server) = test_ws_pair().await; + let mut state = BgState::new(); + let channel_id = Uuid::new_v4(); + let filter = ChannelFilter { + kinds: Some(vec![KIND_STREAM_MESSAGE]), + require_mention: true, + }; + state.channel_subscribed_at.insert(channel_id, 1); + + apply_command_to_state( + &mut state, + RelayCommand::Subscribe { + channel_id, + filter: filter.clone(), + replay_since: Some(1_000), + }, ); assert!( - dedup.contains(&new_id), - "new ID should be found after insert" + !state.channel_subscribed_at.contains_key(&channel_id), + "disconnected subscription intent must not retain a prior live REQ boundary" ); - } - #[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. + let before_send = unix_now_secs(); assert!( - !dedup.insert("id-0".to_string()), - "cross-generation duplicate must be rejected" + send_subscribe( + &mut client, + &mut state, + channel_id, + "agent-pubkey", + Some(1_000), + &filter, + ) + .await ); + let after_send = unix_now_secs(); + let boundary = state + .channel_subscribed_at + .get(&channel_id) + .copied() + .expect("successful live REQ records its completion boundary"); + assert!((before_send..=after_send).contains(&boundary)); + let frame = next_test_frame(&mut server).await; + assert_eq!(frame[0], "REQ"); + assert_eq!(frame[1], channel_sub_id(channel_id)); } #[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(); - - // 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)); - } + fn privileged_control_owner_requires_verified_nip_oa_attestation() { + let owner_keys = Keys::generate(); + let agent_keys = Keys::generate(); + let other_agent_keys = Keys::generate(); + let configured_owner_keys = Keys::generate(); + let configured_owner = configured_owner_keys.public_key().to_hex(); + let auth_json = + buzz_sdk::nip_oa::compute_auth_tag(&owner_keys, &agent_keys.public_key(), "") + .expect("compute owner attestation"); + let auth_tag = + buzz_sdk::nip_oa::parse_auth_tag(&auth_json).expect("parse owner attestation"); + let wrong_agent_auth_json = + buzz_sdk::nip_oa::compute_auth_tag(&owner_keys, &other_agent_keys.public_key(), "") + .expect("compute wrong-agent owner attestation"); + let wrong_agent_auth_tag = buzz_sdk::nip_oa::parse_auth_tag(&wrong_agent_auth_json) + .expect("parse wrong-agent owner attestation"); - // 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" + assert_eq!( + resolve_authorized_owner_pubkey( + Some(&auth_tag), + &agent_keys, + Some(configured_owner.as_str()), + ), + Some(owner_keys.public_key().to_hex()), + "a valid NIP-OA owner remains canonical over configured fallback" ); - - // 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.record_event(channel_id, &event), - "new event after rotation should be accepted" + assert_eq!( + resolve_authorized_owner_pubkey( + Some(&wrong_agent_auth_tag), + &agent_keys, + Some(configured_owner.as_str()), + ), + Some(configured_owner.clone()), + "an invalid attestation for this agent must fall back to trusted configuration" ); - - // The same event must be rejected as a duplicate. - assert!( - !state.record_event(channel_id, &event), - "duplicate event after rotation should be rejected" + assert_eq!( + resolve_authorized_owner_pubkey(None, &agent_keys, Some(configured_owner.as_str())), + Some(configured_owner), + "missing attestation must retain the trusted configured owner" + ); + assert_eq!( + resolve_authorized_owner_pubkey(None, &agent_keys, Some("not-a-pubkey")), + None, + "malformed configured owner must fail closed" ); } - /// Test 8: channel_dropped_since records the OLDEST dropped timestamp. - /// - /// Simulates the backpressure path directly on BgState: - /// - First drop at ts=1000 → entry is 1000 - /// - Second drop at ts=2000 (later) → entry stays 1000 (min) - /// - Third drop at ts=500 (earlier) → entry updates to 500 (min) #[test] - fn acp_records_channel_dropped_since_on_backpressure() { - let mut state = BgState::new(); - let channel_id = Uuid::new_v4(); + fn privileged_control_replay_guard_never_evicts_live_ids() { + let mut replay = PrivilegedControlReplay::new(2); - // Simulate the backpressure path: record ts=1000. - let ts1: u64 = 1_000; - state - .channel_dropped_since - .entry(channel_id) - .and_modify(|d| *d = (*d).min(ts1)) - .or_insert(ts1); assert_eq!( - state.channel_dropped_since.get(&channel_id).copied(), - Some(1_000), - "first drop should record ts=1000" + replay.admit("first".into(), 200, 100), + PrivilegedControlAdmission::Admitted ); - - // Later timestamp (2000) — entry should stay at 1000. - let ts2: u64 = 2_000; - state - .channel_dropped_since - .entry(channel_id) - .and_modify(|d| *d = (*d).min(ts2)) - .or_insert(ts2); assert_eq!( - state.channel_dropped_since.get(&channel_id).copied(), - Some(1_000), - "later drop should not overwrite earlier timestamp" + replay.admit("second".into(), 200, 100), + PrivilegedControlAdmission::Admitted ); - - // Earlier timestamp (500) — entry should update to 500. - let ts3: u64 = 500; - state - .channel_dropped_since - .entry(channel_id) - .and_modify(|d| *d = (*d).min(ts3)) - .or_insert(ts3); assert_eq!( - state.channel_dropped_since.get(&channel_id).copied(), - Some(500), - "earlier drop should update entry to 500" + replay.admit("third".into(), 200, 100), + PrivilegedControlAdmission::Full, + "capacity pressure must fail closed instead of evicting a live ID" + ); + assert_eq!( + replay.admit("first".into(), 200, 200), + PrivilegedControlAdmission::Replay, + "expiry is inclusive because the event remains freshness-admissible at that second" + ); + assert_eq!( + replay.admit("third".into(), 300, 201), + PrivilegedControlAdmission::Admitted, + "expired IDs release bounded capacity" ); + assert_eq!(replay.expires_at.len(), 1); } - /// Test 9: reconnect since filter = min(last_seen, channel_dropped_since) - SINCE_SKEW_SECS. - /// - /// With last_seen=1000 and channel_dropped_since=900, the effective since - /// passed to send_subscribe should be min(1000, 900) - SINCE_SKEW_SECS = 895. - #[test] - fn acp_reconnect_uses_dropped_since_for_replay() { + #[tokio::test] + async fn privileged_controls_require_strict_post_start_nonfuture_timestamps() { + let (mut ws, _server) = test_ws_pair().await; + let (event_tx, mut event_rx) = mpsc::channel(4); + let (observer_control_tx, _observer_control_rx) = mpsc::channel(1); let mut state = BgState::new(); let channel_id = Uuid::new_v4(); + let agent_keys = Keys::generate(); + let owner_keys = Keys::generate(); + let agent_pubkey_hex = agent_keys.public_key().to_hex(); + let now = unix_now_secs(); + state.authorized_owner_pubkey = Some(owner_keys.public_key().to_hex()); + state.privileged_control_start_boundary = now.saturating_sub(1); + seed_test_subscription_with_filter( + &mut state, + channel_id, + ChannelFilter { + kinds: Some(vec![9]), + require_mention: true, + }, + ); - // Set up state: last_seen=1000, channel_dropped_since=900. - state.last_seen.insert(channel_id, 1_000); - state.channel_dropped_since.insert(channel_id, 900); - - // Compute the since value the reconnect path would use. - let since = state.channel_since(&channel_id); - - // The since passed to send_subscribe (which subtracts SINCE_SKEW_SECS internally). - assert_eq!(since, Some(900), "since should be min(1000, 900) = 900"); - - // After subtracting skew (as send_subscribe does), the REQ filter value is: - let req_since = since.unwrap().saturating_sub(SINCE_SKEW_SECS); - assert_eq!( - req_since, 895, - "REQ since filter should be 900 - {} = 895", - SINCE_SKEW_SECS + let boundary = make_test_channel_event( + &owner_keys, + channel_id, + 9, + "!shutdown", + Some(&agent_pubkey_hex), + state.privileged_control_start_boundary, + ); + let boundary_id = boundary.id.to_hex(); + let boundary_frame = Message::Text( + serde_json::to_string(&json!(["EVENT", channel_sub_id(channel_id), boundary])) + .expect("serialize boundary control") + .into(), ); - - // Simulate clearing after resubscribe. - state.channel_dropped_since.remove(&channel_id); assert!( - !state.channel_dropped_since.contains_key(&channel_id), - "channel_dropped_since should be cleared after resubscribe" + handle_ws_message( + boundary_frame, + &mut ws, + &event_tx, + &observer_control_tx, + &mut state, + &agent_keys, + "ws://localhost:3000", + &agent_pubkey_hex, + None, + ) + .await ); - } - - #[test] - fn dynamic_subscribe_records_membership_replay_floor() { - let mut state = BgState::new(); - state.startup_watermark = Some(2_000); - let channel_id = Uuid::new_v4(); - let membership_ts = 10_000; - let filter = ChannelFilter { - kinds: Some(vec![9]), - require_mention: true, - }; - - apply_command_to_state( - &mut state, - RelayCommand::Subscribe { - channel_id, - filter, - replay_since: Some(membership_ts), - }, + assert!( + state.seen_ids.contains(&boundary_id), + "same-second owner text should remain an ordinary message" + ); + let boundary_forwarded = event_rx + .try_recv() + .expect("same-second owner text should be forwarded without authority") + .expect("ordinary event payload"); + assert!( + !boundary_forwarded.privileged_control_admitted, + "second-granularity equality with startup must fail closed for execution" ); - assert_eq!( - state.subscribe_since.get(&channel_id).copied(), - Some(membership_ts), - "dynamic channel subscriptions should replay from the membership notification, not startup" + let skewed_future = make_test_channel_event( + &owner_keys, + channel_id, + 9, + "!cancel", + Some(&agent_pubkey_hex), + now.saturating_add(CHANNEL_EVENT_FUTURE_SKEW_SECS), ); - assert_eq!( - state.channel_since(&channel_id), - Some(membership_ts), - "channel_since should use the dynamic replay floor until an event is seen" + let skewed_future_id = skewed_future.id.to_hex(); + let skewed_future_frame = Message::Text( + serde_json::to_string(&json!(["EVENT", channel_sub_id(channel_id), skewed_future])) + .expect("serialize skewed future control") + .into(), + ); + assert!( + handle_ws_message( + skewed_future_frame, + &mut ws, + &event_tx, + &observer_control_tx, + &mut state, + &agent_keys, + "ws://localhost:3000", + &agent_pubkey_hex, + None, + ) + .await + ); + assert!(state.seen_ids.contains(&skewed_future_id)); + let skewed_forwarded = event_rx + .try_recv() + .expect("bounded future owner text remains an ordinary message") + .expect("ordinary event payload"); + assert!( + !skewed_forwarded.privileged_control_admitted, + "future-skewed owner text must never gain execution authority" + ); + assert!( + state + .last_seen + .get(&channel_id) + .is_some_and(|watermark| *watermark <= unix_now_secs()), + "future-skewed owner text must clamp its replay watermark" ); - } - /// 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(); + let future = make_test_channel_event( + &owner_keys, + channel_id, + 9, + "!cancel", + Some(&agent_pubkey_hex), + now.saturating_add(30), + ); + let future_id = future.id.to_hex(); + let future_frame = Message::Text( + serde_json::to_string(&json!(["EVENT", channel_sub_id(channel_id), future])) + .expect("serialize future control") + .into(), + ); + assert!( + handle_ws_message( + future_frame, + &mut ws, + &event_tx, + &observer_control_tx, + &mut state, + &agent_keys, + "ws://localhost:3000", + &agent_pubkey_hex, + None, + ) + .await + ); + assert!(!state.seen_ids.contains(&future_id)); + assert!(event_rx.try_recv().is_err()); - // 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)); + let fresh = make_test_channel_event( + &owner_keys, + channel_id, + 9, + "!rotate", + Some(&agent_pubkey_hex), + now, + ); + let fresh_id = fresh.id.to_hex(); + let fresh_json = + serde_json::to_string(&json!(["EVENT", channel_sub_id(channel_id), fresh])) + .expect("serialize fresh control"); + assert!( + handle_ws_message( + Message::Text(fresh_json.clone().into()), + &mut ws, + &event_tx, + &observer_control_tx, + &mut state, + &agent_keys, + "ws://localhost:3000", + &agent_pubkey_hex, + None, + ) + .await + ); + let forwarded = event_rx + .try_recv() + .expect("fresh post-start control should be forwarded") + .expect("control payload"); + assert_eq!(forwarded.event.id.to_hex(), fresh_id); + assert!( + forwarded.privileged_control_admitted, + "fresh exact-owner control must carry relay execution authority" + ); + assert!( + handle_ws_message( + Message::Text(fresh_json.clone().into()), + &mut ws, + &event_tx, + &observer_control_tx, + &mut state, + &agent_keys, + "ws://localhost:3000", + &agent_pubkey_hex, + None, + ) + .await + ); + assert!( + event_rx.try_recv().is_err(), + "same-process replay must not be forwarded" + ); - // 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(); + let mut restarted_state = BgState::new(); + restarted_state.authorized_owner_pubkey = Some(owner_keys.public_key().to_hex()); + restarted_state.privileged_control_start_boundary = now; + seed_test_subscription_with_filter( + &mut restarted_state, + channel_id, + ChannelFilter { + kinds: Some(vec![9]), + require_mention: true, + }, + ); assert!( - state.seen_ids.insert(membership_id), - "membership event should be accepted by dedup" + handle_ws_message( + Message::Text(fresh_json.into()), + &mut ws, + &event_tx, + &observer_control_tx, + &mut restarted_state, + &agent_keys, + "ws://localhost:3000", + &agent_pubkey_hex, + None, + ) + .await ); - // 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" + assert!( + restarted_state.seen_ids.contains(&fresh_id), + "pre-start owner text should remain an ordinary message after restart" + ); + let restarted_forwarded = event_rx + .try_recv() + .expect("pre-start owner text should be forwarded without authority") + .expect("ordinary event payload"); + assert_eq!(restarted_forwarded.event.id.to_hex(), fresh_id); + assert!( + !restarted_forwarded.privileged_control_admitted, + "restart replay must never regain execution authority" ); } - /// 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(); + async fn test_ws_pair() -> (WsStream, WebSocketStream) { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("bind test websocket"); + let address = listener.local_addr().expect("read test address"); + let server = tokio::spawn(async move { + let (stream, _) = listener.accept().await.expect("accept test websocket"); + tokio_tungstenite::accept_async(stream) + .await + .expect("complete server websocket handshake") + }); + let (client, _) = connect_async(format!("ws://{address}")) + .await + .expect("connect test websocket"); + (client, server.await.expect("join test websocket server")) + } - // 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)); + async fn next_test_frame( + server: &mut WebSocketStream, + ) -> serde_json::Value { + let message = timeout(Duration::from_secs(1), server.next()) + .await + .expect("timed out waiting for websocket frame") + .expect("test websocket closed") + .expect("read test websocket frame"); + serde_json::from_str(message.to_text().expect("expected text frame")) + .expect("parse test websocket frame") + } - // Simulate backpressure: remove the ID (matching the production code). - state.seen_ids.remove(&event_id_hex); + fn test_channel_filter() -> ChannelFilter { + ChannelFilter { + kinds: Some(vec![9]), + require_mention: false, + } + } - // The ID should now be accepted again on replay. - assert!( - state.seen_ids.insert(event_id_hex), - "after backpressure removal, replay must be accepted" - ); + fn seed_test_subscription(state: &mut BgState, channel_id: Uuid) { + seed_test_subscription_with_filter(state, channel_id, test_channel_filter()); } - /// Subscribe a channel via the production command path so the test exercises - /// real subscription state (active_subscriptions + active_filters + since). - fn subscribe_channel(state: &mut BgState, channel_id: Uuid) { + fn seed_test_subscription_with_filter( + state: &mut BgState, + channel_id: Uuid, + filter: ChannelFilter, + ) { apply_command_to_state( state, RelayCommand::Subscribe { channel_id, - filter: ChannelFilter { - kinds: Some(vec![9]), - require_mention: false, - }, + filter, replay_since: Some(1_000), }, ); + state + .channel_subscribed_at + .insert(channel_id, unix_now_secs().saturating_sub(1)); } - #[test] - fn not_a_channel_member_drops_channel_without_reconnect() { + #[tokio::test] + async fn fresh_reconnect_preserves_gate_until_pending_replay_resumes() { + let (mut client, mut server) = test_ws_pair().await; + let (_cmd_tx, mut cmd_rx) = mpsc::channel(1); let mut state = BgState::new(); let channel_id = Uuid::new_v4(); - subscribe_channel(&mut state, channel_id); + seed_test_subscription(&mut state, channel_id); + state.rate_limit_gate = Some(tokio::time::Instant::now() + Duration::from_millis(150)); - let handled = drop_channel_on_access_denied( - &mut state, - &channel_sub_id(channel_id), - "restricted: not a channel member", - ); + let result = + resubscribe_after_reconnect(&mut client, &mut cmd_rx, &mut state, "agent-pubkey", true) + .await; - assert!(handled, "per-channel denial must be handled (no reconnect)"); - assert!( - !state.active_subscriptions.contains_key(&channel_id), - "the forbidden channel's subscription must be dropped" - ); + assert!(matches!(result, ResubscribeResult::Ok)); + assert!(state.rate_limit_gate.is_some()); + assert!(state.rate_limited_pending.contains_key(&channel_id)); assert!( - !state.active_filters.contains_key(&channel_id), - "channel state must be cleared (Unsubscribe cleanup)" + timeout(Duration::from_millis(50), server.next()) + .await + .is_err(), + "fresh reconnect must not send REQ while the shared quota gate is active" ); - } - - #[test] - fn channel_access_revoked_drops_channel_without_reconnect() { - let mut state = BgState::new(); - let channel_id = Uuid::new_v4(); - subscribe_channel(&mut state, channel_id); - let handled = drop_channel_on_access_denied( - &mut state, - &channel_sub_id(channel_id), - "restricted: channel access revoked", + tokio::time::sleep(Duration::from_millis(125)).await; + assert_eq!( + drain_rate_limited_pending(&mut client, &mut state, "agent-pubkey", 1).await, + 1 ); - - assert!(handled, "per-channel denial must be handled (no reconnect)"); - assert!(!state.active_subscriptions.contains_key(&channel_id)); - assert!(!state.active_filters.contains_key(&channel_id)); + let frame = next_test_frame(&mut server).await; + assert_eq!(frame[0], "REQ"); + assert_eq!(frame[1], channel_sub_id(channel_id)); } - #[test] - fn insufficient_scope_is_not_dropped_and_reconnects() { + #[tokio::test] + async fn subscribe_during_replay_pacing_is_sent_on_live_socket() { + let (client, mut server) = test_ws_pair().await; + let (cmd_tx, mut cmd_rx) = mpsc::channel(4); let mut state = BgState::new(); - let channel_id = Uuid::new_v4(); - subscribe_channel(&mut state, channel_id); - - let handled = drop_channel_on_access_denied( - &mut state, - &channel_sub_id(channel_id), - "restricted: insufficient scope", - ); + let replayed_channel = Uuid::new_v4(); + let deferred_channel = Uuid::new_v4(); + seed_test_subscription(&mut state, replayed_channel); - assert!( - !handled, - "connection-level insufficient-scope must fall through to reconnect, not drop the channel" - ); - assert!( - state.active_subscriptions.contains_key(&channel_id), - "the channel must survive so reconnect can restore it" - ); + let task = tokio::spawn(async move { + let mut client = client; + let result = resubscribe_after_reconnect( + &mut client, + &mut cmd_rx, + &mut state, + "agent-pubkey", + true, + ) + .await; + (result, state) + }); + + let replay = next_test_frame(&mut server).await; + assert_eq!(replay[1], channel_sub_id(replayed_channel)); + cmd_tx + .send(RelayCommand::Subscribe { + channel_id: deferred_channel, + filter: test_channel_filter(), + replay_since: Some(2_000), + }) + .await + .expect("queue subscribe during pacing"); + + let deferred = next_test_frame(&mut server).await; + assert_eq!(deferred[0], "REQ"); + assert_eq!(deferred[1], channel_sub_id(deferred_channel)); + let (result, state) = task.await.expect("join resubscribe task"); + assert!(matches!(result, ResubscribeResult::Ok)); + assert!(state.active_subscriptions.contains_key(&deferred_channel)); } - #[test] - fn auth_required_is_not_dropped_and_reconnects() { + #[tokio::test] + async fn unsubscribe_during_replay_pacing_sends_close_on_live_socket() { + let (client, mut server) = test_ws_pair().await; + let (cmd_tx, mut cmd_rx) = mpsc::channel(4); let mut state = BgState::new(); let channel_id = Uuid::new_v4(); - subscribe_channel(&mut state, channel_id); + seed_test_subscription(&mut state, channel_id); - let handled = drop_channel_on_access_denied( - &mut state, - &channel_sub_id(channel_id), - "auth-required: not authenticated", - ); + let task = tokio::spawn(async move { + let mut client = client; + let result = resubscribe_after_reconnect( + &mut client, + &mut cmd_rx, + &mut state, + "agent-pubkey", + true, + ) + .await; + (result, state) + }); - assert!( - !handled, - "auth-required must fall through to reconnect, not drop the channel" - ); - assert!(state.active_subscriptions.contains_key(&channel_id)); + let replay = next_test_frame(&mut server).await; + assert_eq!(replay[1], channel_sub_id(channel_id)); + cmd_tx + .send(RelayCommand::Unsubscribe { channel_id }) + .await + .expect("queue unsubscribe during pacing"); + + let close = next_test_frame(&mut server).await; + assert_eq!(close, json!(["CLOSE", channel_sub_id(channel_id)])); + let (result, state) = task.await.expect("join resubscribe task"); + assert!(matches!(result, ResubscribeResult::Ok)); + assert!(!state.active_subscriptions.contains_key(&channel_id)); } - #[test] - fn already_removed_channel_is_a_no_op() { + #[tokio::test] + async fn publish_during_replay_pacing_is_sent_on_live_socket() { + let (client, mut server) = test_ws_pair().await; + let (cmd_tx, mut cmd_rx) = mpsc::channel(4); let mut state = BgState::new(); let channel_id = Uuid::new_v4(); - // Channel was never subscribed (or already dropped) — a delayed CLOSED. + seed_test_subscription(&mut state, channel_id); + let event = make_test_event(&nostr::Keys::generate(), 2_000); + let event_id = event.id.to_hex(); - let handled = drop_channel_on_access_denied( - &mut state, - &channel_sub_id(channel_id), - "restricted: not a channel member", - ); + let task = tokio::spawn(async move { + let mut client = client; + let result = resubscribe_after_reconnect( + &mut client, + &mut cmd_rx, + &mut state, + "agent-pubkey", + true, + ) + .await; + result + }); + let replay = next_test_frame(&mut server).await; + assert_eq!(replay[1], channel_sub_id(channel_id)); + cmd_tx + .send(RelayCommand::PublishEvent { + event: Box::new(event), + admission: None, + }) + .await + .expect("queue publish during pacing"); + + let publish = next_test_frame(&mut server).await; + assert_eq!(publish[0], "EVENT"); + assert_eq!(publish[1]["id"], event_id); + assert!(matches!( + task.await.expect("join resubscribe task"), + ResubscribeResult::Ok + )); + } + + #[test] + fn failed_replay_retains_deferred_subscription_intent_in_fifo_order() { + let mut state = BgState::new(); + let kept_channel = Uuid::new_v4(); + let removed_channel = Uuid::new_v4(); + seed_test_subscription(&mut state, removed_channel); + let event = make_test_event(&nostr::Keys::generate(), 2_000); + let mut deferred = VecDeque::from([ + RelayCommand::Subscribe { + channel_id: kept_channel, + filter: test_channel_filter(), + replay_since: Some(2_000), + }, + RelayCommand::Unsubscribe { + channel_id: removed_channel, + }, + RelayCommand::PublishEvent { + event: Box::new(event), + admission: None, + }, + ]); + + retain_deferred_command_intent(&mut state, &mut deferred); + + assert!(deferred.is_empty()); + assert!(state.active_subscriptions.contains_key(&kept_channel)); + assert!(!state.active_subscriptions.contains_key(&removed_channel)); + } + + #[test] + fn bg_state_dedup_first_event_accepted() { + let mut state = BgState::new(); + let channel_id = Uuid::new_v4(); + let keys = nostr::Keys::generate(); + let event = make_test_event(&keys, 1_000_000); assert!( - handled, - "an exact per-channel denial is still handled (keep socket) even if the channel is gone" - ); - assert!( - !state.active_subscriptions.contains_key(&channel_id), - "no-op: nothing to remove and nothing resurrected" + state.record_event(channel_id, &event), + "first event should be accepted" ); } #[test] - fn dropped_channel_is_not_resubscribed_so_loop_cannot_re_form() { + fn bg_state_dedup_duplicate_rejected() { let mut state = BgState::new(); let channel_id = Uuid::new_v4(); - subscribe_channel(&mut state, channel_id); - - drop_channel_on_access_denied( - &mut state, - &channel_sub_id(channel_id), - "restricted: not a channel member", + let keys = nostr::Keys::generate(); + let event = make_test_event(&keys, 1_000_000); + assert!( + state.record_event(channel_id, &event), + "first should be accepted" ); - - // Simulate a reconnect: only channels still in active_subscriptions are - // restored. The dropped channel must not be among them — otherwise the - // forbidden channel would be resubscribed and earn the same CLOSED again. - let resubscribed: Vec = state.active_subscriptions.keys().copied().collect(); assert!( - !resubscribed.contains(&channel_id), - "the dropped channel must not be resubscribed — the loop cannot re-form" + !state.record_event(channel_id, &event), + "duplicate should be rejected" ); } - // ── startup connect retry ──────────────────────────────────────────── + #[test] + fn bg_state_dedup_different_ids_both_accepted() { + let mut state = BgState::new(); + let channel_id = Uuid::new_v4(); + // Two different keys → two different event IDs. + let keys1 = nostr::Keys::generate(); + let keys2 = nostr::Keys::generate(); + let event1 = make_test_event(&keys1, 1_000_000); + let event2 = make_test_event(&keys2, 1_000_001); + assert!(state.record_event(channel_id, &event1)); + assert!(state.record_event(channel_id, &event2)); + } - /// Table-driven coverage of every `RelayError` variant and every - /// `tungstenite::Error` inner variant. Exhaustive — adding a new - /// tungstenite variant without updating this table is a compile error - /// in `is_terminal_ws_error` (no wildcard), and a missing row here - /// is a code-review gap, not a silent misclassification. #[test] - fn connect_error_classification_matches_every_relay_error_variant() { - use tokio_tungstenite::tungstenite::error::{ - CapacityError, Error as WsError, ProtocolError, SubProtocolError, TlsError, UrlError, - }; - use tokio_tungstenite::tungstenite::http; + fn bg_state_last_seen_set_on_first_event() { + let mut state = BgState::new(); + let channel_id = Uuid::new_v4(); + let keys = nostr::Keys::generate(); + let event = make_test_event(&keys, 1_700_000); + state.record_event(channel_id, &event); + assert_eq!(state.last_seen.get(&channel_id).copied(), Some(1_700_000)); + } - fn ws(e: WsError) -> RelayError { - RelayError::WebSocket(Box::new(e)) - } + #[test] + fn bg_state_last_seen_advances_on_newer_event() { + let mut state = BgState::new(); + let channel_id = Uuid::new_v4(); + let keys1 = nostr::Keys::generate(); + let keys2 = nostr::Keys::generate(); + let event1 = make_test_event(&keys1, 1_700_000); + let event2 = make_test_event(&keys2, 1_800_000); + state.record_event(channel_id, &event1); + state.record_event(channel_id, &event2); + assert_eq!(state.last_seen.get(&channel_id).copied(), Some(1_800_000)); + } - let cases: Vec<(&str, RelayError, bool)> = vec![ - // ── outer RelayError variants ── - ("Http: bad URL", RelayError::Http("bad url".into()), true), - ( - "Json: malformed relay frame", - RelayError::Json(serde_json::from_str::<()>("not json").unwrap_err()), - true, - ), - ( - "UnexpectedMessage: unknown frame type", - RelayError::UnexpectedMessage("unknown message type: WAT".into()), - true, - ), + #[test] + fn bg_state_last_seen_does_not_regress_on_older_event() { + let mut state = BgState::new(); + let channel_id = Uuid::new_v4(); + let keys1 = nostr::Keys::generate(); + let keys2 = nostr::Keys::generate(); + let event_new = make_test_event(&keys1, 1_800_000); + let event_old = make_test_event(&keys2, 1_700_000); + state.record_event(channel_id, &event_new); + state.record_event(channel_id, &event_old); + // last_seen should remain at the higher timestamp + assert_eq!(state.last_seen.get(&channel_id).copied(), Some(1_800_000)); + } + + #[test] + fn bg_state_last_seen_independent_per_channel() { + let mut state = BgState::new(); + let ch1 = Uuid::new_v4(); + let ch2 = Uuid::new_v4(); + let keys1 = nostr::Keys::generate(); + let keys2 = nostr::Keys::generate(); + let event1 = make_test_event(&keys1, 1_000_000); + let event2 = make_test_event(&keys2, 2_000_000); + state.record_event(ch1, &event1); + state.record_event(ch2, &event2); + assert_eq!(state.last_seen.get(&ch1).copied(), Some(1_000_000)); + 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(); + + // 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)); + } + + // 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" + ); + + // 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.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" + ); + } + + /// Test 8: channel_dropped_since records the OLDEST dropped timestamp. + /// + /// Simulates the backpressure path directly on BgState: + /// - First drop at ts=1000 → entry is 1000 + /// - Second drop at ts=2000 (later) → entry stays 1000 (min) + /// - Third drop at ts=500 (earlier) → entry updates to 500 (min) + #[test] + fn acp_records_channel_dropped_since_on_backpressure() { + let mut state = BgState::new(); + let channel_id = Uuid::new_v4(); + + // Simulate the backpressure path: record ts=1000. + let ts1: u64 = 1_000; + state + .channel_dropped_since + .entry(channel_id) + .and_modify(|d| *d = (*d).min(ts1)) + .or_insert(ts1); + assert_eq!( + state.channel_dropped_since.get(&channel_id).copied(), + Some(1_000), + "first drop should record ts=1000" + ); + + // Later timestamp (2000) — entry should stay at 1000. + let ts2: u64 = 2_000; + state + .channel_dropped_since + .entry(channel_id) + .and_modify(|d| *d = (*d).min(ts2)) + .or_insert(ts2); + assert_eq!( + state.channel_dropped_since.get(&channel_id).copied(), + Some(1_000), + "later drop should not overwrite earlier timestamp" + ); + + // Earlier timestamp (500) — entry should update to 500. + let ts3: u64 = 500; + state + .channel_dropped_since + .entry(channel_id) + .and_modify(|d| *d = (*d).min(ts3)) + .or_insert(ts3); + assert_eq!( + state.channel_dropped_since.get(&channel_id).copied(), + Some(500), + "earlier drop should update entry to 500" + ); + } + + /// Test 9: reconnect since filter = min(last_seen, channel_dropped_since) - SINCE_SKEW_SECS. + /// + /// With last_seen=1000 and channel_dropped_since=900, the effective since + /// passed to send_subscribe should be min(1000, 900) - SINCE_SKEW_SECS = 895. + #[test] + fn acp_reconnect_uses_dropped_since_for_replay() { + let mut state = BgState::new(); + let channel_id = Uuid::new_v4(); + + // Set up state: last_seen=1000, channel_dropped_since=900. + state.last_seen.insert(channel_id, 1_000); + state.channel_dropped_since.insert(channel_id, 900); + + // Compute the since value the reconnect path would use. + let since = state.channel_since(&channel_id); + + // The since passed to send_subscribe (which subtracts SINCE_SKEW_SECS internally). + assert_eq!(since, Some(900), "since should be min(1000, 900) = 900"); + + // After subtracting skew (as send_subscribe does), the REQ filter value is: + let req_since = since.unwrap().saturating_sub(SINCE_SKEW_SECS); + assert_eq!( + req_since, 895, + "REQ since filter should be 900 - {} = 895", + SINCE_SKEW_SECS + ); + + // Simulate clearing after resubscribe. + state.channel_dropped_since.remove(&channel_id); + assert!( + !state.channel_dropped_since.contains_key(&channel_id), + "channel_dropped_since should be cleared after resubscribe" + ); + } + + #[test] + fn dynamic_subscribe_records_membership_replay_floor() { + let mut state = BgState::new(); + state.startup_watermark = Some(2_000); + let channel_id = Uuid::new_v4(); + let membership_ts = 10_000; + let filter = ChannelFilter { + kinds: Some(vec![9]), + require_mention: true, + }; + + apply_command_to_state( + &mut state, + RelayCommand::Subscribe { + channel_id, + filter, + replay_since: Some(membership_ts), + }, + ); + + assert_eq!( + state.subscribe_since.get(&channel_id).copied(), + Some(membership_ts), + "dynamic channel subscriptions should replay from the membership notification, not startup" + ); + assert_eq!( + state.channel_since(&channel_id), + Some(membership_ts), + "channel_since should use the dynamic replay floor until an event is seen" + ); + } + + /// 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" + ); + } + + /// Subscribe a channel via the production command path so the test exercises + /// real subscription state (active_subscriptions + active_filters + since). + fn subscribe_channel(state: &mut BgState, channel_id: Uuid) { + apply_command_to_state( + state, + RelayCommand::Subscribe { + channel_id, + filter: ChannelFilter { + kinds: Some(vec![9]), + require_mention: false, + }, + replay_since: Some(1_000), + }, + ); + } + + #[test] + fn not_a_channel_member_drops_channel_without_reconnect() { + let mut state = BgState::new(); + let channel_id = Uuid::new_v4(); + subscribe_channel(&mut state, channel_id); + + let handled = drop_channel_on_access_denied( + &mut state, + &channel_sub_id(channel_id), + "restricted: not a channel member", + ); + + assert!(handled, "per-channel denial must be handled (no reconnect)"); + assert!( + !state.active_subscriptions.contains_key(&channel_id), + "the forbidden channel's subscription must be dropped" + ); + assert!( + !state.active_filters.contains_key(&channel_id), + "channel state must be cleared (Unsubscribe cleanup)" + ); + } + + #[test] + fn channel_access_revoked_drops_channel_without_reconnect() { + let mut state = BgState::new(); + let channel_id = Uuid::new_v4(); + subscribe_channel(&mut state, channel_id); + + let handled = drop_channel_on_access_denied( + &mut state, + &channel_sub_id(channel_id), + "restricted: channel access revoked", + ); + + assert!(handled, "per-channel denial must be handled (no reconnect)"); + assert!(!state.active_subscriptions.contains_key(&channel_id)); + assert!(!state.active_filters.contains_key(&channel_id)); + } + + #[test] + fn insufficient_scope_is_not_dropped_and_reconnects() { + let mut state = BgState::new(); + let channel_id = Uuid::new_v4(); + subscribe_channel(&mut state, channel_id); + + let handled = drop_channel_on_access_denied( + &mut state, + &channel_sub_id(channel_id), + "restricted: insufficient scope", + ); + + assert!( + !handled, + "connection-level insufficient-scope must fall through to reconnect, not drop the channel" + ); + assert!( + state.active_subscriptions.contains_key(&channel_id), + "the channel must survive so reconnect can restore it" + ); + } + + #[test] + fn auth_required_is_not_dropped_and_reconnects() { + let mut state = BgState::new(); + let channel_id = Uuid::new_v4(); + subscribe_channel(&mut state, channel_id); + + let handled = drop_channel_on_access_denied( + &mut state, + &channel_sub_id(channel_id), + "auth-required: not authenticated", + ); + + assert!( + !handled, + "auth-required must fall through to reconnect, not drop the channel" + ); + assert!(state.active_subscriptions.contains_key(&channel_id)); + } + + #[test] + fn already_removed_channel_is_a_no_op() { + let mut state = BgState::new(); + let channel_id = Uuid::new_v4(); + // Channel was never subscribed (or already dropped) — a delayed CLOSED. + + let handled = drop_channel_on_access_denied( + &mut state, + &channel_sub_id(channel_id), + "restricted: not a channel member", + ); + + assert!( + handled, + "an exact per-channel denial is still handled (keep socket) even if the channel is gone" + ); + assert!( + !state.active_subscriptions.contains_key(&channel_id), + "no-op: nothing to remove and nothing resurrected" + ); + } + + #[test] + fn dropped_channel_is_not_resubscribed_so_loop_cannot_re_form() { + let mut state = BgState::new(); + let channel_id = Uuid::new_v4(); + subscribe_channel(&mut state, channel_id); + + drop_channel_on_access_denied( + &mut state, + &channel_sub_id(channel_id), + "restricted: not a channel member", + ); + + // Simulate a reconnect: only channels still in active_subscriptions are + // restored. The dropped channel must not be among them — otherwise the + // forbidden channel would be resubscribed and earn the same CLOSED again. + let resubscribed: Vec = state.active_subscriptions.keys().copied().collect(); + assert!( + !resubscribed.contains(&channel_id), + "the dropped channel must not be resubscribed — the loop cannot re-form" + ); + } + + // ── startup connect retry ──────────────────────────────────────────── + + /// Table-driven coverage of every `RelayError` variant and every + /// `tungstenite::Error` inner variant. Exhaustive — adding a new + /// tungstenite variant without updating this table is a compile error + /// in `is_terminal_ws_error` (no wildcard), and a missing row here + /// is a code-review gap, not a silent misclassification. + #[test] + fn connect_error_classification_matches_every_relay_error_variant() { + use tokio_tungstenite::tungstenite::error::{ + CapacityError, Error as WsError, ProtocolError, SubProtocolError, TlsError, UrlError, + }; + use tokio_tungstenite::tungstenite::http; + + fn ws(e: WsError) -> RelayError { + RelayError::WebSocket(Box::new(e)) + } + + let cases: Vec<(&str, RelayError, bool)> = vec![ + // ── outer RelayError variants ── + ("Http: bad URL", RelayError::Http("bad url".into()), true), + ( + "Json: malformed relay frame", + RelayError::Json(serde_json::from_str::<()>("not json").unwrap_err()), + true, + ), + ( + "UnexpectedMessage: unknown frame type", + RelayError::UnexpectedMessage("unknown message type: WAT".into()), + true, + ), ( "AuthFailed: relay dependency fault (NIP-01 `error:` prefix)", RelayError::AuthFailed("error: internal error checking restriction state".into()), @@ -5489,506 +8223,1952 @@ mod tests { ))), false, ), - ( - "WebSocket(ConnectionClosed): link-level closure", - ws(WsError::ConnectionClosed), + ( + "WebSocket(ConnectionClosed): link-level closure", + ws(WsError::ConnectionClosed), + false, + ), + // ── WebSocket(Tls): deterministic config (terminal, pins the arm) ── + // These shapes are constructible but not reachable through our + // rustls production connector — cert failures arrive as Io above. + // Kept to pin the Tls(_) => true arm. + ( + "Tls(Rustls(General)): pins Tls arm terminal", + ws(WsError::Tls( + rustls::Error::General("tls handshake failed".into()).into(), + )), + true, + ), + ( + "Tls(InvalidDnsName): only reachable connect-time Tls variant", + ws(WsError::Tls(TlsError::InvalidDnsName)), + true, + ), + ( + "Tls(Rustls(InvalidCertificate(Expired))): pins Tls arm terminal", + ws(WsError::Tls( + rustls::Error::InvalidCertificate(rustls::CertificateError::Expired).into(), + )), + true, + ), + ( + "WebSocket(AlreadyClosed): unreachable at connect, fail-safe transient", + ws(WsError::AlreadyClosed), + false, + ), + ( + "WebSocket(WriteBufferFull): unreachable at connect, fail-safe transient", + ws(WsError::WriteBufferFull(Box::new( + tokio_tungstenite::tungstenite::Message::Text("x".into()), + ))), + false, + ), + ]; + + for (label, err, want_terminal) in cases { + assert_eq!( + is_terminal_connect_error(&err), + want_terminal, + "{label}: expected terminal={want_terminal}" + ); + } + } + + /// A literal `https://…` URL through production `do_connect()` must fail + /// fast as terminal — the relay endpoint is a plain HTTPS server, not a + /// WebSocket endpoint, and tungstenite returns `Error::Http` (non-101 + /// response) or `Error::Url(UnsupportedUrlScheme)` depending on how far + /// the handshake gets. Either way it must not be retried. + #[tokio::test] + async fn do_connect_wrong_scheme_is_terminal() { + let keys = nostr::Keys::generate(); + let err = do_connect("https://example.com", &keys, None) + .await + .unwrap_err(); + assert!( + is_terminal_connect_error(&err), + "wrong-scheme URL should be terminal, got: {err}" + ); + } + + /// A transient failure (e.g. connection dropped mid-handshake on a spotty + /// link) must be retried and can still succeed once the link recovers. + #[tokio::test(start_paused = true)] + async fn retry_initial_connect_retries_transient_failure_then_succeeds() { + use std::sync::atomic::{AtomicUsize, Ordering}; + + let attempts = AtomicUsize::new(0); + let result: Result<&'static str, RelayError> = retry_initial_connect(|| { + let n = attempts.fetch_add(1, Ordering::SeqCst); + async move { + if n < 2 { + Err(RelayError::ConnectionClosed) + } else { + Ok("connected") + } + } + }) + .await; + + assert_eq!(result.unwrap(), "connected"); + assert_eq!( + attempts.load(Ordering::SeqCst), + 3, + "should succeed on the 3rd attempt (2 transient failures + 1 success)" + ); + } + + /// A terminal error (bad auth, bad config) must not be retried — the + /// same call would fail identically every time, so retrying just delays + /// surfacing a real problem to the caller. + #[tokio::test(start_paused = true)] + async fn retry_initial_connect_does_not_retry_terminal_error() { + use std::sync::atomic::{AtomicUsize, Ordering}; + + let attempts = AtomicUsize::new(0); + let result: Result<(), RelayError> = retry_initial_connect(|| { + attempts.fetch_add(1, Ordering::SeqCst); + async { Err(RelayError::AuthFailed("invalid: bad signature".into())) } + }) + .await; + + assert!(matches!(result, Err(RelayError::AuthFailed(_)))); + assert_eq!( + attempts.load(Ordering::SeqCst), + 1, + "a terminal error must fail on the first attempt with no retries" + ); + } + + /// A relay-side dependency fault (NIP-01 `error:` prefix) is transient — + /// the relay is failing closed on itself, not rejecting this identity — + /// so it must be retried rather than surfaced immediately like a real + /// auth rejection. + #[tokio::test(start_paused = true)] + async fn retry_initial_connect_retries_relay_dependency_fault() { + use std::sync::atomic::{AtomicUsize, Ordering}; + + let attempts = AtomicUsize::new(0); + let result: Result<&'static str, RelayError> = retry_initial_connect(|| { + let n = attempts.fetch_add(1, Ordering::SeqCst); + async move { + if n < 1 { + Err(RelayError::AuthFailed( + "error: internal error checking restriction state".into(), + )) + } else { + Ok("connected") + } + } + }) + .await; + + assert_eq!(result.unwrap(), "connected"); + assert_eq!( + attempts.load(Ordering::SeqCst), + 2, + "a relay dependency fault must be retried, not surfaced immediately" + ); + } + + /// Once every attempt (1 initial + N backoff retries) is exhausted, the + /// last transient error is returned rather than retrying forever — a + /// dead relay must not hang agent startup indefinitely. + #[tokio::test(start_paused = true)] + async fn retry_initial_connect_exhausts_and_returns_last_error() { + use std::sync::atomic::{AtomicUsize, Ordering}; + + let attempts = AtomicUsize::new(0); + let result: Result<(), RelayError> = retry_initial_connect(|| { + attempts.fetch_add(1, Ordering::SeqCst); + async { Err(RelayError::Timeout) } + }) + .await; + + assert!( + matches!(result, Err(RelayError::Timeout)), + "must surface the last attempt's error, not a generic one" + ); + assert_eq!( + attempts.load(Ordering::SeqCst), + STARTUP_CONNECT_BACKOFFS.len() + 1, + "must attempt exactly once plus one retry per backoff entry" + ); + } + + /// Backoff sleeps must actually elapse (not be skipped) — this pins the + /// bounded-but-real-delay contract using `tokio::time::pause` so the + /// test itself stays fast (virtual time, not wall-clock sleeps). + #[tokio::test(start_paused = true)] + async fn retry_initial_connect_sleeps_between_attempts() { + use std::sync::atomic::{AtomicUsize, Ordering}; + + let attempts = AtomicUsize::new(0); + let call = retry_initial_connect(|| { + let n = attempts.fetch_add(1, Ordering::SeqCst); + async move { + if n < 1 { + Err(RelayError::ConnectionClosed) + } else { + Ok(()) + } + } + }); + tokio::pin!(call); + + // Before the first backoff elapses, the retry must still be pending + // (i.e. it actually slept rather than immediately retrying). + tokio::select! { + biased; + _ = tokio::time::sleep(Duration::from_millis(1)) => {} + _ = &mut call => panic!("must not resolve before the backoff sleep elapses"), + } + assert_eq!(attempts.load(Ordering::SeqCst), 1); + + // Advancing past the (jittered, ≤1.2x) first backoff lets it proceed. + let result = call.await; + assert!(result.is_ok()); + assert_eq!(attempts.load(Ordering::SeqCst), 2); + } + + // ── Rate-limit gate, pacing, backoff reset, DNS ────────────────────────── + + /// parse_rate_limit_retry_secs: full hint extracts the N from "retry in Ns". + #[test] + fn parse_rate_limit_retry_secs_with_hint() { + assert_eq!( + parse_rate_limit_retry_secs("rate-limited: quota exceeded; retry in 12s"), + Some(12) + ); + } + + /// parse_rate_limit_retry_secs: message without a hint returns None. + #[test] + fn parse_rate_limit_retry_secs_missing_hint() { + assert_eq!( + parse_rate_limit_retry_secs("rate-limited: too many concurrent requests"), + None + ); + } + + /// parse_rate_limit_retry_secs: explicit zero value is returned as Some(0). + #[test] + fn parse_rate_limit_retry_secs_zero() { + assert_eq!( + parse_rate_limit_retry_secs("rate-limited: quota exceeded; retry in 0s"), + Some(0) + ); + } + + /// parse_rate_limit_retry_secs: garbage input returns None. + #[test] + fn parse_rate_limit_retry_secs_garbage() { + assert_eq!( + parse_rate_limit_retry_secs("not a rate limit message"), + None + ); + } + + /// set_rate_limit_gate arms the gate with jittered expiry from the hint. + /// check_rate_gate returns Some while active and lazily clears on expiry. + #[tokio::test(start_paused = true)] + async fn rate_limit_gate_set_and_expiry() { + let mut state = BgState::new(); + assert!( + state.check_rate_gate().is_none(), + "gate must start inactive" + ); + + // Arm with a 5 s hint. + state.set_rate_limit_gate(5); + assert!( + state.check_rate_gate().is_some(), + "gate must be active immediately after arming" + ); + + // Advance virtual time past the max jitter (1.2 × 5 s = 6 s). + tokio::time::advance(Duration::from_secs(7)).await; + + assert!( + state.check_rate_gate().is_none(), + "gate must have expired after 7s" + ); + assert!( + state.rate_limit_gate.is_none(), + "check_rate_gate must lazily clear the field on expiry" + ); + } + + /// set_rate_limit_gate takes the max of overlapping deadlines. + #[tokio::test(start_paused = true)] + async fn rate_limit_gate_extends_to_max() { + let mut state = BgState::new(); + + // Arm with a long hint first. + state.set_rate_limit_gate(30); + let first_deadline = state.rate_limit_gate.unwrap(); + + // A shorter subsequent hint must NOT shorten the existing gate. + state.set_rate_limit_gate(1); + let second_deadline = state.rate_limit_gate.unwrap(); + + assert_eq!( + first_deadline, second_deadline, + "shorter hint must not overwrite a later existing deadline" + ); + } + + #[test] + fn terminal_confirmation_window_outlives_maximum_accepted_rate_limit_gate() { + assert!( + TERMINAL_OBSERVER_ACK_TIMEOUT + > Duration::from_secs(MAX_RATE_LIMIT_RETRY_SECS).mul_f64(1.2), + "terminal confirmation must remain live beyond the largest accepted retry hint plus maximum jitter" + ); + } + + #[tokio::test(start_paused = true)] + async fn oversized_rate_limit_hint_is_bounded_below_terminal_confirmation_window() { + let mut state = BgState::new(); + let now = tokio::time::Instant::now(); + let deadline = state.set_rate_limit_gate(u64::MAX); + assert!( + deadline.duration_since(now) >= Duration::from_secs(MAX_RATE_LIMIT_RETRY_SECS), + "an explicit reset hint must never be retried before its advertised window" + ); + assert!( + deadline.duration_since(now) + <= Duration::from_secs(MAX_RATE_LIMIT_RETRY_SECS).mul_f64(1.2), + "untrusted retry hints must not extend the local gate beyond its bounded contract" + ); + assert!( + deadline.duration_since(now) < TERMINAL_OBSERVER_ACK_TIMEOUT, + "the accepted gate must expire while the terminal publisher is still awaiting proof" + ); + } + + /// Build a signed observer telemetry frame (kind 24200) for gate tests. + fn make_observer_frame(keys: &Keys) -> Event { + let recipient = Keys::generate(); + let encrypted = buzz_core::observer::encrypt_observer_payload( + keys, + &recipient.public_key(), + &json!({"type": "test"}), + ) + .expect("encrypt test observer payload"); + buzz_sdk::build_agent_observer_frame( + &recipient.public_key().to_hex(), + &keys.public_key().to_hex(), + "telemetry", + &encrypted, + ) + .expect("build test observer frame") + .sign_with_keys(keys) + .expect("sign test observer frame") + } + + fn make_priority_observer_frame(keys: &Keys) -> Event { + let recipient = Keys::generate(); + let encrypted = buzz_core::observer::encrypt_observer_payload( + keys, + &recipient.public_key(), + &json!({"kind": "control_result"}), + ) + .expect("encrypt test priority observer payload"); + buzz_sdk::build_agent_observer_frame( + &recipient.public_key().to_hex(), + &keys.public_key().to_hex(), + "telemetry", + &encrypted, + ) + .expect("build test priority observer frame") + .tag( + Tag::parse([OBSERVER_PRIORITY_TAG, OBSERVER_CONTROL_RESULT_PRIORITY]) + .expect("valid priority tag"), + ) + .sign_with_keys(keys) + .expect("sign test priority observer frame") + } + + fn make_observer_control_event( + owner_keys: &Keys, + agent_keys: &Keys, + created_at_secs: u64, + ) -> Event { + let agent_pubkey_hex = agent_keys.public_key().to_hex(); + let encrypted = buzz_core::observer::encrypt_observer_payload( + owner_keys, + &agent_keys.public_key(), + &json!({"type": "cancel_turn", "channelId": Uuid::new_v4()}), + ) + .expect("encrypt test observer control"); + EventBuilder::new(Kind::Custom(KIND_AGENT_OBSERVER_FRAME as u16), encrypted) + .tags([ + Tag::parse(["p", agent_pubkey_hex.as_str()]).expect("valid recipient tag"), + Tag::parse(["agent", agent_pubkey_hex.as_str()]).expect("valid agent tag"), + Tag::parse(["frame", "control"]).expect("valid frame tag"), + ]) + .custom_created_at(nostr::Timestamp::from(created_at_secs)) + .sign_with_keys(owner_keys) + .expect("sign test observer control") + } + + #[tokio::test] + async fn observer_controls_require_live_strictly_post_start_subscription() { + let (mut ws, _server) = test_ws_pair().await; + let (event_tx, _event_rx) = mpsc::channel(1); + let (observer_control_tx, mut observer_control_rx) = mpsc::channel(2); + let mut state = BgState::new(); + let agent_keys = Keys::generate(); + let owner_keys = Keys::generate(); + let agent_pubkey_hex = agent_keys.public_key().to_hex(); + let now = unix_now_secs(); + state.authorized_owner_pubkey = Some(owner_keys.public_key().to_hex()); + state.privileged_control_start_boundary = now.saturating_sub(2); + let control = make_observer_control_event(&owner_keys, &agent_keys, now); + let control_json = + serde_json::to_string(&json!(["EVENT", OBSERVER_CONTROL_SUB_ID, control.clone()])) + .expect("serialize observer control"); + + assert!( + handle_ws_message( + Message::Text(control_json.clone().into()), + &mut ws, + &event_tx, + &observer_control_tx, + &mut state, + &agent_keys, + "ws://localhost:3000", + &agent_pubkey_hex, + None, + ) + .await + ); + assert!( + observer_control_rx.try_recv().is_err(), + "a matching subscription ID before the REQ is live must not admit controls" + ); + + state.observer_control_sub_active = true; + state.observer_control_subscribed_at = Some(now); + assert!( + handle_ws_message( + Message::Text(control_json.clone().into()), + &mut ws, + &event_tx, + &observer_control_tx, + &mut state, + &agent_keys, + "ws://localhost:3000", + &agent_pubkey_hex, + None, + ) + .await + ); + assert!( + observer_control_rx.try_recv().is_err(), + "same-second controls cannot be proven to post-date subscription startup" + ); + + state.observer_control_subscribed_at = Some(now.saturating_sub(1)); + assert!( + handle_ws_message( + Message::Text(control_json.into()), + &mut ws, + &event_tx, + &observer_control_tx, + &mut state, + &agent_keys, + "ws://localhost:3000", + &agent_pubkey_hex, + None, + ) + .await + ); + assert_eq!( + observer_control_rx + .try_recv() + .expect("valid owner control after subscription should be delivered") + .id, + control.id + ); + } + + #[tokio::test] + async fn invalid_observer_flood_cannot_displace_valid_owner_control() { + let (mut ws, _server) = test_ws_pair().await; + let (event_tx, _event_rx) = mpsc::channel(1); + let (observer_control_tx, mut observer_control_rx) = mpsc::channel(1); + let mut state = BgState::new(); + let agent_keys = Keys::generate(); + let owner_keys = Keys::generate(); + let non_owner_keys = Keys::generate(); + let other_agent_keys = Keys::generate(); + let agent_pubkey_hex = agent_keys.public_key().to_hex(); + let now = unix_now_secs(); + state.authorized_owner_pubkey = Some(owner_keys.public_key().to_hex()); + state.privileged_control_start_boundary = now.saturating_sub(2); + state.observer_control_sub_active = true; + state.observer_control_subscribed_at = Some(now.saturating_sub(1)); + + let valid = make_observer_control_event(&owner_keys, &agent_keys, now); + let mut tampered_json = serde_json::to_value(&valid).expect("serialize valid control"); + let original_content = tampered_json["content"] + .as_str() + .expect("observer ciphertext") + .to_string(); + let replacement = if let Some(rest) = original_content.strip_prefix('A') { + format!("B{rest}") + } else { + let mut chars = original_content.chars(); + let _ = chars.next(); + format!("A{}", chars.as_str()) + }; + tampered_json["content"] = Value::String(replacement); + let tampered: Event = + serde_json::from_value(tampered_json).expect("parse tampered observer control"); + assert!(buzz_core::verify_event(&tampered).is_err()); + + let non_owner = make_observer_control_event(&non_owner_keys, &agent_keys, now); + let wrong_target = make_observer_control_event(&owner_keys, &other_agent_keys, now); + let wrong_kind = make_test_channel_event( + &owner_keys, + Uuid::new_v4(), + KIND_STREAM_MESSAGE, + "not an observer control", + Some(&agent_pubkey_hex), + now, + ); + for invalid in [tampered, non_owner, wrong_target, wrong_kind] { + for _ in 0..16 { + let frame = Message::Text( + serde_json::to_string(&json!(["EVENT", OBSERVER_CONTROL_SUB_ID, invalid])) + .expect("serialize invalid observer frame") + .into(), + ); + assert!( + handle_ws_message( + frame, + &mut ws, + &event_tx, + &observer_control_tx, + &mut state, + &agent_keys, + "ws://localhost:3000", + &agent_pubkey_hex, + None, + ) + .await + ); + } + } + assert!( + observer_control_rx.try_recv().is_err(), + "invalid controls must be filtered before they consume queue capacity" + ); + + let valid_frame = Message::Text( + serde_json::to_string(&json!(["EVENT", OBSERVER_CONTROL_SUB_ID, valid.clone()])) + .expect("serialize valid observer control") + .into(), + ); + assert!( + handle_ws_message( + valid_frame, + &mut ws, + &event_tx, + &observer_control_tx, + &mut state, + &agent_keys, + "ws://localhost:3000", + &agent_pubkey_hex, + None, + ) + .await + ); + assert_eq!( + observer_control_rx + .try_recv() + .expect("valid owner control retains queue capacity") + .id, + valid.id + ); + } + + #[tokio::test] + async fn valid_owner_control_queues_without_blocking_and_delivers_after_recovery() { + let (mut ws, _server) = test_ws_pair().await; + let (event_tx, _event_rx) = mpsc::channel(1); + let (observer_control_tx, mut observer_control_rx) = mpsc::channel(1); + let mut state = BgState::new(); + let agent_keys = Keys::generate(); + let owner_keys = Keys::generate(); + let agent_pubkey_hex = agent_keys.public_key().to_hex(); + let now = unix_now_secs(); + state.authorized_owner_pubkey = Some(owner_keys.public_key().to_hex()); + state.privileged_control_start_boundary = now.saturating_sub(2); + state.observer_control_sub_active = true; + state.observer_control_subscribed_at = Some(now.saturating_sub(1)); + + let queued = make_observer_control_event(&owner_keys, &agent_keys, now); + observer_control_tx + .try_send(queued.clone()) + .expect("prefill observer channel"); + let waiting = make_observer_control_event(&owner_keys, &agent_keys, now); + let waiting_frame = Message::Text( + serde_json::to_string(&json!(["EVENT", OBSERVER_CONTROL_SUB_ID, waiting.clone()])) + .expect("serialize waiting observer control") + .into(), + ); + + let should_continue = timeout( + Duration::from_secs(1), + handle_ws_message( + waiting_frame, + &mut ws, + &event_tx, + &observer_control_tx, + &mut state, + &agent_keys, + "ws://localhost:3000", + &agent_pubkey_hex, + None, + ), + ) + .await + .expect("full downstream queue must not block WebSocket message handling"); + + assert!(should_continue); + assert_eq!( + state.observer_control_pending.len(), + 1, + "authenticated control should be retained in the bounded pending queue" + ); + let first = observer_control_rx + .recv() + .await + .expect("prefilled control remains queued"); + let permit = observer_control_tx + .reserve() + .await + .expect("observer control receiver remains open"); + deliver_next_pending_observer_control(&mut state, permit); + let second = timeout(Duration::from_secs(1), observer_control_rx.recv()) + .await + .expect("valid control should resume after queue capacity returns") + .expect("observer control channel remains open"); + + assert_eq!(first.id, queued.id); + assert_eq!(second.id, waiting.id); + assert!(state.observer_control_pending.is_empty()); + } + + /// While the rate-limit gate is armed, an observer frame (kind 24200) is + /// parked — not silently dropped — and delivered by the drain once the + /// gate clears. A typing indicator in the same window stays dropped. + #[tokio::test] + async fn gated_observer_frame_is_parked_then_drained_not_dropped() { + let (mut client, mut server) = test_ws_pair().await; + let mut state = BgState::new(); + let keys = Keys::generate(); + state.rate_limit_gate = Some(tokio::time::Instant::now() + Duration::from_millis(150)); + + // Observer frame while gated: parked, nothing on the wire. + let observer_frame = make_observer_frame(&keys); + let ok = execute_connected_command( + &mut client, + &mut state, + "agent-pubkey", + RelayCommand::PublishEvent { + event: Box::new(observer_frame.clone()), + admission: None, + }, + ) + .await; + assert!(ok); + assert_eq!( + state.gated_observer_pending.len(), + 1, + "observer frame must be parked while gated" + ); + + // Typing indicator while gated: still dropped, not parked. + let typing = EventBuilder::new(Kind::Custom(KIND_TYPING_INDICATOR as u16), "") + .tags([Tag::parse(["h", &Uuid::new_v4().to_string()]).unwrap()]) + .sign_with_keys(&keys) + .expect("sign typing indicator"); + let ok = execute_connected_command( + &mut client, + &mut state, + "agent-pubkey", + RelayCommand::PublishEvent { + event: Box::new(typing), + admission: None, + }, + ) + .await; + assert!(ok); + assert_eq!( + state.gated_observer_pending.len(), + 1, + "typing indicators must not be parked" + ); + assert!( + timeout(Duration::from_millis(50), server.next()) + .await + .is_err(), + "nothing may reach the wire while the gate is armed" + ); + + // Gate expires — the drain delivers the parked frame. + tokio::time::sleep(Duration::from_millis(160)).await; + assert_eq!( + drain_gated_observer_pending(&mut client, &mut state, 1).await, + GatedObserverDrainOutcome::Sent(1) + ); + assert!(state.gated_observer_pending.is_empty()); + let frame = next_test_frame(&mut server).await; + assert_eq!(frame[0], "EVENT"); + assert_eq!(frame[1]["id"], observer_frame.id.to_hex()); + assert_eq!( + frame[1]["kind"], + u64::from(KIND_AGENT_OBSERVER_FRAME), + "delivered frame must be the parked observer frame" + ); + } + + /// Observer frames arriving while earlier parked frames are still queued + /// are appended behind them (order preserved), even if the gate has + /// already expired. + #[tokio::test] + async fn observer_frames_queue_behind_parked_backlog_in_order() { + let (mut client, mut server) = test_ws_pair().await; + let mut state = BgState::new(); + let keys = Keys::generate(); + state.rate_limit_gate = Some(tokio::time::Instant::now() + Duration::from_millis(50)); + + let first = make_observer_frame(&keys); + let second = make_observer_frame(&keys); + for event in [&first, &second] { + let ok = execute_connected_command( + &mut client, + &mut state, + "agent-pubkey", + RelayCommand::PublishEvent { + event: Box::new(event.clone()), + admission: None, + }, + ) + .await; + assert!(ok); + } + assert_eq!(state.gated_observer_pending.len(), 2); + + // Gate expires but the backlog is not drained yet — a third frame must + // queue behind it rather than jumping ahead on the wire. + tokio::time::sleep(Duration::from_millis(60)).await; + let third = make_observer_frame(&keys); + let ok = execute_connected_command( + &mut client, + &mut state, + "agent-pubkey", + RelayCommand::PublishEvent { + event: Box::new(third.clone()), + admission: None, + }, + ) + .await; + assert!(ok); + assert_eq!( + state.gated_observer_pending.len(), + 3, + "frame must queue behind undrained backlog to preserve order" + ); + + for expected in [&first, &second, &third] { + assert_eq!( + drain_gated_observer_pending(&mut client, &mut state, 1).await, + GatedObserverDrainOutcome::Sent(1) + ); + let frame = next_test_frame(&mut server).await; + assert_eq!(frame[1]["id"], expected.id.to_hex(), "order preserved"); + } + assert!(state.gated_observer_pending.is_empty()); + } + + #[test] + fn observer_notice_requeues_unacknowledged_frames_and_ok_retires_them() { + let mut state = BgState::new(); + let keys = Keys::generate(); + let accepted = make_observer_frame(&keys); + let rejected = make_observer_frame(&keys); + let later = make_observer_frame(&keys); + + state.track_observer_in_flight(Box::new(accepted.clone())); + state.track_observer_in_flight(Box::new(rejected.clone())); + state.acknowledge_observer_frame(&accepted.id.to_hex()); + state.park_gated_observer_frame(Box::new(later.clone())); + state.requeue_observer_in_flight(); + + let ids: Vec<_> = state + .gated_observer_pending + .iter() + .map(|event| event.id) + .collect(); + assert_eq!(ids, [rejected.id, later.id]); + assert!(state.observer_in_flight.is_empty()); + } + + #[test] + fn failed_reconnect_candidate_requeues_all_observer_writes_and_clears_socket_authority() { + let mut state = BgState::new(); + let keys = Keys::generate(); + let ordinary = make_observer_frame(&keys); + let terminal = make_priority_observer_frame(&keys); + let channel_id = Uuid::new_v4(); + + state.track_observer_in_flight(Box::new(ordinary.clone())); + state.track_observer_in_flight(Box::new(terminal.clone())); + state.observer_control_subscribed_at = Some(unix_now_secs()); + state + .channel_subscribed_at + .insert(channel_id, unix_now_secs()); + + state.abandon_failed_reconnect_candidate(); + + assert_eq!( + state + .gated_observer_pending + .iter() + .map(|event| event.id) + .collect::>(), + [ordinary.id] + ); + assert_eq!( + state + .gated_observer_priority_pending + .iter() + .map(|event| event.id) + .collect::>(), + [terminal.id] + ); + assert!(state.observer_in_flight.is_empty()); + assert!(state.observer_priority_in_flight.is_empty()); + assert!(state.observer_control_subscribed_at.is_none()); + assert!(state.channel_subscribed_at.is_empty()); + } + + #[test] + fn delayed_ok_after_notice_resolves_terminal_confirmation_exactly_once() { + let mut state = BgState::new(); + let terminal = make_priority_observer_frame(&Keys::generate()); + let terminal_id = terminal.id.to_hex(); + let (confirmation_tx, mut confirmation_rx) = oneshot::channel(); + + assert!(state.track_observer_in_flight(Box::new(terminal.clone()))); + state.record_observer_confirmation(terminal_id.clone(), true, Some(confirmation_tx), true); + state.requeue_observer_in_flight(); + + state.acknowledge_observer_frame(&terminal_id); + assert_eq!( + confirmation_rx.try_recv(), + Ok(Ok(())), + "the delayed accepted OK must resolve the original publisher" + ); + assert!( + !state.contains_observer_priority_event(&terminal_id), + "accepted terminal proof must not remain queued for duplicate resend" + ); + state.acknowledge_observer_frame(&terminal_id); + assert_eq!( + confirmation_rx.try_recv(), + Err(oneshot::error::TryRecvError::Closed), + "a duplicate OK cannot resolve the publisher twice" + ); + } + + #[test] + fn delayed_terminal_denial_after_notice_resolves_confirmation_without_resend() { + let mut state = BgState::new(); + let terminal = make_priority_observer_frame(&Keys::generate()); + let terminal_id = terminal.id.to_hex(); + let (confirmation_tx, mut confirmation_rx) = oneshot::channel(); + + assert!(state.track_observer_in_flight(Box::new(terminal))); + state.record_observer_confirmation(terminal_id.clone(), true, Some(confirmation_tx), true); + state.requeue_observer_in_flight(); + + assert!(state.record_terminal_observer_denial(&terminal_id, "blocked: policy denied")); + assert_eq!( + confirmation_rx.try_recv(), + Ok(Err( + "relay denied terminal observer frame: blocked: policy denied".to_string() + )) + ); + assert!(!state.contains_observer_priority_event(&terminal_id)); + assert_eq!(state.observer_terminal_denied, 1); + assert!( + !state.record_terminal_observer_denial(&terminal_id, "blocked: policy denied"), + "a duplicate denial must not retire or count the same event twice" + ); + assert_eq!(state.observer_terminal_denied, 1); + } + + #[tokio::test] + async fn rate_limited_ok_false_requeues_exact_observer_frame_and_arms_gate() { + let (mut ws, _server) = test_ws_pair().await; + let (event_tx, _event_rx) = mpsc::channel(1); + let (observer_control_tx, _observer_control_rx) = mpsc::channel(1); + let mut state = BgState::new(); + let keys = Keys::generate(); + let first = make_observer_frame(&keys); + let rejected = make_observer_frame(&keys); + let later = make_observer_frame(&keys); + state.track_observer_in_flight(Box::new(first.clone())); + state.track_observer_in_flight(Box::new(rejected.clone())); + state.park_gated_observer_frame(Box::new(later.clone())); + + let ok_false = Message::Text( + serde_json::to_string(&json!([ + "OK", + rejected.id.to_hex(), false, - ), - // ── WebSocket(Tls): deterministic config (terminal, pins the arm) ── - // These shapes are constructible but not reachable through our - // rustls production connector — cert failures arrive as Io above. - // Kept to pin the Tls(_) => true arm. - ( - "Tls(Rustls(General)): pins Tls arm terminal", - ws(WsError::Tls( - rustls::Error::General("tls handshake failed".into()).into(), - )), - true, - ), - ( - "Tls(InvalidDnsName): only reachable connect-time Tls variant", - ws(WsError::Tls(TlsError::InvalidDnsName)), - true, - ), - ( - "Tls(Rustls(InvalidCertificate(Expired))): pins Tls arm terminal", - ws(WsError::Tls( - rustls::Error::InvalidCertificate(rustls::CertificateError::Expired).into(), - )), - true, - ), - ( - "WebSocket(AlreadyClosed): unreachable at connect, fail-safe transient", - ws(WsError::AlreadyClosed), + "rate-limited: quota exceeded; retry in 5s" + ])) + .expect("serialize OK false") + .into(), + ); + assert!( + handle_ws_message( + ok_false, + &mut ws, + &event_tx, + &observer_control_tx, + &mut state, + &keys, + "ws://localhost:3000", + &keys.public_key().to_hex(), + None, + ) + .await + ); + + assert!(state.check_rate_gate().is_some(), "retry pacing must arm"); + assert_eq!( + state + .gated_observer_pending + .iter() + .map(|event| event.id) + .collect::>(), + [rejected.id, later.id], + "only the exact rejected frame is requeued ahead of later telemetry" + ); + assert_eq!( + state + .observer_in_flight + .iter() + .map(|event| event.id) + .collect::>(), + [first.id], + "unrelated in-flight frames stay pending their own OK" + ); + } + + #[tokio::test(start_paused = true)] + async fn sixty_second_terminal_quota_waits_for_reset_then_accepts_delayed_ok() { + let (mut ws, mut server) = test_ws_pair().await; + let (event_tx, _event_rx) = mpsc::channel(1); + let (observer_control_tx, _observer_control_rx) = mpsc::channel(1); + let mut state = BgState::new(); + let keys = Keys::generate(); + let terminal = make_priority_observer_frame(&keys); + let terminal_id = terminal.id.to_hex(); + let (confirmation_tx, mut confirmation_rx) = oneshot::channel(); + + // Socket readiness uses wall-clock I/O; pause only the quota interval + // whose lower bound this test proves. + tokio::time::resume(); + assert!( + execute_connected_command( + &mut ws, + &mut state, + "agent-pubkey", + RelayCommand::PublishEvent { + event: Box::new(terminal), + admission: Some(confirmation_tx), + }, + ) + .await + ); + assert_eq!(next_test_frame(&mut server).await[0], "EVENT"); + tokio::time::pause(); + + let ok_false = Message::Text( + serde_json::to_string(&json!([ + "OK", + terminal_id, false, - ), - ( - "WebSocket(WriteBufferFull): unreachable at connect, fail-safe transient", - ws(WsError::WriteBufferFull(Box::new( - tokio_tungstenite::tungstenite::Message::Text("x".into()), - ))), + "rate-limited: quota exceeded; retry in 60s" + ])) + .expect("serialize OK false") + .into(), + ); + assert!( + handle_ws_message( + ok_false, + &mut ws, + &event_tx, + &observer_control_tx, + &mut state, + &keys, + "ws://localhost:3000", + &keys.public_key().to_hex(), + None, + ) + .await + ); + let gate = state.rate_limit_gate.expect("sixty-second gate"); + let remaining = gate.duration_since(tokio::time::Instant::now()); + assert!( + remaining >= Duration::from_secs(60), + "the client must never retry before the relay's honest reset horizon" + ); + assert!( + remaining < TERMINAL_OBSERVER_ACK_TIMEOUT, + "the original confirmation owner must outlive the full retry gate" + ); + + tokio::time::advance(Duration::from_secs(59)).await; + assert!(state.check_rate_gate().is_some()); + assert_eq!( + drain_gated_observer_pending(&mut ws, &mut state, 1).await, + GatedObserverDrainOutcome::Sent(0), + "no terminal resend is allowed before the advertised reset" + ); + + tokio::time::advance(Duration::from_secs(14)).await; + assert!(state.check_rate_gate().is_none()); + tokio::time::resume(); + assert_eq!( + drain_gated_observer_pending(&mut ws, &mut state, 1).await, + GatedObserverDrainOutcome::Sent(1) + ); + tokio::task::yield_now().await; + assert_eq!(next_test_frame(&mut server).await[0], "EVENT"); + + state.acknowledge_observer_frame(&terminal_id); + assert_eq!( + confirmation_rx.try_recv(), + Ok(Ok(())), + "the delayed accepted OK must resolve the original terminal publication" + ); + } + + #[tokio::test] + async fn terminal_ok_false_does_not_ack_or_retry_observer_frame() { + let (mut ws, _server) = test_ws_pair().await; + let (event_tx, _event_rx) = mpsc::channel(1); + let (observer_control_tx, _observer_control_rx) = mpsc::channel(1); + let mut state = BgState::new(); + let keys = Keys::generate(); + let denied = make_observer_frame(&keys); + state.track_observer_in_flight(Box::new(denied.clone())); + + let ok_false = Message::Text( + serde_json::to_string(&json!([ + "OK", + denied.id.to_hex(), false, - ), - ]; + "blocked: policy denied" + ])) + .expect("serialize OK false") + .into(), + ); + assert!( + handle_ws_message( + ok_false, + &mut ws, + &event_tx, + &observer_control_tx, + &mut state, + &keys, + "ws://localhost:3000", + &keys.public_key().to_hex(), + None, + ) + .await + ); - for (label, err, want_terminal) in cases { - assert_eq!( - is_terminal_connect_error(&err), - want_terminal, - "{label}: expected terminal={want_terminal}" - ); + assert!( + state.observer_in_flight.is_empty(), + "terminal denial is retired instead of occupying the bounded ACK window forever" + ); + assert_eq!( + state.observer_terminal_denied, 1, + "terminal denial must leave an auditable counter" + ); + assert!( + state.gated_observer_pending.is_empty(), + "terminal denial must not enter an unbounded retry loop" + ); + assert!(state.check_rate_gate().is_none()); + } + + /// The parked-frame queue is bounded: overflow evicts the oldest frame and + /// counts it; the drain resets the counter after logging the summary. + #[tokio::test] + async fn gated_observer_queue_drops_oldest_on_overflow() { + let mut state = BgState::new(); + let keys = Keys::generate(); + let first = make_observer_frame(&keys); + state.park_gated_observer_frame(Box::new(first.clone())); + for _ in 1..GATED_OBSERVER_QUEUE_CAP { + state.park_gated_observer_frame(Box::new(make_observer_frame(&keys))); } + assert_eq!(state.gated_observer_pending.len(), GATED_OBSERVER_QUEUE_CAP); + assert_eq!(state.gated_observer_dropped, 0); + + let overflow = make_observer_frame(&keys); + state.park_gated_observer_frame(Box::new(overflow.clone())); + assert_eq!( + state.gated_observer_pending.len(), + GATED_OBSERVER_QUEUE_CAP, + "queue must stay bounded" + ); + assert_eq!(state.gated_observer_dropped, 1, "loss must be counted"); + assert!( + !state + .gated_observer_pending + .iter() + .any(|e| e.id == first.id), + "oldest frame must be the one evicted" + ); + assert_eq!( + state.gated_observer_pending.back().map(|e| e.id), + Some(overflow.id), + "newest frame must be retained" + ); } - /// A literal `https://…` URL through production `do_connect()` must fail - /// fast as terminal — the relay endpoint is a plain HTTPS server, not a - /// WebSocket endpoint, and tungstenite returns `Error::Http` (non-101 - /// response) or `Error::Url(UnsupportedUrlScheme)` depending on how far - /// the handshake gets. Either way it must not be retried. #[tokio::test] - async fn do_connect_wrong_scheme_is_terminal() { - let keys = nostr::Keys::generate(); - let err = do_connect("https://example.com", &keys, None) - .await - .unwrap_err(); + async fn protected_observer_queue_rejects_newest_without_evicting_terminal_proof() { + let mut state = BgState::new(); + let keys = Keys::generate(); + let first = make_priority_observer_frame(&keys); + assert!(state.park_gated_observer_frame(Box::new(first.clone()))); + for _ in 1..GATED_OBSERVER_PRIORITY_CAP { + assert!(state.park_gated_observer_frame(Box::new(make_priority_observer_frame(&keys)))); + } + + let overflow = make_priority_observer_frame(&keys); assert!( - is_terminal_connect_error(&err), - "wrong-scheme URL should be terminal, got: {err}" + !state.park_gated_observer_frame(Box::new(overflow.clone())), + "a saturated protected lane must reject admission" + ); + assert_eq!( + state.gated_observer_priority_pending.len(), + GATED_OBSERVER_PRIORITY_CAP + ); + assert_eq!(state.gated_observer_dropped, 1); + assert_eq!( + state + .gated_observer_priority_pending + .front() + .map(|event| event.id), + Some(first.id), + "an already-admitted terminal proof must never be evicted" + ); + assert!( + !state + .gated_observer_priority_pending + .iter() + .any(|event| event.id == overflow.id), + "the rejected newest frame must not displace admitted proof" ); } - /// A transient failure (e.g. connection dropped mid-handshake on a spotty - /// link) must be retried and can still succeed once the link recovers. - #[tokio::test(start_paused = true)] - async fn retry_initial_connect_retries_transient_failure_then_succeeds() { - use std::sync::atomic::{AtomicUsize, Ordering}; + #[tokio::test] + async fn terminal_publisher_reports_rejected_priority_admission() { + let (cmd_tx, mut cmd_rx) = mpsc::channel(1); + let publisher = RelayEventPublisher { cmd_tx }; + let event = make_priority_observer_frame(&Keys::generate()); + let publish = tokio::spawn(async move { publisher.publish_terminal_event(event).await }); + + let RelayCommand::PublishEvent { + admission: Some(admission), + .. + } = cmd_rx.recv().await.expect("terminal publish command") + else { + panic!("terminal publish must request confirmed priority admission"); + }; + admission + .send(Err("terminal observer priority state is full".into())) + .expect("publisher must still be waiting for admission"); - let attempts = AtomicUsize::new(0); - let result: Result<&'static str, RelayError> = retry_initial_connect(|| { - let n = attempts.fetch_add(1, Ordering::SeqCst); - async move { - if n < 2 { - Err(RelayError::ConnectionClosed) - } else { - Ok("connected") - } - } - }) - .await; + assert!(matches!( + publish.await.expect("publisher task"), + Err(RelayError::Http(message)) + if message == "terminal observer priority state is full" + )); + } - assert_eq!(result.unwrap(), "connected"); + #[tokio::test] + async fn terminal_publisher_rejects_non_terminal_shape_before_queue_admission() { + let (cmd_tx, mut cmd_rx) = mpsc::channel(4); + let publisher = RelayEventPublisher { cmd_tx }; + let ordinary = make_observer_frame(&Keys::generate()); + + assert!(matches!( + publisher.publish_terminal_event(ordinary).await, + Err(RelayError::Http(message)) + if message == "terminal observer event must be kind 24200 with exactly one control-result priority marker" + )); + assert!( + cmd_rx.try_recv().is_err(), + "invalid terminal shape must never enter the relay command queue" + ); + + let terminal = make_priority_observer_frame(&Keys::generate()); + let mut wrong_kind_value = serde_json::to_value(terminal).expect("serialize terminal"); + wrong_kind_value["kind"] = json!(KIND_TYPING_INDICATOR); + let wrong_kind: Event = + serde_json::from_value(wrong_kind_value).expect("deserialize wrong-kind event"); + assert!(matches!( + publisher.publish_terminal_event(wrong_kind).await, + Err(RelayError::Http(message)) + if message == "terminal observer event must be kind 24200 with exactly one control-result priority marker" + )); + assert!( + cmd_rx.try_recv().is_err(), + "a priority marker on the wrong event kind must never reach a socket" + ); + } + + #[tokio::test] + async fn terminal_publisher_rejects_duplicate_priority_marker_before_queue_admission() { + let (cmd_tx, mut cmd_rx) = mpsc::channel(4); + let publisher = RelayEventPublisher { cmd_tx }; + let terminal = make_priority_observer_frame(&Keys::generate()); + let mut value = serde_json::to_value(terminal).expect("serialize terminal"); + value["tags"] + .as_array_mut() + .expect("terminal tags") + .push(json!([ + OBSERVER_PRIORITY_TAG, + OBSERVER_CONTROL_RESULT_PRIORITY + ])); + let duplicate_marker: Event = + serde_json::from_value(value).expect("deserialize duplicate marker event"); + + assert!(matches!( + publisher.publish_terminal_event(duplicate_marker).await, + Err(RelayError::Http(message)) + if message == "terminal observer event must be kind 24200 with exactly one control-result priority marker" + )); + assert!(cmd_rx.try_recv().is_err()); + } + + #[tokio::test] + async fn terminal_publisher_rejects_extended_priority_marker_before_queue_admission() { + let (cmd_tx, mut cmd_rx) = mpsc::channel(4); + let publisher = RelayEventPublisher { cmd_tx }; + let terminal = make_priority_observer_frame(&Keys::generate()); + let mut value = serde_json::to_value(terminal).expect("serialize terminal"); + let priority = value["tags"] + .as_array_mut() + .expect("terminal tags") + .iter_mut() + .find(|tag| tag[0] == OBSERVER_PRIORITY_TAG) + .expect("priority marker"); + priority + .as_array_mut() + .expect("priority tag") + .push(json!("extra")); + let extended_marker: Event = + serde_json::from_value(value).expect("deserialize extended marker event"); + + assert!(matches!( + publisher.publish_terminal_event(extended_marker).await, + Err(RelayError::Http(message)) + if message == "terminal observer event must be kind 24200 with exactly one control-result priority marker" + )); + assert!( + cmd_rx.try_recv().is_err(), + "an extended marker must not weaken the exact terminal frame contract" + ); + } + + #[tokio::test] + async fn terminal_confirmation_waits_for_relay_ok_not_local_socket_write() { + let (mut client, mut server) = test_ws_pair().await; + let mut state = BgState::new(); + let event = make_priority_observer_frame(&Keys::generate()); + let event_id = event.id.to_hex(); + let (confirmation_tx, mut confirmation_rx) = oneshot::channel(); + + assert!( + execute_connected_command( + &mut client, + &mut state, + "agent-pubkey", + RelayCommand::PublishEvent { + event: Box::new(event), + admission: Some(confirmation_tx), + }, + ) + .await + ); + let frame = next_test_frame(&mut server).await; + assert_eq!(frame[0], "EVENT"); + assert!( + confirmation_rx.try_recv().is_err(), + "a local socket write is not terminal delivery proof" + ); + + state.acknowledge_observer_frame(&event_id); assert_eq!( - attempts.load(Ordering::SeqCst), - 3, - "should succeed on the 3rd attempt (2 transient failures + 1 success)" + confirmation_rx.await.expect("relay OK confirmation"), + Ok(()) ); } - /// A terminal error (bad auth, bad config) must not be retried — the - /// same call would fail identically every time, so retrying just delays - /// surfacing a real problem to the caller. - #[tokio::test(start_paused = true)] - async fn retry_initial_connect_does_not_retry_terminal_error() { - use std::sync::atomic::{AtomicUsize, Ordering}; + #[tokio::test] + async fn terminal_confirmation_timeout_cancels_exact_event_and_next_result_progresses() { + let (mut client, mut server) = test_ws_pair().await; + let mut state = BgState::new(); + let (cmd_tx, mut cmd_rx) = mpsc::channel(4); + let publisher = RelayEventPublisher { cmd_tx }; + let test_ack_timeout = Duration::from_millis(250); - let attempts = AtomicUsize::new(0); - let result: Result<(), RelayError> = retry_initial_connect(|| { - attempts.fetch_add(1, Ordering::SeqCst); - async { Err(RelayError::AuthFailed("invalid: bad signature".into())) } - }) - .await; + let first = make_priority_observer_frame(&Keys::generate()); + let first_id = first.id.to_hex(); + let first_publish = tokio::spawn({ + let publisher = publisher.clone(); + async move { + publisher + .publish_terminal_event_with_ack_timeout(first, test_ack_timeout) + .await + } + }); + let first_command = cmd_rx.recv().await.expect("first terminal publish command"); + assert!( + execute_connected_command(&mut client, &mut state, "agent-pubkey", first_command,) + .await + ); + let first_frame = next_test_frame(&mut server).await; + assert_eq!(first_frame[0], "EVENT"); - assert!(matches!(result, Err(RelayError::AuthFailed(_)))); - assert_eq!( - attempts.load(Ordering::SeqCst), - 1, - "a terminal error must fail on the first attempt with no retries" + assert!(matches!( + first_publish.await.expect("first publisher task"), + Err(RelayError::Timeout) + )); + + let second = make_priority_observer_frame(&Keys::generate()); + let second_id = second.id.to_hex(); + let second_publish = tokio::spawn({ + let publisher = publisher.clone(); + async move { publisher.publish_terminal_event(second).await } + }); + let second_command = cmd_rx + .recv() + .await + .expect("second terminal publish command"); + assert!( + execute_connected_command(&mut client, &mut state, "agent-pubkey", second_command,) + .await + ); + assert!( + !state.contains_observer_priority_event(&first_id), + "the next command must prune only the expired attempt before admitting later work" + ); + let second_frame = next_test_frame(&mut server).await; + assert_eq!(second_frame[0], "EVENT"); + + state.acknowledge_observer_frame(&second_id); + assert!( + second_publish.await.expect("second publisher task").is_ok(), + "a later terminal result must progress after exact timeout cleanup" ); } - /// A relay-side dependency fault (NIP-01 `error:` prefix) is transient — - /// the relay is failing closed on itself, not rejecting this identity — - /// so it must be retried rather than surfaced immediately like a real - /// auth rejection. - #[tokio::test(start_paused = true)] - async fn retry_initial_connect_retries_relay_dependency_fault() { - use std::sync::atomic::{AtomicUsize, Ordering}; + #[tokio::test] + async fn expired_queued_terminal_publish_never_reaches_reconnected_socket() { + let (mut client, mut server) = test_ws_pair().await; + let mut state = BgState::new(); + let (cmd_tx, mut cmd_rx) = mpsc::channel(4); + let publisher = RelayEventPublisher { cmd_tx }; + let event = make_priority_observer_frame(&Keys::generate()); + let event_id = event.id.to_hex(); - let attempts = AtomicUsize::new(0); - let result: Result<&'static str, RelayError> = retry_initial_connect(|| { - let n = attempts.fetch_add(1, Ordering::SeqCst); - async move { - if n < 1 { - Err(RelayError::AuthFailed( - "error: internal error checking restriction state".into(), - )) - } else { - Ok("connected") - } - } - }) - .await; + let publish_result = publisher + .publish_terminal_event_with_ack_timeout(event, Duration::from_millis(100)) + .await; + assert!(matches!(publish_result, Err(RelayError::Timeout))); - assert_eq!(result.unwrap(), "connected"); + let expired_publish = cmd_rx.recv().await.expect("expired publish remains queued"); + assert!( + execute_connected_command(&mut client, &mut state, "agent-pubkey", expired_publish,) + .await + ); + assert!( + cmd_rx.try_recv().is_err(), + "timeout cleanup must not enqueue an ambiguous event-ID cancellation" + ); + assert!( + timeout(Duration::from_millis(100), server.next()) + .await + .is_err(), + "an expired queued terminal result must not reach a reconnected relay" + ); + assert!( + !state.contains_observer_priority_event(&event_id), + "expired reconnect work must leave no pending, in-flight, or confirmation owner" + ); + + let later = make_observer_frame(&Keys::generate()); + let later_id = later.id; + assert!( + execute_connected_command( + &mut client, + &mut state, + "agent-pubkey", + RelayCommand::PublishEvent { + event: Box::new(later), + admission: None, + }, + ) + .await, + "clean pre-send cancellation must leave the socket reusable" + ); + let later_frame = next_test_frame(&mut server).await; assert_eq!( - attempts.load(Ordering::SeqCst), - 2, - "a relay dependency fault must be retried, not surfaced immediately" + later_frame[1]["id"], + later_id.to_hex(), + "later traffic must progress on the clean socket" ); } - /// Once every attempt (1 initial + N backoff retries) is exhausted, the - /// last transient error is returned rather than retrying forever — a - /// dead relay must not hang agent startup indefinitely. - #[tokio::test(start_paused = true)] - async fn retry_initial_connect_exhausts_and_returns_last_error() { - use std::sync::atomic::{AtomicUsize, Ordering}; + #[tokio::test] + async fn timed_out_queued_duplicate_preserves_original_confirmation_owner() { + let (mut client, mut server) = test_ws_pair().await; + let mut state = BgState::new(); + let event = make_priority_observer_frame(&Keys::generate()); + let event_id = event.id.to_hex(); + let (original_tx, mut original_rx) = oneshot::channel(); + assert!( + execute_connected_command( + &mut client, + &mut state, + "agent-pubkey", + RelayCommand::PublishEvent { + event: Box::new(event.clone()), + admission: Some(original_tx), + }, + ) + .await + ); + let original_frame = next_test_frame(&mut server).await; + assert_eq!(original_frame[1]["id"], event_id); - let attempts = AtomicUsize::new(0); - let result: Result<(), RelayError> = retry_initial_connect(|| { - attempts.fetch_add(1, Ordering::SeqCst); - async { Err(RelayError::Timeout) } - }) - .await; + let (cmd_tx, mut cmd_rx) = mpsc::channel(4); + let duplicate_publisher = RelayEventPublisher { cmd_tx }; + let duplicate_result = duplicate_publisher + .publish_terminal_event_with_ack_timeout(event, Duration::from_millis(100)) + .await; + assert!(matches!(duplicate_result, Err(RelayError::Timeout))); + while let Ok(command) = cmd_rx.try_recv() { + assert!( + execute_connected_command(&mut client, &mut state, "agent-pubkey", command,).await + ); + } assert!( - matches!(result, Err(RelayError::Timeout)), - "must surface the last attempt's error, not a generic one" + timeout(Duration::from_millis(100), server.next()) + .await + .is_err(), + "the expired duplicate attempt must never write a second frame" + ); + assert!( + state.contains_observer_priority_event(&event_id), + "duplicate timeout cleanup must preserve the original confirmation owner" ); + assert!(matches!( + original_rx.try_recv(), + Err(oneshot::error::TryRecvError::Empty) + )); + + state.acknowledge_observer_frame(&event_id); assert_eq!( - attempts.load(Ordering::SeqCst), - STARTUP_CONNECT_BACKOFFS.len() + 1, - "must attempt exactly once plus one retry per backoff entry" + original_rx + .await + .expect("original confirmation remains live"), + Ok(()) ); } - /// Backoff sleeps must actually elapse (not be skipped) — this pins the - /// bounded-but-real-delay contract using `tokio::time::pause` so the - /// test itself stays fast (virtual time, not wall-clock sleeps). - #[tokio::test(start_paused = true)] - async fn retry_initial_connect_sleeps_between_attempts() { - use std::sync::atomic::{AtomicUsize, Ordering}; + struct StallAfterStartSendSink { + started: Option>, + buffered: Vec, + flush_allowed: std::sync::Arc, + flushed: std::sync::Arc, + } - let attempts = AtomicUsize::new(0); - let call = retry_initial_connect(|| { - let n = attempts.fetch_add(1, Ordering::SeqCst); - async move { - if n < 1 { - Err(RelayError::ConnectionClosed) - } else { - Ok(()) + impl futures_util::Sink for StallAfterStartSendSink { + type Error = std::convert::Infallible; + + fn poll_ready( + self: std::pin::Pin<&mut Self>, + _cx: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + std::task::Poll::Ready(Ok(())) + } + + fn start_send( + mut self: std::pin::Pin<&mut Self>, + item: Message, + ) -> Result<(), Self::Error> { + self.buffered.push(item); + if let Some(started) = self.started.take() { + let _ = started.send(()); + } + Ok(()) + } + + fn poll_flush( + mut self: std::pin::Pin<&mut Self>, + _cx: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + if self.flush_allowed.load(std::sync::atomic::Ordering::SeqCst) { + if !self.buffered.is_empty() { + self.flushed + .store(true, std::sync::atomic::Ordering::SeqCst); + self.buffered.clear(); } + std::task::Poll::Ready(Ok(())) + } else { + std::task::Poll::Pending } - }); - tokio::pin!(call); + } - // Before the first backoff elapses, the retry must still be pending - // (i.e. it actually slept rather than immediately retrying). - tokio::select! { - biased; - _ = tokio::time::sleep(Duration::from_millis(1)) => {} - _ = &mut call => panic!("must not resolve before the backoff sleep elapses"), + fn poll_close( + self: std::pin::Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + self.poll_flush(cx) } - assert_eq!(attempts.load(Ordering::SeqCst), 1); + } - // Advancing past the (jittered, ≤1.2x) first backoff lets it proceed. - let result = call.await; - assert!(result.is_ok()); - assert_eq!(attempts.load(Ordering::SeqCst), 2); + struct BufferThenStartSendErrorSink { + buffered: Vec, } - // ── Rate-limit gate, pacing, backoff reset, DNS ────────────────────────── + impl futures_util::Sink for BufferThenStartSendErrorSink { + type Error = std::io::Error; - /// parse_rate_limit_retry_secs: full hint extracts the N from "retry in Ns". - #[test] - fn parse_rate_limit_retry_secs_with_hint() { - assert_eq!( - parse_rate_limit_retry_secs("rate-limited: quota exceeded; retry in 12s"), - Some(12) - ); + fn poll_ready( + self: std::pin::Pin<&mut Self>, + _cx: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + std::task::Poll::Ready(Ok(())) + } + + fn start_send( + mut self: std::pin::Pin<&mut Self>, + item: Message, + ) -> Result<(), Self::Error> { + self.buffered.push(item); + Err(std::io::Error::other( + "write failed after buffering the frame", + )) + } + + fn poll_flush( + self: std::pin::Pin<&mut Self>, + _cx: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + std::task::Poll::Ready(Ok(())) + } + + fn poll_close( + self: std::pin::Pin<&mut Self>, + _cx: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + std::task::Poll::Ready(Ok(())) + } } - /// parse_rate_limit_retry_secs: message without a hint returns None. - #[test] - fn parse_rate_limit_retry_secs_missing_hint() { + #[tokio::test] + async fn start_send_error_after_buffering_taints_socket() { + let event = make_priority_observer_frame(&Keys::generate()); + let mut socket = BufferThenStartSendErrorSink { + buffered: Vec::new(), + }; + + let outcome = send_publish_event_frame_until_cancelled(&mut socket, &event, None).await; + assert_eq!( - parse_rate_limit_retry_secs("rate-limited: too many concurrent requests"), - None + outcome, + PublishEventFrameOutcome::FailedSocketTainted, + "invoking start_send must conservatively taint the socket even when it returns an error" + ); + assert!(outcome.socket_tainted()); + assert_eq!( + socket.buffered.len(), + 1, + "a start_send error may leave the EVENT buffered for later traffic" ); } - /// parse_rate_limit_retry_secs: explicit zero value is returned as Some(0). - #[test] - fn parse_rate_limit_retry_secs_zero() { + #[tokio::test] + async fn cancellation_after_start_send_retires_socket_before_later_traffic() { + let event = make_priority_observer_frame(&Keys::generate()); + let (mut confirmation_tx, confirmation_rx) = oneshot::channel(); + let (started_tx, started_rx) = oneshot::channel(); + let old_flush_allowed = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false)); + let old_flushed = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false)); + let mut old_socket = StallAfterStartSendSink { + started: Some(started_tx), + buffered: Vec::new(), + flush_allowed: old_flush_allowed, + flushed: old_flushed.clone(), + }; + let close_attempt = tokio::spawn(async move { + started_rx.await.expect("EVENT reaches start_send"); + drop(confirmation_rx); + }); + + let outcome = send_publish_event_frame_until_cancelled( + &mut old_socket, + &event, + Some(&mut confirmation_tx), + ) + .await; + close_attempt.await.expect("close attempt task"); + + assert_eq!(outcome, PublishEventFrameOutcome::SocketTainted); + assert!(outcome.socket_tainted()); assert_eq!( - parse_rate_limit_retry_secs("rate-limited: quota exceeded; retry in 0s"), - Some(0) + old_socket.buffered.len(), + 1, + "the cancelled send may still be buffered after start_send" + ); + + // The supervisor must retire the tainted transport. Later traffic goes + // only to its replacement, so it cannot flush the cancelled EVENT. + drop(old_socket); + let replacement_flushed = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false)); + let mut replacement = StallAfterStartSendSink { + started: None, + buffered: Vec::new(), + flush_allowed: std::sync::Arc::new(std::sync::atomic::AtomicBool::new(true)), + flushed: replacement_flushed.clone(), + }; + replacement + .send(Message::Ping(Vec::new().into())) + .await + .expect("later traffic uses replacement transport"); + + assert!( + replacement_flushed.load(std::sync::atomic::Ordering::SeqCst), + "the later frame must be sent on the replacement transport" + ); + assert!( + !old_flushed.load(std::sync::atomic::Ordering::SeqCst), + "later traffic must never flush the cancelled EVENT" ); } - /// parse_rate_limit_retry_secs: garbage input returns None. #[test] - fn parse_rate_limit_retry_secs_garbage() { - assert_eq!( - parse_rate_limit_retry_secs("not a rate limit message"), - None + fn receiver_close_purges_only_its_exact_terminal_admission() { + let mut state = BgState::new(); + let keys = Keys::generate(); + let cancelled = make_priority_observer_frame(&keys); + let cancelled_id = cancelled.id.to_hex(); + let survivor = make_priority_observer_frame(&keys); + let survivor_id = survivor.id.to_hex(); + let (cancelled_tx, cancelled_rx) = oneshot::channel(); + let (survivor_tx, mut survivor_rx) = oneshot::channel(); + + assert!(state.park_gated_observer_frame(Box::new(cancelled))); + state.record_observer_confirmation(cancelled_id.clone(), true, Some(cancelled_tx), true); + assert!(state.park_gated_observer_frame(Box::new(survivor))); + state.record_observer_confirmation(survivor_id.clone(), true, Some(survivor_tx), true); + + drop(cancelled_rx); + state.prune_cancelled_terminal_admissions(); + + assert!(!state.contains_observer_priority_event(&cancelled_id)); + assert!( + state.contains_observer_priority_event(&survivor_id), + "receiver-close cleanup must preserve every other terminal owner" ); + assert!(matches!( + survivor_rx.try_recv(), + Err(oneshot::error::TryRecvError::Empty) + )); } - /// set_rate_limit_gate arms the gate with jittered expiry from the hint. - /// check_rate_gate returns Some while active and lazily clears on expiry. - #[tokio::test(start_paused = true)] - async fn rate_limit_gate_set_and_expiry() { + #[tokio::test] + async fn full_command_channel_cannot_extend_terminal_publish_deadline() { + let (cmd_tx, _cmd_rx) = mpsc::channel(1); + cmd_tx + .send(RelayCommand::Reconnect) + .await + .expect("seed full command channel"); + let publisher = RelayEventPublisher { cmd_tx }; + let event = make_priority_observer_frame(&Keys::generate()); + + let bounded = timeout( + Duration::from_secs(1), + publisher.publish_terminal_event_with_ack_timeout(event, Duration::from_millis(100)), + ) + .await + .expect("publisher must return without command-channel capacity"); + assert!(matches!(bounded, Err(RelayError::Timeout))); + } + + #[tokio::test] + async fn duplicate_terminal_event_rejects_second_confirmation_before_socket_write() { + let (mut client, mut server) = test_ws_pair().await; let mut state = BgState::new(); + let keys = Keys::generate(); + let terminal = make_priority_observer_frame(&keys); + let terminal_id = terminal.id.to_hex(); + let (first_tx, first_rx) = oneshot::channel(); assert!( - state.check_rate_gate().is_none(), - "gate must start inactive" + execute_connected_command( + &mut client, + &mut state, + "agent-pubkey", + RelayCommand::PublishEvent { + event: Box::new(terminal.clone()), + admission: Some(first_tx), + }, + ) + .await ); + let first_frame = next_test_frame(&mut server).await; + assert_eq!(first_frame[1]["id"], terminal_id); - // Arm with a 5 s hint. - state.set_rate_limit_gate(5); + let (second_tx, second_rx) = oneshot::channel(); assert!( - state.check_rate_gate().is_some(), - "gate must be active immediately after arming" + execute_connected_command( + &mut client, + &mut state, + "agent-pubkey", + RelayCommand::PublishEvent { + event: Box::new(terminal), + admission: Some(second_tx), + }, + ) + .await + ); + assert_eq!( + second_rx.await.expect("duplicate rejection result"), + Err("terminal observer event is already awaiting relay confirmation".to_string()) ); - // Advance virtual time past the max jitter (1.2 × 5 s = 6 s). - tokio::time::advance(Duration::from_secs(7)).await; - - assert!( - state.check_rate_gate().is_none(), - "gate must have expired after 7s" + state.acknowledge_observer_frame(&terminal_id); + assert_eq!( + first_rx + .await + .expect("first relay confirmation remains live"), + Ok(()) ); + + let next = make_observer_frame(&keys); + let next_id = next.id; assert!( - state.rate_limit_gate.is_none(), - "check_rate_gate must lazily clear the field on expiry" + execute_connected_command( + &mut client, + &mut state, + "agent-pubkey", + RelayCommand::PublishEvent { + event: Box::new(next), + admission: None, + }, + ) + .await ); - } - - /// set_rate_limit_gate takes the max of overlapping deadlines. - #[tokio::test(start_paused = true)] - async fn rate_limit_gate_extends_to_max() { - let mut state = BgState::new(); - - // Arm with a long hint first. - state.set_rate_limit_gate(30); - let first_deadline = state.rate_limit_gate.unwrap(); - - // A shorter subsequent hint must NOT shorten the existing gate. - state.set_rate_limit_gate(1); - let second_deadline = state.rate_limit_gate.unwrap(); - + let next_frame = next_test_frame(&mut server).await; assert_eq!( - first_deadline, second_deadline, - "shorter hint must not overwrite a later existing deadline" + next_frame[1]["id"], + next_id.to_hex(), + "duplicate rejection must not emit a second terminal EVENT" ); } - /// Build a signed observer telemetry frame (kind 24200) for gate tests. - fn make_observer_frame(keys: &Keys) -> Event { - let recipient = Keys::generate(); - let encrypted = buzz_core::observer::encrypt_observer_payload( - keys, - &recipient.public_key(), - &json!({"type": "test"}), - ) - .expect("encrypt test observer payload"); - buzz_sdk::build_agent_observer_frame( - &recipient.public_key().to_hex(), - &keys.public_key().to_hex(), - "telemetry", - &encrypted, - ) - .expect("build test observer frame") - .sign_with_keys(keys) - .expect("sign test observer frame") - } - - /// While the rate-limit gate is armed, an observer frame (kind 24200) is - /// parked — not silently dropped — and delivered by the drain once the - /// gate clears. A typing indicator in the same window stays dropped. #[tokio::test] - async fn gated_observer_frame_is_parked_then_drained_not_dropped() { - let (mut client, mut server) = test_ws_pair().await; + async fn disconnected_duplicate_terminal_event_preserves_first_waiter_and_queue_entry() { let mut state = BgState::new(); - let keys = Keys::generate(); - state.rate_limit_gate = Some(tokio::time::Instant::now() + Duration::from_millis(150)); - - // Observer frame while gated: parked, nothing on the wire. - let observer_frame = make_observer_frame(&keys); - let ok = execute_connected_command( - &mut client, + let terminal = make_priority_observer_frame(&Keys::generate()); + let terminal_id = terminal.id.to_hex(); + let (first_tx, mut first_rx) = oneshot::channel(); + apply_command_to_state( &mut state, - "agent-pubkey", RelayCommand::PublishEvent { - event: Box::new(observer_frame.clone()), + event: Box::new(terminal.clone()), + admission: Some(first_tx), }, - ) - .await; - assert!(ok); - assert_eq!( - state.gated_observer_pending.len(), - 1, - "observer frame must be parked while gated" ); - // Typing indicator while gated: still dropped, not parked. - let typing = EventBuilder::new(Kind::Custom(KIND_TYPING_INDICATOR as u16), "") - .tags([Tag::parse(["h", &Uuid::new_v4().to_string()]).unwrap()]) - .sign_with_keys(&keys) - .expect("sign typing indicator"); - let ok = execute_connected_command( - &mut client, + let (second_tx, second_rx) = oneshot::channel(); + apply_command_to_state( &mut state, - "agent-pubkey", RelayCommand::PublishEvent { - event: Box::new(typing), + event: Box::new(terminal.clone()), + admission: Some(second_tx), }, - ) - .await; - assert!(ok); - assert_eq!( - state.gated_observer_pending.len(), - 1, - "typing indicators must not be parked" ); - assert!( - timeout(Duration::from_millis(50), server.next()) - .await - .is_err(), - "nothing may reach the wire while the gate is armed" + assert_eq!( + second_rx.await.expect("disconnected duplicate rejection"), + Err("terminal observer event is already awaiting relay confirmation".to_string()) ); - // Gate expires — the drain delivers the parked frame. - tokio::time::sleep(Duration::from_millis(160)).await; + let (replay_tx, replay_rx) = oneshot::channel(); + retain_failed_command_intent( + &mut state, + RelayCommand::PublishEvent { + event: Box::new(terminal), + admission: Some(replay_tx), + }, + ); assert_eq!( - drain_gated_observer_pending(&mut client, &mut state, 1).await, - 1 + replay_rx.await.expect("failed-replay duplicate rejection"), + Err("terminal observer event is already awaiting relay confirmation".to_string()) ); - assert!(state.gated_observer_pending.is_empty()); - let frame = next_test_frame(&mut server).await; - assert_eq!(frame[0], "EVENT"); - assert_eq!(frame[1]["id"], observer_frame.id.to_hex()); + assert_eq!(state.gated_observer_priority_pending.len(), 1); assert_eq!( - frame[1]["kind"], - u64::from(KIND_AGENT_OBSERVER_FRAME), - "delivered frame must be the parked observer frame" + state + .gated_observer_priority_pending + .front() + .map(|event| event.id.to_hex()), + Some(terminal_id.clone()) + ); + assert!( + state + .observer_priority_confirmations + .contains_key(&terminal_id), + "the first publisher must retain confirmation ownership" + ); + assert!( + matches!( + first_rx.try_recv(), + Err(oneshot::error::TryRecvError::Empty) + ), + "the first publisher must remain pending instead of observing replacement closure" ); } - /// Observer frames arriving while earlier parked frames are still queued - /// are appended behind them (order preserved), even if the gate has - /// already expired. #[tokio::test] - async fn observer_frames_queue_behind_parked_backlog_in_order() { + async fn saturated_in_flight_priority_window_rejects_before_socket_write() { let (mut client, mut server) = test_ws_pair().await; let mut state = BgState::new(); let keys = Keys::generate(); - state.rate_limit_gate = Some(tokio::time::Instant::now() + Duration::from_millis(50)); + for _ in 0..GATED_OBSERVER_PRIORITY_CAP { + assert!(state.track_observer_in_flight(Box::new(make_priority_observer_frame(&keys)))); + } - let first = make_observer_frame(&keys); - let second = make_observer_frame(&keys); - for event in [&first, &second] { - let ok = execute_connected_command( + let rejected = make_priority_observer_frame(&keys); + let rejected_id = rejected.id; + let (confirmation_tx, confirmation_rx) = oneshot::channel(); + assert!( + execute_connected_command( &mut client, &mut state, "agent-pubkey", RelayCommand::PublishEvent { - event: Box::new(event.clone()), + event: Box::new(rejected), + admission: Some(confirmation_tx), }, ) - .await; - assert!(ok); - } - assert_eq!(state.gated_observer_pending.len(), 2); + .await + ); - // Gate expires but the backlog is not drained yet — a third frame must - // queue behind it rather than jumping ahead on the wire. - tokio::time::sleep(Duration::from_millis(60)).await; - let third = make_observer_frame(&keys); - let ok = execute_connected_command( - &mut client, - &mut state, - "agent-pubkey", - RelayCommand::PublishEvent { - event: Box::new(third.clone()), - }, - ) - .await; - assert!(ok); assert_eq!( - state.gated_observer_pending.len(), - 3, - "frame must queue behind undrained backlog to preserve order" + confirmation_rx.await.expect("priority rejection result"), + Err("terminal observer priority state is full".to_string()) + ); + assert_eq!( + state.observer_priority_in_flight.len(), + GATED_OBSERVER_PRIORITY_CAP + ); + assert!( + !state + .observer_priority_in_flight + .iter() + .any(|event| event.id == rejected_id), + "a rejected terminal result must not enter the acknowledgment window" ); - for expected in [&first, &second, &third] { - assert_eq!( - drain_gated_observer_pending(&mut client, &mut state, 1).await, - 1 - ); - let frame = next_test_frame(&mut server).await; - assert_eq!(frame[1]["id"], expected.id.to_hex(), "order preserved"); - } - assert!(state.gated_observer_pending.is_empty()); - } - - #[test] - fn observer_notice_requeues_unacknowledged_frames_and_ok_retires_them() { - let mut state = BgState::new(); - let keys = Keys::generate(); - let accepted = make_observer_frame(&keys); - let rejected = make_observer_frame(&keys); - let later = make_observer_frame(&keys); - - state.track_observer_in_flight(Box::new(accepted.clone())); - state.track_observer_in_flight(Box::new(rejected.clone())); - state.acknowledge_observer_frame(&accepted.id.to_hex()); - state.park_gated_observer_frame(Box::new(later.clone())); - state.requeue_observer_in_flight(); - - let ids: Vec<_> = state - .gated_observer_pending - .iter() - .map(|event| event.id) - .collect(); - assert_eq!(ids, [rejected.id, later.id]); - assert!(state.observer_in_flight.is_empty()); + let next = make_observer_frame(&keys); + let next_id = next.id; + assert!( + execute_connected_command( + &mut client, + &mut state, + "agent-pubkey", + RelayCommand::PublishEvent { + event: Box::new(next), + admission: None, + }, + ) + .await + ); + let frame = next_test_frame(&mut server).await; + assert_eq!( + frame[1]["id"], + next_id.to_hex(), + "the first socket write after rejection must be the later admitted frame" + ); + assert_eq!(state.gated_observer_dropped, 1); } - /// The parked-frame queue is bounded: overflow evicts the oldest frame and - /// counts it; the drain resets the counter after logging the summary. - #[tokio::test] - async fn gated_observer_queue_drops_oldest_on_overflow() { + #[tokio::test(start_paused = true)] + async fn terminal_control_result_survives_telemetry_saturation_and_drains_within_eight_seconds() + { + let (mut client, mut server) = test_ws_pair().await; let mut state = BgState::new(); let keys = Keys::generate(); - let first = make_observer_frame(&keys); - state.park_gated_observer_frame(Box::new(first.clone())); - for _ in 1..GATED_OBSERVER_QUEUE_CAP { + for _ in 0..GATED_OBSERVER_QUEUE_CAP { state.park_gated_observer_frame(Box::new(make_observer_frame(&keys))); } - assert_eq!(state.gated_observer_pending.len(), GATED_OBSERVER_QUEUE_CAP); - assert_eq!(state.gated_observer_dropped, 0); + let terminal = make_priority_observer_frame(&keys); + state.park_gated_observer_frame(Box::new(terminal.clone())); + state.set_rate_limit_gate(5); + let saturated_at = tokio::time::Instant::now(); - let overflow = make_observer_frame(&keys); - state.park_gated_observer_frame(Box::new(overflow.clone())); + assert_eq!(state.gated_observer_pending.len(), GATED_OBSERVER_QUEUE_CAP); assert_eq!( - state.gated_observer_pending.len(), - GATED_OBSERVER_QUEUE_CAP, - "queue must stay bounded" + state + .gated_observer_priority_pending + .front() + .map(|event| event.id), + Some(terminal.id), + "ordinary telemetry saturation cannot consume protected capacity" ); - assert_eq!(state.gated_observer_dropped, 1, "loss must be counted"); - assert!( - !state - .gated_observer_pending - .iter() - .any(|e| e.id == first.id), - "oldest frame must be the one evicted" + + tokio::time::advance(Duration::from_secs(7)).await; + assert_eq!( + drain_gated_observer_pending(&mut client, &mut state, 1).await, + GatedObserverDrainOutcome::Sent(1) ); + let message = server + .next() + .await + .expect("test websocket stays open") + .expect("read protected observer frame"); + let frame: serde_json::Value = + serde_json::from_str(message.to_text().expect("protected frame is text")) + .expect("parse protected observer frame"); assert_eq!( - state.gated_observer_pending.back().map(|e| e.id), - Some(overflow.id), - "newest frame must be retained" + frame[1]["id"], + terminal.id.to_hex(), + "terminal result drains ahead of the saturated telemetry FIFO" + ); + assert!( + tokio::time::Instant::now().duration_since(saturated_at) < Duration::from_secs(8), + "protected result must beat ModelPicker's eight-second expectation" ); } diff --git a/crates/buzz-acp/tests/startup_signal_cleanup.rs b/crates/buzz-acp/tests/startup_signal_cleanup.rs new file mode 100644 index 0000000000..6a15eb2def --- /dev/null +++ b/crates/buzz-acp/tests/startup_signal_cleanup.rs @@ -0,0 +1,125 @@ +#![cfg(unix)] + +use nix::errno::Errno; +use nix::sys::signal::{kill, killpg, Signal}; +use nix::unistd::Pid; +use std::fs; +use std::os::unix::fs::PermissionsExt; +use std::path::PathBuf; +use std::process::{Child, Command, Stdio}; +use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; + +struct TestProcess { + harness: Child, + adapter_pgid: Option, + temp_dir: PathBuf, +} + +impl Drop for TestProcess { + fn drop(&mut self) { + let _ = self.harness.kill(); + let _ = self.harness.wait(); + if let Some(pgid) = self.adapter_pgid { + let _ = killpg(pgid, Signal::SIGKILL); + } + let _ = fs::remove_dir_all(&self.temp_dir); + } +} + +fn wait_until(mut condition: impl FnMut() -> bool, timeout: Duration) -> bool { + let deadline = Instant::now() + timeout; + while Instant::now() < deadline { + if condition() { + return true; + } + std::thread::sleep(Duration::from_millis(20)); + } + condition() +} + +#[test] +fn repeated_sigterm_during_blocked_eager_initialize_reaps_adapter_process_group() { + let nonce = SystemTime::now() + .duration_since(UNIX_EPOCH) + .expect("system clock") + .as_nanos(); + let temp_dir = + std::env::temp_dir().join(format!("buzz-acp-sigterm-{}-{nonce}", std::process::id())); + fs::create_dir(&temp_dir).expect("create test directory"); + let adapter_script = temp_dir.join("blocked-adapter.sh"); + let adapter_pid_file = temp_dir.join("adapter.pid"); + fs::write( + &adapter_script, + "#!/bin/sh\nset -eu\nprintf '%s\\n' \"$$\" > \"$BUZZ_TEST_ADAPTER_PID_FILE\"\nexec /bin/sleep 300\n", + ) + .expect("write blocked adapter"); + let mut permissions = fs::metadata(&adapter_script) + .expect("adapter metadata") + .permissions(); + permissions.set_mode(0o700); + fs::set_permissions(&adapter_script, permissions).expect("make adapter executable"); + + let harness = Command::new(env!("CARGO_BIN_EXE_buzz-acp")) + .env_clear() + .env("BUZZ_TEST_ADAPTER_PID_FILE", &adapter_pid_file) + .arg("--private-key") + .arg("0000000000000000000000000000000000000000000000000000000000000001") + .arg("--agent-command") + .arg(&adapter_script) + .arg("--agent-args") + .arg("ignored") + .arg("--relay-url") + .arg("ws://127.0.0.1:9") + .arg("--no-presence") + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .spawn() + .expect("spawn buzz-acp harness"); + let mut process = TestProcess { + harness, + adapter_pgid: None, + temp_dir, + }; + + assert!( + wait_until(|| adapter_pid_file.exists(), Duration::from_secs(10)), + "blocked adapter did not publish its PID" + ); + let adapter_pid: i32 = fs::read_to_string(&adapter_pid_file) + .expect("read adapter PID") + .trim() + .parse() + .expect("parse adapter PID"); + let adapter_pgid = Pid::from_raw(adapter_pid); + process.adapter_pgid = Some(adapter_pgid); + assert_eq!( + killpg(adapter_pgid, None::), + Ok(()), + "adapter process group must exist before SIGTERM" + ); + + kill(Pid::from_raw(process.harness.id() as i32), Signal::SIGTERM).expect("signal harness"); + // The listener must remain installed after the graceful request. A second + // signal escalates cleanup without taking the default process-exit path, + // which would abandon process-group verification. + std::thread::sleep(Duration::from_millis(20)); + if process.harness.try_wait().ok().flatten().is_none() { + kill(Pid::from_raw(process.harness.id() as i32), Signal::SIGTERM) + .expect("repeat signal harness during cleanup"); + } + assert!( + wait_until( + || process.harness.try_wait().ok().flatten().is_some(), + Duration::from_secs(10), + ), + "harness did not exit after SIGTERM" + ); + assert!( + wait_until( + || killpg(adapter_pgid, None::) == Err(Errno::ESRCH), + Duration::from_secs(5), + ), + "adapter process group {adapter_pid} survived harness SIGTERM" + ); + process.adapter_pgid = None; +} diff --git a/desktop/src/features/agents/AGENTS.md b/desktop/src/features/agents/AGENTS.md index 35ad4a63af..4a7232e12f 100644 --- a/desktop/src/features/agents/AGENTS.md +++ b/desktop/src/features/agents/AGENTS.md @@ -114,6 +114,16 @@ with a TypeScript lookup table or an id comparison in a component. published or removed. A queued update must stay visibly queued, and the catalog itself must render only relay-confirmed publications — never an optimistic local persona. +11. **Live model-switch success requires request-correlated, per-channel + application proof.** Every switch operation generates an unpredictable + request ID and accepts only signed, fresh `control_result` frames carrying + that exact request ID and model. `sent` and `recycling` are progress only; + each target channel must independently report `switched`, while explicit + terminal failures fail fast. A timeout is `pending`, never success. The + terminal-proof claim journal is relay/community-scoped process state: + `resetCommunityState()` must replace it on every community switch so a + claim observed on the outgoing relay cannot suppress proof on the incoming + relay. ## The tests that enforce this @@ -128,6 +138,9 @@ with a TypeScript lookup table or an id comparison in a component. `isCacheableDiscoveryResponse`, `deriveModelDiscoveryPending`, `isSuccessfulEmptyDiscovery`. If the "reopen to retry" copy becomes inert again, these tests will catch it. +- `lib/liveSwitchOutcome.test.mjs` — exact request correlation, signed + freshness, per-channel `switched` proof, pending-on-timeout behavior, replay + claims, and community-reset replacement. - `desktop/tests/e2e/onboarding-agent-defaults.spec.ts` — onboarding behavior acceptance coverage for readiness, failure states, defaults, navigation, successful-empty vs failed optional-model discovery, and persistence races. diff --git a/desktop/src/features/agents/lib/liveSwitchOutcome.test.mjs b/desktop/src/features/agents/lib/liveSwitchOutcome.test.mjs index 737d84b862..1368a80f92 100644 --- a/desktop/src/features/agents/lib/liveSwitchOutcome.test.mjs +++ b/desktop/src/features/agents/lib/liveSwitchOutcome.test.mjs @@ -1,12 +1,33 @@ import assert from "node:assert/strict"; import test from "node:test"; -import { awaitLiveSwitchOutcome } from "./liveSwitchOutcome.ts"; +import { + awaitLiveSwitchOutcome, + createTerminalClaimJournal, + resetLiveSwitchOutcomeClaims, +} from "./liveSwitchOutcome.ts"; const MODEL = "goose-claude-fable-5"; +const CHANNEL_A = "channel-a"; +const CHANNEL_B = "channel-b"; +const REQUEST_ID = "a".repeat(32); +const REQUEST_BOUNDARY_MS = Date.parse("2026-07-28T12:00:00.000Z"); +let frameId = 0; function frame(status, overrides = {}) { - return { type: "switch_model", status, modelId: MODEL, ...overrides }; + frameId += 1; + return { + type: "switch_model", + status, + modelId: MODEL, + channelId: CHANNEL_A, + requestId: REQUEST_ID, + relayEventId: frameId.toString(16).padStart(64, "0"), + relayCreatedAt: Math.floor((REQUEST_BOUNDARY_MS + 1_000) / 1_000), + observerTimestamp: new Date(REQUEST_BOUNDARY_MS + 1_000).toISOString(), + observerSeq: frameId, + ...overrides, + }; } /** @@ -15,19 +36,22 @@ function frame(status, overrides = {}) { * no-ops, matching `observerRelayStore`), a manual timeout, and a deferred * `sendSwitches` the test resolves explicitly. */ -function harness(channelCount) { +function harness(channelIds, requestId = REQUEST_ID) { let listener = null; let timeoutCb = null; let unsubscribeCalls = 0; let cancelTimeoutCalls = 0; let sendResolve; + let sentRequestId = null; + let subscribedWhenSent = false; const sendStarted = new Promise((resolve) => { sendResolve = resolve; }); const outcome = awaitLiveSwitchOutcome({ - channelCount, + channelIds, modelId: MODEL, + createRequestId: () => requestId, subscribe: (fn) => { listener = fn; return () => { @@ -35,7 +59,9 @@ function harness(channelCount) { listener = null; }; }, - sendSwitches: () => { + sendSwitches: (actualRequestId) => { + sentRequestId = actualRequestId; + subscribedWhenSent = listener !== null; sendResolve(); return Promise.resolve(); }, @@ -45,6 +71,7 @@ function harness(channelCount) { cancelTimeoutCalls += 1; }; }, + now: () => REQUEST_BOUNDARY_MS, }); return { @@ -58,20 +85,26 @@ function harness(channelCount) { get cancelTimeoutCalls() { return cancelTimeoutCalls; }, + get sentRequestId() { + return sentRequestId; + }, + get subscribedWhenSent() { + return subscribedWhenSent; + }, }; } test("awaitLiveSwitchOutcome fast sent on one channel does not mask a later unsupported on another", async () => { - const h = harness(2); - // Channel A acks fast as `sent`; a first-ack-resolves impl would settle "ok" - // here. The fail-fast contract must keep waiting and then reject on B. + const h = harness([CHANNEL_A, CHANNEL_B]); + // Channel A only acknowledges receipt. The fail-fast contract must keep + // waiting and then reject when B proves the model unsupported. h.push(frame("sent")); - h.push(frame("unsupported_model")); + h.push(frame("unsupported_model", { channelId: CHANNEL_B })); assert.equal(await h.outcome, "unsupported"); }); -test("awaitLiveSwitchOutcome resolves ok only after the last channel acks", async () => { - const h = harness(3); +test("awaitLiveSwitchOutcome resolves ok only after every channel proves application", async () => { + const h = harness([CHANNEL_A, CHANNEL_B]); let settled = false; void h.outcome.then(() => { settled = true; @@ -88,20 +121,28 @@ test("awaitLiveSwitchOutcome resolves ok only after the last channel acks", asyn } }; - h.push(frame("sent")); + h.push(frame("sent", { channelId: CHANNEL_A })); await drainMicrotasks(); - assert.equal(settled, false, "must not resolve on the first ack"); + assert.equal(settled, false, "receipt is not proof of application"); - h.push(frame("switched")); + h.push(frame("recycling", { channelId: CHANNEL_B })); + await drainMicrotasks(); + assert.equal( + settled, + false, + "replacement scheduling is not proof of application", + ); + + h.push(frame("switched", { channelId: CHANNEL_A })); await drainMicrotasks(); - assert.equal(settled, false, "must not resolve before the last ack"); + assert.equal(settled, false, "must wait for every target channel"); - h.push(frame("turn_ending")); + h.push(frame("switched", { channelId: CHANNEL_B })); assert.equal(await h.outcome, "ok"); }); test("awaitLiveSwitchOutcome rejects on unsupported immediately and unsubscribes exactly once", async () => { - const h = harness(2); + const h = harness([CHANNEL_A, CHANNEL_B]); h.push(frame("unsupported_model")); assert.equal(await h.outcome, "unsupported"); assert.equal(h.unsubscribeCalls, 1); @@ -114,9 +155,11 @@ test("awaitLiveSwitchOutcome rejects on unsupported immediately and unsubscribes }); test("awaitLiveSwitchOutcome ignores frames for a different model or control type", async () => { - const h = harness(1); + const h = harness([CHANNEL_A]); h.push(frame("sent", { modelId: "some-other-model" })); h.push({ type: "cancel_turn", status: "sent", modelId: MODEL }); + h.push(frame("switched", { channelId: "unrelated-channel" })); + h.push(frame("switched", { requestId: "b".repeat(32) })); let settled = false; void h.outcome.then(() => { settled = true; @@ -128,27 +171,413 @@ test("awaitLiveSwitchOutcome ignores frames for a different model or control typ assert.equal(await h.outcome, "ok"); }); -test("awaitLiveSwitchOutcome resolves ok via the timeout fallback when the harness never replies", async () => { - const h = harness(2); +test("awaitLiveSwitchOutcome stays pending when the harness never proves application", async () => { + const h = harness([CHANNEL_A, CHANNEL_B]); h.fireTimeout(); - assert.equal(await h.outcome, "ok"); + assert.equal(await h.outcome, "pending"); assert.equal(h.unsubscribeCalls, 1, "timeout fallback unsubscribes"); }); +test("awaitLiveSwitchOutcome timeout is not blocked by a hung relay send", async () => { + let fireTimeout = () => {}; + const outcome = awaitLiveSwitchOutcome({ + channelIds: [CHANNEL_A], + modelId: MODEL, + createRequestId: () => REQUEST_ID, + subscribe: () => () => {}, + sendSwitches: () => new Promise(() => {}), + scheduleTimeout: (onTimeout) => { + fireTimeout = onTimeout; + return () => {}; + }, + now: () => REQUEST_BOUNDARY_MS, + }); + + fireTimeout(); + const result = await Promise.race([ + outcome, + new Promise((resolve) => setImmediate(() => resolve("still-blocked"))), + ]); + assert.equal(result, "pending"); +}); + +for (const [status, expected] of [ + ["unsupported_model", "unsupported"], + ["switch_failed", "failed"], + ["switched", "ok"], +]) { + test(`awaitLiveSwitchOutcome ${expected} proof is not blocked by a hung relay send`, async () => { + let listener = () => {}; + const outcome = awaitLiveSwitchOutcome({ + channelIds: [CHANNEL_A], + modelId: MODEL, + createRequestId: () => REQUEST_ID, + subscribe: (onFrame) => { + listener = onFrame; + return () => {}; + }, + sendSwitches: () => new Promise(() => {}), + scheduleTimeout: () => () => {}, + now: () => REQUEST_BOUNDARY_MS, + }); + + listener(frame(status)); + const result = await Promise.race([ + outcome, + new Promise((resolve) => setImmediate(() => resolve("still-blocked"))), + ]); + assert.equal(result, expected); + }); +} + +test("awaitLiveSwitchOutcome rejects when the relay send fails before proof", async () => { + const outcome = awaitLiveSwitchOutcome({ + channelIds: [CHANNEL_A], + modelId: MODEL, + createRequestId: () => REQUEST_ID, + subscribe: () => () => {}, + sendSwitches: () => Promise.reject(new Error("relay send failed")), + scheduleTimeout: () => () => {}, + now: () => REQUEST_BOUNDARY_MS, + }); + + await assert.rejects(outcome, /relay send failed/); +}); + +test("awaitLiveSwitchOutcome observes a relay send rejection after proof", async () => { + let listener = () => {}; + let rejectSend = () => {}; + let unhandled; + const onUnhandled = (reason) => { + unhandled = reason; + }; + process.on("unhandledRejection", onUnhandled); + try { + const outcome = awaitLiveSwitchOutcome({ + channelIds: [CHANNEL_A], + modelId: MODEL, + createRequestId: () => REQUEST_ID, + subscribe: (onFrame) => { + listener = onFrame; + return () => {}; + }, + sendSwitches: () => + new Promise((_resolve, reject) => { + rejectSend = reject; + }), + scheduleTimeout: () => () => {}, + now: () => REQUEST_BOUNDARY_MS, + }); + + listener(frame("switched")); + assert.equal(await outcome, "ok"); + rejectSend(new Error("late relay send failure")); + await new Promise((resolve) => setImmediate(resolve)); + assert.equal(unhandled, undefined); + } finally { + process.off("unhandledRejection", onUnhandled); + } +}); + test("awaitLiveSwitchOutcome fires the per-channel sends after subscribing", async () => { - const h = harness(1); + const h = harness([CHANNEL_A]); // The subscription is registered before the sends fire, so a frame arriving // mid-send is never dropped. Awaiting sendStarted proves sends ran. await h.sendStarted; - h.push(frame("sent")); + assert.equal(h.sentRequestId, REQUEST_ID); + assert.equal(h.subscribedWhenSent, true); + h.push(frame("switched")); assert.equal(await h.outcome, "ok"); }); -test("awaitLiveSwitchOutcome with zero channels resolves ok at the timeout (no acks expected)", async () => { - // No active turns means channelCount 0: remaining starts at 0 but the success - // resolve only fires inside a frame callback, so with no frames the timeout - // fallback is what settles it. This documents the degenerate path. - const h = harness(0); +test("awaitLiveSwitchOutcome with zero channels resolves ok without waiting for proof", async () => { + const h = harness([]); + assert.equal(await h.outcome, "ok"); +}); + +test("awaitLiveSwitchOutcome does not double-count repeated terminal frames", async () => { + const h = harness([CHANNEL_A, CHANNEL_B]); + let settled = false; + void h.outcome.then(() => { + settled = true; + }); + + h.push(frame("switched", { channelId: CHANNEL_A })); + h.push(frame("switched", { channelId: CHANNEL_A })); + await Promise.resolve(); + assert.equal( + settled, + false, + "one channel cannot satisfy two target channels", + ); + + h.push(frame("switched", { channelId: CHANNEL_B })); + assert.equal(await h.outcome, "ok"); +}); + +test("awaitLiveSwitchOutcome accepts exact-request proof emitted 1ms before the desktop clock", async () => { + const h = harness([CHANNEL_A]); + h.push( + frame("switched", { + relayCreatedAt: Math.floor(REQUEST_BOUNDARY_MS / 1_000), + observerTimestamp: new Date(REQUEST_BOUNDARY_MS - 1).toISOString(), + }), + ); h.fireTimeout(); + assert.equal( + await h.outcome, + "ok", + "exact unpredictable request correlation is causal proof despite small cross-node skew", + ); +}); + +test("awaitLiveSwitchOutcome rejects a relay event outside the broad freshness window", async () => { + const h = harness([CHANNEL_A]); + let settled = false; + void h.outcome.then(() => { + settled = true; + }); + + h.push( + frame("switched", { + relayCreatedAt: Math.floor((REQUEST_BOUNDARY_MS - 10 * 60_000) / 1_000), + }), + ); + for (let i = 0; i < 5; i++) { + await Promise.resolve(); + } + assert.equal( + settled, + false, + "an old signed relay envelope is stale even with an exact request ID", + ); + + h.push(frame("switched")); + assert.equal(await h.outcome, "ok"); +}); + +test("awaitLiveSwitchOutcome rejects observer payload time outside the broad freshness window", async () => { + const h = harness([CHANNEL_A]); + let settled = false; + void h.outcome.then(() => { + settled = true; + }); + + h.push( + frame("switched", { + observerTimestamp: new Date( + REQUEST_BOUNDARY_MS - 10 * 60_000, + ).toISOString(), + }), + ); + for (let i = 0; i < 5; i++) { + await Promise.resolve(); + } + assert.equal( + settled, + false, + "a fresh relay envelope cannot make an old observer payload current", + ); + + h.push(frame("switched")); + assert.equal(await h.outcome, "ok"); +}); + +test("awaitLiveSwitchOutcome rejects relay and observer times beyond cross-node future skew", async () => { + for (const overrides of [ + { + relayCreatedAt: Math.floor((REQUEST_BOUNDARY_MS + 10 * 60_000) / 1_000), + }, + { + observerTimestamp: new Date( + REQUEST_BOUNDARY_MS + 10 * 60_000, + ).toISOString(), + }, + ]) { + const h = harness([CHANNEL_A]); + let settled = false; + void h.outcome.then(() => { + settled = true; + }); + h.push(frame("switched", overrides)); + for (let i = 0; i < 5; i++) { + await Promise.resolve(); + } + assert.equal(settled, false, "implausibly future proof must be ignored"); + h.push(frame("switched")); + assert.equal(await h.outcome, "ok"); + } +}); + +test("awaitLiveSwitchOutcome fails closed when signed freshness metadata is absent", async () => { + const h = harness([CHANNEL_A]); + let settled = false; + void h.outcome.then(() => { + settled = true; + }); + + h.push( + frame("switched", { + relayEventId: undefined, + relayCreatedAt: undefined, + observerTimestamp: undefined, + }), + ); + for (let i = 0; i < 5; i++) { + await Promise.resolve(); + } + assert.equal(settled, false); + + h.push(frame("switched")); assert.equal(await h.outcome, "ok"); }); + +test("concurrent same-model requests resolve only their exact request result", async () => { + const listeners = new Set(); + const startRequest = (requestId) => + awaitLiveSwitchOutcome({ + channelIds: [CHANNEL_A], + modelId: MODEL, + createRequestId: () => requestId, + subscribe: (listener) => { + listeners.add(listener); + return () => listeners.delete(listener); + }, + sendSwitches: (sentRequestId) => { + assert.equal(sentRequestId, requestId); + return Promise.resolve(); + }, + scheduleTimeout: () => () => {}, + now: () => REQUEST_BOUNDARY_MS, + }); + + const requestA = "1".repeat(32); + const requestB = "2".repeat(32); + const first = startRequest(requestA); + const second = startRequest(requestB); + + const secondFailure = frame("turn_ending", { requestId: requestB }); + for (const listener of [...listeners]) { + listener(secondFailure); + } + const firstProof = frame("switched", { requestId: requestA }); + for (const listener of [...listeners]) { + listener(firstProof); + } + + assert.equal(await first, "ok"); + assert.equal(await second, "failed"); +}); + +test("more than 4096 sequential completed operations remain live", async () => { + for (let index = 0; index < 4_100; index += 1) { + const requestId = index.toString(16).padStart(32, "0"); + const h = harness([CHANNEL_A], requestId); + h.push( + frame("switched", { + requestId, + relayEventId: (index + 10_000).toString(16).padStart(64, "0"), + }), + ); + h.fireTimeout(); + assert.equal( + await h.outcome, + "ok", + `completed operation ${index} must not be disabled by a global claim cap`, + ); + } +}); + +test("a duplicate relay ID and request pair cannot satisfy another request", async () => { + const requestId = "c".repeat(32); + const relayEventId = "d".repeat(64); + const first = harness([CHANNEL_A], requestId); + first.push(frame("switched", { requestId, relayEventId })); + assert.equal(await first.outcome, "ok"); + + const second = harness([CHANNEL_A], requestId); + second.push(frame("switched", { requestId, relayEventId })); + second.fireTimeout(); + assert.equal(await second.outcome, "pending"); +}); + +test("community reset replaces the relay-scoped terminal claim journal", async () => { + const requestId = "9".repeat(32); + const relayEventId = "8".repeat(64); + const first = harness([CHANNEL_A], requestId); + first.push(frame("switched", { requestId, relayEventId })); + assert.equal(await first.outcome, "ok"); + + resetLiveSwitchOutcomeClaims(); + + const afterCommunitySwitch = harness([CHANNEL_A], requestId); + afterCommunitySwitch.push(frame("switched", { requestId, relayEventId })); + assert.equal( + await afterCommunitySwitch.outcome, + "ok", + "claims from the outgoing relay must not suppress proof in the new community", + ); +}); + +test("terminal claims expire after their whole freshness horizon", () => { + let currentTimeMs = REQUEST_BOUNDARY_MS; + const claim = createTerminalClaimJournal({ + retentionMs: 1_000, + now: () => currentTimeMs, + }); + const proof = frame("switched", { + requestId: "e".repeat(32), + relayEventId: "f".repeat(64), + }); + + assert.equal(claim(proof), true); + assert.equal(claim(proof), false); + currentTimeMs += 1_001; + assert.equal( + claim(proof), + true, + "expired claims must be reclaimed so the journal is freshness-bounded", + ); +}); + +test("the production request ID generator does not reuse an operation ID", async () => { + const requestIds = []; + for (let index = 0; index < 2; index += 1) { + let listener = null; + const outcome = awaitLiveSwitchOutcome({ + channelIds: [CHANNEL_A], + modelId: MODEL, + subscribe: (nextListener) => { + listener = nextListener; + return () => { + listener = null; + }; + }, + sendSwitches: (requestId) => { + requestIds.push(requestId); + listener?.( + frame("switched", { + requestId, + relayEventId: (index + 20_000).toString(16).padStart(64, "0"), + }), + ); + return Promise.resolve(); + }, + scheduleTimeout: () => () => {}, + now: () => REQUEST_BOUNDARY_MS, + }); + assert.equal(await outcome, "ok"); + } + + assert.match(requestIds[0], /^[0-9a-f]{32}$/); + assert.match(requestIds[1], /^[0-9a-f]{32}$/); + assert.notEqual(requestIds[0], requestIds[1]); +}); + +for (const status of ["switch_failed", "turn_ending", "no_active_turn"]) { + test(`awaitLiveSwitchOutcome reports ${status} as a failed application`, async () => { + const h = harness([CHANNEL_A]); + h.push(frame(status)); + assert.equal(await h.outcome, "failed"); + }); +} diff --git a/desktop/src/features/agents/lib/liveSwitchOutcome.ts b/desktop/src/features/agents/lib/liveSwitchOutcome.ts index d12261e596..487e83eb0d 100644 --- a/desktop/src/features/agents/lib/liveSwitchOutcome.ts +++ b/desktop/src/features/agents/lib/liveSwitchOutcome.ts @@ -1,66 +1,274 @@ import type { ControlResultFrame } from "@/shared/api/types"; +const SIGNED_RELAY_EVENT_ID = /^[0-9a-f]{64}$/i; +const CONTROL_REQUEST_ID = /^[0-9a-f]{32}$/; +// Exact unpredictable request correlation is the causal boundary. These +// timestamps are only broad sanity checks: tolerate independently skewed node +// clocks while rejecting old retained events and implausibly future results. +const CONTROL_RESULT_MAX_AGE_MS = 5 * 60_000; +const CONTROL_RESULT_MAX_FUTURE_SKEW_MS = 2 * 60_000; +const CONTROL_RESULT_REPLAY_RETENTION_MS = + CONTROL_RESULT_MAX_AGE_MS + CONTROL_RESULT_MAX_FUTURE_SKEW_MS + 1; + +/** + * Build a freshness-expiring terminal-proof claim journal. + * + * Claims live for the whole interval in which the same signed result could + * pass the broad freshness check. After that horizon the result itself is too + * old to satisfy a request, so retaining its ID would only create an unbounded + * process-lifetime leak. There is deliberately no cardinality cap: a busy + * desktop must not permanently disable live switching after N operations. + */ +export function createTerminalClaimJournal({ + retentionMs = CONTROL_RESULT_REPLAY_RETENTION_MS, + now = Date.now, +}: { + retentionMs?: number; + now?: () => number; +} = {}): (frame: ControlResultFrame, observedAtMs?: number) => boolean { + if (!Number.isFinite(retentionMs) || retentionMs <= 0) { + throw new Error("terminal claim retention must be positive"); + } + const claimedPairs = new Map(); + const expirations: Array<{ + key: string; + expiresAtMs: number; + }> = []; + let expirationCursor = 0; + + const pruneExpired = (observedAtMs: number) => { + while ( + expirationCursor < expirations.length && + expirations[expirationCursor].expiresAtMs <= observedAtMs + ) { + const expired = expirations[expirationCursor]; + if (claimedPairs.get(expired.key) === expired.expiresAtMs) { + claimedPairs.delete(expired.key); + } + expirationCursor += 1; + } + if ( + expirationCursor >= 1_024 && + expirationCursor * 2 >= expirations.length + ) { + expirations.splice(0, expirationCursor); + expirationCursor = 0; + } + }; + + return (frame, observedAtMs = now()) => { + const eventId = frame.relayEventId; + const requestId = frame.requestId; + if ( + !eventId || + !SIGNED_RELAY_EVENT_ID.test(eventId) || + !requestId || + !CONTROL_REQUEST_ID.test(requestId) + ) { + return false; + } + if (!Number.isFinite(observedAtMs)) { + return false; + } + pruneExpired(observedAtMs); + const key = `${eventId}:${requestId}`; + if (claimedPairs.has(key)) { + return false; + } + const expiresAtMs = observedAtMs + retentionMs; + claimedPairs.set(key, expiresAtMs); + expirations.push({ key, expiresAtMs }); + return true; + }; +} + +let claimControlResult = createTerminalClaimJournal(); + +/** + * Replace the relay-scoped terminal-proof claim journal. + * + * Community switches call this after disconnecting the outgoing relay so a + * signed result claimed there cannot suppress an otherwise valid result from + * the newly active relay/community. + */ +export function resetLiveSwitchOutcomeClaims(): void { + claimControlResult = createTerminalClaimJournal(); +} + +function createControlRequestId(): string { + const bytes = globalThis.crypto.getRandomValues(new Uint8Array(16)); + return Array.from(bytes, (byte) => byte.toString(16).padStart(2, "0")).join( + "", + ); +} + +function isFreshControlResult( + frame: ControlResultFrame, + observedAtMs: number, +): boolean { + if ( + !frame.relayEventId || + !SIGNED_RELAY_EVENT_ID.test(frame.relayEventId) || + typeof frame.relayCreatedAt !== "number" || + !Number.isSafeInteger(frame.relayCreatedAt) || + !frame.observerTimestamp + ) { + return false; + } + + const observerTimestampMs = Date.parse(frame.observerTimestamp); + if (!Number.isFinite(observerTimestampMs) || !Number.isFinite(observedAtMs)) { + return false; + } + + const relayCreatedAtMs = frame.relayCreatedAt * 1_000; + const oldestAllowedMs = observedAtMs - CONTROL_RESULT_MAX_AGE_MS; + const newestAllowedMs = observedAtMs + CONTROL_RESULT_MAX_FUTURE_SKEW_MS; + return ( + relayCreatedAtMs >= oldestAllowedMs && + relayCreatedAtMs <= newestAllowedMs && + observerTimestampMs >= oldestAllowedMs && + observerTimestampMs <= newestAllowedMs + ); +} + /** * Resolve the outcome of a live `switch_model` across one or more channels. * * A live switch fires a `switch_model` frame per active channel and learns each - * channel's result asynchronously over the observer relay. The fail-fast rule: - * any single `unsupported_model` result rejects the whole pick immediately; - * every other status must arrive from every channel before resolving success. - * If the harness never replies, the fallback timeout resolves `"ok"` — the - * override still rides the requeued/next session, we just can't confirm it - * synchronously. + * channel's result asynchronously over the observer relay. Receipt (`sent`) and + * replacement scheduling (`recycling`) are non-terminal. Success requires a + * distinct `switched` application proof from every target channel; repeated + * frames from one channel cannot satisfy another. Unsupported and explicit + * failure statuses fail fast. A timeout is reported as `pending`, never as + * success. * * The counting lives here, isolated from React and the relay so it can be unit * tested with synthetic frames and a fake clock. The caller injects the * relay subscription, the per-channel sends, and the timeout scheduler. */ export async function awaitLiveSwitchOutcome({ - channelCount, + channelIds, modelId, + createRequestId = createControlRequestId, subscribe, sendSwitches, scheduleTimeout, + now = Date.now, }: { - /** Number of channels the switch was fired to — the success threshold. */ - channelCount: number; + /** Exact channels the switch was fired to. */ + channelIds: string[]; /** Model being switched to; frames for any other model are ignored. */ modelId: string; + /** Test seam for the unique control request identifier. */ + createRequestId?: () => string; /** Register a control-result listener; returns an unsubscribe function. */ subscribe: (listener: (frame: ControlResultFrame) => void) => () => void; - /** Fire the per-channel `switch_model` sends. Resolves when all are sent. */ - sendSwitches: () => Promise; + /** Fire the per-channel sends with this operation's exact request ID. */ + sendSwitches: (requestId: string) => Promise; /** Schedule the no-reply fallback; returns a cancel function. */ scheduleTimeout: (onTimeout: () => void) => () => void; -}): Promise<"ok" | "unsupported"> { - const settled = new Promise<"ok" | "unsupported">((resolve) => { - let unsubscribe = () => {}; - let cancelTimeout = () => {}; - let remaining = channelCount; - const finish = (outcome: "ok" | "unsupported") => { - cancelTimeout(); - unsubscribe(); - resolve(outcome); - }; - cancelTimeout = scheduleTimeout(() => finish("ok")); - unsubscribe = subscribe((frame) => { - if (frame.type !== "switch_model" || frame.modelId !== modelId) { - return; - } - if (frame.status === "unsupported_model") { - // Any single failure rejects the whole pick immediately. - finish("unsupported"); - return; - } - // sent / switched / turn_ending — count as success for this channel. - remaining -= 1; - if (remaining <= 0) { - finish("ok"); - } - }); - }); + /** Clock seam for receive-time freshness checks and replay-claim expiry. */ + now?: () => number; +}): Promise<"ok" | "unsupported" | "failed" | "pending"> { + const requestId = createRequestId(); + if (!CONTROL_REQUEST_ID.test(requestId)) { + throw new Error( + "live model-switch request ID must be 32 lowercase hex characters", + ); + } + const targetChannels = new Set(channelIds); + if (targetChannels.size === 0) { + await sendSwitches(requestId); + return "ok"; + } + + let unsubscribe = () => {}; + let cancelTimeout = () => {}; + let finished = false; + const appliedChannels = new Set(); + let finish: (outcome: "ok" | "unsupported" | "failed" | "pending") => void = + () => {}; + const settled = new Promise<"ok" | "unsupported" | "failed" | "pending">( + (resolve) => { + finish = (outcome) => { + if (finished) return; + finished = true; + cancelTimeout(); + unsubscribe(); + resolve(outcome); + }; + cancelTimeout = scheduleTimeout(() => finish("pending")); + unsubscribe = subscribe((frame) => { + const observedAtMs = now(); + if ( + frame.type !== "switch_model" || + frame.requestId !== requestId || + frame.modelId !== modelId || + !frame.channelId || + !targetChannels.has(frame.channelId) || + !isFreshControlResult(frame, observedAtMs) + ) { + return; + } + const isFailure = + frame.status === "unsupported_model" || + frame.status === "switch_failed" || + frame.status === "turn_ending" || + frame.status === "no_active_turn"; + const isNewApplication = + frame.status === "switched" && !appliedChannels.has(frame.channelId); + if ( + (isFailure || isNewApplication) && + !claimControlResult(frame, observedAtMs) + ) { + return; + } + if (frame.status === "unsupported_model") { + finish("unsupported"); + return; + } + if ( + frame.status === "switch_failed" || + frame.status === "turn_ending" || + frame.status === "no_active_turn" + ) { + finish("failed"); + return; + } + if (frame.status !== "switched") { + // `sent`, `recycling`, and unknown future non-terminal statuses do not + // prove that the model reached this channel's ACP session. + return; + } + appliedChannels.add(frame.channelId); + if (appliedChannels.size === targetChannels.size) { + finish("ok"); + } + }); + }, + ); - await sendSwitches(); + let sends: Promise; + try { + sends = sendSwitches(requestId); + } catch (error) { + finished = true; + cancelTimeout(); + unsubscribe(); + throw error; + } + const sendFailure = sends.then( + () => new Promise(() => {}), + (error) => { + if (!finished) { + finished = true; + cancelTimeout(); + unsubscribe(); + } + throw error; + }, + ); - return settled; + return Promise.race([settled, sendFailure]); } diff --git a/desktop/src/features/agents/observerRelayControlResults.test.mjs b/desktop/src/features/agents/observerRelayControlResults.test.mjs new file mode 100644 index 0000000000..d595fe6887 --- /dev/null +++ b/desktop/src/features/agents/observerRelayControlResults.test.mjs @@ -0,0 +1,119 @@ +import assert from "node:assert/strict"; +import { beforeEach, test } from "node:test"; + +import { + _testIngestControlResult, + resetAgentObserverStore, + subscribeControlResults, +} from "./observerRelayStore.ts"; + +const AGENT = "a".repeat(64); +const MODEL = "goose-claude-fable-5"; +const CHANNEL = "channel-a"; +const REQUEST_ID = "a".repeat(32); + +function relayEvent(overrides = {}) { + return { + id: "1".repeat(64), + pubkey: AGENT, + created_at: 1_775_000_001, + kind: 24_200, + tags: [ + ["agent", AGENT], + ["frame", "telemetry"], + ], + content: "encrypted", + sig: "2".repeat(128), + ...overrides, + }; +} + +function observerEvent(overrides = {}) { + return { + seq: 7, + timestamp: "2026-04-01T00:00:01.123Z", + kind: "control_result", + agentIndex: 0, + channelId: CHANNEL, + sessionId: null, + turnId: null, + payload: { + type: "switch_model", + status: "switched", + modelId: MODEL, + requestId: REQUEST_ID, + }, + ...overrides, + }; +} + +beforeEach(() => { + resetAgentObserverStore(); +}); + +test("control-result dispatch carries signed relay freshness metadata", () => { + const received = []; + const unsubscribe = subscribeControlResults(AGENT, (frame) => { + received.push(frame); + }); + + _testIngestControlResult(AGENT, relayEvent(), observerEvent()); + unsubscribe(); + + assert.deepEqual(received, [ + { + type: "switch_model", + status: "switched", + modelId: MODEL, + requestId: REQUEST_ID, + channelId: CHANNEL, + relayEventId: "1".repeat(64), + relayCreatedAt: 1_775_000_001, + observerTimestamp: "2026-04-01T00:00:01.123Z", + observerSeq: 7, + }, + ]); +}); + +test("the authenticated observer envelope channel overrides an inner payload claim", () => { + const received = []; + const unsubscribe = subscribeControlResults(AGENT, (frame) => { + received.push(frame); + }); + + _testIngestControlResult( + AGENT, + relayEvent(), + observerEvent({ + payload: { + type: "switch_model", + status: "switched", + modelId: MODEL, + channelId: "spoofed-channel", + }, + }), + ); + unsubscribe(); + + assert.equal(received.length, 1); + assert.equal(received[0].channelId, CHANNEL); +}); + +test("a replay republished under a new signed event id cannot redispatch a deduped observer result", () => { + const received = []; + const unsubscribe = subscribeControlResults(AGENT, (frame) => { + received.push(frame); + }); + const parsed = observerEvent(); + + _testIngestControlResult(AGENT, relayEvent(), parsed); + _testIngestControlResult( + AGENT, + relayEvent({ id: "3".repeat(64), created_at: 1_775_000_099 }), + parsed, + ); + unsubscribe(); + + assert.equal(received.length, 1); + assert.equal(received[0].relayEventId, "1".repeat(64)); +}); diff --git a/desktop/src/features/agents/observerRelayStore.ts b/desktop/src/features/agents/observerRelayStore.ts index 56c69f915a..8685753032 100644 --- a/desktop/src/features/agents/observerRelayStore.ts +++ b/desktop/src/features/agents/observerRelayStore.ts @@ -193,7 +193,7 @@ function observerTag(event: RelayEvent, tagName: string) { return event.tags.find((tag) => tag[0] === tagName)?.[1] ?? null; } -function appendAgentEvent(agentPubkey: string, event: ObserverEvent) { +function appendAgentEvent(agentPubkey: string, event: ObserverEvent): boolean { const key = normalizePubkey(agentPubkey); const current = eventsByAgent.get(key) ?? []; if ( @@ -202,7 +202,7 @@ function appendAgentEvent(agentPubkey: string, event: ObserverEvent) { existing.seq === event.seq && existing.timestamp === event.timestamp, ) ) { - return; + return false; } const sorted = [...current, event].sort(compareObserverEvents); @@ -231,6 +231,7 @@ function appendAgentEvent(agentPubkey: string, event: ObserverEvent) { invalidateSnapshot(key); notifyListeners(); + return true; } /** @@ -391,7 +392,7 @@ async function handleRelayObserverEvent( }); } } - appendAgentEvent(agentPubkey, parsed); + const eventAdded = appendAgentEvent(agentPubkey, parsed); const managementRequest = parseAgentManagementRequest(parsed.payload); if (managementRequest) { for (const listener of agentManagementListeners) { @@ -402,7 +403,9 @@ async function handleRelayObserverEvent( void putAgentSessionConfig(agentPubkey, parsed.payload); onSessionConfigCaptured?.(agentPubkey); } else if (parsed.kind === "control_result") { - dispatchControlResult(agentPubkey, parsed.payload); + if (eventAdded) { + dispatchControlResult(agentPubkey, parsed, event); + } } else if (parsed.kind === "managed_agent_runtime_lifecycle") { void putManagedAgentRuntimeLifecycle(agentPubkey, parsed.payload).catch( (error) => { @@ -495,7 +498,12 @@ function isControlResultFrame(payload: unknown): payload is ControlResultFrame { ); } -function dispatchControlResult(agentPubkey: string, payload: unknown) { +function dispatchControlResult( + agentPubkey: string, + observerEvent: ObserverEvent, + relayEvent: RelayEvent, +) { + const payload = observerEvent.payload; if (!isControlResultFrame(payload)) { return; } @@ -503,8 +511,16 @@ function dispatchControlResult(agentPubkey: string, payload: unknown) { if (!subscribers) { return; } + const frame = { + ...payload, + channelId: observerEvent.channelId ?? payload.channelId ?? undefined, + relayEventId: relayEvent.id, + relayCreatedAt: relayEvent.created_at, + observerTimestamp: observerEvent.timestamp, + observerSeq: observerEvent.seq, + }; for (const subscriber of subscribers) { - subscriber(payload); + subscriber(frame); } } @@ -542,6 +558,21 @@ export function subscribeControlResults( }; } +/** + * Test-only seam for the decoded live `control_result` ingestion path. + * Exercises the same `(seq, timestamp)` replay dedup and signed-envelope + * metadata enrichment used after production decryption. + */ +export function _testIngestControlResult( + agentPubkey: string, + relayEvent: RelayEvent, + observerEvent: ObserverEvent, +): void { + if (appendAgentEvent(agentPubkey, observerEvent)) { + dispatchControlResult(agentPubkey, observerEvent, relayEvent); + } +} + export function getAgentObserverSnapshot( agentPubkey?: string | null, // `_enabled` previously gated store reads — now only gates the relay @@ -749,6 +780,7 @@ export function resetAgentObserverStore() { knownAgentsBySubscription.clear(); pendingUnknownAgentFrames.length = 0; latestLiveSessionByAgentChannel.clear(); + controlResultListeners.clear(); agentManagementListeners.clear(); onSessionConfigCaptured = null; connectionState = "idle"; diff --git a/desktop/src/features/agents/ui/ModelPicker.tsx b/desktop/src/features/agents/ui/ModelPicker.tsx index f7bafde99b..0f3c2122a5 100644 --- a/desktop/src/features/agents/ui/ModelPicker.tsx +++ b/desktop/src/features/agents/ui/ModelPicker.tsx @@ -108,26 +108,32 @@ export function ModelPicker({ }, [configSurface]); // Send a live `switch_model` frame to each channel the agent is working in - // and wait for the harness to acknowledge. Any single `unsupported_model` - // result rejects the whole pick immediately; all other statuses must arrive - // from every channel before resolving success. + // and wait for channel-scoped application proof. Receipt and recycle frames + // remain pending; every channel must report `switched` before success. const sendLiveSwitch = React.useCallback( (modelId: string) => { - const channelIds = activeTurns.map((turn) => turn.channelId); + const channelIds = [ + ...new Set(activeTurns.map((turn) => turn.channelId)), + ]; return awaitLiveSwitchOutcome({ - channelCount: channelIds.length, + channelIds, modelId, subscribe: (listener) => subscribeControlResults(agent.pubkey, listener), - sendSwitches: async () => { + sendSwitches: async (requestId) => { await Promise.all( channelIds.map((channelId) => - switchManagedAgentModel(agent.pubkey, channelId, modelId), + switchManagedAgentModel( + agent.pubkey, + channelId, + modelId, + requestId, + ), ), ); }, - // No reply in time: treat as sent. The override still rides the - // requeued/next session; we just can't confirm synchronously. + // No terminal proof in time remains visibly pending; it is never + // promoted to a successful switch merely because the request sent. scheduleTimeout: (onTimeout) => { const timeout = window.setTimeout(onTimeout, 8_000); return () => window.clearTimeout(timeout); @@ -147,6 +153,14 @@ export function ModelPicker({ toast.error("That model isn't available for this agent."); return; } + if (outcome === "failed") { + toast.error("The model switch wasn't applied to every active turn."); + return; + } + if (outcome === "pending") { + toast.info("Model switch requested; confirmation is still pending."); + return; + } toast.success("Model switched for this session."); onModelChanged?.(); return; diff --git a/desktop/src/features/communities/useCommunityInit.ts b/desktop/src/features/communities/useCommunityInit.ts index abb49485d2..2eab8b6197 100644 --- a/desktop/src/features/communities/useCommunityInit.ts +++ b/desktop/src/features/communities/useCommunityInit.ts @@ -22,6 +22,7 @@ import { restoreActiveAgentTurnsForCommunity, } from "@/features/agents/activeAgentTurnsStore"; import { resetAgentWorkingSignal } from "@/features/agents/agentWorkingSignal"; +import { resetLiveSwitchOutcomeClaims } from "@/features/agents/lib/liveSwitchOutcome"; import { resetAgentObserverStore } from "@/features/agents/observerRelayStore"; import { resetAvatarPresentations } from "@/features/profile/avatarPresentationStore"; import { resetAvatarProfileSync } from "@/features/profile/avatarProfileSync"; @@ -53,6 +54,7 @@ function resetCommunityState({ resetAgentObserverStore(); resetActiveAgentTurnsStore(); resetAgentWorkingSignal(); + resetLiveSwitchOutcomeClaims(); if (resetAvatarState) { resetAvatarProfileSync(); resetAvatarPresentations(); diff --git a/desktop/src/shared/api/agentControl.ts b/desktop/src/shared/api/agentControl.ts index 677f0ffad4..68c7e97173 100644 --- a/desktop/src/shared/api/agentControl.ts +++ b/desktop/src/shared/api/agentControl.ts @@ -22,10 +22,12 @@ export async function switchManagedAgentModel( pubkey: string, channelId: string, modelId: string, + requestId: string, ): Promise { await sendAgentObserverControl(pubkey, { type: "switch_model", channelId, modelId, + requestId, }); } diff --git a/desktop/src/shared/api/agentControlTypes.ts b/desktop/src/shared/api/agentControlTypes.ts new file mode 100644 index 0000000000..ce2156b8b5 --- /dev/null +++ b/desktop/src/shared/api/agentControlTypes.ts @@ -0,0 +1,33 @@ +/** + * Outcome of a live `switch_model` control frame, surfaced asynchronously via + * the agent's `control_result` observer frame. Busy path: `sent` (cancel + + * requeue on the new model) or `turn_ending` (oneshot already consumed this + * turn). Idle-race path can report `recycling`. Only `switched` proves that + * the requested model reached a fresh channel session; `switch_failed`, + * `unsupported_model`, and `no_active_turn` are terminal failures. + */ +export type SwitchManagedAgentModelStatus = + | "sent" + | "turn_ending" + | "recycling" + | "switched" + | "switch_failed" + | "unsupported_model" + | "no_active_turn"; + +export type ControlResultFrame = { + type: "cancel_turn" | "switch_model"; + status: string; + /** Exact desktop-generated identifier echoed by every switch result. */ + requestId?: string; + modelId?: string; + channelId?: string; + /** ID of the signed kind-24200 relay event carrying this result. */ + relayEventId?: string; + /** Signed Nostr `created_at` for the carrying relay event. */ + relayCreatedAt?: number; + /** Timestamp inside the encrypted, signed observer event payload. */ + observerTimestamp?: string; + /** Process-local sequence inside the encrypted, signed observer payload. */ + observerSeq?: number; +}; diff --git a/desktop/src/shared/api/types.ts b/desktop/src/shared/api/types.ts index 689c400b03..dd4e27a4dc 100644 --- a/desktop/src/shared/api/types.ts +++ b/desktop/src/shared/api/types.ts @@ -479,24 +479,10 @@ export type CancelManagedAgentTurnResult = { status: "sent" | "no_active_turn"; }; -/** - * Outcome of a live `switch_model` control frame, surfaced asynchronously via - * the agent's `control_result` observer frame. Busy path: `sent` (cancel + - * requeue on the new model) or `turn_ending` (oneshot already consumed this - * turn). Idle path: `switched`, `unsupported_model`, or `no_active_turn`. - */ -export type SwitchManagedAgentModelStatus = - | "sent" - | "turn_ending" - | "switched" - | "unsupported_model" - | "no_active_turn"; - -export type ControlResultFrame = { - type: "cancel_turn" | "switch_model"; - status: string; - modelId?: string; -}; +export type { + ControlResultFrame, + SwitchManagedAgentModelStatus, +} from "./agentControlTypes"; export type GitBashPrerequisite = { available: boolean; diff --git a/docs/runbooks/acp-session-process-growth.md b/docs/runbooks/acp-session-process-growth.md new file mode 100644 index 0000000000..408f9f69f8 --- /dev/null +++ b/docs/runbooks/acp-session-process-growth.md @@ -0,0 +1,215 @@ +--- +title: ACP session process growth recovery +verified: 2026-07-28 +review_after: 2026-10-27 +topics: + [ + buzz-acp, + acp, + process-growth, + memory, + session-close, + model-switch, + relay-control, + recovery, + ] +references: + - crates/buzz-acp/src/acp.rs + - crates/buzz-acp/src/pool.rs + - crates/buzz-acp/src/lib.rs + - crates/buzz-acp/src/observer.rs + - crates/buzz-acp/src/queue.rs + - crates/buzz-acp/src/relay.rs + - crates/buzz-acp/tests/startup_signal_cleanup.rs + - desktop/src/features/agents/lib/liveSwitchOutcome.ts + - desktop/src/features/agents/ui/ModelPicker.tsx + - desktop/src/shared/api/agentControl.ts + - desktop/src/shared/api/types.ts + - docs/solutions/2026-07-27-acp-session-retirement-leak.md +--- + +# ACP session process growth + +## Trigger + +Use this runbook when a long-lived `buzz-acp` process stays responsive but its +descendant count, task count, or memory grows across repeated cancellations or +session rotations. + +## Observe without mutation + +Identify the supervised service and record: + +- service state, main PID, task count, and memory; +- total descendants grouped by executable; +- ACP `session/new`, `session/close`, cancellation, and rotation counts; +- unique channel count compared with session-creation count; +- the installed `buzz-acp` revision and adapter version. + +Do not infer a session leak from memory alone. The characteristic signature is +repeated `session/new` for a small, stable channel set without corresponding +retirement, plus retained adapter descendants. + +## Contain + +Stopping or restarting a production service is an operator-gated action. Before +firing that gate, capture the evidence above and verify that any relay, control, +or alternative-agent service is a separate unit and will remain running. + +The live containment gate is exactly: + +> `APPROVE BUZZ CONTAINMENT: stop buzz-acp.service only; leave Buzz relay and buzz-acp-codex.service running; verify memory/process recovery; rollback by starting buzz-acp.service.` + +This runbook records that gate; it does not fire it. No stop, start, restart, or +other live command is authorized until the operator supplies that exact approval. + +For the approved affected unit only: + +1. stop the unit; +2. verify its process group and descendants exit; +3. verify host available memory and swap recover; +4. leave unrelated services untouched. + +Rollback of containment is to start the same unit, but do not restart the known +leaking revision merely to restore nominal availability unless the operator +accepts the resource-exhaustion risk. + +## Remediate + +Build and install a reviewed revision that either positively closes a retired +session or replaces the owning adapter process before local state can forget +it. Preserve the previous binary or package as the rollback artifact. Do not +combine the upgrade with credential, relay, channel-membership, or unit +configuration changes. + +## Validate + +After the operator-approved restart: + +1. confirm the running executable and source revision match the reviewed + artifact; +2. exercise repeated clean cancellation, automatic and owner-requested + rotation, model switching, and a controlled membership-removal cycle; +3. confirm every retired session produces a successful `session/close`, or a + capability-negotiated process replacement after the old adapter's direct + child is reaped and the platform boundary is verified: the original process + group is absent on Unix; Windows currently stops at direct-child reaping + until adapters are assigned to Job Objects; +4. force separate eager-startup and graceful-shutdown cleanup failures. Eager + startup must remain visibly degraded without starting an overlapping owner, + and graceful shutdown must not report a clean exit while ownership of any + original direct child or process group remains unverified; +5. force one bounded-cleanup failure and confirm the exact slot enters a + process-lifetime quarantine: maintenance must not refill it after the + ordinary crash cooldown, and the current supervisor must remain alive in a + visible degraded state instead of exiting into an automatic restart; +6. force one checked-out task panic and confirm it uses that same quarantine. + Unwinding loses the typed child owner, so a panic is never sufficient + evidence for an in-process replacement; +7. confirm a stored numeric process-group ID is only signalled while the live + direct-child identity still matches its spawn PID. After that identity is + cleared, the ID may be used for read-only absence probes but never as a + signal target because the operating system may have reused it; +8. confirm poisoned or still-owning adapters never return to the idle pool; + terminal dead-lettering clears the batch's `required_agent` pin; exact-slot + capacity deferral preserves the existing retry budget; cancelled and + ordinary work share one deterministic global FIFO; and model-switched work + resumes on its exact slot after replacement. Enqueue the same signed event + more than once and confirm each occurrence keeps a distinct retry, + dead-letter, cancelled-provenance, and native-steer identity; one steer + acknowledgement must release or remove only its exact occurrence, and + malformed occurrence vectors must fail closed. Give two withheld occurrences + the same receive timestamp, resolve their acknowledgements in reverse, and + confirm enqueue FIFO still determines dispatch and reply-anchor order; +9. confirm an accepted owner cancel or rotate never replays its in-flight batch + when an error, timeout, or task panic wins the completion race; +10. issue two concurrent switches for the same channel and model. Each desktop + request must generate a distinct 32-character lowercase hexadecimal + `requestId`; Rust must receive it and echo it byte-for-byte on every + immediate or asynchronous terminal result; and the desktop must settle a + request only from its exact ID. Receipt and recycle states remain pending, + and success means the requested model was applied to the target session. + Also force a fresh-catalog rejection and confirm the restored prior model + and its prior request identity are actually applied before the replacement + session becomes available; +11. verify a command-shaped relay event is owner-authorized before privileged + classification or replay admission. A control is eligible only for the + exact active subscription and only when its timestamp strictly postdates + the timestamp recorded after that subscription's successful `REQ`; +12. fill the observer replay-dedup set with fresh IDs and confirm it retains all + of them and rejects additional admission until capacity becomes available; + it must not evict a still-fresh ID and reopen its replay window; +13. verify signature, shape, owner, subscription, and freshness failures are + rejected before bounded control enqueue. When the valid-control queue is + full, backpressure must fail closed rather than silently dropping the + control or mutating replay/routing state. Saturate ordinary telemetry and + confirm observer control results still drain first from their protected + FIFO without starving relay reads or pings. Confirm only `accepted=true` + acknowledges a durable frame, exact rate-limit rejections requeue with + pacing no earlier than an advertised sixty-second reset, the ninety-second + confirmation window covers that reset plus positive jitter, one failed + complete publication gets one bounded ledger replay, terminal denials are + surfaced without retry spin, and bounded loss remains visible. Fill the + retained terminal ledger to 1,024 results and verify the newest overflow is + rejected without evicting older proof and forces failed shutdown + verification; +14. sample the service cgroup's full descendant count and memory over a duration + longer than the original reproduction window; process-group proof cannot + detect a descendant that deliberately escaped with `setsid(2)`; +15. verify completed replacement tasks do not accumulate during sustained relay + traffic or a quiet interval. Pre-loop and automatic-exit cleanup failures + must remain alive in non-spawning quarantine until explicit shutdown. + Graceful shutdown must cooperatively cancel checked-out prompt and heartbeat + work, recover each typed adapter owner, and perform bounded reap/probe. Unix + must prove process-group absence; Windows must prove direct-child reaping; + platforms without either declared boundary must not report cleanup as + verified; +16. send a second SIGINT or SIGTERM while graceful cleanup is pending and + confirm it enters the bounded hard-cleanup path. Repeated startup/shutdown + must leave no orphaned signal-handler tasks. Block `session/new` and the + configured initial prompt in separate trials; the first must return the + typed owner for process recycling, and the second must cancel and retire + the partial session without waiting for the turn timeout; +17. queue at least 48 ordinary observer frames, then emit a terminal control + result while replay and live chunk coalescing are active. The terminal + result must interrupt ordinary pacing and arrive inside the desktop + request timeout. Saturate the protected relay lane separately: already + admitted terminal proofs must remain in FIFO order, the newest overflow + must be rejected visibly before any socket write. Submit the same signed + terminal event through two publisher clones and across disconnected/failed + replay paths; the duplicate must not queue or write, and the first + confirmation waiter must remain authoritative. Terminal publication must + remain pending after the local socket write until the exact relay event + receives `OK accepted=true`. During graceful shutdown, a missing + acknowledgement, priority-lane lag, denial, or bounded publisher timeout + must make shutdown verification fail rather than report clean delivery; +18. block child stdin while cancellation cleanup owes both a pending permission + response and `session/cancel`; both writes and response drain must remain + inside one shared grace deadline. Trigger a control immediately after + `session/new` succeeds but before the configured initial prompt is first + polled; Buzz must close the partial idle session and must preserve the batch + without retry/dead-letter accounting while recycling uncertain setup + ownership; +19. verify unrelated relay and alternative-agent services stayed healthy. + +Do not close the runtime incident from a green build or successful restart +alone. Closure requires a bounded post-restart process and memory trend. + +If a slot is quarantined, first prove the prior adapter and its full service +cgroup descendant tree are absent. Restarting the reviewed unit is the recovery +action and remains operator-gated; a cooldown, heartbeat, or local PID probe is +not a substitute for that proof. + +## Roll back + +If the reviewed revision causes a regression: + +1. stop only the affected unit; +2. restore the preserved prior artifact; +3. start the unit; +4. verify service health and message flow; +5. continue process and memory monitoring because rollback reintroduces the + original leak risk. + +Record the rollback and keep containment available until a corrected revision +passes the same validation. diff --git a/docs/solutions/2026-07-27-acp-session-retirement-leak.md b/docs/solutions/2026-07-27-acp-session-retirement-leak.md new file mode 100644 index 0000000000..87a168e1a6 --- /dev/null +++ b/docs/solutions/2026-07-27-acp-session-retirement-leak.md @@ -0,0 +1,310 @@ +--- +title: Retire ACP sessions before releasing adapter ownership +date: 2026-07-27 +category: docs/solutions/incidents +module: buzz-acp +problem_type: incident +component: session-lifecycle +severity: high +applies_when: + - buzz-acp remains active across repeated cancellation or session rotation + - an adapter retains per-session subprocesses until explicit teardown + - local routing state releases a session while its adapter process still owns it +symptoms: + - buzz-acp stays healthy while its descendant count and memory keep growing + - repeated session/new calls serve a small stable set of channels + - retired Claude and MCP subprocess trees remain below the adapter process +root_cause: logic_error +resolution_type: code_fix +related_components: + - process-lifecycle + - work-queue + - model-switch + - relay-control +tags: + - buzz-acp + - acp + - session-lifecycle + - process-leak + - resource-exhaustion +verified: 2026-07-28 +review_after: 2026-10-27 +topics: [buzz-acp, acp, session-lifecycle, process-leak, resource-exhaustion] +references: + - crates/buzz-acp/src/acp.rs + - crates/buzz-acp/src/pool.rs + - crates/buzz-acp/src/lib.rs + - crates/buzz-acp/src/observer.rs + - crates/buzz-acp/src/queue.rs + - crates/buzz-acp/src/relay.rs + - crates/buzz-acp/tests/startup_signal_cleanup.rs + - desktop/src/features/agents/lib/liveSwitchOutcome.ts + - desktop/src/features/agents/ui/ModelPicker.tsx + - desktop/src/shared/api/agentControl.ts + - desktop/src/shared/api/types.ts + - docs/runbooks/acp-session-process-growth.md +--- + +# Retired ACP sessions kept adapter subprocess trees alive + +## What happened + +A long-lived `buzz-acp` deployment accumulated more than 1,800 descendant +processes and more than 70 GB of memory. The service had created 152 ACP +sessions for only two channels, including 150 sessions for one channel, while +its log recorded hundreds of control cancellations. + +The adapter process stayed connected and continued accepting new sessions, so +ordinary health checks remained green while retired Claude and MCP subprocess +trees accumulated underneath it. + +## Root cause + +Buzz removed a session ID from its local routing state after a clean +control-signal cancellation or automatic rotation, but did not send the ACP +`session/close` request first. Local invalidation therefore abandoned the +adapter-owned session rather than retiring it. For adapters that keep a query +and its MCP subprocesses alive until explicit teardown, every later session +created another retained process tree. + +## Resolution + +The ACP client now negotiates and exposes a bounded `session/close` request. +Buzz sends it only when the adapter advertises the optional capability. +Otherwise Buzz retires ownership by shutting down and replacing the adapter +process group. + +The lifecycle now covers: + +- clean control cancellation and automatic turn-count/token rotation; +- rotate/model commands that race with natural prompt completion; +- accepted owner cancel/rotate commands whose drop disposition must dominate a + simultaneous prompt error, timeout, or task panic; +- idle rotate/model commands and channel-membership removal; +- checked-out adapters that return after their channel was removed; +- partial sessions whose Goose system-prompt or initial-message setup fails. + +Main-loop-only paths claim the owning adapter and schedule an asynchronous +process recycle rather than deleting local state. Compatibility recycles are +rate-reserved per slot, do not consume crash budget, and preserve live model +intent. Exact-slot queue affinity holds a switched batch until that same slot +has a fresh session; completed replacement tasks are reaped under both busy and +quiet traffic. + +Close and cancellation cleanup failures are fail-closed. Buzz preserves +retryable work, emits a redacted cleanup error, and replaces the process group. +Replacement itself is also fail-closed: no new adapter starts unless bounded +shutdown verifies the platform's explicit containment boundary. Unix requires +both direct-child reaping and proof that the original process group is absent. +Windows currently requires verified direct-child reaping; adapter descendants +remain a documented residual risk until Buzz places adapters in Job Objects. +Other platforms without a declared containment boundary fail closed. When the +required proof is unavailable, the exact slot enters a process-lifetime +quarantine that maintenance and the ordinary crash cooldown cannot bypass. The +eager-startup path stays alive in a visible degraded state rather than exiting +into an automatic service restart, and graceful shutdown never reports clean +completion while the platform-supported ownership boundary remains unverified. +A checked-out task panic uses the same quarantine because unwinding destroys +the only typed child owner and therefore cannot produce positive reaping +evidence. +An automatic-exit path with any unverified cleanup owner remains alive in that +same non-spawning quarantine until an explicit shutdown request arrives; it +cannot exit into a service-manager restart over a possibly live process group. +Graceful shutdown cooperatively cancels checked-out prompt and heartbeat tasks, +recovers each typed adapter owner, and then performs the bounded adapter +shutdown and absence proof. It does not use task abortion as cleanup evidence. +On platforms without the explicit Unix process-group or interim Windows +direct-child boundary, cleanup is unverified and therefore quarantined rather +than inferred from direct-child exit. + +The stored spawn-time process-group ID remains useful for read-only absence +probes, but Buzz signals it only while Tokio still exposes the matching live +direct-child PID. Once that identity is cleared, the numeric ID may have been +reused and is never treated as a safe signal target. + +Malformed NDJSON and other framing failures stop at the first bad nonempty +frame without logging raw agent output. Quarantined slots retain retryable +queued work and model intent without attempting an overlapping refill; work +that exhausts its retry budget is deliberately dead-lettered and its +`required_agent` exact-slot pin is cleared. Capacity deferral for an exact-slot +batch preserves its existing retry budget rather than resetting it. Cancelled +and ordinary work are ordered through one deterministic global FIFO. +Retry membership is bound to a private UUID created for each enqueue +occurrence, not to the signed Nostr event ID. Separately enqueued copies of the +same signed event therefore remain distinct: a fresh copy cannot inherit an +older occurrence's retry budget, dead-letter fate, or cancelled provenance. +Every batch carries position-aligned occurrence IDs, and missing, partial, or +duplicate identity fails closed instead of being reconstructed from event +content. A process-local enqueue ordinal travels with that identity and breaks +equal-`Instant` ties, so independently resolving native-steer acknowledgements +cannot reverse FIFO order. Native-steer pending state and acknowledgements +carry that same private occurrence ID, so one successful or failed steer +cannot release or remove another enqueue occurrence that happens to share its +signed event ID. + +Relay control events are owner-authorized before privileged classification or +replay admission. They must bind to the exact active subscription and strictly +postdate the timestamp recorded after that subscription's successful `REQ`. +Cheap signature, shape, identity, subscription, and freshness validation runs +before bounded enqueue, and a full valid-control queue fails closed instead of +silently dropping a control. Observer replay deduplication retains every fresh +ID; when its bounded set is full it rejects new admission rather than evicting a +still-fresh ID and reopening that replay window. +Observer result delivery reserves a protected control-result FIFO and +acknowledgement window that drain before ordinary telemetry. Terminal model +results therefore bypass telemetry pacing and survive ordinary observer queue +pressure without blocking relay reads or pings. Protected relay saturation +rejects the newest terminal frame before a socket write without evicting an +already-admitted proof. A duplicate signed terminal event is likewise rejected +before queueing or writing, preserving the first publisher's exact confirmation +ownership, and the publisher treats any local priority-lane lag as delivery +uncertainty. +Terminal publication completes only after the exact relay event receives +`OK accepted=true`; a local socket write or queue admission is not proof. An +exact transient rate-limit rejection requeues the same frame with pacing. The +client honors the relay's full sixty-second reset hint, adds only positive +desynchronization, and keeps the exact confirmation owner alive for ninety +seconds. If that complete relay attempt still fails, the process-local +acknowledgement ledger releases the result for one bounded end-to-end replay. +Terminal denials remain counted and surfaced without an unbounded retry loop. +The ledger retains at most 1,024 accepted results, rejects newest overflow +without evicting older proof, and marks delivery verification failed. Graceful +shutdown closes the observer source and waits a bounded interval for the +publisher; a rejection, missing relay acknowledgement, lag, join failure, or +timeout fails shutdown verification rather than silently discarding an +accepted model-switch result. + +Live model switching now carries one exact request identity end to end. The +desktop generates a distinct 32-character lowercase hexadecimal `requestId` for +each action, Rust receives that same ID, and every immediate or asynchronous +terminal result echoes it byte-for-byte. The desktop accepts a result only for +its exact pending ID, so concurrent requests for the same channel and model +cannot claim each other's result. Receipt and recycle scheduling remain pending; +application-class model rejection is terminal, and success requires +channel-scoped proof that the requested model was applied. +If a fresh adapter catalog rejects a pending model, rollback restores the exact +prior model intent and request identity and applies that restored model to the +fresh session before the session is committed. The rejected model is never +reported as active merely because replacement succeeded. + +Automatic failures detected before the main loop enter the same cleanup +quarantine as later ownership failures. Startup signal handlers are removed +when their owner exits, so repeated startup attempts cannot leave orphaned +signal tasks behind. After the first SIGINT or SIGTERM requests graceful +shutdown, both listeners remain active; a repeated signal enters the bounded +hard-cleanup path instead of being swallowed while cleanup is pending. +Control signals also preempt every agent-owning setup await. A preempted +`session/new` is typed as cooperative setup preemption, so its preserved batch +is requeued as cancelled without retry or dead-letter accounting while the +uncertain adapter is recycled. A control that becomes ready after session +creation but before the initial prompt is first polled still closes that +partial idle session, preventing reuse of a session that skipped its configured +initial message. Cancellation cleanup uses one shared grace deadline for +permission-response and `session/cancel` stdin writes as well as response +drain, including when child stdin is backpressured. + +## Regression proof + +The regression suite drives fake ACP adapters over the real NDJSON transport +and asserts: + +1. clean cancellation sends `session/cancel`, drains the prompt response, then + sends `session/close` before local invalidation; +2. automatic rotation sends `session/close` before the next `session/new`; +3. a rejected close preserves the session ID and queued work and enters + process-group replacement; +4. representative adapter rejection and timeout close failures emit one + operator-visible error and never return the poisoned adapter to the idle + pool; +5. a matching JSON-RPC response without either `result` or `error` is not a + positive acknowledgement; +6. unadvertised close support triggers bounded process recycling without + probing the optional method; +7. idle control, membership removal, partial setup, simultaneous-ready model + switching, panic quarantine, and malformed framing preserve ownership and + retryable work at their production seams, while retry exhaustion + dead-letters the batch and clears its `required_agent` exact-slot pin; +8. successfully delivered cancel/rotate controls are recorded by the + supervisor, so an error-result or panic race drops the owner-discarded batch + instead of replaying it; +9. malformed NDJSON cannot be skipped in favor of a later matching response, + and diagnostics contain no raw wire payload; +10. eager startup remains visibly degraded, and graceful shutdown cannot return + clean success, when direct-child reaping plus original-group absence is + unverified; task panic quarantines the slot, and a cleared live child + identity prevents signalling a stored numeric PGID; +11. terminal dead-letter clears `required_agent`, exact-slot capacity deferral + preserves retry accounting, and interleaved ordinary and cancelled batches + retain deterministic global FIFO order; +12. owner authorization precedes privileged relay classification; stale, + pre-subscription, wrong-subscription, misrouted, invalidly signed, and + structurally invalid controls fail before bounded enqueue; queue saturation + fails closed; and fresh replay IDs remain protected at dedup capacity; +13. two concurrent switches for the same channel and model carry distinct + 32-character lowercase hexadecimal request IDs through Rust, every + immediate and asynchronous terminal frame echoes the originating ID, and + neither desktop request can claim the other's result; +14. separately enqueued copies of one signed event retain distinct occurrence + identity across ordinary retry, cancelled deferral, capacity eviction, + active pruning, and native-steer transitions; a steer acknowledgement + releases or removes only its exact occurrence, and malformed identity + vectors fail closed; equal receive timestamps retain enqueue FIFO even when + acknowledgements resolve in reverse order; +15. a fresh-catalog model rejection applies the restored prior model to the + replacement session and preserves its prior request identity before that + session becomes available; +16. observer control results retain protected priority FIFO delivery under + concurrent relay traffic and ordinary telemetry saturation; only + `accepted=true` acknowledges a frame, exact rate-limit rejections requeue + no earlier than the advertised reset, one failed complete publication gets + one bounded ledger replay, retained proof rejects newest at its hard cap, + and terminal denials are visible without retry spin; +17. pre-loop and automatic-exit cleanup failures remain in non-spawning + quarantine until explicit shutdown; graceful shutdown recovers checked-out + owners for bounded reap/probe, Unix requires process-group absence, + Windows requires direct-child reaping, and undeclared platform boundaries + fail closed; +18. startup signal tasks are removed on shutdown, while a second SIGINT or + SIGTERM during graceful cleanup enters the hard-cleanup path; +19. shutdown preempts blocked session setup, context preparation, and initial + prompts instead of waiting for the final prompt boundary; dropped setup + requests force process recycling, while an in-flight initial prompt is + cancelled and its partial session retired before the typed owner returns; +20. a live terminal control result interrupts both captured telemetry replay + pacing and a live coalescer flush, so an ordinary 48-frame backlog cannot + consume the desktop result timeout before the terminal result publishes; +21. setup preemption preserves an exhausted-budget batch exactly once without + charging retry budget, a control ready before the initial-prompt future is + polled closes the newly created idle session, and backpressured cleanup + writes cannot exceed the shared cancellation grace; +22. protected relay saturation rejects newest without evicting admitted + terminal proof or writing the rejected frame; duplicate terminal event IDs + reject before queue/socket side effects without replacing the first + confirmation waiter; terminal publication remains pending after a local + socket write until exact relay `OK accepted=true`, and rejection or + priority-lane loss makes bounded shutdown verification fail closed. + +Focused `buzz-acp` lifecycle and desktop request-correlation regressions exercise +these boundaries. Rust formatting and lint gates, desktop unit and typecheck +gates, diff hygiene, and changed-path secret scanning are part of the source +verification. The complete desktop Tauri Clippy leg remains host-blocked by the +missing `gdk-3.0` development package and is not represented as passing. + +## Prevention + +The regression cases above keep ownership proof, queue identity, model-switch +correlation, and relay admission fail-closed at their production seams. The +linked [ACP session process growth runbook](../runbooks/acp-session-process-growth.md) +defines the separate operator-gated rollout and runtime proof needed to close +the live incident. + +## Remaining boundary + +This is a source remediation, not proof that a running deployment has adopted +it. Runtime closure requires building and installing the reviewed revision, +restarting only the affected service under its normal operator gate, and +proving that descendant-process and memory counts remain bounded under repeated +cancellation, rotation, model switching, and membership churn. Use the linked +runbook. Process-group verification cannot prove termination of descendants +that deliberately escape with `setsid(2)`, so runtime closure must also inspect +the service cgroup's full descendant and memory trend.