1use agentmux_common::AUTH_KEY_HEADER;
46use std::sync::Arc;
47
48use crate::backend::blockcontroller::health::TurnOrigin;
49use crate::backend::reactive::broadcast_turn_message;
50use crate::backend::rpc::engine::WshRpcEngine;
51use crate::backend::rpc_types::*;
52
53use super::AppState;
54use super::agent_io::stop_one_agent_block;
55use crate::server::agent_handlers::{run_agent_turn, AgentTurnDeps, TurnRegistration};
56
57pub(crate) const BROADCAST_TURN_ORIGIN: TurnOrigin = TurnOrigin::User;
61
62const FLEET_BROADCAST_AUDIT_ACTION: &str = "fleet.broadcast";
63
64const BROADCAST_CONCURRENCY: usize = 8;
69
70pub fn register(engine: &Arc<WshRpcEngine>, state: &AppState) {
71 register_fleet_broadcast(engine, state);
72 register_fleet_bulk_stop(engine, state);
73 register_fleet_group_create(engine, state);
74 register_fleet_group_list(engine, state);
75 register_fleet_group_update(engine, state);
76 register_fleet_group_delete(engine, state);
77}
78
79fn register_fleet_broadcast(engine: &Arc<WshRpcEngine>, state: &AppState) {
80 let state = state.clone();
81 engine.register_typed(
82 COMMAND_FLEET_BROADCAST,
83 move |cmd: CommandFleetBroadcastData, _ctx| {
84 let state = state.clone();
85 async move {
86 let result = fleet_broadcast_impl(&state, cmd.targets, cmd.message).await;
87 Ok(result)
88 }
89 },
90 );
91}
92
93pub(crate) struct BroadcastOutcome {
95 pub block_id: String,
96 pub agent_id: Option<String>,
98 pub result: Result<(), String>,
99}
100
101const NO_REGISTERED_AGENT: &str =
102 "no registered agent for this block (not a live agent pane, or not yet registered)";
103
104pub(crate) async fn run_broadcast<R, D, Fut>(
113 targets: Vec<String>,
114 body: &str,
115 msg_id: &str,
116 concurrency: usize,
117 resolve: R,
118 deliver: D,
119) -> Vec<BroadcastOutcome>
120where
121 R: Fn(&str) -> Option<String>,
122 D: Fn(String, String, TurnOrigin) -> Fut,
123 Fut: std::future::Future<Output = Result<(), String>>,
124{
125 use futures_util::StreamExt;
126
127 let resolved: Vec<(String, Option<String>)> = targets
128 .into_iter()
129 .map(|block_id| {
130 let agent = resolve(&block_id);
131 (block_id, agent)
132 })
133 .collect();
134 let recipients = resolved.iter().filter(|(_, agent)| agent.is_some()).count();
135
136 futures_util::stream::iter(resolved)
137 .map(|(block_id, agent_id)| {
138 let deliver = &deliver;
139 async move {
140 let result = match &agent_id {
141 None => Err(NO_REGISTERED_AGENT.to_string()),
142 Some(agent) => {
143 let text = broadcast_turn_message(agent, recipients, msg_id, body);
144 deliver(block_id.clone(), text, BROADCAST_TURN_ORIGIN).await
145 }
146 };
147 BroadcastOutcome { block_id, agent_id, result }
148 }
149 })
150 .buffered(concurrency.max(1))
151 .collect()
152 .await
153}
154
155#[derive(Debug, PartialEq, Eq)]
157pub(crate) enum BroadcastAction {
158 Done,
160 StartTurn,
162 Fail(String),
163}
164
165pub(crate) fn broadcast_action(route: Result<crate::bootstrap::AgentRoute, String>) -> BroadcastAction {
169 use crate::backend::reactive::SenderDelivery;
170 use crate::bootstrap::AgentRoute;
171 match route {
172 Ok(AgentRoute::Sent(SenderDelivery::Delivered | SenderDelivery::Deferred)) => BroadcastAction::Done,
173 Ok(AgentRoute::Sent(SenderDelivery::Pty)) => BroadcastAction::Fail(
174 "this pane is a terminal, not an agent session; a broadcast is not typed into a shell".to_string(),
175 ),
176 Ok(AgentRoute::StartTurn { .. }) => BroadcastAction::StartTurn,
177 Err(e) => BroadcastAction::Fail(e),
178 }
179}
180
181async fn deliver_broadcast_turn(
189 deps: &AgentTurnDeps,
190 block_id: String,
191 text: String,
192 origin: TurnOrigin,
193) -> Result<(), String> {
194 let input = crate::backend::blockcontroller::health::TurnInput { origin, text: text.clone(), joins_turn: None };
195 match broadcast_action(crate::bootstrap::route_agent_message_from(&block_id, &text, Some(input))) {
196 BroadcastAction::Done => Ok(()),
197 BroadcastAction::Fail(e) => Err(e),
198 BroadcastAction::StartTurn => {
202 run_agent_turn(deps, block_id, text, None, TurnRegistration::Skip, origin, Vec::new()).await
203 }
204 }
205}
206
207enum BroadcastRoute {
209 Local(String),
211 Channel(crate::backend::reactive::registry::AgentEntry),
213}
214
215impl BroadcastRoute {
216 fn agent_id(&self) -> &str {
217 match self {
218 BroadcastRoute::Local(agent) => agent,
219 BroadcastRoute::Channel(entry) => &entry.agent_id,
220 }
221 }
222}
223
224pub(crate) struct BroadcastForward<'a> {
227 pub block_id: &'a str,
228 pub body: &'a str,
229 pub msg_id: &'a str,
230 pub recipients: usize,
232}
233
234pub(crate) async fn forward_broadcast_to_channel(
239 client: &reqwest::Client,
240 local_url: &str,
241 auth_key: &str,
242 req: &BroadcastForward<'_>,
243) -> Result<(), String> {
244 let mut http = client
245 .post(format!("{local_url}/agentmux/agent/broadcast"))
246 .json(&serde_json::json!({
247 "block_id": req.block_id,
248 "body": req.body,
249 "msg_id": req.msg_id,
250 "recipients": req.recipients,
251 }));
252 if !auth_key.is_empty() {
253 http = http.header(AUTH_KEY_HEADER, auth_key);
254 }
255 let resp = http.send().await.map_err(|e| format!("cross-channel forward failed: {e}"))?;
256 if resp.status() == reqwest::StatusCode::NOT_FOUND {
257 return Err("that channel is on an older build and can't receive a broadcast from here".to_string());
258 }
259 if !resp.status().is_success() {
260 return Err(format!("cross-channel forward: HTTP {}", resp.status()));
261 }
262 let body: serde_json::Value = resp
263 .json()
264 .await
265 .map_err(|e| format!("cross-channel forward: response parse failed: {e}"))?;
266 if body.get("success").and_then(|v| v.as_bool()) == Some(true) {
267 Ok(())
268 } else {
269 Err(body.get("error").and_then(|v| v.as_str()).unwrap_or("cross-channel forward failed").to_string())
270 }
271}
272
273pub(crate) async fn deliver_forwarded_broadcast(
278 state: &AppState,
279 block_id: &str,
280 body: &str,
281 msg_id: &str,
282 recipients: usize,
283) -> Result<(), String> {
284 let result = match state.reactive_handler.get_agent_by_block(block_id) {
285 None => Err(NO_REGISTERED_AGENT.to_string()),
286 Some(agent) => {
287 let text = broadcast_turn_message(&agent.agent_id, recipients, msg_id, body);
288 let deps = AgentTurnDeps::from_state(state);
289 deliver_broadcast_turn(&deps, block_id.to_string(), text, BROADCAST_TURN_ORIGIN).await
290 }
291 };
292 let target = state
293 .reactive_handler
294 .get_agent_by_block(block_id)
295 .map(|a| a.agent_id)
296 .unwrap_or_else(|| block_id.to_string());
297 state.reactive_handler.log_fleet_action_audit(
298 None,
299 &target,
300 block_id,
301 FLEET_BROADCAST_AUDIT_ACTION,
302 result.is_ok(),
303 result.as_ref().err().map(String::as_str),
304 msg_id,
305 Some("forwarded from another channel on this machine"),
306 );
307 result
308}
309
310pub(crate) async fn fleet_broadcast_impl(
311 state: &AppState,
312 targets: Vec<String>,
313 message: String,
314) -> FleetActionResult {
315 let deps = AgentTurnDeps::from_state(state);
316 let msg_id = uuid::Uuid::new_v4().to_string();
317 let routes: std::collections::HashMap<String, BroadcastRoute> = targets
320 .iter()
321 .filter_map(|block_id| {
322 let route = match state.reactive_handler.get_agent_by_block(block_id) {
323 Some(agent) => BroadcastRoute::Local(agent.agent_id),
324 None => BroadcastRoute::Channel(shared_channel_for(state, block_id)?),
325 };
326 Some((block_id.clone(), route))
327 })
328 .collect();
329 let routes = Arc::new(routes);
330 let recipients = targets.iter().filter(|b| routes.contains_key(*b)).count();
333 let outcomes = run_broadcast(
334 targets,
335 &message,
336 &msg_id,
337 BROADCAST_CONCURRENCY,
338 |block_id| routes.get(block_id).map(|r| r.agent_id().to_string()),
339 |block_id, text, origin| {
340 let deps = deps.clone();
341 let routes = routes.clone();
342 let client = state.http_client.clone();
343 let (body, msg_id) = (message.clone(), msg_id.clone());
344 async move {
345 match routes.get(&block_id) {
346 Some(BroadcastRoute::Channel(entry)) => {
347 let forward = BroadcastForward { block_id: &block_id, body: &body, msg_id: &msg_id, recipients };
348 forward_broadcast_to_channel(&client, &entry.local_url, &entry.auth_key, &forward).await
349 }
350 _ => deliver_broadcast_turn(&deps, block_id, text, origin).await,
351 }
352 }
353 },
354 )
355 .await;
356
357 let mut result = FleetActionResult::default();
358 for outcome in outcomes {
359 let target_agent = outcome.agent_id.as_deref().unwrap_or(&outcome.block_id);
360 match outcome.result {
361 Ok(()) => {
362 state.reactive_handler.log_fleet_action_audit(
363 None, target_agent, &outcome.block_id, FLEET_BROADCAST_AUDIT_ACTION,
364 true, None, &msg_id, None,
365 );
366 result.succeeded.push(outcome.block_id);
367 }
368 Err(error) => {
369 state.reactive_handler.log_fleet_action_audit(
370 None, target_agent, &outcome.block_id, FLEET_BROADCAST_AUDIT_ACTION,
371 false, Some(&error), &msg_id, None,
372 );
373 result.failed.push(FleetActionFailure { id: outcome.block_id, error });
374 }
375 }
376 }
377 result
378}
379
380
381fn register_fleet_bulk_stop(engine: &Arc<WshRpcEngine>, state: &AppState) {
382 let state = state.clone();
383 engine.register_typed(
384 COMMAND_FLEET_BULK_STOP,
385 move |cmd: CommandFleetBulkStopData, _ctx| {
386 let state = state.clone();
387 async move {
388 let result = fleet_bulk_stop_impl(&state, cmd.targets, cmd.signal.as_deref(), cmd.staged).await;
389 Ok(result)
390 }
391 },
392 );
393}
394
395const FLEET_BULK_STOP_AUDIT_ACTION: &str = "fleet.bulk-stop";
396
397pub(crate) async fn forward_stop_to_shared_channel(
427 state: &AppState,
428 block_id: &str,
429 signal: Option<&str>,
430) -> Option<(String, Result<(), String>)> {
431 let entry = shared_channel_for(state, block_id)?;
432
433 let url = format!("{}/agentmux/agent/stop", entry.local_url);
434 let mut req = state.http_client.post(&url).json(&serde_json::json!({
435 "block_id": block_id,
436 "signal": signal,
437 }));
438 if !entry.auth_key.is_empty() {
439 req = req.header(AUTH_KEY_HEADER, &entry.auth_key);
440 }
441 let outcome = async {
442 let resp = req
443 .send()
444 .await
445 .map_err(|e| format!("cross-channel forward failed: {e}"))?;
446 if !resp.status().is_success() {
447 return Err(format!("cross-channel forward: HTTP {}", resp.status()));
448 }
449 let body: serde_json::Value = resp
450 .json()
451 .await
452 .map_err(|e| format!("cross-channel forward: response parse failed: {e}"))?;
453 if body.get("success").and_then(|v| v.as_bool()) == Some(true) {
454 Ok(())
455 } else {
456 Err(body
457 .get("error")
458 .and_then(|v| v.as_str())
459 .unwrap_or("cross-channel forward failed")
460 .to_string())
461 }
462 }
463 .await;
464
465 Some((entry.agent_id, outcome))
466}
467
468pub(crate) async fn fleet_bulk_stop_impl(
469 state: &AppState,
470 targets: Vec<String>,
471 signal: Option<&str>,
472 staged: Option<StagePlanInput>,
473) -> FleetActionResult {
474 let mut result = FleetActionResult::default();
475 let batch_size = staged.as_ref().map(|s| s.batch_size.max(1)).unwrap_or(targets.len().max(1));
476 let max_fail_percentage = staged.as_ref().map(|s| s.max_fail_percentage);
477
478 let mut iter = targets.into_iter().peekable();
479 'batches: while iter.peek().is_some() {
480 let batch: Vec<String> = (&mut iter).take(batch_size).collect();
481 let batch_len = batch.len();
482 let mut batch_failures = 0usize;
483 for block_id in batch {
484 let request_id = uuid::Uuid::new_v4().to_string();
485 let (target_agent, outcome) = match state.reactive_handler.get_agent_by_block(&block_id) {
491 Some(agent) => (agent.agent_id, stop_one_agent_block(state, &block_id, signal, None).await.map(|_| ())),
492 None => match forward_stop_to_shared_channel(state, &block_id, signal).await {
493 Some((agent_name, outcome)) => (agent_name, outcome),
494 None => (block_id.clone(), stop_one_agent_block(state, &block_id, signal, None).await.map(|_| ())),
495 },
496 };
497 match outcome {
498 Ok(()) => {
499 state.reactive_handler.log_fleet_action_audit(
500 None, &target_agent, &block_id, FLEET_BULK_STOP_AUDIT_ACTION,
501 true, None, &request_id, None,
502 );
503 result.succeeded.push(block_id);
504 }
505 Err(e) => {
506 state.reactive_handler.log_fleet_action_audit(
507 None, &target_agent, &block_id, FLEET_BULK_STOP_AUDIT_ACTION,
508 false, Some(&e), &request_id, None,
509 );
510 batch_failures += 1;
511 result.failed.push(FleetActionFailure { id: block_id, error: e });
512 }
513 }
514 }
515 if let Some(max_pct) = max_fail_percentage {
516 if exceeds_fail_threshold(batch_failures, batch_len, max_pct) {
517 let mut skipped_any = false;
521 for remaining in iter {
522 skipped_any = true;
523 result.failed.push(FleetActionFailure {
524 id: remaining,
525 error: "skipped: staged rollout aborted after a prior batch's failure rate exceeded max_fail_percentage".to_string(),
526 });
527 }
528 result.aborted_early = skipped_any;
532 break 'batches;
533 }
534 }
535 }
536 result
537}
538
539fn shared_channel_for(state: &AppState, block_id: &str) -> Option<crate::backend::reactive::registry::AgentEntry> {
542 let shared_dir = crate::registry::resolve_shared_reactive_dir()?;
543 let entry = crate::backend::reactive::registry::list_all_shared(&shared_dir)
544 .into_iter()
545 .find(|e| e.block_id == block_id)?;
546 let is_loopback = entry.local_url.starts_with("http://127.0.0.1")
547 || entry.local_url.starts_with("http://localhost")
548 || entry.local_url.starts_with("http://[::1]");
549 if !is_loopback || entry.local_url == state.local_web_url {
550 return None;
551 }
552 Some(entry)
553}
554
555enum Forwarded {
557 Pending(crate::sagas::pending_shutdown::PendingView),
559 Immediate(Result<(), String>),
561 Refused(String),
563}
564
565async fn forward_stop_pending(
571 state: &AppState,
572 entry: &crate::backend::reactive::registry::AgentEntry,
573 block_id: &str,
574 by: &str,
575 signal: Option<&str>,
576) -> Forwarded {
577 let url = format!("{}/agentmux/agent/stop-pending", entry.local_url);
578 let mut req = state.http_client.post(&url).json(&serde_json::json!({
579 "block_id": block_id,
580 "signal": signal,
581 "by": by,
582 }));
583 if !entry.auth_key.is_empty() {
584 req = req.header(AUTH_KEY_HEADER, &entry.auth_key);
585 }
586 let resp = match req.send().await {
587 Ok(r) => r,
588 Err(e) => return Forwarded::Refused(format!("cross-channel forward failed: {e}")),
589 };
590 if resp.status() == reqwest::StatusCode::NOT_FOUND {
591 let outcome = forward_stop_to_shared_channel(state, block_id, signal)
592 .await
593 .map(|(_, o)| o)
594 .unwrap_or_else(|| Err("cross-channel forward: target vanished".to_string()));
595 return Forwarded::Immediate(outcome);
596 }
597 if !resp.status().is_success() {
598 return Forwarded::Refused(format!("cross-channel forward: HTTP {}", resp.status()));
599 }
600 let body: serde_json::Value = match resp.json().await {
601 Ok(b) => b,
602 Err(e) => return Forwarded::Refused(format!("cross-channel forward: response parse failed: {e}")),
603 };
604 if body["success"] != true {
605 return Forwarded::Refused(body["error"].as_str().unwrap_or("cross-channel forward failed").to_string());
606 }
607 let p = &body["pending"];
608 let Some(request_id) = p["request_id"].as_str() else {
609 return Forwarded::Refused("cross-channel forward: no request id".to_string());
610 };
611 crate::sagas::pending_shutdown::remember_remote(request_id, &entry.local_url, &entry.auth_key);
612 Forwarded::Pending(crate::sagas::pending_shutdown::PendingView {
613 request_id: request_id.to_string(),
614 block_id: block_id.to_string(),
615 by: p["by"].as_str().unwrap_or(by).to_string(),
616 target: entry.agent_id.clone(),
617 via: p["via"].as_str().unwrap_or("FleetBulkStop").to_string(),
618 reason: String::new(),
619 deadline_ms: p["deadline_ms"].as_i64().unwrap_or_default(),
620 status: "pending",
621 error: None,
622 joined: p["joined"].as_bool().unwrap_or(false),
623 })
624}
625
626pub(crate) async fn fleet_bulk_stop_with_override(
634 state: &AppState,
635 by: &str,
636 targets: Vec<String>,
637 signal: Option<&str>,
638 staged: Option<StagePlanInput>,
639) -> (FleetActionResult, Vec<crate::sagas::pending_shutdown::PendingView>) {
640 let (local, other): (Vec<String>, Vec<String>) = targets
641 .into_iter()
642 .partition(|b| crate::backend::blockcontroller::get_controller(b).is_some());
643 let mut pending: Vec<crate::sagas::pending_shutdown::PendingView> = local
644 .iter()
645 .map(|block_id| {
646 let mut v = crate::sagas::pending_shutdown::request(
647 state,
648 block_id,
649 by,
650 "FleetBulkStop",
651 "",
652 crate::sagas::pending_shutdown::Action::Stop { signal: signal.map(str::to_string) },
653 );
654 let target = state
659 .reactive_handler
660 .get_agent_by_block(block_id)
661 .map(|a| a.agent_id)
662 .unwrap_or_else(|| block_id.clone());
663 state.reactive_handler.log_fleet_action_audit(
664 Some(by), &target, block_id, FLEET_BULK_STOP_AUDIT_ACTION,
665 true, None, &v.request_id, Some(&crate::sagas::pending_shutdown::audit_note(&v, by, "FleetBulkStop")),
666 );
667 v.block_id = block_id.clone();
668 v.target = target;
669 v
670 })
671 .collect();
672 let mut now = FleetActionResult::default();
673 let mut remote_pending = Vec::new();
674 let mut nowhere = Vec::new();
675 let mut forwards = Vec::new();
679 for block_id in other {
680 match shared_channel_for(state, &block_id) {
681 Some(entry) => forwards.push((block_id, entry)),
682 None => nowhere.push(block_id),
683 }
684 }
685 let answers = futures_util::future::join_all(
686 forwards.iter().map(|(block_id, entry)| forward_stop_pending(state, entry, block_id, by, signal)),
687 )
688 .await;
689 for ((block_id, entry), answer) in forwards.into_iter().zip(answers) {
690 let request_id = uuid::Uuid::new_v4().to_string();
691 match answer {
692 Forwarded::Pending(v) => {
693 state.reactive_handler.log_fleet_action_audit(
694 Some(by), &entry.agent_id, &block_id, FLEET_BULK_STOP_AUDIT_ACTION,
695 true, None, &v.request_id,
696 Some(&format!("pending user override on channel {}", entry.channel)),
697 );
698 remote_pending.push(v);
699 }
700 Forwarded::Immediate(outcome) => {
701 let note = format!("channel {} predates the user override: stopped at once", entry.channel);
702 state.reactive_handler.log_fleet_action_audit(
703 Some(by), &entry.agent_id, &block_id, FLEET_BULK_STOP_AUDIT_ACTION,
704 outcome.is_ok(), outcome.as_ref().err().map(String::as_str), &request_id, Some(¬e),
705 );
706 match outcome {
707 Ok(()) => now.succeeded.push(block_id),
708 Err(error) => now.failed.push(FleetActionFailure { id: block_id, error }),
709 }
710 }
711 Forwarded::Refused(error) => {
712 state.reactive_handler.log_fleet_action_audit(
713 Some(by), &entry.agent_id, &block_id, FLEET_BULK_STOP_AUDIT_ACTION,
714 false, Some(&error), &request_id, None,
715 );
716 now.failed.push(FleetActionFailure { id: block_id, error });
717 }
718 }
719 }
720 pending.extend(remote_pending);
721 if !nowhere.is_empty() {
722 let rest = fleet_bulk_stop_impl(state, nowhere, signal, staged).await;
723 now.succeeded.extend(rest.succeeded);
724 now.failed.extend(rest.failed);
725 now.aborted_early = rest.aborted_early;
726 }
727 (now, pending)
728}
729
730fn exceeds_fail_threshold(batch_failures: usize, batch_len: usize, max_pct: u8) -> bool {
738 batch_failures * 100 > max_pct as usize * batch_len
739}
740
741#[cfg(test)]
742mod threshold_tests {
743 use super::exceeds_fail_threshold;
744
745 #[test]
746 fn one_in_three_at_33_percent_threshold_now_trips() {
747 assert!(exceeds_fail_threshold(1, 3, 33));
751 }
752
753 #[test]
754 fn one_in_three_at_34_percent_threshold_does_not_trip() {
755 assert!(!exceeds_fail_threshold(1, 3, 34));
757 }
758
759 #[test]
760 fn zero_failures_never_trips() {
761 assert!(!exceeds_fail_threshold(0, 10, 0));
762 }
763
764 #[test]
765 fn all_failures_trips_any_threshold_below_100() {
766 assert!(exceeds_fail_threshold(5, 5, 99));
767 assert!(!exceeds_fail_threshold(5, 5, 100));
768 }
769
770 #[test]
771 fn two_in_seven_at_28_percent_threshold() {
772 assert!(exceeds_fail_threshold(2, 7, 28));
774 assert!(!exceeds_fail_threshold(2, 7, 29));
775 }
776}
777
778use agentmux_common::time::now_ms;
779
780fn register_fleet_group_create(engine: &Arc<WshRpcEngine>, state: &AppState) {
781 let mstore = state.mstore.clone();
782 engine.register_typed(
783 COMMAND_FLEET_GROUP_CREATE,
784 move |cmd: CommandFleetGroupCreateData, _ctx| {
785 let mstore = mstore.clone();
786 async move {
787 if cmd.name.trim().is_empty() {
788 return Err("fleet.group.create: name is required".to_string());
789 }
790 let id = uuid::Uuid::new_v4().to_string();
791 let created_at = now_ms();
792 mstore
793 .agent_group_create(&id, cmd.name.trim(), &cmd.member_ids, created_at)
794 .map_err(|e| format!("fleet.group.create: {e}"))?;
795 Ok(FleetGroup {
796 id,
797 name: cmd.name.trim().to_string(),
798 member_ids: cmd.member_ids,
799 created_at,
800 })
801 }
802 },
803 );
804}
805
806fn register_fleet_group_list(engine: &Arc<WshRpcEngine>, state: &AppState) {
807 let mstore = state.mstore.clone();
808 engine.register_typed(
809 COMMAND_FLEET_GROUP_LIST,
810 move |_req: CommandFleetGroupListData, _ctx| {
811 let mstore = mstore.clone();
812 async move {
813 let groups = mstore.agent_group_list().map_err(|e| format!("fleet.group.list: {e}"))?;
814 let groups = groups
815 .into_iter()
816 .map(|g| FleetGroup { id: g.id, name: g.name, member_ids: g.member_ids, created_at: g.created_at })
817 .collect();
818 Ok(FleetGroupListResult { groups })
819 }
820 },
821 );
822}
823
824fn register_fleet_group_update(engine: &Arc<WshRpcEngine>, state: &AppState) {
825 let mstore = state.mstore.clone();
826 engine.register_typed(
827 COMMAND_FLEET_GROUP_UPDATE,
828 move |cmd: CommandFleetGroupUpdateData, _ctx| {
829 let mstore = mstore.clone();
830 async move {
831 let name = cmd.name.as_deref().map(|n| n.trim());
832 if matches!(name, Some("")) {
833 return Err("fleet.group.update: name cannot be blank".to_string());
834 }
835 let updated = mstore
836 .agent_group_update(&cmd.id, name, cmd.member_ids.as_deref())
837 .map_err(|e| format!("fleet.group.update: {e}"))?;
838 if !updated {
839 return Err(format!("fleet.group.update: no group with id {}", cmd.id));
840 }
841 let group = mstore
842 .agent_group_get(&cmd.id)
843 .map_err(|e| format!("fleet.group.update: {e}"))?
844 .ok_or_else(|| format!("fleet.group.update: group {} vanished after update", cmd.id))?;
845 Ok(FleetGroup {
846 id: group.id,
847 name: group.name,
848 member_ids: group.member_ids,
849 created_at: group.created_at,
850 })
851 }
852 },
853 );
854}
855
856fn register_fleet_group_delete(engine: &Arc<WshRpcEngine>, state: &AppState) {
857 let mstore = state.mstore.clone();
858 engine.register_typed(
859 COMMAND_FLEET_GROUP_DELETE,
860 move |cmd: CommandFleetGroupDeleteData, _ctx| {
861 let mstore = mstore.clone();
862 async move {
863 let deleted = mstore.agent_group_delete(&cmd.id).map_err(|e| format!("fleet.group.delete: {e}"))?;
864 Ok(FleetGroupDeleteResult { ok: deleted })
865 }
866 },
867 );
868}
869
870#[cfg(test)]
871mod override_tests {
872 use super::*;
873 use std::sync::atomic::{AtomicBool, Ordering};
874
875 struct Stoppable {
877 block_id: String,
878 stopped: Arc<AtomicBool>,
879 }
880
881 impl crate::backend::blockcontroller::Controller for Stoppable {
882 fn start(&self, _: crate::backend::obj::MetaMapType, _: Option<serde_json::Value>, _: bool) -> Result<(), String> {
883 Ok(())
884 }
885 fn stop(&self, _: bool, _: &str) -> Result<(), String> {
886 self.stopped.store(true, Ordering::SeqCst);
887 Ok(())
888 }
889 fn get_runtime_status(&self) -> crate::backend::blockcontroller::BlockControllerRuntimeStatus {
890 crate::backend::blockcontroller::BlockControllerRuntimeStatus { blockid: self.block_id.clone(), ..Default::default() }
891 }
892 fn send_input(&self, _: crate::backend::blockcontroller::BlockInputUnion, _: Option<u64>) -> Result<(), String> {
893 Ok(())
894 }
895 fn controller_type(&self) -> &str {
896 "stub"
897 }
898 fn block_id(&self) -> &str {
899 &self.block_id
900 }
901 fn as_any(&self) -> &dyn std::any::Any {
902 self
903 }
904 }
905
906 fn running(prefix: &str) -> (String, Arc<AtomicBool>) {
907 let block_id = format!("{prefix}-{}", uuid::Uuid::new_v4());
908 let stopped = Arc::new(AtomicBool::new(false));
909 crate::backend::blockcontroller::register_controller(
910 &block_id,
911 Arc::new(Stoppable { block_id: block_id.clone(), stopped: stopped.clone() }),
912 );
913 (block_id, stopped)
914 }
915
916 #[tokio::test]
917 async fn an_agent_s_bulk_stop_waits_for_each_user_and_stops_when_nobody_keeps_it() {
918 let state = crate::server::tests::test_state();
919 let (a, a_stopped) = running("bulk-override-a");
920 let (b, b_stopped) = running("bulk-override-b");
921 let gone = format!("bulk-override-gone-{}", uuid::Uuid::new_v4());
922
923 let (now, pending) =
924 fleet_bulk_stop_with_override(&state, "Korp", vec![a.clone(), b.clone(), gone.clone()], None, None).await;
925 assert_eq!(now.failed.iter().map(|f| f.id.clone()).collect::<Vec<_>>(), vec![gone], "not running: fails at once");
926 assert_eq!(pending.iter().map(|p| p.block_id.clone()).collect::<Vec<_>>(), vec![a.clone(), b.clone()]);
927 assert!(pending.iter().all(|p| p.by == "Korp" && p.via == "FleetBulkStop"));
928 assert!(!a_stopped.load(Ordering::SeqCst) && !b_stopped.load(Ordering::SeqCst), "nothing stops during the window");
929
930 assert_eq!(crate::sagas::pending_shutdown::keep(&state, &b, &pending[1].request_id), "kept_by_user");
932 crate::sagas::pending_shutdown::expire_now(&pending[0].request_id);
933 assert_eq!(crate::sagas::pending_shutdown::settled(&pending[0].request_id).await.status, "shut_down");
934 assert_eq!(crate::sagas::pending_shutdown::settled(&pending[1].request_id).await.status, "kept_by_user");
935 assert!(a_stopped.load(Ordering::SeqCst));
936 assert!(!b_stopped.load(Ordering::SeqCst), "the user kept it");
937 crate::backend::blockcontroller::delete_controller(&a);
938 crate::backend::blockcontroller::delete_controller(&b);
939 }
940
941 #[tokio::test]
942 async fn a_target_that_joins_its_pane_s_pending_close_is_reported_under_its_own_block() {
943 let state = crate::server::tests::test_state();
944 let (tab, _) = running("bulk-join-tab");
945 let sibling = format!("bulk-join-sibling-{}", uuid::Uuid::new_v4());
946 let close = crate::sagas::pending_shutdown::request(
947 &state,
948 &sibling,
949 "Korp",
950 "ClosePane",
951 "",
952 crate::sagas::pending_shutdown::Action::ClosePane { block_ids: vec![sibling.clone(), tab.clone()] },
953 );
954
955 let (_, pending) = fleet_bulk_stop_with_override(&state, "Posa", vec![tab.clone()], None, None).await;
956 assert_eq!(pending.len(), 1);
957 assert_eq!(pending[0].block_id, tab, "the block the caller named, not the sibling");
958 assert_eq!(pending[0].request_id, close.request_id, "the pane's one window");
959 assert!(pending[0].joined);
960 crate::sagas::pending_shutdown::keep(&state, &sibling, &close.request_id);
961 crate::backend::blockcontroller::delete_controller(&tab);
962 }
963}
964
965#[cfg(test)]
966mod broadcast_core_tests {
967 use super::*;
968 use std::sync::atomic::{AtomicUsize, Ordering};
969 use std::sync::{Arc, Mutex};
970
971 fn ids(n: usize) -> Vec<String> {
972 (0..n).map(|i| format!("blk-{i}")).collect()
973 }
974
975 fn agent_for(block: &str) -> Option<String> {
976 block.strip_prefix("blk-").map(|n| format!("agent-{n}"))
977 }
978
979 #[test]
982 fn broadcast_turns_are_attributed_to_the_user() {
983 assert_eq!(BROADCAST_TURN_ORIGIN, TurnOrigin::User);
984 }
985
986 #[tokio::test]
987 async fn every_target_is_delivered_with_the_broadcast_origin() {
988 let seen: Arc<Mutex<Vec<TurnOrigin>>> = Arc::default();
989 let seen2 = seen.clone();
990 let out = run_broadcast(ids(3), "hi", "m1", 4, agent_for, move |_b, _t, origin| {
991 seen2.lock().unwrap().push(origin);
992 async { Ok(()) }
993 })
994 .await;
995 assert_eq!(out.len(), 3);
996 assert!(seen.lock().unwrap().iter().all(|o| *o == TurnOrigin::User));
997 }
998
999 #[tokio::test]
1003 async fn the_self_quit_gate_accepts_a_broadcast_turn_like_a_typed_one() {
1004 use crate::backend::blockcontroller::health::TurnProvenance;
1005 use crate::sagas::self_quit::{gate, GateRefusal};
1006 let text = broadcast_turn_message("agent-1", 3, "m1", "finish the PR then quit");
1007 let p = |tainted| TurnProvenance { origin: BROADCAST_TURN_ORIGIN, tainted, user_text: Some(text.clone()) };
1008 assert_eq!(gate(Some(&p(false)), "then quit"), Ok(()));
1009 assert_eq!(gate(Some(&p(false)), "delete everything"), Err(GateRefusal::InstructionMismatch));
1010 assert_eq!(gate(Some(&p(true)), "then quit"), Err(GateRefusal::Tainted));
1011 }
1012
1013 #[tokio::test]
1014 async fn each_target_gets_its_own_header_with_a_shared_id_and_the_resolved_count() {
1015 let texts: Arc<Mutex<Vec<(String, String)>>> = Arc::default();
1016 let texts2 = texts.clone();
1017 let resolve = |b: &str| if b == "blk-1" { None } else { agent_for(b) };
1019 let out = run_broadcast(ids(3), "great, merge on approval", "bcast-42", 4, resolve, move |b, t, _| {
1020 texts2.lock().unwrap().push((b, t));
1021 async { Ok(()) }
1022 })
1023 .await;
1024 assert_eq!(out.len(), 3);
1025 let texts = texts.lock().unwrap();
1026 assert_eq!(texts.len(), 2, "the unresolved target is not delivered to");
1027 for (block, text) in texts.iter() {
1028 let agent = agent_for(block).unwrap();
1029 let header = text.lines().next().unwrap();
1030 assert!(header.starts_with(&format!("[BROADCAST:FROM=user VIA=swarm TO={agent} RECIPIENTS=2 MSGID=bcast-42 TS=")), "{header}");
1031 assert_eq!(text.lines().nth(1), Some("great, merge on approval"));
1032 }
1033 }
1034
1035 #[tokio::test]
1036 async fn outcomes_are_per_target_in_order_and_an_unresolved_target_says_why() {
1037 let out = run_broadcast(
1038 vec!["blk-0".into(), "nope".into(), "blk-2".into()],
1039 "x",
1040 "m",
1041 4,
1042 agent_for,
1043 |b, _t, _| async move { if b == "blk-2" { Err("no controller".to_string()) } else { Ok(()) } },
1044 )
1045 .await;
1046 let blocks: Vec<&str> = out.iter().map(|o| o.block_id.as_str()).collect();
1047 assert_eq!(blocks, vec!["blk-0", "nope", "blk-2"]);
1048 assert!(out[0].result.is_ok());
1049 assert!(out[1].result.as_ref().unwrap_err().contains("no registered agent"));
1050 assert!(out[1].agent_id.is_none());
1051 assert_eq!(out[2].result.as_ref().unwrap_err(), "no controller");
1052 assert_eq!(out[2].agent_id.as_deref(), Some("agent-2"));
1053 }
1054
1055 #[tokio::test]
1056 async fn a_large_broadcast_does_not_pause_between_chunks() {
1057 let started = std::time::Instant::now();
1058 let out = run_broadcast(ids(25), "x", "m", BROADCAST_CONCURRENCY, agent_for, |_b, _t, _| async { Ok(()) }).await;
1059 assert_eq!(out.iter().filter(|o| o.result.is_ok()).count(), 25);
1060 assert!(started.elapsed() < std::time::Duration::from_millis(900), "the old limiter pause is gone: {:?}", started.elapsed());
1061 }
1062
1063 #[tokio::test]
1064 async fn at_most_the_concurrency_limit_turns_start_at_once() {
1065 let in_flight = Arc::new(AtomicUsize::new(0));
1066 let peak = Arc::new(AtomicUsize::new(0));
1067 let (inf, pk) = (in_flight.clone(), peak.clone());
1068 let out = run_broadcast(ids(20), "x", "m", 3, agent_for, move |_b, _t, _| {
1069 let (inf, pk) = (inf.clone(), pk.clone());
1070 async move {
1071 let now = inf.fetch_add(1, Ordering::SeqCst) + 1;
1072 pk.fetch_max(now, Ordering::SeqCst);
1073 for _ in 0..5 {
1074 tokio::task::yield_now().await;
1075 }
1076 inf.fetch_sub(1, Ordering::SeqCst);
1077 Ok(())
1078 }
1079 })
1080 .await;
1081 assert_eq!(out.len(), 20);
1082 let peak = peak.load(Ordering::SeqCst);
1083 assert!(peak <= 3, "peak {peak} exceeded the limit");
1084 assert!(peak > 1, "targets should overlap, peak {peak}");
1085 }
1086
1087 #[tokio::test]
1088 async fn no_targets_is_a_no_op() {
1089 let out = run_broadcast(vec![], "x", "m", 4, agent_for, |_b, _t, _| async { Ok(()) }).await;
1090 assert!(out.is_empty());
1091 }
1092}
1093
1094#[cfg(test)]
1095mod broadcast_action_tests {
1096 use super::*;
1097 use crate::backend::reactive::SenderDelivery;
1098 use crate::bootstrap::AgentRoute;
1099
1100 #[test]
1101 fn structured_channels_take_the_broadcast_without_a_new_turn() {
1102 assert_eq!(broadcast_action(Ok(AgentRoute::Sent(SenderDelivery::Delivered))), BroadcastAction::Done);
1105 assert_eq!(broadcast_action(Ok(AgentRoute::Sent(SenderDelivery::Deferred))), BroadcastAction::Done);
1106 }
1107
1108 #[test]
1109 fn a_subprocess_or_unspawned_persistent_agent_gets_a_turn() {
1110 assert_eq!(broadcast_action(Ok(AgentRoute::StartTurn { subprocess: true })), BroadcastAction::StartTurn);
1111 assert_eq!(broadcast_action(Ok(AgentRoute::StartTurn { subprocess: false })), BroadcastAction::StartTurn);
1112 }
1113
1114 #[test]
1115 fn a_terminal_pane_is_refused_not_typed_into() {
1116 let BroadcastAction::Fail(e) = broadcast_action(Ok(AgentRoute::Sent(SenderDelivery::Pty))) else {
1117 panic!("a PTY target must fail");
1118 };
1119 assert!(e.contains("not typed into a shell"), "{e}");
1120 }
1121
1122 #[test]
1123 fn a_structured_controller_failure_is_reported_per_target() {
1124 assert_eq!(
1125 broadcast_action(Err("persistent process not running".into())),
1126 BroadcastAction::Fail("persistent process not running".into())
1127 );
1128 }
1129}
1130
1131#[cfg(test)]
1132mod forward_broadcast_tests {
1133 use super::*;
1134 use axum::routing::post;
1135 use std::sync::{Arc, Mutex};
1136
1137 type Seen = Arc<Mutex<Vec<(Option<String>, serde_json::Value)>>>;
1138
1139 async fn peer(status: axum::http::StatusCode, reply: serde_json::Value) -> (String, Seen) {
1142 let seen: Seen = Default::default();
1143 let s = seen.clone();
1144 let app = axum::Router::new().route(
1145 "/agentmux/agent/broadcast",
1146 post(move |headers: axum::http::HeaderMap, axum::Json(body): axum::Json<serde_json::Value>| {
1147 let (s, reply) = (s.clone(), reply.clone());
1148 async move {
1149 let key = headers.get("X-AuthKey").and_then(|v| v.to_str().ok()).map(str::to_string);
1150 s.lock().unwrap().push((key, body));
1151 (status, axum::Json(reply))
1152 }
1153 }),
1154 );
1155 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
1156 let url = format!("http://127.0.0.1:{}", listener.local_addr().unwrap().port());
1157 tokio::spawn(async move {
1158 let _ = axum::serve(listener, app).await;
1159 });
1160 (url, seen)
1161 }
1162
1163 fn forward<'a>() -> BroadcastForward<'a> {
1164 BroadcastForward { block_id: "blk-loap", body: "merge on approval", msg_id: "m-1", recipients: 3 }
1165 }
1166
1167 #[tokio::test]
1168 async fn the_caller_sends_the_body_and_counts_and_never_a_header() {
1169 let (url, seen) = peer(axum::http::StatusCode::OK, serde_json::json!({ "success": true })).await;
1170 let out = forward_broadcast_to_channel(&reqwest::Client::new(), &url, "chan-key", &forward()).await;
1171 assert_eq!(out, Ok(()));
1172 let seen = seen.lock().unwrap();
1173 assert_eq!(seen[0].0.as_deref(), Some("chan-key"), "that channel's own key");
1174 let mut keys: Vec<&str> = seen[0].1.as_object().unwrap().keys().map(String::as_str).collect();
1175 keys.sort_unstable();
1176 assert_eq!(keys, ["block_id", "body", "msg_id", "recipients"], "no header field for a caller to forge");
1177 assert_eq!(seen[0].1["body"], "merge on approval");
1178 assert_eq!(seen[0].1["recipients"], 3);
1179 }
1180
1181 #[tokio::test]
1182 async fn a_channel_that_cannot_take_it_says_why_for_that_target_only() {
1183 let (url, _) = peer(axum::http::StatusCode::OK, serde_json::json!({ "success": false, "error": "no registered agent" })).await;
1184 assert_eq!(
1185 forward_broadcast_to_channel(&reqwest::Client::new(), &url, "", &forward()).await,
1186 Err("no registered agent".to_string())
1187 );
1188 }
1189
1190 #[tokio::test]
1191 async fn a_channel_from_before_the_route_is_named_as_an_older_build() {
1192 let (url, _) = peer(axum::http::StatusCode::NOT_FOUND, serde_json::json!({})).await;
1193 let Err(e) = forward_broadcast_to_channel(&reqwest::Client::new(), &url, "", &forward()).await else {
1194 panic!("a 404 must fail the target");
1195 };
1196 assert!(e.contains("older build"), "{e}");
1197 }
1198
1199 #[tokio::test]
1200 async fn a_missing_key_sends_no_key_header() {
1201 let (url, seen) = peer(axum::http::StatusCode::OK, serde_json::json!({ "success": true })).await;
1202 forward_broadcast_to_channel(&reqwest::Client::new(), &url, "", &forward()).await.unwrap();
1203 assert_eq!(seen.lock().unwrap()[0].0, None);
1204 }
1205
1206 #[tokio::test]
1207 async fn an_unreachable_channel_fails_the_target() {
1208 let Err(e) = forward_broadcast_to_channel(&reqwest::Client::new(), "http://127.0.0.1:9", "", &forward()).await else {
1209 panic!("nothing listens there");
1210 };
1211 assert!(e.starts_with("cross-channel forward failed"), "{e}");
1212 }
1213
1214 #[tokio::test]
1215 async fn the_receiver_refuses_a_block_it_has_no_agent_for() {
1216 let state = crate::server::tests::test_state();
1217 assert_eq!(
1218 deliver_forwarded_broadcast(&state, "no-such-block", "hi", "m-1", 1).await,
1219 Err(NO_REGISTERED_AGENT.to_string())
1220 );
1221 }
1222}