diff --git a/CHANGELOG.md b/CHANGELOG.md index 7c5e11155..752bfa6fd 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 agent roster status recovers after a node-control reconnect instead of leaving live workers offline. ## [11.5.4] - 2026-08-12 diff --git a/crates/broker/src/runtime/api.rs b/crates/broker/src/runtime/api.rs index 3cd3ff6fd..7c3302f6c 100644 --- a/crates/broker/src/runtime/api.rs +++ b/crates/broker/src/runtime/api.rs @@ -347,8 +347,29 @@ impl BrokerRuntime { // the worker MCP never re-registers over HTTP. If node binding is // unavailable, fall back to HTTP pre-registration so a tokenless // node (e.g. mint failure) still spawns a working agent. + let mut fleet_registration = None; + let session_ref = super::fleet::fleet_initial_session_ref(&spec); let worker_relay_key = if let Some(token) = agent_token { seed_supplied_agent_token(relaycast_http, &name, &token); + match super::fleet::resolve_fleet_agent_token_identity( + relaycast_http, + fleet_delivery_book, + &name, + &token, + ) + .await + { + Ok(registration) => { + fleet_registration = Some((registration, None, session_ref.clone())); + } + Err(error) => { + tracing::warn!( + worker = %name, + error = %error, + "could not resolve supplied agent token for reconnect inventory" + ); + } + } Some(token) } else { // Derive the session ref from the resolved spec the same way @@ -356,13 +377,12 @@ impl BrokerRuntime { // `harnessConfig.session_id` registers as a resumable session // rather than a fresh spawn. No invocation id exists on the // HTTP path. - let session_ref = super::fleet::fleet_initial_session_ref(&spec); match super::fleet::register_node_agent_token( fleet_control_tx, fleet_delivery_book, name.as_str(), None, - session_ref, + session_ref.clone(), ) .await { @@ -371,7 +391,9 @@ impl BrokerRuntime { worker = %name, "bound agent to node via agent.register for HTTP spawn" ); - Some(token.token) + let relay_key = token.token.clone(); + fleet_registration = Some((token, None, session_ref)); + Some(relay_key) } Err(node_error) => { tracing::warn!( @@ -388,15 +410,35 @@ impl BrokerRuntime { // delivery. Bind it to this node so it is // deliverable, surfacing a loud warning if the // bind fails. - if let Some(warning) = - super::relaycast_events::bind_http_registered_agent_to_node( + let bind_warning = super::relaycast_events::bind_http_registered_agent_to_node( + relaycast_http, + fleet_node_name, + &name, + ) + .await; + if let Some(warning) = bind_warning { + preregistration_warning = Some(warning); + } else { + match super::fleet::resolve_fleet_agent_token_identity( relaycast_http, - fleet_node_name, + fleet_delivery_book, &name, + &token, ) .await - { - preregistration_warning = Some(warning); + { + Ok(registration) => { + fleet_registration = + Some((registration, None, session_ref.clone())); + } + Err(error) => { + tracing::warn!( + worker = %name, + error = %error, + "could not resolve HTTP-registered agent for reconnect inventory" + ); + } + } } Some(token) } @@ -578,6 +620,17 @@ impl BrokerRuntime { .await { Ok(effective_spec) => { + if let Some((token, invocation_id, session_ref)) = fleet_registration.take() + { + super::fleet::record_fleet_inventory_agent( + fleet_control_tx, + fleet_inventory, + &token, + invocation_id, + session_ref, + ) + .await; + } // Prepend relay skill text for small-tier models and CLI harnesses that // need explicit tool guidance to reliably call add_agent / remove_agent. // Skip when relay prompt injection is opted out — relay tools are absent. diff --git a/crates/broker/src/runtime/fleet.rs b/crates/broker/src/runtime/fleet.rs index e3c19b0e3..e61015ef8 100644 --- a/crates/broker/src/runtime/fleet.rs +++ b/crates/broker/src/runtime/fleet.rs @@ -947,6 +947,7 @@ impl BrokerRuntime { &mut self.agent_spawn_count, &self.fleet_control_tx, &mut self.fleet_delivery_book, + &mut self.fleet_inventory, &self.fleet_node_name, Some(invoke.invocation_id.clone()), session_ref, @@ -1354,13 +1355,88 @@ pub(super) async fn publish_fleet_inventory_snapshot( fleet_control_tx: &mpsc::Sender, fleet_inventory: &HashMap, ) { - if let Err(error) = fleet_control_tx.try_send(FleetControlCommand::UpdateInventory( - fleet_inventory.values().cloned().collect(), - )) { - tracing::warn!(error = %error, "fleet inventory queue is unavailable; periodic heartbeat will retry"); + if let Err(error) = fleet_control_tx + .send(FleetControlCommand::UpdateInventory( + fleet_inventory.values().cloned().collect(), + )) + .await + { + tracing::warn!(error = %error, "fleet inventory channel closed; reconnect inventory update was not delivered"); } } +/// Add a successfully launched, node-registered worker to the authoritative +/// reconnect inventory and publish the new snapshot immediately. +/// +/// Relaycast marks every agent hosted by a provider offline when that +/// provider's node-control socket disconnects. The reconnect path can restore +/// those still-running workers only from `inventory.sync`; keeping the +/// registration solely in [`FleetDeliveryBook`] is not enough because that +/// book is delivery-local and is never sent to the engine. +pub(super) async fn record_fleet_inventory_agent( + fleet_control_tx: &mpsc::Sender, + fleet_inventory: &mut HashMap, + token: &crate::node_control::AgentRegistrationToken, + invocation_id: Option, + session_ref: Option, +) { + fleet_inventory.insert( + WorkerName::from(token.name.as_str()), + InventoryAgent { + agent_id: token.agent_id.clone(), + name: token.name.clone(), + invocation_id, + session_ref, + }, + ); + publish_fleet_inventory_snapshot(fleet_control_tx, fleet_inventory).await; +} + +/// Resolve an opaque agent token to the authoritative identity required by +/// `inventory.sync` and delivery bookkeeping. +/// +/// Some supported spawn callers pre-mint the worker token and pass only that +/// credential to the broker. The token itself does not encode an agent id, so +/// resolve it through Relaycast before launch rather than silently omitting the +/// worker from reconnect inventory. +pub(super) async fn resolve_fleet_agent_token_identity( + relaycast_http: &RelaycastHttpClient, + fleet_delivery_book: &mut FleetDeliveryBook, + expected_name: &WorkerName, + token: &str, +) -> Result { + let relay = relaycast_http + .relay_client() + .ok_or_else(|| "relaycast_client_unavailable".to_string())?; + let agent = relay + .get_current_agent(token.to_string()) + .await + .map_err(|error| format!("agent_token_identity_lookup_failed: {error}"))?; + registration_token_for_resolved_agent(fleet_delivery_book, expected_name, token, agent) +} + +fn registration_token_for_resolved_agent( + fleet_delivery_book: &mut FleetDeliveryBook, + expected_name: &WorkerName, + token: &str, + agent: relaycast::Agent, +) -> Result { + if agent.name != expected_name.as_str() { + return Err(format!( + "agent_token_identity_mismatch: expected '{}', resolved '{}'", + expected_name, agent.name + )); + } + + fleet_delivery_book.bind_authoritative_identity(agent.name.clone(), agent.id.clone()); + Ok(crate::node_control::AgentRegistrationToken { + name: agent.name, + agent_id: agent.id, + token: token.to_string(), + delivery_ack_seq: None, + }) +} + pub(super) async fn refresh_fleet_inventory_session_ref( fleet_control_tx: &mpsc::Sender, fleet_inventory: &mut HashMap, @@ -2511,6 +2587,155 @@ mod tests { } } + #[tokio::test] + async fn successful_node_registration_is_added_to_reconnect_inventory() { + let (tx, mut rx) = mpsc::channel::(2); + let mut inventory = HashMap::from([( + WorkerName::from("already-running"), + InventoryAgent { + agent_id: "agent-existing-id".to_string(), + name: "already-running".to_string(), + invocation_id: Some("inv-existing".to_string()), + session_ref: None, + }, + )]); + let token = crate::node_control::AgentRegistrationToken { + name: "new-worker".to_string(), + agent_id: "agent-new-id".to_string(), + token: "at_test".to_string(), + delivery_ack_seq: None, + }; + + record_fleet_inventory_agent( + &tx, + &mut inventory, + &token, + Some("inv-new".to_string()), + Some("session-new".to_string()), + ) + .await; + + assert_eq!(inventory.len(), 2); + assert_eq!( + inventory.get(&WorkerName::from("new-worker")), + Some(&InventoryAgent { + agent_id: "agent-new-id".to_string(), + name: "new-worker".to_string(), + invocation_id: Some("inv-new".to_string()), + session_ref: Some("session-new".to_string()), + }) + ); + match rx.recv().await { + Some(FleetControlCommand::UpdateInventory(agents)) => { + assert_eq!(agents.len(), 2); + assert!(agents.iter().any(|agent| agent.name == "already-running")); + assert!(agents.iter().any(|agent| { + agent.name == "new-worker" + && agent.agent_id == "agent-new-id" + && agent.invocation_id.as_deref() == Some("inv-new") + && agent.session_ref.as_deref() == Some("session-new") + })); + } + other => panic!("expected inventory update, got {other:?}"), + } + } + + #[tokio::test] + async fn fleet_inventory_snapshot_waits_for_backpressure_instead_of_dropping() { + let (tx, mut rx) = mpsc::channel::(1); + tx.send(FleetControlCommand::HeartbeatNow) + .await + .expect("prefill fleet control queue"); + let inventory = HashMap::from([( + WorkerName::from("worker-a"), + InventoryAgent { + agent_id: "agent-a-id".to_string(), + name: "worker-a".to_string(), + invocation_id: None, + session_ref: None, + }, + )]); + let publish = tokio::spawn({ + let tx = tx.clone(); + async move { publish_fleet_inventory_snapshot(&tx, &inventory).await } + }); + + tokio::task::yield_now().await; + assert!( + !publish.is_finished(), + "inventory publication must wait while the queue is full" + ); + assert!(matches!( + rx.recv().await, + Some(FleetControlCommand::HeartbeatNow) + )); + publish.await.expect("inventory publisher should complete"); + match rx.recv().await { + Some(FleetControlCommand::UpdateInventory(agents)) => { + assert_eq!(agents.len(), 1); + assert_eq!(agents[0].name, "worker-a"); + } + other => panic!("expected reliable inventory update, got {other:?}"), + } + } + + #[test] + fn supplied_agent_token_resolves_to_authoritative_fleet_identity() { + let mut delivery_book = FleetDeliveryBook::default(); + let token = registration_token_for_resolved_agent( + &mut delivery_book, + &WorkerName::from("worker-a"), + "at_test", + relaycast::Agent { + id: "agent-a-id".to_string(), + workspace_id: Some("workspace-a".to_string()), + name: "worker-a".to_string(), + agent_type: "agent".to_string(), + status: "active".to_string(), + persona: None, + metadata: serde_json::Map::new(), + created_at: None, + last_seen: None, + channels: Vec::new(), + }, + ) + .expect("matching supplied token identity should resolve"); + + assert_eq!(token.name, "worker-a"); + assert_eq!(token.agent_id, "agent-a-id"); + assert_eq!(token.token, "at_test"); + assert_eq!( + delivery_book.active_agent_id("worker-a"), + Some("agent-a-id") + ); + } + + #[test] + fn supplied_agent_token_rejects_a_different_agent_identity() { + let mut delivery_book = FleetDeliveryBook::default(); + let error = registration_token_for_resolved_agent( + &mut delivery_book, + &WorkerName::from("worker-a"), + "at_test", + relaycast::Agent { + id: "agent-b-id".to_string(), + workspace_id: Some("workspace-a".to_string()), + name: "worker-b".to_string(), + agent_type: "agent".to_string(), + status: "active".to_string(), + persona: None, + metadata: serde_json::Map::new(), + created_at: None, + last_seen: None, + channels: Vec::new(), + }, + ) + .expect_err("a supplied token for another agent must not enter inventory"); + + assert!(error.contains("agent_token_identity_mismatch")); + assert_eq!(delivery_book.active_agent_id("worker-a"), None); + } + #[tokio::test] async fn refresh_fleet_inventory_session_ref_publishes_immediate_sync() { let (tx, mut rx) = mpsc::channel(4); diff --git a/crates/broker/src/runtime/relaycast_events.rs b/crates/broker/src/runtime/relaycast_events.rs index bb5825bac..64f0d9e90 100644 --- a/crates/broker/src/runtime/relaycast_events.rs +++ b/crates/broker/src/runtime/relaycast_events.rs @@ -354,6 +354,7 @@ pub(super) async fn spawn_worker_from_request( agent_spawn_count: &mut u32, fleet_control_tx: &mpsc::Sender, fleet_delivery_book: &mut FleetDeliveryBook, + fleet_inventory: &mut HashMap, node_name: &str, invocation_id: Option, session_ref: Option, @@ -488,11 +489,32 @@ pub(super) async fn spawn_worker_from_request( // injected as RELAY_AGENT_TOKEN (which also sets RELAY_SKIP_BOOTSTRAP), so // the worker MCP never re-registers over HTTP. Falls back to HTTP // pre-registration when node binding is unavailable. + let mut fleet_registration = None; let worker_relay_key = { if let Some(token) = relaycast_ws_spawn_token(ws_value) .filter(|_| !require_node_registration && !relaycast_spawn_verifies_ready(ws_value)) { seed_supplied_agent_token(workspace_http, &name, &token); + match super::fleet::resolve_fleet_agent_token_identity( + workspace_http, + fleet_delivery_book, + &name, + &token, + ) + .await + { + Ok(registration) => { + fleet_registration = + Some((registration, invocation_id.clone(), session_ref.clone())); + } + Err(error) => { + tracing::warn!( + worker = %name, + error = %error, + "could not resolve supplied agent token for reconnect inventory" + ); + } + } Some(token) } else { match super::fleet::register_node_agent_token( @@ -509,7 +531,9 @@ pub(super) async fn spawn_worker_from_request( worker = %name, "bound agent to node via agent.register for action.invoke spawn" ); - Some(token.token) + let relay_key = token.token.clone(); + fleet_registration = Some((token, invocation_id.clone(), session_ref.clone())); + Some(relay_key) } Err(node_error) => { if require_node_registration || relaycast_spawn_verifies_ready(ws_value) { @@ -541,8 +565,37 @@ pub(super) async fn spawn_worker_from_request( // node binding; in node-only delivery the engine only // delivers to `via_node` agents. Bind it to this node // so it becomes deliverable. - bind_http_registered_agent_to_node(workspace_http, node_name, &name) - .await; + let bind_warning = bind_http_registered_agent_to_node( + workspace_http, + node_name, + &name, + ) + .await; + if bind_warning.is_none() { + match super::fleet::resolve_fleet_agent_token_identity( + workspace_http, + fleet_delivery_book, + &name, + &token, + ) + .await + { + Ok(registration) => { + fleet_registration = Some(( + registration, + invocation_id.clone(), + session_ref.clone(), + )); + } + Err(error) => { + tracing::warn!( + worker = %name, + error = %error, + "could not resolve HTTP-registered agent for reconnect inventory" + ); + } + } + } Some(token) } Ok(Err(error)) => { @@ -600,6 +653,16 @@ pub(super) async fn spawn_worker_from_request( .await { Ok(effective_spec) => { + if let Some((token, invocation_id, session_ref)) = fleet_registration.take() { + super::fleet::record_fleet_inventory_agent( + fleet_control_tx, + fleet_inventory, + &token, + invocation_id, + session_ref, + ) + .await; + } if let Some(prefix) = super::api::relay_skill_prefix( effective_spec.cli.as_deref().unwrap_or(&cli), effective_spec.model.as_deref(),