systemprompt_agent/services/agent_orchestration/

database.rs

//! Database-backed view of agent service state, reconciled against live
//! processes.
//!
//! [`AgentDatabaseService`] wraps the agent-service repository and the
//! config-driven [`AgentRegistry`], translating stored rows into
//! [`AgentStatus`] while verifying that recorded PIDs still correspond to
//! running processes. It is the single source of truth the lifecycle, monitor,
//! and reconciler services query and mutate.
//!
//! Copyright (c) systemprompt.io — Business Source License 1.1.
//! See <https://systemprompt.io> for licensing details.

use crate::repository::agent_service::AgentServiceRepository;
use crate::services::agent_orchestration::{
    AgentStatus, OrchestrationError, OrchestrationResult, process,
};
use crate::services::registry::AgentRegistry;
use systemprompt_models::services::AgentConfig;

#[derive(Debug)]
pub struct AgentDatabaseService {
    pub repository: AgentServiceRepository,
    pub registry: AgentRegistry,
}

impl AgentDatabaseService {
    pub fn new(repository: AgentServiceRepository) -> OrchestrationResult<Self> {
        let registry = AgentRegistry::new().map_err(|e| {
            OrchestrationError::Database(format!("Failed to load agent registry: {e}"))
        })?;

        Ok(Self {
            repository,
            registry,
        })
    }

    #[must_use]
    pub const fn with_registry(
        repository: AgentServiceRepository,
        registry: AgentRegistry,
    ) -> Self {
        Self {
            repository,
            registry,
        }
    }

    pub async fn register_agent(
        &self,
        name: &str,
        pid: u32,
        port: u16,
    ) -> OrchestrationResult<String> {
        self.repository
            .register_agent(name, pid, port)
            .await
            .map_err(|e| OrchestrationError::Database(e.to_string()))
    }

    pub async fn get_status(&self, agent_name: &str) -> OrchestrationResult<AgentStatus> {
        let row = self
            .repository
            .get_agent_status(agent_name)
            .await
            .map_err(|e| OrchestrationError::Database(e.to_string()))?;

        match row {
            Some(r) => match (r.pid, r.status.as_str()) {
                (Some(pid), "running") => {
                    let pid = pid as u32;
                    if process::process_exists(pid) {
                        Ok(AgentStatus::Running {
                            pid,
                            port: r.port as u16,
                        })
                    } else {
                        self.mark_failed(agent_name).await?;
                        Ok(AgentStatus::Failed {
                            reason: "Process died unexpectedly".to_owned(),
                            last_attempt: None,
                            retry_count: 0,
                        })
                    }
                },
                (_, "starting") => Ok(AgentStatus::Failed {
                    reason: "Agent is starting".to_owned(),
                    last_attempt: None,
                    retry_count: 0,
                }),
                (_, "failed" | "crashed" | "stopped") => {
                    let error_msg = self
                        .get_error_message(agent_name)
                        .await
                        .unwrap_or_else(|_| "Unknown failure".to_owned());
                    Ok(AgentStatus::Failed {
                        reason: error_msg,
                        last_attempt: None,
                        retry_count: 0,
                    })
                },
                _ => {
                    self.mark_failed(agent_name).await?;
                    Ok(AgentStatus::Failed {
                        reason: "Invalid database state".to_owned(),
                        last_attempt: None,
                        retry_count: 0,
                    })
                },
            },
            None => Ok(AgentStatus::Failed {
                reason: "No service record found".to_owned(),
                last_attempt: None,
                retry_count: 0,
            }),
        }
    }

    pub async fn mark_failed(&self, agent_name: &str) -> OrchestrationResult<()> {
        self.repository
            .mark_error(agent_name)
            .await
            .map_err(|e| OrchestrationError::Database(e.to_string()))?;

        self.repository
            .mark_crashed(agent_name)
            .await
            .map_err(|e| OrchestrationError::Database(e.to_string()))
    }

    pub async fn mark_crashed(&self, agent_name: &str) -> OrchestrationResult<()> {
        self.mark_failed(agent_name).await
    }

    pub async fn get_error_message(&self, agent_name: &str) -> OrchestrationResult<String> {
        let row = self
            .repository
            .get_agent_status(agent_name)
            .await
            .map_err(|e| OrchestrationError::Database(e.to_string()))?;

        match row {
            Some(r) => Ok(format!("Status: {}", r.status)),
            None => Ok("No service record".to_owned()),
        }
    }

