use std::sync::Arc; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use std::time::Duration; use codex_config::CloudConfigBundle; use codex_config::CloudConfigBundleLoader; use codex_config::LoaderOverrides; use codex_config::test_support::CloudConfigBundleFixture; use codex_core::config::ConfigBuilder; use codex_core::config::ConfigOverrides; use codex_feedback::CodexFeedback; use codex_http_client::HttpClientFactory; use codex_models_manager::manager::ModelCatalogSettings; use codex_models_manager::manager::ModelsEndpointClient; use codex_models_manager::manager::ModelsEndpointFuture; use codex_models_manager::manager::ModelsEndpointResponse; use codex_models_manager::manager::OpenAiModelsManager; use codex_models_manager::manager::SharedModelsManager; use codex_protocol::error::CodexErr; use codex_protocol::error::Result as CoreResult; use codex_utils_absolute_path::AbsolutePathBuf; use pretty_assertions::assert_eq; use tempfile::tempdir; use tokio::sync::Notify; use tokio::sync::mpsc; use tokio::sync::oneshot; use tracing_subscriber::layer::SubscriberExt; use super::*; fn test_thread_manager( config: &codex_core::config::Config, models_manager: Option, auth_manager: Arc, ) -> Arc { let factory = match models_manager { Some(models) => codex_core::test_support::fixed_model_catalog(models), None => Arc::new(ModelCatalogSettings::disk(config.codex_home.to_path_buf())), }; let startup = codex_core::build_startup_models( config, auth_manager, uuid::Uuid::new_v4().to_string(), factory, ); Arc::new(codex_core::ThreadManager::new( config, startup, codex_core::CodexAppsToolsCache::default(), codex_protocol::protocol::SessionSource::Exec, Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), codex_extension_api::empty_extension_registry(), Arc::new(codex_core::test_support::EmptyUserInstructionsProvider), /*analytics_events_client*/ None, codex_core::passthrough_image_store(), codex_core::ThreadStoreBuilder::from_config(config, /*state_db*/ None).build(), /*agent_graph_store*/ None, /*attestation_provider*/ None, /*external_time_provider*/ None, )) } #[derive(Debug)] struct TestModelsEndpoint { fetch_count: AtomicUsize, fetched: Notify, release_second_fetch: Notify, originator_endpoint: Option>, } impl TestModelsEndpoint { fn new() -> Arc { Arc::new(Self { fetch_count: AtomicUsize::new(0), fetched: Notify::new(), release_second_fetch: Notify::new(), originator_endpoint: None, }) } async fn wait_for_fetch_count(&self, expected: usize) { tokio::time::timeout(Duration::from_secs(1), async { while self.fetch_count.load(Ordering::SeqCst) < expected { self.fetched.notified().await; } }) .await .unwrap_or_else(|_| panic!("expected {expected} model fetches")); } } impl ModelsEndpointClient for TestModelsEndpoint { fn for_originator( &self, _originator: codex_models_manager::ModelCatalogOriginator, ) -> Option> { self.originator_endpoint .clone() .map(|endpoint| endpoint as Arc) } fn identity(&self) -> Option { Some("test-provider".to_string()) } fn has_command_auth(&self) -> bool { true } fn uses_codex_backend(&self) -> ModelsEndpointFuture<'_, bool> { Box::pin(async { false }) } fn list_models<'a>( &'a self, _client_version: &'a str, _http_client_factory: HttpClientFactory, ) -> ModelsEndpointFuture<'a, CoreResult> { Box::pin(async move { let fetch_index = self.fetch_count.fetch_add(1, Ordering::SeqCst); self.fetched.notify_one(); if fetch_index == 0 { return Err(CodexErr::Io(std::io::Error::other("test failure"))); } if fetch_index == 1 { self.release_second_fetch.notified().await; } Ok(ModelsEndpointResponse { models: Vec::new(), etag: None, identity: self.identity().expect("test endpoint identity"), }) }) } } #[tokio::test] async fn refreshes_immediately_periodically_and_stops_when_dropped() { let codex_home = tempdir().expect("temp dir"); let mut endpoint = TestModelsEndpoint::new(); let work_endpoint = TestModelsEndpoint::new(); work_endpoint.release_second_fetch.notify_one(); Arc::get_mut(&mut endpoint).unwrap().originator_endpoint = Some(work_endpoint.clone()); let requirements_path = codex_home.path().join("requirements.toml"); std::fs::write(&requirements_path, "model_provider = 'ollama'").unwrap(); let config_manager = crate::config_manager::ConfigManager::new_for_tests( codex_home.path().to_path_buf(), vec![("model_provider".to_string(), "ollama".into())], LoaderOverrides::with_managed_config_path_for_tests( codex_home.path().join("managed_config.toml"), ), CloudConfigBundleLoader::default(), ); // The TUI forwards its typed --oss selection through the raw CLI overrides too. let config = Arc::new( ConfigBuilder::default() .codex_home(codex_home.path().to_path_buf()) .loader_overrides(LoaderOverrides::without_managed_config_for_tests()) .harness_overrides(ConfigOverrides { model_provider: Some("ollama".to_string()), ..Default::default() }) .build() .await .unwrap(), ); let models_manager: SharedModelsManager = Arc::new(OpenAiModelsManager::new( codex_home.path().to_path_buf(), endpoint.clone(), /*auth_manager*/ None, )); let work_manager = models_manager.for_originator("codex_work_desktop").unwrap(); let mut work_updates = work_manager.model_catalog_updates().unwrap(); let auth_manager = codex_login::AuthManager::shared_from_config( config.as_ref(), /*enable_codex_api_key_env*/ false, ) .await .expect("test auth manager"); let catalog = Arc::new(ModelCatalog::new( config_manager, Some(config.cwd.to_path_buf()), Arc::clone(&config), test_thread_manager(&config, Some(Arc::clone(&models_manager)), auth_manager), )); let worker = spawn_with_interval(&catalog, Duration::from_millis(/*millis*/ 10)); endpoint.wait_for_fetch_count(/*expected*/ 2).await; drop(worker); endpoint.release_second_fetch.notify_one(); work_endpoint.wait_for_fetch_count(/*expected*/ 2).await; work_updates.changed().await.unwrap(); tokio::time::sleep(Duration::from_millis(30)).await; assert_eq!(endpoint.fetch_count.load(Ordering::SeqCst), 2); assert_eq!(work_endpoint.fetch_count.load(Ordering::SeqCst), 2); // Reloading after login preserves an embedded caller's typed --oss override. catalog.invalidate().await; assert_eq!( catalog.current().await.unwrap().0.model_provider_id, "ollama" ); assert!(Arc::ptr_eq( &catalog.current().await.unwrap().1, &models_manager )); // The catalog follows a new managed route without reusing the old endpoint. std::fs::write(&requirements_path, "model_provider = 'openai'").unwrap(); catalog.invalidate().await; catalog.list_models(RefreshStrategy::Offline).await.unwrap(); assert_eq!( catalog.current().await.unwrap().0.model_provider_id, "openai" ); std::fs::remove_file(&requirements_path).unwrap(); catalog.invalidate().await; assert_eq!( catalog.current().await.unwrap().0.model_provider_id, "ollama" ); assert_eq!(endpoint.fetch_count.load(Ordering::SeqCst), 2); } #[tokio::test] async fn background_refresh_blocks_invalid_requirements_then_recovers() -> anyhow::Result<()> { let home = tempdir()?; let config = Arc::new( crate::config_manager::ConfigManager::without_managed_config_for_tests( home.path().to_path_buf(), ) .load_non_project_config() .await?, ); let (checks_tx, mut checks_rx) = mpsc::unbounded_channel(); let loader = CloudConfigBundleLoader::from_getter(move || { let (reply_tx, reply_rx) = oneshot::channel::>(); checks_tx .send(reply_tx) .expect("requirements check receiver"); async move { Ok(reply_rx.await.expect("requirements check response")) } }); let config_manager = crate::config_manager::ConfigManager::new_for_tests( home.path().to_path_buf(), Vec::new(), LoaderOverrides::without_managed_config_for_tests(), loader, ); let endpoint = TestModelsEndpoint::new(); let models_manager: SharedModelsManager = Arc::new(OpenAiModelsManager::new( home.path().to_path_buf(), endpoint.clone(), /*auth_manager*/ None, )); let auth_manager = codex_login::AuthManager::shared_from_config( config.as_ref(), /*enable_codex_api_key_env*/ false, ) .await .expect("test auth manager"); let catalog = Arc::new(ModelCatalog::new( config_manager, Some(config.cwd.to_path_buf()), Arc::clone(&config), test_thread_manager(&config, Some(models_manager), auth_manager), )); let feedback = CodexFeedback::new(); // This test uses Tokio's current-thread runtime, so spawned refreshes share the subscriber. let _subscriber_guard = tracing::subscriber::set_default( tracing_subscriber::registry().with(feedback.logger_layer()), ); let worker = spawn_with_interval(&catalog, Duration::from_millis(/*millis*/ 10)); tokio::time::timeout(Duration::from_secs(/*secs*/ 10), async { for expected in 1..=2 { checks_rx.recv().await.unwrap().send(/*t*/ None).unwrap(); endpoint.wait_for_fetch_count(expected).await; } // Let the in-flight request finish, then control each refresh's requirements. endpoint.release_second_fetch.notify_one(); let mut next_check = checks_rx.recv().await.unwrap(); for requirements in [ "model_provider = 'other'", "model_provider = []", "[model_providers.gateway]\nexperimental_bearer_token = 'private-provider-token-marker' trailing", ] { next_check .send(Some( CloudConfigBundleFixture::enterprise_requirement(requirements).into_bundle(), )) .unwrap(); // Reaching the next check proves the previous refresh completed. Keep // this check blocked while asserting that the previous one did not fetch. next_check = checks_rx.recv().await.unwrap(); assert_eq!(endpoint.fetch_count.load(Ordering::SeqCst), 2); } next_check.send(/*t*/ None).unwrap(); endpoint.wait_for_fetch_count(/*expected*/ 3).await; }) .await?; drop(worker); let logs = String::from_utf8( feedback .snapshot(/*session_id*/ None) .log_attachment(/*logs_override*/ None) .buffer, )?; assert!(logs.contains("model catalog refresh blocked by provider requirements")); assert!(logs.contains("error_kind=InvalidData")); assert!(!logs.contains("private-provider-token-marker")); Ok(()) } /// A transient cloud failure must not pin the fallback catalog after policy recovers. #[tokio::test] async fn catalog_reload_recovers_after_cloud_failure() -> anyhow::Result<()> { let home = tempdir()?; let failed = Arc::new(std::sync::atomic::AtomicBool::new(false)); let loader_failed = Arc::clone(&failed); let loader = CloudConfigBundleLoader::from_getter(move || { let failed = Arc::clone(&loader_failed); async move { if failed.load(Ordering::SeqCst) { Err(codex_config::CloudConfigBundleLoadError::new( codex_config::CloudConfigBundleLoadErrorCode::RequestFailed, /*status_code*/ None, "test cloud timeout", )) } else { Ok(Some( CloudConfigBundleFixture::enterprise_requirement("model_provider = 'ollama'") .into_bundle(), )) } } }); let manager = crate::config_manager::ConfigManager::new_for_tests( home.path().to_path_buf(), Vec::new(), LoaderOverrides::without_managed_config_for_tests(), loader, ); let config = Arc::new(manager.load_non_project_config().await?); let auth = codex_login::AuthManager::shared_from_config( config.as_ref(), /*enable_codex_api_key_env*/ false, ) .await?; let catalog = Arc::new(ModelCatalog::new( manager, /*startup_cwd*/ None, Arc::clone(&config), test_thread_manager(&config, /*models_manager*/ None, auth), )); failed.store(true, Ordering::SeqCst); catalog.invalidate().await; assert_eq!(catalog.current().await?.0.model_provider_id, "openai"); assert!(catalog.list_models(RefreshStrategy::Offline).await.is_err()); failed.store(false, Ordering::SeqCst); // Exercise the same online request used by the periodic background worker. catalog.list_models(RefreshStrategy::Online).await?; assert_eq!(catalog.current().await?.0.model_provider_id, "ollama"); assert!( !catalog .list_models(RefreshStrategy::Offline) .await? .is_empty() ); Ok(()) } /// Repairing a project config after fallback must restore that project's local model catalog. #[tokio::test] async fn catalog_reload_recovers_original_project_after_startup_fallback() -> anyhow::Result<()> { let home = tempdir()?; let project = tempdir()?; let project_path = AbsolutePathBuf::from_absolute_path(project.path())? .canonicalize()? .to_path_buf(); std::fs::create_dir(project_path.join(".codex"))?; std::fs::write( home.path().join("config.toml"), toml::to_string(&serde_json::json!({ "projects": {project_path.to_string_lossy(): {"trust_level": "trusted"}} }))?, )?; let project_config = project_path.join(".codex/config.toml"); std::fs::write(&project_config, "model_catalog_json = [")?; let manager = crate::config_manager::ConfigManager::without_managed_config_for_tests( home.path().to_path_buf(), ); let startup = manager .load_startup_config(Some(project_path.clone())) .await?; assert!(startup.used_defaults); let config = Arc::new(startup.config); let auth = codex_login::AuthManager::shared_from_config( config.as_ref(), /*enable_codex_api_key_env*/ false, ) .await?; let catalog = ModelCatalog::new( manager, Some(project_path.clone()), Arc::clone(&config), test_thread_manager(&config, /*models_manager*/ None, auth), ); let mut model = codex_models_manager::bundled_models_response()? .models .remove(0); model.slug = "project-model".into(); std::fs::write( project_path.join(".codex/catalog.json"), serde_json::to_vec(&serde_json::json!({"models": [model]}))?, )?; std::fs::write(&project_config, "model_catalog_json = 'catalog.json'")?; catalog.invalidate().await; let (config, _) = catalog.current().await?; assert_eq!(config.cwd.as_path(), project_path); let models = catalog.list_models(RefreshStrategy::Offline).await?; assert_eq!( models .iter() .map(|model| model.id.as_str()) .collect::>(), vec!["project-model"] ); Ok(()) }