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

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

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