From 68c005869b8d4faf627a76253f5dddeb37c495ce Mon Sep 17 00:00:00 2001 From: Alex Rosenzweig Date: Mon, 28 Sep 2026 09:42:43 -0500 Subject: [PATCH] refactor(agent): share model endpoint resolution Signed-off-by: Alex Rosenzweig --- crates/buzz-agent/src/config.rs | 10 ++-- crates/buzz-agent/src/config/endpoint.rs | 65 ++++++++++++++++++++++++ 2 files changed, 69 insertions(+), 6 deletions(-) create mode 100644 crates/buzz-agent/src/config/endpoint.rs diff --git a/crates/buzz-agent/src/config.rs b/crates/buzz-agent/src/config.rs index f8139358365..59d9db08425 100644 --- a/crates/buzz-agent/src/config.rs +++ b/crates/buzz-agent/src/config.rs @@ -1,3 +1,5 @@ +pub mod endpoint; + use std::time::Duration; pub const PROTOCOL_VERSION: u32 = 2; @@ -639,7 +641,6 @@ pub struct Config { impl Config { pub fn from_env() -> Result { - let databricks_host = env("DATABRICKS_HOST"); let databricks_model = env("DATABRICKS_MODEL"); let provider = resolve_provider( env("BUZZ_AGENT_PROVIDER").as_deref(), @@ -660,7 +661,7 @@ impl Config { // Databricks borrows api_key as the *optional* `DATABRICKS_TOKEN` escape // hatch — empty means "use OAuth PKCE." Legacy Databricks encodes the // model in the URL path; Databricks v2 keeps it in the request body. - let (api_key, model, base_url, openai_api) = match provider { + let (api_key, model, openai_api) = match provider { Provider::Anthropic => ( req("ANTHROPIC_API_KEY")?, resolve_model( @@ -668,7 +669,6 @@ impl Config { env("ANTHROPIC_MODEL").as_deref(), ) .ok_or_else(|| "config: ANTHROPIC_MODEL required".to_string())?, - env_or("ANTHROPIC_BASE_URL", "https://api.anthropic.com"), OpenAiApi::Auto, // unused for Anthropic ), Provider::OpenAi => ( @@ -678,14 +678,12 @@ impl Config { env("OPENAI_COMPAT_MODEL").as_deref(), ) .ok_or_else(|| "config: OPENAI_COMPAT_MODEL required".to_string())?, - env_or("OPENAI_COMPAT_BASE_URL", "https://api.openai.com/v1"), parse_openai_api(env("OPENAI_COMPAT_API").as_deref())?, ), Provider::Databricks | Provider::DatabricksV2 => ( env("DATABRICKS_TOKEN").unwrap_or_default(), resolve_model(buzz_agent_model.as_deref(), databricks_model.as_deref()) .ok_or_else(|| "config: DATABRICKS_MODEL required".to_string())?, - databricks_host.ok_or_else(|| "config: DATABRICKS_HOST required".to_string())?, OpenAiApi::Chat, // only read by OpenAI/legacy Databricks dispatch ), Provider::OpenRouter => ( @@ -695,10 +693,10 @@ impl Config { env("OPENROUTER_MODEL").as_deref(), ) .ok_or_else(|| "config: OPENROUTER_MODEL required".to_string())?, - env_or("OPENROUTER_BASE_URL", "https://openrouter.ai/api/v1"), OpenAiApi::Chat, // OpenRouter uses Chat Completions only ), }; + let base_url = endpoint::provider_base_url(&provider, env)?; let system_prompt = match (env("BUZZ_AGENT_SYSTEM_PROMPT"), env("BUZZ_AGENT_SYSTEM_PROMPT_FILE")) { (Some(_), Some(_)) => return Err( "config: BUZZ_AGENT_SYSTEM_PROMPT and BUZZ_AGENT_SYSTEM_PROMPT_FILE are mutually exclusive".into()), diff --git a/crates/buzz-agent/src/config/endpoint.rs b/crates/buzz-agent/src/config/endpoint.rs new file mode 100644 index 00000000000..6c75100bac4 --- /dev/null +++ b/crates/buzz-agent/src/config/endpoint.rs @@ -0,0 +1,65 @@ +//! Endpoint resolution shared by agent startup and desktop security suggestions. +use super::Provider; + +pub fn provider_base_url( + provider: &Provider, + lookup: impl Fn(&str) -> Option, +) -> Result { + let (key, default) = match provider { + Provider::Anthropic => ("ANTHROPIC_BASE_URL", Some("https://api.anthropic.com")), + Provider::OpenAi => ("OPENAI_COMPAT_BASE_URL", Some("https://api.openai.com/v1")), + Provider::OpenRouter => ("OPENROUTER_BASE_URL", Some("https://openrouter.ai/api/v1")), + Provider::Databricks | Provider::DatabricksV2 => ("DATABRICKS_HOST", None), + }; + lookup(key) + .or_else(|| default.map(str::to_owned)) + .ok_or_else(|| format!("config: {key} required")) +} + +/// Resolve without fetching credentials, signing in, or making network requests. +pub fn configured_endpoint(lookup: impl Fn(&str) -> Option) -> Option { + let provider = super::resolve_provider( + lookup("BUZZ_AGENT_PROVIDER").as_deref(), + Some("configured"), + Some("configured"), + Some("configured"), + ) + .ok()?; + provider_base_url(&provider, lookup).ok() +} + +#[cfg(test)] +mod tests { + use super::*; + #[test] + fn endpoints_follow_selected_provider_and_overrides() { + for (provider, expected) in [ + ("anthropic", "https://api.anthropic.com"), + ("openai-compat", "https://api.openai.com/v1"), + ("openrouter", "https://openrouter.ai/api/v1"), + ] { + assert_eq!( + configured_endpoint(|key| (key == "BUZZ_AGENT_PROVIDER").then(|| provider.into())) + .as_deref(), + Some(expected) + ); + } + assert_eq!( + configured_endpoint(|key| match key { + "BUZZ_AGENT_PROVIDER" => Some("databricks_v2".into()), + "DATABRICKS_HOST" => Some("https://workspace.example.com".into()), + _ => None, + }) + .as_deref(), + Some("https://workspace.example.com") + ); + assert!(configured_endpoint(|_| None).is_none()); + assert_eq!( + provider_base_url(&Provider::OpenAi, |_| Some( + "http://localhost:1234/v1".into() + )) + .unwrap(), + "http://localhost:1234/v1" + ); + } +}