diff --git a/engine/cli/src/commands.rs b/engine/cli/src/commands.rs index a38e4fd878..155273bfad 100644 --- a/engine/cli/src/commands.rs +++ b/engine/cli/src/commands.rs @@ -2,6 +2,25 @@ use anyhow::Result; use baml_runtime::{cli::RuntimeCliDefaults, BamlRuntime}; use clap::{Parser, Subcommand}; +/// Environment variable name to enable internal CLI commands +pub(crate) const INTERNAL_CLI_ENV_VAR: &str = "BAML_INTERNAL_CLI"; + +/// Check if internal CLI commands are enabled via environment variable +fn is_internal_cli_enabled() -> bool { + std::env::var(INTERNAL_CLI_ENV_VAR).is_ok() +} + +/// Check if internal CLI commands are enabled, returning an error if not +fn require_internal_cli_enabled() -> Result<()> { + if !is_internal_cli_enabled() { + anyhow::bail!( + "This is an internal command. Set the {} environment variable to enable it.", + INTERNAL_CLI_ENV_VAR + ); + } + Ok(()) +} + #[derive(Parser, Debug)] #[command(author, version, about = "A CLI tool for working with BAML. Learn more at https://docs.boundaryml.com.", long_about = None)] #[command(styles = clap_cargo::style::CLAP_STYLING)] @@ -201,6 +220,7 @@ impl RuntimeCli { } } Commands::DumpHIR(args) => { + require_internal_cli_enabled()?; args.from = BamlRuntime::parse_baml_src_path(&args.from)?; match args.run( baml_runtime::cli::dump_intermediate::DumpType::HIR, @@ -214,6 +234,7 @@ impl RuntimeCli { } } Commands::DumpBytecode(args) => { + require_internal_cli_enabled()?; args.from = BamlRuntime::parse_baml_src_path(&args.from)?; match args.run( baml_runtime::cli::dump_intermediate::DumpType::Bytecode, @@ -226,17 +247,23 @@ impl RuntimeCli { } } } - Commands::LanguageServer(args) => match args.run() { - Ok(()) => Ok(crate::ExitCode::Success), - Err(_) => Ok(crate::ExitCode::Other), - }, - Commands::Repl(args) => match t.block_on(async { args.run().await }) { - Ok(()) => Ok(crate::ExitCode::Success), - Err(e) => { - eprintln!("Error: {e}"); - Ok(crate::ExitCode::Other) + Commands::LanguageServer(args) => { + require_internal_cli_enabled()?; + match args.run() { + Ok(()) => Ok(crate::ExitCode::Success), + Err(_) => Ok(crate::ExitCode::Other), } - }, + } + Commands::Repl(args) => { + require_internal_cli_enabled()?; + match t.block_on(async { args.run().await }) { + Ok(()) => Ok(crate::ExitCode::Success), + Err(e) => { + eprintln!("Error: {e}"); + Ok(crate::ExitCode::Other) + } + } + } } } } diff --git a/engine/cli/src/lib.rs b/engine/cli/src/lib.rs index 185d5760ba..17645cf4b2 100644 --- a/engine/cli/src/lib.rs +++ b/engine/cli/src/lib.rs @@ -8,7 +8,6 @@ pub(crate) mod lsp; pub(crate) mod propelauth; pub(crate) mod tui; use anyhow::Result; -use clap::Parser; #[derive(Debug, Clone)] pub enum ExitCode { @@ -63,7 +62,24 @@ pub fn run_cli( argv: Vec, caller_type: baml_runtime::RuntimeCliDefaults, ) -> Result { - let mut cli = commands::RuntimeCli::parse_from(argv); + use clap::{CommandFactory, FromArgMatches}; + + // Build the command and conditionally hide internal subcommands + let mut cmd = commands::RuntimeCli::command(); + + // If the internal CLI env var is not set, hide the internal commands + if std::env::var(commands::INTERNAL_CLI_ENV_VAR).is_err() { + cmd = cmd + .mut_subcommand("dump-hir", |subcmd| subcmd.hide(true)) + .mut_subcommand("dump-bytecode", |subcmd| subcmd.hide(true)) + .mut_subcommand("lsp", |subcmd| subcmd.hide(true)) + .mut_subcommand("repl", |subcmd| subcmd.hide(true)); + } + + // Parse the arguments with our modified command + let matches = cmd.get_matches_from(argv); + let mut cli = commands::RuntimeCli::from_arg_matches(&matches)?; + if !matches!(cli.command, commands::Commands::Test(_)) { // We only need to set the exit handlers if we're not running tests // and the caller is Python.