Merge pull request #342 from buster-so/feature-make-recurions-configurable

feat: add support for custom environment variables in CLI commands

.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"

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::<u32>().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(

api/server/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "buster_server"
-version = "0.1.36"
+version = "0.2.0"
 edition = "2021"
 default-run = "buster_server"
 

cli/cli/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "buster-cli"
-version = "0.1.36"
+version = "0.2.0"
 edition = "2021"
 build = "build.rs"
 

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<String> = Vec::new();
+    let mut updated_vars: HashSet<String> = 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")]

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<PathBuf, BusterError> {
+async fn setup_persistent_app_environment(env_vars: Option<&[(String, String)]>) -> Result<PathBuf, BusterError> {
     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<PathBuf, BusterError> {
         ))
     })?;
 
+    // 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
 }

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<String>,
     },
     /// 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<Vec<(String, String)>, _> = 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),
     };

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(())
+} 

tag_info.json

@@ -1,7 +1,7 @@
 {
-  "api_tag": "api/v0.1.36", "api_version": "0.1.36"
+  "api_tag": "api/v0.2.0", "api_version": "0.2.0"
 ,
-  "web_tag": "web/v0.1.36", "web_version": "0.1.36"
+  "web_tag": "web/v0.2.0", "web_version": "0.2.0"
 ,
-  "cli_tag": "cli/v0.1.36", "cli_version": "0.1.36"
+  "cli_tag": "cli/v0.2.0", "cli_version": "0.2.0"
 }

web/package-lock.json

@@ -1,12 +1,12 @@
 {
   "name": "web",
-  "version": "0.1.36",
+  "version": "0.2.0",
   "lockfileVersion": 3,
   "requires": true,
   "packages": {
     "": {
       "name": "web",
-      "version": "0.1.36",
+      "version": "0.2.0",
       "dependencies": {
         "@dnd-kit/core": "^6.3.1",
         "@dnd-kit/modifiers": "^9.0.0",

web/package.json

@@ -1,6 +1,6 @@
 {
   "name": "web",
-  "version": "0.1.36",
+  "version": "0.2.0",
   "private": true,
   "scripts": {
     "dev": "next dev --turbo",