systemprompt_agent/services/agent_orchestration/
database.rs

1//! Database-backed view of agent service state, reconciled against live
2//! processes.
3//!
4//! [`AgentDatabaseService`] wraps the agent-service repository and the
5//! config-driven [`AgentRegistry`], translating stored rows into
6//! [`AgentStatus`] while verifying that recorded PIDs still correspond to
7//! running processes. It is the single source of truth the lifecycle, monitor,
8//! and reconciler services query and mutate.
9
10use crate::repository::agent_service::AgentServiceRepository;
11use crate::services::agent_orchestration::{
12    AgentStatus, OrchestrationError, OrchestrationResult, process,
13};
14use crate::services::registry::AgentRegistry;
15use systemprompt_models::services::AgentConfig;
16
17#[derive(Debug)]
18pub struct AgentDatabaseService {
19    pub repository: AgentServiceRepository,
20    pub registry: AgentRegistry,
21}
22
23impl AgentDatabaseService {
24    pub fn new(repository: AgentServiceRepository) -> OrchestrationResult<Self> {
25        let registry = AgentRegistry::new().map_err(|e| {
26            OrchestrationError::Database(format!("Failed to load agent registry: {e}"))
27        })?;
28
29        Ok(Self {
30            repository,
31            registry,
32        })
33    }
34
35    pub async fn register_agent(
36        &self,
37        name: &str,
38        pid: u32,
39        port: u16,
40    ) -> OrchestrationResult<String> {
41        self.repository
42            .register_agent(name, pid, port)
43            .await
44            .map_err(|e| OrchestrationError::Database(e.to_string()))
45    }
46
47    pub async fn get_status(&self, agent_name: &str) -> OrchestrationResult<AgentStatus> {
48        let row = self
49            .repository
50            .get_agent_status(agent_name)
51            .await
52            .map_err(|e| OrchestrationError::Database(e.to_string()))?;
53
54        match row {
55            Some(r) => match (r.pid, r.status.as_str()) {
56                (Some(pid), "running") => {
57                    let pid = pid as u32;
58                    if process::process_exists(pid) {
59                        Ok(AgentStatus::Running {
60                            pid,
61                            port: r.port as u16,
62                        })
63                    } else {
64                        self.mark_failed(agent_name).await?;
65                        Ok(AgentStatus::Failed {
66                            reason: "Process died unexpectedly".to_owned(),
67                            last_attempt: None,
68                            retry_count: 0,
69                        })
70                    }
71                },
72                (_, "starting") => Ok(AgentStatus::Failed {
73                    reason: "Agent is starting".to_owned(),
74                    last_attempt: None,
75                    retry_count: 0,
76                }),
77                (_, "failed" | "crashed" | "stopped") => {
78                    let error_msg = self
79                        .get_error_message(agent_name)
80                        .await
81                        .unwrap_or_else(|_| "Unknown failure".to_owned());
82                    Ok(AgentStatus::Failed {
83                        reason: error_msg,
84                        last_attempt: None,
85                        retry_count: 0,
86                    })
87                },
88                _ => {
89                    self.mark_failed(agent_name).await?;
90                    Ok(AgentStatus::Failed {
91                        reason: "Invalid database state".to_owned(),
92                        last_attempt: None,
93                        retry_count: 0,
94                    })
95                },
96            },
97            None => Ok(AgentStatus::Failed {
98                reason: "No service record found".to_owned(),
99                last_attempt: None,
100                retry_count: 0,
101            }),
102        }
103    }
104
105    pub async fn mark_failed(&self, agent_name: &str) -> OrchestrationResult<()> {
106        self.repository
107            .mark_error(agent_name)
108            .await
109            .map_err(|e| OrchestrationError::Database(e.to_string()))?;
110
111        self.repository
112            .mark_crashed(agent_name)
113            .await
114            .map_err(|e| OrchestrationError::Database(e.to_string()))
115    }
116
117    pub async fn mark_crashed(&self, agent_name: &str) -> OrchestrationResult<()> {
118        self.mark_failed(agent_name).await
119    }
120
121    pub async fn get_error_message(&self, agent_name: &str) -> OrchestrationResult<String> {
122        let row = self
123            .repository
124            .get_agent_status(agent_name)
125            .await
126            .map_err(|e| OrchestrationError::Database(e.to_string()))?;
127
128        match row {
129            Some(r) => Ok(format!("Status: {}", r.status)),
130            None => Ok("No service record".to_owned()),
131        }
132    }
133
134    pub async fn mark_error(&self, agent_name: &str) -> OrchestrationResult<()> {
135        self.repository
136            .mark_error(agent_name)
137            .await
138            .map_err(|e| OrchestrationError::Database(e.to_string()))
139    }
140
141    pub async fn list_running_agents(&self) -> OrchestrationResult<Vec<String>> {
142        let rows = self
143            .repository
144            .list_running_agents()
145            .await
146            .map_err(|e| OrchestrationError::Database(e.to_string()))?;
147
148        Ok(rows.into_iter().map(|row| row.name).collect())
149    }
150
151    pub async fn list_all_agents(&self) -> OrchestrationResult<Vec<(String, AgentStatus)>> {
152        let agent_configs = self.registry.list_agents().await.map_err(|e| {
153            OrchestrationError::Database(format!("Failed to list agents from config: {e}"))
154        })?;
155
156        let mut agents = Vec::new();
157
158        for agent_config in agent_configs {
159            let agent_name = &agent_config.name;
160
161            let status = self.get_status(agent_name).await?;
162
163            agents.push((agent_name.clone(), status));
164        }
165
166        Ok(agents)
167    }
168
169    pub async fn agent_exists(&self, agent_name: &str) -> OrchestrationResult<bool> {
170        self.registry
171            .get_agent(agent_name)
172            .await
173            .map(|_| true)
174            .or_else(|_| Ok(false))
175    }
176
177    pub async fn get_agent_config(&self, agent_name: &str) -> OrchestrationResult<AgentConfig> {
178        let agent_config = self.registry.get_agent(agent_name).await.map_err(|e| {
179            OrchestrationError::AgentNotFound(format!(
180                "Agent {} not found in config: {}",
181                agent_name, e
182            ))
183        })?;
184
185        Ok(agent_config)
186    }
187
188    pub async fn cleanup_orphaned_services(&self) -> OrchestrationResult<u64> {
189        let rows = self
190            .repository
191            .list_running_agent_pids()
192            .await
193            .map_err(|e| OrchestrationError::Database(e.to_string()))?;
194
195        let mut cleaned = 0u64;
196
197        for row in rows {
198            let pid = row.pid as u32;
199            if !process::process_exists(pid) {
200                self.mark_crashed(&row.name).await?;
201                cleaned += 1;
202            }
203        }
204
205        Ok(cleaned)
206    }
207
208    pub async fn remove_agent_service(&self, agent_name: &str) -> OrchestrationResult<()> {
209        self.repository
210            .remove_agent_service(agent_name)
211            .await
212            .map_err(|e| OrchestrationError::Database(e.to_string()))
213    }
214
215    pub async fn update_health_status(
216        &self,
217        agent_name: &str,
218        health_status: &str,
219    ) -> OrchestrationResult<()> {
220        self.repository
221            .update_health_status(agent_name, health_status)
222            .await
223            .map_err(|e| OrchestrationError::Database(e.to_string()))
224    }
225
226    pub async fn update_agent_running(
227        &self,
228        agent_name: &str,
229        pid: u32,
230        port: u16,
231    ) -> OrchestrationResult<String> {
232        self.repository
233            .register_agent(agent_name, pid, port)
234            .await
235            .map_err(|e| OrchestrationError::Database(e.to_string()))
236    }
237
238    pub async fn update_agent_stopped(&self, agent_name: &str) -> OrchestrationResult<()> {
239        self.repository
240            .mark_stopped(agent_name)
241            .await
242            .map_err(|e| OrchestrationError::Database(e.to_string()))
243    }
244
245    pub async fn register_agent_starting(
246        &self,
247        agent_name: &str,
248        pid: u32,
249        port: u16,
250    ) -> OrchestrationResult<String> {
251        self.repository
252            .register_agent_starting(agent_name, pid, port)
253            .await
254            .map_err(|e| OrchestrationError::Database(e.to_string()))
255    }
256
257    pub async fn mark_running(&self, agent_name: &str) -> OrchestrationResult<()> {
258        self.repository
259            .mark_running(agent_name)
260            .await
261            .map_err(|e| OrchestrationError::Database(e.to_string()))
262    }
263
264    pub async fn get_unresponsive_agents(&self) -> OrchestrationResult<Vec<(String, Option<u32>)>> {
265        use crate::services::agent_orchestration::monitor::check_a2a_agent_health;
266
267        let agents = self.list_all_agents().await?;
268
269        let mut unresponsive = Vec::new();
270        for (agent_name, status) in agents {
271            if let AgentStatus::Running { pid, port, .. } = status {
272                let is_healthy = check_a2a_agent_health(port, 10).await.unwrap_or(false);
273
274                if !is_healthy {
275                    unresponsive.push((agent_name, Some(pid)));
276                }
277            }
278        }
279
280        Ok(unresponsive)
281    }
282}
systemprompt_agent/services/agent_orchestration/database.rs

systemprompt_agent/services/agent_orchestration/
database.rs
