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::McpManager;
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_string(),
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 =
81 McpManager::new(Arc::clone(ctx.db_pool())).context("Failed to initialize MCP manager")?;
82
83 systemprompt_mcp::services::RegistryManager::validate()?;
84 let servers = systemprompt_mcp::services::RegistryManager::get_enabled_servers()?;
85
86 let mut restarted = 0usize;
87 let mut failed = 0usize;
88
89 for server in servers {
90 if !server.enabled {
91 continue;
92 }
93
94 if !quiet {
95 CliService::info(&format!("Restarting MCP server: {}", server.name));
96 }
97 match mcp_manager
98 .restart_services(Some(server.name.clone()))
99 .await
100 {
101 Ok(()) => {
102 restarted += 1;
103 if !quiet {
104 CliService::success(&format!(" {} restarted", server.name));
105 }
106 },
107 Err(e) => {
108 failed += 1;
109 if !quiet {
110 CliService::error(&format!(" Failed to restart {}: {}", server.name, e));
111 }
112 },
113 }
114 }
115
116 let message = super::format_batch_message("MCP servers", restarted, failed, quiet);
117
118 let output = RestartOutput {
119 service_type: "mcp".to_string(),
120 service_name: None,
121 restarted_count: restarted,
122 failed_count: failed,
123 message,
124 };
125
126 Ok(CommandResult::card(output).with_title("Restart All MCP Servers"))
127}
128
129pub async fn execute_failed(
130 ctx: &Arc<AppContext>,
131 config: &CliConfig,
132) -> Result<CommandResult<RestartOutput>> {
133 let quiet = config.is_json_output();
134
135 if !quiet {
136 CliService::section("Restarting Failed Services");
137 }
138
139 let mut restarted_count = 0usize;
140 let mut failed_count = 0usize;
141
142 restart_failed_agents(ctx, &mut restarted_count, &mut failed_count, quiet).await?;
143 restart_failed_mcp(ctx, &mut restarted_count, &mut failed_count, quiet).await?;
144
145 let message = if restarted_count > 0 {
146 let msg = format!("Restarted {} failed services", restarted_count);
147 if !quiet {
148 CliService::success(&msg);
149 }
150 if failed_count > 0 && !quiet {
151 CliService::warning(&format!("Failed to restart {} services", failed_count));
152 }
153 if failed_count > 0 {
154 format!("{}, {} failed to restart", msg, failed_count)
155 } else {
156 msg
157 }
158 } else {
159 let msg = "No failed services found".to_string();
160 if !quiet {
161 CliService::info(&msg);
162 }
163 msg
164 };
165
166 let output = RestartOutput {
167 service_type: "failed".to_string(),
168 service_name: None,
169 restarted_count,
170 failed_count,
171 message,
172 };
173
174 Ok(CommandResult::card(output).with_title("Restart Failed Services"))
175}
176
177async fn restart_failed_agents(
178 ctx: &Arc<AppContext>,
179 restarted_count: &mut usize,
180 failed_count: &mut usize,
181 quiet: bool,
182) -> Result<()> {
183 let orchestrator = super::create_orchestrator(ctx).await?;
184 let agent_registry = AgentRegistry::new()?;
185
186 let all_agents = orchestrator.list_all().await?;
187 for (agent_id, status) in &all_agents {
188 let Ok(agent_config) = agent_registry.get_agent(agent_id).await else {
189 continue;
190 };
191
192 if !agent_config.enabled {
193 continue;
194 }
195
196 if let systemprompt_agent::services::agent_orchestration::AgentStatus::Failed { .. } =
197 status
198 {
199 if !quiet {
200 CliService::info(&format!("Restarting failed agent: {}", agent_config.name));
201 }
202 match orchestrator.restart_agent(agent_id, None).await {
203 Ok(_) => {
204 *restarted_count += 1;
205 if !quiet {
206 CliService::success(&format!(" {} restarted", agent_config.name));
207 }
208 },
209 Err(e) => {
210 *failed_count += 1;
211 if !quiet {
212 CliService::error(&format!(
213 " Failed to restart {}: {}",
214 agent_config.name, e
215 ));
216 }
217 },
218 }
219 }
220 }
221
222 Ok(())
223}
224
225async fn restart_failed_mcp(
226 ctx: &Arc<AppContext>,
227 restarted_count: &mut usize,
228 failed_count: &mut usize,
229 quiet: bool,
230) -> Result<()> {
231 let mcp_manager =
232 McpManager::new(Arc::clone(ctx.db_pool())).context("Failed to initialize MCP manager")?;
233
234 systemprompt_mcp::services::RegistryManager::validate()?;
235 let servers = systemprompt_mcp::services::RegistryManager::get_enabled_servers()?;
236
237 for server in servers {
238 if !server.enabled {
239 continue;
240 }
241
242 let database = systemprompt_mcp::services::DatabaseManager::new(Arc::clone(ctx.db_pool()));
243 let service_info = database.get_service_by_name(&server.name).await?;
244
245 let needs_restart = match service_info {
246 Some(info) => info.status != "running",
247 None => true,
248 };
249
250 if needs_restart {
251 if !quiet {
252 CliService::info(&format!("Restarting MCP server: {}", server.name));
253 }
254 match mcp_manager
255 .restart_services(Some(server.name.clone()))
256 .await
257 {
258 Ok(()) => {
259 *restarted_count += 1;
260 if !quiet {
261 CliService::success(&format!(" {} restarted", server.name));
262 }
263 },
264 Err(e) => {
265 *failed_count += 1;
266 if !quiet {
267 CliService::error(&format!(" Failed to restart {}: {}", server.name, e));
268 }
269 },
270 }
271 }
272 }
273
274 Ok(())
275}