use chrono::Utc;
use rusqlite::OptionalExtension;
use std::collections::HashMap;
use crate::db::Database;
use crate::error::ServerError;
use crate::models::agent::{Agent, AgentRole, AgentStatus, ModelTier};
#[derive(Clone)]
pub struct AgentStore {
db: Database,
}
impl AgentStore {
pub fn new(db: Database) -> Self {
Self { db }
}
pub async fn save(&self, agent: &Agent) -> Result<(), ServerError> {
let a = agent.clone();
self.db
.with_conn_async(move |conn| {
conn.execute(
"INSERT INTO agents (id, name, role, model_tier, workspace_id, parent_id, status, metadata, created_at, updated_at)
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10)
ON CONFLICT(id) DO UPDATE SET
name = excluded.name,
role = excluded.role,
model_tier = excluded.model_tier,
workspace_id = excluded.workspace_id,
parent_id = excluded.parent_id,
status = excluded.status,
metadata = excluded.metadata,
updated_at = excluded.updated_at",
rusqlite::params![
a.id,
a.name,
a.role.as_str(),
a.model_tier.as_str(),
a.workspace_id,
a.parent_id,
a.status.as_str(),
serde_json::to_string(&a.metadata).unwrap_or_default(),
a.created_at.timestamp_millis(),
a.updated_at.timestamp_millis(),
],
)?;
Ok(())
})
.await
}
pub async fn get(&self, agent_id: &str) -> Result<Option<Agent>, ServerError> {
let id = agent_id.to_string();
self.db
.with_conn_async(move |conn| {
let mut stmt = conn.prepare(
"SELECT id, name, role, model_tier, workspace_id, parent_id, status, metadata, created_at, updated_at
FROM agents WHERE id = ?1",
)?;
stmt.query_row(rusqlite::params![id], |row| Ok(row_to_agent(row)))
.optional()
})
.await
}
pub async fn list_by_workspace(&self, workspace_id: &str) -> Result<Vec<Agent>, ServerError> {
let ws_id = workspace_id.to_string();
self.db
.with_conn_async(move |conn| {
let mut stmt = conn.prepare(
"SELECT id, name, role, model_tier, workspace_id, parent_id, status, metadata, created_at, updated_at
FROM agents WHERE workspace_id = ?1 ORDER BY created_at DESC",
)?;
let rows = stmt
.query_map(rusqlite::params![ws_id], |row| Ok(row_to_agent(row)))?
.collect::<Result<Vec<_>, _>>()?;
Ok(rows)
})
.await
}
pub async fn list_by_parent(&self, parent_id: &str) -> Result<Vec<Agent>, ServerError> {
let pid = parent_id.to_string();
self.db
.with_conn_async(move |conn| {
let mut stmt = conn.prepare(
"SELECT id, name, role, model_tier, workspace_id, parent_id, status, metadata, created_at, updated_at
FROM agents WHERE parent_id = ?1 ORDER BY created_at DESC",
)?;
let rows = stmt
.query_map(rusqlite::params![pid], |row| Ok(row_to_agent(row)))?
.collect::<Result<Vec<_>, _>>()?;
Ok(rows)
})
.await
}
pub async fn list_by_role(
&self,
workspace_id: &str,
role: &AgentRole,
) -> Result<Vec<Agent>, ServerError> {
let ws_id = workspace_id.to_string();
let role_str = role.as_str().to_string();
self.db
.with_conn_async(move |conn| {
let mut stmt = conn.prepare(
"SELECT id, name, role, model_tier, workspace_id, parent_id, status, metadata, created_at, updated_at
FROM agents WHERE workspace_id = ?1 AND role = ?2 ORDER BY created_at DESC",
)?;
let rows = stmt
.query_map(rusqlite::params![ws_id, role_str], |row| Ok(row_to_agent(row)))?
.collect::<Result<Vec<_>, _>>()?;
Ok(rows)
})
.await
}
pub async fn list_by_status(
&self,
workspace_id: &str,
status: &AgentStatus,
) -> Result<Vec<Agent>, ServerError> {
let ws_id = workspace_id.to_string();
let status_str = status.as_str().to_string();
self.db
.with_conn_async(move |conn| {
let mut stmt = conn.prepare(
"SELECT id, name, role, model_tier, workspace_id, parent_id, status, metadata, created_at, updated_at
FROM agents WHERE workspace_id = ?1 AND status = ?2 ORDER BY created_at DESC",
)?;
let rows = stmt
.query_map(rusqlite::params![ws_id, status_str], |row| Ok(row_to_agent(row)))?
.collect::<Result<Vec<_>, _>>()?;
Ok(rows)
})
.await
}
pub async fn delete(&self, agent_id: &str) -> Result<(), ServerError> {
let id = agent_id.to_string();
self.db
.with_conn_async(move |conn| {
conn.execute("DELETE FROM agents WHERE id = ?1", rusqlite::params![id])?;
Ok(())
})
.await
}
pub async fn update_status(
&self,
agent_id: &str,
status: &AgentStatus,
) -> Result<(), ServerError> {
let id = agent_id.to_string();
let status_str = status.as_str().to_string();
let now = Utc::now().timestamp_millis();
self.db
.with_conn_async(move |conn| {
conn.execute(
"UPDATE agents SET status = ?1, updated_at = ?2 WHERE id = ?3",
rusqlite::params![status_str, now, id],
)?;
Ok(())
})
.await
}
}
use rusqlite::Row;
fn row_to_agent(row: &Row<'_>) -> Agent {
let metadata_str: String = row.get(7).unwrap_or_default();
let metadata: HashMap<String, String> = serde_json::from_str(&metadata_str).unwrap_or_default();
let created_ms: i64 = row.get(8).unwrap_or(0);
let updated_ms: i64 = row.get(9).unwrap_or(0);
Agent {
id: row.get(0).unwrap_or_default(),
name: row.get(1).unwrap_or_default(),
role: AgentRole::from_str(&row.get::<_, String>(2).unwrap_or_default())
.unwrap_or(AgentRole::Developer),
model_tier: ModelTier::from_str(&row.get::<_, String>(3).unwrap_or_default())
.unwrap_or(ModelTier::Smart),
workspace_id: row.get(4).unwrap_or_default(),
parent_id: row.get(5).unwrap_or(None),
status: AgentStatus::from_str(&row.get::<_, String>(6).unwrap_or_default())
.unwrap_or(AgentStatus::Pending),
metadata,
created_at: chrono::DateTime::from_timestamp_millis(created_ms).unwrap_or_else(Utc::now),
updated_at: chrono::DateTime::from_timestamp_millis(updated_ms).unwrap_or_else(Utc::now),
}
}