Respect model catalog origins when selecting Guardian reviewers (#57770)

## Why

A bundled reviewer entry does not establish that a custom gateway offers that model. Guardian should fall back to the parent model when the reviewer is neither advertised by the provider nor explicitly configured.

## What changed

Track model metadata origins as `Bundled`, `Provider`, or `ExplicitConfig`, including across catalog merges and cache reloads. Exclude bundled entries from Guardian reviewer selection unless `auto_review_model_override` is set. Continue selecting reviewers from provider and explicitly configured catalogs.

## Testing

Add coverage for gateway reviewer selection with each catalog origin, plus origin tracking across cache reloads, empty provider responses, and authentication changes.

GitOrigin-RevId: f8bf865d406e56ee0ade90087b55f2e33b98181e
This commit is contained in:
ioanbudea-OAI
2026-10-28 16:00:30 +00:00
committed by copyberry
parent 7c11b8074d
commit 5c0e42161e
18 changed files with 478 additions and 87 deletions
@@ -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)
+13 -4
View File
@@ -78,14 +78,23 @@ pub(crate) async fn resolve_review_model(
Arc<codex_protocol::openai_models::ModelInfo>,
) {
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,
+29 -2
View File
@@ -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());
@@ -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<codex_models_manager::CatalogModel>> {
self.inner.model_catalog_with_origins(strategy, factory)
}
fn get_remote_models(&self) -> ModelsManagerFuture<'_, Vec<ModelInfo>> {
@@ -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());
@@ -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;
@@ -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?;
@@ -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())
@@ -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;
@@ -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<codex_models_manager::CatalogModel>> {
self.inner.model_catalog_with_origins(strategy, factory)
}
fn get_remote_models(&self) -> ModelsManagerFuture<'_, Vec<ModelInfo>> {
@@ -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(),
});
@@ -430,10 +430,16 @@ impl ModelProvider for AmazonBedrockModelProvider {
_codex_home: PathBuf,
config_model_catalog: Option<ModelsResponse>,
) -> 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<ModelsResponse>,
) -> 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)
}
}
+6
View File
@@ -149,6 +149,11 @@ impl ModelProvider for GcpModelProvider {
&self,
config_model_catalog: Option<ModelsResponse>,
) -> 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,
))
}
}
+11 -1
View File
@@ -370,10 +370,19 @@ pub trait ModelProvider: fmt::Debug + Send + Sync {
&self,
config_model_catalog: Option<ModelsResponse>,
) -> 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(
+83
View File
@@ -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<String>,
}
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<CatalogModel> {
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()
}
}
@@ -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
}]
);
}
}
+3
View File
@@ -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;
+97 -61
View File
@@ -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<CatalogModel>>;
/// 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<HashMap<ModelCatalogOriginator, SharedModelsManager>>,
remote_models: RwLock<ModelsCacheEntry>,
remote_models: RwLock<CatalogState>,
cache: Option<Arc<dyn ModelsCache>>,
endpoint_client: SharedModelsEndpointClient,
catalog_source: CatalogSource,
@@ -317,6 +340,7 @@ pub struct OpenAiModelsManager {
pub struct StaticModelsManager {
remote_models: Vec<ModelInfo>,
auth_manager: Option<Arc<AuthManager>>,
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<Arc<AuthManager>>, model_catalog: ModelsResponse) -> Self {
/// Construct a static model manager with the source of its supplied catalog.
pub fn new(
auth_manager: Option<Arc<AuthManager>>,
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<CatalogModel>> {
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<SharedModelsManager> {
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<ModelInfo>> {
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<Vec<ModelInfo>, 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::<Utc>::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::<Utc>::MIN_UTC
&& current.identity.as_ref() == Some(&identity)
if current.entry.fetched_at == DateTime::<Utc>::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<CatalogModel>> {
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<ModelInfo>> {
Box::pin(async { self.remote_models.clone() })
}
+18 -4
View File
@@ -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