diff --git a/.env.example b/.env.example index 2c56497e7..214b9151a 100644 --- a/.env.example +++ b/.env.example @@ -11,6 +11,7 @@ SUPABASE_URL="http://kong:8000" SUPABASE_SERVICE_ROLE_KEY="eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.ey AgCiAgICAicm9sZSI6ICJzZXJ2aWNlX3JvbGUiLAogICAgImlzcyI6ICJzdXBhYmFzZS1kZW1vIiwKICAgICJpYXQiOiAxNjQxNzY5MjAwLAogICAgImV4cCI6IDE3OTk1MzU2MDAKfQ.DaYlNEoUrrEn2Ig7tqibS-PHK5vgusbcbo7X36XVt4Q" POSTHOG_TELEMETRY_KEY="phc_zZraCicSTfeXX5b9wWQv2rWG8QB4Z3xlotOT7gFtoNi" TELEMETRY_ENABLED="true" +MAX_RECURSION="15" # AI VARS RERANK_API_KEY="your_rerank_api_key" diff --git a/api/libs/agents/src/agent.rs b/api/libs/agents/src/agent.rs index b09d2f767..905897509 100644 --- a/api/libs/agents/src/agent.rs +++ b/api/libs/agents/src/agent.rs @@ -554,8 +554,13 @@ impl Agent { *current = Some(thread_ref.clone()); } + let max_recursion = std::env::var("MAX_RECURSION") + .ok() + .and_then(|v| v.parse::().ok()) + .unwrap_or(15); + // Limit recursion to a maximum of 15 times - if recursion_depth >= 15 { + if recursion_depth >= max_recursion { let max_depth_msg = format!("Maximum recursion depth ({}) reached.", recursion_depth); warn!("{}", max_depth_msg); let message = AgentMessage::assistant( diff --git a/cli/cli/src/commands/config_utils.rs b/cli/cli/src/commands/config_utils.rs index 9150ab67c..3a21bfb66 100644 --- a/cli/cli/src/commands/config_utils.rs +++ b/cli/cli/src/commands/config_utils.rs @@ -2,9 +2,11 @@ use crate::error::BusterError; use dirs; use serde::{Deserialize, Serialize}; use serde_yaml; +use std::collections::{HashMap, HashSet}; use std::fs; use std::io::{self, Write}; use std::path::{Path, PathBuf}; +use std::str::FromStr; use colored::*; use inquire::{Confirm, Password, Select, Text, PasswordDisplayMode, validator::Validation}; @@ -179,6 +181,60 @@ pub fn update_env_file( }) } +pub fn update_arbitrary_env_vars( + target_dotenv_path: &Path, + env_vars: &[(String, String)], +) -> Result<(), BusterError> { + let mut new_env_lines: Vec = Vec::new(); + let mut updated_vars: HashSet = HashSet::new(); + + // Read existing .env file if it exists + if target_dotenv_path.exists() { + let env_content = fs::read_to_string(target_dotenv_path).map_err(|e| { + BusterError::CommandError(format!( + "Failed to read .env file at {}: {}", + target_dotenv_path.display(), + e + )) + })?; + + for line in env_content.lines() { + let mut line_replaced = false; + + // Check if this line starts with any of our env vars + for (key, value) in env_vars { + if line.starts_with(&format!("{}=", key)) { + new_env_lines.push(format!("{}=\"{}\"", key, value)); + updated_vars.insert(key.clone()); + line_replaced = true; + break; + } + } + + // If no replacement was made, keep the original line + if !line_replaced { + new_env_lines.push(line.to_string()); + } + } + } + + // Add any environment variables that weren't found in the existing file + for (key, value) in env_vars { + if !updated_vars.contains(key) { + new_env_lines.push(format!("{}=\"{}\"", key, value)); + } + } + + // Write the updated content back to the file + fs::write(target_dotenv_path, new_env_lines.join("\n")).map_err(|e| { + BusterError::CommandError(format!( + "Failed to write updated .env file to {}: {}", + target_dotenv_path.display(), + e + )) + }) +} + #[derive(Debug, Deserialize, Serialize, Clone, Default)] pub struct ModelInfo { #[serde(skip_serializing_if = "Option::is_none")] diff --git a/cli/cli/src/commands/run.rs b/cli/cli/src/commands/run.rs index bc72be323..9f101efc0 100644 --- a/cli/cli/src/commands/run.rs +++ b/cli/cli/src/commands/run.rs @@ -23,7 +23,7 @@ use colored::*; #[exclude = "supabase/docker-compose.override.yml"] struct StaticAssets; -async fn setup_persistent_app_environment() -> Result { +async fn setup_persistent_app_environment(env_vars: Option<&[(String, String)]>) -> Result { let app_base_dir = config_utils::get_app_base_dir().map_err(|e| { BusterError::CommandError(format!("Failed to get app base directory: {}", e)) })?; @@ -172,6 +172,25 @@ async fn setup_persistent_app_environment() -> Result { )) })?; + // Update arbitrary environment variables if provided + if let Some(env_vars) = env_vars { + if !env_vars.is_empty() { + println!("Updating custom environment variables..."); + config_utils::update_arbitrary_env_vars(&main_dot_env_path, env_vars) + .map_err(|e| { + BusterError::CommandError(format!( + "Failed to update custom environment variables in {}: {}", + main_dot_env_path.display(), + e + )) + })?; + + for (key, value) in env_vars { + println!("Set {}={}", key, if key.to_lowercase().contains("key") || key.to_lowercase().contains("password") || key.to_lowercase().contains("secret") { "***" } else { value }); + } + } + } + println!("--- Configuration Setup/Check Complete ---"); // --- END Configuration Checks/Updates --- @@ -371,9 +390,10 @@ async fn run_docker_compose_command( } } -pub async fn start(no_track: bool) -> Result<(), BusterError> { +pub async fn start(no_track: bool, env_vars: Vec<(String, String)>) -> Result<(), BusterError> { // First, run the setup/check which includes printing headers/prompts if needed - let app_base_dir = setup_persistent_app_environment().await?; + let env_vars_option = if env_vars.is_empty() { None } else { Some(env_vars.as_slice()) }; + let app_base_dir = setup_persistent_app_environment(env_vars_option).await?; // Then, run the docker command in that directory run_docker_compose_command(&app_base_dir, &["up", "-d"], "Starting", no_track).await } diff --git a/cli/cli/src/main.rs b/cli/cli/src/main.rs index 4dea56478..87345e3c9 100644 --- a/cli/cli/src/main.rs +++ b/cli/cli/src/main.rs @@ -3,10 +3,10 @@ mod error; mod types; mod utils; -use anyhow; use clap::{Parser, Subcommand}; use commands::{auth::check_authentication, auth::AuthArgs, init, run}; use utils::updater::check_for_updates; +use anyhow; pub const APP_NAME: &str = "buster"; pub const VERSION: &str = env!("CARGO_PKG_VERSION"); @@ -95,6 +95,10 @@ pub enum Commands { /// Disable telemetry tracking #[arg(long, default_value_t = false)] no_track: bool, + /// Set environment variables (can be used multiple times) + /// Format: KEY=VALUE + #[arg(long = "env", action = clap::ArgAction::Append)] + env_vars: Vec, }, /// Stop the Buster services Stop, @@ -161,7 +165,22 @@ async fn main() { } => commands::generate::generate_semantic_models_command(path, target_semantic_file).await, Commands::Parse { path } => commands::parse::parse_models_command(path).await, Commands::Config => commands::config::manage_settings_interactive().await.map_err(anyhow::Error::from), - Commands::Start { no_track } => run::start(no_track).await.map_err(anyhow::Error::from), + Commands::Start { no_track, env_vars } => { + // Parse env vars from KEY=VALUE format + let parsed_env_vars: Result, _> = env_vars + .iter() + .map(|env_str| { + env_str.split_once('=') + .map(|(k, v)| (k.to_string(), v.to_string())) + .ok_or_else(|| anyhow::anyhow!("Invalid env var format '{}'. Expected KEY=VALUE", env_str)) + }) + .collect(); + + match parsed_env_vars { + Ok(env_pairs) => run::start(no_track, env_pairs).await.map_err(anyhow::Error::from), + Err(e) => Err(e), + } + } Commands::Stop => run::stop().await.map_err(anyhow::Error::from), Commands::Reset => run::reset().await.map_err(anyhow::Error::from), }; diff --git a/cli/tests/env_vars_test.rs b/cli/tests/env_vars_test.rs new file mode 100644 index 000000000..c12008f35 --- /dev/null +++ b/cli/tests/env_vars_test.rs @@ -0,0 +1,159 @@ +use anyhow::Result; +use std::fs; +use std::path::Path; +use tempfile::TempDir; + +// Import from the buster CLI library +use buster_cli::commands::config_utils; + +/// Test that update_arbitrary_env_vars correctly updates environment variables +#[tokio::test] +async fn test_update_arbitrary_env_vars_new_file() -> Result<()> { + let temp_dir = TempDir::new()?; + let env_file_path = temp_dir.path().join(".env"); + + // Test data + let env_vars = vec![ + ("DATABASE_URL".to_string(), "postgres://localhost:5432/testdb".to_string()), + ("API_KEY".to_string(), "secret123".to_string()), + ("DEBUG".to_string(), "true".to_string()), + ]; + + // Update env vars (should create new file) + config_utils::update_arbitrary_env_vars(&env_file_path, &env_vars)?; + + // Verify file was created and contains correct content + assert!(env_file_path.exists()); + let content = fs::read_to_string(&env_file_path)?; + + assert!(content.contains("DATABASE_URL=\"postgres://localhost:5432/testdb\"")); + assert!(content.contains("API_KEY=\"secret123\"")); + assert!(content.contains("DEBUG=\"true\"")); + + Ok(()) +} + +#[tokio::test] +async fn test_update_arbitrary_env_vars_existing_file() -> Result<()> { + let temp_dir = TempDir::new()?; + let env_file_path = temp_dir.path().join(".env"); + + // Create initial .env file + let initial_content = r#" +# Existing configuration +EXISTING_VAR="keep_this" +OLD_API_KEY="old_secret" +DATABASE_URL="postgres://localhost:5432/olddb" +"#; + fs::write(&env_file_path, initial_content)?; + + // Test data - some new, some updating existing + let env_vars = vec![ + ("DATABASE_URL".to_string(), "postgres://localhost:5432/newdb".to_string()), + ("NEW_VAR".to_string(), "new_value".to_string()), + ("API_SECRET".to_string(), "super_secret".to_string()), + ]; + + // Update env vars + config_utils::update_arbitrary_env_vars(&env_file_path, &env_vars)?; + + // Verify content + let content = fs::read_to_string(&env_file_path)?; + + // Should preserve existing vars not being updated + assert!(content.contains("EXISTING_VAR=\"keep_this\"")); + assert!(content.contains("OLD_API_KEY=\"old_secret\"")); + + // Should update existing var + assert!(content.contains("DATABASE_URL=\"postgres://localhost:5432/newdb\"")); + assert!(!content.contains("postgres://localhost:5432/olddb")); + + // Should add new vars + assert!(content.contains("NEW_VAR=\"new_value\"")); + assert!(content.contains("API_SECRET=\"super_secret\"")); + + Ok(()) +} + +#[tokio::test] +async fn test_update_arbitrary_env_vars_preserves_comments() -> Result<()> { + let temp_dir = TempDir::new()?; + let env_file_path = temp_dir.path().join(".env"); + + // Create initial .env file with comments + let initial_content = r#"# Database configuration +DATABASE_URL="postgres://localhost:5432/olddb" + +# API configuration +API_KEY="old_key" +# This is a comment that should be preserved +KEEP_THIS="unchanged" +"#; + fs::write(&env_file_path, initial_content)?; + + // Update only one variable + let env_vars = vec![ + ("API_KEY".to_string(), "new_key".to_string()), + ]; + + config_utils::update_arbitrary_env_vars(&env_file_path, &env_vars)?; + + let content = fs::read_to_string(&env_file_path)?; + + // Should preserve comments and other variables + assert!(content.contains("# Database configuration")); + assert!(content.contains("# API configuration")); + assert!(content.contains("# This is a comment that should be preserved")); + assert!(content.contains("DATABASE_URL=\"postgres://localhost:5432/olddb\"")); + assert!(content.contains("KEEP_THIS=\"unchanged\"")); + + // Should update the target variable + assert!(content.contains("API_KEY=\"new_key\"")); + assert!(!content.contains("API_KEY=\"old_key\"")); + + Ok(()) +} + +#[tokio::test] +async fn test_update_arbitrary_env_vars_empty_list() -> Result<()> { + let temp_dir = TempDir::new()?; + let env_file_path = temp_dir.path().join(".env"); + + // Create initial .env file + let initial_content = "EXISTING_VAR=\"value\"\n"; + fs::write(&env_file_path, initial_content)?; + + // Update with empty list + let env_vars = vec![]; + config_utils::update_arbitrary_env_vars(&env_file_path, &env_vars)?; + + // Should preserve existing content + let content = fs::read_to_string(&env_file_path)?; + assert_eq!(content, initial_content); + + Ok(()) +} + +#[tokio::test] +async fn test_update_arbitrary_env_vars_handles_special_characters() -> Result<()> { + let temp_dir = TempDir::new()?; + let env_file_path = temp_dir.path().join(".env"); + + // Test data with special characters + let env_vars = vec![ + ("URL_WITH_PARAMS".to_string(), "http://localhost:3000/api?key=value&other=123".to_string()), + ("PASSWORD_WITH_SYMBOLS".to_string(), "p@ssw0rd!#$%^&*()".to_string()), + ("JSON_CONFIG".to_string(), "{\"key\":\"value\",\"number\":42}".to_string()), + ]; + + config_utils::update_arbitrary_env_vars(&env_file_path, &env_vars)?; + + let content = fs::read_to_string(&env_file_path)?; + + // Verify special characters are preserved + assert!(content.contains("URL_WITH_PARAMS=\"http://localhost:3000/api?key=value&other=123\"")); + assert!(content.contains("PASSWORD_WITH_SYMBOLS=\"p@ssw0rd!#$%^&*()\"")); + assert!(content.contains("JSON_CONFIG=\"{\"key\":\"value\",\"number\":42}\"")); + + Ok(()) +} \ No newline at end of file