From 766ca23eca77d7d59cae6414bb4a28d0ac3c14b5 Mon Sep 17 00:00:00 2001 From: YoVinchen Date: Wed, 31 Dec 2025 21:31:20 +0800 Subject: [PATCH] feat(proxy): extract session ID from client requests for logging - 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 --- src-tauri/src/proxy/handler_context.rs | 29 ++- src-tauri/src/proxy/handlers.rs | 15 +- src-tauri/src/proxy/mod.rs | 4 +- src-tauri/src/proxy/response_processor.rs | 13 +- src-tauri/src/proxy/session.rs | 269 ++++++++++++++++++++++ 5 files changed, 317 insertions(+), 13 deletions(-) diff --git a/src-tauri/src/proxy/handler_context.rs b/src-tauri/src/proxy/handler_context.rs index 724618a50..773515f3d 100644 --- a/src-tauri/src/proxy/handler_context.rs +++ b/src-tauri/src/proxy/handler_context.rs @@ -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, - /// 请求开始时的“当前供应商”（用于判断是否需要同步 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, }) } diff --git a/src-tauri/src/proxy/handlers.rs b/src-tauri/src/proxy/handlers.rs index 20512dc70..0e29aff74 100644 --- a/src-tauri/src/proxy/handlers.rs +++ b/src-tauri/src/proxy/handlers.rs @@ -61,7 +61,8 @@ pub async fn handle_messages( headers: axum::http::HeaderMap, Json(body): Json, ) -> Result { - 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 { 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, ) -> Result { - 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, ) -> Result { // 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}"); diff --git a/src-tauri/src/proxy/mod.rs b/src-tauri/src/proxy/mod.rs index d8e0e6deb..5b0ac0216 100644 --- a/src-tauri/src/proxy/mod.rs +++ b/src-tauri/src/proxy/mod.rs @@ -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}; diff --git a/src-tauri/src/proxy/response_processor.rs b/src-tauri/src/proxy/response_processor.rs index 76409054c..0df50e039 100644 --- a/src-tauri/src/proxy/response_processor.rs +++ b/src-tauri/src/proxy/response_processor.rs @@ -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, is_streaming: bool, status_code: u16, + session_id: Option, ) { 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, ) { diff --git a/src-tauri/src/proxy/session.rs b/src-tauri/src/proxy/session.rs index 3e1d14e56..cf3f964fe 100644 --- a/src-tauri/src/proxy/session.rs +++ b/src-tauri/src/proxy/session.rs @@ -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 { + // 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 { + 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 { + // 查找 "_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); + } }