diff --git a/bin/start-local.sh b/bin/start-local.sh new file mode 100755 index 00000000000..60d59cc7052 --- /dev/null +++ b/bin/start-local.sh @@ -0,0 +1,58 @@ +#!/usr/bin/env bash +set -euo pipefail + +usage() { + cat <<'EOF' +Usage: bin/start-local.sh + +Start the local new-api service from the repository root. + +Environment: + GO_VERSION Optional gvm Go version, for example: go1.24.6 + GOCACHE Optional Go build cache path. Defaults to ./.gocache + +Notes: + - The Go application loads .env itself. + - If .env is missing, the service starts with built-in defaults. +EOF +} + +if [[ "${1:-}" == "-h" || "${1:-}" == "--help" ]]; then + usage + exit 0 +fi + +script_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +repo_root="$(cd "${script_dir}/.." && pwd)" +cd "${repo_root}" + +if [[ -f "${HOME}/.gvm/scripts/gvm" ]]; then + # gvm scripts reference optional shell variables such as ZSH_VERSION and + # GVM_DEBUG directly, and may return non-zero under errexit in some shells. + set +eu + # shellcheck source=/dev/null + source "${HOME}/.gvm/scripts/gvm" + if [[ -n "${GO_VERSION:-}" ]]; then + gvm use "${GO_VERSION}" >/dev/null + fi + set -euo pipefail +fi + +if ! command -v go >/dev/null 2>&1; then + echo "go command not found. Install Go or set GO_VERSION after installing gvm." >&2 + exit 1 +fi + +if [[ ! -f ".env" ]]; then + echo "warning: .env not found; using application defaults" >&2 +fi + +export GOCACHE="${GOCACHE:-${repo_root}/.gocache}" +mkdir -p "${GOCACHE}" + +echo "Starting new-api from ${repo_root}" +echo "Go: $(go version)" +echo "GOCACHE: ${GOCACHE}" +echo "Config: ${repo_root}/.env" + +exec go run main.go diff --git a/controller/agent_hub.go b/controller/agent_hub.go new file mode 100644 index 00000000000..aa97a0a8078 --- /dev/null +++ b/controller/agent_hub.go @@ -0,0 +1,105 @@ +package controller + +import ( + "strconv" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/i18n" + "github.com/QuantumNous/new-api/logger" + "github.com/QuantumNous/new-api/model" + "github.com/gin-gonic/gin" +) + +type agentHubUserProvisionRequest struct { + Username string `json:"username"` + Password string `json:"password"` + InitialQuota int `json:"initial_quota"` + DisplayName string `json:"display_name"` + Group string `json:"group"` +} + +type agentHubQuotaAdjustmentRequest struct { + RequestId string `json:"request_id"` + Delta int `json:"delta"` + Reason string `json:"reason"` +} + +func CreateAgentHubUserProvision(c *gin.Context) { + var req agentHubUserProvisionRequest + if err := common.DecodeJson(c.Request.Body, &req); err != nil { + common.ApiErrorI18n(c, i18n.MsgInvalidParams) + return + } + + result, err := model.EnsureAgentHubUserProvision(req.Username, req.Password, req.DisplayName, req.Group, req.InitialQuota) + if err != nil { + common.ApiError(c, err) + return + } + + common.ApiSuccess(c, gin.H{ + "user_id": result.UserId, + "api_key": result.ApiKey, + "quota": result.Quota, + }) +} + +func CreateAgentHubQuotaAdjustment(c *gin.Context) { + userId, err := strconv.Atoi(c.Param("user_id")) + if err != nil { + common.ApiErrorI18n(c, i18n.MsgInvalidParams) + return + } + + var req agentHubQuotaAdjustmentRequest + if err := common.DecodeJson(c.Request.Body, &req); err != nil { + common.ApiErrorI18n(c, i18n.MsgInvalidParams) + return + } + + adjustment, replayed, err := model.ApplyAgentHubQuotaAdjustment(req.RequestId, userId, req.Delta, req.Reason) + if err != nil { + common.ApiError(c, err) + return + } + if !replayed { + action := "user.quota_add" + quota := adjustment.Delta + if adjustment.Delta < 0 { + action = "user.quota_subtract" + quota = -adjustment.Delta + } + recordManageAuditFor(c, adjustment.UserId, action, map[string]interface{}{ + "quota": logger.LogQuota(quota), + }) + } + + common.ApiSuccess(c, gin.H{ + "request_id": adjustment.RequestId, + "user_id": adjustment.UserId, + "delta": adjustment.Delta, + "quota_before": adjustment.QuotaBefore, + "quota_after": adjustment.QuotaAfter, + "replayed": replayed, + }) +} + +func GetAgentHubUserQuota(c *gin.Context) { + userId, err := strconv.Atoi(c.Param("user_id")) + if err != nil { + common.ApiErrorI18n(c, i18n.MsgInvalidParams) + return + } + + quota, err := model.GetAgentHubUserQuota(userId) + if err != nil { + common.ApiError(c, err) + return + } + + common.ApiSuccess(c, gin.H{ + "user_id": quota.UserId, + "quota": quota.Quota, + "used_quota": quota.UsedQuota, + }) +} diff --git a/controller/agent_hub_test.go b/controller/agent_hub_test.go new file mode 100644 index 00000000000..7a815a176c6 --- /dev/null +++ b/controller/agent_hub_test.go @@ -0,0 +1,365 @@ +package controller + +import ( + "fmt" + "net/http" + "strconv" + "strings" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/logger" + "github.com/QuantumNous/new-api/model" + "github.com/gin-gonic/gin" + "github.com/glebarez/sqlite" + "gorm.io/gorm" +) + +type agentHubProvisionResponse struct { + UserId int `json:"user_id"` + ApiKey string `json:"api_key"` + Quota int `json:"quota"` +} + +type agentHubQuotaAdjustmentResponse struct { + RequestId string `json:"request_id"` + UserId int `json:"user_id"` + Delta int `json:"delta"` + QuotaBefore int `json:"quota_before"` + QuotaAfter int `json:"quota_after"` + Replayed bool `json:"replayed"` +} + +type agentHubQuotaResponse struct { + UserId int `json:"user_id"` + Quota int `json:"quota"` + UsedQuota int `json:"used_quota"` +} + +func setupAgentHubControllerTestDB(t *testing.T) *gorm.DB { + t.Helper() + + gin.SetMode(gin.TestMode) + common.SetDatabaseTypes(common.DatabaseTypeSQLite, common.DatabaseTypeSQLite) + common.RedisEnabled = false + common.BatchUpdateEnabled = false + + dsn := fmt.Sprintf("file:%s?mode=memory&cache=shared", strings.ReplaceAll(t.Name(), "/", "_")) + db, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{}) + if err != nil { + t.Fatalf("failed to open sqlite db: %v", err) + } + model.DB = db + model.LOG_DB = db + + if err := db.AutoMigrate( + &model.User{}, + &model.Token{}, + &model.AgentHubQuotaAdjustment{}, + &model.Log{}, + ); err != nil { + t.Fatalf("failed to migrate test db: %v", err) + } + + t.Cleanup(func() { + sqlDB, err := db.DB() + if err == nil { + _ = sqlDB.Close() + } + }) + + return db +} + +func seedAgentHubUser(t *testing.T, db *gorm.DB, username string, quota int, usedQuota int) *model.User { + t.Helper() + + user := &model.User{ + Username: username, + Password: "hashed-password", + DisplayName: username, + Role: common.RoleCommonUser, + Status: common.UserStatusEnabled, + Quota: quota, + UsedQuota: usedQuota, + Group: "default", + AffCode: common.GetRandomString(4), + Setting: "{}", + } + if err := db.Create(user).Error; err != nil { + t.Fatalf("failed to seed user: %v", err) + } + return user +} + +func TestAgentHubUserProvisionCreatesAndReusesToken(t *testing.T) { + db := setupAgentHubControllerTestDB(t) + body := map[string]any{ + "username": "ah_user_1", + "password": "pass12345", + "initial_quota": 123, + "display_name": "Agent Hub User", + "group": "default", + } + + ctx, recorder := newAuthenticatedContext(t, http.MethodPost, "/api/agent-hub/user-provisions", body, 1) + CreateAgentHubUserProvision(ctx) + + response := decodeAPIResponse(t, recorder) + if !response.Success { + t.Fatalf("expected provision to succeed, got message: %s", response.Message) + } + var provision agentHubProvisionResponse + if err := common.Unmarshal(response.Data, &provision); err != nil { + t.Fatalf("failed to decode provision response: %v", err) + } + if provision.UserId == 0 { + t.Fatalf("expected provision user_id") + } + if provision.Quota != 123 { + t.Fatalf("expected quota 123, got %d", provision.Quota) + } + if !strings.HasPrefix(provision.ApiKey, "sk-") { + t.Fatalf("expected api_key to include sk- prefix, got %q", provision.ApiKey) + } + + var tokenCount int64 + if err := db.Model(&model.Token{}).Where("user_id = ?", provision.UserId).Count(&tokenCount).Error; err != nil { + t.Fatalf("failed to count tokens: %v", err) + } + if tokenCount != 1 { + t.Fatalf("expected one token, got %d", tokenCount) + } + + ctx, recorder = newAuthenticatedContext(t, http.MethodPost, "/api/agent-hub/user-provisions", body, 1) + CreateAgentHubUserProvision(ctx) + response = decodeAPIResponse(t, recorder) + if !response.Success { + t.Fatalf("expected repeated provision to succeed, got message: %s", response.Message) + } + var repeated agentHubProvisionResponse + if err := common.Unmarshal(response.Data, &repeated); err != nil { + t.Fatalf("failed to decode repeated provision response: %v", err) + } + if repeated.UserId != provision.UserId || repeated.ApiKey != provision.ApiKey { + t.Fatalf("expected repeated provision to reuse user/key, first=%+v repeated=%+v", provision, repeated) + } + if err := db.Model(&model.Token{}).Where("user_id = ?", provision.UserId).Count(&tokenCount).Error; err != nil { + t.Fatalf("failed to count repeated tokens: %v", err) + } + if tokenCount != 1 { + t.Fatalf("expected repeated provision to keep one token, got %d", tokenCount) + } +} + +func TestAgentHubUserProvisionReusesExistingUserAndCreatesToken(t *testing.T) { + db := setupAgentHubControllerTestDB(t) + user := seedAgentHubUser(t, db, "ah_exists", 100, 0) + + body := map[string]any{ + "username": "ah_exists", + "password": "pass12345", + "initial_quota": 123, + } + ctx, recorder := newAuthenticatedContext(t, http.MethodPost, "/api/agent-hub/user-provisions", body, 1) + CreateAgentHubUserProvision(ctx) + + response := decodeAPIResponse(t, recorder) + if !response.Success { + t.Fatalf("expected provision to reuse existing user, got message: %s", response.Message) + } + var provision agentHubProvisionResponse + if err := common.Unmarshal(response.Data, &provision); err != nil { + t.Fatalf("failed to decode provision response: %v", err) + } + if provision.UserId != user.Id || provision.Quota != 100 { + t.Fatalf("expected existing user/quota to be reused, got %+v", provision) + } + if !strings.HasPrefix(provision.ApiKey, "sk-") { + t.Fatalf("expected api_key to include sk- prefix, got %q", provision.ApiKey) + } + var tokenCount int64 + if err := db.Model(&model.Token{}).Where("user_id = ? AND name = ?", user.Id, "agent-hub").Count(&tokenCount).Error; err != nil { + t.Fatalf("failed to count agent-hub tokens: %v", err) + } + if tokenCount != 1 { + t.Fatalf("expected one agent-hub token, got %d", tokenCount) + } +} + +func TestAgentHubQuotaAdjustmentIdempotencyAndQuery(t *testing.T) { + db := setupAgentHubControllerTestDB(t) + user := seedAgentHubUser(t, db, "ah_quota", 100, 7) + + body := map[string]any{ + "request_id": "media-1-pre", + "delta": -30, + "reason": "pre_deduct", + } + ctx, recorder := newAuthenticatedContext(t, http.MethodPost, "/api/agent-hub/users/"+strconv.Itoa(user.Id)+"/quota-adjustments", body, 1) + ctx.Set("username", "root") + ctx.Set("role", common.RoleRootUser) + ctx.Set("use_access_token", true) + ctx.Params = gin.Params{{Key: "user_id", Value: strconv.Itoa(user.Id)}} + CreateAgentHubQuotaAdjustment(ctx) + + response := decodeAPIResponse(t, recorder) + if !response.Success { + t.Fatalf("expected quota adjustment to succeed, got message: %s", response.Message) + } + var adjustment agentHubQuotaAdjustmentResponse + if err := common.Unmarshal(response.Data, &adjustment); err != nil { + t.Fatalf("failed to decode quota adjustment: %v", err) + } + if adjustment.QuotaBefore != 100 || adjustment.QuotaAfter != 70 || adjustment.Replayed { + t.Fatalf("unexpected first adjustment response: %+v", adjustment) + } + assertAgentHubQuotaAuditLog(t, db, user, "user.quota_subtract", logger.LogQuota(30), 1) + + ctx, recorder = newAuthenticatedContext(t, http.MethodPost, "/api/agent-hub/users/"+strconv.Itoa(user.Id)+"/quota-adjustments", body, 1) + ctx.Set("username", "root") + ctx.Set("role", common.RoleRootUser) + ctx.Set("use_access_token", true) + ctx.Params = gin.Params{{Key: "user_id", Value: strconv.Itoa(user.Id)}} + CreateAgentHubQuotaAdjustment(ctx) + response = decodeAPIResponse(t, recorder) + if !response.Success { + t.Fatalf("expected repeated quota adjustment to succeed, got message: %s", response.Message) + } + var repeated agentHubQuotaAdjustmentResponse + if err := common.Unmarshal(response.Data, &repeated); err != nil { + t.Fatalf("failed to decode repeated quota adjustment: %v", err) + } + if !repeated.Replayed || repeated.QuotaAfter != 70 { + t.Fatalf("expected repeated adjustment replay with quota 70, got %+v", repeated) + } + assertAgentHubQuotaAuditLog(t, db, user, "user.quota_subtract", logger.LogQuota(30), 1) + + conflictBody := map[string]any{ + "request_id": "media-1-pre", + "delta": -20, + "reason": "pre_deduct", + } + ctx, recorder = newAuthenticatedContext(t, http.MethodPost, "/api/agent-hub/users/"+strconv.Itoa(user.Id)+"/quota-adjustments", conflictBody, 1) + ctx.Set("username", "root") + ctx.Set("role", common.RoleRootUser) + ctx.Set("use_access_token", true) + ctx.Params = gin.Params{{Key: "user_id", Value: strconv.Itoa(user.Id)}} + CreateAgentHubQuotaAdjustment(ctx) + response = decodeAPIResponse(t, recorder) + if response.Success { + t.Fatalf("expected conflicting quota adjustment to fail") + } + assertAgentHubQuotaAuditLog(t, db, user, "user.quota_subtract", logger.LogQuota(30), 1) + + insufficientBody := map[string]any{ + "request_id": "media-2-pre", + "delta": -1000, + } + ctx, recorder = newAuthenticatedContext(t, http.MethodPost, "/api/agent-hub/users/"+strconv.Itoa(user.Id)+"/quota-adjustments", insufficientBody, 1) + ctx.Set("username", "root") + ctx.Set("role", common.RoleRootUser) + ctx.Set("use_access_token", true) + ctx.Params = gin.Params{{Key: "user_id", Value: strconv.Itoa(user.Id)}} + CreateAgentHubQuotaAdjustment(ctx) + response = decodeAPIResponse(t, recorder) + if response.Success { + t.Fatalf("expected insufficient quota adjustment to fail") + } + assertAgentHubQuotaAuditLog(t, db, user, "user.quota_subtract", logger.LogQuota(30), 1) + + addBody := map[string]any{ + "request_id": "media-3-refund", + "delta": 15, + "reason": "refund", + } + ctx, recorder = newAuthenticatedContext(t, http.MethodPost, "/api/agent-hub/users/"+strconv.Itoa(user.Id)+"/quota-adjustments", addBody, 1) + ctx.Set("username", "root") + ctx.Set("role", common.RoleRootUser) + ctx.Set("use_access_token", true) + ctx.Params = gin.Params{{Key: "user_id", Value: strconv.Itoa(user.Id)}} + CreateAgentHubQuotaAdjustment(ctx) + response = decodeAPIResponse(t, recorder) + if !response.Success { + t.Fatalf("expected positive quota adjustment to succeed, got message: %s", response.Message) + } + var added agentHubQuotaAdjustmentResponse + if err := common.Unmarshal(response.Data, &added); err != nil { + t.Fatalf("failed to decode positive quota adjustment: %v", err) + } + if added.QuotaBefore != 70 || added.QuotaAfter != 85 || added.Replayed { + t.Fatalf("unexpected positive adjustment response: %+v", added) + } + assertAgentHubQuotaAuditLog(t, db, user, "user.quota_add", logger.LogQuota(15), 2) + + ctx, recorder = newAuthenticatedContext(t, http.MethodGet, "/api/agent-hub/users/"+strconv.Itoa(user.Id)+"/quota", nil, 1) + ctx.Params = gin.Params{{Key: "user_id", Value: strconv.Itoa(user.Id)}} + GetAgentHubUserQuota(ctx) + response = decodeAPIResponse(t, recorder) + if !response.Success { + t.Fatalf("expected quota query to succeed, got message: %s", response.Message) + } + var quota agentHubQuotaResponse + if err := common.Unmarshal(response.Data, "a); err != nil { + t.Fatalf("failed to decode quota response: %v", err) + } + if quota.UserId != user.Id || quota.Quota != 85 || quota.UsedQuota != 7 { + t.Fatalf("unexpected quota response: %+v", quota) + } +} + +func assertAgentHubQuotaAuditLog(t *testing.T, db *gorm.DB, user *model.User, action string, quota string, count int64) { + t.Helper() + + var logs []model.Log + if err := db.Where("type = ?", model.LogTypeManage).Order("id asc").Find(&logs).Error; err != nil { + t.Fatalf("failed to query audit logs: %v", err) + } + if int64(len(logs)) != count { + t.Fatalf("expected %d manage audit logs, got %d", count, len(logs)) + } + if len(logs) == 0 { + return + } + + log := logs[len(logs)-1] + if log.UserId != user.Id || log.Username != user.Username { + t.Fatalf("expected audit log to belong to target user %d/%s, got user_id=%d username=%s", user.Id, user.Username, log.UserId, log.Username) + } + if strings.Contains(log.Content, "/api/agent-hub") { + t.Fatalf("expected audit log content to describe quota adjustment, got route content: %s", log.Content) + } + expectedContent := map[string]string{ + "user.quota_add": "Increased user quota by " + quota, + "user.quota_subtract": "Decreased user quota by " + quota, + }[action] + if log.Content != expectedContent { + t.Fatalf("expected audit log content %q, got %q", expectedContent, log.Content) + } + + other, err := common.StrToMap(log.Other) + if err != nil { + t.Fatalf("failed to decode log other: %v", err) + } + op, ok := other["op"].(map[string]interface{}) + if !ok { + t.Fatalf("expected log other.op, got %+v", other) + } + if op["action"] != action { + t.Fatalf("expected op action %s, got %+v", action, op["action"]) + } + params, ok := op["params"].(map[string]interface{}) + if !ok { + t.Fatalf("expected op params, got %+v", op) + } + if params["quota"] != quota { + t.Fatalf("expected quota param %q, got %+v", quota, params["quota"]) + } + adminInfo, ok := other["admin_info"].(map[string]interface{}) + if !ok { + t.Fatalf("expected admin_info, got %+v", other) + } + if adminInfo["auth_method"] != "access_token" { + t.Fatalf("expected access_token auth method, got %+v", adminInfo["auth_method"]) + } +} diff --git a/model/agent_hub.go b/model/agent_hub.go new file mode 100644 index 00000000000..ac0ef0acfe8 --- /dev/null +++ b/model/agent_hub.go @@ -0,0 +1,284 @@ +package model + +import ( + "errors" + "fmt" + "strings" + + "github.com/QuantumNous/new-api/common" + "gorm.io/gorm" +) + +const agentHubTokenName = "agent-hub" + +var ( + ErrAgentHubQuotaConflict = errors.New("agent-hub quota adjustment request_id conflict") + ErrAgentHubQuotaInsufficient = errors.New("agent-hub quota insufficient") +) + +type AgentHubQuotaAdjustment struct { + Id int `json:"id" gorm:"comment:主键 ID"` + RequestId string `json:"request_id" gorm:"type:varchar(64);uniqueIndex;comment:agent-hub-api 生成的全局唯一幂等请求 ID"` + UserId int `json:"user_id" gorm:"index;comment:被调整配额的 New API 用户 ID"` + Delta int `json:"delta" gorm:"type:int;not null;default:0;comment:配额调整值，正数为增加，负数为扣减"` + QuotaBefore int `json:"quota_before" gorm:"type:int;not null;default:0;comment:本次调整前的用户剩余额度"` + QuotaAfter int `json:"quota_after" gorm:"type:int;not null;default:0;comment:本次调整后的用户剩余额度"` + Reason string `json:"reason" gorm:"type:varchar(255);default:'';comment:配额调整原因"` + Status string `json:"status" gorm:"type:varchar(32);index;default:'succeeded';comment:配额调整状态"` + CreatedAt int64 `json:"created_at" gorm:"bigint;comment:创建时间戳，单位秒"` + UpdatedAt int64 `json:"updated_at" gorm:"bigint;index;comment:更新时间戳，单位秒"` +} + +func (a *AgentHubQuotaAdjustment) BeforeCreate(tx *gorm.DB) error { + now := common.GetTimestamp() + a.CreatedAt = now + a.UpdatedAt = now + if a.Status == "" { + a.Status = "succeeded" + } + return nil +} + +func (a *AgentHubQuotaAdjustment) BeforeUpdate(tx *gorm.DB) error { + a.UpdatedAt = common.GetTimestamp() + return nil +} + +type AgentHubProvisionResult struct { + UserId int + ApiKey string + Quota int +} + +type AgentHubQuota struct { + UserId int + Quota int + UsedQuota int +} + +func EnsureAgentHubUserProvision(username string, password string, displayName string, group string, initialQuota int) (*AgentHubProvisionResult, error) { + username = strings.TrimSpace(username) + displayName = strings.TrimSpace(displayName) + group = strings.TrimSpace(group) + if username == "" || password == "" { + return nil, errors.New("username and password are required") + } + if len(username) > UserNameMaxLength { + return nil, fmt.Errorf("username length must be <= %d", UserNameMaxLength) + } + if len(password) < 8 || len(password) > 20 { + return nil, errors.New("password length must be between 8 and 20") + } + if initialQuota < 0 { + return nil, errors.New("initial_quota must be >= 0") + } + if displayName == "" { + displayName = username + } + if group == "" { + group = "default" + } + + result := &AgentHubProvisionResult{} + err := DB.Transaction(func(tx *gorm.DB) error { + var user User + query := tx.Where("username = ?", username).Limit(1).Find(&user) + if query.Error != nil { + return query.Error + } + if query.RowsAffected == 0 { + passwordHash, err := common.Password2Hash(password) + if err != nil { + return err + } + user = User{ + Username: username, + Password: passwordHash, + DisplayName: displayName, + Role: common.RoleCommonUser, + Status: common.UserStatusEnabled, + Quota: initialQuota, + Group: group, + AffCode: common.GetRandomString(4), + Setting: "{}", + } + if err := tx.Create(&user).Error; err != nil { + return err + } + } + + if user.Status != common.UserStatusEnabled { + return errors.New("agent-hub provision user is disabled") + } + + token, err := getOrCreateAgentHubTokenTx(tx, user.Id, user.Group) + if err != nil { + return err + } + + result.UserId = user.Id + result.ApiKey = agentHubAPIKey(token.Key) + result.Quota = user.Quota + return nil + }) + if err != nil { + return nil, err + } + if common.RedisEnabled { + _ = updateUserQuotaCache(result.UserId, result.Quota) + } + return result, nil +} + +func GetAgentHubUserQuota(userId int) (*AgentHubQuota, error) { + if userId <= 0 { + return nil, errors.New("invalid user_id") + } + var user User + if err := DB.Select("id", "quota", "used_quota").Where("id = ?", userId).First(&user).Error; err != nil { + return nil, err + } + return &AgentHubQuota{ + UserId: user.Id, + Quota: user.Quota, + UsedQuota: user.UsedQuota, + }, nil +} + +func ApplyAgentHubQuotaAdjustment(requestId string, userId int, delta int, reason string) (*AgentHubQuotaAdjustment, bool, error) { + requestId = strings.TrimSpace(requestId) + reason = strings.TrimSpace(reason) + if requestId == "" { + return nil, false, errors.New("request_id is required") + } + if len(requestId) > 64 { + return nil, false, errors.New("request_id length must be <= 64") + } + if userId <= 0 { + return nil, false, errors.New("invalid user_id") + } + if delta == 0 { + return nil, false, errors.New("delta must not be 0") + } + if len(reason) > 255 { + return nil, false, errors.New("reason length must be <= 255") + } + + var result *AgentHubQuotaAdjustment + replayed := false + err := DB.Transaction(func(tx *gorm.DB) error { + var existing AgentHubQuotaAdjustment + query := tx.Where("request_id = ?", requestId).Limit(1).Find(&existing) + if query.Error != nil { + return query.Error + } + if query.RowsAffected > 0 { + if existing.UserId != userId || existing.Delta != delta || existing.Reason != reason { + return ErrAgentHubQuotaConflict + } + result = &existing + replayed = true + return nil + } + + adjustment := AgentHubQuotaAdjustment{ + RequestId: requestId, + UserId: userId, + Delta: delta, + Reason: reason, + Status: "pending", + } + if err := tx.Create(&adjustment).Error; err != nil { + var dup AgentHubQuotaAdjustment + if err2 := tx.Where("request_id = ?", requestId).First(&dup).Error; err2 == nil { + if dup.UserId != userId || dup.Delta != delta || dup.Reason != reason { + return ErrAgentHubQuotaConflict + } + result = &dup + replayed = true + return nil + } + return err + } + + update := tx.Model(&User{}).Where("id = ?", userId) + if delta < 0 { + update = update.Where("quota >= ?", -delta) + } + updateResult := update.Update("quota", gorm.Expr("quota + ?", delta)) + if updateResult.Error != nil { + return updateResult.Error + } + if updateResult.RowsAffected == 0 { + var user User + if err := tx.Select("id").Where("id = ?", userId).First(&user).Error; err != nil { + return err + } + return ErrAgentHubQuotaInsufficient + } + + var user User + if err := tx.Select("quota").Where("id = ?", userId).First(&user).Error; err != nil { + return err + } + adjustment.QuotaAfter = user.Quota + adjustment.QuotaBefore = user.Quota - delta + adjustment.Status = "succeeded" + if err := tx.Save(&adjustment).Error; err != nil { + return err + } + result = &adjustment + return nil + }) + if err != nil { + return nil, false, err + } + if common.RedisEnabled && result != nil { + _ = updateUserQuotaCache(userId, result.QuotaAfter) + } + return result, replayed, nil +} + +func getOrCreateAgentHubTokenTx(tx *gorm.DB, userId int, group string) (*Token, error) { + var token Token + now := common.GetTimestamp() + err := tx.Where( + "user_id = ? AND name = ? AND status = ? AND (expired_time = -1 OR expired_time > ?)", + userId, + agentHubTokenName, + common.TokenStatusEnabled, + now, + ).Order("id asc").First(&token).Error + if err == nil { + return &token, nil + } + if !errors.Is(err, gorm.ErrRecordNotFound) { + return nil, err + } + key, err := common.GenerateKey() + if err != nil { + return nil, err + } + token = Token{ + UserId: userId, + Name: agentHubTokenName, + Key: key, + Status: common.TokenStatusEnabled, + CreatedTime: now, + AccessedTime: now, + ExpiredTime: -1, + UnlimitedQuota: true, + Group: group, + } + if err := tx.Create(&token).Error; err != nil { + return nil, err + } + return &token, nil +} + +func agentHubAPIKey(key string) string { + if strings.HasPrefix(key, "sk-") { + return key + } + return "sk-" + key +} diff --git a/model/main.go b/model/main.go index 76f98a59c30..ab622c7dfa8 100644 --- a/model/main.go +++ b/model/main.go @@ -291,6 +291,7 @@ func migrateDB() error { &SubscriptionOrder{}, &UserSubscription{}, &SubscriptionPreConsumeRecord{}, + &AgentHubQuotaAdjustment{}, &CustomOAuthProvider{}, &UserOAuthBinding{}, &PerfMetric{}, @@ -303,6 +304,9 @@ func migrateDB() error { if err != nil { return err } + if err := ensureAgentHubQuotaAdjustmentTableComment(); err != nil { + return err + } if common.UsingMainDatabase(common.DatabaseTypeSQLite) { if err := ensureSubscriptionPlanTableSQLite(); err != nil { return err @@ -345,6 +349,7 @@ func migrateDBFast() error { {&SubscriptionOrder{}, "SubscriptionOrder"}, {&UserSubscription{}, "UserSubscription"}, {&SubscriptionPreConsumeRecord{}, "SubscriptionPreConsumeRecord"}, + {&AgentHubQuotaAdjustment{}, "AgentHubQuotaAdjustment"}, {&CustomOAuthProvider{}, "CustomOAuthProvider"}, {&UserOAuthBinding{}, "UserOAuthBinding"}, {&PerfMetric{}, "PerfMetric"}, @@ -375,6 +380,9 @@ func migrateDBFast() error { return err } } + if err := ensureAgentHubQuotaAdjustmentTableComment(); err != nil { + return err + } if common.UsingMainDatabase(common.DatabaseTypeSQLite) { if err := ensureSubscriptionPlanTableSQLite(); err != nil { return err @@ -388,6 +396,13 @@ func migrateDBFast() error { return nil } +func ensureAgentHubQuotaAdjustmentTableComment() error { + if !common.UsingMainDatabase(common.DatabaseTypeMySQL) { + return nil + } + return DB.Exec("ALTER TABLE `agent_hub_quota_adjustments` COMMENT = 'Agent Hub 配额调整幂等记录'").Error +} + func migrateLOGDB() error { if common.UsingLogDatabase(common.DatabaseTypeClickHouse) { return migrateClickHouseLogDB() diff --git a/router/agent_hub.go b/router/agent_hub.go new file mode 100644 index 00000000000..a801f21b4b9 --- /dev/null +++ b/router/agent_hub.go @@ -0,0 +1,17 @@ +package router + +import ( + "github.com/QuantumNous/new-api/controller" + "github.com/QuantumNous/new-api/middleware" + "github.com/gin-gonic/gin" +) + +func setAgentHubRouter(apiRouter *gin.RouterGroup) { + agentHubRoute := apiRouter.Group("/agent-hub") + agentHubRoute.Use(middleware.RootAuth()) + { + agentHubRoute.POST("/user-provisions", controller.CreateAgentHubUserProvision) + agentHubRoute.POST("/users/:user_id/quota-adjustments", controller.CreateAgentHubQuotaAdjustment) + agentHubRoute.GET("/users/:user_id/quota", controller.GetAgentHubUserQuota) + } +} diff --git a/router/api-router.go b/router/api-router.go index 83f9259b213..44f317aeef2 100644 --- a/router/api-router.go +++ b/router/api-router.go @@ -19,6 +19,7 @@ func SetApiRouter(router *gin.Engine) { apiRouter.Use(middleware.GlobalAPIRateLimit()) anonymousRequestBodyLimit := middleware.AnonymousRequestBodyLimit() { + setAgentHubRouter(apiRouter) apiRouter.GET("/setup", controller.GetSetup) apiRouter.POST("/setup", anonymousRequestBodyLimit, controller.PostSetup) apiRouter.GET("/status", controller.GetStatus) diff --git a/web/bun.lock b/web/bun.lock index d86f3a8b3ff..88e23207f2d 100644 --- a/web/bun.lock +++ b/web/bun.lock @@ -18,6 +18,8 @@ "@visactor/vchart-semi-theme": "~1.8.8", "axios": "catalog:", "clsx": "catalog:", + "date-fns": "2.30.0", + "date-fns-tz": "1.3.8", "dayjs": "catalog:", "highlight.js": "^11.11.1", "history": "^5.3.0", @@ -1512,7 +1514,7 @@ "dagre-d3-es": ["dagre-d3-es@7.0.14", "", { "dependencies": { "d3": "^7.9.0", "lodash-es": "^4.17.21" } }, "sha512-P4rFMVq9ESWqmOgK+dlXvOtLwYg0i7u0HBGJER0LZDJT2VHIPAMZ/riPxqJceWMStH5+E61QxFra9kIS3AqdMg=="], - "date-fns": ["date-fns@4.4.0", "", {}, "sha512-+1UMbeh68lH1SegH83CGWwpb6OHHbpSgr3+s5Eww5M4CAgswBpoWS0AjTOfEJ33HiYKz1hdj/KTFprzXHmq/6w=="], + "date-fns": ["date-fns@2.30.0", "", { "dependencies": { "@babel/runtime": "^7.21.0" } }, "sha512-fnULvOpxnC5/Vg3NCiWelDsLiUc9bRwAPs/+LfTLNvetFCtCTN+yQz15C/fs4AwX1R9K5GLtLfn8QW+dWisaAw=="], "date-fns-tz": ["date-fns-tz@1.3.8", "", { "peerDependencies": { "date-fns": ">=2.0.0" } }, "sha512-qwNXUFtMHTTU6CFSFjoJ80W8Fzzp24LntbjFFBgL/faqds4e5mo9mftoRLgr3Vi1trISsg4awSpYVsOQCRnapQ=="], @@ -2926,6 +2928,8 @@ "zwitch": ["zwitch@2.0.4", "", {}, "sha512-bXE4cR/kVZhKZX/RjPEflHaKVhUVl85noU3v6b8apfQEc1x4A+zBxjZ4lN8LqGd6WZ3dl98pY4o717VFmoPp+A=="], + "@base-ui/react/date-fns": ["date-fns@4.4.0", "", {}, "sha512-+1UMbeh68lH1SegH83CGWwpb6OHHbpSgr3+s5Eww5M4CAgswBpoWS0AjTOfEJ33HiYKz1hdj/KTFprzXHmq/6w=="], + "@codemirror/autocomplete/@codemirror/view": ["@codemirror/view@6.43.3", "", { "dependencies": { "@codemirror/state": "^6.7.0", "crelt": "^1.0.6", "style-mod": "^4.1.0", "w3c-keyname": "^2.2.4" } }, "sha512-MwEwCAr/o0agJefhC2+reBv5kfOQpMcDRUNQrRYZgWlhH8IwQcerMZrpqWyUFSyO0ebgN2cnh/w87F7G4BGSng=="], "@codemirror/lang-html/@codemirror/view": ["@codemirror/view@6.43.3", "", { "dependencies": { "@codemirror/state": "^6.7.0", "crelt": "^1.0.6", "style-mod": "^4.1.0", "w3c-keyname": "^2.2.4" } }, "sha512-MwEwCAr/o0agJefhC2+reBv5kfOQpMcDRUNQrRYZgWlhH8IwQcerMZrpqWyUFSyO0ebgN2cnh/w87F7G4BGSng=="], @@ -2942,10 +2946,6 @@ "@dotenvx/dotenvx/execa": ["execa@5.1.1", "", { "dependencies": { "cross-spawn": "^7.0.3", "get-stream": "^6.0.0", "human-signals": "^2.1.0", "is-stream": "^2.0.0", "merge-stream": "^2.0.0", "npm-run-path": "^4.0.1", "onetime": "^5.1.2", "signal-exit": "^3.0.3", "strip-final-newline": "^2.0.0" } }, "sha512-8uSpZZocAZRBAPIEINJj3Lo9HyGitllczc27Eh5YYojjMFMn8yHMDMaUHE2Jqfq05D/wucwI4JGURyXt1vchyg=="], - "@douyinfe/semi-foundation/date-fns": ["date-fns@2.30.0", "", { "dependencies": { "@babel/runtime": "^7.21.0" } }, "sha512-fnULvOpxnC5/Vg3NCiWelDsLiUc9bRwAPs/+LfTLNvetFCtCTN+yQz15C/fs4AwX1R9K5GLtLfn8QW+dWisaAw=="], - - "@douyinfe/semi-ui/date-fns": ["date-fns@2.30.0", "", { "dependencies": { "@babel/runtime": "^7.21.0" } }, "sha512-fnULvOpxnC5/Vg3NCiWelDsLiUc9bRwAPs/+LfTLNvetFCtCTN+yQz15C/fs4AwX1R9K5GLtLfn8QW+dWisaAw=="], - "@emoji-mart/react/react": ["react@18.3.1", "", { "dependencies": { "loose-envify": "^1.1.0" } }, "sha512-wS+hAgJShR0KhEvPJArfuPVN1+Hz1t0Y6n5jLrGQbkb4urgPE/0Rve+1kMB1v/oWgHgm4WIcV+i7F2pTVj+2iQ=="], "@emotion/babel-plugin/@emotion/hash": ["@emotion/hash@0.9.2", "", {}, "sha512-MyqliTZGuOm3+5ZRSaaBGP3USLw6+EGykkwZns2EPC5g8jJ4z9OrdZY9apkl3+UP9+sdz76YYkwCKP5gh8iY3g=="], @@ -3202,6 +3202,8 @@ "rc-menu/@rc-component/trigger": ["@rc-component/trigger@2.3.1", "", { "dependencies": { "@babel/runtime": "^7.23.2", "@rc-component/portal": "^1.1.0", "classnames": "^2.3.2", "rc-motion": "^2.0.0", "rc-resize-observer": "^1.3.1", "rc-util": "^5.44.0" }, "peerDependencies": { "react": ">=16.9.0", "react-dom": ">=16.9.0" } }, "sha512-ORENF39PeXTzM+gQEshuk460Z8N4+6DkjpxlpE7Q3gYy1iBpLrx0FOJz3h62ryrJZ/3zCAUIkT1Pb/8hHWpb3A=="], + "react-day-picker/date-fns": ["date-fns@4.4.0", "", {}, "sha512-+1UMbeh68lH1SegH83CGWwpb6OHHbpSgr3+s5Eww5M4CAgswBpoWS0AjTOfEJ33HiYKz1hdj/KTFprzXHmq/6w=="], + "react-i18next/typescript": ["typescript@6.0.3", "", { "bin": { "tsc": "bin/tsc", "tsserver": "bin/tsserver" } }, "sha512-y2TvuxSZPDyQakkFRPZHKFm+KKVqIisdg9/CZwm9ftvKXLP8NRWj38/ODjNbr43SsoXqNuAisEf1GdCxqWcdBw=="], "react-rnd/tslib": ["tslib@2.6.2", "", {}, "sha512-AEYxH93jGFPn/a2iVAwW87VuUIkR1FVUKB77NwMF7nBTDkDrrT/Hpt/IrCJ0QXhW27jTBDcf5ZY7w6RiqTMw2Q=="], @@ -3276,6 +3278,8 @@ "@lobehub/ui/@base-ui/react/@base-ui/utils": ["@base-ui/utils@0.2.9", "", { "dependencies": { "@babel/runtime": "^7.29.2", "@floating-ui/utils": "^0.2.11", "reselect": "^5.1.1", "use-sync-external-store": "^1.6.0" }, "peerDependencies": { "@types/react": "^17 || ^18 || ^19", "react": "^17 || ^18 || ^19", "react-dom": "^17 || ^18 || ^19" }, "optionalPeers": ["@types/react"] }, "sha512-x/PDDCYzoqPpjrdyb3VcyylTI2IjUXEtYDGi5foh7KsnmNJIIaVwA2GLgDH1dps1GgXiJbA60hM+AyuTfQzIvw=="], + "@lobehub/ui/@base-ui/react/date-fns": ["date-fns@4.4.0", "", {}, "sha512-+1UMbeh68lH1SegH83CGWwpb6OHHbpSgr3+s5Eww5M4CAgswBpoWS0AjTOfEJ33HiYKz1hdj/KTFprzXHmq/6w=="], + "@lobehub/ui/@shikijs/core/@shikijs/primitive": ["@shikijs/primitive@4.1.0", "", { "dependencies": { "@shikijs/types": "4.1.0", "@shikijs/vscode-textmate": "^10.0.2", "@types/hast": "^3.0.4" } }, "sha512-zx2/2Uwj2q9X3KSyYREEhXO23xBw5WUhP4orK2lE4r+t9JGITmEe0JH+wPmJhqHpOT2bRRs6lAL945+LDvOAGw=="], "@lobehub/ui/@shikijs/core/@shikijs/types": ["@shikijs/types@4.1.0", "", { "dependencies": { "@shikijs/vscode-textmate": "^10.0.2", "@types/hast": "^3.0.4" } }, "sha512-3EQWX54fMpniOrDblzAhiwiJwpiTMW6+B9DWyUd9ska483tbayFYuw47UxwuPknI31bKnySfVQ/QW+jFL4rFdA=="], diff --git a/web/classic/package.json b/web/classic/package.json index ecc6971280b..c101bf030e7 100644 --- a/web/classic/package.json +++ b/web/classic/package.json @@ -13,6 +13,8 @@ "@visactor/vchart-semi-theme": "~1.8.8", "axios": "catalog:", "clsx": "catalog:", + "date-fns": "2.30.0", + "date-fns-tz": "1.3.8", "dayjs": "catalog:", "history": "^5.3.0", "highlight.js": "^11.11.1", diff --git a/web/classic/rsbuild.config.ts b/web/classic/rsbuild.config.ts index 3ccf1e96df8..4421917dfc1 100644 --- a/web/classic/rsbuild.config.ts +++ b/web/classic/rsbuild.config.ts @@ -10,6 +10,7 @@ const semiUiDir = path.resolve( path.dirname(require.resolve('@douyinfe/semi-ui')), '../..', ) +const dateFnsDir = path.dirname(require.resolve('date-fns/package.json')) export default defineConfig(({ envMode }) => { const env = loadEnv({ mode: envMode, prefixes: ['VITE_'] }) @@ -43,6 +44,7 @@ export default defineConfig(({ envMode }) => { resolve: { alias: { '@': path.resolve(__dirname, './src'), + 'date-fns': dateFnsDir, '@douyinfe/semi-ui/dist/css/semi.css': path.resolve( semiUiDir, 'dist/css/semi.css',