Skip to content

Commit 39c4dba

Browse files
feat: Add template validation system with .valid file
- Modify bootstrap to read .valid file and validate templates before loading - Templates not in .valid file are skipped during bootstrap - Backward compatible: if .valid file missing, all templates are loaded - Enables controlled template loading during bootstrap
1 parent 748fcef commit 39c4dba

File tree

1 file changed

+130
-3
lines changed

1 file changed

+130
-3
lines changed

src/core/bootstrap/mod.rs

Lines changed: 130 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ use crate::package_manager::{InstallMode, PackageManager};
44
use crate::security::command_guard::SafeCommand;
55
use crate::shared::utils::{establish_pg_connection, init_secrets_manager};
66
use anyhow::Result;
7+
use uuid::Uuid;
78

89
#[cfg(feature = "drive")]
910
use aws_sdk_s3::Client;
@@ -18,6 +19,13 @@ use std::fs;
1819
use std::os::unix::fs::PermissionsExt;
1920
use std::path::{Path, PathBuf};
2021

22+
#[derive(diesel::QueryableByName)]
23+
#[diesel(check_for_backend(diesel::pg::Pg))]
24+
struct BotExistsResult {
25+
#[diesel(sql_type = diesel::sql_types::Bool)]
26+
exists: bool,
27+
}
28+
2129
fn safe_pkill(args: &[&str]) {
2230
if let Ok(cmd) = SafeCommand::new("pkill").and_then(|c| c.args(args)) {
2331
let _ = cmd.execute();
@@ -1971,6 +1979,68 @@ VAULT_CACHE_TTL=300
19711979
debug!("Drive feature disabled, skipping template upload");
19721980
Ok(())
19731981
}
1982+
fn create_bot_from_template(conn: &mut diesel::PgConnection, bot_name: &str) -> Result<Uuid> {
1983+
use diesel::sql_query;
1984+
1985+
info!("Creating bot '{}' from template", bot_name);
1986+
1987+
let bot_id = Uuid::new_v4();
1988+
let db_name = format!("bot_{}", bot_name.replace(['-', ' '], "_").to_lowercase());
1989+
1990+
sql_query(
1991+
"INSERT INTO bots (id, name, description, is_active, database_name, created_at, updated_at, llm_provider, llm_config, context_provider, context_config)
1992+
VALUES ($1, $2, $3, true, $4, NOW(), NOW(), $5, $6, $7, $8)",
1993+
)
1994+
.bind::<diesel::sql_types::Uuid, _>(bot_id)
1995+
.bind::<diesel::sql_types::Text, _>(bot_name)
1996+
.bind::<diesel::sql_types::Text, _>(format!("Bot agent: {}", bot_name))
1997+
.bind::<diesel::sql_types::Text, _>(&db_name)
1998+
.bind::<diesel::sql_types::Text, _>("local")
1999+
.bind::<diesel::sql_types::Json, _>(serde_json::json!({}))
2000+
.bind::<diesel::sql_types::Text, _>("postgres")
2001+
.bind::<diesel::sql_types::Json, _>(serde_json::json!({}))
2002+
.execute(conn)
2003+
.map_err(|e| anyhow::anyhow!("Failed to create bot '{}': {}", bot_name, e))?;
2004+
2005+
// Create the bot database
2006+
let safe_db_name: String = db_name
2007+
.chars()
2008+
.filter(|c| c.is_alphanumeric() || *c == '_')
2009+
.collect();
2010+
2011+
if !safe_db_name.is_empty() && safe_db_name.len() <= 63 {
2012+
let create_query = format!("CREATE DATABASE {}", safe_db_name);
2013+
if let Err(e) = sql_query(&create_query).execute(conn) {
2014+
let err_str = e.to_string();
2015+
if !err_str.contains("already exists") {
2016+
warn!("Failed to create database for bot '{}': {}", bot_name, e);
2017+
}
2018+
}
2019+
info!("Created database '{}' for bot '{}'", safe_db_name, bot_name);
2020+
}
2021+
2022+
Ok(bot_id)
2023+
}
2024+
2025+
fn read_valid_templates(templates_dir: &Path) -> std::collections::HashSet<String> {
2026+
let valid_file = templates_dir.join(".valid");
2027+
let mut valid_set = std::collections::HashSet::new();
2028+
2029+
if let Ok(content) = std::fs::read_to_string(&valid_file) {
2030+
for line in content.lines() {
2031+
let line = line.trim();
2032+
if !line.is_empty() && !line.starts_with('#') {
2033+
valid_set.insert(line.to_string());
2034+
}
2035+
}
2036+
info!("Loaded {} valid templates from .valid file", valid_set.len());
2037+
} else {
2038+
info!("No .valid file found, will load all templates");
2039+
}
2040+
2041+
valid_set
2042+
}
2043+
19742044
fn create_bots_from_templates(conn: &mut diesel::PgConnection) -> Result<()> {
19752045
use crate::shared::models::schema::bots;
19762046
use diesel::prelude::*;
@@ -2001,22 +2071,79 @@ VAULT_CACHE_TTL=300
20012071
}
20022072
};
20032073

