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