feat: Phase 5-6 - Proxy Conductor, Token Router, Services, Dockerfile
- Proxy Conductor: retry/fallback/auth-refresh orchestration loop - Token Router: exact/wildcard/regex model routing with channel selection - Services: Cron scheduler, OAuth token refresh, Update center, Announcements - User routes: profile + API key management - DB queries module - Dockerfile: multi-stage build - Proxy routes integrated with Conductor + TokenRouter
This commit is contained in:
@@ -28,6 +28,7 @@ reqwest = { version = "0.12", features = ["stream", "json"] }
|
|||||||
# Serialization
|
# Serialization
|
||||||
serde = { version = "1", features = ["derive"] }
|
serde = { version = "1", features = ["derive"] }
|
||||||
serde_json = "1"
|
serde_json = "1"
|
||||||
|
regex = "1"
|
||||||
|
|
||||||
# Auth
|
# Auth
|
||||||
jsonwebtoken = "9"
|
jsonwebtoken = "9"
|
||||||
|
|||||||
@@ -0,0 +1,16 @@
|
|||||||
|
FROM rust:1.83-alpine AS builder
|
||||||
|
RUN apk add --no-cache musl-dev pkgconfig openssl-dev
|
||||||
|
WORKDIR /app
|
||||||
|
COPY Cargo.toml Cargo.lock ./
|
||||||
|
RUN mkdir src && echo "fn main() {}" > src/main.rs
|
||||||
|
RUN cargo build --release 2>/dev/null || true
|
||||||
|
COPY . .
|
||||||
|
RUN cargo build --release
|
||||||
|
|
||||||
|
FROM alpine:3.20
|
||||||
|
RUN apk add --no-cache ca-certificates tzdata sqlite-libs
|
||||||
|
WORKDIR /app
|
||||||
|
COPY --from=builder /app/target/release/boosapi .
|
||||||
|
COPY --from=builder /app/migrations ./migrations
|
||||||
|
EXPOSE 4000
|
||||||
|
CMD ["./boosapi"]
|
||||||
@@ -0,0 +1,164 @@
|
|||||||
|
# 2026-05-15 技术报告 / Technical Report
|
||||||
|
|
||||||
|
## Executive Summary
|
||||||
|
|
||||||
|
本日完成 BoosAPI 品牌迁移、ComfyUI Agent 对话助手开发、Gitea CI/CD Runner 部署、Notification Service 能力梳理及部署文档重写。核心变更包括:将项目内 `Metapi` 全局替换为 `BoosAPI`,新增 `/comfyui-agent` 对话式 AI 视频生成入口,完成 SSE streaming 前后端链路,并完善生产部署文档与通知服务说明。
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## 1. Metapi → BoosAPI 全局改名 / Global Rename
|
||||||
|
|
||||||
|
### 1.1 变更范围 / Change Scope
|
||||||
|
|
||||||
|
| 区域 / Area | 变更内容 / Changes |
|
||||||
|
|---|---|
|
||||||
|
| UI 界面 / UI | About 页面、侧边栏、登录页品牌名更新为 `BoosAPI` |
|
||||||
|
| i18n 翻译 / Internationalization | 所有语言文件中的 `Metapi` 替换为 `BoosAPI` |
|
||||||
|
| 后端通知 / Backend Notification | `notifyService.ts` 通知前缀由 `[metapi]` 改为 `[BoosAPI]` |
|
||||||
|
| 桌面元数据 / Desktop Metadata | Electron / Docker 描述信息更新 |
|
||||||
|
|
||||||
|
### 1.2 已完成的技术标识替换 / Completed Technical Identifier Changes
|
||||||
|
|
||||||
|
| 项目 / Item | 变更说明 / Description |
|
||||||
|
|---|---|
|
||||||
|
| HTTP Header | `x-metapi-*` → `x-boosapi-*` |
|
||||||
|
| 推理签名前缀 / Inference Signature Prefix | `metapi:` → `boosapi:` |
|
||||||
|
| 运行时属性 / Runtime Properties | `__metapiTruncated` 等 → `__boosapi*` |
|
||||||
|
| 环境变量 / Environment Variables | `METAPI_*` → `BOOSAPI_*` |
|
||||||
|
| Package Name | `package.json` 中 `name` 更新为 `boosapi` |
|
||||||
|
|
||||||
|
### 1.3 验证方式 / Verification
|
||||||
|
|
||||||
|
```bash
|
||||||
|
grep -rn '[mM]etapi' src/ --include='*.ts' --include='*.tsx' --include='*.html' \
|
||||||
|
| grep -v '.test.' | grep -v node_modules | grep -iv 'boosapi'
|
||||||
|
```
|
||||||
|
|
||||||
|
预期结果:无输出。
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## 2. ComfyUI Agent 对话助手 / Conversational Assistant
|
||||||
|
|
||||||
|
### 2.1 功能概述 / Feature Overview
|
||||||
|
|
||||||
|
侧边栏新增「对话助手」入口,路由为 `/comfyui-agent`,提供与 AI 视频生成助手的对话式交互,支持文本、角色提示词、图像、场景、语音和 ComfyUI workflow 等多类型内容的流式渲染。
|
||||||
|
|
||||||
|
### 2.2 技术架构 / Architecture
|
||||||
|
|
||||||
|
```
|
||||||
|
用户 / User → POST /api/comfyui-agent/chat → BoosAPI → LLM 本地 API (streaming)
|
||||||
|
→ SSE 事件流 → BoosAPI 转发 → 浏览器
|
||||||
|
```
|
||||||
|
|
||||||
|
### 2.3 后端实现 / Backend
|
||||||
|
|
||||||
|
| 文件 / File | 作用 / Responsibility |
|
||||||
|
|---|---|
|
||||||
|
| `src/server/routes/api/comfyuiAgent.ts` | SSE 流式聊天接口,调用本地 LLM API |
|
||||||
|
| `src/server/services/boosapi_comfy_nodes.py` | ComfyUI 自定义节点,默认连接 `http://127.0.0.1:4000/v1` |
|
||||||
|
|
||||||
|
### 2.4 前端实现 / Frontend
|
||||||
|
|
||||||
|
| 文件 / File | 作用 / Responsibility |
|
||||||
|
|---|---|
|
||||||
|
| `src/web/pages/ComfyUIAgent.tsx` | 对话界面组件,SSE 流式渲染 |
|
||||||
|
| `src/web/App.tsx` | 注册 `/comfyui-agent` 路由及侧边栏入口 |
|
||||||
|
|
||||||
|
### 2.5 SSE 事件类型 / SSE Event Types
|
||||||
|
|
||||||
|
| `type` | 含义 / Meaning | 前端渲染 / Rendering |
|
||||||
|
|---|---|---|
|
||||||
|
| `text` | 流式文本 | 追加到对话内容 |
|
||||||
|
| `character_prompts` | 四角度角色提示词 | 角色卡片 |
|
||||||
|
| `image` | 生成的角色图像 | 内联图片 |
|
||||||
|
| `scene` | 场景提示词 | 场景块 |
|
||||||
|
| `tts` | 语音音频 | 音频播放器 |
|
||||||
|
| `comfyui_workflow` | 完整 ComfyUI 工作流 JSON | 下载按钮 |
|
||||||
|
| `done` | 流结束 | 停止等待 |
|
||||||
|
|
||||||
|
### 2.6 验证结果 / Verification
|
||||||
|
|
||||||
|
| 命令 | 结果 |
|
||||||
|
|---|---|
|
||||||
|
| `npm install` | ✅ |
|
||||||
|
| `npm run typecheck:web` | ✅ 0 error |
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## 3. Gitea CI/CD Runner 部署 / Runner Deployment
|
||||||
|
|
||||||
|
### 3.1 部署信息 / Deployment Info
|
||||||
|
|
||||||
|
| 项目 | 详情 |
|
||||||
|
|---|---|
|
||||||
|
| 系统 | Alpine Linux, OpenRC |
|
||||||
|
| 部署方式 | 裸机直跑,无 DinD |
|
||||||
|
| Runner 版本 | v0.1.3 |
|
||||||
|
| 注册模式 | `:host` 标签 |
|
||||||
|
| OpenRC 服务 | `/etc/init.d/gitea-runner` |
|
||||||
|
| 代理 | `HTTP_PROXY=http://127.0.0.1:7890` |
|
||||||
|
| `NO_PROXY` | 包含 `git.boos.lat` |
|
||||||
|
| 持久化 | `ephemeral: false` |
|
||||||
|
|
||||||
|
### 3.2 已知限制 / Known Limitations
|
||||||
|
|
||||||
|
- 不支持 Docker in Docker;CI 使用 `runs-on: host`
|
||||||
|
- 无原生 artifact 存储;构建产物需手动备份
|
||||||
|
- `actions/checkout` 偶发 timeout;可重试或使用本地 git 镜像
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## 4. Notification Service
|
||||||
|
|
||||||
|
### 4.1 支持通道 / Supported Channels
|
||||||
|
|
||||||
|
| 通道 | 配置前缀 | 协议 |
|
||||||
|
|---|---|---|
|
||||||
|
| Telegram | `telegram*` | Bot API |
|
||||||
|
| 企业微信 | `wecom*` | Webhook |
|
||||||
|
| 飞书 | `feishu*` | Webhook |
|
||||||
|
| 通用 Webhook | `webhook*` | JSON POST |
|
||||||
|
| Bark | `bark*` | HTTP, iOS |
|
||||||
|
| Server酱 | `serverchan*` | HTTP |
|
||||||
|
| SMTP | `smtp*` | Email |
|
||||||
|
|
||||||
|
### 4.2 Telegram 消息格式
|
||||||
|
|
||||||
|
- 纯文本,不启用 HTML/Markdown 解析
|
||||||
|
- 单条最大 3900 字符,超长截断 + `...(truncated)`
|
||||||
|
- 支持 Telegram Topics(`message_thread_id`)
|
||||||
|
|
||||||
|
### 4.3 节流机制 / Throttling
|
||||||
|
|
||||||
|
- 基于通知签名的全局冷静期:`notifyCooldownSec`
|
||||||
|
- 同一签名在冷静期内重复触发时合并为一次
|
||||||
|
- 冷静期为全局共享,不按通道独立
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## 5. 部署文档完善 / Deployment Documentation
|
||||||
|
|
||||||
|
Commit: `5bd847d`
|
||||||
|
|
||||||
|
对 `docs/deployment.md` 全量重写,修复所有 `metapi` 残留:
|
||||||
|
|
||||||
|
- 完整 `docker-compose.yml` 示例(含 healthcheck)
|
||||||
|
- `.env` 变量插值模式
|
||||||
|
- 全部环境变量参考表(20+ 变量)
|
||||||
|
- 扩展配置:外部 MySQL/PostgreSQL、Redis、资源限制
|
||||||
|
- 生产部署检查清单
|
||||||
|
- Nginx 反向代理(WebSocket + SSE)
|
||||||
|
- 管理后台 IP 白名单限制示例
|
||||||
|
- Docker run 常用选项表
|
||||||
|
- 升级/回滚操作流程
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## 6. 待办事项 / TODO
|
||||||
|
|
||||||
|
| 优先级 | 事项 | 说明 |
|
||||||
|
|---|---|---|
|
||||||
|
| High | 服务未运行 | CI 崩溃后容器未恢复,需确定 Docker vs `npm run dev` |
|
||||||
|
| Medium | Telegram 格式优化 | 待确认 HTML/Markdown/简化 |
|
||||||
|
| Medium | CI pipeline 稳定性 | Docker 镜像拉取超时,需排查网络 |
|
||||||
@@ -8,13 +8,9 @@ use crate::db::models::downstream_api_key::DownstreamApiKey;
|
|||||||
#[derive(Debug, Clone)]
|
#[derive(Debug, Clone)]
|
||||||
pub enum DownstreamAuthResult {
|
pub enum DownstreamAuthResult {
|
||||||
/// Found in downstream_api_keys table.
|
/// Found in downstream_api_keys table.
|
||||||
Managed {
|
Managed { key: DownstreamApiKey },
|
||||||
key: DownstreamApiKey,
|
|
||||||
},
|
|
||||||
/// Matched the global proxy token from config.
|
/// Matched the global proxy token from config.
|
||||||
Global {
|
Global { token: String },
|
||||||
token: String,
|
|
||||||
},
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Authorize a downstream token (sk-xxx or global proxy token).
|
/// Authorize a downstream token (sk-xxx or global proxy token).
|
||||||
@@ -29,7 +25,7 @@ pub async fn authorize_downstream_token(
|
|||||||
// Try managed keys first
|
// Try managed keys first
|
||||||
if token.starts_with("sk-") {
|
if token.starts_with("sk-") {
|
||||||
let key: Option<DownstreamApiKey> = sqlx::query_as::<_, DownstreamApiKey>(
|
let key: Option<DownstreamApiKey> = sqlx::query_as::<_, DownstreamApiKey>(
|
||||||
"SELECT * FROM downstream_api_keys WHERE key = ? AND enabled = 1"
|
"SELECT * FROM downstream_api_keys WHERE key = ? AND enabled = 1",
|
||||||
)
|
)
|
||||||
.bind(token)
|
.bind(token)
|
||||||
.fetch_optional(pool)
|
.fetch_optional(pool)
|
||||||
|
|||||||
@@ -1,6 +1,7 @@
|
|||||||
use sqlx::sqlite::{SqlitePool, SqlitePoolOptions};
|
use sqlx::sqlite::{SqlitePool, SqlitePoolOptions};
|
||||||
|
|
||||||
pub mod models;
|
pub mod models;
|
||||||
|
pub mod queries;
|
||||||
|
|
||||||
pub async fn create_pool(database_url: &str, max_connections: u32) -> anyhow::Result<SqlitePool> {
|
pub async fn create_pool(database_url: &str, max_connections: u32) -> anyhow::Result<SqlitePool> {
|
||||||
let pool = SqlitePoolOptions::new()
|
let pool = SqlitePoolOptions::new()
|
||||||
|
|||||||
@@ -1,13 +1,13 @@
|
|||||||
// Database models
|
// Database models
|
||||||
|
|
||||||
pub mod account;
|
pub mod account;
|
||||||
|
pub mod background_task;
|
||||||
pub mod downstream_api_key;
|
pub mod downstream_api_key;
|
||||||
pub mod event;
|
pub mod event;
|
||||||
|
pub mod organization;
|
||||||
pub mod route_channel;
|
pub mod route_channel;
|
||||||
pub mod setting;
|
pub mod setting;
|
||||||
pub mod site;
|
pub mod site;
|
||||||
pub mod site_announcement;
|
pub mod site_announcement;
|
||||||
pub mod background_task;
|
|
||||||
pub mod organization;
|
|
||||||
pub mod token_route;
|
pub mod token_route;
|
||||||
pub mod user;
|
pub mod user;
|
||||||
|
|||||||
@@ -1 +1,40 @@
|
|||||||
// Database queries
|
use crate::{
|
||||||
|
db::models::{downstream_api_key::DownstreamApiKey, user::User},
|
||||||
|
error::AppResult,
|
||||||
|
};
|
||||||
|
use sqlx::SqlitePool;
|
||||||
|
|
||||||
|
pub async fn find_user_by_id(db: &SqlitePool, id: i64) -> AppResult<Option<User>> {
|
||||||
|
sqlx::query_as("SELECT * FROM users WHERE id = ?")
|
||||||
|
.bind(id)
|
||||||
|
.fetch_optional(db)
|
||||||
|
.await
|
||||||
|
.map_err(Into::into)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn find_user_by_email(db: &SqlitePool, email: &str) -> AppResult<Option<User>> {
|
||||||
|
sqlx::query_as("SELECT * FROM users WHERE email = ?")
|
||||||
|
.bind(email)
|
||||||
|
.fetch_optional(db)
|
||||||
|
.await
|
||||||
|
.map_err(Into::into)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn list_user_api_keys(db: &SqlitePool, user_id: i64) -> AppResult<Vec<DownstreamApiKey>> {
|
||||||
|
sqlx::query_as("SELECT * FROM downstream_api_keys WHERE user_id = ? ORDER BY created_at DESC")
|
||||||
|
.bind(user_id)
|
||||||
|
.fetch_all(db)
|
||||||
|
.await
|
||||||
|
.map_err(Into::into)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn find_downstream_key_by_token(
|
||||||
|
db: &SqlitePool,
|
||||||
|
token: &str,
|
||||||
|
) -> AppResult<Option<DownstreamApiKey>> {
|
||||||
|
sqlx::query_as("SELECT * FROM downstream_api_keys WHERE key = ?")
|
||||||
|
.bind(token)
|
||||||
|
.fetch_optional(db)
|
||||||
|
.await
|
||||||
|
.map_err(Into::into)
|
||||||
|
}
|
||||||
|
|||||||
@@ -1 +1,55 @@
|
|||||||
// Proxy capabilities
|
use crate::proxy::providers::types::ProviderCategory;
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, PartialEq)]
|
||||||
|
pub struct ProviderCapabilities {
|
||||||
|
pub category: ProviderCategory,
|
||||||
|
pub supports_chat: bool,
|
||||||
|
pub supports_responses: bool,
|
||||||
|
pub supports_images: bool,
|
||||||
|
pub supports_audio: bool,
|
||||||
|
pub supports_streaming: bool,
|
||||||
|
pub supports_token_refresh: bool,
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn detect_capabilities(platform: &str) -> ProviderCapabilities {
|
||||||
|
let normalized = platform.to_ascii_lowercase();
|
||||||
|
let category = match normalized.as_str() {
|
||||||
|
"openai" => ProviderCategory::OpenAi,
|
||||||
|
"anthropic" | "claude" => ProviderCategory::Anthropic,
|
||||||
|
"gemini" | "google" => ProviderCategory::Gemini,
|
||||||
|
"codex" => ProviderCategory::Codex,
|
||||||
|
"gemini_cli" | "gemini-cli" => ProviderCategory::GeminiCli,
|
||||||
|
"antigravity" => ProviderCategory::Antigravity,
|
||||||
|
_ => ProviderCategory::GenericOpenAi,
|
||||||
|
};
|
||||||
|
|
||||||
|
match category {
|
||||||
|
ProviderCategory::Anthropic => ProviderCapabilities {
|
||||||
|
category,
|
||||||
|
supports_chat: true,
|
||||||
|
supports_responses: false,
|
||||||
|
supports_images: false,
|
||||||
|
supports_audio: false,
|
||||||
|
supports_streaming: true,
|
||||||
|
supports_token_refresh: true,
|
||||||
|
},
|
||||||
|
ProviderCategory::Gemini | ProviderCategory::GeminiCli => ProviderCapabilities {
|
||||||
|
category,
|
||||||
|
supports_chat: true,
|
||||||
|
supports_responses: false,
|
||||||
|
supports_images: true,
|
||||||
|
supports_audio: false,
|
||||||
|
supports_streaming: true,
|
||||||
|
supports_token_refresh: true,
|
||||||
|
},
|
||||||
|
_ => ProviderCapabilities {
|
||||||
|
category,
|
||||||
|
supports_chat: true,
|
||||||
|
supports_responses: true,
|
||||||
|
supports_images: true,
|
||||||
|
supports_audio: true,
|
||||||
|
supports_streaming: true,
|
||||||
|
supports_token_refresh: true,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -1 +1,43 @@
|
|||||||
// Proxy CLI profiles
|
use serde::{Deserialize, Serialize};
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||||
|
pub enum CliProfileKind {
|
||||||
|
Codex,
|
||||||
|
ClaudeCode,
|
||||||
|
GeminiCli,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||||
|
pub struct CliProfile {
|
||||||
|
pub id: &'static str,
|
||||||
|
pub kind: CliProfileKind,
|
||||||
|
pub display_name: &'static str,
|
||||||
|
pub default_model: &'static str,
|
||||||
|
pub env_key_name: &'static str,
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn builtin_profiles() -> Vec<CliProfile> {
|
||||||
|
vec![
|
||||||
|
CliProfile {
|
||||||
|
id: "codex",
|
||||||
|
kind: CliProfileKind::Codex,
|
||||||
|
display_name: "Codex CLI",
|
||||||
|
default_model: "gpt-5.5",
|
||||||
|
env_key_name: "OPENAI_API_KEY",
|
||||||
|
},
|
||||||
|
CliProfile {
|
||||||
|
id: "claude-code",
|
||||||
|
kind: CliProfileKind::ClaudeCode,
|
||||||
|
display_name: "Claude Code",
|
||||||
|
default_model: "claude-sonnet-4",
|
||||||
|
env_key_name: "ANTHROPIC_API_KEY",
|
||||||
|
},
|
||||||
|
CliProfile {
|
||||||
|
id: "gemini-cli",
|
||||||
|
kind: CliProfileKind::GeminiCli,
|
||||||
|
display_name: "Gemini CLI",
|
||||||
|
default_model: "gemini-2.0-flash",
|
||||||
|
env_key_name: "GEMINI_API_KEY",
|
||||||
|
},
|
||||||
|
]
|
||||||
|
}
|
||||||
|
|||||||
@@ -0,0 +1,411 @@
|
|||||||
|
use super::retry_policy::classify_failure;
|
||||||
|
use super::types::{
|
||||||
|
AttemptFailure, AttemptFailureAction, AttemptResult, ExecuteAttemptContext, ExecuteInput,
|
||||||
|
ExecuteResult, FailureInfo, FailureReason, SelectedChannel, SuccessMetrics,
|
||||||
|
};
|
||||||
|
use futures::future::BoxFuture;
|
||||||
|
use serde_json::Value;
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
const MAX_ATTEMPTS: usize = 5;
|
||||||
|
|
||||||
|
pub struct ProxyConductor {
|
||||||
|
pub select_channel: Arc<
|
||||||
|
dyn Fn(String, Option<Value>) -> BoxFuture<'static, Option<SelectedChannel>> + Send + Sync,
|
||||||
|
>,
|
||||||
|
pub select_next_channel: Arc<
|
||||||
|
dyn Fn(String, Vec<i64>, Option<Value>) -> BoxFuture<'static, Option<SelectedChannel>>
|
||||||
|
+ Send
|
||||||
|
+ Sync,
|
||||||
|
>,
|
||||||
|
pub record_success: Arc<dyn Fn(i64, SuccessMetrics) -> BoxFuture<'static, ()> + Send + Sync>,
|
||||||
|
pub record_failure: Arc<dyn Fn(i64, FailureInfo) -> BoxFuture<'static, ()> + Send + Sync>,
|
||||||
|
pub refresh_auth: Option<
|
||||||
|
Arc<
|
||||||
|
dyn Fn(SelectedChannel, FailureInfo) -> BoxFuture<'static, Option<SelectedChannel>>
|
||||||
|
+ Send
|
||||||
|
+ Sync,
|
||||||
|
>,
|
||||||
|
>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ProxyConductor {
|
||||||
|
pub async fn execute(&self, input: ExecuteInput) -> ExecuteResult {
|
||||||
|
let mut attempts = 0;
|
||||||
|
let mut exclude_channel_ids = Vec::new();
|
||||||
|
let mut selected = match (self.select_channel)(
|
||||||
|
input.requested_model.clone(),
|
||||||
|
input.downstream_policy.clone(),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
Some(selected) => selected,
|
||||||
|
None => {
|
||||||
|
return ExecuteResult::Failure {
|
||||||
|
reason: FailureReason::NoChannel,
|
||||||
|
selected: None,
|
||||||
|
status: None,
|
||||||
|
raw_error_text: None,
|
||||||
|
attempts,
|
||||||
|
};
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
loop {
|
||||||
|
if attempts >= MAX_ATTEMPTS {
|
||||||
|
return exhausted_failure(
|
||||||
|
selected,
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
attempts,
|
||||||
|
input.on_terminal_failure.as_deref(),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
}
|
||||||
|
|
||||||
|
let ctx = ExecuteAttemptContext {
|
||||||
|
selected: selected.clone(),
|
||||||
|
attempt_index: attempts,
|
||||||
|
exclude_channel_ids: exclude_channel_ids.clone(),
|
||||||
|
};
|
||||||
|
attempts += 1;
|
||||||
|
|
||||||
|
match (input.attempt_fn)(ctx).await {
|
||||||
|
AttemptResult::Success {
|
||||||
|
response,
|
||||||
|
latency_ms,
|
||||||
|
cost,
|
||||||
|
} => {
|
||||||
|
(self.record_success)(
|
||||||
|
selected.channel_id,
|
||||||
|
SuccessMetrics {
|
||||||
|
latency_ms,
|
||||||
|
cost,
|
||||||
|
metadata: Default::default(),
|
||||||
|
},
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
return ExecuteResult::Success {
|
||||||
|
selected,
|
||||||
|
response,
|
||||||
|
attempts,
|
||||||
|
};
|
||||||
|
}
|
||||||
|
AttemptResult::Failure {
|
||||||
|
action,
|
||||||
|
status,
|
||||||
|
raw_error_text,
|
||||||
|
} => {
|
||||||
|
let classified = if matches!(action, AttemptFailureAction::Stop) {
|
||||||
|
AttemptFailureAction::Stop
|
||||||
|
} else {
|
||||||
|
classify_failure(status, raw_error_text.as_deref())
|
||||||
|
};
|
||||||
|
let info = FailureInfo {
|
||||||
|
action: classified.as_str().to_string(),
|
||||||
|
status,
|
||||||
|
raw_error_text: raw_error_text.clone(),
|
||||||
|
};
|
||||||
|
(self.record_failure)(selected.channel_id, info.clone()).await;
|
||||||
|
|
||||||
|
match classified {
|
||||||
|
AttemptFailureAction::RetrySameChannel => {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
AttemptFailureAction::RefreshAuth => {
|
||||||
|
if let Some(refresh_auth) = &self.refresh_auth {
|
||||||
|
if let Some(refreshed) =
|
||||||
|
refresh_auth(selected.clone(), info.clone()).await
|
||||||
|
{
|
||||||
|
selected = refreshed;
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
exclude_channel_ids.push(selected.channel_id);
|
||||||
|
match (self.select_next_channel)(
|
||||||
|
input.requested_model.clone(),
|
||||||
|
exclude_channel_ids.clone(),
|
||||||
|
input.downstream_policy.clone(),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
Some(next) => {
|
||||||
|
selected = next;
|
||||||
|
}
|
||||||
|
None => {
|
||||||
|
return terminal_failure(
|
||||||
|
selected,
|
||||||
|
FailureReason::Failed,
|
||||||
|
status,
|
||||||
|
raw_error_text,
|
||||||
|
attempts,
|
||||||
|
input.on_terminal_failure.as_deref(),
|
||||||
|
&info,
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
AttemptFailureAction::Failover => {
|
||||||
|
exclude_channel_ids.push(selected.channel_id);
|
||||||
|
match (self.select_next_channel)(
|
||||||
|
input.requested_model.clone(),
|
||||||
|
exclude_channel_ids.clone(),
|
||||||
|
input.downstream_policy.clone(),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
Some(next) => {
|
||||||
|
selected = next;
|
||||||
|
}
|
||||||
|
None => {
|
||||||
|
return terminal_failure(
|
||||||
|
selected,
|
||||||
|
FailureReason::Failed,
|
||||||
|
status,
|
||||||
|
raw_error_text,
|
||||||
|
attempts,
|
||||||
|
input.on_terminal_failure.as_deref(),
|
||||||
|
&info,
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
AttemptFailureAction::Terminal => {
|
||||||
|
return terminal_failure(
|
||||||
|
selected,
|
||||||
|
FailureReason::Terminal,
|
||||||
|
status,
|
||||||
|
raw_error_text,
|
||||||
|
attempts,
|
||||||
|
input.on_terminal_failure.as_deref(),
|
||||||
|
&info,
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
}
|
||||||
|
AttemptFailureAction::Stop => {
|
||||||
|
return terminal_failure(
|
||||||
|
selected,
|
||||||
|
FailureReason::Failed,
|
||||||
|
status,
|
||||||
|
raw_error_text,
|
||||||
|
attempts,
|
||||||
|
input.on_terminal_failure.as_deref(),
|
||||||
|
&info,
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn terminal_failure(
|
||||||
|
selected: SelectedChannel,
|
||||||
|
reason: FailureReason,
|
||||||
|
status: Option<u16>,
|
||||||
|
raw_error_text: Option<String>,
|
||||||
|
attempts: usize,
|
||||||
|
on_terminal_failure: Option<
|
||||||
|
&(dyn Fn(SelectedChannel, AttemptFailure) -> BoxFuture<'static, ()> + Send + Sync),
|
||||||
|
>,
|
||||||
|
info: &FailureInfo,
|
||||||
|
) -> ExecuteResult {
|
||||||
|
if let Some(callback) = on_terminal_failure {
|
||||||
|
callback(selected.clone(), AttemptFailure::from(info)).await;
|
||||||
|
}
|
||||||
|
ExecuteResult::Failure {
|
||||||
|
reason,
|
||||||
|
selected: Some(selected),
|
||||||
|
status,
|
||||||
|
raw_error_text,
|
||||||
|
attempts,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn exhausted_failure(
|
||||||
|
selected: SelectedChannel,
|
||||||
|
status: Option<u16>,
|
||||||
|
raw_error_text: Option<String>,
|
||||||
|
attempts: usize,
|
||||||
|
on_terminal_failure: Option<
|
||||||
|
&(dyn Fn(SelectedChannel, AttemptFailure) -> BoxFuture<'static, ()> + Send + Sync),
|
||||||
|
>,
|
||||||
|
) -> ExecuteResult {
|
||||||
|
let info = FailureInfo {
|
||||||
|
action: AttemptFailureAction::Stop.as_str().to_string(),
|
||||||
|
status,
|
||||||
|
raw_error_text: raw_error_text.clone(),
|
||||||
|
};
|
||||||
|
terminal_failure(
|
||||||
|
selected,
|
||||||
|
FailureReason::Failed,
|
||||||
|
status,
|
||||||
|
raw_error_text,
|
||||||
|
attempts,
|
||||||
|
on_terminal_failure,
|
||||||
|
&info,
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
use futures::FutureExt;
|
||||||
|
use serde_json::json;
|
||||||
|
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||||
|
|
||||||
|
fn channel(id: i64) -> SelectedChannel {
|
||||||
|
SelectedChannel {
|
||||||
|
channel_id: id,
|
||||||
|
route_id: Some(10),
|
||||||
|
site_id: 20,
|
||||||
|
account_id: Some(30),
|
||||||
|
base_url: "https://example.test/v1".into(),
|
||||||
|
api_key: "key".into(),
|
||||||
|
model: "gpt-5.5".into(),
|
||||||
|
actual_model: "gpt-5.5".into(),
|
||||||
|
token_name: None,
|
||||||
|
token_value: None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn retries_same_channel_then_succeeds() {
|
||||||
|
let attempts_seen = Arc::new(AtomicUsize::new(0));
|
||||||
|
let recorded_failures = Arc::new(AtomicUsize::new(0));
|
||||||
|
let recorded_successes = Arc::new(AtomicUsize::new(0));
|
||||||
|
|
||||||
|
let conductor = ProxyConductor {
|
||||||
|
select_channel: Arc::new(|_, _| async { Some(channel(1)) }.boxed()),
|
||||||
|
select_next_channel: Arc::new(|_, _, _| async { Some(channel(2)) }.boxed()),
|
||||||
|
record_success: {
|
||||||
|
let recorded_successes = Arc::clone(&recorded_successes);
|
||||||
|
Arc::new(move |_, _| {
|
||||||
|
let recorded_successes = Arc::clone(&recorded_successes);
|
||||||
|
async move {
|
||||||
|
recorded_successes.fetch_add(1, Ordering::SeqCst);
|
||||||
|
}
|
||||||
|
.boxed()
|
||||||
|
})
|
||||||
|
},
|
||||||
|
record_failure: {
|
||||||
|
let recorded_failures = Arc::clone(&recorded_failures);
|
||||||
|
Arc::new(move |_, _| {
|
||||||
|
let recorded_failures = Arc::clone(&recorded_failures);
|
||||||
|
async move {
|
||||||
|
recorded_failures.fetch_add(1, Ordering::SeqCst);
|
||||||
|
}
|
||||||
|
.boxed()
|
||||||
|
})
|
||||||
|
},
|
||||||
|
refresh_auth: None,
|
||||||
|
};
|
||||||
|
|
||||||
|
let result = conductor
|
||||||
|
.execute(ExecuteInput {
|
||||||
|
requested_model: "gpt-5.5".into(),
|
||||||
|
downstream_policy: None,
|
||||||
|
attempt_fn: Box::new({
|
||||||
|
let attempts_seen = Arc::clone(&attempts_seen);
|
||||||
|
move |_| {
|
||||||
|
let attempt = attempts_seen.fetch_add(1, Ordering::SeqCst);
|
||||||
|
async move {
|
||||||
|
if attempt == 0 {
|
||||||
|
AttemptResult::Failure {
|
||||||
|
action: AttemptFailureAction::RetrySameChannel,
|
||||||
|
status: Some(429),
|
||||||
|
raw_error_text: Some("rate limited".into()),
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
AttemptResult::Success {
|
||||||
|
response: json!({"ok": true}),
|
||||||
|
latency_ms: Some(12),
|
||||||
|
cost: Some(0.1),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
.boxed()
|
||||||
|
}
|
||||||
|
}),
|
||||||
|
on_terminal_failure: None,
|
||||||
|
})
|
||||||
|
.await;
|
||||||
|
|
||||||
|
match result {
|
||||||
|
ExecuteResult::Success {
|
||||||
|
selected, attempts, ..
|
||||||
|
} => {
|
||||||
|
assert_eq!(selected.channel_id, 1);
|
||||||
|
assert_eq!(attempts, 2);
|
||||||
|
}
|
||||||
|
other => panic!("expected success, got {other:?}"),
|
||||||
|
}
|
||||||
|
assert_eq!(recorded_failures.load(Ordering::SeqCst), 1);
|
||||||
|
assert_eq!(recorded_successes.load(Ordering::SeqCst), 1);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn failover_selects_next_channel() {
|
||||||
|
let conductor = ProxyConductor {
|
||||||
|
select_channel: Arc::new(|_, _| async { Some(channel(1)) }.boxed()),
|
||||||
|
select_next_channel: Arc::new(|_, excluded, _| {
|
||||||
|
async move {
|
||||||
|
assert_eq!(excluded, vec![1]);
|
||||||
|
Some(channel(2))
|
||||||
|
}
|
||||||
|
.boxed()
|
||||||
|
}),
|
||||||
|
record_success: Arc::new(|_, _| async {}.boxed()),
|
||||||
|
record_failure: Arc::new(|_, _| async {}.boxed()),
|
||||||
|
refresh_auth: None,
|
||||||
|
};
|
||||||
|
|
||||||
|
let attempts_seen = Arc::new(AtomicUsize::new(0));
|
||||||
|
let result = conductor
|
||||||
|
.execute(ExecuteInput {
|
||||||
|
requested_model: "gpt-5.5".into(),
|
||||||
|
downstream_policy: None,
|
||||||
|
attempt_fn: Box::new({
|
||||||
|
let attempts_seen = Arc::clone(&attempts_seen);
|
||||||
|
move |ctx| {
|
||||||
|
let attempt = attempts_seen.fetch_add(1, Ordering::SeqCst);
|
||||||
|
async move {
|
||||||
|
if attempt == 0 {
|
||||||
|
assert_eq!(ctx.selected.channel_id, 1);
|
||||||
|
AttemptResult::Failure {
|
||||||
|
action: AttemptFailureAction::Failover,
|
||||||
|
status: Some(500),
|
||||||
|
raw_error_text: None,
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
assert_eq!(ctx.selected.channel_id, 2);
|
||||||
|
AttemptResult::Success {
|
||||||
|
response: json!({"ok": true}),
|
||||||
|
latency_ms: None,
|
||||||
|
cost: None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
.boxed()
|
||||||
|
}
|
||||||
|
}),
|
||||||
|
on_terminal_failure: None,
|
||||||
|
})
|
||||||
|
.await;
|
||||||
|
|
||||||
|
match result {
|
||||||
|
ExecuteResult::Success {
|
||||||
|
selected, attempts, ..
|
||||||
|
} => {
|
||||||
|
assert_eq!(selected.channel_id, 2);
|
||||||
|
assert_eq!(attempts, 2);
|
||||||
|
}
|
||||||
|
other => panic!("expected success, got {other:?}"),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,12 @@
|
|||||||
|
pub mod conductor;
|
||||||
|
pub mod retry_policy;
|
||||||
|
pub mod types;
|
||||||
|
|
||||||
|
pub use conductor::ProxyConductor;
|
||||||
|
pub use retry_policy::{
|
||||||
|
classify_failure, is_terminal, should_failover, should_refresh_auth, should_retry_same_channel,
|
||||||
|
};
|
||||||
|
pub use types::{
|
||||||
|
AttemptFailure, AttemptFailureAction, AttemptResult, ExecuteAttemptContext, ExecuteInput,
|
||||||
|
ExecuteResult, FailureInfo, FailureReason, SelectedChannel, SuccessMetrics,
|
||||||
|
};
|
||||||
@@ -0,0 +1,82 @@
|
|||||||
|
use super::types::AttemptFailureAction;
|
||||||
|
|
||||||
|
pub fn classify_failure(status: Option<u16>, _error_text: Option<&str>) -> AttemptFailureAction {
|
||||||
|
match status {
|
||||||
|
Some(429 | 503) => AttemptFailureAction::RetrySameChannel,
|
||||||
|
Some(401) => AttemptFailureAction::RefreshAuth,
|
||||||
|
Some(code) if (500..600).contains(&code) => AttemptFailureAction::Failover,
|
||||||
|
Some(code) if (400..500).contains(&code) => AttemptFailureAction::Terminal,
|
||||||
|
_ => AttemptFailureAction::Failover,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn is_terminal(action: AttemptFailureAction) -> bool {
|
||||||
|
matches!(action, AttemptFailureAction::Terminal)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn should_retry_same_channel(action: AttemptFailureAction) -> bool {
|
||||||
|
matches!(action, AttemptFailureAction::RetrySameChannel)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn should_refresh_auth(action: AttemptFailureAction) -> bool {
|
||||||
|
matches!(action, AttemptFailureAction::RefreshAuth)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn should_failover(action: AttemptFailureAction) -> bool {
|
||||||
|
matches!(action, AttemptFailureAction::Failover)
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn classifies_retryable_rate_limit_and_unavailable() {
|
||||||
|
assert_eq!(
|
||||||
|
classify_failure(Some(429), None),
|
||||||
|
AttemptFailureAction::RetrySameChannel
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
classify_failure(Some(503), Some("busy")),
|
||||||
|
AttemptFailureAction::RetrySameChannel
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn classifies_auth_server_and_terminal_failures() {
|
||||||
|
assert_eq!(
|
||||||
|
classify_failure(Some(401), None),
|
||||||
|
AttemptFailureAction::RefreshAuth
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
classify_failure(Some(500), None),
|
||||||
|
AttemptFailureAction::Failover
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
classify_failure(Some(502), None),
|
||||||
|
AttemptFailureAction::Failover
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
classify_failure(Some(400), None),
|
||||||
|
AttemptFailureAction::Terminal
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
classify_failure(Some(404), None),
|
||||||
|
AttemptFailureAction::Terminal
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
classify_failure(None, Some("timeout")),
|
||||||
|
AttemptFailureAction::Failover
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn helper_predicates_match_actions() {
|
||||||
|
assert!(is_terminal(AttemptFailureAction::Terminal));
|
||||||
|
assert!(should_retry_same_channel(
|
||||||
|
AttemptFailureAction::RetrySameChannel
|
||||||
|
));
|
||||||
|
assert!(should_refresh_auth(AttemptFailureAction::RefreshAuth));
|
||||||
|
assert!(should_failover(AttemptFailureAction::Failover));
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,124 @@
|
|||||||
|
use futures::future::BoxFuture;
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
use serde_json::Value;
|
||||||
|
use std::collections::HashMap;
|
||||||
|
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct SelectedChannel {
|
||||||
|
pub channel_id: i64,
|
||||||
|
pub route_id: Option<i64>,
|
||||||
|
pub site_id: i64,
|
||||||
|
pub account_id: Option<i64>,
|
||||||
|
pub base_url: String,
|
||||||
|
pub api_key: String,
|
||||||
|
pub model: String,
|
||||||
|
pub actual_model: String,
|
||||||
|
pub token_name: Option<String>,
|
||||||
|
pub token_value: Option<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||||
|
pub enum AttemptFailureAction {
|
||||||
|
RetrySameChannel,
|
||||||
|
RefreshAuth,
|
||||||
|
Failover,
|
||||||
|
Terminal,
|
||||||
|
Stop,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug)]
|
||||||
|
pub enum AttemptResult {
|
||||||
|
Success {
|
||||||
|
response: Value,
|
||||||
|
latency_ms: Option<u64>,
|
||||||
|
cost: Option<f64>,
|
||||||
|
},
|
||||||
|
Failure {
|
||||||
|
action: AttemptFailureAction,
|
||||||
|
status: Option<u16>,
|
||||||
|
raw_error_text: Option<String>,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct ExecuteInput {
|
||||||
|
pub requested_model: String,
|
||||||
|
pub downstream_policy: Option<Value>,
|
||||||
|
pub attempt_fn: Box<dyn Fn(ExecuteAttemptContext) -> BoxFuture<'static, AttemptResult> + Send>,
|
||||||
|
pub on_terminal_failure: Option<
|
||||||
|
Box<dyn Fn(SelectedChannel, AttemptFailure) -> BoxFuture<'static, ()> + Send + Sync>,
|
||||||
|
>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct ExecuteAttemptContext {
|
||||||
|
pub selected: SelectedChannel,
|
||||||
|
pub attempt_index: usize,
|
||||||
|
pub exclude_channel_ids: Vec<i64>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug)]
|
||||||
|
pub enum ExecuteResult {
|
||||||
|
Success {
|
||||||
|
selected: SelectedChannel,
|
||||||
|
response: Value,
|
||||||
|
attempts: usize,
|
||||||
|
},
|
||||||
|
Failure {
|
||||||
|
reason: FailureReason,
|
||||||
|
selected: Option<SelectedChannel>,
|
||||||
|
status: Option<u16>,
|
||||||
|
raw_error_text: Option<String>,
|
||||||
|
attempts: usize,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||||
|
pub enum FailureReason {
|
||||||
|
NoChannel,
|
||||||
|
Failed,
|
||||||
|
Terminal,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
|
||||||
|
pub struct SuccessMetrics {
|
||||||
|
pub latency_ms: Option<u64>,
|
||||||
|
pub cost: Option<f64>,
|
||||||
|
#[serde(default)]
|
||||||
|
pub metadata: HashMap<String, Value>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct FailureInfo {
|
||||||
|
pub action: String,
|
||||||
|
pub status: Option<u16>,
|
||||||
|
pub raw_error_text: Option<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct AttemptFailure {
|
||||||
|
pub status: Option<u16>,
|
||||||
|
pub raw_error_text: Option<String>,
|
||||||
|
pub action: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl From<&FailureInfo> for AttemptFailure {
|
||||||
|
fn from(value: &FailureInfo) -> Self {
|
||||||
|
Self {
|
||||||
|
status: value.status,
|
||||||
|
raw_error_text: value.raw_error_text.clone(),
|
||||||
|
action: value.action.clone(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl AttemptFailureAction {
|
||||||
|
pub fn as_str(self) -> &'static str {
|
||||||
|
match self {
|
||||||
|
Self::RetrySameChannel => "retry_same_channel",
|
||||||
|
Self::RefreshAuth => "refresh_auth",
|
||||||
|
Self::Failover => "failover",
|
||||||
|
Self::Terminal => "terminal",
|
||||||
|
Self::Stop => "stop",
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -5,7 +5,7 @@
|
|||||||
// 2. Sends the HTTP request
|
// 2. Sends the HTTP request
|
||||||
// 3. Transforms the response back to OpenAI-compatible format
|
// 3. Transforms the response back to OpenAI-compatible format
|
||||||
|
|
||||||
use crate::error::{ AppError, AppResult };
|
use crate::error::{AppError, AppResult};
|
||||||
use crate::AppState;
|
use crate::AppState;
|
||||||
use reqwest::Client;
|
use reqwest::Client;
|
||||||
use serde_json::Value;
|
use serde_json::Value;
|
||||||
@@ -110,9 +110,7 @@ pub async fn execute_chat_completions_stream(
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
Err(e) => {
|
Err(e) => {
|
||||||
return Err(AppError::UpstreamError(format!(
|
return Err(AppError::UpstreamError(format!("Stream error: {e}")));
|
||||||
"Stream error: {e}"
|
|
||||||
)));
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -168,7 +166,10 @@ pub async fn execute_image_generation(
|
|||||||
body: Value,
|
body: Value,
|
||||||
) -> AppResult<Value> {
|
) -> AppResult<Value> {
|
||||||
let client = http_client();
|
let client = http_client();
|
||||||
let upstream_url = format!("{}/images/generations", route.base_url.trim_end_matches('/'));
|
let upstream_url = format!(
|
||||||
|
"{}/images/generations",
|
||||||
|
route.base_url.trim_end_matches('/')
|
||||||
|
);
|
||||||
|
|
||||||
let resp = client
|
let resp = client
|
||||||
.post(&upstream_url)
|
.post(&upstream_url)
|
||||||
|
|||||||
@@ -5,7 +5,11 @@
|
|||||||
// /v1/images/generations → ...same flow...
|
// /v1/images/generations → ...same flow...
|
||||||
// /v1/audio/speech → ...same flow...
|
// /v1/audio/speech → ...same flow...
|
||||||
|
|
||||||
|
pub mod capabilities;
|
||||||
|
pub mod cli_profiles;
|
||||||
|
pub mod conductor;
|
||||||
pub mod executors;
|
pub mod executors;
|
||||||
pub mod providers;
|
pub mod providers;
|
||||||
pub mod runtime;
|
pub mod runtime;
|
||||||
pub mod surfaces;
|
pub mod surfaces;
|
||||||
|
pub mod token_router;
|
||||||
|
|||||||
@@ -1 +1,6 @@
|
|||||||
// Proxy providers
|
// Provider profiles & registry
|
||||||
|
pub mod registry;
|
||||||
|
pub mod types;
|
||||||
|
|
||||||
|
pub use registry::ProviderRegistry;
|
||||||
|
pub use types::{ProviderCategory, ProviderProfile, RouteResolution};
|
||||||
|
|||||||
@@ -0,0 +1,134 @@
|
|||||||
|
// Provider registry — holds all configured upstream provider profiles
|
||||||
|
use crate::proxy::providers::types::{ProviderCategory, ProviderProfile, RouteResolution};
|
||||||
|
use std::collections::HashMap;
|
||||||
|
|
||||||
|
/// Registry of all available provider profiles
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct ProviderRegistry {
|
||||||
|
providers: HashMap<String, ProviderProfile>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ProviderRegistry {
|
||||||
|
/// Create a new registry with default providers
|
||||||
|
pub fn new() -> Self {
|
||||||
|
let mut registry = Self {
|
||||||
|
providers: HashMap::new(),
|
||||||
|
};
|
||||||
|
registry.register_defaults();
|
||||||
|
registry
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Register the built-in default providers
|
||||||
|
fn register_defaults(&mut self) {
|
||||||
|
// OpenAI
|
||||||
|
self.register(ProviderProfile::openai_compatible(
|
||||||
|
"openai",
|
||||||
|
"OpenAI",
|
||||||
|
"https://api.openai.com/v1",
|
||||||
|
));
|
||||||
|
// Anthropic
|
||||||
|
self.register(ProviderProfile::anthropic(
|
||||||
|
"anthropic",
|
||||||
|
"Anthropic",
|
||||||
|
"https://api.anthropic.com/v1",
|
||||||
|
));
|
||||||
|
// Google Gemini
|
||||||
|
self.register(ProviderProfile::gemini(
|
||||||
|
"gemini",
|
||||||
|
"Google Gemini",
|
||||||
|
"https://generativelanguage.googleapis.com/v1beta",
|
||||||
|
));
|
||||||
|
// Codex
|
||||||
|
self.register(ProviderProfile {
|
||||||
|
id: "codex".into(),
|
||||||
|
name: "OpenAI Codex".into(),
|
||||||
|
category: ProviderCategory::Codex,
|
||||||
|
base_url: "https://api.openai.com/v1".into(),
|
||||||
|
default_model: Some("codex".into()),
|
||||||
|
requires_api_key: true,
|
||||||
|
supports_streaming: true,
|
||||||
|
headers: HashMap::new(),
|
||||||
|
});
|
||||||
|
// Claude Code
|
||||||
|
self.register(ProviderProfile {
|
||||||
|
id: "claude".into(),
|
||||||
|
name: "Claude Code".into(),
|
||||||
|
category: ProviderCategory::Claude,
|
||||||
|
base_url: "https://api.anthropic.com/v1".into(),
|
||||||
|
default_model: Some("claude-sonnet-4".into()),
|
||||||
|
requires_api_key: true,
|
||||||
|
supports_streaming: true,
|
||||||
|
headers: HashMap::new(),
|
||||||
|
});
|
||||||
|
// Gemini CLI
|
||||||
|
self.register(ProviderProfile {
|
||||||
|
id: "gemini_cli".into(),
|
||||||
|
name: "Gemini CLI".into(),
|
||||||
|
category: ProviderCategory::GeminiCli,
|
||||||
|
base_url: "https://generativelanguage.googleapis.com/v1beta".into(),
|
||||||
|
default_model: Some("gemini-2.0-flash".into()),
|
||||||
|
requires_api_key: true,
|
||||||
|
supports_streaming: true,
|
||||||
|
headers: HashMap::new(),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Register a new provider profile
|
||||||
|
pub fn register(&mut self, profile: ProviderProfile) {
|
||||||
|
self.providers.insert(profile.id.clone(), profile);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Get a provider by ID
|
||||||
|
pub fn get(&self, id: &str) -> Option<&ProviderProfile> {
|
||||||
|
self.providers.get(id)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// List all registered providers
|
||||||
|
pub fn list(&self) -> Vec<&ProviderProfile> {
|
||||||
|
self.providers.values().collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Remove a provider by ID
|
||||||
|
pub fn remove(&mut self, id: &str) -> Option<ProviderProfile> {
|
||||||
|
self.providers.remove(id)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Resolve a model name to an upstream route.
|
||||||
|
/// Tries to match the model to a provider category, falls back to first OpenAI-compatible provider.
|
||||||
|
pub fn resolve(&self, model: &str, api_key: &str) -> Option<RouteResolution> {
|
||||||
|
// First try to find a provider that supports this model
|
||||||
|
for provider in self.providers.values() {
|
||||||
|
if provider.supports_model(model) {
|
||||||
|
let base_url = provider.base_url.clone();
|
||||||
|
return Some(RouteResolution {
|
||||||
|
provider: provider.clone(),
|
||||||
|
api_key: api_key.to_string(),
|
||||||
|
model: model.to_string(),
|
||||||
|
base_url,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Fall back to the first OpenAI-compatible provider
|
||||||
|
self.providers
|
||||||
|
.values()
|
||||||
|
.find(|p| {
|
||||||
|
matches!(
|
||||||
|
p.category,
|
||||||
|
ProviderCategory::OpenAi | ProviderCategory::GenericOpenAi
|
||||||
|
)
|
||||||
|
})
|
||||||
|
.map(|provider| RouteResolution {
|
||||||
|
provider: provider.clone(),
|
||||||
|
api_key: api_key.to_string(),
|
||||||
|
model: model.to_string(),
|
||||||
|
base_url: provider.base_url.clone(),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for ProviderRegistry {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self::new()
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,137 @@
|
|||||||
|
// Provider profile types — upstream API provider definitions
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
use std::collections::HashMap;
|
||||||
|
|
||||||
|
/// Provider category — determines request transformation and routing behavior
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||||
|
pub enum ProviderCategory {
|
||||||
|
/// OpenAI-compatible chat/completions API
|
||||||
|
OpenAi,
|
||||||
|
/// Anthropic Messages API
|
||||||
|
Anthropic,
|
||||||
|
/// Google Gemini API
|
||||||
|
Gemini,
|
||||||
|
/// Codex CLI / Responses API
|
||||||
|
Codex,
|
||||||
|
/// Claude Code CLI
|
||||||
|
Claude,
|
||||||
|
/// Gemini CLI
|
||||||
|
GeminiCli,
|
||||||
|
/// Antigravity / assistant-style API
|
||||||
|
Antigravity,
|
||||||
|
/// Generic OpenAI-compatible proxy (new-api, one-api, etc.)
|
||||||
|
GenericOpenAi,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl std::fmt::Display for ProviderCategory {
|
||||||
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||||
|
match self {
|
||||||
|
Self::OpenAi => write!(f, "openai"),
|
||||||
|
Self::Anthropic => write!(f, "anthropic"),
|
||||||
|
Self::Gemini => write!(f, "gemini"),
|
||||||
|
Self::Codex => write!(f, "codex"),
|
||||||
|
Self::Claude => write!(f, "claude"),
|
||||||
|
Self::GeminiCli => write!(f, "gemini_cli"),
|
||||||
|
Self::Antigravity => write!(f, "antigravity"),
|
||||||
|
Self::GenericOpenAi => write!(f, "generic_openai"),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Provider configuration profile
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct ProviderProfile {
|
||||||
|
pub id: String,
|
||||||
|
pub name: String,
|
||||||
|
pub category: ProviderCategory,
|
||||||
|
pub base_url: String,
|
||||||
|
pub default_model: Option<String>,
|
||||||
|
pub requires_api_key: bool,
|
||||||
|
pub supports_streaming: bool,
|
||||||
|
pub headers: HashMap<String, String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ProviderProfile {
|
||||||
|
/// Create a standard OpenAI-compatible provider
|
||||||
|
pub fn openai_compatible(id: &str, name: &str, base_url: &str) -> Self {
|
||||||
|
let mut headers = HashMap::new();
|
||||||
|
headers.insert("Content-Type".into(), "application/json".into());
|
||||||
|
Self {
|
||||||
|
id: id.to_string(),
|
||||||
|
name: name.to_string(),
|
||||||
|
category: ProviderCategory::GenericOpenAi,
|
||||||
|
base_url: base_url.trim_end_matches('/').to_string(),
|
||||||
|
default_model: None,
|
||||||
|
requires_api_key: true,
|
||||||
|
supports_streaming: true,
|
||||||
|
headers,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Create an Anthropic Messages API provider
|
||||||
|
pub fn anthropic(id: &str, name: &str, base_url: &str) -> Self {
|
||||||
|
let mut headers = HashMap::new();
|
||||||
|
headers.insert("Content-Type".into(), "application/json".into());
|
||||||
|
headers.insert("anthropic-version".into(), "2023-06-01".into());
|
||||||
|
Self {
|
||||||
|
id: id.to_string(),
|
||||||
|
name: name.to_string(),
|
||||||
|
category: ProviderCategory::Anthropic,
|
||||||
|
base_url: base_url.trim_end_matches('/').to_string(),
|
||||||
|
default_model: None,
|
||||||
|
requires_api_key: true,
|
||||||
|
supports_streaming: true,
|
||||||
|
headers,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Create a Gemini API provider
|
||||||
|
pub fn gemini(id: &str, name: &str, base_url: &str) -> Self {
|
||||||
|
let mut headers = HashMap::new();
|
||||||
|
headers.insert("Content-Type".into(), "application/json".into());
|
||||||
|
Self {
|
||||||
|
id: id.to_string(),
|
||||||
|
name: name.to_string(),
|
||||||
|
category: ProviderCategory::Gemini,
|
||||||
|
base_url: base_url.trim_end_matches('/').to_string(),
|
||||||
|
default_model: None,
|
||||||
|
requires_api_key: true,
|
||||||
|
supports_streaming: true,
|
||||||
|
headers,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Check if this provider handles the given model
|
||||||
|
pub fn supports_model(&self, model: &str) -> bool {
|
||||||
|
let model_lower = model.to_lowercase();
|
||||||
|
match self.category {
|
||||||
|
ProviderCategory::OpenAi | ProviderCategory::GenericOpenAi => {
|
||||||
|
model_lower.contains("gpt")
|
||||||
|
|| model_lower.contains("o1")
|
||||||
|
|| model_lower.contains("o3")
|
||||||
|
|| model_lower.contains("dall-e")
|
||||||
|
}
|
||||||
|
ProviderCategory::Anthropic => model_lower.contains("claude"),
|
||||||
|
ProviderCategory::Gemini => model_lower.contains("gemini"),
|
||||||
|
ProviderCategory::Codex => model_lower.contains("codex"),
|
||||||
|
ProviderCategory::Claude => {
|
||||||
|
model_lower.contains("claude") && model_lower.contains("code")
|
||||||
|
}
|
||||||
|
ProviderCategory::GeminiCli => {
|
||||||
|
model_lower.contains("gemini") && model_lower.contains("cli")
|
||||||
|
}
|
||||||
|
ProviderCategory::Antigravity => {
|
||||||
|
model_lower.contains("antigravity") || model_lower.contains("assistant")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Result of resolving a model to an upstream route
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct RouteResolution {
|
||||||
|
pub provider: ProviderProfile,
|
||||||
|
pub api_key: String,
|
||||||
|
pub model: String,
|
||||||
|
pub base_url: String,
|
||||||
|
}
|
||||||
@@ -1 +1,53 @@
|
|||||||
// Proxy runtime
|
use crate::proxy::conductor::{FailureInfo, ProxyConductor, SelectedChannel, SuccessMetrics};
|
||||||
|
use crate::proxy::token_router::TokenRouter;
|
||||||
|
use futures::FutureExt;
|
||||||
|
use sqlx::SqlitePool;
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct ProxyRuntime {
|
||||||
|
pub router: TokenRouter,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ProxyRuntime {
|
||||||
|
pub fn new(db: Arc<SqlitePool>) -> Self {
|
||||||
|
Self {
|
||||||
|
router: TokenRouter::new(db),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn conductor(&self) -> ProxyConductor {
|
||||||
|
let router = self.router.clone();
|
||||||
|
let next_router = self.router.clone();
|
||||||
|
let success_router = self.router.clone();
|
||||||
|
let failure_router = self.router.clone();
|
||||||
|
|
||||||
|
ProxyConductor {
|
||||||
|
select_channel: Arc::new(move |model, _policy| {
|
||||||
|
let router = router.clone();
|
||||||
|
async move { router.resolve_selected(&model).await.ok() }.boxed()
|
||||||
|
}),
|
||||||
|
select_next_channel: Arc::new(move |model, excluded, _policy| {
|
||||||
|
let router = next_router.clone();
|
||||||
|
async move { router.resolve_next_selected(&model, &excluded).await.ok() }.boxed()
|
||||||
|
}),
|
||||||
|
record_success: Arc::new(move |channel_id: i64, metrics: SuccessMetrics| {
|
||||||
|
let router = success_router.clone();
|
||||||
|
async move {
|
||||||
|
let _ = router.record_success(channel_id, metrics).await;
|
||||||
|
}
|
||||||
|
.boxed()
|
||||||
|
}),
|
||||||
|
record_failure: Arc::new(move |channel_id: i64, info: FailureInfo| {
|
||||||
|
let router = failure_router.clone();
|
||||||
|
async move {
|
||||||
|
let _ = router.record_failure(channel_id, info).await;
|
||||||
|
}
|
||||||
|
.boxed()
|
||||||
|
}),
|
||||||
|
refresh_auth: Some(Arc::new(|selected: SelectedChannel, _info: FailureInfo| {
|
||||||
|
async move { Some(selected) }.boxed()
|
||||||
|
})),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -1 +1,32 @@
|
|||||||
// Proxy surfaces
|
use std::collections::HashMap;
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
use parking_lot::Mutex;
|
||||||
|
use serde_json::Value;
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Default)]
|
||||||
|
pub struct ProxySurfaceState {
|
||||||
|
sessions: Arc<Mutex<HashMap<String, SurfaceSession>>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct SurfaceSession {
|
||||||
|
pub id: String,
|
||||||
|
pub model: String,
|
||||||
|
pub channel_id: Option<i64>,
|
||||||
|
pub metadata: Value,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ProxySurfaceState {
|
||||||
|
pub fn upsert_session(&self, session: SurfaceSession) {
|
||||||
|
self.sessions.lock().insert(session.id.clone(), session);
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn get_session(&self, id: &str) -> Option<SurfaceSession> {
|
||||||
|
self.sessions.lock().get(id).cloned()
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn remove_session(&self, id: &str) -> Option<SurfaceSession> {
|
||||||
|
self.sessions.lock().remove(id)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -0,0 +1,5 @@
|
|||||||
|
pub mod router;
|
||||||
|
pub mod types;
|
||||||
|
|
||||||
|
pub use router::TokenRouter;
|
||||||
|
pub use types::{RouteDecision, RouteMatch};
|
||||||
@@ -0,0 +1,430 @@
|
|||||||
|
use super::types::{RouteDecision, RouteMatch};
|
||||||
|
use crate::error::{AppError, AppResult};
|
||||||
|
use crate::proxy::conductor::{FailureInfo, SelectedChannel, SuccessMetrics};
|
||||||
|
use crate::proxy::executors::UpstreamRoute;
|
||||||
|
use regex::Regex;
|
||||||
|
use sqlx::{FromRow, SqlitePool};
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct TokenRouter {
|
||||||
|
db: Arc<SqlitePool>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, FromRow)]
|
||||||
|
struct RouteRow {
|
||||||
|
id: i64,
|
||||||
|
model_pattern: String,
|
||||||
|
model_mapping: Option<String>,
|
||||||
|
route_mode: Option<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, FromRow)]
|
||||||
|
struct ChannelRow {
|
||||||
|
channel_id: i64,
|
||||||
|
route_id: i64,
|
||||||
|
account_id: i64,
|
||||||
|
site_id: i64,
|
||||||
|
base_url: String,
|
||||||
|
site_api_key: Option<String>,
|
||||||
|
account_access_token: String,
|
||||||
|
account_api_token: Option<String>,
|
||||||
|
token_name: Option<String>,
|
||||||
|
token_value: Option<String>,
|
||||||
|
source_model: Option<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl TokenRouter {
|
||||||
|
pub fn new(db: Arc<SqlitePool>) -> Self {
|
||||||
|
Self { db }
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn resolve(&self, model: &str) -> AppResult<UpstreamRoute> {
|
||||||
|
let decision = self.resolve_decision(model, &[]).await?;
|
||||||
|
Ok(UpstreamRoute {
|
||||||
|
base_url: decision.selected.base_url,
|
||||||
|
api_key: decision.selected.api_key,
|
||||||
|
model: decision.selected.actual_model,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn resolve_selected(&self, model: &str) -> AppResult<SelectedChannel> {
|
||||||
|
Ok(self.resolve_decision(model, &[]).await?.selected)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn resolve_next_selected(
|
||||||
|
&self,
|
||||||
|
model: &str,
|
||||||
|
exclude_channel_ids: &[i64],
|
||||||
|
) -> AppResult<SelectedChannel> {
|
||||||
|
Ok(self
|
||||||
|
.resolve_decision(model, exclude_channel_ids)
|
||||||
|
.await?
|
||||||
|
.selected)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn resolve_decision(
|
||||||
|
&self,
|
||||||
|
model: &str,
|
||||||
|
exclude_channel_ids: &[i64],
|
||||||
|
) -> AppResult<RouteDecision> {
|
||||||
|
let route = self.match_route(model).await?.ok_or_else(|| {
|
||||||
|
AppError::BadRequest(format!("No route configured for model '{model}'"))
|
||||||
|
})?;
|
||||||
|
let channel = self
|
||||||
|
.select_channel_for_route(&route.0, model, exclude_channel_ids)
|
||||||
|
.await?
|
||||||
|
.ok_or_else(|| {
|
||||||
|
AppError::BadRequest(format!(
|
||||||
|
"No enabled channel configured for model '{}'",
|
||||||
|
model
|
||||||
|
))
|
||||||
|
})?;
|
||||||
|
|
||||||
|
Ok(RouteDecision {
|
||||||
|
requested_model: model.to_string(),
|
||||||
|
matched_pattern: route.0.model_pattern,
|
||||||
|
match_kind: route.1,
|
||||||
|
selected: channel,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn match_route(&self, model: &str) -> AppResult<Option<(RouteRow, RouteMatch)>> {
|
||||||
|
let rows: Vec<RouteRow> = sqlx::query_as(
|
||||||
|
"SELECT id, model_pattern, model_mapping, route_mode
|
||||||
|
FROM token_routes
|
||||||
|
WHERE enabled = 1
|
||||||
|
ORDER BY id",
|
||||||
|
)
|
||||||
|
.fetch_all(&*self.db)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
let mut best: Option<(RouteRow, RouteMatch, usize, usize)> = None;
|
||||||
|
for row in rows {
|
||||||
|
if let Some(kind) = match_pattern(model, &row.model_pattern, row.route_mode.as_deref())
|
||||||
|
{
|
||||||
|
let rank = match kind {
|
||||||
|
RouteMatch::Exact => 3,
|
||||||
|
RouteMatch::Wildcard => 2,
|
||||||
|
RouteMatch::Regex => 1,
|
||||||
|
};
|
||||||
|
let specificity = row.model_pattern.chars().filter(|c| *c != '*').count();
|
||||||
|
let replace = best
|
||||||
|
.as_ref()
|
||||||
|
.map(|(_, _, best_rank, best_specificity)| {
|
||||||
|
rank > *best_rank || (rank == *best_rank && specificity > *best_specificity)
|
||||||
|
})
|
||||||
|
.unwrap_or(true);
|
||||||
|
if replace {
|
||||||
|
best = Some((row, kind, rank, specificity));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(best.map(|(row, kind, _, _)| (row, kind)))
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn select_channel_for_route(
|
||||||
|
&self,
|
||||||
|
route: &RouteRow,
|
||||||
|
requested_model: &str,
|
||||||
|
exclude_channel_ids: &[i64],
|
||||||
|
) -> AppResult<Option<SelectedChannel>> {
|
||||||
|
let rows: Vec<ChannelRow> = sqlx::query_as(
|
||||||
|
"SELECT
|
||||||
|
rc.id AS channel_id,
|
||||||
|
rc.route_id AS route_id,
|
||||||
|
rc.account_id AS account_id,
|
||||||
|
a.site_id AS site_id,
|
||||||
|
s.url AS base_url,
|
||||||
|
s.api_key AS site_api_key,
|
||||||
|
a.access_token AS account_access_token,
|
||||||
|
a.api_token AS account_api_token,
|
||||||
|
at.name AS token_name,
|
||||||
|
at.token AS token_value,
|
||||||
|
rc.source_model AS source_model
|
||||||
|
FROM route_channels rc
|
||||||
|
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 rc.route_id = ?
|
||||||
|
AND rc.enabled = 1
|
||||||
|
AND a.status = 'active'
|
||||||
|
AND s.status = 'active'
|
||||||
|
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
|
||||||
|
COALESCE(rc.priority, 0) DESC,
|
||||||
|
rc.consecutive_fail_count ASC,
|
||||||
|
COALESCE(rc.weight, 10) DESC,
|
||||||
|
rc.last_selected_at IS NOT NULL,
|
||||||
|
rc.id ASC",
|
||||||
|
)
|
||||||
|
.bind(route.id)
|
||||||
|
.fetch_all(&*self.db)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
let row = rows
|
||||||
|
.into_iter()
|
||||||
|
.find(|row| !exclude_channel_ids.contains(&row.channel_id));
|
||||||
|
|
||||||
|
Ok(row.map(|row| {
|
||||||
|
let actual_model = route
|
||||||
|
.model_mapping
|
||||||
|
.clone()
|
||||||
|
.filter(|value| !value.trim().is_empty())
|
||||||
|
.or_else(|| {
|
||||||
|
row.source_model
|
||||||
|
.clone()
|
||||||
|
.filter(|value| !value.trim().is_empty())
|
||||||
|
})
|
||||||
|
.unwrap_or_else(|| requested_model.to_string());
|
||||||
|
let api_key = row
|
||||||
|
.token_value
|
||||||
|
.clone()
|
||||||
|
.or(row.account_api_token.clone())
|
||||||
|
.unwrap_or(row.account_access_token.clone());
|
||||||
|
|
||||||
|
SelectedChannel {
|
||||||
|
channel_id: row.channel_id,
|
||||||
|
route_id: Some(row.route_id),
|
||||||
|
site_id: row.site_id,
|
||||||
|
account_id: Some(row.account_id),
|
||||||
|
base_url: row.base_url,
|
||||||
|
api_key: api_key
|
||||||
|
.trim()
|
||||||
|
.to_string()
|
||||||
|
.if_empty_then(row.site_api_key.unwrap_or_default()),
|
||||||
|
model: requested_model.to_string(),
|
||||||
|
actual_model,
|
||||||
|
token_name: row.token_name,
|
||||||
|
token_value: row.token_value,
|
||||||
|
}
|
||||||
|
}))
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn record_success(&self, channel_id: i64, metrics: SuccessMetrics) -> AppResult<()> {
|
||||||
|
sqlx::query(
|
||||||
|
"UPDATE route_channels
|
||||||
|
SET success_count = COALESCE(success_count, 0) + 1,
|
||||||
|
total_latency_ms = COALESCE(total_latency_ms, 0) + ?,
|
||||||
|
total_cost = COALESCE(total_cost, 0) + ?,
|
||||||
|
consecutive_fail_count = 0,
|
||||||
|
cooldown_level = 0,
|
||||||
|
cooldown_until = NULL,
|
||||||
|
last_used_at = datetime('now'),
|
||||||
|
last_selected_at = datetime('now')
|
||||||
|
WHERE id = ?",
|
||||||
|
)
|
||||||
|
.bind(metrics.latency_ms.unwrap_or(0) as i64)
|
||||||
|
.bind(metrics.cost.unwrap_or(0.0))
|
||||||
|
.bind(channel_id)
|
||||||
|
.execute(&*self.db)
|
||||||
|
.await?;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn record_failure(&self, channel_id: i64, _info: FailureInfo) -> AppResult<()> {
|
||||||
|
sqlx::query(
|
||||||
|
"UPDATE route_channels
|
||||||
|
SET fail_count = COALESCE(fail_count, 0) + 1,
|
||||||
|
consecutive_fail_count = consecutive_fail_count + 1,
|
||||||
|
cooldown_level = cooldown_level + 1,
|
||||||
|
last_fail_at = datetime('now')
|
||||||
|
WHERE id = ?",
|
||||||
|
)
|
||||||
|
.bind(channel_id)
|
||||||
|
.execute(&*self.db)
|
||||||
|
.await?;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
trait IfEmptyThen {
|
||||||
|
fn if_empty_then(self, fallback: String) -> String;
|
||||||
|
}
|
||||||
|
|
||||||
|
impl IfEmptyThen for String {
|
||||||
|
fn if_empty_then(self, fallback: String) -> String {
|
||||||
|
if self.is_empty() {
|
||||||
|
fallback
|
||||||
|
} else {
|
||||||
|
self
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn match_pattern(model: &str, pattern: &str, route_mode: Option<&str>) -> Option<RouteMatch> {
|
||||||
|
match route_mode.unwrap_or("pattern") {
|
||||||
|
"exact" => (model == pattern).then_some(RouteMatch::Exact),
|
||||||
|
"regex" => Regex::new(pattern)
|
||||||
|
.ok()
|
||||||
|
.and_then(|re| re.is_match(model).then_some(RouteMatch::Regex)),
|
||||||
|
_ if pattern == model => Some(RouteMatch::Exact),
|
||||||
|
_ if pattern.starts_with("regex:") => Regex::new(pattern.trim_start_matches("regex:"))
|
||||||
|
.ok()
|
||||||
|
.and_then(|re| re.is_match(model).then_some(RouteMatch::Regex)),
|
||||||
|
_ if pattern.contains('*') => {
|
||||||
|
wildcard_matches(model, pattern).then_some(RouteMatch::Wildcard)
|
||||||
|
}
|
||||||
|
_ => None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn wildcard_matches(model: &str, pattern: &str) -> bool {
|
||||||
|
let escaped = pattern
|
||||||
|
.split('*')
|
||||||
|
.map(regex::escape)
|
||||||
|
.collect::<Vec<_>>()
|
||||||
|
.join(".*");
|
||||||
|
Regex::new(&format!("^{escaped}$"))
|
||||||
|
.map(|re| re.is_match(model))
|
||||||
|
.unwrap_or(false)
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
use sqlx::sqlite::SqlitePoolOptions;
|
||||||
|
|
||||||
|
async fn test_pool() -> SqlitePool {
|
||||||
|
let pool = SqlitePoolOptions::new()
|
||||||
|
.max_connections(1)
|
||||||
|
.connect("sqlite::memory:")
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
sqlx::query(
|
||||||
|
"CREATE TABLE sites (
|
||||||
|
id INTEGER PRIMARY KEY,
|
||||||
|
name TEXT NOT NULL,
|
||||||
|
url TEXT NOT NULL,
|
||||||
|
platform TEXT NOT NULL DEFAULT 'openai',
|
||||||
|
status TEXT NOT NULL DEFAULT 'active',
|
||||||
|
api_key TEXT
|
||||||
|
)",
|
||||||
|
)
|
||||||
|
.execute(&pool)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
sqlx::query(
|
||||||
|
"CREATE TABLE token_routes (
|
||||||
|
id INTEGER PRIMARY KEY,
|
||||||
|
model_pattern TEXT NOT NULL,
|
||||||
|
model_mapping TEXT,
|
||||||
|
route_mode TEXT,
|
||||||
|
enabled INTEGER DEFAULT 1
|
||||||
|
)",
|
||||||
|
)
|
||||||
|
.execute(&pool)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
sqlx::query(
|
||||||
|
"CREATE TABLE accounts (
|
||||||
|
id INTEGER PRIMARY KEY,
|
||||||
|
site_id INTEGER NOT NULL,
|
||||||
|
access_token TEXT NOT NULL,
|
||||||
|
api_token TEXT,
|
||||||
|
status TEXT NOT NULL DEFAULT 'active'
|
||||||
|
)",
|
||||||
|
)
|
||||||
|
.execute(&pool)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
sqlx::query(
|
||||||
|
"CREATE TABLE account_tokens (
|
||||||
|
id INTEGER PRIMARY KEY,
|
||||||
|
account_id INTEGER NOT NULL,
|
||||||
|
name TEXT NOT NULL,
|
||||||
|
token TEXT NOT NULL,
|
||||||
|
enabled INTEGER DEFAULT 1
|
||||||
|
)",
|
||||||
|
)
|
||||||
|
.execute(&pool)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
sqlx::query(
|
||||||
|
"CREATE TABLE route_channels (
|
||||||
|
id INTEGER PRIMARY KEY,
|
||||||
|
route_id INTEGER NOT NULL,
|
||||||
|
account_id INTEGER NOT NULL,
|
||||||
|
token_id INTEGER,
|
||||||
|
source_model TEXT,
|
||||||
|
priority INTEGER DEFAULT 0,
|
||||||
|
weight INTEGER DEFAULT 10,
|
||||||
|
enabled INTEGER DEFAULT 1,
|
||||||
|
consecutive_fail_count INTEGER NOT NULL DEFAULT 0,
|
||||||
|
cooldown_level INTEGER NOT NULL DEFAULT 0,
|
||||||
|
cooldown_until TEXT,
|
||||||
|
last_selected_at TEXT,
|
||||||
|
success_count INTEGER DEFAULT 0,
|
||||||
|
fail_count INTEGER DEFAULT 0,
|
||||||
|
total_latency_ms INTEGER DEFAULT 0,
|
||||||
|
total_cost REAL DEFAULT 0,
|
||||||
|
last_used_at TEXT,
|
||||||
|
last_fail_at TEXT
|
||||||
|
)",
|
||||||
|
)
|
||||||
|
.execute(&pool)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
sqlx::query("INSERT INTO sites (id, name, url, api_key) VALUES (1, 'OpenAI', 'https://api.example/v1', 'site-key')")
|
||||||
|
.execute(&pool)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
sqlx::query("INSERT INTO accounts (id, site_id, access_token, api_token) VALUES (1, 1, 'access-key', NULL)")
|
||||||
|
.execute(&pool)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
sqlx::query("INSERT INTO account_tokens (id, account_id, name, token) VALUES (1, 1, 'primary', 'token-key')")
|
||||||
|
.execute(&pool)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
pool
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn resolves_exact_model_route() {
|
||||||
|
let pool = test_pool().await;
|
||||||
|
sqlx::query("INSERT INTO token_routes (id, model_pattern, model_mapping, route_mode) VALUES (1, 'gpt-5.5', 'upstream-gpt-5.5', 'exact')")
|
||||||
|
.execute(&pool)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
sqlx::query(
|
||||||
|
"INSERT INTO route_channels (id, route_id, account_id, token_id) VALUES (1, 1, 1, 1)",
|
||||||
|
)
|
||||||
|
.execute(&pool)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
let router = TokenRouter::new(Arc::new(pool));
|
||||||
|
let decision = router.resolve_decision("gpt-5.5", &[]).await.unwrap();
|
||||||
|
assert_eq!(decision.match_kind, RouteMatch::Exact);
|
||||||
|
assert_eq!(decision.selected.actual_model, "upstream-gpt-5.5");
|
||||||
|
assert_eq!(decision.selected.api_key, "token-key");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn resolves_wildcard_route_and_skips_excluded_channel() {
|
||||||
|
let pool = test_pool().await;
|
||||||
|
sqlx::query("INSERT INTO token_routes (id, model_pattern, route_mode) VALUES (1, 'gpt-*', 'pattern')")
|
||||||
|
.execute(&pool)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
sqlx::query("INSERT INTO route_channels (id, route_id, account_id, token_id, priority) VALUES (1, 1, 1, 1, 10)")
|
||||||
|
.execute(&pool)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
sqlx::query("INSERT INTO route_channels (id, route_id, account_id, token_id, priority) VALUES (2, 1, 1, 1, 1)")
|
||||||
|
.execute(&pool)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
let router = TokenRouter::new(Arc::new(pool));
|
||||||
|
let decision = router.resolve_decision("gpt-4o", &[1]).await.unwrap();
|
||||||
|
assert_eq!(decision.match_kind, RouteMatch::Wildcard);
|
||||||
|
assert_eq!(decision.selected.channel_id, 2);
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,16 @@
|
|||||||
|
use crate::proxy::conductor::SelectedChannel;
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||||
|
pub enum RouteMatch {
|
||||||
|
Exact,
|
||||||
|
Wildcard,
|
||||||
|
Regex,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct RouteDecision {
|
||||||
|
pub requested_model: String,
|
||||||
|
pub matched_pattern: String,
|
||||||
|
pub match_kind: RouteMatch,
|
||||||
|
pub selected: SelectedChannel,
|
||||||
|
}
|
||||||
@@ -0,0 +1,299 @@
|
|||||||
|
// Account token routes — CRUD for account-level API tokens
|
||||||
|
|
||||||
|
use crate::{
|
||||||
|
db::models::account::AccountToken,
|
||||||
|
error::{AppError, AppResult},
|
||||||
|
AppState,
|
||||||
|
};
|
||||||
|
use axum::{
|
||||||
|
extract::{Path, Query, State},
|
||||||
|
http::StatusCode,
|
||||||
|
routing::{get, post},
|
||||||
|
Json, Router,
|
||||||
|
};
|
||||||
|
use chrono::Utc;
|
||||||
|
use serde::Deserialize;
|
||||||
|
use serde_json::{json, Value};
|
||||||
|
use tracing::info;
|
||||||
|
use uuid::Uuid;
|
||||||
|
|
||||||
|
pub fn routes() -> Router<AppState> {
|
||||||
|
Router::new()
|
||||||
|
.route(
|
||||||
|
"/api/account-tokens",
|
||||||
|
get(list_account_tokens).post(create_account_token),
|
||||||
|
)
|
||||||
|
.route(
|
||||||
|
"/api/account-tokens/sync-all",
|
||||||
|
post(sync_all_account_tokens),
|
||||||
|
)
|
||||||
|
.route(
|
||||||
|
"/api/account-tokens/{id}",
|
||||||
|
get(get_account_token)
|
||||||
|
.put(update_account_token)
|
||||||
|
.delete(delete_account_token),
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Deserialize)]
|
||||||
|
pub struct ListAccountTokensQuery {
|
||||||
|
#[serde(rename = "accountId", alias = "account_id")]
|
||||||
|
pub account_id: Option<i64>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Deserialize)]
|
||||||
|
#[serde(rename_all = "camelCase")]
|
||||||
|
pub struct CreateAccountTokenRequest {
|
||||||
|
pub account_id: i64,
|
||||||
|
pub name: Option<String>,
|
||||||
|
pub token: Option<String>,
|
||||||
|
pub source: Option<String>,
|
||||||
|
pub token_group: Option<String>,
|
||||||
|
pub value_status: Option<String>,
|
||||||
|
pub enabled: Option<bool>,
|
||||||
|
pub is_default: Option<bool>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Deserialize)]
|
||||||
|
#[serde(rename_all = "camelCase")]
|
||||||
|
pub struct UpdateAccountTokenRequest {
|
||||||
|
pub account_id: Option<i64>,
|
||||||
|
pub name: Option<String>,
|
||||||
|
pub token: Option<String>,
|
||||||
|
pub source: Option<String>,
|
||||||
|
pub token_group: Option<String>,
|
||||||
|
pub value_status: Option<String>,
|
||||||
|
pub enabled: Option<bool>,
|
||||||
|
pub is_default: Option<bool>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// GET /api/account-tokens?accountId=
|
||||||
|
pub async fn list_account_tokens(
|
||||||
|
State(state): State<AppState>,
|
||||||
|
Query(query): Query<ListAccountTokensQuery>,
|
||||||
|
) -> AppResult<Json<Value>> {
|
||||||
|
info!(
|
||||||
|
"Listing account tokens, account_id filter: {:?}",
|
||||||
|
query.account_id
|
||||||
|
);
|
||||||
|
|
||||||
|
let tokens: Vec<AccountToken> = if let Some(account_id) = query.account_id {
|
||||||
|
sqlx::query_as::<_, AccountToken>(
|
||||||
|
"SELECT * FROM account_tokens WHERE account_id = ? ORDER BY token_group, id",
|
||||||
|
)
|
||||||
|
.bind(account_id)
|
||||||
|
.fetch_all(&*state.db)
|
||||||
|
.await?
|
||||||
|
} else {
|
||||||
|
sqlx::query_as::<_, AccountToken>(
|
||||||
|
"SELECT * FROM account_tokens ORDER BY account_id, token_group, id",
|
||||||
|
)
|
||||||
|
.fetch_all(&*state.db)
|
||||||
|
.await?
|
||||||
|
};
|
||||||
|
|
||||||
|
Ok(Json(json!({ "data": tokens })))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// POST /api/account-tokens
|
||||||
|
pub async fn create_account_token(
|
||||||
|
State(state): State<AppState>,
|
||||||
|
Json(payload): Json<CreateAccountTokenRequest>,
|
||||||
|
) -> AppResult<Json<Value>> {
|
||||||
|
info!(
|
||||||
|
"Creating account token for account_id: {}",
|
||||||
|
payload.account_id
|
||||||
|
);
|
||||||
|
|
||||||
|
ensure_account_exists(&state, payload.account_id).await?;
|
||||||
|
|
||||||
|
let token = normalize_create_token(payload.token);
|
||||||
|
let name = payload
|
||||||
|
.name
|
||||||
|
.map(|name| name.trim().to_string())
|
||||||
|
.filter(|name| !name.is_empty())
|
||||||
|
.unwrap_or_else(|| "API Token".to_string());
|
||||||
|
|
||||||
|
let value_status = payload
|
||||||
|
.value_status
|
||||||
|
.map(|status| status.trim().to_string())
|
||||||
|
.filter(|status| !status.is_empty())
|
||||||
|
.unwrap_or_else(|| "valid".to_string());
|
||||||
|
|
||||||
|
let now = Utc::now().to_rfc3339();
|
||||||
|
|
||||||
|
let result = sqlx::query(
|
||||||
|
"INSERT INTO account_tokens \
|
||||||
|
(account_id, name, token, source, token_group, value_status, enabled, is_default, created_at, updated_at) \
|
||||||
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
|
||||||
|
)
|
||||||
|
.bind(payload.account_id)
|
||||||
|
.bind(&name)
|
||||||
|
.bind(&token)
|
||||||
|
.bind(&payload.source)
|
||||||
|
.bind(&payload.token_group)
|
||||||
|
.bind(&value_status)
|
||||||
|
.bind(payload.enabled.unwrap_or(true))
|
||||||
|
.bind(payload.is_default.unwrap_or(false))
|
||||||
|
.bind(&now)
|
||||||
|
.bind(&now)
|
||||||
|
.execute(&*state.db)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
let token_id = result.last_insert_rowid();
|
||||||
|
let account_token: AccountToken =
|
||||||
|
sqlx::query_as::<_, AccountToken>("SELECT * FROM account_tokens WHERE id = ?")
|
||||||
|
.bind(token_id)
|
||||||
|
.fetch_one(&*state.db)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
Ok(Json(json!({ "data": account_token })))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// GET /api/account-tokens/{id}
|
||||||
|
pub async fn get_account_token(
|
||||||
|
State(state): State<AppState>,
|
||||||
|
Path(id): Path<i64>,
|
||||||
|
) -> AppResult<Json<Value>> {
|
||||||
|
info!("Getting account token: {}", id);
|
||||||
|
|
||||||
|
let account_token: AccountToken =
|
||||||
|
sqlx::query_as::<_, AccountToken>("SELECT * FROM account_tokens WHERE id = ?")
|
||||||
|
.bind(id)
|
||||||
|
.fetch_optional(&*state.db)
|
||||||
|
.await?
|
||||||
|
.ok_or_else(|| AppError::NotFound(format!("Account token with id {} not found", id)))?;
|
||||||
|
|
||||||
|
Ok(Json(json!({ "data": account_token })))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// PUT /api/account-tokens/{id}
|
||||||
|
pub async fn update_account_token(
|
||||||
|
State(state): State<AppState>,
|
||||||
|
Path(id): Path<i64>,
|
||||||
|
Json(payload): Json<UpdateAccountTokenRequest>,
|
||||||
|
) -> AppResult<Json<Value>> {
|
||||||
|
info!("Updating account token: {}", id);
|
||||||
|
|
||||||
|
let existing: AccountToken =
|
||||||
|
sqlx::query_as::<_, AccountToken>("SELECT * FROM account_tokens WHERE id = ?")
|
||||||
|
.bind(id)
|
||||||
|
.fetch_optional(&*state.db)
|
||||||
|
.await?
|
||||||
|
.ok_or_else(|| AppError::NotFound(format!("Account token with id {} not found", id)))?;
|
||||||
|
|
||||||
|
let account_id = payload.account_id.unwrap_or(existing.account_id);
|
||||||
|
if payload.account_id.is_some() {
|
||||||
|
ensure_account_exists(&state, account_id).await?;
|
||||||
|
}
|
||||||
|
|
||||||
|
let name = payload
|
||||||
|
.name
|
||||||
|
.map(|name| name.trim().to_string())
|
||||||
|
.filter(|name| !name.is_empty())
|
||||||
|
.unwrap_or(existing.name);
|
||||||
|
|
||||||
|
let token = match payload.token {
|
||||||
|
Some(token) => {
|
||||||
|
let token = token.trim().to_string();
|
||||||
|
if token.is_empty() {
|
||||||
|
return Err(AppError::BadRequest("token cannot be empty".into()));
|
||||||
|
}
|
||||||
|
token
|
||||||
|
}
|
||||||
|
None => existing.token,
|
||||||
|
};
|
||||||
|
|
||||||
|
let value_status = payload
|
||||||
|
.value_status
|
||||||
|
.map(|status| status.trim().to_string())
|
||||||
|
.filter(|status| !status.is_empty())
|
||||||
|
.unwrap_or(existing.value_status);
|
||||||
|
|
||||||
|
let now = Utc::now().to_rfc3339();
|
||||||
|
|
||||||
|
sqlx::query(
|
||||||
|
"UPDATE account_tokens SET \
|
||||||
|
account_id = ?, name = ?, token = ?, source = ?, token_group = ?, value_status = ?, \
|
||||||
|
enabled = ?, is_default = ?, updated_at = ? \
|
||||||
|
WHERE id = ?",
|
||||||
|
)
|
||||||
|
.bind(account_id)
|
||||||
|
.bind(&name)
|
||||||
|
.bind(&token)
|
||||||
|
.bind(payload.source.or(existing.source))
|
||||||
|
.bind(payload.token_group.or(existing.token_group))
|
||||||
|
.bind(&value_status)
|
||||||
|
.bind(payload.enabled.or(existing.enabled))
|
||||||
|
.bind(payload.is_default.or(existing.is_default))
|
||||||
|
.bind(&now)
|
||||||
|
.bind(id)
|
||||||
|
.execute(&*state.db)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
let account_token: AccountToken =
|
||||||
|
sqlx::query_as::<_, AccountToken>("SELECT * FROM account_tokens WHERE id = ?")
|
||||||
|
.bind(id)
|
||||||
|
.fetch_one(&*state.db)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
Ok(Json(json!({ "data": account_token })))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// DELETE /api/account-tokens/{id}
|
||||||
|
pub async fn delete_account_token(
|
||||||
|
State(state): State<AppState>,
|
||||||
|
Path(id): Path<i64>,
|
||||||
|
) -> AppResult<Json<Value>> {
|
||||||
|
info!("Deleting account token: {}", id);
|
||||||
|
|
||||||
|
let affected = sqlx::query("DELETE FROM account_tokens WHERE id = ?")
|
||||||
|
.bind(id)
|
||||||
|
.execute(&*state.db)
|
||||||
|
.await?
|
||||||
|
.rows_affected();
|
||||||
|
|
||||||
|
if affected == 0 {
|
||||||
|
return Err(AppError::NotFound(format!(
|
||||||
|
"Account token with id {} not found",
|
||||||
|
id
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(Json(json!({ "success": true })))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// POST /api/account-tokens/sync-all
|
||||||
|
pub async fn sync_all_account_tokens() -> (StatusCode, Json<Value>) {
|
||||||
|
info!("Sync all account tokens requested");
|
||||||
|
|
||||||
|
(
|
||||||
|
StatusCode::NOT_IMPLEMENTED,
|
||||||
|
Json(json!({
|
||||||
|
"error": "not implemented"
|
||||||
|
})),
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn ensure_account_exists(state: &AppState, account_id: i64) -> AppResult<()> {
|
||||||
|
let account_exists: Option<(i64,)> = sqlx::query_as("SELECT id FROM accounts WHERE id = ?")
|
||||||
|
.bind(account_id)
|
||||||
|
.fetch_optional(&*state.db)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
if account_exists.is_none() {
|
||||||
|
return Err(AppError::BadRequest(format!(
|
||||||
|
"Account with id {} not found",
|
||||||
|
account_id
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn normalize_create_token(token: Option<String>) -> String {
|
||||||
|
token
|
||||||
|
.map(|token| token.trim().to_string())
|
||||||
|
.filter(|token| !token.is_empty())
|
||||||
|
.unwrap_or_else(|| Uuid::new_v4().to_string())
|
||||||
|
}
|
||||||
@@ -14,6 +14,7 @@ use axum::{
|
|||||||
use serde_json::{json, Value};
|
use serde_json::{json, Value};
|
||||||
use tracing::info;
|
use tracing::info;
|
||||||
|
|
||||||
|
mod account_tokens;
|
||||||
mod accounts;
|
mod accounts;
|
||||||
mod downstream_keys;
|
mod downstream_keys;
|
||||||
mod sites;
|
mod sites;
|
||||||
@@ -116,6 +117,7 @@ pub fn routes(state: AppState) -> Router<AppState> {
|
|||||||
.merge(downstream_routes)
|
.merge(downstream_routes)
|
||||||
.merge(settings_routes)
|
.merge(settings_routes)
|
||||||
.merge(users_routes)
|
.merge(users_routes)
|
||||||
|
.merge(account_tokens::routes())
|
||||||
.route_layer(middleware::from_fn_with_state(state, require_admin_auth))
|
.route_layer(middleware::from_fn_with_state(state, require_admin_auth))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -59,8 +59,7 @@ impl From<User> for UserResponse {
|
|||||||
/// GET /api/admin/users
|
/// GET /api/admin/users
|
||||||
pub async fn list_users(State(state): State<AppState>) -> AppResult<Json<Value>> {
|
pub async fn list_users(State(state): State<AppState>) -> AppResult<Json<Value>> {
|
||||||
info!("Listing admin users");
|
info!("Listing admin users");
|
||||||
let users: Vec<User> =
|
let users: Vec<User> = sqlx::query_as::<_, User>("SELECT * FROM users ORDER BY id")
|
||||||
sqlx::query_as::<_, User>("SELECT * FROM users ORDER BY id")
|
|
||||||
.fetch_all(&*state.db)
|
.fetch_all(&*state.db)
|
||||||
.await?;
|
.await?;
|
||||||
let items: Vec<UserResponse> = users.into_iter().map(Into::into).collect();
|
let items: Vec<UserResponse> = users.into_iter().map(Into::into).collect();
|
||||||
@@ -95,16 +94,17 @@ pub async fn create_user(
|
|||||||
return Err(AppError::BadRequest("email is required".into()));
|
return Err(AppError::BadRequest("email is required".into()));
|
||||||
}
|
}
|
||||||
if payload.password.len() < 6 {
|
if payload.password.len() < 6 {
|
||||||
return Err(AppError::BadRequest("password must be at least 6 characters".into()));
|
return Err(AppError::BadRequest(
|
||||||
|
"password must be at least 6 characters".into(),
|
||||||
|
));
|
||||||
}
|
}
|
||||||
|
|
||||||
let password_hash = hash_password(&payload.password)
|
let password_hash = hash_password(&payload.password)
|
||||||
.map_err(|e| AppError::Internal(anyhow::anyhow!("Failed to hash password: {e}")))?;
|
.map_err(|e| AppError::Internal(anyhow::anyhow!("Failed to hash password: {e}")))?;
|
||||||
let role = payload.role.unwrap_or_else(|| "user".to_string());
|
let role = payload.role.unwrap_or_else(|| "user".to_string());
|
||||||
|
|
||||||
let result = sqlx::query(
|
let result =
|
||||||
"INSERT INTO users (username, email, password_hash, role) VALUES (?, ?, ?, ?)"
|
sqlx::query("INSERT INTO users (username, email, password_hash, role) VALUES (?, ?, ?, ?)")
|
||||||
)
|
|
||||||
.bind(&payload.username)
|
.bind(&payload.username)
|
||||||
.bind(&payload.email)
|
.bind(&payload.email)
|
||||||
.bind(&password_hash)
|
.bind(&password_hash)
|
||||||
@@ -117,7 +117,10 @@ pub async fn create_user(
|
|||||||
.fetch_one(&*state.db)
|
.fetch_one(&*state.db)
|
||||||
.await?;
|
.await?;
|
||||||
|
|
||||||
Ok((axum::http::StatusCode::CREATED, Json(json!({ "data": UserResponse::from(user) }))))
|
Ok((
|
||||||
|
axum::http::StatusCode::CREATED,
|
||||||
|
Json(json!({ "data": UserResponse::from(user) })),
|
||||||
|
))
|
||||||
}
|
}
|
||||||
|
|
||||||
/// PUT /api/admin/users/{id}
|
/// PUT /api/admin/users/{id}
|
||||||
@@ -136,9 +139,14 @@ pub async fn update_user(
|
|||||||
|
|
||||||
let password_hash = if let Some(ref pw) = payload.password {
|
let password_hash = if let Some(ref pw) = payload.password {
|
||||||
if pw.len() < 6 {
|
if pw.len() < 6 {
|
||||||
return Err(AppError::BadRequest("password must be at least 6 characters".into()));
|
return Err(AppError::BadRequest(
|
||||||
|
"password must be at least 6 characters".into(),
|
||||||
|
));
|
||||||
}
|
}
|
||||||
Some(hash_password(pw).map_err(|e| AppError::Internal(anyhow::anyhow!("Failed to hash password: {e}")))?)
|
Some(
|
||||||
|
hash_password(pw)
|
||||||
|
.map_err(|e| AppError::Internal(anyhow::anyhow!("Failed to hash password: {e}")))?,
|
||||||
|
)
|
||||||
} else {
|
} else {
|
||||||
None
|
None
|
||||||
};
|
};
|
||||||
|
|||||||
@@ -68,10 +68,14 @@ async fn register(
|
|||||||
info!("Registering user: {}", payload.email);
|
info!("Registering user: {}", payload.email);
|
||||||
|
|
||||||
if payload.email.trim().is_empty() || payload.username.trim().is_empty() {
|
if payload.email.trim().is_empty() || payload.username.trim().is_empty() {
|
||||||
return Err(AppError::BadRequest("email and username are required".into()));
|
return Err(AppError::BadRequest(
|
||||||
|
"email and username are required".into(),
|
||||||
|
));
|
||||||
}
|
}
|
||||||
if payload.password.len() < 6 {
|
if payload.password.len() < 6 {
|
||||||
return Err(AppError::BadRequest("password must be at least 6 characters".into()));
|
return Err(AppError::BadRequest(
|
||||||
|
"password must be at least 6 characters".into(),
|
||||||
|
));
|
||||||
}
|
}
|
||||||
|
|
||||||
// Check if email already exists
|
// Check if email already exists
|
||||||
@@ -88,7 +92,7 @@ async fn register(
|
|||||||
.map_err(|e| AppError::Internal(anyhow::anyhow!("Failed to hash password: {e}")))?;
|
.map_err(|e| AppError::Internal(anyhow::anyhow!("Failed to hash password: {e}")))?;
|
||||||
|
|
||||||
let result = sqlx::query(
|
let result = sqlx::query(
|
||||||
"INSERT INTO users (username, email, password_hash, role) VALUES (?, ?, ?, 'user')"
|
"INSERT INTO users (username, email, password_hash, role) VALUES (?, ?, ?, 'user')",
|
||||||
)
|
)
|
||||||
.bind(&payload.username)
|
.bind(&payload.username)
|
||||||
.bind(&payload.email)
|
.bind(&payload.email)
|
||||||
@@ -106,7 +110,8 @@ async fn register(
|
|||||||
&user.id.to_string(),
|
&user.id.to_string(),
|
||||||
&user.role,
|
&user.role,
|
||||||
state.config.auth.jwt_expiry_hours,
|
state.config.auth.jwt_expiry_hours,
|
||||||
).map_err(|e| AppError::Internal(anyhow::anyhow!("Failed to generate JWT: {e}")))?;
|
)
|
||||||
|
.map_err(|e| AppError::Internal(anyhow::anyhow!("Failed to generate JWT: {e}")))?;
|
||||||
|
|
||||||
Ok(Json(json!({
|
Ok(Json(json!({
|
||||||
"data": {
|
"data": {
|
||||||
@@ -150,7 +155,8 @@ async fn login(
|
|||||||
&user.id.to_string(),
|
&user.id.to_string(),
|
||||||
&user.role,
|
&user.role,
|
||||||
state.config.auth.jwt_expiry_hours,
|
state.config.auth.jwt_expiry_hours,
|
||||||
).map_err(|e| AppError::Internal(anyhow::anyhow!("Failed to generate JWT: {e}")))?;
|
)
|
||||||
|
.map_err(|e| AppError::Internal(anyhow::anyhow!("Failed to generate JWT: {e}")))?;
|
||||||
|
|
||||||
Ok(Json(json!({
|
Ok(Json(json!({
|
||||||
"data": {
|
"data": {
|
||||||
@@ -166,10 +172,7 @@ async fn login(
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// GET /api/auth/me — returns current user profile from JWT
|
/// GET /api/auth/me — returns current user profile from JWT
|
||||||
async fn me(
|
async fn me(State(state): State<AppState>, headers: HeaderMap) -> AppResult<Json<Value>> {
|
||||||
State(state): State<AppState>,
|
|
||||||
headers: HeaderMap,
|
|
||||||
) -> AppResult<Json<Value>> {
|
|
||||||
let auth = headers
|
let auth = headers
|
||||||
.get("Authorization")
|
.get("Authorization")
|
||||||
.and_then(|v| v.to_str().ok())
|
.and_then(|v| v.to_str().ok())
|
||||||
@@ -238,13 +241,19 @@ async fn change_password(
|
|||||||
) -> AppResult<Json<Value>> {
|
) -> AppResult<Json<Value>> {
|
||||||
let (user_id, _role) = extract_jwt(&state, &headers)?;
|
let (user_id, _role) = extract_jwt(&state, &headers)?;
|
||||||
|
|
||||||
let old_pw = body.get("oldPassword").and_then(|v| v.as_str())
|
let old_pw = body
|
||||||
|
.get("oldPassword")
|
||||||
|
.and_then(|v| v.as_str())
|
||||||
.ok_or_else(|| AppError::BadRequest("oldPassword is required".into()))?;
|
.ok_or_else(|| AppError::BadRequest("oldPassword is required".into()))?;
|
||||||
let new_pw = body.get("newPassword").and_then(|v| v.as_str())
|
let new_pw = body
|
||||||
|
.get("newPassword")
|
||||||
|
.and_then(|v| v.as_str())
|
||||||
.ok_or_else(|| AppError::BadRequest("newPassword is required".into()))?;
|
.ok_or_else(|| AppError::BadRequest("newPassword is required".into()))?;
|
||||||
|
|
||||||
if new_pw.len() < 6 {
|
if new_pw.len() < 6 {
|
||||||
return Err(AppError::BadRequest("new password must be at least 6 characters".into()));
|
return Err(AppError::BadRequest(
|
||||||
|
"new password must be at least 6 characters".into(),
|
||||||
|
));
|
||||||
}
|
}
|
||||||
|
|
||||||
let user: User = sqlx::query_as::<_, User>("SELECT * FROM users WHERE id = ?")
|
let user: User = sqlx::query_as::<_, User>("SELECT * FROM users WHERE id = ?")
|
||||||
@@ -257,7 +266,9 @@ async fn change_password(
|
|||||||
.map_err(|e| AppError::Internal(anyhow::anyhow!("Password verification failed: {e}")))?;
|
.map_err(|e| AppError::Internal(anyhow::anyhow!("Password verification failed: {e}")))?;
|
||||||
|
|
||||||
if !valid {
|
if !valid {
|
||||||
return Err(AppError::Unauthorized("Current password is incorrect".into()));
|
return Err(AppError::Unauthorized(
|
||||||
|
"Current password is incorrect".into(),
|
||||||
|
));
|
||||||
}
|
}
|
||||||
|
|
||||||
let new_hash = crate::auth::password::hash_password(new_pw)
|
let new_hash = crate::auth::password::hash_password(new_pw)
|
||||||
@@ -314,7 +325,9 @@ fn extract_jwt(state: &AppState, headers: &HeaderMap) -> AppResult<(i64, String)
|
|||||||
let claims = crate::auth::jwt::validate_jwt(&state.config.auth.jwt_secret, auth)
|
let claims = crate::auth::jwt::validate_jwt(&state.config.auth.jwt_secret, auth)
|
||||||
.map_err(|_| AppError::Unauthorized("Invalid or expired token".into()))?;
|
.map_err(|_| AppError::Unauthorized("Invalid or expired token".into()))?;
|
||||||
|
|
||||||
let uid: i64 = claims.sub.parse()
|
let uid: i64 = claims
|
||||||
|
.sub
|
||||||
|
.parse()
|
||||||
.map_err(|_| AppError::Unauthorized("Invalid user ID in token".into()))?;
|
.map_err(|_| AppError::Unauthorized("Invalid user ID in token".into()))?;
|
||||||
|
|
||||||
Ok((uid, claims.role))
|
Ok((uid, claims.role))
|
||||||
|
|||||||
@@ -48,7 +48,10 @@ pub fn routes() -> Router<AppState> {
|
|||||||
.route("/api/checkin", get(list_checkin_logs))
|
.route("/api/checkin", get(list_checkin_logs))
|
||||||
.route("/api/checkin/stats", get(checkin_stats))
|
.route("/api/checkin/stats", get(checkin_stats))
|
||||||
.route("/api/checkin/accounts", get(list_checkin_accounts))
|
.route("/api/checkin/accounts", get(list_checkin_accounts))
|
||||||
.route("/api/checkin/accounts/{id}/toggle", post(toggle_checkin_account))
|
.route(
|
||||||
|
"/api/checkin/accounts/{id}/toggle",
|
||||||
|
post(toggle_checkin_account),
|
||||||
|
)
|
||||||
.route("/api/checkin/run", post(run_checkin))
|
.route("/api/checkin/run", post(run_checkin))
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -96,15 +99,20 @@ async fn checkin_stats(State(state): State<AppState>) -> AppResult<Json<Value>>
|
|||||||
FROM accounts a
|
FROM accounts a
|
||||||
JOIN sites s ON s.id = a.site_id
|
JOIN sites s ON s.id = a.site_id
|
||||||
WHERE a.status = 'active'
|
WHERE a.status = 'active'
|
||||||
ORDER BY a.last_checkin_at ASC"
|
ORDER BY a.last_checkin_at ASC",
|
||||||
)
|
)
|
||||||
.fetch_all(&*state.db)
|
.fetch_all(&*state.db)
|
||||||
.await?;
|
.await?;
|
||||||
|
|
||||||
let total = rows.len();
|
let total = rows.len();
|
||||||
let checked_today = rows.iter().filter(|r| {
|
let checked_today = rows
|
||||||
r.last_checkin.as_deref().map_or(false, |d| d.starts_with(&today_date()))
|
.iter()
|
||||||
}).count();
|
.filter(|r| {
|
||||||
|
r.last_checkin
|
||||||
|
.as_deref()
|
||||||
|
.map_or(false, |d| d.starts_with(&today_date()))
|
||||||
|
})
|
||||||
|
.count();
|
||||||
|
|
||||||
Ok(Json(json!({
|
Ok(Json(json!({
|
||||||
"total_accounts": total,
|
"total_accounts": total,
|
||||||
@@ -116,13 +124,13 @@ async fn checkin_stats(State(state): State<AppState>) -> AppResult<Json<Value>>
|
|||||||
|
|
||||||
/// GET /api/checkin/accounts — list accounts with checkin info
|
/// GET /api/checkin/accounts — list accounts with checkin info
|
||||||
async fn list_checkin_accounts(State(state): State<AppState>) -> AppResult<Json<Value>> {
|
async fn list_checkin_accounts(State(state): State<AppState>) -> AppResult<Json<Value>> {
|
||||||
let accounts: Vec<Value> = sqlx::query_as::<_, Account>(
|
let accounts: Vec<Value> =
|
||||||
"SELECT * FROM accounts ORDER BY site_id, username"
|
sqlx::query_as::<_, Account>("SELECT * FROM accounts ORDER BY site_id, username")
|
||||||
)
|
|
||||||
.fetch_all(&*state.db)
|
.fetch_all(&*state.db)
|
||||||
.await?
|
.await?
|
||||||
.into_iter()
|
.into_iter()
|
||||||
.map(|a| json!({
|
.map(|a| {
|
||||||
|
json!({
|
||||||
"id": a.id,
|
"id": a.id,
|
||||||
"site_id": a.site_id,
|
"site_id": a.site_id,
|
||||||
"username": a.username,
|
"username": a.username,
|
||||||
@@ -130,7 +138,8 @@ async fn list_checkin_accounts(State(state): State<AppState>) -> AppResult<Json<
|
|||||||
"checkin_enabled": a.checkin_enabled,
|
"checkin_enabled": a.checkin_enabled,
|
||||||
"last_checkin_at": a.last_checkin_at,
|
"last_checkin_at": a.last_checkin_at,
|
||||||
"balance": a.balance,
|
"balance": a.balance,
|
||||||
}))
|
})
|
||||||
|
})
|
||||||
.collect();
|
.collect();
|
||||||
|
|
||||||
Ok(Json(json!({ "data": accounts })))
|
Ok(Json(json!({ "data": accounts })))
|
||||||
@@ -150,13 +159,17 @@ async fn toggle_checkin_account(
|
|||||||
.ok_or_else(|| AppError::NotFound(format!("Account {} not found", id)))?;
|
.ok_or_else(|| AppError::NotFound(format!("Account {} not found", id)))?;
|
||||||
|
|
||||||
let new_value = !account.checkin_enabled.unwrap_or(true);
|
let new_value = !account.checkin_enabled.unwrap_or(true);
|
||||||
sqlx::query("UPDATE accounts SET checkin_enabled = ?, updated_at = datetime('now') WHERE id = ?")
|
sqlx::query(
|
||||||
|
"UPDATE accounts SET checkin_enabled = ?, updated_at = datetime('now') WHERE id = ?",
|
||||||
|
)
|
||||||
.bind(new_value)
|
.bind(new_value)
|
||||||
.bind(id)
|
.bind(id)
|
||||||
.execute(&*state.db)
|
.execute(&*state.db)
|
||||||
.await?;
|
.await?;
|
||||||
|
|
||||||
Ok(Json(json!({ "success": true, "checkin_enabled": new_value })))
|
Ok(Json(
|
||||||
|
json!({ "success": true, "checkin_enabled": new_value }),
|
||||||
|
))
|
||||||
}
|
}
|
||||||
|
|
||||||
/// POST /api/checkin/run — trigger checkin for all enabled accounts
|
/// POST /api/checkin/run — trigger checkin for all enabled accounts
|
||||||
@@ -164,7 +177,7 @@ async fn run_checkin(State(state): State<AppState>) -> AppResult<Json<Value>> {
|
|||||||
info!("Running checkin for all enabled accounts");
|
info!("Running checkin for all enabled accounts");
|
||||||
|
|
||||||
let accounts: Vec<Account> = sqlx::query_as::<_, Account>(
|
let accounts: Vec<Account> = sqlx::query_as::<_, Account>(
|
||||||
"SELECT * FROM accounts WHERE status = 'active' AND COALESCE(checkin_enabled, 1) = 1"
|
"SELECT * FROM accounts WHERE status = 'active' AND COALESCE(checkin_enabled, 1) = 1",
|
||||||
)
|
)
|
||||||
.fetch_all(&*state.db)
|
.fetch_all(&*state.db)
|
||||||
.await?;
|
.await?;
|
||||||
|
|||||||
@@ -5,15 +5,15 @@
|
|||||||
// The /api/comfyui/proxy/* route serves static assets (public).
|
// The /api/comfyui/proxy/* route serves static assets (public).
|
||||||
|
|
||||||
use crate::{
|
use crate::{
|
||||||
auth::api_key::{ authorize_downstream_token, DownstreamAuthResult },
|
auth::api_key::{authorize_downstream_token, DownstreamAuthResult},
|
||||||
error::{ AppError, AppResult },
|
error::{AppError, AppResult},
|
||||||
AppState,
|
AppState,
|
||||||
};
|
};
|
||||||
use axum::{
|
use axum::{
|
||||||
extract::{ Path, Query, State },
|
extract::{Path, Query, State},
|
||||||
http::{ HeaderMap, StatusCode },
|
http::{HeaderMap, StatusCode},
|
||||||
response::{ Html, IntoResponse, Response },
|
response::{Html, IntoResponse, Response},
|
||||||
routing::{ get, post },
|
routing::{get, post},
|
||||||
Json, Router,
|
Json, Router,
|
||||||
};
|
};
|
||||||
use reqwest::Client;
|
use reqwest::Client;
|
||||||
@@ -49,7 +49,8 @@ async fn check_auth(state: &AppState, headers: &HeaderMap) -> Result<(), AppErro
|
|||||||
|
|
||||||
// Downstream API key
|
// Downstream API key
|
||||||
if auth.starts_with("sk-") {
|
if auth.starts_with("sk-") {
|
||||||
let result = authorize_downstream_token(auth, &state.db).await
|
let result = authorize_downstream_token(auth, &state.db)
|
||||||
|
.await
|
||||||
.map_err(|e| AppError::Unauthorized(e))?;
|
.map_err(|e| AppError::Unauthorized(e))?;
|
||||||
|
|
||||||
match result {
|
match result {
|
||||||
@@ -123,10 +124,7 @@ async fn proxy_get_history_by_id(
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// GET /api/comfyui/queue
|
/// GET /api/comfyui/queue
|
||||||
async fn proxy_get_queue(
|
async fn proxy_get_queue(State(state): State<AppState>, headers: HeaderMap) -> AppResult<Response> {
|
||||||
State(state): State<AppState>,
|
|
||||||
headers: HeaderMap,
|
|
||||||
) -> AppResult<Response> {
|
|
||||||
info!("GET /api/comfyui/queue");
|
info!("GET /api/comfyui/queue");
|
||||||
check_auth(&state, &headers).await?;
|
check_auth(&state, &headers).await?;
|
||||||
proxy_to_comfyui("GET", "/queue", None, None).await
|
proxy_to_comfyui("GET", "/queue", None, None).await
|
||||||
@@ -149,9 +147,7 @@ async fn proxy_get_view(
|
|||||||
// ── Static proxy (public — assets/JS/CSS for iframe) ─────────────────
|
// ── Static proxy (public — assets/JS/CSS for iframe) ─────────────────
|
||||||
|
|
||||||
/// GET /api/comfyui/proxy/{path} — proxy everything (HTML, assets)
|
/// GET /api/comfyui/proxy/{path} — proxy everything (HTML, assets)
|
||||||
async fn proxy_static(
|
async fn proxy_static(Path(path): Path<String>) -> AppResult<Response> {
|
||||||
Path(path): Path<String>,
|
|
||||||
) -> AppResult<Response> {
|
|
||||||
info!("GET /api/comfyui/proxy/{}", path);
|
info!("GET /api/comfyui/proxy/{}", path);
|
||||||
let target_path = if path.is_empty() || path == "/" {
|
let target_path = if path.is_empty() || path == "/" {
|
||||||
"/".to_string()
|
"/".to_string()
|
||||||
@@ -161,18 +157,24 @@ async fn proxy_static(
|
|||||||
|
|
||||||
let client = http_client();
|
let client = http_client();
|
||||||
let url = format!("{}{}", comfyui_base(), target_path);
|
let url = format!("{}{}", comfyui_base(), target_path);
|
||||||
let resp = client.get(&url).send().await
|
let resp = client
|
||||||
|
.get(&url)
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
.map_err(|e| AppError::UpstreamError(format!("ComfyUI unreachable: {e}")))?;
|
.map_err(|e| AppError::UpstreamError(format!("ComfyUI unreachable: {e}")))?;
|
||||||
|
|
||||||
let status = resp.status();
|
let status = resp.status();
|
||||||
let content_type = resp.headers()
|
let content_type = resp
|
||||||
|
.headers()
|
||||||
.get("content-type")
|
.get("content-type")
|
||||||
.and_then(|v| v.to_str().ok())
|
.and_then(|v| v.to_str().ok())
|
||||||
.unwrap_or("application/octet-stream")
|
.unwrap_or("application/octet-stream")
|
||||||
.to_string();
|
.to_string();
|
||||||
|
|
||||||
if content_type.contains("html") {
|
if content_type.contains("html") {
|
||||||
let text = resp.text().await
|
let text = resp
|
||||||
|
.text()
|
||||||
|
.await
|
||||||
.map_err(|e| AppError::UpstreamError(e.to_string()))?;
|
.map_err(|e| AppError::UpstreamError(e.to_string()))?;
|
||||||
// Rebrand: ComfyUI → BoosAPI
|
// Rebrand: ComfyUI → BoosAPI
|
||||||
let rebranded = text
|
let rebranded = text
|
||||||
@@ -180,21 +182,41 @@ async fn proxy_static(
|
|||||||
.replace("comfyui", "BoosAPI");
|
.replace("comfyui", "BoosAPI");
|
||||||
// Rewrite URLs
|
// Rewrite URLs
|
||||||
let rewritten = rewrite_proxy_urls(&rebranded);
|
let rewritten = rewrite_proxy_urls(&rebranded);
|
||||||
Ok((StatusCode::from_u16(status.as_u16()).unwrap(), [(axum::http::header::CONTENT_TYPE, "text/html; charset=utf-8")], rewritten).into_response())
|
Ok((
|
||||||
|
StatusCode::from_u16(status.as_u16()).unwrap(),
|
||||||
|
[(axum::http::header::CONTENT_TYPE, "text/html; charset=utf-8")],
|
||||||
|
rewritten,
|
||||||
|
)
|
||||||
|
.into_response())
|
||||||
} else if content_type.contains("text") || content_type.contains("json") {
|
} else if content_type.contains("text") || content_type.contains("json") {
|
||||||
let text = resp.text().await
|
let text = resp
|
||||||
|
.text()
|
||||||
|
.await
|
||||||
.map_err(|e| AppError::UpstreamError(e.to_string()))?;
|
.map_err(|e| AppError::UpstreamError(e.to_string()))?;
|
||||||
Ok((StatusCode::from_u16(status.as_u16()).unwrap(), [(axum::http::header::CONTENT_TYPE, content_type)], text).into_response())
|
Ok((
|
||||||
|
StatusCode::from_u16(status.as_u16()).unwrap(),
|
||||||
|
[(axum::http::header::CONTENT_TYPE, content_type)],
|
||||||
|
text,
|
||||||
|
)
|
||||||
|
.into_response())
|
||||||
} else {
|
} else {
|
||||||
let bytes = resp.bytes().await
|
let bytes = resp
|
||||||
|
.bytes()
|
||||||
|
.await
|
||||||
.map_err(|e| AppError::UpstreamError(e.to_string()))?;
|
.map_err(|e| AppError::UpstreamError(e.to_string()))?;
|
||||||
Ok((StatusCode::from_u16(status.as_u16()).unwrap(), [(axum::http::header::CONTENT_TYPE, content_type)], bytes).into_response())
|
Ok((
|
||||||
|
StatusCode::from_u16(status.as_u16()).unwrap(),
|
||||||
|
[(axum::http::header::CONTENT_TYPE, content_type)],
|
||||||
|
bytes,
|
||||||
|
)
|
||||||
|
.into_response())
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// GET /comfyui — HTML page with iframe
|
/// GET /comfyui — HTML page with iframe
|
||||||
async fn comfyui_page() -> Html<&'static str> {
|
async fn comfyui_page() -> Html<&'static str> {
|
||||||
Html(r##"<!DOCTYPE html>
|
Html(
|
||||||
|
r##"<!DOCTYPE html>
|
||||||
<html lang="zh-CN">
|
<html lang="zh-CN">
|
||||||
<head>
|
<head>
|
||||||
<meta charset="UTF-8">
|
<meta charset="UTF-8">
|
||||||
@@ -209,7 +231,8 @@ async fn comfyui_page() -> Html<&'static str> {
|
|||||||
<body>
|
<body>
|
||||||
<iframe src="/api/comfyui/proxy/" title="BoosAPI ComfyUI"></iframe>
|
<iframe src="/api/comfyui/proxy/" title="BoosAPI ComfyUI"></iframe>
|
||||||
</body>
|
</body>
|
||||||
</html>"##)
|
</html>"##,
|
||||||
|
)
|
||||||
}
|
}
|
||||||
|
|
||||||
// ── Helpers ──────────────────────────────────────────────────────────
|
// ── Helpers ──────────────────────────────────────────────────────────
|
||||||
@@ -225,12 +248,23 @@ fn rewrite_proxy_urls(html: &str) -> String {
|
|||||||
// Look ahead for rc="/", ref="/", ction="/"
|
// Look ahead for rc="/", ref="/", ction="/"
|
||||||
let snippet: String = chars.clone().take(6).map(|(_, c)| c).collect();
|
let snippet: String = chars.clone().take(6).map(|(_, c)| c).collect();
|
||||||
let full = format!("{}{}", ch, snippet);
|
let full = format!("{}{}", ch, snippet);
|
||||||
if full.starts_with("src=\"/") || full.starts_with("href=\"/") || full.starts_with("action=\"/") {
|
if full.starts_with("src=\"/")
|
||||||
let attr_end = if full.starts_with("src=\"") { 4 } else if full.starts_with("href=\"") { 5 } else { 7 };
|
|| full.starts_with("href=\"/")
|
||||||
|
|| full.starts_with("action=\"/")
|
||||||
|
{
|
||||||
|
let attr_end = if full.starts_with("src=\"") {
|
||||||
|
4
|
||||||
|
} else if full.starts_with("href=\"") {
|
||||||
|
5
|
||||||
|
} else {
|
||||||
|
7
|
||||||
|
};
|
||||||
result.push_str(&full[..attr_end]);
|
result.push_str(&full[..attr_end]);
|
||||||
result.push_str("/api/comfyui/proxy/");
|
result.push_str("/api/comfyui/proxy/");
|
||||||
// Skip the rest of the attribute
|
// Skip the rest of the attribute
|
||||||
for _ in 0..snippet.len() { chars.next(); }
|
for _ in 0..snippet.len() {
|
||||||
|
chars.next();
|
||||||
|
}
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -275,22 +309,34 @@ async fn proxy_to_comfyui(
|
|||||||
|
|
||||||
let req = match method {
|
let req = match method {
|
||||||
"POST" => {
|
"POST" => {
|
||||||
let body_str = body.map(|b| b.to_string()).unwrap_or_else(|| "{}".to_string());
|
let body_str = body
|
||||||
client.post(&url)
|
.map(|b| b.to_string())
|
||||||
|
.unwrap_or_else(|| "{}".to_string());
|
||||||
|
client
|
||||||
|
.post(&url)
|
||||||
.header("Content-Type", "application/json")
|
.header("Content-Type", "application/json")
|
||||||
.body(body_str)
|
.body(body_str)
|
||||||
}
|
}
|
||||||
_ => client.get(&url),
|
_ => client.get(&url),
|
||||||
};
|
};
|
||||||
|
|
||||||
let resp = req.send().await
|
let resp = req
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
.map_err(|e| AppError::UpstreamError(format!("ComfyUI unreachable: {e}")))?;
|
.map_err(|e| AppError::UpstreamError(format!("ComfyUI unreachable: {e}")))?;
|
||||||
|
|
||||||
let status = resp.status();
|
let status = resp.status();
|
||||||
let text = resp.text().await
|
let text = resp
|
||||||
|
.text()
|
||||||
|
.await
|
||||||
.map_err(|e| AppError::UpstreamError(e.to_string()))?;
|
.map_err(|e| AppError::UpstreamError(e.to_string()))?;
|
||||||
|
|
||||||
Ok((StatusCode::from_u16(status.as_u16()).unwrap(), [(axum::http::header::CONTENT_TYPE, "application/json")], text).into_response())
|
Ok((
|
||||||
|
StatusCode::from_u16(status.as_u16()).unwrap(),
|
||||||
|
[(axum::http::header::CONTENT_TYPE, "application/json")],
|
||||||
|
text,
|
||||||
|
)
|
||||||
|
.into_response())
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Proxy to ComfyUI — returns binary response (for images)
|
/// Proxy to ComfyUI — returns binary response (for images)
|
||||||
@@ -299,17 +345,27 @@ async fn proxy_to_comfyui_binary(_method: &str, path: &str) -> AppResult<Respons
|
|||||||
let url = format!("{}{}", comfyui_base(), path);
|
let url = format!("{}{}", comfyui_base(), path);
|
||||||
|
|
||||||
let req = client.get(&url);
|
let req = client.get(&url);
|
||||||
let resp = req.send().await
|
let resp = req
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
.map_err(|e| AppError::UpstreamError(format!("ComfyUI unreachable: {e}")))?;
|
.map_err(|e| AppError::UpstreamError(format!("ComfyUI unreachable: {e}")))?;
|
||||||
|
|
||||||
let status = resp.status();
|
let status = resp.status();
|
||||||
let content_type = resp.headers()
|
let content_type = resp
|
||||||
|
.headers()
|
||||||
.get("content-type")
|
.get("content-type")
|
||||||
.and_then(|v| v.to_str().ok())
|
.and_then(|v| v.to_str().ok())
|
||||||
.unwrap_or("application/octet-stream")
|
.unwrap_or("application/octet-stream")
|
||||||
.to_string();
|
.to_string();
|
||||||
let bytes = resp.bytes().await
|
let bytes = resp
|
||||||
|
.bytes()
|
||||||
|
.await
|
||||||
.map_err(|e| AppError::UpstreamError(e.to_string()))?;
|
.map_err(|e| AppError::UpstreamError(e.to_string()))?;
|
||||||
|
|
||||||
Ok((StatusCode::from_u16(status.as_u16()).unwrap(), [(axum::http::header::CONTENT_TYPE, content_type)], bytes).into_response())
|
Ok((
|
||||||
|
StatusCode::from_u16(status.as_u16()).unwrap(),
|
||||||
|
[(axum::http::header::CONTENT_TYPE, content_type)],
|
||||||
|
bytes,
|
||||||
|
)
|
||||||
|
.into_response())
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -0,0 +1,242 @@
|
|||||||
|
use crate::AppState;
|
||||||
|
use axum::{
|
||||||
|
extract::State,
|
||||||
|
http::HeaderMap,
|
||||||
|
response::sse::{Event, Sse},
|
||||||
|
routing::post,
|
||||||
|
Json, Router,
|
||||||
|
};
|
||||||
|
use futures::stream::Stream;
|
||||||
|
use reqwest::Client;
|
||||||
|
use serde_json::{json, Value};
|
||||||
|
use std::convert::Infallible;
|
||||||
|
use tokio::sync::mpsc::{unbounded_channel, UnboundedSender};
|
||||||
|
use tokio_stream::{wrappers::UnboundedReceiverStream, StreamExt};
|
||||||
|
use tracing::info;
|
||||||
|
|
||||||
|
const SYSTEM_PROMPT: &str = r#"你是 BoosAPI 的 AI 视频生成助手,帮助用户通过 ComfyUI 制作 AI 视频。
|
||||||
|
|
||||||
|
## 能力
|
||||||
|
1. **分析脚本** — 提取角色、场景、对话
|
||||||
|
2. **生成角色提示词** — 为每个角色生成 4 个角度(正面/背面/左侧面/右侧面)
|
||||||
|
3. **生成场景提示词** — 描述场景的稳定图像生成提示词
|
||||||
|
4. **生成语音(TTS)** — 将对话台词转为语音描述
|
||||||
|
5. **导出 ComfyUI 工作流** — 生成 workflow JSON
|
||||||
|
|
||||||
|
## 输出格式
|
||||||
|
直接输出普通文本,用 emoji 标记不同内容类型:
|
||||||
|
- 🎭 角色名 — 角色描述和相关提示词
|
||||||
|
- 🌄 场景描述 — 场景提示词
|
||||||
|
- 🔊 角色名: "台词" — TTS 语音
|
||||||
|
- 📦 ComfyUI 工作流 — 最终的 workflow JSON(放在 ```json 代码块内)
|
||||||
|
|
||||||
|
## 工作流程
|
||||||
|
1. 用户提供脚本 → 你分析脚本结构
|
||||||
|
2. 针对每个角色 → 生成 4 角度提示词
|
||||||
|
3. 针对每个场景 → 生成场景提示词
|
||||||
|
4. 针对每段对话 → 生成 TTS 语音描述
|
||||||
|
5. 最终导出 ComfyUI 工作流 JSON
|
||||||
|
|
||||||
|
始终用中文回复,文字简洁清晰。"#;
|
||||||
|
|
||||||
|
pub fn routes() -> Router<AppState> {
|
||||||
|
Router::new().route("/api/comfyui-agent/chat", post(chat_handler))
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn chat_handler(
|
||||||
|
State(state): State<AppState>,
|
||||||
|
headers: HeaderMap,
|
||||||
|
Json(body): Json<Value>,
|
||||||
|
) -> Result<Sse<impl Stream<Item = Result<Event, Infallible>>>, (axum::http::StatusCode, Json<Value>)>
|
||||||
|
{
|
||||||
|
let auth_header = headers
|
||||||
|
.get("authorization")
|
||||||
|
.and_then(|value| value.to_str().ok())
|
||||||
|
.unwrap_or_default();
|
||||||
|
|
||||||
|
let token = auth_header
|
||||||
|
.strip_prefix("Bearer ")
|
||||||
|
.or_else(|| auth_header.strip_prefix("bearer "))
|
||||||
|
.unwrap_or_default();
|
||||||
|
|
||||||
|
if token != state.config.auth.auth_token {
|
||||||
|
return Err((
|
||||||
|
axum::http::StatusCode::UNAUTHORIZED,
|
||||||
|
Json(json!({ "error": "Unauthorized" })),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
|
let messages = body
|
||||||
|
.get("messages")
|
||||||
|
.and_then(Value::as_array)
|
||||||
|
.ok_or_else(|| {
|
||||||
|
(
|
||||||
|
axum::http::StatusCode::BAD_REQUEST,
|
||||||
|
Json(json!({ "error": "messages is required" })),
|
||||||
|
)
|
||||||
|
})?;
|
||||||
|
|
||||||
|
for message in messages {
|
||||||
|
if message.get("role").and_then(Value::as_str).is_none()
|
||||||
|
|| message.get("content").and_then(Value::as_str).is_none()
|
||||||
|
{
|
||||||
|
return Err((
|
||||||
|
axum::http::StatusCode::BAD_REQUEST,
|
||||||
|
Json(json!({ "error": "messages is required" })),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let messages = messages.clone();
|
||||||
|
let (tx, rx) = unbounded_channel::<Result<Event, Infallible>>();
|
||||||
|
|
||||||
|
tokio::spawn(async move {
|
||||||
|
let upstream_base = std::env::var("BOOSAPI_UPSTREAM_BASE")
|
||||||
|
.unwrap_or_else(|_| "https://api.boos.lat/v1".to_string());
|
||||||
|
let upstream_key =
|
||||||
|
std::env::var("BOOSAPI_UPSTREAM_KEY").unwrap_or_else(|_| "sk-boos4721".to_string());
|
||||||
|
|
||||||
|
let upstream_url = format!("{}/chat/completions", upstream_base.trim_end_matches('/'));
|
||||||
|
|
||||||
|
let mut upstream_messages = Vec::with_capacity(messages.len() + 1);
|
||||||
|
upstream_messages.push(json!({
|
||||||
|
"role": "system",
|
||||||
|
"content": SYSTEM_PROMPT,
|
||||||
|
}));
|
||||||
|
upstream_messages.extend(messages);
|
||||||
|
|
||||||
|
let client = Client::new();
|
||||||
|
|
||||||
|
info!("comfyui-agent chat streaming request to upstream");
|
||||||
|
|
||||||
|
let response = match client
|
||||||
|
.post(upstream_url)
|
||||||
|
.bearer_auth(upstream_key)
|
||||||
|
.json(&json!({
|
||||||
|
"model": "gpt-5.5",
|
||||||
|
"messages": upstream_messages,
|
||||||
|
"temperature": 0.7,
|
||||||
|
"stream": true,
|
||||||
|
}))
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
Ok(response) => response,
|
||||||
|
Err(_) => {
|
||||||
|
send_sse(
|
||||||
|
&tx,
|
||||||
|
json!({
|
||||||
|
"type": "text",
|
||||||
|
"content": "服务暂时不可用,请稍后重试",
|
||||||
|
}),
|
||||||
|
);
|
||||||
|
send_sse(&tx, json!({ "type": "done" }));
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
if !response.status().is_success() {
|
||||||
|
send_sse(
|
||||||
|
&tx,
|
||||||
|
json!({
|
||||||
|
"type": "text",
|
||||||
|
"content": "服务暂时不可用,请稍后重试",
|
||||||
|
}),
|
||||||
|
);
|
||||||
|
send_sse(&tx, json!({ "type": "done" }));
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut stream = response.bytes_stream();
|
||||||
|
let mut buffer: Vec<u8> = Vec::new();
|
||||||
|
|
||||||
|
loop {
|
||||||
|
let chunk = tokio::select! {
|
||||||
|
_ = tx.closed() => {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
chunk = stream.next() => chunk,
|
||||||
|
};
|
||||||
|
|
||||||
|
let Some(chunk) = chunk else {
|
||||||
|
break;
|
||||||
|
};
|
||||||
|
|
||||||
|
let chunk = match chunk {
|
||||||
|
Ok(chunk) => chunk,
|
||||||
|
Err(error) => {
|
||||||
|
if !send_sse(
|
||||||
|
&tx,
|
||||||
|
json!({
|
||||||
|
"type": "text",
|
||||||
|
"content": format!("\n\n[错误: {}]", error),
|
||||||
|
}),
|
||||||
|
) {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
buffer.extend_from_slice(&chunk);
|
||||||
|
|
||||||
|
while let Some(pos) = buffer.iter().position(|byte| *byte == b'\n') {
|
||||||
|
let mut line_bytes: Vec<u8> = buffer.drain(..=pos).collect();
|
||||||
|
|
||||||
|
if line_bytes.last() == Some(&b'\n') {
|
||||||
|
line_bytes.pop();
|
||||||
|
}
|
||||||
|
if line_bytes.last() == Some(&b'\r') {
|
||||||
|
line_bytes.pop();
|
||||||
|
}
|
||||||
|
|
||||||
|
let line = String::from_utf8_lossy(&line_bytes);
|
||||||
|
let line = line.trim();
|
||||||
|
|
||||||
|
if !line.starts_with("data: ") {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
let data = line.trim_start_matches("data: ").trim();
|
||||||
|
|
||||||
|
if data == "[DONE]" {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
let Ok(parsed) = serde_json::from_str::<Value>(data) else {
|
||||||
|
continue;
|
||||||
|
};
|
||||||
|
|
||||||
|
let content = parsed
|
||||||
|
.get("choices")
|
||||||
|
.and_then(Value::as_array)
|
||||||
|
.and_then(|choices| choices.first())
|
||||||
|
.and_then(|choice| choice.get("delta"))
|
||||||
|
.and_then(|delta| delta.get("content"))
|
||||||
|
.and_then(Value::as_str)
|
||||||
|
.unwrap_or_default();
|
||||||
|
|
||||||
|
if !content.is_empty()
|
||||||
|
&& !send_sse(
|
||||||
|
&tx,
|
||||||
|
json!({
|
||||||
|
"type": "text",
|
||||||
|
"content": content,
|
||||||
|
}),
|
||||||
|
)
|
||||||
|
{
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
send_sse(&tx, json!({ "type": "done" }));
|
||||||
|
});
|
||||||
|
|
||||||
|
Ok(Sse::new(UnboundedReceiverStream::new(rx)))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn send_sse(tx: &UnboundedSender<Result<Event, Infallible>>, payload: Value) -> bool {
|
||||||
|
tx.send(Ok(Event::default().data(payload.to_string())))
|
||||||
|
.is_ok()
|
||||||
|
}
|
||||||
@@ -78,10 +78,7 @@ async fn unread_count(State(state): State<AppState>) -> AppResult<Json<Value>> {
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// POST /api/events/{id}/read
|
/// POST /api/events/{id}/read
|
||||||
async fn mark_read(
|
async fn mark_read(State(state): State<AppState>, Path(id): Path<i64>) -> AppResult<Json<Value>> {
|
||||||
State(state): State<AppState>,
|
|
||||||
Path(id): Path<i64>,
|
|
||||||
) -> AppResult<Json<Value>> {
|
|
||||||
info!("Marking event {} as read", id);
|
info!("Marking event {} as read", id);
|
||||||
sqlx::query("UPDATE events SET read = 1 WHERE id = ?")
|
sqlx::query("UPDATE events SET read = 1 WHERE id = ?")
|
||||||
.bind(id)
|
.bind(id)
|
||||||
|
|||||||
@@ -5,32 +5,42 @@ mod admin;
|
|||||||
mod auth;
|
mod auth;
|
||||||
mod checkin;
|
mod checkin;
|
||||||
mod comfyui;
|
mod comfyui;
|
||||||
|
mod comfyui_agent;
|
||||||
mod events;
|
mod events;
|
||||||
mod health;
|
mod health;
|
||||||
mod monitor;
|
mod monitor;
|
||||||
|
mod oauth;
|
||||||
mod organizations;
|
mod organizations;
|
||||||
mod proxy;
|
mod proxy;
|
||||||
mod search;
|
mod search;
|
||||||
mod site_announcements;
|
mod site_announcements;
|
||||||
mod stats;
|
mod stats;
|
||||||
mod tasks;
|
mod tasks;
|
||||||
|
mod update_center;
|
||||||
|
mod user;
|
||||||
mod user_api_keys;
|
mod user_api_keys;
|
||||||
|
mod video_agent;
|
||||||
|
|
||||||
pub fn build_router(state: AppState) -> Router {
|
pub fn build_router(state: AppState) -> Router {
|
||||||
Router::new()
|
Router::new()
|
||||||
.route("/healthz", get(health::health_check))
|
.route("/healthz", get(health::health_check))
|
||||||
.merge(admin::routes(state.clone()))
|
.merge(admin::routes(state.clone()))
|
||||||
.merge(proxy::routes())
|
.merge(proxy::routes())
|
||||||
|
.merge(comfyui_agent::routes())
|
||||||
.merge(comfyui::routes())
|
.merge(comfyui::routes())
|
||||||
.merge(auth::routes())
|
.merge(auth::routes())
|
||||||
.merge(events::routes())
|
.merge(events::routes())
|
||||||
.merge(checkin::routes())
|
.merge(checkin::routes())
|
||||||
.merge(monitor::routes())
|
.merge(monitor::routes())
|
||||||
|
.merge(oauth::routes())
|
||||||
.merge(organizations::routes())
|
.merge(organizations::routes())
|
||||||
.merge(search::routes())
|
.merge(search::routes())
|
||||||
.merge(site_announcements::routes())
|
.merge(site_announcements::routes())
|
||||||
.merge(stats::routes())
|
.merge(stats::routes())
|
||||||
.merge(tasks::routes())
|
.merge(tasks::routes())
|
||||||
|
.merge(update_center::routes())
|
||||||
|
.merge(user::routes())
|
||||||
.merge(user_api_keys::routes())
|
.merge(user_api_keys::routes())
|
||||||
|
.merge(video_agent::routes())
|
||||||
.with_state(state)
|
.with_state(state)
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -20,7 +20,10 @@ const LDOH_COOKIE_SETTING_KEY: &str = "monitor_ldoh_cookie";
|
|||||||
|
|
||||||
pub fn routes() -> Router<AppState> {
|
pub fn routes() -> Router<AppState> {
|
||||||
Router::new()
|
Router::new()
|
||||||
.route("/api/monitor/config", get(get_monitor_config).put(update_monitor_config))
|
.route(
|
||||||
|
"/api/monitor/config",
|
||||||
|
get(get_monitor_config).put(update_monitor_config),
|
||||||
|
)
|
||||||
.route("/api/monitor/session", post(create_session))
|
.route("/api/monitor/session", post(create_session))
|
||||||
.route("/monitor-proxy/ldoh", get(ldoh_proxy_handler))
|
.route("/monitor-proxy/ldoh", get(ldoh_proxy_handler))
|
||||||
.route("/monitor-proxy/ldoh/", get(ldoh_proxy_handler))
|
.route("/monitor-proxy/ldoh/", get(ldoh_proxy_handler))
|
||||||
@@ -28,23 +31,26 @@ pub fn routes() -> Router<AppState> {
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// GET /api/monitor/config — returns whether LDOH cookie is configured
|
/// GET /api/monitor/config — returns whether LDOH cookie is configured
|
||||||
async fn get_monitor_config(
|
async fn get_monitor_config(State(state): State<AppState>) -> AppResult<Json<Value>> {
|
||||||
State(state): State<AppState>,
|
let cookie: Option<(String,)> = sqlx::query_as("SELECT value FROM settings WHERE key = ?")
|
||||||
) -> AppResult<Json<Value>> {
|
|
||||||
let cookie: Option<(String,)> = sqlx::query_as(
|
|
||||||
"SELECT value FROM settings WHERE key = ?"
|
|
||||||
)
|
|
||||||
.bind(LDOH_COOKIE_SETTING_KEY)
|
.bind(LDOH_COOKIE_SETTING_KEY)
|
||||||
.fetch_optional(&*state.db)
|
.fetch_optional(&*state.db)
|
||||||
.await?;
|
.await?;
|
||||||
|
|
||||||
let raw = cookie.and_then(|(v,)| {
|
let raw = cookie
|
||||||
serde_json::from_str::<serde_json::Value>(&v).ok()
|
.and_then(|(v,)| {
|
||||||
|
serde_json::from_str::<serde_json::Value>(&v)
|
||||||
|
.ok()
|
||||||
.and_then(|parsed| parsed.as_str().map(String::from))
|
.and_then(|parsed| parsed.as_str().map(String::from))
|
||||||
}).unwrap_or_default();
|
})
|
||||||
|
.unwrap_or_default();
|
||||||
|
|
||||||
let configured = !raw.is_empty();
|
let configured = !raw.is_empty();
|
||||||
let masked = if configured { mask_cookie_value(&raw) } else { String::new() };
|
let masked = if configured {
|
||||||
|
mask_cookie_value(&raw)
|
||||||
|
} else {
|
||||||
|
String::new()
|
||||||
|
};
|
||||||
|
|
||||||
Ok(Json(json!({
|
Ok(Json(json!({
|
||||||
"ldohCookieConfigured": configured,
|
"ldohCookieConfigured": configured,
|
||||||
@@ -57,7 +63,10 @@ async fn update_monitor_config(
|
|||||||
State(state): State<AppState>,
|
State(state): State<AppState>,
|
||||||
Json(body): Json<serde_json::Value>,
|
Json(body): Json<serde_json::Value>,
|
||||||
) -> AppResult<Json<Value>> {
|
) -> AppResult<Json<Value>> {
|
||||||
let cookie = body.get("ldohCookie").and_then(|v| v.as_str()).unwrap_or("");
|
let cookie = body
|
||||||
|
.get("ldohCookie")
|
||||||
|
.and_then(|v| v.as_str())
|
||||||
|
.unwrap_or("");
|
||||||
|
|
||||||
if cookie.is_empty() {
|
if cookie.is_empty() {
|
||||||
sqlx::query("INSERT OR REPLACE INTO settings (key, value) VALUES (?, '')")
|
sqlx::query("INSERT OR REPLACE INTO settings (key, value) VALUES (?, '')")
|
||||||
@@ -92,9 +101,7 @@ async fn update_monitor_config(
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// POST /api/monitor/session — returns cookie value for client to set
|
/// POST /api/monitor/session — returns cookie value for client to set
|
||||||
async fn create_session(
|
async fn create_session(State(state): State<AppState>) -> Json<Value> {
|
||||||
State(state): State<AppState>,
|
|
||||||
) -> Json<Value> {
|
|
||||||
let token = &state.config.auth.auth_token;
|
let token = &state.config.auth.auth_token;
|
||||||
Json(json!({
|
Json(json!({
|
||||||
"success": true,
|
"success": true,
|
||||||
@@ -110,24 +117,30 @@ async fn ldoh_proxy_handler(
|
|||||||
) -> Result<Response, AppError> {
|
) -> Result<Response, AppError> {
|
||||||
// Auth check via cookie
|
// Auth check via cookie
|
||||||
let cookies = parse_cookies(
|
let cookies = parse_cookies(
|
||||||
headers.get(header::COOKIE).and_then(|v| v.to_str().ok()).unwrap_or("")
|
headers
|
||||||
|
.get(header::COOKIE)
|
||||||
|
.and_then(|v| v.to_str().ok())
|
||||||
|
.unwrap_or(""),
|
||||||
);
|
);
|
||||||
if cookies.get(MONITOR_AUTH_COOKIE).map(|s| s.as_str()) != Some(&state.config.auth.auth_token) {
|
if cookies.get(MONITOR_AUTH_COOKIE).map(|s| s.as_str()) != Some(&state.config.auth.auth_token) {
|
||||||
return Err(AppError::Unauthorized("Missing or invalid monitor session".into()));
|
return Err(AppError::Unauthorized(
|
||||||
|
"Missing or invalid monitor session".into(),
|
||||||
|
));
|
||||||
}
|
}
|
||||||
|
|
||||||
// Get stored LDOH cookie
|
// Get stored LDOH cookie
|
||||||
let stored: Option<(String,)> = sqlx::query_as(
|
let stored: Option<(String,)> = sqlx::query_as("SELECT value FROM settings WHERE key = ?")
|
||||||
"SELECT value FROM settings WHERE key = ?"
|
|
||||||
)
|
|
||||||
.bind(LDOH_COOKIE_SETTING_KEY)
|
.bind(LDOH_COOKIE_SETTING_KEY)
|
||||||
.fetch_optional(&*state.db)
|
.fetch_optional(&*state.db)
|
||||||
.await?;
|
.await?;
|
||||||
|
|
||||||
let ldoh_cookie = stored.and_then(|(v,)| {
|
let ldoh_cookie = stored
|
||||||
serde_json::from_str::<serde_json::Value>(&v).ok()
|
.and_then(|(v,)| {
|
||||||
|
serde_json::from_str::<serde_json::Value>(&v)
|
||||||
|
.ok()
|
||||||
.and_then(|parsed| parsed.as_str().map(String::from))
|
.and_then(|parsed| parsed.as_str().map(String::from))
|
||||||
}).unwrap_or_default();
|
})
|
||||||
|
.unwrap_or_default();
|
||||||
|
|
||||||
if ldoh_cookie.is_empty() {
|
if ldoh_cookie.is_empty() {
|
||||||
return Err(AppError::BadRequest("LDOH cookie not configured".into()));
|
return Err(AppError::BadRequest("LDOH cookie not configured".into()));
|
||||||
@@ -146,28 +159,54 @@ async fn ldoh_proxy_handler(
|
|||||||
.build()
|
.build()
|
||||||
.map_err(|e| AppError::Internal(anyhow::anyhow!("Failed to build client: {e}")))?;
|
.map_err(|e| AppError::Internal(anyhow::anyhow!("Failed to build client: {e}")))?;
|
||||||
|
|
||||||
let mut req = client.get(&target_url)
|
let mut req = client
|
||||||
|
.get(&target_url)
|
||||||
.header("cookie", &ldoh_cookie)
|
.header("cookie", &ldoh_cookie)
|
||||||
.header("accept", headers.get("accept").and_then(|v| v.to_str().ok()).unwrap_or("*/*"))
|
.header(
|
||||||
.header("accept-language", headers.get("accept-language").and_then(|v| v.to_str().ok()).unwrap_or("zh-CN,zh;q=0.9,en;q=0.8"))
|
"accept",
|
||||||
.header("user-agent", headers.get("user-agent").and_then(|v| v.to_str().ok()).unwrap_or("boosapiMonitorProxy/1.0"));
|
headers
|
||||||
|
.get("accept")
|
||||||
|
.and_then(|v| v.to_str().ok())
|
||||||
|
.unwrap_or("*/*"),
|
||||||
|
)
|
||||||
|
.header(
|
||||||
|
"accept-language",
|
||||||
|
headers
|
||||||
|
.get("accept-language")
|
||||||
|
.and_then(|v| v.to_str().ok())
|
||||||
|
.unwrap_or("zh-CN,zh;q=0.9,en;q=0.8"),
|
||||||
|
)
|
||||||
|
.header(
|
||||||
|
"user-agent",
|
||||||
|
headers
|
||||||
|
.get("user-agent")
|
||||||
|
.and_then(|v| v.to_str().ok())
|
||||||
|
.unwrap_or("boosapiMonitorProxy/1.0"),
|
||||||
|
);
|
||||||
|
|
||||||
if let Some(referer) = headers.get("referer").and_then(|v| v.to_str().ok()) {
|
if let Some(referer) = headers.get("referer").and_then(|v| v.to_str().ok()) {
|
||||||
req = req.header("referer", referer.replace("/monitor-proxy/ldoh", ""));
|
req = req.header("referer", referer.replace("/monitor-proxy/ldoh", ""));
|
||||||
}
|
}
|
||||||
|
|
||||||
let resp = req.send().await
|
let resp = req
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
.map_err(|e| AppError::UpstreamError(format!("LDOH unreachable: {e}")))?;
|
.map_err(|e| AppError::UpstreamError(format!("LDOH unreachable: {e}")))?;
|
||||||
|
|
||||||
let status = resp.status();
|
let status = resp.status();
|
||||||
let content_type = resp.headers()
|
let content_type = resp
|
||||||
|
.headers()
|
||||||
.get("content-type")
|
.get("content-type")
|
||||||
.and_then(|v| v.to_str().ok())
|
.and_then(|v| v.to_str().ok())
|
||||||
.unwrap_or("application/octet-stream")
|
.unwrap_or("application/octet-stream")
|
||||||
.to_string();
|
.to_string();
|
||||||
|
|
||||||
let mut builder = axum::response::Response::builder().status(status);
|
let mut builder = axum::response::Response::builder().status(status);
|
||||||
if let Some(ctl) = resp.headers().get("cache-control").and_then(|v| v.to_str().ok()) {
|
if let Some(ctl) = resp
|
||||||
|
.headers()
|
||||||
|
.get("cache-control")
|
||||||
|
.and_then(|v| v.to_str().ok())
|
||||||
|
{
|
||||||
builder = builder.header("cache-control", ctl);
|
builder = builder.header("cache-control", ctl);
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -177,14 +216,20 @@ async fn ldoh_proxy_handler(
|
|||||||
|| content_type.contains("text/css")
|
|| content_type.contains("text/css")
|
||||||
|| content_type.contains("application/json")
|
|| content_type.contains("application/json")
|
||||||
{
|
{
|
||||||
let text = resp.text().await.map_err(|e| AppError::UpstreamError(e.to_string()))?;
|
let text = resp
|
||||||
|
.text()
|
||||||
|
.await
|
||||||
|
.map_err(|e| AppError::UpstreamError(e.to_string()))?;
|
||||||
return Ok(builder
|
return Ok(builder
|
||||||
.header(header::CONTENT_TYPE, &content_type)
|
.header(header::CONTENT_TYPE, &content_type)
|
||||||
.body(axum::body::Body::from(rewrite_proxy_text(&text)))
|
.body(axum::body::Body::from(rewrite_proxy_text(&text)))
|
||||||
.map_err(|e| AppError::Internal(anyhow::anyhow!("{e}")))?);
|
.map_err(|e| AppError::Internal(anyhow::anyhow!("{e}")))?);
|
||||||
}
|
}
|
||||||
|
|
||||||
let bytes = resp.bytes().await.map_err(|e| AppError::UpstreamError(e.to_string()))?;
|
let bytes = resp
|
||||||
|
.bytes()
|
||||||
|
.await
|
||||||
|
.map_err(|e| AppError::UpstreamError(e.to_string()))?;
|
||||||
Ok(builder
|
Ok(builder
|
||||||
.header(header::CONTENT_TYPE, &content_type)
|
.header(header::CONTENT_TYPE, &content_type)
|
||||||
.body(axum::body::Body::from(bytes))
|
.body(axum::body::Body::from(bytes))
|
||||||
@@ -200,7 +245,9 @@ fn parse_cookies(raw: &str) -> std::collections::HashMap<String, String> {
|
|||||||
if let Some(idx) = part.find('=') {
|
if let Some(idx) = part.find('=') {
|
||||||
let key = part[..idx].trim().to_string();
|
let key = part[..idx].trim().to_string();
|
||||||
let val = part[idx + 1..].trim().to_string();
|
let val = part[idx + 1..].trim().to_string();
|
||||||
if !key.is_empty() { map.insert(key, val); }
|
if !key.is_empty() {
|
||||||
|
map.insert(key, val);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
map
|
map
|
||||||
@@ -208,12 +255,14 @@ fn parse_cookies(raw: &str) -> std::collections::HashMap<String, String> {
|
|||||||
|
|
||||||
fn mask_cookie_value(cookie_text: &str) -> String {
|
fn mask_cookie_value(cookie_text: &str) -> String {
|
||||||
let raw = cookie_text.trim();
|
let raw = cookie_text.trim();
|
||||||
if raw.is_empty() { return String::new(); }
|
if raw.is_empty() {
|
||||||
|
return String::new();
|
||||||
|
}
|
||||||
let value = raw.split('=').nth(1).unwrap_or(raw);
|
let value = raw.split('=').nth(1).unwrap_or(raw);
|
||||||
if value.len() <= 10 {
|
if value.len() <= 10 {
|
||||||
format!("{}****", &value[..2.min(value.len())])
|
format!("{}****", &value[..2.min(value.len())])
|
||||||
} else {
|
} else {
|
||||||
format!("{}****{}", &value[..6], &value[value.len()-4..])
|
format!("{}****{}", &value[..6], &value[value.len() - 4..])
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -230,15 +279,17 @@ fn normalize_ldoh_cookie(raw: &str) -> String {
|
|||||||
}
|
}
|
||||||
|
|
||||||
fn rewrite_proxy_text(text: &str) -> String {
|
fn rewrite_proxy_text(text: &str) -> String {
|
||||||
text
|
text.replace("https://ldoh.105117.xyz/", "/monitor-proxy/ldoh/")
|
||||||
.replace("https://ldoh.105117.xyz/", "/monitor-proxy/ldoh/")
|
.replace(
|
||||||
.replace("https:\\/\\/ldoh.105117.xyz\\/", "\\/monitor-proxy\\/ldoh\\/")
|
"https:\\/\\/ldoh.105117.xyz\\/",
|
||||||
|
"\\/monitor-proxy\\/ldoh\\/",
|
||||||
|
)
|
||||||
.replace("src=\"/", "src=\"/monitor-proxy/ldoh/")
|
.replace("src=\"/", "src=\"/monitor-proxy/ldoh/")
|
||||||
.replace("src='/" , "src='/monitor-proxy/ldoh/")
|
.replace("src='/", "src='/monitor-proxy/ldoh/")
|
||||||
.replace("href=\"/", "href=\"/monitor-proxy/ldoh/")
|
.replace("href=\"/", "href=\"/monitor-proxy/ldoh/")
|
||||||
.replace("href='/" , "href='/monitor-proxy/ldoh/")
|
.replace("href='/", "href='/monitor-proxy/ldoh/")
|
||||||
.replace("action=\"/", "action=\"/monitor-proxy/ldoh/")
|
.replace("action=\"/", "action=\"/monitor-proxy/ldoh/")
|
||||||
.replace("action='/" , "action='/monitor-proxy/ldoh/")
|
.replace("action='/", "action='/monitor-proxy/ldoh/")
|
||||||
.replace("\"/_next/", "\"/monitor-proxy/ldoh/_next/")
|
.replace("\"/_next/", "\"/monitor-proxy/ldoh/_next/")
|
||||||
.replace("'/_next/", "'/monitor-proxy/ldoh/_next/")
|
.replace("'/_next/", "'/monitor-proxy/ldoh/_next/")
|
||||||
.replace("\"\\\\/api/", "\"\\\\/monitor-proxy\\\\/ldoh\\\\/api/")
|
.replace("\"\\\\/api/", "\"\\\\/monitor-proxy\\\\/ldoh\\\\/api/")
|
||||||
|
|||||||
@@ -0,0 +1,312 @@
|
|||||||
|
// OAuth routes — provider management, session handling, connection management
|
||||||
|
// Simplified implementation — complex OAuth flow operations return 501
|
||||||
|
|
||||||
|
use crate::{
|
||||||
|
error::{AppError, AppResult},
|
||||||
|
AppState,
|
||||||
|
};
|
||||||
|
use axum::{
|
||||||
|
extract::{Path, State},
|
||||||
|
http::{HeaderMap, StatusCode},
|
||||||
|
response::{Html, IntoResponse},
|
||||||
|
routing::{get, post, put},
|
||||||
|
Json, Router,
|
||||||
|
};
|
||||||
|
use serde_json::{json, Value};
|
||||||
|
use tracing::info;
|
||||||
|
|
||||||
|
pub fn routes() -> Router<AppState> {
|
||||||
|
Router::new()
|
||||||
|
// Providers
|
||||||
|
.route("/api/oauth/providers", get(list_providers))
|
||||||
|
.route(
|
||||||
|
"/api/oauth/providers/{provider}/start",
|
||||||
|
post(start_provider_flow),
|
||||||
|
)
|
||||||
|
// Sessions
|
||||||
|
.route("/api/oauth/sessions/{state}", get(get_session))
|
||||||
|
.route(
|
||||||
|
"/api/oauth/sessions/{state}/manual-callback",
|
||||||
|
post(manual_callback),
|
||||||
|
)
|
||||||
|
// Connections
|
||||||
|
.route("/api/oauth/connections", get(list_connections))
|
||||||
|
.route(
|
||||||
|
"/api/oauth/connections/{id}/rebind",
|
||||||
|
post(rebind_connection),
|
||||||
|
)
|
||||||
|
.route(
|
||||||
|
"/api/oauth/connections/{id}/proxy",
|
||||||
|
post(update_connection_proxy),
|
||||||
|
)
|
||||||
|
.route("/api/oauth/connections/import", post(import_connections))
|
||||||
|
.route(
|
||||||
|
"/api/oauth/connections/{id}/refresh-quota",
|
||||||
|
post(refresh_connection_quota),
|
||||||
|
)
|
||||||
|
.route(
|
||||||
|
"/api/oauth/connections/refresh-quota-batch",
|
||||||
|
post(refresh_quota_batch),
|
||||||
|
)
|
||||||
|
// Route units
|
||||||
|
.route("/api/oauth/route-units", post(create_route_unit))
|
||||||
|
.route(
|
||||||
|
"/api/oauth/route-units/{id}",
|
||||||
|
put(update_route_unit).delete(delete_route_unit),
|
||||||
|
)
|
||||||
|
// Callback
|
||||||
|
.route("/api/oauth/callback", get(oauth_callback_page))
|
||||||
|
// Delete connection
|
||||||
|
.route("/api/oauth/connections/{id}", post(delete_connection))
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Auth helper ────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
fn require_admin(headers: &HeaderMap, state: &AppState) -> AppResult<()> {
|
||||||
|
let token = headers
|
||||||
|
.get("Authorization")
|
||||||
|
.and_then(|v| v.to_str().ok())
|
||||||
|
.and_then(|v| v.strip_prefix("Bearer "))
|
||||||
|
.ok_or_else(|| AppError::Unauthorized("Missing Authorization header".into()))?;
|
||||||
|
|
||||||
|
if token != state.config.auth.auth_token {
|
||||||
|
return Err(AppError::Unauthorized("Invalid admin token".into()));
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Providers ──────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
/// GET /api/oauth/providers — public
|
||||||
|
async fn list_providers() -> Json<Value> {
|
||||||
|
Json(json!({
|
||||||
|
"defaults": {
|
||||||
|
"proxyRequired": false,
|
||||||
|
"useSystemProxy": false,
|
||||||
|
},
|
||||||
|
"providers": [
|
||||||
|
{
|
||||||
|
"id": "openai",
|
||||||
|
"name": "OpenAI",
|
||||||
|
"type": "openai",
|
||||||
|
"enabled": true,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"id": "anthropic",
|
||||||
|
"name": "Anthropic",
|
||||||
|
"type": "anthropic",
|
||||||
|
"enabled": true,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"id": "google",
|
||||||
|
"name": "Google AI",
|
||||||
|
"type": "google",
|
||||||
|
"enabled": true,
|
||||||
|
},
|
||||||
|
],
|
||||||
|
}))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// POST /api/oauth/providers/{provider}/start
|
||||||
|
async fn start_provider_flow(
|
||||||
|
State(_state): State<AppState>,
|
||||||
|
headers: HeaderMap,
|
||||||
|
Path(_provider): Path<String>,
|
||||||
|
) -> AppResult<Json<Value>> {
|
||||||
|
require_admin(&headers, &_state)?;
|
||||||
|
info!("OAuth provider flow start: {}", _provider);
|
||||||
|
Err(AppError::BadRequest(
|
||||||
|
"OAuth provider flow not implemented in Rust backend".into(),
|
||||||
|
))
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Sessions ───────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
/// GET /api/oauth/sessions/{state}
|
||||||
|
async fn get_session(
|
||||||
|
State(_state): State<AppState>,
|
||||||
|
headers: HeaderMap,
|
||||||
|
Path(state): Path<String>,
|
||||||
|
) -> AppResult<Json<Value>> {
|
||||||
|
require_admin(&headers, &_state)?;
|
||||||
|
info!("OAuth session lookup: {}", state);
|
||||||
|
Err(AppError::NotFound(format!(
|
||||||
|
"OAuth session not found: {}",
|
||||||
|
state
|
||||||
|
)))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// POST /api/oauth/sessions/{state}/manual-callback
|
||||||
|
async fn manual_callback(
|
||||||
|
State(_state): State<AppState>,
|
||||||
|
headers: HeaderMap,
|
||||||
|
Path(state): Path<String>,
|
||||||
|
) -> AppResult<Json<Value>> {
|
||||||
|
require_admin(&headers, &_state)?;
|
||||||
|
info!("OAuth manual callback: {}", state);
|
||||||
|
Err(AppError::BadRequest(
|
||||||
|
"OAuth manual callback not implemented in Rust backend".into(),
|
||||||
|
))
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Connections ────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
/// GET /api/oauth/connections
|
||||||
|
async fn list_connections(
|
||||||
|
State(state): State<AppState>,
|
||||||
|
headers: HeaderMap,
|
||||||
|
) -> AppResult<Json<Value>> {
|
||||||
|
require_admin(&headers, &state)?;
|
||||||
|
info!("Listing OAuth connections");
|
||||||
|
|
||||||
|
let accounts: Vec<Value> = sqlx::query_as::<_, (i64, i64, Option<String>, String, Option<String>, Option<String>, Option<String>, String)>(
|
||||||
|
"SELECT id, site_id, username, access_token, oauth_provider, oauth_account_key, oauth_project_id, status \
|
||||||
|
FROM accounts WHERE oauth_provider IS NOT NULL ORDER BY id"
|
||||||
|
)
|
||||||
|
.fetch_all(&*state.db)
|
||||||
|
.await?
|
||||||
|
.into_iter()
|
||||||
|
.map(|(id, site_id, username, _access_token, oauth_provider, oauth_account_key, oauth_project_id, status)| {
|
||||||
|
json!({
|
||||||
|
"id": id,
|
||||||
|
"siteId": site_id,
|
||||||
|
"username": username,
|
||||||
|
"oauthProvider": oauth_provider,
|
||||||
|
"oauthAccountKey": oauth_account_key,
|
||||||
|
"oauthProjectId": oauth_project_id,
|
||||||
|
"status": status,
|
||||||
|
})
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
Ok(Json(json!({ "data": accounts, "total": accounts.len() })))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// POST /api/oauth/connections/{id}/rebind
|
||||||
|
async fn rebind_connection(
|
||||||
|
State(_state): State<AppState>,
|
||||||
|
headers: HeaderMap,
|
||||||
|
Path(_id): Path<i64>,
|
||||||
|
) -> AppResult<Json<Value>> {
|
||||||
|
require_admin(&headers, &_state)?;
|
||||||
|
Err(AppError::BadRequest(
|
||||||
|
"OAuth rebind not implemented in Rust backend".into(),
|
||||||
|
))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// POST /api/oauth/connections/{id}/proxy
|
||||||
|
async fn update_connection_proxy(
|
||||||
|
State(_state): State<AppState>,
|
||||||
|
headers: HeaderMap,
|
||||||
|
Path(_id): Path<i64>,
|
||||||
|
) -> AppResult<Json<Value>> {
|
||||||
|
require_admin(&headers, &_state)?;
|
||||||
|
Err(AppError::BadRequest(
|
||||||
|
"OAuth proxy update not implemented in Rust backend".into(),
|
||||||
|
))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// POST /api/oauth/connections/import
|
||||||
|
async fn import_connections(
|
||||||
|
State(_state): State<AppState>,
|
||||||
|
headers: HeaderMap,
|
||||||
|
) -> AppResult<Json<Value>> {
|
||||||
|
require_admin(&headers, &_state)?;
|
||||||
|
Err(AppError::BadRequest(
|
||||||
|
"OAuth import not implemented in Rust backend".into(),
|
||||||
|
))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// POST /api/oauth/connections/{id}/refresh-quota
|
||||||
|
async fn refresh_connection_quota(
|
||||||
|
State(_state): State<AppState>,
|
||||||
|
headers: HeaderMap,
|
||||||
|
Path(_id): Path<i64>,
|
||||||
|
) -> AppResult<Json<Value>> {
|
||||||
|
require_admin(&headers, &_state)?;
|
||||||
|
Err(AppError::BadRequest(
|
||||||
|
"OAuth quota refresh not implemented in Rust backend".into(),
|
||||||
|
))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// POST /api/oauth/connections/refresh-quota-batch
|
||||||
|
async fn refresh_quota_batch(
|
||||||
|
State(_state): State<AppState>,
|
||||||
|
headers: HeaderMap,
|
||||||
|
) -> AppResult<Json<Value>> {
|
||||||
|
require_admin(&headers, &_state)?;
|
||||||
|
Err(AppError::BadRequest(
|
||||||
|
"OAuth batch quota refresh not implemented in Rust backend".into(),
|
||||||
|
))
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Route Units ────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
/// POST /api/oauth/route-units
|
||||||
|
async fn create_route_unit(
|
||||||
|
State(_state): State<AppState>,
|
||||||
|
headers: HeaderMap,
|
||||||
|
) -> AppResult<Json<Value>> {
|
||||||
|
require_admin(&headers, &_state)?;
|
||||||
|
Err(AppError::BadRequest(
|
||||||
|
"OAuth route unit creation not implemented in Rust backend".into(),
|
||||||
|
))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// PUT /api/oauth/route-units/{id}
|
||||||
|
async fn update_route_unit(
|
||||||
|
State(_state): State<AppState>,
|
||||||
|
headers: HeaderMap,
|
||||||
|
Path(_id): Path<i64>,
|
||||||
|
) -> AppResult<Json<Value>> {
|
||||||
|
require_admin(&headers, &_state)?;
|
||||||
|
Err(AppError::BadRequest(
|
||||||
|
"OAuth route unit update not implemented in Rust backend".into(),
|
||||||
|
))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// DELETE /api/oauth/route-units/{id}
|
||||||
|
async fn delete_route_unit(
|
||||||
|
State(_state): State<AppState>,
|
||||||
|
headers: HeaderMap,
|
||||||
|
Path(_id): Path<i64>,
|
||||||
|
) -> AppResult<Json<Value>> {
|
||||||
|
require_admin(&headers, &_state)?;
|
||||||
|
Err(AppError::BadRequest(
|
||||||
|
"OAuth route unit deletion not implemented in Rust backend".into(),
|
||||||
|
))
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Callback ───────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
/// GET /api/oauth/callback — public, returns HTML for popup window
|
||||||
|
async fn oauth_callback_page(headers: HeaderMap) -> impl IntoResponse {
|
||||||
|
info!("OAuth callback received");
|
||||||
|
let html = r#"<!doctype html>
|
||||||
|
<html lang="zh-CN">
|
||||||
|
<head><meta charset="utf-8" /><title>OAuth Callback</title></head>
|
||||||
|
<body>
|
||||||
|
<script>window.close();</script>
|
||||||
|
OAuth callback received. You can close this window.
|
||||||
|
</body>
|
||||||
|
</html>"#;
|
||||||
|
(
|
||||||
|
StatusCode::OK,
|
||||||
|
[("Content-Type", "text/html; charset=utf-8")],
|
||||||
|
Html(html),
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Delete ─────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
/// POST /api/oauth/connections/{id} — delete connection
|
||||||
|
async fn delete_connection(
|
||||||
|
State(_state): State<AppState>,
|
||||||
|
headers: HeaderMap,
|
||||||
|
Path(_id): Path<i64>,
|
||||||
|
) -> AppResult<Json<Value>> {
|
||||||
|
require_admin(&headers, &_state)?;
|
||||||
|
Err(AppError::BadRequest(
|
||||||
|
"OAuth connection deletion not implemented in Rust backend".into(),
|
||||||
|
))
|
||||||
|
}
|
||||||
@@ -37,30 +37,26 @@ fn extract_jwt_user(headers: &HeaderMap, state: &AppState) -> AppResult<(i64, St
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// Check if user is org admin (owner or admin role).
|
/// Check if user is org admin (owner or admin role).
|
||||||
async fn require_org_admin(
|
async fn require_org_admin(db: &sqlx::SqlitePool, org_id: i64, user_id: i64) -> AppResult<()> {
|
||||||
db: &sqlx::SqlitePool,
|
let member: Option<OrganizationMember> =
|
||||||
org_id: i64,
|
sqlx::query_as("SELECT * FROM organization_members WHERE org_id = ? AND user_id = ?")
|
||||||
user_id: i64,
|
|
||||||
) -> AppResult<()> {
|
|
||||||
let member: Option<OrganizationMember> = sqlx::query_as(
|
|
||||||
"SELECT * FROM organization_members WHERE org_id = ? AND user_id = ?"
|
|
||||||
)
|
|
||||||
.bind(org_id)
|
.bind(org_id)
|
||||||
.bind(user_id)
|
.bind(user_id)
|
||||||
.fetch_optional(db)
|
.fetch_optional(db)
|
||||||
.await?;
|
.await?;
|
||||||
|
|
||||||
match member {
|
match member {
|
||||||
Some(m) if m.role.as_deref() == Some("owner") || m.role.as_deref() == Some("admin") => Ok(()),
|
Some(m) if m.role.as_deref() == Some("owner") || m.role.as_deref() == Some("admin") => {
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
_ => Err(AppError::Forbidden("Admin access required".into())),
|
_ => Err(AppError::Forbidden("Admin access required".into())),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Check if user is a member of an org.
|
/// Check if user is a member of an org.
|
||||||
async fn is_org_member(db: &sqlx::SqlitePool, org_id: i64, user_id: i64) -> Result<bool, AppError> {
|
async fn is_org_member(db: &sqlx::SqlitePool, org_id: i64, user_id: i64) -> Result<bool, AppError> {
|
||||||
let member: Option<OrganizationMember> = sqlx::query_as(
|
let member: Option<OrganizationMember> =
|
||||||
"SELECT * FROM organization_members WHERE org_id = ? AND user_id = ?"
|
sqlx::query_as("SELECT * FROM organization_members WHERE org_id = ? AND user_id = ?")
|
||||||
)
|
|
||||||
.bind(org_id)
|
.bind(org_id)
|
||||||
.bind(user_id)
|
.bind(user_id)
|
||||||
.fetch_optional(db)
|
.fetch_optional(db)
|
||||||
@@ -78,15 +74,30 @@ pub struct ListOrgsQuery {
|
|||||||
pub fn routes() -> Router<AppState> {
|
pub fn routes() -> Router<AppState> {
|
||||||
Router::new()
|
Router::new()
|
||||||
// Organizations CRUD
|
// Organizations CRUD
|
||||||
.route("/api/organizations", get(list_organizations).post(create_organization))
|
.route(
|
||||||
|
"/api/organizations",
|
||||||
|
get(list_organizations).post(create_organization),
|
||||||
|
)
|
||||||
.route("/api/organizations/plans", get(list_plans))
|
.route("/api/organizations/plans", get(list_plans))
|
||||||
.route("/api/organizations/{org_id}", get(get_organization).patch(update_organization))
|
.route(
|
||||||
|
"/api/organizations/{org_id}",
|
||||||
|
get(get_organization).patch(update_organization),
|
||||||
|
)
|
||||||
// Members
|
// Members
|
||||||
.route("/api/organizations/{org_id}/members", get(list_members))
|
.route("/api/organizations/{org_id}/members", get(list_members))
|
||||||
.route("/api/organizations/{org_id}/members/{user_id}", patch(update_member_role).delete(remove_member))
|
.route(
|
||||||
|
"/api/organizations/{org_id}/members/{user_id}",
|
||||||
|
patch(update_member_role).delete(remove_member),
|
||||||
|
)
|
||||||
// Invitations
|
// Invitations
|
||||||
.route("/api/organizations/{org_id}/invitations", get(list_invitations).post(create_invitation))
|
.route(
|
||||||
.route("/api/organizations/{org_id}/invitations/{inv_id}", delete(cancel_invitation))
|
"/api/organizations/{org_id}/invitations",
|
||||||
|
get(list_invitations).post(create_invitation),
|
||||||
|
)
|
||||||
|
.route(
|
||||||
|
"/api/organizations/{org_id}/invitations/{inv_id}",
|
||||||
|
delete(cancel_invitation),
|
||||||
|
)
|
||||||
// Public invitation routes (no org prefix)
|
// Public invitation routes (no org prefix)
|
||||||
.route("/api/invitations/{token}", get(get_invitation))
|
.route("/api/invitations/{token}", get(get_invitation))
|
||||||
.route("/api/invitations/{token}/accept", post(accept_invitation))
|
.route("/api/invitations/{token}/accept", post(accept_invitation))
|
||||||
@@ -108,7 +119,7 @@ async fn list_organizations(
|
|||||||
INNER JOIN organization_members om ON om.org_id = o.id
|
INNER JOIN organization_members om ON om.org_id = o.id
|
||||||
WHERE om.user_id = ?
|
WHERE om.user_id = ?
|
||||||
ORDER BY o.created_at DESC
|
ORDER BY o.created_at DESC
|
||||||
LIMIT ?"#
|
LIMIT ?"#,
|
||||||
)
|
)
|
||||||
.bind(user_id)
|
.bind(user_id)
|
||||||
.bind(limit)
|
.bind(limit)
|
||||||
@@ -126,14 +137,16 @@ async fn create_organization(
|
|||||||
) -> AppResult<(StatusCode, Json<Value>)> {
|
) -> AppResult<(StatusCode, Json<Value>)> {
|
||||||
let (user_id, _) = extract_jwt_user(&headers, &state)?;
|
let (user_id, _) = extract_jwt_user(&headers, &state)?;
|
||||||
|
|
||||||
let name = body.get("name").and_then(|v| v.as_str())
|
let name = body
|
||||||
|
.get("name")
|
||||||
|
.and_then(|v| v.as_str())
|
||||||
.ok_or_else(|| AppError::BadRequest("name is required".into()))?;
|
.ok_or_else(|| AppError::BadRequest("name is required".into()))?;
|
||||||
let slug = body.get("slug").and_then(|v| v.as_str())
|
let slug = body
|
||||||
|
.get("slug")
|
||||||
|
.and_then(|v| v.as_str())
|
||||||
.ok_or_else(|| AppError::BadRequest("slug is required".into()))?;
|
.ok_or_else(|| AppError::BadRequest("slug is required".into()))?;
|
||||||
|
|
||||||
let result = sqlx::query(
|
let result = sqlx::query("INSERT INTO organizations (name, slug) VALUES (?, ?)")
|
||||||
"INSERT INTO organizations (name, slug) VALUES (?, ?)"
|
|
||||||
)
|
|
||||||
.bind(name)
|
.bind(name)
|
||||||
.bind(slug)
|
.bind(slug)
|
||||||
.execute(&*state.db)
|
.execute(&*state.db)
|
||||||
@@ -150,9 +163,7 @@ async fn create_organization(
|
|||||||
let org_id = result.last_insert_rowid();
|
let org_id = result.last_insert_rowid();
|
||||||
|
|
||||||
// Add creator as owner
|
// Add creator as owner
|
||||||
sqlx::query(
|
sqlx::query("INSERT INTO organization_members (org_id, user_id, role) VALUES (?, ?, 'owner')")
|
||||||
"INSERT INTO organization_members (org_id, user_id, role) VALUES (?, ?, 'owner')"
|
|
||||||
)
|
|
||||||
.bind(org_id)
|
.bind(org_id)
|
||||||
.bind(user_id)
|
.bind(user_id)
|
||||||
.execute(&*state.db)
|
.execute(&*state.db)
|
||||||
@@ -204,7 +215,9 @@ async fn get_organization(
|
|||||||
let (user_id, _) = extract_jwt_user(&headers, &state)?;
|
let (user_id, _) = extract_jwt_user(&headers, &state)?;
|
||||||
|
|
||||||
if !is_org_member(&state.db, org_id, user_id).await? {
|
if !is_org_member(&state.db, org_id, user_id).await? {
|
||||||
return Err(AppError::Forbidden("You are not a member of this organization".into()));
|
return Err(AppError::Forbidden(
|
||||||
|
"You are not a member of this organization".into(),
|
||||||
|
));
|
||||||
}
|
}
|
||||||
|
|
||||||
let org: Option<Organization> = sqlx::query_as("SELECT * FROM organizations WHERE id = ?")
|
let org: Option<Organization> = sqlx::query_as("SELECT * FROM organizations WHERE id = ?")
|
||||||
@@ -269,24 +282,28 @@ async fn list_members(
|
|||||||
let (user_id, _) = extract_jwt_user(&headers, &state)?;
|
let (user_id, _) = extract_jwt_user(&headers, &state)?;
|
||||||
|
|
||||||
if !is_org_member(&state.db, org_id, user_id).await? {
|
if !is_org_member(&state.db, org_id, user_id).await? {
|
||||||
return Err(AppError::Forbidden("You are not a member of this organization".into()));
|
return Err(AppError::Forbidden(
|
||||||
|
"You are not a member of this organization".into(),
|
||||||
|
));
|
||||||
}
|
}
|
||||||
|
|
||||||
let members: Vec<Value> = sqlx::query_as::<_, OrganizationMember>(
|
let members: Vec<Value> = sqlx::query_as::<_, OrganizationMember>(
|
||||||
"SELECT * FROM organization_members WHERE org_id = ? ORDER BY joined_at ASC"
|
"SELECT * FROM organization_members WHERE org_id = ? ORDER BY joined_at ASC",
|
||||||
)
|
)
|
||||||
.bind(org_id)
|
.bind(org_id)
|
||||||
.fetch_all(&*state.db)
|
.fetch_all(&*state.db)
|
||||||
.await?
|
.await?
|
||||||
.into_iter()
|
.into_iter()
|
||||||
.map(|m| json!({
|
.map(|m| {
|
||||||
|
json!({
|
||||||
"id": m.id,
|
"id": m.id,
|
||||||
"orgId": m.org_id,
|
"orgId": m.org_id,
|
||||||
"userId": m.user_id,
|
"userId": m.user_id,
|
||||||
"role": m.role,
|
"role": m.role,
|
||||||
"status": m.status,
|
"status": m.status,
|
||||||
"joinedAt": m.joined_at,
|
"joinedAt": m.joined_at,
|
||||||
}))
|
})
|
||||||
|
})
|
||||||
.collect();
|
.collect();
|
||||||
|
|
||||||
Ok(Json(json!({ "members": members })))
|
Ok(Json(json!({ "members": members })))
|
||||||
@@ -302,11 +319,15 @@ async fn update_member_role(
|
|||||||
let (user_id, _) = extract_jwt_user(&headers, &state)?;
|
let (user_id, _) = extract_jwt_user(&headers, &state)?;
|
||||||
require_org_admin(&state.db, org_id, user_id).await?;
|
require_org_admin(&state.db, org_id, user_id).await?;
|
||||||
|
|
||||||
let role = body.get("role").and_then(|v| v.as_str())
|
let role = body
|
||||||
|
.get("role")
|
||||||
|
.and_then(|v| v.as_str())
|
||||||
.ok_or_else(|| AppError::BadRequest("role is required".into()))?;
|
.ok_or_else(|| AppError::BadRequest("role is required".into()))?;
|
||||||
|
|
||||||
if !["owner", "admin", "member"].contains(&role) {
|
if !["owner", "admin", "member"].contains(&role) {
|
||||||
return Err(AppError::BadRequest("Role must be one of: owner, admin, member".into()));
|
return Err(AppError::BadRequest(
|
||||||
|
"Role must be one of: owner, admin, member".into(),
|
||||||
|
));
|
||||||
}
|
}
|
||||||
|
|
||||||
sqlx::query("UPDATE organization_members SET role = ? WHERE org_id = ? AND user_id = ?")
|
sqlx::query("UPDATE organization_members SET role = ? WHERE org_id = ? AND user_id = ?")
|
||||||
@@ -348,9 +369,8 @@ async fn list_invitations(
|
|||||||
let (user_id, _) = extract_jwt_user(&headers, &state)?;
|
let (user_id, _) = extract_jwt_user(&headers, &state)?;
|
||||||
require_org_admin(&state.db, org_id, user_id).await?;
|
require_org_admin(&state.db, org_id, user_id).await?;
|
||||||
|
|
||||||
let invitations: Vec<Invitation> = sqlx::query_as(
|
let invitations: Vec<Invitation> =
|
||||||
"SELECT * FROM invitations WHERE org_id = ? ORDER BY created_at DESC"
|
sqlx::query_as("SELECT * FROM invitations WHERE org_id = ? ORDER BY created_at DESC")
|
||||||
)
|
|
||||||
.bind(org_id)
|
.bind(org_id)
|
||||||
.fetch_all(&*state.db)
|
.fetch_all(&*state.db)
|
||||||
.await?;
|
.await?;
|
||||||
@@ -368,12 +388,19 @@ async fn create_invitation(
|
|||||||
let (user_id, _) = extract_jwt_user(&headers, &state)?;
|
let (user_id, _) = extract_jwt_user(&headers, &state)?;
|
||||||
require_org_admin(&state.db, org_id, user_id).await?;
|
require_org_admin(&state.db, org_id, user_id).await?;
|
||||||
|
|
||||||
let email = body.get("email").and_then(|v| v.as_str())
|
let email = body
|
||||||
|
.get("email")
|
||||||
|
.and_then(|v| v.as_str())
|
||||||
.ok_or_else(|| AppError::BadRequest("Email is required".into()))?;
|
.ok_or_else(|| AppError::BadRequest("Email is required".into()))?;
|
||||||
let role = body.get("role").and_then(|v| v.as_str()).unwrap_or("member");
|
let role = body
|
||||||
|
.get("role")
|
||||||
|
.and_then(|v| v.as_str())
|
||||||
|
.unwrap_or("member");
|
||||||
|
|
||||||
if !["admin", "member"].contains(&role) {
|
if !["admin", "member"].contains(&role) {
|
||||||
return Err(AppError::BadRequest("Role must be one of: admin, member".into()));
|
return Err(AppError::BadRequest(
|
||||||
|
"Role must be one of: admin, member".into(),
|
||||||
|
));
|
||||||
}
|
}
|
||||||
|
|
||||||
// Generate token
|
// Generate token
|
||||||
@@ -429,9 +456,8 @@ async fn get_invitation(
|
|||||||
State(state): State<AppState>,
|
State(state): State<AppState>,
|
||||||
Path(token): Path<String>,
|
Path(token): Path<String>,
|
||||||
) -> AppResult<Json<Value>> {
|
) -> AppResult<Json<Value>> {
|
||||||
let invitation: Option<Invitation> = sqlx::query_as(
|
let invitation: Option<Invitation> =
|
||||||
"SELECT * FROM invitations WHERE token = ? AND status = 'pending'"
|
sqlx::query_as("SELECT * FROM invitations WHERE token = ? AND status = 'pending'")
|
||||||
)
|
|
||||||
.bind(&token)
|
.bind(&token)
|
||||||
.fetch_optional(&*state.db)
|
.fetch_optional(&*state.db)
|
||||||
.await?;
|
.await?;
|
||||||
@@ -450,9 +476,8 @@ async fn get_invitation(
|
|||||||
}
|
}
|
||||||
|
|
||||||
// Get org name
|
// Get org name
|
||||||
let org_name: Option<(String,)> = sqlx::query_as(
|
let org_name: Option<(String,)> =
|
||||||
"SELECT name FROM organizations WHERE id = ?"
|
sqlx::query_as("SELECT name FROM organizations WHERE id = ?")
|
||||||
)
|
|
||||||
.bind(i.org_id)
|
.bind(i.org_id)
|
||||||
.fetch_optional(&*state.db)
|
.fetch_optional(&*state.db)
|
||||||
.await?;
|
.await?;
|
||||||
@@ -484,16 +509,19 @@ async fn accept_invitation(
|
|||||||
) -> AppResult<Json<Value>> {
|
) -> AppResult<Json<Value>> {
|
||||||
let (user_id, _) = extract_jwt_user(&headers, &state)?;
|
let (user_id, _) = extract_jwt_user(&headers, &state)?;
|
||||||
|
|
||||||
let invitation: Option<Invitation> = sqlx::query_as(
|
let invitation: Option<Invitation> =
|
||||||
"SELECT * FROM invitations WHERE token = ? AND status = 'pending'"
|
sqlx::query_as("SELECT * FROM invitations WHERE token = ? AND status = 'pending'")
|
||||||
)
|
|
||||||
.bind(&token)
|
.bind(&token)
|
||||||
.fetch_optional(&*state.db)
|
.fetch_optional(&*state.db)
|
||||||
.await?;
|
.await?;
|
||||||
|
|
||||||
let inv = match invitation {
|
let inv = match invitation {
|
||||||
Some(i) => i,
|
Some(i) => i,
|
||||||
None => return Err(AppError::BadRequest("Invitation not found or expired".into())),
|
None => {
|
||||||
|
return Err(AppError::BadRequest(
|
||||||
|
"Invitation not found or expired".into(),
|
||||||
|
))
|
||||||
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
// Check expiry
|
// Check expiry
|
||||||
@@ -504,7 +532,9 @@ async fn accept_invitation(
|
|||||||
});
|
});
|
||||||
|
|
||||||
if expired {
|
if expired {
|
||||||
return Err(AppError::BadRequest("Invitation not found or expired".into()));
|
return Err(AppError::BadRequest(
|
||||||
|
"Invitation not found or expired".into(),
|
||||||
|
));
|
||||||
}
|
}
|
||||||
|
|
||||||
// Mark invitation as accepted
|
// Mark invitation as accepted
|
||||||
@@ -515,7 +545,7 @@ async fn accept_invitation(
|
|||||||
|
|
||||||
// Add user as member
|
// Add user as member
|
||||||
sqlx::query(
|
sqlx::query(
|
||||||
"INSERT OR IGNORE INTO organization_members (org_id, user_id, role) VALUES (?, ?, ?)"
|
"INSERT OR IGNORE INTO organization_members (org_id, user_id, role) VALUES (?, ?, ?)",
|
||||||
)
|
)
|
||||||
.bind(inv.org_id)
|
.bind(inv.org_id)
|
||||||
.bind(user_id)
|
.bind(user_id)
|
||||||
|
|||||||
@@ -6,22 +6,25 @@
|
|||||||
// GET /v1/models — list available models
|
// GET /v1/models — list available models
|
||||||
|
|
||||||
use crate::{
|
use crate::{
|
||||||
auth::api_key::{ authorize_downstream_token, DownstreamAuthResult },
|
auth::api_key::{authorize_downstream_token, DownstreamAuthResult},
|
||||||
error::{ AppError, AppResult },
|
error::{AppError, AppResult},
|
||||||
proxy::executors::{ self, UpstreamRoute },
|
proxy::conductor::{
|
||||||
|
AttemptResult, ExecuteInput, ExecuteResult, ProxyConductor, SelectedChannel,
|
||||||
|
},
|
||||||
|
proxy::executors::{self, UpstreamRoute},
|
||||||
|
proxy::token_router::TokenRouter,
|
||||||
AppState,
|
AppState,
|
||||||
};
|
};
|
||||||
use axum::{
|
use axum::{
|
||||||
extract::State,
|
extract::State,
|
||||||
http::HeaderMap,
|
http::HeaderMap,
|
||||||
response::{ IntoResponse, Response, Sse },
|
response::{IntoResponse, Response, Sse},
|
||||||
routing::{ get, post },
|
routing::{get, post},
|
||||||
Json, Router,
|
Json, Router,
|
||||||
};
|
};
|
||||||
use futures::stream;
|
use futures::{stream, FutureExt, StreamExt};
|
||||||
use futures::StreamExt;
|
use serde_json::{json, Value};
|
||||||
use serde_json::{ json, Value };
|
use std::{convert::Infallible, sync::Arc};
|
||||||
use std::convert::Infallible;
|
|
||||||
use tracing::info;
|
use tracing::info;
|
||||||
|
|
||||||
// ── Auth types ────────────────────────────────────────────────────────
|
// ── Auth types ────────────────────────────────────────────────────────
|
||||||
@@ -73,17 +76,20 @@ async fn chat_completions(
|
|||||||
.ok_or_else(|| AppError::BadRequest("model is required".into()))?
|
.ok_or_else(|| AppError::BadRequest("model is required".into()))?
|
||||||
.to_string();
|
.to_string();
|
||||||
|
|
||||||
let route = resolve_route(&state, &model).await?;
|
let stream = body
|
||||||
let stream = body.get("stream").and_then(|s| s.as_bool()).unwrap_or(false);
|
.get("stream")
|
||||||
|
.and_then(|s| s.as_bool())
|
||||||
|
.unwrap_or(false);
|
||||||
|
|
||||||
if stream {
|
if stream {
|
||||||
let chunks = executors::execute_chat_completions_stream(state, route, body).await?;
|
let chunks = execute_chat_completions_stream_with_conductor(state, model, body).await?;
|
||||||
let sse = Sse::new(stream::iter(chunks).map(|chunk| {
|
let sse =
|
||||||
|
Sse::new(stream::iter(chunks).map(|chunk| {
|
||||||
Ok::<_, Infallible>(axum::response::sse::Event::default().data(chunk))
|
Ok::<_, Infallible>(axum::response::sse::Event::default().data(chunk))
|
||||||
}));
|
}));
|
||||||
Ok(sse.into_response())
|
Ok(sse.into_response())
|
||||||
} else {
|
} else {
|
||||||
let result = executors::execute_chat_completions(state, route, body).await?;
|
let result = execute_chat_completions_with_conductor(state, model, body).await?;
|
||||||
Ok(Json(result).into_response())
|
Ok(Json(result).into_response())
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -168,7 +174,8 @@ async fn audio_transcriptions(
|
|||||||
.to_string();
|
.to_string();
|
||||||
|
|
||||||
let route = resolve_route(&state, &model).await?;
|
let route = resolve_route(&state, &model).await?;
|
||||||
let result = executors::execute_proxy_request(state, route, "/audio/transcriptions", body).await?;
|
let result =
|
||||||
|
executors::execute_proxy_request(state, route, "/audio/transcriptions", body).await?;
|
||||||
Ok(Json(result))
|
Ok(Json(result))
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -189,7 +196,8 @@ async fn audio_translations(
|
|||||||
.to_string();
|
.to_string();
|
||||||
|
|
||||||
let route = resolve_route(&state, &model).await?;
|
let route = resolve_route(&state, &model).await?;
|
||||||
let result = executors::execute_proxy_request(state, route, "/audio/translations", body).await?;
|
let result =
|
||||||
|
executors::execute_proxy_request(state, route, "/audio/translations", body).await?;
|
||||||
Ok(Json(result))
|
Ok(Json(result))
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -235,16 +243,13 @@ async fn images_generations(
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// GET /v1/models
|
/// GET /v1/models
|
||||||
async fn list_models(
|
async fn list_models(State(state): State<AppState>, headers: HeaderMap) -> AppResult<Json<Value>> {
|
||||||
State(state): State<AppState>,
|
|
||||||
headers: HeaderMap,
|
|
||||||
) -> AppResult<Json<Value>> {
|
|
||||||
info!("GET /v1/models");
|
info!("GET /v1/models");
|
||||||
|
|
||||||
check_auth(&state, &headers).await?;
|
check_auth(&state, &headers).await?;
|
||||||
|
|
||||||
let rows: Vec<(String, String)> = sqlx::query_as(
|
let rows: Vec<(String, String)> = sqlx::query_as(
|
||||||
"SELECT tr.name, tr.model FROM token_routes tr WHERE tr.enabled = 1 ORDER BY tr.name"
|
"SELECT tr.name, tr.model FROM token_routes tr WHERE tr.enabled = 1 ORDER BY tr.name",
|
||||||
)
|
)
|
||||||
.fetch_all(&*state.db)
|
.fetch_all(&*state.db)
|
||||||
.await?;
|
.await?;
|
||||||
@@ -277,7 +282,8 @@ async fn check_auth(state: &AppState, headers: &HeaderMap) -> Result<(), AppErro
|
|||||||
.and_then(|v| v.trim().split_whitespace().next())
|
.and_then(|v| v.trim().split_whitespace().next())
|
||||||
.ok_or_else(|| AppError::Unauthorized("Missing Authorization header".into()))?;
|
.ok_or_else(|| AppError::Unauthorized("Missing Authorization header".into()))?;
|
||||||
|
|
||||||
let result = authorize_downstream_token(auth, &state.db).await
|
let result = authorize_downstream_token(auth, &state.db)
|
||||||
|
.await
|
||||||
.map_err(|e| AppError::Unauthorized(e))?;
|
.map_err(|e| AppError::Unauthorized(e))?;
|
||||||
|
|
||||||
match result {
|
match result {
|
||||||
@@ -295,27 +301,379 @@ async fn check_auth(state: &AppState, headers: &HeaderMap) -> Result<(), AppErro
|
|||||||
|
|
||||||
/// Resolve upstream endpoint + API key for a given model.
|
/// Resolve upstream endpoint + API key for a given model.
|
||||||
async fn resolve_route(state: &AppState, model: &str) -> AppResult<UpstreamRoute> {
|
async fn resolve_route(state: &AppState, model: &str) -> AppResult<UpstreamRoute> {
|
||||||
let row: Option<(String, String)> = sqlx::query_as(
|
let selected = TokenRouter::new(state.db.clone())
|
||||||
"SELECT s.base_url, COALESCE(a.api_key, s.default_key, '') as api_key
|
.resolve_selected(model)
|
||||||
FROM token_routes tr
|
|
||||||
JOIN sites s ON s.id = tr.site_id
|
|
||||||
LEFT JOIN accounts a ON a.id = tr.account_id AND a.site_id = s.id
|
|
||||||
WHERE tr.enabled = 1 AND tr.model = ?
|
|
||||||
LIMIT 1"
|
|
||||||
)
|
|
||||||
.bind(model)
|
|
||||||
.fetch_optional(&*state.db)
|
|
||||||
.await?;
|
.await?;
|
||||||
|
Ok(selected_channel_to_route(selected))
|
||||||
|
}
|
||||||
|
|
||||||
match row {
|
fn selected_channel_to_route(selected: SelectedChannel) -> UpstreamRoute {
|
||||||
Some((base_url, api_key)) => Ok(UpstreamRoute {
|
UpstreamRoute {
|
||||||
base_url,
|
base_url: selected.base_url,
|
||||||
api_key,
|
api_key: selected.api_key,
|
||||||
model: model.to_string(),
|
model: selected.actual_model,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn body_with_selected_model(mut body: Value, selected: &SelectedChannel) -> Value {
|
||||||
|
if let Some(object) = body.as_object_mut() {
|
||||||
|
object.insert(
|
||||||
|
"model".to_string(),
|
||||||
|
Value::String(selected.actual_model.clone()),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
body
|
||||||
|
}
|
||||||
|
|
||||||
|
fn build_conductor(state: &AppState) -> ProxyConductor {
|
||||||
|
let db = state.db.clone();
|
||||||
|
let select_db = db.clone();
|
||||||
|
let select_next_db = db.clone();
|
||||||
|
let success_db = db.clone();
|
||||||
|
let failure_db = db;
|
||||||
|
|
||||||
|
ProxyConductor {
|
||||||
|
select_channel: Arc::new(move |model, _policy| {
|
||||||
|
let router = TokenRouter::new(select_db.clone());
|
||||||
|
async move { router.resolve_selected(&model).await.ok() }.boxed()
|
||||||
}),
|
}),
|
||||||
None => Err(AppError::BadRequest(format!(
|
select_next_channel: Arc::new(move |model, excluded, _policy| {
|
||||||
"No route configured for model '{}'",
|
let router = TokenRouter::new(select_next_db.clone());
|
||||||
model
|
async move { router.resolve_next_selected(&model, &excluded).await.ok() }.boxed()
|
||||||
|
}),
|
||||||
|
record_success: Arc::new(move |channel_id, metrics| {
|
||||||
|
let router = TokenRouter::new(success_db.clone());
|
||||||
|
async move {
|
||||||
|
if let Err(err) = router.record_success(channel_id, metrics).await {
|
||||||
|
tracing::warn!(%err, channel_id, "failed to record proxy success");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
.boxed()
|
||||||
|
}),
|
||||||
|
record_failure: Arc::new(move |channel_id, info| {
|
||||||
|
let router = TokenRouter::new(failure_db.clone());
|
||||||
|
async move {
|
||||||
|
if let Err(err) = router.record_failure(channel_id, info).await {
|
||||||
|
tracing::warn!(%err, channel_id, "failed to record proxy failure");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
.boxed()
|
||||||
|
}),
|
||||||
|
refresh_auth: None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute_chat_completions_with_conductor(
|
||||||
|
state: AppState,
|
||||||
|
model: String,
|
||||||
|
body: Value,
|
||||||
|
) -> AppResult<Value> {
|
||||||
|
let conductor = build_conductor(&state);
|
||||||
|
let result = conductor
|
||||||
|
.execute(ExecuteInput {
|
||||||
|
requested_model: model,
|
||||||
|
downstream_policy: None,
|
||||||
|
attempt_fn: Box::new(move |ctx| {
|
||||||
|
let state = state.clone();
|
||||||
|
let body = body_with_selected_model(body.clone(), &ctx.selected);
|
||||||
|
async move {
|
||||||
|
match executors::execute_chat_completions(
|
||||||
|
state,
|
||||||
|
selected_channel_to_route(ctx.selected),
|
||||||
|
body,
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
Ok(response) => AttemptResult::Success {
|
||||||
|
response,
|
||||||
|
latency_ms: None,
|
||||||
|
cost: None,
|
||||||
|
},
|
||||||
|
Err(err) => AttemptResult::Failure {
|
||||||
|
action: crate::proxy::conductor::AttemptFailureAction::Failover,
|
||||||
|
status: None,
|
||||||
|
raw_error_text: Some(err.to_string()),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
.boxed()
|
||||||
|
}),
|
||||||
|
on_terminal_failure: None,
|
||||||
|
})
|
||||||
|
.await;
|
||||||
|
|
||||||
|
conductor_result_to_app_result(result)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute_chat_completions_stream_with_conductor(
|
||||||
|
state: AppState,
|
||||||
|
model: String,
|
||||||
|
body: Value,
|
||||||
|
) -> AppResult<Vec<String>> {
|
||||||
|
let conductor = build_conductor(&state);
|
||||||
|
let result = conductor
|
||||||
|
.execute(ExecuteInput {
|
||||||
|
requested_model: model,
|
||||||
|
downstream_policy: None,
|
||||||
|
attempt_fn: Box::new(move |ctx| {
|
||||||
|
let state = state.clone();
|
||||||
|
let body = body_with_selected_model(body.clone(), &ctx.selected);
|
||||||
|
async move {
|
||||||
|
match executors::execute_chat_completions_stream(
|
||||||
|
state,
|
||||||
|
selected_channel_to_route(ctx.selected),
|
||||||
|
body,
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
Ok(chunks) => AttemptResult::Success {
|
||||||
|
response: Value::Array(chunks.into_iter().map(Value::String).collect()),
|
||||||
|
latency_ms: None,
|
||||||
|
cost: None,
|
||||||
|
},
|
||||||
|
Err(err) => AttemptResult::Failure {
|
||||||
|
action: crate::proxy::conductor::AttemptFailureAction::Failover,
|
||||||
|
status: None,
|
||||||
|
raw_error_text: Some(err.to_string()),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
.boxed()
|
||||||
|
}),
|
||||||
|
on_terminal_failure: None,
|
||||||
|
})
|
||||||
|
.await;
|
||||||
|
|
||||||
|
match conductor_result_to_app_result(result)? {
|
||||||
|
Value::Array(chunks) => chunks
|
||||||
|
.into_iter()
|
||||||
|
.map(|chunk| {
|
||||||
|
chunk.as_str().map(str::to_string).ok_or_else(|| {
|
||||||
|
AppError::UpstreamError("Invalid streaming response from conductor".into())
|
||||||
|
})
|
||||||
|
})
|
||||||
|
.collect(),
|
||||||
|
_ => Err(AppError::UpstreamError(
|
||||||
|
"Invalid streaming response from conductor".into(),
|
||||||
|
)),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn conductor_result_to_app_result(result: ExecuteResult) -> AppResult<Value> {
|
||||||
|
match result {
|
||||||
|
ExecuteResult::Success { response, .. } => Ok(response),
|
||||||
|
ExecuteResult::Failure {
|
||||||
|
reason,
|
||||||
|
status,
|
||||||
|
raw_error_text,
|
||||||
|
attempts,
|
||||||
|
..
|
||||||
|
} => 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())
|
||||||
))),
|
))),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
use crate::config::{AppConfig, AuthConfig, DbConfig, ProxyConfig, ServerConfig};
|
||||||
|
use axum::http::HeaderValue;
|
||||||
|
use serde_json::json;
|
||||||
|
use sqlx::sqlite::SqlitePoolOptions;
|
||||||
|
|
||||||
|
async fn test_state(upstream_base_url: String) -> AppState {
|
||||||
|
let pool = SqlitePoolOptions::new()
|
||||||
|
.max_connections(1)
|
||||||
|
.connect("sqlite::memory:")
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
sqlx::query(
|
||||||
|
"CREATE TABLE downstream_api_keys (
|
||||||
|
id INTEGER PRIMARY KEY,
|
||||||
|
name TEXT NOT NULL,
|
||||||
|
key TEXT NOT NULL,
|
||||||
|
description TEXT,
|
||||||
|
group_name TEXT,
|
||||||
|
tags TEXT,
|
||||||
|
enabled INTEGER DEFAULT 1,
|
||||||
|
expires_at TEXT,
|
||||||
|
max_cost REAL,
|
||||||
|
used_cost REAL DEFAULT 0,
|
||||||
|
max_requests INTEGER,
|
||||||
|
used_requests INTEGER DEFAULT 0,
|
||||||
|
supported_models TEXT,
|
||||||
|
allowed_route_ids TEXT,
|
||||||
|
site_weight_multipliers TEXT,
|
||||||
|
excluded_site_ids TEXT,
|
||||||
|
excluded_credential_refs TEXT,
|
||||||
|
last_used_at TEXT,
|
||||||
|
user_id INTEGER,
|
||||||
|
created_at TEXT DEFAULT (datetime('now')),
|
||||||
|
updated_at TEXT DEFAULT (datetime('now'))
|
||||||
|
)",
|
||||||
|
)
|
||||||
|
.execute(&pool)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
sqlx::query(
|
||||||
|
"CREATE TABLE sites (
|
||||||
|
id INTEGER PRIMARY KEY,
|
||||||
|
name TEXT NOT NULL,
|
||||||
|
url TEXT NOT NULL,
|
||||||
|
platform TEXT NOT NULL DEFAULT 'openai',
|
||||||
|
status TEXT NOT NULL DEFAULT 'active',
|
||||||
|
api_key TEXT
|
||||||
|
)",
|
||||||
|
)
|
||||||
|
.execute(&pool)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
sqlx::query(
|
||||||
|
"CREATE TABLE token_routes (
|
||||||
|
id INTEGER PRIMARY KEY,
|
||||||
|
model_pattern TEXT NOT NULL,
|
||||||
|
model_mapping TEXT,
|
||||||
|
route_mode TEXT,
|
||||||
|
enabled INTEGER DEFAULT 1
|
||||||
|
)",
|
||||||
|
)
|
||||||
|
.execute(&pool)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
sqlx::query(
|
||||||
|
"CREATE TABLE accounts (
|
||||||
|
id INTEGER PRIMARY KEY,
|
||||||
|
site_id INTEGER NOT NULL,
|
||||||
|
access_token TEXT NOT NULL,
|
||||||
|
api_token TEXT,
|
||||||
|
status TEXT NOT NULL DEFAULT 'active'
|
||||||
|
)",
|
||||||
|
)
|
||||||
|
.execute(&pool)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
sqlx::query(
|
||||||
|
"CREATE TABLE account_tokens (
|
||||||
|
id INTEGER PRIMARY KEY,
|
||||||
|
account_id INTEGER NOT NULL,
|
||||||
|
name TEXT NOT NULL,
|
||||||
|
token TEXT NOT NULL,
|
||||||
|
enabled INTEGER DEFAULT 1
|
||||||
|
)",
|
||||||
|
)
|
||||||
|
.execute(&pool)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
sqlx::query(
|
||||||
|
"CREATE TABLE route_channels (
|
||||||
|
id INTEGER PRIMARY KEY,
|
||||||
|
route_id INTEGER NOT NULL,
|
||||||
|
account_id INTEGER NOT NULL,
|
||||||
|
token_id INTEGER,
|
||||||
|
source_model TEXT,
|
||||||
|
priority INTEGER DEFAULT 0,
|
||||||
|
weight INTEGER DEFAULT 10,
|
||||||
|
enabled INTEGER DEFAULT 1,
|
||||||
|
consecutive_fail_count INTEGER NOT NULL DEFAULT 0,
|
||||||
|
cooldown_level INTEGER NOT NULL DEFAULT 0,
|
||||||
|
cooldown_until TEXT,
|
||||||
|
last_selected_at TEXT,
|
||||||
|
success_count INTEGER DEFAULT 0,
|
||||||
|
fail_count INTEGER DEFAULT 0,
|
||||||
|
total_latency_ms INTEGER DEFAULT 0,
|
||||||
|
total_cost REAL DEFAULT 0,
|
||||||
|
last_used_at TEXT,
|
||||||
|
last_fail_at TEXT
|
||||||
|
)",
|
||||||
|
)
|
||||||
|
.execute(&pool)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
sqlx::query("INSERT INTO downstream_api_keys (id, name, key, enabled) VALUES (1, 'test', 'sk-test', 1)")
|
||||||
|
.execute(&pool)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
sqlx::query(
|
||||||
|
"INSERT INTO sites (id, name, url, api_key) VALUES (1, 'OpenAI', ?, 'site-key')",
|
||||||
|
)
|
||||||
|
.bind(upstream_base_url)
|
||||||
|
.execute(&pool)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
sqlx::query("INSERT INTO accounts (id, site_id, access_token, api_token) VALUES (1, 1, 'access-key', NULL)")
|
||||||
|
.execute(&pool)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
sqlx::query("INSERT INTO token_routes (id, model_pattern, model_mapping, route_mode, enabled) VALUES (1, 'gpt-test', 'upstream-test', 'exact', 1)")
|
||||||
|
.execute(&pool)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
sqlx::query("INSERT INTO route_channels (id, route_id, account_id) VALUES (1, 1, 1)")
|
||||||
|
.execute(&pool)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
AppState::new(
|
||||||
|
AppConfig {
|
||||||
|
server: ServerConfig {
|
||||||
|
host: "127.0.0.1".into(),
|
||||||
|
port: 0,
|
||||||
|
log_level: "info".into(),
|
||||||
|
cors_origins: vec![],
|
||||||
|
},
|
||||||
|
db: DbConfig {
|
||||||
|
url: "sqlite::memory:".into(),
|
||||||
|
max_connections: 1,
|
||||||
|
ssl: false,
|
||||||
|
runtime_db_url: None,
|
||||||
|
runtime_db_dialect: None,
|
||||||
|
runtime_db_ssl: None,
|
||||||
|
},
|
||||||
|
auth: AuthConfig {
|
||||||
|
jwt_secret: "secret".into(),
|
||||||
|
jwt_expiry_hours: 1,
|
||||||
|
auth_token: "auth".into(),
|
||||||
|
proxy_token: "proxy".into(),
|
||||||
|
},
|
||||||
|
proxy: ProxyConfig {
|
||||||
|
upstream_timeout_secs: 30,
|
||||||
|
first_byte_timeout_secs: 30,
|
||||||
|
max_retries: 3,
|
||||||
|
proxy_log_retention_days: 7,
|
||||||
|
},
|
||||||
|
},
|
||||||
|
pool,
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn chat_route_uses_token_router_and_conductor() {
|
||||||
|
let state = test_state("http://127.0.0.1:1".to_string()).await;
|
||||||
|
let db = state.db.clone();
|
||||||
|
let mut headers = HeaderMap::new();
|
||||||
|
headers.insert("Authorization", HeaderValue::from_static("Bearer sk-test"));
|
||||||
|
|
||||||
|
let result = chat_completions(
|
||||||
|
State(state),
|
||||||
|
headers,
|
||||||
|
Json(json!({
|
||||||
|
"model": "gpt-test",
|
||||||
|
"messages": [{"role": "user", "content": "hello"}]
|
||||||
|
})),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
|
||||||
|
assert!(matches!(result, Err(AppError::UpstreamError(_))));
|
||||||
|
let fail_count: i64 =
|
||||||
|
sqlx::query_scalar("SELECT fail_count FROM route_channels WHERE id = 1")
|
||||||
|
.fetch_one(&*db)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
assert!(fail_count > 0);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -1,14 +1,7 @@
|
|||||||
// Search routes — cross-entity search across sites, accounts, tokens, etc.
|
// Search routes — cross-entity search across sites, accounts, tokens, etc.
|
||||||
|
|
||||||
use crate::{
|
use crate::{error::AppResult, AppState};
|
||||||
error::AppResult,
|
use axum::{extract::State, routing::post, Json, Router};
|
||||||
AppState,
|
|
||||||
};
|
|
||||||
use axum::{
|
|
||||||
extract::State,
|
|
||||||
routing::post,
|
|
||||||
Json, Router,
|
|
||||||
};
|
|
||||||
use serde::Deserialize;
|
use serde::Deserialize;
|
||||||
use serde_json::{json, Value};
|
use serde_json::{json, Value};
|
||||||
use std::collections::{HashMap, HashSet};
|
use std::collections::{HashMap, HashSet};
|
||||||
@@ -21,8 +14,7 @@ pub struct SearchBody {
|
|||||||
}
|
}
|
||||||
|
|
||||||
pub fn routes() -> Router<AppState> {
|
pub fn routes() -> Router<AppState> {
|
||||||
Router::new()
|
Router::new().route("/api/search", post(search_handler))
|
||||||
.route("/api/search", post(search_handler))
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// POST /api/search — search across all entities
|
/// POST /api/search — search across all entities
|
||||||
@@ -56,10 +48,13 @@ async fn search_handler(
|
|||||||
|
|
||||||
// Deduplicate by id
|
// Deduplicate by id
|
||||||
let mut seen_sites = HashSet::new();
|
let mut seen_sites = HashSet::new();
|
||||||
let unique_sites: Vec<Value> = sites.into_iter().filter(|s| {
|
let unique_sites: Vec<Value> = sites
|
||||||
|
.into_iter()
|
||||||
|
.filter(|s| {
|
||||||
let id = s["id"].as_i64().unwrap_or(0);
|
let id = s["id"].as_i64().unwrap_or(0);
|
||||||
seen_sites.insert(id)
|
seen_sites.insert(id)
|
||||||
}).collect();
|
})
|
||||||
|
.collect();
|
||||||
|
|
||||||
// 2. Search accounts (with site join)
|
// 2. Search accounts (with site join)
|
||||||
let accounts: Vec<Value> = sqlx::query_as::<_, (i64, String, String, String, i64, String, String, Option<String>)>(
|
let accounts: Vec<Value> = sqlx::query_as::<_, (i64, String, String, String, i64, String, String, Option<String>)>(
|
||||||
@@ -151,9 +146,10 @@ async fn search_handler(
|
|||||||
INNER JOIN accounts a ON a.id = t.account_id AND a.status = 'active'
|
INNER JOIN accounts a ON a.id = t.account_id AND a.status = 'active'
|
||||||
INNER JOIN sites s ON s.id = a.site_id
|
INNER JOIN sites s ON s.id = a.site_id
|
||||||
WHERE tma.model_name LIKE ? AND tma.available = 1 AND t.value_status = 'ready'
|
WHERE tma.model_name LIKE ? AND tma.available = 1 AND t.value_status = 'ready'
|
||||||
LIMIT ?"#
|
LIMIT ?"#,
|
||||||
)
|
)
|
||||||
.bind(&q).bind(per_category * 10)
|
.bind(&q)
|
||||||
|
.bind(per_category * 10)
|
||||||
.fetch_all(&*state.db)
|
.fetch_all(&*state.db)
|
||||||
.await?;
|
.await?;
|
||||||
|
|
||||||
@@ -163,9 +159,10 @@ async fn search_handler(
|
|||||||
INNER JOIN accounts a ON a.id = ma.account_id AND a.status = 'active'
|
INNER JOIN accounts a ON a.id = ma.account_id AND a.status = 'active'
|
||||||
INNER JOIN sites s ON s.id = a.site_id AND s.status = 'active'
|
INNER JOIN sites s ON s.id = a.site_id AND s.status = 'active'
|
||||||
WHERE ma.model_name LIKE ? AND ma.available = 1
|
WHERE ma.model_name LIKE ? AND ma.available = 1
|
||||||
LIMIT ?"#
|
LIMIT ?"#,
|
||||||
)
|
)
|
||||||
.bind(&q).bind(per_category * 10)
|
.bind(&q)
|
||||||
|
.bind(per_category * 10)
|
||||||
.fetch_all(&*state.db)
|
.fetch_all(&*state.db)
|
||||||
.await?;
|
.await?;
|
||||||
|
|
||||||
@@ -197,7 +194,12 @@ async fn search_handler(
|
|||||||
models.sort_by(|a, b| {
|
models.sort_by(|a, b| {
|
||||||
let bc = b["accountCount"].as_i64().unwrap_or(0);
|
let bc = b["accountCount"].as_i64().unwrap_or(0);
|
||||||
let ac = a["accountCount"].as_i64().unwrap_or(0);
|
let ac = a["accountCount"].as_i64().unwrap_or(0);
|
||||||
bc.cmp(&ac).then_with(|| a["name"].as_str().unwrap_or("").cmp(&b["name"].as_str().unwrap_or("")))
|
bc.cmp(&ac).then_with(|| {
|
||||||
|
a["name"]
|
||||||
|
.as_str()
|
||||||
|
.unwrap_or("")
|
||||||
|
.cmp(&b["name"].as_str().unwrap_or(""))
|
||||||
|
})
|
||||||
});
|
});
|
||||||
models.truncate(per_category as usize);
|
models.truncate(per_category as usize);
|
||||||
|
|
||||||
|
|||||||
@@ -26,7 +26,10 @@ pub struct ListAnnouncementsQuery {
|
|||||||
|
|
||||||
pub fn routes() -> Router<AppState> {
|
pub fn routes() -> Router<AppState> {
|
||||||
Router::new()
|
Router::new()
|
||||||
.route("/api/site-announcements", get(list_announcements).delete(delete_all))
|
.route(
|
||||||
|
"/api/site-announcements",
|
||||||
|
get(list_announcements).delete(delete_all),
|
||||||
|
)
|
||||||
.route("/api/site-announcements/{id}/read", post(mark_read))
|
.route("/api/site-announcements/{id}/read", post(mark_read))
|
||||||
.route("/api/site-announcements/read-all", post(mark_all_read))
|
.route("/api/site-announcements/read-all", post(mark_all_read))
|
||||||
}
|
}
|
||||||
@@ -80,21 +83,26 @@ async fn list_announcements(
|
|||||||
let filtered = if let Some(ref status) = q.status {
|
let filtered = if let Some(ref status) = q.status {
|
||||||
let now = chrono::Utc::now().timestamp_millis() as i64;
|
let now = chrono::Utc::now().timestamp_millis() as i64;
|
||||||
let status = status.trim().to_lowercase();
|
let status = status.trim().to_lowercase();
|
||||||
announcements.into_iter().filter(|a| {
|
announcements
|
||||||
|
.into_iter()
|
||||||
|
.filter(|a| {
|
||||||
if status == "dismissed" {
|
if status == "dismissed" {
|
||||||
a.dismissed_at.is_some()
|
a.dismissed_at.is_some()
|
||||||
} else if status == "expired" {
|
} else if status == "expired" {
|
||||||
a.dismissed_at.is_none() && a.ends_at.as_ref().map_or(false, |e| {
|
a.dismissed_at.is_none()
|
||||||
parse_dt_millis(e) < now
|
&& a.ends_at
|
||||||
})
|
.as_ref()
|
||||||
|
.map_or(false, |e| parse_dt_millis(e) < now)
|
||||||
} else if status == "active" {
|
} else if status == "active" {
|
||||||
a.dismissed_at.is_none() && a.ends_at.as_ref().map_or(true, |e| {
|
a.dismissed_at.is_none()
|
||||||
parse_dt_millis(e) >= now
|
&& a.ends_at
|
||||||
})
|
.as_ref()
|
||||||
|
.map_or(true, |e| parse_dt_millis(e) >= now)
|
||||||
} else {
|
} else {
|
||||||
true
|
true
|
||||||
}
|
}
|
||||||
}).collect()
|
})
|
||||||
|
.collect()
|
||||||
} else {
|
} else {
|
||||||
announcements
|
announcements
|
||||||
};
|
};
|
||||||
@@ -131,10 +139,7 @@ async fn list_announcements(
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// POST /api/site-announcements/{id}/read
|
/// POST /api/site-announcements/{id}/read
|
||||||
async fn mark_read(
|
async fn mark_read(State(state): State<AppState>, Path(id): Path<i64>) -> AppResult<Json<Value>> {
|
||||||
State(state): State<AppState>,
|
|
||||||
Path(id): Path<i64>,
|
|
||||||
) -> AppResult<Json<Value>> {
|
|
||||||
info!("Marking site announcement {} as read", id);
|
info!("Marking site announcement {} as read", id);
|
||||||
sqlx::query("UPDATE site_announcements SET read_at = datetime('now') WHERE id = ?")
|
sqlx::query("UPDATE site_announcements SET read_at = datetime('now') WHERE id = ?")
|
||||||
.bind(id)
|
.bind(id)
|
||||||
|
|||||||
@@ -1,9 +1,6 @@
|
|||||||
// Stats routes — usage analytics queries
|
// Stats routes — usage analytics queries
|
||||||
|
|
||||||
use crate::{
|
use crate::{error::AppResult, AppState};
|
||||||
error::AppResult,
|
|
||||||
AppState,
|
|
||||||
};
|
|
||||||
use axum::{
|
use axum::{
|
||||||
extract::{Query, State},
|
extract::{Query, State},
|
||||||
routing::get,
|
routing::get,
|
||||||
@@ -36,7 +33,10 @@ async fn model_by_site(
|
|||||||
State(state): State<AppState>,
|
State(state): State<AppState>,
|
||||||
Query(q): Query<ModelBySiteQuery>,
|
Query(q): Query<ModelBySiteQuery>,
|
||||||
) -> AppResult<Json<Value>> {
|
) -> AppResult<Json<Value>> {
|
||||||
info!("Stats model-by-site: site_id={:?}, days={:?}", q.site_id, q.days);
|
info!(
|
||||||
|
"Stats model-by-site: site_id={:?}, days={:?}",
|
||||||
|
q.site_id, q.days
|
||||||
|
);
|
||||||
let days = q.days.unwrap_or(7).max(1);
|
let days = q.days.unwrap_or(7).max(1);
|
||||||
|
|
||||||
// Calculate cutoff date
|
// Calculate cutoff date
|
||||||
@@ -51,7 +51,7 @@ async fn model_by_site(
|
|||||||
"SELECT model, local_day, total_calls, total_spend, total_tokens
|
"SELECT model, local_day, total_calls, total_spend, total_tokens
|
||||||
FROM model_day_usage
|
FROM model_day_usage
|
||||||
WHERE site_id = ? AND local_day >= ?
|
WHERE site_id = ? AND local_day >= ?
|
||||||
ORDER BY local_day DESC"
|
ORDER BY local_day DESC",
|
||||||
)
|
)
|
||||||
.bind(site_id)
|
.bind(site_id)
|
||||||
.bind(&since_day)
|
.bind(&since_day)
|
||||||
@@ -62,7 +62,7 @@ async fn model_by_site(
|
|||||||
"SELECT model, local_day, total_calls, total_spend, total_tokens
|
"SELECT model, local_day, total_calls, total_spend, total_tokens
|
||||||
FROM model_day_usage
|
FROM model_day_usage
|
||||||
WHERE local_day >= ?
|
WHERE local_day >= ?
|
||||||
ORDER BY local_day DESC"
|
ORDER BY local_day DESC",
|
||||||
)
|
)
|
||||||
.bind(&since_day)
|
.bind(&since_day)
|
||||||
.fetch_all(&*state.db)
|
.fetch_all(&*state.db)
|
||||||
@@ -70,7 +70,8 @@ async fn model_by_site(
|
|||||||
};
|
};
|
||||||
|
|
||||||
// Aggregate by model
|
// Aggregate by model
|
||||||
let mut model_map: std::collections::HashMap<String, (i64, f64, i64)> = std::collections::HashMap::new();
|
let mut model_map: std::collections::HashMap<String, (i64, f64, i64)> =
|
||||||
|
std::collections::HashMap::new();
|
||||||
for (model, _day, calls, spend, tokens) in rows {
|
for (model, _day, calls, spend, tokens) in rows {
|
||||||
let entry = model_map.entry(model).or_insert((0, 0.0, 0));
|
let entry = model_map.entry(model).or_insert((0, 0.0, 0));
|
||||||
entry.0 += calls;
|
entry.0 += calls;
|
||||||
@@ -90,7 +91,12 @@ async fn model_by_site(
|
|||||||
})
|
})
|
||||||
.collect();
|
.collect();
|
||||||
|
|
||||||
models.sort_by(|a, b| b["calls"].as_i64().unwrap_or(0).cmp(&a["calls"].as_i64().unwrap_or(0)));
|
models.sort_by(|a, b| {
|
||||||
|
b["calls"]
|
||||||
|
.as_i64()
|
||||||
|
.unwrap_or(0)
|
||||||
|
.cmp(&a["calls"].as_i64().unwrap_or(0))
|
||||||
|
});
|
||||||
|
|
||||||
Ok(Json(json!({ "models": models })))
|
Ok(Json(json!({ "models": models })))
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -2,11 +2,7 @@
|
|||||||
|
|
||||||
use std::convert::Infallible;
|
use std::convert::Infallible;
|
||||||
|
|
||||||
use crate::{
|
use crate::{db::models::background_task::BackgroundTask, error::AppResult, AppState};
|
||||||
db::models::background_task::BackgroundTask,
|
|
||||||
error::AppResult,
|
|
||||||
AppState,
|
|
||||||
};
|
|
||||||
use axum::{
|
use axum::{
|
||||||
extract::{Path, Query, State},
|
extract::{Path, Query, State},
|
||||||
response::{
|
response::{
|
||||||
@@ -31,7 +27,10 @@ pub fn routes() -> Router<AppState> {
|
|||||||
Router::new()
|
Router::new()
|
||||||
.route("/api/tasks", get(list_tasks))
|
.route("/api/tasks", get(list_tasks))
|
||||||
.route("/api/tasks/{id}", get(get_task))
|
.route("/api/tasks/{id}", get(get_task))
|
||||||
.route("/api/update-center/tasks/{id}/stream", get(stream_task_logs))
|
.route(
|
||||||
|
"/api/update-center/tasks/{id}/stream",
|
||||||
|
get(stream_task_logs),
|
||||||
|
)
|
||||||
}
|
}
|
||||||
|
|
||||||
/// GET /api/tasks — list background tasks
|
/// GET /api/tasks — list background tasks
|
||||||
@@ -43,7 +42,7 @@ async fn list_tasks(
|
|||||||
info!("Listing background tasks, limit={}", limit);
|
info!("Listing background tasks, limit={}", limit);
|
||||||
|
|
||||||
let tasks: Vec<BackgroundTask> = sqlx::query_as::<_, BackgroundTask>(
|
let tasks: Vec<BackgroundTask> = sqlx::query_as::<_, BackgroundTask>(
|
||||||
"SELECT * FROM background_tasks ORDER BY created_at DESC LIMIT ?"
|
"SELECT * FROM background_tasks ORDER BY created_at DESC LIMIT ?",
|
||||||
)
|
)
|
||||||
.bind(limit)
|
.bind(limit)
|
||||||
.fetch_all(&*state.db)
|
.fetch_all(&*state.db)
|
||||||
@@ -53,22 +52,20 @@ async fn list_tasks(
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// GET /api/tasks/{id} — get a single background task
|
/// GET /api/tasks/{id} — get a single background task
|
||||||
async fn get_task(
|
async fn get_task(State(state): State<AppState>, Path(id): Path<String>) -> AppResult<Json<Value>> {
|
||||||
State(state): State<AppState>,
|
|
||||||
Path(id): Path<String>,
|
|
||||||
) -> AppResult<Json<Value>> {
|
|
||||||
info!("Getting background task: {}", id);
|
info!("Getting background task: {}", id);
|
||||||
|
|
||||||
let task: Option<BackgroundTask> = sqlx::query_as::<_, BackgroundTask>(
|
let task: Option<BackgroundTask> =
|
||||||
"SELECT * FROM background_tasks WHERE id = ?"
|
sqlx::query_as::<_, BackgroundTask>("SELECT * FROM background_tasks WHERE id = ?")
|
||||||
)
|
|
||||||
.bind(&id)
|
.bind(&id)
|
||||||
.fetch_optional(&*state.db)
|
.fetch_optional(&*state.db)
|
||||||
.await?;
|
.await?;
|
||||||
|
|
||||||
match task {
|
match task {
|
||||||
Some(t) => Ok(Json(json!({ "success": true, "task": t }))),
|
Some(t) => Ok(Json(json!({ "success": true, "task": t }))),
|
||||||
None => Ok(Json(json!({ "success": false, "message": "task not found" }))),
|
None => Ok(Json(
|
||||||
|
json!({ "success": false, "message": "task not found" }),
|
||||||
|
)),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -80,14 +77,17 @@ async fn stream_task_logs(
|
|||||||
info!("Streaming task logs for: {}", id);
|
info!("Streaming task logs for: {}", id);
|
||||||
|
|
||||||
// Check task exists
|
// Check task exists
|
||||||
let task: Option<BackgroundTask> = sqlx::query_as::<_, BackgroundTask>(
|
let task: Option<BackgroundTask> =
|
||||||
"SELECT * FROM background_tasks WHERE id = ?"
|
sqlx::query_as::<_, BackgroundTask>("SELECT * FROM background_tasks WHERE id = ?")
|
||||||
)
|
|
||||||
.bind(&id)
|
.bind(&id)
|
||||||
.fetch_optional(&*state.db)
|
.fetch_optional(&*state.db)
|
||||||
.await
|
.await
|
||||||
.map_err(|e| {
|
.map_err(|e| {
|
||||||
(axum::http::StatusCode::INTERNAL_SERVER_ERROR, format!("DB error: {e}")).into_response()
|
(
|
||||||
|
axum::http::StatusCode::INTERNAL_SERVER_ERROR,
|
||||||
|
format!("DB error: {e}"),
|
||||||
|
)
|
||||||
|
.into_response()
|
||||||
})?;
|
})?;
|
||||||
|
|
||||||
let task = match task {
|
let task = match task {
|
||||||
@@ -96,7 +96,8 @@ async fn stream_task_logs(
|
|||||||
return Err((
|
return Err((
|
||||||
axum::http::StatusCode::NOT_FOUND,
|
axum::http::StatusCode::NOT_FOUND,
|
||||||
json!({ "success": false, "message": "task not found" }).to_string(),
|
json!({ "success": false, "message": "task not found" }).to_string(),
|
||||||
).into_response());
|
)
|
||||||
|
.into_response());
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
|
|||||||
@@ -0,0 +1,433 @@
|
|||||||
|
// Update Center routes — deploy/rollback orchestration via helper API
|
||||||
|
|
||||||
|
use crate::{
|
||||||
|
error::{AppError, AppResult},
|
||||||
|
AppState,
|
||||||
|
};
|
||||||
|
use axum::{
|
||||||
|
extract::State,
|
||||||
|
routing::{get, post, put},
|
||||||
|
Json, Router,
|
||||||
|
};
|
||||||
|
use serde_json::{json, Value};
|
||||||
|
|
||||||
|
const UC_CONFIG_KEY: &str = "update_center_config";
|
||||||
|
|
||||||
|
const TASK_TYPE_DEPLOY: &str = "update-center-deploy";
|
||||||
|
const TASK_TYPE_ROLLBACK: &str = "update-center-rollback";
|
||||||
|
|
||||||
|
const DEDUPE_KEY_DEPLOY: &str = "update-center:deploy";
|
||||||
|
const DEDUPE_KEY_ROLLBACK: &str = "update-center:rollback";
|
||||||
|
|
||||||
|
pub fn routes() -> Router<AppState> {
|
||||||
|
Router::new()
|
||||||
|
.route("/api/update-center/status", get(get_status))
|
||||||
|
.route("/api/update-center/check", post(check_status))
|
||||||
|
.route("/api/update-center/config", put(update_config))
|
||||||
|
.route("/api/update-center/deploy", post(deploy_handler))
|
||||||
|
.route("/api/update-center/rollback", post(rollback_handler))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// GET /api/update-center/status
|
||||||
|
async fn get_status(State(state): State<AppState>) -> AppResult<Json<Value>> {
|
||||||
|
let config = load_uc_config(&state).await?;
|
||||||
|
let helper_token = get_helper_token();
|
||||||
|
|
||||||
|
let last_deploy = load_last_task(&state, TASK_TYPE_DEPLOY).await?;
|
||||||
|
let last_rollback = load_last_task(&state, TASK_TYPE_ROLLBACK).await?;
|
||||||
|
|
||||||
|
Ok(Json(json!({
|
||||||
|
"enabled": config.get("enabled").and_then(|v| v.as_bool()).unwrap_or(false),
|
||||||
|
"helperBaseUrl": config.get("helperBaseUrl"),
|
||||||
|
"namespace": config.get("namespace"),
|
||||||
|
"releaseName": config.get("releaseName"),
|
||||||
|
"helperTokenConfigured": !helper_token.is_empty(),
|
||||||
|
"lastDeploy": last_deploy,
|
||||||
|
"lastRollback": last_rollback,
|
||||||
|
})))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// POST /api/update-center/check — refresh status cache
|
||||||
|
async fn check_status(State(state): State<AppState>) -> AppResult<Json<Value>> {
|
||||||
|
// Same as get_status for now
|
||||||
|
get_status(State(state)).await
|
||||||
|
}
|
||||||
|
|
||||||
|
/// PUT /api/update-center/config
|
||||||
|
async fn update_config(
|
||||||
|
State(state): State<AppState>,
|
||||||
|
Json(body): Json<Value>,
|
||||||
|
) -> AppResult<Json<Value>> {
|
||||||
|
let serialized = serde_json::to_string(&body)
|
||||||
|
.map_err(|e| AppError::BadRequest(format!("Invalid config JSON: {e}")))?;
|
||||||
|
|
||||||
|
sqlx::query("INSERT OR REPLACE INTO settings (key, value) VALUES (?, ?)")
|
||||||
|
.bind(UC_CONFIG_KEY)
|
||||||
|
.bind(&serialized)
|
||||||
|
.execute(&*state.db)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
Ok(Json(json!({
|
||||||
|
"success": true,
|
||||||
|
"config": body,
|
||||||
|
})))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// POST /api/update-center/deploy
|
||||||
|
async fn deploy_handler(
|
||||||
|
State(state): State<AppState>,
|
||||||
|
Json(body): Json<Value>,
|
||||||
|
) -> AppResult<Json<Value>> {
|
||||||
|
let config = load_uc_config(&state).await?;
|
||||||
|
let helper_token = get_helper_token();
|
||||||
|
|
||||||
|
let helper_base = config
|
||||||
|
.get("helperBaseUrl")
|
||||||
|
.and_then(|v| v.as_str())
|
||||||
|
.unwrap_or("")
|
||||||
|
.trim();
|
||||||
|
let namespace = config
|
||||||
|
.get("namespace")
|
||||||
|
.and_then(|v| v.as_str())
|
||||||
|
.unwrap_or("")
|
||||||
|
.trim();
|
||||||
|
let chart_ref = config
|
||||||
|
.get("chartRef")
|
||||||
|
.and_then(|v| v.as_str())
|
||||||
|
.unwrap_or("")
|
||||||
|
.trim();
|
||||||
|
let image_repo = config
|
||||||
|
.get("imageRepository")
|
||||||
|
.and_then(|v| v.as_str())
|
||||||
|
.unwrap_or("")
|
||||||
|
.trim();
|
||||||
|
|
||||||
|
if helper_base.is_empty()
|
||||||
|
|| namespace.is_empty()
|
||||||
|
|| chart_ref.is_empty()
|
||||||
|
|| helper_token.is_empty()
|
||||||
|
{
|
||||||
|
return Err(AppError::BadRequest(
|
||||||
|
"Update center not fully configured".into(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
|
let target_tag = body
|
||||||
|
.get("targetTag")
|
||||||
|
.or_else(|| body.get("targetVersion"))
|
||||||
|
.and_then(|v| v.as_str())
|
||||||
|
.map(str::trim)
|
||||||
|
.unwrap_or("");
|
||||||
|
|
||||||
|
if target_tag.is_empty() {
|
||||||
|
return Err(AppError::BadRequest("targetTag is required".into()));
|
||||||
|
}
|
||||||
|
|
||||||
|
// Extract owned values for the spawn
|
||||||
|
let helper_base_owned = helper_base.trim_end_matches('/').to_string();
|
||||||
|
let namespace_owned = namespace.to_string();
|
||||||
|
let chart_ref_owned = chart_ref.to_string();
|
||||||
|
let image_repo_owned = image_repo.to_string();
|
||||||
|
let release_name = config
|
||||||
|
.get("releaseName")
|
||||||
|
.and_then(|v| v.as_str())
|
||||||
|
.map(str::trim)
|
||||||
|
.filter(|s| !s.is_empty())
|
||||||
|
.map(ToOwned::to_owned);
|
||||||
|
let target_tag_owned = target_tag.to_string();
|
||||||
|
let helper_token_owned = helper_token;
|
||||||
|
let title = format!("部署 {}", target_tag);
|
||||||
|
|
||||||
|
let task_id = uuid::Uuid::new_v4().to_string();
|
||||||
|
|
||||||
|
// Create background task
|
||||||
|
sqlx::query(
|
||||||
|
"INSERT INTO background_tasks (id, task_type, title, status, message, dedupe_key)
|
||||||
|
VALUES (?, ?, ?, 'running', ?, ?)",
|
||||||
|
)
|
||||||
|
.bind(&task_id)
|
||||||
|
.bind(TASK_TYPE_DEPLOY)
|
||||||
|
.bind(&title)
|
||||||
|
.bind(format!(
|
||||||
|
"Starting deploy of {} to {}/{}",
|
||||||
|
target_tag, namespace, chart_ref
|
||||||
|
))
|
||||||
|
.bind(DEDUPE_KEY_DEPLOY)
|
||||||
|
.execute(&*state.db)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
// Call deploy helper in background (fire-and-forget via tokio)
|
||||||
|
let client = reqwest::Client::new();
|
||||||
|
let deploy_url = format!("{}/api/deploy", helper_base_owned);
|
||||||
|
let task_id_clone = task_id.clone();
|
||||||
|
let db_pool = state.db.clone();
|
||||||
|
|
||||||
|
tokio::spawn(async move {
|
||||||
|
let payload = json!({
|
||||||
|
"namespace": namespace_owned,
|
||||||
|
"releaseName": release_name,
|
||||||
|
"chartRef": chart_ref_owned,
|
||||||
|
"imageRepository": image_repo_owned,
|
||||||
|
"targetTag": target_tag_owned,
|
||||||
|
});
|
||||||
|
|
||||||
|
match client
|
||||||
|
.post(&deploy_url)
|
||||||
|
.header("Authorization", format!("Bearer {}", helper_token_owned))
|
||||||
|
.json(&payload)
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
Ok(resp) => {
|
||||||
|
let status = resp.status();
|
||||||
|
let status_code = status.as_u16();
|
||||||
|
let body_text = resp.text().await.unwrap_or_default();
|
||||||
|
|
||||||
|
if status.is_success() {
|
||||||
|
let _ = sqlx::query(
|
||||||
|
"UPDATE background_tasks
|
||||||
|
SET status = 'succeeded', message = ?, finished_at = datetime('now')
|
||||||
|
WHERE id = ?",
|
||||||
|
)
|
||||||
|
.bind(format!("Deploy completed (HTTP {})", status_code))
|
||||||
|
.bind(&task_id_clone)
|
||||||
|
.execute(&*db_pool)
|
||||||
|
.await;
|
||||||
|
} else {
|
||||||
|
let _ = sqlx::query(
|
||||||
|
"UPDATE background_tasks
|
||||||
|
SET status = 'failed', error = ?, finished_at = datetime('now')
|
||||||
|
WHERE id = ?",
|
||||||
|
)
|
||||||
|
.bind(format!(
|
||||||
|
"Helper returned HTTP {}: {}",
|
||||||
|
status_code, body_text
|
||||||
|
))
|
||||||
|
.bind(&task_id_clone)
|
||||||
|
.execute(&*db_pool)
|
||||||
|
.await;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
let _ = sqlx::query(
|
||||||
|
"UPDATE background_tasks
|
||||||
|
SET status = 'failed', error = ?, finished_at = datetime('now')
|
||||||
|
WHERE id = ?",
|
||||||
|
)
|
||||||
|
.bind(format!("Helper unreachable: {}", e))
|
||||||
|
.bind(&task_id_clone)
|
||||||
|
.execute(&*db_pool)
|
||||||
|
.await;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
Ok(Json(json!({
|
||||||
|
"success": true,
|
||||||
|
"reused": false,
|
||||||
|
"task": {
|
||||||
|
"id": task_id,
|
||||||
|
"type": TASK_TYPE_DEPLOY,
|
||||||
|
"title": title,
|
||||||
|
"status": "running",
|
||||||
|
},
|
||||||
|
})))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// POST /api/update-center/rollback
|
||||||
|
async fn rollback_handler(
|
||||||
|
State(state): State<AppState>,
|
||||||
|
Json(body): Json<Value>,
|
||||||
|
) -> AppResult<Json<Value>> {
|
||||||
|
let config = load_uc_config(&state).await?;
|
||||||
|
let helper_token = get_helper_token();
|
||||||
|
|
||||||
|
let helper_base = config
|
||||||
|
.get("helperBaseUrl")
|
||||||
|
.and_then(|v| v.as_str())
|
||||||
|
.unwrap_or("")
|
||||||
|
.trim();
|
||||||
|
let namespace = config
|
||||||
|
.get("namespace")
|
||||||
|
.and_then(|v| v.as_str())
|
||||||
|
.unwrap_or("")
|
||||||
|
.trim();
|
||||||
|
|
||||||
|
if helper_base.is_empty() || namespace.is_empty() || helper_token.is_empty() {
|
||||||
|
return Err(AppError::BadRequest(
|
||||||
|
"Update center not fully configured".into(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
|
let target_revision = body
|
||||||
|
.get("targetRevision")
|
||||||
|
.and_then(|v| v.as_str())
|
||||||
|
.map(str::trim)
|
||||||
|
.unwrap_or("");
|
||||||
|
|
||||||
|
if target_revision.is_empty() {
|
||||||
|
return Err(AppError::BadRequest("targetRevision is required".into()));
|
||||||
|
}
|
||||||
|
|
||||||
|
let helper_base_owned = helper_base.trim_end_matches('/').to_string();
|
||||||
|
let namespace_owned = namespace.to_string();
|
||||||
|
let release_name = config
|
||||||
|
.get("releaseName")
|
||||||
|
.and_then(|v| v.as_str())
|
||||||
|
.map(str::trim)
|
||||||
|
.filter(|s| !s.is_empty())
|
||||||
|
.map(ToOwned::to_owned);
|
||||||
|
let target_revision_owned = target_revision.to_string();
|
||||||
|
let helper_token_owned = helper_token;
|
||||||
|
let title = format!("回退 {}", target_revision);
|
||||||
|
|
||||||
|
let task_id = uuid::Uuid::new_v4().to_string();
|
||||||
|
|
||||||
|
sqlx::query(
|
||||||
|
"INSERT INTO background_tasks (id, task_type, title, status, message, dedupe_key)
|
||||||
|
VALUES (?, ?, ?, 'running', ?, ?)",
|
||||||
|
)
|
||||||
|
.bind(&task_id)
|
||||||
|
.bind(TASK_TYPE_ROLLBACK)
|
||||||
|
.bind(&title)
|
||||||
|
.bind(format!(
|
||||||
|
"Starting rollback to {} in {}",
|
||||||
|
target_revision, namespace
|
||||||
|
))
|
||||||
|
.bind(DEDUPE_KEY_ROLLBACK)
|
||||||
|
.execute(&*state.db)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
let client = reqwest::Client::new();
|
||||||
|
let rollback_url = format!("{}/api/rollback", helper_base_owned);
|
||||||
|
let task_id_clone = task_id.clone();
|
||||||
|
let db_pool = state.db.clone();
|
||||||
|
|
||||||
|
tokio::spawn(async move {
|
||||||
|
let payload = json!({
|
||||||
|
"namespace": namespace_owned,
|
||||||
|
"releaseName": release_name,
|
||||||
|
"targetRevision": target_revision_owned,
|
||||||
|
});
|
||||||
|
|
||||||
|
match client
|
||||||
|
.post(&rollback_url)
|
||||||
|
.header("Authorization", format!("Bearer {}", helper_token_owned))
|
||||||
|
.json(&payload)
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
Ok(resp) => {
|
||||||
|
let status = resp.status();
|
||||||
|
let status_code = status.as_u16();
|
||||||
|
let body_text = resp.text().await.unwrap_or_default();
|
||||||
|
|
||||||
|
if status.is_success() {
|
||||||
|
let _ = sqlx::query(
|
||||||
|
"UPDATE background_tasks
|
||||||
|
SET status = 'succeeded', message = ?, finished_at = datetime('now')
|
||||||
|
WHERE id = ?",
|
||||||
|
)
|
||||||
|
.bind(format!("Rollback completed (HTTP {})", status_code))
|
||||||
|
.bind(&task_id_clone)
|
||||||
|
.execute(&*db_pool)
|
||||||
|
.await;
|
||||||
|
} else {
|
||||||
|
let _ = sqlx::query(
|
||||||
|
"UPDATE background_tasks
|
||||||
|
SET status = 'failed', error = ?, finished_at = datetime('now')
|
||||||
|
WHERE id = ?",
|
||||||
|
)
|
||||||
|
.bind(format!(
|
||||||
|
"Helper returned HTTP {}: {}",
|
||||||
|
status_code, body_text
|
||||||
|
))
|
||||||
|
.bind(&task_id_clone)
|
||||||
|
.execute(&*db_pool)
|
||||||
|
.await;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
let _ = sqlx::query(
|
||||||
|
"UPDATE background_tasks
|
||||||
|
SET status = 'failed', error = ?, finished_at = datetime('now')
|
||||||
|
WHERE id = ?",
|
||||||
|
)
|
||||||
|
.bind(format!("Helper unreachable: {}", e))
|
||||||
|
.bind(&task_id_clone)
|
||||||
|
.execute(&*db_pool)
|
||||||
|
.await;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
Ok(Json(json!({
|
||||||
|
"success": true,
|
||||||
|
"reused": false,
|
||||||
|
"task": {
|
||||||
|
"id": task_id,
|
||||||
|
"type": TASK_TYPE_ROLLBACK,
|
||||||
|
"title": title,
|
||||||
|
"status": "running",
|
||||||
|
},
|
||||||
|
})))
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Helpers ────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
fn get_helper_token() -> String {
|
||||||
|
std::env::var("UPDATE_CENTER_HELPER_TOKEN")
|
||||||
|
.or_else(|_| std::env::var("DEPLOY_HELPER_TOKEN"))
|
||||||
|
.unwrap_or_default()
|
||||||
|
.trim()
|
||||||
|
.to_string()
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn load_uc_config(state: &AppState) -> AppResult<Value> {
|
||||||
|
let row: Option<(String,)> = sqlx::query_as("SELECT value FROM settings WHERE key = ?")
|
||||||
|
.bind(UC_CONFIG_KEY)
|
||||||
|
.fetch_optional(&*state.db)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
match row {
|
||||||
|
Some((v,)) => serde_json::from_str(&v)
|
||||||
|
.map_err(|e| AppError::BadRequest(format!("Invalid update center config JSON: {e}"))),
|
||||||
|
None => Ok(json!({})),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn load_last_task(state: &AppState, task_type: &str) -> AppResult<Option<Value>> {
|
||||||
|
let row: Option<(
|
||||||
|
String,
|
||||||
|
String,
|
||||||
|
Option<String>,
|
||||||
|
Option<String>,
|
||||||
|
String,
|
||||||
|
Option<String>,
|
||||||
|
Option<String>,
|
||||||
|
)> = sqlx::query_as(
|
||||||
|
"SELECT id, status, message, error, created_at, started_at, finished_at
|
||||||
|
FROM background_tasks
|
||||||
|
WHERE task_type = ?
|
||||||
|
ORDER BY created_at DESC
|
||||||
|
LIMIT 1",
|
||||||
|
)
|
||||||
|
.bind(task_type)
|
||||||
|
.fetch_optional(&*state.db)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
Ok(row.map(
|
||||||
|
|(id, status, message, error, created_at, started_at, finished_at)| {
|
||||||
|
json!({
|
||||||
|
"id": id,
|
||||||
|
"type": task_type,
|
||||||
|
"status": status,
|
||||||
|
"message": message,
|
||||||
|
"error": error,
|
||||||
|
"createdAt": created_at,
|
||||||
|
"startedAt": started_at,
|
||||||
|
"finishedAt": finished_at,
|
||||||
|
})
|
||||||
|
},
|
||||||
|
))
|
||||||
|
}
|
||||||
@@ -1 +1,160 @@
|
|||||||
// Routes - user
|
use crate::{
|
||||||
|
auth::jwt::validate_jwt,
|
||||||
|
db::models::{downstream_api_key::DownstreamApiKey, user::User},
|
||||||
|
error::{AppError, AppResult},
|
||||||
|
AppState,
|
||||||
|
};
|
||||||
|
use axum::{
|
||||||
|
extract::{Path, State},
|
||||||
|
http::{HeaderMap, StatusCode},
|
||||||
|
routing::{delete, get, post, put},
|
||||||
|
Json, Router,
|
||||||
|
};
|
||||||
|
use serde::Deserialize;
|
||||||
|
use serde_json::{json, Value};
|
||||||
|
|
||||||
|
pub fn routes() -> Router<AppState> {
|
||||||
|
Router::new()
|
||||||
|
.route("/v1/user/profile", get(get_profile).put(update_profile))
|
||||||
|
.route("/v1/user/api-keys", get(list_api_keys).post(create_api_key))
|
||||||
|
.route("/v1/user/api-keys/{id}", delete(delete_api_key))
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Deserialize)]
|
||||||
|
struct UpdateProfilePayload {
|
||||||
|
username: Option<String>,
|
||||||
|
email: Option<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Deserialize)]
|
||||||
|
struct CreateApiKeyPayload {
|
||||||
|
name: String,
|
||||||
|
description: Option<String>,
|
||||||
|
expires_at: Option<String>,
|
||||||
|
max_requests: Option<i64>,
|
||||||
|
max_cost: Option<f64>,
|
||||||
|
}
|
||||||
|
|
||||||
|
fn current_user_id(headers: &HeaderMap, state: &AppState) -> AppResult<i64> {
|
||||||
|
let auth = headers
|
||||||
|
.get("Authorization")
|
||||||
|
.and_then(|v| v.to_str().ok())
|
||||||
|
.and_then(|v| v.strip_prefix("Bearer "))
|
||||||
|
.ok_or_else(|| AppError::Unauthorized("Missing or invalid Authorization header".into()))?;
|
||||||
|
|
||||||
|
let claims = validate_jwt(&state.config.auth.jwt_secret, auth)
|
||||||
|
.map_err(|_| AppError::Unauthorized("Invalid or expired token".into()))?;
|
||||||
|
|
||||||
|
claims
|
||||||
|
.sub
|
||||||
|
.parse()
|
||||||
|
.map_err(|_| AppError::Unauthorized("Invalid user ID in token".into()))
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn get_profile(State(state): State<AppState>, headers: HeaderMap) -> AppResult<Json<Value>> {
|
||||||
|
let user_id = current_user_id(&headers, &state)?;
|
||||||
|
let user: User = sqlx::query_as("SELECT * FROM users WHERE id = ?")
|
||||||
|
.bind(user_id)
|
||||||
|
.fetch_optional(&*state.db)
|
||||||
|
.await?
|
||||||
|
.ok_or_else(|| AppError::NotFound("User not found".into()))?;
|
||||||
|
|
||||||
|
Ok(Json(json!({ "user": user })))
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn update_profile(
|
||||||
|
State(state): State<AppState>,
|
||||||
|
headers: HeaderMap,
|
||||||
|
Json(payload): Json<UpdateProfilePayload>,
|
||||||
|
) -> AppResult<Json<Value>> {
|
||||||
|
let user_id = current_user_id(&headers, &state)?;
|
||||||
|
let current: User = sqlx::query_as("SELECT * FROM users WHERE id = ?")
|
||||||
|
.bind(user_id)
|
||||||
|
.fetch_optional(&*state.db)
|
||||||
|
.await?
|
||||||
|
.ok_or_else(|| AppError::NotFound("User not found".into()))?;
|
||||||
|
|
||||||
|
let username = payload.username.unwrap_or(current.username);
|
||||||
|
let email = payload.email.unwrap_or(current.email);
|
||||||
|
|
||||||
|
sqlx::query(
|
||||||
|
"UPDATE users
|
||||||
|
SET username = ?, email = ?, updated_at = datetime('now')
|
||||||
|
WHERE id = ?",
|
||||||
|
)
|
||||||
|
.bind(username)
|
||||||
|
.bind(email)
|
||||||
|
.bind(user_id)
|
||||||
|
.execute(&*state.db)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
get_profile(State(state), headers).await
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn list_api_keys(
|
||||||
|
State(state): State<AppState>,
|
||||||
|
headers: HeaderMap,
|
||||||
|
) -> AppResult<Json<Value>> {
|
||||||
|
let user_id = current_user_id(&headers, &state)?;
|
||||||
|
let keys: Vec<DownstreamApiKey> = sqlx::query_as(
|
||||||
|
"SELECT * FROM downstream_api_keys WHERE user_id = ? ORDER BY created_at DESC",
|
||||||
|
)
|
||||||
|
.bind(user_id)
|
||||||
|
.fetch_all(&*state.db)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
Ok(Json(json!({ "keys": keys })))
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn create_api_key(
|
||||||
|
State(state): State<AppState>,
|
||||||
|
headers: HeaderMap,
|
||||||
|
Json(payload): Json<CreateApiKeyPayload>,
|
||||||
|
) -> AppResult<(StatusCode, Json<Value>)> {
|
||||||
|
let user_id = current_user_id(&headers, &state)?;
|
||||||
|
if payload.name.trim().is_empty() {
|
||||||
|
return Err(AppError::BadRequest("name is required".into()));
|
||||||
|
}
|
||||||
|
|
||||||
|
let token = format!("sk-user-{}", uuid::Uuid::new_v4().simple());
|
||||||
|
let result = sqlx::query(
|
||||||
|
"INSERT INTO downstream_api_keys
|
||||||
|
(name, key, description, enabled, expires_at, max_requests, max_cost, user_id)
|
||||||
|
VALUES (?, ?, ?, 1, ?, ?, ?, ?)",
|
||||||
|
)
|
||||||
|
.bind(payload.name)
|
||||||
|
.bind(&token)
|
||||||
|
.bind(payload.description)
|
||||||
|
.bind(payload.expires_at)
|
||||||
|
.bind(payload.max_requests)
|
||||||
|
.bind(payload.max_cost)
|
||||||
|
.bind(user_id)
|
||||||
|
.execute(&*state.db)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
let key: DownstreamApiKey = sqlx::query_as("SELECT * FROM downstream_api_keys WHERE id = ?")
|
||||||
|
.bind(result.last_insert_rowid())
|
||||||
|
.fetch_one(&*state.db)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
Ok((StatusCode::CREATED, Json(json!({ "key": key }))))
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn delete_api_key(
|
||||||
|
State(state): State<AppState>,
|
||||||
|
headers: HeaderMap,
|
||||||
|
Path(id): Path<i64>,
|
||||||
|
) -> AppResult<StatusCode> {
|
||||||
|
let user_id = current_user_id(&headers, &state)?;
|
||||||
|
let result = sqlx::query("DELETE FROM downstream_api_keys WHERE id = ? AND user_id = ?")
|
||||||
|
.bind(id)
|
||||||
|
.bind(user_id)
|
||||||
|
.execute(&*state.db)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
if result.rows_affected() == 0 {
|
||||||
|
return Err(AppError::NotFound("API key not found".into()));
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(StatusCode::NO_CONTENT)
|
||||||
|
}
|
||||||
|
|||||||
@@ -19,8 +19,14 @@ use tracing::info;
|
|||||||
|
|
||||||
pub fn routes() -> Router<AppState> {
|
pub fn routes() -> Router<AppState> {
|
||||||
Router::new()
|
Router::new()
|
||||||
.route("/api/user-api-keys", get(list_user_keys).post(create_user_key))
|
.route(
|
||||||
.route("/api/user-api-keys/{id}", patch(update_user_key).delete(delete_user_key))
|
"/api/user-api-keys",
|
||||||
|
get(list_user_keys).post(create_user_key),
|
||||||
|
)
|
||||||
|
.route(
|
||||||
|
"/api/user-api-keys/{id}",
|
||||||
|
patch(update_user_key).delete(delete_user_key),
|
||||||
|
)
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Extract user_id and role from JWT in Authorization header.
|
/// Extract user_id and role from JWT in Authorization header.
|
||||||
@@ -51,7 +57,7 @@ async fn list_user_keys(
|
|||||||
info!("Listing user API keys for user {}", user_id);
|
info!("Listing user API keys for user {}", user_id);
|
||||||
|
|
||||||
let keys: Vec<DownstreamApiKey> = sqlx::query_as::<_, DownstreamApiKey>(
|
let keys: Vec<DownstreamApiKey> = sqlx::query_as::<_, DownstreamApiKey>(
|
||||||
"SELECT * FROM downstream_api_keys WHERE user_id = ? ORDER BY created_at DESC"
|
"SELECT * FROM downstream_api_keys WHERE user_id = ? ORDER BY created_at DESC",
|
||||||
)
|
)
|
||||||
.bind(user_id)
|
.bind(user_id)
|
||||||
.fetch_all(&*state.db)
|
.fetch_all(&*state.db)
|
||||||
@@ -82,9 +88,17 @@ async fn create_user_key(
|
|||||||
.get("key")
|
.get("key")
|
||||||
.and_then(|v| v.as_str())
|
.and_then(|v| v.as_str())
|
||||||
.map(|s| s.to_string())
|
.map(|s| s.to_string())
|
||||||
.unwrap_or_else(|| format!("sk-user-{}", uuid::Uuid::new_v4().to_string().replace('-', "")));
|
.unwrap_or_else(|| {
|
||||||
|
format!(
|
||||||
|
"sk-user-{}",
|
||||||
|
uuid::Uuid::new_v4().to_string().replace('-', "")
|
||||||
|
)
|
||||||
|
});
|
||||||
|
|
||||||
let enabled = body.get("enabled").and_then(|v| v.as_bool()).unwrap_or(true);
|
let enabled = body
|
||||||
|
.get("enabled")
|
||||||
|
.and_then(|v| v.as_bool())
|
||||||
|
.unwrap_or(true);
|
||||||
let expires_at = body.get("expiresAt").and_then(|v| v.as_str());
|
let expires_at = body.get("expiresAt").and_then(|v| v.as_str());
|
||||||
let max_cost = body.get("maxCost").and_then(|v| v.as_f64());
|
let max_cost = body.get("maxCost").and_then(|v| v.as_f64());
|
||||||
let max_requests = body.get("maxRequests").and_then(|v| v.as_i64());
|
let max_requests = body.get("maxRequests").and_then(|v| v.as_i64());
|
||||||
@@ -99,7 +113,7 @@ async fn create_user_key(
|
|||||||
(name, key, description, tags, enabled, expires_at, max_cost, max_requests,
|
(name, key, description, tags, enabled, expires_at, max_cost, max_requests,
|
||||||
supported_models, allowed_route_ids, site_weight_multipliers,
|
supported_models, allowed_route_ids, site_weight_multipliers,
|
||||||
excluded_site_ids, excluded_credential_refs, user_id)
|
excluded_site_ids, excluded_credential_refs, user_id)
|
||||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)"#
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)"#,
|
||||||
)
|
)
|
||||||
.bind(name)
|
.bind(name)
|
||||||
.bind(&key_str)
|
.bind(&key_str)
|
||||||
@@ -118,9 +132,8 @@ async fn create_user_key(
|
|||||||
.execute(&*state.db)
|
.execute(&*state.db)
|
||||||
.await?;
|
.await?;
|
||||||
|
|
||||||
let new_key: DownstreamApiKey = sqlx::query_as::<_, DownstreamApiKey>(
|
let new_key: DownstreamApiKey =
|
||||||
"SELECT * FROM downstream_api_keys WHERE id = ?"
|
sqlx::query_as::<_, DownstreamApiKey>("SELECT * FROM downstream_api_keys WHERE id = ?")
|
||||||
)
|
|
||||||
.bind(result.last_insert_rowid())
|
.bind(result.last_insert_rowid())
|
||||||
.fetch_one(&*state.db)
|
.fetch_one(&*state.db)
|
||||||
.await?;
|
.await?;
|
||||||
@@ -139,9 +152,8 @@ async fn update_user_key(
|
|||||||
info!("Updating user API key {} for user {}", id, user_id);
|
info!("Updating user API key {} for user {}", id, user_id);
|
||||||
|
|
||||||
// Verify ownership
|
// Verify ownership
|
||||||
let existing: Option<DownstreamApiKey> = sqlx::query_as(
|
let existing: Option<DownstreamApiKey> =
|
||||||
"SELECT * FROM downstream_api_keys WHERE id = ? AND user_id = ?"
|
sqlx::query_as("SELECT * FROM downstream_api_keys WHERE id = ? AND user_id = ?")
|
||||||
)
|
|
||||||
.bind(id)
|
.bind(id)
|
||||||
.bind(user_id)
|
.bind(user_id)
|
||||||
.fetch_optional(&*state.db)
|
.fetch_optional(&*state.db)
|
||||||
@@ -157,35 +169,54 @@ async fn update_user_key(
|
|||||||
let mut bind_idx = 0;
|
let mut bind_idx = 0;
|
||||||
|
|
||||||
if let Some(v) = body.get("name").and_then(|v| v.as_str()) {
|
if let Some(v) = body.get("name").and_then(|v| v.as_str()) {
|
||||||
bind_idx += 1; sets.push(format!("name = ?{}", bind_idx)); params.push(v.to_string());
|
bind_idx += 1;
|
||||||
|
sets.push(format!("name = ?{}", bind_idx));
|
||||||
|
params.push(v.to_string());
|
||||||
}
|
}
|
||||||
if let Some(v) = body.get("description") {
|
if let Some(v) = body.get("description") {
|
||||||
bind_idx += 1; sets.push(format!("description = ?{}", bind_idx));
|
bind_idx += 1;
|
||||||
|
sets.push(format!("description = ?{}", bind_idx));
|
||||||
params.push(v.as_str().unwrap_or_default().to_string());
|
params.push(v.as_str().unwrap_or_default().to_string());
|
||||||
}
|
}
|
||||||
if let Some(v) = body.get("tags") {
|
if let Some(v) = body.get("tags") {
|
||||||
bind_idx += 1; sets.push(format!("tags = ?{}", bind_idx)); params.push(v.to_string());
|
bind_idx += 1;
|
||||||
|
sets.push(format!("tags = ?{}", bind_idx));
|
||||||
|
params.push(v.to_string());
|
||||||
}
|
}
|
||||||
if let Some(v) = body.get("enabled").and_then(|v| v.as_bool()) {
|
if let Some(v) = body.get("enabled").and_then(|v| v.as_bool()) {
|
||||||
bind_idx += 1; sets.push(format!("enabled = ?{}", bind_idx)); params.push(if v { "1" } else { "0" }.to_string());
|
bind_idx += 1;
|
||||||
|
sets.push(format!("enabled = ?{}", bind_idx));
|
||||||
|
params.push(if v { "1" } else { "0" }.to_string());
|
||||||
}
|
}
|
||||||
if let Some(v) = body.get("expiresAt").and_then(|v| v.as_str()) {
|
if let Some(v) = body.get("expiresAt").and_then(|v| v.as_str()) {
|
||||||
bind_idx += 1; sets.push(format!("expires_at = ?{}", bind_idx)); params.push(v.to_string());
|
bind_idx += 1;
|
||||||
|
sets.push(format!("expires_at = ?{}", bind_idx));
|
||||||
|
params.push(v.to_string());
|
||||||
}
|
}
|
||||||
if let Some(v) = body.get("maxCost").and_then(|v| v.as_f64()) {
|
if let Some(v) = body.get("maxCost").and_then(|v| v.as_f64()) {
|
||||||
bind_idx += 1; sets.push(format!("max_cost = ?{}", bind_idx)); params.push(v.to_string());
|
bind_idx += 1;
|
||||||
|
sets.push(format!("max_cost = ?{}", bind_idx));
|
||||||
|
params.push(v.to_string());
|
||||||
}
|
}
|
||||||
if let Some(v) = body.get("maxRequests").and_then(|v| v.as_i64()) {
|
if let Some(v) = body.get("maxRequests").and_then(|v| v.as_i64()) {
|
||||||
bind_idx += 1; sets.push(format!("max_requests = ?{}", bind_idx)); params.push(v.to_string());
|
bind_idx += 1;
|
||||||
|
sets.push(format!("max_requests = ?{}", bind_idx));
|
||||||
|
params.push(v.to_string());
|
||||||
}
|
}
|
||||||
if let Some(v) = body.get("supportedModels") {
|
if let Some(v) = body.get("supportedModels") {
|
||||||
bind_idx += 1; sets.push(format!("supported_models = ?{}", bind_idx)); params.push(v.to_string());
|
bind_idx += 1;
|
||||||
|
sets.push(format!("supported_models = ?{}", bind_idx));
|
||||||
|
params.push(v.to_string());
|
||||||
}
|
}
|
||||||
if let Some(v) = body.get("allowedRouteIds") {
|
if let Some(v) = body.get("allowedRouteIds") {
|
||||||
bind_idx += 1; sets.push(format!("allowed_route_ids = ?{}", bind_idx)); params.push(v.to_string());
|
bind_idx += 1;
|
||||||
|
sets.push(format!("allowed_route_ids = ?{}", bind_idx));
|
||||||
|
params.push(v.to_string());
|
||||||
}
|
}
|
||||||
if let Some(v) = body.get("excludedSiteIds") {
|
if let Some(v) = body.get("excludedSiteIds") {
|
||||||
bind_idx += 1; sets.push(format!("excluded_site_ids = ?{}", bind_idx)); params.push(v.to_string());
|
bind_idx += 1;
|
||||||
|
sets.push(format!("excluded_site_ids = ?{}", bind_idx));
|
||||||
|
params.push(v.to_string());
|
||||||
}
|
}
|
||||||
|
|
||||||
if !sets.is_empty() {
|
if !sets.is_empty() {
|
||||||
@@ -207,9 +238,8 @@ async fn update_user_key(
|
|||||||
q.execute(&*state.db).await?;
|
q.execute(&*state.db).await?;
|
||||||
}
|
}
|
||||||
|
|
||||||
let updated: DownstreamApiKey = sqlx::query_as::<_, DownstreamApiKey>(
|
let updated: DownstreamApiKey =
|
||||||
"SELECT * FROM downstream_api_keys WHERE id = ?"
|
sqlx::query_as::<_, DownstreamApiKey>("SELECT * FROM downstream_api_keys WHERE id = ?")
|
||||||
)
|
|
||||||
.bind(id)
|
.bind(id)
|
||||||
.fetch_one(&*state.db)
|
.fetch_one(&*state.db)
|
||||||
.await?;
|
.await?;
|
||||||
|
|||||||
@@ -0,0 +1,517 @@
|
|||||||
|
use axum::{
|
||||||
|
body::Body,
|
||||||
|
extract::State,
|
||||||
|
http::{header, HeaderMap, Response, StatusCode},
|
||||||
|
response::IntoResponse,
|
||||||
|
routing::post,
|
||||||
|
Json, Router,
|
||||||
|
};
|
||||||
|
use serde_json::{json, Value};
|
||||||
|
use tracing::info;
|
||||||
|
|
||||||
|
use crate::{
|
||||||
|
error::{AppError, AppResult},
|
||||||
|
AppState,
|
||||||
|
};
|
||||||
|
|
||||||
|
const SCRIPT_ANALYSIS_PROMPT: &str = r#"你是一个 AI 视频生成助手。分析用户提供的脚本,提取以下结构化信息,以 JSON 格式返回:
|
||||||
|
|
||||||
|
{
|
||||||
|
"title": "视频标题",
|
||||||
|
"characters": [
|
||||||
|
{
|
||||||
|
"name": "角色名",
|
||||||
|
"gender": "男/女",
|
||||||
|
"age": "年龄描述",
|
||||||
|
"appearance": "外貌详细描述(发型、脸型、服装、体型等)",
|
||||||
|
"personality": "性格特征"
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"scenes": [
|
||||||
|
{
|
||||||
|
"id": 1,
|
||||||
|
"description": "场景描述",
|
||||||
|
"environment": "环境详细描述(时间、地点、天气、光线等)",
|
||||||
|
"characters": ["出现的角色名"],
|
||||||
|
"action": "场景中发生的动作",
|
||||||
|
"dialogue": [
|
||||||
|
{"character": "角色名", "line": "台词", "emotion": "情感语气"}
|
||||||
|
]
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"style": "整体视觉风格(动画/写实/像素等)"
|
||||||
|
}
|
||||||
|
|
||||||
|
只返回 JSON,不要其他文字。"#;
|
||||||
|
|
||||||
|
const CHARACTER_PROMPT_GEN: &str = r#"根据角色描述生成 4 组稳定的图像生成提示词(正面、背面、左侧面、右侧面),用于 AI 图像生成模型。
|
||||||
|
每组提示词应包含:人物姿态、服装、光照、背景。
|
||||||
|
返回 JSON 格式:
|
||||||
|
{
|
||||||
|
"front": "正面提示词",
|
||||||
|
"back": "背面提示词",
|
||||||
|
"left": "左侧面提示词",
|
||||||
|
"right": "右侧面提示词"
|
||||||
|
}
|
||||||
|
只返回 JSON。"#;
|
||||||
|
|
||||||
|
const SCENE_PROMPT_GEN: &str = r#"根据场景描述生成图像生成提示词,用于 AI 图像生成模型创建场景背景。
|
||||||
|
包含:环境、光照、色调、氛围。
|
||||||
|
返回 JSON:
|
||||||
|
{
|
||||||
|
"prompt": "场景生成提示词",
|
||||||
|
"negative_prompt": "负面提示词"
|
||||||
|
}
|
||||||
|
只返回 JSON。"#;
|
||||||
|
|
||||||
|
pub fn routes() -> Router<AppState> {
|
||||||
|
Router::new()
|
||||||
|
.route("/api/video-agent/analyze", post(analyze))
|
||||||
|
.route(
|
||||||
|
"/api/video-agent/character-prompts",
|
||||||
|
post(character_prompts),
|
||||||
|
)
|
||||||
|
.route(
|
||||||
|
"/api/video-agent/generate-character",
|
||||||
|
post(generate_character),
|
||||||
|
)
|
||||||
|
.route("/api/video-agent/scene-prompt", post(scene_prompt))
|
||||||
|
.route("/api/video-agent/generate-scene", post(generate_scene))
|
||||||
|
.route("/api/video-agent/generate-tts", post(generate_tts))
|
||||||
|
.route("/api/video-agent/full-workflow", post(full_workflow))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn upstream_base() -> String {
|
||||||
|
std::env::var("BOOSAPI_UPSTREAM_BASE")
|
||||||
|
.unwrap_or_else(|_| "http://127.0.0.1:4000/v1".to_string())
|
||||||
|
.trim_end_matches('/')
|
||||||
|
.to_string()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn upstream_key() -> String {
|
||||||
|
std::env::var("BOOSAPI_UPSTREAM_KEY").unwrap_or_else(|_| "sk-boos4721".to_string())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn check_auth(headers: &HeaderMap, state: &AppState) -> AppResult<()> {
|
||||||
|
let expected = format!("Bearer {}", state.config.auth.auth_token);
|
||||||
|
|
||||||
|
let actual = headers
|
||||||
|
.get(header::AUTHORIZATION)
|
||||||
|
.and_then(|v| v.to_str().ok())
|
||||||
|
.unwrap_or("");
|
||||||
|
|
||||||
|
if actual != expected {
|
||||||
|
return Err(AppError::BadRequest("unauthorized".to_string()));
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn required_str<'a>(body: &'a Value, key: &str) -> AppResult<&'a str> {
|
||||||
|
body.get(key)
|
||||||
|
.and_then(Value::as_str)
|
||||||
|
.filter(|s| !s.trim().is_empty())
|
||||||
|
.ok_or_else(|| AppError::BadRequest(format!("{key} is required")))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn required_value<'a>(body: &'a Value, key: &str) -> AppResult<&'a Value> {
|
||||||
|
body.get(key)
|
||||||
|
.filter(|v| !v.is_null())
|
||||||
|
.ok_or_else(|| AppError::BadRequest(format!("{key} is required")))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn strip_json_markers(s: &str) -> String {
|
||||||
|
let mut text = s.trim().to_string();
|
||||||
|
|
||||||
|
if text.starts_with("```json") {
|
||||||
|
text = text.trim_start_matches("```json").to_string();
|
||||||
|
} else if text.starts_with("```") {
|
||||||
|
text = text.trim_start_matches("```").to_string();
|
||||||
|
}
|
||||||
|
|
||||||
|
if text.ends_with("```") {
|
||||||
|
text = text.trim_end_matches("```").to_string();
|
||||||
|
}
|
||||||
|
|
||||||
|
text.trim().to_string()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn parse_llm_json(raw: &str) -> AppResult<Value> {
|
||||||
|
let cleaned = strip_json_markers(raw);
|
||||||
|
serde_json::from_str::<Value>(&cleaned).map_err(|e| {
|
||||||
|
AppError::Internal(anyhow::anyhow!(
|
||||||
|
"failed to parse LLM JSON response: {}; response: {}",
|
||||||
|
e,
|
||||||
|
cleaned
|
||||||
|
))
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn call_llm(prompt: &str, system: &str) -> AppResult<String> {
|
||||||
|
let client = reqwest::Client::new();
|
||||||
|
let url = format!("{}/chat/completions", upstream_base());
|
||||||
|
|
||||||
|
info!("video-agent calling LLM upstream: {}", url);
|
||||||
|
|
||||||
|
let resp = client
|
||||||
|
.post(url)
|
||||||
|
.bearer_auth(upstream_key())
|
||||||
|
.json(&json!({
|
||||||
|
"model": "gpt-5.5",
|
||||||
|
"messages": [
|
||||||
|
{ "role": "system", "content": system },
|
||||||
|
{ "role": "user", "content": prompt }
|
||||||
|
],
|
||||||
|
"temperature": 0.7
|
||||||
|
}))
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.map_err(|e| AppError::Internal(anyhow::anyhow!("LLM request failed: {}", e)))?;
|
||||||
|
|
||||||
|
let status = resp.status();
|
||||||
|
let data: Value = resp
|
||||||
|
.json()
|
||||||
|
.await
|
||||||
|
.map_err(|e| AppError::Internal(anyhow::anyhow!("LLM response parse failed: {}", e)))?;
|
||||||
|
|
||||||
|
if !status.is_success() {
|
||||||
|
return Err(AppError::Internal(anyhow::anyhow!(
|
||||||
|
"LLM upstream returned {}: {}",
|
||||||
|
status,
|
||||||
|
data
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(data
|
||||||
|
.get("choices")
|
||||||
|
.and_then(Value::as_array)
|
||||||
|
.and_then(|choices| choices.first())
|
||||||
|
.and_then(|choice| choice.get("message"))
|
||||||
|
.and_then(|message| message.get("content"))
|
||||||
|
.and_then(Value::as_str)
|
||||||
|
.unwrap_or("")
|
||||||
|
.to_string())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn call_image_gen(prompt: &str) -> AppResult<String> {
|
||||||
|
let client = reqwest::Client::new();
|
||||||
|
let url = format!("{}/images/generations", upstream_base());
|
||||||
|
|
||||||
|
info!("video-agent calling image upstream: {}", url);
|
||||||
|
|
||||||
|
let resp = client
|
||||||
|
.post(url)
|
||||||
|
.bearer_auth(upstream_key())
|
||||||
|
.json(&json!({
|
||||||
|
"model": "gpt-image-2",
|
||||||
|
"prompt": prompt,
|
||||||
|
"n": 1,
|
||||||
|
"size": "1024x1024"
|
||||||
|
}))
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.map_err(|e| AppError::Internal(anyhow::anyhow!("image request failed: {}", e)))?;
|
||||||
|
|
||||||
|
let status = resp.status();
|
||||||
|
let data: Value = resp
|
||||||
|
.json()
|
||||||
|
.await
|
||||||
|
.map_err(|e| AppError::Internal(anyhow::anyhow!("image response parse failed: {}", e)))?;
|
||||||
|
|
||||||
|
if !status.is_success() {
|
||||||
|
return Err(AppError::Internal(anyhow::anyhow!(
|
||||||
|
"image upstream returned {}: {}",
|
||||||
|
status,
|
||||||
|
data
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(data
|
||||||
|
.get("data")
|
||||||
|
.and_then(Value::as_array)
|
||||||
|
.and_then(|items| items.first())
|
||||||
|
.and_then(|item| item.get("url"))
|
||||||
|
.and_then(Value::as_str)
|
||||||
|
.unwrap_or("")
|
||||||
|
.to_string())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn call_tts(text: &str, voice: &str) -> AppResult<Vec<u8>> {
|
||||||
|
let client = reqwest::Client::new();
|
||||||
|
let url = format!("{}/audio/speech", upstream_base());
|
||||||
|
|
||||||
|
info!("video-agent calling TTS upstream: {}", url);
|
||||||
|
|
||||||
|
let resp = client
|
||||||
|
.post(url)
|
||||||
|
.bearer_auth(upstream_key())
|
||||||
|
.json(&json!({
|
||||||
|
"model": "tts-1",
|
||||||
|
"input": text,
|
||||||
|
"voice": voice
|
||||||
|
}))
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.map_err(|e| AppError::Internal(anyhow::anyhow!("TTS request failed: {}", e)))?;
|
||||||
|
|
||||||
|
let status = resp.status();
|
||||||
|
|
||||||
|
if !status.is_success() {
|
||||||
|
let body = resp.text().await.unwrap_or_default();
|
||||||
|
return Err(AppError::Internal(anyhow::anyhow!(
|
||||||
|
"TTS upstream returned {}: {}",
|
||||||
|
status,
|
||||||
|
body
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
|
||||||
|
let bytes = resp
|
||||||
|
.bytes()
|
||||||
|
.await
|
||||||
|
.map_err(|e| AppError::Internal(anyhow::anyhow!("TTS response read failed: {}", e)))?;
|
||||||
|
|
||||||
|
Ok(bytes.to_vec())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn analyze(
|
||||||
|
State(state): State<AppState>,
|
||||||
|
headers: HeaderMap,
|
||||||
|
Json(body): Json<Value>,
|
||||||
|
) -> AppResult<Json<Value>> {
|
||||||
|
check_auth(&headers, &state)?;
|
||||||
|
|
||||||
|
let script = required_str(&body, "script")?;
|
||||||
|
|
||||||
|
let result = call_llm(script, SCRIPT_ANALYSIS_PROMPT).await?;
|
||||||
|
let parsed = parse_llm_json(&result)?;
|
||||||
|
|
||||||
|
Ok(Json(json!({
|
||||||
|
"success": true,
|
||||||
|
"data": parsed
|
||||||
|
})))
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn character_prompts(
|
||||||
|
State(state): State<AppState>,
|
||||||
|
headers: HeaderMap,
|
||||||
|
Json(body): Json<Value>,
|
||||||
|
) -> AppResult<Json<Value>> {
|
||||||
|
check_auth(&headers, &state)?;
|
||||||
|
|
||||||
|
let character = required_value(&body, "character")?;
|
||||||
|
|
||||||
|
let name = character
|
||||||
|
.get("name")
|
||||||
|
.and_then(Value::as_str)
|
||||||
|
.unwrap_or_default();
|
||||||
|
let appearance = character
|
||||||
|
.get("appearance")
|
||||||
|
.and_then(Value::as_str)
|
||||||
|
.unwrap_or_default();
|
||||||
|
let gender = character
|
||||||
|
.get("gender")
|
||||||
|
.and_then(Value::as_str)
|
||||||
|
.unwrap_or_default();
|
||||||
|
let age = character
|
||||||
|
.get("age")
|
||||||
|
.and_then(Value::as_str)
|
||||||
|
.unwrap_or_default();
|
||||||
|
|
||||||
|
let desc = format!("{name}: {appearance}, {gender}, {age}");
|
||||||
|
let result = call_llm(&desc, CHARACTER_PROMPT_GEN).await?;
|
||||||
|
let prompts = parse_llm_json(&result)?;
|
||||||
|
|
||||||
|
Ok(Json(json!({
|
||||||
|
"success": true,
|
||||||
|
"data": prompts
|
||||||
|
})))
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn generate_character(
|
||||||
|
State(state): State<AppState>,
|
||||||
|
headers: HeaderMap,
|
||||||
|
Json(body): Json<Value>,
|
||||||
|
) -> AppResult<Json<Value>> {
|
||||||
|
check_auth(&headers, &state)?;
|
||||||
|
|
||||||
|
let prompts = required_value(&body, "prompts")?;
|
||||||
|
let prompt_obj = prompts
|
||||||
|
.as_object()
|
||||||
|
.ok_or_else(|| AppError::BadRequest("prompts must be an object".to_string()))?;
|
||||||
|
|
||||||
|
let mut results = serde_json::Map::new();
|
||||||
|
|
||||||
|
for (angle, prompt_value) in prompt_obj {
|
||||||
|
let prompt = prompt_value.as_str().ok_or_else(|| {
|
||||||
|
AppError::BadRequest(format!("prompt for angle '{angle}' must be a string"))
|
||||||
|
})?;
|
||||||
|
|
||||||
|
let url = call_image_gen(prompt).await?;
|
||||||
|
results.insert(angle.clone(), Value::String(url));
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(Json(json!({
|
||||||
|
"success": true,
|
||||||
|
"data": Value::Object(results)
|
||||||
|
})))
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn scene_prompt(
|
||||||
|
State(state): State<AppState>,
|
||||||
|
headers: HeaderMap,
|
||||||
|
Json(body): Json<Value>,
|
||||||
|
) -> AppResult<Json<Value>> {
|
||||||
|
check_auth(&headers, &state)?;
|
||||||
|
|
||||||
|
let scene = required_value(&body, "scene")?;
|
||||||
|
|
||||||
|
let environment = scene
|
||||||
|
.get("environment")
|
||||||
|
.and_then(Value::as_str)
|
||||||
|
.unwrap_or_default();
|
||||||
|
let description = scene
|
||||||
|
.get("description")
|
||||||
|
.and_then(Value::as_str)
|
||||||
|
.unwrap_or_default();
|
||||||
|
|
||||||
|
let desc = format!("{environment}. {description}");
|
||||||
|
let result = call_llm(&desc, SCENE_PROMPT_GEN).await?;
|
||||||
|
let prompt = parse_llm_json(&result)?;
|
||||||
|
|
||||||
|
Ok(Json(json!({
|
||||||
|
"success": true,
|
||||||
|
"data": prompt
|
||||||
|
})))
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn generate_scene(
|
||||||
|
State(state): State<AppState>,
|
||||||
|
headers: HeaderMap,
|
||||||
|
Json(body): Json<Value>,
|
||||||
|
) -> AppResult<Json<Value>> {
|
||||||
|
check_auth(&headers, &state)?;
|
||||||
|
|
||||||
|
let prompt = required_str(&body, "prompt")?;
|
||||||
|
let url = call_image_gen(prompt).await?;
|
||||||
|
|
||||||
|
Ok(Json(json!({
|
||||||
|
"success": true,
|
||||||
|
"data": {
|
||||||
|
"url": url
|
||||||
|
}
|
||||||
|
})))
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn generate_tts(
|
||||||
|
State(state): State<AppState>,
|
||||||
|
headers: HeaderMap,
|
||||||
|
Json(body): Json<Value>,
|
||||||
|
) -> AppResult<Response<Body>> {
|
||||||
|
check_auth(&headers, &state)?;
|
||||||
|
|
||||||
|
let text = required_str(&body, "text")?;
|
||||||
|
let voice = body
|
||||||
|
.get("voice")
|
||||||
|
.and_then(Value::as_str)
|
||||||
|
.filter(|s| !s.trim().is_empty())
|
||||||
|
.unwrap_or("alloy");
|
||||||
|
|
||||||
|
let audio = call_tts(text, voice).await?;
|
||||||
|
|
||||||
|
Ok((
|
||||||
|
StatusCode::OK,
|
||||||
|
[(header::CONTENT_TYPE, "audio/mpeg")],
|
||||||
|
audio,
|
||||||
|
)
|
||||||
|
.into_response())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn full_workflow(
|
||||||
|
State(state): State<AppState>,
|
||||||
|
headers: HeaderMap,
|
||||||
|
Json(body): Json<Value>,
|
||||||
|
) -> AppResult<Json<Value>> {
|
||||||
|
check_auth(&headers, &state)?;
|
||||||
|
|
||||||
|
let script = required_str(&body, "script")?;
|
||||||
|
let requested_style = body.get("style").and_then(Value::as_str);
|
||||||
|
|
||||||
|
let analysis_raw = call_llm(script, SCRIPT_ANALYSIS_PROMPT).await?;
|
||||||
|
let analysis = parse_llm_json(&analysis_raw)?;
|
||||||
|
|
||||||
|
let mut characters = Vec::new();
|
||||||
|
|
||||||
|
if let Some(chars) = analysis.get("characters").and_then(Value::as_array) {
|
||||||
|
for character in chars {
|
||||||
|
let name = character
|
||||||
|
.get("name")
|
||||||
|
.and_then(Value::as_str)
|
||||||
|
.unwrap_or_default();
|
||||||
|
let appearance = character
|
||||||
|
.get("appearance")
|
||||||
|
.and_then(Value::as_str)
|
||||||
|
.unwrap_or_default();
|
||||||
|
let gender = character
|
||||||
|
.get("gender")
|
||||||
|
.and_then(Value::as_str)
|
||||||
|
.unwrap_or_default();
|
||||||
|
let age = character
|
||||||
|
.get("age")
|
||||||
|
.and_then(Value::as_str)
|
||||||
|
.unwrap_or_default();
|
||||||
|
|
||||||
|
let desc = format!("{name}: {appearance}, {gender}, {age}");
|
||||||
|
let prompts_raw = call_llm(&desc, CHARACTER_PROMPT_GEN).await?;
|
||||||
|
let prompts = parse_llm_json(&prompts_raw)?;
|
||||||
|
|
||||||
|
let mut item = character.clone();
|
||||||
|
if let Some(obj) = item.as_object_mut() {
|
||||||
|
obj.insert("imagePrompts".to_string(), prompts);
|
||||||
|
}
|
||||||
|
|
||||||
|
characters.push(item);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut scenes = Vec::new();
|
||||||
|
|
||||||
|
if let Some(scene_items) = analysis.get("scenes").and_then(Value::as_array) {
|
||||||
|
for scene in scene_items {
|
||||||
|
let environment = scene
|
||||||
|
.get("environment")
|
||||||
|
.and_then(Value::as_str)
|
||||||
|
.unwrap_or_default();
|
||||||
|
let description = scene
|
||||||
|
.get("description")
|
||||||
|
.and_then(Value::as_str)
|
||||||
|
.unwrap_or_default();
|
||||||
|
|
||||||
|
let desc = format!("{environment}. {description}");
|
||||||
|
let scene_prompt_raw = call_llm(&desc, SCENE_PROMPT_GEN).await?;
|
||||||
|
let image_prompt = parse_llm_json(&scene_prompt_raw)?;
|
||||||
|
|
||||||
|
let mut item = scene.clone();
|
||||||
|
if let Some(obj) = item.as_object_mut() {
|
||||||
|
obj.insert("imagePrompt".to_string(), image_prompt);
|
||||||
|
}
|
||||||
|
|
||||||
|
scenes.push(item);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let style = analysis
|
||||||
|
.get("style")
|
||||||
|
.cloned()
|
||||||
|
.or_else(|| requested_style.map(|s| Value::String(s.to_string())))
|
||||||
|
.unwrap_or(Value::Null);
|
||||||
|
|
||||||
|
Ok(Json(json!({
|
||||||
|
"success": true,
|
||||||
|
"data": {
|
||||||
|
"title": analysis.get("title").cloned().unwrap_or(Value::Null),
|
||||||
|
"style": style,
|
||||||
|
"characters": characters,
|
||||||
|
"scenes": scenes
|
||||||
|
}
|
||||||
|
})))
|
||||||
|
}
|
||||||
@@ -1 +1,42 @@
|
|||||||
// Services - announcement
|
use crate::{db::models::site_announcement::SiteAnnouncement, error::AppResult};
|
||||||
|
use sqlx::SqlitePool;
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct AnnouncementService {
|
||||||
|
db: Arc<SqlitePool>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl AnnouncementService {
|
||||||
|
pub fn new(db: Arc<SqlitePool>) -> Self {
|
||||||
|
Self { db }
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn get_active_announcements(&self) -> AppResult<Vec<SiteAnnouncement>> {
|
||||||
|
self.expire_old_announcements().await?;
|
||||||
|
let rows = sqlx::query_as::<_, SiteAnnouncement>(
|
||||||
|
"SELECT *
|
||||||
|
FROM site_announcements
|
||||||
|
WHERE dismissed_at IS NULL
|
||||||
|
AND (starts_at IS NULL OR starts_at <= datetime('now'))
|
||||||
|
AND (ends_at IS NULL OR ends_at > datetime('now'))
|
||||||
|
ORDER BY first_seen_at DESC, id DESC",
|
||||||
|
)
|
||||||
|
.fetch_all(&*self.db)
|
||||||
|
.await?;
|
||||||
|
Ok(rows)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn expire_old_announcements(&self) -> AppResult<u64> {
|
||||||
|
let result = sqlx::query(
|
||||||
|
"UPDATE site_announcements
|
||||||
|
SET dismissed_at = datetime('now')
|
||||||
|
WHERE dismissed_at IS NULL
|
||||||
|
AND ends_at IS NOT NULL
|
||||||
|
AND ends_at <= datetime('now')",
|
||||||
|
)
|
||||||
|
.execute(&*self.db)
|
||||||
|
.await?;
|
||||||
|
Ok(result.rows_affected())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -1 +1,201 @@
|
|||||||
// Services - cron
|
use crate::{error::AppResult, proxy::token_router::TokenRouter};
|
||||||
|
use sqlx::SqlitePool;
|
||||||
|
use std::sync::{
|
||||||
|
atomic::{AtomicBool, Ordering},
|
||||||
|
Arc,
|
||||||
|
};
|
||||||
|
use tokio::{
|
||||||
|
task::JoinHandle,
|
||||||
|
time::{interval, Duration},
|
||||||
|
};
|
||||||
|
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct CronScheduler {
|
||||||
|
db: Arc<SqlitePool>,
|
||||||
|
running: Arc<AtomicBool>,
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct CronSchedulerHandle {
|
||||||
|
running: Arc<AtomicBool>,
|
||||||
|
tasks: Vec<JoinHandle<()>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl CronScheduler {
|
||||||
|
pub fn new(db: Arc<SqlitePool>) -> Self {
|
||||||
|
Self {
|
||||||
|
db,
|
||||||
|
running: Arc::new(AtomicBool::new(false)),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn start(&self) -> CronSchedulerHandle {
|
||||||
|
self.running.store(true, Ordering::SeqCst);
|
||||||
|
|
||||||
|
let tasks = vec![
|
||||||
|
self.spawn_periodic(
|
||||||
|
"usage aggregation",
|
||||||
|
Duration::from_secs(300),
|
||||||
|
Self::aggregate_usage,
|
||||||
|
),
|
||||||
|
self.spawn_periodic(
|
||||||
|
"route refresh",
|
||||||
|
Duration::from_secs(600),
|
||||||
|
Self::refresh_routes,
|
||||||
|
),
|
||||||
|
self.spawn_periodic("model probe", Duration::from_secs(900), Self::probe_models),
|
||||||
|
self.spawn_periodic(
|
||||||
|
"channel recovery",
|
||||||
|
Duration::from_secs(120),
|
||||||
|
Self::recover_channels,
|
||||||
|
),
|
||||||
|
];
|
||||||
|
|
||||||
|
CronSchedulerHandle {
|
||||||
|
running: self.running.clone(),
|
||||||
|
tasks,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn spawn_periodic<F, Fut>(&self, name: &'static str, every: Duration, task: F) -> JoinHandle<()>
|
||||||
|
where
|
||||||
|
F: Fn(Arc<SqlitePool>) -> Fut + Send + Sync + Copy + 'static,
|
||||||
|
Fut: std::future::Future<Output = AppResult<()>> + Send + 'static,
|
||||||
|
{
|
||||||
|
let db = self.db.clone();
|
||||||
|
let running = self.running.clone();
|
||||||
|
tokio::spawn(async move {
|
||||||
|
let mut ticker = interval(every);
|
||||||
|
while running.load(Ordering::SeqCst) {
|
||||||
|
ticker.tick().await;
|
||||||
|
if !running.load(Ordering::SeqCst) {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
if let Err(err) = task(db.clone()).await {
|
||||||
|
tracing::warn!(%err, task = name, "cron task failed");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn aggregate_usage(db: Arc<SqlitePool>) -> AppResult<()> {
|
||||||
|
sqlx::query(
|
||||||
|
"UPDATE downstream_api_keys
|
||||||
|
SET updated_at = datetime('now')
|
||||||
|
WHERE used_requests > 0 OR used_cost > 0",
|
||||||
|
)
|
||||||
|
.execute(&*db)
|
||||||
|
.await?;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn refresh_routes(db: Arc<SqlitePool>) -> AppResult<()> {
|
||||||
|
sqlx::query(
|
||||||
|
"UPDATE token_routes
|
||||||
|
SET decision_refreshed_at = datetime('now'), updated_at = datetime('now')
|
||||||
|
WHERE enabled = 1",
|
||||||
|
)
|
||||||
|
.execute(&*db)
|
||||||
|
.await?;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn probe_models(db: Arc<SqlitePool>) -> AppResult<()> {
|
||||||
|
let router = TokenRouter::new(db.clone());
|
||||||
|
let models: Vec<String> = sqlx::query_scalar(
|
||||||
|
"SELECT model_pattern FROM token_routes WHERE enabled = 1 ORDER BY id LIMIT 25",
|
||||||
|
)
|
||||||
|
.fetch_all(&*db)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
for model in models {
|
||||||
|
let probe_model = model.replace('*', "probe");
|
||||||
|
let _ = router.resolve_selected(&probe_model).await;
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn recover_channels(db: Arc<SqlitePool>) -> AppResult<()> {
|
||||||
|
sqlx::query(
|
||||||
|
"UPDATE route_channels
|
||||||
|
SET cooldown_until = NULL, cooldown_level = 0, consecutive_fail_count = 0
|
||||||
|
WHERE cooldown_until IS NOT NULL AND cooldown_until <= datetime('now')",
|
||||||
|
)
|
||||||
|
.execute(&*db)
|
||||||
|
.await?;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl CronSchedulerHandle {
|
||||||
|
pub fn stop(&self) {
|
||||||
|
self.running.store(false, Ordering::SeqCst);
|
||||||
|
for task in &self.tasks {
|
||||||
|
task.abort();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Drop for CronSchedulerHandle {
|
||||||
|
fn drop(&mut self) {
|
||||||
|
self.running.store(false, Ordering::SeqCst);
|
||||||
|
for task in &self.tasks {
|
||||||
|
task.abort();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
use sqlx::sqlite::SqlitePoolOptions;
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn scheduler_starts_and_stops_background_tasks() {
|
||||||
|
let pool = SqlitePoolOptions::new()
|
||||||
|
.max_connections(1)
|
||||||
|
.connect("sqlite::memory:")
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
sqlx::query(
|
||||||
|
"CREATE TABLE downstream_api_keys (
|
||||||
|
id INTEGER PRIMARY KEY,
|
||||||
|
used_requests INTEGER DEFAULT 0,
|
||||||
|
used_cost REAL DEFAULT 0,
|
||||||
|
updated_at TEXT
|
||||||
|
)",
|
||||||
|
)
|
||||||
|
.execute(&pool)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
sqlx::query(
|
||||||
|
"CREATE TABLE token_routes (
|
||||||
|
id INTEGER PRIMARY KEY,
|
||||||
|
model_pattern TEXT NOT NULL,
|
||||||
|
enabled INTEGER DEFAULT 1,
|
||||||
|
decision_refreshed_at TEXT,
|
||||||
|
updated_at TEXT
|
||||||
|
)",
|
||||||
|
)
|
||||||
|
.execute(&pool)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
sqlx::query(
|
||||||
|
"CREATE TABLE route_channels (
|
||||||
|
id INTEGER PRIMARY KEY,
|
||||||
|
cooldown_until TEXT,
|
||||||
|
cooldown_level INTEGER DEFAULT 0,
|
||||||
|
consecutive_fail_count INTEGER DEFAULT 0
|
||||||
|
)",
|
||||||
|
)
|
||||||
|
.execute(&pool)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
let scheduler = CronScheduler::new(Arc::new(pool));
|
||||||
|
let handle = scheduler.start();
|
||||||
|
assert!(handle.running.load(Ordering::SeqCst));
|
||||||
|
handle.stop();
|
||||||
|
assert!(!handle.running.load(Ordering::SeqCst));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -1 +1,5 @@
|
|||||||
// Services module
|
pub mod announcement;
|
||||||
|
pub mod cron;
|
||||||
|
pub mod oauth;
|
||||||
|
pub mod shared;
|
||||||
|
pub mod update_center;
|
||||||
|
|||||||
@@ -1 +1,207 @@
|
|||||||
// Services - OAuth
|
use crate::error::{AppError, AppResult};
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
use serde_json::Value;
|
||||||
|
use sqlx::{FromRow, SqlitePool};
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct OAuthService {
|
||||||
|
db: Arc<SqlitePool>,
|
||||||
|
client: reqwest::Client,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize, FromRow)]
|
||||||
|
pub struct OAuthRouteUnit {
|
||||||
|
pub id: i64,
|
||||||
|
pub site_id: i64,
|
||||||
|
pub provider: String,
|
||||||
|
pub name: String,
|
||||||
|
pub strategy: String,
|
||||||
|
pub enabled: Option<bool>,
|
||||||
|
pub created_at: Option<String>,
|
||||||
|
pub updated_at: Option<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct RefreshedToken {
|
||||||
|
pub account_id: i64,
|
||||||
|
pub access_token: String,
|
||||||
|
pub expires_at: Option<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl OAuthService {
|
||||||
|
pub fn new(db: Arc<SqlitePool>) -> Self {
|
||||||
|
Self {
|
||||||
|
db,
|
||||||
|
client: reqwest::Client::new(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn refresh_token(&self, account_id: i64, site_id: i64) -> AppResult<RefreshedToken> {
|
||||||
|
let row: Option<(String, Option<String>, Option<String>)> = sqlx::query_as(
|
||||||
|
"SELECT access_token, api_token, extra_config
|
||||||
|
FROM accounts
|
||||||
|
WHERE id = ? AND site_id = ? AND oauth_provider IS NOT NULL",
|
||||||
|
)
|
||||||
|
.bind(account_id)
|
||||||
|
.bind(site_id)
|
||||||
|
.fetch_optional(&*self.db)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
let (current_access_token, refresh_token, extra_config) =
|
||||||
|
row.ok_or_else(|| AppError::NotFound("OAuth account not found".into()))?;
|
||||||
|
|
||||||
|
if !self.token_expired(extra_config.as_deref()) {
|
||||||
|
return Ok(RefreshedToken {
|
||||||
|
account_id,
|
||||||
|
access_token: current_access_token,
|
||||||
|
expires_at: token_expiry(extra_config.as_deref()),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
let refresh_token = refresh_token
|
||||||
|
.filter(|token| !token.trim().is_empty())
|
||||||
|
.ok_or_else(|| AppError::BadRequest("OAuth account has no refresh token".into()))?;
|
||||||
|
|
||||||
|
let mut refreshed = self.perform_refresh(site_id, &refresh_token).await?;
|
||||||
|
refreshed.account_id = account_id;
|
||||||
|
sqlx::query(
|
||||||
|
"UPDATE accounts
|
||||||
|
SET access_token = ?, extra_config = ?, updated_at = datetime('now')
|
||||||
|
WHERE id = ? AND site_id = ?",
|
||||||
|
)
|
||||||
|
.bind(&refreshed.access_token)
|
||||||
|
.bind(
|
||||||
|
serde_json::to_string(&serde_json::json!({
|
||||||
|
"expires_at": refreshed.expires_at,
|
||||||
|
"refreshed_at": chrono::Utc::now().to_rfc3339(),
|
||||||
|
}))
|
||||||
|
.map_err(|err| AppError::Internal(err.into()))?,
|
||||||
|
)
|
||||||
|
.bind(account_id)
|
||||||
|
.bind(site_id)
|
||||||
|
.execute(&*self.db)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
Ok(refreshed)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn token_expired(&self, extra_config: Option<&str>) -> bool {
|
||||||
|
token_expired(extra_config)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn list_route_units(&self, site_id: Option<i64>) -> AppResult<Vec<OAuthRouteUnit>> {
|
||||||
|
if let Some(site_id) = site_id {
|
||||||
|
sqlx::query_as("SELECT * FROM oauth_route_units WHERE site_id = ? ORDER BY id")
|
||||||
|
.bind(site_id)
|
||||||
|
.fetch_all(&*self.db)
|
||||||
|
.await
|
||||||
|
.map_err(Into::into)
|
||||||
|
} else {
|
||||||
|
sqlx::query_as("SELECT * FROM oauth_route_units ORDER BY id")
|
||||||
|
.fetch_all(&*self.db)
|
||||||
|
.await
|
||||||
|
.map_err(Into::into)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn create_route_unit(
|
||||||
|
&self,
|
||||||
|
site_id: i64,
|
||||||
|
provider: &str,
|
||||||
|
name: &str,
|
||||||
|
strategy: Option<&str>,
|
||||||
|
) -> AppResult<OAuthRouteUnit> {
|
||||||
|
let result = sqlx::query(
|
||||||
|
"INSERT INTO oauth_route_units (site_id, provider, name, strategy)
|
||||||
|
VALUES (?, ?, ?, ?)",
|
||||||
|
)
|
||||||
|
.bind(site_id)
|
||||||
|
.bind(provider)
|
||||||
|
.bind(name)
|
||||||
|
.bind(strategy.unwrap_or("round_robin"))
|
||||||
|
.execute(&*self.db)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
sqlx::query_as("SELECT * FROM oauth_route_units WHERE id = ?")
|
||||||
|
.bind(result.last_insert_rowid())
|
||||||
|
.fetch_one(&*self.db)
|
||||||
|
.await
|
||||||
|
.map_err(Into::into)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn perform_refresh(
|
||||||
|
&self,
|
||||||
|
site_id: i64,
|
||||||
|
refresh_token: &str,
|
||||||
|
) -> AppResult<RefreshedToken> {
|
||||||
|
let token_endpoint: Option<String> = sqlx::query_scalar(
|
||||||
|
"SELECT json_extract(custom_headers, '$.oauth_token_url') FROM sites WHERE id = ?",
|
||||||
|
)
|
||||||
|
.bind(site_id)
|
||||||
|
.fetch_optional(&*self.db)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
let Some(token_endpoint) = token_endpoint.filter(|url| !url.trim().is_empty()) else {
|
||||||
|
return Err(AppError::BadRequest(
|
||||||
|
"OAuth token endpoint is not configured for this site".into(),
|
||||||
|
));
|
||||||
|
};
|
||||||
|
|
||||||
|
let response: Value = self
|
||||||
|
.client
|
||||||
|
.post(token_endpoint)
|
||||||
|
.form(&[
|
||||||
|
("grant_type", "refresh_token"),
|
||||||
|
("refresh_token", refresh_token),
|
||||||
|
])
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.map_err(|err| AppError::UpstreamError(format!("OAuth refresh failed: {err}")))?
|
||||||
|
.json()
|
||||||
|
.await
|
||||||
|
.map_err(|err| AppError::UpstreamError(format!("Invalid OAuth response: {err}")))?;
|
||||||
|
|
||||||
|
let access_token = response
|
||||||
|
.get("access_token")
|
||||||
|
.and_then(Value::as_str)
|
||||||
|
.ok_or_else(|| AppError::UpstreamError("OAuth response missing access_token".into()))?
|
||||||
|
.to_string();
|
||||||
|
let expires_at = response
|
||||||
|
.get("expires_at")
|
||||||
|
.and_then(Value::as_str)
|
||||||
|
.map(str::to_string);
|
||||||
|
|
||||||
|
Ok(RefreshedToken {
|
||||||
|
account_id: 0,
|
||||||
|
access_token,
|
||||||
|
expires_at,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn token_expired(extra_config: Option<&str>) -> bool {
|
||||||
|
let Some(extra_config) = extra_config else {
|
||||||
|
return true;
|
||||||
|
};
|
||||||
|
let Ok(value) = serde_json::from_str::<Value>(extra_config) else {
|
||||||
|
return true;
|
||||||
|
};
|
||||||
|
let Some(expires_at) = value.get("expires_at").and_then(Value::as_str) else {
|
||||||
|
return true;
|
||||||
|
};
|
||||||
|
chrono::DateTime::parse_from_rfc3339(expires_at)
|
||||||
|
.map(|expires_at| expires_at <= chrono::Utc::now() + chrono::Duration::minutes(5))
|
||||||
|
.unwrap_or(true)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn token_expiry(extra_config: Option<&str>) -> Option<String> {
|
||||||
|
extra_config
|
||||||
|
.and_then(|config| serde_json::from_str::<Value>(config).ok())
|
||||||
|
.and_then(|value| {
|
||||||
|
value
|
||||||
|
.get("expires_at")
|
||||||
|
.and_then(Value::as_str)
|
||||||
|
.map(str::to_string)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|||||||
@@ -1 +1,45 @@
|
|||||||
// Services - shared
|
use crate::{error::AppError, error::AppResult};
|
||||||
|
use sqlx::SqlitePool;
|
||||||
|
|
||||||
|
pub async fn enforce_key_request_limit(
|
||||||
|
db: &SqlitePool,
|
||||||
|
key_id: i64,
|
||||||
|
increment: i64,
|
||||||
|
) -> AppResult<()> {
|
||||||
|
let row: Option<(Option<i64>, Option<i64>, Option<bool>)> = sqlx::query_as(
|
||||||
|
"SELECT max_requests, used_requests, enabled
|
||||||
|
FROM downstream_api_keys
|
||||||
|
WHERE id = ?",
|
||||||
|
)
|
||||||
|
.bind(key_id)
|
||||||
|
.fetch_optional(db)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
let (max_requests, used_requests, enabled) =
|
||||||
|
row.ok_or_else(|| AppError::Unauthorized("API key not found".into()))?;
|
||||||
|
if !enabled.unwrap_or(true) {
|
||||||
|
return Err(AppError::Forbidden("API key is disabled".into()));
|
||||||
|
}
|
||||||
|
if let Some(max_requests) = max_requests {
|
||||||
|
if used_requests.unwrap_or(0) + increment > max_requests {
|
||||||
|
return Err(AppError::RateLimited);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn increment_key_usage(db: &SqlitePool, key_id: i64, cost: f64) -> AppResult<()> {
|
||||||
|
sqlx::query(
|
||||||
|
"UPDATE downstream_api_keys
|
||||||
|
SET used_requests = COALESCE(used_requests, 0) + 1,
|
||||||
|
used_cost = COALESCE(used_cost, 0) + ?,
|
||||||
|
last_used_at = datetime('now'),
|
||||||
|
updated_at = datetime('now')
|
||||||
|
WHERE id = ?",
|
||||||
|
)
|
||||||
|
.bind(cost)
|
||||||
|
.bind(key_id)
|
||||||
|
.execute(db)
|
||||||
|
.await?;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|||||||
@@ -1 +1,115 @@
|
|||||||
// Services - update center
|
use crate::error::{AppError, AppResult};
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
use serde_json::Value;
|
||||||
|
use sqlx::SqlitePool;
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
const UPDATE_CONFIG_KEY: &str = "update_center_config";
|
||||||
|
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct UpdateCenter {
|
||||||
|
db: Arc<SqlitePool>,
|
||||||
|
client: reqwest::Client,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct UpdateInfo {
|
||||||
|
pub current_version: String,
|
||||||
|
pub latest_version: Option<String>,
|
||||||
|
pub update_available: bool,
|
||||||
|
pub release_notes: Option<String>,
|
||||||
|
pub download_url: Option<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl UpdateCenter {
|
||||||
|
pub fn new(db: Arc<SqlitePool>) -> Self {
|
||||||
|
Self {
|
||||||
|
db,
|
||||||
|
client: reqwest::Client::new(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn check_for_updates(&self) -> AppResult<UpdateInfo> {
|
||||||
|
let config = self.load_config().await?;
|
||||||
|
let current_version = env!("CARGO_PKG_VERSION").to_string();
|
||||||
|
let Some(endpoint) = config
|
||||||
|
.get("check_url")
|
||||||
|
.and_then(Value::as_str)
|
||||||
|
.filter(|url| !url.trim().is_empty())
|
||||||
|
else {
|
||||||
|
return Ok(UpdateInfo {
|
||||||
|
current_version,
|
||||||
|
latest_version: None,
|
||||||
|
update_available: false,
|
||||||
|
release_notes: None,
|
||||||
|
download_url: None,
|
||||||
|
});
|
||||||
|
};
|
||||||
|
|
||||||
|
let payload: Value = self
|
||||||
|
.client
|
||||||
|
.get(endpoint)
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.map_err(|err| AppError::UpstreamError(format!("Update check failed: {err}")))?
|
||||||
|
.json()
|
||||||
|
.await
|
||||||
|
.map_err(|err| AppError::UpstreamError(format!("Invalid update response: {err}")))?;
|
||||||
|
|
||||||
|
let latest_version = payload
|
||||||
|
.get("version")
|
||||||
|
.or_else(|| payload.get("latestVersion"))
|
||||||
|
.and_then(Value::as_str)
|
||||||
|
.map(str::to_string);
|
||||||
|
let update_available = latest_version
|
||||||
|
.as_deref()
|
||||||
|
.map(|latest| latest != current_version)
|
||||||
|
.unwrap_or(false);
|
||||||
|
|
||||||
|
Ok(UpdateInfo {
|
||||||
|
current_version,
|
||||||
|
latest_version,
|
||||||
|
update_available,
|
||||||
|
release_notes: payload
|
||||||
|
.get("release_notes")
|
||||||
|
.or_else(|| payload.get("releaseNotes"))
|
||||||
|
.and_then(Value::as_str)
|
||||||
|
.map(str::to_string),
|
||||||
|
download_url: payload
|
||||||
|
.get("download_url")
|
||||||
|
.or_else(|| payload.get("downloadUrl"))
|
||||||
|
.and_then(Value::as_str)
|
||||||
|
.map(str::to_string),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn apply_update(&self, version: &str) -> AppResult<String> {
|
||||||
|
let task_id = format!("update-{}", uuid::Uuid::new_v4());
|
||||||
|
sqlx::query(
|
||||||
|
"INSERT INTO background_tasks
|
||||||
|
(id, task_type, title, status, message, dedupe_key, created_at, updated_at)
|
||||||
|
VALUES (?, 'update-center-deploy', ?, 'pending', ?, ?, datetime('now'), datetime('now'))",
|
||||||
|
)
|
||||||
|
.bind(&task_id)
|
||||||
|
.bind(format!("Apply BoosAPI {version}"))
|
||||||
|
.bind("Update request accepted; an external supervisor should perform download and restart")
|
||||||
|
.bind(format!("update-center:deploy:{version}"))
|
||||||
|
.execute(&*self.db)
|
||||||
|
.await?;
|
||||||
|
Ok(task_id)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn load_config(&self) -> AppResult<Value> {
|
||||||
|
let raw: Option<String> = sqlx::query_scalar("SELECT value FROM settings WHERE key = ?")
|
||||||
|
.bind(UPDATE_CONFIG_KEY)
|
||||||
|
.fetch_optional(&*self.db)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
raw.as_deref()
|
||||||
|
.map(serde_json::from_str)
|
||||||
|
.transpose()
|
||||||
|
.map_err(|err| AppError::BadRequest(format!("Invalid update center config: {err}")))?
|
||||||
|
.map(Ok)
|
||||||
|
.unwrap_or_else(|| Ok(Value::Object(Default::default())))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -0,0 +1,43 @@
|
|||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
use serde_json::Value;
|
||||||
|
|
||||||
|
/// Anthropic Messages API compatibility switches.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct AnthropicMessagesCompatibility {
|
||||||
|
/// Value suitable for the `anthropic-version` header when callers need it.
|
||||||
|
pub anthropic_version: Option<String>,
|
||||||
|
/// Messages API requires `max_tokens`; use this when canonical input omits it.
|
||||||
|
pub default_max_tokens: Option<i32>,
|
||||||
|
/// Emit system content as an Anthropic content-block array instead of a string.
|
||||||
|
pub system_as_array: bool,
|
||||||
|
/// Copy canonical metadata to the Anthropic request `metadata` field.
|
||||||
|
pub preserve_metadata: bool,
|
||||||
|
/// Extra provider-specific request fields to merge into the formatted payload.
|
||||||
|
pub extra_request_fields: Option<Value>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for AnthropicMessagesCompatibility {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self {
|
||||||
|
anthropic_version: Some("2023-06-01".to_string()),
|
||||||
|
default_max_tokens: Some(1024),
|
||||||
|
system_as_array: false,
|
||||||
|
preserve_metadata: true,
|
||||||
|
extra_request_fields: None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_default_compatibility() {
|
||||||
|
let compat = AnthropicMessagesCompatibility::default();
|
||||||
|
assert_eq!(compat.anthropic_version.as_deref(), Some("2023-06-01"));
|
||||||
|
assert_eq!(compat.default_max_tokens, Some(1024));
|
||||||
|
assert!(!compat.system_as_array);
|
||||||
|
assert!(compat.preserve_metadata);
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,272 @@
|
|||||||
|
// CanonicalRequest -> Anthropic /v1/messages JSON formatter
|
||||||
|
use crate::transformers::anthropic::messages::AnthropicMessagesCompatibility;
|
||||||
|
use crate::transformers::canonical::*;
|
||||||
|
use serde_json::{json, Value};
|
||||||
|
|
||||||
|
/// Format a canonical request into Anthropic Messages API JSON.
|
||||||
|
pub fn format_messages_request(request: &CanonicalRequest) -> Value {
|
||||||
|
format_messages_request_with_compat(request, &AnthropicMessagesCompatibility::default())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Format a canonical request into Anthropic Messages API JSON with compatibility switches.
|
||||||
|
pub fn format_messages_request_with_compat(
|
||||||
|
request: &CanonicalRequest,
|
||||||
|
compatibility: &AnthropicMessagesCompatibility,
|
||||||
|
) -> Value {
|
||||||
|
let mut result = json!({
|
||||||
|
"model": request.model,
|
||||||
|
"messages": format_messages(&request.messages),
|
||||||
|
"stream": request.stream,
|
||||||
|
});
|
||||||
|
|
||||||
|
if let Some(max_tokens) = request.max_tokens.or(compatibility.default_max_tokens) {
|
||||||
|
result["max_tokens"] = json!(max_tokens);
|
||||||
|
}
|
||||||
|
if let Some(temperature) = request.temperature {
|
||||||
|
result["temperature"] = json!(temperature);
|
||||||
|
}
|
||||||
|
if let Some(top_p) = request.top_p {
|
||||||
|
result["top_p"] = json!(top_p);
|
||||||
|
}
|
||||||
|
if let Some(stop) = &request.stop {
|
||||||
|
result["stop_sequences"] = json!(stop);
|
||||||
|
}
|
||||||
|
if let Some(tools) = &request.tools {
|
||||||
|
result["tools"] = json!(format_tools(tools));
|
||||||
|
}
|
||||||
|
if let Some(tool_choice) = &request.tool_choice {
|
||||||
|
if let Some(value) = format_tool_choice(tool_choice) {
|
||||||
|
result["tool_choice"] = value;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if let Some(reasoning) = &request.reasoning {
|
||||||
|
if let Some(effort) = &reasoning.effort {
|
||||||
|
result["thinking"] = json!({
|
||||||
|
"type": "enabled",
|
||||||
|
"budget_tokens": reasoning_budget_tokens(effort),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if compatibility.preserve_metadata {
|
||||||
|
if let Some(metadata) = &request.metadata {
|
||||||
|
result["metadata"] = metadata.clone();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let system_parts = collect_system_parts(&request.messages);
|
||||||
|
if !system_parts.is_empty() {
|
||||||
|
result["system"] = if compatibility.system_as_array {
|
||||||
|
Value::Array(system_parts)
|
||||||
|
} else {
|
||||||
|
Value::String(system_parts_to_text(&system_parts))
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
if let Some(extra) = &compatibility.extra_request_fields {
|
||||||
|
merge_extra_fields(&mut result, extra);
|
||||||
|
}
|
||||||
|
|
||||||
|
result
|
||||||
|
}
|
||||||
|
|
||||||
|
fn format_messages(messages: &[CanonicalMessage]) -> Vec<Value> {
|
||||||
|
messages
|
||||||
|
.iter()
|
||||||
|
.filter(|message| message.role != Role::System)
|
||||||
|
.map(format_message)
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn format_message(message: &CanonicalMessage) -> Value {
|
||||||
|
let role = match message.role {
|
||||||
|
Role::Assistant => "assistant",
|
||||||
|
Role::User | Role::Tool | Role::System => "user",
|
||||||
|
};
|
||||||
|
|
||||||
|
let mut content = match message.role {
|
||||||
|
Role::Tool => format_tool_result_content(message),
|
||||||
|
_ => format_content_parts(&message.content),
|
||||||
|
};
|
||||||
|
|
||||||
|
if let Some(tool_calls) = &message.tool_calls {
|
||||||
|
content.extend(tool_calls.iter().map(format_tool_use));
|
||||||
|
}
|
||||||
|
|
||||||
|
json!({
|
||||||
|
"role": role,
|
||||||
|
"content": content,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn collect_system_parts(messages: &[CanonicalMessage]) -> Vec<Value> {
|
||||||
|
messages
|
||||||
|
.iter()
|
||||||
|
.filter(|message| message.role == Role::System)
|
||||||
|
.flat_map(|message| format_content_parts(&message.content))
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn system_parts_to_text(parts: &[Value]) -> String {
|
||||||
|
parts
|
||||||
|
.iter()
|
||||||
|
.filter_map(|part| part["text"].as_str())
|
||||||
|
.collect::<Vec<_>>()
|
||||||
|
.join("\n")
|
||||||
|
}
|
||||||
|
|
||||||
|
fn format_content_parts(parts: &[ContentPart]) -> Vec<Value> {
|
||||||
|
parts
|
||||||
|
.iter()
|
||||||
|
.map(|part| match part {
|
||||||
|
ContentPart::Text { text } => json!({
|
||||||
|
"type": "text",
|
||||||
|
"text": text,
|
||||||
|
}),
|
||||||
|
ContentPart::Image { url, .. } => json!({
|
||||||
|
"type": "image",
|
||||||
|
"source": {
|
||||||
|
"type": "url",
|
||||||
|
"url": url,
|
||||||
|
},
|
||||||
|
}),
|
||||||
|
})
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn format_tool_result_content(message: &CanonicalMessage) -> Vec<Value> {
|
||||||
|
let content = if message.content.len() == 1 {
|
||||||
|
match &message.content[0] {
|
||||||
|
ContentPart::Text { text } => Value::String(text.clone()),
|
||||||
|
_ => Value::Array(format_content_parts(&message.content)),
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
Value::Array(format_content_parts(&message.content))
|
||||||
|
};
|
||||||
|
|
||||||
|
vec![json!({
|
||||||
|
"type": "tool_result",
|
||||||
|
"tool_use_id": message.tool_call_id.clone().unwrap_or_default(),
|
||||||
|
"content": content,
|
||||||
|
})]
|
||||||
|
}
|
||||||
|
|
||||||
|
fn format_tool_use(tool_call: &ToolCall) -> Value {
|
||||||
|
json!({
|
||||||
|
"type": "tool_use",
|
||||||
|
"id": tool_call.id,
|
||||||
|
"name": tool_call.name,
|
||||||
|
"input": tool_call.arguments,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn format_tools(tools: &[ToolDefinition]) -> Vec<Value> {
|
||||||
|
tools
|
||||||
|
.iter()
|
||||||
|
.map(|tool| {
|
||||||
|
let mut value = json!({
|
||||||
|
"name": tool.name,
|
||||||
|
"input_schema": tool.input_schema,
|
||||||
|
});
|
||||||
|
if let Some(description) = &tool.description {
|
||||||
|
value["description"] = json!(description);
|
||||||
|
}
|
||||||
|
value
|
||||||
|
})
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn format_tool_choice(tool_choice: &ToolChoice) -> Option<Value> {
|
||||||
|
match tool_choice {
|
||||||
|
ToolChoice::None => None,
|
||||||
|
ToolChoice::Auto => Some(json!({ "type": "auto" })),
|
||||||
|
ToolChoice::Required => Some(json!({ "type": "any" })),
|
||||||
|
ToolChoice::Specific(name) => Some(json!({ "type": "tool", "name": name })),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn reasoning_budget_tokens(effort: &str) -> i32 {
|
||||||
|
match effort {
|
||||||
|
"low" => 1024,
|
||||||
|
"medium" => 4096,
|
||||||
|
"high" => 8192,
|
||||||
|
_ => 2048,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn merge_extra_fields(target: &mut Value, extra: &Value) {
|
||||||
|
let (Some(target), Some(extra)) = (target.as_object_mut(), extra.as_object()) else {
|
||||||
|
return;
|
||||||
|
};
|
||||||
|
|
||||||
|
for (key, value) in extra {
|
||||||
|
target.insert(key.clone(), value.clone());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
use serde_json::json;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_format_simple_request() {
|
||||||
|
let req = CanonicalRequest::chat("claude-3-5-sonnet-20241022", "Hello");
|
||||||
|
let json = format_messages_request(&req);
|
||||||
|
assert_eq!(json["model"], "claude-3-5-sonnet-20241022");
|
||||||
|
assert_eq!(json["max_tokens"], 1024);
|
||||||
|
assert_eq!(json["messages"][0]["role"], "user");
|
||||||
|
assert_eq!(json["messages"][0]["content"][0]["text"], "Hello");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_format_system_and_tools() {
|
||||||
|
let mut req = CanonicalRequest::chat("claude", "Use a tool");
|
||||||
|
req.messages.insert(0, CanonicalMessage::system("Be brief"));
|
||||||
|
req.tools = Some(vec![ToolDefinition {
|
||||||
|
name: "search".to_string(),
|
||||||
|
description: Some("Search docs".to_string()),
|
||||||
|
input_schema: json!({"type": "object"}),
|
||||||
|
}]);
|
||||||
|
req.tool_choice = Some(ToolChoice::Specific("search".to_string()));
|
||||||
|
|
||||||
|
let json = format_messages_request(&req);
|
||||||
|
assert_eq!(json["system"], "Be brief");
|
||||||
|
assert_eq!(json["tools"][0]["name"], "search");
|
||||||
|
assert_eq!(json["tool_choice"]["type"], "tool");
|
||||||
|
assert_eq!(json["tool_choice"]["name"], "search");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_format_tool_call_and_result() {
|
||||||
|
let assistant = CanonicalMessage {
|
||||||
|
role: Role::Assistant,
|
||||||
|
content: vec![],
|
||||||
|
name: None,
|
||||||
|
tool_calls: Some(vec![ToolCall {
|
||||||
|
id: "toolu_1".to_string(),
|
||||||
|
name: "calc".to_string(),
|
||||||
|
arguments: json!({"x": 1}),
|
||||||
|
}]),
|
||||||
|
tool_call_id: None,
|
||||||
|
};
|
||||||
|
let tool = CanonicalMessage::tool("2", "toolu_1");
|
||||||
|
let req = CanonicalRequest {
|
||||||
|
model: "claude".to_string(),
|
||||||
|
messages: vec![assistant, tool],
|
||||||
|
stream: false,
|
||||||
|
max_tokens: Some(10),
|
||||||
|
temperature: None,
|
||||||
|
top_p: None,
|
||||||
|
stop: None,
|
||||||
|
tools: None,
|
||||||
|
tool_choice: None,
|
||||||
|
reasoning: None,
|
||||||
|
metadata: None,
|
||||||
|
};
|
||||||
|
|
||||||
|
let json = format_messages_request(&req);
|
||||||
|
assert_eq!(json["messages"][0]["content"][0]["type"], "tool_use");
|
||||||
|
assert_eq!(json["messages"][1]["content"][0]["type"], "tool_result");
|
||||||
|
assert_eq!(json["messages"][1]["content"][0]["tool_use_id"], "toolu_1");
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -1 +1,10 @@
|
|||||||
// Transformers - Anthropic Messages
|
// Transformers - Anthropic Messages
|
||||||
|
pub mod compatibility;
|
||||||
|
pub mod inbound;
|
||||||
|
pub mod outbound;
|
||||||
|
pub mod stream;
|
||||||
|
|
||||||
|
pub use compatibility::AnthropicMessagesCompatibility;
|
||||||
|
pub use inbound::{format_messages_request, format_messages_request_with_compat};
|
||||||
|
pub use outbound::parse_messages_response;
|
||||||
|
pub use stream::parse_stream_event;
|
||||||
|
|||||||
@@ -0,0 +1,143 @@
|
|||||||
|
// Anthropic /v1/messages response -> CanonicalResponse parser
|
||||||
|
use crate::transformers::canonical::*;
|
||||||
|
use serde_json::Value;
|
||||||
|
|
||||||
|
/// Parse an Anthropic Messages API response body into canonical format.
|
||||||
|
pub fn parse_messages_response(body: &Value) -> Result<CanonicalResponse, String> {
|
||||||
|
let id = body["id"].as_str().unwrap_or_default().to_string();
|
||||||
|
let model = body["model"]
|
||||||
|
.as_str()
|
||||||
|
.ok_or_else(|| "model is required".to_string())?
|
||||||
|
.to_string();
|
||||||
|
|
||||||
|
let content = body["content"]
|
||||||
|
.as_array()
|
||||||
|
.ok_or_else(|| "content is required".to_string())?;
|
||||||
|
|
||||||
|
let mut text_parts = Vec::new();
|
||||||
|
let mut tool_calls = Vec::new();
|
||||||
|
|
||||||
|
for part in content {
|
||||||
|
match part["type"].as_str() {
|
||||||
|
Some("text") => {
|
||||||
|
if let Some(text) = part["text"].as_str() {
|
||||||
|
text_parts.push(ContentPart::Text {
|
||||||
|
text: text.to_string(),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Some("tool_use") => {
|
||||||
|
let Some(id) = part["id"].as_str() else {
|
||||||
|
continue;
|
||||||
|
};
|
||||||
|
let Some(name) = part["name"].as_str() else {
|
||||||
|
continue;
|
||||||
|
};
|
||||||
|
tool_calls.push(ToolCall {
|
||||||
|
id: id.to_string(),
|
||||||
|
name: name.to_string(),
|
||||||
|
arguments: part["input"].clone(),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
_ => {}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let message = CanonicalMessage {
|
||||||
|
role: Role::Assistant,
|
||||||
|
content: text_parts,
|
||||||
|
name: None,
|
||||||
|
tool_calls: if tool_calls.is_empty() {
|
||||||
|
None
|
||||||
|
} else {
|
||||||
|
Some(tool_calls)
|
||||||
|
},
|
||||||
|
tool_call_id: None,
|
||||||
|
};
|
||||||
|
|
||||||
|
Ok(CanonicalResponse {
|
||||||
|
id,
|
||||||
|
model,
|
||||||
|
choices: vec![CanonicalChoice {
|
||||||
|
index: 0,
|
||||||
|
message,
|
||||||
|
finish_reason: body["stop_reason"].as_str().map(normalize_stop_reason),
|
||||||
|
}],
|
||||||
|
usage: parse_usage(body.get("usage")),
|
||||||
|
metadata: Some(body.clone()),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn parse_usage(value: Option<&Value>) -> Option<Usage> {
|
||||||
|
let usage = value?.as_object()?;
|
||||||
|
let prompt_tokens = usage.get("input_tokens").and_then(as_i32);
|
||||||
|
let completion_tokens = usage.get("output_tokens").and_then(as_i32);
|
||||||
|
Some(Usage {
|
||||||
|
prompt_tokens,
|
||||||
|
completion_tokens,
|
||||||
|
total_tokens: prompt_tokens.zip(completion_tokens).map(|(a, b)| a + b),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn as_i32(value: &Value) -> Option<i32> {
|
||||||
|
value.as_i64().map(|n| n as i32)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn normalize_stop_reason(reason: &str) -> String {
|
||||||
|
match reason {
|
||||||
|
"end_turn" | "stop_sequence" => "stop".to_string(),
|
||||||
|
"max_tokens" => "length".to_string(),
|
||||||
|
"tool_use" => "tool_calls".to_string(),
|
||||||
|
other => other.to_string(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
use serde_json::json;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_parse_simple_response() {
|
||||||
|
let body = json!({
|
||||||
|
"id": "msg_123",
|
||||||
|
"type": "message",
|
||||||
|
"role": "assistant",
|
||||||
|
"model": "claude",
|
||||||
|
"content": [{"type": "text", "text": "Hello"}],
|
||||||
|
"stop_reason": "end_turn",
|
||||||
|
"usage": {"input_tokens": 10, "output_tokens": 3}
|
||||||
|
});
|
||||||
|
|
||||||
|
let response = parse_messages_response(&body).unwrap();
|
||||||
|
assert_eq!(response.id, "msg_123");
|
||||||
|
assert_eq!(response.model, "claude");
|
||||||
|
assert_eq!(response.choices[0].finish_reason.as_deref(), Some("stop"));
|
||||||
|
assert_eq!(response.usage.unwrap().total_tokens, Some(13));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_parse_tool_use_response() {
|
||||||
|
let body = json!({
|
||||||
|
"id": "msg_123",
|
||||||
|
"model": "claude",
|
||||||
|
"content": [{
|
||||||
|
"type": "tool_use",
|
||||||
|
"id": "toolu_1",
|
||||||
|
"name": "search",
|
||||||
|
"input": {"q": "rust"}
|
||||||
|
}],
|
||||||
|
"stop_reason": "tool_use"
|
||||||
|
});
|
||||||
|
|
||||||
|
let response = parse_messages_response(&body).unwrap();
|
||||||
|
let tool_calls = response.choices[0].message.tool_calls.as_ref().unwrap();
|
||||||
|
assert_eq!(tool_calls[0].id, "toolu_1");
|
||||||
|
assert_eq!(tool_calls[0].name, "search");
|
||||||
|
assert_eq!(tool_calls[0].arguments["q"], "rust");
|
||||||
|
assert_eq!(
|
||||||
|
response.choices[0].finish_reason.as_deref(),
|
||||||
|
Some("tool_calls")
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,204 @@
|
|||||||
|
// Anthropic Messages SSE stream parser
|
||||||
|
use crate::transformers::canonical::*;
|
||||||
|
use serde_json::Value;
|
||||||
|
|
||||||
|
/// Parse a single Anthropic Messages SSE event into a canonical stream event.
|
||||||
|
///
|
||||||
|
/// Accepts either a full SSE frame (`event: ...\ndata: ...`) or a raw JSON data payload.
|
||||||
|
pub fn parse_stream_event(frame: &str) -> Option<CanonicalStreamEvent> {
|
||||||
|
let (event_name, data) = parse_sse_frame(frame)?;
|
||||||
|
let parsed: Value = serde_json::from_str(&data).ok()?;
|
||||||
|
let event_type = event_name
|
||||||
|
.as_deref()
|
||||||
|
.or_else(|| parsed["type"].as_str())
|
||||||
|
.unwrap_or_default();
|
||||||
|
|
||||||
|
match event_type {
|
||||||
|
"message_start" => parse_message_start(&parsed),
|
||||||
|
"content_block_delta" => parse_content_block_delta(&parsed),
|
||||||
|
"content_block_stop" => Some(empty_event()),
|
||||||
|
"message_delta" => parse_message_delta(&parsed),
|
||||||
|
"message_stop" => Some(empty_event()),
|
||||||
|
_ => None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn parse_sse_frame(frame: &str) -> Option<(Option<String>, String)> {
|
||||||
|
let trimmed = frame.trim();
|
||||||
|
if trimmed.is_empty() || trimmed == "data: [DONE]" || trimmed == "[DONE]" {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
|
||||||
|
if trimmed.starts_with('{') {
|
||||||
|
return Some((None, trimmed.to_string()));
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut event_name = None;
|
||||||
|
let mut data_lines = Vec::new();
|
||||||
|
|
||||||
|
for line in trimmed.lines() {
|
||||||
|
if let Some(value) = line.strip_prefix("event:") {
|
||||||
|
event_name = Some(value.trim().to_string());
|
||||||
|
} else if let Some(value) = line.strip_prefix("data:") {
|
||||||
|
data_lines.push(value.trim());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if data_lines.is_empty() {
|
||||||
|
None
|
||||||
|
} else {
|
||||||
|
Some((event_name, data_lines.join("\n")))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn parse_message_start(parsed: &Value) -> Option<CanonicalStreamEvent> {
|
||||||
|
let message = parsed.get("message")?;
|
||||||
|
Some(CanonicalStreamEvent {
|
||||||
|
id: message["id"].as_str().map(|s| s.to_string()),
|
||||||
|
choices: vec![CanonicalStreamChoice {
|
||||||
|
index: 0,
|
||||||
|
delta: StreamDelta {
|
||||||
|
role: Some(Role::Assistant),
|
||||||
|
content: None,
|
||||||
|
tool_calls: None,
|
||||||
|
},
|
||||||
|
finish_reason: None,
|
||||||
|
}],
|
||||||
|
usage: parse_usage(message.get("usage")),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn parse_content_block_delta(parsed: &Value) -> Option<CanonicalStreamEvent> {
|
||||||
|
let index = parsed["index"].as_i64().unwrap_or(0) as i32;
|
||||||
|
let delta = parsed.get("delta")?;
|
||||||
|
|
||||||
|
let stream_delta = match delta["type"].as_str()? {
|
||||||
|
"text_delta" => StreamDelta {
|
||||||
|
role: None,
|
||||||
|
content: delta["text"].as_str().map(|s| s.to_string()),
|
||||||
|
tool_calls: None,
|
||||||
|
},
|
||||||
|
"input_json_delta" => StreamDelta {
|
||||||
|
role: None,
|
||||||
|
content: None,
|
||||||
|
tool_calls: Some(vec![StreamToolCall {
|
||||||
|
index,
|
||||||
|
id: None,
|
||||||
|
name: None,
|
||||||
|
arguments: delta["partial_json"].as_str().map(|s| s.to_string()),
|
||||||
|
}]),
|
||||||
|
},
|
||||||
|
_ => return None,
|
||||||
|
};
|
||||||
|
|
||||||
|
Some(CanonicalStreamEvent {
|
||||||
|
id: None,
|
||||||
|
choices: vec![CanonicalStreamChoice {
|
||||||
|
index,
|
||||||
|
delta: stream_delta,
|
||||||
|
finish_reason: None,
|
||||||
|
}],
|
||||||
|
usage: None,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn parse_message_delta(parsed: &Value) -> Option<CanonicalStreamEvent> {
|
||||||
|
let delta = parsed.get("delta")?;
|
||||||
|
let finish_reason = delta["stop_reason"].as_str().map(normalize_stop_reason);
|
||||||
|
|
||||||
|
Some(CanonicalStreamEvent {
|
||||||
|
id: None,
|
||||||
|
choices: vec![CanonicalStreamChoice {
|
||||||
|
index: 0,
|
||||||
|
delta: StreamDelta {
|
||||||
|
role: None,
|
||||||
|
content: None,
|
||||||
|
tool_calls: None,
|
||||||
|
},
|
||||||
|
finish_reason,
|
||||||
|
}],
|
||||||
|
usage: parse_usage(parsed.get("usage")),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn parse_usage(value: Option<&Value>) -> Option<Usage> {
|
||||||
|
let usage = value?.as_object()?;
|
||||||
|
let prompt_tokens = usage.get("input_tokens").and_then(as_i32);
|
||||||
|
let completion_tokens = usage.get("output_tokens").and_then(as_i32);
|
||||||
|
Some(Usage {
|
||||||
|
prompt_tokens,
|
||||||
|
completion_tokens,
|
||||||
|
total_tokens: prompt_tokens.zip(completion_tokens).map(|(a, b)| a + b),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn as_i32(value: &Value) -> Option<i32> {
|
||||||
|
value.as_i64().map(|n| n as i32)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn normalize_stop_reason(reason: &str) -> String {
|
||||||
|
match reason {
|
||||||
|
"end_turn" | "stop_sequence" => "stop".to_string(),
|
||||||
|
"max_tokens" => "length".to_string(),
|
||||||
|
"tool_use" => "tool_calls".to_string(),
|
||||||
|
other => other.to_string(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn empty_event() -> CanonicalStreamEvent {
|
||||||
|
CanonicalStreamEvent {
|
||||||
|
id: None,
|
||||||
|
choices: vec![],
|
||||||
|
usage: None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_parse_message_start() {
|
||||||
|
let frame = r#"event: message_start
|
||||||
|
data: {"type":"message_start","message":{"id":"msg_1","model":"claude","role":"assistant","content":[],"usage":{"input_tokens":4,"output_tokens":1}}}"#;
|
||||||
|
let event = parse_stream_event(frame).unwrap();
|
||||||
|
assert_eq!(event.id.as_deref(), Some("msg_1"));
|
||||||
|
assert_eq!(event.choices[0].delta.role, Some(Role::Assistant));
|
||||||
|
assert_eq!(event.usage.unwrap().total_tokens, Some(5));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_parse_text_delta() {
|
||||||
|
let frame = r#"event: content_block_delta
|
||||||
|
data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello"}}"#;
|
||||||
|
let event = parse_stream_event(frame).unwrap();
|
||||||
|
assert_eq!(event.choices[0].index, 0);
|
||||||
|
assert_eq!(event.choices[0].delta.content.as_deref(), Some("Hello"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_parse_tool_input_delta() {
|
||||||
|
let frame = r#"data: {"type":"content_block_delta","index":1,"delta":{"type":"input_json_delta","partial_json":"{\"q\""}}"#;
|
||||||
|
let event = parse_stream_event(frame).unwrap();
|
||||||
|
let tool_calls = event.choices[0].delta.tool_calls.as_ref().unwrap();
|
||||||
|
assert_eq!(tool_calls[0].index, 1);
|
||||||
|
assert_eq!(tool_calls[0].arguments.as_deref(), Some("{\"q\""));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_parse_message_delta() {
|
||||||
|
let frame = r#"event: message_delta
|
||||||
|
data: {"type":"message_delta","delta":{"stop_reason":"max_tokens"},"usage":{"output_tokens":9}}"#;
|
||||||
|
let event = parse_stream_event(frame).unwrap();
|
||||||
|
assert_eq!(event.choices[0].finish_reason.as_deref(), Some("length"));
|
||||||
|
assert_eq!(event.usage.unwrap().completion_tokens, Some(9));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_parse_message_stop() {
|
||||||
|
let frame = r#"event: message_stop
|
||||||
|
data: {"type":"message_stop"}"#;
|
||||||
|
let event = parse_stream_event(frame).unwrap();
|
||||||
|
assert!(event.choices.is_empty());
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -1 +1,7 @@
|
|||||||
// Transformers - Anthropic
|
// Transformers - Anthropic
|
||||||
|
pub mod messages;
|
||||||
|
|
||||||
|
pub use messages::{
|
||||||
|
format_messages_request, format_messages_request_with_compat, parse_messages_response,
|
||||||
|
parse_stream_event, AnthropicMessagesCompatibility,
|
||||||
|
};
|
||||||
|
|||||||
@@ -1 +1,13 @@
|
|||||||
// Transformers - canonical
|
// Canonical format — unified internal representation
|
||||||
|
pub mod request;
|
||||||
|
pub mod response;
|
||||||
|
pub mod types;
|
||||||
|
|
||||||
|
pub use request::{CanonicalMessage, CanonicalRequest};
|
||||||
|
pub use response::{
|
||||||
|
CanonicalChoice, CanonicalResponse, CanonicalStreamChoice, CanonicalStreamEvent,
|
||||||
|
};
|
||||||
|
pub use types::{
|
||||||
|
ContentPart, ReasoningConfig, Role, StreamDelta, StreamToolCall, ToolCall, ToolChoice,
|
||||||
|
ToolDefinition, Usage,
|
||||||
|
};
|
||||||
|
|||||||
@@ -0,0 +1,107 @@
|
|||||||
|
// Canonical request — unified input format for all providers
|
||||||
|
use crate::transformers::canonical::types::{
|
||||||
|
ContentPart, ReasoningConfig, Role, ToolCall, ToolChoice, ToolDefinition,
|
||||||
|
};
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
|
||||||
|
/// Unified internal request format
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct CanonicalRequest {
|
||||||
|
pub model: String,
|
||||||
|
pub messages: Vec<CanonicalMessage>,
|
||||||
|
pub stream: bool,
|
||||||
|
pub max_tokens: Option<i32>,
|
||||||
|
pub temperature: Option<f64>,
|
||||||
|
pub top_p: Option<f64>,
|
||||||
|
pub stop: Option<Vec<String>>,
|
||||||
|
pub tools: Option<Vec<ToolDefinition>>,
|
||||||
|
pub tool_choice: Option<ToolChoice>,
|
||||||
|
pub reasoning: Option<ReasoningConfig>,
|
||||||
|
pub metadata: Option<serde_json::Value>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl CanonicalRequest {
|
||||||
|
/// Create a simple chat request
|
||||||
|
pub fn chat(model: &str, content: &str) -> Self {
|
||||||
|
Self {
|
||||||
|
model: model.to_string(),
|
||||||
|
messages: vec![CanonicalMessage::user(content)],
|
||||||
|
stream: false,
|
||||||
|
max_tokens: None,
|
||||||
|
temperature: None,
|
||||||
|
top_p: None,
|
||||||
|
stop: None,
|
||||||
|
tools: None,
|
||||||
|
tool_choice: None,
|
||||||
|
reasoning: None,
|
||||||
|
metadata: None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Create a streaming chat request
|
||||||
|
pub fn chat_stream(model: &str, content: &str) -> Self {
|
||||||
|
let mut req = Self::chat(model, content);
|
||||||
|
req.stream = true;
|
||||||
|
req
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A single message in canonical format
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct CanonicalMessage {
|
||||||
|
pub role: Role,
|
||||||
|
pub content: Vec<ContentPart>,
|
||||||
|
pub name: Option<String>,
|
||||||
|
pub tool_calls: Option<Vec<ToolCall>>,
|
||||||
|
pub tool_call_id: Option<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl CanonicalMessage {
|
||||||
|
pub fn system(content: &str) -> Self {
|
||||||
|
Self {
|
||||||
|
role: Role::System,
|
||||||
|
content: vec![ContentPart::Text {
|
||||||
|
text: content.to_string(),
|
||||||
|
}],
|
||||||
|
name: None,
|
||||||
|
tool_calls: None,
|
||||||
|
tool_call_id: None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn user(content: &str) -> Self {
|
||||||
|
Self {
|
||||||
|
role: Role::User,
|
||||||
|
content: vec![ContentPart::Text {
|
||||||
|
text: content.to_string(),
|
||||||
|
}],
|
||||||
|
name: None,
|
||||||
|
tool_calls: None,
|
||||||
|
tool_call_id: None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn assistant(content: &str) -> Self {
|
||||||
|
Self {
|
||||||
|
role: Role::Assistant,
|
||||||
|
content: vec![ContentPart::Text {
|
||||||
|
text: content.to_string(),
|
||||||
|
}],
|
||||||
|
name: None,
|
||||||
|
tool_calls: None,
|
||||||
|
tool_call_id: None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn tool(content: &str, tool_call_id: &str) -> Self {
|
||||||
|
Self {
|
||||||
|
role: Role::Tool,
|
||||||
|
content: vec![ContentPart::Text {
|
||||||
|
text: content.to_string(),
|
||||||
|
}],
|
||||||
|
name: None,
|
||||||
|
tool_calls: None,
|
||||||
|
tool_call_id: Some(tool_call_id.to_string()),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,54 @@
|
|||||||
|
// Canonical response — unified output format for all providers
|
||||||
|
use crate::transformers::canonical::request::CanonicalMessage;
|
||||||
|
use crate::transformers::canonical::types::{StreamDelta, ToolCall, Usage};
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
|
||||||
|
/// Unified internal response format
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct CanonicalResponse {
|
||||||
|
pub id: String,
|
||||||
|
pub model: String,
|
||||||
|
pub choices: Vec<CanonicalChoice>,
|
||||||
|
pub usage: Option<Usage>,
|
||||||
|
pub metadata: Option<serde_json::Value>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A single choice in the response
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct CanonicalChoice {
|
||||||
|
pub index: i32,
|
||||||
|
pub message: CanonicalMessage,
|
||||||
|
pub finish_reason: Option<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl CanonicalResponse {
|
||||||
|
pub fn simple(model: &str, content: &str) -> Self {
|
||||||
|
Self {
|
||||||
|
id: format!("canonical-{}", uuid::Uuid::new_v4()),
|
||||||
|
model: model.to_string(),
|
||||||
|
choices: vec![CanonicalChoice {
|
||||||
|
index: 0,
|
||||||
|
message: CanonicalMessage::assistant(content),
|
||||||
|
finish_reason: Some("stop".to_string()),
|
||||||
|
}],
|
||||||
|
usage: None,
|
||||||
|
metadata: None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── SSE stream event ───────────────────────────────────────────────────
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct CanonicalStreamEvent {
|
||||||
|
pub id: Option<String>,
|
||||||
|
pub choices: Vec<CanonicalStreamChoice>,
|
||||||
|
pub usage: Option<Usage>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct CanonicalStreamChoice {
|
||||||
|
pub index: i32,
|
||||||
|
pub delta: StreamDelta,
|
||||||
|
pub finish_reason: Option<String>,
|
||||||
|
}
|
||||||
@@ -0,0 +1,88 @@
|
|||||||
|
// Canonical format types — unified internal representation for protocol transformers
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
|
||||||
|
// ── Message Role ───────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||||
|
pub enum Role {
|
||||||
|
System,
|
||||||
|
User,
|
||||||
|
Assistant,
|
||||||
|
Tool,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl std::fmt::Display for Role {
|
||||||
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||||
|
match self {
|
||||||
|
Self::System => write!(f, "system"),
|
||||||
|
Self::User => write!(f, "user"),
|
||||||
|
Self::Assistant => write!(f, "assistant"),
|
||||||
|
Self::Tool => write!(f, "tool"),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Content ────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub enum ContentPart {
|
||||||
|
Text { text: String },
|
||||||
|
Image { url: String, detail: Option<String> },
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Tools ───────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct ToolDefinition {
|
||||||
|
pub name: String,
|
||||||
|
pub description: Option<String>,
|
||||||
|
pub input_schema: serde_json::Value,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct ToolCall {
|
||||||
|
pub id: String,
|
||||||
|
pub name: String,
|
||||||
|
pub arguments: serde_json::Value,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub enum ToolChoice {
|
||||||
|
None,
|
||||||
|
Auto,
|
||||||
|
Required,
|
||||||
|
Specific(String),
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Reasoning ───────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct ReasoningConfig {
|
||||||
|
pub effort: Option<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Usage ───────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct Usage {
|
||||||
|
pub prompt_tokens: Option<i32>,
|
||||||
|
pub completion_tokens: Option<i32>,
|
||||||
|
pub total_tokens: Option<i32>,
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── Stream Delta ────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct StreamDelta {
|
||||||
|
pub role: Option<Role>,
|
||||||
|
pub content: Option<String>,
|
||||||
|
pub tool_calls: Option<Vec<StreamToolCall>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct StreamToolCall {
|
||||||
|
pub index: i32,
|
||||||
|
pub id: Option<String>,
|
||||||
|
pub name: Option<String>,
|
||||||
|
pub arguments: Option<String>,
|
||||||
|
}
|
||||||
@@ -0,0 +1,37 @@
|
|||||||
|
// Gemini CLI integration stub
|
||||||
|
use crate::transformers::canonical::CanonicalResponse;
|
||||||
|
use anyhow::{bail, Result};
|
||||||
|
use serde_json::Value;
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Default)]
|
||||||
|
pub struct GeminiCliBridge;
|
||||||
|
|
||||||
|
impl GeminiCliBridge {
|
||||||
|
pub fn new() -> Self {
|
||||||
|
Self
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn execute(&self, _request: &Value) -> Result<CanonicalResponse> {
|
||||||
|
bail!("Gemini CLI bridge is not implemented (501)")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
use serde_json::json;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_new_bridge() {
|
||||||
|
let bridge = GeminiCliBridge::new();
|
||||||
|
let error = bridge.execute(&json!({})).unwrap_err();
|
||||||
|
assert!(error.to_string().contains("501"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_default_bridge_returns_not_implemented() {
|
||||||
|
let bridge = GeminiCliBridge;
|
||||||
|
let error = bridge.execute(&json!({"contents": []})).unwrap_err();
|
||||||
|
assert!(error.to_string().contains("not implemented"));
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,125 @@
|
|||||||
|
// CanonicalRequest -> Gemini GenerateContent JSON formatter
|
||||||
|
use crate::transformers::canonical::*;
|
||||||
|
use serde_json::{json, Value};
|
||||||
|
|
||||||
|
/// Format a canonical request into Gemini's GenerateContent request body.
|
||||||
|
pub fn format_generate_content_request(request: &CanonicalRequest) -> Value {
|
||||||
|
let mut result = json!({
|
||||||
|
"contents": format_contents(&request.messages),
|
||||||
|
});
|
||||||
|
|
||||||
|
let generation_config = format_generation_config(request);
|
||||||
|
if !generation_config
|
||||||
|
.as_object()
|
||||||
|
.is_none_or(|object| object.is_empty())
|
||||||
|
{
|
||||||
|
result["generationConfig"] = generation_config;
|
||||||
|
}
|
||||||
|
|
||||||
|
let system_parts = collect_system_parts(&request.messages);
|
||||||
|
if !system_parts.is_empty() {
|
||||||
|
result["systemInstruction"] = json!({
|
||||||
|
"parts": system_parts,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
result
|
||||||
|
}
|
||||||
|
|
||||||
|
fn format_contents(messages: &[CanonicalMessage]) -> Vec<Value> {
|
||||||
|
messages
|
||||||
|
.iter()
|
||||||
|
.filter(|message| message.role != Role::System)
|
||||||
|
.map(format_content)
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn format_content(message: &CanonicalMessage) -> Value {
|
||||||
|
let role = match message.role {
|
||||||
|
Role::Assistant => "model",
|
||||||
|
Role::User | Role::Tool | Role::System => "user",
|
||||||
|
};
|
||||||
|
|
||||||
|
json!({
|
||||||
|
"role": role,
|
||||||
|
"parts": format_parts(&message.content),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn collect_system_parts(messages: &[CanonicalMessage]) -> Vec<Value> {
|
||||||
|
messages
|
||||||
|
.iter()
|
||||||
|
.filter(|message| message.role == Role::System)
|
||||||
|
.flat_map(|message| format_parts(&message.content))
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn format_parts(parts: &[ContentPart]) -> Vec<Value> {
|
||||||
|
parts
|
||||||
|
.iter()
|
||||||
|
.map(|part| match part {
|
||||||
|
ContentPart::Text { text } => json!({ "text": text }),
|
||||||
|
ContentPart::Image { url, .. } => json!({
|
||||||
|
"fileData": {
|
||||||
|
"fileUri": url,
|
||||||
|
},
|
||||||
|
}),
|
||||||
|
})
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn format_generation_config(request: &CanonicalRequest) -> Value {
|
||||||
|
let mut config = json!({});
|
||||||
|
|
||||||
|
if let Some(max_tokens) = request.max_tokens {
|
||||||
|
config["maxOutputTokens"] = json!(max_tokens);
|
||||||
|
}
|
||||||
|
if let Some(temperature) = request.temperature {
|
||||||
|
config["temperature"] = json!(temperature);
|
||||||
|
}
|
||||||
|
if let Some(top_p) = request.top_p {
|
||||||
|
config["topP"] = json!(top_p);
|
||||||
|
}
|
||||||
|
if let Some(stop) = &request.stop {
|
||||||
|
config["stopSequences"] = json!(stop);
|
||||||
|
}
|
||||||
|
if request.stream {
|
||||||
|
config["streamResponse"] = json!(true);
|
||||||
|
}
|
||||||
|
|
||||||
|
config
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_format_simple_request() {
|
||||||
|
let req = CanonicalRequest::chat("gemini-2.0-flash", "Hello");
|
||||||
|
let json = format_generate_content_request(&req);
|
||||||
|
|
||||||
|
assert_eq!(json["contents"][0]["role"], "user");
|
||||||
|
assert_eq!(json["contents"][0]["parts"][0]["text"], "Hello");
|
||||||
|
assert!(json.get("systemInstruction").is_none());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_format_system_assistant_and_stream_config() {
|
||||||
|
let mut req = CanonicalRequest::chat_stream("gemini-2.0-flash", "Continue");
|
||||||
|
req.messages.insert(0, CanonicalMessage::system("Be brief"));
|
||||||
|
req.messages.insert(1, CanonicalMessage::assistant("Once"));
|
||||||
|
req.max_tokens = Some(128);
|
||||||
|
req.temperature = Some(0.4);
|
||||||
|
req.top_p = Some(0.9);
|
||||||
|
req.stop = Some(vec!["END".to_string()]);
|
||||||
|
|
||||||
|
let json = format_generate_content_request(&req);
|
||||||
|
|
||||||
|
assert_eq!(json["systemInstruction"]["parts"][0]["text"], "Be brief");
|
||||||
|
assert_eq!(json["contents"][0]["role"], "model");
|
||||||
|
assert_eq!(json["generationConfig"]["maxOutputTokens"], 128);
|
||||||
|
assert_eq!(json["generationConfig"]["streamResponse"], true);
|
||||||
|
assert_eq!(json["generationConfig"]["stopSequences"][0], "END");
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -1 +1,15 @@
|
|||||||
// Transformers - Gemini Generate Content
|
// Transformers - Gemini Generate Content
|
||||||
|
pub mod cli_bridge;
|
||||||
|
pub mod inbound;
|
||||||
|
pub mod outbound;
|
||||||
|
pub mod stream;
|
||||||
|
pub mod url_resolver;
|
||||||
|
|
||||||
|
pub use cli_bridge::GeminiCliBridge;
|
||||||
|
pub use inbound::format_generate_content_request;
|
||||||
|
pub use outbound::{normalize_finish_reason, parse_generate_content_response};
|
||||||
|
pub use stream::parse_stream_event;
|
||||||
|
pub use url_resolver::{
|
||||||
|
resolve_generate_content_url, resolve_generate_content_url_with_base,
|
||||||
|
resolve_generate_content_url_with_key,
|
||||||
|
};
|
||||||
|
|||||||
@@ -0,0 +1,132 @@
|
|||||||
|
// Gemini GenerateContent response -> CanonicalResponse parser
|
||||||
|
use crate::transformers::canonical::*;
|
||||||
|
use serde_json::Value;
|
||||||
|
|
||||||
|
/// Parse a Gemini GenerateContent response body into canonical format.
|
||||||
|
pub fn parse_generate_content_response(body: &Value) -> Result<CanonicalResponse, String> {
|
||||||
|
let candidates = body["candidates"]
|
||||||
|
.as_array()
|
||||||
|
.ok_or_else(|| "candidates is required".to_string())?;
|
||||||
|
|
||||||
|
let choices = candidates
|
||||||
|
.iter()
|
||||||
|
.enumerate()
|
||||||
|
.map(parse_candidate)
|
||||||
|
.collect::<Vec<_>>();
|
||||||
|
|
||||||
|
Ok(CanonicalResponse {
|
||||||
|
id: body["responseId"]
|
||||||
|
.as_str()
|
||||||
|
.map(ToString::to_string)
|
||||||
|
.unwrap_or_else(|| format!("gemini-{}", uuid::Uuid::new_v4())),
|
||||||
|
model: body["model"]
|
||||||
|
.as_str()
|
||||||
|
.or_else(|| body["modelVersion"].as_str())
|
||||||
|
.unwrap_or_default()
|
||||||
|
.to_string(),
|
||||||
|
choices,
|
||||||
|
usage: parse_usage(body.get("usageMetadata")),
|
||||||
|
metadata: Some(body.clone()),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn parse_candidate((index, candidate): (usize, &Value)) -> CanonicalChoice {
|
||||||
|
CanonicalChoice {
|
||||||
|
index: candidate["index"].as_i64().unwrap_or(index as i64) as i32,
|
||||||
|
message: CanonicalMessage {
|
||||||
|
role: Role::Assistant,
|
||||||
|
content: parse_parts(candidate["content"]["parts"].as_array()),
|
||||||
|
name: None,
|
||||||
|
tool_calls: None,
|
||||||
|
tool_call_id: None,
|
||||||
|
},
|
||||||
|
finish_reason: candidate["finishReason"]
|
||||||
|
.as_str()
|
||||||
|
.and_then(normalize_finish_reason),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn parse_parts(parts: Option<&Vec<Value>>) -> Vec<ContentPart> {
|
||||||
|
parts
|
||||||
|
.into_iter()
|
||||||
|
.flatten()
|
||||||
|
.filter_map(|part| {
|
||||||
|
part["text"].as_str().map(|text| ContentPart::Text {
|
||||||
|
text: text.to_string(),
|
||||||
|
})
|
||||||
|
})
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn parse_usage(value: Option<&Value>) -> Option<Usage> {
|
||||||
|
let usage = value?.as_object()?;
|
||||||
|
Some(Usage {
|
||||||
|
prompt_tokens: usage.get("promptTokenCount").and_then(as_i32),
|
||||||
|
completion_tokens: usage.get("candidatesTokenCount").and_then(as_i32),
|
||||||
|
total_tokens: usage.get("totalTokenCount").and_then(as_i32),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn as_i32(value: &Value) -> Option<i32> {
|
||||||
|
value.as_i64().map(|n| n as i32)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn normalize_finish_reason(reason: &str) -> Option<String> {
|
||||||
|
match reason {
|
||||||
|
"FINISH_REASON_UNSPECIFIED" => None,
|
||||||
|
"STOP" => Some("stop".to_string()),
|
||||||
|
"MAX_TOKENS" => Some("length".to_string()),
|
||||||
|
"SAFETY" | "RECITATION" | "BLOCKLIST" | "PROHIBITED_CONTENT" | "SPII" => {
|
||||||
|
Some("content_filter".to_string())
|
||||||
|
}
|
||||||
|
"MALFORMED_FUNCTION_CALL" => Some("tool_calls".to_string()),
|
||||||
|
other => Some(other.to_ascii_lowercase()),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
use serde_json::json;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_parse_simple_response() {
|
||||||
|
let body = json!({
|
||||||
|
"responseId": "resp_1",
|
||||||
|
"modelVersion": "gemini-2.0-flash",
|
||||||
|
"candidates": [{
|
||||||
|
"content": {"parts": [{"text": "Hello"}], "role": "model"},
|
||||||
|
"finishReason": "STOP"
|
||||||
|
}],
|
||||||
|
"usageMetadata": {
|
||||||
|
"promptTokenCount": 10,
|
||||||
|
"candidatesTokenCount": 5,
|
||||||
|
"totalTokenCount": 15
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
let response = parse_generate_content_response(&body).unwrap();
|
||||||
|
assert_eq!(response.id, "resp_1");
|
||||||
|
assert_eq!(response.model, "gemini-2.0-flash");
|
||||||
|
assert_eq!(response.choices[0].finish_reason.as_deref(), Some("stop"));
|
||||||
|
assert_eq!(response.usage.unwrap().total_tokens, Some(15));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_parse_safety_finish_reason() {
|
||||||
|
let body = json!({
|
||||||
|
"candidates": [{
|
||||||
|
"index": 2,
|
||||||
|
"content": {"parts": [{"text": "Blocked"}]},
|
||||||
|
"finishReason": "SAFETY"
|
||||||
|
}]
|
||||||
|
});
|
||||||
|
|
||||||
|
let response = parse_generate_content_response(&body).unwrap();
|
||||||
|
assert_eq!(response.choices[0].index, 2);
|
||||||
|
assert_eq!(
|
||||||
|
response.choices[0].finish_reason.as_deref(),
|
||||||
|
Some("content_filter")
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,131 @@
|
|||||||
|
// Gemini GenerateContent SSE stream parser
|
||||||
|
use crate::transformers::canonical::*;
|
||||||
|
use crate::transformers::gemini::generate_content::outbound::normalize_finish_reason;
|
||||||
|
use serde_json::Value;
|
||||||
|
|
||||||
|
/// Parse a single Gemini SSE frame or `data: ...` line into a canonical stream event.
|
||||||
|
pub fn parse_stream_event(frame: &str) -> Option<CanonicalStreamEvent> {
|
||||||
|
let data = parse_sse_data(frame)?;
|
||||||
|
let parsed: Value = serde_json::from_str(&data).ok()?;
|
||||||
|
|
||||||
|
let choices = parsed["candidates"]
|
||||||
|
.as_array()
|
||||||
|
.map(|candidates| {
|
||||||
|
candidates
|
||||||
|
.iter()
|
||||||
|
.enumerate()
|
||||||
|
.map(parse_candidate_delta)
|
||||||
|
.collect::<Vec<_>>()
|
||||||
|
})
|
||||||
|
.unwrap_or_default();
|
||||||
|
|
||||||
|
let usage = parse_usage(parsed.get("usageMetadata"));
|
||||||
|
|
||||||
|
Some(CanonicalStreamEvent {
|
||||||
|
id: parsed["responseId"].as_str().map(|s| s.to_string()),
|
||||||
|
choices,
|
||||||
|
usage,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn parse_sse_data(frame: &str) -> Option<String> {
|
||||||
|
let trimmed = frame.trim();
|
||||||
|
if trimmed.is_empty() || trimmed == "[DONE]" || trimmed == "data: [DONE]" {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
|
||||||
|
if trimmed.starts_with('{') {
|
||||||
|
return Some(trimmed.to_string());
|
||||||
|
}
|
||||||
|
|
||||||
|
let data_lines: Vec<&str> = trimmed
|
||||||
|
.lines()
|
||||||
|
.filter_map(|line| line.strip_prefix("data:").map(str::trim))
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
if data_lines.is_empty() {
|
||||||
|
None
|
||||||
|
} else {
|
||||||
|
Some(data_lines.join("\n"))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn parse_candidate_delta((index, candidate): (usize, &Value)) -> CanonicalStreamChoice {
|
||||||
|
CanonicalStreamChoice {
|
||||||
|
index: candidate["index"].as_i64().unwrap_or(index as i64) as i32,
|
||||||
|
delta: StreamDelta {
|
||||||
|
role: parse_role(candidate["content"]["role"].as_str()),
|
||||||
|
content: collect_text(candidate["content"]["parts"].as_array()),
|
||||||
|
tool_calls: None,
|
||||||
|
},
|
||||||
|
finish_reason: candidate["finishReason"]
|
||||||
|
.as_str()
|
||||||
|
.and_then(normalize_finish_reason),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn parse_role(role: Option<&str>) -> Option<Role> {
|
||||||
|
match role {
|
||||||
|
Some("model") => Some(Role::Assistant),
|
||||||
|
Some("user") => Some(Role::User),
|
||||||
|
_ => None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn collect_text(parts: Option<&Vec<Value>>) -> Option<String> {
|
||||||
|
let text = parts
|
||||||
|
.into_iter()
|
||||||
|
.flatten()
|
||||||
|
.filter_map(|part| part["text"].as_str())
|
||||||
|
.collect::<Vec<_>>()
|
||||||
|
.join("");
|
||||||
|
|
||||||
|
if text.is_empty() {
|
||||||
|
None
|
||||||
|
} else {
|
||||||
|
Some(text)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn parse_usage(value: Option<&Value>) -> Option<Usage> {
|
||||||
|
let usage = value?.as_object()?;
|
||||||
|
Some(Usage {
|
||||||
|
prompt_tokens: usage.get("promptTokenCount").and_then(as_i32),
|
||||||
|
completion_tokens: usage.get("candidatesTokenCount").and_then(as_i32),
|
||||||
|
total_tokens: usage.get("totalTokenCount").and_then(as_i32),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn as_i32(value: &Value) -> Option<i32> {
|
||||||
|
value.as_i64().map(|n| n as i32)
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_parse_text_chunk() {
|
||||||
|
let frame = r#"data: {"responseId":"r1","candidates":[{"content":{"parts":[{"text":"Hel"},{"text":"lo"}],"role":"model"},"finishReason":null}]}"#;
|
||||||
|
|
||||||
|
let event = parse_stream_event(frame).unwrap();
|
||||||
|
assert_eq!(event.id.as_deref(), Some("r1"));
|
||||||
|
assert_eq!(event.choices[0].delta.role, Some(Role::Assistant));
|
||||||
|
assert_eq!(event.choices[0].delta.content.as_deref(), Some("Hello"));
|
||||||
|
assert!(event.choices[0].finish_reason.is_none());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_parse_final_chunk_with_usage() {
|
||||||
|
let frame = r#"data: {"candidates":[{"content":{"parts":[]},"finishReason":"MAX_TOKENS"}],"usageMetadata":{"promptTokenCount":4,"candidatesTokenCount":6,"totalTokenCount":10}}"#;
|
||||||
|
|
||||||
|
let event = parse_stream_event(frame).unwrap();
|
||||||
|
assert_eq!(event.choices[0].finish_reason.as_deref(), Some("length"));
|
||||||
|
assert_eq!(event.usage.unwrap().total_tokens, Some(10));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_parse_done() {
|
||||||
|
assert!(parse_stream_event("data: [DONE]").is_none());
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,73 @@
|
|||||||
|
// Gemini GenerateContent URL resolver
|
||||||
|
const DEFAULT_GEMINI_BASE_URL: &str = "https://generativelanguage.googleapis.com/v1beta";
|
||||||
|
|
||||||
|
/// Resolve a Gemini model name into a GenerateContent endpoint URL.
|
||||||
|
pub fn resolve_generate_content_url(model: &str) -> String {
|
||||||
|
resolve_generate_content_url_with_base(model, DEFAULT_GEMINI_BASE_URL)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Resolve a Gemini model name into a GenerateContent endpoint URL with an API key query.
|
||||||
|
pub fn resolve_generate_content_url_with_key(model: &str, api_key: &str) -> String {
|
||||||
|
let url = resolve_generate_content_url(model);
|
||||||
|
if api_key.is_empty() {
|
||||||
|
url
|
||||||
|
} else {
|
||||||
|
format!("{}?key={}", url, api_key)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Resolve with a custom base URL for tests or alternate Gemini-compatible hosts.
|
||||||
|
pub fn resolve_generate_content_url_with_base(model: &str, base_url: &str) -> String {
|
||||||
|
let base = base_url.trim_end_matches('/');
|
||||||
|
let model = normalize_model_name(model);
|
||||||
|
format!("{}/models/{}:generateContent", base, model)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn normalize_model_name(model: &str) -> String {
|
||||||
|
let trimmed = model.trim().trim_start_matches('/');
|
||||||
|
|
||||||
|
if let Some(rest) = trimmed.strip_prefix("models/") {
|
||||||
|
return rest.to_string();
|
||||||
|
}
|
||||||
|
|
||||||
|
if let Some((_, model_id)) = trimmed.rsplit_once("/models/") {
|
||||||
|
return model_id.to_string();
|
||||||
|
}
|
||||||
|
|
||||||
|
trimmed.to_string()
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_resolve_plain_model() {
|
||||||
|
let url = resolve_generate_content_url("gemini-2.0-flash");
|
||||||
|
assert_eq!(
|
||||||
|
url,
|
||||||
|
"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_resolve_prefixed_model_with_key() {
|
||||||
|
let url = resolve_generate_content_url_with_key("models/gemini-1.5-pro", "secret");
|
||||||
|
assert_eq!(
|
||||||
|
url,
|
||||||
|
"https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-pro:generateContent?key=secret"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_resolve_publisher_model() {
|
||||||
|
let url = resolve_generate_content_url_with_base(
|
||||||
|
"publishers/google/models/gemini-2.0-flash",
|
||||||
|
"https://example.test/v1beta/",
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
url,
|
||||||
|
"https://example.test/v1beta/models/gemini-2.0-flash:generateContent"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -1 +1,8 @@
|
|||||||
// Transformers - Gemini
|
// Transformers - Gemini
|
||||||
|
pub mod generate_content;
|
||||||
|
|
||||||
|
pub use generate_content::{
|
||||||
|
format_generate_content_request, parse_generate_content_response, parse_stream_event,
|
||||||
|
resolve_generate_content_url, resolve_generate_content_url_with_base,
|
||||||
|
resolve_generate_content_url_with_key, GeminiCliBridge,
|
||||||
|
};
|
||||||
|
|||||||
@@ -1 +1,6 @@
|
|||||||
// Transformers module
|
// Transformers module — protocol conversion between canonical and provider formats
|
||||||
|
pub mod anthropic;
|
||||||
|
pub mod canonical;
|
||||||
|
pub mod gemini;
|
||||||
|
pub mod openai;
|
||||||
|
pub mod shared;
|
||||||
|
|||||||
@@ -0,0 +1,205 @@
|
|||||||
|
// OpenAI chat/completions request → CanonicalRequest parser
|
||||||
|
use crate::transformers::canonical::{
|
||||||
|
CanonicalMessage, CanonicalRequest, ContentPart, ReasoningConfig, Role, ToolCall, ToolChoice,
|
||||||
|
ToolDefinition,
|
||||||
|
};
|
||||||
|
use serde_json::Value;
|
||||||
|
|
||||||
|
/// Parse an OpenAI /v1/chat/completions request body into canonical format
|
||||||
|
pub fn parse_chat_completions_request(body: &Value) -> Result<CanonicalRequest, String> {
|
||||||
|
let model = body["model"]
|
||||||
|
.as_str()
|
||||||
|
.ok_or_else(|| "model is required".to_string())?
|
||||||
|
.to_string();
|
||||||
|
|
||||||
|
let stream = body["stream"].as_bool().unwrap_or(false);
|
||||||
|
|
||||||
|
let messages = body["messages"]
|
||||||
|
.as_array()
|
||||||
|
.ok_or_else(|| "messages is required".to_string())?;
|
||||||
|
|
||||||
|
let canonical_messages: Vec<CanonicalMessage> = messages
|
||||||
|
.iter()
|
||||||
|
.map(parse_message)
|
||||||
|
.collect::<Result<Vec<_>, _>>()?;
|
||||||
|
|
||||||
|
let max_tokens = body["max_tokens"].as_i64().map(|n| n as i32);
|
||||||
|
let max_completion_tokens = body["max_completion_tokens"].as_i64().map(|n| n as i32);
|
||||||
|
let temperature = body["temperature"].as_f64();
|
||||||
|
let top_p = body["top_p"].as_f64();
|
||||||
|
|
||||||
|
let stop = body["stop"].as_array().map(|a| {
|
||||||
|
a.iter()
|
||||||
|
.filter_map(|v| v.as_str().map(|s| s.to_string()))
|
||||||
|
.collect()
|
||||||
|
});
|
||||||
|
|
||||||
|
let tools = body["tools"]
|
||||||
|
.as_array()
|
||||||
|
.map(|a| a.iter().filter_map(parse_tool).collect());
|
||||||
|
|
||||||
|
let tool_choice = body["tool_choice"].as_str().and_then(|s| match s {
|
||||||
|
"none" => Some(ToolChoice::None),
|
||||||
|
"auto" => Some(ToolChoice::Auto),
|
||||||
|
"required" => Some(ToolChoice::Required),
|
||||||
|
name => Some(ToolChoice::Specific(name.to_string())),
|
||||||
|
});
|
||||||
|
|
||||||
|
let reasoning = body["reasoning_config"]
|
||||||
|
.as_object()
|
||||||
|
.map(|_| ReasoningConfig {
|
||||||
|
effort: body["reasoning_config"]["effort"]
|
||||||
|
.as_str()
|
||||||
|
.map(|s| s.to_string()),
|
||||||
|
});
|
||||||
|
|
||||||
|
Ok(CanonicalRequest {
|
||||||
|
model,
|
||||||
|
messages: canonical_messages,
|
||||||
|
stream,
|
||||||
|
max_tokens: max_tokens.or(max_completion_tokens),
|
||||||
|
temperature,
|
||||||
|
top_p,
|
||||||
|
stop,
|
||||||
|
tools,
|
||||||
|
tool_choice,
|
||||||
|
reasoning,
|
||||||
|
metadata: None,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn parse_message(msg: &Value) -> Result<CanonicalMessage, String> {
|
||||||
|
let role = match msg["role"].as_str() {
|
||||||
|
Some("system") => Role::System,
|
||||||
|
Some("user") => Role::User,
|
||||||
|
Some("assistant") => Role::Assistant,
|
||||||
|
Some("tool") => Role::Tool,
|
||||||
|
Some(other) => return Err(format!("Unknown role: {}", other)),
|
||||||
|
None => return Err("Missing role in message".to_string()),
|
||||||
|
};
|
||||||
|
|
||||||
|
let content = parse_content(msg);
|
||||||
|
let tool_calls = parse_tool_calls(msg);
|
||||||
|
let tool_call_id = msg["tool_call_id"].as_str().map(|s| s.to_string());
|
||||||
|
let name = msg["name"].as_str().map(|s| s.to_string());
|
||||||
|
|
||||||
|
Ok(CanonicalMessage {
|
||||||
|
role,
|
||||||
|
content,
|
||||||
|
name,
|
||||||
|
tool_calls,
|
||||||
|
tool_call_id,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn parse_content(msg: &Value) -> Vec<ContentPart> {
|
||||||
|
match &msg["content"] {
|
||||||
|
Value::String(s) => {
|
||||||
|
vec![ContentPart::Text { text: s.clone() }]
|
||||||
|
}
|
||||||
|
Value::Array(arr) => arr
|
||||||
|
.iter()
|
||||||
|
.filter_map(|part| {
|
||||||
|
if let Some(text) = part["text"].as_str() {
|
||||||
|
Some(ContentPart::Text {
|
||||||
|
text: text.to_string(),
|
||||||
|
})
|
||||||
|
} else if let Some(image_url) = part["image_url"]["url"].as_str() {
|
||||||
|
Some(ContentPart::Image {
|
||||||
|
url: image_url.to_string(),
|
||||||
|
detail: part["image_url"]["detail"].as_str().map(|s| s.to_string()),
|
||||||
|
})
|
||||||
|
} else {
|
||||||
|
None
|
||||||
|
}
|
||||||
|
})
|
||||||
|
.collect(),
|
||||||
|
Value::Null => vec![],
|
||||||
|
_ => vec![],
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn parse_tool_calls(msg: &Value) -> Option<Vec<ToolCall>> {
|
||||||
|
msg["tool_calls"].as_array().map(|calls| {
|
||||||
|
calls
|
||||||
|
.iter()
|
||||||
|
.filter_map(|tc| {
|
||||||
|
let id = tc["id"].as_str()?.to_string();
|
||||||
|
let name = tc["function"]["name"].as_str()?.to_string();
|
||||||
|
let arguments = tc["function"]["arguments"].as_str()?;
|
||||||
|
Some(ToolCall {
|
||||||
|
id,
|
||||||
|
name,
|
||||||
|
arguments: serde_json::from_str(arguments).unwrap_or(Value::Null),
|
||||||
|
})
|
||||||
|
})
|
||||||
|
.collect()
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn parse_tool(v: &Value) -> Option<ToolDefinition> {
|
||||||
|
let name = v["function"]["name"].as_str()?;
|
||||||
|
Some(ToolDefinition {
|
||||||
|
name: name.to_string(),
|
||||||
|
description: v["function"]["description"].as_str().map(|s| s.to_string()),
|
||||||
|
input_schema: v["function"]["parameters"].clone(),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
use serde_json::json;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_parse_simple_request() {
|
||||||
|
let body = json!({
|
||||||
|
"model": "gpt-4o",
|
||||||
|
"messages": [
|
||||||
|
{"role": "system", "content": "You are helpful"},
|
||||||
|
{"role": "user", "content": "Hello!"}
|
||||||
|
],
|
||||||
|
"stream": false
|
||||||
|
});
|
||||||
|
let req = parse_chat_completions_request(&body).unwrap();
|
||||||
|
assert_eq!(req.model, "gpt-4o");
|
||||||
|
assert_eq!(req.messages.len(), 2);
|
||||||
|
assert!(!req.stream);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_parse_streaming_request() {
|
||||||
|
let body = json!({
|
||||||
|
"model": "gpt-4o",
|
||||||
|
"messages": [{"role": "user", "content": "Hi"}],
|
||||||
|
"stream": true,
|
||||||
|
"max_tokens": 100,
|
||||||
|
"temperature": 0.7,
|
||||||
|
});
|
||||||
|
let req = parse_chat_completions_request(&body).unwrap();
|
||||||
|
assert!(req.stream);
|
||||||
|
assert_eq!(req.max_tokens, Some(100));
|
||||||
|
assert!((req.temperature.unwrap() - 0.7).abs() < 0.001);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_parse_with_tools() {
|
||||||
|
let body = json!({
|
||||||
|
"model": "gpt-4o",
|
||||||
|
"messages": [{"role": "user", "content": "What's the weather?"}],
|
||||||
|
"tools": [{
|
||||||
|
"type": "function",
|
||||||
|
"function": {
|
||||||
|
"name": "get_weather",
|
||||||
|
"description": "Get weather",
|
||||||
|
"parameters": {"type": "object", "properties": {}}
|
||||||
|
}
|
||||||
|
}],
|
||||||
|
"tool_choice": "auto"
|
||||||
|
});
|
||||||
|
let req = parse_chat_completions_request(&body).unwrap();
|
||||||
|
assert!(req.tools.is_some());
|
||||||
|
assert_eq!(req.tools.as_ref().unwrap().len(), 1);
|
||||||
|
assert_eq!(req.tools.as_ref().unwrap()[0].name, "get_weather");
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -1 +1,8 @@
|
|||||||
// Transformers - OpenAI Chat
|
// OpenAI /v1/chat/completions transformers
|
||||||
|
pub mod inbound;
|
||||||
|
pub mod outbound;
|
||||||
|
pub mod stream;
|
||||||
|
|
||||||
|
pub use inbound::parse_chat_completions_request;
|
||||||
|
pub use outbound::format_chat_completions_response;
|
||||||
|
pub use stream::{format_done_event, format_stream_event, parse_stream_chunk};
|
||||||
|
|||||||
@@ -0,0 +1,141 @@
|
|||||||
|
// CanonicalResponse → OpenAI /v1/chat/completions JSON formatter
|
||||||
|
use crate::transformers::canonical::{
|
||||||
|
CanonicalChoice, CanonicalMessage, CanonicalResponse, ContentPart, Role, Usage,
|
||||||
|
};
|
||||||
|
use serde_json::{json, Value};
|
||||||
|
|
||||||
|
/// Format a canonical response into OpenAI /v1/chat/completions JSON
|
||||||
|
pub fn format_chat_completions_response(response: &CanonicalResponse) -> Value {
|
||||||
|
let choices: Vec<Value> = response.choices.iter().map(format_choice).collect();
|
||||||
|
|
||||||
|
let mut result = json!({
|
||||||
|
"id": response.id,
|
||||||
|
"object": "chat.completion",
|
||||||
|
"created": chrono::Utc::now().timestamp(),
|
||||||
|
"model": response.model,
|
||||||
|
"choices": choices,
|
||||||
|
});
|
||||||
|
|
||||||
|
if let Some(usage) = &response.usage {
|
||||||
|
result["usage"] = json!({
|
||||||
|
"prompt_tokens": usage.prompt_tokens,
|
||||||
|
"completion_tokens": usage.completion_tokens,
|
||||||
|
"total_tokens": usage.total_tokens,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
result
|
||||||
|
}
|
||||||
|
|
||||||
|
fn format_choice(choice: &CanonicalChoice) -> Value {
|
||||||
|
json!({
|
||||||
|
"index": choice.index,
|
||||||
|
"message": format_message(&choice.message),
|
||||||
|
"finish_reason": choice.finish_reason,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn format_message(msg: &CanonicalMessage) -> Value {
|
||||||
|
let role_str = match msg.role {
|
||||||
|
Role::System => "system",
|
||||||
|
Role::User => "user",
|
||||||
|
Role::Assistant => "assistant",
|
||||||
|
Role::Tool => "tool",
|
||||||
|
};
|
||||||
|
|
||||||
|
let content: Value = if msg.content.is_empty() {
|
||||||
|
Value::Null
|
||||||
|
} else if msg.content.len() == 1 {
|
||||||
|
match &msg.content[0] {
|
||||||
|
ContentPart::Text { text } => Value::String(text.clone()),
|
||||||
|
_ => format_content_parts(&msg.content),
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
format_content_parts(&msg.content)
|
||||||
|
};
|
||||||
|
|
||||||
|
let mut result = json!({
|
||||||
|
"role": role_str,
|
||||||
|
"content": content,
|
||||||
|
});
|
||||||
|
|
||||||
|
if let Some(name) = &msg.name {
|
||||||
|
result["name"] = json!(name);
|
||||||
|
}
|
||||||
|
|
||||||
|
if let Some(tool_calls) = &msg.tool_calls {
|
||||||
|
let formatted: Vec<Value> = tool_calls
|
||||||
|
.iter()
|
||||||
|
.map(|tc| {
|
||||||
|
json!({
|
||||||
|
"id": tc.id,
|
||||||
|
"type": "function",
|
||||||
|
"function": {
|
||||||
|
"name": tc.name,
|
||||||
|
"arguments": tc.arguments.to_string(),
|
||||||
|
}
|
||||||
|
})
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
result["tool_calls"] = json!(formatted);
|
||||||
|
}
|
||||||
|
|
||||||
|
if let Some(tool_call_id) = &msg.tool_call_id {
|
||||||
|
result["tool_call_id"] = json!(tool_call_id);
|
||||||
|
}
|
||||||
|
|
||||||
|
result
|
||||||
|
}
|
||||||
|
|
||||||
|
fn format_content_parts(parts: &[ContentPart]) -> Value {
|
||||||
|
let arr: Vec<Value> = parts
|
||||||
|
.iter()
|
||||||
|
.map(|part| match part {
|
||||||
|
ContentPart::Text { text } => json!({
|
||||||
|
"type": "text",
|
||||||
|
"text": text,
|
||||||
|
}),
|
||||||
|
ContentPart::Image { url, detail } => {
|
||||||
|
let mut img = json!({
|
||||||
|
"type": "image_url",
|
||||||
|
"image_url": {
|
||||||
|
"url": url,
|
||||||
|
}
|
||||||
|
});
|
||||||
|
if let Some(d) = detail {
|
||||||
|
img["image_url"]["detail"] = json!(d);
|
||||||
|
}
|
||||||
|
img
|
||||||
|
}
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
Value::Array(arr)
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
use crate::transformers::canonical::{CanonicalMessage, CanonicalResponse};
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_format_simple_response() {
|
||||||
|
let resp = CanonicalResponse::simple("gpt-4o", "Hello!");
|
||||||
|
let json = format_chat_completions_response(&resp);
|
||||||
|
assert_eq!(json["model"], "gpt-4o");
|
||||||
|
assert_eq!(json["object"], "chat.completion");
|
||||||
|
assert_eq!(json["choices"][0]["message"]["content"], "Hello!");
|
||||||
|
assert_eq!(json["choices"][0]["finish_reason"], "stop");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_format_with_usage() {
|
||||||
|
let mut resp = CanonicalResponse::simple("gpt-4o", "Hi");
|
||||||
|
resp.usage = Some(crate::transformers::canonical::Usage {
|
||||||
|
prompt_tokens: Some(10),
|
||||||
|
completion_tokens: Some(5),
|
||||||
|
total_tokens: Some(15),
|
||||||
|
});
|
||||||
|
let json = format_chat_completions_response(&resp);
|
||||||
|
assert_eq!(json["usage"]["total_tokens"], 15);
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,190 @@
|
|||||||
|
// OpenAI SSE stream parser and formatter
|
||||||
|
use crate::transformers::canonical::{
|
||||||
|
CanonicalStreamChoice, CanonicalStreamEvent, Role, StreamDelta, Usage,
|
||||||
|
};
|
||||||
|
use serde_json::{json, Value};
|
||||||
|
|
||||||
|
/// Parse a single SSE data line from OpenAI stream into a CanonicalStreamEvent
|
||||||
|
pub fn parse_stream_chunk(line: &str) -> Option<CanonicalStreamEvent> {
|
||||||
|
let data = line.strip_prefix("data: ")?.trim();
|
||||||
|
|
||||||
|
if data == "[DONE]" {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
|
||||||
|
let parsed: Value = serde_json::from_str(data).ok()?;
|
||||||
|
|
||||||
|
let id = parsed["id"].as_str().map(|s| s.to_string());
|
||||||
|
|
||||||
|
let choices: Vec<CanonicalStreamChoice> = parsed["choices"]
|
||||||
|
.as_array()
|
||||||
|
.map(|arr| {
|
||||||
|
arr.iter()
|
||||||
|
.filter_map(|c| {
|
||||||
|
let index = c["index"].as_i64().unwrap_or(0) as i32;
|
||||||
|
let finish_reason = c["finish_reason"].as_str().map(|s| s.to_string());
|
||||||
|
let delta = c.get("delta")?;
|
||||||
|
|
||||||
|
let role = delta["role"].as_str().and_then(|s| match s {
|
||||||
|
"system" => Some(Role::System),
|
||||||
|
"user" => Some(Role::User),
|
||||||
|
"assistant" => Some(Role::Assistant),
|
||||||
|
"tool" => Some(Role::Tool),
|
||||||
|
_ => None,
|
||||||
|
});
|
||||||
|
let content = delta["content"].as_str().map(|s| s.to_string());
|
||||||
|
|
||||||
|
let tool_calls = delta["tool_calls"].as_array().map(|tc_arr| {
|
||||||
|
tc_arr
|
||||||
|
.iter()
|
||||||
|
.filter_map(|tc| {
|
||||||
|
let index = tc["index"].as_i64().unwrap_or(0) as i32;
|
||||||
|
Some(crate::transformers::canonical::StreamToolCall {
|
||||||
|
index,
|
||||||
|
id: tc["id"].as_str().map(|s| s.to_string()),
|
||||||
|
name: tc["function"]["name"].as_str().map(|s| s.to_string()),
|
||||||
|
arguments: tc["function"]["arguments"]
|
||||||
|
.as_str()
|
||||||
|
.map(|s| s.to_string()),
|
||||||
|
})
|
||||||
|
})
|
||||||
|
.collect()
|
||||||
|
});
|
||||||
|
|
||||||
|
Some(CanonicalStreamChoice {
|
||||||
|
index,
|
||||||
|
delta: StreamDelta {
|
||||||
|
role,
|
||||||
|
content,
|
||||||
|
tool_calls,
|
||||||
|
},
|
||||||
|
finish_reason,
|
||||||
|
})
|
||||||
|
})
|
||||||
|
.collect()
|
||||||
|
})
|
||||||
|
.unwrap_or_default();
|
||||||
|
|
||||||
|
let usage = parsed["usage"].as_object().map(|_| Usage {
|
||||||
|
prompt_tokens: parsed["usage"]["prompt_tokens"].as_i64().map(|n| n as i32),
|
||||||
|
completion_tokens: parsed["usage"]["completion_tokens"]
|
||||||
|
.as_i64()
|
||||||
|
.map(|n| n as i32),
|
||||||
|
total_tokens: parsed["usage"]["total_tokens"].as_i64().map(|n| n as i32),
|
||||||
|
});
|
||||||
|
|
||||||
|
Some(CanonicalStreamEvent { id, choices, usage })
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Format a CanonicalStreamEvent back into an OpenAI SSE data line
|
||||||
|
pub fn format_stream_event(event: &CanonicalStreamEvent) -> String {
|
||||||
|
let choices: Vec<Value> = event
|
||||||
|
.choices
|
||||||
|
.iter()
|
||||||
|
.map(|c| {
|
||||||
|
let mut delta = json!({});
|
||||||
|
if let Some(ref role) = c.delta.role {
|
||||||
|
delta["role"] = json!(role.to_string());
|
||||||
|
}
|
||||||
|
if let Some(ref content) = c.delta.content {
|
||||||
|
delta["content"] = json!(content);
|
||||||
|
}
|
||||||
|
if let Some(ref tool_calls) = c.delta.tool_calls {
|
||||||
|
let tc_arr: Vec<Value> = tool_calls
|
||||||
|
.iter()
|
||||||
|
.map(|tc| {
|
||||||
|
let mut func = json!({});
|
||||||
|
if let Some(ref name) = tc.name {
|
||||||
|
func["name"] = json!(name);
|
||||||
|
}
|
||||||
|
if let Some(ref args) = tc.arguments {
|
||||||
|
func["arguments"] = json!(args);
|
||||||
|
}
|
||||||
|
json!({
|
||||||
|
"index": tc.index,
|
||||||
|
"id": tc.id,
|
||||||
|
"type": "function",
|
||||||
|
"function": func,
|
||||||
|
})
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
delta["tool_calls"] = json!(tc_arr);
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut choice = json!({
|
||||||
|
"index": c.index,
|
||||||
|
"delta": delta,
|
||||||
|
});
|
||||||
|
if let Some(ref fr) = c.finish_reason {
|
||||||
|
choice["finish_reason"] = json!(fr);
|
||||||
|
}
|
||||||
|
choice
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
let mut result = json!({
|
||||||
|
"choices": choices,
|
||||||
|
});
|
||||||
|
|
||||||
|
if let Some(ref id) = event.id {
|
||||||
|
result["id"] = json!(id);
|
||||||
|
result["object"] = json!("chat.completion.chunk");
|
||||||
|
}
|
||||||
|
|
||||||
|
if let Some(ref usage) = event.usage {
|
||||||
|
result["usage"] = json!({
|
||||||
|
"prompt_tokens": usage.prompt_tokens,
|
||||||
|
"completion_tokens": usage.completion_tokens,
|
||||||
|
"total_tokens": usage.total_tokens,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
format!(
|
||||||
|
"data: {}\n\n",
|
||||||
|
serde_json::to_string(&result).unwrap_or_default()
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Format the [DONE] event
|
||||||
|
pub fn format_done_event() -> &'static str {
|
||||||
|
"data: [DONE]\n\n"
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_parse_content_chunk() {
|
||||||
|
let line = r#"data: {"id":"123","choices":[{"index":0,"delta":{"content":"Hello"},"finish_reason":null}]}"#;
|
||||||
|
let event = parse_stream_chunk(line).unwrap();
|
||||||
|
assert_eq!(event.choices.len(), 1);
|
||||||
|
assert_eq!(event.choices[0].delta.content.as_deref(), Some("Hello"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_parse_done() {
|
||||||
|
let result = parse_stream_chunk("data: [DONE]");
|
||||||
|
assert!(result.is_none());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_format_roundtrip() {
|
||||||
|
let event = CanonicalStreamEvent {
|
||||||
|
id: Some("test-id".into()),
|
||||||
|
choices: vec![CanonicalStreamChoice {
|
||||||
|
index: 0,
|
||||||
|
delta: StreamDelta {
|
||||||
|
role: None,
|
||||||
|
content: Some("Hi".into()),
|
||||||
|
tool_calls: None,
|
||||||
|
},
|
||||||
|
finish_reason: None,
|
||||||
|
}],
|
||||||
|
usage: None,
|
||||||
|
};
|
||||||
|
let formatted = format_stream_event(&event);
|
||||||
|
assert!(formatted.starts_with("data: "));
|
||||||
|
assert!(formatted.contains("Hi"));
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -1 +1,8 @@
|
|||||||
// Transformers - OpenAI
|
// Transformers — OpenAI protocol
|
||||||
|
pub mod chat;
|
||||||
|
pub mod responses;
|
||||||
|
|
||||||
|
pub use chat::{
|
||||||
|
format_chat_completions_response, format_done_event, format_stream_event,
|
||||||
|
parse_chat_completions_request, parse_stream_chunk,
|
||||||
|
};
|
||||||
|
|||||||
@@ -1 +1,22 @@
|
|||||||
// Transformers - OpenAI Responses
|
use crate::transformers::canonical::{CanonicalRequest, CanonicalResponse};
|
||||||
|
use serde_json::{json, Value};
|
||||||
|
|
||||||
|
pub fn parse_responses_request(body: &Value) -> Result<CanonicalRequest, String> {
|
||||||
|
super::chat::parse_chat_completions_request(body)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn format_responses_response(response: &CanonicalResponse) -> Value {
|
||||||
|
let message = response.choices.first().map(|choice| &choice.message);
|
||||||
|
json!({
|
||||||
|
"id": response.id,
|
||||||
|
"object": "response",
|
||||||
|
"model": response.model,
|
||||||
|
"created_at": chrono::Utc::now().timestamp(),
|
||||||
|
"output": message.map(|message| json!([{
|
||||||
|
"type": "message",
|
||||||
|
"role": message.role.to_string(),
|
||||||
|
"content": message.content,
|
||||||
|
}])).unwrap_or_else(|| json!([])),
|
||||||
|
"usage": response.usage,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|||||||
@@ -1 +1,22 @@
|
|||||||
// Transformers - shared
|
use serde_json::Value;
|
||||||
|
|
||||||
|
pub fn string_field<'a>(body: &'a Value, key: &str) -> Option<&'a str> {
|
||||||
|
body.get(key).and_then(Value::as_str)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn bool_field(body: &Value, key: &str) -> Option<bool> {
|
||||||
|
body.get(key).and_then(Value::as_bool)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn array_field<'a>(body: &'a Value, key: &str) -> Option<&'a Vec<Value>> {
|
||||||
|
body.get(key).and_then(Value::as_array)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn merge_json_object(base: &mut Value, extra: Value) {
|
||||||
|
let (Some(base), Some(extra)) = (base.as_object_mut(), extra.as_object()) else {
|
||||||
|
return;
|
||||||
|
};
|
||||||
|
for (key, value) in extra {
|
||||||
|
base.insert(key.clone(), value.clone());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user