Skip to main content

mermaid_cli/providers/
factory.rs

1//! Runtime provider construction.
2//!
3//! `ProviderFactory` turns `(Config, model_id)` into the right
4//! `Arc<dyn ModelProvider>`. The effect runner holds one of these
5//! and asks it to build a provider the first time a new model is
6//! referenced; subsequent lookups hit the cache.
7
8use 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
17const GEMINI_API_KEY_ENV: &str = "GOOGLE_API_KEY";
18const GEMINI_LEGACY_API_KEY_ENV: &str = "GEMINI_API_KEY";
19
20/// Resolve an API key (env first, then the OS keyring via
21/// `mermaid login <provider>`) or return a clear `ModelError`. A
22/// per-provider `override_env` is authoritative — no keyring fallback.
23fn require_key(provider: &str, default_env: &str, override_env: Option<&str>) -> Result<String> {
24    resolve_provider_key(provider, default_env, override_env).ok_or_else(|| {
25        let env = override_env.unwrap_or(default_env);
26        ModelError::Authentication(format!(
27            "{provider} requires env var {env} (or `mermaid login {provider}`)"
28        ))
29    })
30}
31
32fn require_key_with_fallback(
33    provider: &str,
34    env_var: &str,
35    fallback_env_var: &str,
36) -> Result<String> {
37    resolve_provider_key_with_fallback(provider, env_var, fallback_env_var, None).ok_or_else(|| {
38        ModelError::Authentication(format!(
39            "{provider} requires env var {env_var} (or legacy {fallback_env_var}, or \
40                 `mermaid login {provider}`)"
41        ))
42    })
43}
44
45/// Resolve an API key, or `None` when the key is absent AND the endpoint is a
46/// loopback/LAN host — local OpenAI-compatible servers (llama.cpp, vLLM, LM
47/// Studio) legitimately run without auth. A missing key for a public endpoint
48/// is a hard error carrying the provider's `hint` so the message is actionable.
49fn resolve_optional_key(
50    provider: &str,
51    default_env: &str,
52    override_env: Option<&str>,
53    base_url: &str,
54    hint: Option<&str>,
55) -> Result<Option<String>> {
56    if let Some(key) = resolve_provider_key(provider, default_env, override_env) {
57        return Ok(Some(key));
58    }
59    if base_url_is_local(base_url) {
60        return Ok(None);
61    }
62    let env = override_env.unwrap_or(default_env);
63    let mut msg = format!("{provider} requires env var {env} (or `mermaid login {provider}`)");
64    if let Some(h) = hint {
65        msg.push_str(" — ");
66        msg.push_str(h);
67    }
68    Err(ModelError::Authentication(msg))
69}
70
71/// True when `base_url`'s host is a loopback or private/LAN address — where a
72/// local model server legitimately runs without auth.
73fn base_url_is_local(base_url: &str) -> bool {
74    reqwest::Url::parse(base_url)
75        .ok()
76        .and_then(|u| {
77            u.host_str()
78                .map(|h| crate::utils::classify_host(h).is_internal())
79        })
80        .unwrap_or(false)
81}
82
83/// The header set sent on every request: the profile's static `extra_headers`
84/// first, then user `extra_headers` overrides, then env-sourced `env_headers`
85/// (header name -> env var, resolved now since env is process-static). This
86/// restores the profile's static headers (e.g. OpenRouter analytics) that the
87/// prior registry path dropped.
88fn merged_headers(
89    profile: &crate::models::ProviderProfile,
90    user_cfg: Option<&crate::app::UserProviderConfig>,
91) -> std::collections::HashMap<String, String> {
92    let mut headers: std::collections::HashMap<String, String> = profile
93        .extra_headers
94        .iter()
95        .map(|(k, v)| ((*k).to_string(), (*v).to_string()))
96        .collect();
97    if let Some(cfg) = user_cfg {
98        for (k, v) in &cfg.extra_headers {
99            headers.insert(k.clone(), v.clone());
100        }
101        for (header, env_var) in &cfg.env_headers {
102            if let Ok(val) = std::env::var(env_var) {
103                headers.insert(header.clone(), val);
104            }
105        }
106    }
107    headers
108}
109
110use super::model::{
111    AnthropicProvider, GeminiProvider, MetaProvider, ModelProvider, OllamaProvider,
112    OpenAICompatProvider,
113};
114
115/// A lazily-built, shared provider. `OnceCell` gives single-flight construction
116/// (#84); the outer `Arc` lets `resolve` clone the cell out from under the cache
117/// lock and initialize it without holding the lock across the build.
118type ProviderCell = Arc<tokio::sync::OnceCell<Arc<dyn ModelProvider>>>;
119
120/// Per-process provider cache. Providers are expensive to construct
121/// (HTTP client, connection pool, capability lookup) so the effect
122/// runner asks for them lazily and reuses across turns.
123pub struct ProviderFactory {
124    config: Arc<Config>,
125    /// Per-key cache of built providers, keyed by normalized model id (#83). The
126    /// `Mutex` is only held to get-or-insert a cell, never across the build.
127    cache: Mutex<std::collections::HashMap<String, ProviderCell>>,
128}
129
130impl ProviderFactory {
131    pub fn new(config: Config) -> Self {
132        Self {
133            config: Arc::new(config),
134            cache: Mutex::new(std::collections::HashMap::new()),
135        }
136    }
137
138    pub fn config(&self) -> &Config {
139        &self.config
140    }
141
142    /// Resolve (or lazily construct) a provider for the given model
143    /// ID. Hits the cache on the second and subsequent calls for the
144    /// same ID.
145    pub async fn resolve(&self, model_id: &str) -> Result<Arc<dyn ModelProvider>> {
146        let key = normalize_cache_key(model_id);
147        // Get-or-insert the per-key cell under a brief lock, then initialize it
148        // exactly once outside the lock (#84). A failed build isn't cached, so a
149        // transient error can be retried on the next call.
150        let cell = {
151            let mut cache = self.cache.lock().await;
152            Arc::clone(
153                cache
154                    .entry(key)
155                    .or_insert_with(|| Arc::new(tokio::sync::OnceCell::new())),
156            )
157        };
158        let provider = cell
159            .get_or_try_init(|| async {
160                let p = build_provider(&self.config, model_id).await?;
161                Ok::<Arc<dyn ModelProvider>, ModelError>(Arc::from(p))
162            })
163            .await?;
164        Ok(Arc::clone(provider))
165    }
166}
167
168/// Build a provider for the given `model_id`:
169///   1. `ollama/<model>` → OllamaProvider.
170///   2. `anthropic/<model>` → AnthropicProvider.
171///   3. `gemini/<model>` → GeminiProvider.
172///   4. `meta/<model>` → MetaProvider using the Responses API.
173///   5. Other builtin providers (openai, openrouter, groq, …) → OpenAICompatProvider.
174///   6. User-defined `[providers.<name>]` → custom OpenAICompatProvider.
175///   7. Bare model name → OllamaProvider.
176async fn build_provider(config: &Config, model_id: &str) -> Result<Box<dyn ModelProvider>> {
177    let (provider, model_name) = parse_model_id(model_id);
178    let provider_lc = provider.to_lowercase();
179
180    // 1. Ollama (and bare names). F11: pass Arc<Config> so the wrapper
181    // can forward Ollama hardware options to the adapter.
182    if provider_lc == "ollama" {
183        let backend = ollama_backend_config(config);
184        let p = OllamaProvider::with_app_config(
185            model_name,
186            Arc::new(backend),
187            Arc::new(config.clone()),
188        )
189        .await?;
190        return Ok(Box::new(p));
191    }
192
193    // 2. Anthropic — bespoke API shape.
194    if provider_lc == "anthropic" {
195        let user_cfg = config.providers.get("anthropic");
196        let base_url = resolve_overridable_base_url(
197            "anthropic",
198            user_cfg.and_then(|c| c.base_url.clone()),
199            "https://api.anthropic.com/v1",
200        )?;
201        let api_key = require_key(
202            "anthropic",
203            "ANTHROPIC_API_KEY",
204            user_cfg.and_then(|c| c.api_key_env.as_deref()),
205        )?;
206        let p = AnthropicProvider::new(api_key, model_name.to_string(), base_url)?;
207        return Ok(Box::new(p));
208    }
209
210    // 3. Gemini — GCP AI Studio shape.
211    if provider_lc == "gemini" {
212        let user_cfg = config.providers.get("gemini");
213        let base_url = resolve_overridable_base_url(
214            "gemini",
215            user_cfg.and_then(|c| c.base_url.clone()),
216            "https://generativelanguage.googleapis.com/v1beta",
217        )?;
218        let api_key = match user_cfg.and_then(|c| c.api_key_env.as_deref()) {
219            Some(api_key_env) => require_key("gemini", GEMINI_API_KEY_ENV, Some(api_key_env))?,
220            None => {
221                require_key_with_fallback("gemini", GEMINI_API_KEY_ENV, GEMINI_LEGACY_API_KEY_ENV)?
222            },
223        };
224        let p = GeminiProvider::new(api_key, model_name.to_string(), base_url)?;
225        return Ok(Box::new(p));
226    }
227
228    // 4. Meta — Responses is required for encrypted reasoning continuity across
229    // Mermaid's tool loop even though Meta also exposes Chat Completions.
230    if provider_lc == "meta" {
231        let user_cfg = config.providers.get("meta");
232        let base_url = resolve_overridable_base_url(
233            "meta",
234            user_cfg.and_then(|cfg| cfg.base_url.clone()),
235            super::model::meta::DEFAULT_BASE_URL,
236        )?;
237        let api_key = require_key(
238            "meta",
239            super::model::meta::DEFAULT_API_KEY_ENV,
240            user_cfg.and_then(|cfg| cfg.api_key_env.as_deref()),
241        )?;
242        let mut extra_headers = std::collections::HashMap::new();
243        if let Some(cfg) = user_cfg {
244            extra_headers.extend(cfg.extra_headers.clone());
245            for (header, env_var) in &cfg.env_headers {
246                if let Ok(value) = std::env::var(env_var) {
247                    extra_headers.insert(header.clone(), value);
248                }
249            }
250        }
251        let p = MetaProvider::new(api_key, model_name.to_string(), base_url, extra_headers)?;
252        return Ok(Box::new(p));
253    }
254
255    // 4.5. Cloudflare Workers AI — OpenAI-compatible, but the endpoint URL embeds a
256    // per-account id, so the base_url is synthesized at runtime from
257    // CLOUDFLARE_ACCOUNT_ID (or a full [providers.cloudflare].base_url override, e.g.
258    // AI Gateway). Must precede the generic registry branch below, which would
259    // otherwise route it through the placeholder profile.base_url.
260    if provider_lc == "cloudflare" {
261        let user_cfg = config.providers.get("cloudflare");
262        let profile = lookup_provider("cloudflare").expect("cloudflare is in the registry");
263        let override_env = user_cfg.and_then(|c| c.api_key_env.as_deref());
264        let api_key_env = override_env.unwrap_or(profile.api_key_env);
265        let base_url = match user_cfg.and_then(|c| c.base_url.clone()) {
266            // Override present (AI Gateway / proxy): validate + warn like any built-in
267            // override. The account id isn't needed in this case.
268            Some(url) => {
269                validate_provider_base_url(&url)?;
270                warn_overridden_provider_host("cloudflare", &url);
271                url
272            },
273            // Standard path: synthesize the account-scoped endpoint from env. A fresh
274            // setup missing the token as well gets one error naming both vars, not
275            // two fix-and-retry round-trips.
276            None => match require_cloudflare_account_id() {
277                Ok(id) => cloudflare_base_url(&id),
278                Err(_)
279                    if resolve_provider_key("cloudflare", profile.api_key_env, override_env)
280                        .is_none() =>
281                {
282                    return Err(ModelError::Authentication(format!(
283                        "cloudflare requires env vars {api_key_env} and CLOUDFLARE_ACCOUNT_ID — \
284                         create a token at https://dash.cloudflare.com/profile/api-tokens; the \
285                         account id is on your Cloudflare dashboard (or set \
286                         [providers.cloudflare].base_url)"
287                    )));
288                },
289                Err(e) => return Err(e),
290            },
291        };
292        let api_key = resolve_optional_key(
293            &provider_lc,
294            profile.api_key_env,
295            override_env,
296            &base_url,
297            profile.key_hint,
298        )?;
299        let extra_headers = merged_headers(profile, user_cfg);
300        let p = OpenAICompatProvider::new(
301            profile,
302            base_url,
303            api_key,
304            model_name.to_string(),
305            extra_headers,
306        )?;
307        return Ok(Box::new(p));
308    }
309
310    // 5 + 6. OpenAI-compatible registry or user-custom.
311    if let Some(profile) = lookup_provider(&provider_lc) {
312        let user_cfg = config.providers.get(&provider_lc);
313        let base_url = resolve_overridable_base_url(
314            &provider_lc,
315            user_cfg.and_then(|c| c.base_url.clone()),
316            profile.base_url,
317        )?;
318        // Keyless when the key is unset and the endpoint is local (a registry
319        // provider pointed at a loopback/LAN base_url); otherwise a clear,
320        // hint-carrying auth error.
321        let api_key = resolve_optional_key(
322            &provider_lc,
323            profile.api_key_env,
324            user_cfg.and_then(|c| c.api_key_env.as_deref()),
325            &base_url,
326            profile.key_hint,
327        )?;
328        let extra_headers = merged_headers(profile, user_cfg);
329        let p = OpenAICompatProvider::new(
330            profile,
331            base_url,
332            api_key,
333            model_name.to_string(),
334            extra_headers,
335        )?;
336        return Ok(Box::new(p));
337    }
338
339    // User-custom: no registry entry, but the user has [providers.<name>]
340    // in config with a declared `compat` field.
341    if let Some(user_cfg) = config.providers.get(&provider_lc)
342        && let Some(profile) = user_profile_to_static(&provider_lc, user_cfg)
343    {
344        let base_url = user_cfg.base_url.clone().ok_or_else(|| {
345            ModelError::InvalidRequest(format!(
346                "custom provider '{}' requires base_url in config",
347                provider_lc
348            ))
349        })?;
350        // api_key_env is optional: a local (loopback/LAN) endpoint may run
351        // keyless. When a key IS used, harden the URL so it can't be sent in
352        // cleartext; a keyless endpoint must be local (no secret to leak).
353        // For a CUSTOM provider its api_key_env is the default (not an
354        // override of a registry default), so the keyring may fill the gap;
355        // a stored key alone also works with no api_key_env at all.
356        let api_key_env = user_cfg.api_key_env.as_deref();
357        let resolved = match api_key_env {
358            Some(env) => resolve_provider_key(&provider_lc, env, None),
359            None => crate::utils::default_store().get(&provider_lc),
360        };
361        let api_key = match resolved {
362            Some(key) => {
363                validate_provider_base_url(&base_url)?;
364                Some(key)
365            },
366            None if base_url_is_local(&base_url) => None,
367            None => {
368                let reason = match api_key_env {
369                    Some(env) => format!(
370                        "requires env var {env} (or `mermaid login {provider_lc}`, or a \
371                         loopback/LAN base_url)"
372                    ),
373                    None => "requires api_key_env, or a loopback/LAN base_url".to_string(),
374                };
375                return Err(ModelError::Authentication(format!(
376                    "custom provider '{provider_lc}' {reason}"
377                )));
378            },
379        };
380        let extra_headers = merged_headers(profile, Some(user_cfg));
381        let p = OpenAICompatProvider::new(
382            profile,
383            base_url,
384            api_key,
385            model_name.to_string(),
386            extra_headers,
387        )?;
388        return Ok(Box::new(p));
389    }
390
391    Err(ModelError::InvalidRequest(format!(
392        "Unknown provider '{}' (model_id: {})",
393        provider, model_id
394    )))
395}
396
397/// Cache key for a model id: the provider segment lowercased (matching
398/// `build_provider`'s own normalization) so `Anthropic/x` and `anthropic/x`
399/// resolve to one cached provider instead of building two identical ones (#83).
400fn normalize_cache_key(model_id: &str) -> String {
401    let (provider, model) = parse_model_id(model_id);
402    format!("{}/{}", provider.to_lowercase(), model)
403}
404
405/// Parse `provider/model` → `(provider, model)`. Bare strings are
406/// Ollama by convention.
407fn parse_model_id(model_id: &str) -> (String, &str) {
408    match model_id.split_once('/') {
409        Some((p, m)) => (p.to_string(), m),
410        None => ("ollama".to_string(), model_id),
411    }
412}
413
414pub(crate) fn ollama_backend_config(config: &Config) -> BackendConfig {
415    BackendConfig {
416        // Scheme-less: `normalize_url` in the adapter picks http (loopback/LAN)
417        // vs https (public) by host class (#86).
418        ollama_url: format!("{}:{}", config.ollama.host, config.ollama.port),
419        max_idle_per_host: 10,
420        timeout_secs: 10,
421        ollama_autostart: config.ollama.auto_start,
422    }
423}
424
425/// Process-wide memo of leaked custom-provider profiles, keyed by the
426/// profile-determining inputs (see [`user_profile_to_static`]). Guarantees each
427/// distinct custom provider leaks its `&'static ProviderProfile` at most once.
428static PROFILE_CACHE: std::sync::LazyLock<
429    std::sync::Mutex<std::collections::HashMap<String, &'static crate::models::ProviderProfile>>,
430> = std::sync::LazyLock::new(|| std::sync::Mutex::new(std::collections::HashMap::new()));
431
432/// Convert a user-defined `[providers.<name>]` entry into a `&'static
433/// ProviderProfile`. `ProviderProfile`'s fields are `&'static` (tied to the
434/// registry constants), so a custom provider needs a leaked, owned copy to
435/// participate without redesigning the profile type.
436///
437/// The leak is memoized (F67). `build_provider` runs once per distinct *model
438/// id*, so without a cache this leaked a fresh profile for every custom
439/// `provider/model` pair — a permanent, per-distinct-model_id growth, not the
440/// "0-3" the old comment claimed. The profile's content depends only on
441/// (provider name, base_url, api_key_env, compat) and NOT on the model, so we
442/// key the cache on exactly those and leak at most once per distinct
443/// combination; repeated resolves (including different models of the same custom
444/// provider) reuse the same `&'static`.
445fn user_profile_to_static(
446    name: &str,
447    user_cfg: &crate::app::UserProviderConfig,
448) -> Option<&'static crate::models::ProviderProfile> {
449    use crate::models::{ProviderProfile, ReasoningExtraction, ReasoningStrategy};
450
451    let compat = user_cfg.compat.as_deref().unwrap_or("openai");
452    let base_url = user_cfg.base_url.clone().unwrap_or_default();
453    let api_key_env = user_cfg.api_key_env.clone().unwrap_or_default();
454
455    // Stable key over exactly the fields baked into the profile below. NUL
456    // separates components so distinct inputs can't collide into one key.
457    let cache_key = format!("{name}\u{0}{base_url}\u{0}{api_key_env}\u{0}{compat}");
458
459    let mut cache = PROFILE_CACHE
460        .lock()
461        .unwrap_or_else(|poisoned| poisoned.into_inner());
462    if let Some(profile) = cache.get(cache_key.as_str()) {
463        // `&'static ProviderProfile` is Copy, so this hands back the same leaked
464        // allocation as the first resolve — no new leak.
465        return Some(*profile);
466    }
467
468    let strategy = match compat {
469        "openai" => ReasoningStrategy::None,
470        "openai-effort" => ReasoningStrategy::Effort,
471        "openrouter" => ReasoningStrategy::OpenRouterShape,
472        _ => ReasoningStrategy::None,
473    };
474
475    let profile = Box::new(ProviderProfile {
476        name: Box::leak(name.to_string().into_boxed_str()),
477        base_url: Box::leak(base_url.into_boxed_str()),
478        api_key_env: Box::leak(api_key_env.into_boxed_str()),
479        key_hint: None,
480        extra_headers: &[],
481        reasoning_strategy: strategy,
482        reasoning_extraction: ReasoningExtraction::None,
483        max_tokens_param: crate::models::MaxTokensParam::MaxTokens,
484        disable_parallel_tool_calls_for: &[],
485    });
486    let leaked: &'static ProviderProfile = Box::leak(profile);
487    cache.insert(cache_key, leaked);
488    Some(leaked)
489}
490
491/// Build Cloudflare Workers AI's account-scoped OpenAI-compatible base URL. The
492/// adapter appends `/chat/completions`, so this ends at `/ai/v1`. Kept pure and
493/// separate so it's directly unit-testable (a built provider exposes no base_url
494/// getter).
495fn cloudflare_base_url(account_id: &str) -> String {
496    format!(
497        "https://api.cloudflare.com/client/v4/accounts/{}/ai/v1",
498        account_id.trim()
499    )
500}
501
502/// Best-effort `base_url` for *discovery* surfaces (`doctor`'s provider list,
503/// the `/models` probe) — chat requests never use this; `build_provider` has its
504/// own resolution. A user override wins for any provider; cloudflare synthesizes
505/// its account-scoped URL from `CLOUDFLARE_ACCOUNT_ID` and yields `None` when
506/// the var is unset (there is no real endpoint then — probing the registry
507/// placeholder would just guarantee a 404); everything else uses the profile's
508/// static default.
509pub(crate) fn discovery_base_url(
510    profile: &crate::models::ProviderProfile,
511    override_url: Option<String>,
512) -> Option<String> {
513    if override_url.is_some() {
514        return override_url;
515    }
516    if profile.name == "cloudflare" {
517        return require_cloudflare_account_id()
518            .ok()
519            .map(|id| cloudflare_base_url(&id));
520    }
521    Some(profile.base_url.to_string())
522}
523
524/// Resolve the Cloudflare account id from `CLOUDFLARE_ACCOUNT_ID` (trimmed,
525/// non-empty), or a clear actionable error. It isn't a secret, but it's required
526/// to construct the account-scoped Workers AI endpoint; `resolve_api_key` is
527/// reused only for its empty→None handling.
528fn require_cloudflare_account_id() -> Result<String> {
529    resolve_api_key("CLOUDFLARE_ACCOUNT_ID", None)
530        .map(|s| s.trim().to_string())
531        .filter(|s| !s.is_empty())
532        .ok_or_else(|| {
533            ModelError::Authentication(
534                "cloudflare requires env var CLOUDFLARE_ACCOUNT_ID (your Cloudflare account id) — \
535                 find it on your Cloudflare dashboard, or set [providers.cloudflare].base_url"
536                    .to_string(),
537            )
538        })
539}
540
541/// Validate a provider `base_url` before it's handed an API key. Requires
542/// http/https, and **requires https for any non-loopback, non-private host** —
543/// a typo'd or hostile `http://` endpoint would otherwise receive the bearer
544/// key in cleartext. Plain http stays allowed for loopback / RFC-1918 hosts so
545/// local model servers (Ollama, vLLM) keep working.
546fn validate_provider_base_url(url: &str) -> Result<()> {
547    let parsed = reqwest::Url::parse(url).map_err(|e| {
548        ModelError::InvalidRequest(format!("invalid provider base_url '{url}': {e}"))
549    })?;
550    match parsed.scheme() {
551        "https" => Ok(()),
552        // Plaintext http is only safe to LOOPBACK: a key sent over http to any
553        // other host — including a LAN/private one — crosses the wire in
554        // cleartext. Every caller here is a key-bearing provider; keyless local
555        // Ollama uses a separate path that doesn't validate.
556        "http"
557            if crate::utils::classify_host(parsed.host_str().unwrap_or_default()).is_loopback() =>
558        {
559            Ok(())
560        },
561        "http" => Err(ModelError::InvalidRequest(format!(
562            "provider base_url '{url}' uses http:// to a non-loopback host — refusing to send the \
563             API key in cleartext. Use https, or http://localhost for a local server."
564        ))),
565        other => Err(ModelError::InvalidRequest(format!(
566            "provider base_url '{url}' has unsupported scheme '{other}' (use http or https)"
567        ))),
568    }
569}
570
571/// Resolve a *built-in* provider's `base_url`, honoring a user override but
572/// hardening it (F66). Built-in providers (anthropic, gemini, and the
573/// registry-backed OpenAI-compatible ones) ship a trusted default endpoint; a
574/// `[providers.<name>] base_url` override redirects that provider's API key to a
575/// host the user chose. The override lives in the user's own config, so we allow
576/// it — but we also:
577///   * validate the scheme via [`validate_provider_base_url`], which already
578///     requires https for any non-loopback host, so an `http://` override can't
579///     leak the key in cleartext (http stays allowed only for
580///     localhost/127.0.0.1/::1 local dev), and
581///   * emit a one-time warning naming the host the key will be sent to, so a
582///     redirect to an unexpected host (e.g. `https://attacker.example`) is
583///     visible rather than silent.
584///
585/// With no override the trusted default is returned unchanged (no warning).
586fn resolve_overridable_base_url(
587    provider: &str,
588    override_url: Option<String>,
589    default_url: &str,
590) -> Result<String> {
591    match override_url {
592        Some(url) => {
593            validate_provider_base_url(&url)?;
594            warn_overridden_provider_host(provider, &url);
595            Ok(url)
596        },
597        None => Ok(default_url.to_string()),
598    }
599}
600
601/// Hosts already warned about (per provider) for a built-in `base_url` override,
602/// so the F66 warning fires once per process rather than on every `resolve`.
603/// Keyed by `"<provider>@<host>"`.
604static WARNED_OVERRIDE_HOSTS: std::sync::LazyLock<
605    std::sync::Mutex<std::collections::HashSet<String>>,
606> = std::sync::LazyLock::new(|| std::sync::Mutex::new(std::collections::HashSet::new()));
607
608/// Emit a one-time `tracing::warn!` (deduped per provider+host) naming the host a
609/// built-in provider's API key will be sent to because its trusted default
610/// `base_url` was overridden in config (F66).
611fn warn_overridden_provider_host(provider: &str, base_url: &str) {
612    let host = provider_host(base_url);
613    if should_warn_once(&format!("{provider}@{host}")) {
614        tracing::warn!(
615            "built-in provider '{}' base_url overridden in config: the {} API key will be sent to \
616             host '{}' instead of the trusted default endpoint",
617            provider,
618            provider,
619            host
620        );
621    }
622}
623
624/// Host portion of a `base_url` for the override warning, or `"<unknown>"` if it
625/// can't be parsed (the caller already validated it, so this is belt-and-braces).
626fn provider_host(base_url: &str) -> String {
627    reqwest::Url::parse(base_url)
628        .ok()
629        .and_then(|u| u.host_str().map(str::to_string))
630        .unwrap_or_else(|| "<unknown>".to_string())
631}
632
633/// Record `key` in the process-wide warned-set, returning `true` the first time
634/// (i.e. "warn now") and `false` thereafter. Poison-tolerant.
635fn should_warn_once(key: &str) -> bool {
636    let mut warned = WARNED_OVERRIDE_HOSTS
637        .lock()
638        .unwrap_or_else(|poisoned| poisoned.into_inner());
639    warned.insert(key.to_string())
640}
641
642#[cfg(test)]
643mod tests {
644    use super::*;
645
646    #[test]
647    fn base_url_is_local_classifies_hosts() {
648        assert!(base_url_is_local("http://127.0.0.1:8000/v1"));
649        assert!(base_url_is_local("http://localhost:1234/v1"));
650        assert!(base_url_is_local("http://192.168.1.5:8000/v1"));
651        assert!(!base_url_is_local("https://api.openai.com/v1"));
652        assert!(!base_url_is_local("not a url"));
653    }
654
655    #[test]
656    fn merged_headers_keeps_static_profile_headers_and_user_overrides() {
657        let profile = crate::models::lookup_provider("openrouter").unwrap();
658        // No user config: the static profile headers survive (the latent-bug fix).
659        let base = merged_headers(profile, None);
660        assert_eq!(
661            base.get("X-OpenRouter-Title").map(String::as_str),
662            Some("Mermaid")
663        );
664        assert!(base.contains_key("HTTP-Referer"));
665        // User extra_headers merge on top and can override a static one.
666        let mut cfg = crate::app::UserProviderConfig::default();
667        cfg.extra_headers.insert("X-Custom".into(), "v".into());
668        cfg.extra_headers
669            .insert("X-OpenRouter-Title".into(), "Override".into());
670        let merged = merged_headers(profile, Some(&cfg));
671        assert_eq!(merged.get("X-Custom").map(String::as_str), Some("v"));
672        assert_eq!(
673            merged.get("X-OpenRouter-Title").map(String::as_str),
674            Some("Override")
675        );
676        assert!(merged.contains_key("HTTP-Referer"));
677    }
678
679    #[test]
680    fn merged_headers_resolves_env_headers_and_skips_missing() {
681        let profile = crate::models::lookup_provider("openai").unwrap();
682        let mut cfg = crate::app::UserProviderConfig::default();
683        cfg.env_headers
684            .insert("X-Gateway-Token".into(), "MERMAID_TEST_GW_TOKEN".into());
685        temp_env::with_var("MERMAID_TEST_GW_TOKEN", Some("secret123"), || {
686            let merged = merged_headers(profile, Some(&cfg));
687            assert_eq!(
688                merged.get("X-Gateway-Token").map(String::as_str),
689                Some("secret123")
690            );
691        });
692        temp_env::with_var("MERMAID_TEST_GW_TOKEN", None::<&str>, || {
693            assert!(!merged_headers(profile, Some(&cfg)).contains_key("X-Gateway-Token"));
694        });
695    }
696
697    #[test]
698    fn base_url_requires_https_for_remote_hosts() {
699        // Remote http is refused (would leak the bearer key in cleartext).
700        assert!(validate_provider_base_url("http://api.example.com/v1").is_err());
701        assert!(validate_provider_base_url("ftp://example.com").is_err());
702        // https anywhere and http to LOOPBACK are fine.
703        assert!(validate_provider_base_url("https://api.example.com/v1").is_ok());
704        assert!(validate_provider_base_url("http://localhost:11434/v1").is_ok());
705        assert!(validate_provider_base_url("http://127.0.0.1:8000").is_ok());
706        assert!(validate_provider_base_url("http://[::1]:8000").is_ok());
707        // #26: http to a non-loopback host (even a private LAN one) is refused —
708        // the API key would otherwise cross the wire in cleartext.
709        assert!(validate_provider_base_url("http://192.168.1.5:8080").is_err());
710        assert!(validate_provider_base_url("http://169.254.169.254").is_err());
711    }
712
713    #[test]
714    fn cloudflare_base_url_synthesizes_account_scoped_endpoint() {
715        assert_eq!(
716            cloudflare_base_url("acct123"),
717            "https://api.cloudflare.com/client/v4/accounts/acct123/ai/v1"
718        );
719        // Trims surrounding whitespace (e.g. a trailing newline from `export`).
720        assert_eq!(
721            cloudflare_base_url("  acct123\n"),
722            "https://api.cloudflare.com/client/v4/accounts/acct123/ai/v1"
723        );
724    }
725
726    #[test]
727    fn cloudflare_account_id_required_and_non_blank() {
728        // Unset → clear, actionable error naming the env var.
729        temp_env::with_var("CLOUDFLARE_ACCOUNT_ID", None::<&str>, || {
730            let err = require_cloudflare_account_id().expect_err("must error when unset");
731            assert!(format!("{err}").contains("CLOUDFLARE_ACCOUNT_ID"));
732        });
733        // Whitespace-only is treated as unset.
734        temp_env::with_var("CLOUDFLARE_ACCOUNT_ID", Some("   "), || {
735            assert!(require_cloudflare_account_id().is_err());
736        });
737        // A real value resolves, trimmed.
738        temp_env::with_var("CLOUDFLARE_ACCOUNT_ID", Some(" acct123 "), || {
739            assert_eq!(require_cloudflare_account_id().unwrap(), "acct123");
740        });
741    }
742
743    #[test]
744    fn discovery_base_url_resolves_per_provider() {
745        let cf = lookup_provider("cloudflare").expect("cloudflare is in the registry");
746        let openai = lookup_provider("openai").expect("openai is in the registry");
747        // A user override always wins, for any provider.
748        assert_eq!(
749            discovery_base_url(cf, Some("https://gw.example/v1".to_string())),
750            Some("https://gw.example/v1".to_string())
751        );
752        // Non-cloudflare: the static registry default.
753        assert_eq!(
754            discovery_base_url(openai, None),
755            Some(openai.base_url.to_string())
756        );
757        // Cloudflare synthesizes the account-scoped URL from env...
758        temp_env::with_var("CLOUDFLARE_ACCOUNT_ID", Some("acct123"), || {
759            assert_eq!(
760                discovery_base_url(cf, None),
761                Some("https://api.cloudflare.com/client/v4/accounts/acct123/ai/v1".to_string())
762            );
763        });
764        // ...and yields None when it's unset — nothing real to probe.
765        temp_env::with_var("CLOUDFLARE_ACCOUNT_ID", None::<&str>, || {
766            assert_eq!(discovery_base_url(cf, None), None);
767        });
768    }
769
770    #[tokio::test]
771    async fn cloudflare_missing_both_env_vars_is_one_combined_error() {
772        temp_env::async_with_vars(
773            [
774                ("CLOUDFLARE_ACCOUNT_ID", None::<&str>),
775                ("CLOUDFLARE_API_TOKEN", None),
776            ],
777            async {
778                let f = ProviderFactory::new(Config::default());
779                let err = match f.resolve("cloudflare/@cf/zai-org/glm-5.2").await {
780                    Ok(_) => panic!("must fail with neither env var set"),
781                    Err(e) => e,
782                };
783                let msg = format!("{err}");
784                assert!(
785                    msg.contains("CLOUDFLARE_API_TOKEN") && msg.contains("CLOUDFLARE_ACCOUNT_ID"),
786                    "one error must name both missing vars, got: {msg}"
787                );
788            },
789        )
790        .await;
791    }
792
793    use std::sync::atomic::{AtomicUsize, Ordering};
794
795    fn unique_env(prefix: &str) -> String {
796        static N: AtomicUsize = AtomicUsize::new(0);
797        format!(
798            "{}_{}_{}",
799            prefix,
800            std::process::id(),
801            N.fetch_add(1, Ordering::SeqCst)
802        )
803    }
804
805    #[test]
806    fn parse_bare_name_defaults_to_ollama() {
807        let (p, m) = parse_model_id("qwen3-coder:30b");
808        assert_eq!(p, "ollama");
809        assert_eq!(m, "qwen3-coder:30b");
810    }
811
812    #[test]
813    fn parse_prefixed() {
814        let (p, m) = parse_model_id("anthropic/claude-opus-4-7");
815        assert_eq!(p, "anthropic");
816        assert_eq!(m, "claude-opus-4-7");
817    }
818
819    #[tokio::test]
820    async fn meta_requires_its_documented_api_key_env() {
821        temp_env::async_with_vars(
822            [(
823                crate::providers::model::meta::DEFAULT_API_KEY_ENV,
824                None::<&str>,
825            )],
826            async {
827                let factory = ProviderFactory::new(Config::default());
828                let error = match factory.resolve("meta/muse-spark-1.1").await {
829                    Ok(_) => panic!("Meta must require an API key"),
830                    Err(error) => error,
831                };
832                assert!(
833                    error
834                        .to_string()
835                        .contains(crate::providers::model::meta::DEFAULT_API_KEY_ENV)
836                );
837            },
838        )
839        .await;
840    }
841
842    #[tokio::test]
843    async fn meta_routes_to_responses_provider_with_muse_capabilities() {
844        temp_env::async_with_vars(
845            [(
846                crate::providers::model::meta::DEFAULT_API_KEY_ENV,
847                Some("test-key"),
848            )],
849            async {
850                let factory = ProviderFactory::new(Config::default());
851                let provider = factory.resolve("meta/muse-spark-1.1").await.unwrap();
852                let capabilities = provider.capabilities();
853                assert!(capabilities.supports_tools);
854                assert!(capabilities.supports_vision);
855                assert!(capabilities.emits_provider_continuation);
856                assert_eq!(
857                    capabilities.max_context_tokens,
858                    Some(crate::constants::META_MUSE_SPARK_CONTEXT_WINDOW)
859                );
860                assert_eq!(
861                    capabilities.max_output_tokens,
862                    Some(crate::constants::META_MUSE_SPARK_MAX_OUTPUT_TOKENS)
863                );
864            },
865        )
866        .await;
867    }
868
869    #[test]
870    fn gemini_key_resolution_accepts_legacy_fallback() {
871        let primary = unique_env("MERMAID_FACTORY_GEMINI_PRIMARY");
872        let legacy = unique_env("MERMAID_FACTORY_GEMINI_LEGACY");
873        temp_env::with_vars(
874            [(primary.as_str(), None), (legacy.as_str(), Some("legacy"))],
875            || {
876                let resolved = require_key_with_fallback("gemini", &primary, &legacy)
877                    .expect("legacy fallback should resolve");
878                assert_eq!(resolved, "legacy");
879            },
880        );
881    }
882
883    #[test]
884    fn gemini_key_resolution_prefers_google_primary() {
885        let primary = unique_env("MERMAID_FACTORY_GEMINI_PRIMARY2");
886        let legacy = unique_env("MERMAID_FACTORY_GEMINI_LEGACY2");
887        temp_env::with_vars(
888            [
889                (primary.as_str(), Some("google")),
890                (legacy.as_str(), Some("legacy")),
891            ],
892            || {
893                let resolved = require_key_with_fallback("gemini", &primary, &legacy)
894                    .expect("primary should resolve");
895                assert_eq!(resolved, "google");
896            },
897        );
898    }
899
900    #[tokio::test]
901    async fn factory_reports_unknown_provider_clearly() {
902        let cfg = Config::default();
903        let f = ProviderFactory::new(cfg);
904        match f.resolve("totally-made-up/model").await {
905            Ok(_) => panic!("expected error"),
906            Err(e) => {
907                let msg = format!("{}", e);
908                assert!(
909                    msg.contains("totally-made-up") || msg.contains("Unknown provider"),
910                    "error message: {}",
911                    msg
912                );
913            },
914        }
915    }
916
917    #[test]
918    fn normalize_cache_key_lowercases_provider_only() {
919        // #83: provider segment is lowercased; the model segment is preserved.
920        assert_eq!(
921            normalize_cache_key("Anthropic/Claude-X"),
922            "anthropic/Claude-X"
923        );
924        assert_eq!(
925            normalize_cache_key("anthropic/Claude-X"),
926            "anthropic/Claude-X"
927        );
928        // Bare names default to the ollama provider.
929        assert_eq!(normalize_cache_key("qwen3:30b"), "ollama/qwen3:30b");
930    }
931
932    #[tokio::test]
933    async fn resolve_is_single_flight_and_cached() {
934        // Ollama is keyless and builds no network connection, so this resolves
935        // offline. Two concurrent resolves with different provider casing must
936        // return the same cached instance (#83 normalization + #84 single-flight).
937        let f = ProviderFactory::new(Config::default());
938        let (a, b) = tokio::join!(
939            f.resolve("ollama/test-model"),
940            f.resolve("Ollama/test-model"),
941        );
942        let a = a.expect("resolve a");
943        let b = b.expect("resolve b");
944        assert!(
945            Arc::ptr_eq(&a, &b),
946            "expected one cached provider for casing variants + concurrent resolve"
947        );
948    }
949
950    // F66: a built-in provider's base_url override is hardened — validated for
951    // scheme and otherwise honored (the warning is a side effect we don't assert
952    // here, but the dedup helper is tested separately below).
953    #[test]
954    fn builtin_base_url_override_validated_and_resolved() {
955        // No override → trusted default, unchanged, no validation needed.
956        assert_eq!(
957            resolve_overridable_base_url("anthropic", None, "https://api.anthropic.com/v1")
958                .unwrap(),
959            "https://api.anthropic.com/v1"
960        );
961        // https override is honored (the key would go there — warned, not blocked).
962        assert_eq!(
963            resolve_overridable_base_url(
964                "anthropic",
965                Some("https://proxy.internal/v1".to_string()),
966                "https://api.anthropic.com/v1",
967            )
968            .unwrap(),
969            "https://proxy.internal/v1"
970        );
971        // http override to a NON-loopback host is refused (would leak the key in
972        // cleartext) — the F66 https requirement.
973        assert!(
974            resolve_overridable_base_url(
975                "anthropic",
976                Some("http://attacker.example/v1".to_string()),
977                "https://api.anthropic.com/v1",
978            )
979            .is_err()
980        );
981        // http override to loopback stays allowed for local dev (don't break it).
982        assert!(
983            resolve_overridable_base_url(
984                "openai",
985                Some("http://localhost:8080/v1".to_string()),
986                "https://api.openai.com/v1",
987            )
988            .is_ok()
989        );
990    }
991
992    #[test]
993    fn provider_host_extracts_host_or_unknown() {
994        assert_eq!(
995            provider_host("https://attacker.example/v1"),
996            "attacker.example"
997        );
998        assert_eq!(provider_host("http://127.0.0.1:8080"), "127.0.0.1");
999        assert_eq!(provider_host("not a url"), "<unknown>");
1000    }
1001
1002    #[test]
1003    fn override_host_warning_is_deduped() {
1004        // The F66 warning must be one-time per key: first call fires, the rest
1005        // are suppressed. Use a process-unique key so this test doesn't race the
1006        // shared warned-set with any other test.
1007        let key = unique_env("MERMAID_FACTORY_WARN_KEY");
1008        assert!(should_warn_once(&key), "first warn for a key must fire");
1009        assert!(
1010            !should_warn_once(&key),
1011            "subsequent warns for the same key must be suppressed"
1012        );
1013    }
1014
1015    // F67: identical custom-provider inputs must reuse one leaked &'static
1016    // profile, and distinct inputs must each leak exactly one — no per-model_id
1017    // growth.
1018    #[test]
1019    fn custom_profile_is_memoized_per_key() {
1020        use crate::app::UserProviderConfig;
1021        let cfg = UserProviderConfig {
1022            base_url: Some("https://api.custom.test/v1".to_string()),
1023            api_key_env: Some("CUSTOM_KEY".to_string()),
1024            compat: Some("openai".to_string()),
1025            ..Default::default()
1026        };
1027        // Two resolves with identical inputs (as happens for two different models
1028        // of the same custom provider) return the SAME leaked pointer.
1029        let a = user_profile_to_static("mermaid_test_customx", &cfg).unwrap();
1030        let b = user_profile_to_static("mermaid_test_customx", &cfg).unwrap();
1031        assert!(
1032            std::ptr::eq(a, b),
1033            "identical custom-provider inputs must reuse one leaked &'static profile"
1034        );
1035        assert_eq!(a.base_url, "https://api.custom.test/v1");
1036        assert_eq!(a.api_key_env, "CUSTOM_KEY");
1037
1038        // A different base_url is a different key → a distinct leaked profile.
1039        let cfg2 = UserProviderConfig {
1040            base_url: Some("https://api.custom.test/v2".to_string()),
1041            ..cfg.clone()
1042        };
1043        let c = user_profile_to_static("mermaid_test_customx", &cfg2).unwrap();
1044        assert!(
1045            !std::ptr::eq(a, c),
1046            "a different base_url must leak a distinct profile"
1047        );
1048    }
1049}