3 Commits

Author SHA1 Message Date
boos 4d093548de feat(proxy): wire Gemini native v1beta endpoints through the conductor
Implement the four Gemini native inbound endpoints so Gemini CLI / Google
GenAI SDKs can point their base URL at the gateway:

- GET  /v1beta/models — directory listing in Gemini shape, built from
  configured routes (like /v1/models); no upstream call, no billing.
- POST /v1beta/models/{model}:generateContent — non-streaming, routed
  through ProxyConductor with failover + usage_logs + billing.
- POST /v1beta/models/{model}:streamGenerateContent — SSE passthrough via
  open_stream_with_failover + sse_passthrough_response.
- POST /v1beta/models/{model}:countTokens — metadata call, routed for a
  channel but not billed.

Auth accepts the gateway key from ?key= or x-goog-api-key (Gemini
convention) and reuses the downstream-key authorization + limits.

Billing closes the loop: gemini_extract_usage reads usageMetadata
(promptTokenCount / candidatesTokenCount) for non-streaming, and
gemini_stream_usage_from_line feeds the shared StreamUsageAccumulator so
streamed usageMetadata is billed too.

Unit tests cover model:action parsing, usage extraction, the stream usage
parser, and query forwarding (key/alt dropped).
2026-06-22 08:39:38 +08:00
boos aaa593ac31 feat(proxy): add Gemini native upstream executors
Add Gemini-shaped upstream calls that mirror AQBot's native passthrough
and BoosAPI's /v1/responses forwarding: the model lives in the URL path
(/v1beta/models/{model}:{op}) and the key in ?key=, so the body is sent
verbatim with no model rewrite.

- gemini_base_url / gemini_operation_url normalize the v1beta root and
  build {base}/v1beta/models/{model}:{op}?key=...&{extra}, stripping a
  leading models/ prefix and handling the no-key case.
- execute_gemini_operation runs non-streaming ops (generateContent /
  countTokens) and returns upstream JSON verbatim.
- execute_gemini_stream_response opens streamGenerateContent for SSE
  passthrough, enforcing alt=sse via merge_alt_sse.

Unit tests cover URL building and alt=sse merging.
2026-06-22 08:39:26 +08:00
boos 24a8dec151 fix: define DEFAULT_NEW_USER_MAX_COST_USD and repair test AppConfig initializers
Pre-existing breakage on this branch (introduced by the billing commits),
unrelated to the Gemini work but blocking both `cargo build` and
`cargo test`:

- admin/users.rs imported routes::auth::DEFAULT_NEW_USER_MAX_COST_USD,
  which was never defined. Add it to auth.rs as $1.0 (the billing design's
  "new users default to $1" starter allowance).
- AppConfig gained `payment` and `pixelle_video` fields but several test
  helpers still used struct literals missing them. Initialize both with
  Default::default() at each affected test site.
