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