Add opt-in model discovery for OpenAI API keys (#44392)

## What changed

Add the default-off `api_key_model_discovery` feature and expose it through app-server experimental feature enablement, respecting user configuration overrides.

When enabled, fetch OpenAI API-key model metadata from the Codex backend, honoring explicit provider base URLs. Treat remote catalogs containing visible models as authoritative. When disabled at startup, skip discovery requests and cached catalogs and retain bundled models.

## Testing

Add coverage for feature enablement and user overrides, cached catalogs across sessions, provider routing, and API-key `model/list` results. Extend coverage for remote instructions and auth-specific cache isolation.

GitOrigin-RevId: 4b189c37c615a7f808cb6ff5f0858a9a55f36ebc
This commit is contained in:
andrewgu-oai
2026-09-10 03:18:28 +00:00
committed by copyberry
parent 5a9eb145c4
commit ea53c8d4f7
13 changed files with 465 additions and 54 deletions
@@ -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);
}
+43 -7
View File
@@ -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<dyn ModelsEndpointClient>;
/// 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<ModelsCacheEntry>,
cache: Option<Arc<dyn ModelsCache>>,
endpoint_client: SharedModelsEndpointClient,
api_key_model_discovery_enabled: AtomicBool,
auth_manager: Option<Arc<AuthManager>>,
}
@@ -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 {
+27 -15
View File
@@ -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<String> {
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<String> {
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)