From d2dcd69943aeb66fa203062d1b865d67d7698a96 Mon Sep 17 00:00:00 2001 From: Richard Lee Date: Thu, 9 Jul 2026 18:20:21 -0700 Subject: [PATCH] [codex] keep cloud models cache in memory --- codex-rs/core/src/thread_manager.rs | 3 ++ codex-rs/core/src/thread_manager_tests.rs | 38 ++++++++++++++++++++ codex-rs/models-manager/src/cache.rs | 32 ++++++++--------- codex-rs/models-manager/src/manager_tests.rs | 8 ++++- codex-rs/protocol/src/protocol.rs | 11 ++++++ 5 files changed, 74 insertions(+), 18 deletions(-) diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 94a05e3fd3..2351c42b45 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -406,6 +406,9 @@ impl ThreadManager { Some(Arc::clone(&auth_manager)), ); let models_manager = match local_runtime_paths.as_ref() { + Some(_) if session_source.is_codex_cloud_agent() => { + model_provider.models_manager_without_disk_cache(config.model_catalog.clone()) + } Some(paths) => model_provider .models_manager(paths.codex_home.to_path_buf(), config.model_catalog.clone()), None => model_provider.models_manager_without_disk_cache(config.model_catalog.clone()), diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index b9e1e94593..0e3f336f66 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -21,6 +21,7 @@ use codex_protocol::models::ReasoningItemReasoningSummary; use codex_protocol::models::ResponseItem; use codex_protocol::openai_models::ModelsResponse; use codex_protocol::protocol::AgentMessageEvent; +use codex_protocol::protocol::CODEX_CLOUD_AGENT_SESSION_SOURCE; use codex_protocol::protocol::InitialHistory; use codex_protocol::protocol::InternalSessionSource; use codex_protocol::protocol::ResumedHistory; @@ -1845,6 +1846,43 @@ async fn new_uses_active_provider_for_model_refresh() { assert_eq!(models_mock.requests().len(), 1); } +#[tokio::test] +async fn cloud_agent_uses_memory_only_models_cache_with_local_runtime_paths() { + let server = MockServer::start().await; + let models_mock = mount_models_once(&server, ModelsResponse { models: vec![] }).await; + + let temp_dir = tempdir().expect("tempdir"); + let mut config = test_config().await; + config.codex_home = temp_dir.path().join("codex-home").abs(); + config.cwd = config.codex_home.abs(); + std::fs::create_dir_all(&config.codex_home).expect("create codex home"); + config.model_catalog = None; + config.model_provider.base_url = Some(server.uri()); + + let auth_manager = + AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); + let manager = ThreadManager::new( + &config, + auth_manager, + SessionSource::Custom(CODEX_CLOUD_AGENT_SESSION_SOURCE.to_string()), + Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), + empty_extension_registry(), + Arc::new(crate::test_support::EmptyUserInstructionsProvider), + /*analytics_events_client*/ None, + thread_store_from_config(&config, /*state_db*/ None), + /*agent_graph_store*/ None, + Some(TEST_INSTALLATION_ID.to_string()), + /*attestation_provider*/ None, + /*external_time_provider*/ None, + ); + + let _ = manager.list_models(RefreshStrategy::OnlineIfUncached).await; + let _ = manager.list_models(RefreshStrategy::OnlineIfUncached).await; + + assert_eq!(models_mock.requests().len(), 1); + assert!(!config.codex_home.join("models_cache.json").exists()); +} + #[test] fn interrupted_fork_snapshot_appends_interrupt_boundary() { let committed_history = diff --git a/codex-rs/models-manager/src/cache.rs b/codex-rs/models-manager/src/cache.rs index a9d1e62f1a..60ca68e945 100644 --- a/codex-rs/models-manager/src/cache.rs +++ b/codex-rs/models-manager/src/cache.rs @@ -8,13 +8,15 @@ use std::io::ErrorKind; use std::path::PathBuf; use std::time::Duration; use tokio::fs; +use tokio::sync::RwLock; use tracing::error; use tracing::info; -/// Manages loading and saving of models cache to disk. +/// Manages loading and saving model cache state in memory or on disk. #[derive(Debug)] pub(crate) struct ModelsCacheManager { cache_path: Option, + memory_cache: RwLock>, cache_ttl: Duration, } @@ -23,6 +25,7 @@ impl ModelsCacheManager { pub(crate) fn new(cache_path: PathBuf, cache_ttl: Duration) -> Self { Self { cache_path: Some(cache_path), + memory_cache: RwLock::new(None), cache_ttl, } } @@ -31,16 +34,16 @@ impl ModelsCacheManager { pub(crate) fn without_disk_cache(cache_ttl: Duration) -> Self { Self { cache_path: 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 { - let cache_path = self.cache_path.as_ref()?; info!( - cache_path = %cache_path.display(), - expected_version, + cache_path = ?self.cache_path.as_ref(), + expected_version, "models cache: attempting load_fresh" ); let cache = match self.load().await { @@ -51,14 +54,14 @@ impl ModelsCacheManager { } }; info!( - cache_path = %cache_path.display(), + cache_path = ?self.cache_path.as_ref(), cached_version = ?cache.client_version, fetched_at = %cache.fetched_at, - "models cache: loaded cache file" + "models cache: loaded cache entry" ); if cache.client_version.as_deref() != Some(expected_version) { info!( - cache_path = %cache_path.display(), + cache_path = ?self.cache_path.as_ref(), expected_version, cached_version = ?cache.client_version, "models cache: cache version mismatch" @@ -67,7 +70,7 @@ impl ModelsCacheManager { } if !cache.is_fresh(self.cache_ttl) { info!( - cache_path = %cache_path.display(), + cache_path = ?self.cache_path.as_ref(), cache_ttl_secs = self.cache_ttl.as_secs(), fetched_at = %cache.fetched_at, "models cache: cache is stale" @@ -75,23 +78,20 @@ impl ModelsCacheManager { return None; } info!( - cache_path = %cache_path.display(), + cache_path = ?self.cache_path.as_ref(), cache_ttl_secs = self.cache_ttl.as_secs(), "models cache: cache hit" ); Some(cache) } - /// Persist the cache to disk, creating parent directories as needed. + /// 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() { - return; - } let cache = ModelsCache { fetched_at: Utc::now(), etag, @@ -105,9 +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() { - return Ok(()); - } let mut cache = match self.load().await? { Some(cache) => cache, None => return Err(io::Error::new(ErrorKind::NotFound, "cache not found")), @@ -118,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) => { @@ -133,6 +130,7 @@ impl ModelsCacheManager { async fn save_internal(&self, cache: &ModelsCache) -> io::Result<()> { let Some(cache_path) = self.cache_path.as_ref() else { + *self.memory_cache.write().await = Some(cache.clone()); return Ok(()); }; if let Some(parent) = cache_path.parent() { diff --git a/codex-rs/models-manager/src/manager_tests.rs b/codex-rs/models-manager/src/manager_tests.rs index 51cbc66526..7347b3b139 100644 --- a/codex-rs/models-manager/src/manager_tests.rs +++ b/codex-rs/models-manager/src/manager_tests.rs @@ -221,9 +221,15 @@ async fn manager_without_disk_cache_fetches_and_retains_models_in_memory() { )), ); - let catalog = manager.raw_model_catalog(RefreshStrategy::Online).await; + let catalog = manager + .raw_model_catalog(RefreshStrategy::OnlineIfUncached) + .await; + let cached_catalog = manager + .raw_model_catalog(RefreshStrategy::OnlineIfUncached) + .await; assert_eq!(catalog.models, remote_models); + assert_eq!(cached_catalog, catalog); assert_eq!(manager.get_remote_models().await, remote_models); assert_eq!(endpoint.fetch_count(), 1); } diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index 57a36c3312..359c8a701e 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -2705,6 +2705,9 @@ fn session_cwd_from_items(items: &[RolloutItem]) -> Option { }) } +/// Custom session-source name used by Codex Cloud Agent runtimes. +pub const CODEX_CLOUD_AGENT_SESSION_SOURCE: &str = "codex-cloud-agent"; + #[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, JsonSchema, TS, Default)] #[serde(rename_all = "lowercase")] #[ts(rename_all = "lowercase")] @@ -2840,6 +2843,14 @@ impl SessionSource { matches!(self, SessionSource::Internal(_)) } + /// Returns whether this is a Codex Cloud Agent runtime's base session source. + pub fn is_codex_cloud_agent(&self) -> bool { + matches!( + self, + SessionSource::Custom(source) if source == CODEX_CLOUD_AGENT_SESSION_SOURCE + ) + } + pub fn is_non_root_agent(&self) -> bool { matches!( self,