From 2063ef3a6cdb16b761e478a9c34d1f6f7168d32d Mon Sep 17 00:00:00 2001 From: Ken Jiang Date: Wed, 19 Aug 2026 10:41:35 -0400 Subject: [PATCH 1/3] add gateway cli option for bt --- src/gateway.rs | 237 +++++++++++++++++++++++++++++++++++++++++ src/main.rs | 6 ++ src/setup/mod.rs | 7 ++ src/utils/fs_atomic.rs | 6 ++ src/utils/mod.rs | 1 + 5 files changed, 257 insertions(+) create mode 100644 src/gateway.rs diff --git a/src/gateway.rs b/src/gateway.rs new file mode 100644 index 00000000..7e0cf04c --- /dev/null +++ b/src/gateway.rs @@ -0,0 +1,237 @@ +//! Configure coding agents to send model requests through Braintrust Gateway. + +use std::fs; +use std::path::Path; + +use anyhow::{anyhow, Context, Result}; +use clap::{Args, Subcommand, ValueEnum}; +use serde_json::{Map, Value}; +use toml::Value as TomlValue; + +use crate::args::{BaseArgs, LoginBaseArgs}; +use crate::setup; +use crate::utils::{write_json_atomic_private, write_text_atomic_private}; + +const DEFAULT_GATEWAY_URL: &str = "https://gateway.braintrust.dev"; + +#[derive(Debug, Clone, Args)] +pub struct GatewayArgs { + #[command(subcommand)] + command: GatewayCommand, +} + +/// Gateway setup has no project routing option. Project selection belongs to +/// individual Gateway requests, not persistent agent credentials. +#[derive(Debug, Clone, Args)] +pub struct GatewayBaseArgs { + #[command(flatten)] + pub(crate) login: LoginBaseArgs, + + /// Override active org + #[arg(short = 'o', long = "org", env = "BRAINTRUST_ORG_NAME", global = true)] + pub(crate) org_name: Option, +} + +#[derive(Debug, Clone, Subcommand)] +enum GatewayCommand { + /// Configure a coding agent to use Braintrust Gateway + Setup(GatewaySetupArgs), +} + +#[derive(Debug, Clone, Args)] +struct GatewaySetupArgs { + #[arg(value_enum)] + agent: GatewayAgent, + + /// Gateway URL to use for model requests + #[arg(long, env = "BRAINTRUST_GATEWAY_URL", default_value = DEFAULT_GATEWAY_URL)] + gateway_url: String, +} + +#[derive(Debug, Clone, Copy, ValueEnum)] +enum GatewayAgent { + /// Configure Claude Code + Claude, + /// Configure Codex + Codex, +} + +pub async fn run(base: GatewayBaseArgs, args: GatewayArgs) -> Result<()> { + let GatewayCommand::Setup(args) = args.command; + let mut auth_base = BaseArgs { + login: base.login, + org_name: base.org_name, + project: None, + }; + let api_key = setup::durable_setup_api_key(&mut auth_base).await?; + let home = dirs::home_dir().ok_or_else(|| anyhow!("failed to resolve HOME/USERPROFILE"))?; + let gateway_url = normalize_gateway_url(&args.gateway_url)?; + + let path = match args.agent { + GatewayAgent::Claude => { + let path = home.join(".claude/settings.json"); + configure_claude(&path, &gateway_url, &api_key)?; + path + } + GatewayAgent::Codex => { + let path = home.join(".codex/config.toml"); + configure_codex(&path, &gateway_url, &api_key)?; + path + } + }; + + if auth_base.json { + println!( + "{}", + serde_json::to_string_pretty(&serde_json::json!({ + "agent": match args.agent { GatewayAgent::Claude => "claude", GatewayAgent::Codex => "codex" }, + "gateway_url": gateway_url, + "settings_path": path, + }))? + ); + } else if auth_base.verbose { + eprintln!( + "Configured {} to use Braintrust Gateway ({})", + args.agent.display_name(), + path.display() + ); + } + + Ok(()) +} + +impl GatewayAgent { + fn display_name(self) -> &'static str { + match self { + Self::Claude => "Claude Code", + Self::Codex => "Codex", + } + } +} + +fn normalize_gateway_url(gateway_url: &str) -> Result { + let gateway_url = gateway_url.trim().trim_end_matches('/'); + if gateway_url.is_empty() { + return Err(anyhow!("--gateway-url cannot be empty")); + } + let parsed = reqwest::Url::parse(gateway_url) + .with_context(|| format!("invalid --gateway-url '{gateway_url}'"))?; + if !matches!(parsed.scheme(), "http" | "https") { + return Err(anyhow!("--gateway-url must use http or https")); + } + Ok(gateway_url.to_string()) +} + +fn configure_claude(path: &Path, gateway_url: &str, api_key: &str) -> Result<()> { + let mut root = load_json_object(path)?; + let env = root + .entry("env".to_string()) + .or_insert_with(|| Value::Object(Map::new())) + .as_object_mut() + .ok_or_else(|| anyhow!("field 'env' in {} must be a JSON object", path.display()))?; + env.insert( + "ANTHROPIC_BASE_URL".to_string(), + Value::String(gateway_url.to_string()), + ); + env.insert( + "ANTHROPIC_AUTH_TOKEN".to_string(), + Value::String(api_key.to_string()), + ); + write_json_atomic_private(path, &root) +} + +fn configure_codex(path: &Path, gateway_url: &str, api_key: &str) -> Result<()> { + let mut root = load_toml_table(path)?; + let env = root + .entry("env".to_string()) + .or_insert_with(|| TomlValue::Table(toml::map::Map::new())) + .as_table_mut() + .ok_or_else(|| anyhow!("field 'env' in {} must be a TOML table", path.display()))?; + env.insert( + "OPENAI_BASE_URL".to_string(), + TomlValue::String(gateway_url.to_string()), + ); + env.insert( + "OPENAI_API_KEY".to_string(), + TomlValue::String(api_key.to_string()), + ); + let content = format!("{}\n", toml::to_string_pretty(&TomlValue::Table(root))?); + write_text_atomic_private(path, &content) +} + +fn load_json_object(path: &Path) -> Result> { + match fs::read_to_string(path) { + Ok(content) => serde_json::from_str::(&content) + .with_context(|| format!("failed to parse JSON file {}", path.display()))? + .as_object() + .cloned() + .ok_or_else(|| anyhow!("{} must contain a JSON object", path.display())), + Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(Map::new()), + Err(error) => Err(error).with_context(|| format!("failed to read {}", path.display())), + } +} + +fn load_toml_table(path: &Path) -> Result> { + match fs::read_to_string(path) { + Ok(content) => content + .parse::() + .with_context(|| format!("failed to parse TOML file {}", path.display()))? + .as_table() + .cloned() + .ok_or_else(|| anyhow!("{} must contain a TOML table", path.display())), + Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(toml::map::Map::new()), + Err(error) => Err(error).with_context(|| format!("failed to read {}", path.display())), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use tempfile::tempdir; + + #[test] + fn claude_gateway_configuration_preserves_existing_settings() { + let dir = tempdir().unwrap(); + let path = dir.path().join("settings.json"); + fs::write( + &path, + r#"{"permissions":{"allow":["Read"]},"env":{"KEEP":"yes"}}"#, + ) + .unwrap(); + + configure_claude(&path, "https://gateway.test.example", "sk-test-gateway-key").unwrap(); + + let settings: Value = serde_json::from_str(&fs::read_to_string(path).unwrap()).unwrap(); + assert_eq!(settings["permissions"]["allow"][0], "Read"); + assert_eq!(settings["env"]["KEEP"], "yes"); + assert_eq!( + settings["env"]["ANTHROPIC_BASE_URL"], + "https://gateway.test.example" + ); + assert_eq!( + settings["env"]["ANTHROPIC_AUTH_TOKEN"], + "sk-test-gateway-key" + ); + } + + #[test] + fn codex_gateway_configuration_preserves_existing_settings() { + let dir = tempdir().unwrap(); + let path = dir.path().join("config.toml"); + fs::write(&path, "model = \"test-model\"\n[env]\nKEEP = \"yes\"\n").unwrap(); + + configure_codex(&path, "https://gateway.test.example", "sk-test-gateway-key").unwrap(); + + let settings: TomlValue = fs::read_to_string(path).unwrap().parse().unwrap(); + assert_eq!(settings["model"].as_str(), Some("test-model")); + assert_eq!(settings["env"]["KEEP"].as_str(), Some("yes")); + assert_eq!( + settings["env"]["OPENAI_BASE_URL"].as_str(), + Some("https://gateway.test.example") + ); + assert_eq!( + settings["env"]["OPENAI_API_KEY"].as_str(), + Some("sk-test-gateway-key") + ); + } +} diff --git a/src/main.rs b/src/main.rs index 25bf4006..3b67e2a7 100644 --- a/src/main.rs +++ b/src/main.rs @@ -12,6 +12,7 @@ mod env; mod eval; mod experiments; mod functions; +mod gateway; mod http; mod init; mod js_runner; @@ -176,6 +177,8 @@ enum Commands { Status(CLIArgs), /// Manage coding-agent tracing Trace(CLIArgs), + /// Configure coding agents to use Braintrust Gateway + Gateway(CLIArgs), // /// View and modify config // Config(CLIArgs), } @@ -208,6 +211,7 @@ impl Commands { Commands::Switch(cmd) => &cmd.base, Commands::Status(cmd) => &cmd.base, Commands::Trace(cmd) => &cmd.base, + Commands::Gateway(cmd) => &cmd.base.login, } } @@ -238,6 +242,7 @@ impl Commands { Commands::Switch(cmd) => &mut cmd.base, Commands::Status(cmd) => &mut cmd.base, Commands::Trace(cmd) => &mut cmd.base, + Commands::Gateway(cmd) => &mut cmd.base.login, } } @@ -371,6 +376,7 @@ fn try_main() -> Result<()> { Commands::Trace(cmd) => { bt_daemon::run_trace(cmd.args, trace_host::context(cmd.base)).await? } + Commands::Gateway(cmd) => gateway::run(cmd.base, cmd.args).await?, } Ok(()) }); diff --git a/src/setup/mod.rs b/src/setup/mod.rs index 860692d4..5dadd123 100644 --- a/src/setup/mod.rs +++ b/src/setup/mod.rs @@ -1842,6 +1842,13 @@ async fn ensure_profile_or_setup_browser_auth_context( .await } +/// Resolve setup credentials, creating a durable API key when the selected +/// profile uses OAuth. Consumers write the returned key into static tool +/// configuration, where an expiring OAuth access token would not work. +pub(crate) async fn durable_setup_api_key(base: &mut BaseArgs) -> Result { + Ok(ensure_setup_auth(base, false, true).await?.api_key) +} + async fn ensure_setup_auth( base: &mut BaseArgs, prompt_for_profile_choice: bool, diff --git a/src/utils/fs_atomic.rs b/src/utils/fs_atomic.rs index dfbaaad6..82521dfd 100644 --- a/src/utils/fs_atomic.rs +++ b/src/utils/fs_atomic.rs @@ -25,6 +25,12 @@ pub fn write_text_atomic(path: &Path, contents: &str) -> Result<()> { write_bytes_atomic(path, contents.as_bytes()) } +/// Atomically write text which contains credentials with owner-only permissions +/// on platforms that support Unix-style file modes. +pub fn write_text_atomic_private(path: &Path, contents: &str) -> Result<()> { + write_atomic(path, contents.as_bytes(), Visibility::Private) +} + pub fn write_json_atomic(path: &Path, value: &T) -> Result<()> { write_json(path, value, Visibility::Default) } diff --git a/src/utils/mod.rs b/src/utils/mod.rs index 1429bebe..f5ab34b9 100644 --- a/src/utils/mod.rs +++ b/src/utils/mod.rs @@ -13,6 +13,7 @@ pub(crate) use app_url::{app_project_url, app_project_url_with_encoded_path}; pub use duration::parse_duration_to_seconds; pub use fs_atomic::{ write_bytes_atomic, write_json_atomic, write_json_atomic_private, write_text_atomic, + write_text_atomic_private, }; pub use git::GitRepo; pub(crate) use ids::new_uuid_id; From 2547d30279e1f9c0863fa47dc3b28a99bb4e6e00 Mon Sep 17 00:00:00 2001 From: Ken Jiang Date: Wed, 19 Aug 2026 10:46:23 -0400 Subject: [PATCH 2/3] fix bug --- src/setup/mod.rs | 27 +++++++++++++++++++++++++-- 1 file changed, 25 insertions(+), 2 deletions(-) diff --git a/src/setup/mod.rs b/src/setup/mod.rs index 5dadd123..5fcfa36b 100644 --- a/src/setup/mod.rs +++ b/src/setup/mod.rs @@ -1846,20 +1846,43 @@ async fn ensure_profile_or_setup_browser_auth_context( /// profile uses OAuth. Consumers write the returned key into static tool /// configuration, where an expiring OAuth access token would not work. pub(crate) async fn durable_setup_api_key(base: &mut BaseArgs) -> Result { - Ok(ensure_setup_auth(base, false, true).await?.api_key) + Ok(ensure_setup_auth_without_config_context(base, false, true) + .await? + .api_key) } async fn ensure_setup_auth( base: &mut BaseArgs, prompt_for_profile_choice: bool, needs_api_key: bool, +) -> Result { + ensure_setup_auth_with_config_context(base, prompt_for_profile_choice, needs_api_key, true) + .await +} + +async fn ensure_setup_auth_without_config_context( + base: &mut BaseArgs, + prompt_for_profile_choice: bool, + needs_api_key: bool, +) -> Result { + ensure_setup_auth_with_config_context(base, prompt_for_profile_choice, needs_api_key, false) + .await +} + +async fn ensure_setup_auth_with_config_context( + base: &mut BaseArgs, + prompt_for_profile_choice: bool, + needs_api_key: bool, + use_config_context: bool, ) -> Result { let project_was_explicit = base .project .as_deref() .map(str::trim) .is_some_and(|value| !value.is_empty()); - apply_setup_config_fallbacks(base); + if use_config_context { + apply_setup_config_fallbacks(base); + } let explicit_api_key = base .api_key From 2862a83addb2bdde2c893db7f2c27dc7a0ce3f18 Mon Sep 17 00:00:00 2001 From: Ken Jiang Date: Wed, 19 Aug 2026 11:17:07 -0400 Subject: [PATCH 3/3] address comments --- src/gateway.rs | 86 ++++++++++++++++++++++++++++++++++++++++++-------- src/main.rs | 1 + 2 files changed, 74 insertions(+), 13 deletions(-) diff --git a/src/gateway.rs b/src/gateway.rs index 7e0cf04c..330c909a 100644 --- a/src/gateway.rs +++ b/src/gateway.rs @@ -40,8 +40,13 @@ enum GatewayCommand { #[derive(Debug, Clone, Args)] struct GatewaySetupArgs { + /// Agent to configure (positional form) #[arg(value_enum)] - agent: GatewayAgent, + agent: Option, + + /// Agent to configure + #[arg(long = "agent", env = "BRAINTRUST_GATEWAY_AGENT", value_enum)] + agent_flag: Option, /// Gateway URL to use for model requests #[arg(long, env = "BRAINTRUST_GATEWAY_URL", default_value = DEFAULT_GATEWAY_URL)] @@ -58,16 +63,20 @@ enum GatewayAgent { pub async fn run(base: GatewayBaseArgs, args: GatewayArgs) -> Result<()> { let GatewayCommand::Setup(args) = args.command; + let gateway_url = normalize_gateway_url(&args.gateway_url)?; + let home = dirs::home_dir().ok_or_else(|| anyhow!("failed to resolve HOME/USERPROFILE"))?; + let agent = args + .agent + .or(args.agent_flag) + .ok_or_else(|| anyhow!("an agent is required; pass `claude`/`codex` or --agent "))?; let mut auth_base = BaseArgs { login: base.login, org_name: base.org_name, project: None, }; let api_key = setup::durable_setup_api_key(&mut auth_base).await?; - let home = dirs::home_dir().ok_or_else(|| anyhow!("failed to resolve HOME/USERPROFILE"))?; - let gateway_url = normalize_gateway_url(&args.gateway_url)?; - let path = match args.agent { + let path = match agent { GatewayAgent::Claude => { let path = home.join(".claude/settings.json"); configure_claude(&path, &gateway_url, &api_key)?; @@ -84,7 +93,7 @@ pub async fn run(base: GatewayBaseArgs, args: GatewayArgs) -> Result<()> { println!( "{}", serde_json::to_string_pretty(&serde_json::json!({ - "agent": match args.agent { GatewayAgent::Claude => "claude", GatewayAgent::Codex => "codex" }, + "agent": agent.as_str(), "gateway_url": gateway_url, "settings_path": path, }))? @@ -92,7 +101,7 @@ pub async fn run(base: GatewayBaseArgs, args: GatewayArgs) -> Result<()> { } else if auth_base.verbose { eprintln!( "Configured {} to use Braintrust Gateway ({})", - args.agent.display_name(), + agent.display_name(), path.display() ); } @@ -101,6 +110,13 @@ pub async fn run(base: GatewayBaseArgs, args: GatewayArgs) -> Result<()> { } impl GatewayAgent { + fn as_str(self) -> &'static str { + match self { + Self::Claude => "claude", + Self::Codex => "codex", + } + } + fn display_name(self) -> &'static str { match self { Self::Claude => "Claude Code", @@ -148,12 +164,48 @@ fn configure_codex(path: &Path, gateway_url: &str, api_key: &str) -> Result<()> .as_table_mut() .ok_or_else(|| anyhow!("field 'env' in {} must be a TOML table", path.display()))?; env.insert( - "OPENAI_BASE_URL".to_string(), + "BRAINTRUST_GATEWAY_API_KEY".to_string(), + TomlValue::String(api_key.to_string()), + ); + root.insert( + "model_provider".to_string(), + TomlValue::String("braintrust_gateway".to_string()), + ); + let providers = root + .entry("model_providers".to_string()) + .or_insert_with(|| TomlValue::Table(toml::map::Map::new())) + .as_table_mut() + .ok_or_else(|| { + anyhow!( + "field 'model_providers' in {} must be a TOML table", + path.display() + ) + })?; + let provider = providers + .entry("braintrust_gateway".to_string()) + .or_insert_with(|| TomlValue::Table(toml::map::Map::new())) + .as_table_mut() + .ok_or_else(|| { + anyhow!( + "field 'model_providers.braintrust_gateway' in {} must be a TOML table", + path.display() + ) + })?; + provider.insert( + "name".to_string(), + TomlValue::String("Braintrust Gateway".to_string()), + ); + provider.insert( + "base_url".to_string(), TomlValue::String(gateway_url.to_string()), ); - env.insert( - "OPENAI_API_KEY".to_string(), - TomlValue::String(api_key.to_string()), + provider.insert( + "env_key".to_string(), + TomlValue::String("BRAINTRUST_GATEWAY_API_KEY".to_string()), + ); + provider.insert( + "wire_api".to_string(), + TomlValue::String("responses".to_string()), ); let content = format!("{}\n", toml::to_string_pretty(&TomlValue::Table(root))?); write_text_atomic_private(path, &content) @@ -226,12 +278,20 @@ mod tests { assert_eq!(settings["model"].as_str(), Some("test-model")); assert_eq!(settings["env"]["KEEP"].as_str(), Some("yes")); assert_eq!( - settings["env"]["OPENAI_BASE_URL"].as_str(), - Some("https://gateway.test.example") + settings["model_provider"].as_str(), + Some("braintrust_gateway") ); assert_eq!( - settings["env"]["OPENAI_API_KEY"].as_str(), + settings["env"]["BRAINTRUST_GATEWAY_API_KEY"].as_str(), Some("sk-test-gateway-key") ); + assert_eq!( + settings["model_providers"]["braintrust_gateway"]["base_url"].as_str(), + Some("https://gateway.test.example") + ); + assert_eq!( + settings["model_providers"]["braintrust_gateway"]["env_key"].as_str(), + Some("BRAINTRUST_GATEWAY_API_KEY") + ); } } diff --git a/src/main.rs b/src/main.rs index 3b67e2a7..8a4e944a 100644 --- a/src/main.rs +++ b/src/main.rs @@ -84,6 +84,7 @@ Data & evaluation Additional docs Manage workflow docs for coding agents + gateway Configure coding agents to use Braintrust Gateway trace Manage coding-agent tracing setup Configure Braintrust setup flows (deprecated: use curl -fsSL https://braintrust.dev/wizard/setup.sh | sh) status Show current identity, org, and project context