mod auth;
mod state;
use std::path::{Path, PathBuf};
use std::sync::{Arc, PoisonError, RwLock, RwLockReadGuard, RwLockWriteGuard};
use indexmap::IndexMap;
use oino_ai::{AssistantMessage, JsonValue, Message, Model, ThinkingLevel};
pub use self::auth::{
ApiKeyError, AuthCheck, AuthOverrides, AuthResolveError, AuthStatus, AuthStatusSource,
ResolvedAuth,
};
use self::auth::{AuthPlan, Resolver};
use self::state::{RegisteredProvider, State};
use crate::auth::credential::CredentialType;
use crate::auth::oauth::OAuthRefresher;
use crate::auth::storage::AuthStorage;
use crate::cache::{CacheError, ModelCache};
use crate::compose::{configured_headers, merge_headers, raw_model_headers};
use crate::config_value::resolve_headers_strict;
use crate::env::{EnvSource, ProcessEnv};
use crate::models_json::ModelsJson;
use crate::registration::{ProviderConfigOverride, ProviderModels, ProviderRegistration};
use crate::virtual_models::{
ModelRoute, RegisteredVirtualModel, RouteError, RouteFailed, RoutePrevious, RouteReason,
RouteRequest, RouteResult, VirtualModelDefinition, VirtualModelError, find_latest_response,
is_virtual_api,
};
#[derive(Clone)]
pub struct ModelRegistryOptions {
pub models_json_path: Option<PathBuf>,
pub model_cache: Option<ModelCache>,
pub env: Arc<dyn EnvSource>,
}
impl Default for ModelRegistryOptions {
fn default() -> Self {
Self {
models_json_path: None,
model_cache: None,
env: Arc::new(ProcessEnv),
}
}
}
impl std::fmt::Debug for ModelRegistryOptions {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ModelRegistryOptions")
.field("models_json_path", &self.models_json_path)
.field("model_cache", &self.model_cache)
.finish_non_exhaustive()
}
}
impl ModelRegistryOptions {
#[must_use]
pub fn for_agent_dir(agent_dir: &Path) -> Self {
Self {
models_json_path: Some(agent_dir.join("models.json")),
model_cache: Some(ModelCache::for_agent_dir(agent_dir)),
env: Arc::new(ProcessEnv),
}
}
}
#[derive(Debug, thiserror::Error)]
pub enum RegistryError {
#[error("Provider id must not be empty.")]
EmptyProviderId,
#[error("Provider {0} is not registered.")]
UnknownProvider(String),
#[error("Provider {0} has a static model list.")]
StaticModels(String),
#[error("{0}")]
InvalidConfig(String),
#[error(transparent)]
VirtualModel(#[from] VirtualModelError),
#[error(transparent)]
Cache(#[from] CacheError),
}
#[derive(Clone, Default, PartialEq, Eq)]
pub struct RequestAuth {
pub api_key: Option<String>,
pub headers: Option<IndexMap<String, String>>,
}
impl std::fmt::Debug for RequestAuth {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("RequestAuth")
.field("api_key", &self.api_key.as_ref().map(|_| "<redacted>"))
.field(
"headers",
&self.headers.as_ref().map(|h| h.keys().collect::<Vec<_>>()),
)
.finish()
}
}
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
#[error("{message}")]
pub struct RequestAuthError {
pub message: String,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RegistryDiagnostic {
pub provider: String,
pub message: String,
}
struct Inner {
state: RwLock<State>,
auth: Arc<AuthStorage>,
env: Arc<dyn EnvSource>,
models_json_path: Option<PathBuf>,
cache: Option<ModelCache>,
refresher: RwLock<Option<Arc<dyn OAuthRefresher>>>,
}
#[derive(Clone)]
pub struct ModelRegistry {
inner: Arc<Inner>,
}
impl std::fmt::Debug for ModelRegistry {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ModelRegistry")
.field("models_json_path", &self.inner.models_json_path)
.finish_non_exhaustive()
}
}
impl ModelRegistry {
#[must_use]
pub fn new(auth: Arc<AuthStorage>, options: ModelRegistryOptions) -> Self {
let mut state = State {
models_json: ModelsJson::load(options.models_json_path.as_deref()),
..State::default()
};
state.recompose();
Self {
inner: Arc::new(Inner {
state: RwLock::new(state),
auth,
env: options.env,
models_json_path: options.models_json_path,
cache: options.model_cache,
refresher: RwLock::new(None),
}),
}
}
#[must_use]
pub fn in_memory(env: Arc<dyn EnvSource>) -> Self {
let auth = Arc::new(AuthStorage::in_memory(IndexMap::new()).with_env(env.clone()));
Self::new(
auth,
ModelRegistryOptions {
env,
..ModelRegistryOptions::default()
},
)
}
fn read(&self) -> RwLockReadGuard<'_, State> {
self.inner
.state
.read()
.unwrap_or_else(PoisonError::into_inner)
}
fn write(&self) -> RwLockWriteGuard<'_, State> {
self.inner
.state
.write()
.unwrap_or_else(PoisonError::into_inner)
}
#[must_use]
pub fn auth_storage(&self) -> &Arc<AuthStorage> {
&self.inner.auth
}
pub fn set_oauth_refresher(&self, refresher: Arc<dyn OAuthRefresher>) {
*self
.inner
.refresher
.write()
.unwrap_or_else(PoisonError::into_inner) = Some(refresher);
}
fn refresher(&self) -> Option<Arc<dyn OAuthRefresher>> {
self.inner
.refresher
.read()
.unwrap_or_else(PoisonError::into_inner)
.clone()
}
pub fn refresh(&self) {
let models_json = ModelsJson::load(self.inner.models_json_path.as_deref());
let mut state = self.write();
state.models_json = models_json;
state.recompose();
}
pub fn set_models_json(&self, models_json: ModelsJson) {
let mut state = self.write();
state.models_json = models_json;
state.recompose();
}
#[must_use]
pub fn error(&self) -> Option<String> {
let state = self.read();
let mut errors: Vec<String> = Vec::new();
if let Some(error) = state.models_json.error() {
errors.push(error.to_owned());
}
for (provider, error) in &state.composition_errors {
errors.push(format!("Provider \"{provider}\": {error}"));
}
(!errors.is_empty()).then(|| errors.join("\n\n"))
}
#[must_use]
pub fn diagnostics(&self) -> Vec<RegistryDiagnostic> {
let state = self.read();
let mut diagnostics: Vec<RegistryDiagnostic> = Vec::new();
for provider in &state.providers {
for model in &provider.models {
if is_virtual_api(&model.api) || state.route_target(model).is_some() {
continue;
}
let message = format!("No provider extension implements api \"{}\"", model.api);
if !diagnostics
.iter()
.any(|d| d.provider == provider.id && d.message == message)
{
diagnostics.push(RegistryDiagnostic {
provider: provider.id.clone(),
message,
});
}
}
}
diagnostics
}
pub fn register_provider(
&self,
extension_id: &str,
registration: ProviderRegistration,
) -> Result<(), RegistryError> {
if registration.id.trim().is_empty() {
return Err(RegistryError::EmptyProviderId);
}
let dynamic_models = match (®istration.models, &self.inner.cache) {
(ProviderModels::Dynamic, Some(cache)) => match cache.read(®istration.id) {
Ok(entry) => entry.map(|e| e.models).unwrap_or_default(),
Err(error) => {
tracing::warn!(%error, provider = %registration.id, "ignoring model cache");
Vec::new()
}
},
_ => Vec::new(),
};
let mut state = self.write();
state.registrations.insert(
registration.id.clone(),
RegisteredProvider {
registration,
extension_id: extension_id.to_owned(),
dynamic_models,
},
);
state.recompose();
Ok(())
}
pub fn register_api_implementor(&self, api: &str, extension_id: &str) {
self.write()
.api_implementors
.insert(api.to_owned(), extension_id.to_owned());
}
pub fn set_provider_models(
&self,
provider: &str,
models: &[Model],
) -> Result<(), RegistryError> {
{
let mut state = self.write();
let registered = state
.registrations
.get_mut(provider)
.ok_or_else(|| RegistryError::UnknownProvider(provider.to_owned()))?;
if !matches!(registered.registration.models, ProviderModels::Dynamic) {
return Err(RegistryError::StaticModels(provider.to_owned()));
}
registered.dynamic_models = models.to_vec();
state.recompose();
}
if let Some(cache) = &self.inner.cache {
cache.write(provider, models)?;
}
Ok(())
}
pub fn register_provider_config(
&self,
provider: &str,
config: ProviderConfigOverride,
) -> Result<(), RegistryError> {
if provider.trim().is_empty() {
return Err(RegistryError::EmptyProviderId);
}
let mut state = self.write();
let base = state
.registrations
.get(provider)
.map(RegisteredProvider::base_models)
.unwrap_or_default();
crate::compose::apply_override(
provider,
crate::compose::apply_models_json(
provider,
&base,
state.models_json.provider(provider),
)
.map_err(RegistryError::InvalidConfig)?,
Some(&config),
)
.map_err(RegistryError::InvalidConfig)?;
let merged = match state.overrides.shift_remove(provider) {
Some(previous) => previous.merged_with(config),
None => config,
};
state.overrides.insert(provider.to_owned(), merged);
state.recompose();
Ok(())
}
pub fn unregister_provider(&self, provider: &str) {
let mut state = self.write();
state.registrations.shift_remove(provider);
state.overrides.shift_remove(provider);
state.recompose();
}
pub fn unregister_extension(&self, extension_id: &str) {
let mut state = self.write();
state
.registrations
.retain(|_, registered| registered.extension_id != extension_id);
state.api_implementors.retain(|_, ext| ext != extension_id);
state.virtual_models.remove_extension(extension_id);
state.recompose();
}
#[must_use]
pub fn registered_providers(&self) -> Vec<String> {
let state = self.read();
let mut ids: Vec<String> = state.overrides.keys().cloned().collect();
for id in state.registrations.keys() {
if !ids.contains(id) {
ids.push(id.clone());
}
}
ids
}
#[must_use]
pub fn registration(&self, provider: &str) -> Option<ProviderRegistration> {
self.read()
.registrations
.get(provider)
.map(|r| r.registration.clone())
}
#[must_use]
pub fn registered_provider_config(&self, provider: &str) -> Option<ProviderConfigOverride> {
self.read().overrides.get(provider).cloned()
}
#[must_use]
pub fn all(&self) -> Vec<Model> {
self.read()
.providers
.iter()
.flat_map(|p| p.models.iter().cloned())
.collect()
}
#[must_use]
pub fn available(&self) -> Vec<Model> {
let state = self.read();
let mut available = Vec::new();
for provider in &state.providers {
if self.check_plan(&provider.id, &provider.auth).is_none() {
continue;
}
available.extend(
provider
.models
.iter()
.filter(|m| is_virtual_api(&m.api) || state.route_target(m).is_some())
.cloned(),
);
}
available
}
#[must_use]
pub fn find(&self, provider: &str, id: &str) -> Option<Model> {
self.read().find(provider, id).cloned()
}
#[must_use]
pub fn find_physical(&self, provider: &str, id: &str) -> Option<Model> {
self.find(provider, id).filter(|m| !is_virtual_api(&m.api))
}
#[must_use]
pub fn provider_ids(&self) -> Vec<String> {
self.read().providers.iter().map(|p| p.id.clone()).collect()
}
#[must_use]
pub fn provider_display_name(&self, provider: &str) -> String {
self.read()
.provider(provider)
.map_or_else(|| provider.to_owned(), |p| p.name.clone())
}
#[must_use]
pub fn route_target(&self, model: &Model) -> Option<String> {
self.read().route_target(model).map(str::to_owned)
}
fn check_plan(&self, provider: &str, plan: &AuthPlan) -> Option<AuthCheck> {
let credential = match self.inner.auth.get(provider) {
Ok(credential) => credential,
Err(error) => {
tracing::warn!(%error, provider, "auth check failed");
return None;
}
};
auth::check(plan, credential.as_ref(), self.inner.env.as_ref())
}
fn plan(&self, provider: &str) -> Option<AuthPlan> {
self.read().provider(provider).map(|p| p.auth.clone())
}
#[must_use]
pub fn check_auth(&self, provider: &str) -> Option<AuthCheck> {
let plan = self.plan(provider)?;
self.check_plan(provider, &plan)
}
#[must_use]
pub fn has_configured_auth(&self, provider: &str) -> bool {
self.check_auth(provider).is_some()
}
#[must_use]
pub fn is_using_oauth(&self, provider: &str) -> bool {
self.check_auth(provider)
.is_some_and(|c| c.kind == CredentialType::Oauth)
}
#[must_use]
pub fn auth_status(&self, provider: &str) -> AuthStatus {
let plan = self
.plan(provider)
.unwrap_or(AuthPlan::Standard(auth::StandardAuth {
api_key: None,
oauth: false,
raw_headers: None,
auth_header: false,
}));
let check = self.check_plan(provider, &plan);
auth::status(
&plan,
&self.inner.auth,
provider,
check.as_ref(),
self.inner.env.as_ref(),
)
}
pub fn set_runtime_api_key(&self, provider: &str, key: impl Into<String>) {
self.inner.auth.set_runtime_api_key(provider, key);
}
pub fn remove_runtime_api_key(&self, provider: &str) {
self.inner.auth.remove_runtime_api_key(provider);
}
pub async fn get_auth(
&self,
provider: &str,
overrides: &AuthOverrides,
) -> Result<Option<ResolvedAuth>, AuthResolveError> {
let Some(plan) = self.plan(provider) else {
return Ok(None);
};
let refresher = self.refresher();
Resolver {
provider,
plan: &plan,
storage: &self.inner.auth,
env: self.inner.env.as_ref(),
refresher: refresher.as_deref(),
}
.resolve(overrides)
.await
}
pub async fn get_api_key(&self, provider: &str) -> Option<String> {
self.get_auth(provider, &AuthOverrides::default())
.await
.ok()
.flatten()
.and_then(|auth| auth.api_key)
}
pub async fn get_api_key_and_headers(
&self,
model: &Model,
overrides: &AuthOverrides,
) -> Result<RequestAuth, RequestAuthError> {
let provider = model.provider.as_str();
let to_error = |error: &AuthResolveError| RequestAuthError {
message: error.request_message(provider),
};
let resolution = self
.get_auth(provider, overrides)
.await
.map_err(|e| to_error(&e))?;
let (config_headers, model_headers, auth_header) = {
let state = self.read();
let config = state.models_json.provider(provider);
let override_config = state.overrides.get(provider);
(
configured_headers(config, override_config),
raw_model_headers(&model.id, config, override_config),
override_config
.and_then(|o| o.auth_header)
.or_else(|| config.and_then(|c| c.auth_header))
.unwrap_or(false),
)
};
let description = format!("model \"{provider}/{}\"", model.id);
let env = self.inner.env.as_ref();
if let Some(resolution) = resolution {
let headers = resolve_headers_strict(
model_headers.as_ref(),
&description,
resolution.env.as_ref(),
env,
)
.map_err(|e| to_error(&AuthResolveError::ModelHeaders(e)))?;
return Ok(RequestAuth {
api_key: resolution.api_key,
headers: merge_headers(resolution.headers, headers),
});
}
let raw = match (config_headers, model_headers) {
(None, None) => None,
(base, over) => {
let mut merged = base.unwrap_or_default();
merged.extend(over.unwrap_or_default());
Some(merged)
}
};
let configured = resolve_headers_strict(raw.as_ref(), &description, None, env)
.map_err(|e| to_error(&AuthResolveError::ModelHeaders(e)))?;
if auth_header {
return Err(RequestAuthError {
message: format!("No API key found for \"{provider}\""),
});
}
let headers = match (model.headers.clone(), configured) {
(None, None) => None,
(base, over) => {
let mut merged = base.unwrap_or_default();
merged.extend(over.unwrap_or_default());
Some(merged)
}
};
Ok(RequestAuth {
api_key: None,
headers,
})
}
pub fn register_virtual_model(
&self,
extension_id: &str,
definition: VirtualModelDefinition,
) -> Result<(), RegistryError> {
let mut state = self.write();
let is_physical = state
.find(&definition.provider, &definition.id)
.is_some_and(|m| !is_virtual_api(&m.api));
state
.virtual_models
.register(extension_id, definition, is_physical)?;
state.recompose();
Ok(())
}
pub fn unregister_virtual_model(&self, provider: &str, id: &str) {
let mut state = self.write();
if state.virtual_models.unregister(provider, id) {
state.recompose();
}
}
#[must_use]
pub fn virtual_model(&self, provider: &str, id: &str) -> Option<RegisteredVirtualModel> {
self.read().virtual_models.get(provider, id).cloned()
}
pub fn route_request(
&self,
model: &Model,
messages: Vec<Message>,
reason: RouteReason,
thinking_level: ThinkingLevel,
failed: Option<AssistantMessage>,
state: Option<JsonValue>,
) -> Result<RouteRequest, RouteError> {
if self.virtual_model(&model.provider, &model.id).is_none() {
return Err(RouteError::NotRegistered {
provider: model.provider.clone(),
id: model.id.clone(),
});
}
let previous = find_latest_response(&messages).and_then(|latest| {
self.find_physical(&latest.provider, &latest.model)
.map(|model| RoutePrevious {
model,
thinking_level: latest.thinking_level,
})
});
let failed = failed.and_then(|message| {
self.find_physical(&message.provider, &message.model)
.map(|model| RouteFailed {
model,
thinking_level: message.thinking_level,
message,
})
});
Ok(RouteRequest {
model: model.clone(),
thinking_level,
reason,
previous,
failed,
state: if reason == RouteReason::Direct {
None
} else {
state
},
messages,
})
}
pub fn validate_route(
&self,
model: &Model,
result: RouteResult,
) -> Result<ModelRoute, RouteError> {
let target = format!("{}/{}", result.model.provider, result.model.id);
let physical = self
.find_physical(&result.model.provider, &result.model.id)
.ok_or_else(|| RouteError::NotPhysical {
provider: model.provider.clone(),
id: model.id.clone(),
target: target.clone(),
})?;
if !self.has_configured_auth(&physical.provider) {
return Err(RouteError::NoCredentials {
provider: model.provider.clone(),
id: model.id.clone(),
target,
});
}
let thinking_level = result
.thinking_level
.clamp(&physical.available_thinking_levels());
Ok(ModelRoute {
model: physical,
thinking_level,
state: result.state,
})
}
}
#[cfg(test)]
mod tests;