From 86d6e1133d4e6ba160e3485cd1a14465245acb88 Mon Sep 17 00:00:00 2001 From: kjgbot Date: Wed, 7 Oct 2026 12:31:44 -0700 Subject: [PATCH 1/9] fix(broker): retry pressured control frames --- CHANGELOG.md | 6 +- crates/broker/src/node_control.rs | 839 ++++++++++++++++++++++++++++-- 2 files changed, 801 insertions(+), 44 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 6056d3697..3748267eb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,7 +5,11 @@ All notable changes to Agent Relay will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). -## [Unreleased] +## [Unreleased - Patch] + +### Fixed + +- Fleet spawns and other safely replayable node-control operations now ride out transient Relaycast `d1_pressure` on the existing WebSocket with bounded exponential backoff, while persistent pressure fails closed with its named root cause. ## [13.1.3] - 2026-10-07 diff --git a/crates/broker/src/node_control.rs b/crates/broker/src/node_control.rs index b6ae4ce16..89f12d647 100644 --- a/crates/broker/src/node_control.rs +++ b/crates/broker/src/node_control.rs @@ -58,6 +58,12 @@ const INITIAL_RECONNECT_DELAY: Duration = Duration::from_secs(1); const MAX_RECONNECT_DELAY: Duration = Duration::from_secs(30); const D1_PRESSURE_RETRY_INITIAL_DELAY: Duration = Duration::from_secs(1); const D1_PRESSURE_RETRY_MAX_DELAY: Duration = Duration::from_secs(8); +// Five retries consume 1 + 2 + 4 + 8 + 8 = 23 seconds. Keep the wall-clock +// budget below the 30-second fleet registration waiter so persistent pressure +// reaches the spawn caller as a named d1_pressure failure, not a generic +// agent_register_timeout. +const CONTROL_D1_PRESSURE_MAX_RETRIES: u32 = 5; +const CONTROL_D1_PRESSURE_RETRY_BUDGET: Duration = Duration::from_secs(25); const REGISTER_AGENT_PENDING_TTL: Duration = Duration::from_secs(300); const RELAYCAST_DEFAULT_BASE_URL: &str = "https://cast.agentrelay.com"; const CREATE_NODE_RETRY_BACKOFFS_MS: [u64; 3] = [200, 400, 800]; @@ -724,8 +730,12 @@ pub(crate) struct AgentRegistrationToken { struct PendingAgentRegistration { isolates_channels: bool, name: String, + request: AgentRegister, reply: oneshot::Sender>, created_at: Instant, + d1_pressure_retry_attempts: u32, + d1_pressure_retry_at: Option, + last_d1_pressure: Option, } #[derive(Debug)] @@ -2669,10 +2679,26 @@ struct PendingInventorySync { superseded: bool, } +struct PendingRetryableControlFrame { + frame: BrokerToRelaycast, + created_at: Instant, + retry_attempts: u32, + retry_at: Option, + last_d1_pressure: Option, +} + +#[derive(Debug, PartialEq, Eq)] +enum ControlPressureRetry { + Unmatched, + Scheduled(Duration), + Exhausted, +} + struct ApplicationLiveness { deadline: Duration, last_acknowledged: Instant, pending_inventory_syncs: VecDeque, + pending_retryable_control_frames: HashMap, ready: bool, } @@ -2682,10 +2708,103 @@ impl ApplicationLiveness { deadline, last_acknowledged: Instant::now(), pending_inventory_syncs: VecDeque::new(), + pending_retryable_control_frames: HashMap::new(), ready: false, } } + fn track_retryable_control_frame(&mut self, id: String, frame: BrokerToRelaycast) { + self.pending_retryable_control_frames.insert( + id, + PendingRetryableControlFrame { + frame, + created_at: Instant::now(), + retry_attempts: 0, + retry_at: None, + last_d1_pressure: None, + }, + ); + } + + fn complete_retryable_control_frame(&mut self, id: &str) { + self.pending_retryable_control_frames.remove(id); + } + + fn schedule_control_d1_pressure_retry( + &mut self, + id: &str, + reason: &str, + now: Instant, + ) -> ControlPressureRetry { + let Some(pending) = self.pending_retryable_control_frames.get_mut(id) else { + return ControlPressureRetry::Unmatched; + }; + pending.last_d1_pressure = Some(reason.to_string()); + let deadline = pending.created_at + CONTROL_D1_PRESSURE_RETRY_BUDGET; + if pending.retry_attempts >= CONTROL_D1_PRESSURE_MAX_RETRIES || now >= deadline { + return ControlPressureRetry::Exhausted; + } + let next_attempt = pending.retry_attempts.saturating_add(1); + let delay = d1_pressure_retry_delay(next_attempt); + if now + .checked_add(delay) + .is_none_or(|retry_at| retry_at > deadline) + { + return ControlPressureRetry::Exhausted; + } + pending.retry_attempts = next_attempt; + pending.retry_at = Some(now + delay); + ControlPressureRetry::Scheduled(delay) + } + + fn next_control_retry_event_delay(&self, now: Instant) -> Option { + self.pending_retryable_control_frames + .values() + .flat_map(|pending| { + [ + pending.retry_at, + Some(pending.created_at + CONTROL_D1_PRESSURE_RETRY_BUDGET), + ] + }) + .flatten() + .map(|event_at| event_at.saturating_duration_since(now)) + .min() + } + + fn expire_retryable_control_frames(&mut self, now: Instant) -> Vec<(String, String)> { + let expired: Vec = self + .pending_retryable_control_frames + .iter() + .filter_map(|(id, pending)| { + (now.saturating_duration_since(pending.created_at) + >= CONTROL_D1_PRESSURE_RETRY_BUDGET) + .then(|| id.clone()) + }) + .collect(); + expired + .into_iter() + .filter_map(|id| { + let pending = self.pending_retryable_control_frames.remove(&id)?; + pending.last_d1_pressure.map(|reason| { + ( + id, + format!("d1_pressure: {reason} (retry budget exhausted)"), + ) + }) + }) + .collect() + } + + fn take_due_control_retry(&mut self, now: Instant) -> Option { + let pending = self + .pending_retryable_control_frames + .values_mut() + .filter(|pending| pending.retry_at.is_some_and(|retry_at| retry_at <= now)) + .min_by_key(|pending| pending.retry_at)?; + pending.retry_at = None; + Some(pending.frame.clone()) + } + fn track_inventory_sync(&mut self, frame: InventorySync) { // A newer authoritative snapshot supersedes every older retry payload. // Keep the probes for correlation/liveness, but never let an old frame @@ -3021,13 +3140,16 @@ where } Continue(()) } - Some(FleetControlCommand::Send(message)) => { + Some(FleetControlCommand::Send(mut message)) => { // A `delivery_ack` is the engine's only evidence that a // frame was consumed. The runtime records the *decision* // to ack before handing it here and cannot wait for the // wire, so the probe learns the outcome at the one place // that knows it. See `NodeDeliveryProbe::record_ack_sent`. let is_ack = matches!(message, BrokerToRelaycast::DeliveryAck(_)); + if let Some((request_id, _)) = prepare_retryable_control_frame(&mut message) { + application_liveness.track_retryable_control_frame(request_id, message.clone()); + } let sent = send_wire(sink, &message).await; if let (true, Some(probe)) = (is_ack, config.probe.as_ref()) { if sent.is_ok() { @@ -3047,7 +3169,12 @@ where } Continue(()) } - Some(FleetControlCommand::RetainedSend(completion)) => { + Some(FleetControlCommand::RetainedSend(mut completion)) => { + if let Some((request_id, _)) = prepare_retryable_control_frame(&mut completion.message) + { + application_liveness + .track_retryable_control_frame(request_id, completion.message.clone()); + } if send_wire(sink, &completion.message).await.is_err() { *retained_completion = Some(completion); return Break(ControlRunResult::Disconnected { @@ -3061,11 +3188,10 @@ where let request_id = format!("agent_deregister_{}", Uuid::new_v4().simple()); request.id = Some(request_id.clone()); pending_deregistrations.retain(|_, pending| !pending.is_closed()); - pending_deregistrations.insert(request_id, reply); - if send_wire(sink, &BrokerToRelaycast::AgentDeregister(request)) - .await - .is_err() - { + pending_deregistrations.insert(request_id.clone(), reply); + let message = BrokerToRelaycast::AgentDeregister(request); + application_liveness.track_retryable_control_frame(request_id, message.clone()); + if send_wire(sink, &message).await.is_err() { return Break(ControlRunResult::Disconnected { application_ready: application_liveness.ready, }); @@ -3078,16 +3204,21 @@ where .clone() .unwrap_or_else(|| format!("agent_register_{}", Uuid::new_v4().simple())); request.id = Some(request_id.clone()); + let wire_request = request.clone(); pending_agent_registrations.insert( request_id, PendingAgentRegistration { isolates_channels: request.auto_join_general == Some(false), name: request.name.clone(), + request, reply, created_at: Instant::now(), + d1_pressure_retry_attempts: 0, + d1_pressure_retry_at: None, + last_d1_pressure: None, }, ); - if send_wire(sink, &BrokerToRelaycast::AgentRegister(request)) + if send_wire(sink, &BrokerToRelaycast::AgentRegister(wire_request)) .await .is_err() { @@ -3344,7 +3475,56 @@ async fn run_connected_once_with_completions( let inventory_retry = tokio::time::sleep(inventory_retry_delay.unwrap_or(application_liveness.deadline)); tokio::pin!(inventory_retry); + let agent_registration_retry_delay = + next_agent_registration_retry_event_delay(&pending_agent_registrations, Instant::now()); + let agent_registration_retry = tokio::time::sleep( + agent_registration_retry_delay.unwrap_or(application_liveness.deadline), + ); + tokio::pin!(agent_registration_retry); + let control_retry_delay = + application_liveness.next_control_retry_event_delay(Instant::now()); + let control_retry = + tokio::time::sleep(control_retry_delay.unwrap_or(application_liveness.deadline)); + tokio::pin!(control_retry); tokio::select! { + _ = &mut control_retry, if control_retry_delay.is_some() => { + let now = Instant::now(); + for (id, reason) in application_liveness.expire_retryable_control_frames(now) { + if let Some(pending) = pending_deregistrations.remove(&id) { + let _ = pending.send(Err(reason.clone())); + } + if id.starts_with(crate::runtime::task_request_prefix()) { + let _ = event_tx.send(FleetControlEvent::Message(RelaycastToBroker::Error( + crate::fleet_wire::Error { + v: FLEET_WIRE_VERSION, + id, + ok: false, + code: "d1_pressure".to_string(), + message: reason, + }, + ))).await; + } + } + if let Some(frame) = application_liveness.take_due_control_retry(now) { + if send_wire(&mut sink, &frame).await.is_err() { + drain_agent_registrations(&mut pending_agent_registrations, "node_control_disconnected"); + return ControlRunResult::Disconnected { application_ready: application_liveness.ready }; + } + } + } + _ = &mut agent_registration_retry, if agent_registration_retry_delay.is_some() => { + let now = Instant::now(); + expire_d1_pressure_agent_registrations(&mut pending_agent_registrations, now); + if let Some(frame) = take_due_agent_registration_retry( + &mut pending_agent_registrations, + now, + ) { + if send_wire(&mut sink, &BrokerToRelaycast::AgentRegister(frame)).await.is_err() { + drain_agent_registrations(&mut pending_agent_registrations, "node_control_disconnected"); + return ControlRunResult::Disconnected { application_ready: application_liveness.ready }; + } + } + } _ = &mut inventory_retry, if inventory_retry_delay.is_some() => { let Some(frame) = application_liveness.take_due_retry(Instant::now()) else { continue; @@ -3547,6 +3727,7 @@ where } match frame { RelaycastToBroker::Reply(reply) => { + application_liveness.complete_retryable_control_frame(&reply.id); if reply.id.starts_with(crate::runtime::task_request_prefix()) { return event_tx .send(FleetControlEvent::Message(RelaycastToBroker::Reply( @@ -3596,7 +3777,62 @@ where } } } - RelaycastToBroker::Error(error) => { + RelaycastToBroker::Error(mut error) => { + if error.code == "d1_pressure" { + match application_liveness.schedule_control_d1_pressure_retry( + &error.id, + &error.message, + Instant::now(), + ) { + ControlPressureRetry::Scheduled(delay) => { + tracing::warn!( + target = "relay_broker::fleet", + node_id, + id = %error.id, + retry_delay_ms = delay.as_millis(), + "retryable d1_pressure on a control frame; retry scheduled on the live control socket" + ); + return true; + } + ControlPressureRetry::Exhausted => { + application_liveness + .complete_retryable_control_frame(&error.id); + let reason = format!( + "d1_pressure: {} (retry budget exhausted)", + error.message, + ); + if let Some(pending) = + pending_deregistrations.remove(&error.id) + { + let _ = pending.send(Err(reason)); + return true; + } + if error + .id + .starts_with(crate::runtime::task_request_prefix()) + { + error.code = "d1_pressure".to_string(); + error.message = reason; + return event_tx + .send(FleetControlEvent::Message( + RelaycastToBroker::Error(error), + )) + .await + .is_ok(); + } + tracing::error!( + target = "relay_broker::fleet", + node_id, + id = %error.id, + reason, + "control frame d1_pressure retry budget exhausted" + ); + return true; + } + ControlPressureRetry::Unmatched => {} + } + } + application_liveness.complete_retryable_control_frame(&error.id); if error.id.starts_with(crate::runtime::task_request_prefix()) { return event_tx .send(FleetControlEvent::Message(RelaycastToBroker::Error( @@ -3627,12 +3863,39 @@ where id = %error.id, "engine rejected a node control frame" ); - fail_agent_registration( - &error.id, - format!("{}: {}", error.code, error.message), - pending_agent_registrations, - ); if error.code == "d1_pressure" { + let reason = error.message.clone(); + match schedule_agent_registration_d1_pressure_retry( + &error.id, + &reason, + pending_agent_registrations, + Instant::now(), + ) { + AgentRegistrationPressureRetry::Scheduled(delay) => { + tracing::warn!( + target = "relay_broker::fleet", + node_id, + id = %error.id, + retry_delay_ms = delay.as_millis(), + "retryable d1_pressure on agent registration; retry scheduled on the live control socket" + ); + return true; + } + AgentRegistrationPressureRetry::Exhausted => { + tracing::error!( + target = "relay_broker::fleet", + node_id, + id = %error.id, + "agent registration d1_pressure retry budget exhausted" + ); + fail_agent_registration_d1_pressure_exhausted( + &error.id, + pending_agent_registrations, + ); + return true; + } + AgentRegistrationPressureRetry::Unmatched => {} + } if application_liveness.discard_superseded(&error.id) { tracing::debug!( target = "relay_broker::fleet", @@ -3655,6 +3918,11 @@ where return true; } } + fail_agent_registration( + &error.id, + format!("{}: {}", error.code, error.message), + pending_agent_registrations, + ); let rejected_liveness_probe = application_liveness.reject(&error.id); // A reply proves the application is responsive, but rejecting // the authoritative inventory probe means the control session @@ -3812,6 +4080,108 @@ fn fail_unsupported_channel_isolation( } } +#[derive(Debug, PartialEq, Eq)] +enum AgentRegistrationPressureRetry { + Unmatched, + Scheduled(Duration), + Exhausted, +} + +fn schedule_agent_registration_d1_pressure_retry( + id: &str, + reason: &str, + pending_agent_registrations: &mut HashMap, + now: Instant, +) -> AgentRegistrationPressureRetry { + let Some(pending) = pending_agent_registrations.get_mut(id) else { + return AgentRegistrationPressureRetry::Unmatched; + }; + pending.last_d1_pressure = Some(reason.to_string()); + let deadline = pending.created_at + CONTROL_D1_PRESSURE_RETRY_BUDGET; + if pending.d1_pressure_retry_attempts >= CONTROL_D1_PRESSURE_MAX_RETRIES || now >= deadline { + return AgentRegistrationPressureRetry::Exhausted; + } + + let next_attempt = pending.d1_pressure_retry_attempts.saturating_add(1); + let delay = d1_pressure_retry_delay(next_attempt); + if now + .checked_add(delay) + .is_none_or(|retry_at| retry_at > deadline) + { + return AgentRegistrationPressureRetry::Exhausted; + } + pending.d1_pressure_retry_attempts = next_attempt; + pending.d1_pressure_retry_at = Some(now + delay); + AgentRegistrationPressureRetry::Scheduled(delay) +} + +fn next_agent_registration_retry_event_delay( + pending_agent_registrations: &HashMap, + now: Instant, +) -> Option { + pending_agent_registrations + .values() + .filter(|pending| pending.last_d1_pressure.is_some()) + .flat_map(|pending| { + [ + pending.d1_pressure_retry_at, + Some(pending.created_at + CONTROL_D1_PRESSURE_RETRY_BUDGET), + ] + }) + .flatten() + .map(|event_at| event_at.saturating_duration_since(now)) + .min() +} + +fn expire_d1_pressure_agent_registrations( + pending_agent_registrations: &mut HashMap, + now: Instant, +) { + let expired: Vec = pending_agent_registrations + .iter() + .filter_map(|(id, pending)| { + (pending.last_d1_pressure.is_some() + && now.saturating_duration_since(pending.created_at) + >= CONTROL_D1_PRESSURE_RETRY_BUDGET) + .then(|| id.clone()) + }) + .collect(); + for id in expired { + fail_agent_registration_d1_pressure_exhausted(&id, pending_agent_registrations); + } +} + +fn fail_agent_registration_d1_pressure_exhausted( + id: &str, + pending_agent_registrations: &mut HashMap, +) { + let Some(pending) = pending_agent_registrations.remove(id) else { + return; + }; + let reason = pending + .last_d1_pressure + .unwrap_or_else(|| "Node liveness retry pending".to_string()); + let _ = pending.reply.send(Err(format!( + "d1_pressure: {reason} (retry budget exhausted)" + ))); +} + +fn take_due_agent_registration_retry( + pending_agent_registrations: &mut HashMap, + now: Instant, +) -> Option { + let pending = pending_agent_registrations + .values_mut() + .filter(|pending| { + pending + .d1_pressure_retry_at + .is_some_and(|retry_at| retry_at <= now) + }) + .min_by_key(|pending| pending.d1_pressure_retry_at)?; + pending.d1_pressure_retry_at = None; + Some(pending.request.clone()) +} + fn expire_agent_registrations( pending_agent_registrations: &mut HashMap, now: Instant, @@ -3884,6 +4254,48 @@ fn fail_deferred_commands( } } +/// Gives every safely replayable steady-state control frame a stable +/// correlation id. The engine's `d1_pressure` contract guarantees that the +/// rejected operation was not applied, so exact replay is safe for: +/// +/// - `agent.deregister` (engine teardown is idempotent), +/// - cumulative `delivery.ack`, and +/// - `action.accept` / `action.result` (invocation-ledger idempotency covers +/// spawn results, release results, and task receipts). +/// +/// `node.register`, `inventory.sync`, and `agent.register` have dedicated +/// pending state because their successful replies open liveness or return a +/// token. Heartbeats intentionally are not retried: a newer periodic snapshot +/// supersedes them. No separate channel-join frame exists; the only join sent +/// by this broker is part of the exact `agent.register` frame. +fn prepare_retryable_control_frame( + message: &mut BrokerToRelaycast, +) -> Option<(String, &'static str)> { + let (id, kind) = match message { + BrokerToRelaycast::AgentDeregister(request) => ( + request + .id + .get_or_insert_with(|| format!("agent_deregister_{}", Uuid::new_v4().simple())), + "agent.deregister", + ), + BrokerToRelaycast::DeliveryAck(request) => ( + request + .id + .get_or_insert_with(|| format!("delivery_ack_{}", Uuid::new_v4().simple())), + "delivery.ack", + ), + BrokerToRelaycast::ActionResult(request) => ( + request + .id + .get_or_insert_with(|| format!("action_result_{}", Uuid::new_v4().simple())), + "action.result", + ), + BrokerToRelaycast::ActionAccept(request) => (&mut request.id, "action.accept"), + _ => return None, + }; + Some((id.clone(), kind)) +} + async fn send_wire(sink: &mut S, message: &BrokerToRelaycast) -> Result<()> where S: Sink + Unpin, @@ -4140,6 +4552,33 @@ mod tests { } } + fn test_pending_agent_registration( + name: &str, + isolates_channels: bool, + id: &str, + reply: oneshot::Sender>, + created_at: Instant, + ) -> PendingAgentRegistration { + PendingAgentRegistration { + isolates_channels, + name: name.to_string(), + request: AgentRegister { + auto_join_general: isolates_channels.then_some(false), + v: FLEET_WIRE_VERSION, + id: Some(id.to_string()), + name: name.to_string(), + invocation_id: None, + session_ref: None, + resumable: None, + }, + reply, + created_at, + d1_pressure_retry_attempts: 0, + d1_pressure_retry_at: None, + last_d1_pressure: None, + } + } + async fn read_http_request(socket: &mut TcpStream) -> String { let mut buffer = [0_u8; 8192]; let read = socket.read(&mut buffer).await.unwrap(); @@ -5254,21 +5693,23 @@ mod tests { let mut pending = HashMap::from([ ( "isolated".to_string(), - PendingAgentRegistration { - isolates_channels: true, - name: "isolated".to_string(), - reply: isolated_tx, - created_at: Instant::now(), - }, + test_pending_agent_registration( + "isolated", + true, + "isolated", + isolated_tx, + Instant::now(), + ), ), ( "legacy".to_string(), - PendingAgentRegistration { - isolates_channels: false, - name: "legacy".to_string(), - reply: legacy_tx, - created_at: Instant::now(), - }, + test_pending_agent_registration( + "legacy", + false, + "legacy", + legacy_tx, + Instant::now(), + ), ), ]); let (events, _) = mpsc::channel(1); @@ -5449,6 +5890,63 @@ mod tests { assert!(!liveness.ready); } + #[test] + fn d1_pressure_safe_control_frame_audit_assigns_ids_and_replays_exact_frames() { + let frames = vec![ + BrokerToRelaycast::AgentDeregister(AgentDeregister { + v: FLEET_WIRE_VERSION, + id: None, + agent_id: "agt-1".to_string(), + name: Some("agent-a".to_string()), + }), + delivery_ack("agent-a", 7), + BrokerToRelaycast::ActionResult(ActionResult { + v: FLEET_WIRE_VERSION, + id: None, + invocation_id: "inv-1".to_string(), + result: ActionResultPayload::Output(ActionResultOutput { + output: json!({"ok": true}), + }), + task: None, + }), + BrokerToRelaycast::ActionAccept(crate::fleet_wire::ActionAccept { + v: FLEET_WIRE_VERSION, + id: "task_request_1".to_string(), + invocation_id: "inv-task".to_string(), + execution_id: "exec-1".to_string(), + worker_generation: "generation-1".to_string(), + }), + ]; + + for mut frame in frames { + let (id, _) = prepare_retryable_control_frame(&mut frame) + .expect("audited frame must be retryable"); + let mut state = ApplicationLiveness::new(Duration::from_secs(30)); + state.track_retryable_control_frame(id.clone(), frame.clone()); + let now = Instant::now(); + assert_eq!( + state.schedule_control_d1_pressure_retry(&id, "Node liveness retry pending", now,), + ControlPressureRetry::Scheduled(Duration::from_secs(1)) + ); + assert_eq!( + state.take_due_control_retry(now + Duration::from_secs(1)), + Some(frame), + "the retry must preserve every field and correlation id" + ); + } + + let mut heartbeat = BrokerToRelaycast::NodeHeartbeat( + FleetLoadSnapshot::default().heartbeat(&build_node_register( + &test_manifest(), + "node-test", + "host-test", + "broker/test", + None, + )), + ); + assert_eq!(prepare_retryable_control_frame(&mut heartbeat), None); + } + #[test] fn newer_inventory_snapshot_cancels_an_older_retry_payload() { let mut liveness = ApplicationLiveness::new(Duration::from_secs(120)); @@ -5530,12 +6028,13 @@ mod tests { let (reply_tx, mut reply_rx) = oneshot::channel(); let mut pending = HashMap::from([( "agent_register_1".to_string(), - PendingAgentRegistration { - isolates_channels: false, - name: "agent-a".to_string(), - reply: reply_tx, + test_pending_agent_registration( + "agent-a", + false, + "agent_register_1", + reply_tx, created_at, - }, + ), )]); expire_agent_registrations( @@ -5561,12 +6060,13 @@ mod tests { let (reply_tx, reply_rx) = oneshot::channel(); let mut pending = HashMap::from([( "agent_register_req".to_string(), - PendingAgentRegistration { - isolates_channels: false, - name: "agent-a".to_string(), - reply: reply_tx, - created_at: Instant::now(), - }, + test_pending_agent_registration( + "agent-a", + false, + "agent_register_req", + reply_tx, + Instant::now(), + ), )]); let mut sink = futures_util::sink::drain(); let reply = crate::fleet_wire::Reply { @@ -5605,12 +6105,13 @@ mod tests { let (reply_tx, mut reply_rx) = oneshot::channel(); let mut pending = HashMap::from([( "agent_register_req".to_string(), - PendingAgentRegistration { - isolates_channels: false, - name: "agent-a".to_string(), - reply: reply_tx, - created_at: Instant::now(), - }, + test_pending_agent_registration( + "agent-a", + false, + "agent_register_req", + reply_tx, + Instant::now(), + ), )]); let mut sink = futures_util::sink::drain(); let reply = crate::fleet_wire::Reply { @@ -6230,6 +6731,258 @@ mod tests { let _ = command_tx.send(FleetControlCommand::Shutdown).await; } + #[tokio::test] + async fn node_control_agent_register_retries_d1_pressure_then_succeeds() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let ws_url = format!("ws://{}/v1/node/ws", listener.local_addr().unwrap()); + let (command_tx, command_rx) = mpsc::channel(32); + let (event_tx, mut event_rx) = mpsc::channel(32); + + tokio::spawn(run_node_control_client( + FleetControlConfig { + ws_url, + node_token: Some("nt_test".to_string()), + node_id: "node-test".to_string(), + node_name: "host-test".to_string(), + broker_version: "broker/test".to_string(), + token_minter: None, + session_token: None, + read_idle_timeout: None, + probe: None, + terminal_reconnect_tx: None, + }, + command_rx, + event_tx, + )); + + let server = tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + let mut ws = accept_async(stream).await.unwrap(); + + assert!(matches!( + next_node_to_server(&mut ws).await, + BrokerToRelaycast::NodeRegister(_) + )); + assert!(matches!( + next_node_to_server(&mut ws).await, + BrokerToRelaycast::InventorySync(_) + )); + assert!(matches!( + next_node_to_server(&mut ws).await, + BrokerToRelaycast::NodeHeartbeat(_) + )); + + let initial = match next_non_heartbeat_node_to_server(&mut ws).await { + BrokerToRelaycast::AgentRegister(request) => request, + other => panic!("expected agent.register, got {other:?}"), + }; + let request_id = initial.id.clone().expect("agent.register id"); + + for (attempt, no_retry_before, retry_within) in [ + (1, Duration::from_millis(500), Duration::from_secs(2)), + (2, Duration::from_millis(1_500), Duration::from_secs(3)), + ] { + ws.send(Message::Text( + json!({ + "v": 1, + "id": request_id, + "type": "error", + "ok": false, + "code": "d1_pressure", + "message": "Node liveness retry pending" + }) + .to_string(), + )) + .await + .unwrap(); + assert!( + tokio::time::timeout( + no_retry_before, + next_non_heartbeat_node_to_server(&mut ws), + ) + .await + .is_err(), + "agent.register retry {attempt} must honor exponential backoff" + ); + let retried = + tokio::time::timeout(retry_within, next_non_heartbeat_node_to_server(&mut ws)) + .await + .expect( + "retryable agent.register d1_pressure should retry on the live socket", + ); + assert_eq!(retried, BrokerToRelaycast::AgentRegister(initial.clone())); + } + + ws.send(Message::Text( + serde_json::to_string(&RelaycastToBroker::Reply(crate::fleet_wire::Reply { + v: FLEET_WIRE_VERSION, + id: request_id, + ok: true, + data: json!({ + "name": "agent-a", + "agent_id": "agt-1", + "token": "at_test" + }), + })) + .unwrap(), + )) + .await + .unwrap(); + }); + + command_tx + .send(FleetControlCommand::RegisterNode { + manifest: test_manifest(), + resume_cursor: None, + }) + .await + .unwrap(); + assert_eq!(event_rx.recv().await.unwrap(), FleetControlEvent::Connected); + + let (reply_tx, reply_rx) = oneshot::channel(); + command_tx + .send(FleetControlCommand::RegisterAgent { + request: AgentRegister { + auto_join_general: Some(false), + v: FLEET_WIRE_VERSION, + id: None, + name: "agent-a".to_string(), + invocation_id: Some("inv-1".to_string()), + session_ref: Some("session-1".to_string()), + resumable: Some(true), + }, + reply: reply_tx, + }) + .await + .unwrap(); + + let token = tokio::time::timeout(Duration::from_secs(8), reply_rx) + .await + .expect("agent.register retry budget elapsed") + .unwrap() + .expect("transient d1_pressure must not fail the spawn"); + assert_eq!(token.agent_id, "agt-1"); + assert_eq!(token.token, "at_test"); + + tokio::time::timeout(Duration::from_secs(8), server) + .await + .unwrap() + .unwrap(); + let _ = command_tx.send(FleetControlCommand::Shutdown).await; + } + + #[tokio::test(start_paused = true)] + async fn node_control_agent_register_persistent_d1_pressure_fails_closed_with_named_error() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let ws_url = format!("ws://{}/v1/node/ws", listener.local_addr().unwrap()); + let (command_tx, command_rx) = mpsc::channel(32); + let (event_tx, mut event_rx) = mpsc::channel(32); + + tokio::spawn(run_node_control_client( + FleetControlConfig { + ws_url, + node_token: Some("nt_test".to_string()), + node_id: "node-test".to_string(), + node_name: "host-test".to_string(), + broker_version: "broker/test".to_string(), + token_minter: None, + session_token: None, + read_idle_timeout: None, + probe: None, + terminal_reconnect_tx: None, + }, + command_rx, + event_tx, + )); + + let server = tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + let mut ws = accept_async(stream).await.unwrap(); + assert!(matches!( + next_node_to_server(&mut ws).await, + BrokerToRelaycast::NodeRegister(_) + )); + assert!(matches!( + next_node_to_server(&mut ws).await, + BrokerToRelaycast::InventorySync(_) + )); + assert!(matches!( + next_node_to_server(&mut ws).await, + BrokerToRelaycast::NodeHeartbeat(_) + )); + + let mut initial = None; + for _ in 0..=CONTROL_D1_PRESSURE_MAX_RETRIES { + let request = match next_non_heartbeat_node_to_server(&mut ws).await { + BrokerToRelaycast::AgentRegister(request) => request, + other => panic!("expected agent.register, got {other:?}"), + }; + if let Some(initial) = initial.as_ref() { + assert_eq!( + &request, initial, + "every retry must preserve the exact frame" + ); + } else { + initial = Some(request.clone()); + } + ws.send(Message::Text( + json!({ + "v": 1, + "id": request.id, + "type": "error", + "ok": false, + "code": "d1_pressure", + "message": "Node liveness retry pending" + }) + .to_string(), + )) + .await + .unwrap(); + } + }); + + command_tx + .send(FleetControlCommand::RegisterNode { + manifest: test_manifest(), + resume_cursor: None, + }) + .await + .unwrap(); + assert_eq!(event_rx.recv().await.unwrap(), FleetControlEvent::Connected); + + let (reply_tx, reply_rx) = oneshot::channel(); + command_tx + .send(FleetControlCommand::RegisterAgent { + request: AgentRegister { + auto_join_general: Some(false), + v: FLEET_WIRE_VERSION, + id: None, + name: "agent-a".to_string(), + invocation_id: Some("inv-1".to_string()), + session_ref: Some("session-1".to_string()), + resumable: Some(true), + }, + reply: reply_tx, + }) + .await + .unwrap(); + + let error = tokio::time::timeout(Duration::from_secs(30), reply_rx) + .await + .expect("persistent pressure must fail inside the registration deadline") + .unwrap() + .expect_err("persistent d1_pressure must fail closed"); + assert!( + error.starts_with("d1_pressure:"), + "spawn caller must see the retryable root cause, got {error}" + ); + tokio::time::timeout(Duration::from_secs(1), server) + .await + .unwrap() + .unwrap(); + let _ = command_tx.send(FleetControlCommand::Shutdown).await; + } + #[tokio::test] async fn node_control_client_mints_initial_token_when_none_supplied() { // Broker startup no longer mints the node token on the API-readiness From 55f413b2d3bdb3ffaa8cf108f5c3bd4a845a5efb Mon Sep 17 00:00:00 2001 From: kjgbot Date: Wed, 7 Oct 2026 12:39:13 -0700 Subject: [PATCH 2/9] fix(broker): retain pressured registration cleanup --- crates/broker/src/node_control.rs | 178 +++++++++++++++++++----------- 1 file changed, 116 insertions(+), 62 deletions(-) diff --git a/crates/broker/src/node_control.rs b/crates/broker/src/node_control.rs index 89f12d647..5081e83f4 100644 --- a/crates/broker/src/node_control.rs +++ b/crates/broker/src/node_control.rs @@ -731,7 +731,7 @@ struct PendingAgentRegistration { isolates_channels: bool, name: String, request: AgentRegister, - reply: oneshot::Sender>, + reply: Option>>, created_at: Instant, d1_pressure_retry_attempts: u32, d1_pressure_retry_at: Option, @@ -3211,7 +3211,7 @@ where isolates_channels: request.auto_join_general == Some(false), name: request.name.clone(), request, - reply, + reply: Some(reply), created_at: Instant::now(), d1_pressure_retry_attempts: 0, d1_pressure_retry_at: None, @@ -3997,9 +3997,9 @@ where let data = match data { Some(data) => data, None => { - let _ = pending - .reply - .send(Err("invalid_agent_register_reply_data".to_string())); + if let Some(reply) = pending.reply { + let _ = reply.send(Err("invalid_agent_register_reply_data".to_string())); + } return true; } }; @@ -4009,30 +4009,32 @@ where token: data.token, delivery_ack_seq: data.delivery_ack_seq, }; - match pending.reply.send(Ok(token.clone())) { - Ok(()) => true, - Err(Ok(token)) => { - tracing::warn!( - target = "relay_broker::fleet", - id = %request_id, - name = %pending.name, - agent_id = %token.agent_id, - "late agent.register success after caller stopped waiting; sending compensating agent.deregister" - ); - send_wire( - sink, - &BrokerToRelaycast::AgentDeregister(AgentDeregister { - v: FLEET_WIRE_VERSION, - id: Some(request_id), - agent_id: token.agent_id, - name: Some(pending.name), - }), - ) - .await - .is_ok() - } - Err(Err(_)) => true, - } + let late_token = match pending.reply { + Some(reply) => match reply.send(Ok(token.clone())) { + Ok(()) => return true, + Err(Ok(token)) => token, + Err(Err(_)) => return true, + }, + None => token, + }; + tracing::warn!( + target = "relay_broker::fleet", + id = %request_id, + name = %pending.name, + agent_id = %late_token.agent_id, + "late agent.register success after caller stopped waiting; sending compensating agent.deregister" + ); + send_wire( + sink, + &BrokerToRelaycast::AgentDeregister(AgentDeregister { + v: FLEET_WIRE_VERSION, + id: Some(request_id), + agent_id: late_token.agent_id, + name: Some(pending.name), + }), + ) + .await + .is_ok() } fn fail_agent_registration( @@ -4041,7 +4043,9 @@ fn fail_agent_registration( pending_agent_registrations: &mut HashMap, ) { if let Some(pending) = pending_agent_registrations.remove(id) { - let _ = pending.reply.send(Err(reason)); + if let Some(reply) = pending.reply { + let _ = reply.send(Err(reason)); + } } } @@ -4155,15 +4159,19 @@ fn fail_agent_registration_d1_pressure_exhausted( id: &str, pending_agent_registrations: &mut HashMap, ) { - let Some(pending) = pending_agent_registrations.remove(id) else { + let Some(pending) = pending_agent_registrations.get_mut(id) else { return; }; let reason = pending .last_d1_pressure + .take() .unwrap_or_else(|| "Node liveness retry pending".to_string()); - let _ = pending.reply.send(Err(format!( - "d1_pressure: {reason} (retry budget exhausted)" - ))); + pending.d1_pressure_retry_at = None; + if let Some(reply) = pending.reply.take() { + let _ = reply.send(Err(format!( + "d1_pressure: {reason} (retry budget exhausted)" + ))); + } } fn take_due_agent_registration_retry( @@ -4202,9 +4210,9 @@ fn expire_agent_registrations( name = %pending.name, "agent.register pending reply expired without engine response" ); - let _ = pending - .reply - .send(Err("agent_register_pending_expired".to_string())); + if let Some(reply) = pending.reply { + let _ = reply.send(Err("agent_register_pending_expired".to_string())); + } } } } @@ -4214,7 +4222,9 @@ fn drain_agent_registrations( reason: &str, ) { for (_, pending) in pending_agent_registrations.drain() { - let _ = pending.reply.send(Err(reason.to_string())); + if let Some(reply) = pending.reply { + let _ = reply.send(Err(reason.to_string())); + } } } @@ -4254,42 +4264,36 @@ fn fail_deferred_commands( } } -/// Gives every safely replayable steady-state control frame a stable -/// correlation id. The engine's `d1_pressure` contract guarantees that the -/// rejected operation was not applied, so exact replay is safe for: +/// Tracks safely replayable steady-state control frames which also have a +/// concrete recovery owner. The engine's `d1_pressure` contract guarantees +/// that the rejected operation was not applied, so exact replay is safe for: /// -/// - `agent.deregister` (engine teardown is idempotent), -/// - cumulative `delivery.ack`, and -/// - `action.accept` / `action.result` (invocation-ledger idempotency covers -/// spawn results, release results, and task receipts). +/// - acknowledged `agent.deregister`, +/// - cumulative `delivery.ack` (a missed terminal ack is recovered by durable +/// redelivery), and +/// - task `action.accept` / `action.result`, whose existing request id maps +/// retry exhaustion back to the runtime task owner. /// /// `node.register`, `inventory.sync`, and `agent.register` have dedicated /// pending state because their successful replies open liveness or return a /// token. Heartbeats intentionally are not retried: a newer periodic snapshot /// supersedes them. No separate channel-join frame exists; the only join sent -/// by this broker is part of the exact `agent.register` frame. +/// by this broker is part of the exact `agent.register` frame. One-way +/// spawn/release results and deregistrations remain untracked: acknowledging +/// their initial socket write and later consuming retry exhaustion would lose +/// the terminal outcome, so they require a separate retained-owner protocol. fn prepare_retryable_control_frame( message: &mut BrokerToRelaycast, ) -> Option<(String, &'static str)> { let (id, kind) = match message { - BrokerToRelaycast::AgentDeregister(request) => ( - request - .id - .get_or_insert_with(|| format!("agent_deregister_{}", Uuid::new_v4().simple())), - "agent.deregister", - ), + BrokerToRelaycast::AgentDeregister(request) => (request.id.as_mut()?, "agent.deregister"), BrokerToRelaycast::DeliveryAck(request) => ( request .id .get_or_insert_with(|| format!("delivery_ack_{}", Uuid::new_v4().simple())), "delivery.ack", ), - BrokerToRelaycast::ActionResult(request) => ( - request - .id - .get_or_insert_with(|| format!("action_result_{}", Uuid::new_v4().simple())), - "action.result", - ), + BrokerToRelaycast::ActionResult(request) => (request.id.as_mut()?, "action.result"), BrokerToRelaycast::ActionAccept(request) => (&mut request.id, "action.accept"), _ => return None, }; @@ -4571,7 +4575,7 @@ mod tests { session_ref: None, resumable: None, }, - reply, + reply: Some(reply), created_at, d1_pressure_retry_attempts: 0, d1_pressure_retry_at: None, @@ -5895,14 +5899,14 @@ mod tests { let frames = vec![ BrokerToRelaycast::AgentDeregister(AgentDeregister { v: FLEET_WIRE_VERSION, - id: None, + id: Some("agent_deregister_1".to_string()), agent_id: "agt-1".to_string(), name: Some("agent-a".to_string()), }), delivery_ack("agent-a", 7), BrokerToRelaycast::ActionResult(ActionResult { v: FLEET_WIRE_VERSION, - id: None, + id: Some("task_request_result_1".to_string()), invocation_id: "inv-1".to_string(), result: ActionResultPayload::Output(ActionResultOutput { output: json!({"ok": true}), @@ -5945,6 +5949,27 @@ mod tests { )), ); assert_eq!(prepare_retryable_control_frame(&mut heartbeat), None); + + let mut one_way_result = BrokerToRelaycast::ActionResult(ActionResult { + v: FLEET_WIRE_VERSION, + id: None, + invocation_id: "inv-spawn".to_string(), + result: ActionResultPayload::Output(ActionResultOutput { + output: json!({"ok": true}), + }), + task: None, + }); + assert_eq!(prepare_retryable_control_frame(&mut one_way_result), None); + let mut one_way_deregister = BrokerToRelaycast::AgentDeregister(AgentDeregister { + v: FLEET_WIRE_VERSION, + id: None, + agent_id: "agt-one-way".to_string(), + name: Some("agent-one-way".to_string()), + }); + assert_eq!( + prepare_retryable_control_frame(&mut one_way_deregister), + None + ); } #[test] @@ -6872,11 +6897,13 @@ mod tests { } #[tokio::test(start_paused = true)] - async fn node_control_agent_register_persistent_d1_pressure_fails_closed_with_named_error() { + async fn node_control_agent_register_persistent_d1_pressure_fails_closed_and_cleans_late_success( + ) { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let ws_url = format!("ws://{}/v1/node/ws", listener.local_addr().unwrap()); let (command_tx, command_rx) = mpsc::channel(32); let (event_tx, mut event_rx) = mpsc::channel(32); + let (late_success_tx, late_success_rx) = oneshot::channel(); tokio::spawn(run_node_control_client( FleetControlConfig { @@ -6939,6 +6966,32 @@ mod tests { .await .unwrap(); } + let initial = initial.expect("initial agent.register frame"); + let request_id = initial.id.expect("agent.register request id"); + late_success_rx.await.unwrap(); + ws.send(Message::Text( + serde_json::to_string(&RelaycastToBroker::Reply(crate::fleet_wire::Reply { + v: FLEET_WIRE_VERSION, + id: request_id.clone(), + ok: true, + data: json!({ + "name": "agent-a", + "agent_id": "agt-late", + "token": "at-late" + }), + })) + .unwrap(), + )) + .await + .unwrap(); + match next_non_heartbeat_node_to_server(&mut ws).await { + BrokerToRelaycast::AgentDeregister(deregister) => { + assert_eq!(deregister.id.as_deref(), Some(request_id.as_str())); + assert_eq!(deregister.agent_id, "agt-late"); + assert_eq!(deregister.name.as_deref(), Some("agent-a")); + } + other => panic!("expected compensating agent.deregister, got {other:?}"), + } }); command_tx @@ -6976,6 +7029,7 @@ mod tests { error.starts_with("d1_pressure:"), "spawn caller must see the retryable root cause, got {error}" ); + late_success_tx.send(()).unwrap(); tokio::time::timeout(Duration::from_secs(1), server) .await .unwrap() From 1bb1e1875681a1861da55eeb9e3f1cd15709e2cb Mon Sep 17 00:00:00 2001 From: kjgbot Date: Wed, 7 Oct 2026 13:14:29 -0700 Subject: [PATCH 3/9] fix(broker): harden pressure retry ownership --- CHANGELOG.md | 3 +- crates/broker/src/node_control.rs | 274 ++++++++++++++++++++++++++---- 2 files changed, 246 insertions(+), 31 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3748267eb..a055f99f5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,7 +9,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Fixed -- Fleet spawns and other safely replayable node-control operations now ride out transient Relaycast `d1_pressure` on the existing WebSocket with bounded exponential backoff, while persistent pressure fails closed with its named root cause. +- Fleet spawns now recover from transient Relaycast `d1_pressure`; persistent registration pressure returns a named `d1_pressure` error within 25 seconds. +- Acknowledged `agent.deregister`, cumulative `delivery.ack`, and task action accept/results now recover from transient Relaycast `d1_pressure`. ## [13.1.3] - 2026-10-07 diff --git a/crates/broker/src/node_control.rs b/crates/broker/src/node_control.rs index 5081e83f4..aafe67679 100644 --- a/crates/broker/src/node_control.rs +++ b/crates/broker/src/node_control.rs @@ -2699,6 +2699,7 @@ struct ApplicationLiveness { last_acknowledged: Instant, pending_inventory_syncs: VecDeque, pending_retryable_control_frames: HashMap, + exhausted_retryable_control_frames: HashMap, ready: bool, } @@ -2709,11 +2710,13 @@ impl ApplicationLiveness { last_acknowledged: Instant::now(), pending_inventory_syncs: VecDeque::new(), pending_retryable_control_frames: HashMap::new(), + exhausted_retryable_control_frames: HashMap::new(), ready: false, } } fn track_retryable_control_frame(&mut self, id: String, frame: BrokerToRelaycast) { + self.exhausted_retryable_control_frames.remove(&id); self.pending_retryable_control_frames.insert( id, PendingRetryableControlFrame { @@ -2730,6 +2733,10 @@ impl ApplicationLiveness { self.pending_retryable_control_frames.remove(id); } + fn consume_exhausted_retryable_control_frame(&mut self, id: &str) -> bool { + self.exhausted_retryable_control_frames.remove(id).is_some() + } + fn schedule_control_d1_pressure_retry( &mut self, id: &str, @@ -2772,6 +2779,10 @@ impl ApplicationLiveness { } fn expire_retryable_control_frames(&mut self, now: Instant) -> Vec<(String, String)> { + self.exhausted_retryable_control_frames + .retain(|_, exhausted_at| { + now.saturating_duration_since(*exhausted_at) < CONTROL_D1_PRESSURE_RETRY_BUDGET + }); let expired: Vec = self .pending_retryable_control_frames .iter() @@ -2786,6 +2797,12 @@ impl ApplicationLiveness { .filter_map(|id| { let pending = self.pending_retryable_control_frames.remove(&id)?; pending.last_d1_pressure.map(|reason| { + // At most one response can still be in flight because a retry is + // scheduled only after the previous pressure response arrives. + // Keep a one-shot tombstone so that late terminal response cannot + // produce a second task outcome after the timeout path below. + self.exhausted_retryable_control_frames + .insert(id.clone(), now); ( id, format!("d1_pressure: {reason} (retry budget exhausted)"), @@ -3490,19 +3507,31 @@ async fn run_connected_once_with_completions( _ = &mut control_retry, if control_retry_delay.is_some() => { let now = Instant::now(); for (id, reason) in application_liveness.expire_retryable_control_frames(now) { + let mut surfaced_to_owner = false; if let Some(pending) = pending_deregistrations.remove(&id) { let _ = pending.send(Err(reason.clone())); + surfaced_to_owner = true; } if id.starts_with(crate::runtime::task_request_prefix()) { let _ = event_tx.send(FleetControlEvent::Message(RelaycastToBroker::Error( crate::fleet_wire::Error { v: FLEET_WIRE_VERSION, - id, + id: id.clone(), ok: false, code: "d1_pressure".to_string(), - message: reason, + message: reason.clone(), }, ))).await; + surfaced_to_owner = true; + } + if !surfaced_to_owner { + tracing::error!( + target = "relay_broker::fleet", + node_id = %config.node_id, + id, + reason, + "control frame d1_pressure retry budget exhausted" + ); } } if let Some(frame) = application_liveness.take_due_control_retry(now) { @@ -3727,6 +3756,17 @@ where } match frame { RelaycastToBroker::Reply(reply) => { + if application_liveness + .consume_exhausted_retryable_control_frame(&reply.id) + { + tracing::debug!( + target = "relay_broker::fleet", + node_id, + id = %reply.id, + "consumed late reply for an exhausted control-frame retry" + ); + return true; + } application_liveness.complete_retryable_control_frame(&reply.id); if reply.id.starts_with(crate::runtime::task_request_prefix()) { return event_tx @@ -3771,6 +3811,7 @@ where complete_agent_registration( reply, pending_agent_registrations, + application_liveness, sink, ) .await @@ -3832,6 +3873,17 @@ where ControlPressureRetry::Unmatched => {} } } + if application_liveness + .consume_exhausted_retryable_control_frame(&error.id) + { + tracing::debug!( + target = "relay_broker::fleet", + node_id, + id = %error.id, + "consumed late error for an exhausted control-frame retry" + ); + return true; + } application_liveness.complete_retryable_control_frame(&error.id); if error.id.starts_with(crate::runtime::task_request_prefix()) { return event_tx @@ -3852,17 +3904,6 @@ where pending_agent_registrations, ); } - // Surface every engine rejection at error level. A node.register or - // heartbeat rejection (e.g. node_name_conflict) matches no pending - // agent registration below, so without this it vanishes silently — - // leaving the node half-registered with dead heartbeats and no signal. - tracing::error!( - target = "relay_broker::fleet", - code = %error.code, - message = %error.message, - id = %error.id, - "engine rejected a node control frame" - ); if error.code == "d1_pressure" { let reason = error.message.clone(); match schedule_agent_registration_d1_pressure_retry( @@ -3918,6 +3959,17 @@ where return true; } } + // Surface every unhandled engine rejection at error level. A + // node.register or heartbeat rejection (e.g. node_name_conflict) + // matches no pending operation below, so without this it vanishes + // silently. Expected, retryable pressure returned above at WARN. + tracing::error!( + target = "relay_broker::fleet", + code = %error.code, + message = %error.message, + id = %error.id, + "engine rejected a node control frame" + ); fail_agent_registration( &error.id, format!("{}: {}", error.code, error.message), @@ -3956,6 +4008,7 @@ where async fn complete_agent_registration( reply: crate::fleet_wire::Reply, pending_agent_registrations: &mut HashMap, + application_liveness: &mut ApplicationLiveness, sink: &mut S, ) -> bool where @@ -3980,7 +4033,7 @@ where let name = data.as_ref()?.name.as_deref()?; let key = pending_agent_registrations .iter() - .find(|(_, pending)| pending.name == name) + .find(|(_, pending)| pending.reply.is_some() && pending.name == name) .map(|(key, _)| key.clone())?; pending_agent_registrations .remove(&key) @@ -4024,17 +4077,14 @@ where agent_id = %late_token.agent_id, "late agent.register success after caller stopped waiting; sending compensating agent.deregister" ); - send_wire( - sink, - &BrokerToRelaycast::AgentDeregister(AgentDeregister { - v: FLEET_WIRE_VERSION, - id: Some(request_id), - agent_id: late_token.agent_id, - name: Some(pending.name), - }), - ) - .await - .is_ok() + let cleanup = BrokerToRelaycast::AgentDeregister(AgentDeregister { + v: FLEET_WIRE_VERSION, + id: Some(request_id.clone()), + agent_id: late_token.agent_id, + name: Some(pending.name), + }); + application_liveness.track_retryable_control_frame(request_id, cleanup.clone()); + send_wire(sink, &cleanup).await.is_ok() } fn fail_agent_registration( @@ -6047,6 +6097,62 @@ mod tests { ); } + #[tokio::test] + async fn late_d1_pressure_after_control_retry_timeout_is_consumed_once() { + let mut liveness = ApplicationLiveness::new(Duration::from_secs(120)); + let id = format!("{}late-pressure", crate::runtime::task_request_prefix()); + let frame = BrokerToRelaycast::ActionAccept(crate::fleet_wire::ActionAccept { + v: FLEET_WIRE_VERSION, + id: id.clone(), + invocation_id: "inv-late-pressure".to_string(), + execution_id: "exec-late-pressure".to_string(), + worker_generation: "generation-late-pressure".to_string(), + }); + liveness.track_retryable_control_frame(id.clone(), frame); + let now = Instant::now(); + assert!(matches!( + liveness.schedule_control_d1_pressure_retry(&id, "Node liveness retry pending", now,), + ControlPressureRetry::Scheduled(_) + )); + assert_eq!( + liveness + .expire_retryable_control_frames(now + CONTROL_D1_PRESSURE_RETRY_BUDGET) + .len(), + 1 + ); + + let (events, mut receiver) = mpsc::channel(1); + let healthy = handle_server_message( + Message::Text( + json!({ + "v": 1, + "id": id.clone(), + "type": "error", + "ok": false, + "code": "d1_pressure", + "message": "Node liveness retry pending" + }) + .to_string(), + ), + &events, + &mut HashMap::new(), + &mut HashMap::new(), + &mut liveness, + "node-test", + &mut futures_util::sink::drain(), + None, + None, + ) + .await; + + assert!(healthy); + assert!(matches!( + receiver.try_recv(), + Err(mpsc::error::TryRecvError::Empty) + )); + assert!(!liveness.consume_exhausted_retryable_control_frame(&id)); + } + #[test] fn expire_agent_registrations_bounds_pending_map() { let created_at = Instant::now(); @@ -6105,7 +6211,8 @@ mod tests { }), }; - assert!(complete_agent_registration(reply, &mut pending, &mut sink).await); + let mut liveness = ApplicationLiveness::new(Duration::from_secs(30)); + assert!(complete_agent_registration(reply, &mut pending, &mut liveness, &mut sink).await); assert!( pending.is_empty(), "the pending registration must be consumed by the name-based fallback" @@ -6121,6 +6228,103 @@ mod tests { ); } + #[tokio::test] + async fn complete_agent_registration_name_fallback_skips_exhausted_same_name_request() { + let (stale_tx, _stale_rx) = oneshot::channel(); + let mut stale = test_pending_agent_registration( + "agent-a", + false, + "agent_register_stale", + stale_tx, + Instant::now(), + ); + stale.reply = None; + let (active_tx, active_rx) = oneshot::channel(); + let active = test_pending_agent_registration( + "agent-a", + false, + "agent_register_active", + active_tx, + Instant::now(), + ); + let mut pending = HashMap::from([ + ("agent_register_stale".to_string(), stale), + ("agent_register_active".to_string(), active), + ]); + let reply = crate::fleet_wire::Reply { + v: FLEET_WIRE_VERSION, + id: "196331520553345024".to_string(), + ok: true, + data: json!({ + "name": "agent-a", + "agent_id": "agt-active", + "token": "at_active" + }), + }; + let mut liveness = ApplicationLiveness::new(Duration::from_secs(30)); + + assert!( + complete_agent_registration( + reply, + &mut pending, + &mut liveness, + &mut futures_util::sink::drain(), + ) + .await + ); + assert!(pending.contains_key("agent_register_stale")); + assert!(!pending.contains_key("agent_register_active")); + assert_eq!(active_rx.await.unwrap().unwrap().agent_id, "agt-active"); + } + + #[tokio::test] + async fn late_agent_registration_cleanup_is_tracked_for_pressure_retry() { + let (reply_tx, reply_rx) = oneshot::channel(); + drop(reply_rx); + let mut pending = HashMap::from([( + "agent_register_late".to_string(), + test_pending_agent_registration( + "agent-a", + false, + "agent_register_late", + reply_tx, + Instant::now(), + ), + )]); + let reply = crate::fleet_wire::Reply { + v: FLEET_WIRE_VERSION, + id: "agent_register_late".to_string(), + ok: true, + data: json!({ + "name": "agent-a", + "agent_id": "agt-late", + "token": "at_late" + }), + }; + let mut liveness = ApplicationLiveness::new(Duration::from_secs(30)); + + assert!( + complete_agent_registration( + reply, + &mut pending, + &mut liveness, + &mut futures_util::sink::drain(), + ) + .await + ); + let tracked = liveness + .pending_retryable_control_frames + .get("agent_register_late") + .expect("compensating deregistration must retain a retry owner"); + assert!(matches!( + &tracked.frame, + BrokerToRelaycast::AgentDeregister(request) + if request.id.as_deref() == Some("agent_register_late") + && request.agent_id == "agt-late" + && request.name.as_deref() == Some("agent-a") + )); + } + #[tokio::test] async fn complete_agent_registration_ignores_non_agent_reply() { // A `node.register`/`inventory.sync` reply carries a fresh engine @@ -6146,7 +6350,8 @@ mod tests { data: json!({ "node_id": "node-test", "online": true }), }; - assert!(complete_agent_registration(reply, &mut pending, &mut sink).await); + let mut liveness = ApplicationLiveness::new(Duration::from_secs(30)); + assert!(complete_agent_registration(reply, &mut pending, &mut liveness, &mut sink).await); assert_eq!( pending.len(), 1, @@ -6296,8 +6501,17 @@ mod tests { )) .await .unwrap(); - let ack = next_non_heartbeat_node_to_server(&mut ws).await; - assert_eq!(ack, delivery_ack("agent-a", 1)); + match next_non_heartbeat_node_to_server(&mut ws).await { + BrokerToRelaycast::DeliveryAck(ack) => { + assert_eq!(ack.agent, "agent-a"); + assert_eq!(ack.up_to_seq, 1); + assert!(ack + .id + .as_deref() + .is_some_and(|id| id.starts_with("delivery_ack_"))); + } + other => panic!("expected delivery.ack, got {other:?}"), + } ws.send(Message::Text( serde_json::to_string(&RelaycastToBroker::ActionInvoke(ActionInvoke { @@ -6896,7 +7110,7 @@ mod tests { let _ = command_tx.send(FleetControlCommand::Shutdown).await; } - #[tokio::test(start_paused = true)] + #[tokio::test] async fn node_control_agent_register_persistent_d1_pressure_fails_closed_and_cleans_late_success( ) { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); From bd5ebb0cbb350d3f4237ad0be5f7f4fc090cdbc2 Mon Sep 17 00:00:00 2001 From: kjgbot Date: Wed, 7 Oct 2026 15:54:07 -0700 Subject: [PATCH 4/9] fix(broker): retain pressured results across reconnects --- CHANGELOG.md | 2 +- crates/broker/src/node_control.rs | 890 +++++++++++++++--- .../cli/lib/fleet-spawn-confirmation.test.ts | 19 + .../sdk/src/messaging/relaycast-placement.ts | 2 + packages/sdk/src/messaging/relaycast.ts | 8 +- 5 files changed, 811 insertions(+), 110 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a055f99f5..ab5c8a7d8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,7 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Fixed - Fleet spawns now recover from transient Relaycast `d1_pressure`; persistent registration pressure returns a named `d1_pressure` error within 25 seconds. -- Acknowledged `agent.deregister`, cumulative `delivery.ack`, and task action accept/results now recover from transient Relaycast `d1_pressure`. +- Acknowledged `agent.deregister`, cumulative `delivery.ack`, and action accept/results now recover from transient Relaycast `d1_pressure`; exhausted non-task results are retained across reconnect, and the CLI reports registration pressure as retryable. ## [13.1.3] - 2026-10-07 diff --git a/crates/broker/src/node_control.rs b/crates/broker/src/node_control.rs index aafe67679..e6c520f07 100644 --- a/crates/broker/src/node_control.rs +++ b/crates/broker/src/node_control.rs @@ -76,6 +76,28 @@ fn d1_pressure_retry_delay(retry_attempts: u32) -> Duration { (D1_PRESSURE_RETRY_INITIAL_DELAY * multiplier).min(D1_PRESSURE_RETRY_MAX_DELAY) } +#[derive(Clone, Copy)] +struct ControlRetryPolicy { + initial_delay: Duration, + max_delay: Duration, + max_retries: u32, + budget: Duration, +} + +impl ControlRetryPolicy { + const DEFAULT: Self = Self { + initial_delay: D1_PRESSURE_RETRY_INITIAL_DELAY, + max_delay: D1_PRESSURE_RETRY_MAX_DELAY, + max_retries: CONTROL_D1_PRESSURE_MAX_RETRIES, + budget: CONTROL_D1_PRESSURE_RETRY_BUDGET, + }; + + fn retry_delay(self, retry_attempts: u32) -> Duration { + let multiplier = 1u32 << retry_attempts.saturating_sub(1).min(3); + (self.initial_delay * multiplier).min(self.max_delay) + } +} + /// Whether a fresh re-mint should still be attempted at this consecutive-401 /// count. The count is incremented for the current 401 *before* this is called, /// so the very first 401 (count 1) is still within budget. Crucially this is @@ -2331,6 +2353,7 @@ pub(crate) async fn run_node_control_client_with_completions( // token once, but never follow an unbounded chain of cache changes. let mut proof_conflict_cache_retry_used = false; let mut retained_completion: Option = None; + let mut retained_pressure_frames: VecDeque = VecDeque::new(); loop { // Retained completion frames stay in their unbounded lane until a @@ -2498,11 +2521,13 @@ pub(crate) async fn run_node_control_client_with_completions( &mut command_rx, &mut completion_rx, &mut retained_completion, + &mut retained_pressure_frames, &event_tx, &mut registration, &mut inventory, &mut load, INVENTORY_REFRESH_INTERVAL, + ControlRetryPolicy::DEFAULT, ) .await; if matches!(result, ControlRunResult::Shutdown) { @@ -2696,6 +2721,7 @@ enum ControlPressureRetry { struct ApplicationLiveness { deadline: Duration, + control_retry_policy: ControlRetryPolicy, last_acknowledged: Instant, pending_inventory_syncs: VecDeque, pending_retryable_control_frames: HashMap, @@ -2704,9 +2730,15 @@ struct ApplicationLiveness { } impl ApplicationLiveness { + #[cfg(test)] fn new(deadline: Duration) -> Self { + Self::new_with_retry_policy(deadline, ControlRetryPolicy::DEFAULT) + } + + fn new_with_retry_policy(deadline: Duration, control_retry_policy: ControlRetryPolicy) -> Self { Self { deadline, + control_retry_policy, last_acknowledged: Instant::now(), pending_inventory_syncs: VecDeque::new(), pending_retryable_control_frames: HashMap::new(), @@ -2747,12 +2779,12 @@ impl ApplicationLiveness { return ControlPressureRetry::Unmatched; }; pending.last_d1_pressure = Some(reason.to_string()); - let deadline = pending.created_at + CONTROL_D1_PRESSURE_RETRY_BUDGET; - if pending.retry_attempts >= CONTROL_D1_PRESSURE_MAX_RETRIES || now >= deadline { + let deadline = pending.created_at + self.control_retry_policy.budget; + if pending.retry_attempts >= self.control_retry_policy.max_retries || now >= deadline { return ControlPressureRetry::Exhausted; } let next_attempt = pending.retry_attempts.saturating_add(1); - let delay = d1_pressure_retry_delay(next_attempt); + let delay = self.control_retry_policy.retry_delay(next_attempt); if now .checked_add(delay) .is_none_or(|retry_at| retry_at > deadline) @@ -2765,32 +2797,45 @@ impl ApplicationLiveness { } fn next_control_retry_event_delay(&self, now: Instant) -> Option { + // Successful delivery acknowledgements and non-task action results are + // intentionally one-way: the engine sends no positive reply. Their + // entries therefore live until this bounded budget elapses. Memory is + // bounded to roughly one budget window of control traffic, and this + // O(n) deadline scan covers only that window rather than session age. self.pending_retryable_control_frames .values() .flat_map(|pending| { [ pending.retry_at, - Some(pending.created_at + CONTROL_D1_PRESSURE_RETRY_BUDGET), + Some(pending.created_at + self.control_retry_policy.budget), ] }) .flatten() + .chain( + self.exhausted_retryable_control_frames + .values() + .map(|exhausted_at| *exhausted_at + self.control_retry_policy.budget), + ) .map(|event_at| event_at.saturating_duration_since(now)) .min() } - fn expire_retryable_control_frames(&mut self, now: Instant) -> Vec<(String, String)> { + fn expire_retryable_control_frames( + &mut self, + now: Instant, + ) -> Vec<(String, String, BrokerToRelaycast)> { self.exhausted_retryable_control_frames .retain(|_, exhausted_at| { - now.saturating_duration_since(*exhausted_at) < CONTROL_D1_PRESSURE_RETRY_BUDGET + now.saturating_duration_since(*exhausted_at) < self.control_retry_policy.budget }); let expired: Vec = self .pending_retryable_control_frames .iter() - .filter_map(|(id, pending)| { - (now.saturating_duration_since(pending.created_at) - >= CONTROL_D1_PRESSURE_RETRY_BUDGET) - .then(|| id.clone()) + .filter(|(_, pending)| { + now.saturating_duration_since(pending.created_at) + >= self.control_retry_policy.budget }) + .map(|(id, _)| id.clone()) .collect(); expired .into_iter() @@ -2806,12 +2851,39 @@ impl ApplicationLiveness { ( id, format!("d1_pressure: {reason} (retry budget exhausted)"), + pending.frame, ) }) }) .collect() } + fn take_retryable_control_frame(&mut self, id: &str) -> Option { + self.pending_retryable_control_frames + .remove(id) + .map(|pending| pending.frame) + } + + fn retain_non_task_action_results( + &mut self, + retained_pressure_frames: &mut VecDeque, + ) { + let retained_ids: Vec = self + .pending_retryable_control_frames + .iter() + .filter(|(id, pending)| { + !id.starts_with(crate::runtime::task_request_prefix()) + && matches!(&pending.frame, BrokerToRelaycast::ActionResult(_)) + }) + .map(|(id, _)| id.clone()) + .collect(); + for id in retained_ids { + if let Some(pending) = self.pending_retryable_control_frames.remove(&id) { + retained_pressure_frames.push_back(pending.frame); + } + } + } + fn take_due_control_retry(&mut self, now: Instant) -> Option { let pending = self .pending_retryable_control_frames @@ -3178,7 +3250,7 @@ where if sent.is_err() { if let BrokerToRelaycast::ActionResult(result) = &message { tracing::warn!(invocation_id = %result.invocation_id, frame_kind = "action.result", - "dropping fleet frame after wire send failure"); + "retaining fleet result for replay after reconnect"); } return Break(ControlRunResult::Disconnected { application_ready: application_liveness.ready, @@ -3187,12 +3259,16 @@ where Continue(()) } Some(FleetControlCommand::RetainedSend(mut completion)) => { - if let Some((request_id, _)) = prepare_retryable_control_frame(&mut completion.message) - { + let retry_id = prepare_retryable_control_frame(&mut completion.message) + .map(|(request_id, _)| request_id); + if let Some(request_id) = retry_id.as_ref() { application_liveness - .track_retryable_control_frame(request_id, completion.message.clone()); + .track_retryable_control_frame(request_id.clone(), completion.message.clone()); } if send_wire(sink, &completion.message).await.is_err() { + if let Some(request_id) = retry_id { + application_liveness.complete_retryable_control_frame(&request_id); + } *retained_completion = Some(completion); return Break(ControlRunResult::Disconnected { application_ready: application_liveness.ready, @@ -3262,6 +3338,38 @@ async fn run_connected_once( inventory: &mut Vec, load: &mut FleetLoadSnapshot, inventory_refresh_interval: Duration, +) -> ControlRunResult { + let (_completion_tx, mut completion_rx) = mpsc::unbounded_channel(); + let mut retained_completion = None; + let mut retained_pressure_frames = VecDeque::new(); + run_connected_once_with_completions( + config, + command_rx, + &mut completion_rx, + &mut retained_completion, + &mut retained_pressure_frames, + event_tx, + registration, + inventory, + load, + inventory_refresh_interval, + ControlRetryPolicy::DEFAULT, + ) + .await +} + +#[cfg(test)] +#[allow(clippy::too_many_arguments)] +async fn run_connected_once_with_retry_policy( + config: &FleetControlConfig, + command_rx: &mut mpsc::Receiver, + event_tx: &mpsc::Sender, + registration: &mut Option, + inventory: &mut Vec, + load: &mut FleetLoadSnapshot, + inventory_refresh_interval: Duration, + retry_policy: ControlRetryPolicy, + retained_pressure_frames: &mut VecDeque, ) -> ControlRunResult { let (_completion_tx, mut completion_rx) = mpsc::unbounded_channel(); let mut retained_completion = None; @@ -3270,11 +3378,13 @@ async fn run_connected_once( command_rx, &mut completion_rx, &mut retained_completion, + retained_pressure_frames, event_tx, registration, inventory, load, inventory_refresh_interval, + retry_policy, ) .await } @@ -3285,11 +3395,13 @@ async fn run_connected_once_with_completions( command_rx: &mut mpsc::Receiver, completion_rx: &mut mpsc::UnboundedReceiver, retained_completion: &mut Option, + retained_pressure_frames: &mut VecDeque, event_tx: &mpsc::Sender, registration: &mut Option, inventory: &mut Vec, load: &mut FleetLoadSnapshot, inventory_refresh_interval: Duration, + control_retry_policy: ControlRetryPolicy, ) -> ControlRunResult { let Some(mut node_register) = registration.clone() else { return ControlRunResult::ConnectFailed; @@ -3385,7 +3497,10 @@ async fn run_connected_once_with_completions( .checked_mul(2) .unwrap_or(Duration::MAX) .max(read_idle_timeout); - let mut application_liveness = ApplicationLiveness::new(application_liveness_timeout); + let mut application_liveness = ApplicationLiveness::new_with_retry_policy( + application_liveness_timeout, + control_retry_policy, + ); let deferred_commands = match register_node_session( &mut sink, @@ -3433,6 +3548,18 @@ async fn run_connected_once_with_completions( application_ready: false, }; } + while let Some(mut frame) = retained_pressure_frames.pop_front() { + let request_id = prepare_retryable_control_frame(&mut frame) + .map(|(id, _)| id) + .expect("pressure-retained control frames must remain retryable"); + application_liveness.track_retryable_control_frame(request_id, frame.clone()); + if send_wire(&mut sink, &frame).await.is_err() { + application_liveness.retain_non_task_action_results(retained_pressure_frames); + return ControlRunResult::Disconnected { + application_ready: false, + }; + } + } // The idle check runs on the heartbeat tick, so the tick must be shorter // than the window it polices; an overridden (test) window keeps that ratio. @@ -3483,6 +3610,9 @@ async fn run_connected_once_with_completions( "node_control_disconnected" }; fail_deferred_commands(deferred_commands.collect(), reason, inventory, load); + if matches!(result, ControlRunResult::Disconnected { .. }) { + application_liveness.retain_non_task_action_results(retained_pressure_frames); + } return result; } } @@ -3492,8 +3622,11 @@ async fn run_connected_once_with_completions( let inventory_retry = tokio::time::sleep(inventory_retry_delay.unwrap_or(application_liveness.deadline)); tokio::pin!(inventory_retry); - let agent_registration_retry_delay = - next_agent_registration_retry_event_delay(&pending_agent_registrations, Instant::now()); + let agent_registration_retry_delay = next_agent_registration_retry_event_delay( + &pending_agent_registrations, + Instant::now(), + application_liveness.control_retry_policy, + ); let agent_registration_retry = tokio::time::sleep( agent_registration_retry_delay.unwrap_or(application_liveness.deadline), ); @@ -3506,7 +3639,8 @@ async fn run_connected_once_with_completions( tokio::select! { _ = &mut control_retry, if control_retry_delay.is_some() => { let now = Instant::now(); - for (id, reason) in application_liveness.expire_retryable_control_frames(now) { + let mut reconnect_for_retained_result = false; + for (id, reason, frame) in application_liveness.expire_retryable_control_frames(now) { let mut surfaced_to_owner = false; if let Some(pending) = pending_deregistrations.remove(&id) { let _ = pending.send(Err(reason.clone())); @@ -3524,6 +3658,20 @@ async fn run_connected_once_with_completions( ))).await; surfaced_to_owner = true; } + if !id.starts_with(crate::runtime::task_request_prefix()) + && matches!(&frame, BrokerToRelaycast::ActionResult(_)) + { + retained_pressure_frames.push_back(frame); + reconnect_for_retained_result = true; + surfaced_to_owner = true; + tracing::warn!( + target = "relay_broker::fleet", + node_id = %config.node_id, + id, + reason, + "retaining action.result after d1_pressure exhaustion; reconnecting for redelivery" + ); + } if !surfaced_to_owner { tracing::error!( target = "relay_broker::fleet", @@ -3534,8 +3682,14 @@ async fn run_connected_once_with_completions( ); } } + if reconnect_for_retained_result { + application_liveness.retain_non_task_action_results(retained_pressure_frames); + drain_agent_registrations(&mut pending_agent_registrations, "node_control_disconnected"); + return ControlRunResult::Disconnected { application_ready: application_liveness.ready }; + } if let Some(frame) = application_liveness.take_due_control_retry(now) { if send_wire(&mut sink, &frame).await.is_err() { + application_liveness.retain_non_task_action_results(retained_pressure_frames); drain_agent_registrations(&mut pending_agent_registrations, "node_control_disconnected"); return ControlRunResult::Disconnected { application_ready: application_liveness.ready }; } @@ -3543,12 +3697,17 @@ async fn run_connected_once_with_completions( } _ = &mut agent_registration_retry, if agent_registration_retry_delay.is_some() => { let now = Instant::now(); - expire_d1_pressure_agent_registrations(&mut pending_agent_registrations, now); + expire_d1_pressure_agent_registrations( + &mut pending_agent_registrations, + now, + application_liveness.control_retry_policy, + ); if let Some(frame) = take_due_agent_registration_retry( &mut pending_agent_registrations, now, ) { if send_wire(&mut sink, &BrokerToRelaycast::AgentRegister(frame)).await.is_err() { + application_liveness.retain_non_task_action_results(retained_pressure_frames); drain_agent_registrations(&mut pending_agent_registrations, "node_control_disconnected"); return ControlRunResult::Disconnected { application_ready: application_liveness.ready }; } @@ -3559,6 +3718,7 @@ async fn run_connected_once_with_completions( continue; }; if send_wire(&mut sink, &BrokerToRelaycast::InventorySync(frame)).await.is_err() { + application_liveness.retain_non_task_action_results(retained_pressure_frames); drain_agent_registrations(&mut pending_agent_registrations, "node_control_disconnected"); return ControlRunResult::Disconnected { application_ready: application_liveness.ready }; } @@ -3580,6 +3740,9 @@ async fn run_connected_once_with_completions( ) .await { + if matches!(result, ControlRunResult::Disconnected { .. }) { + application_liveness.retain_non_task_action_results(retained_pressure_frames); + } return result; } } @@ -3601,6 +3764,7 @@ async fn run_connected_once_with_completions( idle_secs = idle.as_secs(), "no inbound node-control frame within the read-idle window; reconnecting" ); + application_liveness.retain_non_task_action_results(retained_pressure_frames); drain_agent_registrations(&mut pending_agent_registrations, "node_control_disconnected"); return ControlRunResult::Disconnected { application_ready: application_liveness.ready }; } @@ -3615,16 +3779,19 @@ async fn run_connected_once_with_completions( pending_inventory_syncs = application_liveness.pending_inventory_syncs.len(), "node-control application acknowledgement deadline exceeded while transport remained active; reconnecting" ); + application_liveness.retain_non_task_action_results(retained_pressure_frames); drain_agent_registrations(&mut pending_agent_registrations, "node_control_disconnected"); return ControlRunResult::Disconnected { application_ready: application_liveness.ready }; } if send_wire(&mut sink, &BrokerToRelaycast::NodeHeartbeat(load.heartbeat(&node_register))).await.is_err() { + application_liveness.retain_non_task_action_results(retained_pressure_frames); drain_agent_registrations(&mut pending_agent_registrations, "node_control_disconnected"); return ControlRunResult::Disconnected { application_ready: application_liveness.ready }; } // Guarantees the peer owes us a frame every interval, so an idle // engine is distinguishable from a dead connection. if send_ws_frame(&mut sink, Message::Ping(Vec::new())).await.is_err() { + application_liveness.retain_non_task_action_results(retained_pressure_frames); drain_agent_registrations(&mut pending_agent_registrations, "node_control_disconnected"); return ControlRunResult::Disconnected { application_ready: application_liveness.ready }; } @@ -3638,11 +3805,13 @@ async fn run_connected_once_with_completions( ) .await { + application_liveness.retain_non_task_action_results(retained_pressure_frames); return ControlRunResult::Disconnected { application_ready: application_liveness.ready }; } } message = stream.next() => { let Some(message) = message else { + application_liveness.retain_non_task_action_results(retained_pressure_frames); drain_agent_registrations(&mut pending_agent_registrations, "node_control_disconnected"); return ControlRunResult::Disconnected { application_ready: application_liveness.ready }; }; @@ -3650,6 +3819,7 @@ async fn run_connected_once_with_completions( Ok(message) => message, Err(error) => { tracing::warn!(target = "relay_broker::fleet", error = %error, "fleet node ws read failed"); + application_liveness.retain_non_task_action_results(retained_pressure_frames); drain_agent_registrations(&mut pending_agent_registrations, "node_control_disconnected"); return ControlRunResult::Disconnected { application_ready: application_liveness.ready }; } @@ -3664,6 +3834,7 @@ async fn run_connected_once_with_completions( &mut pending_agent_registrations, &mut pending_deregistrations, &mut application_liveness, + retained_pressure_frames, &config.node_id, &mut sink, config.probe.as_ref(), @@ -3671,6 +3842,7 @@ async fn run_connected_once_with_completions( ) .await { + application_liveness.retain_non_task_action_results(retained_pressure_frames); drain_agent_registrations(&mut pending_agent_registrations, "node_control_disconnected"); return ControlRunResult::Disconnected { application_ready: application_liveness.ready }; } @@ -3714,6 +3886,7 @@ async fn handle_server_message( pending_agent_registrations: &mut HashMap, pending_deregistrations: &mut HashMap>>, application_liveness: &mut ApplicationLiveness, + retained_pressure_frames: &mut VecDeque, node_id: &str, sink: &mut S, probe: Option<&std::sync::Arc>, @@ -3836,12 +4009,31 @@ where return true; } ControlPressureRetry::Exhausted => { - application_liveness - .complete_retryable_control_frame(&error.id); + let exhausted_frame = application_liveness + .take_retryable_control_frame(&error.id); let reason = format!( "d1_pressure: {} (retry budget exhausted)", error.message, ); + if !error + .id + .starts_with(crate::runtime::task_request_prefix()) + && exhausted_frame.as_ref().is_some_and(|frame| { + matches!(frame, BrokerToRelaycast::ActionResult(_)) + }) + { + retained_pressure_frames.push_back( + exhausted_frame.expect("checked action.result"), + ); + tracing::warn!( + target = "relay_broker::fleet", + node_id, + id = %error.id, + reason, + "retaining action.result after d1_pressure exhaustion; reconnecting for redelivery" + ); + return false; + } if let Some(pending) = pending_deregistrations.remove(&error.id) { @@ -3911,6 +4103,7 @@ where &reason, pending_agent_registrations, Instant::now(), + application_liveness.control_retry_policy, ) { AgentRegistrationPressureRetry::Scheduled(delay) => { tracing::warn!( @@ -4034,6 +4227,11 @@ where let key = pending_agent_registrations .iter() .find(|(_, pending)| pending.reply.is_some() && pending.name == name) + .or_else(|| { + pending_agent_registrations + .iter() + .find(|(_, pending)| pending.name == name) + }) .map(|(key, _)| key.clone())?; pending_agent_registrations .remove(&key) @@ -4146,18 +4344,19 @@ fn schedule_agent_registration_d1_pressure_retry( reason: &str, pending_agent_registrations: &mut HashMap, now: Instant, + retry_policy: ControlRetryPolicy, ) -> AgentRegistrationPressureRetry { let Some(pending) = pending_agent_registrations.get_mut(id) else { return AgentRegistrationPressureRetry::Unmatched; }; pending.last_d1_pressure = Some(reason.to_string()); - let deadline = pending.created_at + CONTROL_D1_PRESSURE_RETRY_BUDGET; - if pending.d1_pressure_retry_attempts >= CONTROL_D1_PRESSURE_MAX_RETRIES || now >= deadline { + let deadline = pending.created_at + retry_policy.budget; + if pending.d1_pressure_retry_attempts >= retry_policy.max_retries || now >= deadline { return AgentRegistrationPressureRetry::Exhausted; } let next_attempt = pending.d1_pressure_retry_attempts.saturating_add(1); - let delay = d1_pressure_retry_delay(next_attempt); + let delay = retry_policy.retry_delay(next_attempt); if now .checked_add(delay) .is_none_or(|retry_at| retry_at > deadline) @@ -4172,6 +4371,7 @@ fn schedule_agent_registration_d1_pressure_retry( fn next_agent_registration_retry_event_delay( pending_agent_registrations: &HashMap, now: Instant, + retry_policy: ControlRetryPolicy, ) -> Option { pending_agent_registrations .values() @@ -4179,7 +4379,7 @@ fn next_agent_registration_retry_event_delay( .flat_map(|pending| { [ pending.d1_pressure_retry_at, - Some(pending.created_at + CONTROL_D1_PRESSURE_RETRY_BUDGET), + Some(pending.created_at + retry_policy.budget), ] }) .flatten() @@ -4190,15 +4390,15 @@ fn next_agent_registration_retry_event_delay( fn expire_d1_pressure_agent_registrations( pending_agent_registrations: &mut HashMap, now: Instant, + retry_policy: ControlRetryPolicy, ) { let expired: Vec = pending_agent_registrations .iter() - .filter_map(|(id, pending)| { - (pending.last_d1_pressure.is_some() - && now.saturating_duration_since(pending.created_at) - >= CONTROL_D1_PRESSURE_RETRY_BUDGET) - .then(|| id.clone()) + .filter(|(_, pending)| { + pending.last_d1_pressure.is_some() + && now.saturating_duration_since(pending.created_at) >= retry_policy.budget }) + .map(|(id, _)| id.clone()) .collect(); for id in expired { fail_agent_registration_d1_pressure_exhausted(&id, pending_agent_registrations); @@ -4321,17 +4521,17 @@ fn fail_deferred_commands( /// - acknowledged `agent.deregister`, /// - cumulative `delivery.ack` (a missed terminal ack is recovered by durable /// redelivery), and -/// - task `action.accept` / `action.result`, whose existing request id maps -/// retry exhaustion back to the runtime task owner. +/// - task `action.accept` / `action.result`, whose request id maps retry +/// exhaustion back to the runtime task owner, and +/// - non-task `action.result`, which gets a stable id and is retained for +/// redelivery after reconnect if the bounded pressure budget is exhausted. /// /// `node.register`, `inventory.sync`, and `agent.register` have dedicated /// pending state because their successful replies open liveness or return a /// token. Heartbeats intentionally are not retried: a newer periodic snapshot /// supersedes them. No separate channel-join frame exists; the only join sent /// by this broker is part of the exact `agent.register` frame. One-way -/// spawn/release results and deregistrations remain untracked: acknowledging -/// their initial socket write and later consuming retry exhaustion would lose -/// the terminal outcome, so they require a separate retained-owner protocol. +/// deregistrations remain untracked because they have no retained owner. fn prepare_retryable_control_frame( message: &mut BrokerToRelaycast, ) -> Option<(String, &'static str)> { @@ -4343,7 +4543,12 @@ fn prepare_retryable_control_frame( .get_or_insert_with(|| format!("delivery_ack_{}", Uuid::new_v4().simple())), "delivery.ack", ), - BrokerToRelaycast::ActionResult(request) => (request.id.as_mut()?, "action.result"), + BrokerToRelaycast::ActionResult(request) => ( + request + .id + .get_or_insert_with(|| format!("action_result_{}", Uuid::new_v4().simple())), + "action.result", + ), BrokerToRelaycast::ActionAccept(request) => (&mut request.id, "action.accept"), _ => return None, }; @@ -4530,7 +4735,6 @@ mod tests { let mut application_liveness = ApplicationLiveness::new(Duration::from_secs(1)); let (delivered, mut delivery_ack) = oneshot::channel(); let mut retained_completion = None; - let result = handle_connected_command( Some(FleetControlCommand::RetainedSend(RetainedFleetCompletion { message: BrokerToRelaycast::ActionResult(ActionResult { @@ -5778,6 +5982,7 @@ mod tests { &mut pending, &mut deregistrations, &mut ApplicationLiveness::new(Duration::from_secs(1)), + &mut VecDeque::new(), "node-test", &mut sink, None, @@ -5829,6 +6034,7 @@ mod tests { &mut HashMap::new(), &mut HashMap::new(), &mut liveness, + &mut VecDeque::new(), "node-test", &mut futures_util::sink::drain(), None, @@ -5869,6 +6075,7 @@ mod tests { &mut HashMap::new(), &mut HashMap::new(), &mut ApplicationLiveness::new(Duration::from_secs(1)), + &mut VecDeque::new(), "node-test", &mut futures_util::sink::drain(), None, @@ -6009,7 +6216,14 @@ mod tests { }), task: None, }); - assert_eq!(prepare_retryable_control_frame(&mut one_way_result), None); + let (result_id, kind) = prepare_retryable_control_frame(&mut one_way_result) + .expect("non-task action.result must gain a retained retry owner"); + assert_eq!(kind, "action.result"); + assert!(result_id.starts_with("action_result_")); + assert!(matches!( + one_way_result, + BrokerToRelaycast::ActionResult(ActionResult { id: Some(id), .. }) if id == result_id + )); let mut one_way_deregister = BrokerToRelaycast::AgentDeregister(AgentDeregister { v: FLEET_WIRE_VERSION, id: None, @@ -6071,6 +6285,7 @@ mod tests { &mut HashMap::new(), &mut HashMap::new(), &mut liveness, + &mut VecDeque::new(), "node-test", &mut futures_util::sink::drain(), None, @@ -6138,6 +6353,7 @@ mod tests { &mut HashMap::new(), &mut HashMap::new(), &mut liveness, + &mut VecDeque::new(), "node-test", &mut futures_util::sink::drain(), None, @@ -6153,6 +6369,48 @@ mod tests { assert!(!liveness.consume_exhausted_retryable_control_frame(&id)); } + #[test] + fn exhausted_control_retry_tombstone_has_its_own_cleanup_deadline() { + let policy = ControlRetryPolicy { + initial_delay: Duration::from_millis(5), + max_delay: Duration::from_millis(5), + max_retries: 1, + budget: Duration::from_millis(50), + }; + let mut liveness = + ApplicationLiveness::new_with_retry_policy(Duration::from_secs(1), policy); + let id = format!("{}tombstone", crate::runtime::task_request_prefix()); + liveness.track_retryable_control_frame( + id.clone(), + BrokerToRelaycast::ActionAccept(crate::fleet_wire::ActionAccept { + v: FLEET_WIRE_VERSION, + id: id.clone(), + invocation_id: "inv-tombstone".to_string(), + execution_id: "exec-tombstone".to_string(), + worker_generation: "generation-tombstone".to_string(), + }), + ); + let now = Instant::now(); + assert!(matches!( + liveness.schedule_control_d1_pressure_retry(&id, "Node liveness retry pending", now,), + ControlPressureRetry::Scheduled(_) + )); + assert_eq!( + liveness + .expire_retryable_control_frames(now + policy.budget) + .len(), + 1 + ); + assert_eq!( + liveness.next_control_retry_event_delay(now + policy.budget), + Some(policy.budget) + ); + assert!(liveness + .expire_retryable_control_frames(now + policy.budget + policy.budget) + .is_empty()); + assert!(liveness.exhausted_retryable_control_frames.is_empty()); + } + #[test] fn expire_agent_registrations_bounds_pending_map() { let created_at = Instant::now(); @@ -6279,21 +6537,19 @@ mod tests { #[tokio::test] async fn late_agent_registration_cleanup_is_tracked_for_pressure_retry() { - let (reply_tx, reply_rx) = oneshot::channel(); - drop(reply_rx); - let mut pending = HashMap::from([( - "agent_register_late".to_string(), - test_pending_agent_registration( - "agent-a", - false, - "agent_register_late", - reply_tx, - Instant::now(), - ), - )]); + let (reply_tx, _reply_rx) = oneshot::channel(); + let mut exhausted = test_pending_agent_registration( + "agent-a", + false, + "agent_register_late", + reply_tx, + Instant::now(), + ); + exhausted.reply = None; + let mut pending = HashMap::from([("agent_register_late".to_string(), exhausted)]); let reply = crate::fleet_wire::Reply { v: FLEET_WIRE_VERSION, - id: "agent_register_late".to_string(), + id: "196331520553345024".to_string(), ok: true, data: json!({ "name": "agent-a", @@ -6746,6 +7002,399 @@ mod tests { server.abort(); } + #[derive(Clone, Copy, Debug, PartialEq, Eq)] + enum SafeControlRetryCase { + DeliveryAck, + ActionResult, + AgentDeregister, + } + + async fn assert_safe_control_frame_socket_retry( + case: SafeControlRetryCase, + persistent: bool, + disconnect_during_backoff: bool, + ) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let ws_url = format!("ws://{}/v1/node/ws", listener.local_addr().unwrap()); + let (command_tx, mut command_rx) = mpsc::channel(8); + let (event_tx, mut event_rx) = mpsc::channel(8); + let (server_done_tx, server_done_rx) = oneshot::channel(); + let retry_policy = ControlRetryPolicy { + initial_delay: Duration::from_millis(100), + max_delay: Duration::from_millis(100), + max_retries: 2, + budget: Duration::from_millis(350), + }; + let server_max_retries = retry_policy.max_retries; + let server = tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + let mut ws = accept_async(stream).await.unwrap(); + assert!(matches!( + next_node_to_server(&mut ws).await, + BrokerToRelaycast::NodeRegister(_) + )); + assert!(matches!( + next_node_to_server(&mut ws).await, + BrokerToRelaycast::InventorySync(_) + )); + assert!(matches!( + next_node_to_server(&mut ws).await, + BrokerToRelaycast::NodeHeartbeat(_) + )); + + let initial = next_non_heartbeat_node_to_server(&mut ws).await; + match (&case, &initial) { + (SafeControlRetryCase::DeliveryAck, BrokerToRelaycast::DeliveryAck(_)) + | (SafeControlRetryCase::ActionResult, BrokerToRelaycast::ActionResult(_)) + | (SafeControlRetryCase::AgentDeregister, BrokerToRelaycast::AgentDeregister(_)) => { + } + _ => panic!("unexpected frame for {case:?}: {initial:?}"), + } + let id = match &initial { + BrokerToRelaycast::DeliveryAck(frame) => frame.id.clone(), + BrokerToRelaycast::ActionResult(frame) => frame.id.clone(), + BrokerToRelaycast::AgentDeregister(frame) => frame.id.clone(), + _ => unreachable!(), + } + .expect("retryable control frame must carry a stable id"); + + let pressure_responses = if persistent { + server_max_retries + 1 + } else { + 1 + }; + for response_index in 0..pressure_responses { + ws.send(Message::Text( + json!({ + "v": 1, + "id": id.clone(), + "type": "error", + "ok": false, + "code": "d1_pressure", + "message": "Node liveness retry pending" + }) + .to_string(), + )) + .await + .unwrap(); + if !disconnect_during_backoff + && (response_index < pressure_responses - 1 || !persistent) + { + if response_index == 0 { + assert!( + tokio::time::timeout( + Duration::from_millis(20), + next_non_heartbeat_node_to_server(&mut ws), + ) + .await + .is_err(), + "{case:?} retried without backoff" + ); + } + let retry = tokio::time::timeout( + Duration::from_millis(500), + next_non_heartbeat_node_to_server(&mut ws), + ) + .await + .expect("retry should arrive inside the short test budget"); + assert_eq!(retry, initial, "{case:?} retry changed the frame"); + } + } + + if disconnect_during_backoff { + ws.close(None).await.unwrap(); + } else if persistent { + if case == SafeControlRetryCase::DeliveryAck { + assert!( + tokio::time::timeout( + Duration::from_millis(150), + next_non_heartbeat_node_to_server(&mut ws), + ) + .await + .is_err(), + "delivery.ack exceeded its retry bound" + ); + } + } else { + ws.send(Message::Text( + json!({"v": 1, "id": id.clone(), "type": "reply", "ok": true, "data": {}}) + .to_string(), + )) + .await + .unwrap(); + } + server_done_tx.send(()).unwrap(); + if !disconnect_during_backoff { + while ws.next().await.is_some() {} + } + }); + + let config = FleetControlConfig { + ws_url, + node_token: Some("nt_test".to_string()), + node_id: "node-test".to_string(), + node_name: "host-test".to_string(), + broker_version: "broker/test".to_string(), + token_minter: None, + session_token: None, + read_idle_timeout: None, + probe: None, + terminal_reconnect_tx: None, + }; + let mut registration = Some(build_node_register( + &test_manifest(), + "node-test", + "host-test", + "broker/test", + None, + )); + let mut inventory = Vec::new(); + let mut load = FleetLoadSnapshot::default(); + let mut retained_pressure_frames = VecDeque::new(); + let session = run_connected_once_with_retry_policy( + &config, + &mut command_rx, + &event_tx, + &mut registration, + &mut inventory, + &mut load, + Duration::from_secs(3_600), + retry_policy, + &mut retained_pressure_frames, + ); + // Keep the command lane open while a pressure-triggered reconnect is + // expected; otherwise dropping the driver's sender can race the close + // frame and make the session report an intentional shutdown instead. + let _command_lane_guard = command_tx.clone(); + let driver = async move { + assert_eq!(event_rx.recv().await.unwrap(), FleetControlEvent::Connected); + let deregistration_reply = match case { + SafeControlRetryCase::DeliveryAck => { + command_tx + .send(FleetControlCommand::Send(delivery_ack("agent-a", 7))) + .await + .unwrap(); + None + } + SafeControlRetryCase::ActionResult => { + let message = BrokerToRelaycast::ActionResult(ActionResult { + v: FLEET_WIRE_VERSION, + id: None, + invocation_id: "inv-pressure".to_string(), + result: ActionResultPayload::Output(ActionResultOutput { + output: json!({"spawned": true}), + }), + task: None, + }); + if persistent { + command_tx + .send(FleetControlCommand::Send(message)) + .await + .unwrap(); + } else { + let (delivered, _delivery_rx) = oneshot::channel(); + command_tx + .send(FleetControlCommand::RetainedSend(RetainedFleetCompletion { + message, + delivered, + })) + .await + .unwrap(); + } + None + } + SafeControlRetryCase::AgentDeregister => { + let (reply, receiver) = oneshot::channel(); + command_tx + .send(FleetControlCommand::DeregisterAgent { + request: AgentDeregister { + v: FLEET_WIRE_VERSION, + id: None, + agent_id: "agt-pressure".to_string(), + name: Some("agent-pressure".to_string()), + }, + reply, + }) + .await + .unwrap(); + Some(receiver) + } + }; + server_done_rx.await.unwrap(); + let deregistration_result = match deregistration_reply { + Some(receiver) => Some(receiver.await.unwrap()), + None => None, + }; + if !((persistent || disconnect_during_backoff) + && case == SafeControlRetryCase::ActionResult) + { + command_tx + .send(FleetControlCommand::Shutdown) + .await + .unwrap(); + } + deregistration_result + }; + + let ((result, deregistration_result), ()) = + tokio::time::timeout(Duration::from_secs(5), async { + let (result, deregistration_result) = tokio::join!(session, driver); + ((result, deregistration_result), server.await.unwrap()) + }) + .await + .expect("safe control retry case must finish inside its short test budget"); + + if (persistent || disconnect_during_backoff) && case == SafeControlRetryCase::ActionResult { + assert!(matches!(result, ControlRunResult::Disconnected { .. })); + assert_eq!(retained_pressure_frames.len(), 1); + assert!(matches!( + retained_pressure_frames.front(), + Some(BrokerToRelaycast::ActionResult(result)) + if result.id.as_deref().is_some_and(|id| id.starts_with("action_result_")) + && result.invocation_id == "inv-pressure" + )); + + let expected = retained_pressure_frames + .front() + .cloned() + .expect("retained action.result"); + let reconnect_listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let reconnect_ws_url = format!( + "ws://{}/v1/node/ws", + reconnect_listener.local_addr().unwrap() + ); + let (reconnect_done_tx, reconnect_done_rx) = oneshot::channel(); + let reconnect_server = tokio::spawn(async move { + let (stream, _) = reconnect_listener.accept().await.unwrap(); + let mut ws = accept_async(stream).await.unwrap(); + assert!(matches!( + next_node_to_server(&mut ws).await, + BrokerToRelaycast::NodeRegister(_) + )); + assert!(matches!( + next_node_to_server(&mut ws).await, + BrokerToRelaycast::InventorySync(_) + )); + assert!(matches!( + next_node_to_server(&mut ws).await, + BrokerToRelaycast::NodeHeartbeat(_) + )); + let replayed = next_non_heartbeat_node_to_server(&mut ws).await; + assert_eq!(replayed, expected, "reconnect must replay the exact result"); + let id = match replayed { + BrokerToRelaycast::ActionResult(result) => result.id.unwrap(), + _ => unreachable!(), + }; + ws.send(Message::Text( + json!({"v": 1, "id": id, "type": "reply", "ok": true, "data": {}}).to_string(), + )) + .await + .unwrap(); + reconnect_done_tx.send(()).unwrap(); + while ws.next().await.is_some() {} + }); + let reconnect_config = FleetControlConfig { + ws_url: reconnect_ws_url, + node_token: Some("nt_test".to_string()), + node_id: "node-test".to_string(), + node_name: "host-test".to_string(), + broker_version: "broker/test".to_string(), + token_minter: None, + session_token: None, + read_idle_timeout: None, + probe: None, + terminal_reconnect_tx: None, + }; + let (reconnect_command_tx, mut reconnect_command_rx) = mpsc::channel(2); + let (reconnect_event_tx, mut reconnect_event_rx) = mpsc::channel(2); + let reconnect_session = run_connected_once_with_retry_policy( + &reconnect_config, + &mut reconnect_command_rx, + &reconnect_event_tx, + &mut registration, + &mut inventory, + &mut load, + Duration::from_secs(3_600), + retry_policy, + &mut retained_pressure_frames, + ); + let reconnect_driver = async move { + assert_eq!( + reconnect_event_rx.recv().await.unwrap(), + FleetControlEvent::Connected + ); + reconnect_done_rx.await.unwrap(); + reconnect_command_tx + .send(FleetControlCommand::Shutdown) + .await + .unwrap(); + }; + let (reconnect_result, ()) = tokio::time::timeout(Duration::from_secs(5), async { + let (result, ()) = tokio::join!(reconnect_session, reconnect_driver); + (result, reconnect_server.await.unwrap()) + }) + .await + .expect("retained result must replay after reconnect"); + assert_eq!(reconnect_result, ControlRunResult::Shutdown); + assert!(retained_pressure_frames.is_empty()); + } else { + assert_eq!(result, ControlRunResult::Shutdown); + } + match (case, persistent, deregistration_result) { + (SafeControlRetryCase::AgentDeregister, false, Some(Ok(()))) => {} + (SafeControlRetryCase::AgentDeregister, true, Some(Err(reason))) => { + assert!(reason.starts_with("d1_pressure:"), "got {reason}"); + } + (SafeControlRetryCase::AgentDeregister, _, other) => { + panic!("unexpected deregistration result: {other:?}") + } + (_, _, None) => {} + (_, _, other) => panic!("unexpected reply owner: {other:?}"), + } + } + + #[tokio::test] + async fn delivery_ack_retries_transient_pressure_on_same_socket() { + assert_safe_control_frame_socket_retry(SafeControlRetryCase::DeliveryAck, false, false) + .await; + } + + #[tokio::test] + async fn delivery_ack_stops_after_persistent_pressure_budget() { + assert_safe_control_frame_socket_retry(SafeControlRetryCase::DeliveryAck, true, false) + .await; + } + + #[tokio::test] + async fn action_result_retries_transient_pressure_on_same_socket() { + assert_safe_control_frame_socket_retry(SafeControlRetryCase::ActionResult, false, false) + .await; + } + + #[tokio::test] + async fn action_result_is_retained_after_persistent_pressure_budget() { + assert_safe_control_frame_socket_retry(SafeControlRetryCase::ActionResult, true, false) + .await; + } + + #[tokio::test] + async fn action_result_is_retained_when_socket_drops_during_pressure_backoff() { + assert_safe_control_frame_socket_retry(SafeControlRetryCase::ActionResult, false, true) + .await; + } + + #[tokio::test] + async fn agent_deregister_retries_transient_pressure_on_same_socket() { + assert_safe_control_frame_socket_retry(SafeControlRetryCase::AgentDeregister, false, false) + .await; + } + + #[tokio::test] + async fn agent_deregister_fails_named_after_persistent_pressure_budget() { + assert_safe_control_frame_socket_retry(SafeControlRetryCase::AgentDeregister, true, false) + .await; + } + #[tokio::test] async fn registration_d1_pressure_backoff_remains_shutdown_interruptible() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); @@ -7115,26 +7764,17 @@ mod tests { ) { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let ws_url = format!("ws://{}/v1/node/ws", listener.local_addr().unwrap()); - let (command_tx, command_rx) = mpsc::channel(32); + let (command_tx, mut command_rx) = mpsc::channel(32); let (event_tx, mut event_rx) = mpsc::channel(32); let (late_success_tx, late_success_rx) = oneshot::channel(); - - tokio::spawn(run_node_control_client( - FleetControlConfig { - ws_url, - node_token: Some("nt_test".to_string()), - node_id: "node-test".to_string(), - node_name: "host-test".to_string(), - broker_version: "broker/test".to_string(), - token_minter: None, - session_token: None, - read_idle_timeout: None, - probe: None, - terminal_reconnect_tx: None, - }, - command_rx, - event_tx, - )); + let (cleanup_seen_tx, cleanup_seen_rx) = oneshot::channel(); + let retry_policy = ControlRetryPolicy { + initial_delay: Duration::from_millis(10), + max_delay: Duration::from_millis(20), + max_retries: 2, + budget: Duration::from_millis(100), + }; + let server_max_retries = retry_policy.max_retries; let server = tokio::spawn(async move { let (stream, _) = listener.accept().await.unwrap(); @@ -7153,7 +7793,7 @@ mod tests { )); let mut initial = None; - for _ in 0..=CONTROL_D1_PRESSURE_MAX_RETRIES { + for _ in 0..=server_max_retries { let request = match next_non_heartbeat_node_to_server(&mut ws).await { BrokerToRelaycast::AgentRegister(request) => request, other => panic!("expected agent.register, got {other:?}"), @@ -7206,49 +7846,84 @@ mod tests { } other => panic!("expected compensating agent.deregister, got {other:?}"), } + cleanup_seen_tx.send(()).unwrap(); + while ws.next().await.is_some() {} }); - command_tx - .send(FleetControlCommand::RegisterNode { - manifest: test_manifest(), - resume_cursor: None, - }) - .await - .unwrap(); - assert_eq!(event_rx.recv().await.unwrap(), FleetControlEvent::Connected); - - let (reply_tx, reply_rx) = oneshot::channel(); - command_tx - .send(FleetControlCommand::RegisterAgent { - request: AgentRegister { - auto_join_general: Some(false), - v: FLEET_WIRE_VERSION, - id: None, - name: "agent-a".to_string(), - invocation_id: Some("inv-1".to_string()), - session_ref: Some("session-1".to_string()), - resumable: Some(true), - }, - reply: reply_tx, - }) - .await - .unwrap(); - - let error = tokio::time::timeout(Duration::from_secs(30), reply_rx) - .await - .expect("persistent pressure must fail inside the registration deadline") - .unwrap() - .expect_err("persistent d1_pressure must fail closed"); + let config = FleetControlConfig { + ws_url, + node_token: Some("nt_test".to_string()), + node_id: "node-test".to_string(), + node_name: "host-test".to_string(), + broker_version: "broker/test".to_string(), + token_minter: None, + session_token: None, + read_idle_timeout: None, + probe: None, + terminal_reconnect_tx: None, + }; + let mut registration = Some(build_node_register( + &test_manifest(), + "node-test", + "host-test", + "broker/test", + None, + )); + let mut inventory = Vec::new(); + let mut load = FleetLoadSnapshot::default(); + let mut retained_pressure_frames = VecDeque::new(); + let session = run_connected_once_with_retry_policy( + &config, + &mut command_rx, + &event_tx, + &mut registration, + &mut inventory, + &mut load, + Duration::from_secs(3_600), + retry_policy, + &mut retained_pressure_frames, + ); + let driver = async move { + assert_eq!(event_rx.recv().await.unwrap(), FleetControlEvent::Connected); + let (reply_tx, reply_rx) = oneshot::channel(); + command_tx + .send(FleetControlCommand::RegisterAgent { + request: AgentRegister { + auto_join_general: Some(false), + v: FLEET_WIRE_VERSION, + id: None, + name: "agent-a".to_string(), + invocation_id: Some("inv-1".to_string()), + session_ref: Some("session-1".to_string()), + resumable: Some(true), + }, + reply: reply_tx, + }) + .await + .unwrap(); + let error = reply_rx + .await + .unwrap() + .expect_err("persistent d1_pressure must fail closed"); + late_success_tx.send(()).unwrap(); + cleanup_seen_rx.await.unwrap(); + command_tx + .send(FleetControlCommand::Shutdown) + .await + .unwrap(); + error + }; + let ((result, error), ()) = tokio::time::timeout(Duration::from_secs(5), async { + let (result, error) = tokio::join!(session, driver); + ((result, error), server.await.unwrap()) + }) + .await + .expect("short persistent-pressure policy must finish reliably"); + assert_eq!(result, ControlRunResult::Shutdown); assert!( error.starts_with("d1_pressure:"), "spawn caller must see the retryable root cause, got {error}" ); - late_success_tx.send(()).unwrap(); - tokio::time::timeout(Duration::from_secs(1), server) - .await - .unwrap() - .unwrap(); - let _ = command_tx.send(FleetControlCommand::Shutdown).await; } #[tokio::test] @@ -10027,6 +10702,7 @@ mod tests { &mut registrations, &mut deregistrations, &mut liveness, + &mut VecDeque::new(), "node-test", &mut sink, None, diff --git a/packages/cli/src/cli/lib/fleet-spawn-confirmation.test.ts b/packages/cli/src/cli/lib/fleet-spawn-confirmation.test.ts index a32f44558..a6a3cbaa5 100644 --- a/packages/cli/src/cli/lib/fleet-spawn-confirmation.test.ts +++ b/packages/cli/src/cli/lib/fleet-spawn-confirmation.test.ts @@ -176,6 +176,25 @@ describe('fleet spawn confirmation is observable from the requester (#1430)', () expect((error as Error).message).toContain('/tmp/worker-1430.log'); }); + it('maps exhausted agent registration pressure to a retryable spawn error', async () => { + const { client } = createClient(async (name, invocationId) => ({ + invocation_id: invocationId, + action_name: name, + status: 'failed', + error: + "node agent.register failed for agent 'worker-pressure': d1_pressure: Node liveness retry pending (retry budget exhausted)", + })); + + const error = await client.placement + .spawn(spawnInput({ confirm: true, confirmTimeoutMs: 1_000, confirmPollIntervalMs: 10 })) + .catch((caught: unknown) => caught); + + expect(error).toBeInstanceOf(RelayPlacementError); + expect((error as RelayPlacementError).code).toBe('spawn_retryable'); + expect((error as Error).message).toContain('d1_pressure'); + expect((error as Error).message).toContain('retry the spawn'); + }); + // MUST-NOT-FIRE — a healthy node. This is the arm a repaired node represents: // confirmation must not turn a working spawn into an error. it('resolves when the node confirms the spawn completed', async () => { diff --git a/packages/sdk/src/messaging/relaycast-placement.ts b/packages/sdk/src/messaging/relaycast-placement.ts index 26d76ce9f..50a8872d7 100644 --- a/packages/sdk/src/messaging/relaycast-placement.ts +++ b/packages/sdk/src/messaging/relaycast-placement.ts @@ -93,6 +93,8 @@ export class RelayPlacementError extends Error { | 'unmapped_repo' /** The node ran the action and reported a failure. */ | 'spawn_failed' + /** The node rejected spawn setup before mutation and a caller may retry. */ + | 'spawn_retryable' /** * The node accepted the invocation but never reported a terminal result. * A node running an obsolete broker advertises `spawn:` capacity diff --git a/packages/sdk/src/messaging/relaycast.ts b/packages/sdk/src/messaging/relaycast.ts index 74c4f115e..8e89a59d4 100644 --- a/packages/sdk/src/messaging/relaycast.ts +++ b/packages/sdk/src/messaging/relaycast.ts @@ -1070,9 +1070,13 @@ export class RelaycastMessagingClient implements RelayMessagingClient { return invocation; } if (status && CONFIRM_FAILURE_STATUSES.has(status)) { + const nodeError = invocation?.error?.trim(); + const retryablePressure = nodeError?.includes('d1_pressure') === true; throw new RelayPlacementError( - 'spawn_failed', - invocation?.error?.trim() || `node '${context.node}' reported ${status} for ${actionName}`, + retryablePressure ? 'spawn_retryable' : 'spawn_failed', + retryablePressure + ? `${nodeError}. No registration mutation was applied; retry the spawn.` + : nodeError || `node '${context.node}' reported ${status} for ${actionName}`, { ...errorContext, state: 'failed', From 1d56a11946882f24054dbb0118dac264861a6632 Mon Sep 17 00:00:00 2001 From: kjgbot Date: Wed, 7 Oct 2026 16:40:44 -0700 Subject: [PATCH 5/9] fix(broker): bound pressured result replay --- CHANGELOG.md | 2 +- crates/broker/src/node_control.rs | 211 +++++++++++++++--- .../cli/lib/fleet-spawn-confirmation.test.ts | 17 ++ packages/sdk/src/messaging/relaycast.ts | 4 +- 4 files changed, 195 insertions(+), 39 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 740ec8acf..8890d4c51 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,7 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Fixed - Fleet spawns now recover from transient Relaycast `d1_pressure`; persistent registration pressure returns a named `d1_pressure` error within 25 seconds. -- Acknowledged `agent.deregister`, cumulative `delivery.ack`, and action accept/results now recover from transient Relaycast `d1_pressure`; exhausted non-task results are retained across reconnect, and the CLI reports registration pressure as retryable. +- Acknowledged `agent.deregister`, cumulative `delivery.ack`, and non-task action results now recover from transient Relaycast `d1_pressure`; exhausted results are retained for bounded reconnect replay, task-frame pressure returns to its shorter-lived owner, and the CLI reports only proven pre-mutation registration pressure as retryable. - Workers started with `agent-relay fleet spawn` now appear in Agent Relay Desktop as your agents, on the machine that runs them, with their CLI's icon, instead of under "Other live agents". ## [13.1.4] - 2026-10-07 diff --git a/crates/broker/src/node_control.rs b/crates/broker/src/node_control.rs index e6c520f07..558709376 100644 --- a/crates/broker/src/node_control.rs +++ b/crates/broker/src/node_control.rs @@ -2710,6 +2710,7 @@ struct PendingRetryableControlFrame { retry_attempts: u32, retry_at: Option, last_d1_pressure: Option, + replayed_after_pressure: bool, } #[derive(Debug, PartialEq, Eq)] @@ -2757,10 +2758,17 @@ impl ApplicationLiveness { retry_attempts: 0, retry_at: None, last_d1_pressure: None, + replayed_after_pressure: false, }, ); } + fn mark_retryable_control_frame_replayed(&mut self, id: &str) { + if let Some(pending) = self.pending_retryable_control_frames.get_mut(id) { + pending.replayed_after_pressure = true; + } + } + fn complete_retryable_control_frame(&mut self, id: &str) { self.pending_retryable_control_frames.remove(id); } @@ -2823,7 +2831,7 @@ impl ApplicationLiveness { fn expire_retryable_control_frames( &mut self, now: Instant, - ) -> Vec<(String, String, BrokerToRelaycast)> { + ) -> Vec<(String, String, BrokerToRelaycast, bool)> { self.exhausted_retryable_control_frames .retain(|_, exhausted_at| { now.saturating_duration_since(*exhausted_at) < self.control_retry_policy.budget @@ -2852,16 +2860,15 @@ impl ApplicationLiveness { id, format!("d1_pressure: {reason} (retry budget exhausted)"), pending.frame, + pending.replayed_after_pressure, ) }) }) .collect() } - fn take_retryable_control_frame(&mut self, id: &str) -> Option { - self.pending_retryable_control_frames - .remove(id) - .map(|pending| pending.frame) + fn take_retryable_control_frame(&mut self, id: &str) -> Option { + self.pending_retryable_control_frames.remove(id) } fn retain_non_task_action_results( @@ -2874,6 +2881,7 @@ impl ApplicationLiveness { .filter(|(id, pending)| { !id.starts_with(crate::runtime::task_request_prefix()) && matches!(&pending.frame, BrokerToRelaycast::ActionResult(_)) + && pending.last_d1_pressure.is_some() }) .map(|(id, _)| id.clone()) .collect(); @@ -3162,6 +3170,7 @@ async fn handle_connected_command( pending_deregistrations: &mut HashMap>>, application_liveness: &mut ApplicationLiveness, retained_completion: &mut Option, + retained_pressure_frames: &mut VecDeque, ) -> std::ops::ControlFlow where S: Sink + Unpin, @@ -3251,6 +3260,17 @@ where if let BrokerToRelaycast::ActionResult(result) = &message { tracing::warn!(invocation_id = %result.invocation_id, frame_kind = "action.result", "retaining fleet result for replay after reconnect"); + if result.task.is_none() + && !result + .id + .as_deref() + .is_some_and(|id| id.starts_with(crate::runtime::task_request_prefix())) + { + if let Some(id) = result.id.as_deref() { + application_liveness.complete_retryable_control_frame(id); + } + retained_pressure_frames.push_back(message); + } } return Break(ControlRunResult::Disconnected { application_ready: application_liveness.ready, @@ -3552,9 +3572,11 @@ async fn run_connected_once_with_completions( let request_id = prepare_retryable_control_frame(&mut frame) .map(|(id, _)| id) .expect("pressure-retained control frames must remain retryable"); - application_liveness.track_retryable_control_frame(request_id, frame.clone()); + application_liveness.track_retryable_control_frame(request_id.clone(), frame.clone()); + application_liveness.mark_retryable_control_frame_replayed(&request_id); if send_wire(&mut sink, &frame).await.is_err() { - application_liveness.retain_non_task_action_results(retained_pressure_frames); + application_liveness.complete_retryable_control_frame(&request_id); + retained_pressure_frames.push_front(frame); return ControlRunResult::Disconnected { application_ready: false, }; @@ -3596,6 +3618,7 @@ async fn run_connected_once_with_completions( &mut pending_deregistrations, &mut application_liveness, retained_completion, + retained_pressure_frames, ) .await { @@ -3640,7 +3663,7 @@ async fn run_connected_once_with_completions( _ = &mut control_retry, if control_retry_delay.is_some() => { let now = Instant::now(); let mut reconnect_for_retained_result = false; - for (id, reason, frame) in application_liveness.expire_retryable_control_frames(now) { + for (id, reason, frame, replayed_after_pressure) in application_liveness.expire_retryable_control_frames(now) { let mut surfaced_to_owner = false; if let Some(pending) = pending_deregistrations.remove(&id) { let _ = pending.send(Err(reason.clone())); @@ -3662,14 +3685,15 @@ async fn run_connected_once_with_completions( && matches!(&frame, BrokerToRelaycast::ActionResult(_)) { retained_pressure_frames.push_back(frame); - reconnect_for_retained_result = true; + reconnect_for_retained_result |= !replayed_after_pressure; surfaced_to_owner = true; tracing::warn!( target = "relay_broker::fleet", node_id = %config.node_id, id, reason, - "retaining action.result after d1_pressure exhaustion; reconnecting for redelivery" + replayed_after_pressure, + "retaining action.result after d1_pressure exhaustion" ); } if !surfaced_to_owner { @@ -3737,6 +3761,7 @@ async fn run_connected_once_with_completions( &mut pending_deregistrations, &mut application_liveness, retained_completion, + retained_pressure_frames, ) .await { @@ -4009,7 +4034,7 @@ where return true; } ControlPressureRetry::Exhausted => { - let exhausted_frame = application_liveness + let exhausted = application_liveness .take_retryable_control_frame(&error.id); let reason = format!( "d1_pressure: {} (retry budget exhausted)", @@ -4018,21 +4043,27 @@ where if !error .id .starts_with(crate::runtime::task_request_prefix()) - && exhausted_frame.as_ref().is_some_and(|frame| { - matches!(frame, BrokerToRelaycast::ActionResult(_)) + && exhausted.as_ref().is_some_and(|pending| { + matches!( + &pending.frame, + BrokerToRelaycast::ActionResult(_) + ) }) { - retained_pressure_frames.push_back( - exhausted_frame.expect("checked action.result"), - ); + let exhausted = + exhausted.expect("checked action.result"); + let replayed_after_pressure = + exhausted.replayed_after_pressure; + retained_pressure_frames.push_back(exhausted.frame); tracing::warn!( target = "relay_broker::fleet", node_id, id = %error.id, reason, - "retaining action.result after d1_pressure exhaustion; reconnecting for redelivery" + replayed_after_pressure, + "retaining action.result after d1_pressure exhaustion" ); - return false; + return replayed_after_pressure; } if let Some(pending) = pending_deregistrations.remove(&error.id) @@ -4521,8 +4552,6 @@ fn fail_deferred_commands( /// - acknowledged `agent.deregister`, /// - cumulative `delivery.ack` (a missed terminal ack is recovered by durable /// redelivery), and -/// - task `action.accept` / `action.result`, whose request id maps retry -/// exhaustion back to the runtime task owner, and /// - non-task `action.result`, which gets a stable id and is retained for /// redelivery after reconnect if the bounded pressure budget is exhausted. /// @@ -4543,13 +4572,26 @@ fn prepare_retryable_control_frame( .get_or_insert_with(|| format!("delivery_ack_{}", Uuid::new_v4().simple())), "delivery.ack", ), + BrokerToRelaycast::ActionResult(request) + if request.task.is_some() + || request + .id + .as_deref() + .is_some_and(|id| id.starts_with(crate::runtime::task_request_prefix())) => + { + return None; + } BrokerToRelaycast::ActionResult(request) => ( request .id .get_or_insert_with(|| format!("action_result_{}", Uuid::new_v4().simple())), "action.result", ), - BrokerToRelaycast::ActionAccept(request) => (&mut request.id, "action.accept"), + // Task owners have a shorter acknowledgement deadline than this + // socket's pressure budget. Surface their correlated d1_pressure + // immediately so the owner, rather than a stale socket timer, decides + // whether to retry the task. + BrokerToRelaycast::ActionAccept(_) => return None, _ => return None, }; Some((id.clone(), kind)) @@ -4735,6 +4777,7 @@ mod tests { let mut application_liveness = ApplicationLiveness::new(Duration::from_secs(1)); let (delivered, mut delivery_ack) = oneshot::channel(); let mut retained_completion = None; + let mut retained_pressure_frames = VecDeque::new(); let result = handle_connected_command( Some(FleetControlCommand::RetainedSend(RetainedFleetCompletion { message: BrokerToRelaycast::ActionResult(ActionResult { @@ -4759,6 +4802,7 @@ mod tests { &mut pending_deregistrations, &mut application_liveness, &mut retained_completion, + &mut retained_pressure_frames, ) .await; @@ -6161,22 +6205,6 @@ mod tests { name: Some("agent-a".to_string()), }), delivery_ack("agent-a", 7), - BrokerToRelaycast::ActionResult(ActionResult { - v: FLEET_WIRE_VERSION, - id: Some("task_request_result_1".to_string()), - invocation_id: "inv-1".to_string(), - result: ActionResultPayload::Output(ActionResultOutput { - output: json!({"ok": true}), - }), - task: None, - }), - BrokerToRelaycast::ActionAccept(crate::fleet_wire::ActionAccept { - v: FLEET_WIRE_VERSION, - id: "task_request_1".to_string(), - invocation_id: "inv-task".to_string(), - execution_id: "exec-1".to_string(), - worker_generation: "generation-1".to_string(), - }), ]; for mut frame in frames { @@ -6207,6 +6235,25 @@ mod tests { ); assert_eq!(prepare_retryable_control_frame(&mut heartbeat), None); + let mut task_result = BrokerToRelaycast::ActionResult(ActionResult { + v: FLEET_WIRE_VERSION, + id: Some("task_request_result_1".to_string()), + invocation_id: "inv-1".to_string(), + result: ActionResultPayload::Output(ActionResultOutput { + output: json!({"ok": true}), + }), + task: None, + }); + assert_eq!(prepare_retryable_control_frame(&mut task_result), None); + let mut task_accept = BrokerToRelaycast::ActionAccept(crate::fleet_wire::ActionAccept { + v: FLEET_WIRE_VERSION, + id: "task_request_1".to_string(), + invocation_id: "inv-task".to_string(), + execution_id: "exec-1".to_string(), + worker_generation: "generation-1".to_string(), + }); + assert_eq!(prepare_retryable_control_frame(&mut task_accept), None); + let mut one_way_result = BrokerToRelaycast::ActionResult(ActionResult { v: FLEET_WIRE_VERSION, id: None, @@ -6411,6 +6458,96 @@ mod tests { assert!(liveness.exhausted_retryable_control_frames.is_empty()); } + #[test] + fn disconnect_retains_only_action_results_that_observed_pressure() { + let mut liveness = ApplicationLiveness::new(Duration::from_secs(1)); + let id = "action_result_disconnect".to_string(); + let frame = BrokerToRelaycast::ActionResult(ActionResult { + v: FLEET_WIRE_VERSION, + id: Some(id.clone()), + invocation_id: "inv-disconnect".to_string(), + result: ActionResultPayload::Output(ActionResultOutput { + output: json!({"ok": true}), + }), + task: None, + }); + liveness.track_retryable_control_frame(id.clone(), frame.clone()); + let mut retained = VecDeque::new(); + liveness.retain_non_task_action_results(&mut retained); + assert!( + retained.is_empty(), + "a successful write has no replay proof" + ); + + assert!(matches!( + liveness.schedule_control_d1_pressure_retry( + &id, + "Node liveness retry pending", + Instant::now(), + ), + ControlPressureRetry::Scheduled(_) + )); + liveness.retain_non_task_action_results(&mut retained); + assert_eq!(retained, VecDeque::from([frame])); + } + + #[tokio::test] + async fn replayed_action_result_exhaustion_waits_for_an_organic_reconnect() { + let policy = ControlRetryPolicy { + initial_delay: Duration::from_millis(1), + max_delay: Duration::from_millis(1), + max_retries: 0, + budget: Duration::from_millis(10), + }; + let mut liveness = + ApplicationLiveness::new_with_retry_policy(Duration::from_secs(1), policy); + let id = "action_result_replayed".to_string(); + liveness.track_retryable_control_frame( + id.clone(), + BrokerToRelaycast::ActionResult(ActionResult { + v: FLEET_WIRE_VERSION, + id: Some(id.clone()), + invocation_id: "inv-replayed".to_string(), + result: ActionResultPayload::Output(ActionResultOutput { + output: json!({"ok": true}), + }), + task: None, + }), + ); + liveness.mark_retryable_control_frame_replayed(&id); + let (events, _receiver) = mpsc::channel(1); + let mut retained = VecDeque::new(); + let stays_connected = handle_server_message( + Message::Text( + json!({ + "v": 1, + "id": id, + "type": "error", + "ok": false, + "code": "d1_pressure", + "message": "Node liveness retry pending" + }) + .to_string(), + ), + &events, + &mut HashMap::new(), + &mut HashMap::new(), + &mut liveness, + &mut retained, + "node-test", + &mut futures_util::sink::drain(), + None, + None, + ) + .await; + assert!(stays_connected, "a replay must not start a reconnect loop"); + assert_eq!( + retained.len(), + 1, + "the result remains queued for a later reconnect" + ); + } + #[test] fn expire_agent_registrations_bounds_pending_map() { let created_at = Instant::now(); diff --git a/packages/cli/src/cli/lib/fleet-spawn-confirmation.test.ts b/packages/cli/src/cli/lib/fleet-spawn-confirmation.test.ts index a6a3cbaa5..498134efb 100644 --- a/packages/cli/src/cli/lib/fleet-spawn-confirmation.test.ts +++ b/packages/cli/src/cli/lib/fleet-spawn-confirmation.test.ts @@ -195,6 +195,23 @@ describe('fleet spawn confirmation is observable from the requester (#1430)', () expect((error as Error).message).toContain('retry the spawn'); }); + it('does not call a later free-form d1_pressure detail pre-mutation', async () => { + const { client } = createClient(async (name, invocationId) => ({ + invocation_id: invocationId, + action_name: name, + status: 'failed', + error: 'spawn cleanup failed after mutation: d1_pressure remained elevated', + })); + + const error = await client.placement + .spawn(spawnInput({ confirm: true, confirmTimeoutMs: 1_000, confirmPollIntervalMs: 10 })) + .catch((caught: unknown) => caught); + + expect(error).toBeInstanceOf(RelayPlacementError); + expect((error as RelayPlacementError).code).toBe('spawn_failed'); + expect((error as Error).message).not.toContain('No registration mutation was applied'); + }); + // MUST-NOT-FIRE — a healthy node. This is the arm a repaired node represents: // confirmation must not turn a working spawn into an error. it('resolves when the node confirms the spawn completed', async () => { diff --git a/packages/sdk/src/messaging/relaycast.ts b/packages/sdk/src/messaging/relaycast.ts index 8e89a59d4..3497b2385 100644 --- a/packages/sdk/src/messaging/relaycast.ts +++ b/packages/sdk/src/messaging/relaycast.ts @@ -1071,7 +1071,9 @@ export class RelaycastMessagingClient implements RelayMessagingClient { } if (status && CONFIRM_FAILURE_STATUSES.has(status)) { const nodeError = invocation?.error?.trim(); - const retryablePressure = nodeError?.includes('d1_pressure') === true; + const retryablePressure = + errorContext.capability.startsWith('spawn:') && + /^node agent\.register failed for agent '[^']+': d1_pressure:/.test(nodeError ?? ''); throw new RelayPlacementError( retryablePressure ? 'spawn_retryable' : 'spawn_failed', retryablePressure From bd1ccb8ad6fec47e646859520f9cf48f9228800d Mon Sep 17 00:00:00 2001 From: kjgbot Date: Wed, 7 Oct 2026 16:47:36 -0700 Subject: [PATCH 6/9] test(broker): harden pressure ownership regressions --- crates/broker/src/node_control.rs | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/crates/broker/src/node_control.rs b/crates/broker/src/node_control.rs index 558709376..d0c062742 100644 --- a/crates/broker/src/node_control.rs +++ b/crates/broker/src/node_control.rs @@ -6237,7 +6237,10 @@ mod tests { let mut task_result = BrokerToRelaycast::ActionResult(ActionResult { v: FLEET_WIRE_VERSION, - id: Some("task_request_result_1".to_string()), + id: Some(format!( + "{}result_1", + crate::runtime::task_request_prefix() + )), invocation_id: "inv-1".to_string(), result: ActionResultPayload::Output(ActionResultOutput { output: json!({"ok": true}), @@ -8050,7 +8053,7 @@ mod tests { .unwrap(); error }; - let ((result, error), ()) = tokio::time::timeout(Duration::from_secs(5), async { + let ((result, error), ()) = tokio::time::timeout(Duration::from_secs(15), async { let (result, error) = tokio::join!(session, driver); ((result, error), server.await.unwrap()) }) From 8f5b34b6b598747aa8eacc207116de03c2b05173 Mon Sep 17 00:00:00 2001 From: kjgbot Date: Wed, 7 Oct 2026 16:47:52 -0700 Subject: [PATCH 7/9] style(broker): format pressure regression --- crates/broker/src/node_control.rs | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/crates/broker/src/node_control.rs b/crates/broker/src/node_control.rs index d0c062742..bf241d2e8 100644 --- a/crates/broker/src/node_control.rs +++ b/crates/broker/src/node_control.rs @@ -6237,10 +6237,7 @@ mod tests { let mut task_result = BrokerToRelaycast::ActionResult(ActionResult { v: FLEET_WIRE_VERSION, - id: Some(format!( - "{}result_1", - crate::runtime::task_request_prefix() - )), + id: Some(format!("{}result_1", crate::runtime::task_request_prefix())), invocation_id: "inv-1".to_string(), result: ActionResultPayload::Output(ActionResultOutput { output: json!({"ok": true}), From 84506441c9bc283aea91822493a454638c2901d2 Mon Sep 17 00:00:00 2001 From: kjgbot Date: Wed, 7 Oct 2026 17:01:05 -0700 Subject: [PATCH 8/9] test(broker): allow CI slack for pressure retry watchdog --- crates/broker/src/node_control.rs | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/crates/broker/src/node_control.rs b/crates/broker/src/node_control.rs index bf241d2e8..7c61f1339 100644 --- a/crates/broker/src/node_control.rs +++ b/crates/broker/src/node_control.rs @@ -8050,7 +8050,10 @@ mod tests { .unwrap(); error }; - let ((result, error), ()) = tokio::time::timeout(Duration::from_secs(15), async { + // This is only a deadlock watchdog: the retry policy above still bounds + // the behavior under test. Leave enough wall-clock slack for saturated + // shared CI runners executing this test alongside the full shard. + let ((result, error), ()) = tokio::time::timeout(Duration::from_secs(60), async { let (result, error) = tokio::join!(session, driver); ((result, error), server.await.unwrap()) }) From 54c441210ca3cdd07ac66530819bec300902f077 Mon Sep 17 00:00:00 2001 From: kjgbot Date: Wed, 7 Oct 2026 23:56:52 -0700 Subject: [PATCH 9/9] test(broker): make pressure exhaustion deterministic --- CHANGELOG.md | 2 +- crates/broker/src/node_control.rs | 134 +++++++++++++++++++++--------- 2 files changed, 96 insertions(+), 40 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8890d4c51..b299255a8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,7 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Fixed -- Fleet spawns now recover from transient Relaycast `d1_pressure`; persistent registration pressure returns a named `d1_pressure` error within 25 seconds. +- Fleet spawns now recover from transient Relaycast `d1_pressure`; persistent registration pressure returns a named `d1_pressure` error within 25 seconds, and any success arriving after retry-count or wall-clock exhaustion is explicitly deregistered. - Acknowledged `agent.deregister`, cumulative `delivery.ack`, and non-task action results now recover from transient Relaycast `d1_pressure`; exhausted results are retained for bounded reconnect replay, task-frame pressure returns to its shorter-lived owner, and the CLI reports only proven pre-mutation registration pressure as retryable. - Workers started with `agent-relay fleet spawn` now appear in Agent Relay Desktop as your agents, on the machine that runs them, with their CLI's icon, instead of under "Other live agents". diff --git a/crates/broker/src/node_control.rs b/crates/broker/src/node_control.rs index 7c61f1339..314cb3294 100644 --- a/crates/broker/src/node_control.rs +++ b/crates/broker/src/node_control.rs @@ -7896,21 +7896,16 @@ mod tests { let _ = command_tx.send(FleetControlCommand::Shutdown).await; } - #[tokio::test] - async fn node_control_agent_register_persistent_d1_pressure_fails_closed_and_cleans_late_success( + async fn assert_agent_register_pressure_exhaustion_fails_closed_and_cleans_late_success( + retry_policy: ControlRetryPolicy, + expected_attempts: Option, ) { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let ws_url = format!("ws://{}/v1/node/ws", listener.local_addr().unwrap()); let (command_tx, mut command_rx) = mpsc::channel(32); let (event_tx, mut event_rx) = mpsc::channel(32); - let (late_success_tx, late_success_rx) = oneshot::channel(); + let (late_success_tx, mut late_success_rx) = oneshot::channel(); let (cleanup_seen_tx, cleanup_seen_rx) = oneshot::channel(); - let retry_policy = ControlRetryPolicy { - initial_delay: Duration::from_millis(10), - max_delay: Duration::from_millis(20), - max_retries: 2, - budget: Duration::from_millis(100), - }; let server_max_retries = retry_policy.max_retries; let server = tokio::spawn(async move { @@ -7929,37 +7924,66 @@ mod tests { BrokerToRelaycast::NodeHeartbeat(_) )); - let mut initial = None; - for _ in 0..=server_max_retries { - let request = match next_non_heartbeat_node_to_server(&mut ws).await { - BrokerToRelaycast::AgentRegister(request) => request, - other => panic!("expected agent.register, got {other:?}"), + let mut initial: Option = None; + let mut attempts = 0u32; + loop { + enum ServerEvent { + LateSuccess, + Frame(BrokerToRelaycast), + } + let event = tokio::select! { + result = &mut late_success_rx => { + result.unwrap(); + ServerEvent::LateSuccess + } + frame = next_non_heartbeat_node_to_server(&mut ws) => { + ServerEvent::Frame(frame) + } }; - if let Some(initial) = initial.as_ref() { - assert_eq!( - &request, initial, - "every retry must preserve the exact frame" - ); - } else { - initial = Some(request.clone()); + match event { + ServerEvent::LateSuccess => { + assert!(attempts > 0, "pressure must be observed before exhaustion"); + if let Some(expected) = expected_attempts { + assert_eq!(attempts, expected, "retry-count exhaustion attempt count"); + } else { + assert!( + attempts < server_max_retries + 1, + "budget exhaustion must preempt retry-count exhaustion" + ); + } + break; + } + ServerEvent::Frame(BrokerToRelaycast::AgentRegister(request)) => { + attempts = attempts.saturating_add(1); + if let Some(initial) = initial.as_ref() { + assert_eq!( + &request, initial, + "every retry must preserve the exact frame" + ); + } else { + initial = Some(request.clone()); + } + ws.send(Message::Text( + json!({ + "v": 1, + "id": request.id, + "type": "error", + "ok": false, + "code": "d1_pressure", + "message": "Node liveness retry pending" + }) + .to_string(), + )) + .await + .unwrap(); + } + ServerEvent::Frame(other) => { + panic!("expected agent.register, got {other:?}") + } } - ws.send(Message::Text( - json!({ - "v": 1, - "id": request.id, - "type": "error", - "ok": false, - "code": "d1_pressure", - "message": "Node liveness retry pending" - }) - .to_string(), - )) - .await - .unwrap(); } let initial = initial.expect("initial agent.register frame"); let request_id = initial.id.expect("agent.register request id"); - late_success_rx.await.unwrap(); ws.send(Message::Text( serde_json::to_string(&RelaycastToBroker::Reply(crate::fleet_wire::Reply { v: FLEET_WIRE_VERSION, @@ -8050,10 +8074,7 @@ mod tests { .unwrap(); error }; - // This is only a deadlock watchdog: the retry policy above still bounds - // the behavior under test. Leave enough wall-clock slack for saturated - // shared CI runners executing this test alongside the full shard. - let ((result, error), ()) = tokio::time::timeout(Duration::from_secs(60), async { + let ((result, error), ()) = tokio::time::timeout(Duration::from_secs(5), async { let (result, error) = tokio::join!(session, driver); ((result, error), server.await.unwrap()) }) @@ -8066,6 +8087,41 @@ mod tests { ); } + #[tokio::test] + async fn node_control_agent_register_persistent_d1_pressure_fails_closed_and_cleans_late_success( + ) { + let retry_policy = ControlRetryPolicy { + initial_delay: Duration::from_millis(10), + max_delay: Duration::from_millis(20), + max_retries: 2, + // This case proves retry-count exhaustion. Keep the independent + // wall-clock guard nonbinding even on a heavily loaded runner. + budget: Duration::from_secs(30), + }; + assert_agent_register_pressure_exhaustion_fails_closed_and_cleans_late_success( + retry_policy, + Some(retry_policy.max_retries + 1), + ) + .await; + } + + #[tokio::test] + async fn node_control_agent_register_budget_exhaustion_fails_closed_and_cleans_late_success() { + let retry_policy = ControlRetryPolicy { + initial_delay: Duration::from_millis(10), + max_delay: Duration::from_millis(20), + max_retries: 100, + // Shorter than the first retry delay, so the budget path wins by + // construction rather than depending on runner scheduling. + budget: Duration::from_millis(1), + }; + assert_agent_register_pressure_exhaustion_fails_closed_and_cleans_late_success( + retry_policy, + None, + ) + .await; + } + #[tokio::test] async fn node_control_client_mints_initial_token_when_none_supplied() { // Broker startup no longer mints the node token on the API-readiness