systemprompt_cli/commands/infrastructure/services/restart/
batch.rs1use crate::cli_settings::CliConfig;
2use crate::shared::CommandOutput;
3use anyhow::Result;
4use std::collections::HashMap;
5use std::sync::Arc;
6use systemprompt_agent::services::agent_orchestration::{AgentOrchestrator, AgentStatus};
7use systemprompt_agent::services::registry::AgentRegistry;
8use systemprompt_logging::CliService;
9use systemprompt_mcp::HealthStatus;
10use systemprompt_runtime::AppContext;
11use systemprompt_scheduler::{
12 RestartPlan, RestartScope, RestartTarget, ServiceSnapshot, ServiceType,
13};
14
15use super::super::lifecycle;
16use super::super::types::RestartOutput;
17
18pub async fn execute_all_agents(
19 ctx: &Arc<AppContext>,
20 config: &CliConfig,
21) -> Result<CommandOutput> {
22 let quiet = config.is_json_output();
23
24 if !quiet {
25 CliService::section("Restarting All Agents");
26 }
27
28 let orchestrator = lifecycle::agent_orchestrator(ctx).await?;
29 let snapshots = agent_snapshots(&orchestrator).await?;
30 let plan = RestartPlan::compute(RestartScope::AllAgents, &snapshots);
31
32 let mut restarted = 0usize;
33 let mut failed = 0usize;
34
35 for target in &plan.targets {
36 if !quiet {
37 CliService::info(&format!("Restarting agent: {}", target.name));
38 }
39 restart_agent_target(&orchestrator, target, &mut restarted, &mut failed, quiet).await;
40 }
41
42 let message = super::format_batch_message("agents", restarted, failed, quiet);
43
44 let output = RestartOutput {
45 service_type: "agents".to_owned(),
46 service_name: None,
47 restarted_count: restarted,
48 failed_count: failed,
49 message,
50 };
51
52 Ok(CommandOutput::card_value("Restart All Agents", &output))
53}
54
55pub async fn execute_all_mcp(ctx: &Arc<AppContext>, config: &CliConfig) -> Result<CommandOutput> {
56 let quiet = config.is_json_output();
57
58 if !quiet {
59 CliService::section("Restarting All MCP Servers");
60 }
61
62 let mcp_manager = lifecycle::mcp_orchestrator(ctx)?;
63 let snapshots = mcp_snapshots(ctx, false).await?;
64 let plan = RestartPlan::compute(RestartScope::AllMcp, &snapshots);
65
66 let mut restarted = 0usize;
67 let mut failed = 0usize;
68
69 for target in &plan.targets {
70 if !quiet {
71 CliService::info(&format!("Restarting MCP server: {}", target.name));
72 }
73 match mcp_manager.restart_services(Some(target.id.clone())).await {
74 Ok(()) => {
75 restarted += 1;
76 if !quiet {
77 CliService::success(&format!(" {} restarted", target.name));
78 }
79 },
80 Err(e) => {
81 failed += 1;
82 if !quiet {
83 CliService::error(&format!(" Failed to restart {}: {}", target.name, e));
84 }
85 },
86 }
87 }
88
89 let message = super::format_batch_message("MCP servers", restarted, failed, quiet);
90
91 let output = RestartOutput {
92 service_type: "mcp".to_owned(),
93 service_name: None,
94 restarted_count: restarted,
95 failed_count: failed,
96 message,
97 };
98
99 Ok(CommandOutput::card_value(
100 "Restart All MCP Servers",
101 &output,
102 ))
103}
104
105pub async fn execute_failed(ctx: &Arc<AppContext>, config: &CliConfig) -> Result<CommandOutput> {
106 let quiet = config.is_json_output();
107
108 if !quiet {
109 CliService::section("Restarting Failed Services");
110 }
111
112 let handles = lifecycle::OrchestratorHandles::build(ctx).await?;
113
114 let mut snapshots = agent_snapshots(&handles.agents).await?;
115 snapshots.extend(mcp_snapshots(ctx, true).await?);
116 let plan = RestartPlan::compute(RestartScope::Failed, &snapshots);
117
118 let mut restarted_count = 0usize;
119 let mut failed_count = 0usize;
120
121 for target in &plan.targets {
122 match target.service_type {
123 ServiceType::Agent => {
124 if !quiet {
125 CliService::info(&format!("Restarting failed agent: {}", target.name));
126 }
127 restart_agent_target(
128 &handles.agents,
129 target,
130 &mut restarted_count,
131 &mut failed_count,
132 quiet,
133 )
134 .await;
135 },
136 ServiceType::Mcp => {
137 if !quiet {
138 CliService::info(&format!("Restarting MCP server: {}", target.name));
139 }
140 restart_mcp_target(
141 &handles.mcp,
142 target,
143 &mut restarted_count,
144 &mut failed_count,
145 quiet,
146 )
147 .await;
148 },
149 ServiceType::Api => {},
150 }
151 }
152
153 let message = if restarted_count > 0 {
154 let msg = format!("Restarted {} failed services", restarted_count);
155 if !quiet {
156 CliService::success(&msg);
157 }
158 if failed_count > 0 && !quiet {
159 CliService::warning(&format!("Failed to restart {} services", failed_count));
160 }
161 if failed_count > 0 {
162 format!("{}, {} failed to restart", msg, failed_count)
163 } else {
164 msg
165 }
166 } else {
167 let msg = "No failed services found".to_owned();
168 if !quiet {
169 CliService::info(&msg);
170 }
171 msg
172 };
173
174 let output = RestartOutput {
175 service_type: "failed".to_owned(),
176 service_name: None,
177 restarted_count,
178 failed_count,
179 message,
180 };
181
182 Ok(CommandOutput::card_value(
183 "Restart Failed Services",
184 &output,
185 ))
186}
187
188async fn restart_mcp_target(
189 orchestrator: &systemprompt_mcp::services::McpOrchestrator,
190 target: &RestartTarget,
191 restarted: &mut usize,
192 failed: &mut usize,
193 quiet: bool,
194) {
195 match orchestrator.restart_services(Some(target.id.clone())).await {
196 Ok(()) => {
197 *restarted += 1;
198 if !quiet {
199 CliService::success(&format!(" {} restarted", target.name));
200 }
201 },
202 Err(e) => {
203 *failed += 1;
204 if !quiet {
205 CliService::error(&format!(" Failed to restart {}: {}", target.name, e));
206 }
207 },
208 }
209}
210
211async fn restart_agent_target(
212 orchestrator: &AgentOrchestrator,
213 target: &RestartTarget,
214 restarted: &mut usize,
215 failed: &mut usize,
216 quiet: bool,
217) {
218 match orchestrator.restart_agent(&target.id, None).await {
219 Ok(_) => {
220 *restarted += 1;
221 if !quiet {
222 CliService::success(&format!(" {} restarted", target.name));
223 }
224 },
225 Err(e) => {
226 *failed += 1;
227 if !quiet {
228 CliService::error(&format!(" Failed to restart {}: {}", target.name, e));
229 }
230 },
231 }
232}
233
234async fn agent_snapshots(orchestrator: &AgentOrchestrator) -> Result<Vec<ServiceSnapshot>> {
235 let agent_registry = AgentRegistry::new()?;
236 let all_agents = orchestrator.list_all().await?;
237
238 let mut snapshots = Vec::with_capacity(all_agents.len());
239 for (agent_id, status) in &all_agents {
240 let Ok(agent_config) = agent_registry.get_agent(agent_id).await else {
241 continue;
242 };
243
244 snapshots.push(ServiceSnapshot {
245 service_type: ServiceType::Agent,
246 id: agent_id.clone(),
247 name: agent_config.name,
248 enabled: agent_config.enabled,
249 healthy: !matches!(status, AgentStatus::Failed { .. }),
250 });
251 }
252 Ok(snapshots)
253}
254
255async fn mcp_snapshots(ctx: &Arc<AppContext>, probe_health: bool) -> Result<Vec<ServiceSnapshot>> {
256 ctx.mcp_registry().validate()?;
257 let servers = ctx.mcp_registry().get_managed_servers()?;
258
259 let health_by_name: HashMap<String, HealthStatus> = if probe_health {
260 let manager = systemprompt_mcp::services::McpOrchestrator::new(
261 Arc::clone(ctx.db_pool()),
262 Arc::clone(ctx.app_paths_arc()),
263 ctx.mcp_registry().clone(),
264 )?;
265 manager
266 .service_statuses()
267 .await?
268 .into_iter()
269 .map(|status| (status.name, status.health))
270 .collect()
271 } else {
272 HashMap::new()
273 };
274
275 let mut snapshots = Vec::with_capacity(servers.len());
276 for server in servers {
277 let healthy = if probe_health {
278 health_by_name
279 .get(&server.name)
280 .is_some_and(|h| matches!(h, HealthStatus::Healthy | HealthStatus::Degraded))
281 } else {
282 true
283 };
284
285 snapshots.push(ServiceSnapshot {
286 service_type: ServiceType::Mcp,
287 id: server.name.clone(),
288 name: server.name.clone(),
289 enabled: true,
290 healthy,
291 });
292 }
293 Ok(snapshots)
294}