diff --git a/codex-rs/app-server/src/request_processors/config_processor.rs b/codex-rs/app-server/src/request_processors/config_processor.rs index bff9b0db39..ef7d769084 100644 --- a/codex-rs/app-server/src/request_processors/config_processor.rs +++ b/codex-rs/app-server/src/request_processors/config_processor.rs @@ -63,6 +63,7 @@ const BACKGROUND_PAGINATED_ROLLOUT_MIGRATION_FEATURE: &str = "background_paginated_rollout_migration"; const SUPPORTED_EXPERIMENTAL_FEATURE_ENABLEMENT: &[&str] = &[ + "api_key_model_discovery", "auth_elicitation", BACKGROUND_PAGINATED_ROLLOUT_MIGRATION_FEATURE, "codex_apps_mcp_2026_07_28", @@ -315,6 +316,11 @@ impl ConfigRequestProcessor { .map_err(|_| internal_error("failed to update feature enablement"))?; let config = self.load_latest_config(/*fallback_cwd*/ None).await?; + self.thread_manager + .get_models_manager() + .set_api_key_model_discovery_enabled( + config.features.enabled(Feature::ApiKeyModelDiscovery), + ); if should_start_background_rollout_migration && config.features.enabled(feature) { self.thread_manager.start_background_rollout_migration(); } diff --git a/codex-rs/app-server/tests/suite/v2/model_list.rs b/codex-rs/app-server/tests/suite/v2/model_list.rs index fbc262a308..d01454a356 100644 --- a/codex-rs/app-server/tests/suite/v2/model_list.rs +++ b/codex-rs/app-server/tests/suite/v2/model_list.rs @@ -7,6 +7,8 @@ use app_test_support::TestAppServer; use app_test_support::write_chatgpt_auth; use app_test_support::write_models_cache; use codex_app_server_protocol::ClientRequest; +use codex_app_server_protocol::ExperimentalFeatureEnablementSetParams; +use codex_app_server_protocol::ExperimentalFeatureEnablementSetResponse; use codex_app_server_protocol::JSONRPCError; use codex_app_server_protocol::Model; use codex_app_server_protocol::ModelListParams; @@ -16,20 +18,129 @@ use codex_app_server_protocol::ModelUpgradeInfo; use codex_app_server_protocol::ReasoningEffortOption; use codex_app_server_protocol::RequestId; use codex_config::types::AuthCredentialsStoreMode; +use codex_login::AuthKeyringBackendKind; +use codex_login::login_with_api_key; use codex_protocol::openai_models::MODEL_SPECIALTY_CYBER; use codex_protocol::openai_models::ModelInfo; use codex_protocol::openai_models::ModelPreset; use codex_protocol::openai_models::ModelsResponse; -use core_test_support::responses::mount_models_once; use pretty_assertions::assert_eq; use serde_json::json; use tempfile::TempDir; +use test_case::test_case; use tokio::time::timeout; +use wiremock::Mock; use wiremock::MockServer; +use wiremock::ResponseTemplate; +use wiremock::matchers::header; +use wiremock::matchers::method; +use wiremock::matchers::path; const DEFAULT_TIMEOUT: Duration = Duration::from_secs(10); const INVALID_REQUEST_ERROR_CODE: i64 = -32600; +#[test_case(None, false; "default off")] +#[test_case(None, true; "app rollout")] +#[test_case(Some(false), true; "user opt out")] +#[test_case(Some(true), false; "user opt in")] +#[tokio::test] +async fn api_key_model_discovery_startup_enablement_respects_user_config( + user_enablement: Option, + app_enablement: bool, +) -> Result<()> { + let server = MockServer::start().await; + let mut remote_model = codex_models_manager::bundled_models_response()? + .models + .remove(0); + remote_model.slug = "rollout-model".into(); + remote_model.visibility = codex_protocol::openai_models::ModelVisibility::List; + remote_model.supported_in_api = true; + Mock::given(method("GET")) + .and(path("/v1/models")) + .respond_with( + ResponseTemplate::new(/*s*/ 200).set_body_json(ModelsResponse { + models: vec![remote_model.clone()], + }), + ) + .mount(&server) + .await; + let codex_home = TempDir::new()?; + let server_uri = server.uri(); + let feature_config = user_enablement + .map(|enabled| format!("[features]\napi_key_model_discovery = {enabled}\n")) + .unwrap_or_default(); + std::fs::write( + codex_home.path().join("config.toml"), + format!("openai_base_url = \"{server_uri}/v1\"\n{feature_config}"), + )?; + login_with_api_key( + codex_home.path(), + "test-key", + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + )?; + let mut mcp = TestAppServer::builder() + .with_codex_home(codex_home.path()) + .with_env_overrides(&[("OPENAI_API_KEY", None), ("CODEX_API_KEY", None)]) + .build_initialized() + .await?; + let mut bundled = codex_models_manager::bundled_models_response()?.models; + bundled.sort_by_key(|model| model.priority); + let mut bundled = ModelPreset::filter_by_auth( + bundled.into_iter().map(Into::into).collect(), + /*chatgpt_mode*/ false, + ); + ModelPreset::mark_default_by_picker_visibility(&mut bundled); + let mut remote = vec![ModelPreset::from(remote_model)]; + ModelPreset::mark_default_by_picker_visibility(&mut remote); + let _: ExperimentalFeatureEnablementSetResponse = mcp + .request( + |request_id| ClientRequest::ExperimentalFeatureEnablementSet { + request_id, + params: ExperimentalFeatureEnablementSetParams { + enablement: [("api_key_model_discovery".to_string(), app_enablement)].into(), + }, + }, + ) + .await?; + let response: ModelListResponse = mcp + .request(|request_id| ClientRequest::ModelList { + request_id, + params: ModelListParams { + limit: Some(100), + include_hidden: Some(true), + cursor: None, + }, + }) + .await?; + let enabled = user_enablement.unwrap_or(app_enablement); + let expected = if enabled { &remote } else { &bundled }; + assert_eq!( + response, + ModelListResponse { + data: expected + .iter() + .map(|preset| Model { + // These catalogs retain personality metadata; the cache fixture does not. + supports_personality: preset.supports_personality, + ..model_from_preset(preset) + }) + .collect(), + next_cursor: None, + } + ); + if !enabled { + assert!( + server + .received_requests() + .await + .expect("request recording is enabled") + .is_empty() + ); + } + Ok(()) +} + fn model_from_preset(preset: &ModelPreset) -> Model { Model { id: preset.id.clone(), @@ -156,16 +267,21 @@ async fn list_models_includes_hidden_models() -> Result<()> { Ok(()) } +#[test_case("chatgpt-access-token", None; "chatgpt")] +#[test_case("test-api-key", Some("test-api-key"); "api key")] #[tokio::test] -async fn list_models_uses_chatgpt_remote_catalog_as_source_of_truth() -> Result<()> { +async fn list_models_uses_remote_catalog_as_source_of_truth( + bearer_token: &str, + api_key: Option<&str>, +) -> Result<()> { let server = MockServer::start().await; let remote_models = [json!("2030-01-01T00:00:00Z"), serde_json::Value::Null] .into_iter() .enumerate() .map(|(priority, retirement_at)| { serde_json::from_value::(json!({ - "slug": format!("chatgpt-remote-only-{priority}"), - "display_name": "ChatGPT Remote Only", + "slug": format!("remote-only-{priority}"), + "display_name": "Remote Only", "description": "Remote-only model for app-server model/list coverage", "model_specialty": MODEL_SPECIALTY_CYBER, "default_reasoning_level": "max", @@ -196,13 +312,18 @@ async fn list_models_uses_chatgpt_remote_catalog_as_source_of_truth() -> Result< })) }) .collect::, _>>()?; - let models_mock = mount_models_once( - &server, - ModelsResponse { - models: remote_models.clone(), - }, - ) - .await; + // The startup refresh worker and model/list can both fetch before the cache is populated. + let _models_mock = Mock::given(method("GET")) + .and(path("/v1/models")) + .and(header("authorization", format!("Bearer {bearer_token}"))) + .respond_with( + ResponseTemplate::new(/*s*/ 200).set_body_json(ModelsResponse { + models: remote_models.clone(), + }), + ) + .expect(1..) + .mount_as_scoped(&server) + .await; let codex_home = TempDir::new()?; let server_uri = server.uri(); @@ -214,19 +335,30 @@ model = "mock-model" approval_policy = "never" sandbox_mode = "read-only" openai_base_url = "{server_uri}/v1" +[features] +api_key_model_discovery = true "# ), )?; - write_chatgpt_auth( - codex_home.path(), - ChatGptAuthFixture::new("chatgpt-access-token").plan_type("pro"), - AuthCredentialsStoreMode::File, - )?; + if let Some(api_key) = api_key { + login_with_api_key( + codex_home.path(), + api_key, + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + )?; + } else { + write_chatgpt_auth( + codex_home.path(), + ChatGptAuthFixture::new(bearer_token).plan_type("pro"), + AuthCredentialsStoreMode::File, + )?; + } let mut mcp = TestAppServer::builder() .with_codex_home(codex_home.path()) .without_auto_env() - .with_env_overrides(&[("OPENAI_API_KEY", None)]) + .with_env_overrides(&[("OPENAI_API_KEY", None), ("CODEX_API_KEY", None)]) .build_initialized() .await?; let request_id = mcp @@ -275,11 +407,6 @@ openai_base_url = "{server_uri}/v1" assert_eq!(items, expected_items); assert!(next_cursor.is_none()); - assert_eq!( - models_mock.requests().len(), - 1, - "expected a single /models request" - ); Ok(()) } diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index f84fac728a..9e0e2eb685 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -631,6 +631,9 @@ "default": null, "description": "Optional feature toggles scoped to this profile.", "properties": { + "api_key_model_discovery": { + "type": "boolean" + }, "apply_patch_freeform": { "type": "boolean" }, @@ -6123,6 +6126,9 @@ "default": null, "description": "Centralized feature flags (new). Prefer this over individual toggles.", "properties": { + "api_key_model_discovery": { + "type": "boolean" + }, "apply_patch_freeform": { "type": "boolean" }, diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 1ed26b95e4..b202e15fc6 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -383,10 +383,14 @@ pub fn build_models_manager( auth_manager: Arc, ) -> SharedModelsManager { let provider = create_model_provider(config.model_provider.clone(), Some(auth_manager)); - provider.models_manager( + let manager = provider.models_manager( config.codex_home.to_path_buf(), config.model_catalog.clone(), - ) + ); + manager.set_api_key_model_discovery_enabled( + config.features.enabled(Feature::ApiKeyModelDiscovery), + ); + manager } pub fn thread_store_from_config( diff --git a/codex-rs/core/tests/suite/models_cache_auth.rs b/codex-rs/core/tests/suite/models_cache_auth.rs index 8fe37ccd0d..a47b82a97b 100644 --- a/codex-rs/core/tests/suite/models_cache_auth.rs +++ b/codex-rs/core/tests/suite/models_cache_auth.rs @@ -72,11 +72,28 @@ async fn auth_and_provider_switches_do_not_reuse_chatgpt_catalog() -> Result<()> .models, bundled ); + assert_eq!(models_mock.requests().len(), 1); drop(chatgpt); + server.reset().await; + let mut api_model = model.clone(); + api_model.default_service_tier = None; + let api_models_mock = responses::mount_models_once( + &server, + ModelsResponse { + models: vec![api_model], + }, + ) + .await; let api = test_codex() .with_home(home.clone()) .with_auth(CodexAuth::from_api_key("api-key")) + .with_config(|config| { + config + .features + .enable(codex_features::Feature::ApiKeyModelDiscovery) + .expect("enable API-key model discovery"); + }) .with_model("gpt-5.5") .build_with_auto_env(&server) .await?; @@ -108,7 +125,7 @@ async fn auth_and_provider_switches_do_not_reuse_chatgpt_catalog() -> Result<()> explicit.single_request().body_json()["service_tier"], "priority" ); - assert_eq!(models_mock.requests().len(), 1); + assert_eq!(api_models_mock.requests().len(), 1); drop(api); let other_server = wiremock::MockServer::start().await; diff --git a/codex-rs/core/tests/suite/remote_models.rs b/codex-rs/core/tests/suite/remote_models.rs index 8e5eb64eae..a8e7a41535 100644 --- a/codex-rs/core/tests/suite/remote_models.rs +++ b/codex-rs/core/tests/suite/remote_models.rs @@ -1,6 +1,7 @@ #![cfg(not(target_os = "windows"))] use anyhow::Result; use codex_core::TurnInputRequest; +use codex_features::Feature; use codex_login::CodexAuth; use codex_model_provider_info::ModelProviderInfo; use codex_model_provider_info::built_in_model_providers; @@ -852,8 +853,10 @@ async fn remote_models_truncation_policy_with_tool_output_override() -> Result<( Ok(()) } +#[test_case(CodexAuth::create_dummy_chatgpt_auth_for_testing(); "chatgpt")] +#[test_case(CodexAuth::from_api_key("test-api-key"); "api key")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn remote_models_apply_legacy_instructions() -> Result<()> { +async fn remote_models_apply_legacy_instructions(auth: CodexAuth) -> Result<()> { skip_if_no_network!(Ok(())); skip_if_sandbox!(Ok(())); @@ -959,12 +962,14 @@ async fn remote_models_apply_legacy_instructions() -> Result<()> { ) .await; - let mut builder = test_codex() - .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) - .with_config(|config| { - config.update_plan_enabled = true; - config.model = Some("gpt-5.2".to_string()); - }); + let mut builder = test_codex().with_auth(auth).with_config(|config| { + config + .features + .enable(Feature::ApiKeyModelDiscovery) + .expect("enable API-key model discovery"); + config.update_plan_enabled = true; + config.model = Some("gpt-5.2".to_string()); + }); let TestCodex { codex, cwd, diff --git a/codex-rs/core/tests/suite/responses_system_proxy.rs b/codex-rs/core/tests/suite/responses_system_proxy.rs index daa385d09b..2581f8c6ae 100644 --- a/codex-rs/core/tests/suite/responses_system_proxy.rs +++ b/codex-rs/core/tests/suite/responses_system_proxy.rs @@ -1,5 +1,6 @@ use anyhow::Result; use codex_features::Feature; +use codex_models_manager::bundled_models_response; use core_test_support::responses::ev_completed; use core_test_support::responses::ev_response_created; use core_test_support::responses::sse; @@ -105,6 +106,9 @@ async fn regular_responses_turn_honors_respect_system_proxy() -> Result<()> { let server = MockServer::start().await; let mut builder = test_codex().with_config(|config| { + // The proxy serves one inference response; model discovery must not consume it. + config.model_catalog = + Some(bundled_models_response().expect("bundled models.json should parse")); config.model_provider.base_url = Some("http://responses-proxy.invalid/v1".to_string()); config .features diff --git a/codex-rs/core/tests/suite/retry_after.rs b/codex-rs/core/tests/suite/retry_after.rs index c24d510049..991afbe92d 100644 --- a/codex-rs/core/tests/suite/retry_after.rs +++ b/codex-rs/core/tests/suite/retry_after.rs @@ -5,6 +5,7 @@ use codex_client::run_with_retry; use codex_http_client::Request; use codex_http_client::TransportError; use codex_login::CodexAuth; +use codex_models_manager::bundled_models_response; use codex_protocol::protocol::CodexErrorInfo; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::Op; @@ -1628,6 +1629,9 @@ async fn websocket_overload_with_nested_retry_after_is_terminal() -> Result<()> .await; let test = test_codex() .with_config(|config| { + // Capture inference retries without unrelated startup model-discovery retries. + config.model_catalog = + Some(bundled_models_response().expect("bundled models.json should parse")); config.model_provider.request_max_retries = Some(2); config.model_provider.stream_max_retries = Some(2); }) @@ -1709,6 +1713,9 @@ async fn websocket_overload_without_retry_after_is_terminal() -> Result<()> { .await; let test = test_codex() .with_config(|config| { + // Capture inference retries without unrelated startup model-discovery retries. + config.model_catalog = + Some(bundled_models_response().expect("bundled models.json should parse")); config.model_provider.request_max_retries = Some(2); config.model_provider.stream_max_retries = Some(2); }) diff --git a/codex-rs/features/src/lib.rs b/codex-rs/features/src/lib.rs index e8c8d1381e..1a31316b79 100644 --- a/codex-rs/features/src/lib.rs +++ b/codex-rs/features/src/lib.rs @@ -91,6 +91,8 @@ impl Stage { /// Unique features toggled via configuration. #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] pub enum Feature { + /// Discover model catalogs for OpenAI API-key authentication. + ApiKeyModelDiscovery, /// Enable the interactive transcript composer and turn-selection UI. TranscriptV2, // Stable. @@ -1219,6 +1221,12 @@ pub const FEATURES: &[FeatureSpec] = &[ stage: Stage::Removed, default_enabled: false, }, + FeatureSpec { + id: Feature::ApiKeyModelDiscovery, + key: "api_key_model_discovery", + stage: Stage::UnderDevelopment, + default_enabled: false, + }, FeatureSpec { id: Feature::EnableRequestCompression, key: "enable_request_compression", diff --git a/codex-rs/model-provider/src/models_endpoint.rs b/codex-rs/model-provider/src/models_endpoint.rs index 12a2917bc5..01b4ee590b 100644 --- a/codex-rs/model-provider/src/models_endpoint.rs +++ b/codex-rs/model-provider/src/models_endpoint.rs @@ -20,6 +20,7 @@ use codex_login::AuthManager; use codex_login::CodexAuth; use codex_login::collect_auth_env_telemetry; use codex_login::default_client::create_client_for_route_async; +use codex_model_provider_info::CHATGPT_CODEX_BASE_URL; use codex_model_provider_info::ModelProviderInfo; use codex_models_manager::manager::ModelsEndpointClient; use codex_models_manager::manager::ModelsEndpointFuture; @@ -84,6 +85,13 @@ impl OpenAiModelsEndpoint { let identity = crate::models_identity::identity(&self.provider_info, auth.as_ref())?; let auth_mode = auth.as_ref().map(CodexAuth::auth_mode); let mut api_provider = self.provider_info.to_api_provider(auth_mode)?; + if auth.as_ref().is_some_and(CodexAuth::is_api_key_auth) + && self.supports_api_key_models() + && self.provider_info.base_url.is_none() + { + // Codex metadata is served by the Codex backend, not the public /v1/models API. + api_provider.base_url = CHATGPT_CODEX_BASE_URL.to_string(); + } enforce_managed_residency(&mut api_provider); let api_auth = resolve_provider_auth(auth.as_ref(), &self.provider_info)?; let request_url = @@ -132,6 +140,10 @@ impl OpenAiModelsEndpoint { } impl ModelsEndpointClient for OpenAiModelsEndpoint { + fn supports_api_key_models(&self) -> bool { + self.provider_info.is_openai() + } + fn identity(&self) -> Option { let auth = self .auth_manager @@ -300,6 +312,10 @@ mod tests { use codex_login::default_client::ResidencyRequirement; use codex_login::default_client::create_client; use codex_login::default_client::set_default_client_residency_requirement; + use codex_models_manager::manager::ModelsManager; + use codex_models_manager::manager::OpenAiModelsManager; + use codex_models_manager::manager::RefreshStrategy; + use codex_protocol::auth::AuthMode; use codex_protocol::config_types::ModelProviderAuthInfo; use codex_protocol::openai_models::ModelsResponse; use pretty_assertions::assert_eq; @@ -333,6 +349,76 @@ mod tests { } } + #[derive(Debug)] + struct CaptureModelsUrl(Mutex>); + + impl ModelsTransportBuilder for CaptureModelsUrl { + fn build( + &self, + _http_client_factory: HttpClientFactory, + request_url: String, + ) -> ModelsTransportFuture<'_> { + *self.0.lock().unwrap() = Some(request_url); + Box::pin(async { Err(std::io::Error::other("transport intentionally unavailable")) }) + } + } + + #[tokio::test] + async fn api_key_discovery_respects_provider_routing() { + let client_version = codex_models_manager::client_version_to_whole(); + for (name, base_url, models_url, inference_url) in [ + ( + "OpenAI", + None, + Some("https://chatgpt.com/backend-api/codex/models"), + "https://api.openai.com/v1", + ), + ( + "OpenAI", + Some("https://example.com/codex"), + Some("https://example.com/codex/models"), + "https://example.com/codex", + ), + ( + "Azure", + Some("https://example.openai.azure.com/openai/v1"), + None, + "https://example.openai.azure.com/openai/v1", + ), + ] { + let capture = Arc::new(CaptureModelsUrl(Mutex::new(/*t*/ None))); + let auth = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("test-api-key")); + let endpoint = Arc::new(OpenAiModelsEndpoint { + provider_info: ModelProviderInfo { + name: name.to_string(), + ..ModelProviderInfo::create_openai_provider(base_url.map(str::to_string)) + }, + auth_manager: Some(auth.clone()), + transport_builder: capture.clone(), + }); + let manager = OpenAiModelsManager::new_without_cache(endpoint.clone(), Some(auth)); + manager.set_api_key_model_discovery_enabled(/*enabled*/ true); + manager + .raw_model_catalog( + RefreshStrategy::Online, + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + ) + .await; + assert_eq!( + *capture.0.lock().unwrap(), + models_url.map(|url| format!("{url}?client_version={client_version}")) + ); + assert_eq!( + endpoint + .provider_info + .to_api_provider(Some(AuthMode::ApiKey)) + .unwrap() + .base_url, + inference_url + ); + } + } + fn provider_info_with_command_auth() -> ModelProviderInfo { ModelProviderInfo { auth: Some(ModelProviderAuthInfo { diff --git a/codex-rs/models-manager/src/api_key_discovery_tests.rs b/codex-rs/models-manager/src/api_key_discovery_tests.rs new file mode 100644 index 0000000000..f0c62fae08 --- /dev/null +++ b/codex-rs/models-manager/src/api_key_discovery_tests.rs @@ -0,0 +1,93 @@ +//! API-key discovery rollout must gate network requests and cached catalog authority. + +use super::*; +use pretty_assertions::assert_eq; + +#[tokio::test] +async fn api_key_discovery_disabled_preserves_command_auth_discovery_and_merging() { + let models = vec![remote_model( + "command-auth", + "Command Auth", + /*priority*/ 0, + )]; + let endpoint = Arc::new(TestModelsEndpoint { + has_command_auth: true, + uses_codex_backend: false, + responses: Mutex::new(vec![models.clone()].into()), + fetch_count: AtomicUsize::new(0), + observed_proxy_policy: Mutex::new(None), + }); + let manager = OpenAiModelsManager::new_without_cache( + endpoint.clone(), + Some(AuthManager::from_auth_for_testing(CodexAuth::from_api_key( + "test-key", + ))), + ); + let mut merged = load_remote_models_from_file().unwrap(); + merged.extend(models.clone()); + assert_eq!( + manager + .raw_model_catalog(RefreshStrategy::Online, DEFAULT_HTTP_CLIENT_FACTORY) + .await + .models, + merged + ); + assert_eq!(endpoint.fetch_count(), 1); +} + +#[tokio::test] +async fn api_key_discovery_startup_flag_controls_fetches_and_cached_catalogs() { + let home = tempdir().unwrap(); + let models = vec![remote_model("dynamic", "Dynamic", /*priority*/ 0)]; + let endpoint = TestModelsEndpoint::without_refresh(vec![models.clone()]); + let auth = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("test-key")); + let manager = + OpenAiModelsManager::new(home.path().into(), endpoint.clone(), Some(auth.clone())); + let bundled = load_remote_models_from_file().unwrap(); + + for strategy in [ + RefreshStrategy::Online, + RefreshStrategy::OnlineIfUncached, + RefreshStrategy::Offline, + ] { + assert_eq!( + manager + .raw_model_catalog(strategy, DEFAULT_HTTP_CLIENT_FACTORY) + .await + .models, + bundled + ); + } + assert_eq!(endpoint.fetch_count(), 0); + + let manager = + OpenAiModelsManager::new(home.path().into(), endpoint.clone(), Some(auth.clone())); + manager.set_api_key_model_discovery_enabled(/*enabled*/ true); + assert_eq!( + manager + .raw_model_catalog(RefreshStrategy::Online, DEFAULT_HTTP_CLIENT_FACTORY) + .await + .models, + models + ); + assert_eq!(endpoint.fetch_count(), 1); + + // A new session must honor its startup flag even when a matching disk cache exists. + for enabled in [false, true] { + let restarted = + OpenAiModelsManager::new(home.path().into(), endpoint.clone(), Some(auth.clone())); + restarted.set_api_key_model_discovery_enabled(enabled); + let expected = if enabled { &models } else { &bundled }; + for strategy in [RefreshStrategy::Offline, RefreshStrategy::OnlineIfUncached] { + assert_eq!( + &restarted + .raw_model_catalog(strategy, DEFAULT_HTTP_CLIENT_FACTORY) + .await + .models, + expected + ); + } + assert_eq!(&restarted.try_get_remote_models().unwrap(), expected); + } + assert_eq!(endpoint.fetch_count(), 1); +} diff --git a/codex-rs/models-manager/src/manager.rs b/codex-rs/models-manager/src/manager.rs index cb38919900..c7b1a19ad5 100644 --- a/codex-rs/models-manager/src/manager.rs +++ b/codex-rs/models-manager/src/manager.rs @@ -19,6 +19,8 @@ use std::future::Future; use std::path::PathBuf; use std::pin::Pin; use std::sync::Arc; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; use std::time::Duration; use tokio::sync::RwLock; use tokio::sync::TryLockError; @@ -47,6 +49,11 @@ pub trait ModelsEndpointClient: fmt::Debug + Send + Sync { /// Returns whether the currently resolved auth can use Codex backend-only models. fn uses_codex_backend(&self) -> ModelsEndpointFuture<'_, bool>; + /// Returns whether this provider supports an authoritative catalog with OpenAI API keys. + fn supports_api_key_models(&self) -> bool { + false + } + /// Fetches the latest remote model catalog and optional ETag. fn list_models<'a>( &'a self, @@ -98,6 +105,10 @@ type SharedModelsEndpointClient = Arc; /// Coordinates model discovery plus cached metadata on disk. pub trait ModelsManager: fmt::Debug + Send + Sync { + /// Supply startup API-key discovery policy; live changes require a new session. + /// Static catalogs ignore this setting. + fn set_api_key_model_discovery_enabled(&self, _enabled: bool) {} + /// List all available models, refreshing according to the specified strategy. /// /// Returns model presets sorted by priority and filtered by auth mode and visibility. @@ -235,6 +246,7 @@ pub struct OpenAiModelsManager { remote_models: RwLock, cache: Option>, endpoint_client: SharedModelsEndpointClient, + api_key_model_discovery_enabled: AtomicBool, auth_manager: Option>, } @@ -298,6 +310,7 @@ impl OpenAiModelsManager { models: remote_models, }), cache, + api_key_model_discovery_enabled: AtomicBool::new(false), endpoint_client, auth_manager, } @@ -315,6 +328,11 @@ impl StaticModelsManager { } impl ModelsManager for OpenAiModelsManager { + fn set_api_key_model_discovery_enabled(&self, enabled: bool) { + self.api_key_model_discovery_enabled + .store(enabled, Ordering::SeqCst); + } + fn raw_model_catalog( &self, refresh_strategy: RefreshStrategy, @@ -419,6 +437,12 @@ impl OpenAiModelsManager { refresh_strategy: RefreshStrategy, http_client_factory: &HttpClientFactory, ) -> CoreResult<()> { + // Gate cache loading as well as requests; disabled sessions use bundled API-key models. + if self.supports_api_key_discovery() + && !self.api_key_model_discovery_enabled.load(Ordering::SeqCst) + { + return Ok(()); + } if !self.should_refresh_models().await { if matches!( refresh_strategy, @@ -475,8 +499,19 @@ impl OpenAiModelsManager { Ok(()) } + fn supports_api_key_discovery(&self) -> bool { + self.endpoint_client.supports_api_key_models() + && !self.endpoint_client.has_command_auth() + && self + .auth_manager + .as_ref() + .is_some_and(|auth_manager| auth_manager.auth_mode() == Some(AuthMode::ApiKey)) + } + async fn should_refresh_models(&self) -> bool { - self.endpoint_client.uses_codex_backend().await || self.endpoint_client.has_command_auth() + self.endpoint_client.uses_codex_backend().await + || self.endpoint_client.has_command_auth() + || self.supports_api_key_discovery() } /// Publish only while the request identity still matches, including after async storage. @@ -485,16 +520,17 @@ impl OpenAiModelsManager { if entry.identity != self.endpoint_client.identity() { return false; } - // Visible ChatGPT catalogs are authoritative; other catalogs augment bundled metadata. + // Visible ChatGPT and OpenAI API-key catalogs are authoritative. let remote_only = entry .models .iter() .any(|model| model.visibility == ModelVisibility::List) - && self.auth_manager.as_ref().is_some_and(|auth_manager| { - auth_manager - .auth_mode() - .is_some_and(AuthMode::has_chatgpt_account) - }); + && (self.supports_api_key_discovery() + || self.auth_manager.as_ref().is_some_and(|auth_manager| { + auth_manager + .auth_mode() + .is_some_and(AuthMode::has_chatgpt_account) + })); if !remote_only { let mut models = load_remote_models_from_file().unwrap_or_default(); for model in entry.models { diff --git a/codex-rs/models-manager/src/manager_tests.rs b/codex-rs/models-manager/src/manager_tests.rs index b3f732c0ae..1408ad8928 100644 --- a/codex-rs/models-manager/src/manager_tests.rs +++ b/codex-rs/models-manager/src/manager_tests.rs @@ -27,6 +27,9 @@ use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use tempfile::tempdir; +#[path = "api_key_discovery_tests.rs"] +mod api_key_discovery_tests; + #[path = "cache_identity_tests.rs"] mod cache_identity_tests; @@ -268,6 +271,10 @@ impl ExternalAuth for TestUnresolvedExternalApiKeyAuth { } impl ModelsEndpointClient for TestModelsEndpoint { + fn supports_api_key_models(&self) -> bool { + true + } + fn identity(&self) -> Option { Some("test-provider".to_string()) } @@ -557,6 +564,7 @@ async fn injected_cache_ttl_refresh_preserves_cached_payload() { "test-api-key", ))), ); + manager.set_api_key_model_discovery_enabled(/*enabled*/ true); manager .raw_model_catalog( @@ -1025,7 +1033,7 @@ async fn refresh_available_models_merges_hidden_only_chatgpt_remote_with_bundled } #[tokio::test] -async fn refresh_available_models_keeps_merging_for_api_auth() { +async fn refresh_available_models_keeps_merging_for_custom_api_auth() { let remote_models = vec![remote_model( "api-auth-visible-remote", "API Auth Visible", @@ -1227,7 +1235,7 @@ async fn refresh_available_models_drops_removed_remote_models() { } #[tokio::test] -async fn refresh_available_models_skips_network_without_chatgpt_auth() { +async fn refresh_available_models_skips_network_without_auth() { let dynamic_slug = "dynamic-model-only-for-test-noauth"; let codex_home = tempdir().expect("temp dir"); let endpoint = TestModelsEndpoint::without_refresh(vec![vec![remote_model( @@ -1244,13 +1252,13 @@ async fn refresh_available_models_skips_network_without_chatgpt_auth() { manager .refresh_available_models(RefreshStrategy::Online, &DEFAULT_HTTP_CLIENT_FACTORY) .await - .expect("refresh should no-op without chatgpt auth"); + .expect("refresh should no-op without auth"); let cached_remote = manager.get_remote_models().await; assert!( !cached_remote .iter() .any(|candidate| candidate.slug == dynamic_slug), - "remote refresh should be skipped without chatgpt auth" + "remote refresh should be skipped without auth" ); assert_eq!( endpoint.fetch_count(), @@ -1307,6 +1315,10 @@ impl TestAuthAwareModelsEndpoint { } impl ModelsEndpointClient for TestAuthAwareModelsEndpoint { + fn supports_api_key_models(&self) -> bool { + true + } + fn identity(&self) -> Option { self.auth_manager .as_ref() @@ -1332,7 +1344,7 @@ impl ModelsEndpointClient for TestAuthAwareModelsEndpoint { } #[tokio::test] -async fn refresh_available_models_skips_network_when_external_api_key_overrides_chatgpt_auth() { +async fn refresh_available_models_fetches_when_external_api_key_overrides_chatgpt_auth() { let dynamic_slug = "dynamic-model-only-for-test-external-api-key"; let codex_home = tempdir().expect("temp dir"); let auth_manager = @@ -1354,24 +1366,23 @@ async fn refresh_available_models_skips_network_when_external_api_key_overrides_ endpoint.clone(), Some(auth_manager), ); + manager.set_api_key_model_discovery_enabled(/*enabled*/ true); manager .refresh_available_models(RefreshStrategy::Online, &DEFAULT_HTTP_CLIENT_FACTORY) .await - .expect("refresh should no-op with API key auth"); + .expect("refresh should fetch with API key auth"); let cached_remote = manager.get_remote_models().await; - assert!( - !cached_remote - .iter() - .any(|candidate| candidate.slug == dynamic_slug), - "remote refresh should be skipped when external API key auth is active" - ); assert_eq!( - endpoint.fetch_count(), - 0, - "endpoint should avoid model fetches when external API key auth is active" + cached_remote, + vec![remote_model( + dynamic_slug, + "External API Key", + /*priority*/ 1 + )] ); + assert_eq!(endpoint.fetch_count(), 1); } #[tokio::test] @@ -1397,6 +1408,7 @@ async fn refresh_available_models_uses_cached_chatgpt_when_external_api_key_is_u endpoint.clone(), Some(auth_manager), ); + manager.set_api_key_model_discovery_enabled(/*enabled*/ true); manager .refresh_available_models(RefreshStrategy::Online, &DEFAULT_HTTP_CLIENT_FACTORY)