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:
Jason
2026-05-14 09:23:21 +08:00
parent 85131d37d8
commit 206125b4e3
6 changed files with 188 additions and 47 deletions
+1 -1
View File
@@ -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();
+80 -25
View File
@@ -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` +1Drop 时在 tokio runtime 上调度
/// 一个异步任务执行 -1,从而支持把 guard move 进流式 body futurestream 自然结束
/// 时 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;
+24 -4
View File
@@ -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"),
);
}
}
+63 -12
View File
@@ -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(
+16 -4
View File
@@ -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;
+4 -1
View File
@@ -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);
}