feat: add support for custom environment variables in CLI commands

- Introduced `env_vars` option in the `start` command to allow users to set environment variables in KEY=VALUE format.
- Implemented `update_arbitrary_env_vars` function to update the .env file with provided variables.
- Adjusted `setup_persistent_app_environment` to handle custom environment variables.
- Added `MAX_RECURSION` variable to .env.example and updated recursion limit logic in the agent implementation.
This commit is contained in:
dal 2025-05-31 07:15:50 -06:00
parent 86c275e70d
commit e9a269983b
No known key found for this signature in database
GPG Key ID: 16F4B0E1E9F61122
6 changed files with 266 additions and 6 deletions

View File

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

View File

@ -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(

View File

@ -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")]

View File

@ -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
}

View File

@ -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),
};

159
cli/tests/env_vars_test.rs Normal file
View File

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