forked from boos/BoosAPI
refactor: purge Node.js, Rust is the one true backend
- Moved boosapi-rs/* to repo root - Deleted all Node.js source (src/server/, src/web/, scripts/) - Deleted node_modules, package.json, tsconfig, vite.config - Updated .gitignore and AGENTS.md for Rust-only stack - Fixed static file path for frontend dist/web/ - Dockerfile stays the same (already Rust-only) - Kept dist/web/ for Rust static serving
This commit is contained in:
boos
@@ -0,0 +1,282 @@
|
||||
use crate::{
|
||||
error::{AppError, AppResult},
|
||||
services::shared::{enforce_key_request_limit, increment_key_usage},
|
||||
};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use sqlx::{FromRow, SqlitePool};
|
||||
use std::sync::Arc;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct UsageTracker {
|
||||
db: Arc<SqlitePool>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, FromRow, PartialEq)]
|
||||
pub struct UsageSummary {
|
||||
pub downstream_key_id: i64,
|
||||
pub request_count: i64,
|
||||
pub prompt_tokens: i64,
|
||||
pub completion_tokens: i64,
|
||||
pub total_tokens: i64,
|
||||
pub total_cost: f64,
|
||||
}
|
||||
|
||||
impl UsageTracker {
|
||||
pub fn new(db: Arc<SqlitePool>) -> Self {
|
||||
Self { db }
|
||||
}
|
||||
|
||||
pub async fn record_usage(
|
||||
&self,
|
||||
key_id: i64,
|
||||
model: impl Into<String>,
|
||||
prompt_tokens: i64,
|
||||
completion_tokens: i64,
|
||||
cost: f64,
|
||||
) -> AppResult<()> {
|
||||
Self::ensure_schema(&self.db).await?;
|
||||
let prompt_tokens = prompt_tokens.max(0);
|
||||
let completion_tokens = completion_tokens.max(0);
|
||||
let total_tokens = prompt_tokens + completion_tokens;
|
||||
|
||||
sqlx::query(
|
||||
"INSERT INTO usage_logs
|
||||
(downstream_key_id, model, prompt_tokens, completion_tokens, total_tokens, cost)
|
||||
VALUES (?, ?, ?, ?, ?, ?)",
|
||||
)
|
||||
.bind(key_id)
|
||||
.bind(model.into())
|
||||
.bind(prompt_tokens)
|
||||
.bind(completion_tokens)
|
||||
.bind(total_tokens)
|
||||
.bind(cost.max(0.0))
|
||||
.execute(&*self.db)
|
||||
.await?;
|
||||
|
||||
increment_key_usage(&self.db, key_id, cost.max(0.0)).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn get_key_usage(
|
||||
&self,
|
||||
key_id: i64,
|
||||
since: Option<&str>,
|
||||
until: Option<&str>,
|
||||
) -> AppResult<UsageSummary> {
|
||||
Self::ensure_schema(&self.db).await?;
|
||||
|
||||
let summary = sqlx::query_as::<_, UsageSummary>(
|
||||
"SELECT
|
||||
? AS downstream_key_id,
|
||||
COUNT(*) AS request_count,
|
||||
COALESCE(SUM(prompt_tokens), 0) AS prompt_tokens,
|
||||
COALESCE(SUM(completion_tokens), 0) AS completion_tokens,
|
||||
COALESCE(SUM(total_tokens), 0) AS total_tokens,
|
||||
COALESCE(SUM(cost), 0.0) AS total_cost
|
||||
FROM usage_logs
|
||||
WHERE downstream_key_id = ?
|
||||
AND (? IS NULL OR created_at >= ?)
|
||||
AND (? IS NULL OR created_at <= ?)",
|
||||
)
|
||||
.bind(key_id)
|
||||
.bind(key_id)
|
||||
.bind(since)
|
||||
.bind(since)
|
||||
.bind(until)
|
||||
.bind(until)
|
||||
.fetch_one(&*self.db)
|
||||
.await?;
|
||||
|
||||
Ok(summary)
|
||||
}
|
||||
|
||||
pub async fn check_rate_limit(&self, key_id: i64) -> AppResult<()> {
|
||||
enforce_key_request_limit(&self.db, key_id, 1).await?;
|
||||
|
||||
let row: Option<(Option<f64>, Option<f64>, Option<bool>)> = sqlx::query_as(
|
||||
"SELECT max_cost, used_cost, enabled
|
||||
FROM downstream_api_keys
|
||||
WHERE id = ?",
|
||||
)
|
||||
.bind(key_id)
|
||||
.fetch_optional(&*self.db)
|
||||
.await?;
|
||||
|
||||
let (max_cost, used_cost, enabled) =
|
||||
row.ok_or_else(|| AppError::Unauthorized("API key not found".into()))?;
|
||||
if !enabled.unwrap_or(true) {
|
||||
return Err(AppError::Forbidden("API key is disabled".into()));
|
||||
}
|
||||
if let Some(max_cost) = max_cost {
|
||||
if used_cost.unwrap_or(0.0) >= max_cost {
|
||||
return Err(AppError::RateLimited);
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn create_daily_usage_snapshot(&self) -> AppResult<()> {
|
||||
Self::create_daily_usage_snapshot_for_db(&self.db).await
|
||||
}
|
||||
|
||||
pub async fn create_daily_usage_snapshot_for_db(db: &SqlitePool) -> AppResult<()> {
|
||||
Self::ensure_schema(db).await?;
|
||||
sqlx::query(
|
||||
"INSERT INTO usage_daily_snapshots
|
||||
(local_day, downstream_key_id, request_count, prompt_tokens,
|
||||
completion_tokens, total_tokens, total_cost, updated_at)
|
||||
SELECT
|
||||
date(created_at, 'localtime') AS local_day,
|
||||
downstream_key_id,
|
||||
COUNT(*) AS request_count,
|
||||
COALESCE(SUM(prompt_tokens), 0) AS prompt_tokens,
|
||||
COALESCE(SUM(completion_tokens), 0) AS completion_tokens,
|
||||
COALESCE(SUM(total_tokens), 0) AS total_tokens,
|
||||
COALESCE(SUM(cost), 0.0) AS total_cost,
|
||||
datetime('now') AS updated_at
|
||||
FROM usage_logs
|
||||
GROUP BY date(created_at, 'localtime'), downstream_key_id
|
||||
ON CONFLICT(local_day, downstream_key_id) DO UPDATE SET
|
||||
request_count = excluded.request_count,
|
||||
prompt_tokens = excluded.prompt_tokens,
|
||||
completion_tokens = excluded.completion_tokens,
|
||||
total_tokens = excluded.total_tokens,
|
||||
total_cost = excluded.total_cost,
|
||||
updated_at = datetime('now')",
|
||||
)
|
||||
.execute(db)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn ensure_schema(db: &SqlitePool) -> AppResult<()> {
|
||||
sqlx::query(
|
||||
"CREATE TABLE IF NOT EXISTS usage_logs (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
downstream_key_id INTEGER,
|
||||
model TEXT,
|
||||
prompt_tokens INTEGER DEFAULT 0,
|
||||
completion_tokens INTEGER DEFAULT 0,
|
||||
total_tokens INTEGER DEFAULT 0,
|
||||
cost REAL DEFAULT 0.0,
|
||||
created_at TEXT DEFAULT (datetime('now'))
|
||||
)",
|
||||
)
|
||||
.execute(db)
|
||||
.await?;
|
||||
sqlx::query(
|
||||
"CREATE INDEX IF NOT EXISTS usage_logs_downstream_key_created_idx
|
||||
ON usage_logs (downstream_key_id, created_at)",
|
||||
)
|
||||
.execute(db)
|
||||
.await?;
|
||||
sqlx::query(
|
||||
"CREATE TABLE IF NOT EXISTS usage_daily_snapshots (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
local_day TEXT NOT NULL,
|
||||
downstream_key_id INTEGER,
|
||||
request_count INTEGER NOT NULL DEFAULT 0,
|
||||
prompt_tokens INTEGER NOT NULL DEFAULT 0,
|
||||
completion_tokens INTEGER NOT NULL DEFAULT 0,
|
||||
total_tokens INTEGER NOT NULL DEFAULT 0,
|
||||
total_cost REAL NOT NULL DEFAULT 0.0,
|
||||
created_at TEXT DEFAULT (datetime('now')),
|
||||
updated_at TEXT DEFAULT (datetime('now')),
|
||||
UNIQUE(local_day, downstream_key_id)
|
||||
)",
|
||||
)
|
||||
.execute(db)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use sqlx::sqlite::SqlitePoolOptions;
|
||||
|
||||
async fn pool() -> SqlitePool {
|
||||
let pool = SqlitePoolOptions::new()
|
||||
.max_connections(1)
|
||||
.connect("sqlite::memory:")
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
sqlx::query(
|
||||
"CREATE TABLE downstream_api_keys (
|
||||
id INTEGER PRIMARY KEY,
|
||||
name TEXT NOT NULL,
|
||||
key TEXT NOT NULL,
|
||||
enabled INTEGER DEFAULT 1,
|
||||
max_cost REAL,
|
||||
used_cost REAL DEFAULT 0,
|
||||
max_requests INTEGER,
|
||||
used_requests INTEGER DEFAULT 0,
|
||||
last_used_at TEXT,
|
||||
updated_at TEXT DEFAULT (datetime('now'))
|
||||
)",
|
||||
)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
sqlx::query(
|
||||
"INSERT INTO downstream_api_keys
|
||||
(id, name, key, enabled, max_cost, used_cost, max_requests, used_requests)
|
||||
VALUES (1, 'test', 'sk-test', 1, 1.0, 0.0, 5, 0)",
|
||||
)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
pool
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn records_and_aggregates_usage() {
|
||||
let tracker = UsageTracker::new(Arc::new(pool().await));
|
||||
|
||||
tracker
|
||||
.record_usage(1, "gpt-test", 10, 20, 0.25)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let summary = tracker.get_key_usage(1, None, None).await.unwrap();
|
||||
assert_eq!(summary.request_count, 1);
|
||||
assert_eq!(summary.prompt_tokens, 10);
|
||||
assert_eq!(summary.completion_tokens, 20);
|
||||
assert_eq!(summary.total_tokens, 30);
|
||||
assert_eq!(summary.total_cost, 0.25);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn rate_limit_checks_request_and_cost_limits() {
|
||||
let tracker = UsageTracker::new(Arc::new(pool().await));
|
||||
tracker.check_rate_limit(1).await.unwrap();
|
||||
|
||||
sqlx::query("UPDATE downstream_api_keys SET used_requests = 5 WHERE id = 1")
|
||||
.execute(&*tracker.db)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(matches!(
|
||||
tracker.check_rate_limit(1).await,
|
||||
Err(AppError::RateLimited)
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn creates_daily_snapshot() {
|
||||
let tracker = UsageTracker::new(Arc::new(pool().await));
|
||||
tracker
|
||||
.record_usage(1, "gpt-test", 2, 3, 0.1)
|
||||
.await
|
||||
.unwrap();
|
||||
tracker.create_daily_usage_snapshot().await.unwrap();
|
||||
|
||||
let count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM usage_daily_snapshots")
|
||||
.fetch_one(&*tracker.db)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(count, 1);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user