From 664eed08099ce70647b552d3b8897b3edcb2e4ce Mon Sep 17 00:00:00 2001 From: Sayan Sisodiya Date: Tue, 10 Feb 2026 20:51:18 -0800 Subject: [PATCH] apply overrides to picker, add test for ModelInfo/ModelInfoPatch drift --- codex-rs/core/config.schema.json | 4 ++ codex-rs/core/src/models_manager/manager.rs | 68 +++++++++++++++++-- .../core/src/models_manager/model_info.rs | 3 + codex-rs/protocol/src/openai_models.rs | 34 +++++++++- 4 files changed, 101 insertions(+), 8 deletions(-) diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index 30f33cd6cb..4d7966a225 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -535,6 +535,10 @@ ], "default": null }, + "prefer_websockets": { + "default": null, + "type": "boolean" + }, "priority": { "default": null, "format": "int32", diff --git a/codex-rs/core/src/models_manager/manager.rs b/codex-rs/core/src/models_manager/manager.rs index f62fbac446..0688f19a00 100644 --- a/codex-rs/core/src/models_manager/manager.rs +++ b/codex-rs/core/src/models_manager/manager.rs @@ -85,7 +85,7 @@ impl ModelsManager { error!("failed to refresh available models: {err}"); } let remote_models = self.get_remote_models(config).await; - self.build_available_models(remote_models) + self.build_available_models(remote_models, config) } /// List collaboration mode presets. @@ -100,7 +100,7 @@ impl ModelsManager { /// Returns an error if the internal lock cannot be acquired. pub fn try_list_models(&self, config: &Config) -> Result, TryLockError> { let remote_models = self.try_get_remote_models(config)?; - Ok(self.build_available_models(remote_models)) + Ok(self.build_available_models(remote_models, config)) } // todo(aibrahim): should be visible to core only and sent on session_configured event @@ -124,7 +124,7 @@ impl ModelsManager { error!("failed to refresh available models: {err}"); } let remote_models = self.get_remote_models(config).await; - let available = self.build_available_models(remote_models); + let available = self.build_available_models(remote_models, config); available .iter() .find(|model| model.is_default) @@ -298,8 +298,19 @@ impl ModelsManager { } /// Merge remote model metadata into picker-ready presets, preserving existing entries. - fn build_available_models(&self, mut remote_models: Vec) -> Vec { + fn build_available_models( + &self, + mut remote_models: Vec, + config: &Config, + ) -> Vec { remote_models.sort_by(|a, b| a.priority.cmp(&b.priority)); + remote_models = remote_models + .into_iter() + .map(|remote_model| { + let remote_slug = remote_model.slug.clone(); + model_info::with_model_info_patch(remote_model, &remote_slug, config) + }) + .collect(); let remote_presets: Vec = remote_models.into_iter().map(Into::into).collect(); let existing_presets = self.local_models.clone(); @@ -390,6 +401,8 @@ mod tests { use crate::features::Feature; use crate::model_provider_info::WireApi; use chrono::Utc; + use codex_protocol::openai_models::ModelInfoPatch; + use codex_protocol::openai_models::ModelVisibility; use codex_protocol::openai_models::ModelsResponse; use core_test_support::responses::mount_models_once; use pretty_assertions::assert_eq; @@ -827,8 +840,8 @@ mod tests { ); } - #[test] - fn build_available_models_picks_default_after_hiding_hidden_models() { + #[tokio::test] + async fn build_available_models_picks_default_after_hiding_hidden_models() { let codex_home = tempdir().expect("temp dir"); let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("Test API Key")); @@ -847,11 +860,52 @@ mod tests { let mut expected_visible = ModelPreset::from(visible_model.clone()); expected_visible.is_default = true; - let available = manager.build_available_models(vec![hidden_model, visible_model]); + let config = ConfigBuilder::default() + .codex_home(codex_home.path().to_path_buf()) + .build() + .await + .expect("load default test config"); + let available = manager.build_available_models(vec![hidden_model, visible_model], &config); assert_eq!(available, vec![expected_hidden, expected_visible]); } + #[tokio::test] + async fn build_available_models_applies_model_info_overrides() { + let codex_home = tempdir().expect("temp dir"); + let auth_manager = + AuthManager::from_auth_for_testing(CodexAuth::from_api_key("Test API Key")); + let provider = provider_for("http://example.test".to_string()); + let mut manager = + ModelsManager::with_provider(codex_home.path().to_path_buf(), auth_manager, provider); + manager.local_models = Vec::new(); + + let mut config = ConfigBuilder::default() + .codex_home(codex_home.path().to_path_buf()) + .build() + .await + .expect("load default test config"); + config.model_info_overrides.insert( + "visible".to_string(), + ModelInfoPatch { + display_name: Some("Visible Local Override".to_string()), + visibility: Some(ModelVisibility::Hide), + ..Default::default() + }, + ); + + let visible_model = remote_model_with_visibility("visible", "Visible", 1, "list"); + let hidden_model = remote_model_with_visibility("hidden", "Hidden", 0, "hide"); + let available = manager.build_available_models(vec![hidden_model, visible_model], &config); + + let visible = available + .iter() + .find(|preset| preset.model == "visible") + .expect("visible model should exist"); + assert_eq!(visible.display_name, "Visible Local Override"); + assert!(!visible.show_in_picker); + } + #[test] fn bundled_models_json_roundtrips() { let file_contents = include_str!("../../models.json"); diff --git a/codex-rs/core/src/models_manager/model_info.rs b/codex-rs/core/src/models_manager/model_info.rs index 17a8b7bcef..9f4a1a4205 100644 --- a/codex-rs/core/src/models_manager/model_info.rs +++ b/codex-rs/core/src/models_manager/model_info.rs @@ -98,6 +98,9 @@ pub(crate) fn with_model_info_patch( if let Some(input_modalities) = &model_info_patch.input_modalities { model.input_modalities = input_modalities.clone(); } + if let Some(prefer_websockets) = model_info_patch.prefer_websockets { + model.prefer_websockets = prefer_websockets; + } model } diff --git a/codex-rs/protocol/src/openai_models.rs b/codex-rs/protocol/src/openai_models.rs index b2a63877d0..f791964250 100644 --- a/codex-rs/protocol/src/openai_models.rs +++ b/codex-rs/protocol/src/openai_models.rs @@ -256,7 +256,7 @@ pub struct ModelInfo { /// User-provided patch for overriding model metadata in local config. /// -/// Every field is optional so users can override only the parts they need. +/// Every field is optional so users can override only the parts of [`ModelInfo`] they need. /// The target model slug is provided by the surrounding map key. #[derive(Debug, Serialize, Deserialize, Clone, Default, PartialEq, Eq, TS, JsonSchema)] #[serde(default)] @@ -283,6 +283,7 @@ pub struct ModelInfoPatch { pub effective_context_window_percent: Option, pub experimental_supported_tools: Option>, pub input_modalities: Option>, + pub prefer_websockets: Option, } impl ModelInfo { @@ -514,6 +515,7 @@ fn nearest_effort(target: ReasoningEffort, supported: &[ReasoningEffort]) -> Rea mod tests { use super::*; use pretty_assertions::assert_eq; + use std::collections::BTreeSet; fn test_model(spec: Option) -> ModelInfo { ModelInfo { @@ -697,4 +699,34 @@ mod tests { ); assert_eq!(personality_variables.get_personality_message(None), None); } + + #[test] + fn model_info_patch_field_coverage_is_explicit() { + fn schema_fields() -> BTreeSet { + let schema = schemars::schema_for!(T); + let object = schema + .schema + .object + .as_ref() + .expect("expected object schema"); + object.properties.keys().cloned().collect() + } + + let model_info_fields = schema_fields::(); + let patch_fields = schema_fields::(); + let intentionally_non_patchable: BTreeSet = + ["slug".to_string()].into_iter().collect(); + let expected_patch_fields: BTreeSet = model_info_fields + .difference(&intentionally_non_patchable) + .cloned() + .collect(); + + assert_eq!(patch_fields, expected_patch_fields); + for field in &intentionally_non_patchable { + assert!( + model_info_fields.contains(field), + "intentionally_non_patchable contains unknown field: {field}" + ); + } + } }