feat(gcp-vertex): add model list with org policy filtering (#6393)

Signed-off-by: rabi <ramishra@redhat.com>

crates/goose-cli/src/commands/configure.rs

@@ -436,12 +436,12 @@ fn select_model_from_list(
     provider_meta: &goose::providers::base::ProviderMetadata,
 ) -> anyhow::Result<String> {
     const MAX_MODELS: usize = 10;
+    const UNLISTED_MODEL_KEY: &str = "__unlisted__";
+
     // Smart model selection:
     // If we have more than MAX_MODELS models, show the recommended models with additional search option.
     // Otherwise, show all models without search.
-
     if models.len() > MAX_MODELS {
-        // Get recommended models from provider metadata
         let recommended_models: Vec<String> = provider_meta
             .known_models
             .iter()
@@ -464,12 +464,22 @@ fn select_model_from_list(
                 ),
             );
 
+            if provider_meta.allows_unlisted_models {
+                model_items.push((
+                    UNLISTED_MODEL_KEY.to_string(),
+                    "Enter a model not listed...".to_string(),
+                    "",
+                ));
+            }
+
             let selection = cliclack::select("Select a model:")
                 .items(&model_items)
                 .interact()?;
 
             if selection == "search_all" {
                 Ok(interactive_model_search(models)?)
+            } else if selection == UNLISTED_MODEL_KEY {
+                prompt_unlisted_model(provider_meta)
             } else {
                 Ok(selection)
             }
@@ -477,19 +487,45 @@ fn select_model_from_list(
             Ok(interactive_model_search(models)?)
         }
     } else {
-        // just a few models, show all without search for better UX
+        let mut model_items: Vec<(String, String, &str)> =
-        Ok(cliclack::select("Select a model:")
+            models.iter().map(|m| (m.clone(), m.clone(), "")).collect();
-            .items(
+
-                &models
+        if provider_meta.allows_unlisted_models {
-                    .iter()
+            model_items.push((
-                    .map(|m| (m, m.as_str(), ""))
+                UNLISTED_MODEL_KEY.to_string(),
-                    .collect::<Vec<_>>(),
+                "Enter a model not listed...".to_string(),
-            )
+                "",
-            .interact()?
+            ));
-            .to_string())
+        }
+
+        let selection = cliclack::select("Select a model:")
+            .items(&model_items)
+            .interact()?;
+
+        if selection == UNLISTED_MODEL_KEY {
+            prompt_unlisted_model(provider_meta)
+        } else {
+            Ok(selection)
+        }
     }
 }
 
+fn prompt_unlisted_model(
+    provider_meta: &goose::providers::base::ProviderMetadata,
+) -> anyhow::Result<String> {
+    let model: String = cliclack::input("Enter the model name:")
+        .placeholder(&provider_meta.default_model)
+        .validate(|input: &String| {
+            if input.trim().is_empty() {
+                Err("Please enter a model name")
+            } else {
+                Ok(())
+            }
+        })
+        .interact()?;
+    Ok(model.trim().to_string())
+}
+
 fn try_store_secret(config: &Config, key_name: &str, value: String) -> anyhow::Result<bool> {
     match config.set_secret(key_name, &value) {
         Ok(_) => Ok(true),

crates/goose/src/providers/base.rs

@@ -108,6 +108,9 @@ pub struct ProviderMetadata {
     pub model_doc_link: String,
     /// Required configuration keys
     pub config_keys: Vec<ConfigKey>,
+    /// Whether this provider allows entering model names not in the fetched list
+    #[serde(default)]
+    pub allows_unlisted_models: bool,
 }
 
 impl ProviderMetadata {
@@ -138,6 +141,7 @@ impl ProviderMetadata {
                 .collect(),
             model_doc_link: model_doc_link.to_string(),
             config_keys,
+            allows_unlisted_models: false,
         }
     }
 
@@ -158,6 +162,7 @@ impl ProviderMetadata {
             known_models: models,
             model_doc_link: model_doc_link.to_string(),
             config_keys,
+            allows_unlisted_models: false,
         }
     }
 
@@ -170,8 +175,15 @@ impl ProviderMetadata {
             known_models: vec![],
             model_doc_link: "".to_string(),
             config_keys: vec![],
+            allows_unlisted_models: false,
         }
     }
