From b80f8cb377faf3731ce7bf3dcdbc7c0f2a5a5a77 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=B8=81=E5=B8=85=E5=B8=85=28=E4=B8=81=E5=B8=85=E5=B8=85?= =?UTF-8?q?=29?= Date: Mon, 10 Aug 2026 08:07:34 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20OriginFlow=20backend=20=E2=80=94=20regi?= =?UTF-8?q?on=20routing,=20Model=20Market=20&=20sales-leads=20APIs,=20orig?= =?UTF-8?q?inflow=20infra?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Region-aware channel routing (model/region_route.go, channel_cache, ability, service/channel_select). - Model Market backend: product table, admin CRUD, public read API, i18n/currency. - Sales-leads backend + analytics events model. - Originflow infra: common/redirect.go (IsSafeRedirect), model/analytics_event.go, deploy/nginx/originflow.conf, OriginFlow SystemName branding. - Build/Docker: build from source, configurable npm registry; nginx serves www/app/api hosts. - Note: upstream test files intentionally retained (not deleted) vs the merged branch. --- .dockerignore | 12 +- Dockerfile | 18 +- common/body_storage.go | 70 +--- common/constants.go | 2 +- common/redirect.go | 28 ++ constant/context_key.go | 20 +- controller/distributor.go | 302 ++++++++++++++++++ controller/market.go | 279 ++++++++++++++++ controller/model.go | 45 ++- controller/model_list_test.go | 70 +--- controller/model_owned_by_test.go | 33 -- controller/public_site.go | 134 ++++++++ controller/region_route.go | 150 +++++++++ controller/sla.go | 163 ++++++++++ controller/team.go | 258 +++++++++++++++ controller/token.go | 132 +------- controller/token_test.go | 40 --- controller/user.go | 13 +- deploy/nginx/originflow.conf | 156 +++++++++ docker-compose.yml | 19 +- i18n/keys.go | 3 - i18n/locales/en.yaml | 3 - i18n/locales/zh-CN.yaml | 3 - i18n/locales/zh-TW.yaml | 3 - middleware/auth.go | 10 - middleware/distributor.go | 2 +- model/ability.go | 72 ++++- model/analytics_event.go | 38 +++ model/channel_cache.go | 47 ++- model/distributor.go | 194 ++++++++++++ model/log.go | 8 + model/main.go | 15 + model/market_model.go | 119 +++++++ model/option.go | 6 - model/public_model_category.go | 15 + model/public_pricing.go | 17 + model/public_site_seed.go | 106 +++++++ model/region_route.go | 281 +++++++++++++++++ model/region_route_test.go | 157 +++++++++ model/sales_lead.go | 20 ++ model/sla.go | 180 +++++++++++ model/team.go | 201 ++++++++++++ model/token.go | 45 +-- model/user.go | 90 +++--- model/user_cache.go | 26 +- model/user_update_test.go | 83 +---- relay/alpha_search_handler.go | 3 +- relay/channel/api_request.go | 62 ++-- relay/channel/aws/relay-aws.go | 78 ++--- relay/channel/aws/relay_aws_test.go | 403 ------------------------ relay/channel/jimeng/adaptor.go | 1 - relay/channel/task/sora/adaptor.go | 2 +- relay/chat_completions_via_responses.go | 3 +- relay/claude_handler.go | 6 +- relay/common/outbound_body.go | 14 +- relay/common/relay_info.go | 7 + relay/compatible_handler.go | 5 +- relay/embedding_handler.go | 3 +- relay/gemini_handler.go | 8 +- relay/image_handler.go | 5 +- relay/rerank_handler.go | 5 +- relay/responses_handler.go | 5 +- router/api-router.go | 25 +- router/p2-router.go | 74 +++++ service/billing_usage.go | 26 +- service/channel_select.go | 34 +- service/channel_select_test.go | 60 ++++ service/group.go | 58 +--- service/text_quota_test.go | 15 +- setting/auto_group.go | 33 -- 70 files changed, 3412 insertions(+), 1211 deletions(-) create mode 100644 common/redirect.go create mode 100644 controller/distributor.go create mode 100644 controller/market.go create mode 100644 controller/public_site.go create mode 100644 controller/region_route.go create mode 100644 controller/sla.go create mode 100644 controller/team.go create mode 100644 deploy/nginx/originflow.conf create mode 100644 model/analytics_event.go create mode 100644 model/distributor.go create mode 100644 model/market_model.go create mode 100644 model/public_model_category.go create mode 100644 model/public_pricing.go create mode 100644 model/public_site_seed.go create mode 100644 model/region_route.go create mode 100644 model/region_route_test.go create mode 100644 model/sales_lead.go create mode 100644 model/sla.go create mode 100644 model/team.go create mode 100644 router/p2-router.go create mode 100644 service/channel_select_test.go diff --git a/.dockerignore b/.dockerignore index 0204d2e89809..2d2d7ed788ad 100644 --- a/.dockerignore +++ b/.dockerignore @@ -7,6 +7,16 @@ Makefile docs .eslintcache .gocache -/web/node_modules +# 依赖一律在镜像内 bun install，任何层级的 node_modules 都不需要进构建上下文 +# (Dependencies are installed inside the image; no node_modules at any depth.) +**/node_modules +# web/dist 由 builder 阶段产出并 COPY --from，构建上下文里的产物一律忽略 +# (web/dist is produced by the builder stage; ignore any prebuilt output here.) /web/dist +/web/classic/dist +/web/default/dist +# 运行时目录，由 compose 以 bind mount 挂载 +# (Runtime directories, provided by compose bind mounts.) +/data +/logs !THIRD-PARTY-LICENSES.md diff --git a/Dockerfile b/Dockerfile index 5be311b10e43..037e218395fa 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,8 +1,22 @@ FROM oven/bun:1@sha256:0733e50325078969732ebe3b15ce4c4be5082f18c4ac1a0f0ca4839c2e4e42a7 AS builder WORKDIR /build/web -COPY web/package.json web/bun.lock ./ -RUN bun install --frozen-lockfile +# .npmrc must land before `bun install`, otherwise the registry it configures is +# ignored (it would only arrive with the later `COPY ./web ./`). +COPY web/package.json web/bun.lock web/.npmrc ./ +# Point at a nearby mirror when the default registry is slow, e.g. +# docker build --build-arg NPM_REGISTRY=https://repo.huaweicloud.com/repository/npm/ . +ARG NPM_REGISTRY= +# Bun's default fetch concurrency corrupts large tarballs (react-icons is 22MB, +# @hugeicons/core-free-icons larger still) on constrained links, surfacing as +# "Integrity check failed" on a different package each run. Cap concurrency and +# retry so a single dropped connection doesn't fail the whole build. +RUN if [ -n "$NPM_REGISTRY" ]; then echo "registry=$NPM_REGISTRY" > .npmrc; fi \ + && for attempt in 1 2 3; do \ + echo "bun install attempt $attempt" \ + && BUN_CONFIG_MAX_HTTP_REQUESTS=8 bun install --frozen-lockfile && break \ + || { [ "$attempt" = 3 ] && exit 1; echo "retrying..."; sleep 5; }; \ + done COPY ./web ./ COPY ./VERSION /build/VERSION RUN DISABLE_ESLINT_PLUGIN='true' VITE_REACT_APP_VERSION=$(cat /build/VERSION) bun run build diff --git a/common/body_storage.go b/common/body_storage.go index 515d749e10c9..094dbda36d3f 100644 --- a/common/body_storage.go +++ b/common/body_storage.go @@ -20,21 +20,6 @@ type BodyStorage interface { Size() int64 // IsDisk 是否是磁盘存储 IsDisk() bool - // NewReader returns an independent reader positioned at the start of the - // stored payload. Each call returns a reader with its own cursor, so - // callers (e.g. http.Request.GetBody) can replay the body concurrently - // with, or after, other readers without sharing seek state. Closing the - // returned reader releases only that reader, never the storage itself; - // after the storage has been closed, NewReader returns ErrStorageClosed. - NewReader() (io.ReadCloser, error) -} - -// ReplayableBody is an outbound request body that can report its byte size and -// create independent readers for transport-level retries. -type ReplayableBody interface { - io.Reader - Size() int64 - NewReader() (io.ReadCloser, error) } // ErrStorageClosed 存储已关闭错误 @@ -95,18 +80,6 @@ func (m *memoryStorage) Bytes() ([]byte, error) { return m.data, nil } -func (m *memoryStorage) NewReader() (io.ReadCloser, error) { - m.mu.Lock() - defer m.mu.Unlock() - if atomic.LoadInt32(&m.closed) == 1 { - return nil, ErrStorageClosed - } - // A fresh bytes.Reader over the shared immutable backing array: an - // independent cursor at zero copy cost. NopCloser keeps Close a no-op, so - // the storage lifecycle stays owned by whoever holds the storage itself. - return io.NopCloser(bytes.NewReader(m.data)), nil -} - func (m *memoryStorage) Size() int64 { return m.size } @@ -256,24 +229,6 @@ func (d *diskStorage) Bytes() ([]byte, error) { return data, nil } -func (d *diskStorage) NewReader() (io.ReadCloser, error) { - d.mu.Lock() - defer d.mu.Unlock() - if atomic.LoadInt32(&d.closed) == 1 { - return nil, ErrStorageClosed - } - // A separate file descriptor over the same cache file: an independent - // cursor at zero copy cost. Closing the returned reader closes only that - // descriptor; the storage keeps owning the primary descriptor and the - // file's lifetime. Readers opened before Close stay usable even after the - // file is unlinked, as the descriptor keeps the inode alive. - file, err := os.Open(d.filePath) - if err != nil { - return nil, fmt.Errorf("failed to open body cache file for replay: %w", err) - } - return file, nil -} - func (d *diskStorage) Size() int64 { return d.size } @@ -347,27 +302,10 @@ func CreateBodyStorageFromReader(reader io.Reader, contentLength int64, maxBytes return storage, nil } -type replayableBodyReader struct { - storage BodyStorage -} - -func (r replayableBodyReader) Read(p []byte) (int, error) { - return r.storage.Read(p) -} - -func (r replayableBodyReader) Size() int64 { - return r.storage.Size() -} - -func (r replayableBodyReader) NewReader() (io.ReadCloser, error) { - return r.storage.NewReader() -} - -// NewReplayableBodyReader exposes the replay capabilities of storage without -// exposing io.Closer. This keeps ownership of the storage lifecycle with the -// caller instead of allowing net/http to close it as the request body. -func NewReplayableBodyReader(storage BodyStorage) ReplayableBody { - return replayableBodyReader{storage: storage} +// ReaderOnly wraps an io.Reader to hide io.Closer, preventing http.NewRequest +// from type-asserting io.ReadCloser and closing the underlying BodyStorage. +func ReaderOnly(r io.Reader) io.Reader { + return struct{ io.Reader }{r} } // CleanupOldCacheFiles 清理旧的缓存文件（用于启动时清理残留） diff --git a/common/constants.go b/common/constants.go index d6b4fb52284c..5e8063419ff0 100644 --- a/common/constants.go +++ b/common/constants.go @@ -12,7 +12,7 @@ import ( var StartTime = time.Now().Unix() // unit: second var Version = "v0.0.0" // this hard coding will be replaced automatically when building, no need to manually change -var SystemName = "New API" +var SystemName = "OriginFlow" var Footer = "" var Logo = "" var TopUpLink = "" diff --git a/common/redirect.go b/common/redirect.go new file mode 100644 index 000000000000..6e14657dbdb5 --- /dev/null +++ b/common/redirect.go @@ -0,0 +1,28 @@ +package common + +import ( + "net/url" + "strings" +) + +// IsSafeRedirect 校验重定向目标是否安全，防止开放重定向（Open Redirect）。 +// 允许：空串（无重定向）、同源相对路径（以 / 开头且非协议相对 //host）、指向已知官方域名。 +func IsSafeRedirect(u string) bool { + u = strings.TrimSpace(u) + if u == "" { + return true + } + // 相对路径，但排除协议相对地址（//evil.com） + if strings.HasPrefix(u, "/") { + return !strings.HasPrefix(u, "//") + } + parsed, err := url.Parse(u) + if err != nil { + return false + } + if parsed.Scheme != "http" && parsed.Scheme != "https" { + return false + } + host := parsed.Host + return host == "91flow.com" || strings.HasSuffix(host, ".91flow.com") +} diff --git a/constant/context_key.go b/constant/context_key.go index ccb8010f9476..0a18ef2bb3c6 100644 --- a/constant/context_key.go +++ b/constant/context_key.go @@ -19,7 +19,6 @@ const ( ContextKeyTokenModelLimitEnabled ContextKey = "token_model_limit_enabled" ContextKeyTokenModelLimit ContextKey = "token_model_limit" ContextKeyTokenCrossGroupRetry ContextKey = "token_cross_group_retry" - ContextKeyTokenAutoGroups ContextKey = "token_auto_groups" /* channel related keys */ ContextKeyChannelId ContextKey = "channel_id" @@ -73,4 +72,23 @@ const ( // fallback in authHelper (finishAdminAudit) skips its record to avoid // duplicate entries. ContextKeyAuditLogged ContextKey = "audit_logged" + + // ContextKeyRequestRegion stores the normalized region of the current request, + // used by region-based channel routing (see model.ResolveRegionRouting). + // It is distinct from the vertex-specific "region" key set in the distributor. + ContextKeyRequestRegion ContextKey = "request_region" + + // ContextKeyUserRegionPreference stores the authenticated user's preferred + // region (see User.RegionPreference). Region-based channel routing falls back + // to it when no X-Region header or request-region context is present. + ContextKeyUserRegionPreference ContextKey = "user_region_preference" + + // ContextKeyUserTeamId stores the authenticated user's enterprise team id + // (see User.TeamId). Consume logs are stamped with it for per-team billing + // aggregation. + ContextKeyUserTeamId ContextKey = "user_team_id" ) + +// HeaderRegion is the request header declaring the caller's region for +// region-based channel routing, e.g. `X-Region: cn`. +const HeaderRegion = "X-Region" diff --git a/controller/distributor.go b/controller/distributor.go new file mode 100644 index 000000000000..dda3cb88c272 --- /dev/null +++ b/controller/distributor.go @@ -0,0 +1,302 @@ +package controller + +import ( + "strconv" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/model" + "github.com/gin-gonic/gin" +) + +// ListDistributors 分页列出分销商。 +func ListDistributors(c *gin.Context) { + page, pageSize := parsePage(c) + keyword := c.Query("keyword") + items, total, err := model.SearchDistributors(page, pageSize, keyword) + if err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, gin.H{"items": items, "total": total}) +} + +// CreateDistributor 创建分销商。 +func CreateDistributor(c *gin.Context) { + var req struct { + UserId int64 `json:"user_id"` + Name string `json:"name"` + Tier string `json:"tier"` + CommissionRate int `json:"commission_rate"` + Status int `json:"status"` + } + if err := common.DecodeJson(c.Request.Body, &req); err != nil { + common.ApiError(c, err) + return + } + if req.UserId <= 0 { + common.ApiErrorMsg(c, "user_id is required") + return + } + if req.Name == "" || len(req.Name) > 64 { + common.ApiErrorMsg(c, "name is required and must be <= 64 characters") + return + } + tier := req.Tier + if tier == "" { + tier = "standard" + } + if !model.AllowedDistributorTiers[tier] { + common.ApiErrorMsg(c, "invalid tier (expected standard|gold|platinum)") + return + } + status := req.Status + if status == 0 { + status = model.DistributorStatusActive + } + if !model.AllowedDistributorStatuses[status] { + common.ApiErrorMsg(c, "invalid status (expected 1 active|2 disabled)") + return + } + m := &model.Distributor{ + UserId: req.UserId, + Name: req.Name, + Tier: tier, + CommissionRate: req.CommissionRate, + Status: status, + } + if err := model.CreateDistributor(m); err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, gin.H{"id": m.Id}) +} + +// GetDistributor 获取单个分销商。 +func GetDistributor(c *gin.Context) { + id, err := strconv.ParseInt(c.Param("id"), 10, 64) + if err != nil { + common.ApiErrorMsg(c, "invalid id") + return + } + m, err := model.GetDistributorById(id) + if err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, m) +} + +// UpdateDistributor 更新分销商。 +func UpdateDistributor(c *gin.Context) { + id, err := strconv.ParseInt(c.Param("id"), 10, 64) + if err != nil { + common.ApiErrorMsg(c, "invalid id") + return + } + var req struct { + Name string `json:"name"` + Tier string `json:"tier"` + CommissionRate int `json:"commission_rate"` + Status int `json:"status"` + } + if err := common.DecodeJson(c.Request.Body, &req); err != nil { + common.ApiError(c, err) + return + } + if !model.AllowedDistributorTiers[req.Tier] { + common.ApiErrorMsg(c, "invalid tier (expected standard|gold|platinum)") + return + } + if !model.AllowedDistributorStatuses[req.Status] { + common.ApiErrorMsg(c, "invalid status (expected 1 active|2 disabled)") + return + } + m := &model.Distributor{Id: id, Name: req.Name, Tier: req.Tier, CommissionRate: req.CommissionRate, Status: req.Status} + if err := model.UpdateDistributor(m); err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, gin.H{"ok": true}) +} + +// DeleteDistributor 删除分销商（级联价格覆盖）。 +func DeleteDistributor(c *gin.Context) { + id, err := strconv.ParseInt(c.Param("id"), 10, 64) + if err != nil { + common.ApiErrorMsg(c, "invalid id") + return + } + if err := model.DeleteDistributor(id); err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, gin.H{"ok": true}) +} + +// ListDistributorSubUsers 列出分销商下级用户（基于邀请链）。 +func ListDistributorSubUsers(c *gin.Context) { + id, err := strconv.ParseInt(c.Param("id"), 10, 64) + if err != nil { + common.ApiErrorMsg(c, "invalid id") + return + } + d, err := model.GetDistributorById(id) + if err != nil { + common.ApiError(c, err) + return + } + page, pageSize := parsePage(c) + items, total, err := model.GetUsersByInviterId(int(d.UserId), page, pageSize) + if err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, gin.H{"items": items, "total": total}) +} + +// ListDistributorPrices 列出分销商价格覆盖。 +func ListDistributorPrices(c *gin.Context) { + id, err := strconv.ParseInt(c.Param("id"), 10, 64) + if err != nil { + common.ApiErrorMsg(c, "invalid id") + return + } + items, err := model.SearchDistributorPrices(id) + if err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, gin.H{"items": items}) +} + +// CreateDistributorPrice 创建价格覆盖。 +func CreateDistributorPrice(c *gin.Context) { + id, err := strconv.ParseInt(c.Param("id"), 10, 64) + if err != nil { + common.ApiErrorMsg(c, "invalid id") + return + } + var req struct { + Model string `json:"model"` + InputPrice int64 `json:"input_price"` + OutputPrice int64 `json:"output_price"` + Currency string `json:"currency"` + Unit string `json:"unit"` + } + if err := common.DecodeJson(c.Request.Body, &req); err != nil { + common.ApiError(c, err) + return + } + if req.Model == "" { + common.ApiErrorMsg(c, "model is required") + return + } + if req.InputPrice < 0 || req.OutputPrice < 0 { + common.ApiErrorMsg(c, "prices must be >= 0") + return + } + currency := req.Currency + if currency == "" { + currency = "CNY" + } + if !model.AllowedDistributorPriceCurrencies[currency] { + common.ApiErrorMsg(c, "invalid currency (expected CNY|USD)") + return + } + unit := req.Unit + if unit == "" { + unit = "token" + } + if !model.AllowedDistributorPriceUnits[unit] { + common.ApiErrorMsg(c, "invalid unit (expected token|image|second|char)") + return + } + m := &model.DistributorPrice{ + DistributorId: id, + Model: req.Model, + InputPrice: req.InputPrice, + OutputPrice: req.OutputPrice, + Currency: currency, + Unit: unit, + } + if err := model.CreateDistributorPrice(m); err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, gin.H{"id": m.Id}) +} + +// UpdateDistributorPrice 更新价格覆盖。 +func UpdateDistributorPrice(c *gin.Context) { + priceId, err := strconv.ParseInt(c.Param("price_id"), 10, 64) + if err != nil { + common.ApiErrorMsg(c, "invalid price_id") + return + } + var req struct { + Model string `json:"model"` + InputPrice int64 `json:"input_price"` + OutputPrice int64 `json:"output_price"` + Currency string `json:"currency"` + Unit string `json:"unit"` + } + if err := common.DecodeJson(c.Request.Body, &req); err != nil { + common.ApiError(c, err) + return + } + if !model.AllowedDistributorPriceCurrencies[req.Currency] { + common.ApiErrorMsg(c, "invalid currency (expected CNY|USD)") + return + } + if !model.AllowedDistributorPriceUnits[req.Unit] { + common.ApiErrorMsg(c, "invalid unit (expected token|image|second|char)") + return + } + m := &model.DistributorPrice{ + Id: priceId, + Model: req.Model, + InputPrice: req.InputPrice, + OutputPrice: req.OutputPrice, + Currency: req.Currency, + Unit: req.Unit, + } + if err := model.UpdateDistributorPrice(m); err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, gin.H{"ok": true}) +} + +// DeleteDistributorPrice 删除价格覆盖。 +func DeleteDistributorPrice(c *gin.Context) { + priceId, err := strconv.ParseInt(c.Param("price_id"), 10, 64) + if err != nil { + common.ApiErrorMsg(c, "invalid price_id") + return + } + if err := model.DeleteDistributorPrice(priceId); err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, gin.H{"ok": true}) +} + +// GetDistributorBilling 分销商下级账单汇总。 +func GetDistributorBilling(c *gin.Context) { + id, err := strconv.ParseInt(c.Param("id"), 10, 64) + if err != nil { + common.ApiErrorMsg(c, "invalid id") + return + } + d, err := model.GetDistributorById(id) + if err != nil { + common.ApiError(c, err) + return + } + billing, err := model.GetDistributorBilling(d.UserId) + if err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, billing) +} diff --git a/controller/market.go b/controller/market.go new file mode 100644 index 000000000000..0a199365925d --- /dev/null +++ b/controller/market.go @@ -0,0 +1,279 @@ +package controller + +import ( + "encoding/json" + "strconv" + "strings" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/model" + "github.com/gin-gonic/gin" +) + +// MarketModel statuses (mirror model.AllowedMarketModelStatuses). +const ( + MarketModelStatusAvailable = 1 + MarketModelStatusComingSoon = 2 + MarketModelStatusDisabled = 3 +) + +var allowedMarketModelStatuses = map[int]bool{ + MarketModelStatusAvailable: true, + MarketModelStatusComingSoon: true, + MarketModelStatusDisabled: true, +} + +var allowedMarketModelUnits = map[string]bool{ + "token": true, + "image": true, + "second": true, + "char": true, +} + +var allowedMarketModelCurrencies = map[string]bool{ + "CNY": true, + "USD": true, +} + +// MarketModelRequest 上架模型的新建/更新请求体。 +type MarketModelRequest struct { + Model string `json:"model"` + Provider string `json:"provider"` + Category string `json:"category"` + Tags string `json:"tags"` + InputPrice int64 `json:"input_price"` + OutputPrice int64 `json:"output_price"` + Currency string `json:"currency"` + Unit string `json:"unit"` + Metadata string `json:"metadata"` + TrialQuota int64 `json:"trial_quota"` + Status int `json:"status"` + Featured bool `json:"featured"` + Sort int `json:"sort"` +} + +func validateMarketModelRequest(req *MarketModelRequest, isCreate bool) (bool, string) { + if isCreate { + if req.Model == "" || len(req.Model) > 255 { + return false, "model is required and must be <= 255 characters" + } + } + if req.Category == "" || len(req.Category) > 32 { + return false, "category is required and must be <= 32 characters" + } + if len(req.Provider) > 64 { + return false, "provider must be <= 64 characters" + } + if len(req.Tags) > 255 { + return false, "tags must be <= 255 characters" + } + if req.InputPrice < 0 || req.OutputPrice < 0 { + return false, "prices must be >= 0" + } + if req.TrialQuota < 0 { + return false, "trial_quota must be >= 0" + } + unit := req.Unit + if unit == "" { + unit = "token" + } + if !allowedMarketModelUnits[unit] { + return false, "invalid unit (expected token|image|second|char)" + } + currency := req.Currency + if currency == "" { + currency = "CNY" + } + if !allowedMarketModelCurrencies[currency] { + return false, "invalid currency (expected CNY|USD)" + } + if req.Metadata != "" { + var tmp interface{} + if err := json.Unmarshal([]byte(req.Metadata), &tmp); err != nil { + return false, "metadata must be valid JSON" + } + } + if !allowedMarketModelStatuses[req.Status] { + return false, "invalid status (expected 1 available|2 coming_soon|3 disabled)" + } + return true, "" +} + +// CreateMarketModel 管理员新建上架模型。 +func CreateMarketModel(c *gin.Context) { + var req MarketModelRequest + if err := common.DecodeJson(c.Request.Body, &req); err != nil { + common.ApiError(c, err) + return + } + if ok, msg := validateMarketModelRequest(&req, true); !ok { + common.ApiErrorMsg(c, msg) + return + } + // 唯一性校验 + if cnt, err := model.CountMarketModelsByModel(req.Model); err != nil { + common.ApiError(c, err) + return + } else if cnt > 0 { + common.ApiErrorMsg(c, "model already listed in market") + return + } + m := &model.MarketModel{ + Model: req.Model, + Provider: req.Provider, + Category: req.Category, + Tags: req.Tags, + InputPrice: req.InputPrice, + OutputPrice: req.OutputPrice, + Unit: orDefault(req.Unit, "token"), + TrialQuota: req.TrialQuota, + Status: req.Status, + Featured: req.Featured, + Sort: req.Sort, + } + if m.Status == 0 { + m.Status = MarketModelStatusAvailable + } + if m.Currency == "" { + m.Currency = "CNY" + } + if err := model.CreateMarketModel(m); err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, gin.H{"id": m.Id}) +} + +// GetMarketModel 管理员查看单条上架记录。 +func GetMarketModel(c *gin.Context) { + id, err := strconv.ParseInt(c.Param("id"), 10, 64) + if err != nil { + common.ApiErrorMsg(c, "invalid id") + return + } + m, err := model.GetMarketModelById(id) + if err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, m) +} + +// ListMarketModels 管理员列出上架模型，可选 status/category 过滤。 +func ListMarketModels(c *gin.Context) { + status := c.Query("status") + category := c.Query("category") + items, err := model.SearchMarketModels(status, category) + if err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, gin.H{"items": items}) +} + +// UpdateMarketModel 管理员更新上架记录（Model 唯一键不可变）。 +func UpdateMarketModel(c *gin.Context) { + id, err := strconv.ParseInt(c.Param("id"), 10, 64) + if err != nil { + common.ApiErrorMsg(c, "invalid id") + return + } + var req MarketModelRequest + if err := common.DecodeJson(c.Request.Body, &req); err != nil { + common.ApiError(c, err) + return + } + if ok, msg := validateMarketModelRequest(&req, false); !ok { + common.ApiErrorMsg(c, msg) + return + } + m := &model.MarketModel{ + Id: id, + Provider: req.Provider, + Category: req.Category, + Tags: req.Tags, + InputPrice: req.InputPrice, + OutputPrice: req.OutputPrice, + Currency: orDefault(req.Currency, "CNY"), + Unit: orDefault(req.Unit, "token"), + Metadata: req.Metadata, + TrialQuota: req.TrialQuota, + Status: req.Status, + Featured: req.Featured, + Sort: req.Sort, + } + if err := model.UpdateMarketModel(m); err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, gin.H{"ok": true}) +} + +// DeleteMarketModel 管理员删除上架记录。 +func DeleteMarketModel(c *gin.Context) { + id, err := strconv.ParseInt(c.Param("id"), 10, 64) + if err != nil { + common.ApiErrorMsg(c, "invalid id") + return + } + if err := model.DeleteMarketModel(id); err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, gin.H{"ok": true}) +} + +// GetPublicMarketModels 门店公开读取：仅返回已上架（status=available）模型，按分类与排序。 +// 支持 ?locale=zh|en：从 Metadata 解析对应 locale 的展示覆盖（name/description）叠加到 i18n 字段。 +func GetPublicMarketModels(c *gin.Context) { + items, err := model.SearchMarketModels(strconv.Itoa(MarketModelStatusAvailable), "") + if err != nil { + common.ApiError(c, err) + return + } + locale := orDefault(c.Query("locale"), "zh") + result := make([]gin.H, 0, len(items)) + for _, m := range items { + result = append(result, gin.H{ + "model": m, + "i18n": resolveMarketModelI18n(m, locale), + }) + } + common.ApiSuccess(c, gin.H{"items": result}) +} + +// resolveMarketModelI18n 从 Metadata JSON 解析指定 locale 的展示覆盖（name/description）。 +// 无 Metadata 或无对应 locale 时返回 nil。 +func resolveMarketModelI18n(m *model.MarketModel, locale string) map[string]string { + if strings.TrimSpace(m.Metadata) == "" { + return nil + } + var data map[string]map[string]string + if err := json.Unmarshal([]byte(m.Metadata), &data); err != nil { + return nil + } + entry, ok := data[locale] + if !ok { + return nil + } + out := map[string]string{} + if v, ok := entry["name"]; ok && v != "" { + out["name"] = v + } + if v, ok := entry["description"]; ok && v != "" { + out["description"] = v + } + if len(out) == 0 { + return nil + } + out["locale"] = locale + return out +} + +// orDefault 返回非空值或默认值。 +func orDefault(v, def string) string { + if strings.TrimSpace(v) == "" { + return def + } + return v +} diff --git a/controller/model.go b/controller/model.go index 1d759301bc7e..b32eebd7daac 100644 --- a/controller/model.go +++ b/controller/model.go @@ -20,7 +20,6 @@ import ( "github.com/QuantumNous/new-api/relaykit/types" "github.com/QuantumNous/new-api/service" "github.com/QuantumNous/new-api/setting/operation_setting" - "github.com/QuantumNous/new-api/setting/ratio_setting" "github.com/gin-gonic/gin" "github.com/samber/lo" ) @@ -191,7 +190,7 @@ func getModelListGroups(c *gin.Context) (modelListGroups, error) { return modelListGroups{ userGroup: userGroup, tokenGroup: tokenGroup, - ownerGroups: service.GetRequestAutoGroups(c, userGroup), + ownerGroups: service.GetUserAutoGroup(userGroup), }, nil } @@ -229,28 +228,32 @@ func ListModels(c *gin.Context, modelType int) { } ownerGroups := groups.ownerGroups modelLimitEnable := common.GetContextKeyBool(c, constant.ContextKeyTokenModelLimitEnabled) - var tokenModelLimit map[string]bool if modelLimitEnable { s, ok := common.GetContextKey(c, constant.ContextKeyTokenModelLimit) + var tokenModelLimit map[string]bool if ok { - tokenModelLimit, _ = s.(map[string]bool) - } - if tokenModelLimit == nil { + tokenModelLimit = s.(map[string]bool) + } else { tokenModelLimit = map[string]bool{} } - } - models := service.GetGroupsEnabledModels(ownerGroups) - for _, modelName := range models { - if modelLimitEnable { - matchingName := ratio_setting.FormatMatchingModelName(modelName) - if !tokenModelLimit[modelName] && !tokenModelLimit[matchingName] { - continue + for allowModel, _ := range tokenModelLimit { + if !acceptUnsetRatioModel { + if !helper.HasModelBillingConfig(allowModel) { + continue + } } + userModelNames = append(userModelNames, allowModel) } - if !acceptUnsetRatioModel && !helper.HasModelBillingConfig(modelName) { - continue + } else { + models := service.GetGroupsEnabledModels(ownerGroups) + for _, modelName := range models { + if !acceptUnsetRatioModel { + if !helper.HasModelBillingConfig(modelName) { + continue + } + } + userModelNames = append(userModelNames, modelName) } - userModelNames = append(userModelNames, modelName) } ownerByModel := map[string]string{} @@ -273,17 +276,11 @@ func ListModels(c *gin.Context, modelType int) { Type: "model", } } - firstID := "" - lastID := "" - if len(useranthropicModels) > 0 { - firstID = useranthropicModels[0].ID - lastID = useranthropicModels[len(useranthropicModels)-1].ID - } c.JSON(200, gin.H{ "data": useranthropicModels, - "first_id": firstID, + "first_id": useranthropicModels[0].ID, "has_more": false, - "last_id": lastID, + "last_id": useranthropicModels[len(useranthropicModels)-1].ID, }) case constant.ChannelTypeGemini: userGeminiModels := make([]dto.GeminiModel, len(userOpenAiModels)) diff --git a/controller/model_list_test.go b/controller/model_list_test.go index 812207b8fd44..b1fa9b956ce3 100644 --- a/controller/model_list_test.go +++ b/controller/model_list_test.go @@ -402,13 +402,7 @@ func TestListModelsTokenLimitIncludesTieredBillingModel(t *testing.T) { "zz-token-tiered-visible-model": `tier("base", p * 1 + c * 2)`, "zz-token-tiered-empty-expr-model": "", }) - db := setupModelListControllerTestDB(t) - require.NoError(t, db.Create(&[]model.Ability{ - {Group: "default", Model: "zz-token-tiered-visible-model", ChannelId: 1, Enabled: true}, - {Group: "default", Model: "zz-token-tiered-empty-expr-model", ChannelId: 1, Enabled: true}, - {Group: "default", Model: "zz-token-tiered-missing-expr-model", ChannelId: 1, Enabled: true}, - {Group: "default", Model: "zz-token-unpriced-model", ChannelId: 1, Enabled: true}, - }).Error) + setupModelListControllerTestDB(t) recorder := httptest.NewRecorder() ctx, _ := gin.CreateTestContext(recorder) @@ -431,68 +425,6 @@ func TestListModelsTokenLimitIncludesTieredBillingModel(t *testing.T) { require.NotContains(t, ids, "zz-token-unpriced-model") } -func TestListModelsTokenLimitUsesResolvedCustomAutoGroups(t *testing.T) { - withSelfUseModeEnabled(t) - originalMax := setting.GetMaxTokenAutoGroups() - originalUsableGroups := setting.UserUsableGroups2JSONString() - originalRatios := ratio_setting.GroupRatio2JSONString() - require.NoError(t, setting.UpdateMaxTokenAutoGroups("5")) - require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default","vip":"VIP"}`)) - require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":1}`)) - t.Cleanup(func() { - require.NoError(t, setting.UpdateMaxTokenAutoGroups(fmt.Sprintf("%d", originalMax))) - require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(originalUsableGroups)) - require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(originalRatios)) - }) - - db := setupModelListControllerTestDB(t) - require.NoError(t, db.Create(&[]model.Ability{ - {Group: "vip", Model: "zz-vip-allowed", ChannelId: 1, Enabled: true}, - {Group: "vip", Model: "zz-vip-denied", ChannelId: 1, Enabled: true}, - {Group: "default", Model: "zz-default-outside-snapshot", ChannelId: 1, Enabled: true}, - }).Error) - - recorder := httptest.NewRecorder() - ctx, _ := gin.CreateTestContext(recorder) - ctx.Request = httptest.NewRequest(http.MethodGet, "/v1/models", nil) - common.SetContextKey(ctx, constant.ContextKeyUserGroup, "default") - common.SetContextKey(ctx, constant.ContextKeyTokenGroup, "auto") - common.SetContextKey(ctx, constant.ContextKeyTokenAutoGroups, []string{"vip"}) - common.SetContextKey(ctx, constant.ContextKeyTokenModelLimitEnabled, true) - common.SetContextKey(ctx, constant.ContextKeyTokenModelLimit, map[string]bool{ - "zz-vip-allowed": true, - "zz-default-outside-snapshot": true, - "zz-not-enabled": true, - }) - - ListModels(ctx, constant.ChannelTypeOpenAI) - ids := decodeListModelsResponse(t, recorder) - require.Equal(t, map[string]struct{}{"zz-vip-allowed": {}}, ids) - - require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default"}`)) - emptyRecorder := httptest.NewRecorder() - emptyCtx, _ := gin.CreateTestContext(emptyRecorder) - emptyCtx.Request = httptest.NewRequest(http.MethodGet, "/v1/models", nil) - common.SetContextKey(emptyCtx, constant.ContextKeyUserGroup, "default") - common.SetContextKey(emptyCtx, constant.ContextKeyTokenGroup, "auto") - common.SetContextKey(emptyCtx, constant.ContextKeyTokenAutoGroups, []string{"vip"}) - common.SetContextKey(emptyCtx, constant.ContextKeyTokenModelLimitEnabled, true) - common.SetContextKey(emptyCtx, constant.ContextKeyTokenModelLimit, map[string]bool{"zz-vip-allowed": true}) - - require.NotPanics(t, func() { - ListModels(emptyCtx, constant.ChannelTypeAnthropic) - }) - var anthropicResponse struct { - Data []dto.AnthropicModel `json:"data"` - FirstID string `json:"first_id"` - LastID string `json:"last_id"` - } - require.NoError(t, common.Unmarshal(emptyRecorder.Body.Bytes(), &anthropicResponse)) - require.Empty(t, anthropicResponse.Data) - require.Empty(t, anthropicResponse.FirstID) - require.Empty(t, anthropicResponse.LastID) -} - func TestCheckUpdatePasswordRequiresCurrentPassword(t *testing.T) { db := setupModelListControllerTestDB(t) hashedPassword, err := common.Password2Hash("CurrentPassword123") diff --git a/controller/model_owned_by_test.go b/controller/model_owned_by_test.go index da9bb1a0ff0e..bc2ef32f135c 100644 --- a/controller/model_owned_by_test.go +++ b/controller/model_owned_by_test.go @@ -1,14 +1,11 @@ package controller import ( - "fmt" "net/http/httptest" "testing" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/constant" - "github.com/QuantumNous/new-api/setting" - "github.com/QuantumNous/new-api/setting/ratio_setting" "github.com/gin-gonic/gin" "github.com/stretchr/testify/require" ) @@ -86,33 +83,3 @@ func TestGetModelListGroupsUsesExplicitTokenGroup(t *testing.T) { require.Equal(t, "vip", groups.tokenGroup) require.Equal(t, []string{"vip"}, groups.ownerGroups) } - -func TestGetModelListGroupsUsesFilteredTokenAutoGroupsSnapshot(t *testing.T) { - originalMax := setting.GetMaxTokenAutoGroups() - originalUsableGroups := setting.UserUsableGroups2JSONString() - originalRatios := ratio_setting.GroupRatio2JSONString() - require.NoError(t, setting.UpdateMaxTokenAutoGroups("1")) - require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default","vip":"VIP"}`)) - require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":1}`)) - t.Cleanup(func() { - require.NoError(t, setting.UpdateMaxTokenAutoGroups(fmt.Sprintf("%d", originalMax))) - require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(originalUsableGroups)) - require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(originalRatios)) - }) - - gin.SetMode(gin.TestMode) - ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) - common.SetContextKey(ctx, constant.ContextKeyUserGroup, "default") - common.SetContextKey(ctx, constant.ContextKeyTokenGroup, "auto") - common.SetContextKey(ctx, constant.ContextKeyTokenAutoGroups, []string{"vip", "default"}) - - groups, err := getModelListGroups(ctx) - require.NoError(t, err) - require.Equal(t, []string{"vip"}, groups.ownerGroups) - - common.SetContextKey(ctx, constant.ContextKeyTokenAutoGroups, []string{"vip"}) - require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default"}`)) - groups, err = getModelListGroups(ctx) - require.NoError(t, err) - require.Empty(t, groups.ownerGroups) -} diff --git a/controller/public_site.go b/controller/public_site.go new file mode 100644 index 000000000000..79fde50810eb --- /dev/null +++ b/controller/public_site.go @@ -0,0 +1,134 @@ +package controller + +import ( + "strings" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/model" + + "github.com/gin-gonic/gin" +) + +// resolvePublicLocale 从 query 读取 locale，仅允许 en/zh，缺省 en。 +func resolvePublicLocale(c *gin.Context) string { + locale := strings.ToLower(strings.TrimSpace(c.Query("locale"))) + if locale != "zh" { + return "en" + } + return "zh" +} + +// GetPublicSiteConfig 返回站点品牌信息，供营销站页头/页脚使用。 +func GetPublicSiteConfig(c *gin.Context) { + common.ApiSuccess(c, gin.H{ + "system_name": common.SystemName, + "logo": common.Logo, + }) +} + +// GetPublicPricing 返回公开定价方案（按 locale）。 +func GetPublicPricing(c *gin.Context) { + locale := resolvePublicLocale(c) + var items []model.PublicPricing + if err := model.DB.Where("locale = ? AND enabled = ?", locale, true). + Order("sort asc").Find(&items).Error; err != nil { + common.ApiErrorMsg(c, err.Error()) + return + } + common.ApiSuccess(c, items) +} + +// GetPublicModelCategories 返回公开模型目录（按 locale）。 +func GetPublicModelCategories(c *gin.Context) { + locale := resolvePublicLocale(c) + var items []model.PublicModelCategory + if err := model.DB.Where("locale = ? AND enabled = ?", locale, true). + Order("sort asc").Find(&items).Error; err != nil { + common.ApiErrorMsg(c, err.Error()) + return + } + common.ApiSuccess(c, items) +} + +type salesLeadRequest struct { + Name string `json:"name"` + Email string `json:"email"` + Company string `json:"company"` + Region string `json:"region"` + UseCase string `json:"use_case"` + MonthlyVolume string `json:"monthly_volume"` + RequiredModels string `json:"required_models"` + Message string `json:"message"` + Source string `json:"source"` + Redirect string `json:"redirect"` +} + +// PostPublicSalesLead 接收「联系销售」表单，写入销售线索。 +func PostPublicSalesLead(c *gin.Context) { + var req salesLeadRequest + if err := c.ShouldBindJSON(&req); err != nil { + common.ApiErrorMsg(c, "请求格式错误") + return + } + + req.Name = strings.TrimSpace(req.Name) + req.Email = strings.TrimSpace(req.Email) + req.Region = strings.TrimSpace(req.Region) + req.UseCase = strings.TrimSpace(req.UseCase) + + if req.Name == "" || len(req.Name) > 64 { + common.ApiErrorMsg(c, "请填写有效的姓名（≤64 字符）") + return + } + if !strings.Contains(req.Email, "@") || len(req.Email) > 128 { + common.ApiErrorMsg(c, "请填写有效的邮箱") + return + } + if req.Region == "" { + common.ApiErrorMsg(c, "请选择所在区域") + return + } + if req.UseCase == "" || len(req.UseCase) > 256 { + common.ApiErrorMsg(c, "请填写使用场景（≤256 字符）") + return + } + if len(req.Company) > 128 || len(req.MonthlyVolume) > 64 || + len(req.RequiredModels) > 512 || len(req.Message) > 2000 { + common.ApiErrorMsg(c, "字段超出长度限制") + return + } + + // 开放重定向防护：redirect 仅允许本站相对路径 + safeRedirect := "" + if r := strings.TrimSpace(req.Redirect); r != "" { + if strings.HasPrefix(r, "/") && !strings.HasPrefix(r, "//") && !strings.Contains(r, ":") { + safeRedirect = r + } + } + + now := time.Now().Unix() + lead := model.SalesLead{ + Name: req.Name, + Email: req.Email, + Company: req.Company, + Region: req.Region, + UseCase: req.UseCase, + MonthlyVolume: req.MonthlyVolume, + RequiredModels: req.RequiredModels, + Message: req.Message, + Status: "new", + Source: req.Source, + CreatedAt: now, + UpdatedAt: now, + } + if err := model.DB.Create(&lead).Error; err != nil { + common.ApiErrorMsg(c, "提交失败，请稍后重试") + return + } + + common.ApiSuccess(c, gin.H{ + "id": lead.Id, + "redirect": safeRedirect, + }) +} diff --git a/controller/region_route.go b/controller/region_route.go new file mode 100644 index 000000000000..3f0e65fcd4ff --- /dev/null +++ b/controller/region_route.go @@ -0,0 +1,150 @@ +package controller + +import ( + "strconv" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/model" + "github.com/gin-gonic/gin" +) + +// ListRegionRoutes 分页列出区域路由策略。 +func ListRegionRoutes(c *gin.Context) { + page, pageSize := parsePage(c) + region := c.Query("region") + modelName := c.Query("model") + items, total, err := model.SearchRegionRoutes(page, pageSize, region, modelName) + if err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, gin.H{"items": items, "total": total}) +} + +// CreateRegionRoute 创建区域路由策略。 +func CreateRegionRoute(c *gin.Context) { + var req struct { + Region string `json:"region"` + Model string `json:"model"` + ChannelIds string `json:"channel_ids"` + Tag string `json:"tag"` + Strategy string `json:"strategy"` + Priority int `json:"priority"` + Weight int `json:"weight"` + Enabled *bool `json:"enabled"` + } + if err := common.DecodeJson(c.Request.Body, &req); err != nil { + common.ApiError(c, err) + return + } + if req.Region == "" { + common.ApiErrorMsg(c, "region is required") + return + } + if req.Model == "" { + req.Model = "*" + } + strategy := req.Strategy + if strategy == "" { + strategy = "availability" + } + if !model.AllowedRegionRouteStrategies[strategy] { + common.ApiErrorMsg(c, "invalid strategy (expected cost|latency|availability|fixed)") + return + } + enabled := true + if req.Enabled != nil { + enabled = *req.Enabled + } + m := &model.RegionRoute{ + Region: req.Region, + Model: req.Model, + ChannelIds: req.ChannelIds, + Tag: req.Tag, + Strategy: strategy, + Priority: req.Priority, + Weight: req.Weight, + Enabled: enabled, + } + if err := model.CreateRegionRoute(m); err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, gin.H{"id": m.Id}) +} + +// GetRegionRoute 获取单个策略。 +func GetRegionRoute(c *gin.Context) { + id, err := strconv.ParseInt(c.Param("id"), 10, 64) + if err != nil { + common.ApiErrorMsg(c, "invalid id") + return + } + m, err := model.GetRegionRouteById(id) + if err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, m) +} + +// UpdateRegionRoute 更新策略。 +func UpdateRegionRoute(c *gin.Context) { + id, err := strconv.ParseInt(c.Param("id"), 10, 64) + if err != nil { + common.ApiErrorMsg(c, "invalid id") + return + } + var req struct { + Region string `json:"region"` + Model string `json:"model"` + ChannelIds string `json:"channel_ids"` + Tag string `json:"tag"` + Strategy string `json:"strategy"` + Priority int `json:"priority"` + Weight int `json:"weight"` + Enabled *bool `json:"enabled"` + } + if err := common.DecodeJson(c.Request.Body, &req); err != nil { + common.ApiError(c, err) + return + } + if !model.AllowedRegionRouteStrategies[req.Strategy] { + common.ApiErrorMsg(c, "invalid strategy (expected cost|latency|availability|fixed)") + return + } + enabled := true + if req.Enabled != nil { + enabled = *req.Enabled + } + m := &model.RegionRoute{ + Id: id, + Region: req.Region, + Model: req.Model, + ChannelIds: req.ChannelIds, + Tag: req.Tag, + Strategy: req.Strategy, + Priority: req.Priority, + Weight: req.Weight, + Enabled: enabled, + } + if err := model.UpdateRegionRoute(m); err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, gin.H{"ok": true}) +} + +// DeleteRegionRoute 删除策略。 +func DeleteRegionRoute(c *gin.Context) { + id, err := strconv.ParseInt(c.Param("id"), 10, 64) + if err != nil { + common.ApiErrorMsg(c, "invalid id") + return + } + if err := model.DeleteRegionRoute(id); err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, gin.H{"ok": true}) +} diff --git a/controller/sla.go b/controller/sla.go new file mode 100644 index 000000000000..276c1b73b273 --- /dev/null +++ b/controller/sla.go @@ -0,0 +1,163 @@ +package controller + +import ( + "strconv" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/model" + "github.com/gin-gonic/gin" +) + +// ListSlaIncidents 管理端分页列出事件。 +func ListSlaIncidents(c *gin.Context) { + page, pageSize := parsePage(c) + status := c.Query("status") + items, total, err := model.SearchSlaIncidents(page, pageSize, status) + if err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, gin.H{"items": items, "total": total}) +} + +// CreateSlaIncident 创建事件。 +func CreateSlaIncident(c *gin.Context) { + var req struct { + Title string `json:"title"` + Description string `json:"description"` + Status int `json:"status"` + Severity string `json:"severity"` + StartedAt int64 `json:"started_at"` + ResolvedAt int64 `json:"resolved_at"` + } + if err := common.DecodeJson(c.Request.Body, &req); err != nil { + common.ApiError(c, err) + return + } + if req.Title == "" || len(req.Title) > 128 { + common.ApiErrorMsg(c, "title is required and must be <= 128 characters") + return + } + if !model.AllowedSlaIncidentStatuses[req.Status] { + common.ApiErrorMsg(c, "invalid status (expected 1 investigating|2 identified|3 monitoring|4 resolved)") + return + } + severity := req.Severity + if severity == "" { + severity = "minor" + } + if !model.AllowedSlaIncidentSeverities[severity] { + common.ApiErrorMsg(c, "invalid severity (expected minor|major|critical)") + return + } + m := &model.SlaIncident{ + Title: req.Title, + Description: req.Description, + Status: req.Status, + Severity: severity, + StartedAt: req.StartedAt, + ResolvedAt: req.ResolvedAt, + } + if err := model.CreateSlaIncident(m); err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, gin.H{"id": m.Id}) +} + +// GetSlaIncident 获取单个事件。 +func GetSlaIncident(c *gin.Context) { + id, err := strconv.ParseInt(c.Param("id"), 10, 64) + if err != nil { + common.ApiErrorMsg(c, "invalid id") + return + } + m, err := model.GetSlaIncidentById(id) + if err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, m) +} + +// UpdateSlaIncident 更新事件。 +func UpdateSlaIncident(c *gin.Context) { + id, err := strconv.ParseInt(c.Param("id"), 10, 64) + if err != nil { + common.ApiErrorMsg(c, "invalid id") + return + } + var req struct { + Title string `json:"title"` + Description string `json:"description"` + Status int `json:"status"` + Severity string `json:"severity"` + StartedAt int64 `json:"started_at"` + ResolvedAt int64 `json:"resolved_at"` + } + if err := common.DecodeJson(c.Request.Body, &req); err != nil { + common.ApiError(c, err) + return + } + if !model.AllowedSlaIncidentStatuses[req.Status] { + common.ApiErrorMsg(c, "invalid status (expected 1 investigating|2 identified|3 monitoring|4 resolved)") + return + } + if !model.AllowedSlaIncidentSeverities[req.Severity] { + common.ApiErrorMsg(c, "invalid severity (expected minor|major|critical)") + return + } + m := &model.SlaIncident{ + Id: id, + Title: req.Title, + Description: req.Description, + Status: req.Status, + Severity: req.Severity, + StartedAt: req.StartedAt, + ResolvedAt: req.ResolvedAt, + } + if err := model.UpdateSlaIncident(m); err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, gin.H{"ok": true}) +} + +// DeleteSlaIncident 删除事件。 +func DeleteSlaIncident(c *gin.Context) { + id, err := strconv.ParseInt(c.Param("id"), 10, 64) + if err != nil { + common.ApiErrorMsg(c, "invalid id") + return + } + if err := model.DeleteSlaIncident(id); err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, gin.H{"ok": true}) +} + +// GetPublicSlaIncidents 公开事件列表（状态页展示，匿名可读）。 +func GetPublicSlaIncidents(c *gin.Context) { + // 公开页展示最近事件（含已解决），限制条数避免过大。 + items, _, err := model.SearchSlaIncidents(1, 50, "") + if err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, gin.H{"items": items}) +} + +// GetPublicSlaStatus 公开服务状态摘要（匿名可读），聚合渠道与性能数据。 +func GetPublicSlaStatus(c *gin.Context) { + window := 24 + if w, err := strconv.Atoi(c.Query("window_hours")); err == nil && w > 0 && w <= 720 { + window = w + } + summary, err := model.GetSlaStatusSummary(window) + if err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, summary) +} diff --git a/controller/team.go b/controller/team.go new file mode 100644 index 000000000000..5fed41106961 --- /dev/null +++ b/controller/team.go @@ -0,0 +1,258 @@ +package controller + +import ( + "strconv" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/model" + "github.com/gin-gonic/gin" +) + +// parsePage 解析分页参数，缺省 page=1, pageSize=10。 +func parsePage(c *gin.Context) (int, int) { + page, _ := strconv.Atoi(c.Query("page")) + pageSize, _ := strconv.Atoi(c.Query("page_size")) + if page <= 0 { + page = 1 + } + if pageSize <= 0 { + pageSize = 10 + } + return page, pageSize +} + +// ListTeams 分页列出团队。 +func ListTeams(c *gin.Context) { + page, pageSize := parsePage(c) + keyword := c.Query("keyword") + items, total, err := model.ListTeams(page, pageSize, keyword) + if err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, gin.H{"items": items, "total": total}) +} + +// CreateTeam 创建团队。 +func CreateTeam(c *gin.Context) { + var req struct { + Name string `json:"name"` + Description string `json:"description"` + OwnerId int64 `json:"owner_id"` + } + if err := common.DecodeJson(c.Request.Body, &req); err != nil { + common.ApiError(c, err) + return + } + if req.Name == "" || len(req.Name) > 64 { + common.ApiErrorMsg(c, "name is required and must be <= 64 characters") + return + } + if req.OwnerId <= 0 { + common.ApiErrorMsg(c, "owner_id is required") + return + } + t := &model.Team{Name: req.Name, Description: req.Description, OwnerId: req.OwnerId} + if err := model.CreateTeam(t); err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, gin.H{"id": t.Id}) +} + +// GetTeam 获取单个团队。 +func GetTeam(c *gin.Context) { + id, err := strconv.ParseInt(c.Param("id"), 10, 64) + if err != nil { + common.ApiErrorMsg(c, "invalid id") + return + } + t, err := model.GetTeamById(id) + if err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, t) +} + +// UpdateTeam 更新团队。 +func UpdateTeam(c *gin.Context) { + id, err := strconv.ParseInt(c.Param("id"), 10, 64) + if err != nil { + common.ApiErrorMsg(c, "invalid id") + return + } + var req struct { + Description string `json:"description"` + OwnerId int64 `json:"owner_id"` + } + if err := common.DecodeJson(c.Request.Body, &req); err != nil { + common.ApiError(c, err) + return + } + if err := model.UpdateTeam(&model.Team{Id: id, Description: req.Description, OwnerId: req.OwnerId}); err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, gin.H{"ok": true}) +} + +// DeleteTeam 删除团队（级联成员与项目）。 +func DeleteTeam(c *gin.Context) { + id, err := strconv.ParseInt(c.Param("id"), 10, 64) + if err != nil { + common.ApiErrorMsg(c, "invalid id") + return + } + if err := model.DeleteTeam(id); err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, gin.H{"ok": true}) +} + +// AddTeamMember 添加团队成员。 +func AddTeamMember(c *gin.Context) { + teamId, err := strconv.ParseInt(c.Param("id"), 10, 64) + if err != nil { + common.ApiErrorMsg(c, "invalid team id") + return + } + var req struct { + UserId int64 `json:"user_id"` + Role string `json:"role"` + } + if err := common.DecodeJson(c.Request.Body, &req); err != nil { + common.ApiError(c, err) + return + } + if req.UserId <= 0 { + common.ApiErrorMsg(c, "user_id is required") + return + } + role := req.Role + if role == "" { + role = "member" + } + if !model.AllowedTeamMemberRoles[role] { + common.ApiErrorMsg(c, "invalid role (expected admin|member)") + return + } + m := &model.TeamMember{TeamId: teamId, UserId: req.UserId, Role: role} + if err := model.CreateTeamMember(m); err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, gin.H{"id": m.Id}) +} + +// ListTeamMembers 分页列出团队成员。 +func ListTeamMembers(c *gin.Context) { + teamId, err := strconv.ParseInt(c.Param("id"), 10, 64) + if err != nil { + common.ApiErrorMsg(c, "invalid team id") + return + } + page, pageSize := parsePage(c) + items, total, err := model.ListTeamMembers(teamId, page, pageSize) + if err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, gin.H{"items": items, "total": total}) +} + +// RemoveTeamMember 移除成员。 +func RemoveTeamMember(c *gin.Context) { + teamId, err := strconv.ParseInt(c.Param("id"), 10, 64) + if err != nil { + common.ApiErrorMsg(c, "invalid team id") + return + } + userId, err := strconv.ParseInt(c.Param("user_id"), 10, 64) + if err != nil { + common.ApiErrorMsg(c, "invalid user_id") + return + } + if err := model.DeleteTeamMember(teamId, userId); err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, gin.H{"ok": true}) +} + +// AddTeamProject 添加团队项目。 +func AddTeamProject(c *gin.Context) { + teamId, err := strconv.ParseInt(c.Param("id"), 10, 64) + if err != nil { + common.ApiErrorMsg(c, "invalid team id") + return + } + var req struct { + Name string `json:"name"` + Description string `json:"description"` + } + if err := common.DecodeJson(c.Request.Body, &req); err != nil { + common.ApiError(c, err) + return + } + if req.Name == "" || len(req.Name) > 64 { + common.ApiErrorMsg(c, "name is required and must be <= 64 characters") + return + } + p := &model.TeamProject{TeamId: teamId, Name: req.Name, Description: req.Description} + if err := model.CreateTeamProject(p); err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, gin.H{"id": p.Id}) +} + +// ListTeamProjects 列出团队项目。 +func ListTeamProjects(c *gin.Context) { + teamId, err := strconv.ParseInt(c.Param("id"), 10, 64) + if err != nil { + common.ApiErrorMsg(c, "invalid team id") + return + } + items, err := model.ListTeamProjects(teamId) + if err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, gin.H{"items": items}) +} + +// RemoveTeamProject 删除团队项目。 +func RemoveTeamProject(c *gin.Context) { + teamId, err := strconv.ParseInt(c.Param("id"), 10, 64) + if err != nil { + common.ApiErrorMsg(c, "invalid team id") + return + } + projectId, err := strconv.ParseInt(c.Param("pid"), 10, 64) + if err != nil { + common.ApiErrorMsg(c, "invalid pid") + return + } + if err := model.DeleteTeamProject(teamId, projectId); err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, gin.H{"ok": true}) +} + +// GetTeamBilling 部门账单汇总。 +func GetTeamBilling(c *gin.Context) { + teamId, err := strconv.ParseInt(c.Param("id"), 10, 64) + if err != nil { + common.ApiErrorMsg(c, "invalid team id") + return + } + billing, err := model.GetTeamBilling(teamId) + if err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, billing) +} diff --git a/controller/token.go b/controller/token.go index c26d82e3dee1..836e9b2952ac 100644 --- a/controller/token.go +++ b/controller/token.go @@ -7,115 +7,30 @@ import ( "strings" "github.com/QuantumNous/new-api/common" - "github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/i18n" "github.com/QuantumNous/new-api/model" - "github.com/QuantumNous/new-api/service" - "github.com/QuantumNous/new-api/setting" "github.com/QuantumNous/new-api/setting/operation_setting" "github.com/gin-gonic/gin" ) -type tokenAutoGroupsInput struct { - Set bool - Groups []string -} - -func (input *tokenAutoGroupsInput) UnmarshalJSON(data []byte) error { - input.Set = true - if strings.TrimSpace(string(data)) == "null" { - input.Groups = nil - return nil - } - return common.Unmarshal(data, &input.Groups) -} - -type tokenRequest struct { - model.Token - AutoGroups tokenAutoGroupsInput `json:"auto_groups"` -} - -type tokenResponse struct { - *model.Token - AutoGroups []string `json:"auto_groups"` -} - -func buildMaskedTokenResponse(token *model.Token) *tokenResponse { +func buildMaskedTokenResponse(token *model.Token) *model.Token { if token == nil { return nil } maskedToken := *token maskedToken.Key = token.GetMaskedKey() - autoGroups, err := token.GetAutoGroups() - if err != nil { - common.SysError(fmt.Sprintf("failed to parse auto groups for token %d: %v", token.Id, err)) - autoGroups = nil - } - if len(autoGroups) == 0 { - autoGroups = nil - } - return &tokenResponse{Token: &maskedToken, AutoGroups: autoGroups} + return &maskedToken } -func buildMaskedTokenResponses(tokens []*model.Token) []*tokenResponse { - maskedTokens := make([]*tokenResponse, 0, len(tokens)) +func buildMaskedTokenResponses(tokens []*model.Token) []*model.Token { + maskedTokens := make([]*model.Token, 0, len(tokens)) for _, token := range tokens { maskedTokens = append(maskedTokens, buildMaskedTokenResponse(token)) } return maskedTokens } -func getTokenRequestUserGroup(c *gin.Context) (string, error) { - if userGroup := common.GetContextKeyString(c, constant.ContextKeyUserGroup); userGroup != "" { - return userGroup, nil - } - if userGroup := c.GetString("group"); userGroup != "" { - return userGroup, nil - } - return model.GetUserGroup(c.GetInt("id"), false) -} - -func setTokenAutoGroups(c *gin.Context, token *model.Token, groups []string) bool { - if len(groups) == 0 { - if err := token.SetAutoGroups(nil); err != nil { - common.ApiError(c, err) - return false - } - return true - } - - maxCount := setting.GetMaxTokenAutoGroups() - if len(groups) > maxCount { - common.ApiErrorI18n(c, i18n.MsgTokenAutoGroupsTooMany, map[string]any{"Max": maxCount}) - return false - } - - userGroup, err := getTokenRequestUserGroup(c) - if err != nil { - common.ApiError(c, err) - return false - } - seen := make(map[string]struct{}, len(groups)) - for _, group := range groups { - if _, ok := seen[group]; ok { - common.ApiErrorI18n(c, i18n.MsgTokenAutoGroupsDuplicate, map[string]any{"Group": group}) - return false - } - seen[group] = struct{}{} - if !service.IsUserSelectableGroup(userGroup, group) { - common.ApiErrorI18n(c, i18n.MsgTokenAutoGroupsInvalid, map[string]any{"Group": group}) - return false - } - } - - if err := token.SetAutoGroups(groups); err != nil { - common.ApiError(c, err) - return false - } - return true -} - func GetAllTokens(c *gin.Context) { userId := c.GetInt("id") pageInfo := common.GetPageQuery(c) @@ -162,18 +77,6 @@ func GetToken(c *gin.Context) { common.ApiSuccess(c, buildMaskedTokenResponse(token)) } -func GetTokenAutoGroups(c *gin.Context) { - userGroup, err := getTokenRequestUserGroup(c) - if err != nil { - common.ApiError(c, err) - return - } - common.ApiSuccess(c, gin.H{ - "groups": service.GetUserAutoGroup(userGroup), - "max_count": setting.GetMaxTokenAutoGroups(), - }) -} - func GetTokenKey(c *gin.Context) { id, err := strconv.Atoi(c.Param("id")) userId := c.GetInt("id") @@ -262,13 +165,12 @@ func GetTokenUsage(c *gin.Context) { } func AddToken(c *gin.Context) { - request := tokenRequest{} - err := c.ShouldBindJSON(&request) + token := model.Token{} + err := c.ShouldBindJSON(&token) if err != nil { common.ApiError(c, err) return } - token := request.Token if len(token.Name) > 50 { common.ApiErrorI18n(c, i18n.MsgTokenNameTooLong) return @@ -299,14 +201,6 @@ func AddToken(c *gin.Context) { }) return } - if token.Group == "auto" { - if !setTokenAutoGroups(c, &token, request.AutoGroups.Groups) { - return - } - } else { - token.CrossGroupRetry = false - _ = token.SetAutoGroups(nil) - } key, err := common.GenerateKey() if err != nil { common.ApiErrorI18n(c, i18n.MsgTokenGenerateFailed) @@ -327,7 +221,6 @@ func AddToken(c *gin.Context) { AllowIps: token.AllowIps, Group: token.Group, CrossGroupRetry: token.CrossGroupRetry, - AutoGroups: token.AutoGroups, } err = cleanToken.Insert() if err != nil { @@ -357,13 +250,12 @@ func DeleteToken(c *gin.Context) { func UpdateToken(c *gin.Context) { userId := c.GetInt("id") statusOnly := c.Query("status_only") - request := tokenRequest{} - err := c.ShouldBindJSON(&request) + token := model.Token{} + err := c.ShouldBindJSON(&token) if err != nil { common.ApiError(c, err) return } - token := request.Token if len(token.Name) > 50 { common.ApiErrorI18n(c, i18n.MsgTokenNameTooLong) return @@ -407,14 +299,6 @@ func UpdateToken(c *gin.Context) { cleanToken.AllowIps = token.AllowIps cleanToken.Group = token.Group cleanToken.CrossGroupRetry = token.CrossGroupRetry - if token.Group != "auto" { - cleanToken.CrossGroupRetry = false - _ = cleanToken.SetAutoGroups(nil) - } else if request.AutoGroups.Set { - if !setTokenAutoGroups(c, cleanToken, request.AutoGroups.Groups) { - return - } - } } err = cleanToken.Update() if err != nil { diff --git a/controller/token_test.go b/controller/token_test.go index 9cca168ab100..12b1cbdd84fb 100644 --- a/controller/token_test.go +++ b/controller/token_test.go @@ -273,34 +273,6 @@ func getTokenKeyColumnType(t *testing.T, db *gorm.DB, dialect string) string { } } -func getTokenAutoGroupsColumnType(t *testing.T, db *gorm.DB, dialect string) string { - t.Helper() - - switch dialect { - case "sqlite": - return getSQLiteColumnType(t, db, "tokens", "auto_groups") - case "mysql": - var columnType string - if err := db.Raw(`SELECT DATA_TYPE FROM information_schema.columns - WHERE table_schema = DATABASE() AND table_name = ? AND column_name = ?`, - "tokens", "auto_groups").Scan(&columnType).Error; err != nil { - t.Fatalf("failed to inspect mysql token auto_groups column: %v", err) - } - return strings.ToLower(columnType) - case "postgres": - var dataType string - if err := db.Raw(`SELECT data_type FROM information_schema.columns - WHERE table_schema = current_schema() AND table_name = ? AND column_name = ?`, - "tokens", "auto_groups").Scan(&dataType).Error; err != nil { - t.Fatalf("failed to inspect postgres token auto_groups column: %v", err) - } - return strings.ToLower(dataType) - default: - t.Fatalf("unsupported dialect %q", dialect) - return "" - } -} - func runTokenMigrationCompatibilityTest(t *testing.T, db *gorm.DB, dialect string, managedTokensTable *bool) { t.Helper() @@ -342,12 +314,6 @@ func runTokenMigrationCompatibilityTest(t *testing.T, db *gorm.DB, dialect strin if got := getTokenKeyColumnType(t, db, dialect); got != "varchar(128)" { t.Fatalf("expected migrated key column type varchar(128), got %q", got) } - if !db.Migrator().HasColumn(&model.Token{}, "auto_groups") { - t.Fatal("expected migration to add auto_groups column") - } - if got := getTokenAutoGroupsColumnType(t, db, dialect); got != "text" { - t.Fatalf("expected migrated auto_groups column type text, got %q", got) - } var migratedToken model.Token if err := db.First(&migratedToken, "name = ?", "legacy-token").Error; err != nil { @@ -359,9 +325,6 @@ func runTokenMigrationCompatibilityTest(t *testing.T, db *gorm.DB, dialect strin if migratedToken.Name != "legacy-token" { t.Fatalf("expected migrated token name to be preserved, got %q", migratedToken.Name) } - if migratedToken.AutoGroups != "" { - t.Fatalf("expected legacy token to inherit global Auto groups, got %q", migratedToken.AutoGroups) - } inserted := model.Token{ UserId: 8, @@ -399,9 +362,6 @@ func TestTokenAutoMigrateUsesVarchar128KeyColumn(t *testing.T) { if got := getTokenKeyColumnType(t, db, "sqlite"); got != "varchar(128)" { t.Fatalf("expected key column type varchar(128), got %q", got) } - if got := getSQLiteColumnType(t, db, "tokens", "auto_groups"); got != "text" { - t.Fatalf("expected auto_groups column type text, got %q", got) - } } func TestTokenMigrationFromChar48ToVarchar128(t *testing.T) { diff --git a/controller/user.go b/controller/user.go index 9b8d931ec1f8..59aaf5847238 100644 --- a/controller/user.go +++ b/controller/user.go @@ -399,6 +399,11 @@ func GetUser(c *gin.Context) { func GenerateAccessToken(c *gin.Context) { id := c.GetInt("id") + user, err := model.GetUserById(id, true) + if err != nil { + common.ApiError(c, err) + return + } // get rand int 28-32 randI := common.GetRandomInt(4) key, err := common.GenerateRandomKey(29 + randI) @@ -407,12 +412,14 @@ func GenerateAccessToken(c *gin.Context) { common.SysLog("failed to generate key: " + err.Error()) return } - if model.DB.Where("access_token = ?", key).First(&model.User{}).RowsAffected != 0 { + user.SetAccessToken(key) + + if model.DB.Where("access_token = ?", user.AccessToken).First(user).RowsAffected != 0 { common.ApiErrorI18n(c, i18n.MsgUuidDuplicate) return } - if err := model.UpdateUserAccessToken(id, key); err != nil { + if err := user.Update(false); err != nil { common.ApiError(c, err) return } @@ -420,7 +427,7 @@ func GenerateAccessToken(c *gin.Context) { c.JSON(http.StatusOK, gin.H{ "success": true, "message": "", - "data": key, + "data": user.AccessToken, }) return } diff --git a/deploy/nginx/originflow.conf b/deploy/nginx/originflow.conf new file mode 100644 index 000000000000..697036f94b38 --- /dev/null +++ b/deploy/nginx/originflow.conf @@ -0,0 +1,156 @@ +# ============================================================================= +# OriginFlow (元点流商) — nginx 站点配置 +# ----------------------------------------------------------------------------- +# 域名规划（ADR-1 域名隔离）： +# www.91flow.com -> 营销前台（由 SPA 的营销路由树渲染） +# app.91flow.com -> New-API 原生控制台 +# api.91flow.com -> Relay / OpenAI 兼容 API（对外调用入口） +# +# 架构说明： +# 生产构建为单个内嵌 SPA（web/default/dist 由 main.go 的 //go:embed 提供， +# SPA fallback 在 router/web-router.go 的 NoRoute）。同一份构建经两个 server +# block 提供服务。前端按 location.host 以 'www.' 开头切换营销/控制台路由， +# 因此 nginx 只需把全部请求（SPA 静态 + /api + /v1）反代到 Go 后端即可。 +# +# 后端地址： +# 裸机/VM：`server 127.0.0.1:3000;` +# Docker Compose：把下面 upstream 的地址改为 `server new-api:3000;` +# +# 证书（Let's Encrypt / certbot）： +# certbot certonly --webroot -w /var/www/certbot \ +# -d www.91flow.com -d app.91flow.com -d api.91flow.com +# 证书路径：/etc/letsencrypt/live//fullchain.pem 与 privkey.pem +# ============================================================================= + +# WebSocket / SSE 升级：仅当客户端发起 Upgrade 时才置为 upgrade +map $http_upgrade $connection_upgrade { + default upgrade; + '' close; +} + +upstream originflow_backend { + # 裸机：127.0.0.1:3000 ；Docker：new-api:3000 + server 127.0.0.1:3000; + keepalive 32; +} + +# 通用代理参数（被各 server block 复用） +proxy_read_timeout 300s; +proxy_send_timeout 300s; +proxy_set_header Host $host; +proxy_set_header X-Real-IP $remote_addr; +proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for; +proxy_set_header X-Forwarded-Proto $scheme; +proxy_set_header X-Forwarded-Host $host; +proxy_set_header Upgrade $http_upgrade; +proxy_set_header Connection $connection_upgrade; + +# ----------------------------------------------------------------------------- +# HTTP -> HTTPS（三个域名统一跳转） +# ----------------------------------------------------------------------------- +server { + listen 80; + server_name www.91flow.com app.91flow.com api.91flow.com; + + # certbot 验证用 + location /.well-known/acme-challenge/ { + root /var/www/certbot; + } + + location / { + return 301 https://$host$request_uri; + } +} + +# ----------------------------------------------------------------------------- +# www.91flow.com — 营销前台 +# ----------------------------------------------------------------------------- +server { + listen 443 ssl; + http2 on; + server_name www.91flow.com; + + ssl_certificate /etc/letsencrypt/live/www.91flow.com/fullchain.pem; + ssl_certificate_key /etc/letsencrypt/live/www.91flow.com/privkey.pem; + ssl_protocols TLSv1.2 TLSv1.3; + ssl_ciphers HIGH:!aNULL:!MD5; + ssl_session_cache shared:SSL:10m; + ssl_session_timeout 1d; + ssl_session_tickets off; + + # 安全响应头 + add_header Strict-Transport-Security "max-age=63072000; includeSubDomains" always; + add_header X-Content-Type-Options "nosniff" always; + add_header X-Frame-Options "SAMEORIGIN" always; + add_header Referrer-Policy "strict-origin-when-cross-origin" always; + + client_max_body_size 10m; # 联系销售等表单提交 + + location / { + proxy_pass http://originflow_backend; + } +} + +# ----------------------------------------------------------------------------- +# app.91flow.com — 原生控制台 +# ----------------------------------------------------------------------------- +server { + listen 443 ssl; + http2 on; + server_name app.91flow.com; + + ssl_certificate /etc/letsencrypt/live/app.91flow.com/fullchain.pem; + ssl_certificate_key /etc/letsencrypt/live/app.91flow.com/privkey.pem; + ssl_protocols TLSv1.2 TLSv1.3; + ssl_ciphers HIGH:!aNULL:!MD5; + ssl_session_cache shared:SSL:10m; + ssl_session_timeout 1d; + ssl_session_tickets off; + + add_header Strict-Transport-Security "max-age=63072000; includeSubDomains" always; + add_header X-Content-Type-Options "nosniff" always; + add_header X-Frame-Options "SAMEORIGIN" always; + add_header Referrer-Policy "strict-origin-when-cross-origin" always; + + client_max_body_size 50m; # 控制台含文件上传/充值等 + + location / { + proxy_pass http://originflow_backend; + } +} + +# ----------------------------------------------------------------------------- +# api.91flow.com — Relay / OpenAI 兼容 API（对外调用入口） +# 开启 CORS 以便第三方前端跨域调用；其余行为与另外两个域名一致（共用同一后端）。 +# ----------------------------------------------------------------------------- +server { + listen 443 ssl; + http2 on; + server_name api.91flow.com; + + ssl_certificate /etc/letsencrypt/live/api.91flow.com/fullchain.pem; + ssl_certificate_key /etc/letsencrypt/live/api.91flow.com/privkey.pem; + ssl_protocols TLSv1.2 TLSv1.3; + ssl_ciphers HIGH:!aNULL:!MD5; + ssl_session_cache shared:SSL:10m; + ssl_session_timeout 1d; + ssl_session_tickets off; + + add_header Strict-Transport-Security "max-age=63072000; includeSubDomains" always; + add_header X-Content-Type-Options "nosniff" always; + + # 对外 API：开放 CORS（使用 Bearer Token，不涉及 Cookie，故可用 "*"） + add_header Access-Control-Allow-Origin "*" always; + add_header Access-Control-Allow-Methods "GET, POST, PUT, DELETE, OPTIONS" always; + add_header Access-Control-Allow-Headers "Authorization, Content-Type" always; + + client_max_body_size 50m; + + # 预检请求直接返回 204 + location / { + if ($request_method = OPTIONS) { + return 204; + } + proxy_pass http://originflow_backend; + } +} diff --git a/docker-compose.yml b/docker-compose.yml index 8e6fe4b57b6d..c6a13204d2f0 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -1,9 +1,13 @@ # New-API Docker Compose Configuration # # Quick Start: -# 1. docker-compose up -d +# 1. docker compose up -d --build # 2. Access at http://localhost:3000 # +# The new-api image is built from this repository's source (see Dockerfile). +# After changing Go or web code, rebuild with: +# docker compose up -d --build new-api +# # Using MySQL instead of PostgreSQL: # 1. Comment out the postgres service and SQL_DSN line 15 # 2. Uncomment the mysql service and SQL_DSN line 16 @@ -16,7 +20,18 @@ version: '3.4' # For compatibility with older Docker versions services: new-api: - image: calciumion/new-api:latest + # 使用本仓库源码构建，而非上游发行镜像，否则二次开发的功能不会包含在内 + # (Build from this repository's source instead of the upstream release image, + # otherwise none of the in-house features are included.) + build: + context: . + dockerfile: Dockerfile + args: + # 网络较慢时用就近镜像源，避免大体积依赖下载超时导致 integrity 校验失败 + # (Use a nearby mirror on slow links; large deps otherwise time out and + # fail bun's integrity check.) Override via NPM_REGISTRY in your shell. + NPM_REGISTRY: ${NPM_REGISTRY:-} + image: new-api:local container_name: new-api restart: always command: --log-dir /app/logs diff --git a/i18n/keys.go b/i18n/keys.go index 64a835e1a942..8e9a4b5694d4 100644 --- a/i18n/keys.go +++ b/i18n/keys.go @@ -55,9 +55,6 @@ const ( MsgTokenExhausted = "token.exhausted" MsgTokenStatusUnavailable = "token.status_unavailable" MsgTokenDbError = "token.db_error" - MsgTokenAutoGroupsTooMany = "token.auto_groups_too_many" - MsgTokenAutoGroupsDuplicate = "token.auto_groups_duplicate" - MsgTokenAutoGroupsInvalid = "token.auto_groups_invalid" ) // Redemption related messages diff --git a/i18n/locales/en.yaml b/i18n/locales/en.yaml index c533daecc32d..3f1fd03cb090 100644 --- a/i18n/locales/en.yaml +++ b/i18n/locales/en.yaml @@ -47,9 +47,6 @@ token.expired: "This token has expired" token.exhausted: "This token quota is exhausted TokenStatusExhausted[sk-{{.Prefix}}***{{.Suffix}}]" token.status_unavailable: "This token status is unavailable" token.db_error: "Invalid token, database query error, please contact administrator" -token.auto_groups_too_many: "A token can select at most {{.Max}} Auto groups" -token.auto_groups_duplicate: "Auto group {{.Group}} is duplicated" -token.auto_groups_invalid: "Auto group {{.Group}} is unavailable or unauthorized" # Redemption messages redemption.name_length: "Redemption code name length must be between 1-20" diff --git a/i18n/locales/zh-CN.yaml b/i18n/locales/zh-CN.yaml index a2f5275be9a8..fe982e59a0f3 100644 --- a/i18n/locales/zh-CN.yaml +++ b/i18n/locales/zh-CN.yaml @@ -48,9 +48,6 @@ token.expired: "该令牌已过期" token.exhausted: "该令牌额度已用尽 TokenStatusExhausted[sk-{{.Prefix}}***{{.Suffix}}]" token.status_unavailable: "该令牌状态不可用" token.db_error: "无效的令牌，数据库查询出错，请联系管理员" -token.auto_groups_too_many: "每个令牌最多可选择 {{.Max}} 个 Auto 分组" -token.auto_groups_duplicate: "Auto 分组 {{.Group}} 重复" -token.auto_groups_invalid: "Auto 分组 {{.Group}} 不可用或无权访问" # Redemption messages redemption.name_length: "兑换码名称长度必须在1-20之间" diff --git a/i18n/locales/zh-TW.yaml b/i18n/locales/zh-TW.yaml index 84ebd57ed587..27759d07f378 100644 --- a/i18n/locales/zh-TW.yaml +++ b/i18n/locales/zh-TW.yaml @@ -48,9 +48,6 @@ token.expired: "該令牌已過期" token.exhausted: "該令牌額度已用盡 TokenStatusExhausted[sk-{{.Prefix}}***{{.Suffix}}]" token.status_unavailable: "該令牌狀態不可用" token.db_error: "無效的令牌，資料庫查詢出錯，請聯繫管理員" -token.auto_groups_too_many: "每個令牌最多可選擇 {{.Max}} 個 Auto 分組" -token.auto_groups_duplicate: "Auto 分組 {{.Group}} 重複" -token.auto_groups_invalid: "Auto 分組 {{.Group}} 不可用或無權存取" # Redemption messages redemption.name_length: "兌換碼名稱長度必須在1-20之間" diff --git a/middleware/auth.go b/middleware/auth.go index 4e1436f38466..2ad09a7a8d2d 100644 --- a/middleware/auth.go +++ b/middleware/auth.go @@ -502,16 +502,6 @@ func SetupContextForToken(c *gin.Context, token *model.Token, parts ...string) e } common.SetContextKey(c, constant.ContextKeyTokenGroup, token.Group) common.SetContextKey(c, constant.ContextKeyTokenCrossGroupRetry, token.CrossGroupRetry) - if token.AutoGroups != "" { - autoGroups, err := token.GetAutoGroups() - if err != nil { - common.SysError(fmt.Sprintf("failed to parse auto groups for token %d: %v", token.Id, err)) - autoGroups = []string{} - common.SetContextKey(c, constant.ContextKeyTokenAutoGroups, autoGroups) - } else if len(autoGroups) > 0 { - common.SetContextKey(c, constant.ContextKeyTokenAutoGroups, autoGroups) - } - } if len(parts) > 1 { if model.IsAdmin(token.UserId) { c.Set("specific_channel_id", parts[1]) diff --git a/middleware/distributor.go b/middleware/distributor.go index 7decf0e28728..bde639ddee5b 100644 --- a/middleware/distributor.go +++ b/middleware/distributor.go @@ -109,7 +109,7 @@ func Distribute() func(c *gin.Context) { channelSupportsRequestPath(preferred, c.Request.URL.Path, modelRequest.Model) { if usingGroup == "auto" { userGroup := common.GetContextKeyString(c, constant.ContextKeyUserGroup) - autoGroups := service.GetRequestAutoGroups(c, userGroup) + autoGroups := service.GetUserAutoGroup(userGroup) for _, g := range autoGroups { if model.IsChannelEnabledForGroupModel(g, modelRequest.Model, preferred.Id) { selectGroup = g diff --git a/model/ability.go b/model/ability.go index d950a6adbfc4..7e53e688df9f 100644 --- a/model/ability.go +++ b/model/ability.go @@ -105,7 +105,9 @@ func getChannelQuery(group string, model string, retry int) (*gorm.DB, error) { return channelQuery, nil } -func GetChannel(group string, model string, retry int, requestPath string) (*Channel, error) { +// GetChannel 在未开启内存缓存时直接查库选择渠道。 +// routing 为区域路由解析结果，语义与 GetRandomSatisfiedChannel 一致。 +func GetChannel(group string, model string, retry int, requestPath string, routing RegionRouting) (*Channel, error) { var abilities []Ability var err error = nil @@ -122,6 +124,7 @@ func GetChannel(group string, model string, retry int, requestPath string) (*Cha return nil, err } abilities = filterAbilitiesByRequestPathAndModel(abilities, requestPath, model) + abilities = filterAbilitiesByRegion(abilities, routing) channel := Channel{} if len(abilities) > 0 { // Randomly choose one @@ -146,6 +149,73 @@ func GetChannel(group string, model string, retry int, requestPath string) (*Cha return &channel, err } +// filterAbilitiesByRegion 按区域路由收窄候选能力：先用白名单过滤渠道 id， +// 再在区域策略生效时只保留策略分数最高的一批渠道。 +// 任一步骤过滤为空时都回退到过滤前的候选集，避免配置问题导致请求失败。 +func filterAbilitiesByRegion(abilities []Ability, routing RegionRouting) []Ability { + if !routing.Active || len(abilities) == 0 { + return abilities + } + + if len(routing.AllowedIds) > 0 { + allowed := make(map[int]bool, len(routing.AllowedIds)) + for _, id := range routing.AllowedIds { + allowed[int(id)] = true + } + filtered := make([]Ability, 0, len(abilities)) + for _, ability := range abilities { + if allowed[ability.ChannelId] { + filtered = append(filtered, ability) + } + } + if len(filtered) > 0 { + abilities = filtered + } + } + + if routing.Strategy == "" || len(abilities) < 2 { + return abilities + } + + channelIds := make([]int, 0, len(abilities)) + seen := make(map[int]struct{}, len(abilities)) + for _, ability := range abilities { + if _, ok := seen[ability.ChannelId]; ok { + continue + } + seen[ability.ChannelId] = struct{}{} + channelIds = append(channelIds, ability.ChannelId) + } + var channels []*Channel + if err := DB.Where("id IN ?", channelIds).Find(&channels).Error; err != nil { + return abilities + } + scores := make(map[int]int64, len(channels)) + var bestScore int64 + first := true + for _, channel := range channels { + score := regionStrategyScore(channel, routing.Strategy) + scores[channel.Id] = score + if first || score > bestScore { + bestScore = score + first = false + } + } + if first { + return abilities + } + best := make([]Ability, 0, len(abilities)) + for _, ability := range abilities { + if score, ok := scores[ability.ChannelId]; ok && score == bestScore { + best = append(best, ability) + } + } + if len(best) == 0 { + return abilities + } + return best +} + // filterAbilitiesByRequestPathAndModel restricts candidates by request path and // model for the DB (non-memory-cache) selection path. Only Advanced Custom // (type 58) channels are path-checked: kept only when one of their routes matches diff --git a/model/analytics_event.go b/model/analytics_event.go new file mode 100644 index 000000000000..a218cd15f7df --- /dev/null +++ b/model/analytics_event.go @@ -0,0 +1,38 @@ +package model + +import ( + "github.com/QuantumNous/new-api/common" +) + +// AnalyticsEvent records first-party product/Marketing funnel events +// (visit, signup, pricing_click, lead_submit, ...). Lightweight, no PII beyond +// what the client voluntarily sends; used for conversion analytics (P1-06). +type AnalyticsEvent struct { + Id int64 `json:"id" gorm:"primaryKey;autoIncrement"` + Event string `json:"event" gorm:"type:varchar(32);index"` // visit|signup|pricing_click|lead_submit|... + Path string `json:"path" gorm:"type:varchar(512)"` + Locale string `json:"locale" gorm:"type:varchar(16)"` + Referrer string `json:"referrer" gorm:"type:varchar(512)"` + UserId int64 `json:"user_id" gorm:"index;default:0"` + CreatedAt int64 `json:"created_at" gorm:"bigint;index"` +} + +// Allowed analytics event names (defense against junk/abuse). +var allowedAnalyticsEvents = map[string]bool{ + "visit": true, + "signup": true, + "pricing_click": true, + "lead_submit": true, + "page_view": true, +} + +func IsValidAnalyticsEvent(name string) bool { + return allowedAnalyticsEvents[name] +} + +func CreateAnalyticsEvent(event *AnalyticsEvent) error { + if event.CreatedAt == 0 { + event.CreatedAt = common.GetTimestamp() + } + return DB.Create(event).Error +} diff --git a/model/channel_cache.go b/model/channel_cache.go index 86c594384d50..430aeab9195c 100644 --- a/model/channel_cache.go +++ b/model/channel_cache.go @@ -111,28 +111,41 @@ func SyncChannelCache(frequency int) { } } -func GetRandomSatisfiedChannel(group string, model string, retry int, requestPath string) (*Channel, error) { +// GetRandomSatisfiedChannel 选择一个满足条件的渠道。 +// routing 为区域路由解析结果：Active 时先按白名单收窄候选渠道， +// 并用区域策略打分替代默认的渠道优先级；未命中时行为与无区域路由完全一致。 +func GetRandomSatisfiedChannel(group string, model string, retry int, requestPath string, routing RegionRouting) (*Channel, error) { // if memory cache is disabled, get channel directly from database if !common.MemoryCacheEnabled { - return GetChannel(group, model, retry, requestPath) + return GetChannel(group, model, retry, requestPath, routing) } channelSyncLock.RLock() defer channelSyncLock.RUnlock() // First, try to find channels with the exact model name. - channels := filterChannelsByRequestPathAndModel(group2model2channels[group][model], requestPath, model) + channels := filterChannelsByRegion( + filterChannelsByRequestPathAndModel(group2model2channels[group][model], requestPath, model), routing) // If no channels found, try to find channels with the normalized model name. if len(channels) == 0 { normalizedModel := ratio_setting.FormatMatchingModelName(model) - channels = filterChannelsByRequestPathAndModel(group2model2channels[group][normalizedModel], requestPath, model) + channels = filterChannelsByRegion( + filterChannelsByRequestPathAndModel(group2model2channels[group][normalizedModel], requestPath, model), routing) } if len(channels) == 0 { return nil, nil } + // 区域策略生效时用策略分数替代渠道优先级参与分层与筛选 + scoreOf := func(channel *Channel) int64 { + if routing.Active && routing.Strategy != "" { + return regionStrategyScore(channel, routing.Strategy) + } + return channel.GetPriority() + } + if len(channels) == 1 { if channel, ok := channelsIDM[channels[0]]; ok { return channel, nil @@ -143,7 +156,7 @@ func GetRandomSatisfiedChannel(group string, model string, retry int, requestPat uniquePriorities := make(map[int]bool) for _, channelId := range channels { if channel, ok := channelsIDM[channelId]; ok { - uniquePriorities[int(channel.GetPriority())] = true + uniquePriorities[int(scoreOf(channel))] = true } else { return nil, fmt.Errorf("数据库一致性错误，渠道# %d 不存在，请联系管理员修复", channelId) } @@ -164,7 +177,7 @@ func GetRandomSatisfiedChannel(group string, model string, retry int, requestPat var targetChannels []*Channel for _, channelId := range channels { if channel, ok := channelsIDM[channelId]; ok { - if channel.GetPriority() == targetPriority { + if scoreOf(channel) == targetPriority { sumWeight += channel.GetWeight() targetChannels = append(targetChannels, channel) } @@ -236,6 +249,28 @@ func filterChannelsByRequestPathAndModel(channels []int, requestPath string, mod return filtered } +// filterChannelsByRegion 按区域路由白名单收窄候选渠道，保持原有优先级顺序。 +// 交集为空时返回原候选集：宁可让区域策略降级失效，也不因配置错误让请求整体失败。 +func filterChannelsByRegion(channels []int, routing RegionRouting) []int { + if !routing.Active || len(channels) == 0 || len(routing.AllowedIds) == 0 { + return channels + } + allowed := make(map[int]bool, len(routing.AllowedIds)) + for _, id := range routing.AllowedIds { + allowed[int(id)] = true + } + filtered := make([]int, 0, len(channels)) + for _, channelId := range channels { + if allowed[channelId] { + filtered = append(filtered, channelId) + } + } + if len(filtered) == 0 { + return channels + } + return filtered +} + func CacheGetChannel(id int) (*Channel, error) { if !common.MemoryCacheEnabled { return GetChannelById(id, true) diff --git a/model/distributor.go b/model/distributor.go new file mode 100644 index 000000000000..dd0daefdb3d9 --- /dev/null +++ b/model/distributor.go @@ -0,0 +1,194 @@ +package model + +import ( + "strings" + "time" + + "gorm.io/gorm" +) + +// Distributor 分销商（渠道商）实体，关联一个管理员账号（UserId）。 +type Distributor struct { + Id int64 `json:"id" gorm:"primaryKey;autoIncrement"` + UserId int64 `json:"user_id" gorm:"not null;uniqueIndex"` // 分销商管理员账号 + Name string `json:"name" gorm:"type:varchar(64);not null"` + Tier string `json:"tier" gorm:"type:varchar(16);default:'standard'"` // standard | gold | platinum + CommissionRate int `json:"commission_rate" gorm:"default:0"` // 佣金比例（百分比） + Status int `json:"status" gorm:"not null;default:1;index"` // 1 启用, 2 停用 + CreatedAt int64 `json:"created_at" gorm:"bigint;index"` + UpdatedAt int64 `json:"updated_at" gorm:"bigint"` +} + +// 分销商等级 / 状态白名单。 +var AllowedDistributorTiers = map[string]bool{ + "standard": true, + "gold": true, + "platinum": true, +} + +const ( + DistributorStatusActive = 1 + DistributorStatusDisabled = 2 +) + +var AllowedDistributorStatuses = map[int]bool{ + DistributorStatusActive: true, + DistributorStatusDisabled: true, +} + +// DistributorPrice 分销商价格覆盖（下级用户的模型售价）。 +type DistributorPrice struct { + Id int64 `json:"id" gorm:"primaryKey;autoIncrement"` + DistributorId int64 `json:"distributor_id" gorm:"not null;index"` + Model string `json:"model" gorm:"type:varchar(255);not null"` + InputPrice int64 `json:"input_price" gorm:"not null;default:0"` // 每 1M 单位最小货币单位 + OutputPrice int64 `json:"output_price" gorm:"not null;default:0"` // 每 1M 单位最小货币单位 + Currency string `json:"currency" gorm:"type:varchar(8);not null;default:'CNY'"` + Unit string `json:"unit" gorm:"type:varchar(16);default:'token'"` + CreatedAt int64 `json:"created_at" gorm:"bigint;index"` + UpdatedAt int64 `json:"updated_at" gorm:"bigint"` +} + +// 价格覆盖单位 / 货币白名单（复用市场模型约定）。 +var AllowedDistributorPriceUnits = map[string]bool{ + "token": true, + "image": true, + "second": true, + "char": true, +} + +var AllowedDistributorPriceCurrencies = map[string]bool{ + "CNY": true, + "USD": true, +} + +// CreateDistributor 创建分销商。 +func CreateDistributor(m *Distributor) error { + now := time.Now().Unix() + m.CreatedAt = now + m.UpdatedAt = now + return DB.Create(m).Error +} + +// GetDistributorById 按 id 获取分销商。 +func GetDistributorById(id int64) (*Distributor, error) { + var m Distributor + if err := DB.Where("id = ?", id).First(&m).Error; err != nil { + return nil, err + } + return &m, nil +} + +// GetDistributorByUserId 按关联账号获取分销商。 +func GetDistributorByUserId(userId int64) (*Distributor, error) { + var m Distributor + if err := DB.Where("user_id = ?", userId).First(&m).Error; err != nil { + return nil, err + } + return &m, nil +} + +// SearchDistributors 分页列出分销商；keyword 为空时不过滤。 +func SearchDistributors(page, pageSize int, keyword string) ([]*Distributor, int64, error) { + var items []*Distributor + var total int64 + q := DB.Model(&Distributor{}) + if strings.TrimSpace(keyword) != "" { + q = q.Where("name LIKE ?", "%"+keyword+"%") + } + if err := q.Count(&total).Error; err != nil { + return nil, 0, err + } + offset := (page - 1) * pageSize + if err := q.Order("id DESC").Offset(offset).Limit(pageSize).Find(&items).Error; err != nil { + return nil, 0, err + } + return items, total, nil +} + +// UpdateDistributor 更新分销商可编辑字段。 +func UpdateDistributor(m *Distributor) error { + updates := map[string]interface{}{ + "name": m.Name, + "tier": m.Tier, + "commission_rate": m.CommissionRate, + "status": m.Status, + "updated_at": time.Now().Unix(), + } + return DB.Model(&Distributor{}).Where("id = ?", m.Id).Updates(updates).Error +} + +// DeleteDistributor 删除分销商。 +func DeleteDistributor(id int64) error { + return DB.Transaction(func(tx *gorm.DB) error { + if err := tx.Where("distributor_id = ?", id).Delete(&DistributorPrice{}).Error; err != nil { + return err + } + return tx.Where("id = ?", id).Delete(&Distributor{}).Error + }) +} + +// CreateDistributorPrice 创建价格覆盖。 +func CreateDistributorPrice(m *DistributorPrice) error { + now := time.Now().Unix() + m.CreatedAt = now + m.UpdatedAt = now + return DB.Create(m).Error +} + +// SearchDistributorPrices 列出某分销商的全部价格覆盖。 +func SearchDistributorPrices(distributorId int64) ([]*DistributorPrice, error) { + var items []*DistributorPrice + if err := DB.Where("distributor_id = ?", distributorId).Order("id ASC").Find(&items).Error; err != nil { + return nil, err + } + return items, nil +} + +// UpdateDistributorPrice 更新价格覆盖。 +func UpdateDistributorPrice(m *DistributorPrice) error { + updates := map[string]interface{}{ + "model": m.Model, + "input_price": m.InputPrice, + "output_price": m.OutputPrice, + "currency": m.Currency, + "unit": m.Unit, + "updated_at": time.Now().Unix(), + } + return DB.Model(&DistributorPrice{}).Where("id = ?", m.Id).Updates(updates).Error +} + +// DeleteDistributorPrice 删除价格覆盖。 +func DeleteDistributorPrice(id int64) error { + return DB.Where("id = ?", id).Delete(&DistributorPrice{}).Error +} + +// DistributorBilling 分销商下级账单汇总（基于下级用户额度近似）。 +type DistributorBilling struct { + DistributorId int64 `json:"distributor_id"` + SubUserCount int64 `json:"sub_user_count"` + Allocated int64 `json:"allocated"` + Used int64 `json:"used"` +} + +// GetDistributorBilling 汇总下级用户（直接邀请）的额度与用量。 +func GetDistributorBilling(distributorUserId int64) (*DistributorBilling, error) { + var subUserIds []int64 + if err := DB.Model(&User{}).Where("inviter_id = ?", distributorUserId).Pluck("id", &subUserIds).Error; err != nil { + return nil, err + } + billing := &DistributorBilling{DistributorId: distributorUserId, SubUserCount: int64(len(subUserIds))} + if len(subUserIds) == 0 { + return billing, nil + } + var allocated, used int64 + if err := DB.Model(&User{}). + Where("id IN ?", subUserIds). + Select("COALESCE(SUM(quota),0), COALESCE(SUM(used_quota),0)"). + Row().Scan(&allocated, &used); err != nil { + return nil, err + } + billing.Allocated = allocated + billing.Used = used + return billing, nil +} diff --git a/model/log.go b/model/log.go index 1d2b38fc7c1c..abc4c6c29754 100644 --- a/model/log.go +++ b/model/log.go @@ -8,6 +8,7 @@ import ( "time" "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/logger" "github.com/QuantumNous/new-api/types" @@ -74,6 +75,7 @@ type Log struct { ChannelName string `json:"channel_name" gorm:"->"` TokenId int `json:"token_id" gorm:"default:0;index"` Group string `json:"group" gorm:"index"` + TeamId int64 `json:"team_id" gorm:"type:bigint;index:idx_logs_team_id;default:0"` // 所属企业团队，用于团队级用量/计费聚合 Ip string `json:"ip" gorm:"index;default:''"` RequestId string `json:"request_id,omitempty" gorm:"type:varchar(64);index:idx_logs_request_id;default:''"` UpstreamRequestId string `json:"upstream_request_id,omitempty" gorm:"type:varchar(128);index:idx_logs_upstream_request_id;default:''"` @@ -373,6 +375,12 @@ func RecordConsumeLog(c *gin.Context, userId int, params RecordConsumeLogParams) UseTime: params.UseTimeSeconds, IsStream: params.IsStream, Group: params.Group, + TeamId: func() int64 { + if v, ok := common.GetContextKeyType[int64](c, constant.ContextKeyUserTeamId); ok { + return v + } + return 0 + }(), Ip: func() string { if needRecordIp { return c.ClientIP() diff --git a/model/main.go b/model/main.go index 21445593e54e..1b89f55299c0 100644 --- a/model/main.go +++ b/model/main.go @@ -292,10 +292,25 @@ func migrateDB() error { &SystemTaskLock{}, &CasbinRule{}, &AuthzRole{}, + &MarketModel{}, + // P2 平台化能力：企业团队空间 / SLA / 区域路由 / 分销商 + &Team{}, + &TeamMember{}, + &TeamProject{}, + &SlaIncident{}, + &RegionRoute{}, + &Distributor{}, + &DistributorPrice{}, + // MVP 营销站公开数据模型 + &SalesLead{}, + &PublicPricing{}, + &PublicModelCategory{}, ) if err != nil { return err } + // 写入营销站默认定价与模型目录（仅在表为空时） + InitPublicSiteDefaults() if err := InitializeUserAuthVersions(); err != nil { return err } diff --git a/model/market_model.go b/model/market_model.go new file mode 100644 index 000000000000..23220e49d5a4 --- /dev/null +++ b/model/market_model.go @@ -0,0 +1,119 @@ +package model + +import ( + "time" +) + +// MarketModel 模型商店（Model Market）的商品条目：面向客户的单模型商业上架信息。 +// 与内部计费表 Pricing（成本比例，随默认刷新）解耦——此处为管理员可维护的“门店价格”， +// 同时关联实际可路由的模型名（= Pricing.ModelName）与营销分类（= PublicModelCategory.Category）。 +type MarketModel struct { + Id int64 `json:"id" gorm:"primaryKey;autoIncrement"` + Model string `json:"model" gorm:"type:varchar(255);not null;uniqueIndex"` // 实际模型名，= Pricing.ModelName + Provider string `json:"provider" gorm:"type:varchar(64);index"` // OpenAI / Anthropic / ... + Category string `json:"category" gorm:"type:varchar(32);not null;index"` // = PublicModelCategory.Category + Tags string `json:"tags" gorm:"type:varchar(255)"` // vision,reasoning,streaming + InputPrice int64 `json:"input_price" gorm:"not null;default:0"` // 客户价：每 1M 个计价单位的最小货币单位（如 CNY 分 / USD 美分），见 Currency + OutputPrice int64 `json:"output_price" gorm:"not null;default:0"` // 客户价：同上，输出侧 + Currency string `json:"currency" gorm:"type:varchar(8);not null;default:'CNY'"` // CNY / USD，计价货币 + Unit string `json:"unit" gorm:"type:varchar(16);default:'token'"` // token|image|second|char（计价数量单位） + Metadata string `json:"metadata" gorm:"type:text"` // JSON：按 locale 的展示覆盖，如 {"zh":{"name":...},"en":{"name":...}} + TrialQuota int64 `json:"trial_quota" gorm:"default:0"` // 首次激活赠送的试用额度（v1 仅展示，未启用激活） + Status int `json:"status" gorm:"not null;default:1;index"` // 1 available, 2 coming_soon, 3 disabled + Featured bool `json:"featured" gorm:"default:false;index"` + Sort int `json:"sort" gorm:"default:0"` + CreatedAt int64 `json:"created_at" gorm:"bigint;index"` + UpdatedAt int64 `json:"updated_at" gorm:"bigint"` +} + +// MarketModel 状态常量。 +const ( + MarketModelStatusAvailable = 1 + MarketModelStatusComingSoon = 2 + MarketModelStatusDisabled = 3 +) + +// AllowedMarketModelStatuses 上架状态白名单。 +var AllowedMarketModelStatuses = map[int]bool{ + MarketModelStatusAvailable: true, + MarketModelStatusComingSoon: true, + MarketModelStatusDisabled: true, +} + +// AllowedMarketModelUnits 计价单位白名单。 +var AllowedMarketModelUnits = map[string]bool{ + "token": true, + "image": true, + "second": true, + "char": true, +} + +// AllowedMarketModelCurrencies 计价货币白名单。 +var AllowedMarketModelCurrencies = map[string]bool{ + "CNY": true, + "USD": true, +} + +// CreateMarketModel 写入一条模型上架记录（由调用方负责填充时间戳与默认值）。 +func CreateMarketModel(m *MarketModel) error { + now := time.Now().Unix() + m.CreatedAt = now + m.UpdatedAt = now + return DB.Create(m).Error +} + +// GetMarketModelById 按 id 获取单条上架记录。 +func GetMarketModelById(id int64) (*MarketModel, error) { + var m MarketModel + err := DB.Where("id = ?", id).First(&m).Error + if err != nil { + return nil, err + } + return &m, nil +} + +// SearchMarketModels 列出上架模型；status/category 为空时不过滤。按 category、sort 升序。 +func SearchMarketModels(status, category string) ([]*MarketModel, error) { + var items []*MarketModel + q := DB.Order("category ASC").Order("sort ASC") + if status != "" { + q = q.Where("status = ?", status) + } + if category != "" { + q = q.Where("category = ?", category) + } + err := q.Find(&items).Error + return items, err +} + +// UpdateMarketModel 更新上架记录的可编辑字段（Model 为唯一键，不可变）。 +func UpdateMarketModel(m *MarketModel) error { + updates := map[string]interface{}{ + "provider": m.Provider, + "category": m.Category, + "tags": m.Tags, + "input_price": m.InputPrice, + "output_price": m.OutputPrice, + "currency": m.Currency, + "unit": m.Unit, + "metadata": m.Metadata, + "trial_quota": m.TrialQuota, + "status": m.Status, + "featured": m.Featured, + "sort": m.Sort, + "updated_at": time.Now().Unix(), + } + return DB.Model(&MarketModel{}).Where("id = ?", m.Id).Updates(updates).Error +} + +// DeleteMarketModel 按 id 删除上架记录。 +func DeleteMarketModel(id int64) error { + return DB.Where("id = ?", id).Delete(&MarketModel{}).Error +} + +// CountMarketModelsByModel 判断某模型名是否已上架（用于唯一性校验）。 +func CountMarketModelsByModel(model string) (int64, error) { + var cnt int64 + err := DB.Model(&MarketModel{}).Where("model = ?", model).Count(&cnt).Error + return cnt, err +} diff --git a/model/option.go b/model/option.go index e7fda5231be7..7ab64e0ded80 100644 --- a/model/option.go +++ b/model/option.go @@ -120,7 +120,6 @@ func InitOptionMap() { common.OptionMap["Chats"] = setting.Chats2JsonString() common.OptionMap["AutoGroups"] = setting.AutoGroups2JsonString() common.OptionMap["DefaultUseAutoGroup"] = strconv.FormatBool(setting.DefaultUseAutoGroup) - common.OptionMap["MaxTokenAutoGroups"] = strconv.Itoa(setting.GetMaxTokenAutoGroups()) common.OptionMap["PayMethods"] = operation_setting.PayMethods2JsonString() common.OptionMap["GitHubClientId"] = "" common.OptionMap["GitHubClientSecret"] = "" @@ -209,9 +208,6 @@ func validateOptionValue(key string, value string) error { if key == operation_setting.ToolPriceOptionKey { return operation_setting.ValidateToolPricesJSON(value) } - if key == "MaxTokenAutoGroups" { - return setting.ValidateMaxTokenAutoGroups(value) - } return nil } @@ -417,8 +413,6 @@ func updateOptionMap(key string, value string) (err error) { err = setting.UpdateChatsByJsonString(value) case "AutoGroups": err = setting.UpdateAutoGroupsByJsonString(value) - case "MaxTokenAutoGroups": - err = setting.UpdateMaxTokenAutoGroups(value) case "CustomCallbackAddress": operation_setting.CustomCallbackAddress = value case "EpayId": diff --git a/model/public_model_category.go b/model/public_model_category.go new file mode 100644 index 000000000000..0b65a34c968d --- /dev/null +++ b/model/public_model_category.go @@ -0,0 +1,15 @@ +package model + +// PublicModelCategory 营销站公开模型目录（按 locale 提供中英双语）。 +type PublicModelCategory struct { + Id uint `gorm:"primaryKey" json:"id"` + Category string `gorm:"type:varchar(32);not null;index" json:"category"` // chinese/global/image/video/audio/embedding + Locale string `gorm:"type:varchar(8);not null;index" json:"locale"` + Title string `gorm:"type:varchar(128);not null" json:"title"` + Description string `gorm:"type:text" json:"description"` + Models string `gorm:"type:text" json:"models"` // JSON 数组: [{name,capability_tags,note}] + Sort int `gorm:"default:0" json:"sort"` + Enabled bool `gorm:"default:true" json:"enabled"` +} + +func (PublicModelCategory) TableName() string { return "public_model_categories" } diff --git a/model/public_pricing.go b/model/public_pricing.go new file mode 100644 index 000000000000..bcc82218a41f --- /dev/null +++ b/model/public_pricing.go @@ -0,0 +1,17 @@ +package model + +// PublicPricing 营销站公开定价方案（按 locale 提供中英双语）。 +type PublicPricing struct { + Id uint `gorm:"primaryKey" json:"id"` + PlanKey string `gorm:"type:varchar(32);not null;index" json:"plan_key"` + Locale string `gorm:"type:varchar(8);not null;index" json:"locale"` + Title string `gorm:"type:varchar(128);not null" json:"title"` + Description string `gorm:"type:text" json:"description"` + BillingMode string `gorm:"type:varchar(32)" json:"billing_mode"` // payg/subscription/custom + PriceText string `gorm:"type:varchar(128)" json:"price_text"` + Features string `gorm:"type:text" json:"features"` // JSON 数组字符串 + Sort int `gorm:"default:0" json:"sort"` + Enabled bool `gorm:"default:true" json:"enabled"` +} + +func (PublicPricing) TableName() string { return "public_pricings" } diff --git a/model/public_site_seed.go b/model/public_site_seed.go new file mode 100644 index 000000000000..fdf1bc4a7722 --- /dev/null +++ b/model/public_site_seed.go @@ -0,0 +1,106 @@ +package model + +import "encoding/json" + +// InitPublicSiteDefaults 在首次启动（表为空）时写入营销站默认的定价方案与模型目录， +// 提供中英双语。仅在表为空时插入，避免重复写入。 +func InitPublicSiteDefaults() { + if DB == nil { + return + } + var pricingCount int64 + if err := DB.Model(&PublicPricing{}).Count(&pricingCount).Error; err != nil { + return + } + if pricingCount == 0 { + pricings := defaultPublicPricings() + _ = DB.Create(&pricings).Error + } + + var catCount int64 + if err := DB.Model(&PublicModelCategory{}).Count(&catCount).Error; err != nil { + return + } + if catCount == 0 { + cats := defaultPublicModelCategories() + _ = DB.Create(&cats).Error + } +} + +func defaultPublicPricings() []PublicPricing { + paygFeatures, _ := json.Marshal([]string{ + "按实际 token 用量计费，无月费", + "全模型统一接入，无需逐家签约", + "即时开通，按量结算", + }) + proFeatures, _ := json.Marshal([]string{ + "每月固定额度，单价更优", + "优先队列与更高并发", + "团队子账号与用量看板", + }) + entFeatures, _ := json.Marshal([]string{ + "专属 SLA 与区域路由", + "定制模型与私有部署", + "专属客户成功经理", + }) + return []PublicPricing{ + {PlanKey: "payg", Locale: "en", Title: "Pay as you go", Description: "Token-based pricing, no monthly fee.", BillingMode: "payg", PriceText: "Per-token", Features: string(paygFeatures), Sort: 1, Enabled: true}, + {PlanKey: "payg", Locale: "zh", Title: "按量计费", Description: "按实际 token 用量计费，无月费。", BillingMode: "payg", PriceText: "按 token", Features: string(paygFeatures), Sort: 1, Enabled: true}, + {PlanKey: "pro", Locale: "en", Title: "Pro", Description: "Monthly quota with better unit price.", BillingMode: "subscription", PriceText: "$49 / month", Features: string(proFeatures), Sort: 2, Enabled: true}, + {PlanKey: "pro", Locale: "zh", Title: "专业版", Description: "每月固定额度，单价更优。", BillingMode: "subscription", PriceText: "￥299 / 月", Features: string(proFeatures), Sort: 2, Enabled: true}, + {PlanKey: "enterprise", Locale: "en", Title: "Enterprise", Description: "Dedicated SLA, regional routing, custom models.", BillingMode: "custom", PriceText: "Contact sales", Features: string(entFeatures), Sort: 3, Enabled: true}, + {PlanKey: "enterprise", Locale: "zh", Title: "企业版", Description: "专属 SLA、区域路由与定制模型。", BillingMode: "custom", PriceText: "联系销售", Features: string(entFeatures), Sort: 3, Enabled: true}, + } +} + +func defaultPublicModelCategories() []PublicModelCategory { + chinese, _ := json.Marshal([]map[string]string{ + {"name": "DeepSeek", "capability_tags": "reasoning, chat", "note": "国产强推理模型"}, + {"name": "Qwen", "capability_tags": "chat, vision", "note": "通义千问全系列"}, + {"name": "GLM", "capability_tags": "chat, agent", "note": "智谱 GLM 系列"}, + {"name": "Doubao", "capability_tags": "chat, vision", "note": "字节豆包"}, + {"name": "Kimi", "capability_tags": "long-context", "note": "月之暗面长上下文"}, + {"name": "Hunyuan", "capability_tags": "chat", "note": "腾讯混元"}, + }) + global, _ := json.Marshal([]map[string]string{ + {"name": "OpenAI", "capability_tags": "chat, vision, audio", "note": "GPT 全系列"}, + {"name": "Anthropic", "capability_tags": "chat, agent", "note": "Claude 全系列"}, + {"name": "Google", "capability_tags": "chat, vision", "note": "Gemini 全系列"}, + {"name": "Meta", "capability_tags": "chat", "note": "Llama 开源系列"}, + {"name": "Mistral", "capability_tags": "chat", "note": "欧洲开源模型"}, + }) + image, _ := json.Marshal([]map[string]string{ + {"name": "DALL·E", "capability_tags": "image", "note": "OpenAI 文生图"}, + {"name": "Midjourney", "capability_tags": "image", "note": "艺术风格图像"}, + {"name": "Stable Diffusion", "capability_tags": "image", "note": "开源文生图"}, + {"name": "Flux", "capability_tags": "image", "note": "高质感图像"}, + }) + video, _ := json.Marshal([]map[string]string{ + {"name": "Runway", "capability_tags": "video", "note": "文生视频"}, + {"name": "Kling", "capability_tags": "video", "note": "可灵文生视频"}, + {"name": "Sora", "capability_tags": "video", "note": "OpenAI 视频生成"}, + }) + audio, _ := json.Marshal([]map[string]string{ + {"name": "Whisper", "capability_tags": "audio", "note": "语音识别"}, + {"name": "TTS", "capability_tags": "audio", "note": "语音合成"}, + {"name": "Suno", "capability_tags": "audio", "note": "音乐生成"}, + }) + embedding, _ := json.Marshal([]map[string]string{ + {"name": "text-embedding", "capability_tags": "embedding", "note": "文本向量"}, + {"name": "bge", "capability_tags": "embedding", "note": "开源嵌入模型"}, + }) + return []PublicModelCategory{ + {Category: "chinese", Locale: "en", Title: "Chinese LLMs", Description: "Top open & commercial large models from China.", Models: string(chinese), Sort: 1, Enabled: true}, + {Category: "chinese", Locale: "zh", Title: "中国大模型", Description: "国内头部开源与商业大模型统一接入。", Models: string(chinese), Sort: 1, Enabled: true}, + {Category: "global", Locale: "en", Title: "Global LLMs", Description: "Leading models from OpenAI, Anthropic, Google and more.", Models: string(global), Sort: 2, Enabled: true}, + {Category: "global", Locale: "zh", Title: "海外大模型", Description: "OpenAI、Anthropic、Google 等全球领先模型。", Models: string(global), Sort: 2, Enabled: true}, + {Category: "image", Locale: "en", Title: "Image Generation", Description: "Text-to-image models for product & creative use.", Models: string(image), Sort: 3, Enabled: true}, + {Category: "image", Locale: "zh", Title: "图像生成", Description: "面向产品与创作的文生图模型。", Models: string(image), Sort: 3, Enabled: true}, + {Category: "video", Locale: "en", Title: "Video Generation", Description: "Text-to-video for marketing and storytelling.", Models: string(video), Sort: 4, Enabled: true}, + {Category: "video", Locale: "zh", Title: "视频生成", Description: "面向营销与叙事的文生视频。", Models: string(video), Sort: 4, Enabled: true}, + {Category: "audio", Locale: "en", Title: "Audio", Description: "Speech recognition, TTS and music generation.", Models: string(audio), Sort: 5, Enabled: true}, + {Category: "audio", Locale: "zh", Title: "语音与音乐", Description: "语音识别、语音合成与音乐生成。", Models: string(audio), Sort: 5, Enabled: true}, + {Category: "embedding", Locale: "en", Title: "Embeddings", Description: "Vector embeddings for RAG & search.", Models: string(embedding), Sort: 6, Enabled: true}, + {Category: "embedding", Locale: "zh", Title: "向量嵌入", Description: "面向 RAG 与检索的向量嵌入。", Models: string(embedding), Sort: 6, Enabled: true}, + } +} diff --git a/model/region_route.go b/model/region_route.go new file mode 100644 index 000000000000..72f574a1ffc7 --- /dev/null +++ b/model/region_route.go @@ -0,0 +1,281 @@ +package model + +import ( + "strconv" + "strings" + "sync" + "time" + + "github.com/QuantumNous/new-api/common" +) + +// RegionRoute 区域路由策略配置。 +// 由 ResolveRegionRouting 解析后接入 model/channel_cache.go、model/ability.go 的选渠道链路。 +type RegionRoute struct { + Id int64 `json:"id" gorm:"primaryKey;autoIncrement"` + Region string `json:"region" gorm:"type:varchar(16);not null;index"` // cn | us | eu | global ... + Model string `json:"model" gorm:"type:varchar(64);not null;default:'*'"` // '*' 表示全部模型 + ChannelIds string `json:"channel_ids" gorm:"type:varchar(512)"` // 逗号分隔的渠道 id 列表 + Tag string `json:"tag" gorm:"type:varchar(64)"` // 或按 tag 选择渠道 + Strategy string `json:"strategy" gorm:"type:varchar(16);not null;default:'availability'"` // cost | latency | availability | fixed + Priority int `json:"priority" gorm:"default:0"` + Weight int `json:"weight" gorm:"default:0"` + // 不能使用 gorm default:true：GORM 在 Create 时会忽略零值字段，导致显式禁用的策略被写成启用。 + // 缺省启用由 controller 层负责（请求未带 enabled 时置 true）。 + Enabled bool `json:"enabled" gorm:"index"` + CreatedAt int64 `json:"created_at" gorm:"bigint;index"` + UpdatedAt int64 `json:"updated_at" gorm:"bigint"` +} + +// 区域路由策略白名单。 +var AllowedRegionRouteStrategies = map[string]bool{ + "cost": true, + "latency": true, + "availability": true, + "fixed": true, +} + +// CreateRegionRoute 创建区域路由策略。 +func CreateRegionRoute(m *RegionRoute) error { + now := time.Now().Unix() + m.CreatedAt = now + m.UpdatedAt = now + err := DB.Create(m).Error + if err == nil { + InvalidateRegionRoutingCache() + } + return err +} + +// GetRegionRouteById 按 id 获取策略。 +func GetRegionRouteById(id int64) (*RegionRoute, error) { + var m RegionRoute + if err := DB.Where("id = ?", id).First(&m).Error; err != nil { + return nil, err + } + return &m, nil +} + +// SearchRegionRoutes 分页列出策略；region/model 为空时不过滤。 +func SearchRegionRoutes(page, pageSize int, region, model string) ([]*RegionRoute, int64, error) { + var items []*RegionRoute + var total int64 + q := DB.Model(&RegionRoute{}) + if strings.TrimSpace(region) != "" { + q = q.Where("region = ?", region) + } + if strings.TrimSpace(model) != "" { + q = q.Where("model = ?", model) + } + if err := q.Count(&total).Error; err != nil { + return nil, 0, err + } + offset := (page - 1) * pageSize + if err := q.Order("id DESC").Offset(offset).Limit(pageSize).Find(&items).Error; err != nil { + return nil, 0, err + } + return items, total, nil +} + +// UpdateRegionRoute 更新策略可编辑字段。 +func UpdateRegionRoute(m *RegionRoute) error { + updates := map[string]interface{}{ + "region": m.Region, + "model": m.Model, + "channel_ids": m.ChannelIds, + "tag": m.Tag, + "strategy": m.Strategy, + "priority": m.Priority, + "weight": m.Weight, + "enabled": m.Enabled, + "updated_at": time.Now().Unix(), + } + err := DB.Model(&RegionRoute{}).Where("id = ?", m.Id).Updates(updates).Error + if err == nil { + InvalidateRegionRoutingCache() + } + return err +} + +// DeleteRegionRoute 删除策略。 +func DeleteRegionRoute(id int64) error { + err := DB.Where("id = ?", id).Delete(&RegionRoute{}).Error + if err == nil { + InvalidateRegionRoutingCache() + } + return err +} + +// GetEnabledRegionRoutes 取某区域下启用的策略。 +// region 同时匹配 'global' 兜底策略；model 支持具体模型或 '*'（通配）匹配。 +func GetEnabledRegionRoutes(region, model string) ([]*RegionRoute, error) { + var items []*RegionRoute + err := DB.Where("enabled = ? AND region IN (?) AND (model = ? OR model = '*')", + true, []string{region, RegionGlobal}, model). + Order("priority DESC").Order("id ASC").Find(&items).Error + if err != nil { + return nil, err + } + return items, nil +} + +// RegionGlobal 是兜底区域标识：配置在该区域下的策略对所有区域生效。 +const RegionGlobal = "global" + +// maxRegionLength 限制区域标识长度，避免异常请求头把超长字符串带进 SQL 查询。 +const maxRegionLength = 16 + +// NormalizeRegion 归一化区域标识：去空格、转小写、截断超长值。 +func NormalizeRegion(region string) string { + region = strings.ToLower(strings.TrimSpace(region)) + if len(region) > maxRegionLength { + region = region[:maxRegionLength] + } + return region +} + +// RegionRouting 是区域路由的解析结果，供选渠道链路使用。 +type RegionRouting struct { + // Active 为 false 时调用方应完全退回默认选渠道逻辑。 + Active bool + // Region 归一化后的区域标识。 + Region string + // Strategy 命中的排序策略（cost | latency | availability | fixed），可能为空。 + Strategy string + // AllowedIds 允许使用的渠道 id 白名单，按配置顺序去重。 + AllowedIds []int64 +} + +// 区域路由解析结果缓存。选渠道位于转发热路径，逐请求查库不可接受， +// 因此在开启内存缓存时按 region|model 缓存解析结果，TTL 内管理端改动最多延迟生效 regionRoutingCacheTTL。 +type regionRoutingCacheEntry struct { + routing RegionRouting + expiresAt time.Time +} + +const regionRoutingCacheTTL = 60 * time.Second + +var regionRoutingCache sync.Map // string -> regionRoutingCacheEntry + +// ResolveRegionRouting 解析某区域 + 模型命中的路由策略，返回渠道白名单与排序策略。 +// region 为空、无启用策略、或策略未圈定任何渠道时返回 Active=false， +// 此时调用方保持原有选渠道行为不变。 +func ResolveRegionRouting(region, modelName string) RegionRouting { + region = NormalizeRegion(region) + if region == "" || DB == nil { + return RegionRouting{} + } + if !common.MemoryCacheEnabled { + return resolveRegionRouting(region, modelName) + } + key := region + "|" + modelName + if cached, ok := regionRoutingCache.Load(key); ok { + if entry, ok := cached.(regionRoutingCacheEntry); ok && time.Now().Before(entry.expiresAt) { + return entry.routing + } + } + routing := resolveRegionRouting(region, modelName) + regionRoutingCache.Store(key, regionRoutingCacheEntry{ + routing: routing, + expiresAt: time.Now().Add(regionRoutingCacheTTL), + }) + return routing +} + +// InvalidateRegionRoutingCache 清空区域路由解析缓存，管理端增删改策略后调用。 +func InvalidateRegionRoutingCache() { + regionRoutingCache.Range(func(key, _ any) bool { + regionRoutingCache.Delete(key) + return true + }) +} + +func resolveRegionRouting(region, modelName string) RegionRouting { + result := RegionRouting{} + routes, err := GetEnabledRegionRoutes(region, modelName) + if err != nil || len(routes) == 0 { + return result + } + + seen := make(map[int64]bool) + ids := make([]int64, 0, len(routes)) + tags := make([]string, 0, len(routes)) + strategy := "" + for _, r := range routes { + // routes 已按 priority DESC, id ASC 排序，首个合法策略即最高优先级策略 + if strategy == "" && AllowedRegionRouteStrategies[r.Strategy] { + strategy = r.Strategy + } + for _, id := range r.ChannelIdsAsSlice() { + if !seen[id] { + seen[id] = true + ids = append(ids, id) + } + } + if tag := strings.TrimSpace(r.Tag); tag != "" { + tags = append(tags, tag) + } + } + + if len(tags) > 0 { + var tagged []int64 + if err := DB.Model(&Channel{}).Where("tag IN (?)", tags).Pluck("id", &tagged).Error; err == nil { + for _, id := range tagged { + if !seen[id] { + seen[id] = true + ids = append(ids, id) + } + } + } + } + + if len(ids) == 0 { + return result + } + result.Active = true + result.Region = region + result.Strategy = strategy + result.AllowedIds = ids + return result +} + +// regionStrategyScore 按区域策略给渠道打分，分数越大越优先被选中。 +// 返回值直接替代默认的 channel.GetPriority()，因此需与优先级同量纲（int64）。 +func regionStrategyScore(channel *Channel, strategy string) int64 { + if channel == nil { + return 0 + } + switch strategy { + case "latency": + // 响应时间越短越优先 + return -int64(channel.ResponseTime) + case "availability": + // 可用渠道优先，其次比较响应时间 + var score int64 + if channel.Status == common.ChannelStatusEnabled { + score = 1_000_000_000 + } + return score - int64(channel.ResponseTime) + case "cost": + // 约定：优先级越低的渠道为越便宜的备用渠道，因此反向取优先级 + return -channel.GetPriority() + default: + // fixed / 未知策略：沿用渠道自身优先级 + return channel.GetPriority() + } +} + +// ChannelIdsAsSlice 将逗号分隔的渠道 id 解析为 int64 切片。 +func (r *RegionRoute) ChannelIdsAsSlice() []int64 { + ids := make([]int64, 0) + for _, s := range strings.Split(r.ChannelIds, ",") { + s = strings.TrimSpace(s) + if s == "" { + continue + } + if id, err := strconv.ParseInt(s, 10, 64); err == nil { + ids = append(ids, id) + } + } + return ids +} diff --git a/model/region_route_test.go b/model/region_route_test.go new file mode 100644 index 000000000000..0eb8df40388f --- /dev/null +++ b/model/region_route_test.go @@ -0,0 +1,157 @@ +package model + +import ( + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func setupRegionRouteTest(t *testing.T) { + t.Helper() + require.NoError(t, DB.AutoMigrate(&RegionRoute{})) + require.NoError(t, DB.Exec("DELETE FROM region_routes").Error) + require.NoError(t, DB.Exec("DELETE FROM channels").Error) + InvalidateRegionRoutingCache() + t.Cleanup(func() { + DB.Exec("DELETE FROM region_routes") + DB.Exec("DELETE FROM channels") + InvalidateRegionRoutingCache() + }) +} + +func newTestChannel(t *testing.T, id int, tag string, priority int64, responseTime int, status int) *Channel { + t.Helper() + channel := &Channel{ + Id: id, + Name: tag, + Key: "sk-test", + Status: status, + Priority: &priority, + ResponseTime: responseTime, + } + if tag != "" { + channel.Tag = &tag + } + return channel +} + +func TestNormalizeRegion(t *testing.T) { + assert.Equal(t, "cn", NormalizeRegion(" CN ")) + assert.Equal(t, "", NormalizeRegion(" ")) + assert.Equal(t, "ap-southeast-1", NormalizeRegion("AP-Southeast-1")) + // 超长输入被截断到 maxRegionLength，避免异常请求头进入 SQL + assert.Len(t, NormalizeRegion("abcdefghijklmnopqrstuvwxyz"), maxRegionLength) +} + +func TestResolveRegionRoutingInactiveCases(t *testing.T) { + setupRegionRouteTest(t) + + // 无区域标识 + assert.False(t, ResolveRegionRouting("", "gpt-4").Active) + // 有区域但无任何策略 + assert.False(t, ResolveRegionRouting("cn", "gpt-4").Active) + + // 策略未圈定任何渠道 -> 不生效 + require.NoError(t, CreateRegionRoute(&RegionRoute{ + Region: "cn", Model: "*", Strategy: "latency", Enabled: true, + })) + assert.False(t, ResolveRegionRouting("cn", "gpt-4").Active) + + // 策略被禁用 -> 不生效 + require.NoError(t, CreateRegionRoute(&RegionRoute{ + Region: "cn", Model: "*", ChannelIds: "1,2", Strategy: "latency", Enabled: false, + })) + assert.False(t, ResolveRegionRouting("cn", "gpt-4").Active) +} + +func TestResolveRegionRoutingMergesChannelsAndPicksStrategy(t *testing.T) { + setupRegionRouteTest(t) + + require.NoError(t, DB.Create(newTestChannel(t, 5, "cn-pool", 0, 100, common.ChannelStatusEnabled)).Error) + require.NoError(t, DB.Create(newTestChannel(t, 9, "us-pool", 0, 100, common.ChannelStatusEnabled)).Error) + + // priority 最高的策略决定排序策略 + require.NoError(t, CreateRegionRoute(&RegionRoute{ + Region: "cn", Model: "*", ChannelIds: "1,2", Strategy: "latency", Priority: 10, Enabled: true, + })) + // global 策略对所有区域生效，渠道 id 与上一条去重合并 + require.NoError(t, CreateRegionRoute(&RegionRoute{ + Region: RegionGlobal, Model: "gpt-4", ChannelIds: "2,3", Strategy: "cost", Priority: 1, Enabled: true, + })) + // 按 tag 圈定渠道 + require.NoError(t, CreateRegionRoute(&RegionRoute{ + Region: "cn", Model: "*", Tag: "cn-pool", Strategy: "fixed", Priority: 0, Enabled: true, + })) + // 其他区域的策略不应命中 + require.NoError(t, CreateRegionRoute(&RegionRoute{ + Region: "us", Model: "*", Tag: "us-pool", Strategy: "fixed", Priority: 99, Enabled: true, + })) + + routing := ResolveRegionRouting(" CN ", "gpt-4") + require.True(t, routing.Active) + assert.Equal(t, "cn", routing.Region) + assert.Equal(t, "latency", routing.Strategy) + assert.ElementsMatch(t, []int64{1, 2, 3, 5}, routing.AllowedIds) + + // 模型不匹配 global 那条时，只保留通配策略的渠道 + routing = ResolveRegionRouting("cn", "claude-3") + require.True(t, routing.Active) + assert.ElementsMatch(t, []int64{1, 2, 5}, routing.AllowedIds) +} + +func TestRegionStrategyScore(t *testing.T) { + fast := newTestChannel(t, 1, "", 5, 100, common.ChannelStatusEnabled) + slow := newTestChannel(t, 2, "", 9, 800, common.ChannelStatusEnabled) + down := newTestChannel(t, 3, "", 9, 50, common.ChannelStatusAutoDisabled) + + // latency：响应时间短的分数高 + assert.Greater(t, regionStrategyScore(fast, "latency"), regionStrategyScore(slow, "latency")) + // availability：可用渠道永远优先于不可用渠道，即使后者更快 + assert.Greater(t, regionStrategyScore(slow, "availability"), regionStrategyScore(down, "availability")) + // cost：优先级低的渠道视为低成本备用渠道，分数更高 + assert.Greater(t, regionStrategyScore(fast, "cost"), regionStrategyScore(slow, "cost")) + // fixed / 未知策略：沿用渠道自身优先级 + assert.Equal(t, int64(9), regionStrategyScore(slow, "fixed")) + assert.Equal(t, int64(9), regionStrategyScore(slow, "")) + assert.Equal(t, int64(0), regionStrategyScore(nil, "latency")) +} + +func TestFilterChannelsByRegion(t *testing.T) { + candidates := []int{1, 2, 3} + + // 未命中区域路由时原样返回 + assert.Equal(t, candidates, filterChannelsByRegion(candidates, RegionRouting{})) + + // 命中时按白名单收窄，且保持原有顺序 + routing := RegionRouting{Active: true, AllowedIds: []int64{3, 1}} + assert.Equal(t, []int{1, 3}, filterChannelsByRegion(candidates, routing)) + + // 交集为空时降级为不过滤，避免请求整体失败 + routing = RegionRouting{Active: true, AllowedIds: []int64{99}} + assert.Equal(t, candidates, filterChannelsByRegion(candidates, routing)) +} + +func TestFilterAbilitiesByRegion(t *testing.T) { + setupRegionRouteTest(t) + + require.NoError(t, DB.Create(newTestChannel(t, 1, "", 0, 500, common.ChannelStatusEnabled)).Error) + require.NoError(t, DB.Create(newTestChannel(t, 2, "", 0, 50, common.ChannelStatusEnabled)).Error) + require.NoError(t, DB.Create(newTestChannel(t, 3, "", 0, 10, common.ChannelStatusEnabled)).Error) + + abilities := []Ability{{ChannelId: 1}, {ChannelId: 2}, {ChannelId: 3}} + + // 未命中区域路由时原样返回 + assert.Len(t, filterAbilitiesByRegion(abilities, RegionRouting{}), 3) + + // 白名单收窄 + latency 策略只保留最快的渠道 + routing := RegionRouting{Active: true, AllowedIds: []int64{1, 2}, Strategy: "latency"} + filtered := filterAbilitiesByRegion(abilities, routing) + require.Len(t, filtered, 1) + assert.Equal(t, 2, filtered[0].ChannelId) + + // 白名单与候选无交集时降级为不过滤 + routing = RegionRouting{Active: true, AllowedIds: []int64{99}} + assert.Len(t, filterAbilitiesByRegion(abilities, routing), 3) +} diff --git a/model/sales_lead.go b/model/sales_lead.go new file mode 100644 index 000000000000..3031afe9d959 --- /dev/null +++ b/model/sales_lead.go @@ -0,0 +1,20 @@ +package model + +// SalesLead 存储营销站「联系销售」表单提交的销售线索。 +type SalesLead struct { + Id uint `gorm:"primaryKey" json:"id"` + Name string `gorm:"type:varchar(64);not null" json:"name"` + Email string `gorm:"type:varchar(128);not null;index" json:"email"` + Company string `gorm:"type:varchar(128)" json:"company"` + Region string `gorm:"type:varchar(64);not null" json:"region"` + UseCase string `gorm:"type:varchar(256);not null" json:"use_case"` + MonthlyVolume string `gorm:"type:varchar(64)" json:"monthly_volume"` + RequiredModels string `gorm:"type:varchar(512)" json:"required_models"` + Message string `gorm:"type:text" json:"message"` + Status string `gorm:"type:varchar(32);default:'new'" json:"status"` + Source string `gorm:"type:varchar(128)" json:"source"` + CreatedAt int64 `json:"created_at"` + UpdatedAt int64 `json:"updated_at"` +} + +func (SalesLead) TableName() string { return "sales_leads" } diff --git a/model/sla.go b/model/sla.go new file mode 100644 index 000000000000..727fc9fbfd37 --- /dev/null +++ b/model/sla.go @@ -0,0 +1,180 @@ +package model + +import ( + "time" +) + +// SlaIncident 服务事件记录（故障 / 维护公告）。 +type SlaIncident struct { + Id int64 `json:"id" gorm:"primaryKey;autoIncrement"` + Title string `json:"title" gorm:"type:varchar(128);not null"` + Description string `json:"description" gorm:"type:text"` + Status int `json:"status" gorm:"not null;default:1;index"` // 1 investigating, 2 identified, 3 monitoring, 4 resolved + Severity string `json:"severity" gorm:"type:varchar(16);default:'minor'"` // minor | major | critical + StartedAt int64 `json:"started_at" gorm:"bigint"` + ResolvedAt int64 `json:"resolved_at" gorm:"bigint"` + CreatedAt int64 `json:"created_at" gorm:"bigint;index"` + UpdatedAt int64 `json:"updated_at" gorm:"bigint"` +} + +// SLA 事件状态常量。 +const ( + SlaIncidentStatusInvestigating = 1 + SlaIncidentStatusIdentified = 2 + SlaIncidentStatusMonitoring = 3 + SlaIncidentStatusResolved = 4 +) + +// AllowedSlaIncidentStatuses 事件状态白名单。 +var AllowedSlaIncidentStatuses = map[int]bool{ + SlaIncidentStatusInvestigating: true, + SlaIncidentStatusIdentified: true, + SlaIncidentStatusMonitoring: true, + SlaIncidentStatusResolved: true, +} + +// AllowedSlaIncidentSeverities 事件严重度白名单。 +var AllowedSlaIncidentSeverities = map[string]bool{ + "minor": true, + "major": true, + "critical": true, +} + +// CreateSlaIncident 创建事件（填充时间戳）。 +func CreateSlaIncident(m *SlaIncident) error { + now := time.Now().Unix() + m.CreatedAt = now + m.UpdatedAt = now + if m.StartedAt == 0 { + m.StartedAt = now + } + return DB.Create(m).Error +} + +// GetSlaIncidentById 按 id 获取事件。 +func GetSlaIncidentById(id int64) (*SlaIncident, error) { + var m SlaIncident + if err := DB.Where("id = ?", id).First(&m).Error; err != nil { + return nil, err + } + return &m, nil +} + +// SearchSlaIncidents 分页列出事件；status 为空时返回全部，按创建时间倒序。 +func SearchSlaIncidents(page, pageSize int, status string) ([]*SlaIncident, int64, error) { + var items []*SlaIncident + var total int64 + q := DB.Model(&SlaIncident{}) + if status != "" { + q = q.Where("status = ?", status) + } + if err := q.Count(&total).Error; err != nil { + return nil, 0, err + } + offset := (page - 1) * pageSize + if err := q.Order("id DESC").Offset(offset).Limit(pageSize).Find(&items).Error; err != nil { + return nil, 0, err + } + return items, total, nil +} + +// UpdateSlaIncident 更新事件可编辑字段。 +func UpdateSlaIncident(m *SlaIncident) error { + updates := map[string]interface{}{ + "title": m.Title, + "description": m.Description, + "status": m.Status, + "severity": m.Severity, + "started_at": m.StartedAt, + "resolved_at": m.ResolvedAt, + "updated_at": time.Now().Unix(), + } + return DB.Model(&SlaIncident{}).Where("id = ?", m.Id).Updates(updates).Error +} + +// DeleteSlaIncident 删除事件。 +func DeleteSlaIncident(id int64) error { + return DB.Where("id = ?", id).Delete(&SlaIncident{}).Error +} + +// SlaNodeStatus 节点（渠道）状态摘要，用于状态页展示。 +type SlaNodeStatus struct { + Id int `json:"id"` + Name string `json:"name"` + Status int `json:"status"` + ResponseTime int `json:"response_time"` +} + +// SlaStatusSummary 状态页整体摘要（复用既有 Channel 与 PerfMetric 数据，不新增存储）。 +type SlaStatusSummary struct { + Availability float64 `json:"availability"` // 0~1，近 24h 成功/总请求 + WindowHours int `json:"window_hours"` // 统计窗口（小时） + NodeCount int64 `json:"node_count"` // 节点总数 + OkNodeCount int64 `json:"ok_node_count"` // 正常节点数 + ActiveIncidents int64 `json:"active_incidents"` // 未解决事件数 + Nodes []SlaNodeStatus `json:"nodes"` +} + +// GetSlaStatusSummary 聚合整体可用率、节点状态与活跃事件。 +func GetSlaStatusSummary(windowHours int) (*SlaStatusSummary, error) { + if windowHours <= 0 { + windowHours = 24 + } + startTs := time.Now().Unix() - int64(windowHours)*3600 + + // 整体可用率：近窗口内 PerfMetric 的成功/总请求。 + var reqCount, succCount int64 + if err := DB.Model(&PerfMetric{}). + Where("bucket_ts >= ?", startTs). + Select("COALESCE(SUM(request_count),0), COALESCE(SUM(success_count),0)"). + Row().Scan(&reqCount, &succCount); err != nil { + return nil, err + } + availability := 1.0 + if reqCount > 0 { + availability = float64(succCount) / float64(reqCount) + } + + // 节点状态：来自 channels。 + var channels []Channel + if err := DB.Find(&channels).Error; err != nil { + return nil, err + } + nodes := make([]SlaNodeStatus, 0, len(channels)) + var okNodeCount int64 + for _, ch := range channels { + if ch.Status == 1 { + okNodeCount++ + } + nodes = append(nodes, SlaNodeStatus{ + Id: ch.Id, + Name: ch.Name, + Status: ch.Status, + ResponseTime: ch.ResponseTime, + }) + } + + // 活跃事件：未解决（status != resolved）。 + var activeIncidents int64 + if err := DB.Model(&SlaIncident{}). + Where("status <> ?", SlaIncidentStatusResolved). + Count(&activeIncidents).Error; err != nil { + return nil, err + } + + return &SlaStatusSummary{ + Availability: availability, + WindowHours: windowHours, + NodeCount: int64(len(channels)), + OkNodeCount: okNodeCount, + ActiveIncidents: activeIncidents, + Nodes: nodes, + }, nil +} + +// CountActiveSlaIncidents 未解决事件数（公开状态页使用）。 +func CountActiveSlaIncidents() (int64, error) { + var n int64 + err := DB.Model(&SlaIncident{}).Where("status <> ?", SlaIncidentStatusResolved).Count(&n).Error + return n, err +} diff --git a/model/team.go b/model/team.go new file mode 100644 index 000000000000..c0582a959624 --- /dev/null +++ b/model/team.go @@ -0,0 +1,201 @@ +package model + +import ( + "strings" + "time" + + "gorm.io/gorm" +) + +// Team 企业团队空间：团队（部门 / 组织单元）实体。 +type Team struct { + Id int64 `json:"id" gorm:"primaryKey;autoIncrement"` + Name string `json:"name" gorm:"type:varchar(64);not null;uniqueIndex"` + Description string `json:"description" gorm:"type:varchar(255)"` + OwnerId int64 `json:"owner_id" gorm:"not null;index"` + CreatedAt int64 `json:"created_at" gorm:"bigint;index"` + UpdatedAt int64 `json:"updated_at" gorm:"bigint"` +} + +// TeamMember 团队成员与角色。 +type TeamMember struct { + Id int64 `json:"id" gorm:"primaryKey;autoIncrement"` + TeamId int64 `json:"team_id" gorm:"not null;index"` + UserId int64 `json:"user_id" gorm:"not null;index"` + Role string `json:"role" gorm:"type:varchar(16);not null;default:'member'"` // admin | member + CreatedAt int64 `json:"created_at" gorm:"bigint"` +} + +// TeamProject 团队内项目。 +type TeamProject struct { + Id int64 `json:"id" gorm:"primaryKey;autoIncrement"` + TeamId int64 `json:"team_id" gorm:"not null;index"` + Name string `json:"name" gorm:"type:varchar(64);not null"` + Description string `json:"description" gorm:"type:varchar(255)"` + CreatedAt int64 `json:"created_at" gorm:"bigint;index"` + UpdatedAt int64 `json:"updated_at" gorm:"bigint"` +} + +// 团队角色白名单。 +var AllowedTeamMemberRoles = map[string]bool{ + "admin": true, + "member": true, +} + +// CreateTeam 创建团队（调用方负责填充时间戳）。 +func CreateTeam(t *Team) error { + now := time.Now().Unix() + t.CreatedAt = now + t.UpdatedAt = now + return DB.Create(t).Error +} + +// GetTeamById 按 id 获取团队。 +func GetTeamById(id int64) (*Team, error) { + var t Team + err := DB.Where("id = ?", id).First(&t).Error + if err != nil { + return nil, err + } + return &t, nil +} + +// ListTeams 分页列出团队；keyword 为空时不过滤（按名称模糊匹配）。 +func ListTeams(page, pageSize int, keyword string) ([]*Team, int64, error) { + var items []*Team + var total int64 + q := DB.Model(&Team{}) + if strings.TrimSpace(keyword) != "" { + q = q.Where("name LIKE ?", "%"+keyword+"%") + } + if err := q.Count(&total).Error; err != nil { + return nil, 0, err + } + offset := (page - 1) * pageSize + if err := q.Order("id DESC").Offset(offset).Limit(pageSize).Find(&items).Error; err != nil { + return nil, 0, err + } + return items, total, nil +} + +// UpdateTeam 更新团队可编辑字段（Name 唯一键不可变概念上，此处允许改名）。 +func UpdateTeam(t *Team) error { + updates := map[string]interface{}{ + "description": t.Description, + "owner_id": t.OwnerId, + "updated_at": time.Now().Unix(), + } + return DB.Model(&Team{}).Where("id = ?", t.Id).Updates(updates).Error +} + +// DeleteTeam 删除团队并级联删除其成员与项目。 +func DeleteTeam(id int64) error { + return DB.Transaction(func(tx *gorm.DB) error { + if err := tx.Where("team_id = ?", id).Delete(&TeamMember{}).Error; err != nil { + return err + } + if err := tx.Where("team_id = ?", id).Delete(&TeamProject{}).Error; err != nil { + return err + } + return tx.Where("id = ?", id).Delete(&Team{}).Error + }) +} + +// CreateTeamMember 添加团队成员。 +func CreateTeamMember(m *TeamMember) error { + m.CreatedAt = time.Now().Unix() + return DB.Create(m).Error +} + +// ListTeamMembers 分页列出团队成员。 +func ListTeamMembers(teamId int64, page, pageSize int) ([]*TeamMember, int64, error) { + var items []*TeamMember + var total int64 + q := DB.Model(&TeamMember{}).Where("team_id = ?", teamId) + if err := q.Count(&total).Error; err != nil { + return nil, 0, err + } + offset := (page - 1) * pageSize + if err := q.Order("id ASC").Offset(offset).Limit(pageSize).Find(&items).Error; err != nil { + return nil, 0, err + } + return items, total, nil +} + +// DeleteTeamMember 按 team + user 移除成员。 +func DeleteTeamMember(teamId, userId int64) error { + return DB.Where("team_id = ? AND user_id = ?", teamId, userId).Delete(&TeamMember{}).Error +} + +// CreateTeamProject 添加团队项目。 +func CreateTeamProject(p *TeamProject) error { + now := time.Now().Unix() + p.CreatedAt = now + p.UpdatedAt = now + return DB.Create(p).Error +} + +// ListTeamProjects 列出团队全部项目。 +func ListTeamProjects(teamId int64) ([]*TeamProject, error) { + var items []*TeamProject + err := DB.Where("team_id = ?", teamId).Order("id ASC").Find(&items).Error + if err != nil { + return nil, err + } + return items, nil +} + +// DeleteTeamProject 删除团队项目。 +func DeleteTeamProject(teamId, projectId int64) error { + return DB.Where("team_id = ? AND id = ?", teamId, projectId).Delete(&TeamProject{}).Error +} + +// TeamBilling 部门账单汇总：成员额度近似 + 团队实际用量（按 team_id 聚合消耗日志）。 +type TeamBilling struct { + TeamId int64 `json:"team_id"` + MemberCount int64 `json:"member_count"` + Allocated int64 `json:"allocated"` // 成员额度总和 + Used int64 `json:"used"` // 成员已用额度总和 + UsageQuota int64 `json:"usage_quota"` // 团队实际消耗配额（来自消耗日志，按 team_id 聚合） + PromptTokens int64 `json:"prompt_tokens"` // 团队实际 prompt tokens + CompletionTokens int64 `json:"completion_tokens"` // 团队实际 completion tokens + RequestCount int64 `json:"request_count"` // 团队实际请求数 +} + +// GetTeamBilling 汇总团队成员的额度与用量。 +func GetTeamBilling(teamId int64) (*TeamBilling, error) { + var memberIds []int64 + if err := DB.Model(&TeamMember{}).Where("team_id = ?", teamId).Pluck("user_id", &memberIds).Error; err != nil { + return nil, err + } + billing := &TeamBilling{TeamId: teamId, MemberCount: int64(len(memberIds))} + if len(memberIds) == 0 { + return billing, nil + } + var allocated, used int64 + if err := DB.Model(&User{}). + Where("id IN ?", memberIds). + Select("COALESCE(SUM(quota),0), COALESCE(SUM(used_quota),0)"). + Row().Scan(&allocated, &used); err != nil { + return nil, err + } + billing.Allocated = allocated + billing.Used = used + + // 实际用量：按 team_id 聚合消耗日志（LOG_DB 未启用时跳过，避免空指针） + if LOG_DB != nil { + var usageQuota, promptTokens, completionTokens, requestCount int64 + if err := LOG_DB.Table("logs"). + Select("COALESCE(SUM(quota),0), COALESCE(SUM(prompt_tokens),0), COALESCE(SUM(completion_tokens),0), COUNT(*)"). + Where("team_id = ? AND type = ?", teamId, LogTypeConsume). + Row().Scan(&usageQuota, &promptTokens, &completionTokens, &requestCount); err != nil { + return nil, err + } + billing.UsageQuota = usageQuota + billing.PromptTokens = promptTokens + billing.CompletionTokens = completionTokens + billing.RequestCount = requestCount + } + + return billing, nil +} diff --git a/model/token.go b/model/token.go index 5aa8b3d56e0e..5d62258e7920 100644 --- a/model/token.go +++ b/model/token.go @@ -28,34 +28,9 @@ type Token struct { UsedQuota int `json:"used_quota" gorm:"default:0"` // used quota Group string `json:"group" gorm:"default:''"` CrossGroupRetry bool `json:"cross_group_retry"` // 跨分组重试，仅auto分组有效 - AutoGroups string `json:"-" gorm:"type:text"` DeletedAt gorm.DeletedAt `gorm:"index"` } -func (token *Token) GetAutoGroups() ([]string, error) { - if token.AutoGroups == "" { - return nil, nil - } - var groups []string - if err := common.UnmarshalJsonStr(token.AutoGroups, &groups); err != nil { - return nil, err - } - return groups, nil -} - -func (token *Token) SetAutoGroups(groups []string) error { - if len(groups) == 0 { - token.AutoGroups = "" - return nil - } - data, err := common.Marshal(groups) - if err != nil { - return err - } - token.AutoGroups = string(data) - return nil -} - func (token *Token) Clean() { token.Key = "" } @@ -316,16 +291,18 @@ func (token *Token) Insert() error { // Update Make sure your token's fields is completed, because this will update non-zero values func (token *Token) Update() (err error) { - err = DB.Model(token).Select("name", "status", "expired_time", "remain_quota", "unlimited_quota", - "model_limits_enabled", "model_limits", "allow_ips", "group", "cross_group_retry", "auto_groups").Updates(token).Error - if shouldUpdateRedis(true, err) { - if cacheErr := cacheSetToken(*token); cacheErr != nil { - common.SysLog("failed to update token cache: " + cacheErr.Error()) - if deleteErr := cacheDeleteToken(token.Key); deleteErr != nil { - common.SysLog("failed to invalidate token cache after update: " + deleteErr.Error()) - } + defer func() { + if shouldUpdateRedis(true, err) { + gopool.Go(func() { + err := cacheSetToken(*token) + if err != nil { + common.SysLog("failed to update token cache: " + err.Error()) + } + }) } - } + }() + err = DB.Model(token).Select("name", "status", "expired_time", "remain_quota", "unlimited_quota", + "model_limits_enabled", "model_limits", "allow_ips", "group", "cross_group_retry").Updates(token).Error return err } diff --git a/model/user.go b/model/user.go index eb4ea08642d5..2c83d0f1a8cf 100644 --- a/model/user.go +++ b/model/user.go @@ -101,6 +101,8 @@ type User struct { AffQuota int `json:"aff_quota" gorm:"type:int;default:0;column:aff_quota"` // 邀请剩余额度 AffHistoryQuota int `json:"aff_history_quota" gorm:"type:int;default:0;column:aff_history"` // 邀请历史额度 InviterId int `json:"inviter_id" gorm:"type:int;column:inviter_id;index"` + TeamId int64 `json:"team_id" gorm:"type:bigint;column:team_id;index"` // 所属企业团队（部门账单聚合） + RegionPreference string `json:"region_preference" gorm:"type:varchar(16);column:region_preference;default:''"` // 区域路由偏好（用户级），空表示不限定，交由 X-Region / 请求上下文决定 DeletedAt gorm.DeletedAt `gorm:"index"` LinuxDOId string `json:"linux_do_id" gorm:"column:linux_do_id;index"` Setting string `json:"setting" gorm:"type:text;column:setting"` @@ -114,16 +116,18 @@ type User struct { func (user *User) ToBaseUser() *UserBase { cache := &UserBase{ - Id: user.Id, - Group: user.Group, - Quota: user.Quota, - Status: user.Status, - Role: user.Role, - Username: user.Username, - Setting: user.Setting, - Email: user.Email, - AuthVersion: user.AuthVersion, - CacheSchema: userCacheSchemaVersion, + Id: user.Id, + Group: user.Group, + Quota: user.Quota, + Status: user.Status, + Role: user.Role, + Username: user.Username, + Setting: user.Setting, + Email: user.Email, + RegionPreference: user.RegionPreference, + TeamId: user.TeamId, + AuthVersion: user.AuthVersion, + CacheSchema: userCacheSchemaVersion, } return cache } @@ -139,22 +143,6 @@ func (user *User) SetAccessToken(token string) { user.AccessToken = &token } -// UpdateUserAccessToken rotates a dashboard personal access token without -// writing a stale user snapshot back over concurrently updated fields. -func UpdateUserAccessToken(id int, token string) error { - if id == 0 { - return errors.New("id 为空！") - } - result := DB.Model(&User{}).Where("id = ?", id).Update("access_token", token) - if result.Error != nil { - return result.Error - } - if result.RowsAffected == 0 { - return gorm.ErrRecordNotFound - } - return nil -} - func (user *User) GetSetting() dto.UserSetting { setting := dto.UserSetting{} if user.Setting != "" { @@ -505,19 +493,15 @@ func HardDeleteUserById(id int) error { return user.HardDelete() } -func inviteUser(inviterId int) error { - result := DB.Model(&User{}).Where("id = ?", inviterId).Updates(map[string]interface{}{ - "aff_count": gorm.Expr("aff_count + ?", 1), - "aff_quota": gorm.Expr("aff_quota + ?", common.QuotaForInviter), - "aff_history": gorm.Expr("aff_history + ?", common.QuotaForInviter), - }) - if result.Error != nil { - return result.Error - } - if result.RowsAffected == 0 { - return gorm.ErrRecordNotFound +func inviteUser(inviterId int) (err error) { + user, err := GetUserById(inviterId, true) + if err != nil { + return err } - return nil + user.AffCount++ + user.AffQuota += common.QuotaForInviter + user.AffHistoryQuota += common.QuotaForInviter + return DB.Save(user).Error } func (user *User) TransferAffQuotaToQuota(quota int) error { @@ -534,7 +518,7 @@ func (user *User) TransferAffQuotaToQuota(quota int) error { defer tx.Rollback() // 确保在函数退出时事务能回滚 // 加锁查询用户以确保数据一致性 - err := lockForUpdate(tx).First(user, user.Id).Error + err := lockForUpdate(tx).First(&user, user.Id).Error if err != nil { return err } @@ -768,16 +752,7 @@ func (user *User) UpdateWithTx(tx *gorm.DB, updatePassword bool) error { return err } } - if err = tx.Model(¤t).Omit( - "access_token", - "quota", - "used_quota", - "request_count", - "aff_count", - "aff_quota", - "aff_history", - "auth_version", - ).Updates(newUser).Error; err != nil { + if err = tx.Model(¤t).Omit("quota", "used_quota", "request_count", "auth_version").Updates(newUser).Error; err != nil { return err } return tx.First(user, user.Id).Error @@ -1447,3 +1422,20 @@ func RootUserExists() bool { } return true } + +// GetUsersByInviterId 列出由指定邀请人直接发展的下级用户（分销商下级用户查询）。 +func GetUsersByInviterId(inviterId int, page, pageSize int) ([]*User, int64, error) { + var items []*User + var total int64 + q := DB.Model(&User{}).Where("inviter_id = ?", inviterId) + if err := q.Count(&total).Error; err != nil { + return nil, 0, err + } + offset := (page - 1) * pageSize + if err := q.Order("id DESC").Offset(offset).Limit(pageSize). + Select("id, username, email, quota, used_quota, status, group, created_at, inviter_id, team_id"). + Find(&items).Error; err != nil { + return nil, 0, err + } + return items, total, nil +} diff --git a/model/user_cache.go b/model/user_cache.go index aa72fc26ae5f..10b8d5cc0caf 100644 --- a/model/user_cache.go +++ b/model/user_cache.go @@ -11,19 +11,21 @@ import ( "github.com/gin-gonic/gin" ) -const userCacheSchemaVersion = 2 +const userCacheSchemaVersion = 3 type UserBase struct { - Id int `json:"id"` - Group string `json:"group"` - Email string `json:"email"` - Quota int `json:"quota"` - Status int `json:"status"` - Role int `json:"role"` - Username string `json:"username"` - Setting string `json:"setting"` - AuthVersion int64 `json:"-"` - CacheSchema int `json:"-"` + Id int `json:"id"` + Group string `json:"group"` + Email string `json:"email"` + Quota int `json:"quota"` + Status int `json:"status"` + Role int `json:"role"` + Username string `json:"username"` + Setting string `json:"setting"` + RegionPreference string `json:"region_preference"` + TeamId int64 `json:"team_id"` + AuthVersion int64 `json:"-"` + CacheSchema int `json:"-"` } func (user *UserBase) WriteContext(c *gin.Context) { @@ -33,6 +35,8 @@ func (user *UserBase) WriteContext(c *gin.Context) { common.SetContextKey(c, constant.ContextKeyUserEmail, user.Email) common.SetContextKey(c, constant.ContextKeyUserName, user.Username) common.SetContextKey(c, constant.ContextKeyUserSetting, user.GetSetting()) + common.SetContextKey(c, constant.ContextKeyUserRegionPreference, user.RegionPreference) + common.SetContextKey(c, constant.ContextKeyUserTeamId, user.TeamId) } func (user *UserBase) GetSetting() dto.UserSetting { diff --git a/model/user_update_test.go b/model/user_update_test.go index 8e69ba1ebba2..c68d8f33c85a 100644 --- a/model/user_update_test.go +++ b/model/user_update_test.go @@ -27,23 +27,19 @@ func setupUserUpdateTestState(t *testing.T) { }) } -func TestUserUpdateDoesNotOverwriteConcurrentAccountingOrTokenChanges(t *testing.T) { +func TestUserUpdateDoesNotOverwriteAccountingFields(t *testing.T) { setupUserUpdateTestState(t) user := User{ - Id: 1, - Username: "quota-race-user", - Password: "password", - DisplayName: "before", - Status: common.UserStatusEnabled, - Quota: 1000, - UsedQuota: 20, - RequestCount: 3, - AffCount: 2, - AffQuota: 800, - AffHistoryQuota: 1200, + Id: 1, + Username: "quota-race-user", + Password: "password", + DisplayName: "before", + Status: common.UserStatusEnabled, + Quota: 1000, + UsedQuota: 20, + RequestCount: 3, } - user.SetAccessToken("old-token") require.NoError(t, DB.Create(&user).Error) staleUser, err := GetUserById(user.Id, true) @@ -53,10 +49,6 @@ func TestUserUpdateDoesNotOverwriteConcurrentAccountingOrTokenChanges(t *testing "quota": gorm.Expr("quota - ?", 400), "used_quota": gorm.Expr("used_quota + ?", 400), "request_count": gorm.Expr("request_count + ?", 1), - "aff_count": gorm.Expr("aff_count + ?", 1), - "aff_quota": gorm.Expr("aff_quota - ?", 500), - "aff_history": gorm.Expr("aff_history + ?", 500), - "access_token": "rotated-token", }).Error) staleUser.DisplayName = "after" @@ -68,63 +60,6 @@ func TestUserUpdateDoesNotOverwriteConcurrentAccountingOrTokenChanges(t *testing assert.Equal(t, 600, got.Quota) assert.Equal(t, 420, got.UsedQuota) assert.Equal(t, 4, got.RequestCount) - assert.Equal(t, 3, got.AffCount) - assert.Equal(t, 300, got.AffQuota) - assert.Equal(t, 1700, got.AffHistoryQuota) - assert.Equal(t, "rotated-token", got.GetAccessToken()) -} - -func TestUpdateUserAccessTokenOnlyUpdatesAccessToken(t *testing.T) { - setupUserUpdateTestState(t) - - user := User{ - Id: 2, - Username: "token-rotation-user", - Password: "password", - DisplayName: "before", - Status: common.UserStatusEnabled, - Quota: 1000, - AffQuota: 800, - AffHistoryQuota: 1200, - } - require.NoError(t, DB.Create(&user).Error) - - require.NoError(t, DB.Model(&User{}).Where("id = ?", user.Id).Updates(map[string]interface{}{ - "quota": gorm.Expr("quota + ?", 500), - "aff_quota": gorm.Expr("aff_quota - ?", 500), - "display_name": "concurrent-update", - }).Error) - - require.NoError(t, UpdateUserAccessToken(user.Id, "rotated-token")) - - var got User - require.NoError(t, DB.First(&got, user.Id).Error) - assert.Equal(t, "rotated-token", got.GetAccessToken()) - assert.Equal(t, "concurrent-update", got.DisplayName) - assert.Equal(t, 1500, got.Quota) - assert.Equal(t, 300, got.AffQuota) - assert.Equal(t, 1200, got.AffHistoryQuota) -} - -func TestUpdateUserAccessTokenRejectsSoftDeletedUser(t *testing.T) { - setupUserUpdateTestState(t) - - user := User{ - Id: 3, - Username: "deleted-token-rotation-user", - Password: "password", - Status: common.UserStatusEnabled, - } - user.SetAccessToken("old-token") - require.NoError(t, DB.Create(&user).Error) - require.NoError(t, DB.Delete(&user).Error) - - err := UpdateUserAccessToken(user.Id, "orphaned-token") - require.ErrorIs(t, err, gorm.ErrRecordNotFound) - - var got User - require.NoError(t, DB.Unscoped().First(&got, user.Id).Error) - assert.Equal(t, "old-token", got.GetAccessToken()) } func TestUpdateUserSettingOnlyUpdatesSetting(t *testing.T) { diff --git a/relay/alpha_search_handler.go b/relay/alpha_search_handler.go index 246eebe864ee..23ebafea6883 100644 --- a/relay/alpha_search_handler.go +++ b/relay/alpha_search_handler.go @@ -62,11 +62,12 @@ func AlphaSearchHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError } logger.LogDebug(c, "requestBody: %s", jsonData) - body, closer, err := relaycommon.NewOutboundJSONBody(jsonData) + body, size, closer, err := relaycommon.NewOutboundJSONBody(jsonData) if err != nil { return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) } defer closer.Close() + info.UpstreamRequestBodySize = size adaptor := GetAdaptor(info.ApiType) if adaptor == nil { diff --git a/relay/channel/api_request.go b/relay/channel/api_request.go index 48241b14a5e9..fc2855a85bca 100644 --- a/relay/channel/api_request.go +++ b/relay/channel/api_request.go @@ -25,26 +25,20 @@ import ( "github.com/gorilla/websocket" ) -// ApplyUpstreamBodyMetadata restores metadata that net/http cannot infer from -// a ReplayableBody. Callers must pass the original body because NewRequest -// hides its dynamic type behind req.Body's io.ReadCloser wrapper. -func ApplyUpstreamBodyMetadata(req *http.Request, body io.Reader) { - replayable, ok := body.(common2.ReplayableBody) - if !ok { +// applyUpstreamContentLength populates req.ContentLength when the upstream +// body is wrapped in a BodyStorage (see relay/common/outbound_body.go). +// +// net/http.NewRequest only auto-detects ContentLength for *bytes.Reader, +// *bytes.Buffer and *strings.Reader. When the body is a type-erased io.Reader +// (which is the case for ReaderOnly(BodyStorage)), the Content-Length header +// would otherwise be omitted, forcing chunked transfer encoding and breaking +// some upstreams that require an explicit Content-Length. +func applyUpstreamContentLength(req *http.Request, info *common.RelayInfo) { + if info == nil { return } - - // BodyStorage structurally satisfies ReplayableBody, but it also exposes - // io.Closer. If a caller passes the storage directly instead of using - // NewReplayableBodyReader, hide Close before the transport takes ownership - // of req.Body so the shared replay source remains available to GetBody. - if _, rawStorage := body.(common2.BodyStorage); rawStorage { - req.Body = io.NopCloser(body) - } - - req.ContentLength = replayable.Size() - if req.GetBody == nil { - req.GetBody = replayable.NewReader + if info.UpstreamRequestBodySize > 0 && req.ContentLength <= 0 { + req.ContentLength = info.UpstreamRequestBodySize } } @@ -320,7 +314,7 @@ func DoApiRequest(a Adaptor, c *gin.Context, info *common.RelayInfo, requestBody if err != nil { return nil, fmt.Errorf("new request failed: %w", err) } - ApplyUpstreamBodyMetadata(req, requestBody) + applyUpstreamContentLength(req, info) headers := req.Header err = a.SetupRequestHeader(c, &headers, info) if err != nil { @@ -350,7 +344,7 @@ func DoFormRequest(a Adaptor, c *gin.Context, info *common.RelayInfo, requestBod if err != nil { return nil, fmt.Errorf("new request failed: %w", err) } - ApplyUpstreamBodyMetadata(req, requestBody) + applyUpstreamContentLength(req, info) // set form data req.Header.Set("Content-Type", c.Request.Header.Get("Content-Type")) headers := req.Header @@ -480,24 +474,11 @@ func sendPingData(c *gin.Context, mutex *sync.Mutex) error { func DoRequest(c *gin.Context, req *http.Request, info *common.RelayInfo) (*http.Response, error) { return doRequest(c, req, info) } - -// keepUpstreamRedirectResponse stops net/http from following redirects while -// returning the upstream 3xx response to the relay without an extra error. -func keepUpstreamRedirectResponse(_ *http.Request, _ []*http.Request) error { - return http.ErrUseLastResponse -} - func doRequest(c *gin.Context, req *http.Request, info *common.RelayInfo) (*http.Response, error) { client, err := service.GetHttpClientWithProxySettings(info.ChannelSetting.Proxy, info.ChannelSetting) if err != nil { return nil, fmt.Errorf("new proxy http client failed: %w", err) } - // Clients are cached and shared across channels, so override redirect - // behavior on a shallow copy instead of mutating the cached client. This - // still reuses its transport and connection pools, including HTTP/2's - // transparent stream retries. - relayClient := *client - relayClient.CheckRedirect = keepUpstreamRedirectResponse if common2.DebugEnabled && req != nil && req.URL != nil { policy := service.NormalizeHTTPTransportPolicy(info.ChannelSetting) logger.LogDebug(c, fmt.Sprintf( @@ -529,7 +510,7 @@ func doRequest(c *gin.Context, req *http.Request, info *common.RelayInfo) (*http } } - resp, err := relayClient.Do(req) + resp, err := client.Do(req) if err != nil { logger.LogError(c, "do request failed: "+err.Error()) return nil, types.NewError(err, types.ErrorCodeDoRequestFailed, types.ErrOptionWithHideErrMsg("upstream error: do request failed")) @@ -567,15 +548,10 @@ func DoTaskApiRequest(a TaskAdaptor, c *gin.Context, info *common.RelayInfo, req if err != nil { return nil, fmt.Errorf("new request failed: %w", err) } - ApplyUpstreamBodyMetadata(req, requestBody) - // Do NOT wrap requestBody in a GetBody closure here: returning the same - // (already consumed) reader would make any transport-level retry silently - // replay an empty body. http.NewRequest already derives a correct, - // snapshot-based GetBody for *bytes.Reader/Buffer/strings.Reader bodies - // (which most task adaptors pass in); ApplyUpstreamBodyMetadata wires the - // same contract for bodies that explicitly implement ReplayableBody. - // Otherwise GetBody stays nil so the transport fails the retry instead of - // sending a corrupted request. + applyUpstreamContentLength(req, info) + req.GetBody = func() (io.ReadCloser, error) { + return io.NopCloser(requestBody), nil + } err = a.BuildRequestHeader(c, req, info) if err != nil { diff --git a/relay/channel/aws/relay-aws.go b/relay/channel/aws/relay-aws.go index c4751b5af855..c502364c6adf 100644 --- a/relay/channel/aws/relay-aws.go +++ b/relay/channel/aws/relay-aws.go @@ -40,24 +40,11 @@ func getAwsErrorStatusCode(err error) int { return http.StatusInternalServerError } -func newAwsInvokeContext(parent context.Context) (context.Context, context.CancelFunc) { +func newAwsInvokeContext() (context.Context, context.CancelFunc) { if common.RelayTimeout <= 0 { - return context.WithCancel(parent) + return context.Background(), func() {} } - return context.WithTimeout(parent, time.Duration(common.RelayTimeout)*time.Second) -} - -func newAwsInvokeError(requestContext context.Context, err error, operation string) *types.NewAPIError { - options := make([]types.NewAPIErrorOptions, 0, 1) - if requestContext.Err() != nil { - options = append(options, types.ErrOptionWithSkipRetry()) - } - return types.NewOpenAIError( - errors.Wrap(err, operation), - types.ErrorCodeAwsInvokeError, - getAwsErrorStatusCode(err), - options..., - ) + return context.WithTimeout(context.Background(), time.Duration(common.RelayTimeout)*time.Second) } func newAwsClient(c *gin.Context, info *relaycommon.RelayInfo) (*bedrockruntime.Client, error) { @@ -228,13 +215,13 @@ func getAwsModelID(requestModel string) string { func awsHandler(c *gin.Context, info *relaycommon.RelayInfo, a *Adaptor) (*types.NewAPIError, *dto.Usage) { - requestContext := c.Request.Context() - ctx, cancel := newAwsInvokeContext(requestContext) + ctx, cancel := newAwsInvokeContext() defer cancel() awsResp, err := a.AwsClient.InvokeModel(ctx, a.AwsReq.(*bedrockruntime.InvokeModelInput)) if err != nil { - return newAwsInvokeError(requestContext, err, "InvokeModel"), nil + statusCode := getAwsErrorStatusCode(err) + return types.NewOpenAIError(errors.Wrap(err, "InvokeModel"), types.ErrorCodeAwsInvokeError, statusCode), nil } claudeInfo := &claude.ClaudeResponseInfo{ @@ -258,13 +245,13 @@ func awsHandler(c *gin.Context, info *relaycommon.RelayInfo, a *Adaptor) (*types } func awsStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, a *Adaptor) (*types.NewAPIError, *dto.Usage) { - requestContext := c.Request.Context() - ctx, cancel := newAwsInvokeContext(requestContext) + ctx, cancel := newAwsInvokeContext() defer cancel() awsResp, err := a.AwsClient.InvokeModelWithResponseStream(ctx, a.AwsReq.(*bedrockruntime.InvokeModelWithResponseStreamInput)) if err != nil { - return newAwsInvokeError(requestContext, err, "InvokeModelWithResponseStream"), nil + statusCode := getAwsErrorStatusCode(err) + return types.NewOpenAIError(errors.Wrap(err, "InvokeModelWithResponseStream"), types.ErrorCodeAwsInvokeError, statusCode), nil } stream := awsResp.GetStream() defer stream.Close() @@ -277,38 +264,23 @@ func awsStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, a *Adaptor) ( Usage: &dto.Usage{}, } - events := stream.Events() -streamLoop: - for { - select { - case <-ctx.Done(): - break streamLoop - case event, ok := <-events: - if !ok { - break streamLoop - } - if ctx.Err() != nil { - break streamLoop - } - - switch v := event.(type) { - case *bedrockruntimeTypes.ResponseStreamMemberChunk: - info.SetFirstResponseTime() - respErr := claude.HandleStreamResponseData(c, info, claudeInfo, string(v.Value.Bytes)) - if respErr != nil { - return respErr, nil - } - case *bedrockruntimeTypes.UnknownUnionMember: - fmt.Println("unknown tag:", v.Tag) - return types.NewError(errors.New("unknown response type"), types.ErrorCodeInvalidRequest), nil - default: - fmt.Println("union is nil or unknown type") - return types.NewError(errors.New("nil or unknown response type"), types.ErrorCodeInvalidRequest), nil + for event := range stream.Events() { + switch v := event.(type) { + case *bedrockruntimeTypes.ResponseStreamMemberChunk: + info.SetFirstResponseTime() + respErr := claude.HandleStreamResponseData(c, info, claudeInfo, string(v.Value.Bytes)) + if respErr != nil { + return respErr, nil } + case *bedrockruntimeTypes.UnknownUnionMember: + fmt.Println("unknown tag:", v.Tag) + return types.NewError(errors.New("unknown response type"), types.ErrorCodeInvalidRequest), nil + default: + fmt.Println("union is nil or unknown type") + return types.NewError(errors.New("nil or unknown response type"), types.ErrorCodeInvalidRequest), nil } } - _ = stream.Close() claude.HandleStreamFinalResponse(c, info, claudeInfo) return nil, claudeInfo.Usage } @@ -316,13 +288,13 @@ streamLoop: // Nova模型处理函数 func handleNovaRequest(c *gin.Context, info *relaycommon.RelayInfo, a *Adaptor) (*types.NewAPIError, *dto.Usage) { - requestContext := c.Request.Context() - ctx, cancel := newAwsInvokeContext(requestContext) + ctx, cancel := newAwsInvokeContext() defer cancel() awsResp, err := a.AwsClient.InvokeModel(ctx, a.AwsReq.(*bedrockruntime.InvokeModelInput)) if err != nil { - return newAwsInvokeError(requestContext, err, "InvokeModel"), nil + statusCode := getAwsErrorStatusCode(err) + return types.NewOpenAIError(errors.Wrap(err, "InvokeModel"), types.ErrorCodeAwsInvokeError, statusCode), nil } // 解析Nova响应 diff --git a/relay/channel/aws/relay_aws_test.go b/relay/channel/aws/relay_aws_test.go index 22d8373873ed..92745ff40929 100644 --- a/relay/channel/aws/relay_aws_test.go +++ b/relay/channel/aws/relay_aws_test.go @@ -2,145 +2,17 @@ package aws import ( "bytes" - "context" - "errors" - "io" "net/http" "net/http/httptest" - "sync" "testing" - "time" "github.com/QuantumNous/new-api/common" relaycommon "github.com/QuantumNous/new-api/relay/common" - "github.com/QuantumNous/new-api/relaykit/dto" - relaytypes "github.com/QuantumNous/new-api/relaykit/types" - "github.com/aws/aws-sdk-go-v2/aws" - "github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream" - "github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream/eventstreamapi" - "github.com/aws/aws-sdk-go-v2/credentials" "github.com/aws/aws-sdk-go-v2/service/bedrockruntime" "github.com/gin-gonic/gin" - "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) -const awsTestModel = "anthropic.claude-3-5-sonnet-20240620-v1:0" - -type awsHTTPClientFunc func(*http.Request) (*http.Response, error) - -func (f awsHTTPClientFunc) Do(request *http.Request) (*http.Response, error) { - return f(request) -} - -type awsNotifyingResponseWriter struct { - *httptest.ResponseRecorder - notifyOn []byte - notified chan int - once sync.Once -} - -func newAwsNotifyingResponseWriter(notifyOn string) *awsNotifyingResponseWriter { - return &awsNotifyingResponseWriter{ - ResponseRecorder: httptest.NewRecorder(), - notifyOn: []byte(notifyOn), - notified: make(chan int, 1), - } -} - -func (w *awsNotifyingResponseWriter) Write(data []byte) (int, error) { - return w.ResponseRecorder.Write(data) -} - -func (w *awsNotifyingResponseWriter) Flush() { - w.ResponseRecorder.Flush() - if bytes.Contains(w.Body.Bytes(), w.notifyOn) { - w.once.Do(func() { - w.notified <- w.Body.Len() - }) - } -} - -func newAwsTestClient(httpClient bedrockruntime.HTTPClient) *bedrockruntime.Client { - return bedrockruntime.New(bedrockruntime.Options{ - Region: "us-east-1", - BaseEndpoint: aws.String("https://bedrock.test"), - Credentials: aws.NewCredentialsCache(credentials.NewStaticCredentialsProvider( - "access-key", "secret-key", "", - )), - HTTPClient: httpClient, - Retryer: aws.NopRetryer{}, - }) -} - -func newAwsTestContext(writer http.ResponseWriter, requestContext context.Context) *gin.Context { - c, _ := gin.CreateTestContext(writer) - c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil).WithContext(requestContext) - return c -} - -func newAwsTestRelayInfo() *relaycommon.RelayInfo { - return &relaycommon.RelayInfo{ - StartTime: time.Now(), - IsStream: true, - OriginModelName: awsTestModel, - RelayFormat: relaytypes.RelayFormatOpenAI, - ShouldIncludeUsage: true, - ChannelMeta: &relaycommon.ChannelMeta{ - UpstreamModelName: awsTestModel, - }, - } -} - -func newAwsInvokeModelInput() *bedrockruntime.InvokeModelInput { - return &bedrockruntime.InvokeModelInput{ - ModelId: aws.String(awsTestModel), - Body: []byte(`{}`), - Accept: aws.String("application/json"), - ContentType: aws.String("application/json"), - } -} - -func newAwsStreamInput() *bedrockruntime.InvokeModelWithResponseStreamInput { - return &bedrockruntime.InvokeModelWithResponseStreamInput{ - ModelId: aws.String(awsTestModel), - Body: []byte(`{}`), - Accept: aws.String("application/json"), - ContentType: aws.String("application/json"), - } -} - -func writeAwsStreamEvent(writer io.Writer, data string) error { - payload, err := common.Marshal(struct { - Bytes []byte `json:"bytes"` - }{Bytes: []byte(data)}) - if err != nil { - return err - } - - return eventstream.NewEncoder().Encode(writer, eventstream.Message{ - Headers: eventstream.Headers{ - {Name: eventstreamapi.MessageTypeHeader, Value: eventstream.StringValue(eventstreamapi.EventMessageType)}, - {Name: eventstreamapi.EventTypeHeader, Value: eventstream.StringValue("chunk")}, - {Name: eventstreamapi.ContentTypeHeader, Value: eventstream.StringValue("application/json")}, - }, - Payload: payload, - }) -} - -func newAwsStreamResponse(request *http.Request, body io.ReadCloser) *http.Response { - return &http.Response{ - StatusCode: http.StatusOK, - Status: "200 OK", - Header: http.Header{ - "Content-Type": []string{"application/vnd.amazon.eventstream"}, - "X-Amzn-Bedrock-Content-Type": []string{"application/json"}, - }, - Body: body, - Request: request, - } -} - func TestDoAwsClientRequest_AppliesRuntimeHeaderOverrideToAnthropicBeta(t *testing.T) { t.Parallel() @@ -181,278 +53,3 @@ func TestDoAwsClientRequest_AppliesRuntimeHeaderOverrideToAnthropicBeta(t *testi require.True(t, ok) require.Equal(t, []any{"computer-use-2025-01-24"}, values) } - -func TestNewAwsInvokeContextInheritsParent(t *testing.T) { - originalRelayTimeout := common.RelayTimeout - t.Cleanup(func() { - common.RelayTimeout = originalRelayTimeout - }) - - tests := []struct { - name string - relayTimeout int - wantDeadline bool - }{ - {name: "without relay timeout", relayTimeout: 0, wantDeadline: false}, - {name: "with relay timeout", relayTimeout: 30, wantDeadline: true}, - } - - for _, test := range tests { - t.Run(test.name, func(t *testing.T) { - common.RelayTimeout = test.relayTimeout - parent, cancelParent := context.WithCancel(context.Background()) - invokeContext, cancelInvoke := newAwsInvokeContext(parent) - defer cancelInvoke() - - _, hasDeadline := invokeContext.Deadline() - assert.Equal(t, test.wantDeadline, hasDeadline) - - cancelParent() - require.ErrorIs(t, invokeContext.Err(), context.Canceled) - }) - } -} - -func TestNewAwsInvokeErrorSkipsRetryOnlyForClientCancellation(t *testing.T) { - canceledContext, cancel := context.WithCancel(context.Background()) - cancel() - - tests := []struct { - name string - requestContext context.Context - err error - wantSkipRetry bool - }{ - { - name: "client context canceled", - requestContext: canceledContext, - err: context.Canceled, - wantSkipRetry: true, - }, - { - name: "relay timeout with live client context", - requestContext: context.Background(), - err: context.DeadlineExceeded, - wantSkipRetry: false, - }, - { - name: "upstream error with live client context", - requestContext: context.Background(), - err: errors.New("upstream failed"), - wantSkipRetry: false, - }, - } - - for _, test := range tests { - t.Run(test.name, func(t *testing.T) { - err := newAwsInvokeError(test.requestContext, test.err, "InvokeModel") - assert.Equal(t, test.wantSkipRetry, relaytypes.IsSkipRetryError(err)) - }) - } -} - -func TestAwsHandlersCancelSdkRequestAndSkipRetry(t *testing.T) { - originalRelayTimeout := common.RelayTimeout - common.RelayTimeout = 0 - t.Cleanup(func() { - common.RelayTimeout = originalRelayTimeout - }) - - tests := []struct { - name string - request any - handle func(*gin.Context, *relaycommon.RelayInfo, *Adaptor) (*relaytypes.NewAPIError, *dto.Usage) - }{ - {name: "non-stream", request: newAwsInvokeModelInput(), handle: awsHandler}, - {name: "stream", request: newAwsStreamInput(), handle: awsStreamHandler}, - {name: "nova", request: newAwsInvokeModelInput(), handle: handleNovaRequest}, - } - - for _, test := range tests { - t.Run(test.name, func(t *testing.T) { - requestContext, cancelRequest := context.WithCancel(context.Background()) - t.Cleanup(cancelRequest) - - upstreamContexts := make(chan context.Context, 1) - client := newAwsTestClient(awsHTTPClientFunc(func(request *http.Request) (*http.Response, error) { - upstreamContexts <- request.Context() - <-request.Context().Done() - return nil, request.Context().Err() - })) - adaptor := &Adaptor{AwsClient: client, AwsReq: test.request} - c := newAwsTestContext(httptest.NewRecorder(), requestContext) - info := newAwsTestRelayInfo() - - type handlerResult struct { - err *relaytypes.NewAPIError - usage *dto.Usage - } - results := make(chan handlerResult, 1) - go func() { - err, usage := test.handle(c, info, adaptor) - results <- handlerResult{err: err, usage: usage} - }() - - var upstreamContext context.Context - select { - case upstreamContext = <-upstreamContexts: - case result := <-results: - t.Fatalf("handler returned before issuing AWS request: %v", result.err) - case <-time.After(5 * time.Second): - t.Fatal("AWS request did not start") - } - - cancelRequest() - - var result handlerResult - select { - case result = <-results: - case <-time.After(5 * time.Second): - t.Fatal("handler did not stop after client cancellation") - } - - require.ErrorIs(t, upstreamContext.Err(), context.Canceled) - require.NotNil(t, result.err) - assert.True(t, relaytypes.IsSkipRetryError(result.err)) - assert.Nil(t, result.usage) - }) - } -} - -func TestAwsStreamHandlerUsesFinalUpstreamUsage(t *testing.T) { - originalRelayTimeout := common.RelayTimeout - common.RelayTimeout = 0 - t.Cleanup(func() { - common.RelayTimeout = originalRelayTimeout - }) - - events := []string{ - `{"type":"message_start","message":{"id":"msg_test","type":"message","role":"assistant","model":"claude-test","content":[],"usage":{"input_tokens":100,"output_tokens":1}}}`, - `{"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}`, - `{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"partial"}}`, - `{"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":423}}`, - `{"type":"message_stop"}`, - } - client := newAwsTestClient(awsHTTPClientFunc(func(request *http.Request) (*http.Response, error) { - var body bytes.Buffer - for _, event := range events { - if err := writeAwsStreamEvent(&body, event); err != nil { - return nil, err - } - } - return newAwsStreamResponse(request, io.NopCloser(bytes.NewReader(body.Bytes()))), nil - })) - adaptor := &Adaptor{AwsClient: client, AwsReq: newAwsStreamInput()} - recorder := httptest.NewRecorder() - c := newAwsTestContext(recorder, context.Background()) - - handlerErr, usage := awsStreamHandler(c, newAwsTestRelayInfo(), adaptor) - - require.Nil(t, handlerErr) - require.NotNil(t, usage) - require.NotNil(t, usage.BillingUsage) - require.NotNil(t, usage.BillingUsage.ClaudeUsage) - assert.Equal(t, 100, usage.BillingUsage.ClaudeUsage.InputTokens) - assert.Equal(t, 423, usage.BillingUsage.ClaudeUsage.OutputTokens) - assert.Contains(t, recorder.Body.String(), "[DONE]") -} - -func TestAwsStreamHandlerStopsAtClientCancellationAndKeepsPartialBillingUsage(t *testing.T) { - originalRelayTimeout := common.RelayTimeout - common.RelayTimeout = 0 - t.Cleanup(func() { - common.RelayTimeout = originalRelayTimeout - }) - - requestContext, cancelRequest := context.WithCancel(context.Background()) - t.Cleanup(cancelRequest) - releaseFinal := make(chan struct{}) - var releaseFinalOnce sync.Once - release := func() { - releaseFinalOnce.Do(func() { - close(releaseFinal) - }) - } - t.Cleanup(release) - - producerResults := make(chan error, 1) - upstreamContexts := make(chan context.Context, 1) - client := newAwsTestClient(awsHTTPClientFunc(func(request *http.Request) (*http.Response, error) { - upstreamContexts <- request.Context() - reader, writer := io.Pipe() - go func() { - defer writer.Close() - initialEvents := []string{ - `{"type":"message_start","message":{"id":"msg_test","type":"message","role":"assistant","model":"claude-test","content":[],"usage":{"input_tokens":100,"output_tokens":1}}}`, - `{"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}`, - `{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"partial"}}`, - } - for _, event := range initialEvents { - if err := writeAwsStreamEvent(writer, event); err != nil { - producerResults <- err - return - } - } - - <-releaseFinal - producerResults <- writeAwsStreamEvent(writer, `{"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":423}}`) - }() - return newAwsStreamResponse(request, reader), nil - })) - - responseWriter := newAwsNotifyingResponseWriter("partial") - c := newAwsTestContext(responseWriter, requestContext) - adaptor := &Adaptor{AwsClient: client, AwsReq: newAwsStreamInput()} - - type handlerResult struct { - err *relaytypes.NewAPIError - usage *dto.Usage - } - results := make(chan handlerResult, 1) - go func() { - err, usage := awsStreamHandler(c, newAwsTestRelayInfo(), adaptor) - results <- handlerResult{err: err, usage: usage} - }() - - var upstreamContext context.Context - select { - case upstreamContext = <-upstreamContexts: - case <-time.After(5 * time.Second): - t.Fatal("AWS stream request did not start") - } - - var bodyLengthBeforeCancel int - select { - case bodyLengthBeforeCancel = <-responseWriter.notified: - case <-time.After(5 * time.Second): - t.Fatal("partial response was not written") - } - cancelRequest() - - var result handlerResult - select { - case result = <-results: - case <-time.After(5 * time.Second): - t.Fatal("stream handler did not stop after client cancellation") - } - - require.ErrorIs(t, upstreamContext.Err(), context.Canceled) - require.Nil(t, result.err) - require.NotNil(t, result.usage) - require.NotNil(t, result.usage.BillingUsage) - require.NotNil(t, result.usage.BillingUsage.ClaudeUsage) - assert.Equal(t, dto.BillingUsageSourceClaudeMessages, result.usage.BillingUsage.Source) - assert.Equal(t, dto.BillingUsageSemanticAnthropic, result.usage.BillingUsage.Semantic) - assert.Equal(t, 100, result.usage.BillingUsage.ClaudeUsage.InputTokens) - assert.Equal(t, 1, result.usage.BillingUsage.ClaudeUsage.OutputTokens) - assert.Equal(t, bodyLengthBeforeCancel, responseWriter.Body.Len()) - assert.NotContains(t, responseWriter.Body.String(), "[DONE]") - - release() - select { - case producerErr := <-producerResults: - require.Error(t, producerErr) - case <-time.After(5 * time.Second): - t.Fatal("upstream producer did not observe the closed stream") - } -} diff --git a/relay/channel/jimeng/adaptor.go b/relay/channel/jimeng/adaptor.go index 776c1829a03e..658d9236420d 100644 --- a/relay/channel/jimeng/adaptor.go +++ b/relay/channel/jimeng/adaptor.go @@ -112,7 +112,6 @@ func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, request if err != nil { return nil, fmt.Errorf("new request failed: %w", err) } - channel.ApplyUpstreamBodyMetadata(req, requestBody) err = Sign(c, req, info.ApiKey) if err != nil { return nil, fmt.Errorf("setup request header failed: %w", err) diff --git a/relay/channel/task/sora/adaptor.go b/relay/channel/task/sora/adaptor.go index 7f81e5335ebb..e9029aa20d46 100644 --- a/relay/channel/task/sora/adaptor.go +++ b/relay/channel/task/sora/adaptor.go @@ -216,7 +216,7 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn return &buf, nil } - return common.NewReplayableBodyReader(storage), nil + return common.ReaderOnly(storage), nil } // DoRequest delegates to common helper. diff --git a/relay/chat_completions_via_responses.go b/relay/chat_completions_via_responses.go index b8a6fc875872..ccc7fb281220 100644 --- a/relay/chat_completions_via_responses.go +++ b/relay/chat_completions_via_responses.go @@ -128,12 +128,13 @@ func chatCompletionsViaResponses(c *gin.Context, info *relaycommon.RelayInfo, ad return nil, types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) } - body, closer, err := relaycommon.NewOutboundJSONBody(jsonData) + body, size, closer, err := relaycommon.NewOutboundJSONBody(jsonData) if err != nil { return nil, types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) } defer closer.Close() jsonData = nil + info.UpstreamRequestBodySize = size var requestBody io.Reader = body var httpResp *http.Response diff --git a/relay/claude_handler.go b/relay/claude_handler.go index c8b01f7ac8dc..63bc11f5889c 100644 --- a/relay/claude_handler.go +++ b/relay/claude_handler.go @@ -159,7 +159,8 @@ func ClaudeHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ if err != nil { return types.NewErrorWithStatusCode(err, types.ErrorCodeReadRequestBodyFailed, http.StatusBadRequest, types.ErrOptionWithSkipRetry()) } - requestBody = common.NewReplayableBodyReader(storage) + info.UpstreamRequestBodySize = storage.Size() + requestBody = common.ReaderOnly(storage) } else { convertedRequest, err := adaptor.ConvertClaudeRequest(c, info, request) if err != nil { @@ -186,12 +187,13 @@ func ClaudeHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ } logger.LogDebug(c, "requestBody: %s", jsonData) - body, closer, err := relaycommon.NewOutboundJSONBody(jsonData) + body, size, closer, err := relaycommon.NewOutboundJSONBody(jsonData) if err != nil { return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) } defer closer.Close() jsonData = nil + info.UpstreamRequestBodySize = size requestBody = body } diff --git a/relay/common/outbound_body.go b/relay/common/outbound_body.go index ac9e110e0eba..94ef8dde1da6 100644 --- a/relay/common/outbound_body.go +++ b/relay/common/outbound_body.go @@ -18,14 +18,14 @@ import ( // The caller MUST invoke closer.Close() once the upstream call has finished // (typically via defer) to release the disk file / memory accounting. // -// The returned body exposes its size and replay capability without exposing -// io.Closer. Request construction uses that metadata to populate ContentLength -// and GetBody, while the caller retains ownership of the underlying storage -// through the separately returned closer. -func NewOutboundJSONBody(data []byte) (body common.ReplayableBody, closer io.Closer, err error) { +// The returned reader is wrapped with common.ReaderOnly to prevent the HTTP +// transport from prematurely closing the underlying BodyStorage. The returned +// size is meant to be propagated to http.Request.ContentLength because the +// type-erased io.Reader prevents net/http from auto-detecting it. +func NewOutboundJSONBody(data []byte) (body io.Reader, size int64, closer io.Closer, err error) { storage, err := common.CreateBodyStorage(data) if err != nil { - return nil, nil, err + return nil, 0, nil, err } - return common.NewReplayableBodyReader(storage), storage, nil + return common.ReaderOnly(storage), storage.Size(), storage, nil } diff --git a/relay/common/relay_info.go b/relay/common/relay_info.go index 594e2640182d..45ae30bf9fe8 100644 --- a/relay/common/relay_info.go +++ b/relay/common/relay_info.go @@ -149,6 +149,13 @@ type RelayInfo struct { UseRuntimeHeadersOverride bool ParamOverrideAudit []string + // UpstreamRequestBodySize is the byte size of the marshaled upstream request + // body. It is set when the body is wrapped in a BodyStorage (see + // relay/common/outbound_body.go), so that DoApiRequest can populate + // http.Request.ContentLength manually (net/http only auto-detects it for + // *bytes.Reader/Buffer/strings.Reader). 0 means "let net/http decide". + UpstreamRequestBodySize int64 + PriceData hosttypes.PriceData // QuotaClamp is set (non-nil) when a quota conversion saturated at the diff --git a/relay/compatible_handler.go b/relay/compatible_handler.go index 8edb3362b6e0..b1e25e036f8d 100644 --- a/relay/compatible_handler.go +++ b/relay/compatible_handler.go @@ -104,7 +104,7 @@ func TextHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *types logger.LogDebug(c, "requestBody: %s", debugBytes) } } - requestBody = common.NewReplayableBodyReader(storage) + requestBody = common.ReaderOnly(storage) } else { convertedRequest, err := adaptor.ConvertOpenAIRequest(c, info, request) if err != nil { @@ -175,12 +175,13 @@ func TextHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *types logger.LogDebug(c, "text request body: %s", jsonData) - body, closer, err := relaycommon.NewOutboundJSONBody(jsonData) + body, size, closer, err := relaycommon.NewOutboundJSONBody(jsonData) if err != nil { return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) } defer closer.Close() jsonData = nil + info.UpstreamRequestBodySize = size requestBody = body } diff --git a/relay/embedding_handler.go b/relay/embedding_handler.go index 44d3de486dc4..3be543bd91a1 100644 --- a/relay/embedding_handler.go +++ b/relay/embedding_handler.go @@ -58,12 +58,13 @@ func EmbeddingHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError * } logger.LogDebug(c, "converted embedding request body: %s", jsonData) - body, closer, err := relaycommon.NewOutboundJSONBody(jsonData) + body, size, closer, err := relaycommon.NewOutboundJSONBody(jsonData) if err != nil { return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) } defer closer.Close() jsonData = nil + info.UpstreamRequestBodySize = size var requestBody io.Reader = body statusCodeMappingStr := c.GetString("status_code_mapping") resp, err := adaptor.DoRequest(c, info, requestBody) diff --git a/relay/gemini_handler.go b/relay/gemini_handler.go index 57010d87c380..786feb8ba345 100644 --- a/relay/gemini_handler.go +++ b/relay/gemini_handler.go @@ -141,7 +141,7 @@ func GeminiHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ if err != nil { return types.NewErrorWithStatusCode(err, types.ErrorCodeReadRequestBodyFailed, http.StatusBadRequest, types.ErrOptionWithSkipRetry()) } - requestBody = common.NewReplayableBodyReader(storage) + requestBody = common.ReaderOnly(storage) } else { // 使用 ConvertGeminiRequest 转换请求格式 convertedRequest, err := adaptor.ConvertGeminiRequest(c, info, request) @@ -164,12 +164,13 @@ func GeminiHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ logger.LogDebug(c, "Gemini request body: %s", jsonData) - body, closer, err := relaycommon.NewOutboundJSONBody(jsonData) + body, size, closer, err := relaycommon.NewOutboundJSONBody(jsonData) if err != nil { return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) } defer closer.Close() jsonData = nil + info.UpstreamRequestBodySize = size requestBody = body } @@ -268,12 +269,13 @@ func GeminiEmbeddingHandler(c *gin.Context, info *relaycommon.RelayInfo) (newAPI } } logger.LogDebug(c, "Gemini embedding request body: %s", jsonData) - body, closer, err := relaycommon.NewOutboundJSONBody(jsonData) + body, size, closer, err := relaycommon.NewOutboundJSONBody(jsonData) if err != nil { return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) } defer closer.Close() jsonData = nil + info.UpstreamRequestBodySize = size requestBody = body resp, err := adaptor.DoRequest(c, info, requestBody) diff --git a/relay/image_handler.go b/relay/image_handler.go index 690f229f3ff0..f970a43e8e20 100644 --- a/relay/image_handler.go +++ b/relay/image_handler.go @@ -51,7 +51,7 @@ func ImageHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *type if err != nil { return types.NewErrorWithStatusCode(err, types.ErrorCodeReadRequestBodyFailed, http.StatusBadRequest, types.ErrOptionWithSkipRetry()) } - requestBody = common.NewReplayableBodyReader(storage) + requestBody = common.ReaderOnly(storage) } else { convertedRequest, err := adaptor.ConvertImageRequest(c, info, *request) if err != nil { @@ -77,12 +77,13 @@ func ImageHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *type } logger.LogDebug(c, "image request body: %s", jsonData) - body, closer, err := relaycommon.NewOutboundJSONBody(jsonData) + body, size, closer, err := relaycommon.NewOutboundJSONBody(jsonData) if err != nil { return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) } defer closer.Close() jsonData = nil + info.UpstreamRequestBodySize = size requestBody = body } } diff --git a/relay/rerank_handler.go b/relay/rerank_handler.go index 460fe3056801..769feb227415 100644 --- a/relay/rerank_handler.go +++ b/relay/rerank_handler.go @@ -47,7 +47,7 @@ func RerankHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ if err != nil { return types.NewErrorWithStatusCode(err, types.ErrorCodeReadRequestBodyFailed, http.StatusBadRequest, types.ErrOptionWithSkipRetry()) } - requestBody = common.NewReplayableBodyReader(storage) + requestBody = common.ReaderOnly(storage) } else { convertedRequest, err := adaptor.ConvertRerankRequest(c, info.RelayMode, *request) if err != nil { @@ -68,12 +68,13 @@ func RerankHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ } logger.LogDebug(c, "Rerank request body: %s", jsonData) - body, closer, err := relaycommon.NewOutboundJSONBody(jsonData) + body, size, closer, err := relaycommon.NewOutboundJSONBody(jsonData) if err != nil { return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) } defer closer.Close() jsonData = nil + info.UpstreamRequestBodySize = size requestBody = body } diff --git a/relay/responses_handler.go b/relay/responses_handler.go index f8b4dec2be7b..c736c6a546b5 100644 --- a/relay/responses_handler.go +++ b/relay/responses_handler.go @@ -82,7 +82,7 @@ func ResponsesHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError * if err != nil { return types.NewError(err, types.ErrorCodeReadRequestBodyFailed, types.ErrOptionWithSkipRetry()) } - requestBody = common.NewReplayableBodyReader(storage) + requestBody = common.ReaderOnly(storage) } else { convertedRequest, err := adaptor.ConvertOpenAIResponsesRequest(c, info, *request) if err != nil { @@ -109,12 +109,13 @@ func ResponsesHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError * } logger.LogDebug(c, "requestBody: %s", jsonData) - body, closer, err := relaycommon.NewOutboundJSONBody(jsonData) + body, size, closer, err := relaycommon.NewOutboundJSONBody(jsonData) if err != nil { return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) } defer closer.Close() jsonData = nil + info.UpstreamRequestBodySize = size requestBody = body } diff --git a/router/api-router.go b/router/api-router.go index 2b0bfd6c5dd3..1da47e43b850 100644 --- a/router/api-router.go +++ b/router/api-router.go @@ -22,6 +22,24 @@ func SetApiRouter(router *gin.Engine) { apiRouter.GET("/setup", controller.GetSetup) apiRouter.POST("/setup", anonymousRequestBodyLimit, controller.PostSetup) apiRouter.GET("/status", controller.GetStatus) + // 营销站公开接口（无需登录，复用全局限流 + 关键接口关键限流） + publicRouter := apiRouter.Group("/public") + { + publicRouter.GET("/site-config", controller.GetPublicSiteConfig) + publicRouter.GET("/pricing", controller.GetPublicPricing) + publicRouter.GET("/model-categories", controller.GetPublicModelCategories) + publicRouter.POST("/sales-lead", middleware.CriticalRateLimit(), anonymousRequestBodyLimit, controller.PostPublicSalesLead) + } + // 模型商店后台管理（管理员，需 AdminAuth） + marketModelRoute := apiRouter.Group("/admin/market-models") + marketModelRoute.Use(middleware.AdminAuth()) + marketModelRoute.GET("/", controller.ListMarketModels) + marketModelRoute.GET("/:id", controller.GetMarketModel) + marketModelRoute.POST("/", controller.CreateMarketModel) + marketModelRoute.PUT("/:id", controller.UpdateMarketModel) + marketModelRoute.DELETE("/:id", controller.DeleteMarketModel) + // 模型商店公开读取（门店展示，仅已上架） + apiRouter.GET("/market-models", controller.GetPublicMarketModels) apiRouter.GET("/uptime/status", controller.GetUptimeKumaStatus) apiRouter.GET("/models", middleware.UserAuth(), controller.DashboardListModels) apiRouter.GET("/status/test", middleware.AdminAuth(), controller.TestStatus) @@ -90,7 +108,7 @@ func SetApiRouter(router *gin.Engine) { selfRoute.GET("/models", controller.GetUserModels) selfRoute.PUT("/self", middleware.CriticalRateLimit(), middleware.DisableCache(), controller.UpdateSelf) selfRoute.DELETE("/self", controller.DeleteSelf) - selfRoute.GET("/token", middleware.CriticalRateLimit(), middleware.DisableCache(), controller.GenerateAccessToken) + selfRoute.GET("/token", middleware.DisableCache(), controller.GenerateAccessToken) selfRoute.GET("/passkey", controller.PasskeyStatus) selfRoute.POST("/passkey/register/begin", middleware.DisableCache(), controller.PasskeyRegisterBegin) selfRoute.POST("/passkey/register/finish", middleware.DisableCache(), controller.PasskeyRegisterFinish) @@ -233,12 +251,15 @@ func SetApiRouter(router *gin.Engine) { } registerChannelRoutes(apiRouter) registerAuthzRoutes(apiRouter) + registerTeamRoutes(apiRouter) + registerSlaRoutes(apiRouter) + registerRegionRouteRoutes(apiRouter) + registerDistributorRoutes(apiRouter) tokenRoute := apiRouter.Group("/token") tokenRoute.Use(middleware.UserAuth()) { tokenRoute.GET("/", controller.GetAllTokens) tokenRoute.GET("/search", middleware.SearchRateLimit(), controller.SearchTokens) - tokenRoute.GET("/auto-groups", controller.GetTokenAutoGroups) tokenRoute.GET("/:id", controller.GetToken) tokenRoute.POST("/:id/key", middleware.CriticalRateLimit(), middleware.DisableCache(), controller.GetTokenKey) tokenRoute.POST("/", controller.AddToken) diff --git a/router/p2-router.go b/router/p2-router.go new file mode 100644 index 000000000000..3f7773f664a7 --- /dev/null +++ b/router/p2-router.go @@ -0,0 +1,74 @@ +package router + +import ( + "github.com/QuantumNous/new-api/controller" + "github.com/QuantumNous/new-api/middleware" + "github.com/gin-gonic/gin" +) + +// P2 平台化能力：四个功能的管理端路由 + SLA 公开状态页。 +// 前端页面与区域路由真实选渠道逻辑不在本轮范围。 + +func registerTeamRoutes(apiRouter *gin.RouterGroup) { + r := apiRouter.Group("/admin/teams") + r.Use(middleware.AdminAuth()) + { + r.GET("", controller.ListTeams) + r.POST("", controller.CreateTeam) + r.GET("/:id", controller.GetTeam) + r.PUT("/:id", controller.UpdateTeam) + r.DELETE("/:id", controller.DeleteTeam) + r.POST("/:id/members", controller.AddTeamMember) + r.GET("/:id/members", controller.ListTeamMembers) + r.DELETE("/:id/members/:user_id", controller.RemoveTeamMember) + r.POST("/:id/projects", controller.AddTeamProject) + r.GET("/:id/projects", controller.ListTeamProjects) + r.DELETE("/:id/projects/:pid", controller.RemoveTeamProject) + r.GET("/:id/billing", controller.GetTeamBilling) + } +} + +func registerSlaRoutes(apiRouter *gin.RouterGroup) { + admin := apiRouter.Group("/admin/sla-incidents") + admin.Use(middleware.AdminAuth()) + { + admin.GET("", controller.ListSlaIncidents) + admin.POST("", controller.CreateSlaIncident) + admin.GET("/:id", controller.GetSlaIncident) + admin.PUT("/:id", controller.UpdateSlaIncident) + admin.DELETE("/:id", controller.DeleteSlaIncident) + } + // 公开匿名：状态页事件列表与服务状态摘要。 + apiRouter.GET("/sla/incidents", controller.GetPublicSlaIncidents) + apiRouter.GET("/sla/status", controller.GetPublicSlaStatus) +} + +func registerRegionRouteRoutes(apiRouter *gin.RouterGroup) { + r := apiRouter.Group("/admin/region-routes") + r.Use(middleware.AdminAuth()) + { + r.GET("", controller.ListRegionRoutes) + r.POST("", controller.CreateRegionRoute) + r.GET("/:id", controller.GetRegionRoute) + r.PUT("/:id", controller.UpdateRegionRoute) + r.DELETE("/:id", controller.DeleteRegionRoute) + } +} + +func registerDistributorRoutes(apiRouter *gin.RouterGroup) { + r := apiRouter.Group("/admin/distributors") + r.Use(middleware.AdminAuth()) + { + r.GET("", controller.ListDistributors) + r.POST("", controller.CreateDistributor) + r.GET("/:id", controller.GetDistributor) + r.PUT("/:id", controller.UpdateDistributor) + r.DELETE("/:id", controller.DeleteDistributor) + r.GET("/:id/sub-users", controller.ListDistributorSubUsers) + r.GET("/:id/billing", controller.GetDistributorBilling) + r.GET("/:id/prices", controller.ListDistributorPrices) + r.POST("/:id/prices", controller.CreateDistributorPrice) + r.PUT("/:id/prices/:price_id", controller.UpdateDistributorPrice) + r.DELETE("/:id/prices/:price_id", controller.DeleteDistributorPrice) + } +} diff --git a/service/billing_usage.go b/service/billing_usage.go index 12656e693708..2766178d68c4 100644 --- a/service/billing_usage.go +++ b/service/billing_usage.go @@ -25,32 +25,36 @@ func effectiveBillingUsage(usage *dto.Usage) *dto.Usage { } func usageBillingPathForLog(isLocalCountTokens bool, usage *dto.Usage) string { - effectiveUsage, ok := usageFromBillingUsage(usage) - if !ok { - if isLocalCountTokens { - return usageBillingPathLocal - } + if isLocalCountTokens { + return usageBillingPathLocal + } + if usage == nil || usage.BillingUsage == nil { return usageBillingPathUpstream } - - switch effectiveUsage.UsageSemantic { - case dto.BillingUsageSemanticOpenAI: + source := strings.TrimSpace(usage.BillingUsage.Source) + semantic := strings.TrimSpace(usage.BillingUsage.Semantic) + if strings.EqualFold(source, dto.BillingUsageSourceOAIChat) || + strings.EqualFold(source, dto.BillingUsageSourceOAIResponses) || + strings.EqualFold(semantic, dto.BillingUsageSemanticOpenAI) { if usage.BillingUsage.Estimated { return usageBillingPathOpenAIEstimated } return usageBillingPathOpenAI - case dto.BillingUsageSemanticAnthropic: + } + if strings.EqualFold(source, dto.BillingUsageSourceClaudeMessages) || + strings.EqualFold(semantic, dto.BillingUsageSemanticAnthropic) { if usage.BillingUsage.Estimated { return usageBillingPathAnthropicEstimated } return usageBillingPathAnthropic - case dto.BillingUsageSemanticGemini: + } + if strings.EqualFold(source, dto.BillingUsageSourceGeminiChat) || + strings.EqualFold(semantic, dto.BillingUsageSemanticGemini) { if usage.BillingUsage.Estimated { return usageBillingPathGeminiEstimated } return usageBillingPathGemini } - return usageBillingPathUpstream } diff --git a/service/channel_select.go b/service/channel_select.go index 0ab88dc84ff2..5fc96dca9944 100644 --- a/service/channel_select.go +++ b/service/channel_select.go @@ -7,6 +7,7 @@ import ( "github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/logger" "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/setting" "github.com/gin-gonic/gin" ) @@ -86,11 +87,19 @@ func CacheGetRandomSatisfiedChannel(param *RetryParam) (*model.Channel, string, selectGroup := param.TokenGroup userGroup := common.GetContextKeyString(param.Ctx, constant.ContextKeyUserGroup) + // 区域路由：命中策略时收窄候选渠道并按策略排序，未命中时保持原有选渠道行为 + routing := model.ResolveRegionRouting(detectRegion(param.Ctx), param.ModelName) + if routing.Active { + common.SetContextKey(param.Ctx, constant.ContextKeyRequestRegion, routing.Region) + logger.LogDebug(param.Ctx, "region routing hit: region=%s, strategy=%s, allowed channels=%v", + routing.Region, routing.Strategy, routing.AllowedIds) + } + if param.TokenGroup == "auto" { - autoGroups := GetRequestAutoGroups(param.Ctx, userGroup) - if len(autoGroups) == 0 { + if len(setting.GetAutoGroups()) == 0 { return nil, selectGroup, errors.New("auto groups is not enabled") } + autoGroups := GetUserAutoGroup(userGroup) // startGroupIndex: the group index to start searching from // startGroupIndex: 开始搜索的分组索引 @@ -115,7 +124,7 @@ func CacheGetRandomSatisfiedChannel(param *RetryParam) (*model.Channel, string, } logger.LogDebug(param.Ctx, "Auto selecting group: %s, priorityRetry: %d", autoGroup, priorityRetry) - channel, _ = model.GetRandomSatisfiedChannel(autoGroup, param.ModelName, priorityRetry, param.RequestPath) + channel, _ = model.GetRandomSatisfiedChannel(autoGroup, param.ModelName, priorityRetry, param.RequestPath, routing) if channel == nil { // Current group has no available channel for this model, try next group // 当前分组没有该模型的可用渠道，尝试下一个分组 @@ -153,10 +162,27 @@ func CacheGetRandomSatisfiedChannel(param *RetryParam) (*model.Channel, string, break } } else { - channel, err = model.GetRandomSatisfiedChannel(param.TokenGroup, param.ModelName, param.GetRetry(), param.RequestPath) + channel, err = model.GetRandomSatisfiedChannel(param.TokenGroup, param.ModelName, param.GetRetry(), param.RequestPath, routing) if err != nil { return nil, param.TokenGroup, err } } return channel, selectGroup, nil } + +// detectRegion 解析当前请求的区域标识： +// 优先取请求头 X-Region，其次取上下文中已解析过的区域（重试时复用）， +// 最后退回已登录用户的区域偏好（User.RegionPreference）。 +// 三者皆空时返回 ""，由调用方保持原有选渠道行为。 +func detectRegion(ctx *gin.Context) string { + if ctx == nil { + return "" + } + if region := model.NormalizeRegion(ctx.GetHeader(constant.HeaderRegion)); region != "" { + return region + } + if region := model.NormalizeRegion(common.GetContextKeyString(ctx, constant.ContextKeyRequestRegion)); region != "" { + return region + } + return model.NormalizeRegion(common.GetContextKeyString(ctx, constant.ContextKeyUserRegionPreference)) +} diff --git a/service/channel_select_test.go b/service/channel_select_test.go new file mode 100644 index 000000000000..e41b172c0464 --- /dev/null +++ b/service/channel_select_test.go @@ -0,0 +1,60 @@ +package service + +import ( + "net/http" + "net/http/httptest" + "testing" + + "github.com/QuantumNous/new-api/constant" + "github.com/gin-gonic/gin" +) + +// TestDetectRegionFallbackToUserPreference 验证区域解析的优先级： +// X-Region 请求头 > 请求上下文区域 > 用户区域偏好 > 空。 +func TestDetectRegionFallbackToUserPreference(t *testing.T) { + // 仅设置用户区域偏好 + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest(http.MethodGet, "/", nil) + c.Set(string(constant.ContextKeyUserRegionPreference), "eu") + if got := detectRegion(c); got != "eu" { + t.Fatalf("expected user preference 'eu', got %q", got) + } + + // 请求头优先于用户偏好 + w2 := httptest.NewRecorder() + c2, _ := gin.CreateTestContext(w2) + c2.Request = httptest.NewRequest(http.MethodGet, "/", nil) + c2.Request.Header.Set(constant.HeaderRegion, "us") + c2.Set(string(constant.ContextKeyUserRegionPreference), "eu") + if got := detectRegion(c2); got != "us" { + t.Fatalf("expected header 'us' to win, got %q", got) + } + + // 请求上下文区域优先于用户偏好，但让位于请求头 + w3 := httptest.NewRecorder() + c3, _ := gin.CreateTestContext(w3) + c3.Request = httptest.NewRequest(http.MethodGet, "/", nil) + c3.Set(string(constant.ContextKeyRequestRegion), "cn") + c3.Set(string(constant.ContextKeyUserRegionPreference), "eu") + if got := detectRegion(c3); got != "cn" { + t.Fatalf("expected request-region 'cn' to win over preference, got %q", got) + } + + // 全部为空时返回空，保持原有选渠道行为 + w4 := httptest.NewRecorder() + c4, _ := gin.CreateTestContext(w4) + c4.Request = httptest.NewRequest(http.MethodGet, "/", nil) + if got := detectRegion(c4); got != "" { + t.Fatalf("expected empty region, got %q", got) + } + + // 用户偏好被归一化（去空格、转小写、截断超长） + w5 := httptest.NewRecorder() + c5, _ := gin.CreateTestContext(w5) + c5.Request = httptest.NewRequest(http.MethodGet, "/", nil) + c5.Set(string(constant.ContextKeyUserRegionPreference), " EU ") + if got := detectRegion(c5); got != "eu" { + t.Fatalf("expected normalized 'eu', got %q", got) + } +} diff --git a/service/group.go b/service/group.go index e792083e4a09..8cb359bcff63 100644 --- a/service/group.go +++ b/service/group.go @@ -3,12 +3,9 @@ package service import ( "strings" - "github.com/QuantumNous/new-api/common" - "github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/model" "github.com/QuantumNous/new-api/setting" "github.com/QuantumNous/new-api/setting/ratio_setting" - "github.com/gin-gonic/gin" ) func GetUserUsableGroups(userGroup string) map[string]string { @@ -45,67 +42,18 @@ func GroupInUserUsableGroups(userGroup, groupName string) bool { return ok } -func IsUserSelectableGroup(userGroup, groupName string) bool { - if groupName == "" || groupName == "auto" { - return false - } - return GroupInUserUsableGroups(userGroup, groupName) && ratio_setting.ContainsGroupRatio(groupName) -} - // GetUserAutoGroup 根据用户分组获取自动分组设置 func GetUserAutoGroup(userGroup string) []string { + groups := GetUserUsableGroups(userGroup) autoGroups := make([]string, 0) - seen := make(map[string]struct{}) for _, group := range setting.GetAutoGroups() { - if !IsUserSelectableGroup(userGroup, group) { - continue - } - if _, ok := seen[group]; ok { - continue + if _, ok := groups[group]; ok { + autoGroups = append(autoGroups, group) } - seen[group] = struct{}{} - autoGroups = append(autoGroups, group) } return autoGroups } -// FilterUserTokenAutoGroups applies current permissions before the current -// per-token limit. It intentionally does not fall back to the global Auto list. -func FilterUserTokenAutoGroups(userGroup string, groups []string) []string { - maxCount := setting.GetMaxTokenAutoGroups() - filtered := make([]string, 0, min(len(groups), maxCount)) - seen := make(map[string]struct{}) - for _, group := range groups { - if !IsUserSelectableGroup(userGroup, group) { - continue - } - if _, ok := seen[group]; ok { - continue - } - seen[group] = struct{}{} - filtered = append(filtered, group) - if len(filtered) == maxCount { - break - } - } - return filtered -} - -// GetRequestAutoGroups resolves the ordered Auto groups for the current token. -// The absence of the context value means that the token inherits the complete -// global Auto list; a present (even empty) value is an explicit token snapshot. -func GetRequestAutoGroups(c *gin.Context, userGroup string) []string { - value, ok := common.GetContextKey(c, constant.ContextKeyTokenAutoGroups) - if !ok { - return GetUserAutoGroup(userGroup) - } - groups, ok := value.([]string) - if !ok { - return []string{} - } - return FilterUserTokenAutoGroups(userGroup, groups) -} - // GetGroupsEnabledModels 按 groups 顺序获取各分组启用的模型并去重 func GetGroupsEnabledModels(groups []string) []string { seen := make(map[string]struct{}) diff --git a/service/text_quota_test.go b/service/text_quota_test.go index c9e958e7bf2b..5a908af2d485 100644 --- a/service/text_quota_test.go +++ b/service/text_quota_test.go @@ -284,18 +284,9 @@ func TestCalculateTextQuotaSummaryUsesOpenAIBillingUsageBeforeTopLevelUsage(t *t } func TestUsageBillingPathForLog(t *testing.T) { - require.Equal(t, usageBillingPathAnthropic, usageBillingPathForLog(true, &dto.Usage{ + require.Equal(t, usageBillingPathLocal, usageBillingPathForLog(true, &dto.Usage{ BillingUsage: dto.NewClaudeMessagesBillingUsage(&dto.ClaudeUsage{InputTokens: 1}), })) - invalidBillingUsage := &dto.Usage{ - PromptTokens: 1, - BillingUsage: &dto.BillingUsage{ - Source: dto.BillingUsageSourceClaudeMessages, - Semantic: dto.BillingUsageSemanticAnthropic, - }, - } - require.Equal(t, usageBillingPathLocal, usageBillingPathForLog(true, invalidBillingUsage)) - require.Equal(t, usageBillingPathUpstream, usageBillingPathForLog(false, invalidBillingUsage)) require.Equal(t, usageBillingPathUpstream, usageBillingPathForLog(false, &dto.Usage{})) require.Equal(t, usageBillingPathOpenAI, usageBillingPathForLog(false, &dto.Usage{ BillingUsage: dto.NewOpenAIChatBillingUsage(&dto.Usage{PromptTokens: 1}), @@ -306,7 +297,7 @@ func TestUsageBillingPathForLog(t *testing.T) { require.Equal(t, usageBillingPathGemini, usageBillingPathForLog(false, &dto.Usage{ BillingUsage: dto.NewGeminiChatBillingUsage(&dto.GeminiUsageMetadata{PromptTokenCount: 1}), })) - require.Equal(t, usageBillingPathGeminiEstimated, usageBillingPathForLog(true, &dto.Usage{ + require.Equal(t, usageBillingPathGeminiEstimated, usageBillingPathForLog(false, &dto.Usage{ BillingUsage: dto.NewEstimatedGeminiChatBillingUsage(&dto.Usage{PromptTokens: 1}), })) } @@ -315,7 +306,7 @@ func TestAppendUsageBillingPathForLogWritesAdminInfo(t *testing.T) { other := map[string]interface{}{ "admin_info": map[string]interface{}{}, } - appendUsageBillingPathForLog(other, true, &dto.Usage{ + appendUsageBillingPathForLog(other, false, &dto.Usage{ BillingUsage: dto.NewClaudeMessagesBillingUsage(&dto.ClaudeUsage{InputTokens: 1}), }) diff --git a/setting/auto_group.go b/setting/auto_group.go index 3b509b6026e6..9261286bca93 100644 --- a/setting/auto_group.go +++ b/setting/auto_group.go @@ -1,27 +1,15 @@ package setting import ( - "fmt" - "strconv" - "sync/atomic" - "github.com/QuantumNous/new-api/common" ) -const DefaultMaxTokenAutoGroups = 5 - var autoGroups = []string{ "default", } var DefaultUseAutoGroup = false -var maxTokenAutoGroups atomic.Int64 - -func init() { - maxTokenAutoGroups.Store(DefaultMaxTokenAutoGroups) -} - func ContainsAutoGroup(group string) bool { for _, autoGroup := range autoGroups { if autoGroup == group { @@ -47,24 +35,3 @@ func AutoGroups2JsonString() string { func GetAutoGroups() []string { return autoGroups } - -func GetMaxTokenAutoGroups() int { - return int(maxTokenAutoGroups.Load()) -} - -func ValidateMaxTokenAutoGroups(value string) error { - maxCount, err := strconv.Atoi(value) - if err != nil || maxCount <= 0 { - return fmt.Errorf("MaxTokenAutoGroups must be a positive integer") - } - return nil -} - -func UpdateMaxTokenAutoGroups(value string) error { - if err := ValidateMaxTokenAutoGroups(value); err != nil { - return err - } - maxCount, _ := strconv.Atoi(value) - maxTokenAutoGroups.Store(int64(maxCount)) - return nil -}