cc_switch_lib/services/stream_check/
service.rs

1use std::time::Instant;
2
3use crate::{app_config::AppType, error::AppError, provider::Provider};
4
5use super::types::{HealthStatus, StreamCheckConfig, StreamCheckResult};
6
7/// 流式健康检查服务
8pub struct StreamCheckService;
9
10impl StreamCheckService {
11    /// 执行流式健康检查（带重试）
12    pub async fn check_with_retry(
13        app_type: &AppType,
14        provider: &Provider,
15        config: &StreamCheckConfig,
16    ) -> Result<StreamCheckResult, AppError> {
17        let effective_config = Self::merge_provider_config(provider, config);
18        let mut last_result = None;
19
20        for attempt in 0..=effective_config.max_retries {
21            let result = Self::check_once(app_type, provider, &effective_config).await;
22
23            match &result {
24                Ok(r) if r.success => {
25                    return Ok(StreamCheckResult {
26                        retry_count: attempt,
27                        ..r.clone()
28                    });
29                }
30                Ok(r) => {
31                    if Self::should_retry(&r.message) && attempt < effective_config.max_retries {
32                        last_result = Some(r.clone());
33                        continue;
34                    }
35                    return Ok(StreamCheckResult {
36                        retry_count: attempt,
37                        ..r.clone()
38                    });
39                }
40                Err(err) => {
41                    if Self::should_retry(&err.to_string())
42                        && attempt < effective_config.max_retries
43                    {
44                        continue;
45                    }
46                    return Err(AppError::Message(err.to_string()));
47                }
48            }
49        }
50
51        Ok(last_result.unwrap_or_else(|| StreamCheckResult {
52            status: HealthStatus::Failed,
53            success: false,
54            message: "Check failed".to_string(),
55            response_time_ms: None,
56            http_status: None,
57            model_used: String::new(),
58            tested_at: chrono::Utc::now().timestamp(),
59            retry_count: effective_config.max_retries,
60        }))
61    }
62
63    pub(crate) fn merge_provider_config(
64        provider: &Provider,
65        global_config: &StreamCheckConfig,
66    ) -> StreamCheckConfig {
67        let test_config = provider
68            .meta
69            .as_ref()
70            .and_then(|meta| meta.test_config.as_ref())
71            .filter(|cfg| cfg.enabled);
72
73        match test_config {
74            Some(cfg) => StreamCheckConfig {
75                timeout_secs: cfg.timeout_secs.unwrap_or(global_config.timeout_secs),
76                max_retries: cfg.max_retries.unwrap_or(global_config.max_retries),
77                degraded_threshold_ms: cfg
78                    .degraded_threshold_ms
79                    .unwrap_or(global_config.degraded_threshold_ms),
80                claude_model: cfg
81                    .test_model
82                    .clone()
83                    .unwrap_or_else(|| global_config.claude_model.clone()),
84                codex_model: cfg
85                    .test_model
86                    .clone()
87                    .unwrap_or_else(|| global_config.codex_model.clone()),
88                gemini_model: cfg
89                    .test_model
90                    .clone()
91                    .unwrap_or_else(|| global_config.gemini_model.clone()),
92                test_prompt: cfg
93                    .test_prompt
94                    .clone()
95                    .unwrap_or_else(|| global_config.test_prompt.clone()),
96            },
97            None => global_config.clone(),
98        }
99    }
100
101    pub(crate) async fn check_once(
102        app_type: &AppType,
103        provider: &Provider,
104        config: &StreamCheckConfig,
105    ) -> Result<StreamCheckResult, AppError> {
106        if matches!(app_type, AppType::OpenClaw) {
107            return Err(AppError::Message("OpenClaw 暂不支持流式检查".to_string()));
108        }
109
110        let start = Instant::now();
111        let base_url = Self::extract_base_url(provider, app_type)?;
112        let auth = Self::extract_auth(provider, app_type, &base_url)?;
113        let client = Self::build_client_for_provider(provider)?;
114        let request_timeout = std::time::Duration::from_secs(config.timeout_secs);
115        let model_to_test = Self::resolve_test_model(app_type, provider, config);
116        let test_prompt = &config.test_prompt;
117
118        let result = match app_type {
119            AppType::Claude => {
120                Self::check_claude_stream(
121                    &client,
122                    &base_url,
123                    &auth,
124                    &model_to_test,
125                    test_prompt,
126                    request_timeout,
127                    provider,
128                )
129                .await
130            }
131            AppType::Codex => {
132                Self::check_codex_stream(
133                    &client,
134                    &base_url,
135                    &auth,
136                    &model_to_test,
137                    test_prompt,
138                    request_timeout,
139                )
140                .await
141            }
142            AppType::Gemini => {
143                Self::check_gemini_stream(
144                    &client,
145                    &base_url,
146                    &auth,
147                    &model_to_test,
148                    test_prompt,
149                    request_timeout,
150                )
151                .await
152            }
153            AppType::OpenCode => {
154                Self::check_codex_stream(
155                    &client,
156                    &base_url,
157                    &auth,
158                    &model_to_test,
159                    test_prompt,
160                    request_timeout,
161                )
162                .await
163            }
164            AppType::OpenClaw => unreachable!("OpenClaw should return unsupported earlier"),
165            AppType::Hermes => {
166                Self::check_codex_stream(
167                    &client,
168                    &base_url,
169                    &auth,
170                    &model_to_test,
171                    test_prompt,
172                    request_timeout,
173                )
174                .await
175            }
176        };
177
178        let response_time = start.elapsed().as_millis() as u64;
179        let tested_at = chrono::Utc::now().timestamp();
180
181        match result {
182            Ok((status_code, model)) => Ok(StreamCheckResult {
183                status: Self::determine_status(response_time, config.degraded_threshold_ms),
184                success: true,
185                message: "Check succeeded".to_string(),
186                response_time_ms: Some(response_time),
187                http_status: Some(status_code),
188                model_used: model,
189                tested_at,
190                retry_count: 0,
191            }),
192            Err(err) => Ok(StreamCheckResult {
193                status: HealthStatus::Failed,
194                success: false,
195                message: err.to_string(),
196                response_time_ms: Some(response_time),
197                http_status: None,
198                model_used: String::new(),
199                tested_at,
200                retry_count: 0,
201            }),
202        }
203    }
204
205    pub(crate) fn determine_status(latency_ms: u64, threshold: u64) -> HealthStatus {
206        if latency_ms <= threshold {
207            HealthStatus::Operational
208        } else {
209            HealthStatus::Degraded
210        }
211    }
212
213    pub(crate) fn should_retry(msg: &str) -> bool {
214        let lower = msg.to_lowercase();
215        lower.contains("timeout") || lower.contains("abort") || lower.contains("timed out")
216    }
217}
cc_switch_lib/services/stream_check/service.rs

cc_switch_lib/services/stream_check/
service.rs
