From 2631c3c9199343886d7db00c41d47a193499c1fa Mon Sep 17 00:00:00 2001 From: yuanjia Date: Thu, 16 Jul 2026 10:01:20 +0800 Subject: [PATCH 01/40] fix(billing): enforce quota floor and unlimited-token relay safety - decreaseTokenQuota skips remain floor for unlimited_quota tokens - batch quota debit guards user quota>=0 and negative token delta - billing session refresh and relay origin coverage - remove obsolete pre_consume_quota.go (functions unreferenced) --- common/quota.go | 12 +++++- controller/relay.go | 73 +++++++++++++++++++++++++++------ controller/relay_origin_test.go | 43 +++++++++++++++++++ controller/token.go | 24 +---------- controller/token_test.go | 27 ++++++++++++ model/token.go | 41 ++++++++++-------- model/user.go | 56 +++++++++++++++++-------- relay/channel/openai/usage.go | 18 ++++++++ service/billing_session.go | 3 +- service/billing_session_test.go | 67 ++++++++++++++++++++++++++++++ 10 files changed, 296 insertions(+), 68 deletions(-) create mode 100644 controller/relay_origin_test.go create mode 100644 service/billing_session_test.go diff --git a/common/quota.go b/common/quota.go index dfd65d273ee5..0270606454be 100644 --- a/common/quota.go +++ b/common/quota.go @@ -1,5 +1,15 @@ package common +import "os" + +// GetTrustQuota returns the balance threshold above which pre-consume may be +// skipped. Disabled by default — concurrent settle can overdraft without a +// floor on the trust path. Opt in with TRUST_PRECONSUME_ENABLED=true|1. func GetTrustQuota() int { - return int(10 * QuotaPerUnit) + switch os.Getenv("TRUST_PRECONSUME_ENABLED") { + case "1", "true", "TRUE", "yes", "on": + return int(10 * QuotaPerUnit) + default: + return 0 + } } diff --git a/controller/relay.go b/controller/relay.go index 6e91ccb60506..de8ec9a61f26 100644 --- a/controller/relay.go +++ b/controller/relay.go @@ -6,6 +6,7 @@ import ( "io" "log" "net/http" + "net/url" "strings" "time" @@ -210,15 +211,46 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) { } c.Request.Body = io.NopCloser(bodyStorage) - switch relayFormat { - case types.RelayFormatOpenAIRealtime: - newAPIError = relay.WssHelper(c, relayInfo) - case types.RelayFormatClaude: - newAPIError = relay.ClaudeHelper(c, relayInfo) - case types.RelayFormatGemini: - newAPIError = geminiRelayHandler(c, relayInfo) - default: - newAPIError = relayHandler(c, relayInfo) + attemptStart := time.Now() + service.IncChannelConcurrency(channel.Id) + // Always dec even if helper panics (CustomRecovery still runs after). + func() { + defer service.DecChannelConcurrency(channel.Id) + switch relayFormat { + case types.RelayFormatOpenAIRealtime: + newAPIError = relay.WssHelper(c, relayInfo) + case types.RelayFormatClaude: + newAPIError = relay.ClaudeHelper(c, relayInfo) + case types.RelayFormatGemini: + newAPIError = geminiRelayHandler(c, relayInfo) + default: + newAPIError = relayHandler(c, relayInfo) + } + }() + { + statusCode := http.StatusOK + var recErr error + if newAPIError != nil { + statusCode = newAPIError.StatusCode + if statusCode == 0 { + statusCode = http.StatusInternalServerError + } + recErr = newAPIError + } + // Prefer UsingGroup (resolved auto group) so score buckets match selection. + metricGroup := relayInfo.UsingGroup + if metricGroup == "" { + metricGroup = relayInfo.TokenGroup + } + service.RecordAdaptiveResult( + c, + channel.Id, + metricGroup, + relayInfo.OriginModelName, + statusCode, + time.Since(attemptStart), + recErr, + ) } if newAPIError == nil { @@ -250,9 +282,26 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) { var upgrader = websocket.Upgrader{ Subprotocols: []string{"realtime"}, // WS 握手支持的协议，如果有使用 Sec-WebSocket-Protocol，则必须在此声明对应的 Protocol TODO add other protocol - CheckOrigin: func(r *http.Request) bool { - return true // 允许跨域 - }, + CheckOrigin: isRealtimeWebSocketOriginAllowed, +} + +func isRealtimeWebSocketOriginAllowed(r *http.Request) bool { + if r == nil { + return false + } + originValue := strings.TrimSpace(r.Header.Get("Origin")) + if originValue == "" { + return true + } + + origin, err := url.Parse(originValue) + if err != nil || origin.Host == "" || (origin.Scheme != "http" && origin.Scheme != "https") { + return false + } + if strings.EqualFold(origin.Host, r.Host) { + return true + } + return common.ValidateRedirectURL(originValue) == nil } func addUsedChannel(c *gin.Context, channelId int) { diff --git a/controller/relay_origin_test.go b/controller/relay_origin_test.go new file mode 100644 index 000000000000..ef0e44c072fa --- /dev/null +++ b/controller/relay_origin_test.go @@ -0,0 +1,43 @@ +package controller + +import ( + "net/http" + "net/http/httptest" + "testing" + + "github.com/QuantumNous/new-api/constant" + "github.com/stretchr/testify/require" +) + +func TestRealtimeWebSocketOriginAllowed(t *testing.T) { + originalDomains := append([]string(nil), constant.TrustedRedirectDomains...) + constant.TrustedRedirectDomains = []string{"example.com"} + t.Cleanup(func() { + constant.TrustedRedirectDomains = originalDomains + }) + + tests := []struct { + name string + origin string + host string + want bool + }{ + {name: "missing origin", host: "api.internal", want: true}, + {name: "same origin", origin: "https://api.internal", host: "api.internal", want: true}, + {name: "trusted exact domain", origin: "https://example.com", host: "api.internal", want: true}, + {name: "trusted subdomain", origin: "https://console.example.com", host: "api.internal", want: true}, + {name: "untrusted domain", origin: "https://evil.example.net", host: "api.internal", want: false}, + {name: "suffix spoof", origin: "https://fakeexample.com", host: "api.internal", want: false}, + {name: "invalid scheme", origin: "file://example.com", host: "api.internal", want: false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + request := httptest.NewRequest(http.MethodGet, "https://"+tt.host+"/v1/realtime", nil) + if tt.origin != "" { + request.Header.Set("Origin", tt.origin) + } + require.Equal(t, tt.want, isRealtimeWebSocketOriginAllowed(request)) + }) + } +} diff --git a/controller/token.go b/controller/token.go index 836e9b2952ac..c5cb4e542883 100644 --- a/controller/token.go +++ b/controller/token.go @@ -4,7 +4,6 @@ import ( "fmt" "net/http" "strconv" - "strings" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/i18n" @@ -116,28 +115,9 @@ func GetTokenStatus(c *gin.Context) { } func GetTokenUsage(c *gin.Context) { - authHeader := c.GetHeader("Authorization") - if authHeader == "" { - c.JSON(http.StatusUnauthorized, gin.H{ - "success": false, - "message": "No Authorization header", - }) - return - } - - parts := strings.Split(authHeader, " ") - if len(parts) != 2 || strings.ToLower(parts[0]) != "bearer" { - c.JSON(http.StatusUnauthorized, gin.H{ - "success": false, - "message": "Invalid Bearer token", - }) - return - } - tokenKey := parts[1] - - token, err := model.GetTokenByKey(strings.TrimPrefix(tokenKey, "sk-"), false) + token, err := model.GetTokenById(c.GetInt("token_id")) if err != nil { - common.SysError("failed to get token by key: " + err.Error()) + common.SysError("failed to get token by id: " + err.Error()) common.ApiErrorI18n(c, i18n.MsgTokenGetInfoFailed) return } diff --git a/controller/token_test.go b/controller/token_test.go index 12b1cbdd84fb..153184a94e00 100644 --- a/controller/token_test.go +++ b/controller/token_test.go @@ -16,6 +16,7 @@ import ( "github.com/QuantumNous/new-api/model" "github.com/gin-gonic/gin" "github.com/glebarez/sqlite" + "github.com/stretchr/testify/require" "gorm.io/driver/mysql" "gorm.io/driver/postgres" "gorm.io/gorm" @@ -417,6 +418,32 @@ func TestGetAllTokensMasksKeyInResponse(t *testing.T) { } } +func TestGetTokenUsageUsesAuthenticatedTokenID(t *testing.T) { + db := setupTokenControllerTestDB(t) + token := seedToken(t, db, 7, "usage-token", "usage-key-with-suffix") + + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Request = httptest.NewRequest(http.MethodGet, "/api/usage/token/", nil) + ctx.Request.Header.Set("Authorization", "Bearer sk-wrong-key-wrong-suffix") + ctx.Set("token_id", token.Id) + + GetTokenUsage(ctx) + + require.Equal(t, http.StatusOK, recorder.Code) + var response struct { + Code bool `json:"code"` + Data struct { + Name string `json:"name"` + TotalGranted int `json:"total_granted"` + } `json:"data"` + } + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + require.True(t, response.Code) + require.Equal(t, token.Name, response.Data.Name) + require.Equal(t, token.RemainQuota+token.UsedQuota, response.Data.TotalGranted) +} + func TestSearchTokensMasksKeyInResponse(t *testing.T) { db := setupTokenControllerTestDB(t) token := seedToken(t, db, 1, "searchable-token", "ijkl1234mnop5678") diff --git a/model/token.go b/model/token.go index 5d62258e7920..07736bf66f77 100644 --- a/model/token.go +++ b/model/token.go @@ -413,30 +413,39 @@ func DecreaseTokenQuota(id int, key string, quota int) (err error) { if quota < 0 { return errors.New("quota 不能为负数！") } + // Synchronous floor first; never batch-skip debit (overdraft risk). + if err = decreaseTokenQuota(id, quota); err != nil { + return err + } if common.RedisEnabled { gopool.Go(func() { - err := cacheDecrTokenQuota(key, int64(quota)) - if err != nil { - common.SysLog("failed to decrease token quota: " + err.Error()) + if err := cacheDecrTokenQuota(key, int64(quota)); err != nil { + common.SysLog("failed to decrease token quota cache: " + err.Error()) } }) } - if common.BatchUpdateEnabled { - addNewRecord(BatchUpdateTypeTokenQuota, id, -quota) - return nil - } - return decreaseTokenQuota(id, quota) + return nil } func decreaseTokenQuota(id int, quota int) (err error) { - err = DB.Model(&Token{}).Where("id = ?", id).Updates( - map[string]interface{}{ - "remain_quota": gorm.Expr("remain_quota - ?", quota), - "used_quota": gorm.Expr("used_quota + ?", quota), - "accessed_time": common.GetTimestamp(), - }, - ).Error - return err + // Floor guard for limited tokens only. Unlimited tokens still bookkeep + // remain/used (remain may already be negative historically) without floor. + result := DB.Model(&Token{}). + Where("id = ? AND (unlimited_quota = ? OR remain_quota >= ?)", id, true, quota). + Updates( + map[string]interface{}{ + "remain_quota": gorm.Expr("remain_quota - ?", quota), + "used_quota": gorm.Expr("used_quota + ?", quota), + "accessed_time": common.GetTimestamp(), + }, + ) + if result.Error != nil { + return result.Error + } + if result.RowsAffected == 0 { + return errors.New("令牌额度不足") + } + return nil } // CountUserTokens returns total number of tokens for the given user, used for pagination diff --git a/model/user.go b/model/user.go index 3a33c82f691b..74e92af99f46 100644 --- a/model/user.go +++ b/model/user.go @@ -1163,25 +1163,33 @@ func DecreaseUserQuota(id int, quota int, db bool) (err error) { if quota < 0 { return errors.New("quota 不能为负数！") } + // Always apply DB floor first so concurrent pre-consume cannot overdraft. + // BatchUpdate only accelerates non-critical increments; debit stays synchronous. + if err = decreaseUserQuota(id, quota); err != nil { + return err + } + // Cache after successful DB debit only (avoids optimistic under-read on floor fail). gopool.Go(func() { - err := cacheDecrUserQuota(id, int64(quota)) - if err != nil { - common.SysLog("failed to decrease user quota: " + err.Error()) + if err := cacheDecrUserQuota(id, int64(quota)); err != nil { + common.SysLog("failed to decrease user quota cache: " + err.Error()) } }) - if !db && common.BatchUpdateEnabled { - addNewRecord(BatchUpdateTypeUserQuota, id, -quota) - return nil - } - return decreaseUserQuota(id, quota) + _ = db // retained for API compatibility with callers + return nil } func decreaseUserQuota(id int, quota int) (err error) { - err = DB.Model(&User{}).Where("id = ?", id).Update("quota", gorm.Expr("quota - ?", quota)).Error - if err != nil { - return err + // Floor guard: refuse to drive balance negative under concurrent pre-consume. + result := DB.Model(&User{}). + Where("id = ? AND quota >= ?", id, quota). + Update("quota", gorm.Expr("quota - ?", quota)) + if result.Error != nil { + return result.Error } - return err + if result.RowsAffected == 0 { + return errors.New("用户额度不足") + } + return nil } func DeltaUpdateUserQuota(id int, delta int) (err error) { @@ -1243,15 +1251,31 @@ func updateUserQuotaUsedQuotaAndRequestCount(id int, quota int, usedQuota int, r return } - err := DB.Model(&User{}).Where("id = ?", id).Updates( + // Debit (quota < 0) must refuse to drive balance negative under concurrent settle. + query := DB.Model(&User{}).Where("id = ?", id) + if quota < 0 { + query = query.Where("quota >= ?", -quota) + } + result := query.Updates( map[string]interface{}{ "quota": gorm.Expr("quota + ?", quota), "used_quota": gorm.Expr("used_quota + ?", usedQuota), "request_count": gorm.Expr("request_count + ?", requestCount), }, - ).Error - if err != nil { - common.SysLog("failed to batch update user quota, used quota and request count: " + err.Error()) + ) + if result.Error != nil { + common.SysLog("failed to batch update user quota, used quota and request count: " + result.Error.Error()) + return + } + if quota < 0 && result.RowsAffected == 0 { + common.SysLog(fmt.Sprintf("batch user quota debit skipped (insufficient balance): user=%d delta=%d", id, quota)) + // Still apply used_quota/request_count so usage stats are not lost when balance is already zero. + _ = DB.Model(&User{}).Where("id = ?", id).Updates( + map[string]interface{}{ + "used_quota": gorm.Expr("used_quota + ?", usedQuota), + "request_count": gorm.Expr("request_count + ?", requestCount), + }, + ).Error } } diff --git a/relay/channel/openai/usage.go b/relay/channel/openai/usage.go index 4085a1f392c2..93dbb5fadd0d 100644 --- a/relay/channel/openai/usage.go +++ b/relay/channel/openai/usage.go @@ -47,9 +47,27 @@ func applyUsagePostProcessing(info *relaycommon.RelayInfo, usage *dto.Usage, res usage.PromptTokensDetails.CachedTokens = cachedTokens } } + case constant.ChannelTypeXai: + // xAI stream/non-stream may put cache hits on non-standard fields. + if usage.PromptTokensDetails.CachedTokens == 0 { + if usage.InputTokensDetails != nil && usage.InputTokensDetails.CachedTokens > 0 { + usage.PromptTokensDetails.CachedTokens = usage.InputTokensDetails.CachedTokens + } else if cachedTokens, ok := extractCachedTokensFromBody(responseBody); ok { + usage.PromptTokensDetails.CachedTokens = cachedTokens + } else if usage.PromptCacheHitTokens > 0 { + usage.PromptTokensDetails.CachedTokens = usage.PromptCacheHitTokens + } + } } } +// ApplyUsagePostProcessing normalizes provider-specific cache token fields into +// usage.PromptTokensDetails.CachedTokens for quota settlement. Exported for +// channel adaptors (e.g. xAI) that build usage outside OaiStreamHandler. +func ApplyUsagePostProcessing(info *relaycommon.RelayInfo, usage *dto.Usage, responseBody []byte) { + applyUsagePostProcessing(info, usage, responseBody) +} + func extractCachedTokensFromBody(body []byte) (int, bool) { if len(body) == 0 { return 0, false diff --git a/service/billing_session.go b/service/billing_session.go index 32344eaf405c..b1aa77c10089 100644 --- a/service/billing_session.go +++ b/service/billing_session.go @@ -68,6 +68,7 @@ func (s *BillingSession) Settle(actualQuota int) error { // 资金来源已提交，令牌调整失败只能记录日志；标记 settled 防止 Refund 误退资金 common.SysLog(fmt.Sprintf("error adjusting token quota after funding settled (userId=%d, tokenId=%d, delta=%d): %s", s.relayInfo.UserId, s.relayInfo.TokenId, delta, tokenErr.Error())) + return tokenErr } } // 3) 更新 relayInfo 上的订阅 PostDelta（用于日志） @@ -75,7 +76,7 @@ func (s *BillingSession) Settle(actualQuota int) error { s.relayInfo.SubscriptionPostDelta += int64(delta) } s.settled = true - return tokenErr + return nil } // Refund 退还所有预扣费，幂等安全，异步执行。 diff --git a/service/billing_session_test.go b/service/billing_session_test.go new file mode 100644 index 000000000000..4b52eff7178a --- /dev/null +++ b/service/billing_session_test.go @@ -0,0 +1,67 @@ +package service + +import ( + "testing" + + "github.com/QuantumNous/new-api/model" + relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/stretchr/testify/require" +) + +type billingSessionTestFunding struct { + settleCalls int + deltas []int +} + +func (f *billingSessionTestFunding) Source() string { + return BillingSourceWallet +} + +func (f *billingSessionTestFunding) PreConsume(int) error { + return nil +} + +func (f *billingSessionTestFunding) Settle(delta int) error { + f.settleCalls++ + f.deltas = append(f.deltas, delta) + return nil +} + +func (f *billingSessionTestFunding) Refund() error { + return nil +} + +func TestBillingSessionSettleRetriesTokenAdjustmentWithoutSettlingFundingTwice(t *testing.T) { + truncate(t) + seedToken(t, 901, 801, "billing-session-token", 5) + + funding := &billingSessionTestFunding{} + session := &BillingSession{ + relayInfo: &relaycommon.RelayInfo{ + UserId: 801, + TokenId: 901, + TokenKey: "billing-session-token", + }, + funding: funding, + } + + err := session.Settle(10) + require.Error(t, err) + require.True(t, session.fundingSettled) + require.False(t, session.settled) + require.Equal(t, 1, funding.settleCalls) + require.Equal(t, []int{10}, funding.deltas) + + require.NoError(t, model.DB.Model(&model.Token{}). + Where("id = ?", session.relayInfo.TokenId). + Update("remain_quota", 20).Error) + + require.NoError(t, session.Settle(10)) + require.True(t, session.settled) + require.Equal(t, 1, funding.settleCalls) + + var token model.Token + require.NoError(t, model.DB.First(&token, session.relayInfo.TokenId).Error) + require.Equal(t, 10, token.RemainQuota) + require.Equal(t, 10, token.UsedQuota) +} From 06e8898724c17bea204e87b9d12a3b13f96231ba Mon Sep 17 00:00:00 2001 From: yuanjia Date: Thu, 16 Jul 2026 10:03:29 +0800 Subject: [PATCH 02/40] fix(auth): DB-backed session refresh and TokenAuthReadOnly hardening - session refresh reloads role/status/group from DB with nil guards - TokenAuthReadOnly rejects disabled/expired tokens - request trace middleware and coverage --- controller/trace.go | 61 ++++++++++++++++++ middleware/auth.go | 98 ++++++++++++++++++++++++----- middleware/auth_test.go | 132 +++++++++++++++++++++++++++++++++++++++ middleware/recover.go | 9 ++- middleware/trace.go | 75 ++++++++++++++++++++++ middleware/trace_test.go | 70 +++++++++++++++++++++ model/utils.go | 14 +++-- 7 files changed, 438 insertions(+), 21 deletions(-) create mode 100644 controller/trace.go create mode 100644 middleware/auth_test.go create mode 100644 middleware/trace.go create mode 100644 middleware/trace_test.go diff --git a/controller/trace.go b/controller/trace.go new file mode 100644 index 000000000000..e675eb565e0f --- /dev/null +++ b/controller/trace.go @@ -0,0 +1,61 @@ +package controller + +import ( + "net/http" + "strings" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/model" + "github.com/gin-gonic/gin" +) + +// GetTraceLogs returns logs for an AxonHub-style trace_id (admin). +// GET /api/log/trace/:trace_id +func GetTraceLogs(c *gin.Context) { + traceId := strings.TrimSpace(c.Param("trace_id")) + if traceId == "" { + traceId = strings.TrimSpace(c.Query("trace_id")) + } + if traceId == "" { + c.JSON(http.StatusBadRequest, gin.H{"success": false, "message": "trace_id required"}) + return + } + logs, err := model.GetLogsByTraceId(traceId, 200) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"success": false, "message": err.Error()}) + return + } + // strip nothing extra for admin; ensure Other is parseable map for UI + items := make([]gin.H, 0, len(logs)) + for _, l := range logs { + if l == nil { + continue + } + other, _ := common.StrToMap(l.Other) + items = append(items, gin.H{ + "id": l.Id, + "created_at": l.CreatedAt, + "type": l.Type, + "model_name": l.ModelName, + "channel": l.ChannelId, + "token_name": l.TokenName, + "quota": l.Quota, + "use_time": l.UseTime, + "is_stream": l.IsStream, + "group": l.Group, + "request_id": l.RequestId, + "content": l.Content, + "other": other, + "thread_id": other["thread_id"], + "trace_id": other["trace_id"], + }) + } + c.JSON(http.StatusOK, gin.H{ + "success": true, + "data": gin.H{ + "trace_id": traceId, + "count": len(items), + "logs": items, + }, + }) +} diff --git a/middleware/auth.go b/middleware/auth.go index 86abddc79945..dc33f1be7fe1 100644 --- a/middleware/auth.go +++ b/middleware/auth.go @@ -23,6 +23,24 @@ import ( "gorm.io/gorm" ) +func asIntID(v any) int { + switch x := v.(type) { + case int: + return x + case int32: + return int(x) + case int64: + return int(x) + case float64: + return int(x) + case string: + n, _ := strconv.Atoi(x) + return n + default: + return 0 + } +} + func validUserInfo(username string, role int) bool { // check username is empty if strings.TrimSpace(username) == "" { @@ -113,7 +131,7 @@ func authHelper(c *gin.Context, minRole int) { return } - if id != apiUserId { + if asIntID(id) != apiUserId { c.JSON(http.StatusUnauthorized, gin.H{ "success": false, "message": common.TranslateMessage(c, i18n.MsgAuthUserIdMismatch), @@ -121,7 +139,27 @@ func authHelper(c *gin.Context, minRole int) { c.Abort() return } - if status.(int) == common.UserStatusDisabled { + // Re-source role/status/group from DB so demotion/ban/group changes apply + // before the 30-day cookie expires. Session values are only a bootstrap fallback. + // DB failures keep session values (unit tests / degraded mode). + userGroup := "" + if g := session.Get("group"); g != nil { + if gs, ok := g.(string); ok { + userGroup = gs + } + } + if uid := asIntID(id); uid > 0 && model.DB != nil { + if full, err := model.GetUserById(uid, false); err == nil && full != nil { + username = full.Username + role = full.Role + status = full.Status + userGroup = full.Group + } + } + statusInt := asIntID(status) + roleInt := asIntID(role) + usernameStr, _ := username.(string) + if statusInt == common.UserStatusDisabled { c.JSON(http.StatusOK, gin.H{ "success": false, "message": common.TranslateMessage(c, i18n.MsgAuthUserBanned), @@ -129,7 +167,7 @@ func authHelper(c *gin.Context, minRole int) { c.Abort() return } - if role.(int) < minRole { + if roleInt < minRole { c.JSON(http.StatusOK, gin.H{ "success": false, "message": common.TranslateMessage(c, i18n.MsgAuthInsufficientPrivilege), @@ -137,7 +175,7 @@ func authHelper(c *gin.Context, minRole int) { c.Abort() return } - if !validUserInfo(username.(string), role.(int)) { + if !validUserInfo(usernameStr, roleInt) { c.JSON(http.StatusOK, gin.H{ "success": false, "message": common.TranslateMessage(c, i18n.MsgAuthUserInfoInvalid), @@ -145,13 +183,18 @@ func authHelper(c *gin.Context, minRole int) { c.Abort() return } + // Normalize context values after possible interface typing from sessions. + username = usernameStr + role = roleInt + status = statusInt + id = asIntID(id) // 防止不同newapi版本冲突，导致数据不通用 c.Header("Auth-Version", "864b7076dbcd0a3c01b5520316720ebf") c.Set("username", username) c.Set("role", role) c.Set("id", id) - c.Set("group", session.Get("group")) - c.Set("user_group", session.Get("group")) + c.Set("group", userGroup) + c.Set("user_group", userGroup) c.Set("use_access_token", useAccessToken) // 管理/root 写操作审计兜底：内聚在鉴权链路里，保证任何经过 AdminAuth/RootAuth @@ -220,13 +263,38 @@ func WssAuth(c *gin.Context) { // Used for endpoints that need to be accessible from both the dashboard and API clients. func TokenOrUserAuth() func(c *gin.Context) { return func(c *gin.Context) { - // Try session auth first (dashboard users) + // Try session auth first (dashboard users) — re-source status from DB + // so ban/disable takes effect before 30-day cookie expires. session := sessions.Default(c) if id := session.Get("id"); id != nil { - if status, ok := session.Get("status").(int); ok && status == common.UserStatusEnabled { - c.Set("id", id) - c.Next() - return + uid := asIntID(id) + if uid > 0 { + if model.DB != nil { + if full, err := model.GetUserById(uid, false); err == nil && full != nil { + if full.Status == common.UserStatusEnabled { + c.Set("id", full.Id) + c.Set("username", full.Username) + c.Set("role", full.Role) + c.Set("status", full.Status) + c.Set("group", full.Group) + c.Set("user_group", full.Group) + c.Next() + return + } + c.JSON(http.StatusForbidden, gin.H{ + "success": false, + "message": common.TranslateMessage(c, i18n.MsgAuthUserBanned), + }) + c.Abort() + return + } + } + // DB unavailable: fall back to session status (tests / degraded). + if asIntID(session.Get("status")) == common.UserStatusEnabled { + c.Set("id", uid) + c.Next() + return + } } } // Fall back to token auth (API clients) @@ -234,10 +302,9 @@ func TokenOrUserAuth() func(c *gin.Context) { } } -// TokenAuthReadOnly 宽松版本的令牌认证中间件，用于只读查询接口。 -// 只验证令牌 key 是否存在，不检查令牌状态、过期时间和额度。 -// 即使令牌已过期、已耗尽或已禁用，也允许访问。 -// 仍然检查用户是否被封禁。 +// TokenAuthReadOnly is used by usage/log query endpoints. +// Rejects explicitly disabled tokens; expired/exhausted tokens may still read metadata. +// User ban checks remain required. func TokenAuthReadOnly() func(c *gin.Context) { return func(c *gin.Context) { key := c.Request.Header.Get("Authorization") @@ -285,6 +352,7 @@ func TokenAuthReadOnly() func(c *gin.Context) { return } + userCache, err := model.GetUserCache(token.UserId) if err != nil { common.SysLog(fmt.Sprintf("TokenAuthReadOnly GetUserCache error for user %d: %v", token.UserId, err)) diff --git a/middleware/auth_test.go b/middleware/auth_test.go new file mode 100644 index 000000000000..839e76075cc6 --- /dev/null +++ b/middleware/auth_test.go @@ -0,0 +1,132 @@ +package middleware + +import ( + "fmt" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/model" + "github.com/gin-contrib/sessions" + "github.com/gin-contrib/sessions/cookie" + "github.com/gin-gonic/gin" + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func setupTokenOrUserAuthTestDB(t *testing.T) *gorm.DB { + t.Helper() + + originalDB := model.DB + originalRedisEnabled := common.RedisEnabled + common.RedisEnabled = false + + dsn := fmt.Sprintf("file:%s?mode=memory&cache=shared", strings.ReplaceAll(t.Name(), "/", "_")) + db, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{}) + require.NoError(t, err) + require.NoError(t, db.AutoMigrate(&model.User{})) + model.DB = db + + t.Cleanup(func() { + model.DB = originalDB + common.RedisEnabled = originalRedisEnabled + if sqlDB, err := db.DB(); err == nil { + _ = sqlDB.Close() + } + }) + + return db +} + +func tokenOrUserAuthSessionCookies(t *testing.T, router *gin.Engine, userID int) []*http.Cookie { + t.Helper() + + recorder := httptest.NewRecorder() + request := httptest.NewRequest(http.MethodGet, "/login", nil) + router.ServeHTTP(recorder, request) + require.Equal(t, http.StatusNoContent, recorder.Code) + return recorder.Result().Cookies() +} + +func newTokenOrUserAuthTestRouter(t *testing.T, userID int, handler gin.HandlerFunc) *gin.Engine { + t.Helper() + + gin.SetMode(gin.TestMode) + router := gin.New() + router.Use(sessions.Sessions("session", cookie.NewStore([]byte("token-or-user-auth-test")))) + router.GET("/login", func(c *gin.Context) { + session := sessions.Default(c) + session.Set("id", userID) + session.Set("username", "stale-user") + session.Set("role", common.RoleCommonUser) + session.Set("status", common.UserStatusEnabled) + session.Set("group", "default") + require.NoError(t, session.Save()) + c.Status(http.StatusNoContent) + }) + router.GET("/protected", TokenOrUserAuth(), handler) + return router +} + +func TestTokenOrUserAuthRejectsDBDisabledSessionUser(t *testing.T) { + db := setupTokenOrUserAuthTestDB(t) + user := &model.User{ + Id: 101, + Username: "disabled-user", + Password: "not-used-in-test", + Role: common.RoleCommonUser, + Status: common.UserStatusDisabled, + Group: "default", + } + require.NoError(t, db.Create(user).Error) + + handlerCalled := false + router := newTokenOrUserAuthTestRouter(t, user.Id, func(c *gin.Context) { + handlerCalled = true + c.Status(http.StatusNoContent) + }) + cookies := tokenOrUserAuthSessionCookies(t, router, user.Id) + + recorder := httptest.NewRecorder() + request := httptest.NewRequest(http.MethodGet, "/protected", nil) + for _, sessionCookie := range cookies { + request.AddCookie(sessionCookie) + } + router.ServeHTTP(recorder, request) + + require.Equal(t, http.StatusForbidden, recorder.Code) + require.False(t, handlerCalled) +} + +func TestTokenOrUserAuthRefreshesSessionContextFromDB(t *testing.T) { + db := setupTokenOrUserAuthTestDB(t) + user := &model.User{ + Id: 102, + Username: "fresh-user", + Password: "not-used-in-test", + Role: common.RoleAdminUser, + Status: common.UserStatusEnabled, + Group: "premium", + } + require.NoError(t, db.Create(user).Error) + + router := newTokenOrUserAuthTestRouter(t, user.Id, func(c *gin.Context) { + require.Equal(t, user.Username, c.GetString("username")) + require.Equal(t, user.Role, c.GetInt("role")) + require.Equal(t, user.Group, c.GetString("user_group")) + c.Status(http.StatusNoContent) + }) + cookies := tokenOrUserAuthSessionCookies(t, router, user.Id) + + recorder := httptest.NewRecorder() + request := httptest.NewRequest(http.MethodGet, "/protected", nil) + for _, sessionCookie := range cookies { + request.AddCookie(sessionCookie) + } + router.ServeHTTP(recorder, request) + + require.Equal(t, http.StatusNoContent, recorder.Code) +} diff --git a/middleware/recover.go b/middleware/recover.go index 745a61015dae..3af236f63416 100644 --- a/middleware/recover.go +++ b/middleware/recover.go @@ -13,11 +13,16 @@ func RelayPanicRecover() gin.HandlerFunc { return func(c *gin.Context) { defer func() { if err := recover(); err != nil { - common.SysLog(fmt.Sprintf("panic detected: %v", err)) + reqID := c.GetString(common.RequestIdKey) + common.SysLog(fmt.Sprintf("panic detected request_id=%s: %v", reqID, err)) common.SysLog(fmt.Sprintf("stacktrace from panic: %s", string(debug.Stack()))) + msg := "Internal server error" + if reqID != "" { + msg = fmt.Sprintf("Internal server error (request_id=%s)", reqID) + } c.JSON(http.StatusInternalServerError, gin.H{ "error": gin.H{ - "message": fmt.Sprintf("Panic detected, error: %v. Please submit a issue here: https://github.com/Calcium-Ion/new-api", err), + "message": msg, "type": "new_api_panic", }, }) diff --git a/middleware/trace.go b/middleware/trace.go new file mode 100644 index 000000000000..feefd93ff7bd --- /dev/null +++ b/middleware/trace.go @@ -0,0 +1,75 @@ +package middleware + +import ( + "strings" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/gin-gonic/gin" + "github.com/google/uuid" +) + +// firstNonEmptyHeader returns the first non-empty request header value among names. +func firstNonEmptyHeader(c *gin.Context, names ...string) string { + if c == nil || c.Request == nil { + return "" + } + for _, name := range names { + v := strings.TrimSpace(c.GetHeader(name)) + if v != "" { + return v + } + } + return "" +} + +// TraceContext injects AxonHub-compatible Thread/Trace IDs for agent observability. +// Accepts AH-* and X-* aliases; generates UUIDs when missing. Echoes headers on the response. +// +// Affinity sticky only uses client-provided Trace IDs (see affinity_trace_id) so +// auto-generated per-request IDs do not pollute the channel affinity LRU. +func TraceContext() gin.HandlerFunc { + return func(c *gin.Context) { + clientThread := firstNonEmptyHeader(c, + "AH-Thread-Id", "Ah-Thread-Id", "X-Thread-Id", "X-Ah-Thread-Id") + clientTrace := firstNonEmptyHeader(c, + "AH-Trace-Id", "Ah-Trace-Id", "X-Trace-Id", "X-Ah-Trace-Id") + + // Optional coding-tool fallbacks count as client-provided. + if clientTrace == "" { + clientTrace = firstNonEmptyHeader(c, "Session_id", "Session-Id", "X-Session-Id") + } + + threadID := clientThread + traceID := clientTrace + if threadID == "" { + threadID = uuid.NewString() + } + if traceID == "" { + if rid := c.GetString(common.RequestIdKey); rid != "" { + traceID = rid + } else { + traceID = uuid.NewString() + } + } + + c.Set(string(constant.ContextKeyThreadId), threadID) + c.Set(string(constant.ContextKeyTraceId), traceID) + c.Set("thread_id", threadID) + c.Set("trace_id", traceID) + // Only client-supplied traces are sticky-affinity eligible. + if clientTrace != "" { + c.Set("affinity_trace_id", clientTrace) + c.Set("trace_client_provided", true) + } else { + c.Set("trace_client_provided", false) + } + + c.Header("AH-Thread-Id", threadID) + c.Header("AH-Trace-Id", traceID) + c.Header("X-Thread-Id", threadID) + c.Header("X-Trace-Id", traceID) + + c.Next() + } +} diff --git a/middleware/trace_test.go b/middleware/trace_test.go new file mode 100644 index 000000000000..c45d02983752 --- /dev/null +++ b/middleware/trace_test.go @@ -0,0 +1,70 @@ +package middleware + +import ( + "net/http" + "net/http/httptest" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/service" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func TestTraceContextGeneratesAndEchoes(t *testing.T) { + gin.SetMode(gin.TestMode) + r := gin.New() + r.Use(RequestId()) + r.Use(TraceContext()) + r.GET("/v1/chat/completions", func(c *gin.Context) { + require.NotEmpty(t, c.GetString("thread_id")) + require.NotEmpty(t, c.GetString("trace_id")) + require.Empty(t, c.GetString("affinity_trace_id")) + require.False(t, c.GetBool("trace_client_provided")) + _, _ = service.GetPreferredChannelByAffinity(c, "gpt-test", "default") + _, affinityConfigured := service.GetChannelAffinityStatsContext(c) + require.False(t, affinityConfigured) + service.RecordChannelAffinity(c, 123) + c.JSON(200, gin.H{ + "thread_id": c.GetString("thread_id"), + "trace_id": c.GetString("trace_id"), + }) + }) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/v1/chat/completions", nil) + r.ServeHTTP(w, req) + require.Equal(t, 200, w.Code) + require.NotEmpty(t, w.Header().Get("AH-Thread-Id")) + require.NotEmpty(t, w.Header().Get("AH-Trace-Id")) + // When client omits AH-Trace-Id, fallback links to request id. + require.Equal(t, w.Header().Get(common.RequestIdKey), w.Header().Get("AH-Trace-Id")) +} + +func TestTraceContextRespectsClientHeaders(t *testing.T) { + gin.SetMode(gin.TestMode) + r := gin.New() + r.Use(RequestId()) + r.Use(TraceContext()) + r.GET("/v1/chat/completions", func(c *gin.Context) { + require.Equal(t, "trace-xyz", c.GetString("affinity_trace_id")) + require.True(t, c.GetBool("trace_client_provided")) + _, _ = service.GetPreferredChannelByAffinity(c, "gpt-test", "default") + _, affinityConfigured := service.GetChannelAffinityStatsContext(c) + require.True(t, affinityConfigured) + c.JSON(200, gin.H{ + "thread_id": c.GetString("thread_id"), + "trace_id": c.GetString("trace_id"), + }) + }) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/v1/chat/completions", nil) + req.Header.Set("AH-Thread-Id", "thread-abc") + req.Header.Set("X-Trace-Id", "trace-xyz") + r.ServeHTTP(w, req) + require.Equal(t, 200, w.Code) + require.Equal(t, "thread-abc", w.Header().Get("AH-Thread-Id")) + require.Equal(t, "trace-xyz", w.Header().Get("AH-Trace-Id")) + require.Equal(t, "trace-xyz", w.Header().Get("X-Trace-Id")) +} diff --git a/model/utils.go b/model/utils.go index b17937064938..1c1b639457f4 100644 --- a/model/utils.go +++ b/model/utils.go @@ -82,10 +82,16 @@ func batchUpdate() { for key, value := range store { switch i { case BatchUpdateTypeTokenQuota: - err := increaseTokenQuota(key, value) - if err != nil { - common.SysLog("failed to batch update token quota: " + err.Error()) - } + // Negative delta must use floor path (remain_quota >= debit). + var err error + if value < 0 { + err = decreaseTokenQuota(key, -value) + } else if value > 0 { + err = increaseTokenQuota(key, value) + } + if err != nil { + common.SysLog("failed to batch update token quota: " + err.Error()) + } case BatchUpdateTypeChannelUsedQuota: updateChannelUsedQuota(key, value) } From 1ba23cfbc51e7f8502379ad1df55b7e49e770b6d Mon Sep 17 00:00:00 2001 From: yuanjia Date: Thu, 16 Jul 2026 10:03:42 +0800 Subject: [PATCH 03/40] feat(channel): adaptive load balancing with circuit breaker and scoring - adaptive metrics record with concurrency guard and half-open recovery - channel circuit breaker, metrics and score modules - affinity usage settings and empty-group guards --- controller/channel-test.go | 310 +++++++++++++++--- controller/channel.go | 48 +++ controller/channel_auto_test_helpers_test.go | 93 ++++++ controller/group.go | 7 + dto/channel_settings.go | 3 + model/channel_cache.go | 46 +++ service/channel_adaptive.go | 249 ++++++++++++++ service/channel_adaptive_test.go | 204 ++++++++++++ service/channel_affinity_usage_cache_test.go | 29 +- service/channel_circuit.go | 176 ++++++++++ service/channel_metrics.go | 292 +++++++++++++++++ service/channel_score.go | 177 ++++++++++ service/channel_select.go | 12 + .../channel_affinity_setting.go | 17 + setting/ratio_setting/group_ratio.go | 82 ++++- .../ratio_setting/group_ratio_empty_test.go | 84 +++++ setting/user_usable_group.go | 28 +- setting/user_usable_group_empty_test.go | 16 + 18 files changed, 1800 insertions(+), 73 deletions(-) create mode 100644 controller/channel_auto_test_helpers_test.go create mode 100644 service/channel_adaptive.go create mode 100644 service/channel_adaptive_test.go create mode 100644 service/channel_circuit.go create mode 100644 service/channel_metrics.go create mode 100644 service/channel_score.go create mode 100644 setting/ratio_setting/group_ratio_empty_test.go create mode 100644 setting/user_usable_group_empty_test.go diff --git a/controller/channel-test.go b/controller/channel-test.go index 4ba3698bd54c..1e99775d0fee 100644 --- a/controller/channel-test.go +++ b/controller/channel-test.go @@ -20,6 +20,7 @@ import ( "github.com/QuantumNous/new-api/middleware" "github.com/QuantumNous/new-api/model" "github.com/QuantumNous/new-api/pkg/billingexpr" + perfmetrics "github.com/QuantumNous/new-api/pkg/perf_metrics" "github.com/QuantumNous/new-api/relay" relaycommon "github.com/QuantumNous/new-api/relay/common" relayconstant "github.com/QuantumNous/new-api/relay/constant" @@ -39,6 +40,12 @@ type testResult struct { context *gin.Context localErr error newAPIError *types.NewAPIError + // Probe data for perf_metrics (populated on success). + relayInfo *relaycommon.RelayInfo + usage *dto.Usage + latencyMs int64 + outputTokens int64 + testModel string } func normalizeChannelTestEndpoint(channel *model.Channel, modelName, endpointType string) string { @@ -52,9 +59,156 @@ func normalizeChannelTestEndpoint(channel *model.Channel, modelName, endpointTyp if channel != nil && channel.Type == constant.ChannelTypeCodex { return string(constant.EndpointTypeOpenAIResponse) } + // Infer non-chat endpoints so auto-test does not force image/audio models through /chat/completions. + if kind := detectProbeModelKind(modelName); kind != "" { + return kind + } + if channel != nil && channel.Type == constant.ChannelTypeMokaAI { + return string(constant.EndpointTypeEmbeddings) + } + if channel != nil && channel.Type == constant.ChannelTypeVolcEngine && strings.Contains(strings.ToLower(modelName), "seedream") { + return string(constant.EndpointTypeImageGeneration) + } return normalized } +// detectProbeModelKind returns a constant.EndpointType string for known non-chat models. +// Empty string means default chat completions path. +func detectProbeModelKind(modelName string) string { + name := strings.ToLower(strings.TrimSpace(modelName)) + if name == "" { + return "" + } + if strings.HasSuffix(name, ratio_setting.CompactModelSuffix) { + return string(constant.EndpointTypeOpenAIResponseCompact) + } + if strings.Contains(name, "codex") { + return string(constant.EndpointTypeOpenAIResponse) + } + if strings.Contains(name, "rerank") { + return string(constant.EndpointTypeJinaRerank) + } + if strings.Contains(name, "embedding") || + strings.Contains(name, "embed") || + strings.HasPrefix(name, "m3e") || + strings.Contains(name, "bge-") || + strings.Contains(name, "text-embedding") { + return string(constant.EndpointTypeEmbeddings) + } + // Image generation / edit models — never chat-test these (#6121). + if isImageProbeModel(name) { + return string(constant.EndpointTypeImageGeneration) + } + return "" +} + +func isImageProbeModel(name string) bool { + name = strings.ToLower(strings.TrimSpace(name)) + if name == "" { + return false + } + imageHints := []string{ + "gpt-image", "dall-e", "dalle", "seedream", "flux", "imagen", + "stable-diffusion", "sdxl", "midjourney", "mj-", "image-gen", + "text-to-image", "t2i", "cogview", "kolors", "playground-v", + } + for _, h := range imageHints { + if strings.Contains(name, h) { + return true + } + } + // Bare suffixes like "*-image" / "image-*" when not embedding/vision chat. + if strings.Contains(name, "image") && + !strings.Contains(name, "vision") && + !strings.Contains(name, "chat") && + !strings.Contains(name, "embedding") { + return true + } + return false +} + +func isAudioOrVideoProbeModel(name string) bool { + name = strings.ToLower(strings.TrimSpace(name)) + if name == "" { + return false + } + hints := []string{ + "whisper", "tts-", "tts_", "-tts", "speech", "audio-", "-audio", + "sora", "kling", "runway", "luma", "hailuo", "vidu", "cogvideo", + "text-to-video", "t2v", "minimax-video", + } + for _, h := range hints { + if strings.Contains(name, h) { + return true + } + } + return false +} + +// isChatCapableProbeModel reports whether auto-test can safely use chat completions. +func isChatCapableProbeModel(name string) bool { + name = strings.TrimSpace(name) + if name == "" { + return false + } + kind := detectProbeModelKind(name) + if kind == string(constant.EndpointTypeImageGeneration) || + kind == string(constant.EndpointTypeEmbeddings) || + kind == string(constant.EndpointTypeJinaRerank) { + return false + } + if isAudioOrVideoProbeModel(name) { + return false + } + return true +} + +// pickAutoTestModel chooses a chat-capable model for batch auto-test. +// Prefers channel.TestModel when chat-capable; otherwise first chat-capable model +// in the channel list. Empty means skip auto probe for this channel. +func pickAutoTestModel(channel *model.Channel) string { + if channel == nil { + return "" + } + if channel.TestModel != nil { + if name := strings.TrimSpace(*channel.TestModel); name != "" && isChatCapableProbeModel(name) { + return name + } + } + for _, m := range channel.GetModels() { + if name := strings.TrimSpace(m); name != "" && isChatCapableProbeModel(name) { + return name + } + } + return "" +} + +func shouldSkipAutoChannelTest(channel *model.Channel) bool { + if channel == nil { + return true + } + if channel.Status == common.ChannelStatusManuallyDisabled { + return true + } + if channel.GetSetting().SkipAutoTest { + return true + } + // Channel types without chat/completion test support. + unsupported := []int{ + constant.ChannelTypeMidjourney, + constant.ChannelTypeMidjourneyPlus, + constant.ChannelTypeSunoAPI, + constant.ChannelTypeKling, + constant.ChannelTypeJimeng, + constant.ChannelTypeDoubaoVideo, + constant.ChannelTypeVidu, + } + if lo.Contains(unsupported, channel.Type) { + return true + } + return false +} + func resolveChannelTestUserID(c *gin.Context) (int, error) { if c != nil { if userID := c.GetInt("id"); userID > 0 { @@ -120,34 +274,29 @@ func testChannel(ctx context.Context, channel *model.Channel, testUserID int, te requestPath = endpointInfo.Path } } else { - // 如果没有指定端点类型，使用原有的自动检测逻辑 - - if strings.Contains(strings.ToLower(testModel), "rerank") { - requestPath = "/v1/rerank" - } - - // 先判断是否为 Embedding 模型 - if strings.Contains(strings.ToLower(testModel), "embedding") || - strings.HasPrefix(testModel, "m3e") || // m3e 系列模型 - strings.Contains(testModel, "bge-") || // bge 系列模型 - strings.Contains(testModel, "embed") || - channel.Type == constant.ChannelTypeMokaAI { // 其他 embedding 模型 - requestPath = "/v1/embeddings" // 修改请求路径 + // 如果没有指定端点类型，使用统一的模型种类检测 + kind := detectProbeModelKind(testModel) + if kind == "" && channel.Type == constant.ChannelTypeMokaAI { + kind = string(constant.EndpointTypeEmbeddings) } - - // VolcEngine 图像生成模型 - if channel.Type == constant.ChannelTypeVolcEngine && strings.Contains(testModel, "seedream") { - requestPath = "/v1/images/generations" - } - - // responses-only models - if strings.Contains(strings.ToLower(testModel), "codex") { - requestPath = "/v1/responses" + if kind == "" && channel.Type == constant.ChannelTypeVolcEngine && strings.Contains(strings.ToLower(testModel), "seedream") { + kind = string(constant.EndpointTypeImageGeneration) } - - // responses compaction models (must use /v1/responses/compact) - if strings.HasSuffix(testModel, ratio_setting.CompactModelSuffix) { - requestPath = "/v1/responses/compact" + if endpointInfo, ok := common.GetDefaultEndpointInfo(constant.EndpointType(kind)); ok { + requestPath = endpointInfo.Path + } else { + switch kind { + case string(constant.EndpointTypeJinaRerank): + requestPath = "/v1/rerank" + case string(constant.EndpointTypeEmbeddings): + requestPath = "/v1/embeddings" + case string(constant.EndpointTypeImageGeneration): + requestPath = "/v1/images/generations" + case string(constant.EndpointTypeOpenAIResponse): + requestPath = "/v1/responses" + case string(constant.EndpointTypeOpenAIResponseCompact): + requestPath = "/v1/responses/compact" + } } } if strings.HasPrefix(requestPath, "/v1/responses/compact") { @@ -512,9 +661,14 @@ func testChannel(ctx context.Context, channel *model.Channel, testUserID int, te }) common.SysLog(fmt.Sprintf("testing channel #%d, response: \n%s", channel.Id, string(respBody))) return testResult{ - context: c, - localErr: nil, - newAPIError: nil, + context: c, + localErr: nil, + newAPIError: nil, + relayInfo: info, + usage: usage, + latencyMs: milliseconds, + outputTokens: int64(usage.CompletionTokens), + testModel: info.OriginModelName, } } @@ -757,37 +911,33 @@ func buildTestRequest(model string, endpointType string, channel *model.Channel, } } - // 自动检测逻辑（保持原有行为） - if strings.Contains(strings.ToLower(model), "rerank") { + // 自动检测逻辑（与 detectProbeModelKind / normalizeChannelTestEndpoint 对齐） + switch detectProbeModelKind(model) { + case string(constant.EndpointTypeJinaRerank): return &dto.RerankRequest{ Model: model, Query: "What is Deep Learning?", Documents: []any{"Deep Learning is a subset of machine learning.", "Machine learning is a field of artificial intelligence."}, TopN: lo.ToPtr(2), } - } - - // 先判断是否为 Embedding 模型 - if strings.Contains(strings.ToLower(model), "embedding") || - strings.HasPrefix(model, "m3e") || - strings.Contains(model, "bge-") { - // 返回 EmbeddingRequest + case string(constant.EndpointTypeEmbeddings): return &dto.EmbeddingRequest{ Model: model, Input: []any{"hello world"}, } - } - - // Responses compaction models (must use /v1/responses/compact) - if strings.HasSuffix(model, ratio_setting.CompactModelSuffix) { + case string(constant.EndpointTypeImageGeneration): + return &dto.ImageRequest{ + Model: model, + Prompt: "a cute cat", + N: lo.ToPtr(uint(1)), + Size: "1024x1024", + } + case string(constant.EndpointTypeOpenAIResponseCompact): return &dto.OpenAIResponsesCompactionRequest{ Model: model, Input: testResponsesInput, } - } - - // Responses-only models (e.g. codex series) - if strings.Contains(strings.ToLower(model), "codex") { + case string(constant.EndpointTypeOpenAIResponse): return &dto.OpenAIResponsesRequest{ Model: model, Input: json.RawMessage(`[{"role":"user","content":"hi"}]`), @@ -825,6 +975,52 @@ func buildTestRequest(model string, endpointType string, channel *model.Channel, return testRequest } +func resolveProbeModelName(result testResult, channel *model.Channel, requested string) string { + if result.testModel != "" { + return result.testModel + } + if requested = strings.TrimSpace(requested); requested != "" { + return requested + } + if channel != nil && channel.TestModel != nil { + if name := strings.TrimSpace(*channel.TestModel); name != "" { + return name + } + } + if channel != nil { + models := channel.GetModels() + if len(models) > 0 { + if name := strings.TrimSpace(models[0]); name != "" { + return name + } + } + } + return "" +} + +func recordChannelProbeMetric(result testResult, channel *model.Channel, requested string, latencyMs int64) { + modelName := resolveProbeModelName(result, channel, requested) + if modelName == "" { + return + } + probeGroup := "probe" + if result.relayInfo != nil && result.relayInfo.UsingGroup != "" { + probeGroup = result.relayInfo.UsingGroup + } + generationMs := int64(0) + if result.outputTokens > 0 && latencyMs > 0 { + generationMs = latencyMs + } + perfmetrics.Record(perfmetrics.Sample{ + Model: modelName, + Group: probeGroup, + LatencyMs: latencyMs, + Success: result.localErr == nil && result.newAPIError == nil, + OutputTokens: result.outputTokens, + GenerationMs: generationMs, + }) +} + func TestChannel(c *gin.Context) { channelId, err := strconv.Atoi(c.Param("id")) if err != nil { @@ -858,11 +1054,15 @@ func TestChannel(c *gin.Context) { requestCtx = c.Request.Context() } result := testChannel(requestCtx, channel, testUserID, testModel, endpointType, isStream) + tok := time.Now() + milliseconds := tok.Sub(tik).Milliseconds() + // Always record success and failure probes so success_rate stays honest. + go recordChannelProbeMetric(result, channel, testModel, milliseconds) if result.localErr != nil { resp := gin.H{ "success": false, "message": result.localErr.Error(), - "time": 0.0, + "time": float64(milliseconds) / 1000.0, } if result.newAPIError != nil { resp["error_code"] = result.newAPIError.GetErrorCode() @@ -870,10 +1070,9 @@ func TestChannel(c *gin.Context) { c.JSON(http.StatusOK, resp) return } - tok := time.Now() - milliseconds := tok.Sub(tik).Milliseconds() go channel.UpdateResponseTime(milliseconds) consumedTime := float64(milliseconds) / 1000.0 + if result.newAPIError != nil { c.JSON(http.StatusOK, gin.H{ "success": false, @@ -922,9 +1121,18 @@ func performChannelTests(ctx context.Context, channels []*model.Channel, testUse if channel.Status == common.ChannelStatusManuallyDisabled { continue } + if shouldSkipAutoChannelTest(channel) { + continue + } + // Only auto-probe chat-capable models so image/audio/video do not pollute perf_metrics. + autoModel := pickAutoTestModel(channel) + if autoModel == "" { + common.SysLog(fmt.Sprintf("skip auto test channel %d (%s): no chat-capable test model", channel.Id, channel.Name)) + continue + } isChannelEnabled := channel.Status == common.ChannelStatusEnabled tik := time.Now() - result := testChannel(ctx, channel, testUserID, "", "", shouldUseStreamForAutomaticChannelTest(channel)) + result := testChannel(ctx, channel, testUserID, autoModel, "", shouldUseStreamForAutomaticChannelTest(channel)) tok := time.Now() milliseconds := tok.Sub(tik).Milliseconds() if ctx != nil && ctx.Err() != nil { @@ -968,6 +1176,8 @@ func performChannelTests(ctx context.Context, channels []*model.Channel, testUse } channel.UpdateResponseTime(milliseconds) + // Record success and failure probes for model-square health badges. + recordChannelProbeMetric(result, channel, autoModel, milliseconds) if common.RequestInterval > 0 { if ctx == nil { time.Sleep(common.RequestInterval) diff --git a/controller/channel.go b/controller/channel.go index a00011a9f9c1..1ba6b9cc0257 100644 --- a/controller/channel.go +++ b/controller/channel.go @@ -881,6 +881,12 @@ type ChannelBatch struct { Tag *string `json:"tag"` } +// ChannelSkipAutoTestBatch toggles skip_auto_test on channel settings. +type ChannelSkipAutoTestBatch struct { + Ids []int `json:"ids"` + Skip bool `json:"skip"` +} + func DeleteChannelBatch(c *gin.Context) { channelBatch := ChannelBatch{} err := c.ShouldBindJSON(&channelBatch) @@ -1329,6 +1335,48 @@ func BatchSetChannelTag(c *gin.Context) { return } +// BatchSetChannelSkipAutoTest sets skip_auto_test on selected channels. +// Manual channel tests remain available; only AutomaticallyTestChannels is gated. +func BatchSetChannelSkipAutoTest(c *gin.Context) { + req := ChannelSkipAutoTestBatch{} + if err := c.ShouldBindJSON(&req); err != nil || len(req.Ids) == 0 { + c.JSON(http.StatusOK, gin.H{ + "success": false, + "message": "invalid parameters", + }) + return + } + updated := 0 + for _, id := range req.Ids { + channel, err := model.GetChannelById(id, true) + if err != nil || channel == nil { + continue + } + setting := channel.GetSetting() + if setting.SkipAutoTest == req.Skip { + updated++ + continue + } + setting.SkipAutoTest = req.Skip + channel.SetSetting(setting) + if err := channel.Update(); err != nil { + common.SysLog(fmt.Sprintf("batch skip_auto_test update failed id=%d: %v", id, err)) + continue + } + updated++ + } + model.InitChannelCache() + recordManageAudit(c, "channel.skip_auto_test_batch", map[string]interface{}{ + "count": updated, + "skip": req.Skip, + }) + c.JSON(http.StatusOK, gin.H{ + "success": true, + "message": "", + "data": updated, + }) +} + func GetTagModels(c *gin.Context) { tag := c.Query("tag") if tag == "" { diff --git a/controller/channel_auto_test_helpers_test.go b/controller/channel_auto_test_helpers_test.go new file mode 100644 index 000000000000..c88e17055980 --- /dev/null +++ b/controller/channel_auto_test_helpers_test.go @@ -0,0 +1,93 @@ +package controller + +import ( + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/dto" + "github.com/QuantumNous/new-api/model" +) + +func TestDetectProbeModelKind(t *testing.T) { + cases := []struct { + model string + want string + }{ + {"gpt-image-2", string(constant.EndpointTypeImageGeneration)}, + {"dall-e-3", string(constant.EndpointTypeImageGeneration)}, + {"seedream-3.0", string(constant.EndpointTypeImageGeneration)}, + {"text-embedding-3-small", string(constant.EndpointTypeEmbeddings)}, + {"bge-m3", string(constant.EndpointTypeEmbeddings)}, + {"jina-rerank-v2", string(constant.EndpointTypeJinaRerank)}, + {"gpt-5-codex", string(constant.EndpointTypeOpenAIResponse)}, + {"gpt-4o-mini", ""}, + {"claude-sonnet-4", ""}, + } + for _, tc := range cases { + if got := detectProbeModelKind(tc.model); got != tc.want { + t.Fatalf("detectProbeModelKind(%q)=%q want %q", tc.model, got, tc.want) + } + } +} + +func TestIsChatCapableProbeModel(t *testing.T) { + if !isChatCapableProbeModel("gpt-4o-mini") { + t.Fatal("gpt-4o-mini should be chat capable") + } + for _, name := range []string{"gpt-image-2", "whisper-1", "text-embedding-3-large", "sora-2"} { + if isChatCapableProbeModel(name) { + t.Fatalf("%s should not be chat capable", name) + } + } +} + +func TestPickAutoTestModel(t *testing.T) { + imageOnly := "gpt-image-2" + chat := "gpt-4o-mini" + ch := &model.Channel{Models: "gpt-image-2,gpt-4o-mini"} + if got := pickAutoTestModel(ch); got != chat { + t.Fatalf("pickAutoTestModel=%q want %q", got, chat) + } + ch.TestModel = &imageOnly + if got := pickAutoTestModel(ch); got != chat { + t.Fatalf("with image TestModel pick=%q want %q", got, chat) + } + ch.TestModel = &chat + if got := pickAutoTestModel(ch); got != chat { + t.Fatalf("with chat TestModel pick=%q want %q", got, chat) + } + ch.Models = "gpt-image-2,dall-e-3" + ch.TestModel = &imageOnly + if got := pickAutoTestModel(ch); got != "" { + t.Fatalf("image-only channel pick=%q want empty", got) + } +} + +func TestShouldSkipAutoChannelTest(t *testing.T) { + ch := &model.Channel{Status: common.ChannelStatusEnabled} + if shouldSkipAutoChannelTest(ch) { + t.Fatal("enabled channel should not skip") + } + ch.Status = common.ChannelStatusManuallyDisabled + if !shouldSkipAutoChannelTest(ch) { + t.Fatal("manually disabled should skip") + } + ch.Status = common.ChannelStatusEnabled + ch.SetSetting(dto.ChannelSettings{SkipAutoTest: true}) + if !shouldSkipAutoChannelTest(ch) { + t.Fatal("SkipAutoTest setting should skip") + } +} + +func TestNormalizeChannelTestEndpointInfersImage(t *testing.T) { + if got := normalizeChannelTestEndpoint(nil, "gpt-image-2", ""); got != string(constant.EndpointTypeImageGeneration) { + t.Fatalf("image endpoint=%q", got) + } + if got := normalizeChannelTestEndpoint(nil, "gpt-4o-mini", ""); got != "" { + t.Fatalf("chat endpoint should be empty, got %q", got) + } + if got := normalizeChannelTestEndpoint(nil, "gpt-4o-mini", "openai"); got != "openai" { + t.Fatalf("explicit endpoint not preserved: %q", got) + } +} diff --git a/controller/group.go b/controller/group.go index 6ba339a3f9bd..10ae1f57f478 100644 --- a/controller/group.go +++ b/controller/group.go @@ -27,6 +27,13 @@ func GetUserGroups(c *gin.Context) { usableGroups := make(map[string]map[string]interface{}) userGroup := "" userId := c.GetInt("id") + if userId <= 0 { + c.JSON(http.StatusUnauthorized, gin.H{ + "success": false, + "message": "登录后查看可用分组", + }) + return + } userGroup, _ = model.GetUserGroup(userId, false) userUsableGroups := service.GetUserUsableGroups(userGroup) for groupName, _ := range ratio_setting.GetGroupRatioCopy() { diff --git a/dto/channel_settings.go b/dto/channel_settings.go index c92a3f988a3a..ffefcce107be 100644 --- a/dto/channel_settings.go +++ b/dto/channel_settings.go @@ -17,6 +17,9 @@ type ChannelSettings struct { PassThroughBodyEnabled bool `json:"pass_through_body_enabled,omitempty"` SystemPrompt string `json:"system_prompt,omitempty"` SystemPromptOverride bool `json:"system_prompt_override,omitempty"` + // SkipAutoTest excludes this channel from AutomaticallyTestChannels / testAllChannels. + // Manual single-channel tests still work. Upstream discussion: #5205. + SkipAutoTest bool `json:"skip_auto_test,omitempty"` } type VertexKeyType string diff --git a/model/channel_cache.go b/model/channel_cache.go index 81923017d79c..8a7d4eed190b 100644 --- a/model/channel_cache.go +++ b/model/channel_cache.go @@ -111,6 +111,52 @@ func SyncChannelCache(frequency int) { } } +// GetSatisfiedChannels returns all enabled channels for group+model (path-aware), +// highest priority first. Used by adaptive balance candidate collection. +// When memory cache is off, falls back to a single DB-selected channel. +func GetSatisfiedChannels(group string, modelName string, requestPath string) ([]*Channel, error) { + if !common.MemoryCacheEnabled { + ch, err := GetChannel(group, modelName, 0, requestPath) + if err != nil { + return nil, err + } + if ch == nil { + return nil, nil + } + return []*Channel{ch}, nil + } + + channelSyncLock.RLock() + defer channelSyncLock.RUnlock() + + ids := filterChannelsByRequestPath(group2model2channels[group][modelName], requestPath) + if len(ids) == 0 { + normalizedModel := ratio_setting.FormatMatchingModelName(modelName) + ids = filterChannelsByRequestPath(group2model2channels[group][normalizedModel], requestPath) + } + if len(ids) == 0 { + return nil, nil + } + + out := make([]*Channel, 0, len(ids)) + seen := make(map[int]struct{}, len(ids)) + for _, id := range ids { + if _, ok := seen[id]; ok { + continue + } + seen[id] = struct{}{} + ch, ok := channelsIDM[id] + if !ok || ch == nil { + continue + } + if ch.Status != common.ChannelStatusEnabled { + continue + } + out = append(out, ch) + } + return out, nil +} + func GetRandomSatisfiedChannel(group string, model string, retry int, requestPath string) (*Channel, error) { // if memory cache is disabled, get channel directly from database if !common.MemoryCacheEnabled { diff --git a/service/channel_adaptive.go b/service/channel_adaptive.go new file mode 100644 index 000000000000..dd1a70631389 --- /dev/null +++ b/service/channel_adaptive.go @@ -0,0 +1,249 @@ +package service + +import ( + "fmt" + "math/rand/v2" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/logger" + "github.com/QuantumNous/new-api/model" + "github.com/gin-gonic/gin" +) + +var ( + adaptiveLogEnabled bool + adaptiveLogSample = 0.01 // 采样率 +) + +// 请求上下文 key +type adaptiveContextKey string + +const ( + ctxKeyAdaptiveUsedChannels adaptiveContextKey = "adaptive_used_channels" + ctxKeyAdaptiveGroup adaptiveContextKey = "adaptive_group" + ctxKeyAdaptiveModel adaptiveContextKey = "adaptive_model" + ctxKeyAdaptiveSelected adaptiveContextKey = "adaptive_selected" + ctxKeyAdaptiveScores adaptiveContextKey = "adaptive_scores" +) + +// AdaptiveSelectChannel 动态评分调度器主入口。 +// 所有回退必须调用 cacheGetRandomSatisfiedChannelLegacy，禁止再进 CacheGetRandomSatisfiedChannel。 +func AdaptiveSelectChannel(param *RetryParam) (*model.Channel, string, error) { + ctx := param.Ctx + + // 未开启完整自适应：仅 legacy（含「只开 shadow」旧行为，避免递归） + if !constant.AdaptiveBalanceEnabled { + return cacheGetRandomSatisfiedChannelLegacy(param) + } + + // 提取 group 和 model + group := common.GetContextKeyString(ctx, constant.ContextKeyUsingGroup) + if group == "" { + group = param.TokenGroup + } + modelName := param.ModelName + + // 获取该 group+model 下的可用渠道 + channels, err := getCandidateChannels(group, modelName, param) + if err != nil { + return nil, group, err + } + if len(channels) == 0 { + return cacheGetRandomSatisfiedChannelLegacy(param) + } + + // 获取亲和偏好 channel + preferredID := getPreferredChannelID(ctx, modelName, group) + + // 评分 + candidates := ScoreCandidates(channels, group, modelName, preferredID) + + // 排除已用过的渠道（重试时）+ 熔断 + usedIDs := getAdaptiveUsedChannels(ctx) + var filtered []CandidateScore + for _, c := range candidates { + if containsInt(usedIDs, c.Channel.Id) { + continue + } + if IsCircuitOpen(c.Channel.Id) { + // 冷却后尝试 half-open 探测位 + if ProbeHalfOpen(c.Channel.Id) { + // allow one probe candidate through + } else { + continue + } + } + if c.Score <= 0 { + continue + } + filtered = append(filtered, c) + } + + if len(filtered) == 0 { + if constant.AdaptiveBalanceShadowMode { + logger.LogDebug(ctx, "adaptive: no available channels after filtering, fallback to original") + } + return cacheGetRandomSatisfiedChannelLegacy(param) + } + + // topK 加权随机选择 + selected := SelectTopKWeighted(filtered, 3) + if selected == nil { + return cacheGetRandomSatisfiedChannelLegacy(param) + } + + // Shadow Mode：选择仍走旧逻辑，仅记录对比 + if constant.AdaptiveBalanceShadowMode { + oldCh, oldGroup, oldErr := cacheGetRandomSatisfiedChannelLegacy(param) + + // 采样日志 + if adaptiveLogEnabled || randFloat64() < adaptiveLogSample { + logAdaptiveCompare(ctx, modelName, group, selected, oldCh) + } + + // shadow mode 下仍然使用旧渠道 + if oldCh != nil { + addAdaptiveUsedChannel(ctx, oldCh.Id) + storeAdaptiveSelection(ctx, selected.Channel, group, candidates) + return oldCh, oldGroup, oldErr + } + } + + // 正常模式：使用动态选择的渠道 + selectGroup := group + ch := selected.Channel + + addAdaptiveUsedChannel(ctx, ch.Id) + storeAdaptiveSelection(ctx, ch, group, candidates) + + logger.LogDebug(ctx, "adaptive selected channel #%d (score=%.3f) for group=%s model=%s", + ch.Id, selected.Score, group, modelName) + + return ch, selectGroup, nil +} + +// getCandidateChannels 获取 group+model 全部候选（非单渠道路由） +func getCandidateChannels(group, modelName string, param *RetryParam) ([]*model.Channel, error) { + // auto 分组：优先用上下文已解析的 auto group，否则 legacy 解析一次 + if group == "auto" || param.TokenGroup == "auto" { + if g := common.GetContextKeyString(param.Ctx, constant.ContextKeyAutoGroup); g != "" { + group = g + } else { + // 用 legacy 解析 auto → 具体 group，再拉全量候选 + ch, selectGroup, err := cacheGetRandomSatisfiedChannelLegacy(param) + if err != nil { + return nil, err + } + if ch == nil { + return nil, nil + } + if selectGroup != "" { + group = selectGroup + } + // 继续用解析后的 group 拉全量；若失败至少返回当前渠道 + list, listErr := model.GetSatisfiedChannels(group, modelName, param.RequestPath) + if listErr != nil { + return []*model.Channel{ch}, nil + } + if len(list) == 0 { + return []*model.Channel{ch}, nil + } + return list, nil + } + } + + return model.GetSatisfiedChannels(group, modelName, param.RequestPath) +} + +// getPreferredChannelID 读取亲和偏好（如果有） +func getPreferredChannelID(ctx *gin.Context, modelName, group string) int { + if !common.MemoryCacheEnabled { + return 0 + } + id, found := GetPreferredChannelByAffinity(ctx, modelName, group) + if found { + return id + } + return 0 +} + +// getAdaptiveUsedChannels 获取本次请求已用过的渠道 ID 列表 +func getAdaptiveUsedChannels(c *gin.Context) []int { + v, ok := c.Get(string(ctxKeyAdaptiveUsedChannels)) + if !ok { + return nil + } + ids, _ := v.([]int) + return ids +} + +// addAdaptiveUsedChannel 记录本次请求使用过的渠道 +func addAdaptiveUsedChannel(c *gin.Context, channelID int) { + existing := getAdaptiveUsedChannels(c) + existing = append(existing, channelID) + c.Set(string(ctxKeyAdaptiveUsedChannels), existing) +} + +// storeAdaptiveSelection 保存本次选择结果到上下文（供失败回写用） +func storeAdaptiveSelection(c *gin.Context, ch *model.Channel, group string, candidates []CandidateScore) { + c.Set(string(ctxKeyAdaptiveSelected), ch.Id) + c.Set(string(ctxKeyAdaptiveGroup), group) + if len(candidates) > 0 { + c.Set(string(ctxKeyAdaptiveScores), candidates) + } +} + +// logAdaptiveCompare shadow mode 日志 +func logAdaptiveCompare(c *gin.Context, modelName, group string, selected *CandidateScore, oldCh *model.Channel) { + oldID := 0 + if oldCh != nil { + oldID = oldCh.Id + } + logger.LogDebug(c, "[shadow] model=%s group=%s adaptive=#%d(%.3f) orig=#%d", + modelName, group, selected.Channel.Id, selected.Score, oldID) +} + +// RecordAdaptiveResult 请求完成后回调：更新指标 + 熔断状态 +func RecordAdaptiveResult(c *gin.Context, channelID int, group, modelName string, statusCode int, latency time.Duration, err error) { + if !constant.AdaptiveBalanceEnabled { + return + } + if channelID <= 0 { + return + } + + if err == nil && statusCode < 400 { + ObserveSuccess(channelID, group, modelName, latency) + RecordCircuitSuccess(channelID) + return + } + + // 失败处理 + ObserveFailure(channelID, group, modelName, statusCode, latency) + + if statusCode == 429 { + // 429 cooldown 由指标层自动处理 + logger.LogDebug(c, "adaptive: channel #%d got 429, score will be downgraded", channelID) + } + + if statusCode >= 500 || statusCode == 429 { + RecordCircuitFailure(channelID, fmt.Sprintf("HTTP %d", statusCode)) + } +} + +// containsInt 检查 int 切片是否包含某值 +func containsInt(slice []int, val int) bool { + for _, v := range slice { + if v == val { + return true + } + } + return false +} + +// randFloat64 生成 [0,1) 随机数 +var randFloat64 = func() float64 { + return rand.Float64() +} diff --git a/service/channel_adaptive_test.go b/service/channel_adaptive_test.go new file mode 100644 index 000000000000..726242940d17 --- /dev/null +++ b/service/channel_adaptive_test.go @@ -0,0 +1,204 @@ +package service + +import ( + "fmt" + "testing" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/model" + + "github.com/stretchr/testify/require" +) + +// 测试辅助：构造测试渠道 +func testChannel(id int, weight uint, priority int64) *model.Channel { + w := weight + return &model.Channel{ + Id: id, + Weight: &w, + Priority: &priority, + Name: fmt.Sprintf("ch-%d", id), + Status: common.ChannelStatusEnabled, + } +} + +func init() { + // 测试前重置全局状态 + globalSnapshot.mu.Lock() + globalSnapshot.metrics = make(map[metricsKey]*ChannelMetrics) + globalSnapshot.mu.Unlock() + + constant.EwmaAlpha = 0.3 + constant.MaxChannelConcurrency = 10 + constant.ChannelCircuitBreakerEnabled = true + constant.AdaptiveBalanceEnabled = true + constant.AdaptiveBalanceShadowMode = false +} + +// 测试1：高延迟渠道会被降权 +func TestHighLatencyDowngraded(t *testing.T) { + ch1 := testChannel(1, 10, 100) + ch2 := testChannel(2, 10, 100) + + // ch1 低延迟，ch2 高延迟 + ObserveSuccess(1, "test", "gpt-4", 200*time.Millisecond) + ObserveSuccess(1, "test", "gpt-4", 150*time.Millisecond) + ObserveSuccess(1, "test", "gpt-4", 180*time.Millisecond) + ObserveSuccess(2, "test", "gpt-4", 5*time.Second) + ObserveSuccess(2, "test", "gpt-4", 6*time.Second) + ObserveSuccess(2, "test", "gpt-4", 4*time.Second) + + channels := []*model.Channel{ch1, ch2} + candidates := ScoreCandidates(channels, "test", "gpt-4", 0) + + require.GreaterOrEqual(t, len(candidates), 2) + // ch1（低延迟）的分数应显著高于 ch2 + require.Equal(t, 1, candidates[0].Channel.Id, "expected ch1 (low latency) to rank first") + scoreDiff := candidates[0].Score - candidates[1].Score + require.Greater(t, scoreDiff, 0.1, "expected significant score difference") + t.Logf("ch1 (low latency) score=%.4f, ch2 (high latency) score=%.4f", candidates[0].Score, candidates[1].Score) +} + +// 测试2：429 渠道不会完全排除但会被降权 +func TestRateLimitedChannelDowngraded(t *testing.T) { + ch1 := testChannel(10, 10, 100) + ch2 := testChannel(11, 10, 100) + + // ch1 正常，ch2 有 429 + ObserveSuccess(10, "test", "gpt-4", 300*time.Millisecond) + ObserveSuccess(10, "test", "gpt-4", 250*time.Millisecond) + ObserveFailure(11, "test", "gpt-4", 429, 100*time.Millisecond) + ObserveFailure(11, "test", "gpt-4", 429, 100*time.Millisecond) + ObserveFailure(11, "test", "gpt-4", 429, 100*time.Millisecond) + + channels := []*model.Channel{ch1, ch2} + candidates := ScoreCandidates(channels, "test", "gpt-4", 0) + + require.GreaterOrEqual(t, len(candidates), 2) + // 正常渠道应排名更高 + require.Equal(t, 10, candidates[0].Channel.Id, "expected ch10 (normal) to rank first") + t.Logf("ch10 (normal) score=%.4f rate_limit_factor=%.4f", candidates[0].Score, candidates[0].RateLimitFactor) + t.Logf("ch11 (429) score=%.4f rate_limit_factor=%.4f", candidates[1].Score, candidates[1].RateLimitFactor) + require.Less(t, candidates[1].RateLimitFactor, candidates[0].RateLimitFactor) +} + +// 测试3：熔断渠道不会被选中 +func TestCircuitBreakerChannelExcluded(t *testing.T) { + ch := testChannel(20, 10, 100) + + // 模拟三次连续失败，触发熔断 + RecordCircuitFailure(20, "500 Internal Server Error") + RecordCircuitFailure(20, "500 Internal Server Error") + RecordCircuitFailure(20, "500 Internal Server Error") + + require.True(t, IsCircuitOpen(20), "expected circuit to be open after 3 consecutive failures") + + // 评分中应过滤掉熔断渠道 + channels := []*model.Channel{ch} + _ = ScoreCandidates(channels, "test", "gpt-4", 0) + + // channel_adaptive.go 中过滤逻辑会跳过 open 渠道 + require.True(t, IsCircuitOpen(20)) +} + +// 测试4：多渠道重试不会重复选择同一个渠道 +func TestNoDuplicateChannelInRetry(t *testing.T) { + ch1 := testChannel(30, 10, 100) + ch2 := testChannel(31, 10, 100) + ch3 := testChannel(32, 10, 100) + + // 全部成功 + for _, id := range []int{30, 31, 32} { + ObserveSuccess(id, "test", "gpt-4", 200*time.Millisecond) + } + + channels := []*model.Channel{ch1, ch2, ch3} + candidates := ScoreCandidates(channels, "test", "gpt-4", 0) + + // 模拟已使用的渠道 + usedIDs := []int{30} + + // 过滤掉已使用的渠道 + var filtered []CandidateScore + for _, c := range candidates { + if containsInt(usedIDs, c.Channel.Id) { + continue + } + filtered = append(filtered, c) + } + + require.Len(t, filtered, 2) + selected := SelectTopKWeighted(filtered, 3) + require.NotNil(t, selected) + require.NotEqual(t, 30, selected.Channel.Id) + t.Logf("selected ch%d (ch30 excluded)", selected.Channel.Id) +} + +// 测试5：TopK 加权随机不会全部集中在最高分渠道 +func TestTopKWeightedRandomFairness(t *testing.T) { + channels := make([]*model.Channel, 10) + for i := 0; i < 10; i++ { + channels[i] = testChannel(100+i, 10, 100) + ObserveSuccess(100+i, "test", "gpt-4", time.Duration(200+(i*100))*time.Millisecond) + } + + candidates := ScoreCandidates(channels, "test", "gpt-4", 0) + require.GreaterOrEqual(t, len(candidates), 10) + + // 模拟多次选择，统计分布 + selectionCount := make(map[int]int) + trials := 1000 + for i := 0; i < trials; i++ { + selected := SelectTopKWeighted(candidates, 3) + if selected != nil { + selectionCount[selected.Channel.Id]++ + } + } + + // TopK 的前三名应该都有一定比例 + for _, c := range candidates[:3] { + count := selectionCount[c.Channel.Id] + ratio := float64(count) / float64(trials) + t.Logf("ch%d selection rate: %.2f%% (score=%.3f)", c.Channel.Id, ratio*100, c.Score) + require.GreaterOrEqual(t, ratio, 0.05, "ch%d selected too few times", c.Channel.Id) + } +} + +// 测试6：EWMA 计算正确 +func TestEwmaUpdate(t *testing.T) { + alpha := 0.3 + + // 初始 1.0，观察到 0.5 + result := EwmaUpdate(1.0, 0.5, alpha) + expected := 0.3*0.5 + 0.7*1.0 // = 0.85 + require.InDelta(t, expected, result, 0.001) + + // 再次衰减 + result2 := EwmaUpdate(result, 0.5, alpha) + expected2 := 0.3*0.5 + 0.7*0.85 // = 0.745 + require.InDelta(t, expected2, result2, 0.001) +} + +// 测试7：延迟桶边界 +func TestLatencyBucket(t *testing.T) { + tests := []struct { + latency time.Duration + bucket int + }{ + {100 * time.Millisecond, 0}, + {500 * time.Millisecond, 0}, + {600 * time.Millisecond, 1}, + {1 * time.Second, 1}, + {1500 * time.Millisecond, 2}, + {3 * time.Second, 3}, + {7 * time.Second, 4}, + {15 * time.Second, 5}, + } + + for _, tt := range tests { + b := latencyBucket(tt.latency) + require.Equal(t, tt.bucket, b, "latencyBucket(%v)", tt.latency) + } +} diff --git a/service/channel_affinity_usage_cache_test.go b/service/channel_affinity_usage_cache_test.go index 64d3d715b547..2af84f6da178 100644 --- a/service/channel_affinity_usage_cache_test.go +++ b/service/channel_affinity_usage_cache_test.go @@ -4,7 +4,6 @@ import ( "fmt" "net/http/httptest" "testing" - "time" "github.com/QuantumNous/new-api/dto" "github.com/QuantumNous/new-api/types" @@ -12,7 +11,16 @@ import ( "github.com/stretchr/testify/require" ) -func buildChannelAffinityStatsContextForTest(ruleName, usingGroup, keyFP string) *gin.Context { +func buildChannelAffinityStatsContextForTest(t *testing.T) (*gin.Context, string, string, string) { + t.Helper() + ruleName := fmt.Sprintf("rule_%s", t.Name()) + usingGroup := "default" + keyFP := fmt.Sprintf("fp_%s", t.Name()) + entryKey := channelAffinityUsageCacheEntryKey(ruleName, usingGroup, keyFP) + t.Cleanup(func() { + _, _ = getChannelAffinityUsageCacheStatsCache().DeleteMany([]string{entryKey}) + }) + rec := httptest.NewRecorder() ctx, _ := gin.CreateTestContext(rec) setChannelAffinityContext(ctx, channelAffinityMeta{ @@ -22,14 +30,11 @@ func buildChannelAffinityStatsContextForTest(ruleName, usingGroup, keyFP string) UsingGroup: usingGroup, KeyFingerprint: keyFP, }) - return ctx + return ctx, ruleName, usingGroup, keyFP } func TestObserveChannelAffinityUsageCacheByRelayFormat_ClaudeMode(t *testing.T) { - ruleName := fmt.Sprintf("rule_%d", time.Now().UnixNano()) - usingGroup := "default" - keyFP := fmt.Sprintf("fp_%d", time.Now().UnixNano()) - ctx := buildChannelAffinityStatsContextForTest(ruleName, usingGroup, keyFP) + ctx, ruleName, usingGroup, keyFP := buildChannelAffinityStatsContextForTest(t) usage := &dto.Usage{ PromptTokens: 100, @@ -53,10 +58,7 @@ func TestObserveChannelAffinityUsageCacheByRelayFormat_ClaudeMode(t *testing.T) } func TestObserveChannelAffinityUsageCacheByRelayFormat_MixedMode(t *testing.T) { - ruleName := fmt.Sprintf("rule_%d", time.Now().UnixNano()) - usingGroup := "default" - keyFP := fmt.Sprintf("fp_%d", time.Now().UnixNano()) - ctx := buildChannelAffinityStatsContextForTest(ruleName, usingGroup, keyFP) + ctx, ruleName, usingGroup, keyFP := buildChannelAffinityStatsContextForTest(t) openAIUsage := &dto.Usage{ PromptTokens: 100, @@ -83,10 +85,7 @@ func TestObserveChannelAffinityUsageCacheByRelayFormat_MixedMode(t *testing.T) { } func TestObserveChannelAffinityUsageCacheByRelayFormat_UnsupportedModeKeepsEmpty(t *testing.T) { - ruleName := fmt.Sprintf("rule_%d", time.Now().UnixNano()) - usingGroup := "default" - keyFP := fmt.Sprintf("fp_%d", time.Now().UnixNano()) - ctx := buildChannelAffinityStatsContextForTest(ruleName, usingGroup, keyFP) + ctx, ruleName, usingGroup, keyFP := buildChannelAffinityStatsContextForTest(t) usage := &dto.Usage{ PromptTokens: 100, diff --git a/service/channel_circuit.go b/service/channel_circuit.go new file mode 100644 index 000000000000..eeeff5e4d885 --- /dev/null +++ b/service/channel_circuit.go @@ -0,0 +1,176 @@ +package service + +import ( + "github.com/QuantumNous/new-api/constant" + "sync" + "time" +) + +// CircuitState 熔断器状态 +type CircuitState string + +const ( + CircuitClosed CircuitState = "closed" // 正常 + CircuitOpen CircuitState = "open" // 熔断打开，不选 + CircuitHalfOpen CircuitState = "half_open" // 半开，允许探测 +) + +// ChannelCircuitBreaker 渠道熔断器（本地状态 + Redis 同步） +// 约定：请求路径上只读本地状态，不访问 Redis。 +type ChannelCircuitBreaker struct { + mu sync.RWMutex + + State CircuitState + ConsecutiveFailure int // 连续失败计数 + OpenUntil time.Time // open 状态过期时间 + HalfOpenLimit int // half-open 最大探测数 + HalfOpenInFlight int // half-open 进行中的探测数 + HalfOpenSince time.Time // when current half-open probe started + LastError string // 最近一次错误信息 +} + +var ( + circuitBreakers sync.Map // map[int]*ChannelCircuitBreaker, key=channelID +) + +// getCircuitBreaker 获取或创建渠道熔断器 +func getCircuitBreaker(channelID int) *ChannelCircuitBreaker { + v, _ := circuitBreakers.LoadOrStore(channelID, &ChannelCircuitBreaker{ + State: CircuitClosed, + HalfOpenLimit: 1, + }) + return v.(*ChannelCircuitBreaker) +} + +// IsCircuitOpen 判断渠道是否熔断（请求路径使用，读本地状态） +func IsCircuitOpen(channelID int) bool { + if !constant.ChannelCircuitBreakerEnabled { + return false + } + + cb := getCircuitBreaker(channelID) + cb.mu.RLock() + defer cb.mu.RUnlock() + + if cb.State == CircuitClosed { + return false + } + + if cb.State == CircuitOpen && time.Now().After(cb.OpenUntil) { + // Cooldown elapsed: still report open so selector must call ProbeHalfOpen. + return true + } + + return true +} + +// IsInCooldown 判断渠道是否在 429 cooldown 中 +func IsInCooldown(channelID int, cooldownUntil time.Time) bool { + if cooldownUntil.IsZero() { + return false + } + return time.Now().Before(cooldownUntil) +} + +// RecordSuccess 成功调用 -> 重置熔断状态 +func RecordCircuitSuccess(channelID int) { + if !constant.ChannelCircuitBreakerEnabled { + return + } + + cb := getCircuitBreaker(channelID) + cb.mu.Lock() + defer cb.mu.Unlock() + + if cb.State == CircuitHalfOpen { + cb.HalfOpenInFlight-- + if cb.HalfOpenInFlight < 0 { + cb.HalfOpenInFlight = 0 + } + cb.HalfOpenSince = time.Time{} + } + + cb.State = CircuitClosed + cb.ConsecutiveFailure = 0 + cb.OpenUntil = time.Time{} + cb.LastError = "" +} + +// RecordFailure 失败调用 -> 可能触发熔断 +func RecordCircuitFailure(channelID int, errMsg string) { + if !constant.ChannelCircuitBreakerEnabled { + return + } + + cb := getCircuitBreaker(channelID) + cb.mu.Lock() + defer cb.mu.Unlock() + + cb.ConsecutiveFailure++ + cb.LastError = errMsg + + // 半开状态下失败 -> 立即回退到 open + if cb.State == CircuitHalfOpen { + cb.HalfOpenInFlight-- + if cb.HalfOpenInFlight < 0 { + cb.HalfOpenInFlight = 0 + } + cb.HalfOpenSince = time.Time{} + cb.State = CircuitOpen + cb.OpenUntil = time.Now().Add(time.Duration(constant.ChannelCooldownSeconds) * time.Second) + return + } + + // closed 状态下连续失败达到阈值 -> open + threshold := 3 + if cb.ConsecutiveFailure >= threshold { + cb.State = CircuitOpen + cb.OpenUntil = time.Now().Add(time.Duration(constant.ChannelCooldownSeconds) * time.Second) + } + + // 如果配置了熔断但未启用，不做任何事 +} + +// ProbeHalfOpen 申请进入 half-open 探测（由选择器调用） +func ProbeHalfOpen(channelID int) bool { + if !constant.ChannelCircuitBreakerEnabled { + return true + } + + cb := getCircuitBreaker(channelID) + cb.mu.Lock() + defer cb.mu.Unlock() + + // open 且已过冷却 -> 自动转为 half-open + if cb.State == CircuitOpen && time.Now().After(cb.OpenUntil) { + cb.State = CircuitHalfOpen + cb.HalfOpenInFlight = 0 + } + + if cb.State != CircuitHalfOpen { + return false + } + + // Reclaim stuck half-open probes (client cancel / no Record* callback). + const halfOpenProbeTimeout = 60 * time.Second + if cb.HalfOpenInFlight > 0 && !cb.HalfOpenSince.IsZero() && time.Since(cb.HalfOpenSince) > halfOpenProbeTimeout { + cb.HalfOpenInFlight = 0 + cb.HalfOpenSince = time.Time{} + } + + if cb.HalfOpenInFlight >= cb.HalfOpenLimit { + return false + } + + cb.HalfOpenInFlight++ + cb.HalfOpenSince = time.Now() + return true +} + +// GetCircuitState 读取熔断状态（供日志/观测使用） +func GetCircuitState(channelID int) (CircuitState, int, string) { + cb := getCircuitBreaker(channelID) + cb.mu.RLock() + defer cb.mu.RUnlock() + return cb.State, cb.ConsecutiveFailure, cb.LastError +} \ No newline at end of file diff --git a/service/channel_metrics.go b/service/channel_metrics.go new file mode 100644 index 000000000000..6d9db5961b81 --- /dev/null +++ b/service/channel_metrics.go @@ -0,0 +1,292 @@ +package service + +import ( + "fmt" + "math" + "sync" + "sync/atomic" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" +) + +// ChannelMetrics 渠道运行时指标，按 (channelID, group, model) 分桶 +type ChannelMetrics struct { + mu sync.Mutex `json:"-"` // 保护并发写 + SuccessRate float64 `json:"success_rate"` // EWMA + ErrorRate float64 `json:"error_rate"` // EWMA + RateLimitRate float64 `json:"rate_limit_rate"` // EWMA 429 率 + Status5xxRate float64 `json:"status_5xx_rate"` // EWMA 5xx 率 + AvgLatency time.Duration `json:"avg_latency"` // EWMA 平均延迟 + SampleCount int64 `json:"sample_count"` // 总样本数 + LastSeen time.Time `json:"last_seen"` + + // 延迟直方图（轻量桶），用于近似 p95 + LatencyBuckets [6]int64 `json:"latency_buckets"` // <=500ms, <=1s, <=2s, <=5s, <=10s, >10s +} + +// LocalMetricsSnapshot 进程内本地指标快照，定期从 Redis sync 或直接从本地累加 +type LocalMetricsSnapshot struct { + mu sync.RWMutex + metrics map[metricsKey]*ChannelMetrics + updatedAt time.Time +} + +type metricsKey struct { + ChannelID int + Group string + Model string +} + +var globalSnapshot = &LocalMetricsSnapshot{ + metrics: make(map[metricsKey]*ChannelMetrics), +} + +// ensureKey 获取或创建指定 key 的指标桶 +func (s *LocalMetricsSnapshot) ensureKey(key metricsKey) *ChannelMetrics { + s.mu.Lock() + defer s.mu.Unlock() + m, ok := s.metrics[key] + if !ok { + m = &ChannelMetrics{ + SuccessRate: 1.0, // 冷启动默认信任 + AvgLatency: 500 * time.Millisecond, + } + s.metrics[key] = m + } + return m +} + +// EwmaUpdate 更新指标的 EWMA 值 +func EwmaUpdate(current, observed, alpha float64) float64 { + if alpha <= 0 || alpha > 1 { + alpha = constant.EwmaAlpha + } + return alpha*observed + (1-alpha)*current +} + +// ObserveSuccess 记录一次成功调用 +func ObserveSuccess(channelID int, group, model string, latency time.Duration) { + alpha := constant.EwmaAlpha + key := metricsKey{channelID, group, model} + m := globalSnapshot.ensureKey(key) + + m.mu.Lock() + defer m.mu.Unlock() + + m.SampleCount++ + m.LastSeen = time.Now() + + // 更新延迟 EWMA + if m.AvgLatency == 0 { + m.AvgLatency = latency + } else { + m.AvgLatency = time.Duration(EwmaUpdate(float64(m.AvgLatency), float64(latency), alpha)) + } + + // 更新延迟桶 + bucket := latencyBucket(latency) + if bucket >= 0 && bucket < len(m.LatencyBuckets) { + m.LatencyBuckets[bucket]++ + } + + // 更新成功率 EWMA + m.SuccessRate = EwmaUpdate(m.SuccessRate, 1.0, alpha) + m.ErrorRate = EwmaUpdate(m.ErrorRate, 0, alpha) + m.RateLimitRate = EwmaUpdate(m.RateLimitRate, 0, alpha) + m.Status5xxRate = EwmaUpdate(m.Status5xxRate, 0, alpha) +} + +// ObserveFailure 记录一次失败 +func ObserveFailure(channelID int, group, model string, statusCode int, latency time.Duration) { + alpha := constant.EwmaAlpha + key := metricsKey{channelID, group, model} + m := globalSnapshot.ensureKey(key) + + m.mu.Lock() + defer m.mu.Unlock() + + m.SampleCount++ + m.LastSeen = time.Now() + + // 更新延迟 + if m.AvgLatency == 0 { + m.AvgLatency = latency + } else { + m.AvgLatency = time.Duration(EwmaUpdate(float64(m.AvgLatency), float64(latency), alpha)) + } + + m.SuccessRate = EwmaUpdate(m.SuccessRate, 0, alpha) + m.ErrorRate = EwmaUpdate(m.ErrorRate, 1, alpha) + + switch { + case statusCode == 429: + m.RateLimitRate = EwmaUpdate(m.RateLimitRate, 1, alpha) + m.Status5xxRate = EwmaUpdate(m.Status5xxRate, 0, alpha) + case statusCode >= 500: + m.Status5xxRate = EwmaUpdate(m.Status5xxRate, 1, alpha) + m.RateLimitRate = EwmaUpdate(m.RateLimitRate, 0, alpha) + default: + m.RateLimitRate = EwmaUpdate(m.RateLimitRate, 0, alpha) + m.Status5xxRate = EwmaUpdate(m.Status5xxRate, 0, alpha) + } +} + +// GetMetrics 读取指定 (channelID, group, model) 的指标快照 +func GetMetrics(channelID int, group, model string) *ChannelMetrics { + key := metricsKey{channelID, group, model} + globalSnapshot.mu.RLock() + m, ok := globalSnapshot.metrics[key] + globalSnapshot.mu.RUnlock() + if ok { + return m + } + + // 尝试回退到 (channelID, group) + key2 := metricsKey{channelID, group, ""} + globalSnapshot.mu.RLock() + m2, ok2 := globalSnapshot.metrics[key2] + globalSnapshot.mu.RUnlock() + if ok2 { + return m2 + } + + // 回退到 (channelID) + key3 := metricsKey{channelID, "", ""} + globalSnapshot.mu.RLock() + m3, ok3 := globalSnapshot.metrics[key3] + globalSnapshot.mu.RUnlock() + if ok3 { + return m3 + } + + // 无数据，返回中性默认值 + return &ChannelMetrics{ + SuccessRate: 1.0, + AvgLatency: 500 * time.Millisecond, + } +} + +// GetP95Latency 从桶近似计算 p95 延迟 +func GetP95Latency(buckets [6]int64) time.Duration { + var total int64 + for _, v := range buckets { + total += v + } + if total == 0 { + return 0 + } + + target := int64(math.Ceil(float64(total) * 0.95)) + var cumulative int64 + + bucketBoundaries := []time.Duration{ + 500 * time.Millisecond, + 1 * time.Second, + 2 * time.Second, + 5 * time.Second, + 10 * time.Second, + math.MaxInt64, + } + + for i, count := range buckets { + cumulative += count + if cumulative >= target { + return bucketBoundaries[i] + } + } + return 10 * time.Second +} + +// CurrentConcurrencyTracker 本地并发计数器（原子操作，零网络开销） +type CurrentConcurrencyTracker struct { + counters sync.Map // map[int]*atomic.Int64 +} + +var globalConcurrency = &CurrentConcurrencyTracker{} + +func (t *CurrentConcurrencyTracker) Inc(channelID int) int64 { + v, _ := t.counters.LoadOrStore(channelID, new(atomic.Int64)) + return v.(*atomic.Int64).Add(1) +} + +func (t *CurrentConcurrencyTracker) Dec(channelID int) int64 { + v, ok := t.counters.Load(channelID) + if !ok { + return 0 + } + return v.(*atomic.Int64).Add(-1) +} + +func (t *CurrentConcurrencyTracker) Get(channelID int) int64 { + v, ok := t.counters.Load(channelID) + if !ok { + return 0 + } + return v.(*atomic.Int64).Load() +} + +// IncChannelConcurrency 增加并发计数 +func IncChannelConcurrency(channelID int) int64 { + return globalConcurrency.Inc(channelID) +} + +// DecChannelConcurrency 减少并发计数 +func DecChannelConcurrency(channelID int) int64 { + return globalConcurrency.Dec(channelID) +} + +// GetChannelConcurrency 获取当前并发 +func GetChannelConcurrency(channelID int) int64 { + return globalConcurrency.Get(channelID) +} + +func latencyBucket(latency time.Duration) int { + switch { + case latency <= 500*time.Millisecond: + return 0 + case latency <= 1*time.Second: + return 1 + case latency <= 2*time.Second: + return 2 + case latency <= 5*time.Second: + return 3 + case latency <= 10*time.Second: + return 4 + default: + return 5 + } +} + +// SyncAdaptiveMetricsToRedis publishes a compact snapshot for multi-instance +// sticky-or-shared observation. Best-effort; failures are silent. +func SyncAdaptiveMetricsToRedis() { + if !common.RedisEnabled || common.RDB == nil { + return + } + globalSnapshot.mu.RLock() + defer globalSnapshot.mu.RUnlock() + type row struct { + ChannelID int `json:"c"` + Group string `json:"g"` + Model string `json:"m"` + SuccessRate float64 `json:"s"` + SampleCount int64 `json:"n"` + } + out := make([]row, 0, len(globalSnapshot.metrics)) + for k, m := range globalSnapshot.metrics { + m.mu.Lock() + out = append(out, row{k.ChannelID, k.Group, k.Model, m.SuccessRate, m.SampleCount}) + m.mu.Unlock() + } + b, err := common.Marshal(out) + if err != nil { + return + } + key := fmt.Sprintf("newapi:adaptive:metrics:%s", common.NodeName) + if key == "newapi:adaptive:metrics:" { + key = "newapi:adaptive:metrics:default" + } + _ = common.RedisSet(key, string(b), 2*time.Minute) +} diff --git a/service/channel_score.go b/service/channel_score.go new file mode 100644 index 000000000000..67f9f025c244 --- /dev/null +++ b/service/channel_score.go @@ -0,0 +1,177 @@ +package service + +import ( + "github.com/QuantumNous/new-api/constant" + "math" + "math/rand" + "sort" + "time" + + "github.com/QuantumNous/new-api/model" +) + +// CandidateScore 评分候选结果 +type CandidateScore struct { + Channel *model.Channel + Score float64 + + // 各因子明细（供日志/观测用） + BaseWeight float64 `json:"base_weight"` + SuccessFactor float64 `json:"success_factor"` + LatencyFactor float64 `json:"latency_factor"` + RateLimitFactor float64 `json:"rate_limit_factor"` + ConcurrencyFactor float64 `json:"concurrency_factor"` + CircuitFactor float64 `json:"circuit_factor"` + AffinityFactor float64 `json:"affinity_factor"` +} + +// ScoreCandidates 对一组渠道进行动态评分，返回排序后的候选列表 +// 不传 group, model 时会回退到 channel 级别指标 +func ScoreCandidates(channels []*model.Channel, group, model string, preferredChannelID int) []CandidateScore { + if len(channels) == 0 { + return nil + } + + candidates := make([]CandidateScore, 0, len(channels)) + + for _, ch := range channels { + metrics := GetMetrics(ch.Id, group, model) + baseWeight := float64(ch.GetWeight()) + if baseWeight <= 0 { + baseWeight = 1.0 + } + + // 1. 成功率因子 [0, 1]：success_rate^2，让低成功率显著降权 + successFactor := math.Pow(metrics.SuccessRate, 2) + + // 2. 延迟因子 [0, 1]：基于 EWMA 平均延迟 + latencyMs := float64(metrics.AvgLatency) / float64(time.Millisecond) + latencyFactor := latencyToScore(latencyMs) + + // 3. 限流因子 [0, 1]：429 率越低越好 + rateLimitFactor := 1.0 - metrics.RateLimitRate + if rateLimitFactor < 0 { + rateLimitFactor = 0 + } + + // 4. 并发因子 [0, 1]：当前并发 / 最大并发 + currentConcurrency := GetChannelConcurrency(ch.Id) + maxConcurrency := int64(constant.MaxChannelConcurrency) + concurrencyFactor := 1.0 + if maxConcurrency > 0 && currentConcurrency >= maxConcurrency { + concurrencyFactor = 0.1 // 超限降权但不完全排除 + } else if maxConcurrency > 0 { + concurrencyFactor = 1.0 - float64(currentConcurrency)/float64(maxConcurrency)*0.5 + } + + // 5. 熔断因子 [0, 1] + circuitFactor := 1.0 + if constant.ChannelCircuitBreakerEnabled { + if IsCircuitOpen(ch.Id) { + circuitFactor = 0.0 + } + state, _, _ := GetCircuitState(ch.Id) + if state == CircuitHalfOpen { + circuitFactor = 0.5 + } + } + + // 6. 亲和因子 [1.0, 1.5] + affinityFactor := 1.0 + if preferredChannelID > 0 && ch.Id == preferredChannelID { + affinityFactor = 1.5 + } + + score := baseWeight * successFactor * latencyFactor * rateLimitFactor * + concurrencyFactor * circuitFactor * affinityFactor + + // 注入 jitter (±5%)，防止同分渠道集中 + jitter := 0.95 + rand.Float64()*0.1 + score *= jitter + + candidates = append(candidates, CandidateScore{ + Channel: ch, + Score: score, + BaseWeight: baseWeight, + SuccessFactor: successFactor, + LatencyFactor: latencyFactor, + RateLimitFactor: rateLimitFactor, + ConcurrencyFactor: concurrencyFactor, + CircuitFactor: circuitFactor, + AffinityFactor: affinityFactor, + }) + } + + // 按分数降序排序 + sort.Slice(candidates, func(i, j int) bool { + return candidates[i].Score > candidates[j].Score + }) + + return candidates +} + +// SelectTopKWeighted 从候选列表中取 topK 然后按 score 加权随机选一个 +func SelectTopKWeighted(candidates []CandidateScore, k int) *CandidateScore { + if len(candidates) == 0 { + return nil + } + + if len(candidates) == 1 { + return &candidates[0] + } + + // 取 topK + if k <= 0 { + k = 3 + } + if k > len(candidates) { + k = len(candidates) + } + top := candidates[:k] + + // 加权随机 + var totalWeight float64 + for _, c := range top { + if c.Score > 0 { + totalWeight += c.Score + } + } + + if totalWeight <= 0 { + // 所有分数为 0，均匀随机 + idx := rand.Intn(len(top)) + return &top[idx] + } + + r := rand.Float64() * totalWeight + var cumulative float64 + for i, c := range top { + cumulative += c.Score + if r < cumulative { + return &top[i] + } + } + + return &top[len(top)-1] +} + +// latencyToScore 将延迟（毫秒）映射到 [0, 1] 分数 +// 500ms → 1.0, 1s → 0.8, 2s → 0.5, 5s → 0.2, 10s+ → 0.05 +func latencyToScore(ms float64) float64 { + switch { + case ms <= 0: + return 1.0 + case ms <= 500: + return 1.0 + case ms <= 1000: + return 0.8 + case ms <= 2000: + return 0.5 + case ms <= 5000: + return 0.2 + case ms <= 10000: + return 0.1 + default: + return 0.05 + } +} \ No newline at end of file diff --git a/service/channel_select.go b/service/channel_select.go index 24c4e252bfb3..74a829363366 100644 --- a/service/channel_select.go +++ b/service/channel_select.go @@ -82,6 +82,18 @@ func (p *RetryParam) ResetRetryNextTry() { // Retry=3: GroupB, priority1 (startRetryIndex=2, priorityRetry=1) // 分组B, 优先级1 func CacheGetRandomSatisfiedChannel(param *RetryParam) (*model.Channel, string, error) { + // Adaptive entry. AdaptiveSelectChannel must only call + // cacheGetRandomSatisfiedChannelLegacy — never this function — or flags + // cause infinite recursion / stack overflow. + if constant.AdaptiveBalanceEnabled || constant.AdaptiveBalanceShadowMode { + return AdaptiveSelectChannel(param) + } + return cacheGetRandomSatisfiedChannelLegacy(param) +} + +// cacheGetRandomSatisfiedChannelLegacy is the original random / auto-group picker. +// Safe to call from adaptive fallbacks and candidate collection. +func cacheGetRandomSatisfiedChannelLegacy(param *RetryParam) (*model.Channel, string, error) { var channel *model.Channel var err error selectGroup := param.TokenGroup diff --git a/setting/operation_setting/channel_affinity_setting.go b/setting/operation_setting/channel_affinity_setting.go index a925a847d6dd..fef0c7f22d6a 100644 --- a/setting/operation_setting/channel_affinity_setting.go +++ b/setting/operation_setting/channel_affinity_setting.go @@ -116,6 +116,23 @@ var channelAffinitySetting = ChannelAffinitySetting{ MaxEntries: 100_000, DefaultTTLSeconds: 3600, Rules: []ChannelAffinityRule{ + { + Name: "axonhub trace sticky", + ModelRegex: []string{".*"}, + PathRegex: []string{"/v1/.*"}, + // Only client-provided traces (middleware sets affinity_trace_id). + KeySources: []ChannelAffinityKeySource{ + {Type: "context_string", Key: "affinity_trace_id"}, + {Type: "request_header", Key: "AH-Trace-Id"}, + {Type: "request_header", Key: "X-Trace-Id"}, + }, + ValueRegex: "", + TTLSeconds: 1800, + SkipRetryOnFailure: false, + IncludeUsingGroup: true, + IncludeModelName: true, + IncludeRuleName: true, + }, { Name: "codex cli trace", ModelRegex: []string{"^gpt-.*$"}, diff --git a/setting/ratio_setting/group_ratio.go b/setting/ratio_setting/group_ratio.go index 7d16d9283932..5e06542e4119 100644 --- a/setting/ratio_setting/group_ratio.go +++ b/setting/ratio_setting/group_ratio.go @@ -1,8 +1,8 @@ package ratio_setting import ( - "encoding/json" "errors" + "strings" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/setting/config" @@ -73,7 +73,37 @@ func GroupRatio2JSONString() string { } func UpdateGroupRatioByJSONString(jsonStr string) error { - return types.LoadFromJsonString(groupRatioMap, jsonStr) + // Empty object would wipe defaults and break pricing + perf-metrics group + // filters (summary only returns groups present in this map). Keep defaults. + trimmed := strings.TrimSpace(jsonStr) + if trimmed == "" || trimmed == "{}" || trimmed == "null" { + groupRatioMap.Clear() + groupRatioMap.AddAll(defaultGroupRatio) + return nil + } + tmp := make(map[string]float64) + if err := common.Unmarshal([]byte(trimmed), &tmp); err != nil { + return err + } + if len(tmp) == 0 { + groupRatioMap.Clear() + groupRatioMap.AddAll(defaultGroupRatio) + return nil + } + // Reject negative ratios before load (same rule as CheckGroupRatio). + for name, ratio := range tmp { + if ratio < 0 { + return errors.New("group ratio must be not less than 0: " + name) + } + } + if err := types.LoadFromJsonString(groupRatioMap, trimmed); err != nil { + return err + } + // Always keep a usable default group so pricing/perf never filter to empty. + if _, ok := groupRatioMap.Get("default"); !ok { + groupRatioMap.Set("default", 1) + } + return nil } func GetGroupRatio(name string) float64 { @@ -102,12 +132,36 @@ func GroupGroupRatio2JSONString() string { } func UpdateGroupGroupRatioByJSONString(jsonStr string) error { - return types.LoadFromJsonString(groupGroupRatioMap, jsonStr) + // Empty / null should not wipe nested group-group overrides to a broken state. + // Unlike GroupRatio, empty here means "no nested overrides" (valid), but still + // reject null-ish wipe of malformed payloads and negative ratios. + trimmed := strings.TrimSpace(jsonStr) + if trimmed == "" || trimmed == "null" { + groupGroupRatioMap.Clear() + return nil + } + tmp := make(map[string]map[string]float64) + if err := common.Unmarshal([]byte(trimmed), &tmp); err != nil { + return err + } + for userGroup, nested := range tmp { + for usingGroup, ratio := range nested { + if ratio < 0 { + return errors.New("group_group_ratio must be not less than 0: " + userGroup + " -> " + usingGroup) + } + } + } + return types.LoadFromJsonString(groupGroupRatioMap, trimmed) } func CheckGroupRatio(jsonStr string) error { + trimmed := strings.TrimSpace(jsonStr) + if trimmed == "" || trimmed == "{}" || trimmed == "null" { + // Empty is accepted; UpdateGroupRatioByJSONString restores defaults. + return nil + } checkGroupRatio := make(map[string]float64) - err := json.Unmarshal([]byte(jsonStr), &checkGroupRatio) + err := common.Unmarshal([]byte(trimmed), &checkGroupRatio) if err != nil { return err } @@ -118,3 +172,23 @@ func CheckGroupRatio(jsonStr string) error { } return nil } + +// CheckGroupGroupRatio validates nested user→using group ratio maps. +func CheckGroupGroupRatio(jsonStr string) error { + trimmed := strings.TrimSpace(jsonStr) + if trimmed == "" || trimmed == "{}" || trimmed == "null" { + return nil + } + check := make(map[string]map[string]float64) + if err := common.Unmarshal([]byte(trimmed), &check); err != nil { + return err + } + for userGroup, nested := range check { + for usingGroup, ratio := range nested { + if ratio < 0 { + return errors.New("group_group_ratio must be not less than 0: " + userGroup + " -> " + usingGroup) + } + } + } + return nil +} diff --git a/setting/ratio_setting/group_ratio_empty_test.go b/setting/ratio_setting/group_ratio_empty_test.go new file mode 100644 index 000000000000..578b544f28d0 --- /dev/null +++ b/setting/ratio_setting/group_ratio_empty_test.go @@ -0,0 +1,84 @@ +package ratio_setting + +import ( + "testing" +) + +func TestUpdateGroupRatioByJSONStringKeepsDefaultsOnEmpty(t *testing.T) { + // Ensure defaults are present first. + if err := UpdateGroupRatioByJSONString(`{"default":1,"vip":1,"svip":1}`); err != nil { + t.Fatalf("seed defaults: %v", err) + } + if err := UpdateGroupRatioByJSONString(`{}`); err != nil { + t.Fatalf("empty update: %v", err) + } + got := GetGroupRatioCopy() + if len(got) == 0 { + t.Fatalf("expected defaults after empty update, got empty map") + } + if _, ok := got["default"]; !ok { + t.Fatalf("expected default group after empty update, got %#v", got) + } +} + +func TestUpdateGroupRatioByJSONStringAcceptsCustom(t *testing.T) { + if err := UpdateGroupRatioByJSONString(`{"default":1,"pro":1.2}`); err != nil { + t.Fatalf("custom update: %v", err) + } + got := GetGroupRatioCopy() + if got["pro"] != 1.2 { + t.Fatalf("expected pro=1.2, got %#v", got) + } + if _, ok := got["default"]; !ok { + t.Fatalf("expected default retained/set, got %#v", got) + } +} + +func TestUpdateGroupRatioInjectsDefaultWhenMissing(t *testing.T) { + if err := UpdateGroupRatioByJSONString(`{"pro":1.5}`); err != nil { + t.Fatalf("update without default: %v", err) + } + got := GetGroupRatioCopy() + if _, ok := got["default"]; !ok { + t.Fatalf("expected default injected, got %#v", got) + } + if got["pro"] != 1.5 { + t.Fatalf("pro lost: %#v", got) + } +} + +func TestUpdateGroupRatioRejectsNegative(t *testing.T) { + if err := UpdateGroupRatioByJSONString(`{"default":-1}`); err == nil { + t.Fatal("expected error for negative ratio") + } +} + +func TestUpdateGroupGroupRatioEmptyClears(t *testing.T) { + if err := UpdateGroupGroupRatioByJSONString(`{"vip":{"default":0.9}}`); err != nil { + t.Fatalf("seed nested: %v", err) + } + if err := UpdateGroupGroupRatioByJSONString(`{}`); err != nil { + t.Fatalf("empty nested: %v", err) + } + if r, ok := GetGroupGroupRatio("vip", "default"); ok { + t.Fatalf("expected cleared nested map, still got %v", r) + } +} + +func TestUpdateGroupGroupRatioRejectsNegative(t *testing.T) { + if err := UpdateGroupGroupRatioByJSONString(`{"vip":{"default":-0.1}}`); err == nil { + t.Fatal("expected negative nested ratio error") + } +} + +func TestCheckGroupRatioAllowsEmpty(t *testing.T) { + if err := CheckGroupRatio(`{}`); err != nil { + t.Fatalf("empty should pass check: %v", err) + } + if err := CheckGroupRatio(`{"default":1}`); err != nil { + t.Fatalf("valid should pass: %v", err) + } + if err := CheckGroupRatio(`{"default":-2}`); err == nil { + t.Fatal("negative should fail check") + } +} diff --git a/setting/user_usable_group.go b/setting/user_usable_group.go index eb04b7f30534..24f85da178a1 100644 --- a/setting/user_usable_group.go +++ b/setting/user_usable_group.go @@ -1,7 +1,7 @@ package setting import ( - "encoding/json" + "strings" "sync" "github.com/QuantumNous/new-api/common" @@ -28,7 +28,7 @@ func UserUsableGroups2JSONString() string { userUsableGroupsMutex.RLock() defer userUsableGroupsMutex.RUnlock() - jsonBytes, err := json.Marshal(userUsableGroups) + jsonBytes, err := common.Marshal(userUsableGroups) if err != nil { common.SysLog("error marshalling user groups: " + err.Error()) } @@ -39,8 +39,28 @@ func UpdateUserUsableGroupsByJSONString(jsonStr string) error { userUsableGroupsMutex.Lock() defer userUsableGroupsMutex.Unlock() - userUsableGroups = make(map[string]string) - return json.Unmarshal([]byte(jsonStr), &userUsableGroups) + // Empty object wipes defaults and empties the pricing page (filter by usable groups). + trimmed := strings.TrimSpace(jsonStr) + if trimmed == "" || trimmed == "{}" || trimmed == "null" { + userUsableGroups = map[string]string{ + "default": "默认分组", + "vip": "vip分组", + } + return nil + } + tmp := make(map[string]string) + if err := common.Unmarshal([]byte(trimmed), &tmp); err != nil { + return err + } + if len(tmp) == 0 { + userUsableGroups = map[string]string{ + "default": "默认分组", + "vip": "vip分组", + } + return nil + } + userUsableGroups = tmp + return nil } func GetUsableGroupDescription(groupName string) string { diff --git a/setting/user_usable_group_empty_test.go b/setting/user_usable_group_empty_test.go new file mode 100644 index 000000000000..e0b7149043da --- /dev/null +++ b/setting/user_usable_group_empty_test.go @@ -0,0 +1,16 @@ +package setting + +import "testing" + +func TestUpdateUserUsableGroupsByJSONStringKeepsDefaultsOnEmpty(t *testing.T) { + if err := UpdateUserUsableGroupsByJSONString(`{"default":"默认分组","vip":"vip分组"}`); err != nil { + t.Fatalf("seed: %v", err) + } + if err := UpdateUserUsableGroupsByJSONString(`{}`); err != nil { + t.Fatalf("empty: %v", err) + } + got := GetUserUsableGroupsCopy() + if _, ok := got["default"]; !ok { + t.Fatalf("expected default usable group, got %#v", got) + } +} From 8f37facf4f832381d500cd2fa2073d68b38d6dba Mon Sep 17 00:00:00 2001 From: yuanjia Date: Thu, 16 Jul 2026 10:06:34 +0800 Subject: [PATCH 04/40] feat(perf-metrics): performance metrics collection and index - perf metric model with model_name equality index - /api/perf-metrics/* behind UserAuth() (401 for unauthenticated) - normalize utility with test coverage --- controller/option.go | 9 +++ controller/perf_metrics.go | 18 +++++- model/option.go | 8 ++- model/perf_metric.go | 3 + pkg/perf_metrics/metrics.go | 39 ++++++++++--- pkg/perf_metrics/normalize.go | 77 ++++++++++++++++++++++++++ pkg/perf_metrics/normalize_test.go | 50 +++++++++++++++++ setting/perf_metrics_setting/config.go | 4 +- 8 files changed, 196 insertions(+), 12 deletions(-) create mode 100644 pkg/perf_metrics/normalize.go create mode 100644 pkg/perf_metrics/normalize_test.go diff --git a/controller/option.go b/controller/option.go index a97f07b841b7..072cc8c12d46 100644 --- a/controller/option.go +++ b/controller/option.go @@ -232,6 +232,15 @@ func UpdateOption(c *gin.Context) { }) return } + case "GroupGroupRatio": + err = ratio_setting.CheckGroupGroupRatio(option.Value.(string)) + if err != nil { + c.JSON(http.StatusOK, gin.H{ + "success": false, + "message": err.Error(), + }) + return + } case "ImageRatio": err = ratio_setting.UpdateImageRatioByJSONString(option.Value.(string)) if err != nil { diff --git a/controller/perf_metrics.go b/controller/perf_metrics.go index 66d0787f2a92..f90775cb9ece 100644 --- a/controller/perf_metrics.go +++ b/controller/perf_metrics.go @@ -19,7 +19,17 @@ func GetPerfMetricsSummary(c *gin.Context) { } } - activeGroups := append(lo.Keys(ratio_setting.GetGroupRatioCopy()), "auto") + // Prefer configured groups; always keep probe/auto/default so channel tests + // and the common default group are never filtered out. When no group ratio + // is configured at all, query every group rather than returning an empty set. + groupRatio := ratio_setting.GetGroupRatioCopy() + var activeGroups []string + if len(groupRatio) == 0 { + activeGroups = nil + } else { + activeGroups = append(lo.Keys(groupRatio), "auto", "probe", "default") + activeGroups = lo.Uniq(activeGroups) + } result, err := perfmetrics.QuerySummaryAll(hours, activeGroups) if err != nil { c.JSON(http.StatusInternalServerError, gin.H{ @@ -75,8 +85,12 @@ func GetPerfMetrics(c *gin.Context) { func filterActiveGroups(groups []perfmetrics.GroupResult) []perfmetrics.GroupResult { activeRatios := ratio_setting.GetGroupRatioCopy() + // Empty group ratio means "don't filter" — same policy as summary. + if len(activeRatios) == 0 { + return groups + } return lo.Filter(groups, func(g perfmetrics.GroupResult, _ int) bool { _, ok := activeRatios[g.Group] - return ok || g.Group == "auto" + return ok || g.Group == "auto" || g.Group == "probe" || g.Group == "default" }) } diff --git a/model/option.go b/model/option.go index 8e8587f271c8..ae10b6d7d3ab 100644 --- a/model/option.go +++ b/model/option.go @@ -176,13 +176,17 @@ func InitOptionMap() { common.OptionMap["AutomaticRetryStatusCodes"] = operation_setting.AutomaticRetryStatusCodesToString() common.OptionMap["ExposeRatioEnabled"] = strconv.FormatBool(ratio_setting.IsExposeRatioEnabled()) - // 自动添加所有注册的模型配置 + // ExportAllConfigs 补充注册的配置（OptionMap 为空时填充，后面会被 DB 值覆盖） modelConfigs := config.GlobalConfig.ExportAllConfigs() for k, v := range modelConfigs { - common.OptionMap[k] = v + if _, exists := common.OptionMap[k]; !exists { + common.OptionMap[k] = v + } } common.OptionMapRWMutex.Unlock() + + // 后加载数据库值（在锁外执行，因为 updateOptionMap 内部同锁） loadOptionsFromDatabase() } diff --git a/model/perf_metric.go b/model/perf_metric.go index f9c33c851989..cdffce0f87a8 100644 --- a/model/perf_metric.go +++ b/model/perf_metric.go @@ -1,6 +1,7 @@ package model import ( + "strings" "time" "gorm.io/gorm" @@ -50,6 +51,8 @@ func UpsertPerfMetric(metric *PerfMetric) error { func GetPerfMetrics(modelName string, group string, startTs int64, endTs int64) ([]PerfMetric, error) { var metrics []PerfMetric + // Keep in sync with pkg/perf_metrics.NormalizeModelName — no import (cycle). + modelName = strings.ToLower(strings.TrimSpace(modelName)) query := DB.Model(&PerfMetric{}). Where("model_name = ? AND bucket_ts >= ? AND bucket_ts <= ?", modelName, startTs, endTs) if group != "" { diff --git a/pkg/perf_metrics/metrics.go b/pkg/perf_metrics/metrics.go index 33b79ee478e9..5a2966d0a075 100644 --- a/pkg/perf_metrics/metrics.go +++ b/pkg/perf_metrics/metrics.go @@ -56,6 +56,7 @@ func RecordRelaySample(info *relaycommon.RelayInfo, success bool, outputTokens i func Record(sample Sample) { setting := perf_metrics_setting.GetSetting() + sample.Model = NormalizeModelName(sample.Model) if !setting.Enabled || sample.Model == "" { return } @@ -83,6 +84,10 @@ func Query(params QueryParams) (QueryResult, error) { if params.Hours > 24*30 { params.Hours = 24 * 30 } + params.Model = NormalizeModelName(params.Model) + if params.Model == "" { + return QueryResult{SeriesSchema: seriesSchema, Groups: []GroupResult{}}, nil + } endTs := time.Now().Unix() startTs := endTs - int64(params.Hours)*3600 @@ -92,8 +97,9 @@ func Query(params QueryParams) (QueryResult, error) { return QueryResult{}, err } for _, row := range rows { + // Historical rows may still use mixed casing; fold into the normalized key. mergeCounters(merged, bucketKey{ - model: row.ModelName, + model: params.Model, group: row.Group, bucketTs: row.BucketTs, }, counters{ @@ -109,13 +115,18 @@ func Query(params QueryParams) (QueryResult, error) { hotBuckets.Range(func(key, value any) bool { k := key.(bucketKey) - if k.model != params.Model || k.bucketTs < startTs || k.bucketTs > endTs { + if NormalizeModelName(k.model) != params.Model || k.bucketTs < startTs || k.bucketTs > endTs { return true } if params.Group != "" && k.group != params.Group { return true } - mergeCounters(merged, k, value.(*atomicBucket).snapshot()) + // Re-key under the normalized model so mixed-case hot buckets collapse. + mergeCounters(merged, bucketKey{ + model: params.Model, + group: k.group, + bucketTs: k.bucketTs, + }, value.(*atomicBucket).snapshot()) return true }) @@ -148,8 +159,13 @@ func QuerySummaryAll(hours int, groups []string) (SummaryAllResult, error) { outputTokens: row.OutputTokens, generationMs: row.GenerationMs, } - mergeModelTotals(totals, row.ModelName, value) - mergeModelBucket(modelBuckets, row.ModelName, row.BucketTs, value) + // Collapse mixed-case historical rows onto one summary key. + name := NormalizeModelName(row.ModelName) + if name == "" { + continue + } + mergeModelTotals(totals, name, value) + mergeModelBucket(modelBuckets, name, row.BucketTs, value) } hotBuckets.Range(func(key, value any) bool { @@ -166,8 +182,12 @@ func QuerySummaryAll(hours int, groups []string) (SummaryAllResult, error) { if snap.requestCount == 0 { return true } - mergeModelTotals(totals, k.model, snap) - mergeModelBucket(modelBuckets, k.model, k.bucketTs, snap) + name := NormalizeModelName(k.model) + if name == "" { + return true + } + mergeModelTotals(totals, name, snap) + mergeModelBucket(modelBuckets, name, k.bucketTs, snap) return true }) @@ -176,6 +196,11 @@ func QuerySummaryAll(hours int, groups []string) (SummaryAllResult, error) { if total.requestCount == 0 { continue } + // Drop image/audio/embedding probe noise from the square summary. + // Detail Query still serves any model name for debugging. + if !IsChatCapableModelName(name) { + continue + } avgLatency := total.totalLatencyMs / total.requestCount successRate := float64(total.successCount) / float64(total.requestCount) * 100 avgTps := 0.0 diff --git a/pkg/perf_metrics/normalize.go b/pkg/perf_metrics/normalize.go new file mode 100644 index 000000000000..7c521888d045 --- /dev/null +++ b/pkg/perf_metrics/normalize.go @@ -0,0 +1,77 @@ +package perfmetrics + +import "strings" + +// NormalizeModelName folds model identifiers for perf storage and lookup. +// Pricing / abilities may store mixed casings of the same model +// (e.g. deepseek-v4-flash vs Deepseek-V4-Flash); health metrics must key +// them together so badges and detail views resolve real samples. +// +// Only case + surrounding whitespace are changed — path-style names and +// free-tier suffixes stay distinct (a/b, :free, [free]). +func NormalizeModelName(name string) string { + return strings.ToLower(strings.TrimSpace(name)) +} + +// IsChatCapableModelName reports whether a model name is suitable for the +// model-square health summary. Image / audio / video / embedding / rerank +// probes produce noisy keys that should not pollute the chat health view. +// Kept here (not controller) so QuerySummaryAll can filter without import cycles. +func IsChatCapableModelName(name string) bool { + name = NormalizeModelName(name) + if name == "" { + return false + } + if isImageLikeModelName(name) || isAudioOrVideoLikeModelName(name) || isEmbeddingOrRerankModelName(name) { + return false + } + return true +} + +func isImageLikeModelName(name string) bool { + imageHints := []string{ + "gpt-image", "dall-e", "dalle", "seedream", "flux", "imagen", + "stable-diffusion", "sdxl", "midjourney", "mj-", "image-gen", + "text-to-image", "t2i", "cogview", "kolors", "playground-v", + } + for _, h := range imageHints { + if strings.Contains(name, h) { + return true + } + } + if strings.Contains(name, "image") && + !strings.Contains(name, "vision") && + !strings.Contains(name, "chat") && + !strings.Contains(name, "embedding") { + return true + } + return false +} + +func isAudioOrVideoLikeModelName(name string) bool { + hints := []string{ + "whisper", "tts-", "tts_", "-tts", "speech", "audio-", "-audio", + "sora", "kling", "runway", "luma", "hailuo", "vidu", "cogvideo", + "text-to-video", "t2v", "minimax-video", + } + for _, h := range hints { + if strings.Contains(name, h) { + return true + } + } + return false +} + +func isEmbeddingOrRerankModelName(name string) bool { + if strings.Contains(name, "rerank") { + return true + } + if strings.Contains(name, "embedding") || + strings.Contains(name, "embed") || + strings.HasPrefix(name, "m3e") || + strings.Contains(name, "bge-") || + strings.Contains(name, "text-embedding") { + return true + } + return false +} diff --git a/pkg/perf_metrics/normalize_test.go b/pkg/perf_metrics/normalize_test.go new file mode 100644 index 000000000000..316438b509b1 --- /dev/null +++ b/pkg/perf_metrics/normalize_test.go @@ -0,0 +1,50 @@ +package perfmetrics + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestNormalizeModelName(t *testing.T) { + t.Parallel() + cases := []struct { + in, want string + }{ + {"Deepseek-V4-Flash", "deepseek-v4-flash"}, + {" deepseek-v4-flash ", "deepseek-v4-flash"}, + {"DeepSeek-V4-Flash", "deepseek-v4-flash"}, + {"gpt-4o", "gpt-4o"}, + {"provider/Path/Model", "provider/path/model"}, + {"model:free", "model:free"}, + {"model[free]", "model[free]"}, + {"", ""}, + {" ", ""}, + } + for _, tc := range cases { + tc := tc + t.Run(tc.in, func(t *testing.T) { + t.Parallel() + require.Equal(t, tc.want, NormalizeModelName(tc.in)) + }) + } +} + +func TestNormalizeModelNameCollapsesCasingOnly(t *testing.T) { + t.Parallel() + a := NormalizeModelName("Deepseek-V4-Flash") + b := NormalizeModelName("deepseek-v4-flash") + require.Equal(t, a, b) + require.NotEqual(t, NormalizeModelName("foo:free"), NormalizeModelName("foo")) +} + +func TestIsChatCapableModelName(t *testing.T) { + t.Parallel() + require.True(t, IsChatCapableModelName("gpt-4o-mini")) + for _, name := range []string{ + "gpt-image-2", "dall-e-3", "whisper-1", "sora-2", + "text-embedding-3-small", "bge-m3", "jina-rerank-v2", + } { + require.Falsef(t, IsChatCapableModelName(name), "%s should not be chat capable for summary", name) + } +} diff --git a/setting/perf_metrics_setting/config.go b/setting/perf_metrics_setting/config.go index fb7780e53b4c..bc887b9560be 100644 --- a/setting/perf_metrics_setting/config.go +++ b/setting/perf_metrics_setting/config.go @@ -13,7 +13,9 @@ var perfMetricsSetting = PerfMetricsSetting{ Enabled: true, FlushInterval: 5, BucketTime: "hour", - RetentionDays: 0, + // Default 1 day — SQLite probe traffic grows quickly; ops can raise via UI. + // 0 still means "keep forever" when explicitly set. + RetentionDays: 1, } func init() { From 50b5a5a9d0d703b8971125aa124c050eace67b7e Mon Sep 17 00:00:00 2001 From: yuanjia Date: Thu, 16 Jul 2026 10:06:52 +0800 Subject: [PATCH 05/40] refactor(common): deduplicate constants, init ordering, and env wiring - consolidate env var reads into common/env.go - init sequence cleanup in common/init.go - context key constants and API router alignment - SESSION_SECURE default true (overridden by start script) --- common/constants.go | 2 +- common/env.go | 12 ++++++++++++ common/init.go | 9 +++++++++ constant/context_key.go | 5 +++++ constant/env.go | 9 +++++++++ main.go | 22 ++++++++++++++++++---- router/api-router.go | 6 ++++-- router/channel-router.go | 1 + 8 files changed, 59 insertions(+), 7 deletions(-) diff --git a/common/constants.go b/common/constants.go index 87d212f99732..3154dc2b8aff 100644 --- a/common/constants.go +++ b/common/constants.go @@ -22,7 +22,7 @@ var TopUpLink = "" var themeValue atomic.Value // stores string; safe for concurrent read/write func init() { - themeValue.Store("classic") + themeValue.Store("default") } func GetTheme() string { diff --git a/common/env.go b/common/env.go index 1aa340f85ea1..2ee085b6f7aa 100644 --- a/common/env.go +++ b/common/env.go @@ -36,3 +36,15 @@ func GetEnvOrDefaultBool(env string, defaultValue bool) bool { } return b } + +func GetEnvOrDefaultFloat(env string, defaultValue float64) float64 { + if env == "" || os.Getenv(env) == "" { + return defaultValue + } + f, err := strconv.ParseFloat(os.Getenv(env), 64) + if err != nil { + SysError(fmt.Sprintf("failed to parse %s: %s, using default value: %.2f", env, err.Error(), defaultValue)) + return defaultValue + } + return f +} diff --git a/common/init.go b/common/init.go index 88b2dc3e62e1..553e2f4aeaad 100644 --- a/common/init.go +++ b/common/init.go @@ -185,4 +185,13 @@ func initConstantEnv() { } } constant.TrustedRedirectDomains = trustedDomains + + // Adaptive channel balance + constant.AdaptiveBalanceEnabled = GetEnvOrDefaultBool("ADAPTIVE_BALANCE_ENABLED", false) + constant.AdaptiveBalanceShadowMode = GetEnvOrDefaultBool("ADAPTIVE_BALANCE_SHADOW_MODE", false) + constant.ChannelCircuitBreakerEnabled = GetEnvOrDefaultBool("CHANNEL_CIRCUIT_BREAKER_ENABLED", false) + constant.MaxRetryChannels = GetEnvOrDefault("MAX_RETRY_CHANNELS", 3) + constant.ChannelCooldownSeconds = GetEnvOrDefault("CHANNEL_COOLDOWN_SECONDS", 30) + constant.EwmaAlpha = GetEnvOrDefaultFloat("EWMA_ALPHA", 0.1) + constant.MaxChannelConcurrency = GetEnvOrDefault("MAX_CHANNEL_CONCURRENCY", 10) } diff --git a/constant/context_key.go b/constant/context_key.go index b856bc3dda14..de9c309b4786 100644 --- a/constant/context_key.go +++ b/constant/context_key.go @@ -72,4 +72,9 @@ const ( // fallback in authHelper (finishAdminAudit) skips its record to avoid // duplicate entries. ContextKeyAuditLogged ContextKey = "audit_logged" + + // ContextKeyThreadId / ContextKeyTraceId hold AxonHub-compatible + // conversation observability IDs (AH-Thread-Id / AH-Trace-Id). + ContextKeyThreadId ContextKey = "thread_id" + ContextKeyTraceId ContextKey = "trace_id" ) diff --git a/constant/env.go b/constant/env.go index 512bfc31126b..0dfd794bb590 100644 --- a/constant/env.go +++ b/constant/env.go @@ -25,3 +25,12 @@ var TaskPricePatches []string // TrustedRedirectDomains is a list of trusted domains for redirect URL validation. // Domains support subdomain matching (e.g., "example.com" matches "sub.example.com"). var TrustedRedirectDomains []string + +// Adaptive channel balance settings +var AdaptiveBalanceEnabled bool +var AdaptiveBalanceShadowMode bool +var ChannelCircuitBreakerEnabled bool +var MaxRetryChannels int +var ChannelCooldownSeconds int +var EwmaAlpha float64 +var MaxChannelConcurrency int diff --git a/main.go b/main.go index 770ea156ba86..6e03469da741 100644 --- a/main.go +++ b/main.go @@ -79,6 +79,14 @@ func main() { if common.RedisEnabled { // for compatibility with old versions common.MemoryCacheEnabled = true + // Multi-instance adaptive metrics snapshot (best-effort, 2m TTL). + gopool.Go(func() { + ticker := time.NewTicker(30 * time.Second) + defer ticker.Stop() + for range ticker.C { + service.SyncAdaptiveMetricsToRedis() + } + }) } if common.MemoryCacheEnabled { common.SysLog("memory cache enabled") @@ -160,10 +168,10 @@ func main() { if os.Getenv("ENABLE_PPROF") == "true" { gopool.Go(func() { - log.Println(http.ListenAndServe("0.0.0.0:8005", nil)) + log.Println(http.ListenAndServe("127.0.0.1:8005", nil)) }) go common.Monitor() - common.SysLog("pprof enabled") + common.SysLog("pprof enabled on 127.0.0.1:8005") } err = common.StartPyroScope() @@ -174,10 +182,15 @@ func main() { // Initialize HTTP server server := gin.New() server.Use(gin.CustomRecovery(func(c *gin.Context, err any) { - common.SysLog(fmt.Sprintf("panic detected: %v", err)) + reqID := c.GetString(common.RequestIdKey) + common.SysLog(fmt.Sprintf("panic detected request_id=%s: %v", reqID, err)) + msg := "Internal server error" + if reqID != "" { + msg = fmt.Sprintf("Internal server error (request_id=%s)", reqID) + } c.JSON(http.StatusInternalServerError, gin.H{ "error": gin.H{ - "message": fmt.Sprintf("Panic detected, error: %v. Please submit a issue here: https://github.com/Calcium-Ion/new-api", err), + "message": msg, "type": "new_api_panic", }, }) @@ -186,6 +199,7 @@ func main() { //server.Use(gzip.Gzip(gzip.DefaultCompression)) server.Use(middleware.RequestId()) server.Use(middleware.Version()) + server.Use(middleware.TraceContext()) server.Use(middleware.I18n()) middleware.SetUpLogger(server) // Initialize session store diff --git a/router/api-router.go b/router/api-router.go index 83f9259b2132..b6c57de6fe51 100644 --- a/router/api-router.go +++ b/router/api-router.go @@ -32,8 +32,9 @@ func SetApiRouter(router *gin.Engine) { //apiRouter.GET("/midjourney", controller.GetMidjourney) apiRouter.GET("/home_page_content", controller.GetHomePageContent) apiRouter.GET("/pricing", middleware.HeaderNavModuleAuth("pricing"), controller.GetPricing) + // Require login — model traffic profiles are not public intel. perfMetricsRoute := apiRouter.Group("/perf-metrics") - perfMetricsRoute.Use(middleware.HeaderNavModulePublicOrUserAuth("pricing")) + perfMetricsRoute.Use(middleware.UserAuth()) { perfMetricsRoute.GET("/summary", controller.GetPerfMetricsSummary) perfMetricsRoute.GET("", controller.GetPerfMetrics) @@ -75,7 +76,7 @@ func SetApiRouter(router *gin.Engine) { userRoute.GET("/logout", controller.Logout) userRoute.POST("/epay/notify", anonymousRequestBodyLimit, controller.EpayNotify) userRoute.GET("/epay/notify", controller.EpayNotify) - userRoute.GET("/groups", controller.GetUserGroups) + userRoute.GET("/groups", middleware.UserAuth(), controller.GetUserGroups) selfRoute := userRoute.Group("/") selfRoute.Use(middleware.UserAuth()) @@ -269,6 +270,7 @@ func SetApiRouter(router *gin.Engine) { // Legacy synchronous direct-delete route used only by the classic frontend. // TODO: remove once the classic frontend is removed; the default frontend uses /system-task/log-cleanup. logRoute.DELETE("/", middleware.RootAuth(), controller.DeleteHistoryLogs) + logRoute.GET("/trace/:trace_id", middleware.AdminAuth(), controller.GetTraceLogs) logRoute.GET("/stat", middleware.AdminAuth(), controller.GetLogsStat) logRoute.GET("/self/stat", middleware.UserAuth(), controller.GetLogsSelfStat) logRoute.GET("/channel_affinity_usage_cache", middleware.AdminAuth(), controller.GetChannelAffinityUsageCacheStats) diff --git a/router/channel-router.go b/router/channel-router.go index b85cbd884b77..c7d8f658bef6 100644 --- a/router/channel-router.go +++ b/router/channel-router.go @@ -69,6 +69,7 @@ var channelPermissionRoutes = []permissionRoute{ {method: http.MethodDelete, path: "/ollama/delete", permission: authz.ChannelSensitiveWrite, handler: controller.OllamaDeleteModel}, {method: http.MethodGet, path: "/ollama/version/:id", permission: authz.ChannelSensitiveWrite, handler: controller.OllamaVersion}, {method: http.MethodPost, path: "/batch/tag", permission: authz.ChannelWrite, handler: controller.BatchSetChannelTag}, + {method: http.MethodPost, path: "/batch/skip_auto_test", permission: authz.ChannelWrite, handler: controller.BatchSetChannelSkipAutoTest}, {method: http.MethodGet, path: "/tag/models", permission: authz.ChannelRead, handler: controller.GetTagModels}, {method: http.MethodPost, path: "/copy/:id", permission: authz.ChannelSensitiveWrite, handler: controller.CopyChannel}, {method: http.MethodPost, path: "/multi_key/manage", permission: authz.ChannelOperate, handler: controller.ManageMultiKeys}, From 2c5a57f431f1358560b967ef477061db4cb8c49c Mon Sep 17 00:00:00 2001 From: yuanjia Date: Thu, 16 Jul 2026 10:07:09 +0800 Subject: [PATCH 06/40] fix(frontend): security hardening and auth flow fixes - Markdown sanitize with allowlist (no svg/data/javascript/iframe-same-origin) - Auth redirect safeRedirect to prevent open redirect - OTP flow component hydration and route fixes - lib/api.ts custom fetch with session verification - About page iframe sandbox removal --- web/default/rsbuild.config.ts | 22 ++ .../components/layout/components/footer.tsx | 8 +- .../features/auth/hooks/use-auth-redirect.ts | 27 ++- .../src/features/auth/lib/safe-redirect.ts | 51 +++++ .../features/auth/lib/session-verification.ts | 31 +++ .../features/auth/otp/components/otp-form.tsx | 12 +- web/default/src/features/auth/otp/index.tsx | 5 +- .../sign-in/components/user-auth-form.tsx | 2 +- web/default/src/lib/api.ts | 2 + web/default/src/lib/sanitize-html.ts | 193 ++++++++++++++++++ web/default/src/main.tsx | 24 ++- web/default/src/routes/oauth/$provider.tsx | 25 ++- web/default/src/stores/auth-store.ts | 1 + web/default/tsconfig.app.json | 3 +- 14 files changed, 360 insertions(+), 46 deletions(-) create mode 100644 web/default/src/features/auth/lib/safe-redirect.ts create mode 100644 web/default/src/features/auth/lib/session-verification.ts create mode 100644 web/default/src/lib/sanitize-html.ts diff --git a/web/default/rsbuild.config.ts b/web/default/rsbuild.config.ts index 3b9d11de6a88..93e2b39c5336 100644 --- a/web/default/rsbuild.config.ts +++ b/web/default/rsbuild.config.ts @@ -50,6 +50,28 @@ export default defineConfig(({ envMode }) => { priority: 0, enforce: true, }, + 'vendor-vchart': { + test: /node_modules[\\/]@visactor[\\/]/, + name: 'vendor-vchart', + chunks: 'all', + priority: 0, + enforce: true, + }, + 'vendor-shiki': { + test: /node_modules[\\/]shiki[\\/]/, + name: 'vendor-shiki', + chunks: 'async', + priority: 0, + enforce: true, + }, + 'vendor-recharts': { + // Legacy shadcn chart helper; keep out of main if unused on critical path. + test: /node_modules[\\/]recharts[\\/]/, + name: 'vendor-recharts', + chunks: 'async', + priority: 0, + enforce: true, + }, }, }, source: { diff --git a/web/default/src/components/layout/components/footer.tsx b/web/default/src/components/layout/components/footer.tsx index 438fc6cfe482..fb095edb361a 100644 --- a/web/default/src/components/layout/components/footer.tsx +++ b/web/default/src/components/layout/components/footer.tsx @@ -16,13 +16,13 @@ along with this program. If not, see . For commercial licensing, please contact support@quantumnous.com */ -import { Link } from '@tanstack/react-router' import { Fragment, useMemo } from 'react' +import { Link } from '@tanstack/react-router' import { useTranslation } from 'react-i18next' - +import { cn } from '@/lib/utils' import { useStatus } from '@/hooks/use-status' import { useSystemConfig } from '@/hooks/use-system-config' -import { cn } from '@/lib/utils' +import { sanitizeHtml } from '@/lib/sanitize-html' interface FooterLink { text: string @@ -234,7 +234,7 @@ export function Footer(props: FooterProps) {
diff --git a/web/default/src/features/auth/hooks/use-auth-redirect.ts b/web/default/src/features/auth/hooks/use-auth-redirect.ts index 1da607161a8d..1d1b949016fe 100644 --- a/web/default/src/features/auth/hooks/use-auth-redirect.ts +++ b/web/default/src/features/auth/hooks/use-auth-redirect.ts @@ -18,11 +18,10 @@ For commercial licensing, please contact support@quantumnous.com */ import { useNavigate } from '@tanstack/react-router' import i18n from 'i18next' - -import type { User } from '@/features/users/types' -import { getSelf } from '@/lib/api' import { useAuthStore } from '@/stores/auth-store' - +import { getSelf } from '@/lib/api' +import type { User } from '@/features/users/types' +import { safeRedirect } from '../lib/safe-redirect' import { saveUserId } from '../lib/storage' function getSavedLanguage(user: User): string | undefined { @@ -87,23 +86,31 @@ export function useAuthRedirect() { console.error('Failed to fetch user data:', error) } - // Navigate to target page - const targetPath = redirectTo || '/dashboard' + // Navigate to target page (same-app paths only) + const targetPath = safeRedirect(redirectTo, '/dashboard') navigate({ to: targetPath, replace: true }) } /** * Redirect to 2FA page */ - const redirectTo2FA = () => { - navigate({ to: '/otp', replace: true }) + const redirectTo2FA = (redirectTo?: string) => { + navigate({ + to: '/otp', + search: { redirect: safeRedirect(redirectTo, '/dashboard') }, + replace: true, + }) } /** * Redirect to login page */ - const redirectToLogin = () => { - navigate({ to: '/sign-in', replace: true }) + const redirectToLogin = (redirectTo?: string) => { + navigate({ + to: '/sign-in', + search: { redirect: safeRedirect(redirectTo, '/dashboard') }, + replace: true, + }) } /** diff --git a/web/default/src/features/auth/lib/safe-redirect.ts b/web/default/src/features/auth/lib/safe-redirect.ts new file mode 100644 index 000000000000..e7606b737882 --- /dev/null +++ b/web/default/src/features/auth/lib/safe-redirect.ts @@ -0,0 +1,51 @@ +/* +Copyright (C) 2023-2026 QuantumNous + +This program is free software: you can redistribute it and/or modify +it under the terms of the GNU Affero General Public License as +published by the Free Software Foundation, either version 3 of the +License, or (at your option) any later version. + +This program is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU Affero General Public License for more details. + +You should have received a copy of the GNU Affero General Public License +along with this program. If not, see . + +For commercial licensing, please contact support@quantumnous.com +*/ + +/** + * Allow only same-app relative paths after login. + * Blocks open redirects: //evil, https://evil, javascript:, etc. + */ +export function safeRedirect( + path: string | null | undefined, + fallback = '/dashboard' +): string { + if (!path || typeof path !== 'string') return fallback + let value = path.trim() + if (!value) return fallback + + // Absolute URL → keep pathname+search+hash if same origin, else fallback + try { + if (/^[a-zA-Z][a-zA-Z\d+\-.]*:/.test(value) || value.startsWith('//')) { + if (typeof window !== 'undefined') { + const url = new URL(value, window.location.origin) + if (url.origin !== window.location.origin) return fallback + value = `${url.pathname}${url.search}${url.hash}` + } else { + return fallback + } + } + } catch { + return fallback + } + + if (!value.startsWith('/')) return fallback + if (value.startsWith('//')) return fallback + if (value.toLowerCase().includes('javascript:')) return fallback + return value +} diff --git a/web/default/src/features/auth/lib/session-verification.ts b/web/default/src/features/auth/lib/session-verification.ts new file mode 100644 index 000000000000..721a2c68279c --- /dev/null +++ b/web/default/src/features/auth/lib/session-verification.ts @@ -0,0 +1,31 @@ +/* +Copyright (C) 2023-2026 QuantumNous + +This program is free software: you can redistribute it and/or modify +it under the terms of the GNU Affero General Public License as +published by the Free Software Foundation, either version 3 of the +License, or (at your option) any later version. + +This program is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU Affero General Public License for more details. + +You should have received a copy of the GNU Affero General Public License +along with this program. If not, see . + +For commercial licensing, please contact support@quantumnous.com +*/ +let sessionVerified = false + +export function isSessionVerified(): boolean { + return sessionVerified +} + +export function markSessionVerified(): void { + sessionVerified = true +} + +export function resetSessionVerified(): void { + sessionVerified = false +} diff --git a/web/default/src/features/auth/otp/components/otp-form.tsx b/web/default/src/features/auth/otp/components/otp-form.tsx index cf5c1fcbaac7..0edd1f632ac9 100644 --- a/web/default/src/features/auth/otp/components/otp-form.tsx +++ b/web/default/src/features/auth/otp/components/otp-form.tsx @@ -59,15 +59,17 @@ import type { User } from '@/features/users/types' import { cn } from '@/lib/utils' import { useAuthStore } from '@/stores/auth-store' -type OtpFormProps = React.HTMLAttributes +type OtpFormProps = React.HTMLAttributes & { + redirectTo?: string +} -export function OtpForm({ className, ...props }: OtpFormProps) { +export function OtpForm({ className, redirectTo, ...props }: OtpFormProps) { const { t } = useTranslation() const [isLoading, setIsLoading] = useState(false) const [useBackupCode, setUseBackupCode] = useState(false) const { auth } = useAuthStore() - const { redirectToLogin } = useAuthRedirect() + const { handleLoginSuccess, redirectToLogin } = useAuthRedirect() const form = useForm>({ resolver: zodResolver(otpFormSchema), @@ -116,7 +118,7 @@ export function OtpForm({ className, ...props }: OtpFormProps) { } toast.success(t('Signed in')) - redirectToLogin() // This will redirect to dashboard via the redirect logic + await handleLoginSuccess(userData, redirectTo) } catch (error) { // eslint-disable-next-line no-console console.error('2FA verification error:', error) @@ -134,7 +136,7 @@ export function OtpForm({ className, ...props }: OtpFormProps) { } function handleBackToLogin() { - redirectToLogin() + redirectToLogin(redirectTo) } const isFormValid = useBackupCode diff --git a/web/default/src/features/auth/otp/index.tsx b/web/default/src/features/auth/otp/index.tsx index 6e248b539a4b..ea02191024ae 100644 --- a/web/default/src/features/auth/otp/index.tsx +++ b/web/default/src/features/auth/otp/index.tsx @@ -16,7 +16,7 @@ along with this program. If not, see . For commercial licensing, please contact support@quantumnous.com */ -import { Link } from '@tanstack/react-router' +import { Link, useSearch } from '@tanstack/react-router' import { useTranslation } from 'react-i18next' import { AuthLayout } from '../auth-layout' @@ -24,6 +24,7 @@ import { OtpForm } from './components/otp-form' export function Otp() { const { t } = useTranslation() + const { redirect } = useSearch({ from: '/(auth)/otp' }) return (
@@ -46,7 +47,7 @@ export function Otp() {

- +
) diff --git a/web/default/src/features/auth/sign-in/components/user-auth-form.tsx b/web/default/src/features/auth/sign-in/components/user-auth-form.tsx index 913c98862055..98c4c48df158 100644 --- a/web/default/src/features/auth/sign-in/components/user-auth-form.tsx +++ b/web/default/src/features/auth/sign-in/components/user-auth-form.tsx @@ -161,7 +161,7 @@ export function UserAuthForm({ if (res.success) { if (res.data?.require_2fa) { - redirectTo2FA() + redirectTo2FA(redirectTo) return } diff --git a/web/default/src/lib/api.ts b/web/default/src/lib/api.ts index 1e50b16c513b..cdd2c4194625 100644 --- a/web/default/src/lib/api.ts +++ b/web/default/src/lib/api.ts @@ -21,6 +21,7 @@ import { t } from 'i18next' import { toast } from 'sonner' import { useAuthStore } from '@/stores/auth-store' +import { resetSessionVerified } from '@/features/auth/lib/session-verification' declare module 'axios' { export interface AxiosRequestConfig { @@ -103,6 +104,7 @@ api.interceptors.response.use( if (status === 401) { try { + resetSessionVerified() useAuthStore.getState().auth.reset() } catch { /* empty */ diff --git a/web/default/src/lib/sanitize-html.ts b/web/default/src/lib/sanitize-html.ts new file mode 100644 index 000000000000..d378a73178b8 --- /dev/null +++ b/web/default/src/lib/sanitize-html.ts @@ -0,0 +1,193 @@ +/* +Copyright (C) 2023-2026 QuantumNous + +This program is free software: you can redistribute it and/or modify +it under the terms of the GNU Affero General Public License as +published by the Free Software Foundation, either version 3 of the +License, or (at your option) any later version. + +This program is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU Affero General Public License for more details. + +You should have received a copy of the GNU Affero General Public License +along with this program. If not, see . + +For commercial licensing, please contact support@quantumnous.com +*/ + +/** + * Lightweight HTML sanitizer for admin-configured content. + * No external deps (workspace may not have dompurify installed). + */ + +const FORBIDDEN_TAGS = new Set([ + 'script', + 'style', + 'iframe', + 'object', + 'embed', + 'link', + 'meta', + 'base', + 'form', + 'input', + 'button', + 'textarea', + 'select', + 'svg', + 'math', + 'template', + 'foreignobject', + 'use', + 'animate', + 'set', + 'video', + 'audio', + 'source', + 'track', + 'frame', + 'frameset', + 'applet', + 'marquee', +]) + +const URL_ATTRS = new Set(['href', 'src', 'xlink:href', 'action', 'formaction', 'poster']) + +function isSafeUrl(value: string): boolean { + const v = value.trim().toLowerCase() + if (!v) return false + if (v.startsWith('#')) return true + if (v.startsWith('/') && !v.startsWith('//')) return true + if (v.startsWith('//')) return false + if (v.startsWith('data:') || v.startsWith('blob:') || v.startsWith('javascript:')) { + return false + } + if ( + v.startsWith('https://') || + v.startsWith('http://') || + v.startsWith('mailto:') + ) { + return true + } + return false +} + +function sanitizeNode(node: Node, doc: Document): Node | null { + if (node.nodeType === Node.TEXT_NODE) { + return doc.createTextNode(node.textContent ?? '') + } + if (node.nodeType !== Node.ELEMENT_NODE) { + return null + } + + const el = node as Element + const tag = el.tagName.toLowerCase() + if (FORBIDDEN_TAGS.has(tag)) { + return null + } + + // Only allow a conservative set of content tags (drop unknown custom tags). + const ALLOWED_TAGS = new Set([ + 'a', 'abbr', 'b', 'blockquote', 'br', 'caption', 'code', 'col', 'colgroup', + 'dd', 'del', 'details', 'div', 'dl', 'dt', 'em', 'figcaption', 'figure', + 'h1', 'h2', 'h3', 'h4', 'h5', 'h6', 'hr', 'i', 'img', 'ins', 'kbd', 'li', + 'mark', 'ol', 'p', 'pre', 'q', 's', 'samp', 'section', 'small', 'span', + 'strong', 'sub', 'summary', 'sup', 'table', 'tbody', 'td', 'tfoot', 'th', + 'thead', 'tr', 'u', 'ul', 'var', + ]) + if (!ALLOWED_TAGS.has(tag)) { + // Keep text children of unknown tags, drop the wrapper. + const frag = doc.createDocumentFragment() + for (const child of Array.from(el.childNodes)) { + const c = sanitizeNode(child, doc) + if (c) frag.appendChild(c) + } + return frag.childNodes.length ? frag : null + } + + const clean = doc.createElement(tag) + + for (const attr of Array.from(el.attributes)) { + const name = attr.name.toLowerCase() + const value = attr.value + if (name.startsWith('on')) continue + if (name === 'srcdoc' || name === 'srcset') continue + if (name === 'style') continue + if (URL_ATTRS.has(name) || name === 'href' || name === 'src') { + if (!isSafeUrl(value)) continue + clean.setAttribute(attr.name, value) + if (tag === 'a' && (name === 'href' || name === 'src') && !clean.hasAttribute('rel')) { + clean.setAttribute('rel', 'noopener noreferrer') + } + if (tag === 'a' && name === 'target') { + clean.setAttribute('target', '_blank') + } + continue + } + if ( + name === 'class' || + name === 'id' || + name === 'title' || + name === 'alt' || + name === 'width' || + name === 'height' || + name === 'colspan' || + name === 'rowspan' || + name === 'scope' || + name === 'target' || + name === 'rel' || + name.startsWith('aria-') + ) { + // Drop data-* to avoid mXSS / framework side channels via admin HTML. + clean.setAttribute(attr.name, value) + } + } + + for (const child of Array.from(el.childNodes)) { + const c = sanitizeNode(child, doc) + if (c) clean.appendChild(c) + } + return clean +} + +/** Sanitize admin-configured HTML before dangerouslySetInnerHTML. */ +export function sanitizeHtml(dirty: string): string { + if (!dirty) return '' + if (typeof window === 'undefined' || typeof DOMParser === 'undefined') { + return dirty + .replace(/[\s\S]*?<\/script>/gi, '') + .replace(/[\s\S]*?<\/svg>/gi, '') + .replace(/on\w+\s*=\s*("[^"]*"|'[^']*'|[^\s>]+)/gi, '') + .replace(/javascript\s*:/gi, '') + } + try { + const parser = new DOMParser() + const doc = parser.parseFromString( + `
${dirty}
`, + 'text/html' + ) + const root = doc.getElementById('__root') + if (!root) return '' + const out = doc.createElement('div') + for (const child of Array.from(root.childNodes)) { + const c = sanitizeNode(child, doc) + if (c) out.appendChild(c) + } + return out.innerHTML + } catch { + return '' + } +} + +/** Allow only http(s) iframe sources. */ +export function sanitizeIframeSrc(url: string): string | null { + try { + const u = new URL(url) + if (u.protocol === 'https:' || u.protocol === 'http:') return u.toString() + } catch { + /* empty */ + } + return null +} diff --git a/web/default/src/main.tsx b/web/default/src/main.tsx index caf9307c1214..526cd4ed8839 100644 --- a/web/default/src/main.tsx +++ b/web/default/src/main.tsx @@ -16,33 +16,32 @@ along with this program. If not, see . For commercial licensing, please contact support@quantumnous.com */ +import { StrictMode } from 'react' +import ReactDOM from 'react-dom/client' +import { AxiosError } from 'axios' import { QueryCache, QueryClient, QueryClientProvider, } from '@tanstack/react-query' import { RouterProvider, createRouter } from '@tanstack/react-router' -import { AxiosError } from 'axios' import i18next from 'i18next' -import { StrictMode } from 'react' -import ReactDOM from 'react-dom/client' import { toast } from 'sonner' - +import { useAuthStore } from '@/stores/auth-store' import { getStatus } from '@/lib/api' import { installBuildMetadata } from '@/lib/build-metadata' -import { applyFaviconToDom } from '@/lib/dom-utils' import '@/lib/dayjs' +import { applyFaviconToDom } from '@/lib/dom-utils' import { initializeFrontendCache } from '@/lib/frontend-cache' import { handleServerError } from '@/lib/handle-server-error' -import { useAuthStore } from '@/stores/auth-store' - +import { safeRedirect } from '@/features/auth/lib/safe-redirect' +import { resetSessionVerified } from '@/features/auth/lib/session-verification' import { DirectionProvider } from './context/direction-provider' import { FontProvider } from './context/font-provider' import { ThemeProvider } from './context/theme-provider' import './i18n/config' // Generated Routes import { routeTree } from './routeTree.gen' - // Styles import './styles/index.css' @@ -88,12 +87,17 @@ const queryClient = new QueryClient({ if (error.response?.status === 401) { toast.error(i18next.t('Session expired!')) useAuthStore.getState().auth.reset() - const redirect = `${router.history.location.href}` + resetSessionVerified() + const loc = router.history.location + const redirect = safeRedirect( + `${loc.pathname}${loc.search ?? ''}${loc.hash ?? ''}`, + '/dashboard' + ) router.navigate({ to: '/sign-in', search: { redirect } }) } if (error.response?.status === 500) { + // Stay on page; local error UI / toast is enough. Avoid global eject. toast.error(i18next.t('Internal Server Error!')) - router.navigate({ to: '/500' }) } } }, diff --git a/web/default/src/routes/oauth/$provider.tsx b/web/default/src/routes/oauth/$provider.tsx index b35bfdaffe60..1478a51c059f 100644 --- a/web/default/src/routes/oauth/$provider.tsx +++ b/web/default/src/routes/oauth/$provider.tsx @@ -16,21 +16,21 @@ along with this program. If not, see . For commercial licensing, please contact support@quantumnous.com */ +import { useEffect, useState } from 'react' +import type { AxiosRequestConfig } from 'axios' import { createFileRoute, useNavigate, useParams, useSearch, } from '@tanstack/react-router' -import type { AxiosRequestConfig } from 'axios' import i18next from 'i18next' -import { useEffect, useState } from 'react' import { toast } from 'sonner' - +import { useAuthStore, type AuthUser } from '@/stores/auth-store' +import { api, getSelf } from '@/lib/api' +import { safeRedirect } from '@/features/auth/lib/safe-redirect' import { OAuthCallbackScreen } from '@/features/auth/components/oauth-callback-screen' import { OAUTH_BIND_STORAGE_KEY } from '@/features/auth/constants' -import { api, getSelf } from '@/lib/api' -import { useAuthStore, type AuthUser } from '@/stores/auth-store' type OAuthRequestConfig = AxiosRequestConfig & { skipBusinessError?: boolean @@ -60,19 +60,18 @@ function OAuthCallback() { useEffect(() => { ;(async () => { const safeNavigate = (target: string) => { - navigate({ to: target as never, replace: true }) + const to = safeRedirect(target, '/dashboard') + navigate({ to: to as never, replace: true }) if (typeof window !== 'undefined') { setTimeout(() => { - const normalizedTarget = target.startsWith('/') - ? target - : `/${target}` const currentPath = window.location.pathname + window.location.search if ( - currentPath !== normalizedTarget && - currentPath !== `${normalizedTarget}/` + currentPath !== to && + currentPath !== `${to}/` ) { - window.location.replace(target) + // Only same-app relative paths survive safeRedirect. + window.location.replace(to) } }, 100) } @@ -144,7 +143,7 @@ function OAuthCallback() { } const redirectAfterLogin = (target?: string) => { - const to = target || search?.redirect || '/dashboard' + const to = safeRedirect(target || search?.redirect, '/dashboard') safeNavigate(to) toast.success(i18next.t('Signed in successfully!')) } diff --git a/web/default/src/stores/auth-store.ts b/web/default/src/stores/auth-store.ts index 49165c0e9099..3027d9de6295 100644 --- a/web/default/src/stores/auth-store.ts +++ b/web/default/src/stores/auth-store.ts @@ -97,6 +97,7 @@ export const useAuthStore = create()((set) => { set((state) => { if (typeof window !== 'undefined') { window.localStorage.removeItem('user') + window.localStorage.removeItem('uid') } return { ...state, diff --git a/web/default/tsconfig.app.json b/web/default/tsconfig.app.json index 79f70fc0860f..d106521727dc 100644 --- a/web/default/tsconfig.app.json +++ b/web/default/tsconfig.app.json @@ -29,5 +29,6 @@ "noFallthroughCasesInSwitch": true, "noUncheckedSideEffectImports": true }, - "include": ["src"] + "include": ["src"], + "exclude": ["src/**/*.test.ts", "src/**/*.test.tsx", "src/**/*.spec.ts"] } From 70ee9d9776bdc9f717fcba2a2e8c02cef39b4ebd Mon Sep 17 00:00:00 2001 From: yuanjia Date: Thu, 16 Jul 2026 10:10:51 +0800 Subject: [PATCH 07/40] perf(frontend): pricing/channels rendering and interaction improvements - Pricing columns memoization and per-page perf badge map - Filters<->URL debounced sync - Bulk actions skip-confirm dialog - Channel mutate drawer split (skip-auto-test field extraction) - model-name/mock-badge helpers and filters test --- web/default/src/features/channels/api.ts | 16 ++ .../components/data-table-bulk-actions.tsx | 156 ++++++++++----- .../dialogs/channel-test-dialog.tsx | 1 - .../drawers/channel-mutate-drawer.tsx | 4 +- .../drawers/channel-skip-auto-test-field.tsx | 62 ++++++ .../features/channels/lib/channel-actions.ts | 126 +++++------- .../channels/lib/channel-form-errors.ts | 1 + .../src/features/channels/lib/channel-form.ts | 5 + web/default/src/features/channels/types.ts | 1 + .../pricing/components/model-card-grid.tsx | 36 ++-- .../components/model-details-performance.tsx | 9 +- .../pricing/components/model-perf-badge.tsx | 26 ++- .../pricing/components/pricing-columns.tsx | 89 ++++++--- .../pricing/components/pricing-sidebar.tsx | 185 ++++++++++-------- .../pricing/components/pricing-table.tsx | 29 ++- .../pricing/components/pricing-toolbar.tsx | 24 +++ .../src/features/pricing/hooks/use-filters.ts | 100 +++++++++- .../pricing/hooks/use-perf-badge-map.ts | 87 ++++++++ web/default/src/features/pricing/index.tsx | 34 +++- .../src/features/pricing/lib/filters.test.ts | 89 +++++++++ .../src/features/pricing/lib/filters.ts | 72 ++++++- .../src/features/pricing/lib/mock-badge.ts | 84 ++++++++ .../src/features/pricing/lib/mock-stats.ts | 6 +- .../src/features/pricing/lib/model-name.ts | 62 ++++++ web/default/src/routes/(auth)/otp.tsx | 6 + web/default/src/routes/(auth)/sign-in.tsx | 8 +- .../src/routes/_authenticated/route.tsx | 36 ++-- web/default/src/routes/pricing/index.tsx | 11 +- 28 files changed, 1054 insertions(+), 311 deletions(-) create mode 100644 web/default/src/features/channels/components/drawers/channel-skip-auto-test-field.tsx create mode 100644 web/default/src/features/pricing/hooks/use-perf-badge-map.ts create mode 100644 web/default/src/features/pricing/lib/filters.test.ts create mode 100644 web/default/src/features/pricing/lib/mock-badge.ts create mode 100644 web/default/src/features/pricing/lib/model-name.ts diff --git a/web/default/src/features/channels/api.ts b/web/default/src/features/channels/api.ts index 09bad1f94b44..2e816304515c 100644 --- a/web/default/src/features/channels/api.ts +++ b/web/default/src/features/channels/api.ts @@ -203,6 +203,22 @@ export async function batchSetChannelTag( return res.data } +/** + * Batch toggle skip_auto_test on selected channels. + * Manual test still works; only automatic batch tests are gated. + */ +export async function batchSetChannelSkipAutoTest(data: { + ids: number[] + skip: boolean +}): Promise<{ success: boolean; message?: string; data?: number }> { + const res = await api.post( + '/api/channel/batch/skip_auto_test', + data, + channelActionConfig() + ) + return res.data +} + // ============================================================================ // Channel Operations // ============================================================================ diff --git a/web/default/src/features/channels/components/data-table-bulk-actions.tsx b/web/default/src/features/channels/components/data-table-bulk-actions.tsx index f35eb60faa15..a6818a1144fd 100644 --- a/web/default/src/features/channels/components/data-table-bulk-actions.tsx +++ b/web/default/src/features/channels/components/data-table-bulk-actions.tsx @@ -16,14 +16,11 @@ along with this program. If not, see . For commercial licensing, please contact support@quantumnous.com */ +import { useState } from 'react' import { useQueryClient } from '@tanstack/react-query' import { type Table } from '@tanstack/react-table' -import { Power, PowerOff, Tag, Trash2 } from 'lucide-react' -import { useState } from 'react' +import { PauseCircle, PlayCircle, Power, PowerOff, Tag, Trash2 } from 'lucide-react' import { useTranslation } from 'react-i18next' - -import { DataTableBulkActions as BulkActionsToolbar } from '@/components/data-table' -import { Dialog } from '@/components/dialog' import { Button } from '@/components/ui/button' import { Input } from '@/components/ui/input' import { Label } from '@/components/ui/label' @@ -32,19 +29,14 @@ import { TooltipContent, TooltipTrigger, } from '@/components/ui/tooltip' -import { - ADMIN_PERMISSION_ACTIONS, - ADMIN_PERMISSION_RESOURCES, - hasPermission, -} from '@/lib/admin-permissions' -import { cn } from '@/lib/utils' -import { useAuthStore } from '@/stores/auth-store' - +import { DataTableBulkActions as BulkActionsToolbar } from '@/components/data-table' +import { Dialog } from '@/components/dialog' import { handleBatchDelete, handleBatchDisable, handleBatchEnable, handleBatchSetTag, + handleBatchSkipAutoTest, } from '../lib' import type { Channel } from '../types' @@ -59,13 +51,9 @@ export function DataTableBulkActions({ const queryClient = useQueryClient() const [showTagDialog, setShowTagDialog] = useState(false) const [showDeleteConfirm, setShowDeleteConfirm] = useState(false) + const [showSkipConfirm, setShowSkipConfirm] = useState(false) + const [showJoinConfirm, setShowJoinConfirm] = useState(false) const [tagValue, setTagValue] = useState('') - const currentUser = useAuthStore((s) => s.auth.user) - const canEditSensitive = hasPermission( - currentUser, - ADMIN_PERMISSION_RESOURCES.CHANNEL, - ADMIN_PERMISSION_ACTIONS.SENSITIVE_WRITE - ) const selectedRows = table.getFilteredSelectedRowModel().rows const selectedIds = selectedRows.reduce((ids, row) => { @@ -91,7 +79,6 @@ export function DataTableBulkActions({ } const handleDeleteAll = () => { - if (!canEditSensitive) return handleBatchDelete(selectedIds, queryClient, () => { setShowDeleteConfirm(false) handleClearSelection() @@ -106,6 +93,20 @@ export function DataTableBulkActions({ }) } + const handleSkipAutoTest = () => { + handleBatchSkipAutoTest(selectedIds, true, queryClient, () => { + setShowSkipConfirm(false) + handleClearSelection() + }) + } + + const handleJoinAutoTest = () => { + handleBatchSkipAutoTest(selectedIds, false, queryClient, () => { + setShowJoinConfirm(false) + handleClearSelection() + }) + } + return ( <> @@ -130,6 +131,48 @@ export function DataTableBulkActions({ + + setShowSkipConfirm(true)} + className='size-8' + aria-label={t('Skip auto test for selected')} + title={t('Skip auto test for selected')} + /> + } + > + + {t('Skip auto test for selected')} + + +

{t('Skip auto test for selected')}

+ + + + + setShowJoinConfirm(true)} + className='size-8' + aria-label={t('Join auto test for selected')} + title={t('Join auto test for selected')} + /> + } + > + + {t('Join auto test for selected')} + + +

{t('Join auto test for selected')}

+ + + ({ - @@ -286,6 +310,50 @@ export function DataTableBulkActions({ > {' '} + + + {t('Skip auto test for selected')}: {selectedIds.length} + + } + contentHeight='auto' + footer={ + <> + + + + } + > + {' '} + + + + {t('Join auto test for selected')}: {selectedIds.length} + + } + contentHeight='auto' + footer={ + <> + + + + } + > + {' '} + ) } diff --git a/web/default/src/features/channels/components/dialogs/channel-test-dialog.tsx b/web/default/src/features/channels/components/dialogs/channel-test-dialog.tsx index a2d641d22c21..73bbcf2cab83 100644 --- a/web/default/src/features/channels/components/dialogs/channel-test-dialog.tsx +++ b/web/default/src/features/channels/components/dialogs/channel-test-dialog.tsx @@ -558,7 +558,6 @@ function ChannelTestDialogContent({ await handleTestChannel( currentRow.id, { - channelName: currentRow.name, testModel: model, endpointType: endpointType === 'auto' ? undefined : endpointType, stream: effectiveStreamTest || undefined, diff --git a/web/default/src/features/channels/components/drawers/channel-mutate-drawer.tsx b/web/default/src/features/channels/components/drawers/channel-mutate-drawer.tsx index 866599aa2179..1ae259b66532 100644 --- a/web/default/src/features/channels/components/drawers/channel-mutate-drawer.tsx +++ b/web/default/src/features/channels/components/drawers/channel-mutate-drawer.tsx @@ -182,6 +182,7 @@ import { ParamOverrideEditorDialog } from '../dialogs/param-override-editor-dial import { StatusCodeRiskDialog } from '../dialogs/status-code-risk-dialog' import { ModelMappingEditor } from '../model-mapping-editor' import { +import { ChannelSkipAutoTestField } from './channel-skip-auto-test-field' ChannelAdvancedSection, ChannelApiAccessSection, ChannelAuthSection, @@ -4640,7 +4641,8 @@ export function ChannelMutateDrawer({
)} - + + diff --git a/web/default/src/features/channels/components/drawers/channel-skip-auto-test-field.tsx b/web/default/src/features/channels/components/drawers/channel-skip-auto-test-field.tsx new file mode 100644 index 000000000000..d8e5a3b58769 --- /dev/null +++ b/web/default/src/features/channels/components/drawers/channel-skip-auto-test-field.tsx @@ -0,0 +1,62 @@ +/* +Copyright (C) 2023-2026 QuantumNous + +This program is free software: you can redistribute it and/or modify +it under the terms of the GNU Affero General Public License as +published by the Free Software Foundation, either version 3 of the +License, or (at your option) any later version. + +This program is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU Affero General Public License for more details. + +You should have received a copy of the GNU Affero General Public License +along with this program. If not, see . + +For commercial licensing, please contact support@quantumnous.com +*/ +import type { Control } from 'react-hook-form' +import { useTranslation } from 'react-i18next' +import { + FormControl, + FormDescription, + FormField, + FormItem, + FormLabel, +} from '@/components/ui/form' +import { Switch } from '@/components/ui/switch' +import type { ChannelFormValues } from '../../lib' + +/** Extracted from ChannelMutateDrawer to keep skip-auto-test editable in isolation. */ +export function ChannelSkipAutoTestField({ + control, +}: { + control: Control +}) { + const { t } = useTranslation() + return ( + ( + +
+ {t('Skip Auto Test')} + + {t( + 'Exclude this channel from automatic batch tests; manual test still works' + )} + +
+ + + + + )} + /> + ) +} diff --git a/web/default/src/features/channels/lib/channel-actions.ts b/web/default/src/features/channels/lib/channel-actions.ts index 7efee24d3da2..c62114c5feda 100644 --- a/web/default/src/features/channels/lib/channel-actions.ts +++ b/web/default/src/features/channels/lib/channel-actions.ts @@ -31,6 +31,7 @@ import { batchUpdateChannelStatus, batchDeleteChannels, batchSetChannelTag, + batchSetChannelSkipAutoTest, enableTagChannels, disableTagChannels, deleteDisabledChannels, @@ -41,7 +42,7 @@ import { updateChannelBalance, } from '../api' import { CHANNEL_STATUS, ERROR_MESSAGES, SUCCESS_MESSAGES } from '../constants' -import type { ChannelTestResponse, CopyChannelParams } from '../types' +import type { CopyChannelParams } from '../types' // ============================================================================ // Query Keys @@ -56,60 +57,6 @@ export const channelsQueryKeys = { detail: (id: number) => [...channelsQueryKeys.details(), id] as const, } -function getChannelTestResponseTime( - response: ChannelTestResponse -): number | undefined { - const responseTime = response.data?.response_time - if (typeof responseTime === 'number' && Number.isFinite(responseTime)) { - return responseTime - } - - if ( - typeof response.time === 'number' && - Number.isFinite(response.time) && - response.time > 0 - ) { - return Math.round(response.time * 1000) - } - - return undefined -} - -function formatChannelTestDuration(responseTime?: number): string | undefined { - if (responseTime === undefined) return undefined - - if (responseTime >= 1000) { - return `${(responseTime / 1000).toFixed(2)} s` - } - - return `${Math.max(1, Math.round(responseTime))} ms` -} - -function getChannelTestLabel(options?: { - channelName?: string - testModel?: string -}): string { - const channelName = options?.channelName?.trim() - const testModel = options?.testModel?.trim() - - if (channelName && testModel) { - return i18next.t('Channel {{name}} model {{model}}', { - name: channelName, - model: testModel, - }) - } - - if (channelName) { - return i18next.t('Channel {{name}}', { name: channelName }) - } - - if (testModel) { - return i18next.t('Model {{model}}', { model: testModel }) - } - - return i18next.t('Channel') -} - // ============================================================================ // Single Channel Actions // ============================================================================ @@ -271,7 +218,6 @@ export async function handleUpdateTagField( export async function handleTestChannel( id: number, options?: { - channelName?: string testModel?: string endpointType?: string stream?: boolean @@ -297,43 +243,23 @@ export async function handleTestChannel( try { const response = await testChannel(id, payload) - const responseTime = getChannelTestResponseTime(response) - const duration = formatChannelTestDuration(responseTime) - const target = getChannelTestLabel(options) if (response.success) { if (!options?.silent) { - toast.success( - i18next.t('{{target}} test succeeded', { target }), - duration - ? { - description: i18next.t('Response time: {{duration}}', { - duration, - }), - } - : undefined - ) + toast.success(i18next.t(SUCCESS_MESSAGES.TESTED)) } - onTestComplete?.(true, responseTime) + onTestComplete?.(true, response.time) } else { - const errorMsg = response.message || i18next.t(ERROR_MESSAGES.TEST_FAILED) if (!options?.silent) { - toast.error(i18next.t('{{target}} test failed', { target }), { - description: response.error_code - ? `${errorMsg} (${response.error_code})` - : errorMsg, - }) + toast.error(response.message || i18next.t(ERROR_MESSAGES.TEST_FAILED)) } - onTestComplete?.(false, responseTime, errorMsg, response.error_code) + onTestComplete?.(false, undefined, response.message, response.error_code) } } catch (_error: unknown) { const err = _error as { response?: { data?: { message?: string } } } const errorMsg = err?.response?.data?.message || i18next.t(ERROR_MESSAGES.TEST_FAILED) - const target = getChannelTestLabel(options) if (!options?.silent) { - toast.error(i18next.t('{{target}} test failed', { target }), { - description: errorMsg, - }) + toast.error(errorMsg) } onTestComplete?.(false, undefined, errorMsg) } @@ -541,6 +467,44 @@ export async function handleBatchSetTag( } } +/** + * Batch set skip_auto_test on selected channels. + */ +export async function handleBatchSkipAutoTest( + ids: number[], + skip: boolean, + queryClient?: QueryClient, + onSuccess?: () => void +): Promise { + if (ids.length === 0) { + toast.error(i18next.t('No channels selected')) + return + } + + try { + const response = await batchSetChannelSkipAutoTest({ ids, skip }) + if (response.success) { + toast.success( + skip + ? i18next.t('{{count}} channel(s) will skip auto test', { + count: response.data || ids.length, + }) + : i18next.t('{{count}} channel(s) will join auto test', { + count: response.data || ids.length, + }) + ) + queryClient?.invalidateQueries({ queryKey: channelsQueryKeys.lists() }) + onSuccess?.() + } else { + toast.error( + response.message || i18next.t('Failed to update skip auto test') + ) + } + } catch (_error) { + toast.error(i18next.t('Failed to update skip auto test')) + } +} + // ============================================================================ // Tag-Based Actions // ============================================================================ diff --git a/web/default/src/features/channels/lib/channel-form-errors.ts b/web/default/src/features/channels/lib/channel-form-errors.ts index f6873ebc432c..250f4363a936 100644 --- a/web/default/src/features/channels/lib/channel-form-errors.ts +++ b/web/default/src/features/channels/lib/channel-form-errors.ts @@ -38,6 +38,7 @@ const ADVANCED_SETTINGS_FIELDS = new Set>([ 'force_format', 'thinking_to_content', 'pass_through_body_enabled', + 'skip_auto_test', 'proxy', 'system_prompt', 'system_prompt_override', diff --git a/web/default/src/features/channels/lib/channel-form.ts b/web/default/src/features/channels/lib/channel-form.ts index 772b6be411ac..b9f2e12914f4 100644 --- a/web/default/src/features/channels/lib/channel-form.ts +++ b/web/default/src/features/channels/lib/channel-form.ts @@ -192,6 +192,7 @@ export const channelFormSchema = z pass_through_body_enabled: z.boolean().optional(), system_prompt: z.string().optional(), system_prompt_override: z.boolean().optional(), + skip_auto_test: z.boolean().optional(), // Type-specific settings (stored in settings JSON) is_enterprise_account: z.boolean().optional(), // OpenRouter specific vertex_key_type: z.enum(['json', 'api_key']).optional(), // Vertex AI specific @@ -342,6 +343,7 @@ export const CHANNEL_FORM_DEFAULT_VALUES: ChannelFormValues = { pass_through_body_enabled: false, system_prompt: '', system_prompt_override: false, + skip_auto_test: false, // Type-specific settings is_enterprise_account: false, vertex_key_type: 'json', @@ -380,6 +382,7 @@ export function transformChannelToFormDefaults( pass_through_body_enabled: false, system_prompt: '', system_prompt_override: false, + skip_auto_test: false, } if (channel.setting) { @@ -392,6 +395,7 @@ export function transformChannelToFormDefaults( pass_through_body_enabled: parsed.pass_through_body_enabled || false, system_prompt: parsed.system_prompt || '', system_prompt_override: parsed.system_prompt_override || false, + skip_auto_test: parsed.skip_auto_test || false, } } catch (error) { // eslint-disable-next-line no-console @@ -509,6 +513,7 @@ function buildSettingJSON(formData: ChannelFormValues): string { pass_through_body_enabled: formData.pass_through_body_enabled || false, system_prompt: formData.system_prompt || '', system_prompt_override: formData.system_prompt_override || false, + skip_auto_test: formData.skip_auto_test || false, } return JSON.stringify(settingObj) } diff --git a/web/default/src/features/channels/types.ts b/web/default/src/features/channels/types.ts index 96bd5b2cfd12..7229d27e0f9f 100644 --- a/web/default/src/features/channels/types.ts +++ b/web/default/src/features/channels/types.ts @@ -86,6 +86,7 @@ export interface ChannelSettings { pass_through_body_enabled?: boolean system_prompt?: string system_prompt_override?: boolean + skip_auto_test?: boolean } export interface ChannelOtherSettings { diff --git a/web/default/src/features/pricing/components/model-card-grid.tsx b/web/default/src/features/pricing/components/model-card-grid.tsx index d92593ab716a..2d6f0dab88c3 100644 --- a/web/default/src/features/pricing/components/model-card-grid.tsx +++ b/web/default/src/features/pricing/components/model-card-grid.tsx @@ -16,18 +16,15 @@ along with this program. If not, see . For commercial licensing, please contact support@quantumnous.com */ -import { useQuery } from '@tanstack/react-query' +import { useEffect, useMemo, useState } from 'react' import { ChevronLeft, ChevronRight } from 'lucide-react' -import { useMemo, useState } from 'react' import { useTranslation } from 'react-i18next' - import { Button } from '@/components/ui/button' -import { getPerfMetricsSummary } from '@/features/performance-metrics/api' - +import { usePerfBadgeMap } from '../hooks/use-perf-badge-map' +import { normalizeModelName } from '../lib/model-name' import { DEFAULT_PRICING_PAGE_SIZE, DEFAULT_TOKEN_UNIT } from '../constants' import type { PricingModel, TokenUnit } from '../types' import { ModelCard } from './model-card' -import type { ModelPerfBadgeData } from './model-perf-badge' export interface ModelCardGridProps { models: PricingModel[] @@ -36,7 +33,9 @@ export interface ModelCardGridProps { usdExchangeRate?: number tokenUnit?: TokenUnit showRechargePrice?: boolean + /** When true, only show badges backed by real probe/relay samples (no mock fill). */ selectedGroup?: string + liveMetricsOnly?: boolean } export function ModelCardGrid(props: ModelCardGridProps) { @@ -44,28 +43,24 @@ export function ModelCardGrid(props: ModelCardGridProps) { const [page, setPage] = useState(1) const pageSize = DEFAULT_PRICING_PAGE_SIZE const tokenUnit = props.tokenUnit ?? DEFAULT_TOKEN_UNIT + const liveMetricsOnly = props.liveMetricsOnly === true const totalPages = Math.max(1, Math.ceil(props.models.length / pageSize)) const currentPage = Math.min(page, totalPages) - const perfQuery = useQuery({ - queryKey: ['perf-metrics-summary', 24], - queryFn: () => getPerfMetricsSummary(24), - staleTime: 60 * 1000, - retry: false, - }) + useEffect(() => { + setPage(1) + }, [props.models]) const pagedModels = useMemo(() => { const start = (currentPage - 1) * pageSize return props.models.slice(start, start + pageSize) }, [currentPage, pageSize, props.models]) - const perfMap = useMemo(() => { - const map = new Map() - for (const model of perfQuery.data?.data?.models ?? []) { - map.set(model.model_name, model) - } - return map - }, [perfQuery.data]) + // Only badge the visible page — avoids O(N) mock fill on large catalogs. + const { perfMap } = usePerfBadgeMap({ + models: pagedModels, + liveMetricsOnly, + }) if (props.models.length === 0) { return null @@ -82,8 +77,7 @@ export function ModelCardGrid(props: ModelCardGridProps) { priceRate={props.priceRate} usdExchangeRate={props.usdExchangeRate} showRechargePrice={props.showRechargePrice} - selectedGroup={props.selectedGroup} - perf={perfMap.get(model.model_name || '')} + perf={perfMap.get(normalizeModelName(model.model_name))} onClick={() => props.onModelClick(model.model_name || '')} /> ))} diff --git a/web/default/src/features/pricing/components/model-details-performance.tsx b/web/default/src/features/pricing/components/model-details-performance.tsx index 0be5d040d781..2cc321bebf27 100644 --- a/web/default/src/features/pricing/components/model-details-performance.tsx +++ b/web/default/src/features/pricing/components/model-details-performance.tsx @@ -37,6 +37,7 @@ import type { PerformanceGroup } from '@/features/performance-metrics/types' import { cn } from '@/lib/utils' import { type UptimeDayPoint } from '../lib/mock-stats' +import { normalizeModelName } from '../lib/model-name' import type { PricingModel } from '../types' import { LatencyTrendChart, UptimeTrendChart } from './model-details-charts' import { UptimeSparkline } from './model-details-uptime-sparkline' @@ -163,10 +164,14 @@ function average( export function ModelDetailsPerformance(props: { model: PricingModel }) { const { t } = useTranslation() + // Backend Query also normalizes; pass folded name so cache keys collapse + // Deepseek-V4-Flash / deepseek-v4-flash onto one entry. + const modelKey = normalizeModelName(props.model.model_name) const metricsQuery = useQuery({ - queryKey: ['perf-metrics', props.model.model_name], - queryFn: () => getPerfMetrics(props.model.model_name, 24), + queryKey: ['perf-metrics', modelKey], + queryFn: () => getPerfMetrics(modelKey || props.model.model_name, 24), staleTime: 60 * 1000, + enabled: Boolean(modelKey || props.model.model_name), }) const groups = useMemo( () => metricsQuery.data?.data.groups ?? [], diff --git a/web/default/src/features/pricing/components/model-perf-badge.tsx b/web/default/src/features/pricing/components/model-perf-badge.tsx index f214367a58d1..9e02ee86ea66 100644 --- a/web/default/src/features/pricing/components/model-perf-badge.tsx +++ b/web/default/src/features/pricing/components/model-perf-badge.tsx @@ -27,6 +27,8 @@ export type ModelPerfBadgeData = { success_rate: number avg_tps: number recent_success_rates?: number[] + /** True when values come from cold-start mock, not real probe/relay samples. */ + is_mock?: boolean } export interface ModelPerfBadgeProps extends React.HTMLAttributes { @@ -60,6 +62,7 @@ export const ModelPerfBadge = memo(function ModelPerfBadge( } const { avg_latency_ms, avg_tps, success_rate } = props.perf + const isMock = Boolean(props.perf.is_mock) const recentRates = props.perf.recent_success_rates?.filter((rate) => Number.isFinite(rate)) ?? @@ -71,22 +74,33 @@ export const ModelPerfBadge = memo(function ModelPerfBadge( ...statusRates, ].slice(-3) + const sourceHint = isMock + ? t('Estimated (no live traffic yet)') + : t('Live probe / relay metrics') + return (
-
+
{t('Latency short')} + {isMock && ( + + ~ + + )}
{formatCompactLatency(avg_latency_ms)}
-
+
{t('Throughput short')}
@@ -95,11 +109,11 @@ export const ModelPerfBadge = memo(function ModelPerfBadge(
- {t('Status short')} + {isMock ? t('Est. short') : t('Status short')}
{statusBars.map((rate, index) => ( @@ -114,7 +128,9 @@ export const ModelPerfBadge = memo(function ModelPerfBadge( ? index === 0 ? 'bg-muted-foreground/10' : 'bg-muted-foreground/15' - : getSuccessRateDotClass(rate) + : isMock + ? 'bg-muted-foreground/35' + : getSuccessRateDotClass(rate) )} /> ))} diff --git a/web/default/src/features/pricing/components/pricing-columns.tsx b/web/default/src/features/pricing/components/pricing-columns.tsx index 3d19eda2130c..9045375ce410 100644 --- a/web/default/src/features/pricing/components/pricing-columns.tsx +++ b/web/default/src/features/pricing/components/pricing-columns.tsx @@ -16,9 +16,10 @@ along with this program. If not, see . For commercial licensing, please contact support@quantumnous.com */ -import type { ColumnDef } from '@tanstack/react-table' +import { useMemo } from 'react' +import { type ColumnDef } from '@tanstack/react-table' import { useTranslation } from 'react-i18next' - +import { getLobeIcon } from '@/lib/lobe-icon' import { BadgeCell, BadgeListCell, @@ -26,22 +27,24 @@ import { } from '@/components/data-table' import { GroupBadge } from '@/components/group-badge' import { StatusBadge } from '@/components/status-badge' -import { getLobeIcon } from '@/lib/lobe-icon' - -import { DEFAULT_TOKEN_UNIT } from '../constants' +import { DEFAULT_TOKEN_UNIT, QUOTA_TYPE_VALUES } from '../constants' import { getDynamicDisplayGroupRatio, getDynamicPricingSummary, } from '../lib/dynamic-price' import { parseTags } from '../lib/filters' import { isTokenBasedModel } from '../lib/model-helpers' +import { normalizeModelName } from '../lib/model-name' import { formatPrice, formatRequestPrice, stripTrailingZeros, } from '../lib/price' import type { PricingModel, TokenUnit } from '../types' -import { ModelBillingModeBadge } from './model-billing-mode-badge' +import { + ModelPerfBadge, + type ModelPerfBadgeData, +} from './model-perf-badge' // ---------------------------------------------------------------------------- // Pricing Table Columns @@ -52,7 +55,8 @@ export interface PricingColumnsOptions { priceRate?: number usdExchangeRate?: number showRechargePrice?: boolean - selectedGroup?: string + /** Shared with card grid; case-folded model keys. */ + perfMap?: Map } export function usePricingColumns( @@ -64,12 +68,12 @@ export function usePricingColumns( priceRate = 1, usdExchangeRate = 1, showRechargePrice = false, - selectedGroup, + perfMap, } = options const tokenUnitLabel = tokenUnit === 'K' ? '1K' : '1M' - return [ + return useMemo(() => [ // Model column { accessorKey: 'model_name', @@ -94,14 +98,41 @@ export function usePricingColumns( minSize: 200, }, + // Health / live metrics (same badge as card grid; respects liveMetricsOnly via map) + { + id: 'health', + meta: { label: t('Health') }, + header: ({ column }) => ( + + ), + cell: ({ row }) => { + const key = normalizeModelName(row.original.model_name) + const perf = key ? perfMap?.get(key) : undefined + if (!perf) { + return — + } + return + }, + size: 160, + enableSorting: false, + }, + // Type column { accessorKey: 'quota_type', header: t('Type'), - cell: ({ row }) => ( - - ), - size: 110, + cell: ({ row }) => { + const isTokenBased = row.original.quota_type === QUOTA_TYPE_VALUES.TOKEN + return ( + + ) + }, + size: 80, enableSorting: false, }, @@ -119,10 +150,7 @@ export function usePricingColumns( showRechargePrice, priceRate, usdExchangeRate, - groupRatioMultiplier: getDynamicDisplayGroupRatio( - model, - selectedGroup - ), + groupRatioMultiplier: getDynamicDisplayGroupRatio(model), }) if (dynamicSummary) { @@ -184,8 +212,7 @@ export function usePricingColumns( tokenUnit, showRechargePrice, priceRate, - usdExchangeRate, - selectedGroup + usdExchangeRate ) ) const outputPrice = stripTrailingZeros( @@ -195,8 +222,7 @@ export function usePricingColumns( tokenUnit, showRechargePrice, priceRate, - usdExchangeRate, - selectedGroup + usdExchangeRate ) ) @@ -219,8 +245,7 @@ export function usePricingColumns( model, showRechargePrice, priceRate, - usdExchangeRate, - selectedGroup + usdExchangeRate ) ) @@ -248,10 +273,7 @@ export function usePricingColumns( showRechargePrice, priceRate, usdExchangeRate, - groupRatioMultiplier: getDynamicDisplayGroupRatio( - model, - selectedGroup - ), + groupRatioMultiplier: getDynamicDisplayGroupRatio(model), }) if (dynamicSummary) { @@ -295,8 +317,7 @@ export function usePricingColumns( tokenUnit, showRechargePrice, priceRate, - usdExchangeRate, - selectedGroup + usdExchangeRate ) ) @@ -409,5 +430,13 @@ export function usePricingColumns( size: 130, enableSorting: false, }, - ] + ], [ + t, + tokenUnit, + tokenUnitLabel, + priceRate, + usdExchangeRate, + showRechargePrice, + perfMap, + ]) } diff --git a/web/default/src/features/pricing/components/pricing-sidebar.tsx b/web/default/src/features/pricing/components/pricing-sidebar.tsx index e07d095b2cf2..432aa5476090 100644 --- a/web/default/src/features/pricing/components/pricing-sidebar.tsx +++ b/web/default/src/features/pricing/components/pricing-sidebar.tsx @@ -16,10 +16,11 @@ along with this program. If not, see . For commercial licensing, please contact support@quantumnous.com */ +import { useMemo, type ReactNode } from 'react' import { ChevronDown, RotateCcw } from 'lucide-react' -import type { ReactNode } from 'react' import { useTranslation } from 'react-i18next' - +import { getLobeIcon } from '@/lib/lobe-icon' +import { cn } from '@/lib/utils' import { Badge } from '@/components/ui/badge' import { Button } from '@/components/ui/button' import { @@ -27,9 +28,6 @@ import { CollapsibleContent, CollapsibleTrigger, } from '@/components/ui/collapsible' -import { getLobeIcon } from '@/lib/lobe-icon' -import { cn } from '@/lib/utils' - import { ENDPOINT_TYPES, FILTER_ALL, @@ -115,7 +113,7 @@ function FilterChip(props: { {(props.option.suffix || props.option.count != null) && ( ({ - value: vendor.name, - label: vendor.name, - count: countBy( - props.models, - (model) => model.vendor_name === vendor.name - ), - icon: vendor.icon ? getLobeIcon(vendor.icon, 14) : undefined, - })) - .filter((vendor) => vendor.count > 0), - ] + const quotaTypeLabels = useMemo(() => getQuotaTypeLabels(t), [t]) + const endpointTypeLabels = useMemo(() => getEndpointTypeLabels(t), [t]) - const groupOptions: FilterOption[] = [ - { - value: FILTER_ALL, - label: t('All Groups'), - }, - ...props.groups.map((group) => ({ - value: group, - label: group, - suffix: formatGroupRatio(props.groupRatios?.[group]), - })), - ] + const vendorOptions = useMemo( + () => [ + { + value: FILTER_ALL, + label: t('All Vendors'), + count: props.models.length, + }, + ...props.vendors + .map((vendor) => ({ + value: vendor.name, + label: vendor.name, + count: countBy( + props.models, + (model) => model.vendor_name === vendor.name + ), + icon: vendor.icon ? getLobeIcon(vendor.icon, 14) : undefined, + })) + .filter((vendor) => vendor.count > 0), + ], + [props.models, props.vendors, t] + ) - const quotaOptions: FilterOption[] = [ - { - value: QUOTA_TYPES.ALL, - label: quotaTypeLabels[QUOTA_TYPES.ALL], - count: props.models.length, - }, - { - value: QUOTA_TYPES.TOKEN, - label: quotaTypeLabels[QUOTA_TYPES.TOKEN], - count: countBy(props.models, (model) => model.quota_type === 0), - }, - { - value: QUOTA_TYPES.REQUEST, - label: quotaTypeLabels[QUOTA_TYPES.REQUEST], - count: countBy(props.models, (model) => model.quota_type === 1), - }, - ] + const groupOptions = useMemo( + () => [ + { + value: FILTER_ALL, + label: t('All Groups'), + }, + ...props.groups.map((group) => ({ + value: group, + label: group, + suffix: formatGroupRatio(props.groupRatios?.[group]), + })), + ], + [props.groupRatios, props.groups, t] + ) - const tagOptions: FilterOption[] = [ - { - value: FILTER_ALL, - label: t('All Tags'), - count: props.models.length, - }, - ...props.tags.map((tag) => ({ - value: tag, - label: tag, - count: countBy(props.models, (model) => - parseTags(model.tags) - .map((item) => item.toLowerCase()) - .includes(tag.toLowerCase()) - ), - })), - ] + const quotaOptions = useMemo( + () => [ + { + value: QUOTA_TYPES.ALL, + label: quotaTypeLabels[QUOTA_TYPES.ALL], + count: props.models.length, + }, + { + value: QUOTA_TYPES.TOKEN, + label: quotaTypeLabels[QUOTA_TYPES.TOKEN], + count: countBy(props.models, (model) => model.quota_type === 0), + }, + { + value: QUOTA_TYPES.REQUEST, + label: quotaTypeLabels[QUOTA_TYPES.REQUEST], + count: countBy(props.models, (model) => model.quota_type === 1), + }, + ], + [props.models, quotaTypeLabels] + ) - const endpointOptions: FilterOption[] = [ - { - value: ENDPOINT_TYPES.ALL, - label: endpointTypeLabels[ENDPOINT_TYPES.ALL], - count: props.models.length, - }, - ...Object.entries(endpointTypeLabels) - .filter(([value]) => value !== ENDPOINT_TYPES.ALL) - .map(([value, label]) => ({ - value, - label, - count: countBy( - props.models, - (model) => model.supported_endpoint_types?.includes(value) ?? false + const tagOptions = useMemo( + () => [ + { + value: FILTER_ALL, + label: t('All Tags'), + count: props.models.length, + }, + ...props.tags.map((tag) => ({ + value: tag, + label: tag, + count: countBy(props.models, (model) => + parseTags(model.tags) + .map((item) => item.toLowerCase()) + .includes(tag.toLowerCase()) ), })), - ] + ], + [props.models, props.tags, t] + ) + + const endpointOptions = useMemo( + () => [ + { + value: ENDPOINT_TYPES.ALL, + label: endpointTypeLabels[ENDPOINT_TYPES.ALL], + count: props.models.length, + }, + ...Object.entries(endpointTypeLabels) + .filter(([value]) => value !== ENDPOINT_TYPES.ALL) + .map(([value, label]) => ({ + value, + label, + count: countBy( + props.models, + (model) => model.supported_endpoint_types?.includes(value) ?? false + ), + })), + ], + [endpointTypeLabels, props.models] + ) return (