+
+    /// Set allows_unlisted_models flag (builder pattern)
+    pub fn with_unlisted_models(mut self) -> Self {
+        self.allows_unlisted_models = true;
+        self
+    }
 }
 
 /// Configuration key metadata for provider setup

crates/goose/src/providers/formats/gcpvertexai.rs

@@ -70,77 +70,49 @@ pub enum ModelError {
     UnsupportedLocation(String),
 }
 
+/// Default model for GCP Vertex AI.
+pub const DEFAULT_MODEL: &str = "gemini-2.5-flash";
+
+pub const KNOWN_MODELS: &[&str] = &[
+    "claude-opus-4-5@20251101",
+    "claude-sonnet-4-5@20250929",
+    "claude-opus-4-1@20250805",
+    "claude-haiku-4-5@20251001",
+    "claude-opus-4@20250514",
+    "claude-sonnet-4@20250514",
+    "claude-3-5-haiku@20241022",
+    "claude-3-haiku@20240307",
+    "gemini-3-pro",
+    "gemini-3-flash",
+    "gemini-2.5-pro",
+    "gemini-2.5-flash",
+    "gemini-2.5-flash-lite",
+    "gemini-2.0-flash",
+    "gemini-2.0-flash-lite",
+];
+
 /// Represents available GCP Vertex AI models for goose.
 ///
-/// This enum encompasses different model families and their versions
+/// This enum encompasses different model families that are supported
-/// that are supported in the GCP Vertex AI platform.
+/// in the GCP Vertex AI platform.
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub enum GcpVertexAIModel {
-    /// Claude model family with specific versions
+    /// Claude model family
-    Claude(ClaudeVersion),
+    Claude(String),
-    /// Gemini model family with specific versions
+    /// Gemini model family
-    Gemini(GeminiVersion),
+    Gemini(String),
     /// MaaS (Model as a Service) models from Model Garden
     /// Contains (publisher, full_model_name)
     MaaS(String, String),
 }
 
-/// Represents available versions of the Claude model for goose.
-#[derive(Debug, Clone, PartialEq, Eq)]
-pub enum ClaudeVersion {
-    /// Claude Sonnet 4
-    Sonnet4,
-    /// Claude Opus 4
-    Opus4,
-    /// Generic Claude model for custom or new versions
-    Generic(String),
-}
-
-/// Represents available versions of the Gemini model for goose.
-#[derive(Debug, Clone, PartialEq, Eq)]
-pub enum GeminiVersion {
-    /// Gemini 1.5 Pro version
-    Pro15,
-    /// Gemini 2.0 Flash version
-    Flash20,
-    /// Gemini 2.0 Pro Experimental version
-    Pro20Exp,
-    /// Gemini 2.5 Pro Experimental version
-    Pro25Exp,
-    /// Gemini 2.5 Flash Preview version
-    Flash25Preview,
-    /// Gemini 2.5 Pro Preview version
-    Pro25Preview,
-    /// Gemini 2.5 Flash version
-    Flash25,
-    /// Gemini 2.5 Pro version
-    Pro25,
-    /// Generic Gemini model for custom or new versions
-    Generic(String),
-}
-
 impl fmt::Display for GcpVertexAIModel {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        let model_id = match self {
+        match self {
-            Self::Claude(version) => match version {
+            Self::Claude(name) => write!(f, "{name}"),
-                ClaudeVersion::Sonnet4 => "claude-sonnet-4@20250514",
+            Self::Gemini(name) => write!(f, "{name}"),
-                ClaudeVersion::Opus4 => "claude-opus-4@20250514",
+            Self::MaaS(_, name) => write!(f, "{name}"),
-                ClaudeVersion::Generic(name) => name,
+        }
-            },
-            Self::Gemini(version) => match version {
-                GeminiVersion::Pro15 => "gemini-1.5-pro-002",
-                GeminiVersion::Flash20 => "gemini-2.0-flash-001",
-                GeminiVersion::Pro20Exp => "gemini-2.0-pro-exp-02-05",
-                GeminiVersion::Pro25Exp => "gemini-2.5-pro-exp-03-25",
-                GeminiVersion::Flash25Preview => "gemini-2.5-flash-preview-05-20",
-                GeminiVersion::Pro25Preview => "gemini-2.5-pro-preview-05-06",
-                GeminiVersion::Flash25 => "gemini-2.5-flash",
-                GeminiVersion::Pro25 => "gemini-2.5-pro",
-                GeminiVersion::Generic(name) => name,
-            },
-            Self::MaaS(_, model_name) => model_name,
-        };
-        write!(f, "{model_id}")
     }
 }
 
