1mod auth;
11mod state;
12
13use std::path::{Path, PathBuf};
14use std::sync::{Arc, PoisonError, RwLock, RwLockReadGuard, RwLockWriteGuard};
15
16use indexmap::IndexMap;
17use oino_ai::{AssistantMessage, JsonValue, Message, Model, ThinkingLevel};
18
19pub use self::auth::{
20 ApiKeyError, AuthCheck, AuthOverrides, AuthResolveError, AuthStatus, AuthStatusSource,
21 ResolvedAuth,
22};
23use self::auth::{AuthPlan, Resolver};
24use self::state::{RegisteredProvider, State};
25use crate::auth::credential::CredentialType;
26use crate::auth::oauth::OAuthRefresher;
27use crate::auth::storage::AuthStorage;
28use crate::cache::{CacheError, ModelCache};
29use crate::compose::{configured_headers, merge_headers, raw_model_headers};
30use crate::config_value::resolve_headers_strict;
31use crate::env::{EnvSource, ProcessEnv};
32use crate::models_json::ModelsJson;
33use crate::registration::{ProviderConfigOverride, ProviderModels, ProviderRegistration};
34use crate::virtual_models::{
35 ModelRoute, RegisteredVirtualModel, RouteError, RouteFailed, RoutePrevious, RouteReason,
36 RouteRequest, RouteResult, VirtualModelDefinition, VirtualModelError, find_latest_response,
37 is_virtual_api,
38};
39
40#[derive(Clone)]
42pub struct ModelRegistryOptions {
43 pub models_json_path: Option<PathBuf>,
45 pub model_cache: Option<ModelCache>,
47 pub env: Arc<dyn EnvSource>,
49}
50
51impl Default for ModelRegistryOptions {
52 fn default() -> Self {
53 Self {
54 models_json_path: None,
55 model_cache: None,
56 env: Arc::new(ProcessEnv),
57 }
58 }
59}
60
61impl std::fmt::Debug for ModelRegistryOptions {
62 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
63 f.debug_struct("ModelRegistryOptions")
64 .field("models_json_path", &self.models_json_path)
65 .field("model_cache", &self.model_cache)
66 .finish_non_exhaustive()
67 }
68}
69
70impl ModelRegistryOptions {
71 #[must_use]
73 pub fn for_agent_dir(agent_dir: &Path) -> Self {
74 Self {
75 models_json_path: Some(agent_dir.join("models.json")),
76 model_cache: Some(ModelCache::for_agent_dir(agent_dir)),
77 env: Arc::new(ProcessEnv),
78 }
79 }
80}
81
82#[derive(Debug, thiserror::Error)]
84pub enum RegistryError {
85 #[error("Provider id must not be empty.")]
87 EmptyProviderId,
88 #[error("Provider {0} is not registered.")]
90 UnknownProvider(String),
91 #[error("Provider {0} has a static model list.")]
93 StaticModels(String),
94 #[error("{0}")]
96 InvalidConfig(String),
97 #[error(transparent)]
99 VirtualModel(#[from] VirtualModelError),
100 #[error(transparent)]
102 Cache(#[from] CacheError),
103}
104
105#[derive(Clone, Default, PartialEq, Eq)]
107pub struct RequestAuth {
108 pub api_key: Option<String>,
110 pub headers: Option<IndexMap<String, String>>,
112}
113
114impl std::fmt::Debug for RequestAuth {
115 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
116 f.debug_struct("RequestAuth")
117 .field("api_key", &self.api_key.as_ref().map(|_| "<redacted>"))
118 .field(
119 "headers",
120 &self.headers.as_ref().map(|h| h.keys().collect::<Vec<_>>()),
121 )
122 .finish()
123 }
124}
125
126#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
128#[error("{message}")]
129pub struct RequestAuthError {
130 pub message: String,
132}
133
134#[derive(Debug, Clone, PartialEq, Eq)]
136pub struct RegistryDiagnostic {
137 pub provider: String,
139 pub message: String,
141}
142
143struct Inner {
144 state: RwLock<State>,
145 auth: Arc<AuthStorage>,
146 env: Arc<dyn EnvSource>,
147 models_json_path: Option<PathBuf>,
148 cache: Option<ModelCache>,
149 refresher: RwLock<Option<Arc<dyn OAuthRefresher>>>,
150}
151
152#[derive(Clone)]
154pub struct ModelRegistry {
155 inner: Arc<Inner>,
156}
157
158impl std::fmt::Debug for ModelRegistry {
159 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
160 f.debug_struct("ModelRegistry")
161 .field("models_json_path", &self.inner.models_json_path)
162 .finish_non_exhaustive()
163 }
164}
165
166impl ModelRegistry {
167 #[must_use]
169 pub fn new(auth: Arc<AuthStorage>, options: ModelRegistryOptions) -> Self {
170 let mut state = State {
171 models_json: ModelsJson::load(options.models_json_path.as_deref()),
172 ..State::default()
173 };
174 state.recompose();
175 Self {
176 inner: Arc::new(Inner {
177 state: RwLock::new(state),
178 auth,
179 env: options.env,
180 models_json_path: options.models_json_path,
181 cache: options.model_cache,
182 refresher: RwLock::new(None),
183 }),
184 }
185 }
186
187 #[must_use]
189 pub fn in_memory(env: Arc<dyn EnvSource>) -> Self {
190 let auth = Arc::new(AuthStorage::in_memory(IndexMap::new()).with_env(env.clone()));
191 Self::new(
192 auth,
193 ModelRegistryOptions {
194 env,
195 ..ModelRegistryOptions::default()
196 },
197 )
198 }
199
200 fn read(&self) -> RwLockReadGuard<'_, State> {
201 self.inner
202 .state
203 .read()
204 .unwrap_or_else(PoisonError::into_inner)
205 }
206
207 fn write(&self) -> RwLockWriteGuard<'_, State> {
208 self.inner
209 .state
210 .write()
211 .unwrap_or_else(PoisonError::into_inner)
212 }
213
214 #[must_use]
216 pub fn auth_storage(&self) -> &Arc<AuthStorage> {
217 &self.inner.auth
218 }
219
220 pub fn set_oauth_refresher(&self, refresher: Arc<dyn OAuthRefresher>) {
222 *self
223 .inner
224 .refresher
225 .write()
226 .unwrap_or_else(PoisonError::into_inner) = Some(refresher);
227 }
228
229 fn refresher(&self) -> Option<Arc<dyn OAuthRefresher>> {
230 self.inner
231 .refresher
232 .read()
233 .unwrap_or_else(PoisonError::into_inner)
234 .clone()
235 }
236
237 pub fn refresh(&self) {
239 let models_json = ModelsJson::load(self.inner.models_json_path.as_deref());
240 let mut state = self.write();
241 state.models_json = models_json;
242 state.recompose();
243 }
244
245 pub fn set_models_json(&self, models_json: ModelsJson) {
247 let mut state = self.write();
248 state.models_json = models_json;
249 state.recompose();
250 }
251
252 #[must_use]
255 pub fn error(&self) -> Option<String> {
256 let state = self.read();
257 let mut errors: Vec<String> = Vec::new();
258 if let Some(error) = state.models_json.error() {
259 errors.push(error.to_owned());
260 }
261 for (provider, error) in &state.composition_errors {
262 errors.push(format!("Provider \"{provider}\": {error}"));
263 }
264 (!errors.is_empty()).then(|| errors.join("\n\n"))
265 }
266
267 #[must_use]
270 pub fn diagnostics(&self) -> Vec<RegistryDiagnostic> {
271 let state = self.read();
272 let mut diagnostics: Vec<RegistryDiagnostic> = Vec::new();
273 for provider in &state.providers {
274 for model in &provider.models {
275 if is_virtual_api(&model.api) || state.route_target(model).is_some() {
276 continue;
277 }
278 let message = format!("No provider extension implements api \"{}\"", model.api);
279 if !diagnostics
280 .iter()
281 .any(|d| d.provider == provider.id && d.message == message)
282 {
283 diagnostics.push(RegistryDiagnostic {
284 provider: provider.id.clone(),
285 message,
286 });
287 }
288 }
289 }
290 diagnostics
291 }
292
293 pub fn register_provider(
302 &self,
303 extension_id: &str,
304 registration: ProviderRegistration,
305 ) -> Result<(), RegistryError> {
306 if registration.id.trim().is_empty() {
307 return Err(RegistryError::EmptyProviderId);
308 }
309 let dynamic_models = match (®istration.models, &self.inner.cache) {
310 (ProviderModels::Dynamic, Some(cache)) => match cache.read(®istration.id) {
311 Ok(entry) => entry.map(|e| e.models).unwrap_or_default(),
312 Err(error) => {
313 tracing::warn!(%error, provider = %registration.id, "ignoring model cache");
314 Vec::new()
315 }
316 },
317 _ => Vec::new(),
318 };
319 let mut state = self.write();
320 state.registrations.insert(
321 registration.id.clone(),
322 RegisteredProvider {
323 registration,
324 extension_id: extension_id.to_owned(),
325 dynamic_models,
326 },
327 );
328 state.recompose();
329 Ok(())
330 }
331
332 pub fn register_api_implementor(&self, api: &str, extension_id: &str) {
335 self.write()
336 .api_implementors
337 .insert(api.to_owned(), extension_id.to_owned());
338 }
339
340 pub fn set_provider_models(
346 &self,
347 provider: &str,
348 models: &[Model],
349 ) -> Result<(), RegistryError> {
350 {
351 let mut state = self.write();
352 let registered = state
353 .registrations
354 .get_mut(provider)
355 .ok_or_else(|| RegistryError::UnknownProvider(provider.to_owned()))?;
356 if !matches!(registered.registration.models, ProviderModels::Dynamic) {
357 return Err(RegistryError::StaticModels(provider.to_owned()));
358 }
359 registered.dynamic_models = models.to_vec();
360 state.recompose();
361 }
362 if let Some(cache) = &self.inner.cache {
363 cache.write(provider, models)?;
364 }
365 Ok(())
366 }
367
368 pub fn register_provider_config(
375 &self,
376 provider: &str,
377 config: ProviderConfigOverride,
378 ) -> Result<(), RegistryError> {
379 if provider.trim().is_empty() {
380 return Err(RegistryError::EmptyProviderId);
381 }
382 let mut state = self.write();
383 let base = state
384 .registrations
385 .get(provider)
386 .map(RegisteredProvider::base_models)
387 .unwrap_or_default();
388 crate::compose::apply_override(
389 provider,
390 crate::compose::apply_models_json(
391 provider,
392 &base,
393 state.models_json.provider(provider),
394 )
395 .map_err(RegistryError::InvalidConfig)?,
396 Some(&config),
397 )
398 .map_err(RegistryError::InvalidConfig)?;
399 let merged = match state.overrides.shift_remove(provider) {
400 Some(previous) => previous.merged_with(config),
401 None => config,
402 };
403 state.overrides.insert(provider.to_owned(), merged);
404 state.recompose();
405 Ok(())
406 }
407
408 pub fn unregister_provider(&self, provider: &str) {
411 let mut state = self.write();
412 state.registrations.shift_remove(provider);
413 state.overrides.shift_remove(provider);
414 state.recompose();
415 }
416
417 pub fn unregister_extension(&self, extension_id: &str) {
420 let mut state = self.write();
421 state
422 .registrations
423 .retain(|_, registered| registered.extension_id != extension_id);
424 state.api_implementors.retain(|_, ext| ext != extension_id);
425 state.virtual_models.remove_extension(extension_id);
426 state.recompose();
427 }
428
429 #[must_use]
431 pub fn registered_providers(&self) -> Vec<String> {
432 let state = self.read();
433 let mut ids: Vec<String> = state.overrides.keys().cloned().collect();
434 for id in state.registrations.keys() {
435 if !ids.contains(id) {
436 ids.push(id.clone());
437 }
438 }
439 ids
440 }
441
442 #[must_use]
444 pub fn registration(&self, provider: &str) -> Option<ProviderRegistration> {
445 self.read()
446 .registrations
447 .get(provider)
448 .map(|r| r.registration.clone())
449 }
450
451 #[must_use]
453 pub fn registered_provider_config(&self, provider: &str) -> Option<ProviderConfigOverride> {
454 self.read().overrides.get(provider).cloned()
455 }
456
457 #[must_use]
461 pub fn all(&self) -> Vec<Model> {
462 self.read()
463 .providers
464 .iter()
465 .flat_map(|p| p.models.iter().cloned())
466 .collect()
467 }
468
469 #[must_use]
472 pub fn available(&self) -> Vec<Model> {
473 let state = self.read();
474 let mut available = Vec::new();
475 for provider in &state.providers {
476 if self.check_plan(&provider.id, &provider.auth).is_none() {
477 continue;
478 }
479 available.extend(
480 provider
481 .models
482 .iter()
483 .filter(|m| is_virtual_api(&m.api) || state.route_target(m).is_some())
484 .cloned(),
485 );
486 }
487 available
488 }
489
490 #[must_use]
492 pub fn find(&self, provider: &str, id: &str) -> Option<Model> {
493 self.read().find(provider, id).cloned()
494 }
495
496 #[must_use]
498 pub fn find_physical(&self, provider: &str, id: &str) -> Option<Model> {
499 self.find(provider, id).filter(|m| !is_virtual_api(&m.api))
500 }
501
502 #[must_use]
504 pub fn provider_ids(&self) -> Vec<String> {
505 self.read().providers.iter().map(|p| p.id.clone()).collect()
506 }
507
508 #[must_use]
510 pub fn provider_display_name(&self, provider: &str) -> String {
511 self.read()
512 .provider(provider)
513 .map_or_else(|| provider.to_owned(), |p| p.name.clone())
514 }
515
516 #[must_use]
518 pub fn route_target(&self, model: &Model) -> Option<String> {
519 self.read().route_target(model).map(str::to_owned)
520 }
521
522 fn check_plan(&self, provider: &str, plan: &AuthPlan) -> Option<AuthCheck> {
525 let credential = match self.inner.auth.get(provider) {
526 Ok(credential) => credential,
527 Err(error) => {
528 tracing::warn!(%error, provider, "auth check failed");
529 return None;
530 }
531 };
532 auth::check(plan, credential.as_ref(), self.inner.env.as_ref())
533 }
534
535 fn plan(&self, provider: &str) -> Option<AuthPlan> {
536 self.read().provider(provider).map(|p| p.auth.clone())
537 }
538
539 #[must_use]
541 pub fn check_auth(&self, provider: &str) -> Option<AuthCheck> {
542 let plan = self.plan(provider)?;
543 self.check_plan(provider, &plan)
544 }
545
546 #[must_use]
548 pub fn has_configured_auth(&self, provider: &str) -> bool {
549 self.check_auth(provider).is_some()
550 }
551
552 #[must_use]
554 pub fn is_using_oauth(&self, provider: &str) -> bool {
555 self.check_auth(provider)
556 .is_some_and(|c| c.kind == CredentialType::Oauth)
557 }
558
559 #[must_use]
561 pub fn auth_status(&self, provider: &str) -> AuthStatus {
562 let plan = self
563 .plan(provider)
564 .unwrap_or(AuthPlan::Standard(auth::StandardAuth {
565 api_key: None,
566 oauth: false,
567 raw_headers: None,
568 auth_header: false,
569 }));
570 let check = self.check_plan(provider, &plan);
571 auth::status(
572 &plan,
573 &self.inner.auth,
574 provider,
575 check.as_ref(),
576 self.inner.env.as_ref(),
577 )
578 }
579
580 pub fn set_runtime_api_key(&self, provider: &str, key: impl Into<String>) {
582 self.inner.auth.set_runtime_api_key(provider, key);
583 }
584
585 pub fn remove_runtime_api_key(&self, provider: &str) {
587 self.inner.auth.remove_runtime_api_key(provider);
588 }
589
590 pub async fn get_auth(
596 &self,
597 provider: &str,
598 overrides: &AuthOverrides,
599 ) -> Result<Option<ResolvedAuth>, AuthResolveError> {
600 let Some(plan) = self.plan(provider) else {
601 return Ok(None);
602 };
603 let refresher = self.refresher();
604 Resolver {
605 provider,
606 plan: &plan,
607 storage: &self.inner.auth,
608 env: self.inner.env.as_ref(),
609 refresher: refresher.as_deref(),
610 }
611 .resolve(overrides)
612 .await
613 }
614
615 pub async fn get_api_key(&self, provider: &str) -> Option<String> {
617 self.get_auth(provider, &AuthOverrides::default())
618 .await
619 .ok()
620 .flatten()
621 .and_then(|auth| auth.api_key)
622 }
623
624 pub async fn get_api_key_and_headers(
631 &self,
632 model: &Model,
633 overrides: &AuthOverrides,
634 ) -> Result<RequestAuth, RequestAuthError> {
635 let provider = model.provider.as_str();
636 let to_error = |error: &AuthResolveError| RequestAuthError {
637 message: error.request_message(provider),
638 };
639 let resolution = self
640 .get_auth(provider, overrides)
641 .await
642 .map_err(|e| to_error(&e))?;
643 let (config_headers, model_headers, auth_header) = {
644 let state = self.read();
645 let config = state.models_json.provider(provider);
646 let override_config = state.overrides.get(provider);
647 (
648 configured_headers(config, override_config),
649 raw_model_headers(&model.id, config, override_config),
650 override_config
651 .and_then(|o| o.auth_header)
652 .or_else(|| config.and_then(|c| c.auth_header))
653 .unwrap_or(false),
654 )
655 };
656 let description = format!("model \"{provider}/{}\"", model.id);
657 let env = self.inner.env.as_ref();
658 if let Some(resolution) = resolution {
659 let headers = resolve_headers_strict(
660 model_headers.as_ref(),
661 &description,
662 resolution.env.as_ref(),
663 env,
664 )
665 .map_err(|e| to_error(&AuthResolveError::ModelHeaders(e)))?;
666 return Ok(RequestAuth {
667 api_key: resolution.api_key,
668 headers: merge_headers(resolution.headers, headers),
669 });
670 }
671 let raw = match (config_headers, model_headers) {
673 (None, None) => None,
674 (base, over) => {
675 let mut merged = base.unwrap_or_default();
676 merged.extend(over.unwrap_or_default());
677 Some(merged)
678 }
679 };
680 let configured = resolve_headers_strict(raw.as_ref(), &description, None, env)
681 .map_err(|e| to_error(&AuthResolveError::ModelHeaders(e)))?;
682 if auth_header {
683 return Err(RequestAuthError {
684 message: format!("No API key found for \"{provider}\""),
685 });
686 }
687 let headers = match (model.headers.clone(), configured) {
688 (None, None) => None,
689 (base, over) => {
690 let mut merged = base.unwrap_or_default();
691 merged.extend(over.unwrap_or_default());
692 Some(merged)
693 }
694 };
695 Ok(RequestAuth {
696 api_key: None,
697 headers,
698 })
699 }
700
701 pub fn register_virtual_model(
709 &self,
710 extension_id: &str,
711 definition: VirtualModelDefinition,
712 ) -> Result<(), RegistryError> {
713 let mut state = self.write();
714 let is_physical = state
715 .find(&definition.provider, &definition.id)
716 .is_some_and(|m| !is_virtual_api(&m.api));
717 state
718 .virtual_models
719 .register(extension_id, definition, is_physical)?;
720 state.recompose();
721 Ok(())
722 }
723
724 pub fn unregister_virtual_model(&self, provider: &str, id: &str) {
726 let mut state = self.write();
727 if state.virtual_models.unregister(provider, id) {
728 state.recompose();
729 }
730 }
731
732 #[must_use]
734 pub fn virtual_model(&self, provider: &str, id: &str) -> Option<RegisteredVirtualModel> {
735 self.read().virtual_models.get(provider, id).cloned()
736 }
737
738 pub fn route_request(
745 &self,
746 model: &Model,
747 messages: Vec<Message>,
748 reason: RouteReason,
749 thinking_level: ThinkingLevel,
750 failed: Option<AssistantMessage>,
751 state: Option<JsonValue>,
752 ) -> Result<RouteRequest, RouteError> {
753 if self.virtual_model(&model.provider, &model.id).is_none() {
754 return Err(RouteError::NotRegistered {
755 provider: model.provider.clone(),
756 id: model.id.clone(),
757 });
758 }
759 let previous = find_latest_response(&messages).and_then(|latest| {
760 self.find_physical(&latest.provider, &latest.model)
761 .map(|model| RoutePrevious {
762 model,
763 thinking_level: latest.thinking_level,
764 })
765 });
766 let failed = failed.and_then(|message| {
767 self.find_physical(&message.provider, &message.model)
768 .map(|model| RouteFailed {
769 model,
770 thinking_level: message.thinking_level,
771 message,
772 })
773 });
774 Ok(RouteRequest {
775 model: model.clone(),
776 thinking_level,
777 reason,
778 previous,
779 failed,
780 state: if reason == RouteReason::Direct {
781 None
782 } else {
783 state
784 },
785 messages,
786 })
787 }
788
789 pub fn validate_route(
796 &self,
797 model: &Model,
798 result: RouteResult,
799 ) -> Result<ModelRoute, RouteError> {
800 let target = format!("{}/{}", result.model.provider, result.model.id);
801 let physical = self
802 .find_physical(&result.model.provider, &result.model.id)
803 .ok_or_else(|| RouteError::NotPhysical {
804 provider: model.provider.clone(),
805 id: model.id.clone(),
806 target: target.clone(),
807 })?;
808 if !self.has_configured_auth(&physical.provider) {
809 return Err(RouteError::NoCredentials {
810 provider: model.provider.clone(),
811 id: model.id.clone(),
812 target,
813 });
814 }
815 let thinking_level = result
816 .thinking_level
817 .clamp(&physical.available_thinking_levels());
818 Ok(ModelRoute {
819 model: physical,
820 thinking_level,
821 state: result.state,
822 })
823 }
824}
825
826#[cfg(test)]
827mod tests;