diff --git a/src-tauri/src/claude_desktop_config.rs b/src-tauri/src/claude_desktop_config.rs index 2d777dbe3..70af4cfce 100644 --- a/src-tauri/src/claude_desktop_config.rs +++ b/src-tauri/src/claude_desktop_config.rs @@ -575,7 +575,7 @@ pub fn proxy_model_routes(provider: &Provider) -> Result .collect::>(); let mut result = Vec::new(); let mut entries = routes.iter().collect::>(); - entries.sort_by(|(left, _), (right, _)| left.cmp(right)); + entries.sort_by_key(|(left, _)| *left); for (route_id, route) in entries { let supports_1m = route.supports_1m.unwrap_or(false); let route_id = route_id.trim(); diff --git a/src-tauri/src/proxy/forwarder.rs b/src-tauri/src/proxy/forwarder.rs index f19306f06..f0ba93066 100644 --- a/src-tauri/src/proxy/forwarder.rs +++ b/src-tauri/src/proxy/forwarder.rs @@ -35,6 +35,9 @@ pub struct ForwardResult { pub response: ProxyResponse, pub provider: Provider, pub claude_api_format: Option, + /// 活跃连接 RAII guard：随响应一起流转到 response_processor / handle_claude_transform， + /// 最终被 move 进流式 body future（或非流式响应作用域），覆盖整个响应生命周期。 + pub(crate) connection_guard: Option, } pub struct ForwardError { @@ -42,6 +45,44 @@ pub struct ForwardError { pub provider: Option, } +/// 活跃连接 RAII guard +/// +/// 构造时把 `ProxyStatus.active_connections` +1；Drop 时在 tokio runtime 上调度 +/// 一个异步任务执行 -1，从而支持把 guard move 进流式 body future（stream 自然结束 +/// 时 guard 与 future 一起 drop）。 +/// +/// 设计动机：之前在 `forward_with_retry` 出口处同步 -1，但流式响应的 body 实际 +/// 在 `create_logged_passthrough_stream` 内还会继续 yield 字节流，导致 UI 的 +/// `active_connections` 计数过早归零。RAII guard 让"减量"由 Rust 类型系统驱动， +/// 不需要每条出口路径都手动调用。 +pub(crate) struct ActiveConnectionGuard { + status: Arc>, +} + +impl ActiveConnectionGuard { + pub(crate) async fn acquire(status: Arc>) -> Self { + { + let mut s = status.write().await; + s.active_connections = s.active_connections.saturating_add(1); + } + Self { status } + } +} + +impl Drop for ActiveConnectionGuard { + fn drop(&mut self) { + // Drop 不能 await：把减量操作调度到 tokio runtime + let status = self.status.clone(); + if let Ok(handle) = tokio::runtime::Handle::try_current() { + handle.spawn(async move { + let mut s = status.write().await; + s.active_connections = s.active_connections.saturating_sub(1); + }); + } + // 没有 runtime 时静默丢失计数（仅 UI 展示用，可接受最终一致性） + } +} + pub struct RequestForwarder { /// 共享的 ProviderRouter（持有熔断器状态） router: Arc, @@ -235,10 +276,10 @@ impl RequestForwarder { extensions: Extensions, providers: Vec, ) -> Result { + let guard = ActiveConnectionGuard::acquire(self.status.clone()).await; { let mut s = self.status.write().await; s.total_requests = s.total_requests.saturating_add(1); - s.active_connections = s.active_connections.saturating_add(1); s.last_request_at = Some(chrono::Utc::now().to_rfc3339()); } let result = self @@ -246,11 +287,13 @@ impl RequestForwarder { app_type, method, endpoint, body, headers, extensions, providers, ) .await; - { - let mut s = self.status.write().await; - s.active_connections = s.active_connections.saturating_sub(1); - } - result + // 把 guard 注入到 Ok 结果，让它随响应一起流转到 response_processor， + // 在流式 body 的 future 内才真正 drop。 + // Err 路径：guard 在函数 scope 内随返回值落地时自动 drop。 + result.map(|mut fr| { + fr.connection_guard = Some(guard); + fr + }) } /// 实际转发逻辑（不包含客户端维度的入口/出口计数） @@ -288,15 +331,16 @@ impl RequestForwarder { let mut last_provider = None; let mut attempted_providers = 0usize; - // 整流器重试标记：确保整流最多触发一次 - let mut rectifier_retried = false; - let mut budget_rectifier_retried = false; - // 单 Provider 场景下跳过熔断器检查（故障转移关闭时） let bypass_circuit_breaker = providers.len() == 1; // 依次尝试每个供应商 for provider in providers.iter() { + // 整流器重试标记：每个 provider 独立持有，避免标记跨 provider 短路故障转移 + // —— 首家 provider 整流后被 5xx/timeout 击落时，下家仍能用整流后的请求体走整流流程 + let mut rectifier_retried = false; + let mut budget_rectifier_retried = false; + // 上限检查：尊重用户在 AppProxyConfig.max_retries 上配置的「重试次数」。 // 放在熔断器 allow 检查之前，避免在已经超限时还占用 HalfOpen 探测名额。 if attempted_providers >= self.max_attempts { @@ -415,6 +459,7 @@ impl RequestForwarder { response, provider: provider.clone(), claude_api_format, + connection_guard: None, }); } Err(e) => { @@ -547,6 +592,7 @@ impl RequestForwarder { response, provider: provider.clone(), claude_api_format, + connection_guard: None, }); } Err(retry_err) => { @@ -705,6 +751,7 @@ impl RequestForwarder { response, provider: provider.clone(), claude_api_format, + connection_guard: None, }); } Err(retry_err) => { @@ -753,24 +800,25 @@ impl RequestForwarder { }); } - // 失败：记录失败并更新熔断器 - let _ = self - .router - .record_result( - &provider.id, - app_type_str, - used_half_open_permit, - false, - Some(e.to_string()), - ) - .await; - - // 分类错误 + // 先分类错误，决定是否计入 provider 健康度 + // —— NonRetryable / ClientAbort 是客户端层错误，无论换哪家 provider 都会被拒绝， + // 不应污染熔断器和数据库健康度（与 release_permit_neutral 同语义）。 let category = self.categorize_proxy_error(&e); match category { ErrorCategory::Retryable => { - // 可重试：更新错误信息，继续尝试下一个供应商 + // 可重试：真正的 provider 故障 → 记录失败并更新熔断器/DB 健康度 + let _ = self + .router + .record_result( + &provider.id, + app_type_str, + used_half_open_permit, + false, + Some(e.to_string()), + ) + .await; + { let mut status = self.status.write().await; status.last_error = @@ -791,7 +839,14 @@ impl RequestForwarder { continue; } ErrorCategory::NonRetryable | ErrorCategory::ClientAbort => { - // 不可重试：直接返回错误 + // 不可重试：客户端层错误或客户端断连 → 不污染健康度，仅释放 HalfOpen permit + self.router + .release_permit_neutral( + &provider.id, + app_type_str, + used_half_open_permit, + ) + .await; { let mut status = self.status.write().await; status.failed_requests += 1; diff --git a/src-tauri/src/proxy/handler_context.rs b/src-tauri/src/proxy/handler_context.rs index c1f5ffe91..987eed406 100644 --- a/src-tauri/src/proxy/handler_context.rs +++ b/src-tauri/src/proxy/handler_context.rs @@ -175,10 +175,9 @@ impl RequestContext { /// Gemini API 的模型名称在 URI 中，格式如： /// `/v1beta/models/gemini-pro:generateContent` pub fn with_model_from_uri(mut self, uri: &axum::http::Uri) -> Self { - let endpoint = uri - .path_and_query() - .map(|pq| pq.as_str()) - .unwrap_or(uri.path()); + // 用 path() 而不是 path_and_query()：模型名必须从路径段中解析， + // 否则 GET /v1beta/models/?key=... 会把 query 拼到 request_model 上。 + let endpoint = uri.path(); self.request_model = extract_gemini_model_from_path(endpoint).unwrap_or_else(|| "unknown".to_string()); @@ -285,6 +284,8 @@ pub(crate) fn extract_gemini_model_from_path(endpoint: &str) -> Option { .iter() .position(|s| *s == "models") .and_then(|i| segments.get(i + 1).copied()) + // 防御性裁剪：即便调用方传入带 ? 或 :action 的字符串，也只保留 model id 本身 + .map(|s| s.split('?').next().unwrap_or(s)) .map(|s| s.split(':').next().unwrap_or(s)) .filter(|s| !s.is_empty()) .map(|s| s.to_string()) @@ -347,4 +348,23 @@ mod tests { // `/v1beta/models` (list endpoint) has no following segment → None. assert_eq!(extract_gemini_model_from_path("/v1beta/models"), None); } + + #[test] + fn extract_model_get_with_query_only() { + // GET /v1beta/models/?key=... 无 action verb，仅靠 ':' 拆分会把 query 带进 model 名。 + // 修复后应该把 query 剥掉。 + assert_eq!( + extract_gemini_model_from_path("/v1beta/models/gemini-pro?key=abc").as_deref(), + Some("gemini-pro"), + ); + } + + #[test] + fn extract_model_get_with_proxy_prefix_and_query() { + assert_eq!( + extract_gemini_model_from_path("/gemini/v1beta/models/gemini-2.0-flash?key=abc") + .as_deref(), + Some("gemini-2.0-flash"), + ); + } } diff --git a/src-tauri/src/proxy/handlers.rs b/src-tauri/src/proxy/handlers.rs index 71a95e946..85fdf0677 100644 --- a/src-tauri/src/proxy/handlers.rs +++ b/src-tauri/src/proxy/handlers.rs @@ -9,6 +9,7 @@ use super::{ error_mapper::{get_error_message, map_proxy_error_to_status}, + forwarder::ActiveConnectionGuard, handler_config::{ claude_stream_usage_event_filter, CLAUDE_PARSER_CONFIG, CODEX_PARSER_CONFIG, GEMINI_PARSER_CONFIG, OPENAI_PARSER_CONFIG, @@ -145,7 +146,7 @@ async fn handle_messages_for_app( // 转发请求 let forwarder = ctx.create_forwarder(&state); - let result = match forwarder + let mut result = match forwarder .forward_with_retry( &app_type, method, @@ -167,6 +168,7 @@ async fn handle_messages_for_app( } }; + let connection_guard = result.connection_guard.take(); ctx.provider = result.provider; let api_format = result .claude_api_format @@ -181,12 +183,27 @@ async fn handle_messages_for_app( // Claude 特有：格式转换处理 if needs_transform { - return handle_claude_transform(response, &ctx, &state, &body, is_stream, &api_format) - .await; + return handle_claude_transform( + response, + &ctx, + &state, + &body, + is_stream, + &api_format, + connection_guard, + ) + .await; } // 通用响应处理（透传模式） - process_response(response, &ctx, &state, &CLAUDE_PARSER_CONFIG).await + process_response( + response, + &ctx, + &state, + &CLAUDE_PARSER_CONFIG, + connection_guard, + ) + .await } fn validate_claude_desktop_gateway_auth( @@ -226,6 +243,7 @@ async fn handle_claude_transform( original_body: &Value, is_stream: bool, api_format: &str, + connection_guard: Option, ) -> Result { let status = response.status(); let is_codex_oauth = ctx @@ -326,6 +344,7 @@ async fn handle_claude_transform( "Claude/OpenRouter", usage_collector, timeout_config, + connection_guard, ); let mut headers = axum::http::HeaderMap::new(); @@ -477,7 +496,7 @@ pub async fn handle_chat_completions( .unwrap_or(false); let forwarder = ctx.create_forwarder(&state); - let result = match forwarder + let mut result = match forwarder .forward_with_retry( &AppType::Codex, method, @@ -499,10 +518,18 @@ pub async fn handle_chat_completions( } }; + let connection_guard = result.connection_guard.take(); ctx.provider = result.provider; let response = result.response; - process_response(response, &ctx, &state, &OPENAI_PARSER_CONFIG).await + process_response( + response, + &ctx, + &state, + &OPENAI_PARSER_CONFIG, + connection_guard, + ) + .await } /// 处理 /v1/responses 请求（OpenAI Responses API - Codex CLI 透传） @@ -533,7 +560,7 @@ pub async fn handle_responses( .unwrap_or(false); let forwarder = ctx.create_forwarder(&state); - let result = match forwarder + let mut result = match forwarder .forward_with_retry( &AppType::Codex, method, @@ -555,10 +582,18 @@ pub async fn handle_responses( } }; + let connection_guard = result.connection_guard.take(); ctx.provider = result.provider; let response = result.response; - process_response(response, &ctx, &state, &CODEX_PARSER_CONFIG).await + process_response( + response, + &ctx, + &state, + &CODEX_PARSER_CONFIG, + connection_guard, + ) + .await } /// 处理 /v1/responses/compact 请求（OpenAI Responses Compact API - Codex CLI 透传） @@ -589,7 +624,7 @@ pub async fn handle_responses_compact( .unwrap_or(false); let forwarder = ctx.create_forwarder(&state); - let result = match forwarder + let mut result = match forwarder .forward_with_retry( &AppType::Codex, method, @@ -611,10 +646,18 @@ pub async fn handle_responses_compact( } }; + let connection_guard = result.connection_guard.take(); ctx.provider = result.provider; let response = result.response; - process_response(response, &ctx, &state, &CODEX_PARSER_CONFIG).await + process_response( + response, + &ctx, + &state, + &CODEX_PARSER_CONFIG, + connection_guard, + ) + .await } // ============================================================================ @@ -662,7 +705,7 @@ pub async fn handle_gemini( .unwrap_or(false); let forwarder = ctx.create_forwarder(&state); - let result = match forwarder + let mut result = match forwarder .forward_with_retry( &AppType::Gemini, method, @@ -684,10 +727,18 @@ pub async fn handle_gemini( } }; + let connection_guard = result.connection_guard.take(); ctx.provider = result.provider; let response = result.response; - process_response(response, &ctx, &state, &GEMINI_PARSER_CONFIG).await + process_response( + response, + &ctx, + &state, + &GEMINI_PARSER_CONFIG, + connection_guard, + ) + .await } fn should_use_claude_transform_streaming( diff --git a/src-tauri/src/proxy/response_processor.rs b/src-tauri/src/proxy/response_processor.rs index 9f00f8860..94de7d326 100644 --- a/src-tauri/src/proxy/response_processor.rs +++ b/src-tauri/src/proxy/response_processor.rs @@ -3,6 +3,7 @@ //! 统一处理流式和非流式 API 响应 use super::{ + forwarder::ActiveConnectionGuard, handler_config::{StreamUsageEventFilter, UsageParserConfig}, handler_context::{RequestContext, StreamingTimeoutConfig}, hyper_client::ProxyResponse, @@ -181,6 +182,7 @@ pub async fn handle_streaming( ctx: &RequestContext, state: &ProxyState, parser_config: &UsageParserConfig, + connection_guard: Option, ) -> Response { let status = response.status(); log::debug!( @@ -218,8 +220,13 @@ pub async fn handle_streaming( let timeout_config = ctx.streaming_timeout_config(); // 创建带日志和超时的透传流 - let logged_stream = - create_logged_passthrough_stream(stream, ctx.tag, usage_collector, timeout_config); + let logged_stream = create_logged_passthrough_stream( + stream, + ctx.tag, + usage_collector, + timeout_config, + connection_guard, + ); let body = axum::body::Body::from_stream(logged_stream); match builder.body(body) { @@ -237,6 +244,8 @@ pub async fn handle_non_streaming( ctx: &RequestContext, state: &ProxyState, parser_config: &UsageParserConfig, + // guard 在函数 scope 内持有，整包响应读取完成后随函数返回一并 drop + _connection_guard: Option, ) -> Result { // 整包超时：仅在故障转移开启且配置值非零时生效 let body_timeout = @@ -339,11 +348,12 @@ pub async fn process_response( ctx: &RequestContext, state: &ProxyState, parser_config: &UsageParserConfig, + connection_guard: Option, ) -> Result { if is_sse_response(&response) { - Ok(handle_streaming(response, ctx, state, parser_config).await) + Ok(handle_streaming(response, ctx, state, parser_config, connection_guard).await) } else { - handle_non_streaming(response, ctx, state, parser_config).await + handle_non_streaming(response, ctx, state, parser_config, connection_guard).await } } @@ -631,8 +641,10 @@ pub fn create_logged_passthrough_stream( tag: &'static str, usage_collector: Option, timeout_config: StreamingTimeoutConfig, + connection_guard: Option, ) -> impl Stream> + Send { async_stream::stream! { + let _conn_guard = connection_guard; let mut buffer = String::new(); let mut utf8_remainder: Vec = Vec::new(); let mut collector = usage_collector; diff --git a/src-tauri/src/services/stream_check.rs b/src-tauri/src/services/stream_check.rs index 9acfb51d4..6e721e8b3 100644 --- a/src-tauri/src/services/stream_check.rs +++ b/src-tauri/src/services/stream_check.rs @@ -444,7 +444,10 @@ impl StreamCheckService { // - AuthStrategy::ClaudeAuth → Authorization: Bearer // - AuthStrategy::Bearer → Authorization: Bearer // 避免之前"无条件 Bearer + 条件 x-api-key 双发"导致的假阴性 / auth conflict。 - for (name, value) in ClaudeAdapter::new().get_auth_headers(auth) { + let auth_headers = ClaudeAdapter::new() + .get_auth_headers(auth) + .map_err(|e| AppError::Message(format!("stream check 构造鉴权头失败: {e}")))?; + for (name, value) in auth_headers { request_builder = request_builder.header(name, value); }