diff --git a/src-tauri/src/services/provider/usage.rs b/src-tauri/src/services/provider/usage.rs index 5d567dc7f3..4b50467b85 100644 --- a/src-tauri/src/services/provider/usage.rs +++ b/src-tauri/src/services/provider/usage.rs @@ -85,25 +85,118 @@ pub(crate) async fn execute_and_format_usage_result( fn extract_api_key_from_provider(provider: &crate::provider::Provider) -> Option { if let Some(env) = provider.settings_config.get("env") { // Try multiple possible API key fields - env.get("ANTHROPIC_AUTH_TOKEN") + if let Some(api_key) = env + .get("ANTHROPIC_AUTH_TOKEN") .or_else(|| env.get("ANTHROPIC_API_KEY")) .or_else(|| env.get("OPENROUTER_API_KEY")) .or_else(|| env.get("GOOGLE_API_KEY")) + .or_else(|| env.get("GEMINI_API_KEY")) .and_then(|v| v.as_str()) - .map(|s| s.to_string()) - } else { - None + .filter(|s| !s.trim().is_empty()) + { + return Some(api_key.to_string()); + } } + + provider + .settings_config + .get("auth") + .and_then(|auth| auth.get("OPENAI_API_KEY")) + .or_else(|| { + provider + .settings_config + .get("options") + .and_then(|options| options.get("apiKey")) + }) + .or_else(|| provider.settings_config.get("apiKey")) + .or_else(|| provider.settings_config.get("api_key")) + .and_then(|v| v.as_str()) + .filter(|s| !s.trim().is_empty()) + .map(|s| s.to_string()) } /// Extract base URL from provider configuration fn extract_base_url_from_provider(provider: &crate::provider::Provider) -> Option { if let Some(env) = provider.settings_config.get("env") { // Try multiple possible base URL fields - env.get("ANTHROPIC_BASE_URL") + if let Some(base_url) = env + .get("ANTHROPIC_BASE_URL") .or_else(|| env.get("GOOGLE_GEMINI_BASE_URL")) + .or_else(|| env.get("OPENAI_BASE_URL")) .and_then(|v| v.as_str()) - .map(|s| s.trim_end_matches('/').to_string()) + .filter(|s| !s.trim().is_empty()) + { + return Some(base_url.trim_end_matches('/').to_string()); + } + } + + provider + .settings_config + .get("config") + .and_then(|v| v.as_str()) + .and_then(extract_codex_base_url_from_toml) + .or_else(|| { + provider + .settings_config + .get("options") + .and_then(|options| options.get("baseURL").or_else(|| options.get("baseUrl"))) + .and_then(|v| v.as_str()) + .filter(|s| !s.trim().is_empty()) + .map(|s| s.trim_end_matches('/').to_string()) + }) + .or_else(|| { + provider + .settings_config + .get("baseUrl") + .or_else(|| provider.settings_config.get("base_url")) + .and_then(|v| v.as_str()) + .filter(|s| !s.trim().is_empty()) + .map(|s| s.trim_end_matches('/').to_string()) + }) +} + +fn extract_codex_base_url_from_toml(config_text: &str) -> Option { + let parsed = config_text.parse::().ok()?; + + if let Some(model_provider) = parsed.get("model_provider").and_then(|v| v.as_str()) { + if let Some(base_url) = parsed + .get("model_providers") + .and_then(|providers| providers.get(model_provider)) + .and_then(|provider| provider.get("base_url")) + .and_then(|v| v.as_str()) + .filter(|s| !s.trim().is_empty()) + { + return Some(base_url.trim_end_matches('/').to_string()); + } + } + + if let Some(base_url) = parsed + .get("base_url") + .and_then(|v| v.as_str()) + .filter(|s| !s.trim().is_empty()) + { + return Some(base_url.trim_end_matches('/').to_string()); + } + + let base_urls: Vec = parsed + .get("model_providers") + .and_then(|providers| providers.as_table()) + .map(|providers| { + providers + .values() + .filter_map(|provider| { + provider + .get("base_url") + .and_then(|v| v.as_str()) + .filter(|s| !s.trim().is_empty()) + .map(|s| s.trim_end_matches('/').to_string()) + }) + .collect() + }) + .unwrap_or_default(); + + if base_urls.len() == 1 { + base_urls.into_iter().next() } else { None } @@ -185,9 +278,9 @@ pub async fn query_usage( /// Test usage script (using temporary script content, not saved) #[allow(clippy::too_many_arguments)] pub async fn test_usage_script( - _state: &AppState, - _app_type: AppType, - _provider_id: &str, + state: &AppState, + app_type: AppType, + provider_id: &str, script_code: &str, timeout: u64, api_key: Option<&str>, @@ -196,11 +289,35 @@ pub async fn test_usage_script( user_id: Option<&str>, template_type: Option<&str>, ) -> Result { - // Use provided credential parameters directly for testing + let provider_credentials = || -> Option<(Option, Option)> { + let providers = state.db.get_all_providers(app_type.as_str()).ok()?; + let provider = providers.get(provider_id)?; + Some(( + extract_api_key_from_provider(provider), + extract_base_url_from_provider(provider), + )) + }; + + let (provider_api_key, provider_base_url) = + if api_key.is_none_or(|s| s.is_empty()) || base_url.is_none_or(|s| s.is_empty()) { + provider_credentials().unwrap_or((None, None)) + } else { + (None, None) + }; + + let resolved_api_key = api_key + .filter(|s| !s.is_empty()) + .or(provider_api_key.as_deref()) + .unwrap_or(""); + let resolved_base_url = base_url + .filter(|s| !s.is_empty()) + .or(provider_base_url.as_deref()) + .unwrap_or(""); + execute_and_format_usage_result( script_code, - api_key.unwrap_or(""), - base_url.unwrap_or(""), + resolved_api_key, + resolved_base_url, timeout, access_token, user_id, @@ -226,3 +343,58 @@ pub(crate) fn validate_usage_script(script: &UsageScript) -> Result<(), AppError Ok(()) } + +#[cfg(test)] +mod tests { + use super::*; + use crate::provider::Provider; + use serde_json::json; + + fn provider_with_settings(settings_config: serde_json::Value) -> Provider { + Provider::with_id( + "test-provider".to_string(), + "Test Provider".to_string(), + settings_config, + None, + ) + } + + #[test] + fn extracts_codex_credentials_from_auth_and_active_model_provider() { + let provider = provider_with_settings(json!({ + "auth": { + "OPENAI_API_KEY": "sk-test" + }, + "config": r#" +model_provider = "cubence" + +[model_providers.cubence] +base_url = "https://api.cubence.me/v1" + +[model_providers.other] +base_url = "https://api.other.example/v1" +"# + })); + + assert_eq!( + extract_api_key_from_provider(&provider), + Some("sk-test".to_string()) + ); + assert_eq!( + extract_base_url_from_provider(&provider), + Some("https://api.cubence.me/v1".to_string()) + ); + } + + #[test] + fn extracts_codex_base_url_from_single_provider_without_active_provider() { + let base_url = extract_codex_base_url_from_toml( + r#" +[model_providers.cubence] +base_url = "https://api.cubence.me/v1/" +"#, + ); + + assert_eq!(base_url, Some("https://api.cubence.me/v1".to_string())); + } +} diff --git a/src-tauri/src/usage_script.rs b/src-tauri/src/usage_script.rs index 7d121e4faf..3ed0b06c8e 100644 --- a/src-tauri/src/usage_script.rs +++ b/src-tauri/src/usage_script.rs @@ -405,20 +405,36 @@ fn build_script_with_vars( access_token: Option<&str>, user_id: Option<&str>, ) -> String { - let mut replaced = script_code - .replace("{{apiKey}}", api_key) - .replace("{{baseUrl}}", base_url); + let mut replaced = replace_template_aliases( + script_code, + &[("{{apiKey}}", api_key), ("{{apikey}}", api_key)], + ); + replaced = replace_template_aliases( + &replaced, + &[("{{baseUrl}}", base_url), ("{{baseurl}}", base_url)], + ); if let Some(token) = access_token { - replaced = replaced.replace("{{accessToken}}", token); + replaced = replace_template_aliases( + &replaced, + &[("{{accessToken}}", token), ("{{accesstoken}}", token)], + ); } if let Some(uid) = user_id { - replaced = replaced.replace("{{userId}}", uid); + replaced = replace_template_aliases(&replaced, &[("{{userId}}", uid), ("{{userid}}", uid)]); } replaced } +fn replace_template_aliases(input: &str, aliases: &[(&str, &str)]) -> String { + aliases + .iter() + .fold(input.to_string(), |acc, (alias, value)| { + acc.replace(alias, value) + }) +} + /// 验证 base_url 的基本安全性 fn validate_base_url(base_url: &str) -> Result<(), AppError> { if base_url.is_empty() { @@ -641,4 +657,39 @@ mod tests { } } } + + #[test] + fn test_build_script_with_vars_supports_legacy_lowercase_aliases() { + let script = r#" + ({ + request: { + url: "{{baseurl}}/usage?key={{apikey}}", + method: "GET", + headers: { + Authorization: "Bearer {{apikey}}", + "X-Access-Token": "{{accesstoken}}", + "X-User-Id": "{{userid}}" + } + }, + extractor: () => ({ key: "{{apikey}}" }) + }) + "#; + + let replaced = build_script_with_vars( + script, + "test-api-key", + "https://api.example.com/v1", + Some("token-123"), + Some("user-456"), + ); + + assert!(replaced.contains("https://api.example.com/v1/usage?key=test-api-key")); + assert!(replaced.contains("Bearer test-api-key")); + assert!(replaced.contains("\"X-Access-Token\": \"token-123\"")); + assert!(replaced.contains("\"X-User-Id\": \"user-456\"")); + assert!(!replaced.contains("{{apikey}}")); + assert!(!replaced.contains("{{baseurl}}")); + assert!(!replaced.contains("{{accesstoken}}")); + assert!(!replaced.contains("{{userid}}")); + } }