diff --git a/codex-rs/app-server/tests/suite/v2/analytics.rs b/codex-rs/app-server/tests/suite/v2/analytics.rs index f20e9bf190..9a280c992d 100644 --- a/codex-rs/app-server/tests/suite/v2/analytics.rs +++ b/codex-rs/app-server/tests/suite/v2/analytics.rs @@ -132,8 +132,17 @@ async fn guardian_review_turns_and_tools_reach_analytics() -> Result<()> { let server = responses::start_mock_server().await; let codex_home = TempDir::new()?; mount_analytics_capture(&server, codex_home.path()).await?; + let catalog_path = codex_home.path().join("guardian-models.json"); + std::fs::write( + &catalog_path, + serde_json::to_vec(&bundled_models_response()?)?, + )?; MockResponsesConfig::new(&server.uri()) - .with_root_config(&format!("chatgpt_base_url = \"{}\"", server.uri())) + .with_root_config(&format!( + "chatgpt_base_url = \"{}\"\nmodel_catalog_json = {}", + server.uri(), + serde_json::to_string(&catalog_path)? + )) .with_provider_config("supports_websockets = false") .enable_feature(Feature::GuardianApproval) .disable_feature(Feature::Apps) diff --git a/codex-rs/core/src/guardian/reviewer_config.rs b/codex-rs/core/src/guardian/reviewer_config.rs index 7d0676ac59..c89b2da437 100644 --- a/codex-rs/core/src/guardian/reviewer_config.rs +++ b/codex-rs/core/src/guardian/reviewer_config.rs @@ -78,14 +78,23 @@ pub(crate) async fn resolve_review_model( Arc, ) { let turn = context.turn(); - let available_models = session - .services - .models_manager - .list_models( + let models_manager = &session.services.models_manager; + let catalog = models_manager + .model_catalog_with_origins( codex_models_manager::manager::RefreshStrategy::Offline, turn.config.http_client_factory(), ) .await; + let available_models = models_manager.build_available_models( + catalog + .into_iter() + .filter(|entry| { + context.model_info.auto_review_model_override.is_some() + || entry.origin != codex_models_manager::ModelOrigin::Bundled + }) + .map(|entry| entry.model) + .collect(), + ); let default_review_model_id = turn.provider.approval_review_preferred_model(); let review_model = codex_guardian_reviewer::select_review_model( &context.model_info, diff --git a/codex-rs/core/src/guardian/tests.rs b/codex-rs/core/src/guardian/tests.rs index be96a54d5a..bd0ef336ab 100644 --- a/codex-rs/core/src/guardian/tests.rs +++ b/codex-rs/core/src/guardian/tests.rs @@ -1742,6 +1742,7 @@ fn build_guardian_transcript_preserves_recent_tool_context_when_user_history_is_ enum GuardianTestCatalog { Bundled, + Explicit, ParentOnly, } @@ -1788,6 +1789,7 @@ async fn guardian_reuse_respects_effective_policy_and_personality_opt_out( .models_manager = Arc::new(StaticModelsManager::new( /*auth_manager*/ None, ModelsResponse { models: vec![] }, + codex_models_manager::ModelOrigin::ExplicitConfig, )); seed_guardian_parent_history(&session, &turn).await; @@ -1874,6 +1876,18 @@ async fn guardian_request_model_for_auto_review( let (mut session, mut turn) = guardian_test_session_and_turn(&server).await; match catalog { GuardianTestCatalog::Bundled => {} + GuardianTestCatalog::Explicit => { + let auth_manager = Arc::clone(&session.services.auth_manager); + Arc::get_mut(&mut session) + .expect("unique session") + .services + .models_manager = Arc::new(StaticModelsManager::new( + Some(auth_manager), + codex_models_manager::bundled_models_response().expect("valid catalog"), + codex_models_manager::ModelOrigin::ExplicitConfig, + )); + crate::guardian::test_host::install(&session, &turn.config); + } GuardianTestCatalog::ParentOnly => { let parent_model = turn.model_info().as_ref().clone(); let auth_manager = Arc::clone(&session.services.auth_manager); @@ -1882,6 +1896,7 @@ async fn guardian_request_model_for_auto_review( ModelsResponse { models: vec![parent_model], }, + codex_models_manager::ModelOrigin::ExplicitConfig, ); Arc::get_mut(&mut session) .expect("session should be unique") @@ -1993,7 +2008,7 @@ async fn guardian_review_uses_preferred_review_model_without_model_catalog_overr let (request_model, parent_model, preferred_model, analytics_result) = guardian_request_model_for_auto_review( /*auto_review_model_override*/ None, - GuardianTestCatalog::Bundled, + GuardianTestCatalog::Explicit, ) .await?; @@ -2387,6 +2402,7 @@ async fn guardian_reuses_prompt_cache_key_and_appends_prior_reviews() -> anyhow: ModelsResponse { models: vec![reviewer_model], }, + codex_models_manager::ModelOrigin::ExplicitConfig, )); let turn_mut = Arc::get_mut(&mut turn).expect("turn should be unique"); update_turn_settings_for_test(turn_mut, |settings| { @@ -3832,7 +3848,17 @@ async fn guardian_review_session_config_preserves_parent_network_proxy() { #[tokio::test] async fn guardian_review_session_config_clears_context_overrides_for_distinct_effective_model() { let server = start_mock_server().await; - let (session, mut turn) = guardian_test_session_and_turn(&server).await; + let (mut session, mut turn) = guardian_test_session_and_turn(&server).await; + let auth_manager = Arc::clone(&session.services.auth_manager); + Arc::get_mut(&mut session) + .expect("unique session") + .services + .models_manager = Arc::new(StaticModelsManager::new( + Some(auth_manager), + codex_models_manager::bundled_models_response().expect("valid catalog"), + codex_models_manager::ModelOrigin::ExplicitConfig, + )); + crate::guardian::test_host::install(&session, &turn.config); let mut config = (*turn.config).clone(); config.model = Some("codex-auto-review".to_string()); config.model_context_window = Some(900_000); @@ -3870,6 +3896,7 @@ async fn guardian_review_session_config_preserves_context_overrides_for_same_eff ModelsResponse { models: vec![parent_model], }, + codex_models_manager::ModelOrigin::ExplicitConfig, )); let mut config = (*turn.config).clone(); config.model = Some("stale-parent-model".to_string()); diff --git a/codex-rs/core/src/session/step_activation_tests.rs b/codex-rs/core/src/session/step_activation_tests.rs index 9f558273b6..094b59de58 100644 --- a/codex-rs/core/src/session/step_activation_tests.rs +++ b/codex-rs/core/src/session/step_activation_tests.rs @@ -348,6 +348,7 @@ impl GatedModelsManager { inner: StaticModelsManager::new( /*auth_manager*/ None, ModelsResponse { models }, + codex_models_manager::ModelOrigin::ExplicitConfig, ), first_b_lookup: StdMutex::new(Some(Arc::clone(&lookup))), }), @@ -357,12 +358,12 @@ impl GatedModelsManager { } impl ModelsManager for GatedModelsManager { - fn raw_model_catalog( + fn model_catalog_with_origins( &self, strategy: RefreshStrategy, factory: HttpClientFactory, - ) -> ModelsManagerFuture<'_, ModelsResponse> { - self.inner.raw_model_catalog(strategy, factory) + ) -> ModelsManagerFuture<'_, Vec> { + self.inner.model_catalog_with_origins(strategy, factory) } fn get_remote_models(&self) -> ModelsManagerFuture<'_, Vec> { diff --git a/codex-rs/core/src/session/step_settings_tests.rs b/codex-rs/core/src/session/step_settings_tests.rs index 86b8ac13dd..049ca1b428 100644 --- a/codex-rs/core/src/session/step_settings_tests.rs +++ b/codex-rs/core/src/session/step_settings_tests.rs @@ -436,7 +436,11 @@ async fn model_resolution_preserves_startup_overrides_and_instruction_provenance let catalog = ModelsResponse { models: vec![model], }; - let models_manager = StaticModelsManager::new(/*auth_manager*/ None, catalog.clone()); + let models_manager = StaticModelsManager::new( + /*auth_manager*/ None, + catalog.clone(), + codex_models_manager::ModelOrigin::ExplicitConfig, + ); let mut config = test_config().await; config.model = Some("model-a".to_string()); diff --git a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs index 9582a977e0..d5a8dcc9eb 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs @@ -573,6 +573,7 @@ async fn spawn_agent_service_tier_inheritance_uses_root_preference_and_child_mod session.services.models_manager = Arc::new(StaticModelsManager::new( /*auth_manager*/ None, service_tier_test_catalog(), + codex_models_manager::ModelOrigin::ExplicitConfig, )); session.thread_id = root.thread_id; @@ -626,6 +627,7 @@ async fn spawn_agent_service_tier_inheritance_uses_root_preference_and_child_mod session.services.models_manager = Arc::new(StaticModelsManager::new( /*auth_manager*/ None, service_tier_test_catalog(), + codex_models_manager::ModelOrigin::ExplicitConfig, )); session.thread_id = root.thread_id; diff --git a/codex-rs/core/tests/suite/guardian_review.rs b/codex-rs/core/tests/suite/guardian_review.rs index 4ca00c69f8..eef8bb9183 100644 --- a/codex-rs/core/tests/suite/guardian_review.rs +++ b/codex-rs/core/tests/suite/guardian_review.rs @@ -26,6 +26,7 @@ use codex_features::Feature; use codex_history::RolloutItem; use codex_login::AuthHeaders; use codex_login::CodexAuth; +use codex_models_manager::ModelOrigin; use codex_models_manager::manager::RefreshStrategy; use codex_otel::MetricsClient; use codex_otel::MetricsConfig; @@ -96,6 +97,7 @@ use wiremock::Mock; use wiremock::ResponseTemplate; use wiremock::http::Method; use wiremock::matchers::method; +use wiremock::matchers::path; use wiremock::matchers::path_regex; use super::guardian_mcp_elicitation::guardian_blocking_samples; @@ -208,6 +210,8 @@ async fn guardian_session_inherits_parent_http_fallback( config.approvals_reviewer = ApprovalsReviewer::User; config.guardian_policy_config = Some(configured_policy.to_string()); config.guardian_policy_template = Some(configured_template.to_string()); + config.model_catalog = + Some(codex_models_manager::bundled_models_response().expect("valid catalog")); }); let test = builder.build_with_auto_env(&server).await?; @@ -313,6 +317,83 @@ async fn guardian_session_inherits_parent_http_fallback( Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[test_case(ModelOrigin::Bundled, "gateway-main"; "bundled_reviewer_falls_back")] +#[test_case(ModelOrigin::Provider, "codex-auto-review"; "advertised_reviewer")] +#[test_case(ModelOrigin::ExplicitConfig, "codex-auto-review"; "configured_reviewer")] +async fn guardian_custom_gateway_selects_reviewer_by_catalog_origin( + origin: ModelOrigin, + expected_model: &str, +) -> Result<()> { + skip_if_no_network!(Ok(())); + let server = start_mock_server().await; + let mut reviewer_model = codex_models_manager::bundled_models_response()? + .models + .into_iter() + .find(|model| model.slug == "codex-auto-review") + .expect("reviewer fixture"); + reviewer_model.use_responses_lite = false; + let catalog = ModelsResponse { + models: vec![ + codex_models_manager::model_info::model_info_from_slug("gateway-main"), + reviewer_model, + ], + }; + Mock::given(method("GET")) + .and(path("/catalog")) + .respond_with(ResponseTemplate::new(200).set_body_json(&catalog)) + .mount(&server) + .await; + let catalog_url = format!("{}/catalog", server.uri()); + let responses = mount_sse_sequence(&server, vec![ + sse(vec![ev_response_created("parent"), ev_function_call( + "action", "exec_command", + r#"{"cmd":"true","sandbox_permissions":"require_escalated","justification":"test"}"#, + ), ev_completed("parent")]), + sse(vec![ev_response_created("review"), ev_assistant_message( + "decision", r#"{"outcome":"deny"}"#, + ), ev_completed("review")]), + sse(vec![ev_completed("done")]), + ]).await; + let test = test_codex() + .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) + .with_config(move |config| { + config.model = Some("gateway-main".to_string()); + config.model_catalog = match origin { + ModelOrigin::ExplicitConfig => Some(catalog), + ModelOrigin::Bundled | ModelOrigin::Provider => None, + }; + config.model_provider.model_catalog_url = match origin { + ModelOrigin::Provider => Some(catalog_url.into()), + ModelOrigin::Bundled | ModelOrigin::ExplicitConfig => None, + }; + config.model_provider.name = "Gateway".to_string(); + config.model_provider.supports_websockets = false; + config.permissions.approval_policy = Constrained::allow_any(AskForApproval::OnRequest); + config.approvals_reviewer = ApprovalsReviewer::AutoReview; + }) + .build_with_auto_env(&server) + .await?; + test.codex + .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text { + text: "run a command".to_string(), + text_elements: vec![], + }])) + .await?; + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + let requests = responses.requests(); + let reviewer = requests + .iter() + .find(|request| request.body_json()["client_metadata"]["x-openai-subagent"] == "guardian") + .expect("Guardian must review the action"); + assert_eq!(reviewer.body_json()["model"], expected_model); + test.codex.shutdown_and_wait().await?; + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn guardian_parent_id_survives_code_mode_response_handoff() -> Result<()> { skip_if_no_network!(Ok(())); @@ -421,6 +502,8 @@ for (const phase of ["before", "after"]) { config.model_provider.supports_websockets = false; config.permissions.approval_policy = Constrained::allow_any(AskForApproval::OnRequest); config.approvals_reviewer = ApprovalsReviewer::User; + config.model_catalog = + Some(codex_models_manager::bundled_models_response().expect("valid catalog")); }) .build_with_auto_env(&server) .await?; diff --git a/codex-rs/core/tests/suite/guardian_subagent_authorization.rs b/codex-rs/core/tests/suite/guardian_subagent_authorization.rs index a11d977bfd..74c876f4f4 100644 --- a/codex-rs/core/tests/suite/guardian_subagent_authorization.rs +++ b/codex-rs/core/tests/suite/guardian_subagent_authorization.rs @@ -373,6 +373,8 @@ async fn code_mode_messaging_fixture( } config.permissions.approval_policy = Constrained::allow_any(AskForApproval::OnRequest); config.approvals_reviewer = ApprovalsReviewer::AutoReview; + config.model_catalog = + Some(codex_models_manager::bundled_models_response().expect("valid catalog")); config .permissions .set_permission_profile(PermissionProfile::workspace_write()) diff --git a/codex-rs/core/tests/suite/prompt_caching.rs b/codex-rs/core/tests/suite/prompt_caching.rs index 6cf6691cee..1042bf8b6e 100644 --- a/codex-rs/core/tests/suite/prompt_caching.rs +++ b/codex-rs/core/tests/suite/prompt_caching.rs @@ -158,6 +158,7 @@ async fn prompt_tools_are_consistent_across_requests( ModelsResponse { models: vec![model], }, + codex_models_manager::ModelOrigin::ExplicitConfig, )); let server = start_mock_server().await; diff --git a/codex-rs/core/tests/suite/wake_reservation.rs b/codex-rs/core/tests/suite/wake_reservation.rs index e6a44d0395..6ede021843 100644 --- a/codex-rs/core/tests/suite/wake_reservation.rs +++ b/codex-rs/core/tests/suite/wake_reservation.rs @@ -24,7 +24,6 @@ use codex_models_manager::manager::StaticModelsManager; use codex_protocol::AgentPath; use codex_protocol::config_types::CollaborationModeMask; use codex_protocol::openai_models::ModelInfo; -use codex_protocol::openai_models::ModelsResponse; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::InterAgentCommunication; use codex_protocol::protocol::Op; @@ -50,12 +49,12 @@ struct GatedModelsManager { } impl ModelsManager for GatedModelsManager { - fn raw_model_catalog( + fn model_catalog_with_origins( &self, strategy: RefreshStrategy, factory: HttpClientFactory, - ) -> ModelsManagerFuture<'_, ModelsResponse> { - self.inner.raw_model_catalog(strategy, factory) + ) -> ModelsManagerFuture<'_, Vec> { + self.inner.model_catalog_with_origins(strategy, factory) } fn get_remote_models(&self) -> ModelsManagerFuture<'_, Vec> { @@ -145,7 +144,11 @@ async fn interrupt_during_wake_context_construction_does_not_start_turn() -> any let server = responses::start_mock_server().await; let requests = responses::mount_sse_once(&server, responses::sse_completed("initial")).await; let models = Arc::new(GatedModelsManager { - inner: StaticModelsManager::new(/*auth_manager*/ None, bundled_models_response()?), + inner: StaticModelsManager::new( + /*auth_manager*/ None, + bundled_models_response()?, + codex_models_manager::ModelOrigin::ExplicitConfig, + ), gate_next_lookup: AtomicBool::new(/*v*/ false), context_build: Gate::default(), }); diff --git a/codex-rs/model-provider/src/amazon_bedrock/mod.rs b/codex-rs/model-provider/src/amazon_bedrock/mod.rs index 3bd18668b8..a441d2704a 100644 --- a/codex-rs/model-provider/src/amazon_bedrock/mod.rs +++ b/codex-rs/model-provider/src/amazon_bedrock/mod.rs @@ -430,10 +430,16 @@ impl ModelProvider for AmazonBedrockModelProvider { _codex_home: PathBuf, config_model_catalog: Option, ) -> SharedModelsManager { + let origin = if config_model_catalog.is_some() { + codex_models_manager::ModelOrigin::ExplicitConfig + } else { + codex_models_manager::ModelOrigin::Provider + }; Arc::new(StaticModelsManager::new( /*auth_manager*/ None, config_model_catalog .map_or_else(|| self.default_model_catalog(), normalize_bedrock_catalog), + origin, )) } @@ -441,11 +447,7 @@ impl ModelProvider for AmazonBedrockModelProvider { &self, config_model_catalog: Option, ) -> SharedModelsManager { - Arc::new(StaticModelsManager::new( - /*auth_manager*/ None, - config_model_catalog - .map_or_else(|| self.default_model_catalog(), normalize_bedrock_catalog), - )) + self.models_manager(PathBuf::new(), config_model_catalog) } } diff --git a/codex-rs/model-provider/src/gcp/mod.rs b/codex-rs/model-provider/src/gcp/mod.rs index 2b9df8e62a..1367b24033 100644 --- a/codex-rs/model-provider/src/gcp/mod.rs +++ b/codex-rs/model-provider/src/gcp/mod.rs @@ -149,6 +149,11 @@ impl ModelProvider for GcpModelProvider { &self, config_model_catalog: Option, ) -> SharedModelsManager { + let origin = if config_model_catalog.is_some() { + codex_models_manager::ModelOrigin::ExplicitConfig + } else { + codex_models_manager::ModelOrigin::Provider + }; let mut model_catalog = config_model_catalog .or_else(|| codex_models_manager::bundled_models_response().ok()) .unwrap_or_default(); @@ -168,6 +173,7 @@ impl ModelProvider for GcpModelProvider { Arc::new(StaticModelsManager::new( self.auth_manager.clone(), model_catalog, + origin, )) } } diff --git a/codex-rs/model-provider/src/provider.rs b/codex-rs/model-provider/src/provider.rs index 5bd223b949..f455ab60b6 100644 --- a/codex-rs/model-provider/src/provider.rs +++ b/codex-rs/model-provider/src/provider.rs @@ -370,10 +370,19 @@ pub trait ModelProvider: fmt::Debug + Send + Sync { &self, config_model_catalog: Option, ) -> SharedModelsManager { + let origin = if config_model_catalog.is_some() { + codex_models_manager::ModelOrigin::ExplicitConfig + } else { + codex_models_manager::ModelOrigin::Bundled + }; let model_catalog = config_model_catalog .or_else(|| codex_models_manager::bundled_models_response().ok()) .unwrap_or_default(); - Arc::new(StaticModelsManager::new(self.auth_manager(), model_catalog)) + Arc::new(StaticModelsManager::new( + self.auth_manager(), + model_catalog, + origin, + )) } /// Creates a model manager that can use a caller-provided cache for remote catalogs. @@ -546,6 +555,7 @@ impl ConfiguredModelProvider { return Arc::new(StaticModelsManager::new( self.auth_manager.clone(), model_catalog, + codex_models_manager::ModelOrigin::ExplicitConfig, )); } let endpoint = Arc::new( diff --git a/codex-rs/models-manager/src/catalog.rs b/codex-rs/models-manager/src/catalog.rs new file mode 100644 index 0000000000..65873678db --- /dev/null +++ b/codex-rs/models-manager/src/catalog.rs @@ -0,0 +1,83 @@ +//! Tracks the source of model metadata without changing the catalog wire or cache formats. + +use std::borrow::Cow; +use std::collections::HashSet; + +use codex_protocol::openai_models::ModelInfo; + +use crate::cache::ModelsCacheEntry; + +/// Where an effective catalog entry was obtained. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ModelOrigin { + Bundled, + Provider, + ExplicitConfig, +} + +/// Model metadata and its source captured from the same catalog snapshot. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CatalogModel { + pub model: ModelInfo, + pub origin: ModelOrigin, +} + +/// Catalog authority is independent of whether its model metadata is borrowed or owned. +pub(crate) enum SelectedCatalog<'a> { + CurrentSnapshot(Cow<'a, [ModelInfo]>), + Fallback(Cow<'a, [ModelInfo]>), +} + +impl<'a> SelectedCatalog<'a> { + pub fn into_models(self) -> Cow<'a, [ModelInfo]> { + match self { + Self::CurrentSnapshot(models) | Self::Fallback(models) => models, + } + } +} + +/// The effective catalog and the IDs in the unmerged provider response share one lock. +#[derive(Debug)] +pub(crate) struct CatalogState { + pub entry: ModelsCacheEntry, + pub provider_model_ids: HashSet, +} + +impl CatalogState { + pub fn bundled(entry: ModelsCacheEntry) -> Self { + Self { + entry, + provider_model_ids: HashSet::new(), + } + } + + pub fn provider(entry: ModelsCacheEntry) -> Self { + // Capture raw response IDs before merging, including when loading from disk cache. + let provider_model_ids = entry + .models + .iter() + .map(|model| model.slug.clone()) + .collect(); + Self { + entry, + provider_model_ids, + } + } + + pub fn model_catalog(&self, selected: SelectedCatalog<'_>) -> Vec { + let uses_snapshot_ids = matches!(selected, SelectedCatalog::CurrentSnapshot(_)); + selected + .into_models() + .into_owned() + .into_iter() + .map(|model| { + let origin = if uses_snapshot_ids && self.provider_model_ids.contains(&model.slug) { + ModelOrigin::Provider + } else { + ModelOrigin::Bundled + }; + CatalogModel { model, origin } + }) + .collect() + } +} diff --git a/codex-rs/models-manager/src/catalog_origin_tests.rs b/codex-rs/models-manager/src/catalog_origin_tests.rs new file mode 100644 index 0000000000..eb7f2fa3f9 --- /dev/null +++ b/codex-rs/models-manager/src/catalog_origin_tests.rs @@ -0,0 +1,96 @@ +#![cfg(test)] +//! Catalog origins survive merges, cache reloads, and changes of provider identity. + +use super::*; +use pretty_assertions::assert_eq; + +#[tokio::test] +async fn catalog_origins_survive_reload_and_auth_change() { + for (slug, reviewer_origin) in [ + ("codex-auto-review", ModelOrigin::Provider), + ("gateway-model", ModelOrigin::Bundled), + ] { + let home = tempdir().unwrap(); + let model = remote_model_with_visibility(slug, "Gateway", /*priority*/ 0, "hide"); + let auth = + AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); + let endpoint = + TestAuthAwareModelsEndpoint::new(Some(auth.clone()), vec![vec![model.clone()]]); + let manager = + OpenAiModelsManager::new(home.path().into(), endpoint.clone(), Some(auth.clone())); + manager.set_api_key_model_discovery_enabled(/*enabled*/ true); + let bundled = manager + .model_catalog_with_origins(RefreshStrategy::Offline, DEFAULT_HTTP_CLIENT_FACTORY) + .await; + let catalog = manager + .model_catalog_with_origins(RefreshStrategy::Online, DEFAULT_HTTP_CLIENT_FACTORY) + .await; + assert!(catalog.contains(&CatalogModel { + model, + origin: ModelOrigin::Provider + })); + assert_eq!( + catalog + .iter() + .find(|entry| entry.model.slug == "codex-auto-review") + .map(|entry| entry.origin), + Some(reviewer_origin) + ); + + let restarted = openai_manager_for_tests(home.path().to_path_buf(), endpoint.clone()); + assert_eq!( + restarted + .model_catalog_with_origins(RefreshStrategy::Offline, DEFAULT_HTTP_CLIENT_FACTORY) + .await, + catalog + ); + assert_eq!(endpoint.fetch_count(), 1); + + // An empty subsequent response restores bundled origins instead of retaining old IDs. + assert_eq!( + restarted + .model_catalog_with_origins(RefreshStrategy::Online, DEFAULT_HTTP_CLIENT_FACTORY) + .await, + bundled + ); + + // The original manager still holds provider IDs, but its identity is no longer current. + auth.set_external_auth(Arc::new(TestExternalApiKeyAuth)) + .await + .unwrap(); + assert_eq!( + manager + .model_catalog_with_origins(RefreshStrategy::Offline, DEFAULT_HTTP_CLIENT_FACTORY) + .await, + bundled + ); + assert_eq!(endpoint.fetch_count(), 2); + } +} + +#[tokio::test] +async fn static_catalog_records_its_supplied_origin() { + let model = remote_model("reviewer", "Reviewer", /*priority*/ 0); + for origin in [ + ModelOrigin::Bundled, + ModelOrigin::Provider, + ModelOrigin::ExplicitConfig, + ] { + let manager = StaticModelsManager::new( + /*auth_manager*/ None, + ModelsResponse { + models: vec![model.clone()], + }, + origin, + ); + assert_eq!( + manager + .model_catalog_with_origins(RefreshStrategy::Offline, DEFAULT_HTTP_CLIENT_FACTORY) + .await, + vec![CatalogModel { + model: model.clone(), + origin + }] + ); + } +} diff --git a/codex-rs/models-manager/src/lib.rs b/codex-rs/models-manager/src/lib.rs index f3da87e511..a35deb042f 100644 --- a/codex-rs/models-manager/src/lib.rs +++ b/codex-rs/models-manager/src/lib.rs @@ -1,4 +1,5 @@ pub mod cache; +mod catalog; pub mod collaboration_mode_presets; pub(crate) mod config; pub mod manager; @@ -7,6 +8,8 @@ pub mod model_presets; mod originator; pub mod test_support; +pub use catalog::CatalogModel; +pub use catalog::ModelOrigin; pub use codex_protocol::auth::AuthMode; pub use config::ModelsManagerConfig; pub use originator::ModelCatalogOriginator; diff --git a/codex-rs/models-manager/src/manager.rs b/codex-rs/models-manager/src/manager.rs index 6fba03dab8..3df77b5cff 100644 --- a/codex-rs/models-manager/src/manager.rs +++ b/codex-rs/models-manager/src/manager.rs @@ -2,6 +2,10 @@ use super::cache::FileModelsCache; use crate::ModelCatalogOriginator; use crate::cache::ModelsCache; use crate::cache::ModelsCacheEntry; +use crate::catalog::CatalogModel; +use crate::catalog::CatalogState; +use crate::catalog::ModelOrigin; +use crate::catalog::SelectedCatalog; use crate::collaboration_mode_presets::builtin_collaboration_mode_presets; use crate::config::ModelsManagerConfig; use crate::model_info; @@ -158,12 +162,31 @@ pub trait ModelsManager: fmt::Debug + Send + Sync { ) } - /// Return the active raw model catalog, refreshing according to the specified strategy. + /// Project model metadata from the source-aware read, using its refresh and fallback policy. fn raw_model_catalog( &self, refresh_strategy: RefreshStrategy, http_client_factory: HttpClientFactory, - ) -> ModelsManagerFuture<'_, ModelsResponse>; + ) -> ModelsManagerFuture<'_, ModelsResponse> { + Box::pin(async move { + ModelsResponse { + models: self + .model_catalog_with_origins(refresh_strategy, http_client_factory) + .await + .into_iter() + .map(|entry| entry.model) + .collect(), + } + }) + } + + /// Return model metadata and its locally determined source from the same snapshot. + /// Implement refresh, error handling, and fallback here for both catalog views. + fn model_catalog_with_origins( + &self, + refresh_strategy: RefreshStrategy, + http_client_factory: HttpClientFactory, + ) -> ModelsManagerFuture<'_, Vec>; /// Best-effort refresh when the in-memory catalog belongs to different credentials. /// Static catalogs need no refresh. Failures follow the catalog's fallback policy. @@ -304,7 +327,7 @@ impl CatalogSource { #[derive(Debug)] pub struct OpenAiModelsManager { originator_managers: std::sync::Mutex>, - remote_models: RwLock, + remote_models: RwLock, cache: Option>, endpoint_client: SharedModelsEndpointClient, catalog_source: CatalogSource, @@ -317,6 +340,7 @@ pub struct OpenAiModelsManager { pub struct StaticModelsManager { remote_models: Vec, auth_manager: Option>, + origin: ModelOrigin, } impl OpenAiModelsManager { @@ -364,13 +388,13 @@ impl OpenAiModelsManager { ) -> Self { Self { originator_managers: std::sync::Mutex::new(HashMap::new()), - remote_models: RwLock::new(ModelsCacheEntry { + remote_models: RwLock::new(CatalogState::bundled(ModelsCacheEntry { fetched_at: Utc::now(), etag: None, client_version: Some(crate::client_version_to_whole()), identity: endpoint_client.identity(), models: load_remote_models_from_file().unwrap_or_default(), - }), + })), cache, api_key_model_discovery_enabled: Arc::new(AtomicBool::new(false)), endpoint_client, @@ -383,22 +407,43 @@ impl OpenAiModelsManager { /// Clear the bundled seed before sharing it; the existing auth identity must be retained. pub fn with_provider_catalog(mut self) -> Self { self.catalog_source = CatalogSource::ExplicitProvider(Mutex::new(())); - self.remote_models.get_mut().models.clear(); + self.remote_models.get_mut().entry.models.clear(); self } } impl StaticModelsManager { - /// Construct a static model manager from an authoritative catalog. - pub fn new(auth_manager: Option>, model_catalog: ModelsResponse) -> Self { + /// Construct a static model manager with the source of its supplied catalog. + pub fn new( + auth_manager: Option>, + model_catalog: ModelsResponse, + origin: ModelOrigin, + ) -> Self { Self { remote_models: model_catalog.models, auth_manager, + origin, } } } impl ModelsManager for OpenAiModelsManager { + fn model_catalog_with_origins( + &self, + refresh_strategy: RefreshStrategy, + http_client_factory: HttpClientFactory, + ) -> ModelsManagerFuture<'_, Vec> { + Box::pin(async move { + if let Err(err) = self + .refresh_available_models(refresh_strategy, &http_client_factory) + .await + { + error!("failed to refresh available models: {err}"); + } + let state = self.remote_models.read().await; + state.model_catalog(self.models_for_entry(&state.entry)) + }) + } fn for_originator(&self, originator: &str) -> Option { let originator = ModelCatalogOriginator::from_originator(originator)?; let mut managers = self @@ -429,18 +474,6 @@ impl ModelsManager for OpenAiModelsManager { .store(enabled, Ordering::SeqCst); } - fn raw_model_catalog( - &self, - refresh_strategy: RefreshStrategy, - http_client_factory: HttpClientFactory, - ) -> ModelsManagerFuture<'_, ModelsResponse> { - Box::pin(OpenAiModelsManager::raw_model_catalog( - self, - refresh_strategy, - http_client_factory, - )) - } - fn refresh_after_auth_change( &self, http_client_factory: HttpClientFactory, @@ -452,7 +485,8 @@ impl ModelsManager for OpenAiModelsManager { return Ok(()); } let identity = self.endpoint_client.identity(); - if identity.is_some() && self.remote_models.read().await.identity == identity { + if identity.is_some() && self.remote_models.read().await.entry.identity == identity + { return Ok(()); } self.refresh_available_models( @@ -474,13 +508,18 @@ impl ModelsManager for OpenAiModelsManager { fn get_remote_models(&self) -> ModelsManagerFuture<'_, Vec> { Box::pin(async move { let entry = self.remote_models.read().await; - self.models_for_entry(&entry).into_owned() + self.models_for_entry(&entry.entry) + .into_models() + .into_owned() }) } fn try_get_remote_models(&self) -> Result, TryLockError> { let entry = self.remote_models.try_read()?; - Ok(self.models_for_entry(&entry).into_owned()) + Ok(self + .models_for_entry(&entry.entry) + .into_models() + .into_owned()) } fn auth_manager(&self) -> Option<&AuthManager> { @@ -496,7 +535,7 @@ impl ModelsManager for OpenAiModelsManager { Box::pin(async move { let remote = { let entry = self.remote_models.read().await; - let models = self.models_for_entry(&entry); + let models = self.models_for_entry(&entry.entry).into_models(); match &self.catalog_source { CatalogSource::Default => find_model_from_candidates(model, &models), CatalogSource::ExplicitProvider(_) => models @@ -528,30 +567,16 @@ impl ModelsManager for OpenAiModelsManager { impl OpenAiModelsManager { /// Selects the catalog available under the current provider and auth identity. - fn models_for_entry<'a>(&self, entry: &'a ModelsCacheEntry) -> Cow<'a, [ModelInfo]> { + fn models_for_entry<'a>(&self, entry: &'a ModelsCacheEntry) -> SelectedCatalog<'a> { if !self.api_key_discovery_disabled() && entry.identity.is_some() && entry.identity == self.endpoint_client.identity() { - Cow::Borrowed(&entry.models) + SelectedCatalog::CurrentSnapshot(Cow::Borrowed(&entry.models)) } else { - Cow::Owned(self.catalog_source.fallback_models().unwrap_or_default()) - } - } - - async fn raw_model_catalog( - &self, - refresh_strategy: RefreshStrategy, - http_client_factory: HttpClientFactory, - ) -> ModelsResponse { - if let Err(err) = self - .refresh_available_models(refresh_strategy, &http_client_factory) - .await - { - error!("failed to refresh available models: {err}"); - } - ModelsResponse { - models: self.get_remote_models().await, + SelectedCatalog::Fallback(Cow::Owned( + self.catalog_source.fallback_models().unwrap_or_default(), + )) } } @@ -559,7 +584,7 @@ impl OpenAiModelsManager { let refresh_guard = self.catalog_source.refresh_lock().await; let (identity, current_etag) = { let entry = self.remote_models.read().await; - (entry.identity.clone(), entry.etag.clone()) + (entry.entry.identity.clone(), entry.entry.etag.clone()) }; if let Some(identity) = identity && Some(&identity) == self.endpoint_client.identity().as_ref() @@ -652,12 +677,17 @@ impl OpenAiModelsManager { if is_provider_catalog && (access_denied || credentials_rejected) { { let mut current = self.remote_models.write().await; + let CatalogState { + entry: current, + provider_model_ids, + } = &mut *current; if request_identity != self.endpoint_client.identity() { return Err(error); } // Keep this failure in memory even if persistent invalidation fails. current.fetched_at = DateTime::::MIN_UTC; current.models.clear(); + provider_model_ids.clear(); current.etag = None; current.identity.clone_from(&request_identity); } @@ -733,11 +763,13 @@ impl OpenAiModelsManager { } /// Publish only while the request identity still matches, including after async storage. - async fn apply_remote_models(&self, mut entry: ModelsCacheEntry) -> bool { + async fn apply_remote_models(&self, entry: ModelsCacheEntry) -> bool { let mut current = self.remote_models.write().await; if entry.identity != self.endpoint_client.identity() { return false; } + let mut catalog = CatalogState::provider(entry); + let entry = &mut catalog.entry; let remote_only = entry .models .iter() @@ -749,7 +781,7 @@ impl OpenAiModelsManager { .is_some_and(AuthMode::has_chatgpt_account) })); if !remote_only && let Some(mut models) = self.catalog_source.fallback_models() { - for model in entry.models { + for model in std::mem::take(&mut entry.models) { if let Some(index) = models .iter() .position(|existing| existing.slug == model.slug) @@ -761,7 +793,7 @@ impl OpenAiModelsManager { } entry.models = models; } - *current = entry; + *current = catalog; true } @@ -780,8 +812,8 @@ impl OpenAiModelsManager { { let current = self.remote_models.read().await; // An access rejection must not be replaced by stale cache data if invalidation failed. - if current.fetched_at == DateTime::::MIN_UTC - && current.identity.as_ref() == Some(&identity) + if current.entry.fetched_at == DateTime::::MIN_UTC + && current.entry.identity.as_ref() == Some(&identity) { return false; } @@ -819,6 +851,22 @@ impl OpenAiModelsManager { } impl ModelsManager for StaticModelsManager { + fn model_catalog_with_origins( + &self, + _refresh_strategy: RefreshStrategy, + _http_client_factory: HttpClientFactory, + ) -> ModelsManagerFuture<'_, Vec> { + Box::pin(async move { + self.remote_models + .iter() + .cloned() + .map(|model| CatalogModel { + model, + origin: self.origin, + }) + .collect() + }) + } #[tracing::instrument(skip_all, fields(model = model))] fn get_model_info<'a>( &'a self, @@ -866,18 +914,6 @@ impl ModelsManager for StaticModelsManager { ) } - fn raw_model_catalog( - &self, - _refresh_strategy: RefreshStrategy, - _http_client_factory: HttpClientFactory, - ) -> ModelsManagerFuture<'_, ModelsResponse> { - Box::pin(async move { - ModelsResponse { - models: self.get_remote_models().await, - } - }) - } - fn get_remote_models(&self) -> ModelsManagerFuture<'_, Vec> { Box::pin(async { self.remote_models.clone() }) } diff --git a/codex-rs/models-manager/src/manager_tests.rs b/codex-rs/models-manager/src/manager_tests.rs index 26a9ccc1a1..dce9bfeb77 100644 --- a/codex-rs/models-manager/src/manager_tests.rs +++ b/codex-rs/models-manager/src/manager_tests.rs @@ -39,6 +39,9 @@ mod cache_identity_tests; #[path = "model_info_overrides_tests.rs"] mod model_info_overrides_tests; +#[path = "catalog_origin_tests.rs"] +mod catalog_origin_tests; + const DEFAULT_HTTP_CLIENT_FACTORY: HttpClientFactory = HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault); @@ -347,7 +350,11 @@ where } fn static_manager_for_tests(model_catalog: ModelsResponse) -> StaticModelsManager { - StaticModelsManager::new(/*auth_manager*/ None, model_catalog) + StaticModelsManager::new( + /*auth_manager*/ None, + model_catalog, + ModelOrigin::ExplicitConfig, + ) } #[tokio::test] @@ -1248,7 +1255,7 @@ async fn authoritative_catalog_refresh_failure_preserves_snapshot_and_freshness( .await, initial_info ); - assert_eq!(manager.remote_models.read().await.clone(), stored[0]); + assert_eq!(manager.remote_models.read().await.entry, stored[0]); assert_eq!(cache.stored_entries(), stored); let restarted = OpenAiModelsManager::new_with_cache( @@ -1382,7 +1389,7 @@ async fn online_refresh_updates_access_programs_with_unchanged_etag() { ); assert_eq!(manager.get_remote_models().await, vec![model]); assert_eq!( - manager.remote_models.read().await.etag.as_deref(), + manager.remote_models.read().await.entry.etag.as_deref(), Some("stable-catalog-etag") ); @@ -1401,7 +1408,13 @@ async fn online_refresh_updates_access_programs_with_unchanged_etag() { ); assert_eq!(cache_endpoint.fetch_count(), 0); assert_eq!( - cache_manager.remote_models.read().await.etag.as_deref(), + cache_manager + .remote_models + .read() + .await + .entry + .etag + .as_deref(), Some("stable-catalog-etag") ); } @@ -1806,6 +1819,7 @@ async fn static_manager_reads_latest_auth_mode() { ModelsResponse { models: vec![chatgpt_only_model, api_model], }, + ModelOrigin::ExplicitConfig, ); let chatgpt_models = manager