@@ -164,35 +136,19 @@ impl TryFrom<&str> for GcpVertexAIModel {
     type Error = ModelError;
 
     fn try_from(s: &str) -> Result<Self, Self::Error> {
-        // Known models
+        if s.starts_with("claude-") {
-        match s {
+            Ok(Self::Claude(s.to_string()))
-            "claude-sonnet-4@20250514" => Ok(Self::Claude(ClaudeVersion::Sonnet4)),
+        } else if s.starts_with("gemini-") {
-            "claude-opus-4@20250514" => Ok(Self::Claude(ClaudeVersion::Opus4)),
+            Ok(Self::Gemini(s.to_string()))
-            "gemini-1.5-pro-002" => Ok(Self::Gemini(GeminiVersion::Pro15)),
+        } else if s.ends_with("-maas") {
-            "gemini-2.0-flash-001" => Ok(Self::Gemini(GeminiVersion::Flash20)),
+            let publisher = s
-            "gemini-2.0-pro-exp-02-05" => Ok(Self::Gemini(GeminiVersion::Pro20Exp)),
+                .split('-')
-            "gemini-2.5-pro-exp-03-25" => Ok(Self::Gemini(GeminiVersion::Pro25Exp)),
+                .next()
-            "gemini-2.5-flash-preview-05-20" => Ok(Self::Gemini(GeminiVersion::Flash25Preview)),
+                .ok_or_else(|| ModelError::UnsupportedModel(s.to_string()))?
-            "gemini-2.5-pro-preview-05-06" => Ok(Self::Gemini(GeminiVersion::Pro25Preview)),
+                .to_string();
-            "gemini-2.5-flash" => Ok(Self::Gemini(GeminiVersion::Flash25)),
+            Ok(Self::MaaS(publisher, s.to_string()))
-            "gemini-2.5-pro" => Ok(Self::Gemini(GeminiVersion::Pro25)),
+        } else {
-            // MaaS models (Model as a Service from Model Garden)
+            Err(ModelError::UnsupportedModel(s.to_string()))
-            _ if s.ends_with("-maas") => {
-                let publisher = s
-                    .split('-')
-                    .next()
-                    .ok_or_else(|| ModelError::UnsupportedModel(s.to_string()))?
-                    .to_string();
-                Ok(Self::MaaS(publisher, s.to_string()))
-            }
-            // Generic models based on prefix matching
-            _ if s.starts_with("claude-") => {
-                Ok(Self::Claude(ClaudeVersion::Generic(s.to_string())))
-            }
-            _ if s.starts_with("gemini-") => {
-                Ok(Self::Gemini(GeminiVersion::Generic(s.to_string())))
-            }
-            _ => Err(ModelError::UnsupportedModel(s.to_string())),
         }
     }
 }
@@ -397,21 +353,16 @@ mod tests {
 
     #[test]
     fn test_model_parsing() -> Result<()> {
-        let valid_models = [
+        let claude = GcpVertexAIModel::try_from("claude-sonnet-4@20250514")?;
-            "claude-sonnet-4-20250514",
+        assert!(matches!(claude, GcpVertexAIModel::Claude(_)));
-            "claude-sonnet-4@20250514",
+        assert_eq!(claude.to_string(), "claude-sonnet-4@20250514");
-            "gemini-1.5-pro-002",
-            "gemini-2.0-flash-001",
-            "gemini-2.0-pro-exp-02-05",
-            "gemini-2.5-pro-exp-03-25",
-            "gemini-2.5-flash-preview-05-20",
-            "gemini-2.5-pro-preview-05-06",
-        ];
 
-        for model_id in valid_models {
+        let gemini = GcpVertexAIModel::try_from("gemini-2.5-flash")?;
-            let model = GcpVertexAIModel::try_from(model_id)?;
+        assert!(matches!(gemini, GcpVertexAIModel::Gemini(_)));
-            assert_eq!(model.to_string(), model_id);
+        assert_eq!(gemini.to_string(), "gemini-2.5-flash");
-        }
+
+        let maas = GcpVertexAIModel::try_from("qwen-maas")?;
+        assert!(matches!(maas, GcpVertexAIModel::MaaS(_, _)));
 
         assert!(GcpVertexAIModel::try_from("unsupported-model").is_err());
         Ok(())
@@ -419,71 +370,24 @@ mod tests {
 
     #[test]
     fn test_default_locations() -> Result<()> {
-        let test_cases = [
+        let claude_model = GcpVertexAIModel::try_from("claude-sonnet-4@20250514")?;
-            ("claude-sonnet-4-20250514", GcpLocation::Ohio),
+        assert_eq!(claude_model.known_location(), GcpLocation::Ohio);
-            ("claude-sonnet-4@20250514", GcpLocation::Ohio),
-            ("gemini-1.5-pro-002", GcpLocation::Iowa),
-            ("gemini-2.0-flash-001", GcpLocation::Iowa),
-            ("gemini-2.0-pro-exp-02-05", GcpLocation::Iowa),
-            ("gemini-2.5-pro-exp-03-25", GcpLocation::Iowa),
-            ("gemini-2.5-flash-preview-05-20", GcpLocation::Iowa),
-            ("gemini-2.5-pro-preview-05-06", GcpLocation::Iowa),
-        ];
 
-        for (model_id, expected_location) in test_cases {
+        let gemini_model = GcpVertexAIModel::try_from("gemini-2.5-flash")?;
-            let model = GcpVertexAIModel::try_from(model_id)?;
+        assert_eq!(gemini_model.known_location(), GcpLocation::Iowa);
-            assert_eq!(
-                model.known_location(),
-                expected_location,
-                "Model {model_id} should have default location {expected_location:?}",
-            );
-
-            let context = RequestContext::new(model_id)?;
-            assert_eq!(
-                context.model.known_location(),
-                expected_location,
-                "RequestContext for {model_id} should have default location {expected_location:?}",
-            );
-        }
 
         Ok(())
     }
 
     #[test]
-    fn test_generic_model_parsing() -> Result<()> {
+    fn test_unknown_model_parsing() -> Result<()> {
-        // Test generic Claude models
+        let model = GcpVertexAIModel::try_from("claude-future-version")?;
-        let claude_models = [
+        assert!(matches!(model, GcpVertexAIModel::Claude(_)));
-            "claude-3-8-apex@20250301",
+        assert_eq!(model.to_string(), "claude-future-version");
-            "claude-new-version",
-            "claude-experimental",
-        ];
 
-        for model_id in claude_models {
+        let model = GcpVertexAIModel::try_from("gemini-4.0-ultra")?;
-            let model = GcpVertexAIModel::try_from(model_id)?;
+        assert!(matches!(model, GcpVertexAIModel::Gemini(_)));
-            match model {
+        assert_eq!(model.to_string(), "gemini-4.0-ultra");
-                GcpVertexAIModel::Claude(ClaudeVersion::Generic(ref name)) => {
-                    assert_eq!(name, model_id);
-                }
-                _ => panic!("Expected Claude generic model for {model_id}"),
-            }
-            assert_eq!(model.to_string(), model_id);
-            assert_eq!(model.known_location(), GcpLocation::Ohio);
-        }
-
-        // Test generic Gemini models
-        let gemini_models = ["gemini-3-pro", "gemini-2.0-flash", "gemini-experimental"];
-
-        for model_id in gemini_models {
-            let model = GcpVertexAIModel::try_from(model_id)?;
-            match model {
-                GcpVertexAIModel::Gemini(GeminiVersion::Generic(ref name)) => {
-                    assert_eq!(name, model_id);
-                }
-                _ => panic!("Expected Gemini generic model for {model_id}"),
-            }
-            assert_eq!(model.to_string(), model_id);
-            assert_eq!(model.known_location(), GcpLocation::Iowa);
-        }
 
         Ok(())
     }

crates/goose/src/providers/gcpvertexai.rs

@@ -19,11 +19,9 @@ use crate::providers::base::{ConfigKey, MessageStream, Provider, ProviderMetadat
 
 use crate::providers::errors::ProviderError;
 use crate::providers::formats::gcpvertexai::{
-    create_request, get_usage, response_to_message, response_to_streaming_message, ClaudeVersion,
+    create_request, get_usage, response_to_message, response_to_streaming_message, GcpLocation,
-    GcpVertexAIModel, GeminiVersion, ModelProvider, RequestContext,
+    ModelProvider, RequestContext, DEFAULT_MODEL, KNOWN_MODELS,
 };
-
-use crate::providers::formats::gcpvertexai::GcpLocation::Iowa;
 use crate::providers::gcpauth::GcpAuth;
 use crate::providers::retry::RetryConfig;
 use crate::providers::utils::RequestLog;
@@ -225,7 +223,7 @@ impl GcpVertexAIProvider {
             .get_param("GCP_LOCATION")
             .ok()
             .filter(|location: &String| !location.trim().is_empty())
-            .unwrap_or_else(|| Iowa.to_string()))
+            .unwrap_or_else(|| GcpLocation::Iowa.to_string()))
     }
 
     /// Retrieves an authentication token for API requests.
@@ -430,43 +428,118 @@ impl GcpVertexAIProvider {
             _ => result,
         }
     }
+
+    async fn filter_by_org_policy(&self, models: Vec<String>) -> Vec<String> {
+        let Ok(auth_header) = self.get_auth_header().await else {
+            tracing::debug!("Could not get auth header for org policy check, returning all models");
+            return models;
+        };
+
+        let url = format!(
+            "https://cloudresourcemanager.googleapis.com/v1/projects/{}:getEffectiveOrgPolicy",
+            self.project_id
+        );
+
+        let payload = serde_json::json!({
+            "constraint": "constraints/vertexai.allowedModels"
+        });
+
+        let response = match self
+            .client
+            .post(&url)
+            .header("Authorization", &auth_header)
+            .json(&payload)
+            .send()
+            .await
+        {
+            Ok(r) => r,
+            Err(e) => {
+                tracing::debug!("Failed to fetch org policy: {e}, returning all models");
+                return models;
+            }
+        };
+
+        let json = match response.json::<Value>().await {
+            Ok(j) => j,
+            Err(e) => {
+                tracing::debug!("Failed to parse org policy response: {e}, returning all models");
+                return models;
+            }
+        };
+
+        let allowed_patterns: Vec<String> = json
+            .get("listPolicy")
+            .and_then(|lp| lp.get("allowedValues"))
+            .and_then(|av| av.as_array())
+            .map(|arr| {
+                arr.iter()
+                    .filter_map(|v| v.as_str())
+                    .map(|s| s.to_string())
+                    .collect()
+            })
+            .unwrap_or_default();
+
+        if allowed_patterns.is_empty() {
+            return models;
+        }
+
+        models
+            .into_iter()
+            .filter(|model| Self::is_model_allowed(model, &allowed_patterns))
+            .collect()
+    }
+
+    fn is_model_allowed(model: &str, allowed_patterns: &[String]) -> bool {
+        let publisher = if model.starts_with("claude-") {
+            "anthropic"
+        } else if model.starts_with("gemini-") {
+            "google"
+        } else {
+            return true;
+        };
+
+        for pattern in allowed_patterns {
+            if pattern.contains(&format!("publishers/{publisher}/models/*")) {
+                return true;
+            }
+
+            let pattern_model = pattern
+                .split("/models/")
+                .nth(1)
+                .map(|s| s.trim_end_matches(":predict").trim_end_matches(":*"));
+
+            if let Some(pattern_model) = pattern_model {
+                if model == pattern_model || model.starts_with(&format!("{pattern_model}@")) {
+                    return true;
+                }
+            }
+        }
+
+        false
+    }
 }
 
 #[async_trait]
 impl Provider for GcpVertexAIProvider {
-    /// Returns metadata about the GCP Vertex AI provider.
     fn metadata() -> ProviderMetadata
     where
         Self: Sized,
     {
-        let model_strings: Vec<String> = [
-            GcpVertexAIModel::Claude(ClaudeVersion::Sonnet4),
-            GcpVertexAIModel::Claude(ClaudeVersion::Opus4),
-            GcpVertexAIModel::Gemini(GeminiVersion::Pro15),
-            GcpVertexAIModel::Gemini(GeminiVersion::Flash20),
-            GcpVertexAIModel::Gemini(GeminiVersion::Pro20Exp),
-            GcpVertexAIModel::Gemini(GeminiVersion::Pro25Exp),
-            GcpVertexAIModel::Gemini(GeminiVersion::Flash25Preview),
-            GcpVertexAIModel::Gemini(GeminiVersion::Pro25Preview),
-            GcpVertexAIModel::Gemini(GeminiVersion::Flash25),
-            GcpVertexAIModel::Gemini(GeminiVersion::Pro25),
-        ]
-        .iter()
-        .map(|model| model.to_string())
-        .collect();
-
-        let known_models: Vec<&str> = model_strings.iter().map(|s| s.as_str()).collect();
-
         ProviderMetadata::new(
             "gcp_vertex_ai",
             "GCP Vertex AI",
             "Access variety of AI models such as Claude, Gemini through Vertex AI",
-            "gemini-2.5-flash",
+            DEFAULT_MODEL,
-            known_models,
+            KNOWN_MODELS.to_vec(),
             GCP_VERTEX_AI_DOC_URL,
             vec![
                 ConfigKey::new("GCP_PROJECT_ID", true, false, None),
-                ConfigKey::new("GCP_LOCATION", true, false, Some(Iowa.to_string().as_str())),
+                ConfigKey::new(
+                    "GCP_LOCATION",
+                    true,
+                    false,
+                    Some(&GcpLocation::Iowa.to_string()),
+                ),
                 ConfigKey::new(
                     "GCP_MAX_RETRIES",
                     false,
@@ -493,6 +566,7 @@ impl Provider for GcpVertexAIProvider {
                 ),
             ],
         )
+        .with_unlisted_models()
     }
 
     fn get_name(&self) -> &str {
@@ -587,6 +661,12 @@ impl Provider for GcpVertexAIProvider {
             }
         }))
     }
+
+    async fn fetch_supported_models(&self) -> Result<Option<Vec<String>>, ProviderError> {
+        let models: Vec<String> = KNOWN_MODELS.iter().map(|s| s.to_string()).collect();
+        let filtered = self.filter_by_org_policy(models).await;
+        Ok(Some(filtered))
+    }
 }
 
 #[cfg(test)]
