systemprompt_cli/commands/infrastructure/services/
restart.rs1use crate::cli_settings::CliConfig;
2use crate::shared::CommandResult;
3use anyhow::{Context, Result};
4use std::sync::Arc;
5use systemprompt_agent::services::agent_orchestration::AgentOrchestrator;
6use systemprompt_agent::services::registry::AgentRegistry;
7use systemprompt_agent::AgentState;
8use systemprompt_logging::CliService;
9use systemprompt_mcp::services::McpManager;
10use systemprompt_models::ProfileBootstrap;
11use systemprompt_oauth::JwtValidationProviderImpl;
12use systemprompt_runtime::AppContext;
13use systemprompt_scheduler::ProcessCleanup;
14
15use super::types::RestartOutput;
16
17const DEFAULT_API_PORT: u16 = 8080;
18
19fn create_agent_state(ctx: &AppContext) -> Result<Arc<AgentState>> {
20 let jwt_provider = Arc::new(
21 JwtValidationProviderImpl::from_config().context("Failed to create JWT provider")?,
22 );
23 Ok(Arc::new(AgentState::new(
24 Arc::clone(ctx.db_pool()),
25 Arc::new(ctx.config().clone()),
26 jwt_provider,
27 )))
28}
29
30fn get_api_port() -> u16 {
31 ProfileBootstrap::get()
32 .map(|p| p.server.port)
33 .unwrap_or(DEFAULT_API_PORT)
34}
35
36async fn resolve_name(agent_identifier: &str) -> Result<String> {
37 let registry = AgentRegistry::new().await?;
38 let agent = registry.get_agent(agent_identifier).await?;
39 Ok(agent.name)
40}
41
42pub async fn execute_api(config: &CliConfig) -> Result<CommandResult<RestartOutput>> {
43 let quiet = config.is_json_output();
44
45 if !quiet {
46 CliService::section("Restarting API Server");
47 }
48
49 let port = get_api_port();
50 let Some(pid) = ProcessCleanup::check_port(port) else {
51 if !quiet {
52 CliService::warning("API server is not running");
53 CliService::info("Starting API server...");
54 }
55 super::serve::execute(true, false, config).await?;
56 let output = RestartOutput {
57 service_type: "api".to_string(),
58 service_name: None,
59 restarted_count: 1,
60 failed_count: 0,
61 message: "API server started (was not running)".to_string(),
62 };
63 return Ok(CommandResult::card(output).with_title("Restart API Server"));
64 };
65
66 if !quiet {
67 CliService::info(&format!("Stopping API server (PID: {})...", pid));
68 }
69
70 ProcessCleanup::terminate_gracefully(pid, 100).await;
71 ProcessCleanup::kill_port(port);
72
73 ProcessCleanup::wait_for_port_free(port, 5, 500).await?;
74
75 if !quiet {
76 CliService::success("API server stopped");
77 CliService::info("Starting API server...");
78 }
79
80 super::serve::execute(true, false, config).await?;
81
82 let message = "API server restarted successfully".to_string();
83 if !quiet {
84 CliService::success(&message);
85 }
86
87 let output = RestartOutput {
88 service_type: "api".to_string(),
89 service_name: None,
90 restarted_count: 1,
91 failed_count: 0,
92 message,
93 };
94
95 Ok(CommandResult::card(output).with_title("Restart API Server"))
96}
97
98pub async fn execute_agent(
99 ctx: &Arc<AppContext>,
100 agent_id: &str,
101 config: &CliConfig,
102) -> Result<CommandResult<RestartOutput>> {
103 let quiet = config.is_json_output();
104
105 if !quiet {
106 CliService::section(&format!("Restarting Agent: {}", agent_id));
107 }
108
109 let agent_state = create_agent_state(ctx)?;
110 let orchestrator = AgentOrchestrator::new(agent_state, None)
111 .await
112 .context("Failed to initialize agent orchestrator")?;
113
114 let name = resolve_name(agent_id).await?;
115 let service_id = orchestrator.restart_agent(&name, None).await?;
116
117 let message = format!(
118 "Agent {} restarted successfully (service ID: {})",
119 agent_id, service_id
120 );
121 if !quiet {
122 CliService::success(&message);
123 }
124
125 let output = RestartOutput {
126 service_type: "agent".to_string(),
127 service_name: Some(agent_id.to_string()),
128 restarted_count: 1,
129 failed_count: 0,
130 message,
131 };
132
133 Ok(CommandResult::card(output).with_title("Restart Agent"))
134}
135
136pub async fn execute_mcp(
137 ctx: &Arc<AppContext>,
138 server_name: &str,
139 build: bool,
140 config: &CliConfig,
141) -> Result<CommandResult<RestartOutput>> {
142 let quiet = config.is_json_output();
143 let action = if build {
144 "Building and restarting"
145 } else {
146 "Restarting"
147 };
148
149 if !quiet {
150 CliService::section(&format!("{} MCP Server: {}", action, server_name));
151 }
152
153 let manager =
154 McpManager::new(Arc::clone(ctx.db_pool())).context("Failed to initialize MCP manager")?;
155
156 if build {
157 manager
158 .build_and_restart_services(Some(server_name.to_string()))
159 .await?;
160 } else {
161 manager
162 .restart_services_sync(Some(server_name.to_string()))
163 .await?;
164 }
165
166 let message = format!("MCP server {} restarted successfully", server_name);
167 if !quiet {
168 CliService::success(&message);
169 }
170
171 let output = RestartOutput {
172 service_type: "mcp".to_string(),
173 service_name: Some(server_name.to_string()),
174 restarted_count: 1,
175 failed_count: 0,
176 message,
177 };
178
179 Ok(CommandResult::card(output).with_title("Restart MCP Server"))
180}
181
182pub async fn execute_all_agents(
183 ctx: &Arc<AppContext>,
184 config: &CliConfig,
185) -> Result<CommandResult<RestartOutput>> {
186 let quiet = config.is_json_output();
187
188 if !quiet {
189 CliService::section("Restarting All Agents");
190 }
191
192 let agent_state = create_agent_state(ctx)?;
193 let orchestrator = AgentOrchestrator::new(agent_state, None)
194 .await
195 .context("Failed to initialize agent orchestrator")?;
196
197 let agent_registry = AgentRegistry::new().await?;
198 let all_agents = orchestrator.list_all().await?;
199
200 let mut restarted = 0usize;
201 let mut failed = 0usize;
202
203 for (agent_id, _status) in &all_agents {
204 let Ok(agent_config) = agent_registry.get_agent(agent_id).await else {
205 continue;
206 };
207
208 if !agent_config.enabled {
209 continue;
210 }
211
212 if !quiet {
213 CliService::info(&format!("Restarting agent: {}", agent_config.name));
214 }
215 match orchestrator.restart_agent(agent_id, None).await {
216 Ok(_) => {
217 restarted += 1;
218 if !quiet {
219 CliService::success(&format!(" {} restarted", agent_config.name));
220 }
221 },
222 Err(e) => {
223 failed += 1;
224 if !quiet {
225 CliService::error(&format!(" Failed to restart {}: {}", agent_config.name, e));
226 }
227 },
228 }
229 }
230
231 let message = match (restarted, failed) {
232 (0, 0) => {
233 if !quiet {
234 CliService::info("No enabled agents found");
235 }
236 "No enabled agents found".to_string()
237 },
238 (r, 0) => {
239 let msg = format!("Restarted {} agents", r);
240 if !quiet {
241 CliService::success(&msg);
242 }
243 msg
244 },
245 (0, f) => {
246 let msg = format!("Failed to restart {} agents", f);
247 if !quiet {
248 CliService::warning(&msg);
249 }
250 msg
251 },
252 (r, f) => {
253 if !quiet {
254 CliService::success(&format!("Restarted {} agents", r));
255 CliService::warning(&format!("Failed to restart {} agents", f));
256 }
257 format!("Restarted {} agents, {} failed", r, f)
258 },
259 };
260
261 let output = RestartOutput {
262 service_type: "agents".to_string(),
263 service_name: None,
264 restarted_count: restarted,
265 failed_count: failed,
266 message,
267 };
268
269 Ok(CommandResult::card(output).with_title("Restart All Agents"))
270}
271
272pub async fn execute_all_mcp(
273 ctx: &Arc<AppContext>,
274 config: &CliConfig,
275) -> Result<CommandResult<RestartOutput>> {
276 let quiet = config.is_json_output();
277
278 if !quiet {
279 CliService::section("Restarting All MCP Servers");
280 }
281
282 let mcp_manager =
283 McpManager::new(Arc::clone(ctx.db_pool())).context("Failed to initialize MCP manager")?;
284
285 systemprompt_mcp::services::RegistryManager::validate()?;
286 let servers = systemprompt_mcp::services::RegistryManager::get_enabled_servers()?;
287
288 let mut restarted = 0usize;
289 let mut failed = 0usize;
290
291 for server in servers {
292 if !server.enabled {
293 continue;
294 }
295
296 if !quiet {
297 CliService::info(&format!("Restarting MCP server: {}", server.name));
298 }
299 match mcp_manager
300 .restart_services(Some(server.name.clone()))
301 .await
302 {
303 Ok(()) => {
304 restarted += 1;
305 if !quiet {
306 CliService::success(&format!(" {} restarted", server.name));
307 }
308 },
309 Err(e) => {
310 failed += 1;
311 if !quiet {
312 CliService::error(&format!(" Failed to restart {}: {}", server.name, e));
313 }
314 },
315 }
316 }
317
318 let message = match (restarted, failed) {
319 (0, 0) => {
320 if !quiet {
321 CliService::info("No enabled MCP servers found");
322 }
323 "No enabled MCP servers found".to_string()
324 },
325 (r, 0) => {
326 let msg = format!("Restarted {} MCP servers", r);
327 if !quiet {
328 CliService::success(&msg);
329 }
330 msg
331 },
332 (0, f) => {
333 let msg = format!("Failed to restart {} MCP servers", f);
334 if !quiet {
335 CliService::warning(&msg);
336 }
337 msg
338 },
339 (r, f) => {
340 if !quiet {
341 CliService::success(&format!("Restarted {} MCP servers", r));
342 CliService::warning(&format!("Failed to restart {} MCP servers", f));
343 }
344 format!("Restarted {} MCP servers, {} failed", r, f)
345 },
346 };
347
348 let output = RestartOutput {
349 service_type: "mcp".to_string(),
350 service_name: None,
351 restarted_count: restarted,
352 failed_count: failed,
353 message,
354 };
355
356 Ok(CommandResult::card(output).with_title("Restart All MCP Servers"))
357}
358
359pub async fn execute_failed(
360 ctx: &Arc<AppContext>,
361 config: &CliConfig,
362) -> Result<CommandResult<RestartOutput>> {
363 let quiet = config.is_json_output();
364
365 if !quiet {
366 CliService::section("Restarting Failed Services");
367 }
368
369 let mut restarted_count = 0usize;
370 let mut failed_count = 0usize;
371
372 restart_failed_agents(ctx, &mut restarted_count, &mut failed_count, quiet).await?;
373 restart_failed_mcp(ctx, &mut restarted_count, &mut failed_count, quiet).await?;
374
375 let message = if restarted_count > 0 {
376 let msg = format!("Restarted {} failed services", restarted_count);
377 if !quiet {
378 CliService::success(&msg);
379 }
380 if failed_count > 0 && !quiet {
381 CliService::warning(&format!("Failed to restart {} services", failed_count));
382 }
383 if failed_count > 0 {
384 format!("{}, {} failed to restart", msg, failed_count)
385 } else {
386 msg
387 }
388 } else {
389 let msg = "No failed services found".to_string();
390 if !quiet {
391 CliService::info(&msg);
392 }
393 msg
394 };
395
396 let output = RestartOutput {
397 service_type: "failed".to_string(),
398 service_name: None,
399 restarted_count,
400 failed_count,
401 message,
402 };
403
404 Ok(CommandResult::card(output).with_title("Restart Failed Services"))
405}
406
407async fn restart_failed_agents(
408 ctx: &Arc<AppContext>,
409 restarted_count: &mut usize,
410 failed_count: &mut usize,
411 quiet: bool,
412) -> Result<()> {
413 let agent_state = create_agent_state(ctx)?;
414 let orchestrator = AgentOrchestrator::new(agent_state, None)
415 .await
416 .context("Failed to initialize agent orchestrator")?;
417
418 let agent_registry = AgentRegistry::new().await?;
419
420 let all_agents = orchestrator.list_all().await?;
421 for (agent_id, status) in &all_agents {
422 let Ok(agent_config) = agent_registry.get_agent(agent_id).await else {
423 continue;
424 };
425
426 if !agent_config.enabled {
427 continue;
428 }
429
430 if let systemprompt_agent::services::agent_orchestration::AgentStatus::Failed { .. } =
431 status
432 {
433 if !quiet {
434 CliService::info(&format!("Restarting failed agent: {}", agent_config.name));
435 }
436 match orchestrator.restart_agent(agent_id, None).await {
437 Ok(_) => {
438 *restarted_count += 1;
439 if !quiet {
440 CliService::success(&format!(" {} restarted", agent_config.name));
441 }
442 },
443 Err(e) => {
444 *failed_count += 1;
445 if !quiet {
446 CliService::error(&format!(
447 " Failed to restart {}: {}",
448 agent_config.name, e
449 ));
450 }
451 },
452 }
453 }
454 }
455
456 Ok(())
457}
458
459async fn restart_failed_mcp(
460 ctx: &Arc<AppContext>,
461 restarted_count: &mut usize,
462 failed_count: &mut usize,
463 quiet: bool,
464) -> Result<()> {
465 let mcp_manager =
466 McpManager::new(Arc::clone(ctx.db_pool())).context("Failed to initialize MCP manager")?;
467
468 systemprompt_mcp::services::RegistryManager::validate()?;
469 let servers = systemprompt_mcp::services::RegistryManager::get_enabled_servers()?;
470
471 for server in servers {
472 if !server.enabled {
473 continue;
474 }
475
476 let database = systemprompt_mcp::services::DatabaseManager::new(Arc::clone(ctx.db_pool()));
477 let service_info = database.get_service_by_name(&server.name).await?;
478
479 let needs_restart = match service_info {
480 Some(info) => info.status != "running",
481 None => true,
482 };
483
484 if needs_restart {
485 if !quiet {
486 CliService::info(&format!("Restarting MCP server: {}", server.name));
487 }
488 match mcp_manager
489 .restart_services(Some(server.name.clone()))
490 .await
491 {
492 Ok(()) => {
493 *restarted_count += 1;
494 if !quiet {
495 CliService::success(&format!(" {} restarted", server.name));
496 }
497 },
498 Err(e) => {
499 *failed_count += 1;
500 if !quiet {
501 CliService::error(&format!(" Failed to restart {}: {}", server.name, e));
502 }
503 },
504 }
505 }
506 }
507
508 Ok(())
509}