use std::collections::{BTreeSet, HashMap, HashSet};
use std::path::{Path, PathBuf};
use std::sync::RwLock;
use std::time::Duration;
use serde_json::Value;
use crate::subscription::{SubscriptionProvider, SubscriptionReader, SubscriptionToken};
#[path = "model_catalog_parse.rs"]
mod parse;
#[cfg(test)]
use parse::parse_catalog;
use parse::{next_catalog_cursor, parse_catalog_records, provider_protocols};
#[path = "model_catalog_errors.rs"]
mod errors;
#[cfg(test)]
use errors::resource_error_code;
pub use errors::{is_credential_rejection, is_permission_refusal};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CatalogAcceptance {
Accepted,
MissingSubscription,
CredentialRejected,
Unverified,
}
#[must_use]
pub fn classify_catalog_acceptance(result: &Result<Vec<String>, String>) -> CatalogAcceptance {
match result {
Ok(models) if !models.is_empty() => CatalogAcceptance::Accepted,
Ok(_) => CatalogAcceptance::MissingSubscription,
Err(error) if is_credential_rejection(error) => CatalogAcceptance::CredentialRejected,
Err(_) => CatalogAcceptance::Unverified,
}
}
pub const CATALOG_TTL: Duration = Duration::from_secs(5 * 60);
const FETCH_TIMEOUT: Duration = Duration::from_secs(15);
const MAX_CATALOG_PAGES: usize = 100;
#[derive(Debug, Clone, PartialEq, Eq, serde::Deserialize, serde::Serialize)]
pub struct CatalogRecord {
pub provider: SubscriptionProvider,
pub account: String,
pub canonical_id: String,
pub raw: serde_json::Map<String, Value>,
pub source_order: u64,
pub fetched_at: i64,
pub health_generation: String,
pub protocols: BTreeSet<crate::client_policy::ClientProtocol>,
}
impl CatalogRecord {
fn synthetic(
provider: SubscriptionProvider,
account: &str,
canonical_id: String,
source_order: usize,
fetched_at: i64,
) -> Self {
let mut raw = serde_json::Map::new();
raw.insert("id".into(), Value::String(canonical_id.clone()));
Self {
provider,
account: account.to_string(),
canonical_id,
raw,
source_order: source_order as u64,
fetched_at,
health_generation: format!("{provider}:{account}:{fetched_at}"),
protocols: provider_protocols(provider),
}
}
}
#[derive(Debug, Clone, Default, PartialEq, Eq, serde::Deserialize, serde::Serialize)]
pub struct CatalogStatus {
pub models: Vec<String>,
pub records: Vec<CatalogRecord>,
pub account: Option<String>,
pub refreshed_at: Option<i64>,
pub last_error: Option<String>,
pub discovered: bool,
pub credential_healthy: bool,
}
impl CatalogStatus {
#[must_use]
pub fn routable_models(&self) -> &[String] {
if self.discovered && self.credential_healthy {
&self.models
} else {
&[]
}
}
#[must_use]
pub fn routable_records(&self) -> &[CatalogRecord] {
if self.discovered && self.credential_healthy {
&self.records
} else {
&[]
}
}
#[must_use]
pub const fn is_degraded(&self) -> bool {
!self.discovered || !self.credential_healthy
}
}
pub struct ModelCatalogCache {
entries: RwLock<HashMap<(SubscriptionProvider, String), CatalogStatus>>,
persistence: Option<CatalogPersistence>,
}
#[derive(Debug, Clone)]
struct CatalogPersistence {
path: PathBuf,
invalidations: PathBuf,
}
#[derive(Debug, serde::Deserialize, serde::Serialize)]
struct PersistedCatalogs {
version: u8,
entries: Vec<PersistedCatalogEntry>,
}
#[derive(Debug, serde::Deserialize, serde::Serialize)]
struct PersistedCatalogEntry {
provider: SubscriptionProvider,
router_account: String,
status: CatalogStatus,
}
const PERSISTED_CATALOG_VERSION: u8 = 1;
const PERSISTED_CATALOG_FILE: &str = "model-catalogs.json";
const CATALOG_INVALIDATION_DIR: &str = "model-catalog-invalidations";
impl Default for ModelCatalogCache {
fn default() -> Self {
Self::new()
}
}
impl ModelCatalogCache {
#[must_use]
pub fn new() -> Self {
Self {
entries: RwLock::new(HashMap::new()),
persistence: None,
}
}
#[must_use]
pub fn persistent(data_dir: &Path) -> Self {
let persistence = CatalogPersistence {
path: data_dir.join(PERSISTED_CATALOG_FILE),
invalidations: data_dir.join(CATALOG_INVALIDATION_DIR),
};
let mut entries = load_persisted_catalogs(&persistence.path).unwrap_or_else(|error| {
tracing::warn!("could not load the persisted model catalog: {error}");
HashMap::new()
});
for status in entries.values_mut() {
if status.discovered {
status.credential_healthy = false;
status.last_error =
Some("awaiting authenticated catalog refresh after router restart".to_string());
}
}
Self {
entries: RwLock::new(entries),
persistence: Some(persistence),
}
}
pub fn invalidate_persisted(
data_dir: &Path,
provider: SubscriptionProvider,
router_account: &str,
) -> Result<(), String> {
let directory = data_dir.join(CATALOG_INVALIDATION_DIR);
secure_directory(&directory).map_err(|error| {
format!(
"could not create model-catalog invalidation directory {}: {error}",
directory.display()
)
})?;
let path = invalidation_path(&directory, provider, router_account);
crate::durable_file::atomic_write_owner_only(
&path,
chrono::Utc::now().timestamp_millis().to_string().as_bytes(),
)
.map_err(|error| {
format!(
"could not invalidate the {provider} model catalog for {router_account}: {error}"
)
})
}
#[must_use]
pub fn models(&self, provider: SubscriptionProvider) -> Vec<String> {
let mut models = {
let entries = self
.entries
.read()
.unwrap_or_else(std::sync::PoisonError::into_inner);
entries
.iter()
.filter(|((entry_provider, _), _)| *entry_provider == provider)
.flat_map(|((_, account), status)| {
self.routable_status(provider, account, status)
.routable_models()
.to_vec()
})
.collect::<Vec<_>>()
};
models.sort();
models.dedup();
models
}
#[must_use]
#[allow(clippy::significant_drop_tightening)]
pub fn records(&self, provider: SubscriptionProvider) -> Vec<CatalogRecord> {
let entries = self
.entries
.read()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let mut matching = entries
.iter()
.filter(|((entry_provider, _), _)| *entry_provider == provider)
.collect::<Vec<_>>();
matching.sort_by(|((_, left), _), ((_, right), _)| left.cmp(right));
matching
.into_iter()
.flat_map(|((_, account), status)| {
self.routable_status(provider, account, status)
.routable_records()
.to_vec()
})
.collect()
}
pub(crate) fn records_for_accounts(
&self,
provider: SubscriptionProvider,
accounts: &[String],
) -> Vec<CatalogRecord> {
accounts
.iter()
.flat_map(|account| {
self.status_for(provider, account)
.routable_records()
.to_vec()
})
.collect()
}
#[must_use]
pub fn provider_is_degraded(&self, provider: SubscriptionProvider) -> bool {
let statuses = {
let entries = self
.entries
.read()
.unwrap_or_else(std::sync::PoisonError::into_inner);
entries
.iter()
.filter(|((entry_provider, _), _)| *entry_provider == provider)
.map(|((_, account), status)| self.routable_status(provider, account, status))
.collect::<Vec<_>>()
};
let mut matching = statuses.iter();
let Some(first) = matching.next() else {
return true;
};
first.is_degraded() && matching.all(CatalogStatus::is_degraded)
}
#[must_use]
pub fn provider_has_observation(&self, provider: SubscriptionProvider) -> bool {
self.entries
.read()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.iter()
.any(|((entry_provider, _), status)| {
*entry_provider == provider && (status.discovered || status.last_error.is_some())
})
}
#[must_use]
pub fn status(&self, provider: SubscriptionProvider) -> CatalogStatus {
self.status_for(provider, crate::credential_recovery_store::PRIMARY_ACCOUNT)
}
#[must_use]
pub fn status_for(&self, provider: SubscriptionProvider, account: &str) -> CatalogStatus {
let (status, effective_account) = {
let entries = self
.entries
.read()
.unwrap_or_else(std::sync::PoisonError::into_inner);
entries.get(&(provider, account.to_string())).map_or_else(
|| {
let fallback = (account != crate::credential_recovery_store::PRIMARY_ACCOUNT)
.then(|| {
entries.get(&(
provider,
crate::credential_recovery_store::PRIMARY_ACCOUNT.to_string(),
))
})
.flatten()
.filter(|primary| primary.account.is_none())
.cloned();
(fallback, crate::credential_recovery_store::PRIMARY_ACCOUNT)
},
|status| (Some(status.clone()), account),
)
};
status
.map(|status| self.routable_status(provider, effective_account, &status))
.unwrap_or_default()
}
#[must_use]
pub fn statuses(&self) -> Vec<(SubscriptionProvider, CatalogStatus)> {
let mut entries: Vec<_> = self
.entries
.read()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.iter()
.filter(|((_, account), _)| {
account == crate::credential_recovery_store::PRIMARY_ACCOUNT
})
.map(|((provider, account), status)| {
(*provider, self.routable_status(*provider, account, status))
})
.collect();
entries.sort_by_key(|(provider, _)| provider.to_string());
entries
}
pub fn record_success(&self, provider: SubscriptionProvider, models: Vec<String>) {
self.record_success_for(provider, None, models);
}
pub fn record_success_for(
&self,
provider: SubscriptionProvider,
account: Option<String>,
models: Vec<String>,
) {
self.record_success_for_account(
provider,
crate::credential_recovery_store::PRIMARY_ACCOUNT,
account,
models,
);
}
pub fn record_success_for_account(
&self,
provider: SubscriptionProvider,
router_account: &str,
account: Option<String>,
mut models: Vec<String>,
) {
models.sort();
models.dedup();
let fetched_at = chrono::Utc::now().timestamp();
let record_account = account.as_deref().unwrap_or(router_account);
let records = models
.iter()
.cloned()
.enumerate()
.map(|(index, model)| {
CatalogRecord::synthetic(provider, record_account, model, index, fetched_at)
})
.collect();
self.record_records_for_account(provider, router_account, account, records);
}
pub fn record_records_for_account(
&self,
provider: SubscriptionProvider,
router_account: &str,
account: Option<String>,
mut records: Vec<CatalogRecord>,
) {
let mut seen = HashSet::new();
records.retain(|record| seen.insert(record.canonical_id.clone()));
let mut models = records
.iter()
.map(|record| record.canonical_id.clone())
.collect::<Vec<_>>();
models.sort();
models.dedup();
let mut entries = self
.entries
.write()
.unwrap_or_else(std::sync::PoisonError::into_inner);
entries.insert(
(provider, router_account.to_string()),
CatalogStatus {
models,
records,
account,
refreshed_at: Some(chrono::Utc::now().timestamp()),
last_error: None,
discovered: true,
credential_healthy: true,
},
);
drop(entries);
if self.persist_entries() {
self.clear_invalidation(provider, router_account);
}
}
#[cfg(test)]
pub(crate) fn record_failure(
&self,
provider: SubscriptionProvider,
error: &str,
credential_rejected: bool,
) {
self.record_failure_for_account(
provider,
crate::credential_recovery_store::PRIMARY_ACCOUNT,
error,
credential_rejected,
);
}
pub(crate) fn record_failure_for_account(
&self,
provider: SubscriptionProvider,
account: &str,
error: &str,
credential_rejected: bool,
) {
let mut entries = self
.entries
.write()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let entry = entries.entry((provider, account.to_string())).or_default();
entry.last_error = Some(error.to_string());
if credential_rejected {
entry.credential_healthy = false;
}
drop(entries);
self.persist_entries();
}
fn routable_status(
&self,
provider: SubscriptionProvider,
router_account: &str,
status: &CatalogStatus,
) -> CatalogStatus {
let mut status = status.clone();
if self.is_invalidated(provider, router_account) {
status.credential_healthy = false;
status.last_error =
Some("authorization changed; awaiting authenticated catalog refresh".to_string());
}
status
}
fn is_invalidated(&self, provider: SubscriptionProvider, router_account: &str) -> bool {
self.persistence.as_ref().is_some_and(|persistence| {
invalidation_path(&persistence.invalidations, provider, router_account)
.try_exists()
.unwrap_or(true)
})
}
fn clear_invalidation(&self, provider: SubscriptionProvider, router_account: &str) {
let Some(persistence) = &self.persistence else {
return;
};
let path = invalidation_path(&persistence.invalidations, provider, router_account);
if let Err(error) = std::fs::remove_file(&path)
&& error.kind() != std::io::ErrorKind::NotFound
{
tracing::warn!(
"could not clear the {provider} model-catalog invalidation for {router_account}: {error}"
);
}
}
fn persist_entries(&self) -> bool {
let Some(persistence) = &self.persistence else {
return true;
};
let mut entries = self
.entries
.read()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.iter()
.map(
|((provider, router_account), status)| PersistedCatalogEntry {
provider: *provider,
router_account: router_account.clone(),
status: status.clone(),
},
)
.collect::<Vec<_>>();
entries.sort_by(|left, right| {
left.provider
.as_str()
.cmp(right.provider.as_str())
.then_with(|| left.router_account.cmp(&right.router_account))
});
let document = PersistedCatalogs {
version: PERSISTED_CATALOG_VERSION,
entries,
};
let result = serde_json::to_vec_pretty(&document)
.map_err(std::io::Error::other)
.and_then(|bytes| {
crate::durable_file::atomic_write_owner_only(&persistence.path, &bytes)
});
if let Err(error) = result {
tracing::warn!("could not persist the live model catalog: {error}");
return false;
}
true
}
}
fn load_persisted_catalogs(
path: &Path,
) -> Result<HashMap<(SubscriptionProvider, String), CatalogStatus>, String> {
let bytes = match std::fs::read(path) {
Ok(bytes) => bytes,
Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(HashMap::new()),
Err(error) => return Err(error.to_string()),
};
let persisted: PersistedCatalogs =
serde_json::from_slice(&bytes).map_err(|error| error.to_string())?;
if persisted.version != PERSISTED_CATALOG_VERSION {
return Err(format!(
"unsupported persisted catalog version {}",
persisted.version
));
}
Ok(persisted
.entries
.into_iter()
.map(|entry| ((entry.provider, entry.router_account), entry.status))
.collect())
}
fn invalidation_path(
directory: &Path,
provider: SubscriptionProvider,
router_account: &str,
) -> PathBuf {
use sha2::Digest as _;
let digest = sha2::Sha256::digest(format!("{provider}\0{router_account}").as_bytes());
directory.join(format!(
"{}-{}.invalidated",
provider.as_str(),
hex::encode(digest)
))
}
fn secure_directory(path: &Path) -> std::io::Result<()> {
std::fs::create_dir_all(path)?;
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt as _;
std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o700))?;
}
Ok(())
}
pub async fn refresh_catalogs(
client: &reqwest::Client,
readers: &[SubscriptionReader],
token_cache: &crate::refresh::TokenCache,
cache: &ModelCatalogCache,
) {
let readers = readers
.iter()
.cloned()
.map(|reader| {
(
crate::credential_recovery_store::PRIMARY_ACCOUNT.to_string(),
reader,
)
})
.collect::<Vec<_>>();
refresh_catalogs_for_accounts(client, &readers, token_cache, cache).await;
}
pub async fn refresh_catalogs_for_accounts(
client: &reqwest::Client,
readers: &[(String, SubscriptionReader)],
token_cache: &crate::refresh::TokenCache,
cache: &ModelCatalogCache,
) {
let now_ms = chrono::Utc::now().timestamp_millis();
for (account, reader) in readers {
token_cache.register_reader(account, reader);
}
let refreshes = readers.iter().map(|(router_account, reader)| async move {
let provider = reader.provider();
let token = match token_cache
.get_fresh_registered(client, provider, router_account, now_ms)
.await
{
Ok(token) => token,
Err(error) => return (provider, router_account, false, Err(error)),
};
let stamped_expired = token.is_expired(now_ms);
let mut account = token.account_id.clone();
let mut result = fetch_provider_catalog_records(client, provider, &token, None).await;
if result
.as_ref()
.is_err_and(|error| is_credential_rejection(error))
&& let Some(refreshed) = token_cache
.refresh_rejected(client, provider, router_account, token, now_ms)
.await
{
tracing::info!(
"{provider} rejected an unexpired catalog token; re-probing once after refresh"
);
account = refreshed.account_id.clone();
result = fetch_provider_catalog_records(client, provider, &refreshed, None).await;
}
let result = result.map(|models| (account, models));
(provider, router_account, stamped_expired, result)
});
for (provider, router_account, stamped_expired, result) in
futures_util::future::join_all(refreshes).await
{
match result {
Ok((account, records)) => {
tracing::info!(
"refreshed {provider} model catalog with {} model(s)",
records.len()
);
token_cache.record_credential_working_for(provider, router_account);
cache.record_records_for_account(provider, router_account, account, records);
}
Err(error) => {
let rejected = is_credential_rejection(&error);
if rejected {
token_cache.record_credential_rejected_for(provider, router_account);
token_cache.announce_unusable_for(provider, router_account, &error);
}
let permission_refusal = is_permission_refusal(&error);
let error = if stamped_expired {
format!("{error} (credential is stamped expired; last known catalog retained)")
} else {
error
};
if permission_refusal {
tracing::warn!(
"{provider} refused the catalog request on permission grounds, not \
credential grounds; the stored token is unchanged and will be retried \
on the next tick: {error}"
);
} else if cache
.status_for(provider, router_account)
.last_error
.as_deref()
== Some(error.as_str())
{
tracing::debug!("{provider} model catalog is still failing: {error}");
} else {
tracing::warn!("failed to refresh {provider} model catalog: {error}");
}
cache.record_failure_for_account(provider, router_account, &error, rejected);
}
}
}
}
pub async fn refresh_catalogs_forever(
client: reqwest::Client,
readers: Vec<SubscriptionReader>,
token_cache: std::sync::Arc<crate::refresh::TokenCache>,
cache: std::sync::Arc<ModelCatalogCache>,
) {
loop {
refresh_catalogs(&client, &readers, &token_cache, &cache).await;
tokio::time::sleep(CATALOG_TTL).await;
}
}
pub async fn refresh_catalogs_for_accounts_forever(
client: reqwest::Client,
readers: Vec<(String, SubscriptionReader)>,
token_cache: std::sync::Arc<crate::refresh::TokenCache>,
cache: std::sync::Arc<ModelCatalogCache>,
) {
loop {
refresh_catalogs_for_accounts(&client, &readers, &token_cache, &cache).await;
tokio::time::sleep(CATALOG_TTL).await;
}
}
pub async fn fetch_provider_catalog(
client: &reqwest::Client,
provider: SubscriptionProvider,
token: &SubscriptionToken,
base_url_override: Option<&str>,
) -> Result<Vec<String>, String> {
fetch_provider_catalog_records(client, provider, token, base_url_override)
.await
.map(|records| {
records
.into_iter()
.map(|record| record.canonical_id)
.collect()
})
}
pub async fn fetch_provider_catalog_records(
client: &reqwest::Client,
provider: SubscriptionProvider,
token: &SubscriptionToken,
base_url_override: Option<&str>,
) -> Result<Vec<CatalogRecord>, String> {
let client =
crate::upstream_client::subscription_client(client, provider, base_url_override.is_some());
let base = base_url_override.map_or_else(
|| catalog_base_url(provider, token),
|value| value.trim_end_matches('/').to_string(),
);
let client_version = crate::codex_identity::client_version();
let url = match provider {
SubscriptionProvider::Claude => format!("{base}/v1/models"),
SubscriptionProvider::Codex | SubscriptionProvider::Qwen => format!("{base}/models"),
SubscriptionProvider::Gemini => format!("{base}/v1beta/models"),
};
let base_url =
reqwest::Url::parse(&url).map_err(|error| format!("invalid catalog URL: {error}"))?;
let account = token.account_id.clone().unwrap_or_else(|| "primary".into());
let fetched_at = chrono::Utc::now().timestamp();
let generation = format!("{provider}:{account}:{}", uuid::Uuid::new_v4());
let mut cursor: Option<(String, String)> = None;
let mut visited = HashSet::new();
let mut records = Vec::new();
for page in 0..MAX_CATALOG_PAGES {
if let Some((key, value)) = cursor.as_ref()
&& !visited.insert(format!("{key}={value}"))
{
return Err(format!("repeated pagination cursor for {provider}: {key}"));
}
let mut page_url = base_url.clone();
{
let mut query = page_url.query_pairs_mut();
match provider {
SubscriptionProvider::Claude => {
query.append_pair("limit", "1000");
}
SubscriptionProvider::Gemini => {
query.append_pair("pageSize", "1000");
}
SubscriptionProvider::Codex => {
query.append_pair("client_version", &client_version);
}
SubscriptionProvider::Qwen => {}
}
if let Some((key, value)) = cursor.as_ref() {
query.append_pair(key, value);
}
}
let mut request = client
.get(page_url)
.bearer_auth(&token.access_token)
.timeout(FETCH_TIMEOUT);
match provider {
SubscriptionProvider::Claude => {
request = request
.header("anthropic-version", "2023-06-01")
.header("anthropic-beta", "oauth-2025-04-20");
}
SubscriptionProvider::Codex => {
request =
request.headers(crate::codex_identity::headers(token.account_id.as_deref()));
}
SubscriptionProvider::Gemini | SubscriptionProvider::Qwen => {}
}
let response = request
.send()
.await
.map_err(|error| format!("request failed: {error}"))?;
let status = response.status();
if !status.is_success() {
let body = response.text().await.unwrap_or_default();
let detail = body.chars().take(240).collect::<String>();
return Err(format!("HTTP {status}: {detail}"));
}
let body: Value = response
.json()
.await
.map_err(|error| format!("invalid JSON response: {error}"))?;
let page_records = parse_catalog_records(
provider,
&body,
&account,
fetched_at,
&generation,
records.len(),
)?;
cursor = next_catalog_cursor(provider, &body, &page_records)?;
records.extend(page_records);
if cursor.is_none() {
return Ok(records);
}
if page + 1 == MAX_CATALOG_PAGES {
return Err(format!(
"{provider} catalog exceeded the {MAX_CATALOG_PAGES}-page safety limit"
));
}
}
unreachable!("bounded catalog loop returns on its final iteration")
}
fn catalog_base_url(provider: SubscriptionProvider, token: &SubscriptionToken) -> String {
match provider {
SubscriptionProvider::Gemini => "https://generativelanguage.googleapis.com".to_string(),
_ => token.base_url(provider).trim_end_matches('/').to_string(),
}
}
#[cfg(test)]
#[path = "model_catalog_tests.rs"]
mod tests;
#[cfg(test)]
#[path = "model_catalog_invalidation_tests.rs"]
mod invalidation_tests;