1use 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
35const OAUTH_MINIMUM_VALIDITY_MS: u64 = 5 * 60 * 1000;
37const OAUTH_REFRESH_TIMEOUT: Duration = Duration::from_secs(15);
38
39#[derive(Clone, Debug, Default, PartialEq)]
41pub struct Auth {
42 pub api_key: Option<String>,
44 pub headers: IndexMap<String, Option<String>>,
46 pub source: Option<String>,
48 pub base_url: Option<String>,
50 pub error: Option<String>,
52 pub env: Option<ProviderEnv>,
55}
56
57impl Auth {
58 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 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 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
109fn 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
125async 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
134fn 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
148fn 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
165fn 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#[derive(Clone, Copy, Debug, PartialEq, Eq)]
187pub enum LoginKind {
188 OAuth,
190 ApiKey,
192}
193
194#[derive(Clone, Debug, Default, PartialEq, serde::Deserialize)]
197pub struct Availability {
198 pub source: Option<String>,
200 pub chat: Option<HashSet<String>>,
203 pub all: Option<HashMap<String, HashSet<String>>>,
206}
207
208#[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 extension_keys: HashSet<String>,
218 extension_providers: HashSet<String>,
220 replaced: HashMap<String, Option<ProviderConfig>>,
223 error: Option<String>,
224 classifiers: Vec<ClassifierModel>,
226 images: Vec<ImageModel>,
228 dynamic: IndexMap<String, ProviderModels>,
230 llama: bool,
233 extension_models: IndexMap<String, Vec<Model>>,
235 extension_typed: IndexMap<String, Vec<ModelDefinition>>,
238 operations: IndexMap<(String, String), Arc<dyn ModelOperations>>,
241 key_auths: IndexMap<String, (Arc<dyn KeyAuth>, Option<String>)>,
244 filters: IndexMap<String, Availability>,
247 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
262pub fn strip_json_comments(input: &str) -> String {
265 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
388fn is_typed(definition: &ModelDefinition) -> bool {
391 definition
392 .kind
393 .as_deref()
394 .is_some_and(|kind| kind != "chat")
395}
396
397fn 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 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 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 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 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 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 pub fn enable_llama(&mut self) {
657 if self.llama {
658 return;
659 }
660 self.llama = true;
661 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 pub fn llama_enabled(&self) -> bool {
679 self.llama
680 }
681
682 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 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 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 (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 pub fn models_store(&self) -> Option<&ModelsStore> {
738 self.models_store.as_ref()
739 }
740
741 pub fn apply_catalogs(&mut self, models: IndexMap<String, ProviderModels>) {
743 self.dynamic.extend(models);
744 self.rebuild();
745 }
746
747 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 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 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 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 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 pub fn error(&self) -> Option<&str> {
896 self.error.as_deref()
897 }
898
899 pub fn models(&self) -> &[Model] {
901 &self.models
902 }
903
904 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 pub fn classifiers(&self) -> &[ClassifierModel] {
913 &self.classifiers
914 }
915
916 pub fn image_models(&self) -> &[ImageModel] {
918 &self.images
919 }
920
921 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 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 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 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 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 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 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 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 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 pub fn is_extension_provider(&self, provider: &str) -> bool {
1087 self.extension_providers.contains(provider)
1088 }
1089
1090 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 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 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 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 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 pub fn key_auth(&self, provider: &str) -> Option<&Arc<dyn KeyAuth>> {
1168 self.key_auths.get(provider).map(|(auth, _)| auth)
1169 }
1170
1171 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 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 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 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 pub fn register_oauth(&mut self, provider: &str, flow: Arc<dyn OAuthProvider>) {
1221 self.oauth.insert(provider.to_owned(), flow);
1222 }
1223
1224 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 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 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 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 pub fn store(&self) -> &Arc<CredentialStore> {
1273 &self.store
1274 }
1275
1276 pub fn has_auth(&self, provider: &str) -> bool {
1279 self.auth_source(provider).is_some()
1280 }
1281
1282 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 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 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 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 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 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 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 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 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 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 pub fn stored_credentials(&self) -> Vec<(String, CredentialKind)> {
1415 self.store.list()
1416 }
1417
1418 pub fn available(&self) -> Vec<&Model> {
1422 self.available_chat(None)
1423 }
1424
1425 pub fn available_of_all_types(&self) -> Vec<&Model> {
1429 self.available_chat(Some("chat"))
1430 }
1431
1432 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 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 pub async fn auth(&self, model: &Model) -> Auth {
1495 self.auth_valid_for(model, None).await
1496 }
1497
1498 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 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 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 _ => {
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 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 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 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(¤t) {
1681 return Ok(None);
1682 }
1683 let refresh = flow.refresh(¤t, 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 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 pub fn provider_config(&self, provider: &str) -> Option<&ProviderConfig> {
1721 self.config.providers.get(provider)
1722 }
1723
1724 pub fn registered_by_extension(&self, provider: &str) -> bool {
1726 self.extension_providers.contains(provider)
1727 }
1728
1729 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 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 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 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 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 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 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 #[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(®istry, "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(®istry, "y"), ["a"]);
1903 }
1904
1905 #[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 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 #[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(®istry, "anthropic").as_deref(), Some("stored"));
2010 assert_eq!(
2011 status(®istry, "openai").as_deref(),
2012 Some("models_json_command")
2013 );
2014 assert_eq!(status(®istry, "mistral").as_deref(), Some("HOME"));
2015 assert_eq!(status(®istry, "groq"), None);
2016 assert_eq!(status(®istry, "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(®istry, "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}