@@ -705,15 +785,9 @@ mod tests {
     #[test]
     fn test_provider_metadata() {
         let metadata = GcpVertexAIProvider::metadata();
-        let model_names: Vec<String> = metadata
+        assert!(!metadata.known_models.is_empty());
-            .known_models
+        assert_eq!(metadata.default_model, "gemini-2.5-flash");
-            .iter()
-            .map(|m| m.name.clone())
-            .collect();
-        assert!(model_names.contains(&"claude-sonnet-4@20250514".to_string()));
-        assert!(model_names.contains(&"gemini-1.5-pro-002".to_string()));
-        assert!(model_names.contains(&"gemini-2.5-pro".to_string()));
-        // Should contain the original 2 config keys plus 4 new retry-related ones
         assert_eq!(metadata.config_keys.len(), 6);
+        assert!(metadata.allows_unlisted_models);
     }
 }

crates/goose/src/providers/provider_registry.rs

@@ -114,6 +114,7 @@ impl ProviderRegistry {
             known_models,
             model_doc_link: base_metadata.model_doc_link,
             config_keys,
+            allows_unlisted_models: false,
         };
 
         self.entries.insert(

crates/goose/tests/agent.rs

@@ -397,6 +397,7 @@ mod tests {
                     known_models: vec![],
                     model_doc_link: "".to_string(),
                     config_keys: vec![],
+                    allows_unlisted_models: false,
                 }
             }
 

