1use std::collections::HashMap;
15use std::sync::{Arc, LazyLock, Mutex};
16
17use serde::Serialize;
18
19use crate::server::AppState;
20
21pub const OVERRIDE_WINDOW: std::time::Duration = std::time::Duration::from_secs(15);
22pub const EVENT_SHUTDOWN_PENDING: &str = "agent:shutdown-pending";
23pub const EVENT_SHUTDOWN_PENDING_CLEARED: &str = "agent:shutdown-pending-cleared";
24const KEEP_FINISHED_MS: i64 = 10 * 60 * 1000;
26
27#[derive(Clone, Debug)]
29pub enum Action {
30 SelfQuit { detail: String },
33 ClosePane { block_ids: Vec<String> },
35 Stop { signal: Option<String> },
37}
38
39impl Action {
40 fn rank(&self) -> u8 {
43 match self {
44 Action::Stop { .. } => 0,
45 Action::SelfQuit { .. } => 1,
46 Action::ClosePane { .. } => 2,
47 }
48 }
49
50 fn covers(&self, block_id: &str) -> Vec<String> {
52 let mut blocks = vec![block_id.to_string()];
53 if let Action::ClosePane { block_ids } = self {
54 blocks.extend(block_ids.iter().filter(|b| b.as_str() != block_id).cloned());
55 }
56 blocks
57 }
58
59 fn merge(&mut self, other: Action) {
63 match (self, other) {
64 (Action::Stop { signal }, Action::Stop { signal: theirs }) => {
65 let forceful = |s: &Option<String>| matches!(s.as_deref(), Some("SIGKILL") | Some("SIGTERM"));
66 if theirs.as_deref() == Some("SIGKILL") || (signal.is_none() && forceful(&theirs)) {
67 *signal = theirs;
68 }
69 }
70 (Action::ClosePane { block_ids }, Action::ClosePane { block_ids: theirs }) => {
71 for b in theirs {
72 if !block_ids.contains(&b) {
73 block_ids.push(b);
74 }
75 }
76 }
77 _ => {}
78 }
79 }
80}
81
82#[derive(Clone, Debug, Serialize, PartialEq, Eq, ts_rs::TS)]
84#[ts(export, export_to = "../../../frontend/types/rpc/")]
85pub struct PendingView {
86 pub request_id: String,
87 pub block_id: String,
88 pub by: String,
90 pub target: String,
92 pub via: String,
94 pub reason: String,
95 #[ts(type = "number")]
96 pub deadline_ms: i64,
97 #[ts(type = r#""pending" | "kept_by_user" | "proceeding" | "shut_down" | "failed" | "superseded""#)]
100 pub status: &'static str,
101 #[serde(skip_serializing_if = "Option::is_none")]
102 #[ts(optional)]
103 pub error: Option<String>,
104 #[serde(skip)]
107 #[ts(skip)]
108 pub joined: bool,
109}
110
111#[derive(Clone, Debug, Serialize, ts_rs::TS)]
113#[ts(export, export_to = "../../../frontend/types/rpc/")]
114pub struct PendingClearedEvent {
115 pub block_id: String,
116 pub request_id: String,
117 pub outcome: String,
118}
119
120struct Entry {
121 view: PendingView,
122 kept: Arc<tokio::sync::Notify>,
123 finished_at_ms: Option<i64>,
124 action: Option<Action>,
126 requesters: Vec<(String, String)>,
129 covers: Vec<String>,
136}
137
138static ENTRIES: LazyLock<Mutex<HashMap<String, Entry>>> = LazyLock::new(Default::default);
139
140#[derive(Clone, Debug, PartialEq, Eq)]
144pub struct RemoteRequest {
145 pub local_url: String,
146 pub auth_key: String,
147 remembered_at_ms: i64,
148}
149
150static REMOTE: LazyLock<Mutex<HashMap<String, RemoteRequest>>> = LazyLock::new(Default::default);
151
152pub fn remember_remote(request_id: &str, local_url: &str, auth_key: &str) {
154 let now = now_ms();
155 let mut map = REMOTE.lock().unwrap_or_else(|e| e.into_inner());
156 map.retain(|_, r| now - r.remembered_at_ms < KEEP_FINISHED_MS + OVERRIDE_WINDOW.as_millis() as i64);
158 map.insert(
159 request_id.to_string(),
160 RemoteRequest { local_url: local_url.to_string(), auth_key: auth_key.to_string(), remembered_at_ms: now },
161 );
162}
163
164pub fn remote_for(request_id: &str) -> Option<RemoteRequest> {
166 REMOTE.lock().unwrap_or_else(|e| e.into_inner()).get(request_id).cloned()
167}
168
169use agentmux_common::time::now_ms;
170
171fn entries() -> std::sync::MutexGuard<'static, HashMap<String, Entry>> {
172 ENTRIES.lock().unwrap_or_else(|e| e.into_inner())
173}
174
175fn prune(map: &mut HashMap<String, Entry>, now: i64) {
176 map.retain(|_, e| e.finished_at_ms.is_none_or(|t| now - t < KEEP_FINISHED_MS));
177}
178
179fn is_active(status: &str) -> bool {
180 matches!(status, "pending" | "proceeding")
181}
182
183pub fn active_for(block_id: &str) -> Option<PendingView> {
186 entries()
187 .values()
188 .find(|e| is_active(e.view.status) && e.covers.iter().any(|b| b == block_id))
189 .map(|e| e.view.clone())
190}
191
192fn covers_of(request_id: &str) -> Vec<String> {
194 entries().get(request_id).map(|e| e.covers.clone()).unwrap_or_default()
195}
196
197pub fn status(request_id: &str) -> Option<PendingView> {
198 entries().get(request_id).map(|e| e.view.clone())
199}
200
201fn begin(request_id: &str) -> Option<Action> {
206 let mut map = entries();
207 match map.get_mut(request_id) {
208 Some(e) if e.view.status == "pending" => {
209 e.view.status = "proceeding";
210 e.action.take()
211 }
212 _ => None,
213 }
214}
215
216fn set_status(request_id: &str, status: &'static str, error: Option<String>) -> Option<PendingView> {
217 let mut map = entries();
218 let e = map.get_mut(request_id)?;
219 e.view.status = status;
220 e.view.error = error;
221 if status != "pending" && status != "proceeding" {
222 e.finished_at_ms = Some(now_ms());
223 }
224 Some(e.view.clone())
225}
226
227fn publish_pending(state: &AppState, block_id: &str, view: &PendingView) {
229 state.broker.publish_event::<crate::backend::mps_events::AgentShutdownPending>(view, vec![format!("block:{block_id}")], 1);
230}
231
232fn notify_router(state: &AppState) -> Option<Arc<crate::backend::notify::router::Router>> {
233 crate::backend::notify::router::get(&state.broker)
234}
235
236pub fn request(state: &AppState, block_id: &str, by: &str, via: &str, reason: &str, action: Action) -> PendingView {
245 request_within(state, block_id, by, via, reason, action, OVERRIDE_WINDOW)
246}
247
248fn request_within(
249 state: &AppState,
250 block_id: &str,
251 by: &str,
252 via: &str,
253 reason: &str,
254 action: Action,
255 window: std::time::Duration,
256) -> PendingView {
257 let new_covers = action.covers(block_id);
258 let view = {
259 let mut map = entries();
260 prune(&mut map, now_ms());
261 if let Some(existing) = map
262 .values_mut()
263 .find(|e| is_active(e.view.status) && e.covers.iter().any(|b| new_covers.contains(b)))
264 {
265 let mut grew = false;
266 for b in &new_covers {
267 if !existing.covers.contains(b) {
268 existing.covers.push(b.clone());
269 grew = true;
270 }
271 }
272 if !existing.requesters.iter().any(|(b, v)| b == by && v == via) {
273 existing.requesters.push((by.to_string(), via.to_string()));
274 }
275 let pending = existing.view.status == "pending";
276 let stronger = pending && existing.action.as_ref().is_some_and(|cur| action.rank() > cur.rank());
277 if pending && !stronger {
278 if let Some(cur) = existing.action.as_mut().filter(|cur| cur.rank() == action.rank()) {
279 cur.merge(action.clone());
280 }
281 }
282 if stronger {
283 existing.action = Some(action);
284 existing.view.by = by.to_string();
285 existing.view.via = via.to_string();
286 if !reason.trim().is_empty() {
287 existing.view.reason = reason.trim().to_string();
288 }
289 }
290 let joined = PendingView { joined: true, ..existing.view.clone() };
291 let covers = existing.covers.clone();
292 drop(map);
293 if pending && (stronger || grew) {
295 let view = PendingView { joined: false, ..joined.clone() };
298 for b in &covers {
299 publish_pending(state, b, &view);
300 }
301 }
302 return joined;
303 }
304 let view = PendingView {
305 request_id: uuid::Uuid::new_v4().to_string(),
306 block_id: block_id.to_string(),
307 by: by.to_string(),
308 target: state
309 .reactive_handler
310 .get_agent_by_block(block_id)
311 .map(|a| a.agent_id)
312 .unwrap_or_else(|| block_id.to_string()),
313 via: via.to_string(),
314 reason: reason.trim().to_string(),
315 deadline_ms: now_ms() + window.as_millis() as i64,
316 status: "pending",
317 error: None,
318 joined: false,
319 };
320 map.insert(
321 view.request_id.clone(),
322 Entry {
323 view: view.clone(),
324 kept: Arc::new(tokio::sync::Notify::new()),
325 finished_at_ms: None,
326 action: Some(action),
327 requesters: vec![(by.to_string(), via.to_string())],
328 covers: new_covers.clone(),
329 },
330 );
331 view
332 };
333 for b in &new_covers {
336 publish_pending(state, b, &view);
337 }
338 if let Some(router) = notify_router(state) {
339 let why = if view.reason.is_empty() { String::new() } else { format!(": {}", view.reason) };
340 let body = format!("{} asked to shut it down{why}. Open it to keep it running.", view.by);
341 let block = block_id.to_string();
342 tokio::task::spawn_blocking(move || {
344 router.emit_fixed(crate::backend::notify::policy::NotifyKind::ShutdownPending, &block, Some(body))
345 });
346 }
347 tracing::info!(block_id, by, via, request_id = %view.request_id, "shutdown pending: user has 15 s to keep it");
348
349 let kept = entries().get(&view.request_id).map(|e| e.kept.clone());
350 let (state, v) = (state.clone(), view.clone());
351 tokio::spawn(async move {
352 if let Some(kept) = kept {
353 tokio::select! {
354 _ = tokio::time::sleep(window) => {}
355 _ = kept.notified() => {}
356 }
357 }
358 let Some(action) = begin(&v.request_id) else {
359 return; };
361 cleared(&state, &v, "proceeding");
362 let outcome = run_action(&state, &v.block_id, action).await;
363 let (status, error) = match outcome {
364 Ok(()) => ("shut_down", None),
365 Err(e) if e.contains("block not found") || e.starts_with("NOT_RUNNING") => ("superseded", None),
367 Err(e) => ("failed", Some(e)),
368 };
369 audit(&state, &v.request_id, status, error.as_deref());
370 set_status(&v.request_id, status, error);
371 });
372 view
373}
374
375fn audit(state: &AppState, request_id: &str, outcome: &str, error: Option<&str>) {
377 let Some((v, requesters)) = entries().get(request_id).map(|e| (e.view.clone(), e.requesters.clone())) else {
378 return;
379 };
380 for (by, via) in requesters {
381 let ran = if via == v.via { String::new() } else { format!(" (ran as {} by {})", v.via, v.by) };
382 state.reactive_handler.log_fleet_action_audit(
383 Some(&by),
384 &v.target,
385 &v.block_id,
386 "agent.shutdown_override",
387 error.is_none(),
388 error,
389 &v.request_id,
390 Some(&format!("{outcome} | via {via}{ran} | {}", v.reason)),
391 );
392 }
393}
394
395pub fn audit_note(v: &PendingView, by: &str, via: &str) -> String {
398 if !v.joined {
399 "pending user override".to_string()
400 } else if v.by == by && v.via == via {
401 format!("joined its own pending request {}", v.request_id)
402 } else {
403 format!("joined {} by {}'s pending request", v.via, v.by)
404 }
405}
406
407fn cleared(state: &AppState, v: &PendingView, outcome: &str) {
408 let data = PendingClearedEvent { block_id: v.block_id.clone(), request_id: v.request_id.clone(), outcome: outcome.to_string() };
409 for b in covers_of(&v.request_id) {
410 state.broker.publish_event::<crate::backend::mps_events::AgentShutdownPendingCleared>(&data, vec![format!("block:{b}")], 1);
411 }
412 if let Some(router) = notify_router(state) {
413 router.resolve_nonblocking(&v.block_id, crate::backend::notify::policy::Family::Shutdown);
414 }
415}
416
417async fn run_action(state: &AppState, block_id: &str, action: Action) -> Result<(), String> {
418 match action {
419 Action::SelfQuit { detail } => {
420 let deadline = tokio::time::Instant::now() + super::self_quit::SELF_QUIT_TURN_GRACE;
422 while tokio::time::Instant::now() < deadline
423 && crate::backend::blockcontroller::get_controller(block_id)
424 .is_some_and(|c| c.get_runtime_status().turn_active)
425 {
426 tokio::time::sleep(std::time::Duration::from_millis(250)).await;
427 }
428 super::self_quit::run_with(state, block_id, super::self_quit::QuitOrigin::AgentTool, Some(&detail))
429 .await
430 .map(|_| ())
431 }
432 Action::ClosePane { block_ids } => super::close_pane::run(state, block_ids).await.map(|_| ()),
433 Action::Stop { signal } => {
434 crate::server::app_api::agent_io::stop_one_agent_block(state, block_id, signal.as_deref(), None)
435 .await
436 .map(|_| ())
437 }
438 }
439}
440
441#[cfg(test)]
443pub fn expire_now(request_id: &str) {
444 if let Some(e) = entries().get(request_id) {
445 e.kept.notify_one();
446 }
447}
448
449#[cfg(test)]
451pub async fn settled(request_id: &str) -> PendingView {
452 for _ in 0..400 {
453 let v = status(request_id).expect("known request");
454 if !matches!(v.status, "pending" | "proceeding") {
455 return v;
456 }
457 tokio::time::sleep(std::time::Duration::from_millis(10)).await;
458 }
459 panic!("request {request_id} never settled");
460}
461
462pub fn keep(state: &AppState, block_id: &str, request_id: &str) -> &'static str {
465 let view = {
466 let mut map = entries();
467 match map.get_mut(request_id) {
468 Some(e) if e.covers.iter().any(|b| b == block_id) && e.view.status == "pending" => {
469 e.view.status = "kept_by_user";
472 e.finished_at_ms = Some(now_ms());
473 e.kept.notify_one();
474 e.view.clone()
475 }
476 _ => return "too_late",
477 }
478 };
479 cleared(state, &view, "kept_by_user");
480 audit(state, request_id, "kept_by_user", None);
481 tracing::info!(block_id, request_id, by = %view.by, "shutdown pending: kept by the user");
482 "kept_by_user"
483}
484
485#[cfg(test)]
486mod tests {
487 use super::*;
488 use crate::server::tests::test_state;
489
490 #[tokio::test]
491 async fn a_request_is_pending_then_a_second_one_keeps_the_same_clock() {
492 let state = test_state();
493 let a = request(&state, "blk-p1", "Korp", "QuitSelf", "done", Action::SelfQuit { detail: String::new() });
494 assert_eq!(a.status, "pending");
495 let b = request(&state, "blk-p1", "cron", "QuitSelf", "again", Action::SelfQuit { detail: String::new() });
496 assert_eq!((b.request_id.as_str(), b.deadline_ms), (a.request_id.as_str(), a.deadline_ms), "never restarts the clock");
497 assert_eq!(active_for("blk-p1").map(|v| v.request_id), Some(a.request_id.clone()));
498 assert_eq!(keep(&state, "blk-p1", &a.request_id), "kept_by_user");
499 }
500
501 #[tokio::test]
502 async fn a_request_under_way_is_returned_too_so_no_second_window_opens() {
503 let state = test_state();
504 let a = request(&state, "blk-p3", "Korp", "QuitSelf", "done", Action::SelfQuit { detail: String::new() });
505 assert!(begin(&a.request_id).is_some(), "the window ran out");
506 let b = request(&state, "blk-p3", "Korp", "QuitSelf", "again", Action::SelfQuit { detail: String::new() });
507 assert_eq!((b.request_id.as_str(), b.status), (a.request_id.as_str(), "proceeding"));
508 assert_eq!(keep(&state, "blk-p3", &a.request_id), "too_late");
509 assert!(begin(&a.request_id).is_none(), "the action runs once");
510 }
511
512 #[tokio::test]
513 async fn a_stronger_request_joins_and_takes_over_the_pending_action() {
514 let state = test_state();
515 let stop = request(&state, "blk-p4", "an agent", "FleetBulkStop", "", Action::Stop { signal: None });
516 let close = request(
517 &state,
518 "blk-p4",
519 "Korp",
520 "ClosePane",
521 "stuck",
522 Action::ClosePane { block_ids: vec!["blk-p4".into()] },
523 );
524 assert_eq!(close.request_id, stop.request_id, "one window per block");
525 assert_eq!((close.by.as_str(), close.via.as_str(), close.deadline_ms), ("Korp", "ClosePane", stop.deadline_ms));
526 assert!(matches!(begin(&stop.request_id), Some(Action::ClosePane { .. })), "the pane closes: it does all a stop does");
527
528 let quit = request(&state, "blk-p5", "Camper", "QuitSelf", "done", Action::SelfQuit { detail: String::new() });
530 let again = request(&state, "blk-p5", "an agent", "FleetBulkStop", "", Action::Stop { signal: None });
531 assert_eq!((again.request_id.as_str(), again.via.as_str()), (quit.request_id.as_str(), "QuitSelf"));
532 assert!(matches!(begin(&quit.request_id), Some(Action::SelfQuit { .. })));
533 let requesters = entries().get(&quit.request_id).unwrap().requesters.clone();
534 assert_eq!(requesters.len(), 2, "both are audited when it settles");
535 }
536
537 #[tokio::test]
538 async fn a_joiner_of_the_same_kind_is_folded_in_not_dropped() {
539 let state = test_state();
540 let a = request(&state, "blk-p6", "Korp", "FleetBulkStop", "", Action::Stop { signal: None });
541 assert!(!a.joined);
542 let b = request(&state, "blk-p6", "Korp", "FleetBulkStop", "", Action::Stop { signal: Some("SIGKILL".into()) });
543 assert!(b.joined && b.request_id == a.request_id);
544 assert_eq!(audit_note(&b, "Korp", "FleetBulkStop"), format!("joined its own pending request {}", a.request_id));
545 assert!(matches!(begin(&a.request_id), Some(Action::Stop { signal: Some(s) }) if s == "SIGKILL"), "the forceful stop wins");
546
547 let c = request(&state, "blk-p7", "Korp", "ClosePane", "", Action::ClosePane { block_ids: vec!["x".into()] });
548 request(&state, "blk-p7", "Posa", "ClosePane", "", Action::ClosePane { block_ids: vec!["x".into(), "y".into()] });
549 assert!(matches!(begin(&c.request_id), Some(Action::ClosePane { block_ids }) if block_ids == vec!["x".to_string(), "y".to_string()]));
550 }
551
552 #[tokio::test]
553 async fn closes_naming_two_tabs_of_one_pane_share_one_window() {
554 let state = test_state();
555 let pane = || Action::ClosePane { block_ids: vec!["tab-a".into(), "tab-b".into()] };
556 let a = request(&state, "tab-a", "Korp", "ClosePane", "", pane());
557 let b = request(&state, "tab-b", "Posa", "ClosePane", "", pane());
558 assert!(b.joined && b.request_id == a.request_id, "one pane, one window");
559 assert_eq!(active_for("tab-b").map(|v| v.request_id), Some(a.request_id.clone()));
560 assert_eq!(keep(&state, "tab-b", &a.request_id), "kept_by_user", "kept from the other tab's banner");
561 assert!(begin(&a.request_id).is_none(), "nothing closes");
562
563 let stop = request(&state, "tab-d", "an agent", "FleetBulkStop", "", Action::Stop { signal: None });
565 let close = request(&state, "tab-c", "Korp", "ClosePane", "", Action::ClosePane { block_ids: vec!["tab-c".into(), "tab-d".into()] });
566 assert_eq!(close.request_id, stop.request_id);
567 assert_eq!(active_for("tab-c").map(|v| v.via), Some("ClosePane".to_string()));
568 assert!(matches!(begin(&stop.request_id), Some(Action::ClosePane { .. })));
569 }
570
571 #[tokio::test]
572 async fn keep_running_cancels_it_and_a_second_keep_is_too_late() {
573 let state = test_state();
574 let v = request(&state, "blk-p2", "Korp", "QuitSelf", "done", Action::SelfQuit { detail: String::new() });
575 assert_eq!(keep(&state, "blk-p2", &v.request_id), "kept_by_user");
576 assert_eq!(status(&v.request_id).unwrap().status, "kept_by_user");
577 assert_eq!(active_for("blk-p2"), None);
578 assert_eq!(keep(&state, "blk-p2", &v.request_id), "too_late");
579 assert_eq!(keep(&state, "some-other-block", &v.request_id), "too_late", "only the block it was for");
580 }
581
582 #[tokio::test]
583 async fn keep_and_the_window_running_out_never_both_win() {
584 let state = test_state();
585 for i in 0..50 {
586 let block = format!("blk-race-{i}");
587 let v = request_within(&state, &block, "Korp", "QuitSelf", "r", Action::SelfQuit { detail: String::new() }, std::time::Duration::ZERO);
588 let kept = keep(&state, &block, &v.request_id);
589 for _ in 0..100 {
590 if !matches!(status(&v.request_id).unwrap().status, "pending" | "proceeding") {
591 break;
592 }
593 tokio::time::sleep(std::time::Duration::from_millis(5)).await;
594 }
595 let end = status(&v.request_id).unwrap().status;
596 if kept == "kept_by_user" {
597 assert_eq!(end, "kept_by_user", "a kept request must never go ahead");
598 } else {
599 assert_ne!(end, "kept_by_user");
600 }
601 }
602 }
603
604 #[tokio::test]
605 async fn with_no_override_it_goes_ahead_after_the_window() {
606 let state = test_state();
607 let window = std::time::Duration::from_millis(50);
608 let v = request_within(&state, "blk-gone", "Korp", "QuitSelf", "done", Action::SelfQuit { detail: String::new() }, window);
609 for _ in 0..100 {
610 if !matches!(status(&v.request_id).unwrap().status, "pending" | "proceeding") {
611 break;
612 }
613 tokio::time::sleep(std::time::Duration::from_millis(20)).await;
614 }
615 assert_eq!(status(&v.request_id).unwrap().status, "superseded");
618 }
619}