Skip to main content

yapi_ai/
registry.rs

1//! The models yapi can use and the credentials to use them.
2//!
3//! Layers, as in pi's provider composer: the built-in catalog, then `models.json`
4//! (provider base URLs, headers and compat, added or replaced models, overrides),
5//! then extension providers. Credentials resolve in pi's order: a runtime key
6//! (`--api-key`), `auth.json`, `models.json` `apiKey`, then environment variables.
7
8use std::collections::{HashMap, HashSet};
9use std::path::Path;
10use std::sync::{Arc, LazyLock};
11use std::time::Duration;
12
13use indexmap::IndexMap;
14use regex_lite::{Captures, Regex};
15use serde_json::{Map, Value};
16use tokio_util::sync::CancellationToken;
17use yapi_types::auth::{ApiKeyCredential, Credential, OAuthCredential};
18use yapi_types::classify::{AssistantImages, ClassifierContext, ClassifierResult, ImagesContext};
19use yapi_types::model::{ClassifierModel, ImageModel, Model, Pricing};
20use yapi_types::models::{ModelDefinition, ModelOverride, ModelsConfig, ProviderConfig};
21
22use crate::api::ModelOperations;
23use crate::api::classify::ClassifyOptions;
24use crate::api::images::ImagesOptions;
25use crate::auth::{
26    AuthError, AuthPrompt, CredentialKind, CredentialStore, Interaction, KeyAuth, LoginOptions,
27    OAuthAuth, OAuthProvider, builtin_oauth, copilot,
28};
29use crate::catalog;
30use crate::credentials::{self, BEARER_TOKEN_ENV, ProviderEnv};
31use crate::key_auth::{self, Ambient};
32use crate::model_catalog::{self, ModelsStore, ProviderModels, Source, Target};
33use crate::providers;
34
35/// Tokens with less validity left than this are refreshed before use.
36const OAUTH_MINIMUM_VALIDITY_MS: u64 = 5 * 60 * 1000;
37const OAUTH_REFRESH_TIMEOUT: Duration = Duration::from_secs(15);
38
39/// Credentials for one request.
40#[derive(Clone, Debug, Default, PartialEq)]
41pub struct Auth {
42    /// The API key, if the provider takes one.
43    pub api_key: Option<String>,
44    /// Headers that carry or accompany the credential. `None` removes a header.
45    pub headers: IndexMap<String, Option<String>>,
46    /// Where the credential came from, for messages.
47    pub source: Option<String>,
48    /// Replaces the model's base URL, for accounts with their own endpoint.
49    pub base_url: Option<String>,
50    /// Why credentials could not be resolved. A request with it fails.
51    pub error: Option<String>,
52    /// Provider settings from the credential, such as a Cloudflare account
53    /// id, that the request reads ahead of the process environment.
54    pub env: Option<ProviderEnv>,
55}
56
57impl Auth {
58    /// Takes the credentials of a provider-specific resolution.
59    fn use_resolved(&mut self, resolved: key_auth::Resolved) {
60        self.api_key = resolved.api_key;
61        self.headers.extend(resolved.headers);
62        self.env = resolved.env;
63        self.source = Some(resolved.source);
64        if resolved.base_url.is_some() {
65            self.base_url = resolved.base_url;
66        }
67    }
68
69    /// Applies the credentials beneath a classifier or image request's own:
70    /// the caller's key and headers win, and the account's base URL replaces
71    /// `base_url`.
72    fn apply_under(
73        self,
74        base_url: &mut String,
75        api_key: &mut Option<String>,
76        headers: &mut IndexMap<String, Option<String>>,
77    ) {
78        if let Some(url) = self.base_url {
79            *base_url = url;
80        }
81        if api_key.is_none() {
82            *api_key = self.api_key;
83        }
84        let caller = std::mem::replace(headers, self.headers);
85        headers.extend(caller);
86    }
87
88    /// Applies the credentials to a request: the key when there is one,
89    /// headers over the request's, and the account's base URL. Fails
90    /// with [`Auth::error`] when credentials could not be resolved.
91    pub fn apply(self, request: &mut crate::stream::Request) -> Result<(), String> {
92        if let Some(error) = self.error {
93            return Err(error);
94        }
95        if self.api_key.is_some() {
96            request.options.api_key = self.api_key;
97        }
98        request.options.headers.extend(self.headers);
99        if let Some(base_url) = self.base_url {
100            request.model.base_url = base_url;
101        }
102        if self.env.is_some() {
103            request.options.env = self.env;
104        }
105        Ok(())
106    }
107}
108
109/// Where `provider`'s credential comes from when nothing is stored or
110/// configured: its key variables, else the provider's ambient sources.
111fn environment_source(provider: &str) -> Option<String> {
112    if key_auth::is_custom(provider) {
113        return key_auth::resolve(provider, None, &Ambient::process()).map(|r| r.source);
114    }
115    credentials::env_api_key(provider, None)
116        .map(|(name, _)| name.to_owned())
117        .or_else(|| {
118            (provider == "anthropic")
119                .then(|| key_auth::anthropic_federation(&Ambient::process()))
120                .flatten()
121                .map(|resolved| resolved.source)
122        })
123}
124
125/// `credential` with its key resolved, or with no key when it does not resolve.
126async fn resolve_stored(credential: ApiKeyCredential) -> ApiKeyCredential {
127    let key = match &credential.key {
128        Some(key) => credentials::resolve(key, credential.env.as_ref(), true).await,
129        None => None,
130    };
131    ApiKeyCredential { key, ..credential }
132}
133
134/// Replaces the models `key` matches by id with `refreshed` ones, and
135/// appends the new ones (pi's `mergeModels`).
136fn overlay<T: Clone>(models: &mut Vec<T>, refreshed: &[T], key: impl Fn(&T) -> Option<&str>) {
137    for model in refreshed {
138        let Some(id) = key(model) else {
139            continue;
140        };
141        match models.iter().position(|existing| key(existing) == Some(id)) {
142            Some(index) => models[index] = model.clone(),
143            None => models.push(model.clone()),
144        }
145    }
146}
147
148/// A model that stands for a provider in provider-level auth.
149fn placeholder_model(provider: &str) -> Option<Model> {
150    serde_json::from_value(serde_json::json!({
151        "id": "",
152        "name": "",
153        "api": "pi-messages",
154        "provider": provider,
155        "baseUrl": "",
156        "reasoning": false,
157        "input": ["text"],
158        "cost": {"input": 0, "output": 0, "cacheRead": 0, "cacheWrite": 0},
159        "contextWindow": 0,
160        "maxTokens": 0,
161    }))
162    .ok()
163}
164
165/// The models of a Radius gateway config that older sign-ins, from before
166/// the models store, kept in their credential.
167fn legacy_radius_models(provider: &str, config: &Value) -> Vec<Model> {
168    let Some(base_url) = config["baseUrl"].as_str() else {
169        return Vec::new();
170    };
171    config["models"]
172        .as_array()
173        .into_iter()
174        .flatten()
175        .filter_map(|model| {
176            let mut model = model.clone();
177            model["api"] = Value::String("pi-messages".into());
178            model["provider"] = Value::String(provider.to_owned());
179            model["baseUrl"] = Value::String(base_url.to_owned());
180            serde_json::from_value(model).ok()
181        })
182        .collect()
183}
184
185/// How `/login` authenticates a provider.
186#[derive(Clone, Copy, Debug, PartialEq, Eq)]
187pub enum LoginKind {
188    /// An account sign-in.
189    OAuth,
190    /// A stored API key.
191    ApiKey,
192}
193
194/// What a pi-ai `Provider` from an extension makes available for the stored
195/// credential, as pi's `Models` finds it each time availability refreshes.
196#[derive(Clone, Debug, Default, PartialEq, serde::Deserialize)]
197pub struct Availability {
198    /// Where its API key `check` finds a key, or `None` when it finds none.
199    pub source: Option<String>,
200    /// The ids of the chat models its `filterModels` keeps. `None` keeps
201    /// every one.
202    pub chat: Option<HashSet<String>>,
203    /// The ids of the models of each type kept across every type, as pi's
204    /// `getAllAvailable` keeps them. `None` keeps every one.
205    pub all: Option<HashMap<String, HashSet<String>>>,
206}
207
208/// All known models with their configuration.
209#[derive(Clone, Default)]
210pub struct ModelRegistry {
211    models: Vec<Model>,
212    config: ModelsConfig,
213    store: Arc<CredentialStore>,
214    oauth: IndexMap<String, Arc<dyn OAuthProvider>>,
215    runtime_keys: IndexMap<String, String>,
216    /// Providers whose `apiKey` an extension configures.
217    extension_keys: HashSet<String>,
218    /// Providers extensions register.
219    extension_providers: HashSet<String>,
220    /// The `models.json` entries extension providers replaced, by provider,
221    /// with `None` where there was none.
222    replaced: HashMap<String, Option<ProviderConfig>>,
223    error: Option<String>,
224    /// Classifier models, built-in and refreshed.
225    classifiers: Vec<ClassifierModel>,
226    /// Image-generation models, built-in and refreshed.
227    images: Vec<ImageModel>,
228    /// Models from refreshed catalogs, by provider, over the built-in ones.
229    dynamic: IndexMap<String, ProviderModels>,
230    /// Whether pi's built-in `llama.cpp` extension runs, which provides the
231    /// llama.cpp provider.
232    llama: bool,
233    /// Models extension providers register, by provider.
234    extension_models: IndexMap<String, Vec<Model>>,
235    /// The image and classifier models of extension providers that define
236    /// models, by provider. They replace the provider's own, as in pi.
237    extension_typed: IndexMap<String, Vec<ModelDefinition>>,
238    /// Classifier and image APIs extension providers implement, by
239    /// provider and API.
240    operations: IndexMap<(String, String), Arc<dyn ModelOperations>>,
241    /// API-key auth of extension providers, with where the credential
242    /// comes from when nothing is stored, `None` when it is not configured.
243    key_auths: IndexMap<String, (Arc<dyn KeyAuth>, Option<String>)>,
244    /// What extension providers' API key checks and model filters make
245    /// available for the stored credential, by provider.
246    filters: IndexMap<String, Availability>,
247    /// Where refreshed catalogs persist.
248    models_store: Option<ModelsStore>,
249}
250
251impl std::fmt::Debug for ModelRegistry {
252    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
253        f.debug_struct("ModelRegistry")
254            .field("models", &self.models.len())
255            .field("store", &self.store)
256            .field("oauth", &self.oauth.keys().collect::<Vec<_>>())
257            .field("error", &self.error)
258            .finish_non_exhaustive()
259    }
260}
261
262/// Removes `//` comments and trailing commas outside strings, as pi's
263/// `stripJsonComments` does for `models.json`.
264pub fn strip_json_comments(input: &str) -> String {
265    // pi's two regular expressions, with JavaScript's `\s` spelled out.
266    static PATTERNS: LazyLock<[Regex; 2]> = LazyLock::new(|| {
267        [
268            r#""(?:\\.|[^"\\])*"|//[^\n]*"#,
269            r#""(?:\\.|[^"\\])*"|,([\t\n\x0B\f\r \x{A0}\x{1680}\x{2000}-\x{200A}\x{2028}\x{2029}\x{202F}\x{205F}\x{3000}\x{FEFF}]*[}\]])"#,
270        ]
271        .map(|pattern| {
272            Regex::new(pattern).expect("constant pattern; `strips_comments_and_commas`")
273        })
274    });
275    let [comments, commas] = &*PATTERNS;
276    let stripped = comments.replace_all(input, |found: &Captures| {
277        if found[0].starts_with('"') {
278            found[0].to_owned()
279        } else {
280            String::new()
281        }
282    });
283    commas
284        .replace_all(&stripped, |found: &Captures| {
285            found
286                .get(1)
287                .map_or(&found[0], |tail| tail.as_str())
288                .to_owned()
289        })
290        .into_owned()
291}
292
293fn merge_compat(
294    base: Option<&Map<String, Value>>,
295    over: Option<&Map<String, Value>>,
296) -> Option<Map<String, Value>> {
297    let Some(over) = over else {
298        return base.cloned();
299    };
300    let mut merged = base.cloned().unwrap_or_default();
301    for (key, value) in over {
302        merged.insert(key.clone(), value.clone());
303    }
304    for key in [
305        "openRouterRouting",
306        "vercelGatewayRouting",
307        "chatTemplateKwargs",
308        "chatTemplateArgs",
309    ] {
310        let base_value = base
311            .and_then(|base| base.get(key))
312            .and_then(Value::as_object);
313        let over_value = over.get(key).and_then(Value::as_object);
314        if base_value.is_some() || over_value.is_some() {
315            let mut nested = base_value.cloned().unwrap_or_default();
316            for (k, v) in over_value.into_iter().flatten() {
317                nested.insert(k.clone(), v.clone());
318            }
319            merged.insert(key.into(), Value::Object(nested));
320        }
321    }
322    Some(merged)
323}
324
325fn model_from_definition(
326    provider: &str,
327    definition: &ModelDefinition,
328    config: &ProviderConfig,
329    defaults: Option<&Model>,
330) -> Result<Model, String> {
331    let api = definition
332        .api
333        .clone()
334        .or_else(|| config.api.clone())
335        .or_else(|| defaults.map(|model| model.api.clone()))
336        .ok_or_else(|| {
337            format!(
338                "Provider {provider}, model {}: no \"api\" specified. Set at provider or model level.",
339                definition.id
340            )
341        })?;
342    let base_url = definition
343        .base_url
344        .clone()
345        .or_else(|| config.base_url.clone())
346        .or_else(|| defaults.map(|model| model.base_url.clone()))
347        .ok_or_else(|| {
348            format!("Provider {provider}: \"baseUrl\" is required when defining custom models.")
349        })?;
350    let options = &definition.options;
351    let cost = options
352        .cost
353        .as_ref()
354        .map_or_else(Pricing::default, |cost| Pricing {
355            input: cost.input.unwrap_or(0.0),
356            output: cost.output.unwrap_or(0.0),
357            cache_read: cost.cache_read.unwrap_or(0.0),
358            cache_write: cost.cache_write.unwrap_or(0.0),
359            tiers: cost.tiers.clone(),
360        });
361    Ok(Model {
362        id: definition.id.clone(),
363        name: definition
364            .name
365            .clone()
366            .unwrap_or_else(|| definition.id.clone()),
367        api,
368        provider: provider.to_owned(),
369        base_url,
370        reasoning: options.reasoning.unwrap_or(false),
371        input: options
372            .input
373            .clone()
374            .unwrap_or_else(|| vec![yapi_types::models::InputKind::Text]),
375        cost,
376        context_window: options.context_window.unwrap_or(128_000),
377        max_tokens: options.max_tokens.unwrap_or(16_384),
378        thinking_level_map: options.thinking_level_map.clone(),
379        prompt_cache: options.prompt_cache.clone(),
380        sampling_params: options.sampling_params.clone(),
381        compat: merge_compat(config.compat.as_ref(), options.compat.as_ref()),
382        input_limits: options.input_limits.clone(),
383        headers: None,
384        kind: None,
385    })
386}
387
388/// Whether `definition` is an image or classifier model, or of a type
389/// yapi does not know, rather than a chat model.
390fn is_typed(definition: &ModelDefinition) -> bool {
391    definition
392        .kind
393        .as_deref()
394        .is_some_and(|kind| kind != "chat")
395}
396
397/// The image and classifier models of an extension provider's
398/// `definitions`, over `classifiers` and `images`, the provider's own: pi's
399/// `extensionModelFromDefinition`. Each takes its API and base URL from the
400/// definition, the provider's base URL, else an own model of its type with
401/// its id or API, else the first. Definitions of other types are skipped.
402fn typed_models(
403    provider: &str,
404    config: &ProviderConfig,
405    definitions: &[ModelDefinition],
406    classifiers: &[ClassifierModel],
407    images: &[ImageModel],
408) -> Result<(Vec<ClassifierModel>, Vec<ImageModel>), String> {
409    let mut typed = (Vec::new(), Vec::new());
410    for definition in definitions {
411        let own: Vec<(&str, &str, &str)> = match definition.kind.as_deref() {
412            Some("classifier") => classifiers
413                .iter()
414                .filter(|model| model.provider == provider)
415                .map(|model| {
416                    (
417                        model.id.as_str(),
418                        model.api.as_str(),
419                        model.base_url.as_str(),
420                    )
421                })
422                .collect(),
423            Some("image") => images
424                .iter()
425                .filter(|model| model.provider == provider)
426                .map(|model| {
427                    (
428                        model.id.as_str(),
429                        model.api.as_str(),
430                        model.base_url.as_str(),
431                    )
432                })
433                .collect(),
434            _ => continue,
435        };
436        let defaults = own
437            .iter()
438            .find(|(id, ..)| *id == definition.id)
439            .or_else(|| {
440                own.iter()
441                    .find(|(_, api, _)| definition.api.as_deref() == Some(*api))
442            })
443            .or_else(|| own.first());
444        let api = definition
445            .api
446            .as_deref()
447            .or(defaults.map(|(_, api, _)| *api))
448            .ok_or_else(|| {
449                format!(
450                    "Provider {provider}, model {}: no \"api\" specified. Set it at model level.",
451                    definition.id
452                )
453            })?;
454        let base_url = definition
455            .base_url
456            .as_deref()
457            .or(config.base_url.as_deref())
458            .or(defaults.map(|(.., base_url)| *base_url))
459            .ok_or_else(|| {
460                format!("Provider {provider}: \"baseUrl\" is required when defining custom models.")
461            })?;
462        let invalid =
463            |err: serde_json::Error| format!("Provider {provider}, model {}: {err}", definition.id);
464        let mut model = serde_json::to_value(definition).map_err(invalid)?;
465        model["api"] = Value::from(api);
466        model["provider"] = Value::from(provider);
467        model["baseUrl"] = Value::from(base_url);
468        if definition.kind.as_deref() == Some("classifier") {
469            typed
470                .0
471                .push(serde_json::from_value(model).map_err(invalid)?);
472        } else {
473            typed
474                .1
475                .push(serde_json::from_value(model).map_err(invalid)?);
476        }
477    }
478    Ok(typed)
479}
480
481fn apply_override(model: &mut Model, over: &ModelOverride) {
482    let options = &over.options;
483    if let Some(name) = &over.name {
484        model.name = name.clone();
485    }
486    if let Some(reasoning) = options.reasoning {
487        model.reasoning = reasoning;
488    }
489    if let Some(map) = &options.thinking_level_map {
490        let mut merged = model.thinking_level_map.clone().unwrap_or_default();
491        merged.extend(map.clone());
492        model.thinking_level_map = Some(merged);
493    }
494    if let Some(input) = &options.input {
495        model.input = input.clone();
496    }
497    if let Some(limits) = &options.input_limits {
498        let mut merged = model.input_limits.clone().unwrap_or_default();
499        if limits.max_request_bytes.is_some() {
500            merged.max_request_bytes = limits.max_request_bytes;
501        }
502        if let Some(images) = &limits.images {
503            let mut base = merged.images.clone().unwrap_or_default();
504            if images.max_per_message.is_some() {
505                base.max_per_message = images.max_per_message;
506            }
507            if images.max_per_request.is_some() {
508                base.max_per_request = images.max_per_request;
509            }
510            if images.resize.is_some() {
511                base.resize = images.resize.clone();
512            }
513            merged.images = Some(base);
514        }
515        model.input_limits = Some(merged);
516    }
517    if let Some(cost) = &options.cost {
518        model.cost = Pricing {
519            input: cost.input.unwrap_or(model.cost.input),
520            output: cost.output.unwrap_or(model.cost.output),
521            cache_read: cost.cache_read.unwrap_or(model.cost.cache_read),
522            cache_write: cost.cache_write.unwrap_or(model.cost.cache_write),
523            tiers: cost.tiers.clone().or_else(|| model.cost.tiers.clone()),
524        };
525    }
526    if let Some(cache) = &options.prompt_cache {
527        let mut merged = model.prompt_cache.clone().unwrap_or_default();
528        if cache.short.is_some() {
529            merged.short = cache.short;
530        }
531        if cache.long.is_some() {
532            merged.long = cache.long;
533        }
534        model.prompt_cache = Some(merged);
535    }
536    if let Some(window) = options.context_window {
537        model.context_window = window;
538    }
539    if let Some(max) = options.max_tokens {
540        model.max_tokens = max;
541    }
542    if let Some(params) = &options.sampling_params {
543        let mut merged = model.sampling_params.clone().unwrap_or_default();
544        merged.extend(params.clone());
545        model.sampling_params = Some(merged);
546    }
547    model.compat = merge_compat(model.compat.as_ref(), options.compat.as_ref());
548}
549
550impl ModelRegistry {
551    /// Loads the catalog with `models.json` and `auth.json` from `agent_dir`. It
552    /// ignores a broken `models.json` and reports it through
553    /// [`ModelRegistry::error`].
554    pub fn load(agent_dir: &Path) -> ModelRegistry {
555        let mut registry = ModelRegistry {
556            store: Arc::new(CredentialStore::open(agent_dir.join("auth.json"))),
557            ..ModelRegistry::default()
558        };
559        let models_path = agent_dir.join("models.json");
560        let file = models_path.display();
561        // pi's three failures: reading, parsing and the schema.
562        match std::fs::read_to_string(&models_path) {
563            Ok(text) => {
564                let text = text.strip_prefix('\u{feff}').unwrap_or(&text);
565                match serde_json::from_str::<ModelsConfig>(&strip_json_comments(text)) {
566                    Ok(config) => registry.config = config,
567                    Err(err) if err.is_data() => {
568                        registry.error = Some(format!(
569                            "Invalid models.json schema:\n  - {err}\n\nFile: {file}"
570                        ));
571                    }
572                    Err(err) => {
573                        registry.error = Some(format!(
574                            "Failed to parse models.json: {err}\n\nFile: {file}"
575                        ));
576                    }
577                }
578            }
579            Err(err) if err.kind() == std::io::ErrorKind::NotFound => {}
580            Err(err) => {
581                registry.error = Some(format!("Failed to load models.json: {err}\n\nFile: {file}"));
582            }
583        }
584        registry.models_store = Some(ModelsStore::new(agent_dir.join("models-store.json")));
585        registry.restore_catalogs();
586        registry.rebuild();
587        registry
588    }
589
590    /// Restores the catalogs earlier refreshes stored, and catalogs older
591    /// Radius sign-ins kept in their credential.
592    fn restore_catalogs(&mut self) {
593        let Some(store) = &self.models_store else {
594            return;
595        };
596        let stored = store.read_all();
597        for (provider, source) in self.catalog_sources() {
598            if let Some(entry) = stored.get(&provider) {
599                self.dynamic.insert(
600                    provider.clone(),
601                    model_catalog::restore(&provider, &source, entry),
602                );
603            } else if matches!(source, Source::Radius { .. })
604                && let Some(Credential::OAuth(credential)) = self.store.get(&provider)
605                && let Some(config) = credential.extra.get("gatewayConfig")
606            {
607                let models = legacy_radius_models(&provider, config);
608                if !models.is_empty() {
609                    self.dynamic.insert(provider, ProviderModels::chat(models));
610                }
611            }
612        }
613    }
614
615    /// Providers whose catalogs refresh, and from where: pi.dev for the
616    /// built-ins, the gateway for Radius and `models.json` Radius gateways.
617    fn catalog_sources(&self) -> Vec<(String, Source)> {
618        let mut sources: Vec<(String, Source)> = catalog::builtin_providers()
619            .filter(|provider| *provider != "radius")
620            .map(|provider| (provider.to_owned(), Source::Remote))
621            .collect();
622        sources.push((
623            "radius".into(),
624            Source::Radius {
625                gateway: crate::auth::radius::DEFAULT_GATEWAY.into(),
626            },
627        ));
628        for provider in self.config.providers.keys() {
629            if let Some(gateway) = self.radius_gateway(provider) {
630                let gateway = crate::auth::radius::normalize_gateway(&gateway);
631                sources.retain(|(id, _)| id != provider);
632                sources.push((provider.clone(), Source::Radius { gateway }));
633            }
634        }
635        if self.llama {
636            // pi refreshes llama.cpp only for a stored server.
637            let server = match self.store.get(crate::llama::PROVIDER_ID) {
638                Some(Credential::ApiKey(credential)) => credential
639                    .env
640                    .as_ref()
641                    .and_then(|env| env.get(crate::llama::BASE_URL_ENV))
642                    .and_then(|url| crate::llama::normalize_server_url(url).ok()),
643                _ => None,
644            };
645            sources.push((
646                crate::llama::PROVIDER_ID.into(),
647                Source::Llama {
648                    server: server.unwrap_or_default(),
649                },
650            ));
651        }
652        sources
653    }
654
655    /// Adds pi's llama.cpp provider, as its built-in extension does.
656    pub fn enable_llama(&mut self) {
657        if self.llama {
658            return;
659        }
660        self.llama = true;
661        // Without a stored catalog the provider has no models yet, so the
662        // model list stays as it is.
663        if let Some(store) = &self.models_store
664            && let Some(entry) = store.read(crate::llama::PROVIDER_ID)
665        {
666            let source = Source::Llama {
667                server: String::new(),
668            };
669            self.dynamic.insert(
670                crate::llama::PROVIDER_ID.into(),
671                model_catalog::restore(crate::llama::PROVIDER_ID, &source, &entry),
672            );
673            self.rebuild();
674        }
675    }
676
677    /// Whether the llama.cpp provider is present.
678    pub fn llama_enabled(&self) -> bool {
679        self.llama
680    }
681
682    /// The llama.cpp server and its key, when the provider is configured.
683    pub async fn llama_server(&self) -> Option<(String, Option<String>)> {
684        if !self.llama {
685            return None;
686        }
687        let credential = match self.store.get(crate::llama::PROVIDER_ID) {
688            Some(Credential::ApiKey(credential)) => Some(credential),
689            _ => None,
690        };
691        let resolved = key_auth::resolve(
692            crate::llama::PROVIDER_ID,
693            credential.as_ref(),
694            &Ambient::process(),
695        )?;
696        let server = resolved
697            .env
698            .as_ref()
699            .and_then(|env| env.get(crate::llama::BASE_URL_ENV))
700            .cloned()?;
701        // A stored key, or the environment's. pi sends no "local" placeholder.
702        let key = credential
703            .and_then(|credential| credential.key)
704            .or_else(|| std::env::var("LLAMA_API_KEY").ok());
705        Some((server, key))
706    }
707
708    /// What a catalog refresh needs: each refreshable provider, whether it
709    /// is configured, and its token for gateways that require one.
710    pub async fn catalog_targets(&self) -> Vec<Target> {
711        let mut targets = Vec::new();
712        for (provider, source) in self.catalog_sources() {
713            let configured = self.has_auth(&provider);
714            let token = match (&source, configured) {
715                (Source::Radius { .. }, true) => self.provider_token(&provider).await,
716                // pi sends the stored key only.
717                (Source::Llama { .. }, true) => match self.store.get(&provider) {
718                    Some(Credential::ApiKey(credential)) => credential.key,
719                    _ => None,
720                },
721                _ => None,
722            };
723            let configured =
724                configured && !matches!(&source, Source::Llama { server } if server.is_empty());
725            targets.push(Target {
726                provider,
727                source,
728                token,
729                configured,
730            });
731        }
732        targets
733    }
734
735    /// Where refreshed catalogs persist, when the registry was loaded from an
736    /// agent directory.
737    pub fn models_store(&self) -> Option<&ModelsStore> {
738        self.models_store.as_ref()
739    }
740
741    /// Replaces the refreshed catalogs of the providers in `models`.
742    pub fn apply_catalogs(&mut self, models: IndexMap<String, ProviderModels>) {
743        self.dynamic.extend(models);
744        self.rebuild();
745    }
746
747    /// A token for provider-level calls such as catalog refreshes: the API key
748    /// or bearer token a request would send, refreshed first when it is an
749    /// expiring OAuth token.
750    pub async fn provider_token(&self, provider: &str) -> Option<String> {
751        let auth = self.provider_auth(provider).await?;
752        if auth.error.is_some() {
753            return None;
754        }
755        auth.api_key.or_else(|| {
756            auth.headers
757                .iter()
758                .find(|(name, _)| name.eq_ignore_ascii_case("authorization"))
759                .and_then(|(_, value)| value.as_deref())
760                .and_then(|value| value.strip_prefix("Bearer "))
761                .map(str::to_owned)
762        })
763    }
764
765    /// Credentials for a provider-level call: those of the provider's first
766    /// model, or of a placeholder model of it.
767    async fn provider_auth(&self, provider: &str) -> Option<Auth> {
768        let model = self
769            .models
770            .iter()
771            .find(|model| model.provider == provider)
772            .cloned()
773            .or_else(|| placeholder_model(provider))?;
774        Some(self.auth(&model).await)
775    }
776
777    /// Credentials for `provider` as a whole, with its `models.json` headers
778    /// but no model's, as pi's `getAuth(providerId)`. `None` for a provider
779    /// the registry does not know.
780    pub async fn provider_level_auth(&self, provider: &str) -> Option<Auth> {
781        let known = providers::info(provider).is_some()
782            || self.config.providers.contains_key(provider)
783            || self.oauth.contains_key(provider)
784            || self.models.iter().any(|model| model.provider == provider)
785            || (self.llama && provider == crate::llama::PROVIDER_ID);
786        if !known {
787            return None;
788        }
789        Some(self.auth(&placeholder_model(provider)?).await)
790    }
791
792    /// The catalog alone, with no files.
793    pub fn builtin() -> ModelRegistry {
794        let mut registry = ModelRegistry::default();
795        registry.rebuild();
796        registry
797    }
798
799    fn rebuild(&mut self) {
800        let mut models = catalog::all_builtin_models();
801        let mut classifiers = catalog::all_builtin_classifiers();
802        let mut images = catalog::all_builtin_image_models();
803        for (provider, refreshed) in &self.dynamic {
804            overlay(&mut models, &refreshed.chat, |model| {
805                (model.provider == *provider).then_some(model.id.as_str())
806            });
807            overlay(&mut classifiers, &refreshed.classifiers, |model| {
808                (model.provider == *provider).then_some(model.id.as_str())
809            });
810            overlay(&mut images, &refreshed.images, |model| {
811                (model.provider == *provider).then_some(model.id.as_str())
812            });
813        }
814        let mut errors = Vec::new();
815        let unconfigured = ProviderConfig::default();
816        for (provider, definitions) in &self.extension_typed {
817            let config = self.config.providers.get(provider).unwrap_or(&unconfigured);
818            match typed_models(provider, config, definitions, &classifiers, &images) {
819                Ok((own_classifiers, own_images)) => {
820                    classifiers.retain(|model| model.provider != *provider);
821                    classifiers.extend(own_classifiers);
822                    images.retain(|model| model.provider != *provider);
823                    images.extend(own_images);
824                }
825                Err(err) => errors.push(err),
826            }
827        }
828        self.classifiers = classifiers;
829        self.images = images;
830        let mut providers: Vec<String> = catalog::builtin_providers().map(str::to_owned).collect();
831        for id in self.config.providers.keys() {
832            if !providers.contains(id) {
833                providers.push(id.clone());
834            }
835        }
836        if self.llama && !providers.iter().any(|id| id == crate::llama::PROVIDER_ID) {
837            providers.push(crate::llama::PROVIDER_ID.into());
838        }
839        for (provider, config) in self.config.providers.clone() {
840            let (mut own, others): (Vec<Model>, Vec<Model>) = models
841                .into_iter()
842                .partition(|model| model.provider == provider);
843            models = others;
844            for model in &mut own {
845                if config.oauth.as_deref() != Some("radius")
846                    && let Some(base_url) = &config.base_url
847                {
848                    model.base_url = base_url.clone();
849                }
850                model.compat = merge_compat(model.compat.as_ref(), config.compat.as_ref());
851            }
852            for definition in config.models.iter().flatten() {
853                let defaults = own
854                    .iter()
855                    .find(|model| model.id == definition.id)
856                    .or_else(|| {
857                        let api = definition.api.as_ref().or(config.api.as_ref());
858                        api.and_then(|api| own.iter().find(|model| &model.api == api))
859                    })
860                    .or_else(|| own.iter().find(|model| model.api == "openai-completions"))
861                    .or_else(|| own.first())
862                    .cloned();
863                match model_from_definition(&provider, definition, &config, defaults.as_ref()) {
864                    Ok(model) => match own.iter().position(|existing| existing.id == model.id) {
865                        Some(index) => own[index] = model,
866                        None => own.push(model),
867                    },
868                    Err(err) => errors.push(err),
869                }
870            }
871            for model in &mut own {
872                if let Some(over) = config
873                    .model_overrides
874                    .as_ref()
875                    .and_then(|o| o.get(&model.id))
876                {
877                    apply_override(model, over);
878                }
879            }
880            models.extend(own);
881        }
882        for (provider, registered) in &self.extension_models {
883            models.retain(|model| model.provider != *provider);
884            models.extend(registered.iter().cloned());
885        }
886        // Keep provider order: built-ins in catalog order, then custom providers.
887        models.sort_by_cached_key(|model| providers.iter().position(|p| *p == model.provider));
888        self.models = models;
889        if !errors.is_empty() && self.error.is_none() {
890            self.error = Some(errors.join("\n"));
891        }
892    }
893
894    /// Why `models.json` could not be applied, if it could not.
895    pub fn error(&self) -> Option<&str> {
896        self.error.as_deref()
897    }
898
899    /// Every chat model.
900    pub fn models(&self) -> &[Model] {
901        &self.models
902    }
903
904    /// A model by provider and id.
905    pub fn find(&self, provider: &str, id: &str) -> Option<&Model> {
906        self.models
907            .iter()
908            .find(|model| model.provider == provider && model.id == id)
909    }
910
911    /// Every classifier model, built-in and refreshed.
912    pub fn classifiers(&self) -> &[ClassifierModel] {
913        &self.classifiers
914    }
915
916    /// Every image-generation model, built-in and refreshed.
917    pub fn image_models(&self) -> &[ImageModel] {
918        &self.images
919    }
920
921    /// Classifier models whose providers are configured and keep them.
922    pub fn available_classifiers(&self) -> Vec<&ClassifierModel> {
923        let mut models = self.configured(&self.classifiers, |model| &model.provider);
924        models.retain(|model| self.keeps(&model.provider, Some("classifier"), &model.id));
925        models
926    }
927
928    /// Image models whose providers are configured and keep them.
929    pub fn available_image_models(&self) -> Vec<&ImageModel> {
930        let mut models = self.configured(&self.images, |model| &model.provider);
931        models.retain(|model| self.keeps(&model.provider, Some("image"), &model.id));
932        models
933    }
934
935    /// The `models` whose providers are configured.
936    fn configured<'a, T>(&self, models: &'a [T], provider: fn(&T) -> &String) -> Vec<&'a T> {
937        let mut configured: HashMap<&str, bool> = HashMap::new();
938        models
939            .iter()
940            .filter(|model| {
941                *configured
942                    .entry(provider(model))
943                    .or_insert_with(|| self.has_auth(provider(model)))
944            })
945            .collect()
946    }
947
948    /// Whether `provider` keeps its model `id` for the stored credential: of
949    /// `kind` across every type, or as a chat model when `kind` is `None`.
950    fn keeps(&self, provider: &str, kind: Option<&str>, id: &str) -> bool {
951        let Some(filter) = self.filters.get(provider) else {
952            return true;
953        };
954        match kind {
955            None => filter.chat.as_ref().is_none_or(|ids| ids.contains(id)),
956            Some(kind) => filter
957                .all
958                .as_ref()
959                .is_none_or(|all| all.get(kind).is_some_and(|ids| ids.contains(id))),
960        }
961    }
962
963    /// Credentials for a provider-level request with `base_url`, `api` and
964    /// static `headers`: pi's `applyAuth` for any model type. Fails with pi's
965    /// message when the provider is not configured.
966    async fn model_auth(
967        &self,
968        provider: &str,
969        id: &str,
970        api: &str,
971        base_url: &str,
972        headers: Option<&IndexMap<String, String>>,
973    ) -> Result<Auth, String> {
974        let mut model = placeholder_model(provider).ok_or("invalid model")?;
975        model.id = id.to_owned();
976        model.api = api.to_owned();
977        model.base_url = base_url.to_owned();
978        model.headers = headers.cloned();
979        let auth = self.auth(&model).await;
980        if let Some(error) = auth.error {
981            return Err(error);
982        }
983        if auth.api_key.is_none() && auth.source.is_none() {
984            return Err(format!("Provider is not configured: {provider}"));
985        }
986        Ok(auth)
987    }
988
989    /// Classifies `context` with `model`, resolving the provider's
990    /// credentials (pi's `Models.classify`). Failures are in the result.
991    pub async fn classify(
992        &self,
993        model: &ClassifierModel,
994        context: &ClassifierContext,
995        mut options: ClassifyOptions,
996    ) -> ClassifierResult {
997        let mut model = model.clone();
998        let auth = self
999            .model_auth(
1000                &model.provider,
1001                &model.id,
1002                &model.api,
1003                &model.base_url,
1004                model.headers.as_ref(),
1005            )
1006            .await;
1007        match auth {
1008            Ok(auth) => auth.apply_under(
1009                &mut model.base_url,
1010                &mut options.api_key,
1011                &mut options.headers,
1012            ),
1013            Err(message) => {
1014                return crate::api::classify::error_result(&model, message, &options.cancel);
1015            }
1016        }
1017        match self.operations_for(&model.provider, &model.api) {
1018            Some(own) => own.classify(&model, context, &options).await,
1019            None => crate::api::classify::classify(&model, context, &options).await,
1020        }
1021    }
1022
1023    /// The extension implementation of `provider`'s classifier or image API `api`.
1024    fn operations_for(&self, provider: &str, api: &str) -> Option<&Arc<dyn ModelOperations>> {
1025        self.operations.get(&(provider.to_owned(), api.to_owned()))
1026    }
1027
1028    /// Generates images with `model`, resolving the provider's credentials
1029    /// (pi's `Models.generateImages`). Failures are in the result.
1030    pub async fn generate_images(
1031        &self,
1032        model: &ImageModel,
1033        context: &ImagesContext,
1034        mut options: ImagesOptions,
1035    ) -> AssistantImages {
1036        let mut model = model.clone();
1037        let auth = self
1038            .model_auth(
1039                &model.provider,
1040                &model.id,
1041                &model.api,
1042                &model.base_url,
1043                model.headers.as_ref(),
1044            )
1045            .await;
1046        match auth {
1047            Ok(auth) => auth.apply_under(
1048                &mut model.base_url,
1049                &mut options.api_key,
1050                &mut options.headers,
1051            ),
1052            Err(message) => {
1053                return crate::api::images::error_result(&model, message, &options.cancel);
1054            }
1055        }
1056        match self.operations_for(&model.provider, &model.api) {
1057            Some(own) => own.generate_images(&model, context, &options).await,
1058            None => crate::api::images::generate_images(&model, context, &options).await,
1059        }
1060    }
1061
1062    /// Adds a provider an extension configures, as a `models.json` entry
1063    /// that replaces any of the same name (pi's `registerProvider`).
1064    pub fn register_config(&mut self, provider: &str, mut config: ProviderConfig) {
1065        if config.api_key.is_some() {
1066            self.extension_keys.insert(provider.to_owned());
1067        }
1068        if self.extension_providers.insert(provider.to_owned()) {
1069            let replaced = self.config.providers.get(provider).cloned();
1070            self.replaced.insert(provider.to_owned(), replaced);
1071        }
1072        match config.models.as_mut() {
1073            Some(models) => {
1074                let typed = models.extract_if(.., |model| is_typed(model)).collect();
1075                self.extension_typed.insert(provider.to_owned(), typed);
1076            }
1077            None => {
1078                self.extension_typed.shift_remove(provider);
1079            }
1080        }
1081        self.config.providers.insert(provider.to_owned(), config);
1082        self.rebuild();
1083    }
1084
1085    /// Whether an extension registered `provider` and has not unregistered it.
1086    pub fn is_extension_provider(&self, provider: &str) -> bool {
1087        self.extension_providers.contains(provider)
1088    }
1089
1090    /// Removes the provider an extension registered as `provider`, with its
1091    /// sign-ins, availability and implementations, and restores the
1092    /// built-in or `models.json` provider it replaced (pi's
1093    /// `unregisterProvider`). Does nothing for a provider no extension
1094    /// registered.
1095    pub fn unregister_provider(&mut self, provider: &str) {
1096        if !self.extension_providers.remove(provider) {
1097            return;
1098        }
1099        self.extension_keys.remove(provider);
1100        self.extension_typed.shift_remove(provider);
1101        self.extension_models.shift_remove(provider);
1102        self.oauth.shift_remove(provider);
1103        self.key_auths.shift_remove(provider);
1104        self.filters.shift_remove(provider);
1105        self.operations.retain(|(owner, _), _| owner != provider);
1106        match self.replaced.remove(provider).flatten() {
1107            Some(config) => self.config.providers.insert(provider.to_owned(), config),
1108            None => self.config.providers.shift_remove(provider),
1109        };
1110        self.rebuild();
1111    }
1112
1113    /// Replaces the models an extension provider configures with ones its
1114    /// `refreshModels` returned. Fails, changing nothing, when a model lacks
1115    /// an API or base URL, as pi validates a refreshed list before using it.
1116    pub fn replace_models(
1117        &mut self,
1118        provider: &str,
1119        mut models: Vec<ModelDefinition>,
1120    ) -> Result<(), String> {
1121        let config = self
1122            .config
1123            .providers
1124            .get_mut(provider)
1125            .ok_or_else(|| format!("Unknown provider: {provider}"))?;
1126        // Any built-in model supplies the API and base URL a definition leaves out.
1127        let defaults = catalog::builtin_models(provider).into_iter().next();
1128        let typed: Vec<ModelDefinition> = models.extract_if(.., |model| is_typed(model)).collect();
1129        for definition in &models {
1130            model_from_definition(provider, definition, config, defaults.as_ref())?;
1131        }
1132        typed_models(
1133            provider,
1134            config,
1135            &typed,
1136            &catalog::all_builtin_classifiers(),
1137            &catalog::all_builtin_image_models(),
1138        )?;
1139        config.models = Some(models);
1140        self.extension_typed.insert(provider.to_owned(), typed);
1141        self.rebuild();
1142        Ok(())
1143    }
1144
1145    /// Uses `implementation` for `provider`'s classifier and image models of
1146    /// `apis` (pi's `classifiers` and `images` of an extension's provider).
1147    pub fn register_operations(
1148        &mut self,
1149        provider: &str,
1150        apis: impl IntoIterator<Item = String>,
1151        implementation: Arc<dyn ModelOperations>,
1152    ) {
1153        for api in apis {
1154            let key = (provider.to_owned(), api);
1155            self.operations.insert(key, Arc::clone(&implementation));
1156        }
1157    }
1158
1159    /// Resolves `provider`'s API key with `auth` instead of the environment,
1160    /// and signs in with its `login`. The provider is not configured without
1161    /// a stored credential until [`ModelRegistry::set_availability`] finds a key.
1162    pub fn register_key_auth(&mut self, provider: &str, auth: Arc<dyn KeyAuth>) {
1163        self.key_auths.insert(provider.to_owned(), (auth, None));
1164    }
1165
1166    /// `provider`'s API-key auth from an extension, if it has one.
1167    pub fn key_auth(&self, provider: &str) -> Option<&Arc<dyn KeyAuth>> {
1168        self.key_auths.get(provider).map(|(auth, _)| auth)
1169    }
1170
1171    /// Applies what extension provider `provider` makes available for the
1172    /// stored credential: where its key auth finds a key, and the models
1173    /// its filters keep.
1174    pub fn set_availability(&mut self, provider: &str, availability: Availability) {
1175        if let Some((_, source)) = self.key_auths.get_mut(provider) {
1176            source.clone_from(&availability.source);
1177        }
1178        self.filters.insert(provider.to_owned(), availability);
1179    }
1180
1181    /// Uses `models` as `provider`'s chat models, as an extension's legacy
1182    /// `oauth.modifyModels` made them for its signed-in account. `None`
1183    /// restores the provider's own.
1184    pub fn modify_models(&mut self, provider: &str, models: Option<Vec<Model>>) {
1185        match models {
1186            Some(models) => self.register_provider(provider, models),
1187            None => {
1188                if self.extension_models.shift_remove(provider).is_some() {
1189                    self.rebuild();
1190                }
1191            }
1192        }
1193    }
1194
1195    /// The credential a catalog refresh of `provider` passes to the network,
1196    /// as pi's `resolveRefreshCredential`: a stored OAuth credential,
1197    /// refreshed when it has expired, else the API key a request would use.
1198    pub async fn refresh_credential(&self, provider: &str) -> Option<Credential> {
1199        if let Some(Credential::OAuth(stored)) = self.credential(provider) {
1200            let flow = self.oauth_flow(provider)?;
1201            let refreshed = self.refreshed_oauth(provider, flow.as_ref(), stored, 0);
1202            return refreshed.await.ok().flatten().map(Credential::OAuth);
1203        }
1204        let auth = self.provider_auth(provider).await?;
1205        Some(Credential::ApiKey(ApiKeyCredential {
1206            key: Some(auth.api_key?),
1207            env: auth.env,
1208        }))
1209    }
1210
1211    /// Adds models from an extension provider, replacing the provider's models.
1212    pub fn register_provider(&mut self, provider: &str, models: Vec<Model>) {
1213        self.extension_models
1214            .insert(provider.to_owned(), models.clone());
1215        self.models.retain(|model| model.provider != provider);
1216        self.models.extend(models);
1217    }
1218
1219    /// Uses `flow` for `provider`'s sign-in and refresh, replacing the built-in.
1220    pub fn register_oauth(&mut self, provider: &str, flow: Arc<dyn OAuthProvider>) {
1221        self.oauth.insert(provider.to_owned(), flow);
1222    }
1223
1224    /// The sign-in for `provider`: a registered one, a `models.json` Radius
1225    /// gateway's, else yapi's built-in.
1226    pub fn oauth_flow(&self, provider: &str) -> Option<Arc<dyn OAuthProvider>> {
1227        if let Some(flow) = self.oauth.get(provider) {
1228            return Some(Arc::clone(flow));
1229        }
1230        if let Some(gateway) = self.radius_gateway(provider) {
1231            let name = self.provider_name(provider);
1232            return Some(Arc::new(crate::auth::radius::RadiusOAuth::new(
1233                &name, &gateway,
1234            )));
1235        }
1236        builtin_oauth(provider)
1237    }
1238
1239    /// The gateway of a `models.json` provider with `"oauth": "radius"`: its
1240    /// base URL without a trailing `/v1`.
1241    pub fn radius_gateway(&self, provider: &str) -> Option<String> {
1242        let config = self.config.providers.get(provider)?;
1243        if config.oauth.as_deref() != Some("radius") {
1244            return None;
1245        }
1246        let base = config.base_url.as_deref()?.trim_end_matches('/');
1247        Some(base.strip_suffix("/v1").unwrap_or(base).to_owned())
1248    }
1249
1250    /// Uses `key` for `provider` for this process, ahead of every stored credential.
1251    pub fn set_runtime_key(&mut self, provider: &str, key: String) {
1252        self.runtime_keys.insert(provider.to_owned(), key);
1253    }
1254
1255    fn credential(&self, provider: &str) -> Option<Credential> {
1256        self.store.get(provider)
1257    }
1258
1259    /// `provider`'s credential as pi's `RuntimeCredentials.read` gives it to
1260    /// providers: a runtime key as an API key credential, else the stored one.
1261    pub fn read_credential(&self, provider: &str) -> Option<Credential> {
1262        match self.runtime_keys.get(provider) {
1263            Some(key) => Some(Credential::ApiKey(ApiKeyCredential {
1264                key: Some(key.clone()),
1265                env: None,
1266            })),
1267            None => self.credential(provider),
1268        }
1269    }
1270
1271    /// The credential store behind this registry.
1272    pub fn store(&self) -> &Arc<CredentialStore> {
1273        &self.store
1274    }
1275
1276    /// Whether some credential is configured for `provider`, without running
1277    /// commands.
1278    pub fn has_auth(&self, provider: &str) -> bool {
1279        self.auth_source(provider).is_some()
1280    }
1281
1282    /// pi's failure for `provider` while `auth.json` cannot be read: requests
1283    /// that would consult it fail instead of using other credentials.
1284    pub fn store_error(&self, provider: &str) -> Option<String> {
1285        if self.runtime_keys.contains_key(provider) {
1286            return None;
1287        }
1288        self.store
1289            .read_error()
1290            .map(|error| format!("Credential store read failed for {provider}: {error}"))
1291    }
1292
1293    /// Whether `provider`'s credential comes from a `!command`, whose output
1294    /// decides whether there is a key at all.
1295    pub fn uses_command_key(&self, provider: &str) -> bool {
1296        if self.runtime_keys.contains_key(provider) {
1297            return false;
1298        }
1299        let stored = match self.credential(provider) {
1300            Some(Credential::ApiKey(credential)) => credential.key,
1301            Some(Credential::OAuth(_)) => return false,
1302            None => None,
1303        };
1304        let configured = self
1305            .config
1306            .providers
1307            .get(provider)
1308            .and_then(|config| config.api_key.clone());
1309        [stored, configured]
1310            .iter()
1311            .flatten()
1312            .any(|key| credentials::is_command(key))
1313    }
1314
1315    /// Where `provider`'s credential comes from, without running commands, or
1316    /// `None` when it has none.
1317    fn auth_source(&self, provider: &str) -> Option<String> {
1318        if self.runtime_keys.contains_key(provider) {
1319            return Some("--api-key".into());
1320        }
1321        match self.credential(provider) {
1322            Some(Credential::ApiKey(credential)) => {
1323                if key_auth::is_custom(provider) {
1324                    // Whether the key resolves is unknown without running its
1325                    // command, so a command counts as a key here.
1326                    let key = credential.key.clone().filter(|key| {
1327                        credentials::is_command(key)
1328                            || credentials::is_configured(key, credential.env.as_ref())
1329                    });
1330                    let credential = ApiKeyCredential { key, ..credential };
1331                    return key_auth::resolve(provider, Some(&credential), &Ambient::process())
1332                        .map(|_| "stored credential".into());
1333                }
1334                // An empty stored key counts as none, as in pi.
1335                if let Some(key) = credential.key.as_ref().filter(|key| !key.is_empty())
1336                    && (credentials::is_command(key)
1337                        || credentials::is_configured(key, credential.env.as_ref()))
1338                {
1339                    return Some("stored credential".into());
1340                }
1341            }
1342            Some(Credential::OAuth(_)) => return Some("OAuth".into()),
1343            None => {}
1344        }
1345        // A `models.json` key replaces the provider's own resolution, so the
1346        // environment is not consulted when it cannot be resolved.
1347        if let Some(key) = self.configured_key(provider) {
1348            return (credentials::is_command(key) || credentials::is_configured(key, None))
1349                .then(|| "configured API key".into());
1350        }
1351        self.ambient_source(provider)
1352    }
1353
1354    /// How `provider` is configured, as pi's `getProviderAuthStatus` labels
1355    /// it in `/login`: `runtime`, `stored`, `models_json_command`,
1356    /// `models_json_key`, `fallback` (an extension's key), or the environment
1357    /// variables the key comes from. `None` when it is not configured.
1358    pub fn login_status(&self, provider: &str) -> Option<String> {
1359        if self.runtime_keys.contains_key(provider) {
1360            return Some("runtime".into());
1361        }
1362        if self.store.list().iter().any(|(id, _)| id == provider) {
1363            return Some("stored".into());
1364        }
1365        let configured = self
1366            .config
1367            .providers
1368            .get(provider)
1369            .and_then(|config| config.api_key.as_ref());
1370        if let Some(key) = configured {
1371            if credentials::is_command(key) {
1372                return Some("models_json_command".into());
1373            }
1374            let names = credentials::env_var_names(key);
1375            if !names.is_empty() {
1376                return credentials::is_configured(key, None).then(|| names.join(", "));
1377            }
1378            let source = if self.extension_keys.contains(provider) {
1379                "fallback"
1380            } else {
1381                "models_json_key"
1382            };
1383            return Some(source.into());
1384        }
1385        self.ambient_source(provider)
1386    }
1387
1388    /// Where `provider`'s credential comes from when nothing is stored or
1389    /// configured: its extension's key auth, else [`environment_source`].
1390    fn ambient_source(&self, provider: &str) -> Option<String> {
1391        match self.key_auths.get(provider) {
1392            Some((_, source)) => source.clone(),
1393            None => environment_source(provider),
1394        }
1395    }
1396
1397    /// The `apiKey` that `models.json` configures for `provider`, if any.
1398    fn configured_key(&self, provider: &str) -> Option<&String> {
1399        self.config
1400            .providers
1401            .get(provider)
1402            .and_then(|config| config.api_key.as_ref())
1403            .filter(|key| !key.is_empty())
1404    }
1405
1406    /// Whether `provider` authenticates with a stored OAuth credential. A
1407    /// runtime key takes its place.
1408    pub fn is_using_oauth(&self, provider: &str) -> bool {
1409        !self.runtime_keys.contains_key(provider)
1410            && matches!(self.credential(provider), Some(Credential::OAuth(_)))
1411    }
1412
1413    /// Stored credentials by provider, in file order.
1414    pub fn stored_credentials(&self) -> Vec<(String, CredentialKind)> {
1415        self.store.list()
1416    }
1417
1418    /// Models with credentials, as pi's `getAvailable`. A GitHub Copilot
1419    /// sign-in limits its models to those the account enables, and an
1420    /// extension provider's `filterModels` to those it keeps.
1421    pub fn available(&self) -> Vec<&Model> {
1422        self.available_chat(None)
1423    }
1424
1425    /// Chat models with credentials as pi's `getAvailableOfType("chat")`
1426    /// lists them, where a provider's `filterAllModels` decides over its
1427    /// `filterModels`.
1428    pub fn available_of_all_types(&self) -> Vec<&Model> {
1429        self.available_chat(Some("chat"))
1430    }
1431
1432    /// Chat models with credentials that their providers keep, as chat
1433    /// models when `kind` is `None`, else across every type.
1434    fn available_chat(&self, kind: Option<&str>) -> Vec<&Model> {
1435        let copilot_ids = match self.credential("github-copilot") {
1436            Some(Credential::OAuth(credential)) => copilot::available_model_ids(&credential),
1437            _ => None,
1438        };
1439        let mut models = self.configured(&self.models, |model| &model.provider);
1440        models.retain(|model| {
1441            (model.provider != "github-copilot"
1442                || copilot_ids
1443                    .as_ref()
1444                    .is_none_or(|ids| ids.contains(&model.id)))
1445                && self.keeps(&model.provider, kind, &model.id)
1446        });
1447        models
1448    }
1449
1450    /// The API key a request to `provider` would use, when it is known
1451    /// without running a command or signing in: a runtime key, a stored or
1452    /// configured template value, or the environment.
1453    pub fn known_api_key(&self, provider: &str) -> Option<String> {
1454        if let Some(key) = self.runtime_keys.get(provider) {
1455            return Some(key.clone());
1456        }
1457        let mut env = None;
1458        match self.credential(provider) {
1459            Some(Credential::ApiKey(credential)) if key_auth::is_custom(provider) => {
1460                let key = credential
1461                    .key
1462                    .as_deref()
1463                    .and_then(|key| credentials::resolve_template(key, credential.env.as_ref()));
1464                let credential = ApiKeyCredential { key, ..credential };
1465                return key_auth::resolve(provider, Some(&credential), &Ambient::process())
1466                    .and_then(|resolved| resolved.api_key);
1467            }
1468            Some(Credential::ApiKey(credential)) => {
1469                if let Some(value) = credential
1470                    .key
1471                    .as_deref()
1472                    .and_then(|key| credentials::resolve_template(key, credential.env.as_ref()))
1473                    .filter(|value| !value.is_empty())
1474                {
1475                    return Some(value);
1476                }
1477                env = credential.env;
1478            }
1479            Some(Credential::OAuth(_)) => return None,
1480            None => {}
1481        }
1482        if let Some(key) = self.configured_key(provider) {
1483            return credentials::resolve_template(key, None).filter(|value| !value.is_empty());
1484        }
1485        if key_auth::is_custom(provider) {
1486            return key_auth::resolve(provider, None, &Ambient::process())
1487                .and_then(|resolved| resolved.api_key);
1488        }
1489        credentials::env_api_key(provider, env.as_ref()).map(|(_, value)| value)
1490    }
1491
1492    /// Credentials for a request to `model`, with its configured headers. A
1493    /// stored OAuth token about to expire is refreshed first.
1494    pub async fn auth(&self, model: &Model) -> Auth {
1495        self.auth_valid_for(model, None).await
1496    }
1497
1498    /// [`ModelRegistry::auth`] with pi's `minOAuthValidityMs`: a stored OAuth
1499    /// token valid for less than five minutes or `min_validity_ms` is
1500    /// refreshed first, and with `min_validity_ms` a refreshed token valid for
1501    /// less than that is an error.
1502    pub async fn auth_valid_for(&self, model: &Model, min_validity_ms: Option<u64>) -> Auth {
1503        let provider = model.provider.as_str();
1504        let mut auth = Auth {
1505            headers: self.headers(model).await,
1506            ..Auth::default()
1507        };
1508        if let Some(key) = self.runtime_keys.get(provider) {
1509            auth.api_key = Some(key.clone());
1510            // pi reads a runtime key as a stored credential.
1511            auth.source = Some("stored credential".into());
1512            return auth;
1513        }
1514        if let Some(error) = self.store_error(provider) {
1515            auth.error = Some(error);
1516            return auth;
1517        }
1518        let mut env: Option<ProviderEnv> = None;
1519        let stored = self.credential(provider);
1520        if let Some((own, _)) = self.key_auths.get(provider)
1521            && !matches!(stored, Some(Credential::OAuth(_)))
1522        {
1523            let credential = match stored {
1524                Some(Credential::ApiKey(credential)) => Some(resolve_stored(credential).await),
1525                _ => None,
1526            };
1527            match own.resolve(credential.as_ref()).await {
1528                Ok(Some(resolved)) => auth.use_resolved(resolved),
1529                Ok(None) => {}
1530                Err(message) => {
1531                    auth.error = Some(format!(
1532                        "API key auth failed for provider {provider}: {message}"
1533                    ));
1534                }
1535            }
1536            return auth;
1537        }
1538        match stored {
1539            Some(Credential::ApiKey(credential)) if key_auth::is_custom(provider) => {
1540                // The stored credential owns the provider: an unresolved key
1541                // leaves the provider's own fallbacks, not other sources.
1542                let credential = resolve_stored(credential).await;
1543                if let Some(resolved) =
1544                    key_auth::resolve(provider, Some(&credential), &Ambient::process())
1545                {
1546                    auth.use_resolved(resolved);
1547                }
1548                return auth;
1549            }
1550            Some(Credential::ApiKey(credential)) => {
1551                env = credential.env.clone();
1552                if let Some(key) = &credential.key
1553                    && let Some(value) = credentials::resolve(key, env.as_ref(), true)
1554                        .await
1555                        .filter(|value| !value.is_empty())
1556                {
1557                    auth.api_key = Some(value);
1558                    auth.source = Some("stored credential".into());
1559                    return auth;
1560                }
1561            }
1562            Some(Credential::OAuth(credential)) => {
1563                match self.oauth_auth(provider, credential, min_validity_ms).await {
1564                    Ok(Some(oauth)) => {
1565                        auth.api_key = oauth.api_key;
1566                        auth.headers.extend(oauth.headers);
1567                        auth.base_url = oauth.base_url;
1568                        auth.source = Some("OAuth".into());
1569                    }
1570                    Ok(None) => {}
1571                    Err(message) => auth.error = Some(message),
1572                }
1573                return auth;
1574            }
1575            None => {}
1576        }
1577        if let Some(key) = self.configured_key(provider) {
1578            match credentials::resolve(key, None, false).await {
1579                Some(value) if !value.is_empty() => {
1580                    auth.api_key = Some(value);
1581                    auth.source = Some("configured API key".into());
1582                }
1583                // pi resolves the configured key or fails, without falling
1584                // back to the environment.
1585                _ => {
1586                    let description = format!("API key for provider \"{provider}\"");
1587                    auth.error = Some(format!(
1588                        "API key auth failed for provider {provider}: {}",
1589                        credentials::unresolved_message(key, &description, None)
1590                    ));
1591                }
1592            }
1593            return auth;
1594        }
1595        if key_auth::is_custom(provider) {
1596            if let Some(resolved) = key_auth::resolve(provider, None, &Ambient::process()) {
1597                auth.use_resolved(resolved);
1598            }
1599            return auth;
1600        }
1601        if let Some((name, value)) = credentials::env_api_key(provider, env.as_ref()) {
1602            if name == BEARER_TOKEN_ENV {
1603                auth.headers
1604                    .insert("Authorization".into(), Some(format!("Bearer {value}")));
1605            } else {
1606                auth.api_key = Some(value);
1607            }
1608            auth.source = Some(name.to_owned());
1609        } else if provider == "anthropic"
1610            && let Some(resolved) = key_auth::anthropic_federation(&Ambient::process())
1611        {
1612            auth.use_resolved(resolved);
1613        }
1614        auth
1615    }
1616
1617    /// Request credentials from a stored OAuth token, refreshed as
1618    /// [`ModelRegistry::auth_valid_for`] says. `Ok(None)` means the provider
1619    /// was logged out meanwhile.
1620    async fn oauth_auth(
1621        &self,
1622        provider: &str,
1623        stored: OAuthCredential,
1624        min_validity_ms: Option<u64>,
1625    ) -> Result<Option<OAuthAuth>, String> {
1626        let Some(flow) = self.oauth_flow(provider) else {
1627            // yapi has no sign-in for this provider: use the token as stored.
1628            return Ok(Some(OAuthAuth {
1629                api_key: Some(stored.access),
1630                ..OAuthAuth::default()
1631            }));
1632        };
1633        let minimum = min_validity_ms.map_or(OAUTH_MINIMUM_VALIDITY_MS, |min| {
1634            min.max(OAUTH_MINIMUM_VALIDITY_MS)
1635        });
1636        let refreshed = self.refreshed_oauth(provider, flow.as_ref(), stored, minimum);
1637        let Some(credential) = refreshed.await? else {
1638            return Ok(None);
1639        };
1640        if min_validity_ms.is_some()
1641            && (crate::auth::now_ms() + minimum) as f64 >= credential.expires
1642        {
1643            return Err(format!(
1644                "OAuth refresh returned a token that expires too soon for {provider}"
1645            ));
1646        }
1647        flow.to_auth(&credential)
1648            .await
1649            .map(Some)
1650            .map_err(|err| format!("OAuth auth derivation failed for {provider}: {err}"))
1651    }
1652
1653    /// `stored`, or when it expires within `min_validity_ms`, the credential
1654    /// `flow` refreshes it to under the `auth.json` lock, after checking again
1655    /// that no other process refreshed it. `Ok(None)` means the provider was
1656    /// logged out meanwhile.
1657    async fn refreshed_oauth(
1658        &self,
1659        provider: &str,
1660        flow: &dyn OAuthProvider,
1661        stored: OAuthCredential,
1662        min_validity_ms: u64,
1663    ) -> Result<Option<OAuthCredential>, String> {
1664        let expires_soon = |credential: &OAuthCredential| {
1665            (crate::auth::now_ms() + min_validity_ms) as f64 >= credential.expires
1666        };
1667        if !expires_soon(&stored) {
1668            return Ok(Some(stored));
1669        }
1670        let cancel = CancellationToken::new();
1671        let cancel_ref = &cancel;
1672        let refreshed = self
1673            .store
1674            .modify(
1675                provider,
1676                |current| async move {
1677                    let Some(Credential::OAuth(current)) = current else {
1678                        return Ok(None);
1679                    };
1680                    if !expires_soon(&current) {
1681                        return Ok(None);
1682                    }
1683                    let refresh = flow.refresh(&current, cancel_ref);
1684                    match tokio::time::timeout(OAUTH_REFRESH_TIMEOUT, refresh).await {
1685                        Ok(result) => result.map(|next| Some(Credential::OAuth(next))),
1686                        Err(_) => Err(AuthError::failed(
1687                            "The operation was aborted due to timeout",
1688                        )),
1689                    }
1690                },
1691                &cancel,
1692            )
1693            .await
1694            .map_err(|err| format!("OAuth refresh failed for {provider}: {err}"))?;
1695        Ok(match refreshed {
1696            Some(Credential::OAuth(next)) => Some(next),
1697            _ => None,
1698        })
1699    }
1700
1701    /// The display name of a provider: its `models.json` or extension
1702    /// `name`, else the built-in name, else the name of its registered
1703    /// sign-in, else its id.
1704    pub fn provider_name(&self, provider: &str) -> String {
1705        if let Some(name) = self
1706            .config
1707            .providers
1708            .get(provider)
1709            .and_then(|config| config.name.as_ref())
1710        {
1711            return name.clone();
1712        }
1713        providers::info(provider)
1714            .map(|info| info.name.to_owned())
1715            .or_else(|| self.oauth.get(provider).map(|flow| flow.name().to_owned()))
1716            .unwrap_or_else(|| provider.to_owned())
1717    }
1718
1719    /// The `models.json` or extension configuration of `provider`.
1720    pub fn provider_config(&self, provider: &str) -> Option<&ProviderConfig> {
1721        self.config.providers.get(provider)
1722    }
1723
1724    /// Whether an extension registered `provider`.
1725    pub fn registered_by_extension(&self, provider: &str) -> bool {
1726        self.extension_providers.contains(provider)
1727    }
1728
1729    /// Whether `models.json` or an extension sets `authHeader` for `provider`.
1730    pub fn auth_header(&self, provider: &str) -> bool {
1731        self.provider_config(provider)
1732            .and_then(|config| config.auth_header)
1733            .unwrap_or(false)
1734    }
1735
1736    /// Whether `models.json` or an extension configures an `apiKey` for
1737    /// `provider`, even an empty one.
1738    pub fn configures_api_key(&self, provider: &str) -> bool {
1739        self.config
1740            .providers
1741            .get(provider)
1742            .is_some_and(|config| config.api_key.is_some())
1743    }
1744
1745    /// Runs a sign-in and stores the credential. An API key login asks for the
1746    /// key, with an extension provider's own `login` when it has one, and an
1747    /// OAuth login runs the provider's flow.
1748    pub async fn login(
1749        &self,
1750        provider: &str,
1751        kind: LoginKind,
1752        interaction: &Interaction,
1753        options: &LoginOptions,
1754    ) -> Result<Credential, AuthError> {
1755        interaction.check()?;
1756        let credential = match (kind, self.key_auth(provider)) {
1757            (LoginKind::OAuth, _) => {
1758                let flow = self.oauth_flow(provider).ok_or_else(|| {
1759                    AuthError::Failed(format!(
1760                        "{} does not support oauth login",
1761                        self.provider_name(provider)
1762                    ))
1763                })?;
1764                Credential::OAuth(flow.login(interaction, options).await?)
1765            }
1766            (LoginKind::ApiKey, Some(own)) => Credential::ApiKey(own.login(interaction).await?),
1767            (LoginKind::ApiKey, None) if key_auth::is_custom(provider) => Credential::ApiKey(
1768                key_auth::login(provider, interaction)
1769                    .await
1770                    .unwrap_or_else(|| Err(AuthError::failed("No login for provider")))?,
1771            ),
1772            (LoginKind::ApiKey, None) => {
1773                let name = providers::info(provider)
1774                    .and_then(|info| info.api_key)
1775                    .map_or("API key", |method| method.name);
1776                let key = interaction
1777                    .prompt(AuthPrompt::Secret {
1778                        message: format!("Enter {name}"),
1779                    })
1780                    .await?;
1781                interaction.check()?;
1782                Credential::ApiKey(ApiKeyCredential {
1783                    key: Some(key),
1784                    env: None,
1785                })
1786            }
1787        };
1788        self.store
1789            .set(provider, credential.clone(), interaction.cancel())
1790            .await
1791            .map_err(|err| match err {
1792                AuthError::Cancelled => AuthError::Cancelled,
1793                AuthError::Failed(message) => AuthError::Failed(format!(
1794                    "Credential store modify failed for {provider}: {message}"
1795                )),
1796            })?;
1797        Ok(credential)
1798    }
1799
1800    /// Removes `provider`'s stored credential.
1801    pub async fn logout(&self, provider: &str) -> Result<(), AuthError> {
1802        self.store
1803            .delete(provider, &CancellationToken::new())
1804            .await
1805            .map_err(|err| {
1806                AuthError::Failed(format!(
1807                    "Credential store delete failed for {provider}: {err}"
1808                ))
1809            })
1810    }
1811
1812    /// `models.json` headers for a model: provider headers, then the model's
1813    /// override and definition headers. Values use the credential forms.
1814    async fn headers(&self, model: &Model) -> IndexMap<String, Option<String>> {
1815        let mut headers = IndexMap::new();
1816        let Some(config) = self.config.providers.get(&model.provider) else {
1817            return headers;
1818        };
1819        let override_headers = config
1820            .model_overrides
1821            .as_ref()
1822            .and_then(|overrides| overrides.get(&model.id))
1823            .and_then(|over| over.options.headers.as_ref());
1824        let definition_headers = config
1825            .models
1826            .iter()
1827            .flatten()
1828            .find(|definition| definition.id == model.id)
1829            .and_then(|definition| definition.options.headers.as_ref());
1830        for source in [
1831            config.headers.as_ref(),
1832            override_headers,
1833            definition_headers,
1834        ]
1835        .into_iter()
1836        .flatten()
1837        {
1838            for (name, value) in source {
1839                headers.insert(name.clone(), credentials::resolve(value, None, false).await);
1840            }
1841        }
1842        headers
1843    }
1844
1845    /// Where `auth.json` lives.
1846    pub fn auth_path(&self) -> Option<&Path> {
1847        self.store.path()
1848    }
1849}
1850
1851#[cfg(test)]
1852mod tests {
1853    use super::*;
1854
1855    #[test]
1856    fn strips_comments_and_commas() {
1857        assert_eq!(
1858            strip_json_comments("{\"a\": \"x//y\", // note\n \"b\": [1, 2,],\n}"),
1859            "{\"a\": \"x//y\", \n \"b\": [1, 2]\n}"
1860        );
1861        // JavaScript's `\s` before the bracket. An escaped quote stays in the string.
1862        assert_eq!(
1863            strip_json_comments("[1,\u{a0}\u{feff}] \"a\\\"// b\""),
1864            "[1\u{a0}\u{feff}] \"a\\\"// b\""
1865        );
1866    }
1867
1868    /// A refreshed list replaces an extension provider's models only when
1869    /// every model has an API and a base URL.
1870    #[test]
1871    fn replace_models_validates_first() {
1872        let config = |value| serde_json::from_value::<ProviderConfig>(value).unwrap();
1873        let models = |value| serde_json::from_value::<Vec<ModelDefinition>>(value).unwrap();
1874        let ids = |registry: &ModelRegistry, provider: &str| -> Vec<String> {
1875            registry
1876                .models()
1877                .iter()
1878                .filter(|model| model.provider == provider)
1879                .map(|model| model.id.clone())
1880                .collect()
1881        };
1882        let mut registry = ModelRegistry::builtin();
1883        registry.register_config(
1884            "x",
1885            config(serde_json::json!({"baseUrl": "http://x", "api": "openai-completions", "models": [{"id": "a"}]})),
1886        );
1887        registry
1888            .replace_models("x", models(serde_json::json!([{"id": "b"}, {"id": "c"}])))
1889            .unwrap();
1890        assert_eq!(ids(&registry, "x"), ["b", "c"]);
1891        registry.register_config(
1892            "y",
1893            config(serde_json::json!({"baseUrl": "http://y", "models": [{"id": "a", "api": "openai-completions"}]})),
1894        );
1895        let error = registry
1896            .replace_models("y", models(serde_json::json!([{"id": "b"}])))
1897            .unwrap_err();
1898        assert_eq!(
1899            error,
1900            "Provider y, model b: no \"api\" specified. Set at provider or model level."
1901        );
1902        assert_eq!(ids(&registry, "y"), ["a"]);
1903    }
1904
1905    /// An extension's image and classifier definitions replace the
1906    /// provider's models of those types, taking a missing API and base URL
1907    /// from its own, and never list as chat models.
1908    #[test]
1909    fn typed_definitions_override() {
1910        let config = |value| serde_json::from_value::<ProviderConfig>(value).unwrap();
1911        let cost = serde_json::json!({"input": 0, "output": 0, "cacheRead": 0, "cacheWrite": 0});
1912        let mut registry = ModelRegistry::builtin();
1913        assert!(
1914            registry
1915                .image_models()
1916                .iter()
1917                .filter(|model| model.provider == "openrouter")
1918                .count()
1919                > 1
1920        );
1921        registry.register_config(
1922            "openrouter",
1923            config(serde_json::json!({"models": [
1924                {"type": "image", "id": "draw", "name": "Draw", "input": ["text"], "output": ["image"], "cost": cost},
1925                {"type": "classifier", "id": "judge", "name": "Judge", "api": "x-classify", "baseUrl": "http://x", "input": ["text"], "cost": cost, "contextWindow": 10},
1926                {"type": "embedding", "id": "embed"},
1927            ]})),
1928        );
1929        let images: Vec<_> = registry
1930            .image_models()
1931            .iter()
1932            .filter(|model| model.provider == "openrouter")
1933            .collect();
1934        assert_eq!(images.len(), 1);
1935        assert_eq!(
1936            (
1937                images[0].id.as_str(),
1938                images[0].api.as_str(),
1939                images[0].base_url.as_str()
1940            ),
1941            ("draw", "openrouter-images", "https://openrouter.ai/api/v1")
1942        );
1943        let classifiers: Vec<_> = registry
1944            .classifiers()
1945            .iter()
1946            .filter(|model| model.provider == "openrouter")
1947            .collect();
1948        assert_eq!(classifiers.len(), 1);
1949        assert_eq!(
1950            (classifiers[0].id.as_str(), classifiers[0].context_window),
1951            ("judge", 10)
1952        );
1953        assert!(
1954            !registry
1955                .models()
1956                .iter()
1957                .any(|model| ["draw", "judge", "embed"].contains(&model.id.as_str()))
1958        );
1959
1960        // A refreshed list is checked first, then replaces them in turn.
1961        let models = |value| serde_json::from_value::<Vec<ModelDefinition>>(value).unwrap();
1962        registry.register_config(
1963            "x",
1964            config(serde_json::json!({"baseUrl": "http://x", "models": []})),
1965        );
1966        let error = registry
1967            .replace_models(
1968                "x",
1969                models(serde_json::json!([{"type": "classifier", "id": "c", "name": "C"}])),
1970            )
1971            .unwrap_err();
1972        assert_eq!(
1973            error,
1974            "Provider x, model c: no \"api\" specified. Set it at model level."
1975        );
1976        registry.replace_models("openrouter", Vec::new()).unwrap();
1977        assert!(
1978            !registry
1979                .image_models()
1980                .iter()
1981                .any(|model| model.provider == "openrouter")
1982        );
1983    }
1984
1985    /// pi's `/login` labels for each credential source, and names from
1986    /// `models.json`.
1987    #[test]
1988    fn login_status_follows_pi() {
1989        let dir = std::env::temp_dir().join(format!("yapi-login-status-{}", std::process::id()));
1990        std::fs::create_dir_all(&dir).unwrap();
1991        std::fs::write(
1992            dir.join("models.json"),
1993            r#"{"providers": {
1994                "openai": {"apiKey": "!printf key", "name": "OpenAI Renamed"},
1995                "mistral": {"apiKey": "$HOME"},
1996                "groq": {"apiKey": "${YAPI_LOGIN_STATUS_UNSET}"},
1997                "xai": {"apiKey": "literal"}
1998            }}"#,
1999        )
2000        .unwrap();
2001        std::fs::write(
2002            dir.join("auth.json"),
2003            r#"{"anthropic": {"type": "api_key", "key": "$YAPI_LOGIN_STATUS_UNSET"},
2004                "github-copilot": {"type": "oauth", "access": "a", "refresh": "r", "expires": 0}}"#,
2005        )
2006        .unwrap();
2007        let mut registry = ModelRegistry::load(&dir);
2008        let status = |registry: &ModelRegistry, id: &str| registry.login_status(id);
2009        assert_eq!(status(&registry, "anthropic").as_deref(), Some("stored"));
2010        assert_eq!(
2011            status(&registry, "openai").as_deref(),
2012            Some("models_json_command")
2013        );
2014        assert_eq!(status(&registry, "mistral").as_deref(), Some("HOME"));
2015        assert_eq!(status(&registry, "groq"), None);
2016        assert_eq!(status(&registry, "xai").as_deref(), Some("models_json_key"));
2017        assert_eq!(registry.provider_name("openai"), "OpenAI Renamed");
2018        assert_eq!(registry.provider_name("mistral"), "Mistral");
2019
2020        assert!(registry.is_using_oauth("github-copilot"));
2021        registry.set_runtime_key("github-copilot", "key".into());
2022        assert_eq!(
2023            status(&registry, "github-copilot").as_deref(),
2024            Some("runtime")
2025        );
2026        assert!(!registry.is_using_oauth("github-copilot"));
2027        std::fs::remove_dir_all(&dir).unwrap();
2028    }
2029}