ui/desktop/openapi.json

@@ -4735,6 +4735,10 @@
           "config_keys"
         ],
         "properties": {
+          "allows_unlisted_models": {
+            "type": "boolean",
+            "description": "Whether this provider allows entering model names not in the fetched list"
+          },
           "config_keys": {
             "type": "array",
             "items": {

ui/desktop/src/api/types.gen.ts

@@ -630,6 +630,10 @@ export type ProviderEngine = 'openai' | 'ollama' | 'anthropic';
  * Metadata about a provider's configuration requirements and capabilities
  */
 export type ProviderMetadata = {
+    /**
+     * Whether this provider allows entering model names not in the fetched list
+     */
+    allows_unlisted_models?: boolean;
     /**
      * Required configuration keys
      */

ui/desktop/src/components/settings/models/subcomponents/LeadWorkerSettings.tsx

@@ -119,12 +119,17 @@ export function LeadWorkerSettings({ isOpen, onClose }: LeadWorkerSettingsProps)
                 });
               });
             }
+            // Add custom model option for all non-Custom providers
+            if (p.provider_type !== 'Custom') {
+              options.push({
+                value: `__custom__:${p.name}`,
+                label: 'Enter a model not listed...',
+                provider: p.name,
+              });
+            }
           });
         }
 
