mirror of
https://github.com/farion1231/cc-switch.git
synced 2026-07-30 10:25:05 +08:00
Merge branch 'main' into feat/gemini-proxy-integration
# Conflicts: # src-tauri/src/proxy/providers/claude.rs # src-tauri/src/proxy/sse.rs # src-tauri/src/services/stream_check.rs
This commit is contained in:
YoVinchen
@@ -0,0 +1,411 @@
|
||||
//! 供应商余额查询服务
|
||||
//!
|
||||
//! 支持 DeepSeek、StepFun、SiliconFlow、OpenRouter、Novita AI 的账户余额查询。
|
||||
//! 返回 UsageResult 格式，与现有用量系统无缝对接。
|
||||
|
||||
use crate::provider::{UsageData, UsageResult};
|
||||
use std::time::Duration;
|
||||
|
||||
// ── 供应商检测 ──────────────────────────────────────────────
|
||||
|
||||
enum BalanceProvider {
|
||||
DeepSeek,
|
||||
StepFun,
|
||||
SiliconFlow,
|
||||
SiliconFlowEn,
|
||||
OpenRouter,
|
||||
NovitaAI,
|
||||
}
|
||||
|
||||
fn detect_provider(base_url: &str) -> Option<BalanceProvider> {
|
||||
let url = base_url.to_lowercase();
|
||||
if url.contains("api.deepseek.com") {
|
||||
Some(BalanceProvider::DeepSeek)
|
||||
} else if url.contains("api.stepfun.ai") || url.contains("api.stepfun.com") {
|
||||
Some(BalanceProvider::StepFun)
|
||||
} else if url.contains("api.siliconflow.cn") {
|
||||
Some(BalanceProvider::SiliconFlow)
|
||||
} else if url.contains("api.siliconflow.com") {
|
||||
Some(BalanceProvider::SiliconFlowEn)
|
||||
} else if url.contains("openrouter.ai") {
|
||||
Some(BalanceProvider::OpenRouter)
|
||||
} else if url.contains("api.novita.ai") {
|
||||
Some(BalanceProvider::NovitaAI)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
fn make_error(msg: String) -> UsageResult {
|
||||
UsageResult {
|
||||
success: false,
|
||||
data: None,
|
||||
error: Some(msg),
|
||||
}
|
||||
}
|
||||
|
||||
fn make_auth_error(status: reqwest::StatusCode) -> UsageResult {
|
||||
UsageResult {
|
||||
success: false,
|
||||
data: Some(vec![UsageData {
|
||||
plan_name: None,
|
||||
remaining: None,
|
||||
total: None,
|
||||
used: None,
|
||||
unit: None,
|
||||
is_valid: Some(false),
|
||||
invalid_message: Some(format!("Authentication failed (HTTP {status})")),
|
||||
extra: None,
|
||||
}]),
|
||||
error: Some(format!("Authentication failed (HTTP {status})")),
|
||||
}
|
||||
}
|
||||
|
||||
// ── DeepSeek ────────────────────────────────────────────────
|
||||
// GET https://api.deepseek.com/user/balance
|
||||
// Response: { balance_infos: [{ currency, total_balance, granted_balance, topped_up_balance }], is_available }
|
||||
|
||||
async fn query_deepseek(api_key: &str) -> UsageResult {
|
||||
let client = crate::proxy::http_client::get();
|
||||
|
||||
let resp = client
|
||||
.get("https://api.deepseek.com/user/balance")
|
||||
.header("Authorization", format!("Bearer {api_key}"))
|
||||
.header("Accept", "application/json")
|
||||
.timeout(Duration::from_secs(10))
|
||||
.send()
|
||||
.await;
|
||||
|
||||
let resp = match resp {
|
||||
Ok(r) => r,
|
||||
Err(e) => return make_error(format!("Network error: {e}")),
|
||||
};
|
||||
|
||||
let status = resp.status();
|
||||
if status == reqwest::StatusCode::UNAUTHORIZED || status == reqwest::StatusCode::FORBIDDEN {
|
||||
return make_auth_error(status);
|
||||
}
|
||||
if !status.is_success() {
|
||||
let body = resp.text().await.unwrap_or_default();
|
||||
return make_error(format!("API error (HTTP {status}): {body}"));
|
||||
}
|
||||
|
||||
let body: serde_json::Value = match resp.json().await {
|
||||
Ok(v) => v,
|
||||
Err(e) => return make_error(format!("Failed to parse response: {e}")),
|
||||
};
|
||||
|
||||
let is_available = body
|
||||
.get("is_available")
|
||||
.and_then(|v| v.as_bool())
|
||||
.unwrap_or(true);
|
||||
let mut data = Vec::new();
|
||||
|
||||
if let Some(infos) = body.get("balance_infos").and_then(|v| v.as_array()) {
|
||||
for info in infos {
|
||||
let currency = info
|
||||
.get("currency")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("CNY");
|
||||
let total = parse_f64_field(info, "total_balance");
|
||||
|
||||
data.push(UsageData {
|
||||
plan_name: Some(currency.to_string()),
|
||||
remaining: total,
|
||||
total: None,
|
||||
used: None,
|
||||
unit: Some(currency.to_string()),
|
||||
is_valid: Some(is_available),
|
||||
invalid_message: if !is_available {
|
||||
Some("Insufficient balance".to_string())
|
||||
} else {
|
||||
None
|
||||
},
|
||||
extra: None,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
UsageResult {
|
||||
success: true,
|
||||
data: if data.is_empty() { None } else { Some(data) },
|
||||
error: None,
|
||||
}
|
||||
}
|
||||
|
||||
// ── StepFun ─────────────────────────────────────────────────
|
||||
// GET https://api.stepfun.com/v1/accounts
|
||||
// Response: { object, type, balance, total_cash_balance, total_voucher_balance }
|
||||
|
||||
async fn query_stepfun(api_key: &str) -> UsageResult {
|
||||
let client = crate::proxy::http_client::get();
|
||||
|
||||
let resp = client
|
||||
.get("https://api.stepfun.com/v1/accounts")
|
||||
.header("Authorization", format!("Bearer {api_key}"))
|
||||
.header("Accept", "application/json")
|
||||
.timeout(Duration::from_secs(10))
|
||||
.send()
|
||||
.await;
|
||||
|
||||
let resp = match resp {
|
||||
Ok(r) => r,
|
||||
Err(e) => return make_error(format!("Network error: {e}")),
|
||||
};
|
||||
|
||||
let status = resp.status();
|
||||
if status == reqwest::StatusCode::UNAUTHORIZED || status == reqwest::StatusCode::FORBIDDEN {
|
||||
return make_auth_error(status);
|
||||
}
|
||||
if !status.is_success() {
|
||||
let body = resp.text().await.unwrap_or_default();
|
||||
return make_error(format!("API error (HTTP {status}): {body}"));
|
||||
}
|
||||
|
||||
let body: serde_json::Value = match resp.json().await {
|
||||
Ok(v) => v,
|
||||
Err(e) => return make_error(format!("Failed to parse response: {e}")),
|
||||
};
|
||||
|
||||
let balance = parse_f64_field(&body, "balance").unwrap_or(0.0);
|
||||
|
||||
UsageResult {
|
||||
success: true,
|
||||
data: Some(vec![UsageData {
|
||||
plan_name: Some("StepFun".to_string()),
|
||||
remaining: Some(balance),
|
||||
total: None,
|
||||
used: None,
|
||||
unit: Some("CNY".to_string()),
|
||||
is_valid: Some(true),
|
||||
invalid_message: None,
|
||||
extra: None,
|
||||
}]),
|
||||
error: None,
|
||||
}
|
||||
}
|
||||
|
||||
// ── SiliconFlow ─────────────────────────────────────────────
|
||||
// GET https://api.siliconflow.cn/v1/user/info (or .com for EN)
|
||||
// Response: { code, data: { balance, chargeBalance, totalBalance, status } }
|
||||
|
||||
async fn query_siliconflow(api_key: &str, is_cn: bool) -> UsageResult {
|
||||
let client = crate::proxy::http_client::get();
|
||||
|
||||
let domain = if is_cn {
|
||||
"api.siliconflow.cn"
|
||||
} else {
|
||||
"api.siliconflow.com"
|
||||
};
|
||||
let url = format!("https://{domain}/v1/user/info");
|
||||
|
||||
let resp = client
|
||||
.get(&url)
|
||||
.header("Authorization", format!("Bearer {api_key}"))
|
||||
.header("Accept", "application/json")
|
||||
.timeout(Duration::from_secs(10))
|
||||
.send()
|
||||
.await;
|
||||
|
||||
let resp = match resp {
|
||||
Ok(r) => r,
|
||||
Err(e) => return make_error(format!("Network error: {e}")),
|
||||
};
|
||||
|
||||
let status = resp.status();
|
||||
if status == reqwest::StatusCode::UNAUTHORIZED || status == reqwest::StatusCode::FORBIDDEN {
|
||||
return make_auth_error(status);
|
||||
}
|
||||
if !status.is_success() {
|
||||
let body = resp.text().await.unwrap_or_default();
|
||||
return make_error(format!("API error (HTTP {status}): {body}"));
|
||||
}
|
||||
|
||||
let body: serde_json::Value = match resp.json().await {
|
||||
Ok(v) => v,
|
||||
Err(e) => return make_error(format!("Failed to parse response: {e}")),
|
||||
};
|
||||
|
||||
let data = match body.get("data") {
|
||||
Some(d) => d,
|
||||
None => return make_error("Missing 'data' field in response".to_string()),
|
||||
};
|
||||
|
||||
let total_balance = parse_f64_field(data, "totalBalance").unwrap_or(0.0);
|
||||
|
||||
UsageResult {
|
||||
success: true,
|
||||
data: Some(vec![UsageData {
|
||||
plan_name: Some("SiliconFlow".to_string()),
|
||||
remaining: Some(total_balance),
|
||||
total: None,
|
||||
used: None,
|
||||
unit: Some("CNY".to_string()),
|
||||
is_valid: Some(true),
|
||||
invalid_message: None,
|
||||
extra: None,
|
||||
}]),
|
||||
error: None,
|
||||
}
|
||||
}
|
||||
|
||||
// ── OpenRouter ──────────────────────────────────────────────
|
||||
// GET https://openrouter.ai/api/v1/credits
|
||||
// Response: { data: { total_credits, total_usage } }
|
||||
|
||||
async fn query_openrouter(api_key: &str) -> UsageResult {
|
||||
let client = crate::proxy::http_client::get();
|
||||
|
||||
let resp = client
|
||||
.get("https://openrouter.ai/api/v1/credits")
|
||||
.header("Authorization", format!("Bearer {api_key}"))
|
||||
.header("Accept", "application/json")
|
||||
.timeout(Duration::from_secs(10))
|
||||
.send()
|
||||
.await;
|
||||
|
||||
let resp = match resp {
|
||||
Ok(r) => r,
|
||||
Err(e) => return make_error(format!("Network error: {e}")),
|
||||
};
|
||||
|
||||
let status = resp.status();
|
||||
if status == reqwest::StatusCode::UNAUTHORIZED || status == reqwest::StatusCode::FORBIDDEN {
|
||||
return make_auth_error(status);
|
||||
}
|
||||
if !status.is_success() {
|
||||
let body = resp.text().await.unwrap_or_default();
|
||||
return make_error(format!("API error (HTTP {status}): {body}"));
|
||||
}
|
||||
|
||||
let body: serde_json::Value = match resp.json().await {
|
||||
Ok(v) => v,
|
||||
Err(e) => return make_error(format!("Failed to parse response: {e}")),
|
||||
};
|
||||
|
||||
let data = body.get("data").unwrap_or(&body);
|
||||
let total_credits = parse_f64_field(data, "total_credits").unwrap_or(0.0);
|
||||
let total_usage = parse_f64_field(data, "total_usage").unwrap_or(0.0);
|
||||
let remaining = total_credits - total_usage;
|
||||
|
||||
UsageResult {
|
||||
success: true,
|
||||
data: Some(vec![UsageData {
|
||||
plan_name: Some("OpenRouter".to_string()),
|
||||
remaining: Some(remaining),
|
||||
total: Some(total_credits),
|
||||
used: Some(total_usage),
|
||||
unit: Some("USD".to_string()),
|
||||
is_valid: Some(remaining > 0.0),
|
||||
invalid_message: if remaining <= 0.0 {
|
||||
Some("No credits remaining".to_string())
|
||||
} else {
|
||||
None
|
||||
},
|
||||
extra: None,
|
||||
}]),
|
||||
error: None,
|
||||
}
|
||||
}
|
||||
|
||||
// ── Novita AI ───────────────────────────────────────────────
|
||||
// GET https://api.novita.ai/v3/user/balance
|
||||
// Response: { availableBalance, cashBalance, creditLimit, outstandingInvoices }
|
||||
// 金额单位：0.0001 USD
|
||||
|
||||
async fn query_novita(api_key: &str) -> UsageResult {
|
||||
let client = crate::proxy::http_client::get();
|
||||
|
||||
let resp = client
|
||||
.get("https://api.novita.ai/v3/user/balance")
|
||||
.header("Authorization", format!("Bearer {api_key}"))
|
||||
.header("Accept", "application/json")
|
||||
.timeout(Duration::from_secs(10))
|
||||
.send()
|
||||
.await;
|
||||
|
||||
let resp = match resp {
|
||||
Ok(r) => r,
|
||||
Err(e) => return make_error(format!("Network error: {e}")),
|
||||
};
|
||||
|
||||
let status = resp.status();
|
||||
if status == reqwest::StatusCode::UNAUTHORIZED || status == reqwest::StatusCode::FORBIDDEN {
|
||||
return make_auth_error(status);
|
||||
}
|
||||
if !status.is_success() {
|
||||
let body = resp.text().await.unwrap_or_default();
|
||||
return make_error(format!("API error (HTTP {status}): {body}"));
|
||||
}
|
||||
|
||||
let body: serde_json::Value = match resp.json().await {
|
||||
Ok(v) => v,
|
||||
Err(e) => return make_error(format!("Failed to parse response: {e}")),
|
||||
};
|
||||
|
||||
// Novita 金额单位为 0.0001 USD，需除以 10000 转为 USD
|
||||
let available = parse_f64_field(&body, "availableBalance").unwrap_or(0.0) / 10000.0;
|
||||
|
||||
UsageResult {
|
||||
success: true,
|
||||
data: Some(vec![UsageData {
|
||||
plan_name: Some("Novita AI".to_string()),
|
||||
remaining: Some(available),
|
||||
total: None,
|
||||
used: None,
|
||||
unit: Some("USD".to_string()),
|
||||
is_valid: Some(available > 0.0),
|
||||
invalid_message: if available <= 0.0 {
|
||||
Some("No balance remaining".to_string())
|
||||
} else {
|
||||
None
|
||||
},
|
||||
extra: None,
|
||||
}]),
|
||||
error: None,
|
||||
}
|
||||
}
|
||||
|
||||
// ── 工具函数 ────────────────────────────────────────────────
|
||||
|
||||
/// 解析 JSON 字段为 f64，兼容数字和字符串格式
|
||||
fn parse_f64_field(obj: &serde_json::Value, field: &str) -> Option<f64> {
|
||||
obj.get(field).and_then(|v| {
|
||||
v.as_f64()
|
||||
.or_else(|| v.as_str().and_then(|s| s.parse().ok()))
|
||||
})
|
||||
}
|
||||
|
||||
// ── 公开入口 ────────────────────────────────────────────────
|
||||
|
||||
pub async fn get_balance(base_url: &str, api_key: &str) -> Result<UsageResult, String> {
|
||||
if api_key.trim().is_empty() {
|
||||
return Ok(UsageResult {
|
||||
success: false,
|
||||
data: None,
|
||||
error: Some("API key is empty".to_string()),
|
||||
});
|
||||
}
|
||||
|
||||
let provider = match detect_provider(base_url) {
|
||||
Some(p) => p,
|
||||
None => {
|
||||
return Ok(UsageResult {
|
||||
success: false,
|
||||
data: None,
|
||||
error: Some("Unknown balance provider".to_string()),
|
||||
})
|
||||
}
|
||||
};
|
||||
|
||||
let result = match provider {
|
||||
BalanceProvider::DeepSeek => query_deepseek(api_key).await,
|
||||
BalanceProvider::StepFun => query_stepfun(api_key).await,
|
||||
BalanceProvider::SiliconFlow => query_siliconflow(api_key, true).await,
|
||||
BalanceProvider::SiliconFlowEn => query_siliconflow(api_key, false).await,
|
||||
BalanceProvider::OpenRouter => query_openrouter(api_key).await,
|
||||
BalanceProvider::NovitaAI => query_novita(api_key).await,
|
||||
};
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
@@ -1,3 +1,4 @@
|
||||
pub mod balance;
|
||||
pub mod coding_plan;
|
||||
pub mod config;
|
||||
pub mod env_checker;
|
||||
@@ -8,6 +9,9 @@ pub mod omo;
|
||||
pub mod prompt;
|
||||
pub mod provider;
|
||||
pub mod proxy;
|
||||
pub mod session_usage;
|
||||
pub mod session_usage_codex;
|
||||
pub mod session_usage_gemini;
|
||||
pub mod skill;
|
||||
pub mod speedtest;
|
||||
pub mod stream_check;
|
||||
|
||||
@@ -999,11 +999,12 @@ pub fn import_default_config(state: &AppState, app_type: AppType) -> Result<bool
|
||||
return Ok(false);
|
||||
}
|
||||
|
||||
{
|
||||
let providers = state.db.get_all_providers(app_type.as_str())?;
|
||||
if !providers.is_empty() {
|
||||
return Ok(false); // 已有供应商，跳过
|
||||
}
|
||||
// 允许 "只有官方 seed 预设" 的情况下继续导入 live：
|
||||
// - 启动编排顺序是先 import 后 seed，新用户启动时 providers 为空，导入照常
|
||||
// - 老用户已有非 seed provider，跳过导入（正确）
|
||||
// - 用户手动点 ProviderEmptyState 的导入按钮时，与官方 seed 共存而不被阻塞
|
||||
if state.db.has_non_official_seed_provider(app_type.as_str())? {
|
||||
return Ok(false);
|
||||
}
|
||||
|
||||
let settings_config = match app_type {
|
||||
@@ -1208,11 +1209,11 @@ pub fn import_opencode_providers_from_live(state: &AppState) -> Result<usize, Ap
|
||||
}
|
||||
|
||||
let mut imported = 0;
|
||||
let existing = state.db.get_all_providers("opencode")?;
|
||||
let existing_ids = state.db.get_provider_ids("opencode")?;
|
||||
|
||||
for (id, config) in providers {
|
||||
// Skip if already exists in database
|
||||
if existing.contains_key(&id) {
|
||||
if existing_ids.contains(&id) {
|
||||
log::debug!("OpenCode provider '{id}' already exists in database, skipping");
|
||||
continue;
|
||||
}
|
||||
@@ -1265,7 +1266,7 @@ pub fn import_openclaw_providers_from_live(state: &AppState) -> Result<usize, Ap
|
||||
}
|
||||
|
||||
let mut imported = 0;
|
||||
let existing = state.db.get_all_providers("openclaw")?;
|
||||
let existing_ids = state.db.get_provider_ids("openclaw")?;
|
||||
|
||||
for (id, config) in providers {
|
||||
// Validate: skip entries with empty id or no models
|
||||
@@ -1279,7 +1280,7 @@ pub fn import_openclaw_providers_from_live(state: &AppState) -> Result<usize, Ap
|
||||
}
|
||||
|
||||
// Skip if already exists in database
|
||||
if existing.contains_key(&id) {
|
||||
if existing_ids.contains(&id) {
|
||||
log::debug!("OpenClaw provider '{id}' already exists in database, skipping");
|
||||
continue;
|
||||
}
|
||||
|
||||
@@ -0,0 +1,634 @@
|
||||
//! Claude Code 会话日志使用追踪
|
||||
//!
|
||||
//! 从 ~/.claude/projects/ 下的 JSONL 会话文件中提取 token 使用数据，
|
||||
//! 实现无代理模式下的使用统计。
|
||||
//!
|
||||
//! ## 数据流
|
||||
//! ```text
|
||||
//! ~/.claude/projects/*/*.jsonl → 增量解析 → 去重 → 费用计算 → proxy_request_logs 表
|
||||
//! ```
|
||||
|
||||
use crate::config::get_claude_config_dir;
|
||||
use crate::database::{lock_conn, Database};
|
||||
use crate::error::AppError;
|
||||
use crate::proxy::usage::calculator::{CostCalculator, ModelPricing};
|
||||
use crate::proxy::usage::parser::TokenUsage;
|
||||
use rust_decimal::Decimal;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::HashMap;
|
||||
use std::fs;
|
||||
use std::io::{BufRead, BufReader};
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::time::SystemTime;
|
||||
|
||||
/// 同步结果
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct SessionSyncResult {
|
||||
pub imported: u32,
|
||||
pub skipped: u32,
|
||||
pub files_scanned: u32,
|
||||
pub errors: Vec<String>,
|
||||
}
|
||||
|
||||
/// 数据来源分布
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct DataSourceSummary {
|
||||
pub data_source: String,
|
||||
pub request_count: u32,
|
||||
pub total_cost_usd: String,
|
||||
}
|
||||
|
||||
/// 从 JSONL 中解析出的 assistant 消息使用数据
|
||||
#[derive(Debug)]
|
||||
struct ParsedAssistantUsage {
|
||||
message_id: String,
|
||||
model: String,
|
||||
input_tokens: u32,
|
||||
output_tokens: u32,
|
||||
cache_read_tokens: u32,
|
||||
cache_creation_tokens: u32,
|
||||
stop_reason: Option<String>,
|
||||
timestamp: Option<String>,
|
||||
session_id: Option<String>,
|
||||
}
|
||||
|
||||
/// 同步 Claude Code 会话日志到使用统计数据库
|
||||
pub fn sync_claude_session_logs(db: &Database) -> Result<SessionSyncResult, AppError> {
|
||||
let projects_dir = get_claude_config_dir().join("projects");
|
||||
if !projects_dir.exists() {
|
||||
return Ok(SessionSyncResult {
|
||||
imported: 0,
|
||||
skipped: 0,
|
||||
files_scanned: 0,
|
||||
errors: vec![],
|
||||
});
|
||||
}
|
||||
|
||||
let mut result = SessionSyncResult {
|
||||
imported: 0,
|
||||
skipped: 0,
|
||||
files_scanned: 0,
|
||||
errors: vec![],
|
||||
};
|
||||
|
||||
// 收集所有 .jsonl 文件
|
||||
let jsonl_files = collect_jsonl_files(&projects_dir);
|
||||
|
||||
for file_path in &jsonl_files {
|
||||
result.files_scanned += 1;
|
||||
|
||||
match sync_single_file(db, file_path) {
|
||||
Ok((imported, skipped)) => {
|
||||
result.imported += imported;
|
||||
result.skipped += skipped;
|
||||
}
|
||||
Err(e) => {
|
||||
let msg = format!("{}: {e}", file_path.display());
|
||||
log::warn!("[SESSION-SYNC] 文件解析失败: {msg}");
|
||||
result.errors.push(msg);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if result.imported > 0 {
|
||||
log::info!(
|
||||
"[SESSION-SYNC] 同步完成: 导入 {} 条, 跳过 {} 条, 扫描 {} 个文件",
|
||||
result.imported,
|
||||
result.skipped,
|
||||
result.files_scanned
|
||||
);
|
||||
}
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
/// 收集目录下所有 .jsonl 文件
|
||||
fn collect_jsonl_files(projects_dir: &Path) -> Vec<PathBuf> {
|
||||
let mut files = Vec::new();
|
||||
|
||||
let entries = match fs::read_dir(projects_dir) {
|
||||
Ok(e) => e,
|
||||
Err(_) => return files,
|
||||
};
|
||||
|
||||
for entry in entries.flatten() {
|
||||
let path = entry.path();
|
||||
if !path.is_dir() {
|
||||
continue;
|
||||
}
|
||||
// 每个项目目录下的 .jsonl 文件
|
||||
if let Ok(sub_entries) = fs::read_dir(&path) {
|
||||
for sub_entry in sub_entries.flatten() {
|
||||
let sub_path = sub_entry.path();
|
||||
if sub_path.extension().and_then(|e| e.to_str()) == Some("jsonl") {
|
||||
files.push(sub_path);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
files
|
||||
}
|
||||
|
||||
/// 同步单个 JSONL 文件，返回 (imported, skipped)
|
||||
fn sync_single_file(db: &Database, file_path: &Path) -> Result<(u32, u32), AppError> {
|
||||
let file_path_str = file_path.to_string_lossy().to_string();
|
||||
|
||||
// 获取文件元数据
|
||||
let metadata = fs::metadata(file_path)
|
||||
.map_err(|e| AppError::Config(format!("无法读取文件元数据: {e}")))?;
|
||||
let file_modified = metadata
|
||||
.modified()
|
||||
.ok()
|
||||
.and_then(|t| t.duration_since(SystemTime::UNIX_EPOCH).ok())
|
||||
.map(|d| d.as_secs() as i64)
|
||||
.unwrap_or(0);
|
||||
|
||||
// 检查同步状态
|
||||
let (last_modified, last_offset) = get_sync_state(db, &file_path_str)?;
|
||||
|
||||
// 文件未变化则跳过
|
||||
if file_modified <= last_modified {
|
||||
return Ok((0, 0));
|
||||
}
|
||||
|
||||
// 从上次偏移位置开始增量解析
|
||||
let file =
|
||||
fs::File::open(file_path).map_err(|e| AppError::Config(format!("无法打开文件: {e}")))?;
|
||||
let reader = BufReader::new(file);
|
||||
|
||||
let mut line_offset: i64 = 0;
|
||||
let mut messages: HashMap<String, ParsedAssistantUsage> = HashMap::new();
|
||||
let mut current_session_id: Option<String> = None;
|
||||
|
||||
for line_result in reader.lines() {
|
||||
line_offset += 1;
|
||||
|
||||
// 跳过已处理的行
|
||||
if line_offset <= last_offset {
|
||||
continue;
|
||||
}
|
||||
|
||||
let line = match line_result {
|
||||
Ok(l) => l,
|
||||
Err(_) => continue, // 容忍不完整的最后一行
|
||||
};
|
||||
|
||||
if line.trim().is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
let value: serde_json::Value = match serde_json::from_str(&line) {
|
||||
Ok(v) => v,
|
||||
Err(_) => continue,
|
||||
};
|
||||
|
||||
// 提取 session ID (从 system 或首条消息)
|
||||
if current_session_id.is_none() {
|
||||
if let Some(sid) = value.get("sessionId").and_then(|v| v.as_str()) {
|
||||
current_session_id = Some(sid.to_string());
|
||||
}
|
||||
}
|
||||
|
||||
// 只处理 assistant 类型的消息
|
||||
if value.get("type").and_then(|t| t.as_str()) != Some("assistant") {
|
||||
continue;
|
||||
}
|
||||
|
||||
let message = match value.get("message") {
|
||||
Some(m) => m,
|
||||
None => continue,
|
||||
};
|
||||
|
||||
let msg_id = match message.get("id").and_then(|v| v.as_str()) {
|
||||
Some(id) => id.to_string(),
|
||||
None => continue,
|
||||
};
|
||||
|
||||
let usage = match message.get("usage") {
|
||||
Some(u) => u,
|
||||
None => continue,
|
||||
};
|
||||
|
||||
let parsed = ParsedAssistantUsage {
|
||||
message_id: msg_id.clone(),
|
||||
model: message
|
||||
.get("model")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("unknown")
|
||||
.to_string(),
|
||||
input_tokens: usage
|
||||
.get("input_tokens")
|
||||
.and_then(|v| v.as_u64())
|
||||
.unwrap_or(0) as u32,
|
||||
output_tokens: usage
|
||||
.get("output_tokens")
|
||||
.and_then(|v| v.as_u64())
|
||||
.unwrap_or(0) as u32,
|
||||
cache_read_tokens: usage
|
||||
.get("cache_read_input_tokens")
|
||||
.and_then(|v| v.as_u64())
|
||||
.unwrap_or(0) as u32,
|
||||
cache_creation_tokens: usage
|
||||
.get("cache_creation_input_tokens")
|
||||
.and_then(|v| v.as_u64())
|
||||
.unwrap_or(0) as u32,
|
||||
stop_reason: message
|
||||
.get("stop_reason")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|s| s.to_string()),
|
||||
timestamp: value
|
||||
.get("timestamp")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|s| s.to_string()),
|
||||
session_id: current_session_id.clone(),
|
||||
};
|
||||
|
||||
// 按 message.id 去重：优先保留有 stop_reason 的条目，否则保留最新的
|
||||
let should_replace = match messages.get(&msg_id) {
|
||||
None => true,
|
||||
Some(existing) => {
|
||||
// 新条目有 stop_reason 而旧条目没有 → 替换
|
||||
if parsed.stop_reason.is_some() && existing.stop_reason.is_none() {
|
||||
true
|
||||
}
|
||||
// 两个都有或都没有 stop_reason → 取 output_tokens 更大的
|
||||
else if parsed.stop_reason.is_some() == existing.stop_reason.is_some() {
|
||||
parsed.output_tokens > existing.output_tokens
|
||||
} else {
|
||||
false
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
if should_replace {
|
||||
messages.insert(msg_id, parsed);
|
||||
}
|
||||
}
|
||||
|
||||
// 写入数据库
|
||||
let mut imported: u32 = 0;
|
||||
let mut skipped: u32 = 0;
|
||||
|
||||
for msg in messages.values() {
|
||||
// 只导入有 stop_reason 的最终条目（完整的 API 调用）
|
||||
if msg.stop_reason.is_none() {
|
||||
continue;
|
||||
}
|
||||
|
||||
let request_id = format!("session:{}", msg.message_id);
|
||||
|
||||
// 跳过 output_tokens 为 0 的无意义条目
|
||||
if msg.output_tokens == 0 {
|
||||
continue;
|
||||
}
|
||||
|
||||
match insert_session_log_entry(db, &request_id, msg) {
|
||||
Ok(true) => imported += 1,
|
||||
Ok(false) => skipped += 1,
|
||||
Err(e) => {
|
||||
log::warn!("[SESSION-SYNC] 插入失败 ({}): {e}", msg.message_id);
|
||||
skipped += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 更新同步状态
|
||||
update_sync_state(db, &file_path_str, file_modified, line_offset)?;
|
||||
|
||||
Ok((imported, skipped))
|
||||
}
|
||||
|
||||
/// 获取文件的同步状态
|
||||
fn get_sync_state(db: &Database, file_path: &str) -> Result<(i64, i64), AppError> {
|
||||
let conn = lock_conn!(db.conn);
|
||||
let result = conn.query_row(
|
||||
"SELECT last_modified, last_line_offset FROM session_log_sync WHERE file_path = ?1",
|
||||
rusqlite::params![file_path],
|
||||
|row| Ok((row.get::<_, i64>(0)?, row.get::<_, i64>(1)?)),
|
||||
);
|
||||
Ok(result.unwrap_or((0, 0)))
|
||||
}
|
||||
|
||||
/// 更新文件的同步状态
|
||||
fn update_sync_state(
|
||||
db: &Database,
|
||||
file_path: &str,
|
||||
last_modified: i64,
|
||||
last_offset: i64,
|
||||
) -> Result<(), AppError> {
|
||||
let now = SystemTime::now()
|
||||
.duration_since(SystemTime::UNIX_EPOCH)
|
||||
.map(|d| d.as_secs() as i64)
|
||||
.unwrap_or(0);
|
||||
|
||||
let conn = lock_conn!(db.conn);
|
||||
conn.execute(
|
||||
"INSERT OR REPLACE INTO session_log_sync (file_path, last_modified, last_line_offset, last_synced_at)
|
||||
VALUES (?1, ?2, ?3, ?4)",
|
||||
rusqlite::params![file_path, last_modified, last_offset, now],
|
||||
)
|
||||
.map_err(|e| AppError::Database(format!("更新同步状态失败: {e}")))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 插入单条会话日志到 proxy_request_logs，返回是否成功插入 (true=新插入, false=已存在)
|
||||
fn insert_session_log_entry(
|
||||
db: &Database,
|
||||
request_id: &str,
|
||||
msg: &ParsedAssistantUsage,
|
||||
) -> Result<bool, AppError> {
|
||||
let conn = lock_conn!(db.conn);
|
||||
|
||||
// 检查是否已存在
|
||||
let exists: bool = conn
|
||||
.query_row(
|
||||
"SELECT COUNT(*) FROM proxy_request_logs WHERE request_id = ?1",
|
||||
rusqlite::params![request_id],
|
||||
|row| row.get::<_, i64>(0).map(|c| c > 0),
|
||||
)
|
||||
.unwrap_or(false);
|
||||
|
||||
if exists {
|
||||
return Ok(false);
|
||||
}
|
||||
|
||||
// 解析时间戳
|
||||
let created_at = msg
|
||||
.timestamp
|
||||
.as_ref()
|
||||
.and_then(|ts| {
|
||||
// 尝试解析 ISO 8601 时间戳
|
||||
chrono::DateTime::parse_from_rfc3339(ts)
|
||||
.ok()
|
||||
.map(|dt| dt.timestamp())
|
||||
})
|
||||
.unwrap_or_else(|| {
|
||||
SystemTime::now()
|
||||
.duration_since(SystemTime::UNIX_EPOCH)
|
||||
.map(|d| d.as_secs() as i64)
|
||||
.unwrap_or(0)
|
||||
});
|
||||
|
||||
// 计算费用
|
||||
let usage = TokenUsage {
|
||||
input_tokens: msg.input_tokens,
|
||||
output_tokens: msg.output_tokens,
|
||||
cache_read_tokens: msg.cache_read_tokens,
|
||||
cache_creation_tokens: msg.cache_creation_tokens,
|
||||
model: Some(msg.model.clone()),
|
||||
};
|
||||
|
||||
let pricing = find_model_pricing_for_session(&conn, &msg.model);
|
||||
let multiplier = Decimal::from(1);
|
||||
let (input_cost, output_cost, cache_read_cost, cache_creation_cost, total_cost) = match pricing
|
||||
{
|
||||
Some(p) => {
|
||||
let cost = CostCalculator::calculate(&usage, &p, multiplier);
|
||||
(
|
||||
cost.input_cost.to_string(),
|
||||
cost.output_cost.to_string(),
|
||||
cost.cache_read_cost.to_string(),
|
||||
cost.cache_creation_cost.to_string(),
|
||||
cost.total_cost.to_string(),
|
||||
)
|
||||
}
|
||||
None => (
|
||||
"0".to_string(),
|
||||
"0".to_string(),
|
||||
"0".to_string(),
|
||||
"0".to_string(),
|
||||
"0".to_string(),
|
||||
),
|
||||
};
|
||||
|
||||
conn.execute(
|
||||
"INSERT OR IGNORE INTO proxy_request_logs (
|
||||
request_id, provider_id, app_type, model, request_model,
|
||||
input_tokens, output_tokens, cache_read_tokens, cache_creation_tokens,
|
||||
input_cost_usd, output_cost_usd, cache_read_cost_usd, cache_creation_cost_usd, total_cost_usd,
|
||||
latency_ms, first_token_ms, status_code, error_message, session_id,
|
||||
provider_type, is_streaming, cost_multiplier, created_at, data_source
|
||||
) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17, ?18, ?19, ?20, ?21, ?22, ?23, ?24)",
|
||||
rusqlite::params![
|
||||
request_id,
|
||||
"_session", // provider_id: 标记为会话来源
|
||||
"claude", // app_type
|
||||
msg.model,
|
||||
msg.model, // request_model = model
|
||||
msg.input_tokens,
|
||||
msg.output_tokens,
|
||||
msg.cache_read_tokens,
|
||||
msg.cache_creation_tokens,
|
||||
input_cost,
|
||||
output_cost,
|
||||
cache_read_cost,
|
||||
cache_creation_cost,
|
||||
total_cost,
|
||||
0i64, // latency_ms: 会话日志无此数据
|
||||
Option::<i64>::None, // first_token_ms
|
||||
200i64, // status_code: 有 stop_reason 说明请求成功
|
||||
Option::<String>::None, // error_message
|
||||
msg.session_id,
|
||||
Some("session_log"), // provider_type
|
||||
1i64, // is_streaming: Claude Code 通常使用流式
|
||||
"1.0", // cost_multiplier
|
||||
created_at,
|
||||
"session_log", // data_source
|
||||
],
|
||||
)
|
||||
.map_err(|e| AppError::Database(format!("插入会话日志失败: {e}")))?;
|
||||
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
/// 从 model_pricing 表查找模型定价（支持模糊匹配）
|
||||
fn find_model_pricing_for_session(
|
||||
conn: &rusqlite::Connection,
|
||||
model_id: &str,
|
||||
) -> Option<ModelPricing> {
|
||||
// 精确匹配
|
||||
if let Ok(Some(pricing)) = try_find_pricing(conn, model_id) {
|
||||
return Some(pricing);
|
||||
}
|
||||
|
||||
// 模糊匹配：去掉日期后缀
|
||||
// 例如 "claude-opus-4-6-20260206" -> "claude-opus-4-6"
|
||||
let parts: Vec<&str> = model_id.rsplitn(2, '-').collect();
|
||||
if parts.len() == 2 {
|
||||
if let Some(suffix) = parts.first() {
|
||||
if suffix.len() == 8 && suffix.chars().all(|c| c.is_ascii_digit()) {
|
||||
if let Ok(Some(pricing)) = try_find_pricing(conn, parts[1]) {
|
||||
return Some(pricing);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 尝试 LIKE 匹配
|
||||
let pattern = format!("{model_id}%");
|
||||
let result = conn.query_row(
|
||||
"SELECT input_cost_per_million, output_cost_per_million,
|
||||
cache_read_cost_per_million, cache_creation_cost_per_million
|
||||
FROM model_pricing WHERE model_id LIKE ?1 LIMIT 1",
|
||||
rusqlite::params![pattern],
|
||||
|row| {
|
||||
Ok((
|
||||
row.get::<_, String>(0)?,
|
||||
row.get::<_, String>(1)?,
|
||||
row.get::<_, String>(2)?,
|
||||
row.get::<_, String>(3)?,
|
||||
))
|
||||
},
|
||||
);
|
||||
|
||||
match result {
|
||||
Ok((input, output, cache_read, cache_creation)) => {
|
||||
ModelPricing::from_strings(&input, &output, &cache_read, &cache_creation).ok()
|
||||
}
|
||||
Err(_) => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn try_find_pricing(
|
||||
conn: &rusqlite::Connection,
|
||||
model_id: &str,
|
||||
) -> Result<Option<ModelPricing>, AppError> {
|
||||
let result = conn.query_row(
|
||||
"SELECT input_cost_per_million, output_cost_per_million,
|
||||
cache_read_cost_per_million, cache_creation_cost_per_million
|
||||
FROM model_pricing WHERE model_id = ?1",
|
||||
rusqlite::params![model_id],
|
||||
|row| {
|
||||
Ok((
|
||||
row.get::<_, String>(0)?,
|
||||
row.get::<_, String>(1)?,
|
||||
row.get::<_, String>(2)?,
|
||||
row.get::<_, String>(3)?,
|
||||
))
|
||||
},
|
||||
);
|
||||
|
||||
match result {
|
||||
Ok((input, output, cache_read, cache_creation)) => {
|
||||
ModelPricing::from_strings(&input, &output, &cache_read, &cache_creation)
|
||||
.map(Some)
|
||||
.map_err(|e| AppError::Database(format!("解析定价失败: {e}")))
|
||||
}
|
||||
Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None),
|
||||
Err(e) => Err(AppError::Database(format!("查询定价失败: {e}"))),
|
||||
}
|
||||
}
|
||||
|
||||
/// 查询数据来源分布统计
|
||||
pub fn get_data_source_breakdown(db: &Database) -> Result<Vec<DataSourceSummary>, AppError> {
|
||||
let conn = lock_conn!(db.conn);
|
||||
|
||||
let mut stmt = conn.prepare(
|
||||
"SELECT COALESCE(data_source, 'proxy') as ds, COUNT(*) as cnt,
|
||||
COALESCE(SUM(CAST(total_cost_usd AS REAL)), 0) as cost
|
||||
FROM proxy_request_logs
|
||||
GROUP BY ds
|
||||
ORDER BY cnt DESC",
|
||||
)?;
|
||||
|
||||
let rows = stmt.query_map([], |row| {
|
||||
Ok(DataSourceSummary {
|
||||
data_source: row.get(0)?,
|
||||
request_count: row.get::<_, i64>(1)? as u32,
|
||||
total_cost_usd: format!("{:.6}", row.get::<_, f64>(2)?),
|
||||
})
|
||||
})?;
|
||||
|
||||
let mut summaries = Vec::new();
|
||||
for row in rows {
|
||||
summaries.push(row.map_err(|e| AppError::Database(e.to_string()))?);
|
||||
}
|
||||
|
||||
Ok(summaries)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_parse_usage_from_jsonl_line() {
|
||||
let line = r#"{"type":"assistant","message":{"id":"msg_test123","model":"claude-opus-4-6","usage":{"input_tokens":3,"output_tokens":150,"cache_read_input_tokens":5000,"cache_creation_input_tokens":10000},"stop_reason":"end_turn"},"timestamp":"2026-04-05T12:00:00Z","sessionId":"session-abc"}"#;
|
||||
|
||||
let value: serde_json::Value = serde_json::from_str(line).unwrap();
|
||||
assert_eq!(
|
||||
value.get("type").and_then(|t| t.as_str()),
|
||||
Some("assistant")
|
||||
);
|
||||
|
||||
let message = value.get("message").unwrap();
|
||||
let usage = message.get("usage").unwrap();
|
||||
|
||||
assert_eq!(usage.get("input_tokens").unwrap().as_u64().unwrap(), 3);
|
||||
assert_eq!(usage.get("output_tokens").unwrap().as_u64().unwrap(), 150);
|
||||
assert_eq!(
|
||||
usage
|
||||
.get("cache_read_input_tokens")
|
||||
.unwrap()
|
||||
.as_u64()
|
||||
.unwrap(),
|
||||
5000
|
||||
);
|
||||
assert_eq!(
|
||||
usage
|
||||
.get("cache_creation_input_tokens")
|
||||
.unwrap()
|
||||
.as_u64()
|
||||
.unwrap(),
|
||||
10000
|
||||
);
|
||||
assert_eq!(
|
||||
message.get("stop_reason").unwrap().as_str().unwrap(),
|
||||
"end_turn"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_dedup_by_message_id() {
|
||||
// 同一个 message.id 有多条，应该取 stop_reason 有值的那条
|
||||
let mut messages: HashMap<String, ParsedAssistantUsage> = HashMap::new();
|
||||
|
||||
// 中间条目（无 stop_reason）
|
||||
let intermediate = ParsedAssistantUsage {
|
||||
message_id: "msg_1".to_string(),
|
||||
model: "claude-opus-4-6".to_string(),
|
||||
input_tokens: 3,
|
||||
output_tokens: 26,
|
||||
cache_read_tokens: 5000,
|
||||
cache_creation_tokens: 10000,
|
||||
stop_reason: None,
|
||||
timestamp: Some("2026-04-05T12:00:00Z".to_string()),
|
||||
session_id: None,
|
||||
};
|
||||
messages.insert("msg_1".to_string(), intermediate);
|
||||
|
||||
// 最终条目（有 stop_reason）
|
||||
let final_entry = ParsedAssistantUsage {
|
||||
message_id: "msg_1".to_string(),
|
||||
model: "claude-opus-4-6".to_string(),
|
||||
input_tokens: 3,
|
||||
output_tokens: 1349,
|
||||
cache_read_tokens: 5000,
|
||||
cache_creation_tokens: 10000,
|
||||
stop_reason: Some("end_turn".to_string()),
|
||||
timestamp: Some("2026-04-05T12:00:00Z".to_string()),
|
||||
session_id: None,
|
||||
};
|
||||
|
||||
// 应该替换
|
||||
let should_replace = final_entry.stop_reason.is_some()
|
||||
&& messages.get("msg_1").unwrap().stop_reason.is_none();
|
||||
assert!(should_replace);
|
||||
|
||||
messages.insert("msg_1".to_string(), final_entry);
|
||||
assert_eq!(messages.get("msg_1").unwrap().output_tokens, 1349);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,814 @@
|
||||
//! Codex 会话日志使用追踪
|
||||
//!
|
||||
//! 从 ~/.codex/sessions/ 下的 JSONL 会话文件中提取精确 token 使用数据，
|
||||
//! 替代原有的 state_5.sqlite 估算方案。
|
||||
//!
|
||||
//! ## 数据流
|
||||
//! ```text
|
||||
//! ~/.codex/sessions/YYYY/MM/DD/*.jsonl → 增量解析 → delta 计算 → 费用计算 → proxy_request_logs 表
|
||||
//! ```
|
||||
//!
|
||||
//! ## 解析的事件类型
|
||||
//! - `session_meta` → 提取 session_id
|
||||
//! - `turn_context` → 提取当前 model
|
||||
//! - `event_msg` (type=token_count) → 提取累计 token 用量，计算 delta
|
||||
|
||||
use crate::codex_config::get_codex_config_dir;
|
||||
use crate::database::{lock_conn, Database};
|
||||
use crate::error::AppError;
|
||||
use crate::proxy::usage::calculator::{CostCalculator, ModelPricing};
|
||||
use crate::proxy::usage::parser::TokenUsage;
|
||||
use crate::services::session_usage::SessionSyncResult;
|
||||
use rust_decimal::Decimal;
|
||||
use std::fs;
|
||||
use std::io::{BufRead, BufReader};
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::time::SystemTime;
|
||||
|
||||
/// 累计 token 用量（跟踪 total_token_usage 字段）
|
||||
#[derive(Debug, Clone, Default)]
|
||||
struct CumulativeTokens {
|
||||
input: u64,
|
||||
cached_input: u64,
|
||||
output: u64,
|
||||
}
|
||||
|
||||
/// 单次 API 调用的 token 增量
|
||||
#[derive(Debug)]
|
||||
struct DeltaTokens {
|
||||
input: u32,
|
||||
cached_input: u32,
|
||||
output: u32,
|
||||
}
|
||||
|
||||
impl DeltaTokens {
|
||||
fn is_zero(&self) -> bool {
|
||||
self.input == 0 && self.cached_input == 0 && self.output == 0
|
||||
}
|
||||
}
|
||||
|
||||
/// 单文件解析时的运行状态
|
||||
struct FileParseState {
|
||||
session_id: Option<String>,
|
||||
current_model: String,
|
||||
prev_total: Option<CumulativeTokens>,
|
||||
event_index: u32,
|
||||
}
|
||||
|
||||
/// 归一化 Codex 模型名
|
||||
///
|
||||
/// 处理规则（按顺序）：
|
||||
/// 1. 转小写：`GLM-4.6` → `glm-4.6`
|
||||
/// 2. 剥离 provider 前缀：`openai/gpt-5.4` → `gpt-5.4`
|
||||
/// 3. 剥离 ISO 日期后缀：`gpt-5.4-2026-03-05` → `gpt-5.4`
|
||||
/// 4. 剥离紧凑日期后缀：`gpt-5.4-20260305` → `gpt-5.4`
|
||||
fn normalize_codex_model(raw: &str) -> String {
|
||||
// Step 1: 小写
|
||||
let mut name = raw.to_lowercase();
|
||||
|
||||
// Step 2: 剥离 "provider/" 前缀（如 openai/, azure/）
|
||||
if let Some(pos) = name.rfind('/') {
|
||||
name = name[pos + 1..].to_string();
|
||||
}
|
||||
|
||||
// Step 3: 剥离 ISO 日期后缀 -YYYY-MM-DD（正好 11 字符）
|
||||
if name.len() > 11 {
|
||||
let suffix = &name[name.len() - 11..];
|
||||
if suffix.as_bytes()[0] == b'-'
|
||||
&& suffix[1..5].chars().all(|c| c.is_ascii_digit())
|
||||
&& suffix.as_bytes()[5] == b'-'
|
||||
&& suffix[6..8].chars().all(|c| c.is_ascii_digit())
|
||||
&& suffix.as_bytes()[8] == b'-'
|
||||
&& suffix[9..11].chars().all(|c| c.is_ascii_digit())
|
||||
{
|
||||
name.truncate(name.len() - 11);
|
||||
}
|
||||
}
|
||||
|
||||
// Step 4: 剥离紧凑日期后缀 -YYYYMMDD（正好 9 字符）
|
||||
if name.len() > 9 {
|
||||
let parts: Vec<&str> = name.rsplitn(2, '-').collect();
|
||||
if parts.len() == 2 {
|
||||
if let Some(suffix) = parts.first() {
|
||||
if suffix.len() == 8 && suffix.chars().all(|c| c.is_ascii_digit()) {
|
||||
name = parts[1].to_string();
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
name
|
||||
}
|
||||
|
||||
/// 计算两次累计值之间的 delta
|
||||
fn compute_delta(prev: &Option<CumulativeTokens>, current: &CumulativeTokens) -> DeltaTokens {
|
||||
match prev {
|
||||
None => DeltaTokens {
|
||||
input: current.input as u32,
|
||||
cached_input: current.cached_input as u32,
|
||||
output: current.output as u32,
|
||||
},
|
||||
Some(p) => DeltaTokens {
|
||||
input: current.input.saturating_sub(p.input) as u32,
|
||||
cached_input: current.cached_input.saturating_sub(p.cached_input) as u32,
|
||||
output: current.output.saturating_sub(p.output) as u32,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
/// 从 JSON Value 中提取累计 token 用量
|
||||
fn parse_cumulative_tokens(total_usage: &serde_json::Value) -> Option<CumulativeTokens> {
|
||||
if total_usage.is_null() || !total_usage.is_object() {
|
||||
return None;
|
||||
}
|
||||
Some(CumulativeTokens {
|
||||
input: total_usage
|
||||
.get("input_tokens")
|
||||
.and_then(|v| v.as_u64())
|
||||
.unwrap_or(0),
|
||||
cached_input: total_usage
|
||||
.get("cached_input_tokens")
|
||||
.or_else(|| total_usage.get("cache_read_input_tokens"))
|
||||
.and_then(|v| v.as_u64())
|
||||
.unwrap_or(0),
|
||||
output: total_usage
|
||||
.get("output_tokens")
|
||||
.and_then(|v| v.as_u64())
|
||||
.unwrap_or(0),
|
||||
})
|
||||
}
|
||||
|
||||
/// 同步 Codex 使用数据（从 JSONL 会话日志）
|
||||
pub fn sync_codex_usage(db: &Database) -> Result<SessionSyncResult, AppError> {
|
||||
let codex_dir = get_codex_config_dir();
|
||||
|
||||
let files = collect_codex_session_files(&codex_dir);
|
||||
|
||||
let mut result = SessionSyncResult {
|
||||
imported: 0,
|
||||
skipped: 0,
|
||||
files_scanned: files.len() as u32,
|
||||
errors: vec![],
|
||||
};
|
||||
|
||||
if files.is_empty() {
|
||||
return Ok(result);
|
||||
}
|
||||
|
||||
for file_path in &files {
|
||||
match sync_single_codex_file(db, file_path) {
|
||||
Ok((imported, skipped)) => {
|
||||
result.imported += imported;
|
||||
result.skipped += skipped;
|
||||
}
|
||||
Err(e) => {
|
||||
let msg = format!("Codex 会话文件解析失败 {}: {e}", file_path.display());
|
||||
log::warn!("[CODEX-SYNC] {msg}");
|
||||
result.errors.push(msg);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if result.imported > 0 {
|
||||
log::info!(
|
||||
"[CODEX-SYNC] 同步完成: 导入 {} 条, 跳过 {} 条, ��描 {} 个文件",
|
||||
result.imported,
|
||||
result.skipped,
|
||||
result.files_scanned
|
||||
);
|
||||
}
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
/// 收集所有 Codex 会话 JSONL 文件
|
||||
fn collect_codex_session_files(codex_dir: &Path) -> Vec<PathBuf> {
|
||||
let mut files = Vec::new();
|
||||
|
||||
// 1. 扫描 sessions/YYYY/MM/DD/*.jsonl（日期分区目录��
|
||||
let sessions_dir = codex_dir.join("sessions");
|
||||
if sessions_dir.is_dir() {
|
||||
collect_jsonl_recursive(&sessions_dir, &mut files, 0, 3);
|
||||
}
|
||||
|
||||
// 2. 扫描 archived_sessions/*.jsonl（扁平归档目录）
|
||||
let archived_dir = codex_dir.join("archived_sessions");
|
||||
if archived_dir.is_dir() {
|
||||
if let Ok(entries) = fs::read_dir(&archived_dir) {
|
||||
for entry in entries.flatten() {
|
||||
let path = entry.path();
|
||||
if path.extension().and_then(|e| e.to_str()) == Some("jsonl") {
|
||||
files.push(path);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
files
|
||||
}
|
||||
|
||||
/// 递归扫描目录下的 .jsonl 文件（限制最大深度）
|
||||
fn collect_jsonl_recursive(dir: &Path, files: &mut Vec<PathBuf>, depth: u32, max_depth: u32) {
|
||||
let entries = match fs::read_dir(dir) {
|
||||
Ok(e) => e,
|
||||
Err(_) => return,
|
||||
};
|
||||
|
||||
for entry in entries.flatten() {
|
||||
let path = entry.path();
|
||||
if path.is_dir() && depth < max_depth {
|
||||
collect_jsonl_recursive(&path, files, depth + 1, max_depth);
|
||||
} else if path.extension().and_then(|e| e.to_str()) == Some("jsonl") {
|
||||
files.push(path);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 同步单�� Codex JSONL 文件，返回 (imported, skipped)
|
||||
fn sync_single_codex_file(db: &Database, file_path: &Path) -> Result<(u32, u32), AppError> {
|
||||
let file_path_str = file_path.to_string_lossy().to_string();
|
||||
|
||||
// 获取文件元数据
|
||||
let metadata = fs::metadata(file_path)
|
||||
.map_err(|e| AppError::Config(format!("无法读取文���元数据: {e}")))?;
|
||||
let file_modified = metadata
|
||||
.modified()
|
||||
.ok()
|
||||
.and_then(|t| t.duration_since(SystemTime::UNIX_EPOCH).ok())
|
||||
.map(|d| d.as_secs() as i64)
|
||||
.unwrap_or(0);
|
||||
|
||||
// 检查同步状态
|
||||
let (last_modified, last_offset) = get_sync_state(db, &file_path_str)?;
|
||||
|
||||
// 文件未变化则跳过
|
||||
if file_modified <= last_modified {
|
||||
return Ok((0, 0));
|
||||
}
|
||||
|
||||
// 打开文件逐行解析
|
||||
let file =
|
||||
fs::File::open(file_path).map_err(|e| AppError::Config(format!("无法打开文件: {e}")))?;
|
||||
let reader = BufReader::new(file);
|
||||
|
||||
let mut state = FileParseState {
|
||||
session_id: None,
|
||||
current_model: "unknown".to_string(),
|
||||
prev_total: None,
|
||||
event_index: 0,
|
||||
};
|
||||
|
||||
let mut line_offset: i64 = 0;
|
||||
let mut imported: u32 = 0;
|
||||
let mut skipped: u32 = 0;
|
||||
|
||||
for line_result in reader.lines() {
|
||||
line_offset += 1;
|
||||
|
||||
let line = match line_result {
|
||||
Ok(l) => l,
|
||||
Err(_) => continue, // 容忍不完整的最后一行
|
||||
};
|
||||
|
||||
if line.trim().is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
// 快速过滤：在 JSON 反序列化前跳过无关行
|
||||
let is_event_msg = line.contains("\"event_msg\"");
|
||||
let is_turn_context = line.contains("\"turn_context\"");
|
||||
let is_session_meta = line.contains("\"session_meta\"");
|
||||
|
||||
if !is_event_msg && !is_turn_context && !is_session_meta {
|
||||
continue;
|
||||
}
|
||||
if is_event_msg && !line.contains("\"token_count\"") {
|
||||
continue;
|
||||
}
|
||||
|
||||
let value: serde_json::Value = match serde_json::from_str(&line) {
|
||||
Ok(v) => v,
|
||||
Err(_) => continue,
|
||||
};
|
||||
|
||||
let event_type = match value.get("type").and_then(|t| t.as_str()) {
|
||||
Some(t) => t,
|
||||
None => continue,
|
||||
};
|
||||
|
||||
match event_type {
|
||||
"session_meta" => {
|
||||
if state.session_id.is_none() {
|
||||
let payload = value.get("payload");
|
||||
state.session_id = payload
|
||||
.and_then(|p| {
|
||||
p.get("session_id")
|
||||
.or_else(|| p.get("sessionId"))
|
||||
.or_else(|| p.get("id"))
|
||||
})
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|s| s.to_string());
|
||||
}
|
||||
}
|
||||
"turn_context" => {
|
||||
if let Some(payload) = value.get("payload") {
|
||||
// model 可能在 payload.model 或 payload.info.model
|
||||
if let Some(model) = payload
|
||||
.get("model")
|
||||
.or_else(|| payload.get("info").and_then(|info| info.get("model")))
|
||||
.and_then(|v| v.as_str())
|
||||
{
|
||||
state.current_model = normalize_codex_model(model);
|
||||
}
|
||||
}
|
||||
}
|
||||
"event_msg" => {
|
||||
let payload = match value.get("payload") {
|
||||
Some(p) => p,
|
||||
None => continue,
|
||||
};
|
||||
|
||||
// 只处理 token_count 类型
|
||||
if payload.get("type").and_then(|t| t.as_str()) != Some("token_count") {
|
||||
continue;
|
||||
}
|
||||
|
||||
let info = match payload.get("info") {
|
||||
Some(i) if !i.is_null() => i,
|
||||
_ => continue, // info 为 null 的首个事件跳��
|
||||
};
|
||||
|
||||
// 提取模型（token_count 事件也可能携带 model）
|
||||
if let Some(model) = info
|
||||
.get("model")
|
||||
.or_else(|| info.get("model_name"))
|
||||
.or_else(|| payload.get("model"))
|
||||
.and_then(|v| v.as_str())
|
||||
{
|
||||
state.current_model = normalize_codex_model(model);
|
||||
}
|
||||
|
||||
// 优先用 total_token_usage（累计值），fallback 到 last_token_usage（增量值）
|
||||
let (cumulative, is_total) = if let Some(total) = info.get("total_token_usage") {
|
||||
(parse_cumulative_tokens(total), true)
|
||||
} else if let Some(last) = info.get("last_token_usage") {
|
||||
(parse_cumulative_tokens(last), false)
|
||||
} else {
|
||||
continue;
|
||||
};
|
||||
|
||||
let cumulative = match cumulative {
|
||||
Some(c) => c,
|
||||
None => continue,
|
||||
};
|
||||
|
||||
let delta = if is_total {
|
||||
// 累计值模式：计算与上次的 delta
|
||||
let d = compute_delta(&state.prev_total, &cumulative);
|
||||
state.prev_total = Some(cumulative);
|
||||
d
|
||||
} else {
|
||||
// 增量值模式：直接使用 last_token_usage 的值
|
||||
DeltaTokens {
|
||||
input: cumulative.input as u32,
|
||||
cached_input: cumulative.cached_input as u32,
|
||||
output: cumulative.output as u32,
|
||||
}
|
||||
};
|
||||
|
||||
// 钳制：cached 不应超过 input（防护异常数据）
|
||||
let delta = DeltaTokens {
|
||||
cached_input: delta.cached_input.min(delta.input),
|
||||
..delta
|
||||
};
|
||||
|
||||
if delta.is_zero() {
|
||||
continue; // 跳过 task 边界的零 delta 事件
|
||||
}
|
||||
|
||||
state.event_index += 1;
|
||||
|
||||
// 跳过已处理的行（但仍需解析以恢复状态）
|
||||
if line_offset <= last_offset {
|
||||
continue;
|
||||
}
|
||||
|
||||
// 生成唯一 request_id
|
||||
let session_id_str = state.session_id.as_deref().unwrap_or("unknown");
|
||||
let request_id = format!("codex_session:{}:{}", session_id_str, state.event_index);
|
||||
|
||||
// 提取时间戳
|
||||
let timestamp = value
|
||||
.get("timestamp")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|s| s.to_string());
|
||||
|
||||
match insert_codex_session_entry(
|
||||
db,
|
||||
&request_id,
|
||||
&delta,
|
||||
&state.current_model,
|
||||
state.session_id.as_deref(),
|
||||
timestamp.as_deref(),
|
||||
) {
|
||||
Ok(true) => imported += 1,
|
||||
Ok(false) => skipped += 1,
|
||||
Err(e) => {
|
||||
log::warn!("[CODEX-SYNC] 插入失败 ({}): {e}", request_id);
|
||||
skipped += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
// 更新同步状态
|
||||
update_sync_state(db, &file_path_str, file_modified, line_offset)?;
|
||||
|
||||
Ok((imported, skipped))
|
||||
}
|
||||
|
||||
/// 插入单条 Codex 会话记录到 proxy_request_logs
|
||||
fn insert_codex_session_entry(
|
||||
db: &Database,
|
||||
request_id: &str,
|
||||
delta: &DeltaTokens,
|
||||
model: &str,
|
||||
session_id: Option<&str>,
|
||||
timestamp: Option<&str>,
|
||||
) -> Result<bool, AppError> {
|
||||
let conn = lock_conn!(db.conn);
|
||||
|
||||
// 检查是否已存在
|
||||
let exists: bool = conn
|
||||
.query_row(
|
||||
"SELECT COUNT(*) FROM proxy_request_logs WHERE request_id = ?1",
|
||||
rusqlite::params![request_id],
|
||||
|row| row.get::<_, i64>(0).map(|c| c > 0),
|
||||
)
|
||||
.unwrap_or(false);
|
||||
|
||||
if exists {
|
||||
return Ok(false);
|
||||
}
|
||||
|
||||
// 解析时间戳
|
||||
let created_at = timestamp
|
||||
.and_then(|ts| {
|
||||
chrono::DateTime::parse_from_rfc3339(ts)
|
||||
.ok()
|
||||
.map(|dt| dt.timestamp())
|
||||
})
|
||||
.unwrap_or_else(|| {
|
||||
SystemTime::now()
|
||||
.duration_since(SystemTime::UNIX_EPOCH)
|
||||
.map(|d| d.as_secs() as i64)
|
||||
.unwrap_or(0)
|
||||
});
|
||||
|
||||
// 计算费用
|
||||
let usage = TokenUsage {
|
||||
input_tokens: delta.input,
|
||||
output_tokens: delta.output,
|
||||
cache_read_tokens: delta.cached_input,
|
||||
cache_creation_tokens: 0,
|
||||
model: Some(model.to_string()),
|
||||
};
|
||||
|
||||
let pricing = find_codex_pricing(&conn, model);
|
||||
let multiplier = Decimal::from(1);
|
||||
let (input_cost, output_cost, cache_read_cost, cache_creation_cost, total_cost) = match pricing
|
||||
{
|
||||
Some(p) => {
|
||||
let cost = CostCalculator::calculate(&usage, &p, multiplier);
|
||||
(
|
||||
cost.input_cost.to_string(),
|
||||
cost.output_cost.to_string(),
|
||||
cost.cache_read_cost.to_string(),
|
||||
cost.cache_creation_cost.to_string(),
|
||||
cost.total_cost.to_string(),
|
||||
)
|
||||
}
|
||||
None => (
|
||||
"0".to_string(),
|
||||
"0".to_string(),
|
||||
"0".to_string(),
|
||||
"0".to_string(),
|
||||
"0".to_string(),
|
||||
),
|
||||
};
|
||||
|
||||
conn.execute(
|
||||
"INSERT OR IGNORE INTO proxy_request_logs (
|
||||
request_id, provider_id, app_type, model, request_model,
|
||||
input_tokens, output_tokens, cache_read_tokens, cache_creation_tokens,
|
||||
input_cost_usd, output_cost_usd, cache_read_cost_usd, cache_creation_cost_usd, total_cost_usd,
|
||||
latency_ms, first_token_ms, status_code, error_message, session_id,
|
||||
provider_type, is_streaming, cost_multiplier, created_at, data_source
|
||||
) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17, ?18, ?19, ?20, ?21, ?22, ?23, ?24)",
|
||||
rusqlite::params![
|
||||
request_id,
|
||||
"_codex_session", // provider_id
|
||||
"codex", // app_type
|
||||
model,
|
||||
model, // request_model = model
|
||||
delta.input,
|
||||
delta.output,
|
||||
delta.cached_input,
|
||||
0i64, // cache_creation_tokens: Codex 日志无此数据
|
||||
input_cost,
|
||||
output_cost,
|
||||
cache_read_cost,
|
||||
cache_creation_cost,
|
||||
total_cost,
|
||||
0i64, // latency_ms
|
||||
Option::<i64>::None, // first_token_ms
|
||||
200i64, // status_code
|
||||
Option::<String>::None, // error_message
|
||||
session_id.map(|s| s.to_string()),
|
||||
Some("codex_session"), // provider_type
|
||||
1i64, // is_streaming
|
||||
"1.0", // cost_multiplier
|
||||
created_at,
|
||||
"codex_session", // data_source
|
||||
],
|
||||
)
|
||||
.map_err(|e| AppError::Database(format!("插入 Codex 会话日志失败: {e}")))?;
|
||||
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
/// 获取文件的同步状态
|
||||
fn get_sync_state(db: &Database, file_path: &str) -> Result<(i64, i64), AppError> {
|
||||
let conn = lock_conn!(db.conn);
|
||||
let result = conn.query_row(
|
||||
"SELECT last_modified, last_line_offset FROM session_log_sync WHERE file_path = ?1",
|
||||
rusqlite::params![file_path],
|
||||
|row| Ok((row.get::<_, i64>(0)?, row.get::<_, i64>(1)?)),
|
||||
);
|
||||
Ok(result.unwrap_or((0, 0)))
|
||||
}
|
||||
|
||||
/// 更新文件的同步状态
|
||||
fn update_sync_state(
|
||||
db: &Database,
|
||||
file_path: &str,
|
||||
last_modified: i64,
|
||||
last_offset: i64,
|
||||
) -> Result<(), AppError> {
|
||||
let now = SystemTime::now()
|
||||
.duration_since(SystemTime::UNIX_EPOCH)
|
||||
.map(|d| d.as_secs() as i64)
|
||||
.unwrap_or(0);
|
||||
|
||||
let conn = lock_conn!(db.conn);
|
||||
conn.execute(
|
||||
"INSERT OR REPLACE INTO session_log_sync (file_path, last_modified, last_line_offset, last_synced_at)
|
||||
VALUES (?1, ?2, ?3, ?4)",
|
||||
rusqlite::params![file_path, last_modified, last_offset, now],
|
||||
)
|
||||
.map_err(|e| AppError::Database(format!("更新同步状态失败: {e}")))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// ��找 Codex 模型定价（带归一化）
|
||||
fn find_codex_pricing(conn: &rusqlite::Connection, model_id: &str) -> Option<ModelPricing> {
|
||||
let normalized = normalize_codex_model(model_id);
|
||||
|
||||
// 1. 精确匹配（归一化后的名称）
|
||||
if let Some(pricing) = try_find_pricing(conn, &normalized) {
|
||||
return Some(pricing);
|
||||
}
|
||||
|
||||
// 2. LIKE 模糊匹配（兜底）
|
||||
let pattern = format!("{normalized}%");
|
||||
conn.query_row(
|
||||
"SELECT input_cost_per_million, output_cost_per_million,
|
||||
cache_read_cost_per_million, cache_creation_cost_per_million
|
||||
FROM model_pricing WHERE model_id LIKE ?1 LIMIT 1",
|
||||
rusqlite::params![pattern],
|
||||
|row| {
|
||||
Ok((
|
||||
row.get::<_, String>(0)?,
|
||||
row.get::<_, String>(1)?,
|
||||
row.get::<_, String>(2)?,
|
||||
row.get::<_, String>(3)?,
|
||||
))
|
||||
},
|
||||
)
|
||||
.ok()
|
||||
.and_then(|(i, o, cr, cc)| ModelPricing::from_strings(&i, &o, &cr, &cc).ok())
|
||||
}
|
||||
|
||||
/// 精确匹配定价查询
|
||||
fn try_find_pricing(conn: &rusqlite::Connection, model_id: &str) -> Option<ModelPricing> {
|
||||
conn.query_row(
|
||||
"SELECT input_cost_per_million, output_cost_per_million,
|
||||
cache_read_cost_per_million, cache_creation_cost_per_million
|
||||
FROM model_pricing WHERE model_id = ?1",
|
||||
rusqlite::params![model_id],
|
||||
|row| {
|
||||
Ok((
|
||||
row.get::<_, String>(0)?,
|
||||
row.get::<_, String>(1)?,
|
||||
row.get::<_, String>(2)?,
|
||||
row.get::<_, String>(3)?,
|
||||
))
|
||||
},
|
||||
)
|
||||
.ok()
|
||||
.and_then(|(i, o, cr, cc)| ModelPricing::from_strings(&i, &o, &cr, &cc).ok())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_delta_first_event() {
|
||||
let prev = None;
|
||||
let current = CumulativeTokens {
|
||||
input: 17934,
|
||||
cached_input: 9600,
|
||||
output: 454,
|
||||
};
|
||||
let delta = compute_delta(&prev, ¤t);
|
||||
assert_eq!(delta.input, 17934);
|
||||
assert_eq!(delta.cached_input, 9600);
|
||||
assert_eq!(delta.output, 454);
|
||||
assert!(!delta.is_zero());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_delta_subsequent_event() {
|
||||
let prev = Some(CumulativeTokens {
|
||||
input: 17934,
|
||||
cached_input: 9600,
|
||||
output: 454,
|
||||
});
|
||||
let current = CumulativeTokens {
|
||||
input: 36722,
|
||||
cached_input: 27904,
|
||||
output: 804,
|
||||
};
|
||||
let delta = compute_delta(&prev, ¤t);
|
||||
assert_eq!(delta.input, 36722 - 17934);
|
||||
assert_eq!(delta.cached_input, 27904 - 9600);
|
||||
assert_eq!(delta.output, 804 - 454);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_delta_zero_at_task_boundary() {
|
||||
let prev = Some(CumulativeTokens {
|
||||
input: 58346,
|
||||
cached_input: 46976,
|
||||
output: 1045,
|
||||
});
|
||||
// task 边界：相同的累计值
|
||||
let current = CumulativeTokens {
|
||||
input: 58346,
|
||||
cached_input: 46976,
|
||||
output: 1045,
|
||||
};
|
||||
let delta = compute_delta(&prev, ¤t);
|
||||
assert!(delta.is_zero());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_delta_saturating_sub() {
|
||||
// 异常情况：当前值小于前值（不应发生，但需防护）
|
||||
let prev = Some(CumulativeTokens {
|
||||
input: 100,
|
||||
cached_input: 50,
|
||||
output: 30,
|
||||
});
|
||||
let current = CumulativeTokens {
|
||||
input: 80,
|
||||
cached_input: 40,
|
||||
output: 20,
|
||||
};
|
||||
let delta = compute_delta(&prev, ¤t);
|
||||
assert_eq!(delta.input, 0);
|
||||
assert_eq!(delta.cached_input, 0);
|
||||
assert_eq!(delta.output, 0);
|
||||
assert!(delta.is_zero());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_cumulative_tokens_valid() {
|
||||
let json: serde_json::Value = serde_json::json!({
|
||||
"input_tokens": 17934,
|
||||
"cached_input_tokens": 9600,
|
||||
"output_tokens": 454,
|
||||
"reasoning_output_tokens": 233,
|
||||
"total_tokens": 18388
|
||||
});
|
||||
let tokens = parse_cumulative_tokens(&json).unwrap();
|
||||
assert_eq!(tokens.input, 17934);
|
||||
assert_eq!(tokens.cached_input, 9600);
|
||||
assert_eq!(tokens.output, 454);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_cumulative_tokens_null() {
|
||||
let json = serde_json::Value::Null;
|
||||
assert!(parse_cumulative_tokens(&json).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_cumulative_tokens_alt_field_names() {
|
||||
// 某些版本可能使用 cache_read_input_tokens 而非 cached_input_tokens
|
||||
let json: serde_json::Value = serde_json::json!({
|
||||
"input_tokens": 1000,
|
||||
"cache_read_input_tokens": 500,
|
||||
"output_tokens": 200
|
||||
});
|
||||
let tokens = parse_cumulative_tokens(&json).unwrap();
|
||||
assert_eq!(tokens.cached_input, 500);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_collect_codex_session_files_nonexistent() {
|
||||
let files = collect_codex_session_files(Path::new("/nonexistent/path"));
|
||||
assert!(files.is_empty());
|
||||
}
|
||||
|
||||
// ── 模型名归一化测试 ──
|
||||
|
||||
#[test]
|
||||
fn test_normalize_codex_model_lowercase() {
|
||||
assert_eq!(normalize_codex_model("GLM-4.6"), "glm-4.6");
|
||||
assert_eq!(normalize_codex_model("DeepSeek-Chat"), "deepseek-chat");
|
||||
assert_eq!(normalize_codex_model("GPT-5.4"), "gpt-5.4");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_normalize_codex_model_strip_prefix() {
|
||||
assert_eq!(normalize_codex_model("openai/gpt-5.4"), "gpt-5.4");
|
||||
assert_eq!(
|
||||
normalize_codex_model("azure/gpt-5.2-codex"),
|
||||
"gpt-5.2-codex"
|
||||
);
|
||||
assert_eq!(normalize_codex_model("OPENAI/GPT-5.4"), "gpt-5.4");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_normalize_codex_model_strip_iso_date() {
|
||||
assert_eq!(normalize_codex_model("gpt-5.4-2026-03-05"), "gpt-5.4");
|
||||
assert_eq!(
|
||||
normalize_codex_model("gpt-5.4-pro-2026-03-05"),
|
||||
"gpt-5.4-pro"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_normalize_codex_model_strip_compact_date() {
|
||||
assert_eq!(normalize_codex_model("gpt-5.4-20260305"), "gpt-5.4");
|
||||
assert_eq!(
|
||||
normalize_codex_model("claude-opus-4-6-20260206"),
|
||||
"claude-opus-4-6"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_normalize_codex_model_no_change() {
|
||||
assert_eq!(normalize_codex_model("gpt-5.4"), "gpt-5.4");
|
||||
assert_eq!(normalize_codex_model("gpt-5.2-codex"), "gpt-5.2-codex");
|
||||
assert_eq!(normalize_codex_model("o3"), "o3");
|
||||
assert_eq!(normalize_codex_model("deepseek-chat"), "deepseek-chat");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_normalize_codex_model_combined() {
|
||||
// prefix + uppercase + ISO date
|
||||
assert_eq!(
|
||||
normalize_codex_model("openai/GPT-5.4-2026-03-05"),
|
||||
"gpt-5.4"
|
||||
);
|
||||
// prefix + compact date
|
||||
assert_eq!(normalize_codex_model("openai/gpt-5.4-20260305"), "gpt-5.4");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_cached_clamped_to_input() {
|
||||
// cached > input 的异常场景应被 min() 钳制
|
||||
let prev = Some(CumulativeTokens {
|
||||
input: 100,
|
||||
cached_input: 0,
|
||||
output: 50,
|
||||
});
|
||||
let current = CumulativeTokens {
|
||||
input: 110, // delta = 10
|
||||
cached_input: 80, // delta = 80（异常：大于 input delta）
|
||||
output: 60,
|
||||
};
|
||||
let delta = compute_delta(&prev, ¤t);
|
||||
// 钳制前：cached_input = 80, input = 10
|
||||
assert_eq!(delta.cached_input, 80);
|
||||
assert_eq!(delta.input, 10);
|
||||
// 实际钳制在调用侧：delta.cached_input.min(delta.input)
|
||||
let clamped = delta.cached_input.min(delta.input);
|
||||
assert_eq!(clamped, 10);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,503 @@
|
||||
//! Gemini CLI 会话日志使用追踪
|
||||
//!
|
||||
//! 从 ~/.gemini/tmp/<project_hash>/chats/session-*.json 中提取精确 token 使用数据。
|
||||
//!
|
||||
//! ## 数据流
|
||||
//! ```text
|
||||
//! ~/.gemini/tmp/*/chats/session-*.json → 全量解析 → 费用计算 → proxy_request_logs 表
|
||||
//! ```
|
||||
//!
|
||||
//! ## 与 Claude/Codex 解析器的差异
|
||||
//! - JSON 格式（非 JSONL）：每个文件是单个 JSON 对象，包含 messages 数组
|
||||
//! - 无需 delta 计算：tokens 字段是 per-message 独立值
|
||||
//! - 无需状态恢复：不依赖前一条消息的累计值
|
||||
//! - 天然去重：每条消息有唯一 id 字段
|
||||
|
||||
use crate::database::{lock_conn, Database};
|
||||
use crate::error::AppError;
|
||||
use crate::gemini_config::get_gemini_dir;
|
||||
use crate::proxy::usage::calculator::{CostCalculator, ModelPricing};
|
||||
use crate::proxy::usage::parser::TokenUsage;
|
||||
use crate::services::session_usage::SessionSyncResult;
|
||||
use rust_decimal::Decimal;
|
||||
use std::fs;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::time::SystemTime;
|
||||
|
||||
/// 从 Gemini message 中提取的 token 数据
|
||||
#[derive(Debug)]
|
||||
struct GeminiTokens {
|
||||
input: u32,
|
||||
output: u32,
|
||||
cached: u32,
|
||||
thoughts: u32,
|
||||
}
|
||||
|
||||
/// 同步 Gemini 使用数据（从 JSON 会话日志）
|
||||
pub fn sync_gemini_usage(db: &Database) -> Result<SessionSyncResult, AppError> {
|
||||
let gemini_dir = get_gemini_dir();
|
||||
|
||||
let files = collect_gemini_session_files(&gemini_dir);
|
||||
|
||||
let mut result = SessionSyncResult {
|
||||
imported: 0,
|
||||
skipped: 0,
|
||||
files_scanned: files.len() as u32,
|
||||
errors: vec![],
|
||||
};
|
||||
|
||||
if files.is_empty() {
|
||||
return Ok(result);
|
||||
}
|
||||
|
||||
for file_path in &files {
|
||||
match sync_single_gemini_file(db, file_path) {
|
||||
Ok((imported, skipped)) => {
|
||||
result.imported += imported;
|
||||
result.skipped += skipped;
|
||||
}
|
||||
Err(e) => {
|
||||
let msg = format!("Gemini 会话文件解析失败 {}: {e}", file_path.display());
|
||||
log::warn!("[GEMINI-SYNC] {msg}");
|
||||
result.errors.push(msg);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if result.imported > 0 {
|
||||
log::info!(
|
||||
"[GEMINI-SYNC] 同步完成: 导入 {} 条, 跳过 {} 条, 扫描 {} 个文件",
|
||||
result.imported,
|
||||
result.skipped,
|
||||
result.files_scanned
|
||||
);
|
||||
}
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
/// 收集所有 Gemini 会话 JSON 文件
|
||||
fn collect_gemini_session_files(gemini_dir: &Path) -> Vec<PathBuf> {
|
||||
let mut files = Vec::new();
|
||||
|
||||
let tmp_dir = gemini_dir.join("tmp");
|
||||
if !tmp_dir.is_dir() {
|
||||
return files;
|
||||
}
|
||||
|
||||
// 遍历 tmp/<project_hash>/chats/session-*.json
|
||||
let project_dirs = match fs::read_dir(&tmp_dir) {
|
||||
Ok(entries) => entries,
|
||||
Err(_) => return files,
|
||||
};
|
||||
|
||||
for entry in project_dirs.flatten() {
|
||||
let chats_dir = entry.path().join("chats");
|
||||
if !chats_dir.is_dir() {
|
||||
continue;
|
||||
}
|
||||
|
||||
let chat_files = match fs::read_dir(&chats_dir) {
|
||||
Ok(entries) => entries,
|
||||
Err(_) => continue,
|
||||
};
|
||||
|
||||
for file_entry in chat_files.flatten() {
|
||||
let path = file_entry.path();
|
||||
let is_session = path
|
||||
.file_name()
|
||||
.and_then(|n| n.to_str())
|
||||
.map(|n| n.starts_with("session-") && n.ends_with(".json"))
|
||||
.unwrap_or(false);
|
||||
if is_session {
|
||||
files.push(path);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
files
|
||||
}
|
||||
|
||||
/// 同步单个 Gemini 会话 JSON 文件，返回 (imported, skipped)
|
||||
fn sync_single_gemini_file(db: &Database, file_path: &Path) -> Result<(u32, u32), AppError> {
|
||||
let file_path_str = file_path.to_string_lossy().to_string();
|
||||
|
||||
// 获取文件元数据
|
||||
let metadata = fs::metadata(file_path)
|
||||
.map_err(|e| AppError::Config(format!("无法读取文件元数据: {e}")))?;
|
||||
let file_modified = metadata
|
||||
.modified()
|
||||
.ok()
|
||||
.and_then(|t| t.duration_since(SystemTime::UNIX_EPOCH).ok())
|
||||
.map(|d| d.as_secs() as i64)
|
||||
.unwrap_or(0);
|
||||
|
||||
// 检查同步状态
|
||||
let (last_modified, _last_offset) = get_sync_state(db, &file_path_str)?;
|
||||
|
||||
// 文件未变化则跳过
|
||||
if file_modified <= last_modified {
|
||||
return Ok((0, 0));
|
||||
}
|
||||
|
||||
// 读取并解析整个 JSON 文件
|
||||
let content = fs::read_to_string(file_path)
|
||||
.map_err(|e| AppError::Config(format!("无法读取文件: {e}")))?;
|
||||
let value: serde_json::Value = serde_json::from_str(&content)
|
||||
.map_err(|e| AppError::Config(format!("JSON 解析失败: {e}")))?;
|
||||
|
||||
// 提取顶层 sessionId
|
||||
let session_id = value
|
||||
.get("sessionId")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|s| s.to_string());
|
||||
|
||||
// 遍历 messages 数组
|
||||
let messages = match value.get("messages").and_then(|v| v.as_array()) {
|
||||
Some(msgs) => msgs,
|
||||
None => return Ok((0, 0)),
|
||||
};
|
||||
|
||||
let mut imported: u32 = 0;
|
||||
let mut skipped: u32 = 0;
|
||||
let mut gemini_msg_count: i64 = 0;
|
||||
|
||||
for msg in messages {
|
||||
// 只处理 type == "gemini" 的消息
|
||||
if msg.get("type").and_then(|t| t.as_str()) != Some("gemini") {
|
||||
continue;
|
||||
}
|
||||
|
||||
// 提取 tokens 对象
|
||||
let tokens_obj = match msg.get("tokens") {
|
||||
Some(t) if t.is_object() => t,
|
||||
_ => continue,
|
||||
};
|
||||
|
||||
let tokens = parse_gemini_tokens(tokens_obj);
|
||||
if tokens.input == 0 && tokens.output == 0 && tokens.thoughts == 0 && tokens.cached == 0 {
|
||||
continue; // 跳过全零的空 token 消息
|
||||
}
|
||||
|
||||
gemini_msg_count += 1;
|
||||
|
||||
// 提取消息 ID 和模型
|
||||
let message_id = msg.get("id").and_then(|v| v.as_str()).unwrap_or("unknown");
|
||||
let model = msg
|
||||
.get("model")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("unknown");
|
||||
let timestamp = msg.get("timestamp").and_then(|v| v.as_str());
|
||||
|
||||
// 生成唯一 request_id
|
||||
let session_id_str = session_id.as_deref().unwrap_or("unknown");
|
||||
let request_id = format!("gemini_session:{session_id_str}:{message_id}");
|
||||
|
||||
match insert_gemini_session_entry(
|
||||
db,
|
||||
&request_id,
|
||||
&tokens,
|
||||
model,
|
||||
session_id.as_deref(),
|
||||
timestamp,
|
||||
) {
|
||||
Ok(true) => imported += 1,
|
||||
Ok(false) => skipped += 1,
|
||||
Err(e) => {
|
||||
log::warn!("[GEMINI-SYNC] 插入失败 ({}): {e}", request_id);
|
||||
skipped += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 更新同步状态
|
||||
update_sync_state(db, &file_path_str, file_modified, gemini_msg_count)?;
|
||||
|
||||
Ok((imported, skipped))
|
||||
}
|
||||
|
||||
/// 从 tokens JSON 对象中提取 token 数据
|
||||
fn parse_gemini_tokens(tokens: &serde_json::Value) -> GeminiTokens {
|
||||
GeminiTokens {
|
||||
input: tokens.get("input").and_then(|v| v.as_u64()).unwrap_or(0) as u32,
|
||||
output: tokens.get("output").and_then(|v| v.as_u64()).unwrap_or(0) as u32,
|
||||
cached: tokens.get("cached").and_then(|v| v.as_u64()).unwrap_or(0) as u32,
|
||||
thoughts: tokens.get("thoughts").and_then(|v| v.as_u64()).unwrap_or(0) as u32,
|
||||
}
|
||||
}
|
||||
|
||||
/// 插入单条 Gemini 会话记录到 proxy_request_logs
|
||||
fn insert_gemini_session_entry(
|
||||
db: &Database,
|
||||
request_id: &str,
|
||||
tokens: &GeminiTokens,
|
||||
model: &str,
|
||||
session_id: Option<&str>,
|
||||
timestamp: Option<&str>,
|
||||
) -> Result<bool, AppError> {
|
||||
let conn = lock_conn!(db.conn);
|
||||
|
||||
// 解析时间戳
|
||||
let created_at = timestamp
|
||||
.and_then(|ts| {
|
||||
chrono::DateTime::parse_from_rfc3339(ts)
|
||||
.ok()
|
||||
.map(|dt| dt.timestamp())
|
||||
})
|
||||
.unwrap_or_else(|| {
|
||||
SystemTime::now()
|
||||
.duration_since(SystemTime::UNIX_EPOCH)
|
||||
.map(|d| d.as_secs() as i64)
|
||||
.unwrap_or(0)
|
||||
});
|
||||
|
||||
// 合并 thoughts 到 output（思考 token 按输出计费）
|
||||
let output_tokens = tokens.output + tokens.thoughts;
|
||||
|
||||
// 计算费用
|
||||
let usage = TokenUsage {
|
||||
input_tokens: tokens.input,
|
||||
output_tokens,
|
||||
cache_read_tokens: tokens.cached,
|
||||
cache_creation_tokens: 0,
|
||||
model: Some(model.to_string()),
|
||||
};
|
||||
|
||||
let pricing = find_gemini_pricing(&conn, model);
|
||||
let multiplier = Decimal::from(1);
|
||||
let (input_cost, output_cost, cache_read_cost, cache_creation_cost, total_cost) = match pricing
|
||||
{
|
||||
Some(p) => {
|
||||
let cost = CostCalculator::calculate(&usage, &p, multiplier);
|
||||
(
|
||||
cost.input_cost.to_string(),
|
||||
cost.output_cost.to_string(),
|
||||
cost.cache_read_cost.to_string(),
|
||||
cost.cache_creation_cost.to_string(),
|
||||
cost.total_cost.to_string(),
|
||||
)
|
||||
}
|
||||
None => (
|
||||
"0".to_string(),
|
||||
"0".to_string(),
|
||||
"0".to_string(),
|
||||
"0".to_string(),
|
||||
"0".to_string(),
|
||||
),
|
||||
};
|
||||
|
||||
// 使用 UPSERT：新记录插入，已存在记录更新 token 和费用（Gemini 全量重读可能携带更新值）
|
||||
conn.execute(
|
||||
"INSERT INTO proxy_request_logs (
|
||||
request_id, provider_id, app_type, model, request_model,
|
||||
input_tokens, output_tokens, cache_read_tokens, cache_creation_tokens,
|
||||
input_cost_usd, output_cost_usd, cache_read_cost_usd, cache_creation_cost_usd, total_cost_usd,
|
||||
latency_ms, first_token_ms, status_code, error_message, session_id,
|
||||
provider_type, is_streaming, cost_multiplier, created_at, data_source
|
||||
) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17, ?18, ?19, ?20, ?21, ?22, ?23, ?24)
|
||||
ON CONFLICT(request_id) DO UPDATE SET
|
||||
model = excluded.model,
|
||||
input_tokens = excluded.input_tokens,
|
||||
output_tokens = excluded.output_tokens,
|
||||
cache_read_tokens = excluded.cache_read_tokens,
|
||||
input_cost_usd = excluded.input_cost_usd,
|
||||
output_cost_usd = excluded.output_cost_usd,
|
||||
cache_read_cost_usd = excluded.cache_read_cost_usd,
|
||||
cache_creation_cost_usd = excluded.cache_creation_cost_usd,
|
||||
total_cost_usd = excluded.total_cost_usd
|
||||
WHERE input_tokens != excluded.input_tokens
|
||||
OR output_tokens != excluded.output_tokens
|
||||
OR cache_read_tokens != excluded.cache_read_tokens
|
||||
OR model != excluded.model",
|
||||
rusqlite::params![
|
||||
request_id,
|
||||
"_gemini_session", // provider_id
|
||||
"gemini", // app_type
|
||||
model,
|
||||
model, // request_model = model
|
||||
tokens.input,
|
||||
output_tokens,
|
||||
tokens.cached,
|
||||
0i64, // cache_creation_tokens
|
||||
input_cost,
|
||||
output_cost,
|
||||
cache_read_cost,
|
||||
cache_creation_cost,
|
||||
total_cost,
|
||||
0i64, // latency_ms
|
||||
Option::<i64>::None, // first_token_ms
|
||||
200i64, // status_code
|
||||
Option::<String>::None, // error_message
|
||||
session_id.map(|s| s.to_string()),
|
||||
Some("gemini_session"), // provider_type
|
||||
1i64, // is_streaming
|
||||
"1.0", // cost_multiplier
|
||||
created_at,
|
||||
"gemini_session", // data_source
|
||||
],
|
||||
)
|
||||
.map_err(|e| AppError::Database(format!("插入 Gemini 会话日志失败: {e}")))?;
|
||||
|
||||
// changes() > 0 表示新插入或已更新，== 0 表示值完全相同（无实际变更）
|
||||
Ok(conn.changes() > 0)
|
||||
}
|
||||
|
||||
/// 获取文件的同步状态
|
||||
fn get_sync_state(db: &Database, file_path: &str) -> Result<(i64, i64), AppError> {
|
||||
let conn = lock_conn!(db.conn);
|
||||
let result = conn.query_row(
|
||||
"SELECT last_modified, last_line_offset FROM session_log_sync WHERE file_path = ?1",
|
||||
rusqlite::params![file_path],
|
||||
|row| Ok((row.get::<_, i64>(0)?, row.get::<_, i64>(1)?)),
|
||||
);
|
||||
Ok(result.unwrap_or((0, 0)))
|
||||
}
|
||||
|
||||
/// 更新文件的同步状态
|
||||
fn update_sync_state(
|
||||
db: &Database,
|
||||
file_path: &str,
|
||||
last_modified: i64,
|
||||
last_offset: i64,
|
||||
) -> Result<(), AppError> {
|
||||
let now = SystemTime::now()
|
||||
.duration_since(SystemTime::UNIX_EPOCH)
|
||||
.map(|d| d.as_secs() as i64)
|
||||
.unwrap_or(0);
|
||||
|
||||
let conn = lock_conn!(db.conn);
|
||||
conn.execute(
|
||||
"INSERT OR REPLACE INTO session_log_sync (file_path, last_modified, last_line_offset, last_synced_at)
|
||||
VALUES (?1, ?2, ?3, ?4)",
|
||||
rusqlite::params![file_path, last_modified, last_offset, now],
|
||||
)
|
||||
.map_err(|e| AppError::Database(format!("更新同步状态失败: {e}")))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 查找 Gemini 模型定价
|
||||
fn find_gemini_pricing(conn: &rusqlite::Connection, model_id: &str) -> Option<ModelPricing> {
|
||||
// 精确匹配
|
||||
if let Some(pricing) = try_find_pricing(conn, model_id) {
|
||||
return Some(pricing);
|
||||
}
|
||||
|
||||
// LIKE 模糊匹配（兜底）
|
||||
let pattern = format!("{model_id}%");
|
||||
conn.query_row(
|
||||
"SELECT input_cost_per_million, output_cost_per_million,
|
||||
cache_read_cost_per_million, cache_creation_cost_per_million
|
||||
FROM model_pricing WHERE model_id LIKE ?1 LIMIT 1",
|
||||
rusqlite::params![pattern],
|
||||
|row| {
|
||||
Ok((
|
||||
row.get::<_, String>(0)?,
|
||||
row.get::<_, String>(1)?,
|
||||
row.get::<_, String>(2)?,
|
||||
row.get::<_, String>(3)?,
|
||||
))
|
||||
},
|
||||
)
|
||||
.ok()
|
||||
.and_then(|(i, o, cr, cc)| ModelPricing::from_strings(&i, &o, &cr, &cc).ok())
|
||||
}
|
||||
|
||||
/// 精确匹配定价查询
|
||||
fn try_find_pricing(conn: &rusqlite::Connection, model_id: &str) -> Option<ModelPricing> {
|
||||
conn.query_row(
|
||||
"SELECT input_cost_per_million, output_cost_per_million,
|
||||
cache_read_cost_per_million, cache_creation_cost_per_million
|
||||
FROM model_pricing WHERE model_id = ?1",
|
||||
rusqlite::params![model_id],
|
||||
|row| {
|
||||
Ok((
|
||||
row.get::<_, String>(0)?,
|
||||
row.get::<_, String>(1)?,
|
||||
row.get::<_, String>(2)?,
|
||||
row.get::<_, String>(3)?,
|
||||
))
|
||||
},
|
||||
)
|
||||
.ok()
|
||||
.and_then(|(i, o, cr, cc)| ModelPricing::from_strings(&i, &o, &cr, &cc).ok())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_collect_gemini_session_files_nonexistent() {
|
||||
let files = collect_gemini_session_files(Path::new("/nonexistent/path"));
|
||||
assert!(files.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_gemini_tokens() {
|
||||
let json: serde_json::Value = serde_json::json!({
|
||||
"input": 8522,
|
||||
"output": 29,
|
||||
"cached": 3138,
|
||||
"thoughts": 405,
|
||||
"tool": 0,
|
||||
"total": 8956
|
||||
});
|
||||
let tokens = parse_gemini_tokens(&json);
|
||||
assert_eq!(tokens.input, 8522);
|
||||
assert_eq!(tokens.output, 29);
|
||||
assert_eq!(tokens.cached, 3138);
|
||||
assert_eq!(tokens.thoughts, 405);
|
||||
// output + thoughts = 29 + 405 = 434（用于计费）
|
||||
assert_eq!(tokens.output + tokens.thoughts, 434);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_gemini_tokens_missing_fields() {
|
||||
// 缺少某些字段时应返回 0
|
||||
let json: serde_json::Value = serde_json::json!({
|
||||
"input": 100,
|
||||
"output": 50
|
||||
});
|
||||
let tokens = parse_gemini_tokens(&json);
|
||||
assert_eq!(tokens.input, 100);
|
||||
assert_eq!(tokens.output, 50);
|
||||
assert_eq!(tokens.cached, 0);
|
||||
assert_eq!(tokens.thoughts, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_gemini_tokens_all_zero() {
|
||||
let json: serde_json::Value = serde_json::json!({
|
||||
"input": 0,
|
||||
"output": 0,
|
||||
"cached": 0,
|
||||
"thoughts": 0,
|
||||
"tool": 0,
|
||||
"total": 0
|
||||
});
|
||||
let tokens = parse_gemini_tokens(&json);
|
||||
assert_eq!(tokens.input, 0);
|
||||
assert_eq!(tokens.output, 0);
|
||||
// 全零（包括 cached=0）会被 sync 逻辑跳过
|
||||
assert!(
|
||||
tokens.input == 0 && tokens.output == 0 && tokens.thoughts == 0 && tokens.cached == 0
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_gemini_tokens_cache_only_not_skipped() {
|
||||
// 纯缓存命中消息（input/output/thoughts=0 但 cached>0）不应被跳过
|
||||
let json: serde_json::Value = serde_json::json!({
|
||||
"input": 0,
|
||||
"output": 0,
|
||||
"cached": 5000,
|
||||
"thoughts": 0
|
||||
});
|
||||
let tokens = parse_gemini_tokens(&json);
|
||||
assert_eq!(tokens.cached, 5000);
|
||||
// 跳过条件：所有四个字段都为 0 才跳过
|
||||
let should_skip =
|
||||
tokens.input == 0 && tokens.output == 0 && tokens.thoughts == 0 && tokens.cached == 0;
|
||||
assert!(!should_skip, "纯缓存命中记录不应被跳过");
|
||||
}
|
||||
}
|
||||
@@ -34,6 +34,17 @@ pub enum SyncMethod {
|
||||
Copy,
|
||||
}
|
||||
|
||||
/// Skill 存储位置（SSOT 目录选择）
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum SkillStorageLocation {
|
||||
/// CC Switch 管理目录 (~/.cc-switch/skills/)
|
||||
#[default]
|
||||
CcSwitch,
|
||||
/// Agent Skills 统一标准目录 (~/.agents/skills/)
|
||||
Unified,
|
||||
}
|
||||
|
||||
/// 可发现的技能（来自仓库）
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct DiscoverableSkill {
|
||||
@@ -160,6 +171,81 @@ pub struct SkillUninstallResult {
|
||||
pub backup_path: Option<String>,
|
||||
}
|
||||
|
||||
/// Skill 更新检测结果
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct SkillUpdateInfo {
|
||||
/// Skill ID
|
||||
pub id: String,
|
||||
/// Skill 名称
|
||||
pub name: String,
|
||||
/// 当前本地哈希
|
||||
pub current_hash: Option<String>,
|
||||
/// 远程最新哈希
|
||||
pub remote_hash: String,
|
||||
}
|
||||
|
||||
/// Skill 存储位置迁移结果
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct MigrationResult {
|
||||
pub migrated_count: usize,
|
||||
pub skipped_count: usize,
|
||||
pub errors: Vec<String>,
|
||||
}
|
||||
|
||||
// ========== skills.sh API 类型 ==========
|
||||
|
||||
/// skills.sh API 原始响应
|
||||
///
|
||||
/// 注意：API 命名不一致（searchType 是 camelCase，duration_ms 是 snake_case），
|
||||
/// 因此不能用 rename_all，需要逐字段指定。
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
struct SkillsShApiResponse {
|
||||
pub query: String,
|
||||
#[serde(rename = "searchType")]
|
||||
#[allow(dead_code)]
|
||||
pub search_type: String,
|
||||
pub skills: Vec<SkillsShApiSkill>,
|
||||
pub count: usize,
|
||||
#[allow(dead_code)]
|
||||
pub duration_ms: u64,
|
||||
}
|
||||
|
||||
/// skills.sh API 原始技能条目
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
struct SkillsShApiSkill {
|
||||
pub id: String,
|
||||
#[serde(rename = "skillId")]
|
||||
pub skill_id: String,
|
||||
pub name: String,
|
||||
pub installs: u64,
|
||||
pub source: String,
|
||||
}
|
||||
|
||||
/// skills.sh 搜索结果（返回给前端）
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct SkillsShSearchResult {
|
||||
pub skills: Vec<SkillsShDiscoverableSkill>,
|
||||
pub total_count: usize,
|
||||
pub query: String,
|
||||
}
|
||||
|
||||
/// skills.sh 可安装技能（返回给前端）
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct SkillsShDiscoverableSkill {
|
||||
pub key: String,
|
||||
pub name: String,
|
||||
pub directory: String,
|
||||
pub repo_owner: String,
|
||||
pub repo_name: String,
|
||||
pub repo_branch: String,
|
||||
pub installs: u64,
|
||||
pub readme_url: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct SkillBackupEntry {
|
||||
@@ -389,9 +475,20 @@ impl SkillService {
|
||||
|
||||
// ========== 路径管理 ==========
|
||||
|
||||
/// 获取 SSOT 目录（~/.cc-switch/skills/）
|
||||
/// 获取 SSOT 目录（根据设置返回 ~/.cc-switch/skills/ 或 ~/.agents/skills/）
|
||||
pub fn get_ssot_dir() -> Result<PathBuf> {
|
||||
let dir = get_app_config_dir().join("skills");
|
||||
let location = crate::settings::get_skill_storage_location();
|
||||
let dir = match location {
|
||||
SkillStorageLocation::CcSwitch => get_app_config_dir().join("skills"),
|
||||
SkillStorageLocation::Unified => {
|
||||
let home = dirs::home_dir().context(format_skill_error(
|
||||
"GET_HOME_DIR_FAILED",
|
||||
&[],
|
||||
Some("checkPermission"),
|
||||
))?;
|
||||
home.join(".agents").join("skills")
|
||||
}
|
||||
};
|
||||
fs::create_dir_all(&dir)?;
|
||||
Ok(dir)
|
||||
}
|
||||
@@ -569,14 +666,28 @@ impl SkillService {
|
||||
repo_branch = used_branch;
|
||||
|
||||
// 复制到 SSOT
|
||||
let source = temp_dir.join(&source_rel);
|
||||
let mut source = temp_dir.join(&source_rel);
|
||||
if !source.exists() {
|
||||
let _ = fs::remove_dir_all(&temp_dir);
|
||||
return Err(anyhow!(format_skill_error(
|
||||
"SKILL_DIR_NOT_FOUND",
|
||||
&[("path", &source.display().to_string())],
|
||||
Some("checkRepoUrl"),
|
||||
)));
|
||||
// 回退：在 temp_dir 中递归查找名称匹配的目录（含 SKILL.md）
|
||||
let target_name = source_rel
|
||||
.file_name()
|
||||
.map(|n| n.to_string_lossy().to_string())
|
||||
.unwrap_or_default();
|
||||
if let Some(found) = Self::find_skill_dir_by_name(&temp_dir, &target_name) {
|
||||
log::info!(
|
||||
"Skill directory '{}' not found at direct path, using fallback: {}",
|
||||
target_name,
|
||||
found.display()
|
||||
);
|
||||
source = found;
|
||||
} else {
|
||||
let _ = fs::remove_dir_all(&temp_dir);
|
||||
return Err(anyhow!(format_skill_error(
|
||||
"SKILL_DIR_NOT_FOUND",
|
||||
&[("path", &source.display().to_string())],
|
||||
Some("checkRepoUrl"),
|
||||
)));
|
||||
}
|
||||
}
|
||||
|
||||
let canonical_temp = temp_dir.canonicalize().unwrap_or_else(|_| temp_dir.clone());
|
||||
@@ -632,6 +743,12 @@ impl SkillService {
|
||||
));
|
||||
|
||||
// 创建 InstalledSkill 记录
|
||||
// 计算内容哈希
|
||||
let content_hash = Self::compute_dir_hash(&dest).map(Some).unwrap_or_else(|e| {
|
||||
log::warn!("Failed to compute content hash for {}: {e}", install_name);
|
||||
None
|
||||
});
|
||||
|
||||
let installed_skill = InstalledSkill {
|
||||
id: skill.key.clone(),
|
||||
name: skill.name.clone(),
|
||||
@@ -647,6 +764,8 @@ impl SkillService {
|
||||
readme_url,
|
||||
apps: SkillApps::only(current_app),
|
||||
installed_at: chrono::Utc::now().timestamp(),
|
||||
content_hash,
|
||||
updated_at: 0,
|
||||
};
|
||||
|
||||
// 保存到数据库
|
||||
@@ -706,6 +825,412 @@ impl SkillService {
|
||||
Ok(SkillUninstallResult { backup_path })
|
||||
}
|
||||
|
||||
// ========== 更新检测 ==========
|
||||
|
||||
/// 计算目录内容的 SHA-256 哈希
|
||||
///
|
||||
/// 递归遍历目录下所有非隐藏文件，按相对路径字典序排列，
|
||||
/// 将 "相对路径\0内容\0" 逐文件 feed 给同一个 hasher。
|
||||
pub fn compute_dir_hash(dir: &Path) -> Result<String> {
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
let mut files: Vec<PathBuf> = Vec::new();
|
||||
Self::collect_files_for_hash(dir, dir, &mut files)?;
|
||||
files.sort();
|
||||
|
||||
let mut hasher = Sha256::new();
|
||||
for file_path in &files {
|
||||
let relative = file_path.strip_prefix(dir).unwrap_or(file_path);
|
||||
let rel_str = relative.to_string_lossy().replace('\\', "/");
|
||||
hasher.update(rel_str.as_bytes());
|
||||
hasher.update(b"\0");
|
||||
let content = fs::read(file_path)
|
||||
.with_context(|| format!("读取文件失败: {}", file_path.display()))?;
|
||||
hasher.update(&content);
|
||||
hasher.update(b"\0");
|
||||
}
|
||||
|
||||
Ok(format!("{:x}", hasher.finalize()))
|
||||
}
|
||||
|
||||
/// 递归收集目录下所有非隐藏文件
|
||||
#[allow(clippy::only_used_in_recursion)]
|
||||
fn collect_files_for_hash(base: &Path, current: &Path, files: &mut Vec<PathBuf>) -> Result<()> {
|
||||
let entries = fs::read_dir(current)
|
||||
.with_context(|| format!("读取目录失败: {}", current.display()))?;
|
||||
for entry in entries {
|
||||
let entry = entry?;
|
||||
let name = entry.file_name().to_string_lossy().to_string();
|
||||
if name.starts_with('.') {
|
||||
continue;
|
||||
}
|
||||
let path = entry.path();
|
||||
if path.is_dir() {
|
||||
Self::collect_files_for_hash(base, &path, files)?;
|
||||
} else {
|
||||
files.push(path);
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 检查所有已安装 Skill 的更新
|
||||
///
|
||||
/// 仅检查有 repo_owner 的 Skill（本地 Skill 跳过），
|
||||
/// 按仓库分组下载，避免重复下载同一仓库。
|
||||
pub async fn check_updates(&self, db: &Arc<Database>) -> Result<Vec<SkillUpdateInfo>> {
|
||||
let skills = db.get_all_installed_skills()?;
|
||||
let mut updates = Vec::new();
|
||||
|
||||
// 按 (owner, name, branch) 分组
|
||||
let mut repo_groups: HashMap<(String, String, String), Vec<InstalledSkill>> =
|
||||
HashMap::new();
|
||||
|
||||
for skill in skills.into_values() {
|
||||
let (owner, name, branch) =
|
||||
match (&skill.repo_owner, &skill.repo_name, &skill.repo_branch) {
|
||||
(Some(o), Some(n), Some(b)) => (o.clone(), n.clone(), b.clone()),
|
||||
(Some(o), Some(n), None) => (o.clone(), n.clone(), "main".to_string()),
|
||||
_ => continue,
|
||||
};
|
||||
repo_groups
|
||||
.entry((owner, name, branch))
|
||||
.or_default()
|
||||
.push(skill);
|
||||
}
|
||||
|
||||
let ssot_dir = Self::get_ssot_dir()?;
|
||||
|
||||
for ((owner, name, branch), group_skills) in &repo_groups {
|
||||
let repo = SkillRepo {
|
||||
owner: owner.clone(),
|
||||
name: name.clone(),
|
||||
branch: branch.clone(),
|
||||
enabled: true,
|
||||
};
|
||||
|
||||
// 下载仓库 ZIP
|
||||
let (temp_dir, _used_branch) = match timeout(
|
||||
std::time::Duration::from_secs(60),
|
||||
self.download_repo(&repo),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Ok(result)) => result,
|
||||
Ok(Err(e)) => {
|
||||
log::warn!("检查更新时下载 {}/{} 失败: {e}", owner, name);
|
||||
continue;
|
||||
}
|
||||
Err(_) => {
|
||||
log::warn!("检查更新时下载 {}/{} 超时", owner, name);
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
// 扫描仓库中的所有 Skill 目录
|
||||
let mut remote_skills: Vec<DiscoverableSkill> = Vec::new();
|
||||
let _ = self.scan_dir_recursive(&temp_dir, &temp_dir, &repo, &mut remote_skills);
|
||||
|
||||
for skill in group_skills {
|
||||
// 在远程仓库中找到匹配的 Skill 目录
|
||||
let remote_match = remote_skills.iter().find(|rs| {
|
||||
// 匹配方式：安装名称的最后一段
|
||||
let remote_install_name =
|
||||
rs.directory.rsplit('/').next().unwrap_or(&rs.directory);
|
||||
remote_install_name.eq_ignore_ascii_case(&skill.directory)
|
||||
});
|
||||
|
||||
let remote_skill_dir = match remote_match {
|
||||
Some(rs) => temp_dir.join(&rs.directory),
|
||||
None => continue,
|
||||
};
|
||||
|
||||
if !remote_skill_dir.exists() {
|
||||
continue;
|
||||
}
|
||||
|
||||
let remote_hash = match Self::compute_dir_hash(&remote_skill_dir) {
|
||||
Ok(h) => h,
|
||||
Err(e) => {
|
||||
log::warn!("计算远程哈希失败 {}: {e}", skill.id);
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
// 本地哈希：优先数据库，否则实时计算
|
||||
let local_hash = match &skill.content_hash {
|
||||
Some(h) => Some(h.clone()),
|
||||
None => {
|
||||
let local_dir = ssot_dir.join(&skill.directory);
|
||||
if local_dir.exists() {
|
||||
match Self::compute_dir_hash(&local_dir) {
|
||||
Ok(h) => {
|
||||
let _ = db.update_skill_hash(&skill.id, &h, 0);
|
||||
Some(h)
|
||||
}
|
||||
Err(_) => None,
|
||||
}
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
if local_hash.as_deref() != Some(&remote_hash) {
|
||||
updates.push(SkillUpdateInfo {
|
||||
id: skill.id.clone(),
|
||||
name: skill.name.clone(),
|
||||
current_hash: local_hash,
|
||||
remote_hash,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
let _ = fs::remove_dir_all(&temp_dir);
|
||||
}
|
||||
|
||||
Ok(updates)
|
||||
}
|
||||
|
||||
/// 更新单个 Skill（重新下载并替换本地文件）
|
||||
pub async fn update_skill(&self, db: &Arc<Database>, skill_id: &str) -> Result<InstalledSkill> {
|
||||
let skill = db
|
||||
.get_installed_skill(skill_id)?
|
||||
.ok_or_else(|| anyhow!("Skill not found: {skill_id}"))?;
|
||||
|
||||
let (owner, name, branch) = match (&skill.repo_owner, &skill.repo_name) {
|
||||
(Some(o), Some(n)) => (
|
||||
o.clone(),
|
||||
n.clone(),
|
||||
skill
|
||||
.repo_branch
|
||||
.clone()
|
||||
.unwrap_or_else(|| "main".to_string()),
|
||||
),
|
||||
_ => return Err(anyhow!("Cannot update local skill: {skill_id}")),
|
||||
};
|
||||
|
||||
let repo = SkillRepo {
|
||||
owner: owner.clone(),
|
||||
name: name.clone(),
|
||||
branch: branch.clone(),
|
||||
enabled: true,
|
||||
};
|
||||
|
||||
let ssot_dir = Self::get_ssot_dir()?;
|
||||
|
||||
// 下载仓库
|
||||
let (temp_dir, used_branch) = timeout(
|
||||
std::time::Duration::from_secs(60),
|
||||
self.download_repo(&repo),
|
||||
)
|
||||
.await
|
||||
.map_err(|_| {
|
||||
anyhow!(format_skill_error(
|
||||
"DOWNLOAD_TIMEOUT",
|
||||
&[("owner", &owner), ("name", &name), ("timeout", "60")],
|
||||
Some("checkNetwork"),
|
||||
))
|
||||
})??;
|
||||
|
||||
// 在解压的仓库中查找 Skill 源目录
|
||||
let mut remote_skills: Vec<DiscoverableSkill> = Vec::new();
|
||||
let _ = self.scan_dir_recursive(&temp_dir, &temp_dir, &repo, &mut remote_skills);
|
||||
|
||||
let remote_match = remote_skills
|
||||
.iter()
|
||||
.find(|rs| {
|
||||
let remote_install_name = rs.directory.rsplit('/').next().unwrap_or(&rs.directory);
|
||||
remote_install_name.eq_ignore_ascii_case(&skill.directory)
|
||||
})
|
||||
.ok_or_else(|| {
|
||||
let _ = fs::remove_dir_all(&temp_dir);
|
||||
anyhow!(format_skill_error(
|
||||
"SKILL_DIR_NOT_FOUND",
|
||||
&[("path", &skill.directory)],
|
||||
Some("checkRepoUrl"),
|
||||
))
|
||||
})?;
|
||||
|
||||
let source = temp_dir.join(&remote_match.directory);
|
||||
if !source.exists() {
|
||||
let _ = fs::remove_dir_all(&temp_dir);
|
||||
return Err(anyhow!(format_skill_error(
|
||||
"SKILL_DIR_NOT_FOUND",
|
||||
&[("path", &source.display().to_string())],
|
||||
Some("checkRepoUrl"),
|
||||
)));
|
||||
}
|
||||
|
||||
// 备份旧文件
|
||||
let _ = Self::create_uninstall_backup(&skill);
|
||||
|
||||
// 删除旧 SSOT 目录并复制新文件
|
||||
let dest = ssot_dir.join(&skill.directory);
|
||||
if dest.exists() {
|
||||
fs::remove_dir_all(&dest)?;
|
||||
}
|
||||
Self::copy_dir_recursive(&source, &dest)?;
|
||||
let _ = fs::remove_dir_all(&temp_dir);
|
||||
|
||||
// 计算新哈希 + 解析新元数据
|
||||
let new_hash = Self::compute_dir_hash(&dest).ok();
|
||||
let skill_md = dest.join("SKILL.md");
|
||||
let (new_name, new_description) = Self::read_skill_name_desc(&skill_md, &skill.directory);
|
||||
|
||||
// 更新 readme_url
|
||||
let doc_path = skill
|
||||
.readme_url
|
||||
.as_deref()
|
||||
.and_then(Self::extract_doc_path_from_url)
|
||||
.unwrap_or_else(|| format!("{}/SKILL.md", skill.directory.trim_end_matches('/')));
|
||||
let readme_url = Some(Self::build_skill_doc_url(
|
||||
&owner,
|
||||
&name,
|
||||
&used_branch,
|
||||
&doc_path,
|
||||
));
|
||||
|
||||
let updated_skill = InstalledSkill {
|
||||
id: skill.id.clone(),
|
||||
name: new_name,
|
||||
description: new_description,
|
||||
directory: skill.directory.clone(),
|
||||
repo_owner: skill.repo_owner.clone(),
|
||||
repo_name: skill.repo_name.clone(),
|
||||
repo_branch: Some(used_branch),
|
||||
readme_url,
|
||||
apps: skill.apps.clone(),
|
||||
installed_at: skill.installed_at,
|
||||
content_hash: new_hash,
|
||||
updated_at: chrono::Utc::now().timestamp(),
|
||||
};
|
||||
|
||||
db.save_skill(&updated_skill)?;
|
||||
|
||||
// 同步到所有已启用的应用目录
|
||||
for app in updated_skill.apps.enabled_apps() {
|
||||
if let Err(e) = Self::sync_to_app_dir(&updated_skill.directory, &app) {
|
||||
log::warn!("同步更新后的 skill 到 {:?} 失败: {e}", app);
|
||||
}
|
||||
}
|
||||
|
||||
log::info!("Skill {} 更新成功", updated_skill.name);
|
||||
Ok(updated_skill)
|
||||
}
|
||||
|
||||
/// 为缺少 content_hash 的已安装 Skill 补算哈希
|
||||
pub fn backfill_content_hashes(db: &Arc<Database>) -> Result<usize> {
|
||||
let skills = db.get_all_installed_skills()?;
|
||||
let ssot_dir = Self::get_ssot_dir()?;
|
||||
let mut count = 0;
|
||||
|
||||
for skill in skills.values() {
|
||||
if skill.content_hash.is_some() {
|
||||
continue;
|
||||
}
|
||||
let skill_dir = ssot_dir.join(&skill.directory);
|
||||
if !skill_dir.exists() {
|
||||
continue;
|
||||
}
|
||||
match Self::compute_dir_hash(&skill_dir) {
|
||||
Ok(hash) => {
|
||||
let _ = db.update_skill_hash(&skill.id, &hash, 0);
|
||||
count += 1;
|
||||
}
|
||||
Err(e) => {
|
||||
log::warn!("补算哈希失败 {}: {e}", skill.id);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if count > 0 {
|
||||
log::info!("已为 {count} 个 Skill 补算内容哈希");
|
||||
}
|
||||
Ok(count)
|
||||
}
|
||||
|
||||
/// 迁移 Skill 存储位置（在两个 SSOT 目录间移动文件）
|
||||
///
|
||||
/// 安全策略：先移文件，后改设置。中途崩溃时设置仍指向旧目录。
|
||||
pub fn migrate_storage(
|
||||
db: &Arc<Database>,
|
||||
target: SkillStorageLocation,
|
||||
) -> Result<MigrationResult> {
|
||||
let current = crate::settings::get_skill_storage_location();
|
||||
if current == target {
|
||||
return Ok(MigrationResult {
|
||||
migrated_count: 0,
|
||||
skipped_count: 0,
|
||||
errors: vec![],
|
||||
});
|
||||
}
|
||||
|
||||
// 1. 解析旧目录和新目录（不改设置）
|
||||
let old_dir = Self::get_ssot_dir()?;
|
||||
let new_dir = match target {
|
||||
SkillStorageLocation::CcSwitch => get_app_config_dir().join("skills"),
|
||||
SkillStorageLocation::Unified => {
|
||||
let home = dirs::home_dir().context("Cannot determine home directory")?;
|
||||
home.join(".agents").join("skills")
|
||||
}
|
||||
};
|
||||
fs::create_dir_all(&new_dir)?;
|
||||
|
||||
// 2. 逐个移动 skill 目录
|
||||
let skills = db.get_all_installed_skills()?;
|
||||
let mut result = MigrationResult {
|
||||
migrated_count: 0,
|
||||
skipped_count: 0,
|
||||
errors: vec![],
|
||||
};
|
||||
|
||||
for skill in skills.values() {
|
||||
let src = old_dir.join(&skill.directory);
|
||||
let dst = new_dir.join(&skill.directory);
|
||||
|
||||
if !src.exists() {
|
||||
result.skipped_count += 1;
|
||||
continue;
|
||||
}
|
||||
if dst.exists() {
|
||||
result.skipped_count += 1;
|
||||
continue;
|
||||
}
|
||||
|
||||
// 优先 rename（同文件系统原子操作），失败则 copy+delete
|
||||
match fs::rename(&src, &dst) {
|
||||
Ok(()) => result.migrated_count += 1,
|
||||
Err(_) => match Self::copy_dir_recursive(&src, &dst) {
|
||||
Ok(()) => {
|
||||
let _ = fs::remove_dir_all(&src);
|
||||
result.migrated_count += 1;
|
||||
}
|
||||
Err(e) => {
|
||||
result.errors.push(format!("{}: {e}", skill.directory));
|
||||
}
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// 3. 文件移动完成后才持久化设置
|
||||
crate::settings::set_skill_storage_location(target)?;
|
||||
|
||||
// 4. 刷新所有应用目录的 symlink（指向新 SSOT）
|
||||
for app in AppType::all() {
|
||||
let _ = Self::sync_to_app(db, &app);
|
||||
}
|
||||
|
||||
log::info!(
|
||||
"Skill 存储迁移完成: {} 迁移, {} 跳过, {} 错误",
|
||||
result.migrated_count,
|
||||
result.skipped_count,
|
||||
result.errors.len()
|
||||
);
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
pub fn list_backups() -> Result<Vec<SkillBackupEntry>> {
|
||||
let backup_dir = Self::get_backup_dir()?;
|
||||
let mut entries = Vec::new();
|
||||
@@ -800,9 +1325,13 @@ impl SkillService {
|
||||
let mut restored_skill = metadata.skill;
|
||||
restored_skill.installed_at = Utc::now().timestamp();
|
||||
restored_skill.apps = SkillApps::only(current_app);
|
||||
restored_skill.updated_at = 0;
|
||||
|
||||
Self::copy_dir_recursive(&backup_skill_dir, &restore_path)?;
|
||||
|
||||
// 重新计算内容哈希
|
||||
restored_skill.content_hash = Self::compute_dir_hash(&restore_path).ok();
|
||||
|
||||
if let Err(err) = db.save_skill(&restored_skill) {
|
||||
let _ = fs::remove_dir_all(&restore_path);
|
||||
return Err(err.into());
|
||||
@@ -991,6 +1520,10 @@ impl SkillService {
|
||||
let (id, repo_owner, repo_name, repo_branch, readme_url) =
|
||||
build_repo_info_from_lock(&agents_lock, &dir_name);
|
||||
|
||||
// 计算内容哈希
|
||||
let ssot_skill_dir = ssot_dir.join(&dir_name);
|
||||
let content_hash = Self::compute_dir_hash(&ssot_skill_dir).ok();
|
||||
|
||||
// 创建记录
|
||||
let skill = InstalledSkill {
|
||||
id,
|
||||
@@ -1003,6 +1536,8 @@ impl SkillService {
|
||||
readme_url,
|
||||
apps,
|
||||
installed_at: chrono::Utc::now().timestamp(),
|
||||
content_hash,
|
||||
updated_at: 0,
|
||||
};
|
||||
|
||||
// 保存到数据库
|
||||
@@ -1514,6 +2049,38 @@ impl SkillService {
|
||||
}
|
||||
}
|
||||
|
||||
/// 在目录树中查找名称匹配且包含 SKILL.md 的子目录
|
||||
///
|
||||
/// 用于 skills.sh 安装回退：API 只返回 skillId（如 "find-skills"），
|
||||
/// 但实际文件可能在仓库子目录中（如 "skills/find-skills"）。
|
||||
fn find_skill_dir_by_name(root: &Path, target_name: &str) -> Option<PathBuf> {
|
||||
fn walk(dir: &Path, target: &str, depth: usize) -> Option<PathBuf> {
|
||||
if depth > 3 {
|
||||
return None;
|
||||
}
|
||||
let entries = fs::read_dir(dir).ok()?;
|
||||
for entry in entries.flatten() {
|
||||
let path = entry.path();
|
||||
if !path.is_dir() {
|
||||
continue;
|
||||
}
|
||||
let name = entry.file_name();
|
||||
let name_str = name.to_string_lossy();
|
||||
if name_str.starts_with('.') {
|
||||
continue;
|
||||
}
|
||||
if name_str.eq_ignore_ascii_case(target) && path.join("SKILL.md").exists() {
|
||||
return Some(path);
|
||||
}
|
||||
if let Some(found) = walk(&path, target, depth + 1) {
|
||||
return Some(found);
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
walk(root, target_name, 0)
|
||||
}
|
||||
|
||||
/// 去重技能列表（基于完整 key，不同仓库的同名 skill 分开显示）
|
||||
fn deduplicate_discoverable_skills(skills: &mut Vec<DiscoverableSkill>) {
|
||||
let mut seen = HashMap::new();
|
||||
@@ -1965,6 +2532,9 @@ impl SkillService {
|
||||
}
|
||||
Self::copy_dir_recursive(&skill_dir, &dest)?;
|
||||
|
||||
// 计算内容哈希
|
||||
let content_hash = Self::compute_dir_hash(&dest).ok();
|
||||
|
||||
// 创建 InstalledSkill 记录
|
||||
let skill = InstalledSkill {
|
||||
id: format!("local:{install_name}"),
|
||||
@@ -1977,6 +2547,8 @@ impl SkillService {
|
||||
readme_url: None,
|
||||
apps: SkillApps::only(current_app),
|
||||
installed_at: chrono::Utc::now().timestamp(),
|
||||
content_hash,
|
||||
updated_at: 0,
|
||||
};
|
||||
|
||||
// 保存到数据库
|
||||
@@ -2116,6 +2688,67 @@ impl SkillService {
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// ========== skills.sh 搜索 ==========
|
||||
|
||||
/// 搜索 skills.sh 公共目录
|
||||
pub async fn search_skills_sh(
|
||||
query: &str,
|
||||
limit: usize,
|
||||
offset: usize,
|
||||
) -> Result<SkillsShSearchResult> {
|
||||
let client = crate::proxy::http_client::get();
|
||||
|
||||
let url = url::Url::parse_with_params(
|
||||
"https://skills.sh/api/search",
|
||||
&[
|
||||
("q", query),
|
||||
("limit", &limit.to_string()),
|
||||
("offset", &offset.to_string()),
|
||||
],
|
||||
)?;
|
||||
|
||||
let resp = client
|
||||
.get(url)
|
||||
.timeout(std::time::Duration::from_secs(10))
|
||||
.send()
|
||||
.await?
|
||||
.error_for_status()?
|
||||
.json::<SkillsShApiResponse>()
|
||||
.await?;
|
||||
|
||||
let skills = resp
|
||||
.skills
|
||||
.into_iter()
|
||||
.filter_map(|s| {
|
||||
let parts: Vec<&str> = s.source.splitn(2, '/').collect();
|
||||
if parts.len() != 2 {
|
||||
return None;
|
||||
}
|
||||
let (owner, repo) = (parts[0].to_string(), parts[1].to_string());
|
||||
// 过滤非 GitHub 来源（如 "skills.volces.com"、"mcp-hub.momenta.works"）
|
||||
if owner.contains('.') || repo.contains('.') {
|
||||
return None;
|
||||
}
|
||||
Some(SkillsShDiscoverableSkill {
|
||||
key: s.id,
|
||||
name: s.name,
|
||||
directory: s.skill_id.clone(),
|
||||
repo_owner: owner.clone(),
|
||||
repo_name: repo.clone(),
|
||||
repo_branch: "main".to_string(),
|
||||
installs: s.installs,
|
||||
readme_url: Some(format!("https://github.com/{}/{}", owner, repo)),
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
|
||||
Ok(SkillsShSearchResult {
|
||||
skills,
|
||||
total_count: resp.count,
|
||||
query: resp.query,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// ========== 迁移支持 ==========
|
||||
@@ -2288,6 +2921,8 @@ pub fn migrate_skills_to_ssot(db: &Arc<Database>) -> Result<usize> {
|
||||
let (id, repo_owner, repo_name, repo_branch, readme_url) =
|
||||
build_repo_info_from_lock(&agents_lock, &directory);
|
||||
|
||||
let content_hash = SkillService::compute_dir_hash(&ssot_path).ok();
|
||||
|
||||
let skill = InstalledSkill {
|
||||
id,
|
||||
name,
|
||||
@@ -2299,6 +2934,8 @@ pub fn migrate_skills_to_ssot(db: &Arc<Database>) -> Result<usize> {
|
||||
readme_url,
|
||||
apps,
|
||||
installed_at: chrono::Utc::now().timestamp(),
|
||||
content_hash,
|
||||
updated_at: 0,
|
||||
};
|
||||
|
||||
db.save_skill(&skill)?;
|
||||
|
||||
@@ -199,6 +199,14 @@ impl StreamCheckService {
|
||||
claude_api_format_override: Option<String>,
|
||||
) -> Result<StreamCheckResult, AppError> {
|
||||
let start = Instant::now();
|
||||
|
||||
// OpenCode / OpenClaw 的 settings_config 结构与 Claude/Codex/Gemini 不同
|
||||
// （baseUrl / apiKey 直接作为根字段而非嵌套在 env），并且协议由 `api`
|
||||
// 或 `npm` 字段显式指定。它们不走 get_adapter 路径，而是直接分发。
|
||||
if matches!(app_type, AppType::OpenCode | AppType::OpenClaw) {
|
||||
return Self::check_once_without_adapter(app_type, provider, config, start).await;
|
||||
}
|
||||
|
||||
let adapter = get_adapter(app_type);
|
||||
|
||||
let base_url = match base_url_override {
|
||||
@@ -231,6 +239,7 @@ impl StreamCheckService {
|
||||
request_timeout,
|
||||
provider,
|
||||
claude_api_format_override.as_deref(),
|
||||
None,
|
||||
)
|
||||
.await
|
||||
}
|
||||
@@ -254,24 +263,13 @@ impl StreamCheckService {
|
||||
&model_to_test,
|
||||
test_prompt,
|
||||
request_timeout,
|
||||
None,
|
||||
)
|
||||
.await
|
||||
}
|
||||
AppType::OpenCode => {
|
||||
// OpenCode doesn't support stream check yet
|
||||
return Err(AppError::localized(
|
||||
"opencode_no_stream_check",
|
||||
"OpenCode 暂不支持健康检查",
|
||||
"OpenCode does not support health check yet",
|
||||
));
|
||||
}
|
||||
AppType::OpenClaw => {
|
||||
// OpenClaw doesn't support stream check yet
|
||||
return Err(AppError::localized(
|
||||
"openclaw_no_stream_check",
|
||||
"OpenClaw 暂不支持健康检查",
|
||||
"OpenClaw does not support health check yet",
|
||||
));
|
||||
AppType::OpenCode | AppType::OpenClaw => {
|
||||
// Already handled via early dispatch above
|
||||
unreachable!("OpenCode/OpenClaw 已通过 check_once_without_adapter 处理")
|
||||
}
|
||||
};
|
||||
|
||||
@@ -313,6 +311,10 @@ impl StreamCheckService {
|
||||
/// - "openai_chat": OpenAI Chat Completions API (/v1/chat/completions)
|
||||
/// - "openai_responses": OpenAI Responses API (/v1/responses)
|
||||
/// - "gemini_native": Gemini Native streamGenerateContent
|
||||
///
|
||||
/// `extra_headers` 是一个可选的供应商级自定义 header 集合（从 OpenClaw
|
||||
/// 的 `settings_config.headers` 或 OpenCode 的 `settings_config.options.headers`
|
||||
/// 读取），在所有内置 header 之后追加，用于覆盖或补充（例如自定义 User-Agent）。
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
async fn check_claude_stream(
|
||||
client: &Client,
|
||||
@@ -323,6 +325,7 @@ impl StreamCheckService {
|
||||
timeout: std::time::Duration,
|
||||
provider: &Provider,
|
||||
claude_api_format_override: Option<&str>,
|
||||
extra_headers: Option<&serde_json::Map<String, serde_json::Value>>,
|
||||
) -> Result<(u16, String), AppError> {
|
||||
let base = base_url.trim_end_matches('/');
|
||||
let is_github_copilot = auth.strategy == AuthStrategy::GitHubCopilot;
|
||||
@@ -367,8 +370,16 @@ impl StreamCheckService {
|
||||
"messages": [{ "role": "user", "content": test_prompt }],
|
||||
"stream": true
|
||||
});
|
||||
// Codex OAuth (ChatGPT Plus/Pro 反代) 需要 store:false + include 标记，
|
||||
// 否则 Stream Check 会和生产路径一样被服务端 400 拒绝。
|
||||
let is_codex_oauth = provider
|
||||
.meta
|
||||
.as_ref()
|
||||
.and_then(|m| m.provider_type.as_deref())
|
||||
== Some("codex_oauth");
|
||||
|
||||
let body = if is_openai_responses {
|
||||
anthropic_to_responses(anthropic_body, Some(&provider.id))
|
||||
anthropic_to_responses(anthropic_body, Some(&provider.id), is_codex_oauth)
|
||||
.map_err(|e| AppError::Message(format!("Failed to build test request: {e}")))?
|
||||
} else if is_gemini_native {
|
||||
anthropic_to_gemini(anthropic_body)
|
||||
@@ -475,6 +486,15 @@ impl StreamCheckService {
|
||||
.header("connection", "keep-alive");
|
||||
}
|
||||
|
||||
// 供应商自定义 headers 最后追加，允许覆盖内置默认值（例如 user-agent）
|
||||
if let Some(headers) = extra_headers {
|
||||
for (key, value) in headers {
|
||||
if let Some(v) = value.as_str() {
|
||||
request_builder = request_builder.header(key.as_str(), v);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let response = request_builder
|
||||
.timeout(timeout)
|
||||
.json(&body)
|
||||
@@ -595,6 +615,7 @@ impl StreamCheckService {
|
||||
model: &str,
|
||||
test_prompt: &str,
|
||||
timeout: std::time::Duration,
|
||||
extra_headers: Option<&serde_json::Map<String, serde_json::Value>>,
|
||||
) -> Result<(u16, String), AppError> {
|
||||
let base = base_url.trim_end_matches('/');
|
||||
// Gemini 原生 API: /v1beta/models/{model}:streamGenerateContent?alt=sse
|
||||
@@ -614,11 +635,22 @@ impl StreamCheckService {
|
||||
}]
|
||||
});
|
||||
|
||||
let response = client
|
||||
let mut request_builder = client
|
||||
.post(&url)
|
||||
.header("x-goog-api-key", &auth.api_key)
|
||||
.header("Content-Type", "application/json")
|
||||
.header("Accept", "text/event-stream")
|
||||
.header("Accept", "text/event-stream");
|
||||
|
||||
// 供应商自定义 headers 最后追加
|
||||
if let Some(headers) = extra_headers {
|
||||
for (key, value) in headers {
|
||||
if let Some(v) = value.as_str() {
|
||||
request_builder = request_builder.header(key.as_str(), v);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let response = request_builder
|
||||
.timeout(timeout)
|
||||
.json(&body)
|
||||
.send()
|
||||
@@ -643,6 +675,451 @@ impl StreamCheckService {
|
||||
}
|
||||
}
|
||||
|
||||
/// OpenCode / OpenClaw 的独立分发入口（绕过 `get_adapter`）
|
||||
///
|
||||
/// 这两个应用的 `settings_config` 与 Claude/Codex/Gemini 完全不同：
|
||||
/// - OpenClaw: `{ baseUrl, apiKey, api, models: [...] }`，`api` 字段标识协议
|
||||
/// - OpenCode: `{ npm, options: { baseURL, apiKey }, models: {...} }`，`npm` 字段标识协议
|
||||
///
|
||||
/// 因此不能复用 `get_adapter`（会 fallback 到 CodexAdapter 而提取失败），
|
||||
/// 改为独立解析 base_url/api_key/协议，再分发到现有的 check_*_stream 函数。
|
||||
async fn check_once_without_adapter(
|
||||
app_type: &AppType,
|
||||
provider: &Provider,
|
||||
config: &StreamCheckConfig,
|
||||
start: Instant,
|
||||
) -> Result<StreamCheckResult, AppError> {
|
||||
// 获取 HTTP 客户端：优先使用供应商单独代理配置，否则使用全局客户端
|
||||
let proxy_config = provider.meta.as_ref().and_then(|m| m.proxy_config.as_ref());
|
||||
let client = crate::proxy::http_client::get_for_provider(proxy_config);
|
||||
let request_timeout = std::time::Duration::from_secs(config.timeout_secs);
|
||||
|
||||
let model_to_test = Self::resolve_test_model(app_type, provider, config);
|
||||
let test_prompt = &config.test_prompt;
|
||||
|
||||
let result = match app_type {
|
||||
AppType::OpenClaw => {
|
||||
Self::check_openclaw_stream(
|
||||
&client,
|
||||
provider,
|
||||
&model_to_test,
|
||||
test_prompt,
|
||||
request_timeout,
|
||||
)
|
||||
.await
|
||||
}
|
||||
AppType::OpenCode => {
|
||||
Self::check_opencode_stream(
|
||||
&client,
|
||||
provider,
|
||||
&model_to_test,
|
||||
test_prompt,
|
||||
request_timeout,
|
||||
)
|
||||
.await
|
||||
}
|
||||
_ => unreachable!("check_once_without_adapter 只处理 OpenCode/OpenClaw"),
|
||||
};
|
||||
|
||||
let response_time = start.elapsed().as_millis() as u64;
|
||||
Ok(Self::build_stream_check_result(
|
||||
result,
|
||||
response_time,
|
||||
config.degraded_threshold_ms,
|
||||
))
|
||||
}
|
||||
|
||||
/// 将 check_*_stream 的原始结果包装成 StreamCheckResult
|
||||
///
|
||||
/// 抽取自 check_once 的末尾逻辑，以便 OpenCode/OpenClaw 的独立分支复用。
|
||||
fn build_stream_check_result(
|
||||
result: Result<(u16, String), AppError>,
|
||||
response_time: u64,
|
||||
degraded_threshold_ms: u64,
|
||||
) -> StreamCheckResult {
|
||||
let tested_at = chrono::Utc::now().timestamp();
|
||||
match result {
|
||||
Ok((status_code, model)) => StreamCheckResult {
|
||||
status: Self::determine_status(response_time, degraded_threshold_ms),
|
||||
success: true,
|
||||
message: "Check succeeded".to_string(),
|
||||
response_time_ms: Some(response_time),
|
||||
http_status: Some(status_code),
|
||||
model_used: model,
|
||||
tested_at,
|
||||
retry_count: 0,
|
||||
},
|
||||
Err(e) => StreamCheckResult {
|
||||
status: HealthStatus::Failed,
|
||||
success: false,
|
||||
message: e.to_string(),
|
||||
response_time_ms: Some(response_time),
|
||||
http_status: None,
|
||||
model_used: String::new(),
|
||||
tested_at,
|
||||
retry_count: 0,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
/// OpenClaw 流式检查分发器
|
||||
///
|
||||
/// 根据 `settings_config.api` 字段分发到对应协议的检查器。
|
||||
/// 取值参见 `openclawApiProtocols` (前端 openclawProviderPresets.ts):
|
||||
/// - `openai-completions` → check_claude_stream + api_format="openai_chat"
|
||||
/// - `openai-responses` → check_claude_stream + api_format="openai_responses"
|
||||
/// - `anthropic-messages` → check_claude_stream + api_format="anthropic" (ClaudeAuth 策略)
|
||||
/// - `google-generative-ai` → check_gemini_stream (Google API Key 策略)
|
||||
/// - `bedrock-converse-stream` → 不支持（需要 AWS SigV4 签名）
|
||||
async fn check_openclaw_stream(
|
||||
client: &Client,
|
||||
provider: &Provider,
|
||||
model: &str,
|
||||
test_prompt: &str,
|
||||
timeout: std::time::Duration,
|
||||
) -> Result<(u16, String), AppError> {
|
||||
// 自定义认证头（如 Longcat 的 `apikey` 头）不走标准 Bearer，
|
||||
// 具体头名由 OpenClaw 网关内部决定，cc-switch 无法准确构造，
|
||||
// 因此直接返回友好错误而不是让用户看到一个误导性的 401。
|
||||
if Self::openclaw_uses_auth_header(provider) {
|
||||
return Err(AppError::localized(
|
||||
"openclaw_auth_header_not_supported",
|
||||
"该供应商使用自定义认证头，暂不支持流式健康检查。建议直接通过 OpenClaw 测试。",
|
||||
"This provider uses a custom auth header; stream health check is not supported. Please test it directly via OpenClaw.",
|
||||
));
|
||||
}
|
||||
|
||||
let base_url = Self::extract_openclaw_base_url(provider)?;
|
||||
let api_key = Self::extract_openclaw_api_key(provider)?;
|
||||
let api = Self::extract_openclaw_protocol(provider);
|
||||
let extra_headers = Self::extract_openclaw_headers(provider);
|
||||
|
||||
match api.as_deref() {
|
||||
Some("openai-completions") => {
|
||||
let auth = AuthInfo::new(api_key, AuthStrategy::Bearer);
|
||||
Self::check_claude_stream(
|
||||
client,
|
||||
&base_url,
|
||||
&auth,
|
||||
model,
|
||||
test_prompt,
|
||||
timeout,
|
||||
provider,
|
||||
Some("openai_chat"),
|
||||
extra_headers,
|
||||
)
|
||||
.await
|
||||
}
|
||||
Some("openai-responses") => {
|
||||
let auth = AuthInfo::new(api_key, AuthStrategy::Bearer);
|
||||
Self::check_claude_stream(
|
||||
client,
|
||||
&base_url,
|
||||
&auth,
|
||||
model,
|
||||
test_prompt,
|
||||
timeout,
|
||||
provider,
|
||||
Some("openai_responses"),
|
||||
extra_headers,
|
||||
)
|
||||
.await
|
||||
}
|
||||
Some("anthropic-messages") => {
|
||||
// 使用 ClaudeAuth（Bearer-only）以兼容 Claude 中转服务。
|
||||
// 某些中转同时收到 Authorization 和 x-api-key 会报错，ClaudeAuth
|
||||
// 策略保证只下发 Bearer。官方 Anthropic 也接受纯 Bearer。
|
||||
let auth = AuthInfo::new(api_key, AuthStrategy::ClaudeAuth);
|
||||
Self::check_claude_stream(
|
||||
client,
|
||||
&base_url,
|
||||
&auth,
|
||||
model,
|
||||
test_prompt,
|
||||
timeout,
|
||||
provider,
|
||||
Some("anthropic"),
|
||||
extra_headers,
|
||||
)
|
||||
.await
|
||||
}
|
||||
Some("google-generative-ai") => {
|
||||
let auth = AuthInfo::new(api_key, AuthStrategy::Google);
|
||||
Self::check_gemini_stream(
|
||||
client,
|
||||
&base_url,
|
||||
&auth,
|
||||
model,
|
||||
test_prompt,
|
||||
timeout,
|
||||
extra_headers,
|
||||
)
|
||||
.await
|
||||
}
|
||||
Some("bedrock-converse-stream") => Err(AppError::localized(
|
||||
"openclaw_bedrock_not_supported",
|
||||
"AWS Bedrock 需要 SigV4 签名，当前不支持健康检查。请通过 AWS 控制台或 OpenClaw 验证连通性。",
|
||||
"AWS Bedrock requires SigV4 signing and is not supported by stream health check. Please verify connectivity via AWS console or OpenClaw.",
|
||||
)),
|
||||
Some(other) => Err(AppError::localized(
|
||||
"openclaw_protocol_not_yet_supported",
|
||||
format!("OpenClaw 暂不支持协议: {other}"),
|
||||
format!("OpenClaw protocol not yet supported: {other}"),
|
||||
)),
|
||||
None => Err(AppError::localized(
|
||||
"openclaw_protocol_missing",
|
||||
"OpenClaw 供应商缺少 api 字段",
|
||||
"OpenClaw provider is missing the `api` field",
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
/// 判断 OpenClaw 供应商是否使用自定义认证头（`authHeader: true`）
|
||||
fn openclaw_uses_auth_header(provider: &Provider) -> bool {
|
||||
provider
|
||||
.settings_config
|
||||
.get("authHeader")
|
||||
.and_then(|v| v.as_bool())
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
/// 提取 OpenClaw 供应商的自定义 headers（来自 `settings_config.headers`）
|
||||
fn extract_openclaw_headers(
|
||||
provider: &Provider,
|
||||
) -> Option<&serde_json::Map<String, serde_json::Value>> {
|
||||
provider
|
||||
.settings_config
|
||||
.get("headers")
|
||||
.and_then(|v| v.as_object())
|
||||
.filter(|m| !m.is_empty())
|
||||
}
|
||||
|
||||
fn extract_openclaw_base_url(provider: &Provider) -> Result<String, AppError> {
|
||||
provider
|
||||
.settings_config
|
||||
.get("baseUrl")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|s| s.trim().to_string())
|
||||
.filter(|s| !s.is_empty())
|
||||
.ok_or_else(|| {
|
||||
AppError::localized(
|
||||
"openclaw_base_url_missing",
|
||||
"OpenClaw 供应商缺少 baseUrl",
|
||||
"OpenClaw provider is missing `baseUrl`",
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
fn extract_openclaw_api_key(provider: &Provider) -> Result<String, AppError> {
|
||||
provider
|
||||
.settings_config
|
||||
.get("apiKey")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|s| s.trim().to_string())
|
||||
.filter(|s| !s.is_empty())
|
||||
.ok_or_else(|| {
|
||||
AppError::localized(
|
||||
"openclaw_api_key_missing",
|
||||
"OpenClaw 供应商缺少 apiKey",
|
||||
"OpenClaw provider is missing `apiKey`",
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
fn extract_openclaw_protocol(provider: &Provider) -> Option<String> {
|
||||
provider
|
||||
.settings_config
|
||||
.get("api")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|s| s.trim().to_string())
|
||||
.filter(|s| !s.is_empty())
|
||||
}
|
||||
|
||||
/// OpenCode 流式检查分发器
|
||||
///
|
||||
/// OpenCode 用 `npm` 字段（AI SDK 包名）隐式指定协议。映射关系参见
|
||||
/// `opencodeNpmPackages` (前端 opencodeProviderPresets.ts):
|
||||
/// - `@ai-sdk/openai-compatible` → check_claude_stream + api_format="openai_chat"
|
||||
/// - `@ai-sdk/openai` → check_claude_stream + api_format="openai_responses"
|
||||
/// - `@ai-sdk/anthropic` → check_claude_stream + api_format="anthropic"
|
||||
/// - `@ai-sdk/google` → check_gemini_stream (Google API Key 策略)
|
||||
/// - `@ai-sdk/amazon-bedrock` → 不支持（需要 AWS SigV4 签名）
|
||||
///
|
||||
/// URL/API Key 存放在 `settings_config.options.{baseURL,apiKey}`，注意
|
||||
/// `baseURL` 大写 L（与 OpenClaw 的 `baseUrl` 首字母小写 u 不同）。
|
||||
async fn check_opencode_stream(
|
||||
client: &Client,
|
||||
provider: &Provider,
|
||||
model: &str,
|
||||
test_prompt: &str,
|
||||
timeout: std::time::Duration,
|
||||
) -> Result<(u16, String), AppError> {
|
||||
let npm = Self::extract_opencode_npm(provider);
|
||||
// 若用户未显式填 baseURL，则根据 npm 回退到 AI SDK 包自带的默认端点
|
||||
let base_url = Self::resolve_opencode_base_url(provider, npm.as_deref())?;
|
||||
let api_key = Self::extract_opencode_api_key(provider)?;
|
||||
let extra_headers = Self::extract_opencode_headers(provider);
|
||||
|
||||
match npm.as_deref() {
|
||||
Some("@ai-sdk/openai-compatible") => {
|
||||
let auth = AuthInfo::new(api_key, AuthStrategy::Bearer);
|
||||
Self::check_claude_stream(
|
||||
client,
|
||||
&base_url,
|
||||
&auth,
|
||||
model,
|
||||
test_prompt,
|
||||
timeout,
|
||||
provider,
|
||||
Some("openai_chat"),
|
||||
extra_headers,
|
||||
)
|
||||
.await
|
||||
}
|
||||
Some("@ai-sdk/openai") => {
|
||||
let auth = AuthInfo::new(api_key, AuthStrategy::Bearer);
|
||||
Self::check_claude_stream(
|
||||
client,
|
||||
&base_url,
|
||||
&auth,
|
||||
model,
|
||||
test_prompt,
|
||||
timeout,
|
||||
provider,
|
||||
Some("openai_responses"),
|
||||
extra_headers,
|
||||
)
|
||||
.await
|
||||
}
|
||||
Some("@ai-sdk/anthropic") => {
|
||||
// 见 check_openclaw_stream 对 anthropic-messages 的注释：
|
||||
// 用 ClaudeAuth（Bearer-only）兼容中转服务。
|
||||
let auth = AuthInfo::new(api_key, AuthStrategy::ClaudeAuth);
|
||||
Self::check_claude_stream(
|
||||
client,
|
||||
&base_url,
|
||||
&auth,
|
||||
model,
|
||||
test_prompt,
|
||||
timeout,
|
||||
provider,
|
||||
Some("anthropic"),
|
||||
extra_headers,
|
||||
)
|
||||
.await
|
||||
}
|
||||
Some("@ai-sdk/google") => {
|
||||
let auth = AuthInfo::new(api_key, AuthStrategy::Google);
|
||||
Self::check_gemini_stream(
|
||||
client,
|
||||
&base_url,
|
||||
&auth,
|
||||
model,
|
||||
test_prompt,
|
||||
timeout,
|
||||
extra_headers,
|
||||
)
|
||||
.await
|
||||
}
|
||||
Some("@ai-sdk/amazon-bedrock") => Err(AppError::localized(
|
||||
"opencode_bedrock_not_supported",
|
||||
"AWS Bedrock 需要 SigV4 签名，当前不支持健康检查。请通过 AWS 控制台或 OpenCode 验证连通性。",
|
||||
"AWS Bedrock requires SigV4 signing and is not supported by stream health check. Please verify connectivity via AWS console or OpenCode.",
|
||||
)),
|
||||
Some(other) => Err(AppError::localized(
|
||||
"opencode_npm_not_yet_supported",
|
||||
format!("OpenCode 暂不支持 SDK 包: {other}"),
|
||||
format!("OpenCode SDK package not yet supported: {other}"),
|
||||
)),
|
||||
None => Err(AppError::localized(
|
||||
"opencode_npm_missing",
|
||||
"OpenCode 供应商缺少 npm 字段",
|
||||
"OpenCode provider is missing the `npm` field",
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
/// 按 OpenCode 的实际 SDK 包特性确定 baseURL：
|
||||
/// - 用户显式填写的 `options.baseURL` 总是优先
|
||||
/// - 否则根据 `npm` 返回 AI SDK 包自带的默认端点
|
||||
/// - `@ai-sdk/openai-compatible` 没有默认端点，必须显式填
|
||||
///
|
||||
/// 注意：这里的默认端点对应 AI SDK 包的行为（例如 `@ai-sdk/openai`
|
||||
/// 自带 `/v1` 路径后缀），与 `proxy/providers/mod.rs` 里的
|
||||
/// `ProviderType::default_endpoint()` 语义不同——后者是代理层的上游
|
||||
/// 默认值，不带 `/v1`。两者维护的是不同系统的默认值，不能简单共享。
|
||||
fn resolve_opencode_base_url(
|
||||
provider: &Provider,
|
||||
npm: Option<&str>,
|
||||
) -> Result<String, AppError> {
|
||||
if let Some(explicit) = Self::extract_opencode_base_url(provider) {
|
||||
return Ok(explicit);
|
||||
}
|
||||
|
||||
let fallback = match npm {
|
||||
Some("@ai-sdk/openai") => Some("https://api.openai.com/v1"),
|
||||
Some("@ai-sdk/anthropic") => Some("https://api.anthropic.com"),
|
||||
Some("@ai-sdk/google") => Some("https://generativelanguage.googleapis.com"),
|
||||
_ => None,
|
||||
};
|
||||
|
||||
fallback.map(|s| s.to_string()).ok_or_else(|| {
|
||||
AppError::localized(
|
||||
"opencode_base_url_missing",
|
||||
"OpenCode 供应商缺少 options.baseURL，且当前 SDK 包没有默认端点",
|
||||
"OpenCode provider is missing `options.baseURL` and the SDK package has no default endpoint",
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
fn extract_opencode_base_url(provider: &Provider) -> Option<String> {
|
||||
provider
|
||||
.settings_config
|
||||
.get("options")
|
||||
.and_then(|v| v.get("baseURL"))
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|s| s.trim().to_string())
|
||||
.filter(|s| !s.is_empty())
|
||||
}
|
||||
|
||||
/// 提取 OpenCode 供应商的自定义 headers（来自 `settings_config.options.headers`）
|
||||
fn extract_opencode_headers(
|
||||
provider: &Provider,
|
||||
) -> Option<&serde_json::Map<String, serde_json::Value>> {
|
||||
provider
|
||||
.settings_config
|
||||
.get("options")
|
||||
.and_then(|v| v.get("headers"))
|
||||
.and_then(|v| v.as_object())
|
||||
.filter(|m| !m.is_empty())
|
||||
}
|
||||
|
||||
fn extract_opencode_api_key(provider: &Provider) -> Result<String, AppError> {
|
||||
provider
|
||||
.settings_config
|
||||
.get("options")
|
||||
.and_then(|v| v.get("apiKey"))
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|s| s.trim().to_string())
|
||||
.filter(|s| !s.is_empty())
|
||||
.ok_or_else(|| {
|
||||
AppError::localized(
|
||||
"opencode_api_key_missing",
|
||||
"OpenCode 供应商缺少 options.apiKey",
|
||||
"OpenCode provider is missing `options.apiKey`",
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
fn extract_opencode_npm(provider: &Provider) -> Option<String> {
|
||||
provider
|
||||
.settings_config
|
||||
.get("npm")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|s| s.trim().to_string())
|
||||
.filter(|s| !s.is_empty())
|
||||
}
|
||||
|
||||
fn determine_status(latency_ms: u64, threshold: u64) -> HealthStatus {
|
||||
if latency_ms <= threshold {
|
||||
HealthStatus::Operational
|
||||
@@ -846,6 +1323,127 @@ impl StreamCheckService {
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn make_provider(settings_config: serde_json::Value) -> Provider {
|
||||
Provider::with_id(
|
||||
"test".to_string(),
|
||||
"Test".to_string(),
|
||||
settings_config,
|
||||
None,
|
||||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_openclaw_uses_auth_header_true() {
|
||||
let p = make_provider(serde_json::json!({
|
||||
"baseUrl": "https://api.longcat.chat/v1",
|
||||
"apiKey": "k",
|
||||
"api": "openai-completions",
|
||||
"authHeader": true,
|
||||
}));
|
||||
assert!(StreamCheckService::openclaw_uses_auth_header(&p));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_openclaw_uses_auth_header_default_false() {
|
||||
let p = make_provider(serde_json::json!({
|
||||
"baseUrl": "https://api.deepseek.com/v1",
|
||||
"apiKey": "k",
|
||||
"api": "openai-completions",
|
||||
}));
|
||||
assert!(!StreamCheckService::openclaw_uses_auth_header(&p));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_resolve_opencode_base_url_explicit_wins() {
|
||||
let p = make_provider(serde_json::json!({
|
||||
"npm": "@ai-sdk/openai",
|
||||
"options": { "baseURL": "https://proxy.local/v1", "apiKey": "k" },
|
||||
"models": {},
|
||||
}));
|
||||
let resolved =
|
||||
StreamCheckService::resolve_opencode_base_url(&p, Some("@ai-sdk/openai")).unwrap();
|
||||
assert_eq!(resolved, "https://proxy.local/v1");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_resolve_opencode_base_url_falls_back_for_known_npm() {
|
||||
let p = make_provider(serde_json::json!({
|
||||
"npm": "@ai-sdk/openai",
|
||||
"options": { "apiKey": "k" },
|
||||
"models": {},
|
||||
}));
|
||||
let resolved =
|
||||
StreamCheckService::resolve_opencode_base_url(&p, Some("@ai-sdk/openai")).unwrap();
|
||||
assert_eq!(resolved, "https://api.openai.com/v1");
|
||||
|
||||
let p2 = make_provider(serde_json::json!({
|
||||
"npm": "@ai-sdk/anthropic",
|
||||
"options": { "apiKey": "k" },
|
||||
"models": {},
|
||||
}));
|
||||
let resolved2 =
|
||||
StreamCheckService::resolve_opencode_base_url(&p2, Some("@ai-sdk/anthropic")).unwrap();
|
||||
assert_eq!(resolved2, "https://api.anthropic.com");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_resolve_opencode_base_url_errors_for_openai_compatible_without_url() {
|
||||
// @ai-sdk/openai-compatible 没有默认端点，必须显式填
|
||||
let p = make_provider(serde_json::json!({
|
||||
"npm": "@ai-sdk/openai-compatible",
|
||||
"options": { "apiKey": "k" },
|
||||
"models": {},
|
||||
}));
|
||||
let result =
|
||||
StreamCheckService::resolve_opencode_base_url(&p, Some("@ai-sdk/openai-compatible"));
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_openclaw_headers_preserves_map() {
|
||||
let p = make_provider(serde_json::json!({
|
||||
"baseUrl": "https://example.com/v1",
|
||||
"apiKey": "k",
|
||||
"api": "openai-completions",
|
||||
"headers": { "User-Agent": "MyBot/1.0", "X-Trace": "abc" },
|
||||
}));
|
||||
let headers = StreamCheckService::extract_openclaw_headers(&p).unwrap();
|
||||
assert_eq!(
|
||||
headers.get("User-Agent").and_then(|v| v.as_str()),
|
||||
Some("MyBot/1.0")
|
||||
);
|
||||
assert_eq!(headers.get("X-Trace").and_then(|v| v.as_str()), Some("abc"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_openclaw_headers_ignores_empty_map() {
|
||||
let p = make_provider(serde_json::json!({
|
||||
"baseUrl": "https://example.com/v1",
|
||||
"apiKey": "k",
|
||||
"api": "openai-completions",
|
||||
"headers": {},
|
||||
}));
|
||||
assert!(StreamCheckService::extract_openclaw_headers(&p).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_opencode_headers_from_options() {
|
||||
let p = make_provider(serde_json::json!({
|
||||
"npm": "@ai-sdk/openai-compatible",
|
||||
"options": {
|
||||
"baseURL": "https://example.com/v1",
|
||||
"apiKey": "k",
|
||||
"headers": { "X-Custom": "yes" },
|
||||
},
|
||||
"models": {},
|
||||
}));
|
||||
let headers = StreamCheckService::extract_opencode_headers(&p).unwrap();
|
||||
assert_eq!(
|
||||
headers.get("X-Custom").and_then(|v| v.as_str()),
|
||||
Some("yes")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_determine_status() {
|
||||
assert_eq!(
|
||||
|
||||
@@ -60,7 +60,7 @@ pub struct SubscriptionQuota {
|
||||
}
|
||||
|
||||
impl SubscriptionQuota {
|
||||
fn not_found(tool: &str) -> Self {
|
||||
pub(crate) fn not_found(tool: &str) -> Self {
|
||||
Self {
|
||||
tool: tool.to_string(),
|
||||
credential_status: CredentialStatus::NotFound,
|
||||
@@ -73,7 +73,7 @@ impl SubscriptionQuota {
|
||||
}
|
||||
}
|
||||
|
||||
fn error(tool: &str, status: CredentialStatus, message: String) -> Self {
|
||||
pub(crate) fn error(tool: &str, status: CredentialStatus, message: String) -> Self {
|
||||
Self {
|
||||
tool: tool.to_string(),
|
||||
credential_status: status,
|
||||
@@ -621,8 +621,17 @@ fn unix_ts_to_iso(ts: i64) -> Option<String> {
|
||||
chrono::DateTime::from_timestamp(ts, 0).map(|dt| dt.to_rfc3339())
|
||||
}
|
||||
|
||||
/// 查询 Codex 官方订阅额度
|
||||
async fn query_codex_quota(access_token: &str, account_id: Option<&str>) -> SubscriptionQuota {
|
||||
/// 查询 Codex / ChatGPT 反代订阅额度
|
||||
///
|
||||
/// 参数化 `tool_label` 和 `expired_message` 让该函数可被两个调用点共用：
|
||||
/// - `"codex"` + "Please re-login with Codex CLI."（CLI 凭据路径）
|
||||
/// - `"codex_oauth"` + "Please re-login via cc-switch."（cc-switch 自管 OAuth 路径）
|
||||
pub(crate) async fn query_codex_quota(
|
||||
access_token: &str,
|
||||
account_id: Option<&str>,
|
||||
tool_label: &str,
|
||||
expired_message: &str,
|
||||
) -> SubscriptionQuota {
|
||||
let client = crate::proxy::http_client::get();
|
||||
|
||||
let mut req = client
|
||||
@@ -639,7 +648,7 @@ async fn query_codex_quota(access_token: &str, account_id: Option<&str>) -> Subs
|
||||
Ok(r) => r,
|
||||
Err(e) => {
|
||||
return SubscriptionQuota::error(
|
||||
"codex",
|
||||
tool_label,
|
||||
CredentialStatus::Valid,
|
||||
format!("Network error: {e}"),
|
||||
);
|
||||
@@ -650,16 +659,16 @@ async fn query_codex_quota(access_token: &str, account_id: Option<&str>) -> Subs
|
||||
|
||||
if status == reqwest::StatusCode::UNAUTHORIZED || status == reqwest::StatusCode::FORBIDDEN {
|
||||
return SubscriptionQuota::error(
|
||||
"codex",
|
||||
tool_label,
|
||||
CredentialStatus::Expired,
|
||||
format!("Authentication failed (HTTP {status}). Please re-login with Codex CLI."),
|
||||
format!("{expired_message} (HTTP {status})"),
|
||||
);
|
||||
}
|
||||
|
||||
if !status.is_success() {
|
||||
let body = resp.text().await.unwrap_or_default();
|
||||
return SubscriptionQuota::error(
|
||||
"codex",
|
||||
tool_label,
|
||||
CredentialStatus::Valid,
|
||||
format!("API error (HTTP {status}): {body}"),
|
||||
);
|
||||
@@ -669,7 +678,7 @@ async fn query_codex_quota(access_token: &str, account_id: Option<&str>) -> Subs
|
||||
Ok(v) => v,
|
||||
Err(e) => {
|
||||
return SubscriptionQuota::error(
|
||||
"codex",
|
||||
tool_label,
|
||||
CredentialStatus::Valid,
|
||||
format!("Failed to parse API response: {e}"),
|
||||
);
|
||||
@@ -697,7 +706,7 @@ async fn query_codex_quota(access_token: &str, account_id: Option<&str>) -> Subs
|
||||
}
|
||||
|
||||
SubscriptionQuota {
|
||||
tool: "codex".to_string(),
|
||||
tool: tool_label.to_string(),
|
||||
credential_status: CredentialStatus::Valid,
|
||||
credential_message: None,
|
||||
success: true,
|
||||
@@ -1221,7 +1230,13 @@ pub async fn get_subscription_quota(tool: &str) -> Result<SubscriptionQuota, Str
|
||||
CredentialStatus::Expired => {
|
||||
// 即使可能过期也尝试调用 API
|
||||
if let Some(token) = token {
|
||||
let result = query_codex_quota(&token, account_id.as_deref()).await;
|
||||
let result = query_codex_quota(
|
||||
&token,
|
||||
account_id.as_deref(),
|
||||
"codex",
|
||||
"Authentication failed. Please re-login with Codex CLI.",
|
||||
)
|
||||
.await;
|
||||
if result.success {
|
||||
return Ok(result);
|
||||
}
|
||||
@@ -1234,7 +1249,13 @@ pub async fn get_subscription_quota(tool: &str) -> Result<SubscriptionQuota, Str
|
||||
}
|
||||
CredentialStatus::Valid => {
|
||||
let token = token.expect("token must be Some when status is Valid");
|
||||
Ok(query_codex_quota(&token, account_id.as_deref()).await)
|
||||
Ok(query_codex_quota(
|
||||
&token,
|
||||
account_id.as_deref(),
|
||||
"codex",
|
||||
"Authentication failed. Please re-login with Codex CLI.",
|
||||
)
|
||||
.await)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -113,6 +113,21 @@ pub struct RequestLogDetail {
|
||||
pub status_code: u16,
|
||||
pub error_message: Option<String>,
|
||||
pub created_at: i64,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub data_source: Option<String>,
|
||||
}
|
||||
|
||||
/// SQL fragment: resolve provider_name with fallback for session-based entries.
|
||||
/// Session logs use placeholder provider_ids (_session, _codex_session, _gemini_session)
|
||||
/// that don't exist in the providers table — this COALESCE gives them readable names.
|
||||
fn provider_name_coalesce(log_alias: &str, provider_alias: &str) -> String {
|
||||
format!(
|
||||
"COALESCE({provider_alias}.name, CASE {log_alias}.provider_id \
|
||||
WHEN '_session' THEN 'Claude (Session)' \
|
||||
WHEN '_codex_session' THEN 'Codex (Session)' \
|
||||
WHEN '_gemini_session' THEN 'Gemini (Session)' \
|
||||
ELSE {log_alias}.provider_id END)"
|
||||
)
|
||||
}
|
||||
|
||||
impl Database {
|
||||
@@ -121,44 +136,54 @@ impl Database {
|
||||
&self,
|
||||
start_date: Option<i64>,
|
||||
end_date: Option<i64>,
|
||||
app_type: Option<&str>,
|
||||
) -> Result<UsageSummary, AppError> {
|
||||
let conn = lock_conn!(self.conn);
|
||||
|
||||
let (where_clause, params_vec) = if start_date.is_some() || end_date.is_some() {
|
||||
let mut conditions = Vec::new();
|
||||
let mut params = Vec::new();
|
||||
// Build detail WHERE clause
|
||||
let mut conditions = Vec::new();
|
||||
let mut params_vec: Vec<Box<dyn rusqlite::ToSql>> = Vec::new();
|
||||
|
||||
if let Some(start) = start_date {
|
||||
conditions.push("created_at >= ?");
|
||||
params.push(start);
|
||||
}
|
||||
if let Some(end) = end_date {
|
||||
conditions.push("created_at <= ?");
|
||||
params.push(end);
|
||||
}
|
||||
if let Some(start) = start_date {
|
||||
conditions.push("created_at >= ?");
|
||||
params_vec.push(Box::new(start));
|
||||
}
|
||||
if let Some(end) = end_date {
|
||||
conditions.push("created_at <= ?");
|
||||
params_vec.push(Box::new(end));
|
||||
}
|
||||
if let Some(at) = app_type {
|
||||
conditions.push("app_type = ?");
|
||||
params_vec.push(Box::new(at.to_string()));
|
||||
}
|
||||
|
||||
(format!("WHERE {}", conditions.join(" AND ")), params)
|
||||
let where_clause = if conditions.is_empty() {
|
||||
String::new()
|
||||
} else {
|
||||
(String::new(), Vec::new())
|
||||
format!("WHERE {}", conditions.join(" AND "))
|
||||
};
|
||||
|
||||
// Build rollup WHERE clause using date strings (use ? for sequential binding)
|
||||
let (rollup_where, rollup_params) = if start_date.is_some() || end_date.is_some() {
|
||||
let mut conditions: Vec<String> = Vec::new();
|
||||
let mut params = Vec::new();
|
||||
// Build rollup WHERE clause using date strings
|
||||
let mut rollup_conditions: Vec<String> = Vec::new();
|
||||
let mut rollup_params: Vec<Box<dyn rusqlite::ToSql>> = Vec::new();
|
||||
|
||||
if let Some(start) = start_date {
|
||||
conditions.push("date >= date(?, 'unixepoch', 'localtime')".to_string());
|
||||
params.push(start);
|
||||
}
|
||||
if let Some(end) = end_date {
|
||||
conditions.push("date <= date(?, 'unixepoch', 'localtime')".to_string());
|
||||
params.push(end);
|
||||
}
|
||||
if let Some(start) = start_date {
|
||||
rollup_conditions.push("date >= date(?, 'unixepoch', 'localtime')".to_string());
|
||||
rollup_params.push(Box::new(start));
|
||||
}
|
||||
if let Some(end) = end_date {
|
||||
rollup_conditions.push("date <= date(?, 'unixepoch', 'localtime')".to_string());
|
||||
rollup_params.push(Box::new(end));
|
||||
}
|
||||
if let Some(at) = app_type {
|
||||
rollup_conditions.push("app_type = ?".to_string());
|
||||
rollup_params.push(Box::new(at.to_string()));
|
||||
}
|
||||
|
||||
(format!("WHERE {}", conditions.join(" AND ")), params)
|
||||
let rollup_where = if rollup_conditions.is_empty() {
|
||||
String::new()
|
||||
} else {
|
||||
(String::new(), Vec::new())
|
||||
format!("WHERE {}", rollup_conditions.join(" AND "))
|
||||
};
|
||||
|
||||
let sql = format!(
|
||||
@@ -192,10 +217,11 @@ impl Database {
|
||||
);
|
||||
|
||||
// Combine params: detail params first, then rollup params
|
||||
let mut all_params: Vec<i64> = params_vec;
|
||||
let mut all_params: Vec<Box<dyn rusqlite::ToSql>> = params_vec;
|
||||
all_params.extend(rollup_params);
|
||||
let param_refs: Vec<&dyn rusqlite::ToSql> = all_params.iter().map(|p| p.as_ref()).collect();
|
||||
|
||||
let result = conn.query_row(&sql, rusqlite::params_from_iter(all_params), |row| {
|
||||
let result = conn.query_row(&sql, param_refs.as_slice(), |row| {
|
||||
let total_requests: i64 = row.get(0)?;
|
||||
let total_cost: f64 = row.get(1)?;
|
||||
let total_input_tokens: i64 = row.get(2)?;
|
||||
@@ -229,6 +255,7 @@ impl Database {
|
||||
&self,
|
||||
start_date: Option<i64>,
|
||||
end_date: Option<i64>,
|
||||
app_type: Option<&str>,
|
||||
) -> Result<Vec<DailyStats>, AppError> {
|
||||
let conn = lock_conn!(self.conn);
|
||||
|
||||
@@ -260,9 +287,15 @@ impl Database {
|
||||
bucket_count = 1;
|
||||
}
|
||||
|
||||
let app_type_filter = if app_type.is_some() {
|
||||
"AND app_type = ?4"
|
||||
} else {
|
||||
""
|
||||
};
|
||||
|
||||
// Query detail logs
|
||||
let sql = "
|
||||
SELECT
|
||||
let sql = format!(
|
||||
"SELECT
|
||||
CAST((created_at - ?1) / ?3 AS INTEGER) as bucket_idx,
|
||||
COUNT(*) as request_count,
|
||||
COALESCE(SUM(CAST(total_cost_usd AS REAL)), 0) as total_cost,
|
||||
@@ -272,12 +305,13 @@ impl Database {
|
||||
COALESCE(SUM(cache_creation_tokens), 0) as total_cache_creation_tokens,
|
||||
COALESCE(SUM(cache_read_tokens), 0) as total_cache_read_tokens
|
||||
FROM proxy_request_logs
|
||||
WHERE created_at >= ?1 AND created_at <= ?2
|
||||
WHERE created_at >= ?1 AND created_at <= ?2 {app_type_filter}
|
||||
GROUP BY bucket_idx
|
||||
ORDER BY bucket_idx ASC";
|
||||
ORDER BY bucket_idx ASC"
|
||||
);
|
||||
|
||||
let mut stmt = conn.prepare(sql)?;
|
||||
let rows = stmt.query_map(params![start_ts, end_ts, bucket_seconds], |row| {
|
||||
let mut stmt = conn.prepare(&sql)?;
|
||||
let row_mapper = |row: &rusqlite::Row| {
|
||||
Ok((
|
||||
row.get::<_, i64>(0)?,
|
||||
DailyStats {
|
||||
@@ -291,24 +325,33 @@ impl Database {
|
||||
total_cache_read_tokens: row.get::<_, i64>(7)? as u64,
|
||||
},
|
||||
))
|
||||
})?;
|
||||
};
|
||||
|
||||
let mut map: HashMap<i64, DailyStats> = HashMap::new();
|
||||
for row in rows {
|
||||
let (mut bucket_idx, stat) = row?;
|
||||
if bucket_idx < 0 {
|
||||
continue;
|
||||
|
||||
// Collect rows into map (need to handle both param variants)
|
||||
{
|
||||
let rows = if let Some(at) = app_type {
|
||||
stmt.query_map(params![start_ts, end_ts, bucket_seconds, at], row_mapper)?
|
||||
} else {
|
||||
stmt.query_map(params![start_ts, end_ts, bucket_seconds], row_mapper)?
|
||||
};
|
||||
for row in rows {
|
||||
let (mut bucket_idx, stat) = row?;
|
||||
if bucket_idx < 0 {
|
||||
continue;
|
||||
}
|
||||
if bucket_idx >= bucket_count {
|
||||
bucket_idx = bucket_count - 1;
|
||||
}
|
||||
map.insert(bucket_idx, stat);
|
||||
}
|
||||
if bucket_idx >= bucket_count {
|
||||
bucket_idx = bucket_count - 1;
|
||||
}
|
||||
map.insert(bucket_idx, stat);
|
||||
}
|
||||
|
||||
// Also query rollup data (daily granularity, only useful for daily buckets)
|
||||
if bucket_seconds >= 86400 {
|
||||
let rollup_sql = "
|
||||
SELECT
|
||||
let rollup_sql = format!(
|
||||
"SELECT
|
||||
CAST((CAST(strftime('%s', date) AS INTEGER) - ?1) / ?3 AS INTEGER) as bucket_idx,
|
||||
COALESCE(SUM(request_count), 0),
|
||||
COALESCE(SUM(CAST(total_cost_usd AS REAL)), 0),
|
||||
@@ -318,12 +361,12 @@ impl Database {
|
||||
COALESCE(SUM(cache_creation_tokens), 0),
|
||||
COALESCE(SUM(cache_read_tokens), 0)
|
||||
FROM usage_daily_rollups
|
||||
WHERE date >= date(?1, 'unixepoch', 'localtime') AND date <= date(?2, 'unixepoch', 'localtime')
|
||||
WHERE date >= date(?1, 'unixepoch', 'localtime') AND date <= date(?2, 'unixepoch', 'localtime') {app_type_filter}
|
||||
GROUP BY bucket_idx
|
||||
ORDER BY bucket_idx ASC";
|
||||
ORDER BY bucket_idx ASC"
|
||||
);
|
||||
|
||||
let mut rstmt = conn.prepare(rollup_sql)?;
|
||||
let rrows = rstmt.query_map(params![start_ts, end_ts, bucket_seconds], |row| {
|
||||
let rollup_row_mapper = |row: &rusqlite::Row| {
|
||||
Ok((
|
||||
row.get::<_, i64>(0)?,
|
||||
(
|
||||
@@ -336,7 +379,17 @@ impl Database {
|
||||
row.get::<_, i64>(7)? as u64,
|
||||
),
|
||||
))
|
||||
})?;
|
||||
};
|
||||
|
||||
let mut rstmt = conn.prepare(&rollup_sql)?;
|
||||
let rrows = if let Some(at) = app_type {
|
||||
rstmt.query_map(
|
||||
params![start_ts, end_ts, bucket_seconds, at],
|
||||
rollup_row_mapper,
|
||||
)?
|
||||
} else {
|
||||
rstmt.query_map(params![start_ts, end_ts, bucket_seconds], rollup_row_mapper)?
|
||||
};
|
||||
|
||||
for row in rrows {
|
||||
let (mut bucket_idx, (req, cost, tok, inp, out, cc, cr)) = row?;
|
||||
@@ -398,11 +451,23 @@ impl Database {
|
||||
}
|
||||
|
||||
/// 获取 Provider 统计
|
||||
pub fn get_provider_stats(&self) -> Result<Vec<ProviderStats>, AppError> {
|
||||
pub fn get_provider_stats(
|
||||
&self,
|
||||
app_type: Option<&str>,
|
||||
) -> Result<Vec<ProviderStats>, AppError> {
|
||||
let conn = lock_conn!(self.conn);
|
||||
|
||||
let (detail_where, rollup_where) = if app_type.is_some() {
|
||||
("WHERE l.app_type = ?1", "WHERE r.app_type = ?2")
|
||||
} else {
|
||||
("", "")
|
||||
};
|
||||
|
||||
// UNION detail logs + rollup data, then aggregate
|
||||
let sql = "SELECT
|
||||
let detail_pname = provider_name_coalesce("l", "p");
|
||||
let rollup_pname = provider_name_coalesce("r", "p2");
|
||||
let sql = format!(
|
||||
"SELECT
|
||||
provider_id, app_type, provider_name,
|
||||
SUM(request_count) as request_count,
|
||||
SUM(total_tokens) as total_tokens,
|
||||
@@ -413,7 +478,7 @@ impl Database {
|
||||
ELSE 0 END as avg_latency
|
||||
FROM (
|
||||
SELECT l.provider_id, l.app_type,
|
||||
p.name as provider_name,
|
||||
{detail_pname} as provider_name,
|
||||
COUNT(*) as request_count,
|
||||
COALESCE(SUM(l.input_tokens + l.output_tokens), 0) as total_tokens,
|
||||
COALESCE(SUM(CAST(l.total_cost_usd AS REAL)), 0) as total_cost,
|
||||
@@ -421,10 +486,11 @@ impl Database {
|
||||
COALESCE(SUM(l.latency_ms), 0) as latency_sum
|
||||
FROM proxy_request_logs l
|
||||
LEFT JOIN providers p ON l.provider_id = p.id AND l.app_type = p.app_type
|
||||
{detail_where}
|
||||
GROUP BY l.provider_id, l.app_type
|
||||
UNION ALL
|
||||
SELECT r.provider_id, r.app_type,
|
||||
p2.name as provider_name,
|
||||
{rollup_pname} as provider_name,
|
||||
COALESCE(SUM(r.request_count), 0),
|
||||
COALESCE(SUM(r.input_tokens + r.output_tokens), 0),
|
||||
COALESCE(SUM(CAST(r.total_cost_usd AS REAL)), 0),
|
||||
@@ -432,13 +498,15 @@ impl Database {
|
||||
COALESCE(SUM(r.avg_latency_ms * r.request_count), 0)
|
||||
FROM usage_daily_rollups r
|
||||
LEFT JOIN providers p2 ON r.provider_id = p2.id AND r.app_type = p2.app_type
|
||||
{rollup_where}
|
||||
GROUP BY r.provider_id, r.app_type
|
||||
)
|
||||
GROUP BY provider_id, app_type
|
||||
ORDER BY total_cost DESC";
|
||||
ORDER BY total_cost DESC"
|
||||
);
|
||||
|
||||
let mut stmt = conn.prepare(sql)?;
|
||||
let rows = stmt.query_map([], |row| {
|
||||
let mut stmt = conn.prepare(&sql)?;
|
||||
let row_mapper = |row: &rusqlite::Row| {
|
||||
let request_count: i64 = row.get(3)?;
|
||||
let success_count: i64 = row.get(6)?;
|
||||
let success_rate = if request_count > 0 {
|
||||
@@ -449,16 +517,20 @@ impl Database {
|
||||
|
||||
Ok(ProviderStats {
|
||||
provider_id: row.get(0)?,
|
||||
provider_name: row
|
||||
.get::<_, Option<String>>(2)?
|
||||
.unwrap_or_else(|| "Unknown".to_string()),
|
||||
provider_name: row.get(2)?,
|
||||
request_count: request_count as u64,
|
||||
total_tokens: row.get::<_, i64>(4)? as u64,
|
||||
total_cost: format!("{:.6}", row.get::<_, f64>(5)?),
|
||||
success_rate,
|
||||
avg_latency_ms: row.get::<_, f64>(7)? as u64,
|
||||
})
|
||||
})?;
|
||||
};
|
||||
|
||||
let rows = if let Some(at) = app_type {
|
||||
stmt.query_map(params![at, at], row_mapper)?
|
||||
} else {
|
||||
stmt.query_map([], row_mapper)?
|
||||
};
|
||||
|
||||
let mut stats = Vec::new();
|
||||
for row in rows {
|
||||
@@ -469,11 +541,18 @@ impl Database {
|
||||
}
|
||||
|
||||
/// 获取模型统计
|
||||
pub fn get_model_stats(&self) -> Result<Vec<ModelStats>, AppError> {
|
||||
pub fn get_model_stats(&self, app_type: Option<&str>) -> Result<Vec<ModelStats>, AppError> {
|
||||
let conn = lock_conn!(self.conn);
|
||||
|
||||
let (detail_where, rollup_where) = if app_type.is_some() {
|
||||
("WHERE app_type = ?1", "WHERE app_type = ?2")
|
||||
} else {
|
||||
("", "")
|
||||
};
|
||||
|
||||
// UNION detail logs + rollup data
|
||||
let sql = "SELECT
|
||||
let sql = format!(
|
||||
"SELECT
|
||||
model,
|
||||
SUM(request_count) as request_count,
|
||||
SUM(total_tokens) as total_tokens,
|
||||
@@ -484,6 +563,7 @@ impl Database {
|
||||
COALESCE(SUM(input_tokens + output_tokens), 0) as total_tokens,
|
||||
COALESCE(SUM(CAST(total_cost_usd AS REAL)), 0) as total_cost
|
||||
FROM proxy_request_logs
|
||||
{detail_where}
|
||||
GROUP BY model
|
||||
UNION ALL
|
||||
SELECT model,
|
||||
@@ -491,13 +571,15 @@ impl Database {
|
||||
COALESCE(SUM(input_tokens + output_tokens), 0),
|
||||
COALESCE(SUM(CAST(total_cost_usd AS REAL)), 0)
|
||||
FROM usage_daily_rollups
|
||||
{rollup_where}
|
||||
GROUP BY model
|
||||
)
|
||||
GROUP BY model
|
||||
ORDER BY total_cost DESC";
|
||||
ORDER BY total_cost DESC"
|
||||
);
|
||||
|
||||
let mut stmt = conn.prepare(sql)?;
|
||||
let rows = stmt.query_map([], |row| {
|
||||
let mut stmt = conn.prepare(&sql)?;
|
||||
let row_mapper = |row: &rusqlite::Row| {
|
||||
let request_count: i64 = row.get(1)?;
|
||||
let total_cost: f64 = row.get(3)?;
|
||||
let avg_cost = if request_count > 0 {
|
||||
@@ -513,7 +595,13 @@ impl Database {
|
||||
total_cost: format!("{total_cost:.6}"),
|
||||
avg_cost_per_request: format!("{avg_cost:.6}"),
|
||||
})
|
||||
})?;
|
||||
};
|
||||
|
||||
let rows = if let Some(at) = app_type {
|
||||
stmt.query_map(params![at, at], row_mapper)?
|
||||
} else {
|
||||
stmt.query_map([], row_mapper)?
|
||||
};
|
||||
|
||||
let mut stats = Vec::new();
|
||||
for row in rows {
|
||||
@@ -582,13 +670,14 @@ impl Database {
|
||||
params.push(Box::new(page_size as i64));
|
||||
params.push(Box::new(offset as i64));
|
||||
|
||||
let logs_pname = provider_name_coalesce("l", "p");
|
||||
let sql = format!(
|
||||
"SELECT l.request_id, l.provider_id, p.name as provider_name, l.app_type, l.model,
|
||||
"SELECT l.request_id, l.provider_id, {logs_pname} as provider_name, l.app_type, l.model,
|
||||
l.request_model, l.cost_multiplier,
|
||||
l.input_tokens, l.output_tokens, l.cache_read_tokens, l.cache_creation_tokens,
|
||||
l.input_cost_usd, l.output_cost_usd, l.cache_read_cost_usd, l.cache_creation_cost_usd, l.total_cost_usd,
|
||||
l.is_streaming, l.latency_ms, l.first_token_ms, l.duration_ms,
|
||||
l.status_code, l.error_message, l.created_at
|
||||
l.status_code, l.error_message, l.created_at, l.data_source
|
||||
FROM proxy_request_logs l
|
||||
LEFT JOIN providers p ON l.provider_id = p.id AND l.app_type = p.app_type
|
||||
{where_clause}
|
||||
@@ -625,6 +714,7 @@ impl Database {
|
||||
status_code: row.get::<_, i64>(20)? as u16,
|
||||
error_message: row.get(21)?,
|
||||
created_at: row.get(22)?,
|
||||
data_source: row.get(23)?,
|
||||
})
|
||||
})?;
|
||||
|
||||
@@ -658,45 +748,48 @@ impl Database {
|
||||
) -> Result<Option<RequestLogDetail>, AppError> {
|
||||
let conn = lock_conn!(self.conn);
|
||||
|
||||
let result = conn.query_row(
|
||||
"SELECT l.request_id, l.provider_id, p.name as provider_name, l.app_type, l.model,
|
||||
let detail_pname = provider_name_coalesce("l", "p");
|
||||
let detail_sql = format!(
|
||||
"SELECT l.request_id, l.provider_id, {detail_pname} as provider_name, l.app_type, l.model,
|
||||
l.request_model, l.cost_multiplier,
|
||||
input_tokens, output_tokens, cache_read_tokens, cache_creation_tokens,
|
||||
input_cost_usd, output_cost_usd, cache_read_cost_usd, cache_creation_cost_usd, total_cost_usd,
|
||||
is_streaming, latency_ms, first_token_ms, duration_ms,
|
||||
status_code, error_message, created_at
|
||||
status_code, error_message, created_at, l.data_source
|
||||
FROM proxy_request_logs l
|
||||
LEFT JOIN providers p ON l.provider_id = p.id AND l.app_type = p.app_type
|
||||
WHERE l.request_id = ?",
|
||||
[request_id],
|
||||
|row| {
|
||||
Ok(RequestLogDetail {
|
||||
request_id: row.get(0)?,
|
||||
provider_id: row.get(1)?,
|
||||
provider_name: row.get(2)?,
|
||||
app_type: row.get(3)?,
|
||||
model: row.get(4)?,
|
||||
request_model: row.get(5)?,
|
||||
cost_multiplier: row.get::<_, Option<String>>(6)?.unwrap_or_else(|| "1".to_string()),
|
||||
input_tokens: row.get::<_, i64>(7)? as u32,
|
||||
output_tokens: row.get::<_, i64>(8)? as u32,
|
||||
cache_read_tokens: row.get::<_, i64>(9)? as u32,
|
||||
cache_creation_tokens: row.get::<_, i64>(10)? as u32,
|
||||
input_cost_usd: row.get(11)?,
|
||||
output_cost_usd: row.get(12)?,
|
||||
cache_read_cost_usd: row.get(13)?,
|
||||
cache_creation_cost_usd: row.get(14)?,
|
||||
total_cost_usd: row.get(15)?,
|
||||
is_streaming: row.get::<_, i64>(16)? != 0,
|
||||
latency_ms: row.get::<_, i64>(17)? as u64,
|
||||
first_token_ms: row.get::<_, Option<i64>>(18)?.map(|v| v as u64),
|
||||
duration_ms: row.get::<_, Option<i64>>(19)?.map(|v| v as u64),
|
||||
status_code: row.get::<_, i64>(20)? as u16,
|
||||
error_message: row.get(21)?,
|
||||
created_at: row.get(22)?,
|
||||
})
|
||||
},
|
||||
WHERE l.request_id = ?"
|
||||
);
|
||||
let result = conn.query_row(&detail_sql, [request_id], |row| {
|
||||
Ok(RequestLogDetail {
|
||||
request_id: row.get(0)?,
|
||||
provider_id: row.get(1)?,
|
||||
provider_name: row.get(2)?,
|
||||
app_type: row.get(3)?,
|
||||
model: row.get(4)?,
|
||||
request_model: row.get(5)?,
|
||||
cost_multiplier: row
|
||||
.get::<_, Option<String>>(6)?
|
||||
.unwrap_or_else(|| "1".to_string()),
|
||||
input_tokens: row.get::<_, i64>(7)? as u32,
|
||||
output_tokens: row.get::<_, i64>(8)? as u32,
|
||||
cache_read_tokens: row.get::<_, i64>(9)? as u32,
|
||||
cache_creation_tokens: row.get::<_, i64>(10)? as u32,
|
||||
input_cost_usd: row.get(11)?,
|
||||
output_cost_usd: row.get(12)?,
|
||||
cache_read_cost_usd: row.get(13)?,
|
||||
cache_creation_cost_usd: row.get(14)?,
|
||||
total_cost_usd: row.get(15)?,
|
||||
is_streaming: row.get::<_, i64>(16)? != 0,
|
||||
latency_ms: row.get::<_, i64>(17)? as u64,
|
||||
first_token_ms: row.get::<_, Option<i64>>(18)?.map(|v| v as u64),
|
||||
duration_ms: row.get::<_, Option<i64>>(19)?.map(|v| v as u64),
|
||||
status_code: row.get::<_, i64>(20)? as u16,
|
||||
error_message: row.get(21)?,
|
||||
created_at: row.get(22)?,
|
||||
data_source: row.get(23)?,
|
||||
})
|
||||
});
|
||||
|
||||
match result {
|
||||
Ok(mut detail) => {
|
||||
@@ -1040,7 +1133,7 @@ mod tests {
|
||||
)?;
|
||||
}
|
||||
|
||||
let summary = db.get_usage_summary(None, None)?;
|
||||
let summary = db.get_usage_summary(None, None, None)?;
|
||||
assert_eq!(summary.total_requests, 2);
|
||||
assert_eq!(summary.success_rate, 100.0);
|
||||
|
||||
@@ -1075,7 +1168,7 @@ mod tests {
|
||||
)?;
|
||||
}
|
||||
|
||||
let stats = db.get_model_stats()?;
|
||||
let stats = db.get_model_stats(None)?;
|
||||
assert_eq!(stats.len(), 1);
|
||||
assert_eq!(stats[0].model, "claude-3-sonnet");
|
||||
assert_eq!(stats[0].request_count, 1);
|
||||
|
||||
Reference in New Issue
Block a user