codex/codex-rs/app-server/src/models_refresh_worker_tests.rs

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<SharedModelsManager>,
    auth_manager: Arc<codex_login::AuthManager>,
) -> Arc<codex_core::ThreadManager> {
    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<Arc<TestModelsEndpoint>>,
}

impl TestModelsEndpoint {
    fn new() -> Arc<Self> {
        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<Arc<dyn ModelsEndpointClient>> {
        self.originator_endpoint
            .clone()
            .map(|endpoint| endpoint as Arc<dyn ModelsEndpointClient>)
    }

    fn identity(&self) -> Option<String> {
        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<ModelsEndpointResponse>> {
        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::<Option<CloudConfigBundle>>();
        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<_>>(),
        vec!["project-model"]
    );
    Ok(())
}