342 lines
19 KiB
Go
342 lines
19 KiB
Go
package service
|
|
|
|
import (
|
|
"fmt"
|
|
"os"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/QuantumNous/new-api/common"
|
|
"github.com/QuantumNous/new-api/constant"
|
|
"github.com/QuantumNous/new-api/model"
|
|
fingerprint "github.com/QuantumNous/new-api/pkg/model_fingerprint"
|
|
"github.com/QuantumNous/new-api/setting/operation_setting"
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/glebarez/sqlite"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
"gorm.io/driver/mysql"
|
|
"gorm.io/driver/postgres"
|
|
"gorm.io/gorm"
|
|
"gorm.io/gorm/schema"
|
|
"net/http/httptest"
|
|
)
|
|
|
|
func TestGroupDegradationDatabaseMatrix(t *testing.T) {
|
|
for _, dialect := range []string{"sqlite", "mysql", "postgres"} {
|
|
t.Run(dialect, func(t *testing.T) {
|
|
var driver gorm.Dialector
|
|
switch dialect {
|
|
case "sqlite":
|
|
driver = sqlite.Open(":memory:")
|
|
case "mysql":
|
|
if os.Getenv("TEST_MYSQL_DSN") == "" {
|
|
t.Skip("TEST_MYSQL_DSN unset")
|
|
}
|
|
driver = mysql.Open(os.Getenv("TEST_MYSQL_DSN"))
|
|
case "postgres":
|
|
if os.Getenv("TEST_POSTGRES_DSN") == "" {
|
|
t.Skip("TEST_POSTGRES_DSN unset")
|
|
}
|
|
driver = postgres.Open(os.Getenv("TEST_POSTGRES_DSN"))
|
|
}
|
|
db, err := gorm.Open(driver, &gorm.Config{NamingStrategy: schema.NamingStrategy{TablePrefix: fmt.Sprintf("degradation_%d_", time.Now().UnixNano())}})
|
|
require.NoError(t, err)
|
|
oldDB := model.DB
|
|
oldType, oldLogType := common.MainDatabaseType(), common.LogDatabaseType()
|
|
model.DB = db
|
|
common.SetDatabaseTypes(common.DatabaseType(dialect), oldLogType)
|
|
rows := []any{&model.Option{}, &model.Channel{}, &model.SystemTask{}, &model.DetectionSchedule{}, &model.DetectionRun{}, &model.DetectionCall{}, &model.Token{}}
|
|
require.NoError(t, db.AutoMigrate(rows...))
|
|
t.Cleanup(func() {
|
|
for _, row := range rows {
|
|
require.NoError(t, db.Migrator().DropTable(row))
|
|
}
|
|
model.DB = oldDB
|
|
common.SetDatabaseTypes(oldType, oldLogType)
|
|
groupDegradationSnapshot.updated = time.Time{}
|
|
groupDegradationSnapshot.latest = nil
|
|
groupDegradationSnapshot.vision = nil
|
|
groupDegradationSnapshot.ciallo = nil
|
|
conn, err := db.DB()
|
|
require.NoError(t, err)
|
|
require.NoError(t, conn.Close())
|
|
})
|
|
var version string
|
|
query := "select version()"
|
|
if dialect == "sqlite" {
|
|
query = "select sqlite_version()"
|
|
}
|
|
require.NoError(t, db.Raw(query).Scan(&version).Error)
|
|
t.Log("database version:", version)
|
|
now := time.Now().Unix()
|
|
require.NoError(t, db.Create(&model.Channel{Id: 1, Group: "default", Models: fingerprint.Model, Status: 1, Key: "test-only"}).Error)
|
|
body, err := common.Marshal(AstraSamplingState{Rounds: []AstraSamplingRound{{ID: "mismatch", ChannelID: 1, StartedAt: now - 30, FinishedAt: now - 10, Result: fingerprint.Result{Engine: fingerprint.Version, Status: "mismatch", Closest: "gpt-5.5", ValidSamples: 2}}}})
|
|
require.NoError(t, err)
|
|
require.NoError(t, db.Create(&model.SystemTask{TaskID: "degradation-test", Type: AstraSamplingTaskType, Status: model.SystemTaskStatusSucceeded, CreatedAt: now - 30, State: string(body)}).Error)
|
|
policy := &model.TokenRoutingPolicy{Version: 1, Mode: "custom", SkipDegradedGroups: true, Groups: []model.TokenGroupEntry{{Group: "default", Priority: 0}}}
|
|
encoded, err := common.Marshal(policy)
|
|
require.NoError(t, err)
|
|
token := model.Token{UserId: 1, Key: "degradation-test-key", RoutingPolicy: string(encoded)}
|
|
require.NoError(t, db.Create(&token).Error)
|
|
require.NoError(t, db.First(&token, token.Id).Error)
|
|
stored, err := token.GetRoutingPolicy()
|
|
require.NoError(t, err)
|
|
require.True(t, stored.SkipDegradedGroups)
|
|
ctx := func() *gin.Context {
|
|
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
|
c.Request = httptest.NewRequest("POST", "/v1/responses", nil)
|
|
c.Set(model.TokenRoutingContextKey, stored)
|
|
return c
|
|
}
|
|
groupDegradationSnapshot.updated = time.Time{}
|
|
assert.True(t, requestGroupDegraded(ctx(), fingerprint.Model, "default"))
|
|
smart := ctx()
|
|
smartPolicy := &model.TokenRoutingPolicy{Version: 1, Mode: "smart"}
|
|
smart.Set(model.TokenRoutingContextKey, smartPolicy)
|
|
assert.True(t, requestGroupDegraded(smart, fingerprint.Model, "default"))
|
|
smartPolicy.AllowDegraded = true
|
|
assert.False(t, requestGroupDegraded(smart, fingerprint.Model, "default"))
|
|
smartPolicy.AllowDegraded = false
|
|
assert.False(t, requestGroupDegraded(smart, "unmeasured-model", "default"))
|
|
manualResult, err := common.Marshal(fingerprint.Result{Engine: fingerprint.Version, Status: "match", Closest: fingerprint.Model, ValidSamples: 2})
|
|
require.NoError(t, err)
|
|
manual := model.DetectionRun{ID: "manual-recovery", RequestKey: "manual-recovery", ChannelID: 1, Kind: "fingerprint", UserID: 1, CreatedAt: now, Settled: true, Status: "succeeded", Result: string(manualResult)}
|
|
require.NoError(t, db.Create(&manual).Error)
|
|
groupDegradationSnapshot.updated = time.Time{}
|
|
assert.False(t, requestGroupDegraded(ctx(), fingerprint.Model, "default"))
|
|
pelican := model.DetectionRun{ID: "pelican-degraded", RequestKey: "pelican-degraded", ChannelID: 1, Kind: "pelican", CreatedAt: now, Settled: true, Status: "succeeded", Quality: "degraded"}
|
|
require.NoError(t, db.Create(&pelican).Error)
|
|
groupDegradationSnapshot.updated = time.Time{}
|
|
assert.False(t, requestGroupDegraded(ctx(), fingerprint.Model, "default"))
|
|
require.NoError(t, db.Model(&pelican).Update("quality", "normal").Error)
|
|
groupDegradationSnapshot.updated = time.Time{}
|
|
assert.False(t, requestGroupDegraded(ctx(), fingerprint.Model, "default"))
|
|
|
|
vision := model.DetectionRun{ID: "vision-degraded", RequestKey: "vision-degraded", ChannelID: 1, Kind: "vision", CreatedAt: now - 5, Settled: true, Status: "succeeded", Quality: "degraded", Result: `{"closest":"石头"}`}
|
|
require.NoError(t, db.Create(&vision).Error)
|
|
InvalidateGroupQualityCache()
|
|
assert.True(t, requestGroupDegraded(ctx(), fingerprint.Model, "default"), "fingerprint recovery must not clear a failed vision test")
|
|
recent, err := RecentDegradedGroups(ctx().Request.Context())
|
|
require.NoError(t, err)
|
|
assert.True(t, recent["default"])
|
|
failedVision := model.DetectionRun{ID: "vision-error", RequestKey: "vision-error", ChannelID: 1, Kind: "vision", CreatedAt: now - 1, Settled: true, Status: "failed", Quality: "normal"}
|
|
require.NoError(t, db.Create(&failedVision).Error)
|
|
InvalidateGroupQualityCache()
|
|
assert.True(t, requestGroupDegraded(ctx(), fingerprint.Model, "default"), "transport errors cannot clear the latest valid vision result")
|
|
require.NoError(t, db.Model(&failedVision).Update("quality", "degraded").Error)
|
|
require.NoError(t, db.Model(&vision).Updates(map[string]any{"quality": "normal", "result": `{"closest":"这是音箱"}`}).Error)
|
|
InvalidateGroupQualityCache()
|
|
assert.False(t, requestGroupDegraded(ctx(), fingerprint.Model, "default"))
|
|
recent, err = RecentDegradedGroups(ctx().Request.Context())
|
|
require.NoError(t, err)
|
|
assert.False(t, recent["default"])
|
|
require.NoError(t, db.Model(&vision).Updates(map[string]any{"quality": "degraded", "created_at": now - 8*3600, "result": `{"closest":"石头"}`}).Error)
|
|
InvalidateGroupQualityCache()
|
|
assert.True(t, requestGroupDegraded(ctx(), fingerprint.Model, "default"), "routing keeps the latest valid result")
|
|
recent, err = RecentDegradedGroups(ctx().Request.Context())
|
|
require.NoError(t, err)
|
|
assert.False(t, recent["default"], "recent tag expires after the existing six-hour window")
|
|
legacyAnswer := `{"closest":"我这边没收到可查看的图片或文件，暂时无法判断是什么。请重新上传图片"}`
|
|
require.NoError(t, db.Model(&vision).Updates(map[string]any{"result": legacyAnswer, "status": "failed", "error": "vision_keyword_missing"}).Error)
|
|
InvalidateGroupQualityCache()
|
|
unavailable, err := VisionUnavailableGroups(ctx().Request.Context())
|
|
require.NoError(t, err)
|
|
assert.True(t, unavailable["default"], "legacy answers remain effective beyond the recent window and later transport failures")
|
|
assert.False(t, requestGroupDegraded(ctx(), fingerprint.Model, "default"))
|
|
recent, err = RecentDegradedGroups(ctx().Request.Context())
|
|
require.NoError(t, err)
|
|
assert.False(t, recent["default"])
|
|
require.NoError(t, db.Model(&vision).Update("created_at", now-2).Error)
|
|
summaries, err := GroupDetectionSummaries(ctx().Request.Context(), time.Unix(now, 0), nil, nil)
|
|
require.NoError(t, err)
|
|
summary := summaries["default"].Vision
|
|
require.NotNil(t, summary)
|
|
assert.Equal(t, 1, summary.Incomplete)
|
|
assert.Equal(t, "incomplete", summary.LatestValid.Quality)
|
|
var unchanged model.DetectionRun
|
|
require.NoError(t, db.First(&unchanged, "id = ?", vision.ID).Error)
|
|
assert.Equal(t, "degraded", unchanged.Quality, "reading history must not rewrite stored results")
|
|
ciallo := model.DetectionRun{ID: "ciallo-degraded", RequestKey: "ciallo-degraded", ChannelID: 1, Kind: "ciallo", CreatedAt: now, Settled: true, Status: "succeeded", Quality: "degraded"}
|
|
require.NoError(t, db.Create(&ciallo).Error)
|
|
InvalidateGroupQualityCache()
|
|
assert.True(t, requestGroupDegraded(ctx(), fingerprint.Model, "default"), "incomplete vision must not clear Ciallo degradation")
|
|
require.NoError(t, db.Model(&vision).Updates(map[string]any{"status": "succeeded", "result": `{"closest":"这是音箱"}`, "created_at": now}).Error)
|
|
InvalidateGroupQualityCache()
|
|
unavailable, err = VisionUnavailableGroups(ctx().Request.Context())
|
|
require.NoError(t, err)
|
|
assert.False(t, unavailable["default"])
|
|
assert.True(t, requestGroupDegraded(ctx(), fingerprint.Model, "default"), "vision recovery must not clear Ciallo degradation")
|
|
require.NoError(t, db.Model(&ciallo).Update("quality", "normal").Error)
|
|
InvalidateGroupQualityCache()
|
|
assert.False(t, requestGroupDegraded(ctx(), fingerprint.Model, "default"))
|
|
require.NoError(t, db.Delete(&manual).Error)
|
|
require.NoError(t, model.SaveStatusPageGroup(ctx().Request.Context(), "default", nil))
|
|
groupDegradationSnapshot.updated = time.Time{}
|
|
assert.False(t, requestGroupDegraded(ctx(), fingerprint.Model, "default"))
|
|
unavailable, err = VisionUnavailableGroups(ctx().Request.Context())
|
|
require.NoError(t, err)
|
|
assert.False(t, unavailable["default"], "administrator reset removes the vision conclusion")
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestGroupDegradationRetainsFullHistoryForNotifications(t *testing.T) {
|
|
now := time.Now()
|
|
channels := []model.Channel{{Id: 1, Group: "default", Models: fingerprint.Model, Status: 1}}
|
|
rounds := []AstraSamplingRound{{ID: "valid", ChannelID: 1, StartedAt: now.Unix() - 100, FinishedAt: now.Unix() - 90, Result: fingerprint.Result{Engine: fingerprint.Version, Status: "mismatch", Closest: "gpt-5.5", ValidSamples: 1}}}
|
|
for i := range 15 {
|
|
rounds = append(rounds, AstraSamplingRound{ChannelID: 1, StartedAt: now.Unix() - int64(i), FinishedAt: now.Unix(), Result: fingerprint.Result{Engine: fingerprint.Version, Status: "error"}})
|
|
}
|
|
body, err := common.Marshal(AstraSamplingState{Rounds: rounds})
|
|
require.NoError(t, err)
|
|
tasks := []*model.SystemTask{{Status: model.SystemTaskStatusSucceeded, State: string(body)}}
|
|
summary := AggregateAstraSampling(channels, tasks, now)["default"]
|
|
require.NotNil(t, summary.LatestValid)
|
|
assert.Equal(t, "gpt-5.5", summary.LatestValid.Closest)
|
|
assert.Len(t, summary.History, 16)
|
|
rounds = append(rounds, AstraSamplingRound{ID: "recovered", ChannelID: 1, StartedAt: now.Unix(), FinishedAt: now.Unix(), Result: fingerprint.Result{Engine: fingerprint.Version, Status: "match", Closest: fingerprint.Model, ValidSamples: 2}})
|
|
body, err = common.Marshal(AstraSamplingState{Rounds: rounds})
|
|
require.NoError(t, err)
|
|
tasks[0].State = string(body)
|
|
summary = AggregateAstraSampling(channels, tasks, now)["default"]
|
|
require.NotNil(t, summary.LatestValid)
|
|
assert.Equal(t, fingerprint.Model, summary.LatestValid.Closest)
|
|
reset := aggregateAstraSampling(channels, tasks, now, map[string]model.StatusPageGroup{"default": {ResetBefore: now.Unix()}})
|
|
assert.Nil(t, reset["default"].LatestValid)
|
|
}
|
|
|
|
func TestGroupDegradationOverridesExistingPoolAffinity(t *testing.T) {
|
|
db := setupChannelSelectAutoGroupsTest(t)
|
|
createChannelSelectAutoGroupsChannel(t, db, 6401, "default", fingerprint.Model)
|
|
createChannelSelectAutoGroupsChannel(t, db, 6402, "vip", fingerprint.Model)
|
|
model.InitChannelCache()
|
|
ResetGroupAvailability()
|
|
t.Cleanup(ResetGroupAvailability)
|
|
previousRedis := common.RedisEnabled
|
|
common.RedisEnabled = false
|
|
t.Cleanup(func() { common.RedisEnabled = previousRedis; _ = poolBindings().Purge() })
|
|
require.NoError(t, poolBindings().Purge())
|
|
config := operation_setting.GetChannelAffinitySetting()
|
|
previous := *config
|
|
*config = operation_setting.ChannelAffinitySetting{Enabled: true, DefaultTTLSeconds: 3600, SessionMode: "strict", Rules: []operation_setting.ChannelAffinityRule{{Name: "degradation-test", ModelRegex: []string{".*"}, KeySources: []operation_setting.ChannelAffinityKeySource{{Type: "request_header", Key: "Session-Id"}}, SessionMode: "strict"}}}
|
|
t.Cleanup(func() { *config = previous })
|
|
ctx := func(degraded bool) *gin.Context {
|
|
c := setupGroupPoolContext(t)
|
|
c.Request = httptest.NewRequest("POST", "/v1/responses", nil)
|
|
c.Request.Header.Set("Session-Id", "degradation-session")
|
|
c.Set("id", 100)
|
|
c.Set("token_id", 100)
|
|
c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: "custom", SkipDegradedGroups: true, Groups: []model.TokenGroupEntry{{Group: "default", Priority: 0}, {Group: "vip", Priority: 1}}})
|
|
rounds := map[string]AstraSamplingHistory{}
|
|
if degraded {
|
|
rounds["default"] = AstraSamplingHistory{At: time.Now().Unix(), Closest: "gpt-5.5", ValidSamples: 1}
|
|
}
|
|
c.Set("pool_group_fingerprints", rounds)
|
|
return c
|
|
}
|
|
first := ctx(false)
|
|
selected, group, err := SelectChannelForRequest(first, fingerprint.Model, &RetryParam{Ctx: first, TokenGroup: "default", ModelName: fingerprint.Model})
|
|
require.Nil(t, err)
|
|
require.NotNil(t, selected)
|
|
require.Equal(t, 6401, selected.Id)
|
|
first.Set("channel_id", selected.Id)
|
|
common.SetContextKey(first, constant.ContextKeyAutoGroup, group)
|
|
RequestPolicy(first).Successful = true
|
|
RecordPoolAffinity(first)
|
|
bound, _ := PreferredPoolChannel(ctx(false), fingerprint.Model)
|
|
require.NotNil(t, bound)
|
|
image := ctx(false)
|
|
image.Request = httptest.NewRequest("POST", "/v1/responses", strings.NewReader(`{"input":[{"role":"user","content":[{"type":"input_image","image_url":"data:image/png;base64,eA=="}]}]}`))
|
|
image.Request.Header.Set("Session-Id", "degradation-session")
|
|
image.Set("pool_group_visions", map[string]DetectionHistory{"default": {Quality: "incomplete"}})
|
|
selected, group, err = SelectChannelForRequest(image, fingerprint.Model, &RetryParam{Ctx: image, TokenGroup: "default", ModelName: fingerprint.Model})
|
|
require.Nil(t, err)
|
|
require.NotNil(t, selected)
|
|
assert.Equal(t, 6401, selected.Id, "incomplete vision must preserve the existing session affinity")
|
|
assert.Equal(t, "default", group)
|
|
next := ctx(true)
|
|
selected, group, err = SelectChannelForRequest(next, fingerprint.Model, &RetryParam{Ctx: next, TokenGroup: "default", ModelName: fingerprint.Model})
|
|
require.Nil(t, err)
|
|
require.NotNil(t, selected)
|
|
assert.Equal(t, 6402, selected.Id)
|
|
assert.Equal(t, "vip", group)
|
|
}
|
|
|
|
func TestVisionIncompleteRoutingAcrossModelsAndPolicies(t *testing.T) {
|
|
db := setupChannelSelectAutoGroupsTest(t)
|
|
for i, group := range []string{"default", "vip"} {
|
|
createChannelSelectAutoGroupsChannel(t, db, 6601+i, group, fingerprint.Model)
|
|
require.NoError(t, db.Model(&model.Channel{}).Where("id = ?", 6601+i).Update("models", fingerprint.Model+",gpt-6-sol").Error)
|
|
require.NoError(t, db.Create(&model.Ability{Group: group, Model: "gpt-6-sol", ChannelId: 6601 + i, Enabled: true, Weight: 100}).Error)
|
|
}
|
|
model.InitChannelCache()
|
|
ResetGroupAvailability()
|
|
t.Cleanup(ResetGroupAvailability)
|
|
for _, mode := range []string{"smart", "custom"} {
|
|
for _, allow := range []bool{false, true} {
|
|
for _, modelName := range []string{fingerprint.Model, "gpt-6-sol"} {
|
|
for _, image := range []bool{false, true} {
|
|
c := setupGroupPoolContext(t)
|
|
c.Request = httptest.NewRequest("POST", "/v1/responses", nil)
|
|
c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: mode, AllowDegraded: allow, SkipDegradedGroups: !allow, Groups: []model.TokenGroupEntry{{Group: "default"}, {Group: "vip", Priority: 1}}})
|
|
if image {
|
|
c.Request = httptest.NewRequest("POST", "/v1/responses", strings.NewReader(`{"input":[{"role":"user","content":[{"type":"input_image","image_url":"data:image/png;base64,eA=="}]}]}`))
|
|
}
|
|
c.Set("pool_group_fingerprints", map[string]AstraSamplingHistory{})
|
|
visions := map[string]DetectionHistory{"default": {Quality: "incomplete"}, "vip": {Quality: "incomplete"}}
|
|
c.Set("pool_group_visions", visions)
|
|
channel, _, err := selectTokenRoutingPool(&RetryParam{Ctx: c, TokenGroup: "default", ModelName: modelName})
|
|
require.NoError(t, err)
|
|
require.NotNil(t, channel)
|
|
assert.Equal(t, 6601, channel.Id, "mode=%s allow=%v model=%s image=%v", mode, allow, modelName, image)
|
|
assert.False(t, poolGroupSkipped(c, modelName, "default"), "incomplete vision must not exclude retries or cooldown recovery")
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestGroupDegradationSkipsPoolGroupsAndRespectsUnknownRecoveryAndOptOut(t *testing.T) {
|
|
db := setupChannelSelectAutoGroupsTest(t)
|
|
createChannelSelectAutoGroupsChannel(t, db, 6301, "default", fingerprint.Model)
|
|
createChannelSelectAutoGroupsChannel(t, db, 6302, "vip", fingerprint.Model)
|
|
model.InitChannelCache()
|
|
ResetGroupAvailability()
|
|
t.Cleanup(ResetGroupAvailability)
|
|
for _, test := range []struct {
|
|
name string
|
|
enabled bool
|
|
closest string
|
|
age time.Duration
|
|
degradeBackup bool
|
|
want int
|
|
}{
|
|
{"latest mismatch switches", true, "gpt-5.5", time.Minute, false, 6302},
|
|
{"recovered group returns", true, fingerprint.Model, time.Minute, false, 6301},
|
|
{"unknown stays eligible", true, "", time.Minute, false, 6301},
|
|
{"valid result persists beyond window", true, "gpt-5.5", 7 * time.Hour, false, 6302},
|
|
{"disabled keeps routing", false, "gpt-5.5", time.Minute, false, 6301},
|
|
{"all degraded has no route", true, "gpt-5.5", time.Minute, true, 0},
|
|
} {
|
|
t.Run(test.name, func(t *testing.T) {
|
|
c := setupGroupPoolContext(t)
|
|
c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: "custom", SkipDegradedGroups: test.enabled, Groups: []model.TokenGroupEntry{{Group: "default", Priority: 0}, {Group: "vip", Priority: 1}}})
|
|
rounds := map[string]AstraSamplingHistory{"default": {At: time.Now().Add(-test.age).Unix(), Closest: test.closest, ValidSamples: 1}}
|
|
if test.degradeBackup {
|
|
rounds["vip"] = rounds["default"]
|
|
}
|
|
c.Set("pool_group_fingerprints", rounds)
|
|
selected, _, err := selectTokenRoutingPool(&RetryParam{Ctx: c, TokenGroup: "default", ModelName: fingerprint.Model})
|
|
require.NoError(t, err)
|
|
if test.want == 0 {
|
|
assert.Nil(t, selected)
|
|
return
|
|
}
|
|
require.NotNil(t, selected)
|
|
assert.Equal(t, test.want, selected.Id)
|
|
})
|
|
}
|
|
}
|