mirror of
https://github.com/farion1231/cc-switch.git
synced 2026-07-24 12:44:18 +08:00
fix(proxy): patch P0-P3 routing/lifecycle issues across forwarder paths
* stream_check: thread Result from get_auth_headers via map_err so the workspace builds again * forwarder: scope rectifier / budget-rectifier flags per-provider so failover can still apply rectification on the next attempt * forwarder: categorize before record_result; route NonRetryable and ClientAbort through release_permit_neutral so client-side failures don't pollute circuit breaker or DB health * handler_context: parse Gemini model from uri.path() and strip both ?query and :action verb defensively in extract_gemini_model_from_path * forwarder + response_processor + handlers: introduce ActiveConnectionGuard (RAII) so active_connections decrement covers the full streaming body lifetime, not just response headers * claude_desktop_config: use sort_by_key to clear the clippy gate
This commit is contained in:
@@ -575,7 +575,7 @@ pub fn proxy_model_routes(provider: &Provider) -> Result<Vec<ResolvedModelRoute>
|
||||
.collect::<std::collections::HashSet<_>>();
|
||||
let mut result = Vec::new();
|
||||
let mut entries = routes.iter().collect::<Vec<_>>();
|
||||
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();
|
||||
|
||||
@@ -35,6 +35,9 @@ pub struct ForwardResult {
|
||||
pub response: ProxyResponse,
|
||||
pub provider: Provider,
|
||||
pub claude_api_format: Option<String>,
|
||||
/// 活跃连接 RAII guard:随响应一起流转到 response_processor / handle_claude_transform,
|
||||
/// 最终被 move 进流式 body future(或非流式响应作用域),覆盖整个响应生命周期。
|
||||
pub(crate) connection_guard: Option<ActiveConnectionGuard>,
|
||||
}
|
||||
|
||||
pub struct ForwardError {
|
||||
@@ -42,6 +45,44 @@ pub struct ForwardError {
|
||||
pub provider: Option<Provider>,
|
||||
}
|
||||
|
||||
/// 活跃连接 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<RwLock<ProxyStatus>>,
|
||||
}
|
||||
|
||||
impl ActiveConnectionGuard {
|
||||
pub(crate) async fn acquire(status: Arc<RwLock<ProxyStatus>>) -> 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<ProviderRouter>,
|
||||
@@ -235,10 +276,10 @@ impl RequestForwarder {
|
||||
extensions: Extensions,
|
||||
providers: Vec<Provider>,
|
||||
) -> Result<ForwardResult, ForwardError> {
|
||||
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;
|
||||
|
||||
@@ -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/<id>?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<String> {
|
||||
.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/<id>?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"),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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<ActiveConnectionGuard>,
|
||||
) -> Result<axum::response::Response, ProxyError> {
|
||||
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(
|
||||
|
||||
@@ -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<ActiveConnectionGuard>,
|
||||
) -> 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<ActiveConnectionGuard>,
|
||||
) -> Result<Response, ProxyError> {
|
||||
// 整包超时:仅在故障转移开启且配置值非零时生效
|
||||
let body_timeout =
|
||||
@@ -339,11 +348,12 @@ pub async fn process_response(
|
||||
ctx: &RequestContext,
|
||||
state: &ProxyState,
|
||||
parser_config: &UsageParserConfig,
|
||||
connection_guard: Option<ActiveConnectionGuard>,
|
||||
) -> Result<Response, ProxyError> {
|
||||
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<SseUsageCollector>,
|
||||
timeout_config: StreamingTimeoutConfig,
|
||||
connection_guard: Option<ActiveConnectionGuard>,
|
||||
) -> impl Stream<Item = Result<Bytes, std::io::Error>> + Send {
|
||||
async_stream::stream! {
|
||||
let _conn_guard = connection_guard;
|
||||
let mut buffer = String::new();
|
||||
let mut utf8_remainder: Vec<u8> = Vec::new();
|
||||
let mut collector = usage_collector;
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user