2074+
let valid_templates = Self::read_valid_templates(&templates_dir);
2075+
let load_all = valid_templates.is_empty();
2076+
20042077
let default_bot: Option<(uuid::Uuid, String)> = bots::table
20052078
.filter(bots::is_active.eq(true))
20062079
.select((bots::id, bots::name))
20072080
.first(conn)
20082081
.optional()?;
20092082

2010-
let Some((default_bot_id, default_bot_name)) = default_bot else {
2011-
error!("No active bot found in database - cannot sync template configs");
2012-
return Ok(());
2083+
let (default_bot_id, default_bot_name) = match default_bot {
2084+
Some(bot) => bot,
2085+
None => {
2086+
// Create default bot if it doesn't exist
2087+
info!("No active bot found, creating 'default' bot from template");
2088+
let bot_id = Self::create_bot_from_template(conn, "default")?;
2089+
(bot_id, "default".to_string())
2090+
}
20132091
};
20142092

20152093
info!(
20162094
"Syncing template configs to bot '{}' ({})",
20172095
default_bot_name, default_bot_id
20182096
);
20192097

2098+
// Scan for .gbai template files and create bots if they don't exist
2099+
let entries = std::fs::read_dir(&templates_dir)
2100+
.map_err(|e| anyhow::anyhow!("Failed to read templates directory: {}", e))?;
2101+
2102+
for entry in entries.flatten() {
2103+
let file_name = entry.file_name();
2104+
let file_name_str = match file_name.to_str() {
2105+
Some(name) => name,
2106+
None => continue,
2107+
};
2108+
2109+
if !file_name_str.ends_with(".gbai") {
2110+
continue;
2111+
}
2112+
2113+
if !load_all && !valid_templates.contains(file_name_str) {
2114+
debug!("Skipping template '{}' (not in .valid file)", file_name_str);
2115+
continue;
2116+
}
2117+
2118+
let bot_name = file_name_str.trim_end_matches(".gbai");
2119+
2120+
// Check if bot already exists
2121+
let bot_exists: bool =
2122+
diesel::sql_query("SELECT EXISTS(SELECT 1 FROM bots WHERE name = $1) as exists")
2123+
.bind::<diesel::sql_types::Text, _>(bot_name)
2124+
.get_result::<BotExistsResult>(conn)
2125+
.map(|r| r.exists)
2126+
.unwrap_or(false);
2127+
2128+
if bot_exists {
2129+
info!("Bot '{}' already exists, skipping creation", bot_name);
2130+
continue;
2131+
}
2132+
2133+
// Create bot from template
2134+
match Self::create_bot_from_template(conn, bot_name) {
2135+
Ok(bot_id) => {
2136+
info!(
2137+
"Successfully created bot '{}' ({}) from template",
2138+
bot_name, bot_id
2139+
);
2140+
}
2141+
Err(e) => {
2142+
error!("Failed to create bot '{}' from template: {:#}", bot_name, e);
2143+
}
2144+
}
2145+
}
2146+
20202147
let default_template = templates_dir.join("default.gbai");
20212148
info!(
20222149
"Looking for default template at: {}",

0 commit comments

Comments
 (0)