1use std::sync::Arc;
9
10use tokio::sync::Mutex;
11
12use crate::app::Config;
13use crate::models::config::BackendConfig;
14use crate::models::{ModelError, Result, lookup_provider};
15use crate::utils::{resolve_api_key, resolve_provider_key, resolve_provider_key_with_fallback};
16
17use super::model::{anthropic, gemini, meta};
18
19fn require_key(provider: &str, default_env: &str, override_env: Option<&str>) -> Result<String> {
23 resolve_provider_key(provider, default_env, override_env).ok_or_else(|| {
24 let env = override_env.unwrap_or(default_env);
25 ModelError::Authentication(format!(
26 "{provider} requires env var {env} (or `mermaid login {provider}`)"
27 ))
28 })
29}
30
31fn require_key_with_fallback(
32 provider: &str,
33 env_var: &str,
34 fallback_env_var: &str,
35) -> Result<String> {
36 resolve_provider_key_with_fallback(provider, env_var, fallback_env_var, None).ok_or_else(|| {
37 ModelError::Authentication(format!(
38 "{provider} requires env var {env_var} (or legacy {fallback_env_var}, or \
39 `mermaid login {provider}`)"
40 ))
41 })
42}
43
44fn key_env_source(default_env: &str, override_env: Option<&str>) -> Option<String> {
50 let env = override_env.unwrap_or(default_env);
54 resolve_api_key(env, None).map(|_| env.to_string())
55}
56
57fn key_env_source_with_fallback(default_env: &str, fallback_env: &str) -> Option<String> {
59 if resolve_api_key(default_env, None).is_some() {
60 return Some(default_env.to_string());
61 }
62 if resolve_api_key(fallback_env, None).is_some() {
63 return Some(fallback_env.to_string());
64 }
65 None
66}
67
68#[derive(Debug, Clone, PartialEq, Eq)]
78pub(crate) struct ProviderEndpoint {
79 pub(crate) base_url: String,
81 pub(crate) api_key: Option<String>,
84 pub(crate) key_env: Option<String>,
87}
88
89impl ProviderEndpoint {
90 fn require_key(&self, provider: &str) -> Result<String> {
96 self.api_key
97 .clone()
98 .ok_or_else(|| ModelError::Authentication(format!("{provider} requires an API key")))
99 }
100}
101
102fn is_known_provider(config: &Config, provider_lc: &str) -> bool {
105 matches!(provider_lc, "anthropic" | "gemini" | "meta")
106 || lookup_provider(provider_lc).is_some()
107 || config.providers.contains_key(provider_lc)
108}
109
110pub(crate) fn resolve_provider_endpoint(
118 config: &Config,
119 provider: &str,
120) -> Result<ProviderEndpoint> {
121 let provider_lc = provider.to_lowercase();
122 let user_cfg = config.providers.get(&provider_lc);
123 let override_env = user_cfg.and_then(|c| c.api_key_env.as_deref());
124 let override_url = user_cfg.and_then(|c| c.base_url.clone());
125
126 let bespoke = match provider_lc.as_str() {
129 "anthropic" => Some((
130 anthropic::DEFAULT_API_KEY_ENV,
131 None,
132 anthropic::DEFAULT_BASE_URL,
133 )),
134 "gemini" => Some((
135 gemini::DEFAULT_API_KEY_ENV,
136 Some(gemini::LEGACY_API_KEY_ENV),
137 gemini::DEFAULT_BASE_URL,
138 )),
139 "meta" => Some((meta::DEFAULT_API_KEY_ENV, None, meta::DEFAULT_BASE_URL)),
140 _ => None,
141 };
142 if let Some((default_env, legacy_env, default_url)) = bespoke {
143 let base_url = resolve_overridable_base_url(&provider_lc, override_url, default_url)?;
144 let (api_key, key_env) = match legacy_env {
145 Some(legacy) if override_env.is_none() => (
148 require_key_with_fallback(&provider_lc, default_env, legacy)?,
149 key_env_source_with_fallback(default_env, legacy),
150 ),
151 _ => (
152 require_key(&provider_lc, default_env, override_env)?,
153 key_env_source(default_env, override_env),
154 ),
155 };
156 return Ok(ProviderEndpoint {
157 base_url,
158 api_key: Some(api_key),
159 key_env,
160 });
161 }
162
163 if provider_lc == "cloudflare" {
169 let profile = lookup_provider("cloudflare").expect("cloudflare is in the registry");
170 let api_key_env = override_env.unwrap_or(profile.api_key_env);
171 let base_url = match override_url {
172 Some(url) => {
175 validate_provider_base_url(&url)?;
176 warn_overridden_provider_host("cloudflare", &url);
177 url
178 },
179 None => match require_cloudflare_account_id() {
183 Ok(id) => cloudflare_base_url(&id),
184 Err(_)
185 if resolve_provider_key("cloudflare", profile.api_key_env, override_env)
186 .is_none() =>
187 {
188 return Err(ModelError::Authentication(format!(
189 "cloudflare requires env vars {api_key_env} and CLOUDFLARE_ACCOUNT_ID — \
190 create a token at https://dash.cloudflare.com/profile/api-tokens; the \
191 account id is on your Cloudflare dashboard (or set \
192 [providers.cloudflare].base_url)"
193 )));
194 },
195 Err(e) => return Err(e),
196 },
197 };
198 let api_key = resolve_optional_key(
199 &provider_lc,
200 profile.api_key_env,
201 override_env,
202 &base_url,
203 profile.key_hint,
204 )?;
205 let key_env = api_key
206 .is_some()
207 .then(|| key_env_source(profile.api_key_env, override_env))
208 .flatten();
209 return Ok(ProviderEndpoint {
210 base_url,
211 api_key,
212 key_env,
213 });
214 }
215
216 if let Some(profile) = lookup_provider(&provider_lc) {
218 let base_url = resolve_overridable_base_url(&provider_lc, override_url, profile.base_url)?;
219 let api_key = resolve_optional_key(
223 &provider_lc,
224 profile.api_key_env,
225 override_env,
226 &base_url,
227 profile.key_hint,
228 )?;
229 let key_env = api_key
230 .is_some()
231 .then(|| key_env_source(profile.api_key_env, override_env))
232 .flatten();
233 return Ok(ProviderEndpoint {
234 base_url,
235 api_key,
236 key_env,
237 });
238 }
239
240 if user_cfg.is_some() {
243 let base_url = override_url.ok_or_else(|| {
244 ModelError::InvalidRequest(format!(
245 "custom provider '{provider_lc}' requires base_url in config"
246 ))
247 })?;
248 let resolved = match override_env {
255 Some(env) => resolve_provider_key(&provider_lc, env, None),
256 None => crate::utils::default_store().get(&provider_lc),
257 };
258 let (api_key, key_env) = match resolved {
259 Some(key) => {
260 validate_provider_base_url(&base_url)?;
261 let env = override_env.and_then(|env| key_env_source(env, None));
262 (Some(key), env)
263 },
264 None if base_url_is_local(&base_url) => (None, None),
265 None => {
266 let reason = match override_env {
267 Some(env) => format!(
268 "requires env var {env} (or `mermaid login {provider_lc}`, or a \
269 loopback/LAN base_url)"
270 ),
271 None => "requires api_key_env, or a loopback/LAN base_url".to_string(),
272 };
273 return Err(ModelError::Authentication(format!(
274 "custom provider '{provider_lc}' {reason}"
275 )));
276 },
277 };
278 return Ok(ProviderEndpoint {
279 base_url,
280 api_key,
281 key_env,
282 });
283 }
284
285 Err(ModelError::InvalidRequest(format!(
286 "Unknown provider '{provider}'"
287 )))
288}
289
290fn resolve_optional_key(
295 provider: &str,
296 default_env: &str,
297 override_env: Option<&str>,
298 base_url: &str,
299 hint: Option<&str>,
300) -> Result<Option<String>> {
301 if let Some(key) = resolve_provider_key(provider, default_env, override_env) {
302 return Ok(Some(key));
303 }
304 if base_url_is_local(base_url) {
305 return Ok(None);
306 }
307 let env = override_env.unwrap_or(default_env);
308 let mut msg = format!("{provider} requires env var {env} (or `mermaid login {provider}`)");
309 if let Some(h) = hint {
310 msg.push_str(" — ");
311 msg.push_str(h);
312 }
313 Err(ModelError::Authentication(msg))
314}
315
316fn base_url_is_local(base_url: &str) -> bool {
319 reqwest::Url::parse(base_url)
320 .ok()
321 .and_then(|u| {
322 u.host_str()
323 .map(|h| crate::utils::classify_host(h).is_internal())
324 })
325 .unwrap_or(false)
326}
327
328fn merged_headers(
334 profile: &crate::models::ProviderProfile,
335 user_cfg: Option<&crate::app::UserProviderConfig>,
336) -> std::collections::HashMap<String, String> {
337 let mut headers: std::collections::HashMap<String, String> = profile
338 .extra_headers
339 .iter()
340 .map(|(k, v)| ((*k).to_string(), (*v).to_string()))
341 .collect();
342 if let Some(cfg) = user_cfg {
343 for (k, v) in &cfg.extra_headers {
344 headers.insert(k.clone(), v.clone());
345 }
346 for (header, env_var) in &cfg.env_headers {
347 if let Ok(val) = std::env::var(env_var) {
348 headers.insert(header.clone(), val);
349 }
350 }
351 }
352 headers
353}
354
355use super::model::{
356 AnthropicProvider, GeminiProvider, MetaProvider, ModelProvider, OllamaProvider,
357 OpenAICompatProvider,
358};
359
360type ProviderCell = Arc<tokio::sync::OnceCell<Arc<dyn ModelProvider>>>;
364
365pub struct ProviderFactory {
369 config: Arc<Config>,
370 cache: Mutex<std::collections::HashMap<String, ProviderCell>>,
373}
374
375impl ProviderFactory {
376 pub fn new(config: Config) -> Self {
377 Self {
378 config: Arc::new(config),
379 cache: Mutex::new(std::collections::HashMap::new()),
380 }
381 }
382
383 pub fn with_seeded_providers(
394 config: Config,
395 seeds: impl IntoIterator<Item = (String, Arc<dyn ModelProvider>)>,
396 ) -> Self {
397 let cache = seeds
398 .into_iter()
399 .map(|(model_id, provider)| {
400 let cell = tokio::sync::OnceCell::new_with(Some(provider));
401 (normalize_cache_key(&model_id), Arc::new(cell))
402 })
403 .collect();
404 Self {
405 config: Arc::new(config),
406 cache: Mutex::new(cache),
407 }
408 }
409
410 pub fn config(&self) -> &Config {
411 &self.config
412 }
413
414 pub async fn resolve(&self, model_id: &str) -> Result<Arc<dyn ModelProvider>> {
418 let key = normalize_cache_key(model_id);
419 let cell = {
423 let mut cache = self.cache.lock().await;
424 Arc::clone(
425 cache
426 .entry(key)
427 .or_insert_with(|| Arc::new(tokio::sync::OnceCell::new())),
428 )
429 };
430 let provider = cell
431 .get_or_try_init(|| async {
432 let p = build_provider(&self.config, model_id).await?;
433 Ok::<Arc<dyn ModelProvider>, ModelError>(Arc::from(p))
434 })
435 .await?;
436 Ok(Arc::clone(provider))
437 }
438}
439
440async fn build_provider(config: &Config, model_id: &str) -> Result<Box<dyn ModelProvider>> {
449 let (provider, model_name) = parse_model_id(model_id);
450 let provider_lc = provider.to_lowercase();
451
452 if provider_lc == "ollama" {
455 let backend = ollama_backend_config(config);
456 let p = OllamaProvider::with_app_config(
457 model_name,
458 Arc::new(backend),
459 Arc::new(config.clone()),
460 )
461 .await?;
462 return Ok(Box::new(p));
463 }
464
465 if !is_known_provider(config, &provider_lc) {
470 return Err(ModelError::InvalidRequest(format!(
471 "Unknown provider '{}' (model_id: {})",
472 provider, model_id
473 )));
474 }
475 let endpoint = resolve_provider_endpoint(config, &provider_lc)?;
476 let user_cfg = config.providers.get(&provider_lc);
477
478 if provider_lc == "anthropic" {
480 let p = AnthropicProvider::new(
481 endpoint.require_key(&provider_lc)?,
482 model_name.to_string(),
483 endpoint.base_url,
484 )?;
485 return Ok(Box::new(p));
486 }
487
488 if provider_lc == "gemini" {
490 let p = GeminiProvider::new(
491 endpoint.require_key(&provider_lc)?,
492 model_name.to_string(),
493 endpoint.base_url,
494 )?;
495 return Ok(Box::new(p));
496 }
497
498 if provider_lc == "meta" {
501 let api_key = endpoint.require_key(&provider_lc)?;
502 let mut extra_headers = std::collections::HashMap::new();
503 if let Some(cfg) = user_cfg {
504 extra_headers.extend(cfg.extra_headers.clone());
505 for (header, env_var) in &cfg.env_headers {
506 if let Ok(value) = std::env::var(env_var) {
507 extra_headers.insert(header.clone(), value);
508 }
509 }
510 }
511 let p = MetaProvider::new(
512 api_key,
513 model_name.to_string(),
514 endpoint.base_url,
515 extra_headers,
516 )?;
517 return Ok(Box::new(p));
518 }
519
520 let profile = match lookup_provider(&provider_lc) {
524 Some(profile) => profile,
525 None => user_cfg
528 .and_then(|cfg| user_profile_to_static(&provider_lc, cfg))
529 .ok_or_else(|| {
530 ModelError::InvalidRequest(format!(
531 "Unknown provider '{}' (model_id: {})",
532 provider, model_id
533 ))
534 })?,
535 };
536 let extra_headers = merged_headers(profile, user_cfg);
537 let p = OpenAICompatProvider::new(
538 profile,
539 endpoint.base_url,
540 endpoint.api_key,
541 model_name.to_string(),
542 extra_headers,
543 )?;
544 Ok(Box::new(p))
545}
546
547fn normalize_cache_key(model_id: &str) -> String {
551 let (provider, model) = parse_model_id(model_id);
552 format!("{}/{}", provider.to_lowercase(), model)
553}
554
555fn parse_model_id(model_id: &str) -> (String, &str) {
558 match model_id.split_once('/') {
559 Some((p, m)) => (p.to_string(), m),
560 None => ("ollama".to_string(), model_id),
561 }
562}
563
564pub(crate) fn ollama_backend_config(config: &Config) -> BackendConfig {
565 BackendConfig {
566 ollama_url: format!("{}:{}", config.ollama.host, config.ollama.port),
569 max_idle_per_host: 10,
570 timeout_secs: 10,
571 ollama_autostart: config.ollama.auto_start,
572 }
573}
574
575static PROFILE_CACHE: std::sync::LazyLock<
579 std::sync::Mutex<std::collections::HashMap<String, &'static crate::models::ProviderProfile>>,
580> = std::sync::LazyLock::new(|| std::sync::Mutex::new(std::collections::HashMap::new()));
581
582fn user_profile_to_static(
596 name: &str,
597 user_cfg: &crate::app::UserProviderConfig,
598) -> Option<&'static crate::models::ProviderProfile> {
599 use crate::models::{ProviderProfile, ReasoningExtraction, ReasoningStrategy};
600
601 let compat = user_cfg.compat.as_deref().unwrap_or("openai");
602 let base_url = user_cfg.base_url.clone().unwrap_or_default();
603 let api_key_env = user_cfg.api_key_env.clone().unwrap_or_default();
604
605 let cache_key = format!("{name}\u{0}{base_url}\u{0}{api_key_env}\u{0}{compat}");
608
609 let mut cache = PROFILE_CACHE
610 .lock()
611 .unwrap_or_else(|poisoned| poisoned.into_inner());
612 if let Some(profile) = cache.get(cache_key.as_str()) {
613 return Some(*profile);
616 }
617
618 let strategy = match compat {
619 "openai" => ReasoningStrategy::None,
620 "openai-effort" => ReasoningStrategy::Effort,
621 "openrouter" => ReasoningStrategy::OpenRouterShape,
622 _ => ReasoningStrategy::None,
623 };
624
625 let profile = Box::new(ProviderProfile {
626 name: Box::leak(name.to_string().into_boxed_str()),
627 base_url: Box::leak(base_url.into_boxed_str()),
628 api_key_env: Box::leak(api_key_env.into_boxed_str()),
629 key_hint: None,
630 extra_headers: &[],
631 reasoning_strategy: strategy,
632 reasoning_extraction: ReasoningExtraction::None,
633 max_tokens_param: crate::models::MaxTokensParam::MaxTokens,
634 disable_parallel_tool_calls_for: &[],
635 });
636 let leaked: &'static ProviderProfile = Box::leak(profile);
637 cache.insert(cache_key, leaked);
638 Some(leaked)
639}
640
641fn cloudflare_base_url(account_id: &str) -> String {
646 format!(
647 "https://api.cloudflare.com/client/v4/accounts/{}/ai/v1",
648 account_id.trim()
649 )
650}
651
652pub(crate) fn discovery_base_url(
660 profile: &crate::models::ProviderProfile,
661 override_url: Option<String>,
662) -> Option<String> {
663 if override_url.is_some() {
664 return override_url;
665 }
666 if profile.name == "cloudflare" {
667 return require_cloudflare_account_id()
668 .ok()
669 .map(|id| cloudflare_base_url(&id));
670 }
671 Some(profile.base_url.to_string())
672}
673
674fn require_cloudflare_account_id() -> Result<String> {
679 resolve_api_key("CLOUDFLARE_ACCOUNT_ID", None)
680 .map(|s| s.trim().to_string())
681 .filter(|s| !s.is_empty())
682 .ok_or_else(|| {
683 ModelError::Authentication(
684 "cloudflare requires env var CLOUDFLARE_ACCOUNT_ID (your Cloudflare account id) — \
685 find it on your Cloudflare dashboard, or set [providers.cloudflare].base_url"
686 .to_string(),
687 )
688 })
689}
690
691fn validate_provider_base_url(url: &str) -> Result<()> {
697 let parsed = reqwest::Url::parse(url).map_err(|e| {
698 ModelError::InvalidRequest(format!("invalid provider base_url '{url}': {e}"))
699 })?;
700 match parsed.scheme() {
701 "https" => Ok(()),
702 "http"
707 if crate::utils::classify_host(parsed.host_str().unwrap_or_default()).is_loopback() =>
708 {
709 Ok(())
710 },
711 "http" => Err(ModelError::InvalidRequest(format!(
712 "provider base_url '{url}' uses http:// to a non-loopback host — refusing to send the \
713 API key in cleartext. Use https, or http://localhost for a local server."
714 ))),
715 other => Err(ModelError::InvalidRequest(format!(
716 "provider base_url '{url}' has unsupported scheme '{other}' (use http or https)"
717 ))),
718 }
719}
720
721fn resolve_overridable_base_url(
737 provider: &str,
738 override_url: Option<String>,
739 default_url: &str,
740) -> Result<String> {
741 match override_url {
742 Some(url) => {
743 validate_provider_base_url(&url)?;
744 warn_overridden_provider_host(provider, &url);
745 Ok(url)
746 },
747 None => Ok(default_url.to_string()),
748 }
749}
750
751static WARNED_OVERRIDE_HOSTS: std::sync::LazyLock<
755 std::sync::Mutex<std::collections::HashSet<String>>,
756> = std::sync::LazyLock::new(|| std::sync::Mutex::new(std::collections::HashSet::new()));
757
758fn warn_overridden_provider_host(provider: &str, base_url: &str) {
762 let host = provider_host(base_url);
763 if should_warn_once(&format!("{provider}@{host}")) {
764 tracing::warn!(
765 "built-in provider '{}' base_url overridden in config: the {} API key will be sent to \
766 host '{}' instead of the trusted default endpoint",
767 provider,
768 provider,
769 host
770 );
771 }
772}
773
774fn provider_host(base_url: &str) -> String {
777 reqwest::Url::parse(base_url)
778 .ok()
779 .and_then(|u| u.host_str().map(str::to_string))
780 .unwrap_or_else(|| "<unknown>".to_string())
781}
782
783fn should_warn_once(key: &str) -> bool {
786 let mut warned = WARNED_OVERRIDE_HOSTS
787 .lock()
788 .unwrap_or_else(|poisoned| poisoned.into_inner());
789 warned.insert(key.to_string())
790}
791
792#[cfg(test)]
793mod tests {
794 use super::*;
795
796 #[test]
797 fn base_url_is_local_classifies_hosts() {
798 assert!(base_url_is_local("http://127.0.0.1:8000/v1"));
799 assert!(base_url_is_local("http://localhost:1234/v1"));
800 assert!(base_url_is_local("http://192.168.1.5:8000/v1"));
801 assert!(!base_url_is_local("https://api.openai.com/v1"));
802 assert!(!base_url_is_local("not a url"));
803 }
804
805 #[test]
806 fn merged_headers_keeps_static_profile_headers_and_user_overrides() {
807 let profile = crate::models::lookup_provider("openrouter").unwrap();
808 let base = merged_headers(profile, None);
810 assert_eq!(
811 base.get("X-OpenRouter-Title").map(String::as_str),
812 Some("Mermaid")
813 );
814 assert!(base.contains_key("HTTP-Referer"));
815 let mut cfg = crate::app::UserProviderConfig::default();
817 cfg.extra_headers.insert("X-Custom".into(), "v".into());
818 cfg.extra_headers
819 .insert("X-OpenRouter-Title".into(), "Override".into());
820 let merged = merged_headers(profile, Some(&cfg));
821 assert_eq!(merged.get("X-Custom").map(String::as_str), Some("v"));
822 assert_eq!(
823 merged.get("X-OpenRouter-Title").map(String::as_str),
824 Some("Override")
825 );
826 assert!(merged.contains_key("HTTP-Referer"));
827 }
828
829 #[test]
830 fn merged_headers_resolves_env_headers_and_skips_missing() {
831 let profile = crate::models::lookup_provider("openai").unwrap();
832 let mut cfg = crate::app::UserProviderConfig::default();
833 cfg.env_headers
834 .insert("X-Gateway-Token".into(), "MERMAID_TEST_GW_TOKEN".into());
835 temp_env::with_var("MERMAID_TEST_GW_TOKEN", Some("secret123"), || {
836 let merged = merged_headers(profile, Some(&cfg));
837 assert_eq!(
838 merged.get("X-Gateway-Token").map(String::as_str),
839 Some("secret123")
840 );
841 });
842 temp_env::with_var("MERMAID_TEST_GW_TOKEN", None::<&str>, || {
843 assert!(!merged_headers(profile, Some(&cfg)).contains_key("X-Gateway-Token"));
844 });
845 }
846
847 #[test]
848 fn base_url_requires_https_for_remote_hosts() {
849 assert!(validate_provider_base_url("http://api.example.com/v1").is_err());
851 assert!(validate_provider_base_url("ftp://example.com").is_err());
852 assert!(validate_provider_base_url("https://api.example.com/v1").is_ok());
854 assert!(validate_provider_base_url("http://localhost:11434/v1").is_ok());
855 assert!(validate_provider_base_url("http://127.0.0.1:8000").is_ok());
856 assert!(validate_provider_base_url("http://[::1]:8000").is_ok());
857 assert!(validate_provider_base_url("http://192.168.1.5:8080").is_err());
860 assert!(validate_provider_base_url("http://169.254.169.254").is_err());
861 }
862
863 #[test]
864 fn cloudflare_base_url_synthesizes_account_scoped_endpoint() {
865 assert_eq!(
866 cloudflare_base_url("acct123"),
867 "https://api.cloudflare.com/client/v4/accounts/acct123/ai/v1"
868 );
869 assert_eq!(
871 cloudflare_base_url(" acct123\n"),
872 "https://api.cloudflare.com/client/v4/accounts/acct123/ai/v1"
873 );
874 }
875
876 #[test]
877 fn cloudflare_account_id_required_and_non_blank() {
878 temp_env::with_var("CLOUDFLARE_ACCOUNT_ID", None::<&str>, || {
880 let err = require_cloudflare_account_id().expect_err("must error when unset");
881 assert!(format!("{err}").contains("CLOUDFLARE_ACCOUNT_ID"));
882 });
883 temp_env::with_var("CLOUDFLARE_ACCOUNT_ID", Some(" "), || {
885 assert!(require_cloudflare_account_id().is_err());
886 });
887 temp_env::with_var("CLOUDFLARE_ACCOUNT_ID", Some(" acct123 "), || {
889 assert_eq!(require_cloudflare_account_id().unwrap(), "acct123");
890 });
891 }
892
893 #[test]
894 fn discovery_base_url_resolves_per_provider() {
895 let cf = lookup_provider("cloudflare").expect("cloudflare is in the registry");
896 let openai = lookup_provider("openai").expect("openai is in the registry");
897 assert_eq!(
899 discovery_base_url(cf, Some("https://gw.example/v1".to_string())),
900 Some("https://gw.example/v1".to_string())
901 );
902 assert_eq!(
904 discovery_base_url(openai, None),
905 Some(openai.base_url.to_string())
906 );
907 temp_env::with_var("CLOUDFLARE_ACCOUNT_ID", Some("acct123"), || {
909 assert_eq!(
910 discovery_base_url(cf, None),
911 Some("https://api.cloudflare.com/client/v4/accounts/acct123/ai/v1".to_string())
912 );
913 });
914 temp_env::with_var("CLOUDFLARE_ACCOUNT_ID", None::<&str>, || {
916 assert_eq!(discovery_base_url(cf, None), None);
917 });
918 }
919
920 #[tokio::test]
921 async fn cloudflare_missing_both_env_vars_is_one_combined_error() {
922 temp_env::async_with_vars(
923 [
924 ("CLOUDFLARE_ACCOUNT_ID", None::<&str>),
925 ("CLOUDFLARE_API_TOKEN", None),
926 ],
927 async {
928 let f = ProviderFactory::new(Config::default());
929 let err = match f.resolve("cloudflare/@cf/zai-org/glm-5.2").await {
930 Ok(_) => panic!("must fail with neither env var set"),
931 Err(e) => e,
932 };
933 let msg = format!("{err}");
934 assert!(
935 msg.contains("CLOUDFLARE_API_TOKEN") && msg.contains("CLOUDFLARE_ACCOUNT_ID"),
936 "one error must name both missing vars, got: {msg}"
937 );
938 },
939 )
940 .await;
941 }
942
943 use std::sync::atomic::{AtomicUsize, Ordering};
944
945 fn unique_env(prefix: &str) -> String {
946 static N: AtomicUsize = AtomicUsize::new(0);
947 format!(
948 "{}_{}_{}",
949 prefix,
950 std::process::id(),
951 N.fetch_add(1, Ordering::SeqCst)
952 )
953 }
954
955 #[test]
956 fn parse_bare_name_defaults_to_ollama() {
957 let (p, m) = parse_model_id("qwen3-coder:30b");
958 assert_eq!(p, "ollama");
959 assert_eq!(m, "qwen3-coder:30b");
960 }
961
962 #[test]
963 fn parse_prefixed() {
964 let (p, m) = parse_model_id("anthropic/claude-opus-4-7");
965 assert_eq!(p, "anthropic");
966 assert_eq!(m, "claude-opus-4-7");
967 }
968
969 #[tokio::test]
970 async fn meta_requires_its_documented_api_key_env() {
971 temp_env::async_with_vars(
972 [(
973 crate::providers::model::meta::DEFAULT_API_KEY_ENV,
974 None::<&str>,
975 )],
976 async {
977 let factory = ProviderFactory::new(Config::default());
978 let error = match factory.resolve("meta/muse-spark-1.1").await {
979 Ok(_) => panic!("Meta must require an API key"),
980 Err(error) => error,
981 };
982 assert!(
983 error
984 .to_string()
985 .contains(crate::providers::model::meta::DEFAULT_API_KEY_ENV)
986 );
987 },
988 )
989 .await;
990 }
991
992 #[tokio::test]
993 async fn meta_routes_to_responses_provider_with_muse_capabilities() {
994 temp_env::async_with_vars(
995 [(
996 crate::providers::model::meta::DEFAULT_API_KEY_ENV,
997 Some("test-key"),
998 )],
999 async {
1000 let factory = ProviderFactory::new(Config::default());
1001 let provider = factory.resolve("meta/muse-spark-1.1").await.unwrap();
1002 let capabilities = provider.capabilities();
1003 assert!(capabilities.supports_tools);
1004 assert!(capabilities.supports_vision);
1005 assert!(capabilities.emits_provider_continuation);
1006 assert_eq!(
1007 capabilities.max_context_tokens,
1008 Some(crate::constants::META_MUSE_SPARK_CONTEXT_WINDOW)
1009 );
1010 assert_eq!(
1011 capabilities.max_output_tokens,
1012 Some(crate::constants::META_MUSE_SPARK_MAX_OUTPUT_TOKENS)
1013 );
1014 },
1015 )
1016 .await;
1017 }
1018
1019 #[test]
1020 fn gemini_key_resolution_accepts_legacy_fallback() {
1021 let primary = unique_env("MERMAID_FACTORY_GEMINI_PRIMARY");
1022 let legacy = unique_env("MERMAID_FACTORY_GEMINI_LEGACY");
1023 temp_env::with_vars(
1024 [(primary.as_str(), None), (legacy.as_str(), Some("legacy"))],
1025 || {
1026 let resolved = require_key_with_fallback("gemini", &primary, &legacy)
1027 .expect("legacy fallback should resolve");
1028 assert_eq!(resolved, "legacy");
1029 },
1030 );
1031 }
1032
1033 #[test]
1034 fn gemini_key_resolution_prefers_google_primary() {
1035 let primary = unique_env("MERMAID_FACTORY_GEMINI_PRIMARY2");
1036 let legacy = unique_env("MERMAID_FACTORY_GEMINI_LEGACY2");
1037 temp_env::with_vars(
1038 [
1039 (primary.as_str(), Some("google")),
1040 (legacy.as_str(), Some("legacy")),
1041 ],
1042 || {
1043 let resolved = require_key_with_fallback("gemini", &primary, &legacy)
1044 .expect("primary should resolve");
1045 assert_eq!(resolved, "google");
1046 },
1047 );
1048 }
1049
1050 #[tokio::test]
1051 async fn factory_reports_unknown_provider_clearly() {
1052 let cfg = Config::default();
1053 let f = ProviderFactory::new(cfg);
1054 match f.resolve("totally-made-up/model").await {
1055 Ok(_) => panic!("expected error"),
1056 Err(e) => {
1057 let msg = format!("{}", e);
1058 assert!(
1059 msg.contains("totally-made-up") || msg.contains("Unknown provider"),
1060 "error message: {}",
1061 msg
1062 );
1063 },
1064 }
1065 }
1066
1067 #[test]
1068 fn normalize_cache_key_lowercases_provider_only() {
1069 assert_eq!(
1071 normalize_cache_key("Anthropic/Claude-X"),
1072 "anthropic/Claude-X"
1073 );
1074 assert_eq!(
1075 normalize_cache_key("anthropic/Claude-X"),
1076 "anthropic/Claude-X"
1077 );
1078 assert_eq!(normalize_cache_key("qwen3:30b"), "ollama/qwen3:30b");
1080 }
1081
1082 #[tokio::test]
1083 async fn resolve_is_single_flight_and_cached() {
1084 let f = ProviderFactory::new(Config::default());
1088 let (a, b) = tokio::join!(
1089 f.resolve("ollama/test-model"),
1090 f.resolve("Ollama/test-model"),
1091 );
1092 let a = a.expect("resolve a");
1093 let b = b.expect("resolve b");
1094 assert!(
1095 Arc::ptr_eq(&a, &b),
1096 "expected one cached provider for casing variants + concurrent resolve"
1097 );
1098 }
1099
1100 #[test]
1104 fn builtin_base_url_override_validated_and_resolved() {
1105 assert_eq!(
1107 resolve_overridable_base_url("anthropic", None, "https://api.anthropic.com/v1")
1108 .unwrap(),
1109 "https://api.anthropic.com/v1"
1110 );
1111 assert_eq!(
1113 resolve_overridable_base_url(
1114 "anthropic",
1115 Some("https://proxy.internal/v1".to_string()),
1116 "https://api.anthropic.com/v1",
1117 )
1118 .unwrap(),
1119 "https://proxy.internal/v1"
1120 );
1121 assert!(
1124 resolve_overridable_base_url(
1125 "anthropic",
1126 Some("http://attacker.example/v1".to_string()),
1127 "https://api.anthropic.com/v1",
1128 )
1129 .is_err()
1130 );
1131 assert!(
1133 resolve_overridable_base_url(
1134 "openai",
1135 Some("http://localhost:8080/v1".to_string()),
1136 "https://api.openai.com/v1",
1137 )
1138 .is_ok()
1139 );
1140 }
1141
1142 #[test]
1143 fn provider_host_extracts_host_or_unknown() {
1144 assert_eq!(
1145 provider_host("https://attacker.example/v1"),
1146 "attacker.example"
1147 );
1148 assert_eq!(provider_host("http://127.0.0.1:8080"), "127.0.0.1");
1149 assert_eq!(provider_host("not a url"), "<unknown>");
1150 }
1151
1152 #[test]
1153 fn override_host_warning_is_deduped() {
1154 let key = unique_env("MERMAID_FACTORY_WARN_KEY");
1158 assert!(should_warn_once(&key), "first warn for a key must fire");
1159 assert!(
1160 !should_warn_once(&key),
1161 "subsequent warns for the same key must be suppressed"
1162 );
1163 }
1164
1165 #[test]
1169 fn custom_profile_is_memoized_per_key() {
1170 use crate::app::UserProviderConfig;
1171 let cfg = UserProviderConfig {
1172 base_url: Some("https://api.custom.test/v1".to_string()),
1173 api_key_env: Some("CUSTOM_KEY".to_string()),
1174 compat: Some("openai".to_string()),
1175 ..Default::default()
1176 };
1177 let a = user_profile_to_static("mermaid_test_customx", &cfg).unwrap();
1180 let b = user_profile_to_static("mermaid_test_customx", &cfg).unwrap();
1181 assert!(
1182 std::ptr::eq(a, b),
1183 "identical custom-provider inputs must reuse one leaked &'static profile"
1184 );
1185 assert_eq!(a.base_url, "https://api.custom.test/v1");
1186 assert_eq!(a.api_key_env, "CUSTOM_KEY");
1187
1188 let cfg2 = UserProviderConfig {
1190 base_url: Some("https://api.custom.test/v2".to_string()),
1191 ..cfg.clone()
1192 };
1193 let c = user_profile_to_static("mermaid_test_customx", &cfg2).unwrap();
1194 assert!(
1195 !std::ptr::eq(a, c),
1196 "a different base_url must leak a distinct profile"
1197 );
1198 }
1199}