diff --git a/src/app.rs b/src/app.rs index 0689c8b..55a05ee 100644 --- a/src/app.rs +++ b/src/app.rs @@ -157,6 +157,23 @@ fn has_single_user_message(chat: &Chat) -> bool { == 1 } +fn models_dialog_provider_ids() -> Option> { + let mut signature = crate::persistence::AuthDAO::new() + .and_then(|dao| dao.load()) + .ok()? + .into_keys() + .map(|provider_id| format!("auth:{provider_id}")) + .collect::>(); + + if let Ok(discovery) = crate::model::discovery::Discovery::new() { + signature.extend(discovery.custom_provider_dialog_signature()); + } + + signature.sort(); + signature.dedup(); + Some(signature) +} + #[derive(Debug, Clone, PartialEq, Eq)] pub struct StreamingRetryStatus { pub attempt: usize, @@ -315,6 +332,7 @@ enum ModelsTaskKind { struct ModelsTaskMessage { kind: ModelsTaskKind, result: crate::command::registry::CommandResult, + provider_signature: Option>, } #[derive(Debug)] @@ -1063,7 +1081,6 @@ impl App { ); let agent_steps = agent_registry.max_steps_map(); let provider_timeouts = loaded_config.merged_config.provider_timeouts.clone(); - let theme_for_colors = themes .get(current_theme_index) .or_else(|| themes.first()) @@ -1082,7 +1099,10 @@ impl App { .with_permission_rules(loaded_config.merged_config.permission_rules.clone()) .with_agent_permission_rules(agent_registry.permission_rules_map()); - let discovery = crate::model::discovery::Discovery::new().ok(); + let discovery = crate::model::discovery::Discovery::new_with_custom(Some( + loaded_config.merged_config.custom_providers.clone(), + )) + .ok(); let now = std::time::Instant::now(); Ok(Self { @@ -6741,6 +6761,16 @@ impl App { Ok(providers) => providers, Err(_) => return, }; + let connected_provider_ids = connected_providers + .keys() + .cloned() + .collect::>(); + + let discovery = Discovery::new(); + let configured_provider_ids = discovery + .as_ref() + .map(Discovery::custom_provider_ids) + .unwrap_or_default(); let include_runtime = crate::model::extensions::ModelExtensions::runtime() .iter() @@ -6760,17 +6790,22 @@ impl App { ) }); - if connected_providers.is_empty() && !include_runtime && !include_unauthenticated_free { + if connected_providers.is_empty() + && configured_provider_ids.is_empty() + && !include_runtime + && !include_unauthenticated_free + { return; } let has_persistent = connected_providers.keys().any(|provider_id| { !crate::model::extensions::ModelExtensions::is_runtime_provider(provider_id) - }) || include_unauthenticated_free + }) || !configured_provider_ids.is_empty() + || include_unauthenticated_free || connected_providers.is_empty(); let models = if has_persistent { - match Discovery::new() { + match discovery.as_ref() { Ok(discovery) => match tokio::task::block_in_place(|| { let rt = tokio::runtime::Handle::current(); rt.block_on(discovery.fetch_models()) @@ -6806,6 +6841,19 @@ impl App { } else { return; }; + let mut models = models; + if include_runtime { + let runtime_models = tokio::task::block_in_place(|| { + let rt = tokio::runtime::Handle::current(); + rt.block_on( + crate::model::extensions::ModelExtensions::runtime_models_for_dialog_cached_or_empty(), + ) + }); + crate::model::discovery::merge_dialog_models(&mut models, runtime_models); + } + if let Ok(discovery) = discovery.as_ref() { + discovery.apply_custom_models_to_dialog(&mut models); + } self.model_reasoning_options = models .iter() @@ -6826,8 +6874,11 @@ impl App { std::collections::HashMap::new(); let is_model_selectable = |model: &ModelType| { - connected_providers.contains_key(&model.provider_id) - || crate::model::extensions::ModelExtensions::is_available_without_connection(model) + crate::model::discovery::is_model_selectable( + model, + &connected_provider_ids, + &configured_provider_ids, + ) }; for model in &models { @@ -7081,14 +7132,7 @@ impl App { return true; } - let connected_provider_ids = crate::persistence::AuthDAO::new() - .and_then(|dao| dao.load()) - .map(|providers| { - let mut ids = providers.into_keys().collect::>(); - ids.sort(); - ids - }) - .ok(); + let connected_provider_ids = models_dialog_provider_ids(); if kind == ModelsTaskKind::Load && parsed.args.is_empty() @@ -7129,12 +7173,17 @@ impl App { } let parsed = parsed.clone(); + let provider_signature = models_dialog_provider_ids(); tokio::spawn(async move { let result = match kind { ModelsTaskKind::Load => crate::command::handlers::load_models(parsed).await, ModelsTaskKind::Refresh => crate::command::handlers::refresh_models().await, }; - let _ = sender.send(ModelsTaskMessage { kind, result }); + let _ = sender.send(ModelsTaskMessage { + kind, + result, + provider_signature, + }); }); true } @@ -8076,7 +8125,12 @@ impl App { } for event in events { - match (event.kind, event.result) { + let ModelsTaskMessage { + kind, + result, + provider_signature, + } = event; + match (kind, result) { ( ModelsTaskKind::Load, crate::command::registry::CommandResult::ShowDialog { title, items }, @@ -8094,14 +8148,10 @@ impl App { }) .collect(); self.models_dialog_state.finish_loading(); - self.models_dialog_provider_ids = crate::persistence::AuthDAO::new() - .and_then(|dao| dao.load()) - .map(|providers| { - let mut ids = providers.into_keys().collect::>(); - ids.sort(); - ids - }) - .ok(); + let current_signature = models_dialog_provider_ids(); + self.models_dialog_provider_ids = (provider_signature == current_signature) + .then_some(provider_signature) + .flatten(); if self.overlay_focus == OverlayFocus::ModelsDialog { self.show_models_dialog(title, dialog_items); } @@ -11872,6 +11922,7 @@ mod tests { active: false, }], }, + provider_signature: models_dialog_provider_ids(), }) .unwrap(); app.models_receiver = Some(receiver); @@ -11900,15 +11951,8 @@ mod tests { active: false, }], ); - // Match the connected-provider snapshot used by the reopen cache check. - app.models_dialog_provider_ids = crate::persistence::AuthDAO::new() - .and_then(|dao| dao.load()) - .map(|providers| { - let mut ids = providers.into_keys().collect::>(); - ids.sort(); - ids - }) - .ok(); + // Match the auth/config snapshot used by the reopen cache check. + app.models_dialog_provider_ids = models_dialog_provider_ids(); app.models_dialog_state.dialog.hide(); app.overlay_focus = OverlayFocus::None; @@ -11945,6 +11989,7 @@ mod tests { .send(ModelsTaskMessage { kind: ModelsTaskKind::Refresh, result: crate::command::registry::CommandResult::Success(String::new()), + provider_signature: models_dialog_provider_ids(), }) .unwrap(); app.models_receiver = Some(receiver); diff --git a/src/command/handlers.rs b/src/command/handlers.rs index 5fd16b9..98b304b 100644 --- a/src/command/handlers.rs +++ b/src/command/handlers.rs @@ -106,7 +106,6 @@ pub fn handle_connect<'a>( Ok(providers) => providers, Err(e) => return CommandResult::Error(format!("Failed to load providers: {}", e)), }; - fn fallback_providers( ) -> std::collections::HashMap { use crate::model::discovery::Provider; @@ -274,6 +273,10 @@ pub async fn load_models(parsed: ParsedCommand) -> CommandResult { Ok(providers) => providers, Err(e) => return CommandResult::Error(format!("Failed to load providers: {}", e)), }; + let connected_provider_ids = connected_providers + .keys() + .cloned() + .collect::>(); let provider_filter_matches_runtime = provider_filter.as_deref().is_some_and(|filter| { let filter = filter.to_ascii_lowercase(); @@ -288,6 +291,16 @@ pub async fn load_models(parsed: ParsedCommand) -> CommandResult { }) }); + let discovery = Discovery::new(); + let configured_provider_ids = discovery + .as_ref() + .map(Discovery::custom_provider_ids) + .unwrap_or_default(); + let provider_filter_matches_configured = provider_filter.as_deref().is_some_and(|filter| { + discovery + .as_ref() + .is_ok_and(|discovery| discovery.custom_provider_matches_filter(filter)) + }); let provider_filter_matches_unauthenticated_free = provider_filter.as_deref().is_some_and( crate::model::extensions::ModelExtensions::unauthenticated_free_provider_matches_filter, ); @@ -300,16 +313,16 @@ pub async fn load_models(parsed: ParsedCommand) -> CommandResult { let has_persistent = connected_providers.keys().any(|provider_id| { !crate::model::extensions::ModelExtensions::is_runtime_provider(provider_id) }) || provider_filter.is_none() + || provider_filter_matches_configured || provider_filter_matches_unauthenticated_free; let snapshot_models = crate::model::effective_catalog::models_for_dialog() .ok() .flatten(); - let discovery = Discovery::new(); let mut models: Vec = if let Some(models) = snapshot_models.as_ref() { models.clone() } else if has_persistent { - match discovery { + match discovery.as_ref() { Ok(d) => match d.fetch_models().await { Ok(models) => models .into_iter() @@ -352,23 +365,31 @@ pub async fn load_models(parsed: ParsedCommand) -> CommandResult { Vec::new() }; + if let Ok(discovery) = discovery.as_ref() { + discovery.apply_custom_models_to_dialog(&mut models); + } + let mut runtime_errors = Vec::new(); - if snapshot_models.is_none() && has_runtime { + if has_runtime { let runtime_result = crate::model::extensions::ModelExtensions::runtime_models_for_dialog_cached().await; - models.extend(runtime_result.models); + crate::model::discovery::merge_dialog_models(&mut models, runtime_result.models); runtime_errors = runtime_result.errors; } if snapshot_models.is_none() && !models.is_empty() { - if let Err(err) = - crate::model::effective_catalog::publish_refreshed_models(models.clone()) - { - push_toast(Toast::new( - format!("Failed to seed model catalog cache: {}", err), - ToastLevel::Warning, - Some(std::time::Duration::from_secs(3)), - )); + if let Ok(discovery) = Discovery::new_with_custom(None) { + if let Ok(snapshot_models) = discovery.fetch_models().await { + if let Err(err) = + crate::model::effective_catalog::publish_refreshed_models(snapshot_models) + { + push_toast(Toast::new( + format!("Failed to seed model catalog cache: {}", err), + ToastLevel::Warning, + Some(std::time::Duration::from_secs(3)), + )); + } + } } } @@ -378,14 +399,14 @@ pub async fn load_models(parsed: ParsedCommand) -> CommandResult { std::collections::HashMap::new(); let is_model_selectable = |model: &ModelType| { - (connected_providers.contains_key(&model.provider_id) - || crate::model::extensions::ModelExtensions::is_available_without_connection( - model, - )) - && crate::model::extensions::ModelExtensions::model_matches_provider_filter( - model, - provider_filter.as_deref(), - ) + crate::model::discovery::is_model_selectable( + model, + &connected_provider_ids, + &configured_provider_ids, + ) && crate::model::extensions::ModelExtensions::model_matches_provider_filter( + model, + provider_filter.as_deref(), + ) }; for model in &models { @@ -798,7 +819,7 @@ pub async fn refresh_models() -> CommandResult { .sum::() + runtime_model_count; - let models = match crate::model::discovery::Discovery::new() { + let models = match crate::model::discovery::Discovery::new_with_custom(None) { Ok(discovery) => match discovery.fetch_models().await { Ok(models) => models, Err(err) => { diff --git a/src/config/configuration.rs b/src/config/configuration.rs index d7aaadf..e6bce66 100644 --- a/src/config/configuration.rs +++ b/src/config/configuration.rs @@ -8,6 +8,9 @@ use std::collections::{BTreeMap, BTreeSet, HashMap}; use std::fs; use std::path::{Path, PathBuf}; +// Re-export json5 for use in load_config_value +use json5; + pub fn discover_themes( xdg_config_home: &Path, project_root: &Path, @@ -103,6 +106,18 @@ fn list_json_files(dir: &Path) -> Vec { out } +fn parse_string_array(value: Option<&Value>) -> Vec { + value + .and_then(Value::as_array) + .into_iter() + .flatten() + .filter_map(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(str::to_ascii_lowercase) + .collect() +} + #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum SourceKind { OpenCode, @@ -377,12 +392,72 @@ pub struct MergedConfig { pub agent_permission_rules: HashMap, pub agent_steps: HashMap, pub provider_timeouts: HashMap, + pub custom_providers: HashMap, pub notifications: NotificationsConfig, pub images: ImagesConfig, pub websearch: WebsearchConfig, pub mcp: McpConfig, } +#[derive(Debug, Clone)] +pub struct CustomModelConfig { + pub name: Option, + pub context_window: Option, + pub max_tokens: Option, + pub attachment: Option, + pub reasoning: Option, + pub temperature: Option, + pub tool_call: Option, + pub modalities: Option, + pub launch: bool, +} + +#[derive(Debug, Clone)] +pub struct CustomModelModalities { + pub input: Vec, + pub output: Vec, +} + +#[derive(Debug, Clone)] +pub struct CustomProviderConfig { + pub name: Option, + pub npm: Option, + pub base_url: Option, + pub api_key: Option, + pub models: HashMap, +} + +impl CustomProviderConfig { + pub fn resolved_api_key(&self) -> Option { + resolve_api_key_value(self.api_key.as_deref(), |variable| { + std::env::var(variable).ok() + }) + } +} + +fn resolve_api_key_value( + value: Option<&str>, + get_env: impl FnOnce(&str) -> Option, +) -> Option { + let value = value?.trim(); + if value.is_empty() { + return None; + } + + if let Some(variable) = value + .strip_prefix("{env:") + .and_then(|value| value.strip_suffix('}')) + .map(str::trim) + .filter(|value| !value.is_empty()) + { + return get_env(variable) + .map(|value| value.trim().to_string()) + .filter(|value| !value.is_empty()); + } + + Some(value.to_string()) +} + #[derive(Debug, Clone)] pub struct LoadedConfig { pub merged_config: MergedConfig, @@ -824,20 +899,11 @@ fn load_config_value(path: &Path) -> Result { let content = fs::read_to_string(path) .with_context(|| format!("Failed to read config file {}", path.display()))?; - match path.extension().and_then(|s| s.to_str()) { - Some(ext) if ext.eq_ignore_ascii_case("jsonc") => { - let v: Value = json5::from_str(&content) - .with_context(|| format!("Invalid JSONC in {}", path.display()))?; - Ok(v) - } - _ => { - let v: Value = serde_json::from_str(&content) - .with_context(|| format!("Invalid JSON in {}", path.display()))?; - Ok(v) - } - } + // Use json5 for lenient parsing (handles trailing commas, comments, etc.) + let v: Value = json5::from_str(&content) + .with_context(|| format!("Invalid JSON/JSONC in {}", path.display()))?; + Ok(v) } - fn filter_top_level(value: Value, kind: SourceKind) -> Value { let mut map = match value { Value::Object(m) => m, @@ -1155,6 +1221,7 @@ fn parse_merged_config(merged: &Value, diagnostics: &mut ConfigDiagnostics) -> M ); out.sync_agent_derived_fields(); out.provider_timeouts = parse_provider_timeouts(obj.get("provider"), diagnostics); + out.custom_providers = parse_custom_providers(obj.get("provider"), diagnostics); let mut notifications = NotificationsConfig::default(); apply_notifications(obj.get("notifications"), &mut notifications, diagnostics); @@ -1785,6 +1852,133 @@ fn parse_provider_timeouts( out } +fn parse_custom_providers( + value: Option<&Value>, + _diagnostics: &mut ConfigDiagnostics, +) -> HashMap { + let mut out = HashMap::new(); + let Some(Value::Object(providers)) = value else { + return out; + }; + + for (provider_id, provider_val) in providers { + let Some(provider_obj) = provider_val.as_object() else { + continue; + }; + + let name = provider_obj + .get("name") + .and_then(Value::as_str) + .map(str::to_string); + + let npm = provider_obj + .get("npm") + .and_then(Value::as_str) + .map(|s| s.trim().to_string()) + .filter(|s| !s.is_empty()); + + let options = provider_obj.get("options").and_then(Value::as_object); + let base_url = options + .and_then(|options| options.get("baseURL")) + .and_then(Value::as_str) + .map(|s| s.trim().to_string()) + .filter(|s| !s.is_empty()); + let api_key = options + .and_then(|options| options.get("apiKey")) + .and_then(Value::as_str) + .map(|s| s.trim().to_string()) + .filter(|s| !s.is_empty()); + + let mut models = HashMap::new(); + if let Some(models_val) = provider_obj.get("models") { + if let Value::Object(model_map) = models_val { + for (model_id, model_val) in model_map { + let Some(model_obj) = model_val.as_object() else { + continue; + }; + + let model_name = model_obj + .get("name") + .and_then(Value::as_str) + .map(str::to_string); + + let limit = model_obj.get("limit").and_then(Value::as_object); + + let context_window = model_obj + .get("contextWindow") + .and_then(Value::as_u64) + .or_else(|| { + limit + .and_then(|limit| limit.get("context")) + .and_then(Value::as_u64) + }) + .and_then(|value| u32::try_from(value).ok()); + + let max_tokens = model_obj + .get("maxTokens") + .and_then(Value::as_u64) + .or_else(|| { + limit + .and_then(|limit| limit.get("output")) + .and_then(Value::as_u64) + }) + .and_then(|value| u32::try_from(value).ok()); + + let modalities = + model_obj + .get("modalities") + .and_then(Value::as_object) + .map(|modalities| CustomModelModalities { + input: parse_string_array(modalities.get("input")), + output: parse_string_array(modalities.get("output")), + }); + + let launch = model_obj + .get("_launch") + .and_then(Value::as_bool) + .unwrap_or(false); + + models.insert( + model_id.clone(), + CustomModelConfig { + name: model_name, + context_window, + max_tokens, + attachment: model_obj.get("attachment").and_then(Value::as_bool), + reasoning: model_obj.get("reasoning").and_then(Value::as_bool), + temperature: model_obj.get("temperature").and_then(Value::as_bool), + tool_call: model_obj.get("tool_call").and_then(Value::as_bool), + modalities, + launch, + }, + ); + } + } + } + + if name.is_none() + && npm.is_none() + && base_url.is_none() + && api_key.is_none() + && models.is_empty() + { + continue; + } + + out.insert( + provider_id.trim().to_ascii_lowercase(), + CustomProviderConfig { + name, + npm, + base_url, + api_key, + models, + }, + ); + } + + out +} fn parse_images(value: Option<&Value>, diagnostics: &mut ConfigDiagnostics) -> ImagesConfig { let mut images = ImagesConfig::default(); @@ -2708,4 +2902,106 @@ mod tests { assert_eq!(config.agent_steps.get("build"), Some(&42)); assert!(diagnostics.warnings.is_empty()); } + + #[test] + fn parses_standard_custom_provider_options_and_model_metadata() { + let mut diagnostics = ConfigDiagnostics::default(); + let config = parse_merged_config( + &json!({ + "provider": { + "custom": { + "name": "Custom Provider", + "npm": "@ai-sdk/openai-compatible", + "options": { + "baseURL": "https://example.com/v1", + "apiKey": "{env:CUSTOM_PROVIDER_KEY}" + }, + "models": { + "vision-model": { + "name": "Vision Model", + "attachment": true, + "reasoning": true, + "temperature": true, + "tool_call": true, + "limit": { + "context": 128000, + "output": 8192 + }, + "modalities": { + "input": ["text", "image"], + "output": ["text"] + } + } + } + }, + "anthropic": { + "options": { + "timeout": 300000 + } + } + } + }), + &mut diagnostics, + ); + + let provider = config + .custom_providers + .get("custom") + .expect("custom provider"); + assert_eq!(provider.name.as_deref(), Some("Custom Provider")); + assert_eq!(provider.npm.as_deref(), Some("@ai-sdk/openai-compatible")); + assert_eq!(provider.base_url.as_deref(), Some("https://example.com/v1")); + assert_eq!( + provider.api_key.as_deref(), + Some("{env:CUSTOM_PROVIDER_KEY}") + ); + + let model = provider.models.get("vision-model").expect("custom model"); + assert_eq!(model.name.as_deref(), Some("Vision Model")); + assert_eq!(model.context_window, Some(128000)); + assert_eq!(model.max_tokens, Some(8192)); + assert_eq!(model.attachment, Some(true)); + assert_eq!(model.reasoning, Some(true)); + assert_eq!(model.temperature, Some(true)); + assert_eq!(model.tool_call, Some(true)); + let modalities = model.modalities.as_ref().expect("model modalities"); + assert_eq!(modalities.input, ["text", "image"]); + assert_eq!(modalities.output, ["text"]); + + assert!(!config.custom_providers.contains_key("anthropic")); + assert_eq!( + config.provider_timeouts.get("anthropic"), + Some(&ProviderTimeout::Millis(300000)) + ); + assert!(diagnostics.warnings.is_empty()); + } + + #[test] + fn resolves_literal_custom_provider_api_key() { + let provider = CustomProviderConfig { + name: None, + npm: None, + base_url: None, + api_key: Some(" secret-key ".to_string()), + models: HashMap::new(), + }; + + assert_eq!(provider.resolved_api_key().as_deref(), Some("secret-key")); + } + + #[test] + fn resolves_custom_provider_api_key_from_environment_reference() { + assert_eq!( + resolve_api_key_value(Some("{env:CUSTOM_PROVIDER_KEY}"), |variable| { + assert_eq!(variable, "CUSTOM_PROVIDER_KEY"); + Some(" env-secret ".to_string()) + }) + .as_deref(), + Some("env-secret") + ); + assert_eq!( + resolve_api_key_value(Some("{env:MISSING_KEY}"), |_| None), + None + ); + } } diff --git a/src/config/mod.rs b/src/config/mod.rs index acef81e..ae3beaa 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -1,13 +1,13 @@ pub mod configuration; pub use configuration::{ - ConfigLoader, ImageOpenCommandConfig, ImageOpenWith, ImagesConfig, McpConfig, McpServerConfig, - NotificationEventConfig, NotificationsConfig, ProviderTimeout, TerminalNotificationCondition, - TerminalNotificationMode, + ConfigLoader, CustomProviderConfig, ImageOpenCommandConfig, ImageOpenWith, ImagesConfig, + McpConfig, McpServerConfig, NotificationEventConfig, NotificationsConfig, ProviderTimeout, + TerminalNotificationCondition, TerminalNotificationMode, }; #[cfg(test)] -pub use configuration::{McpLocalConfig, McpRemoteConfig}; +pub use configuration::McpLocalConfig; #[cfg(target_os = "macos")] pub use configuration::MacosNotificationBackend; diff --git a/src/llm/client.rs b/src/llm/client.rs index 53577f9..f550c76 100644 --- a/src/llm/client.rs +++ b/src/llm/client.rs @@ -631,6 +631,13 @@ pub async fn build_subagent_llm_session( }) } +fn resolve_api_key( + auth_config: Option<&crate::persistence::AuthConfig>, + custom_provider_api_key: Option, +) -> Option { + configured_api_key(auth_config).or(custom_provider_api_key) +} + pub async fn summarize_for_compaction( provider_name: String, model: String, @@ -764,12 +771,13 @@ async fn prepare_request_config( let auth_dao = crate::persistence::AuthDAO::new()?; let auth_config = auth_dao.get_provider(provider_name)?; + let discovery = crate::model::discovery::Discovery::new()?; + let custom_provider_api_key = discovery.custom_provider_api_key(provider_name); let provider = if let Some(provider) = crate::model::extensions::ModelExtensions::provider_for_request(provider_name) { provider } else { - let discovery = crate::model::discovery::Discovery::new()?; let providers = discovery.fetch_providers().await?; providers @@ -792,7 +800,7 @@ async fn prepare_request_config( provider.name.clone(), base_url, model_route.model_name.clone(), - configured_api_key(auth_config.as_ref()), + resolve_api_key(auth_config.as_ref(), custom_provider_api_key), reasoning_effort, supports_image_input, ); @@ -2090,11 +2098,31 @@ mod tests { use super::{ apply_provider_request_defaults, convert_messages, convert_messages_for_model, is_openai_oauth_model_allowed, maybe_apply_unauthenticated_free_provider_key, - model_supports_image_input, openai_request_instructions, resolve_model_route, - vlm_agent_has_model, AisdkMessage, OpenAIRequestOptions, ProviderKind, + model_supports_image_input, openai_request_instructions, resolve_api_key, + resolve_model_route, vlm_agent_has_model, AisdkMessage, OpenAIRequestOptions, ProviderKind, ProviderRequestConfig, }; + use crate::persistence::AuthConfig; + + #[test] + fn stored_auth_takes_precedence_over_custom_provider_api_key() { + assert_eq!( + resolve_api_key( + Some(&AuthConfig::Api { + key: "stored-key".to_string(), + }), + Some("config-key".to_string()), + ) + .as_deref(), + Some("stored-key") + ); + assert_eq!( + resolve_api_key(Some(&AuthConfig::Local), Some("config-key".to_string())).as_deref(), + Some("config-key") + ); + } + #[test] fn openai_oauth_instructions_preserve_stripped_system_prompt() { let options = OpenAIRequestOptions { diff --git a/src/mcp.rs b/src/mcp.rs index 0e5f7e6..4525c4d 100644 --- a/src/mcp.rs +++ b/src/mcp.rs @@ -11,6 +11,7 @@ use rmcp::ServiceExt; use serde_json::Value; use std::collections::{BTreeMap, HashMap}; use std::path::{Path, PathBuf}; +use std::process::Stdio; use std::sync::Arc; use std::time::Duration; use tokio::sync::Mutex; @@ -189,13 +190,15 @@ impl McpManager { .map(|cwd| resolve_path(&self.workspace, cwd)) .unwrap_or_else(|| self.workspace.clone()); let env = local.environment.clone(); - let transport = TokioChildProcess::new( + let (transport, _) = TokioChildProcess::builder( tokio::process::Command::new(command).configure(move |cmd| { cmd.args(args); cmd.current_dir(cwd); cmd.envs(env); }), - )?; + ) + .stderr(Stdio::null()) + .spawn()?; Ok(().serve(transport).await?) } McpServerConfig::Remote(remote) => { diff --git a/src/model/discovery.rs b/src/model/discovery.rs index 6dc7112..08098f0 100644 --- a/src/model/discovery.rs +++ b/src/model/discovery.rs @@ -30,7 +30,8 @@ pub struct Provider { static HTTP_CLIENT: OnceLock = OnceLock::new(); static MEMORY_CACHE: OnceLock>>> = OnceLock::new(); -static MEMORY_MODEL_CACHE: OnceLock>> = OnceLock::new(); +static MEMORY_MODEL_CACHE: OnceLock), CachedModels>>> = + OnceLock::new(); #[derive(Clone)] struct CachedModels { @@ -55,7 +56,7 @@ fn memory_cache() -> &'static Mutex>> { MEMORY_CACHE.get_or_init(|| Mutex::new(HashMap::new())) } -fn memory_model_cache() -> &'static Mutex> { +fn memory_model_cache() -> &'static Mutex), CachedModels>> { MEMORY_MODEL_CACHE.get_or_init(|| Mutex::new(HashMap::new())) } @@ -150,10 +151,163 @@ struct CacheEntry { pub struct Discovery { client: Client, cache_path: PathBuf, + custom_providers: + Option>, +} + +pub fn is_model_selectable( + model: &crate::model::types::Model, + connected_provider_ids: &std::collections::HashSet, + configured_provider_ids: &std::collections::HashSet, +) -> bool { + connected_provider_ids.contains(&model.provider_id) + || configured_provider_ids.contains(&model.provider_id) + || crate::model::extensions::ModelExtensions::is_available_without_connection(model) +} + +pub fn merge_dialog_models( + models: &mut Vec, + additional_models: impl IntoIterator, +) { + let mut known = models + .iter() + .map(|model| (model.provider_id.clone(), model.id.clone())) + .collect::>(); + for model in additional_models { + if known.insert((model.provider_id.clone(), model.id.clone())) { + models.push(model); + } + } } impl Discovery { + pub fn custom_provider_ids(&self) -> std::collections::HashSet { + self.custom_providers + .as_ref() + .map(|providers| providers.keys().cloned().collect()) + .unwrap_or_default() + } + + pub fn custom_provider_dialog_signature(&self) -> Vec { + let Some(custom_providers) = &self.custom_providers else { + return Vec::new(); + }; + + let mut signature = Vec::new(); + for (provider_id, provider) in custom_providers { + signature.push(format!( + "provider:{provider_id}:{}", + provider.name.as_deref().unwrap_or_default() + )); + for (model_id, model) in &provider.models { + let mut input_modalities = model + .modalities + .as_ref() + .map(|modalities| modalities.input.clone()) + .unwrap_or_default(); + input_modalities.sort(); + let mut output_modalities = model + .modalities + .as_ref() + .map(|modalities| modalities.output.clone()) + .unwrap_or_default(); + output_modalities.sort(); + signature.push(format!( + "model:{provider_id}:{model_id}:{}:{:?}:{}:{}", + model.name.as_deref().unwrap_or_default(), + model.attachment, + input_modalities.join(","), + output_modalities.join(",") + )); + } + } + signature.sort(); + signature + } + + pub fn custom_provider_matches_filter(&self, filter: &str) -> bool { + let filter = filter.trim().to_ascii_lowercase(); + self.custom_providers.as_ref().is_some_and(|providers| { + providers.iter().any(|(provider_id, provider)| { + provider_id.contains(&filter) + || provider + .name + .as_deref() + .is_some_and(|name| name.to_ascii_lowercase().contains(&filter)) + }) + }) + } + + pub fn apply_custom_models_to_dialog(&self, models: &mut Vec) { + let Some(custom_providers) = &self.custom_providers else { + return; + }; + + for (provider_id, provider) in custom_providers { + let provider_name = provider.name.as_deref().unwrap_or(provider_id); + for (model_id, custom_model) in &provider.models { + let model = models + .iter_mut() + .find(|model| model.provider_id == *provider_id && model.id == *model_id); + + if let Some(model) = model { + model.provider_name = provider_name.to_string(); + if let Some(name) = &custom_model.name { + model.name.clone_from(name); + } + if let Some(modalities) = &custom_model.modalities { + model.attachment = modalities.input.iter().any(|input| input == "image"); + } + if let Some(attachment) = custom_model.attachment { + model.attachment = attachment; + } + continue; + } + + let attachment = custom_model.attachment.unwrap_or_else(|| { + custom_model.modalities.as_ref().is_some_and(|modalities| { + modalities.input.iter().any(|input| input == "image") + }) + }); + models.push(crate::model::types::Model { + id: model_id.clone(), + name: custom_model + .name + .clone() + .unwrap_or_else(|| model_id.clone()), + family: String::new(), + provider_id: provider_id.clone(), + provider_name: provider_name.to_string(), + attachment, + structured_output: false, + free: false, + local: false, + reasoning_options: Vec::new(), + }); + } + } + } + + pub fn custom_provider_api_key(&self, provider_id: &str) -> Option { + self.custom_providers + .as_ref()? + .get(&provider_id.trim().to_ascii_lowercase())? + .resolved_api_key() + } + pub fn new() -> Result { + // Try to load custom providers from config file + let custom_providers = crate::config::ConfigLoader::load() + .map(|loaded| loaded.merged_config.custom_providers) + .ok(); + Self::new_with_custom(custom_providers) + } + + pub fn new_with_custom( + custom_providers: Option< + std::collections::HashMap, + >, + ) -> Result { if cfg!(test) || env::var("CRABCODE_TEST_MODE").is_ok() { let cache_dir = PathBuf::from("/tmp/crabcode_test_cache"); fs::create_dir_all(&cache_dir).context("Failed to create test cache directory")?; @@ -163,6 +317,7 @@ impl Discovery { Ok(Self { client: shared_http_client()?, cache_path, + custom_providers, }) } else { crate::persistence::ensure_cache_dir().context("Failed to create cache directory")?; @@ -173,6 +328,7 @@ impl Discovery { Ok(Self { client: shared_http_client()?, cache_path, + custom_providers, }) } } @@ -295,7 +451,7 @@ impl Discovery { cache.insert(cache_path.clone(), Arc::new(entry)); } if let Ok(mut cache) = memory_model_cache().lock() { - cache.remove(cache_path); + cache.retain(|(path, _), _| path != cache_path); } Ok(()) @@ -360,6 +516,7 @@ impl Discovery { }; crate::model::extensions::ModelExtensions::augment_runtime_catalog(&mut providers); + self.apply_custom_provider_overlays(&mut providers); Ok(providers) } @@ -369,15 +526,147 @@ impl Discovery { let mut providers = self.fetch_with_internal_providers(cached.as_ref()).await?; self.save_to_cache(&providers)?; crate::model::extensions::ModelExtensions::augment_runtime_catalog(&mut providers); + self.apply_custom_provider_overlays(&mut providers); Ok(providers) } + fn apply_custom_provider_overlays(&self, providers: &mut HashMap) { + let Some(custom_providers) = &self.custom_providers else { + return; + }; + + for (provider_id, custom_provider) in custom_providers { + let provider = providers + .entry(provider_id.clone()) + .or_insert_with(|| Provider { + id: provider_id.clone(), + name: custom_provider + .name + .clone() + .unwrap_or_else(|| provider_id.clone()), + api: String::new(), + doc: String::new(), + env: Vec::new(), + npm: String::new(), + models: HashMap::new(), + }); + + if let Some(name) = &custom_provider.name { + provider.name.clone_from(name); + } + if let Some(base_url) = &custom_provider.base_url { + provider.api.clone_from(base_url); + } + if let Some(npm) = &custom_provider.npm { + provider.npm.clone_from(npm); + } + + for (model_id, custom_model) in &custom_provider.models { + let model = provider + .models + .entry(model_id.clone()) + .or_insert_with(|| Model { + id: model_id.clone(), + name: custom_model + .name + .clone() + .unwrap_or_else(|| model_id.clone()), + family: String::new(), + attachment: false, + reasoning: false, + reasoning_options: Vec::new(), + tool_call: false, + structured_output: false, + temperature: false, + knowledge: String::new(), + release_date: String::new(), + last_updated: String::new(), + status: None, + modalities: Some(Modalities { + input: vec!["text".to_string()], + output: vec!["text".to_string()], + }), + open_weights: false, + cost: None, + limit: None, + provider: None, + }); + + if let Some(name) = &custom_model.name { + model.name.clone_from(name); + } + if let Some(reasoning) = custom_model.reasoning { + model.reasoning = reasoning; + } + if let Some(temperature) = custom_model.temperature { + model.temperature = temperature; + } + if let Some(tool_call) = custom_model.tool_call { + model.tool_call = tool_call; + } + if let Some(modalities) = &custom_model.modalities { + model.modalities = Some(Modalities { + input: modalities.input.clone(), + output: modalities.output.clone(), + }); + model.attachment = modalities.input.iter().any(|input| input == "image"); + } + if let Some(attachment) = custom_model.attachment { + model.attachment = attachment; + if attachment { + let modalities = model.modalities.get_or_insert_with(|| Modalities { + input: vec!["text".to_string()], + output: vec!["text".to_string()], + }); + if !modalities.input.iter().any(|input| input == "image") { + modalities.input.push("image".to_string()); + } + } else if let Some(modalities) = model.modalities.as_mut() { + modalities.input.retain(|input| input != "image"); + } + } + if custom_model.context_window.is_some() || custom_model.max_tokens.is_some() { + let current = model.limit.as_ref(); + if let Some(context) = custom_model + .context_window + .or_else(|| current.map(|limit| limit.context)) + { + model.limit = Some(Limit { + context, + output: custom_model + .max_tokens + .or_else(|| current.map(|limit| limit.output)) + .unwrap_or(context), + }); + } + } + + if custom_provider.npm.is_some() || custom_provider.base_url.is_some() { + let model_provider = model.provider.get_or_insert_with(|| ModelProvider { + npm: None, + api: None, + }); + if let Some(npm) = &custom_provider.npm { + model_provider.npm = Some(npm.clone()); + } + if let Some(base_url) = &custom_provider.base_url { + model_provider.api = Some(base_url.clone()); + } + } + } + } + } + pub async fn fetch_models(&self) -> Result> { let mut models = crate::model::extensions::ModelExtensions::runtime_models_from_cache(); + let cache_key = ( + self.get_cache_path().clone(), + self.custom_provider_dialog_signature(), + ); if let Some(cached) = memory_model_cache() .lock() .ok() - .and_then(|cache| cache.get(self.get_cache_path()).cloned()) + .and_then(|cache| cache.get(&cache_key).cloned()) .filter(|cached| cached.cached_at.elapsed().as_secs() <= CACHE_TTL_SECONDS) { models.extend(cached.models); @@ -432,7 +721,7 @@ impl Discovery { if let Ok(mut cache) = memory_model_cache().lock() { cache.insert( - self.get_cache_path().clone(), + cache_key, CachedModels { models: persistent_models.clone(), cached_at: std::time::Instant::now(), @@ -557,6 +846,9 @@ impl Default for Discovery { #[cfg(test)] mod tests { use super::*; + use crate::config::configuration::{ + CustomModelConfig, CustomModelModalities, CustomProviderConfig, + }; fn unique_test_cache_path(name: &str) -> PathBuf { let nanos = std::time::SystemTime::now() @@ -577,6 +869,279 @@ mod tests { assert!(discovery.is_ok()); } + #[test] + fn configured_provider_models_are_selectable_without_auth() { + let model = crate::model::types::Model { + id: "configured-model".to_string(), + name: "Configured Model".to_string(), + family: String::new(), + provider_id: "configured-provider".to_string(), + provider_name: "Configured Provider".to_string(), + attachment: false, + structured_output: false, + free: false, + local: false, + reasoning_options: Vec::new(), + }; + let connected_provider_ids = std::collections::HashSet::new(); + let configured_provider_ids = + std::collections::HashSet::from(["configured-provider".to_string()]); + + assert!(is_model_selectable( + &model, + &connected_provider_ids, + &configured_provider_ids, + )); + assert!(!is_model_selectable( + &model, + &connected_provider_ids, + &std::collections::HashSet::new(), + )); + + let mut models = vec![model.clone()]; + merge_dialog_models(&mut models, [model]); + assert_eq!(models.len(), 1); + } + + #[test] + fn custom_provider_ids_and_names_match_model_filters() { + let providers = HashMap::from([( + "mygateway".to_string(), + CustomProviderConfig { + name: Some("My Gateway".to_string()), + npm: None, + base_url: None, + api_key: None, + models: HashMap::from([( + "configured-model".to_string(), + CustomModelConfig { + name: Some("Configured Model".to_string()), + context_window: None, + max_tokens: None, + attachment: Some(true), + reasoning: None, + temperature: None, + tool_call: None, + modalities: None, + launch: false, + }, + )]), + }, + )]); + let discovery = Discovery::new_with_custom(Some(providers)).expect("discovery"); + + assert_eq!( + discovery.custom_provider_ids(), + std::collections::HashSet::from(["mygateway".to_string()]) + ); + assert!(discovery.custom_provider_matches_filter("mygateway")); + assert!(discovery.custom_provider_matches_filter("gateway")); + assert!(!discovery.custom_provider_matches_filter("openai")); + assert_eq!( + discovery.custom_provider_dialog_signature(), + vec![ + "model:mygateway:configured-model:Configured Model:Some(true)::".to_string(), + "provider:mygateway:My Gateway".to_string(), + ] + ); + + let mut models = Vec::new(); + discovery.apply_custom_models_to_dialog(&mut models); + assert_eq!(models.len(), 1); + assert_eq!(models[0].provider_id, "mygateway"); + assert_eq!(models[0].provider_name, "My Gateway"); + assert_eq!(models[0].id, "configured-model"); + assert_eq!(models[0].name, "Configured Model"); + assert!(models[0].attachment); + } + + #[test] + fn custom_model_overlay_preserves_unspecified_catalog_metadata() { + let mut providers = HashMap::from([( + "custom".to_string(), + Provider { + id: "custom".to_string(), + name: "Catalog Provider".to_string(), + api: "https://catalog.example/v1".to_string(), + doc: "https://catalog.example/docs".to_string(), + env: vec!["CATALOG_KEY".to_string()], + npm: "@ai-sdk/openai-compatible".to_string(), + models: HashMap::from([( + "vision-model".to_string(), + Model { + id: "vision-model".to_string(), + name: "Catalog Model".to_string(), + family: "catalog-family".to_string(), + attachment: true, + reasoning: true, + reasoning_options: Vec::new(), + tool_call: true, + structured_output: true, + temperature: true, + knowledge: "2025-01".to_string(), + release_date: "2025-01-01".to_string(), + last_updated: "2025-02-01".to_string(), + status: None, + modalities: Some(Modalities { + input: vec!["text".to_string(), "image".to_string()], + output: vec!["text".to_string()], + }), + open_weights: false, + cost: None, + limit: Some(Limit { + context: 64000, + output: 4096, + }), + provider: Some(ModelProvider { + npm: Some("@ai-sdk/openai-compatible".to_string()), + api: Some("https://catalog.example/v1".to_string()), + }), + }, + )]), + }, + )]); + let custom_providers = HashMap::from([( + "custom".to_string(), + CustomProviderConfig { + name: None, + npm: None, + base_url: None, + api_key: None, + models: HashMap::from([( + "vision-model".to_string(), + CustomModelConfig { + name: Some("Configured Model".to_string()), + context_window: Some(128000), + max_tokens: None, + attachment: None, + reasoning: None, + temperature: None, + tool_call: None, + modalities: None, + launch: false, + }, + )]), + }, + )]); + let discovery = Discovery::new_with_custom(Some(custom_providers)).expect("discovery"); + + discovery.apply_custom_provider_overlays(&mut providers); + + let provider = providers.get("custom").expect("provider"); + assert_eq!(provider.name, "Catalog Provider"); + assert_eq!(provider.api, "https://catalog.example/v1"); + assert_eq!(provider.npm, "@ai-sdk/openai-compatible"); + let model = provider.models.get("vision-model").expect("model"); + assert_eq!(model.name, "Configured Model"); + assert!(model.attachment); + assert!(model.reasoning); + assert!(model.tool_call); + assert!(model.structured_output); + assert!(model.temperature); + assert_eq!(model.family, "catalog-family"); + assert_eq!( + model + .modalities + .as_ref() + .map(|value| value.input.as_slice()), + Some(["text".to_string(), "image".to_string()].as_slice()) + ); + assert_eq!( + model.limit.as_ref().map(|limit| limit.context), + Some(128000) + ); + assert_eq!(model.limit.as_ref().map(|limit| limit.output), Some(4096)); + } + + #[test] + fn custom_model_modalities_enable_image_input() { + let mut providers = HashMap::new(); + let custom_providers = HashMap::from([( + "custom".to_string(), + CustomProviderConfig { + name: None, + npm: Some("@ai-sdk/openai-compatible".to_string()), + base_url: Some("https://example.com/v1".to_string()), + api_key: None, + models: HashMap::from([( + "vision-model".to_string(), + CustomModelConfig { + name: None, + context_window: Some(128000), + max_tokens: Some(8192), + attachment: None, + reasoning: Some(true), + temperature: Some(true), + tool_call: Some(true), + modalities: Some(CustomModelModalities { + input: vec!["text".to_string(), "image".to_string()], + output: vec!["text".to_string()], + }), + launch: false, + }, + )]), + }, + )]); + let discovery = Discovery::new_with_custom(Some(custom_providers)).expect("discovery"); + + discovery.apply_custom_provider_overlays(&mut providers); + + let model = providers["custom"] + .models + .get("vision-model") + .expect("model"); + assert!(model.attachment); + assert!(model.reasoning); + assert!(model.temperature); + assert!(model.tool_call); + assert_eq!( + model.limit.as_ref().map(|limit| limit.context), + Some(128000) + ); + assert_eq!(model.limit.as_ref().map(|limit| limit.output), Some(8192)); + } + + #[test] + fn custom_model_attachment_flag_updates_modalities() { + let mut providers = HashMap::new(); + let custom_providers = HashMap::from([( + "custom".to_string(), + CustomProviderConfig { + name: None, + npm: None, + base_url: None, + api_key: None, + models: HashMap::from([( + "vision-model".to_string(), + CustomModelConfig { + name: None, + context_window: None, + max_tokens: None, + attachment: Some(true), + reasoning: None, + temperature: None, + tool_call: None, + modalities: None, + launch: false, + }, + )]), + }, + )]); + let discovery = Discovery::new_with_custom(Some(custom_providers)).expect("discovery"); + + discovery.apply_custom_provider_overlays(&mut providers); + + let model = providers["custom"] + .models + .get("vision-model") + .expect("model"); + assert!(model.attachment); + assert!(model + .modalities + .as_ref() + .is_some_and(|modalities| modalities.input.iter().any(|input| input == "image"))); + } + #[tokio::test] async fn test_fetch_providers() { let discovery = Discovery::new().unwrap(); diff --git a/src/model/effective_catalog.rs b/src/model/effective_catalog.rs index 3ed2db1..890e662 100644 --- a/src/model/effective_catalog.rs +++ b/src/model/effective_catalog.rs @@ -13,7 +13,7 @@ use std::path::PathBuf; use std::time::{SystemTime, UNIX_EPOCH}; const SNAPSHOT_FILE: &str = "effective_catalog.json"; -const SNAPSHOT_SCHEMA_VERSION: u32 = 1; +const SNAPSHOT_SCHEMA_VERSION: u32 = 2; #[derive(Clone, Serialize, Deserialize)] struct SnapshotModel {