Skip to main content

ares/db/
tenant_agents.rs

1use crate::types::{AppError, Result};
2use serde::{Deserialize, Serialize};
3use sqlx::{PgPool, Row};
4use std::time::{SystemTime, UNIX_EPOCH};
5
6fn now_ts() -> i64 {
7    SystemTime::now()
8        .duration_since(UNIX_EPOCH)
9        .unwrap()
10        .as_secs() as i64
11}
12
13// =============================================================================
14// Structs
15// =============================================================================
16
17#[derive(Debug, Clone, Serialize, Deserialize)]
18pub struct TenantAgent {
19    pub id: String,
20    pub tenant_id: String,
21    pub agent_name: String,
22    pub display_name: String,
23    pub description: Option<String>,
24    pub config: serde_json::Value,
25    pub enabled: bool,
26    pub created_at: i64,
27    pub updated_at: i64,
28}
29
30#[derive(Debug, Clone, Serialize, Deserialize)]
31pub struct AgentTemplate {
32    pub id: String,
33    pub product_type: String,
34    pub agent_name: String,
35    pub display_name: String,
36    pub description: Option<String>,
37    pub config: serde_json::Value,
38    pub created_at: i64,
39}
40
41#[derive(Debug, Deserialize)]
42pub struct CreateTenantAgentRequest {
43    pub agent_name: String,
44    pub display_name: String,
45    pub description: Option<String>,
46    pub config: serde_json::Value,
47}
48
49#[derive(Debug, Deserialize)]
50pub struct UpdateTenantAgentRequest {
51    pub display_name: Option<String>,
52    pub description: Option<String>,
53    pub config: Option<serde_json::Value>,
54    pub enabled: Option<bool>,
55}
56
57// =============================================================================
58// Tenant Agent CRUD
59// =============================================================================
60
61pub async fn list_tenant_agents(pool: &PgPool, tenant_id: &str) -> Result<Vec<TenantAgent>> {
62    let rows = sqlx::query(
63        "SELECT id, tenant_id, agent_name, display_name, description, config, enabled, created_at, updated_at
64         FROM tenant_agents WHERE tenant_id = $1 ORDER BY agent_name"
65    )
66    .bind(tenant_id)
67    .fetch_all(pool)
68    .await
69    .map_err(|e| AppError::Database(e.to_string()))?;
70
71    rows.iter()
72        .map(|row| {
73            Ok(TenantAgent {
74                id: row.get("id"),
75                tenant_id: row.get("tenant_id"),
76                agent_name: row.get("agent_name"),
77                display_name: row.get("display_name"),
78                description: row.get("description"),
79                config: row.get::<serde_json::Value, _>("config"),
80                enabled: row.get("enabled"),
81                created_at: row.get("created_at"),
82                updated_at: row.get("updated_at"),
83            })
84        })
85        .collect()
86}
87
88pub async fn get_tenant_agent(
89    pool: &PgPool,
90    tenant_id: &str,
91    agent_name: &str,
92) -> Result<TenantAgent> {
93    let row = sqlx::query(
94        "SELECT id, tenant_id, agent_name, display_name, description, config, enabled, created_at, updated_at
95         FROM tenant_agents WHERE tenant_id = $1 AND agent_name = $2"
96    )
97    .bind(tenant_id)
98    .bind(agent_name)
99    .fetch_optional(pool)
100    .await
101    .map_err(|e| AppError::Database(e.to_string()))?
102    .ok_or_else(|| AppError::NotFound(format!("Agent '{}' not found for tenant '{}'", agent_name, tenant_id)))?;
103
104    Ok(TenantAgent {
105        id: row.get("id"),
106        tenant_id: row.get("tenant_id"),
107        agent_name: row.get("agent_name"),
108        display_name: row.get("display_name"),
109        description: row.get("description"),
110        config: row.get::<serde_json::Value, _>("config"),
111        enabled: row.get("enabled"),
112        created_at: row.get("created_at"),
113        updated_at: row.get("updated_at"),
114    })
115}
116
117pub async fn create_tenant_agent(
118    pool: &PgPool,
119    tenant_id: &str,
120    req: CreateTenantAgentRequest,
121) -> Result<TenantAgent> {
122    let id = uuid::Uuid::new_v4().to_string();
123    let now = now_ts();
124
125    sqlx::query(
126        "INSERT INTO tenant_agents (id, tenant_id, agent_name, display_name, description, config, enabled, created_at, updated_at)
127         VALUES ($1, $2, $3, $4, $5, $6, true, $7, $7)"
128    )
129    .bind(&id)
130    .bind(tenant_id)
131    .bind(&req.agent_name)
132    .bind(&req.display_name)
133    .bind(&req.description)
134    .bind(&req.config)
135    .bind(now)
136    .execute(pool)
137    .await
138    .map_err(|e| AppError::Database(e.to_string()))?;
139
140    get_tenant_agent(pool, tenant_id, &req.agent_name).await
141}
142
143pub async fn update_tenant_agent(
144    pool: &PgPool,
145    tenant_id: &str,
146    agent_name: &str,
147    req: UpdateTenantAgentRequest,
148) -> Result<TenantAgent> {
149    let now = now_ts();
150
151    // Fetch current state
152    let current = get_tenant_agent(pool, tenant_id, agent_name).await?;
153
154    let display_name = req.display_name.unwrap_or(current.display_name);
155    let description = req.description.or(current.description);
156    let config = req.config.unwrap_or(current.config);
157    let enabled = req.enabled.unwrap_or(current.enabled);
158
159    sqlx::query(
160        "UPDATE tenant_agents SET display_name = $1, description = $2, config = $3, enabled = $4, updated_at = $5
161         WHERE tenant_id = $6 AND agent_name = $7"
162    )
163    .bind(&display_name)
164    .bind(&description)
165    .bind(&config)
166    .bind(enabled)
167    .bind(now)
168    .bind(tenant_id)
169    .bind(agent_name)
170    .execute(pool)
171    .await
172    .map_err(|e| AppError::Database(e.to_string()))?;
173
174    get_tenant_agent(pool, tenant_id, agent_name).await
175}
176
177pub async fn delete_tenant_agent(pool: &PgPool, tenant_id: &str, agent_name: &str) -> Result<()> {
178    let result = sqlx::query("DELETE FROM tenant_agents WHERE tenant_id = $1 AND agent_name = $2")
179        .bind(tenant_id)
180        .bind(agent_name)
181        .execute(pool)
182        .await
183        .map_err(|e| AppError::Database(e.to_string()))?;
184
185    if result.rows_affected() == 0 {
186        return Err(AppError::NotFound(format!(
187            "Agent '{}' not found for tenant '{}'",
188            agent_name, tenant_id
189        )));
190    }
191    Ok(())
192}
193
194// =============================================================================
195// Template operations
196// =============================================================================
197
198pub async fn list_agent_templates(
199    pool: &PgPool,
200    product_type: Option<&str>,
201) -> Result<Vec<AgentTemplate>> {
202    let rows = if let Some(pt) = product_type {
203        sqlx::query(
204            "SELECT id, product_type, agent_name, display_name, description, config, created_at
205             FROM agent_templates WHERE product_type = $1 ORDER BY agent_name",
206        )
207        .bind(pt)
208        .fetch_all(pool)
209        .await
210        .map_err(|e| AppError::Database(e.to_string()))?
211    } else {
212        sqlx::query(
213            "SELECT id, product_type, agent_name, display_name, description, config, created_at
214             FROM agent_templates ORDER BY product_type, agent_name",
215        )
216        .fetch_all(pool)
217        .await
218        .map_err(|e| AppError::Database(e.to_string()))?
219    };
220
221    rows.iter()
222        .map(|row| {
223            Ok(AgentTemplate {
224                id: row.get("id"),
225                product_type: row.get("product_type"),
226                agent_name: row.get("agent_name"),
227                display_name: row.get("display_name"),
228                description: row.get("description"),
229                config: row.get::<serde_json::Value, _>("config"),
230                created_at: row.get("created_at"),
231            })
232        })
233        .collect()
234}
235
236/// Clones all agent templates for a product type into a tenant's agent list.
237/// Idempotent — skips agents that already exist (ON CONFLICT DO NOTHING).
238pub async fn clone_templates_for_tenant(
239    pool: &PgPool,
240    tenant_id: &str,
241    product_type: &str,
242) -> Result<Vec<TenantAgent>> {
243    let templates = list_agent_templates(pool, Some(product_type)).await?;
244    let now = now_ts();
245
246    for tpl in &templates {
247        let id = uuid::Uuid::new_v4().to_string();
248        sqlx::query(
249            "INSERT INTO tenant_agents (id, tenant_id, agent_name, display_name, description, config, enabled, created_at, updated_at)
250             VALUES ($1, $2, $3, $4, $5, $6, true, $7, $7)
251             ON CONFLICT (tenant_id, agent_name) DO NOTHING"
252        )
253        .bind(&id)
254        .bind(tenant_id)
255        .bind(&tpl.agent_name)
256        .bind(&tpl.display_name)
257        .bind(&tpl.description)
258        .bind(&tpl.config)
259        .bind(now)
260        .execute(pool)
261        .await
262        .map_err(|e| AppError::Database(e.to_string()))?;
263    }
264
265    list_tenant_agents(pool, tenant_id).await
266}
267
268// =============================================================================
269// Seed default templates
270// =============================================================================
271
272/// Seeds default agent templates. Idempotent — uses ON CONFLICT DO NOTHING.
273/// Called once on ARES startup after migrations.
274pub async fn seed_default_templates(pool: &PgPool) -> Result<()> {
275    let now = now_ts();
276
277    struct TemplateSpec {
278        product_type: &'static str,
279        agent_name: &'static str,
280        display_name: &'static str,
281        description: &'static str,
282        model: &'static str,
283        system_prompt: &'static str,
284    }
285
286    let templates: &[TemplateSpec] = &[
287        // Generic
288        TemplateSpec {
289            product_type: "generic",
290            agent_name: "assistant",
291            display_name: "General Assistant",
292            description: "Default conversational agent",
293            model: "fast",
294            system_prompt: "You are a helpful AI assistant. Answer questions clearly and concisely. If you don't know something, say so. Be direct and useful.",
295        },
296        // Client-specific agent templates are loaded by the managed platform crate
297        // from TOON config files, not hardcoded in the OSS layer.
298    ];
299
300    for tpl in templates {
301        let id = uuid::Uuid::new_v4().to_string();
302        let config = serde_json::json!({
303            "model": tpl.model,
304            "system_prompt": tpl.system_prompt,
305            "tools": [],
306            "max_tool_iterations": 3
307        });
308
309        sqlx::query(
310            "INSERT INTO agent_templates (id, product_type, agent_name, display_name, description, config, created_at)
311             VALUES ($1, $2, $3, $4, $5, $6, $7)
312             ON CONFLICT (product_type, agent_name) DO NOTHING"
313        )
314        .bind(&id)
315        .bind(tpl.product_type)
316        .bind(tpl.agent_name)
317        .bind(tpl.display_name)
318        .bind(tpl.description)
319        .bind(&config)
320        .bind(now)
321        .execute(pool)
322        .await
323        .map_err(|e| AppError::Database(format!("Failed to seed template {}/{}: {}", tpl.product_type, tpl.agent_name, e)))?;
324    }
325
326    tracing::info!("Agent templates seeded ({} templates)", templates.len());
327    Ok(())
328}