diff --git a/codex-rs/app-server/src/models_refresh_worker_tests.rs b/codex-rs/app-server/src/models_refresh_worker_tests.rs index 4d7ff48d31..4a1a181f4e 100644 --- a/codex-rs/app-server/src/models_refresh_worker_tests.rs +++ b/codex-rs/app-server/src/models_refresh_worker_tests.rs @@ -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, )); diff --git a/codex-rs/core/src/test_support.rs b/codex-rs/core/src/test_support.rs index 1db714fada..6b27ca462d 100644 --- a/codex-rs/core/src/test_support.rs +++ b/codex-rs/core/src/test_support.rs @@ -147,8 +147,13 @@ pub fn models_manager_with_provider( auth_manager: Arc, 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 { diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 4ea86cb33d..512f6cfdb3 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -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, diff --git a/codex-rs/memories/write/src/startup_tests.rs b/codex-rs/memories/write/src/startup_tests.rs index fad1ce14a4..54c57a1e11 100644 --- a/codex-rs/memories/write/src/startup_tests.rs +++ b/codex-rs/memories/write/src/startup_tests.rs @@ -627,10 +627,11 @@ impl ModelProvider for MockMemoryModelProvider { fn models_manager( &self, codex_home: PathBuf, + model_provider_id: String, config_model_catalog: Option, ) -> codex_models_manager::manager::SharedModelsManager { self.delegate - .models_manager(codex_home, config_model_catalog) + .models_manager(codex_home, model_provider_id, config_model_catalog) } } diff --git a/codex-rs/model-provider/src/amazon_bedrock/mod.rs b/codex-rs/model-provider/src/amazon_bedrock/mod.rs index 108308c1bf..6bfe8d77a4 100644 --- a/codex-rs/model-provider/src/amazon_bedrock/mod.rs +++ b/codex-rs/model-provider/src/amazon_bedrock/mod.rs @@ -164,6 +164,7 @@ impl ModelProvider for AmazonBedrockModelProvider { fn models_manager( &self, _codex_home: PathBuf, + _model_provider_id: String, config_model_catalog: Option, ) -> SharedModelsManager { Arc::new(StaticModelsManager::new( diff --git a/codex-rs/model-provider/src/provider.rs b/codex-rs/model-provider/src/provider.rs index 5f101e52cd..b7501c6598 100644 --- a/codex-rs/model-provider/src/provider.rs +++ b/codex-rs/model-provider/src/provider.rs @@ -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, ) -> SharedModelsManager; } @@ -310,6 +311,7 @@ impl ModelProvider for ConfiguredModelProvider { fn models_manager( &self, codex_home: PathBuf, + model_provider_id: String, config_model_catalog: Option, ) -> 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!( diff --git a/codex-rs/models-manager/src/cache.rs b/codex-rs/models-manager/src/cache.rs index 812716d29c..e5749c8e7f 100644 --- a/codex-rs/models-manager/src/cache.rs +++ b/codex-rs/models-manager/src/cache.rs @@ -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 { + pub(crate) async fn load_fresh( + &self, + expected_version: &str, + expected_cache_key: &ModelsCacheKey, + ) -> Option { 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, 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, + #[serde(default, skip_serializing_if = "Option::is_none")] + account_id: Option, +} + +impl ModelsCacheKey { + pub(crate) fn new( + provider_id: String, + auth_mode: Option, + account_id: Option, + ) -> 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, #[serde(default, skip_serializing_if = "Option::is_none")] pub(crate) client_version: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub(crate) cache_key: Option, pub(crate) models: Vec, } diff --git a/codex-rs/models-manager/src/manager.rs b/codex-rs/models-manager/src/manager.rs index 67b939051b..065c1c162c 100644 --- a/codex-rs/models-manager/src/manager.rs +++ b/codex-rs/models-manager/src/manager.rs @@ -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; /// OpenAI-compatible model manager backed by bundled models, cache, and `/models`. #[derive(Debug)] pub struct OpenAiModelsManager { + model_provider_id: String, remote_models: RwLock>, etag: RwLock>, 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, auth_manager: Option>, ) -> 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"); diff --git a/codex-rs/models-manager/src/manager_tests.rs b/codex-rs/models-manager/src/manager_tests.rs index e7d1691a3a..e0406d7706 100644 --- a/codex-rs/models-manager/src/manager_tests.rs +++ b/codex-rs/models-manager/src/manager_tests.rs @@ -179,8 +179,9 @@ fn openai_manager_for_tests( codex_home: std::path::PathBuf, endpoint_client: Arc, ) -> 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, auth_manager: Option>, ) -> 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, + auth_manager: Option>, +) -> 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(),