Skip to main content

agent_engine/daemon/
lifecycle.rs

1//! Daemon shutdown: end every session concurrently inside ONE wall budget,
2//! then shut the shared extension host down.
3
4use std::sync::Arc;
5use std::time::Duration;
6
7use super::DaemonState;
8use crate::session::{EndReason, SessionCommand};
9
10/// `session::budgets::SESSION_END_TIMEOUT_SECS`: the worst case of one
11/// session's `finish` (cancelled-turn drain, save, hooks, observability
12/// flush); sessions end concurrently, so it is shared, not ×N.
13pub const SESSION_END_BUDGET: Duration =
14    Duration::from_secs(crate::session::budgets::SESSION_END_TIMEOUT_SECS);
15/// `--force`: give the actors this long, then drop them.
16pub const FORCE_BUDGET: Duration = Duration::from_millis(500);
17
18pub async fn shutdown_all(state: &Arc<DaemonState>, force: bool) {
19    let sessions = state.live_sessions();
20    let n = sessions.len();
21    tracing::info!(sessions = n, force, "daemon: ending sessions");
22    let budget = if force { FORCE_BUDGET } else { SESSION_END_BUDGET };
23    let ends = sessions.iter().map(|h| async move {
24        let _ = h.send(SessionCommand::End { reason: EndReason::HostShutdown }).await;
25        h.closed().await;
26    });
27    if tokio::time::timeout(budget, futures::future::join_all(ends)).await.is_err() {
28        tracing::warn!(budget = ?budget, "daemon: session shutdown budget exceeded; continuing");
29    }
30    for h in &sessions {
31        state.host.remove_session(&h.id);
32    }
33
34    let ext = Arc::clone(state.host.ext_manager());
35    let ext_shutdown = async move { ext.write().await.shutdown_all().await };
36    if tokio::time::timeout(Duration::from_secs(5), ext_shutdown).await.is_err() {
37        tracing::warn!("daemon: extension shutdown budget exceeded");
38    }
39    state.host.flush_logs();
40}