    pub async fn mark_error(&self, agent_name: &str) -> OrchestrationResult<()> {
        self.repository
            .mark_error(agent_name)
            .await
            .map_err(|e| OrchestrationError::Database(e.to_string()))
    }

    pub async fn list_running_agents(&self) -> OrchestrationResult<Vec<String>> {
        let rows = self
            .repository
            .list_running_agents()
            .await
            .map_err(|e| OrchestrationError::Database(e.to_string()))?;

        Ok(rows.into_iter().map(|row| row.name).collect())
    }

    pub async fn list_all_agents(&self) -> OrchestrationResult<Vec<(String, AgentStatus)>> {
        let agent_configs = self.registry.list_agents().await.map_err(|e| {
            OrchestrationError::Database(format!("Failed to list agents from config: {e}"))
        })?;

        let mut agents = Vec::new();

        for agent_config in agent_configs {
            let agent_name = &agent_config.name;

            let status = self.get_status(agent_name).await?;

            agents.push((agent_name.clone(), status));
        }

        Ok(agents)
    }

    pub async fn agent_exists(&self, agent_name: &str) -> OrchestrationResult<bool> {
        self.registry
            .get_agent(agent_name)
            .await
            .map(|_| true)
            .or_else(|_| Ok(false))
    }

    pub async fn get_agent_config(&self, agent_name: &str) -> OrchestrationResult<AgentConfig> {
        let agent_config = self.registry.get_agent(agent_name).await.map_err(|e| {
            OrchestrationError::AgentNotFound(format!(
                "Agent {} not found in config: {}",
                agent_name, e
            ))
        })?;

        Ok(agent_config)
    }

    pub async fn cleanup_orphaned_services(&self) -> OrchestrationResult<u64> {
        let rows = self
            .repository
            .list_running_agent_pids()
            .await
            .map_err(|e| OrchestrationError::Database(e.to_string()))?;

        let mut cleaned = 0u64;

        for row in rows {
            let pid = row.pid as u32;
            if !process::process_exists(pid) {
                self.mark_crashed(&row.name).await?;
                cleaned += 1;
            }
        }

        Ok(cleaned)
    }

    pub async fn remove_agent_service(&self, agent_name: &str) -> OrchestrationResult<()> {
        self.repository
            .remove_agent_service(agent_name)
            .await
            .map_err(|e| OrchestrationError::Database(e.to_string()))
    }

    pub async fn update_health_status(
        &self,
        agent_name: &str,
        health_status: &str,
    ) -> OrchestrationResult<()> {
        self.repository
            .update_health_status(agent_name, health_status)
            .await
            .map_err(|e| OrchestrationError::Database(e.to_string()))
    }

    pub async fn update_agent_running(
        &self,
        agent_name: &str,
        pid: u32,
        port: u16,
    ) -> OrchestrationResult<String> {
        self.repository
            .register_agent(agent_name, pid, port)
            .await
            .map_err(|e| OrchestrationError::Database(e.to_string()))
    }

    pub async fn update_agent_stopped(&self, agent_name: &str) -> OrchestrationResult<()> {
        self.repository
            .mark_stopped(agent_name)
            .await
            .map_err(|e| OrchestrationError::Database(e.to_string()))
    }

    pub async fn register_agent_starting(
        &self,
        agent_name: &str,
        pid: u32,
        port: u16,
    ) -> OrchestrationResult<String> {
        self.repository
            .register_agent_starting(agent_name, pid, port)
            .await
            .map_err(|e| OrchestrationError::Database(e.to_string()))
    }

    pub async fn mark_running(&self, agent_name: &str) -> OrchestrationResult<()> {
        self.repository
            .mark_running(agent_name)
            .await
            .map_err(|e| OrchestrationError::Database(e.to_string()))
    }

    pub async fn get_unresponsive_agents(&self) -> OrchestrationResult<Vec<(String, Option<u32>)>> {
        use crate::services::agent_orchestration::monitor::check_a2a_agent_health;

        let agents = self.list_all_agents().await?;

        let mut unresponsive = Vec::new();
        for (agent_name, status) in agents {
            if let AgentStatus::Running { pid, port, .. } = status {
                let is_healthy = check_a2a_agent_health(port, 10).await.unwrap_or(false);

                if !is_healthy {
                    unresponsive.push((agent_name, Some(pid)));
                }
            }
        }

        Ok(unresponsive)
    }
}