diff --git a/README.md b/README.md index 6a30bc0e7..be083699c 100644 --- a/README.md +++ b/README.md @@ -16,7 +16,7 @@ ## 🤔 这个工具能帮你做什么？ **场景一：换个更好用的 IDE** -> 我有 Kiro 账号，每月免费用 Claude Sonnet 4.5，但 Kiro IDE 不太顺手。我想用 Claude Code 或 Cursor 来写代码，但又不想额外付费买 API。 +> 我有 Kiro 账号，可以用 Claude 系列模型，但 Kiro IDE 不太顺手。我想用 Claude Code 或 Cursor 来写代码，但又不想额外付费买 API。 **场景二：把额度分享给其他工具** > Claude Code 这个月额度还剩很多，与其浪费不如转给 Cherry Studio 聊天用，或者给我的 AI Agent 项目提供 API 接口。 @@ -24,7 +24,7 @@ **场景三：统一管理多个 AI 账号** > 我有 Kiro、Gemini CLI、通义千问好几个账号，想统一管理，哪个有额度就用哪个。 -**ProxyCast 就是解决这些问题的工具** —— 它把你已有的 AI 客户端凭证转换成标准 OpenAI API，让任何支持 OpenAI 接口的工具都能用上你的免费额度。 +**ProxyCast 就是解决这些问题的工具** —— 它把你已有的 AI 客户端凭证转换成标准 OpenAI API，让任何支持 OpenAI 接口的工具都能使用。 --- @@ -49,12 +49,11 @@ ## ✨ 核心特性 ### 🎯 多 Provider 统一管理 -- **Kiro Claude** - 通过 OAuth 免费使用 Claude Sonnet 4.5 -- **Gemini CLI** - 通过 OAuth 突破 Gemini 免费限制 +- **Kiro** - 通过 OAuth 使用 Claude 系列模型（Opus 4.5、Sonnet 4.5、Sonnet 4、Haiku 4.5） +- **Gemini CLI** - 通过 OAuth 使用 Gemini 模型 - **Gemini API Key** - 多账号负载均衡，支持模型排除 - **通义千问** - 通过 OAuth 使用 Qwen3 Coder Plus -- **OpenAI Codex** - 通过 OAuth 使用 GPT 模型 -- **iFlow** - 支持 OAuth 和 Cookie 两种认证方式 +- **Antigravity** - 通过 OAuth 使用 Claude 模型 - **Vertex AI** - Google Cloud AI 平台，支持模型别名 - **OpenAI 自定义** - 配置自定义 OpenAI 兼容 API - **Claude 自定义** - 配置自定义 Claude API @@ -79,9 +78,9 @@ - **远程管理 API** - 通过 API 远程管理配置和凭证 - **访问控制** - 支持 localhost 限制和密钥认证 -### 🔌 Amp CLI 集成 +### 🔌 多路由支持 - 支持 `/api/provider/{provider}/v1/*` 路由模式 -- 模型映射 - 将不可用模型映射到可用替代 +- 模型映射 - 将请求模型映射到 Provider 支持的模型 - 管理端点代理 - 代理认证和账户功能 ### 🌐 完整 API 兼容 @@ -91,7 +90,7 @@ - `/v1/messages/count_tokens` - Token 计数 - `/health` - 健康检查 - `/ready` - 就绪检查 -- `/api/provider/{provider}/v1/*` - Amp CLI 路由 +- `/api/provider/{provider}/v1/*` - Provider 路由 - `/v0/management/*` - 远程管理 API - `/v0/management/backup` - 触发数据库备份 - `/v0/management/restore` - 从备份恢复 @@ -100,26 +99,26 @@ ## 📸 界面截图 -### Dashboard - 服务控制与 API 测试 -![Dashboard](docs/images/420984ac-8287-44c6-b209-bbffd59dc0eb.png) +### 仪表盘 - 系统状态与监控 +![Dashboard](docs/images/943663ed-b17c-4b32-a74c-c0243ffb3dea.png) -### 凭证池 - 多凭证管理 -![Provider Pool](docs/images/44dfe29b-ef1e-4fab-82cc-8fbb3d5ae673.png) +### 凭证池 - 多凭证管理与配额查询 +![Provider Pool](docs/images/aee62eb5-3aeb-4454-b14d-24b1d5f9a0fe.png) -### API Server - 路由与日志 -![API Server](docs/images/54223543-bfb1-4a96-a9d8-a1b9b2f9c154.png) +### 路由管理 - 智能路由规则和容错策略 +![Router](docs/images/067c7d64-e116-4a30-b533-748873166f37.png) -### 设置页面 - 服务器配置 -![Settings](docs/images/83270842-209e-4a36-98e3-607642273806.png) +### 配置管理 - 客户端配置切换 +![Config](docs/images/25eb018a-5be2-4f82-ba22-e68f39160cac.png) -### AI Clients - 客户端配置 -![AI Clients](docs/images/c9d36815-a29a-437e-9f26-872ed0949921.png) +### 扩展 - MCP/Prompts/Skills 管理 +![Extensions](docs/images/ffc70018-aa5f-4738-883d-045614488608.png) -### MCP 服务器管理 -![MCP](docs/images/fc54f689-b092-4180-a531-a41e310c42cb.png) +### API Server - 服务控制与 API 测试 +![API Server](docs/images/151b4355-821c-4bda-a731-c4367b6b8716.png) -### Prompts 管理 -![Prompts](docs/images/ff6d0f51-a703-4642-a4ab-2b9512b332af.png) +### 设置 - 应用参数和偏好 +![Settings](docs/images/c7d8236b-ea6c-4496-ada5-288cd0a01738.png) --- @@ -131,16 +130,7 @@ - **macOS (Apple Silicon)**: `ProxyCast_x.x.x_aarch64.dmg` - **Windows (x64)**: `ProxyCast_x.x.x_x64-setup.exe` - -### 凭证文件位置 - -| Provider | 默认路径 | 说明 | -|----------|---------|------| -| **Kiro** | `~/.aws/sso/cache/kiro-auth-token.json` | Kiro OAuth Token | -| **Gemini** | `~/.gemini/oauth_creds.json` | Gemini CLI OAuth | -| **Qwen** | `~/.qwen/oauth_creds.json` | 通义千问 OAuth | - -> **提示**: `~` 表示用户主目录（macOS: `/Users/用户名`，Windows: `C:\Users\用户名`） +- **Ubuntu/Debian (x64)**: `ProxyCast_x.x.x_amd64.deb` ### 使用步骤 diff --git a/docs/images/067c7d64-e116-4a30-b533-748873166f37.png b/docs/images/067c7d64-e116-4a30-b533-748873166f37.png new file mode 100644 index 000000000..15e989ccd Binary files /dev/null and b/docs/images/067c7d64-e116-4a30-b533-748873166f37.png differ diff --git a/docs/images/151b4355-821c-4bda-a731-c4367b6b8716.png b/docs/images/151b4355-821c-4bda-a731-c4367b6b8716.png new file mode 100644 index 000000000..d38507ddc Binary files /dev/null and b/docs/images/151b4355-821c-4bda-a731-c4367b6b8716.png differ diff --git a/docs/images/20251218001105_248_33.jpg b/docs/images/20251218001105_248_33.jpg deleted file mode 100644 index 60456b1a1..000000000 Binary files a/docs/images/20251218001105_248_33.jpg and /dev/null differ diff --git a/docs/images/25eb018a-5be2-4f82-ba22-e68f39160cac.png b/docs/images/25eb018a-5be2-4f82-ba22-e68f39160cac.png new file mode 100644 index 000000000..8fe5b9470 Binary files /dev/null and b/docs/images/25eb018a-5be2-4f82-ba22-e68f39160cac.png differ diff --git a/docs/images/420984ac-8287-44c6-b209-bbffd59dc0eb.png b/docs/images/420984ac-8287-44c6-b209-bbffd59dc0eb.png deleted file mode 100644 index c94006981..000000000 Binary files a/docs/images/420984ac-8287-44c6-b209-bbffd59dc0eb.png and /dev/null differ diff --git a/docs/images/44dfe29b-ef1e-4fab-82cc-8fbb3d5ae673.png b/docs/images/44dfe29b-ef1e-4fab-82cc-8fbb3d5ae673.png deleted file mode 100644 index 552e65096..000000000 Binary files a/docs/images/44dfe29b-ef1e-4fab-82cc-8fbb3d5ae673.png and /dev/null differ diff --git a/docs/images/54223543-bfb1-4a96-a9d8-a1b9b2f9c154.png b/docs/images/54223543-bfb1-4a96-a9d8-a1b9b2f9c154.png deleted file mode 100644 index db7adf974..000000000 Binary files a/docs/images/54223543-bfb1-4a96-a9d8-a1b9b2f9c154.png and /dev/null differ diff --git a/docs/images/83270842-209e-4a36-98e3-607642273806.png b/docs/images/83270842-209e-4a36-98e3-607642273806.png deleted file mode 100644 index 382806474..000000000 Binary files a/docs/images/83270842-209e-4a36-98e3-607642273806.png and /dev/null differ diff --git a/docs/images/943663ed-b17c-4b32-a74c-c0243ffb3dea.png b/docs/images/943663ed-b17c-4b32-a74c-c0243ffb3dea.png new file mode 100644 index 000000000..71b933c8b Binary files /dev/null and b/docs/images/943663ed-b17c-4b32-a74c-c0243ffb3dea.png differ diff --git a/docs/images/aee62eb5-3aeb-4454-b14d-24b1d5f9a0fe.png b/docs/images/aee62eb5-3aeb-4454-b14d-24b1d5f9a0fe.png new file mode 100644 index 000000000..98fb49fd3 Binary files /dev/null and b/docs/images/aee62eb5-3aeb-4454-b14d-24b1d5f9a0fe.png differ diff --git a/docs/images/c7d8236b-ea6c-4496-ada5-288cd0a01738.png b/docs/images/c7d8236b-ea6c-4496-ada5-288cd0a01738.png new file mode 100644 index 000000000..20f6d97e5 Binary files /dev/null and b/docs/images/c7d8236b-ea6c-4496-ada5-288cd0a01738.png differ diff --git a/docs/images/c9d36815-a29a-437e-9f26-872ed0949921.png b/docs/images/c9d36815-a29a-437e-9f26-872ed0949921.png deleted file mode 100644 index a64bd531f..000000000 Binary files a/docs/images/c9d36815-a29a-437e-9f26-872ed0949921.png and /dev/null differ diff --git a/docs/images/fc54f689-b092-4180-a531-a41e310c42cb.png b/docs/images/fc54f689-b092-4180-a531-a41e310c42cb.png deleted file mode 100644 index ae395f58a..000000000 Binary files a/docs/images/fc54f689-b092-4180-a531-a41e310c42cb.png and /dev/null differ diff --git a/docs/images/ff6d0f51-a703-4642-a4ab-2b9512b332af.png b/docs/images/ff6d0f51-a703-4642-a4ab-2b9512b332af.png deleted file mode 100644 index 639a8b446..000000000 Binary files a/docs/images/ff6d0f51-a703-4642-a4ab-2b9512b332af.png and /dev/null differ diff --git a/docs/images/ffc70018-aa5f-4738-883d-045614488608.png b/docs/images/ffc70018-aa5f-4738-883d-045614488608.png new file mode 100644 index 000000000..2b4113451 Binary files /dev/null and b/docs/images/ffc70018-aa5f-4738-883d-045614488608.png differ diff --git a/docs/images/微信图片_20251218001131_249_33.jpg b/docs/images/微信图片_20251218001131_249_33.jpg deleted file mode 100644 index 6451c13d4..000000000 Binary files a/docs/images/微信图片_20251218001131_249_33.jpg and /dev/null differ diff --git a/package.json b/package.json index 645ad08ae..f82ee3b93 100644 --- a/package.json +++ b/package.json @@ -1,7 +1,7 @@ { "name": "proxycast", "private": true, - "version": "0.14.6", + "version": "0.15.0", "type": "module", "repository": { "type": "git", diff --git a/src-tauri/Cargo.lock b/src-tauri/Cargo.lock index ff94b028e..2413a19a3 100644 --- a/src-tauri/Cargo.lock +++ b/src-tauri/Cargo.lock @@ -3377,7 +3377,7 @@ dependencies = [ [[package]] name = "proxycast" -version = "0.14.6" +version = "0.15.1" dependencies = [ "anyhow", "async-stream", @@ -3394,6 +3394,7 @@ dependencies = [ "indexmap 2.12.1", "md5", "notify", + "once_cell", "open", "parking_lot", "proptest", diff --git a/src-tauri/Cargo.toml b/src-tauri/Cargo.toml index 3b3f59eb6..2ecc06af1 100644 --- a/src-tauri/Cargo.toml +++ b/src-tauri/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "proxycast" -version = "0.14.6" +version = "0.15.1" description = "AI API Proxy Desktop App" authors = ["you"] edition = "2021" @@ -58,6 +58,7 @@ sha2 = "0.10" serde_urlencoded = "0.7" open = "5" url = "2" +once_cell = "1" [dev-dependencies] proptest = "1" diff --git a/src-tauri/proptest-regressions/providers/tests.txt b/src-tauri/proptest-regressions/providers/tests.txt new file mode 100644 index 000000000..21ac60ee8 --- /dev/null +++ b/src-tauri/proptest-regressions/providers/tests.txt @@ -0,0 +1,7 @@ +# Seeds for failure cases proptest has generated in the past. It is +# automatically read and these particular cases re-run before any +# novel cases are generated. +# +# It is recommended to check this file in to source control so that +# everyone who runs the test benefits from these saved cases. +cc 10c52f014c7e9b4cd049a9802452d417b5774e6f42173c2c9329544d8ac4340c # shrinks to lead_time_mins = 21, time_offset_secs = 1260 diff --git a/src-tauri/proptest-regressions/proxy/tests.txt b/src-tauri/proptest-regressions/proxy/tests.txt new file mode 100644 index 000000000..75b28b39e --- /dev/null +++ b/src-tauri/proptest-regressions/proxy/tests.txt @@ -0,0 +1,7 @@ +# Seeds for failure cases proptest has generated in the past. It is +# automatically read and these particular cases re-run before any +# novel cases are generated. +# +# It is recommended to check this file in to source control so that +# everyone who runs the test benefits from these saved cases. +cc 99af5a6dad66f0b5a2650223417a2e83a7a3bfa8b6eab8ad57a88023367e739c # shrinks to url = "http://08:1024" diff --git a/src-tauri/src/commands/provider_pool_cmd.rs b/src-tauri/src/commands/provider_pool_cmd.rs index c12009ab6..5b3770400 100644 --- a/src-tauri/src/commands/provider_pool_cmd.rs +++ b/src-tauri/src/commands/provider_pool_cmd.rs @@ -12,7 +12,7 @@ use chrono::Utc; use std::fs; use std::path::{Path, PathBuf}; use std::sync::Arc; -use tauri::State; +use tauri::{Emitter, State}; use uuid::Uuid; pub struct ProviderPoolServiceState(pub Arc); @@ -1068,3 +1068,775 @@ pub struct MigrationResultResponse { /// 错误信息列表 pub errors: Vec, } + +/// 获取 Antigravity OAuth 授权 URL 并等待回调（不自动打开浏览器） +/// +/// 启动服务器后通过事件发送授权 URL，然后等待回调 +/// 成功后返回凭证 +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +pub struct AntigravityAuthUrlResponse { + pub auth_url: String, +} + +#[tauri::command] +pub async fn get_antigravity_auth_url_and_wait( + app: tauri::AppHandle, + db: State<'_, DbConnection>, + pool_service: State<'_, ProviderPoolServiceState>, + name: Option, + skip_project_id_fetch: Option, +) -> Result { + use crate::providers::antigravity; + + tracing::info!("[Antigravity OAuth] 启动服务器并获取授权 URL"); + + // 启动服务器并获取授权 URL + let (auth_url, wait_future) = + antigravity::start_oauth_server_and_get_url(skip_project_id_fetch.unwrap_or(false)) + .await + .map_err(|e| format!("启动 OAuth 服务器失败: {}", e))?; + + tracing::info!("[Antigravity OAuth] 授权 URL: {}", auth_url); + + // 通过事件发送授权 URL 给前端 + let _ = app.emit( + "antigravity-auth-url", + AntigravityAuthUrlResponse { + auth_url: auth_url.clone(), + }, + ); + + // 等待回调 + let result = wait_future.await.map_err(|e| e.to_string())?; + + tracing::info!( + "[Antigravity OAuth] 登录成功，凭证保存到: {}", + result.creds_file_path + ); + + // 从凭证中获取 project_id + let project_id = result.credentials.projectId.clone(); + + // 添加到凭证池 + let credential = pool_service.0.add_credential( + &db, + "antigravity", + CredentialData::AntigravityOAuth { + creds_file_path: result.creds_file_path, + project_id, + }, + name, + Some(true), + None, + )?; + + tracing::info!( + "[Antigravity OAuth] 凭证已添加到凭证池: {}", + credential.uuid + ); + + Ok(credential) +} + +/// 启动 Antigravity OAuth 登录流程 +/// +/// 打开浏览器让用户登录 Google 账号，获取 Antigravity 凭证 +#[tauri::command] +pub async fn start_antigravity_oauth_login( + db: State<'_, DbConnection>, + pool_service: State<'_, ProviderPoolServiceState>, + name: Option, + skip_project_id_fetch: Option, +) -> Result { + use crate::providers::antigravity; + + tracing::info!("[Antigravity OAuth] 开始 OAuth 登录流程"); + + // 启动 OAuth 登录 + let result = antigravity::start_oauth_login(skip_project_id_fetch.unwrap_or(false)) + .await + .map_err(|e| format!("Antigravity OAuth 登录失败: {}", e))?; + + tracing::info!( + "[Antigravity OAuth] 登录成功，凭证保存到: {}", + result.creds_file_path + ); + + // 从凭证中获取 project_id + let project_id = result.credentials.projectId.clone(); + + // 添加到凭证池 + let credential = pool_service.0.add_credential( + &db, + "antigravity", + CredentialData::AntigravityOAuth { + creds_file_path: result.creds_file_path, + project_id, + }, + name, + Some(true), + None, + )?; + + tracing::info!( + "[Antigravity OAuth] 凭证已添加到凭证池: {}", + credential.uuid + ); + + Ok(credential) +} + +/// Codex OAuth 授权 URL 响应 +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +pub struct CodexAuthUrlResponse { + pub auth_url: String, +} + +/// 获取 Codex OAuth 授权 URL 并等待回调（不自动打开浏览器） +/// +/// 启动服务器后通过事件发送授权 URL，然后等待回调 +/// 成功后返回凭证 +#[tauri::command] +pub async fn get_codex_auth_url_and_wait( + app: tauri::AppHandle, + db: State<'_, DbConnection>, + pool_service: State<'_, ProviderPoolServiceState>, + name: Option, +) -> Result { + use crate::providers::codex; + + tracing::info!("[Codex OAuth] 启动服务器并获取授权 URL"); + + // 启动服务器并获取授权 URL + let (auth_url, wait_future) = codex::start_codex_oauth_server_and_get_url() + .await + .map_err(|e| format!("启动 OAuth 服务器失败: {}", e))?; + + tracing::info!("[Codex OAuth] 授权 URL: {}", auth_url); + + // 通过事件发送授权 URL 给前端 + let _ = app.emit( + "codex-auth-url", + CodexAuthUrlResponse { + auth_url: auth_url.clone(), + }, + ); + + // 等待回调 + let result = wait_future.await.map_err(|e| e.to_string())?; + + tracing::info!( + "[Codex OAuth] 登录成功，凭证保存到: {}", + result.creds_file_path + ); + + // 添加到凭证池 + let credential = pool_service.0.add_credential( + &db, + "codex", + CredentialData::CodexOAuth { + creds_file_path: result.creds_file_path, + }, + name, + Some(true), + None, + )?; + + tracing::info!("[Codex OAuth] 凭证已添加到凭证池: {}", credential.uuid); + + Ok(credential) +} + +/// 启动 Codex OAuth 登录流程 +/// +/// 打开浏览器让用户登录 OpenAI 账号，获取 Codex 凭证 +#[tauri::command] +pub async fn start_codex_oauth_login( + db: State<'_, DbConnection>, + pool_service: State<'_, ProviderPoolServiceState>, + name: Option, +) -> Result { + use crate::providers::codex; + + tracing::info!("[Codex OAuth] 开始 OAuth 登录流程"); + + // 启动 OAuth 登录 + let result = codex::start_codex_oauth_login() + .await + .map_err(|e| format!("Codex OAuth 登录失败: {}", e))?; + + tracing::info!( + "[Codex OAuth] 登录成功，凭证保存到: {}", + result.creds_file_path + ); + + // 添加到凭证池 + let credential = pool_service.0.add_credential( + &db, + "codex", + CredentialData::CodexOAuth { + creds_file_path: result.creds_file_path, + }, + name, + Some(true), + None, + )?; + + tracing::info!("[Codex OAuth] 凭证已添加到凭证池: {}", credential.uuid); + + Ok(credential) +} + +/// Claude OAuth 授权 URL 响应 +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +pub struct ClaudeOAuthAuthUrlResponse { + pub auth_url: String, +} + +/// 获取 Claude OAuth 授权 URL 并等待回调（不自动打开浏览器） +/// +/// 启动服务器后通过事件发送授权 URL，然后等待回调 +/// 成功后返回凭证 +#[tauri::command] +pub async fn get_claude_oauth_auth_url_and_wait( + app: tauri::AppHandle, + db: State<'_, DbConnection>, + pool_service: State<'_, ProviderPoolServiceState>, + name: Option, +) -> Result { + use crate::providers::claude_oauth; + + tracing::info!("[Claude OAuth] 启动服务器并获取授权 URL"); + + // 启动服务器并获取授权 URL + let (auth_url, wait_future) = claude_oauth::start_claude_oauth_server_and_get_url() + .await + .map_err(|e| format!("启动 OAuth 服务器失败: {}", e))?; + + tracing::info!("[Claude OAuth] 授权 URL: {}", auth_url); + + // 通过事件发送授权 URL 给前端 + let _ = app.emit( + "claude-oauth-auth-url", + ClaudeOAuthAuthUrlResponse { + auth_url: auth_url.clone(), + }, + ); + + // 等待回调 + let result = wait_future.await.map_err(|e| e.to_string())?; + + tracing::info!( + "[Claude OAuth] 登录成功，凭证保存到: {}", + result.creds_file_path + ); + + // 添加到凭证池 + let credential = pool_service.0.add_credential( + &db, + "claude_oauth", + CredentialData::ClaudeOAuth { + creds_file_path: result.creds_file_path, + }, + name, + Some(true), + None, + )?; + + tracing::info!("[Claude OAuth] 凭证已添加到凭证池: {}", credential.uuid); + + Ok(credential) +} + +/// 启动 Claude OAuth 登录流程 +/// +/// 打开浏览器让用户登录 Claude 账号，获取凭证 +#[tauri::command] +pub async fn start_claude_oauth_login( + db: State<'_, DbConnection>, + pool_service: State<'_, ProviderPoolServiceState>, + name: Option, +) -> Result { + use crate::providers::claude_oauth; + + tracing::info!("[Claude OAuth] 开始 OAuth 登录流程"); + + // 启动 OAuth 登录 + let result = claude_oauth::start_claude_oauth_login() + .await + .map_err(|e| format!("Claude OAuth 登录失败: {}", e))?; + + tracing::info!( + "[Claude OAuth] 登录成功，凭证保存到: {}", + result.creds_file_path + ); + + // 添加到凭证池 + let credential = pool_service.0.add_credential( + &db, + "claude_oauth", + CredentialData::ClaudeOAuth { + creds_file_path: result.creds_file_path, + }, + name, + Some(true), + None, + )?; + + tracing::info!("[Claude OAuth] 凭证已添加到凭证池: {}", credential.uuid); + + Ok(credential) +} + +/// Qwen Device Code 响应 +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +pub struct QwenDeviceCodeResponse { + pub user_code: String, + pub verification_uri: String, + pub verification_uri_complete: Option, + pub expires_in: i64, +} + +/// 获取 Qwen Device Code 并等待用户授权 +/// +/// 启动 Device Code Flow 后通过事件发送设备码信息，然后轮询等待授权 +/// 成功后返回凭证 +#[tauri::command] +pub async fn get_qwen_device_code_and_wait( + app: tauri::AppHandle, + db: State<'_, DbConnection>, + pool_service: State<'_, ProviderPoolServiceState>, + name: Option, +) -> Result { + use crate::providers::qwen; + + tracing::info!("[Qwen] 启动 Device Code Flow"); + + // 启动 Device Code Flow 并获取设备码信息 + let (device_response, wait_future) = qwen::start_qwen_device_code_and_get_info() + .await + .map_err(|e| format!("启动 Device Code Flow 失败: {}", e))?; + + tracing::info!( + "[Qwen] Device Code: user_code={}, verification_uri={}", + device_response.user_code, + device_response.verification_uri + ); + + // 通过事件发送设备码信息给前端 + let _ = app.emit( + "qwen-device-code", + QwenDeviceCodeResponse { + user_code: device_response.user_code.clone(), + verification_uri: device_response.verification_uri.clone(), + verification_uri_complete: device_response.verification_uri_complete.clone(), + expires_in: device_response.expires_in, + }, + ); + + // 等待用户授权 + let result = wait_future.await.map_err(|e| e.to_string())?; + + tracing::info!("[Qwen] 登录成功，凭证保存到: {}", result.creds_file_path); + + // 添加到凭证池 + let credential = pool_service.0.add_credential( + &db, + "qwen", + CredentialData::QwenOAuth { + creds_file_path: result.creds_file_path, + }, + name, + Some(true), + None, + )?; + + tracing::info!("[Qwen] 凭证已添加到凭证池: {}", credential.uuid); + + Ok(credential) +} + +/// 启动 Qwen Device Code Flow 登录流程 +/// +/// 自动打开浏览器让用户完成授权 +#[tauri::command] +pub async fn start_qwen_device_code_login( + db: State<'_, DbConnection>, + pool_service: State<'_, ProviderPoolServiceState>, + name: Option, +) -> Result { + use crate::providers::qwen; + + tracing::info!("[Qwen] 开始 Device Code Flow 登录流程"); + + // 启动 Device Code Flow 登录 + let result = qwen::start_qwen_device_code_login() + .await + .map_err(|e| format!("Qwen Device Code Flow 登录失败: {}", e))?; + + tracing::info!("[Qwen] 登录成功，凭证保存到: {}", result.creds_file_path); + + // 添加到凭证池 + let credential = pool_service.0.add_credential( + &db, + "qwen", + CredentialData::QwenOAuth { + creds_file_path: result.creds_file_path, + }, + name, + Some(true), + None, + )?; + + tracing::info!("[Qwen] 凭证已添加到凭证池: {}", credential.uuid); + + Ok(credential) +} + +/// iFlow OAuth 授权 URL 响应 +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +pub struct IFlowAuthUrlResponse { + pub auth_url: String, +} + +/// 获取 iFlow OAuth 授权 URL 并等待回调（不自动打开浏览器） +/// +/// 启动服务器后通过事件发送授权 URL，然后等待回调 +/// 成功后返回凭证 +#[tauri::command] +pub async fn get_iflow_auth_url_and_wait( + app: tauri::AppHandle, + db: State<'_, DbConnection>, + pool_service: State<'_, ProviderPoolServiceState>, + name: Option, +) -> Result { + use crate::providers::iflow; + + tracing::info!("[iFlow OAuth] 启动服务器并获取授权 URL"); + + // 启动服务器并获取授权 URL + let (auth_url, wait_future) = iflow::start_iflow_oauth_server_and_get_url() + .await + .map_err(|e| format!("启动 OAuth 服务器失败: {}", e))?; + + tracing::info!("[iFlow OAuth] 授权 URL: {}", auth_url); + + // 通过事件发送授权 URL 给前端 + let _ = app.emit( + "iflow-auth-url", + IFlowAuthUrlResponse { + auth_url: auth_url.clone(), + }, + ); + + // 等待回调 + let result = wait_future.await.map_err(|e| e.to_string())?; + + tracing::info!( + "[iFlow OAuth] 登录成功，凭证保存到: {}", + result.creds_file_path + ); + + // 添加到凭证池 + let credential = pool_service.0.add_credential( + &db, + "iflow", + CredentialData::IFlowOAuth { + creds_file_path: result.creds_file_path, + }, + name, + Some(true), + None, + )?; + + tracing::info!("[iFlow OAuth] 凭证已添加到凭证池: {}", credential.uuid); + + Ok(credential) +} + +/// 启动 iFlow OAuth 登录流程 +/// +/// 打开浏览器让用户登录 iFlow 账号，获取凭证 +#[tauri::command] +pub async fn start_iflow_oauth_login( + db: State<'_, DbConnection>, + pool_service: State<'_, ProviderPoolServiceState>, + name: Option, +) -> Result { + use crate::providers::iflow; + + tracing::info!("[iFlow OAuth] 开始 OAuth 登录流程"); + + // 启动 OAuth 登录 + let result = iflow::start_iflow_oauth_login() + .await + .map_err(|e| format!("iFlow OAuth 登录失败: {}", e))?; + + tracing::info!( + "[iFlow OAuth] 登录成功，凭证保存到: {}", + result.creds_file_path + ); + + // 添加到凭证池 + let credential = pool_service.0.add_credential( + &db, + "iflow", + CredentialData::IFlowOAuth { + creds_file_path: result.creds_file_path, + }, + name, + Some(true), + None, + )?; + + tracing::info!("[iFlow OAuth] 凭证已添加到凭证池: {}", credential.uuid); + + Ok(credential) +} + +/// 获取 Kiro 凭证的 Machine ID 指纹信息 +/// +/// 返回凭证的唯一设备指纹，用于在 UI 中展示 +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +pub struct KiroFingerprintInfo { + /// Machine ID（SHA256 哈希，64 字符） + pub machine_id: String, + /// Machine ID 的短格式（前 16 字符） + pub machine_id_short: String, + /// 指纹来源（profileArn / clientId / system） + pub source: String, + /// 认证方式 + pub auth_method: String, +} + +#[tauri::command] +pub async fn get_kiro_credential_fingerprint( + db: State<'_, DbConnection>, + uuid: String, +) -> Result { + use crate::database::dao::provider_pool::ProviderPoolDao; + use crate::providers::kiro::{generate_machine_id_from_credentials, KiroProvider}; + + // 获取凭证文件路径（在锁释放前完成） + let creds_file_path = { + let conn = db.lock().map_err(|e| e.to_string())?; + let credential = ProviderPoolDao::get_by_uuid(&conn, &uuid) + .map_err(|e| e.to_string())? + .ok_or_else(|| format!("凭证不存在: {}", uuid))?; + + // 检查是否为 Kiro 凭证 + match &credential.credential { + CredentialData::KiroOAuth { creds_file_path } => creds_file_path.clone(), + _ => return Err("只有 Kiro 凭证支持获取指纹信息".to_string()), + } + }; // conn 在这里释放 + + // 加载凭证文件（异步操作，锁已释放） + let mut provider = KiroProvider::new(); + provider + .load_credentials_from_path(&creds_file_path) + .await + .map_err(|e| format!("加载凭证失败: {}", e))?; + + // 确定指纹来源 + let (source, profile_arn, client_id) = if provider.credentials.profile_arn.is_some() { + ( + "profileArn".to_string(), + provider.credentials.profile_arn.as_deref(), + None, + ) + } else if provider.credentials.client_id.is_some() { + ( + "clientId".to_string(), + None, + provider.credentials.client_id.as_deref(), + ) + } else { + ("system".to_string(), None, None) + }; + + // 生成 Machine ID + let machine_id = generate_machine_id_from_credentials(profile_arn, client_id); + let machine_id_short = machine_id[..16].to_string(); + + // 获取认证方式 + let auth_method = provider + .credentials + .auth_method + .clone() + .unwrap_or_else(|| "social".to_string()); + + Ok(KiroFingerprintInfo { + machine_id, + machine_id_short, + source, + auth_method, + }) +} + +/// Gemini OAuth 授权 URL 响应 +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +pub struct GeminiAuthUrlResponse { + pub auth_url: String, + pub session_id: String, +} + +use once_cell::sync::Lazy; +/// Gemini OAuth 会话存储（用于存储 code_verifier） +use std::collections::HashMap; +use tokio::sync::RwLock; + +static GEMINI_OAUTH_SESSIONS: Lazy< + RwLock>, +> = Lazy::new(|| RwLock::new(HashMap::new())); + +/// 获取 Gemini OAuth 授权 URL（不等待回调） +/// +/// 生成授权 URL 和 session_id，通过事件发送给前端 +/// 用户需要手动复制授权码回来，然后调用 exchange_gemini_code +#[tauri::command] +pub async fn get_gemini_auth_url_and_wait( + app: tauri::AppHandle, + _db: State<'_, DbConnection>, + _pool_service: State<'_, ProviderPoolServiceState>, + _name: Option, +) -> Result { + use crate::providers::gemini; + + tracing::info!("[Gemini OAuth] 生成授权 URL"); + + // 生成授权 URL 和会话信息 + let (auth_url, session) = gemini::generate_gemini_auth_url_with_session(); + let session_id = session.session_id.clone(); + + tracing::info!("[Gemini OAuth] 授权 URL: {}", auth_url); + tracing::info!("[Gemini OAuth] Session ID: {}", session_id); + + // 存储会话信息（用于后续交换 token） + { + let mut sessions = GEMINI_OAUTH_SESSIONS.write().await; + sessions.insert(session_id.clone(), session); + + // 清理过期的会话（超过 10 分钟） + let now = chrono::Utc::now().timestamp(); + sessions.retain(|_, s| now - s.created_at < 600); + } + + // 通过事件发送授权 URL 给前端 + let _ = app.emit( + "gemini-auth-url", + GeminiAuthUrlResponse { + auth_url: auth_url.clone(), + session_id: session_id.clone(), + }, + ); + + // 返回错误，让前端知道需要用户手动输入授权码 + // 这不是真正的错误，只是流程需要用户交互 + Err(format!("AUTH_URL:{}", auth_url)) +} + +/// 用 Gemini 授权码交换 Token 并添加凭证 +#[tauri::command] +pub async fn exchange_gemini_code( + db: State<'_, DbConnection>, + pool_service: State<'_, ProviderPoolServiceState>, + code: String, + session_id: Option, + name: Option, +) -> Result { + use crate::providers::gemini; + + tracing::info!("[Gemini OAuth] 开始交换授权码"); + + // 获取 code_verifier + let code_verifier = if let Some(ref sid) = session_id { + let sessions = GEMINI_OAUTH_SESSIONS.read().await; + sessions + .get(sid) + .map(|s| s.code_verifier.clone()) + .ok_or_else(|| "会话已过期，请重新获取授权 URL".to_string())? + } else { + // 如果没有 session_id，尝试使用最近的会话 + let sessions = GEMINI_OAUTH_SESSIONS.read().await; + sessions + .values() + .max_by_key(|s| s.created_at) + .map(|s| s.code_verifier.clone()) + .ok_or_else(|| "没有可用的会话，请先获取授权 URL".to_string())? + }; + + // 交换 token 并创建凭证 + let result = gemini::exchange_gemini_code_and_create_credentials(&code, &code_verifier) + .await + .map_err(|e| format!("交换授权码失败: {}", e))?; + + tracing::info!( + "[Gemini OAuth] 登录成功，凭证保存到: {}", + result.creds_file_path + ); + + // 清理使用过的会话 + if let Some(ref sid) = session_id { + let mut sessions = GEMINI_OAUTH_SESSIONS.write().await; + sessions.remove(sid); + } + + // 添加到凭证池 + let credential = pool_service.0.add_credential( + &db, + "gemini", + CredentialData::GeminiOAuth { + creds_file_path: result.creds_file_path, + project_id: None, // 项目 ID 会在健康检查时自动获取 + }, + name, + Some(true), + None, + )?; + + tracing::info!("[Gemini OAuth] 凭证已添加到凭证池: {}", credential.uuid); + + Ok(credential) +} + +/// 启动 Gemini OAuth 登录流程 +/// +/// 打开浏览器让用户登录 Google 账号，获取 Gemini 凭证 +#[tauri::command] +pub async fn start_gemini_oauth_login( + db: State<'_, DbConnection>, + pool_service: State<'_, ProviderPoolServiceState>, + name: Option, +) -> Result { + use crate::providers::gemini; + + tracing::info!("[Gemini OAuth] 开始 OAuth 登录流程"); + + // 启动 OAuth 登录 + let result = gemini::start_gemini_oauth_login() + .await + .map_err(|e| format!("Gemini OAuth 登录失败: {}", e))?; + + tracing::info!( + "[Gemini OAuth] 登录成功，凭证保存到: {}", + result.creds_file_path + ); + + // 添加到凭证池 + let credential = pool_service.0.add_credential( + &db, + "gemini", + CredentialData::GeminiOAuth { + creds_file_path: result.creds_file_path, + project_id: None, + }, + name, + Some(true), + None, + )?; + + tracing::info!("[Gemini OAuth] 凭证已添加到凭证池: {}", credential.uuid); + + Ok(credential) +} diff --git a/src-tauri/src/converter/openai_to_antigravity.rs b/src-tauri/src/converter/openai_to_antigravity.rs index 26a7ec977..7753bbb9e 100644 --- a/src-tauri/src/converter/openai_to_antigravity.rs +++ b/src-tauri/src/converter/openai_to_antigravity.rs @@ -1,9 +1,11 @@ //! OpenAI 格式转换为 Antigravity (Gemini) 格式 use crate::models::openai::*; use serde::{Deserialize, Serialize}; +use uuid::Uuid; /// Antigravity/Gemini 内容部分 #[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] pub struct GeminiPart { #[serde(skip_serializing_if = "Option::is_none")] pub text: Option, @@ -16,6 +18,7 @@ pub struct GeminiPart { } #[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] pub struct InlineData { pub mime_type: String, pub data: String, @@ -23,12 +26,16 @@ pub struct InlineData { #[derive(Debug, Clone, Serialize, Deserialize)] pub struct GeminiFunctionCall { + #[serde(skip_serializing_if = "Option::is_none")] + pub id: Option, pub name: String, pub args: serde_json::Value, } #[derive(Debug, Clone, Serialize, Deserialize)] pub struct GeminiFunctionResponse { + #[serde(skip_serializing_if = "Option::is_none")] + pub id: Option, pub name: String, pub response: serde_json::Value, } @@ -42,6 +49,7 @@ pub struct GeminiContent { /// Antigravity/Gemini 工具定义 #[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] pub struct GeminiTool { pub function_declarations: Vec, } @@ -69,12 +77,23 @@ pub struct GeminiGenerationConfig { pub top_k: Option, #[serde(skip_serializing_if = "Option::is_none")] pub stop_sequences: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + pub candidate_count: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub thinking_config: Option, } -/// Antigravity 请求体 #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] -pub struct AntigravityRequestBody { +pub struct ThinkingConfig { + pub include_thoughts: bool, + pub thinking_budget: i32, +} + +/// Antigravity 请求体内部结构 +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct AntigravityRequestInner { pub contents: Vec, #[serde(skip_serializing_if = "Option::is_none")] pub system_instruction: Option, @@ -82,10 +101,59 @@ pub struct AntigravityRequestBody { pub generation_config: Option, #[serde(skip_serializing_if = "Option::is_none")] pub tools: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + pub tool_config: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub session_id: Option, +} + +/// 生成随机请求 ID +fn generate_request_id() -> String { + format!("agent-{}", Uuid::new_v4()) +} + +/// 生成随机会话 ID +fn generate_session_id() -> String { + let uuid = Uuid::new_v4(); + let bytes = uuid.as_bytes(); + let n: u64 = u64::from_le_bytes([ + bytes[0], bytes[1], bytes[2], bytes[3], bytes[4], bytes[5], bytes[6], bytes[7], + ]) % 9_000_000_000_000_000_000; + format!("-{}", n) +} + +/// 模型名称映射 +fn model_mapping(model: &str) -> &str { + match model { + "claude-sonnet-4-5-thinking" => "claude-sonnet-4-5", + "claude-opus-4-5" => "claude-opus-4-5-thinking", + "gemini-2.5-flash-thinking" => "gemini-2.5-flash", + "gemini-2.5-computer-use-preview-10-2025" => "rev19-uic3-1p", + "gemini-3-pro-image-preview" => "gemini-3-pro-image", + "gemini-3-pro-preview" => "gemini-3-pro-high", + "gemini-claude-sonnet-4-5" => "claude-sonnet-4-5", + "gemini-claude-sonnet-4-5-thinking" => "claude-sonnet-4-5-thinking", + _ => model, + } +} + +/// 是否启用思维链 +fn is_enable_thinking(model: &str) -> bool { + model.ends_with("-thinking") + || model == "gemini-2.5-pro" + || model.starts_with("gemini-3-pro-") + || model == "rev19-uic3-1p" + || model == "gpt-oss-120b-medium" } /// 将 OpenAI ChatCompletionRequest 转换为 Antigravity 请求体 -pub fn convert_openai_to_antigravity(request: &ChatCompletionRequest) -> serde_json::Value { +pub fn convert_openai_to_antigravity_with_context( + request: &ChatCompletionRequest, + project_id: &str, +) -> serde_json::Value { + let actual_model = model_mapping(&request.model); + let enable_thinking = is_enable_thinking(&request.model); + let mut contents: Vec = Vec::new(); let mut system_instruction: Option = None; @@ -116,35 +184,67 @@ pub fn convert_openai_to_antigravity(request: &ChatCompletionRequest) -> serde_j } } "assistant" => { - let parts = convert_assistant_content(msg); + let parts = convert_assistant_content(msg, &contents); if !parts.is_empty() { - contents.push(GeminiContent { - role: "model".to_string(), - parts, - }); + // 检查是否需要合并到上一条 model 消息 + let should_merge = if let Some(last) = contents.last() { + last.role == "model" + && msg.tool_calls.is_some() + && msg.get_content_text().is_empty() + } else { + false + }; + + if should_merge { + if let Some(last) = contents.last_mut() { + last.parts.extend(parts); + } + } else { + contents.push(GeminiContent { + role: "model".to_string(), + parts, + }); + } } } "tool" => { - // Tool 响应需要合并到 user 消息 + // Tool 响应 let tool_id = msg.tool_call_id.clone().unwrap_or_default(); let content = msg.get_content_text(); - // 尝试解析为 JSON，否则包装为对象 - let response_value = serde_json::from_str(&content) - .unwrap_or_else(|_| serde_json::json!({ "result": content })); + // 从之前的 model 消息中找到对应的 functionCall name + let function_name = find_function_name(&contents, &tool_id); - contents.push(GeminiContent { - role: "user".to_string(), - parts: vec![GeminiPart { - text: None, - inline_data: None, - function_call: None, - function_response: Some(GeminiFunctionResponse { - name: tool_id, - response: response_value, - }), - }], - }); + let response_value = serde_json::json!({ "output": content }); + + let function_response = GeminiPart { + text: None, + inline_data: None, + function_call: None, + function_response: Some(GeminiFunctionResponse { + id: Some(tool_id), + name: function_name, + response: response_value, + }), + }; + + // 检查是否需要合并到上一条 user 消息 + let should_merge = if let Some(last) = contents.last() { + last.role == "user" && last.parts.iter().any(|p| p.function_response.is_some()) + } else { + false + }; + + if should_merge { + if let Some(last) = contents.last_mut() { + last.parts.push(function_response); + } + } else { + contents.push(GeminiContent { + role: "user".to_string(), + parts: vec![function_response], + }); + } } _ => {} } @@ -152,40 +252,120 @@ pub fn convert_openai_to_antigravity(request: &ChatCompletionRequest) -> serde_j // 构建生成配置 let generation_config = Some(GeminiGenerationConfig { - temperature: request.temperature, - max_output_tokens: request.max_tokens.map(|t| t as i32), - top_p: None, - top_k: None, - stop_sequences: None, + temperature: request.temperature.or(Some(1.0)), + max_output_tokens: request.max_tokens.map(|t| t as i32).or(Some(8096)), + top_p: Some(0.85), + top_k: Some(50), + stop_sequences: Some(vec![ + "<|user|>".to_string(), + "<|bot|>".to_string(), + "<|context_request|>".to_string(), + "<|endoftext|>".to_string(), + "<|end_of_turn|>".to_string(), + ]), + candidate_count: Some(1), + thinking_config: Some(ThinkingConfig { + include_thoughts: enable_thinking, + thinking_budget: if enable_thinking { 1024 } else { 0 }, + }), }); // 转换工具 let tools = request.tools.as_ref().map(|tools| { - vec![GeminiTool { - function_declarations: tools - .iter() - .map(|t| GeminiFunctionDeclaration { + tools + .iter() + .map(|t| GeminiTool { + function_declarations: vec![GeminiFunctionDeclaration { name: t.function.name.clone(), description: t.function.description.clone(), - parameters: t.function.parameters.clone(), - }) - .collect(), - }] + parameters: clean_parameters(t.function.parameters.clone()), + }], + }) + .collect() }); - let body = AntigravityRequestBody { + let tool_config = if tools.is_some() { + Some(serde_json::json!({ + "functionCallingConfig": { + "mode": "VALIDATED" + } + })) + } else { + None + }; + + let inner = AntigravityRequestInner { contents, system_instruction, generation_config, tools, + tool_config, + session_id: Some(generate_session_id()), }; - // 包装为 Antigravity 请求格式 + // 构建完整的 Antigravity 请求体 serde_json::json!({ - "request": body + "project": project_id, + "requestId": generate_request_id(), + "request": inner, + "model": actual_model, + "userAgent": "antigravity" }) } +/// 从之前的 model 消息中找到对应的 functionCall name +fn find_function_name(contents: &[GeminiContent], tool_id: &str) -> String { + for content in contents.iter().rev() { + if content.role == "model" { + for part in &content.parts { + if let Some(fc) = &part.function_call { + if fc.id.as_deref() == Some(tool_id) { + return fc.name.clone(); + } + } + } + } + } + String::new() +} + +/// 清理参数中不需要的字段 +fn clean_parameters(params: Option) -> Option { + params.map(|v| clean_value(v)) +} + +fn clean_value(value: serde_json::Value) -> serde_json::Value { + const EXCLUDED_KEYS: &[&str] = &[ + "$schema", + "additionalProperties", + "minLength", + "maxLength", + "minItems", + "maxItems", + "uniqueItems", + ]; + + match value { + serde_json::Value::Object(map) => { + let cleaned: serde_json::Map = map + .into_iter() + .filter(|(k, _)| !EXCLUDED_KEYS.contains(&k.as_str())) + .map(|(k, v)| (k, clean_value(v))) + .collect(); + serde_json::Value::Object(cleaned) + } + serde_json::Value::Array(arr) => { + serde_json::Value::Array(arr.into_iter().map(clean_value).collect()) + } + other => other, + } +} + +/// 兼容旧接口 +pub fn convert_openai_to_antigravity(request: &ChatCompletionRequest) -> serde_json::Value { + convert_openai_to_antigravity_with_context(request, "") +} + /// 转换用户消息内容 fn convert_user_content(msg: &ChatMessage) -> Vec { let mut parts = Vec::new(); @@ -234,14 +414,14 @@ fn convert_user_content(msg: &ChatMessage) -> Vec { } /// 转换助手消息内容 -fn convert_assistant_content(msg: &ChatMessage) -> Vec { +fn convert_assistant_content(msg: &ChatMessage, _contents: &[GeminiContent]) -> Vec { let mut parts = Vec::new(); // 文本内容 let text = msg.get_content_text(); if !text.is_empty() { parts.push(GeminiPart { - text: Some(text), + text: Some(text.trim_end().to_string()), inline_data: None, function_call: None, function_response: None, @@ -258,8 +438,9 @@ fn convert_assistant_content(msg: &ChatMessage) -> Vec { text: None, inline_data: None, function_call: Some(GeminiFunctionCall { + id: Some(tc.id.clone()), name: tc.function.name.clone(), - args, + args: serde_json::json!({ "query": args }), }), function_response: None, }); diff --git a/src-tauri/src/converter/openai_to_cw.rs b/src-tauri/src/converter/openai_to_cw.rs index 8586e5730..334a73cc5 100644 --- a/src-tauri/src/converter/openai_to_cw.rs +++ b/src-tauri/src/converter/openai_to_cw.rs @@ -5,18 +5,31 @@ use std::collections::HashMap; use uuid::Uuid; /// 模型映射表 +/// +/// 参考 AIClient-2-API 的 provider-models.js 和 claude-kiro.js +/// 支持的模型列表： +/// - claude-opus-4-5, claude-opus-4-5-20251101 +/// - claude-haiku-4-5, claude-haiku-4-5-20251001 +/// - claude-sonnet-4-5, claude-sonnet-4-5-20250929 +/// - claude-sonnet-4-20250514 +/// - claude-3-7-sonnet-20250219, claude-3-5-sonnet-20241022, claude-3-5-sonnet-latest pub fn get_model_map() -> HashMap<&'static str, &'static str> { let mut map = HashMap::new(); + // Opus 4.5 系列 map.insert("claude-opus-4-5", "claude-opus-4.5"); map.insert("claude-opus-4-5-20251101", "claude-opus-4.5"); + // Haiku 4.5 系列 map.insert("claude-haiku-4-5", "claude-haiku-4.5"); map.insert("claude-haiku-4-5-20251001", "claude-haiku-4.5"); + // Sonnet 4.5 系列 map.insert("claude-sonnet-4-5", "CLAUDE_SONNET_4_5_20250929_V1_0"); map.insert( "claude-sonnet-4-5-20250929", "CLAUDE_SONNET_4_5_20250929_V1_0", ); + // Sonnet 4 系列 map.insert("claude-sonnet-4-20250514", "CLAUDE_SONNET_4_20250514_V1_0"); + // Sonnet 3.7/3.5 系列（兼容旧版本） map.insert( "claude-3-7-sonnet-20250219", "CLAUDE_3_7_SONNET_20250219_V1_0", @@ -32,6 +45,22 @@ pub fn get_model_map() -> HashMap<&'static str, &'static str> { map } +/// 获取 Kiro 支持的模型列表 +/// +/// 返回所有支持的模型 ID，用于前端展示和健康检查 +pub fn get_supported_models() -> Vec<&'static str> { + vec![ + "claude-opus-4-5", + "claude-opus-4-5-20251101", + "claude-haiku-4-5", + "claude-haiku-4-5-20251001", + "claude-sonnet-4-5", + "claude-sonnet-4-5-20250929", + "claude-sonnet-4-20250514", + "claude-3-7-sonnet-20250219", + ] +} + pub const DEFAULT_MODEL: &str = "CLAUDE_SONNET_4_5_20250929_V1_0"; /// 预处理消息：合并连续的 tool 消息到前一个 assistant 消息后的 user 消息 diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index 22a47cab7..674c1f0d1 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -14,6 +14,7 @@ pub mod proxy; pub mod resilience; pub mod router; mod server; +mod server_utils; mod services; pub mod telemetry; pub mod tray; @@ -58,6 +59,14 @@ pub enum ProviderType { /// Gemini API Key (multi-account load balancing) #[serde(rename = "gemini_api_key")] GeminiApiKey, + /// Codex (OpenAI OAuth) + Codex, + /// Claude OAuth (Anthropic OAuth) + #[serde(rename = "claude_oauth")] + ClaudeOAuth, + /// iFlow + #[serde(rename = "iflow")] + IFlow, } impl std::fmt::Display for ProviderType { @@ -71,6 +80,9 @@ impl std::fmt::Display for ProviderType { ProviderType::Antigravity => write!(f, "antigravity"), ProviderType::Vertex => write!(f, "vertex"), ProviderType::GeminiApiKey => write!(f, "gemini_api_key"), + ProviderType::Codex => write!(f, "codex"), + ProviderType::ClaudeOAuth => write!(f, "claude_oauth"), + ProviderType::IFlow => write!(f, "iflow"), } } } @@ -88,6 +100,9 @@ impl std::str::FromStr for ProviderType { "antigravity" => Ok(ProviderType::Antigravity), "vertex" => Ok(ProviderType::Vertex), "gemini_api_key" => Ok(ProviderType::GeminiApiKey), + "codex" => Ok(ProviderType::Codex), + "claude_oauth" => Ok(ProviderType::ClaudeOAuth), + "iflow" => Ok(ProviderType::IFlow), _ => Err(format!("Invalid provider: {s}")), } } @@ -1030,6 +1045,12 @@ async fn check_api_compatibility( ("gemini-2.5-flash", "basic"), ("gemini-2.5-flash", "tool_call"), ], + ProviderType::Codex => vec![("gpt-4.1", "basic"), ("gpt-4.1", "tool_call")], + ProviderType::ClaudeOAuth => vec![ + ("claude-sonnet-4-5", "basic"), + ("claude-sonnet-4-5", "tool_call"), + ], + ProviderType::IFlow => vec![("gpt-4o", "basic"), ("gpt-4o", "tool_call")], ProviderType::OpenAI | ProviderType::Claude => vec![], }; @@ -1751,6 +1772,20 @@ pub fn run() { commands::provider_pool_cmd::debug_kiro_credentials, commands::provider_pool_cmd::test_user_credentials, commands::provider_pool_cmd::migrate_private_config_to_pool, + commands::provider_pool_cmd::start_antigravity_oauth_login, + commands::provider_pool_cmd::get_antigravity_auth_url_and_wait, + commands::provider_pool_cmd::get_codex_auth_url_and_wait, + commands::provider_pool_cmd::start_codex_oauth_login, + commands::provider_pool_cmd::get_claude_oauth_auth_url_and_wait, + commands::provider_pool_cmd::start_claude_oauth_login, + commands::provider_pool_cmd::get_qwen_device_code_and_wait, + commands::provider_pool_cmd::start_qwen_device_code_login, + commands::provider_pool_cmd::get_iflow_auth_url_and_wait, + commands::provider_pool_cmd::start_iflow_oauth_login, + commands::provider_pool_cmd::get_gemini_auth_url_and_wait, + commands::provider_pool_cmd::start_gemini_oauth_login, + commands::provider_pool_cmd::exchange_gemini_code, + commands::provider_pool_cmd::get_kiro_credential_fingerprint, // Route commands commands::route_cmd::get_available_routes, commands::route_cmd::get_route_curl_examples, diff --git a/src-tauri/src/models/provider_pool_model.rs b/src-tauri/src/models/provider_pool_model.rs index c18c2c9c6..0e15aca6d 100644 --- a/src-tauri/src/models/provider_pool_model.rs +++ b/src-tauri/src/models/provider_pool_model.rs @@ -21,69 +21,11 @@ pub enum CredentialSource { Private, } -/// Provider 类型枚举 -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] -#[serde(rename_all = "lowercase")] -pub enum PoolProviderType { - Kiro, - Gemini, - Qwen, - #[serde(rename = "openai")] - OpenAI, - Claude, - Antigravity, - Vertex, - /// Gemini API Key (multi-account load balancing) - #[serde(rename = "gemini_api_key")] - GeminiApiKey, - /// Codex (OpenAI OAuth) - Codex, - /// Claude OAuth (Anthropic OAuth) - #[serde(rename = "claude_oauth")] - ClaudeOAuth, - /// iFlow - #[serde(rename = "iflow")] - IFlow, -} - -impl std::fmt::Display for PoolProviderType { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - PoolProviderType::Kiro => write!(f, "kiro"), - PoolProviderType::Gemini => write!(f, "gemini"), - PoolProviderType::Qwen => write!(f, "qwen"), - PoolProviderType::OpenAI => write!(f, "openai"), - PoolProviderType::Claude => write!(f, "claude"), - PoolProviderType::Antigravity => write!(f, "antigravity"), - PoolProviderType::Vertex => write!(f, "vertex"), - PoolProviderType::GeminiApiKey => write!(f, "gemini_api_key"), - PoolProviderType::Codex => write!(f, "codex"), - PoolProviderType::ClaudeOAuth => write!(f, "claude_oauth"), - PoolProviderType::IFlow => write!(f, "iflow"), - } - } -} - -impl std::str::FromStr for PoolProviderType { - type Err = String; - - fn from_str(s: &str) -> Result { - match s.to_lowercase().as_str() { - "kiro" => Ok(PoolProviderType::Kiro), - "gemini" => Ok(PoolProviderType::Gemini), - "qwen" => Ok(PoolProviderType::Qwen), - "openai" => Ok(PoolProviderType::OpenAI), - "claude" => Ok(PoolProviderType::Claude), - "antigravity" => Ok(PoolProviderType::Antigravity), - "vertex" => Ok(PoolProviderType::Vertex), - "gemini_api_key" => Ok(PoolProviderType::GeminiApiKey), - "codex" => Ok(PoolProviderType::Codex), - "claude_oauth" => Ok(PoolProviderType::ClaudeOAuth), - "iflow" => Ok(PoolProviderType::IFlow), - _ => Err(format!("Invalid provider type: {s}")), - } - } -} +/// Provider 类型别名 +/// +/// 为了向后兼容，PoolProviderType 是 crate::ProviderType 的类型别名。 +/// 所有 Provider 类型定义已统一到 lib.rs 中的 ProviderType。 +pub type PoolProviderType = crate::ProviderType; /// 凭证数据，根据 Provider 类型不同而不同 #[derive(Debug, Clone, Serialize, Deserialize)] @@ -334,6 +276,7 @@ impl ProviderCredential { /// 检查两个来源的排除列表： /// 1. `not_supported_models` - 通用的不支持模型列表（精确匹配） /// 2. `excluded_models` - 来自 CredentialData::GeminiApiKey 的排除列表（支持通配符） + /// 3. Antigravity 凭证只支持特定的模型列表 pub fn supports_model(&self, model: &str) -> bool { // 检查通用的不支持模型列表（精确匹配） if self.not_supported_models.contains(&model.to_string()) { @@ -352,6 +295,19 @@ impl ProviderCredential { } } + // Antigravity 凭证只支持特定的模型 + if let CredentialData::AntigravityOAuth { .. } = &self.credential { + // Antigravity 支持的模型列表 + const ANTIGRAVITY_SUPPORTED_MODELS: &[&str] = &[ + "gemini-3-pro-preview", + "gemini-3-pro-image-preview", + "gemini-2.5-computer-use-preview-10-2025", + "gemini-claude-sonnet-4-5", + "gemini-claude-sonnet-4-5-thinking", + ]; + return ANTIGRAVITY_SUPPORTED_MODELS.contains(&model); + } + true } diff --git a/src-tauri/src/providers/antigravity.rs b/src-tauri/src/providers/antigravity.rs index ad03af118..83079714b 100644 --- a/src-tauri/src/providers/antigravity.rs +++ b/src-tauri/src/providers/antigravity.rs @@ -2,10 +2,14 @@ //! //! 支持 Gemini 3 Pro 等高级模型，通过 Google 内部 API 访问。 +use super::traits::{CredentialProvider, ProviderResult}; +use async_trait::async_trait; use reqwest::Client; use serde::{Deserialize, Serialize}; use std::error::Error; use std::path::PathBuf; +use std::sync::Arc; +use tokio::sync::oneshot; use uuid::Uuid; // Constants @@ -16,9 +20,18 @@ const CREDENTIALS_DIR: &str = ".antigravity"; const CREDENTIALS_FILE: &str = "oauth_creds.json"; // OAuth credentials - 与 Antigravity CLI 相同 -const OAUTH_CLIENT_ID: &str = +pub const OAUTH_CLIENT_ID: &str = "1071006060591-tmhssin2h21lcre235vtolojh4g403ep.apps.googleusercontent.com"; -const OAUTH_CLIENT_SECRET: &str = "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"; +pub const OAUTH_CLIENT_SECRET: &str = "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"; + +// OAuth scopes +const OAUTH_SCOPES: &[&str] = &[ + "https://www.googleapis.com/auth/cloud-platform", + "https://www.googleapis.com/auth/userinfo.email", + "https://www.googleapis.com/auth/userinfo.profile", + "https://www.googleapis.com/auth/cclog", + "https://www.googleapis.com/auth/experimentsandconfigs", +]; // Token 刷新提前量（秒） const REFRESH_SKEW: i64 = 3000; @@ -27,9 +40,12 @@ const REFRESH_SKEW: i64 = 3000; pub const ANTIGRAVITY_MODELS: &[&str] = &[ "gemini-3-pro-preview", "gemini-3-pro-image-preview", + "gemini-3-flash-preview", + "gemini-2.5-flash", "gemini-2.5-computer-use-preview-10-2025", "gemini-claude-sonnet-4-5", "gemini-claude-sonnet-4-5-thinking", + "gemini-claude-opus-4-5-thinking", ]; /// 模型别名映射（用户友好名称 -> 内部名称） @@ -38,8 +54,11 @@ fn alias_to_model_name(model: &str) -> &str { "gemini-2.5-computer-use-preview-10-2025" => "rev19-uic3-1p", "gemini-3-pro-image-preview" => "gemini-3-pro-image", "gemini-3-pro-preview" => "gemini-3-pro-high", + "gemini-3-flash-preview" => "gemini-3-flash", + "gemini-2.5-flash" => "gemini-2.5-flash", "gemini-claude-sonnet-4-5" => "claude-sonnet-4-5", "gemini-claude-sonnet-4-5-thinking" => "claude-sonnet-4-5-thinking", + "gemini-claude-opus-4-5-thinking" => "claude-opus-4-5-thinking", _ => model, } } @@ -51,8 +70,11 @@ fn model_name_to_alias(model: &str) -> &str { "rev19-uic3-1p" => "gemini-2.5-computer-use-preview-10-2025", "gemini-3-pro-image" => "gemini-3-pro-image-preview", "gemini-3-pro-high" => "gemini-3-pro-preview", + "gemini-3-flash" => "gemini-3-flash-preview", + "gemini-2.5-flash" => "gemini-2.5-flash", "claude-sonnet-4-5" => "gemini-claude-sonnet-4-5", "claude-sonnet-4-5-thinking" => "gemini-claude-sonnet-4-5-thinking", + "claude-opus-4-5-thinking" => "gemini-claude-opus-4-5-thinking", _ => model, } } @@ -93,8 +115,7 @@ pub struct AntigravityCredentials { /// 过期时间戳（毫秒）- 兼容旧格式 #[serde(skip_serializing_if = "Option::is_none")] pub expiry_date: Option, - /// 过期时间（RFC3339 格式）- 与 CLIProxyAPI 兼容 - #[serde(skip_serializing_if = "Option::is_none")] + /// 过期时间（RFC3339 格式）- pub expire: Option, pub scope: Option, /// 最后刷新时间（RFC3339 格式） @@ -103,6 +124,21 @@ pub struct AntigravityCredentials { /// 凭证类型标识 #[serde(default = "default_antigravity_type", rename = "type")] pub cred_type: String, + /// Token 有效期（秒） + #[serde(skip_serializing_if = "Option::is_none")] + pub expires_in: Option, + /// Token 获取时间戳（毫秒）- + #[serde(skip_serializing_if = "Option::is_none")] + pub timestamp: Option, + /// 是否启用 - + #[serde(skip_serializing_if = "Option::is_none")] + pub enable: Option, + /// 项目 ID + #[serde(skip_serializing_if = "Option::is_none", alias = "project_id")] + pub projectId: Option, + /// 用户邮箱 + #[serde(skip_serializing_if = "Option::is_none")] + pub email: Option, } fn default_antigravity_type() -> String { @@ -120,6 +156,11 @@ impl Default for AntigravityCredentials { scope: None, last_refresh: None, cred_type: default_antigravity_type(), + expires_in: None, + timestamp: None, + enable: None, + projectId: None, + email: None, } } } @@ -180,9 +221,32 @@ impl AntigravityProvider { path: &str, ) -> Result<(), Box> { let content = tokio::fs::read_to_string(path).await?; - let creds: AntigravityCredentials = serde_json::from_str(&content)?; - self.credentials = creds; - Ok(()) + + // 尝试解析为单个凭证对象 + if let Ok(creds) = serde_json::from_str::(&content) { + self.credentials = creds; + // 如果凭证中有 projectId，设置到 provider + if let Some(ref pid) = self.credentials.projectId { + self.project_id = Some(pid.clone()); + } + return Ok(()); + } + + // 尝试解析为数组格式（兼容 antigravity2api-nodejs 的 accounts.json） + if let Ok(creds_array) = serde_json::from_str::>(&content) { + // 找到第一个启用的凭证 + if let Some(creds) = creds_array.into_iter().find(|c| c.enable != Some(false)) { + self.credentials = creds; + // 如果凭证中有 projectId，设置到 provider + if let Some(ref pid) = self.credentials.projectId { + self.project_id = Some(pid.clone()); + } + return Ok(()); + } + return Err("凭证文件中没有可用的账号（所有账号都被禁用）".into()); + } + + Err("无法解析凭证文件，请确保是有效的 JSON 格式".into()) } pub async fn save_credentials(&self) -> Result<(), Box> { @@ -200,6 +264,11 @@ impl AntigravityProvider { return false; } + // 检查是否被禁用 + if self.credentials.enable == Some(false) { + return false; + } + // 优先检查 RFC3339 格式的过期时间 if let Some(expire_str) = &self.credentials.expire { if let Ok(expires) = chrono::DateTime::parse_from_rfc3339(expire_str) { @@ -215,6 +284,17 @@ impl AntigravityProvider { // Token valid if more than 5 minutes until expiry return expiry > now + 300_000; } + + // 兼容 antigravity2api-nodejs 格式：timestamp + expires_in + if let (Some(timestamp), Some(expires_in)) = + (self.credentials.timestamp, self.credentials.expires_in) + { + let expiry = timestamp + (expires_in * 1000); + let now = chrono::Utc::now().timestamp_millis(); + // Token valid if more than 5 minutes until expiry + return expiry > now + 300_000; + } + true } @@ -234,6 +314,17 @@ impl AntigravityProvider { let refresh_skew_ms = REFRESH_SKEW * 1000; return expiry <= now + refresh_skew_ms; } + + // 兼容 antigravity2api-nodejs 格式：timestamp + expires_in + if let (Some(timestamp), Some(expires_in)) = + (self.credentials.timestamp, self.credentials.expires_in) + { + let expiry = timestamp + (expires_in * 1000); + let now = chrono::Utc::now().timestamp_millis(); + let refresh_skew_ms = REFRESH_SKEW * 1000; + return expiry <= now + refresh_skew_ms; + } + true } @@ -272,11 +363,18 @@ impl AntigravityProvider { self.credentials.access_token = Some(new_token.to_string()); - // 更新过期时间（同时保存两种格式以兼容） + // 更新过期时间（同时保存多种格式以兼容） if let Some(expires_in) = data["expires_in"].as_i64() { - let expires_at = chrono::Utc::now() + chrono::Duration::seconds(expires_in); + let now = chrono::Utc::now(); + let expires_at = now + chrono::Duration::seconds(expires_in); + + // RFC3339 格式 self.credentials.expire = Some(expires_at.to_rfc3339()); + // 毫秒时间戳格式 self.credentials.expiry_date = Some(expires_at.timestamp_millis()); + // antigravity2api-nodejs 格式 + self.credentials.expires_in = Some(expires_in); + self.credentials.timestamp = Some(now.timestamp_millis()); } // 如果返回了新的 refresh_token，也更新它 @@ -520,3 +618,952 @@ impl AntigravityProvider { self.available_models.iter().any(|m| m == model) } } + +// ============================================================================ +// OAuth 登录功能 +// ============================================================================ + +/// OAuth 回调结果 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct OAuthCallbackResult { + pub code: String, + pub state: String, +} + +/// OAuth 登录成功后的凭证信息 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct AntigravityOAuthResult { + pub credentials: AntigravityCredentials, + pub creds_file_path: String, +} + +/// 生成 OAuth 授权 URL +pub fn generate_auth_url(port: u16, state: &str) -> String { + let scopes = OAUTH_SCOPES.join(" "); + let redirect_uri = format!("http://localhost:{}/oauth-callback", port); + + let params = [ + ("access_type", "offline"), + ("client_id", OAUTH_CLIENT_ID), + ("prompt", "consent"), + ("redirect_uri", &redirect_uri), + ("response_type", "code"), + ("scope", &scopes), + ("state", state), + ]; + + let query = params + .iter() + .map(|(k, v)| format!("{}={}", k, urlencoding::encode(v))) + .collect::>() + .join("&"); + + format!("https://accounts.google.com/o/oauth2/v2/auth?{}", query) +} + +/// 用授权码交换 Token +pub async fn exchange_code_for_token( + client: &Client, + code: &str, + redirect_uri: &str, +) -> Result> { + let params = [ + ("code", code), + ("client_id", OAUTH_CLIENT_ID), + ("client_secret", OAUTH_CLIENT_SECRET), + ("redirect_uri", redirect_uri), + ("grant_type", "authorization_code"), + ]; + + let resp = client + .post("https://oauth2.googleapis.com/token") + .form(¶ms) + .send() + .await?; + + if !resp.status().is_success() { + let status = resp.status(); + let body = resp.text().await.unwrap_or_default(); + return Err(format!("Token 交换失败: {} - {}", status, body).into()); + } + + let data: serde_json::Value = resp.json().await?; + Ok(data) +} + +/// 获取用户邮箱 +pub async fn fetch_user_email( + client: &Client, + access_token: &str, +) -> Result, Box> { + let resp = client + .get("https://www.googleapis.com/oauth2/v2/userinfo") + .header("Authorization", format!("Bearer {}", access_token)) + .send() + .await?; + + if resp.status().is_success() { + let data: serde_json::Value = resp.json().await?; + Ok(data["email"].as_str().map(|s| s.to_string())) + } else { + Ok(None) + } +} + +/// 获取项目 ID（验证账号资格） +/// 返回值说明： +/// - Ok(Some(FetchedProjectId::HasProject(id))) - 有资格，且有 projectId +/// - Ok(Some(FetchedProjectId::NoProject)) - 有资格，但 projectId 为空（需要生成随机 ID） +/// - Ok(None) - 无资格（字段不存在，即 undefined） +/// - Err(_) - 请求失败 +#[derive(Debug, Clone)] +pub enum FetchedProjectId { + /// 有 projectId + HasProject(String), + /// projectId 为空字符串（有资格但无 projectId） + NoProject, +} + +pub async fn fetch_project_id_for_oauth( + client: &Client, + access_token: &str, +) -> Result, Box> { + tracing::info!("[Antigravity OAuth] 正在获取 projectId..."); + + let resp = client + .post("https://daily-cloudcode-pa.sandbox.googleapis.com/v1internal:loadCodeAssist") + .header("Authorization", format!("Bearer {}", access_token)) + .header("User-Agent", "antigravity/1.11.9 windows/amd64") + .header("Content-Type", "application/json") + .json(&serde_json::json!({ "metadata": { "ideType": "ANTIGRAVITY" } })) + .send() + .await?; + + let status = resp.status(); + tracing::info!("[Antigravity OAuth] loadCodeAssist 响应状态: {}", status); + + if status.is_success() { + let body_text = resp.text().await?; + tracing::info!("[Antigravity OAuth] loadCodeAssist 响应体: {}", body_text); + + let data: serde_json::Value = serde_json::from_str(&body_text)?; + // 检查字段是否存在 + // - 如果字段不存在（undefined）-> None（无资格） + // - 如果字段存在但为空字符串 -> Some(NoProject)（有资格但无 projectId） + // - 如果字段存在且有值 -> Some(HasProject(id))（有资格且有 projectId） + match data.get("cloudaicompanionProject") { + None => { + tracing::warn!("[Antigravity OAuth] cloudaicompanionProject 字段不存在"); + Ok(None) // 字段不存在，无资格 + } + Some(value) => { + if value.is_null() { + tracing::warn!("[Antigravity OAuth] cloudaicompanionProject 为 null"); + Ok(None) // null 也视为无资格 + } else if let Some(s) = value.as_str() { + if s.is_empty() { + tracing::info!("[Antigravity OAuth] cloudaicompanionProject 为空字符串，有资格但无 projectId"); + Ok(Some(FetchedProjectId::NoProject)) // 空字符串，有资格但无 projectId + } else { + tracing::info!("[Antigravity OAuth] 获取到 projectId: {}", s); + Ok(Some(FetchedProjectId::HasProject(s.to_string()))) // 有 projectId + } + } else { + tracing::warn!( + "[Antigravity OAuth] cloudaicompanionProject 不是字符串类型: {:?}", + value + ); + Ok(None) // 非字符串类型，视为无资格 + } + } + } + } else { + let body = resp.text().await.unwrap_or_default(); + tracing::error!( + "[Antigravity OAuth] loadCodeAssist 请求失败: {} - {}", + status, + body + ); + Err(format!("loadCodeAssist 请求失败: {} - {}", status, body).into()) + } +} + +/// OAuth 成功页面 HTML +const OAUTH_SUCCESS_HTML: &str = r#" + + + + 授权成功 + + + +
+

✓ 授权成功

+

账号已添加到 ProxyCast

+ +

可以关闭此页面

+
+ +"#; + +/// OAuth 失败页面 HTML +const OAUTH_ERROR_HTML: &str = r#" + + + + 授权失败 + + + +
+

✗ 授权失败

+

ERROR_PLACEHOLDER

+

请关闭此页面后重试

+
+ +"#; + +/// OAuth 授权 URL 结果 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct OAuthAuthUrlResult { + pub auth_url: String, + pub port: u16, + pub state: String, +} + +/// 启动 OAuth 服务器并返回授权 URL（不打开浏览器） +/// 服务器会在后台等待回调，成功后返回凭证 +pub async fn start_oauth_server_and_get_url( + skip_project_id_fetch: bool, +) -> Result< + ( + String, + impl std::future::Future>>, + ), + Box, +> { + use axum::{extract::Query, response::Html, routing::get, Router}; + use std::collections::HashMap; + use tokio::net::TcpListener; + + let client = Client::builder() + .timeout(std::time::Duration::from_secs(30)) + .build()?; + + // 生成随机 state + let state = Uuid::new_v4().to_string(); + let state_clone = state.clone(); + + // 创建 channel 用于接收回调结果 + let (tx, rx) = oneshot::channel::>(); + let tx = Arc::new(tokio::sync::Mutex::new(Some(tx))); + + // 绑定到随机端口 + let listener = TcpListener::bind("127.0.0.1:0").await?; + let port = listener.local_addr()?.port(); + + let redirect_uri = format!("http://localhost:{}/oauth-callback", port); + let redirect_uri_clone = redirect_uri.clone(); + + // 生成授权 URL + let auth_url = generate_auth_url(port, &state); + + tracing::info!( + "[Antigravity OAuth] 服务器启动在端口 {}, 授权 URL: {}", + port, + auth_url + ); + + // 构建路由 + let app = Router::new().route( + "/oauth-callback", + get(move |Query(params): Query>| { + let tx = tx.clone(); + let client = client.clone(); + let state_expected = state_clone.clone(); + let redirect_uri = redirect_uri_clone.clone(); + + async move { + let code = params.get("code"); + let returned_state = params.get("state"); + let error = params.get("error"); + + // 检查错误 + if let Some(err) = error { + let html = OAUTH_ERROR_HTML.replace("ERROR_PLACEHOLDER", err); + if let Some(sender) = tx.lock().await.take() { + let _ = sender.send(Err(format!("OAuth 错误: {}", err))); + } + return Html(html); + } + + // 检查 state + if returned_state.map(|s| s.as_str()) != Some(&state_expected) { + let html = OAUTH_ERROR_HTML.replace("ERROR_PLACEHOLDER", "State 验证失败"); + if let Some(sender) = tx.lock().await.take() { + let _ = sender.send(Err("State 验证失败".to_string())); + } + return Html(html); + } + + // 检查 code + let code = match code { + Some(c) => c, + None => { + let html = OAUTH_ERROR_HTML.replace("ERROR_PLACEHOLDER", "未收到授权码"); + if let Some(sender) = tx.lock().await.take() { + let _ = sender.send(Err("未收到授权码".to_string())); + } + return Html(html); + } + }; + + // 交换 Token + let token_result = exchange_code_for_token(&client, code, &redirect_uri).await; + let token_data = match token_result { + Ok(data) => data, + Err(e) => { + let html = OAUTH_ERROR_HTML.replace("ERROR_PLACEHOLDER", &e.to_string()); + if let Some(sender) = tx.lock().await.take() { + let _ = sender.send(Err(e.to_string())); + } + return Html(html); + } + }; + + let access_token = token_data["access_token"].as_str().unwrap_or_default(); + let refresh_token = token_data["refresh_token"].as_str().map(|s| s.to_string()); + let expires_in = token_data["expires_in"].as_i64(); + + // 获取用户邮箱 + let email = fetch_user_email(&client, access_token).await.ok().flatten(); + + // 获取项目 ID + let project_id = if skip_project_id_fetch { + tracing::info!("[Antigravity OAuth] 跳过 projectId 获取，使用随机生成的 ID"); + Some(generate_project_id()) + } else { + match fetch_project_id_for_oauth(&client, access_token).await { + Ok(Some(FetchedProjectId::HasProject(pid))) => Some(pid), + Ok(Some(FetchedProjectId::NoProject)) => { + tracing::info!("[Antigravity OAuth] projectId 为空，使用随机生成的 ID"); + Some(generate_project_id()) + } + Ok(None) => { + tracing::warn!("[Antigravity OAuth] 无法获取 projectId（字段不存在），使用随机生成的 ID"); + Some(generate_project_id()) + } + Err(e) => { + tracing::warn!("[Antigravity OAuth] 获取 projectId 失败: {}, 使用随机 ID", e); + Some(generate_project_id()) + } + } + }; + + // 构建凭证 + let now = chrono::Utc::now(); + let credentials = AntigravityCredentials { + access_token: Some(access_token.to_string()), + refresh_token, + token_type: Some("Bearer".to_string()), + expiry_date: expires_in.map(|e| (now + chrono::Duration::seconds(e)).timestamp_millis()), + expire: expires_in.map(|e| (now + chrono::Duration::seconds(e)).to_rfc3339()), + scope: Some(OAUTH_SCOPES.join(" ")), + last_refresh: Some(now.to_rfc3339()), + cred_type: "antigravity".to_string(), + expires_in, + timestamp: Some(now.timestamp_millis()), + enable: Some(true), + projectId: project_id, + email: email.clone(), + }; + + // 保存凭证到应用数据目录 + let creds_dir = dirs::data_dir() + .unwrap_or_else(|| PathBuf::from(".")) + .join("proxycast") + .join("credentials") + .join("antigravity"); + + if let Err(e) = std::fs::create_dir_all(&creds_dir) { + let html = OAUTH_ERROR_HTML.replace("ERROR_PLACEHOLDER", &format!("创建目录失败: {}", e)); + if let Some(sender) = tx.lock().await.take() { + let _ = sender.send(Err(e.to_string())); + } + return Html(html); + } + + // 使用 UUID 作为文件名 + let file_name = format!("{}.json", Uuid::new_v4()); + let creds_path = creds_dir.join(&file_name); + + let creds_json = match serde_json::to_string_pretty(&credentials) { + Ok(json) => json, + Err(e) => { + let html = OAUTH_ERROR_HTML.replace("ERROR_PLACEHOLDER", &format!("序列化失败: {}", e)); + if let Some(sender) = tx.lock().await.take() { + let _ = sender.send(Err(e.to_string())); + } + return Html(html); + } + }; + + if let Err(e) = std::fs::write(&creds_path, &creds_json) { + let html = OAUTH_ERROR_HTML.replace("ERROR_PLACEHOLDER", &format!("保存凭证失败: {}", e)); + if let Some(sender) = tx.lock().await.take() { + let _ = sender.send(Err(e.to_string())); + } + return Html(html); + } + + let creds_path_str = creds_path.to_string_lossy().to_string(); + tracing::info!("[Antigravity OAuth] 凭证已保存到: {}", creds_path_str); + + // 发送成功结果 + let result = AntigravityOAuthResult { + credentials, + creds_file_path: creds_path_str, + }; + + if let Some(sender) = tx.lock().await.take() { + let _ = sender.send(Ok(result)); + } + + let email_display = email.unwrap_or_else(|| "未知邮箱".to_string()); + let html = OAUTH_SUCCESS_HTML.replace("EMAIL_PLACEHOLDER", &email_display); + Html(html) + } + }), + ); + + // 创建等待回调的 Future + let wait_future = async move { + // 启动服务器 + let server = axum::serve(listener, app); + + // 同时运行服务器和等待回调结果 + tokio::select! { + result = async { + tokio::time::timeout( + std::time::Duration::from_secs(300), + rx + ).await + } => { + match result { + Ok(Ok(Ok(r))) => Ok(r), + Ok(Ok(Err(e))) => Err(e.into()), + Ok(Err(_)) => Err("OAuth 回调通道关闭".into()), + Err(_) => Err("OAuth 登录超时（5分钟）".into()), + } + } + server_result = server => { + match server_result { + Ok(_) => Err("服务器意外关闭".into()), + Err(e) => Err(format!("服务器错误: {}", e).into()), + } + } + } + }; + + Ok((auth_url, wait_future)) +} + +/// 启动 OAuth 登录流程（使用指定端口） +/// 用于配合 get_oauth_auth_url 使用 +pub async fn start_oauth_login_with_port( + port: u16, + state: String, + skip_project_id_fetch: bool, +) -> Result> { + use axum::{extract::Query, response::Html, routing::get, Router}; + use std::collections::HashMap; + use tokio::net::TcpListener; + + let client = Client::builder() + .timeout(std::time::Duration::from_secs(30)) + .build()?; + + let state_clone = state.clone(); + + // 创建 channel 用于接收回调结果 + let (tx, rx) = oneshot::channel::>(); + let tx = Arc::new(tokio::sync::Mutex::new(Some(tx))); + + // 绑定到指定端口 + let listener = TcpListener::bind(format!("127.0.0.1:{}", port)).await?; + + let redirect_uri = format!("http://localhost:{}/oauth-callback", port); + let redirect_uri_clone = redirect_uri.clone(); + + // 构建路由 + let app = Router::new().route( + "/oauth-callback", + get(move |Query(params): Query>| { + let tx = tx.clone(); + let client = client.clone(); + let state_expected = state_clone.clone(); + let redirect_uri = redirect_uri_clone.clone(); + + async move { + let code = params.get("code"); + let returned_state = params.get("state"); + let error = params.get("error"); + + // 检查错误 + if let Some(err) = error { + let html = OAUTH_ERROR_HTML.replace("ERROR_PLACEHOLDER", err); + if let Some(sender) = tx.lock().await.take() { + let _ = sender.send(Err(format!("OAuth 错误: {}", err))); + } + return Html(html); + } + + // 检查 state + if returned_state.map(|s| s.as_str()) != Some(&state_expected) { + let html = OAUTH_ERROR_HTML.replace("ERROR_PLACEHOLDER", "State 验证失败"); + if let Some(sender) = tx.lock().await.take() { + let _ = sender.send(Err("State 验证失败".to_string())); + } + return Html(html); + } + + // 检查 code + let code = match code { + Some(c) => c, + None => { + let html = OAUTH_ERROR_HTML.replace("ERROR_PLACEHOLDER", "未收到授权码"); + if let Some(sender) = tx.lock().await.take() { + let _ = sender.send(Err("未收到授权码".to_string())); + } + return Html(html); + } + }; + + // 交换 Token + let token_result = exchange_code_for_token(&client, code, &redirect_uri).await; + let token_data = match token_result { + Ok(data) => data, + Err(e) => { + let html = OAUTH_ERROR_HTML.replace("ERROR_PLACEHOLDER", &e.to_string()); + if let Some(sender) = tx.lock().await.take() { + let _ = sender.send(Err(e.to_string())); + } + return Html(html); + } + }; + + let access_token = token_data["access_token"].as_str().unwrap_or_default(); + let refresh_token = token_data["refresh_token"].as_str().map(|s| s.to_string()); + let expires_in = token_data["expires_in"].as_i64(); + + // 获取用户邮箱 + let email = fetch_user_email(&client, access_token).await.ok().flatten(); + + // 获取项目 ID + let project_id = if skip_project_id_fetch { + tracing::info!("[Antigravity OAuth] 跳过 projectId 获取，使用随机生成的 ID"); + Some(generate_project_id()) + } else { + match fetch_project_id_for_oauth(&client, access_token).await { + Ok(Some(FetchedProjectId::HasProject(pid))) => Some(pid), + Ok(Some(FetchedProjectId::NoProject)) => { + tracing::info!("[Antigravity OAuth] projectId 为空，使用随机生成的 ID"); + Some(generate_project_id()) + } + Ok(None) => { + tracing::warn!("[Antigravity OAuth] 无法获取 projectId（字段不存在），使用随机生成的 ID"); + Some(generate_project_id()) + } + Err(e) => { + tracing::warn!("[Antigravity OAuth] 获取 projectId 失败: {}, 使用随机 ID", e); + Some(generate_project_id()) + } + } + }; + + // 构建凭证 + let now = chrono::Utc::now(); + let credentials = AntigravityCredentials { + access_token: Some(access_token.to_string()), + refresh_token, + token_type: Some("Bearer".to_string()), + expiry_date: expires_in.map(|e| (now + chrono::Duration::seconds(e)).timestamp_millis()), + expire: expires_in.map(|e| (now + chrono::Duration::seconds(e)).to_rfc3339()), + scope: Some(OAUTH_SCOPES.join(" ")), + last_refresh: Some(now.to_rfc3339()), + cred_type: "antigravity".to_string(), + expires_in, + timestamp: Some(now.timestamp_millis()), + enable: Some(true), + projectId: project_id, + email: email.clone(), + }; + + // 保存凭证到应用数据目录 + let creds_dir = dirs::data_dir() + .unwrap_or_else(|| PathBuf::from(".")) + .join("proxycast") + .join("credentials") + .join("antigravity"); + + if let Err(e) = std::fs::create_dir_all(&creds_dir) { + let html = OAUTH_ERROR_HTML.replace("ERROR_PLACEHOLDER", &format!("创建目录失败: {}", e)); + if let Some(sender) = tx.lock().await.take() { + let _ = sender.send(Err(e.to_string())); + } + return Html(html); + } + + // 使用 UUID 作为文件名 + let file_name = format!("{}.json", Uuid::new_v4()); + let creds_path = creds_dir.join(&file_name); + + let creds_json = match serde_json::to_string_pretty(&credentials) { + Ok(json) => json, + Err(e) => { + let html = OAUTH_ERROR_HTML.replace("ERROR_PLACEHOLDER", &format!("序列化失败: {}", e)); + if let Some(sender) = tx.lock().await.take() { + let _ = sender.send(Err(e.to_string())); + } + return Html(html); + } + }; + + if let Err(e) = std::fs::write(&creds_path, &creds_json) { + let html = OAUTH_ERROR_HTML.replace("ERROR_PLACEHOLDER", &format!("保存凭证失败: {}", e)); + if let Some(sender) = tx.lock().await.take() { + let _ = sender.send(Err(e.to_string())); + } + return Html(html); + } + + let creds_path_str = creds_path.to_string_lossy().to_string(); + tracing::info!("[Antigravity OAuth] 凭证已保存到: {}", creds_path_str); + + // 发送成功结果 + let result = AntigravityOAuthResult { + credentials, + creds_file_path: creds_path_str, + }; + + if let Some(sender) = tx.lock().await.take() { + let _ = sender.send(Ok(result)); + } + + let email_display = email.unwrap_or_else(|| "未知邮箱".to_string()); + let html = OAUTH_SUCCESS_HTML.replace("EMAIL_PLACEHOLDER", &email_display); + Html(html) + } + }), + ); + + // 启动服务器 + let server = axum::serve(listener, app); + + // 同时运行服务器和等待回调结果 + tokio::select! { + result = async { + tokio::time::timeout( + std::time::Duration::from_secs(300), + rx + ).await + } => { + match result { + Ok(Ok(Ok(r))) => Ok(r), + Ok(Ok(Err(e))) => Err(e.into()), + Ok(Err(_)) => Err("OAuth 回调通道关闭".into()), + Err(_) => Err("OAuth 登录超时（5分钟）".into()), + } + } + server_result = server => { + match server_result { + Ok(_) => Err("服务器意外关闭".into()), + Err(e) => Err(format!("服务器错误: {}", e).into()), + } + } + } +} + +/// 启动 OAuth 登录流程 +/// 返回 (auth_url, credentials_file_path) +pub async fn start_oauth_login( + skip_project_id_fetch: bool, +) -> Result> { + use axum::{extract::Query, response::Html, routing::get, Router}; + use std::collections::HashMap; + use tokio::net::TcpListener; + + let client = Client::builder() + .timeout(std::time::Duration::from_secs(30)) + .build()?; + + // 生成随机 state + let state = Uuid::new_v4().to_string(); + let state_clone = state.clone(); + + // 创建 channel 用于接收回调结果 + let (tx, rx) = oneshot::channel::>(); + let tx = Arc::new(tokio::sync::Mutex::new(Some(tx))); + + // 绑定到随机端口 + let listener = TcpListener::bind("127.0.0.1:0").await?; + let port = listener.local_addr()?.port(); + + let redirect_uri = format!("http://localhost:{}/oauth-callback", port); + let redirect_uri_clone = redirect_uri.clone(); + + // 构建路由 + let app = Router::new().route( + "/oauth-callback", + get(move |Query(params): Query>| { + let tx = tx.clone(); + let client = client.clone(); + let state_expected = state_clone.clone(); + let redirect_uri = redirect_uri_clone.clone(); + + async move { + let code = params.get("code"); + let returned_state = params.get("state"); + let error = params.get("error"); + + // 检查错误 + if let Some(err) = error { + let html = OAUTH_ERROR_HTML.replace("ERROR_PLACEHOLDER", err); + if let Some(sender) = tx.lock().await.take() { + let _ = sender.send(Err(format!("OAuth 错误: {}", err))); + } + return Html(html); + } + + // 检查 state + if returned_state.map(|s| s.as_str()) != Some(&state_expected) { + let html = OAUTH_ERROR_HTML.replace("ERROR_PLACEHOLDER", "State 验证失败"); + if let Some(sender) = tx.lock().await.take() { + let _ = sender.send(Err("State 验证失败".to_string())); + } + return Html(html); + } + + // 检查 code + let code = match code { + Some(c) => c, + None => { + let html = OAUTH_ERROR_HTML.replace("ERROR_PLACEHOLDER", "未收到授权码"); + if let Some(sender) = tx.lock().await.take() { + let _ = sender.send(Err("未收到授权码".to_string())); + } + return Html(html); + } + }; + + // 交换 Token + let token_result = exchange_code_for_token(&client, code, &redirect_uri).await; + let token_data = match token_result { + Ok(data) => data, + Err(e) => { + let html = OAUTH_ERROR_HTML.replace("ERROR_PLACEHOLDER", &e.to_string()); + if let Some(sender) = tx.lock().await.take() { + let _ = sender.send(Err(e.to_string())); + } + return Html(html); + } + }; + + let access_token = token_data["access_token"].as_str().unwrap_or_default(); + let refresh_token = token_data["refresh_token"].as_str().map(|s| s.to_string()); + let expires_in = token_data["expires_in"].as_i64(); + + // 获取用户邮箱 + let email = fetch_user_email(&client, access_token).await.ok().flatten(); + + // 获取项目 ID + // 参考 antigravity2api-nodejs 的逻辑： + // - projectId === undefined -> 无资格（但我们改为使用随机 ID，因为很多账号都没有 projectId） + // - projectId === "" -> 有资格但无 projectId，使用随机生成的 + // - projectId 有值 -> 有资格且有 projectId + let project_id = if skip_project_id_fetch { + tracing::info!("[Antigravity OAuth] 跳过 projectId 获取，使用随机生成的 ID"); + Some(generate_project_id()) + } else { + match fetch_project_id_for_oauth(&client, access_token).await { + Ok(Some(FetchedProjectId::HasProject(pid))) => { + // 有资格且有 projectId + Some(pid) + } + Ok(Some(FetchedProjectId::NoProject)) => { + // 有资格但 projectId 为空，使用随机生成的 + tracing::info!("[Antigravity OAuth] projectId 为空，使用随机生成的 ID"); + Some(generate_project_id()) + } + Ok(None) => { + // 字段不存在，也使用随机 ID（很多账号都是这种情况） + tracing::warn!("[Antigravity OAuth] 无法获取 projectId（字段不存在），使用随机生成的 ID"); + Some(generate_project_id()) + } + Err(e) => { + tracing::warn!("[Antigravity OAuth] 获取 projectId 失败: {}, 使用随机 ID", e); + Some(generate_project_id()) + } + } + }; + + // 构建凭证 + let now = chrono::Utc::now(); + let mut credentials = AntigravityCredentials { + access_token: Some(access_token.to_string()), + refresh_token, + token_type: Some("Bearer".to_string()), + expiry_date: expires_in.map(|e| (now + chrono::Duration::seconds(e)).timestamp_millis()), + expire: expires_in.map(|e| (now + chrono::Duration::seconds(e)).to_rfc3339()), + scope: Some(OAUTH_SCOPES.join(" ")), + last_refresh: Some(now.to_rfc3339()), + cred_type: "antigravity".to_string(), + expires_in, + timestamp: Some(now.timestamp_millis()), + enable: Some(true), + projectId: project_id, + email: email.clone(), + }; + + // 保存凭证到应用数据目录 + let creds_dir = dirs::data_dir() + .unwrap_or_else(|| PathBuf::from(".")) + .join("proxycast") + .join("credentials") + .join("antigravity"); + + if let Err(e) = std::fs::create_dir_all(&creds_dir) { + let html = OAUTH_ERROR_HTML.replace("ERROR_PLACEHOLDER", &format!("创建目录失败: {}", e)); + if let Some(sender) = tx.lock().await.take() { + let _ = sender.send(Err(e.to_string())); + } + return Html(html); + } + + // 使用 UUID 作为文件名 + let file_name = format!("{}.json", Uuid::new_v4()); + let creds_path = creds_dir.join(&file_name); + + let creds_json = match serde_json::to_string_pretty(&credentials) { + Ok(json) => json, + Err(e) => { + let html = OAUTH_ERROR_HTML.replace("ERROR_PLACEHOLDER", &format!("序列化失败: {}", e)); + if let Some(sender) = tx.lock().await.take() { + let _ = sender.send(Err(e.to_string())); + } + return Html(html); + } + }; + + if let Err(e) = std::fs::write(&creds_path, &creds_json) { + let html = OAUTH_ERROR_HTML.replace("ERROR_PLACEHOLDER", &format!("保存凭证失败: {}", e)); + if let Some(sender) = tx.lock().await.take() { + let _ = sender.send(Err(e.to_string())); + } + return Html(html); + } + + let creds_path_str = creds_path.to_string_lossy().to_string(); + tracing::info!("[Antigravity OAuth] 凭证已保存到: {}", creds_path_str); + + // 发送成功结果 + let result = AntigravityOAuthResult { + credentials, + creds_file_path: creds_path_str, + }; + + if let Some(sender) = tx.lock().await.take() { + let _ = sender.send(Ok(result)); + } + + let email_display = email.unwrap_or_else(|| "未知邮箱".to_string()); + let html = OAUTH_SUCCESS_HTML.replace("EMAIL_PLACEHOLDER", &email_display); + Html(html) + } + }), + ); + + // 生成授权 URL + let auth_url = generate_auth_url(port, &state); + + // 打开浏览器 + tracing::info!("[Antigravity OAuth] 打开浏览器进行授权: {}", auth_url); + if let Err(e) = open::that(&auth_url) { + tracing::warn!("[Antigravity OAuth] 无法自动打开浏览器: {}", e); + } + + // 启动服务器 + let server = axum::serve(listener, app); + + // 同时运行服务器和等待回调结果 + tokio::select! { + // 等待回调结果（带超时） + result = async { + tokio::time::timeout( + std::time::Duration::from_secs(300), + rx + ).await + } => { + match result { + Ok(Ok(Ok(r))) => Ok(r), + Ok(Ok(Err(e))) => Err(e.into()), + Ok(Err(_)) => Err("OAuth 回调通道关闭".into()), + Err(_) => Err("OAuth 登录超时（5分钟）".into()), + } + } + // 服务器运行（不会主动结束，除非出错） + server_result = server => { + match server_result { + Ok(_) => Err("服务器意外关闭".into()), + Err(e) => Err(format!("服务器错误: {}", e).into()), + } + } + } +} + +// ============================================================================ +// CredentialProvider Trait 实现 +// ============================================================================ + +#[async_trait] +impl CredentialProvider for AntigravityProvider { + async fn load_credentials_from_path(&mut self, path: &str) -> ProviderResult<()> { + AntigravityProvider::load_credentials_from_path(self, path).await + } + + async fn save_credentials(&self) -> ProviderResult<()> { + AntigravityProvider::save_credentials(self).await + } + + fn is_token_valid(&self) -> bool { + AntigravityProvider::is_token_valid(self) + } + + fn is_token_expiring_soon(&self) -> bool { + AntigravityProvider::is_token_expiring_soon(self) + } + + async fn refresh_token(&mut self) -> ProviderResult { + AntigravityProvider::refresh_token(self).await + } + + fn get_access_token(&self) -> Option<&str> { + self.credentials.access_token.as_deref() + } + + fn provider_type(&self) -> &'static str { + "antigravity" + } +} diff --git a/src-tauri/src/providers/claude_oauth.rs b/src-tauri/src/providers/claude_oauth.rs index 8ce364f55..b00434bce 100644 --- a/src-tauri/src/providers/claude_oauth.rs +++ b/src-tauri/src/providers/claude_oauth.rs @@ -347,3 +347,366 @@ impl ClaudeOAuthProvider { format!("http://localhost:{}/callback", self.callback_port) } } + +// ============================================================================ +// OAuth 登录功能 +// ============================================================================ + +use std::sync::Arc; +use tokio::sync::oneshot; +use uuid::Uuid; + +/// OAuth 登录成功后的凭证信息 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ClaudeOAuthResult { + pub credentials: ClaudeOAuthCredentials, + pub creds_file_path: String, +} + +/// 生成 Claude OAuth 授权 URL +pub fn generate_claude_auth_url(port: u16, state: &str, code_challenge: &str) -> String { + let redirect_uri = format!("http://localhost:{}/oauth-callback", port); + + let params = [ + ("client_id", CLAUDE_CLIENT_ID), + ("response_type", "code"), + ("redirect_uri", redirect_uri.as_str()), + ("scope", "user:inference user:profile"), + ("state", state), + ("code_challenge", code_challenge), + ("code_challenge_method", "S256"), + ]; + + let query = params + .iter() + .map(|(k, v)| format!("{}={}", k, urlencoding::encode(v))) + .collect::>() + .join("&"); + + format!("{}?{}", CLAUDE_AUTH_URL, query) +} + +/// 用授权码交换 Token +pub async fn exchange_claude_code_for_token( + client: &Client, + code: &str, + code_verifier: &str, + redirect_uri: &str, +) -> Result> { + let body = serde_json::json!({ + "grant_type": "authorization_code", + "client_id": CLAUDE_CLIENT_ID, + "code": code, + "redirect_uri": redirect_uri, + "code_verifier": code_verifier + }); + + let resp = client + .post(CLAUDE_TOKEN_URL) + .header("Content-Type", "application/json") + .header("Accept", "application/json") + .json(&body) + .send() + .await?; + + if !resp.status().is_success() { + let status = resp.status(); + let body = resp.text().await.unwrap_or_default(); + return Err(format!("Token 交换失败: {} - {}", status, body).into()); + } + + let data: serde_json::Value = resp.json().await?; + Ok(data) +} + +/// OAuth 成功页面 HTML +const CLAUDE_OAUTH_SUCCESS_HTML: &str = r#" + + + + 授权成功 + + + +
+

✓ 授权成功

+

Claude 账号已添加到 ProxyCast

+ +

可以关闭此页面

+
+ +"#; + +/// OAuth 失败页面 HTML +const CLAUDE_OAUTH_ERROR_HTML: &str = r#" + + + + 授权失败 + + + +
+

✗ 授权失败

+

ERROR_PLACEHOLDER

+

请关闭此页面后重试

+
+ +"#; + +/// 启动 OAuth 服务器并返回授权 URL（不打开浏览器） +pub async fn start_claude_oauth_server_and_get_url() -> Result< + ( + String, + impl std::future::Future>>, + ), + Box, +> { + use axum::{extract::Query, response::Html, routing::get, Router}; + use std::collections::HashMap; + use tokio::net::TcpListener; + + let client = Client::builder() + .timeout(std::time::Duration::from_secs(30)) + .build()?; + + // 生成 PKCE codes + let pkce_codes = PKCECodes::generate()?; + let code_verifier = pkce_codes.code_verifier.clone(); + let code_challenge = pkce_codes.code_challenge.clone(); + + // 生成随机 state + let state = Uuid::new_v4().to_string(); + let state_clone = state.clone(); + + // 创建 channel 用于接收回调结果 + let (tx, rx) = oneshot::channel::>(); + let tx = Arc::new(tokio::sync::Mutex::new(Some(tx))); + + // 绑定到随机端口 + let listener = TcpListener::bind("127.0.0.1:0").await?; + let port = listener.local_addr()?.port(); + + let redirect_uri = format!("http://localhost:{}/oauth-callback", port); + let redirect_uri_clone = redirect_uri.clone(); + + // 生成授权 URL + let auth_url = generate_claude_auth_url(port, &state, &code_challenge); + + tracing::info!( + "[Claude OAuth] 服务器启动在端口 {}, 授权 URL: {}", + port, + auth_url + ); + + // 构建路由 + let app = Router::new().route( + "/oauth-callback", + get(move |Query(params): Query>| { + let tx = tx.clone(); + let client = client.clone(); + let state_expected = state_clone.clone(); + let redirect_uri = redirect_uri_clone.clone(); + let code_verifier = code_verifier.clone(); + + async move { + let code = params.get("code"); + let returned_state = params.get("state"); + let error = params.get("error"); + + // 检查错误 + if let Some(err) = error { + let html = CLAUDE_OAUTH_ERROR_HTML.replace("ERROR_PLACEHOLDER", err); + if let Some(sender) = tx.lock().await.take() { + let _ = sender.send(Err(format!("OAuth 错误: {}", err))); + } + return Html(html); + } + + // 检查 state + if returned_state.map(|s| s.as_str()) != Some(&state_expected) { + let html = + CLAUDE_OAUTH_ERROR_HTML.replace("ERROR_PLACEHOLDER", "State 验证失败"); + if let Some(sender) = tx.lock().await.take() { + let _ = sender.send(Err("State 验证失败".to_string())); + } + return Html(html); + } + + // 检查 code + let code = match code { + Some(c) => c, + None => { + let html = + CLAUDE_OAUTH_ERROR_HTML.replace("ERROR_PLACEHOLDER", "未收到授权码"); + if let Some(sender) = tx.lock().await.take() { + let _ = sender.send(Err("未收到授权码".to_string())); + } + return Html(html); + } + }; + + // 交换 Token + let token_result = + exchange_claude_code_for_token(&client, code, &code_verifier, &redirect_uri) + .await; + let token_data = match token_result { + Ok(data) => data, + Err(e) => { + let html = + CLAUDE_OAUTH_ERROR_HTML.replace("ERROR_PLACEHOLDER", &e.to_string()); + if let Some(sender) = tx.lock().await.take() { + let _ = sender.send(Err(e.to_string())); + } + return Html(html); + } + }; + + let access_token = token_data["access_token"].as_str().unwrap_or_default(); + let refresh_token = token_data["refresh_token"].as_str().map(|s| s.to_string()); + let expires_in = token_data["expires_in"].as_i64(); + + // 从响应中提取用户邮箱 + let email = token_data["account"]["email_address"] + .as_str() + .map(|s| s.to_string()); + + // 构建凭证 + let now = chrono::Utc::now(); + let credentials = ClaudeOAuthCredentials { + access_token: Some(access_token.to_string()), + refresh_token, + email: email.clone(), + expire: expires_in.map(|e| (now + chrono::Duration::seconds(e)).to_rfc3339()), + last_refresh: Some(now.to_rfc3339()), + cred_type: "claude_oauth".to_string(), + }; + + // 保存凭证到应用数据目录 + let creds_dir = dirs::data_dir() + .unwrap_or_else(|| std::path::PathBuf::from(".")) + .join("proxycast") + .join("credentials") + .join("claude_oauth"); + + if let Err(e) = std::fs::create_dir_all(&creds_dir) { + let html = CLAUDE_OAUTH_ERROR_HTML + .replace("ERROR_PLACEHOLDER", &format!("创建目录失败: {}", e)); + if let Some(sender) = tx.lock().await.take() { + let _ = sender.send(Err(format!("创建目录失败: {}", e))); + } + return Html(html); + } + + // 生成唯一文件名 + let uuid = Uuid::new_v4().to_string(); + let timestamp = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_secs(); + let filename = format!("claude_oauth_{}_{}.json", &uuid[..8], timestamp); + let creds_file_path = creds_dir.join(&filename); + + // 保存凭证 + let creds_json = match serde_json::to_string_pretty(&credentials) { + Ok(json) => json, + Err(e) => { + let html = CLAUDE_OAUTH_ERROR_HTML + .replace("ERROR_PLACEHOLDER", &format!("序列化凭证失败: {}", e)); + if let Some(sender) = tx.lock().await.take() { + let _ = sender.send(Err(format!("序列化凭证失败: {}", e))); + } + return Html(html); + } + }; + + if let Err(e) = std::fs::write(&creds_file_path, &creds_json) { + let html = CLAUDE_OAUTH_ERROR_HTML + .replace("ERROR_PLACEHOLDER", &format!("保存凭证失败: {}", e)); + if let Some(sender) = tx.lock().await.take() { + let _ = sender.send(Err(format!("保存凭证失败: {}", e))); + } + return Html(html); + } + + tracing::info!("[Claude OAuth] 凭证已保存到: {:?}", creds_file_path); + + // 发送成功结果 + let result = ClaudeOAuthResult { + credentials, + creds_file_path: creds_file_path.to_string_lossy().to_string(), + }; + + if let Some(sender) = tx.lock().await.take() { + let _ = sender.send(Ok(result)); + } + + // 返回成功页面 + let html = CLAUDE_OAUTH_SUCCESS_HTML.replace( + "EMAIL_PLACEHOLDER", + &email.unwrap_or_else(|| "未知邮箱".to_string()), + ); + Html(html) + } + }), + ); + + // 启动服务器 + let server = axum::serve(listener, app); + + // 创建等待 future + let wait_future = async move { + // 设置超时（5 分钟） + let timeout = tokio::time::timeout(std::time::Duration::from_secs(300), async { + // 启动服务器（在后台运行） + tokio::spawn(async move { + if let Err(e) = server.await { + tracing::error!("[Claude OAuth] 服务器错误: {}", e); + } + }); + + // 等待回调结果 + match rx.await { + Ok(result) => result.map_err(|e| { + Box::new(std::io::Error::new(std::io::ErrorKind::Other, e)) + as Box + }), + Err(_) => Err("OAuth 回调通道关闭".into()), + } + }); + + match timeout.await { + Ok(result) => result, + Err(_) => Err("OAuth 登录超时（5分钟）".into()), + } + }; + + Ok((auth_url, wait_future)) +} + +/// 启动 Claude OAuth 登录流程（自动打开浏览器） +pub async fn start_claude_oauth_login() -> Result> { + let (auth_url, wait_future) = start_claude_oauth_server_and_get_url().await?; + + tracing::info!("[Claude OAuth] 打开浏览器进行授权: {}", auth_url); + + // 打开浏览器 + if let Err(e) = open::that(&auth_url) { + tracing::warn!("[Claude OAuth] 无法打开浏览器: {}. 请手动打开 URL.", e); + } + + // 等待回调 + wait_future.await +} diff --git a/src-tauri/src/providers/codex.rs b/src-tauri/src/providers/codex.rs index 64761b639..4d5b9f647 100644 --- a/src-tauri/src/providers/codex.rs +++ b/src-tauri/src/providers/codex.rs @@ -1047,6 +1047,14 @@ impl CodexProvider { // Build the Codex API URL let url = match mode { AuthMode::ApiKey => { + let has_custom_base_url = self + .credentials + .api_base_url + .as_deref() + .map(|s| s.trim()) + .filter(|s| !s.is_empty()) + .is_some(); + let base_url = self .credentials .api_base_url @@ -1054,6 +1062,17 @@ impl CodexProvider { .map(|s| s.trim()) .filter(|s| !s.is_empty()) .unwrap_or(DEFAULT_API_BASE_URL); + + // Warn if API key doesn't look like OpenAI format but no custom base URL is set + if !has_custom_base_url && !token.starts_with("sk-") { + tracing::warn!( + "[CODEX] API key does not appear to be an OpenAI key (doesn't start with 'sk-'), \ + but no api_base_url is configured. Requests will be sent to {}. \ + If you're using a third-party API provider, please add 'api_base_url' to ~/.codex/auth.json", + DEFAULT_API_BASE_URL + ); + } + Self::build_responses_url(base_url) } AuthMode::OAuth => format!("{}/responses", CODEX_API_BASE_URL), @@ -1692,3 +1711,404 @@ mod tests { assert_eq!(result.unwrap(), "expired_access_token"); } } + +// ============================================================================ +// OAuth 登录功能（参考 Antigravity 实现） +// ============================================================================ + +use std::sync::Arc; +use tokio::sync::oneshot; +use uuid::Uuid; + +/// OAuth 登录成功后的凭证信息 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct CodexOAuthResult { + pub credentials: CodexCredentials, + pub creds_file_path: String, +} + +/// OpenAI OAuth 固定回调端口（必须与 client_id 注册的回调地址一致） +const OPENAI_OAUTH_CALLBACK_PORT: u16 = 1455; + +/// OpenAI OAuth 固定回调路径（必须与 client_id 注册的回调地址一致） +const OPENAI_OAUTH_CALLBACK_PATH: &str = "/auth/callback"; + +/// 生成 OAuth 授权 URL（用于外部浏览器登录） +/// +/// 注意：OpenAI OAuth 要求 redirect_uri 必须是预先注册的固定地址 +/// Codex CLI 的 client_id 只注册了 http://localhost:1455/auth/callback +pub fn generate_codex_auth_url(state: &str, code_challenge: &str) -> String { + let redirect_uri = format!( + "http://localhost:{}{}", + OPENAI_OAUTH_CALLBACK_PORT, OPENAI_OAUTH_CALLBACK_PATH + ); + + let params = [ + ("client_id", OPENAI_CLIENT_ID), + ("response_type", "code"), + ("redirect_uri", redirect_uri.as_str()), + ("scope", "openid email profile offline_access"), + ("state", state), + ("code_challenge", code_challenge), + ("code_challenge_method", "S256"), + ("prompt", "login"), + ("id_token_add_organizations", "true"), + ("codex_cli_simplified_flow", "true"), + ]; + + let query = params + .iter() + .map(|(k, v)| format!("{}={}", k, urlencoding::encode(v))) + .collect::>() + .join("&"); + + format!("{}?{}", OPENAI_AUTH_URL, query) +} + +/// 用授权码交换 Token +pub async fn exchange_codex_code_for_token( + client: &Client, + code: &str, + code_verifier: &str, + redirect_uri: &str, +) -> Result> { + let params = [ + ("grant_type", "authorization_code"), + ("client_id", OPENAI_CLIENT_ID), + ("code", code), + ("redirect_uri", redirect_uri), + ("code_verifier", code_verifier), + ]; + + let resp = client + .post(OPENAI_TOKEN_URL) + .header("Content-Type", "application/x-www-form-urlencoded") + .header("Accept", "application/json") + .form(¶ms) + .send() + .await?; + + if !resp.status().is_success() { + let status = resp.status(); + let body = resp.text().await.unwrap_or_default(); + return Err(format!("Token 交换失败: {} - {}", status, body).into()); + } + + let data: serde_json::Value = resp.json().await?; + Ok(data) +} + +/// OAuth 成功页面 HTML +const CODEX_OAUTH_SUCCESS_HTML: &str = r#" + + + + 授权成功 + + + +
+

✓ 授权成功

+

Codex 账号已添加到 ProxyCast

+ +

可以关闭此页面

+
+ +"#; + +/// OAuth 失败页面 HTML +const CODEX_OAUTH_ERROR_HTML: &str = r#" + + + + 授权失败 + + + +
+

✗ 授权失败

+

ERROR_PLACEHOLDER

+

请关闭此页面后重试

+
+ +"#; + +/// 启动 OAuth 服务器并返回授权 URL（不打开浏览器） +/// 服务器会在后台等待回调，成功后返回凭证 +/// +/// 注意：OpenAI OAuth 要求使用固定的回调地址 http://localhost:1455/auth/callback +pub async fn start_codex_oauth_server_and_get_url() -> Result< + ( + String, + impl std::future::Future>>, + ), + Box, +> { + use axum::{extract::Query, response::Html, routing::get, Router}; + use std::collections::HashMap; + use tokio::net::TcpListener; + + let client = Client::builder() + .timeout(std::time::Duration::from_secs(30)) + .build()?; + + // 生成 PKCE codes + let pkce_codes = PKCECodes::generate()?; + let code_verifier = pkce_codes.code_verifier.clone(); + let code_challenge = pkce_codes.code_challenge.clone(); + + // 生成随机 state + let state = Uuid::new_v4().to_string(); + let state_clone = state.clone(); + + // 创建 channel 用于接收回调结果 + let (tx, rx) = oneshot::channel::>(); + let tx = Arc::new(tokio::sync::Mutex::new(Some(tx))); + + // 使用固定端口 1455（OpenAI OAuth 要求） + let port = OPENAI_OAUTH_CALLBACK_PORT; + let listener = TcpListener::bind(format!("127.0.0.1:{}", port)).await.map_err(|e| { + if e.kind() == std::io::ErrorKind::AddrInUse { + format!( + "端口 {} 已被占用。OpenAI OAuth 要求使用固定端口 1455，请关闭占用该端口的应用后重试。", + port + ) + } else { + format!("绑定端口 {} 失败: {}", port, e) + } + })?; + + let redirect_uri = format!( + "http://localhost:{}{}", + OPENAI_OAUTH_CALLBACK_PORT, OPENAI_OAUTH_CALLBACK_PATH + ); + let redirect_uri_clone = redirect_uri.clone(); + + // 生成授权 URL（不再传入 port 参数） + let auth_url = generate_codex_auth_url(&state, &code_challenge); + + tracing::info!( + "[Codex OAuth] 服务器启动在端口 {}, 授权 URL: {}", + port, + auth_url + ); + + // 构建路由（使用固定的回调路径 /auth/callback） + let app = Router::new().route( + OPENAI_OAUTH_CALLBACK_PATH, + get(move |Query(params): Query>| { + let tx = tx.clone(); + let client = client.clone(); + let state_expected = state_clone.clone(); + let redirect_uri = redirect_uri_clone.clone(); + let code_verifier = code_verifier.clone(); + + async move { + let code = params.get("code"); + let returned_state = params.get("state"); + let error = params.get("error"); + + // 检查错误 + if let Some(err) = error { + let html = CODEX_OAUTH_ERROR_HTML.replace("ERROR_PLACEHOLDER", err); + if let Some(sender) = tx.lock().await.take() { + let _ = sender.send(Err(format!("OAuth 错误: {}", err))); + } + return Html(html); + } + + // 检查 state + if returned_state.map(|s| s.as_str()) != Some(&state_expected) { + let html = + CODEX_OAUTH_ERROR_HTML.replace("ERROR_PLACEHOLDER", "State 验证失败"); + if let Some(sender) = tx.lock().await.take() { + let _ = sender.send(Err("State 验证失败".to_string())); + } + return Html(html); + } + + // 检查 code + let code = match code { + Some(c) => c, + None => { + let html = + CODEX_OAUTH_ERROR_HTML.replace("ERROR_PLACEHOLDER", "未收到授权码"); + if let Some(sender) = tx.lock().await.take() { + let _ = sender.send(Err("未收到授权码".to_string())); + } + return Html(html); + } + }; + + // 交换 Token + let token_result = + exchange_codex_code_for_token(&client, code, &code_verifier, &redirect_uri) + .await; + let token_data = match token_result { + Ok(data) => data, + Err(e) => { + let html = + CODEX_OAUTH_ERROR_HTML.replace("ERROR_PLACEHOLDER", &e.to_string()); + if let Some(sender) = tx.lock().await.take() { + let _ = sender.send(Err(e.to_string())); + } + return Html(html); + } + }; + + let access_token = token_data["access_token"].as_str().unwrap_or_default(); + let refresh_token = token_data["refresh_token"].as_str().map(|s| s.to_string()); + let id_token = token_data["id_token"].as_str().map(|s| s.to_string()); + let expires_in = token_data["expires_in"].as_i64(); + + // 解析 ID Token 获取用户信息 + let (account_id, email) = if let Some(ref id_token) = id_token { + parse_jwt_claims(id_token) + } else { + (None, None) + }; + + // 构建凭证 + let now = chrono::Utc::now(); + let credentials = CodexCredentials { + id_token, + access_token: Some(access_token.to_string()), + refresh_token, + api_key: None, + api_base_url: None, + account_id, + last_refresh: Some(now.to_rfc3339()), + email: email.clone(), + r#type: "codex".to_string(), + expires_at: expires_in + .map(|e| (now + chrono::Duration::seconds(e)).to_rfc3339()), + }; + + // 保存凭证到应用数据目录 + let creds_dir = dirs::data_dir() + .unwrap_or_else(|| PathBuf::from(".")) + .join("proxycast") + .join("credentials") + .join("codex"); + + if let Err(e) = std::fs::create_dir_all(&creds_dir) { + let html = CODEX_OAUTH_ERROR_HTML + .replace("ERROR_PLACEHOLDER", &format!("创建目录失败: {}", e)); + if let Some(sender) = tx.lock().await.take() { + let _ = sender.send(Err(format!("创建目录失败: {}", e))); + } + return Html(html); + } + + // 生成唯一文件名 + let uuid = Uuid::new_v4().to_string(); + let timestamp = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_secs(); + let filename = format!("codex_{}_{}.json", &uuid[..8], timestamp); + let creds_file_path = creds_dir.join(&filename); + + // 保存凭证 + let creds_json = match serde_json::to_string_pretty(&credentials) { + Ok(json) => json, + Err(e) => { + let html = CODEX_OAUTH_ERROR_HTML + .replace("ERROR_PLACEHOLDER", &format!("序列化凭证失败: {}", e)); + if let Some(sender) = tx.lock().await.take() { + let _ = sender.send(Err(format!("序列化凭证失败: {}", e))); + } + return Html(html); + } + }; + + if let Err(e) = std::fs::write(&creds_file_path, &creds_json) { + let html = CODEX_OAUTH_ERROR_HTML + .replace("ERROR_PLACEHOLDER", &format!("保存凭证失败: {}", e)); + if let Some(sender) = tx.lock().await.take() { + let _ = sender.send(Err(format!("保存凭证失败: {}", e))); + } + return Html(html); + } + + tracing::info!("[Codex OAuth] 凭证已保存到: {:?}", creds_file_path); + + // 发送成功结果 + let result = CodexOAuthResult { + credentials, + creds_file_path: creds_file_path.to_string_lossy().to_string(), + }; + + if let Some(sender) = tx.lock().await.take() { + let _ = sender.send(Ok(result)); + } + + // 返回成功页面 + let html = CODEX_OAUTH_SUCCESS_HTML.replace( + "EMAIL_PLACEHOLDER", + &email.unwrap_or_else(|| "未知邮箱".to_string()), + ); + Html(html) + } + }), + ); + + // 启动服务器 + let server = axum::serve(listener, app); + + // 创建等待 future + let wait_future = async move { + // 设置超时（5 分钟） + let timeout = tokio::time::timeout(std::time::Duration::from_secs(300), async { + // 启动服务器（在后台运行） + tokio::spawn(async move { + if let Err(e) = server.await { + tracing::error!("[Codex OAuth] 服务器错误: {}", e); + } + }); + + // 等待回调结果 + match rx.await { + Ok(result) => result.map_err(|e| { + Box::new(std::io::Error::new(std::io::ErrorKind::Other, e)) + as Box + }), + Err(_) => Err("OAuth 回调通道关闭".into()), + } + }); + + match timeout.await { + Ok(result) => result, + Err(_) => Err("OAuth 登录超时（5分钟）".into()), + } + }; + + Ok((auth_url, wait_future)) +} + +/// 启动 Codex OAuth 登录流程（自动打开浏览器） +pub async fn start_codex_oauth_login() -> Result> { + let (auth_url, wait_future) = start_codex_oauth_server_and_get_url().await?; + + tracing::info!("[Codex OAuth] 打开浏览器进行授权: {}", auth_url); + + // 打开浏览器 + if let Err(e) = open::that(&auth_url) { + tracing::warn!("[Codex OAuth] 无法打开浏览器: {}. 请手动打开 URL.", e); + } + + // 等待回调 + wait_future.await +} diff --git a/src-tauri/src/providers/gemini.rs b/src-tauri/src/providers/gemini.rs index 52892842f..04c27477f 100644 --- a/src-tauri/src/providers/gemini.rs +++ b/src-tauri/src/providers/gemini.rs @@ -6,6 +6,8 @@ use super::error::{ create_auth_error, create_config_error, create_token_refresh_error, ProviderError, }; +use super::traits::{CredentialProvider, ProviderResult}; +use async_trait::async_trait; use reqwest::Client; use serde::{Deserialize, Serialize}; use std::error::Error; @@ -747,3 +749,661 @@ mod gemini_api_key_tests { let _ = provider; } } + +// ============================================================================ +// Gemini OAuth 登录功能 +// ============================================================================ + +use std::sync::Arc; +use tokio::sync::oneshot; +use uuid::Uuid; + +// Gemini CLI OAuth 配置 - 与 claude-relay-service 对齐 +pub const GEMINI_OAUTH_CLIENT_ID: &str = + "681255809395-oo8ft2oprdrnp9e3aqf6av3hmdib135j.apps.googleusercontent.com"; +pub const GEMINI_OAUTH_CLIENT_SECRET: &str = "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"; +pub const GEMINI_OAUTH_SCOPES: &[&str] = &["https://www.googleapis.com/auth/cloud-platform"]; +pub const GEMINI_OAUTH_REDIRECT_URI: &str = "https://codeassist.google.com/authcode"; + +/// OAuth 登录成功后的凭证信息 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct GeminiOAuthResult { + pub credentials: GeminiCredentials, + pub creds_file_path: String, +} + +/// 生成 PKCE code_verifier 和 code_challenge +fn generate_pkce() -> (String, String) { + use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine}; + use sha2::{Digest, Sha256}; + + // 生成 43-128 字符的随机字符串作为 code_verifier + let code_verifier: String = (0..64) + .map(|_| { + let idx = rand::random::() % 66; + let chars = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-._~"; + chars[idx as usize] as char + }) + .collect(); + + // 计算 code_challenge = BASE64URL(SHA256(code_verifier)) + let mut hasher = Sha256::new(); + hasher.update(code_verifier.as_bytes()); + let hash = hasher.finalize(); + let code_challenge = URL_SAFE_NO_PAD.encode(hash); + + (code_verifier, code_challenge) +} + +/// 生成 OAuth 授权 URL（使用 PKCE） +pub fn generate_gemini_auth_url(state: &str, code_challenge: &str) -> String { + let scopes = GEMINI_OAUTH_SCOPES.join(" "); + + let params = [ + ("access_type", "offline"), + ("client_id", GEMINI_OAUTH_CLIENT_ID), + ("code_challenge", code_challenge), + ("code_challenge_method", "S256"), + ("prompt", "select_account"), + ("redirect_uri", GEMINI_OAUTH_REDIRECT_URI), + ("response_type", "code"), + ("scope", &scopes), + ("state", state), + ]; + + let query = params + .iter() + .map(|(k, v)| format!("{}={}", k, urlencoding::encode(v))) + .collect::>() + .join("&"); + + format!("https://accounts.google.com/o/oauth2/v2/auth?{}", query) +} + +/// Gemini OAuth 会话信息（用于存储 PKCE code_verifier） +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct GeminiOAuthSession { + pub session_id: String, + pub code_verifier: String, + pub state: String, + pub created_at: i64, +} + +/// 生成 Gemini OAuth 授权 URL 和会话信息 +/// +/// 返回 (auth_url, session) 元组 +/// - auth_url: 用户需要在浏览器中打开的授权 URL +/// - session: 包含 code_verifier 的会话信息，用于后续交换 token +pub fn generate_gemini_auth_url_with_session() -> (String, GeminiOAuthSession) { + let (code_verifier, code_challenge) = generate_pkce(); + let state = Uuid::new_v4().to_string(); + let session_id = Uuid::new_v4().to_string(); + + let auth_url = generate_gemini_auth_url(&state, &code_challenge); + + let session = GeminiOAuthSession { + session_id, + code_verifier, + state, + created_at: chrono::Utc::now().timestamp(), + }; + + (auth_url, session) +} + +/// 用授权码交换 Token 并创建凭证 +/// +/// 完整流程： +/// 1. 用 code + code_verifier 交换 tokens +/// 2. 获取用户邮箱 +/// 3. 获取项目 ID +/// 4. 保存凭证到文件 +pub async fn exchange_gemini_code_and_create_credentials( + code: &str, + code_verifier: &str, +) -> Result> { + let client = Client::builder() + .timeout(std::time::Duration::from_secs(30)) + .build()?; + + tracing::info!("[Gemini OAuth] 正在用授权码交换 Token..."); + + // 交换 Token + let token_data = exchange_gemini_code_for_token(&client, code, code_verifier).await?; + + let access_token = token_data["access_token"] + .as_str() + .ok_or("响应中没有 access_token")? + .to_string(); + let refresh_token = token_data["refresh_token"].as_str().map(|s| s.to_string()); + let expires_in = token_data["expires_in"].as_i64(); + + tracing::info!("[Gemini OAuth] Token 交换成功"); + + // 获取用户邮箱 + let email = fetch_gemini_user_email(&client, &access_token) + .await + .ok() + .flatten(); + + tracing::info!("[Gemini OAuth] 用户邮箱: {:?}", email); + + // 获取项目 ID + let _project_id = fetch_gemini_project_id(&client, &access_token) + .await + .ok() + .flatten(); + + // 构建凭证 + let now = chrono::Utc::now(); + let expires_at = expires_in.map(|secs| now + chrono::Duration::seconds(secs)); + + let credentials = GeminiCredentials { + access_token: Some(access_token), + refresh_token, + token_type: Some("Bearer".to_string()), + expiry_date: expires_at.map(|t| t.timestamp_millis()), + expire: expires_at.map(|t| t.to_rfc3339()), + scope: Some(GEMINI_OAUTH_SCOPES.join(" ")), + email, + last_refresh: Some(now.to_rfc3339()), + cred_type: "gemini".to_string(), + token: None, + }; + + // 保存凭证到文件 + let file_path = save_gemini_credentials_to_file(&credentials).await?; + + tracing::info!("[Gemini OAuth] 凭证已保存到: {}", file_path); + + Ok(GeminiOAuthResult { + credentials, + creds_file_path: file_path, + }) +} + +/// 用授权码交换 Token（使用 PKCE） +pub async fn exchange_gemini_code_for_token( + client: &Client, + code: &str, + code_verifier: &str, +) -> Result> { + let params = [ + ("code", code), + ("client_id", GEMINI_OAUTH_CLIENT_ID), + ("client_secret", GEMINI_OAUTH_CLIENT_SECRET), + ("code_verifier", code_verifier), + ("redirect_uri", GEMINI_OAUTH_REDIRECT_URI), + ("grant_type", "authorization_code"), + ]; + + let resp = client.post(GEMINI_TOKEN_URL).form(¶ms).send().await?; + + if !resp.status().is_success() { + let status = resp.status(); + let body = resp.text().await.unwrap_or_default(); + return Err(format!("Token 交换失败: {} - {}", status, body).into()); + } + + let data: serde_json::Value = resp.json().await?; + Ok(data) +} + +/// 获取用户邮箱 +pub async fn fetch_gemini_user_email( + client: &Client, + access_token: &str, +) -> Result, Box> { + let resp = client + .get("https://www.googleapis.com/oauth2/v2/userinfo") + .header("Authorization", format!("Bearer {}", access_token)) + .send() + .await?; + + if resp.status().is_success() { + let data: serde_json::Value = resp.json().await?; + Ok(data["email"].as_str().map(|s| s.to_string())) + } else { + Ok(None) + } +} + +/// 获取项目 ID（通过 loadCodeAssist 接口） +pub async fn fetch_gemini_project_id( + client: &Client, + access_token: &str, +) -> Result, Box> { + tracing::info!("[Gemini OAuth] 正在获取 projectId..."); + + let resp = client + .post(format!( + "{}/{CODE_ASSIST_API_VERSION}:loadCodeAssist", + CODE_ASSIST_ENDPOINT + )) + .header("Authorization", format!("Bearer {}", access_token)) + .header("Content-Type", "application/json") + .json(&serde_json::json!({ + "cloudaicompanionProject": "", + "metadata": { + "ideType": "IDE_UNSPECIFIED", + "platform": "PLATFORM_UNSPECIFIED", + "pluginType": "GEMINI", + "duetProject": "" + } + })) + .send() + .await?; + + let status = resp.status(); + tracing::info!("[Gemini OAuth] loadCodeAssist 响应状态: {}", status); + + if status.is_success() { + let data: serde_json::Value = resp.json().await?; + if let Some(project) = data["cloudaicompanionProject"].as_str() { + if !project.is_empty() { + tracing::info!("[Gemini OAuth] 获取到 projectId: {}", project); + return Ok(Some(project.to_string())); + } + } + tracing::info!("[Gemini OAuth] cloudaicompanionProject 为空"); + Ok(None) + } else { + let body = resp.text().await.unwrap_or_default(); + tracing::warn!( + "[Gemini OAuth] loadCodeAssist 请求失败: {} - {}", + status, + body + ); + Ok(None) + } +} + +/// OAuth 成功页面 HTML +const GEMINI_OAUTH_SUCCESS_HTML: &str = r#" + + + + 授权成功 + + + +
+

✓ 授权成功

+

Gemini 账号已添加到 ProxyCast

+ +

可以关闭此页面

+
+ +"#; + +/// OAuth 失败页面 HTML +const GEMINI_OAUTH_ERROR_HTML: &str = r#" + + + + 授权失败 + + + +
+

✗ 授权失败

+

ERROR_PLACEHOLDER

+

请关闭此页面后重试

+
+ +"#; + +/// 保存 Gemini 凭证到文件 +async fn save_gemini_credentials_to_file( + credentials: &GeminiCredentials, +) -> Result> { + // 生成唯一文件名 + let uuid = Uuid::new_v4().to_string(); + let timestamp = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_secs(); + + let filename = format!("gemini_{}_{}_gemini.json", &uuid[..8], timestamp); + + // 获取凭证存储目录 + let credentials_dir = dirs::data_dir() + .ok_or_else(|| "无法获取应用数据目录")? + .join("proxycast") + .join("credentials"); + + // 确保目录存在 + tokio::fs::create_dir_all(&credentials_dir).await?; + + let file_path = credentials_dir.join(&filename); + + // 写入凭证 + let content = serde_json::to_string_pretty(credentials)?; + tokio::fs::write(&file_path, content).await?; + + Ok(file_path.to_string_lossy().to_string()) +} + +/// 启动 OAuth 服务器并返回授权 URL（不打开浏览器） +/// 服务器会在后台等待回调，成功后返回凭证 +pub async fn start_gemini_oauth_server_and_get_url() -> Result< + ( + String, + impl std::future::Future< + Output = Result>, + >, + ), + Box, +> { + use axum::{extract::Query, response::Html, routing::get, Router}; + use std::collections::HashMap; + use tokio::net::TcpListener; + + let client = Client::builder() + .timeout(std::time::Duration::from_secs(30)) + .build()?; + + // 生成 PKCE + let (code_verifier, code_challenge) = generate_pkce(); + + // 生成随机 state + let state = Uuid::new_v4().to_string(); + let state_clone = state.clone(); + let code_verifier_clone = code_verifier.clone(); + + // 创建 channel 用于接收回调结果 + let (tx, rx) = oneshot::channel::>(); + let tx = Arc::new(tokio::sync::Mutex::new(Some(tx))); + + // 尝试绑定到多个端口 + let ports_to_try = [11451, 11452, 11453, 11454, 11455, 0]; + let mut listener = None; + let mut bound_port = 0; + + for port in ports_to_try { + match TcpListener::bind(format!("127.0.0.1:{}", port)).await { + Ok(l) => { + bound_port = l.local_addr()?.port(); + listener = Some(l); + tracing::info!("[Gemini OAuth] 成功绑定到端口 {}", bound_port); + break; + } + Err(e) => { + tracing::warn!("[Gemini OAuth] 端口 {} 绑定失败: {}", port, e); + continue; + } + } + } + + let listener = listener.ok_or("无法绑定到任何可用端口")?; + + // 生成授权 URL + let auth_url = generate_gemini_auth_url(&state, &code_challenge); + + tracing::info!( + "[Gemini OAuth] 服务器启动在端口 {}, 授权 URL: {}", + bound_port, + auth_url + ); + + // 构建路由 + let app = Router::new().route( + "/oauth-callback", + get(move |Query(params): Query>| { + let tx = tx.clone(); + let client = client.clone(); + let state_expected = state_clone.clone(); + let code_verifier = code_verifier_clone.clone(); + + async move { + let code = params.get("code"); + let returned_state = params.get("state"); + let error = params.get("error"); + + // 检查错误 + if let Some(err) = error { + let error_desc = params + .get("error_description") + .map(|s| s.as_str()) + .unwrap_or("未知错误"); + let error_msg = format!("{}: {}", err, error_desc); + tracing::error!("[Gemini OAuth] 授权失败: {}", error_msg); + + if let Some(tx) = tx.lock().await.take() { + let _ = tx.send(Err(error_msg.clone())); + } + + let html = GEMINI_OAUTH_ERROR_HTML.replace("ERROR_PLACEHOLDER", &error_msg); + return Html(html); + } + + // 验证 state + if returned_state.map(|s| s.as_str()) != Some(&state_expected) { + let error_msg = "State 验证失败"; + tracing::error!("[Gemini OAuth] {}", error_msg); + + if let Some(tx) = tx.lock().await.take() { + let _ = tx.send(Err(error_msg.to_string())); + } + + let html = GEMINI_OAUTH_ERROR_HTML.replace("ERROR_PLACEHOLDER", error_msg); + return Html(html); + } + + // 获取授权码 + let code = match code { + Some(c) => c, + None => { + let error_msg = "未收到授权码"; + tracing::error!("[Gemini OAuth] {}", error_msg); + + if let Some(tx) = tx.lock().await.take() { + let _ = tx.send(Err(error_msg.to_string())); + } + + let html = GEMINI_OAUTH_ERROR_HTML.replace("ERROR_PLACEHOLDER", error_msg); + return Html(html); + } + }; + + tracing::info!("[Gemini OAuth] 收到授权码，正在交换 Token..."); + + // 交换 Token + let token_result = + exchange_gemini_code_for_token(&client, code, &code_verifier).await; + + match token_result { + Ok(token_data) => { + let access_token = token_data["access_token"] + .as_str() + .unwrap_or("") + .to_string(); + let refresh_token = + token_data["refresh_token"].as_str().map(|s| s.to_string()); + let expires_in = token_data["expires_in"].as_i64(); + + // 获取用户邮箱 + let email = fetch_gemini_user_email(&client, &access_token) + .await + .ok() + .flatten(); + + // 获取项目 ID + let project_id = fetch_gemini_project_id(&client, &access_token) + .await + .ok() + .flatten(); + + // 构建凭证 + let now = chrono::Utc::now(); + let expires_at = + expires_in.map(|secs| now + chrono::Duration::seconds(secs)); + + let credentials = GeminiCredentials { + access_token: Some(access_token), + refresh_token, + token_type: Some("Bearer".to_string()), + expiry_date: expires_at.map(|t| t.timestamp_millis()), + expire: expires_at.map(|t| t.to_rfc3339()), + scope: Some(GEMINI_OAUTH_SCOPES.join(" ")), + email: email.clone(), + last_refresh: Some(now.to_rfc3339()), + cred_type: "gemini".to_string(), + token: None, + }; + + // 保存凭证到文件 + match save_gemini_credentials_to_file(&credentials).await { + Ok(file_path) => { + tracing::info!("[Gemini OAuth] 凭证已保存到: {}", file_path); + + let result = GeminiOAuthResult { + credentials: credentials.clone(), + creds_file_path: file_path, + }; + + if let Some(tx) = tx.lock().await.take() { + let _ = tx.send(Ok(result)); + } + + let email_display = email.unwrap_or_else(|| "未知邮箱".to_string()); + let project_display = project_id + .map(|p| format!("

Project ID: {}

", p)) + .unwrap_or_default(); + let html = GEMINI_OAUTH_SUCCESS_HTML + .replace("EMAIL_PLACEHOLDER", &email_display) + .replace( + "\n", + &format!("{}\n", project_display), + ); + Html(html) + } + Err(e) => { + let error_msg = format!("保存凭证失败: {}", e); + tracing::error!("[Gemini OAuth] {}", error_msg); + + if let Some(tx) = tx.lock().await.take() { + let _ = tx.send(Err(error_msg.clone())); + } + + let html = GEMINI_OAUTH_ERROR_HTML + .replace("ERROR_PLACEHOLDER", &error_msg); + Html(html) + } + } + } + Err(e) => { + let error_msg = format!("Token 交换失败: {}", e); + tracing::error!("[Gemini OAuth] {}", error_msg); + + if let Some(tx) = tx.lock().await.take() { + let _ = tx.send(Err(error_msg.clone())); + } + + let html = GEMINI_OAUTH_ERROR_HTML.replace("ERROR_PLACEHOLDER", &error_msg); + Html(html) + } + } + } + }), + ); + + // 启动服务器 + let server_future = async move { + axum::serve(listener, app) + .await + .map_err(|e| format!("服务器错误: {}", e)) + }; + + // 启动服务器任务 + tokio::spawn(server_future); + + // 返回授权 URL 和等待结果的 Future + let wait_future = async move { + match rx.await { + Ok(result) => result.map_err(|e| e.into()), + Err(_) => Err("OAuth 回调通道关闭".into()), + } + }; + + Ok((auth_url, wait_future)) +} + +/// 启动 Gemini OAuth 登录流程（自动打开浏览器） +pub async fn start_gemini_oauth_login( +) -> Result> { + let (auth_url, wait_future) = start_gemini_oauth_server_and_get_url().await?; + + // 打开浏览器 + tracing::info!("[Gemini OAuth] 正在打开浏览器..."); + if let Err(e) = open::that(&auth_url) { + tracing::warn!("[Gemini OAuth] 无法自动打开浏览器: {}", e); + } + + // 等待回调 + wait_future.await +} + +// ============================================================================ +// CredentialProvider Trait 实现 +// ============================================================================ + +#[async_trait] +impl CredentialProvider for GeminiProvider { + async fn load_credentials_from_path(&mut self, path: &str) -> ProviderResult<()> { + GeminiProvider::load_credentials_from_path(self, path).await + } + + async fn save_credentials(&self) -> ProviderResult<()> { + GeminiProvider::save_credentials(self).await + } + + fn is_token_valid(&self) -> bool { + GeminiProvider::is_token_valid(self) + } + + fn is_token_expiring_soon(&self) -> bool { + // Gemini 使用与 is_token_valid 相同的逻辑，但阈值为 10 分钟 + if self.credentials.access_token.is_none() { + return true; + } + + if let Some(expire_str) = &self.credentials.expire { + if let Ok(expires) = chrono::DateTime::parse_from_rfc3339(expire_str) { + let now = chrono::Utc::now(); + return expires <= now + chrono::Duration::minutes(10); + } + } + + if let Some(expiry) = self.credentials.expiry_date { + let now = chrono::Utc::now().timestamp_millis(); + return expiry <= now + 600_000; // 10 分钟 + } + + false + } + + async fn refresh_token(&mut self) -> ProviderResult { + GeminiProvider::refresh_token(self).await + } + + fn get_access_token(&self) -> Option<&str> { + self.credentials.access_token.as_deref() + } + + fn provider_type(&self) -> &'static str { + "gemini" + } +} diff --git a/src-tauri/src/providers/iflow.rs b/src-tauri/src/providers/iflow.rs index c75f254be..7d654d91a 100644 --- a/src-tauri/src/providers/iflow.rs +++ b/src-tauri/src/providers/iflow.rs @@ -1575,3 +1575,431 @@ mod tests { assert!(provider.are_cookies_expired()); } } + +// ============================================================================ +// OAuth 登录功能（与 CLIProxyAPI 对齐） +// ============================================================================ + +use std::sync::Arc; +use tokio::sync::oneshot; +use uuid::Uuid; + +/// iFlow OAuth 登录结果 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct IFlowOAuthResult { + pub credentials: IFlowCredentials, + pub creds_file_path: String, +} + +/// OAuth 成功页面 HTML +const IFLOW_OAUTH_SUCCESS_HTML: &str = r#" + + + + 授权成功 + + + +
+

✓ 授权成功

+

iFlow 账号已添加到 ProxyCast

+ +

可以关闭此页面

+
+ +"#; + +/// OAuth 失败页面 HTML +const IFLOW_OAUTH_ERROR_HTML: &str = r#" + + + + 授权失败 + + + +
+

✗ 授权失败

+

ERROR_PLACEHOLDER

+

请关闭此页面后重试

+
+ +"#; + +/// 生成 iFlow OAuth 授权 URL +pub fn generate_iflow_auth_url(port: u16, state: &str) -> String { + let redirect_uri = format!("http://localhost:{}/oauth2callback", port); + + let params = [ + ("loginMethod", "phone"), + ("type", "phone"), + ("redirect", redirect_uri.as_str()), + ("state", state), + ("client_id", IFLOW_CLIENT_ID), + ]; + + let query = params + .iter() + .map(|(k, v)| format!("{}={}", k, urlencoding::encode(v))) + .collect::>() + .join("&"); + + format!("{}?{}", IFLOW_AUTH_URL, query) +} + +/// 用授权码交换 Token +pub async fn exchange_iflow_code_for_token( + client: &Client, + code: &str, + redirect_uri: &str, +) -> Result> { + // 构建 Basic Auth 头 + let basic_auth = BASE64_STANDARD.encode(format!("{}:{}", IFLOW_CLIENT_ID, IFLOW_CLIENT_SECRET)); + + let params = [ + ("grant_type", "authorization_code"), + ("code", code), + ("redirect_uri", redirect_uri), + ("client_id", IFLOW_CLIENT_ID), + ("client_secret", IFLOW_CLIENT_SECRET), + ]; + + let resp = client + .post(IFLOW_TOKEN_URL) + .header("Content-Type", "application/x-www-form-urlencoded") + .header("Accept", "application/json") + .header("Authorization", format!("Basic {}", basic_auth)) + .form(¶ms) + .send() + .await?; + + if !resp.status().is_success() { + let status = resp.status(); + let body = resp.text().await.unwrap_or_default(); + return Err(format!("Token 交换失败: {} - {}", status, body).into()); + } + + let data: serde_json::Value = resp.json().await?; + + let access_token = data["access_token"] + .as_str() + .ok_or("响应中没有 access_token")? + .to_string(); + let refresh_token = data["refresh_token"].as_str().map(|s| s.to_string()); + let token_type = data["token_type"].as_str().map(|s| s.to_string()); + let scope = data["scope"].as_str().map(|s| s.to_string()); + let expires_in = data["expires_in"].as_i64().unwrap_or(3600); + + let expires_at = chrono::Utc::now() + chrono::Duration::seconds(expires_in); + + // 获取用户信息和 API Key + let user_info_url = format!( + "{}?accessToken={}", + IFLOW_USER_INFO_URL, + urlencoding::encode(&access_token) + ); + + let user_resp = client + .get(&user_info_url) + .header("Accept", "application/json") + .send() + .await; + + let (api_key, email) = if let Ok(resp) = user_resp { + if resp.status().is_success() { + if let Ok(user_data) = resp.json::().await { + if user_data["success"].as_bool().unwrap_or(false) { + let api_key = user_data["data"]["apiKey"].as_str().map(|s| s.to_string()); + let email = user_data["data"]["email"] + .as_str() + .or_else(|| user_data["data"]["phone"].as_str()) + .map(|s| s.to_string()); + (api_key, email) + } else { + (None, None) + } + } else { + (None, None) + } + } else { + (None, None) + } + } else { + (None, None) + }; + + let credentials = IFlowCredentials { + auth_type: "oauth".to_string(), + access_token: Some(access_token), + refresh_token, + expire: Some(expires_at.to_rfc3339()), + expires_at: Some(expires_at.to_rfc3339()), + cookies: None, + cookie_expires_at: None, + email, + user_id: None, + last_refresh: Some(chrono::Utc::now().to_rfc3339()), + api_key, + token_type, + scope, + cred_type: "iflow".to_string(), + }; + + Ok(credentials) +} + +/// 启动 OAuth 服务器并返回授权 URL（不打开浏览器） +pub async fn start_iflow_oauth_server_and_get_url() -> Result< + ( + String, + impl std::future::Future>>, + ), + Box, +> { + use axum::{extract::Query, response::Html, routing::get, Router}; + use std::collections::HashMap; + use tokio::net::TcpListener; + + let client = Client::builder() + .timeout(std::time::Duration::from_secs(30)) + .build()?; + + // 生成随机 state + let state = Uuid::new_v4().to_string(); + let state_clone = state.clone(); + + // 创建 channel 用于接收回调结果 + let (tx, rx) = oneshot::channel::>(); + let tx = Arc::new(tokio::sync::Mutex::new(Some(tx))); + + // 尝试绑定到端口，如果默认端口被占用则尝试其他端口 + let ports_to_try = [DEFAULT_CALLBACK_PORT, 11452, 11453, 11454, 11455, 0]; // 0 表示让系统分配 + let mut listener = None; + let mut last_error = None; + + for port in ports_to_try { + match TcpListener::bind(format!("127.0.0.1:{}", port)).await { + Ok(l) => { + listener = Some(l); + break; + } + Err(e) => { + last_error = Some(e); + continue; + } + } + } + + let listener = listener.ok_or_else(|| { + format!( + "启动 OAuth 服务器失败: {}", + last_error + .map(|e| e.to_string()) + .unwrap_or_else(|| "无法绑定端口".to_string()) + ) + })?; + let port = listener.local_addr()?.port(); + + let redirect_uri = format!("http://localhost:{}/oauth2callback", port); + let redirect_uri_clone = redirect_uri.clone(); + + // 生成授权 URL + let auth_url = generate_iflow_auth_url(port, &state); + + tracing::info!( + "[iFlow OAuth] 服务器启动在端口 {}, 授权 URL: {}", + port, + auth_url + ); + + // 构建路由 + let app = Router::new().route( + "/oauth2callback", + get(move |Query(params): Query>| { + let tx = tx.clone(); + let client = client.clone(); + let state_expected = state_clone.clone(); + let redirect_uri = redirect_uri_clone.clone(); + + async move { + let code = params.get("code"); + let returned_state = params.get("state"); + let error = params.get("error"); + + // 检查错误 + if let Some(err) = error { + let html = IFLOW_OAUTH_ERROR_HTML.replace("ERROR_PLACEHOLDER", err); + if let Some(sender) = tx.lock().await.take() { + let _ = sender.send(Err(format!("OAuth 错误: {}", err))); + } + return Html(html); + } + + // 检查 state + if returned_state.map(|s| s.as_str()) != Some(&state_expected) { + let html = + IFLOW_OAUTH_ERROR_HTML.replace("ERROR_PLACEHOLDER", "State 验证失败"); + if let Some(sender) = tx.lock().await.take() { + let _ = sender.send(Err("State 验证失败".to_string())); + } + return Html(html); + } + + // 检查 code + let code = match code { + Some(c) => c, + None => { + let html = + IFLOW_OAUTH_ERROR_HTML.replace("ERROR_PLACEHOLDER", "未收到授权码"); + if let Some(sender) = tx.lock().await.take() { + let _ = sender.send(Err("未收到授权码".to_string())); + } + return Html(html); + } + }; + + // 交换 Token + let token_result = + exchange_iflow_code_for_token(&client, code, &redirect_uri).await; + let credentials = match token_result { + Ok(creds) => creds, + Err(e) => { + let html = + IFLOW_OAUTH_ERROR_HTML.replace("ERROR_PLACEHOLDER", &e.to_string()); + if let Some(sender) = tx.lock().await.take() { + let _ = sender.send(Err(e.to_string())); + } + return Html(html); + } + }; + + let email = credentials.email.clone(); + + // 保存凭证到应用数据目录 + let creds_dir = dirs::data_dir() + .unwrap_or_else(|| std::path::PathBuf::from(".")) + .join("proxycast") + .join("credentials") + .join("iflow"); + + if let Err(e) = std::fs::create_dir_all(&creds_dir) { + let html = IFLOW_OAUTH_ERROR_HTML + .replace("ERROR_PLACEHOLDER", &format!("创建目录失败: {}", e)); + if let Some(sender) = tx.lock().await.take() { + let _ = sender.send(Err(format!("创建目录失败: {}", e))); + } + return Html(html); + } + + // 生成唯一文件名 + let uuid = Uuid::new_v4().to_string(); + let timestamp = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_secs(); + let filename = format!("iflow_{}_{}.json", &uuid[..8], timestamp); + let creds_file_path = creds_dir.join(&filename); + + // 保存凭证 + let creds_json = match serde_json::to_string_pretty(&credentials) { + Ok(json) => json, + Err(e) => { + let html = IFLOW_OAUTH_ERROR_HTML + .replace("ERROR_PLACEHOLDER", &format!("序列化凭证失败: {}", e)); + if let Some(sender) = tx.lock().await.take() { + let _ = sender.send(Err(format!("序列化凭证失败: {}", e))); + } + return Html(html); + } + }; + + if let Err(e) = std::fs::write(&creds_file_path, &creds_json) { + let html = IFLOW_OAUTH_ERROR_HTML + .replace("ERROR_PLACEHOLDER", &format!("保存凭证失败: {}", e)); + if let Some(sender) = tx.lock().await.take() { + let _ = sender.send(Err(format!("保存凭证失败: {}", e))); + } + return Html(html); + } + + tracing::info!("[iFlow OAuth] 凭证已保存到: {:?}", creds_file_path); + + // 发送成功结果 + let result = IFlowOAuthResult { + credentials, + creds_file_path: creds_file_path.to_string_lossy().to_string(), + }; + + if let Some(sender) = tx.lock().await.take() { + let _ = sender.send(Ok(result)); + } + + // 返回成功页面 + let html = IFLOW_OAUTH_SUCCESS_HTML.replace( + "EMAIL_PLACEHOLDER", + &email.unwrap_or_else(|| "未知用户".to_string()), + ); + Html(html) + } + }), + ); + + // 启动服务器 + let server = axum::serve(listener, app); + + // 创建等待 future + let wait_future = async move { + // 设置超时（5 分钟） + let timeout = tokio::time::timeout(std::time::Duration::from_secs(300), async { + // 启动服务器（在后台运行） + tokio::spawn(async move { + if let Err(e) = server.await { + tracing::error!("[iFlow OAuth] 服务器错误: {}", e); + } + }); + + // 等待回调结果 + match rx.await { + Ok(result) => result.map_err(|e| { + Box::new(std::io::Error::new(std::io::ErrorKind::Other, e)) + as Box + }), + Err(_) => Err("OAuth 回调通道关闭".into()), + } + }); + + match timeout.await { + Ok(result) => result, + Err(_) => Err("OAuth 登录超时（5分钟）".into()), + } + }; + + Ok((auth_url, wait_future)) +} + +/// 启动 iFlow OAuth 登录流程（自动打开浏览器） +pub async fn start_iflow_oauth_login() -> Result> { + let (auth_url, wait_future) = start_iflow_oauth_server_and_get_url().await?; + + tracing::info!("[iFlow OAuth] 打开浏览器进行授权: {}", auth_url); + + // 打开浏览器 + if let Err(e) = open::that(&auth_url) { + tracing::warn!("[iFlow OAuth] 无法打开浏览器: {}. 请手动打开 URL.", e); + } + + // 等待回调 + wait_future.await +} diff --git a/src-tauri/src/providers/kiro.rs b/src-tauri/src/providers/kiro.rs index fb0150ee6..f2af8f961 100644 --- a/src-tauri/src/providers/kiro.rs +++ b/src-tauri/src/providers/kiro.rs @@ -1,12 +1,80 @@ //! Kiro/CodeWhisperer Provider use crate::converter::openai_to_cw::convert_openai_to_codewhisperer; use crate::models::openai::*; +use crate::providers::traits::{CredentialProvider, ProviderResult}; +use async_trait::async_trait; use reqwest::Client; use serde::{Deserialize, Serialize}; use std::error::Error; use std::path::PathBuf; -/// 生成设备指纹 (Machine ID 的 SHA256) +/// 根据凭证信息生成唯一的 Machine ID（参考 AIClient-2-API 实现） +/// +/// 关键改进：每个凭证生成独立的 Machine ID，避免多账号共用同一指纹被检测 +/// 优先级：profileArn > clientId > 系统硬件 ID +/// +/// 这样每个 OAuth 凭证都有自己独立的指纹，模拟不同设备登录 +pub fn generate_machine_id_from_credentials( + profile_arn: Option<&str>, + client_id: Option<&str>, +) -> String { + use sha2::{Digest, Sha256}; + + // 优先使用凭证相关的唯一标识，确保每个账号有独立的 Machine ID + let unique_key = profile_arn + .filter(|s| !s.is_empty()) + .or(client_id.filter(|s| !s.is_empty())) + .map(|s| s.to_string()) + .unwrap_or_else(|| { + // 回退到系统硬件 ID + get_raw_machine_id().unwrap_or_else(|| "KIRO_DEFAULT_MACHINE".to_string()) + }); + + let mut hasher = Sha256::new(); + hasher.update(unique_key.as_bytes()); + let result = hasher.finalize(); + format!("{:x}", result) +} + +/// 获取系统运行时信息 +/// +/// 返回真实的操作系统名称和版本，用于构建更真实的 User-Agent +fn get_system_runtime_info() -> (String, String) { + let os_name = if cfg!(target_os = "macos") { + // macOS: 获取真实版本号 + let version = std::process::Command::new("sw_vers") + .arg("-productVersion") + .output() + .ok() + .and_then(|o| String::from_utf8(o.stdout).ok()) + .map(|s| s.trim().to_string()) + .unwrap_or_else(|| "14.0".to_string()); + format!("macos#{}", version) + } else if cfg!(target_os = "linux") { + // Linux: 获取内核版本 + let version = std::process::Command::new("uname") + .arg("-r") + .output() + .ok() + .and_then(|o| String::from_utf8(o.stdout).ok()) + .map(|s| s.trim().to_string()) + .unwrap_or_else(|| "5.15.0".to_string()); + format!("linux#{}", version) + } else if cfg!(target_os = "windows") { + // Windows: 使用固定版本（实际应该获取真实版本） + "windows#10.0".to_string() + } else { + "other#1.0".to_string() + }; + + // Node.js 版本模拟（Kiro IDE 使用 Electron，内置 Node.js） + // 使用常见的 LTS 版本 + let node_version = "20.18.0".to_string(); + + (os_name, node_version) +} + +/// 生成设备指纹 (Machine ID 的 SHA256) - 保留用于兼容 /// /// 与 Kiro IDE 保持一致的指纹生成方式（参考 Kir-Manager）： /// - macOS: 使用 IOPlatformUUID（硬件级别唯一标识） @@ -635,7 +703,11 @@ impl KiroProvider { ); // 获取设备指纹和版本号（用于 Social 认证的 User-Agent） - let device_fp = get_device_fingerprint(); + // 使用基于凭证的 Machine ID，确保每个账号有独立的指纹 + let machine_id = generate_machine_id_from_credentials( + self.credentials.profile_arn.as_deref(), + self.credentials.client_id.as_deref(), + ); let kiro_version = get_kiro_version(); let resp = if auth_method == "idc" { @@ -668,11 +740,11 @@ impl KiroProvider { .header("Host", "oidc.us-east-1.amazonaws.com") .header( "x-amz-user-agent", - "aws-sdk-js/3.738.0 ua/2.1 os/other lang/js api/sso-oidc#3.738.0 m/E KiroIDE", + format!("aws-sdk-js/3.738.0 ua/2.1 os/other lang/js api/sso-oidc#3.738.0 m/E KiroIDE-{}-{}", kiro_version, machine_id), ) .header("User-Agent", "node") .header("Accept", "*/*") - .header("Connection", "keep-alive") + .header("Connection", "close") .json(&body) .send() .await? @@ -685,13 +757,14 @@ impl KiroProvider { .post(&refresh_url) .header( "User-Agent", - format!("KiroIDE-{}-{}", kiro_version, device_fp), + format!("KiroIDE-{}-{}", kiro_version, machine_id), ) .header("Accept", "application/json, text/plain, */*") .header("Accept-Encoding", "br, gzip, deflate") .header("Content-Type", "application/json") .header("Accept-Language", "*") .header("Sec-Fetch-Mode", "cors") + .header("Connection", "close") .json(&body) .send() .await? @@ -859,7 +932,7 @@ impl KiroProvider { None }; - let cw_request = convert_openai_to_codewhisperer(request, profile_arn); + let cw_request = convert_openai_to_codewhisperer(request, profile_arn.clone()); let url = self.get_base_url(); // 安全修复：仅在 PROXYCAST_DEBUG=1 时写入请求调试文件，避免泄露敏感信息 @@ -905,9 +978,20 @@ impl KiroProvider { current_has_tools ); - // 生成设备指纹用于伪装 Kiro IDE - let device_fp = get_device_fingerprint(); + // 生成基于凭证的唯一 Machine ID（关键改进：每个账号独立指纹） + let machine_id = generate_machine_id_from_credentials( + profile_arn.as_deref(), + self.credentials.client_id.as_deref(), + ); let kiro_version = get_kiro_version(); + let (os_name, node_version) = get_system_runtime_info(); + + tracing::debug!( + "[KIRO_FINGERPRINT] machine_id={} (based on profile_arn={}, client_id={})", + &machine_id[..16], + profile_arn.is_some(), + self.credentials.client_id.is_some() + ); let resp = self .client @@ -917,17 +1001,20 @@ impl KiroProvider { .header("Accept", "application/json") .header("amz-sdk-invocation-id", uuid::Uuid::new_v4().to_string()) .header("amz-sdk-request", "attempt=1; max=1") + .header("x-amzn-kiro-agent-mode", "vibe") + // 关键指纹头：使用基于凭证的唯一 Machine ID .header( "x-amz-user-agent", - format!("aws-sdk-js/1.0.7 KiroIDE-{kiro_version}-{device_fp}"), + format!("aws-sdk-js/1.0.0 KiroIDE-{kiro_version}-{machine_id}"), ) .header( "user-agent", format!( - "aws-sdk-js/1.0.7 ua/2.1 os/macos#14.0 lang/js md/nodejs#20.16.0 api/codewhispererstreaming#1.0.7 m/E KiroIDE-{kiro_version}-{device_fp}" + "aws-sdk-js/1.0.0 ua/2.1 os/{os_name} lang/js md/nodejs#{node_version} api/codewhispererruntime#1.0.0 m/E KiroIDE-{kiro_version}-{machine_id}" ), ) - .header("x-amzn-kiro-agent-mode", "vibe") + // 添加 Connection: close 避免连接复用被检测 + .header("Connection", "close") .json(&cw_request) .send() .await?; @@ -972,3 +1059,39 @@ fn merge_credentials(target: &mut KiroCredentials, source: &KiroCredentials) { } // cred_type 使用默认值，不需要合并 } + +// ============================================================================ +// CredentialProvider Trait 实现 +// ============================================================================ + +#[async_trait] +impl CredentialProvider for KiroProvider { + async fn load_credentials_from_path(&mut self, path: &str) -> ProviderResult<()> { + // 调用已有的实现 + KiroProvider::load_credentials_from_path(self, path).await + } + + async fn save_credentials(&self) -> ProviderResult<()> { + KiroProvider::save_credentials(self).await + } + + fn is_token_valid(&self) -> bool { + !self.is_token_expired() + } + + fn is_token_expiring_soon(&self) -> bool { + KiroProvider::is_token_expiring_soon(self) + } + + async fn refresh_token(&mut self) -> ProviderResult { + KiroProvider::refresh_token(self).await + } + + fn get_access_token(&self) -> Option<&str> { + self.credentials.access_token.as_deref() + } + + fn provider_type(&self) -> &'static str { + "kiro" + } +} diff --git a/src-tauri/src/providers/mod.rs b/src-tauri/src/providers/mod.rs index de0973332..8e6288f43 100644 --- a/src-tauri/src/providers/mod.rs +++ b/src-tauri/src/providers/mod.rs @@ -8,11 +8,16 @@ pub mod iflow; pub mod kiro; pub mod openai_custom; pub mod qwen; +pub mod traits; pub mod vertex; #[cfg(test)] mod tests; +// Trait exports +#[allow(unused_imports)] +pub use traits::{CredentialProvider, ProviderResult, TokenManager}; + #[allow(unused_imports)] pub use antigravity::AntigravityProvider; #[allow(unused_imports)] @@ -22,7 +27,7 @@ pub use claude_oauth::ClaudeOAuthProvider; #[allow(unused_imports)] pub use codex::CodexProvider; #[allow(unused_imports)] -pub use error::{ProviderError, ProviderResult}; +pub use error::ProviderError; #[allow(unused_imports)] pub use gemini::{GeminiApiKeyCredential, GeminiApiKeyProvider, GeminiProvider}; #[allow(unused_imports)] diff --git a/src-tauri/src/providers/qwen.rs b/src-tauri/src/providers/qwen.rs index e1ca9779a..7c8ad5951 100644 --- a/src-tauri/src/providers/qwen.rs +++ b/src-tauri/src/providers/qwen.rs @@ -6,6 +6,8 @@ use super::error::{ create_auth_error, create_config_error, create_token_refresh_error, ProviderError, }; +use super::traits::{CredentialProvider, ProviderResult}; +use async_trait::async_trait; use reqwest::Client; use serde::{Deserialize, Serialize}; use std::error::Error; @@ -319,3 +321,387 @@ impl QwenProvider { Ok(resp) } } + +// ============================================================================ +// Device Code Flow 登录功能（与 CLIProxyAPI 对齐） +// ============================================================================ + +use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine}; +use rand::RngCore; +use sha2::{Digest, Sha256}; +use uuid::Uuid; + +// Device Code Flow 端点 +const QWEN_DEVICE_CODE_URL: &str = "https://chat.qwen.ai/api/v1/oauth2/device/code"; +const QWEN_OAUTH_SCOPE: &str = "openid profile email model.completion"; +const QWEN_DEVICE_GRANT_TYPE: &str = "urn:ietf:params:oauth:grant-type:device_code"; + +/// Device Code Flow 响应 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct DeviceCodeResponse { + /// 设备码（用于轮询） + #[serde(alias = "deviceCode")] + pub device_code: String, + /// 用户码（用户在浏览器中输入） + #[serde(alias = "userCode")] + pub user_code: String, + /// 验证 URL + #[serde(alias = "verificationUri")] + pub verification_uri: String, + /// 完整验证 URL（包含 user_code） + #[serde(default, alias = "verificationUriComplete")] + pub verification_uri_complete: Option, + /// 过期时间（秒） + #[serde(alias = "expiresIn")] + pub expires_in: i64, + /// 轮询间隔（秒），默认 5 秒 + #[serde(default = "default_interval")] + pub interval: i64, +} + +fn default_interval() -> i64 { + 5 +} + +/// Qwen OAuth 登录结果 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct QwenOAuthResult { + pub credentials: QwenCredentials, + pub creds_file_path: String, +} + +/// PKCE 代码生成 +fn generate_pkce_pair() -> Result<(String, String), Box> { + let mut bytes = [0u8; 32]; + rand::thread_rng().fill_bytes(&mut bytes); + let code_verifier = URL_SAFE_NO_PAD.encode(bytes); + + let mut hasher = Sha256::new(); + hasher.update(code_verifier.as_bytes()); + let hash = hasher.finalize(); + let code_challenge = URL_SAFE_NO_PAD.encode(hash); + + Ok((code_verifier, code_challenge)) +} + +/// 发起 Device Code Flow +pub async fn initiate_device_flow( + client: &Client, +) -> Result<(DeviceCodeResponse, String), Box> { + let (code_verifier, code_challenge) = generate_pkce_pair()?; + + let params = [ + ("client_id", QWEN_CLIENT_ID), + ("scope", QWEN_OAUTH_SCOPE), + ("code_challenge", code_challenge.as_str()), + ("code_challenge_method", "S256"), + ]; + + let resp = client + .post(QWEN_DEVICE_CODE_URL) + .header("Content-Type", "application/x-www-form-urlencoded") + .header("Accept", "application/json") + .form(¶ms) + .send() + .await?; + + if !resp.status().is_success() { + let status = resp.status(); + let body = resp.text().await.unwrap_or_default(); + return Err(format!("Device code 请求失败: {} - {}", status, body).into()); + } + + // 先获取响应体文本，以便在解析失败时提供详细错误信息 + let body = resp.text().await?; + tracing::debug!("[QWEN] Device Code 响应: {}", body); + + let device_response: DeviceCodeResponse = serde_json::from_str(&body) + .map_err(|e| format!("解析 Device Code 响应失败: {} - 响应内容: {}", e, body))?; + + if device_response.device_code.is_empty() { + return Err("Device code 响应中没有 device_code".into()); + } + + tracing::info!( + "[QWEN] Device Code Flow 已启动，user_code: {}, verification_uri: {}", + device_response.user_code, + device_response.verification_uri + ); + + Ok((device_response, code_verifier)) +} + +/// 轮询 Token 端点 +pub async fn poll_for_token( + client: &Client, + device_code: &str, + code_verifier: &str, + interval: u64, + max_attempts: u32, +) -> Result> { + let poll_interval = std::time::Duration::from_secs(interval.max(5)); + + for attempt in 0..max_attempts { + if attempt > 0 { + tokio::time::sleep(poll_interval).await; + } + + tracing::debug!("[QWEN] 轮询 Token，第 {} 次尝试", attempt + 1); + + let params = [ + ("grant_type", QWEN_DEVICE_GRANT_TYPE), + ("client_id", QWEN_CLIENT_ID), + ("device_code", device_code), + ("code_verifier", code_verifier), + ]; + + let resp = client + .post(QWEN_TOKEN_URL) + .header("Content-Type", "application/x-www-form-urlencoded") + .header("Accept", "application/json") + .form(¶ms) + .send() + .await; + + let resp = match resp { + Ok(r) => r, + Err(e) => { + tracing::warn!("[QWEN] 轮询请求失败: {}", e); + continue; + } + }; + + let status = resp.status(); + let body = resp.text().await.unwrap_or_default(); + + if status.is_success() { + // 成功获取 Token + let data: serde_json::Value = serde_json::from_str(&body)?; + + let access_token = data["access_token"] + .as_str() + .ok_or("响应中没有 access_token")? + .to_string(); + let refresh_token = data["refresh_token"].as_str().map(|s| s.to_string()); + let token_type = data["token_type"].as_str().map(|s| s.to_string()); + let resource_url = data["resource_url"].as_str().map(|s| s.to_string()); + let expires_in = data["expires_in"].as_i64().unwrap_or(3600); + + let expires_at = chrono::Utc::now() + chrono::Duration::seconds(expires_in); + + let credentials = QwenCredentials { + access_token: Some(access_token), + refresh_token, + token_type, + resource_url, + expiry_date: Some(expires_at.timestamp_millis()), + expire: Some(expires_at.to_rfc3339()), + last_refresh: Some(chrono::Utc::now().to_rfc3339()), + cred_type: "qwen".to_string(), + }; + + tracing::info!("[QWEN] Token 获取成功"); + return Ok(credentials); + } + + // 解析错误响应 + if let Ok(error_data) = serde_json::from_str::(&body) { + let error_type = error_data["error"].as_str().unwrap_or(""); + + match error_type { + "authorization_pending" => { + // 用户尚未完成授权，继续轮询 + tracing::debug!("[QWEN] 等待用户授权..."); + continue; + } + "slow_down" => { + // 轮询太频繁，增加间隔 + tracing::debug!("[QWEN] 服务器要求降低轮询频率"); + tokio::time::sleep(std::time::Duration::from_secs(5)).await; + continue; + } + "expired_token" => { + return Err("Device code 已过期，请重新开始授权流程".into()); + } + "access_denied" => { + return Err("用户拒绝了授权请求".into()); + } + _ => { + let error_desc = error_data["error_description"] + .as_str() + .unwrap_or("未知错误"); + return Err(format!("Token 获取失败: {} - {}", error_type, error_desc).into()); + } + } + } + + // 其他错误 + if status.as_u16() != 400 { + return Err(format!("Token 请求失败: {} - {}", status, body).into()); + } + } + + Err("授权超时，请重新开始授权流程".into()) +} + +/// 启动 Qwen Device Code Flow 登录 +pub async fn start_qwen_device_code_login() -> Result> +{ + let client = Client::builder() + .timeout(std::time::Duration::from_secs(30)) + .build()?; + + // 发起 Device Code Flow + let (device_response, code_verifier) = initiate_device_flow(&client).await?; + + // 打开浏览器 + let verification_url = device_response + .verification_uri_complete + .as_ref() + .unwrap_or(&device_response.verification_uri); + + tracing::info!("[QWEN] 打开浏览器进行授权: {}", verification_url); + + if let Err(e) = open::that(verification_url) { + tracing::warn!("[QWEN] 无法打开浏览器: {}. 请手动打开 URL.", e); + } + + // 轮询 Token + let credentials = poll_for_token( + &client, + &device_response.device_code, + &code_verifier, + device_response.interval as u64, + 60, // 最多轮询 60 次（约 5 分钟） + ) + .await?; + + // 保存凭证到应用数据目录 + let creds_dir = dirs::data_dir() + .unwrap_or_else(|| std::path::PathBuf::from(".")) + .join("proxycast") + .join("credentials") + .join("qwen"); + + std::fs::create_dir_all(&creds_dir)?; + + let uuid = Uuid::new_v4().to_string(); + let timestamp = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_secs(); + let filename = format!("qwen_{}_{}.json", &uuid[..8], timestamp); + let creds_file_path = creds_dir.join(&filename); + + let creds_json = serde_json::to_string_pretty(&credentials)?; + std::fs::write(&creds_file_path, &creds_json)?; + + tracing::info!("[QWEN] 凭证已保存到: {:?}", creds_file_path); + + Ok(QwenOAuthResult { + credentials, + creds_file_path: creds_file_path.to_string_lossy().to_string(), + }) +} + +/// 启动 Qwen Device Code Flow 并返回设备码信息（不自动打开浏览器） +pub async fn start_qwen_device_code_and_get_info() -> Result< + ( + DeviceCodeResponse, + impl std::future::Future>>, + ), + Box, +> { + let client = Client::builder() + .timeout(std::time::Duration::from_secs(30)) + .build()?; + + // 发起 Device Code Flow + let (device_response, code_verifier) = initiate_device_flow(&client).await?; + + let device_code = device_response.device_code.clone(); + let interval = device_response.interval as u64; + + // 创建等待 future + let wait_future = async move { + // 轮询 Token + let credentials = + poll_for_token(&client, &device_code, &code_verifier, interval, 60).await?; + + // 保存凭证到应用数据目录 + let creds_dir = dirs::data_dir() + .unwrap_or_else(|| std::path::PathBuf::from(".")) + .join("proxycast") + .join("credentials") + .join("qwen"); + + std::fs::create_dir_all(&creds_dir)?; + + let uuid = Uuid::new_v4().to_string(); + let timestamp = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_secs(); + let filename = format!("qwen_{}_{}.json", &uuid[..8], timestamp); + let creds_file_path = creds_dir.join(&filename); + + let creds_json = serde_json::to_string_pretty(&credentials)?; + std::fs::write(&creds_file_path, &creds_json)?; + + tracing::info!("[QWEN] 凭证已保存到: {:?}", creds_file_path); + + Ok(QwenOAuthResult { + credentials, + creds_file_path: creds_file_path.to_string_lossy().to_string(), + }) + }; + + Ok((device_response, wait_future)) +} + +// ============================================================================ +// CredentialProvider Trait 实现 +// ============================================================================ + +#[async_trait] +impl CredentialProvider for QwenProvider { + async fn load_credentials_from_path(&mut self, path: &str) -> ProviderResult<()> { + QwenProvider::load_credentials_from_path(self, path).await + } + + async fn save_credentials(&self) -> ProviderResult<()> { + QwenProvider::save_credentials(self).await + } + + fn is_token_valid(&self) -> bool { + QwenProvider::is_token_valid(self) + } + + fn is_token_expiring_soon(&self) -> bool { + // Qwen 使用与 is_token_valid 相同的逻辑，但阈值为 10 分钟 + if self.credentials.access_token.is_none() { + return true; + } + + if let Some(expire_str) = &self.credentials.expire { + if let Ok(expires) = chrono::DateTime::parse_from_rfc3339(expire_str) { + let now = chrono::Utc::now(); + return expires <= now + chrono::Duration::minutes(10); + } + } + + false + } + + async fn refresh_token(&mut self) -> ProviderResult { + QwenProvider::refresh_token(self).await + } + + fn get_access_token(&self) -> Option<&str> { + self.credentials.access_token.as_deref() + } + + fn provider_type(&self) -> &'static str { + "qwen" + } +} diff --git a/src-tauri/src/providers/traits.rs b/src-tauri/src/providers/traits.rs new file mode 100644 index 000000000..7f8024f4b --- /dev/null +++ b/src-tauri/src/providers/traits.rs @@ -0,0 +1,182 @@ +//! Provider Trait 定义 +//! +//! 统一的 Provider 接口，用于凭证管理和 Token 生命周期管理。 + +use async_trait::async_trait; +use std::error::Error; + +/// Provider 结果类型别名（与现有方法签名兼容） +pub type ProviderResult = Result>; + +/// 凭证管理 Trait +/// +/// 定义所有 OAuth Provider 必须实现的凭证管理接口 +#[async_trait] +pub trait CredentialProvider: Send + Sync { + /// 从指定路径加载凭证 + async fn load_credentials_from_path(&mut self, path: &str) -> ProviderResult<()>; + + /// 保存凭证到文件 + async fn save_credentials(&self) -> ProviderResult<()>; + + /// 检查 Token 是否有效（未过期） + fn is_token_valid(&self) -> bool; + + /// 检查 Token 是否即将过期（通常提前 5 分钟） + fn is_token_expiring_soon(&self) -> bool; + + /// 刷新 Token + /// + /// 返回新的 access_token + async fn refresh_token(&mut self) -> ProviderResult; + + /// 获取当前 access_token + fn get_access_token(&self) -> Option<&str>; + + /// 获取 Provider 类型名称 + fn provider_type(&self) -> &'static str; +} + +/// Token 管理辅助 Trait +/// +/// 提供带重试的 Token 刷新功能 +#[async_trait] +pub trait TokenManager: CredentialProvider { + /// 带重试的 Token 刷新 + /// + /// # Arguments + /// * `max_retries` - 最大重试次数 + /// * `retry_delay_ms` - 重试间隔（毫秒） + async fn refresh_token_with_retry( + &mut self, + max_retries: u32, + retry_delay_ms: u64, + ) -> ProviderResult { + let mut last_error = None; + + for attempt in 0..=max_retries { + match self.refresh_token().await { + Ok(token) => return Ok(token), + Err(e) => { + tracing::warn!( + "[{}] Token refresh attempt {} failed: {}", + self.provider_type(), + attempt + 1, + e + ); + last_error = Some(e); + if attempt < max_retries { + tokio::time::sleep(tokio::time::Duration::from_millis(retry_delay_ms)) + .await; + } + } + } + } + + Err(last_error.unwrap_or_else(|| "Token refresh failed".into())) + } + + /// 确保 Token 有效（如需要则刷新） + async fn ensure_valid_token(&mut self) -> ProviderResult { + if !self.is_token_valid() || self.is_token_expiring_soon() { + self.refresh_token().await + } else { + self.get_access_token() + .map(|s| s.to_string()) + .ok_or_else(|| "No access token available".into()) + } + } +} + +// 为所有实现了 CredentialProvider 的类型自动实现 TokenManager +impl TokenManager for T {} + +#[cfg(test)] +mod tests { + use super::*; + + // Mock Provider for testing + struct MockProvider { + token: Option, + valid: bool, + expiring_soon: bool, + refresh_count: u32, + } + + #[async_trait] + impl CredentialProvider for MockProvider { + async fn load_credentials_from_path(&mut self, _path: &str) -> ProviderResult<()> { + Ok(()) + } + + async fn save_credentials(&self) -> ProviderResult<()> { + Ok(()) + } + + fn is_token_valid(&self) -> bool { + self.valid + } + + fn is_token_expiring_soon(&self) -> bool { + self.expiring_soon + } + + async fn refresh_token(&mut self) -> ProviderResult { + self.refresh_count += 1; + self.token = Some(format!("new_token_{}", self.refresh_count)); + self.valid = true; + self.expiring_soon = false; + Ok(self.token.clone().unwrap()) + } + + fn get_access_token(&self) -> Option<&str> { + self.token.as_deref() + } + + fn provider_type(&self) -> &'static str { + "mock" + } + } + + #[tokio::test] + async fn test_ensure_valid_token_when_valid() { + let mut provider = MockProvider { + token: Some("existing_token".to_string()), + valid: true, + expiring_soon: false, + refresh_count: 0, + }; + + let token = provider.ensure_valid_token().await.unwrap(); + assert_eq!(token, "existing_token"); + assert_eq!(provider.refresh_count, 0); + } + + #[tokio::test] + async fn test_ensure_valid_token_when_expiring() { + let mut provider = MockProvider { + token: Some("old_token".to_string()), + valid: true, + expiring_soon: true, + refresh_count: 0, + }; + + let token = provider.ensure_valid_token().await.unwrap(); + assert_eq!(token, "new_token_1"); + assert_eq!(provider.refresh_count, 1); + } + + #[tokio::test] + async fn test_ensure_valid_token_when_invalid() { + let mut provider = MockProvider { + token: Some("invalid_token".to_string()), + valid: false, + expiring_soon: false, + refresh_count: 0, + }; + + let token = provider.ensure_valid_token().await.unwrap(); + assert_eq!(token, "new_token_1"); + assert_eq!(provider.refresh_count, 1); + } +} diff --git a/src-tauri/src/server.rs b/src-tauri/src/server.rs deleted file mode 100644 index 1582b851b..000000000 --- a/src-tauri/src/server.rs +++ /dev/null @@ -1,5779 +0,0 @@ -//! HTTP API 服务器 -use crate::config::{ - Config, ConfigChangeEvent, ConfigChangeKind, ConfigManager, FileWatcher, HotReloadManager, - ReloadResult, -}; -use crate::converter::anthropic_to_openai::convert_anthropic_to_openai; -use crate::converter::openai_to_antigravity::{ - convert_antigravity_to_openai_response, convert_openai_to_antigravity, -}; -use crate::credential::CredentialSyncService; -use crate::database::dao::provider_pool::ProviderPoolDao; -use crate::database::DbConnection; -use crate::injection::Injector; -use crate::logger::LogStore; -use crate::models::anthropic::*; -use crate::models::openai::*; -use crate::models::route_model::{RouteInfo, RouteListResponse}; -use crate::processor::{RequestContext, RequestProcessor}; -use crate::providers::antigravity::AntigravityProvider; -use crate::providers::claude_custom::ClaudeCustomProvider; -use crate::providers::gemini::GeminiProvider; -use crate::providers::kiro::KiroProvider; -use crate::providers::openai_custom::OpenAICustomProvider; -use crate::providers::qwen::QwenProvider; -use crate::providers::vertex::VertexProvider; -use crate::services::backup_service::BackupService; -use crate::services::provider_pool_service::ProviderPoolService; -use crate::services::token_cache_service::TokenCacheService; -use crate::telemetry::{RequestLog, RequestStatus}; -use crate::websocket::{WsConfig, WsConnectionManager, WsStats}; -use axum::{ - body::Body, - extract::{DefaultBodyLimit, Path, State}, - http::{header, HeaderMap, StatusCode}, - response::{IntoResponse, Response}, - routing::{get, post}, - Json, Router, -}; -use chrono::Utc; -use fs2::available_space; -use futures::stream; -use serde::{Deserialize, Serialize}; -use std::collections::HashMap; -use std::path::PathBuf; -use std::sync::Arc; -use subtle::ConstantTimeEq; -use tokio::sync::{mpsc, oneshot, RwLock}; - -/// 安全截断字符串到指定字符数，避免 UTF-8 边界问题 -fn safe_truncate(s: &str, max_chars: usize) -> String { - let chars: Vec = s.chars().collect(); - if chars.len() <= max_chars { - s.to_string() - } else { - chars[..max_chars].iter().collect() - } -} - -/// 计算 MessageContent 的字符长度 -fn message_content_len(content: &crate::models::openai::MessageContent) -> usize { - use crate::models::openai::{ContentPart, MessageContent}; - match content { - MessageContent::Text(s) => s.len(), - MessageContent::Parts(parts) => parts - .iter() - .filter_map(|p| { - if let ContentPart::Text { text } = p { - Some(text.len()) - } else { - None - } - }) - .sum(), - } -} - -fn api_key_matches(provided_key: &str, expected_key: &str) -> bool { - provided_key - .as_bytes() - .ct_eq(expected_key.as_bytes()) - .into() -} - -/// 记录请求统计到遥测系统 -fn record_request_telemetry( - state: &AppState, - ctx: &RequestContext, - status: crate::telemetry::RequestStatus, - error_message: Option, -) { - use crate::telemetry::RequestLog; - - let provider = ctx.provider.unwrap_or(crate::ProviderType::Kiro); - let mut log = RequestLog::new( - ctx.request_id.clone(), - provider, - ctx.resolved_model.clone(), - ctx.is_stream, - ); - - // 设置状态和持续时间 - match status { - crate::telemetry::RequestStatus::Success => log.mark_success(ctx.elapsed_ms(), 200), - crate::telemetry::RequestStatus::Failed => log.mark_failed( - ctx.elapsed_ms(), - None, - error_message.clone().unwrap_or_default(), - ), - crate::telemetry::RequestStatus::Timeout => log.mark_timeout(ctx.elapsed_ms()), - crate::telemetry::RequestStatus::Cancelled => log.mark_cancelled(ctx.elapsed_ms()), - crate::telemetry::RequestStatus::Retrying => { - log.duration_ms = ctx.elapsed_ms(); - } - } - - // 设置凭证 ID - if let Some(cred_id) = &ctx.credential_id { - log.set_credential_id(cred_id.clone()); - } - - // 设置重试次数 - log.retry_count = ctx.retry_count; - - // 记录到统计聚合器 - { - let stats = state.processor.stats.write(); - stats.record(log.clone()); - } - - // 记录到请求日志记录器（用于前端日志列表显示） - if let Some(logger) = &state.request_logger { - let _ = logger.record(log.clone()); - } - - tracing::info!( - "[TELEMETRY] request_id={} provider={:?} model={} status={:?} duration_ms={}", - ctx.request_id, - provider, - ctx.resolved_model, - status, - ctx.elapsed_ms() - ); -} - -/// 记录 Token 使用量到遥测系统 -fn record_token_usage( - state: &AppState, - ctx: &RequestContext, - input_tokens: Option, - output_tokens: Option, -) { - use crate::telemetry::{TokenSource, TokenUsageRecord}; - - // 只有当至少有一个 Token 值时才记录 - if input_tokens.is_none() && output_tokens.is_none() { - return; - } - - let provider = ctx.provider.unwrap_or(crate::ProviderType::Kiro); - let record = TokenUsageRecord::new( - uuid::Uuid::new_v4().to_string(), - provider, - ctx.resolved_model.clone(), - input_tokens.unwrap_or(0), - output_tokens.unwrap_or(0), - TokenSource::Actual, - ) - .with_request_id(ctx.request_id.clone()); - - // 记录到 Token 追踪器 - { - let tokens = state.processor.tokens.write(); - tokens.record(record); - } - - tracing::debug!( - "[TOKEN] request_id={} input={} output={}", - ctx.request_id, - input_tokens.unwrap_or(0), - output_tokens.unwrap_or(0) - ); -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct ServerStatus { - pub running: bool, - pub host: String, - pub port: u16, - pub requests: u64, - pub uptime_secs: u64, -} - -pub struct ServerState { - pub config: Config, - pub running: bool, - pub requests: u64, - pub start_time: Option, - pub kiro_provider: KiroProvider, - pub gemini_provider: GeminiProvider, - pub qwen_provider: QwenProvider, - pub openai_custom_provider: OpenAICustomProvider, - pub claude_custom_provider: ClaudeCustomProvider, - pub default_provider_ref: Arc>, - shutdown_tx: Option>, - /// 服务器运行时使用的 API key（启动时从配置复制） - /// 用于 test_api 命令，确保测试使用的 API key 和服务器一致 - pub running_api_key: Option, -} - -impl ServerState { - pub fn new(config: Config) -> Self { - let kiro = KiroProvider::new(); - let gemini = GeminiProvider::new(); - let qwen = QwenProvider::new(); - let openai_custom = OpenAICustomProvider::new(); - let claude_custom = ClaudeCustomProvider::new(); - let default_provider_ref = Arc::new(RwLock::new(config.default_provider.clone())); - - Self { - config, - running: false, - requests: 0, - start_time: None, - kiro_provider: kiro, - gemini_provider: gemini, - qwen_provider: qwen, - openai_custom_provider: openai_custom, - claude_custom_provider: claude_custom, - default_provider_ref, - shutdown_tx: None, - running_api_key: None, - } - } - - pub fn status(&self) -> ServerStatus { - ServerStatus { - running: self.running, - host: self.config.server.host.clone(), - port: self.config.server.port, - requests: self.requests, - uptime_secs: self.start_time.map(|t| t.elapsed().as_secs()).unwrap_or(0), - } - } - - pub async fn start( - &mut self, - logs: Arc>, - pool_service: Arc, - token_cache: Arc, - db: Option, - ) -> Result<(), Box> { - self.start_with_telemetry(logs, pool_service, token_cache, db, None, None, None) - .await - } - - /// 启动服务器（使用共享的遥测实例） - /// - /// 这允许服务器与 TelemetryState 共享同一个 StatsAggregator、TokenTracker 和 RequestLogger， - /// 使得请求处理过程中记录的统计数据能够在前端监控页面中显示。 - pub async fn start_with_telemetry( - &mut self, - logs: Arc>, - pool_service: Arc, - token_cache: Arc, - db: Option, - shared_stats: Option>>, - shared_tokens: Option>>, - shared_logger: Option>, - ) -> Result<(), Box> { - if self.running { - return Ok(()); - } - - let (tx, rx) = oneshot::channel(); - self.shutdown_tx = Some(tx); - - let host = self.config.server.host.clone(); - let port = self.config.server.port; - let api_key = self.config.server.api_key.clone(); - let api_key_for_state = api_key.clone(); // 用于保存到 running_api_key - let default_provider_ref = self.default_provider_ref.clone(); - - if api_key.trim().is_empty() { - return Err("API Key 不能为空".into()); - } - - if !is_localhost_host(&host) { - return Err("当前版本仅支持本地监听，请使用 127.0.0.1/localhost/::1".into()); - } - - if (!is_localhost_host(&host) || self.config.remote_management.allow_remote) - && crate::config::is_default_api_key(&api_key) - { - return Err("非本地访问场景下禁止使用默认 API Key，请设置强口令".into()); - } - - if self.config.server.tls.enable { - return Err("当前版本暂不支持 TLS，请关闭 TLS 配置".into()); - } - - if self.config.remote_management.allow_remote { - return Err("当前版本未启用 TLS，禁止开启远程管理".into()); - } - - tracing::warn!("当前未启用 TLS，生产环境请使用反向代理终止 HTTPS"); - - // 重新加载凭证 - let _ = self.kiro_provider.load_credentials().await; - let kiro = self.kiro_provider.clone(); - - // 创建参数注入器 - let injection_enabled = self.config.injection.enabled; - let injector = Injector::with_rules( - self.config - .injection - .rules - .iter() - .map(|r| r.clone().into()) - .collect(), - ); - - // 获取配置和配置路径用于热重载 - let config = self.config.clone(); - let config_path = crate::config::ConfigManager::default_config_path(); - - tokio::spawn(async move { - if let Err(e) = run_server( - &host, - port, - &api_key, - default_provider_ref, - kiro, - logs, - rx, - pool_service, - token_cache, - db, - injector, - injection_enabled, - shared_stats, - shared_tokens, - shared_logger, - Some(config), - Some(config_path), - ) - .await - { - tracing::error!("Server error: {}", e); - } - }); - - self.running = true; - self.start_time = Some(std::time::Instant::now()); - // 保存服务器运行时使用的 API key，用于 test_api 命令 - self.running_api_key = Some(api_key_for_state); - Ok(()) - } - - pub async fn stop(&mut self) { - if let Some(tx) = self.shutdown_tx.take() { - let _ = tx.send(()); - } - self.running = false; - self.start_time = None; - self.running_api_key = None; - } -} - -fn is_localhost_host(host: &str) -> bool { - if host == "localhost" { - return true; - } - host.parse::() - .map(|addr| addr.is_loopback()) - .unwrap_or(false) -} - -impl Clone for KiroProvider { - fn clone(&self) -> Self { - Self { - credentials: self.credentials.clone(), - client: reqwest::Client::new(), - creds_path: self.creds_path.clone(), - } - } -} - -#[derive(Clone)] -#[allow(dead_code)] -struct AppState { - api_key: String, - base_url: String, - default_provider: Arc>, - config: Option, - config_manager: Option>>, - start_time: std::time::Instant, - kiro: Arc>, - logs: Arc>, - kiro_refresh_lock: Arc>, - gemini_refresh_lock: Arc>, - qwen_refresh_lock: Arc>, - pool_service: Arc, - token_cache: Arc, - db: Option, - /// 参数注入器 - injector: Arc>, - /// 是否启用参数注入 - injection_enabled: Arc>, - /// 请求处理器 - processor: Arc, - /// WebSocket 连接管理器 - ws_manager: Arc, - /// WebSocket 统计信息 - ws_stats: Arc, - /// 热重载管理器 - hot_reload_manager: Option>, - /// 请求日志记录器（与 TelemetryState 共享） - request_logger: Option>, - /// Amp CLI 路由器 - amp_router: Arc, - /// 备份服务 - backup_service: Option>, -} - -/// 启动配置文件监控 -/// -/// 监控配置文件变化并触发热重载。 -/// -/// # 连接保持 -/// -/// 热重载过程不会中断现有连接： -/// - 配置更新在独立的 tokio 任务中异步执行 -/// - 使用 RwLock 进行原子性更新，不会阻塞正在处理的请求 -/// - 服务器继续运行，不需要重启 -/// - HTTP 和 WebSocket 连接保持活跃 -async fn start_config_watcher( - config_path: PathBuf, - hot_reload_manager: Option>, - processor: Arc, - logs: Arc>, - db: Option, - config_manager: Option>>, -) -> Option { - let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::(); - - // 创建文件监控器 - let mut watcher = match FileWatcher::new(&config_path, tx) { - Ok(w) => w, - Err(e) => { - tracing::error!("[HOT_RELOAD] 创建文件监控器失败: {}", e); - return None; - } - }; - - // 启动监控 - if let Err(e) = watcher.start() { - tracing::error!("[HOT_RELOAD] 启动文件监控失败: {}", e); - return None; - } - - tracing::info!("[HOT_RELOAD] 配置文件监控已启动: {:?}", config_path); - - // 启动事件处理任务 - let hot_reload_manager_clone = hot_reload_manager.clone(); - let processor_clone = processor.clone(); - let logs_clone = logs.clone(); - let db_clone = db.clone(); - let config_manager_clone = config_manager.clone(); - - tokio::spawn(async move { - while let Some(event) = rx.recv().await { - // 只处理修改事件 - if event.kind != ConfigChangeKind::Modified { - continue; - } - - tracing::info!("[HOT_RELOAD] 检测到配置文件变更: {:?}", event.path); - logs_clone.write().await.add( - "info", - &format!("[HOT_RELOAD] 检测到配置文件变更: {:?}", event.path), - ); - - // 执行热重载 - if let Some(ref manager) = hot_reload_manager_clone { - let result = manager.reload(); - match &result { - ReloadResult::Success { .. } => { - tracing::info!("[HOT_RELOAD] 配置热重载成功"); - logs_clone - .write() - .await - .add("info", "[HOT_RELOAD] 配置热重载成功"); - - // 更新处理器中的组件 - let new_config = manager.config(); - update_processor_config(&processor_clone, &new_config).await; - - // 同步凭证池 - if let (Some(ref db), Some(ref cfg_manager)) = - (&db_clone, &config_manager_clone) - { - match sync_credential_pool_from_config(db, cfg_manager, &logs_clone) - .await - { - Ok(count) => { - tracing::info!( - "[HOT_RELOAD] 凭证池同步完成，共 {} 个凭证", - count - ); - logs_clone.write().await.add( - "info", - &format!( - "[HOT_RELOAD] 凭证池同步完成，共 {} 个凭证", - count - ), - ); - } - Err(e) => { - tracing::warn!("[HOT_RELOAD] 凭证池同步失败: {}", e); - logs_clone.write().await.add( - "warn", - &format!("[HOT_RELOAD] 凭证池同步失败: {}", e), - ); - } - } - } - } - ReloadResult::RolledBack { error, .. } => { - tracing::warn!("[HOT_RELOAD] 配置热重载失败，已回滚: {}", error); - logs_clone.write().await.add( - "warn", - &format!("[HOT_RELOAD] 配置热重载失败，已回滚: {}", error), - ); - } - ReloadResult::Failed { - error, - rollback_error, - .. - } => { - tracing::error!( - "[HOT_RELOAD] 配置热重载失败: {}, 回滚错误: {:?}", - error, - rollback_error - ); - logs_clone.write().await.add( - "error", - &format!( - "[HOT_RELOAD] 配置热重载失败: {}, 回滚错误: {:?}", - error, rollback_error - ), - ); - } - } - } - } - }); - - Some(watcher) -} - -/// 更新处理器配置 -/// -/// 当配置热重载成功后，更新 RequestProcessor 中的各个组件。 -/// -/// # 原子性更新 -/// -/// 每个组件的更新都是原子性的，使用 RwLock 确保： -/// - 正在处理的请求不会看到部分更新的状态 -/// - 更新过程不会阻塞新请求的处理 -/// - 现有连接不受影响 -async fn update_processor_config(processor: &RequestProcessor, config: &Config) { - let _reload_guard = processor.reload_lock.write().await; - // 更新注入器规则 - { - let mut injector = processor.injector.write().await; - injector.clear(); - for rule in &config.injection.rules { - injector.add_rule(rule.clone().into()); - } - tracing::debug!( - "[HOT_RELOAD] 注入器规则已更新: {} 条规则", - config.injection.rules.len() - ); - } - - // 更新路由器规则 - { - let mut router = processor.router.write().await; - router.clear_rules(); - for rule in &config.routing.rules { - // 解析 provider 字符串为 ProviderType - if let Ok(provider_type) = rule.provider.parse::() { - router.add_rule(crate::router::RoutingRule { - pattern: rule.pattern.clone(), - target_provider: provider_type, - priority: rule.priority, - enabled: true, - }); - } else { - tracing::warn!("[HOT_RELOAD] 无法解析 provider: {}", rule.provider); - } - } - tracing::debug!( - "[HOT_RELOAD] 路由规则已更新: {} 条规则", - config.routing.rules.len() - ); - } - - // 更新模型映射器 - { - let mut mapper = processor.mapper.write().await; - mapper.clear(); - for (alias, model) in &config.routing.model_aliases { - mapper.add_alias(alias, model); - } - tracing::debug!( - "[HOT_RELOAD] 模型别名已更新: {} 个别名", - config.routing.model_aliases.len() - ); - } - - // 注意：重试配置目前不支持热更新，因为 Retrier 是不可变的 - // 如果需要更新重试配置，需要重启服务器 - tracing::debug!( - "[HOT_RELOAD] 重试配置: max_retries={}, base_delay={}ms (需重启生效)", - config.retry.max_retries, - config.retry.base_delay_ms - ); - - tracing::info!("[HOT_RELOAD] 处理器配置更新完成"); -} - -/// 从配置同步凭证池 -/// -/// 当配置热重载成功后，从 YAML 配置中加载凭证并同步到数据库。 -/// -/// # 同步策略 -/// -/// - 从配置中加载所有凭证 -/// - 对于配置中存在但数据库中不存在的凭证，添加到数据库 -/// - 对于配置中存在且数据库中也存在的凭证，更新数据库中的记录 -/// - 对于数据库中存在但配置中不存在的凭证，保留（不删除，避免丢失运行时状态） -async fn sync_credential_pool_from_config( - db: &DbConnection, - config_manager: &Arc>, - _logs: &Arc>, -) -> Result { - // 创建凭证同步服务 - let sync_service = CredentialSyncService::new(config_manager.clone()); - - // 从配置加载凭证 - let credentials = sync_service.load_from_config().map_err(|e| e.to_string())?; - - let conn = db.lock().map_err(|e| e.to_string())?; - let mut synced_count = 0; - - for cred in &credentials { - // 检查凭证是否已存在 - let existing = - ProviderPoolDao::get_by_uuid(&conn, &cred.uuid).map_err(|e| e.to_string())?; - - if existing.is_some() { - // 更新现有凭证 - ProviderPoolDao::update(&conn, cred).map_err(|e| e.to_string())?; - tracing::debug!( - "[HOT_RELOAD] 更新凭证: {} ({})", - cred.uuid, - cred.provider_type - ); - } else { - // 添加新凭证 - ProviderPoolDao::insert(&conn, cred).map_err(|e| e.to_string())?; - tracing::debug!( - "[HOT_RELOAD] 添加凭证: {} ({})", - cred.uuid, - cred.provider_type - ); - } - synced_count += 1; - } - - Ok(synced_count) -} - -async fn run_server( - host: &str, - port: u16, - api_key: &str, - default_provider: Arc>, - kiro: KiroProvider, - logs: Arc>, - shutdown: oneshot::Receiver<()>, - pool_service: Arc, - token_cache: Arc, - db: Option, - injector: Injector, - injection_enabled: bool, - shared_stats: Option>>, - shared_tokens: Option>>, - shared_logger: Option>, - config: Option, - config_path: Option, -) -> Result<(), Box> { - let base_url = format!("http://{}:{}", host, port); - - // 创建请求处理器（使用共享的遥测实例或默认实例） - let processor = match (shared_stats, shared_tokens) { - (Some(stats), Some(tokens)) => Arc::new(RequestProcessor::with_shared_telemetry( - pool_service.clone(), - stats, - tokens, - )), - _ => Arc::new(RequestProcessor::with_defaults(pool_service.clone())), - }; - - // 将注入器规则同步到处理器 - { - let mut proc_injector = processor.injector.write().await; - for rule in injector.rules() { - proc_injector.add_rule(rule.clone()); - } - } - - // 初始化 WebSocket 管理器 - let ws_manager = Arc::new(WsConnectionManager::new(WsConfig::default())); - let ws_stats = ws_manager.stats().clone(); - - // 初始化热重载管理器 - let hot_reload_manager = match (&config, &config_path) { - (Some(cfg), Some(path)) => Some(Arc::new(HotReloadManager::new(cfg.clone(), path.clone()))), - _ => None, - }; - - // 初始化配置管理器（用于凭证池同步） - let config_manager: Option>> = - match (&config, &config_path) { - (Some(cfg), Some(path)) => Some(Arc::new(std::sync::RwLock::new( - ConfigManager::with_config(cfg.clone(), path.clone()), - ))), - _ => None, - }; - - let logs_clone = logs.clone(); - let db_clone = db.clone(); - - // 初始化 Amp CLI 路由器 - let amp_router = Arc::new(crate::router::AmpRouter::new( - config - .as_ref() - .map(|c| c.ampcode.clone()) - .unwrap_or_default(), - )); - - let backup_service = match BackupService::with_defaults() { - Ok(service) => { - tracing::info!( - "[BACKUP] 备份服务初始化成功，备份目录: {:?}", - service.backup_dir() - ); - Some(Arc::new(service)) - } - Err(e) => { - tracing::warn!("[BACKUP] 备份服务初始化失败，自动备份将不可用: {}", e); - None - } - }; - if let Some(service) = backup_service.clone() { - let db_for_backup = db.clone(); - tokio::spawn(async move { - let mut ticker = tokio::time::interval(std::time::Duration::from_secs(24 * 60 * 60)); - loop { - ticker.tick().await; - let result = match &db_for_backup { - Some(db) => service.backup_database_with_connection(db), - None => service.backup_database(), - }; - match result { - Ok(path) => tracing::info!("[BACKUP] 自动备份成功: {:?}", path), - Err(err) => tracing::warn!("[BACKUP] 自动备份失败: {}", err), - } - } - }); - } - - let state = AppState { - api_key: api_key.to_string(), - base_url, - default_provider, - config: config.clone(), - config_manager: config_manager.clone(), - start_time: std::time::Instant::now(), - kiro: Arc::new(RwLock::new(kiro)), - logs, - kiro_refresh_lock: Arc::new(tokio::sync::Mutex::new(())), - gemini_refresh_lock: Arc::new(tokio::sync::Mutex::new(())), - qwen_refresh_lock: Arc::new(tokio::sync::Mutex::new(())), - pool_service, - token_cache, - db, - injector: Arc::new(RwLock::new(injector)), - injection_enabled: Arc::new(RwLock::new(injection_enabled)), - processor: processor.clone(), - ws_manager, - ws_stats, - hot_reload_manager: hot_reload_manager.clone(), - request_logger: shared_logger, - amp_router, - backup_service, - }; - - // 启动配置文件监控 - let _file_watcher = if let Some(path) = config_path { - start_config_watcher( - path, - hot_reload_manager, - processor, - logs_clone, - db_clone, - config_manager, - ) - .await - } else { - None - }; - - // P1 安全修复：降低默认请求体大小限制，防止 DoS - // 从 100MB 降低到 20MB，对于需要更大请求的端点单独配置 - let body_limit = 20 * 1024 * 1024; // 20MB - - // 创建管理 API 路由（带认证中间件） - let management_config = config - .as_ref() - .map(|c| c.remote_management.clone()) - .unwrap_or_default(); - - let management_routes = Router::new() - .route("/v0/management/status", get(management_status)) - .route("/v0/management/backup", post(management_backup)) - .route("/v0/management/restore", post(management_restore)) - .route( - "/v0/management/credentials", - get(management_list_credentials), - ) - .route( - "/v0/management/credentials", - post(management_add_credential), - ) - .route("/v0/management/config", get(management_get_config)) - .route( - "/v0/management/config", - axum::routing::put(management_update_config), - ) - .layer(crate::middleware::ManagementAuthLayer::new( - management_config, - )); - - let app = Router::new() - .route("/health", get(health)) - .route("/ready", get(readiness)) - .route("/v1/models", get(models)) - .route("/v1/routes", get(list_routes)) - .route("/v1/chat/completions", post(chat_completions)) - .route("/v1/messages", post(anthropic_messages)) - .route("/v1/messages/count_tokens", post(count_tokens)) - // WebSocket 路由 - .route("/v1/ws", get(ws_upgrade_handler)) - .route("/ws", get(ws_upgrade_handler)) - // 多供应商路由 - .route( - "/:selector/v1/messages", - post(anthropic_messages_with_selector), - ) - .route( - "/:selector/v1/chat/completions", - post(chat_completions_with_selector), - ) - // Amp CLI 路由 - .route( - "/api/provider/:provider/v1/chat/completions", - post(amp_chat_completions), - ) - .route("/api/provider/:provider/v1/messages", post(amp_messages)) - // Amp CLI 管理代理路由 - .route( - "/api/auth/*path", - axum::routing::any(amp_management_proxy_auth), - ) - .route( - "/api/user/*path", - axum::routing::any(amp_management_proxy_user), - ) - // 管理 API 路由 - .merge(management_routes) - .layer(DefaultBodyLimit::max(body_limit)) - .with_state(state); - - let addr: std::net::SocketAddr = format!("{host}:{port}").parse()?; - let listener = tokio::net::TcpListener::bind(addr).await?; - - tracing::info!("Server listening on {}", addr); - - axum::serve( - listener, - app.into_make_service_with_connect_info::(), - ) - .with_graceful_shutdown(async move { - let _ = shutdown.await; - }) - .await?; - - Ok(()) -} - -#[derive(Debug, Serialize)] -struct CheckResult { - status: String, - message: Option, - latency_ms: Option, -} - -#[derive(Debug, Serialize)] -struct HealthStatus { - status: String, - timestamp: chrono::DateTime, - version: String, - checks: HashMap, -} - -async fn health(State(state): State) -> impl IntoResponse { - let mut checks = HashMap::new(); - - let db_check = check_database(&state).await; - checks.insert("database".to_string(), db_check); - - let pool_check = check_credential_pool(&state).await; - checks.insert("credential_pool".to_string(), pool_check); - - let disk_check = check_disk_space(&state).await; - checks.insert("disk_space".to_string(), disk_check); - - let log_check = check_log_directory(&state).await; - checks.insert("log_directory".to_string(), log_check); - - let overall_status = if checks.values().all(|c| c.status == "healthy") { - "healthy" - } else if checks.values().any(|c| c.status == "unhealthy") { - "unhealthy" - } else { - "degraded" - }; - - let status_code = match overall_status { - "healthy" | "degraded" => StatusCode::OK, - "unhealthy" => StatusCode::SERVICE_UNAVAILABLE, - _ => StatusCode::INTERNAL_SERVER_ERROR, - }; - - ( - status_code, - Json(HealthStatus { - status: overall_status.to_string(), - timestamp: Utc::now(), - version: env!("CARGO_PKG_VERSION").to_string(), - checks, - }), - ) -} - -async fn readiness(State(state): State) -> impl IntoResponse { - let db_ok = matches!(check_database(&state).await.status.as_str(), "healthy"); - let pool_ok = matches!( - check_credential_pool(&state).await.status.as_str(), - "healthy" - ); - - if !db_ok || !pool_ok { - return ( - StatusCode::SERVICE_UNAVAILABLE, - Json(serde_json::json!({ - "ready": false, - "reason": "Database or credential pool not ready" - })), - ); - } - - ( - StatusCode::OK, - Json(serde_json::json!({ - "ready": true - })), - ) -} - -async fn check_database(state: &AppState) -> CheckResult { - let start = std::time::Instant::now(); - let Some(db) = &state.db else { - return CheckResult { - status: "unhealthy".to_string(), - message: Some("database not initialized".to_string()), - latency_ms: None, - }; - }; - let ok = db - .lock() - .map(|conn| conn.query_row::("SELECT 1", [], |row| row.get(0))) - .is_ok(); - - CheckResult { - status: if ok { "healthy" } else { "unhealthy" }.to_string(), - message: if ok { - None - } else { - Some("database query failed".to_string()) - }, - latency_ms: Some(start.elapsed().as_millis() as u64), - } -} - -async fn check_credential_pool(state: &AppState) -> CheckResult { - let start = std::time::Instant::now(); - let Some(db) = &state.db else { - return CheckResult { - status: "unhealthy".to_string(), - message: Some("database not initialized".to_string()), - latency_ms: None, - }; - }; - - let stats = match state.pool_service.get_overview(db) { - Ok(items) => { - let mut total = 0usize; - let mut healthy = 0usize; - let mut disabled = 0usize; - for item in items { - total += item.stats.total_count; - healthy += item.stats.healthy_count; - disabled += item.stats.disabled_count; - } - (total, healthy, disabled) - } - Err(_) => (0, 0, 0), - }; - - let (total, healthy, disabled) = stats; - let status = if healthy > 0 { - "healthy" - } else if total > 0 { - "degraded" - } else { - "unhealthy" - }; - - CheckResult { - status: status.to_string(), - message: Some(format!( - "total={} healthy={} disabled={}", - total, healthy, disabled - )), - latency_ms: Some(start.elapsed().as_millis() as u64), - } -} - -async fn check_disk_space(state: &AppState) -> CheckResult { - let start = std::time::Instant::now(); - let Some(log_path) = state.logs.read().await.get_log_file_path() else { - return CheckResult { - status: "degraded".to_string(), - message: Some("log path not available".to_string()), - latency_ms: None, - }; - }; - let log_path = PathBuf::from(log_path); - let dir = log_path.parent().unwrap_or(log_path.as_path()); - - let available = available_space(dir).unwrap_or(0); - let available_gb = available / (1024 * 1024 * 1024); - let status = if available_gb >= 10 { - "healthy" - } else if available_gb >= 1 { - "degraded" - } else { - "unhealthy" - }; - - CheckResult { - status: status.to_string(), - message: Some(format!("available_gb={}", available_gb)), - latency_ms: Some(start.elapsed().as_millis() as u64), - } -} - -async fn check_log_directory(state: &AppState) -> CheckResult { - let start = std::time::Instant::now(); - let Some(log_path) = state.logs.read().await.get_log_file_path() else { - return CheckResult { - status: "degraded".to_string(), - message: Some("log path not available".to_string()), - latency_ms: None, - }; - }; - let log_path = PathBuf::from(log_path); - let dir = log_path.parent().unwrap_or(log_path.as_path()); - - let test_file = dir.join(".proxycast_write_check"); - let writable = std::fs::OpenOptions::new() - .create(true) - .write(true) - .open(&test_file) - .and_then(|_| std::fs::remove_file(&test_file)) - .is_ok(); - - CheckResult { - status: if writable { "healthy" } else { "unhealthy" }.to_string(), - message: if writable { - None - } else { - Some("log directory not writable".to_string()) - }, - latency_ms: Some(start.elapsed().as_millis() as u64), - } -} - -async fn models() -> impl IntoResponse { - Json(serde_json::json!({ - "object": "list", - "data": [ - // Kiro/Claude models - {"id": "claude-sonnet-4-5", "object": "model", "owned_by": "anthropic"}, - {"id": "claude-sonnet-4-5-20250929", "object": "model", "owned_by": "anthropic"}, - {"id": "claude-3-7-sonnet-20250219", "object": "model", "owned_by": "anthropic"}, - {"id": "claude-3-5-sonnet-latest", "object": "model", "owned_by": "anthropic"}, - // Gemini models - {"id": "gemini-2.5-flash", "object": "model", "owned_by": "google"}, - {"id": "gemini-2.5-flash-lite", "object": "model", "owned_by": "google"}, - {"id": "gemini-2.5-pro", "object": "model", "owned_by": "google"}, - {"id": "gemini-2.5-pro-preview-06-05", "object": "model", "owned_by": "google"}, - {"id": "gemini-3-pro-preview", "object": "model", "owned_by": "google"}, - // Qwen models - {"id": "qwen3-coder-plus", "object": "model", "owned_by": "alibaba"}, - {"id": "qwen3-coder-flash", "object": "model", "owned_by": "alibaba"} - ] - })) -} - -async fn verify_api_key( - headers: &HeaderMap, - expected_key: &str, -) -> Result<(), (StatusCode, Json)> { - let auth = headers - .get("authorization") - .or_else(|| headers.get("x-api-key")) - .and_then(|v| v.to_str().ok()); - - let key = match auth { - Some(s) if s.starts_with("Bearer ") => &s[7..], - Some(s) => s, - None => { - return Err(( - StatusCode::UNAUTHORIZED, - Json(serde_json::json!({"error": {"message": "No API key provided"}})), - )) - } - }; - - if !api_key_matches(key, expected_key) { - return Err(( - StatusCode::UNAUTHORIZED, - Json(serde_json::json!({"error": {"message": "Invalid API key"}})), - )); - } - - Ok(()) -} - -/// Anthropic 格式的 API key 验证 -/// 返回 Anthropic 标准错误格式：{"type": "error", "error": {"type": "...", "message": "..."}} -async fn verify_api_key_anthropic( - headers: &HeaderMap, - expected_key: &str, -) -> Result<(), (StatusCode, Json)> { - let auth = headers - .get("x-api-key") - .or_else(|| headers.get("authorization")) - .and_then(|v| v.to_str().ok()); - - let key = match auth { - Some(s) if s.starts_with("Bearer ") => &s[7..], - Some(s) => s, - None => { - return Err(( - StatusCode::UNAUTHORIZED, - Json(serde_json::json!({ - "type": "error", - "error": { - "type": "authentication_error", - "message": "No API key provided. Please set the x-api-key header." - } - })), - )) - } - }; - - if !api_key_matches(key, expected_key) { - return Err(( - StatusCode::UNAUTHORIZED, - Json(serde_json::json!({ - "type": "error", - "error": { - "type": "authentication_error", - "message": "Invalid API key" - } - })), - )); - } - - Ok(()) -} - -async fn chat_completions( - State(state): State, - headers: HeaderMap, - Json(mut request): Json, -) -> Response { - let _reload_guard = state.processor.reload_lock.read().await; - if let Err(e) = verify_api_key(&headers, &state.api_key).await { - state - .logs - .write() - .await - .add("warn", "Unauthorized request to /v1/chat/completions"); - return e.into_response(); - } - - // 创建请求上下文 - let mut ctx = RequestContext::new(request.model.clone()).with_stream(request.stream); - - state.logs.write().await.add( - "info", - &format!( - "POST /v1/chat/completions request_id={} model={} stream={}", - ctx.request_id, request.model, request.stream - ), - ); - - // 使用 RequestProcessor 解析模型别名和路由 - let provider = state.processor.resolve_and_route(&mut ctx).await; - - // 更新请求中的模型名为解析后的模型 - if ctx.resolved_model != ctx.original_model { - request.model = ctx.resolved_model.clone(); - state.logs.write().await.add( - "info", - &format!( - "[MAPPER] request_id={} alias={} -> model={}", - ctx.request_id, ctx.original_model, ctx.resolved_model - ), - ); - } - - // 应用参数注入 - let injection_enabled = *state.injection_enabled.read().await; - if injection_enabled { - let injector = state.processor.injector.read().await; - let mut payload = serde_json::to_value(&request).unwrap_or_default(); - let result = injector.inject(&request.model, &mut payload); - if result.has_injections() { - state.logs.write().await.add( - "info", - &format!( - "[INJECT] request_id={} applied_rules={:?} injected_params={:?}", - ctx.request_id, result.applied_rules, result.injected_params - ), - ); - // 更新请求 - if let Ok(updated) = serde_json::from_value(payload) { - request = updated; - } - } - } - - // 获取当前默认 provider（用于凭证池选择） - let default_provider = state.default_provider.read().await.clone(); - - // 记录路由结果 - state.logs.write().await.add( - "info", - &format!( - "[ROUTE] request_id={} model={} provider={}", - ctx.request_id, ctx.resolved_model, provider - ), - ); - - // 尝试从凭证池中选择凭证 - let credential = match &state.db { - Some(db) => state - .pool_service - .select_credential(db, &default_provider, Some(&request.model)) - .ok() - .flatten(), - None => None, - }; - - // 如果找到凭证池中的凭证，使用它 - if let Some(cred) = credential { - state.logs.write().await.add( - "info", - &format!( - "[ROUTE] Using pool credential: type={} name={:?} uuid={}", - cred.provider_type, - cred.name, - &cred.uuid[..8] - ), - ); - let response = call_provider_openai(&state, &cred, &request).await; - - // 记录请求统计 - let is_success = response.status().is_success(); - let status = if is_success { - crate::telemetry::RequestStatus::Success - } else { - crate::telemetry::RequestStatus::Failed - }; - record_request_telemetry(&state, &ctx, status, None); - - // 如果成功，记录估算的 Token 使用量 - if is_success { - let estimated_input_tokens = request - .messages - .iter() - .map(|m| { - let content_len = match &m.content { - Some(c) => message_content_len(c), - None => 0, - }; - content_len / 4 - }) - .sum::() as u32; - // 输出 Token 使用估算值（假设平均响应长度） - let estimated_output_tokens = 100u32; - record_token_usage( - &state, - &ctx, - Some(estimated_input_tokens), - Some(estimated_output_tokens), - ); - } - - return response; - } - - // 回退到旧的单凭证模式 - state.logs.write().await.add( - "debug", - &format!( - "[ROUTE] No pool credential found for '{}', using legacy mode", - default_provider - ), - ); - - // 检查是否需要刷新 token（无 token 或即将过期） - { - let _guard = state.kiro_refresh_lock.lock().await; - let mut kiro = state.kiro.write().await; - let needs_refresh = - kiro.credentials.access_token.is_none() || kiro.is_token_expiring_soon(); - if needs_refresh { - if let Err(e) = kiro.refresh_token().await { - state - .logs - .write() - .await - .add("error", &format!("Token refresh failed: {e}")); - return ( - StatusCode::UNAUTHORIZED, - Json(serde_json::json!({"error": {"message": format!("Token refresh failed: {e}")}})), - ).into_response(); - } - } - } - - let kiro = state.kiro.read().await; - - match kiro.call_api(&request).await { - Ok(resp) => { - let status = resp.status(); - if status.is_success() { - match resp.text().await { - Ok(body) => { - let parsed = parse_cw_response(&body); - let has_tool_calls = !parsed.tool_calls.is_empty(); - - state.logs.write().await.add( - "info", - &format!( - "Request completed: content_len={}, tool_calls={}", - parsed.content.len(), - parsed.tool_calls.len() - ), - ); - - // 构建消息 - let message = if has_tool_calls { - serde_json::json!({ - "role": "assistant", - "content": if parsed.content.is_empty() { serde_json::Value::Null } else { serde_json::json!(parsed.content) }, - "tool_calls": parsed.tool_calls.iter().map(|tc| { - serde_json::json!({ - "id": tc.id, - "type": "function", - "function": { - "name": tc.function.name, - "arguments": tc.function.arguments - } - }) - }).collect::>() - }) - } else { - serde_json::json!({ - "role": "assistant", - "content": parsed.content - }) - }; - - // 估算 Token 数量（基于字符数，约 4 字符 = 1 token） - let estimated_output_tokens = (parsed.content.len() / 4) as u32; - // 估算输入 Token（基于请求消息） - let estimated_input_tokens = request - .messages - .iter() - .map(|m| { - let content_len = match &m.content { - Some(c) => message_content_len(c), - None => 0, - }; - content_len / 4 - }) - .sum::() - as u32; - - let response = serde_json::json!({ - "id": format!("chatcmpl-{}", uuid::Uuid::new_v4()), - "object": "chat.completion", - "created": std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap_or_default() - .as_secs(), - "model": request.model, - "choices": [{ - "index": 0, - "message": message, - "finish_reason": if has_tool_calls { "tool_calls" } else { "stop" } - }], - "usage": { - "prompt_tokens": estimated_input_tokens, - "completion_tokens": estimated_output_tokens, - "total_tokens": estimated_input_tokens + estimated_output_tokens - } - }); - // 记录成功请求统计 - record_request_telemetry( - &state, - &ctx, - crate::telemetry::RequestStatus::Success, - None, - ); - // 记录 Token 使用量 - record_token_usage( - &state, - &ctx, - Some(estimated_input_tokens), - Some(estimated_output_tokens), - ); - Json(response).into_response() - } - Err(e) => { - // 记录失败请求统计 - record_request_telemetry( - &state, - &ctx, - crate::telemetry::RequestStatus::Failed, - Some(e.to_string()), - ); - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({"error": {"message": e.to_string()}})), - ) - .into_response() - } - } - } else if status.as_u16() == 403 || status.as_u16() == 402 { - // Token 过期或账户问题，尝试重新加载凭证并刷新 - drop(kiro); - let _guard = state.kiro_refresh_lock.lock().await; - let mut kiro = state.kiro.write().await; - state.logs.write().await.add( - "warn", - &format!( - "[AUTH] Got {}, reloading credentials and attempting token refresh...", - status.as_u16() - ), - ); - - // 先重新加载凭证文件（可能用户换了账户） - if let Err(e) = kiro.load_credentials().await { - state.logs.write().await.add( - "error", - &format!("[AUTH] Failed to reload credentials: {e}"), - ); - } - - match kiro.refresh_token().await { - Ok(_) => { - state - .logs - .write() - .await - .add("info", "[AUTH] Token refreshed successfully after reload"); - // 重试请求 - drop(kiro); - let kiro = state.kiro.read().await; - match kiro.call_api(&request).await { - Ok(retry_resp) => { - if retry_resp.status().is_success() { - match retry_resp.text().await { - Ok(body) => { - let parsed = parse_cw_response(&body); - let has_tool_calls = !parsed.tool_calls.is_empty(); - - let message = if has_tool_calls { - serde_json::json!({ - "role": "assistant", - "content": if parsed.content.is_empty() { serde_json::Value::Null } else { serde_json::json!(parsed.content) }, - "tool_calls": parsed.tool_calls.iter().map(|tc| { - serde_json::json!({ - "id": tc.id, - "type": "function", - "function": { - "name": tc.function.name, - "arguments": tc.function.arguments - } - }) - }).collect::>() - }) - } else { - serde_json::json!({ - "role": "assistant", - "content": parsed.content - }) - }; - - let response = serde_json::json!({ - "id": format!("chatcmpl-{}", uuid::Uuid::new_v4()), - "object": "chat.completion", - "created": std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap_or_default() - .as_secs(), - "model": request.model, - "choices": [{ - "index": 0, - "message": message, - "finish_reason": if has_tool_calls { "tool_calls" } else { "stop" } - }], - "usage": { - "prompt_tokens": 0, - "completion_tokens": 0, - "total_tokens": 0 - } - }); - return Json(response).into_response(); - } - Err(e) => return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({"error": {"message": e.to_string()}})), - ).into_response(), - } - } - let body = retry_resp.text().await.unwrap_or_default(); - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({"error": {"message": format!("Retry failed: {}", body)}})), - ).into_response() - } - Err(e) => ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({"error": {"message": e.to_string()}})), - ) - .into_response(), - } - } - Err(e) => { - state - .logs - .write() - .await - .add("error", &format!("[AUTH] Token refresh failed: {e}")); - ( - StatusCode::UNAUTHORIZED, - Json(serde_json::json!({"error": {"message": format!("Token refresh failed: {e}")}})), - ) - .into_response() - } - } - } else { - let body = resp.text().await.unwrap_or_default(); - state.logs.write().await.add( - "error", - &format!("Upstream error {}: {}", status, safe_truncate(&body, 200)), - ); - ( - StatusCode::from_u16(status.as_u16()).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR), - Json(serde_json::json!({"error": {"message": format!("Upstream error: {}", body)}})) - ).into_response() - } - } - Err(e) => { - state - .logs - .write() - .await - .add("error", &format!("API call failed: {e}")); - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({"error": {"message": e.to_string()}})), - ) - .into_response() - } - } -} - -async fn anthropic_messages( - State(state): State, - headers: HeaderMap, - Json(mut request): Json, -) -> Response { - let _reload_guard = state.processor.reload_lock.read().await; - // 使用 Anthropic 格式的认证验证（优先检查 x-api-key） - if let Err(e) = verify_api_key_anthropic(&headers, &state.api_key).await { - state - .logs - .write() - .await - .add("warn", "Unauthorized request to /v1/messages"); - return e.into_response(); - } - - // 创建请求上下文 - let mut ctx = RequestContext::new(request.model.clone()).with_stream(request.stream); - - // 详细记录请求信息 - let msg_count = request.messages.len(); - let has_tools = request.tools.as_ref().map(|t| t.len()).unwrap_or(0); - let has_system = request.system.is_some(); - state.logs.write().await.add( - "info", - &format!( - "[REQ] POST /v1/messages request_id={} model={} stream={} messages={} tools={} has_system={}", - ctx.request_id, request.model, request.stream, msg_count, has_tools, has_system - ), - ); - - // 使用 RequestProcessor 解析模型别名和路由 - let provider = state.processor.resolve_and_route(&mut ctx).await; - - // 更新请求中的模型名为解析后的模型 - if ctx.resolved_model != ctx.original_model { - request.model = ctx.resolved_model.clone(); - state.logs.write().await.add( - "info", - &format!( - "[MAPPER] request_id={} alias={} -> model={}", - ctx.request_id, ctx.original_model, ctx.resolved_model - ), - ); - } - - // 记录最后一条消息的角色和内容预览 - if let Some(last_msg) = request.messages.last() { - let content_preview = match &last_msg.content { - serde_json::Value::String(s) => s.chars().take(100).collect::(), - serde_json::Value::Array(arr) => { - if let Some(first) = arr.first() { - if let Some(text) = first.get("text").and_then(|t| t.as_str()) { - text.chars().take(100).collect::() - } else { - format!("[{} blocks]", arr.len()) - } - } else { - "[empty]".to_string() - } - } - _ => "[unknown]".to_string(), - }; - state.logs.write().await.add( - "debug", - &format!( - "[REQ] request_id={} last_message: role={} content={}", - ctx.request_id, last_msg.role, content_preview - ), - ); - } - - // 应用参数注入 - let injection_enabled = *state.injection_enabled.read().await; - if injection_enabled { - let injector = state.processor.injector.read().await; - let mut payload = serde_json::to_value(&request).unwrap_or_default(); - let result = injector.inject(&request.model, &mut payload); - if result.has_injections() { - state.logs.write().await.add( - "info", - &format!( - "[INJECT] request_id={} applied_rules={:?} injected_params={:?}", - ctx.request_id, result.applied_rules, result.injected_params - ), - ); - // 更新请求 - if let Ok(updated) = serde_json::from_value(payload) { - request = updated; - } - } - } - - // 获取当前默认 provider（用于凭证池选择） - let default_provider = state.default_provider.read().await.clone(); - - // 记录路由结果 - state.logs.write().await.add( - "info", - &format!( - "[ROUTE] request_id={} model={} provider={}", - ctx.request_id, ctx.resolved_model, provider - ), - ); - - // 尝试从凭证池中选择凭证 - let credential = match &state.db { - Some(db) => { - // 根据 default_provider 配置选择凭证 - state - .pool_service - .select_credential(db, &default_provider, Some(&request.model)) - .ok() - .flatten() - } - None => None, - }; - - // 如果找到凭证池中的凭证，使用它 - if let Some(cred) = credential { - state.logs.write().await.add( - "info", - &format!( - "[ROUTE] Using pool credential: type={} name={:?} uuid={}", - cred.provider_type, - cred.name, - &cred.uuid[..8] - ), - ); - let response = call_provider_anthropic(&state, &cred, &request).await; - - // 记录请求统计 - let is_success = response.status().is_success(); - let status = if is_success { - crate::telemetry::RequestStatus::Success - } else { - crate::telemetry::RequestStatus::Failed - }; - record_request_telemetry(&state, &ctx, status, None); - - // 如果成功，记录估算的 Token 使用量 - if is_success { - let estimated_input_tokens = request - .messages - .iter() - .map(|m| { - let content_len = match &m.content { - serde_json::Value::String(s) => s.len(), - serde_json::Value::Array(arr) => arr - .iter() - .filter_map(|v| v.get("text").and_then(|t| t.as_str())) - .map(|s| s.len()) - .sum(), - _ => 0, - }; - content_len / 4 - }) - .sum::() as u32; - // 输出 Token 使用估算值 - let estimated_output_tokens = 100u32; - record_token_usage( - &state, - &ctx, - Some(estimated_input_tokens), - Some(estimated_output_tokens), - ); - } - - return response; - } - - // 回退到旧的单凭证模式 - state.logs.write().await.add( - "debug", - &format!( - "[ROUTE] No pool credential found for '{}', using legacy mode", - default_provider - ), - ); - - // 检查是否需要刷新 token（无 token 或即将过期） - { - let _guard = state.kiro_refresh_lock.lock().await; - let mut kiro = state.kiro.write().await; - let needs_refresh = - kiro.credentials.access_token.is_none() || kiro.is_token_expiring_soon(); - if needs_refresh { - state.logs.write().await.add( - "info", - "[AUTH] No access token or token expiring soon, attempting refresh...", - ); - if let Err(e) = kiro.refresh_token().await { - state - .logs - .write() - .await - .add("error", &format!("[AUTH] Token refresh failed: {e}")); - return ( - StatusCode::UNAUTHORIZED, - Json(serde_json::json!({"error": {"message": format!("Token refresh failed: {e}")}})), - ) - .into_response(); - } - state - .logs - .write() - .await - .add("info", "[AUTH] Token refreshed successfully"); - } - } - - // 转换为 OpenAI 格式 - let openai_request = convert_anthropic_to_openai(&request); - - // 记录转换后的请求信息 - state.logs.write().await.add( - "debug", - &format!( - "[CONVERT] OpenAI format: messages={} tools={} stream={}", - openai_request.messages.len(), - openai_request.tools.as_ref().map(|t| t.len()).unwrap_or(0), - openai_request.stream - ), - ); - - let kiro = state.kiro.read().await; - - match kiro.call_api(&openai_request).await { - Ok(resp) => { - let status = resp.status(); - state - .logs - .write() - .await - .add("info", &format!("[RESP] Upstream status: {status}")); - - if status.is_success() { - match resp.bytes().await { - Ok(bytes) => { - // 使用 lossy 转换，避免无效 UTF-8 导致崩溃 - let body = String::from_utf8_lossy(&bytes).to_string(); - - // 记录原始响应长度 - state.logs.write().await.add( - "debug", - &format!("[RESP] Raw body length: {} bytes", bytes.len()), - ); - - // 保存原始响应到文件用于调试 - let request_id = uuid::Uuid::new_v4().to_string()[..8].to_string(); - state.logs.read().await.log_raw_response(&request_id, &body); - state.logs.write().await.add( - "debug", - &format!("[RESP] Raw response saved to raw_response_{request_id}.txt"), - ); - - // 记录响应的前200字符用于调试（减少日志量） - let preview: String = - body.chars().filter(|c| !c.is_control()).take(200).collect(); - state - .logs - .write() - .await - .add("debug", &format!("[RESP] Body preview: {preview}")); - - let parsed = parse_cw_response(&body); - - // 详细记录解析结果 - state.logs.write().await.add( - "info", - &format!( - "[RESP] Parsed: content_len={}, tool_calls={}, content_preview={}", - parsed.content.len(), - parsed.tool_calls.len(), - parsed.content.chars().take(100).collect::() - ), - ); - - // 记录 tool calls 详情 - for (i, tc) in parsed.tool_calls.iter().enumerate() { - state.logs.write().await.add( - "debug", - &format!( - "[RESP] Tool call {}: name={} id={}", - i, tc.function.name, tc.id - ), - ); - } - - // 如果请求流式响应，返回 SSE 格式 - if request.stream { - return build_anthropic_stream_response(&request.model, &parsed); - } - - // 非流式响应 - build_anthropic_response(&request.model, &parsed) - } - Err(e) => { - state - .logs - .write() - .await - .add("error", &format!("[ERROR] Response body read failed: {e}")); - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({"error": {"message": e.to_string()}})), - ) - .into_response() - } - } - } else if status.as_u16() == 403 || status.as_u16() == 402 { - // Token 过期或账户问题，尝试重新加载凭证并刷新 - drop(kiro); - let _guard = state.kiro_refresh_lock.lock().await; - let mut kiro = state.kiro.write().await; - state.logs.write().await.add( - "warn", - &format!( - "[AUTH] Got {}, reloading credentials and attempting token refresh...", - status.as_u16() - ), - ); - - // 先重新加载凭证文件（可能用户换了账户） - if let Err(e) = kiro.load_credentials().await { - state.logs.write().await.add( - "error", - &format!("[AUTH] Failed to reload credentials: {e}"), - ); - } - - match kiro.refresh_token().await { - Ok(_) => { - state.logs.write().await.add( - "info", - "[AUTH] Token refreshed successfully, retrying request...", - ); - drop(kiro); - let kiro = state.kiro.read().await; - match kiro.call_api(&openai_request).await { - Ok(retry_resp) => { - let retry_status = retry_resp.status(); - state.logs.write().await.add( - "info", - &format!("[RETRY] Response status: {retry_status}"), - ); - if retry_resp.status().is_success() { - match retry_resp.bytes().await { - Ok(bytes) => { - let body = String::from_utf8_lossy(&bytes).to_string(); - let parsed = parse_cw_response(&body); - state.logs.write().await.add( - "info", - &format!( - "[RETRY] Success: content_len={}, tool_calls={}", - parsed.content.len(), parsed.tool_calls.len() - ), - ); - if request.stream { - return build_anthropic_stream_response( - &request.model, - &parsed, - ); - } - return build_anthropic_response( - &request.model, - &parsed, - ); - } - Err(e) => { - state.logs.write().await.add( - "error", - &format!("[RETRY] Body read failed: {e}"), - ); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({"error": {"message": e.to_string()}})), - ) - .into_response(); - } - } - } - let body = retry_resp - .bytes() - .await - .map(|b| String::from_utf8_lossy(&b).to_string()) - .unwrap_or_default(); - state.logs.write().await.add( - "error", - &format!( - "[RETRY] Failed with status {retry_status}: {}", - safe_truncate(&body, 500) - ), - ); - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({"error": {"message": format!("Retry failed: {}", body)}})), - ) - .into_response() - } - Err(e) => { - state - .logs - .write() - .await - .add("error", &format!("[RETRY] Request failed: {e}")); - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({"error": {"message": e.to_string()}})), - ) - .into_response() - } - } - } - Err(e) => { - state - .logs - .write() - .await - .add("error", &format!("[AUTH] Token refresh failed: {e}")); - ( - StatusCode::UNAUTHORIZED, - Json(serde_json::json!({"error": {"message": format!("Token refresh failed: {e}")}})), - ) - .into_response() - } - } - } else { - let body = resp.text().await.unwrap_or_default(); - state.logs.write().await.add( - "error", - &format!( - "[ERROR] Upstream error HTTP {}: {}", - status, - safe_truncate(&body, 500) - ), - ); - ( - StatusCode::from_u16(status.as_u16()) - .unwrap_or(StatusCode::INTERNAL_SERVER_ERROR), - Json( - serde_json::json!({"error": {"message": format!("Upstream error: {}", body)}}), - ), - ) - .into_response() - } - } - Err(e) => { - // 详细记录网络/连接错误 - let error_details = format!("{e:?}"); - state - .logs - .write() - .await - .add("error", &format!("[ERROR] Kiro API call failed: {e}")); - state.logs.write().await.add( - "debug", - &format!("[ERROR] Full error details: {error_details}"), - ); - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({"error": {"message": e.to_string()}})), - ) - .into_response() - } - } -} - -/// 构建 Anthropic 非流式响应 -fn build_anthropic_response(model: &str, parsed: &CWParsedResponse) -> Response { - let has_tool_calls = !parsed.tool_calls.is_empty(); - let mut content_array: Vec = Vec::new(); - - if !parsed.content.is_empty() { - content_array.push(serde_json::json!({ - "type": "text", - "text": parsed.content - })); - } - - for tc in &parsed.tool_calls { - let input: serde_json::Value = - serde_json::from_str(&tc.function.arguments).unwrap_or(serde_json::json!({})); - content_array.push(serde_json::json!({ - "type": "tool_use", - "id": tc.id, - "name": tc.function.name, - "input": input - })); - } - - if content_array.is_empty() { - content_array.push(serde_json::json!({"type": "text", "text": ""})); - } - - // 估算 output tokens: 基于响应内容长度 (约 4 字符 = 1 token) - let mut output_tokens: u32 = (parsed.content.len() / 4) as u32; - for tc in &parsed.tool_calls { - output_tokens += (tc.function.arguments.len() / 4) as u32; - } - // 从 context_usage_percentage 估算 input tokens - // 假设 100% = 200k tokens (Claude 的上下文窗口) - let input_tokens = ((parsed.context_usage_percentage / 100.0) * 200000.0) as u32; - - let response = serde_json::json!({ - "id": format!("msg_{}", uuid::Uuid::new_v4()), - "type": "message", - "role": "assistant", - "content": content_array, - "model": model, - "stop_reason": if has_tool_calls { "tool_use" } else { "end_turn" }, - "stop_sequence": null, - "usage": { - "input_tokens": input_tokens, - "output_tokens": output_tokens - } - }); - Json(response).into_response() -} - -/// 构建 Anthropic 流式响应 (SSE) -fn build_anthropic_stream_response(model: &str, parsed: &CWParsedResponse) -> Response { - let has_tool_calls = !parsed.tool_calls.is_empty(); - let message_id = format!("msg_{}", uuid::Uuid::new_v4()); - let model = model.to_string(); - let content = parsed.content.clone(); - let tool_calls = parsed.tool_calls.clone(); - - // 估算 output tokens: 基于响应内容长度 (约 4 字符 = 1 token) - let mut output_tokens: u32 = (parsed.content.len() / 4) as u32; - for tc in &parsed.tool_calls { - output_tokens += (tc.function.arguments.len() / 4) as u32; - } - // 从 context_usage_percentage 估算 input tokens - let input_tokens = ((parsed.context_usage_percentage / 100.0) * 200000.0) as u32; - - // 构建 SSE 事件流 - let mut events: Vec = Vec::new(); - - // 1. message_start - let message_start = serde_json::json!({ - "type": "message_start", - "message": { - "id": message_id, - "type": "message", - "role": "assistant", - "model": model, - "content": [], - "stop_reason": null, - "stop_sequence": null, - "usage": {"input_tokens": input_tokens, "output_tokens": 0} - } - }); - events.push(format!("event: message_start\ndata: {message_start}\n\n")); - - let mut block_index = 0; - - // 2. 文本内容块 - 即使为空也要发送，Claude Code 需要至少一个 content block - // content_block_start - let block_start = serde_json::json!({ - "type": "content_block_start", - "index": block_index, - "content_block": {"type": "text", "text": ""} - }); - events.push(format!( - "event: content_block_start\ndata: {block_start}\n\n" - )); - - if !content.is_empty() { - // content_block_delta - 发送完整内容 - let block_delta = serde_json::json!({ - "type": "content_block_delta", - "index": block_index, - "delta": {"type": "text_delta", "text": content} - }); - events.push(format!( - "event: content_block_delta\ndata: {block_delta}\n\n" - )); - } - - // content_block_stop - let block_stop = serde_json::json!({ - "type": "content_block_stop", - "index": block_index - }); - events.push(format!("event: content_block_stop\ndata: {block_stop}\n\n")); - - block_index += 1; - - // 3. Tool use 块 - for tc in &tool_calls { - // content_block_start - let block_start = serde_json::json!({ - "type": "content_block_start", - "index": block_index, - "content_block": { - "type": "tool_use", - "id": tc.id, - "name": tc.function.name, - "input": {} - } - }); - events.push(format!( - "event: content_block_start\ndata: {block_start}\n\n" - )); - - // content_block_delta - input_json_delta - // 注意：partial_json 应该是原始 JSON 字符串，不是再次序列化的 - let partial_json = if tc.function.arguments.is_empty() { - "{}".to_string() - } else { - tc.function.arguments.clone() - }; - let block_delta = serde_json::json!({ - "type": "content_block_delta", - "index": block_index, - "delta": { - "type": "input_json_delta", - "partial_json": partial_json - } - }); - events.push(format!( - "event: content_block_delta\ndata: {block_delta}\n\n" - )); - - // content_block_stop - let block_stop = serde_json::json!({ - "type": "content_block_stop", - "index": block_index - }); - events.push(format!("event: content_block_stop\ndata: {block_stop}\n\n")); - - block_index += 1; - } - - // 4. message_delta - let message_delta = serde_json::json!({ - "type": "message_delta", - "delta": { - "stop_reason": if has_tool_calls { "tool_use" } else { "end_turn" }, - "stop_sequence": null - }, - "usage": {"output_tokens": output_tokens} - }); - events.push(format!("event: message_delta\ndata: {message_delta}\n\n")); - - // 5. message_stop - let message_stop = serde_json::json!({"type": "message_stop"}); - events.push(format!("event: message_stop\ndata: {message_stop}\n\n")); - - // 创建 SSE 响应 - let body_stream = stream::iter(events.into_iter().map(Ok::<_, std::convert::Infallible>)); - let body = Body::from_stream(body_stream); - - Response::builder() - .status(StatusCode::OK) - .header(header::CONTENT_TYPE, "text/event-stream") - .header(header::CACHE_CONTROL, "no-cache") - .header(header::CONNECTION, "keep-alive") - .body(body) - .unwrap_or_else(|e| { - tracing::error!("Failed to build SSE response: {}", e); - Response::builder() - .status(StatusCode::INTERNAL_SERVER_ERROR) - .body(Body::empty()) - .unwrap_or_default() - }) -} - -async fn count_tokens( - State(state): State, - headers: HeaderMap, - Json(_request): Json, -) -> Response { - let _reload_guard = state.processor.reload_lock.read().await; - if let Err(e) = verify_api_key(&headers, &state.api_key).await { - return e.into_response(); - } - - // Claude Code 需要这个端点，返回估算值 - Json(serde_json::json!({ - "input_tokens": 100 - })) - .into_response() -} - -/// CodeWhisperer 响应解析结果 -#[derive(Debug, Default)] -struct CWParsedResponse { - content: String, - tool_calls: Vec, - usage_credits: f64, - context_usage_percentage: f64, -} - -/// 解析 CodeWhisperer AWS Event Stream 响应 -/// AWS Event Stream 是二进制格式，JSON payload 嵌入在二进制头部之间 -fn parse_cw_response(body: &str) -> CWParsedResponse { - let mut result = CWParsedResponse::default(); - // 使用 HashMap 来跟踪多个并发的 tool calls - // key: toolUseId, value: (name, input_accumulated) - let mut tool_map: std::collections::HashMap = - std::collections::HashMap::new(); - - // 将字符串转换为字节，因为 AWS Event Stream 包含二进制数据 - let bytes = body.as_bytes(); - - // 搜索所有 JSON 对象的模式 - // AWS Event Stream 格式: [binary headers]{"content":"..."}[binary trailer] - let json_patterns: &[&[u8]] = &[ - b"{\"content\":", - b"{\"name\":", - b"{\"input\":", - b"{\"stop\":", - b"{\"followupPrompt\":", - b"{\"toolUseId\":", - b"{\"unit\":", // meteringEvent - b"{\"contextUsagePercentage\":", // contextUsageEvent - ]; - - let mut pos = 0; - while pos < bytes.len() { - // 找到下一个 JSON 对象的开始 - let mut next_start: Option = None; - - for pattern in json_patterns { - if let Some(idx) = find_subsequence(&bytes[pos..], pattern) { - let abs_pos = pos + idx; - if next_start.is_none_or(|start| abs_pos < start) { - next_start = Some(abs_pos); - } - } - } - - let start = match next_start { - Some(s) => s, - None => break, - }; - - // 从 start 位置提取完整的 JSON 对象 - if let Some(json_str) = extract_json_from_bytes(&bytes[start..]) { - if let Ok(value) = serde_json::from_str::(&json_str) { - // 处理 content 事件 - if let Some(content) = value.get("content").and_then(|v| v.as_str()) { - // 跳过 followupPrompt - if value.get("followupPrompt").is_none() { - result.content.push_str(content); - } - } - // 处理 tool use 事件 (包含 toolUseId) - else if let Some(tool_use_id) = value.get("toolUseId").and_then(|v| v.as_str()) { - let name = value - .get("name") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let input_chunk = value - .get("input") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let is_stop = value.get("stop").and_then(|v| v.as_bool()).unwrap_or(false); - - // 获取或创建 tool entry - let entry = tool_map - .entry(tool_use_id.to_string()) - .or_insert_with(|| (String::new(), String::new())); - - // 更新 name（如果有） - if !name.is_empty() { - entry.0 = name; - } - - // 累积 input - entry.1.push_str(&input_chunk); - - // 如果是 stop 事件，完成这个 tool call - if is_stop { - if let Some((name, input)) = tool_map.remove(tool_use_id) { - if !name.is_empty() { - result.tool_calls.push(ToolCall { - id: tool_use_id.to_string(), - call_type: "function".to_string(), - function: FunctionCall { - name, - arguments: input, - }, - }); - } - } - } - } - // 处理独立的 stop 事件（没有 toolUseId） - else if value.get("stop").and_then(|v| v.as_bool()).unwrap_or(false) { - // 这种情况不应该发生，但以防万一 - } - // 处理 meteringEvent: {"unit":"credit","unitPlural":"credits","usage":0.34} - else if let Some(usage) = value.get("usage").and_then(|v| v.as_f64()) { - result.usage_credits = usage; - } - // 处理 contextUsageEvent: {"contextUsagePercentage":54.36} - else if let Some(ctx_usage) = - value.get("contextUsagePercentage").and_then(|v| v.as_f64()) - { - result.context_usage_percentage = ctx_usage; - } - } - pos = start + json_str.len(); - } else { - pos = start + 1; - } - } - - // 处理未完成的 tool calls（没有收到 stop 事件的） - for (id, (name, input)) in tool_map { - if !name.is_empty() { - result.tool_calls.push(ToolCall { - id, - call_type: "function".to_string(), - function: FunctionCall { - name, - arguments: input, - }, - }); - } - } - - // 解析 bracket 格式的 tool calls: [Called xxx with args: {...}] - parse_bracket_tool_calls(&mut result); - - result -} - -/// 在字节数组中查找子序列 -fn find_subsequence(haystack: &[u8], needle: &[u8]) -> Option { - haystack - .windows(needle.len()) - .position(|window| window == needle) -} - -/// 从字节数组中提取 JSON 对象字符串 -fn extract_json_from_bytes(bytes: &[u8]) -> Option { - if bytes.is_empty() || bytes[0] != b'{' { - return None; - } - - let mut brace_count = 0; - let mut in_string = false; - let mut escape_next = false; - let mut end_pos = None; - - for (i, &b) in bytes.iter().enumerate() { - if escape_next { - escape_next = false; - continue; - } - - match b { - b'\\' if in_string => escape_next = true, - b'"' => in_string = !in_string, - b'{' if !in_string => brace_count += 1, - b'}' if !in_string => { - brace_count -= 1; - if brace_count == 0 { - end_pos = Some(i + 1); - break; - } - } - _ => {} - } - } - - end_pos.and_then(|end| String::from_utf8(bytes[..end].to_vec()).ok()) -} - -/// 从字符串中提取完整的 JSON 对象 (保留用于兼容) -#[allow(dead_code)] -fn extract_json_object(s: &str) -> Option<&str> { - if !s.starts_with('{') { - return None; - } - - let mut brace_count = 0; - let mut in_string = false; - let mut escape_next = false; - - for (i, c) in s.char_indices() { - if escape_next { - escape_next = false; - continue; - } - - match c { - '\\' if in_string => escape_next = true, - '"' => in_string = !in_string, - '{' if !in_string => brace_count += 1, - '}' if !in_string => { - brace_count -= 1; - if brace_count == 0 { - return Some(&s[..=i]); - } - } - _ => {} - } - } - None -} - -/// 解析 bracket 格式的 tool calls -fn parse_bracket_tool_calls(result: &mut CWParsedResponse) { - let re = - regex::Regex::new(r"\[Called\s+(\w+)\s+with\s+args:\s*(\{[^}]*(?:\{[^}]*\}[^}]*)*\})\]") - .ok(); - - if let Some(re) = re { - let mut to_remove = Vec::new(); - for cap in re.captures_iter(&result.content) { - if let (Some(name), Some(args)) = (cap.get(1), cap.get(2)) { - let tool_id = format!( - "call_{}", - &uuid::Uuid::new_v4().to_string().replace('-', "")[..8] - ); - result.tool_calls.push(ToolCall { - id: tool_id, - call_type: "function".to_string(), - function: FunctionCall { - name: name.as_str().to_string(), - arguments: args.as_str().to_string(), - }, - }); - if let Some(full_match) = cap.get(0) { - to_remove.push(full_match.as_str().to_string()); - } - } - } - // 从 content 中移除 tool call 文本 - for s in to_remove { - result.content = result.content.replace(&s, ""); - } - result.content = result.content.trim().to_string(); - } -} - -/// 列出所有可用路由 -/// P1 安全修复：添加 API Key 鉴权，防止信息泄露 -async fn list_routes(State(state): State, headers: HeaderMap) -> Response { - // 验证 API Key - if let Err(e) = verify_api_key(&headers, &state.api_key).await { - return e.into_response(); - } - - let routes = match &state.db { - Some(db) => state - .pool_service - .get_available_routes(db, &state.base_url) - .unwrap_or_default(), - None => Vec::new(), - }; - - // 添加默认路由 - let mut all_routes = vec![RouteInfo { - selector: "default".to_string(), - provider_type: "kiro".to_string(), - credential_count: 1, - endpoints: vec![ - crate::models::route_model::RouteEndpoint { - path: "/v1/messages".to_string(), - protocol: "claude".to_string(), - url: format!("{}/v1/messages", state.base_url), - }, - crate::models::route_model::RouteEndpoint { - path: "/v1/chat/completions".to_string(), - protocol: "openai".to_string(), - url: format!("{}/v1/chat/completions", state.base_url), - }, - ], - tags: vec!["默认".to_string()], - enabled: true, - }]; - all_routes.extend(routes); - - let response = RouteListResponse { - base_url: state.base_url.clone(), - default_provider: "kiro".to_string(), - routes: all_routes, - }; - - Json(response).into_response() -} - -/// 带选择器的 Anthropic messages 处理 -async fn anthropic_messages_with_selector( - State(state): State, - Path(selector): Path, - headers: HeaderMap, - Json(request): Json, -) -> Response { - let _reload_guard = state.processor.reload_lock.read().await; - // 使用 Anthropic 格式的认证验证 - if let Err(e) = verify_api_key_anthropic(&headers, &state.api_key).await { - state.logs.write().await.add( - "warn", - &format!("Unauthorized request to /{}/v1/messages", selector), - ); - return e.into_response(); - } - - state.logs.write().await.add( - "info", - &format!( - "[REQ] POST /{}/v1/messages model={} stream={}", - selector, request.model, request.stream - ), - ); - - // 尝试解析凭证 - let credential = match &state.db { - Some(db) => { - // 首先尝试按名称查找 - if let Ok(Some(cred)) = state.pool_service.get_by_name(db, &selector) { - Some(cred) - } - // 然后尝试按 UUID 查找 - else if let Ok(Some(cred)) = state.pool_service.get_by_uuid(db, &selector) { - Some(cred) - } - // 最后尝试按 provider 类型轮询 - else if let Ok(Some(cred)) = - state - .pool_service - .select_credential(db, &selector, Some(&request.model)) - { - Some(cred) - } else { - None - } - } - None => None, - }; - - match credential { - Some(cred) => { - state.logs.write().await.add( - "info", - &format!( - "[ROUTE] Using credential: type={} name={:?} uuid={}", - cred.provider_type, - cred.name, - &cred.uuid[..8] - ), - ); - - // 根据凭证类型调用相应的 Provider - call_provider_anthropic(&state, &cred, &request).await - } - None => { - // 回退到默认 Kiro provider - state.logs.write().await.add( - "warn", - &format!( - "[ROUTE] Credential not found for selector '{}', falling back to default", - selector - ), - ); - // 调用原有的 Kiro 处理逻辑 - anthropic_messages_internal(&state, &request).await - } - } -} - -/// 带选择器的 OpenAI chat completions 处理 -async fn chat_completions_with_selector( - State(state): State, - Path(selector): Path, - headers: HeaderMap, - Json(request): Json, -) -> Response { - let _reload_guard = state.processor.reload_lock.read().await; - if let Err(e) = verify_api_key(&headers, &state.api_key).await { - state.logs.write().await.add( - "warn", - &format!("Unauthorized request to /{}/v1/chat/completions", selector), - ); - return e.into_response(); - } - - state.logs.write().await.add( - "info", - &format!( - "[REQ] POST /{}/v1/chat/completions model={} stream={}", - selector, request.model, request.stream - ), - ); - - // 尝试解析凭证 - let credential = match &state.db { - Some(db) => { - if let Ok(Some(cred)) = state.pool_service.get_by_name(db, &selector) { - Some(cred) - } else if let Ok(Some(cred)) = state.pool_service.get_by_uuid(db, &selector) { - Some(cred) - } else if let Ok(Some(cred)) = - state - .pool_service - .select_credential(db, &selector, Some(&request.model)) - { - Some(cred) - } else { - None - } - } - None => None, - }; - - match credential { - Some(cred) => { - state.logs.write().await.add( - "info", - &format!( - "[ROUTE] Using credential: type={} name={:?} uuid={}", - cred.provider_type, - cred.name, - &cred.uuid[..8] - ), - ); - - call_provider_openai(&state, &cred, &request).await - } - None => { - state.logs.write().await.add( - "warn", - &format!( - "[ROUTE] Credential not found for selector '{}', falling back to default", - selector - ), - ); - chat_completions_internal(&state, &request).await - } - } -} - -// ============ Amp CLI 路由处理 ============ - -/// Amp CLI chat completions 处理 -/// -/// 处理 `/api/provider/:provider/v1/chat/completions` 路由 -/// 支持模型映射，将不可用模型映射到可用替代 -async fn amp_chat_completions( - State(state): State, - Path(provider): Path, - headers: HeaderMap, - Json(mut request): Json, -) -> Response { - let _reload_guard = state.processor.reload_lock.read().await; - if let Err(e) = verify_api_key(&headers, &state.api_key).await { - state.logs.write().await.add( - "warn", - &format!( - "Unauthorized request to /api/provider/{}/v1/chat/completions", - provider - ), - ); - return e.into_response(); - } - - // 应用模型映射 - let original_model = request.model.clone(); - let mapped_model = state.amp_router.apply_model_mapping(&request.model); - if mapped_model != original_model { - state.logs.write().await.add( - "info", - &format!( - "[AMP] Model mapping applied: {} -> {}", - original_model, mapped_model - ), - ); - request.model = mapped_model; - } - - state.logs.write().await.add( - "info", - &format!( - "[AMP] POST /api/provider/{}/v1/chat/completions model={} stream={}", - provider, request.model, request.stream - ), - ); - - // 尝试根据 provider 名称选择凭证 - let credential = match &state.db { - Some(db) => { - // 首先尝试按 provider 类型选择 - if let Ok(Some(cred)) = - state - .pool_service - .select_credential(db, &provider, Some(&request.model)) - { - Some(cred) - } - // 然后尝试按名称查找 - else if let Ok(Some(cred)) = state.pool_service.get_by_name(db, &provider) { - Some(cred) - } - // 最后尝试按 UUID 查找 - else if let Ok(Some(cred)) = state.pool_service.get_by_uuid(db, &provider) { - Some(cred) - } else { - None - } - } - None => None, - }; - - match credential { - Some(cred) => { - state.logs.write().await.add( - "info", - &format!( - "[AMP] Using credential: type={} name={:?} uuid={}", - cred.provider_type, - cred.name, - &cred.uuid[..8] - ), - ); - call_provider_openai(&state, &cred, &request).await - } - None => { - state.logs.write().await.add( - "warn", - &format!( - "[AMP] Credential not found for provider '{}', falling back to default", - provider - ), - ); - chat_completions_internal(&state, &request).await - } - } -} - -/// Amp CLI messages 处理 -/// -/// 处理 `/api/provider/:provider/v1/messages` 路由 -/// 支持模型映射，将不可用模型映射到可用替代 -async fn amp_messages( - State(state): State, - Path(provider): Path, - headers: HeaderMap, - Json(mut request): Json, -) -> Response { - let _reload_guard = state.processor.reload_lock.read().await; - // 使用 Anthropic 格式的认证验证 - if let Err(e) = verify_api_key_anthropic(&headers, &state.api_key).await { - state.logs.write().await.add( - "warn", - &format!( - "Unauthorized request to /api/provider/{}/v1/messages", - provider - ), - ); - return e.into_response(); - } - - // 应用模型映射 - let original_model = request.model.clone(); - let mapped_model = state.amp_router.apply_model_mapping(&request.model); - if mapped_model != original_model { - state.logs.write().await.add( - "info", - &format!( - "[AMP] Model mapping applied: {} -> {}", - original_model, mapped_model - ), - ); - request.model = mapped_model; - } - - state.logs.write().await.add( - "info", - &format!( - "[AMP] POST /api/provider/{}/v1/messages model={} stream={}", - provider, request.model, request.stream - ), - ); - - // 尝试根据 provider 名称选择凭证 - let credential = match &state.db { - Some(db) => { - // 首先尝试按 provider 类型选择 - if let Ok(Some(cred)) = - state - .pool_service - .select_credential(db, &provider, Some(&request.model)) - { - Some(cred) - } - // 然后尝试按名称查找 - else if let Ok(Some(cred)) = state.pool_service.get_by_name(db, &provider) { - Some(cred) - } - // 最后尝试按 UUID 查找 - else if let Ok(Some(cred)) = state.pool_service.get_by_uuid(db, &provider) { - Some(cred) - } else { - None - } - } - None => None, - }; - - match credential { - Some(cred) => { - state.logs.write().await.add( - "info", - &format!( - "[AMP] Using credential: type={} name={:?} uuid={}", - cred.provider_type, - cred.name, - &cred.uuid[..8] - ), - ); - call_provider_anthropic(&state, &cred, &request).await - } - None => { - state.logs.write().await.add( - "warn", - &format!( - "[AMP] Credential not found for provider '{}', falling back to default", - provider - ), - ); - anthropic_messages_internal(&state, &request).await - } - } -} - -/// Amp CLI 管理代理 - auth 路由 -/// -/// 处理 `/api/auth/*` 路由，将请求代理到上游 URL -async fn amp_management_proxy_auth( - State(state): State, - Path(path): Path, - headers: HeaderMap, - method: axum::http::Method, - body: axum::body::Bytes, -) -> Response { - amp_management_proxy_internal(state, &format!("auth/{}", path), headers, method, body).await -} - -/// Amp CLI 管理代理 - user 路由 -/// -/// 处理 `/api/user/*` 路由，将请求代理到上游 URL -async fn amp_management_proxy_user( - State(state): State, - Path(path): Path, - headers: HeaderMap, - method: axum::http::Method, - body: axum::body::Bytes, -) -> Response { - amp_management_proxy_internal(state, &format!("user/{}", path), headers, method, body).await -} - -/// Amp CLI 管理代理内部实现 -/// -/// 处理 `/api/auth/*` 和 `/api/user/*` 路由 -/// 将请求代理到上游 URL -/// -/// # 参数 -/// - `path`: 请求路径（不含 /api/ 前缀，如 "auth/login" 或 "user/profile"） -async fn amp_management_proxy_internal( - state: AppState, - path: &str, - headers: HeaderMap, - method: axum::http::Method, - body: axum::body::Bytes, -) -> Response { - let full_path = format!("/api/{}", path); - - // 检查是否是管理路由 - if !state.amp_router.is_management_route(&full_path) { - state.logs.write().await.add( - "warn", - &format!("[AMP] Invalid management route: {}", full_path), - ); - return ( - StatusCode::NOT_FOUND, - Json(serde_json::json!({"error": {"message": "Not found"}})), - ) - .into_response(); - } - - // P0 安全修复：不信任任何请求头判断 localhost - // 此函数应该接收 ConnectInfo 来获取真实连接 IP - // 由于当前函数签名限制，暂时禁用基于头的 localhost 判断 - // 检查 localhost 限制 - if state.amp_router.restrict_management_to_localhost() { - // 安全警告：此处应使用 ConnectInfo 获取真实 IP - // 当前实现拒绝所有非本地请求，因为无法可靠验证来源 - state.logs.write().await.add( - "warn", - "[AMP] Management proxy requires ConnectInfo for secure localhost verification", - ); - return ( - StatusCode::FORBIDDEN, - Json(serde_json::json!({"error": {"message": "Management endpoints require secure localhost verification. Please access directly without proxy."}})), - ) - .into_response(); - } - - // 获取上游 URL - let upstream_url = match state.amp_router.get_management_upstream_path(&full_path) { - Some(url) => url, - None => { - state.logs.write().await.add( - "warn", - &format!("[AMP] No upstream URL configured for management proxy"), - ); - return ( - StatusCode::SERVICE_UNAVAILABLE, - Json(serde_json::json!({"error": {"message": "Upstream URL not configured"}})), - ) - .into_response(); - } - }; - - state.logs.write().await.add( - "info", - &format!( - "[AMP] Proxying management request: {} {} -> {}", - method, full_path, upstream_url - ), - ); - - // 创建 HTTP 客户端 - let client = reqwest::Client::new(); - - // 构建请求 - let mut request_builder = match method { - axum::http::Method::GET => client.get(&upstream_url), - axum::http::Method::POST => client.post(&upstream_url), - axum::http::Method::PUT => client.put(&upstream_url), - axum::http::Method::DELETE => client.delete(&upstream_url), - axum::http::Method::PATCH => client.patch(&upstream_url), - axum::http::Method::HEAD => client.head(&upstream_url), - axum::http::Method::OPTIONS => client.request(reqwest::Method::OPTIONS, &upstream_url), - _ => { - return ( - StatusCode::METHOD_NOT_ALLOWED, - Json(serde_json::json!({"error": {"message": "Method not allowed"}})), - ) - .into_response(); - } - }; - - // 复制请求头（排除 host 和 content-length） - for (name, value) in headers.iter() { - let name_str = name.as_str().to_lowercase(); - if name_str != "host" && name_str != "content-length" { - if let Ok(value_str) = value.to_str() { - request_builder = request_builder.header(name.as_str(), value_str); - } - } - } - - // 添加请求体 - if !body.is_empty() { - request_builder = request_builder.body(body.to_vec()); - } - - // 发送请求 - match request_builder.send().await { - Ok(response) => { - let status = response.status(); - let response_headers = response.headers().clone(); - - match response.bytes().await { - Ok(response_body) => { - let mut builder = Response::builder().status(status.as_u16()); - - // 复制响应头 - for (name, value) in response_headers.iter() { - let name_str = name.as_str().to_lowercase(); - // 排除 transfer-encoding 和 content-length（axum 会自动处理） - if name_str != "transfer-encoding" && name_str != "content-length" { - builder = builder.header(name.as_str(), value.to_str().unwrap_or("")); - } - } - - builder - .body(Body::from(response_body.to_vec())) - .unwrap_or_else(|_| { - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({"error": {"message": "Failed to build response"}})), - ) - .into_response() - }) - } - Err(e) => { - state.logs.write().await.add( - "error", - &format!("[AMP] Failed to read upstream response: {}", e), - ); - ( - StatusCode::BAD_GATEWAY, - Json(serde_json::json!({"error": {"message": format!("Failed to read upstream response: {}", e)}})), - ) - .into_response() - } - } - } - Err(e) => { - state.logs.write().await.add( - "error", - &format!("[AMP] Failed to proxy request to upstream: {}", e), - ); - ( - StatusCode::BAD_GATEWAY, - Json(serde_json::json!({"error": {"message": format!("Failed to connect to upstream: {}", e)}})), - ) - .into_response() - } - } -} - -/// 内部 Anthropic messages 处理 (使用默认 Kiro) -async fn anthropic_messages_internal( - state: &AppState, - request: &AnthropicMessagesRequest, -) -> Response { - // 检查 token - { - let _guard = state.kiro_refresh_lock.lock().await; - let mut kiro = state.kiro.write().await; - let needs_refresh = - kiro.credentials.access_token.is_none() || kiro.is_token_expiring_soon(); - if needs_refresh { - if let Err(e) = kiro.refresh_token().await { - state - .logs - .write() - .await - .add("error", &format!("[AUTH] Token refresh failed: {e}")); - return ( - StatusCode::UNAUTHORIZED, - Json(serde_json::json!({"error": {"message": format!("Token refresh failed: {e}")}})), - ) - .into_response(); - } - } - } - - let openai_request = convert_anthropic_to_openai(request); - let kiro = state.kiro.read().await; - - match kiro.call_api(&openai_request).await { - Ok(resp) => { - let status = resp.status(); - if status.is_success() { - match resp.bytes().await { - Ok(bytes) => { - let body = String::from_utf8_lossy(&bytes).to_string(); - let parsed = parse_cw_response(&body); - if request.stream { - build_anthropic_stream_response(&request.model, &parsed) - } else { - build_anthropic_response(&request.model, &parsed) - } - } - Err(e) => ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({"error": {"message": e.to_string()}})), - ) - .into_response(), - } - } else { - let body = resp.text().await.unwrap_or_default(); - ( - StatusCode::from_u16(status.as_u16()).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR), - Json(serde_json::json!({"error": {"message": format!("Upstream error: {}", body)}})), - ) - .into_response() - } - } - Err(e) => ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({"error": {"message": e.to_string()}})), - ) - .into_response(), - } -} - -/// 内部 OpenAI chat completions 处理 (使用默认 Kiro) -async fn chat_completions_internal(state: &AppState, request: &ChatCompletionRequest) -> Response { - { - let _guard = state.kiro_refresh_lock.lock().await; - let mut kiro = state.kiro.write().await; - let needs_refresh = - kiro.credentials.access_token.is_none() || kiro.is_token_expiring_soon(); - if needs_refresh { - if let Err(e) = kiro.refresh_token().await { - return ( - StatusCode::UNAUTHORIZED, - Json(serde_json::json!({"error": {"message": format!("Token refresh failed: {e}")}})), - ) - .into_response(); - } - } - } - - let kiro = state.kiro.read().await; - match kiro.call_api(request).await { - Ok(resp) => { - let status = resp.status(); - if status.is_success() { - match resp.text().await { - Ok(body) => { - let parsed = parse_cw_response(&body); - let has_tool_calls = !parsed.tool_calls.is_empty(); - - let message = if has_tool_calls { - serde_json::json!({ - "role": "assistant", - "content": if parsed.content.is_empty() { serde_json::Value::Null } else { serde_json::json!(parsed.content) }, - "tool_calls": parsed.tool_calls.iter().map(|tc| { - serde_json::json!({ - "id": tc.id, - "type": "function", - "function": { - "name": tc.function.name, - "arguments": tc.function.arguments - } - }) - }).collect::>() - }) - } else { - serde_json::json!({ - "role": "assistant", - "content": parsed.content - }) - }; - - let response = serde_json::json!({ - "id": format!("chatcmpl-{}", uuid::Uuid::new_v4()), - "object": "chat.completion", - "created": std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap_or_default() - .as_secs(), - "model": request.model, - "choices": [{ - "index": 0, - "message": message, - "finish_reason": if has_tool_calls { "tool_calls" } else { "stop" } - }], - "usage": { - "prompt_tokens": 0, - "completion_tokens": 0, - "total_tokens": 0 - } - }); - Json(response).into_response() - } - Err(e) => ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({"error": {"message": e.to_string()}})), - ) - .into_response(), - } - } else { - let body = resp.text().await.unwrap_or_default(); - ( - StatusCode::from_u16(status.as_u16()).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR), - Json(serde_json::json!({"error": {"message": format!("Upstream error: {}", body)}})), - ) - .into_response() - } - } - Err(e) => ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({"error": {"message": e.to_string()}})), - ) - .into_response(), - } -} - -use crate::models::provider_pool_model::{CredentialData, ProviderCredential}; - -/// 根据凭证调用 Provider (Anthropic 格式) -async fn call_provider_anthropic( - state: &AppState, - credential: &ProviderCredential, - request: &AnthropicMessagesRequest, -) -> Response { - match &credential.credential { - CredentialData::KiroOAuth { creds_file_path } => { - // 使用 TokenCacheService 获取有效 token - let db = match &state.db { - Some(db) => db, - None => { - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({"error": {"message": "Database not available"}})), - ) - .into_response(); - } - }; - - // 获取缓存的 token - let token = match state - .token_cache - .get_valid_token(db, &credential.uuid) - .await - { - Ok(t) => t, - Err(e) => { - tracing::warn!("[POOL] Token cache miss, loading from source: {}", e); - // 回退到从源文件加载 - let mut kiro = KiroProvider::new(); - if let Err(e) = kiro.load_credentials_from_path(creds_file_path).await { - // 记录凭证加载失败 - let _ = state.pool_service.mark_unhealthy( - db, - &credential.uuid, - Some(&format!("Failed to load credentials: {}", e)), - ); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({"error": {"message": format!("Failed to load Kiro credentials: {}", e)}})), - ) - .into_response(); - } - if let Err(e) = kiro.refresh_token().await { - // 记录 Token 刷新失败 - let _ = state.pool_service.mark_unhealthy( - db, - &credential.uuid, - Some(&format!("Token refresh failed: {}", e)), - ); - return ( - StatusCode::UNAUTHORIZED, - Json(serde_json::json!({"error": {"message": format!("Token refresh failed: {}", e)}})), - ) - .into_response(); - } - kiro.credentials.access_token.unwrap_or_default() - } - }; - - // 使用获取到的 token 创建 KiroProvider - let mut kiro = KiroProvider::new(); - kiro.credentials.access_token = Some(token); - // 从源文件加载其他配置（region, profile_arn 等） - let _ = kiro.load_credentials_from_path(creds_file_path).await; - - let openai_request = convert_anthropic_to_openai(request); - let resp = match kiro.call_api(&openai_request).await { - Ok(r) => r, - Err(e) => { - // 记录 API 调用失败 - let _ = state.pool_service.mark_unhealthy( - db, - &credential.uuid, - Some(&e.to_string()), - ); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({"error": {"message": e.to_string()}})), - ) - .into_response(); - } - }; - - let status = resp.status(); - if status.is_success() { - match resp.bytes().await { - Ok(bytes) => { - let body = String::from_utf8_lossy(&bytes).to_string(); - let parsed = parse_cw_response(&body); - // 记录成功 - let _ = state.pool_service.mark_healthy( - db, - &credential.uuid, - Some(&request.model), - ); - let _ = state.pool_service.record_usage(db, &credential.uuid); - if request.stream { - build_anthropic_stream_response(&request.model, &parsed) - } else { - build_anthropic_response(&request.model, &parsed) - } - } - Err(e) => { - let _ = state.pool_service.mark_unhealthy( - db, - &credential.uuid, - Some(&e.to_string()), - ); - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({"error": {"message": e.to_string()}})), - ) - .into_response() - } - } - } else if status.as_u16() == 401 || status.as_u16() == 403 { - // Token 过期，强制刷新并重试 - tracing::info!( - "[POOL] Got {}, forcing token refresh for {}", - status, - &credential.uuid[..8] - ); - - let new_token = match state - .token_cache - .refresh_and_cache(db, &credential.uuid, true) - .await - { - Ok(t) => t, - Err(e) => { - // 记录 Token 刷新失败 - let _ = state.pool_service.mark_unhealthy( - db, - &credential.uuid, - Some(&format!("Token refresh failed: {}", e)), - ); - return ( - StatusCode::UNAUTHORIZED, - Json(serde_json::json!({"error": {"message": format!("Token refresh failed: {}", e)}})), - ) - .into_response(); - } - }; - - // 使用新 token 重试 - kiro.credentials.access_token = Some(new_token); - match kiro.call_api(&openai_request).await { - Ok(retry_resp) => { - if retry_resp.status().is_success() { - match retry_resp.bytes().await { - Ok(bytes) => { - let body = String::from_utf8_lossy(&bytes).to_string(); - let parsed = parse_cw_response(&body); - // 记录重试成功 - let _ = state.pool_service.mark_healthy( - db, - &credential.uuid, - Some(&request.model), - ); - let _ = state.pool_service.record_usage(db, &credential.uuid); - if request.stream { - build_anthropic_stream_response(&request.model, &parsed) - } else { - build_anthropic_response(&request.model, &parsed) - } - } - Err(e) => { - let _ = state.pool_service.mark_unhealthy( - db, - &credential.uuid, - Some(&e.to_string()), - ); - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({"error": {"message": e.to_string()}})), - ) - .into_response() - } - } - } else { - let body = retry_resp.text().await.unwrap_or_default(); - let _ = state.pool_service.mark_unhealthy( - db, - &credential.uuid, - Some(&format!("Retry failed: {}", body)), - ); - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({"error": {"message": format!("Retry failed: {}", body)}})), - ) - .into_response() - } - } - Err(e) => { - let _ = state.pool_service.mark_unhealthy( - db, - &credential.uuid, - Some(&e.to_string()), - ); - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({"error": {"message": e.to_string()}})), - ) - .into_response() - } - } - } else { - let body = resp.text().await.unwrap_or_default(); - let _ = state - .pool_service - .mark_unhealthy(db, &credential.uuid, Some(&body)); - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({"error": {"message": body}})), - ) - .into_response() - } - } - CredentialData::GeminiOAuth { .. } => { - // Gemini OAuth 路由暂不支持 - ( - StatusCode::NOT_IMPLEMENTED, - Json(serde_json::json!({"error": {"message": "Gemini OAuth routing not yet implemented. Use /v1/messages with Gemini models instead."}})), - ) - .into_response() - } - CredentialData::QwenOAuth { .. } => { - // Qwen OAuth 路由暂不支持 - ( - StatusCode::NOT_IMPLEMENTED, - Json(serde_json::json!({"error": {"message": "Qwen OAuth routing not yet implemented. Use /v1/messages with Qwen models instead."}})), - ) - .into_response() - } - CredentialData::AntigravityOAuth { - creds_file_path, - project_id, - } => { - let mut antigravity = AntigravityProvider::new(); - if let Err(e) = antigravity - .load_credentials_from_path(creds_file_path) - .await - { - // 记录凭证加载失败 - if let Some(db) = &state.db { - let _ = state.pool_service.mark_unhealthy( - db, - &credential.uuid, - Some(&format!("Failed to load credentials: {}", e)), - ); - } - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({"error": {"message": format!("Failed to load Antigravity credentials: {}", e)}})), - ) - .into_response(); - } - - // 检查并刷新 token - if antigravity.is_token_expiring_soon() { - if let Err(e) = antigravity.refresh_token().await { - // 记录 Token 刷新失败 - if let Some(db) = &state.db { - let _ = state.pool_service.mark_unhealthy( - db, - &credential.uuid, - Some(&format!("Token refresh failed: {}", e)), - ); - } - return ( - StatusCode::UNAUTHORIZED, - Json(serde_json::json!({"error": {"message": format!("Token refresh failed: {}", e)}})), - ) - .into_response(); - } - } - - // 设置项目 ID - if let Some(pid) = project_id { - antigravity.project_id = Some(pid.clone()); - } else if let Err(e) = antigravity.discover_project().await { - tracing::warn!("[Antigravity] Failed to discover project: {}", e); - } - - // 先转换为 OpenAI 格式，再转换为 Antigravity 格式 - let openai_request = convert_anthropic_to_openai(request); - let antigravity_request = convert_openai_to_antigravity(&openai_request); - - match antigravity - .generate_content(&request.model, &antigravity_request) - .await - { - Ok(resp) => { - // 转换为 OpenAI 格式，再构建 Anthropic 响应 - let content = resp["candidates"][0]["content"]["parts"][0]["text"] - .as_str() - .unwrap_or(""); - let parsed = CWParsedResponse { - content: content.to_string(), - tool_calls: Vec::new(), - usage_credits: 0.0, - context_usage_percentage: 0.0, - }; - // 记录成功 - if let Some(db) = &state.db { - let _ = state.pool_service.mark_healthy( - db, - &credential.uuid, - Some(&request.model), - ); - let _ = state.pool_service.record_usage(db, &credential.uuid); - } - if request.stream { - build_anthropic_stream_response(&request.model, &parsed) - } else { - build_anthropic_response(&request.model, &parsed) - } - } - Err(e) => { - // 记录 API 调用失败 - if let Some(db) = &state.db { - let _ = state.pool_service.mark_unhealthy( - db, - &credential.uuid, - Some(&e.to_string()), - ); - } - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({"error": {"message": e.to_string()}})), - ) - .into_response() - } - } - } - CredentialData::OpenAIKey { api_key, base_url } => { - let openai = OpenAICustomProvider::with_config(api_key.clone(), base_url.clone()); - let openai_request = convert_anthropic_to_openai(request); - match openai.call_api(&openai_request).await { - Ok(resp) => { - if resp.status().is_success() { - match resp.text().await { - Ok(body) => { - if let Ok(openai_resp) = - serde_json::from_str::(&body) - { - let content = openai_resp["choices"][0]["message"]["content"] - .as_str() - .unwrap_or(""); - let parsed = CWParsedResponse { - content: content.to_string(), - tool_calls: Vec::new(), - usage_credits: 0.0, - context_usage_percentage: 0.0, - }; - // 记录成功 - if let Some(db) = &state.db { - let _ = state.pool_service.mark_healthy( - db, - &credential.uuid, - Some(&request.model), - ); - let _ = - state.pool_service.record_usage(db, &credential.uuid); - } - if request.stream { - build_anthropic_stream_response(&request.model, &parsed) - } else { - build_anthropic_response(&request.model, &parsed) - } - } else { - // 记录解析失败 - if let Some(db) = &state.db { - let _ = state.pool_service.mark_unhealthy( - db, - &credential.uuid, - Some("Failed to parse OpenAI response"), - ); - } - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({"error": {"message": "Failed to parse OpenAI response"}})), - ) - .into_response() - } - } - Err(e) => { - if let Some(db) = &state.db { - let _ = state.pool_service.mark_unhealthy( - db, - &credential.uuid, - Some(&e.to_string()), - ); - } - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({"error": {"message": e.to_string()}})), - ) - .into_response() - } - } - } else { - let body = resp.text().await.unwrap_or_default(); - if let Some(db) = &state.db { - let _ = state.pool_service.mark_unhealthy( - db, - &credential.uuid, - Some(&body), - ); - } - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({"error": {"message": body}})), - ) - .into_response() - } - } - Err(e) => { - if let Some(db) = &state.db { - let _ = state.pool_service.mark_unhealthy( - db, - &credential.uuid, - Some(&e.to_string()), - ); - } - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({"error": {"message": e.to_string()}})), - ) - .into_response() - } - } - } - CredentialData::ClaudeKey { api_key, base_url } => { - // 打印 Claude 代理 URL 用于调试 - let actual_base_url = base_url.as_deref().unwrap_or("https://api.anthropic.com"); - let claude = ClaudeCustomProvider::with_config(api_key.clone(), base_url.clone()); - let request_url = claude.get_base_url(); - state.logs.write().await.add( - "info", - &format!( - "[CLAUDE] 使用 Claude API 代理: base_url={} -> {}/v1/messages credential_uuid={}", - actual_base_url, - request_url, - &credential.uuid[..8] - ), - ); - // 打印请求参数 - let request_json = serde_json::to_string(request).unwrap_or_default(); - state.logs.write().await.add( - "debug", - &format!( - "[CLAUDE] 请求参数: {}", - &request_json.chars().take(500).collect::() - ), - ); - match claude.call_api(request).await { - Ok(resp) => { - let status = resp.status(); - // 打印响应状态 - state.logs.write().await.add( - "info", - &format!( - "[CLAUDE] 响应状态: status={} model={}", - status, - request.model - ), - ); - match resp.text().await { - Ok(body) => { - if status.is_success() { - // 打印响应内容预览 - state.logs.write().await.add( - "debug", - &format!( - "[CLAUDE] 响应内容: {}", - &body.chars().take(500).collect::() - ), - ); - // 记录成功 - if let Some(db) = &state.db { - let _ = state.pool_service.mark_healthy( - db, - &credential.uuid, - Some(&request.model), - ); - let _ = state.pool_service.record_usage(db, &credential.uuid); - } - Response::builder() - .status(StatusCode::OK) - .header(header::CONTENT_TYPE, "application/json") - .body(Body::from(body)) - .unwrap_or_else(|_| { - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({"error": {"message": "Failed to build response"}})), - ) - .into_response() - }) - } else { - state.logs.write().await.add( - "error", - &format!( - "[CLAUDE] 请求失败: status={} body={}", - status, - &body.chars().take(200).collect::() - ), - ); - if let Some(db) = &state.db { - let _ = state.pool_service.mark_unhealthy( - db, - &credential.uuid, - Some(&body), - ); - } - ( - StatusCode::from_u16(status.as_u16()) - .unwrap_or(StatusCode::INTERNAL_SERVER_ERROR), - Json(serde_json::json!({"error": {"message": body}})), - ) - .into_response() - } - } - Err(e) => { - state.logs.write().await.add( - "error", - &format!("[CLAUDE] 读取响应失败: {}", e), - ); - if let Some(db) = &state.db { - let _ = state.pool_service.mark_unhealthy( - db, - &credential.uuid, - Some(&e.to_string()), - ); - } - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({"error": {"message": e.to_string()}})), - ) - .into_response() - } - } - } - Err(e) => { - if let Some(db) = &state.db { - let _ = state.pool_service.mark_unhealthy( - db, - &credential.uuid, - Some(&e.to_string()), - ); - } - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({"error": {"message": e.to_string()}})), - ) - .into_response() - } - } - } - CredentialData::VertexKey { api_key, base_url, .. } => { - // Vertex AI uses Gemini-compatible API, convert Anthropic to OpenAI format first - let openai_request = convert_anthropic_to_openai(request); - let vertex = VertexProvider::with_config(api_key.clone(), base_url.clone()); - match vertex.chat_completions(&serde_json::to_value(&openai_request).unwrap_or_default()).await { - Ok(resp) => { - let status = resp.status(); - match resp.text().await { - Ok(body) => { - if status.is_success() { - if let Some(db) = &state.db { - let _ = state.pool_service.mark_healthy(db, &credential.uuid, Some(&request.model)); - let _ = state.pool_service.record_usage(db, &credential.uuid); - } - Response::builder() - .status(StatusCode::OK) - .header(header::CONTENT_TYPE, "application/json") - .body(Body::from(body)) - .unwrap_or_else(|_| { - (StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({"error": {"message": "Failed to build response"}}))).into_response() - }) - } else { - if let Some(db) = &state.db { - let _ = state.pool_service.mark_unhealthy(db, &credential.uuid, Some(&body)); - } - (StatusCode::from_u16(status.as_u16()).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR), Json(serde_json::json!({"error": {"message": body}}))).into_response() - } - } - Err(e) => { - if let Some(db) = &state.db { - let _ = state.pool_service.mark_unhealthy(db, &credential.uuid, Some(&e.to_string())); - } - (StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({"error": {"message": e.to_string()}}))).into_response() - } - } - } - Err(e) => { - if let Some(db) = &state.db { - let _ = state.pool_service.mark_unhealthy(db, &credential.uuid, Some(&e.to_string())); - } - (StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({"error": {"message": e.to_string()}}))).into_response() - } - } - } - // Gemini API Key credentials - not supported for Anthropic format - CredentialData::GeminiApiKey { .. } => { - ( - StatusCode::BAD_REQUEST, - Json(serde_json::json!({"error": {"message": "Gemini API Key credentials do not support Anthropic format"}})), - ) - .into_response() - } - // 新增的凭证类型暂不支持 Anthropic 格式 - CredentialData::CodexOAuth { .. } - | CredentialData::ClaudeOAuth { .. } - | CredentialData::IFlowOAuth { .. } - | CredentialData::IFlowCookie { .. } => { - ( - StatusCode::BAD_REQUEST, - Json(serde_json::json!({"error": {"message": "This credential type does not support Anthropic format yet"}})), - ) - .into_response() - } - } -} - -/// 根据凭证调用 Provider (OpenAI 格式) -async fn call_provider_openai( - state: &AppState, - credential: &ProviderCredential, - request: &ChatCompletionRequest, -) -> Response { - let start_time = std::time::Instant::now(); - - match &credential.credential { - CredentialData::KiroOAuth { creds_file_path } => { - let mut kiro = KiroProvider::new(); - if let Err(e) = kiro.load_credentials_from_path(creds_file_path).await { - // 记录凭证加载失败 - if let Some(db) = &state.db { - let _ = state.pool_service.mark_unhealthy(db, &credential.uuid, Some(&e.to_string())); - } - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({"error": {"message": format!("Failed to load Kiro credentials: {}", e)}})), - ) - .into_response(); - } - if let Err(e) = kiro.refresh_token().await { - // 记录 Token 刷新失败 - if let Some(db) = &state.db { - let _ = state.pool_service.mark_unhealthy(db, &credential.uuid, Some(&format!("Token refresh failed: {}", e))); - } - return ( - StatusCode::UNAUTHORIZED, - Json(serde_json::json!({"error": {"message": format!("Token refresh failed: {}", e)}})), - ) - .into_response(); - } - - match kiro.call_api(request).await { - Ok(resp) => { - let status = resp.status(); - if status.is_success() { - // 记录成功 - if let Some(db) = &state.db { - let _ = state.pool_service.mark_healthy(db, &credential.uuid, Some(&request.model)); - let _ = state.pool_service.record_usage(db, &credential.uuid); - } - match resp.text().await { - Ok(body) => { - let parsed = parse_cw_response(&body); - let has_tool_calls = !parsed.tool_calls.is_empty(); - - let message = if has_tool_calls { - serde_json::json!({ - "role": "assistant", - "content": if parsed.content.is_empty() { serde_json::Value::Null } else { serde_json::json!(parsed.content) }, - "tool_calls": parsed.tool_calls.iter().map(|tc| { - serde_json::json!({ - "id": tc.id, - "type": "function", - "function": { - "name": tc.function.name, - "arguments": tc.function.arguments - } - }) - }).collect::>() - }) - } else { - serde_json::json!({ - "role": "assistant", - "content": parsed.content - }) - }; - - Json(serde_json::json!({ - "id": format!("chatcmpl-{}", uuid::Uuid::new_v4()), - "object": "chat.completion", - "created": std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap_or_default() - .as_secs(), - "model": request.model, - "choices": [{ - "index": 0, - "message": message, - "finish_reason": if has_tool_calls { "tool_calls" } else { "stop" } - }], - "usage": { - "prompt_tokens": 0, - "completion_tokens": 0, - "total_tokens": 0 - } - })) - .into_response() - } - Err(e) => ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({"error": {"message": e.to_string()}})), - ) - .into_response(), - } - } else { - // 记录 API 调用失败 - let body = resp.text().await.unwrap_or_default(); - if let Some(db) = &state.db { - let _ = state.pool_service.mark_unhealthy(db, &credential.uuid, Some(&format!("HTTP {}: {}", status, safe_truncate(&body, 100)))); - } - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({"error": {"message": body}})), - ) - .into_response() - } - } - Err(e) => { - // 记录请求错误 - if let Some(db) = &state.db { - let _ = state.pool_service.mark_unhealthy(db, &credential.uuid, Some(&e.to_string())); - } - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({"error": {"message": e.to_string()}})), - ) - .into_response() - } - } - } - CredentialData::GeminiOAuth { .. } => { - ( - StatusCode::NOT_IMPLEMENTED, - Json(serde_json::json!({"error": {"message": "Gemini OAuth routing not yet implemented."}})), - ) - .into_response() - } - CredentialData::QwenOAuth { .. } => { - ( - StatusCode::NOT_IMPLEMENTED, - Json(serde_json::json!({"error": {"message": "Qwen OAuth routing not yet implemented."}})), - ) - .into_response() - } - CredentialData::AntigravityOAuth { creds_file_path, project_id } => { - let mut antigravity = AntigravityProvider::new(); - if let Err(e) = antigravity.load_credentials_from_path(creds_file_path).await { - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({"error": {"message": format!("Failed to load Antigravity credentials: {}", e)}})), - ) - .into_response(); - } - - // 检查并刷新 token - if antigravity.is_token_expiring_soon() { - if let Err(e) = antigravity.refresh_token().await { - return ( - StatusCode::UNAUTHORIZED, - Json(serde_json::json!({"error": {"message": format!("Token refresh failed: {}", e)}})), - ) - .into_response(); - } - } - - // 设置项目 ID - if let Some(pid) = project_id { - antigravity.project_id = Some(pid.clone()); - } else if let Err(e) = antigravity.discover_project().await { - tracing::warn!("[Antigravity] Failed to discover project: {}", e); - } - - // 转换请求格式 - let antigravity_request = convert_openai_to_antigravity(request); - - match antigravity.generate_content(&request.model, &antigravity_request).await { - Ok(resp) => { - let openai_response = convert_antigravity_to_openai_response(&resp, &request.model); - Json(openai_response).into_response() - } - Err(e) => ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({"error": {"message": e.to_string()}})), - ) - .into_response(), - } - } - CredentialData::OpenAIKey { api_key, base_url } => { - let openai = OpenAICustomProvider::with_config(api_key.clone(), base_url.clone()); - match openai.call_api(request).await { - Ok(resp) => { - if resp.status().is_success() { - match resp.text().await { - Ok(body) => { - if let Ok(json) = serde_json::from_str::(&body) { - Json(json).into_response() - } else { - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({"error": {"message": "Invalid JSON response"}})), - ) - .into_response() - } - } - Err(e) => ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({"error": {"message": e.to_string()}})), - ) - .into_response(), - } - } else { - let body = resp.text().await.unwrap_or_default(); - ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({"error": {"message": body}})), - ) - .into_response() - } - } - Err(e) => ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({"error": {"message": e.to_string()}})), - ) - .into_response(), - } - } - CredentialData::ClaudeKey { api_key, base_url } => { - // 打印 Claude 代理 URL 用于调试 - let actual_base_url = base_url.as_deref().unwrap_or("https://api.anthropic.com"); - tracing::info!( - "[CLAUDE] 使用 Claude API 代理: base_url={} credential_uuid={}", - actual_base_url, - &credential.uuid[..8] - ); - let claude = ClaudeCustomProvider::with_config(api_key.clone(), base_url.clone()); - match claude.call_openai_api(request).await { - Ok(resp) => Json(resp).into_response(), - Err(e) => ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(serde_json::json!({"error": {"message": e.to_string()}})), - ) - .into_response(), - } - } - CredentialData::VertexKey { api_key, base_url, model_aliases } => { - // Resolve model alias if present - let resolved_model = model_aliases.get(&request.model).cloned().unwrap_or_else(|| request.model.clone()); - let mut modified_request = request.clone(); - modified_request.model = resolved_model; - - let vertex = VertexProvider::with_config(api_key.clone(), base_url.clone()); - match vertex.chat_completions(&serde_json::to_value(&modified_request).unwrap_or_default()).await { - Ok(resp) => { - if resp.status().is_success() { - match resp.text().await { - Ok(body) => { - if let Ok(json) = serde_json::from_str::(&body) { - Json(json).into_response() - } else { - (StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({"error": {"message": "Invalid JSON response"}}))).into_response() - } - } - Err(e) => (StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({"error": {"message": e.to_string()}}))).into_response(), - } - } else { - let body = resp.text().await.unwrap_or_default(); - (StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({"error": {"message": body}}))).into_response() - } - } - Err(e) => (StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({"error": {"message": e.to_string()}}))).into_response(), - } - } - // Gemini API Key credentials - not supported for OpenAI format yet - CredentialData::GeminiApiKey { .. } => { - ( - StatusCode::BAD_REQUEST, - Json(serde_json::json!({"error": {"message": "Gemini API Key credentials do not support OpenAI format yet"}})), - ) - .into_response() - } - // 新增的凭证类型暂不支持 OpenAI 格式 - CredentialData::CodexOAuth { .. } - | CredentialData::ClaudeOAuth { .. } - | CredentialData::IFlowOAuth { .. } - | CredentialData::IFlowCookie { .. } => { - ( - StatusCode::BAD_REQUEST, - Json(serde_json::json!({"error": {"message": "This credential type does not support OpenAI format yet"}})), - ) - .into_response() - } - } -} - -// ========== WebSocket 处理 ========== - -use crate::websocket::{ - WsApiRequest, WsApiResponse, WsEndpoint, WsError, WsMessage as WsProtoMessage, -}; -use axum::extract::ws::{Message as WsMessage, WebSocket, WebSocketUpgrade}; -use futures::{SinkExt, StreamExt as FuturesStreamExt}; - -/// WebSocket 升级处理器 -async fn ws_upgrade_handler( - ws: WebSocketUpgrade, - State(state): State, - headers: HeaderMap, -) -> impl IntoResponse { - // 验证 API 密钥 - let auth = headers - .get("authorization") - .or_else(|| headers.get("x-api-key")) - .and_then(|v| v.to_str().ok()); - - let key = match auth { - Some(s) if s.starts_with("Bearer ") => &s[7..], - Some(s) => s, - None => { - return axum::http::Response::builder() - .status(401) - .body(Body::from("No API key provided")) - .unwrap() - .into_response(); - } - }; - - if !api_key_matches(key, &state.api_key) { - return axum::http::Response::builder() - .status(401) - .body(Body::from("Invalid API key")) - .unwrap() - .into_response(); - } - - // 获取客户端信息 - let client_info = headers - .get("user-agent") - .and_then(|v| v.to_str().ok()) - .map(|s| s.to_string()); - - ws.on_upgrade(move |socket| handle_websocket(socket, state, client_info)) -} - -/// 处理 WebSocket 连接 -async fn handle_websocket(socket: WebSocket, state: AppState, client_info: Option) { - let conn_id = uuid::Uuid::new_v4().to_string(); - - // 注册连接 - if let Err(e) = state - .ws_manager - .register(conn_id.clone(), client_info.clone()) - { - state.logs.write().await.add( - "error", - &format!("[WS] Failed to register connection: {}", e.message), - ); - return; - } - - state.logs.write().await.add( - "info", - &format!( - "[WS] New connection: {} (client: {:?})", - &conn_id[..8], - client_info - ), - ); - - let (mut sender, mut receiver) = socket.split(); - - // 消息处理循环 - while let Some(msg) = receiver.next().await { - match msg { - Ok(WsMessage::Text(text)) => { - state.ws_manager.on_message(); - state.ws_manager.increment_request_count(&conn_id); - - match serde_json::from_str::(&text) { - Ok(ws_msg) => { - let response = handle_ws_message(&state, &conn_id, ws_msg).await; - if let Some(resp) = response { - let resp_text = serde_json::to_string(&resp).unwrap_or_default(); - if sender - .send(WsMessage::Text(resp_text.into())) - .await - .is_err() - { - break; - } - } - } - Err(e) => { - state.ws_manager.on_error(); - let error = WsProtoMessage::Error(WsError::invalid_message(format!( - "Failed to parse message: {}", - e - ))); - let error_text = serde_json::to_string(&error).unwrap_or_default(); - if sender - .send(WsMessage::Text(error_text.into())) - .await - .is_err() - { - break; - } - } - } - } - Ok(WsMessage::Binary(_)) => { - state.ws_manager.on_error(); - let error = WsProtoMessage::Error(WsError::invalid_message( - "Binary messages not supported", - )); - let error_text = serde_json::to_string(&error).unwrap_or_default(); - if sender - .send(WsMessage::Text(error_text.into())) - .await - .is_err() - { - break; - } - } - Ok(WsMessage::Ping(data)) => { - if sender.send(WsMessage::Pong(data)).await.is_err() { - break; - } - } - Ok(WsMessage::Pong(_)) => { - // 收到 pong，连接正常 - } - Ok(WsMessage::Close(_)) => { - break; - } - Err(e) => { - state.logs.write().await.add( - "error", - &format!("[WS] Connection {} error: {}", &conn_id[..8], e), - ); - break; - } - } - } - - // 清理连接 - state.ws_manager.unregister(&conn_id); - state.logs.write().await.add( - "info", - &format!("[WS] Connection closed: {}", &conn_id[..8]), - ); -} - -/// 处理 WebSocket 消息 -async fn handle_ws_message( - state: &AppState, - conn_id: &str, - msg: WsProtoMessage, -) -> Option { - match msg { - WsProtoMessage::Ping { timestamp } => Some(WsProtoMessage::Pong { timestamp }), - WsProtoMessage::Pong { .. } => None, - WsProtoMessage::Request(request) => { - state.logs.write().await.add( - "info", - &format!( - "[WS] Request from {}: id={} endpoint={:?}", - &conn_id[..8], - request.request_id, - request.endpoint - ), - ); - - // 处理 API 请求 - let response = handle_ws_api_request(state, &request).await; - Some(response) - } - WsProtoMessage::Response(_) - | WsProtoMessage::StreamChunk(_) - | WsProtoMessage::StreamEnd(_) => Some(WsProtoMessage::Error(WsError::invalid_request( - None, - "Invalid message type from client", - ))), - WsProtoMessage::Error(_) => None, - } -} - -/// 处理 WebSocket API 请求 -async fn handle_ws_api_request(state: &AppState, request: &WsApiRequest) -> WsProtoMessage { - match request.endpoint { - WsEndpoint::Models => { - // 返回模型列表 - let models = serde_json::json!({ - "object": "list", - "data": [ - {"id": "claude-sonnet-4-5", "object": "model", "owned_by": "anthropic"}, - {"id": "claude-sonnet-4-5-20250929", "object": "model", "owned_by": "anthropic"}, - {"id": "claude-3-7-sonnet-20250219", "object": "model", "owned_by": "anthropic"}, - {"id": "gemini-2.5-flash", "object": "model", "owned_by": "google"}, - {"id": "gemini-2.5-pro", "object": "model", "owned_by": "google"}, - {"id": "qwen3-coder-plus", "object": "model", "owned_by": "alibaba"}, - ] - }); - WsProtoMessage::Response(WsApiResponse { - request_id: request.request_id.clone(), - payload: models, - }) - } - WsEndpoint::ChatCompletions => { - // 解析 ChatCompletionRequest - match serde_json::from_value::(request.payload.clone()) { - Ok(chat_request) => { - handle_ws_chat_completions(state, &request.request_id, chat_request).await - } - Err(e) => WsProtoMessage::Error(WsError::invalid_request( - Some(request.request_id.clone()), - format!("Invalid chat completion request: {}", e), - )), - } - } - WsEndpoint::Messages => { - // 解析 AnthropicMessagesRequest - match serde_json::from_value::(request.payload.clone()) { - Ok(messages_request) => { - handle_ws_anthropic_messages(state, &request.request_id, messages_request).await - } - Err(e) => WsProtoMessage::Error(WsError::invalid_request( - Some(request.request_id.clone()), - format!("Invalid messages request: {}", e), - )), - } - } - } -} - -/// 处理 WebSocket chat completions 请求 -async fn handle_ws_chat_completions( - state: &AppState, - request_id: &str, - mut request: ChatCompletionRequest, -) -> WsProtoMessage { - let _reload_guard = state.processor.reload_lock.read().await; - // 创建请求上下文 - let mut ctx = RequestContext::new(request.model.clone()).with_stream(request.stream); - - // 使用 RequestProcessor 解析模型别名和路由 - let _provider = state.processor.resolve_and_route(&mut ctx).await; - - // 更新请求中的模型名为解析后的模型 - if ctx.resolved_model != ctx.original_model { - request.model = ctx.resolved_model.clone(); - } - - // 应用参数注入 - let injection_enabled = *state.injection_enabled.read().await; - if injection_enabled { - let injector = state.processor.injector.read().await; - let mut payload = serde_json::to_value(&request).unwrap_or_default(); - let result = injector.inject(&request.model, &mut payload); - if result.has_injections() { - if let Ok(updated) = serde_json::from_value(payload) { - request = updated; - } - } - } - - // 获取默认 provider - let default_provider = state.default_provider.read().await.clone(); - - // 尝试从凭证池中选择凭证 - let credential = match &state.db { - Some(db) => state - .pool_service - .select_credential(db, &default_provider, Some(&request.model)) - .ok() - .flatten(), - None => None, - }; - - // 如果找到凭证，使用它调用 API - if let Some(cred) = credential { - // 简化实现：直接调用 provider 并返回结果 - // 实际实现应该复用 call_provider_openai 的逻辑 - match call_provider_openai_for_ws(state, &cred, &request).await { - Ok(response) => WsProtoMessage::Response(WsApiResponse { - request_id: request_id.to_string(), - payload: response, - }), - Err(e) => WsProtoMessage::Error(WsError::upstream(Some(request_id.to_string()), e)), - } - } else { - // 回退到 Kiro provider - let kiro = state.kiro.read().await; - match kiro.call_api(&request).await { - Ok(resp) => { - if resp.status().is_success() { - match resp.text().await { - Ok(body) => { - let parsed = parse_cw_response(&body); - let has_tool_calls = !parsed.tool_calls.is_empty(); - - let message = if has_tool_calls { - serde_json::json!({ - "role": "assistant", - "content": if parsed.content.is_empty() { serde_json::Value::Null } else { serde_json::json!(parsed.content) }, - "tool_calls": parsed.tool_calls.iter().map(|tc| { - serde_json::json!({ - "id": tc.id, - "type": "function", - "function": { - "name": tc.function.name, - "arguments": tc.function.arguments - } - }) - }).collect::>() - }) - } else { - serde_json::json!({ - "role": "assistant", - "content": parsed.content - }) - }; - - let response = serde_json::json!({ - "id": format!("chatcmpl-{}", uuid::Uuid::new_v4()), - "object": "chat.completion", - "created": std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap_or_default() - .as_secs(), - "model": request.model, - "choices": [{ - "index": 0, - "message": message, - "finish_reason": if has_tool_calls { "tool_calls" } else { "stop" } - }], - "usage": { - "prompt_tokens": 0, - "completion_tokens": 0, - "total_tokens": 0 - } - }); - - WsProtoMessage::Response(WsApiResponse { - request_id: request_id.to_string(), - payload: response, - }) - } - Err(e) => WsProtoMessage::Error(WsError::internal( - Some(request_id.to_string()), - e.to_string(), - )), - } - } else { - let body = resp.text().await.unwrap_or_default(); - WsProtoMessage::Error(WsError::upstream( - Some(request_id.to_string()), - format!("Upstream error: {}", body), - )) - } - } - Err(e) => WsProtoMessage::Error(WsError::internal( - Some(request_id.to_string()), - e.to_string(), - )), - } - } -} - -/// 处理 WebSocket anthropic messages 请求 -async fn handle_ws_anthropic_messages( - state: &AppState, - request_id: &str, - mut request: AnthropicMessagesRequest, -) -> WsProtoMessage { - let _reload_guard = state.processor.reload_lock.read().await; - // 创建请求上下文 - let mut ctx = RequestContext::new(request.model.clone()).with_stream(request.stream); - - // 使用 RequestProcessor 解析模型别名和路由 - let _provider = state.processor.resolve_and_route(&mut ctx).await; - - // 更新请求中的模型名为解析后的模型 - if ctx.resolved_model != ctx.original_model { - request.model = ctx.resolved_model.clone(); - } - - // 应用参数注入 - let injection_enabled = *state.injection_enabled.read().await; - if injection_enabled { - let injector = state.processor.injector.read().await; - let mut payload = serde_json::to_value(&request).unwrap_or_default(); - let result = injector.inject(&request.model, &mut payload); - if result.has_injections() { - if let Ok(updated) = serde_json::from_value(payload) { - request = updated; - } - } - } - - // 获取默认 provider - let default_provider = state.default_provider.read().await.clone(); - - // 尝试从凭证池中选择凭证 - let credential = match &state.db { - Some(db) => state - .pool_service - .select_credential(db, &default_provider, Some(&request.model)) - .ok() - .flatten(), - None => None, - }; - - // 如果找到凭证，使用它调用 API - if let Some(cred) = credential { - match call_provider_anthropic_for_ws(state, &cred, &request).await { - Ok(response) => WsProtoMessage::Response(WsApiResponse { - request_id: request_id.to_string(), - payload: response, - }), - Err(e) => WsProtoMessage::Error(WsError::upstream(Some(request_id.to_string()), e)), - } - } else { - // 回退到 Kiro provider - let kiro = state.kiro.read().await; - - // 转换为 OpenAI 格式 - let openai_request = convert_anthropic_to_openai(&request); - - match kiro.call_api(&openai_request).await { - Ok(resp) => { - if resp.status().is_success() { - match resp.text().await { - Ok(body) => { - let parsed = parse_cw_response(&body); - - // 转换为 Anthropic 格式响应 - let response = serde_json::json!({ - "id": format!("msg_{}", uuid::Uuid::new_v4()), - "type": "message", - "role": "assistant", - "content": [{ - "type": "text", - "text": parsed.content - }], - "model": request.model, - "stop_reason": "end_turn", - "usage": { - "input_tokens": 0, - "output_tokens": 0 - } - }); - - WsProtoMessage::Response(WsApiResponse { - request_id: request_id.to_string(), - payload: response, - }) - } - Err(e) => WsProtoMessage::Error(WsError::internal( - Some(request_id.to_string()), - e.to_string(), - )), - } - } else { - let body = resp.text().await.unwrap_or_default(); - WsProtoMessage::Error(WsError::upstream( - Some(request_id.to_string()), - format!("Upstream error: {}", body), - )) - } - } - Err(e) => WsProtoMessage::Error(WsError::internal( - Some(request_id.to_string()), - e.to_string(), - )), - } - } -} - -/// WebSocket 专用的 OpenAI 格式 Provider 调用 -async fn call_provider_openai_for_ws( - state: &AppState, - credential: &ProviderCredential, - request: &ChatCompletionRequest, -) -> Result { - use crate::models::provider_pool_model::CredentialData; - - match &credential.credential { - CredentialData::KiroOAuth { creds_file_path } => { - let mut kiro = KiroProvider::new(); - if let Err(e) = kiro.load_credentials_from_path(creds_file_path).await { - if let Some(db) = &state.db { - let _ = state.pool_service.mark_unhealthy( - db, - &credential.uuid, - Some(&format!("Failed to load credentials: {}", e)), - ); - } - return Err(e.to_string()); - } - if let Err(e) = kiro.refresh_token().await { - if let Some(db) = &state.db { - let _ = state.pool_service.mark_unhealthy( - db, - &credential.uuid, - Some(&format!("Token refresh failed: {}", e)), - ); - } - return Err(e.to_string()); - } - - let resp = match kiro.call_api(request).await { - Ok(r) => r, - Err(e) => { - if let Some(db) = &state.db { - let _ = state.pool_service.mark_unhealthy( - db, - &credential.uuid, - Some(&e.to_string()), - ); - } - return Err(e.to_string()); - } - }; - if resp.status().is_success() { - let body = resp.text().await.map_err(|e| e.to_string())?; - let parsed = parse_cw_response(&body); - let has_tool_calls = !parsed.tool_calls.is_empty(); - - // 记录成功 - if let Some(db) = &state.db { - let _ = - state - .pool_service - .mark_healthy(db, &credential.uuid, Some(&request.model)); - let _ = state.pool_service.record_usage(db, &credential.uuid); - } - - let message = if has_tool_calls { - serde_json::json!({ - "role": "assistant", - "content": if parsed.content.is_empty() { serde_json::Value::Null } else { serde_json::json!(parsed.content) }, - "tool_calls": parsed.tool_calls.iter().map(|tc| { - serde_json::json!({ - "id": tc.id, - "type": "function", - "function": { - "name": tc.function.name, - "arguments": tc.function.arguments - } - }) - }).collect::>() - }) - } else { - serde_json::json!({ - "role": "assistant", - "content": parsed.content - }) - }; - - Ok(serde_json::json!({ - "id": format!("chatcmpl-{}", uuid::Uuid::new_v4()), - "object": "chat.completion", - "created": std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap_or_default() - .as_secs(), - "model": request.model, - "choices": [{ - "index": 0, - "message": message, - "finish_reason": if has_tool_calls { "tool_calls" } else { "stop" } - }], - "usage": { - "prompt_tokens": 0, - "completion_tokens": 0, - "total_tokens": 0 - } - })) - } else { - let body = resp.text().await.unwrap_or_default(); - if let Some(db) = &state.db { - let _ = state - .pool_service - .mark_unhealthy(db, &credential.uuid, Some(&body)); - } - Err(format!("Upstream error: {}", body)) - } - } - CredentialData::OpenAIKey { api_key, base_url } => { - let provider = OpenAICustomProvider::with_config(api_key.clone(), base_url.clone()); - let resp = match provider.call_api(request).await { - Ok(r) => r, - Err(e) => { - if let Some(db) = &state.db { - let _ = state.pool_service.mark_unhealthy( - db, - &credential.uuid, - Some(&e.to_string()), - ); - } - return Err(e.to_string()); - } - }; - if resp.status().is_success() { - // 记录成功 - if let Some(db) = &state.db { - let _ = - state - .pool_service - .mark_healthy(db, &credential.uuid, Some(&request.model)); - let _ = state.pool_service.record_usage(db, &credential.uuid); - } - resp.json::() - .await - .map_err(|e| e.to_string()) - } else { - let body = resp.text().await.unwrap_or_default(); - if let Some(db) = &state.db { - let _ = state - .pool_service - .mark_unhealthy(db, &credential.uuid, Some(&body)); - } - Err(format!("Upstream error: {}", body)) - } - } - CredentialData::ClaudeKey { api_key, base_url } => { - // 打印 Claude 代理 URL 用于调试 - let actual_base_url = base_url.as_deref().unwrap_or("https://api.anthropic.com"); - tracing::info!( - "[CLAUDE] 使用 Claude API 代理: base_url={} credential_uuid={}", - actual_base_url, - &credential.uuid[..8] - ); - let provider = ClaudeCustomProvider::with_config(api_key.clone(), base_url.clone()); - match provider.call_openai_api(request).await { - Ok(result) => { - // 记录成功 - if let Some(db) = &state.db { - let _ = state.pool_service.mark_healthy( - db, - &credential.uuid, - Some(&request.model), - ); - let _ = state.pool_service.record_usage(db, &credential.uuid); - } - Ok(result) - } - Err(e) => { - if let Some(db) = &state.db { - let _ = state.pool_service.mark_unhealthy( - db, - &credential.uuid, - Some(&e.to_string()), - ); - } - Err(e.to_string()) - } - } - } - CredentialData::AntigravityOAuth { - creds_file_path, .. - } => { - let mut antigravity = AntigravityProvider::new(); - if let Err(e) = antigravity - .load_credentials_from_path(creds_file_path) - .await - { - if let Some(db) = &state.db { - let _ = state.pool_service.mark_unhealthy( - db, - &credential.uuid, - Some(&format!("Failed to load credentials: {}", e)), - ); - } - return Err(e.to_string()); - } - if !antigravity.is_token_valid() { - if let Err(e) = antigravity.refresh_token().await { - if let Some(db) = &state.db { - let _ = state.pool_service.mark_unhealthy( - db, - &credential.uuid, - Some(&format!("Token refresh failed: {}", e)), - ); - } - return Err(e.to_string()); - } - } - let antigravity_request = convert_openai_to_antigravity(request); - match antigravity - .call_api("generateContent", &antigravity_request) - .await - { - Ok(resp) => { - // 记录成功 - if let Some(db) = &state.db { - let _ = state.pool_service.mark_healthy( - db, - &credential.uuid, - Some(&request.model), - ); - let _ = state.pool_service.record_usage(db, &credential.uuid); - } - Ok(convert_antigravity_to_openai_response( - &resp, - &request.model, - )) - } - Err(e) => { - if let Some(db) = &state.db { - let _ = state.pool_service.mark_unhealthy( - db, - &credential.uuid, - Some(&e.to_string()), - ); - } - Err(e.to_string()) - } - } - } - // GeminiOAuth 和 QwenOAuth 暂不支持 WebSocket，需要使用 HTTP 端点 - _ => Err( - "This credential type is not yet supported via WebSocket. Please use HTTP endpoints." - .to_string(), - ), - } -} - -/// WebSocket 专用的 Anthropic 格式 Provider 调用 -async fn call_provider_anthropic_for_ws( - state: &AppState, - credential: &ProviderCredential, - request: &AnthropicMessagesRequest, -) -> Result { - use crate::models::provider_pool_model::CredentialData; - - match &credential.credential { - CredentialData::ClaudeKey { api_key, base_url } => { - // 打印 Claude 代理 URL 用于调试 - let actual_base_url = base_url.as_deref().unwrap_or("https://api.anthropic.com"); - tracing::info!( - "[CLAUDE] 使用 Claude API 代理: base_url={} credential_uuid={}", - actual_base_url, - &credential.uuid[..8] - ); - let provider = ClaudeCustomProvider::with_config(api_key.clone(), base_url.clone()); - let resp = match provider.call_api(request).await { - Ok(r) => r, - Err(e) => { - if let Some(db) = &state.db { - let _ = state.pool_service.mark_unhealthy( - db, - &credential.uuid, - Some(&e.to_string()), - ); - } - return Err(e.to_string()); - } - }; - if resp.status().is_success() { - // 记录成功 - if let Some(db) = &state.db { - let _ = - state - .pool_service - .mark_healthy(db, &credential.uuid, Some(&request.model)); - let _ = state.pool_service.record_usage(db, &credential.uuid); - } - resp.json::() - .await - .map_err(|e| e.to_string()) - } else { - let body = resp.text().await.unwrap_or_default(); - if let Some(db) = &state.db { - let _ = state - .pool_service - .mark_unhealthy(db, &credential.uuid, Some(&body)); - } - Err(format!("Upstream error: {}", body)) - } - } - _ => { - // 转换为 OpenAI 格式并调用（健康状态更新在 call_provider_openai_for_ws 中处理） - let openai_request = convert_anthropic_to_openai(request); - let result = call_provider_openai_for_ws(state, credential, &openai_request).await?; - - // 转换响应为 Anthropic 格式 - Ok(serde_json::json!({ - "id": format!("msg_{}", uuid::Uuid::new_v4()), - "type": "message", - "role": "assistant", - "content": [{ - "type": "text", - "text": result.get("choices") - .and_then(|c| c.get(0)) - .and_then(|c| c.get("message")) - .and_then(|m| m.get("content")) - .and_then(|c| c.as_str()) - .unwrap_or("") - }], - "model": request.model, - "stop_reason": "end_turn", - "usage": { - "input_tokens": 0, - "output_tokens": 0 - } - })) - } - } -} - -// ============ Management API Types and Handlers ============ - -/// 管理 API 状态响应 -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct ManagementStatusResponse { - /// 服务器是否运行中 - pub running: bool, - /// 监听地址 - pub host: String, - /// 监听端口 - pub port: u16, - /// 处理的请求数 - pub requests: u64, - /// 运行时间（秒） - pub uptime_secs: u64, - /// 版本号 - pub version: String, - /// TLS 是否启用 - pub tls_enabled: bool, - /// 默认 Provider - pub default_provider: String, -} - -/// 凭证信息（用于列表显示） -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct CredentialInfo { - /// 凭证 ID - pub id: String, - /// Provider 类型 - pub provider_type: String, - /// 是否禁用 - pub disabled: bool, - /// 是否有效 - pub is_valid: bool, -} - -/// 凭证列表响应 -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct CredentialsListResponse { - /// 凭证列表 - pub credentials: Vec, - /// 总数 - pub total: usize, -} - -/// 添加凭证请求 -#[derive(Debug, Clone, Deserialize)] -pub struct AddCredentialRequest { - /// Provider 类型 - pub provider_type: String, - /// 凭证 ID - pub id: String, - /// API Key（用于 API Key 类型的凭证） - #[serde(default)] - pub api_key: Option, - /// Token 文件路径（用于 OAuth 类型的凭证） - #[serde(default)] - pub token_file: Option, - /// Base URL - #[serde(default)] - pub base_url: Option, - /// 代理 URL - #[serde(default)] - pub proxy_url: Option, -} - -/// 添加凭证响应 -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct AddCredentialResponse { - /// 是否成功 - pub success: bool, - /// 消息 - pub message: String, - /// 凭证 ID - pub id: Option, -} - -/// 配置响应（简化版，不包含敏感信息） -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct ManagementConfigResponse { - /// 服务器配置 - pub server: ManagementServerConfigInfo, - /// 路由配置 - pub routing: ManagementRoutingConfigInfo, - /// 重试配置 - pub retry: ManagementRetryConfigInfo, - /// 远程管理配置（不包含 secret_key） - pub remote_management: ManagementRemoteInfo, -} - -/// 服务器配置信息 -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct ManagementServerConfigInfo { - pub host: String, - pub port: u16, - pub tls_enabled: bool, -} - -/// 路由配置信息 -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct ManagementRoutingConfigInfo { - pub default_provider: String, - pub rules_count: usize, -} - -/// 重试配置信息 -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct ManagementRetryConfigInfo { - pub max_retries: u32, - pub base_delay_ms: u64, - pub max_delay_ms: u64, -} - -/// 远程管理配置信息 -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct ManagementRemoteInfo { - pub allow_remote: bool, - pub has_secret_key: bool, - pub disable_control_panel: bool, -} - -/// 更新配置请求 -#[derive(Debug, Clone, Deserialize)] -pub struct UpdateConfigRequest { - /// 默认 Provider - #[serde(default)] - pub default_provider: Option, - /// 是否允许远程访问 - #[serde(default)] - pub allow_remote: Option, -} - -/// 更新配置响应 -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct UpdateConfigResponse { - pub success: bool, - pub message: String, -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct BackupResponse { - pub success: bool, - pub message: String, - pub backup_path: Option, -} - -#[derive(Debug, Clone, Deserialize)] -pub struct RestoreRequest { - pub backup_path: String, -} - -fn snapshot_config(state: &AppState) -> Option { - if let Some(manager) = &state.config_manager { - if let Ok(guard) = manager.read() { - return Some(guard.config().clone()); - } - } - state.config.clone() -} - -/// GET /v0/management/status - 获取服务器状态 -pub async fn management_status(State(state): State) -> impl IntoResponse { - let default_provider = state.default_provider.read().await.clone(); - - // 获取请求数量 - let requests = state.processor.stats.read().len() as u64; - let config = snapshot_config(&state).unwrap_or_default(); - let uptime_secs = state.start_time.elapsed().as_secs(); - - let response = ManagementStatusResponse { - running: true, - host: config.server.host, - port: config.server.port, - requests, - uptime_secs, - version: env!("CARGO_PKG_VERSION").to_string(), - tls_enabled: config.server.tls.enable, - default_provider, - }; - - Json(response) -} - -/// POST /v0/management/backup - 触发数据库备份 -pub async fn management_backup(State(state): State) -> impl IntoResponse { - let Some(service) = &state.backup_service else { - return ( - StatusCode::SERVICE_UNAVAILABLE, - Json(BackupResponse { - success: false, - message: "Backup service not available".to_string(), - backup_path: None, - }), - ); - }; - - let result = match &state.db { - Some(db) => service.backup_database_with_connection(db), - None => service.backup_database(), - }; - - match result { - Ok(path) => ( - StatusCode::OK, - Json(BackupResponse { - success: true, - message: "Backup created".to_string(), - backup_path: Some(path.to_string_lossy().to_string()), - }), - ), - Err(err) => ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(BackupResponse { - success: false, - message: err, - backup_path: None, - }), - ), - } -} - -/// POST /v0/management/restore - 从备份恢复数据库 -pub async fn management_restore( - State(state): State, - Json(request): Json, -) -> impl IntoResponse { - let Some(service) = &state.backup_service else { - return ( - StatusCode::SERVICE_UNAVAILABLE, - Json(BackupResponse { - success: false, - message: "Backup service not available".to_string(), - backup_path: None, - }), - ); - }; - - let backup_path = PathBuf::from(request.backup_path); - let result = match &state.db { - Some(db) => service.restore_database_with_connection(db, &backup_path), - None => service.restore_database(&backup_path), - }; - - match result { - Ok(()) => ( - StatusCode::OK, - Json(BackupResponse { - success: true, - message: "Restore completed".to_string(), - backup_path: None, - }), - ), - Err(err) => ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(BackupResponse { - success: false, - message: err, - backup_path: None, - }), - ), - } -} - -/// GET /v0/management/credentials - 获取凭证列表 -pub async fn management_list_credentials(State(state): State) -> impl IntoResponse { - let mut credentials = Vec::new(); - - // 从数据库获取凭证列表 - if let Some(ref db) = state.db { - if let Ok(conn) = db.lock() { - if let Ok(pool_credentials) = ProviderPoolDao::get_all(&conn) { - for cred in pool_credentials { - credentials.push(CredentialInfo { - id: cred.uuid.clone(), - provider_type: cred.provider_type.to_string(), - disabled: cred.is_disabled, - is_valid: cred.is_healthy, - }); - } - } - } - } - - let total = credentials.len(); - Json(CredentialsListResponse { credentials, total }) -} - -/// POST /v0/management/credentials - 添加凭证 -pub async fn management_add_credential( - State(state): State, - Json(request): Json, -) -> impl IntoResponse { - use crate::models::provider_pool_model::{ - CredentialData, PoolProviderType, ProviderCredential, - }; - - // 验证请求 - if request.id.is_empty() { - return ( - StatusCode::BAD_REQUEST, - Json(AddCredentialResponse { - success: false, - message: "Credential ID is required".to_string(), - id: None, - }), - ); - } - - if request.provider_type.is_empty() { - return ( - StatusCode::BAD_REQUEST, - Json(AddCredentialResponse { - success: false, - message: "Provider type is required".to_string(), - id: None, - }), - ); - } - - // 解析 provider 类型 - let provider_type: PoolProviderType = match request.provider_type.parse() { - Ok(pt) => pt, - Err(_) => { - return ( - StatusCode::BAD_REQUEST, - Json(AddCredentialResponse { - success: false, - message: format!("Invalid provider type: {}", request.provider_type), - id: None, - }), - ); - } - }; - - // 根据 provider 类型创建凭证数据 - let credential_data = match provider_type { - PoolProviderType::OpenAI => { - if let Some(api_key) = request.api_key { - CredentialData::OpenAIKey { - api_key, - base_url: request.base_url, - } - } else { - return ( - StatusCode::BAD_REQUEST, - Json(AddCredentialResponse { - success: false, - message: "API key is required for OpenAI provider".to_string(), - id: None, - }), - ); - } - } - PoolProviderType::Claude => { - if let Some(api_key) = request.api_key { - CredentialData::ClaudeKey { - api_key, - base_url: request.base_url, - } - } else { - return ( - StatusCode::BAD_REQUEST, - Json(AddCredentialResponse { - success: false, - message: "API key is required for Claude provider".to_string(), - id: None, - }), - ); - } - } - PoolProviderType::Vertex => { - if let Some(api_key) = request.api_key { - CredentialData::VertexKey { - api_key, - base_url: request.base_url, - model_aliases: std::collections::HashMap::new(), - } - } else { - return ( - StatusCode::BAD_REQUEST, - Json(AddCredentialResponse { - success: false, - message: "API key is required for Vertex provider".to_string(), - id: None, - }), - ); - } - } - PoolProviderType::Kiro => { - if let Some(token_file) = request.token_file { - CredentialData::KiroOAuth { - creds_file_path: token_file, - } - } else { - return ( - StatusCode::BAD_REQUEST, - Json(AddCredentialResponse { - success: false, - message: "Token file is required for Kiro provider".to_string(), - id: None, - }), - ); - } - } - PoolProviderType::Gemini => { - if let Some(token_file) = request.token_file { - CredentialData::GeminiOAuth { - creds_file_path: token_file, - project_id: None, - } - } else { - return ( - StatusCode::BAD_REQUEST, - Json(AddCredentialResponse { - success: false, - message: "Token file is required for Gemini provider".to_string(), - id: None, - }), - ); - } - } - PoolProviderType::Qwen => { - if let Some(token_file) = request.token_file { - CredentialData::QwenOAuth { - creds_file_path: token_file, - } - } else { - return ( - StatusCode::BAD_REQUEST, - Json(AddCredentialResponse { - success: false, - message: "Token file is required for Qwen provider".to_string(), - id: None, - }), - ); - } - } - PoolProviderType::Antigravity => { - if let Some(token_file) = request.token_file { - CredentialData::AntigravityOAuth { - creds_file_path: token_file, - project_id: None, - } - } else { - return ( - StatusCode::BAD_REQUEST, - Json(AddCredentialResponse { - success: false, - message: "Token file is required for Antigravity provider".to_string(), - id: None, - }), - ); - } - } - PoolProviderType::GeminiApiKey => { - if let Some(api_key) = request.api_key { - CredentialData::GeminiApiKey { - api_key, - base_url: request.base_url, - excluded_models: Vec::new(), - } - } else { - return ( - StatusCode::BAD_REQUEST, - Json(AddCredentialResponse { - success: false, - message: "API key is required for Gemini API Key provider".to_string(), - id: None, - }), - ); - } - } - PoolProviderType::Codex => { - if let Some(token_file) = request.token_file { - CredentialData::CodexOAuth { - creds_file_path: token_file, - } - } else { - return ( - StatusCode::BAD_REQUEST, - Json(AddCredentialResponse { - success: false, - message: "Token file is required for Codex provider".to_string(), - id: None, - }), - ); - } - } - PoolProviderType::ClaudeOAuth => { - if let Some(token_file) = request.token_file { - CredentialData::ClaudeOAuth { - creds_file_path: token_file, - } - } else { - return ( - StatusCode::BAD_REQUEST, - Json(AddCredentialResponse { - success: false, - message: "Token file is required for Claude OAuth provider".to_string(), - id: None, - }), - ); - } - } - PoolProviderType::IFlow => { - if let Some(token_file) = request.token_file { - // 默认使用 OAuth 类型，Cookie 类型需要通过其他方式添加 - CredentialData::IFlowOAuth { - creds_file_path: token_file, - } - } else { - return ( - StatusCode::BAD_REQUEST, - Json(AddCredentialResponse { - success: false, - message: "Token file is required for iFlow provider".to_string(), - id: None, - }), - ); - } - } - }; - - // 创建凭证 - let mut credential = ProviderCredential::new(provider_type, credential_data); - credential.uuid = request.id.clone(); - credential.name = Some(request.id.clone()); - - // 添加凭证到数据库 - if let Some(ref db) = state.db { - if let Ok(conn) = db.lock() { - match ProviderPoolDao::insert(&conn, &credential) { - Ok(_) => { - tracing::info!( - "[MANAGEMENT] Added credential: {} ({})", - request.id, - request.provider_type - ); - return ( - StatusCode::CREATED, - Json(AddCredentialResponse { - success: true, - message: "Credential added successfully".to_string(), - id: Some(request.id), - }), - ); - } - Err(e) => { - tracing::error!("[MANAGEMENT] Failed to add credential: {}", e); - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(AddCredentialResponse { - success: false, - message: format!("Failed to add credential: {}", e), - id: None, - }), - ); - } - } - } - } - - ( - StatusCode::SERVICE_UNAVAILABLE, - Json(AddCredentialResponse { - success: false, - message: "Database not available".to_string(), - id: None, - }), - ) -} - -/// GET /v0/management/config - 获取配置 -pub async fn management_get_config(State(state): State) -> impl IntoResponse { - let default_provider = state.default_provider.read().await.clone(); - - // 获取路由规则数量 - let rules_count = state.processor.router.read().await.rules().len(); - let config = snapshot_config(&state).unwrap_or_default(); - - let response = ManagementConfigResponse { - server: ManagementServerConfigInfo { - host: config.server.host, - port: config.server.port, - tls_enabled: config.server.tls.enable, - }, - routing: ManagementRoutingConfigInfo { - default_provider, - rules_count, - }, - retry: ManagementRetryConfigInfo { - max_retries: config.retry.max_retries, - base_delay_ms: config.retry.base_delay_ms, - max_delay_ms: config.retry.max_delay_ms, - }, - remote_management: ManagementRemoteInfo { - allow_remote: config.remote_management.allow_remote, - has_secret_key: config - .remote_management - .secret_key - .as_ref() - .map(|key| !key.is_empty()) - .unwrap_or(false), - disable_control_panel: config.remote_management.disable_control_panel, - }, - }; - - Json(response) -} - -/// PUT /v0/management/config - 更新配置 -pub async fn management_update_config( - State(state): State, - Json(request): Json, -) -> impl IntoResponse { - let mut updated = false; - let mut needs_restart = false; - - // 更新默认 Provider - if let Some(provider) = request.default_provider { - // 验证 provider 类型 - if provider.parse::().is_ok() { - let mut dp = state.default_provider.write().await; - *dp = provider.clone(); - tracing::info!("[MANAGEMENT] Updated default_provider to: {}", provider); - if let Some(manager) = &state.config_manager { - if let Ok(mut guard) = manager.write() { - guard.config_mut().default_provider = provider.clone(); - guard.config_mut().routing.default_provider = provider.clone(); - if let Err(err) = guard.save() { - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(UpdateConfigResponse { - success: false, - message: format!("Failed to save config: {}", err), - }), - ); - } - } else { - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(UpdateConfigResponse { - success: false, - message: "Failed to lock config manager".to_string(), - }), - ); - } - } - updated = true; - } else { - return ( - StatusCode::BAD_REQUEST, - Json(UpdateConfigResponse { - success: false, - message: format!("Invalid provider type: {}", provider), - }), - ); - } - } - - // 更新是否允许远程访问（需要重启生效） - if let Some(allow_remote) = request.allow_remote { - if allow_remote { - return ( - StatusCode::BAD_REQUEST, - Json(UpdateConfigResponse { - success: false, - message: "当前版本未启用 TLS，禁止开启远程管理".to_string(), - }), - ); - } - let manager = match &state.config_manager { - Some(manager) => manager, - None => { - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(UpdateConfigResponse { - success: false, - message: "Config manager is not available".to_string(), - }), - ); - } - }; - if let Ok(mut guard) = manager.write() { - guard.config_mut().remote_management.allow_remote = allow_remote; - if let Err(err) = guard.save() { - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(UpdateConfigResponse { - success: false, - message: format!("Failed to save config: {}", err), - }), - ); - } - } else { - return ( - StatusCode::INTERNAL_SERVER_ERROR, - Json(UpdateConfigResponse { - success: false, - message: "Failed to lock config manager".to_string(), - }), - ); - } - tracing::info!( - "[MANAGEMENT] Updated remote_management.allow_remote to: {}", - allow_remote - ); - updated = true; - needs_restart = true; - } - - if updated { - ( - StatusCode::OK, - Json(UpdateConfigResponse { - success: true, - message: if needs_restart { - "Configuration updated. Restart required to apply all changes.".to_string() - } else { - "Configuration updated successfully".to_string() - }, - }), - ) - } else { - ( - StatusCode::OK, - Json(UpdateConfigResponse { - success: true, - message: "No changes applied".to_string(), - }), - ) - } -} diff --git a/src-tauri/src/server/handlers/api.rs b/src-tauri/src/server/handlers/api.rs new file mode 100644 index 000000000..159f0a9bc --- /dev/null +++ b/src-tauri/src/server/handlers/api.rs @@ -0,0 +1,982 @@ +//! API 端点处理器 +//! +//! 处理 OpenAI 和 Anthropic 格式的 API 请求 + +use axum::{ + body::Body, + extract::State, + http::{header, HeaderMap, StatusCode}, + response::{IntoResponse, Response}, + Json, +}; +use futures::stream; + +use crate::converter::anthropic_to_openai::convert_anthropic_to_openai; +use crate::converter::openai_to_antigravity::{ + convert_antigravity_to_openai_response, convert_openai_to_antigravity_with_context, +}; +use crate::models::anthropic::AnthropicMessagesRequest; +use crate::models::openai::ChatCompletionRequest; +use crate::processor::RequestContext; +use crate::providers::{AntigravityProvider, GeminiProvider, KiroProvider, QwenProvider}; +use crate::server::{record_request_telemetry, record_token_usage, AppState}; +use crate::server_utils::{ + build_anthropic_response, build_anthropic_stream_response, message_content_len, + parse_cw_response, safe_truncate, +}; +use crate::telemetry::RequestStatus; +use crate::ProviderType; + +use super::{call_provider_anthropic, call_provider_openai}; + +/// OpenAI 格式的 API key 验证 +pub async fn verify_api_key( + headers: &HeaderMap, + expected_key: &str, +) -> Result<(), (StatusCode, Json)> { + let auth = headers + .get("authorization") + .or_else(|| headers.get("x-api-key")) + .and_then(|v| v.to_str().ok()); + + let key = match auth { + Some(s) if s.starts_with("Bearer ") => &s[7..], + Some(s) => s, + None => { + return Err(( + StatusCode::UNAUTHORIZED, + Json(serde_json::json!({"error": {"message": "No API key provided"}})), + )) + } + }; + + if key != expected_key { + return Err(( + StatusCode::UNAUTHORIZED, + Json(serde_json::json!({"error": {"message": "Invalid API key"}})), + )); + } + + Ok(()) +} + +/// Anthropic 格式的 API key 验证 +pub async fn verify_api_key_anthropic( + headers: &HeaderMap, + expected_key: &str, +) -> Result<(), (StatusCode, Json)> { + let auth = headers + .get("x-api-key") + .or_else(|| headers.get("authorization")) + .and_then(|v| v.to_str().ok()); + + let key = match auth { + Some(s) if s.starts_with("Bearer ") => &s[7..], + Some(s) => s, + None => { + return Err(( + StatusCode::UNAUTHORIZED, + Json(serde_json::json!({ + "type": "error", + "error": { + "type": "authentication_error", + "message": "No API key provided. Please set the x-api-key header." + } + })), + )) + } + }; + + if key != expected_key { + return Err(( + StatusCode::UNAUTHORIZED, + Json(serde_json::json!({ + "type": "error", + "error": { + "type": "authentication_error", + "message": "Invalid API key" + } + })), + )); + } + + Ok(()) +} + +pub async fn chat_completions( + State(state): State, + headers: HeaderMap, + Json(mut request): Json, +) -> Response { + if let Err(e) = verify_api_key(&headers, &state.api_key).await { + state + .logs + .write() + .await + .add("warn", "Unauthorized request to /v1/chat/completions"); + return e.into_response(); + } + + // 创建请求上下文 + let mut ctx = RequestContext::new(request.model.clone()).with_stream(request.stream); + + state.logs.write().await.add( + "info", + &format!( + "POST /v1/chat/completions request_id={} model={} stream={}", + ctx.request_id, request.model, request.stream + ), + ); + + // 使用 RequestProcessor 解析模型别名和路由 + let provider = state.processor.resolve_and_route(&mut ctx).await; + + // 更新请求中的模型名为解析后的模型 + if ctx.resolved_model != ctx.original_model { + request.model = ctx.resolved_model.clone(); + state.logs.write().await.add( + "info", + &format!( + "[MAPPER] request_id={} alias={} -> model={}", + ctx.request_id, ctx.original_model, ctx.resolved_model + ), + ); + } + + // 应用参数注入 + let injection_enabled = *state.injection_enabled.read().await; + if injection_enabled { + let injector = state.processor.injector.read().await; + let mut payload = serde_json::to_value(&request).unwrap_or_default(); + let result = injector.inject(&request.model, &mut payload); + if result.has_injections() { + state.logs.write().await.add( + "info", + &format!( + "[INJECT] request_id={} applied_rules={:?} injected_params={:?}", + ctx.request_id, result.applied_rules, result.injected_params + ), + ); + // 更新请求 + if let Ok(updated) = serde_json::from_value(payload) { + request = updated; + } + } + } + + // 获取当前默认 provider（用于凭证池选择） + let default_provider = state.default_provider.read().await.clone(); + + // 记录路由结果 + state.logs.write().await.add( + "info", + &format!( + "[ROUTE] request_id={} model={} provider={}", + ctx.request_id, ctx.resolved_model, provider + ), + ); + + // 尝试从凭证池中选择凭证 + let credential = match &state.db { + Some(db) => state + .pool_service + .select_credential(db, &default_provider, Some(&request.model)) + .ok() + .flatten(), + None => None, + }; + + // 如果找到凭证池中的凭证，使用它 + if let Some(cred) = credential { + state.logs.write().await.add( + "info", + &format!( + "[ROUTE] Using pool credential: type={} name={:?} uuid={}", + cred.provider_type, + cred.name, + &cred.uuid[..8] + ), + ); + let response = call_provider_openai(&state, &cred, &request).await; + + // 记录请求统计 + let is_success = response.status().is_success(); + let status = if is_success { + crate::telemetry::RequestStatus::Success + } else { + crate::telemetry::RequestStatus::Failed + }; + record_request_telemetry(&state, &ctx, status, None); + + // 如果成功，记录估算的 Token 使用量 + if is_success { + let estimated_input_tokens = request + .messages + .iter() + .map(|m| { + let content_len = match &m.content { + Some(c) => message_content_len(c), + None => 0, + }; + content_len / 4 + }) + .sum::() as u32; + // 输出 Token 使用估算值（假设平均响应长度） + let estimated_output_tokens = 100u32; + record_token_usage( + &state, + &ctx, + Some(estimated_input_tokens), + Some(estimated_output_tokens), + ); + } + + return response; + } + + // 回退到旧的单凭证模式 + state.logs.write().await.add( + "debug", + &format!( + "[ROUTE] No pool credential found for '{}', using legacy mode", + default_provider + ), + ); + + // 检查是否需要刷新 token（无 token 或即将过期） + { + let _guard = state.kiro_refresh_lock.lock().await; + let mut kiro = state.kiro.write().await; + let needs_refresh = + kiro.credentials.access_token.is_none() || kiro.is_token_expiring_soon(); + if needs_refresh { + if let Err(e) = kiro.refresh_token().await { + state + .logs + .write() + .await + .add("error", &format!("Token refresh failed: {e}")); + return ( + StatusCode::UNAUTHORIZED, + Json(serde_json::json!({"error": {"message": format!("Token refresh failed: {e}")}})), + ).into_response(); + } + } + } + + let kiro = state.kiro.read().await; + + match kiro.call_api(&request).await { + Ok(resp) => { + let status = resp.status(); + if status.is_success() { + match resp.text().await { + Ok(body) => { + let parsed = parse_cw_response(&body); + let has_tool_calls = !parsed.tool_calls.is_empty(); + + state.logs.write().await.add( + "info", + &format!( + "Request completed: content_len={}, tool_calls={}", + parsed.content.len(), + parsed.tool_calls.len() + ), + ); + + // 构建消息 + let message = if has_tool_calls { + serde_json::json!({ + "role": "assistant", + "content": if parsed.content.is_empty() { serde_json::Value::Null } else { serde_json::json!(parsed.content) }, + "tool_calls": parsed.tool_calls.iter().map(|tc| { + serde_json::json!({ + "id": tc.id, + "type": "function", + "function": { + "name": tc.function.name, + "arguments": tc.function.arguments + } + }) + }).collect::>() + }) + } else { + serde_json::json!({ + "role": "assistant", + "content": parsed.content + }) + }; + + // 估算 Token 数量（基于字符数，约 4 字符 = 1 token） + let estimated_output_tokens = (parsed.content.len() / 4) as u32; + // 估算输入 Token（基于请求消息） + let estimated_input_tokens = request + .messages + .iter() + .map(|m| { + let content_len = match &m.content { + Some(c) => message_content_len(c), + None => 0, + }; + content_len / 4 + }) + .sum::() + as u32; + + let response = serde_json::json!({ + "id": format!("chatcmpl-{}", uuid::Uuid::new_v4()), + "object": "chat.completion", + "created": std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_secs(), + "model": request.model, + "choices": [{ + "index": 0, + "message": message, + "finish_reason": if has_tool_calls { "tool_calls" } else { "stop" } + }], + "usage": { + "prompt_tokens": estimated_input_tokens, + "completion_tokens": estimated_output_tokens, + "total_tokens": estimated_input_tokens + estimated_output_tokens + } + }); + // 记录成功请求统计 + record_request_telemetry( + &state, + &ctx, + crate::telemetry::RequestStatus::Success, + None, + ); + // 记录 Token 使用量 + record_token_usage( + &state, + &ctx, + Some(estimated_input_tokens), + Some(estimated_output_tokens), + ); + Json(response).into_response() + } + Err(e) => { + // 记录失败请求统计 + record_request_telemetry( + &state, + &ctx, + crate::telemetry::RequestStatus::Failed, + Some(e.to_string()), + ); + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"error": {"message": e.to_string()}})), + ) + .into_response() + } + } + } else if status.as_u16() == 403 || status.as_u16() == 402 { + // Token 过期或账户问题，尝试重新加载凭证并刷新 + drop(kiro); + let _guard = state.kiro_refresh_lock.lock().await; + let mut kiro = state.kiro.write().await; + state.logs.write().await.add( + "warn", + &format!( + "[AUTH] Got {}, reloading credentials and attempting token refresh...", + status.as_u16() + ), + ); + + // 先重新加载凭证文件（可能用户换了账户） + if let Err(e) = kiro.load_credentials().await { + state.logs.write().await.add( + "error", + &format!("[AUTH] Failed to reload credentials: {e}"), + ); + } + + match kiro.refresh_token().await { + Ok(_) => { + state + .logs + .write() + .await + .add("info", "[AUTH] Token refreshed successfully after reload"); + // 重试请求 + drop(kiro); + let kiro = state.kiro.read().await; + match kiro.call_api(&request).await { + Ok(retry_resp) => { + if retry_resp.status().is_success() { + match retry_resp.text().await { + Ok(body) => { + let parsed = parse_cw_response(&body); + let has_tool_calls = !parsed.tool_calls.is_empty(); + + let message = if has_tool_calls { + serde_json::json!({ + "role": "assistant", + "content": if parsed.content.is_empty() { serde_json::Value::Null } else { serde_json::json!(parsed.content) }, + "tool_calls": parsed.tool_calls.iter().map(|tc| { + serde_json::json!({ + "id": tc.id, + "type": "function", + "function": { + "name": tc.function.name, + "arguments": tc.function.arguments + } + }) + }).collect::>() + }) + } else { + serde_json::json!({ + "role": "assistant", + "content": parsed.content + }) + }; + + let response = serde_json::json!({ + "id": format!("chatcmpl-{}", uuid::Uuid::new_v4()), + "object": "chat.completion", + "created": std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_secs(), + "model": request.model, + "choices": [{ + "index": 0, + "message": message, + "finish_reason": if has_tool_calls { "tool_calls" } else { "stop" } + }], + "usage": { + "prompt_tokens": 0, + "completion_tokens": 0, + "total_tokens": 0 + } + }); + return Json(response).into_response(); + } + Err(e) => return ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"error": {"message": e.to_string()}})), + ).into_response(), + } + } + let body = retry_resp.text().await.unwrap_or_default(); + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"error": {"message": format!("Retry failed: {}", body)}})), + ).into_response() + } + Err(e) => ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"error": {"message": e.to_string()}})), + ) + .into_response(), + } + } + Err(e) => { + state + .logs + .write() + .await + .add("error", &format!("[AUTH] Token refresh failed: {e}")); + ( + StatusCode::UNAUTHORIZED, + Json(serde_json::json!({"error": {"message": format!("Token refresh failed: {e}")}})), + ) + .into_response() + } + } + } else { + let body = resp.text().await.unwrap_or_default(); + state.logs.write().await.add( + "error", + &format!("Upstream error {}: {}", status, safe_truncate(&body, 200)), + ); + ( + StatusCode::from_u16(status.as_u16()).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR), + Json(serde_json::json!({"error": {"message": format!("Upstream error: {}", body)}})) + ).into_response() + } + } + Err(e) => { + state + .logs + .write() + .await + .add("error", &format!("API call failed: {e}")); + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"error": {"message": e.to_string()}})), + ) + .into_response() + } + } +} + +pub async fn anthropic_messages( + State(state): State, + headers: HeaderMap, + Json(mut request): Json, +) -> Response { + // 使用 Anthropic 格式的认证验证（优先检查 x-api-key） + if let Err(e) = verify_api_key_anthropic(&headers, &state.api_key).await { + state + .logs + .write() + .await + .add("warn", "Unauthorized request to /v1/messages"); + return e.into_response(); + } + + // 创建请求上下文 + let mut ctx = RequestContext::new(request.model.clone()).with_stream(request.stream); + + // 详细记录请求信息 + let msg_count = request.messages.len(); + let has_tools = request.tools.as_ref().map(|t| t.len()).unwrap_or(0); + let has_system = request.system.is_some(); + state.logs.write().await.add( + "info", + &format!( + "[REQ] POST /v1/messages request_id={} model={} stream={} messages={} tools={} has_system={}", + ctx.request_id, request.model, request.stream, msg_count, has_tools, has_system + ), + ); + + // 使用 RequestProcessor 解析模型别名和路由 + let provider = state.processor.resolve_and_route(&mut ctx).await; + + // 更新请求中的模型名为解析后的模型 + if ctx.resolved_model != ctx.original_model { + request.model = ctx.resolved_model.clone(); + state.logs.write().await.add( + "info", + &format!( + "[MAPPER] request_id={} alias={} -> model={}", + ctx.request_id, ctx.original_model, ctx.resolved_model + ), + ); + } + + // 记录最后一条消息的角色和内容预览 + if let Some(last_msg) = request.messages.last() { + let content_preview = match &last_msg.content { + serde_json::Value::String(s) => s.chars().take(100).collect::(), + serde_json::Value::Array(arr) => { + if let Some(first) = arr.first() { + if let Some(text) = first.get("text").and_then(|t| t.as_str()) { + text.chars().take(100).collect::() + } else { + format!("[{} blocks]", arr.len()) + } + } else { + "[empty]".to_string() + } + } + _ => "[unknown]".to_string(), + }; + state.logs.write().await.add( + "debug", + &format!( + "[REQ] request_id={} last_message: role={} content={}", + ctx.request_id, last_msg.role, content_preview + ), + ); + } + + // 应用参数注入 + let injection_enabled = *state.injection_enabled.read().await; + if injection_enabled { + let injector = state.processor.injector.read().await; + let mut payload = serde_json::to_value(&request).unwrap_or_default(); + let result = injector.inject(&request.model, &mut payload); + if result.has_injections() { + state.logs.write().await.add( + "info", + &format!( + "[INJECT] request_id={} applied_rules={:?} injected_params={:?}", + ctx.request_id, result.applied_rules, result.injected_params + ), + ); + // 更新请求 + if let Ok(updated) = serde_json::from_value(payload) { + request = updated; + } + } + } + + // 获取当前默认 provider（用于凭证池选择） + let default_provider = state.default_provider.read().await.clone(); + + // 记录路由结果 + state.logs.write().await.add( + "info", + &format!( + "[ROUTE] request_id={} model={} provider={}", + ctx.request_id, ctx.resolved_model, provider + ), + ); + + // 尝试从凭证池中选择凭证 + let credential = match &state.db { + Some(db) => { + // 根据 default_provider 配置选择凭证 + state + .pool_service + .select_credential(db, &default_provider, Some(&request.model)) + .ok() + .flatten() + } + None => None, + }; + + // 如果找到凭证池中的凭证，使用它 + if let Some(cred) = credential { + state.logs.write().await.add( + "info", + &format!( + "[ROUTE] Using pool credential: type={} name={:?} uuid={}", + cred.provider_type, + cred.name, + &cred.uuid[..8] + ), + ); + let response = call_provider_anthropic(&state, &cred, &request).await; + + // 记录请求统计 + let is_success = response.status().is_success(); + let status = if is_success { + crate::telemetry::RequestStatus::Success + } else { + crate::telemetry::RequestStatus::Failed + }; + record_request_telemetry(&state, &ctx, status, None); + + // 如果成功，记录估算的 Token 使用量 + if is_success { + let estimated_input_tokens = request + .messages + .iter() + .map(|m| { + let content_len = match &m.content { + serde_json::Value::String(s) => s.len(), + serde_json::Value::Array(arr) => arr + .iter() + .filter_map(|v| v.get("text").and_then(|t| t.as_str())) + .map(|s| s.len()) + .sum(), + _ => 0, + }; + content_len / 4 + }) + .sum::() as u32; + // 输出 Token 使用估算值 + let estimated_output_tokens = 100u32; + record_token_usage( + &state, + &ctx, + Some(estimated_input_tokens), + Some(estimated_output_tokens), + ); + } + + return response; + } + + // 回退到旧的单凭证模式 + state.logs.write().await.add( + "debug", + &format!( + "[ROUTE] No pool credential found for '{}', using legacy mode", + default_provider + ), + ); + + // 检查是否需要刷新 token（无 token 或即将过期） + { + let _guard = state.kiro_refresh_lock.lock().await; + let mut kiro = state.kiro.write().await; + let needs_refresh = + kiro.credentials.access_token.is_none() || kiro.is_token_expiring_soon(); + if needs_refresh { + state.logs.write().await.add( + "info", + "[AUTH] No access token or token expiring soon, attempting refresh...", + ); + if let Err(e) = kiro.refresh_token().await { + state + .logs + .write() + .await + .add("error", &format!("[AUTH] Token refresh failed: {e}")); + return ( + StatusCode::UNAUTHORIZED, + Json(serde_json::json!({"error": {"message": format!("Token refresh failed: {e}")}})), + ) + .into_response(); + } + state + .logs + .write() + .await + .add("info", "[AUTH] Token refreshed successfully"); + } + } + + // 转换为 OpenAI 格式 + let openai_request = convert_anthropic_to_openai(&request); + + // 记录转换后的请求信息 + state.logs.write().await.add( + "debug", + &format!( + "[CONVERT] OpenAI format: messages={} tools={} stream={}", + openai_request.messages.len(), + openai_request.tools.as_ref().map(|t| t.len()).unwrap_or(0), + openai_request.stream + ), + ); + + let kiro = state.kiro.read().await; + + match kiro.call_api(&openai_request).await { + Ok(resp) => { + let status = resp.status(); + state + .logs + .write() + .await + .add("info", &format!("[RESP] Upstream status: {status}")); + + if status.is_success() { + match resp.bytes().await { + Ok(bytes) => { + // 使用 lossy 转换，避免无效 UTF-8 导致崩溃 + let body = String::from_utf8_lossy(&bytes).to_string(); + + // 记录原始响应长度 + state.logs.write().await.add( + "debug", + &format!("[RESP] Raw body length: {} bytes", bytes.len()), + ); + + // 保存原始响应到文件用于调试 + let request_id = uuid::Uuid::new_v4().to_string()[..8].to_string(); + state.logs.read().await.log_raw_response(&request_id, &body); + state.logs.write().await.add( + "debug", + &format!("[RESP] Raw response saved to raw_response_{request_id}.txt"), + ); + + // 记录响应的前200字符用于调试（减少日志量） + let preview: String = + body.chars().filter(|c| !c.is_control()).take(200).collect(); + state + .logs + .write() + .await + .add("debug", &format!("[RESP] Body preview: {preview}")); + + let parsed = parse_cw_response(&body); + + // 详细记录解析结果 + state.logs.write().await.add( + "info", + &format!( + "[RESP] Parsed: content_len={}, tool_calls={}, content_preview={}", + parsed.content.len(), + parsed.tool_calls.len(), + parsed.content.chars().take(100).collect::() + ), + ); + + // 记录 tool calls 详情 + for (i, tc) in parsed.tool_calls.iter().enumerate() { + state.logs.write().await.add( + "debug", + &format!( + "[RESP] Tool call {}: name={} id={}", + i, tc.function.name, tc.id + ), + ); + } + + // 如果请求流式响应，返回 SSE 格式 + if request.stream { + return build_anthropic_stream_response(&request.model, &parsed); + } + + // 非流式响应 + build_anthropic_response(&request.model, &parsed) + } + Err(e) => { + state + .logs + .write() + .await + .add("error", &format!("[ERROR] Response body read failed: {e}")); + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"error": {"message": e.to_string()}})), + ) + .into_response() + } + } + } else if status.as_u16() == 403 || status.as_u16() == 402 { + // Token 过期或账户问题，尝试重新加载凭证并刷新 + drop(kiro); + let _guard = state.kiro_refresh_lock.lock().await; + let mut kiro = state.kiro.write().await; + state.logs.write().await.add( + "warn", + &format!( + "[AUTH] Got {}, reloading credentials and attempting token refresh...", + status.as_u16() + ), + ); + + // 先重新加载凭证文件（可能用户换了账户） + if let Err(e) = kiro.load_credentials().await { + state.logs.write().await.add( + "error", + &format!("[AUTH] Failed to reload credentials: {e}"), + ); + } + + match kiro.refresh_token().await { + Ok(_) => { + state.logs.write().await.add( + "info", + "[AUTH] Token refreshed successfully, retrying request...", + ); + drop(kiro); + let kiro = state.kiro.read().await; + match kiro.call_api(&openai_request).await { + Ok(retry_resp) => { + let retry_status = retry_resp.status(); + state.logs.write().await.add( + "info", + &format!("[RETRY] Response status: {retry_status}"), + ); + if retry_resp.status().is_success() { + match retry_resp.bytes().await { + Ok(bytes) => { + let body = String::from_utf8_lossy(&bytes).to_string(); + let parsed = parse_cw_response(&body); + state.logs.write().await.add( + "info", + &format!( + "[RETRY] Success: content_len={}, tool_calls={}", + parsed.content.len(), parsed.tool_calls.len() + ), + ); + if request.stream { + return build_anthropic_stream_response( + &request.model, + &parsed, + ); + } + return build_anthropic_response( + &request.model, + &parsed, + ); + } + Err(e) => { + state.logs.write().await.add( + "error", + &format!("[RETRY] Body read failed: {e}"), + ); + return ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"error": {"message": e.to_string()}})), + ) + .into_response(); + } + } + } + let body = retry_resp + .bytes() + .await + .map(|b| String::from_utf8_lossy(&b).to_string()) + .unwrap_or_default(); + state.logs.write().await.add( + "error", + &format!( + "[RETRY] Failed with status {retry_status}: {}", + safe_truncate(&body, 500) + ), + ); + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"error": {"message": format!("Retry failed: {}", body)}})), + ) + .into_response() + } + Err(e) => { + state + .logs + .write() + .await + .add("error", &format!("[RETRY] Request failed: {e}")); + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"error": {"message": e.to_string()}})), + ) + .into_response() + } + } + } + Err(e) => { + state + .logs + .write() + .await + .add("error", &format!("[AUTH] Token refresh failed: {e}")); + ( + StatusCode::UNAUTHORIZED, + Json(serde_json::json!({"error": {"message": format!("Token refresh failed: {e}")}})), + ) + .into_response() + } + } + } else { + let body = resp.text().await.unwrap_or_default(); + state.logs.write().await.add( + "error", + &format!( + "[ERROR] Upstream error HTTP {}: {}", + status, + safe_truncate(&body, 500) + ), + ); + ( + StatusCode::from_u16(status.as_u16()) + .unwrap_or(StatusCode::INTERNAL_SERVER_ERROR), + Json( + serde_json::json!({"error": {"message": format!("Upstream error: {}", body)}}), + ), + ) + .into_response() + } + } + Err(e) => { + // 详细记录网络/连接错误 + let error_details = format!("{e:?}"); + state + .logs + .write() + .await + .add("error", &format!("[ERROR] Kiro API call failed: {e}")); + state.logs.write().await.add( + "debug", + &format!("[ERROR] Full error details: {error_details}"), + ); + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"error": {"message": e.to_string()}})), + ) + .into_response() + } + } +} diff --git a/src-tauri/src/server/handlers/management.rs b/src-tauri/src/server/handlers/management.rs new file mode 100644 index 000000000..2ca09fb80 --- /dev/null +++ b/src-tauri/src/server/handlers/management.rs @@ -0,0 +1,557 @@ +//! Management API 处理器 +//! +//! 提供服务器状态查询、凭证管理、配置管理等功能 + +use axum::{extract::State, http::StatusCode, response::IntoResponse, Json}; +use serde::{Deserialize, Serialize}; + +use crate::database::dao::provider_pool::ProviderPoolDao; +use crate::server::AppState; + +// ============ Types ============ + +/// 管理 API 状态响应 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ManagementStatusResponse { + /// 服务器是否运行中 + pub running: bool, + /// 监听地址 + pub host: String, + /// 监听端口 + pub port: u16, + /// 处理的请求数 + pub requests: u64, + /// 运行时间（秒） + pub uptime_secs: u64, + /// 版本号 + pub version: String, + /// TLS 是否启用 + pub tls_enabled: bool, + /// 默认 Provider + pub default_provider: String, +} + +/// 凭证信息（用于列表显示） +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct CredentialInfo { + /// 凭证 ID + pub id: String, + /// Provider 类型 + pub provider_type: String, + /// 是否禁用 + pub disabled: bool, + /// 是否有效 + pub is_valid: bool, +} + +/// 凭证列表响应 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct CredentialsListResponse { + /// 凭证列表 + pub credentials: Vec, + /// 总数 + pub total: usize, +} + +/// 添加凭证请求 +#[derive(Debug, Clone, Deserialize)] +pub struct AddCredentialRequest { + /// Provider 类型 + pub provider_type: String, + /// 凭证 ID + pub id: String, + /// API Key（用于 API Key 类型的凭证） + #[serde(default)] + pub api_key: Option, + /// Token 文件路径（用于 OAuth 类型的凭证） + #[serde(default)] + pub token_file: Option, + /// Base URL + #[serde(default)] + pub base_url: Option, + /// 代理 URL + #[serde(default)] + pub proxy_url: Option, +} + +/// 添加凭证响应 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct AddCredentialResponse { + /// 是否成功 + pub success: bool, + /// 消息 + pub message: String, + /// 凭证 ID + pub id: Option, +} + +/// 配置响应（简化版，不包含敏感信息） +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ManagementConfigResponse { + /// 服务器配置 + pub server: ManagementServerConfigInfo, + /// 路由配置 + pub routing: ManagementRoutingConfigInfo, + /// 重试配置 + pub retry: ManagementRetryConfigInfo, + /// 远程管理配置（不包含 secret_key） + pub remote_management: ManagementRemoteInfo, +} + +/// 服务器配置信息 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ManagementServerConfigInfo { + pub host: String, + pub port: u16, + pub tls_enabled: bool, +} + +/// 路由配置信息 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ManagementRoutingConfigInfo { + pub default_provider: String, + pub rules_count: usize, +} + +/// 重试配置信息 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ManagementRetryConfigInfo { + pub max_retries: u32, + pub base_delay_ms: u64, + pub max_delay_ms: u64, +} + +/// 远程管理配置信息 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ManagementRemoteInfo { + pub allow_remote: bool, + pub has_secret_key: bool, + pub disable_control_panel: bool, +} + +/// 更新配置请求 +#[derive(Debug, Clone, Deserialize)] +pub struct UpdateConfigRequest { + /// 默认 Provider + #[serde(default)] + pub default_provider: Option, + /// 是否允许远程访问 + #[serde(default)] + pub allow_remote: Option, +} + +/// 更新配置响应 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct UpdateConfigResponse { + pub success: bool, + pub message: String, +} + +// ============ Handlers ============ + +/// GET /v0/management/status - 获取服务器状态 +pub async fn management_status(State(state): State) -> impl IntoResponse { + let default_provider = state.default_provider.read().await.clone(); + + // 获取请求数量 + let requests = state.processor.stats.read().len() as u64; + + let response = ManagementStatusResponse { + running: true, + host: "0.0.0.0".to_string(), + port: 8999, + requests, + uptime_secs: 0, // TODO: Track actual uptime + version: env!("CARGO_PKG_VERSION").to_string(), + tls_enabled: false, + default_provider, + }; + + Json(response) +} + +/// GET /v0/management/credentials - 获取凭证列表 +pub async fn management_list_credentials(State(state): State) -> impl IntoResponse { + let mut credentials = Vec::new(); + + // 从数据库获取凭证列表 + if let Some(ref db) = state.db { + if let Ok(conn) = db.lock() { + if let Ok(pool_credentials) = ProviderPoolDao::get_all(&conn) { + for cred in pool_credentials { + credentials.push(CredentialInfo { + id: cred.uuid.clone(), + provider_type: cred.provider_type.to_string(), + disabled: cred.is_disabled, + is_valid: cred.is_healthy, + }); + } + } + } + } + + let total = credentials.len(); + Json(CredentialsListResponse { credentials, total }) +} + +/// POST /v0/management/credentials - 添加凭证 +pub async fn management_add_credential( + State(state): State, + Json(request): Json, +) -> impl IntoResponse { + use crate::models::provider_pool_model::{ + CredentialData, PoolProviderType, ProviderCredential, + }; + + // 验证请求 + if request.id.is_empty() { + return ( + StatusCode::BAD_REQUEST, + Json(AddCredentialResponse { + success: false, + message: "Credential ID is required".to_string(), + id: None, + }), + ); + } + + if request.provider_type.is_empty() { + return ( + StatusCode::BAD_REQUEST, + Json(AddCredentialResponse { + success: false, + message: "Provider type is required".to_string(), + id: None, + }), + ); + } + + // 解析 provider 类型 + let provider_type: PoolProviderType = match request.provider_type.parse() { + Ok(pt) => pt, + Err(_) => { + return ( + StatusCode::BAD_REQUEST, + Json(AddCredentialResponse { + success: false, + message: format!("Invalid provider type: {}", request.provider_type), + id: None, + }), + ); + } + }; + + // 根据 provider 类型创建凭证数据 + let credential_data = match provider_type { + PoolProviderType::OpenAI => { + if let Some(api_key) = request.api_key { + CredentialData::OpenAIKey { + api_key, + base_url: request.base_url, + } + } else { + return ( + StatusCode::BAD_REQUEST, + Json(AddCredentialResponse { + success: false, + message: "API key is required for OpenAI provider".to_string(), + id: None, + }), + ); + } + } + PoolProviderType::Claude => { + if let Some(api_key) = request.api_key { + CredentialData::ClaudeKey { + api_key, + base_url: request.base_url, + } + } else { + return ( + StatusCode::BAD_REQUEST, + Json(AddCredentialResponse { + success: false, + message: "API key is required for Claude provider".to_string(), + id: None, + }), + ); + } + } + PoolProviderType::Vertex => { + if let Some(api_key) = request.api_key { + CredentialData::VertexKey { + api_key, + base_url: request.base_url, + model_aliases: std::collections::HashMap::new(), + } + } else { + return ( + StatusCode::BAD_REQUEST, + Json(AddCredentialResponse { + success: false, + message: "API key is required for Vertex provider".to_string(), + id: None, + }), + ); + } + } + PoolProviderType::Kiro => { + if let Some(token_file) = request.token_file { + CredentialData::KiroOAuth { + creds_file_path: token_file, + } + } else { + return ( + StatusCode::BAD_REQUEST, + Json(AddCredentialResponse { + success: false, + message: "Token file is required for Kiro provider".to_string(), + id: None, + }), + ); + } + } + PoolProviderType::Gemini => { + if let Some(token_file) = request.token_file { + CredentialData::GeminiOAuth { + creds_file_path: token_file, + project_id: None, + } + } else { + return ( + StatusCode::BAD_REQUEST, + Json(AddCredentialResponse { + success: false, + message: "Token file is required for Gemini provider".to_string(), + id: None, + }), + ); + } + } + PoolProviderType::Qwen => { + if let Some(token_file) = request.token_file { + CredentialData::QwenOAuth { + creds_file_path: token_file, + } + } else { + return ( + StatusCode::BAD_REQUEST, + Json(AddCredentialResponse { + success: false, + message: "Token file is required for Qwen provider".to_string(), + id: None, + }), + ); + } + } + PoolProviderType::Antigravity => { + if let Some(token_file) = request.token_file { + CredentialData::AntigravityOAuth { + creds_file_path: token_file, + project_id: None, + } + } else { + return ( + StatusCode::BAD_REQUEST, + Json(AddCredentialResponse { + success: false, + message: "Token file is required for Antigravity provider".to_string(), + id: None, + }), + ); + } + } + PoolProviderType::GeminiApiKey => { + if let Some(api_key) = request.api_key { + CredentialData::GeminiApiKey { + api_key, + base_url: request.base_url, + excluded_models: Vec::new(), + } + } else { + return ( + StatusCode::BAD_REQUEST, + Json(AddCredentialResponse { + success: false, + message: "API key is required for Gemini API Key provider".to_string(), + id: None, + }), + ); + } + } + PoolProviderType::Codex => { + if let Some(token_file) = request.token_file { + CredentialData::CodexOAuth { + creds_file_path: token_file, + } + } else { + return ( + StatusCode::BAD_REQUEST, + Json(AddCredentialResponse { + success: false, + message: "Token file is required for Codex provider".to_string(), + id: None, + }), + ); + } + } + PoolProviderType::ClaudeOAuth => { + if let Some(token_file) = request.token_file { + CredentialData::ClaudeOAuth { + creds_file_path: token_file, + } + } else { + return ( + StatusCode::BAD_REQUEST, + Json(AddCredentialResponse { + success: false, + message: "Token file is required for Claude OAuth provider".to_string(), + id: None, + }), + ); + } + } + PoolProviderType::IFlow => { + if let Some(token_file) = request.token_file { + // 默认使用 OAuth 类型，Cookie 类型需要通过其他方式添加 + CredentialData::IFlowOAuth { + creds_file_path: token_file, + } + } else { + return ( + StatusCode::BAD_REQUEST, + Json(AddCredentialResponse { + success: false, + message: "Token file is required for iFlow provider".to_string(), + id: None, + }), + ); + } + } + }; + + // 创建凭证 + let mut credential = ProviderCredential::new(provider_type, credential_data); + credential.uuid = request.id.clone(); + credential.name = Some(request.id.clone()); + + // 添加凭证到数据库 + if let Some(ref db) = state.db { + if let Ok(conn) = db.lock() { + match ProviderPoolDao::insert(&conn, &credential) { + Ok(_) => { + tracing::info!( + "[MANAGEMENT] Added credential: {} ({})", + request.id, + request.provider_type + ); + return ( + StatusCode::CREATED, + Json(AddCredentialResponse { + success: true, + message: "Credential added successfully".to_string(), + id: Some(request.id), + }), + ); + } + Err(e) => { + tracing::error!("[MANAGEMENT] Failed to add credential: {}", e); + return ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(AddCredentialResponse { + success: false, + message: format!("Failed to add credential: {}", e), + id: None, + }), + ); + } + } + } + } + + ( + StatusCode::SERVICE_UNAVAILABLE, + Json(AddCredentialResponse { + success: false, + message: "Database not available".to_string(), + id: None, + }), + ) +} + +/// GET /v0/management/config - 获取配置 +pub async fn management_get_config(State(state): State) -> impl IntoResponse { + let default_provider = state.default_provider.read().await.clone(); + + // 获取路由规则数量 + let rules_count = state.processor.router.read().await.rules().len(); + + let response = ManagementConfigResponse { + server: ManagementServerConfigInfo { + host: "0.0.0.0".to_string(), + port: 8999, + tls_enabled: false, + }, + routing: ManagementRoutingConfigInfo { + default_provider, + rules_count, + }, + retry: ManagementRetryConfigInfo { + max_retries: 3, + base_delay_ms: 1000, + max_delay_ms: 30000, + }, + remote_management: ManagementRemoteInfo { + allow_remote: false, + has_secret_key: true, + disable_control_panel: false, + }, + }; + + Json(response) +} + +/// PUT /v0/management/config - 更新配置 +pub async fn management_update_config( + State(state): State, + Json(request): Json, +) -> impl IntoResponse { + let mut updated = false; + + // 更新默认 Provider + if let Some(provider) = request.default_provider { + // 验证 provider 类型 + if provider.parse::().is_ok() { + let mut dp = state.default_provider.write().await; + *dp = provider.clone(); + tracing::info!("[MANAGEMENT] Updated default_provider to: {}", provider); + updated = true; + } else { + return ( + StatusCode::BAD_REQUEST, + Json(UpdateConfigResponse { + success: false, + message: format!("Invalid provider type: {}", provider), + }), + ); + } + } + + if updated { + ( + StatusCode::OK, + Json(UpdateConfigResponse { + success: true, + message: "Configuration updated successfully".to_string(), + }), + ) + } else { + ( + StatusCode::OK, + Json(UpdateConfigResponse { + success: true, + message: "No changes applied".to_string(), + }), + ) + } +} diff --git a/src-tauri/src/server/handlers/mod.rs b/src-tauri/src/server/handlers/mod.rs new file mode 100644 index 000000000..4279943d5 --- /dev/null +++ b/src-tauri/src/server/handlers/mod.rs @@ -0,0 +1,13 @@ +//! HTTP 请求处理器模块 +//! +//! 将 server 中的各类处理器拆分到独立文件 + +pub mod api; +pub mod management; +pub mod provider_calls; +pub mod websocket; + +pub use api::*; +pub use management::*; +pub use provider_calls::*; +pub use websocket::*; diff --git a/src-tauri/src/server/handlers/provider_calls.rs b/src-tauri/src/server/handlers/provider_calls.rs new file mode 100644 index 000000000..b1483abe2 --- /dev/null +++ b/src-tauri/src/server/handlers/provider_calls.rs @@ -0,0 +1,925 @@ +//! Provider 调用处理器 +//! +//! 根据凭证类型调用不同的 Provider API + +use axum::{ + body::Body, + http::{header, StatusCode}, + response::{IntoResponse, Response}, + Json, +}; +use futures::stream; + +use crate::converter::anthropic_to_openai::convert_anthropic_to_openai; +use crate::converter::openai_to_antigravity::{ + convert_antigravity_to_openai_response, convert_openai_to_antigravity_with_context, +}; +use crate::models::anthropic::AnthropicMessagesRequest; +use crate::models::openai::ChatCompletionRequest; +use crate::models::provider_pool_model::{CredentialData, ProviderCredential}; +use crate::providers::{ + AntigravityProvider, ClaudeCustomProvider, GeminiProvider, KiroProvider, OpenAICustomProvider, + QwenProvider, VertexProvider, +}; +use crate::server::AppState; +use crate::server_utils::{ + build_anthropic_response, build_anthropic_stream_response, parse_cw_response, safe_truncate, + CWParsedResponse, +}; +/// 根据凭证调用 Provider (Anthropic 格式) +pub async fn call_provider_anthropic( + state: &AppState, + credential: &ProviderCredential, + request: &AnthropicMessagesRequest, +) -> Response { + match &credential.credential { + CredentialData::KiroOAuth { creds_file_path } => { + // 使用 TokenCacheService 获取有效 token + let db = match &state.db { + Some(db) => db, + None => { + return ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"error": {"message": "Database not available"}})), + ) + .into_response(); + } + }; + // 获取缓存的 token + let token = match state + .token_cache + .get_valid_token(db, &credential.uuid) + .await + { + Ok(t) => t, + Err(e) => { + tracing::warn!("[POOL] Token cache miss, loading from source: {}", e); + // 回退到从源文件加载 + let mut kiro = KiroProvider::new(); + if let Err(e) = kiro.load_credentials_from_path(creds_file_path).await { + // 记录凭证加载失败 + let _ = state.pool_service.mark_unhealthy( + db, + &credential.uuid, + Some(&format!("Failed to load credentials: {}", e)), + ); + return ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"error": {"message": format!("Failed to load Kiro credentials: {}", e)}})), + ) + .into_response(); + } + if let Err(e) = kiro.refresh_token().await { + // 记录 Token 刷新失败 + let _ = state.pool_service.mark_unhealthy( + db, + &credential.uuid, + Some(&format!("Token refresh failed: {}", e)), + ); + return ( + StatusCode::UNAUTHORIZED, + Json(serde_json::json!({"error": {"message": format!("Token refresh failed: {}", e)}})), + ) + .into_response(); + } + kiro.credentials.access_token.unwrap_or_default() + } + }; + // 使用获取到的 token 创建 KiroProvider + let mut kiro = KiroProvider::new(); + kiro.credentials.access_token = Some(token); + // 从源文件加载其他配置（region, profile_arn 等） + let _ = kiro.load_credentials_from_path(creds_file_path).await; + let openai_request = convert_anthropic_to_openai(request); + let resp = match kiro.call_api(&openai_request).await { + Ok(r) => r, + Err(e) => { + // 记录 API 调用失败 + let _ = state.pool_service.mark_unhealthy( + db, + &credential.uuid, + Some(&e.to_string()), + ); + return ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"error": {"message": e.to_string()}})), + ) + .into_response(); + } + }; + let status = resp.status(); + if status.is_success() { + match resp.bytes().await { + Ok(bytes) => { + let body = String::from_utf8_lossy(&bytes).to_string(); + let parsed = parse_cw_response(&body); + // 记录成功 + let _ = state.pool_service.mark_healthy( + db, + &credential.uuid, + Some(&request.model), + ); + let _ = state.pool_service.record_usage(db, &credential.uuid); + if request.stream { + build_anthropic_stream_response(&request.model, &parsed) + } else { + build_anthropic_response(&request.model, &parsed) + } + } + Err(e) => { + let _ = state.pool_service.mark_unhealthy( + db, + &credential.uuid, + Some(&e.to_string()), + ); + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"error": {"message": e.to_string()}})), + ) + .into_response() + } + } + } else if status.as_u16() == 401 || status.as_u16() == 403 { + // Token 过期，强制刷新并重试 + tracing::info!( + "[POOL] Got {}, forcing token refresh for {}", + status, + &credential.uuid[..8] + ); + let new_token = match state + .token_cache + .refresh_and_cache(db, &credential.uuid, true) + .await + { + Ok(t) => t, + Err(e) => { + // 记录 Token 刷新失败 + let _ = state.pool_service.mark_unhealthy( + db, + &credential.uuid, + Some(&format!("Token refresh failed: {}", e)), + ); + return ( + StatusCode::UNAUTHORIZED, + Json(serde_json::json!({"error": {"message": format!("Token refresh failed: {}", e)}})), + ) + .into_response(); + } + }; + // 使用新 token 重试 + kiro.credentials.access_token = Some(new_token); + match kiro.call_api(&openai_request).await { + Ok(retry_resp) => { + if retry_resp.status().is_success() { + match retry_resp.bytes().await { + Ok(bytes) => { + let body = String::from_utf8_lossy(&bytes).to_string(); + let parsed = parse_cw_response(&body); + // 记录重试成功 + let _ = state.pool_service.mark_healthy( + db, + &credential.uuid, + Some(&request.model), + ); + let _ = state.pool_service.record_usage(db, &credential.uuid); + if request.stream { + build_anthropic_stream_response(&request.model, &parsed) + } else { + build_anthropic_response(&request.model, &parsed) + } + } + Err(e) => { + let _ = state.pool_service.mark_unhealthy( + db, + &credential.uuid, + Some(&e.to_string()), + ); + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"error": {"message": e.to_string()}})), + ) + .into_response() + } + } + } else { + let body = retry_resp.text().await.unwrap_or_default(); + let _ = state.pool_service.mark_unhealthy( + db, + &credential.uuid, + Some(&format!("Retry failed: {}", body)), + ); + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"error": {"message": format!("Retry failed: {}", body)}})), + ) + .into_response() + } + } + Err(e) => { + let _ = state.pool_service.mark_unhealthy( + db, + &credential.uuid, + Some(&e.to_string()), + ); + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"error": {"message": e.to_string()}})), + ) + .into_response() + } + } + } else { + let body = resp.text().await.unwrap_or_default(); + let _ = state + .pool_service + .mark_unhealthy(db, &credential.uuid, Some(&body)); + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"error": {"message": body}})), + ) + .into_response() + } + } + CredentialData::GeminiOAuth { .. } => { + // Gemini OAuth 路由暂不支持 + ( + StatusCode::NOT_IMPLEMENTED, + Json(serde_json::json!({"error": {"message": "Gemini OAuth routing not yet implemented. Use /v1/messages with Gemini models instead."}})), + ) + .into_response() + } + CredentialData::QwenOAuth { .. } => { + // Qwen OAuth 路由暂不支持 + ( + StatusCode::NOT_IMPLEMENTED, + Json(serde_json::json!({"error": {"message": "Qwen OAuth routing not yet implemented. Use /v1/messages with Qwen models instead."}})), + ) + .into_response() + } + CredentialData::AntigravityOAuth { + creds_file_path, + project_id, + } => { + let mut antigravity = AntigravityProvider::new(); + if let Err(e) = antigravity + .load_credentials_from_path(creds_file_path) + .await + { + // 记录凭证加载失败 + if let Some(db) = &state.db { + let _ = state.pool_service.mark_unhealthy( + db, + &credential.uuid, + Some(&format!("Failed to load credentials: {}", e)), + ); + } + return ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"error": {"message": format!("Failed to load Antigravity credentials: {}", e)}})), + ) + .into_response(); + } + // 检查并刷新 token + if antigravity.is_token_expiring_soon() { + if let Err(e) = antigravity.refresh_token().await { + // 记录 Token 刷新失败 + if let Some(db) = &state.db { + let _ = state.pool_service.mark_unhealthy( + db, + &credential.uuid, + Some(&format!("Token refresh failed: {}", e)), + ); + } + return ( + StatusCode::UNAUTHORIZED, + Json(serde_json::json!({"error": {"message": format!("Token refresh failed: {}", e)}})), + ) + .into_response(); + } + } + // 设置项目 ID + if let Some(pid) = project_id { + antigravity.project_id = Some(pid.clone()); + } else if let Err(e) = antigravity.discover_project().await { + tracing::warn!("[Antigravity] Failed to discover project: {}", e); + } + // 获取 project_id 用于请求 + let proj_id = antigravity.project_id.clone().unwrap_or_default(); + // 先转换为 OpenAI 格式，再转换为 Antigravity 格式 + let openai_request = convert_anthropic_to_openai(request); + let antigravity_request = convert_openai_to_antigravity_with_context(&openai_request, &proj_id); + match antigravity + .generate_content(&request.model, &antigravity_request) + .await + { + Ok(resp) => { + // 转换为 OpenAI 格式，再构建 Anthropic 响应 + let content = resp["candidates"][0]["content"]["parts"][0]["text"] + .as_str() + .unwrap_or(""); + let parsed = CWParsedResponse { + content: content.to_string(), + tool_calls: Vec::new(), + usage_credits: 0.0, + context_usage_percentage: 0.0, + }; + // 记录成功 + if let Some(db) = &state.db { + let _ = state.pool_service.mark_healthy( + db, + &credential.uuid, + Some(&request.model), + ); + let _ = state.pool_service.record_usage(db, &credential.uuid); + } + if request.stream { + build_anthropic_stream_response(&request.model, &parsed) + } else { + build_anthropic_response(&request.model, &parsed) + } + } + Err(e) => { + // 记录 API 调用失败 + if let Some(db) = &state.db { + let _ = state.pool_service.mark_unhealthy( + db, + &credential.uuid, + Some(&e.to_string()), + ); + } + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"error": {"message": e.to_string()}})), + ) + .into_response() + } + } + } + CredentialData::OpenAIKey { api_key, base_url } => { + let openai = OpenAICustomProvider::with_config(api_key.clone(), base_url.clone()); + let openai_request = convert_anthropic_to_openai(request); + match openai.call_api(&openai_request).await { + Ok(resp) => { + if resp.status().is_success() { + match resp.text().await { + Ok(body) => { + if let Ok(openai_resp) = + serde_json::from_str::(&body) + { + let content = openai_resp["choices"][0]["message"]["content"] + .as_str() + .unwrap_or(""); + let parsed = CWParsedResponse { + content: content.to_string(), + tool_calls: Vec::new(), + usage_credits: 0.0, + context_usage_percentage: 0.0, + }; + // 记录成功 + if let Some(db) = &state.db { + let _ = state.pool_service.mark_healthy( + db, + &credential.uuid, + Some(&request.model), + ); + let _ = + state.pool_service.record_usage(db, &credential.uuid); + } + if request.stream { + build_anthropic_stream_response(&request.model, &parsed) + } else { + build_anthropic_response(&request.model, &parsed) + } + } else { + // 记录解析失败 + if let Some(db) = &state.db { + let _ = state.pool_service.mark_unhealthy( + db, + &credential.uuid, + Some("Failed to parse OpenAI response"), + ); + } + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"error": {"message": "Failed to parse OpenAI response"}})), + ) + .into_response() + } + } + Err(e) => { + if let Some(db) = &state.db { + let _ = state.pool_service.mark_unhealthy( + db, + &credential.uuid, + Some(&e.to_string()), + ); + } + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"error": {"message": e.to_string()}})), + ) + .into_response() + } + } + } else { + let body = resp.text().await.unwrap_or_default(); + if let Some(db) = &state.db { + let _ = state.pool_service.mark_unhealthy( + db, + &credential.uuid, + Some(&body), + ); + } + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"error": {"message": body}})), + ) + .into_response() + } + } + Err(e) => { + if let Some(db) = &state.db { + let _ = state.pool_service.mark_unhealthy( + db, + &credential.uuid, + Some(&e.to_string()), + ); + } + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"error": {"message": e.to_string()}})), + ) + .into_response() + } + } + } + CredentialData::ClaudeKey { api_key, base_url } => { + // 打印 Claude 代理 URL 用于调试 + let actual_base_url = base_url.as_deref().unwrap_or("https://api.anthropic.com"); + let claude = ClaudeCustomProvider::with_config(api_key.clone(), base_url.clone()); + let request_url = claude.get_base_url(); + state.logs.write().await.add( + "info", + &format!( + "[CLAUDE] 使用 Claude API 代理: base_url={} -> {}/v1/messages credential_uuid={}", + actual_base_url, + request_url, + &credential.uuid[..8] + ), + ); + // 打印请求参数 + let request_json = serde_json::to_string(request).unwrap_or_default(); + state.logs.write().await.add( + "debug", + &format!( + "[CLAUDE] 请求参数: {}", + &request_json.chars().take(500).collect::() + ), + ); + match claude.call_api(request).await { + Ok(resp) => { + let status = resp.status(); + // 打印响应状态 + state.logs.write().await.add( + "info", + &format!( + "[CLAUDE] 响应状态: status={} model={}", + status, + request.model + ), + ); + match resp.text().await { + Ok(body) => { + if status.is_success() { + // 打印响应内容预览 + state.logs.write().await.add( + "debug", + &format!( + "[CLAUDE] 响应内容: {}", + &body.chars().take(500).collect::() + ), + ); + // 记录成功 + if let Some(db) = &state.db { + let _ = state.pool_service.mark_healthy( + db, + &credential.uuid, + Some(&request.model), + ); + let _ = state.pool_service.record_usage(db, &credential.uuid); + } + Response::builder() + .status(StatusCode::OK) + .header(header::CONTENT_TYPE, "application/json") + .body(Body::from(body)) + .unwrap_or_else(|_| { + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"error": {"message": "Failed to build response"}})), + ) + .into_response() + }) + } else { + state.logs.write().await.add( + "error", + &format!( + "[CLAUDE] 请求失败: status={} body={}", + status, + &body.chars().take(200).collect::() + ), + ); + if let Some(db) = &state.db { + let _ = state.pool_service.mark_unhealthy( + db, + &credential.uuid, + Some(&body), + ); + } + ( + StatusCode::from_u16(status.as_u16()) + .unwrap_or(StatusCode::INTERNAL_SERVER_ERROR), + Json(serde_json::json!({"error": {"message": body}})), + ) + .into_response() + } + } + Err(e) => { + state.logs.write().await.add( + "error", + &format!("[CLAUDE] 读取响应失败: {}", e), + ); + if let Some(db) = &state.db { + let _ = state.pool_service.mark_unhealthy( + db, + &credential.uuid, + Some(&e.to_string()), + ); + } + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"error": {"message": e.to_string()}})), + ) + .into_response() + } + } + } + Err(e) => { + if let Some(db) = &state.db { + let _ = state.pool_service.mark_unhealthy( + db, + &credential.uuid, + Some(&e.to_string()), + ); + } + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"error": {"message": e.to_string()}})), + ) + .into_response() + } + } + } + CredentialData::VertexKey { api_key, base_url, .. } => { + // Vertex AI uses Gemini-compatible API, convert Anthropic to OpenAI format first + let openai_request = convert_anthropic_to_openai(request); + let vertex = VertexProvider::with_config(api_key.clone(), base_url.clone()); + match vertex.chat_completions(&serde_json::to_value(&openai_request).unwrap_or_default()).await { + Ok(resp) => { + let status = resp.status(); + match resp.text().await { + Ok(body) => { + if status.is_success() { + if let Some(db) = &state.db { + let _ = state.pool_service.mark_healthy(db, &credential.uuid, Some(&request.model)); + let _ = state.pool_service.record_usage(db, &credential.uuid); + } + Response::builder() + .status(StatusCode::OK) + .header(header::CONTENT_TYPE, "application/json") + .body(Body::from(body)) + .unwrap_or_else(|_| { + (StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({"error": {"message": "Failed to build response"}}))).into_response() + }) + } else { + if let Some(db) = &state.db { + let _ = state.pool_service.mark_unhealthy(db, &credential.uuid, Some(&body)); + } + (StatusCode::from_u16(status.as_u16()).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR), Json(serde_json::json!({"error": {"message": body}}))).into_response() + } + } + Err(e) => { + if let Some(db) = &state.db { + let _ = state.pool_service.mark_unhealthy(db, &credential.uuid, Some(&e.to_string())); + } + (StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({"error": {"message": e.to_string()}}))).into_response() + } + } + } + Err(e) => { + if let Some(db) = &state.db { + let _ = state.pool_service.mark_unhealthy(db, &credential.uuid, Some(&e.to_string())); + } + (StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({"error": {"message": e.to_string()}}))).into_response() + } + } + } + // Gemini API Key credentials - not supported for Anthropic format + CredentialData::GeminiApiKey { .. } => { + ( + StatusCode::BAD_REQUEST, + Json(serde_json::json!({"error": {"message": "Gemini API Key credentials do not support Anthropic format"}})), + ) + .into_response() + } + // 新增的凭证类型暂不支持 Anthropic 格式 + CredentialData::CodexOAuth { .. } + | CredentialData::ClaudeOAuth { .. } + | CredentialData::IFlowOAuth { .. } + | CredentialData::IFlowCookie { .. } => { + ( + StatusCode::BAD_REQUEST, + Json(serde_json::json!({"error": {"message": "This credential type does not support Anthropic format yet"}})), + ) + .into_response() + } + } +} +/// 根据凭证调用 Provider (OpenAI 格式) +pub async fn call_provider_openai( + state: &AppState, + credential: &ProviderCredential, + request: &ChatCompletionRequest, +) -> Response { + let start_time = std::time::Instant::now(); + match &credential.credential { + CredentialData::KiroOAuth { creds_file_path } => { + let mut kiro = KiroProvider::new(); + if let Err(e) = kiro.load_credentials_from_path(creds_file_path).await { + // 记录凭证加载失败 + if let Some(db) = &state.db { + let _ = state.pool_service.mark_unhealthy(db, &credential.uuid, Some(&e.to_string())); + } + return ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"error": {"message": format!("Failed to load Kiro credentials: {}", e)}})), + ) + .into_response(); + } + if let Err(e) = kiro.refresh_token().await { + // 记录 Token 刷新失败 + if let Some(db) = &state.db { + let _ = state.pool_service.mark_unhealthy(db, &credential.uuid, Some(&format!("Token refresh failed: {}", e))); + } + return ( + StatusCode::UNAUTHORIZED, + Json(serde_json::json!({"error": {"message": format!("Token refresh failed: {}", e)}})), + ) + .into_response(); + } + match kiro.call_api(request).await { + Ok(resp) => { + let status = resp.status(); + if status.is_success() { + // 记录成功 + if let Some(db) = &state.db { + let _ = state.pool_service.mark_healthy(db, &credential.uuid, Some(&request.model)); + let _ = state.pool_service.record_usage(db, &credential.uuid); + } + match resp.text().await { + Ok(body) => { + let parsed = parse_cw_response(&body); + let has_tool_calls = !parsed.tool_calls.is_empty(); + let message = if has_tool_calls { + serde_json::json!({ + "role": "assistant", + "content": if parsed.content.is_empty() { serde_json::Value::Null } else { serde_json::json!(parsed.content) }, + "tool_calls": parsed.tool_calls.iter().map(|tc| { + serde_json::json!({ + "id": tc.id, + "type": "function", + "function": { + "name": tc.function.name, + "arguments": tc.function.arguments + } + }) + }).collect::>() + }) + } else { + serde_json::json!({ + "role": "assistant", + "content": parsed.content + }) + }; + Json(serde_json::json!({ + "id": format!("chatcmpl-{}", uuid::Uuid::new_v4()), + "object": "chat.completion", + "created": std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_secs(), + "model": request.model, + "choices": [{ + "index": 0, + "message": message, + "finish_reason": if has_tool_calls { "tool_calls" } else { "stop" } + }], + "usage": { + "prompt_tokens": 0, + "completion_tokens": 0, + "total_tokens": 0 + } + })) + .into_response() + } + Err(e) => ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"error": {"message": e.to_string()}})), + ) + .into_response(), + } + } else { + // 记录 API 调用失败 + let body = resp.text().await.unwrap_or_default(); + if let Some(db) = &state.db { + let _ = state.pool_service.mark_unhealthy(db, &credential.uuid, Some(&format!("HTTP {}: {}", status, safe_truncate(&body, 100)))); + } + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"error": {"message": body}})), + ) + .into_response() + } + } + Err(e) => { + // 记录请求错误 + if let Some(db) = &state.db { + let _ = state.pool_service.mark_unhealthy(db, &credential.uuid, Some(&e.to_string())); + } + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"error": {"message": e.to_string()}})), + ) + .into_response() + } + } + } + CredentialData::GeminiOAuth { .. } => { + ( + StatusCode::NOT_IMPLEMENTED, + Json(serde_json::json!({"error": {"message": "Gemini OAuth routing not yet implemented."}})), + ) + .into_response() + } + CredentialData::QwenOAuth { .. } => { + ( + StatusCode::NOT_IMPLEMENTED, + Json(serde_json::json!({"error": {"message": "Qwen OAuth routing not yet implemented."}})), + ) + .into_response() + } + CredentialData::AntigravityOAuth { creds_file_path, project_id } => { + let mut antigravity = AntigravityProvider::new(); + if let Err(e) = antigravity.load_credentials_from_path(creds_file_path).await { + return ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"error": {"message": format!("Failed to load Antigravity credentials: {}", e)}})), + ) + .into_response(); + } + // 检查并刷新 token + if antigravity.is_token_expiring_soon() { + if let Err(e) = antigravity.refresh_token().await { + return ( + StatusCode::UNAUTHORIZED, + Json(serde_json::json!({"error": {"message": format!("Token refresh failed: {}", e)}})), + ) + .into_response(); + } + } + // 设置项目 ID + if let Some(pid) = project_id { + antigravity.project_id = Some(pid.clone()); + } else if let Err(e) = antigravity.discover_project().await { + tracing::warn!("[Antigravity] Failed to discover project: {}", e); + } + // 获取 project_id 用于请求 + let proj_id = antigravity.project_id.clone().unwrap_or_default(); + // 转换请求格式 + let antigravity_request = convert_openai_to_antigravity_with_context(request, &proj_id); + match antigravity.generate_content(&request.model, &antigravity_request).await { + Ok(resp) => { + let openai_response = convert_antigravity_to_openai_response(&resp, &request.model); + Json(openai_response).into_response() + } + Err(e) => ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"error": {"message": e.to_string()}})), + ) + .into_response(), + } + } + CredentialData::OpenAIKey { api_key, base_url } => { + let openai = OpenAICustomProvider::with_config(api_key.clone(), base_url.clone()); + match openai.call_api(request).await { + Ok(resp) => { + if resp.status().is_success() { + match resp.text().await { + Ok(body) => { + if let Ok(json) = serde_json::from_str::(&body) { + Json(json).into_response() + } else { + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"error": {"message": "Invalid JSON response"}})), + ) + .into_response() + } + } + Err(e) => ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"error": {"message": e.to_string()}})), + ) + .into_response(), + } + } else { + let body = resp.text().await.unwrap_or_default(); + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"error": {"message": body}})), + ) + .into_response() + } + } + Err(e) => ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"error": {"message": e.to_string()}})), + ) + .into_response(), + } + } + CredentialData::ClaudeKey { api_key, base_url } => { + // 打印 Claude 代理 URL 用于调试 + let actual_base_url = base_url.as_deref().unwrap_or("https://api.anthropic.com"); + tracing::info!( + "[CLAUDE] 使用 Claude API 代理: base_url={} credential_uuid={}", + actual_base_url, + &credential.uuid[..8] + ); + let claude = ClaudeCustomProvider::with_config(api_key.clone(), base_url.clone()); + match claude.call_openai_api(request).await { + Ok(resp) => Json(resp).into_response(), + Err(e) => ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"error": {"message": e.to_string()}})), + ) + .into_response(), + } + } + CredentialData::VertexKey { api_key, base_url, model_aliases } => { + // Resolve model alias if present + let resolved_model = model_aliases.get(&request.model).cloned().unwrap_or_else(|| request.model.clone()); + let mut modified_request = request.clone(); + modified_request.model = resolved_model; + let vertex = VertexProvider::with_config(api_key.clone(), base_url.clone()); + match vertex.chat_completions(&serde_json::to_value(&modified_request).unwrap_or_default()).await { + Ok(resp) => { + if resp.status().is_success() { + match resp.text().await { + Ok(body) => { + if let Ok(json) = serde_json::from_str::(&body) { + Json(json).into_response() + } else { + (StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({"error": {"message": "Invalid JSON response"}}))).into_response() + } + } + Err(e) => (StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({"error": {"message": e.to_string()}}))).into_response(), + } + } else { + let body = resp.text().await.unwrap_or_default(); + (StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({"error": {"message": body}}))).into_response() + } + } + Err(e) => (StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({"error": {"message": e.to_string()}}))).into_response(), + } + } + // Gemini API Key credentials - not supported for OpenAI format yet + CredentialData::GeminiApiKey { .. } => { + ( + StatusCode::BAD_REQUEST, + Json(serde_json::json!({"error": {"message": "Gemini API Key credentials do not support OpenAI format yet"}})), + ) + .into_response() + } + // 新增的凭证类型暂不支持 OpenAI 格式 + CredentialData::CodexOAuth { .. } + | CredentialData::ClaudeOAuth { .. } + | CredentialData::IFlowOAuth { .. } + | CredentialData::IFlowCookie { .. } => { + ( + StatusCode::BAD_REQUEST, + Json(serde_json::json!({"error": {"message": "This credential type does not support OpenAI format yet"}})), + ) + .into_response() + } + } +} diff --git a/src-tauri/src/server/handlers/websocket.rs b/src-tauri/src/server/handlers/websocket.rs new file mode 100644 index 000000000..beedf1876 --- /dev/null +++ b/src-tauri/src/server/handlers/websocket.rs @@ -0,0 +1,845 @@ +//! WebSocket 连接处理器 +//! +//! 处理 WebSocket 连接的建立、消息收发和 API 请求转发 + +use axum::{ + body::Body, + extract::{ + ws::{Message as WsMessage, WebSocket, WebSocketUpgrade}, + State, + }, + http::HeaderMap, + response::IntoResponse, +}; +use futures::{SinkExt, StreamExt as FuturesStreamExt}; + +use crate::converter::anthropic_to_openai::convert_anthropic_to_openai; +use crate::converter::openai_to_antigravity::{ + convert_antigravity_to_openai_response, convert_openai_to_antigravity_with_context, +}; +use crate::models::anthropic::AnthropicMessagesRequest; +use crate::models::openai::ChatCompletionRequest; +use crate::models::provider_pool_model::ProviderCredential; +use crate::processor::RequestContext; +use crate::providers::{ + AntigravityProvider, ClaudeCustomProvider, KiroProvider, OpenAICustomProvider, +}; +use crate::server::AppState; +use crate::server_utils::parse_cw_response; +use crate::websocket::{ + WsApiRequest, WsApiResponse, WsEndpoint, WsError, WsMessage as WsProtoMessage, +}; + +/// WebSocket 升级处理器 +pub async fn ws_upgrade_handler( + ws: WebSocketUpgrade, + State(state): State, + headers: HeaderMap, +) -> impl IntoResponse { + // 验证 API 密钥 + let auth = headers + .get("authorization") + .or_else(|| headers.get("x-api-key")) + .and_then(|v| v.to_str().ok()); + + let key = match auth { + Some(s) if s.starts_with("Bearer ") => &s[7..], + Some(s) => s, + None => { + return axum::http::Response::builder() + .status(401) + .body(Body::from("No API key provided")) + .unwrap() + .into_response(); + } + }; + + if key != state.api_key { + return axum::http::Response::builder() + .status(401) + .body(Body::from("Invalid API key")) + .unwrap() + .into_response(); + } + + // 获取客户端信息 + let client_info = headers + .get("user-agent") + .and_then(|v| v.to_str().ok()) + .map(|s| s.to_string()); + + ws.on_upgrade(move |socket| handle_websocket(socket, state, client_info)) +} + +/// 处理 WebSocket 连接 +pub async fn handle_websocket(socket: WebSocket, state: AppState, client_info: Option) { + let conn_id = uuid::Uuid::new_v4().to_string(); + + // 注册连接 + if let Err(e) = state + .ws_manager + .register(conn_id.clone(), client_info.clone()) + { + state.logs.write().await.add( + "error", + &format!("[WS] Failed to register connection: {}", e.message), + ); + return; + } + + state.logs.write().await.add( + "info", + &format!( + "[WS] New connection: {} (client: {:?})", + &conn_id[..8], + client_info + ), + ); + + let (mut sender, mut receiver) = socket.split(); + + // 消息处理循环 + while let Some(msg) = receiver.next().await { + match msg { + Ok(WsMessage::Text(text)) => { + state.ws_manager.on_message(); + state.ws_manager.increment_request_count(&conn_id); + + match serde_json::from_str::(&text) { + Ok(ws_msg) => { + let response = handle_ws_message(&state, &conn_id, ws_msg).await; + if let Some(resp) = response { + let resp_text = serde_json::to_string(&resp).unwrap_or_default(); + if sender + .send(WsMessage::Text(resp_text.into())) + .await + .is_err() + { + break; + } + } + } + Err(e) => { + state.ws_manager.on_error(); + let error = WsProtoMessage::Error(WsError::invalid_message(format!( + "Failed to parse message: {}", + e + ))); + let error_text = serde_json::to_string(&error).unwrap_or_default(); + if sender + .send(WsMessage::Text(error_text.into())) + .await + .is_err() + { + break; + } + } + } + } + Ok(WsMessage::Binary(_)) => { + state.ws_manager.on_error(); + let error = WsProtoMessage::Error(WsError::invalid_message( + "Binary messages not supported", + )); + let error_text = serde_json::to_string(&error).unwrap_or_default(); + if sender + .send(WsMessage::Text(error_text.into())) + .await + .is_err() + { + break; + } + } + Ok(WsMessage::Ping(data)) => { + if sender.send(WsMessage::Pong(data)).await.is_err() { + break; + } + } + Ok(WsMessage::Pong(_)) => { + // 收到 pong，连接正常 + } + Ok(WsMessage::Close(_)) => { + break; + } + Err(e) => { + state.logs.write().await.add( + "error", + &format!("[WS] Connection {} error: {}", &conn_id[..8], e), + ); + break; + } + } + } + + // 清理连接 + state.ws_manager.unregister(&conn_id); + state.logs.write().await.add( + "info", + &format!("[WS] Connection closed: {}", &conn_id[..8]), + ); +} + +/// 处理 WebSocket 消息 +async fn handle_ws_message( + state: &AppState, + conn_id: &str, + msg: WsProtoMessage, +) -> Option { + match msg { + WsProtoMessage::Ping { timestamp } => Some(WsProtoMessage::Pong { timestamp }), + WsProtoMessage::Pong { .. } => None, + WsProtoMessage::Request(request) => { + state.logs.write().await.add( + "info", + &format!( + "[WS] Request from {}: id={} endpoint={:?}", + &conn_id[..8], + request.request_id, + request.endpoint + ), + ); + + // 处理 API 请求 + let response = handle_ws_api_request(state, &request).await; + Some(response) + } + WsProtoMessage::Response(_) + | WsProtoMessage::StreamChunk(_) + | WsProtoMessage::StreamEnd(_) => Some(WsProtoMessage::Error(WsError::invalid_request( + None, + "Invalid message type from client", + ))), + WsProtoMessage::Error(_) => None, + } +} + +/// 处理 WebSocket API 请求 +async fn handle_ws_api_request(state: &AppState, request: &WsApiRequest) -> WsProtoMessage { + match request.endpoint { + WsEndpoint::Models => { + // 返回模型列表 + let models = serde_json::json!({ + "object": "list", + "data": [ + {"id": "claude-sonnet-4-5", "object": "model", "owned_by": "anthropic"}, + {"id": "claude-sonnet-4-5-20250929", "object": "model", "owned_by": "anthropic"}, + {"id": "claude-3-7-sonnet-20250219", "object": "model", "owned_by": "anthropic"}, + {"id": "gemini-2.5-flash", "object": "model", "owned_by": "google"}, + {"id": "gemini-2.5-pro", "object": "model", "owned_by": "google"}, + {"id": "qwen3-coder-plus", "object": "model", "owned_by": "alibaba"}, + ] + }); + WsProtoMessage::Response(WsApiResponse { + request_id: request.request_id.clone(), + payload: models, + }) + } + WsEndpoint::ChatCompletions => { + // 解析 ChatCompletionRequest + match serde_json::from_value::(request.payload.clone()) { + Ok(chat_request) => { + handle_ws_chat_completions(state, &request.request_id, chat_request).await + } + Err(e) => WsProtoMessage::Error(WsError::invalid_request( + Some(request.request_id.clone()), + format!("Invalid chat completion request: {}", e), + )), + } + } + WsEndpoint::Messages => { + // 解析 AnthropicMessagesRequest + match serde_json::from_value::(request.payload.clone()) { + Ok(messages_request) => { + handle_ws_anthropic_messages(state, &request.request_id, messages_request).await + } + Err(e) => WsProtoMessage::Error(WsError::invalid_request( + Some(request.request_id.clone()), + format!("Invalid messages request: {}", e), + )), + } + } + } +} + +/// 处理 WebSocket chat completions 请求 +async fn handle_ws_chat_completions( + state: &AppState, + request_id: &str, + mut request: ChatCompletionRequest, +) -> WsProtoMessage { + // 创建请求上下文 + let mut ctx = RequestContext::new(request.model.clone()).with_stream(request.stream); + + // 使用 RequestProcessor 解析模型别名和路由 + let _provider = state.processor.resolve_and_route(&mut ctx).await; + + // 更新请求中的模型名为解析后的模型 + if ctx.resolved_model != ctx.original_model { + request.model = ctx.resolved_model.clone(); + } + + // 应用参数注入 + let injection_enabled = *state.injection_enabled.read().await; + if injection_enabled { + let injector = state.processor.injector.read().await; + let mut payload = serde_json::to_value(&request).unwrap_or_default(); + let result = injector.inject(&request.model, &mut payload); + if result.has_injections() { + if let Ok(updated) = serde_json::from_value(payload) { + request = updated; + } + } + } + + // 获取默认 provider + let default_provider = state.default_provider.read().await.clone(); + + // 尝试从凭证池中选择凭证 + let credential = match &state.db { + Some(db) => state + .pool_service + .select_credential(db, &default_provider, Some(&request.model)) + .ok() + .flatten(), + None => None, + }; + + // 如果找到凭证，使用它调用 API + if let Some(cred) = credential { + // 简化实现：直接调用 provider 并返回结果 + // 实际实现应该复用 call_provider_openai 的逻辑 + match call_provider_openai_for_ws(state, &cred, &request).await { + Ok(response) => WsProtoMessage::Response(WsApiResponse { + request_id: request_id.to_string(), + payload: response, + }), + Err(e) => WsProtoMessage::Error(WsError::upstream(Some(request_id.to_string()), e)), + } + } else { + // 回退到 Kiro provider + let kiro = state.kiro.read().await; + match kiro.call_api(&request).await { + Ok(resp) => { + if resp.status().is_success() { + match resp.text().await { + Ok(body) => { + let parsed = parse_cw_response(&body); + let has_tool_calls = !parsed.tool_calls.is_empty(); + + let message = if has_tool_calls { + serde_json::json!({ + "role": "assistant", + "content": if parsed.content.is_empty() { serde_json::Value::Null } else { serde_json::json!(parsed.content) }, + "tool_calls": parsed.tool_calls.iter().map(|tc| { + serde_json::json!({ + "id": tc.id, + "type": "function", + "function": { + "name": tc.function.name, + "arguments": tc.function.arguments + } + }) + }).collect::>() + }) + } else { + serde_json::json!({ + "role": "assistant", + "content": parsed.content + }) + }; + + let response = serde_json::json!({ + "id": format!("chatcmpl-{}", uuid::Uuid::new_v4()), + "object": "chat.completion", + "created": std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_secs(), + "model": request.model, + "choices": [{ + "index": 0, + "message": message, + "finish_reason": if has_tool_calls { "tool_calls" } else { "stop" } + }], + "usage": { + "prompt_tokens": 0, + "completion_tokens": 0, + "total_tokens": 0 + } + }); + + WsProtoMessage::Response(WsApiResponse { + request_id: request_id.to_string(), + payload: response, + }) + } + Err(e) => WsProtoMessage::Error(WsError::internal( + Some(request_id.to_string()), + e.to_string(), + )), + } + } else { + let body = resp.text().await.unwrap_or_default(); + WsProtoMessage::Error(WsError::upstream( + Some(request_id.to_string()), + format!("Upstream error: {}", body), + )) + } + } + Err(e) => WsProtoMessage::Error(WsError::internal( + Some(request_id.to_string()), + e.to_string(), + )), + } + } +} + +/// 处理 WebSocket anthropic messages 请求 +async fn handle_ws_anthropic_messages( + state: &AppState, + request_id: &str, + mut request: AnthropicMessagesRequest, +) -> WsProtoMessage { + // 创建请求上下文 + let mut ctx = RequestContext::new(request.model.clone()).with_stream(request.stream); + + // 使用 RequestProcessor 解析模型别名和路由 + let _provider = state.processor.resolve_and_route(&mut ctx).await; + + // 更新请求中的模型名为解析后的模型 + if ctx.resolved_model != ctx.original_model { + request.model = ctx.resolved_model.clone(); + } + + // 应用参数注入 + let injection_enabled = *state.injection_enabled.read().await; + if injection_enabled { + let injector = state.processor.injector.read().await; + let mut payload = serde_json::to_value(&request).unwrap_or_default(); + let result = injector.inject(&request.model, &mut payload); + if result.has_injections() { + if let Ok(updated) = serde_json::from_value(payload) { + request = updated; + } + } + } + + // 获取默认 provider + let default_provider = state.default_provider.read().await.clone(); + + // 尝试从凭证池中选择凭证 + let credential = match &state.db { + Some(db) => state + .pool_service + .select_credential(db, &default_provider, Some(&request.model)) + .ok() + .flatten(), + None => None, + }; + + // 如果找到凭证，使用它调用 API + if let Some(cred) = credential { + match call_provider_anthropic_for_ws(state, &cred, &request).await { + Ok(response) => WsProtoMessage::Response(WsApiResponse { + request_id: request_id.to_string(), + payload: response, + }), + Err(e) => WsProtoMessage::Error(WsError::upstream(Some(request_id.to_string()), e)), + } + } else { + // 回退到 Kiro provider + let kiro = state.kiro.read().await; + + // 转换为 OpenAI 格式 + let openai_request = convert_anthropic_to_openai(&request); + + match kiro.call_api(&openai_request).await { + Ok(resp) => { + if resp.status().is_success() { + match resp.text().await { + Ok(body) => { + let parsed = parse_cw_response(&body); + + // 转换为 Anthropic 格式响应 + let response = serde_json::json!({ + "id": format!("msg_{}", uuid::Uuid::new_v4()), + "type": "message", + "role": "assistant", + "content": [{ + "type": "text", + "text": parsed.content + }], + "model": request.model, + "stop_reason": "end_turn", + "usage": { + "input_tokens": 0, + "output_tokens": 0 + } + }); + + WsProtoMessage::Response(WsApiResponse { + request_id: request_id.to_string(), + payload: response, + }) + } + Err(e) => WsProtoMessage::Error(WsError::internal( + Some(request_id.to_string()), + e.to_string(), + )), + } + } else { + let body = resp.text().await.unwrap_or_default(); + WsProtoMessage::Error(WsError::upstream( + Some(request_id.to_string()), + format!("Upstream error: {}", body), + )) + } + } + Err(e) => WsProtoMessage::Error(WsError::internal( + Some(request_id.to_string()), + e.to_string(), + )), + } + } +} + +/// WebSocket 专用的 OpenAI 格式 Provider 调用 +pub async fn call_provider_openai_for_ws( + state: &AppState, + credential: &ProviderCredential, + request: &ChatCompletionRequest, +) -> Result { + use crate::models::provider_pool_model::CredentialData; + + match &credential.credential { + CredentialData::KiroOAuth { creds_file_path } => { + let mut kiro = KiroProvider::new(); + if let Err(e) = kiro.load_credentials_from_path(creds_file_path).await { + if let Some(db) = &state.db { + let _ = state.pool_service.mark_unhealthy( + db, + &credential.uuid, + Some(&format!("Failed to load credentials: {}", e)), + ); + } + return Err(e.to_string()); + } + if let Err(e) = kiro.refresh_token().await { + if let Some(db) = &state.db { + let _ = state.pool_service.mark_unhealthy( + db, + &credential.uuid, + Some(&format!("Token refresh failed: {}", e)), + ); + } + return Err(e.to_string()); + } + + let resp = match kiro.call_api(request).await { + Ok(r) => r, + Err(e) => { + if let Some(db) = &state.db { + let _ = state.pool_service.mark_unhealthy( + db, + &credential.uuid, + Some(&e.to_string()), + ); + } + return Err(e.to_string()); + } + }; + if resp.status().is_success() { + let body = resp.text().await.map_err(|e| e.to_string())?; + let parsed = parse_cw_response(&body); + let has_tool_calls = !parsed.tool_calls.is_empty(); + + // 记录成功 + if let Some(db) = &state.db { + let _ = + state + .pool_service + .mark_healthy(db, &credential.uuid, Some(&request.model)); + let _ = state.pool_service.record_usage(db, &credential.uuid); + } + + let message = if has_tool_calls { + serde_json::json!({ + "role": "assistant", + "content": if parsed.content.is_empty() { serde_json::Value::Null } else { serde_json::json!(parsed.content) }, + "tool_calls": parsed.tool_calls.iter().map(|tc| { + serde_json::json!({ + "id": tc.id, + "type": "function", + "function": { + "name": tc.function.name, + "arguments": tc.function.arguments + } + }) + }).collect::>() + }) + } else { + serde_json::json!({ + "role": "assistant", + "content": parsed.content + }) + }; + + Ok(serde_json::json!({ + "id": format!("chatcmpl-{}", uuid::Uuid::new_v4()), + "object": "chat.completion", + "created": std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_secs(), + "model": request.model, + "choices": [{ + "index": 0, + "message": message, + "finish_reason": if has_tool_calls { "tool_calls" } else { "stop" } + }], + "usage": { + "prompt_tokens": 0, + "completion_tokens": 0, + "total_tokens": 0 + } + })) + } else { + let body = resp.text().await.unwrap_or_default(); + if let Some(db) = &state.db { + let _ = state + .pool_service + .mark_unhealthy(db, &credential.uuid, Some(&body)); + } + Err(format!("Upstream error: {}", body)) + } + } + CredentialData::OpenAIKey { api_key, base_url } => { + let provider = OpenAICustomProvider::with_config(api_key.clone(), base_url.clone()); + let resp = match provider.call_api(request).await { + Ok(r) => r, + Err(e) => { + if let Some(db) = &state.db { + let _ = state.pool_service.mark_unhealthy( + db, + &credential.uuid, + Some(&e.to_string()), + ); + } + return Err(e.to_string()); + } + }; + if resp.status().is_success() { + // 记录成功 + if let Some(db) = &state.db { + let _ = + state + .pool_service + .mark_healthy(db, &credential.uuid, Some(&request.model)); + let _ = state.pool_service.record_usage(db, &credential.uuid); + } + resp.json::() + .await + .map_err(|e| e.to_string()) + } else { + let body = resp.text().await.unwrap_or_default(); + if let Some(db) = &state.db { + let _ = state + .pool_service + .mark_unhealthy(db, &credential.uuid, Some(&body)); + } + Err(format!("Upstream error: {}", body)) + } + } + CredentialData::ClaudeKey { api_key, base_url } => { + // 打印 Claude 代理 URL 用于调试 + let actual_base_url = base_url.as_deref().unwrap_or("https://api.anthropic.com"); + tracing::info!( + "[CLAUDE] 使用 Claude API 代理: base_url={} credential_uuid={}", + actual_base_url, + &credential.uuid[..8] + ); + let provider = ClaudeCustomProvider::with_config(api_key.clone(), base_url.clone()); + match provider.call_openai_api(request).await { + Ok(result) => { + // 记录成功 + if let Some(db) = &state.db { + let _ = state.pool_service.mark_healthy( + db, + &credential.uuid, + Some(&request.model), + ); + let _ = state.pool_service.record_usage(db, &credential.uuid); + } + Ok(result) + } + Err(e) => { + if let Some(db) = &state.db { + let _ = state.pool_service.mark_unhealthy( + db, + &credential.uuid, + Some(&e.to_string()), + ); + } + Err(e.to_string()) + } + } + } + CredentialData::AntigravityOAuth { + creds_file_path, + project_id, + } => { + let mut antigravity = AntigravityProvider::new(); + if let Err(e) = antigravity + .load_credentials_from_path(creds_file_path) + .await + { + if let Some(db) = &state.db { + let _ = state.pool_service.mark_unhealthy( + db, + &credential.uuid, + Some(&format!("Failed to load credentials: {}", e)), + ); + } + return Err(e.to_string()); + } + if !antigravity.is_token_valid() { + if let Err(e) = antigravity.refresh_token().await { + if let Some(db) = &state.db { + let _ = state.pool_service.mark_unhealthy( + db, + &credential.uuid, + Some(&format!("Token refresh failed: {}", e)), + ); + } + return Err(e.to_string()); + } + } + + // 设置项目 ID + if let Some(pid) = project_id { + antigravity.project_id = Some(pid.clone()); + } + let proj_id = antigravity.project_id.clone().unwrap_or_default(); + + let antigravity_request = convert_openai_to_antigravity_with_context(request, &proj_id); + match antigravity + .call_api("generateContent", &antigravity_request) + .await + { + Ok(resp) => { + // 记录成功 + if let Some(db) = &state.db { + let _ = state.pool_service.mark_healthy( + db, + &credential.uuid, + Some(&request.model), + ); + let _ = state.pool_service.record_usage(db, &credential.uuid); + } + Ok(convert_antigravity_to_openai_response( + &resp, + &request.model, + )) + } + Err(e) => { + if let Some(db) = &state.db { + let _ = state.pool_service.mark_unhealthy( + db, + &credential.uuid, + Some(&e.to_string()), + ); + } + Err(e.to_string()) + } + } + } + // GeminiOAuth 和 QwenOAuth 暂不支持 WebSocket，需要使用 HTTP 端点 + _ => Err( + "This credential type is not yet supported via WebSocket. Please use HTTP endpoints." + .to_string(), + ), + } +} + +/// WebSocket 专用的 Anthropic 格式 Provider 调用 +pub async fn call_provider_anthropic_for_ws( + state: &AppState, + credential: &ProviderCredential, + request: &AnthropicMessagesRequest, +) -> Result { + use crate::models::provider_pool_model::CredentialData; + + match &credential.credential { + CredentialData::ClaudeKey { api_key, base_url } => { + // 打印 Claude 代理 URL 用于调试 + let actual_base_url = base_url.as_deref().unwrap_or("https://api.anthropic.com"); + tracing::info!( + "[CLAUDE] 使用 Claude API 代理: base_url={} credential_uuid={}", + actual_base_url, + &credential.uuid[..8] + ); + let provider = ClaudeCustomProvider::with_config(api_key.clone(), base_url.clone()); + let resp = match provider.call_api(request).await { + Ok(r) => r, + Err(e) => { + if let Some(db) = &state.db { + let _ = state.pool_service.mark_unhealthy( + db, + &credential.uuid, + Some(&e.to_string()), + ); + } + return Err(e.to_string()); + } + }; + if resp.status().is_success() { + // 记录成功 + if let Some(db) = &state.db { + let _ = + state + .pool_service + .mark_healthy(db, &credential.uuid, Some(&request.model)); + let _ = state.pool_service.record_usage(db, &credential.uuid); + } + resp.json::() + .await + .map_err(|e| e.to_string()) + } else { + let body = resp.text().await.unwrap_or_default(); + if let Some(db) = &state.db { + let _ = state + .pool_service + .mark_unhealthy(db, &credential.uuid, Some(&body)); + } + Err(format!("Upstream error: {}", body)) + } + } + _ => { + // 转换为 OpenAI 格式并调用（健康状态更新在 call_provider_openai_for_ws 中处理） + let openai_request = convert_anthropic_to_openai(request); + let result = call_provider_openai_for_ws(state, credential, &openai_request).await?; + + // 转换响应为 Anthropic 格式 + Ok(serde_json::json!({ + "id": format!("msg_{}", uuid::Uuid::new_v4()), + "type": "message", + "role": "assistant", + "content": [{ + "type": "text", + "text": result.get("choices") + .and_then(|c| c.get(0)) + .and_then(|c| c.get("message")) + .and_then(|m| m.get("content")) + .and_then(|c| c.as_str()) + .unwrap_or("") + }], + "model": request.model, + "stop_reason": "end_turn", + "usage": { + "input_tokens": 0, + "output_tokens": 0 + } + })) + } + } +} diff --git a/src-tauri/src/server/mod.rs b/src-tauri/src/server/mod.rs new file mode 100644 index 000000000..f4d9cd639 --- /dev/null +++ b/src-tauri/src/server/mod.rs @@ -0,0 +1,1801 @@ +//! HTTP API 服务器 +use crate::config::{ + Config, ConfigChangeEvent, ConfigChangeKind, ConfigManager, FileWatcher, HotReloadManager, + ReloadResult, +}; +use crate::converter::anthropic_to_openai::convert_anthropic_to_openai; +use crate::converter::openai_to_antigravity::{ + convert_antigravity_to_openai_response, convert_openai_to_antigravity_with_context, +}; +use crate::credential::CredentialSyncService; +use crate::database::dao::provider_pool::ProviderPoolDao; +use crate::database::DbConnection; +use crate::injection::Injector; +use crate::logger::LogStore; +use crate::models::anthropic::*; +use crate::models::openai::*; +use crate::models::provider_pool_model::CredentialData; +use crate::models::route_model::{RouteInfo, RouteListResponse}; +use crate::processor::{RequestContext, RequestProcessor}; +use crate::providers::antigravity::AntigravityProvider; +use crate::providers::claude_custom::ClaudeCustomProvider; +use crate::providers::gemini::GeminiProvider; +use crate::providers::kiro::KiroProvider; +use crate::providers::openai_custom::OpenAICustomProvider; +use crate::providers::qwen::QwenProvider; +use crate::providers::vertex::VertexProvider; +use crate::server_utils::{ + build_anthropic_response, build_anthropic_stream_response, build_gemini_native_request, health, + message_content_len, models, parse_cw_response, safe_truncate, CWParsedResponse, +}; +use crate::services::provider_pool_service::ProviderPoolService; +use crate::services::token_cache_service::TokenCacheService; +use crate::telemetry::{RequestLog, RequestStatus}; +use crate::websocket::{WsConfig, WsConnectionManager, WsStats}; +use axum::{ + body::Body, + extract::{DefaultBodyLimit, Path, State}, + http::{header, HeaderMap, StatusCode}, + response::{IntoResponse, Response}, + routing::{get, post}, + Json, Router, +}; +use futures::stream; +use serde::{Deserialize, Serialize}; +use std::path::PathBuf; +use std::sync::Arc; +use tokio::sync::{mpsc, oneshot, RwLock}; + +/// 记录请求统计到遥测系统 +pub fn record_request_telemetry( + state: &AppState, + ctx: &RequestContext, + status: crate::telemetry::RequestStatus, + error_message: Option, +) { + use crate::telemetry::RequestLog; + + let provider = ctx.provider.unwrap_or(crate::ProviderType::Kiro); + let mut log = RequestLog::new( + ctx.request_id.clone(), + provider, + ctx.resolved_model.clone(), + ctx.is_stream, + ); + + // 设置状态和持续时间 + match status { + crate::telemetry::RequestStatus::Success => log.mark_success(ctx.elapsed_ms(), 200), + crate::telemetry::RequestStatus::Failed => log.mark_failed( + ctx.elapsed_ms(), + None, + error_message.clone().unwrap_or_default(), + ), + crate::telemetry::RequestStatus::Timeout => log.mark_timeout(ctx.elapsed_ms()), + crate::telemetry::RequestStatus::Cancelled => log.mark_cancelled(ctx.elapsed_ms()), + crate::telemetry::RequestStatus::Retrying => { + log.duration_ms = ctx.elapsed_ms(); + } + } + + // 设置凭证 ID + if let Some(cred_id) = &ctx.credential_id { + log.set_credential_id(cred_id.clone()); + } + + // 设置重试次数 + log.retry_count = ctx.retry_count; + + // 记录到统计聚合器 + { + let stats = state.processor.stats.write(); + stats.record(log.clone()); + } + + // 记录到请求日志记录器（用于前端日志列表显示） + if let Some(logger) = &state.request_logger { + let _ = logger.record(log.clone()); + } + + tracing::info!( + "[TELEMETRY] request_id={} provider={:?} model={} status={:?} duration_ms={}", + ctx.request_id, + provider, + ctx.resolved_model, + status, + ctx.elapsed_ms() + ); +} + +/// 记录 Token 使用量到遥测系统 +pub fn record_token_usage( + state: &AppState, + ctx: &RequestContext, + input_tokens: Option, + output_tokens: Option, +) { + use crate::telemetry::{TokenSource, TokenUsageRecord}; + + // 只有当至少有一个 Token 值时才记录 + if input_tokens.is_none() && output_tokens.is_none() { + return; + } + + let provider = ctx.provider.unwrap_or(crate::ProviderType::Kiro); + let record = TokenUsageRecord::new( + uuid::Uuid::new_v4().to_string(), + provider, + ctx.resolved_model.clone(), + input_tokens.unwrap_or(0), + output_tokens.unwrap_or(0), + TokenSource::Actual, + ) + .with_request_id(ctx.request_id.clone()); + + // 记录到 Token 追踪器 + { + let tokens = state.processor.tokens.write(); + tokens.record(record); + } + + tracing::debug!( + "[TOKEN] request_id={} input={} output={}", + ctx.request_id, + input_tokens.unwrap_or(0), + output_tokens.unwrap_or(0) + ); +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ServerStatus { + pub running: bool, + pub host: String, + pub port: u16, + pub requests: u64, + pub uptime_secs: u64, +} + +pub struct ServerState { + pub config: Config, + pub running: bool, + pub requests: u64, + pub start_time: Option, + pub kiro_provider: KiroProvider, + pub gemini_provider: GeminiProvider, + pub qwen_provider: QwenProvider, + pub openai_custom_provider: OpenAICustomProvider, + pub claude_custom_provider: ClaudeCustomProvider, + pub default_provider_ref: Arc>, + shutdown_tx: Option>, + /// 服务器运行时使用的 API key（启动时从配置复制） + /// 用于 test_api 命令，确保测试使用的 API key 和服务器一致 + pub running_api_key: Option, +} + +impl ServerState { + pub fn new(config: Config) -> Self { + let kiro = KiroProvider::new(); + let gemini = GeminiProvider::new(); + let qwen = QwenProvider::new(); + let openai_custom = OpenAICustomProvider::new(); + let claude_custom = ClaudeCustomProvider::new(); + let default_provider_ref = Arc::new(RwLock::new(config.default_provider.clone())); + + Self { + config, + running: false, + requests: 0, + start_time: None, + kiro_provider: kiro, + gemini_provider: gemini, + qwen_provider: qwen, + openai_custom_provider: openai_custom, + claude_custom_provider: claude_custom, + default_provider_ref, + shutdown_tx: None, + running_api_key: None, + } + } + + pub fn status(&self) -> ServerStatus { + ServerStatus { + running: self.running, + host: self.config.server.host.clone(), + port: self.config.server.port, + requests: self.requests, + uptime_secs: self.start_time.map(|t| t.elapsed().as_secs()).unwrap_or(0), + } + } + + pub async fn start( + &mut self, + logs: Arc>, + pool_service: Arc, + token_cache: Arc, + db: Option, + ) -> Result<(), Box> { + self.start_with_telemetry(logs, pool_service, token_cache, db, None, None, None) + .await + } + + /// 启动服务器（使用共享的遥测实例） + /// + /// 这允许服务器与 TelemetryState 共享同一个 StatsAggregator、TokenTracker 和 RequestLogger， + /// 使得请求处理过程中记录的统计数据能够在前端监控页面中显示。 + pub async fn start_with_telemetry( + &mut self, + logs: Arc>, + pool_service: Arc, + token_cache: Arc, + db: Option, + shared_stats: Option>>, + shared_tokens: Option>>, + shared_logger: Option>, + ) -> Result<(), Box> { + if self.running { + return Ok(()); + } + + let (tx, rx) = oneshot::channel(); + self.shutdown_tx = Some(tx); + + let host = self.config.server.host.clone(); + let port = self.config.server.port; + let api_key = self.config.server.api_key.clone(); + let api_key_for_state = api_key.clone(); // 用于保存到 running_api_key + let default_provider_ref = self.default_provider_ref.clone(); + + // 重新加载凭证 + let _ = self.kiro_provider.load_credentials().await; + let kiro = self.kiro_provider.clone(); + + // 创建参数注入器 + let injection_enabled = self.config.injection.enabled; + let injector = Injector::with_rules( + self.config + .injection + .rules + .iter() + .map(|r| r.clone().into()) + .collect(), + ); + + // 获取配置和配置路径用于热重载 + let config = self.config.clone(); + let config_path = crate::config::ConfigManager::default_config_path(); + + tokio::spawn(async move { + if let Err(e) = run_server( + &host, + port, + &api_key, + default_provider_ref, + kiro, + logs, + rx, + pool_service, + token_cache, + db, + injector, + injection_enabled, + shared_stats, + shared_tokens, + shared_logger, + Some(config), + Some(config_path), + ) + .await + { + tracing::error!("Server error: {}", e); + } + }); + + self.running = true; + self.start_time = Some(std::time::Instant::now()); + // 保存服务器运行时使用的 API key，用于 test_api 命令 + self.running_api_key = Some(api_key_for_state); + Ok(()) + } + + pub async fn stop(&mut self) { + if let Some(tx) = self.shutdown_tx.take() { + let _ = tx.send(()); + } + self.running = false; + self.start_time = None; + self.running_api_key = None; + } +} + +impl Clone for KiroProvider { + fn clone(&self) -> Self { + Self { + credentials: self.credentials.clone(), + client: reqwest::Client::new(), + creds_path: self.creds_path.clone(), + } + } +} + +pub mod handlers; + +#[derive(Clone)] +#[allow(dead_code)] +pub struct AppState { + pub api_key: String, + pub base_url: String, + pub default_provider: Arc>, + pub kiro: Arc>, + pub logs: Arc>, + pub kiro_refresh_lock: Arc>, + pub gemini_refresh_lock: Arc>, + pub qwen_refresh_lock: Arc>, + pub pool_service: Arc, + pub token_cache: Arc, + pub db: Option, + /// 参数注入器 + pub injector: Arc>, + /// 是否启用参数注入 + pub injection_enabled: Arc>, + /// 请求处理器 + pub processor: Arc, + /// WebSocket 连接管理器 + pub ws_manager: Arc, + /// WebSocket 统计信息 + pub ws_stats: Arc, + /// 热重载管理器 + pub hot_reload_manager: Option>, + /// 请求日志记录器（与 TelemetryState 共享） + pub request_logger: Option>, + /// Amp CLI 路由器 + pub amp_router: Arc, +} + +/// 启动配置文件监控 +/// +/// 监控配置文件变化并触发热重载。 +/// +/// # 连接保持 +/// +/// 热重载过程不会中断现有连接： +/// - 配置更新在独立的 tokio 任务中异步执行 +/// - 使用 RwLock 进行原子性更新，不会阻塞正在处理的请求 +/// - 服务器继续运行，不需要重启 +/// - HTTP 和 WebSocket 连接保持活跃 +async fn start_config_watcher( + config_path: PathBuf, + hot_reload_manager: Option>, + processor: Arc, + logs: Arc>, + db: Option, + config_manager: Option>>, +) -> Option { + let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::(); + + // 创建文件监控器 + let mut watcher = match FileWatcher::new(&config_path, tx) { + Ok(w) => w, + Err(e) => { + tracing::error!("[HOT_RELOAD] 创建文件监控器失败: {}", e); + return None; + } + }; + + // 启动监控 + if let Err(e) = watcher.start() { + tracing::error!("[HOT_RELOAD] 启动文件监控失败: {}", e); + return None; + } + + tracing::info!("[HOT_RELOAD] 配置文件监控已启动: {:?}", config_path); + + // 启动事件处理任务 + let hot_reload_manager_clone = hot_reload_manager.clone(); + let processor_clone = processor.clone(); + let logs_clone = logs.clone(); + let db_clone = db.clone(); + let config_manager_clone = config_manager.clone(); + + tokio::spawn(async move { + while let Some(event) = rx.recv().await { + // 只处理修改事件 + if event.kind != ConfigChangeKind::Modified { + continue; + } + + tracing::info!("[HOT_RELOAD] 检测到配置文件变更: {:?}", event.path); + logs_clone.write().await.add( + "info", + &format!("[HOT_RELOAD] 检测到配置文件变更: {:?}", event.path), + ); + + // 执行热重载 + if let Some(ref manager) = hot_reload_manager_clone { + let result = manager.reload(); + match &result { + ReloadResult::Success { .. } => { + tracing::info!("[HOT_RELOAD] 配置热重载成功"); + logs_clone + .write() + .await + .add("info", "[HOT_RELOAD] 配置热重载成功"); + + // 更新处理器中的组件 + let new_config = manager.config(); + update_processor_config(&processor_clone, &new_config).await; + + // 同步凭证池 + if let (Some(ref db), Some(ref cfg_manager)) = + (&db_clone, &config_manager_clone) + { + match sync_credential_pool_from_config(db, cfg_manager, &logs_clone) + .await + { + Ok(count) => { + tracing::info!( + "[HOT_RELOAD] 凭证池同步完成，共 {} 个凭证", + count + ); + logs_clone.write().await.add( + "info", + &format!( + "[HOT_RELOAD] 凭证池同步完成，共 {} 个凭证", + count + ), + ); + } + Err(e) => { + tracing::warn!("[HOT_RELOAD] 凭证池同步失败: {}", e); + logs_clone.write().await.add( + "warn", + &format!("[HOT_RELOAD] 凭证池同步失败: {}", e), + ); + } + } + } + } + ReloadResult::RolledBack { error, .. } => { + tracing::warn!("[HOT_RELOAD] 配置热重载失败，已回滚: {}", error); + logs_clone.write().await.add( + "warn", + &format!("[HOT_RELOAD] 配置热重载失败，已回滚: {}", error), + ); + } + ReloadResult::Failed { + error, + rollback_error, + .. + } => { + tracing::error!( + "[HOT_RELOAD] 配置热重载失败: {}, 回滚错误: {:?}", + error, + rollback_error + ); + logs_clone.write().await.add( + "error", + &format!( + "[HOT_RELOAD] 配置热重载失败: {}, 回滚错误: {:?}", + error, rollback_error + ), + ); + } + } + } + } + }); + + Some(watcher) +} + +/// 更新处理器配置 +/// +/// 当配置热重载成功后，更新 RequestProcessor 中的各个组件。 +/// +/// # 原子性更新 +/// +/// 每个组件的更新都是原子性的，使用 RwLock 确保： +/// - 正在处理的请求不会看到部分更新的状态 +/// - 更新过程不会阻塞新请求的处理 +/// - 现有连接不受影响 +async fn update_processor_config(processor: &RequestProcessor, config: &Config) { + // 更新注入器规则 + { + let mut injector = processor.injector.write().await; + injector.clear(); + for rule in &config.injection.rules { + injector.add_rule(rule.clone().into()); + } + tracing::debug!( + "[HOT_RELOAD] 注入器规则已更新: {} 条规则", + config.injection.rules.len() + ); + } + + // 更新路由器规则 + { + let mut router = processor.router.write().await; + router.clear_rules(); + for rule in &config.routing.rules { + // 解析 provider 字符串为 ProviderType + if let Ok(provider_type) = rule.provider.parse::() { + router.add_rule(crate::router::RoutingRule { + pattern: rule.pattern.clone(), + target_provider: provider_type, + priority: rule.priority, + enabled: true, + }); + } else { + tracing::warn!("[HOT_RELOAD] 无法解析 provider: {}", rule.provider); + } + } + tracing::debug!( + "[HOT_RELOAD] 路由规则已更新: {} 条规则", + config.routing.rules.len() + ); + } + + // 更新模型映射器 + { + let mut mapper = processor.mapper.write().await; + mapper.clear(); + for (alias, model) in &config.routing.model_aliases { + mapper.add_alias(alias, model); + } + tracing::debug!( + "[HOT_RELOAD] 模型别名已更新: {} 个别名", + config.routing.model_aliases.len() + ); + } + + // 注意：重试配置目前不支持热更新，因为 Retrier 是不可变的 + // 如果需要更新重试配置，需要重启服务器 + tracing::debug!( + "[HOT_RELOAD] 重试配置: max_retries={}, base_delay={}ms (需重启生效)", + config.retry.max_retries, + config.retry.base_delay_ms + ); + + tracing::info!("[HOT_RELOAD] 处理器配置更新完成"); +} + +/// 从配置同步凭证池 +/// +/// 当配置热重载成功后，从 YAML 配置中加载凭证并同步到数据库。 +/// +/// # 同步策略 +/// +/// - 从配置中加载所有凭证 +/// - 对于配置中存在但数据库中不存在的凭证，添加到数据库 +/// - 对于配置中存在且数据库中也存在的凭证，更新数据库中的记录 +/// - 对于数据库中存在但配置中不存在的凭证，保留（不删除，避免丢失运行时状态） +async fn sync_credential_pool_from_config( + db: &DbConnection, + config_manager: &Arc>, + _logs: &Arc>, +) -> Result { + // 创建凭证同步服务 + let sync_service = CredentialSyncService::new(config_manager.clone()); + + // 从配置加载凭证 + let credentials = sync_service.load_from_config().map_err(|e| e.to_string())?; + + let conn = db.lock().map_err(|e| e.to_string())?; + let mut synced_count = 0; + + for cred in &credentials { + // 检查凭证是否已存在 + let existing = + ProviderPoolDao::get_by_uuid(&conn, &cred.uuid).map_err(|e| e.to_string())?; + + if existing.is_some() { + // 更新现有凭证 + ProviderPoolDao::update(&conn, cred).map_err(|e| e.to_string())?; + tracing::debug!( + "[HOT_RELOAD] 更新凭证: {} ({})", + cred.uuid, + cred.provider_type + ); + } else { + // 添加新凭证 + ProviderPoolDao::insert(&conn, cred).map_err(|e| e.to_string())?; + tracing::debug!( + "[HOT_RELOAD] 添加凭证: {} ({})", + cred.uuid, + cred.provider_type + ); + } + synced_count += 1; + } + + Ok(synced_count) +} + +async fn run_server( + host: &str, + port: u16, + api_key: &str, + default_provider: Arc>, + kiro: KiroProvider, + logs: Arc>, + shutdown: oneshot::Receiver<()>, + pool_service: Arc, + token_cache: Arc, + db: Option, + injector: Injector, + injection_enabled: bool, + shared_stats: Option>>, + shared_tokens: Option>>, + shared_logger: Option>, + config: Option, + config_path: Option, +) -> Result<(), Box> { + let base_url = format!("http://{}:{}", host, port); + + // 创建请求处理器（使用共享的遥测实例或默认实例） + let processor = match (shared_stats, shared_tokens) { + (Some(stats), Some(tokens)) => Arc::new(RequestProcessor::with_shared_telemetry( + pool_service.clone(), + stats, + tokens, + )), + _ => Arc::new(RequestProcessor::with_defaults(pool_service.clone())), + }; + + // 将注入器规则同步到处理器 + { + let mut proc_injector = processor.injector.write().await; + for rule in injector.rules() { + proc_injector.add_rule(rule.clone()); + } + } + + // 初始化 WebSocket 管理器 + let ws_manager = Arc::new(WsConnectionManager::new(WsConfig::default())); + let ws_stats = ws_manager.stats().clone(); + + // 初始化热重载管理器 + let hot_reload_manager = match (&config, &config_path) { + (Some(cfg), Some(path)) => Some(Arc::new(HotReloadManager::new(cfg.clone(), path.clone()))), + _ => None, + }; + + // 初始化配置管理器（用于凭证池同步） + let config_manager: Option>> = + match (&config, &config_path) { + (Some(cfg), Some(path)) => Some(Arc::new(std::sync::RwLock::new( + ConfigManager::with_config(cfg.clone(), path.clone()), + ))), + _ => None, + }; + + let logs_clone = logs.clone(); + let db_clone = db.clone(); + + // 初始化 Amp CLI 路由器 + let amp_router = Arc::new(crate::router::AmpRouter::new( + config + .as_ref() + .map(|c| c.ampcode.clone()) + .unwrap_or_default(), + )); + + let state = AppState { + api_key: api_key.to_string(), + base_url, + default_provider, + kiro: Arc::new(RwLock::new(kiro)), + logs, + kiro_refresh_lock: Arc::new(tokio::sync::Mutex::new(())), + gemini_refresh_lock: Arc::new(tokio::sync::Mutex::new(())), + qwen_refresh_lock: Arc::new(tokio::sync::Mutex::new(())), + pool_service, + token_cache, + db, + injector: Arc::new(RwLock::new(injector)), + injection_enabled: Arc::new(RwLock::new(injection_enabled)), + processor: processor.clone(), + ws_manager, + ws_stats, + hot_reload_manager: hot_reload_manager.clone(), + request_logger: shared_logger, + amp_router, + }; + + // 启动配置文件监控 + let _file_watcher = if let Some(path) = config_path { + start_config_watcher( + path, + hot_reload_manager, + processor, + logs_clone, + db_clone, + config_manager, + ) + .await + } else { + None + }; + + // 设置请求体大小限制为 100MB，支持大型上下文请求（如 Claude Code 的 /compact 命令） + let body_limit = 100 * 1024 * 1024; // 100MB + + // 创建管理 API 路由（带认证中间件） + let management_config = config + .as_ref() + .map(|c| c.remote_management.clone()) + .unwrap_or_default(); + + let management_routes = Router::new() + .route("/v0/management/status", get(handlers::management_status)) + .route( + "/v0/management/credentials", + get(handlers::management_list_credentials), + ) + .route( + "/v0/management/credentials", + post(handlers::management_add_credential), + ) + .route( + "/v0/management/config", + get(handlers::management_get_config), + ) + .route( + "/v0/management/config", + axum::routing::put(handlers::management_update_config), + ) + .layer(crate::middleware::ManagementAuthLayer::new( + management_config, + )); + + let app = Router::new() + .route("/health", get(health)) + .route("/v1/models", get(models)) + .route("/v1/routes", get(list_routes)) + .route("/v1/chat/completions", post(handlers::chat_completions)) + .route("/v1/messages", post(handlers::anthropic_messages)) + .route("/v1/messages/count_tokens", post(count_tokens)) + // Gemini 原生协议路由 + .route("/v1/gemini/*path", post(gemini_generate_content)) + // WebSocket 路由 + .route("/v1/ws", get(handlers::ws_upgrade_handler)) + .route("/ws", get(handlers::ws_upgrade_handler)) + // 多供应商路由 + .route( + "/:selector/v1/messages", + post(anthropic_messages_with_selector), + ) + .route( + "/:selector/v1/chat/completions", + post(chat_completions_with_selector), + ) + // Amp CLI 路由 + .route( + "/api/provider/:provider/v1/chat/completions", + post(amp_chat_completions), + ) + .route("/api/provider/:provider/v1/messages", post(amp_messages)) + // Amp CLI 管理代理路由 + .route( + "/api/auth/*path", + axum::routing::any(amp_management_proxy_auth), + ) + .route( + "/api/user/*path", + axum::routing::any(amp_management_proxy_user), + ) + // 管理 API 路由 + .merge(management_routes) + .layer(DefaultBodyLimit::max(body_limit)) + .with_state(state); + + let addr: std::net::SocketAddr = format!("{host}:{port}").parse()?; + let listener = tokio::net::TcpListener::bind(addr).await?; + + tracing::info!("Server listening on {}", addr); + + axum::serve(listener, app) + .with_graceful_shutdown(async move { + let _ = shutdown.await; + }) + .await?; + + Ok(()) +} + +async fn count_tokens( + State(state): State, + headers: HeaderMap, + Json(_request): Json, +) -> Response { + if let Err(e) = handlers::verify_api_key(&headers, &state.api_key).await { + return e.into_response(); + } + + // Claude Code 需要这个端点，返回估算值 + Json(serde_json::json!({ + "input_tokens": 100 + })) + .into_response() +} + +/// Gemini 原生协议处理 +/// 路由: POST /v1/gemini/{model}:{method} +/// 例如: /v1/gemini/gemini-3-pro-preview:generateContent +async fn gemini_generate_content( + State(state): State, + headers: HeaderMap, + Path(path): Path, + Json(request): Json, +) -> Response { + if let Err(e) = handlers::verify_api_key(&headers, &state.api_key).await { + return e.into_response(); + } + + // 解析路径: {model}:{method} + // 例如: gemini-3-pro-preview:generateContent + let parts: Vec<&str> = path.splitn(2, ':').collect(); + if parts.len() != 2 { + return ( + StatusCode::BAD_REQUEST, + Json(serde_json::json!({ + "error": { + "message": format!("无效的路径格式: {}，期望格式: model:method", path) + } + })), + ) + .into_response(); + } + + let model = parts[0]; + let method = parts[1]; + + state.logs.write().await.add( + "info", + &format!( + "[GEMINI] POST /v1/gemini/{} model={} method={}", + path, model, method + ), + ); + + // 目前只支持 generateContent 方法 + if method != "generateContent" && method != "streamGenerateContent" { + return ( + StatusCode::BAD_REQUEST, + Json(serde_json::json!({ + "error": { + "message": format!("不支持的方法: {}，目前只支持 generateContent", method) + } + })), + ) + .into_response(); + } + + let is_stream = method == "streamGenerateContent"; + + // 获取默认 provider + let default_provider = state.default_provider.read().await.clone(); + + // 尝试从凭证池中选择 Antigravity 凭证 + let credential = match &state.db { + Some(db) => state + .pool_service + .select_credential(db, &default_provider, Some(model)) + .ok() + .flatten(), + None => None, + }; + + let cred = match credential { + Some(c) => c, + None => { + return ( + StatusCode::NOT_FOUND, + Json(serde_json::json!({ + "error": { + "message": "没有可用的 Antigravity 凭证，请先添加凭证" + } + })), + ) + .into_response(); + } + }; + + state.logs.write().await.add( + "info", + &format!( + "[GEMINI] 使用凭证: type={} name={:?} uuid={}", + cred.provider_type, + cred.name, + &cred.uuid[..8] + ), + ); + + // 调用 Antigravity Provider + match &cred.credential { + CredentialData::AntigravityOAuth { + creds_file_path, + project_id, + } => { + let mut antigravity = AntigravityProvider::new(); + if let Err(e) = antigravity + .load_credentials_from_path(creds_file_path) + .await + { + return ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({ + "error": { + "message": format!("加载 Antigravity 凭证失败: {}", e) + } + })), + ) + .into_response(); + } + + // 检查并刷新 token + if antigravity.is_token_expiring_soon() { + if let Err(e) = antigravity.refresh_token().await { + return ( + StatusCode::UNAUTHORIZED, + Json(serde_json::json!({ + "error": { + "message": format!("Token 刷新失败: {}", e) + } + })), + ) + .into_response(); + } + } + + // 设置项目 ID + if let Some(pid) = project_id { + antigravity.project_id = Some(pid.clone()); + } else if antigravity.project_id.is_none() { + // 如果凭证中没有 project_id，尝试从 API 获取或生成随机 ID + if let Err(e) = antigravity.discover_project().await { + tracing::warn!("[Antigravity] 获取项目 ID 失败: {}，使用随机生成的 ID", e); + // 生成随机项目 ID + let uuid = uuid::Uuid::new_v4(); + let bytes = uuid.as_bytes(); + let adjectives = ["useful", "bright", "swift", "calm", "bold"]; + let nouns = ["fuze", "wave", "spark", "flow", "core"]; + let adj = adjectives[(bytes[0] as usize) % adjectives.len()]; + let noun = nouns[(bytes[1] as usize) % nouns.len()]; + let random_part: String = uuid.to_string()[..5].to_lowercase(); + antigravity.project_id = Some(format!("{}-{}-{}", adj, noun, random_part)); + } + } + + let proj_id = antigravity.project_id.clone().unwrap_or_else(|| { + // 最后的后备：生成随机 ID + let uuid = uuid::Uuid::new_v4(); + format!("proxycast-{}", &uuid.to_string()[..8]) + }); + + state + .logs + .write() + .await + .add("debug", &format!("[GEMINI] 使用 project_id: {}", proj_id)); + + // 构建 Antigravity 请求体 + // 直接使用用户传入的 Gemini 格式请求，只添加必要的字段 + let antigravity_request = build_gemini_native_request(&request, model, &proj_id); + + state.logs.write().await.add( + "debug", + &format!( + "[GEMINI] 请求体: {}", + serde_json::to_string(&antigravity_request).unwrap_or_default() + ), + ); + + if is_stream { + // 流式响应 - 暂不支持，返回错误 + return ( + StatusCode::NOT_IMPLEMENTED, + Json(serde_json::json!({ + "error": { + "message": "流式响应暂不支持，请使用 generateContent" + } + })), + ) + .into_response(); + } + + // 非流式响应 + match antigravity + .call_api("generateContent", &antigravity_request) + .await + { + Ok(resp) => { + state.logs.write().await.add( + "info", + &format!( + "[GEMINI] 响应成功: {}", + serde_json::to_string(&resp) + .unwrap_or_default() + .chars() + .take(200) + .collect::() + ), + ); + + // 直接返回 Gemini 格式响应 + Json(resp).into_response() + } + Err(e) => { + state + .logs + .write() + .await + .add("error", &format!("[GEMINI] 请求失败: {}", e)); + + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({ + "error": { + "message": e.to_string() + } + })), + ) + .into_response() + } + } + } + _ => ( + StatusCode::BAD_REQUEST, + Json(serde_json::json!({ + "error": { + "message": "Gemini 原生协议只支持 Antigravity 凭证" + } + })), + ) + .into_response(), + } +} + +/// 列出所有可用路由 +async fn list_routes(State(state): State) -> impl IntoResponse { + let routes = match &state.db { + Some(db) => state + .pool_service + .get_available_routes(db, &state.base_url) + .unwrap_or_default(), + None => Vec::new(), + }; + + // 添加默认路由 + let mut all_routes = vec![RouteInfo { + selector: "default".to_string(), + provider_type: "kiro".to_string(), + credential_count: 1, + endpoints: vec![ + crate::models::route_model::RouteEndpoint { + path: "/v1/messages".to_string(), + protocol: "claude".to_string(), + url: format!("{}/v1/messages", state.base_url), + }, + crate::models::route_model::RouteEndpoint { + path: "/v1/chat/completions".to_string(), + protocol: "openai".to_string(), + url: format!("{}/v1/chat/completions", state.base_url), + }, + ], + tags: vec!["默认".to_string()], + enabled: true, + }]; + all_routes.extend(routes); + + let response = RouteListResponse { + base_url: state.base_url.clone(), + default_provider: "kiro".to_string(), + routes: all_routes, + }; + + Json(response) +} + +/// 带选择器的 Anthropic messages 处理 +async fn anthropic_messages_with_selector( + State(state): State, + Path(selector): Path, + headers: HeaderMap, + Json(request): Json, +) -> Response { + // 使用 Anthropic 格式的认证验证 + if let Err(e) = handlers::verify_api_key_anthropic(&headers, &state.api_key).await { + state.logs.write().await.add( + "warn", + &format!("Unauthorized request to /{}/v1/messages", selector), + ); + return e.into_response(); + } + + state.logs.write().await.add( + "info", + &format!( + "[REQ] POST /{}/v1/messages model={} stream={}", + selector, request.model, request.stream + ), + ); + + // 尝试解析凭证 + let credential = match &state.db { + Some(db) => { + // 首先尝试按名称查找 + if let Ok(Some(cred)) = state.pool_service.get_by_name(db, &selector) { + Some(cred) + } + // 然后尝试按 UUID 查找 + else if let Ok(Some(cred)) = state.pool_service.get_by_uuid(db, &selector) { + Some(cred) + } + // 最后尝试按 provider 类型轮询 + else if let Ok(Some(cred)) = + state + .pool_service + .select_credential(db, &selector, Some(&request.model)) + { + Some(cred) + } else { + None + } + } + None => None, + }; + + match credential { + Some(cred) => { + state.logs.write().await.add( + "info", + &format!( + "[ROUTE] Using credential: type={} name={:?} uuid={}", + cred.provider_type, + cred.name, + &cred.uuid[..8] + ), + ); + + // 根据凭证类型调用相应的 Provider + handlers::call_provider_anthropic(&state, &cred, &request).await + } + None => { + // 回退到默认 Kiro provider + state.logs.write().await.add( + "warn", + &format!( + "[ROUTE] Credential not found for selector '{}', falling back to default", + selector + ), + ); + // 调用原有的 Kiro 处理逻辑 + anthropic_messages_internal(&state, &request).await + } + } +} + +/// 带选择器的 OpenAI chat completions 处理 +async fn chat_completions_with_selector( + State(state): State, + Path(selector): Path, + headers: HeaderMap, + Json(request): Json, +) -> Response { + if let Err(e) = handlers::verify_api_key(&headers, &state.api_key).await { + state.logs.write().await.add( + "warn", + &format!("Unauthorized request to /{}/v1/chat/completions", selector), + ); + return e.into_response(); + } + + state.logs.write().await.add( + "info", + &format!( + "[REQ] POST /{}/v1/chat/completions model={} stream={}", + selector, request.model, request.stream + ), + ); + + // 尝试解析凭证 + let credential = match &state.db { + Some(db) => { + if let Ok(Some(cred)) = state.pool_service.get_by_name(db, &selector) { + Some(cred) + } else if let Ok(Some(cred)) = state.pool_service.get_by_uuid(db, &selector) { + Some(cred) + } else if let Ok(Some(cred)) = + state + .pool_service + .select_credential(db, &selector, Some(&request.model)) + { + Some(cred) + } else { + None + } + } + None => None, + }; + + match credential { + Some(cred) => { + state.logs.write().await.add( + "info", + &format!( + "[ROUTE] Using credential: type={} name={:?} uuid={}", + cred.provider_type, + cred.name, + &cred.uuid[..8] + ), + ); + + handlers::call_provider_openai(&state, &cred, &request).await + } + None => { + state.logs.write().await.add( + "warn", + &format!( + "[ROUTE] Credential not found for selector '{}', falling back to default", + selector + ), + ); + chat_completions_internal(&state, &request).await + } + } +} + +// ============ Amp CLI 路由处理 ============ + +/// Amp CLI chat completions 处理 +/// +/// 处理 `/api/provider/:provider/v1/chat/completions` 路由 +/// 支持模型映射，将不可用模型映射到可用替代 +async fn amp_chat_completions( + State(state): State, + Path(provider): Path, + headers: HeaderMap, + Json(mut request): Json, +) -> Response { + if let Err(e) = handlers::verify_api_key(&headers, &state.api_key).await { + state.logs.write().await.add( + "warn", + &format!( + "Unauthorized request to /api/provider/{}/v1/chat/completions", + provider + ), + ); + return e.into_response(); + } + + // 应用模型映射 + let original_model = request.model.clone(); + let mapped_model = state.amp_router.apply_model_mapping(&request.model); + if mapped_model != original_model { + state.logs.write().await.add( + "info", + &format!( + "[AMP] Model mapping applied: {} -> {}", + original_model, mapped_model + ), + ); + request.model = mapped_model; + } + + state.logs.write().await.add( + "info", + &format!( + "[AMP] POST /api/provider/{}/v1/chat/completions model={} stream={}", + provider, request.model, request.stream + ), + ); + + // 尝试根据 provider 名称选择凭证 + let credential = match &state.db { + Some(db) => { + // 首先尝试按 provider 类型选择 + if let Ok(Some(cred)) = + state + .pool_service + .select_credential(db, &provider, Some(&request.model)) + { + Some(cred) + } + // 然后尝试按名称查找 + else if let Ok(Some(cred)) = state.pool_service.get_by_name(db, &provider) { + Some(cred) + } + // 最后尝试按 UUID 查找 + else if let Ok(Some(cred)) = state.pool_service.get_by_uuid(db, &provider) { + Some(cred) + } else { + None + } + } + None => None, + }; + + match credential { + Some(cred) => { + state.logs.write().await.add( + "info", + &format!( + "[AMP] Using credential: type={} name={:?} uuid={}", + cred.provider_type, + cred.name, + &cred.uuid[..8] + ), + ); + handlers::call_provider_openai(&state, &cred, &request).await + } + None => { + state.logs.write().await.add( + "warn", + &format!( + "[AMP] Credential not found for provider '{}', falling back to default", + provider + ), + ); + chat_completions_internal(&state, &request).await + } + } +} + +/// Amp CLI messages 处理 +/// +/// 处理 `/api/provider/:provider/v1/messages` 路由 +/// 支持模型映射，将不可用模型映射到可用替代 +async fn amp_messages( + State(state): State, + Path(provider): Path, + headers: HeaderMap, + Json(mut request): Json, +) -> Response { + // 使用 Anthropic 格式的认证验证 + if let Err(e) = handlers::verify_api_key_anthropic(&headers, &state.api_key).await { + state.logs.write().await.add( + "warn", + &format!( + "Unauthorized request to /api/provider/{}/v1/messages", + provider + ), + ); + return e.into_response(); + } + + // 应用模型映射 + let original_model = request.model.clone(); + let mapped_model = state.amp_router.apply_model_mapping(&request.model); + if mapped_model != original_model { + state.logs.write().await.add( + "info", + &format!( + "[AMP] Model mapping applied: {} -> {}", + original_model, mapped_model + ), + ); + request.model = mapped_model; + } + + state.logs.write().await.add( + "info", + &format!( + "[AMP] POST /api/provider/{}/v1/messages model={} stream={}", + provider, request.model, request.stream + ), + ); + + // 尝试根据 provider 名称选择凭证 + let credential = match &state.db { + Some(db) => { + // 首先尝试按 provider 类型选择 + if let Ok(Some(cred)) = + state + .pool_service + .select_credential(db, &provider, Some(&request.model)) + { + Some(cred) + } + // 然后尝试按名称查找 + else if let Ok(Some(cred)) = state.pool_service.get_by_name(db, &provider) { + Some(cred) + } + // 最后尝试按 UUID 查找 + else if let Ok(Some(cred)) = state.pool_service.get_by_uuid(db, &provider) { + Some(cred) + } else { + None + } + } + None => None, + }; + + match credential { + Some(cred) => { + state.logs.write().await.add( + "info", + &format!( + "[AMP] Using credential: type={} name={:?} uuid={}", + cred.provider_type, + cred.name, + &cred.uuid[..8] + ), + ); + handlers::call_provider_anthropic(&state, &cred, &request).await + } + None => { + state.logs.write().await.add( + "warn", + &format!( + "[AMP] Credential not found for provider '{}', falling back to default", + provider + ), + ); + anthropic_messages_internal(&state, &request).await + } + } +} + +/// Amp CLI 管理代理 - auth 路由 +/// +/// 处理 `/api/auth/*` 路由，将请求代理到上游 URL +async fn amp_management_proxy_auth( + State(state): State, + Path(path): Path, + headers: HeaderMap, + method: axum::http::Method, + body: axum::body::Bytes, +) -> Response { + amp_management_proxy_internal(state, &format!("auth/{}", path), headers, method, body).await +} + +/// Amp CLI 管理代理 - user 路由 +/// +/// 处理 `/api/user/*` 路由，将请求代理到上游 URL +async fn amp_management_proxy_user( + State(state): State, + Path(path): Path, + headers: HeaderMap, + method: axum::http::Method, + body: axum::body::Bytes, +) -> Response { + amp_management_proxy_internal(state, &format!("user/{}", path), headers, method, body).await +} + +/// Amp CLI 管理代理内部实现 +/// +/// 处理 `/api/auth/*` 和 `/api/user/*` 路由 +/// 将请求代理到上游 URL +/// +/// # 参数 +/// - `path`: 请求路径（不含 /api/ 前缀，如 "auth/login" 或 "user/profile"） +async fn amp_management_proxy_internal( + state: AppState, + path: &str, + headers: HeaderMap, + method: axum::http::Method, + body: axum::body::Bytes, +) -> Response { + let full_path = format!("/api/{}", path); + + // 检查是否是管理路由 + if !state.amp_router.is_management_route(&full_path) { + state.logs.write().await.add( + "warn", + &format!("[AMP] Invalid management route: {}", full_path), + ); + return ( + StatusCode::NOT_FOUND, + Json(serde_json::json!({"error": {"message": "Not found"}})), + ) + .into_response(); + } + + // 检查 localhost 限制 + if state.amp_router.restrict_management_to_localhost() { + // 从 headers 中获取客户端 IP + let client_ip = headers + .get("x-forwarded-for") + .and_then(|v| v.to_str().ok()) + .map(|s| s.split(',').next().unwrap_or("").trim().to_string()) + .or_else(|| { + headers + .get("x-real-ip") + .and_then(|v| v.to_str().ok()) + .map(|s| s.to_string()) + }); + + if let Some(ip) = &client_ip { + let is_localhost = ip == "127.0.0.1" || ip == "::1" || ip == "localhost"; + if !is_localhost { + state.logs.write().await.add( + "warn", + &format!("[AMP] Management proxy blocked from non-localhost: {}", ip), + ); + return ( + StatusCode::FORBIDDEN, + Json(serde_json::json!({"error": {"message": "Management endpoints are restricted to localhost"}})), + ) + .into_response(); + } + } + } + + // 获取上游 URL + let upstream_url = match state.amp_router.get_management_upstream_path(&full_path) { + Some(url) => url, + None => { + state.logs.write().await.add( + "warn", + &format!("[AMP] No upstream URL configured for management proxy"), + ); + return ( + StatusCode::SERVICE_UNAVAILABLE, + Json(serde_json::json!({"error": {"message": "Upstream URL not configured"}})), + ) + .into_response(); + } + }; + + state.logs.write().await.add( + "info", + &format!( + "[AMP] Proxying management request: {} {} -> {}", + method, full_path, upstream_url + ), + ); + + // 创建 HTTP 客户端 + let client = reqwest::Client::new(); + + // 构建请求 + let mut request_builder = match method { + axum::http::Method::GET => client.get(&upstream_url), + axum::http::Method::POST => client.post(&upstream_url), + axum::http::Method::PUT => client.put(&upstream_url), + axum::http::Method::DELETE => client.delete(&upstream_url), + axum::http::Method::PATCH => client.patch(&upstream_url), + axum::http::Method::HEAD => client.head(&upstream_url), + axum::http::Method::OPTIONS => client.request(reqwest::Method::OPTIONS, &upstream_url), + _ => { + return ( + StatusCode::METHOD_NOT_ALLOWED, + Json(serde_json::json!({"error": {"message": "Method not allowed"}})), + ) + .into_response(); + } + }; + + // 复制请求头（排除 host 和 content-length） + for (name, value) in headers.iter() { + let name_str = name.as_str().to_lowercase(); + if name_str != "host" && name_str != "content-length" { + if let Ok(value_str) = value.to_str() { + request_builder = request_builder.header(name.as_str(), value_str); + } + } + } + + // 添加请求体 + if !body.is_empty() { + request_builder = request_builder.body(body.to_vec()); + } + + // 发送请求 + match request_builder.send().await { + Ok(response) => { + let status = response.status(); + let response_headers = response.headers().clone(); + + match response.bytes().await { + Ok(response_body) => { + let mut builder = Response::builder().status(status.as_u16()); + + // 复制响应头 + for (name, value) in response_headers.iter() { + let name_str = name.as_str().to_lowercase(); + // 排除 transfer-encoding 和 content-length（axum 会自动处理） + if name_str != "transfer-encoding" && name_str != "content-length" { + builder = builder.header(name.as_str(), value.to_str().unwrap_or("")); + } + } + + builder + .body(Body::from(response_body.to_vec())) + .unwrap_or_else(|_| { + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"error": {"message": "Failed to build response"}})), + ) + .into_response() + }) + } + Err(e) => { + state.logs.write().await.add( + "error", + &format!("[AMP] Failed to read upstream response: {}", e), + ); + ( + StatusCode::BAD_GATEWAY, + Json(serde_json::json!({"error": {"message": format!("Failed to read upstream response: {}", e)}})), + ) + .into_response() + } + } + } + Err(e) => { + state.logs.write().await.add( + "error", + &format!("[AMP] Failed to proxy request to upstream: {}", e), + ); + ( + StatusCode::BAD_GATEWAY, + Json(serde_json::json!({"error": {"message": format!("Failed to connect to upstream: {}", e)}})), + ) + .into_response() + } + } +} + +/// 内部 Anthropic messages 处理 (使用默认 Kiro) +async fn anthropic_messages_internal( + state: &AppState, + request: &AnthropicMessagesRequest, +) -> Response { + // 检查 token + { + let _guard = state.kiro_refresh_lock.lock().await; + let mut kiro = state.kiro.write().await; + let needs_refresh = + kiro.credentials.access_token.is_none() || kiro.is_token_expiring_soon(); + if needs_refresh { + if let Err(e) = kiro.refresh_token().await { + state + .logs + .write() + .await + .add("error", &format!("[AUTH] Token refresh failed: {e}")); + return ( + StatusCode::UNAUTHORIZED, + Json(serde_json::json!({"error": {"message": format!("Token refresh failed: {e}")}})), + ) + .into_response(); + } + } + } + + let openai_request = convert_anthropic_to_openai(request); + let kiro = state.kiro.read().await; + + match kiro.call_api(&openai_request).await { + Ok(resp) => { + let status = resp.status(); + if status.is_success() { + match resp.bytes().await { + Ok(bytes) => { + let body = String::from_utf8_lossy(&bytes).to_string(); + let parsed = parse_cw_response(&body); + if request.stream { + build_anthropic_stream_response(&request.model, &parsed) + } else { + build_anthropic_response(&request.model, &parsed) + } + } + Err(e) => ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"error": {"message": e.to_string()}})), + ) + .into_response(), + } + } else { + let body = resp.text().await.unwrap_or_default(); + ( + StatusCode::from_u16(status.as_u16()).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR), + Json(serde_json::json!({"error": {"message": format!("Upstream error: {}", body)}})), + ) + .into_response() + } + } + Err(e) => ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"error": {"message": e.to_string()}})), + ) + .into_response(), + } +} + +/// 内部 OpenAI chat completions 处理 (使用默认 Kiro) +async fn chat_completions_internal(state: &AppState, request: &ChatCompletionRequest) -> Response { + { + let _guard = state.kiro_refresh_lock.lock().await; + let mut kiro = state.kiro.write().await; + let needs_refresh = + kiro.credentials.access_token.is_none() || kiro.is_token_expiring_soon(); + if needs_refresh { + if let Err(e) = kiro.refresh_token().await { + return ( + StatusCode::UNAUTHORIZED, + Json(serde_json::json!({"error": {"message": format!("Token refresh failed: {e}")}})), + ) + .into_response(); + } + } + } + + let kiro = state.kiro.read().await; + match kiro.call_api(request).await { + Ok(resp) => { + let status = resp.status(); + if status.is_success() { + match resp.text().await { + Ok(body) => { + let parsed = parse_cw_response(&body); + let has_tool_calls = !parsed.tool_calls.is_empty(); + + let message = if has_tool_calls { + serde_json::json!({ + "role": "assistant", + "content": if parsed.content.is_empty() { serde_json::Value::Null } else { serde_json::json!(parsed.content) }, + "tool_calls": parsed.tool_calls.iter().map(|tc| { + serde_json::json!({ + "id": tc.id, + "type": "function", + "function": { + "name": tc.function.name, + "arguments": tc.function.arguments + } + }) + }).collect::>() + }) + } else { + serde_json::json!({ + "role": "assistant", + "content": parsed.content + }) + }; + + let response = serde_json::json!({ + "id": format!("chatcmpl-{}", uuid::Uuid::new_v4()), + "object": "chat.completion", + "created": std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_secs(), + "model": request.model, + "choices": [{ + "index": 0, + "message": message, + "finish_reason": if has_tool_calls { "tool_calls" } else { "stop" } + }], + "usage": { + "prompt_tokens": 0, + "completion_tokens": 0, + "total_tokens": 0 + } + }); + Json(response).into_response() + } + Err(e) => ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"error": {"message": e.to_string()}})), + ) + .into_response(), + } + } else { + let body = resp.text().await.unwrap_or_default(); + ( + StatusCode::from_u16(status.as_u16()).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR), + Json(serde_json::json!({"error": {"message": format!("Upstream error: {}", body)}})), + ) + .into_response() + } + } + Err(e) => ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"error": {"message": e.to_string()}})), + ) + .into_response(), + } +} + +use crate::models::provider_pool_model::ProviderCredential; diff --git a/src-tauri/src/server_utils.rs b/src-tauri/src/server_utils.rs new file mode 100644 index 000000000..ee3c8f790 --- /dev/null +++ b/src-tauri/src/server_utils.rs @@ -0,0 +1,656 @@ +//! 服务器工具函数 +//! +//! 包含响应解析、字符串处理、响应构建等公共工具函数。 + +use crate::models::openai::{ContentPart, FunctionCall, MessageContent, ToolCall}; +use axum::{ + body::Body, + http::{header, StatusCode}, + response::{IntoResponse, Response}, + Json, +}; +use futures::stream; +use std::collections::HashMap; + +/// CodeWhisperer 响应解析结果 +#[derive(Debug, Default)] +pub struct CWParsedResponse { + pub content: String, + pub tool_calls: Vec, + pub usage_credits: f64, + pub context_usage_percentage: f64, +} + +/// 安全截断字符串到指定字符数，避免 UTF-8 边界问题 +pub fn safe_truncate(s: &str, max_chars: usize) -> String { + let chars: Vec = s.chars().collect(); + if chars.len() <= max_chars { + s.to_string() + } else { + chars[..max_chars].iter().collect() + } +} + +/// 计算 MessageContent 的字符长度 +pub fn message_content_len(content: &MessageContent) -> usize { + match content { + MessageContent::Text(s) => s.len(), + MessageContent::Parts(parts) => parts + .iter() + .filter_map(|p| { + if let ContentPart::Text { text } = p { + Some(text.len()) + } else { + None + } + }) + .sum(), + } +} + +/// 解析 CodeWhisperer AWS Event Stream 响应 +/// +/// AWS Event Stream 是二进制格式，JSON payload 嵌入在二进制头部之间 +pub fn parse_cw_response(body: &str) -> CWParsedResponse { + let mut result = CWParsedResponse::default(); + // 使用 HashMap 来跟踪多个并发的 tool calls + // key: toolUseId, value: (name, input_accumulated) + let mut tool_map: HashMap = HashMap::new(); + + // 将字符串转换为字节，因为 AWS Event Stream 包含二进制数据 + let bytes = body.as_bytes(); + + // 搜索所有 JSON 对象的模式 + // AWS Event Stream 格式: [binary headers]{"content":"..."}[binary trailer] + let json_patterns: &[&[u8]] = &[ + b"{\"content\":", + b"{\"name\":", + b"{\"input\":", + b"{\"stop\":", + b"{\"followupPrompt\":", + b"{\"toolUseId\":", + b"{\"unit\":", // meteringEvent + b"{\"contextUsagePercentage\":", // contextUsageEvent + ]; + + let mut pos = 0; + while pos < bytes.len() { + // 找到下一个 JSON 对象的开始 + let mut next_start: Option = None; + + for pattern in json_patterns { + if let Some(idx) = find_subsequence(&bytes[pos..], pattern) { + let abs_pos = pos + idx; + if next_start.is_none_or(|start| abs_pos < start) { + next_start = Some(abs_pos); + } + } + } + + let start = match next_start { + Some(s) => s, + None => break, + }; + + // 从 start 位置提取完整的 JSON 对象 + if let Some(json_str) = extract_json_from_bytes(&bytes[start..]) { + if let Ok(value) = serde_json::from_str::(&json_str) { + // 处理 content 事件 + if let Some(content) = value.get("content").and_then(|v| v.as_str()) { + // 跳过 followupPrompt + if value.get("followupPrompt").is_none() { + result.content.push_str(content); + } + } + // 处理 tool use 事件 (包含 toolUseId) + else if let Some(tool_use_id) = value.get("toolUseId").and_then(|v| v.as_str()) { + let name = value + .get("name") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let input_chunk = value + .get("input") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let is_stop = value.get("stop").and_then(|v| v.as_bool()).unwrap_or(false); + + // 获取或创建 tool entry + let entry = tool_map + .entry(tool_use_id.to_string()) + .or_insert_with(|| (String::new(), String::new())); + + // 更新 name（如果有） + if !name.is_empty() { + entry.0 = name; + } + + // 累积 input + entry.1.push_str(&input_chunk); + + // 如果是 stop 事件，完成这个 tool call + if is_stop { + if let Some((name, input)) = tool_map.remove(tool_use_id) { + if !name.is_empty() { + result.tool_calls.push(ToolCall { + id: tool_use_id.to_string(), + call_type: "function".to_string(), + function: FunctionCall { + name, + arguments: input, + }, + }); + } + } + } + } + // 处理独立的 stop 事件（没有 toolUseId）- 这种情况不应该发生，但以防万一 + else if value.get("stop").and_then(|v| v.as_bool()).unwrap_or(false) { + // no-op + } + // 处理 meteringEvent: {"unit":"credit","unitPlural":"credits","usage":0.34} + else if let Some(usage) = value.get("usage").and_then(|v| v.as_f64()) { + result.usage_credits = usage; + } + // 处理 contextUsageEvent: {"contextUsagePercentage":54.36} + else if let Some(ctx_usage) = + value.get("contextUsagePercentage").and_then(|v| v.as_f64()) + { + result.context_usage_percentage = ctx_usage; + } + } + pos = start + json_str.len(); + } else { + pos = start + 1; + } + } + + // 处理未完成的 tool calls（没有收到 stop 事件的） + for (id, (name, input)) in tool_map { + if !name.is_empty() { + result.tool_calls.push(ToolCall { + id, + call_type: "function".to_string(), + function: FunctionCall { + name, + arguments: input, + }, + }); + } + } + + // 解析 bracket 格式的 tool calls: [Called xxx with args: {...}] + parse_bracket_tool_calls(&mut result); + + result +} + +/// 在字节数组中查找子序列 +pub fn find_subsequence(haystack: &[u8], needle: &[u8]) -> Option { + haystack + .windows(needle.len()) + .position(|window| window == needle) +} + +/// 从字节数组中提取 JSON 对象字符串 +pub fn extract_json_from_bytes(bytes: &[u8]) -> Option { + if bytes.is_empty() || bytes[0] != b'{' { + return None; + } + + let mut brace_count = 0; + let mut in_string = false; + let mut escape_next = false; + let mut end_pos = None; + + for (i, &b) in bytes.iter().enumerate() { + if escape_next { + escape_next = false; + continue; + } + + match b { + b'\\' if in_string => escape_next = true, + b'"' => in_string = !in_string, + b'{' if !in_string => brace_count += 1, + b'}' if !in_string => { + brace_count -= 1; + if brace_count == 0 { + end_pos = Some(i + 1); + break; + } + } + _ => {} + } + } + + end_pos.and_then(|end| String::from_utf8(bytes[..end].to_vec()).ok()) +} + +/// 从字符串中提取完整的 JSON 对象 (保留用于兼容) +#[allow(dead_code)] +pub fn extract_json_object(s: &str) -> Option<&str> { + if !s.starts_with('{') { + return None; + } + + let mut brace_count = 0; + let mut in_string = false; + let mut escape_next = false; + + for (i, c) in s.char_indices() { + if escape_next { + escape_next = false; + continue; + } + + match c { + '\\' if in_string => escape_next = true, + '"' => in_string = !in_string, + '{' if !in_string => brace_count += 1, + '}' if !in_string => { + brace_count -= 1; + if brace_count == 0 { + return Some(&s[..=i]); + } + } + _ => {} + } + } + None +} + +/// 解析 bracket 格式的 tool calls +/// +/// 格式: [Called xxx with args: {...}] +pub fn parse_bracket_tool_calls(result: &mut CWParsedResponse) { + let re = + regex::Regex::new(r"\[Called\s+(\w+)\s+with\s+args:\s*(\{[^}]*(?:\{[^}]*\}[^}]*)*\})\]") + .ok(); + + if let Some(re) = re { + let mut to_remove = Vec::new(); + for cap in re.captures_iter(&result.content) { + if let (Some(name), Some(args)) = (cap.get(1), cap.get(2)) { + let tool_id = format!( + "call_{}", + &uuid::Uuid::new_v4().to_string().replace('-', "")[..8] + ); + result.tool_calls.push(ToolCall { + id: tool_id, + call_type: "function".to_string(), + function: FunctionCall { + name: name.as_str().to_string(), + arguments: args.as_str().to_string(), + }, + }); + if let Some(full_match) = cap.get(0) { + to_remove.push(full_match.as_str().to_string()); + } + } + } + // 从 content 中移除 tool call 文本 + for s in to_remove { + result.content = result.content.replace(&s, ""); + } + result.content = result.content.trim().to_string(); + } +} + +/// 构建 Anthropic 非流式响应 +pub fn build_anthropic_response(model: &str, parsed: &CWParsedResponse) -> Response { + let has_tool_calls = !parsed.tool_calls.is_empty(); + let mut content_array: Vec = Vec::new(); + + if !parsed.content.is_empty() { + content_array.push(serde_json::json!({ + "type": "text", + "text": parsed.content + })); + } + + for tc in &parsed.tool_calls { + let input: serde_json::Value = + serde_json::from_str(&tc.function.arguments).unwrap_or(serde_json::json!({})); + content_array.push(serde_json::json!({ + "type": "tool_use", + "id": tc.id, + "name": tc.function.name, + "input": input + })); + } + + if content_array.is_empty() { + content_array.push(serde_json::json!({"type": "text", "text": ""})); + } + + // 估算 output tokens: 基于响应内容长度 (约 4 字符 = 1 token) + let mut output_tokens: u32 = (parsed.content.len() / 4) as u32; + for tc in &parsed.tool_calls { + output_tokens += (tc.function.arguments.len() / 4) as u32; + } + // 从 context_usage_percentage 估算 input tokens + // 假设 100% = 200k tokens (Claude 的上下文窗口) + let input_tokens = ((parsed.context_usage_percentage / 100.0) * 200000.0) as u32; + + let response = serde_json::json!({ + "id": format!("msg_{}", uuid::Uuid::new_v4()), + "type": "message", + "role": "assistant", + "content": content_array, + "model": model, + "stop_reason": if has_tool_calls { "tool_use" } else { "end_turn" }, + "stop_sequence": null, + "usage": { + "input_tokens": input_tokens, + "output_tokens": output_tokens + } + }); + Json(response).into_response() +} + +/// 构建 Anthropic 流式响应 (SSE) +pub fn build_anthropic_stream_response(model: &str, parsed: &CWParsedResponse) -> Response { + let has_tool_calls = !parsed.tool_calls.is_empty(); + let message_id = format!("msg_{}", uuid::Uuid::new_v4()); + let model = model.to_string(); + let content = parsed.content.clone(); + let tool_calls = parsed.tool_calls.clone(); + + // 估算 output tokens: 基于响应内容长度 (约 4 字符 = 1 token) + let mut output_tokens: u32 = (parsed.content.len() / 4) as u32; + for tc in &parsed.tool_calls { + output_tokens += (tc.function.arguments.len() / 4) as u32; + } + // 从 context_usage_percentage 估算 input tokens + let input_tokens = ((parsed.context_usage_percentage / 100.0) * 200000.0) as u32; + + // 构建 SSE 事件流 + let mut events: Vec = Vec::new(); + + // 1. message_start + let message_start = serde_json::json!({ + "type": "message_start", + "message": { + "id": message_id, + "type": "message", + "role": "assistant", + "model": model, + "content": [], + "stop_reason": null, + "stop_sequence": null, + "usage": {"input_tokens": input_tokens, "output_tokens": 0} + } + }); + events.push(format!("event: message_start\ndata: {message_start}\n\n")); + + let mut block_index = 0; + + // 2. 文本内容块 - 即使为空也要发送，Claude Code 需要至少一个 content block + // content_block_start + let block_start = serde_json::json!({ + "type": "content_block_start", + "index": block_index, + "content_block": {"type": "text", "text": ""} + }); + events.push(format!( + "event: content_block_start\ndata: {block_start}\n\n" + )); + + if !content.is_empty() { + // content_block_delta - 发送完整内容 + let block_delta = serde_json::json!({ + "type": "content_block_delta", + "index": block_index, + "delta": {"type": "text_delta", "text": content} + }); + events.push(format!( + "event: content_block_delta\ndata: {block_delta}\n\n" + )); + } + + // content_block_stop + let block_stop = serde_json::json!({ + "type": "content_block_stop", + "index": block_index + }); + events.push(format!("event: content_block_stop\ndata: {block_stop}\n\n")); + + block_index += 1; + + // 3. Tool use 块 + for tc in &tool_calls { + // content_block_start + let block_start = serde_json::json!({ + "type": "content_block_start", + "index": block_index, + "content_block": { + "type": "tool_use", + "id": tc.id, + "name": tc.function.name, + "input": {} + } + }); + events.push(format!( + "event: content_block_start\ndata: {block_start}\n\n" + )); + + // content_block_delta - input_json_delta + let partial_json = if tc.function.arguments.is_empty() { + "{}".to_string() + } else { + tc.function.arguments.clone() + }; + let block_delta = serde_json::json!({ + "type": "content_block_delta", + "index": block_index, + "delta": { + "type": "input_json_delta", + "partial_json": partial_json + } + }); + events.push(format!( + "event: content_block_delta\ndata: {block_delta}\n\n" + )); + + // content_block_stop + let block_stop = serde_json::json!({ + "type": "content_block_stop", + "index": block_index + }); + events.push(format!("event: content_block_stop\ndata: {block_stop}\n\n")); + + block_index += 1; + } + + // 4. message_delta + let message_delta = serde_json::json!({ + "type": "message_delta", + "delta": { + "stop_reason": if has_tool_calls { "tool_use" } else { "end_turn" }, + "stop_sequence": null + }, + "usage": {"output_tokens": output_tokens} + }); + events.push(format!("event: message_delta\ndata: {message_delta}\n\n")); + + // 5. message_stop + let message_stop = serde_json::json!({"type": "message_stop"}); + events.push(format!("event: message_stop\ndata: {message_stop}\n\n")); + + // 创建 SSE 响应 + let body_stream = stream::iter(events.into_iter().map(Ok::<_, std::convert::Infallible>)); + let body = Body::from_stream(body_stream); + + Response::builder() + .status(StatusCode::OK) + .header(header::CONTENT_TYPE, "text/event-stream") + .header(header::CACHE_CONTROL, "no-cache") + .header(header::CONNECTION, "keep-alive") + .body(body) + .unwrap_or_else(|e| { + tracing::error!("Failed to build SSE response: {}", e); + Response::builder() + .status(StatusCode::INTERNAL_SERVER_ERROR) + .body(Body::empty()) + .unwrap_or_default() + }) +} + +/// 构建 Gemini 原生请求体 +/// +/// 将用户传入的 Gemini 格式请求转换为 Antigravity 请求格式 +pub fn build_gemini_native_request( + request: &serde_json::Value, + model: &str, + project_id: &str, +) -> serde_json::Value { + // 模型名称映射 + let actual_model = match model { + "gemini-2.5-computer-use-preview-10-2025" => "rev19-uic3-1p", + "gemini-3-pro-image-preview" => "gemini-3-pro-image", + "gemini-3-pro-preview" => "gemini-3-pro-high", + "gemini-claude-sonnet-4-5" => "claude-sonnet-4-5", + "gemini-claude-sonnet-4-5-thinking" => "claude-sonnet-4-5-thinking", + _ => model, + }; + + // 是否启用思维链 + let enable_thinking = model.ends_with("-thinking") + || model == "gemini-2.5-pro" + || model.starts_with("gemini-3-pro-") + || model == "rev19-uic3-1p" + || model == "gpt-oss-120b-medium"; + + // 生成请求 ID 和会话 ID + let request_id = format!("agent-{}", uuid::Uuid::new_v4()); + let session_id = { + let uuid = uuid::Uuid::new_v4(); + let bytes = uuid.as_bytes(); + let n: u64 = u64::from_le_bytes([ + bytes[0], bytes[1], bytes[2], bytes[3], bytes[4], bytes[5], bytes[6], bytes[7], + ]) % 9_000_000_000_000_000_000; + format!("-{}", n) + }; + + // 构建内部请求 + let mut inner_request = request.clone(); + + // 添加会话 ID + inner_request["sessionId"] = serde_json::json!(session_id); + + // 确保有 generationConfig + if inner_request.get("generationConfig").is_none() { + inner_request["generationConfig"] = serde_json::json!({ + "temperature": 1.0, + "maxOutputTokens": 8096, + "topP": 0.85, + "topK": 50, + "candidateCount": 1, + "stopSequences": [ + "<|user|>", + "<|bot|>", + "<|context_request|>", + "<|endoftext|>", + "<|end_of_turn|>" + ], + "thinkingConfig": { + "includeThoughts": enable_thinking, + "thinkingBudget": if enable_thinking { 1024 } else { 0 } + } + }); + } else { + // 确保有 thinkingConfig + if inner_request["generationConfig"] + .get("thinkingConfig") + .is_none() + { + inner_request["generationConfig"]["thinkingConfig"] = serde_json::json!({ + "includeThoughts": enable_thinking, + "thinkingBudget": if enable_thinking { 1024 } else { 0 } + }); + } + } + + // 删除安全设置（Antigravity 不支持） + if let Some(obj) = inner_request.as_object_mut() { + obj.remove("safetySettings"); + } + + // 构建完整的 Antigravity 请求体 + serde_json::json!({ + "project": project_id, + "requestId": request_id, + "request": inner_request, + "model": actual_model, + "userAgent": "antigravity" + }) +} + +/// 健康检查端点响应 +pub async fn health() -> impl IntoResponse { + Json(serde_json::json!({ + "status": "healthy", + "version": env!("CARGO_PKG_VERSION") + })) +} + +/// 模型列表端点响应 +pub async fn models() -> impl IntoResponse { + Json(serde_json::json!({ + "object": "list", + "data": [ + // Kiro/Claude models + {"id": "claude-sonnet-4-5", "object": "model", "owned_by": "anthropic"}, + {"id": "claude-sonnet-4-5-20250929", "object": "model", "owned_by": "anthropic"}, + {"id": "claude-3-7-sonnet-20250219", "object": "model", "owned_by": "anthropic"}, + {"id": "claude-3-5-sonnet-latest", "object": "model", "owned_by": "anthropic"}, + // Gemini models + {"id": "gemini-2.5-flash", "object": "model", "owned_by": "google"}, + {"id": "gemini-2.5-flash-lite", "object": "model", "owned_by": "google"}, + {"id": "gemini-2.5-pro", "object": "model", "owned_by": "google"}, + {"id": "gemini-2.5-pro-preview-06-05", "object": "model", "owned_by": "google"}, + {"id": "gemini-3-pro-preview", "object": "model", "owned_by": "google"}, + // Qwen models + {"id": "qwen3-coder-plus", "object": "model", "owned_by": "alibaba"}, + {"id": "qwen3-coder-flash", "object": "model", "owned_by": "alibaba"} + ] + })) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_safe_truncate() { + assert_eq!(safe_truncate("hello", 10), "hello"); + assert_eq!(safe_truncate("hello world", 5), "hello"); + assert_eq!(safe_truncate("你好世界", 2), "你好"); + } + + #[test] + fn test_find_subsequence() { + let haystack = b"hello world"; + assert_eq!(find_subsequence(haystack, b"world"), Some(6)); + assert_eq!(find_subsequence(haystack, b"foo"), None); + } + + #[test] + fn test_extract_json_from_bytes() { + let json = b"{\"key\":\"value\"}"; + assert_eq!( + extract_json_from_bytes(json), + Some("{\"key\":\"value\"}".to_string()) + ); + + let nested = b"{\"outer\":{\"inner\":\"value\"}}"; + assert_eq!( + extract_json_from_bytes(nested), + Some("{\"outer\":{\"inner\":\"value\"}}".to_string()) + ); + + assert_eq!(extract_json_from_bytes(b"not json"), None); + } +} diff --git a/src-tauri/src/services/provider_pool_service.rs b/src-tauri/src/services/provider_pool_service.rs index 7c916fd2b..40c64950f 100644 --- a/src-tauri/src/services/provider_pool_service.rs +++ b/src-tauri/src/services/provider_pool_service.rs @@ -300,6 +300,8 @@ impl ProviderPoolService { } /// 执行单个凭证的健康检查 + /// + /// 如果遇到 401 错误，会自动尝试刷新 token 后重试 pub async fn check_credential_health( &self, db: &DbConnection, @@ -335,6 +337,71 @@ impl ProviderPoolService { }) } Err(e) => { + // 如果是 401 错误，尝试刷新 token 后重试 + if e.contains("401") || e.contains("Unauthorized") { + tracing::info!("[健康检查] 检测到 401 错误，尝试刷新 token: {}", uuid); + + // 尝试刷新 token + match self.refresh_credential_token(db, uuid).await { + Ok(_) => { + tracing::info!("[健康检查] Token 刷新成功，重新检查健康状态"); + + // 重新获取凭证（token 已更新） + let updated_cred = { + let conn = db.lock().map_err(|e| e.to_string())?; + ProviderPoolDao::get_by_uuid(&conn, uuid) + .map_err(|e| e.to_string())? + .ok_or_else(|| format!("Credential not found: {}", uuid))? + }; + + // 重新执行健康检查 + let retry_start = std::time::Instant::now(); + let retry_result = self + .perform_health_check(&updated_cred.credential, &check_model) + .await; + let retry_duration_ms = retry_start.elapsed().as_millis() as u64; + + match retry_result { + Ok(_) => { + self.mark_healthy(db, uuid, Some(&check_model))?; + return Ok(HealthCheckResult { + uuid: uuid.to_string(), + success: true, + model: Some(check_model), + message: Some( + "Health check passed after token refresh".to_string(), + ), + duration_ms: duration_ms + retry_duration_ms, + }); + } + Err(retry_e) => { + tracing::warn!("[健康检查] Token 刷新后仍然失败: {}", retry_e); + self.mark_unhealthy(db, uuid, Some(&retry_e))?; + return Ok(HealthCheckResult { + uuid: uuid.to_string(), + success: false, + model: Some(check_model), + message: Some(retry_e), + duration_ms: duration_ms + retry_duration_ms, + }); + } + } + } + Err(refresh_err) => { + tracing::warn!("[健康检查] Token 刷新失败: {}", refresh_err); + // Token 刷新失败，返回原始错误 + self.mark_unhealthy(db, uuid, Some(&e))?; + return Ok(HealthCheckResult { + uuid: uuid.to_string(), + success: false, + model: Some(check_model), + message: Some(format!("{} (Token 刷新失败: {})", e, refresh_err)), + duration_ms, + }); + } + } + } + self.mark_unhealthy(db, uuid, Some(&e))?; Ok(HealthCheckResult { uuid: uuid.to_string(), @@ -550,11 +617,13 @@ impl ProviderPoolService { } // Gemini OAuth 健康检查 + // 使用 cloudcode-pa.googleapis.com API（与 Gemini CLI 兼容） + // 使用 loadCodeAssist 接口进行健康检查，这是最简单可靠的方式 async fn check_gemini_health( &self, creds_path: &str, _project_id: Option<&str>, - model: &str, + _model: &str, ) -> Result<(), String> { let creds_content = std::fs::read_to_string(creds_path).map_err(|e| format!("读取凭证文件失败: {}", e))?; @@ -565,24 +634,25 @@ impl ProviderPoolService { .as_str() .ok_or_else(|| "凭证中缺少 access_token".to_string())?; - let url = format!( - "https://generativelanguage.googleapis.com/v1beta/models/{}:generateContent", - model - ); + // 使用 loadCodeAssist 接口进行健康检查 + // 这个接口用于获取项目信息，是最简单可靠的健康检查方式 + let url = "https://cloudcode-pa.googleapis.com/v1internal:loadCodeAssist"; let request_body = serde_json::json!({ - "contents": [{ - "parts": [{"text": "Say OK"}] - }], - "generationConfig": { - "maxOutputTokens": 10 + "cloudaicompanionProject": "", + "metadata": { + "ideType": "IDE_UNSPECIFIED", + "platform": "PLATFORM_UNSPECIFIED", + "pluginType": "GEMINI", + "duetProject": "" } }); let response = self .client - .post(&url) + .post(url) .bearer_auth(access_token) + .header("Content-Type", "application/json") .json(&request_body) .timeout(self.health_check_timeout) .send() @@ -592,7 +662,9 @@ impl ProviderPoolService { if response.status().is_success() { Ok(()) } else { - Err(format!("HTTP {}", response.status())) + let status = response.status(); + let body = response.text().await.unwrap_or_default(); + Err(format!("HTTP {} - {}", status, body)) } } @@ -607,16 +679,31 @@ impl ProviderPoolService { .as_str() .ok_or_else(|| "凭证中缺少 access_token".to_string())?; + // 获取 base_url，优先使用 resource_url，否则使用默认值 + let base_url = if let Some(resource_url) = creds["resource_url"].as_str() { + if resource_url.starts_with("http") { + format!("{}/v1", resource_url.trim_end_matches('/')) + } else { + format!("https://{}/v1", resource_url) + } + } else { + "https://portal.qwen.ai/v1".to_string() + }; + let request_body = serde_json::json!({ "model": model, "messages": [{"role": "user", "content": "Say OK"}], "max_tokens": 10 }); + let url = format!("{}/chat/completions", base_url); + let response = self .client - .post("https://chat.qwen.ai/api/v1/chat/completions") + .post(&url) .bearer_auth(access_token) + .header("Content-Type", "application/json") + .header("Accept", "application/json") .json(&request_body) .timeout(self.health_check_timeout) .send() diff --git a/src-tauri/tauri.conf.json b/src-tauri/tauri.conf.json index ddf44545a..ee03e3378 100644 --- a/src-tauri/tauri.conf.json +++ b/src-tauri/tauri.conf.json @@ -1,7 +1,7 @@ { "$schema": "https://schema.tauri.app/config/2", "productName": "ProxyCast", - "version": "0.14.6", + "version": "0.14.10", "identifier": "com.proxycast.app", "build": { "beforeDevCommand": "npm run dev", diff --git a/src/components/api-server/ApiServerPage.tsx b/src/components/api-server/ApiServerPage.tsx index 434407f6f..72d09685a 100644 --- a/src/components/api-server/ApiServerPage.tsx +++ b/src/components/api-server/ApiServerPage.tsx @@ -239,6 +239,45 @@ export function ApiServerPage() { : `http://localhost:${config?.server.port ?? 8999}`; const apiKey = config?.server.api_key ?? ""; + // 根据 Provider 类型获取测试模型 + const getTestModel = (provider: string): string => { + switch (provider) { + case "antigravity": + return "gemini-3-pro-preview"; + case "gemini": + return "gemini-2.0-flash"; + case "qwen": + return "qwen-max"; + case "openai": + return "gpt-4o"; + case "claude": + return "claude-sonnet-4-20250514"; + case "kiro": + default: + return "claude-opus-4-5-20251101"; + } + }; + + const testModel = getTestModel(defaultProvider); + + // 根据 Provider 类型获取 Gemini 测试模型 + const getGeminiTestModel = (provider: string): string => { + switch (provider) { + case "antigravity": + return "gemini-3-pro-preview"; + case "gemini": + return "gemini-2.0-flash"; + default: + return "gemini-2.0-flash"; + } + }; + + const geminiTestModel = getGeminiTestModel(defaultProvider); + + // 是否显示 Gemini 测试端点 + const showGeminiTest = + defaultProvider === "antigravity" || defaultProvider === "gemini"; + // Test endpoints const testEndpoints = [ { @@ -264,7 +303,7 @@ export function ApiServerPage() { path: "/v1/chat/completions", needsAuth: true, body: JSON.stringify({ - model: "claude-opus-4-5-20251101", + model: testModel, messages: [{ role: "user", content: "Say hi in one word" }], }), }, @@ -275,7 +314,7 @@ export function ApiServerPage() { path: "/v1/messages", needsAuth: true, body: JSON.stringify({ - model: "claude-opus-4-5-20251101", + model: testModel, max_tokens: 100, messages: [ { @@ -285,6 +324,31 @@ export function ApiServerPage() { ], }), }, + // Gemini 原生协议测试（仅在 Antigravity 或 Gemini Provider 时显示） + ...(showGeminiTest + ? [ + { + id: "gemini", + name: "Gemini Generate", + method: "POST", + path: `/v1/gemini/${geminiTestModel}:generateContent`, + needsAuth: true, + body: JSON.stringify({ + contents: [ + { + role: "user", + parts: [ + { text: "What is 2+2? Answer with just the number." }, + ], + }, + ], + generationConfig: { + maxOutputTokens: 100, + }, + }), + }, + ] + : []), ]; const runTest = async (endpoint: (typeof testEndpoints)[0]) => { diff --git a/src/components/provider-pool/AddCredentialModal.tsx b/src/components/provider-pool/AddCredentialModal.tsx index b4cb89d71..d5c35f187 100644 --- a/src/components/provider-pool/AddCredentialModal.tsx +++ b/src/components/provider-pool/AddCredentialModal.tsx @@ -1,7 +1,19 @@ +/** + * 添加凭证模态框 + * 根据 Provider 类型显示不同的表单 + */ + import { useState } from "react"; import { X, Key, FolderOpen } from "lucide-react"; import { open } from "@tauri-apps/plugin-dialog"; import { providerPoolApi, PoolProviderType } from "@/lib/api/providerPool"; +import { AntigravityForm } from "./credential-forms/AntigravityForm"; +import { CodexForm } from "./credential-forms/CodexForm"; +import { ClaudeOAuthForm } from "./credential-forms/ClaudeOAuthForm"; +import { QwenForm } from "./credential-forms/QwenForm"; +import { IFlowForm } from "./credential-forms/IFlowForm"; +import { GeminiForm } from "./credential-forms/GeminiForm"; +import { defaultCredsPath, providerLabels } from "./credential-forms/types"; interface AddCredentialModalProps { providerType: PoolProviderType; @@ -9,17 +21,6 @@ interface AddCredentialModalProps { onSuccess: () => void; } -// Default credential paths -const defaultCredsPath: Record = { - kiro: "~/.aws/sso/cache/kiro-auth-token.json", - gemini: "~/.gemini/oauth_creds.json", - qwen: "~/.qwen/oauth_creds.json", - antigravity: "~/.antigravity/oauth_creds.json", - codex: "~/.codex/oauth.json", - claude_oauth: "~/.claude/oauth.json", - iflow: "~/.iflow/oauth_creds.json", -}; - export function AddCredentialModal({ providerType, onClose, @@ -29,37 +30,19 @@ export function AddCredentialModal({ const [loading, setLoading] = useState(false); const [error, setError] = useState(null); - // OAuth fields - initialize with default path + // OAuth 字段 const [credsFilePath, setCredsFilePath] = useState( defaultCredsPath[providerType] || "", ); const [projectId, setProjectId] = useState(""); - // API Key fields + // API Key 字段 const [apiKey, setApiKey] = useState(""); const [baseUrl, setBaseUrl] = useState(""); - const isOAuth = [ - "kiro", - "gemini", - "qwen", - "antigravity", - "codex", - "claude_oauth", - "iflow", - ].includes(providerType); - - const providerLabels: Record = { - kiro: "Kiro (AWS)", - gemini: "Gemini (Google)", - qwen: "Qwen (阿里)", - openai: "OpenAI", - claude: "Claude (Anthropic)", - antigravity: "Antigravity (Gemini 3 Pro)", - codex: "Codex (OpenAI OAuth)", - claude_oauth: "Claude OAuth", - iflow: "iFlow", - }; + // 判断是否为 OAuth 类型（不包括有特殊表单的 antigravity、codex、claude_oauth、qwen、iflow、gemini） + const isSimpleOAuth = ["kiro"].includes(providerType); + const isApiKey = ["openai", "claude"].includes(providerType); const handleSelectFile = async () => { try { @@ -75,6 +58,83 @@ export function AddCredentialModal({ } }; + // Antigravity 表单 + const antigravityForm = AntigravityForm({ + name, + credsFilePath, + setCredsFilePath, + projectId, + setProjectId, + onSelectFile: handleSelectFile, + loading, + setLoading, + setError, + onSuccess, + }); + + // Codex 表单 + const codexForm = CodexForm({ + name, + credsFilePath, + setCredsFilePath, + onSelectFile: handleSelectFile, + loading, + setLoading, + setError, + onSuccess, + }); + + // Claude OAuth 表单 + const claudeOAuthForm = ClaudeOAuthForm({ + name, + credsFilePath, + setCredsFilePath, + onSelectFile: handleSelectFile, + loading, + setLoading, + setError, + onSuccess, + }); + + // Qwen 表单 + const qwenForm = QwenForm({ + name, + credsFilePath, + setCredsFilePath, + onSelectFile: handleSelectFile, + loading, + setLoading, + setError, + onSuccess, + }); + + // iFlow 表单 + const iflowForm = IFlowForm({ + name, + credsFilePath, + setCredsFilePath, + onSelectFile: handleSelectFile, + loading, + setLoading, + setError, + onSuccess, + }); + + // Gemini 表单 + const geminiForm = GeminiForm({ + name, + credsFilePath, + setCredsFilePath, + projectId, + setProjectId, + onSelectFile: handleSelectFile, + loading, + setLoading, + setError, + onSuccess, + }); + + // 简单 OAuth 和 API Key 的提交处理 const handleSubmit = async () => { setLoading(true); setError(null); @@ -82,9 +142,10 @@ export function AddCredentialModal({ try { const trimmedName = name.trim() || undefined; - if (isOAuth) { + if (isSimpleOAuth) { if (!credsFilePath) { setError("请选择凭证文件"); + setLoading(false); return; } @@ -99,29 +160,11 @@ export function AddCredentialModal({ trimmedName, ); break; - case "qwen": - await providerPoolApi.addQwenOAuth(credsFilePath, trimmedName); - break; - case "antigravity": - await providerPoolApi.addAntigravityOAuth( - credsFilePath, - projectId.trim() || undefined, - trimmedName, - ); - break; - case "codex": - await providerPoolApi.addCodexOAuth(credsFilePath, trimmedName); - break; - case "claude_oauth": - await providerPoolApi.addClaudeOAuth(credsFilePath, trimmedName); - break; - case "iflow": - await providerPoolApi.addIFlowOAuth(credsFilePath, trimmedName); - break; } - } else { + } else if (isApiKey) { if (!apiKey) { setError("请输入 API Key"); + setLoading(false); return; } @@ -151,6 +194,283 @@ export function AddCredentialModal({ } }; + // 渲染简单 OAuth 表单 + const renderSimpleOAuthForm = () => ( + <> +
+ +
+ setCredsFilePath(e.target.value)} + placeholder="输入凭证文件的完整路径..." + className="flex-1 rounded-lg border bg-background px-3 py-2 text-sm" + /> + +
+

+ {providerType === "kiro" && + "默认路径: ~/.aws/sso/cache/kiro-auth-token.json"} + {providerType === "gemini" && "默认路径: ~/.gemini/oauth_creds.json"} +

+
+ + {providerType === "gemini" && ( +
+ + setProjectId(e.target.value)} + placeholder="Google Cloud Project ID..." + className="w-full rounded-lg border bg-background px-3 py-2 text-sm" + /> +
+ )} + + ); + + // 渲染 API Key 表单 + const renderApiKeyForm = () => ( + <> +
+ +
+ + setApiKey(e.target.value)} + placeholder="sk-..." + className="w-full rounded-lg border bg-background pl-10 pr-3 py-2 text-sm" + /> +
+
+ +
+ + setBaseUrl(e.target.value)} + placeholder={ + providerType === "openai" + ? "https://api.openai.com/v1" + : "https://api.anthropic.com/v1" + } + className="w-full rounded-lg border bg-background px-3 py-2 text-sm" + /> +

+ 留空使用默认 URL，或输入自定义代理地址 +

+
+ + ); + + // 渲染底部按钮 + const renderFooterButton = () => { + // Antigravity 登录模式 + if (providerType === "antigravity" && antigravityForm.mode === "login") { + if (!antigravityForm.authUrl) { + return ( + + ); + } + return null; + } + + // Antigravity 文件模式 + if (providerType === "antigravity" && antigravityForm.mode === "file") { + return ( + + ); + } + + // Codex 登录模式 + if (providerType === "codex" && codexForm.mode === "login") { + if (!codexForm.authUrl) { + return ( + + ); + } + return null; + } + + // Codex 文件模式 + if (providerType === "codex" && codexForm.mode === "file") { + return ( + + ); + } + + // Claude OAuth 登录模式 + if (providerType === "claude_oauth" && claudeOAuthForm.mode === "login") { + if (!claudeOAuthForm.authUrl) { + return ( + + ); + } + return null; + } + + // Claude OAuth 文件模式 + if (providerType === "claude_oauth" && claudeOAuthForm.mode === "file") { + return ( + + ); + } + + // Qwen 登录模式 + if (providerType === "qwen" && qwenForm.mode === "login") { + if (!qwenForm.deviceCode) { + return ( + + ); + } + return null; + } + + // Qwen 文件模式 + if (providerType === "qwen" && qwenForm.mode === "file") { + return ( + + ); + } + + // iFlow 登录模式 + if (providerType === "iflow" && iflowForm.mode === "login") { + if (!iflowForm.authUrl) { + return ( + + ); + } + return null; + } + + // iFlow 文件模式 + if (providerType === "iflow" && iflowForm.mode === "file") { + return ( + + ); + } + + // Gemini 登录模式 + if (providerType === "gemini" && geminiForm.mode === "login") { + if (!geminiForm.authUrl) { + return ( + + ); + } + return null; + } + + // Gemini 文件模式 + if (providerType === "gemini" && geminiForm.mode === "file") { + return ( + + ); + } + + // 其他类型 + return ( + + ); + }; + return (
@@ -166,7 +486,7 @@ export function AddCredentialModal({ {/* Content */}
- {/* Name field */} + {/* 名称字段 */}
- {isOAuth ? ( - <> - {/* Credential File */} -
- -
- setCredsFilePath(e.target.value)} - placeholder="输入凭证文件的完整路径..." - className="flex-1 rounded-lg border bg-background px-3 py-2 text-sm" - /> - -
-

- {providerType === "kiro" && - "默认路径: ~/.aws/sso/cache/kiro-auth-token.json"} - {providerType === "gemini" && - "默认路径: ~/.gemini/oauth_creds.json"} - {providerType === "qwen" && - "默认路径: ~/.qwen/oauth_creds.json"} - {providerType === "antigravity" && - "默认路径: ~/.antigravity/oauth_creds.json"} - {providerType === "codex" && "默认路径: ~/.codex/oauth.json"} - {providerType === "claude_oauth" && - "默认路径: ~/.claude/oauth.json"} - {providerType === "iflow" && - "默认路径: ~/.iflow/oauth_creds.json"} -

-
+ {/* 根据类型渲染不同表单 */} + {providerType === "antigravity" && antigravityForm.render()} + {providerType === "codex" && codexForm.render()} + {providerType === "claude_oauth" && claudeOAuthForm.render()} + {providerType === "qwen" && qwenForm.render()} + {providerType === "iflow" && iflowForm.render()} + {providerType === "gemini" && geminiForm.render()} + {isSimpleOAuth && renderSimpleOAuthForm()} + {isApiKey && renderApiKeyForm()} - {/* Gemini/Antigravity specific: Project ID */} - {(providerType === "gemini" || - providerType === "antigravity") && ( -
- - setProjectId(e.target.value)} - placeholder="Google Cloud Project ID..." - className="w-full rounded-lg border bg-background px-3 py-2 text-sm" - /> -
- )} - - ) : ( - <> - {/* API Key */} -
- -
- - setApiKey(e.target.value)} - placeholder="sk-..." - className="w-full rounded-lg border bg-background pl-10 pr-3 py-2 text-sm" - /> -
-
- - {/* Base URL */} -
- - setBaseUrl(e.target.value)} - placeholder={ - providerType === "openai" - ? "https://api.openai.com/v1" - : "https://api.anthropic.com/v1" - } - className="w-full rounded-lg border bg-background px-3 py-2 text-sm" - /> -

- 留空使用默认 URL，或输入自定义代理地址 -

-
- - )} - - {/* Error */} + {/* 错误提示 */} {error && (
{error} @@ -296,13 +526,7 @@ export function AddCredentialModal({ > 取消 - + {renderFooterButton()}
diff --git a/src/components/provider-pool/CredentialCard.tsx b/src/components/provider-pool/CredentialCard.tsx index edcd5b052..dc22b6521 100644 --- a/src/components/provider-pool/CredentialCard.tsx +++ b/src/components/provider-pool/CredentialCard.tsx @@ -17,11 +17,18 @@ import { Globe, BarChart3, ChevronUp, + Fingerprint, + Copy, + Check, } from "lucide-react"; import type { CredentialDisplay, CredentialSource, } from "@/lib/api/providerPool"; +import { + getKiroCredentialFingerprint, + type KiroFingerprintInfo, +} from "@/lib/api/providerPool"; import { usageApi, type UsageInfo } from "@/lib/api/usage"; import { UsageDisplay } from "./UsageDisplay"; @@ -59,6 +66,46 @@ export function CredentialCard({ const [usageInfo, setUsageInfo] = useState(null); const [usageError, setUsageError] = useState(null); + // 指纹信息状态（仅 Kiro 凭证） + const [fingerprintInfo, setFingerprintInfo] = + useState(null); + const [fingerprintLoading, setFingerprintLoading] = useState(false); + const [fingerprintExpanded, setFingerprintExpanded] = useState(false); + const [fingerprintCopied, setFingerprintCopied] = useState(false); + + // 查询指纹信息 + const handleCheckFingerprint = async () => { + if (fingerprintExpanded && fingerprintInfo) { + // 已展开且有数据，直接折叠 + setFingerprintExpanded(false); + return; + } + + setFingerprintExpanded(true); + setFingerprintLoading(true); + + try { + const info = await getKiroCredentialFingerprint(credential.uuid); + setFingerprintInfo(info); + } catch (e) { + console.error("获取指纹信息失败:", e); + } finally { + setFingerprintLoading(false); + } + }; + + // 复制 Machine ID + const handleCopyMachineId = async () => { + if (!fingerprintInfo) return; + try { + await navigator.clipboard.writeText(fingerprintInfo.machine_id); + setFingerprintCopied(true); + setTimeout(() => setFingerprintCopied(false), 2000); + } catch (e) { + console.error("复制失败:", e); + } + }; + // 查询用量 const handleCheckUsage = async () => { if (usageExpanded && usageInfo) { @@ -295,6 +342,24 @@ export function CredentialCard({ )} + {/* 指纹信息按钮 - 仅 Kiro 凭证显示 */} + {isKiroCredential && ( + + )} + {/* 用量查询按钮 - 仅 Kiro 凭证显示 */} {isKiroCredential && ( +
+ + {fingerprintLoading ? ( +
+
+ 加载中... +
+ ) : fingerprintInfo ? ( +
+
+ + Machine ID: + + + {fingerprintInfo.machine_id_short}... + + +
+
+ + 来源: + + {fingerprintInfo.source} + + + + 认证: + + {fingerprintInfo.auth_method} + + +
+
+ ) : ( +
+ 无法获取指纹信息 +
+ )} +
+ )} + {/* 用量信息展示区域 - 仅 Kiro 凭证 */} {isKiroCredential && usageExpanded && (
diff --git a/src/components/provider-pool/credential-forms/AntigravityForm.tsx b/src/components/provider-pool/credential-forms/AntigravityForm.tsx new file mode 100644 index 000000000..51bd01e09 --- /dev/null +++ b/src/components/provider-pool/credential-forms/AntigravityForm.tsx @@ -0,0 +1,155 @@ +/** + * Antigravity 凭证添加表单 + * 支持 Google OAuth 登录和文件导入两种模式 + */ + +import { useState, useEffect } from "react"; +import { listen } from "@tauri-apps/api/event"; +import { providerPoolApi } from "@/lib/api/providerPool"; +import { ModeSelector } from "./ModeSelector"; +import { FileImportForm } from "./FileImportForm"; +import { OAuthUrlDisplay } from "./OAuthUrlDisplay"; + +interface AntigravityFormProps { + name: string; + credsFilePath: string; + setCredsFilePath: (path: string) => void; + projectId: string; + setProjectId: (id: string) => void; + onSelectFile: () => void; + loading: boolean; + setLoading: (loading: boolean) => void; + setError: (error: string | null) => void; + onSuccess: () => void; +} + +export function AntigravityForm({ + name, + credsFilePath, + setCredsFilePath, + projectId, + setProjectId, + onSelectFile, + loading: _loading, + setLoading, + setError, + onSuccess, +}: AntigravityFormProps) { + const [mode, setMode] = useState<"login" | "file">("login"); + const [authUrl, setAuthUrl] = useState(null); + const [waitingForCallback, setWaitingForCallback] = useState(false); + + // 监听后端发送的授权 URL 事件 + useEffect(() => { + let unlisten: (() => void) | undefined; + + const setupListener = async () => { + unlisten = await listen<{ auth_url: string }>( + "antigravity-auth-url", + (event) => { + setAuthUrl(event.payload.auth_url); + }, + ); + }; + + setupListener(); + + return () => { + if (unlisten) unlisten(); + }; + }, []); + + // 获取授权 URL 并启动服务器等待回调 + const handleGetAuthUrl = async () => { + setLoading(true); + setError(null); + setAuthUrl(null); + setWaitingForCallback(true); + + try { + const trimmedName = name.trim() || undefined; + await providerPoolApi.getAntigravityAuthUrlAndWait(trimmedName, false); + onSuccess(); + } catch (e) { + const errorMsg = e instanceof Error ? e.message : String(e); + setError(errorMsg); + setWaitingForCallback(false); + } finally { + setLoading(false); + } + }; + + // 文件导入提交 + const handleFileSubmit = async () => { + if (!credsFilePath) { + setError("请选择凭证文件"); + return; + } + + setLoading(true); + setError(null); + + try { + const trimmedName = name.trim() || undefined; + await providerPoolApi.addAntigravityOAuth( + credsFilePath, + projectId.trim() || undefined, + trimmedName, + ); + onSuccess(); + } catch (e) { + setError(e instanceof Error ? e.message : String(e)); + } finally { + setLoading(false); + } + }; + + return { + mode, + authUrl, + waitingForCallback, + handleGetAuthUrl, + handleFileSubmit, + render: () => ( + <> + + + {mode === "login" ? ( +
+
+

+ 点击下方按钮获取授权 + URL，然后复制到浏览器（支持指纹浏览器）完成登录。 +

+

+ 授权成功后，凭证将自动保存并添加到凭证池。 +

+
+ + +
+ ) : ( + + )} + + ), + }; +} diff --git a/src/components/provider-pool/credential-forms/ClaudeOAuthForm.tsx b/src/components/provider-pool/credential-forms/ClaudeOAuthForm.tsx new file mode 100644 index 000000000..7e3dc9f77 --- /dev/null +++ b/src/components/provider-pool/credential-forms/ClaudeOAuthForm.tsx @@ -0,0 +1,144 @@ +/** + * Claude OAuth 凭证添加表单 + * 支持 Claude OAuth 登录和文件导入两种模式 + */ + +import { useState, useEffect } from "react"; +import { listen } from "@tauri-apps/api/event"; +import { providerPoolApi } from "@/lib/api/providerPool"; +import { ModeSelector } from "./ModeSelector"; +import { FileImportForm } from "./FileImportForm"; +import { OAuthUrlDisplay } from "./OAuthUrlDisplay"; + +interface ClaudeOAuthFormProps { + name: string; + credsFilePath: string; + setCredsFilePath: (path: string) => void; + onSelectFile: () => void; + loading: boolean; + setLoading: (loading: boolean) => void; + setError: (error: string | null) => void; + onSuccess: () => void; +} + +export function ClaudeOAuthForm({ + name, + credsFilePath, + setCredsFilePath, + onSelectFile, + loading: _loading, + setLoading, + setError, + onSuccess, +}: ClaudeOAuthFormProps) { + const [mode, setMode] = useState<"login" | "file">("login"); + const [authUrl, setAuthUrl] = useState(null); + const [waitingForCallback, setWaitingForCallback] = useState(false); + + // 监听后端发送的授权 URL 事件 + useEffect(() => { + let unlisten: (() => void) | undefined; + + const setupListener = async () => { + unlisten = await listen<{ auth_url: string }>( + "claude-oauth-auth-url", + (event) => { + setAuthUrl(event.payload.auth_url); + }, + ); + }; + + setupListener(); + + return () => { + if (unlisten) unlisten(); + }; + }, []); + + // 获取授权 URL 并启动服务器等待回调 + const handleGetAuthUrl = async () => { + setLoading(true); + setError(null); + setAuthUrl(null); + setWaitingForCallback(true); + + try { + const trimmedName = name.trim() || undefined; + await providerPoolApi.getClaudeOAuthAuthUrlAndWait(trimmedName); + onSuccess(); + } catch (e) { + const errorMsg = e instanceof Error ? e.message : String(e); + setError(errorMsg); + setWaitingForCallback(false); + } finally { + setLoading(false); + } + }; + + // 文件导入提交 + const handleFileSubmit = async () => { + if (!credsFilePath) { + setError("请选择凭证文件"); + return; + } + + setLoading(true); + setError(null); + + try { + const trimmedName = name.trim() || undefined; + await providerPoolApi.addClaudeOAuth(credsFilePath, trimmedName); + onSuccess(); + } catch (e) { + setError(e instanceof Error ? e.message : String(e)); + } finally { + setLoading(false); + } + }; + + return { + mode, + authUrl, + waitingForCallback, + handleGetAuthUrl, + handleFileSubmit, + render: () => ( + <> + + + {mode === "login" ? ( +
+
+

+ 点击下方按钮获取授权 URL，然后复制到浏览器（支持指纹浏览器）完成 + Claude 登录。 +

+

+ 授权成功后，凭证将自动保存并添加到凭证池。 +

+
+ + +
+ ) : ( + + )} + + ), + }; +} diff --git a/src/components/provider-pool/credential-forms/CodexForm.tsx b/src/components/provider-pool/credential-forms/CodexForm.tsx new file mode 100644 index 000000000..94f3c1846 --- /dev/null +++ b/src/components/provider-pool/credential-forms/CodexForm.tsx @@ -0,0 +1,144 @@ +/** + * Codex 凭证添加表单 + * 支持 OpenAI OAuth 登录和文件导入两种模式 + */ + +import { useState, useEffect } from "react"; +import { listen } from "@tauri-apps/api/event"; +import { providerPoolApi } from "@/lib/api/providerPool"; +import { ModeSelector } from "./ModeSelector"; +import { FileImportForm } from "./FileImportForm"; +import { OAuthUrlDisplay } from "./OAuthUrlDisplay"; + +interface CodexFormProps { + name: string; + credsFilePath: string; + setCredsFilePath: (path: string) => void; + onSelectFile: () => void; + loading: boolean; + setLoading: (loading: boolean) => void; + setError: (error: string | null) => void; + onSuccess: () => void; +} + +export function CodexForm({ + name, + credsFilePath, + setCredsFilePath, + onSelectFile, + loading: _loading, + setLoading, + setError, + onSuccess, +}: CodexFormProps) { + const [mode, setMode] = useState<"login" | "file">("login"); + const [authUrl, setAuthUrl] = useState(null); + const [waitingForCallback, setWaitingForCallback] = useState(false); + + // 监听后端发送的授权 URL 事件 + useEffect(() => { + let unlisten: (() => void) | undefined; + + const setupListener = async () => { + unlisten = await listen<{ auth_url: string }>( + "codex-auth-url", + (event) => { + setAuthUrl(event.payload.auth_url); + }, + ); + }; + + setupListener(); + + return () => { + if (unlisten) unlisten(); + }; + }, []); + + // 获取授权 URL 并启动服务器等待回调 + const handleGetAuthUrl = async () => { + setLoading(true); + setError(null); + setAuthUrl(null); + setWaitingForCallback(true); + + try { + const trimmedName = name.trim() || undefined; + await providerPoolApi.getCodexAuthUrlAndWait(trimmedName); + onSuccess(); + } catch (e) { + const errorMsg = e instanceof Error ? e.message : String(e); + setError(errorMsg); + setWaitingForCallback(false); + } finally { + setLoading(false); + } + }; + + // 文件导入提交 + const handleFileSubmit = async () => { + if (!credsFilePath) { + setError("请选择凭证文件"); + return; + } + + setLoading(true); + setError(null); + + try { + const trimmedName = name.trim() || undefined; + await providerPoolApi.addCodexOAuth(credsFilePath, trimmedName); + onSuccess(); + } catch (e) { + setError(e instanceof Error ? e.message : String(e)); + } finally { + setLoading(false); + } + }; + + return { + mode, + authUrl, + waitingForCallback, + handleGetAuthUrl, + handleFileSubmit, + render: () => ( + <> + + + {mode === "login" ? ( +
+
+

+ 点击下方按钮获取授权 URL，然后复制到浏览器（支持指纹浏览器）完成 + OpenAI 登录。 +

+

+ 授权成功后，凭证将自动保存并添加到凭证池。 +

+
+ + +
+ ) : ( + + )} + + ), + }; +} diff --git a/src/components/provider-pool/credential-forms/FileImportForm.tsx b/src/components/provider-pool/credential-forms/FileImportForm.tsx new file mode 100644 index 000000000..a67f3cda3 --- /dev/null +++ b/src/components/provider-pool/credential-forms/FileImportForm.tsx @@ -0,0 +1,70 @@ +/** + * 文件导入表单组件 + */ + +import { FolderOpen } from "lucide-react"; + +interface FileImportFormProps { + credsFilePath: string; + setCredsFilePath: (path: string) => void; + onSelectFile: () => void; + placeholder?: string; + hint?: string; + projectId?: string; + setProjectId?: (id: string) => void; + showProjectId?: boolean; +} + +export function FileImportForm({ + credsFilePath, + setCredsFilePath, + onSelectFile, + placeholder = "选择凭证文件...", + hint, + projectId, + setProjectId, + showProjectId = false, +}: FileImportFormProps) { + return ( + <> +
+ +
+ setCredsFilePath(e.target.value)} + placeholder={placeholder} + className="flex-1 rounded-lg border bg-background px-3 py-2 text-sm" + /> + +
+ {hint &&

{hint}

} +
+ + {showProjectId && setProjectId && ( +
+ + setProjectId(e.target.value)} + placeholder="Google Cloud Project ID..." + className="w-full rounded-lg border bg-background px-3 py-2 text-sm" + /> +
+ )} + + ); +} diff --git a/src/components/provider-pool/credential-forms/GeminiForm.tsx b/src/components/provider-pool/credential-forms/GeminiForm.tsx new file mode 100644 index 000000000..ed9ba2d8b --- /dev/null +++ b/src/components/provider-pool/credential-forms/GeminiForm.tsx @@ -0,0 +1,270 @@ +/** + * Gemini 凭证添加表单 + * 支持 Google OAuth 登录和文件导入两种模式 + * + * Gemini OAuth 流程： + * 1. 生成授权 URL（包含 PKCE） + * 2. 用户在浏览器中打开 URL 并授权 + * 3. 浏览器跳转到 codeassist.google.com/authcode 显示 code + * 4. 用户复制 code 回应用 + * 5. 应用用 code 交换 tokens + */ + +import { useState, useEffect } from "react"; +import { listen } from "@tauri-apps/api/event"; +import { providerPoolApi } from "@/lib/api/providerPool"; +import { ModeSelector } from "./ModeSelector"; +import { FileImportForm } from "./FileImportForm"; +import { Copy, Check, Loader2 } from "lucide-react"; + +interface GeminiFormProps { + name: string; + credsFilePath: string; + setCredsFilePath: (path: string) => void; + projectId: string; + setProjectId: (id: string) => void; + onSelectFile: () => void; + loading: boolean; + setLoading: (loading: boolean) => void; + setError: (error: string | null) => void; + onSuccess: () => void; +} + +export function GeminiForm({ + name, + credsFilePath, + setCredsFilePath, + projectId, + setProjectId, + onSelectFile, + loading: _loading, + setLoading, + setError, + onSuccess, +}: GeminiFormProps) { + const [mode, setMode] = useState<"login" | "file">("login"); + const [authUrl, setAuthUrl] = useState(null); + const [sessionId, setSessionId] = useState(null); + const [authCode, setAuthCode] = useState(""); + const [copied, setCopied] = useState(false); + const [exchanging, setExchanging] = useState(false); + + // 监听后端发送的授权 URL 事件 + useEffect(() => { + let unlisten: (() => void) | undefined; + + const setupListener = async () => { + unlisten = await listen<{ auth_url: string; session_id: string }>( + "gemini-auth-url", + (event) => { + console.log("[Gemini OAuth] 收到授权 URL 事件:", event.payload); + setAuthUrl(event.payload.auth_url); + setSessionId(event.payload.session_id); + }, + ); + }; + + setupListener(); + + return () => { + if (unlisten) unlisten(); + }; + }, []); + + // 获取授权 URL + const handleGetAuthUrl = async () => { + setLoading(true); + setError(null); + setAuthUrl(null); + setSessionId(null); + setAuthCode(""); + + try { + // 调用后端生成授权 URL + await providerPoolApi.getGeminiAuthUrlAndWait(name.trim() || undefined); + // URL 会通过事件返回 + } catch (e) { + const errorMsg = e instanceof Error ? e.message : String(e); + // 检查是否包含 AUTH_URL 前缀（后端返回的授权 URL） + if (errorMsg.includes("AUTH_URL:")) { + const urlMatch = errorMsg.match(/AUTH_URL:(.+?)(?:\s|$)/); + if (urlMatch) { + setAuthUrl(urlMatch[1]); + } + } else { + setError(errorMsg); + } + } finally { + setLoading(false); + } + }; + + // 用 code 交换 token + const handleExchangeCode = async () => { + if (!authCode.trim()) { + setError("请输入授权码"); + return; + } + + setExchanging(true); + setError(null); + + try { + const trimmedName = name.trim() || undefined; + await providerPoolApi.exchangeGeminiCode( + authCode.trim(), + sessionId || undefined, + trimmedName, + ); + onSuccess(); + } catch (e) { + setError(e instanceof Error ? e.message : String(e)); + } finally { + setExchanging(false); + } + }; + + // 复制 URL + const handleCopyUrl = async () => { + if (authUrl) { + await navigator.clipboard.writeText(authUrl); + setCopied(true); + setTimeout(() => setCopied(false), 2000); + } + }; + + // 文件导入提交 + const handleFileSubmit = async () => { + if (!credsFilePath) { + setError("请选择凭证文件"); + return; + } + + setLoading(true); + setError(null); + + try { + const trimmedName = name.trim() || undefined; + await providerPoolApi.addGeminiOAuth( + credsFilePath, + projectId.trim() || undefined, + trimmedName, + ); + onSuccess(); + } catch (e) { + setError(e instanceof Error ? e.message : String(e)); + } finally { + setLoading(false); + } + }; + + return { + mode, + authUrl, + sessionId, + authCode, + exchanging, + handleGetAuthUrl, + handleExchangeCode, + handleFileSubmit, + render: () => ( + <> + + + {mode === "login" ? ( +
+
+

+ 点击下方按钮获取授权 URL，然后复制到浏览器完成 Google 登录。 +

+

+ 授权成功后，复制页面显示的授权码粘贴到下方输入框。 +

+
+ + {/* 授权 URL 显示 */} + {authUrl && ( +
+
+ 授权 URL + +
+
+

+ {authUrl.length > 100 + ? `${authUrl.slice(0, 100)}...` + : authUrl} +

+
+ + {/* 授权码输入 */} +
+ + setAuthCode(e.target.value)} + placeholder="粘贴浏览器页面显示的授权码..." + className="w-full rounded-lg border bg-background px-3 py-2 text-sm" + /> +

+ 在浏览器中完成授权后，复制页面显示的授权码 +

+
+ + {/* 提交按钮 */} + +
+ )} +
+ ) : ( + + )} + + ), + }; +} diff --git a/src/components/provider-pool/credential-forms/IFlowForm.tsx b/src/components/provider-pool/credential-forms/IFlowForm.tsx new file mode 100644 index 000000000..a13c11483 --- /dev/null +++ b/src/components/provider-pool/credential-forms/IFlowForm.tsx @@ -0,0 +1,144 @@ +/** + * iFlow 凭证添加表单 + * 支持 OAuth 登录和文件导入两种模式 + */ + +import { useState, useEffect } from "react"; +import { listen } from "@tauri-apps/api/event"; +import { providerPoolApi } from "@/lib/api/providerPool"; +import { ModeSelector } from "./ModeSelector"; +import { FileImportForm } from "./FileImportForm"; +import { OAuthUrlDisplay } from "./OAuthUrlDisplay"; + +interface IFlowFormProps { + name: string; + credsFilePath: string; + setCredsFilePath: (path: string) => void; + onSelectFile: () => void; + loading: boolean; + setLoading: (loading: boolean) => void; + setError: (error: string | null) => void; + onSuccess: () => void; +} + +export function IFlowForm({ + name, + credsFilePath, + setCredsFilePath, + onSelectFile, + loading: _loading, + setLoading, + setError, + onSuccess, +}: IFlowFormProps) { + const [mode, setMode] = useState<"login" | "file">("login"); + const [authUrl, setAuthUrl] = useState(null); + const [waitingForCallback, setWaitingForCallback] = useState(false); + + // 监听后端发送的授权 URL 事件 + useEffect(() => { + let unlisten: (() => void) | undefined; + + const setupListener = async () => { + unlisten = await listen<{ auth_url: string }>( + "iflow-auth-url", + (event) => { + setAuthUrl(event.payload.auth_url); + }, + ); + }; + + setupListener(); + + return () => { + if (unlisten) unlisten(); + }; + }, []); + + // 获取授权 URL 并启动服务器等待回调 + const handleGetAuthUrl = async () => { + setLoading(true); + setError(null); + setAuthUrl(null); + setWaitingForCallback(true); + + try { + const trimmedName = name.trim() || undefined; + await providerPoolApi.getIFlowAuthUrlAndWait(trimmedName); + onSuccess(); + } catch (e) { + const errorMsg = e instanceof Error ? e.message : String(e); + setError(errorMsg); + setWaitingForCallback(false); + } finally { + setLoading(false); + } + }; + + // 文件导入提交 + const handleFileSubmit = async () => { + if (!credsFilePath) { + setError("请选择凭证文件"); + return; + } + + setLoading(true); + setError(null); + + try { + const trimmedName = name.trim() || undefined; + await providerPoolApi.addIFlowOAuth(credsFilePath, trimmedName); + onSuccess(); + } catch (e) { + setError(e instanceof Error ? e.message : String(e)); + } finally { + setLoading(false); + } + }; + + return { + mode, + authUrl, + waitingForCallback, + handleGetAuthUrl, + handleFileSubmit, + render: () => ( + <> + + + {mode === "login" ? ( +
+
+

+ 点击下方按钮获取授权 URL，然后复制到浏览器（支持指纹浏览器）完成 + iFlow 登录。 +

+

+ 授权成功后，凭证将自动保存并添加到凭证池。 +

+
+ + +
+ ) : ( + + )} + + ), + }; +} diff --git a/src/components/provider-pool/credential-forms/ModeSelector.tsx b/src/components/provider-pool/credential-forms/ModeSelector.tsx new file mode 100644 index 000000000..67839199e --- /dev/null +++ b/src/components/provider-pool/credential-forms/ModeSelector.tsx @@ -0,0 +1,48 @@ +/** + * 登录/文件导入模式选择器 + */ + +import { LogIn, FolderOpen } from "lucide-react"; + +interface ModeSelectorProps { + mode: "login" | "file"; + setMode: (mode: "login" | "file") => void; + loginLabel?: string; + fileLabel?: string; +} + +export function ModeSelector({ + mode, + setMode, + loginLabel = "登录", + fileLabel = "导入文件", +}: ModeSelectorProps) { + return ( +
+ + +
+ ); +} diff --git a/src/components/provider-pool/credential-forms/OAuthUrlDisplay.tsx b/src/components/provider-pool/credential-forms/OAuthUrlDisplay.tsx new file mode 100644 index 000000000..dff245565 --- /dev/null +++ b/src/components/provider-pool/credential-forms/OAuthUrlDisplay.tsx @@ -0,0 +1,72 @@ +/** + * OAuth 授权 URL 显示组件 + * 用于显示授权 URL 和等待回调状态 + */ + +import { useState } from "react"; +import { Copy, Check, Loader2 } from "lucide-react"; + +interface OAuthUrlDisplayProps { + authUrl: string | null; + waitingForCallback: boolean; + colorScheme?: "blue" | "green" | "purple" | "amber"; +} + +export function OAuthUrlDisplay({ + authUrl, + waitingForCallback, + colorScheme: _colorScheme = "blue", +}: OAuthUrlDisplayProps) { + const [urlCopied, setUrlCopied] = useState(false); + + const handleCopyUrl = () => { + if (authUrl) { + navigator.clipboard.writeText(authUrl); + setUrlCopied(true); + setTimeout(() => setUrlCopied(false), 2000); + } + }; + + if (!authUrl) return null; + + return ( +
+
+
+ 授权 URL + +
+

+ {authUrl.slice(0, 100)}... +

+
+ + {waitingForCallback && ( +
+
+ +

+ 请复制上方 URL 到浏览器完成登录，正在等待授权回调... +

+
+
+ )} +
+ ); +} diff --git a/src/components/provider-pool/credential-forms/QwenForm.tsx b/src/components/provider-pool/credential-forms/QwenForm.tsx new file mode 100644 index 000000000..d032a49d5 --- /dev/null +++ b/src/components/provider-pool/credential-forms/QwenForm.tsx @@ -0,0 +1,215 @@ +/** + * Qwen 凭证添加表单 + * 支持 Device Code Flow 登录和文件导入两种模式 + */ + +import { useState, useEffect } from "react"; +import { listen } from "@tauri-apps/api/event"; +import { open } from "@tauri-apps/plugin-shell"; +import { providerPoolApi } from "@/lib/api/providerPool"; +import { ModeSelector } from "./ModeSelector"; +import { FileImportForm } from "./FileImportForm"; +import { Copy, Check, ExternalLink } from "lucide-react"; + +interface QwenFormProps { + name: string; + credsFilePath: string; + setCredsFilePath: (path: string) => void; + onSelectFile: () => void; + loading: boolean; + setLoading: (loading: boolean) => void; + setError: (error: string | null) => void; + onSuccess: () => void; +} + +interface DeviceCodeInfo { + user_code: string; + verification_uri: string; + verification_uri_complete?: string; + expires_in: number; +} + +export function QwenForm({ + name, + credsFilePath, + setCredsFilePath, + onSelectFile, + loading: _loading, + setLoading, + setError, + onSuccess, +}: QwenFormProps) { + const [mode, setMode] = useState<"login" | "file">("login"); + const [deviceCode, setDeviceCode] = useState(null); + const [waitingForAuth, setWaitingForAuth] = useState(false); + const [copied, setCopied] = useState(false); + + // 监听后端发送的设备码事件 + useEffect(() => { + let unlisten: (() => void) | undefined; + + const setupListener = async () => { + unlisten = await listen("qwen-device-code", (event) => { + setDeviceCode(event.payload); + }); + }; + + setupListener(); + + return () => { + if (unlisten) unlisten(); + }; + }, []); + + // 复制用户码 + const handleCopyCode = async () => { + if (deviceCode?.user_code) { + await navigator.clipboard.writeText(deviceCode.user_code); + setCopied(true); + setTimeout(() => setCopied(false), 2000); + } + }; + + // 打开验证链接 + const handleOpenVerificationUrl = async () => { + if (deviceCode) { + const url = + deviceCode.verification_uri_complete || deviceCode.verification_uri; + try { + await open(url); + } catch (e) { + console.error("Failed to open URL:", e); + // 如果 Tauri shell 失败，尝试使用 window.open 作为后备 + window.open(url, "_blank"); + } + } + }; + + // 获取设备码并启动轮询 + const handleGetDeviceCode = async () => { + setLoading(true); + setError(null); + setDeviceCode(null); + setWaitingForAuth(true); + + try { + const trimmedName = name.trim() || undefined; + await providerPoolApi.getQwenDeviceCodeAndWait(trimmedName); + onSuccess(); + } catch (e) { + const errorMsg = e instanceof Error ? e.message : String(e); + setError(errorMsg); + setWaitingForAuth(false); + } finally { + setLoading(false); + } + }; + + // 文件导入提交 + const handleFileSubmit = async () => { + if (!credsFilePath) { + setError("请选择凭证文件"); + return; + } + + setLoading(true); + setError(null); + + try { + const trimmedName = name.trim() || undefined; + await providerPoolApi.addQwenOAuth(credsFilePath, trimmedName); + onSuccess(); + } catch (e) { + setError(e instanceof Error ? e.message : String(e)); + } finally { + setLoading(false); + } + }; + + return { + mode, + deviceCode, + waitingForAuth, + handleGetDeviceCode, + handleFileSubmit, + render: () => ( + <> + + + {mode === "login" ? ( +
+
+

+ 点击下方按钮获取设备码，然后在浏览器中完成 Qwen 登录授权。 +

+

+ 授权成功后，凭证将自动保存并添加到凭证池。 +

+
+ + {deviceCode && ( +
+ {/* 用户码显示 */} +
+

+ 请在浏览器中输入以下验证码： +

+
+ + {deviceCode.user_code} + + +
+
+ + {/* 验证链接 */} + + + {waitingForAuth && ( +
+
+ 等待授权中... +
+ )} +
+ )} +
+ ) : ( + + )} + + ), + }; +} diff --git a/src/components/provider-pool/credential-forms/index.ts b/src/components/provider-pool/credential-forms/index.ts new file mode 100644 index 000000000..867fd7c52 --- /dev/null +++ b/src/components/provider-pool/credential-forms/index.ts @@ -0,0 +1,14 @@ +/** + * 凭证表单组件导出 + */ + +export * from "./types"; +export * from "./ModeSelector"; +export * from "./OAuthUrlDisplay"; +export * from "./FileImportForm"; +export * from "./AntigravityForm"; +export * from "./CodexForm"; +export * from "./ClaudeOAuthForm"; +export * from "./QwenForm"; +export * from "./IFlowForm"; +export * from "./GeminiForm"; diff --git a/src/components/provider-pool/credential-forms/types.ts b/src/components/provider-pool/credential-forms/types.ts new file mode 100644 index 000000000..cad5e3290 --- /dev/null +++ b/src/components/provider-pool/credential-forms/types.ts @@ -0,0 +1,59 @@ +/** + * 凭证表单共享类型定义 + */ + +import { PoolProviderType } from "@/lib/api/providerPool"; + +/** OAuth 登录表单的通用 Props */ +export interface OAuthLoginFormProps { + name: string; + loading: boolean; + error: string | null; + onSuccess: () => void; + setLoading: (loading: boolean) => void; + setError: (error: string | null) => void; +} + +/** 文件导入表单的通用 Props */ +export interface FileImportFormProps { + credsFilePath: string; + setCredsFilePath: (path: string) => void; + projectId?: string; + setProjectId?: (id: string) => void; + onSelectFile: () => void; + defaultPathHint?: string; + fileHint?: string; +} + +/** API Key 表单的通用 Props */ +export interface ApiKeyFormProps { + apiKey: string; + setApiKey: (key: string) => void; + baseUrl: string; + setBaseUrl: (url: string) => void; + providerType: PoolProviderType; +} + +/** 默认凭证文件路径 */ +export const defaultCredsPath: Record = { + kiro: "~/.aws/sso/cache/kiro-auth-token.json", + gemini: "~/.gemini/oauth_creds.json", + qwen: "~/.qwen/oauth_creds.json", + antigravity: "", + codex: "~/.codex/auth.json", + claude_oauth: "~/.claude/oauth.json", + iflow: "~/.iflow/oauth_creds.json", +}; + +/** Provider 显示名称 */ +export const providerLabels: Record = { + kiro: "Kiro (AWS)", + gemini: "Gemini (Google)", + qwen: "Qwen (阿里)", + openai: "OpenAI", + claude: "Claude (Anthropic)", + antigravity: "Antigravity (Gemini 3 Pro)", + codex: "Codex (OpenAI OAuth)", + claude_oauth: "Claude OAuth", + iflow: "iFlow", +}; diff --git a/src/lib/api/providerPool.ts b/src/lib/api/providerPool.ts index f17dc5f67..ae4894b00 100644 --- a/src/lib/api/providerPool.ts +++ b/src/lib/api/providerPool.ts @@ -359,6 +359,96 @@ export const providerPoolApi = { return invoke("add_iflow_cookie_credential", { credsFilePath, name }); }, + // Antigravity OAuth 登录（打开浏览器授权） + async startAntigravityOAuthLogin( + name?: string, + skipProjectIdFetch?: boolean, + ): Promise { + return invoke("start_antigravity_oauth_login", { + name, + skipProjectIdFetch, + }); + }, + + // 获取 Antigravity OAuth 授权 URL 并等待回调（不自动打开浏览器） + // 服务器会在后台等待回调，成功后返回凭证 + // 如果需要显示 URL，错误信息会包含 AUTH_URL: 前缀 + async getAntigravityAuthUrlAndWait( + name?: string, + skipProjectIdFetch?: boolean, + ): Promise { + return invoke("get_antigravity_auth_url_and_wait", { + name, + skipProjectIdFetch, + }); + }, + + // Codex OAuth 登录（打开浏览器授权） + async startCodexOAuthLogin(name?: string): Promise { + return invoke("start_codex_oauth_login", { name }); + }, + + // 获取 Codex OAuth 授权 URL 并等待回调（不自动打开浏览器） + // 服务器会在后台等待回调，成功后返回凭证 + async getCodexAuthUrlAndWait(name?: string): Promise { + return invoke("get_codex_auth_url_and_wait", { name }); + }, + + // Claude OAuth 登录（打开浏览器授权） + async startClaudeOAuthLogin(name?: string): Promise { + return invoke("start_claude_oauth_login", { name }); + }, + + // 获取 Claude OAuth 授权 URL 并等待回调（不自动打开浏览器） + // 服务器会在后台等待回调，成功后返回凭证 + async getClaudeOAuthAuthUrlAndWait( + name?: string, + ): Promise { + return invoke("get_claude_oauth_auth_url_and_wait", { name }); + }, + + // Qwen Device Code Flow 登录（打开浏览器授权） + async startQwenDeviceCodeLogin(name?: string): Promise { + return invoke("start_qwen_device_code_login", { name }); + }, + + // 获取 Qwen Device Code 并等待用户授权（不自动打开浏览器） + // 服务器会在后台轮询等待授权，成功后返回凭证 + async getQwenDeviceCodeAndWait(name?: string): Promise { + return invoke("get_qwen_device_code_and_wait", { name }); + }, + + // iFlow OAuth 登录（打开浏览器授权） + async startIFlowOAuthLogin(name?: string): Promise { + return invoke("start_iflow_oauth_login", { name }); + }, + + // 获取 iFlow OAuth 授权 URL 并等待回调（不自动打开浏览器） + // 服务器会在后台等待回调，成功后返回凭证 + async getIFlowAuthUrlAndWait(name?: string): Promise { + return invoke("get_iflow_auth_url_and_wait", { name }); + }, + + // Gemini OAuth 登录（打开浏览器授权） + async startGeminiOAuthLogin(name?: string): Promise { + return invoke("start_gemini_oauth_login", { name }); + }, + + // 获取 Gemini OAuth 授权 URL 并等待回调（不自动打开浏览器） + // 服务器会在后台等待回调，成功后返回凭证 + async getGeminiAuthUrlAndWait(name?: string): Promise { + return invoke("get_gemini_auth_url_and_wait", { name }); + }, + + // 用 Gemini 授权码交换 token + async exchangeGeminiCode( + code: string, + sessionId?: string, + name?: string, + ): Promise { + return invoke("exchange_gemini_code", { code, sessionId, name }); + }, + // OAuth token management async refreshCredentialToken(uuid: string): Promise { return invoke("refresh_pool_credential_token", { uuid }); @@ -380,3 +470,22 @@ export interface MigrationResult { skipped_count: number; errors: string[]; } + +// Kiro 凭证指纹信息 +export interface KiroFingerprintInfo { + /** Machine ID（SHA256 哈希，64 字符） */ + machine_id: string; + /** Machine ID 的短格式（前 16 字符） */ + machine_id_short: string; + /** 指纹来源（profileArn / clientId / system） */ + source: string; + /** 认证方式 */ + auth_method: string; +} + +// 获取 Kiro 凭证的指纹信息 +export async function getKiroCredentialFingerprint( + uuid: string, +): Promise { + return invoke("get_kiro_credential_fingerprint", { uuid }); +}