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_api_key_with_fallback};
16
17const GEMINI_API_KEY_ENV: &str = "GOOGLE_API_KEY";
18const GEMINI_LEGACY_API_KEY_ENV: &str = "GEMINI_API_KEY";
19
20fn require_key(provider: &str, env_var: &str) -> Result<String> {
25 resolve_api_key(env_var, None).ok_or_else(|| {
26 ModelError::Authentication(format!("{} requires env var {}", provider, env_var))
27 })
28}
29
30fn require_key_with_fallback(
31 provider: &str,
32 env_var: &str,
33 fallback_env_var: &str,
34) -> Result<String> {
35 resolve_api_key_with_fallback(env_var, fallback_env_var, None).ok_or_else(|| {
36 ModelError::Authentication(format!(
37 "{} requires env var {} (or legacy {})",
38 provider, env_var, fallback_env_var
39 ))
40 })
41}
42
43use super::model::{
44 AnthropicProvider, GeminiProvider, ModelProvider, OllamaProvider, OpenAICompatProvider,
45};
46
47type ProviderCell = Arc<tokio::sync::OnceCell<Arc<dyn ModelProvider>>>;
51
52pub struct ProviderFactory {
56 config: Arc<Config>,
57 cache: Mutex<std::collections::HashMap<String, ProviderCell>>,
60}
61
62impl ProviderFactory {
63 pub fn new(config: Config) -> Self {
64 Self {
65 config: Arc::new(config),
66 cache: Mutex::new(std::collections::HashMap::new()),
67 }
68 }
69
70 pub fn config(&self) -> &Config {
71 &self.config
72 }
73
74 pub async fn resolve(&self, model_id: &str) -> Result<Arc<dyn ModelProvider>> {
78 let key = normalize_cache_key(model_id);
79 let cell = {
83 let mut cache = self.cache.lock().await;
84 Arc::clone(
85 cache
86 .entry(key)
87 .or_insert_with(|| Arc::new(tokio::sync::OnceCell::new())),
88 )
89 };
90 let provider = cell
91 .get_or_try_init(|| async {
92 let p = build_provider(&self.config, model_id).await?;
93 Ok::<Arc<dyn ModelProvider>, ModelError>(Arc::from(p))
94 })
95 .await?;
96 Ok(Arc::clone(provider))
97 }
98}
99
100async fn build_provider(config: &Config, model_id: &str) -> Result<Box<dyn ModelProvider>> {
108 let (provider, model_name) = parse_model_id(model_id);
109 let provider_lc = provider.to_lowercase();
110
111 if provider_lc == "ollama" {
114 let backend = ollama_backend_config(config);
115 let p = OllamaProvider::with_app_config(
116 model_name,
117 Arc::new(backend),
118 Arc::new(config.clone()),
119 )
120 .await?;
121 return Ok(Box::new(p));
122 }
123
124 if provider_lc == "anthropic" {
126 let user_cfg = config.providers.get("anthropic");
127 let base_url = resolve_overridable_base_url(
128 "anthropic",
129 user_cfg.and_then(|c| c.base_url.clone()),
130 "https://api.anthropic.com/v1",
131 )?;
132 let api_key_env = user_cfg
133 .and_then(|c| c.api_key_env.as_deref())
134 .unwrap_or("ANTHROPIC_API_KEY");
135 let api_key = require_key("anthropic", api_key_env)?;
136 let p = AnthropicProvider::new(api_key, model_name.to_string(), base_url)?;
137 return Ok(Box::new(p));
138 }
139
140 if provider_lc == "gemini" {
142 let user_cfg = config.providers.get("gemini");
143 let base_url = resolve_overridable_base_url(
144 "gemini",
145 user_cfg.and_then(|c| c.base_url.clone()),
146 "https://generativelanguage.googleapis.com/v1beta",
147 )?;
148 let api_key = match user_cfg.and_then(|c| c.api_key_env.as_deref()) {
149 Some(api_key_env) => require_key("gemini", api_key_env)?,
150 None => {
151 require_key_with_fallback("gemini", GEMINI_API_KEY_ENV, GEMINI_LEGACY_API_KEY_ENV)?
152 },
153 };
154 let p = GeminiProvider::new(api_key, model_name.to_string(), base_url)?;
155 return Ok(Box::new(p));
156 }
157
158 if let Some(profile) = lookup_provider(&provider_lc) {
160 let user_cfg = config.providers.get(&provider_lc);
161 let base_url = resolve_overridable_base_url(
162 &provider_lc,
163 user_cfg.and_then(|c| c.base_url.clone()),
164 profile.base_url,
165 )?;
166 let api_key_env = user_cfg
167 .and_then(|c| c.api_key_env.as_deref())
168 .unwrap_or(profile.api_key_env);
169 let api_key = require_key(&provider_lc, api_key_env)?;
170 let extra_headers = user_cfg
171 .map(|c| c.extra_headers.clone())
172 .unwrap_or_default();
173 let p = OpenAICompatProvider::new(
174 profile,
175 base_url,
176 api_key,
177 model_name.to_string(),
178 extra_headers,
179 )?;
180 return Ok(Box::new(p));
181 }
182
183 if let Some(user_cfg) = config.providers.get(&provider_lc)
186 && let Some(profile) = user_profile_to_static(&provider_lc, user_cfg)
187 {
188 let base_url = user_cfg.base_url.clone().ok_or_else(|| {
189 ModelError::InvalidRequest(format!(
190 "custom provider '{}' requires base_url in config",
191 provider_lc
192 ))
193 })?;
194 validate_provider_base_url(&base_url)?;
195 let api_key_env = user_cfg.api_key_env.as_deref().ok_or_else(|| {
196 ModelError::InvalidRequest(format!(
197 "custom provider '{}' requires api_key_env in config",
198 provider_lc
199 ))
200 })?;
201 let api_key = require_key(&provider_lc, api_key_env)?;
202 let p = OpenAICompatProvider::new(
203 profile,
204 base_url,
205 api_key,
206 model_name.to_string(),
207 user_cfg.extra_headers.clone(),
208 )?;
209 return Ok(Box::new(p));
210 }
211
212 Err(ModelError::InvalidRequest(format!(
213 "Unknown provider '{}' (model_id: {})",
214 provider, model_id
215 )))
216}
217
218fn normalize_cache_key(model_id: &str) -> String {
222 let (provider, model) = parse_model_id(model_id);
223 format!("{}/{}", provider.to_lowercase(), model)
224}
225
226fn parse_model_id(model_id: &str) -> (String, &str) {
229 match model_id.split_once('/') {
230 Some((p, m)) => (p.to_string(), m),
231 None => ("ollama".to_string(), model_id),
232 }
233}
234
235pub(crate) fn ollama_backend_config(config: &Config) -> BackendConfig {
236 BackendConfig {
237 ollama_url: format!("{}:{}", config.ollama.host, config.ollama.port),
240 max_idle_per_host: 10,
241 timeout_secs: 10,
242 }
243}
244
245static PROFILE_CACHE: std::sync::LazyLock<
249 std::sync::Mutex<std::collections::HashMap<String, &'static crate::models::ProviderProfile>>,
250> = std::sync::LazyLock::new(|| std::sync::Mutex::new(std::collections::HashMap::new()));
251
252fn user_profile_to_static(
266 name: &str,
267 user_cfg: &crate::app::UserProviderConfig,
268) -> Option<&'static crate::models::ProviderProfile> {
269 use crate::models::{ProviderProfile, ReasoningExtraction, ReasoningStrategy};
270
271 let compat = user_cfg.compat.as_deref().unwrap_or("openai");
272 let base_url = user_cfg.base_url.clone().unwrap_or_default();
273 let api_key_env = user_cfg.api_key_env.clone().unwrap_or_default();
274
275 let cache_key = format!("{name}\u{0}{base_url}\u{0}{api_key_env}\u{0}{compat}");
278
279 let mut cache = PROFILE_CACHE
280 .lock()
281 .unwrap_or_else(|poisoned| poisoned.into_inner());
282 if let Some(profile) = cache.get(cache_key.as_str()) {
283 return Some(*profile);
286 }
287
288 let strategy = match compat {
289 "openai" => ReasoningStrategy::None,
290 "openai-effort" => ReasoningStrategy::Effort,
291 "openrouter" => ReasoningStrategy::OpenRouterShape,
292 _ => ReasoningStrategy::None,
293 };
294
295 let profile = Box::new(ProviderProfile {
296 name: Box::leak(name.to_string().into_boxed_str()),
297 base_url: Box::leak(base_url.into_boxed_str()),
298 api_key_env: Box::leak(api_key_env.into_boxed_str()),
299 extra_headers: &[],
300 reasoning_strategy: strategy,
301 reasoning_extraction: ReasoningExtraction::None,
302 max_tokens_param: crate::models::MaxTokensParam::MaxTokens,
303 disable_parallel_tool_calls_for: &[],
304 });
305 let leaked: &'static ProviderProfile = Box::leak(profile);
306 cache.insert(cache_key, leaked);
307 Some(leaked)
308}
309
310fn validate_provider_base_url(url: &str) -> Result<()> {
316 let parsed = reqwest::Url::parse(url).map_err(|e| {
317 ModelError::InvalidRequest(format!("invalid provider base_url '{url}': {e}"))
318 })?;
319 match parsed.scheme() {
320 "https" => Ok(()),
321 "http"
326 if crate::utils::classify_host(parsed.host_str().unwrap_or_default()).is_loopback() =>
327 {
328 Ok(())
329 },
330 "http" => Err(ModelError::InvalidRequest(format!(
331 "provider base_url '{url}' uses http:// to a non-loopback host — refusing to send the \
332 API key in cleartext. Use https, or http://localhost for a local server."
333 ))),
334 other => Err(ModelError::InvalidRequest(format!(
335 "provider base_url '{url}' has unsupported scheme '{other}' (use http or https)"
336 ))),
337 }
338}
339
340fn resolve_overridable_base_url(
356 provider: &str,
357 override_url: Option<String>,
358 default_url: &str,
359) -> Result<String> {
360 match override_url {
361 Some(url) => {
362 validate_provider_base_url(&url)?;
363 warn_overridden_provider_host(provider, &url);
364 Ok(url)
365 },
366 None => Ok(default_url.to_string()),
367 }
368}
369
370static WARNED_OVERRIDE_HOSTS: std::sync::LazyLock<
374 std::sync::Mutex<std::collections::HashSet<String>>,
375> = std::sync::LazyLock::new(|| std::sync::Mutex::new(std::collections::HashSet::new()));
376
377fn warn_overridden_provider_host(provider: &str, base_url: &str) {
381 let host = provider_host(base_url);
382 if should_warn_once(&format!("{provider}@{host}")) {
383 tracing::warn!(
384 "built-in provider '{}' base_url overridden in config: the {} API key will be sent to \
385 host '{}' instead of the trusted default endpoint",
386 provider,
387 provider,
388 host
389 );
390 }
391}
392
393fn provider_host(base_url: &str) -> String {
396 reqwest::Url::parse(base_url)
397 .ok()
398 .and_then(|u| u.host_str().map(str::to_string))
399 .unwrap_or_else(|| "<unknown>".to_string())
400}
401
402fn should_warn_once(key: &str) -> bool {
405 let mut warned = WARNED_OVERRIDE_HOSTS
406 .lock()
407 .unwrap_or_else(|poisoned| poisoned.into_inner());
408 warned.insert(key.to_string())
409}
410
411#[cfg(test)]
412mod tests {
413 use super::*;
414
415 #[test]
416 fn base_url_requires_https_for_remote_hosts() {
417 assert!(validate_provider_base_url("http://api.example.com/v1").is_err());
419 assert!(validate_provider_base_url("ftp://example.com").is_err());
420 assert!(validate_provider_base_url("https://api.example.com/v1").is_ok());
422 assert!(validate_provider_base_url("http://localhost:11434/v1").is_ok());
423 assert!(validate_provider_base_url("http://127.0.0.1:8000").is_ok());
424 assert!(validate_provider_base_url("http://[::1]:8000").is_ok());
425 assert!(validate_provider_base_url("http://192.168.1.5:8080").is_err());
428 assert!(validate_provider_base_url("http://169.254.169.254").is_err());
429 }
430 use std::sync::atomic::{AtomicUsize, Ordering};
431
432 fn unique_env(prefix: &str) -> String {
433 static N: AtomicUsize = AtomicUsize::new(0);
434 format!(
435 "{}_{}_{}",
436 prefix,
437 std::process::id(),
438 N.fetch_add(1, Ordering::SeqCst)
439 )
440 }
441
442 #[test]
443 fn parse_bare_name_defaults_to_ollama() {
444 let (p, m) = parse_model_id("qwen3-coder:30b");
445 assert_eq!(p, "ollama");
446 assert_eq!(m, "qwen3-coder:30b");
447 }
448
449 #[test]
450 fn parse_prefixed() {
451 let (p, m) = parse_model_id("anthropic/claude-opus-4-7");
452 assert_eq!(p, "anthropic");
453 assert_eq!(m, "claude-opus-4-7");
454 }
455
456 #[test]
457 fn gemini_key_resolution_accepts_legacy_fallback() {
458 let primary = unique_env("MERMAID_FACTORY_GEMINI_PRIMARY");
459 let legacy = unique_env("MERMAID_FACTORY_GEMINI_LEGACY");
460 temp_env::with_vars(
461 [(primary.as_str(), None), (legacy.as_str(), Some("legacy"))],
462 || {
463 let resolved = require_key_with_fallback("gemini", &primary, &legacy)
464 .expect("legacy fallback should resolve");
465 assert_eq!(resolved, "legacy");
466 },
467 );
468 }
469
470 #[test]
471 fn gemini_key_resolution_prefers_google_primary() {
472 let primary = unique_env("MERMAID_FACTORY_GEMINI_PRIMARY2");
473 let legacy = unique_env("MERMAID_FACTORY_GEMINI_LEGACY2");
474 temp_env::with_vars(
475 [
476 (primary.as_str(), Some("google")),
477 (legacy.as_str(), Some("legacy")),
478 ],
479 || {
480 let resolved = require_key_with_fallback("gemini", &primary, &legacy)
481 .expect("primary should resolve");
482 assert_eq!(resolved, "google");
483 },
484 );
485 }
486
487 #[tokio::test]
488 async fn factory_reports_unknown_provider_clearly() {
489 let cfg = Config::default();
490 let f = ProviderFactory::new(cfg);
491 match f.resolve("totally-made-up/model").await {
492 Ok(_) => panic!("expected error"),
493 Err(e) => {
494 let msg = format!("{}", e);
495 assert!(
496 msg.contains("totally-made-up") || msg.contains("Unknown provider"),
497 "error message: {}",
498 msg
499 );
500 },
501 }
502 }
503
504 #[test]
505 fn normalize_cache_key_lowercases_provider_only() {
506 assert_eq!(
508 normalize_cache_key("Anthropic/Claude-X"),
509 "anthropic/Claude-X"
510 );
511 assert_eq!(
512 normalize_cache_key("anthropic/Claude-X"),
513 "anthropic/Claude-X"
514 );
515 assert_eq!(normalize_cache_key("qwen3:30b"), "ollama/qwen3:30b");
517 }
518
519 #[tokio::test]
520 async fn resolve_is_single_flight_and_cached() {
521 let f = ProviderFactory::new(Config::default());
525 let (a, b) = tokio::join!(
526 f.resolve("ollama/test-model"),
527 f.resolve("Ollama/test-model"),
528 );
529 let a = a.expect("resolve a");
530 let b = b.expect("resolve b");
531 assert!(
532 Arc::ptr_eq(&a, &b),
533 "expected one cached provider for casing variants + concurrent resolve"
534 );
535 }
536
537 #[test]
541 fn builtin_base_url_override_validated_and_resolved() {
542 assert_eq!(
544 resolve_overridable_base_url("anthropic", None, "https://api.anthropic.com/v1")
545 .unwrap(),
546 "https://api.anthropic.com/v1"
547 );
548 assert_eq!(
550 resolve_overridable_base_url(
551 "anthropic",
552 Some("https://proxy.internal/v1".to_string()),
553 "https://api.anthropic.com/v1",
554 )
555 .unwrap(),
556 "https://proxy.internal/v1"
557 );
558 assert!(
561 resolve_overridable_base_url(
562 "anthropic",
563 Some("http://attacker.example/v1".to_string()),
564 "https://api.anthropic.com/v1",
565 )
566 .is_err()
567 );
568 assert!(
570 resolve_overridable_base_url(
571 "openai",
572 Some("http://localhost:8080/v1".to_string()),
573 "https://api.openai.com/v1",
574 )
575 .is_ok()
576 );
577 }
578
579 #[test]
580 fn provider_host_extracts_host_or_unknown() {
581 assert_eq!(
582 provider_host("https://attacker.example/v1"),
583 "attacker.example"
584 );
585 assert_eq!(provider_host("http://127.0.0.1:8080"), "127.0.0.1");
586 assert_eq!(provider_host("not a url"), "<unknown>");
587 }
588
589 #[test]
590 fn override_host_warning_is_deduped() {
591 let key = unique_env("MERMAID_FACTORY_WARN_KEY");
595 assert!(should_warn_once(&key), "first warn for a key must fire");
596 assert!(
597 !should_warn_once(&key),
598 "subsequent warns for the same key must be suppressed"
599 );
600 }
601
602 #[test]
606 fn custom_profile_is_memoized_per_key() {
607 use crate::app::UserProviderConfig;
608 let cfg = UserProviderConfig {
609 base_url: Some("https://api.custom.test/v1".to_string()),
610 api_key_env: Some("CUSTOM_KEY".to_string()),
611 compat: Some("openai".to_string()),
612 ..Default::default()
613 };
614 let a = user_profile_to_static("mermaid_test_customx", &cfg).unwrap();
617 let b = user_profile_to_static("mermaid_test_customx", &cfg).unwrap();
618 assert!(
619 std::ptr::eq(a, b),
620 "identical custom-provider inputs must reuse one leaked &'static profile"
621 );
622 assert_eq!(a.base_url, "https://api.custom.test/v1");
623 assert_eq!(a.api_key_env, "CUSTOM_KEY");
624
625 let cfg2 = UserProviderConfig {
627 base_url: Some("https://api.custom.test/v2".to_string()),
628 ..cfg.clone()
629 };
630 let c = user_profile_to_static("mermaid_test_customx", &cfg2).unwrap();
631 assert!(
632 !std::ptr::eq(a, c),
633 "a different base_url must leak a distinct profile"
634 );
635 }
636}