diff --git a/crates/provider-setup/src/service/implementation/credentials.rs b/crates/provider-setup/src/service/implementation/credentials.rs index 1cd8782573..8de77ec656 100644 --- a/crates/provider-setup/src/service/implementation/credentials.rs +++ b/crates/provider-setup/src/service/implementation/credentials.rs @@ -232,13 +232,10 @@ impl LiveProviderSetupService { .and_then(|v| v.as_str()) .ok_or_else(|| "missing 'provider' parameter".to_string())?; - let models: Vec = params - .get("models") - .and_then(|v| v.as_array()) - .ok_or_else(|| "missing 'models' array parameter".to_string())? - .iter() - .filter_map(|v| v.as_str().map(String::from)) - .collect(); + if !params.get("models").is_some_and(Value::is_array) { + return Err("missing 'models' array parameter".into()); + } + let models = parse_models_param(¶ms); // Validate provider exists (known or custom). if !is_custom_provider(provider_name) { @@ -248,6 +245,12 @@ impl LiveProviderSetupService { } } + let previous_models = self + .key_store + .load_config(provider_name) + .map(|config| config.models) + .unwrap_or_default(); + self.key_store .save_config(provider_name, None, None, Some(models.clone())) .map_err(ServiceError::message)?; @@ -255,6 +258,9 @@ impl LiveProviderSetupService { // Update the cross-provider priority list. if let Some(ref priority) = self.priority_models { let mut list = priority.write().await; + for previous in previous_models { + list.retain(|existing| existing != &previous); + } for m in models.iter().rev() { list.retain(|existing| existing != m); list.insert(0, m.clone()); @@ -271,3 +277,99 @@ impl LiveProviderSetupService { Ok(serde_json::json!({ "ok": true })) } } + +#[allow(clippy::unwrap_used, clippy::expect_used)] +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use tokio::sync::RwLock; + + use { + super::LiveProviderSetupService, crate::key_store::KeyStore, + moltis_config::schema::ProvidersConfig, moltis_providers::ProviderRegistry, + }; + + fn test_service( + key_store: KeyStore, + priority: Arc>>, + ) -> LiveProviderSetupService { + let mut service = LiveProviderSetupService::new( + Arc::new(RwLock::new(ProviderRegistry::empty())), + ProvidersConfig::default(), + None, + ); + service.key_store = key_store; + service.set_priority_models(priority); + service + } + + #[tokio::test] + async fn save_models_replaces_previous_provider_priorities() { + let dir = tempfile::tempdir().unwrap(); + let key_store = KeyStore::with_path(dir.path().join("provider_keys.json")); + key_store + .save_config( + "openai", + Some("sk-test".to_string()), + None, + Some(vec!["old-a".to_string(), "old-b".to_string()]), + ) + .unwrap(); + let priority = Arc::new(RwLock::new(vec![ + "other-provider-model".to_string(), + "old-a".to_string(), + "old-b".to_string(), + ])); + let service = test_service(key_store.clone(), Arc::clone(&priority)); + + service + .save_models_inner(serde_json::json!({ + "provider": "openai", + "models": ["openai::new-a"] + })) + .await + .unwrap(); + + assert_eq!(key_store.load_config("openai").unwrap().models, vec![ + "new-a" + ]); + assert_eq!(*priority.read().await, vec![ + "new-a".to_string(), + "other-provider-model".to_string() + ]); + } + + #[tokio::test] + async fn save_models_empty_selection_clears_previous_provider_priorities() { + let dir = tempfile::tempdir().unwrap(); + let key_store = KeyStore::with_path(dir.path().join("provider_keys.json")); + key_store + .save_config( + "openai", + Some("sk-test".to_string()), + None, + Some(vec!["old-a".to_string(), "old-b".to_string()]), + ) + .unwrap(); + let priority = Arc::new(RwLock::new(vec![ + "old-a".to_string(), + "other-provider-model".to_string(), + "old-b".to_string(), + ])); + let service = test_service(key_store.clone(), Arc::clone(&priority)); + + service + .save_models_inner(serde_json::json!({ + "provider": "openai", + "models": [] + })) + .await + .unwrap(); + + assert!(key_store.load_config("openai").unwrap().models.is_empty()); + assert_eq!(*priority.read().await, vec![ + "other-provider-model".to_string() + ]); + } +} diff --git a/crates/web/ui/e2e/specs/providers.spec.js b/crates/web/ui/e2e/specs/providers.spec.js index 6155ee58d6..e4efaa305f 100644 --- a/crates/web/ui/e2e/specs/providers.spec.js +++ b/crates/web/ui/e2e/specs/providers.spec.js @@ -186,4 +186,90 @@ test.describe("Provider setup page", () => { .toContain("Custom.gguf configured successfully!"); expect(pageErrors).toEqual([]); }); + + test("preferred model selector replaces existing selections", async ({ page }) => { + const pageErrors = watchPageErrors(page); + await openProvidersPage(page); + + await page.evaluate(async () => { + const [providers, state] = await Promise.all([import("/assets/js/providers.js"), import("/assets/js/state.js")]); + window.__providerSaveModelsCalls = []; + + state.setWs({ + readyState: WebSocket.OPEN, + send(raw) { + const frame = JSON.parse(raw); + const respond = (payload) => { + queueMicrotask(() => { + const pending = state.pending[frame.id]; + if (!pending) return; + pending(payload); + delete state.pending[frame.id]; + }); + }; + + if (frame.method === "models.list") { + respond({ + ok: true, + payload: [ + { + id: "openai::gpt-preferred-a", + provider: "openai", + displayName: "GPT Preferred A", + supportsTools: true, + preferred: true, + }, + { + id: "openai::gpt-preferred-b", + provider: "openai", + displayName: "GPT Preferred B", + supportsTools: true, + preferred: true, + }, + ], + }); + return; + } + + if (frame.method === "providers.available") { + respond({ + ok: true, + payload: [ + { + name: "openai", + displayName: "OpenAI", + authType: "api-key", + models: ["gpt-preferred-a", "gpt-preferred-b"], + }, + ], + }); + return; + } + + if (frame.method === "providers.save_models") { + window.__providerSaveModelsCalls.push(frame.params); + respond({ ok: true, payload: { ok: true } }); + return; + } + + respond({ ok: true, payload: {} }); + }, + }); + + providers.openModelSelectorForProvider("openai", "OpenAI"); + }); + + await expect(page.locator("#providerModal")).toBeVisible(); + await expect(page.locator("#providerModalBody .model-card.selected")).toHaveCount(2); + + await page.locator("#providerModalBody .model-card", { hasText: "GPT Preferred B" }).click(); + await expect(page.locator("#providerModalBody .model-card.selected")).toHaveCount(1); + await page.getByRole("button", { name: "Save", exact: true }).click(); + + await expect + .poll(() => page.evaluate(() => window.__providerSaveModelsCalls || [])) + .toEqual([{ provider: "openai", models: ["openai::gpt-preferred-a"] }]); + + expect(pageErrors).toEqual([]); + }); }); diff --git a/crates/web/ui/src/providers/auth-flow.ts b/crates/web/ui/src/providers/auth-flow.ts index 49239a7d6f..02faff0e5c 100644 --- a/crates/web/ui/src/providers/auth-flow.ts +++ b/crates/web/ui/src/providers/auth-flow.ts @@ -802,7 +802,12 @@ function showMultiModelSelector( m.title.textContent = `${providerDisplayName} \u2014 Preferred Models`; m.body.textContent = ""; - const selectedIds: Set = new Set(savedModels); + const savedModelTokens = new Set(Array.from(savedModels, stripModelNamespace)); + const selectedIds: Set = new Set( + models + .filter((model) => savedModels.has(model.id) || savedModelTokens.has(stripModelNamespace(model.id))) + .map((model) => model.id), + ); // Track per-model probe state: "probing" | "ok" | { error: string } const probeResults: Map = new Map();