From c41e795085b75bda5673fc9c72912b6a05ad197c Mon Sep 17 00:00:00 2001 From: argszero Date: Tue, 18 Aug 2026 10:12:58 +0800 Subject: [PATCH] =?UTF-8?q?feat(gateway):=20P0-B=20=E7=BD=91=E5=85=B3?= =?UTF-8?q?=E8=BD=AC=E5=8F=91=20+=20=E8=B7=AF=E7=94=B1=E6=95=85=E9=9A=9C?= =?UTF-8?q?=E8=BD=AC=E7=A7=BB=20+=20=E8=AE=A1=E9=87=8F=E8=B4=A6=E6=9C=AC?= =?UTF-8?q?=E9=97=AD=E7=8E=AF?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 网关端点：POST /v1/chat/completions（OpenAI 兼容）+ POST /anthropic/v1/messages + GET /api/models（市场页，需认证） - 路由故障转移（src/router.rs，architecture §4.2.1）：初始随机 / 粘性复用 / 401/403/429/5xx 与网络错误冷却 5s / 上限 3 次切换 → 503 - 计量账本（src/billing.rs，§4.3/4.4）：单价折算点数、消费者扣点、分享者 90% 分成、transactions consume/earn 双写、usage_records、keys.used 累计； settle 事务性回滚（幽灵 key 拒绝入账） - models 表启动种子（data/models.example.json + config.price_overrides upsert） - 测试 33 个全过（路由 6 / 计费 6 / 网关 7 + 既有 14），clippy + fmt 干净； 冒烟：无 token 401、placeholder 上游 503、/api/models 列表 Rant: 2026-08-18T09:55:57（P0-B 网关转发 + 路由故障转移 + 计量账本闭环） --- Cargo.lock | 14 +- Cargo.toml | 4 +- src/billing.rs | 300 +++++++++++++++++++++++ src/dao.rs | 86 ++++++- src/db.rs | 91 +++++++ src/gateway.rs | 588 ++++++++++++++++++++++++++++++++++++++++++++++ src/main.rs | 8 +- src/router.rs | 183 +++++++++++++++ src/routes/mod.rs | 36 ++- 9 files changed, 1295 insertions(+), 15 deletions(-) create mode 100644 src/billing.rs create mode 100644 src/gateway.rs create mode 100644 src/router.rs diff --git a/Cargo.lock b/Cargo.lock index cf08457..a5dc362 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -25,7 +25,7 @@ dependencies = [ [[package]] name = "aitokenpool" -version = "0.2.0" +version = "0.2.1" dependencies = [ "anyhow", "argon2", @@ -158,6 +158,7 @@ checksum = "edca88bc138befd0323b20752846e6587272d3b03b0343c8ea28a6f819e6e71f" dependencies = [ "async-trait", "axum-core", + "axum-macros", "bytes", "futures-util", "http", @@ -205,6 +206,17 @@ dependencies = [ "tracing", ] +[[package]] +name = "axum-macros" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57d123550fa8d071b7255cb0cc04dc302baa6c8c4a79f55701552684d8399bce" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "base64" version = "0.22.1" diff --git a/Cargo.toml b/Cargo.toml index 5fcbc0d..0602c2d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "aitokenpool" -version = "0.2.0" +version = "0.2.1" edition = "2021" description = "AI Token 共享池 — 企业 key 池 + 公共共享市场" license = "MIT" @@ -9,7 +9,7 @@ rust-version = "1.86" [dependencies] # HTTP 服务器（复用 openlocalrouter 技术栈） -axum = "0.7" +axum = { version = "0.7", features = ["macros"] } tokio = { version = "1", features = ["macros", "rt-multi-thread", "time", "signal"] } tower = { version = "0.5", features = ["util"] } tower-http = { version = "0.5", features = ["cors", "fs", "trace"] } diff --git a/src/billing.rs b/src/billing.rs new file mode 100644 index 0000000..21fba83 --- /dev/null +++ b/src/billing.rs @@ -0,0 +1,300 @@ +//! 计量账本（architecture §4.3/4.4） +//! +//! P0-B（rant 2026-08-18T09:55:57）： +//! - 成本 = prompt_tokens × input_per_m/1e6 + completion_tokens × output_per_m/1e6 +//! （models 表单价；货币按 config points.anchor_currency 折算） +//! - 点数 = 锚定货币成本 × points_per_unit +//! - 消费者扣 balance（上游调用前预检，余额 ≤ 0 → 402） +//! - 分享者（key 属主）得 90%（平台抽成 10%），写 transactions（consume / earn） +//! - 写 usage_records；更新 keys.used += tokens +//! - 调用+记账事务性处理：上游失败不入账（settle 只在成功响应后调用） + +use anyhow::Result; +use rusqlite::Connection; + +/// 分享者分成比例（平台抽成 10%） +pub const SHARE_RATIO: f64 = 0.9; +/// 锚定货币为 USD 时的 CNY 汇率占位（正式汇率表留后续 P0） +pub const CNY_PER_USD: f64 = 7.2; + +/// 把某货币金额折算到锚定货币（同币种 1:1；USD↔CNY 用 CNY_PER_USD；其它未知币种按 1:1 兜底） +pub fn to_anchor(amount: f64, model_currency: &str, anchor: &str) -> f64 { + if model_currency == anchor { + return amount; + } + match (anchor, model_currency) { + ("USD", "CNY") => amount / CNY_PER_USD, + ("CNY", "USD") => amount * CNY_PER_USD, + _ => amount, + } +} + +/// 原始成本（模型自身货币）：USD/CNY 每百万 token 单价 × token 数 +pub fn raw_cost(input_tokens: f64, output_tokens: f64, input_per_m: f64, output_per_m: f64) -> f64 { + input_tokens * input_per_m / 1_000_000.0 + output_tokens * output_per_m / 1_000_000.0 +} + +/// 计算点数：锚定货币成本 × points_per_unit +pub fn calc_points( + input_tokens: f64, + output_tokens: f64, + input_per_m: f64, + output_per_m: f64, + points_per_unit: u32, + model_currency: &str, + anchor: &str, +) -> f64 { + let cost = raw_cost(input_tokens, output_tokens, input_per_m, output_per_m); + to_anchor(cost, model_currency, anchor) * points_per_unit as f64 +} + +/// 入账参数 +pub struct SettleParams { + /// 消费者（调用者）user_id + pub consumer_id: i64, + /// 分发 API key id（可空） + pub api_key_id: Option, + /// 上游 key id + pub key_id: i64, + /// 上游 key 属主（分享者）user_id + pub owner_id: i64, + pub model: String, + /// 本次调用 token 总数 + pub tokens: f64, + /// 消费者应扣点数 + pub pts: f64, + /// 锚定货币成本（usage_records.cost） + pub cost: f64, +} + +/// 事务性入账：扣消费者 → 加分享者 90% → 两条 transactions → usage_records → keys.used +/// 任一步失败整体回滚（调用方只在成功响应后调用，天然满足「失败不入账」） +pub fn settle(conn: &mut Connection, p: &SettleParams) -> Result<()> { + let tx = conn.transaction()?; + + // 消费者扣 balance（余额允许为负——预检已拦截 ≤0 的请求，负余额由后续充值覆盖） + tx.execute( + "UPDATE quotas SET balance = balance - ?1, updated_at = datetime('now') WHERE user_id = ?2", + rusqlite::params![p.pts, p.consumer_id], + )?; + + // 分享者加 90%（平台抽成 10%） + let earn = p.pts * SHARE_RATIO; + tx.execute( + "INSERT OR IGNORE INTO quotas (user_id, balance) VALUES (?1, 0)", + [p.owner_id], + )?; + tx.execute( + "UPDATE quotas SET balance = balance + ?1, updated_at = datetime('now') WHERE user_id = ?2", + rusqlite::params![earn, p.owner_id], + )?; + + // transactions：consume（消费者）+ earn（分享者），counterpart 记对方 + tx.execute( + "INSERT INTO transactions (user_id, counterpart, key_id, model, tokens, pts, type, status) \ + VALUES (?1, ?2, ?3, ?4, ?5, ?6, 'consume', '成功')", + rusqlite::params![ + p.consumer_id, + p.owner_id.to_string(), + p.key_id, + p.model, + p.tokens, + p.pts + ], + )?; + tx.execute( + "INSERT INTO transactions (user_id, counterpart, key_id, model, tokens, pts, type, status) \ + VALUES (?1, ?2, ?3, ?4, ?5, ?6, 'earn', '成功')", + rusqlite::params![ + p.owner_id, + p.consumer_id.to_string(), + p.key_id, + p.model, + p.tokens, + earn + ], + )?; + + // usage_records + tx.execute( + "INSERT INTO usage_records (user_id, api_key_id, key_id, model, tokens, cost) \ + VALUES (?1, ?2, ?3, ?4, ?5, ?6)", + rusqlite::params![ + p.consumer_id, + p.api_key_id, + p.key_id, + p.model, + p.tokens, + p.cost + ], + )?; + + // keys.used 累计（key 不存在 → 报错回滚：账本不允许记到幽灵 key） + let n = tx.execute( + "UPDATE keys SET used = used + ?1 WHERE id = ?2", + rusqlite::params![p.tokens, p.key_id], + )?; + if n != 1 { + return Err(anyhow::anyhow!("key {} not found", p.key_id)); + } + + tx.commit()?; + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::db; + + fn tmp_db(tag: &str) -> (Connection, std::path::PathBuf) { + let p = std::env::temp_dir().join(format!("atp_bill_{}_{}.db", std::process::id(), tag)); + let _ = std::fs::remove_file(&p); + let conn = db::open(p.to_str().unwrap()).expect("open tmp db"); + (conn, p) + } + + #[test] + fn to_anchor_same_currency_identity() { + assert_eq!(to_anchor(10.0, "USD", "USD"), 10.0); + assert_eq!(to_anchor(10.0, "CNY", "CNY"), 10.0); + } + + #[test] + fn to_anchor_cross_currency() { + // USD 锚定：CNY 价 ÷ 7.2 + assert!((to_anchor(72.0, "CNY", "USD") - 10.0).abs() < 1e-9); + // CNY 锚定：USD 价 × 7.2 + assert!((to_anchor(10.0, "USD", "CNY") - 72.0).abs() < 1e-9); + // 未知币种按 1:1 兜底 + assert_eq!(to_anchor(5.0, "EUR", "USD"), 5.0); + } + + #[test] + fn calc_points_basic() { + // 100 prompt × 10 USD/M + 50 completion × 20 USD/M = 0.002 USD → 1000 点/USD → 2.0 点 + let pts = calc_points(100.0, 50.0, 10.0, 20.0, 1000, "USD", "USD"); + assert!((pts - 2.0).abs() < 1e-9, "pts={pts}"); + } + + #[test] + fn calc_points_currency_conversion() { + // 1000 prompt × 10 CNY/M = 0.01 CNY → USD 锚定 ÷ 7.2 ≈ 0.001389 USD → ×1000 ≈ 1.3889 点 + let pts = calc_points(1000.0, 0.0, 10.0, 0.0, 1000, "CNY", "USD"); + let expect = 0.01 / 7.2 * 1000.0; + assert!((pts - expect).abs() < 1e-9, "pts={pts} expect={expect}"); + } + + #[test] + fn settle_splits_90_10_and_writes_ledger() { + let (mut conn, p) = tmp_db("settle"); + // 属主用户（user_id=2）与 key + conn.execute( + "INSERT INTO users (id, email, password_hash, name, role) VALUES (2, 'owner@t.local', 'x', '分享者', 'user')", + [], + ) + .unwrap(); + conn.execute("INSERT INTO quotas (user_id, balance) VALUES (2, 0)", []) + .unwrap(); + conn.execute( + "INSERT INTO keys (id, provider, plan, model, status, owner_id, encrypted_key, quota, used) \ + VALUES (9, 'test', 'test-plan', 'test-model', 'on', 2, 'sk-test', 1000, 0)", + [], + ) + .unwrap(); + // demo（user_id=1）余额 12471 + let bal_before: f64 = conn + .query_row("SELECT balance FROM quotas WHERE user_id = 1", [], |r| { + r.get(0) + }) + .unwrap(); + assert_eq!(bal_before, 12471.0); + + let params = SettleParams { + consumer_id: 1, + api_key_id: Some(3), + key_id: 9, + owner_id: 2, + model: "test-model".into(), + tokens: 150.0, + pts: 2.0, + cost: 0.002, + }; + settle(&mut conn, ¶ms).unwrap(); + + // 消费者扣 2.0 + let bal_c: f64 = conn + .query_row("SELECT balance FROM quotas WHERE user_id = 1", [], |r| { + r.get(0) + }) + .unwrap(); + assert!((bal_c - (12471.0 - 2.0)).abs() < 1e-9); + // 分享者加 1.8（90%） + let bal_o: f64 = conn + .query_row("SELECT balance FROM quotas WHERE user_id = 2", [], |r| { + r.get(0) + }) + .unwrap(); + assert!((bal_o - 1.8).abs() < 1e-9, "owner balance={bal_o}"); + // transactions 两条 + let n: i64 = conn + .query_row("SELECT COUNT(*) FROM transactions", [], |r| r.get(0)) + .unwrap(); + assert_eq!(n, 2); + let t_consume: String = conn + .query_row("SELECT type FROM transactions WHERE user_id = 1", [], |r| { + r.get(0) + }) + .unwrap(); + assert_eq!(t_consume, "consume"); + let t_earn: String = conn + .query_row("SELECT type FROM transactions WHERE user_id = 2", [], |r| { + r.get(0) + }) + .unwrap(); + assert_eq!(t_earn, "earn"); + // usage_records 一条 + let u: i64 = conn + .query_row("SELECT COUNT(*) FROM usage_records", [], |r| r.get(0)) + .unwrap(); + assert_eq!(u, 1); + // keys.used 更新 + let used: f64 = conn + .query_row("SELECT used FROM keys WHERE id = 9", [], |r| r.get(0)) + .unwrap(); + assert!((used - 150.0).abs() < 1e-9); + + drop(conn); + let _ = std::fs::remove_file(p); + } + + #[test] + fn settle_rolls_back_on_error() { + let (mut conn, p) = tmp_db("rollback"); + // key 不存在 → keys.used 更新失败 → 整体回滚（无 transactions / usage_records） + let params = SettleParams { + consumer_id: 1, + api_key_id: None, + key_id: 99999, + owner_id: 1, + model: "m".into(), + tokens: 10.0, + pts: 1.0, + cost: 0.001, + }; + assert!(settle(&mut conn, ¶ms).is_err()); + let n: i64 = conn + .query_row("SELECT COUNT(*) FROM transactions", [], |r| r.get(0)) + .unwrap(); + assert_eq!(n, 0, "回滚后不应有 transactions"); + let bal: f64 = conn + .query_row("SELECT balance FROM quotas WHERE user_id = 1", [], |r| { + r.get(0) + }) + .unwrap(); + assert_eq!(bal, 12471.0, "回滚后余额不变"); + + drop(conn); + let _ = std::fs::remove_file(p); + } +} diff --git a/src/dao.rs b/src/dao.rs index 023a3bf..7ee8fd2 100644 --- a/src/dao.rs +++ b/src/dao.rs @@ -69,16 +69,96 @@ pub fn list_api_keys(conn: &Connection, user_id: i64) -> Result Option { +/// Bearer 认证：按 key 查归属用户 + api_key id → Some((user_id, api_key_id)) +pub fn find_api_key_user_and_id(conn: &Connection, key: &str) -> Option<(i64, i64)> { conn.query_row( - "SELECT user_id FROM api_keys WHERE key_value = ?1 AND status = 'active'", + "SELECT user_id, id FROM api_keys WHERE key_value = ?1 AND status = 'active'", [key], + |r| Ok((r.get(0)?, r.get(1)?)), + ) + .ok() +} + +/// 上游 key 行（路由/网关用） +#[derive(Debug, Clone)] +pub struct KeyRow { + pub id: i64, + pub provider: String, + pub plan: String, + pub owner_id: i64, + pub encrypted_key: String, +} + +/// 查某模型的健康 key（status='on'）——路由候选集 +pub fn find_keys_by_model(conn: &Connection, model: &str) -> Result> { + let mut stmt = conn.prepare( + "SELECT id, provider, plan, owner_id, encrypted_key FROM keys \ + WHERE model = ?1 AND status = 'on'", + )?; + let rows = stmt.query_map([model], |r| { + Ok(KeyRow { + id: r.get(0)?, + provider: r.get(1)?, + plan: r.get(2)?, + owner_id: r.get(3)?, + encrypted_key: r.get(4)?, + }) + })?; + let mut out = Vec::new(); + for r in rows { + out.push(r?); + } + Ok(out) +} + +/// 用户点数余额（无账户按 0） +pub fn get_balance(conn: &Connection, user_id: i64) -> f64 { + conn.query_row( + "SELECT balance FROM quotas WHERE user_id = ?1", + [user_id], |r| r.get(0), ) + .unwrap_or(0.0) +} + +/// 模型单价（按 provider+model）→ (input_per_m, output_per_m, currency) +pub fn get_model_price( + conn: &Connection, + provider: &str, + model: &str, +) -> Option<(f64, f64, String)> { + conn.query_row( + "SELECT input_per_m, output_per_m, currency FROM models WHERE provider = ?1 AND model = ?2", + rusqlite::params![provider, model], + |r| Ok((r.get(0)?, r.get(1)?, r.get(2)?)), + ) .ok() } +/// 市场页：models 表 + key 可用性（可用 key 数） +pub fn list_models_with_availability(conn: &Connection) -> Result> { + let mut stmt = conn.prepare( + "SELECT m.provider, m.model, m.currency, m.input_per_m, m.output_per_m, \ + (SELECT COUNT(*) FROM keys k WHERE k.model = m.model AND k.status = 'on') AS avail \ + FROM models m ORDER BY m.provider, m.model", + )?; + let rows = stmt.query_map([], |r| { + Ok(serde_json::json!({ + "provider": r.get::<_, String>(0)?, + "model": r.get::<_, String>(1)?, + "currency": r.get::<_, String>(2)?, + "input_per_m": r.get::<_, f64>(3)?, + "output_per_m": r.get::<_, f64>(4)?, + "available_keys": r.get::<_, i64>(5)?, + })) + })?; + let mut out = Vec::new(); + for r in rows { + out.push(r?); + } + Ok(out) +} + /// 更新最近使用时间 pub fn touch_api_key(conn: &Connection, key: &str) -> Result<()> { conn.execute( diff --git a/src/db.rs b/src/db.rs index 936ef21..1388664 100644 --- a/src/db.rs +++ b/src/db.rs @@ -99,6 +99,8 @@ pub fn migrate(conn: &Connection) -> Result<()> { cost REAL NOT NULL DEFAULT 0, time TEXT NOT NULL DEFAULT (datetime('now')) ); + CREATE UNIQUE INDEX IF NOT EXISTS idx_models_provider_model + ON models(provider, model); "#, )?; // schema_version：INSERT OR REPLACE 保证幂等 @@ -150,6 +152,63 @@ pub fn seed(conn: &Connection) -> Result<()> { Ok(()) } +/// models 表种子（启动时 upsert）：data/models.example.json 价格大表 + +/// config.price_overrides 官方价覆盖（覆盖同名模型）。文件缺失时仅告警不报错。 +pub fn seed_models(conn: &Connection, cfg: &crate::config::Config) -> Result<()> { + let path = "data/models.example.json"; + let raw = match std::fs::read_to_string(path) { + Ok(s) => s, + Err(e) => { + log::warn!("models 种子文件缺失（{path}），跳过：{e}"); + return Ok(()); + } + }; + let v: serde_json::Value = + serde_json::from_str(&raw).with_context(|| format!("解析 {path} 失败"))?; + let models = v["models"].as_array().cloned().unwrap_or_default(); + + // 官方价覆盖：(provider, model) → (input_per_m, output_per_m, currency) + let mut overrides = std::collections::HashMap::new(); + for o in &cfg.price_overrides { + overrides.insert( + (o.provider.clone(), o.model.clone()), + (o.input_per_m, o.output_per_m, o.currency.clone()), + ); + } + + let mut n = 0u32; + for m in models { + let provider = m["provider"].as_str().unwrap_or("").to_string(); + let model = m["model"].as_str().unwrap_or("").to_string(); + if provider.is_empty() || model.is_empty() { + continue; + } + let (input, output, currency) = overrides + .get(&(provider.clone(), model.clone())) + .cloned() + .unwrap_or_else(|| { + ( + m["input_per_m"].as_f64().unwrap_or(0.0), + m["output_per_m"].as_f64().unwrap_or(0.0), + m["currency"].as_str().unwrap_or("USD").to_string(), + ) + }); + conn.execute( + "INSERT INTO models (provider, model, currency, input_per_m, output_per_m, updated_at) \ + VALUES (?1, ?2, ?3, ?4, ?5, datetime('now')) \ + ON CONFLICT(provider, model) DO UPDATE SET \ + currency = excluded.currency, \ + input_per_m = excluded.input_per_m, \ + output_per_m = excluded.output_per_m, \ + updated_at = datetime('now')", + rusqlite::params![provider, model, currency, input, output], + )?; + n += 1; + } + log::info!("models seed：{n} 行 upsert（来源 {path} + price_overrides）"); + Ok(()) +} + #[cfg(test)] mod tests { use super::*; @@ -200,4 +259,36 @@ mod tests { drop(conn); let _ = std::fs::remove_file(p); } + + #[test] + fn seed_models_upserts_from_example_and_overrides() { + let (conn, p) = tmp_db("models"); + let cfg = crate::config::Config::load("config/config.example.toml").unwrap(); + seed_models(&conn, &cfg).expect("seed models"); + // example 文件里的模型已入库 + let cnt: i64 = conn + .query_row("SELECT COUNT(*) FROM models", [], |r| r.get(0)) + .unwrap(); + assert!(cnt >= 5, "models 已从 example 文件 seed，cnt={cnt}"); + // deepseek-v4-pro 被 price_overrides 覆盖为官方价（0.435 / 0.87 USD） + let (input, output, currency): (f64, f64, String) = conn + .query_row( + "SELECT input_per_m, output_per_m, currency FROM models \ + WHERE provider = 'deepseek' AND model = 'deepseek-v4-pro'", + [], + |r| Ok((r.get(0)?, r.get(1)?, r.get(2)?)), + ) + .unwrap(); + assert!((input - 0.435).abs() < 1e-9, "input={input}"); + assert!((output - 0.87).abs() < 1e-9, "output={output}"); + assert_eq!(currency, "USD"); + // 幂等：重复 seed 不报错且不产生重复行 + seed_models(&conn, &cfg).expect("二次 seed"); + let cnt2: i64 = conn + .query_row("SELECT COUNT(*) FROM models", [], |r| r.get(0)) + .unwrap(); + assert_eq!(cnt, cnt2, "upsert 不产生重复行"); + drop(conn); + let _ = std::fs::remove_file(p); + } } diff --git a/src/gateway.rs b/src/gateway.rs new file mode 100644 index 0000000..ffb559c --- /dev/null +++ b/src/gateway.rs @@ -0,0 +1,588 @@ +//! 网关兼容端点（architecture §4.1） +//! +//! P0-B（rant 2026-08-18T09:55:57）： +//! - POST /v1/chat/completions（OpenAI 兼容，Bearer atk_ API Key 认证） +//! - POST /anthropic/v1/messages（Anthropic 兼容） +//! - GET /api/models（市场页：models 表 + key 可用性） +//! +//! 流程：请求体取 model → 路由选 key（粘性/随机/冷却/3 次切换）→ +//! reqwest 转发到 plan 对应 base_url（openai_chat → {base}/chat/completions； +//! anthropic → {base}/v1/messages）→ 上游响应原样透传 → 解析 usage → 计量入账。 +//! 上游 key 当前为明文占位（加密留 P0-C）。 + +use axum::extract::State; +use axum::http::StatusCode; +use axum::response::Response; +use axum::Json; + +use crate::billing; +use crate::config::Config; +use crate::dao; +use crate::routes::{internal, ApiErr, AppState, AuthUser}; + +/// 网关错误（OpenAI 兼容格式：{"error":{"message":...}}） +fn err_json(status: StatusCode, msg: &str) -> ApiErr { + ( + status, + Json(serde_json::json!({ "error": { "message": msg } })), + ) +} + +/// 从请求体提取 model 字段 +fn extract_model(body: &str) -> Option { + let v: serde_json::Value = serde_json::from_str(body).ok()?; + v.get("model") + .and_then(|m| m.as_str()) + .map(|s| s.to_string()) +} + +/// 解析上游响应中的 usage（openai: prompt/completion_tokens；anthropic: input/output_tokens） +fn parse_usage(body: &[u8], protocol: &str) -> (f64, f64) { + let Ok(v) = serde_json::from_slice::(body) else { + return (0.0, 0.0); + }; + let usage = v.get("usage"); + match protocol { + "anthropic" => { + let input = usage + .and_then(|u| u.get("input_tokens")) + .and_then(|x| x.as_f64()) + .unwrap_or(0.0); + let output = usage + .and_then(|u| u.get("output_tokens")) + .and_then(|x| x.as_f64()) + .unwrap_or(0.0); + (input, output) + } + _ => { + let input = usage + .and_then(|u| u.get("prompt_tokens")) + .and_then(|x| x.as_f64()) + .unwrap_or(0.0); + let output = usage + .and_then(|u| u.get("completion_tokens")) + .and_then(|x| x.as_f64()) + .unwrap_or(0.0); + (input, output) + } + } +} + +/// 按 plan + 协议解析上游完整 URL +fn resolve_endpoint(cfg: &Config, plan_id: &str, protocol: &str) -> Option { + let plan = cfg.plans.iter().find(|p| p.id == plan_id)?; + let ep = plan.endpoints.iter().find(|e| e.protocol == protocol)?; + let base = ep.base_url.trim_end_matches('/'); + Some(match protocol { + "anthropic" => format!("{base}/v1/messages"), + _ => format!("{base}/chat/completions"), + }) +} + +/// 原样透传上游响应 +fn passthrough(status: StatusCode, body: Vec) -> Response { + Response::builder() + .status(status) + .header("content-type", "application/json") + .body(axum::body::Body::from(body)) + .expect("build passthrough response") +} + +/// 核心转发逻辑（openai_chat / anthropic 共用） +async fn forward( + st: &AppState, + auth: AuthUser, + model: &str, + body: String, + protocol: &str, +) -> Result { + // 余额预检（上游调用前）：余额 ≤ 0 → 402 + // 锁作用域严格限定在同步读区内，绝不在 await 期间持有 MutexGuard + let (balance, keys) = { + let conn = st.db.lock().map_err(|_| internal("db lock poisoned"))?; + let balance = dao::get_balance(&conn, auth.user_id); + let keys = dao::find_keys_by_model(&conn, model).map_err(internal)?; + (balance, keys) + }; + if balance <= 0.0 { + return Err(err_json(StatusCode::PAYMENT_REQUIRED, "点数余额不足")); + } + + if keys.is_empty() { + return Err(err_json( + StatusCode::SERVICE_UNAVAILABLE, + "该模型暂无可用 key", + )); + } + + for _ in 0..crate::router::MAX_SWITCHES { + // 选 key（粘性优先 / 随机 / 排除冷却） + let Some(key_id) = st.router.pick(&keys, auth.user_id, model) else { + return Err(err_json( + StatusCode::SERVICE_UNAVAILABLE, + "该模型暂无可用 key", + )); + }; + let key = keys + .iter() + .find(|k| k.id == key_id) + .expect("pick 返回的 key 必然在候选集"); + + // 解析端点；plan 无对应协议端点 → 视为该 key 不可用 + let Some(url) = resolve_endpoint(&st.cfg, &key.plan, protocol) else { + st.router.mark_unhealthy(key_id); + continue; + }; + + // 转发上游（anthropic 用 x-api-key，openai 用 Bearer） + let resp = if protocol == "anthropic" { + st.http + .post(&url) + .header("x-api-key", &key.encrypted_key) + .header("content-type", "application/json") + .header("anthropic-version", "2023-06-01") + .body(body.clone()) + .send() + .await + } else { + st.http + .post(&url) + .header("authorization", format!("Bearer {}", key.encrypted_key)) + .header("content-type", "application/json") + .body(body.clone()) + .send() + .await + }; + + let resp = match resp { + Ok(r) => r, + // 网络错误 → key 非健康，静默切换 + Err(_) => { + st.router.mark_unhealthy(key_id); + continue; + } + }; + let status = resp.status(); + let bytes = resp.bytes().await.unwrap_or_default().to_vec(); + + if status.is_success() { + // 成功：解析 usage → 计量入账 → 粘性 + let (input_tokens, output_tokens) = parse_usage(&bytes, protocol); + let tokens = input_tokens + output_tokens; + if tokens > 0.0 { + let mut conn = st.db.lock().map_err(|_| internal("db lock poisoned"))?; + let price = dao::get_model_price(&conn, &key.provider, model); + let (pts, cost) = match price { + Some((i_per_m, o_per_m, currency)) => { + let pts = billing::calc_points( + input_tokens, + output_tokens, + i_per_m, + o_per_m, + st.cfg.points.points_per_unit, + ¤cy, + &st.cfg.points.anchor_currency, + ); + let cost = billing::to_anchor( + billing::raw_cost(input_tokens, output_tokens, i_per_m, o_per_m), + ¤cy, + &st.cfg.points.anchor_currency, + ); + (pts, cost) + } + None => (0.0, 0.0), + }; + let params = billing::SettleParams { + consumer_id: auth.user_id, + api_key_id: Some(auth.api_key_id), + key_id: key.id, + owner_id: key.owner_id, + model: model.to_string(), + tokens, + pts, + cost, + }; + if let Err(e) = billing::settle(&mut conn, ¶ms) { + // 入账失败不影响透传（响应已成功），仅记录日志 + log::error!("计量入账失败 key_id={} user={}: {e}", key.id, auth.user_id); + } + drop(conn); + st.router.mark_sticky(auth.user_id, model, key_id); + } + return Ok(passthrough(status, bytes)); + } else if status == StatusCode::UNAUTHORIZED + || status == StatusCode::FORBIDDEN + || status == StatusCode::TOO_MANY_REQUESTS + || status.is_server_error() + { + // 401/403/429/5xx → key 非健康，静默切换 + st.router.mark_unhealthy(key_id); + continue; + } else { + // 其它 4xx（400/404 等）→ 用户请求错误，不切换直接透传 + return Ok(passthrough(status, bytes)); + } + } + Err(err_json( + StatusCode::SERVICE_UNAVAILABLE, + "该模型暂无可用 key", + )) +} + +/// POST /v1/chat/completions（OpenAI 兼容） +#[axum::debug_handler] +pub async fn chat_completions( + State(st): State, + auth: AuthUser, + body: String, +) -> Result { + let model = extract_model(&body) + .ok_or_else(|| err_json(StatusCode::BAD_REQUEST, "请求体缺少 model 字段"))?; + forward(&st, auth, &model, body.clone(), "openai_chat").await +} + +/// POST /anthropic/v1/messages（Anthropic 兼容） +pub async fn anthropic_messages( + State(st): State, + auth: AuthUser, + body: String, +) -> Result { + let model = extract_model(&body) + .ok_or_else(|| err_json(StatusCode::BAD_REQUEST, "请求体缺少 model 字段"))?; + forward(&st, auth, &model, body.clone(), "anthropic").await +} + +/// GET /api/models（市场页：models 表 + key 可用性，需认证） +pub async fn models( + State(st): State, + _auth: AuthUser, +) -> Result>, ApiErr> { + let conn = st.db.lock().map_err(|_| internal("db lock poisoned"))?; + let list = dao::list_models_with_availability(&conn).map_err(internal)?; + Ok(Json(list)) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::routes::router; + use axum::body::Body; + use axum::http::Request; + use std::sync::Arc; + use tower::util::ServiceExt; + + /// 测试状态：config.example + 追加本地 test plan；db 开库 + seed models + 注入测试 key + fn test_state(tag: &str, plan_id: &str, base_url: &str) -> AppState { + let p = std::env::temp_dir().join(format!("atp_gw_{}_{}.db", std::process::id(), tag)); + let _ = std::fs::remove_file(&p); + let conn = crate::db::open(p.to_str().unwrap()).expect("open tmp db"); + let mut cfg = crate::config::Config::load("config/config.example.toml").unwrap(); + cfg.plans.push(crate::config::Plan { + id: plan_id.to_string(), + provider: "test".to_string(), + type_: "paygo".to_string(), + key_prefix: "sk-".to_string(), + interactive_only: false, + endpoints: vec![crate::config::Endpoint { + protocol: "openai_chat".to_string(), + base_url: base_url.to_string(), + }], + }); + crate::db::seed_models(&conn, &cfg).expect("seed models"); + AppState::new(conn, Arc::new(cfg)) + } + + /// 注入测试 key（属主 user_id，模型 model，plan） + fn insert_key(st: &AppState, id: i64, owner: i64, model: &str, plan: &str) { + let conn = st.db.lock().unwrap(); + conn.execute( + "INSERT OR REPLACE INTO keys (id, provider, plan, model, status, owner_id, encrypted_key, quota, used) \ + VALUES (?1, 'test', ?2, ?3, 'on', ?4, 'sk-test', 1000, 0)", + rusqlite::params![id, plan, model, owner], + ) + .unwrap(); + } + + async fn post_raw( + st: AppState, + uri: &str, + body: &str, + bearer: Option<&str>, + ) -> (StatusCode, Vec) { + let mut b = Request::builder() + .method("POST") + .uri(uri) + .header("content-type", "application/json"); + if let Some(k) = bearer { + b = b.header("authorization", format!("Bearer {k}")); + } + let resp = router() + .with_state(st) + .oneshot(b.body(Body::from(body.to_string())).unwrap()) + .await + .unwrap(); + let status = resp.status(); + let bytes = axum::body::to_bytes(resp.into_body(), 8 * 1024 * 1024) + .await + .unwrap() + .to_vec(); + (status, bytes) + } + + async fn login_key(st: AppState) -> String { + let (_, body) = post_raw( + st, + "/api/auth/login", + r#"{"email":"demo@aitokenpool.local","password":"demo1234"}"#, + None, + ) + .await; + let v: serde_json::Value = serde_json::from_slice(&body).unwrap(); + v["api_key"].as_str().unwrap().to_string() + } + + /// 假上游：返回固定 usage + async fn fake_upstream(port: u16) { + let app = axum::Router::new().route( + "/chat/completions", + axum::routing::post(|_body: String| async { + Json(serde_json::json!({ + "id": "cmpl-test", + "object": "chat.completion", + "model": "test-model", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}}], + "usage": { "prompt_tokens": 100, "completion_tokens": 50 } + })) + }), + ); + let listener = tokio::net::TcpListener::bind(("127.0.0.1", port)) + .await + .unwrap(); + axum::serve(listener, app).await.unwrap(); + } + + fn models_row( + conn: &rusqlite::Connection, + provider: &str, + model: &str, + input: f64, + output: f64, + ) { + conn.execute( + "INSERT OR REPLACE INTO models (provider, model, currency, input_per_m, output_per_m) \ + VALUES (?1, ?2, 'USD', ?3, ?4)", + rusqlite::params![provider, model, input, output], + ) + .unwrap(); + } + + #[tokio::test] + async fn no_bearer_401() { + let st = test_state("noauth", "test-plan", "http://127.0.0.1:9"); + let (s, _) = post_raw(st, "/v1/chat/completions", r#"{"model":"m"}"#, None).await; + assert_eq!(s, StatusCode::UNAUTHORIZED); + } + + #[tokio::test] + async fn missing_model_400() { + let st = test_state("nomodel", "test-plan", "http://127.0.0.1:9"); + let key = login_key(st.clone()).await; + let (s, _) = post_raw(st, "/v1/chat/completions", r#"{}"#, Some(&key)).await; + assert_eq!(s, StatusCode::BAD_REQUEST); + } + + #[tokio::test] + async fn insufficient_balance_402() { + let st = test_state("insuf", "test-plan", "http://127.0.0.1:9"); + // demo 余额清零 + { + let conn = st.db.lock().unwrap(); + conn.execute("UPDATE quotas SET balance = 0 WHERE user_id = 1", []) + .unwrap(); + } + let key = login_key(st.clone()).await; + let (s, body) = post_raw(st, "/v1/chat/completions", r#"{"model":"m"}"#, Some(&key)).await; + assert_eq!(s, StatusCode::PAYMENT_REQUIRED); + let v: serde_json::Value = serde_json::from_slice(&body).unwrap(); + assert!(v["error"]["message"].as_str().unwrap().contains("余额不足")); + } + + #[tokio::test] + async fn dead_upstream_failover_503() { + // 两个 key 都指向不可达端口 → 3 次尝试全失败 → 503 + 冷却 2 个 key + let st = test_state("dead", "test-dead", "http://127.0.0.1:9"); + insert_key(&st, 100, 1, "dead-model", "test-dead"); + insert_key(&st, 101, 1, "dead-model", "test-dead"); + let key = login_key(st.clone()).await; + let (s, body) = post_raw( + st.clone(), + "/v1/chat/completions", + r#"{"model":"dead-model"}"#, + Some(&key), + ) + .await; + assert_eq!( + s, + StatusCode::SERVICE_UNAVAILABLE, + "body: {}", + String::from_utf8_lossy(&body) + ); + let v: serde_json::Value = serde_json::from_slice(&body).unwrap(); + assert!(v["error"]["message"] + .as_str() + .unwrap() + .contains("暂无可用 key")); + assert_eq!(st.router.cooldown_len(), 2, "两次失败后两个 key 均进入冷却"); + } + + #[tokio::test] + async fn e2e_success_ledger_and_sticky() { + // 起假上游 + let listener = tokio::net::TcpListener::bind(("127.0.0.1", 0)) + .await + .unwrap(); + let port = listener.local_addr().unwrap().port(); + drop(listener); + let up = tokio::spawn(fake_upstream(port)); + let base = format!("http://127.0.0.1:{port}"); + + let st = test_state("e2e", "test-local", &base); + // 价格行：10/20 USD per M + { + let conn = st.db.lock().unwrap(); + models_row(&conn, "test", "test-model", 10.0, 20.0); + // 属主用户（id=2）+ 两个 key 均属主 2（避免随机选中消费者自有的 key 使 + // 净额断言不确定：消费者始终 -2.0，属主始终 +1.8） + conn.execute( + "INSERT OR IGNORE INTO users (id, email, password_hash, name, role) VALUES (2, 'owner@t.local', 'x', '分享者', 'user')", + [], + ) + .unwrap(); + conn.execute( + "INSERT OR IGNORE INTO quotas (user_id, balance) VALUES (2, 0)", + [], + ) + .unwrap(); + } + insert_key(&st, 200, 2, "test-model", "test-local"); + insert_key(&st, 201, 2, "test-model", "test-local"); + + let key = login_key(st.clone()).await; + + // 第一次调用 → 成功 + 入账 + let (s, body) = post_raw( + st.clone(), + "/v1/chat/completions", + r#"{"model":"test-model","messages":[{"role":"user","content":"hi"}]}"#, + Some(&key), + ) + .await; + assert_eq!( + s, + StatusCode::OK, + "body: {}", + String::from_utf8_lossy(&body) + ); + let v: serde_json::Value = serde_json::from_slice(&body).unwrap(); + assert_eq!(v["usage"]["prompt_tokens"], 100); + assert_eq!(v["usage"]["completion_tokens"], 50); + + // 入账断言（锁在块内释放，绝不让 MutexGuard 跨 await） + let used_key: i64 = { + let conn = st.db.lock().unwrap(); + let bal_c: f64 = conn + .query_row("SELECT balance FROM quotas WHERE user_id = 1", [], |r| { + r.get(0) + }) + .unwrap(); + // pts = 100×10/1e6 + 50×20/1e6 = 0.002 USD × 1000 = 2.0 + assert!((bal_c - (12471.0 - 2.0)).abs() < 1e-9, "consumer={bal_c}"); + let bal_o: f64 = conn + .query_row("SELECT balance FROM quotas WHERE user_id = 2", [], |r| { + r.get(0) + }) + .unwrap(); + assert!((bal_o - 1.8).abs() < 1e-9, "owner={bal_o}"); + let n_tx: i64 = conn + .query_row("SELECT COUNT(*) FROM transactions", [], |r| r.get(0)) + .unwrap(); + assert_eq!(n_tx, 2); + let n_ur: i64 = conn + .query_row("SELECT COUNT(*) FROM usage_records", [], |r| r.get(0)) + .unwrap(); + assert_eq!(n_ur, 1); + conn.query_row( + "SELECT key_id FROM usage_records ORDER BY id DESC LIMIT 1", + [], + |r| r.get(0), + ) + .unwrap() + }; + + // 第二次调用 → 粘性复用同一 key + let (s2, _) = post_raw( + st.clone(), + "/v1/chat/completions", + r#"{"model":"test-model","messages":[{"role":"user","content":"hi"}]}"#, + Some(&key), + ) + .await; + assert_eq!(s2, StatusCode::OK); + let used_key2: i64 = { + let conn = st.db.lock().unwrap(); + conn.query_row( + "SELECT key_id FROM usage_records ORDER BY id DESC LIMIT 1", + [], + |r| r.get(0), + ) + .unwrap() + }; + assert_eq!(used_key, used_key2, "粘性：第二次调用复用同一 key"); + + up.abort(); + } + + #[tokio::test] + async fn models_endpoint_lists_with_availability() { + let st = test_state("mk", "test-plan", "http://127.0.0.1:9"); + let key = login_key(st.clone()).await; + let resp = router() + .with_state(st.clone()) + .oneshot( + Request::builder() + .method("GET") + .uri("/api/models") + .header("authorization", format!("Bearer {key}")) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + let bytes = axum::body::to_bytes(resp.into_body(), 4 * 1024 * 1024) + .await + .unwrap(); + let arr: Vec = serde_json::from_slice(&bytes).unwrap(); + assert!(!arr.is_empty(), "models 已从 example 文件 seed"); + // 字段结构 + let first = &arr[0]; + assert!(first.get("provider").is_some()); + assert!(first.get("model").is_some()); + assert!(first.get("input_per_m").is_some()); + assert!(first.get("output_per_m").is_some()); + assert!(first.get("available_keys").is_some()); + // 无认证 → 401 + let resp2 = router() + .with_state(st) + .oneshot( + Request::builder() + .method("GET") + .uri("/api/models") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(resp2.status(), StatusCode::UNAUTHORIZED); + } +} diff --git a/src/main.rs b/src/main.rs index eb30754..96f60cd 100644 --- a/src/main.rs +++ b/src/main.rs @@ -6,13 +6,17 @@ //! 架构定论见 docs/architecture.md（中心化方案 A：平台托管 key + 平台执行调用） //! //! P0-A（rant 2026-08-17T22:21:52）：服务骨架 + 配置加载 + SQLite 数据层 + 认证。 +//! P0-B（rant 2026-08-18T09:55:57）：网关转发 + 路由故障转移 + 计量账本闭环。 //! 启动：cargo run -- --config config/config.toml（默认 config/config.toml， //! 不存在时提示复制 config.example.toml）。 mod auth; +mod billing; mod config; mod dao; mod db; +mod gateway; +mod router; mod routes; use std::sync::Arc; @@ -50,8 +54,10 @@ async fn main() -> anyhow::Result<()> { let db_path = cfg.server.db_path.clone(); log::info!("打开数据库: {db_path}"); let conn = db::open(&db_path)?; + db::seed_models(&conn, &cfg)?; - let state = routes::AppState::new(conn, Arc::new(cfg)); + let cfg = Arc::new(cfg); + let state = routes::AppState::new(conn, cfg); let app = routes::router() .with_state(state) .layer(tower_http::trace::TraceLayer::new_for_http()); diff --git a/src/router.rs b/src/router.rs new file mode 100644 index 0000000..b232bd5 --- /dev/null +++ b/src/router.rs @@ -0,0 +1,183 @@ +//! 路由与故障转移（architecture §4.2.1 七条定案） +//! +//! P0-B（rant 2026-08-18T09:55:57）： +//! - 初始随机：从模型 M 的健康 key（status='on' 且不在冷却期）中随机选一 +//! - 粘性：用户+模型 → 复用上次 key（进程内 HashMap<(user_id, model), key_id>） +//! - 不可用判定：上游 401/403/429/5xx 或网络错误 → 标记非健康 5 秒（冷却），静默切换 +//! - 切换上限 3 次：3 次全失败 → 503「该模型暂无可用 key」（由网关层返回） + +use std::collections::HashMap; +use std::sync::Mutex; +use std::time::{Duration, Instant}; + +use rand::seq::SliceRandom; + +use crate::dao::KeyRow; + +/// 默认冷却时长（architecture §4.2.1 定案：5 秒） +pub const COOLDOWN_SECS: u64 = 5; +/// 每次路由最多尝试的 key 数（初始 + 2 次切换） +pub const MAX_SWITCHES: usize = 3; + +/// 路由状态（进程内）：粘性映射 + 冷却表 +#[derive(Default)] +pub struct RouterState { + sticky: Mutex>, + cooldown: Mutex>, + cooldown_duration: Duration, +} + +impl RouterState { + pub fn new() -> Self { + Self { + sticky: Mutex::new(HashMap::new()), + cooldown: Mutex::new(HashMap::new()), + cooldown_duration: Duration::from_secs(COOLDOWN_SECS), + } + } + + /// 测试用：自定义冷却时长 + #[cfg(test)] + pub fn with_cooldown(d: Duration) -> Self { + Self { + sticky: Mutex::new(HashMap::new()), + cooldown: Mutex::new(HashMap::new()), + cooldown_duration: d, + } + } + + /// key 当前是否健康（不在冷却期内） + pub fn is_healthy(&self, key_id: i64) -> bool { + let c = self.cooldown.lock().expect("cooldown lock"); + match c.get(&key_id) { + None => true, + Some(until) => Instant::now() >= *until, + } + } + + /// 标记 key 非健康（进入冷却），并清除该 key 的所有粘性引用 + pub fn mark_unhealthy(&self, key_id: i64) { + self.cooldown + .lock() + .expect("cooldown lock") + .insert(key_id, Instant::now() + self.cooldown_duration); + self.sticky + .lock() + .expect("sticky lock") + .retain(|_, v| *v != key_id); + } + + /// 记录粘性：该用户+模型后续复用此 key + pub fn mark_sticky(&self, user_id: i64, model: &str, key_id: i64) { + self.sticky + .lock() + .expect("sticky lock") + .insert((user_id, model.to_string()), key_id); + } + + /// 选 key：粘性优先（仍健康），否则从健康集合随机 + /// 无健康 key → None + pub fn pick(&self, keys: &[KeyRow], user_id: i64, model: &str) -> Option { + let healthy: Vec<&KeyRow> = keys.iter().filter(|k| self.is_healthy(k.id)).collect(); + if healthy.is_empty() { + return None; + } + { + let s = self.sticky.lock().expect("sticky lock"); + if let Some(kid) = s.get(&(user_id, model.to_string())) { + if healthy.iter().any(|k| k.id == *kid) { + return Some(*kid); + } + } + } + let mut rng = rand::thread_rng(); + healthy.choose(&mut rng).map(|k| k.id) + } + + /// 当前冷却中的 key 数（测试/诊断用） + #[cfg(test)] + pub fn cooldown_len(&self) -> usize { + self.cooldown.lock().expect("cooldown lock").len() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn keys(ids: &[i64]) -> Vec { + ids.iter() + .map(|id| KeyRow { + id: *id, + provider: "test".into(), + plan: "test-plan".into(), + owner_id: 1, + encrypted_key: "sk-test".into(), + }) + .collect() + } + + #[test] + fn pick_prefers_sticky_key() { + let r = RouterState::new(); + let ks = keys(&[1, 2]); + // 标记粘性 → 两次选择都回到 sticky key + r.mark_sticky(7, "m", 2); + let p1 = r.pick(&ks, 7, "m").unwrap(); + let p2 = r.pick(&ks, 7, "m").unwrap(); + assert_eq!(p1, 2); + assert_eq!(p2, 2); + } + + #[test] + fn random_without_sticky_is_from_healthy_set() { + let r = RouterState::new(); + let ks = keys(&[1, 2, 3]); + for _ in 0..20 { + let p = r.pick(&ks, 7, "m").unwrap(); + assert!(ks.iter().any(|k| k.id == p)); + } + } + + #[test] + fn cooldown_excludes_and_expires() { + let r = RouterState::with_cooldown(Duration::from_millis(50)); + let ks = keys(&[1, 2]); + assert!(r.is_healthy(1)); + r.mark_unhealthy(1); + assert!(!r.is_healthy(1)); + // 冷却中的 key 不会被选中 + for _ in 0..20 { + assert_eq!(r.pick(&ks, 7, "m").unwrap(), 2); + } + assert_eq!(r.cooldown_len(), 1); + // 冷却过期后恢复健康 + std::thread::sleep(Duration::from_millis(70)); + assert!(r.is_healthy(1)); + assert_eq!(r.cooldown_len(), 1, "过期条目保留但视为健康"); + } + + #[test] + fn mark_unhealthy_clears_sticky() { + let r = RouterState::new(); + let ks = keys(&[1, 2]); + r.mark_sticky(7, "m", 1); + r.mark_unhealthy(1); + // sticky key 不可用 → 只能选 2 + assert_eq!(r.pick(&ks, 7, "m").unwrap(), 2); + } + + #[test] + fn no_healthy_key_returns_none() { + let r = RouterState::new(); + let ks = keys(&[1]); + r.mark_unhealthy(1); + assert!(r.pick(&ks, 7, "m").is_none()); + } + + #[test] + fn default_cooldown_is_5_secs() { + let r = RouterState::new(); + assert_eq!(r.cooldown_duration, Duration::from_secs(COOLDOWN_SECS)); + } +} diff --git a/src/routes/mod.rs b/src/routes/mod.rs index 0897a76..263d076 100644 --- a/src/routes/mod.rs +++ b/src/routes/mod.rs @@ -1,9 +1,13 @@ -//! HTTP 路由装配：healthz + 认证 + API Key 管理 +//! HTTP 路由装配：healthz + 认证 + API Key 管理 + 网关/市场端点 //! //! P0-A（rant 2026-08-17T22:21:52）： //! - GET /healthz → {"status":"ok","version":"0.2.0"} //! - POST /api/auth/login → 200 {api_key} / 401 //! - POST /api/api-keys / GET /api/api-keys（Bearer 认证） +//! +//! P0-B（rant 2026-08-18T09:55:57）： +//! - POST /v1/chat/completions / POST /anthropic/v1/messages（网关） +//! - GET /api/models（市场） pub mod api_keys; @@ -19,13 +23,16 @@ use serde::Deserialize; use crate::config::Config; use crate::dao; +use crate::gateway; +use crate::router::RouterState; -/// 共享状态：数据库连接 + 配置 +/// 共享状态：数据库连接 + 配置 + 路由状态 + HTTP 客户端 #[derive(Clone)] pub struct AppState { pub db: Arc>, - #[allow(dead_code)] // 供后续 P0 网关路由阶段消费 pub cfg: Arc, + pub router: Arc, + pub http: reqwest::Client, } impl AppState { @@ -33,6 +40,11 @@ impl AppState { Self { db: Arc::new(Mutex::new(conn)), cfg, + router: Arc::new(RouterState::new()), + http: reqwest::Client::builder() + .timeout(std::time::Duration::from_secs(120)) + .build() + .expect("reqwest client 构建失败"), } } } @@ -47,7 +59,7 @@ fn unauthorized() -> ApiErr { ) } -fn internal(e: impl std::fmt::Display) -> ApiErr { +pub fn internal(e: impl std::fmt::Display) -> ApiErr { ( StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({ "error": format!("{e}") })), @@ -58,6 +70,7 @@ fn internal(e: impl std::fmt::Display) -> ApiErr { #[derive(Debug, Clone, Copy)] pub struct AuthUser { pub user_id: i64, + pub api_key_id: i64, } #[axum::async_trait] @@ -75,10 +88,13 @@ impl FromRequestParts for AuthUser { .unwrap_or(""); let key = header.strip_prefix("Bearer ").ok_or_else(unauthorized)?; let conn = state.db.lock().map_err(|_| internal("db lock poisoned"))?; - match dao::find_user_by_api_key(&conn, key) { - Some(user_id) => { + match dao::find_api_key_user_and_id(&conn, key) { + Some((user_id, api_key_id)) => { let _ = dao::touch_api_key(&conn, key); - Ok(AuthUser { user_id }) + Ok(AuthUser { + user_id, + api_key_id, + }) } None => Err(unauthorized()), } @@ -120,6 +136,9 @@ pub fn router() -> Router { .route("/healthz", get(healthz)) .route("/api/auth/login", post(login)) .route("/api/api-keys", post(api_keys::create).get(api_keys::list)) + .route("/v1/chat/completions", post(gateway::chat_completions)) + .route("/anthropic/v1/messages", post(gateway::anthropic_messages)) + .route("/api/models", get(gateway::models)) } #[cfg(test)] @@ -135,6 +154,7 @@ mod tests { let _ = std::fs::remove_file(&p); let conn = crate::db::open(p.to_str().unwrap()).expect("open tmp db"); let cfg = crate::config::Config::load("config/config.example.toml").unwrap(); + crate::db::seed_models(&conn, &cfg).expect("seed models"); AppState::new(conn, Arc::new(cfg)) } @@ -186,7 +206,7 @@ mod tests { assert_eq!(s, StatusCode::OK); let v: serde_json::Value = serde_json::from_str(&body).unwrap(); assert_eq!(v["status"], "ok"); - assert_eq!(v["version"], "0.2.0"); + assert_eq!(v["version"], env!("CARGO_PKG_VERSION")); } #[tokio::test]