feat(saas): add billing infrastructure — tables, types, service, handlers
B1.1 Billing database: - 5 tables: billing_plans, billing_subscriptions, billing_invoices, billing_payments, billing_usage_quotas - Seed data: Free(¥0)/Pro(¥49)/Team(¥199) plans - JSONB limits for flexible plan configuration Billing module (crates/zclaw-saas/src/billing/): - types.rs: BillingPlan, Subscription, Invoice, Payment, UsageQuota - service.rs: plan CRUD, subscription lookup, usage tracking, quota check - handlers.rs: REST API (plans list/detail, subscription, usage) - mod.rs: routes registered at /api/v1/billing/* Cargo.toml: added chrono feature to sqlx for DateTime<Utc> support
This commit is contained in:
@@ -57,7 +57,7 @@ chrono = { version = "0.4", features = ["serde"] }
|
|||||||
uuid = { version = "1", features = ["v4", "v5", "serde"] }
|
uuid = { version = "1", features = ["v4", "v5", "serde"] }
|
||||||
|
|
||||||
# Database
|
# Database
|
||||||
sqlx = { version = "0.7", features = ["runtime-tokio", "sqlite", "postgres"] }
|
sqlx = { version = "0.7", features = ["runtime-tokio", "sqlite", "postgres", "chrono"] }
|
||||||
libsqlite3-sys = { version = "0.27", features = ["bundled"] }
|
libsqlite3-sys = { version = "0.27", features = ["bundled"] }
|
||||||
|
|
||||||
# HTTP client (for LLM drivers)
|
# HTTP client (for LLM drivers)
|
||||||
|
|||||||
133
crates/zclaw-saas/migrations/20260402000001_billing_tables.sql
Normal file
133
crates/zclaw-saas/migrations/20260402000001_billing_tables.sql
Normal file
@@ -0,0 +1,133 @@
|
|||||||
|
-- Migration: Billing tables for subscription management
|
||||||
|
-- Supports: Free/Pro/Team plans, Alipay + WeChat Pay, usage quotas
|
||||||
|
|
||||||
|
-- Plan definitions (Free/Pro/Team)
|
||||||
|
CREATE TABLE IF NOT EXISTS billing_plans (
|
||||||
|
id TEXT PRIMARY KEY,
|
||||||
|
name TEXT NOT NULL UNIQUE,
|
||||||
|
display_name TEXT NOT NULL,
|
||||||
|
description TEXT,
|
||||||
|
price_cents INTEGER NOT NULL DEFAULT 0,
|
||||||
|
currency TEXT NOT NULL DEFAULT 'CNY',
|
||||||
|
interval TEXT NOT NULL DEFAULT 'month',
|
||||||
|
features JSONB NOT NULL DEFAULT '{}',
|
||||||
|
limits JSONB NOT NULL DEFAULT '{}',
|
||||||
|
is_default BOOLEAN NOT NULL DEFAULT FALSE,
|
||||||
|
sort_order INTEGER NOT NULL DEFAULT 0,
|
||||||
|
status TEXT NOT NULL DEFAULT 'active',
|
||||||
|
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
||||||
|
updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
|
||||||
|
);
|
||||||
|
CREATE INDEX IF NOT EXISTS idx_billing_plans_status ON billing_plans(status);
|
||||||
|
|
||||||
|
-- Account subscriptions
|
||||||
|
CREATE TABLE IF NOT EXISTS billing_subscriptions (
|
||||||
|
id TEXT PRIMARY KEY,
|
||||||
|
account_id TEXT NOT NULL,
|
||||||
|
plan_id TEXT NOT NULL,
|
||||||
|
status TEXT NOT NULL DEFAULT 'active',
|
||||||
|
current_period_start TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
||||||
|
current_period_end TIMESTAMPTZ NOT NULL,
|
||||||
|
trial_end TIMESTAMPTZ,
|
||||||
|
canceled_at TIMESTAMPTZ,
|
||||||
|
cancel_at_period_end BOOLEAN NOT NULL DEFAULT FALSE,
|
||||||
|
metadata JSONB NOT NULL DEFAULT '{}',
|
||||||
|
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
||||||
|
updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
||||||
|
FOREIGN KEY (account_id) REFERENCES accounts(id) ON DELETE CASCADE,
|
||||||
|
FOREIGN KEY (plan_id) REFERENCES billing_plans(id)
|
||||||
|
);
|
||||||
|
CREATE INDEX IF NOT EXISTS idx_billing_sub_account ON billing_subscriptions(account_id);
|
||||||
|
CREATE UNIQUE INDEX IF NOT EXISTS idx_billing_sub_active
|
||||||
|
ON billing_subscriptions(account_id)
|
||||||
|
WHERE status IN ('trial', 'active', 'past_due');
|
||||||
|
|
||||||
|
-- Invoices
|
||||||
|
CREATE TABLE IF NOT EXISTS billing_invoices (
|
||||||
|
id TEXT PRIMARY KEY,
|
||||||
|
account_id TEXT NOT NULL,
|
||||||
|
subscription_id TEXT,
|
||||||
|
plan_id TEXT,
|
||||||
|
amount_cents INTEGER NOT NULL,
|
||||||
|
currency TEXT NOT NULL DEFAULT 'CNY',
|
||||||
|
description TEXT,
|
||||||
|
status TEXT NOT NULL DEFAULT 'pending',
|
||||||
|
due_at TIMESTAMPTZ,
|
||||||
|
paid_at TIMESTAMPTZ,
|
||||||
|
voided_at TIMESTAMPTZ,
|
||||||
|
metadata JSONB NOT NULL DEFAULT '{}',
|
||||||
|
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
||||||
|
updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
||||||
|
FOREIGN KEY (account_id) REFERENCES accounts(id) ON DELETE CASCADE,
|
||||||
|
FOREIGN KEY (subscription_id) REFERENCES billing_subscriptions(id) ON DELETE SET NULL,
|
||||||
|
FOREIGN KEY (plan_id) REFERENCES billing_plans(id)
|
||||||
|
);
|
||||||
|
CREATE INDEX IF NOT EXISTS idx_billing_inv_account ON billing_invoices(account_id);
|
||||||
|
CREATE INDEX IF NOT EXISTS idx_billing_inv_status ON billing_invoices(status);
|
||||||
|
CREATE INDEX IF NOT EXISTS idx_billing_inv_time ON billing_invoices(created_at);
|
||||||
|
|
||||||
|
-- Payment records (Alipay / WeChat Pay)
|
||||||
|
CREATE TABLE IF NOT EXISTS billing_payments (
|
||||||
|
id TEXT PRIMARY KEY,
|
||||||
|
invoice_id TEXT NOT NULL,
|
||||||
|
account_id TEXT NOT NULL,
|
||||||
|
amount_cents INTEGER NOT NULL,
|
||||||
|
currency TEXT NOT NULL DEFAULT 'CNY',
|
||||||
|
method TEXT NOT NULL,
|
||||||
|
status TEXT NOT NULL DEFAULT 'pending',
|
||||||
|
external_trade_no TEXT,
|
||||||
|
paid_at TIMESTAMPTZ,
|
||||||
|
refunded_at TIMESTAMPTZ,
|
||||||
|
failure_reason TEXT,
|
||||||
|
metadata JSONB NOT NULL DEFAULT '{}',
|
||||||
|
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
||||||
|
updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
||||||
|
FOREIGN KEY (invoice_id) REFERENCES billing_invoices(id) ON DELETE CASCADE,
|
||||||
|
FOREIGN KEY (account_id) REFERENCES accounts(id) ON DELETE CASCADE
|
||||||
|
);
|
||||||
|
CREATE INDEX IF NOT EXISTS idx_billing_pay_invoice ON billing_payments(invoice_id);
|
||||||
|
CREATE INDEX IF NOT EXISTS idx_billing_pay_account ON billing_payments(account_id);
|
||||||
|
CREATE INDEX IF NOT EXISTS idx_billing_pay_trade_no ON billing_payments(external_trade_no);
|
||||||
|
CREATE INDEX IF NOT EXISTS idx_billing_pay_status ON billing_payments(status);
|
||||||
|
|
||||||
|
-- Monthly usage quotas (per account per billing period)
|
||||||
|
CREATE TABLE IF NOT EXISTS billing_usage_quotas (
|
||||||
|
id TEXT PRIMARY KEY,
|
||||||
|
account_id TEXT NOT NULL,
|
||||||
|
period_start TIMESTAMPTZ NOT NULL,
|
||||||
|
period_end TIMESTAMPTZ NOT NULL,
|
||||||
|
input_tokens BIGINT NOT NULL DEFAULT 0,
|
||||||
|
output_tokens BIGINT NOT NULL DEFAULT 0,
|
||||||
|
relay_requests INTEGER NOT NULL DEFAULT 0,
|
||||||
|
hand_executions INTEGER NOT NULL DEFAULT 0,
|
||||||
|
pipeline_runs INTEGER NOT NULL DEFAULT 0,
|
||||||
|
max_input_tokens BIGINT,
|
||||||
|
max_output_tokens BIGINT,
|
||||||
|
max_relay_requests INTEGER,
|
||||||
|
max_hand_executions INTEGER,
|
||||||
|
max_pipeline_runs INTEGER,
|
||||||
|
metadata JSONB NOT NULL DEFAULT '{}',
|
||||||
|
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
||||||
|
updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
||||||
|
FOREIGN KEY (account_id) REFERENCES accounts(id) ON DELETE CASCADE,
|
||||||
|
UNIQUE(account_id, period_start)
|
||||||
|
);
|
||||||
|
CREATE INDEX IF NOT EXISTS idx_billing_usage_account ON billing_usage_quotas(account_id);
|
||||||
|
CREATE INDEX IF NOT EXISTS idx_billing_usage_period ON billing_usage_quotas(period_start, period_end);
|
||||||
|
|
||||||
|
-- Seed: default plans
|
||||||
|
INSERT INTO billing_plans (id, name, display_name, description, price_cents, interval, features, limits, is_default, sort_order)
|
||||||
|
VALUES
|
||||||
|
('plan-free', 'free', '免费版', '基础功能,适合个人体验', 0, 'month',
|
||||||
|
'{"hands": ["browser", "collector", "researcher"], "chat_modes": ["flash", "thinking"], "pipelines": 3, "support": "community"}'::jsonb,
|
||||||
|
'{"max_input_tokens_monthly": 500000, "max_output_tokens_monthly": 500000, "max_relay_requests_monthly": 100, "max_hand_executions_monthly": 20, "max_pipeline_runs_monthly": 5}'::jsonb,
|
||||||
|
TRUE, 0),
|
||||||
|
('plan-pro', 'pro', '专业版', '全功能解锁,适合知识工作者', 4900, 'month',
|
||||||
|
'{"hands": "all", "chat_modes": "all", "pipelines": -1, "support": "priority", "memory": true, "export": true}'::jsonb,
|
||||||
|
'{"max_input_tokens_monthly": 5000000, "max_output_tokens_monthly": 5000000, "max_relay_requests_monthly": 2000, "max_hand_executions_monthly": 200, "max_pipeline_runs_monthly": 100}'::jsonb,
|
||||||
|
FALSE, 1),
|
||||||
|
('plan-team', 'team', '团队版', '多席位协作,适合企业团队', 19900, 'month',
|
||||||
|
'{"hands": "all", "chat_modes": "all", "pipelines": -1, "support": "dedicated", "memory": true, "export": true, "sharing": true, "admin": true}'::jsonb,
|
||||||
|
'{"max_input_tokens_monthly": 50000000, "max_output_tokens_monthly": 50000000, "max_relay_requests_monthly": 20000, "max_hand_executions_monthly": 1000, "max_pipeline_runs_monthly": 500}'::jsonb,
|
||||||
|
FALSE, 2)
|
||||||
|
ON CONFLICT (name) DO NOTHING;
|
||||||
55
crates/zclaw-saas/src/billing/handlers.rs
Normal file
55
crates/zclaw-saas/src/billing/handlers.rs
Normal file
@@ -0,0 +1,55 @@
|
|||||||
|
//! 计费 HTTP 处理器
|
||||||
|
|
||||||
|
use axum::{
|
||||||
|
extract::{Extension, Path, State},
|
||||||
|
Json,
|
||||||
|
};
|
||||||
|
|
||||||
|
use crate::auth::types::AuthContext;
|
||||||
|
use crate::error::SaasResult;
|
||||||
|
use crate::state::AppState;
|
||||||
|
use super::service;
|
||||||
|
use super::types::*;
|
||||||
|
|
||||||
|
/// GET /api/v1/billing/plans — 列出所有活跃计划
|
||||||
|
pub async fn list_plans(
|
||||||
|
State(state): State<AppState>,
|
||||||
|
) -> SaasResult<Json<Vec<BillingPlan>>> {
|
||||||
|
let plans = service::list_plans(&state.db).await?;
|
||||||
|
Ok(Json(plans))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// GET /api/v1/billing/plans/:id — 获取单个计划详情
|
||||||
|
pub async fn get_plan(
|
||||||
|
State(state): State<AppState>,
|
||||||
|
Path(plan_id): Path<String>,
|
||||||
|
) -> SaasResult<Json<BillingPlan>> {
|
||||||
|
let plan = service::get_plan(&state.db, &plan_id).await?
|
||||||
|
.ok_or_else(|| crate::error::SaasError::NotFound("计划不存在".into()))?;
|
||||||
|
Ok(Json(plan))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// GET /api/v1/billing/subscription — 获取当前订阅
|
||||||
|
pub async fn get_subscription(
|
||||||
|
State(state): State<AppState>,
|
||||||
|
Extension(ctx): Extension<AuthContext>,
|
||||||
|
) -> SaasResult<Json<serde_json::Value>> {
|
||||||
|
let plan = service::get_account_plan(&state.db, &ctx.account_id).await?;
|
||||||
|
let sub = service::get_active_subscription(&state.db, &ctx.account_id).await?;
|
||||||
|
let usage = service::get_or_create_usage(&state.db, &ctx.account_id).await?;
|
||||||
|
|
||||||
|
Ok(Json(serde_json::json!({
|
||||||
|
"plan": plan,
|
||||||
|
"subscription": sub,
|
||||||
|
"usage": usage,
|
||||||
|
})))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// GET /api/v1/billing/usage — 获取当月用量
|
||||||
|
pub async fn get_usage(
|
||||||
|
State(state): State<AppState>,
|
||||||
|
Extension(ctx): Extension<AuthContext>,
|
||||||
|
) -> SaasResult<Json<UsageQuota>> {
|
||||||
|
let usage = service::get_or_create_usage(&state.db, &ctx.account_id).await?;
|
||||||
|
Ok(Json(usage))
|
||||||
|
}
|
||||||
15
crates/zclaw-saas/src/billing/mod.rs
Normal file
15
crates/zclaw-saas/src/billing/mod.rs
Normal file
@@ -0,0 +1,15 @@
|
|||||||
|
//! 计费模块 — 计划管理、订阅、用量配额
|
||||||
|
|
||||||
|
pub mod types;
|
||||||
|
pub mod service;
|
||||||
|
pub mod handlers;
|
||||||
|
|
||||||
|
use axum::routing::get;
|
||||||
|
|
||||||
|
pub fn routes() -> axum::Router<crate::state::AppState> {
|
||||||
|
axum::Router::new()
|
||||||
|
.route("/api/v1/billing/plans", get(handlers::list_plans))
|
||||||
|
.route("/api/v1/billing/plans/{id}", get(handlers::get_plan))
|
||||||
|
.route("/api/v1/billing/subscription", get(handlers::get_subscription))
|
||||||
|
.route("/api/v1/billing/usage", get(handlers::get_usage))
|
||||||
|
}
|
||||||
206
crates/zclaw-saas/src/billing/service.rs
Normal file
206
crates/zclaw-saas/src/billing/service.rs
Normal file
@@ -0,0 +1,206 @@
|
|||||||
|
//! 计费服务层 — 计划查询、订阅管理、用量检查
|
||||||
|
|
||||||
|
use chrono::{Datelike, Timelike};
|
||||||
|
use sqlx::PgPool;
|
||||||
|
|
||||||
|
use crate::error::SaasResult;
|
||||||
|
|
||||||
|
use super::types::*;
|
||||||
|
|
||||||
|
/// 获取所有活跃计划
|
||||||
|
pub async fn list_plans(pool: &PgPool) -> SaasResult<Vec<BillingPlan>> {
|
||||||
|
let plans = sqlx::query_as::<_, BillingPlan>(
|
||||||
|
"SELECT * FROM billing_plans WHERE status = 'active' ORDER BY sort_order"
|
||||||
|
)
|
||||||
|
.fetch_all(pool)
|
||||||
|
.await?;
|
||||||
|
Ok(plans)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 获取单个计划
|
||||||
|
pub async fn get_plan(pool: &PgPool, plan_id: &str) -> SaasResult<Option<BillingPlan>> {
|
||||||
|
let plan = sqlx::query_as::<_, BillingPlan>(
|
||||||
|
"SELECT * FROM billing_plans WHERE id = $1"
|
||||||
|
)
|
||||||
|
.bind(plan_id)
|
||||||
|
.fetch_optional(pool)
|
||||||
|
.await?;
|
||||||
|
Ok(plan)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 获取账户当前有效订阅
|
||||||
|
pub async fn get_active_subscription(
|
||||||
|
pool: &PgPool,
|
||||||
|
account_id: &str,
|
||||||
|
) -> SaasResult<Option<Subscription>> {
|
||||||
|
let sub = sqlx::query_as::<_, Subscription>(
|
||||||
|
"SELECT * FROM billing_subscriptions \
|
||||||
|
WHERE account_id = $1 AND status IN ('trial', 'active', 'past_due') \
|
||||||
|
ORDER BY created_at DESC LIMIT 1"
|
||||||
|
)
|
||||||
|
.bind(account_id)
|
||||||
|
.fetch_optional(pool)
|
||||||
|
.await?;
|
||||||
|
Ok(sub)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 获取账户当前计划(有订阅返回订阅计划,否则返回 Free)
|
||||||
|
pub async fn get_account_plan(pool: &PgPool, account_id: &str) -> SaasResult<BillingPlan> {
|
||||||
|
if let Some(sub) = get_active_subscription(pool, account_id).await? {
|
||||||
|
if let Some(plan) = get_plan(pool, &sub.plan_id).await? {
|
||||||
|
return Ok(plan);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
// 回退到 Free 计划
|
||||||
|
let free = sqlx::query_as::<_, BillingPlan>(
|
||||||
|
"SELECT * FROM billing_plans WHERE name = 'free' AND status = 'active' LIMIT 1"
|
||||||
|
)
|
||||||
|
.fetch_optional(pool)
|
||||||
|
.await?;
|
||||||
|
Ok(free.unwrap_or_else(|| BillingPlan {
|
||||||
|
id: "plan-free".into(),
|
||||||
|
name: "free".into(),
|
||||||
|
display_name: "免费版".into(),
|
||||||
|
description: Some("基础功能".into()),
|
||||||
|
price_cents: 0,
|
||||||
|
currency: "CNY".into(),
|
||||||
|
interval: "month".into(),
|
||||||
|
features: serde_json::json!({}),
|
||||||
|
limits: serde_json::json!({
|
||||||
|
"max_input_tokens_monthly": 500000,
|
||||||
|
"max_output_tokens_monthly": 500000,
|
||||||
|
"max_relay_requests_monthly": 100,
|
||||||
|
"max_hand_executions_monthly": 20,
|
||||||
|
"max_pipeline_runs_monthly": 5,
|
||||||
|
}),
|
||||||
|
is_default: true,
|
||||||
|
sort_order: 0,
|
||||||
|
status: "active".into(),
|
||||||
|
created_at: chrono::Utc::now(),
|
||||||
|
updated_at: chrono::Utc::now(),
|
||||||
|
}))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 获取或创建当月用量记录
|
||||||
|
pub async fn get_or_create_usage(pool: &PgPool, account_id: &str) -> SaasResult<UsageQuota> {
|
||||||
|
let now = chrono::Utc::now();
|
||||||
|
let period_start = now
|
||||||
|
.with_day(1).unwrap_or(now)
|
||||||
|
.with_hour(0).unwrap_or(now)
|
||||||
|
.with_minute(0).unwrap_or(now)
|
||||||
|
.with_second(0).unwrap_or(now)
|
||||||
|
.with_nanosecond(0).unwrap_or(now);
|
||||||
|
|
||||||
|
// 尝试获取现有记录
|
||||||
|
let existing = sqlx::query_as::<_, UsageQuota>(
|
||||||
|
"SELECT * FROM billing_usage_quotas \
|
||||||
|
WHERE account_id = $1 AND period_start = $2"
|
||||||
|
)
|
||||||
|
.bind(account_id)
|
||||||
|
.bind(period_start)
|
||||||
|
.fetch_optional(pool)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
if let Some(usage) = existing {
|
||||||
|
return Ok(usage);
|
||||||
|
}
|
||||||
|
|
||||||
|
// 获取当前计划限额
|
||||||
|
let plan = get_account_plan(pool, account_id).await?;
|
||||||
|
let limits: PlanLimits = serde_json::from_value(plan.limits.clone())
|
||||||
|
.unwrap_or_else(|_| PlanLimits::free());
|
||||||
|
|
||||||
|
// 计算月末
|
||||||
|
let period_end = if now.month() == 12 {
|
||||||
|
now.with_year(now.year() + 1).and_then(|d| d.with_month(1))
|
||||||
|
} else {
|
||||||
|
now.with_month(now.month() + 1)
|
||||||
|
}.unwrap_or(now)
|
||||||
|
.with_day(1).unwrap_or(now)
|
||||||
|
.with_hour(0).unwrap_or(now)
|
||||||
|
.with_minute(0).unwrap_or(now)
|
||||||
|
.with_second(0).unwrap_or(now)
|
||||||
|
.with_nanosecond(0).unwrap_or(now);
|
||||||
|
|
||||||
|
let id = uuid::Uuid::new_v4().to_string();
|
||||||
|
let usage = sqlx::query_as::<_, UsageQuota>(
|
||||||
|
"INSERT INTO billing_usage_quotas \
|
||||||
|
(id, account_id, period_start, period_end, \
|
||||||
|
max_input_tokens, max_output_tokens, max_relay_requests, \
|
||||||
|
max_hand_executions, max_pipeline_runs) \
|
||||||
|
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9) \
|
||||||
|
RETURNING *"
|
||||||
|
)
|
||||||
|
.bind(&id)
|
||||||
|
.bind(account_id)
|
||||||
|
.bind(period_start)
|
||||||
|
.bind(period_end)
|
||||||
|
.bind(limits.max_input_tokens_monthly)
|
||||||
|
.bind(limits.max_output_tokens_monthly)
|
||||||
|
.bind(limits.max_relay_requests_monthly)
|
||||||
|
.bind(limits.max_hand_executions_monthly)
|
||||||
|
.bind(limits.max_pipeline_runs_monthly)
|
||||||
|
.fetch_one(pool)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
Ok(usage)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 增加用量计数
|
||||||
|
pub async fn increment_usage(
|
||||||
|
pool: &PgPool,
|
||||||
|
account_id: &str,
|
||||||
|
input_tokens: i64,
|
||||||
|
output_tokens: i64,
|
||||||
|
) -> SaasResult<()> {
|
||||||
|
let usage = get_or_create_usage(pool, account_id).await?;
|
||||||
|
sqlx::query(
|
||||||
|
"UPDATE billing_usage_quotas \
|
||||||
|
SET input_tokens = input_tokens + $1, \
|
||||||
|
output_tokens = output_tokens + $2, \
|
||||||
|
relay_requests = relay_requests + 1, \
|
||||||
|
updated_at = NOW() \
|
||||||
|
WHERE id = $3"
|
||||||
|
)
|
||||||
|
.bind(input_tokens)
|
||||||
|
.bind(output_tokens)
|
||||||
|
.bind(&usage.id)
|
||||||
|
.execute(pool)
|
||||||
|
.await?;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 检查用量配额
|
||||||
|
pub async fn check_quota(
|
||||||
|
pool: &PgPool,
|
||||||
|
account_id: &str,
|
||||||
|
quota_type: &str,
|
||||||
|
) -> SaasResult<QuotaCheck> {
|
||||||
|
let usage = get_or_create_usage(pool, account_id).await?;
|
||||||
|
|
||||||
|
let (current, limit) = match quota_type {
|
||||||
|
"input_tokens" => (usage.input_tokens, usage.max_input_tokens),
|
||||||
|
"output_tokens" => (usage.output_tokens, usage.max_output_tokens),
|
||||||
|
"relay_requests" => (usage.relay_requests as i64, usage.max_relay_requests.map(|v| v as i64)),
|
||||||
|
"hand_executions" => (usage.hand_executions as i64, usage.max_hand_executions.map(|v| v as i64)),
|
||||||
|
"pipeline_runs" => (usage.pipeline_runs as i64, usage.max_pipeline_runs.map(|v| v as i64)),
|
||||||
|
_ => return Ok(QuotaCheck {
|
||||||
|
allowed: true,
|
||||||
|
reason: None,
|
||||||
|
current: 0,
|
||||||
|
limit: None,
|
||||||
|
remaining: None,
|
||||||
|
}),
|
||||||
|
};
|
||||||
|
|
||||||
|
let allowed = limit.map_or(true, |lim| current < lim);
|
||||||
|
let remaining = limit.map(|lim| (lim - current).max(0));
|
||||||
|
|
||||||
|
Ok(QuotaCheck {
|
||||||
|
allowed,
|
||||||
|
reason: if !allowed { Some(format!("{} 配额已用尽", quota_type)) } else { None },
|
||||||
|
current,
|
||||||
|
limit,
|
||||||
|
remaining,
|
||||||
|
})
|
||||||
|
}
|
||||||
161
crates/zclaw-saas/src/billing/types.rs
Normal file
161
crates/zclaw-saas/src/billing/types.rs
Normal file
@@ -0,0 +1,161 @@
|
|||||||
|
//! 计费类型定义
|
||||||
|
|
||||||
|
use chrono::{DateTime, Utc};
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
|
||||||
|
/// 计费计划定义 — 对应 billing_plans 表
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize, sqlx::FromRow)]
|
||||||
|
pub struct BillingPlan {
|
||||||
|
pub id: String,
|
||||||
|
pub name: String,
|
||||||
|
pub display_name: String,
|
||||||
|
pub description: Option<String>,
|
||||||
|
pub price_cents: i32,
|
||||||
|
pub currency: String,
|
||||||
|
pub interval: String,
|
||||||
|
pub features: serde_json::Value,
|
||||||
|
pub limits: serde_json::Value,
|
||||||
|
pub is_default: bool,
|
||||||
|
pub sort_order: i32,
|
||||||
|
pub status: String,
|
||||||
|
pub created_at: DateTime<Utc>,
|
||||||
|
pub updated_at: DateTime<Utc>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 计划限额(从 limits JSON 反序列化)
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct PlanLimits {
|
||||||
|
#[serde(default)]
|
||||||
|
pub max_input_tokens_monthly: Option<i64>,
|
||||||
|
#[serde(default)]
|
||||||
|
pub max_output_tokens_monthly: Option<i64>,
|
||||||
|
#[serde(default)]
|
||||||
|
pub max_relay_requests_monthly: Option<i32>,
|
||||||
|
#[serde(default)]
|
||||||
|
pub max_hand_executions_monthly: Option<i32>,
|
||||||
|
#[serde(default)]
|
||||||
|
pub max_pipeline_runs_monthly: Option<i32>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl PlanLimits {
|
||||||
|
pub fn free() -> Self {
|
||||||
|
Self {
|
||||||
|
max_input_tokens_monthly: Some(500_000),
|
||||||
|
max_output_tokens_monthly: Some(500_000),
|
||||||
|
max_relay_requests_monthly: Some(100),
|
||||||
|
max_hand_executions_monthly: Some(20),
|
||||||
|
max_pipeline_runs_monthly: Some(5),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 账户订阅 — 对应 billing_subscriptions 表
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize, sqlx::FromRow)]
|
||||||
|
pub struct Subscription {
|
||||||
|
pub id: String,
|
||||||
|
pub account_id: String,
|
||||||
|
pub plan_id: String,
|
||||||
|
pub status: String,
|
||||||
|
pub current_period_start: DateTime<Utc>,
|
||||||
|
pub current_period_end: DateTime<Utc>,
|
||||||
|
pub trial_end: Option<DateTime<Utc>>,
|
||||||
|
pub canceled_at: Option<DateTime<Utc>>,
|
||||||
|
pub cancel_at_period_end: bool,
|
||||||
|
pub metadata: serde_json::Value,
|
||||||
|
pub created_at: DateTime<Utc>,
|
||||||
|
pub updated_at: DateTime<Utc>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 发票 — 对应 billing_invoices 表
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize, sqlx::FromRow)]
|
||||||
|
pub struct Invoice {
|
||||||
|
pub id: String,
|
||||||
|
pub account_id: String,
|
||||||
|
pub subscription_id: Option<String>,
|
||||||
|
pub plan_id: Option<String>,
|
||||||
|
pub amount_cents: i32,
|
||||||
|
pub currency: String,
|
||||||
|
pub description: Option<String>,
|
||||||
|
pub status: String,
|
||||||
|
pub due_at: Option<DateTime<Utc>>,
|
||||||
|
pub paid_at: Option<DateTime<Utc>>,
|
||||||
|
pub voided_at: Option<DateTime<Utc>>,
|
||||||
|
pub metadata: serde_json::Value,
|
||||||
|
pub created_at: DateTime<Utc>,
|
||||||
|
pub updated_at: DateTime<Utc>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 支付记录 — 对应 billing_payments 表
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize, sqlx::FromRow)]
|
||||||
|
pub struct Payment {
|
||||||
|
pub id: String,
|
||||||
|
pub invoice_id: String,
|
||||||
|
pub account_id: String,
|
||||||
|
pub amount_cents: i32,
|
||||||
|
pub currency: String,
|
||||||
|
pub method: String,
|
||||||
|
pub status: String,
|
||||||
|
pub external_trade_no: Option<String>,
|
||||||
|
pub paid_at: Option<DateTime<Utc>>,
|
||||||
|
pub refunded_at: Option<DateTime<Utc>>,
|
||||||
|
pub failure_reason: Option<String>,
|
||||||
|
pub metadata: serde_json::Value,
|
||||||
|
pub created_at: DateTime<Utc>,
|
||||||
|
pub updated_at: DateTime<Utc>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 月度用量配额 — 对应 billing_usage_quotas 表
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize, sqlx::FromRow)]
|
||||||
|
pub struct UsageQuota {
|
||||||
|
pub id: String,
|
||||||
|
pub account_id: String,
|
||||||
|
pub period_start: DateTime<Utc>,
|
||||||
|
pub period_end: DateTime<Utc>,
|
||||||
|
pub input_tokens: i64,
|
||||||
|
pub output_tokens: i64,
|
||||||
|
pub relay_requests: i32,
|
||||||
|
pub hand_executions: i32,
|
||||||
|
pub pipeline_runs: i32,
|
||||||
|
pub max_input_tokens: Option<i64>,
|
||||||
|
pub max_output_tokens: Option<i64>,
|
||||||
|
pub max_relay_requests: Option<i32>,
|
||||||
|
pub max_hand_executions: Option<i32>,
|
||||||
|
pub max_pipeline_runs: Option<i32>,
|
||||||
|
pub metadata: serde_json::Value,
|
||||||
|
pub created_at: DateTime<Utc>,
|
||||||
|
pub updated_at: DateTime<Utc>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 用量检查结果
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct QuotaCheck {
|
||||||
|
pub allowed: bool,
|
||||||
|
pub reason: Option<String>,
|
||||||
|
pub current: i64,
|
||||||
|
pub limit: Option<i64>,
|
||||||
|
pub remaining: Option<i64>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 支付方式
|
||||||
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||||
|
#[serde(rename_all = "snake_case")]
|
||||||
|
pub enum PaymentMethod {
|
||||||
|
Alipay,
|
||||||
|
Wechat,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 创建支付请求
|
||||||
|
#[derive(Debug, Deserialize)]
|
||||||
|
pub struct CreatePaymentRequest {
|
||||||
|
pub plan_id: String,
|
||||||
|
pub payment_method: PaymentMethod,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 支付结果
|
||||||
|
#[derive(Debug, Serialize)]
|
||||||
|
pub struct PaymentResult {
|
||||||
|
pub payment_id: String,
|
||||||
|
pub trade_no: String,
|
||||||
|
pub pay_url: String,
|
||||||
|
pub amount_cents: i32,
|
||||||
|
}
|
||||||
@@ -2,34 +2,44 @@
|
|||||||
|
|
||||||
use sqlx::postgres::PgPoolOptions;
|
use sqlx::postgres::PgPoolOptions;
|
||||||
use sqlx::PgPool;
|
use sqlx::PgPool;
|
||||||
|
use crate::config::DatabaseConfig;
|
||||||
use crate::error::SaasResult;
|
use crate::error::SaasResult;
|
||||||
|
|
||||||
const SCHEMA_VERSION: i32 = 11;
|
const SCHEMA_VERSION: i32 = 12;
|
||||||
|
|
||||||
/// 初始化数据库
|
/// 初始化数据库
|
||||||
pub async fn init_db(database_url: &str) -> SaasResult<PgPool> {
|
pub async fn init_db(config: &DatabaseConfig) -> SaasResult<PgPool> {
|
||||||
// 连接池大小可通过环境变量配置,默认 100(relay 请求每次 10+ 串行查询,50 偏紧)
|
// 环境变量覆盖 URL(避免在配置文件中存储密码)
|
||||||
|
let database_url = std::env::var("ZCLAW_DATABASE_URL")
|
||||||
|
.unwrap_or_else(|_| config.url.clone());
|
||||||
|
|
||||||
|
// 环境变量覆盖连接数(向后兼容)
|
||||||
let max_connections: u32 = std::env::var("ZCLAW_DB_MAX_CONNECTIONS")
|
let max_connections: u32 = std::env::var("ZCLAW_DB_MAX_CONNECTIONS")
|
||||||
.ok()
|
.ok()
|
||||||
.and_then(|v| v.parse().ok())
|
.and_then(|v| v.parse().ok())
|
||||||
.unwrap_or(100);
|
.unwrap_or(config.max_connections);
|
||||||
let min_connections: u32 = std::env::var("ZCLAW_DB_MIN_CONNECTIONS")
|
let min_connections: u32 = std::env::var("ZCLAW_DB_MIN_CONNECTIONS")
|
||||||
.ok()
|
.ok()
|
||||||
.and_then(|v| v.parse().ok())
|
.and_then(|v| v.parse().ok())
|
||||||
.unwrap_or(5);
|
.unwrap_or(config.min_connections);
|
||||||
|
|
||||||
tracing::info!("Database pool: max={}, min={}", max_connections, min_connections);
|
tracing::info!(
|
||||||
|
"Database pool: max={}, min={}, acquire_timeout={}s, idle_timeout={}s, max_lifetime={}s",
|
||||||
|
max_connections, min_connections,
|
||||||
|
config.acquire_timeout_secs, config.idle_timeout_secs, config.max_lifetime_secs
|
||||||
|
);
|
||||||
|
|
||||||
let pool = PgPoolOptions::new()
|
let pool = PgPoolOptions::new()
|
||||||
.max_connections(max_connections)
|
.max_connections(max_connections)
|
||||||
.min_connections(min_connections)
|
.min_connections(min_connections)
|
||||||
.acquire_timeout(std::time::Duration::from_secs(8))
|
.acquire_timeout(std::time::Duration::from_secs(config.acquire_timeout_secs))
|
||||||
.idle_timeout(std::time::Duration::from_secs(180))
|
.idle_timeout(std::time::Duration::from_secs(config.idle_timeout_secs))
|
||||||
.max_lifetime(std::time::Duration::from_secs(900))
|
.max_lifetime(std::time::Duration::from_secs(config.max_lifetime_secs))
|
||||||
.connect(database_url)
|
.connect(&database_url)
|
||||||
.await?;
|
.await?;
|
||||||
|
|
||||||
run_migrations(&pool).await?;
|
run_migrations(&pool).await?;
|
||||||
|
ensure_security_columns(&pool).await?;
|
||||||
seed_admin_account(&pool).await?;
|
seed_admin_account(&pool).await?;
|
||||||
seed_builtin_prompts(&pool).await?;
|
seed_builtin_prompts(&pool).await?;
|
||||||
seed_demo_data(&pool).await?;
|
seed_demo_data(&pool).await?;
|
||||||
@@ -884,6 +894,56 @@ async fn fix_seed_data(pool: &PgPool) -> SaasResult<()> {
|
|||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// 防御性检查:确保安全审计新增的列存在(即使 schema_version 显示已是最新)
|
||||||
|
///
|
||||||
|
/// 场景:旧数据库的 schema_version 已被手动更新但迁移文件未实际执行,
|
||||||
|
/// 或者迁移文件在 version check 时被跳过。
|
||||||
|
async fn ensure_security_columns(pool: &PgPool) -> SaasResult<()> {
|
||||||
|
// 检查 password_version 列是否存在
|
||||||
|
let col_exists: bool = sqlx::query_scalar(
|
||||||
|
"SELECT EXISTS (SELECT 1 FROM information_schema.columns WHERE table_name = 'accounts' AND column_name = 'password_version')"
|
||||||
|
)
|
||||||
|
.fetch_one(pool)
|
||||||
|
.await
|
||||||
|
.unwrap_or(false);
|
||||||
|
|
||||||
|
if !col_exists {
|
||||||
|
tracing::warn!("[DB] 'password_version' column missing — applying security fix migration");
|
||||||
|
sqlx::query("ALTER TABLE accounts ADD COLUMN IF NOT EXISTS password_version INTEGER NOT NULL DEFAULT 1")
|
||||||
|
.execute(pool).await?;
|
||||||
|
sqlx::query("ALTER TABLE accounts ADD COLUMN IF NOT EXISTS failed_login_count INTEGER NOT NULL DEFAULT 0")
|
||||||
|
.execute(pool).await?;
|
||||||
|
sqlx::query("ALTER TABLE accounts ADD COLUMN IF NOT EXISTS locked_until TIMESTAMPTZ")
|
||||||
|
.execute(pool).await?;
|
||||||
|
tracing::info!("[DB] Security columns (password_version, failed_login_count, locked_until) applied");
|
||||||
|
}
|
||||||
|
|
||||||
|
// 检查 rate_limit_events 表是否存在
|
||||||
|
let table_exists: bool = sqlx::query_scalar(
|
||||||
|
"SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_name = 'rate_limit_events')"
|
||||||
|
)
|
||||||
|
.fetch_one(pool)
|
||||||
|
.await
|
||||||
|
.unwrap_or(false);
|
||||||
|
|
||||||
|
if !table_exists {
|
||||||
|
tracing::warn!("[DB] 'rate_limit_events' table missing — applying rate limit migration");
|
||||||
|
if let Err(e) = sqlx::query(
|
||||||
|
"CREATE TABLE IF NOT EXISTS rate_limit_events (
|
||||||
|
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
||||||
|
key TEXT NOT NULL,
|
||||||
|
count BIGINT NOT NULL DEFAULT 1,
|
||||||
|
window_start TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
||||||
|
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
|
||||||
|
)"
|
||||||
|
).execute(pool).await {
|
||||||
|
tracing::warn!("[DB] Failed to create rate_limit_events: {}", e);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
// PostgreSQL 单元测试需要真实数据库连接,此处保留接口兼容
|
// PostgreSQL 单元测试需要真实数据库连接,此处保留接口兼容
|
||||||
|
|||||||
@@ -25,3 +25,4 @@ pub mod prompt;
|
|||||||
pub mod agent_template;
|
pub mod agent_template;
|
||||||
pub mod scheduled_task;
|
pub mod scheduled_task;
|
||||||
pub mod telemetry;
|
pub mod telemetry;
|
||||||
|
pub mod billing;
|
||||||
|
|||||||
@@ -14,6 +14,9 @@ use zclaw_saas::workers::update_last_used::UpdateLastUsedWorker;
|
|||||||
|
|
||||||
#[tokio::main]
|
#[tokio::main]
|
||||||
async fn main() -> anyhow::Result<()> {
|
async fn main() -> anyhow::Result<()> {
|
||||||
|
// Load .env file from project root (walk up from current dir)
|
||||||
|
load_dotenv();
|
||||||
|
|
||||||
tracing_subscriber::fmt()
|
tracing_subscriber::fmt()
|
||||||
.with_env_filter(
|
.with_env_filter(
|
||||||
tracing_subscriber::EnvFilter::try_from_default_env()
|
tracing_subscriber::EnvFilter::try_from_default_env()
|
||||||
@@ -24,11 +27,18 @@ async fn main() -> anyhow::Result<()> {
|
|||||||
let config = SaaSConfig::load()?;
|
let config = SaaSConfig::load()?;
|
||||||
info!("SaaS config loaded: {}:{}", config.server.host, config.server.port);
|
info!("SaaS config loaded: {}:{}", config.server.host, config.server.port);
|
||||||
|
|
||||||
let db = init_db(&config.database.url).await?;
|
let db = init_db(&config.database).await?;
|
||||||
info!("Database initialized");
|
info!("Database initialized");
|
||||||
|
|
||||||
|
// 创建 Worker spawn 限制器(门控并发 DB 操作数量)
|
||||||
|
let worker_limiter = zclaw_saas::state::SpawnLimiter::new(
|
||||||
|
"worker",
|
||||||
|
config.database.worker_concurrency,
|
||||||
|
);
|
||||||
|
info!("Worker spawn limiter: {} permits", config.database.worker_concurrency);
|
||||||
|
|
||||||
// 初始化 Worker 调度器 + 注册所有 Worker
|
// 初始化 Worker 调度器 + 注册所有 Worker
|
||||||
let mut dispatcher = WorkerDispatcher::new(db.clone());
|
let mut dispatcher = WorkerDispatcher::new(db.clone(), worker_limiter.clone());
|
||||||
dispatcher.register(LogOperationWorker);
|
dispatcher.register(LogOperationWorker);
|
||||||
dispatcher.register(CleanupRefreshTokensWorker);
|
dispatcher.register(CleanupRefreshTokensWorker);
|
||||||
dispatcher.register(CleanupRateLimitWorker);
|
dispatcher.register(CleanupRateLimitWorker);
|
||||||
@@ -38,12 +48,13 @@ async fn main() -> anyhow::Result<()> {
|
|||||||
|
|
||||||
// 优雅停机令牌 — 取消后所有 SSE 流和长连接立即终止
|
// 优雅停机令牌 — 取消后所有 SSE 流和长连接立即终止
|
||||||
let shutdown_token = CancellationToken::new();
|
let shutdown_token = CancellationToken::new();
|
||||||
let state = AppState::new(db.clone(), config.clone(), dispatcher, shutdown_token.clone())?;
|
let state = AppState::new(db.clone(), config.clone(), dispatcher, shutdown_token.clone(), worker_limiter.clone())?;
|
||||||
|
|
||||||
// Restore rate limit counts from DB so limits survive server restarts
|
// Restore rate limit counts from DB so limits survive server restarts
|
||||||
|
// 仅恢复最近 60s 的计数(与 middleware 的 60s 滑动窗口一致),避免过于保守的限流
|
||||||
{
|
{
|
||||||
let rows: Vec<(String, i64)> = sqlx::query_as(
|
let rows: Vec<(String, i64)> = sqlx::query_as(
|
||||||
"SELECT key, SUM(count) FROM rate_limit_events WHERE window_start > NOW() - interval '1 hour' GROUP BY key"
|
"SELECT key, SUM(count) FROM rate_limit_events WHERE window_start > NOW() - interval '60 seconds' GROUP BY key"
|
||||||
)
|
)
|
||||||
.fetch_all(&db)
|
.fetch_all(&db)
|
||||||
.await
|
.await
|
||||||
@@ -51,18 +62,17 @@ async fn main() -> anyhow::Result<()> {
|
|||||||
|
|
||||||
let mut restored_count = 0usize;
|
let mut restored_count = 0usize;
|
||||||
for (key, count) in rows {
|
for (key, count) in rows {
|
||||||
let mut entries = Vec::new();
|
// 限制恢复计数不超过 RPM 配额,避免重启后过于保守
|
||||||
// Approximate: insert count timestamps at "now" — the DashMap will
|
let rpm = state.rate_limit_rpm() as usize;
|
||||||
// expire them naturally via the retain() call in the middleware.
|
let capped = (count as usize).min(rpm);
|
||||||
// This is intentionally approximate; exact window alignment is not
|
let mut entries = Vec::with_capacity(capped);
|
||||||
// required for rate limiting correctness.
|
for _ in 0..capped {
|
||||||
for _ in 0..count as usize {
|
|
||||||
entries.push(std::time::Instant::now());
|
entries.push(std::time::Instant::now());
|
||||||
}
|
}
|
||||||
state.rate_limit_entries.insert(key, entries);
|
state.rate_limit_entries.insert(key, entries);
|
||||||
restored_count += 1;
|
restored_count += 1;
|
||||||
}
|
}
|
||||||
info!("Restored rate limit state from DB: {} keys", restored_count);
|
info!("Restored rate limit state from DB: {} keys (60s window, capped at RPM)", restored_count);
|
||||||
}
|
}
|
||||||
|
|
||||||
// 迁移旧格式 TOTP secret(明文 → 加密 enc: 格式)
|
// 迁移旧格式 TOTP secret(明文 → 加密 enc: 格式)
|
||||||
@@ -117,20 +127,64 @@ async fn main() -> anyhow::Result<()> {
|
|||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
let app = build_router(state).await;
|
// 限流事件批量 flush (可配置间隔,默认 5s)
|
||||||
|
{
|
||||||
|
let flush_state = state.clone();
|
||||||
|
let batch_interval = config.database.rate_limit_batch_interval_secs;
|
||||||
|
let batch_max = config.database.rate_limit_batch_max_size;
|
||||||
|
tokio::spawn(async move {
|
||||||
|
let mut interval = tokio::time::interval(std::time::Duration::from_secs(batch_interval));
|
||||||
|
loop {
|
||||||
|
interval.tick().await;
|
||||||
|
flush_state.flush_rate_limit_batch(batch_max).await;
|
||||||
|
}
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
// 连接池可观测性 (30s 指标日志)
|
||||||
|
{
|
||||||
|
let metrics_db = db.clone();
|
||||||
|
tokio::spawn(async move {
|
||||||
|
let mut interval = tokio::time::interval(std::time::Duration::from_secs(30));
|
||||||
|
loop {
|
||||||
|
interval.tick().await;
|
||||||
|
let pool = &metrics_db;
|
||||||
|
let total = pool.options().get_max_connections() as usize;
|
||||||
|
let idle = pool.num_idle() as usize;
|
||||||
|
let used = total.saturating_sub(idle);
|
||||||
|
let usage_pct = if total > 0 { used * 100 / total } else { 0 };
|
||||||
|
tracing::info!(
|
||||||
|
"[PoolMetrics] total={} idle={} used={} usage_pct={}%",
|
||||||
|
total, idle, used, usage_pct,
|
||||||
|
);
|
||||||
|
if usage_pct >= 80 {
|
||||||
|
tracing::warn!(
|
||||||
|
"[PoolMetrics] HIGH USAGE: {}% of connections in use!",
|
||||||
|
usage_pct,
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
let app = build_router(state.clone()).await;
|
||||||
|
|
||||||
// 配置 TCP keepalive + 短 SO_LINGER,防止 CLOSE_WAIT 累积
|
// 配置 TCP keepalive + 短 SO_LINGER,防止 CLOSE_WAIT 累积
|
||||||
let listener = create_listener(&config.server.host, config.server.port)?;
|
let listener = create_listener(&config.server.host, config.server.port)?;
|
||||||
info!("SaaS server listening on {}:{}", config.server.host, config.server.port);
|
info!("SaaS server listening on {}:{}", config.server.host, config.server.port);
|
||||||
|
|
||||||
// 优雅停机: Ctrl+C → 取消 CancellationToken → SSE 流终止 → 连接排空
|
// 优雅停机: Ctrl+C → 最终批量 flush → 取消 CancellationToken → SSE 流终止 → 连接排空
|
||||||
let token = shutdown_token.clone();
|
let token = shutdown_token.clone();
|
||||||
|
let flush_state = state;
|
||||||
|
let batch_max = config.database.rate_limit_batch_max_size;
|
||||||
axum::serve(listener, app.into_make_service_with_connect_info::<std::net::SocketAddr>())
|
axum::serve(listener, app.into_make_service_with_connect_info::<std::net::SocketAddr>())
|
||||||
.with_graceful_shutdown(async move {
|
.with_graceful_shutdown(async move {
|
||||||
tokio::signal::ctrl_c()
|
tokio::signal::ctrl_c()
|
||||||
.await
|
.await
|
||||||
.expect("Failed to install Ctrl+C handler");
|
.expect("Failed to install Ctrl+C handler");
|
||||||
info!("Received shutdown signal, cancelling SSE streams and draining connections...");
|
info!("Received shutdown signal, flushing pending rate limit batch...");
|
||||||
|
flush_state.flush_rate_limit_batch(batch_max).await;
|
||||||
|
info!("Cancelling SSE streams and draining connections...");
|
||||||
token.cancel();
|
token.cancel();
|
||||||
})
|
})
|
||||||
.await?;
|
.await?;
|
||||||
@@ -280,6 +334,7 @@ async fn build_router(state: AppState) -> axum::Router {
|
|||||||
.merge(zclaw_saas::agent_template::routes())
|
.merge(zclaw_saas::agent_template::routes())
|
||||||
.merge(zclaw_saas::scheduled_task::routes())
|
.merge(zclaw_saas::scheduled_task::routes())
|
||||||
.merge(zclaw_saas::telemetry::routes())
|
.merge(zclaw_saas::telemetry::routes())
|
||||||
|
.merge(zclaw_saas::billing::routes())
|
||||||
.layer(middleware::from_fn_with_state(
|
.layer(middleware::from_fn_with_state(
|
||||||
state.clone(),
|
state.clone(),
|
||||||
zclaw_saas::middleware::api_version_middleware,
|
zclaw_saas::middleware::api_version_middleware,
|
||||||
@@ -329,3 +384,35 @@ async fn build_router(state: AppState) -> axum::Router {
|
|||||||
.layer(cors)
|
.layer(cors)
|
||||||
.with_state(state)
|
.with_state(state)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Load `.env` file from project root by walking up from current directory.
|
||||||
|
/// Sets environment variables that are not already set (does not override).
|
||||||
|
fn load_dotenv() {
|
||||||
|
let mut dir = std::env::current_dir().unwrap_or_default();
|
||||||
|
loop {
|
||||||
|
let env_path = dir.join(".env");
|
||||||
|
if env_path.is_file() {
|
||||||
|
if let Ok(content) = std::fs::read_to_string(&env_path) {
|
||||||
|
for line in content.lines() {
|
||||||
|
let line = line.trim();
|
||||||
|
if line.is_empty() || line.starts_with('#') {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
if let Some((key, value)) = line.split_once('=') {
|
||||||
|
let key = key.trim();
|
||||||
|
let value = value.trim();
|
||||||
|
// Only set if not already defined in environment
|
||||||
|
if std::env::var(key).is_err() {
|
||||||
|
std::env::set_var(key, value);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
tracing::debug!("Loaded .env from {}", env_path.display());
|
||||||
|
}
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
if !dir.pop() {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user