mirror of
https://github.com/openai/codex.git
synced 2026-10-27 18:48:37 +00:00
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:
@@ -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();
|
||||
}
|
||||
|
||||
@@ -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<bool>,
|
||||
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::<ModelInfo>(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::<Result<Vec<_>, _>>()?;
|
||||
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(())
|
||||
}
|
||||
|
||||
|
||||
@@ -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"
|
||||
},
|
||||
|
||||
@@ -383,10 +383,14 @@ pub fn build_models_manager(
|
||||
auth_manager: Arc<AuthManager>,
|
||||
) -> 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(
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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);
|
||||
})
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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<String> {
|
||||
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<Option<String>>);
|
||||
|
||||
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 {
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
@@ -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,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)
|
||||
|
||||
Reference in New Issue
Block a user