1525 lines
52 KiB
Rust
1525 lines
52 KiB
Rust
use super::{
|
|
current_core_status, management_authorization, management_endpoint, management_http_client,
|
|
CoreProcessState, GuiConfigFile, GuiConfigState,
|
|
};
|
|
use chrono::{DateTime, Local};
|
|
use rusqlite::{
|
|
params, params_from_iter, types::Value as SqlValue, Connection, OptionalExtension, Row,
|
|
Transaction,
|
|
};
|
|
use serde::{Deserialize, Serialize};
|
|
use serde_json::Value;
|
|
use sha2::{Digest, Sha256};
|
|
#[cfg(test)]
|
|
use std::time::{SystemTime, UNIX_EPOCH};
|
|
use std::{
|
|
fs,
|
|
path::{Path, PathBuf},
|
|
sync::Mutex,
|
|
time::Duration,
|
|
};
|
|
use tauri::{Emitter, Manager};
|
|
use tokio_util::sync::CancellationToken;
|
|
|
|
const USAGE_DIR_NAME: &str = "usage-records";
|
|
const USAGE_DATABASE_FILE: &str = "usage.db";
|
|
const LEGACY_USAGE_EVENTS_DIR: &str = "events";
|
|
const LEGACY_USAGE_INBOX_DIR: &str = "inbox";
|
|
const LEGACY_JSON_MIGRATION_KEY: &str = "legacy_json_v1";
|
|
const USAGE_UPDATED_EVENT: &str = "usage-records-updated";
|
|
const USAGE_SCHEMA_VERSION: u8 = 1;
|
|
const USAGE_QUEUE_BATCH_SIZE: usize = 500;
|
|
const SQLITE_BUSY_TIMEOUT_SECONDS: u64 = 5;
|
|
|
|
pub(crate) struct UsageCollectorState {
|
|
inner: Mutex<UsageCollectorInner>,
|
|
}
|
|
|
|
struct UsageCollectorInner {
|
|
token: Option<CancellationToken>,
|
|
status: UsageCollectorStatus,
|
|
}
|
|
|
|
impl Default for UsageCollectorState {
|
|
fn default() -> Self {
|
|
Self {
|
|
inner: Mutex::new(UsageCollectorInner {
|
|
token: None,
|
|
status: UsageCollectorStatus::waiting(),
|
|
}),
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Clone, Serialize)]
|
|
#[serde(rename_all = "camelCase")]
|
|
pub(crate) struct UsageCollectorStatus {
|
|
state: String,
|
|
message: String,
|
|
last_collected_at: Option<String>,
|
|
total_records: u64,
|
|
}
|
|
|
|
impl UsageCollectorStatus {
|
|
fn waiting() -> Self {
|
|
Self {
|
|
state: "waiting-core".to_string(),
|
|
message: "等待内核启动".to_string(),
|
|
last_collected_at: None,
|
|
total_records: 0,
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Clone, Default, Serialize, Deserialize)]
|
|
struct UsageTokenStats {
|
|
#[serde(default)]
|
|
input_tokens: u64,
|
|
#[serde(default)]
|
|
output_tokens: u64,
|
|
#[serde(default)]
|
|
reasoning_tokens: u64,
|
|
#[serde(default)]
|
|
cache_read_tokens: u64,
|
|
#[serde(default)]
|
|
cache_creation_tokens: u64,
|
|
#[serde(default)]
|
|
total_tokens: u64,
|
|
}
|
|
|
|
#[derive(Clone, Serialize, Deserialize)]
|
|
pub(crate) struct UsageRecord {
|
|
id: String,
|
|
timestamp: String,
|
|
#[serde(default)]
|
|
latency_ms: u64,
|
|
#[serde(default)]
|
|
ttft_ms: Option<u64>,
|
|
#[serde(default)]
|
|
source: String,
|
|
#[serde(default)]
|
|
auth_index: String,
|
|
#[serde(default)]
|
|
failed: bool,
|
|
#[serde(default)]
|
|
provider: String,
|
|
#[serde(default)]
|
|
model: String,
|
|
#[serde(default)]
|
|
alias: String,
|
|
#[serde(default)]
|
|
reasoning_effort: String,
|
|
#[serde(default)]
|
|
service_tier: String,
|
|
#[serde(default)]
|
|
response_service_tier: String,
|
|
#[serde(default)]
|
|
executor_type: String,
|
|
#[serde(default)]
|
|
endpoint: String,
|
|
#[serde(default)]
|
|
auth_type: String,
|
|
#[serde(default)]
|
|
api_key_hash: String,
|
|
#[serde(default)]
|
|
api_key_display: String,
|
|
#[serde(default)]
|
|
api_key_remark: String,
|
|
#[serde(default)]
|
|
request_id: String,
|
|
#[serde(default)]
|
|
tokens: UsageTokenStats,
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
#[serde(rename_all = "camelCase")]
|
|
struct LegacyUsageHourFile {
|
|
schema_version: u8,
|
|
records: Vec<UsageRecord>,
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
#[serde(rename_all = "camelCase")]
|
|
struct LegacyUsageInboxFile {
|
|
schema_version: u8,
|
|
records: Vec<UsageRecord>,
|
|
}
|
|
|
|
#[derive(Clone, Default, Deserialize)]
|
|
pub(crate) struct UsageQuery {
|
|
#[serde(default)]
|
|
start: Option<String>,
|
|
#[serde(default)]
|
|
end: Option<String>,
|
|
#[serde(default)]
|
|
model: Option<String>,
|
|
#[serde(default)]
|
|
provider: Option<String>,
|
|
#[serde(default)]
|
|
source: Option<String>,
|
|
#[serde(default)]
|
|
api_key_hash: Option<String>,
|
|
#[serde(default)]
|
|
failed: Option<bool>,
|
|
#[serde(default)]
|
|
page: Option<usize>,
|
|
#[serde(default)]
|
|
page_size: Option<usize>,
|
|
}
|
|
|
|
#[derive(Default, Serialize)]
|
|
#[serde(rename_all = "camelCase")]
|
|
pub(crate) struct UsageOverview {
|
|
total_requests: u64,
|
|
success_count: u64,
|
|
failure_count: u64,
|
|
success_rate: f64,
|
|
input_tokens: u64,
|
|
output_tokens: u64,
|
|
reasoning_tokens: u64,
|
|
cache_read_tokens: u64,
|
|
cache_creation_tokens: u64,
|
|
total_tokens: u64,
|
|
rpm: f64,
|
|
tpm: f64,
|
|
average_latency_ms: f64,
|
|
timeline: Vec<UsageTimelinePoint>,
|
|
}
|
|
|
|
#[derive(Default, Serialize)]
|
|
#[serde(rename_all = "camelCase")]
|
|
struct UsageTimelinePoint {
|
|
hour: String,
|
|
requests: u64,
|
|
success: u64,
|
|
failure: u64,
|
|
tokens: u64,
|
|
}
|
|
|
|
#[derive(Default, Serialize)]
|
|
#[serde(rename_all = "camelCase")]
|
|
pub(crate) struct UsageAnalysis {
|
|
models: Vec<UsageCategory>,
|
|
providers: Vec<UsageCategory>,
|
|
sources: Vec<UsageCategory>,
|
|
api_keys: Vec<UsageCategory>,
|
|
}
|
|
|
|
#[derive(Clone, Default, Serialize)]
|
|
#[serde(rename_all = "camelCase")]
|
|
struct UsageCategory {
|
|
key: String,
|
|
label: String,
|
|
requests: u64,
|
|
failures: u64,
|
|
tokens: u64,
|
|
}
|
|
|
|
#[derive(Serialize)]
|
|
#[serde(rename_all = "camelCase")]
|
|
pub(crate) struct UsageEventPage {
|
|
items: Vec<UsageRecord>,
|
|
total: usize,
|
|
page: usize,
|
|
page_size: usize,
|
|
total_pages: usize,
|
|
}
|
|
|
|
struct UsageSqlFilter {
|
|
clause: String,
|
|
params: Vec<SqlValue>,
|
|
}
|
|
|
|
impl UsageCollectorState {
|
|
fn start(&self) -> Option<CancellationToken> {
|
|
let mut inner = self.inner.lock().ok()?;
|
|
if inner.token.is_some() {
|
|
return None;
|
|
}
|
|
let token = CancellationToken::new();
|
|
inner.token = Some(token.clone());
|
|
Some(token)
|
|
}
|
|
|
|
fn stop(&self) {
|
|
if let Ok(mut inner) = self.inner.lock() {
|
|
if let Some(token) = inner.token.take() {
|
|
token.cancel();
|
|
}
|
|
}
|
|
}
|
|
|
|
fn set_status(&self, status: UsageCollectorStatus) {
|
|
if let Ok(mut inner) = self.inner.lock() {
|
|
inner.status = status;
|
|
}
|
|
}
|
|
|
|
fn set_total_records(&self, total_records: u64) {
|
|
if let Ok(mut inner) = self.inner.lock() {
|
|
inner.status.total_records = total_records;
|
|
}
|
|
}
|
|
|
|
fn increment_total_records(&self, added: usize) {
|
|
if let Ok(mut inner) = self.inner.lock() {
|
|
inner.status.total_records = inner.status.total_records.saturating_add(added as u64);
|
|
}
|
|
}
|
|
|
|
fn status(&self) -> Result<UsageCollectorStatus, String> {
|
|
self.inner
|
|
.lock()
|
|
.map(|inner| inner.status.clone())
|
|
.map_err(|_| "使用记录采集状态锁已损坏".to_string())
|
|
}
|
|
}
|
|
|
|
pub(crate) fn initialize_usage_storage() -> Result<(), String> {
|
|
initialize_usage_storage_at(&usage_root_dir()?)
|
|
}
|
|
|
|
fn initialize_usage_storage_at(root: &Path) -> Result<(), String> {
|
|
fs::create_dir_all(root).map_err(|error| format!("创建使用记录目录失败: {error}"))?;
|
|
let mut connection = open_usage_database_at(root)?;
|
|
connection
|
|
.pragma_update(None, "journal_mode", "WAL")
|
|
.map_err(|error| format!("启用 SQLite WAL 失败: {error}"))?;
|
|
initialize_usage_schema(&connection)?;
|
|
migrate_legacy_json_storage(&mut connection, root)
|
|
}
|
|
|
|
fn initialize_usage_schema(connection: &Connection) -> Result<(), String> {
|
|
connection
|
|
.execute_batch(
|
|
r#"
|
|
CREATE TABLE IF NOT EXISTS usage_metadata (
|
|
key TEXT PRIMARY KEY NOT NULL,
|
|
value TEXT NOT NULL
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS usage_events (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
event_key TEXT NOT NULL UNIQUE,
|
|
timestamp TEXT NOT NULL,
|
|
timestamp_ms INTEGER NOT NULL,
|
|
local_hour TEXT NOT NULL,
|
|
latency_ms INTEGER NOT NULL DEFAULT 0,
|
|
ttft_ms INTEGER,
|
|
source TEXT NOT NULL DEFAULT '',
|
|
auth_index TEXT NOT NULL DEFAULT '',
|
|
failed INTEGER NOT NULL DEFAULT 0,
|
|
provider TEXT NOT NULL DEFAULT '',
|
|
model TEXT NOT NULL DEFAULT '',
|
|
alias TEXT NOT NULL DEFAULT '',
|
|
reasoning_effort TEXT NOT NULL DEFAULT '',
|
|
service_tier TEXT NOT NULL DEFAULT '',
|
|
response_service_tier TEXT NOT NULL DEFAULT '',
|
|
executor_type TEXT NOT NULL DEFAULT '',
|
|
endpoint TEXT NOT NULL DEFAULT '',
|
|
auth_type TEXT NOT NULL DEFAULT '',
|
|
api_key_hash TEXT NOT NULL DEFAULT '',
|
|
api_key_display TEXT NOT NULL DEFAULT '',
|
|
api_key_remark TEXT NOT NULL DEFAULT '',
|
|
request_id TEXT NOT NULL DEFAULT '',
|
|
input_tokens INTEGER NOT NULL DEFAULT 0,
|
|
output_tokens INTEGER NOT NULL DEFAULT 0,
|
|
reasoning_tokens INTEGER NOT NULL DEFAULT 0,
|
|
cache_read_tokens INTEGER NOT NULL DEFAULT 0,
|
|
cache_creation_tokens INTEGER NOT NULL DEFAULT 0,
|
|
total_tokens INTEGER NOT NULL DEFAULT 0,
|
|
created_at TEXT NOT NULL
|
|
);
|
|
|
|
CREATE INDEX IF NOT EXISTS idx_usage_events_timestamp
|
|
ON usage_events(timestamp_ms DESC, id DESC);
|
|
CREATE INDEX IF NOT EXISTS idx_usage_events_local_hour
|
|
ON usage_events(local_hour, timestamp_ms DESC);
|
|
CREATE INDEX IF NOT EXISTS idx_usage_events_model_timestamp
|
|
ON usage_events(model, timestamp_ms DESC);
|
|
CREATE INDEX IF NOT EXISTS idx_usage_events_provider_timestamp
|
|
ON usage_events(provider, timestamp_ms DESC);
|
|
CREATE INDEX IF NOT EXISTS idx_usage_events_source_timestamp
|
|
ON usage_events(source, timestamp_ms DESC);
|
|
CREATE INDEX IF NOT EXISTS idx_usage_events_api_key_timestamp
|
|
ON usage_events(api_key_hash, timestamp_ms DESC);
|
|
CREATE INDEX IF NOT EXISTS idx_usage_events_failed_timestamp
|
|
ON usage_events(failed, timestamp_ms DESC);
|
|
|
|
PRAGMA user_version = 1;
|
|
"#,
|
|
)
|
|
.map_err(|error| format!("初始化 SQLite 使用记录结构失败: {error}"))
|
|
}
|
|
|
|
fn open_usage_database() -> Result<Connection, String> {
|
|
open_usage_database_at(&usage_root_dir()?)
|
|
}
|
|
|
|
fn open_usage_database_at(root: &Path) -> Result<Connection, String> {
|
|
fs::create_dir_all(root).map_err(|error| format!("创建使用记录目录失败: {error}"))?;
|
|
let path = root.join(USAGE_DATABASE_FILE);
|
|
let connection = Connection::open(&path)
|
|
.map_err(|error| format!("打开 SQLite 使用记录数据库失败 {}: {error}", path.display()))?;
|
|
connection
|
|
.busy_timeout(Duration::from_secs(SQLITE_BUSY_TIMEOUT_SECONDS))
|
|
.map_err(|error| format!("设置 SQLite busy timeout 失败: {error}"))?;
|
|
connection
|
|
.pragma_update(None, "foreign_keys", "ON")
|
|
.map_err(|error| format!("启用 SQLite foreign keys 失败: {error}"))?;
|
|
connection
|
|
.pragma_update(None, "synchronous", "NORMAL")
|
|
.map_err(|error| format!("设置 SQLite synchronous 模式失败: {error}"))?;
|
|
Ok(connection)
|
|
}
|
|
|
|
fn migrate_legacy_json_storage(connection: &mut Connection, root: &Path) -> Result<(), String> {
|
|
let migrated = connection
|
|
.query_row(
|
|
"SELECT value FROM usage_metadata WHERE key = ?1",
|
|
params![LEGACY_JSON_MIGRATION_KEY],
|
|
|row| row.get::<_, String>(0),
|
|
)
|
|
.optional()
|
|
.map_err(|error| format!("读取旧使用记录迁移状态失败: {error}"))?
|
|
.is_some();
|
|
if migrated {
|
|
return Ok(());
|
|
}
|
|
|
|
let transaction = connection
|
|
.transaction()
|
|
.map_err(|error| format!("开始旧使用记录迁移事务失败: {error}"))?;
|
|
let mut migrated_records = 0_usize;
|
|
|
|
for path in sorted_json_files(&root.join(LEGACY_USAGE_EVENTS_DIR))? {
|
|
let content = fs::read_to_string(&path)
|
|
.map_err(|error| format!("读取旧使用记录失败 {}: {error}", path.display()))?;
|
|
let file = serde_json::from_str::<LegacyUsageHourFile>(&content)
|
|
.map_err(|error| format!("解析旧使用记录失败 {}: {error}", path.display()))?;
|
|
validate_legacy_schema(file.schema_version, &path)?;
|
|
migrated_records = migrated_records.saturating_add(insert_usage_records_in_transaction(
|
|
&transaction,
|
|
&file.records,
|
|
)?);
|
|
}
|
|
|
|
for path in sorted_json_files(&root.join(LEGACY_USAGE_INBOX_DIR))? {
|
|
let content = fs::read_to_string(&path)
|
|
.map_err(|error| format!("读取旧使用记录收件箱失败 {}: {error}", path.display()))?;
|
|
let file = serde_json::from_str::<LegacyUsageInboxFile>(&content)
|
|
.map_err(|error| format!("解析旧使用记录收件箱失败 {}: {error}", path.display()))?;
|
|
validate_legacy_schema(file.schema_version, &path)?;
|
|
migrated_records = migrated_records.saturating_add(insert_usage_records_in_transaction(
|
|
&transaction,
|
|
&file.records,
|
|
)?);
|
|
}
|
|
|
|
transaction
|
|
.execute(
|
|
"INSERT INTO usage_metadata (key, value) VALUES (?1, ?2)",
|
|
params![LEGACY_JSON_MIGRATION_KEY, migrated_records.to_string()],
|
|
)
|
|
.map_err(|error| format!("记录旧使用记录迁移状态失败: {error}"))?;
|
|
transaction
|
|
.commit()
|
|
.map_err(|error| format!("提交旧使用记录迁移失败: {error}"))?;
|
|
Ok(())
|
|
}
|
|
|
|
fn sorted_json_files(directory: &Path) -> Result<Vec<PathBuf>, String> {
|
|
if !directory.is_dir() {
|
|
return Ok(Vec::new());
|
|
}
|
|
let mut paths = fs::read_dir(directory)
|
|
.map_err(|error| format!("读取旧使用记录目录失败 {}: {error}", directory.display()))?
|
|
.filter_map(Result::ok)
|
|
.map(|entry| entry.path())
|
|
.filter(|path| path.extension().and_then(|value| value.to_str()) == Some("json"))
|
|
.collect::<Vec<_>>();
|
|
paths.sort();
|
|
Ok(paths)
|
|
}
|
|
|
|
fn validate_legacy_schema(version: u8, path: &Path) -> Result<(), String> {
|
|
if version == USAGE_SCHEMA_VERSION {
|
|
Ok(())
|
|
} else {
|
|
Err(format!(
|
|
"不支持的旧使用记录版本 {version}: {}",
|
|
path.display()
|
|
))
|
|
}
|
|
}
|
|
|
|
pub(crate) fn start_usage_collector(app: tauri::AppHandle) {
|
|
let state = app.state::<UsageCollectorState>();
|
|
if let Ok(total_records) = total_usage_records() {
|
|
state.set_total_records(total_records);
|
|
}
|
|
let Some(token) = state.start() else {
|
|
return;
|
|
};
|
|
tauri::async_runtime::spawn(async move {
|
|
usage_collector_loop(app, token).await;
|
|
});
|
|
}
|
|
|
|
pub(crate) fn stop_usage_collector(app: &tauri::AppHandle) {
|
|
app.state::<UsageCollectorState>().stop();
|
|
}
|
|
|
|
async fn usage_collector_loop(app: tauri::AppHandle, token: CancellationToken) {
|
|
let root = match usage_root_dir() {
|
|
Ok(root) => root,
|
|
Err(error) => {
|
|
set_collector_error(&app, error);
|
|
return;
|
|
}
|
|
};
|
|
|
|
let mut retry_seconds = 1_u64;
|
|
loop {
|
|
if token.is_cancelled() {
|
|
return;
|
|
}
|
|
let config = match app.state::<GuiConfigState>().snapshot() {
|
|
Ok(config) => config,
|
|
Err(error) => {
|
|
set_collector_error(&app, error);
|
|
wait_or_cancel(&token, retry_seconds).await;
|
|
retry_seconds = (retry_seconds * 2).min(10);
|
|
continue;
|
|
}
|
|
};
|
|
let process_state = app.state::<CoreProcessState>();
|
|
let core_running = current_core_status(Some(process_state.inner()), Some(config.port))
|
|
.map(|status| status.running)
|
|
.unwrap_or(false);
|
|
if !core_running {
|
|
set_collector_status(&app, "waiting-core", "等待内核启动", None);
|
|
retry_seconds = 1;
|
|
wait_or_cancel(&token, 1).await;
|
|
continue;
|
|
}
|
|
|
|
match fetch_usage_queue(&config).await {
|
|
Ok(items) if items.is_empty() => {
|
|
set_collector_status(&app, "collecting", "使用记录采集中", None);
|
|
retry_seconds = 1;
|
|
wait_or_cancel(&token, 1).await;
|
|
}
|
|
Ok(items) => match persist_queue_items(&root, items, &config) {
|
|
Ok(saved) => {
|
|
let collected_at = Local::now().to_rfc3339();
|
|
if saved > 0 {
|
|
app.state::<UsageCollectorState>()
|
|
.increment_total_records(saved);
|
|
}
|
|
set_collector_status(
|
|
&app,
|
|
"collecting",
|
|
&format!("已保存 {saved} 条新记录"),
|
|
Some(collected_at.clone()),
|
|
);
|
|
if saved > 0 {
|
|
let _ = app.emit(USAGE_UPDATED_EVENT, collected_at);
|
|
}
|
|
retry_seconds = 1;
|
|
}
|
|
Err(error) => {
|
|
set_collector_error(&app, error);
|
|
wait_or_cancel(&token, retry_seconds).await;
|
|
retry_seconds = (retry_seconds * 2).min(10);
|
|
}
|
|
},
|
|
Err(error) => {
|
|
set_collector_error(&app, error);
|
|
wait_or_cancel(&token, retry_seconds).await;
|
|
retry_seconds = (retry_seconds * 2).min(10);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
async fn wait_or_cancel(token: &CancellationToken, seconds: u64) {
|
|
tokio::select! {
|
|
_ = token.cancelled() => {},
|
|
_ = tokio::time::sleep(Duration::from_secs(seconds)) => {},
|
|
}
|
|
}
|
|
|
|
async fn fetch_usage_queue(config: &GuiConfigFile) -> Result<Vec<Value>, String> {
|
|
let client = management_http_client()?;
|
|
let response = client
|
|
.get(management_endpoint(config, "usage-queue")?)
|
|
.header("Authorization", management_authorization(config)?)
|
|
.query(&[("count", USAGE_QUEUE_BATCH_SIZE)])
|
|
.send()
|
|
.await
|
|
.map_err(|error| format!("读取 CPA 使用记录队列失败: {error}"))?;
|
|
let status = response.status();
|
|
let text = response
|
|
.text()
|
|
.await
|
|
.map_err(|error| format!("读取 CPA 使用记录响应失败: {error}"))?;
|
|
if !status.is_success() {
|
|
return Err(format!(
|
|
"CPA 使用记录队列返回 HTTP {}: {}",
|
|
status.as_u16(),
|
|
text.trim()
|
|
));
|
|
}
|
|
serde_json::from_str::<Vec<Value>>(&text)
|
|
.map_err(|error| format!("解析 CPA 使用记录失败: {error}"))
|
|
}
|
|
|
|
fn set_collector_error(app: &tauri::AppHandle, error: String) {
|
|
set_collector_status(app, "error", &error, None);
|
|
}
|
|
|
|
fn set_collector_status(
|
|
app: &tauri::AppHandle,
|
|
state_name: &str,
|
|
message: &str,
|
|
last_collected_at: Option<String>,
|
|
) {
|
|
let state = app.state::<UsageCollectorState>();
|
|
let previous = state.status().ok();
|
|
let total_records = previous
|
|
.as_ref()
|
|
.map(|value| value.total_records)
|
|
.unwrap_or(0);
|
|
state.set_status(UsageCollectorStatus {
|
|
state: state_name.to_string(),
|
|
message: message.to_string(),
|
|
last_collected_at: last_collected_at
|
|
.or_else(|| previous.and_then(|value| value.last_collected_at)),
|
|
total_records,
|
|
});
|
|
}
|
|
|
|
fn persist_queue_items(
|
|
root: &Path,
|
|
items: Vec<Value>,
|
|
config: &GuiConfigFile,
|
|
) -> Result<usize, String> {
|
|
let records = items
|
|
.into_iter()
|
|
.filter(Value::is_object)
|
|
.filter(is_generated_usage_event)
|
|
.map(|item| normalize_usage_record(item, config))
|
|
.collect::<Result<Vec<_>, _>>()?;
|
|
if records.is_empty() {
|
|
return Ok(0);
|
|
}
|
|
let mut connection = open_usage_database_at(root)?;
|
|
insert_usage_records(&mut connection, &records)
|
|
}
|
|
|
|
fn insert_usage_records(
|
|
connection: &mut Connection,
|
|
records: &[UsageRecord],
|
|
) -> Result<usize, String> {
|
|
let transaction = connection
|
|
.transaction()
|
|
.map_err(|error| format!("开始 SQLite 使用记录事务失败: {error}"))?;
|
|
let inserted = insert_usage_records_in_transaction(&transaction, records)?;
|
|
transaction
|
|
.commit()
|
|
.map_err(|error| format!("提交 SQLite 使用记录失败: {error}"))?;
|
|
Ok(inserted)
|
|
}
|
|
|
|
fn insert_usage_records_in_transaction(
|
|
transaction: &Transaction<'_>,
|
|
records: &[UsageRecord],
|
|
) -> Result<usize, String> {
|
|
if records.is_empty() {
|
|
return Ok(0);
|
|
}
|
|
let mut statement = transaction
|
|
.prepare(
|
|
r#"
|
|
INSERT OR IGNORE INTO usage_events (
|
|
event_key, timestamp, timestamp_ms, local_hour, latency_ms, ttft_ms,
|
|
source, auth_index, failed, provider, model, alias, reasoning_effort,
|
|
service_tier, response_service_tier, executor_type, endpoint, auth_type,
|
|
api_key_hash, api_key_display, api_key_remark, request_id,
|
|
input_tokens, output_tokens, reasoning_tokens, cache_read_tokens,
|
|
cache_creation_tokens, total_tokens, created_at
|
|
) 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, ?25, ?26, ?27, ?28, ?29
|
|
)
|
|
"#,
|
|
)
|
|
.map_err(|error| format!("准备 SQLite 使用记录写入失败: {error}"))?;
|
|
let created_at = Local::now().to_rfc3339();
|
|
let mut inserted = 0_usize;
|
|
for record in records {
|
|
inserted = inserted.saturating_add(
|
|
statement
|
|
.execute(params![
|
|
record.id,
|
|
record.timestamp,
|
|
record_timestamp_millis(record),
|
|
record_local_hour(record),
|
|
to_sql_i64(record.latency_ms),
|
|
record.ttft_ms.map(to_sql_i64),
|
|
record.source,
|
|
record.auth_index,
|
|
record.failed,
|
|
record.provider,
|
|
record.model,
|
|
record.alias,
|
|
record.reasoning_effort,
|
|
record.service_tier,
|
|
record.response_service_tier,
|
|
record.executor_type,
|
|
record.endpoint,
|
|
record.auth_type,
|
|
record.api_key_hash,
|
|
record.api_key_display,
|
|
record.api_key_remark,
|
|
record.request_id,
|
|
to_sql_i64(record.tokens.input_tokens),
|
|
to_sql_i64(record.tokens.output_tokens),
|
|
to_sql_i64(record.tokens.reasoning_tokens),
|
|
to_sql_i64(record.tokens.cache_read_tokens),
|
|
to_sql_i64(record.tokens.cache_creation_tokens),
|
|
to_sql_i64(record.tokens.total_tokens),
|
|
created_at,
|
|
])
|
|
.map_err(|error| format!("写入 SQLite 使用记录失败: {error}"))?,
|
|
);
|
|
}
|
|
Ok(inserted)
|
|
}
|
|
|
|
fn is_generated_usage_event(value: &Value) -> bool {
|
|
if value.get("generate").and_then(Value::as_bool) == Some(false) {
|
|
return false;
|
|
}
|
|
let legacy_websocket_prewarm = value
|
|
.get("executor_type")
|
|
.and_then(Value::as_str)
|
|
.is_some_and(|value| value.trim() == "CodexWebsocketsExecutor")
|
|
&& value
|
|
.get("tokens")
|
|
.and_then(Value::as_object)
|
|
.is_some_and(|tokens| {
|
|
[
|
|
"input_tokens",
|
|
"output_tokens",
|
|
"reasoning_tokens",
|
|
"cached_tokens",
|
|
"cache_read_tokens",
|
|
"cache_creation_tokens",
|
|
"total_tokens",
|
|
]
|
|
.into_iter()
|
|
.all(|key| u64_field(tokens, key) == 0)
|
|
});
|
|
!legacy_websocket_prewarm
|
|
}
|
|
|
|
fn normalize_usage_record(value: Value, config: &GuiConfigFile) -> Result<UsageRecord, String> {
|
|
let object = value
|
|
.as_object()
|
|
.ok_or_else(|| "CPA 使用记录必须是 JSON 对象".to_string())?;
|
|
let timestamp = string_field(object, "timestamp")
|
|
.filter(|value| DateTime::parse_from_rfc3339(value).is_ok())
|
|
.unwrap_or_else(|| Local::now().to_rfc3339());
|
|
let request_id = string_field(object, "request_id").unwrap_or_default();
|
|
let api_key = string_field(object, "api_key").unwrap_or_default();
|
|
let api_key_hash = hash_text(&api_key);
|
|
let api_key_remark = config
|
|
.api_keys
|
|
.iter()
|
|
.find(|entry| entry.key == api_key)
|
|
.map(|entry| entry.remark.clone())
|
|
.unwrap_or_default();
|
|
let tokens_object = object.get("tokens").and_then(Value::as_object);
|
|
let mut tokens = UsageTokenStats {
|
|
input_tokens: token_u64(tokens_object, "input_tokens"),
|
|
output_tokens: token_u64(tokens_object, "output_tokens"),
|
|
reasoning_tokens: token_u64(tokens_object, "reasoning_tokens"),
|
|
cache_read_tokens: token_u64(tokens_object, "cache_read_tokens")
|
|
.max(token_u64(tokens_object, "cached_tokens")),
|
|
cache_creation_tokens: token_u64(tokens_object, "cache_creation_tokens"),
|
|
total_tokens: token_u64(tokens_object, "total_tokens"),
|
|
};
|
|
if tokens.total_tokens == 0 {
|
|
tokens.total_tokens = tokens.input_tokens.saturating_add(tokens.output_tokens);
|
|
}
|
|
let mut canonical = object.clone();
|
|
canonical.remove("response_headers");
|
|
canonical.remove("api_key");
|
|
let id = if request_id.is_empty() {
|
|
hash_text(&serde_json::to_string(&canonical).unwrap_or_default())
|
|
} else {
|
|
request_id.clone()
|
|
};
|
|
Ok(UsageRecord {
|
|
id,
|
|
timestamp,
|
|
latency_ms: u64_field(object, "latency_ms"),
|
|
ttft_ms: optional_u64_field(object, "ttft_ms"),
|
|
source: string_field(object, "source").unwrap_or_default(),
|
|
auth_index: string_field(object, "auth_index").unwrap_or_default(),
|
|
failed: object
|
|
.get("failed")
|
|
.and_then(Value::as_bool)
|
|
.unwrap_or(false),
|
|
provider: string_field(object, "provider").unwrap_or_default(),
|
|
model: string_field(object, "model").unwrap_or_else(|| "unknown".to_string()),
|
|
alias: string_field(object, "alias").unwrap_or_default(),
|
|
reasoning_effort: string_field(object, "reasoning_effort").unwrap_or_default(),
|
|
service_tier: string_field(object, "service_tier").unwrap_or_default(),
|
|
response_service_tier: string_field(object, "response_service_tier").unwrap_or_default(),
|
|
executor_type: string_field(object, "executor_type").unwrap_or_default(),
|
|
endpoint: string_field(object, "endpoint").unwrap_or_default(),
|
|
auth_type: string_field(object, "auth_type").unwrap_or_default(),
|
|
api_key_hash,
|
|
api_key_display: mask_api_key(&api_key),
|
|
api_key_remark,
|
|
request_id,
|
|
tokens,
|
|
})
|
|
}
|
|
|
|
fn build_usage_filter(query: &UsageQuery) -> UsageSqlFilter {
|
|
let mut clauses = Vec::<String>::new();
|
|
let mut params = Vec::<SqlValue>::new();
|
|
if let Some(start) = query.start.as_deref().and_then(parse_timestamp_millis) {
|
|
clauses.push("timestamp_ms >= ?".to_string());
|
|
params.push(SqlValue::Integer(start));
|
|
}
|
|
if let Some(end) = query.end.as_deref().and_then(parse_timestamp_millis) {
|
|
clauses.push("timestamp_ms <= ?".to_string());
|
|
params.push(SqlValue::Integer(end));
|
|
}
|
|
add_text_filter(&mut clauses, &mut params, "model", query.model.as_deref());
|
|
add_text_filter(
|
|
&mut clauses,
|
|
&mut params,
|
|
"provider",
|
|
query.provider.as_deref(),
|
|
);
|
|
add_text_filter(&mut clauses, &mut params, "source", query.source.as_deref());
|
|
add_text_filter(
|
|
&mut clauses,
|
|
&mut params,
|
|
"api_key_hash",
|
|
query.api_key_hash.as_deref(),
|
|
);
|
|
if let Some(failed) = query.failed {
|
|
clauses.push("failed = ?".to_string());
|
|
params.push(SqlValue::Integer(i64::from(failed)));
|
|
}
|
|
UsageSqlFilter {
|
|
clause: if clauses.is_empty() {
|
|
String::new()
|
|
} else {
|
|
format!(" WHERE {}", clauses.join(" AND "))
|
|
},
|
|
params,
|
|
}
|
|
}
|
|
|
|
fn add_text_filter(
|
|
clauses: &mut Vec<String>,
|
|
params: &mut Vec<SqlValue>,
|
|
column: &str,
|
|
value: Option<&str>,
|
|
) {
|
|
if let Some(value) = value.map(str::trim).filter(|value| !value.is_empty()) {
|
|
clauses.push(format!("{column} = ? COLLATE NOCASE"));
|
|
params.push(SqlValue::Text(value.to_string()));
|
|
}
|
|
}
|
|
|
|
#[tauri::command]
|
|
pub(crate) fn get_usage_collector_status(
|
|
state: tauri::State<'_, UsageCollectorState>,
|
|
) -> Result<UsageCollectorStatus, String> {
|
|
state.status()
|
|
}
|
|
|
|
#[tauri::command]
|
|
pub(crate) fn get_usage_overview(query: UsageQuery) -> Result<UsageOverview, String> {
|
|
let connection = open_usage_database()?;
|
|
load_usage_overview(&connection, &query)
|
|
}
|
|
|
|
fn load_usage_overview(
|
|
connection: &Connection,
|
|
query: &UsageQuery,
|
|
) -> Result<UsageOverview, String> {
|
|
let filter = build_usage_filter(query);
|
|
let summary_sql = format!(
|
|
r#"
|
|
SELECT
|
|
COUNT(*),
|
|
COALESCE(SUM(CASE WHEN failed = 0 THEN 1 ELSE 0 END), 0),
|
|
COALESCE(SUM(CASE WHEN failed != 0 THEN 1 ELSE 0 END), 0),
|
|
COALESCE(SUM(input_tokens), 0),
|
|
COALESCE(SUM(output_tokens), 0),
|
|
COALESCE(SUM(reasoning_tokens), 0),
|
|
COALESCE(SUM(cache_read_tokens), 0),
|
|
COALESCE(SUM(cache_creation_tokens), 0),
|
|
COALESCE(SUM(total_tokens), 0),
|
|
COALESCE(SUM(latency_ms), 0),
|
|
MIN(timestamp_ms),
|
|
MAX(timestamp_ms)
|
|
FROM usage_events{}
|
|
"#,
|
|
filter.clause
|
|
);
|
|
let summary = connection
|
|
.query_row(
|
|
&summary_sql,
|
|
params_from_iter(filter.params.iter()),
|
|
|row| {
|
|
Ok((
|
|
row.get::<_, i64>(0)?,
|
|
row.get::<_, i64>(1)?,
|
|
row.get::<_, i64>(2)?,
|
|
row.get::<_, i64>(3)?,
|
|
row.get::<_, i64>(4)?,
|
|
row.get::<_, i64>(5)?,
|
|
row.get::<_, i64>(6)?,
|
|
row.get::<_, i64>(7)?,
|
|
row.get::<_, i64>(8)?,
|
|
row.get::<_, i64>(9)?,
|
|
row.get::<_, Option<i64>>(10)?,
|
|
row.get::<_, Option<i64>>(11)?,
|
|
))
|
|
},
|
|
)
|
|
.map_err(|error| format!("统计 SQLite 使用记录失败: {error}"))?;
|
|
|
|
let timeline_sql = format!(
|
|
r#"
|
|
SELECT
|
|
local_hour,
|
|
COUNT(*),
|
|
COALESCE(SUM(CASE WHEN failed = 0 THEN 1 ELSE 0 END), 0),
|
|
COALESCE(SUM(CASE WHEN failed != 0 THEN 1 ELSE 0 END), 0),
|
|
COALESCE(SUM(total_tokens), 0)
|
|
FROM usage_events{}
|
|
GROUP BY local_hour
|
|
ORDER BY local_hour ASC
|
|
"#,
|
|
filter.clause
|
|
);
|
|
let mut statement = connection
|
|
.prepare(&timeline_sql)
|
|
.map_err(|error| format!("准备 SQLite 使用趋势查询失败: {error}"))?;
|
|
let timeline = statement
|
|
.query_map(params_from_iter(filter.params.iter()), |row| {
|
|
Ok(UsageTimelinePoint {
|
|
hour: row.get(0)?,
|
|
requests: from_sql_i64(row.get(1)?),
|
|
success: from_sql_i64(row.get(2)?),
|
|
failure: from_sql_i64(row.get(3)?),
|
|
tokens: from_sql_i64(row.get(4)?),
|
|
})
|
|
})
|
|
.map_err(|error| format!("查询 SQLite 使用趋势失败: {error}"))?
|
|
.collect::<Result<Vec<_>, _>>()
|
|
.map_err(|error| format!("读取 SQLite 使用趋势失败: {error}"))?;
|
|
|
|
let mut overview = UsageOverview {
|
|
total_requests: from_sql_i64(summary.0),
|
|
success_count: from_sql_i64(summary.1),
|
|
failure_count: from_sql_i64(summary.2),
|
|
input_tokens: from_sql_i64(summary.3),
|
|
output_tokens: from_sql_i64(summary.4),
|
|
reasoning_tokens: from_sql_i64(summary.5),
|
|
cache_read_tokens: from_sql_i64(summary.6),
|
|
cache_creation_tokens: from_sql_i64(summary.7),
|
|
total_tokens: from_sql_i64(summary.8),
|
|
timeline,
|
|
..UsageOverview::default()
|
|
};
|
|
if overview.total_requests > 0 {
|
|
overview.success_rate =
|
|
overview.success_count as f64 * 100.0 / overview.total_requests as f64;
|
|
overview.average_latency_ms =
|
|
from_sql_i64(summary.9) as f64 / overview.total_requests as f64;
|
|
let minutes = query_window_minutes(query, summary.10, summary.11);
|
|
overview.rpm = overview.total_requests as f64 / minutes;
|
|
overview.tpm = overview.total_tokens as f64 / minutes;
|
|
}
|
|
Ok(overview)
|
|
}
|
|
|
|
#[tauri::command]
|
|
pub(crate) fn get_usage_analysis(query: UsageQuery) -> Result<UsageAnalysis, String> {
|
|
let connection = open_usage_database()?;
|
|
load_usage_analysis(&connection, &query)
|
|
}
|
|
|
|
fn load_usage_analysis(
|
|
connection: &Connection,
|
|
query: &UsageQuery,
|
|
) -> Result<UsageAnalysis, String> {
|
|
Ok(UsageAnalysis {
|
|
models: load_simple_categories(connection, query, "model", "unknown")?,
|
|
providers: load_simple_categories(connection, query, "provider", "未知 Provider")?,
|
|
sources: load_simple_categories(connection, query, "source", "未知来源")?,
|
|
api_keys: load_api_key_categories(connection, query)?,
|
|
})
|
|
}
|
|
|
|
fn load_simple_categories(
|
|
connection: &Connection,
|
|
query: &UsageQuery,
|
|
column: &str,
|
|
fallback: &str,
|
|
) -> Result<Vec<UsageCategory>, String> {
|
|
let filter = build_usage_filter(query);
|
|
let sql = format!(
|
|
r#"
|
|
SELECT
|
|
COALESCE(NULLIF(TRIM({column}), ''), ?),
|
|
COUNT(*),
|
|
COALESCE(SUM(CASE WHEN failed != 0 THEN 1 ELSE 0 END), 0),
|
|
COALESCE(SUM(total_tokens), 0)
|
|
FROM usage_events{}
|
|
GROUP BY 1
|
|
ORDER BY 4 DESC, 2 DESC
|
|
"#,
|
|
filter.clause
|
|
);
|
|
let mut values = Vec::with_capacity(filter.params.len() + 1);
|
|
values.push(SqlValue::Text(fallback.to_string()));
|
|
values.extend(filter.params);
|
|
let mut statement = connection
|
|
.prepare(&sql)
|
|
.map_err(|error| format!("准备 SQLite 使用分析查询失败: {error}"))?;
|
|
let categories = statement
|
|
.query_map(params_from_iter(values.iter()), |row| {
|
|
let key = row.get::<_, String>(0)?;
|
|
Ok(UsageCategory {
|
|
label: key.clone(),
|
|
key,
|
|
requests: from_sql_i64(row.get(1)?),
|
|
failures: from_sql_i64(row.get(2)?),
|
|
tokens: from_sql_i64(row.get(3)?),
|
|
})
|
|
})
|
|
.map_err(|error| format!("查询 SQLite 使用分析失败: {error}"))?
|
|
.collect::<Result<Vec<_>, _>>()
|
|
.map_err(|error| format!("读取 SQLite 使用分析失败: {error}"))?;
|
|
Ok(categories)
|
|
}
|
|
|
|
fn load_api_key_categories(
|
|
connection: &Connection,
|
|
query: &UsageQuery,
|
|
) -> Result<Vec<UsageCategory>, String> {
|
|
let filter = build_usage_filter(query);
|
|
let sql = format!(
|
|
r#"
|
|
SELECT
|
|
COALESCE(NULLIF(TRIM(api_key_hash), ''), '未记录密钥'),
|
|
MAX(TRIM(api_key_remark)),
|
|
MAX(TRIM(api_key_display)),
|
|
COUNT(*),
|
|
COALESCE(SUM(CASE WHEN failed != 0 THEN 1 ELSE 0 END), 0),
|
|
COALESCE(SUM(total_tokens), 0)
|
|
FROM usage_events{}
|
|
GROUP BY 1
|
|
ORDER BY 6 DESC, 4 DESC
|
|
"#,
|
|
filter.clause
|
|
);
|
|
let mut statement = connection
|
|
.prepare(&sql)
|
|
.map_err(|error| format!("准备 SQLite API Key 使用分析查询失败: {error}"))?;
|
|
let categories = statement
|
|
.query_map(params_from_iter(filter.params.iter()), |row| {
|
|
let key = row.get::<_, String>(0)?;
|
|
let remark = row.get::<_, String>(1)?;
|
|
let display = row.get::<_, String>(2)?;
|
|
let label = if !remark.is_empty() {
|
|
if display.is_empty() {
|
|
remark
|
|
} else {
|
|
format!("{remark} · {display}")
|
|
}
|
|
} else if !display.is_empty() {
|
|
display
|
|
} else {
|
|
"未记录密钥".to_string()
|
|
};
|
|
Ok(UsageCategory {
|
|
key,
|
|
label,
|
|
requests: from_sql_i64(row.get(3)?),
|
|
failures: from_sql_i64(row.get(4)?),
|
|
tokens: from_sql_i64(row.get(5)?),
|
|
})
|
|
})
|
|
.map_err(|error| format!("查询 SQLite API Key 使用分析失败: {error}"))?
|
|
.collect::<Result<Vec<_>, _>>()
|
|
.map_err(|error| format!("读取 SQLite API Key 使用分析失败: {error}"))?;
|
|
Ok(categories)
|
|
}
|
|
|
|
#[tauri::command]
|
|
pub(crate) fn get_usage_events(query: UsageQuery) -> Result<UsageEventPage, String> {
|
|
let connection = open_usage_database()?;
|
|
load_usage_events(&connection, &query)
|
|
}
|
|
|
|
fn load_usage_events(
|
|
connection: &Connection,
|
|
query: &UsageQuery,
|
|
) -> Result<UsageEventPage, String> {
|
|
let filter = build_usage_filter(query);
|
|
let total_sql = format!("SELECT COUNT(*) FROM usage_events{}", filter.clause);
|
|
let total = connection
|
|
.query_row(&total_sql, params_from_iter(filter.params.iter()), |row| {
|
|
row.get::<_, i64>(0)
|
|
})
|
|
.map(from_sql_i64)
|
|
.map_err(|error| format!("统计 SQLite 使用事件失败: {error}"))?
|
|
.min(usize::MAX as u64) as usize;
|
|
let page_size = query.page_size.unwrap_or(50).clamp(20, 200);
|
|
let total_pages = total.div_ceil(page_size).max(1);
|
|
let page = query.page.unwrap_or(1).clamp(1, total_pages);
|
|
let offset = (page - 1).saturating_mul(page_size);
|
|
|
|
let sql = format!(
|
|
r#"
|
|
SELECT
|
|
event_key, timestamp, latency_ms, ttft_ms, source, auth_index, failed,
|
|
provider, model, alias, reasoning_effort, service_tier,
|
|
response_service_tier, executor_type, endpoint, auth_type,
|
|
api_key_hash, api_key_display, api_key_remark, request_id,
|
|
input_tokens, output_tokens, reasoning_tokens, cache_read_tokens,
|
|
cache_creation_tokens, total_tokens
|
|
FROM usage_events{}
|
|
ORDER BY timestamp_ms DESC, id DESC
|
|
LIMIT ? OFFSET ?
|
|
"#,
|
|
filter.clause
|
|
);
|
|
let mut values = filter.params;
|
|
values.push(SqlValue::Integer(page_size as i64));
|
|
values.push(SqlValue::Integer(offset.min(i64::MAX as usize) as i64));
|
|
let mut statement = connection
|
|
.prepare(&sql)
|
|
.map_err(|error| format!("准备 SQLite 使用事件查询失败: {error}"))?;
|
|
let items = statement
|
|
.query_map(params_from_iter(values.iter()), usage_record_from_row)
|
|
.map_err(|error| format!("查询 SQLite 使用事件失败: {error}"))?
|
|
.collect::<Result<Vec<_>, _>>()
|
|
.map_err(|error| format!("读取 SQLite 使用事件失败: {error}"))?;
|
|
Ok(UsageEventPage {
|
|
items,
|
|
total,
|
|
page,
|
|
page_size,
|
|
total_pages,
|
|
})
|
|
}
|
|
|
|
fn usage_record_from_row(row: &Row<'_>) -> rusqlite::Result<UsageRecord> {
|
|
Ok(UsageRecord {
|
|
id: row.get(0)?,
|
|
timestamp: row.get(1)?,
|
|
latency_ms: from_sql_i64(row.get(2)?),
|
|
ttft_ms: row.get::<_, Option<i64>>(3)?.map(from_sql_i64),
|
|
source: row.get(4)?,
|
|
auth_index: row.get(5)?,
|
|
failed: row.get::<_, i64>(6)? != 0,
|
|
provider: row.get(7)?,
|
|
model: row.get(8)?,
|
|
alias: row.get(9)?,
|
|
reasoning_effort: row.get(10)?,
|
|
service_tier: row.get(11)?,
|
|
response_service_tier: row.get(12)?,
|
|
executor_type: row.get(13)?,
|
|
endpoint: row.get(14)?,
|
|
auth_type: row.get(15)?,
|
|
api_key_hash: row.get(16)?,
|
|
api_key_display: row.get(17)?,
|
|
api_key_remark: row.get(18)?,
|
|
request_id: row.get(19)?,
|
|
tokens: UsageTokenStats {
|
|
input_tokens: from_sql_i64(row.get(20)?),
|
|
output_tokens: from_sql_i64(row.get(21)?),
|
|
reasoning_tokens: from_sql_i64(row.get(22)?),
|
|
cache_read_tokens: from_sql_i64(row.get(23)?),
|
|
cache_creation_tokens: from_sql_i64(row.get(24)?),
|
|
total_tokens: from_sql_i64(row.get(25)?),
|
|
},
|
|
})
|
|
}
|
|
|
|
fn total_usage_records() -> Result<u64, String> {
|
|
let connection = open_usage_database()?;
|
|
connection
|
|
.query_row("SELECT COUNT(*) FROM usage_events", [], |row| {
|
|
row.get::<_, i64>(0)
|
|
})
|
|
.map(from_sql_i64)
|
|
.map_err(|error| format!("统计 SQLite 使用记录总数失败: {error}"))
|
|
}
|
|
|
|
fn query_window_minutes(query: &UsageQuery, first: Option<i64>, last: Option<i64>) -> f64 {
|
|
let start = query
|
|
.start
|
|
.as_deref()
|
|
.and_then(parse_timestamp_millis)
|
|
.or(first)
|
|
.unwrap_or(0);
|
|
let end = query
|
|
.end
|
|
.as_deref()
|
|
.and_then(parse_timestamp_millis)
|
|
.or(last)
|
|
.unwrap_or(start);
|
|
((end.saturating_sub(start)) as f64 / 60_000.0).max(1.0)
|
|
}
|
|
|
|
fn usage_root_dir() -> Result<PathBuf, String> {
|
|
let executable =
|
|
std::env::current_exe().map_err(|error| format!("读取程序路径失败: {error}"))?;
|
|
let directory = executable
|
|
.parent()
|
|
.ok_or_else(|| "程序路径没有父目录".to_string())?;
|
|
Ok(directory.join(USAGE_DIR_NAME))
|
|
}
|
|
|
|
fn record_local_hour(record: &UsageRecord) -> String {
|
|
local_hour_from_timestamp(&record.timestamp)
|
|
.unwrap_or_else(|| Local::now().format("%Y-%m-%d-%H").to_string())
|
|
}
|
|
|
|
fn local_hour_from_timestamp(value: &str) -> Option<String> {
|
|
DateTime::parse_from_rfc3339(value).ok().map(|timestamp| {
|
|
timestamp
|
|
.with_timezone(&Local)
|
|
.format("%Y-%m-%d-%H")
|
|
.to_string()
|
|
})
|
|
}
|
|
|
|
fn record_timestamp_millis(record: &UsageRecord) -> i64 {
|
|
parse_timestamp_millis(&record.timestamp).unwrap_or(0)
|
|
}
|
|
|
|
fn parse_timestamp_millis(value: &str) -> Option<i64> {
|
|
DateTime::parse_from_rfc3339(value)
|
|
.ok()
|
|
.map(|timestamp| timestamp.timestamp_millis())
|
|
}
|
|
|
|
fn to_sql_i64(value: u64) -> i64 {
|
|
value.min(i64::MAX as u64) as i64
|
|
}
|
|
|
|
fn from_sql_i64(value: i64) -> u64 {
|
|
value.max(0) as u64
|
|
}
|
|
|
|
#[cfg(test)]
|
|
fn unique_file_stamp() -> u128 {
|
|
SystemTime::now()
|
|
.duration_since(UNIX_EPOCH)
|
|
.map(|duration| duration.as_nanos())
|
|
.unwrap_or(0)
|
|
}
|
|
|
|
fn string_field(object: &serde_json::Map<String, Value>, key: &str) -> Option<String> {
|
|
object
|
|
.get(key)
|
|
.and_then(|value| match value {
|
|
Value::String(value) => Some(value.trim().to_string()),
|
|
Value::Number(value) => Some(value.to_string()),
|
|
_ => None,
|
|
})
|
|
.filter(|value| !value.is_empty())
|
|
}
|
|
|
|
fn u64_field(object: &serde_json::Map<String, Value>, key: &str) -> u64 {
|
|
optional_u64_field(object, key).unwrap_or(0)
|
|
}
|
|
|
|
fn optional_u64_field(object: &serde_json::Map<String, Value>, key: &str) -> Option<u64> {
|
|
object.get(key).and_then(|value| {
|
|
value.as_u64().or_else(|| {
|
|
value
|
|
.as_i64()
|
|
.map(|number| number.max(0) as u64)
|
|
.or_else(|| value.as_str().and_then(|text| text.parse::<u64>().ok()))
|
|
})
|
|
})
|
|
}
|
|
|
|
fn token_u64(object: Option<&serde_json::Map<String, Value>>, key: &str) -> u64 {
|
|
object.map(|object| u64_field(object, key)).unwrap_or(0)
|
|
}
|
|
|
|
fn hash_text(value: &str) -> String {
|
|
if value.is_empty() {
|
|
return String::new();
|
|
}
|
|
format!("{:x}", Sha256::digest(value.as_bytes()))
|
|
}
|
|
|
|
fn mask_api_key(value: &str) -> String {
|
|
let value = value.trim();
|
|
if value.is_empty() {
|
|
return String::new();
|
|
}
|
|
if value.chars().count() <= 8 {
|
|
return format!("{}••••", value.chars().take(2).collect::<String>());
|
|
}
|
|
let start = value.chars().take(4).collect::<String>();
|
|
let end = value
|
|
.chars()
|
|
.rev()
|
|
.take(4)
|
|
.collect::<String>()
|
|
.chars()
|
|
.rev()
|
|
.collect::<String>();
|
|
format!("{start}••••{end}")
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn collector_total_records_are_cached_and_incremented() {
|
|
let state = UsageCollectorState::default();
|
|
state.set_total_records(40);
|
|
state.increment_total_records(2);
|
|
|
|
assert_eq!(state.status().unwrap().total_records, 42);
|
|
}
|
|
|
|
fn test_root(name: &str) -> PathBuf {
|
|
let root = std::env::temp_dir().join(format!(
|
|
"cpa-gui-usage-{name}-{}-{}",
|
|
std::process::id(),
|
|
unique_file_stamp()
|
|
));
|
|
fs::create_dir_all(&root).unwrap();
|
|
root
|
|
}
|
|
|
|
fn open_test_database(root: &Path) -> Connection {
|
|
initialize_usage_storage_at(root).unwrap();
|
|
open_usage_database_at(root).unwrap()
|
|
}
|
|
|
|
fn sample_record(id: &str, timestamp: &str, model: &str) -> UsageRecord {
|
|
UsageRecord {
|
|
id: id.to_string(),
|
|
timestamp: timestamp.to_string(),
|
|
latency_ms: 100,
|
|
ttft_ms: Some(20),
|
|
source: "source".to_string(),
|
|
auth_index: "auth".to_string(),
|
|
failed: false,
|
|
provider: "openai".to_string(),
|
|
model: model.to_string(),
|
|
alias: String::new(),
|
|
reasoning_effort: "high".to_string(),
|
|
service_tier: String::new(),
|
|
response_service_tier: String::new(),
|
|
executor_type: String::new(),
|
|
endpoint: "POST /v1/responses".to_string(),
|
|
auth_type: "oauth".to_string(),
|
|
api_key_hash: "hash".to_string(),
|
|
api_key_display: "12••••".to_string(),
|
|
api_key_remark: "内置密钥".to_string(),
|
|
request_id: id.to_string(),
|
|
tokens: UsageTokenStats {
|
|
input_tokens: 10,
|
|
output_tokens: 20,
|
|
reasoning_tokens: 5,
|
|
cache_read_tokens: 2,
|
|
cache_creation_tokens: 0,
|
|
total_tokens: 30,
|
|
},
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn masks_api_keys_without_exposing_the_full_value() {
|
|
assert_eq!(mask_api_key("123456"), "12••••");
|
|
assert_eq!(mask_api_key("sk-1234567890"), "sk-1••••7890");
|
|
}
|
|
|
|
#[test]
|
|
fn local_hour_key_uses_year_month_day_and_hour() {
|
|
let record = sample_record("id", "2026-07-17T20:30:00+08:00", "model");
|
|
assert_eq!(record_local_hour(&record).len(), "2026-07-17-20".len());
|
|
}
|
|
|
|
#[test]
|
|
fn normalizes_queue_records_without_persisting_secrets_or_headers() {
|
|
let config = GuiConfigFile {
|
|
locale: "zh-CN".to_string(),
|
|
port: 8317,
|
|
allow_lan: false,
|
|
run_on_startup: false,
|
|
auth_dir: String::new(),
|
|
api_keys: Vec::new(),
|
|
management_secret_key: "123456".to_string(),
|
|
plugins_enabled: false,
|
|
routing_strategy: "round-robin".to_string(),
|
|
};
|
|
let record = normalize_usage_record(
|
|
serde_json::json!({
|
|
"timestamp": "2026-07-17T20:30:00+08:00",
|
|
"request_id": "request-1",
|
|
"api_key": "secret-client-key",
|
|
"response_headers": { "authorization": ["secret-upstream-token"] },
|
|
"model": "gpt-test",
|
|
"tokens": { "input_tokens": 10, "output_tokens": 20 }
|
|
}),
|
|
&config,
|
|
)
|
|
.unwrap();
|
|
let rendered = serde_json::to_string(&record).unwrap();
|
|
|
|
assert!(!rendered.contains("secret-client-key"));
|
|
assert!(!rendered.contains("secret-upstream-token"));
|
|
assert_eq!(record.tokens.total_tokens, 30);
|
|
assert!(!record.api_key_hash.is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn sqlite_storage_uses_wal_and_reference_busy_timeout() {
|
|
let root = test_root("sqlite-pragmas");
|
|
let connection = open_test_database(&root);
|
|
let journal_mode = connection
|
|
.query_row("PRAGMA journal_mode", [], |row| row.get::<_, String>(0))
|
|
.unwrap();
|
|
let busy_timeout = connection
|
|
.query_row("PRAGMA busy_timeout", [], |row| row.get::<_, i64>(0))
|
|
.unwrap();
|
|
let foreign_keys = connection
|
|
.query_row("PRAGMA foreign_keys", [], |row| row.get::<_, i64>(0))
|
|
.unwrap();
|
|
|
|
assert_eq!(journal_mode.to_ascii_lowercase(), "wal");
|
|
assert_eq!(busy_timeout, 5_000);
|
|
assert_eq!(foreign_keys, 1);
|
|
assert!(root.join(USAGE_DATABASE_FILE).is_file());
|
|
fs::remove_dir_all(root).unwrap();
|
|
}
|
|
|
|
#[test]
|
|
fn sqlite_storage_deduplicates_event_keys_transactionally() {
|
|
let root = test_root("sqlite-deduplicate");
|
|
let mut connection = open_test_database(&root);
|
|
let record = sample_record("request-1", "2026-07-17T20:30:00+08:00", "gpt-test");
|
|
|
|
assert_eq!(
|
|
insert_usage_records(&mut connection, std::slice::from_ref(&record)).unwrap(),
|
|
1
|
|
);
|
|
assert_eq!(insert_usage_records(&mut connection, &[record]).unwrap(), 0);
|
|
let count = connection
|
|
.query_row("SELECT COUNT(*) FROM usage_events", [], |row| {
|
|
row.get::<_, i64>(0)
|
|
})
|
|
.unwrap();
|
|
|
|
assert_eq!(count, 1);
|
|
fs::remove_dir_all(root).unwrap();
|
|
}
|
|
|
|
#[test]
|
|
fn legacy_hour_and_inbox_json_are_migrated_once() {
|
|
let root = test_root("legacy-migration");
|
|
let events_dir = root.join(LEGACY_USAGE_EVENTS_DIR);
|
|
let inbox_dir = root.join(LEGACY_USAGE_INBOX_DIR);
|
|
fs::create_dir_all(&events_dir).unwrap();
|
|
fs::create_dir_all(&inbox_dir).unwrap();
|
|
let first = sample_record("request-1", "2026-07-17T20:30:00+08:00", "gpt-a");
|
|
let second = sample_record("request-2", "2026-07-17T20:31:00+08:00", "gpt-b");
|
|
fs::write(
|
|
events_dir.join("2026-07-17-20.json"),
|
|
serde_json::to_vec(&serde_json::json!({
|
|
"schemaVersion": USAGE_SCHEMA_VERSION,
|
|
"hour": "2026-07-17-20",
|
|
"timezone": "+08:00",
|
|
"records": [first]
|
|
}))
|
|
.unwrap(),
|
|
)
|
|
.unwrap();
|
|
fs::write(
|
|
inbox_dir.join("pending.json"),
|
|
serde_json::to_vec(&serde_json::json!({
|
|
"schemaVersion": USAGE_SCHEMA_VERSION,
|
|
"records": [second]
|
|
}))
|
|
.unwrap(),
|
|
)
|
|
.unwrap();
|
|
|
|
initialize_usage_storage_at(&root).unwrap();
|
|
initialize_usage_storage_at(&root).unwrap();
|
|
let connection = open_usage_database_at(&root).unwrap();
|
|
let count = connection
|
|
.query_row("SELECT COUNT(*) FROM usage_events", [], |row| {
|
|
row.get::<_, i64>(0)
|
|
})
|
|
.unwrap();
|
|
let marker = connection
|
|
.query_row(
|
|
"SELECT value FROM usage_metadata WHERE key = ?1",
|
|
params![LEGACY_JSON_MIGRATION_KEY],
|
|
|row| row.get::<_, String>(0),
|
|
)
|
|
.unwrap();
|
|
|
|
assert_eq!(count, 2);
|
|
assert_eq!(marker, "2");
|
|
fs::remove_dir_all(root).unwrap();
|
|
}
|
|
|
|
#[test]
|
|
fn sqlite_queries_filter_aggregate_and_paginate() {
|
|
let root = test_root("sqlite-query");
|
|
let mut connection = open_test_database(&root);
|
|
let success = sample_record("request-1", "2026-07-17T20:30:00+08:00", "gpt-a");
|
|
let mut failed = sample_record("request-2", "2026-07-17T21:30:00+08:00", "gpt-b");
|
|
failed.failed = true;
|
|
insert_usage_records(&mut connection, &[success, failed]).unwrap();
|
|
let query = UsageQuery {
|
|
model: Some("GPT-B".to_string()),
|
|
failed: Some(true),
|
|
page_size: Some(20),
|
|
..UsageQuery::default()
|
|
};
|
|
|
|
let overview = load_usage_overview(&connection, &query).unwrap();
|
|
let analysis = load_usage_analysis(&connection, &query).unwrap();
|
|
let events = load_usage_events(&connection, &query).unwrap();
|
|
|
|
assert_eq!(overview.total_requests, 1);
|
|
assert_eq!(overview.failure_count, 1);
|
|
assert_eq!(analysis.models[0].key, "gpt-b");
|
|
assert_eq!(events.total, 1);
|
|
assert_eq!(events.items[0].id, "request-2");
|
|
fs::remove_dir_all(root).unwrap();
|
|
}
|
|
}
|