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

Scope model cache by provider and account

Browse Source
This commit is contained in:
won
2026-07-03 14:24:49 +09:00
parent da4c8ca57d
commit 07041404f5
9 changed files with 282 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/app-server/src/models_refresh_worker_tests.rs
+1
View File
@@ -76,6 +76,7 @@ async fn refreshes_immediately_periodically_and_stops_when_dropped() {
let endpoint = TestModelsEndpoint::new();
let models_manager: SharedModelsManager = Arc::new(OpenAiModelsManager::new(
codex_home.path().to_path_buf(),
"test-provider".to_string(),
endpoint.clone(),
/*auth_manager*/ None,
));
codex-rs/core/src/test_support.rs
+6 -1
View File
@@ -147,8 +147,13 @@ pub fn models_manager_with_provider(
auth_manager: Arc<AuthManager>,
provider: ModelProviderInfo,
) -> SharedModelsManager {
let model_provider_id = provider.name.clone();
let provider = create_model_provider(provider, Some(auth_manager));
provider.models_manager(codex_home, /*config_model_catalog*/ None)
provider.models_manager(
codex_home,
model_provider_id,
/*config_model_catalog*/ None,
)
}
pub fn get_model_offline(model: Option<&str>) -> String {
codex-rs/core/src/thread_manager.rs
+6 -1
View File
@@ -266,6 +266,7 @@ pub fn build_models_manager(
let provider = create_model_provider(config.model_provider.clone(), Some(auth_manager));
provider.models_manager(
config.codex_home.to_path_buf(),
config.model_provider_id.clone(),
config.model_catalog.clone(),
)
}
@@ -447,7 +448,11 @@ impl ThreadManager {
threads: Arc::new(RwLock::new(HashMap::new())),
thread_created_tx,
models_manager: create_model_provider(provider, Some(auth_manager.clone()))
.models_manager(codex_home, /*config_model_catalog*/ None),
.models_manager(
codex_home,
OPENAI_PROVIDER_ID.to_string(),
/*config_model_catalog*/ None,
),
environment_manager,
skills_service,
plugins_manager,
codex-rs/memories/write/src/startup_tests.rs
+2 -1
View File
@@ -627,10 +627,11 @@ impl ModelProvider for MockMemoryModelProvider {
fn models_manager(
&self,
codex_home: PathBuf,
model_provider_id: String,
config_model_catalog: Option<ModelsResponse>,
) -> codex_models_manager::manager::SharedModelsManager {
self.delegate
.models_manager(codex_home, config_model_catalog)
.models_manager(codex_home, model_provider_id, config_model_catalog)
}
}
codex-rs/model-provider/src/amazon_bedrock/mod.rs
+1
View File
@@ -164,6 +164,7 @@ impl ModelProvider for AmazonBedrockModelProvider {
fn models_manager(
&self,
_codex_home: PathBuf,
_model_provider_id: String,
config_model_catalog: Option<ModelsResponse>,
) -> SharedModelsManager {
Arc::new(StaticModelsManager::new(
codex-rs/model-provider/src/provider.rs
+14 -4
View File
@@ -197,6 +197,7 @@ pub trait ModelProvider: fmt::Debug + Send + Sync {
fn models_manager(
&self,
codex_home: PathBuf,
model_provider_id: String,
config_model_catalog: Option<ModelsResponse>,
) -> SharedModelsManager;
}
@@ -310,6 +311,7 @@ impl ModelProvider for ConfiguredModelProvider {
fn models_manager(
&self,
codex_home: PathBuf,
model_provider_id: String,
config_model_catalog: Option<ModelsResponse>,
) -> SharedModelsManager {
match config_model_catalog {
@@ -324,6 +326,7 @@ impl ModelProvider for ConfiguredModelProvider {
));
Arc::new(OpenAiModelsManager::new(
codex_home,
model_provider_id,
endpoint,
self.auth_manager.clone(),
))
@@ -649,8 +652,11 @@ mod tests {
ModelProviderInfo::create_amazon_bedrock_provider(/*aws*/ None),
/*auth_manager*/ None,
);
let manager =
provider.models_manager(test_codex_home(), /*config_model_catalog*/ None);
let manager = provider.models_manager(
test_codex_home(),
"amazon-bedrock".to_string(),
/*config_model_catalog*/ None,
);
let catalog = manager.raw_model_catalog(RefreshStrategy::Online).await;
let model_ids = catalog
@@ -697,6 +703,7 @@ mod tests {
);
let manager = provider.models_manager(
test_codex_home(),
"amazon-bedrock".to_string(),
Some(ModelsResponse {
models: vec![configured_model],
}),
@@ -742,8 +749,11 @@ mod tests {
)),
);
let manager =
provider.models_manager(test_codex_home(), /*config_model_catalog*/ None);
let manager = provider.models_manager(
test_codex_home(),
"test-provider".to_string(),
/*config_model_catalog*/ None,
);
let catalog = manager.raw_model_catalog(RefreshStrategy::Online).await;
assert!(
codex-rs/models-manager/src/cache.rs
+51 -2
View File
@@ -1,5 +1,6 @@
use chrono::DateTime;
use chrono::Utc;
use codex_protocol::auth::AuthMode;
use codex_protocol::openai_models::ModelInfo;
use serde::Deserialize;
use serde::Serialize;
@@ -28,7 +29,11 @@ impl ModelsCacheManager {
}
/// 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> {
pub(crate) async fn load_fresh(
&self,
expected_version: &str,
expected_cache_key: &ModelsCacheKey,
) -> Option<ModelsCache> {
info!(
cache_path = %self.cache_path.display(),
expected_version,
@@ -47,6 +52,13 @@ impl ModelsCacheManager {
fetched_at = %cache.fetched_at,
"models cache: loaded cache file"
);
if cache.cache_key.as_ref() != Some(expected_cache_key) {
info!(
cache_path = %self.cache_path.display(),
"models cache: cache key mismatch"
);
return None;
}
if cache.client_version.as_deref() != Some(expected_version) {
info!(
cache_path = %self.cache_path.display(),
@@ -79,11 +91,13 @@ impl ModelsCacheManager {
models: &[ModelInfo],
etag: Option<String>,
client_version: String,
cache_key: ModelsCacheKey,
) {
let cache = ModelsCache {
fetched_at: Utc::now(),
etag,
client_version: Some(client_version),
cache_key: Some(cache_key),
models: models.to_vec(),
};
if let Err(err) = self.save_internal(&cache).await {
@@ -92,11 +106,20 @@ impl ModelsCacheManager {
}
/// Renew the cache TTL by updating the fetched_at timestamp to now.
pub(crate) async fn renew_cache_ttl(&self) -> io::Result<()> {
pub(crate) async fn renew_cache_ttl(
&self,
expected_cache_key: &ModelsCacheKey,
) -> 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.cache_key.as_ref() != Some(expected_cache_key) {
return Err(io::Error::new(
ErrorKind::InvalidData,
"models cache key mismatch",
));
}
cache.fetched_at = Utc::now();
self.save_internal(&cache).await
}
@@ -157,6 +180,30 @@ impl ModelsCacheManager {
}
}
/// Identity of the provider and account that produced a model catalog.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub(crate) struct ModelsCacheKey {
provider_id: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
auth_mode: Option<AuthMode>,
#[serde(default, skip_serializing_if = "Option::is_none")]
account_id: Option<String>,
}
impl ModelsCacheKey {
pub(crate) fn new(
provider_id: String,
auth_mode: Option<AuthMode>,
account_id: Option<String>,
) -> Self {
Self {
provider_id,
auth_mode,
account_id,
}
}
}
/// Serialized snapshot of models and metadata cached on disk.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub(crate) struct ModelsCache {
@@ -165,6 +212,8 @@ pub(crate) struct ModelsCache {
pub(crate) etag: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub(crate) client_version: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub(crate) cache_key: Option<ModelsCacheKey>,
pub(crate) models: Vec<ModelInfo>,
}
codex-rs/models-manager/src/manager.rs
+30 -8
View File
@@ -1,3 +1,4 @@
use super::cache::ModelsCacheKey;
use super::cache::ModelsCacheManager;
use crate::collaboration_mode_presets::builtin_collaboration_mode_presets;
use crate::config::ModelsManagerConfig;
@@ -199,6 +200,7 @@ pub type SharedModelsManager = Arc<dyn ModelsManager>;
/// OpenAI-compatible model manager backed by bundled models, cache, and `/models`.
#[derive(Debug)]
pub struct OpenAiModelsManager {
model_provider_id: String,
remote_models: RwLock<Vec<ModelInfo>>,
etag: RwLock<Option<String>>,
cache_manager: ModelsCacheManager,
@@ -217,6 +219,7 @@ impl OpenAiModelsManager {
/// Construct an OpenAI-compatible remote model manager.
pub fn new(
codex_home: PathBuf,
model_provider_id: String,
endpoint_client: Arc<dyn ModelsEndpointClient>,
auth_manager: Option<Arc<AuthManager>>,
) -> Self {
@@ -224,6 +227,7 @@ impl OpenAiModelsManager {
let cache_manager = ModelsCacheManager::new(cache_path, DEFAULT_MODEL_CACHE_TTL);
let remote_models = load_remote_models_from_file().unwrap_or_default();
Self {
model_provider_id,
remote_models: RwLock::new(remote_models),
etag: RwLock::new(None),
cache_manager,
@@ -231,6 +235,19 @@ impl OpenAiModelsManager {
auth_manager,
}
}
fn cache_key(&self) -> ModelsCacheKey {
let auth_mode = self
.auth_manager
.as_ref()
.and_then(|auth_manager| auth_manager.auth_mode());
let account_id = auth_mode
.filter(|auth_mode| auth_mode.uses_codex_backend())
.and(self.auth_manager.as_ref())
.and_then(|auth_manager| auth_manager.auth_cached())
.and_then(|auth| auth.get_account_id());
ModelsCacheKey::new(self.model_provider_id.clone(), auth_mode, account_id)
}
}
impl StaticModelsManager {
@@ -288,10 +305,10 @@ impl OpenAiModelsManager {
async fn refresh_if_new_etag(&self, etag: String) {
let current_etag = self.get_etag().await;
if current_etag.clone().is_some() && current_etag.as_deref() == Some(etag.as_str()) {
if let Err(err) = self.cache_manager.renew_cache_ttl().await {
error!("failed to renew cache TTL: {err}");
match self.cache_manager.renew_cache_ttl(&self.cache_key()).await {
Ok(()) => return,
Err(err) => error!("failed to renew cache TTL: {err}"),
}
return;
}
if let Err(err) = self.refresh_available_models(RefreshStrategy::Online).await {
error!("failed to refresh available models: {err}");
@@ -334,11 +351,12 @@ impl OpenAiModelsManager {
async fn fetch_and_update_models(&self) -> CoreResult<()> {
let client_version = crate::client_version_to_whole();
let cache_key = self.cache_key();
let (models, etag) = self.endpoint_client.list_models(&client_version).await?;
self.apply_remote_models(models.clone()).await;
*self.etag.write().await = etag.clone();
self.cache_manager
.persist_cache(&models, etag, client_version)
.persist_cache(&models, etag, client_version, cache_key)
.await;
Ok(())
}
@@ -383,15 +401,19 @@ impl OpenAiModelsManager {
*self.remote_models.write().await = existing_models;
}
/// Attempt to satisfy the refresh from the cache when it matches the provider and TTL.
/// Attempt to satisfy the refresh from the cache when it matches the provider, account, and
/// TTL.
async fn try_load_cache(&self) -> bool {
let _timer =
codex_otel::start_global_timer("codex.remote_models.load_cache.duration_ms", &[]);
let client_version = crate::client_version_to_whole();
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 self.cache_manager.load_fresh(&client_version).await {
let cache_key = self.cache_key();
let cache = match self
.cache_manager
.load_fresh(&client_version, &cache_key)
.await
{
Some(cache) => cache,
None => {
info!("models cache: no usable cache entry");
codex-rs/models-manager/src/manager_tests.rs
+171 -5
View File
@@ -179,8 +179,9 @@ fn openai_manager_for_tests(
codex_home: std::path::PathBuf,
endpoint_client: Arc<dyn ModelsEndpointClient>,
) -> OpenAiModelsManager {
openai_manager_for_tests_with_auth(
openai_manager_for_tests_with_provider_and_auth(
codex_home,
"openai",
endpoint_client,
Some(AuthManager::from_auth_for_testing(
CodexAuth::create_dummy_chatgpt_auth_for_testing(),
@@ -193,14 +194,33 @@ fn openai_manager_for_tests_with_auth(
endpoint_client: Arc<dyn ModelsEndpointClient>,
auth_manager: Option<Arc<AuthManager>>,
) -> OpenAiModelsManager {
OpenAiModelsManager::new(codex_home, endpoint_client, auth_manager)
openai_manager_for_tests_with_provider_and_auth(
codex_home,
"openai",
endpoint_client,
auth_manager,
)
}
fn openai_manager_for_tests_with_provider_and_auth(
codex_home: std::path::PathBuf,
model_provider_id: &str,
endpoint_client: Arc<dyn ModelsEndpointClient>,
auth_manager: Option<Arc<AuthManager>>,
) -> OpenAiModelsManager {
OpenAiModelsManager::new(
codex_home,
model_provider_id.to_string(),
endpoint_client,
auth_manager,
)
}
fn static_manager_for_tests(model_catalog: ModelsResponse) -> StaticModelsManager {
StaticModelsManager::new(/*auth_manager*/ None, model_catalog)
}
async fn chatgpt_auth_tokens_for_tests(codex_home: &Path) -> CodexAuth {
async fn chatgpt_auth_tokens_for_tests(codex_home: &Path, account_id: &str) -> CodexAuth {
let auth_dot_json = codex_login::AuthDotJson {
auth_mode: Some(AuthMode::ChatgptAuthTokens),
openai_api_key: None,
@@ -213,7 +233,7 @@ c2ln",
.expect("fake id token should parse"),
access_token: "Access Token".to_string(),
refresh_token: "test".to_string(),
account_id: Some("account_id".to_string()),
account_id: Some(account_id.to_string()),
}),
last_refresh: Some(Utc::now()),
agent_identity: None,
@@ -529,6 +549,152 @@ async fn refresh_available_models_uses_cached_remote_only_catalog_for_chatgpt_au
);
}
#[tokio::test]
async fn offline_catalog_does_not_reuse_cache_written_by_another_provider() {
let auto_review_model = remote_model_with_visibility(
"codex-auto-review",
"Codex Auto Review",
/*priority*/ 1,
"hide",
);
let openai_models = vec![
remote_model("openai-visible", "OpenAI Visible", /*priority*/ 0),
auto_review_model.clone(),
];
let foreign_model = remote_model("foreign-visible", "Foreign Visible", /*priority*/ 0);
let codex_home = tempdir().expect("temp dir");
let openai_endpoint = TestModelsEndpoint::new(vec![openai_models]);
let openai_manager = openai_manager_for_tests(codex_home.path().to_path_buf(), openai_endpoint);
openai_manager
.refresh_available_models(RefreshStrategy::Online)
.await
.expect("OpenAI refresh succeeds");
let foreign_endpoint = TestModelsEndpoint::new(vec![vec![foreign_model.clone()]]);
let foreign_manager = openai_manager_for_tests_with_provider_and_auth(
codex_home.path().to_path_buf(),
"foreign",
foreign_endpoint,
Some(AuthManager::from_auth_for_testing(
CodexAuth::create_dummy_chatgpt_auth_for_testing(),
)),
);
foreign_manager
.refresh_available_models(RefreshStrategy::Online)
.await
.expect("foreign provider refresh succeeds");
let available = openai_manager.list_models(RefreshStrategy::Offline).await;
assert!(
available
.iter()
.any(|model| model.model == auto_review_model.slug),
"OpenAI's cached catalog should retain codex-auto-review"
);
assert!(
!available
.iter()
.any(|model| model.model == foreign_model.slug),
"OpenAI's cached catalog should not contain another provider's models"
);
}
#[tokio::test]
async fn offline_catalog_does_not_reuse_cache_written_by_another_account() {
let auto_review_model = remote_model_with_visibility(
"codex-auto-review",
"Codex Auto Review",
/*priority*/ 1,
"hide",
);
let first_account_models = vec![
remote_model("first-visible", "First Visible", /*priority*/ 0),
auto_review_model.clone(),
];
let second_account_model =
remote_model("second-visible", "Second Visible", /*priority*/ 0);
let cache_home = tempdir().expect("cache home");
let first_auth_home = tempdir().expect("first auth home");
let second_auth_home = tempdir().expect("second auth home");
let first_auth = chatgpt_auth_tokens_for_tests(first_auth_home.path(), "first-account").await;
let second_auth =
chatgpt_auth_tokens_for_tests(second_auth_home.path(), "second-account").await;
let first_endpoint = TestModelsEndpoint::new(vec![first_account_models]);
let first_manager = openai_manager_for_tests_with_provider_and_auth(
cache_home.path().to_path_buf(),
"openai",
first_endpoint,
Some(AuthManager::from_auth_for_testing(first_auth)),
);
first_manager
.refresh_available_models(RefreshStrategy::Online)
.await
.expect("first account refresh succeeds");
let second_endpoint = TestModelsEndpoint::new(vec![vec![second_account_model.clone()]]);
let second_manager = openai_manager_for_tests_with_provider_and_auth(
cache_home.path().to_path_buf(),
"openai",
second_endpoint,
Some(AuthManager::from_auth_for_testing(second_auth)),
);
second_manager
.refresh_available_models(RefreshStrategy::Online)
.await
.expect("second account refresh succeeds");
let available = first_manager.list_models(RefreshStrategy::Offline).await;
assert!(
available
.iter()
.any(|model| model.model == auto_review_model.slug),
"the first account's catalog should retain codex-auto-review"
);
assert!(
!available
.iter()
.any(|model| model.model == second_account_model.slug),
"the first account's catalog should not contain another account's models"
);
}
#[tokio::test]
async fn matching_etag_refetches_when_cache_key_mismatches() {
let initial_model = remote_model("initial-model", "Initial Model", /*priority*/ 0);
let refreshed_model = remote_model("refreshed-model", "Refreshed Model", /*priority*/ 0);
let codex_home = tempdir().expect("temp dir");
let endpoint =
TestModelsEndpoint::new(vec![vec![initial_model], vec![refreshed_model.clone()]]);
let manager = openai_manager_for_tests(codex_home.path().to_path_buf(), endpoint.clone());
manager
.refresh_available_models(RefreshStrategy::Online)
.await
.expect("initial refresh succeeds");
*manager.etag.write().await = Some("same-etag".to_string());
manager
.cache_manager
.mutate_cache_for_test(|cache| cache.cache_key = None)
.await
.expect("cache key should be removed");
manager.refresh_if_new_etag("same-etag".to_string()).await;
assert_eq!(endpoint.fetch_count(), 2, "cache mismatch should refetch");
assert!(
manager
.get_remote_models()
.await
.iter()
.any(|model| model.slug == refreshed_model.slug),
"refetched catalog should be applied"
);
}
#[tokio::test]
async fn get_model_info_uses_fallback_for_bundled_models_when_chatgpt_remote_is_authoritative() {
let remote_models = vec![remote_model(
@@ -951,7 +1117,7 @@ async fn refresh_available_models_fetches_with_chatgpt_auth_tokens() {
"ChatGPT Auth Tokens",
/*priority*/ 1,
)]]);
let auth = chatgpt_auth_tokens_for_tests(codex_home.path()).await;
let auth = chatgpt_auth_tokens_for_tests(codex_home.path(), "account_id").await;
let manager = openai_manager_for_tests_with_auth(
codex_home.path().to_path_buf(),
endpoint.clone(),