diff --git a/src-tauri/src/commands/stream_check.rs b/src-tauri/src/commands/stream_check.rs new file mode 100644 index 000000000..e70d5ea2d --- /dev/null +++ b/src-tauri/src/commands/stream_check.rs @@ -0,0 +1,94 @@ +//! 流式健康检查命令 + +use crate::app_config::AppType; +use crate::error::AppError; +use crate::services::stream_check::{ + HealthStatus, StreamCheckConfig, StreamCheckResult, StreamCheckService, +}; +use crate::store::AppState; +use tauri::State; + +/// 流式健康检查（单个供应商） +#[tauri::command] +pub async fn stream_check_provider( + state: State<'_, AppState>, + app_type: AppType, + provider_id: String, +) -> Result { + let config = state.db.get_stream_check_config()?; + + let providers = state.db.get_all_providers(app_type.as_str())?; + let provider = providers + .get(&provider_id) + .ok_or_else(|| AppError::Message(format!("供应商 {provider_id} 不存在")))?; + + let result = StreamCheckService::check_with_retry(&app_type, provider, &config).await?; + + // 记录日志 + let _ = state.db.save_stream_check_log( + &provider_id, + &provider.name, + app_type.as_str(), + &result, + ); + + Ok(result) +} + +/// 批量流式健康检查 +#[tauri::command] +pub async fn stream_check_all_providers( + state: State<'_, AppState>, + app_type: AppType, + proxy_targets_only: bool, +) -> Result, AppError> { + let config = state.db.get_stream_check_config()?; + let providers = state.db.get_all_providers(app_type.as_str())?; + + let mut results = Vec::new(); + + for (id, provider) in providers { + if proxy_targets_only && !provider.is_proxy_target.unwrap_or(false) { + continue; + } + + let result = StreamCheckService::check_with_retry(&app_type, &provider, &config) + .await + .unwrap_or_else(|e| StreamCheckResult { + status: HealthStatus::Failed, + success: false, + message: e.to_string(), + response_time_ms: None, + http_status: None, + model_used: String::new(), + tested_at: chrono::Utc::now().timestamp(), + retry_count: 0, + }); + + let _ = state.db.save_stream_check_log( + &id, + &provider.name, + app_type.as_str(), + &result, + ); + + results.push((id, result)); + } + + Ok(results) +} + +/// 获取流式检查配置 +#[tauri::command] +pub fn get_stream_check_config(state: State<'_, AppState>) -> Result { + state.db.get_stream_check_config() +} + +/// 保存流式检查配置 +#[tauri::command] +pub fn save_stream_check_config( + state: State<'_, AppState>, + config: StreamCheckConfig, +) -> Result<(), AppError> { + state.db.save_stream_check_config(&config) +} diff --git a/src-tauri/src/database/dao/stream_check.rs b/src-tauri/src/database/dao/stream_check.rs new file mode 100644 index 000000000..0f57e7cdb --- /dev/null +++ b/src-tauri/src/database/dao/stream_check.rs @@ -0,0 +1,57 @@ +//! 流式健康检查日志 DAO + +use crate::database::{lock_conn, Database}; +use crate::error::AppError; +use crate::services::stream_check::{StreamCheckConfig, StreamCheckResult}; + +impl Database { + /// 保存流式检查日志 + pub fn save_stream_check_log( + &self, + provider_id: &str, + provider_name: &str, + app_type: &str, + result: &StreamCheckResult, + ) -> Result { + let conn = lock_conn!(self.conn); + + conn.execute( + "INSERT INTO stream_check_logs + (provider_id, provider_name, app_type, status, success, message, + response_time_ms, http_status, model_used, retry_count, tested_at) + VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11)", + rusqlite::params![ + provider_id, + provider_name, + app_type, + format!("{:?}", result.status).to_lowercase(), + result.success, + result.message, + result.response_time_ms.map(|t| t as i64), + result.http_status.map(|s| s as i64), + result.model_used, + result.retry_count as i64, + result.tested_at, + ], + ) + .map_err(|e| AppError::Database(e.to_string()))?; + + Ok(conn.last_insert_rowid()) + } + + /// 获取流式检查配置 + pub fn get_stream_check_config(&self) -> Result { + match self.get_setting("stream_check_config")? { + Some(json) => serde_json::from_str(&json) + .map_err(|e| AppError::Message(format!("解析配置失败: {e}"))), + None => Ok(StreamCheckConfig::default()), + } + } + + /// 保存流式检查配置 + pub fn save_stream_check_config(&self, config: &StreamCheckConfig) -> Result<(), AppError> { + let json = serde_json::to_string(config) + .map_err(|e| AppError::Message(format!("序列化配置失败: {e}")))?; + self.set_setting("stream_check_config", &json) + } +} diff --git a/src-tauri/src/services/stream_check.rs b/src-tauri/src/services/stream_check.rs new file mode 100644 index 000000000..eaf2517b8 --- /dev/null +++ b/src-tauri/src/services/stream_check.rs @@ -0,0 +1,384 @@ +//! 流式健康检查服务 +//! +//! 使用流式 API 进行快速健康检查，只需接收首个 chunk 即判定成功。 + +use futures::StreamExt; +use reqwest::Client; +use serde::{Deserialize, Serialize}; +use serde_json::json; +use std::time::{Duration, Instant}; + +use crate::app_config::AppType; +use crate::error::AppError; +use crate::provider::Provider; +use crate::proxy::providers::{get_adapter, AuthInfo}; + +/// 健康状态枚举 +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +#[serde(rename_all = "lowercase")] +pub enum HealthStatus { + Operational, + Degraded, + Failed, +} + +/// 流式检查配置 +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct StreamCheckConfig { + pub timeout_secs: u64, + pub max_retries: u32, + pub degraded_threshold_ms: u64, +} + +impl Default for StreamCheckConfig { + fn default() -> Self { + Self { + timeout_secs: 45, + max_retries: 2, + degraded_threshold_ms: 6000, + } + } +} + +/// 流式检查结果 +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct StreamCheckResult { + pub status: HealthStatus, + pub success: bool, + pub message: String, + pub response_time_ms: Option, + pub http_status: Option, + pub model_used: String, + pub tested_at: i64, + pub retry_count: u32, +} + +/// 流式健康检查服务 +pub struct StreamCheckService; + +impl StreamCheckService { + /// 执行流式健康检查（带重试） + pub async fn check_with_retry( + app_type: &AppType, + provider: &Provider, + config: &StreamCheckConfig, + ) -> Result { + let mut last_result = None; + + for attempt in 0..=config.max_retries { + let result = Self::check_once(app_type, provider, config).await; + + match &result { + Ok(r) if r.success => { + return Ok(StreamCheckResult { + retry_count: attempt, + ..r.clone() + }); + } + Ok(r) => { + // 失败但非异常，判断是否重试 + if Self::should_retry(&r.message) && attempt < config.max_retries { + last_result = Some(r.clone()); + continue; + } + return Ok(StreamCheckResult { + retry_count: attempt, + ..r.clone() + }); + } + Err(e) => { + if Self::should_retry(&e.to_string()) && attempt < config.max_retries { + continue; + } + return Err(AppError::Message(e.to_string())); + } + } + } + + Ok(last_result.unwrap_or_else(|| StreamCheckResult { + status: HealthStatus::Failed, + success: false, + message: "检查失败".to_string(), + response_time_ms: None, + http_status: None, + model_used: String::new(), + tested_at: chrono::Utc::now().timestamp(), + retry_count: config.max_retries, + })) + } + + /// 单次流式检查 + async fn check_once( + app_type: &AppType, + provider: &Provider, + config: &StreamCheckConfig, + ) -> Result { + let start = Instant::now(); + let adapter = get_adapter(app_type); + + let base_url = adapter + .extract_base_url(provider) + .map_err(|e| AppError::Message(format!("提取 base_url 失败: {e}")))?; + + let auth = adapter + .extract_auth(provider) + .ok_or_else(|| AppError::Message("未找到 API Key".to_string()))?; + + let client = Client::builder() + .timeout(Duration::from_secs(config.timeout_secs)) + .user_agent("cc-switch/1.0") + .build() + .map_err(|e| AppError::Message(format!("创建客户端失败: {e}")))?; + + let result = match app_type { + AppType::Claude => Self::check_claude_stream(&client, &base_url, &auth).await, + AppType::Codex => Self::check_codex_stream(&client, &base_url, &auth).await, + AppType::Gemini => Self::check_gemini_stream(&client, &base_url, &auth).await, + }; + + let response_time = start.elapsed().as_millis() as u64; + let tested_at = chrono::Utc::now().timestamp(); + + match result { + Ok((status_code, model)) => { + let health_status = + Self::determine_status(response_time, config.degraded_threshold_ms); + Ok(StreamCheckResult { + status: health_status, + success: true, + message: "检查成功".to_string(), + response_time_ms: Some(response_time), + http_status: Some(status_code), + model_used: model, + tested_at, + retry_count: 0, + }) + } + Err(e) => Ok(StreamCheckResult { + status: HealthStatus::Failed, + success: false, + message: e.to_string(), + response_time_ms: Some(response_time), + http_status: None, + model_used: String::new(), + tested_at, + retry_count: 0, + }), + } + } + + /// Claude 流式检查 + async fn check_claude_stream( + client: &Client, + base_url: &str, + auth: &AuthInfo, + ) -> Result<(u16, String), AppError> { + let base = base_url.trim_end_matches('/'); + let url = if base.ends_with("/v1") { + format!("{base}/messages") + } else { + format!("{base}/v1/messages") + }; + + let model = "claude-3-5-haiku-latest"; + + let body = json!({ + "model": model, + "max_tokens": 1, + "messages": [{ "role": "user", "content": "hi" }], + "stream": true + }); + + let response = client + .post(&url) + .header("x-api-key", &auth.api_key) + .header("anthropic-version", "2023-06-01") + .header("Content-Type", "application/json") + .json(&body) + .send() + .await + .map_err(Self::map_request_error)?; + + let status = response.status().as_u16(); + + if !response.status().is_success() { + let error_text = response.text().await.unwrap_or_default(); + return Err(AppError::Message(format!("HTTP {status}: {error_text}"))); + } + + // 流式读取：只需首个 chunk + let mut stream = response.bytes_stream(); + if let Some(chunk) = stream.next().await { + match chunk { + Ok(_) => Ok((status, model.to_string())), + Err(e) => Err(AppError::Message(format!("读取流失败: {e}"))), + } + } else { + Err(AppError::Message("未收到响应数据".to_string())) + } + } + + /// Codex 流式检查 + async fn check_codex_stream( + client: &Client, + base_url: &str, + auth: &AuthInfo, + ) -> Result<(u16, String), AppError> { + let base = base_url.trim_end_matches('/'); + let url = if base.ends_with("/v1") { + format!("{base}/chat/completions") + } else { + format!("{base}/v1/chat/completions") + }; + + let model = "gpt-4o-mini"; + + let body = json!({ + "model": model, + "messages": [ + { "role": "system", "content": "" }, + { "role": "assistant", "content": "" }, + { "role": "user", "content": "hi" } + ], + "max_tokens": 1, + "temperature": 0, + "stream": true + }); + + let response = client + .post(&url) + .header("Authorization", format!("Bearer {}", auth.api_key)) + .header("Content-Type", "application/json") + .json(&body) + .send() + .await + .map_err(Self::map_request_error)?; + + let status = response.status().as_u16(); + + if !response.status().is_success() { + let error_text = response.text().await.unwrap_or_default(); + return Err(AppError::Message(format!("HTTP {status}: {error_text}"))); + } + + let mut stream = response.bytes_stream(); + if let Some(chunk) = stream.next().await { + match chunk { + Ok(_) => Ok((status, model.to_string())), + Err(e) => Err(AppError::Message(format!("读取流失败: {e}"))), + } + } else { + Err(AppError::Message("未收到响应数据".to_string())) + } + } + + /// Gemini 流式检查 + async fn check_gemini_stream( + client: &Client, + base_url: &str, + auth: &AuthInfo, + ) -> Result<(u16, String), AppError> { + let base = base_url.trim_end_matches('/'); + let url = format!("{base}/v1/chat/completions"); + + let model = "gemini-1.5-flash"; + + let body = json!({ + "model": model, + "messages": [{ "role": "user", "content": "hi" }], + "max_tokens": 1, + "temperature": 0, + "stream": true + }); + + let response = client + .post(&url) + .header("Authorization", format!("Bearer {}", auth.api_key)) + .header("Content-Type", "application/json") + .json(&body) + .send() + .await + .map_err(Self::map_request_error)?; + + let status = response.status().as_u16(); + + if !response.status().is_success() { + let error_text = response.text().await.unwrap_or_default(); + return Err(AppError::Message(format!("HTTP {status}: {error_text}"))); + } + + let mut stream = response.bytes_stream(); + if let Some(chunk) = stream.next().await { + match chunk { + Ok(_) => Ok((status, model.to_string())), + Err(e) => Err(AppError::Message(format!("读取流失败: {e}"))), + } + } else { + Err(AppError::Message("未收到响应数据".to_string())) + } + } + + fn determine_status(latency_ms: u64, threshold: u64) -> HealthStatus { + if latency_ms <= threshold { + HealthStatus::Operational + } else { + HealthStatus::Degraded + } + } + + fn should_retry(msg: &str) -> bool { + let lower = msg.to_lowercase(); + lower.contains("timeout") + || lower.contains("abort") + || lower.contains("中断") + || lower.contains("超时") + } + + fn map_request_error(e: reqwest::Error) -> AppError { + if e.is_timeout() { + AppError::Message("请求超时".to_string()) + } else if e.is_connect() { + AppError::Message(format!("连接失败: {e}")) + } else { + AppError::Message(e.to_string()) + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_determine_status() { + assert_eq!( + StreamCheckService::determine_status(3000, 6000), + HealthStatus::Operational + ); + assert_eq!( + StreamCheckService::determine_status(6000, 6000), + HealthStatus::Operational + ); + assert_eq!( + StreamCheckService::determine_status(6001, 6000), + HealthStatus::Degraded + ); + } + + #[test] + fn test_should_retry() { + assert!(StreamCheckService::should_retry("请求超时")); + assert!(StreamCheckService::should_retry("request timeout")); + assert!(!StreamCheckService::should_retry("API Key 无效")); + } + + #[test] + fn test_default_config() { + let config = StreamCheckConfig::default(); + assert_eq!(config.timeout_secs, 45); + assert_eq!(config.max_retries, 2); + assert_eq!(config.degraded_threshold_ms, 6000); + } +} diff --git a/src/components/providers/HealthStatusIndicator.tsx b/src/components/providers/HealthStatusIndicator.tsx new file mode 100644 index 000000000..8b87601fc --- /dev/null +++ b/src/components/providers/HealthStatusIndicator.tsx @@ -0,0 +1,45 @@ +import React from "react"; +import { cn } from "@/lib/utils"; +import type { HealthStatus } from "@/lib/api/model-test"; + +interface HealthStatusIndicatorProps { + status: HealthStatus; + responseTimeMs?: number; + className?: string; +} + +const statusConfig = { + operational: { + color: "bg-emerald-500", + label: "正常", + textColor: "text-emerald-600 dark:text-emerald-400", + }, + degraded: { + color: "bg-yellow-500", + label: "降级", + textColor: "text-yellow-600 dark:text-yellow-400", + }, + failed: { + color: "bg-red-500", + label: "失败", + textColor: "text-red-600 dark:text-red-400", + }, +}; + +export const HealthStatusIndicator: React.FC = ({ + status, + responseTimeMs, + className, +}) => { + const config = statusConfig[status]; + + return ( +
+
+ + {config.label} + {responseTimeMs !== undefined && ` (${responseTimeMs}ms)`} + +
+ ); +}; diff --git a/src/hooks/useStreamCheck.ts b/src/hooks/useStreamCheck.ts new file mode 100644 index 000000000..0685b6031 --- /dev/null +++ b/src/hooks/useStreamCheck.ts @@ -0,0 +1,77 @@ +import { useState, useCallback } from "react"; +import { toast } from "sonner"; +import { useTranslation } from "react-i18next"; +import { + streamCheckProvider, + type StreamCheckResult, +} from "@/lib/api/model-test"; +import type { AppId } from "@/lib/api"; + +export function useStreamCheck(appId: AppId) { + const { t } = useTranslation(); + const [checkingIds, setCheckingIds] = useState>(new Set()); + + const checkProvider = useCallback( + async ( + providerId: string, + providerName: string, + ): Promise => { + setCheckingIds((prev) => new Set(prev).add(providerId)); + + try { + const result = await streamCheckProvider(appId, providerId); + + if (result.status === "operational") { + toast.success( + t("streamCheck.operational", { + name: providerName, + time: result.responseTimeMs, + defaultValue: `${providerName} 运行正常 (${result.responseTimeMs}ms)`, + }), + ); + } else if (result.status === "degraded") { + toast.warning( + t("streamCheck.degraded", { + name: providerName, + time: result.responseTimeMs, + defaultValue: `${providerName} 响应较慢 (${result.responseTimeMs}ms)`, + }), + ); + } else { + toast.error( + t("streamCheck.failed", { + name: providerName, + error: result.message, + defaultValue: `${providerName} 检查失败: ${result.message}`, + }), + ); + } + + return result; + } catch (e) { + toast.error( + t("streamCheck.error", { + name: providerName, + error: String(e), + defaultValue: `${providerName} 检查出错: ${String(e)}`, + }), + ); + return null; + } finally { + setCheckingIds((prev) => { + const next = new Set(prev); + next.delete(providerId); + return next; + }); + } + }, + [appId, t], + ); + + const isChecking = useCallback( + (providerId: string) => checkingIds.has(providerId), + [checkingIds], + ); + + return { checkProvider, isChecking }; +}