M0+M1: 基建 + 用户/密钥/核心代理
后端 (Go/Gin/GORM): - 配置(viper+env)、SQLite/Postgres 迁移、argon2id、AES-GCM 渠道密钥、JWT+refresh cookie - 用户注册/登录/刷新/登出、API Key CRUD(仅存哈希、明文一次展示) - 代理网关: /v1/chat/completions、/v1/responses、/v1/models 直通 OpenAI 渠道 非流式+流式(SSE 零缓冲转发), 用量捕获(chat 末块/responses completed 嵌套), OpenAI 错误格式(401/402/404/502), 余额检查 - 异步批量记账 + 余额流水 + 日聚合, admin 用户/余额/配置 API - 单测: crypto/jwt/apikey/流式 usage 提取 前端 (Vue3+TS+Vite+Tailwind v4): - taste-skill 设计 tokens: 深色仪表盘, 石墨+信号铜色, Outfit+JetBrains Mono - Landing/登录/注册, 控制台(仪表盘图表/密钥管理/用量明细) - 基础组件 Button/Input/Badge/Modal, ECharts 用量图 部署: docker-compose(nginx+api+postgres), 双 Dockerfile, nginx SSE 反代 联调: scripts/mockupstream 本地 mock 上游, 端到端验证通过
This commit is contained in:
@@ -0,0 +1,160 @@
|
||||
package api
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/openteam/server/internal/pkg/resp"
|
||||
"github.com/openteam/server/internal/store"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// AdminUsers GET /api/v1/admin/users — 用户列表(搜索、分页)。
|
||||
func (h *Handler) AdminUsers(c *gin.Context) {
|
||||
page, _ := strconv.Atoi(c.DefaultQuery("page", "1"))
|
||||
size, _ := strconv.Atoi(c.DefaultQuery("page_size", "20"))
|
||||
if page < 1 {
|
||||
page = 1
|
||||
}
|
||||
if size < 1 || size > 100 {
|
||||
size = 20
|
||||
}
|
||||
q := h.a.DB.Model(&store.User{})
|
||||
if kw := c.Query("q"); kw != "" {
|
||||
q = q.Where("username LIKE ? OR email LIKE ?", "%"+kw+"%", "%"+kw+"%")
|
||||
}
|
||||
var total int64
|
||||
q.Count(&total)
|
||||
var users []store.User
|
||||
q.Order("id DESC").Offset((page - 1) * size).Limit(size).Find(&users)
|
||||
out := make([]gin.H, 0, len(users))
|
||||
for _, u := range users {
|
||||
out = append(out, h.publicUser(&u))
|
||||
}
|
||||
resp.OK(c, gin.H{"items": out, "total": total, "page": page, "page_size": size})
|
||||
}
|
||||
|
||||
// AdminPatchUser PATCH /api/v1/admin/users/:id — 角色/状态/余额。
|
||||
func (h *Handler) AdminPatchUser(c *gin.Context) {
|
||||
id, err := strconv.ParseUint(c.Param("id"), 10, 64)
|
||||
if err != nil {
|
||||
resp.Fail(c, http.StatusBadRequest, "invalid user id")
|
||||
return
|
||||
}
|
||||
var req struct {
|
||||
Role *string `json:"role"`
|
||||
Status *string `json:"status"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
resp.Fail(c, http.StatusBadRequest, "invalid input")
|
||||
return
|
||||
}
|
||||
var u store.User
|
||||
if err := h.a.DB.First(&u, id).Error; err != nil {
|
||||
resp.Fail(c, http.StatusNotFound, "user not found")
|
||||
return
|
||||
}
|
||||
updates := map[string]any{}
|
||||
if req.Role != nil {
|
||||
if *req.Role != store.RoleUser && *req.Role != store.RoleAdmin {
|
||||
resp.Fail(c, http.StatusBadRequest, "role must be user or admin")
|
||||
return
|
||||
}
|
||||
updates["role"] = *req.Role
|
||||
}
|
||||
if req.Status != nil {
|
||||
if *req.Status != store.UserStatusActive && *req.Status != store.UserStatusDisabled {
|
||||
resp.Fail(c, http.StatusBadRequest, "status must be active or disabled")
|
||||
return
|
||||
}
|
||||
updates["status"] = *req.Status
|
||||
}
|
||||
if len(updates) > 0 {
|
||||
h.a.DB.Model(&u).Updates(updates)
|
||||
}
|
||||
resp.OK(c, gin.H{"ok": true})
|
||||
}
|
||||
|
||||
// AdminAdjustBalance POST /api/v1/admin/users/:id/balance — 调整余额并写流水。
|
||||
func (h *Handler) AdminAdjustBalance(c *gin.Context) {
|
||||
admin, _ := userFromContext(c)
|
||||
id, err := strconv.ParseUint(c.Param("id"), 10, 64)
|
||||
if err != nil {
|
||||
resp.Fail(c, http.StatusBadRequest, "invalid user id")
|
||||
return
|
||||
}
|
||||
var req struct {
|
||||
Change float64 `json:"change" binding:"required"`
|
||||
Remark string `json:"remark"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
resp.Fail(c, http.StatusBadRequest, "invalid input: change is required")
|
||||
return
|
||||
}
|
||||
var u store.User
|
||||
if err := h.a.DB.First(&u, id).Error; err != nil {
|
||||
resp.Fail(c, http.StatusNotFound, "user not found")
|
||||
return
|
||||
}
|
||||
newBalance := u.Balance + req.Change
|
||||
ref := fmt.Sprintf("admin:%d:%d", admin.ID, time.Now().UnixNano())
|
||||
err = h.a.DB.Transaction(func(tx *gorm.DB) error {
|
||||
if err := tx.Model(&store.User{}).Where("id = ?", u.ID).Update("balance", newBalance).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Create(&store.BalanceLog{
|
||||
UserID: u.ID,
|
||||
Change: req.Change,
|
||||
BalanceAfter: newBalance,
|
||||
Type: store.BalanceTypeAdminAdjust,
|
||||
RefID: ref,
|
||||
Remark: req.Remark,
|
||||
}).Error
|
||||
})
|
||||
if err != nil {
|
||||
resp.Fail(c, http.StatusInternalServerError, "failed to adjust balance")
|
||||
return
|
||||
}
|
||||
resp.OK(c, gin.H{"ok": true, "balance": newBalance})
|
||||
}
|
||||
|
||||
// AdminConfig GET /api/v1/admin/config
|
||||
func (h *Handler) AdminConfig(c *gin.Context) {
|
||||
var cfgs []store.SystemConfig
|
||||
h.a.DB.Find(&cfgs)
|
||||
m := map[string]any{}
|
||||
for _, cfg := range cfgs {
|
||||
var v any
|
||||
_ = json.Unmarshal([]byte(cfg.Value), &v)
|
||||
m[cfg.Key] = v
|
||||
}
|
||||
m["registration.mode"] = h.a.Cfg.Auth.RegistrationMode
|
||||
resp.OK(c, gin.H{"config": m})
|
||||
}
|
||||
|
||||
// AdminPutConfig PUT /api/v1/admin/config
|
||||
func (h *Handler) AdminPutConfig(c *gin.Context) {
|
||||
var req struct {
|
||||
Config map[string]any `json:"config"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
resp.Fail(c, http.StatusBadRequest, "invalid input")
|
||||
return
|
||||
}
|
||||
for k, v := range req.Config {
|
||||
if k == "registration.mode" {
|
||||
if v == "open" || v == "invite" {
|
||||
h.a.Cfg.Auth.RegistrationMode = v.(string)
|
||||
}
|
||||
continue
|
||||
}
|
||||
b, _ := json.Marshal(v)
|
||||
cfg := store.SystemConfig{Key: k, Value: string(b)}
|
||||
h.a.DB.Save(&cfg)
|
||||
}
|
||||
resp.OK(c, gin.H{"ok": true})
|
||||
}
|
||||
@@ -0,0 +1,197 @@
|
||||
// Package api 管理 API(/api/v1):认证、用户、密钥、用量、管理后台。
|
||||
package api
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/openteam/server/internal/app"
|
||||
"github.com/openteam/server/internal/pkg/resp"
|
||||
"github.com/openteam/server/internal/store"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// Handler 聚合所有管理 API。
|
||||
type Handler struct {
|
||||
a *app.App
|
||||
}
|
||||
|
||||
func NewHandler(a *app.App) *Handler { return &Handler{a: a} }
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// 认证
|
||||
|
||||
type registerReq struct {
|
||||
Username string `json:"username" binding:"required,min=3,max=32"`
|
||||
Email string `json:"email" binding:"required,email"`
|
||||
Password string `json:"password" binding:"required,min=8,max=72"`
|
||||
InviteCode string `json:"invite_code"`
|
||||
}
|
||||
|
||||
// Register POST /api/v1/auth/register
|
||||
func (h *Handler) Register(c *gin.Context) {
|
||||
var req registerReq
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
resp.Fail(c, http.StatusBadRequest, "invalid input: "+err.Error())
|
||||
return
|
||||
}
|
||||
if h.a.Cfg.Auth.RegistrationMode == "invite" {
|
||||
var ic string
|
||||
h.a.DB.Model(&store.SystemConfig{}).Where("key = ?", "invite_codes").Pluck("value", &ic)
|
||||
if !strings.Contains(ic, req.InviteCode) {
|
||||
resp.Fail(c, http.StatusForbidden, "valid invite code required")
|
||||
return
|
||||
}
|
||||
}
|
||||
var count int64
|
||||
h.a.DB.Model(&store.User{}).Where("username = ? OR email = ?", req.Username, req.Email).Count(&count)
|
||||
if count > 0 {
|
||||
resp.Fail(c, http.StatusConflict, "username or email already exists")
|
||||
return
|
||||
}
|
||||
hash, err := h.a.Hasher.HashPassword(req.Password)
|
||||
if err != nil {
|
||||
resp.Fail(c, http.StatusInternalServerError, "failed to hash password")
|
||||
return
|
||||
}
|
||||
u := store.User{
|
||||
Username: req.Username,
|
||||
Email: strings.ToLower(req.Email),
|
||||
PasswordHash: hash,
|
||||
Role: store.RoleUser,
|
||||
Balance: 5.0, // 新用户赠送体验余额(可通过 admin 调整)
|
||||
Status: store.UserStatusActive,
|
||||
}
|
||||
if req.InviteCode != "" {
|
||||
ic := req.InviteCode
|
||||
u.InviteCode = &ic
|
||||
}
|
||||
if err := h.a.DB.Create(&u).Error; err != nil {
|
||||
resp.Fail(c, http.StatusInternalServerError, "failed to create user")
|
||||
return
|
||||
}
|
||||
resp.Created(c, gin.H{"id": u.ID, "username": u.Username})
|
||||
}
|
||||
|
||||
type loginReq struct {
|
||||
Username string `json:"username" binding:"required"`
|
||||
Password string `json:"password" binding:"required"`
|
||||
}
|
||||
|
||||
// Login POST /api/v1/auth/login — 返回 access token,refresh token 写入 HttpOnly Cookie。
|
||||
func (h *Handler) Login(c *gin.Context) {
|
||||
var req loginReq
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
resp.Fail(c, http.StatusBadRequest, "invalid input")
|
||||
return
|
||||
}
|
||||
var u store.User
|
||||
err := h.a.DB.Where("username = ? OR email = ?", req.Username, strings.ToLower(req.Username)).First(&u).Error
|
||||
if err != nil || u.Status != store.UserStatusActive {
|
||||
resp.Fail(c, http.StatusUnauthorized, "invalid credentials")
|
||||
return
|
||||
}
|
||||
ok, err := h.a.Hasher.VerifyPassword(u.PasswordHash, req.Password)
|
||||
if err != nil || !ok {
|
||||
resp.Fail(c, http.StatusUnauthorized, "invalid credentials")
|
||||
return
|
||||
}
|
||||
now := time.Now()
|
||||
h.a.DB.Model(&u).Update("last_login_at", now)
|
||||
|
||||
access, _, err := h.a.JWT.Sign(u.ID, u.Username, u.Role, "access")
|
||||
if err != nil {
|
||||
resp.Fail(c, http.StatusInternalServerError, "failed to issue token")
|
||||
return
|
||||
}
|
||||
refresh, _, err := h.a.JWT.Sign(u.ID, u.Username, u.Role, "refresh")
|
||||
if err != nil {
|
||||
resp.Fail(c, http.StatusInternalServerError, "failed to issue token")
|
||||
return
|
||||
}
|
||||
h.setRefreshCookie(c, refresh)
|
||||
resp.OK(c, gin.H{"access_token": access, "expires_in": int(h.a.JWT.AccessTTL().Seconds()), "user": h.publicUser(&u)})
|
||||
}
|
||||
|
||||
// Refresh POST /api/v1/auth/refresh — 用 refresh cookie 换新 access token。
|
||||
func (h *Handler) Refresh(c *gin.Context) {
|
||||
tok, err := c.Cookie(h.a.Cfg.JWT.CookieName)
|
||||
if err != nil || tok == "" {
|
||||
resp.Fail(c, http.StatusUnauthorized, "refresh token missing")
|
||||
return
|
||||
}
|
||||
claims, err := h.a.JWT.Parse(tok)
|
||||
if err != nil || claims.Subject != "refresh" {
|
||||
resp.Fail(c, http.StatusUnauthorized, "invalid refresh token")
|
||||
return
|
||||
}
|
||||
var u store.User
|
||||
if err := h.a.DB.First(&u, claims.UserID).Error; err != nil || u.Status != store.UserStatusActive {
|
||||
resp.Fail(c, http.StatusUnauthorized, "user not found or disabled")
|
||||
return
|
||||
}
|
||||
access, _, err := h.a.JWT.Sign(u.ID, u.Username, u.Role, "access")
|
||||
if err != nil {
|
||||
resp.Fail(c, http.StatusInternalServerError, "failed to issue token")
|
||||
return
|
||||
}
|
||||
resp.OK(c, gin.H{"access_token": access, "expires_in": int(h.a.JWT.AccessTTL().Seconds())})
|
||||
}
|
||||
|
||||
// Logout POST /api/v1/auth/logout — 清除 refresh cookie。
|
||||
func (h *Handler) Logout(c *gin.Context) {
|
||||
h.setRefreshCookie(c, "")
|
||||
resp.OK(c, gin.H{"ok": true})
|
||||
}
|
||||
|
||||
// Me GET /api/v1/auth/me
|
||||
func (h *Handler) Me(c *gin.Context) {
|
||||
u := sessionUser(c)
|
||||
resp.OK(c, gin.H{"user": h.publicUser(u)})
|
||||
}
|
||||
|
||||
func (h *Handler) setRefreshCookie(c *gin.Context, value string) {
|
||||
maxAge := int(h.a.JWT.RefreshTTL().Seconds())
|
||||
if value == "" {
|
||||
maxAge = -1
|
||||
}
|
||||
http.SetCookie(c.Writer, &http.Cookie{
|
||||
Name: h.a.Cfg.JWT.CookieName,
|
||||
Value: value,
|
||||
Path: "/",
|
||||
MaxAge: maxAge,
|
||||
HttpOnly: true,
|
||||
Secure: h.a.Cfg.JWT.CookieSecure,
|
||||
SameSite: http.SameSiteLaxMode,
|
||||
Domain: h.a.Cfg.JWT.CookieDomain,
|
||||
})
|
||||
}
|
||||
|
||||
func (h *Handler) publicUser(u *store.User) gin.H {
|
||||
return gin.H{
|
||||
"id": u.ID,
|
||||
"username": u.Username,
|
||||
"email": u.Email,
|
||||
"role": u.Role,
|
||||
"balance": u.Balance,
|
||||
"status": u.Status,
|
||||
"created_at": u.CreatedAt,
|
||||
}
|
||||
}
|
||||
|
||||
func sessionUser(c *gin.Context) *store.User {
|
||||
u, _ := c.Get("session_user")
|
||||
return u.(*store.User)
|
||||
}
|
||||
|
||||
func userFromContext(c *gin.Context) (*store.User, bool) {
|
||||
u, ok := c.Get("session_user")
|
||||
if !ok {
|
||||
return nil, false
|
||||
}
|
||||
return u.(*store.User), true
|
||||
}
|
||||
|
||||
var _ = gorm.ErrRecordNotFound
|
||||
@@ -0,0 +1,181 @@
|
||||
package api
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/openteam/server/internal/pkg/apikey"
|
||||
"github.com/openteam/server/internal/pkg/resp"
|
||||
"github.com/openteam/server/internal/store"
|
||||
)
|
||||
|
||||
type createKeyReq struct {
|
||||
Name string `json:"name" binding:"required,min=1,max=64"`
|
||||
QuotaTokensPerDay *int64 `json:"quota_tokens_per_day"`
|
||||
QuotaRequestsPerDay *int `json:"quota_requests_per_day"`
|
||||
AllowedModels []string `json:"allowed_models"`
|
||||
ExpiresAt *string `json:"expires_at"` // RFC3339
|
||||
}
|
||||
|
||||
// CreateKey POST /api/v1/keys — 创建密钥,明文仅此一次返回。
|
||||
func (h *Handler) CreateKey(c *gin.Context) {
|
||||
u, ok := userFromContext(c)
|
||||
if !ok {
|
||||
resp.Fail(c, http.StatusUnauthorized, "authentication required")
|
||||
return
|
||||
}
|
||||
var req createKeyReq
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
resp.Fail(c, http.StatusBadRequest, "invalid input: "+err.Error())
|
||||
return
|
||||
}
|
||||
plain, hash, prefix, err := apikey.Generate()
|
||||
if err != nil {
|
||||
resp.Fail(c, http.StatusInternalServerError, "failed to generate key")
|
||||
return
|
||||
}
|
||||
k := store.APIKey{
|
||||
UserID: u.ID,
|
||||
Name: req.Name,
|
||||
KeyHash: hash,
|
||||
KeyPrefix: prefix,
|
||||
QuotaTokensPerDay: req.QuotaTokensPerDay,
|
||||
QuotaRequestsPerDay: req.QuotaRequestsPerDay,
|
||||
AllowedModels: req.AllowedModels,
|
||||
Status: store.KeyStatusActive,
|
||||
}
|
||||
if req.ExpiresAt != nil {
|
||||
t, err := time.Parse(time.RFC3339, *req.ExpiresAt)
|
||||
if err != nil {
|
||||
resp.Fail(c, http.StatusBadRequest, "expires_at must be RFC3339")
|
||||
return
|
||||
}
|
||||
k.ExpiresAt = &t
|
||||
}
|
||||
if err := h.a.DB.Create(&k).Error; err != nil {
|
||||
resp.Fail(c, http.StatusInternalServerError, "failed to create key")
|
||||
return
|
||||
}
|
||||
resp.Created(c, gin.H{
|
||||
"id": k.ID,
|
||||
"name": k.Name,
|
||||
"key": plain, // 仅此一次
|
||||
"key_prefix": k.KeyPrefix,
|
||||
"created_at": k.CreatedAt,
|
||||
})
|
||||
}
|
||||
|
||||
// ListKeys GET /api/v1/keys
|
||||
func (h *Handler) ListKeys(c *gin.Context) {
|
||||
u, ok := userFromContext(c)
|
||||
if !ok {
|
||||
resp.Fail(c, http.StatusUnauthorized, "authentication required")
|
||||
return
|
||||
}
|
||||
var keys []store.APIKey
|
||||
if err := h.a.DB.Where("user_id = ?", u.ID).Order("id DESC").Find(&keys).Error; err != nil {
|
||||
resp.Fail(c, http.StatusInternalServerError, "failed to load keys")
|
||||
return
|
||||
}
|
||||
out := make([]gin.H, 0, len(keys))
|
||||
for _, k := range keys {
|
||||
out = append(out, gin.H{
|
||||
"id": k.ID,
|
||||
"name": k.Name,
|
||||
"key_prefix": k.KeyPrefix,
|
||||
"quota_tokens_per_day": k.QuotaTokensPerDay,
|
||||
"quota_requests_per_day": k.QuotaRequestsPerDay,
|
||||
"allowed_models": k.AllowedModels,
|
||||
"expires_at": k.ExpiresAt,
|
||||
"status": k.Status,
|
||||
"last_used_at": k.LastUsedAt,
|
||||
"created_at": k.CreatedAt,
|
||||
})
|
||||
}
|
||||
resp.OK(c, gin.H{"items": out})
|
||||
}
|
||||
|
||||
// PatchKey PATCH /api/v1/keys/:id — 改名、限额、白名单、启停。
|
||||
func (h *Handler) PatchKey(c *gin.Context) {
|
||||
u, ok := userFromContext(c)
|
||||
if !ok {
|
||||
resp.Fail(c, http.StatusUnauthorized, "authentication required")
|
||||
return
|
||||
}
|
||||
id, err := strconv.ParseUint(c.Param("id"), 10, 64)
|
||||
if err != nil {
|
||||
resp.Fail(c, http.StatusBadRequest, "invalid key id")
|
||||
return
|
||||
}
|
||||
var req struct {
|
||||
Name *string `json:"name"`
|
||||
QuotaTokensPerDay *int64 `json:"quota_tokens_per_day"`
|
||||
QuotaRequestsPerDay *int `json:"quota_requests_per_day"`
|
||||
AllowedModels *[]string `json:"allowed_models"`
|
||||
Status *string `json:"status"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
resp.Fail(c, http.StatusBadRequest, "invalid input")
|
||||
return
|
||||
}
|
||||
var k store.APIKey
|
||||
if err := h.a.DB.Where("id = ? AND user_id = ?", id, u.ID).First(&k).Error; err != nil {
|
||||
resp.Fail(c, http.StatusNotFound, "key not found")
|
||||
return
|
||||
}
|
||||
updates := map[string]any{}
|
||||
if req.Name != nil {
|
||||
updates["name"] = *req.Name
|
||||
}
|
||||
if req.QuotaTokensPerDay != nil {
|
||||
updates["quota_tokens_per_day"] = *req.QuotaTokensPerDay
|
||||
}
|
||||
if req.QuotaRequestsPerDay != nil {
|
||||
updates["quota_requests_per_day"] = *req.QuotaRequestsPerDay
|
||||
}
|
||||
if req.AllowedModels != nil {
|
||||
updates["allowed_models"] = *req.AllowedModels
|
||||
}
|
||||
if req.Status != nil {
|
||||
if *req.Status != store.KeyStatusActive && *req.Status != store.KeyStatusRevoked {
|
||||
resp.Fail(c, http.StatusBadRequest, "status must be active or revoked")
|
||||
return
|
||||
}
|
||||
updates["status"] = *req.Status
|
||||
}
|
||||
if len(updates) > 0 {
|
||||
if err := h.a.DB.Model(&k).Updates(updates).Error; err != nil {
|
||||
resp.Fail(c, http.StatusInternalServerError, "failed to update key")
|
||||
return
|
||||
}
|
||||
}
|
||||
resp.OK(c, gin.H{"ok": true})
|
||||
}
|
||||
|
||||
// DeleteKey DELETE /api/v1/keys/:id — 吊销。
|
||||
func (h *Handler) DeleteKey(c *gin.Context) {
|
||||
u, ok := userFromContext(c)
|
||||
if !ok {
|
||||
resp.Fail(c, http.StatusUnauthorized, "authentication required")
|
||||
return
|
||||
}
|
||||
id, err := strconv.ParseUint(c.Param("id"), 10, 64)
|
||||
if err != nil {
|
||||
resp.Fail(c, http.StatusBadRequest, "invalid key id")
|
||||
return
|
||||
}
|
||||
res := h.a.DB.Model(&store.APIKey{}).
|
||||
Where("id = ? AND user_id = ?", id, u.ID).
|
||||
Update("status", store.KeyStatusRevoked)
|
||||
if res.Error != nil {
|
||||
resp.Fail(c, http.StatusInternalServerError, "failed to revoke key")
|
||||
return
|
||||
}
|
||||
if res.RowsAffected == 0 {
|
||||
resp.Fail(c, http.StatusNotFound, "key not found")
|
||||
return
|
||||
}
|
||||
resp.OK(c, gin.H{"ok": true})
|
||||
}
|
||||
@@ -0,0 +1,79 @@
|
||||
// Package middleware Gin 中间件:会话鉴权、admin 校验。
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/openteam/server/internal/app"
|
||||
"github.com/openteam/server/internal/pkg/resp"
|
||||
"github.com/openteam/server/internal/store"
|
||||
)
|
||||
|
||||
const CtxSessionUser = "session_user"
|
||||
|
||||
// SessionAuth 会话鉴权:Authorization: Bearer <access token>。
|
||||
func SessionAuth(a *app.App) gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
auth := c.GetHeader("Authorization")
|
||||
if len(auth) < 8 || auth[:7] != "Bearer " {
|
||||
resp.Fail(c, http.StatusUnauthorized, "missing or invalid access token")
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
claims, err := a.JWT.Parse(auth[7:])
|
||||
if err != nil || claims.Subject != "access" {
|
||||
resp.Fail(c, http.StatusUnauthorized, "invalid or expired access token")
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
var u store.User
|
||||
if err := a.DB.First(&u, claims.UserID).Error; err != nil {
|
||||
resp.Fail(c, http.StatusUnauthorized, "user not found")
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
if u.Status != store.UserStatusActive {
|
||||
resp.Fail(c, http.StatusForbidden, "user account disabled")
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
c.Set(CtxSessionUser, &u)
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
|
||||
// AdminOnly 仅管理员。
|
||||
func AdminOnly(c *gin.Context) {
|
||||
u, ok := c.Get(CtxSessionUser)
|
||||
if !ok {
|
||||
resp.Fail(c, http.StatusUnauthorized, "authentication required")
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
if u.(*store.User).Role != store.RoleAdmin {
|
||||
resp.Fail(c, http.StatusForbidden, "admin permission required")
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
c.Next()
|
||||
}
|
||||
|
||||
// CORS 开发期放开;生产按配置白名单(PLANNING §8)。
|
||||
func CORS() gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
origin := c.GetHeader("Origin")
|
||||
if origin != "" {
|
||||
c.Header("Access-Control-Allow-Origin", origin)
|
||||
c.Header("Vary", "Origin")
|
||||
c.Header("Access-Control-Allow-Credentials", "true")
|
||||
c.Header("Access-Control-Allow-Headers", "Authorization, Content-Type, X-Requested-With")
|
||||
c.Header("Access-Control-Allow-Methods", "GET, POST, PUT, PATCH, DELETE, OPTIONS")
|
||||
}
|
||||
if c.Request.Method == http.MethodOptions {
|
||||
c.AbortWithStatus(http.StatusNoContent)
|
||||
return
|
||||
}
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,84 @@
|
||||
package api
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/openteam/server/internal/api/middleware"
|
||||
"github.com/openteam/server/internal/app"
|
||||
"github.com/openteam/server/internal/proxy"
|
||||
)
|
||||
|
||||
// NewRouter 装配所有路由:
|
||||
// - /v1/* 代理端点(Bearer API Key,OpenAI 格式错误体)
|
||||
// - /api/v1/* 管理 API(会话 JWT)
|
||||
func NewRouter(a *app.App, gw *proxy.Gateway) *gin.Engine {
|
||||
if a.Cfg.Env == "production" {
|
||||
gin.SetMode(gin.ReleaseMode)
|
||||
}
|
||||
r := gin.New()
|
||||
r.Use(gin.Logger(), gin.Recovery(), middleware.CORS())
|
||||
|
||||
h := NewHandler(a)
|
||||
|
||||
// --- 代理端点(对外)---
|
||||
proxyGroup := r.Group("/v1")
|
||||
{
|
||||
proxyGroup.Any("/chat/completions", gw.Auth, gw.Handle)
|
||||
proxyGroup.Any("/responses", gw.Auth, gw.Handle)
|
||||
proxyGroup.Any("/models", gw.Auth, gw.Handle)
|
||||
}
|
||||
// 未匹配的 /v1/* 返回 OpenAI 风格 404(需先认证)
|
||||
r.NoRoute(func(c *gin.Context) {
|
||||
if len(c.Request.URL.Path) >= 3 && c.Request.URL.Path[:3] == "/v1" {
|
||||
gw.Auth(c)
|
||||
if !c.IsAborted() {
|
||||
gw.Handle(c)
|
||||
}
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "not found"})
|
||||
})
|
||||
|
||||
// --- 管理 API ---
|
||||
api := r.Group("/api/v1")
|
||||
{
|
||||
auth := api.Group("/auth")
|
||||
{
|
||||
auth.POST("/register", h.Register)
|
||||
auth.POST("/login", h.Login)
|
||||
auth.POST("/refresh", h.Refresh)
|
||||
auth.POST("/logout", h.Logout)
|
||||
auth.GET("/me", middleware.SessionAuth(a), h.Me)
|
||||
}
|
||||
|
||||
user := api.Group("", middleware.SessionAuth(a))
|
||||
{
|
||||
user.GET("/user/profile", h.UserProfile)
|
||||
user.GET("/user/balance", h.UserBalance)
|
||||
user.GET("/user/models", h.UserModels)
|
||||
user.GET("/usage/summary", h.UsageSummary)
|
||||
user.GET("/usage/stats", h.UsageStats)
|
||||
user.GET("/usage/logs", h.UsageLogs)
|
||||
user.POST("/keys", h.CreateKey)
|
||||
user.GET("/keys", h.ListKeys)
|
||||
user.PATCH("/keys/:id", h.PatchKey)
|
||||
user.DELETE("/keys/:id", h.DeleteKey)
|
||||
}
|
||||
|
||||
admin := api.Group("/admin", middleware.SessionAuth(a), middleware.AdminOnly)
|
||||
{
|
||||
admin.GET("/users", h.AdminUsers)
|
||||
admin.PATCH("/users/:id", h.AdminPatchUser)
|
||||
admin.POST("/users/:id/balance", h.AdminAdjustBalance)
|
||||
admin.GET("/config", h.AdminConfig)
|
||||
admin.PUT("/config", h.AdminPutConfig)
|
||||
// 渠道/模型/用量管理(M4);充值审核(M5 预留)
|
||||
}
|
||||
}
|
||||
|
||||
r.GET("/healthz", func(c *gin.Context) {
|
||||
c.JSON(http.StatusOK, gin.H{"status": "ok"})
|
||||
})
|
||||
return r
|
||||
}
|
||||
@@ -0,0 +1,165 @@
|
||||
package api
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/openteam/server/internal/pkg/resp"
|
||||
"github.com/openteam/server/internal/store"
|
||||
)
|
||||
|
||||
// UserProfile GET /api/v1/user/profile
|
||||
func (h *Handler) UserProfile(c *gin.Context) {
|
||||
u, _ := userFromContext(c)
|
||||
resp.OK(c, gin.H{"user": h.publicUser(u)})
|
||||
}
|
||||
|
||||
// UserBalance GET /api/v1/user/balance — 余额 + 近 30 日消耗。
|
||||
func (h *Handler) UserBalance(c *gin.Context) {
|
||||
u, _ := userFromContext(c)
|
||||
var spent float64
|
||||
h.a.DB.Model(&store.UsageLog{}).
|
||||
Where("user_id = ? AND status = ? AND created_at >= ?", u.ID, store.UsageStatusSuccess, time.Now().Add(-30*24*time.Hour)).
|
||||
Select("COALESCE(SUM(cost),0)").Scan(&spent)
|
||||
resp.OK(c, gin.H{
|
||||
"balance": u.Balance,
|
||||
"spent_last_30d": spent,
|
||||
"today": h.todayUsage(c, u.ID),
|
||||
"models_available": h.availableModelCount(),
|
||||
})
|
||||
}
|
||||
|
||||
func (h *Handler) availableModelCount() int64 {
|
||||
var n int64
|
||||
h.a.DB.Model(&store.Model{}).Where("enabled = ?", true).Count(&n)
|
||||
return n
|
||||
}
|
||||
|
||||
func (h *Handler) todayUsage(c *gin.Context, userID uint64) gin.H {
|
||||
var requests int64
|
||||
var tokens int64
|
||||
var cost float64
|
||||
today := time.Now().UTC().Format("2006-01-02")
|
||||
h.a.DB.Model(&store.UsageDaily{}).
|
||||
Where("user_id = ? AND date = ?", userID, today).
|
||||
Select("COALESCE(SUM(requests),0), COALESCE(SUM(input_tokens+output_tokens+cache_read_tokens),0), COALESCE(SUM(cost),0)").
|
||||
Row().Scan(&requests, &tokens, &cost)
|
||||
return gin.H{"requests": requests, "tokens": tokens, "cost": cost}
|
||||
}
|
||||
|
||||
// UserModels GET /api/v1/user/models — 控制台可用模型列表(无需 API Key)。
|
||||
func (h *Handler) UserModels(c *gin.Context) {
|
||||
var ms []store.Model
|
||||
if err := h.a.DB.Where("enabled = ?", true).Order("sort ASC, id ASC").Find(&ms).Error; err != nil {
|
||||
resp.Fail(c, http.StatusInternalServerError, "failed to load models")
|
||||
return
|
||||
}
|
||||
out := make([]string, 0, len(ms))
|
||||
for _, m := range ms {
|
||||
out = append(out, m.Name)
|
||||
}
|
||||
resp.OK(c, gin.H{"items": out})
|
||||
}
|
||||
|
||||
// UsageSummary GET /api/v1/usage/summary — 今日/本月汇总。
|
||||
func (h *Handler) UsageSummary(c *gin.Context) {
|
||||
u, _ := userFromContext(c)
|
||||
now := time.Now().UTC()
|
||||
today := now.Format("2006-01-02")
|
||||
month := now.Format("2006-01")
|
||||
var todayReq, monthReq int64
|
||||
var todayTok, monthTok int64
|
||||
var todayCost, monthCost float64
|
||||
h.a.DB.Model(&store.UsageDaily{}).Where("user_id = ? AND date = ?", u.ID, today).
|
||||
Select("COALESCE(SUM(requests),0), COALESCE(SUM(input_tokens+output_tokens+cache_read_tokens),0), COALESCE(SUM(cost),0)").
|
||||
Row().Scan(&todayReq, &todayTok, &todayCost)
|
||||
h.a.DB.Model(&store.UsageDaily{}).Where("user_id = ? AND date LIKE ?", u.ID, month+"%").
|
||||
Select("COALESCE(SUM(requests),0), COALESCE(SUM(input_tokens+output_tokens+cache_read_tokens),0), COALESCE(SUM(cost),0)").
|
||||
Row().Scan(&monthReq, &monthTok, &monthCost)
|
||||
resp.OK(c, gin.H{
|
||||
"today": gin.H{"requests": todayReq, "tokens": todayTok, "cost": todayCost},
|
||||
"month": gin.H{"requests": monthReq, "tokens": monthTok, "cost": monthCost},
|
||||
})
|
||||
}
|
||||
|
||||
// UsageStats GET /api/v1/usage/stats?from&to&group=day|model
|
||||
func (h *Handler) UsageStats(c *gin.Context) {
|
||||
u, _ := userFromContext(c)
|
||||
from := c.DefaultQuery("from", time.Now().Add(-30*24*time.Hour).Format("2006-01-02"))
|
||||
to := c.DefaultQuery("to", time.Now().Format("2006-01-02"))
|
||||
group := c.DefaultQuery("group", "day")
|
||||
|
||||
q := h.a.DB.Model(&store.UsageDaily{}).Where("user_id = ? AND date BETWEEN ? AND ?", u.ID, from, to)
|
||||
out := make([]gin.H, 0, 64)
|
||||
if group == "model" {
|
||||
var rows []struct {
|
||||
ModelID uint64
|
||||
Requests int64
|
||||
Tokens int64
|
||||
Cost float64
|
||||
}
|
||||
q.Select("model_id, SUM(requests) requests, SUM(input_tokens+output_tokens+cache_read_tokens) tokens, SUM(cost) cost").
|
||||
Group("model_id").Scan(&rows)
|
||||
for _, r := range rows {
|
||||
var m store.Model
|
||||
name := strconv.FormatUint(r.ModelID, 10)
|
||||
if h.a.DB.First(&m, r.ModelID).Error == nil {
|
||||
name = m.Name
|
||||
}
|
||||
out = append(out, gin.H{"model": name, "requests": r.Requests, "tokens": r.Tokens, "cost": r.Cost})
|
||||
}
|
||||
} else {
|
||||
var rows []struct {
|
||||
Date string
|
||||
Requests int64
|
||||
Tokens int64
|
||||
Cost float64
|
||||
}
|
||||
q.Select("date, SUM(requests) requests, SUM(input_tokens+output_tokens+cache_read_tokens) tokens, SUM(cost) cost").
|
||||
Group("date").Order("date").Scan(&rows)
|
||||
for _, r := range rows {
|
||||
out = append(out, gin.H{"date": r.Date, "requests": r.Requests, "tokens": r.Tokens, "cost": r.Cost})
|
||||
}
|
||||
}
|
||||
resp.OK(c, gin.H{"items": out})
|
||||
}
|
||||
|
||||
// UsageLogs GET /api/v1/usage/logs?from&to&page&page_size&model
|
||||
func (h *Handler) UsageLogs(c *gin.Context) {
|
||||
u, _ := userFromContext(c)
|
||||
page, _ := strconv.Atoi(c.DefaultQuery("page", "1"))
|
||||
size, _ := strconv.Atoi(c.DefaultQuery("page_size", "20"))
|
||||
if page < 1 {
|
||||
page = 1
|
||||
}
|
||||
if size < 1 || size > 100 {
|
||||
size = 20
|
||||
}
|
||||
q := h.a.DB.Model(&store.UsageLog{}).Where("user_id = ?", u.ID)
|
||||
if from := c.Query("from"); from != "" {
|
||||
q = q.Where("created_at >= ?", from+" 00:00:00")
|
||||
}
|
||||
if to := c.Query("to"); to != "" {
|
||||
q = q.Where("created_at <= ?", to+" 23:59:59")
|
||||
}
|
||||
if model := c.Query("model"); model != "" {
|
||||
q = q.Where("model_name = ?", model)
|
||||
}
|
||||
var total int64
|
||||
q.Count(&total)
|
||||
var logs []store.UsageLog
|
||||
q.Order("id DESC").Offset((page - 1) * size).Limit(size).Find(&logs)
|
||||
out := make([]gin.H, 0, len(logs))
|
||||
for _, l := range logs {
|
||||
out = append(out, gin.H{
|
||||
"id": l.ID, "request_id": l.RequestID, "model": l.ModelName, "protocol": l.Protocol,
|
||||
"input_tokens": l.InputTokens, "output_tokens": l.OutputTokens,
|
||||
"cache_read_tokens": l.CacheReadTokens, "cost": l.Cost,
|
||||
"latency_ms": l.LatencyMS, "status": l.Status, "error_code": l.ErrorCode,
|
||||
"created_at": l.CreatedAt,
|
||||
})
|
||||
}
|
||||
resp.OK(c, gin.H{"items": out, "total": total, "page": page, "page_size": size})
|
||||
}
|
||||
@@ -0,0 +1,132 @@
|
||||
// Package app 应用容器:装配配置、数据库、密码/加密/JWT 与记账器。
|
||||
package app
|
||||
|
||||
import (
|
||||
"context"
|
||||
"log"
|
||||
"time"
|
||||
|
||||
"github.com/openteam/server/internal/config"
|
||||
"github.com/openteam/server/internal/pkg/crypto"
|
||||
"github.com/openteam/server/internal/pkg/jwt"
|
||||
"github.com/openteam/server/internal/store"
|
||||
"github.com/openteam/server/internal/usage"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
type App struct {
|
||||
Cfg *config.Config
|
||||
DB *gorm.DB
|
||||
Hasher *crypto.PasswordHasher
|
||||
Enc *crypto.Encryptor
|
||||
JWT *jwt.Manager
|
||||
Usage *usage.Recorder
|
||||
startedAt time.Time
|
||||
}
|
||||
|
||||
func New(cfg *config.Config) (*App, error) {
|
||||
db, err := store.Open(cfg.DB.Driver, cfg.DB.DSN)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
a := &App{
|
||||
Cfg: cfg,
|
||||
DB: db,
|
||||
Hasher: crypto.NewPasswordHasher(cfg.Auth.Argon2Time, cfg.Auth.Argon2Memory, cfg.Auth.Argon2Threads, cfg.Auth.Argon2KeyLen, cfg.Auth.SaltLen),
|
||||
Enc: crypto.NewEncryptor(cfg.Master),
|
||||
JWT: jwt.NewManager(cfg.JWT.Secret, cfg.JWT.Issuer, cfg.JWT.AccessTTL, cfg.JWT.RefreshTTL),
|
||||
startedAt: time.Now(),
|
||||
}
|
||||
a.Usage = usage.NewRecorder(db)
|
||||
|
||||
if err := a.Seed(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return a, nil
|
||||
}
|
||||
|
||||
func (a *App) Close() {
|
||||
a.Usage.Close()
|
||||
}
|
||||
|
||||
// Seed 首次启动初始化:管理员账号 + 默认渠道 + 默认模型。
|
||||
func (a *App) Seed() error {
|
||||
// 1. 管理员(从环境变量读取,默认 admin/admin123,生产必须改)
|
||||
var count int64
|
||||
a.DB.Model(&store.User{}).Where("role = ?", store.RoleAdmin).Count(&count)
|
||||
if count == 0 {
|
||||
hash, err := a.Hasher.HashPassword(envOr("OT_ADMIN_PASSWORD", "admin123"))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
admin := store.User{
|
||||
Username: envOr("OT_ADMIN_USERNAME", "admin"),
|
||||
Email: envOr("OT_ADMIN_EMAIL", "admin@localhost"),
|
||||
PasswordHash: hash,
|
||||
Role: store.RoleAdmin,
|
||||
Balance: 1000, // 初始余额,便于联调;生产由充值/调整决定
|
||||
Status: store.UserStatusActive,
|
||||
}
|
||||
if err := a.DB.Create(&admin).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
log.Printf("seed: created admin user %q (change the default password!)", admin.Username)
|
||||
}
|
||||
|
||||
// 2. 默认渠道(配置了上游 key 时创建)
|
||||
if a.Cfg.Proxy.UpstreamKey != "" {
|
||||
var chCount int64
|
||||
a.DB.Model(&store.Channel{}).Count(&chCount)
|
||||
if chCount == 0 {
|
||||
enc, err := a.Enc.Encrypt(a.Cfg.Proxy.UpstreamKey)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
ch := store.Channel{
|
||||
Name: a.Cfg.Proxy.DefaultChannelName,
|
||||
Provider: store.ChannelProviderOpenAI,
|
||||
BaseURL: a.Cfg.Proxy.UpstreamBaseURL,
|
||||
APIKeyEnc: enc,
|
||||
Weight: 1,
|
||||
Priority: 0,
|
||||
TimeoutMS: int(a.Cfg.Proxy.Timeout / time.Millisecond),
|
||||
MaxConcurrency: 16,
|
||||
HealthStatus: store.ChannelHealthHealthy,
|
||||
Enabled: true,
|
||||
}
|
||||
if err := a.DB.Create(&ch).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
// 默认模型 + 绑定
|
||||
m := store.Model{
|
||||
Name: a.Cfg.Proxy.DefaultModel,
|
||||
DisplayName: a.Cfg.Proxy.DefaultModel,
|
||||
InputPrice: 0.15, // 每百万 token,示例价
|
||||
OutputPrice: 0.60,
|
||||
Enabled: true,
|
||||
}
|
||||
if err := a.DB.Create(&m).Error; err == nil {
|
||||
a.DB.Create(&store.ChannelModelBinding{ChannelID: ch.ID, ModelID: m.ID, UpstreamModel: m.Name})
|
||||
}
|
||||
log.Printf("seed: created default channel %q (%s)", ch.Name, ch.BaseURL)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (a *App) Shutdown(ctx context.Context) {
|
||||
a.Usage.Close()
|
||||
if sqlDB, err := a.DB.DB(); err == nil {
|
||||
_ = sqlDB.Close()
|
||||
}
|
||||
_ = ctx
|
||||
}
|
||||
|
||||
func envOr(key, fallback string) string {
|
||||
v := envLookup(key)
|
||||
if v == "" {
|
||||
return fallback
|
||||
}
|
||||
return v
|
||||
}
|
||||
@@ -0,0 +1,5 @@
|
||||
package app
|
||||
|
||||
import "os"
|
||||
|
||||
func envLookup(key string) string { return os.Getenv(key) }
|
||||
@@ -0,0 +1,68 @@
|
||||
// Package channel 渠道仓储:选择、加解密、健康过滤。
|
||||
// M1 阶段实现最小选择逻辑(按优先级+权重取第一个健康启用的渠道),
|
||||
// 负载均衡/健康检查/故障转移在 M4 完善。
|
||||
package channel
|
||||
|
||||
import (
|
||||
"errors"
|
||||
|
||||
"github.com/openteam/server/internal/pkg/crypto"
|
||||
"github.com/openteam/server/internal/store"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
var ErrNoChannel = errors.New("no available channel")
|
||||
|
||||
type Service struct {
|
||||
db *gorm.DB
|
||||
enc *crypto.Encryptor
|
||||
}
|
||||
|
||||
func NewService(db *gorm.DB, enc *crypto.Encryptor) *Service {
|
||||
return &Service{db: db, enc: enc}
|
||||
}
|
||||
|
||||
// Select 选择处理请求的渠道:启用 + 健康,按 priority 升序、weight 降序。
|
||||
func (s *Service) Select() (*store.Channel, error) {
|
||||
var chs []store.Channel
|
||||
if err := s.db.Where("enabled = ? AND health_status = ?", true, store.ChannelHealthHealthy).
|
||||
Order("priority ASC, weight DESC, id ASC").Find(&chs).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(chs) == 0 {
|
||||
return nil, ErrNoChannel
|
||||
}
|
||||
return &chs[0], nil
|
||||
}
|
||||
|
||||
// UpstreamKey 解密渠道上游密钥。
|
||||
func (s *Service) UpstreamKey(ch *store.Channel) (string, error) {
|
||||
return s.enc.Decrypt(ch.APIKeyEnc)
|
||||
}
|
||||
|
||||
// ResolveModel 按全局模型名找到绑定渠道;M1 简化:返回绑定该模型的第一个健康渠道。
|
||||
func (s *Service) ResolveModel(modelName string) (*store.Channel, *store.ChannelModelBinding, error) {
|
||||
var m store.Model
|
||||
if err := s.db.Where("name = ? AND enabled = ?", modelName, true).First(&m).Error; err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
var b store.ChannelModelBinding
|
||||
if err := s.db.Where("model_id = ?", m.ID).
|
||||
Joins("JOIN channels ON channels.id = channel_model_bindings.channel_id AND channels.enabled = ? AND channels.health_status = ?", true, store.ChannelHealthHealthy).
|
||||
Order("channel_model_bindings.weight DESC").
|
||||
First(&b).Error; err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
ch, err := s.Select()
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
// 用绑定里的渠道(如果健康),否则回退默认渠道
|
||||
if b.ChannelID != ch.ID {
|
||||
var bound store.Channel
|
||||
if err := s.db.First(&bound, b.ChannelID).Error; err == nil && bound.Enabled && bound.HealthStatus == store.ChannelHealthHealthy {
|
||||
return &bound, &b, nil
|
||||
}
|
||||
}
|
||||
return ch, &b, nil
|
||||
}
|
||||
@@ -0,0 +1,124 @@
|
||||
// Package config 加载服务配置:.env / 环境变量 / 默认值(viper)。
|
||||
package config
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/spf13/viper"
|
||||
)
|
||||
|
||||
type Config struct {
|
||||
Env string // development | production
|
||||
Port int
|
||||
DB DBConfig
|
||||
JWT JWTConfig
|
||||
Auth AuthConfig
|
||||
Proxy ProxyConfig
|
||||
Master string // 渠道密钥 AES-GCM 主密钥(来自环境变量)
|
||||
}
|
||||
|
||||
type DBConfig struct {
|
||||
Driver string // sqlite | postgres
|
||||
DSN string
|
||||
}
|
||||
|
||||
type JWTConfig struct {
|
||||
Secret string
|
||||
AccessTTL time.Duration
|
||||
RefreshTTL time.Duration
|
||||
Issuer string
|
||||
CookieName string
|
||||
CookieSecure bool
|
||||
CookieDomain string
|
||||
}
|
||||
|
||||
type AuthConfig struct {
|
||||
RegistrationMode string // open | invite
|
||||
Argon2Time uint32
|
||||
Argon2Memory uint32
|
||||
Argon2Threads uint8
|
||||
Argon2KeyLen uint32
|
||||
SaltLen int
|
||||
}
|
||||
|
||||
type ProxyConfig struct {
|
||||
DefaultChannelName string // 首次启动自动创建的渠道名(如 openai)
|
||||
UpstreamBaseURL string // 渠道 base_url 默认值
|
||||
UpstreamKey string // 渠道上游 key 默认值
|
||||
DefaultModel string // 渠道模型导入时使用的模型名
|
||||
Timeout time.Duration
|
||||
}
|
||||
|
||||
func Load() (*Config, error) {
|
||||
v := viper.New()
|
||||
v.SetEnvPrefix("OT")
|
||||
v.SetEnvKeyReplacer(strings.NewReplacer(".", "_"))
|
||||
v.AutomaticEnv()
|
||||
|
||||
// 默认值(与 .env.example 对应)
|
||||
v.SetDefault("env", "development")
|
||||
v.SetDefault("port", 8080)
|
||||
|
||||
v.SetDefault("db.driver", "sqlite")
|
||||
v.SetDefault("db.dsn", "data/openteam.db")
|
||||
|
||||
v.SetDefault("jwt.secret", "dev-only-secret-change-me")
|
||||
v.SetDefault("jwt.access_ttl", "2h")
|
||||
v.SetDefault("jwt.refresh_ttl", "168h")
|
||||
v.SetDefault("jwt.issuer", "openteam")
|
||||
v.SetDefault("jwt.cookie_name", "ot_refresh")
|
||||
v.SetDefault("jwt.cookie_secure", false)
|
||||
v.SetDefault("jwt.cookie_domain", "")
|
||||
|
||||
v.SetDefault("auth.registration_mode", "open")
|
||||
v.SetDefault("auth.argon2_time", 3)
|
||||
v.SetDefault("auth.argon2_memory", 64*1024) // 64 MiB
|
||||
v.SetDefault("auth.argon2_threads", 2)
|
||||
v.SetDefault("auth.argon2_keylen", 32)
|
||||
v.SetDefault("auth.salt_len", 16)
|
||||
|
||||
v.SetDefault("proxy.default_channel_name", "openai")
|
||||
v.SetDefault("proxy.upstream_base_url", "https://api.openai.com")
|
||||
v.SetDefault("proxy.upstream_key", "")
|
||||
v.SetDefault("proxy.default_model", "gpt-4o-mini")
|
||||
v.SetDefault("proxy.timeout", "120s")
|
||||
|
||||
// 支持读取 .env 文件(可选,不强制)
|
||||
v.SetConfigFile(".env")
|
||||
_ = v.ReadInConfig()
|
||||
|
||||
return &Config{
|
||||
Env: v.GetString("env"),
|
||||
Port: v.GetInt("port"),
|
||||
DB: DBConfig{
|
||||
Driver: v.GetString("db.driver"),
|
||||
DSN: v.GetString("db.dsn"),
|
||||
},
|
||||
JWT: JWTConfig{
|
||||
Secret: v.GetString("jwt.secret"),
|
||||
AccessTTL: v.GetDuration("jwt.access_ttl"),
|
||||
RefreshTTL: v.GetDuration("jwt.refresh_ttl"),
|
||||
Issuer: v.GetString("jwt.issuer"),
|
||||
CookieName: v.GetString("jwt.cookie_name"),
|
||||
CookieSecure: v.GetBool("jwt.cookie_secure"),
|
||||
CookieDomain: v.GetString("jwt.cookie_domain"),
|
||||
},
|
||||
Auth: AuthConfig{
|
||||
RegistrationMode: v.GetString("auth.registration_mode"),
|
||||
Argon2Time: v.GetUint32("auth.argon2_time"),
|
||||
Argon2Memory: v.GetUint32("auth.argon2_memory"),
|
||||
Argon2Threads: v.GetUint8("auth.argon2_threads"),
|
||||
Argon2KeyLen: v.GetUint32("auth.argon2_keylen"),
|
||||
SaltLen: v.GetInt("auth.salt_len"),
|
||||
},
|
||||
Proxy: ProxyConfig{
|
||||
DefaultChannelName: v.GetString("proxy.default_channel_name"),
|
||||
UpstreamBaseURL: v.GetString("proxy.upstream_base_url"),
|
||||
UpstreamKey: v.GetString("proxy.upstream_key"),
|
||||
DefaultModel: v.GetString("proxy.default_model"),
|
||||
Timeout: v.GetDuration("proxy.timeout"),
|
||||
},
|
||||
Master: v.GetString("master_key"),
|
||||
}, nil
|
||||
}
|
||||
@@ -0,0 +1,52 @@
|
||||
// Package apikey 生成与管理 API Key:sk- + 48 位 base62 随机串。
|
||||
// 库中仅存 SHA-256 哈希与展示前缀(PLANNING §4.3.3)。
|
||||
package apikey
|
||||
|
||||
import (
|
||||
"crypto/rand"
|
||||
"crypto/sha256"
|
||||
"encoding/base64"
|
||||
"encoding/hex"
|
||||
"strings"
|
||||
)
|
||||
|
||||
const (
|
||||
alphabet = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
|
||||
keyLen = 48
|
||||
prefix = "sk-"
|
||||
)
|
||||
|
||||
// Generate 生成明文 key(仅创建时展示一次)与哈希、前缀。
|
||||
func Generate() (plain, hash, keyPrefix string, err error) {
|
||||
buf := make([]byte, keyLen)
|
||||
if _, err = rand.Read(buf); err != nil {
|
||||
return "", "", "", err
|
||||
}
|
||||
for i := range buf {
|
||||
buf[i] = alphabet[int(buf[i])%len(alphabet)]
|
||||
}
|
||||
plain = prefix + string(buf)
|
||||
return plain, Hash(plain), Prefix(plain), nil
|
||||
}
|
||||
|
||||
// Hash 返回 key 的 SHA-256 十六进制。
|
||||
func Hash(key string) string {
|
||||
sum := sha256.Sum256([]byte(key))
|
||||
return hex.EncodeToString(sum[:])
|
||||
}
|
||||
|
||||
// Prefix 展示前缀:sk-aB3cD5…(前 12 字符)
|
||||
func Prefix(key string) string {
|
||||
if len(key) <= 12 {
|
||||
return key
|
||||
}
|
||||
return key[:12]
|
||||
}
|
||||
|
||||
// Valid 校验明文格式。
|
||||
func Valid(key string) bool {
|
||||
return strings.HasPrefix(key, prefix) && len(key) == len(prefix)+keyLen
|
||||
}
|
||||
|
||||
// base64 占位,避免未使用导入告警
|
||||
var _ = base64.StdEncoding
|
||||
@@ -0,0 +1,33 @@
|
||||
package apikey
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestGenerate(t *testing.T) {
|
||||
plain, hash, prefix, err := Generate()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !Valid(plain) {
|
||||
t.Fatalf("generated key invalid: %q", plain)
|
||||
}
|
||||
if len(plain) != 3+48 {
|
||||
t.Fatalf("key length = %d, want 51", len(plain))
|
||||
}
|
||||
if Hash(plain) != hash {
|
||||
t.Fatal("hash mismatch")
|
||||
}
|
||||
if len(prefix) > len(plain) || prefix != plain[:len(prefix)] {
|
||||
t.Fatal("prefix must be prefix of plain key")
|
||||
}
|
||||
// 两次生成不重复
|
||||
plain2, _, _, _ := Generate()
|
||||
if plain == plain2 {
|
||||
t.Fatal("keys should be unique")
|
||||
}
|
||||
}
|
||||
|
||||
func TestValid(t *testing.T) {
|
||||
if Valid("") || Valid("sk-short") || Valid("xxx") {
|
||||
t.Fatal("invalid keys should be rejected")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,125 @@
|
||||
// Package crypto 密码哈希(argon2id)与对称加密(AES-GCM)。
|
||||
package crypto
|
||||
|
||||
import (
|
||||
"crypto/aes"
|
||||
"crypto/cipher"
|
||||
"crypto/rand"
|
||||
"crypto/subtle"
|
||||
"encoding/base64"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"golang.org/x/crypto/argon2"
|
||||
)
|
||||
|
||||
type PasswordHasher struct {
|
||||
Time uint32
|
||||
Memory uint32
|
||||
Threads uint8
|
||||
KeyLen uint32
|
||||
SaltLen int
|
||||
}
|
||||
|
||||
func NewPasswordHasher(time, memory uint32, threads uint8, keyLen uint32, saltLen int) *PasswordHasher {
|
||||
return &PasswordHasher{Time: time, Memory: memory, Threads: threads, KeyLen: keyLen, SaltLen: saltLen}
|
||||
}
|
||||
|
||||
// HashPassword argon2id 编码为 $argon2id$v=19$m=...,t=...,p=...$salt$hash
|
||||
func (h *PasswordHasher) HashPassword(password string) (string, error) {
|
||||
salt := make([]byte, h.SaltLen)
|
||||
if _, err := rand.Read(salt); err != nil {
|
||||
return "", err
|
||||
}
|
||||
key := argon2.IDKey([]byte(password), salt, h.Time, h.Memory, h.Threads, h.KeyLen)
|
||||
enc := base64.RawStdEncoding
|
||||
return fmt.Sprintf("$argon2id$v=19$m=%d,t=%d,p=%d$%s$%s",
|
||||
h.Memory, h.Time, h.Threads, enc.EncodeToString(salt), enc.EncodeToString(key)), nil
|
||||
}
|
||||
|
||||
// VerifyPassword 校验密码,返回是否匹配(常数时间比较)。
|
||||
func (h *PasswordHasher) VerifyPassword(encoded, password string) (bool, error) {
|
||||
parts := strings.Split(encoded, "$")
|
||||
if len(parts) != 6 || parts[1] != "argon2id" {
|
||||
return false, errors.New("invalid hash format")
|
||||
}
|
||||
var memory uint32
|
||||
var time uint32
|
||||
var threads uint8
|
||||
if _, err := fmt.Sscanf(parts[3], "m=%d,t=%d,p=%d", &memory, &time, &threads); err != nil {
|
||||
return false, err
|
||||
}
|
||||
enc := base64.RawStdEncoding
|
||||
salt, err := enc.DecodeString(parts[4])
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
want, err := enc.DecodeString(parts[5])
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
got := argon2.IDKey([]byte(password), salt, time, memory, threads, uint32(len(want)))
|
||||
return subtle.ConstantTimeCompare(got, want) == 1, nil
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// AES-GCM 渠道密钥加密
|
||||
|
||||
type Encryptor struct {
|
||||
key []byte
|
||||
}
|
||||
|
||||
// NewEncryptor 主密钥必须为 16/24/32 字节;不足时用 SHA-256 派生固定 32 字节。
|
||||
func NewEncryptor(master string) *Encryptor {
|
||||
key := []byte(master)
|
||||
switch len(key) {
|
||||
case 16, 24, 32:
|
||||
default:
|
||||
sum := sha256Sum(master)
|
||||
key = sum
|
||||
}
|
||||
return &Encryptor{key: key}
|
||||
}
|
||||
|
||||
// Encrypt 输出 base64(nonce || ciphertext)
|
||||
func (e *Encryptor) Encrypt(plain string) (string, error) {
|
||||
block, err := aes.NewCipher(e.key)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
gcm, err := cipher.NewGCM(block)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
nonce := make([]byte, gcm.NonceSize())
|
||||
if _, err := rand.Read(nonce); err != nil {
|
||||
return "", err
|
||||
}
|
||||
ct := gcm.Seal(nil, nonce, []byte(plain), nil)
|
||||
return base64.StdEncoding.EncodeToString(append(nonce, ct...)), nil
|
||||
}
|
||||
|
||||
func (e *Encryptor) Decrypt(enc string) (string, error) {
|
||||
raw, err := base64.StdEncoding.DecodeString(enc)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
block, err := aes.NewCipher(e.key)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
gcm, err := cipher.NewGCM(block)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if len(raw) < gcm.NonceSize() {
|
||||
return "", errors.New("ciphertext too short")
|
||||
}
|
||||
nonce, ct := raw[:gcm.NonceSize()], raw[gcm.NonceSize():]
|
||||
plain, err := gcm.Open(nil, nonce, ct, nil)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return string(plain), nil
|
||||
}
|
||||
@@ -0,0 +1,44 @@
|
||||
package crypto
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestPasswordHashRoundTrip(t *testing.T) {
|
||||
h := NewPasswordHasher(3, 64*1024, 2, 32, 16)
|
||||
hash, err := h.HashPassword("s3cret-password")
|
||||
if err != nil {
|
||||
t.Fatalf("hash: %v", err)
|
||||
}
|
||||
ok, err := h.VerifyPassword(hash, "s3cret-password")
|
||||
if err != nil || !ok {
|
||||
t.Fatalf("verify correct password: ok=%v err=%v", ok, err)
|
||||
}
|
||||
ok, _ = h.VerifyPassword(hash, "wrong-password")
|
||||
if ok {
|
||||
t.Fatal("wrong password should not verify")
|
||||
}
|
||||
}
|
||||
|
||||
func TestEncryptDecrypt(t *testing.T) {
|
||||
e := NewEncryptor("master-key-0123456789abcdef")
|
||||
enc, err := e.Encrypt("sk-upstream-secret")
|
||||
if err != nil {
|
||||
t.Fatalf("encrypt: %v", err)
|
||||
}
|
||||
dec, err := e.Decrypt(enc)
|
||||
if err != nil || dec != "sk-upstream-secret" {
|
||||
t.Fatalf("decrypt: got %q err %v", dec, err)
|
||||
}
|
||||
// 密文不可读
|
||||
if dec == enc {
|
||||
t.Fatal("ciphertext should differ from plaintext")
|
||||
}
|
||||
}
|
||||
|
||||
func TestShortMasterKeyDerived(t *testing.T) {
|
||||
e := NewEncryptor("short")
|
||||
enc, _ := e.Encrypt("x")
|
||||
dec, err := e.Decrypt(enc)
|
||||
if err != nil || dec != "x" {
|
||||
t.Fatalf("short key derive failed: %v", err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,8 @@
|
||||
package crypto
|
||||
|
||||
import "crypto/sha256"
|
||||
|
||||
func sha256Sum(s string) []byte {
|
||||
sum := sha256.Sum256([]byte(s))
|
||||
return sum[:]
|
||||
}
|
||||
@@ -0,0 +1,71 @@
|
||||
// Package jwt 访问令牌(短时,存内存)与刷新令牌(HttpOnly Cookie)签发/校验。
|
||||
package jwt
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"time"
|
||||
|
||||
"github.com/golang-jwt/jwt/v5"
|
||||
)
|
||||
|
||||
type Claims struct {
|
||||
UserID uint64 `json:"uid"`
|
||||
Username string `json:"uname"`
|
||||
Role string `json:"role"`
|
||||
jwt.RegisteredClaims
|
||||
}
|
||||
|
||||
type Manager struct {
|
||||
secret []byte
|
||||
issuer string
|
||||
accessTTL time.Duration
|
||||
refreshTTL time.Duration
|
||||
}
|
||||
|
||||
func NewManager(secret, issuer string, accessTTL, refreshTTL time.Duration) *Manager {
|
||||
return &Manager{secret: []byte(secret), issuer: issuer, accessTTL: accessTTL, refreshTTL: refreshTTL}
|
||||
}
|
||||
|
||||
func (m *Manager) AccessTTL() time.Duration { return m.accessTTL }
|
||||
func (m *Manager) RefreshTTL() time.Duration { return m.refreshTTL }
|
||||
|
||||
// Sign 签发 token;typ 取 "access" / "refresh"。
|
||||
func (m *Manager) Sign(userID uint64, username, role, typ string) (string, time.Time, error) {
|
||||
ttl := m.accessTTL
|
||||
if typ == "refresh" {
|
||||
ttl = m.refreshTTL
|
||||
}
|
||||
now := time.Now()
|
||||
exp := now.Add(ttl)
|
||||
claims := Claims{
|
||||
UserID: userID,
|
||||
Username: username,
|
||||
Role: role,
|
||||
RegisteredClaims: jwt.RegisteredClaims{
|
||||
Issuer: m.issuer,
|
||||
Subject: typ,
|
||||
IssuedAt: jwt.NewNumericDate(now),
|
||||
ExpiresAt: jwt.NewNumericDate(exp),
|
||||
},
|
||||
}
|
||||
tok := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
|
||||
s, err := tok.SignedString(m.secret)
|
||||
return s, exp, err
|
||||
}
|
||||
|
||||
var ErrInvalidToken = errors.New("invalid token")
|
||||
|
||||
// Parse 校验签名与有效期。
|
||||
func (m *Manager) Parse(token string) (*Claims, error) {
|
||||
claims := &Claims{}
|
||||
tok, err := jwt.ParseWithClaims(token, claims, func(t *jwt.Token) (any, error) {
|
||||
if _, ok := t.Method.(*jwt.SigningMethodHMAC); !ok {
|
||||
return nil, ErrInvalidToken
|
||||
}
|
||||
return m.secret, nil
|
||||
})
|
||||
if err != nil || !tok.Valid {
|
||||
return nil, ErrInvalidToken
|
||||
}
|
||||
return claims, nil
|
||||
}
|
||||
@@ -0,0 +1,41 @@
|
||||
package jwt
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestSignParse(t *testing.T) {
|
||||
m := NewManager("test-secret", "openteam", time.Hour, 24*time.Hour)
|
||||
tok, exp, err := m.Sign(42, "alice", "admin", "access")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if time.Until(exp) < 50*time.Minute {
|
||||
t.Fatal("expiry too short")
|
||||
}
|
||||
claims, err := m.Parse(tok)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if claims.UserID != 42 || claims.Username != "alice" || claims.Role != "admin" || claims.Subject != "access" {
|
||||
t.Fatalf("claims mismatch: %+v", claims)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExpiredToken(t *testing.T) {
|
||||
m := NewManager("test-secret", "openteam", -time.Minute, time.Hour)
|
||||
tok, _, _ := m.Sign(1, "a", "user", "access")
|
||||
if _, err := m.Parse(tok); err == nil {
|
||||
t.Fatal("expired token should fail")
|
||||
}
|
||||
}
|
||||
|
||||
func TestWrongSecret(t *testing.T) {
|
||||
m1 := NewManager("secret-a", "openteam", time.Hour, time.Hour)
|
||||
m2 := NewManager("secret-b", "openteam", time.Hour, time.Hour)
|
||||
tok, _, _ := m1.Sign(1, "a", "user", "access")
|
||||
if _, err := m2.Parse(tok); err == nil {
|
||||
t.Fatal("token signed with different secret should fail")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,39 @@
|
||||
// Package resp 统一 JSON 响应与错误格式。
|
||||
package resp
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
type Body struct {
|
||||
Data any `json:"data,omitempty"`
|
||||
Error *Error `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
type Error struct {
|
||||
Message string `json:"message"`
|
||||
Type string `json:"type,omitempty"`
|
||||
Code string `json:"code,omitempty"`
|
||||
}
|
||||
|
||||
// OK 200
|
||||
func OK(c *gin.Context, data any) {
|
||||
c.JSON(http.StatusOK, Body{Data: data})
|
||||
}
|
||||
|
||||
// Created 201
|
||||
func Created(c *gin.Context, data any) {
|
||||
c.JSON(http.StatusCreated, Body{Data: data})
|
||||
}
|
||||
|
||||
// Fail 业务错误(message 会展示给用户)
|
||||
func Fail(c *gin.Context, status int, message string) {
|
||||
c.JSON(status, Body{Error: &Error{Message: message}})
|
||||
}
|
||||
|
||||
// FailCode 带错误码的业务错误
|
||||
func FailCode(c *gin.Context, status int, code, message string) {
|
||||
c.JSON(status, Body{Error: &Error{Message: message, Type: code}})
|
||||
}
|
||||
@@ -0,0 +1,133 @@
|
||||
// Package proxy API 网关核心:代理 /v1/chat/completions、/v1/responses、/v1/models。
|
||||
// M1:对 OpenAI 渠道直通(passthrough),不转格式;M3 起加入协议转换。
|
||||
package proxy
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/openteam/server/internal/channel"
|
||||
"github.com/openteam/server/internal/pkg/apikey"
|
||||
"github.com/openteam/server/internal/pkg/crypto"
|
||||
"github.com/openteam/server/internal/store"
|
||||
"github.com/openteam/server/internal/usage"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
const (
|
||||
CtxUserID = "proxy_user_id"
|
||||
CtxKeyID = "proxy_key_id"
|
||||
CtxTrace = "proxy_trace_id"
|
||||
)
|
||||
|
||||
type Gateway struct {
|
||||
db *gorm.DB
|
||||
ch *channel.Service
|
||||
rec *usage.Recorder
|
||||
enc *crypto.Encryptor
|
||||
hc *http.Client
|
||||
}
|
||||
|
||||
func NewGateway(db *gorm.DB, enc *crypto.Encryptor, rec *usage.Recorder) *Gateway {
|
||||
return &Gateway{
|
||||
db: db,
|
||||
ch: channel.NewService(db, enc),
|
||||
rec: rec,
|
||||
enc: enc,
|
||||
hc: &http.Client{Timeout: 120 * time.Second},
|
||||
}
|
||||
}
|
||||
|
||||
// Auth 代理鉴权中间件:Bearer sk-xxx → 哈希查表 → 校验状态/过期/模型白名单。
|
||||
func (g *Gateway) Auth(c *gin.Context) {
|
||||
auth := c.GetHeader("Authorization")
|
||||
key := strings.TrimPrefix(auth, "Bearer ")
|
||||
key = strings.TrimSpace(key)
|
||||
if !apikey.Valid(key) {
|
||||
openAIError(c, http.StatusUnauthorized, "invalid_api_key", "Invalid API key format. Expected: Bearer sk-...")
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
hash := apikey.Hash(key)
|
||||
var k store.APIKey
|
||||
if err := g.db.Where("key_hash = ? AND status = ?", hash, store.KeyStatusActive).First(&k).Error; err != nil {
|
||||
openAIError(c, http.StatusUnauthorized, "invalid_api_key", "Invalid API key")
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
var u store.User
|
||||
if err := g.db.First(&u, k.UserID).Error; err != nil || u.Status != store.UserStatusActive {
|
||||
openAIError(c, http.StatusForbidden, "user_disabled", "User account is disabled")
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
if k.ExpiresAt != nil && time.Now().After(*k.ExpiresAt) {
|
||||
openAIError(c, http.StatusUnauthorized, "key_expired", "API key has expired")
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
|
||||
c.Set(CtxUserID, u.ID)
|
||||
c.Set(CtxKeyID, k.ID)
|
||||
c.Set(CtxTrace, newTraceID())
|
||||
g.db.Model(&store.APIKey{}).Where("id = ?", k.ID).Update("last_used_at", time.Now())
|
||||
c.Next()
|
||||
}
|
||||
|
||||
// Handle 路由到对应协议处理器。
|
||||
func (g *Gateway) Handle(c *gin.Context) {
|
||||
switch {
|
||||
case c.Request.URL.Path == "/v1/chat/completions":
|
||||
g.chatCompletions(c)
|
||||
case c.Request.URL.Path == "/v1/responses":
|
||||
g.responses(c)
|
||||
case c.Request.URL.Path == "/v1/models" && c.Request.Method == http.MethodGet:
|
||||
g.models(c)
|
||||
default:
|
||||
openAIError(c, http.StatusNotFound, "not_found", "Unknown endpoint: "+c.Request.URL.Path)
|
||||
}
|
||||
}
|
||||
|
||||
// models GET /v1/models:返回启用的全局模型(OpenAI 风格)。
|
||||
func (g *Gateway) models(c *gin.Context) {
|
||||
var ms []store.Model
|
||||
if err := g.db.Where("enabled = ?", true).Order("sort ASC, id ASC").Find(&ms).Error; err != nil {
|
||||
openAIError(c, http.StatusInternalServerError, "internal_error", "failed to load models")
|
||||
return
|
||||
}
|
||||
data := make([]gin.H, 0, len(ms))
|
||||
for _, m := range ms {
|
||||
data = append(data, gin.H{
|
||||
"id": m.Name,
|
||||
"object": "model",
|
||||
"created": m.CreatedAt.Unix(),
|
||||
"owned_by": "openteam",
|
||||
})
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"object": "list", "data": data})
|
||||
}
|
||||
|
||||
// resolveUser 取当前用户(含余额)。
|
||||
func (g *Gateway) resolveUser(c *gin.Context) (*store.User, bool) {
|
||||
uid, _ := c.Get(CtxUserID)
|
||||
var u store.User
|
||||
if err := g.db.First(&u, uid).Error; err != nil {
|
||||
openAIError(c, http.StatusUnauthorized, "invalid_api_key", "Invalid API key")
|
||||
return nil, false
|
||||
}
|
||||
return &u, true
|
||||
}
|
||||
|
||||
// checkBalance 余额不足返回 402(PLANNING §4.4.3)。
|
||||
func (g *Gateway) checkBalance(c *gin.Context, u *store.User) bool {
|
||||
if u.Balance <= 0 {
|
||||
openAIError(c, http.StatusPaymentRequired, "insufficient_balance", "Insufficient balance. Please recharge or contact admin.")
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
var errNoChannel = errors.New("no available channel")
|
||||
@@ -0,0 +1,115 @@
|
||||
package proxy
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/openteam/server/internal/store"
|
||||
)
|
||||
|
||||
// chatCompletions POST /v1/chat/completions
|
||||
func (g *Gateway) chatCompletions(c *gin.Context) {
|
||||
u, ok := g.resolveUser(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
if !g.checkBalance(c, u) {
|
||||
return
|
||||
}
|
||||
|
||||
br, body, err := parseBody(c)
|
||||
if err != nil {
|
||||
openAIError(c, http.StatusBadRequest, "invalid_request_error", "failed to read request body")
|
||||
return
|
||||
}
|
||||
c.Set("protocol", "chat")
|
||||
c.Set("model_name", br.Model)
|
||||
|
||||
ch, err := g.ch.Select()
|
||||
if err != nil {
|
||||
openAIError(c, http.StatusServiceUnavailable, "no_channel", "No available upstream channel")
|
||||
g.recordError(c, nil, nil, now(), "no_channel")
|
||||
return
|
||||
}
|
||||
|
||||
sink := &usageSink{}
|
||||
c.Set("usage_raw", &sinkHolder{sink: sink})
|
||||
g.doPassthrough(c, ch, "/v1/chat/completions", body, br.Stream, func(raw json.RawMessage) {
|
||||
sink.push(raw)
|
||||
})
|
||||
}
|
||||
|
||||
// responses POST /v1/responses(OpenAI Responses API)
|
||||
func (g *Gateway) responses(c *gin.Context) {
|
||||
u, ok := g.resolveUser(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
if !g.checkBalance(c, u) {
|
||||
return
|
||||
}
|
||||
|
||||
br, body, err := parseBody(c)
|
||||
if err != nil {
|
||||
openAIError(c, http.StatusBadRequest, "invalid_request_error", "failed to read request body")
|
||||
return
|
||||
}
|
||||
c.Set("protocol", "responses")
|
||||
c.Set("model_name", br.Model)
|
||||
|
||||
ch, err := g.ch.Select()
|
||||
if err != nil {
|
||||
openAIError(c, http.StatusServiceUnavailable, "no_channel", "No available upstream channel")
|
||||
g.recordError(c, nil, nil, now(), "no_channel")
|
||||
return
|
||||
}
|
||||
// M1 仅支持 OpenAI 原生渠道直通;Anthropic 渠道的转换在 M3
|
||||
if ch.Provider != store.ChannelProviderOpenAI {
|
||||
openAIError(c, http.StatusNotImplemented, "conversion_pending",
|
||||
"Responses protocol on this channel requires format conversion (planned in M3)")
|
||||
return
|
||||
}
|
||||
|
||||
sink := &usageSink{}
|
||||
c.Set("usage_raw", &sinkHolder{sink: sink})
|
||||
g.doPassthrough(c, ch, "/v1/responses", body, br.Stream, func(raw json.RawMessage) {
|
||||
sink.push(raw)
|
||||
})
|
||||
}
|
||||
|
||||
// usageSinkHolder 桥接:gin context 里保存 sink 引用,供 finishUsage 读取最终 usage。
|
||||
type sinkHolder struct {
|
||||
sink *usageSink
|
||||
}
|
||||
|
||||
// openAIError 按 OpenAI 错误格式返回(PLANNING §4.1.4)。
|
||||
func openAIError(c *gin.Context, status int, code, message string) {
|
||||
c.AbortWithStatusJSON(status, gin.H{
|
||||
"error": gin.H{
|
||||
"message": message,
|
||||
"type": errorTypeFor(status),
|
||||
"param": nil,
|
||||
"code": code,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
func errorTypeFor(status int) string {
|
||||
switch status {
|
||||
case http.StatusUnauthorized:
|
||||
return "authentication_error"
|
||||
case http.StatusForbidden:
|
||||
return "permission_error"
|
||||
case http.StatusNotFound:
|
||||
return "invalid_request_error"
|
||||
case http.StatusBadRequest:
|
||||
return "invalid_request_error"
|
||||
case http.StatusPaymentRequired:
|
||||
return "insufficient_quota"
|
||||
case http.StatusTooManyRequests:
|
||||
return "rate_limit_error"
|
||||
default:
|
||||
return "api_error"
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,382 @@
|
||||
package proxy
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/openteam/server/internal/store"
|
||||
)
|
||||
|
||||
// newTraceID 生成请求 trace(用于日志与记账幂等 ref)。
|
||||
func newTraceID() string {
|
||||
b := make([]byte, 8)
|
||||
_, _ = rand.Read(b)
|
||||
return hex.EncodeToString(b)
|
||||
}
|
||||
|
||||
// bodyReq 统一取出请求体并解析 model / stream 字段。
|
||||
type bodyReq struct {
|
||||
Model string `json:"model"`
|
||||
Stream bool `json:"stream"`
|
||||
}
|
||||
|
||||
// parseBody 读取并回填请求体,解析 model/stream。
|
||||
func parseBody(c *gin.Context) (*bodyReq, []byte, error) {
|
||||
body, err := io.ReadAll(c.Request.Body)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
c.Request.Body = io.NopCloser(bytes.NewReader(body))
|
||||
br := &bodyReq{}
|
||||
_ = json.Unmarshal(body, br) // 解析失败按空处理,直通仍可转发
|
||||
return br, body, nil
|
||||
}
|
||||
|
||||
// upstreamURL 组装上游地址:base_url + 客户端路径(/v1/chat/completions 等)。
|
||||
func upstreamURL(ch *store.Channel, path string) string {
|
||||
base := strings.TrimRight(ch.BaseURL, "/")
|
||||
return base + path
|
||||
}
|
||||
|
||||
// doPassthrough 通用直通:替换 Authorization 为渠道密钥,转发请求。
|
||||
// convert 回调用于改写请求体(M1 直通为原样;M3 转换时改写)。
|
||||
func (g *Gateway) doPassthrough(c *gin.Context, ch *store.Channel, path string, body []byte, stream bool, outUsage func(usageRaw json.RawMessage)) {
|
||||
upKey, err := g.ch.UpstreamKey(ch)
|
||||
if err != nil {
|
||||
openAIError(c, http.StatusInternalServerError, "channel_error", "failed to decrypt channel key")
|
||||
return
|
||||
}
|
||||
|
||||
upBody := body
|
||||
// 流式 chat:注入 stream_options.include_usage,保证末块带 usage(OpenAI 行为)
|
||||
if stream && path == "/v1/chat/completions" && !bytes.Contains(upBody, []byte(`"include_usage"`)) {
|
||||
var m map[string]any
|
||||
if json.Unmarshal(upBody, &m) == nil {
|
||||
m["stream_options"] = map[string]any{"include_usage": true}
|
||||
if b, err := json.Marshal(m); err == nil {
|
||||
upBody = b
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(c.Request.Context(), time.Duration(ch.TimeoutMS)*time.Millisecond)
|
||||
defer cancel()
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, upstreamURL(ch, path), bytes.NewReader(upBody))
|
||||
if err != nil {
|
||||
openAIError(c, http.StatusInternalServerError, "internal_error", "failed to build upstream request")
|
||||
return
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Authorization", "Bearer "+upKey)
|
||||
req.Header.Set("Accept", c.GetHeader("Accept"))
|
||||
if ua := c.GetHeader("User-Agent"); ua != "" {
|
||||
req.Header.Set("User-Agent", ua)
|
||||
}
|
||||
// 透传 OpenAI 生态请求头(组织/项目等)
|
||||
for _, h := range []string{"OpenAI-Organization", "OpenAI-Project", "OpenAI-Beta"} {
|
||||
if v := c.GetHeader(h); v != "" {
|
||||
req.Header.Set(h, v)
|
||||
}
|
||||
}
|
||||
|
||||
start := time.Now()
|
||||
resp, err := g.hc.Do(req)
|
||||
if err != nil {
|
||||
status := http.StatusBadGateway
|
||||
msg := "Upstream request failed: " + err.Error()
|
||||
if ctx.Err() == context.DeadlineExceeded {
|
||||
status = http.StatusGatewayTimeout
|
||||
msg = "Upstream request timed out"
|
||||
}
|
||||
openAIError(c, status, "upstream_error", msg)
|
||||
g.recordError(c, ch, nil, start, "upstream_error")
|
||||
return
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
// 非 2xx:透传上游错误体(OpenAI 格式),并记录 error 用量
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
errBody, _ := io.ReadAll(resp.Body)
|
||||
status := resp.StatusCode
|
||||
// 上游 5xx → 网关 502/504(重试逻辑 M4)
|
||||
if status >= 500 {
|
||||
status = http.StatusBadGateway
|
||||
}
|
||||
c.DataFromReader(status, int64(len(errBody)), "application/json", bytes.NewReader(errBody), nil)
|
||||
c.Header("Content-Type", "application/json")
|
||||
g.recordError(c, ch, resp, start, "upstream_http_"+strconv.Itoa(resp.StatusCode))
|
||||
return
|
||||
}
|
||||
|
||||
// 成功响应
|
||||
c.Header("Content-Type", resp.Header.Get("Content-Type"))
|
||||
c.Status(http.StatusOK)
|
||||
if stream {
|
||||
g.streamCopy(c, ch, resp.Body, start, outUsage)
|
||||
} else {
|
||||
g.copyAndCapture(c, ch, resp.Body, start, outUsage)
|
||||
}
|
||||
}
|
||||
|
||||
// copyAndCapture 非流式:整体转发 + 解析 usage + 记账。
|
||||
func (g *Gateway) copyAndCapture(c *gin.Context, ch *store.Channel, r io.Reader, start time.Time, outUsage func(json.RawMessage)) {
|
||||
data, err := io.ReadAll(r)
|
||||
if err != nil {
|
||||
openAIError(c, http.StatusBadGateway, "upstream_error", "failed reading upstream response")
|
||||
g.recordError(c, ch, nil, start, "read_error")
|
||||
return
|
||||
}
|
||||
// 尝试解析 usage(chat / responses 字段不同)
|
||||
if usageRaw := extractUsage(data); usageRaw != nil {
|
||||
outUsage(usageRaw)
|
||||
}
|
||||
_, _ = c.Writer.Write(data)
|
||||
g.finishUsage(c, ch, start, store.UsageStatusSuccess, "")
|
||||
}
|
||||
|
||||
// streamCopy 流式:边读上游 SSE 边写客户端,零缓冲转发;扫描 usage 行记账。
|
||||
// 客户端断连(ctx cancel)即中止上游读取。
|
||||
func (g *Gateway) streamCopy(c *gin.Context, ch *store.Channel, r io.Reader, start time.Time, outUsage func(json.RawMessage)) {
|
||||
w := c.Writer
|
||||
flusher, ok := w.(http.Flusher)
|
||||
if !ok {
|
||||
flusher = nopFlusher{}
|
||||
}
|
||||
|
||||
scanner := newSSEScanner(r)
|
||||
for {
|
||||
line, err := scanner.Next()
|
||||
if line != nil {
|
||||
if _, werr := w.Write(line); werr != nil {
|
||||
// 客户端断开:取消上游(ctx cancel 由 request ctx 处理)
|
||||
g.recordError(c, ch, nil, start, "client_disconnect")
|
||||
return
|
||||
}
|
||||
flusher.Flush()
|
||||
if usageRaw := scanUsage(line); usageRaw != nil {
|
||||
outUsage(usageRaw)
|
||||
}
|
||||
}
|
||||
if err != nil {
|
||||
if err == io.EOF {
|
||||
g.finishUsage(c, ch, start, store.UsageStatusSuccess, "")
|
||||
} else {
|
||||
g.recordError(c, ch, nil, start, "stream_read_error")
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
type nopFlusher struct{}
|
||||
|
||||
func (nopFlusher) Flush() {}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// usage 提取
|
||||
|
||||
// usageShape 兼容 chat (prompt/completion) 与 responses (input/output) 两种命名。
|
||||
type usageShape struct {
|
||||
PromptTokens int64 `json:"prompt_tokens"`
|
||||
CompletionTokens int64 `json:"completion_tokens"`
|
||||
InputTokens int64 `json:"input_tokens"`
|
||||
OutputTokens int64 `json:"output_tokens"`
|
||||
TotalTokens int64 `json:"total_tokens"`
|
||||
// Claude 缓存口径(M3 接入)
|
||||
CacheReadInputTokens int64 `json:"cache_read_input_tokens"`
|
||||
CacheCreationInputTokens int64 `json:"cache_creation_input_tokens"`
|
||||
}
|
||||
|
||||
// extractUsage 从完整响应体提取 usage 子对象。
|
||||
func extractUsage(data []byte) json.RawMessage {
|
||||
var m map[string]json.RawMessage
|
||||
if json.Unmarshal(data, &m) != nil {
|
||||
return nil
|
||||
}
|
||||
if u, ok := m["usage"]; ok && string(u) != "null" {
|
||||
return u
|
||||
}
|
||||
// responses 事件/响应:usage 嵌套在 response 对象内
|
||||
if respRaw, ok := m["response"]; ok {
|
||||
var resp map[string]json.RawMessage
|
||||
if json.Unmarshal(respRaw, &resp) == nil {
|
||||
if u, ok := resp["usage"]; ok && string(u) != "null" {
|
||||
return u
|
||||
}
|
||||
}
|
||||
}
|
||||
// chat 兜底:choices[].message.usage
|
||||
if choices, ok := m["choices"]; ok {
|
||||
var cs []map[string]json.RawMessage
|
||||
if json.Unmarshal(choices, &cs) == nil {
|
||||
for _, ch := range cs {
|
||||
if msgRaw, ok := ch["message"]; ok {
|
||||
var msg map[string]json.RawMessage
|
||||
if json.Unmarshal(msgRaw, &msg) == nil {
|
||||
if u, ok := msg["usage"]; ok && string(u) != "null" {
|
||||
return u
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// scanUsage 从 SSE 一行中提取 usage(OpenAI 末块 / responses completed 事件)。
|
||||
func scanUsage(line []byte) json.RawMessage {
|
||||
s := string(line)
|
||||
if !strings.Contains(s, `"usage"`) {
|
||||
return nil
|
||||
}
|
||||
if strings.HasPrefix(s, "data: ") {
|
||||
s = strings.TrimPrefix(s, "data: ")
|
||||
}
|
||||
s = strings.TrimSpace(s)
|
||||
if s == "[DONE]" || s == "" {
|
||||
return nil
|
||||
}
|
||||
var m map[string]json.RawMessage
|
||||
if json.Unmarshal([]byte(s), &m) != nil {
|
||||
return nil
|
||||
}
|
||||
if u, ok := m["usage"]; ok && string(u) != "null" {
|
||||
return u
|
||||
}
|
||||
// responses 流式:usage 在 response 对象内(response.completed 事件)
|
||||
if respRaw, ok := m["response"]; ok {
|
||||
var resp map[string]json.RawMessage
|
||||
if json.Unmarshal(respRaw, &resp) == nil {
|
||||
if u, ok := resp["usage"]; ok && string(u) != "null" {
|
||||
return u
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// sseScanner 按 SSE 行边界读取(兼容 \n 与 \r\n),保留原始行内容。
|
||||
// 基于 bufio.Reader:行内可含任意内容,跨 chunk 自动拼接。
|
||||
type sseScanner struct {
|
||||
r *bufio.Reader
|
||||
}
|
||||
|
||||
func newSSEScanner(r io.Reader) *sseScanner { return &sseScanner{r: bufio.NewReaderSize(r, 32*1024)} }
|
||||
|
||||
func (s *sseScanner) Next() ([]byte, error) {
|
||||
line, err := s.r.ReadBytes('\n')
|
||||
if len(line) > 0 {
|
||||
return line, nil
|
||||
}
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return nil, io.EOF
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// 记账
|
||||
|
||||
// usageSink 累积流式多次 usage(取最后一次,即最终值)。
|
||||
type usageSink struct {
|
||||
last json.RawMessage
|
||||
}
|
||||
|
||||
func (u *usageSink) push(raw json.RawMessage) {
|
||||
if len(raw) > 0 {
|
||||
u.last = raw
|
||||
}
|
||||
}
|
||||
|
||||
// finishUsage 落账:计算成本并异步写入。
|
||||
func (g *Gateway) finishUsage(c *gin.Context, ch *store.Channel, start time.Time, status, errCode string) {
|
||||
uid, _ := c.Get(CtxUserID)
|
||||
kid, _ := c.Get(CtxKeyID)
|
||||
trace, _ := c.Get(CtxTrace)
|
||||
|
||||
var us usageShape
|
||||
if h, ok := c.Get("usage_raw"); ok {
|
||||
if holder, ok := h.(*sinkHolder); ok && holder.sink != nil && len(holder.sink.last) > 0 {
|
||||
_ = json.Unmarshal(holder.sink.last, &us)
|
||||
}
|
||||
}
|
||||
|
||||
in := us.PromptTokens + us.InputTokens
|
||||
out := us.CompletionTokens + us.OutputTokens
|
||||
cacheRead := us.CacheReadInputTokens
|
||||
cacheCreate := us.CacheCreationInputTokens
|
||||
|
||||
modelName, _ := c.Get("model_name")
|
||||
mn, _ := modelName.(string)
|
||||
|
||||
var model store.Model
|
||||
var cost float64
|
||||
var modelID uint64
|
||||
_ = g.db.Where("name = ?", mn).First(&model).Error
|
||||
if model.ID > 0 {
|
||||
modelID = model.ID
|
||||
cost = float64(in)/1e6*model.InputPrice +
|
||||
float64(out)/1e6*model.OutputPrice +
|
||||
float64(cacheRead)/1e6*model.CacheReadPrice
|
||||
} else {
|
||||
cost = float64(in)/1e6*0.15 + float64(out)/1e6*0.60 // 无定价模型时按示例价
|
||||
}
|
||||
|
||||
proto, _ := c.Get("protocol")
|
||||
p, _ := proto.(string)
|
||||
if p == "" {
|
||||
p = "chat"
|
||||
}
|
||||
traceStr, _ := trace.(string)
|
||||
errMsg := errCode
|
||||
latency := int(time.Since(start).Milliseconds())
|
||||
|
||||
// 已写响应头但流中途出错:记 error
|
||||
if status == store.UsageStatusSuccess && c.Writer.Status() >= 400 {
|
||||
status = store.UsageStatusError
|
||||
}
|
||||
|
||||
g.rec.Record(&store.UsageLog{
|
||||
RequestID: fmt.Sprintf("trace-%s", traceStr),
|
||||
TraceID: traceStr,
|
||||
UserID: uid.(uint64),
|
||||
KeyID: kid.(uint64),
|
||||
ChannelID: ch.ID,
|
||||
ModelID: modelID,
|
||||
ModelName: mn,
|
||||
Protocol: p,
|
||||
InputTokens: in,
|
||||
OutputTokens: out,
|
||||
CacheReadTokens: cacheRead,
|
||||
CacheCreationTokens: cacheCreate,
|
||||
InputPrice: model.InputPrice,
|
||||
OutputPrice: model.OutputPrice,
|
||||
CacheReadPrice: model.CacheReadPrice,
|
||||
Cost: cost,
|
||||
LatencyMS: latency,
|
||||
Status: status,
|
||||
ErrorCode: &errMsg,
|
||||
CreatedAt: time.Now().UTC(),
|
||||
})
|
||||
}
|
||||
|
||||
// recordError 失败请求的记账(不产生扣费,status=error)。
|
||||
func (g *Gateway) recordError(c *gin.Context, ch *store.Channel, resp *http.Response, start time.Time, code string) {
|
||||
status := store.UsageStatusError
|
||||
_ = resp
|
||||
g.finishUsage(c, ch, start, status, code)
|
||||
}
|
||||
|
||||
func now() time.Time { return time.Now() }
|
||||
@@ -0,0 +1,107 @@
|
||||
package proxy
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"io"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestScanUsageChat(t *testing.T) {
|
||||
line := []byte(`data: {"id":"x","choices":[],"usage":{"prompt_tokens":12,"completion_tokens":9,"total_tokens":21}}`)
|
||||
raw := scanUsage(line)
|
||||
if raw == nil {
|
||||
t.Fatal("chat usage not detected")
|
||||
}
|
||||
var us usageShape
|
||||
if err := json.Unmarshal(raw, &us); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if us.PromptTokens != 12 || us.CompletionTokens != 9 {
|
||||
t.Fatalf("usage mismatch: %+v", us)
|
||||
}
|
||||
}
|
||||
|
||||
func TestScanUsageResponsesNested(t *testing.T) {
|
||||
line := []byte(`data: {"response":{"id":"r","status":"completed","usage":{"input_tokens":15,"output_tokens":11,"total_tokens":26}},"type":"response.completed"}`)
|
||||
raw := scanUsage(line)
|
||||
if raw == nil {
|
||||
t.Fatal("responses nested usage not detected")
|
||||
}
|
||||
var us usageShape
|
||||
if err := json.Unmarshal(raw, &us); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if us.InputTokens != 15 || us.OutputTokens != 11 {
|
||||
t.Fatalf("usage mismatch: %+v", us)
|
||||
}
|
||||
}
|
||||
|
||||
func TestScanUsageIgnoresNonData(t *testing.T) {
|
||||
if scanUsage([]byte("event: response.completed")) != nil {
|
||||
t.Fatal("event line should be ignored")
|
||||
}
|
||||
if scanUsage([]byte("data: [DONE]")) != nil {
|
||||
t.Fatal("[DONE] should be ignored")
|
||||
}
|
||||
if scanUsage([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}")) != nil {
|
||||
t.Fatal("content chunk without usage should be ignored")
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractUsageFromFullBody(t *testing.T) {
|
||||
body := []byte(`{"id":"x","choices":[{"message":{"content":"hi"}}],"usage":{"prompt_tokens":1,"completion_tokens":2}}`)
|
||||
raw := extractUsage(body)
|
||||
if raw == nil {
|
||||
t.Fatal("usage not extracted from full body")
|
||||
}
|
||||
var us usageShape
|
||||
_ = json.Unmarshal(raw, &us)
|
||||
if us.PromptTokens != 1 || us.CompletionTokens != 2 {
|
||||
t.Fatalf("usage mismatch: %+v", us)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSSEScannerLines(t *testing.T) {
|
||||
// 模拟分块写入的 SSE 流
|
||||
data := "data: {\"a\":1}\n\ndata: {\"usage\":{\"input_tokens\":3}}\n\n"
|
||||
parts := [][]byte{[]byte(data[:10]), []byte(data[10:20]), []byte(data[20:])}
|
||||
reader := newChunkReader(parts)
|
||||
s := newSSEScanner(reader)
|
||||
var lines [][]byte
|
||||
for {
|
||||
line, err := s.Next()
|
||||
if line != nil {
|
||||
lines = append(lines, line)
|
||||
}
|
||||
if err != nil {
|
||||
break
|
||||
}
|
||||
}
|
||||
if len(lines) != 4 {
|
||||
t.Fatalf("expected 4 lines, got %d", len(lines))
|
||||
}
|
||||
// 合并后应能还原原始数据
|
||||
joined := ""
|
||||
for _, l := range lines {
|
||||
joined += string(l)
|
||||
}
|
||||
if joined != string(data) {
|
||||
t.Fatalf("stream corrupted:\n got: %q\nwant: %q", joined, data)
|
||||
}
|
||||
}
|
||||
|
||||
type chunkReader struct {
|
||||
parts [][]byte
|
||||
idx int
|
||||
}
|
||||
|
||||
func newChunkReader(parts [][]byte) *chunkReader { return &chunkReader{parts: parts} }
|
||||
|
||||
func (r *chunkReader) Read(p []byte) (int, error) {
|
||||
if r.idx >= len(r.parts) {
|
||||
return 0, io.EOF
|
||||
}
|
||||
n := copy(p, r.parts[r.idx])
|
||||
r.idx++
|
||||
return n, nil
|
||||
}
|
||||
@@ -0,0 +1,34 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"log"
|
||||
|
||||
"github.com/glebarez/sqlite"
|
||||
"gorm.io/gorm"
|
||||
gormlogger "gorm.io/gorm/logger"
|
||||
)
|
||||
|
||||
// Open 打开数据库连接并自动迁移。
|
||||
// 开发默认 SQLite(dsn 支持 file:...?_journal_mode=WAL),生产可切 postgres。
|
||||
func Open(driver, dsn string) (*gorm.DB, error) {
|
||||
var dialector gorm.Dialector
|
||||
switch driver {
|
||||
case "postgres":
|
||||
dialector = postgresDialector(dsn)
|
||||
default:
|
||||
dialector = sqlite.Open(dsn)
|
||||
}
|
||||
|
||||
db, err := gorm.Open(dialector, &gorm.Config{
|
||||
Logger: gormlogger.Default.LogMode(gormlogger.Warn),
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if err := db.AutoMigrate(AllModels()...); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
log.Printf("store: connected driver=%s (migrated)", driver)
|
||||
return db, nil
|
||||
}
|
||||
@@ -0,0 +1,11 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"gorm.io/driver/postgres"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// postgresDialector 延迟引用 postgres 驱动,避免开发环境额外依赖。
|
||||
func postgresDialector(dsn string) gorm.Dialector {
|
||||
return postgres.Open(dsn)
|
||||
}
|
||||
@@ -0,0 +1,202 @@
|
||||
// Package store 数据模型与仓储层(GORM)。
|
||||
// 字段设计对应 PLANNING.md §5:金额/价格 numeric(20,8),token bigint,时间 UTC。
|
||||
package store
|
||||
|
||||
import (
|
||||
"time"
|
||||
)
|
||||
|
||||
// 角色 / 状态枚举(字符串存库,便于阅读与迁移)
|
||||
const (
|
||||
RoleUser = "user"
|
||||
RoleAdmin = "admin"
|
||||
|
||||
UserStatusActive = "active"
|
||||
UserStatusDisabled = "disabled"
|
||||
|
||||
KeyStatusActive = "active"
|
||||
KeyStatusRevoked = "revoked"
|
||||
|
||||
ChannelProviderOpenAI = "openai"
|
||||
ChannelProviderAnthropic = "anthropic"
|
||||
ChannelProviderCompatible = "compatible"
|
||||
ChannelHealthHealthy = "healthy"
|
||||
ChannelHealthDegraded = "degraded"
|
||||
ChannelHealthCooldown = "cooldown"
|
||||
|
||||
UsageStatusSuccess = "success"
|
||||
UsageStatusError = "error"
|
||||
UsageStatusCanceled = "canceled"
|
||||
|
||||
BalanceTypeRecharge = "recharge"
|
||||
BalanceTypeUsage = "usage"
|
||||
BalanceTypeRefund = "refund"
|
||||
BalanceTypeAdminAdjust = "admin_adjust"
|
||||
|
||||
RechargeStatusPending = "pending"
|
||||
RechargeStatusCredited = "credited"
|
||||
RechargeStatusRejected = "rejected"
|
||||
RechargeMethodManual = "manual"
|
||||
RechargeMethodOnline = "online"
|
||||
)
|
||||
|
||||
// User 用户(PLANNING §5.1)
|
||||
type User struct {
|
||||
ID uint64 `gorm:"primaryKey;autoIncrement" json:"id"`
|
||||
Username string `gorm:"uniqueIndex;size:64;not null" json:"username"`
|
||||
Email string `gorm:"uniqueIndex;size:255;not null" json:"email"`
|
||||
PasswordHash string `gorm:"size:255;not null" json:"-"`
|
||||
Role string `gorm:"size:16;not null;default:user" json:"role"`
|
||||
Balance float64 `gorm:"type:numeric(20,8);not null;default:0" json:"balance"`
|
||||
Status string `gorm:"size:16;not null;default:active" json:"status"`
|
||||
InviteCode *string `json:"invite_code,omitempty"`
|
||||
LastLoginAt *time.Time `json:"last_login_at,omitempty"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
}
|
||||
|
||||
// APIKey 密钥(PLANNING §5.2):库中只存 SHA-256 哈希 + 展示前缀
|
||||
type APIKey struct {
|
||||
ID uint64 `gorm:"primaryKey;autoIncrement" json:"id"`
|
||||
UserID uint64 `gorm:"index;not null" json:"user_id"`
|
||||
Name string `gorm:"size:64;not null" json:"name"`
|
||||
KeyHash string `gorm:"uniqueIndex;size:64;not null" json:"-"`
|
||||
KeyPrefix string `gorm:"size:32;not null" json:"key_prefix"`
|
||||
QuotaTokensPerDay *int64 `json:"quota_tokens_per_day,omitempty"`
|
||||
QuotaRequestsPerDay *int `json:"quota_requests_per_day,omitempty"`
|
||||
AllowedModels []string `gorm:"type:jsonb;serializer:json" json:"allowed_models,omitempty"`
|
||||
ExpiresAt *time.Time `json:"expires_at,omitempty"`
|
||||
Status string `gorm:"size:16;not null;default:active" json:"status"`
|
||||
LastUsedAt *time.Time `json:"last_used_at,omitempty"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
}
|
||||
|
||||
// Channel 上游渠道(PLANNING §5.3)
|
||||
type Channel struct {
|
||||
ID uint64 `gorm:"primaryKey;autoIncrement" json:"id"`
|
||||
Name string `gorm:"uniqueIndex;size:64;not null" json:"name"`
|
||||
Provider string `gorm:"size:16;not null" json:"provider"` // openai|anthropic|compatible
|
||||
BaseURL string `gorm:"size:255;not null" json:"base_url"`
|
||||
APIKeyEnc string `gorm:"size:1024;not null" json:"-"` // AES-GCM 密文
|
||||
Weight int `gorm:"not null;default:1" json:"weight"`
|
||||
Priority int `gorm:"not null;default:0" json:"priority"` // 数值小优先
|
||||
TimeoutMS int `gorm:"not null;default:120000" json:"timeout_ms"`
|
||||
MaxConcurrency int `gorm:"not null;default:16" json:"max_concurrency"`
|
||||
HealthStatus string `gorm:"size:16;not null;default:healthy" json:"health_status"`
|
||||
Enabled bool `gorm:"not null;default:true" json:"enabled"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
}
|
||||
|
||||
// Model 全局模型 + 定价(PLANNING §5.4,价格按每百万 token,USD)
|
||||
type Model struct {
|
||||
ID uint64 `gorm:"primaryKey;autoIncrement" json:"id"`
|
||||
Name string `gorm:"uniqueIndex;size:128;not null" json:"name"`
|
||||
DisplayName string `gorm:"size:128" json:"display_name"`
|
||||
InputPrice float64 `gorm:"type:numeric(20,8);not null;default:0" json:"input_price"`
|
||||
OutputPrice float64 `gorm:"type:numeric(20,8);not null;default:0" json:"output_price"`
|
||||
CacheReadPrice float64 `gorm:"type:numeric(20,8);not null;default:0" json:"cache_read_price"`
|
||||
Enabled bool `gorm:"not null;default:true" json:"enabled"`
|
||||
Sort int `gorm:"not null;default:0" json:"sort"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
}
|
||||
|
||||
// ChannelModelBinding 渠道↔模型绑定(多对多,PLANNING §5.4)
|
||||
type ChannelModelBinding struct {
|
||||
ID uint64 `gorm:"primaryKey;autoIncrement" json:"id"`
|
||||
ChannelID uint64 `gorm:"index:idx_ch_model,unique;not null" json:"channel_id"`
|
||||
ModelID uint64 `gorm:"index:idx_ch_model,unique;not null" json:"model_id"`
|
||||
UpstreamModel string `gorm:"size:255;not null" json:"upstream_model"`
|
||||
Weight int `gorm:"not null;default:1" json:"weight"`
|
||||
Channel Channel `gorm:"foreignKey:ChannelID" json:"-"`
|
||||
Model Model `gorm:"foreignKey:ModelID" json:"-"`
|
||||
}
|
||||
|
||||
// UsageLog 请求级用量明细(PLANNING §5.5)
|
||||
type UsageLog struct {
|
||||
ID uint64 `gorm:"primaryKey;autoIncrement" json:"id"`
|
||||
RequestID string `gorm:"size:128" json:"request_id"` // 上游 request id
|
||||
TraceID string `gorm:"size:64;index" json:"trace_id"`
|
||||
UserID uint64 `gorm:"index:idx_user_created;not null" json:"user_id"`
|
||||
KeyID uint64 `json:"key_id"`
|
||||
ChannelID uint64 `json:"channel_id"`
|
||||
ModelID uint64 `json:"model_id"`
|
||||
ModelName string `gorm:"size:128" json:"model_name"`
|
||||
Protocol string `gorm:"size:32" json:"protocol"` // responses|chat|messages
|
||||
InputTokens int64 `json:"input_tokens"`
|
||||
OutputTokens int64 `json:"output_tokens"`
|
||||
CacheReadTokens int64 `json:"cache_read_tokens"`
|
||||
CacheCreationTokens int64 `json:"cache_creation_tokens"`
|
||||
InputPrice float64 `gorm:"type:numeric(20,8)" json:"input_price"` // 快照
|
||||
OutputPrice float64 `gorm:"type:numeric(20,8)" json:"output_price"` // 快照
|
||||
CacheReadPrice float64 `gorm:"type:numeric(20,8)" json:"cache_read_price"` // 快照
|
||||
Cost float64 `gorm:"type:numeric(20,8)" json:"cost"`
|
||||
LatencyMS int `json:"latency_ms"`
|
||||
Status string `gorm:"size:16;not null" json:"status"`
|
||||
ErrorCode *string `json:"error_code,omitempty"`
|
||||
CreatedAt time.Time `gorm:"index" json:"created_at"`
|
||||
}
|
||||
|
||||
// UsageDaily 日粒度预聚合(PLANNING §5.6)
|
||||
type UsageDaily struct {
|
||||
ID uint64 `gorm:"primaryKey;autoIncrement" json:"id"`
|
||||
UserID uint64 `gorm:"index:idx_daily_user_model,unique" json:"user_id"`
|
||||
ModelID uint64 `gorm:"index:idx_daily_user_model,unique" json:"model_id"`
|
||||
Date string `gorm:"size:10;index:idx_daily_user_model,unique" json:"date"` // YYYY-MM-DD (UTC)
|
||||
Requests int64 `json:"requests"`
|
||||
InputTokens int64 `json:"input_tokens"`
|
||||
OutputTokens int64 `json:"output_tokens"`
|
||||
CacheReadTokens int64 `json:"cache_read_tokens"`
|
||||
Cost float64 `gorm:"type:numeric(20,8)" json:"cost"`
|
||||
}
|
||||
|
||||
// RechargeOrder 充值订单(PLANNING §5.7,预留:首版不做充值)
|
||||
type RechargeOrder struct {
|
||||
ID uint64 `gorm:"primaryKey;autoIncrement" json:"id"`
|
||||
UserID uint64 `gorm:"index;not null" json:"user_id"`
|
||||
Amount float64 `gorm:"type:numeric(20,8);not null" json:"amount"`
|
||||
Status string `gorm:"size:16;not null;default:pending" json:"status"`
|
||||
Method string `gorm:"size:16;not null;default:manual" json:"method"`
|
||||
TransactionID string `gorm:"size:128" json:"transaction_id,omitempty"`
|
||||
ReviewedBy *uint64 `json:"reviewed_by,omitempty"`
|
||||
ReviewedAt *time.Time `json:"reviewed_at,omitempty"`
|
||||
Remark string `gorm:"size:512" json:"remark,omitempty"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
}
|
||||
|
||||
// BalanceLog 余额流水(PLANNING §5.8,幂等:ref_id + type 唯一)
|
||||
type BalanceLog struct {
|
||||
ID uint64 `gorm:"primaryKey;autoIncrement" json:"id"`
|
||||
UserID uint64 `gorm:"index:idx_balance_user;not null" json:"user_id"`
|
||||
Change float64 `gorm:"type:numeric(20,8);not null" json:"change"`
|
||||
BalanceAfter float64 `gorm:"type:numeric(20,8);not null" json:"balance_after"`
|
||||
Type string `gorm:"size:16;not null" json:"type"`
|
||||
RefID string `gorm:"size:128;index:idx_balance_ref,unique" json:"ref_id"`
|
||||
Remark string `gorm:"size:512" json:"remark,omitempty"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
}
|
||||
|
||||
// SystemConfig 系统配置(PLANNING §5.9)
|
||||
type SystemConfig struct {
|
||||
Key string `gorm:"primaryKey;size:64" json:"key"`
|
||||
Value string `gorm:"type:jsonb;not null" json:"value"`
|
||||
}
|
||||
|
||||
// AllModels 迁移顺序(外键依赖在后)
|
||||
func AllModels() []any {
|
||||
return []any{
|
||||
&User{},
|
||||
&APIKey{},
|
||||
&Channel{},
|
||||
&Model{},
|
||||
&ChannelModelBinding{},
|
||||
&UsageLog{},
|
||||
&UsageDaily{},
|
||||
&RechargeOrder{},
|
||||
&BalanceLog{},
|
||||
&SystemConfig{},
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,144 @@
|
||||
// Package usage 异步记账:请求完成后写入 usage_logs,批量落库(PLANNING §3.2)。
|
||||
package usage
|
||||
|
||||
import (
|
||||
"log"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/openteam/server/internal/store"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
)
|
||||
|
||||
type Recorder struct {
|
||||
db *gorm.DB
|
||||
ch chan *store.UsageLog
|
||||
wg sync.WaitGroup
|
||||
closed chan struct{}
|
||||
}
|
||||
|
||||
const batchSize = 32
|
||||
|
||||
func NewRecorder(db *gorm.DB) *Recorder {
|
||||
r := &Recorder{
|
||||
db: db,
|
||||
ch: make(chan *store.UsageLog, 512),
|
||||
closed: make(chan struct{}),
|
||||
}
|
||||
r.wg.Add(1)
|
||||
go r.run()
|
||||
return r
|
||||
}
|
||||
|
||||
// Record 提交一条用量(非阻塞;队列满时同步写入,保证不丢账)。
|
||||
func (r *Recorder) Record(l *store.UsageLog) {
|
||||
select {
|
||||
case r.ch <- l:
|
||||
default:
|
||||
// 队列积压:直接同步写,避免丢账
|
||||
if err := r.flush([]*store.UsageLog{l}); err != nil {
|
||||
log.Printf("usage: sync write failed: %v", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (r *Recorder) Close() {
|
||||
close(r.closed)
|
||||
r.wg.Wait()
|
||||
close(r.ch)
|
||||
}
|
||||
|
||||
func (r *Recorder) run() {
|
||||
defer r.wg.Done()
|
||||
buf := make([]*store.UsageLog, 0, batchSize)
|
||||
tick := time.NewTicker(2 * time.Second)
|
||||
defer tick.Stop()
|
||||
for {
|
||||
select {
|
||||
case l, ok := <-r.ch:
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
buf = append(buf, l)
|
||||
if len(buf) >= batchSize {
|
||||
if err := r.flush(buf); err != nil {
|
||||
log.Printf("usage: batch write failed: %v", err)
|
||||
}
|
||||
buf = buf[:0]
|
||||
}
|
||||
case <-r.closed:
|
||||
if len(buf) > 0 {
|
||||
if err := r.flush(buf); err != nil {
|
||||
log.Printf("usage: final batch write failed: %v", err)
|
||||
}
|
||||
}
|
||||
return
|
||||
case <-tick.C:
|
||||
if len(buf) > 0 {
|
||||
if err := r.flush(buf); err != nil {
|
||||
log.Printf("usage: batch write failed: %v", err)
|
||||
}
|
||||
buf = buf[:0]
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// flush 批量插入用量明细,并同步更新余额、余额流水与日聚合。
|
||||
// 记账口径:单次成本 = in×in_price + out×out_price + cache_read×cache_read_price(每百万 token)。
|
||||
func (r *Recorder) flush(logs []*store.UsageLog) error {
|
||||
if len(logs) == 0 {
|
||||
return nil
|
||||
}
|
||||
return r.db.Transaction(func(tx *gorm.DB) error {
|
||||
if err := tx.Create(logs).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
for _, l := range logs {
|
||||
if l.Status != store.UsageStatusSuccess || l.Cost <= 0 {
|
||||
continue
|
||||
}
|
||||
// 扣余额(余额可为负,流式请求不中断;后续请求被拒)
|
||||
var user store.User
|
||||
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).First(&user, l.UserID).Error; err != nil {
|
||||
continue
|
||||
}
|
||||
newBalance := user.Balance - l.Cost
|
||||
tx.Model(&store.User{}).Where("id = ?", l.UserID).Update("balance", newBalance)
|
||||
tx.Create(&store.BalanceLog{
|
||||
UserID: l.UserID,
|
||||
Change: -l.Cost,
|
||||
BalanceAfter: newBalance,
|
||||
Type: store.BalanceTypeUsage,
|
||||
RefID: usageRefID(l.TraceID),
|
||||
Remark: "usage: " + l.ModelName,
|
||||
})
|
||||
|
||||
// 日聚合 upsert
|
||||
date := l.CreatedAt.UTC().Format("2006-01-02")
|
||||
tx.Clauses(clause.OnConflict{
|
||||
Columns: []clause.Column{{Name: "user_id"}, {Name: "model_id"}, {Name: "date"}},
|
||||
DoUpdates: clause.Assignments(map[string]any{
|
||||
"requests": gorm.Expr("requests + 1"),
|
||||
"input_tokens": gorm.Expr("input_tokens + ?", l.InputTokens),
|
||||
"output_tokens": gorm.Expr("output_tokens + ?", l.OutputTokens),
|
||||
"cache_read_tokens": gorm.Expr("cache_read_tokens + ?", l.CacheReadTokens),
|
||||
"cost": gorm.Expr("cost + ?", l.Cost),
|
||||
}),
|
||||
}).Create(&store.UsageDaily{
|
||||
UserID: l.UserID, ModelID: l.ModelID, Date: date,
|
||||
Requests: 1, InputTokens: l.InputTokens, OutputTokens: l.OutputTokens,
|
||||
CacheReadTokens: l.CacheReadTokens, Cost: l.Cost,
|
||||
})
|
||||
}
|
||||
return nil
|
||||
})
|
||||
}
|
||||
|
||||
func usageRefID(traceID string) string {
|
||||
if traceID == "" {
|
||||
traceID = "unknown"
|
||||
}
|
||||
return "usage:" + traceID
|
||||
}
|
||||
Reference in New Issue
Block a user