diff --git a/codex-rs/models-manager/src/cache.rs b/codex-rs/models-manager/src/cache.rs index 8e2c82db81..82ec68b858 100644 --- a/codex-rs/models-manager/src/cache.rs +++ b/codex-rs/models-manager/src/cache.rs @@ -7,31 +7,25 @@ use std::io; use std::io::ErrorKind; use std::path::PathBuf; use std::time::Duration; -use std::time::Instant; use tokio::fs; use tokio::sync::RwLock; use tracing::error; use tracing::info; -/// Tracks in-process freshness or loads and saves model cache snapshots on disk. +/// Manages loading and saving model cache state in memory or on disk. #[derive(Debug)] pub(crate) struct ModelsCacheManager { cache_path: Option, - memory_cache_fetched_at: RwLock>, + memory_cache: RwLock>, cache_ttl: Duration, } -pub(crate) enum ModelsCacheHit { - InMemory, - Persisted(ModelsCache), -} - 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: Some(cache_path), - memory_cache_fetched_at: RwLock::new(None), + memory_cache: RwLock::new(None), cache_ttl, } } @@ -40,31 +34,18 @@ impl ModelsCacheManager { pub(crate) fn without_disk_cache(cache_ttl: Duration) -> Self { Self { cache_path: None, - memory_cache_fetched_at: RwLock::new(None), + memory_cache: RwLock::new(None), cache_ttl, } } /// 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 { + pub(crate) async fn load_fresh(&self, expected_version: &str) -> Option { info!( cache_path = ?self.cache_path.as_ref(), expected_version, "models cache: attempting load_fresh" ); - if self.cache_path.is_none() { - // An in-memory entry cannot outlive this process or its fixed client version. The - // model manager already owns the corresponding models and ETag. - let fetched_at = *self.memory_cache_fetched_at.read().await; - let is_fresh = fetched_at.is_some_and(|fetched_at| { - !self.cache_ttl.is_zero() && fetched_at.elapsed() <= self.cache_ttl - }); - info!( - cache_ttl_secs = self.cache_ttl.as_secs(), - is_fresh, "models cache: checked in-memory freshness" - ); - return is_fresh.then_some(ModelsCacheHit::InMemory); - } let cache = match self.load().await { Ok(cache) => cache?, Err(err) => { @@ -101,20 +82,16 @@ impl ModelsCacheManager { cache_ttl_secs = self.cache_ttl.as_secs(), "models cache: cache hit" ); - Some(ModelsCacheHit::Persisted(cache)) + Some(cache) } - /// Record fresh cache state, serializing it when disk persistence is configured. + /// Persist the cache to its configured storage, creating disk directories as needed. pub(crate) async fn persist_cache( &self, models: &[ModelInfo], etag: Option, client_version: String, ) { - if self.cache_path.is_none() { - *self.memory_cache_fetched_at.write().await = Some(Instant::now()); - return; - } let cache = ModelsCache { fetched_at: Utc::now(), etag, @@ -128,14 +105,6 @@ impl ModelsCacheManager { /// Renew the cache TTL by updating the fetched_at timestamp to now. pub(crate) async fn renew_cache_ttl(&self) -> io::Result<()> { - if self.cache_path.is_none() { - let mut fetched_at = self.memory_cache_fetched_at.write().await; - if fetched_at.is_none() { - return Err(io::Error::new(ErrorKind::NotFound, "cache not found")); - } - *fetched_at = Some(Instant::now()); - return Ok(()); - } let mut cache = match self.load().await? { Some(cache) => cache, None => return Err(io::Error::new(ErrorKind::NotFound, "cache not found")), @@ -146,7 +115,7 @@ impl ModelsCacheManager { async fn load(&self) -> io::Result> { let Some(cache_path) = self.cache_path.as_ref() else { - return Ok(None); + return Ok(self.memory_cache.read().await.clone()); }; match fs::read(cache_path).await { Ok(contents) => { @@ -161,10 +130,8 @@ impl ModelsCacheManager { async fn save_internal(&self, cache: &ModelsCache) -> io::Result<()> { let Some(cache_path) = self.cache_path.as_ref() else { - return Err(io::Error::new( - ErrorKind::Unsupported, - "memory-only cache does not serialize model data", - )); + *self.memory_cache.write().await = Some(cache.clone()); + return Ok(()); }; if let Some(parent) = cache_path.parent() { fs::create_dir_all(parent).await?; @@ -209,7 +176,7 @@ impl ModelsCacheManager { } } -/// Serialized snapshot of models and metadata cached on disk. +/// Snapshot of models and metadata retained in memory or serialized to disk. #[derive(Debug, Clone, Serialize, Deserialize)] pub(crate) struct ModelsCache { pub(crate) fetched_at: DateTime, diff --git a/codex-rs/models-manager/src/manager.rs b/codex-rs/models-manager/src/manager.rs index 17074362b0..6b949f2b39 100644 --- a/codex-rs/models-manager/src/manager.rs +++ b/codex-rs/models-manager/src/manager.rs @@ -1,4 +1,3 @@ -use super::cache::ModelsCacheHit; use super::cache::ModelsCacheManager; use crate::collaboration_mode_presets::builtin_collaboration_mode_presets; use crate::config::ModelsManagerConfig; @@ -7,7 +6,6 @@ use codex_http_client::HttpClientFactory; use codex_login::AuthManager; use codex_protocol::auth::AuthMode; use codex_protocol::config_types::CollaborationModeMask; -use codex_protocol::error::CodexErr; use codex_protocol::error::Result as CoreResult; use codex_protocol::openai_models::ModelInfo; use codex_protocol::openai_models::ModelPreset; @@ -20,7 +18,6 @@ use std::pin::Pin; use std::sync::Arc; use std::time::Duration; use tokio::sync::RwLock; -use tokio::sync::Semaphore; use tokio::sync::TryLockError; use tracing::Instrument as _; use tracing::error; @@ -218,7 +215,6 @@ pub type SharedModelsManager = Arc; pub struct OpenAiModelsManager { remote_models: RwLock>, etag: RwLock>, - refresh_permit: Semaphore, cache_manager: ModelsCacheManager, endpoint_client: SharedModelsEndpointClient, auth_manager: Option>, @@ -267,7 +263,6 @@ impl OpenAiModelsManager { Self { remote_models: RwLock::new(remote_models), etag: RwLock::new(None), - refresh_permit: Semaphore::new(1), cache_manager, endpoint_client, auth_manager, @@ -380,7 +375,7 @@ impl OpenAiModelsManager { RefreshStrategy::Offline => { // Only try to load from cache, never fetch self.try_load_cache().await; - return Ok(()); + Ok(()) } RefreshStrategy::OnlineIfUncached => { // Try cache first, fall back to online if unavailable @@ -388,24 +383,14 @@ impl OpenAiModelsManager { info!("models cache: using cached models for OnlineIfUncached"); return Ok(()); } + info!("models cache: cache miss, fetching remote models"); + self.fetch_and_update_models(http_client_factory).await + } + RefreshStrategy::Online => { + // Always fetch from network + self.fetch_and_update_models(http_client_factory).await } - RefreshStrategy::Online => {} } - - // Serialize network refreshes so models, ETag, and freshness come from one response. - let _refresh_permit = self - .refresh_permit - .acquire() - .await - .map_err(|_| CodexErr::InternalServerError)?; - if matches!(refresh_strategy, RefreshStrategy::OnlineIfUncached) - && self.try_load_cache().await - { - info!("models cache: another refresh populated the cache"); - return Ok(()); - } - info!("models cache: fetching remote models"); - self.fetch_and_update_models(http_client_factory).await } async fn fetch_and_update_models( @@ -480,10 +465,6 @@ impl OpenAiModelsManager { return false; } }; - let ModelsCacheHit::Persisted(cache) = cache else { - info!("models cache: using current in-memory models"); - return true; - }; let models = cache.models.clone(); *self.etag.write().await = cache.etag.clone(); self.apply_remote_models(models.clone()).await; diff --git a/codex-rs/models-manager/src/manager_tests.rs b/codex-rs/models-manager/src/manager_tests.rs index 6bcdba6283..20ba8ff207 100644 --- a/codex-rs/models-manager/src/manager_tests.rs +++ b/codex-rs/models-manager/src/manager_tests.rs @@ -21,7 +21,6 @@ use std::sync::Mutex; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use tempfile::tempdir; -use tokio::sync::Notify; #[path = "model_info_overrides_tests.rs"] mod model_info_overrides_tests; @@ -83,13 +82,6 @@ struct TestModelsEndpoint { responses: Mutex>>, fetch_count: AtomicUsize, observed_proxy_policy: Mutex>, - first_fetch_gate: Option>, -} - -#[derive(Debug)] -struct FirstFetchGate { - started: Notify, - release: Notify, } impl TestModelsEndpoint { @@ -100,7 +92,6 @@ impl TestModelsEndpoint { responses: Mutex::new(responses.into()), fetch_count: AtomicUsize::new(0), observed_proxy_policy: Mutex::new(None), - first_fetch_gate: None, }) } @@ -111,28 +102,9 @@ impl TestModelsEndpoint { responses: Mutex::new(responses.into()), fetch_count: AtomicUsize::new(0), observed_proxy_policy: Mutex::new(None), - first_fetch_gate: None, }) } - fn with_blocked_first_fetch( - responses: Vec>, - ) -> (Arc, Arc) { - let gate = Arc::new(FirstFetchGate { - started: Notify::new(), - release: Notify::new(), - }); - let endpoint = Arc::new(Self { - has_command_auth: false, - uses_codex_backend: true, - responses: Mutex::new(responses.into()), - fetch_count: AtomicUsize::new(0), - observed_proxy_policy: Mutex::new(None), - first_fetch_gate: Some(Arc::clone(&gate)), - }); - (endpoint, gate) - } - fn fetch_count(&self) -> usize { self.fetch_count.load(Ordering::SeqCst) } @@ -145,13 +117,7 @@ impl TestModelsEndpoint { } async fn list_models(&self) -> CoreResult<(Vec, Option)> { - let fetch_index = self.fetch_count.fetch_add(1, Ordering::SeqCst); - if fetch_index == 0 - && let Some(gate) = self.first_fetch_gate.as_ref() - { - gate.started.notify_one(); - gate.release.notified().await; - } + self.fetch_count.fetch_add(1, Ordering::SeqCst); let models = self .responses .lock() @@ -274,96 +240,6 @@ async fn manager_without_disk_cache_fetches_and_retains_models_in_memory() { assert_eq!(endpoint.fetch_count(), 1); } -#[tokio::test] -async fn manager_without_disk_cache_refetches_when_stale() { - let remote_models = vec![remote_model("remote", "Remote", /*priority*/ 0)]; - let endpoint = TestModelsEndpoint::new(vec![remote_models.clone(), remote_models]); - let mut manager = OpenAiModelsManager::new_without_disk_cache( - endpoint.clone(), - Some(AuthManager::from_auth_for_testing( - CodexAuth::create_dummy_chatgpt_auth_for_testing(), - )), - ); - manager.cache_manager.set_ttl(Duration::ZERO); - - for _ in 0..2 { - let _ = manager - .raw_model_catalog( - RefreshStrategy::OnlineIfUncached, - DEFAULT_HTTP_CLIENT_FACTORY, - ) - .await; - } - - assert_eq!(endpoint.fetch_count(), 2); -} - -#[tokio::test] -async fn manager_without_disk_cache_online_always_refetches() { - let remote_models = vec![remote_model("remote", "Remote", /*priority*/ 0)]; - let endpoint = TestModelsEndpoint::new(vec![remote_models.clone(), remote_models]); - let manager = OpenAiModelsManager::new_without_disk_cache( - endpoint.clone(), - Some(AuthManager::from_auth_for_testing( - CodexAuth::create_dummy_chatgpt_auth_for_testing(), - )), - ); - - for _ in 0..2 { - let _ = manager - .raw_model_catalog(RefreshStrategy::Online, DEFAULT_HTTP_CLIENT_FACTORY) - .await; - } - - assert_eq!(endpoint.fetch_count(), 2); -} - -#[tokio::test] -async fn concurrent_memory_cache_misses_share_one_fetch() { - let remote_models = vec![remote_model("remote", "Remote", /*priority*/ 0)]; - let (endpoint, gate) = - TestModelsEndpoint::with_blocked_first_fetch(vec![remote_models.clone(), remote_models]); - let manager = Arc::new(OpenAiModelsManager::new_without_disk_cache( - endpoint.clone(), - Some(AuthManager::from_auth_for_testing( - CodexAuth::create_dummy_chatgpt_auth_for_testing(), - )), - )); - - let first_refresh = tokio::spawn({ - let manager = Arc::clone(&manager); - async move { - manager - .raw_model_catalog( - RefreshStrategy::OnlineIfUncached, - DEFAULT_HTTP_CLIENT_FACTORY, - ) - .await - } - }); - gate.started.notified().await; - let second_refresh = tokio::spawn({ - let manager = Arc::clone(&manager); - async move { - manager - .raw_model_catalog( - RefreshStrategy::OnlineIfUncached, - DEFAULT_HTTP_CLIENT_FACTORY, - ) - .await - } - }); - tokio::task::yield_now().await; - gate.release.notify_one(); - - let (first_catalog, second_catalog) = tokio::join!(first_refresh, second_refresh); - assert_eq!( - first_catalog.expect("first refresh should complete"), - second_catalog.expect("second refresh should complete") - ); - assert_eq!(endpoint.fetch_count(), 1); -} - async fn chatgpt_auth_tokens_for_tests(codex_home: &Path) -> CodexAuth { let auth_dot_json = codex_login::AuthDotJson { auth_mode: Some(AuthMode::ChatgptAuthTokens), @@ -800,7 +676,6 @@ async fn refresh_available_models_keeps_merging_for_api_auth() { responses: Mutex::new(vec![remote_models.clone()].into()), fetch_count: AtomicUsize::new(0), observed_proxy_policy: Mutex::new(None), - first_fetch_gate: None, }); let manager = openai_manager_for_tests_with_auth( codex_home.path().to_path_buf(),