From d5b5eaa0e7bee0f1725fe62d649d50826a6d3bfc Mon Sep 17 00:00:00 2001 From: Ryan Ward Date: Mon, 14 Sep 2026 13:13:07 -0500 Subject: [PATCH] Make LiteLLM model discovery authoritative --- .../action_model/execute/start_agent.rs | 27 ++ app/src/ai/llms.rs | 343 ++++++++++++------ app/src/ai/llms_tests.rs | 79 ++++ app/src/settings_view/provider_setup_view.rs | 9 +- 4 files changed, 336 insertions(+), 122 deletions(-) diff --git a/app/src/ai/blocklist/action_model/execute/start_agent.rs b/app/src/ai/blocklist/action_model/execute/start_agent.rs index 167e2534..96107a9a 100644 --- a/app/src/ai/blocklist/action_model/execute/start_agent.rs +++ b/app/src/ai/blocklist/action_model/execute/start_agent.rs @@ -19,6 +19,7 @@ use crate::ai::agent::{ }; use crate::ai::blocklist::orchestration_event_streamer::OrchestrationEventStreamer; use crate::ai::blocklist::{BlocklistAIHistoryEvent, BlocklistAIHistoryModel}; +use crate::ai::llms::{LLMId, LLMPreferences}; use crate::ai::local_harness_setup::local_harness_product_disabled_message; /// Per-request outcome of a StartAgent dispatch. @@ -707,6 +708,32 @@ impl StartAgentExecutor { let (prompt, execution_mode) = normalize_legacy_local_child_harness_command(prompt, execution_mode); + if parent_run_id.is_none() { + if let StartAgentExecutionMode::Local { + harness_type: None, + model_id: Some(model_id), + } = &execution_mode + { + if ctx + .try_get_singleton_model_as_ref::() + .is_some_and(|preferences| { + preferences + .get_llm_info(&LLMId::from(model_id.as_str())) + .is_none() + }) + { + let _ = sender.try_send(StartAgentOutcome::Error(format!( + "Model '{model_id}' is no longer available. Refresh the provider models and select an available model before launching this agent." + ))); + return StartAgentDispatch { + request_id, + receiver, + wait_policy, + detached, + }; + } + } + } let prompt = compose_leaf_agent_prompt(&prompt); self.pending.insert( request_id, diff --git a/app/src/ai/llms.rs b/app/src/ai/llms.rs index 9bc98399..63ffcc2e 100644 --- a/app/src/ai/llms.rs +++ b/app/src/ai/llms.rs @@ -17,6 +17,7 @@ use galaxy_core::features::FeatureFlag; use galaxy_core::ui::icons::Icon; use galaxy_core::user_preferences::GetUserPreferences; use galaxyui::r#async::Timer; +use galaxyui::windowing::WindowManager; use galaxyui::{AppContext, Entity, EntityId, ModelContext, SingletonEntity}; use parking_lot::FairMutex; use serde::{de, Deserialize, Serialize}; @@ -37,6 +38,8 @@ use crate::settings::{ OpenAIProviderConfig, OpenAIProviderKind, }; use crate::user_config::{WarpConfig, WarpConfigUpdateEvent}; +use crate::view_components::DismissibleToast; +use crate::workspace::ToastStack; use crate::workspaces::user_workspaces::{UserWorkspaces, UserWorkspacesEvent}; use crate::{report_error, AISettings}; @@ -597,6 +600,10 @@ pub struct LLMPreferences { #[cfg(not(target_family = "wasm"))] fetched_openai_models: Vec, #[cfg(not(target_family = "wasm"))] + litellm_connections: Vec, + #[cfg(not(target_family = "wasm"))] + catalog_refresh_generations: HashMap, u64>, + #[cfg(not(target_family = "wasm"))] chatgpt_subscription_models_refresh_in_flight: bool, #[cfg(not(target_family = "wasm"))] acp_selections: HashMap, @@ -702,6 +709,7 @@ impl LLMPreferences { | AISettingsChangedEvent::AcpAgentId { .. } | AISettingsChangedEvent::BedrockModels { .. } ) { + let previous_ids: Vec<_> = me.openai_provider_routing.keys().cloned().collect(); me.inject_bedrock_models(ctx); me.inject_openai_models(ctx); if matches!( @@ -717,9 +725,28 @@ impl LLMPreferences { { me.refresh_bedrock_models(ctx); } + me.refresh_changed_litellm_connections(ctx); // Safety: ensure the default model is still present in choices. // If all provider models were removed, the default_id would dangle. me.ensure_default_model_present(); + let removed: HashSet<_> = previous_ids + .into_iter() + .map(LLMId::from) + .filter(|id| me.get_llm_info(id).is_none()) + .collect(); + if !removed.is_empty() { + me.base_llm_for_terminal_view + .retain(|_, id| !removed.contains(id)); + if ctx + .try_get_singleton_model_as_ref::() + .is_some() + { + me.reconcile_disabled_model_preferences(ctx); + } + me.trigger_snapshot_save(ctx); + ctx.emit(LLMPreferencesEvent::UpdatedActiveAgentModeLLM); + ctx.emit(LLMPreferencesEvent::UpdatedActiveCodingLLM); + } ctx.emit(LLMPreferencesEvent::UpdatedAvailableLLMs); } }); @@ -738,6 +765,10 @@ impl LLMPreferences { #[cfg(not(target_family = "wasm"))] fetched_openai_models: Vec::new(), #[cfg(not(target_family = "wasm"))] + litellm_connections: Vec::new(), + #[cfg(not(target_family = "wasm"))] + catalog_refresh_generations: HashMap::new(), + #[cfg(not(target_family = "wasm"))] chatgpt_subscription_models_refresh_in_flight: false, #[cfg(not(target_family = "wasm"))] acp_selections: HashMap::new(), @@ -766,6 +797,7 @@ impl LLMPreferences { me.inject_openai_models(ctx); me.ensure_default_model_present(); me.fetch_openai_models_from_endpoint(ctx); + me.refresh_changed_litellm_connections(ctx); me.refresh_chatgpt_subscription_models(ctx); me.schedule_chatgpt_subscription_model_refresh(ctx); } @@ -1507,6 +1539,11 @@ impl LLMPreferences { } }; + let generation = self.catalog_refresh_generations.entry(None).or_default(); + *generation += 1; + let generation = *generation; + let requested_base_url = base_url.clone(); + let requested_api_key = settings.openai_api_key.value().clone(); let _ = ctx.spawn( async move { let base = base_url.trim_end_matches('/'); @@ -1515,35 +1552,89 @@ impl LLMPreferences { .build() .unwrap_or_default(); - // Try LiteLLM /model/info first for rich metadata - if let Some(models) = - fetch_from_litellm_model_info(base, api_key.as_deref(), &client).await - { - return models; - } - - // Fallback to standard OpenAI /models endpoint - fetch_from_openai_models(base, api_key.as_deref(), &client).await + fetch_litellm_catalog(base, api_key.as_deref(), &client).await }, - |me, models, ctx| { - if !models.is_empty() { - me.fetched_openai_models = models.clone(); - AISettings::handle(ctx).update(ctx, |settings, ctx| { - if let Err(err) = settings.openai_models.set_value(models, ctx) { - report_error!(err.context("Failed to persist fetched OpenAI models")); - } - }); - me.inject_openai_models(ctx); - ctx.emit(LLMPreferencesEvent::UpdatedAvailableLLMs); + move |me, result, ctx| { + if me.catalog_refresh_generations.get(&None) != Some(&generation) { + return; + } + let settings = AISettings::as_ref(ctx); + if settings.openai_base_url.value() != &requested_base_url + || settings.openai_api_key.value() != &requested_api_key + || !settings.openai_providers.value().is_empty() + { + return; + } + match result { + Ok(models) => { + let models = + replace_litellm_catalog(settings.openai_models.value(), models); + me.fetched_openai_models = models.clone(); + AISettings::handle(ctx).update(ctx, |settings, ctx| { + if let Err(err) = settings.openai_models.set_value(models, ctx) { + report_error!( + err.context("Failed to persist fetched OpenAI models") + ); + } + }); + } + Err(error) => Self::report_catalog_refresh_error(error, ctx), } }, ); } + #[cfg(not(target_family = "wasm"))] + fn report_catalog_refresh_error(error: String, ctx: &mut ModelContext) { + log::warn!("Model refresh failed; keeping cached catalog: {error}"); + if let Some(window_id) = WindowManager::as_ref(ctx).active_window() { + ToastStack::handle(ctx).update(ctx, |stack, ctx| { + stack.add_ephemeral_toast( + DismissibleToast::error(format!( + "Model refresh failed. Keeping saved models. {error}" + )), + window_id, + ctx, + ); + }); + } + } + + #[cfg(not(target_family = "wasm"))] + fn refresh_changed_litellm_connections(&mut self, ctx: &mut ModelContext) { + let settings = AISettings::as_ref(ctx); + if !settings.is_openai_provider_enabled() { + self.litellm_connections.clear(); + return; + } + let connections: Vec<_> = settings + .openai_providers + .value() + .iter() + .cloned() + .map(|mut provider| { + provider.models.clear(); + provider + }) + .collect(); + let changed: Vec<_> = connections + .iter() + .enumerate() + .filter_map(|(index, provider)| { + (provider.enabled + && provider.kind == OpenAIProviderKind::LiteLLM + && self.litellm_connections.get(index) != Some(provider)) + .then_some(index) + }) + .collect(); + self.litellm_connections = connections; + for index in changed { + self.fetch_openai_provider_models(index, ctx); + } + } + /// Explicitly refreshes the models for one entry in the OpenAI-compatible - /// provider registry. Unlike the legacy endpoint refresh, this is only - /// called from a user action so configured remote endpoints are never - /// contacted merely because Galaxy started. + /// provider registry. LiteLLM also refreshes on startup and connector changes. #[cfg(not(target_family = "wasm"))] pub fn fetch_openai_provider_models( &mut self, @@ -1569,68 +1660,54 @@ impl LLMPreferences { return; } - let provider_kind = provider.kind; - let requested_base_url = provider.base_url; - let requested_provider_kind = provider_kind; - let api_key = provider.api_key.filter(|key| !key.is_empty()); - let request_base_url = requested_base_url.clone(); - + let generation = self + .catalog_refresh_generations + .entry(Some(provider_index)) + .or_default(); + *generation += 1; + let generation = *generation; + let requested_provider = provider.clone(); let _ = ctx.spawn( - async move { - if provider_kind == OpenAIProviderKind::ChatGPTSubscription { - return match Self::discover_chatgpt_subscription_models().await { - Ok(models) => models, - Err(error) => { - log::warn!( - "[chatgpt/models] Failed to discover ChatGPT subscription models: {error}" - ); - Vec::new() - } - }; - } - - let base = request_base_url.trim_end_matches('/'); - let client = reqwest::Client::builder() - .timeout(std::time::Duration::from_secs(10)) - .build() - .unwrap_or_default(); - - if provider_kind == OpenAIProviderKind::LiteLLM { - if let Some(models) = - fetch_from_litellm_model_info(base, api_key.as_deref(), &client).await - { - return models; - } - } - - fetch_from_openai_models(base, api_key.as_deref(), &client).await - }, - move |_, discovered_models, ctx| { - if discovered_models.is_empty() { + async move { Self::discover_openai_provider_models(provider).await }, + move |me, result, ctx| { + if me.catalog_refresh_generations.get(&Some(provider_index)) != Some(&generation) { return; } - + let discovered_models = match result { + Ok(models) => models, + Err(error) => { + Self::report_catalog_refresh_error(error, ctx); + return; + } + }; AISettings::handle(ctx).update(ctx, |settings, ctx| { let mut providers = settings.openai_providers.value().clone(); let Some(provider) = providers.get_mut(provider_index) else { return; }; - - // Do not apply a response to an entry that was edited or - // reordered while its discovery request was in flight. - if provider.kind != requested_provider_kind - || provider.base_url != requested_base_url - { + let mut current_connection = provider.clone(); + current_connection.models.clear(); + let mut requested_connection = requested_provider.clone(); + requested_connection.models.clear(); + if current_connection != requested_connection { return; } - - provider.models = if provider.kind == OpenAIProviderKind::ChatGPTSubscription { - merge_discovered_chatgpt_subscription_models( - &provider.models, - discovered_models, - ) - } else { - merge_discovered_provider_models(&provider.models, discovered_models) + provider.models = match provider.kind { + OpenAIProviderKind::ChatGPTSubscription => { + merge_discovered_chatgpt_subscription_models( + &provider.models, + discovered_models, + ) + } + OpenAIProviderKind::LiteLLM => { + replace_litellm_catalog(&provider.models, discovered_models) + } + OpenAIProviderKind::OpenAI + | OpenAIProviderKind::Anthropic + | OpenAIProviderKind::Gemini + | OpenAIProviderKind::VertexAI => { + merge_discovered_provider_models(&provider.models, discovered_models) + } }; if let Err(err) = settings.openai_providers.set_value(providers, ctx) { report_error!(err.context("Failed to persist discovered provider models")); @@ -1714,28 +1791,11 @@ impl LLMPreferences { .map_err(|error| format!("Could not create the provider client: {error}"))?; let api_key = provider.api_key.as_deref().filter(|key| !key.is_empty()); - let models = if provider.kind == OpenAIProviderKind::LiteLLM { - if let Some(models) = fetch_from_litellm_model_info(&base_url, api_key, &client).await { - models - } else { - fetch_from_openai_models(&base_url, api_key, &client).await - } + if provider.kind == OpenAIProviderKind::LiteLLM { + fetch_litellm_catalog(&base_url, api_key, &client).await } else { fetch_from_openai_models(&base_url, api_key, &client).await - }; - - if models.is_empty() { - let endpoint_description = if provider.kind == OpenAIProviderKind::LiteLLM { - "/model/info or /models" - } else { - "/models" - }; - return Err(format!( - "The provider responded, but no models were found at {endpoint_description}." - )); } - - Ok(models) } #[cfg(not(target_family = "wasm"))] @@ -2690,6 +2750,28 @@ pub(crate) fn merge_discovered_provider_models( merged } +/// A successful LiteLLM /models response owns catalog membership. +#[cfg(not(target_family = "wasm"))] +pub(crate) fn replace_litellm_catalog( + existing: &[OpenAIModelConfig], + discovered: Vec, +) -> Vec { + let retained: Vec<_> = existing + .iter() + .filter(|old| discovered.iter().any(|new| new.model_id == old.model_id)) + .cloned() + .collect(); + let mut models = merge_discovered_provider_models(&retained, discovered); + for model in &mut models { + if let Some(old) = retained.iter().find(|old| old.model_id == model.model_id) { + model + .capability_overrides + .extend(old.capability_overrides.clone()); + } + } + models +} + /// Merges ChatGPT subscription model metadata as a backend-owned catalog. /// /// Unlike generic OpenAI-compatible providers, ChatGPT subscription models come @@ -3092,41 +3174,65 @@ async fn fetch_from_litellm_model_info( /// Fetches models from the standard OpenAI-compatible `/models` endpoint. /// Used as a fallback when `/model/info` is unavailable. +#[cfg(not(target_family = "wasm"))] +async fn fetch_litellm_catalog( + base_url: &str, + api_key: Option<&str>, + client: &reqwest::Client, +) -> Result, String> { + let models = fetch_from_openai_models(base_url, api_key, client).await?; + if models.is_empty() { + return Ok(models); + } + let metadata = fetch_from_litellm_model_info(base_url, api_key, client) + .await + .unwrap_or_default(); + Ok(models + .into_iter() + .map(|model| { + metadata + .iter() + .find(|info| info.model_id == model.model_id) + .cloned() + .unwrap_or(model) + }) + .collect()) +} + #[cfg(not(target_family = "wasm"))] async fn fetch_from_openai_models( base_url: &str, api_key: Option<&str>, client: &reqwest::Client, -) -> Vec { +) -> Result, String> { let url = format!("{base_url}/models"); let mut request = client.get(&url); if let Some(key) = api_key { request = request.header("Authorization", format!("Bearer {key}")); } - - let response = match request.send().await { - Ok(r) => r, - Err(e) => { - log::warn!("[openai/litellm] Failed to fetch models from /models endpoint: {e}"); - return Vec::new(); - } - }; - + let response = request + .send() + .await + .map_err(|error| format!("Could not refresh /models: {error}"))?; if !response.status().is_success() { - log::warn!( - "[openai/litellm] /models returned HTTP {}", + return Err(format!( + "Model refresh failed: /models returned HTTP {}", response.status() - ); - return Vec::new(); + )); + } + let body: serde_json::Value = response + .json() + .await + .map_err(|error| format!("Invalid /models response: {error}"))?; + let entries = body["data"] + .as_array() + .ok_or_else(|| "Invalid /models response: missing data array".to_string())?; + if entries + .iter() + .any(|entry| entry["id"].as_str().is_none_or(|id| id.trim().is_empty())) + { + return Err("Invalid /models response: model ID is missing or empty".to_string()); } - - let body: serde_json::Value = match response.json().await { - Ok(v) => v, - Err(e) => { - log::warn!("[openai/litellm] Failed to parse /models response: {e}"); - return Vec::new(); - } - }; fn u32_from_any(value: &serde_json::Value, keys: &[&str]) -> Option { keys.iter() @@ -3134,10 +3240,7 @@ async fn fetch_from_openai_models( .and_then(|value| u32::try_from(value).ok()) } - let models: Vec = body["data"] - .as_array() - .map(Vec::as_slice) - .unwrap_or_default() + let models: Vec = entries .iter() .filter_map(|m| { let id = m["id"].as_str()?; @@ -3221,7 +3324,7 @@ async fn fetch_from_openai_models( "[openai/litellm] Fetched {} model(s) from /models endpoint", models.len() ); - models + Ok(models) } /// Gets the last cached LLM metadata. diff --git a/app/src/ai/llms_tests.rs b/app/src/ai/llms_tests.rs index e182ff7b..1f2afafb 100644 --- a/app/src/ai/llms_tests.rs +++ b/app/src/ai/llms_tests.rs @@ -229,6 +229,8 @@ fn empty_preferences() -> LLMPreferences { custom_model_routers: Vec::new(), openai_provider_routing: HashMap::new(), fetched_openai_models: Vec::new(), + litellm_connections: Vec::new(), + catalog_refresh_generations: HashMap::new(), chatgpt_subscription_models_refresh_in_flight: false, acp_selections: HashMap::new(), } @@ -394,6 +396,8 @@ fn acp_models_are_injected_only_while_acp_is_enabled() { custom_model_routers: Vec::new(), openai_provider_routing: HashMap::new(), fetched_openai_models: Vec::new(), + litellm_connections: Vec::new(), + catalog_refresh_generations: HashMap::new(), chatgpt_subscription_models_refresh_in_flight: false, acp_selections: HashMap::new(), }; @@ -962,3 +966,78 @@ fn chatgpt_default_window_limits_routing_without_a_separate_input_limit() { 272_000 ); } + +#[test] +fn litellm_refresh_replaces_membership_and_preserves_preferences() { + let mut retained = openai_model("retained"); + retained.enabled = false; + retained.display_name = "My model".to_string(); + retained.capability_overrides.insert( + "vision".to_string(), + crate::settings::ModelCapabilityOverride::Unsupported, + ); + let removed = openai_model("removed"); + let mut updated = openai_model("retained"); + updated.context_size = 300_000; + let models = replace_litellm_catalog(&[retained, removed], vec![updated, openai_model("new")]); + assert_eq!(models.len(), 2); + assert_eq!(models[0].model_id, "retained"); + assert_eq!(models[0].context_size, 300_000); + assert_eq!(models[0].display_name, "My model"); + assert!(!models[0].enabled); + assert_eq!( + models[0].capability_override("vision"), + crate::settings::ModelCapabilityOverride::Unsupported + ); + assert_eq!(models[1].model_id, "new"); + assert!(replace_litellm_catalog(&models, vec![]).is_empty()); +} + +#[tokio::test] +async fn litellm_models_endpoint_owns_membership_even_with_extra_metadata_models() { + let mut server = mockito::Server::new_async().await; + let catalog = server + .mock("GET", "/models") + .with_status(200) + .with_header("content-type", "application/json") + .with_body(r#"{"data":[{"id":"current"}]}"#) + .create_async() + .await; + let metadata = server.mock("GET", "/model/info") + .with_status(200).with_header("content-type", "application/json") + .with_body(r#"{"data":[{"model_name":"removed","model_info":{"max_input_tokens":100000}},{"model_name":"current","model_info":{"max_input_tokens":128000}}]}"#) + .create_async().await; + let models = fetch_litellm_catalog(&server.url(), None, &reqwest::Client::new()) + .await + .unwrap(); + assert_eq!(models.len(), 1); + assert_eq!(models[0].model_id, "current"); + catalog.assert_async().await; + metadata.assert_async().await; +} + +#[tokio::test] +async fn model_refresh_distinguishes_empty_catalog_from_failure() { + let mut server = mockito::Server::new_async().await; + for (status, body, succeeds) in [ + (200, r#"{"data":[]}"#, true), + (200, r#"{"error":"unavailable"}"#, false), + (200, r#"{"data":[{}]}"#, false), + (401, r#"{"error":"unauthorized"}"#, false), + ] { + let response = server + .mock("GET", "/models") + .with_status(status) + .with_header("content-type", "application/json") + .with_body(body) + .create_async() + .await; + let result = fetch_litellm_catalog(&server.url(), None, &reqwest::Client::new()).await; + assert_eq!(result.is_ok(), succeeds); + if let Ok(models) = result { + assert!(models.is_empty()); + } + response.assert_async().await; + response.remove_async().await; + } +} diff --git a/app/src/settings_view/provider_setup_view.rs b/app/src/settings_view/provider_setup_view.rs index 7eb67d8f..8a4ac13f 100644 --- a/app/src/settings_view/provider_setup_view.rs +++ b/app/src/settings_view/provider_setup_view.rs @@ -19,7 +19,7 @@ use pathfinder_color::ColorU; #[cfg(not(target_family = "wasm"))] use crate::ai::chatgpt_auth::{ChatGPTAuthModel, ChatGPTAuthModelEvent, ChatGPTAuthState}; use crate::ai::harness_display; -use crate::ai::llms::{merge_discovered_provider_models, LLMPreferences}; +use crate::ai::llms::{merge_discovered_provider_models, replace_litellm_catalog, LLMPreferences}; use crate::appearance::Appearance; use crate::editor::{ EditorView, Event as EditorEvent, SingleLineEditorOptions, TextColors, TextOptions, @@ -1000,11 +1000,16 @@ impl ProviderSetupView { let provider = self.draft_provider(); let existing_models = self.draft_models.clone(); + let authoritative_catalog = provider.kind == OpenAIProviderKind::LiteLLM; ctx.spawn( async move { LLMPreferences::discover_openai_provider_models(provider).await }, move |me, result, ctx| match result { Ok(models) => { - me.draft_models = merge_discovered_provider_models(&existing_models, models); + me.draft_models = if authoritative_catalog { + replace_litellm_catalog(&existing_models, models) + } else { + merge_discovered_provider_models(&existing_models, models) + }; me.discovery_state = DiscoveryState::Idle; me.step = ProviderSetupStep::Models; me.sync_model_switches(ctx);