mirror of
https://github.com/farion1231/cc-switch.git
synced 2026-08-01 04:02:02 +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:
@@ -17,7 +17,8 @@ impl Database {
|
||||
let conn = lock_conn!(self.conn);
|
||||
conn.query_row(
|
||||
"SELECT listen_address, listen_port, max_retries,
|
||||
request_timeout, enable_logging, live_takeover_active
|
||||
request_timeout, enable_logging, live_takeover_active,
|
||||
streaming_first_byte_timeout, streaming_idle_timeout, non_streaming_timeout
|
||||
FROM proxy_config WHERE id = 1",
|
||||
[],
|
||||
|row| {
|
||||
@@ -28,6 +29,9 @@ impl Database {
|
||||
request_timeout: row.get::<_, i32>(3)? as u64,
|
||||
enable_logging: row.get::<_, i32>(4)? != 0,
|
||||
live_takeover_active: row.get::<_, i32>(5).unwrap_or(0) != 0,
|
||||
streaming_first_byte_timeout: row.get::<_, i32>(6).unwrap_or(30) as u64,
|
||||
streaming_idle_timeout: row.get::<_, i32>(7).unwrap_or(60) as u64,
|
||||
non_streaming_timeout: row.get::<_, i32>(8).unwrap_or(300) as u64,
|
||||
})
|
||||
},
|
||||
)
|
||||
@@ -51,8 +55,11 @@ impl Database {
|
||||
|
||||
conn.execute(
|
||||
"INSERT OR REPLACE INTO proxy_config
|
||||
(id, enabled, listen_address, listen_port, max_retries, request_timeout, enable_logging, live_takeover_active, target_app, created_at, updated_at)
|
||||
VALUES (1, ?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8,
|
||||
(id, enabled, listen_address, listen_port, max_retries, request_timeout,
|
||||
enable_logging, live_takeover_active, target_app,
|
||||
streaming_first_byte_timeout, streaming_idle_timeout, non_streaming_timeout,
|
||||
created_at, updated_at)
|
||||
VALUES (1, ?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11,
|
||||
COALESCE((SELECT created_at FROM proxy_config WHERE id = 1), datetime('now')),
|
||||
datetime('now'))",
|
||||
rusqlite::params![
|
||||
@@ -64,6 +71,9 @@ impl Database {
|
||||
if config.enable_logging { 1 } else { 0 },
|
||||
if config.live_takeover_active { 1 } else { 0 },
|
||||
"claude", // 兼容旧字段,写入默认值
|
||||
config.streaming_first_byte_timeout as i32,
|
||||
config.streaming_idle_timeout as i32,
|
||||
config.non_streaming_timeout as i32,
|
||||
],
|
||||
)
|
||||
.map_err(|e| AppError::Database(e.to_string()))?;
|
||||
|
||||
@@ -135,6 +135,9 @@ impl Database {
|
||||
request_timeout INTEGER NOT NULL DEFAULT 300,
|
||||
enable_logging INTEGER NOT NULL DEFAULT 1,
|
||||
target_app TEXT NOT NULL DEFAULT 'claude',
|
||||
streaming_first_byte_timeout INTEGER NOT NULL DEFAULT 30,
|
||||
streaming_idle_timeout INTEGER NOT NULL DEFAULT 60,
|
||||
non_streaming_timeout INTEGER NOT NULL DEFAULT 300,
|
||||
created_at TEXT NOT NULL DEFAULT (datetime('now')),
|
||||
updated_at TEXT NOT NULL DEFAULT (datetime('now'))
|
||||
)",
|
||||
@@ -313,6 +316,20 @@ impl Database {
|
||||
[],
|
||||
);
|
||||
|
||||
// 尝试添加超时配置列到 proxy_config 表(v3 新增)
|
||||
let _ = conn.execute(
|
||||
"ALTER TABLE proxy_config ADD COLUMN streaming_first_byte_timeout INTEGER NOT NULL DEFAULT 30",
|
||||
[],
|
||||
);
|
||||
let _ = conn.execute(
|
||||
"ALTER TABLE proxy_config ADD COLUMN streaming_idle_timeout INTEGER NOT NULL DEFAULT 60",
|
||||
[],
|
||||
);
|
||||
let _ = conn.execute(
|
||||
"ALTER TABLE proxy_config ADD COLUMN non_streaming_timeout INTEGER NOT NULL DEFAULT 300",
|
||||
[],
|
||||
);
|
||||
|
||||
// 确保 in_failover_queue 列存在(对于已存在的 v2 数据库)
|
||||
Self::add_column_if_missing(
|
||||
conn,
|
||||
@@ -475,6 +492,28 @@ impl Database {
|
||||
"BOOLEAN NOT NULL DEFAULT 0",
|
||||
)?;
|
||||
|
||||
// 添加代理超时配置字段(在 v2 迁移中直接添加)
|
||||
if Self::table_exists(conn, "proxy_config")? {
|
||||
Self::add_column_if_missing(
|
||||
conn,
|
||||
"proxy_config",
|
||||
"streaming_first_byte_timeout",
|
||||
"INTEGER NOT NULL DEFAULT 30",
|
||||
)?;
|
||||
Self::add_column_if_missing(
|
||||
conn,
|
||||
"proxy_config",
|
||||
"streaming_idle_timeout",
|
||||
"INTEGER NOT NULL DEFAULT 60",
|
||||
)?;
|
||||
Self::add_column_if_missing(
|
||||
conn,
|
||||
"proxy_config",
|
||||
"non_streaming_timeout",
|
||||
"INTEGER NOT NULL DEFAULT 300",
|
||||
)?;
|
||||
}
|
||||
|
||||
// 删除旧的 failover_queue 表(如果存在)
|
||||
conn.execute("DROP INDEX IF EXISTS idx_failover_queue_order", [])
|
||||
.map_err(|e| AppError::Database(format!("删除 failover_queue 索引失败: {e}")))?;
|
||||
|
||||
@@ -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