openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-18 17:18:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

Allow injecting model catalog caches (#36992)

Browse Source 
## What changed

- Add the public asynchronous `ModelsCache` contract and cache entry/error types.
- Let model providers and `OpenAiModelsManager` accept caller-provided cache implementations while retaining the existing file-backed cache by default.
- Treat cache read failures as misses, keep cache write failures non-fatal, and support refreshing an entry's TTL after ETag revalidation.
- Continue to prefer an explicitly configured static model catalog over an injected cache.

## Testing

- Cover injected cache hits, read and write failures, TTL refreshes, and end-to-end model selection.
- Verify that the file-backed implementation satisfies the new cache contract and does not serve expired entries before revalidation.

GitOrigin-RevId: eef957a317910b8502a0c20f8fc112f689d5e389
This commit is contained in:
Rasmus Rygaard
2026-08-04 23:04:11 +00:00
committed by copyberry
parent 11e390bb10
commit 44cb66e4ed
8 changed files with 903 additions and 203 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/tests/common/test_codex.rs
+13 -1
View File
@@ -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<Arc<dyn TimeProvider>>,
code_mode_host_program: Option<PathBuf>,
history_mode: Option<ThreadHistoryMode>,
models_manager: Option<SharedModelsManager>,
}
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,
}
}
codex-rs/core/tests/suite/injected_models_cache.rs
+256
View File
@@ -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<ModelsCacheEntry>,
load_error: bool,
stored_entries: Mutex<Vec<ModelsCacheEntry>>,
}
impl ModelsCache for TestModelsCache {
fn load<'a>(
&'a self,
_client_version: &'a str,
) -> ModelsCacheFuture<'a, Result<Option<ModelsCacheEntry>, 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<ModelInfo>,
fetch_count: AtomicUsize,
}
impl TestModelsEndpoint {
fn new(models: Vec<ModelInfo>) -> Arc<Self> {
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<ModelInfo>, Option<String>)>> {
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<dyn ModelsCache>,
endpoint: Arc<TestModelsEndpoint>,
) -> 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(())
}
codex-rs/core/tests/suite/mod.rs
+1
View File
@@ -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;
codex-rs/model-provider/src/provider.rs
+40
View File
@@ -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<ModelsResponse>,
cache: Arc<dyn ModelsCache>,
) -> SharedModelsManager {
drop(cache);
self.models_manager_without_cache(config_model_catalog)
}
}
pub type ModelProviderFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
@@ -401,6 +417,30 @@ impl ModelProvider for ConfiguredModelProvider {
}
}
}
fn models_manager_with_cache(
&self,
config_model_catalog: Option<ModelsResponse>,
cache: Arc<dyn ModelsCache>,
) -> 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)]
codex-rs/models-manager/src/cache.rs
+203 -155
View File
@@ -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<Option<ModelsCacheEntry>, 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<Box<dyn Future<Output = T> + 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<ModelsCache> {
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<String>,
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<Option<ModelsCache>> {
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<F>(&self, f: F) -> io::Result<()>
where
F: FnOnce(&mut DateTime<Utc>),
{
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<F>(&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<Utc>,
/// 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<Utc>,
/// Endpoint ETag used to revalidate the catalog without downloading an unchanged payload.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub(crate) etag: Option<String>,
pub etag: Option<String>,
/// 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<String>,
pub client_version: Option<String>,
/// Models returned by the catalog endpoint.
#[serde(
deserialize_with = "codex_protocol::openai_models::deserialize_model_infos_with_legacy_base"
)]
pub(crate) models: Vec<ModelInfo>,
pub models: Vec<ModelInfo>,
}
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<String>) -> 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<Option<ModelsCacheEntry>, 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<Option<ModelsCacheEntry>, 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<Option<ModelsCacheEntry>> {
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())
}
codex-rs/models-manager/src/lib.rs
+1 -1
View File
@@ -1,4 +1,4 @@
pub(crate) mod cache;
pub mod cache;
pub mod collaboration_mode_presets;
pub(crate) mod config;
pub mod manager;
codex-rs/models-manager/src/manager.rs
+57 -21
View File
@@ -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<dyn ModelsManager>;
pub struct OpenAiModelsManager {
remote_models: RwLock<Vec<ModelInfo>>,
etag: RwLock<Option<String>>,
cache_manager: Option<ModelsCacheManager>,
cache: Option<Arc<dyn ModelsCache>>,
endpoint_client: SharedModelsEndpointClient,
auth_manager: Option<Arc<AuthManager>>,
}
@@ -235,8 +238,11 @@ impl OpenAiModelsManager {
auth_manager: Option<Arc<AuthManager>>,
) -> 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<dyn ModelsEndpointClient>,
auth_manager: Option<Arc<AuthManager>>,
) -> 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<ModelsCacheManager>,
/// 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<dyn ModelsCache>,
endpoint_client: Arc<dyn ModelsEndpointClient>,
auth_manager: Option<Arc<AuthManager>>,
) -> Self {
Self::new_with_optional_cache(Some(cache), endpoint_client, auth_manager)
}
fn new_with_optional_cache(
cache: Option<Arc<dyn ModelsCache>>,
endpoint_client: Arc<dyn ModelsEndpointClient>,
auth_manager: Option<Arc<AuthManager>>,
) -> 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
codex-rs/models-manager/src/manager_tests.rs
+332 -25
View File
@@ -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<Option<OutboundProxyPolicy>>,
}
#[derive(Debug)]
struct TestModelsCache {
entry: Mutex<Option<ModelsCacheEntry>>,
load_error: bool,
store_error: bool,
stored_entries: Mutex<Vec<ModelsCacheEntry>>,
}
impl TestModelsCache {
fn with_entry(entry: ModelsCacheEntry) -> Arc<Self> {
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<Self> {
Arc::new(Self {
entry: Mutex::new(None),
load_error,
store_error,
stored_entries: Mutex::new(Vec::new()),
})
}
fn stored_entries(&self) -> Vec<ModelsCacheEntry> {
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<Option<ModelsCacheEntry>, 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<Vec<ModelInfo>>) -> Arc<Self> {
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<F>(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(