diff --git a/src-tauri/src/proxy/providers/auth.rs b/src-tauri/src/proxy/providers/auth.rs index 9c293e40f..7e161fbf2 100644 --- a/src-tauri/src/proxy/providers/auth.rs +++ b/src-tauri/src/proxy/providers/auth.rs @@ -11,12 +11,27 @@ pub struct AuthInfo { pub api_key: String, /// 认证策略 pub strategy: AuthStrategy, + /// OAuth access_token（用于 GoogleOAuth 策略） + pub access_token: Option, } impl AuthInfo { /// 创建新的认证信息 pub fn new(api_key: String, strategy: AuthStrategy) -> Self { - Self { api_key, strategy } + Self { + api_key, + strategy, + access_token: None, + } + } + + /// 创建带有 access_token 的认证信息（用于 OAuth） + pub fn with_access_token(api_key: String, access_token: String) -> Self { + Self { + api_key, + strategy: AuthStrategy::GoogleOAuth, + access_token: Some(access_token), + } } /// 返回遮蔽后的 API Key（用于日志输出） @@ -34,6 +49,18 @@ impl AuthInfo { "***".to_string() } } + + /// 返回遮蔽后的 access_token（用于日志输出） + #[allow(dead_code)] + pub fn masked_access_token(&self) -> Option { + self.access_token.as_ref().map(|token| { + if token.len() > 8 { + format!("{}...{}", &token[..4], &token[token.len() - 4..]) + } else { + "***".to_string() + } + }) + } } /// 认证策略 @@ -46,13 +73,29 @@ pub enum AuthStrategy { /// - Header: `anthropic-version: 2023-06-01` Anthropic, + /// Claude 中转服务认证方式（仅 Bearer，无 x-api-key） + /// + /// - Header: `Authorization: Bearer ` + /// + /// 用于不支持 x-api-key 的中转服务 + ClaudeAuth, + /// Bearer Token 认证方式（OpenAI 等） + /// /// - Header: `Authorization: Bearer ` Bearer, - /// Google 认证方式 + /// Google API Key 认证方式 + /// /// - Header: `x-goog-api-key: ` Google, + + /// Google OAuth 认证方式 + /// + /// - Header: `Authorization: Bearer ` + /// + /// 用于 Gemini CLI 等需要 OAuth 的场景 + GoogleOAuth, } #[cfg(test)] @@ -89,4 +132,79 @@ mod tests { assert_ne!(AuthStrategy::Anthropic, AuthStrategy::Bearer); assert_ne!(AuthStrategy::Bearer, AuthStrategy::Google); } + + #[test] + fn test_auth_info_new_has_no_access_token() { + let auth = AuthInfo::new("api-key".to_string(), AuthStrategy::Bearer); + assert!(auth.access_token.is_none()); + } + + #[test] + fn test_auth_info_with_access_token() { + let auth = AuthInfo::with_access_token( + "refresh-token".to_string(), + "ya29.access-token-12345".to_string(), + ); + assert_eq!(auth.api_key, "refresh-token"); + assert_eq!(auth.strategy, AuthStrategy::GoogleOAuth); + assert_eq!( + auth.access_token, + Some("ya29.access-token-12345".to_string()) + ); + } + + #[test] + fn test_masked_access_token_long() { + let auth = + AuthInfo::with_access_token("refresh".to_string(), "ya29.1234567890abcdef".to_string()); + assert_eq!(auth.masked_access_token(), Some("ya29...cdef".to_string())); + } + + #[test] + fn test_masked_access_token_short() { + let auth = AuthInfo::with_access_token("refresh".to_string(), "short".to_string()); + assert_eq!(auth.masked_access_token(), Some("***".to_string())); + } + + #[test] + fn test_masked_access_token_none() { + let auth = AuthInfo::new("api-key".to_string(), AuthStrategy::Bearer); + assert!(auth.masked_access_token().is_none()); + } + + #[test] + fn test_claude_auth_strategy() { + let auth = AuthInfo::new("sk-test".to_string(), AuthStrategy::ClaudeAuth); + assert_eq!(auth.strategy, AuthStrategy::ClaudeAuth); + assert_ne!(auth.strategy, AuthStrategy::Anthropic); + assert_ne!(auth.strategy, AuthStrategy::Bearer); + } + + #[test] + fn test_google_oauth_strategy() { + let auth = AuthInfo::new("refresh-token".to_string(), AuthStrategy::GoogleOAuth); + assert_eq!(auth.strategy, AuthStrategy::GoogleOAuth); + assert_ne!(auth.strategy, AuthStrategy::Google); + } + + #[test] + fn test_all_strategies_are_distinct() { + let strategies = [ + AuthStrategy::Anthropic, + AuthStrategy::ClaudeAuth, + AuthStrategy::Bearer, + AuthStrategy::Google, + AuthStrategy::GoogleOAuth, + ]; + + for (i, s1) in strategies.iter().enumerate() { + for (j, s2) in strategies.iter().enumerate() { + if i == j { + assert_eq!(s1, s2); + } else { + assert_ne!(s1, s2); + } + } + } + } } diff --git a/src-tauri/src/proxy/providers/claude.rs b/src-tauri/src/proxy/providers/claude.rs index f01895917..f54321ca5 100644 --- a/src-tauri/src/proxy/providers/claude.rs +++ b/src-tauri/src/proxy/providers/claude.rs @@ -1,8 +1,13 @@ //! Claude (Anthropic) Provider Adapter //! //! 支持透传模式和 OpenRouter 兼容模式 +//! +//! ## 认证模式 +//! - **Claude**: Anthropic 官方 API (x-api-key + anthropic-version) +//! - **ClaudeAuth**: 中转服务 (仅 Bearer 认证，无 x-api-key) +//! - **OpenRouter**: 需要 Anthropic ↔ OpenAI 格式转换 -use super::{AuthInfo, AuthStrategy, ProviderAdapter}; +use super::{AuthInfo, AuthStrategy, ProviderAdapter, ProviderType}; use crate::provider::Provider; use crate::proxy::error::ProxyError; use reqwest::RequestBuilder; @@ -15,6 +20,28 @@ impl ClaudeAdapter { Self } + /// 获取供应商类型 + /// + /// 根据 base_url 和 auth_mode 检测具体的供应商类型： + /// - OpenRouter: base_url 包含 openrouter.ai + /// - ClaudeAuth: auth_mode 为 bearer_only + /// - Claude: 默认 Anthropic 官方 + pub fn provider_type(&self, provider: &Provider) -> ProviderType { + // 检测 OpenRouter + if let Ok(base_url) = self.extract_base_url(provider) { + if base_url.contains("openrouter.ai") { + return ProviderType::OpenRouter; + } + } + + // 检测 ClaudeAuth (仅 Bearer 认证) + if self.is_bearer_only_mode(provider) { + return ProviderType::ClaudeAuth; + } + + ProviderType::Claude + } + /// 检测是否使用 OpenRouter fn is_openrouter(&self, provider: &Provider) -> bool { if let Ok(base_url) = self.extract_base_url(provider) { @@ -23,6 +50,31 @@ impl ClaudeAdapter { false } + /// 检测是否为仅 Bearer 认证模式 + fn is_bearer_only_mode(&self, provider: &Provider) -> bool { + // 检查 settings_config 中的 auth_mode + if let Some(auth_mode) = provider + .settings_config + .get("auth_mode") + .and_then(|v| v.as_str()) + { + if auth_mode == "bearer_only" { + return true; + } + } + + // 检查 env 中的 AUTH_MODE + if let Some(env) = provider.settings_config.get("env") { + if let Some(auth_mode) = env.get("AUTH_MODE").and_then(|v| v.as_str()) { + if auth_mode == "bearer_only" { + return true; + } + } + } + + false + } + /// 从 Provider 配置中提取 API Key fn extract_key(&self, provider: &Provider) -> Option { if let Some(env) = provider.settings_config.get("env") { @@ -122,11 +174,11 @@ impl ProviderAdapter for ClaudeAdapter { } fn extract_auth(&self, provider: &Provider) -> Option { - let is_openrouter = self.is_openrouter(provider); - let strategy = if is_openrouter { - AuthStrategy::Bearer - } else { - AuthStrategy::Anthropic + let provider_type = self.provider_type(provider); + let strategy = match provider_type { + ProviderType::OpenRouter => AuthStrategy::Bearer, + ProviderType::ClaudeAuth => AuthStrategy::ClaudeAuth, + _ => AuthStrategy::Anthropic, }; self.extract_key(provider) @@ -149,9 +201,16 @@ impl ProviderAdapter for ClaudeAdapter { fn add_auth_headers(&self, request: RequestBuilder, auth: &AuthInfo) -> RequestBuilder { match auth.strategy { + // Anthropic 官方: Authorization Bearer + x-api-key + anthropic-version AuthStrategy::Anthropic => request + .header("Authorization", format!("Bearer {}", auth.api_key)) .header("x-api-key", &auth.api_key) .header("anthropic-version", "2023-06-01"), + // ClaudeAuth 中转服务: 仅 Bearer，无 x-api-key + AuthStrategy::ClaudeAuth => { + request.header("Authorization", format!("Bearer {}", auth.api_key)) + } + // OpenRouter: Bearer AuthStrategy::Bearer => { request.header("Authorization", format!("Bearer {}", auth.api_key)) } @@ -241,6 +300,74 @@ mod tests { assert_eq!(auth.strategy, AuthStrategy::Bearer); } + #[test] + fn test_extract_auth_claude_auth_mode() { + let adapter = ClaudeAdapter::new(); + let provider = create_provider(json!({ + "env": { + "ANTHROPIC_BASE_URL": "https://some-proxy.com", + "ANTHROPIC_AUTH_TOKEN": "sk-proxy-key" + }, + "auth_mode": "bearer_only" + })); + + let auth = adapter.extract_auth(&provider).unwrap(); + assert_eq!(auth.api_key, "sk-proxy-key"); + assert_eq!(auth.strategy, AuthStrategy::ClaudeAuth); + } + + #[test] + fn test_extract_auth_claude_auth_env_mode() { + let adapter = ClaudeAdapter::new(); + let provider = create_provider(json!({ + "env": { + "ANTHROPIC_BASE_URL": "https://some-proxy.com", + "ANTHROPIC_AUTH_TOKEN": "sk-proxy-key", + "AUTH_MODE": "bearer_only" + } + })); + + let auth = adapter.extract_auth(&provider).unwrap(); + assert_eq!(auth.api_key, "sk-proxy-key"); + assert_eq!(auth.strategy, AuthStrategy::ClaudeAuth); + } + + #[test] + fn test_provider_type_detection() { + let adapter = ClaudeAdapter::new(); + + // Anthropic 官方 + let anthropic = create_provider(json!({ + "env": { + "ANTHROPIC_BASE_URL": "https://api.anthropic.com", + "ANTHROPIC_AUTH_TOKEN": "sk-ant-test" + } + })); + assert_eq!(adapter.provider_type(&anthropic), ProviderType::Claude); + + // OpenRouter + let openrouter = create_provider(json!({ + "env": { + "ANTHROPIC_BASE_URL": "https://openrouter.ai/api", + "OPENROUTER_API_KEY": "sk-or-test" + } + })); + assert_eq!(adapter.provider_type(&openrouter), ProviderType::OpenRouter); + + // ClaudeAuth + let claude_auth = create_provider(json!({ + "env": { + "ANTHROPIC_BASE_URL": "https://some-proxy.com", + "ANTHROPIC_AUTH_TOKEN": "sk-test" + }, + "auth_mode": "bearer_only" + })); + assert_eq!( + adapter.provider_type(&claude_auth), + ProviderType::ClaudeAuth + ); + } + #[test] fn test_build_url_anthropic() { let adapter = ClaudeAdapter::new(); diff --git a/src-tauri/src/proxy/providers/codex.rs b/src-tauri/src/proxy/providers/codex.rs index 1ece7e2fb..a7599b526 100644 --- a/src-tauri/src/proxy/providers/codex.rs +++ b/src-tauri/src/proxy/providers/codex.rs @@ -1,11 +1,21 @@ //! Codex (OpenAI) Provider Adapter //! //! 仅透传模式，支持直连 OpenAI API +//! +//! ## 客户端检测 +//! 支持检测官方 Codex 客户端 (codex_vscode, codex_cli_rs) use super::{AuthInfo, AuthStrategy, ProviderAdapter}; use crate::provider::Provider; use crate::proxy::error::ProxyError; +use regex::Regex; use reqwest::RequestBuilder; +use std::sync::LazyLock; + +/// 官方 Codex 客户端 User-Agent 正则 +#[allow(dead_code)] +static CODEX_CLIENT_REGEX: LazyLock = + LazyLock::new(|| Regex::new(r"^(codex_vscode|codex_cli_rs)/[\d.]+").unwrap()); /// Codex 适配器 pub struct CodexAdapter; @@ -15,6 +25,14 @@ impl CodexAdapter { Self } + /// 检测是否为官方 Codex 客户端 + /// + /// 匹配 User-Agent 模式: `^(codex_vscode|codex_cli_rs)/[\d.]+` + #[allow(dead_code)] + pub fn is_official_client(user_agent: &str) -> bool { + CODEX_CLIENT_REGEX.is_match(user_agent) + } + /// 从 Provider 配置中提取 API Key fn extract_key(&self, provider: &Provider) -> Option { // 1. 尝试从 env 中获取 @@ -212,4 +230,36 @@ mod tests { let url = adapter.build_url("https://www.packyapi.com/v1", "/v1/responses"); assert_eq!(url, "https://www.packyapi.com/v1/responses"); } + + // 官方客户端检测测试 + #[test] + fn test_is_official_client_vscode() { + assert!(CodexAdapter::is_official_client("codex_vscode/1.0.0")); + assert!(CodexAdapter::is_official_client("codex_vscode/2.3.4")); + assert!(CodexAdapter::is_official_client("codex_vscode/0.1")); + } + + #[test] + fn test_is_official_client_cli() { + assert!(CodexAdapter::is_official_client("codex_cli_rs/1.0.0")); + assert!(CodexAdapter::is_official_client("codex_cli_rs/0.5.2")); + } + + #[test] + fn test_is_not_official_client() { + assert!(!CodexAdapter::is_official_client("Mozilla/5.0")); + assert!(!CodexAdapter::is_official_client("curl/7.68.0")); + assert!(!CodexAdapter::is_official_client("python-requests/2.25.1")); + assert!(!CodexAdapter::is_official_client("codex_other/1.0.0")); + assert!(!CodexAdapter::is_official_client("")); + } + + #[test] + fn test_is_official_client_partial_match() { + // 必须从开头匹配 + assert!(!CodexAdapter::is_official_client("some codex_vscode/1.0.0")); + assert!(!CodexAdapter::is_official_client( + "prefix_codex_cli_rs/1.0.0" + )); + } } diff --git a/src-tauri/src/proxy/providers/gemini.rs b/src-tauri/src/proxy/providers/gemini.rs index 3c3dcf43c..b148f5a48 100644 --- a/src-tauri/src/proxy/providers/gemini.rs +++ b/src-tauri/src/proxy/providers/gemini.rs @@ -1,8 +1,12 @@ //! Gemini (Google) Provider Adapter //! -//! 仅透传模式，支持直连 Google Gemini API +//! 支持 API Key 和 OAuth 两种认证方式 +//! +//! ## 认证模式 +//! - **Gemini**: API Key 认证 (x-goog-api-key) +//! - **GeminiCli**: OAuth Bearer 认证 (用于 Gemini CLI) -use super::{AuthInfo, AuthStrategy, ProviderAdapter}; +use super::{AuthInfo, AuthStrategy, ProviderAdapter, ProviderType}; use crate::provider::Provider; use crate::proxy::error::ProxyError; use reqwest::RequestBuilder; @@ -10,13 +14,111 @@ use reqwest::RequestBuilder; /// Gemini 适配器 pub struct GeminiAdapter; +/// OAuth 凭证结构 +#[derive(Debug, Clone)] +#[allow(dead_code)] +pub struct OAuthCredentials { + pub access_token: String, + pub refresh_token: Option, + pub client_id: Option, + pub client_secret: Option, +} + +#[allow(dead_code)] +impl OAuthCredentials { + /// 检查是否需要刷新 token（有 refresh_token 但没有有效的 access_token） + pub fn needs_refresh(&self) -> bool { + self.refresh_token.is_some() && self.access_token.is_empty() + } + + /// 检查是否可以刷新 token + pub fn can_refresh(&self) -> bool { + self.refresh_token.is_some() && self.client_id.is_some() && self.client_secret.is_some() + } +} + impl GeminiAdapter { pub fn new() -> Self { Self } - /// 从 Provider 配置中提取 API Key - fn extract_key(&self, provider: &Provider) -> Option { + /// 获取供应商类型 + /// + /// 根据 API Key 格式检测： + /// - GeminiCli: access_token (ya29. 开头) 或 JSON 格式凭证 + /// - Gemini: 普通 API Key + pub fn provider_type(&self, provider: &Provider) -> ProviderType { + if let Some(key) = self.extract_key_raw(provider) { + // OAuth access_token 以 ya29. 开头 + if key.starts_with("ya29.") { + return ProviderType::GeminiCli; + } + // JSON 格式的 OAuth 凭证 + if key.starts_with('{') { + return ProviderType::GeminiCli; + } + } + ProviderType::Gemini + } + + /// 检测认证类型 + pub fn detect_auth_type(&self, provider: &Provider) -> AuthStrategy { + match self.provider_type(provider) { + ProviderType::GeminiCli => AuthStrategy::GoogleOAuth, + _ => AuthStrategy::Google, + } + } + + /// 解析 OAuth 凭证 + pub fn parse_oauth_credentials(&self, key: &str) -> Option { + // 直接是 access_token + if key.starts_with("ya29.") { + return Some(OAuthCredentials { + access_token: key.to_string(), + refresh_token: None, + client_id: None, + client_secret: None, + }); + } + + // JSON 格式 + if key.starts_with('{') { + if let Ok(json) = serde_json::from_str::(key) { + let access_token = json + .get("access_token") + .and_then(|v| v.as_str()) + .map(|s| s.to_string()) + .unwrap_or_default(); + let refresh_token = json + .get("refresh_token") + .and_then(|v| v.as_str()) + .map(|s| s.to_string()); + let client_id = json + .get("client_id") + .and_then(|v| v.as_str()) + .map(|s| s.to_string()); + let client_secret = json + .get("client_secret") + .and_then(|v| v.as_str()) + .map(|s| s.to_string()); + + // 如果有 access_token 或 refresh_token，返回凭证 + if !access_token.is_empty() || refresh_token.is_some() { + return Some(OAuthCredentials { + access_token, + refresh_token, + client_id, + client_secret, + }); + } + } + } + + None + } + + /// 从 Provider 配置中提取原始 API Key + fn extract_key_raw(&self, provider: &Provider) -> Option { if let Some(env) = provider.settings_config.get("env") { // 优先使用 GOOGLE_GEMINI_API_KEY if let Some(key) = env.get("GOOGLE_GEMINI_API_KEY").and_then(|v| v.as_str()) { @@ -84,8 +186,21 @@ impl ProviderAdapter for GeminiAdapter { } fn extract_auth(&self, provider: &Provider) -> Option { - self.extract_key(provider) - .map(|key| AuthInfo::new(key, AuthStrategy::Google)) + let key = self.extract_key_raw(provider)?; + let strategy = self.detect_auth_type(provider); + + match strategy { + AuthStrategy::GoogleOAuth => { + // 解析 OAuth 凭证 + if let Some(creds) = self.parse_oauth_credentials(&key) { + Some(AuthInfo::with_access_token(key, creds.access_token)) + } else { + // 回退到普通 API Key + Some(AuthInfo::new(key, AuthStrategy::Google)) + } + } + _ => Some(AuthInfo::new(key, AuthStrategy::Google)), + } } fn build_url(&self, base_url: &str, endpoint: &str) -> String { @@ -107,7 +222,17 @@ impl ProviderAdapter for GeminiAdapter { } fn add_auth_headers(&self, request: RequestBuilder, auth: &AuthInfo) -> RequestBuilder { - request.header("x-goog-api-key", &auth.api_key) + match auth.strategy { + // OAuth Bearer 认证 + AuthStrategy::GoogleOAuth => { + let token = auth.access_token.as_ref().unwrap_or(&auth.api_key); + request + .header("Authorization", format!("Bearer {token}")) + .header("x-goog-api-client", "GeminiCLI/1.0") + } + // API Key 认证 + _ => request.header("x-goog-api-key", &auth.api_key), + } } } @@ -147,7 +272,7 @@ mod tests { } #[test] - fn test_extract_auth() { + fn test_extract_auth_api_key() { let adapter = GeminiAdapter::new(); let provider = create_provider(json!({ "env": { @@ -158,6 +283,76 @@ mod tests { let auth = adapter.extract_auth(&provider).unwrap(); assert_eq!(auth.api_key, "AIza-test-key-12345678"); assert_eq!(auth.strategy, AuthStrategy::Google); + assert!(auth.access_token.is_none()); + } + + #[test] + fn test_extract_auth_oauth_access_token() { + let adapter = GeminiAdapter::new(); + let provider = create_provider(json!({ + "env": { + "GOOGLE_GEMINI_API_KEY": "ya29.test-access-token-12345" + } + })); + + let auth = adapter.extract_auth(&provider).unwrap(); + assert_eq!(auth.strategy, AuthStrategy::GoogleOAuth); + assert_eq!( + auth.access_token, + Some("ya29.test-access-token-12345".to_string()) + ); + } + + #[test] + fn test_extract_auth_oauth_json() { + let adapter = GeminiAdapter::new(); + let provider = create_provider(json!({ + "env": { + "GOOGLE_GEMINI_API_KEY": "{\"access_token\":\"ya29.test-token\",\"refresh_token\":\"1//refresh\"}" + } + })); + + let auth = adapter.extract_auth(&provider).unwrap(); + assert_eq!(auth.strategy, AuthStrategy::GoogleOAuth); + assert_eq!(auth.access_token, Some("ya29.test-token".to_string())); + } + + #[test] + fn test_provider_type_detection() { + let adapter = GeminiAdapter::new(); + + // API Key + let api_key_provider = create_provider(json!({ + "env": { + "GOOGLE_GEMINI_API_KEY": "AIza-test-key" + } + })); + assert_eq!( + adapter.provider_type(&api_key_provider), + ProviderType::Gemini + ); + + // OAuth access_token + let oauth_provider = create_provider(json!({ + "env": { + "GOOGLE_GEMINI_API_KEY": "ya29.test-token" + } + })); + assert_eq!( + adapter.provider_type(&oauth_provider), + ProviderType::GeminiCli + ); + + // OAuth JSON + let oauth_json_provider = create_provider(json!({ + "env": { + "GOOGLE_GEMINI_API_KEY": "{\"access_token\":\"ya29.test\"}" + } + })); + assert_eq!( + adapter.provider_type(&oauth_json_provider), + ProviderType::GeminiCli + ); } #[test] @@ -199,4 +394,33 @@ mod tests { "https://generativelanguage.googleapis.com/v1beta/models/gemini-pro:generateContent" ); } + + #[test] + fn test_parse_oauth_credentials_direct_token() { + let adapter = GeminiAdapter::new(); + let creds = adapter + .parse_oauth_credentials("ya29.test-access-token") + .unwrap(); + assert_eq!(creds.access_token, "ya29.test-access-token"); + assert!(creds.refresh_token.is_none()); + } + + #[test] + fn test_parse_oauth_credentials_json() { + let adapter = GeminiAdapter::new(); + let creds = adapter + .parse_oauth_credentials( + "{\"access_token\":\"ya29.test\",\"refresh_token\":\"1//refresh\"}", + ) + .unwrap(); + assert_eq!(creds.access_token, "ya29.test"); + assert_eq!(creds.refresh_token, Some("1//refresh".to_string())); + } + + #[test] + fn test_parse_oauth_credentials_invalid() { + let adapter = GeminiAdapter::new(); + assert!(adapter.parse_oauth_credentials("AIza-api-key").is_none()); + assert!(adapter.parse_oauth_credentials("invalid-json{").is_none()); + } } diff --git a/src-tauri/src/proxy/providers/mod.rs b/src-tauri/src/proxy/providers/mod.rs index 4b9720e33..f41a03b88 100644 --- a/src-tauri/src/proxy/providers/mod.rs +++ b/src-tauri/src/proxy/providers/mod.rs @@ -21,6 +21,8 @@ pub mod streaming; pub mod transform; use crate::app_config::AppType; +use crate::provider::Provider; +use serde::{Deserialize, Serialize}; // 公开导出 pub use adapter::ProviderAdapter; @@ -29,6 +31,141 @@ pub use claude::ClaudeAdapter; pub use codex::CodexAdapter; pub use gemini::GeminiAdapter; +/// 供应商类型枚举 +/// +/// 区分不同供应商的具体实现方式，决定认证和请求处理逻辑。 +/// 比 AppType 更细粒度，支持同一 AppType 下的多种变体。 +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ProviderType { + /// Anthropic 官方 API (x-api-key + anthropic-version) + Claude, + /// Claude 中转服务 (仅 Bearer 认证，无 x-api-key) + ClaudeAuth, + /// OpenAI Codex Response API + Codex, + /// Google Gemini API (x-goog-api-key) + Gemini, + /// Google Gemini CLI (OAuth Bearer) + GeminiCli, + /// OpenRouter (需要 Anthropic ↔ OpenAI 格式转换) + OpenRouter, +} + +impl ProviderType { + /// 是否需要格式转换 + /// + /// OpenRouter 需要将 Anthropic 格式转换为 OpenAI 格式 + #[allow(dead_code)] + pub fn needs_transform(&self) -> bool { + matches!(self, ProviderType::OpenRouter) + } + + /// 获取默认端点 + #[allow(dead_code)] + pub fn default_endpoint(&self) -> &'static str { + match self { + ProviderType::Claude | ProviderType::ClaudeAuth => "https://api.anthropic.com", + ProviderType::Codex => "https://api.openai.com", + ProviderType::Gemini | ProviderType::GeminiCli => { + "https://generativelanguage.googleapis.com" + } + ProviderType::OpenRouter => "https://openrouter.ai/api", + } + } + + /// 从 AppType 和 Provider 配置推断供应商类型 + /// + /// 根据配置中的 base_url、auth_mode、api_key 格式等信息推断具体的供应商类型 + #[allow(dead_code)] + pub fn from_app_type_and_config(app_type: &AppType, provider: &Provider) -> Self { + match app_type { + AppType::Claude => { + // 检测是否为 OpenRouter + let adapter = ClaudeAdapter::new(); + if let Ok(base_url) = adapter.extract_base_url(provider) { + if base_url.contains("openrouter.ai") { + return ProviderType::OpenRouter; + } + } + // 检测是否为中转服务（仅 Bearer 认证） + // 注意：ProviderMeta 没有直接的 auth_mode 字段， + // 我们通过检查 settings_config 中的配置来判断 + // 检查 settings_config 中的 auth_mode + if let Some(auth_mode) = provider + .settings_config + .get("auth_mode") + .and_then(|v| v.as_str()) + { + if auth_mode == "bearer_only" { + return ProviderType::ClaudeAuth; + } + } + // 检查 env 中的 auth_mode + if let Some(env) = provider.settings_config.get("env") { + if let Some(auth_mode) = env.get("AUTH_MODE").and_then(|v| v.as_str()) { + if auth_mode == "bearer_only" { + return ProviderType::ClaudeAuth; + } + } + } + ProviderType::Claude + } + AppType::Codex => ProviderType::Codex, + AppType::Gemini => { + // 检测是否为 CLI 模式（OAuth） + let adapter = GeminiAdapter::new(); + if let Some(auth) = adapter.extract_auth(provider) { + let key = &auth.api_key; + // OAuth access_token 以 ya29. 开头 + if key.starts_with("ya29.") { + return ProviderType::GeminiCli; + } + // JSON 格式的 OAuth 凭证 + if key.starts_with('{') { + return ProviderType::GeminiCli; + } + } + ProviderType::Gemini + } + } + } + + /// 转换为字符串表示 + pub fn as_str(&self) -> &'static str { + match self { + ProviderType::Claude => "claude", + ProviderType::ClaudeAuth => "claude_auth", + ProviderType::Codex => "codex", + ProviderType::Gemini => "gemini", + ProviderType::GeminiCli => "gemini_cli", + ProviderType::OpenRouter => "openrouter", + } + } +} + +impl std::fmt::Display for ProviderType { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.as_str()) + } +} + +impl std::str::FromStr for ProviderType { + type Err = String; + + fn from_str(s: &str) -> Result { + match s.to_lowercase().as_str() { + "claude" => Ok(ProviderType::Claude), + "claude_auth" | "claude-auth" => Ok(ProviderType::ClaudeAuth), + "codex" => Ok(ProviderType::Codex), + "gemini" => Ok(ProviderType::Gemini), + "gemini_cli" | "gemini-cli" => Ok(ProviderType::GeminiCli), + "openrouter" => Ok(ProviderType::OpenRouter), + _ => Err(format!("Invalid provider type: {s}")), + } + } +} + /// 根据 AppType 获取对应的适配器 pub fn get_adapter(app_type: &AppType) -> Box { match app_type { @@ -37,3 +174,251 @@ pub fn get_adapter(app_type: &AppType) -> Box { AppType::Gemini => Box::new(GeminiAdapter::new()), } } + +/// 根据 ProviderType 获取对应的适配器 +#[allow(dead_code)] +pub fn get_adapter_for_provider_type(provider_type: &ProviderType) -> Box { + match provider_type { + ProviderType::Claude | ProviderType::ClaudeAuth | ProviderType::OpenRouter => { + Box::new(ClaudeAdapter::new()) + } + ProviderType::Codex => Box::new(CodexAdapter::new()), + ProviderType::Gemini | ProviderType::GeminiCli => Box::new(GeminiAdapter::new()), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + fn create_provider(config: serde_json::Value) -> Provider { + Provider { + id: "test".to_string(), + name: "Test Provider".to_string(), + settings_config: config, + website_url: None, + category: None, + created_at: None, + sort_index: None, + notes: None, + meta: None, + icon: None, + icon_color: None, + is_proxy_target: None, + } + } + + #[test] + fn test_provider_type_needs_transform() { + assert!(!ProviderType::Claude.needs_transform()); + assert!(!ProviderType::ClaudeAuth.needs_transform()); + assert!(!ProviderType::Codex.needs_transform()); + assert!(!ProviderType::Gemini.needs_transform()); + assert!(!ProviderType::GeminiCli.needs_transform()); + assert!(ProviderType::OpenRouter.needs_transform()); + } + + #[test] + fn test_provider_type_default_endpoint() { + assert_eq!( + ProviderType::Claude.default_endpoint(), + "https://api.anthropic.com" + ); + assert_eq!( + ProviderType::ClaudeAuth.default_endpoint(), + "https://api.anthropic.com" + ); + assert_eq!( + ProviderType::Codex.default_endpoint(), + "https://api.openai.com" + ); + assert_eq!( + ProviderType::Gemini.default_endpoint(), + "https://generativelanguage.googleapis.com" + ); + assert_eq!( + ProviderType::GeminiCli.default_endpoint(), + "https://generativelanguage.googleapis.com" + ); + assert_eq!( + ProviderType::OpenRouter.default_endpoint(), + "https://openrouter.ai/api" + ); + } + + #[test] + fn test_provider_type_from_str() { + assert_eq!( + "claude".parse::().unwrap(), + ProviderType::Claude + ); + assert_eq!( + "claude_auth".parse::().unwrap(), + ProviderType::ClaudeAuth + ); + assert_eq!( + "claude-auth".parse::().unwrap(), + ProviderType::ClaudeAuth + ); + assert_eq!( + "codex".parse::().unwrap(), + ProviderType::Codex + ); + assert_eq!( + "gemini".parse::().unwrap(), + ProviderType::Gemini + ); + assert_eq!( + "gemini_cli".parse::().unwrap(), + ProviderType::GeminiCli + ); + assert_eq!( + "gemini-cli".parse::().unwrap(), + ProviderType::GeminiCli + ); + assert_eq!( + "openrouter".parse::().unwrap(), + ProviderType::OpenRouter + ); + assert!("invalid".parse::().is_err()); + } + + #[test] + fn test_provider_type_as_str() { + assert_eq!(ProviderType::Claude.as_str(), "claude"); + assert_eq!(ProviderType::ClaudeAuth.as_str(), "claude_auth"); + assert_eq!(ProviderType::Codex.as_str(), "codex"); + assert_eq!(ProviderType::Gemini.as_str(), "gemini"); + assert_eq!(ProviderType::GeminiCli.as_str(), "gemini_cli"); + assert_eq!(ProviderType::OpenRouter.as_str(), "openrouter"); + } + + #[test] + fn test_provider_type_serde() { + // Test serialization + let claude = ProviderType::Claude; + let serialized = serde_json::to_string(&claude).unwrap(); + assert_eq!(serialized, "\"claude\""); + + let claude_auth = ProviderType::ClaudeAuth; + let serialized = serde_json::to_string(&claude_auth).unwrap(); + assert_eq!(serialized, "\"claude_auth\""); + + // Test deserialization + let deserialized: ProviderType = serde_json::from_str("\"claude\"").unwrap(); + assert_eq!(deserialized, ProviderType::Claude); + + let deserialized: ProviderType = serde_json::from_str("\"gemini_cli\"").unwrap(); + assert_eq!(deserialized, ProviderType::GeminiCli); + } + + #[test] + fn test_from_app_type_claude_direct() { + let provider = create_provider(json!({ + "env": { + "ANTHROPIC_BASE_URL": "https://api.anthropic.com", + "ANTHROPIC_AUTH_TOKEN": "sk-ant-test" + } + })); + + let provider_type = ProviderType::from_app_type_and_config(&AppType::Claude, &provider); + assert_eq!(provider_type, ProviderType::Claude); + } + + #[test] + fn test_from_app_type_claude_openrouter() { + let provider = create_provider(json!({ + "env": { + "ANTHROPIC_BASE_URL": "https://openrouter.ai/api", + "OPENROUTER_API_KEY": "sk-or-test" + } + })); + + let provider_type = ProviderType::from_app_type_and_config(&AppType::Claude, &provider); + assert_eq!(provider_type, ProviderType::OpenRouter); + } + + #[test] + fn test_from_app_type_claude_auth() { + let provider = create_provider(json!({ + "env": { + "ANTHROPIC_BASE_URL": "https://some-proxy.com", + "ANTHROPIC_AUTH_TOKEN": "sk-test" + }, + "auth_mode": "bearer_only" + })); + + let provider_type = ProviderType::from_app_type_and_config(&AppType::Claude, &provider); + assert_eq!(provider_type, ProviderType::ClaudeAuth); + } + + #[test] + fn test_from_app_type_codex() { + let provider = create_provider(json!({ + "env": { + "OPENAI_API_KEY": "sk-test" + } + })); + + let provider_type = ProviderType::from_app_type_and_config(&AppType::Codex, &provider); + assert_eq!(provider_type, ProviderType::Codex); + } + + #[test] + fn test_from_app_type_gemini_api_key() { + let provider = create_provider(json!({ + "env": { + "GOOGLE_GEMINI_API_KEY": "AIza-test-key" + } + })); + + let provider_type = ProviderType::from_app_type_and_config(&AppType::Gemini, &provider); + assert_eq!(provider_type, ProviderType::Gemini); + } + + #[test] + fn test_from_app_type_gemini_cli_oauth() { + let provider = create_provider(json!({ + "env": { + "GOOGLE_GEMINI_API_KEY": "ya29.test-access-token" + } + })); + + let provider_type = ProviderType::from_app_type_and_config(&AppType::Gemini, &provider); + assert_eq!(provider_type, ProviderType::GeminiCli); + } + + #[test] + fn test_from_app_type_gemini_cli_json() { + let provider = create_provider(json!({ + "env": { + "GOOGLE_GEMINI_API_KEY": "{\"access_token\":\"ya29.test\",\"refresh_token\":\"1//test\"}" + } + })); + + let provider_type = ProviderType::from_app_type_and_config(&AppType::Gemini, &provider); + assert_eq!(provider_type, ProviderType::GeminiCli); + } + + #[test] + fn test_get_adapter_for_provider_type() { + let adapter = get_adapter_for_provider_type(&ProviderType::Claude); + assert_eq!(adapter.name(), "Claude"); + + let adapter = get_adapter_for_provider_type(&ProviderType::ClaudeAuth); + assert_eq!(adapter.name(), "Claude"); + + let adapter = get_adapter_for_provider_type(&ProviderType::OpenRouter); + assert_eq!(adapter.name(), "Claude"); + + let adapter = get_adapter_for_provider_type(&ProviderType::Codex); + assert_eq!(adapter.name(), "Codex"); + + let adapter = get_adapter_for_provider_type(&ProviderType::Gemini); + assert_eq!(adapter.name(), "Gemini"); + + let adapter = get_adapter_for_provider_type(&ProviderType::GeminiCli); + assert_eq!(adapter.name(), "Gemini"); + } +} diff --git a/src-tauri/src/proxy/response_handler.rs b/src-tauri/src/proxy/response_handler.rs new file mode 100644 index 000000000..7045643a0 --- /dev/null +++ b/src-tauri/src/proxy/response_handler.rs @@ -0,0 +1,214 @@ +//! Response Handler - 统一响应处理 +//! +//! 提供流式和非流式响应的统一处理接口 + +use super::session::ProxySession; +use super::usage::parser::TokenUsage; +use super::ProxyError; +use bytes::Bytes; +use futures::stream::{Stream, StreamExt}; +use serde_json::Value; +use std::sync::Arc; +use std::time::{Duration, Instant}; +use tokio::sync::Mutex; +use tokio::time::timeout; + +/// 响应类型 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[allow(dead_code)] +pub enum ResponseType { + /// 流式响应 (SSE) + Stream, + /// 非流式响应 + NonStream, +} + +impl ResponseType { + /// 从 Content-Type 检测响应类型 + #[allow(dead_code)] + pub fn from_content_type(content_type: &str) -> Self { + if content_type.contains("text/event-stream") { + ResponseType::Stream + } else { + ResponseType::NonStream + } + } +} + +/// 流式响应处理器 +#[allow(dead_code)] +pub struct StreamHandler { + /// 空闲超时时间 + idle_timeout: Duration, + /// 收集的事件 + events: Arc>>, +} + +#[allow(dead_code)] +impl StreamHandler { + /// 创建新的流式处理器 + pub fn new(idle_timeout_secs: u64) -> Self { + Self { + idle_timeout: Duration::from_secs(idle_timeout_secs), + events: Arc::new(Mutex::new(Vec::new())), + } + } + + /// 处理流式响应，返回分流后的客户端流 + /// + /// 客户端流立即返回，内部流在后台收集事件 + pub fn handle_stream( + &self, + stream: S, + ) -> impl Stream> + Send + where + S: Stream> + Send + 'static, + { + let events = self.events.clone(); + let idle_timeout = self.idle_timeout; + + async_stream::stream! { + let mut _last_activity = Instant::now(); + let mut buffer = String::new(); + + tokio::pin!(stream); + + loop { + let chunk_result = timeout(idle_timeout, stream.next()).await; + + match chunk_result { + Ok(Some(Ok(bytes))) => { + _last_activity = Instant::now(); + + // 解析 SSE 事件 + let text = String::from_utf8_lossy(&bytes); + buffer.push_str(&text); + + // 提取完整事件 + while let Some(pos) = buffer.find("\n\n") { + let event_text = buffer[..pos].to_string(); + buffer = buffer[pos + 2..].to_string(); + + for line in event_text.lines() { + if let Some(data) = line.strip_prefix("data: ") { + if data.trim() != "[DONE]" { + if let Ok(json) = serde_json::from_str::(data) { + let mut guard = events.lock().await; + guard.push(json); + } + } + } + } + } + + yield Ok(bytes); + } + Ok(Some(Err(e))) => { + log::error!("流错误: {e}"); + yield Err(std::io::Error::other(e.to_string())); + break; + } + Ok(None) => { + // 流结束 + break; + } + Err(_) => { + // 空闲超时 + log::warn!("流式响应空闲超时: {idle_timeout:?} 无数据"); + yield Err(std::io::Error::other("Stream idle timeout")); + break; + } + } + } + } + } + + /// 获取收集的事件 + pub async fn get_events(&self) -> Vec { + let guard = self.events.lock().await; + guard.clone() + } + + /// 从收集的事件中提取 Token 使用量 + pub async fn extract_usage(&self, session: &ProxySession) -> Option { + let events = self.get_events().await; + + match session.client_format { + super::session::ClientFormat::Claude => TokenUsage::from_claude_stream_events(&events), + super::session::ClientFormat::Codex => TokenUsage::from_codex_stream_events(&events), + super::session::ClientFormat::Gemini | super::session::ClientFormat::GeminiCli => { + TokenUsage::from_gemini_stream_chunks(&events) + } + _ => None, + } + } +} + +/// 非流式响应处理器 +#[allow(dead_code)] +pub struct NonStreamHandler; + +#[allow(dead_code)] +impl NonStreamHandler { + /// 处理非流式响应 + /// + /// 克隆响应体用于后台解析，原始响应立即返回 + pub async fn handle_response( + body: &[u8], + session: &ProxySession, + ) -> Result, ProxyError> { + let json: Value = serde_json::from_slice(body) + .map_err(|e| ProxyError::TransformError(format!("Failed to parse response: {e}")))?; + + let usage = match session.client_format { + super::session::ClientFormat::Claude => TokenUsage::from_claude_response(&json), + super::session::ClientFormat::Codex => TokenUsage::from_codex_response_adjusted(&json), + super::session::ClientFormat::Gemini | super::session::ClientFormat::GeminiCli => { + TokenUsage::from_gemini_response(&json) + } + super::session::ClientFormat::OpenAI => TokenUsage::from_openrouter_response(&json), + _ => None, + }; + + Ok(usage) + } +} + +/// 统一响应分发器 +#[allow(dead_code)] +pub struct ResponseDispatcher; + +#[allow(dead_code)] +impl ResponseDispatcher { + /// 判断响应类型 + pub fn detect_type(content_type: &str) -> ResponseType { + ResponseType::from_content_type(content_type) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_response_type_detection() { + assert_eq!( + ResponseType::from_content_type("text/event-stream"), + ResponseType::Stream + ); + assert_eq!( + ResponseType::from_content_type("text/event-stream; charset=utf-8"), + ResponseType::Stream + ); + assert_eq!( + ResponseType::from_content_type("application/json"), + ResponseType::NonStream + ); + } + + #[test] + fn test_stream_handler_creation() { + let handler = StreamHandler::new(30); + assert_eq!(handler.idle_timeout, Duration::from_secs(30)); + } +} diff --git a/src-tauri/src/proxy/session.rs b/src-tauri/src/proxy/session.rs new file mode 100644 index 000000000..3e1d14e56 --- /dev/null +++ b/src-tauri/src/proxy/session.rs @@ -0,0 +1,298 @@ +//! Proxy Session - 请求会话管理 +//! +//! 为每个代理请求创建会话上下文，在整个请求生命周期中跟踪状态和元数据。 + +use std::time::Instant; +use uuid::Uuid; + +/// 客户端请求格式 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[allow(dead_code)] +pub enum ClientFormat { + /// Claude Messages API (/v1/messages) + Claude, + /// Codex Response API (/v1/responses) + Codex, + /// OpenAI Chat Completions API (/v1/chat/completions) + OpenAI, + /// Gemini API (/v1beta/models/*/generateContent) + Gemini, + /// Gemini CLI API (/v1internal/models/*/generateContent) + GeminiCli, + /// 未知格式 + Unknown, +} + +#[allow(dead_code)] +impl ClientFormat { + /// 从请求路径检测格式 + pub fn from_path(path: &str) -> Self { + if path.contains("/v1/messages") { + ClientFormat::Claude + } else if path.contains("/v1/responses") { + ClientFormat::Codex + } else if path.contains("/v1/chat/completions") { + ClientFormat::OpenAI + } else if path.contains("/v1internal/") && path.contains("generateContent") { + // Gemini CLI 使用 /v1internal/ 路径 + ClientFormat::GeminiCli + } else if (path.contains("/v1beta/") || path.contains("/v1/")) + && path.contains("generateContent") + { + // Gemini API 使用 /v1beta/ 或 /v1/ 路径 + ClientFormat::Gemini + } else if path.contains("generateContent") { + // 通用 Gemini 端点 + ClientFormat::Gemini + } else { + ClientFormat::Unknown + } + } + + /// 从请求体内容检测格式（回退方案） + pub fn from_body(body: &serde_json::Value) -> Self { + // Claude 格式特征: messages 数组 + model 字段 + 无 response_format + if body.get("messages").is_some() + && body.get("model").is_some() + && body.get("response_format").is_none() + && body.get("contents").is_none() + { + // 区分 Claude 和 OpenAI + if body.get("max_tokens").is_some() { + return ClientFormat::Claude; + } + return ClientFormat::OpenAI; + } + + // Codex 格式特征: input 字段 + if body.get("input").is_some() { + return ClientFormat::Codex; + } + + // Gemini 格式特征: contents 数组 + if body.get("contents").is_some() { + return ClientFormat::Gemini; + } + + ClientFormat::Unknown + } + + /// 转换为字符串 + pub fn as_str(&self) -> &'static str { + match self { + ClientFormat::Claude => "claude", + ClientFormat::Codex => "codex", + ClientFormat::OpenAI => "openai", + ClientFormat::Gemini => "gemini", + ClientFormat::GeminiCli => "gemini_cli", + ClientFormat::Unknown => "unknown", + } + } +} + +impl std::fmt::Display for ClientFormat { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.as_str()) + } +} + +/// 代理会话 +/// +/// 包含请求全生命周期的上下文数据 +#[derive(Debug, Clone)] +#[allow(dead_code)] +pub struct ProxySession { + /// 唯一会话 ID + pub session_id: String, + /// 请求开始时间 + pub start_time: Instant, + /// HTTP 方法 + pub method: String, + /// 请求 URL + pub request_url: String, + /// User-Agent + pub user_agent: Option, + /// 客户端请求格式 + pub client_format: ClientFormat, + /// 选定的供应商 ID + pub provider_id: Option, + /// 模型名称 + pub model: Option, + /// 是否为流式请求 + pub is_streaming: bool, +} + +#[allow(dead_code)] +impl ProxySession { + /// 从请求创建会话 + pub fn from_request( + method: &str, + request_url: &str, + user_agent: Option<&str>, + body: Option<&serde_json::Value>, + ) -> Self { + // 检测客户端格式 + let mut client_format = ClientFormat::from_path(request_url); + if client_format == ClientFormat::Unknown { + if let Some(body) = body { + client_format = ClientFormat::from_body(body); + } + } + + // 检测是否为流式请求 + let is_streaming = body + .and_then(|b| b.get("stream")) + .and_then(|v| v.as_bool()) + .unwrap_or(false); + + // 提取模型名称 + let model = body + .and_then(|b| b.get("model")) + .and_then(|v| v.as_str()) + .map(|s| s.to_string()); + + Self { + session_id: Uuid::new_v4().to_string(), + start_time: Instant::now(), + method: method.to_string(), + request_url: request_url.to_string(), + user_agent: user_agent.map(|s| s.to_string()), + client_format, + provider_id: None, + model, + is_streaming, + } + } + + /// 设置供应商 ID + pub fn with_provider(mut self, provider_id: &str) -> Self { + self.provider_id = Some(provider_id.to_string()); + self + } + + /// 获取请求延迟（毫秒） + pub fn latency_ms(&self) -> u64 { + self.start_time.elapsed().as_millis() as u64 + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn test_client_format_from_path_claude() { + assert_eq!( + ClientFormat::from_path("/v1/messages"), + ClientFormat::Claude + ); + assert_eq!( + ClientFormat::from_path("/api/v1/messages"), + ClientFormat::Claude + ); + } + + #[test] + fn test_client_format_from_path_codex() { + assert_eq!( + ClientFormat::from_path("/v1/responses"), + ClientFormat::Codex + ); + } + + #[test] + fn test_client_format_from_path_openai() { + assert_eq!( + ClientFormat::from_path("/v1/chat/completions"), + ClientFormat::OpenAI + ); + } + + #[test] + fn test_client_format_from_path_gemini() { + assert_eq!( + ClientFormat::from_path("/v1beta/models/gemini-pro:generateContent"), + ClientFormat::Gemini + ); + } + + #[test] + fn test_client_format_from_path_gemini_cli() { + assert_eq!( + ClientFormat::from_path("/v1internal/models/gemini-pro:generateContent"), + ClientFormat::GeminiCli + ); + } + + #[test] + fn test_client_format_from_body_claude() { + let body = json!({ + "model": "claude-3-5-sonnet", + "messages": [{"role": "user", "content": "Hello"}], + "max_tokens": 1024 + }); + assert_eq!(ClientFormat::from_body(&body), ClientFormat::Claude); + } + + #[test] + fn test_client_format_from_body_codex() { + let body = json!({ + "input": "Write a function" + }); + assert_eq!(ClientFormat::from_body(&body), ClientFormat::Codex); + } + + #[test] + fn test_client_format_from_body_gemini() { + let body = json!({ + "contents": [{"parts": [{"text": "Hello"}]}] + }); + assert_eq!(ClientFormat::from_body(&body), ClientFormat::Gemini); + } + + #[test] + fn test_session_id_uniqueness() { + let session1 = ProxySession::from_request("POST", "/v1/messages", None, None); + let session2 = ProxySession::from_request("POST", "/v1/messages", None, None); + assert_ne!(session1.session_id, session2.session_id); + } + + #[test] + fn test_session_from_request() { + let body = json!({ + "model": "claude-3-5-sonnet", + "messages": [{"role": "user", "content": "Hello"}], + "max_tokens": 1024, + "stream": true + }); + + let session = + ProxySession::from_request("POST", "/v1/messages", Some("Mozilla/5.0"), Some(&body)); + + assert_eq!(session.method, "POST"); + assert_eq!(session.request_url, "/v1/messages"); + assert_eq!(session.user_agent, Some("Mozilla/5.0".to_string())); + assert_eq!(session.client_format, ClientFormat::Claude); + assert_eq!(session.model, Some("claude-3-5-sonnet".to_string())); + assert!(session.is_streaming); + } + + #[test] + fn test_session_with_provider() { + let session = ProxySession::from_request("POST", "/v1/messages", None, None) + .with_provider("provider-123"); + + assert_eq!(session.provider_id, Some("provider-123".to_string())); + } + + #[test] + fn test_client_format_as_str() { + assert_eq!(ClientFormat::Claude.as_str(), "claude"); + assert_eq!(ClientFormat::Codex.as_str(), "codex"); + assert_eq!(ClientFormat::OpenAI.as_str(), "openai"); + assert_eq!(ClientFormat::Gemini.as_str(), "gemini"); + assert_eq!(ClientFormat::GeminiCli.as_str(), "gemini_cli"); + assert_eq!(ClientFormat::Unknown.as_str(), "unknown"); + } +}