admin/CC-Switch
1
0
Fork 0
mirror of https://github.com/farion1231/cc-switch.git synced 2026-07-24 12:44:18 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(proxy): extract session ID from client requests for logging

Browse Source 
- Add SessionIdExtractor to parse session ID from Claude/Codex requests
- Support extraction from metadata.user_id, headers, previous_response_id
- Pass session_id through RequestContext to usage logger
- Enable request correlation by session in proxy_request_logs
This commit is contained in:
YoVinchen
2025-12-31 21:31:20 +08:00
parent 48a8aabfb8
commit 766ca23eca
5 changed files with 317 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src-tauri/src/proxy/handler_context.rs
+25 -4
View File
@@ -5,8 +5,10 @@
use crate::app_config::AppType;
use crate::provider::Provider;
use crate::proxy::{
forwarder::RequestForwarder, server::ProxyState, types::AppProxyConfig, ProxyError,
extract_session_id, forwarder::RequestForwarder, server::ProxyState, types::AppProxyConfig,
ProxyError,
};
use axum::http::HeaderMap;
use std::time::Instant;
/// 流式超时配置
@@ -26,6 +28,7 @@ pub struct StreamingTimeoutConfig {
/// - 选中的 Provider 列表（用于故障转移）
/// - 请求模型名称
/// - 日志标签
/// - Session ID（用于日志关联）
pub struct RequestContext {
/// 请求开始时间
pub start_time: Instant,
@@ -35,7 +38,7 @@ pub struct RequestContext {
pub provider: Provider,
/// 完整的 Provider 列表（用于故障转移）
providers: Vec<Provider>,
/// 请求开始时的“当前供应商”（用于判断是否需要同步 UI/托盘）
/// 请求开始时的"当前供应商"（用于判断是否需要同步 UI/托盘）
///
/// 这里使用本地 settings 的设备级 current provider。
/// 代理模式下如果实际使用的 provider 与此不一致，会触发切换以确保 UI 始终准确。
@@ -49,6 +52,8 @@ pub struct RequestContext {
/// 应用类型（预留，目前通过 app_type_str 使用）
#[allow(dead_code)]
pub app_type: AppType,
/// Session ID（从客户端请求提取或新生成）
pub session_id: String,
}
impl RequestContext {
@@ -57,6 +62,7 @@ impl RequestContext {
/// # Arguments
/// * `state` - 代理服务器状态
/// * `body` - 请求体 JSON
/// * `headers` - 请求头（用于提取 Session ID）
/// * `app_type` - 应用类型
/// * `tag` - 日志标签
/// * `app_type_str` - 应用类型字符串
@@ -66,6 +72,7 @@ impl RequestContext {
pub async fn new(
state: &ProxyState,
body: &serde_json::Value,
headers: &HeaderMap,
app_type: AppType,
tag: &'static str,
app_type_str: &'static str,
@@ -89,6 +96,18 @@ impl RequestContext {
.unwrap_or("unknown")
.to_string();
// 提取 Session ID
let session_result = extract_session_id(headers, body, app_type_str);
let session_id = session_result.session_id.clone();
log::debug!(
"[{}] Session ID: {} (from {:?}, client_provided: {})",
tag,
session_id,
session_result.source,
session_result.client_provided
);
// 使用共享的 ProviderRouter 选择 Provider（熔断器状态跨请求保持）
// 注意：只在这里调用一次，结果传递给 forwarder，避免重复消耗 HalfOpen 名额
let providers = state
@@ -109,11 +128,12 @@ impl RequestContext {
.ok_or(ProxyError::NoAvailableProvider)?;
log::info!(
"[{}] Provider: {}, model: {}, failover chain: {} providers",
"[{}] Provider: {}, model: {}, failover chain: {} providers, session: {}",
tag,
provider.name,
request_model,
providers.len()
providers.len(),
session_id
);
Ok(Self {
@@ -126,6 +146,7 @@ impl RequestContext {
tag,
app_type_str,
app_type,
session_id,
})
}
src-tauri/src/proxy/handlers.rs
+9 -6
View File
@@ -61,7 +61,8 @@ pub async fn handle_messages(
headers: axum::http::HeaderMap,
Json(body): Json<Value>,
) -> Result<axum::response::Response, ProxyError> {
let mut ctx = RequestContext::new(&state, &body, AppType::Claude, "Claude", "claude").await?;
let mut ctx =
RequestContext::new(&state, &body, &headers, AppType::Claude, "Claude", "claude").await?;
let is_stream = body
.get("stream")
@@ -305,7 +306,8 @@ pub async fn handle_chat_completions(
) -> Result<axum::response::Response, ProxyError> {
log::info!("[Codex] ====== /v1/chat/completions 请求开始 ======");
let mut ctx = RequestContext::new(&state, &body, AppType::Codex, "Codex", "codex").await?;
let mut ctx =
RequestContext::new(&state, &body, &headers, AppType::Codex, "Codex", "codex").await?;
let is_stream = body
.get("stream")
@@ -353,7 +355,8 @@ pub async fn handle_responses(
headers: axum::http::HeaderMap,
Json(body): Json<Value>,
) -> Result<axum::response::Response, ProxyError> {
let mut ctx = RequestContext::new(&state, &body, AppType::Codex, "Codex", "codex").await?;
let mut ctx =
RequestContext::new(&state, &body, &headers, AppType::Codex, "Codex", "codex").await?;
let is_stream = body
.get("stream")
@@ -401,7 +404,7 @@ pub async fn handle_gemini(
Json(body): Json<Value>,
) -> Result<axum::response::Response, ProxyError> {
// Gemini 的模型名称在 URI 中
let mut ctx = RequestContext::new(&state, &body, AppType::Gemini, "Gemini", "gemini")
let mut ctx = RequestContext::new(&state, &body, &headers, AppType::Gemini, "Gemini", "gemini")
.await?
.with_model_from_uri(&uri);
@@ -465,7 +468,7 @@ fn log_forward_error(
let request_id = uuid::Uuid::new_v4().to_string();
if let Err(e) = logger.log_error_with_context(
request_id.clone(),
request_id,
ctx.provider.id.clone(),
ctx.app_type_str.to_string(),
ctx.request_model.clone(),
@@ -473,7 +476,7 @@ fn log_forward_error(
error_message,
ctx.latency_ms(),
is_streaming,
Some(request_id),
Some(ctx.session_id.clone()),
None,
) {
log::warn!("记录失败请求日志失败: {e}");
src-tauri/src/proxy/mod.rs
+3 -1
View File
@@ -34,7 +34,9 @@ pub use provider_router::ProviderRouter;
#[allow(unused_imports)]
pub use response_handler::{NonStreamHandler, ResponseType, StreamHandler};
#[allow(unused_imports)]
pub use session::{ClientFormat, ProxySession};
pub use session::{
extract_session_id, ClientFormat, ProxySession, SessionIdResult, SessionIdSource,
};
#[allow(unused_imports)]
pub use types::{ProxyConfig, ProxyServerInfo, ProxyStatus};
src-tauri/src/proxy/response_processor.rs
+11 -2
View File
@@ -264,6 +264,7 @@ fn create_usage_collector(
let start_time = ctx.start_time;
let stream_parser = parser_config.stream_parser;
let model_extractor = parser_config.model_extractor;
let session_id = ctx.session_id.clone();
SseUsageCollector::new(start_time, move |events, first_token_ms| {
if let Some(usage) = stream_parser(&events) {
@@ -272,6 +273,7 @@ fn create_usage_collector(
let state = state.clone();
let provider_id = provider_id.clone();
let session_id = session_id.clone();
tokio::spawn(async move {
log_usage_internal(
@@ -284,6 +286,7 @@ fn create_usage_collector(
first_token_ms,
true, // is_streaming
status_code,
Some(session_id),
)
.await;
});
@@ -292,6 +295,7 @@ fn create_usage_collector(
let latency_ms = start_time.elapsed().as_millis() as u64;
let state = state.clone();
let provider_id = provider_id.clone();
let session_id = session_id.clone();
tokio::spawn(async move {
log_usage_internal(
@@ -304,6 +308,7 @@ fn create_usage_collector(
first_token_ms,
true, // is_streaming
status_code,
Some(session_id),
)
.await;
});
@@ -326,6 +331,7 @@ fn spawn_log_usage(
let app_type_str = ctx.app_type_str.to_string();
let model = model.to_string();
let latency_ms = ctx.latency_ms();
let session_id = ctx.session_id.clone();
tokio::spawn(async move {
log_usage_internal(
@@ -338,6 +344,7 @@ fn spawn_log_usage(
None,
is_streaming,
status_code,
Some(session_id),
)
.await;
});
@@ -355,6 +362,7 @@ async fn log_usage_internal(
first_token_ms: Option<u64>,
is_streaming: bool,
status_code: u16,
session_id: Option<String>,
) {
use super::usage::logger::UsageLogger;
@@ -379,7 +387,8 @@ async fn log_usage_internal(
let request_id = uuid::Uuid::new_v4().to_string();
log::debug!(
"[{app_type}] 记录请求日志: id={request_id}, provider={provider_id}, model={model}, streaming={is_streaming}, status={status_code}, latency_ms={latency_ms}, first_token_ms={first_token_ms:?}, input={}, output={}, cache_read={}, cache_creation={}",
"[{app_type}] 记录请求日志: id={request_id}, provider={provider_id}, model={model}, streaming={is_streaming}, status={status_code}, latency_ms={latency_ms}, first_token_ms={first_token_ms:?}, session={}, input={}, output={}, cache_read={}, cache_creation={}",
session_id.as_deref().unwrap_or("none"),
usage.input_tokens,
usage.output_tokens,
usage.cache_read_tokens,
@@ -396,7 +405,7 @@ async fn log_usage_internal(
latency_ms,
first_token_ms,
status_code,
None,
session_id,
None, // provider_type
is_streaming,
) {
src-tauri/src/proxy/session.rs
+269
View File
@@ -1,7 +1,15 @@
//! Proxy Session - 请求会话管理
//!
//! 为每个代理请求创建会话上下文，在整个请求生命周期中跟踪状态和元数据。
//!
//! ## Session ID 提取
//!
//! 支持从客户端请求中提取 Session ID，用于关联同一对话的多个请求：
//! - Claude: 从 `metadata.user_id` (格式: `user_xxx_session_yyy`) 或 `metadata.session_id` 提取
//! - Codex: 从 `previous_response_id` 或 headers 中的 `session_id` 提取
//! - 其他: 生成新的 UUID
use axum::http::HeaderMap;
use std::time::Instant;
use uuid::Uuid;
@@ -176,6 +184,179 @@ impl ProxySession {
}
}
// ============================================================================
// Session ID 提取器
// ============================================================================
/// Session ID 来源
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SessionIdSource {
/// 从 metadata.user_id 提取 (Claude)
MetadataUserId,
/// 从 metadata.session_id 提取
MetadataSessionId,
/// 从 headers 提取 (Codex)
Header,
/// 从 previous_response_id 提取 (Codex)
PreviousResponseId,
/// 新生成
Generated,
}
/// Session ID 提取结果
#[derive(Debug, Clone)]
pub struct SessionIdResult {
/// 提取或生成的 Session ID
pub session_id: String,
/// Session ID 来源
pub source: SessionIdSource,
/// 是否为客户端提供的 ID（非新生成）
pub client_provided: bool,
}
/// 从请求中提取或生成 Session ID
///
/// 轻量化实现，仅提取 session_id 用于日志记录，不做复杂的 Session 管理。
///
/// ## 提取优先级
///
/// ### Claude 请求
/// 1. `metadata.user_id` (格式: `user_xxx_session_yyy`) → 提取 `yyy` 部分
/// 2. `metadata.session_id` → 直接使用
/// 3. 生成新 UUID
///
/// ### Codex 请求
/// 1. Headers: `session_id` 或 `x-session-id`
/// 2. `metadata.session_id`
/// 3. `previous_response_id` (对话延续)
/// 4. 生成新 UUID
///
/// ## 示例
///
/// ```ignore
/// let result = extract_session_id(&headers, &body, "claude");
/// println!("Session ID: {} (from {:?})", result.session_id, result.source);
/// ```
pub fn extract_session_id(
headers: &HeaderMap,
body: &serde_json::Value,
client_format: &str,
) -> SessionIdResult {
// Codex 请求特殊处理
if client_format == "codex" || client_format == "openai" {
if let Some(result) = extract_codex_session(headers, body) {
return result;
}
}
// Claude 请求：从 metadata 提取
if let Some(result) = extract_from_metadata(body) {
return result;
}
// 兜底：生成新 Session ID
generate_new_session_id()
}
/// 提取 Codex Session ID
fn extract_codex_session(headers: &HeaderMap, body: &serde_json::Value) -> Option<SessionIdResult> {
// 1. 从 headers 提取
for header_name in &["session_id", "x-session-id"] {
if let Some(value) = headers.get(*header_name) {
if let Ok(session_id) = value.to_str() {
// Codex Session ID 通常较长（UUID 格式）
if session_id.len() > 20 {
return Some(SessionIdResult {
session_id: format!("codex_{session_id}"),
source: SessionIdSource::Header,
client_provided: true,
});
}
}
}
}
// 2. 从 body.metadata.session_id 提取
if let Some(session_id) = body
.get("metadata")
.and_then(|m| m.get("session_id"))
.and_then(|v| v.as_str())
{
if session_id.len() > 10 {
return Some(SessionIdResult {
session_id: format!("codex_{session_id}"),
source: SessionIdSource::MetadataSessionId,
client_provided: true,
});
}
}
// 3. 从 previous_response_id 提取（对话延续）
if let Some(prev_id) = body.get("previous_response_id").and_then(|v| v.as_str()) {
if prev_id.len() > 10 {
return Some(SessionIdResult {
session_id: format!("codex_{prev_id}"),
source: SessionIdSource::PreviousResponseId,
client_provided: true,
});
}
}
None
}
/// 从 metadata 提取 Session ID (Claude)
fn extract_from_metadata(body: &serde_json::Value) -> Option<SessionIdResult> {
let metadata = body.get("metadata")?;
// 1. 从 metadata.user_id 提取（格式: user_xxx_session_yyy）
if let Some(user_id) = metadata.get("user_id").and_then(|v| v.as_str()) {
if let Some(session_id) = parse_session_from_user_id(user_id) {
return Some(SessionIdResult {
session_id,
source: SessionIdSource::MetadataUserId,
client_provided: true,
});
}
}
// 2. 直接从 metadata.session_id 提取
if let Some(session_id) = metadata.get("session_id").and_then(|v| v.as_str()) {
if !session_id.is_empty() {
return Some(SessionIdResult {
session_id: session_id.to_string(),
source: SessionIdSource::MetadataSessionId,
client_provided: true,
});
}
}
None
}
/// 从 user_id 解析 session_id
///
/// 格式: `user_identifier_session_actual_session_id`
fn parse_session_from_user_id(user_id: &str) -> Option<String> {
// 查找 "_session_" 分隔符
if let Some(pos) = user_id.find("_session_") {
let session_id = &user_id[pos + 9..]; // "_session_" 长度为 9
if !session_id.is_empty() {
return Some(session_id.to_string());
}
}
None
}
/// 生成新的 Session ID
fn generate_new_session_id() -> SessionIdResult {
SessionIdResult {
session_id: Uuid::new_v4().to_string(),
source: SessionIdSource::Generated,
client_provided: false,
}
}
#[cfg(test)]
mod tests {
use super::*;
@@ -295,4 +476,92 @@ mod tests {
assert_eq!(ClientFormat::GeminiCli.as_str(), "gemini_cli");
assert_eq!(ClientFormat::Unknown.as_str(), "unknown");
}
// ========== Session ID 提取测试 ==========
#[test]
fn test_extract_session_from_claude_metadata_user_id() {
let headers = HeaderMap::new();
let body = json!({
"model": "claude-3-5-sonnet",
"messages": [{"role": "user", "content": "Hello"}],
"metadata": {
"user_id": "user_john_doe_session_abc123def456"
}
});
let result = extract_session_id(&headers, &body, "claude");
assert_eq!(result.session_id, "abc123def456");
assert_eq!(result.source, SessionIdSource::MetadataUserId);
assert!(result.client_provided);
}
#[test]
fn test_extract_session_from_claude_metadata_session_id() {
let headers = HeaderMap::new();
let body = json!({
"model": "claude-3-5-sonnet",
"messages": [{"role": "user", "content": "Hello"}],
"metadata": {
"session_id": "my-session-123"
}
});
let result = extract_session_id(&headers, &body, "claude");
assert_eq!(result.session_id, "my-session-123");
assert_eq!(result.source, SessionIdSource::MetadataSessionId);
assert!(result.client_provided);
}
#[test]
fn test_extract_session_from_codex_previous_response_id() {
let headers = HeaderMap::new();
let body = json!({
"input": "Write a function",
"previous_response_id": "resp_abc123def456789"
});
let result = extract_session_id(&headers, &body, "codex");
assert_eq!(result.session_id, "codex_resp_abc123def456789");
assert_eq!(result.source, SessionIdSource::PreviousResponseId);
assert!(result.client_provided);
}
#[test]
fn test_extract_session_generates_new_when_not_found() {
let headers = HeaderMap::new();
let body = json!({
"model": "claude-3-5-sonnet",
"messages": [{"role": "user", "content": "Hello"}]
});
let result = extract_session_id(&headers, &body, "claude");
assert!(!result.session_id.is_empty());
assert_eq!(result.source, SessionIdSource::Generated);
assert!(!result.client_provided);
}
#[test]
fn test_parse_session_from_user_id() {
assert_eq!(
parse_session_from_user_id("user_john_session_abc123"),
Some("abc123".to_string())
);
assert_eq!(
parse_session_from_user_id("my_app_session_xyz789"),
Some("xyz789".to_string())
);
// 注意: "_session_" 是分隔符，所以下面的字符串会匹配
assert_eq!(
parse_session_from_user_id("no_session_marker"),
Some("marker".to_string())
);
// 没有 "_session_" 分隔符的情况
assert_eq!(parse_session_from_user_id("user_john_abc123"), None);
assert_eq!(parse_session_from_user_id("_session_"), None);
}
}