mirror of
https://github.com/farion1231/cc-switch.git
synced 2026-07-30 02:14:43 +08:00
feat(proxy): implement streaming timeout control with validation
- Add first byte timeout (0 or 1-180s) for streaming requests - Add idle timeout (0 or 60-600s) for streaming data gaps - Add non-streaming timeout (0 or 60-1800s) for total request - Implement timeout logic in response processor - Add 1800s global timeout fallback when disabled - Add database schema migration for timeout fields - Add i18n translations for timeout settings
This commit is contained in:
@@ -39,7 +39,7 @@ pub struct RequestForwarder {
|
||||
failover_manager: Arc<FailoverSwitchManager>,
|
||||
/// AppHandle,用于发射事件和更新托盘
|
||||
app_handle: Option<tauri::AppHandle>,
|
||||
/// 请求开始时的“当前供应商 ID”(用于判断是否需要同步 UI/托盘)
|
||||
/// 请求开始时的"当前供应商 ID"(用于判断是否需要同步 UI/托盘)
|
||||
current_provider_id_at_start: String,
|
||||
}
|
||||
|
||||
@@ -47,17 +47,27 @@ impl RequestForwarder {
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn new(
|
||||
router: Arc<ProviderRouter>,
|
||||
timeout_secs: u64,
|
||||
non_streaming_timeout: u64,
|
||||
max_retries: u8,
|
||||
status: Arc<RwLock<ProxyStatus>>,
|
||||
current_providers: Arc<RwLock<std::collections::HashMap<String, (String, String)>>>,
|
||||
failover_manager: Arc<FailoverSwitchManager>,
|
||||
app_handle: Option<tauri::AppHandle>,
|
||||
current_provider_id_at_start: String,
|
||||
_streaming_first_byte_timeout: u64,
|
||||
_streaming_idle_timeout: u64,
|
||||
) -> Self {
|
||||
// 全局超时设置为 1800 秒(30 分钟),确保业务层超时配置能正常工作
|
||||
// 参考 Claude Code Hub 的 undici 全局超时设计
|
||||
const GLOBAL_TIMEOUT_SECS: u64 = 1800;
|
||||
|
||||
let mut client_builder = Client::builder();
|
||||
if timeout_secs > 0 {
|
||||
client_builder = client_builder.timeout(Duration::from_secs(timeout_secs));
|
||||
if non_streaming_timeout > 0 {
|
||||
// 使用配置的非流式超时
|
||||
client_builder = client_builder.timeout(Duration::from_secs(non_streaming_timeout));
|
||||
} else {
|
||||
// 禁用超时时使用全局超时作为保底
|
||||
client_builder = client_builder.timeout(Duration::from_secs(GLOBAL_TIMEOUT_SECS));
|
||||
}
|
||||
|
||||
let client = client_builder
|
||||
|
||||
@@ -9,6 +9,15 @@ use crate::proxy::{
|
||||
};
|
||||
use std::time::Instant;
|
||||
|
||||
/// 流式超时配置
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub struct StreamingTimeoutConfig {
|
||||
/// 首字节超时(秒),0 表示禁用
|
||||
pub first_byte_timeout: u64,
|
||||
/// 静默期超时(秒),0 表示禁用
|
||||
pub idle_timeout: u64,
|
||||
}
|
||||
|
||||
/// 请求上下文
|
||||
///
|
||||
/// 贯穿整个请求生命周期,包含:
|
||||
@@ -135,13 +144,15 @@ impl RequestContext {
|
||||
pub fn create_forwarder(&self, state: &ProxyState) -> RequestForwarder {
|
||||
RequestForwarder::new(
|
||||
state.provider_router.clone(),
|
||||
self.config.request_timeout,
|
||||
self.config.non_streaming_timeout,
|
||||
self.config.max_retries,
|
||||
state.status.clone(),
|
||||
state.current_providers.clone(),
|
||||
state.failover_manager.clone(),
|
||||
state.app_handle.clone(),
|
||||
self.current_provider_id.clone(),
|
||||
self.config.streaming_first_byte_timeout,
|
||||
self.config.streaming_idle_timeout,
|
||||
)
|
||||
}
|
||||
|
||||
@@ -157,4 +168,13 @@ impl RequestContext {
|
||||
pub fn latency_ms(&self) -> u64 {
|
||||
self.start_time.elapsed().as_millis() as u64
|
||||
}
|
||||
|
||||
/// 获取流式超时配置
|
||||
#[inline]
|
||||
pub fn streaming_timeout_config(&self) -> StreamingTimeoutConfig {
|
||||
StreamingTimeoutConfig {
|
||||
first_byte_timeout: self.config.streaming_first_byte_timeout,
|
||||
idle_timeout: self.config.streaming_idle_timeout,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -170,10 +170,14 @@ async fn handle_claude_transform(
|
||||
})
|
||||
};
|
||||
|
||||
// 获取流式超时配置
|
||||
let timeout_config = ctx.streaming_timeout_config();
|
||||
|
||||
let logged_stream = create_logged_passthrough_stream(
|
||||
sse_stream,
|
||||
"Claude/OpenRouter",
|
||||
Some(usage_collector),
|
||||
timeout_config,
|
||||
);
|
||||
|
||||
let mut headers = axum::http::HeaderMap::new();
|
||||
|
||||
@@ -3,8 +3,11 @@
|
||||
//! 统一处理流式和非流式 API 响应
|
||||
|
||||
use super::{
|
||||
handler_config::UsageParserConfig, handler_context::RequestContext, server::ProxyState,
|
||||
usage::parser::TokenUsage, ProxyError,
|
||||
handler_config::UsageParserConfig,
|
||||
handler_context::{RequestContext, StreamingTimeoutConfig},
|
||||
server::ProxyState,
|
||||
usage::parser::TokenUsage,
|
||||
ProxyError,
|
||||
};
|
||||
use axum::response::Response;
|
||||
use bytes::Bytes;
|
||||
@@ -17,6 +20,7 @@ use std::{
|
||||
atomic::{AtomicBool, Ordering},
|
||||
Arc,
|
||||
},
|
||||
time::Duration,
|
||||
};
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
@@ -60,8 +64,12 @@ pub async fn handle_streaming(
|
||||
// 创建使用量收集器
|
||||
let usage_collector = create_usage_collector(ctx, state, status.as_u16(), parser_config);
|
||||
|
||||
// 创建带日志的透传流
|
||||
let logged_stream = create_logged_passthrough_stream(stream, ctx.tag, Some(usage_collector));
|
||||
// 获取流式超时配置
|
||||
let timeout_config = ctx.streaming_timeout_config();
|
||||
|
||||
// 创建带日志和超时的透传流
|
||||
let logged_stream =
|
||||
create_logged_passthrough_stream(stream, ctx.tag, Some(usage_collector), timeout_config);
|
||||
|
||||
let body = axum::body::Body::from_stream(logged_stream);
|
||||
builder.body(body).unwrap()
|
||||
@@ -348,21 +356,60 @@ async fn log_usage_internal(
|
||||
}
|
||||
}
|
||||
|
||||
/// 创建带日志记录的透传流
|
||||
/// 创建带日志记录和超时控制的透传流
|
||||
pub fn create_logged_passthrough_stream(
|
||||
stream: impl Stream<Item = Result<Bytes, std::io::Error>> + Send + 'static,
|
||||
tag: &'static str,
|
||||
usage_collector: Option<SseUsageCollector>,
|
||||
timeout_config: StreamingTimeoutConfig,
|
||||
) -> impl Stream<Item = Result<Bytes, std::io::Error>> + Send {
|
||||
async_stream::stream! {
|
||||
let mut buffer = String::new();
|
||||
let mut collector = usage_collector;
|
||||
let mut is_first_chunk = true;
|
||||
|
||||
// 超时配置
|
||||
let first_byte_timeout = if timeout_config.first_byte_timeout > 0 {
|
||||
Some(Duration::from_secs(timeout_config.first_byte_timeout))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let idle_timeout = if timeout_config.idle_timeout > 0 {
|
||||
Some(Duration::from_secs(timeout_config.idle_timeout))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
tokio::pin!(stream);
|
||||
|
||||
while let Some(chunk) = stream.next().await {
|
||||
match chunk {
|
||||
Ok(bytes) => {
|
||||
loop {
|
||||
// 选择超时时间:首字节超时或静默期超时
|
||||
let timeout_duration = if is_first_chunk {
|
||||
first_byte_timeout
|
||||
} else {
|
||||
idle_timeout
|
||||
};
|
||||
|
||||
let chunk_result = match timeout_duration {
|
||||
Some(duration) => {
|
||||
match tokio::time::timeout(duration, stream.next()).await {
|
||||
Ok(Some(chunk)) => Some(chunk),
|
||||
Ok(None) => None, // 流结束
|
||||
Err(_) => {
|
||||
// 超时
|
||||
let timeout_type = if is_first_chunk { "首字节" } else { "静默期" };
|
||||
log::error!("[{tag}] 流式响应{}超时 ({}秒)", timeout_type, duration.as_secs());
|
||||
yield Err(std::io::Error::other(format!("流式响应{timeout_type}超时")));
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
None => stream.next().await, // 无超时限制
|
||||
};
|
||||
|
||||
match chunk_result {
|
||||
Some(Ok(bytes)) => {
|
||||
is_first_chunk = false;
|
||||
let text = String::from_utf8_lossy(&bytes);
|
||||
buffer.push_str(&text);
|
||||
|
||||
@@ -398,11 +445,15 @@ pub fn create_logged_passthrough_stream(
|
||||
|
||||
yield Ok(bytes);
|
||||
}
|
||||
Err(e) => {
|
||||
Some(Err(e)) => {
|
||||
log::error!("[{tag}] 流错误: {e}");
|
||||
yield Err(std::io::Error::other(e.to_string()));
|
||||
break;
|
||||
}
|
||||
None => {
|
||||
// 流正常结束
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -9,13 +9,34 @@ pub struct ProxyConfig {
|
||||
pub listen_port: u16,
|
||||
/// 最大重试次数
|
||||
pub max_retries: u8,
|
||||
/// 请求超时时间(秒)
|
||||
/// 请求超时时间(秒)- 已废弃,保留兼容
|
||||
pub request_timeout: u64,
|
||||
/// 是否启用日志
|
||||
pub enable_logging: bool,
|
||||
/// 是否正在接管 Live 配置
|
||||
#[serde(default)]
|
||||
pub live_takeover_active: bool,
|
||||
/// 流式首字超时(秒)- 等待首个数据块的最大时间
|
||||
#[serde(default = "default_streaming_first_byte_timeout")]
|
||||
pub streaming_first_byte_timeout: u64,
|
||||
/// 流式静默超时(秒)- 两个数据块之间的最大间隔
|
||||
#[serde(default = "default_streaming_idle_timeout")]
|
||||
pub streaming_idle_timeout: u64,
|
||||
/// 非流式总超时(秒)- 非流式请求的总超时时间
|
||||
#[serde(default = "default_non_streaming_timeout")]
|
||||
pub non_streaming_timeout: u64,
|
||||
}
|
||||
|
||||
fn default_streaming_first_byte_timeout() -> u64 {
|
||||
30
|
||||
}
|
||||
|
||||
fn default_streaming_idle_timeout() -> u64 {
|
||||
60
|
||||
}
|
||||
|
||||
fn default_non_streaming_timeout() -> u64 {
|
||||
600
|
||||
}
|
||||
|
||||
impl Default for ProxyConfig {
|
||||
@@ -27,6 +48,9 @@ impl Default for ProxyConfig {
|
||||
request_timeout: 300,
|
||||
enable_logging: true,
|
||||
live_takeover_active: false,
|
||||
streaming_first_byte_timeout: 30,
|
||||
streaming_idle_timeout: 60,
|
||||
non_streaming_timeout: 600,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user