diff --git a/codex-rs/core/tests/common/test_codex.rs b/codex-rs/core/tests/common/test_codex.rs index 762c84832e..6a09c5cdea 100644 --- a/codex-rs/core/tests/common/test_codex.rs +++ b/codex-rs/core/tests/common/test_codex.rs @@ -36,6 +36,7 @@ use codex_login::CodexAuth; use codex_model_provider_info::ModelProviderInfo; use codex_model_provider_info::built_in_model_providers; use codex_models_manager::bundled_models_response; +use codex_models_manager::manager::SharedModelsManager; use codex_protocol::mcp::ClientMcpExtensions; use codex_protocol::mcp::OPENAI_FORM_EXTENSION_ID; use codex_protocol::models::PermissionProfile; @@ -308,6 +309,7 @@ pub struct TestCodexBuilder { external_time_provider: Option>, code_mode_host_program: Option, history_mode: Option, + models_manager: Option, } impl TestCodexBuilder { @@ -324,6 +326,11 @@ impl TestCodexBuilder { self } + pub fn with_models_manager(mut self, models_manager: SharedModelsManager) -> Self { + self.models_manager = Some(models_manager); + self + } + pub fn with_model(self, model: &str) -> Self { let new_model = model.to_string(); self.with_config(move |config| { @@ -643,10 +650,14 @@ impl TestCodexBuilder { )) }); let auth_manager = codex_core::test_support::auth_manager_from_auth(auth.clone()); + let models_manager = self + .models_manager + .clone() + .unwrap_or_else(|| codex_core::build_models_manager(&config, auth_manager.clone())); let thread_manager = ThreadManager::new( &config, auth_manager.clone(), - codex_core::build_models_manager(&config, auth_manager), + models_manager, codex_core::CodexAppsToolsCache::default(), SessionSource::Exec, Arc::clone(&environment_manager), @@ -1281,6 +1292,7 @@ pub fn test_codex() -> TestCodexBuilder { external_time_provider: None, code_mode_host_program: None, history_mode: None, + models_manager: None, } } diff --git a/codex-rs/core/tests/suite/injected_models_cache.rs b/codex-rs/core/tests/suite/injected_models_cache.rs new file mode 100644 index 0000000000..ba8c123bd4 --- /dev/null +++ b/codex-rs/core/tests/suite/injected_models_cache.rs @@ -0,0 +1,256 @@ +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; + +use anyhow::Result; +use chrono::Utc; +use codex_http_client::HttpClientFactory; +use codex_login::AuthManager; +use codex_login::CodexAuth; +use codex_models_manager::cache::ModelsCache; +use codex_models_manager::cache::ModelsCacheEntry; +use codex_models_manager::cache::ModelsCacheError; +use codex_models_manager::cache::ModelsCacheFuture; +use codex_models_manager::manager::ModelsEndpointClient; +use codex_models_manager::manager::ModelsEndpointFuture; +use codex_models_manager::manager::OpenAiModelsManager; +use codex_models_manager::manager::RefreshStrategy; +use codex_models_manager::manager::SharedModelsManager; +use codex_models_manager::model_info::model_info_from_slug; +use codex_protocol::error::Result as CoreResult; +use codex_protocol::openai_models::ModelInfo; +use codex_protocol::openai_models::ModelVisibility; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::Op; +use codex_protocol::protocol::ThreadSettingsOverrides; +use codex_protocol::user_input::UserInput; +use core_test_support::responses; +use core_test_support::responses::ev_assistant_message; +use core_test_support::responses::ev_completed; +use core_test_support::responses::ev_response_created; +use core_test_support::responses::mount_sse_once; +use core_test_support::responses::sse; +use core_test_support::submit_thread_settings; +use core_test_support::test_codex::test_codex; +use core_test_support::wait_for_event; +use pretty_assertions::assert_eq; + +#[derive(Debug)] +struct TestModelsCache { + entry: Option, + load_error: bool, + stored_entries: Mutex>, +} + +impl ModelsCache for TestModelsCache { + fn load<'a>( + &'a self, + _client_version: &'a str, + ) -> ModelsCacheFuture<'a, Result, ModelsCacheError>> { + Box::pin(async move { + if self.load_error { + Err(ModelsCacheError::new("test load failure")) + } else { + Ok(self.entry.clone()) + } + }) + } + + fn store<'a>( + &'a self, + entry: &'a ModelsCacheEntry, + ) -> ModelsCacheFuture<'a, Result<(), ModelsCacheError>> { + Box::pin(async move { + self.stored_entries + .lock() + .expect("stored entries lock should not be poisoned") + .push(entry.clone()); + Ok(()) + }) + } + + fn refresh_ttl<'a>( + &'a self, + _client_version: &'a str, + ) -> ModelsCacheFuture<'a, Result<(), ModelsCacheError>> { + Box::pin(async move { + let mut entry = self + .entry + .clone() + .ok_or_else(|| ModelsCacheError::new("cache not found"))?; + entry.fetched_at = Utc::now(); + self.stored_entries + .lock() + .expect("stored entries lock should not be poisoned") + .push(entry); + Ok(()) + }) + } +} + +#[derive(Debug)] +struct TestModelsEndpoint { + models: Vec, + fetch_count: AtomicUsize, +} + +impl TestModelsEndpoint { + fn new(models: Vec) -> Arc { + Arc::new(Self { + models, + fetch_count: AtomicUsize::new(0), + }) + } +} + +impl ModelsEndpointClient for TestModelsEndpoint { + fn has_command_auth(&self) -> bool { + false + } + + fn uses_codex_backend(&self) -> ModelsEndpointFuture<'_, bool> { + Box::pin(async { true }) + } + + fn list_models<'a>( + &'a self, + _client_version: &'a str, + _http_client_factory: HttpClientFactory, + ) -> ModelsEndpointFuture<'a, CoreResult<(Vec, Option)>> { + Box::pin(async move { + self.fetch_count.fetch_add(1, Ordering::SeqCst); + Ok((self.models.clone(), None)) + }) + } +} + +fn remote_model(slug: &str) -> ModelInfo { + ModelInfo { + visibility: ModelVisibility::List, + used_fallback_model_metadata: false, + ..model_info_from_slug(slug) + } +} + +fn models_manager( + cache: Arc, + endpoint: Arc, +) -> SharedModelsManager { + Arc::new(OpenAiModelsManager::new_with_cache( + cache, + endpoint, + Some(AuthManager::from_auth_for_testing( + CodexAuth::create_dummy_chatgpt_auth_for_testing(), + )), + )) +} + +async fn run_agent_with_model(models_manager: SharedModelsManager, model_slug: &str) -> Result<()> { + let server = responses::start_mock_server().await; + let response_mock = mount_sse_once( + &server, + sse(vec![ + ev_response_created("resp-1"), + ev_assistant_message("msg-1", "done"), + ev_completed("resp-1"), + ]), + ) + .await; + let mut builder = test_codex() + .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) + .with_models_manager(models_manager); + let test = builder.build(&server).await?; + let available_models = test + .thread_manager + .get_models_manager() + .list_models( + RefreshStrategy::OnlineIfUncached, + codex_core::test_support::default_http_client_factory(), + ) + .await; + assert!( + available_models + .iter() + .any(|model| model.model == model_slug) + ); + + submit_thread_settings( + &test.codex, + ThreadSettingsOverrides { + model: Some(model_slug.to_string()), + ..Default::default() + }, + ) + .await?; + test.codex + .submit(Op::UserInput { + items: vec![UserInput::Text { + text: "hello".to_string(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + responsesapi_client_metadata: None, + additional_context: Default::default(), + thread_settings: Default::default(), + }) + .await?; + loop { + if matches!( + wait_for_event(&test.codex, |_| true).await, + EventMsg::TurnComplete(_) + ) { + break; + } + } + + assert_eq!( + response_mock.single_request().body_json()["model"], + model_slug + ); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn injected_cache_hit_drives_agent_model_selection() -> Result<()> { + let model_slug = "injected-cache-model"; + let cache = Arc::new(TestModelsCache { + entry: Some(ModelsCacheEntry { + fetched_at: Utc::now(), + etag: None, + client_version: Some(codex_models_manager::client_version_to_whole()), + models: vec![remote_model(model_slug)], + }), + load_error: false, + stored_entries: Mutex::new(Vec::new()), + }); + let endpoint = TestModelsEndpoint::new(Vec::new()); + + run_agent_with_model(models_manager(cache, endpoint.clone()), model_slug).await?; + + assert_eq!(endpoint.fetch_count.load(Ordering::SeqCst), 0); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn injected_cache_error_falls_back_for_agent_model_selection() -> Result<()> { + let model_slug = "injected-cache-fallback-model"; + let cache = Arc::new(TestModelsCache { + entry: None, + load_error: true, + stored_entries: Mutex::new(Vec::new()), + }); + let endpoint = TestModelsEndpoint::new(vec![remote_model(model_slug)]); + + run_agent_with_model(models_manager(cache.clone(), endpoint.clone()), model_slug).await?; + + assert!(endpoint.fetch_count.load(Ordering::SeqCst) >= 1); + assert!( + !cache + .stored_entries + .lock() + .expect("stored entries lock should not be poisoned") + .is_empty() + ); + Ok(()) +} diff --git a/codex-rs/core/tests/suite/mod.rs b/codex-rs/core/tests/suite/mod.rs index 7f7311e9b6..a1668778e4 100644 --- a/codex-rs/core/tests/suite/mod.rs +++ b/codex-rs/core/tests/suite/mod.rs @@ -72,6 +72,7 @@ mod hooks; #[cfg(not(target_os = "windows"))] mod hooks_mcp; mod image_rollout; +mod injected_models_cache; mod items; mod json_result; mod live_cli; diff --git a/codex-rs/model-provider/src/provider.rs b/codex-rs/model-provider/src/provider.rs index 1964864f2c..8b21cbfdf5 100644 --- a/codex-rs/model-provider/src/provider.rs +++ b/codex-rs/model-provider/src/provider.rs @@ -11,6 +11,7 @@ use codex_api::is_azure_responses_provider; use codex_login::AuthManager; use codex_login::CodexAuth; use codex_model_provider_info::ModelProviderInfo; +use codex_models_manager::cache::ModelsCache; use codex_models_manager::manager::OpenAiModelsManager; use codex_models_manager::manager::SharedModelsManager; use codex_models_manager::manager::StaticModelsManager; @@ -229,6 +230,21 @@ pub trait ModelProvider: fmt::Debug + Send + Sync { .unwrap_or_default(); Arc::new(StaticModelsManager::new(self.auth_manager(), model_catalog)) } + + /// Creates a model manager that can use a caller-provided cache for remote catalogs. + /// + /// Providers with remote catalogs should override this method. The default preserves the + /// authoritative catalog returned by [`ModelProvider::models_manager_without_cache`] and does + /// not consult `cache`. Implementations should likewise ignore the cache when + /// `config_model_catalog` supplies an authoritative static catalog. + fn models_manager_with_cache( + &self, + config_model_catalog: Option, + cache: Arc, + ) -> SharedModelsManager { + drop(cache); + self.models_manager_without_cache(config_model_catalog) + } } pub type ModelProviderFuture<'a, T> = Pin + Send + 'a>>; @@ -401,6 +417,30 @@ impl ModelProvider for ConfiguredModelProvider { } } } + + fn models_manager_with_cache( + &self, + config_model_catalog: Option, + cache: Arc, + ) -> SharedModelsManager { + match config_model_catalog { + Some(model_catalog) => Arc::new(StaticModelsManager::new( + self.auth_manager.clone(), + model_catalog, + )), + None => { + let endpoint = Arc::new(OpenAiModelsEndpoint::new( + self.info.clone(), + self.auth_manager.clone(), + )); + Arc::new(OpenAiModelsManager::new_with_cache( + cache, + endpoint, + self.auth_manager.clone(), + )) + } + } + } } #[cfg(test)] diff --git a/codex-rs/models-manager/src/cache.rs b/codex-rs/models-manager/src/cache.rs index 9fe4667a76..15e8cf4dd8 100644 --- a/codex-rs/models-manager/src/cache.rs +++ b/codex-rs/models-manager/src/cache.rs @@ -3,178 +3,82 @@ use chrono::Utc; use codex_protocol::openai_models::ModelInfo; use serde::Deserialize; use serde::Serialize; +use std::fmt; +use std::future::Future; use std::io; use std::io::ErrorKind; use std::path::PathBuf; +use std::pin::Pin; use std::time::Duration; use tokio::fs; -use tracing::error; use tracing::info; -/// Manages loading and saving of models cache to disk. -#[derive(Debug)] -pub(crate) struct ModelsCacheManager { - cache_path: PathBuf, - cache_ttl: Duration, +/// Asynchronous storage for model catalog snapshots used by the models manager. +/// +/// Implementations own cache freshness and lookup partitioning. [`ModelsCache::load`] must not +/// return stale entries or entries for a different client version. A shared backend must also keep +/// catalogs for different providers and tenants separate, typically through configuration captured +/// by the implementation. +/// +/// Cache failures are non-fatal. The models manager logs them and falls back to the configured +/// models endpoint. +pub trait ModelsCache: fmt::Debug + Send + Sync { + /// Loads a fresh entry for `client_version`. + /// + /// Returns `Ok(None)` for a normal miss, including an absent, stale, or version-mismatched + /// entry. Returns `Err` when the backend could not complete the lookup; the models manager + /// treats that error as a miss and fetches the catalog from the models endpoint. + fn load<'a>( + &'a self, + client_version: &'a str, + ) -> ModelsCacheFuture<'a, Result, ModelsCacheError>>; + + /// Stores `entry`, replacing the snapshot for this cache implementation's lookup identity. + /// + /// The implementation derives that identity from its own configuration and the entry metadata. + /// For example, a shared backend can capture provider and tenant identity when it is constructed. + fn store<'a>( + &'a self, + entry: &'a ModelsCacheEntry, + ) -> ModelsCacheFuture<'a, Result<(), ModelsCacheError>>; + + /// Extends the freshness of the entry for `client_version` without changing its catalog. + /// + /// The models manager calls this after the endpoint confirms that the cached ETag is still + /// current. Implementations must preserve the models, ETag, and client version. This operation + /// cannot be expressed in terms of [`ModelsCache::load`], because `load` intentionally omits + /// expired entries. File-backed implementations can read and revalidate the stored entry + /// directly, while backends with native TTL support can extend the entry in place. + fn refresh_ttl<'a>( + &'a self, + client_version: &'a str, + ) -> ModelsCacheFuture<'a, Result<(), ModelsCacheError>>; } -impl ModelsCacheManager { - /// Create a new cache manager with the given path and TTL. - pub(crate) fn new(cache_path: PathBuf, cache_ttl: Duration) -> Self { - Self { - cache_path, - cache_ttl, - } - } +/// Boxed future returned by [`ModelsCache`] implementations. +pub type ModelsCacheFuture<'a, T> = Pin + Send + 'a>>; - /// Attempt to load a fresh cache entry. Returns `None` if the cache doesn't exist or is stale. - pub(crate) async fn load_fresh(&self, expected_version: &str) -> Option { - info!( - cache_path = %self.cache_path.display(), - expected_version, - "models cache: attempting load_fresh" - ); - let cache = match self.load().await { - Ok(cache) => cache?, - Err(err) => { - error!("failed to load models cache: {err}"); - return None; - } - }; - info!( - cache_path = %self.cache_path.display(), - cached_version = ?cache.client_version, - fetched_at = %cache.fetched_at, - "models cache: loaded cache file" - ); - if cache.client_version.as_deref() != Some(expected_version) { - info!( - cache_path = %self.cache_path.display(), - expected_version, - cached_version = ?cache.client_version, - "models cache: cache version mismatch" - ); - return None; - } - if !cache.is_fresh(self.cache_ttl) { - info!( - cache_path = %self.cache_path.display(), - cache_ttl_secs = self.cache_ttl.as_secs(), - fetched_at = %cache.fetched_at, - "models cache: cache is stale" - ); - return None; - } - info!( - cache_path = %self.cache_path.display(), - cache_ttl_secs = self.cache_ttl.as_secs(), - "models cache: cache hit" - ); - Some(cache) - } - - /// Persist the cache to disk, creating parent directories as needed. - pub(crate) async fn persist_cache( - &self, - models: &[ModelInfo], - etag: Option, - client_version: String, - ) { - let cache = ModelsCache { - fetched_at: Utc::now(), - etag, - client_version: Some(client_version), - models: models.to_vec(), - }; - if let Err(err) = self.save_internal(&cache).await { - error!("failed to write models cache: {err}"); - } - } - - /// Renew the cache TTL once more than half its lifetime has elapsed. - pub(crate) async fn renew_cache_ttl(&self) -> io::Result<()> { - let mut cache = match self.load().await? { - Some(cache) => cache, - None => return Err(io::Error::new(ErrorKind::NotFound, "cache not found")), - }; - if cache.is_fresh(self.cache_ttl / 2) { - return Ok(()); - } - cache.fetched_at = Utc::now(); - self.save_internal(&cache).await - } - - async fn load(&self) -> io::Result> { - match fs::read(&self.cache_path).await { - Ok(contents) => { - let cache = serde_json::from_slice(&contents) - .map_err(|err| io::Error::new(ErrorKind::InvalidData, err.to_string()))?; - Ok(Some(cache)) - } - Err(err) if err.kind() == ErrorKind::NotFound => Ok(None), - Err(err) => Err(err), - } - } - - async fn save_internal(&self, cache: &ModelsCache) -> io::Result<()> { - if let Some(parent) = self.cache_path.parent() { - fs::create_dir_all(parent).await?; - } - let json = serde_json::to_vec_pretty(cache) - .map_err(|err| io::Error::new(ErrorKind::InvalidData, err.to_string()))?; - fs::write(&self.cache_path, json).await - } - - #[cfg(test)] - /// Set the cache TTL. - pub(crate) fn set_ttl(&mut self, ttl: Duration) { - self.cache_ttl = ttl; - } - - #[cfg(test)] - /// Manipulate cache file for testing. Allows setting a custom fetched_at timestamp. - pub(crate) async fn manipulate_cache_for_test(&self, f: F) -> io::Result<()> - where - F: FnOnce(&mut DateTime), - { - let mut cache = match self.load().await? { - Some(cache) => cache, - None => return Err(io::Error::new(ErrorKind::NotFound, "cache not found")), - }; - f(&mut cache.fetched_at); - self.save_internal(&cache).await - } - - #[cfg(test)] - /// Mutate the full cache contents for testing. - pub(crate) async fn mutate_cache_for_test(&self, f: F) -> io::Result<()> - where - F: FnOnce(&mut ModelsCache), - { - let mut cache = match self.load().await? { - Some(cache) => cache, - None => return Err(io::Error::new(ErrorKind::NotFound, "cache not found")), - }; - f(&mut cache); - self.save_internal(&cache).await - } -} - -/// Serialized snapshot of models and metadata cached on disk. -#[derive(Debug, Clone, Serialize, Deserialize)] -pub(crate) struct ModelsCache { - pub(crate) fetched_at: DateTime, +/// Serialized model catalog and validation metadata stored in a [`ModelsCache`]. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct ModelsCacheEntry { + /// Time when the catalog was fetched or last revalidated with the models endpoint. + pub fetched_at: DateTime, + /// Endpoint ETag used to revalidate the catalog without downloading an unchanged payload. #[serde(default, skip_serializing_if = "Option::is_none")] - pub(crate) etag: Option, + pub etag: Option, + /// Client version for which this catalog was returned. + /// + /// The models manager rejects entries whose value is absent or differs from its current version. #[serde(default, skip_serializing_if = "Option::is_none")] - pub(crate) client_version: Option, + pub client_version: Option, + /// Models returned by the catalog endpoint. #[serde( deserialize_with = "codex_protocol::openai_models::deserialize_model_infos_with_legacy_base" )] - pub(crate) models: Vec, + pub models: Vec, } -impl ModelsCache { +impl ModelsCacheEntry { /// Returns `true` when the cache entry has not exceeded the configured TTL. fn is_fresh(&self, ttl: Duration) -> bool { if ttl.is_zero() { @@ -187,3 +91,147 @@ impl ModelsCache { age <= ttl_duration } } + +/// Error returned by a [`ModelsCache`] implementation. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ModelsCacheError { + message: String, +} + +impl ModelsCacheError { + /// Create a cache error without exposing backend-specific error types. + pub fn new(message: impl Into) -> Self { + Self { + message: message.into(), + } + } +} + +impl fmt::Display for ModelsCacheError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(&self.message) + } +} + +impl std::error::Error for ModelsCacheError {} + +/// Built-in file-backed model catalog cache. +#[derive(Debug)] +pub(crate) struct FileModelsCache { + cache_path: PathBuf, + cache_ttl: Duration, +} + +impl FileModelsCache { + /// Create a file-backed cache with the given path and TTL. + pub(crate) fn new(cache_path: PathBuf, cache_ttl: Duration) -> Self { + Self { + cache_path, + cache_ttl, + } + } +} + +impl ModelsCache for FileModelsCache { + fn load<'a>( + &'a self, + client_version: &'a str, + ) -> ModelsCacheFuture<'a, Result, ModelsCacheError>> { + Box::pin( + async move { load_fresh_file(&self.cache_path, self.cache_ttl, client_version).await }, + ) + } + + fn store<'a>( + &'a self, + entry: &'a ModelsCacheEntry, + ) -> ModelsCacheFuture<'a, Result<(), ModelsCacheError>> { + Box::pin(async move { save_file(&self.cache_path, entry).await }) + } + + fn refresh_ttl<'a>( + &'a self, + _client_version: &'a str, + ) -> ModelsCacheFuture<'a, Result<(), ModelsCacheError>> { + Box::pin(async move { + let mut entry = load_file(&self.cache_path) + .await + .map_err(cache_error)? + .ok_or_else(|| ModelsCacheError::new("cache not found"))?; + if entry.is_fresh(self.cache_ttl / 2) { + return Ok(()); + } + entry.fetched_at = Utc::now(); + save_file(&self.cache_path, &entry).await + }) + } +} + +async fn load_fresh_file( + cache_path: &PathBuf, + cache_ttl: Duration, + expected_version: &str, +) -> Result, ModelsCacheError> { + info!( + cache_path = %cache_path.display(), + expected_version, + "models cache: attempting load_fresh" + ); + let Some(cache) = load_file(cache_path).await.map_err(cache_error)? else { + return Ok(None); + }; + info!( + cache_path = %cache_path.display(), + cached_version = ?cache.client_version, + fetched_at = %cache.fetched_at, + "models cache: loaded cache file" + ); + if cache.client_version.as_deref() != Some(expected_version) { + info!( + cache_path = %cache_path.display(), + expected_version, + cached_version = ?cache.client_version, + "models cache: cache version mismatch" + ); + return Ok(None); + } + if !cache.is_fresh(cache_ttl) { + info!( + cache_path = %cache_path.display(), + cache_ttl_secs = cache_ttl.as_secs(), + fetched_at = %cache.fetched_at, + "models cache: cache is stale" + ); + return Ok(None); + } + info!( + cache_path = %cache_path.display(), + cache_ttl_secs = cache_ttl.as_secs(), + "models cache: cache hit" + ); + Ok(Some(cache)) +} + +async fn load_file(cache_path: &PathBuf) -> io::Result> { + match fs::read(cache_path).await { + Ok(contents) => { + let cache = serde_json::from_slice(&contents) + .map_err(|err| io::Error::new(ErrorKind::InvalidData, err.to_string()))?; + Ok(Some(cache)) + } + Err(err) if err.kind() == ErrorKind::NotFound => Ok(None), + Err(err) => Err(err), + } +} + +async fn save_file(cache_path: &PathBuf, cache: &ModelsCacheEntry) -> Result<(), ModelsCacheError> { + if let Some(parent) = cache_path.parent() { + fs::create_dir_all(parent).await.map_err(cache_error)?; + } + let json = serde_json::to_vec_pretty(cache).map_err(cache_error)?; + fs::write(cache_path, json).await.map_err(cache_error) +} + +fn cache_error(error: impl fmt::Display) -> ModelsCacheError { + ModelsCacheError::new(error.to_string()) +} diff --git a/codex-rs/models-manager/src/lib.rs b/codex-rs/models-manager/src/lib.rs index 4a4c590afe..cd7b59a053 100644 --- a/codex-rs/models-manager/src/lib.rs +++ b/codex-rs/models-manager/src/lib.rs @@ -1,4 +1,4 @@ -pub(crate) mod cache; +pub mod cache; pub mod collaboration_mode_presets; pub(crate) mod config; pub mod manager; diff --git a/codex-rs/models-manager/src/manager.rs b/codex-rs/models-manager/src/manager.rs index a0a0348e4e..45835c37cd 100644 --- a/codex-rs/models-manager/src/manager.rs +++ b/codex-rs/models-manager/src/manager.rs @@ -1,7 +1,10 @@ -use super::cache::ModelsCacheManager; +use super::cache::FileModelsCache; +use crate::cache::ModelsCache; +use crate::cache::ModelsCacheEntry; use crate::collaboration_mode_presets::builtin_collaboration_mode_presets; use crate::config::ModelsManagerConfig; use crate::model_info; +use chrono::Utc; use codex_http_client::HttpClientFactory; use codex_login::AuthManager; use codex_protocol::auth::AuthMode; @@ -215,7 +218,7 @@ pub type SharedModelsManager = Arc; pub struct OpenAiModelsManager { remote_models: RwLock>, etag: RwLock>, - cache_manager: Option, + cache: Option>, endpoint_client: SharedModelsEndpointClient, auth_manager: Option>, } @@ -235,8 +238,11 @@ impl OpenAiModelsManager { auth_manager: Option>, ) -> Self { let cache_path = codex_home.join(MODEL_CACHE_FILE); - Self::new_with_cache_manager( - Some(ModelsCacheManager::new(cache_path, DEFAULT_MODEL_CACHE_TTL)), + Self::new_with_optional_cache( + Some(Arc::new(FileModelsCache::new( + cache_path, + DEFAULT_MODEL_CACHE_TTL, + ))), endpoint_client, auth_manager, ) @@ -247,11 +253,23 @@ impl OpenAiModelsManager { endpoint_client: Arc, auth_manager: Option>, ) -> Self { - Self::new_with_cache_manager(/*cache_manager*/ None, endpoint_client, auth_manager) + Self::new_with_optional_cache(/*cache*/ None, endpoint_client, auth_manager) } - fn new_with_cache_manager( - cache_manager: Option, + /// Constructs an OpenAI-compatible model manager with a caller-provided cache. + /// + /// The cache is consulted by cache-aware refresh strategies. Cache misses and backend errors + /// fall back to the models endpoint, and cache write failures do not fail model discovery. + pub fn new_with_cache( + cache: Arc, + endpoint_client: Arc, + auth_manager: Option>, + ) -> Self { + Self::new_with_optional_cache(Some(cache), endpoint_client, auth_manager) + } + + fn new_with_optional_cache( + cache: Option>, endpoint_client: Arc, auth_manager: Option>, ) -> Self { @@ -259,7 +277,7 @@ impl OpenAiModelsManager { Self { remote_models: RwLock::new(remote_models), etag: RwLock::new(None), - cache_manager, + cache, endpoint_client, auth_manager, } @@ -338,8 +356,8 @@ impl OpenAiModelsManager { async fn refresh_if_new_etag(&self, etag: String, http_client_factory: HttpClientFactory) { let current_etag = self.get_etag().await; if current_etag.clone().is_some() && current_etag.as_deref() == Some(etag.as_str()) { - if let Some(cache_manager) = self.cache_manager.as_ref() - && let Err(err) = cache_manager.renew_cache_ttl().await + if let Some(cache) = self.cache.as_ref() + && let Err(err) = cache.refresh_ttl(&crate::client_version_to_whole()).await { error!("failed to renew cache TTL: {err}"); } @@ -402,10 +420,16 @@ impl OpenAiModelsManager { .await?; self.apply_remote_models(models.clone()).await; *self.etag.write().await = etag.clone(); - if let Some(cache_manager) = self.cache_manager.as_ref() { - cache_manager - .persist_cache(&models, etag, client_version) - .await; + if let Some(cache) = self.cache.as_ref() { + let entry = ModelsCacheEntry { + fetched_at: Utc::now(), + etag, + client_version: Some(client_version), + models, + }; + if let Err(err) = cache.store(&entry).await { + error!("failed to write models cache: {err}"); + } } Ok(()) } @@ -452,7 +476,7 @@ impl OpenAiModelsManager { /// Attempt to satisfy the refresh from the cache when it matches the provider and TTL. async fn try_load_cache(&self) -> bool { - let Some(cache_manager) = self.cache_manager.as_ref() else { + let Some(cache) = self.cache.as_ref() else { return false; }; let _timer = @@ -461,19 +485,31 @@ impl OpenAiModelsManager { info!(client_version, "models cache: evaluating cache eligibility"); // TODO(celia-oai): Include provider identity in cache eligibility so switching // providers does not reuse a fresh models_cache.json entry from another provider. - let cache = match cache_manager.load_fresh(&client_version).await { - Some(cache) => cache, - None => { + let cache_entry = match cache.load(&client_version).await { + Ok(Some(cache_entry)) => cache_entry, + Ok(None) => { info!("models cache: no usable cache entry"); return false; } + Err(err) => { + error!("failed to load models cache: {err}"); + return false; + } }; - let models = cache.models.clone(); - *self.etag.write().await = cache.etag.clone(); + if cache_entry.client_version.as_deref() != Some(client_version.as_str()) { + info!( + expected_version = client_version, + cached_version = ?cache_entry.client_version, + "models cache: cache version mismatch" + ); + return false; + } + let models = cache_entry.models.clone(); + *self.etag.write().await = cache_entry.etag.clone(); self.apply_remote_models(models.clone()).await; info!( models_count = models.len(), - etag = ?cache.etag, + etag = ?cache_entry.etag, "models cache: cache entry applied" ); true diff --git a/codex-rs/models-manager/src/manager_tests.rs b/codex-rs/models-manager/src/manager_tests.rs index b63ce1c222..0768d1e4c4 100644 --- a/codex-rs/models-manager/src/manager_tests.rs +++ b/codex-rs/models-manager/src/manager_tests.rs @@ -1,5 +1,10 @@ use super::*; use crate::ModelsManagerConfig; +use crate::cache::FileModelsCache; +use crate::cache::ModelsCache; +use crate::cache::ModelsCacheEntry; +use crate::cache::ModelsCacheError; +use crate::cache::ModelsCacheFuture; use chrono::Utc; use codex_http_client::HttpClientFactory; use codex_http_client::OutboundProxyPolicy; @@ -89,6 +94,95 @@ struct TestModelsEndpoint { observed_proxy_policy: Mutex>, } +#[derive(Debug)] +struct TestModelsCache { + entry: Mutex>, + load_error: bool, + store_error: bool, + stored_entries: Mutex>, +} + +impl TestModelsCache { + fn with_entry(entry: ModelsCacheEntry) -> Arc { + Arc::new(Self { + entry: Mutex::new(Some(entry)), + load_error: false, + store_error: false, + stored_entries: Mutex::new(Vec::new()), + }) + } + + fn failing(load_error: bool, store_error: bool) -> Arc { + Arc::new(Self { + entry: Mutex::new(None), + load_error, + store_error, + stored_entries: Mutex::new(Vec::new()), + }) + } + + fn stored_entries(&self) -> Vec { + self.stored_entries + .lock() + .expect("stored entries lock should not be poisoned") + .clone() + } +} + +impl ModelsCache for TestModelsCache { + fn load<'a>( + &'a self, + _client_version: &'a str, + ) -> ModelsCacheFuture<'a, Result, ModelsCacheError>> { + Box::pin(async move { + if self.load_error { + return Err(ModelsCacheError::new("test load failure")); + } + Ok(self + .entry + .lock() + .expect("cache entry lock should not be poisoned") + .clone()) + }) + } + + fn store<'a>( + &'a self, + entry: &'a ModelsCacheEntry, + ) -> ModelsCacheFuture<'a, Result<(), ModelsCacheError>> { + Box::pin(async move { + if self.store_error { + return Err(ModelsCacheError::new("test store failure")); + } + self.stored_entries + .lock() + .expect("stored entries lock should not be poisoned") + .push(entry.clone()); + Ok(()) + }) + } + + fn refresh_ttl<'a>( + &'a self, + _client_version: &'a str, + ) -> ModelsCacheFuture<'a, Result<(), ModelsCacheError>> { + Box::pin(async move { + let refreshed = { + let mut entry = self + .entry + .lock() + .expect("cache entry lock should not be poisoned"); + let entry = entry + .as_mut() + .ok_or_else(|| ModelsCacheError::new("cache not found"))?; + entry.fetched_at = Utc::now(); + entry.clone() + }; + self.store(&refreshed).await + }) + } +} + impl TestModelsEndpoint { fn new(responses: Vec>) -> Arc { Arc::new(Self { @@ -211,10 +305,101 @@ fn openai_manager_for_tests_with_auth( OpenAiModelsManager::new(codex_home, endpoint_client, auth_manager) } +async fn mutate_file_cache_for_test(codex_home: &Path, f: F) +where + F: FnOnce(&mut ModelsCacheEntry), +{ + let client_version = crate::client_version_to_whole(); + let cache = FileModelsCache::new(codex_home.join(MODEL_CACHE_FILE), DEFAULT_MODEL_CACHE_TTL); + let mut entry = cache + .load(&client_version) + .await + .expect("cache load succeeds") + .expect("cache entry exists"); + f(&mut entry); + cache.store(&entry).await.expect("cache store succeeds"); +} + fn static_manager_for_tests(model_catalog: ModelsResponse) -> StaticModelsManager { StaticModelsManager::new(/*auth_manager*/ None, model_catalog) } +#[tokio::test] +async fn file_cache_implements_models_cache_contract() { + let codex_home = tempdir().expect("temp dir"); + let cache = FileModelsCache::new( + codex_home.path().join(MODEL_CACHE_FILE), + DEFAULT_MODEL_CACHE_TTL, + ); + let client_version = crate::client_version_to_whole(); + let entry = ModelsCacheEntry { + fetched_at: Utc::now(), + etag: Some("file-etag".to_string()), + client_version: Some(client_version.clone()), + models: vec![remote_model( + "file-cached", + "File Cached", + /*priority*/ 0, + )], + }; + + cache.store(&entry).await.expect("cache store succeeds"); + + assert_eq!( + cache + .load(&client_version) + .await + .expect("cache load succeeds"), + Some(entry) + ); +} + +#[tokio::test] +async fn file_cache_refresh_ttl_renews_expired_entry_without_serving_it_stale() { + let codex_home = tempdir().expect("temp dir"); + let cache = FileModelsCache::new( + codex_home.path().join(MODEL_CACHE_FILE), + DEFAULT_MODEL_CACHE_TTL, + ); + let client_version = crate::client_version_to_whole(); + let expired_at = Utc::now() - chrono::Duration::hours(1); + let entry = ModelsCacheEntry { + fetched_at: expired_at, + etag: Some("expired-etag".to_string()), + client_version: Some(client_version.clone()), + models: vec![remote_model( + "expired-file-cache", + "Expired File Cache", + /*priority*/ 0, + )], + }; + cache.store(&entry).await.expect("cache store succeeds"); + + assert_eq!( + cache + .load(&client_version) + .await + .expect("cache load succeeds"), + None, + "an expired entry must not be returned before revalidation" + ); + + cache + .refresh_ttl(&client_version) + .await + .expect("TTL refresh succeeds"); + + let refreshed = cache + .load(&client_version) + .await + .expect("cache load succeeds") + .expect("revalidated entry is fresh"); + assert!(refreshed.fetched_at > expired_at); + let mut expected = entry; + expected.fetched_at = refreshed.fetched_at; + assert_eq!(refreshed, expected); +} + #[tokio::test] async fn manager_without_cache_fetches_on_every_refresh() { let remote_models = vec![remote_model("remote", "Remote", /*priority*/ 0)]; @@ -245,6 +430,134 @@ async fn manager_without_cache_fetches_on_every_refresh() { assert_eq!(endpoint.fetch_count(), 2); } +#[tokio::test] +async fn injected_cache_hit_avoids_remote_fetch() { + let cached_models = vec![remote_model("cached", "Cached", /*priority*/ 0)]; + let cache = TestModelsCache::with_entry(ModelsCacheEntry { + fetched_at: Utc::now(), + etag: Some("cached-etag".to_string()), + client_version: Some(crate::client_version_to_whole()), + models: cached_models.clone(), + }); + let endpoint = TestModelsEndpoint::new(vec![vec![remote_model( + "remote", "Remote", /*priority*/ 0, + )]]); + let manager = OpenAiModelsManager::new_with_cache( + cache, + endpoint.clone(), + Some(AuthManager::from_auth_for_testing( + CodexAuth::create_dummy_chatgpt_auth_for_testing(), + )), + ); + + let catalog = manager + .raw_model_catalog( + RefreshStrategy::OnlineIfUncached, + DEFAULT_HTTP_CLIENT_FACTORY, + ) + .await; + + assert_eq!(catalog.models, cached_models); + assert_eq!(endpoint.fetch_count(), 0); +} + +#[tokio::test] +async fn injected_cache_read_error_falls_back_and_persists_remote_models() { + let remote_models = vec![remote_model("remote", "Remote", /*priority*/ 0)]; + let cache = TestModelsCache::failing(/*load_error*/ true, /*store_error*/ false); + let endpoint = TestModelsEndpoint::new(vec![remote_models.clone()]); + let manager = OpenAiModelsManager::new_with_cache( + cache.clone(), + endpoint.clone(), + Some(AuthManager::from_auth_for_testing( + CodexAuth::create_dummy_chatgpt_auth_for_testing(), + )), + ); + + let catalog = manager + .raw_model_catalog( + RefreshStrategy::OnlineIfUncached, + DEFAULT_HTTP_CLIENT_FACTORY, + ) + .await; + + assert_eq!(catalog.models, remote_models); + assert_eq!(endpoint.fetch_count(), 1); + let stored_entries = cache.stored_entries(); + assert_eq!( + stored_entries, + vec![ModelsCacheEntry { + fetched_at: stored_entries[0].fetched_at, + etag: None, + client_version: Some(crate::client_version_to_whole()), + models: remote_models, + }] + ); +} + +#[tokio::test] +async fn injected_cache_write_error_does_not_fail_remote_refresh() { + let remote_models = vec![remote_model("remote", "Remote", /*priority*/ 0)]; + let cache = TestModelsCache::failing(/*load_error*/ false, /*store_error*/ true); + let endpoint = TestModelsEndpoint::new(vec![remote_models.clone()]); + let manager = OpenAiModelsManager::new_with_cache( + cache, + endpoint.clone(), + Some(AuthManager::from_auth_for_testing( + CodexAuth::create_dummy_chatgpt_auth_for_testing(), + )), + ); + + let catalog = manager + .raw_model_catalog( + RefreshStrategy::OnlineIfUncached, + DEFAULT_HTTP_CLIENT_FACTORY, + ) + .await; + + assert_eq!(catalog.models, remote_models); + assert_eq!(endpoint.fetch_count(), 1); +} + +#[tokio::test] +async fn injected_cache_ttl_refresh_preserves_cached_payload() { + let cached_models = vec![remote_model("cached", "Cached", /*priority*/ 0)]; + let cached_at = Utc::now() - chrono::Duration::minutes(1); + let cache = TestModelsCache::with_entry(ModelsCacheEntry { + fetched_at: cached_at, + etag: Some("cached-etag".to_string()), + client_version: Some(crate::client_version_to_whole()), + models: cached_models.clone(), + }); + let manager = OpenAiModelsManager::new_with_cache( + cache.clone(), + TestModelsEndpoint::new(Vec::new()), + Some(AuthManager::from_auth_for_testing(CodexAuth::from_api_key( + "test-api-key", + ))), + ); + + manager + .raw_model_catalog( + RefreshStrategy::OnlineIfUncached, + DEFAULT_HTTP_CLIENT_FACTORY, + ) + .await; + manager + .refresh_if_new_etag("cached-etag".to_string(), DEFAULT_HTTP_CLIENT_FACTORY) + .await; + + let stored_entries = cache.stored_entries(); + assert_eq!(stored_entries.len(), 1); + assert!(stored_entries[0].fetched_at > cached_at); + assert_eq!(stored_entries[0].etag.as_deref(), Some("cached-etag")); + assert_eq!( + stored_entries[0].client_version, + Some(crate::client_version_to_whole()) + ); + assert_eq!(stored_entries[0].models, cached_models); +} + async fn chatgpt_auth_tokens_for_tests(codex_home: &Path) -> CodexAuth { let auth_dot_json = codex_login::AuthDotJson { auth_mode: Some(AuthMode::ChatgptAuthTokens), @@ -753,15 +1066,10 @@ async fn refresh_available_models_refetches_when_cache_stale() { .expect("initial refresh succeeds"); // Rewrite cache with an old timestamp so it is treated as stale. - manager - .cache_manager - .as_ref() - .expect("cached model manager") - .manipulate_cache_for_test(|fetched_at| { - *fetched_at = Utc::now() - chrono::Duration::hours(1); - }) - .await - .expect("cache manipulation succeeds"); + mutate_file_cache_for_test(codex_home.path(), |cache| { + cache.fetched_at = Utc::now() - chrono::Duration::hours(1); + }) + .await; manager .refresh_available_models( @@ -794,16 +1102,11 @@ async fn refresh_available_models_refetches_when_version_mismatch() { .await .expect("initial refresh succeeds"); - manager - .cache_manager - .as_ref() - .expect("cached model manager") - .mutate_cache_for_test(|cache| { - let client_version = crate::client_version_to_whole(); - cache.client_version = Some(format!("{client_version}-mismatch")); - }) - .await - .expect("cache mutation succeeds"); + mutate_file_cache_for_test(codex_home.path(), |cache| { + let client_version = crate::client_version_to_whole(); + cache.client_version = Some(format!("{client_version}-mismatch")); + }) + .await; manager .refresh_available_models( @@ -834,12 +1137,16 @@ async fn refresh_available_models_drops_removed_remote_models() { /*priority*/ 1, )]; let endpoint = TestModelsEndpoint::new(vec![initial_models, refreshed_models]); - let mut manager = openai_manager_for_tests(codex_home.path().to_path_buf(), endpoint.clone()); - manager - .cache_manager - .as_mut() - .expect("cached model manager") - .set_ttl(Duration::ZERO); + let manager = OpenAiModelsManager::new_with_cache( + Arc::new(FileModelsCache::new( + codex_home.path().join(MODEL_CACHE_FILE), + Duration::ZERO, + )), + endpoint.clone(), + Some(AuthManager::from_auth_for_testing( + CodexAuth::create_dummy_chatgpt_auth_for_testing(), + )), + ); manager .refresh_available_models(