Compare commits
9 Commits
d697146941
...
be7538a1da
| Author | SHA1 | Date | |
|---|---|---|---|
| be7538a1da | |||
| cfe2098b73 | |||
| bd464016a8 | |||
| 7e8430934e | |||
| ab2fd271e8 | |||
| 4dd2be4cc3 | |||
| fd94303433 | |||
| 91dcbc5c53 | |||
| 278b0c66b8 |
@@ -0,0 +1,569 @@
|
||||
# 计费闭环 实现计划
|
||||
|
||||
> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking.
|
||||
|
||||
**Goal:** 为 BoosAPI 建立 ledger-first 的双轨计费闭环(USD 池 + 工作流积分池),所有余额变动写统一流水账 `credit_ledger`。
|
||||
|
||||
**Architecture:** 新建 `credit_ledger`(唯一真相源)+ `user_credits`(积分缓存)+ `recharge_tiers`(充值档位)三张表。USD 轨沿用 `downstream_api_keys.max_cost/used_cost` 作热路径缓存,每笔变动同事务写一条 ledger。积分轨为工作流执行新增扣费/拦截。分 6 期,本计划详写第 1 期(地基),后 5 期为概要待细化。
|
||||
|
||||
**Tech Stack:** Rust (axum + sqlx + SQLite)、React + Vite 前端、in-memory SQLite 测试。
|
||||
|
||||
---
|
||||
|
||||
## 文件结构(第 1 期)
|
||||
|
||||
- 新建 `src/services/billing/ledger.rs` — ledger 写入 + 余额查询的唯一入口(`record_ledger_entry`、`ensure_ledger_schema`)。所有动余额的地方调它,保证"缓存 + ledger 同事务"不变量集中在一处。
|
||||
- 修改 `src/services/billing/mod.rs` — `mod ledger;` 导出;`record_usage` 在写 usage_logs + increment 后,同事务补写 ledger(usd, consume)。
|
||||
- 修改 `src/services/payment.rs` — `credit_user_balance` 加写 ledger(usd, recharge)。
|
||||
- 修改 `src/routes/admin/users.rs` — `adjust_user_balance` 加写 ledger(usd, admin_adjust)。
|
||||
- 测试:`src/services/billing/ledger.rs` 内 `#[cfg(test)] mod tests`(复用 in-memory SQLite 模式)。
|
||||
|
||||
---
|
||||
|
||||
## 第 1 期:地基(credit_ledger + user_credits 表 + USD 轨记账接入)
|
||||
|
||||
### Task 1:ledger 表结构 + 写入入口
|
||||
|
||||
**Files:**
|
||||
- Create: `src/services/billing/ledger.rs`
|
||||
- Modify: `src/services/billing/mod.rs`(顶部加 `pub mod ledger;`)
|
||||
|
||||
- [ ] **Step 1: 写失败测试** — 在 `src/services/billing/ledger.rs` 末尾:
|
||||
|
||||
```rust
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use sqlx::sqlite::SqlitePoolOptions;
|
||||
|
||||
async fn mem_pool() -> SqlitePool {
|
||||
let pool = SqlitePoolOptions::new()
|
||||
.max_connections(1)
|
||||
.connect("sqlite::memory:")
|
||||
.await
|
||||
.unwrap();
|
||||
ensure_ledger_schema(&pool).await.unwrap();
|
||||
pool
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn record_ledger_entry_writes_row_with_balance_snapshot() {
|
||||
let pool = mem_pool().await;
|
||||
let mut tx = pool.begin().await.unwrap();
|
||||
record_ledger_entry(
|
||||
&mut tx,
|
||||
LedgerEntry {
|
||||
user_id: 7,
|
||||
track: "usd",
|
||||
entry_type: "recharge",
|
||||
amount: 10.0,
|
||||
balance_after: 10.0,
|
||||
ref_type: Some("stripe"),
|
||||
ref_id: Some("cs_1"),
|
||||
model: None,
|
||||
operator: Some("system"),
|
||||
note: None,
|
||||
},
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
tx.commit().await.unwrap();
|
||||
|
||||
let row: (i64, String, String, f64, f64) = sqlx::query_as(
|
||||
"SELECT user_id, track, entry_type, amount, balance_after
|
||||
FROM credit_ledger WHERE user_id = 7",
|
||||
)
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(row, (7, "usd".into(), "recharge".into(), 10.0, 10.0));
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
- [ ] **Step 2: 运行测试确认失败**
|
||||
|
||||
Run: `cargo test -p boosapi --lib services::billing::ledger 2>&1 | tail -5`
|
||||
Expected: 编译失败(`ensure_ledger_schema`/`record_ledger_entry`/`LedgerEntry` 未定义)。
|
||||
|
||||
- [ ] **Step 3: 写最小实现** — 在 `src/services/billing/ledger.rs` 顶部(tests 之前):
|
||||
|
||||
```rust
|
||||
use crate::error::AppResult;
|
||||
use sqlx::{SqlitePool, Transaction};
|
||||
|
||||
/// One balance-changing event. Written to `credit_ledger` in the same
|
||||
/// transaction that updates the cached balance, so ledger and cache never diverge.
|
||||
pub struct LedgerEntry<'a> {
|
||||
pub user_id: i64,
|
||||
pub track: &'a str, // "usd" | "points"
|
||||
pub entry_type: &'a str, // "recharge" | "admin_adjust" | "consume" | "refund" | "grant"
|
||||
pub amount: f64, // +credit / -debit
|
||||
pub balance_after: f64,
|
||||
pub ref_type: Option<&'a str>,
|
||||
pub ref_id: Option<&'a str>,
|
||||
pub model: Option<&'a str>,
|
||||
pub operator: Option<&'a str>,
|
||||
pub note: Option<&'a str>,
|
||||
}
|
||||
|
||||
pub async fn ensure_ledger_schema(db: &SqlitePool) -> AppResult<()> {
|
||||
sqlx::query(
|
||||
"CREATE TABLE IF NOT EXISTS credit_ledger (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
|
||||
user_id INTEGER NOT NULL,
|
||||
track TEXT NOT NULL,
|
||||
entry_type TEXT NOT NULL,
|
||||
amount REAL NOT NULL,
|
||||
balance_after REAL NOT NULL,
|
||||
ref_type TEXT,
|
||||
ref_id TEXT,
|
||||
model TEXT,
|
||||
operator TEXT,
|
||||
note TEXT,
|
||||
created_at TEXT DEFAULT (datetime('now'))
|
||||
)",
|
||||
)
|
||||
.execute(db)
|
||||
.await?;
|
||||
sqlx::query(
|
||||
"CREATE INDEX IF NOT EXISTS credit_ledger_user_created_idx
|
||||
ON credit_ledger (user_id, created_at DESC)",
|
||||
)
|
||||
.execute(db)
|
||||
.await?;
|
||||
sqlx::query(
|
||||
"CREATE INDEX IF NOT EXISTS credit_ledger_track_type_idx
|
||||
ON credit_ledger (track, entry_type)",
|
||||
)
|
||||
.execute(db)
|
||||
.await?;
|
||||
sqlx::query(
|
||||
"CREATE TABLE IF NOT EXISTS user_credits (
|
||||
user_id INTEGER PRIMARY KEY NOT NULL,
|
||||
points_balance REAL NOT NULL DEFAULT 0,
|
||||
points_total_recharged REAL NOT NULL DEFAULT 0,
|
||||
updated_at TEXT DEFAULT (datetime('now'))
|
||||
)",
|
||||
)
|
||||
.execute(db)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Append one ledger row inside an existing transaction. Caller is responsible
|
||||
/// for also updating the cached balance in the same `tx`.
|
||||
pub async fn record_ledger_entry(
|
||||
tx: &mut Transaction<'_, sqlx::Sqlite>,
|
||||
entry: LedgerEntry<'_>,
|
||||
) -> AppResult<()> {
|
||||
sqlx::query(
|
||||
"INSERT INTO credit_ledger
|
||||
(user_id, track, entry_type, amount, balance_after,
|
||||
ref_type, ref_id, model, operator, note)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
|
||||
)
|
||||
.bind(entry.user_id)
|
||||
.bind(entry.track)
|
||||
.bind(entry.entry_type)
|
||||
.bind(entry.amount)
|
||||
.bind(entry.balance_after)
|
||||
.bind(entry.ref_type)
|
||||
.bind(entry.ref_id)
|
||||
.bind(entry.model)
|
||||
.bind(entry.operator)
|
||||
.bind(entry.note)
|
||||
.execute(&mut **tx)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
```
|
||||
|
||||
并在 `src/services/billing/mod.rs` 顶部加:`pub mod ledger;`
|
||||
|
||||
- [ ] **Step 4: 运行测试确认通过**
|
||||
|
||||
Run: `cargo test -p boosapi --lib services::billing::ledger 2>&1 | tail -5`
|
||||
Expected: PASS(1 passed)。
|
||||
|
||||
- [ ] **Step 5: 提交**
|
||||
|
||||
```bash
|
||||
git add src/services/billing/ledger.rs src/services/billing/mod.rs
|
||||
git commit -m "feat(billing): credit_ledger schema + record_ledger_entry"
|
||||
```
|
||||
|
||||
### Task 2:USD 消费接入 ledger(record_usage)
|
||||
|
||||
**Files:**
|
||||
- Modify: `src/services/billing/mod.rs`(`record_usage`,行 ~48-77)
|
||||
|
||||
- [ ] **Step 1: 写失败测试** — 在 billing/mod.rs 的 tests mod 加:
|
||||
|
||||
```rust
|
||||
#[tokio::test]
|
||||
async fn record_usage_writes_consume_ledger_entry() {
|
||||
let pool = billing_test_pool().await; // 复用现有测试 pool helper;若无则建 users+downstream_api_keys+usage_logs
|
||||
let tracker = UsageTracker::new(std::sync::Arc::new(pool.clone()));
|
||||
crate::services::billing::ledger::ensure_ledger_schema(&pool).await.unwrap();
|
||||
// key_id=7 属于 user 9,max_cost=10 used_cost=2(由 helper 建好)
|
||||
tracker.record_usage(7, "gpt-4o", 100, 50, 0.42).await.unwrap();
|
||||
|
||||
let row: (String, String, f64, Option<String>) = sqlx::query_as(
|
||||
"SELECT track, entry_type, amount, model FROM credit_ledger
|
||||
WHERE user_id = 9 ORDER BY id DESC LIMIT 1",
|
||||
).fetch_one(&pool).await.unwrap();
|
||||
assert_eq!(row.0, "usd");
|
||||
assert_eq!(row.1, "consume");
|
||||
assert_eq!(row.2, -0.42); // 消费为负
|
||||
assert_eq!(row.3.as_deref(), Some("gpt-4o"));
|
||||
}
|
||||
```
|
||||
|
||||
注:若 billing/mod.rs 无现成 `billing_test_pool`,在 tests mod 加一个建 `users`/`downstream_api_keys`(含 user_id=9,key id=7 max_cost=10 used_cost=2)/`usage_logs` 的 helper(参照 payment.rs 的 payment_test_pool)。
|
||||
|
||||
- [ ] **Step 2: 运行确认失败**
|
||||
|
||||
Run: `cargo test -p boosapi --lib record_usage_writes_consume_ledger_entry 2>&1 | tail -5`
|
||||
Expected: FAIL(credit_ledger 无 user 9 行 / 表不存在)。
|
||||
|
||||
- [ ] **Step 3: 实现** — `record_usage` 改为事务化,末尾补写 ledger。替换 record_usage 函数体为:
|
||||
|
||||
```rust
|
||||
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?;
|
||||
crate::services::billing::ledger::ensure_ledger_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;
|
||||
let cost = cost.max(0.0);
|
||||
let model = model.into();
|
||||
|
||||
let mut tx = self.db.begin().await?;
|
||||
sqlx::query(
|
||||
"INSERT INTO usage_logs
|
||||
(downstream_key_id, model, prompt_tokens, completion_tokens, total_tokens, cost)
|
||||
VALUES (?, ?, ?, ?, ?, ?)",
|
||||
)
|
||||
.bind(key_id).bind(&model)
|
||||
.bind(prompt_tokens).bind(completion_tokens).bind(total_tokens).bind(cost)
|
||||
.execute(&mut *tx).await?;
|
||||
|
||||
sqlx::query(
|
||||
"UPDATE downstream_api_keys
|
||||
SET used_requests = COALESCE(used_requests,0)+1,
|
||||
used_cost = COALESCE(used_cost,0)+?,
|
||||
last_used_at = datetime('now'), updated_at = datetime('now')
|
||||
WHERE id = ?",
|
||||
).bind(cost).bind(key_id).execute(&mut *tx).await?;
|
||||
|
||||
// Resolve owning user + new used_cost for the ledger snapshot.
|
||||
let owner: Option<(Option<i64>, Option<f64>, Option<f64>)> = sqlx::query_as(
|
||||
"SELECT user_id, max_cost, used_cost FROM downstream_api_keys WHERE id = ?",
|
||||
).bind(key_id).fetch_optional(&mut *tx).await?;
|
||||
if let Some((Some(user_id), max_cost, used_cost)) = owner {
|
||||
let used = used_cost.unwrap_or(0.0);
|
||||
let remaining = max_cost.map(|m| (m - used).max(0.0)).unwrap_or(used);
|
||||
crate::services::billing::ledger::record_ledger_entry(
|
||||
&mut tx,
|
||||
crate::services::billing::ledger::LedgerEntry {
|
||||
user_id, track: "usd", entry_type: "consume",
|
||||
amount: -cost, balance_after: remaining,
|
||||
ref_type: Some("usage_log"), ref_id: None,
|
||||
model: Some(&model), operator: Some("system"), note: None,
|
||||
},
|
||||
).await?;
|
||||
}
|
||||
tx.commit().await?;
|
||||
Ok(())
|
||||
}
|
||||
```
|
||||
|
||||
(注:原 `increment_key_usage(&self.db, ...)` 调用被事务内 UPDATE 取代,删除该行。)
|
||||
|
||||
- [ ] **Step 4: 运行确认通过**
|
||||
|
||||
Run: `cargo test -p boosapi --lib record_usage 2>&1 | tail -6`
|
||||
Expected: PASS。
|
||||
|
||||
- [ ] **Step 5: 提交**
|
||||
|
||||
```bash
|
||||
git add src/services/billing/mod.rs
|
||||
git commit -m "feat(billing): USD consume writes credit_ledger in one tx"
|
||||
```
|
||||
|
||||
### Task 3:USD 充值接入 ledger(credit_user_balance)
|
||||
|
||||
**Files:**
|
||||
- Modify: `src/services/payment.rs`(`credit_user_balance`,行 ~2167-2202)
|
||||
|
||||
- [ ] **Step 1: 写失败测试** — payment.rs tests mod(复用 payment_test_pool,user 9/key 7 max_cost=10):
|
||||
|
||||
```rust
|
||||
#[tokio::test]
|
||||
async fn credit_user_balance_writes_recharge_ledger() {
|
||||
let pool = payment_test_pool().await;
|
||||
crate::services::billing::ledger::ensure_ledger_schema(&pool).await.unwrap();
|
||||
let mut tx = pool.begin().await.unwrap();
|
||||
credit_user_balance(&mut tx, 9, 20.0).await.unwrap();
|
||||
tx.commit().await.unwrap();
|
||||
let row: (String, String, f64) = sqlx::query_as(
|
||||
"SELECT track, entry_type, amount FROM credit_ledger WHERE user_id = 9 ORDER BY id DESC LIMIT 1",
|
||||
).fetch_one(&pool).await.unwrap();
|
||||
assert_eq!(row, ("usd".into(), "recharge".into(), 20.0));
|
||||
}
|
||||
```
|
||||
|
||||
- [ ] **Step 2: 运行确认失败**
|
||||
|
||||
Run: `cargo test -p boosapi --lib credit_user_balance_writes_recharge_ledger 2>&1 | tail -5`
|
||||
Expected: FAIL。
|
||||
|
||||
- [ ] **Step 3: 实现** — `credit_user_balance` 在 UPDATE max_cost 之后、`Ok(Some(key_id))` 之前插入 ledger 写入(已在事务 tx 内):
|
||||
|
||||
```rust
|
||||
// After the UPDATE that sets next_max_cost on the key:
|
||||
crate::services::billing::ledger::record_ledger_entry(
|
||||
tx,
|
||||
crate::services::billing::ledger::LedgerEntry {
|
||||
user_id, track: "usd", entry_type: "recharge",
|
||||
amount: amount_usd.max(0.0),
|
||||
balance_after: (next_max_cost - used_cost).max(0.0),
|
||||
ref_type: Some("recharge_record"), ref_id: None,
|
||||
model: None, operator: Some("system"), note: None,
|
||||
},
|
||||
).await?;
|
||||
```
|
||||
|
||||
(`credit_user_balance` 已是 `tx: &mut Transaction`,直接传 `tx`。确保调用方在调用前已 `ensure_ledger_schema`,或在函数开头加一次幂等建表。)
|
||||
|
||||
- [ ] **Step 4: 运行确认通过**
|
||||
|
||||
Run: `cargo test -p boosapi --lib credit_user_balance 2>&1 | tail -6`
|
||||
Expected: PASS(含已有的 recharge completion 测试)。
|
||||
|
||||
- [ ] **Step 5: 提交**
|
||||
|
||||
```bash
|
||||
git add src/services/payment.rs
|
||||
git commit -m "feat(billing): USD recharge writes credit_ledger"
|
||||
```
|
||||
|
||||
### Task 4:USD 管理员调整接入 ledger(adjust_user_balance)
|
||||
|
||||
**Files:**
|
||||
- Modify: `src/routes/admin/users.rs`(`adjust_user_balance`,行 ~820)
|
||||
|
||||
- [ ] **Step 1: 写失败测试** — users.rs tests mod:调 `adjust_user_balance(State, Path(9), Json{delta:5.0})` 后断言 credit_ledger 有一行 `(usd, admin_adjust, 5.0)`。(建表与 user/key 同 Task 3 helper。)
|
||||
|
||||
- [ ] **Step 2: 运行确认失败** — `cargo test -p boosapi --lib adjust_user_balance 2>&1 | tail -5` → FAIL。
|
||||
|
||||
- [ ] **Step 3: 实现** — `adjust_user_balance` 现在直接 `execute(&*state.db)`。改为事务:begin → UPDATE max_cost → `record_ledger_entry(track="usd", entry_type="admin_adjust", amount=delta, balance_after=next_max-used, operator=<管理员,暂用 "admin">)` → commit。函数开头 `ledger::ensure_ledger_schema(&state.db).await?;`。
|
||||
|
||||
- [ ] **Step 4: 运行确认通过** — `cargo test -p boosapi --lib adjust_user_balance 2>&1 | tail -6` → PASS。
|
||||
|
||||
- [ ] **Step 5: 提交**
|
||||
|
||||
```bash
|
||||
git add src/routes/admin/users.rs
|
||||
git commit -m "feat(billing): USD admin adjust writes credit_ledger"
|
||||
```
|
||||
|
||||
### Task 5:全量回归
|
||||
|
||||
- [ ] **Step 1:** `cargo test 2>&1 | grep "test result"` → 全 ok。
|
||||
- [ ] **Step 2:** `cargo build 2>&1 | tail -3` → Finished。
|
||||
- [ ] **Step 3:** 提交(若有遗留)。
|
||||
|
||||
---
|
||||
|
||||
## 后续期次(概要,第 1 期落地后逐期细化为同样的 TDD 任务)
|
||||
|
||||
### 第 2 期:积分轨
|
||||
- `user_credits` 已建表(Task 1)。新增 `adjust_user_credits`(admin 发放/增减,写 ledger points/admin_adjust)。
|
||||
- comfy_cloud 确认到账 `admin_confirm_user_recharge` 改为加 `points_balance` + ledger(points, recharge),受 config 开关 `payment.comfy_recharge_to_points` 控制(默认 true);历史不迁移。
|
||||
- 工作流积分单价 config(默认值 + 按工作流名分级映射)。
|
||||
- `workflow/executor.rs` `run_workflow`:执行**前**查 `points_balance >= 单价`,不足返回 Forbidden「工作流积分不足」;执行**成功后**扣积分 + ledger(points, consume, model=工作流名)。
|
||||
- 测试:拦截、扣费、充值进积分、管理员调整。
|
||||
|
||||
### 第 3 期:A 消费明细
|
||||
- `GET /api/user/billing/ledger`(查本人 credit_ledger,track/type/时间过滤、分页)。
|
||||
- `GET /api/user/billing/summary`(按天/模型聚合,双轨分开)。
|
||||
- 扩展 `UsageRecords.tsx`:双轨切换 + 汇总卡 + 明细列表。
|
||||
|
||||
### 第 4 期:B 充值档位
|
||||
- `recharge_tiers` 表(已在 spec,Task 1 未建,本期建)+ admin CRUD `/api/admin/recharge-tiers`。
|
||||
- 用户充值档位从 DB 读(替代 env `USER_RECHARGE_OPTIONS`)。
|
||||
- admin `RechargeTiers` 配置页 + `UserRecharge.tsx` 接 DB。
|
||||
|
||||
### 第 5 期:C 预警
|
||||
- 余额接口(user_billing_summary 等)加 `lowBalance`(余额 < 阈值,config 默认 0.2)。
|
||||
- Dashboard / UserRecharge / ComfyUIAgent 顶部预警条。
|
||||
|
||||
### 第 6 期:D 审计 + 对账
|
||||
- `GET /api/admin/billing/ledger`(跨用户,过滤分页)+ admin `BillingAudit` 页。
|
||||
- cron:拉 Comfy `/customers/balance` 总余额,与 points consume ledger 比对,偏差超阈值写 log + admin 告警标志(不自动改数据)。
|
||||
|
||||
---
|
||||
|
||||
## 自审
|
||||
|
||||
- **spec 覆盖**:第 1 期覆盖 spec「数据模型」+「USD 轨记账接入」三处(充值/消费/调整);后 5 期概要逐条对应 spec 的积分轨/A/B/C/D + 对账。recharge_tiers 表 spec 列在数据模型,实现挪到第 4 期建(第 1 期用不到),已注明。
|
||||
- **占位扫描**:第 1 期 Task 1-3 含完整代码;Task 4-5 用文字描述但步骤明确(adjust 与 Task 3 同模式,刻意不重复大段代码,执行时参照 Task 3)。后 5 期为概要,执行前细化。
|
||||
- **类型一致**:`LedgerEntry` 字段、`record_ledger_entry`/`ensure_ledger_schema` 签名在 Task 1 定义,Task 2-4 一致引用。`track`/`entry_type` 取值(usd/points、recharge/consume/admin_adjust)全程一致。
|
||||
|
||||
---
|
||||
|
||||
## 第 2 期:积分轨(细化 TDD)
|
||||
|
||||
文件结构:
|
||||
- `src/services/billing/ledger.rs` 加积分轨核心:`points_balance(db,user_id)`(读余额)、`apply_points_delta(db, user_id, delta, entry_type, model, operator, ref_*)`(事务内:UPSERT user_credits 余额 + 写 ledger(points),返回新余额)。所有积分增减走这一个入口。
|
||||
- `src/config.rs`:ComfyUIConfig 加 `workflow_points_cost`(默认单价)。
|
||||
- `src/workflow/executor.rs`:run_workflow_with_progress 前置拦截 + 收尾扣积分。
|
||||
- `src/services/payment.rs`:admin_confirm_user_recharge 改为加积分(受 config 开关)。
|
||||
- `src/routes/admin/users.rs`:新增 `adjust_user_credits` handler + 路由。
|
||||
|
||||
### Task 6:积分轨核心 apply_points_delta + points_balance
|
||||
|
||||
**Files:** Modify `src/services/billing/ledger.rs`
|
||||
|
||||
- [ ] **Step1 失败测试**(ledger.rs tests mod):
|
||||
|
||||
```rust
|
||||
#[tokio::test]
|
||||
async fn apply_points_delta_upserts_balance_and_writes_ledger() {
|
||||
let pool = mem_pool().await;
|
||||
// 首次 +100(充值)
|
||||
let bal = apply_points_delta(&pool, 9, 100.0, "recharge", None, Some("system"), Some("recharge_record"), None).await.unwrap();
|
||||
assert_eq!(bal, 100.0);
|
||||
// 再 -30(消费)
|
||||
let bal2 = apply_points_delta(&pool, 9, -30.0, "consume", Some("wf/video"), Some("system"), Some("workflow"), None).await.unwrap();
|
||||
assert_eq!(bal2, 70.0);
|
||||
// user_credits 缓存 = 70
|
||||
let cached: f64 = sqlx::query_scalar("SELECT points_balance FROM user_credits WHERE user_id = 9").fetch_one(&pool).await.unwrap();
|
||||
assert_eq!(cached, 70.0);
|
||||
// ledger 两行
|
||||
let n: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM credit_ledger WHERE user_id=9 AND track='points'").fetch_one(&pool).await.unwrap();
|
||||
assert_eq!(n, 2);
|
||||
}
|
||||
```
|
||||
|
||||
- [ ] **Step2** `cargo test -p boosapi --lib apply_points_delta_upserts 2>&1 | tail -6` → FAIL。
|
||||
|
||||
- [ ] **Step3 实现**(ledger.rs,tests 之前):
|
||||
|
||||
```rust
|
||||
/// Current points balance for a user (0 if no row yet).
|
||||
pub async fn points_balance(db: &SqlitePool, user_id: i64) -> AppResult<f64> {
|
||||
let bal: Option<f64> = sqlx::query_scalar(
|
||||
"SELECT points_balance FROM user_credits WHERE user_id = ?",
|
||||
)
|
||||
.bind(user_id)
|
||||
.fetch_optional(db)
|
||||
.await?;
|
||||
Ok(bal.unwrap_or(0.0))
|
||||
}
|
||||
|
||||
/// Apply a points delta (+credit / -debit) atomically: upsert the cached
|
||||
/// balance in `user_credits` and append one `credit_ledger(track=points)` row,
|
||||
/// in a single transaction. Returns the new balance.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub async fn apply_points_delta(
|
||||
db: &SqlitePool,
|
||||
user_id: i64,
|
||||
delta: f64,
|
||||
entry_type: &str,
|
||||
model: Option<&str>,
|
||||
operator: Option<&str>,
|
||||
ref_type: Option<&str>,
|
||||
ref_id: Option<&str>,
|
||||
) -> AppResult<f64> {
|
||||
ensure_ledger_schema(db).await?;
|
||||
let mut tx = db.begin().await?;
|
||||
let current: f64 = sqlx::query_scalar(
|
||||
"SELECT points_balance FROM user_credits WHERE user_id = ?",
|
||||
)
|
||||
.bind(user_id)
|
||||
.fetch_optional(&mut *tx)
|
||||
.await?
|
||||
.unwrap_or(0.0);
|
||||
let next = current + delta;
|
||||
let total_recharged_inc = if delta > 0.0 && entry_type == "recharge" { delta } else { 0.0 };
|
||||
sqlx::query(
|
||||
"INSERT INTO user_credits (user_id, points_balance, points_total_recharged, updated_at)
|
||||
VALUES (?, ?, ?, datetime('now'))
|
||||
ON CONFLICT(user_id) DO UPDATE SET
|
||||
points_balance = ?,
|
||||
points_total_recharged = points_total_recharged + ?,
|
||||
updated_at = datetime('now')",
|
||||
)
|
||||
.bind(user_id).bind(next).bind(total_recharged_inc)
|
||||
.bind(next).bind(total_recharged_inc)
|
||||
.execute(&mut *tx).await?;
|
||||
record_ledger_entry(&mut tx, LedgerEntry {
|
||||
user_id, track: "points", entry_type,
|
||||
amount: delta, balance_after: next,
|
||||
ref_type, ref_id, model, operator, note: None,
|
||||
}).await?;
|
||||
tx.commit().await?;
|
||||
Ok(next)
|
||||
}
|
||||
```
|
||||
|
||||
- [ ] **Step4** `cargo test -p boosapi --lib apply_points_delta 2>&1 | tail -6` → PASS;`cargo build -j 2 | tail -2`。
|
||||
- [ ] **Step5** `git add src/services/billing/ledger.rs && git commit -m "feat(billing): points track apply_points_delta + balance"`
|
||||
|
||||
### Task 7:工作流积分单价 config
|
||||
|
||||
**Files:** Modify `src/config.rs`
|
||||
|
||||
- [ ] **Step1** 测试:config 默认值测试(若 config.rs 有 default 测试则加断言;否则跳测试,靠编译)。
|
||||
- [ ] **Step2/3 实现**:ComfyUIConfig 加字段 `#[serde(default = "default_workflow_points_cost")] pub workflow_points_cost: f64,`,default fn 返回 `10.0`;Default impl 里加 `workflow_points_cost: 10.0,`。
|
||||
- [ ] **Step4** `cargo build -j 2 2>&1 | tail -3` → Finished(config 字段加完整个 crate 能编译)。
|
||||
- [ ] **Step5** `git add src/config.rs && git commit -m "feat(billing): workflow_points_cost config"`
|
||||
|
||||
### Task 8:工作流执行前置拦截 + 成功扣积分
|
||||
|
||||
**Files:** Modify `src/workflow/executor.rs`
|
||||
|
||||
- [ ] **Step1 失败测试**:executor.rs tests mod(若无则建)。建 in-memory pool + user_credits/credit_ledger,构造一个最简 graph,`ctx.user_id=Some(9)`:
|
||||
- 余额不足(0 积分,单价 10)→ run_workflow_with_progress 返回 Err(Forbidden 含"积分不足")。
|
||||
- 余额充足(100)→ 执行成功后 points_balance 变 90,credit_ledger 出现一行 (points, consume, -10)。
|
||||
(若 executor 难以单测真实节点,则把"拦截"和"扣费"逻辑抽成小函数 `check_workflow_points(state, user_id, cost)` 和 `charge_workflow_points(state, user_id, cost, workflow_name)` 各自单测,run_workflow 调它们。优先抽函数单测。)
|
||||
|
||||
- [ ] **Step2** 确认 FAIL。
|
||||
|
||||
- [ ] **Step3 实现**:
|
||||
- 在 run_workflow_with_progress 的 `validate_and_sort` 之后、`insert_run` 之前:若 `ctx.user_id` 是 Some 且 `ctx.charge_points`(见下):`let cost = state.config.comfyui.workflow_points_cost; if cost > 0.0 { let bal = ledger::points_balance(&state.db, uid).await?; if bal < cost { return Err(AppError::Forbidden("工作流积分不足,请充值".into())); } }`
|
||||
- 在 `complete_run` 之后、`Ok(run_public_id)` 之前:`if let Some(uid) = ctx.user_id { let cost = state.config.comfyui.workflow_points_cost; if cost > 0.0 { let _ = ledger::apply_points_delta(&state.db, uid, -cost, "consume", Some(&format!("workflow:{}", workflow_id)), Some("system"), Some("workflow_run"), Some(&run_public_id)).await; } }`(扣费 best-effort,失败记 log 不阻断已产出结果)。
|
||||
- 注意:Admin 调用(builtin key,user_id 可能是 Some(admin) 或特殊)是否扣?按设计仅对普通用户扣。可用 RunContext 加 `pub charge_points: bool` 字段,run_context 构造时 User=>true / Admin=>false。models.rs RunContext 加该字段,两处构造(routes.rs run_context、workflow_ui.rs workflow_context)补 `charge_points`。
|
||||
|
||||
- [ ] **Step4** 测试 PASS + `cargo build -j 2`。
|
||||
- [ ] **Step5** `git add src/workflow/executor.rs src/workflow/models.rs src/workflow/routes.rs src/routes/workflow_ui.rs && git commit -m "feat(billing): workflow points gating + charge"`
|
||||
|
||||
### Task 9:comfy 充值改加积分 + admin 积分调整
|
||||
|
||||
**Files:** Modify `src/services/payment.rs`、`src/routes/admin/users.rs`、`src/config.rs`(开关)、`src/routes/admin/mod.rs`(路由)
|
||||
|
||||
- [ ] **Step1 失败测试**:
|
||||
- payment.rs:`admin_confirm_user_recharge` 确认后,credit_ledger 出现 (points, recharge, +amount)(而非 usd);user_credits 增加。
|
||||
- users.rs:`adjust_user_credits` handler 调用后 credit_ledger 出现 (points, admin_adjust)。
|
||||
- [ ] **Step2** FAIL。
|
||||
- [ ] **Step3 实现**:
|
||||
- config 开关 `payment.comfy_recharge_to_points: bool`(默认 true)。
|
||||
- `admin_confirm_user_recharge`:开关为 true 时改调 `ledger::apply_points_delta(db, user_id, amount_usd_as_points, "recharge", None, Some("admin"), Some("recharge_record"), Some(public_id))`(USD→积分换算:本期 1 USD = N 积分,N 用 config `payment.usd_to_points_rate` 默认 100;即 grant_points = amount_usd * rate)。开关 false 时维持原 credit_user_balance(USD)。**历史不迁移**。
|
||||
- 新增 `adjust_user_credits(State, Path(id), Json{delta})` handler(仿 adjust_user_balance)→ `apply_points_delta(db,id,delta,"admin_adjust",None,Some("admin"),Some("manual"),None)`,返回 {success, data{userId, pointsBalance}}。
|
||||
- admin/mod.rs 注册 `POST /api/admin/users/{id}/credits/adjust`。
|
||||
- [ ] **Step4** 相关测试 + payment/users 模块全过 + build。
|
||||
- [ ] **Step5** 分两个 commit 或一个:`git commit -m "feat(billing): comfy recharge to points + admin credit adjust"`
|
||||
|
||||
### Task 10:第 2 期全量回归
|
||||
- [ ] `cargo test 2>&1 | grep "test result"` 全 ok;`cargo build -j 2 | tail -2` Finished。
|
||||
@@ -0,0 +1,149 @@
|
||||
# 计费闭环系统设计
|
||||
|
||||
日期:2026-06-16
|
||||
状态:设计已获批,待写实现计划
|
||||
|
||||
## 背景
|
||||
|
||||
BoosAPI 现有计费是零散拼接的:每个用户的 downstream API key 有 `max_cost`(美元上限)/ `used_cost`(已消耗),proxy 请求路径在 `used_cost >= max_cost` 时硬拦截返回 429;充值通过 Stripe 自动到账或 comfy_cloud 管理员确认到账,调用 `credit_user_balance` 增加 `max_cost`;新用户默认 $1 额度;usage_logs 按 key 记每笔消费。
|
||||
|
||||
缺口:ComfyUI 工作流执行不扣额度、无积分体系、无消费明细可视化、无充值档位配置、无低额度预警、无全量变动审计。
|
||||
|
||||
本设计把上述补齐为一个 **ledger-first(统一流水账)** 的计费闭环。
|
||||
|
||||
## 核心模型:双轨额度
|
||||
|
||||
- **USD 轨**:LLM 对话、生图、视频等所有走 `/v1/*` 的调用。沿用 `downstream_api_keys.max_cost/used_cost`。
|
||||
- **工作流积分轨**:仅 ComfyUI 工作流执行。新建每用户独立的积分池,平台可管理。
|
||||
|
||||
两轨并行、互不换算。comfy_cloud 充值进积分轨;其他充值(Stripe)进 USD 轨。
|
||||
|
||||
## 数据模型
|
||||
|
||||
### 新表 `credit_ledger`(唯一真相源,所有变动记一行)
|
||||
|
||||
| 字段 | 类型 | 说明 |
|
||||
|---|---|---|
|
||||
| `id` | INTEGER PK AUTOINCREMENT | |
|
||||
| `user_id` | INTEGER NOT NULL | 所属用户 |
|
||||
| `track` | TEXT NOT NULL | `usd` / `points` |
|
||||
| `entry_type` | TEXT NOT NULL | `recharge` / `admin_adjust` / `consume` / `refund` / `grant` |
|
||||
| `amount` | REAL NOT NULL | 正=增 负=减(usd 用美元,points 用积分) |
|
||||
| `balance_after` | REAL NOT NULL | 该轨本次变动后余额快照 |
|
||||
| `ref_type` | TEXT | `usage_log` / `recharge_record` / `manual` / `stripe` 等 |
|
||||
| `ref_id` | TEXT | 关联记录 id/public_id |
|
||||
| `model` | TEXT | 消费类:模型名或工作流名(供聚合) |
|
||||
| `operator` | TEXT | 操作者:管理员名 / `system` / `user` |
|
||||
| `note` | TEXT | 备注 |
|
||||
| `created_at` | TEXT DEFAULT (datetime('now')) | |
|
||||
|
||||
索引:`(user_id, created_at DESC)`、`(track, entry_type)`、`(user_id, track)`。
|
||||
|
||||
### 新表 `user_credits`(积分余额缓存)
|
||||
|
||||
| 字段 | 类型 | 说明 |
|
||||
|---|---|---|
|
||||
| `user_id` | INTEGER PK | |
|
||||
| `points_balance` | REAL NOT NULL DEFAULT 0 | 当前积分余额(热路径读写) |
|
||||
| `points_total_recharged` | REAL NOT NULL DEFAULT 0 | 累计充值积分(统计) |
|
||||
| `updated_at` | TEXT | |
|
||||
|
||||
### 新表 `recharge_tiers`(充值档位,admin 可配)
|
||||
|
||||
| 字段 | 类型 | 说明 |
|
||||
|---|---|---|
|
||||
| `id` | INTEGER PK AUTOINCREMENT | |
|
||||
| `track` | TEXT NOT NULL | `usd` / `points` |
|
||||
| `amount_usd` | REAL NOT NULL | 支付金额(USD) |
|
||||
| `grant_amount` | REAL NOT NULL | 到账量(USD 数或积分数) |
|
||||
| `bonus` | REAL NOT NULL DEFAULT 0 | 赠送量 |
|
||||
| `title` | TEXT | 档位标题 |
|
||||
| `sort_order` | INTEGER DEFAULT 0 | 排序 |
|
||||
| `enabled` | INTEGER DEFAULT 1 | 启用 |
|
||||
|
||||
### 余额派生策略:ledger + 缓存双写
|
||||
|
||||
- **USD 轨**:`max_cost/used_cost` 作快速余额缓存(拦截/展示零延迟读),每笔变动**同事务**写一条 `credit_ledger(track=usd)`。
|
||||
- **积分轨**:`user_credits.points_balance` 作缓存,每笔变动**同事务**写 `credit_ledger(track=points)`。
|
||||
- **关键不变量**:任何余额变动 = 缓存字段更新 + 一条 ledger,**同一 DB 事务**,事务失败整体回滚,绝不只改其一。
|
||||
- ledger 是审计/明细/对账真相;缓存字段是热路径性能优化。
|
||||
|
||||
## 双轨记账接入
|
||||
|
||||
原则:所有动余额处,同事务"改缓存 + 写一条 ledger"。
|
||||
|
||||
### USD 轨
|
||||
|
||||
| 变动 | 现有代码 | 改造 |
|
||||
|---|---|---|
|
||||
| 充值到账 | `credit_user_balance`(payment.rs) | 同事务加写 ledger(usd, recharge) |
|
||||
| 消费扣费 | `increment_key_usage`/`record_billing` | 同事务加写 ledger(usd, consume, model=模型名) |
|
||||
| 管理员增减 | `adjust_user_balance` | 同事务加写 ledger(usd, admin_adjust, operator) |
|
||||
| 额度拦截 | proxy `used_cost>=max_cost`→429 | 不变,补低额度预警信号 |
|
||||
|
||||
### 积分轨(基本新建)
|
||||
|
||||
| 变动 | 现状 | 新增 |
|
||||
|---|---|---|
|
||||
| 充值到账 | comfy_cloud 确认到账现加 USD | 改为加 `points_balance` + ledger(points, recharge) |
|
||||
| 消费扣费 | 工作流不扣 | 工作流执行**成功后**按单价扣积分 + ledger(points, consume, model=工作流名);失败不扣 |
|
||||
| 管理员发放/增减 | 无 | 新增 `adjust_user_credits` + ledger(points, admin_adjust) |
|
||||
| 积分拦截 | 无 | 工作流执行**前**检查 `points_balance >= 单价`,不足拒绝;低额度预警 |
|
||||
|
||||
- 积分消费扣费点:`workflow/executor.rs` 的 `run_workflow` 成功收尾处(有 `ctx.user_id`),与现有"工作流按用户 key 记 USD 消费"平行。
|
||||
- **comfy_cloud 充值改向积分轨**:行为变更,历史 USD **不迁移**,仅对切换后新充值生效,用 config 开关控制(可回滚)。
|
||||
|
||||
## 积分消费规则:A+C 混合
|
||||
|
||||
- **A 固定/分级单价**(主):工作流执行扣固定积分(config 默认值,可按工作流类型分级)。可预测、支持并发。
|
||||
- **C 对账校准**(兜底):后台 cron 拉 Comfy 总余额对账,偏差超阈值提示管理员调单价。
|
||||
|
||||
## 四块功能
|
||||
|
||||
### A 消费明细(用户侧,明细+聚合)
|
||||
- 新建 `GET /api/user/billing/ledger`:查该用户 credit_ledger,可按 track/type/时间过滤。
|
||||
- 新建 `GET /api/user/billing/summary`:按天/按模型聚合,双轨分开。
|
||||
- 扩展 `UsageRecords.tsx`:双轨切换 + 聚合汇总卡 + 明细列表。
|
||||
|
||||
### B 充值档位(本期可配置,订阅后期)
|
||||
- 新建 `recharge_tiers` 表 + admin CRUD `GET/POST/PUT/DELETE /api/admin/recharge-tiers`。
|
||||
- 改造用户充值:档位从 DB 读(替代 env 写死的 `USER_RECHARGE_OPTIONS`)。
|
||||
- 扩展 `UserRecharge.tsx`(档位从接口拉)+ 新建 admin `RechargeTiers` 配置页。
|
||||
|
||||
### C 超额 + 预警(两轨硬拦 + 预警)
|
||||
- USD 拦截已有(不动);积分拦截新建(工作流执行前)。
|
||||
- 余额接口加 `lowBalance` 标志(余额 < 阈值,config 可配默认 20%)。
|
||||
- Dashboard/充值页/视频助手:余额低显示预警条。
|
||||
|
||||
### D 审计对账(管理员侧,全量流水)
|
||||
- 新建 `GET /api/admin/billing/ledger`:跨用户查 credit_ledger(按 user/track/type/时间过滤、分页)。
|
||||
- 新建 admin `BillingAudit` 页:全量变动流水表(谁/何时/给谁/加减/原因),站点管理风格。
|
||||
- 对账校准:cron 拉 Comfy 总余额,与积分消费 ledger 比对,偏差超阈值写 log + admin 告警标志。**不自动改数据**,仅提示,管理员手动调单价。
|
||||
|
||||
## 错误处理(不变量)
|
||||
|
||||
- 余额变动 = 缓存 + ledger 同事务;失败整体回滚。
|
||||
- 消费扣费 best-effort 不阻断响应(沿用 record_billing 原则)。
|
||||
- 积分轨:拦截是前置硬检查,扣费是后置;后置扣费失败已产出结果 → 记日志补偿。
|
||||
- comfy 充值改向积分:config 开关控制,可回滚。
|
||||
|
||||
## 测试策略
|
||||
|
||||
- 单元:ledger 写入+缓存同步事务一致性;积分拦截(不足拒绝);单价扣费;充值进积分;管理员调整双轨。
|
||||
- 复用现有 payment/billing 测试框架(in-memory SQLite)。
|
||||
|
||||
## 实现顺序(分 6 期,每期可独立上线)
|
||||
|
||||
1. **地基**:credit_ledger + user_credits 表 + 迁移;USD 轨记账接进 ledger(充值/消费/调整三处双写)。
|
||||
2. **积分轨**:工作流积分扣费 + 拦截 + comfy 充值改进积分 + 管理员积分调整。
|
||||
3. **A 消费明细**:用户侧 ledger/summary 接口 + UsageRecords 扩展。
|
||||
4. **B 充值档位**:recharge_tiers 表 + admin 配置页 + 用户充值页接 DB。
|
||||
5. **C 预警**:余额接口加 lowBalance + 前端预警条。
|
||||
6. **D 审计 + 对账**:admin BillingAudit 页 + cron 对账校准。
|
||||
|
||||
## 新增清单小结
|
||||
|
||||
- **表**:credit_ledger(核心)、user_credits(积分缓存)、recharge_tiers(档位)。
|
||||
- **后端接口**:user billing ledger/summary、admin recharge-tiers CRUD、admin billing ledger、adjust_user_credits、积分拦截/扣费。
|
||||
- **前端页面**:admin RechargeTiers、admin BillingAudit;扩展 UsageRecords / UserRecharge / Dashboard。
|
||||
- **cron**:Comfy 余额对账校准任务。
|
||||
@@ -62,6 +62,8 @@ pub struct ComfyUIConfig {
|
||||
pub money_printer_video_cost: f64,
|
||||
#[serde(default)]
|
||||
pub digital_human_video_cost: f64,
|
||||
#[serde(default = "default_workflow_points_cost")]
|
||||
pub workflow_points_cost: f64,
|
||||
#[serde(default = "default_comfyui_docker_image")]
|
||||
pub docker_image: String,
|
||||
#[serde(default = "default_comfyui_container_name")]
|
||||
@@ -113,6 +115,10 @@ fn default_comfyui_container_name() -> String {
|
||||
"boosapi-comfyui-workflow".to_string()
|
||||
}
|
||||
|
||||
fn default_workflow_points_cost() -> f64 {
|
||||
10.0
|
||||
}
|
||||
|
||||
impl Default for ComfyUIConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
@@ -120,6 +126,7 @@ impl Default for ComfyUIConfig {
|
||||
backend: ComfyUIBackend::Boosapi,
|
||||
money_printer_video_cost: 0.0,
|
||||
digital_human_video_cost: 0.0,
|
||||
workflow_points_cost: 10.0,
|
||||
docker_image: default_comfyui_docker_image(),
|
||||
container_name: default_comfyui_container_name(),
|
||||
install_dir: None,
|
||||
|
||||
+244
-2
@@ -4,6 +4,7 @@ use crate::{
|
||||
auth::password::hash_password,
|
||||
db::models::{downstream_api_key::DownstreamApiKey, user::User},
|
||||
error::{AppError, AppResult},
|
||||
routes::auth::DEFAULT_NEW_USER_MAX_COST_USD,
|
||||
AppState,
|
||||
};
|
||||
use axum::{
|
||||
@@ -238,17 +239,18 @@ pub async fn create_user(
|
||||
|
||||
let user_id = result.last_insert_rowid();
|
||||
|
||||
// Auto-generate downstream API key
|
||||
// Auto-generate downstream API key with the default starter credit.
|
||||
let key = format!("sk-user-{}", uuid::Uuid::new_v4());
|
||||
let key_name = format!("{} 的默认密钥", payload.username.trim());
|
||||
|
||||
sqlx::query(
|
||||
"INSERT INTO downstream_api_keys (name, key, user_id, enabled) VALUES (?, ?, ?, ?)",
|
||||
"INSERT INTO downstream_api_keys (name, key, user_id, enabled, max_cost) VALUES (?, ?, ?, ?, ?)",
|
||||
)
|
||||
.bind(&key_name)
|
||||
.bind(&key)
|
||||
.bind(user_id)
|
||||
.bind(true)
|
||||
.bind(DEFAULT_NEW_USER_MAX_COST_USD)
|
||||
.execute(&mut *tx)
|
||||
.await?;
|
||||
|
||||
@@ -708,6 +710,195 @@ pub async fn update_user_key(
|
||||
Ok(Json(json!({ "data": updated })))
|
||||
}
|
||||
|
||||
// ── User balance management (admin) ────────────────────────────────────
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct ListBalancesQuery {
|
||||
pub page: Option<i64>,
|
||||
pub page_size: Option<i64>,
|
||||
pub q: Option<String>,
|
||||
}
|
||||
|
||||
/// GET /api/admin/user-balances — list users with their credit pool
|
||||
/// (summed max_cost / used_cost across their downstream keys).
|
||||
pub async fn list_user_balances(
|
||||
State(state): State<AppState>,
|
||||
Query(query): Query<ListBalancesQuery>,
|
||||
) -> AppResult<Json<Value>> {
|
||||
let page = query.page.unwrap_or(1).max(1);
|
||||
let page_size = query.page_size.unwrap_or(50).clamp(1, 200);
|
||||
let offset = (page - 1) * page_size;
|
||||
let q = query.q.unwrap_or_default().trim().to_string();
|
||||
let q_like = format!("%{}%", q);
|
||||
|
||||
let rows = sqlx::query(
|
||||
r#"
|
||||
SELECT
|
||||
u.id AS user_id,
|
||||
u.username AS username,
|
||||
u.email AS email,
|
||||
u.role AS role,
|
||||
u.status AS status,
|
||||
COALESCE(SUM(COALESCE(k.used_cost, 0.0)), 0.0) AS used_cost,
|
||||
COALESCE(SUM(CASE WHEN k.max_cost IS NULL THEN 0.0 ELSE k.max_cost END), 0.0) AS limited_max_cost,
|
||||
SUM(CASE WHEN k.id IS NULL THEN 0 WHEN k.max_cost IS NULL THEN 1 ELSE 0 END) AS unlimited_key_count,
|
||||
COUNT(k.id) AS key_count
|
||||
FROM users u
|
||||
LEFT JOIN downstream_api_keys k ON k.user_id = u.id
|
||||
WHERE (? = '' OR u.username LIKE ? OR u.email LIKE ? OR CAST(u.id AS TEXT) = ?)
|
||||
GROUP BY u.id
|
||||
ORDER BY u.id DESC
|
||||
LIMIT ? OFFSET ?
|
||||
"#,
|
||||
)
|
||||
.bind(&q)
|
||||
.bind(&q_like)
|
||||
.bind(&q_like)
|
||||
.bind(&q)
|
||||
.bind(page_size)
|
||||
.bind(offset)
|
||||
.fetch_all(&*state.db)
|
||||
.await?;
|
||||
|
||||
let total: i64 = sqlx::query_scalar(
|
||||
"SELECT COUNT(*) FROM users
|
||||
WHERE (? = '' OR username LIKE ? OR email LIKE ? OR CAST(id AS TEXT) = ?)",
|
||||
)
|
||||
.bind(&q)
|
||||
.bind(&q_like)
|
||||
.bind(&q_like)
|
||||
.bind(&q)
|
||||
.fetch_one(&*state.db)
|
||||
.await?;
|
||||
|
||||
let items: Vec<Value> = rows
|
||||
.iter()
|
||||
.map(|row| {
|
||||
use sqlx::Row;
|
||||
let used: f64 = row.try_get("used_cost").unwrap_or(0.0);
|
||||
let limited_max: f64 = row.try_get("limited_max_cost").unwrap_or(0.0);
|
||||
let unlimited: i64 = row.try_get("unlimited_key_count").unwrap_or(0);
|
||||
let remaining = if unlimited > 0 {
|
||||
Value::Null
|
||||
} else {
|
||||
Value::from((limited_max - used).max(0.0))
|
||||
};
|
||||
json!({
|
||||
"userId": row.try_get::<i64, _>("user_id").unwrap_or(0),
|
||||
"username": row.try_get::<Option<String>, _>("username").ok().flatten(),
|
||||
"email": row.try_get::<Option<String>, _>("email").ok().flatten(),
|
||||
"role": row.try_get::<Option<String>, _>("role").ok().flatten(),
|
||||
"status": row.try_get::<Option<String>, _>("status").ok().flatten(),
|
||||
"keyCount": row.try_get::<i64, _>("key_count").unwrap_or(0),
|
||||
"usedCost": used,
|
||||
"maxCost": if unlimited > 0 { Value::Null } else { Value::from(limited_max) },
|
||||
"remainingCost": remaining,
|
||||
"unlimitedKeyCount": unlimited,
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
|
||||
Ok(Json(json!({
|
||||
"data": items,
|
||||
"meta": {
|
||||
"page": page,
|
||||
"pageSize": page_size,
|
||||
"total": total,
|
||||
"totalPages": ((total + page_size - 1) / page_size).max(1),
|
||||
},
|
||||
})))
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct AdjustBalanceRequest {
|
||||
/// Amount to add (positive) or remove (negative) from the user's credit pool.
|
||||
pub delta: f64,
|
||||
}
|
||||
|
||||
/// POST /api/admin/users/{id}/balance/adjust — add or remove credit on the
|
||||
/// user's primary (lowest-id) downstream key by bumping its max_cost.
|
||||
pub async fn adjust_user_balance(
|
||||
State(state): State<AppState>,
|
||||
Path(id): Path<i64>,
|
||||
Json(payload): Json<AdjustBalanceRequest>,
|
||||
) -> AppResult<Json<Value>> {
|
||||
let delta = payload.delta;
|
||||
if !delta.is_finite() || delta == 0.0 {
|
||||
return Err(AppError::BadRequest("delta must be a non-zero finite number".into()));
|
||||
}
|
||||
|
||||
crate::services::billing::ledger::ensure_ledger_schema(&state.db).await?;
|
||||
|
||||
let key: Option<(i64, Option<f64>, Option<f64>)> = sqlx::query_as(
|
||||
"SELECT id, max_cost, used_cost
|
||||
FROM downstream_api_keys
|
||||
WHERE user_id = ?
|
||||
ORDER BY id ASC
|
||||
LIMIT 1",
|
||||
)
|
||||
.bind(id)
|
||||
.fetch_optional(&*state.db)
|
||||
.await?;
|
||||
|
||||
let Some((key_id, max_cost, used_cost)) = key else {
|
||||
return Err(AppError::BadRequest(
|
||||
"该用户没有可调整额度的 API Key(请先为其创建 Key)".into(),
|
||||
));
|
||||
};
|
||||
|
||||
let used = used_cost.unwrap_or(0.0).max(0.0);
|
||||
let current = max_cost.unwrap_or(used).max(used);
|
||||
// Never let the limit drop below what's already been spent.
|
||||
let next_max = (current + delta).max(used);
|
||||
let remaining = (next_max - used).max(0.0);
|
||||
|
||||
let mut tx = state.db.begin().await?;
|
||||
sqlx::query(
|
||||
"UPDATE downstream_api_keys
|
||||
SET max_cost = ?, updated_at = datetime('now')
|
||||
WHERE id = ?",
|
||||
)
|
||||
.bind(next_max)
|
||||
.bind(key_id)
|
||||
.execute(&mut *tx)
|
||||
.await?;
|
||||
|
||||
crate::services::billing::ledger::record_ledger_entry(
|
||||
&mut tx,
|
||||
crate::services::billing::ledger::LedgerEntry {
|
||||
user_id: id,
|
||||
track: "usd",
|
||||
entry_type: "admin_adjust",
|
||||
amount: delta,
|
||||
balance_after: remaining,
|
||||
ref_type: Some("manual"),
|
||||
ref_id: None,
|
||||
model: None,
|
||||
operator: Some("admin"),
|
||||
note: None,
|
||||
},
|
||||
)
|
||||
.await?;
|
||||
tx.commit().await?;
|
||||
|
||||
info!(
|
||||
user_id = id,
|
||||
key_id, delta, next_max, "admin adjusted user balance"
|
||||
);
|
||||
|
||||
Ok(Json(json!({
|
||||
"success": true,
|
||||
"data": {
|
||||
"userId": id,
|
||||
"keyId": key_id,
|
||||
"maxCost": next_max,
|
||||
"usedCost": used,
|
||||
"remainingCost": remaining,
|
||||
}
|
||||
})))
|
||||
}
|
||||
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
@@ -800,6 +991,8 @@ mod tests {
|
||||
template_proxy_url: None,
|
||||
},
|
||||
comfyui: Default::default(),
|
||||
pixelle_video: Default::default(),
|
||||
payment: Default::default(),
|
||||
},
|
||||
pool,
|
||||
BuiltinKeys {
|
||||
@@ -847,4 +1040,53 @@ mod tests {
|
||||
.expect("read role");
|
||||
assert_eq!(role, "admin");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn adjust_user_balance_writes_admin_adjust_ledger() {
|
||||
let state = test_state().await;
|
||||
crate::services::billing::ledger::ensure_ledger_schema(&state.db)
|
||||
.await
|
||||
.expect("ledger schema");
|
||||
|
||||
let user_id = sqlx::query(
|
||||
"INSERT INTO users (username, email, password_hash, role, status)
|
||||
VALUES (?, ?, ?, 'user', 'active')",
|
||||
)
|
||||
.bind("Spender")
|
||||
.bind("spender@example.test")
|
||||
.bind("hash")
|
||||
.execute(&*state.db)
|
||||
.await
|
||||
.expect("insert user")
|
||||
.last_insert_rowid();
|
||||
|
||||
sqlx::query(
|
||||
"INSERT INTO downstream_api_keys (name, key, user_id, max_cost, used_cost)
|
||||
VALUES (?, ?, ?, 10.0, 0.0)",
|
||||
)
|
||||
.bind("primary")
|
||||
.bind("sk-spender")
|
||||
.bind(user_id)
|
||||
.execute(&*state.db)
|
||||
.await
|
||||
.expect("insert key");
|
||||
|
||||
adjust_user_balance(
|
||||
State(state.clone()),
|
||||
Path(user_id),
|
||||
Json(AdjustBalanceRequest { delta: 5.0 }),
|
||||
)
|
||||
.await
|
||||
.expect("adjust balance");
|
||||
|
||||
let row: (String, String, f64) = sqlx::query_as(
|
||||
"SELECT track, entry_type, amount FROM credit_ledger
|
||||
WHERE user_id = ? ORDER BY id DESC LIMIT 1",
|
||||
)
|
||||
.bind(user_id)
|
||||
.fetch_one(&*state.db)
|
||||
.await
|
||||
.expect("ledger row");
|
||||
assert_eq!(row, ("usd".into(), "admin_adjust".into(), 5.0));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,210 @@
|
||||
use crate::error::AppResult;
|
||||
use sqlx::{SqlitePool, Transaction};
|
||||
|
||||
/// One balance-changing event. Written to `credit_ledger` in the same
|
||||
/// transaction that updates the cached balance, so ledger and cache never diverge.
|
||||
pub struct LedgerEntry<'a> {
|
||||
pub user_id: i64,
|
||||
pub track: &'a str, // "usd" | "points"
|
||||
pub entry_type: &'a str, // "recharge" | "admin_adjust" | "consume" | "refund" | "grant"
|
||||
pub amount: f64, // +credit / -debit
|
||||
pub balance_after: f64,
|
||||
pub ref_type: Option<&'a str>,
|
||||
pub ref_id: Option<&'a str>,
|
||||
pub model: Option<&'a str>,
|
||||
pub operator: Option<&'a str>,
|
||||
pub note: Option<&'a str>,
|
||||
}
|
||||
|
||||
pub async fn ensure_ledger_schema(db: &SqlitePool) -> AppResult<()> {
|
||||
sqlx::query(
|
||||
"CREATE TABLE IF NOT EXISTS credit_ledger (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
|
||||
user_id INTEGER NOT NULL,
|
||||
track TEXT NOT NULL,
|
||||
entry_type TEXT NOT NULL,
|
||||
amount REAL NOT NULL,
|
||||
balance_after REAL NOT NULL,
|
||||
ref_type TEXT,
|
||||
ref_id TEXT,
|
||||
model TEXT,
|
||||
operator TEXT,
|
||||
note TEXT,
|
||||
created_at TEXT DEFAULT (datetime('now'))
|
||||
)",
|
||||
)
|
||||
.execute(db)
|
||||
.await?;
|
||||
sqlx::query(
|
||||
"CREATE INDEX IF NOT EXISTS credit_ledger_user_created_idx
|
||||
ON credit_ledger (user_id, created_at DESC)",
|
||||
)
|
||||
.execute(db)
|
||||
.await?;
|
||||
sqlx::query(
|
||||
"CREATE INDEX IF NOT EXISTS credit_ledger_track_type_idx
|
||||
ON credit_ledger (track, entry_type)",
|
||||
)
|
||||
.execute(db)
|
||||
.await?;
|
||||
sqlx::query(
|
||||
"CREATE TABLE IF NOT EXISTS user_credits (
|
||||
user_id INTEGER PRIMARY KEY NOT NULL,
|
||||
points_balance REAL NOT NULL DEFAULT 0,
|
||||
points_total_recharged REAL NOT NULL DEFAULT 0,
|
||||
updated_at TEXT DEFAULT (datetime('now'))
|
||||
)",
|
||||
)
|
||||
.execute(db)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Append one ledger row inside an existing transaction. Caller is responsible
|
||||
/// for also updating the cached balance in the same `tx`.
|
||||
pub async fn record_ledger_entry(
|
||||
tx: &mut Transaction<'_, sqlx::Sqlite>,
|
||||
entry: LedgerEntry<'_>,
|
||||
) -> AppResult<()> {
|
||||
sqlx::query(
|
||||
"INSERT INTO credit_ledger
|
||||
(user_id, track, entry_type, amount, balance_after,
|
||||
ref_type, ref_id, model, operator, note)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
|
||||
)
|
||||
.bind(entry.user_id)
|
||||
.bind(entry.track)
|
||||
.bind(entry.entry_type)
|
||||
.bind(entry.amount)
|
||||
.bind(entry.balance_after)
|
||||
.bind(entry.ref_type)
|
||||
.bind(entry.ref_id)
|
||||
.bind(entry.model)
|
||||
.bind(entry.operator)
|
||||
.bind(entry.note)
|
||||
.execute(&mut **tx)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Current points balance for a user (0 if no row yet).
|
||||
pub async fn points_balance(db: &SqlitePool, user_id: i64) -> AppResult<f64> {
|
||||
let bal: Option<f64> = sqlx::query_scalar(
|
||||
"SELECT points_balance FROM user_credits WHERE user_id = ?",
|
||||
)
|
||||
.bind(user_id)
|
||||
.fetch_optional(db)
|
||||
.await?;
|
||||
Ok(bal.unwrap_or(0.0))
|
||||
}
|
||||
|
||||
/// Apply a points delta (+credit / -debit) atomically: upsert the cached
|
||||
/// balance in `user_credits` and append one `credit_ledger(track=points)` row,
|
||||
/// in a single transaction. Returns the new balance.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub async fn apply_points_delta(
|
||||
db: &SqlitePool,
|
||||
user_id: i64,
|
||||
delta: f64,
|
||||
entry_type: &str,
|
||||
model: Option<&str>,
|
||||
operator: Option<&str>,
|
||||
ref_type: Option<&str>,
|
||||
ref_id: Option<&str>,
|
||||
) -> AppResult<f64> {
|
||||
ensure_ledger_schema(db).await?;
|
||||
let mut tx = db.begin().await?;
|
||||
let current: f64 = sqlx::query_scalar(
|
||||
"SELECT points_balance FROM user_credits WHERE user_id = ?",
|
||||
)
|
||||
.bind(user_id)
|
||||
.fetch_optional(&mut *tx)
|
||||
.await?
|
||||
.unwrap_or(0.0);
|
||||
let next = current + delta;
|
||||
let total_recharged_inc = if delta > 0.0 && entry_type == "recharge" { delta } else { 0.0 };
|
||||
sqlx::query(
|
||||
"INSERT INTO user_credits (user_id, points_balance, points_total_recharged, updated_at)
|
||||
VALUES (?, ?, ?, datetime('now'))
|
||||
ON CONFLICT(user_id) DO UPDATE SET
|
||||
points_balance = ?,
|
||||
points_total_recharged = points_total_recharged + ?,
|
||||
updated_at = datetime('now')",
|
||||
)
|
||||
.bind(user_id).bind(next).bind(total_recharged_inc)
|
||||
.bind(next).bind(total_recharged_inc)
|
||||
.execute(&mut *tx).await?;
|
||||
record_ledger_entry(&mut tx, LedgerEntry {
|
||||
user_id, track: "points", entry_type,
|
||||
amount: delta, balance_after: next,
|
||||
ref_type, ref_id, model, operator, note: None,
|
||||
}).await?;
|
||||
tx.commit().await?;
|
||||
Ok(next)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use sqlx::sqlite::SqlitePoolOptions;
|
||||
|
||||
async fn mem_pool() -> SqlitePool {
|
||||
let pool = SqlitePoolOptions::new()
|
||||
.max_connections(1)
|
||||
.connect("sqlite::memory:")
|
||||
.await
|
||||
.unwrap();
|
||||
ensure_ledger_schema(&pool).await.unwrap();
|
||||
pool
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn record_ledger_entry_writes_row_with_balance_snapshot() {
|
||||
let pool = mem_pool().await;
|
||||
let mut tx = pool.begin().await.unwrap();
|
||||
record_ledger_entry(
|
||||
&mut tx,
|
||||
LedgerEntry {
|
||||
user_id: 7,
|
||||
track: "usd",
|
||||
entry_type: "recharge",
|
||||
amount: 10.0,
|
||||
balance_after: 10.0,
|
||||
ref_type: Some("stripe"),
|
||||
ref_id: Some("cs_1"),
|
||||
model: None,
|
||||
operator: Some("system"),
|
||||
note: None,
|
||||
},
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
tx.commit().await.unwrap();
|
||||
|
||||
let row: (i64, String, String, f64, f64) = sqlx::query_as(
|
||||
"SELECT user_id, track, entry_type, amount, balance_after
|
||||
FROM credit_ledger WHERE user_id = 7",
|
||||
)
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(row, (7, "usd".into(), "recharge".into(), 10.0, 10.0));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn apply_points_delta_upserts_balance_and_writes_ledger() {
|
||||
let pool = mem_pool().await;
|
||||
// 首次 +100(充值)
|
||||
let bal = apply_points_delta(&pool, 9, 100.0, "recharge", None, Some("system"), Some("recharge_record"), None).await.unwrap();
|
||||
assert_eq!(bal, 100.0);
|
||||
// 再 -30(消费)
|
||||
let bal2 = apply_points_delta(&pool, 9, -30.0, "consume", Some("wf/video"), Some("system"), Some("workflow"), None).await.unwrap();
|
||||
assert_eq!(bal2, 70.0);
|
||||
// user_credits 缓存 = 70
|
||||
let cached: f64 = sqlx::query_scalar("SELECT points_balance FROM user_credits WHERE user_id = 9").fetch_one(&pool).await.unwrap();
|
||||
assert_eq!(cached, 70.0);
|
||||
// ledger 两行
|
||||
let n: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM credit_ledger WHERE user_id=9 AND track='points'").fetch_one(&pool).await.unwrap();
|
||||
assert_eq!(n, 2);
|
||||
}
|
||||
}
|
||||
+148
-11
@@ -1,6 +1,8 @@
|
||||
pub mod ledger;
|
||||
|
||||
use crate::{
|
||||
error::{AppError, AppResult},
|
||||
services::shared::{enforce_key_request_limit, increment_key_usage},
|
||||
services::shared::enforce_key_request_limit,
|
||||
};
|
||||
use chrono::Timelike;
|
||||
use serde::{Deserialize, Serialize};
|
||||
@@ -54,25 +56,68 @@ impl UsageTracker {
|
||||
cost: f64,
|
||||
) -> AppResult<()> {
|
||||
Self::ensure_schema(&self.db).await?;
|
||||
crate::services::billing::ledger::ensure_ledger_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;
|
||||
let cost = cost.max(0.0);
|
||||
let model = model.into();
|
||||
|
||||
let mut tx = self.db.begin().await?;
|
||||
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(&model)
|
||||
.bind(prompt_tokens)
|
||||
.bind(completion_tokens)
|
||||
.bind(total_tokens)
|
||||
.bind(cost.max(0.0))
|
||||
.execute(&*self.db)
|
||||
.bind(cost)
|
||||
.execute(&mut *tx)
|
||||
.await?;
|
||||
|
||||
increment_key_usage(&self.db, key_id, cost.max(0.0)).await?;
|
||||
sqlx::query(
|
||||
"UPDATE downstream_api_keys
|
||||
SET used_requests = COALESCE(used_requests,0)+1,
|
||||
used_cost = COALESCE(used_cost,0)+?,
|
||||
last_used_at = datetime('now'), updated_at = datetime('now')
|
||||
WHERE id = ?",
|
||||
)
|
||||
.bind(cost)
|
||||
.bind(key_id)
|
||||
.execute(&mut *tx)
|
||||
.await?;
|
||||
|
||||
// Resolve owning user + new used_cost for the ledger snapshot.
|
||||
let owner: Option<(Option<i64>, Option<f64>, Option<f64>)> = sqlx::query_as(
|
||||
"SELECT user_id, max_cost, used_cost FROM downstream_api_keys WHERE id = ?",
|
||||
)
|
||||
.bind(key_id)
|
||||
.fetch_optional(&mut *tx)
|
||||
.await?;
|
||||
if let Some((Some(user_id), max_cost, used_cost)) = owner {
|
||||
let used = used_cost.unwrap_or(0.0);
|
||||
let remaining = max_cost.map(|m| (m - used).max(0.0)).unwrap_or(used);
|
||||
crate::services::billing::ledger::record_ledger_entry(
|
||||
&mut tx,
|
||||
crate::services::billing::ledger::LedgerEntry {
|
||||
user_id,
|
||||
track: "usd",
|
||||
entry_type: "consume",
|
||||
amount: -cost,
|
||||
balance_after: remaining,
|
||||
ref_type: Some("usage_log"),
|
||||
ref_id: None,
|
||||
model: Some(&model),
|
||||
operator: Some("system"),
|
||||
note: None,
|
||||
},
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
tx.commit().await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -229,7 +274,12 @@ impl UsageTracker {
|
||||
.bind(event.key_id)
|
||||
.bind(&event.model)
|
||||
.bind(&event.model)
|
||||
.bind(event.error_message.as_deref().unwrap_or("Comfy task failed"))
|
||||
.bind(
|
||||
event
|
||||
.error_message
|
||||
.as_deref()
|
||||
.unwrap_or("Comfy task failed"),
|
||||
)
|
||||
.bind(event.client_family)
|
||||
.bind(event.client_app_id)
|
||||
.bind(event.client_app_name)
|
||||
@@ -825,6 +875,7 @@ mod tests {
|
||||
used_cost REAL DEFAULT 0,
|
||||
max_requests INTEGER,
|
||||
used_requests INTEGER DEFAULT 0,
|
||||
user_id INTEGER,
|
||||
last_used_at TEXT,
|
||||
updated_at TEXT DEFAULT (datetime('now'))
|
||||
)",
|
||||
@@ -844,6 +895,62 @@ mod tests {
|
||||
pool
|
||||
}
|
||||
|
||||
async fn billing_test_pool() -> SqlitePool {
|
||||
let pool = SqlitePoolOptions::new()
|
||||
.max_connections(1)
|
||||
.connect("sqlite::memory:")
|
||||
.await
|
||||
.unwrap();
|
||||
sqlx::query(
|
||||
"CREATE TABLE users (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
username TEXT NOT NULL,
|
||||
email TEXT NOT NULL,
|
||||
password_hash TEXT NOT NULL,
|
||||
role TEXT NOT NULL DEFAULT 'user',
|
||||
status TEXT NOT NULL DEFAULT 'active',
|
||||
created_at TEXT,
|
||||
updated_at TEXT
|
||||
)",
|
||||
)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
sqlx::query(
|
||||
"CREATE TABLE downstream_api_keys (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
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,
|
||||
user_id INTEGER,
|
||||
last_used_at TEXT,
|
||||
updated_at TEXT DEFAULT (datetime('now'))
|
||||
)",
|
||||
)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
sqlx::query(
|
||||
"INSERT INTO users (id, username, email, password_hash, role, status)
|
||||
VALUES (9, 'creator', 'creator@example.test', 'hash', 'user', 'active')",
|
||||
)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
sqlx::query(
|
||||
"INSERT INTO downstream_api_keys (id, name, key, max_cost, used_cost, user_id)
|
||||
VALUES (7, 'default', 'sk-user-test', 10.0, 2.0, 9)",
|
||||
)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
pool
|
||||
}
|
||||
|
||||
async fn create_proxy_logs(pool: &SqlitePool) {
|
||||
sqlx::query(
|
||||
"CREATE TABLE proxy_logs (
|
||||
@@ -911,6 +1018,32 @@ mod tests {
|
||||
assert_eq!(summary.total_cost, 0.25);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn record_usage_writes_consume_ledger_entry() {
|
||||
let pool = billing_test_pool().await;
|
||||
let tracker = UsageTracker::new(std::sync::Arc::new(pool.clone()));
|
||||
crate::services::billing::ledger::ensure_ledger_schema(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
// key_id=7 属于 user 9,max_cost=10 used_cost=2(由 helper 建好)
|
||||
tracker
|
||||
.record_usage(7, "gpt-4o", 100, 50, 0.42)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let row: (String, String, f64, Option<String>) = sqlx::query_as(
|
||||
"SELECT track, entry_type, amount, model FROM credit_ledger
|
||||
WHERE user_id = 9 ORDER BY id DESC LIMIT 1",
|
||||
)
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(row.0, "usd");
|
||||
assert_eq!(row.1, "consume");
|
||||
assert_eq!(row.2, -0.42); // 消费为负
|
||||
assert_eq!(row.3.as_deref(), Some("gpt-4o"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn rate_limit_checks_request_and_cost_limits() {
|
||||
let tracker = UsageTracker::new(Arc::new(pool().await));
|
||||
@@ -1017,7 +1150,10 @@ mod tests {
|
||||
.fetch_one(&*tracker.db)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(proxy_log, ("failed".to_string(), 0.0, "Not found".to_string()));
|
||||
assert_eq!(
|
||||
proxy_log,
|
||||
("failed".to_string(), 0.0, "Not found".to_string())
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -1076,7 +1212,9 @@ mod tests {
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let projected = UsageTracker::project_proxy_usage_for_db(&pool).await.unwrap();
|
||||
let projected = UsageTracker::project_proxy_usage_for_db(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(projected, 1);
|
||||
|
||||
let model_usage: (i64, i64, f64, i64) = sqlx::query_as(
|
||||
@@ -1099,9 +1237,8 @@ mod tests {
|
||||
.unwrap();
|
||||
assert_eq!(site_usage, (1, 1, 0.12));
|
||||
|
||||
let hour_bucket: String = sqlx::query_scalar(
|
||||
"SELECT bucket_start_utc FROM site_hour_usage WHERE site_id = 10",
|
||||
)
|
||||
let hour_bucket: String =
|
||||
sqlx::query_scalar("SELECT bucket_start_utc FROM site_hour_usage WHERE site_id = 10")
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
+2402
-33
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user