2026-06-22 08:39:17 +08:00
13 changed files with 743 additions and 295 deletions
@@ -1,100 +0,0 @@
# NIT AI Studio — 第一阶段(聚合图像/数字人)设计
日期:2026-06-17
分支:`nit_edu`
状态:已通过 brainstorming,待落实现计划
## 背景
NIT 精英班专属的「AI 综合应用软件」是一套面向精英班学员的 AI 创作工具,完整愿景涵盖
文本结构化创作(论文/小说/剧本/分镜/教材)、PPT 生成、图像(P图/制图)、视频
(文生/图生/短视频)、数字人。体量过大,分阶段建设。
**本 spec 仅覆盖 AI Studio 第一阶段:聚合现有图像 + 数字人能力**,为精英班学员提供
统一创作入口,复用 BoosAPI 已有的 `picture_agent`(图像)与 `digital_human`(数字人)。
## 已确认的关键决策
1. 定位:NIT 精英班专属模块,非通用套件。
2. 第一阶段只聚合 `picture_agent`(图像)+ `digital_human`(数字人)。
3. 视频(含 money_printer、工作流视频)留后续阶段,本阶段完全不碰 money_printer。
4. 认证桥接:给精英班学员配后台下游 key + 扩展 surface 的 resolve_auth 接受 student JWT。
5. 计费:学员创作消耗走积分轨(已建的 points track),管理员可给学员充积分。
6. 文本创作 / PPT 留后续阶段。
## 前置依赖
本阶段 gate 在第一期(报考闭环)产物上,实现排在第一期之后:
- student JWT(`audience='student'`,第一期发放)。
- `enrollment.status=approved` 且 `tier=elite`。
- 非精英班 / 未审核 → 403。
## 非目标(本阶段明确不做)
- 视频能力(文生/图生/短视频、money_printer、工作流视频)。
- 文本结构化创作(论文/小说/剧本/分镜/教材)。
- PPT 生成。
- 通用套件(面向非 NIT 用户)。
## 学员下游 key 配置
- `downstream_api_keys` 加可空 `student_id` 列(关联 `students.id`,与现有 `user_id` 并列;
二者互斥:一行要么属于 user,要么属于 student)。
- 审核通过且 `tier=elite` 时,自动 provision 一个学员下游 key,沿用现有 `comfyui-agent`
内置 key 的幂等 `ensure` 模式:缺失则建、被禁用则重新启用、保持唯一。
- 学员从精英班降级或报考被驳回时,禁用该 key。
- key 命名约定:`sk-nit-stu-{uuid_hex}`,`name='nit-student'`,便于识别。
## 认证扩展
仅改两个 surface:`src/routes/picture_agent.rs`、`src/routes/digital_human.rs`。
**money_printer 不动。**
各自的 `resolve_auth` 加一个分支(现有 admin/proxy token、user JWT 逻辑不变):
- 新增:`audience='student'` 的 JWT → 校验该学员存在 `approved` 且 `tier=elite` 的报考
→ 查其 `student_id` 对应的下游 key → 作 `outbound_key`。
- 抽共享 helper `resolve_student_outbound_key(student_id, state) -> AppResult<(key_id, key)>`,
放在 `src/services/nit/`,两处 surface 复用,避免重复逻辑。
- 校验失败(无 elite 报考 / key 缺失)→ 403。
## 计费(复用积分轨)
- 学员用图像/数字人,消耗扣在其下游 key 上,走已建的积分轨(`apply_points_delta`,track=points)。
- 管理员给学员充积分:扩展现有 admin 积分调整端点支持 `student_id`(或新增 NIT 学员积分端点,
实现时择一,优先扩展现有端点)。
- 精英班开通时可送一笔初始积分,数额为配置项,默认 `0`。
## 统一入口(前端)
- `src/web/pages/nit/Studio.tsx`:精英班学员创作工作台首页,卡片式列出「图像 / 数字人」
两个入口 + 显示积分余额。
- 卡片跳进现有 `picture_agent` / `digital_human` 页面(带 student token)。
- 非精英班学员:显示引导升级提示,不展示工具卡片。
- `api.ts` 的 `nitApi` 扩展:查学员积分余额、查 Studio 可用工具。
## 错误处理
- 非精英班 / 未审核学员访问 → 403 + 友好提示。
- 学员下游 key 缺失(provision 失败)→ 明确错误,提示联系管理员。
- 积分余额不足 → 拦截创作,提示充值(复用现有 billing 引导文案范式)。
## 测试策略(TDD)
后端(rstest + 内存 SQLite,沿用现有 `recharge_test_pool` 范式):
- 学员下游 key ensure 幂等性(缺失则建 / 禁用则启用 / 降级则禁用)。
- `resolve_auth` 接受 elite student JWT、拒绝非精英班、拒绝未审核。
- 学员创作扣积分、余额不足拦截。
前端(vitest):
- `nitApi` 积分余额 / Studio 工具查询方法(沿用 `api.test.ts` 范式)。
- Studio 页卡片渲染 + 余额展示 + 非精英班引导(沿用 `UsageRecords.test.tsx` 范式)。
## 验证
- `cargo build -j 2`、`cargo test -p boosapi --lib`(后端 nit / picture_agent / digital_human 相关)。
- `npx vitest run`(相关前端测试)、`npm run typecheck`。
## 提交约定
全部在 `nit_edu` 分支,不混入 `refactor/rust-backend`。提交按逻辑拆分(迁移 / service /
surface 认证 / 前端 分开),沿用项目中文 commit 风格。
@@ -1,195 +0,0 @@
# NIT 报考学习平台 — 第一期(报考闭环 A+B+C)设计
日期:2026-06-17
分支:`nit_edu`
状态:已通过 brainstorming,待落实现计划
## 背景
NIT 报考学习平台面向报考「NIT 人工智能与大模型应用」考试的学员,提供注册、报考、
学习等功能。完整需求(见 `plan/NIT报考学习平台(初版需求备注解析).docx`)涵盖
8 个子系统,体量过大,无法单 spec 完成,故拆分分期建设。
**本 spec 仅覆盖第一期:报考闭环 A+B+C** —— 学员注册 → 报考 → 缴费审核开课。
学习中心(理论课/实操课/练习题)、OCR、校区负责人多级权限、内容保护等留后续期。
## 已确认的关键决策
1. 第一个子项目:A+B+C 报考闭环(注册 → 报考 → 缴费审核开课)。
2. 独立 `students` 表,与现有 `users`(API 网关用户)完全隔离。
3. 缴费:线下转账/扫码给校区 + 学员上传凭证截图 + 后台人工审核开课(非在线支付)。
4. 第一期不做 OCR:学员手动填身份证号/姓名,二寸照手动命名上传。
5. 开课权限只打状态标记(enrollment.status=approved + tier),学习中心留后续期再读此标记。
6. 后台第一期只做平台超管(复用现有 admin 认证),校区负责人登录留后续期。
7. 文件存储复用现有 workspace 机制(multipart + `record_bytes_asset` + `workspace_assets` 表)。
8. 代码组织:NIT 作为 BoosAPI 内的新模块,复用现有 axum/sqlx/JWT/workspace 基建。
9. 学员认证:发 `audience='student'` 的 JWT,与 admin/user token 互不通用。
10. 报名信息表:固定核心字段 + 一个 JSON 扩展字段(应对考试中心模版变化),不做动态表单引擎。
## 非目标(第一期明确不做)
- 学习中心:理论课/实操课视频、二维码、学习进度、练习题、模拟题、错题集。
- 实操课专属 AI 综合应用软件入口。
- 身份证 OCR 自动识别与二寸照自动命名。
- 校区负责人多级权限与数据隔离。
- 内容保护(防截图/录屏/水印)。
- 在线支付。
## 数据模型
三张新表,SQLite,迁移文件放 `migrations/`。
### `students` — 学员实名账户
| 字段 | 类型 | 说明 |
|------|------|------|
| `id` | INTEGER PK | |
| `public_id` | TEXT UNIQUE | 对外 ID,格式 `stu_{uuid_hex}` |
| `name` | TEXT NOT NULL | 姓名(实名,注册后不可改) |
| `phone` | TEXT NOT NULL UNIQUE | 实名手机号,登录用 |
| `id_card_no` | TEXT NOT NULL UNIQUE | 身份证号 |
| `campus_id` | INTEGER NOT NULL | FK → `nit_campuses.id` |
| `major` | TEXT | 所学专业(社会学员可空) |
| `class_name` | TEXT | 班级(可空) |
| `head_teacher` | TEXT | 班主任姓名(可空) |
| `password_hash` | TEXT NOT NULL | 登录密码哈希(沿用现有 users 哈希方案) |
| `privacy_agreed_at` | TEXT NOT NULL | 勾选隐私协议时间 |
| `status` | TEXT NOT NULL DEFAULT 'active' | active / disabled |
| `created_at` / `updated_at` | TEXT | |
### `nit_campuses` — 校区/机构库(超管录入,学员注册下拉选)
| 字段 | 类型 | 说明 |
|------|------|------|
| `id` | INTEGER PK | |
| `public_id` | TEXT UNIQUE | `cmp_{uuid_hex}` |
| `province` | TEXT NOT NULL | 省份 |
| `name` | TEXT NOT NULL | 院校/机构名称 |
| `manager_name` | TEXT | 负责人姓名 |
| `manager_phone` | TEXT | 负责人联系方式 |
| `enabled` | INTEGER NOT NULL DEFAULT 1 | 仅 enabled 的校区可被学员选择 |
| `created_at` / `updated_at` | TEXT | |
### `nit_enrollments` — 报考记录
| 字段 | 类型 | 说明 |
|------|------|------|
| `id` | INTEGER PK | |
| `public_id` | TEXT UNIQUE | `enr_{uuid_hex}` |
| `student_id` | INTEGER NOT NULL | FK → `students.id` |
| `subject` | TEXT NOT NULL | 第一期固定「NIT人工智能与大模型应用」 |
| `tier` | TEXT NOT NULL | `basic`(980)/ `elite`(1980) |
| `payment_time` | TEXT | 学员填写的缴费时间 |
| `payment_proof_asset_id` | TEXT | 缴费凭证截图(workspace asset public_id) |
| `id_card_front_asset_id` | TEXT | 身份证正面 |
| `id_card_back_asset_id` | TEXT | 身份证反面 |
| `photo_asset_id` | TEXT | 二寸蓝底照 |
| `applicant_info` | TEXT (JSON) | 报名材料表可变字段 |
| `status` | TEXT NOT NULL DEFAULT 'pending_review' | pending_review / approved / rejected |
| `reviewed_by` | TEXT | 审核管理员标识 |
| `reviewed_at` | TEXT | |
| `review_note` | TEXT | 驳回原因/备注 |
| `created_at` / `updated_at` | TEXT | |
**关系**:student → campus(多对一);enrollment → student(多对一)。一个学员允许在被驳回后
重新提交报考(可有多条 rejected + 至多一条 active 的 pending_review/approved)。
**班型权限**:第一期权限即 `approved` 后 enrollment 上的 `tier` 标记;`basic` 对应理论课权限、
`elite` 对应理论+实操权限,实际课程内容由后续期读取此标记实现。
**文件**:四类上传文件全部走现有 `workspace_assets`(`source='nit-enrollment'`,
经 `record_bytes_asset` 落本地磁盘),enrollment 仅存其 `public_id` 引用。
## 认证与路由
### 学员认证
复用现有 `auth::jwt` 工具,给 students 发 `audience='student'` 的 JWT。新增
`resolve_student(headers, state) -> AppResult<i64>`(返回 student_id),与 admin/user token
互不通用。登录凭证为「手机号 + 密码」。
### 路由分布(挂进现有 axum router via `routes/mod.rs` merge)
- `src/routes/nit/mod.rs` — 装配 NIT 子路由,导出 `routes()`。
- `src/routes/nit/student_auth.rs`
- `POST /api/nit/register` — 学员注册
- `POST /api/nit/login` — 学员登录
- `src/routes/nit/enrollment.rs`(学员侧,需 student JWT)
- `POST /api/nit/enrollment/uploads` — multipart 上传单个文件,返回 asset public_id
- `POST /api/nit/enrollment` — 创建报考记录
- `GET /api/nit/enrollment/mine` — 查本人报考状态
- `src/routes/nit/admin.rs`(超管侧,复用现有 admin token 校验)
- `GET/POST/PUT /api/admin/nit/campuses` — 校区 CRUD
- `GET /api/admin/nit/enrollments` — 报考列表(按 省份/校区/班型/状态 筛选)
- `POST /api/admin/nit/enrollments/{public_id}/approve`
- `POST /api/admin/nit/enrollments/{public_id}/reject`
- `GET /api/admin/nit/enrollments/export` — CSV 导出
业务逻辑(查询/状态机)抽到 `src/services/nit/`(如 `mod.rs` + `enrollment.rs`),
路由层只做参数解析 + 鉴权 + 调用 service,保持文件聚焦。
## 数据流
### 注册
1. 校验手机号、身份证号未被占用(唯一约束 + 友好报错)。
2. 校验 `campus_id` 存在且 enabled。
3. 必须 `privacy_agreed=true`。
4. 写入 `students`(password 哈希),发 student JWT 返回。
### 报考(两步,避免半截事务)
1. 学员逐个调 `POST /api/nit/enrollment/uploads`(multipart,带 `purpose` 区分四类文件)
→ `record_bytes_asset` → 返回各 asset `public_id`。
2. 学员调 `POST /api/nit/enrollment`,带各 asset id + tier + payment_time + applicant_info
→ 校验本人无在审/已批的活动报考 → 创建 `pending_review` 记录。
### 审核(超管)
- 列表筛选 `pending_review` → 查看凭证/证件图(经 workspace asset file 接口)。
- `approve`:状态机 `pending_review → approved`,落 `tier` 权限标记,记 reviewed_by/at。
- `reject`:`pending_review → rejected`,记 review_note;学员可据此重新提交。
- 状态机只允许从 `pending_review` 流转,杜绝并发重复审核。
### 导出
按 省份/校区/班型/状态 筛选,导出 CSV(实名信息 + 班型),用于核算报考人数/教材数量。
## 前端页面(`src/web/pages/nit/`)
- 学员注册页(实名信息 + 校区下拉 + 隐私协议勾选)
- 学员登录页(手机号 + 密码)
- 报考页(选班型 + 上传四类文件 + 填缴费时间 + 报名信息表 + 提交)
- 「我的报考」状态页(pending/approved/rejected + 驳回原因 + 可重新提交)
- 超管校区管理页(CRUD)
- 超管报考审核页(列表 + 筛选 + 看图 + 通过/驳回 + 导出)
`api.ts` 新增 `nitApi`(学员侧,带 student token 注入)与 `adminApi` NIT 审核扩展。
复用现有 React/Vite/api.ts 范式与 UI 组件。
## 错误处理
- 唯一约束冲突 → 友好提示「该手机号/身份证号已注册」。
- 上传非图片或超大文件 → 拒绝并提示。
- 报考时存在活动报考 → 拒绝重复提交。
- 审核状态机:非 `pending_review` 状态的 approve/reject → 409/冲突提示。
## 测试策略(TDD)
后端(rstest + 内存 SQLite,沿用现有 `recharge_test_pool` 范式):
- 注册手机号/身份证唯一性、校区校验、隐私协议必填。
- 登录成功/失败。
- 文件上传返回 asset id。
- 报考创建 + 重复报考拦截。
- 审核状态机(approve/reject/非法流转)。
- 导出筛选正确性。
前端(vitest):
- `nitApi` 各方法 URL/method/body 正确(沿用 `api.test.ts` 范式)。
- 注册/报考/审核关键页面交互(沿用 `UsageRecords.test.tsx` 的 react-test-renderer 范式)。
## 验证
- `cargo build -j 2`(避免 OOM)、`cargo test -p boosapi --lib nit`(后端)。
- `npx vitest run`(相关前端测试)、`npm run typecheck`。
## 提交约定
NIT 工作全部在 `nit_edu` 分支,不混入 `refactor/rust-backend`。提交按逻辑拆分
(迁移 / service / 路由 / 前端 分开),沿用项目中文 commit 风格。
+212
View File
@@ -1159,3 +1159,215 @@ pub async fn execute_image_edit(
Ok(json) Ok(json)
} }
// ── Gemini native passthrough ───────────────────────────────────────────
//
// Gemini's native protocol carries the model in the URL path
// (`/v1beta/models/{model}:{op}`) and the API key in the `?key=` query, so the
// request body needs no rewrite — we forward it verbatim. This mirrors AQBot's
// `handle_native_request`, which deliberately skips `rewrite_body_model` for
// Gemini, and BoosAPI's own `/v1/responses` passthrough.
/// Normalize an upstream base into a Gemini `v1beta` root. Accepts bases that
/// already end in `/v1beta` (or contain `/v1beta/`) and leaves them as-is;
/// otherwise appends `/v1beta`. Trailing slashes are trimmed.
fn gemini_base_url(base_url: &str) -> String {
let trimmed = base_url.trim().trim_end_matches('/');
if trimmed.ends_with("/v1beta") || trimmed.contains("/v1beta/") {
trimmed.to_string()
} else {
format!("{trimmed}/v1beta")
}
}
/// Build `{base}/v1beta/models/{model}:{op}?key={api_key}[&{extra_query}]`.
/// `extra_query` carries client query params already filtered (e.g. `alt=sse`),
/// without a leading `?`/`&`.
fn gemini_operation_url(
base_url: &str,
model: &str,
operation: &str,
api_key: &str,
extra_query: Option<&str>,
) -> String {
let base = gemini_base_url(base_url);
let model = model.trim().trim_start_matches('/');
let model = model.strip_prefix("models/").unwrap_or(model);
let mut url = format!("{base}/models/{model}:{operation}");
let mut sep = '?';
if !api_key.is_empty() {
url.push(sep);
url.push_str("key=");
url.push_str(api_key);
sep = '&';
}
if let Some(query) = extra_query.filter(|q| !q.is_empty()) {
url.push(sep);
url.push_str(query);
}
url
}
/// Execute a non-streaming Gemini operation (`generateContent` / `countTokens`)
/// and return the upstream JSON verbatim.
pub async fn execute_gemini_operation(
state: AppState,
route: UpstreamRoute,
operation: &str,
body: Value,
extra_query: Option<&str>,
) -> AppResult<Value> {
let client = state.http.clone();
let upstream_url =
gemini_operation_url(&route.base_url, &route.model, operation, &route.api_key, extra_query);
let resp = client
.post(&upstream_url)
.header("Content-Type", "application/json")
.json(&body)
.send()
.await
.map_err(|e| AppError::UpstreamError(format!("Upstream request failed: {e}")))?;
let status = resp.status();
let text = resp
.text()
.await
.map_err(|e| AppError::UpstreamError(format!("Failed to read upstream response: {e}")))?;
if !status.is_success() {
return Err(AppError::UpstreamError(format!(
"Upstream returned {status}: {text}"
)));
}
parse_upstream_json(&text, &upstream_url, status)
}
/// Execute a streaming Gemini operation (`streamGenerateContent`) and return the
/// live upstream response for SSE passthrough. `alt=sse` is enforced in the
/// query so the upstream emits `data:`-framed events the accumulator can parse.
pub async fn execute_gemini_stream_response(
state: AppState,
route: UpstreamRoute,
operation: &str,
body: Value,
extra_query: Option<&str>,
) -> AppResult<reqwest::Response> {
let client = state.http.clone();
let query = merge_alt_sse(extra_query);
let upstream_url = gemini_operation_url(
&route.base_url,
&route.model,
operation,
&route.api_key,
Some(&query),
);
let resp = client
.post(&upstream_url)
.header("Content-Type", "application/json")
.json(&body)
.send()
.await
.map_err(|e| AppError::UpstreamError(format!("Upstream request failed: {e}")))?;
let status = resp.status();
if !status.is_success() {
let text = resp
.text()
.await
.unwrap_or_else(|_| "unknown error".to_string());
return Err(AppError::UpstreamError(format!(
"Upstream returned {status}: {text}"
)));
}
Ok(resp)
}
/// Ensure `alt=sse` is present in the extra query (Gemini streams plain JSON
/// arrays without it). Preserves any other client-supplied params.
fn merge_alt_sse(extra_query: Option<&str>) -> String {
match extra_query.filter(|q| !q.is_empty()) {
Some(q) if q.split('&').any(|p| p == "alt=sse" || p.starts_with("alt=")) => q.to_string(),
Some(q) => format!("{q}&alt=sse"),
None => "alt=sse".to_string(),
}
}
#[cfg(test)]
mod gemini_tests {
use super::*;
#[test]
fn base_url_appends_v1beta_when_missing() {
assert_eq!(
gemini_base_url("https://generativelanguage.googleapis.com"),
"https://generativelanguage.googleapis.com/v1beta"
);
}
#[test]
fn base_url_preserves_existing_v1beta() {
assert_eq!(
gemini_base_url("https://generativelanguage.googleapis.com/v1beta/"),
"https://generativelanguage.googleapis.com/v1beta"
);
}
#[test]
fn operation_url_embeds_model_op_and_key() {
let url = gemini_operation_url(
"https://generativelanguage.googleapis.com",
"gemini-2.0-flash",
"generateContent",
"secret",
None,
);
assert_eq!(
url,
"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent?key=secret"
);
}
#[test]
fn operation_url_strips_models_prefix_and_appends_query() {
let url = gemini_operation_url(
"https://generativelanguage.googleapis.com/v1beta",
"models/gemini-1.5-pro",
"streamGenerateContent",
"k",
Some("alt=sse"),
);
assert_eq!(
url,
"https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-pro:streamGenerateContent?key=k&alt=sse"
);
}
#[test]
fn operation_url_without_key_uses_question_mark_for_query() {
let url = gemini_operation_url(
"https://example.test/v1beta",
"gemini-x",
"countTokens",
"",
Some("foo=bar"),
);
assert_eq!(
url,
"https://example.test/v1beta/models/gemini-x:countTokens?foo=bar"
);
}
#[test]
fn merge_alt_sse_variants() {
assert_eq!(merge_alt_sse(None), "alt=sse");
assert_eq!(merge_alt_sse(Some("")), "alt=sse");
assert_eq!(merge_alt_sse(Some("foo=bar")), "foo=bar&alt=sse");
// already present → unchanged
assert_eq!(merge_alt_sse(Some("alt=sse")), "alt=sse");
assert_eq!(merge_alt_sse(Some("alt=json")), "alt=json");
}
}
+2
View File
@@ -2840,6 +2840,8 @@ mod tests {
template_proxy_url: None, template_proxy_url: None,
}, },
comfyui: Default::default(), comfyui: Default::default(),
payment: Default::default(),
pixelle_video: Default::default(),
}, },
pool, pool,
BuiltinKeys { BuiltinKeys {
+2
View File
@@ -412,6 +412,8 @@ mod tests {
template_proxy_url: None, template_proxy_url: None,
}, },
comfyui: Default::default(), comfyui: Default::default(),
payment: Default::default(),
pixelle_video: Default::default(),
}, },
pool, pool,
BuiltinKeys { BuiltinKeys {
+6
View File
@@ -19,6 +19,11 @@ use serde_json::{json, Value};
use std::net::{IpAddr, SocketAddr}; use std::net::{IpAddr, SocketAddr};
use tracing::{info, warn}; use tracing::{info, warn};
/// Default spending cap (USD) granted to a newly created user's auto-generated
/// API key. Per the billing design (新用户默认 $1 额度), new users start with a
/// $1 allowance on their default downstream key.
pub const DEFAULT_NEW_USER_MAX_COST_USD: f64 = 1.0;
// ── Request / Response types ───────────────────────────────────────── // ── Request / Response types ─────────────────────────────────────────
#[derive(Debug, Deserialize)] #[derive(Debug, Deserialize)]
@@ -1019,6 +1024,7 @@ mod tests {
}, },
comfyui: Default::default(), comfyui: Default::default(),
payment: Default::default(), payment: Default::default(),
pixelle_video: Default::default(),
}, },
pool, pool,
BuiltinKeys { BuiltinKeys {
+2
View File
@@ -127,6 +127,8 @@ mod tests {
template_proxy_url: None, template_proxy_url: None,
}, },
comfyui: ComfyUIConfig::default(), comfyui: ComfyUIConfig::default(),
payment: Default::default(),
pixelle_video: Default::default(),
}, },
pool, pool,
BuiltinKeys { BuiltinKeys {
+2
View File
@@ -1305,6 +1305,8 @@ mod tests {
template_proxy_url: None, template_proxy_url: None,
}, },
comfyui: Default::default(), comfyui: Default::default(),
payment: Default::default(),
pixelle_video: Default::default(),
}, },
pool, pool,
BuiltinKeys { BuiltinKeys {
+2
View File
@@ -536,6 +536,8 @@ mod tests {
template_proxy_url: None, template_proxy_url: None,
}, },
comfyui: Default::default(), comfyui: Default::default(),
payment: Default::default(),
pixelle_video: Default::default(),
}, },
pool, pool,
BuiltinKeys { BuiltinKeys {
+2
View File
@@ -2071,6 +2071,8 @@ mod tests {
template_proxy_url: None, template_proxy_url: None,
}, },
comfyui, comfyui,
payment: Default::default(),
pixelle_video: Default::default(),
}, },
pool, pool,
BuiltinKeys { BuiltinKeys {
+2
View File
@@ -2550,6 +2550,8 @@ mod tests {
template_proxy_url: None, template_proxy_url: None,
}, },
comfyui: ComfyUIConfig::default(), comfyui: ComfyUIConfig::default(),
payment: Default::default(),
pixelle_video: Default::default(),
}, },
pool, pool,
BuiltinKeys { BuiltinKeys {
+509
View File
@@ -73,6 +73,16 @@ pub fn routes() -> Router<AppState> {
.route("/v1/videos", post(videos_handler)) .route("/v1/videos", post(videos_handler))
.route("/v1/videos/{id}", get(video_task_handler)) .route("/v1/videos/{id}", get(video_task_handler))
.route("/v1/models", get(list_models)) .route("/v1/models", get(list_models))
// ── Gemini native inbound (v1beta) ────────────────────────────
// Lets Gemini CLI / Google GenAI SDKs point their base URL straight
// at this gateway. The model lives in the path and the key in the
// `?key=` query (or x-goog-api-key header), so requests pass through
// to a Gemini-shaped upstream without body rewriting.
.route("/v1beta/models", get(gemini_list_models))
.route(
"/v1beta/models/{model_action}",
post(gemini_model_operation),
)
// ── ComfyUI shim aliases ────────────────────────────────────── // ── ComfyUI shim aliases ──────────────────────────────────────
// ComfyUI's OpenAI nodes hit /proxy/openai/... paths (see // ComfyUI's OpenAI nodes hit /proxy/openai/... paths (see
// comfy_api_nodes/nodes_openai.py). Mounting them here as aliases // comfy_api_nodes/nodes_openai.py). Mounting them here as aliases
@@ -157,6 +167,209 @@ async fn responses_handler(
} }
} }
// ── Gemini native inbound handlers (v1beta) ────────────────────────────
//
// Gemini clients (Gemini CLI / Google GenAI SDK) carry the model in the URL
// path (`/v1beta/models/{model}:{op}`) and the gateway key in `?key=` or the
// `x-goog-api-key` header. We authenticate against our downstream keys, route
// the Gemini model through the same ProxyConductor as the OpenAI endpoints, and
// forward the body verbatim to a Gemini-shaped upstream (no body rewrite).
/// Gemini operations carried in the `{model}:{op}` path segment.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum GeminiOperation {
GenerateContent,
StreamGenerateContent,
CountTokens,
}
impl GeminiOperation {
fn as_str(self) -> &'static str {
match self {
GeminiOperation::GenerateContent => "generateContent",
GeminiOperation::StreamGenerateContent => "streamGenerateContent",
GeminiOperation::CountTokens => "countTokens",
}
}
}
/// Parse the `{model}:{operation}` path segment (e.g.
/// `gemini-2.5-pro:streamGenerateContent` or `models/gemini-x:countTokens`).
/// Splits on the last `:` so model names that contain none are rejected.
fn parse_gemini_model_action(model_action: &str) -> Option<(String, GeminiOperation)> {
let (model, action) = model_action.rsplit_once(':')?;
let operation = match action {
"generateContent" => GeminiOperation::GenerateContent,
"streamGenerateContent" => GeminiOperation::StreamGenerateContent,
"countTokens" => GeminiOperation::CountTokens,
_ => return None,
};
if model.trim().is_empty() {
return None;
}
Some((model.to_string(), operation))
}
/// Authenticate a Gemini-native request: the key arrives in `?key=` or the
/// `x-goog-api-key` header (Gemini convention), not `Authorization: Bearer`.
/// Reuses the same downstream-key authorization + limits as `check_auth`.
async fn check_gemini_auth(
state: &AppState,
headers: &HeaderMap,
query: &std::collections::HashMap<String, String>,
) -> Result<Option<i64>, AppError> {
let key = query
.get("key")
.map(|s| s.as_str())
.or_else(|| headers.get("x-goog-api-key").and_then(|v| v.to_str().ok()))
.map(str::trim)
.filter(|s| !s.is_empty())
.ok_or_else(|| {
AppError::Unauthorized("Missing API key (use ?key= or x-goog-api-key)".into())
})?;
let proxy_token = state.runtime_proxy_token();
let result = authorize_downstream_token(key, &state.db, Some(&proxy_token))
.await
.map_err(AppError::Unauthorized)?;
match result {
DownstreamAuthResult::Managed { key } => {
check_downstream_key_limits(&key)?;
Ok(Some(key.id))
}
DownstreamAuthResult::Global { .. } => Ok(None),
}
}
/// Rebuild the upstream query string from client params, dropping the inbound
/// gateway `key` (the upstream key is injected by the executor) and `alt` (the
/// streaming path enforces its own `alt=sse`). Returns `None` when nothing
/// remains. Values are forwarded as-is: real Gemini clients only put `key`/`alt`
/// in the query and carry everything else (generationConfig, etc.) in the body.
fn gemini_forward_query(
query: &std::collections::HashMap<String, String>,
drop_alt: bool,
) -> Option<String> {
let mut parts: Vec<(&String, &String)> = query
.iter()
.filter(|(k, _)| k.as_str() != "key" && !(drop_alt && k.as_str() == "alt"))
.collect();
if parts.is_empty() {
return None;
}
// Stable order so the URL is deterministic (and testable).
parts.sort_by(|a, b| a.0.cmp(b.0));
let encoded = parts
.iter()
.map(|(k, v)| format!("{k}={v}"))
.collect::<Vec<_>>()
.join("&");
Some(encoded)
}
/// GET /v1beta/models — list models in Gemini's native shape.
/// Like `/v1/models`, this is a directory listing from configured routes; no
/// upstream call, no billing.
async fn gemini_list_models(
State(state): State<AppState>,
headers: HeaderMap,
Query(query): Query<std::collections::HashMap<String, String>>,
) -> AppResult<Json<Value>> {
info!("GET /v1beta/models");
let _ = check_gemini_auth(&state, &headers, &query).await?;
let rows: Vec<(String, String)> = sqlx::query_as(
"SELECT DISTINCT
tr.model_pattern,
COALESCE(NULLIF(TRIM(tr.model_mapping), ''), tr.model_pattern)
FROM token_routes tr
JOIN route_channels rc ON rc.route_id = tr.id
JOIN accounts a ON a.id = rc.account_id
JOIN sites s ON s.id = a.site_id
LEFT JOIN account_tokens at ON at.id = rc.token_id AND at.enabled = 1
WHERE tr.enabled = 1
AND rc.enabled = 1
AND a.status = 'active'
AND s.status = 'active'
AND TRIM(COALESCE(s.url, '')) <> ''
AND (rc.cooldown_until IS NULL OR rc.cooldown_until <= datetime('now'))
AND (rc.token_id IS NULL OR at.id IS NOT NULL)
ORDER BY tr.id",
)
.fetch_all(&*state.db)
.await?;
let models: Vec<Value> = rows
.into_iter()
.map(|(_pattern, model)| {
let name = if model.starts_with("models/") {
model.clone()
} else {
format!("models/{model}")
};
json!({
"name": name,
"supportedGenerationMethods": [
"generateContent",
"streamGenerateContent",
"countTokens",
],
})
})
.collect();
Ok(Json(json!({ "models": models })))
}
/// POST /v1beta/models/{model}:{op} — dispatch to the matching Gemini op.
async fn gemini_model_operation(
State(state): State<AppState>,
axum::extract::Path(model_action): axum::extract::Path<String>,
Query(query): Query<std::collections::HashMap<String, String>>,
headers: HeaderMap,
Json(body): Json<Value>,
) -> Result<Response, AppError> {
let Some((model, operation)) = parse_gemini_model_action(&model_action) else {
return Err(AppError::BadRequest(format!(
"Unsupported Gemini model action '{model_action}' \
(expected {{model}}:generateContent|streamGenerateContent|countTokens)"
)));
};
info!("POST /v1beta/models/{model}:{}", operation.as_str());
let key_id = check_gemini_auth(&state, &headers, &query).await?;
match operation {
GeminiOperation::GenerateContent => {
let extra_query = gemini_forward_query(&query, false);
let result =
execute_gemini_with_conductor(state, key_id, model, body, extra_query).await?;
Ok(Json(result).into_response())
}
GeminiOperation::StreamGenerateContent => {
let extra_query = gemini_forward_query(&query, true);
execute_gemini_stream_with_conductor(state, key_id, model, body, extra_query).await
}
GeminiOperation::CountTokens => {
// countTokens is a metadata call (no generation) — route for a
// channel but skip billing, mirroring the lighter helper endpoints.
let extra_query = gemini_forward_query(&query, false);
let route = resolve_route(&state, &model).await?;
let result = executors::execute_gemini_operation(
state,
route,
"countTokens",
body,
extra_query.as_deref(),
)
.await?;
Ok(Json(result).into_response())
}
}
}
/// POST /v1/chat/completions /// POST /v1/chat/completions
async fn chat_completions( async fn chat_completions(
State(state): State<AppState>, State(state): State<AppState>,
@@ -1334,6 +1547,38 @@ fn responses_stream_usage_from_line(line: &str) -> Option<(i64, i64)> {
(prompt != 0 || completion != 0).then_some((prompt, completion)) (prompt != 0 || completion != 0).then_some((prompt, completion))
} }
/// Extract `(prompt, completion)` tokens from a Gemini response body's
/// `usageMetadata` (`promptTokenCount` / `candidatesTokenCount`). Returns
/// `(0, 0)` when the field is absent — mirrors `extract_usage` for OpenAI.
fn gemini_extract_usage(response: &Value) -> (i64, i64) {
let usage = response.get("usageMetadata");
let prompt = usage
.and_then(|u| u.get("promptTokenCount"))
.and_then(|v| v.as_i64())
.unwrap_or(0);
let completion = usage
.and_then(|u| u.get("candidatesTokenCount"))
.and_then(|v| v.as_i64())
.unwrap_or(0);
(prompt, completion)
}
/// Per-line usage parser for Gemini `streamGenerateContent?alt=sse` output.
/// Gemini emits `data:`-framed JSON chunks; the final chunk carries the
/// cumulative `usageMetadata`. Feeds the shared `StreamUsageAccumulator`.
fn gemini_stream_usage_from_line(line: &str) -> Option<(i64, i64)> {
let payload = line
.strip_prefix("data: ")
.or_else(|| line.strip_prefix("data:"))?
.trim();
if payload.is_empty() || payload == "[DONE]" {
return None;
}
let value: Value = serde_json::from_str(payload).ok()?;
let (prompt, completion) = gemini_extract_usage(&value);
(prompt != 0 || completion != 0).then_some((prompt, completion))
}
fn build_conductor(state: &AppState) -> ProxyConductor { fn build_conductor(state: &AppState) -> ProxyConductor {
let db = state.db.clone(); let db = state.db.clone();
let select_db = db.clone(); let select_db = db.clone();
@@ -1881,6 +2126,185 @@ async fn execute_responses_stream_with_conductor(
} }
} }
/// Non-streaming Gemini op (`generateContent`) through the ProxyConductor.
/// Mirrors `execute_responses_with_conductor` but builds a Gemini upstream URL
/// (model in path, key in query) and forwards the body verbatim — Gemini bodies
/// carry no `model` field, so there is no body rewrite.
async fn execute_gemini_with_conductor(
state: AppState,
key_id: Option<i64>,
model: String,
body: Value,
extra_query: Option<String>,
) -> AppResult<Value> {
let started_at = Instant::now();
let request_log_body = body.clone();
let conductor = build_conductor(&state);
let billing_state = state.clone();
let billing_model = model.clone();
let result = conductor
.execute(ExecuteInput {
requested_model: model,
downstream_policy: None,
attempt_fn: Box::new(move |ctx| {
let state = state.clone();
let body = body.clone();
let extra_query = extra_query.clone();
async move {
match executors::execute_gemini_operation(
state,
selected_channel_to_route(ctx.selected),
"generateContent",
body,
extra_query.as_deref(),
)
.await
{
Ok(response) => AttemptResult::Success {
response,
latency_ms: None,
cost: None,
},
Err(err) => AttemptResult::Failure {
action: AttemptFailureAction::Failover,
status: upstream_status_from_error(&err),
raw_error_text: Some(err.to_string()),
},
}
}
.boxed()
}),
on_terminal_failure: None,
})
.await;
let (selected, response, attempts) = match result {
ExecuteResult::Success {
selected,
response,
attempts,
} => (selected, response, attempts),
ExecuteResult::Failure {
reason,
selected,
status,
raw_error_text,
attempts,
} => {
record_proxy_log_failure(
&billing_state,
key_id,
selected.as_ref(),
&billing_model,
status,
Some(elapsed_ms_since(started_at)),
raw_error_text.as_deref(),
attempts,
false,
proxy_log_content_from_request(&request_log_body),
)
.await;
return Err(AppError::UpstreamError(format!(
"Proxy conductor failed after {attempts} attempts ({reason:?}, status: {:?}): {}",
status,
raw_error_text.unwrap_or_else(|| "no upstream error text".to_string())
)));
}
};
let (prompt, completion) = gemini_extract_usage(&response);
let cost = estimate_request_cost(&selected, prompt, completion);
record_proxy_log_success(
&billing_state,
key_id,
&selected,
&billing_model,
Some(elapsed_ms_since(started_at)),
prompt,
completion,
cost,
attempts,
false,
proxy_log_content_from_request_response(&request_log_body, Some(&response)),
)
.await;
record_billing(
&billing_state,
key_id,
&billing_model,
prompt,
completion,
cost,
)
.await;
Ok(response)
}
/// Streaming Gemini op (`streamGenerateContent`) through the ProxyConductor.
/// Mirrors `execute_responses_stream_with_conductor`, opening the upstream with
/// failover and SSE-passing it through with the Gemini usage parser so streamed
/// `usageMetadata` is billed.
async fn execute_gemini_stream_with_conductor(
state: AppState,
key_id: Option<i64>,
model: String,
body: Value,
extra_query: Option<String>,
) -> Result<Response, AppError> {
let started_at = Instant::now();
let request_log_body = body.clone();
let body_for_attempt = body.clone();
let opened = open_stream_with_failover(&state, &model, move |state, selected| {
let body = body_for_attempt.clone();
let extra_query = extra_query.clone();
async move {
executors::execute_gemini_stream_response(
state,
selected_channel_to_route(selected),
"streamGenerateContent",
body,
extra_query.as_deref(),
)
.await
}
})
.await;
match opened {
Ok(opened) => sse_passthrough_response(
state,
opened,
key_id,
model,
started_at,
proxy_log_content_from_request(&request_log_body),
gemini_stream_usage_from_line,
),
Err(failure) => {
record_proxy_log_failure(
&state,
key_id,
failure.selected.as_ref(),
&model,
failure.status,
Some(elapsed_ms_since(started_at)),
failure.raw_error_text.as_deref(),
failure.attempts,
true,
proxy_log_content_from_request(&request_log_body),
)
.await;
Err(AppError::UpstreamError(format!(
"Proxy conductor failed after {} attempts (stream open, status: {:?}): {}",
failure.attempts,
failure.status,
failure
.raw_error_text
.unwrap_or_else(|| "no upstream error text".to_string())
)))
}
}
}
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::*; use super::*;
@@ -2121,6 +2545,8 @@ mod tests {
template_proxy_url: None, template_proxy_url: None,
}, },
comfyui: Default::default(), comfyui: Default::default(),
payment: Default::default(),
pixelle_video: Default::default(),
}, },
pool, pool,
crate::services::builtin_keys::BuiltinKeys { crate::services::builtin_keys::BuiltinKeys {
@@ -2155,4 +2581,87 @@ mod tests {
.unwrap(); .unwrap();
assert!(fail_count > 0); assert!(fail_count > 0);
} }
// ── Gemini native inbound unit tests ───────────────────────────────
#[test]
fn parse_gemini_model_action_known_ops() {
let (m, op) = parse_gemini_model_action("gemini-2.5-pro:generateContent").unwrap();
assert_eq!(m, "gemini-2.5-pro");
assert_eq!(op, GeminiOperation::GenerateContent);
let (m, op) = parse_gemini_model_action("gemini-2.5-pro:streamGenerateContent").unwrap();
assert_eq!(m, "gemini-2.5-pro");
assert_eq!(op, GeminiOperation::StreamGenerateContent);
// model names may carry the `models/` prefix
let (m, op) = parse_gemini_model_action("models/gemini-1.5-pro:countTokens").unwrap();
assert_eq!(m, "models/gemini-1.5-pro");
assert_eq!(op, GeminiOperation::CountTokens);
}
#[test]
fn parse_gemini_model_action_rejects_bad_input() {
// no operation suffix
assert!(parse_gemini_model_action("gemini-2.5-pro").is_none());
// unknown operation
assert!(parse_gemini_model_action("gemini-2.5-pro:embedContent").is_none());
// empty model
assert!(parse_gemini_model_action(":generateContent").is_none());
}
#[test]
fn gemini_extract_usage_reads_usage_metadata() {
let body = json!({
"candidates": [],
"usageMetadata": {
"promptTokenCount": 11,
"candidatesTokenCount": 22,
"totalTokenCount": 33
}
});
assert_eq!(gemini_extract_usage(&body), (11, 22));
// absent usageMetadata → zeros
assert_eq!(gemini_extract_usage(&json!({})), (0, 0));
}
#[test]
fn gemini_stream_usage_from_line_parses_data_frame() {
let line = r#"data: {"candidates":[{"content":{"parts":[]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":4,"candidatesTokenCount":6,"totalTokenCount":10}}"#;
assert_eq!(gemini_stream_usage_from_line(line), Some((4, 6)));
}
#[test]
fn gemini_stream_usage_from_line_ignores_non_usage_and_done() {
// a chunk with no usageMetadata yields None (don't overwrite accumulator)
let no_usage = r#"data: {"candidates":[{"content":{"parts":[{"text":"hi"}]}}]}"#;
assert_eq!(gemini_stream_usage_from_line(no_usage), None);
// sentinels / blanks / non-data lines
assert_eq!(gemini_stream_usage_from_line("data: [DONE]"), None);
assert_eq!(gemini_stream_usage_from_line("data:"), None);
assert_eq!(gemini_stream_usage_from_line(": comment"), None);
}
#[test]
fn gemini_forward_query_drops_key_and_alt() {
let mut q = std::collections::HashMap::new();
q.insert("key".to_string(), "secret".to_string());
q.insert("alt".to_string(), "sse".to_string());
// stream path drops both key and alt → nothing left
assert_eq!(gemini_forward_query(&q, true), None);
// non-stream keeps alt but always drops key
assert_eq!(gemini_forward_query(&q, false), Some("alt=sse".to_string()));
}
#[test]
fn gemini_forward_query_preserves_other_params_sorted() {
let mut q = std::collections::HashMap::new();
q.insert("key".to_string(), "secret".to_string());
q.insert("pageSize".to_string(), "50".to_string());
q.insert("filter".to_string(), "x".to_string());
assert_eq!(
gemini_forward_query(&q, true),
Some("filter=x&pageSize=50".to_string())
);
}
} }
+2
View File
@@ -696,6 +696,8 @@ mod tests {
template_proxy_url: None, template_proxy_url: None,
}, },
comfyui: ComfyUIConfig::default(), comfyui: ComfyUIConfig::default(),
payment: Default::default(),
pixelle_video: Default::default(),
}, },
pool, pool,
BuiltinKeys { BuiltinKeys {