-        // Append a simple "custom" option to enable free-text entry
-        options.push({ value: '__custom__', label: 'Use custom model…', provider: '' });
-
         setModelOptions(options);
       } catch (error) {
         console.error('Error loading configuration:', error);
@@ -241,9 +246,10 @@ export function LeadWorkerSettings({ isOpen, onClose }: LeadWorkerSettingsProps)
                   onChange={(newValue: unknown) => {
                     const option = newValue as { value: string; provider: string } | null;
                     if (option) {
-                      if (option.value === '__custom__') {
+                      if (option.value.startsWith('__custom__')) {
                         setIsLeadCustomModel(true);
                         setLeadModel('');
+                        setLeadProvider(option.provider);
                         return;
                       }
                       setLeadModel(option.value);
@@ -294,9 +300,10 @@ export function LeadWorkerSettings({ isOpen, onClose }: LeadWorkerSettingsProps)
                   onChange={(newValue: unknown) => {
                     const option = newValue as { value: string; provider: string } | null;
                     if (option) {
-                      if (option.value === '__custom__') {
+                      if (option.value.startsWith('__custom__')) {
                         setIsWorkerCustomModel(true);
                         setWorkerModel('');
+                        setWorkerProvider(option.provider);
                         return;
                       }
                       setWorkerModel(option.value);

ui/desktop/src/components/settings/models/subcomponents/SwitchModelModal.tsx

@@ -213,29 +213,34 @@ export const SwitchModelModal = ({
         const errors: string[] = [];
 
         results.forEach(({ provider: p, models, error }) => {
+          const modelList = error
+            ? (p.metadata.known_models?.map(({ name }) => name) || [])
+            : (models || []);
+
           if (error) {
             errors.push(error);
-            // Fallback to metadata known_models on error
+          }
-            if (p.metadata.known_models && p.metadata.known_models.length > 0) {
+
-              groupedOptions.push({
+          const options: { value: string; label: string; provider: string; providerType: ProviderType }[] =
-                options: p.metadata.known_models.map(({ name }) => ({
+            modelList.map((m) => ({
-                  value: name,
+              value: m,
-                  label: name,
+              label: m,
-                  providerType: p.provider_type,
+              provider: p.name,
-                  provider: p.name,
+              providerType: p.provider_type,
-                })),
+            }));
-              });
+
-            }
+          if (p.metadata.allows_unlisted_models && p.provider_type !== 'Custom') {
-          } else if (models && models.length > 0) {
+            options.push({
-            groupedOptions.push({
+              value: 'custom',
-              options: models.map((m) => ({
+              label: 'Enter a model not listed...',
-                value: m,
+              provider: p.name,
-                label: m,
+              providerType: p.provider_type,
-                provider: p.name,
-                providerType: p.provider_type,
-              })),
             });
           }
+
+          if (options.length > 0) {
+            groupedOptions.push({ options });
+          }
         });
 
         // Log errors if any providers failed (don't show to user)
@@ -243,20 +248,6 @@ export const SwitchModelModal = ({
           console.error('Provider model fetch errors:', errors);
         }
 
-        // Add the "Custom model" option to each provider group
-        groupedOptions.forEach((group) => {
-          const option = group.options[0];
-          const providerName = option?.provider;
-          if (providerName && option?.providerType !== 'Custom') {
-            group.options.push({
-              value: 'custom',
-              label: 'Use custom model',
-              provider: providerName,
-              providerType: option?.providerType,
-            });
-          }
-        });
-
         setModelOptions(groupedOptions);
         setOriginalModelOptions(groupedOptions);
       } catch (error: unknown) {
@@ -293,6 +284,7 @@ export const SwitchModelModal = ({
     if (selectedOption?.value === 'custom') {
       setIsCustomModel(true);
       setModel('');
+      setProvider(selectedOption.provider);
       setUserClearedModel(false);
     } else if (selectedOption === null) {
       // User cleared the selection
@@ -302,6 +294,7 @@ export const SwitchModelModal = ({
     } else {
       setIsCustomModel(false);
       setModel(selectedOption?.value || '');
+      setProvider(selectedOption?.provider || '');
       setUserClearedModel(false);
     }
   };