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_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
20/// Resolve an API key or return a clear `ModelError` when the env
21/// var isn't set. Takes `default_env` (the registry-default name)
22/// and allows no override — the factory passes the already-resolved
23/// env name.
24fn 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
47/// A lazily-built, shared provider. `OnceCell` gives single-flight construction
48/// (#84); the outer `Arc` lets `resolve` clone the cell out from under the cache
49/// lock and initialize it without holding the lock across the build.
50type ProviderCell = Arc<tokio::sync::OnceCell<Arc<dyn ModelProvider>>>;
51
52/// Per-process provider cache. Providers are expensive to construct
53/// (HTTP client, connection pool, capability lookup) so the effect
54/// runner asks for them lazily and reuses across turns.
55pub struct ProviderFactory {
56    config: Arc<Config>,
57    /// Per-key cache of built providers, keyed by normalized model id (#83). The
58    /// `Mutex` is only held to get-or-insert a cell, never across the build.
59    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    /// Resolve (or lazily construct) a provider for the given model
75    /// ID. Hits the cache on the second and subsequent calls for the
76    /// same ID.
77    pub async fn resolve(&self, model_id: &str) -> Result<Arc<dyn ModelProvider>> {
78        let key = normalize_cache_key(model_id);
79        // Get-or-insert the per-key cell under a brief lock, then initialize it
80        // exactly once outside the lock (#84). A failed build isn't cached, so a
81        // transient error can be retried on the next call.
82        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
100/// Build a provider for the given `model_id`:
101///   1. `ollama/<model>` → OllamaProvider.
102///   2. `anthropic/<model>` → AnthropicProvider.
103///   3. `gemini/<model>` → GeminiProvider.
104///   4. Other builtin providers (openai, openrouter, groq, …) → OpenAICompatProvider.
105///   5. User-defined `[providers.<name>]` → custom OpenAICompatProvider.
106///   6. Bare model name → OllamaProvider.
107async 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    // 1. Ollama (and bare names). F11: pass Arc<Config> so the wrapper
112    // can forward Ollama hardware options to the adapter.
113    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    // 2. Anthropic — bespoke API shape.
125    if provider_lc == "anthropic" {
126        let user_cfg = config.providers.get("anthropic");
127        let base_url = user_cfg
128            .and_then(|c| c.base_url.clone())
129            .unwrap_or_else(|| "https://api.anthropic.com/v1".to_string());
130        validate_provider_base_url(&base_url)?;
131        let api_key_env = user_cfg
132            .and_then(|c| c.api_key_env.as_deref())
133            .unwrap_or("ANTHROPIC_API_KEY");
134        let api_key = require_key("anthropic", api_key_env)?;
135        let p = AnthropicProvider::new(api_key, model_name.to_string(), base_url)?;
136        return Ok(Box::new(p));
137    }
138
139    // 3. Gemini — GCP AI Studio shape.
140    if provider_lc == "gemini" {
141        let user_cfg = config.providers.get("gemini");
142        let base_url = user_cfg
143            .and_then(|c| c.base_url.clone())
144            .unwrap_or_else(|| "https://generativelanguage.googleapis.com/v1beta".to_string());
145        validate_provider_base_url(&base_url)?;
146        let api_key = match user_cfg.and_then(|c| c.api_key_env.as_deref()) {
147            Some(api_key_env) => require_key("gemini", api_key_env)?,
148            None => {
149                require_key_with_fallback("gemini", GEMINI_API_KEY_ENV, GEMINI_LEGACY_API_KEY_ENV)?
150            },
151        };
152        let p = GeminiProvider::new(api_key, model_name.to_string(), base_url)?;
153        return Ok(Box::new(p));
154    }
155
156    // 4 + 5. OpenAI-compatible registry or user-custom.
157    if let Some(profile) = lookup_provider(&provider_lc) {
158        let user_cfg = config.providers.get(&provider_lc);
159        let base_url = user_cfg
160            .and_then(|c| c.base_url.clone())
161            .unwrap_or_else(|| profile.base_url.to_string());
162        validate_provider_base_url(&base_url)?;
163        let api_key_env = user_cfg
164            .and_then(|c| c.api_key_env.as_deref())
165            .unwrap_or(profile.api_key_env);
166        let api_key = require_key(&provider_lc, api_key_env)?;
167        let extra_headers = user_cfg
168            .map(|c| c.extra_headers.clone())
169            .unwrap_or_default();
170        let p = OpenAICompatProvider::new(
171            profile,
172            base_url,
173            api_key,
174            model_name.to_string(),
175            extra_headers,
176        )?;
177        return Ok(Box::new(p));
178    }
179
180    // User-custom: no registry entry, but the user has [providers.<name>]
181    // in config with a declared `compat` field.
182    if let Some(user_cfg) = config.providers.get(&provider_lc)
183        && let Some(profile) = user_profile_to_static(&provider_lc, user_cfg)
184    {
185        let base_url = user_cfg.base_url.clone().ok_or_else(|| {
186            ModelError::InvalidRequest(format!(
187                "custom provider '{}' requires base_url in config",
188                provider_lc
189            ))
190        })?;
191        validate_provider_base_url(&base_url)?;
192        let api_key_env = user_cfg.api_key_env.as_deref().ok_or_else(|| {
193            ModelError::InvalidRequest(format!(
194                "custom provider '{}' requires api_key_env in config",
195                provider_lc
196            ))
197        })?;
198        let api_key = require_key(&provider_lc, api_key_env)?;
199        let p = OpenAICompatProvider::new(
200            profile,
201            base_url,
202            api_key,
203            model_name.to_string(),
204            user_cfg.extra_headers.clone(),
205        )?;
206        return Ok(Box::new(p));
207    }
208
209    Err(ModelError::InvalidRequest(format!(
210        "Unknown provider '{}' (model_id: {})",
211        provider, model_id
212    )))
213}
214
215/// Cache key for a model id: the provider segment lowercased (matching
216/// `build_provider`'s own normalization) so `Anthropic/x` and `anthropic/x`
217/// resolve to one cached provider instead of building two identical ones (#83).
218fn normalize_cache_key(model_id: &str) -> String {
219    let (provider, model) = parse_model_id(model_id);
220    format!("{}/{}", provider.to_lowercase(), model)
221}
222
223/// Parse `provider/model` → `(provider, model)`. Bare strings are
224/// Ollama by convention.
225fn parse_model_id(model_id: &str) -> (String, &str) {
226    match model_id.split_once('/') {
227        Some((p, m)) => (p.to_string(), m),
228        None => ("ollama".to_string(), model_id),
229    }
230}
231
232pub(crate) fn ollama_backend_config(config: &Config) -> BackendConfig {
233    BackendConfig {
234        // Scheme-less: `normalize_url` in the adapter picks http (loopback/LAN)
235        // vs https (public) by host class (#86).
236        ollama_url: format!("{}:{}", config.ollama.host, config.ollama.port),
237        max_idle_per_host: 10,
238        timeout_secs: 10,
239    }
240}
241
242/// Convert a user-defined `[providers.<name>]` entry into a `&'static
243/// ProviderProfile`. We need `&'static` because `ProviderProfile`'s
244/// lifetime is tied to the registry constants; we leak a tiny owned
245/// copy so custom providers can participate without redesigning the
246/// profile type. Leaked allocations are bounded by the number of
247/// custom providers (typically 0-3).
248fn user_profile_to_static(
249    name: &str,
250    user_cfg: &crate::app::UserProviderConfig,
251) -> Option<&'static crate::models::ProviderProfile> {
252    use crate::models::{ProviderProfile, ReasoningExtraction, ReasoningStrategy};
253
254    let compat = user_cfg.compat.as_deref().unwrap_or("openai");
255    let strategy = match compat {
256        "openai" => ReasoningStrategy::None,
257        "openai-effort" => ReasoningStrategy::Effort,
258        "openrouter" => ReasoningStrategy::OpenRouterShape,
259        _ => ReasoningStrategy::None,
260    };
261
262    let profile = Box::new(ProviderProfile {
263        name: Box::leak(name.to_string().into_boxed_str()),
264        base_url: Box::leak(
265            user_cfg
266                .base_url
267                .clone()
268                .unwrap_or_default()
269                .into_boxed_str(),
270        ),
271        api_key_env: Box::leak(
272            user_cfg
273                .api_key_env
274                .clone()
275                .unwrap_or_default()
276                .into_boxed_str(),
277        ),
278        extra_headers: &[],
279        reasoning_strategy: strategy,
280        reasoning_extraction: ReasoningExtraction::None,
281        max_tokens_param: crate::models::MaxTokensParam::MaxTokens,
282        disable_parallel_tool_calls_for: &[],
283    });
284    Some(Box::leak(profile))
285}
286
287/// Validate a provider `base_url` before it's handed an API key. Requires
288/// http/https, and **requires https for any non-loopback, non-private host** —
289/// a typo'd or hostile `http://` endpoint would otherwise receive the bearer
290/// key in cleartext. Plain http stays allowed for loopback / RFC-1918 hosts so
291/// local model servers (Ollama, vLLM) keep working.
292fn validate_provider_base_url(url: &str) -> Result<()> {
293    let parsed = reqwest::Url::parse(url).map_err(|e| {
294        ModelError::InvalidRequest(format!("invalid provider base_url '{url}': {e}"))
295    })?;
296    match parsed.scheme() {
297        "https" => Ok(()),
298        // Plaintext http is only safe to LOOPBACK: a key sent over http to any
299        // other host — including a LAN/private one — crosses the wire in
300        // cleartext. Every caller here is a key-bearing provider; keyless local
301        // Ollama uses a separate path that doesn't validate.
302        "http"
303            if crate::utils::classify_host(parsed.host_str().unwrap_or_default()).is_loopback() =>
304        {
305            Ok(())
306        },
307        "http" => Err(ModelError::InvalidRequest(format!(
308            "provider base_url '{url}' uses http:// to a non-loopback host — refusing to send the \
309             API key in cleartext. Use https, or http://localhost for a local server."
310        ))),
311        other => Err(ModelError::InvalidRequest(format!(
312            "provider base_url '{url}' has unsupported scheme '{other}' (use http or https)"
313        ))),
314    }
315}
316
317#[cfg(test)]
318mod tests {
319    use super::*;
320
321    #[test]
322    fn base_url_requires_https_for_remote_hosts() {
323        // Remote http is refused (would leak the bearer key in cleartext).
324        assert!(validate_provider_base_url("http://api.example.com/v1").is_err());
325        assert!(validate_provider_base_url("ftp://example.com").is_err());
326        // https anywhere and http to LOOPBACK are fine.
327        assert!(validate_provider_base_url("https://api.example.com/v1").is_ok());
328        assert!(validate_provider_base_url("http://localhost:11434/v1").is_ok());
329        assert!(validate_provider_base_url("http://127.0.0.1:8000").is_ok());
330        assert!(validate_provider_base_url("http://[::1]:8000").is_ok());
331        // #26: http to a non-loopback host (even a private LAN one) is refused —
332        // the API key would otherwise cross the wire in cleartext.
333        assert!(validate_provider_base_url("http://192.168.1.5:8080").is_err());
334        assert!(validate_provider_base_url("http://169.254.169.254").is_err());
335    }
336    use std::sync::atomic::{AtomicUsize, Ordering};
337
338    fn unique_env(prefix: &str) -> String {
339        static N: AtomicUsize = AtomicUsize::new(0);
340        format!(
341            "{}_{}_{}",
342            prefix,
343            std::process::id(),
344            N.fetch_add(1, Ordering::SeqCst)
345        )
346    }
347
348    #[test]
349    fn parse_bare_name_defaults_to_ollama() {
350        let (p, m) = parse_model_id("qwen3-coder:30b");
351        assert_eq!(p, "ollama");
352        assert_eq!(m, "qwen3-coder:30b");
353    }
354
355    #[test]
356    fn parse_prefixed() {
357        let (p, m) = parse_model_id("anthropic/claude-opus-4-7");
358        assert_eq!(p, "anthropic");
359        assert_eq!(m, "claude-opus-4-7");
360    }
361
362    #[test]
363    fn gemini_key_resolution_accepts_legacy_fallback() {
364        let primary = unique_env("MERMAID_FACTORY_GEMINI_PRIMARY");
365        let legacy = unique_env("MERMAID_FACTORY_GEMINI_LEGACY");
366        temp_env::with_vars(
367            [(primary.as_str(), None), (legacy.as_str(), Some("legacy"))],
368            || {
369                let resolved = require_key_with_fallback("gemini", &primary, &legacy)
370                    .expect("legacy fallback should resolve");
371                assert_eq!(resolved, "legacy");
372            },
373        );
374    }
375
376    #[test]
377    fn gemini_key_resolution_prefers_google_primary() {
378        let primary = unique_env("MERMAID_FACTORY_GEMINI_PRIMARY2");
379        let legacy = unique_env("MERMAID_FACTORY_GEMINI_LEGACY2");
380        temp_env::with_vars(
381            [
382                (primary.as_str(), Some("google")),
383                (legacy.as_str(), Some("legacy")),
384            ],
385            || {
386                let resolved = require_key_with_fallback("gemini", &primary, &legacy)
387                    .expect("primary should resolve");
388                assert_eq!(resolved, "google");
389            },
390        );
391    }
392
393    #[tokio::test]
394    async fn factory_reports_unknown_provider_clearly() {
395        let cfg = Config::default();
396        let f = ProviderFactory::new(cfg);
397        match f.resolve("totally-made-up/model").await {
398            Ok(_) => panic!("expected error"),
399            Err(e) => {
400                let msg = format!("{}", e);
401                assert!(
402                    msg.contains("totally-made-up") || msg.contains("Unknown provider"),
403                    "error message: {}",
404                    msg
405                );
406            },
407        }
408    }
409
410    #[test]
411    fn normalize_cache_key_lowercases_provider_only() {
412        // #83: provider segment is lowercased; the model segment is preserved.
413        assert_eq!(
414            normalize_cache_key("Anthropic/Claude-X"),
415            "anthropic/Claude-X"
416        );
417        assert_eq!(
418            normalize_cache_key("anthropic/Claude-X"),
419            "anthropic/Claude-X"
420        );
421        // Bare names default to the ollama provider.
422        assert_eq!(normalize_cache_key("qwen3:30b"), "ollama/qwen3:30b");
423    }
424
425    #[tokio::test]
426    async fn resolve_is_single_flight_and_cached() {
427        // Ollama is keyless and builds no network connection, so this resolves
428        // offline. Two concurrent resolves with different provider casing must
429        // return the same cached instance (#83 normalization + #84 single-flight).
430        let f = ProviderFactory::new(Config::default());
431        let (a, b) = tokio::join!(
432            f.resolve("ollama/test-model"),
433            f.resolve("Ollama/test-model"),
434        );
435        let a = a.expect("resolve a");
436        let b = b.expect("resolve b");
437        assert!(
438            Arc::ptr_eq(&a, &b),
439            "expected one cached provider for casing variants + concurrent resolve"
440        );
441    }
442}
mermaid_cli/providers/factory.rs

mermaid_cli/providers/
factory.rs
