systemprompt_cli/commands/infrastructure/services/restart/
batch.rs

1use crate::cli_settings::CliConfig;
2use crate::shared::CommandOutput;
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<CommandOutput> {
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(CommandOutput::card_value("Restart All Agents", &output))
68}
69
70pub async fn execute_all_mcp(ctx: &Arc<AppContext>, config: &CliConfig) -> Result<CommandOutput> {
71    let quiet = config.is_json_output();
72
73    if !quiet {
74        CliService::section("Restarting All MCP Servers");
75    }
76
77    let mcp_manager = McpOrchestrator::new(
78        Arc::clone(ctx.db_pool()),
79        Arc::clone(ctx.app_paths_arc()),
80        ctx.mcp_registry().clone(),
81    )
82    .context("Failed to initialize MCP manager")?;
83
84    ctx.mcp_registry().validate()?;
85    let servers = ctx.mcp_registry().get_enabled_servers()?;
86
87    let mut restarted = 0usize;
88    let mut failed = 0usize;
89
90    for server in servers {
91        if !server.enabled {
92            continue;
93        }
94
95        if !quiet {
96            CliService::info(&format!("Restarting MCP server: {}", server.name));
97        }
98        match mcp_manager
99            .restart_services(Some(server.name.clone()))
100            .await
101        {
102            Ok(()) => {
103                restarted += 1;
104                if !quiet {
105                    CliService::success(&format!("  {} restarted", server.name));
106                }
107            },
108            Err(e) => {
109                failed += 1;
110                if !quiet {
111                    CliService::error(&format!("  Failed to restart {}: {}", server.name, e));
112                }
113            },
114        }
115    }
116
117    let message = super::format_batch_message("MCP servers", restarted, failed, quiet);
118
119    let output = RestartOutput {
120        service_type: "mcp".to_owned(),
121        service_name: None,
122        restarted_count: restarted,
123        failed_count: failed,
124        message,
125    };
126
127    Ok(CommandOutput::card_value(
128        "Restart All MCP Servers",
129        &output,
130    ))
131}
132
133pub async fn execute_failed(ctx: &Arc<AppContext>, config: &CliConfig) -> Result<CommandOutput> {
134    let quiet = config.is_json_output();
135
136    if !quiet {
137        CliService::section("Restarting Failed Services");
138    }
139
140    let mut restarted_count = 0usize;
141    let mut failed_count = 0usize;
142
143    restart_failed_agents(ctx, &mut restarted_count, &mut failed_count, quiet).await?;
144    restart_failed_mcp(ctx, &mut restarted_count, &mut failed_count, quiet).await?;
145
146    let message = if restarted_count > 0 {
147        let msg = format!("Restarted {} failed services", restarted_count);
148        if !quiet {
149            CliService::success(&msg);
150        }
151        if failed_count > 0 && !quiet {
152            CliService::warning(&format!("Failed to restart {} services", failed_count));
153        }
154        if failed_count > 0 {
155            format!("{}, {} failed to restart", msg, failed_count)
156        } else {
157            msg
158        }
159    } else {
160        let msg = "No failed services found".to_owned();
161        if !quiet {
162            CliService::info(&msg);
163        }
164        msg
165    };
166
167    let output = RestartOutput {
168        service_type: "failed".to_owned(),
169        service_name: None,
170        restarted_count,
171        failed_count,
172        message,
173    };
174
175    Ok(CommandOutput::card_value(
176        "Restart Failed Services",
177        &output,
178    ))
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}
systemprompt_cli/commands/infrastructure/services/restart/batch.rs

systemprompt_cli/commands/infrastructure/services/restart/
batch.rs
