mirror of
https://github.com/farion1231/cc-switch.git
synced 2026-06-16 13:34:04 +08:00
feat(proxy): add ProviderType enum for fine-grained provider detection
Introduce ProviderType enum to distinguish between different provider implementations (Claude, ClaudeAuth, Codex, Gemini, GeminiCli, OpenRouter). This enables proper authentication handling and request transformation based on the actual provider type rather than just AppType. - Add ProviderType enum with detection logic from config - Enhance Claude adapter with OpenRouter detection - Enhance Gemini adapter with CLI mode detection - Add helper methods for provider type inference
This commit is contained in:
@@ -11,12 +11,27 @@ pub struct AuthInfo {
|
||||
pub api_key: String,
|
||||
/// 认证策略
|
||||
pub strategy: AuthStrategy,
|
||||
/// OAuth access_token(用于 GoogleOAuth 策略)
|
||||
pub access_token: Option<String>,
|
||||
}
|
||||
|
||||
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<String> {
|
||||
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 <api_key>`
|
||||
///
|
||||
/// 用于不支持 x-api-key 的中转服务
|
||||
ClaudeAuth,
|
||||
|
||||
/// Bearer Token 认证方式(OpenAI 等)
|
||||
///
|
||||
/// - Header: `Authorization: Bearer <api_key>`
|
||||
Bearer,
|
||||
|
||||
/// Google 认证方式
|
||||
/// Google API Key 认证方式
|
||||
///
|
||||
/// - Header: `x-goog-api-key: <api_key>`
|
||||
Google,
|
||||
|
||||
/// Google OAuth 认证方式
|
||||
///
|
||||
/// - Header: `Authorization: Bearer <access_token>`
|
||||
///
|
||||
/// 用于 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);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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<String> {
|
||||
if let Some(env) = provider.settings_config.get("env") {
|
||||
@@ -122,11 +174,11 @@ impl ProviderAdapter for ClaudeAdapter {
|
||||
}
|
||||
|
||||
fn extract_auth(&self, provider: &Provider) -> Option<AuthInfo> {
|
||||
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();
|
||||
|
||||
@@ -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<Regex> =
|
||||
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<String> {
|
||||
// 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"
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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<String>,
|
||||
pub client_id: Option<String>,
|
||||
pub client_secret: Option<String>,
|
||||
}
|
||||
|
||||
#[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<String> {
|
||||
/// 获取供应商类型
|
||||
///
|
||||
/// 根据 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<OAuthCredentials> {
|
||||
// 直接是 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::<serde_json::Value>(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<String> {
|
||||
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<AuthInfo> {
|
||||
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());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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<Self, Self::Err> {
|
||||
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<dyn ProviderAdapter> {
|
||||
match app_type {
|
||||
@@ -37,3 +174,251 @@ pub fn get_adapter(app_type: &AppType) -> Box<dyn ProviderAdapter> {
|
||||
AppType::Gemini => Box::new(GeminiAdapter::new()),
|
||||
}
|
||||
}
|
||||
|
||||
/// 根据 ProviderType 获取对应的适配器
|
||||
#[allow(dead_code)]
|
||||
pub fn get_adapter_for_provider_type(provider_type: &ProviderType) -> Box<dyn ProviderAdapter> {
|
||||
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::<ProviderType>().unwrap(),
|
||||
ProviderType::Claude
|
||||
);
|
||||
assert_eq!(
|
||||
"claude_auth".parse::<ProviderType>().unwrap(),
|
||||
ProviderType::ClaudeAuth
|
||||
);
|
||||
assert_eq!(
|
||||
"claude-auth".parse::<ProviderType>().unwrap(),
|
||||
ProviderType::ClaudeAuth
|
||||
);
|
||||
assert_eq!(
|
||||
"codex".parse::<ProviderType>().unwrap(),
|
||||
ProviderType::Codex
|
||||
);
|
||||
assert_eq!(
|
||||
"gemini".parse::<ProviderType>().unwrap(),
|
||||
ProviderType::Gemini
|
||||
);
|
||||
assert_eq!(
|
||||
"gemini_cli".parse::<ProviderType>().unwrap(),
|
||||
ProviderType::GeminiCli
|
||||
);
|
||||
assert_eq!(
|
||||
"gemini-cli".parse::<ProviderType>().unwrap(),
|
||||
ProviderType::GeminiCli
|
||||
);
|
||||
assert_eq!(
|
||||
"openrouter".parse::<ProviderType>().unwrap(),
|
||||
ProviderType::OpenRouter
|
||||
);
|
||||
assert!("invalid".parse::<ProviderType>().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");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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<Mutex<Vec<Value>>>,
|
||||
}
|
||||
|
||||
#[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<S>(
|
||||
&self,
|
||||
stream: S,
|
||||
) -> impl Stream<Item = Result<Bytes, std::io::Error>> + Send
|
||||
where
|
||||
S: Stream<Item = Result<Bytes, reqwest::Error>> + 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::<Value>(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<Value> {
|
||||
let guard = self.events.lock().await;
|
||||
guard.clone()
|
||||
}
|
||||
|
||||
/// 从收集的事件中提取 Token 使用量
|
||||
pub async fn extract_usage(&self, session: &ProxySession) -> Option<TokenUsage> {
|
||||
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<Option<TokenUsage>, 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));
|
||||
}
|
||||
}
|
||||
@@ -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<String>,
|
||||
/// 客户端请求格式
|
||||
pub client_format: ClientFormat,
|
||||
/// 选定的供应商 ID
|
||||
pub provider_id: Option<String>,
|
||||
/// 模型名称
|
||||
pub model: Option<String>,
|
||||
/// 是否为流式请求
|
||||
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");
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user