feat: sync deployed v1.0.0-rc.38-custom.42
Publish Docker image (Multi-arch) / Build & push (amd64) (push) Failing after 1m40s
Release (Linux, macOS, Windows) / Linux Release (push) Failing after 42m16s
Publish Docker image (Multi-arch) / Build & push (arm64) (push) Canceled after 0s
Release (Linux, macOS, Windows) / macOS Release (push) Canceled after 0s
Release (Linux, macOS, Windows) / Windows Release (push) Canceled after 0s
Publish Docker image (Multi-arch) / Create multi-arch manifests (push) Canceled after 0s
Publish Docker image (Multi-arch) / Build & push (amd64) (push) Failing after 1m40s
Release (Linux, macOS, Windows) / Linux Release (push) Failing after 42m16s
Publish Docker image (Multi-arch) / Build & push (arm64) (push) Canceled after 0s
Release (Linux, macOS, Windows) / macOS Release (push) Canceled after 0s
Release (Linux, macOS, Windows) / Windows Release (push) Canceled after 0s
Publish Docker image (Multi-arch) / Create multi-arch manifests (push) Canceled after 0s
Sync production source including routing, billing, Astra detection and public uptime. Update custom build instructions and license packaging.
This commit is contained in:
1 parent
c2b44332d6
commit
a5f7b8e586
828 files changed
+122253
-13702
No files matched your search
+19
-2
@@ -131,13 +131,14 @@ func GetChannel(
|
||||
model string,
|
||||
retry int,
|
||||
filters []dto.ChannelFilter,
|
||||
traces ...*SelectionTrace,
|
||||
) (*Channel, error) {
|
||||
var abilities []Ability
|
||||
err := DB.Where(commonGroupCol+" = ? and model = ? and enabled = ?", group, model, true).Order("priority DESC, weight DESC").Find(&abilities).Error
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
abilities = filterAbilitiesByConstraints(abilities, model, filters)
|
||||
abilities = filterAbilitiesByConstraints(abilities, model, filters, traces...)
|
||||
if len(abilities) > 0 {
|
||||
priorities := make([]int64, 0)
|
||||
seen := make(map[int64]bool)
|
||||
@@ -187,7 +188,7 @@ func GetChannel(
|
||||
// filterAbilitiesByConstraints applies the same ChannelSatisfiesFilters
|
||||
// predicate used by the memory-cache path. A failed channel lookup fails
|
||||
// closed when a task-plugin identity is required and fails open otherwise.
|
||||
func filterAbilitiesByConstraints(abilities []Ability, modelName string, filters []dto.ChannelFilter) []Ability {
|
||||
func filterAbilitiesByConstraints(abilities []Ability, modelName string, filters []dto.ChannelFilter, traces ...*SelectionTrace) []Ability {
|
||||
if len(abilities) == 0 {
|
||||
return nil
|
||||
}
|
||||
@@ -218,6 +219,22 @@ func filterAbilitiesByConstraints(abilities []Ability, modelName string, filters
|
||||
filtered := make([]Ability, 0, len(abilities))
|
||||
for _, ability := range abilities {
|
||||
channel := channelsByID[ability.ChannelId]
|
||||
if len(traces) > 0 && traces[0] != nil {
|
||||
reason := "eligible"
|
||||
if ok, kind := ChannelSatisfiesFilters(channel, modelName, filters); !ok {
|
||||
reason = "excluded_" + string(kind)
|
||||
}
|
||||
traces[0].Add(channel, ability.ChannelId, ability.Group, modelName, reason)
|
||||
// DB routing uses the ability's priority and weight.
|
||||
if n := len(traces[0].Candidates); n > 0 && traces[0].Candidates[n-1].ChannelID == ability.ChannelId {
|
||||
candidate := &traces[0].Candidates[n-1]
|
||||
candidate.Weight = int(ability.Weight)
|
||||
candidate.Priority = 0
|
||||
if ability.Priority != nil {
|
||||
candidate.Priority = *ability.Priority
|
||||
}
|
||||
}
|
||||
}
|
||||
if ok, _ := ChannelSatisfiesFilters(channel, modelName, filters); ok {
|
||||
filtered = append(filtered, ability)
|
||||
}
|
||||
|
||||
@@ -149,3 +149,25 @@ func UpdateUserOAuthBindingForSessionWithTx(tx *gorm.DB, identity AuthSessionIde
|
||||
// The existing unique provider/subject index rejects concurrent ownership.
|
||||
return tx.Model(&binding).Update("provider_user_id", subject).Error
|
||||
}
|
||||
|
||||
// MigrateLegacyGitHubBindingWithTx rewrites a login-name GitHub binding to the
|
||||
// numeric account ID once the login flow has confirmed the account. It reports
|
||||
// whether the row was rewritten; a row whose binding changed since it was read
|
||||
// is left untouched.
|
||||
func MigrateLegacyGitHubBindingWithTx(tx *gorm.DB, userID int, legacyID, gitHubID string) (bool, error) {
|
||||
if userID <= 0 || legacyID == "" || gitHubID == "" {
|
||||
return false, ErrAccountBindingChanged
|
||||
}
|
||||
var count int64
|
||||
if err := tx.Model(&User{}).Where("github_id = ? AND id <> ?", gitHubID, userID).Count(&count).Error; err != nil {
|
||||
return false, err
|
||||
}
|
||||
if count != 0 {
|
||||
return false, ErrExternalIdentityAlreadyClaimed
|
||||
}
|
||||
result := tx.Model(&User{}).Where("id = ? AND github_id = ?", userID, legacyID).Update("github_id", gitHubID)
|
||||
if result.Error != nil {
|
||||
return false, result.Error
|
||||
}
|
||||
return result.RowsAffected == 1, nil
|
||||
}
|
||||
+33
-5
@@ -59,6 +59,8 @@ type Channel struct {
|
||||
Keys []string `json:"-" gorm:"-"`
|
||||
}
|
||||
|
||||
const ChannelStatusReasonAllKeysDisabled = "All keys are disabled"
|
||||
|
||||
type ChannelInfo struct {
|
||||
IsMultiKey bool `json:"is_multi_key"` // 是否多Key模式
|
||||
MultiKeySize int `json:"multi_key_size"` // 多Key模式下的Key数量
|
||||
@@ -711,7 +713,7 @@ func handlerMultiKeyUpdate(channel *Channel, usingKey string, status int, reason
|
||||
if !hasEnabledMultiKey(keys, channel.ChannelInfo.MultiKeyStatusList) {
|
||||
channel.Status = common.ChannelStatusAutoDisabled
|
||||
info := channel.GetOtherInfo()
|
||||
info["status_reason"] = "All keys are disabled"
|
||||
info["status_reason"] = ChannelStatusReasonAllKeysDisabled
|
||||
info["status_time"] = common.GetTimestamp()
|
||||
channel.SetOtherInfo(info)
|
||||
} else if status == common.ChannelStatusEnabled {
|
||||
@@ -782,7 +784,12 @@ func UpdateChannelStatus(channelId int, usingKey string, status int, reason stri
|
||||
if err != nil {
|
||||
return false
|
||||
} else {
|
||||
if channel.Status == status {
|
||||
// A manual channel operation must replace the exhaustion reason even
|
||||
// when the status value is already manually disabled.
|
||||
overridesKeyExhaustion := channel.ChannelInfo.IsMultiKey && usingKey == "" &&
|
||||
status == common.ChannelStatusManuallyDisabled && reason != ChannelStatusReasonAllKeysDisabled &&
|
||||
channel.GetOtherInfo()["status_reason"] == ChannelStatusReasonAllKeysDisabled
|
||||
if channel.Status == status && !overridesKeyExhaustion {
|
||||
return false
|
||||
}
|
||||
|
||||
@@ -819,6 +826,19 @@ func EnableChannelByTag(tag string) error {
|
||||
}
|
||||
|
||||
func DisableChannelByTag(tag string) error {
|
||||
// Explicit tag-level disable also cancels automatic restoration for
|
||||
// channels that were already disabled because all keys were unavailable.
|
||||
var channels []Channel
|
||||
if err := DB.Where("tag = ?", tag).Find(&channels).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
for _, channel := range channels {
|
||||
if channel.ChannelInfo.IsMultiKey && channel.GetOtherInfo()["status_reason"] == ChannelStatusReasonAllKeysDisabled {
|
||||
if !UpdateChannelStatus(channel.Id, "", common.ChannelStatusManuallyDisabled, "manual tag operation") {
|
||||
return fmt.Errorf("failed to disable channel #%d by tag", channel.Id)
|
||||
}
|
||||
}
|
||||
}
|
||||
err := DB.Model(&Channel{}).Where("tag = ?", tag).Update("status", common.ChannelStatusManuallyDisabled).Error
|
||||
if err != nil {
|
||||
return err
|
||||
@@ -992,7 +1012,10 @@ func (channel *Channel) ValidateSettings() error {
|
||||
if err := channelOtherSettings.ValidateToolLossPolicy(); err != nil {
|
||||
return err
|
||||
}
|
||||
if channel.Type == constant.ChannelTypeAdvancedCustom {
|
||||
if preset := common.GetAdvancedCustomPreset(channel.Type); preset != nil {
|
||||
channelOtherSettings.AdvancedCustom = preset
|
||||
}
|
||||
if constant.IsAdvancedCustomChannel(channel.Type) {
|
||||
if channelOtherSettings.AdvancedCustom == nil {
|
||||
return fmt.Errorf("advanced_custom is required")
|
||||
}
|
||||
@@ -1002,7 +1025,7 @@ func (channel *Channel) ValidateSettings() error {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if channel.Type == constant.ChannelTypeAdvancedCustom && channelOtherSettings.UpstreamModelUpdateCheckEnabled {
|
||||
if constant.IsAdvancedCustomChannel(channel.Type) && channelOtherSettings.UpstreamModelUpdateCheckEnabled {
|
||||
if _, ok := channelOtherSettings.AdvancedCustom.ModelListRoute(); !ok {
|
||||
return fmt.Errorf("advanced custom channels require a %s route when upstream model update checks are enabled", dto.AdvancedCustomModelListPath)
|
||||
}
|
||||
@@ -1011,7 +1034,9 @@ func (channel *Channel) ValidateSettings() error {
|
||||
}
|
||||
|
||||
func (channel *Channel) GetSetting() dto.ChannelSettings {
|
||||
setting := dto.ChannelSettings{}
|
||||
// Model validation is enabled by default; an explicit false in the JSON
|
||||
// remains the per-channel opt-out.
|
||||
setting := dto.ChannelSettings{ResponsesAstraModelGuard: common.GetPointer(true)}
|
||||
if channel.Setting != nil && *channel.Setting != "" {
|
||||
err := common.Unmarshal([]byte(*channel.Setting), &setting)
|
||||
if err != nil {
|
||||
@@ -1042,6 +1067,9 @@ func (channel *Channel) GetOtherSettings() dto.ChannelOtherSettings {
|
||||
_ = channel.Save() // 保存修改
|
||||
}
|
||||
}
|
||||
if preset := common.GetAdvancedCustomPreset(channel.Type); preset != nil {
|
||||
setting.AdvancedCustom = preset
|
||||
}
|
||||
return setting
|
||||
}
|
||||
|
||||
|
||||
+21
-7
@@ -58,12 +58,17 @@ func InitChannelCache() {
|
||||
}
|
||||
groups := strings.SplitSeq(channel.Group, ",")
|
||||
for group := range groups {
|
||||
model2channels, ok := newGroup2model2channels[group]
|
||||
if !ok {
|
||||
model2channels = make(map[string][]int)
|
||||
newGroup2model2channels[group] = model2channels
|
||||
}
|
||||
models := channel.GetModels()
|
||||
for _, model := range models {
|
||||
if _, ok := newGroup2model2channels[group][model]; !ok {
|
||||
newGroup2model2channels[group][model] = make([]int, 0)
|
||||
if _, ok := model2channels[model]; !ok {
|
||||
model2channels[model] = make([]int, 0)
|
||||
}
|
||||
newGroup2model2channels[group][model] = append(newGroup2model2channels[group][model], channel.Id)
|
||||
model2channels[model] = append(model2channels[model], channel.Id)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -119,22 +124,31 @@ func GetRandomSatisfiedChannel(
|
||||
model string,
|
||||
retry int,
|
||||
filters []dto.ChannelFilter,
|
||||
) (*Channel, error) {
|
||||
traces ...*SelectionTrace,
|
||||
) (selected *Channel, err error) {
|
||||
var trace *SelectionTrace
|
||||
if len(traces) > 0 {
|
||||
trace = traces[0]
|
||||
}
|
||||
// if memory cache is disabled, get channel directly from database
|
||||
if !common.MemoryCacheEnabled {
|
||||
return GetChannel(group, model, retry, filters)
|
||||
defer func() { trace.Finish(selected) }()
|
||||
return GetChannel(group, model, retry, filters, trace)
|
||||
}
|
||||
|
||||
channelSyncLock.RLock()
|
||||
defer channelSyncLock.RUnlock()
|
||||
defer func() { trace.Finish(selected) }()
|
||||
|
||||
// First, try to find channels with the exact model name.
|
||||
channels, _ := filterCandidateIDs(group2model2channels[group][model], model, filters)
|
||||
channels, _ := filterCandidateIDsObserved(group2model2channels[group][model], model, filters, trace, group)
|
||||
|
||||
// If no channels found, try to find channels with the normalized model name.
|
||||
if len(channels) == 0 {
|
||||
normalizedModel := ratio_setting.RoutingMatchModelName(model)
|
||||
channels, _ = filterCandidateIDs(group2model2channels[group][normalizedModel], model, filters)
|
||||
if normalizedModel != model {
|
||||
channels, _ = filterCandidateIDsObserved(group2model2channels[group][normalizedModel], model, filters, trace, group)
|
||||
}
|
||||
}
|
||||
|
||||
if len(channels) == 0 {
|
||||
|
||||
@@ -0,0 +1,60 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestInitChannelCacheHandlesChannelBeforeAbilityRefresh(t *testing.T) {
|
||||
resetPricingEndpointTestTables(t)
|
||||
|
||||
channel := &Channel{
|
||||
Id: 501,
|
||||
Type: constant.ChannelTypeOpenAI,
|
||||
Key: "key-501",
|
||||
Status: common.ChannelStatusEnabled,
|
||||
Name: "channel-before-ability-refresh",
|
||||
Models: "grok-4.6",
|
||||
Group: "groksuper",
|
||||
}
|
||||
require.NoError(t, DB.Create(channel).Error)
|
||||
|
||||
require.NotPanics(t, InitChannelCache)
|
||||
|
||||
selected, err := GetRandomSatisfiedChannel("groksuper", "grok-4.6", 0, nil)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, selected)
|
||||
assert.Equal(t, channel.Id, selected.Id)
|
||||
}
|
||||
|
||||
func TestSelectionTraceSnapshot(t *testing.T) {
|
||||
resetPricingEndpointTestTables(t)
|
||||
previous := common.MemoryCacheEnabled
|
||||
common.MemoryCacheEnabled = true
|
||||
t.Cleanup(func() { common.MemoryCacheEnabled = previous })
|
||||
for i := 1; i <= 2; i++ {
|
||||
ch := &Channel{Id: 700 + i, Type: constant.ChannelTypeOpenAI, Key: "secret", Status: common.ChannelStatusEnabled, Name: "original", Models: "test", Group: "trace", Priority: common.GetPointer(int64(i))}
|
||||
require.NoError(t, DB.Create(ch).Error)
|
||||
require.NoError(t, ch.AddAbilities(nil))
|
||||
}
|
||||
InitChannelCache()
|
||||
trace := &SelectionTrace{}
|
||||
ch, err := GetRandomSatisfiedChannel("trace", "test", 0, nil, trace)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 702, ch.Id)
|
||||
require.Equal(t, 2, trace.Total)
|
||||
require.Len(t, trace.Candidates, 2)
|
||||
for _, candidate := range trace.Candidates {
|
||||
if candidate.ChannelID == 702 {
|
||||
require.Equal(t, "selected", candidate.Reason)
|
||||
} else {
|
||||
require.Equal(t, "different_priority", candidate.Reason)
|
||||
}
|
||||
}
|
||||
ch.Name = "changed"
|
||||
require.Equal(t, "original", trace.Candidates[0].Name)
|
||||
}
|
||||
@@ -8,8 +8,10 @@ import (
|
||||
)
|
||||
|
||||
var filterEvalOrder = []dto.ChannelFilterKind{
|
||||
dto.FilterAttemptedChannel,
|
||||
dto.FilterRequestPath,
|
||||
dto.FilterTaskPluginIdentity,
|
||||
dto.FilterResponsesWebSocket,
|
||||
}
|
||||
|
||||
// ChannelSatisfiesFilters reports whether ch passes every filter.
|
||||
@@ -37,6 +39,27 @@ func ChannelSatisfiesFilters(ch *Channel, modelName string, filters []dto.Channe
|
||||
// A missing id in channelsIDM is kept for request_path (downstream consistency
|
||||
// error) and dropped for task_plugin_identity, matching the previous filters.
|
||||
func filterCandidateIDs(ids []int, modelName string, filters []dto.ChannelFilter) (kept []int, emptiedBy dto.ChannelFilterKind) {
|
||||
return filterCandidateIDsObserved(ids, modelName, filters, nil, "")
|
||||
}
|
||||
|
||||
func filterCandidateIDsObserved(ids []int, modelName string, filters []dto.ChannelFilter, trace *SelectionTrace, group string) (kept []int, emptiedBy dto.ChannelFilterKind) {
|
||||
if trace != nil {
|
||||
for _, id := range ids {
|
||||
ch, exists := channelsIDM[id]
|
||||
reason := "eligible"
|
||||
for _, kind := range filterEvalOrder {
|
||||
fs := filtersByKind(filters, kind)
|
||||
if len(fs) > 0 && !candidatePassesKindFilters(ch, exists, modelName, kind, fs) {
|
||||
reason = "excluded_" + string(kind)
|
||||
break
|
||||
}
|
||||
}
|
||||
if !exists {
|
||||
reason = "missing_channel"
|
||||
}
|
||||
trace.Add(ch, id, group, modelName, reason)
|
||||
}
|
||||
}
|
||||
if len(ids) == 0 {
|
||||
return ids, ""
|
||||
}
|
||||
@@ -88,20 +111,38 @@ func candidatePassesKindFilters(ch *Channel, exists bool, modelName string, kind
|
||||
|
||||
func channelMatchesFilter(ch *Channel, modelName string, filter dto.ChannelFilter) bool {
|
||||
switch filter.Kind {
|
||||
case dto.FilterAttemptedChannel:
|
||||
return !slices.Contains(filter.ExcludedChannelIDs, ch.Id)
|
||||
case dto.FilterRequestPath:
|
||||
if filter.RequestPath == "" {
|
||||
return true
|
||||
}
|
||||
if ch.Type != constant.ChannelTypeAdvancedCustom {
|
||||
if !constant.IsAdvancedCustomChannel(ch.Type) {
|
||||
return true
|
||||
}
|
||||
config := ch.GetOtherSettings().AdvancedCustom
|
||||
return config != nil && config.SupportsPathForModel(filter.RequestPath, modelName)
|
||||
case dto.FilterTaskPluginIdentity:
|
||||
if ch.Type == constant.ChannelTypeTaskPlugin {
|
||||
return filter.TaskPluginKey != "" && ch.GetSetting().TaskPluginKey == filter.TaskPluginKey
|
||||
key := ch.GetSetting().TaskPluginKey
|
||||
return filter.TaskPluginKey != "" && (key == filter.TaskPluginKey || slices.Contains(filter.TaskPluginKeys, key))
|
||||
}
|
||||
return filter.TaskPluginKey == "" || slices.Contains(filter.TaskPluginChannelTypes, ch.Type)
|
||||
case dto.FilterResponsesWebSocket:
|
||||
if !ch.GetSetting().ResponsesWebSocketEnabled {
|
||||
return false
|
||||
}
|
||||
switch ch.Type {
|
||||
case constant.ChannelTypeOpenAI, constant.ChannelTypeCodex, constant.ChannelTypeSub2API, constant.ChannelTypeNewAPI:
|
||||
return true
|
||||
case constant.ChannelTypeAdvancedCustom:
|
||||
// The session forwards native Responses events without protocol
|
||||
// conversion, so only a converter-free /v1/responses route qualifies.
|
||||
route, ok := ch.GetOtherSettings().AdvancedCustom.MatchPathForModel("/v1/responses", modelName)
|
||||
return ok && route.IsNative()
|
||||
default:
|
||||
return false
|
||||
}
|
||||
default:
|
||||
return true
|
||||
}
|
||||
|
||||
@@ -1,12 +1,21 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
filterdto "github.com/QuantumNous/new-api/dto"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
"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"
|
||||
)
|
||||
|
||||
func TestChannelValidateSettingsRejectsInvalidHTTPTransport(t *testing.T) {
|
||||
@@ -41,6 +50,25 @@ func TestChannelValidateSettingsRejectsInvalidHTTPTransport(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestChannelGetSettingDefaultsAstraGuardOn(t *testing.T) {
|
||||
channel := &Channel{}
|
||||
|
||||
assert.True(t, channel.GetSetting().AstraModelGuardEnabled())
|
||||
|
||||
channel.Setting = common.GetPointer(`{"proxy":""}`)
|
||||
assert.True(t, channel.GetSetting().AstraModelGuardEnabled())
|
||||
|
||||
channel.Setting = common.GetPointer(`{"responses_astra_model_guard":false}`)
|
||||
setting := channel.GetSetting()
|
||||
assert.False(t, setting.AstraModelGuardEnabled())
|
||||
|
||||
setting.Proxy = "http://127.0.0.1:8080"
|
||||
channel.SetSetting(setting)
|
||||
reloaded := channel.GetSetting()
|
||||
assert.False(t, reloaded.AstraModelGuardEnabled())
|
||||
assert.Contains(t, *channel.Setting, `"responses_astra_model_guard":false`)
|
||||
}
|
||||
|
||||
func TestAdvancedCustomChannelRequiresModelListRouteOnlyWhenUpdateChecksEnabled(t *testing.T) {
|
||||
inferenceRoute := dto.AdvancedCustomRoute{
|
||||
IncomingPath: "/v1/chat/completions",
|
||||
@@ -98,3 +126,91 @@ func TestAdvancedCustomChannelRequiresModelListRouteOnlyWhenUpdateChecksEnabled(
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestInferencePresetSettingsAndDatabaseRoundTrip(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(filepath.Join(t.TempDir(), "presets.db"))
|
||||
case "mysql":
|
||||
dsn := os.Getenv("TEST_MYSQL_DSN")
|
||||
if dsn == "" {
|
||||
t.Skip("TEST_MYSQL_DSN is not configured")
|
||||
}
|
||||
driver = mysql.Open(dsn)
|
||||
case "postgres":
|
||||
dsn := os.Getenv("TEST_POSTGRES_DSN")
|
||||
if dsn == "" {
|
||||
t.Skip("TEST_POSTGRES_DSN is not configured")
|
||||
}
|
||||
driver = postgres.Open(dsn)
|
||||
}
|
||||
db, err := gorm.Open(driver, &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
sqlDB, err := db.DB()
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(func() { require.NoError(t, sqlDB.Close()) })
|
||||
table := db.Table("inference_preset_channels").Session(&gorm.Session{})
|
||||
require.NoError(t, table.AutoMigrate(&Channel{}))
|
||||
t.Cleanup(func() { require.NoError(t, db.Migrator().DropTable("inference_preset_channels")) })
|
||||
var version string
|
||||
if dialect == "sqlite" {
|
||||
require.NoError(t, db.Raw("select sqlite_version()").Scan(&version).Error)
|
||||
} else {
|
||||
require.NoError(t, db.Raw("select version()").Scan(&version).Error)
|
||||
}
|
||||
t.Logf("%s version: %s", dialect, version)
|
||||
for _, channelType := range []int{constant.ChannelTypeVLLM, constant.ChannelTypeSGLang} {
|
||||
t.Run(fmt.Sprint(channelType), func(t *testing.T) {
|
||||
channel := &Channel{Type: channelType, Key: "EMPTY", Name: "inference", Status: common.ChannelStatusEnabled}
|
||||
require.NoError(t, channel.ValidateSettings())
|
||||
require.NotNil(t, channel.GetOtherSettings().AdvancedCustom)
|
||||
require.NoError(t, table.Create(channel).Error)
|
||||
for range 2 {
|
||||
var loaded Channel
|
||||
require.NoError(t, table.First(&loaded, channel.Id).Error)
|
||||
assert.Equal(t, channelType, loaded.Type)
|
||||
assert.Empty(t, loaded.OtherSettings)
|
||||
defaults := loaded.GetOtherSettings().AdvancedCustom
|
||||
require.NotNil(t, defaults)
|
||||
assert.True(t, defaults.SupportsPath("/v1/messages"))
|
||||
assert.Empty(t, loaded.OtherSettings, "reading defaults must not rewrite saved settings")
|
||||
defaults.Routes[0].UpstreamPath = "/changed-locally"
|
||||
assert.Equal(t, "/v1/chat/completions", loaded.GetOtherSettings().AdvancedCustom.Routes[0].UpstreamPath)
|
||||
}
|
||||
settings := dto.ChannelOtherSettings{AdvancedCustom: &dto.AdvancedCustomConfig{Routes: []dto.AdvancedCustomRoute{{IncomingPath: "/v1/chat/completions", UpstreamPath: "/custom/chat", Models: []string{"allowed"}, Auth: &dto.AdvancedCustomRouteAuth{Type: "none"}}}}}
|
||||
channel.SetOtherSettings(settings)
|
||||
require.NoError(t, channel.ValidateSettings())
|
||||
require.NoError(t, table.Save(channel).Error)
|
||||
for range 2 {
|
||||
var loaded Channel
|
||||
require.NoError(t, table.First(&loaded, channel.Id).Error)
|
||||
actual := loaded.GetOtherSettings().AdvancedCustom
|
||||
require.Equal(t, common.GetAdvancedCustomPreset(channelType), actual, "named channels must ignore editable advanced_custom overrides")
|
||||
assert.Equal(t, channel.OtherSettings, loaded.OtherSettings)
|
||||
for _, tc := range []struct {
|
||||
path, model string
|
||||
allowed bool
|
||||
}{
|
||||
{"/v1/chat/completions", "allowed", true},
|
||||
{"/v1/chat/completions", "other", true},
|
||||
{"/v1/messages", "allowed", true},
|
||||
{"/v1/images/generations", "allowed", false},
|
||||
} {
|
||||
ok, _ := ChannelSatisfiesFilters(&loaded, tc.model, []filterdto.ChannelFilter{{Kind: filterdto.FilterRequestPath, RequestPath: tc.path}})
|
||||
assert.Equal(t, tc.allowed, ok)
|
||||
}
|
||||
endpoints := getPricingEndpointTypesForAbility(AbilityWithChannel{ChannelType: channelType, Ability: Ability{Model: "allowed", ChannelId: loaded.Id}}, map[int]*dto.AdvancedCustomConfig{loaded.Id: actual})
|
||||
expectedEndpoints := []constant.EndpointType{constant.EndpointTypeOpenAI, constant.EndpointTypeOpenAIResponse, constant.EndpointTypeAnthropic, constant.EndpointTypeEmbeddings}
|
||||
if channelType == constant.ChannelTypeSGLang {
|
||||
expectedEndpoints = append(expectedEndpoints, constant.EndpointTypeJinaRerank)
|
||||
}
|
||||
assert.ElementsMatch(t, expectedEndpoints, endpoints)
|
||||
}
|
||||
})
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,229 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
)
|
||||
|
||||
// DetectionRun is both the durable task and wallet ledger. Public projections
|
||||
// are built explicitly; this record must never be serialized to anonymous users.
|
||||
type DetectionRun struct {
|
||||
ID string `gorm:"primaryKey;type:varchar(64)"`
|
||||
UserID int `gorm:"index"`
|
||||
RequestKey string `gorm:"uniqueIndex;type:varchar(128)"`
|
||||
ChannelID int `gorm:"index"`
|
||||
GroupName string
|
||||
Kind string `gorm:"type:varchar(24)"`
|
||||
Status string `gorm:"type:varchar(24);index"`
|
||||
CreatedAt int64 `gorm:"autoCreateTime:false;index"`
|
||||
FinishedAt int64
|
||||
Deadline int64
|
||||
Reserved int
|
||||
Charged int
|
||||
Price string `gorm:"type:text"`
|
||||
RemoteID string
|
||||
Quality string
|
||||
HTML string `gorm:"type:text"`
|
||||
Result string `gorm:"type:text"`
|
||||
Error string
|
||||
Settled bool
|
||||
LogJSON string `gorm:"type:text"`
|
||||
LogDelivered bool `gorm:"index"`
|
||||
}
|
||||
|
||||
type DetectionCall struct {
|
||||
ID string `gorm:"primaryKey;type:varchar(64)"`
|
||||
RunID string `gorm:"index;type:varchar(64)"`
|
||||
SecretHash string `gorm:"uniqueIndex;type:varchar(64)"`
|
||||
ExpiresAt int64
|
||||
StartedAt int64
|
||||
FinishedAt int64
|
||||
State string `gorm:"type:varchar(24)"`
|
||||
Usage string `gorm:"type:text"`
|
||||
Sample string `gorm:"type:text"`
|
||||
Quota int
|
||||
Available *bool
|
||||
Error string
|
||||
}
|
||||
|
||||
func DetectionSecretHash(secret string) string {
|
||||
sum := sha256.Sum256([]byte(secret))
|
||||
return hex.EncodeToString(sum[:])
|
||||
}
|
||||
|
||||
func CreateDetection(run *DetectionRun) error {
|
||||
if run.Reserved < 0 || run.Reserved > common.MaxWalletQuota {
|
||||
return errors.New("invalid reservation")
|
||||
}
|
||||
original := *run
|
||||
err := relayAttemptTransaction(func(tx *gorm.DB) error {
|
||||
*run = original
|
||||
if run.UserID > 0 {
|
||||
var user User
|
||||
if err := lockForUpdate(tx).First(&user, run.UserID).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
var previous DetectionRun
|
||||
if err := tx.Where("request_key = ?", run.RequestKey).First(&previous).Error; err == nil {
|
||||
if previous.UserID != run.UserID || previous.GroupName != run.GroupName || previous.Kind != run.Kind {
|
||||
return errors.New("idempotency key conflict")
|
||||
}
|
||||
*run = previous
|
||||
return nil
|
||||
} else if !errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return err
|
||||
}
|
||||
var active int64
|
||||
if err := tx.Model(&DetectionRun{}).Where("user_id = ? AND settled = ?", run.UserID, false).Count(&active).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if active > 0 {
|
||||
return errors.New("a detection is already running")
|
||||
}
|
||||
if err := tx.Model(&DetectionRun{}).Where("user_id = ? AND group_name = ? AND kind = ? AND created_at > ?", run.UserID, run.GroupName, run.Kind, time.Now().Unix()-300).Count(&active).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if active > 0 {
|
||||
return errors.New("please wait five minutes between detections")
|
||||
}
|
||||
if user.Status != common.UserStatusEnabled || user.Quota < run.Reserved {
|
||||
return errors.New("insufficient wallet balance")
|
||||
}
|
||||
if err := tx.Model(&User{}).Where("id = ?", run.UserID).Update("quota", gorm.Expr("quota - ?", run.Reserved)).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return tx.Create(run).Error
|
||||
})
|
||||
if err == nil && run.UserID > 0 {
|
||||
_ = invalidateUserCache(run.UserID)
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
func ClaimDetectionCall(hash string) (*DetectionCall, error) {
|
||||
var call DetectionCall
|
||||
err := relayAttemptTransaction(func(tx *gorm.DB) error {
|
||||
if err := tx.Where("secret_hash = ?", hash).First(&call).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
var run DetectionRun
|
||||
if err := lockForUpdate(tx).First(&run, "id = ?", call.RunID).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if err := lockForUpdate(tx).Where("id = ?", call.ID).First(&call).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
now := time.Now().Unix()
|
||||
if run.Kind != "pelican" || call.State != "ready" || call.ExpiresAt <= now || run.Settled || run.Status == "failed" || run.Status == "succeeded" || run.Deadline-10 <= now {
|
||||
return errors.New("expired or consumed credential")
|
||||
}
|
||||
update := tx.Model(&DetectionCall{}).Where("id = ? AND state = ?", call.ID, "ready").Updates(map[string]any{"state": "running", "started_at": now})
|
||||
if update.Error != nil {
|
||||
return update.Error
|
||||
}
|
||||
if update.RowsAffected != 1 {
|
||||
return errors.New("credential already consumed")
|
||||
}
|
||||
call.State = "running"
|
||||
call.StartedAt = now
|
||||
return nil
|
||||
})
|
||||
return &call, err
|
||||
}
|
||||
|
||||
func FinishDetectionCall(call *DetectionCall) error {
|
||||
if call.Quota < 0 || call.Quota > common.MaxWalletQuota {
|
||||
return errors.New("invalid charge")
|
||||
}
|
||||
return DB.Model(&DetectionCall{}).Where("id = ? AND state = ?", call.ID, "running").Updates(map[string]any{
|
||||
"state": "finished", "finished_at": time.Now().Unix(), "usage": call.Usage, "sample": call.Sample, "quota": call.Quota, "available": call.Available, "error": call.Error,
|
||||
}).Error
|
||||
}
|
||||
|
||||
func SettleDetection(id string) error {
|
||||
var owner int
|
||||
err := relayAttemptTransaction(func(tx *gorm.DB) error {
|
||||
var run DetectionRun
|
||||
if err := lockForUpdate(tx).First(&run, "id = ?", id).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
owner = run.UserID
|
||||
if run.Settled {
|
||||
return nil
|
||||
}
|
||||
var calls []DetectionCall
|
||||
if err := tx.Where("run_id = ?", id).Find(&calls).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
charge := 0
|
||||
for _, call := range calls {
|
||||
if call.State == "running" {
|
||||
return errors.New("generation still running")
|
||||
}
|
||||
if call.Quota < 0 || call.Quota > run.Reserved-charge {
|
||||
return errors.New("charge exceeds reservation")
|
||||
}
|
||||
charge += call.Quota
|
||||
}
|
||||
logJSON := ""
|
||||
if owner > 0 {
|
||||
var user User
|
||||
if err := lockForUpdate(tx).Unscoped().First(&user, owner).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if user.Quota > common.MaxWalletQuota-(run.Reserved-charge) || user.UsedQuota > common.MaxWalletQuota-charge {
|
||||
return errors.New("wallet settlement exceeds quota limit")
|
||||
}
|
||||
if err := tx.Unscoped().Model(&User{}).Where("id = ?", owner).Updates(map[string]any{"quota": gorm.Expr("quota + ?", run.Reserved-charge), "used_quota": gorm.Expr("used_quota + ?", charge)}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
key := "detection:" + id
|
||||
other, _ := common.Marshal(map[string]any{"billing_source": "wallet", "detection_id": id, "detection_kind": run.Kind, "reserved": run.Reserved, "refunded": run.Reserved - charge})
|
||||
entry := Log{AttemptKey: &key, UserId: owner, Username: user.Username, CreatedAt: time.Now().Unix(), Type: LogTypeSystem, Content: fmt.Sprintf("%s detection: charged %d quota; refunded %d quota", run.Kind, charge, run.Reserved-charge), Quota: charge, Group: run.GroupName, ModelName: "gpt-6-astra", RequestId: id, Other: string(other)}
|
||||
encoded, err := common.Marshal(entry)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
logJSON = string(encoded)
|
||||
} else {
|
||||
charge = 0
|
||||
}
|
||||
if err := tx.Model(&DetectionCall{}).Where("run_id = ? AND state = ?", id, "ready").Update("state", "revoked").Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Model(&DetectionRun{}).Where("id = ?", id).Updates(map[string]any{"settled": true, "charged": charge, "finished_at": time.Now().Unix(), "log_json": logJSON}).Error
|
||||
})
|
||||
if err == nil && owner > 0 {
|
||||
_ = invalidateUserCache(owner)
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
// Detection ledger delivery is idempotent and separate from request metrics.
|
||||
func DeliverDetectionLogs() error {
|
||||
var runs []DetectionRun
|
||||
if err := DB.Where("settled = ? AND log_delivered = ? AND log_json <> ?", true, false, "").Limit(100).Find(&runs).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
for _, run := range runs {
|
||||
var entry Log
|
||||
if err := common.UnmarshalJsonStr(run.LogJSON, &entry); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := LOG_DB.Clauses(clause.OnConflict{Columns: []clause.Column{{Name: "attempt_key"}}, DoNothing: true}).Create(&entry).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if err := DB.Model(&DetectionRun{}).Where("id = ?", run.ID).Update("log_delivered", true).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,61 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
)
|
||||
|
||||
// DetectionSchedule survives task-history cleanup and is never projected directly
|
||||
// to public APIs. Only automatic results advance this channel/kind clock.
|
||||
type DetectionSchedule struct {
|
||||
ChannelID int `gorm:"primaryKey;autoIncrement:false"`
|
||||
Kind string `gorm:"primaryKey;type:varchar(24)"`
|
||||
LastValid string `gorm:"type:varchar(16)"`
|
||||
IntervalMinutes int
|
||||
NextDueAt int64
|
||||
LastFinishedAt int64
|
||||
LastResultID string `gorm:"type:varchar(64)"`
|
||||
}
|
||||
|
||||
func NextDetectionInterval(lastValid, outcome string, current int) int {
|
||||
if outcome != "normal" && outcome != "degraded" {
|
||||
return 15
|
||||
}
|
||||
if lastValid == "" {
|
||||
return 30
|
||||
}
|
||||
if lastValid != outcome {
|
||||
return 15
|
||||
}
|
||||
if current == 15 {
|
||||
return 30
|
||||
}
|
||||
return 60
|
||||
}
|
||||
|
||||
func EnsureDetectionSchedule(ctx context.Context, seed DetectionSchedule) error {
|
||||
return DB.WithContext(ctx).Clauses(clause.OnConflict{DoNothing: true}).Create(&seed).Error
|
||||
}
|
||||
|
||||
func ApplyDetectionScheduleResult(ctx context.Context, channelID int, kind, resultID string, finishedAt int64, outcome string) error {
|
||||
return DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
var state DetectionSchedule
|
||||
if err := lockForUpdate(tx).Where("channel_id = ? AND kind = ?", channelID, kind).First(&state).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
// A result is committed before its clock advances. Reconciliation can
|
||||
// safely replay it after a crash, without advancing the clock twice.
|
||||
if finishedAt <= state.LastFinishedAt || resultID == state.LastResultID {
|
||||
return nil
|
||||
}
|
||||
state.IntervalMinutes = NextDetectionInterval(state.LastValid, outcome, state.IntervalMinutes)
|
||||
if outcome == "normal" || outcome == "degraded" {
|
||||
state.LastValid = outcome
|
||||
}
|
||||
state.LastFinishedAt, state.LastResultID = finishedAt, resultID
|
||||
state.NextDueAt = finishedAt + int64(state.IntervalMinutes)*60
|
||||
return tx.Save(&state).Error
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,138 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"os"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"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"
|
||||
)
|
||||
|
||||
func TestAdaptiveDetectionIntervals(t *testing.T) {
|
||||
state := DetectionSchedule{IntervalMinutes: 30}
|
||||
for _, step := range []struct {
|
||||
result string
|
||||
interval int
|
||||
}{
|
||||
{"normal", 30}, {"normal", 60}, {"normal", 60}, {"degraded", 15}, {"normal", 15}, {"normal", 30}, {"normal", 60},
|
||||
{"unknown", 15}, {"unknown", 15}, {"normal", 30}, {"normal", 60}, {"failed", 15}, {"degraded", 15}, {"degraded", 30}, {"degraded", 60},
|
||||
} {
|
||||
state.IntervalMinutes = NextDetectionInterval(state.LastValid, step.result, state.IntervalMinutes)
|
||||
assert.Equal(t, step.interval, state.IntervalMinutes, step.result)
|
||||
if step.result == "normal" || step.result == "degraded" {
|
||||
state.LastValid = step.result
|
||||
}
|
||||
}
|
||||
assert.Equal(t, 15, NextDetectionInterval("", "unknown", 30))
|
||||
assert.Equal(t, 30, NextDetectionInterval("", "degraded", 15))
|
||||
}
|
||||
|
||||
func TestAdaptiveDetectionDatabaseMatrix(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("adaptive_%d_", time.Now().UnixNano())}})
|
||||
require.NoError(t, err)
|
||||
conn, err := db.DB()
|
||||
require.NoError(t, err)
|
||||
conn.SetMaxOpenConns(1)
|
||||
if dialect != "sqlite" {
|
||||
conn.SetMaxOpenConns(4)
|
||||
}
|
||||
old := DB
|
||||
DB = db
|
||||
oldType := common.MainDatabaseType()
|
||||
common.SetMainDatabaseType(common.DatabaseType(dialect))
|
||||
t.Cleanup(func() { common.SetMainDatabaseType(oldType) })
|
||||
t.Cleanup(func() {
|
||||
for _, row := range []any{&DetectionSchedule{}, &DetectionRun{}, &SystemTask{}} {
|
||||
require.NoError(t, db.Migrator().DropTable(row))
|
||||
}
|
||||
DB = old
|
||||
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.Logf("database version: %s", version)
|
||||
// Fresh schema, then the exact pre-upgrade task models with real rows.
|
||||
require.NoError(t, db.AutoMigrate(&DetectionSchedule{}))
|
||||
require.NoError(t, db.Migrator().DropTable(&DetectionSchedule{}))
|
||||
require.NoError(t, db.AutoMigrate(&DetectionRun{}, &SystemTask{}))
|
||||
oldRun := DetectionRun{ID: "old", RequestKey: "old-unique", ChannelID: 12, Kind: "pelican", Status: "succeeded", Quality: "normal", Settled: true, FinishedAt: 100, Charged: 17}
|
||||
require.NoError(t, db.Create(&oldRun).Error)
|
||||
require.NoError(t, db.Create(&SystemTask{TaskID: "old-task", Type: "astra_fingerprint", State: `{"rounds":[]}`}).Error)
|
||||
for range 2 {
|
||||
require.NoError(t, db.AutoMigrate(&DetectionRun{}, &SystemTask{}, &DetectionSchedule{}))
|
||||
}
|
||||
var retained DetectionRun
|
||||
require.NoError(t, db.First(&retained, "id = ?", "old").Error)
|
||||
assert.Equal(t, 17, retained.Charged)
|
||||
duplicate := oldRun
|
||||
duplicate.ID = "duplicate"
|
||||
require.Error(t, db.Create(&duplicate).Error)
|
||||
ctx := context.Background()
|
||||
seed := DetectionSchedule{ChannelID: 12, Kind: "pelican", IntervalMinutes: 30}
|
||||
require.NoError(t, EnsureDetectionSchedule(ctx, seed))
|
||||
require.NoError(t, EnsureDetectionSchedule(ctx, seed))
|
||||
require.NoError(t, ApplyDetectionScheduleResult(ctx, 12, "pelican", "one", 100, "normal"))
|
||||
require.NoError(t, ApplyDetectionScheduleResult(ctx, 12, "pelican", "two", 1900, "normal"))
|
||||
var wg sync.WaitGroup
|
||||
errs := make(chan error, 2)
|
||||
for range 2 {
|
||||
wg.Go(func() { errs <- ApplyDetectionScheduleResult(ctx, 12, "pelican", "change", 5500, "degraded") })
|
||||
}
|
||||
wg.Wait()
|
||||
close(errs)
|
||||
for err := range errs {
|
||||
require.NoError(t, err)
|
||||
}
|
||||
require.NoError(t, ApplyDetectionScheduleResult(ctx, 12, "pelican", "stale", 1901, "normal"))
|
||||
var loaded DetectionSchedule
|
||||
require.NoError(t, db.First(&loaded, "channel_id = ? AND kind = ?", 12, "pelican").Error)
|
||||
assert.Equal(t, 15, loaded.IntervalMinutes)
|
||||
assert.Equal(t, int64(6400), loaded.NextDueAt)
|
||||
assert.Equal(t, "degraded", loaded.LastValid)
|
||||
seed.Kind = "fingerprint"
|
||||
require.NoError(t, EnsureDetectionSchedule(ctx, seed))
|
||||
require.NoError(t, ApplyDetectionScheduleResult(ctx, 12, "fingerprint", "fp", 5600, "normal"))
|
||||
require.NoError(t, db.Migrator().DropTable(&SystemTask{}))
|
||||
require.NoError(t, db.AutoMigrate(&SystemTask{}, &DetectionSchedule{}))
|
||||
var rows []DetectionSchedule
|
||||
require.NoError(t, db.Find(&rows).Error)
|
||||
require.Len(t, rows, 2)
|
||||
// Reinitializing from history must not overwrite a persisted clock.
|
||||
seed.Kind = "pelican"
|
||||
require.NoError(t, EnsureDetectionSchedule(ctx, seed))
|
||||
require.NoError(t, db.First(&loaded, "channel_id = ? AND kind = ?", 12, "pelican").Error)
|
||||
assert.Equal(t, int64(6400), loaded.NextDueAt)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,142 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"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"
|
||||
"os"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestDetectionDatabaseMatrix(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":
|
||||
dsn := os.Getenv("TEST_MYSQL_DSN")
|
||||
if dsn == "" {
|
||||
t.Skip("TEST_MYSQL_DSN unset")
|
||||
}
|
||||
driver = mysql.Open(dsn)
|
||||
case "postgres":
|
||||
dsn := os.Getenv("TEST_POSTGRES_DSN")
|
||||
if dsn == "" {
|
||||
t.Skip("TEST_POSTGRES_DSN unset")
|
||||
}
|
||||
driver = postgres.Open(dsn)
|
||||
}
|
||||
db, err := gorm.Open(driver, &gorm.Config{NamingStrategy: schema.NamingStrategy{TablePrefix: fmt.Sprintf("det_%d_", time.Now().UnixNano())}})
|
||||
require.NoError(t, err)
|
||||
sqlDB, err := db.DB()
|
||||
require.NoError(t, err)
|
||||
sqlDB.SetMaxOpenConns(1)
|
||||
if dialect != "sqlite" {
|
||||
sqlDB.SetMaxOpenConns(4)
|
||||
}
|
||||
var version string
|
||||
if dialect == "sqlite" {
|
||||
require.NoError(t, db.Raw("select sqlite_version()").Scan(&version).Error)
|
||||
} else {
|
||||
require.NoError(t, db.Raw("select version()").Scan(&version).Error)
|
||||
}
|
||||
t.Logf("database version: %s", version)
|
||||
oldDB, oldLogDB, oldRedis := DB, LOG_DB, common.RedisEnabled
|
||||
DB = db
|
||||
LOG_DB = db
|
||||
common.RedisEnabled = false
|
||||
t.Cleanup(func() {
|
||||
for _, m := range []any{&DetectionCall{}, &DetectionRun{}, &User{}, &Log{}} {
|
||||
require.NoError(t, db.Migrator().DropTable(m))
|
||||
}
|
||||
DB = oldDB
|
||||
LOG_DB = oldLogDB
|
||||
common.RedisEnabled = oldRedis
|
||||
require.NoError(t, sqlDB.Close())
|
||||
})
|
||||
require.NoError(t, db.AutoMigrate(&Log{}))
|
||||
require.NoError(t, db.AutoMigrate(&User{}))
|
||||
user := User{Username: "detection-owner", Quota: 1000, Status: common.UserStatusEnabled}
|
||||
require.NoError(t, db.Create(&user).Error)
|
||||
for range 2 {
|
||||
require.NoError(t, db.AutoMigrate(&User{}, &DetectionRun{}, &DetectionCall{}))
|
||||
}
|
||||
now := time.Now().Unix()
|
||||
run := DetectionRun{ID: "detection", RequestKey: "user-idempotency", UserID: user.Id, Kind: "pelican", GroupName: "a", Reserved: 300, CreatedAt: now, Deadline: now + 600, Status: "pending"}
|
||||
require.NoError(t, CreateDetection(&run))
|
||||
require.NoError(t, CreateDetection(&run))
|
||||
require.NoError(t, db.First(&user, user.Id).Error)
|
||||
assert.Equal(t, 700, user.Quota)
|
||||
second := run
|
||||
second.ID = "second"
|
||||
second.RequestKey = "second"
|
||||
require.Error(t, CreateDetection(&second))
|
||||
call := DetectionCall{ID: "call", RunID: run.ID, SecretHash: DetectionSecretHash("secret"), ExpiresAt: now + 900, State: "ready"}
|
||||
require.NoError(t, db.Create(&call).Error)
|
||||
var wg sync.WaitGroup
|
||||
claimedCalls := make(chan *DetectionCall, 2)
|
||||
for range 2 {
|
||||
wg.Go(func() {
|
||||
result, claimErr := ClaimDetectionCall(call.SecretHash)
|
||||
if claimErr == nil {
|
||||
claimedCalls <- result
|
||||
}
|
||||
})
|
||||
}
|
||||
wg.Wait()
|
||||
close(claimedCalls)
|
||||
require.Len(t, claimedCalls, 1)
|
||||
claimed := <-claimedCalls
|
||||
_, err = ClaimDetectionCall(call.SecretHash)
|
||||
require.Error(t, err)
|
||||
require.Error(t, SettleDetection(run.ID))
|
||||
claimed.Quota = 42
|
||||
claimed.Usage = `{"input_tokens":10,"output_tokens":5}`
|
||||
require.NoError(t, FinishDetectionCall(claimed))
|
||||
claimed.Quota = 99
|
||||
require.NoError(t, FinishDetectionCall(claimed))
|
||||
require.NoError(t, SettleDetection(run.ID))
|
||||
require.NoError(t, SettleDetection(run.ID))
|
||||
require.NoError(t, DeliverDetectionLogs())
|
||||
require.NoError(t, DeliverDetectionLogs())
|
||||
var logCount int64
|
||||
require.NoError(t, db.Model(&Log{}).Count(&logCount).Error)
|
||||
assert.EqualValues(t, 1, logCount)
|
||||
require.NoError(t, db.First(&user, user.Id).Error)
|
||||
assert.Equal(t, 958, user.Quota)
|
||||
assert.Equal(t, 42, user.UsedQuota)
|
||||
_, err = ClaimDetectionCall(call.SecretHash)
|
||||
require.Error(t, err)
|
||||
require.NoError(t, db.AutoMigrate(&DetectionRun{}, &DetectionCall{}))
|
||||
require.NoError(t, db.AutoMigrate(&DetectionRun{}, &DetectionCall{}))
|
||||
var stored DetectionRun
|
||||
require.NoError(t, db.First(&stored, "id = ?", run.ID).Error)
|
||||
assert.Equal(t, 42, stored.Charged)
|
||||
expired := DetectionRun{ID: "expired", RequestKey: "expired", Kind: "pelican", Deadline: now - 1, CreatedAt: now - 901, Status: "pending"}
|
||||
require.NoError(t, CreateDetection(&expired))
|
||||
expCall := DetectionCall{ID: "expired-call", RunID: expired.ID, SecretHash: DetectionSecretHash("expired"), State: "ready", ExpiresAt: now - 1}
|
||||
require.NoError(t, db.Create(&expCall).Error)
|
||||
_, err = ClaimDetectionCall(expCall.SecretHash)
|
||||
require.Error(t, err)
|
||||
require.NoError(t, SettleDetection(expired.ID))
|
||||
deleted := User{Username: "deleted-owner", AffCode: "deleted-owner", Quota: 500, Status: common.UserStatusEnabled}
|
||||
require.NoError(t, db.Create(&deleted).Error)
|
||||
deletedRun := DetectionRun{ID: "deleted", RequestKey: "deleted", UserID: deleted.Id, Kind: "pelican", Reserved: 100, CreatedAt: now}
|
||||
require.NoError(t, CreateDetection(&deletedRun))
|
||||
require.NoError(t, db.Delete(&deleted).Error)
|
||||
require.NoError(t, SettleDetection(deletedRun.ID))
|
||||
require.NoError(t, db.Unscoped().First(&deleted, deleted.Id).Error)
|
||||
assert.Equal(t, 500, deleted.Quota)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,51 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"os"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/driver/postgres"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// Opt-in only: the deployment harness restores a snapshot into an isolated
|
||||
// disposable database. No workers, relay calls or background jobs are started.
|
||||
func TestAdaptiveDeploymentUpgradeFromSnapshot(t *testing.T) {
|
||||
dsn := os.Getenv("ADAPTIVE_UPGRADE_DSN")
|
||||
if dsn == "" {
|
||||
t.Skip("isolated upgrade snapshot unset")
|
||||
}
|
||||
require.True(t, strings.HasSuffix(dsn, "/adaptive_upgrade34?sslmode=disable"))
|
||||
db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
var before, after struct {
|
||||
Users int64
|
||||
Quota int64
|
||||
Used int64
|
||||
Channels int64
|
||||
Runs int64
|
||||
}
|
||||
query := "SELECT (SELECT count(*) FROM users) AS users, (SELECT sum(quota) FROM users) AS quota, (SELECT sum(used_quota) FROM users) AS used, (SELECT count(*) FROM channels) AS channels, (SELECT count(*) FROM detection_runs) AS runs"
|
||||
require.NoError(t, db.Raw(query).Scan(&before).Error)
|
||||
conn, err := db.DB()
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, conn.Close())
|
||||
oldDB, oldMaster, oldMain, oldLog := DB, common.IsMasterNode, common.MainDatabaseType(), common.LogDatabaseType()
|
||||
t.Cleanup(func() { DB = oldDB; common.IsMasterNode = oldMaster; common.SetDatabaseTypes(oldMain, oldLog) })
|
||||
common.IsMasterNode = true
|
||||
t.Setenv("SQL_DSN", dsn)
|
||||
for range 2 {
|
||||
require.NoError(t, InitDB())
|
||||
require.True(t, DB.Migrator().HasTable(&DetectionSchedule{}))
|
||||
require.NoError(t, DB.Raw(query).Scan(&after).Error)
|
||||
assert.Equal(t, before, after)
|
||||
conn, err = DB.DB()
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, conn.Close())
|
||||
}
|
||||
t.Logf("snapshot preserved: users=%d channels=%d runs=%d; wallet totals unchanged", before.Users, before.Channels, before.Runs)
|
||||
}
|
||||
@@ -0,0 +1,37 @@
|
||||
package model
|
||||
|
||||
// Deprecated model compatibility: OpenAI removed DALL·E 2 and DALL·E 3 on
|
||||
// 2026-05-12. Keep conversion of existing prices available for compatible
|
||||
// upstreams, deriving every multiplier from the legacy DTO billing logic.
|
||||
// https://developers.openai.com/api/docs/deprecations#2025-11-14-dalle-model-snapshots
|
||||
|
||||
import (
|
||||
"strings"
|
||||
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
)
|
||||
|
||||
func legacyDallePricingRules(name string) []LegacyPricingRule {
|
||||
if !strings.HasPrefix(name, "dall-e") {
|
||||
return nil
|
||||
}
|
||||
var rules []LegacyPricingRule
|
||||
if name != "dall-e-3" {
|
||||
for _, size := range []string{"256x256", "512x512"} {
|
||||
request := dto.ImageRequest{Model: name, Size: size}
|
||||
rules = append(rules, LegacyPricingRule{`param("size") == "` + size + `"`, request.GetTokenCountMeta().ImagePriceRatio})
|
||||
}
|
||||
}
|
||||
rectangle := dto.ImageRequest{Model: name, Size: "1024x1792"}
|
||||
if name != "dall-e-2" && name != "dall-e" {
|
||||
rules = append(rules, LegacyPricingRule{`param("size") == "1024x1792" || param("size") == "1792x1024"`, rectangle.GetTokenCountMeta().ImagePriceRatio})
|
||||
}
|
||||
if name == "dall-e-3" {
|
||||
squareHD := dto.ImageRequest{Model: name, Size: "1024x1024", Quality: "hd"}
|
||||
rectangleHD := dto.ImageRequest{Model: name, Size: rectangle.Size, Quality: "hd"}
|
||||
rules = append(rules,
|
||||
LegacyPricingRule{`param("quality") == "hd" && param("size") != "1024x1792" && param("size") != "1792x1024"`, squareHD.GetTokenCountMeta().ImagePriceRatio},
|
||||
LegacyPricingRule{`param("quality") == "hd" && (param("size") == "1024x1792" || param("size") == "1792x1024")`, rectangleHD.GetTokenCountMeta().ImagePriceRatio / rectangle.GetTokenCountMeta().ImagePriceRatio})
|
||||
}
|
||||
return rules
|
||||
}
|
||||
+63
-24
@@ -57,27 +57,29 @@ func sanitizeClickHouseLikePattern(input string) (string, error) {
|
||||
}
|
||||
|
||||
type Log struct {
|
||||
Id int `json:"id" gorm:"index:idx_created_at_id,priority:2;index:idx_user_id_id,priority:2"`
|
||||
UserId int `json:"user_id" gorm:"index;index:idx_user_id_id,priority:1"`
|
||||
CreatedAt int64 `json:"created_at" gorm:"bigint;index:idx_created_at_id,priority:1;index:idx_created_at_type"`
|
||||
Type int `json:"type" gorm:"index:idx_created_at_type"`
|
||||
Content string `json:"content"`
|
||||
Username string `json:"username" gorm:"index;index:index_username_model_name,priority:2;default:''"`
|
||||
TokenName string `json:"token_name" gorm:"index;default:''"`
|
||||
ModelName string `json:"model_name" gorm:"index;index:index_username_model_name,priority:1;default:''"`
|
||||
Quota int `json:"quota" gorm:"default:0"`
|
||||
PromptTokens int `json:"prompt_tokens" gorm:"default:0"`
|
||||
CompletionTokens int `json:"completion_tokens" gorm:"default:0"`
|
||||
UseTime int `json:"use_time" gorm:"default:0"`
|
||||
IsStream bool `json:"is_stream"`
|
||||
ChannelId int `json:"channel" gorm:"index"`
|
||||
ChannelName string `json:"channel_name" gorm:"->"`
|
||||
TokenId int `json:"token_id" gorm:"default:0;index"`
|
||||
Group string `json:"group" gorm:"index"`
|
||||
Ip string `json:"ip" gorm:"index;default:''"`
|
||||
RequestId string `json:"request_id,omitempty" gorm:"type:varchar(64);index:idx_logs_request_id;default:''"`
|
||||
UpstreamRequestId string `json:"upstream_request_id,omitempty" gorm:"type:varchar(128);index:idx_logs_upstream_request_id;default:''"`
|
||||
Other string `json:"other"`
|
||||
AttemptKey *string `json:"attempt_key,omitempty" gorm:"type:varchar(96);uniqueIndex"`
|
||||
Id int `json:"id" gorm:"index:idx_created_at_id,priority:2;index:idx_user_id_id,priority:2"`
|
||||
UserId int `json:"user_id" gorm:"index;index:idx_user_id_id,priority:1"`
|
||||
CreatedAt int64 `json:"created_at" gorm:"bigint;index:idx_created_at_id,priority:1;index:idx_created_at_type"`
|
||||
Type int `json:"type" gorm:"index:idx_created_at_type"`
|
||||
Content string `json:"content"`
|
||||
Username string `json:"username" gorm:"index;index:index_username_model_name,priority:2;default:''"`
|
||||
TokenName string `json:"token_name" gorm:"index;default:''"`
|
||||
ModelName string `json:"model_name" gorm:"index;index:index_username_model_name,priority:1;default:''"`
|
||||
Quota int `json:"quota" gorm:"default:0"`
|
||||
PromptTokens int `json:"prompt_tokens" gorm:"default:0"`
|
||||
CompletionTokens int `json:"completion_tokens" gorm:"default:0"`
|
||||
UseTime int `json:"use_time" gorm:"default:0"`
|
||||
IsStream bool `json:"is_stream"`
|
||||
ChannelId int `json:"channel" gorm:"index"`
|
||||
ChannelName string `json:"channel_name" gorm:"->"`
|
||||
TokenId int `json:"token_id" gorm:"default:0;index"`
|
||||
Group string `json:"group" gorm:"index"`
|
||||
Ip string `json:"ip" gorm:"index;default:''"`
|
||||
RequestId string `json:"request_id,omitempty" gorm:"type:varchar(64);index:idx_logs_request_id;default:''"`
|
||||
UpstreamRequestId string `json:"upstream_request_id,omitempty" gorm:"type:varchar(128);index:idx_logs_upstream_request_id;default:''"`
|
||||
CorrelationKey string `json:"-" gorm:"type:varchar(64);index:idx_logs_correlation;default:''"`
|
||||
Other string `json:"other"`
|
||||
}
|
||||
|
||||
// don't use iota, avoid change log type value
|
||||
@@ -116,6 +118,7 @@ func assignDisplayLogIds(logs []*Log, startIdx int) {
|
||||
func formatUserLogs(logs []*Log, startIdx int) {
|
||||
for i := range logs {
|
||||
logs[i].ChannelName = ""
|
||||
normalizeErrorLogGroup(logs[i])
|
||||
logs[i].Other = formatLogOtherJSON(logs[i].Other, logOtherVisibilityUser)
|
||||
}
|
||||
assignDisplayLogIds(logs, startIdx)
|
||||
@@ -125,6 +128,7 @@ func formatUserLogs(logs []*Log, startIdx int) {
|
||||
// admin_info. Root callers must not pass their results through this formatter.
|
||||
func FormatAdminLogs(logs []*Log) {
|
||||
for i := range logs {
|
||||
normalizeErrorLogGroup(logs[i])
|
||||
logs[i].Other = formatLogOtherJSON(logs[i].Other, logOtherVisibilityAdmin)
|
||||
}
|
||||
}
|
||||
@@ -133,6 +137,7 @@ func FormatAdminLogs(logs []*Log) {
|
||||
// without removing root-only diagnostics.
|
||||
func FormatRootLogs(logs []*Log) {
|
||||
for i := range logs {
|
||||
normalizeErrorLogGroup(logs[i])
|
||||
logs[i].Other = formatLogOtherJSON(logs[i].Other, logOtherVisibilityRoot)
|
||||
}
|
||||
}
|
||||
@@ -313,8 +318,10 @@ func RecordErrorLog(c *gin.Context, userId int, channelId int, modelName string,
|
||||
}(),
|
||||
RequestId: requestId,
|
||||
UpstreamRequestId: upstreamRequestId,
|
||||
CorrelationKey: c.GetString("log_correlation_key"),
|
||||
Other: otherStr,
|
||||
}
|
||||
normalizeErrorLogGroup(log)
|
||||
err := createLog(log)
|
||||
if err != nil {
|
||||
logger.LogError(c, "failed to record log: "+err.Error())
|
||||
@@ -377,6 +384,7 @@ func RecordConsumeLog(c *gin.Context, userId int, params RecordConsumeLogParams)
|
||||
}(),
|
||||
RequestId: requestId,
|
||||
UpstreamRequestId: upstreamRequestId,
|
||||
CorrelationKey: c.GetString("log_correlation_key"),
|
||||
Other: otherStr,
|
||||
}
|
||||
err := createLog(log)
|
||||
@@ -461,7 +469,7 @@ func RecordTaskBillingLog(params RecordTaskBillingLogParams) {
|
||||
}
|
||||
}
|
||||
|
||||
func GetAllLogs(logType int, startTimestamp int64, endTimestamp int64, modelName string, username string, tokenName string, startIdx int, num int, channel int, group string, requestId string, upstreamRequestId string) (logs []*Log, total int64, err error) {
|
||||
func GetAllLogs(logType int, startTimestamp int64, endTimestamp int64, modelName string, username string, tokenName string, startIdx int, num int, channel int, group string, requestId string, upstreamRequestId string, correlationKey ...string) (logs []*Log, total int64, err error) {
|
||||
var tx *gorm.DB
|
||||
if logType == LogTypeUnknown {
|
||||
tx = LOG_DB
|
||||
@@ -478,6 +486,12 @@ func GetAllLogs(logType int, startTimestamp int64, endTimestamp int64, modelName
|
||||
if tokenName != "" {
|
||||
tx = tx.Where("logs.token_name = ?", tokenName)
|
||||
}
|
||||
if len(correlationKey) > 1 {
|
||||
tx = applyAttemptLogFilter(tx, correlationKey[1])
|
||||
}
|
||||
if len(correlationKey) > 0 && correlationKey[0] != "" {
|
||||
tx = tx.Where("logs.correlation_key = ?", correlationKey[0])
|
||||
}
|
||||
if requestId != "" {
|
||||
tx = tx.Where("logs.request_id = ?", requestId)
|
||||
}
|
||||
@@ -557,7 +571,7 @@ func GetAllLogs(logType int, startTimestamp int64, endTimestamp int64, modelName
|
||||
|
||||
const logSearchCountLimit = 10000
|
||||
|
||||
func GetUserLogs(userId int, logType int, startTimestamp int64, endTimestamp int64, modelName string, tokenName string, startIdx int, num int, group string, requestId string, upstreamRequestId string) (logs []*Log, total int64, err error) {
|
||||
func GetUserLogs(userId int, logType int, startTimestamp int64, endTimestamp int64, modelName string, tokenName string, startIdx int, num int, group string, requestId string, upstreamRequestId string, attemptStatus ...string) (logs []*Log, total int64, err error) {
|
||||
var tx *gorm.DB
|
||||
if logType == LogTypeUnknown {
|
||||
tx = LOG_DB.Where("logs.user_id = ?", userId)
|
||||
@@ -565,6 +579,10 @@ func GetUserLogs(userId int, logType int, startTimestamp int64, endTimestamp int
|
||||
tx = LOG_DB.Where("logs.user_id = ? and logs.type = ?", userId, logType)
|
||||
}
|
||||
|
||||
if len(attemptStatus) > 0 {
|
||||
tx = applyAttemptLogFilter(tx, attemptStatus[0])
|
||||
}
|
||||
|
||||
if tx, err = applyExplicitLogTextFilter(tx, "logs.model_name", modelName); err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
@@ -611,11 +629,15 @@ type Stat struct {
|
||||
Tpm int `json:"tpm"`
|
||||
}
|
||||
|
||||
func SumUsedQuota(logType int, startTimestamp int64, endTimestamp int64, modelName string, username string, tokenName string, channel int, group string) (stat Stat, err error) {
|
||||
func SumUsedQuota(logType int, startTimestamp int64, endTimestamp int64, modelName string, username string, tokenName string, channel int, group string, correlationKey ...string) (stat Stat, err error) {
|
||||
tx := LOG_DB.Table("logs").Select("COALESCE(sum(quota), 0) quota")
|
||||
|
||||
// 为rpm和tpm创建单独的查询
|
||||
rpmTpmQuery := LOG_DB.Table("logs").Select("count(*) rpm, COALESCE(sum(prompt_tokens), 0) + COALESCE(sum(completion_tokens), 0) tpm")
|
||||
if len(correlationKey) > 0 && correlationKey[0] != "" {
|
||||
tx = tx.Where("correlation_key = ?", correlationKey[0])
|
||||
rpmTpmQuery = rpmTpmQuery.Where("correlation_key = ?", correlationKey[0])
|
||||
}
|
||||
|
||||
if tx, err = applyExplicitLogTextFilter(tx, "username", username); err != nil {
|
||||
return stat, err
|
||||
@@ -651,6 +673,11 @@ func SumUsedQuota(logType int, startTimestamp int64, endTimestamp int64, modelNa
|
||||
tx = tx.Where("type = ?", LogTypeConsume)
|
||||
rpmTpmQuery = rpmTpmQuery.Where("type = ?", LogTypeConsume)
|
||||
|
||||
if len(correlationKey) > 1 {
|
||||
tx = applyAttemptLogFilter(tx, correlationKey[1])
|
||||
rpmTpmQuery = applyAttemptLogFilter(rpmTpmQuery, correlationKey[1])
|
||||
}
|
||||
|
||||
// 只统计最近60秒的rpm和tpm
|
||||
rpmTpmQuery = rpmTpmQuery.Where("created_at >= ?", time.Now().Add(-60*time.Second).Unix())
|
||||
|
||||
@@ -737,3 +764,15 @@ func DeleteOldLogBatch(ctx context.Context, targetTimestamp int64, limit int) (i
|
||||
}
|
||||
return result.RowsAffected, nil
|
||||
}
|
||||
|
||||
// A fixed top-level marker keeps filtering portable across SQL log stores.
|
||||
func applyAttemptLogFilter(tx *gorm.DB, status string) *gorm.DB {
|
||||
switch status {
|
||||
case "verified", "unchecked", "rejected", "charged_rejection", "switched_success", "mismatch_allowed", "transport_error":
|
||||
return tx.Where("other LIKE ?", `%"routing_status":"`+status+`"%`)
|
||||
case "estimated":
|
||||
return tx.Where("other LIKE ?", `%"attempt_usage_source":"estimated"%`)
|
||||
default:
|
||||
return tx
|
||||
}
|
||||
}
|
||||
@@ -9,6 +9,60 @@ import (
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestRequestDetailsUserProjection(t *testing.T) {
|
||||
other := NewLogOther()
|
||||
other.SetPublic("model_ratio", 1)
|
||||
other.SetAdmin("request_details", map[string]any{"session_id": "private-session", "correlation_key": "private-key"})
|
||||
other.SetAdmin("request_policy", []any{map[string]any{"candidates": []string{"private-channel"}}})
|
||||
log := &Log{CorrelationKey: "private-key", Other: other.JSONString()}
|
||||
formatUserLogs([]*Log{log}, 0)
|
||||
encoded, err := common.Marshal(log)
|
||||
require.NoError(t, err)
|
||||
require.NotContains(t, string(encoded), "private-")
|
||||
require.NotContains(t, string(encoded), "request_details")
|
||||
require.NotContains(t, string(encoded), "request_policy")
|
||||
require.Contains(t, log.Other, "model_ratio")
|
||||
}
|
||||
|
||||
func TestAstraRejectionUserProjectionShowsModelWithoutRoute(t *testing.T) {
|
||||
other := NewLogOther()
|
||||
other.SetPublic("error_code", "response_model_rejected")
|
||||
other.SetAdmin("request_details", map[string]any{
|
||||
"session_id": "private-session",
|
||||
"model_guard": map[string]any{"attempts": []map[string]any{
|
||||
{"expected_model": "gpt-6-astra", "returned_model": "gpt-5.6-luna", "reason": "model_mismatch", "channel_id": 1},
|
||||
{"expected_model": "gpt-6-astra", "returned_model": "gpt-5.6-luna", "reason": "model_mismatch", "channel_id": 1},
|
||||
}},
|
||||
})
|
||||
log := &Log{Other: other.JSONString()}
|
||||
formatUserLogs([]*Log{log}, 0)
|
||||
var projected map[string]any
|
||||
require.NoError(t, common.UnmarshalJsonStr(log.Other, &projected))
|
||||
assert.Equal(t, "model_mismatch", projected["model_guard_result"])
|
||||
assert.Equal(t, "gpt-5.6-luna", projected["response_model"].(map[string]any)["returned_model"])
|
||||
assert.NotContains(t, log.Other, "admin_info")
|
||||
assert.NotContains(t, log.Other, "private-session")
|
||||
assert.NotContains(t, log.Other, "channel_id")
|
||||
|
||||
other.SetAdmin("request_details", map[string]any{"model_guard": map[string]any{"attempts": []map[string]any{
|
||||
{"expected_model": "gpt-6-astra", "reason": "confirmation_timeout", "channel_id": 1},
|
||||
}}})
|
||||
log.Other = other.JSONString()
|
||||
formatUserLogs([]*Log{log}, 0)
|
||||
projected = nil
|
||||
require.NoError(t, common.UnmarshalJsonStr(log.Other, &projected))
|
||||
assert.Equal(t, "confirmation_timeout", projected["model_guard_result"])
|
||||
assert.NotContains(t, projected, "response_model")
|
||||
|
||||
other.SetPublic("error_code", "upstream_error")
|
||||
log.Other = other.JSONString()
|
||||
formatUserLogs([]*Log{log}, 0)
|
||||
projected = nil
|
||||
require.NoError(t, common.UnmarshalJsonStr(log.Other, &projected))
|
||||
assert.NotContains(t, projected, "model_guard_result")
|
||||
assert.NotContains(t, projected, "response_model")
|
||||
}
|
||||
|
||||
// TestFormatUserLogsStripsQuotaSaturation verifies the admin-only quota
|
||||
// saturation marker (nested under other.admin_info) is removed for non-admin
|
||||
// log views, since formatUserLogs strips the whole admin_info object.
|
||||
|
||||
+53
-1
@@ -209,6 +209,57 @@ func normalizeLegacyRejectReason(values map[string]json.RawMessage) bool {
|
||||
return true
|
||||
}
|
||||
|
||||
// Expose only the provider's model declaration and a bounded reason to the
|
||||
// log owner. The full attempt trail remains under admin_info.
|
||||
func projectResponseModelRejection(values map[string]json.RawMessage) bool {
|
||||
var code string
|
||||
if err := common.Unmarshal(values["error_code"], &code); err != nil || code != "response_model_rejected" {
|
||||
return false
|
||||
}
|
||||
var admin struct {
|
||||
RequestDetails struct {
|
||||
ModelGuard struct {
|
||||
Attempts []struct {
|
||||
Expected string `json:"expected_model"`
|
||||
Returned string `json:"returned_model"`
|
||||
Reason string `json:"reason"`
|
||||
} `json:"attempts"`
|
||||
} `json:"model_guard"`
|
||||
} `json:"request_details"`
|
||||
}
|
||||
if err := common.Unmarshal(values[logOtherAdminInfoKey], &admin); err != nil || len(admin.RequestDetails.ModelGuard.Attempts) == 0 {
|
||||
return false
|
||||
}
|
||||
attempts := admin.RequestDetails.ModelGuard.Attempts
|
||||
last := attempts[len(attempts)-1]
|
||||
reason := last.Reason
|
||||
if reason != "confirmation_timeout" && reason != "missing_model" && reason != "confirmation_buffer_limit" && reason != "model_mismatch" {
|
||||
return false
|
||||
}
|
||||
result, err := common.Marshal(reason)
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
values["model_guard_result"] = result
|
||||
for i := len(attempts) - 1; i >= 0; i-- {
|
||||
entry := attempts[i]
|
||||
if entry.Returned == "" || entry.Expected == "" {
|
||||
continue
|
||||
}
|
||||
observation, err := common.Marshal(map[string]any{
|
||||
"requested_model": entry.Expected,
|
||||
"upstream_model": entry.Expected,
|
||||
"returned_model": entry.Returned,
|
||||
"mismatch": entry.Returned != entry.Expected,
|
||||
})
|
||||
if err == nil {
|
||||
values["response_model"] = observation
|
||||
}
|
||||
break
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// formatLogOtherJSON applies the role projection while keeping untouched JSON
|
||||
// values as RawMessage. This preserves integers larger than JavaScript's safe
|
||||
// range instead of round-tripping them through float64.
|
||||
@@ -226,6 +277,7 @@ func formatLogOtherJSON(value string, visibility logOtherVisibility) string {
|
||||
}
|
||||
|
||||
changed := false
|
||||
changed = projectResponseModelRejection(values)
|
||||
if visibility == logOtherVisibilityUser {
|
||||
for _, key := range []string{logOtherAdminInfoKey, logOtherRootInfoKey, logOtherAuditInfoKey} {
|
||||
if _, exists := values[key]; exists {
|
||||
@@ -240,7 +292,7 @@ func formatLogOtherJSON(value string, visibility logOtherVisibility) string {
|
||||
}
|
||||
}
|
||||
} else {
|
||||
changed = normalizeLegacyRejectReason(values)
|
||||
changed = normalizeLegacyRejectReason(values) || changed
|
||||
if visibility == logOtherVisibilityAdmin {
|
||||
if _, exists := values[logOtherRootInfoKey]; exists {
|
||||
delete(values, logOtherRootInfoKey)
|
||||
|
||||
@@ -0,0 +1,34 @@
|
||||
package model
|
||||
|
||||
import "encoding/json"
|
||||
|
||||
// Correct historical error-log presentation from the already-stored routing
|
||||
// evidence, without rewriting financial history or exposing private metadata.
|
||||
// Consume records always retain their original billed group.
|
||||
func normalizeErrorLogGroup(log *Log) {
|
||||
if log == nil || log.Type != LogTypeError {
|
||||
return
|
||||
}
|
||||
if log.ChannelId <= 0 {
|
||||
log.Group = ""
|
||||
return
|
||||
}
|
||||
var other struct {
|
||||
Admin struct {
|
||||
Events []struct {
|
||||
ChannelID int `json:"channel_id"`
|
||||
Group string `json:"group"`
|
||||
} `json:"request_policy"`
|
||||
} `json:"admin_info"`
|
||||
}
|
||||
if json.Unmarshal([]byte(log.Other), &other) != nil {
|
||||
return
|
||||
}
|
||||
for i := len(other.Admin.Events) - 1; i >= 0; i-- {
|
||||
event := other.Admin.Events[i]
|
||||
if event.ChannelID == log.ChannelId && event.Group != "" && event.Group != "auto" {
|
||||
log.Group = event.Group
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,38 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"github.com/stretchr/testify/require"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestErrorLogGroupProjection(t *testing.T) {
|
||||
for _, format := range []struct {
|
||||
name string
|
||||
fn func([]*Log)
|
||||
}{
|
||||
{"root", FormatRootLogs}, {"admin", FormatAdminLogs}, {"user", func(logs []*Log) { formatUserLogs(logs, 0) }},
|
||||
} {
|
||||
t.Run(format.name, func(t *testing.T) {
|
||||
metadata := `{"admin_info":{"request_policy":[{"channel_id":101,"group":"example-primary"},{"channel_id":102,"group":"example-preferred"}],"private":"secret"},"root_info":{"private":"secret"}}`
|
||||
logs := []*Log{
|
||||
{Type: LogTypeError, ChannelId: 102, Group: "default", Other: metadata},
|
||||
{Type: LogTypeError, ChannelId: 101, Group: "default", Other: metadata},
|
||||
{Type: LogTypeError, ChannelId: 0, Group: "default", Other: metadata},
|
||||
{Type: LogTypeConsume, ChannelId: 102, Group: "billed-group", Quota: 12345, Other: metadata},
|
||||
{Type: LogTypeError, ChannelId: 99, Group: "legacy", Other: metadata},
|
||||
{Type: LogTypeError, ChannelId: 102, Group: "legacy", Other: `invalid`},
|
||||
}
|
||||
format.fn(logs)
|
||||
for i, want := range []string{"example-preferred", "example-primary", "", "billed-group", "legacy", "legacy"} {
|
||||
require.Equal(t, want, logs[i].Group)
|
||||
}
|
||||
require.Equal(t, 12345, logs[3].Quota)
|
||||
if format.name == "user" {
|
||||
require.NotContains(t, logs[0].Other, "secret")
|
||||
require.NotContains(t, logs[0].Other, "request_policy")
|
||||
}
|
||||
format.fn(logs)
|
||||
require.Equal(t, "example-preferred", logs[0].Group)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -47,8 +47,9 @@ func getUserVerificationState(tx *gorm.DB, userID int, forUpdate bool) (*UserVer
|
||||
|
||||
// CreateUserSessionFromLoginFlow commits the one-time login authorization and
|
||||
// the resulting session together. The user lock serializes credential changes
|
||||
// and session issuance, including the per-user session limits.
|
||||
func CreateUserSessionFromLoginFlow(token string, session *UserSession, validate func(*AuthFlow, *UserVerificationState) error) error {
|
||||
// and session issuance, including the per-user session limits. validate runs
|
||||
// inside the transaction and may write through tx.
|
||||
func CreateUserSessionFromLoginFlow(token string, session *UserSession, validate func(*gorm.DB, *AuthFlow, *UserVerificationState) error) error {
|
||||
if session == nil || validate == nil {
|
||||
return ErrUserSessionInvalid
|
||||
}
|
||||
@@ -63,7 +64,7 @@ func CreateUserSessionFromLoginFlow(token string, session *UserSession, validate
|
||||
if state.Status != common.UserStatusEnabled || state.AuthVersion != session.UserAuthVersion {
|
||||
return ErrUserSessionInactive
|
||||
}
|
||||
if err := validate(flow, state); err != nil {
|
||||
if err := validate(tx, flow, state); err != nil {
|
||||
return err
|
||||
}
|
||||
now := time.Now().Unix()
|
||||
|
||||
@@ -220,6 +220,19 @@ func InitDB() (err error) {
|
||||
}
|
||||
common.SysLog("database migration started")
|
||||
err = migrateDB()
|
||||
if err == nil {
|
||||
err = MigrateTokenAstraAutoSwitch()
|
||||
}
|
||||
if err == nil {
|
||||
err = MigrateTokenAstraDetection()
|
||||
}
|
||||
if err == nil {
|
||||
var count int
|
||||
count, err = MigrateTokenRoutingPolicies()
|
||||
if err == nil {
|
||||
common.SysLog(fmt.Sprintf("token routing migration: %d tokens", count))
|
||||
}
|
||||
}
|
||||
return err
|
||||
} else {
|
||||
common.FatalLog(err)
|
||||
@@ -330,10 +343,14 @@ func migrateDB() error {
|
||||
if err := migrateTokenModelLimitsToText(); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := migrateOptionPrimaryKey(DB); err != nil {
|
||||
common.SysError("failed to migrate options primary key: " + err.Error())
|
||||
}
|
||||
|
||||
err := DB.AutoMigrate(
|
||||
&Channel{},
|
||||
&Token{},
|
||||
&RelayAttempt{},
|
||||
&User{},
|
||||
&UserSession{},
|
||||
&AuthFlow{},
|
||||
@@ -365,6 +382,9 @@ func migrateDB() error {
|
||||
&SystemInstance{},
|
||||
&SystemTask{},
|
||||
&SystemTaskLock{},
|
||||
&DetectionRun{},
|
||||
&DetectionCall{},
|
||||
&DetectionSchedule{},
|
||||
&CasbinRule{},
|
||||
&AuthzRole{},
|
||||
)
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/shopspring/decimal"
|
||||
@@ -10,8 +11,8 @@ import (
|
||||
"gorm.io/gorm/schema"
|
||||
)
|
||||
|
||||
// Embed the concrete dialectors to retain their transaction/savepoint and other
|
||||
// optional GORM interfaces. Only schema comparison needs normalization.
|
||||
// Embed the concrete dialectors to customize schema reconciliation while
|
||||
// retaining transaction/savepoint and other optional GORM interfaces.
|
||||
type mysqlMigrationDialector struct{ mysql.Dialector }
|
||||
|
||||
func (d mysqlMigrationDialector) Migrator(db *gorm.DB) gorm.Migrator {
|
||||
@@ -66,6 +67,46 @@ func (m postgresSchemaMigrator) MigrateColumn(value any, field *schema.Field, co
|
||||
return m.Migrator.MigrateColumn(value, field, column)
|
||||
}
|
||||
|
||||
// Older schemas and database imports can use a different constraint name from
|
||||
// the current naming strategy. Resolve single-column uniqueness from the
|
||||
// catalog instead of asking GORM to drop an inferred uni_<table>_<column>.
|
||||
func (m postgresSchemaMigrator) MigrateColumnUnique(value any, field *schema.Field, column gorm.ColumnType) error {
|
||||
unique, ok := column.Unique()
|
||||
if !ok || field.PrimaryKey || !unique || field.Unique {
|
||||
return m.Migrator.MigrateColumnUnique(value, field, column)
|
||||
}
|
||||
return m.RunWithValue(value, func(stmt *gorm.Statement) error {
|
||||
schemaName, tableName := m.CurrentSchema(stmt, stmt.Table)
|
||||
return m.DB.Transaction(func(tx *gorm.DB) error {
|
||||
if err := tx.Exec("LOCK TABLE ? IN ACCESS EXCLUSIVE MODE", m.CurrentTable(stmt)).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
var constraints []string
|
||||
if err := tx.Raw(`
|
||||
SELECT constraint_meta.conname
|
||||
FROM pg_catalog.pg_constraint AS constraint_meta
|
||||
JOIN pg_catalog.pg_class AS table_meta ON table_meta.oid = constraint_meta.conrelid
|
||||
JOIN pg_catalog.pg_namespace AS schema_meta ON schema_meta.oid = table_meta.relnamespace
|
||||
JOIN pg_catalog.pg_attribute AS attribute_meta
|
||||
ON attribute_meta.attrelid = table_meta.oid
|
||||
AND attribute_meta.attnum = constraint_meta.conkey[1]
|
||||
WHERE schema_meta.nspname = ? AND table_meta.relname = ?
|
||||
AND constraint_meta.contype = 'u'
|
||||
AND cardinality(constraint_meta.conkey) = 1
|
||||
AND attribute_meta.attname = ?
|
||||
ORDER BY constraint_meta.conname`, schemaName, tableName, field.DBName).Scan(&constraints).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
for _, name := range constraints {
|
||||
if err := tx.Migrator().DropConstraint(value, name); err != nil {
|
||||
return fmt.Errorf("drop obsolete unique constraint %q on %s.%s: %w", name, stmt.Table, field.DBName, err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
type charColumnType struct{ migrationColumnType }
|
||||
|
||||
func (charColumnType) DatabaseTypeName() string { return "char" }
|
||||
@@ -10,6 +10,7 @@ import (
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
)
|
||||
|
||||
type MigrationIdentityFields struct {
|
||||
@@ -82,6 +83,32 @@ func TestMigrationSchemaStability(t *testing.T) {
|
||||
t.Cleanup(func() { _ = sqlDB.Close() })
|
||||
recorder := &migrationSQLRecorder{}
|
||||
db = db.Session(&gorm.Session{Logger: recorder})
|
||||
t.Run("request_log_correlation", func(t *testing.T) {
|
||||
const table = "request_log_correlation_test"
|
||||
t.Cleanup(func() { _ = db.Migrator().DropTable(table) })
|
||||
tableDB := db.Table(table).Session(&gorm.Session{})
|
||||
require.NoError(t, tableDB.AutoMigrate(&Log{}))
|
||||
require.NoError(t, tableDB.AutoMigrate(&Log{}))
|
||||
require.True(t, tableDB.Migrator().HasIndex(&Log{}, "idx_logs_correlation"))
|
||||
require.NoError(t, tableDB.Migrator().DropIndex(&Log{}, "idx_logs_correlation"))
|
||||
require.NoError(t, tableDB.Migrator().DropColumn(&Log{}, "correlation_key"))
|
||||
require.NoError(t, tableDB.Omit("correlation_key").Create(&Log{Id: 1, UserId: 1, Quota: 123, Other: "{}"}).Error)
|
||||
for i := 0; i < 2; i++ {
|
||||
require.NoError(t, tableDB.AutoMigrate(&Log{}))
|
||||
}
|
||||
var old Log
|
||||
require.NoError(t, tableDB.First(&old, 1).Error)
|
||||
require.Equal(t, 123, old.Quota)
|
||||
require.Empty(t, old.CorrelationKey)
|
||||
for _, row := range []Log{{Id: 2, UserId: 1, CorrelationKey: "user1-session"}, {Id: 3, UserId: 2, CorrelationKey: "user2-session"}} {
|
||||
require.NoError(t, tableDB.Create(&row).Error)
|
||||
}
|
||||
var found []Log
|
||||
require.NoError(t, tableDB.Where("correlation_key = ?", "user1-session").Find(&found).Error)
|
||||
require.Len(t, found, 1)
|
||||
require.Equal(t, 1, found[0].UserId)
|
||||
require.True(t, tableDB.Migrator().HasIndex(&Log{}, "idx_logs_request_id"))
|
||||
})
|
||||
|
||||
t.Run("identity_and_indexes", func(t *testing.T) {
|
||||
const table = "migration_identity_test"
|
||||
@@ -141,6 +168,48 @@ func TestMigrationSchemaStability(t *testing.T) {
|
||||
require.NoError(t, db.Table(table).Create(&migrationConstraintV1{Name: "existing"}).Error)
|
||||
})
|
||||
|
||||
if dialect == "postgres" {
|
||||
t.Run("renamed_unique_constraints", func(t *testing.T) {
|
||||
const table = "migration_renamed_unique"
|
||||
t.Cleanup(func() { _ = db.Migrator().DropTable(table) })
|
||||
tableDB := db.Table(table).Session(&gorm.Session{})
|
||||
require.NoError(t, tableDB.AutoMigrate(&migrationConstraintV1{}))
|
||||
original := migrationConstraintV1{Name: "existing"}
|
||||
require.NoError(t, tableDB.Create(&original).Error)
|
||||
for _, name := range []string{"models_model_name_key", `imported "model" name`} {
|
||||
require.NoError(t, db.Exec("ALTER TABLE ? ADD CONSTRAINT ? UNIQUE (name)", clause.Table{Name: table}, clause.Column{Name: name}).Error)
|
||||
}
|
||||
require.NoError(t, db.Exec("ALTER TABLE ? ADD CONSTRAINT keep_composite UNIQUE (id, name)", clause.Table{Name: table}).Error)
|
||||
require.NoError(t, db.Exec("CREATE INDEX keep_name_lookup ON ? (name)", clause.Table{Name: table}).Error)
|
||||
// A unique field must keep the old names and still reject duplicates.
|
||||
require.NoError(t, tableDB.AutoMigrate(&migrationConstraintV2{}))
|
||||
require.Error(t, tableDB.Create(&migrationConstraintV2{Name: "existing"}).Error)
|
||||
require.NoError(t, db.Exec("CREATE TABLE migration_unique_reference (name varchar(64) REFERENCES migration_renamed_unique(name))").Error)
|
||||
t.Cleanup(func() { _ = db.Migrator().DropTable("migration_unique_reference") })
|
||||
require.NoError(t, db.Exec("INSERT INTO migration_unique_reference (name) VALUES (?)", "existing").Error)
|
||||
require.Error(t, tableDB.AutoMigrate(&migrationConstraintV1{}), "dependent foreign keys must not be cascaded away")
|
||||
for _, name := range []string{"models_model_name_key", `imported "model" name`} {
|
||||
assert.True(t, tableDB.Migrator().HasConstraint(&migrationConstraintV1{}, name), "all old constraints must survive rollback")
|
||||
}
|
||||
var references int64
|
||||
require.NoError(t, db.Table("migration_unique_reference").Count(&references).Error)
|
||||
assert.EqualValues(t, 1, references)
|
||||
require.NoError(t, db.Migrator().DropTable("migration_unique_reference"))
|
||||
// Removing column uniqueness must resolve both actual constraint names.
|
||||
require.NoError(t, tableDB.AutoMigrate(&migrationConstraintV1{}))
|
||||
recorder.reset()
|
||||
require.NoError(t, tableDB.AutoMigrate(&migrationConstraintV1{}))
|
||||
assert.Empty(t, recorder.schemaMutations())
|
||||
require.NoError(t, tableDB.Create(&migrationConstraintV1{Name: "existing"}).Error)
|
||||
assert.True(t, tableDB.Migrator().HasConstraint(&migrationConstraintV1{}, "keep_composite"))
|
||||
assert.True(t, tableDB.Migrator().HasIndex(&migrationConstraintV1{}, "keep_name_lookup"))
|
||||
var rows []migrationConstraintV1
|
||||
require.NoError(t, tableDB.Order("id").Find(&rows).Error)
|
||||
require.Len(t, rows, 2)
|
||||
assert.Equal(t, original, rows[0])
|
||||
})
|
||||
}
|
||||
|
||||
if dialect == "mysql" {
|
||||
t.Run("decimal_default_and_real_changes", func(t *testing.T) {
|
||||
const table = "migration_decimal_test"
|
||||
|
||||
+253
-44
@@ -4,6 +4,7 @@ import (
|
||||
"crypto/sha256"
|
||||
"errors"
|
||||
"fmt"
|
||||
"maps"
|
||||
"math"
|
||||
"slices"
|
||||
"sort"
|
||||
@@ -11,6 +12,7 @@ import (
|
||||
"sync"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/pkg/billingexpr"
|
||||
"github.com/QuantumNous/new-api/pkg/jsplugin"
|
||||
"github.com/QuantumNous/new-api/setting/billing_setting"
|
||||
"github.com/QuantumNous/new-api/setting/operation_setting"
|
||||
@@ -31,11 +33,24 @@ type ModelPricingChange struct {
|
||||
}
|
||||
|
||||
type ModelPricingEntry struct {
|
||||
ModelName string `json:"model_name"`
|
||||
Version string `json:"version"`
|
||||
Configured PricingValues `json:"configured"`
|
||||
Effective PricingValues `json:"effective"`
|
||||
UsageSchema map[string]jsplugin.UsageFieldSchema `json:"usage_schema,omitempty"`
|
||||
ModelPricingDescription
|
||||
PluginVariants []ModelPricingPluginVariant `json:"plugin_variants,omitempty"`
|
||||
ModelName string `json:"model_name"`
|
||||
Version string `json:"version"`
|
||||
Configured PricingValues `json:"configured"`
|
||||
UsageSchema map[string]jsplugin.UsageFieldSchema `json:"usage_schema,omitempty"`
|
||||
}
|
||||
|
||||
type ModelPricingPluginVariant struct {
|
||||
PluginKey string `json:"plugin_key"`
|
||||
PluginName string `json:"plugin_name"`
|
||||
Icon string `json:"icon,omitempty"`
|
||||
UsageSchema map[string]jsplugin.UsageFieldSchema `json:"usage_schema"`
|
||||
UsageExamples []jsplugin.UsageExample `json:"usage_examples,omitempty"`
|
||||
Configured string `json:"configured"`
|
||||
Effective string `json:"effective"`
|
||||
Compatible bool `json:"compatible"`
|
||||
Stale bool `json:"stale,omitempty"`
|
||||
}
|
||||
|
||||
type ModelPricingSnapshot struct {
|
||||
@@ -51,7 +66,7 @@ var ErrModelPricingConflict = errors.New("model pricing changed; reload before s
|
||||
var modelPricingOptionKeys = []string{
|
||||
"AudioCompletionRatio", "AudioRatio", "CacheRatio", "CompletionRatio",
|
||||
"CreateCacheRatio", "ImageRatio", "ModelPrice", "ModelRatio",
|
||||
"billing_setting.billing_expr", "billing_setting.billing_mode",
|
||||
"billing_setting.billing_expr", "billing_setting.billing_mode", billing_setting.PluginBillingExprOption,
|
||||
}
|
||||
|
||||
var modelPricingMutationMu sync.Mutex
|
||||
@@ -78,28 +93,50 @@ func defaultPricingMaps() map[string]map[string]any {
|
||||
return result
|
||||
}
|
||||
|
||||
func readModelPricingMaps(db *gorm.DB) (map[string]map[string]any, error) {
|
||||
func readModelPricingMaps(db *gorm.DB) (map[string]map[string]any, map[string]bool, []string, error) {
|
||||
var rows []Option
|
||||
if err := db.Where(commonKeyCol+" IN ?", modelPricingOptionKeys).Find(&rows).Error; err != nil {
|
||||
return nil, err
|
||||
return nil, nil, nil, err
|
||||
}
|
||||
values := defaultPricingMaps()
|
||||
existing := make(map[string]bool)
|
||||
counts := make(map[string]int)
|
||||
for _, row := range rows {
|
||||
var entries map[string]any
|
||||
if err := common.UnmarshalJsonStr(row.Value, &entries); err != nil {
|
||||
return nil, fmt.Errorf("%s: %w", row.Key, err)
|
||||
return nil, nil, nil, fmt.Errorf("%s: %w", row.Key, err)
|
||||
}
|
||||
if entries == nil {
|
||||
return nil, fmt.Errorf("%s must be a JSON object", row.Key)
|
||||
return nil, nil, nil, fmt.Errorf("%s must be a JSON object", row.Key)
|
||||
}
|
||||
values[row.Key] = entries
|
||||
existing[row.Key] = true
|
||||
counts[row.Key]++
|
||||
}
|
||||
return values, nil
|
||||
var duplicated []string
|
||||
for _, key := range modelPricingOptionKeys {
|
||||
if counts[key] > 1 {
|
||||
duplicated = append(duplicated, key)
|
||||
}
|
||||
}
|
||||
return values, existing, duplicated, nil
|
||||
}
|
||||
|
||||
func modelPricingValues(values map[string]map[string]any, name string) PricingValues {
|
||||
result := make(PricingValues)
|
||||
for _, key := range modelPricingOptionKeys {
|
||||
if key == billing_setting.PluginBillingExprOption {
|
||||
variants := make(map[string]any)
|
||||
for variant, expression := range values[key] {
|
||||
if plugin, model, ok := billing_setting.SplitPluginBillingExprKey(variant); ok && model == name {
|
||||
variants[plugin] = expression
|
||||
}
|
||||
}
|
||||
if len(variants) > 0 {
|
||||
result[key] = variants
|
||||
}
|
||||
continue
|
||||
}
|
||||
if value, exists := values[key][name]; exists {
|
||||
result[key] = value
|
||||
}
|
||||
@@ -107,11 +144,35 @@ func modelPricingValues(values map[string]map[string]any, name string) PricingVa
|
||||
return result
|
||||
}
|
||||
|
||||
// replaceModelPricing writes one complete model draft into the option maps.
|
||||
// Plugin expressions are grouped in the draft and flattened only in storage.
|
||||
func replaceModelPricing(values map[string]map[string]any, name string, draft PricingValues) {
|
||||
for _, key := range modelPricingOptionKeys {
|
||||
if key == billing_setting.PluginBillingExprOption {
|
||||
for variant := range values[key] {
|
||||
if _, model, ok := billing_setting.SplitPluginBillingExprKey(variant); ok && model == name {
|
||||
delete(values[key], variant)
|
||||
}
|
||||
}
|
||||
variants, _ := draft[key].(map[string]any)
|
||||
for plugin, expr := range variants {
|
||||
values[key][billing_setting.PluginBillingExprKey(plugin, name)] = expr
|
||||
}
|
||||
continue
|
||||
}
|
||||
delete(values[key], name)
|
||||
if value, exists := draft[key]; exists {
|
||||
values[key][name] = value
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func effectiveModelPricing(values map[string]map[string]any, name string) PricingValues {
|
||||
result := modelPricingValues(values, name)
|
||||
// Legacy wildcard aliases are resolved by the same normalization as relay.
|
||||
alias := ratio_setting.FormatMatchingModelName(name)
|
||||
for _, key := range modelPricingOptionKeys[:8] {
|
||||
for _, key := range []string{"ModelPrice", "ModelRatio", "CompletionRatio", "AudioRatio", "AudioCompletionRatio"} {
|
||||
delete(result, key)
|
||||
if value, exists := values[key][alias]; exists {
|
||||
result[key] = value
|
||||
}
|
||||
@@ -141,22 +202,56 @@ func effectiveModelPricing(values map[string]map[string]any, name string) Pricin
|
||||
}
|
||||
// Completion ratios include engine-enforced model defaults. Expose their
|
||||
// effective value without persisting them into the editable configuration.
|
||||
completion := ratio_setting.GetCompletionRatioInfo(name)
|
||||
if _, exists := result["CompletionRatio"]; !exists || completion.Locked {
|
||||
result["CompletionRatio"] = completion.Ratio
|
||||
var configuredCompletion *float64
|
||||
if ratio, exists := result["CompletionRatio"].(float64); exists {
|
||||
configuredCompletion = &ratio
|
||||
}
|
||||
result["CompletionRatio"] = ratio_setting.ResolveCompletionRatio(name, configuredCompletion).Ratio
|
||||
for key, fallback := range map[string]float64{
|
||||
"CacheRatio": ratio_setting.DefaultCacheRatio,
|
||||
"CreateCacheRatio": ratio_setting.DefaultCreateCacheRatio,
|
||||
"ImageRatio": ratio_setting.DefaultImageRatio,
|
||||
} {
|
||||
if _, exists := result[key]; !exists {
|
||||
result[key] = fallback
|
||||
}
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// PreviewModelPricing resolves a complete editable draft using the same defaults
|
||||
// as the saved-price display and conversion. It has no write side effects.
|
||||
func PreviewModelPricing(name string, draft PricingValues) (PricingValues, error) {
|
||||
if draft == nil {
|
||||
return nil, errors.New("pricing draft is required")
|
||||
}
|
||||
values, _, _, err := readModelPricingMaps(DB)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := validateModelPricing(name, draft, modelPricingValues(values, name)); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
replaceModelPricing(values, name, draft)
|
||||
return effectiveModelPricing(values, name), nil
|
||||
}
|
||||
|
||||
func GetModelPricingSnapshot(names []string) (*ModelPricingSnapshot, error) {
|
||||
values, err := readModelPricingMaps(DB)
|
||||
values, _, _, err := readModelPricingMaps(DB)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(names) == 0 {
|
||||
nameSet := make(map[string]bool)
|
||||
for _, entries := range values {
|
||||
for key, entries := range values {
|
||||
for name := range entries {
|
||||
if key == billing_setting.PluginBillingExprOption {
|
||||
_, model, ok := billing_setting.SplitPluginBillingExprKey(name)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
name = model
|
||||
}
|
||||
nameSet[name] = true
|
||||
}
|
||||
}
|
||||
@@ -172,12 +267,55 @@ func GetModelPricingSnapshot(names []string) (*ModelPricingSnapshot, error) {
|
||||
generation := jsplugin.DefaultRegistry.Generation()
|
||||
for _, name := range names {
|
||||
configured := modelPricingValues(values, name)
|
||||
entry := ModelPricingEntry{ModelName: name, Version: ModelPricingVersion(configured), Configured: configured, Effective: effectiveModelPricing(values, name)}
|
||||
entry := ModelPricingEntry{ModelName: name, Version: ModelPricingVersion(configured), Configured: configured,
|
||||
ModelPricingDescription: ModelPricingDescription{Effective: effectiveModelPricing(values, name)}}
|
||||
entry.CacheWriteMode = ResolveCacheWriteMode(name, configured)
|
||||
entry.BillingDetails = ResolveLegacyBillingDetails(name, entry.Effective, configured)
|
||||
if plugin, ok := generation.GetByModel(name); ok {
|
||||
entry.UsageSchema = plugin.Meta.UsageSchema
|
||||
entry.UsageSchema, _ = plugin.Meta.UsageForModel(name)
|
||||
} else if target, ok := ResolveTaskModelAlias(generation, name); ok {
|
||||
if plugin, ok := generation.Get(target.PluginKey); ok {
|
||||
entry.UsageSchema = plugin.Meta.UsageSchema
|
||||
entry.UsageSchema, _ = plugin.Meta.UsageForModel(target.Declared)
|
||||
}
|
||||
}
|
||||
plugins := generation.PluginsByModel(name)
|
||||
configuredVariants, _ := configured[billing_setting.PluginBillingExprOption].(map[string]any)
|
||||
if len(plugins) >= 2 || len(configuredVariants) > 0 {
|
||||
keys := make(map[string]bool, len(plugins)+len(configuredVariants))
|
||||
for _, plugin := range plugins {
|
||||
keys[plugin.Meta.Key] = true
|
||||
}
|
||||
for key := range configuredVariants {
|
||||
keys[key] = true
|
||||
}
|
||||
for _, key := range slices.Sorted(maps.Keys(keys)) {
|
||||
configuredValue, overridden := configuredVariants[key]
|
||||
configuredExpr, _ := configuredValue.(string)
|
||||
plugin, exists := generation.Get(key)
|
||||
if !exists || !slices.Contains(plugin.Meta.Models, name) {
|
||||
variant := ModelPricingPluginVariant{
|
||||
PluginKey: key, PluginName: key, Configured: configuredExpr,
|
||||
UsageSchema: map[string]jsplugin.UsageFieldSchema{}, Stale: true,
|
||||
}
|
||||
if exists {
|
||||
variant.PluginName, variant.Icon = plugin.Meta.Name, plugin.Meta.Icon
|
||||
}
|
||||
entry.PluginVariants = append(entry.PluginVariants, variant)
|
||||
continue
|
||||
}
|
||||
schema, examples := plugin.Meta.UsageForModel(name)
|
||||
if schema == nil {
|
||||
schema = map[string]jsplugin.UsageFieldSchema{}
|
||||
}
|
||||
expression := configuredExpr
|
||||
if !overridden && entry.Effective["billing_setting.billing_mode"] == billing_setting.BillingModeTieredExpr {
|
||||
expression, _ = entry.Effective["billing_setting.billing_expr"].(string)
|
||||
}
|
||||
entry.PluginVariants = append(entry.PluginVariants, ModelPricingPluginVariant{
|
||||
PluginKey: plugin.Meta.Key, PluginName: plugin.Meta.Name, Icon: plugin.Meta.Icon,
|
||||
UsageSchema: schema, UsageExamples: examples, Configured: configuredExpr, Effective: expression,
|
||||
Compatible: billing_setting.TaskExprCompatible(expression, schema),
|
||||
})
|
||||
}
|
||||
}
|
||||
result.Entries = append(result.Entries, entry)
|
||||
@@ -207,10 +345,56 @@ func GetModelPricingSnapshot(names []string) (*ModelPricingSnapshot, error) {
|
||||
}
|
||||
|
||||
func ValidateModelPricing(name string, values PricingValues) error {
|
||||
variants := make(map[string]any)
|
||||
for key, expression := range billing_setting.GetPluginBillingExprCopy() {
|
||||
if plugin, model, ok := billing_setting.SplitPluginBillingExprKey(key); ok && model == name {
|
||||
variants[plugin] = expression
|
||||
}
|
||||
}
|
||||
previous := PricingValues{billing_setting.PluginBillingExprOption: variants}
|
||||
if expression, ok := billing_setting.GetBillingExpr(name); ok {
|
||||
previous["billing_setting.billing_expr"] = expression
|
||||
}
|
||||
return validateModelPricing(name, values, previous)
|
||||
}
|
||||
|
||||
// Writes pass the locked database snapshot here, so allowing an unchanged stale
|
||||
// override cannot bypass validation through an out-of-date process-local cache.
|
||||
func validateModelPricing(name string, values, previous PricingValues) error {
|
||||
if strings.TrimSpace(name) == "" {
|
||||
return errors.New("model name is required")
|
||||
}
|
||||
generation := jsplugin.DefaultRegistry.Generation()
|
||||
previousVariants, _ := previous[billing_setting.PluginBillingExprOption].(map[string]any)
|
||||
variants := map[string]any{}
|
||||
if value, exists := values[billing_setting.PluginBillingExprOption]; exists {
|
||||
var ok bool
|
||||
variants, ok = value.(map[string]any)
|
||||
if !ok || variants == nil {
|
||||
return errors.New("plugin billing expressions must be a plugin-to-expression object")
|
||||
}
|
||||
for key, value := range variants {
|
||||
expression, ok := value.(string)
|
||||
if !ok || strings.TrimSpace(expression) == "" {
|
||||
return fmt.Errorf("model %s: plugin %s: billing expression is required", name, key)
|
||||
}
|
||||
plugin, exists := generation.Get(key)
|
||||
if !exists || !slices.Contains(plugin.Meta.Models, name) {
|
||||
if previousVariants[key] == expression {
|
||||
continue
|
||||
}
|
||||
return fmt.Errorf("model %s: plugin %s does not declare this model", name, key)
|
||||
}
|
||||
schema, _ := plugin.Meta.UsageForModel(name)
|
||||
if err := billing_setting.SmokeTestTaskExpr(expression, schema); err != nil {
|
||||
return fmt.Errorf("model %s: plugin %s: %w", name, key, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
for key, value := range values {
|
||||
if key == billing_setting.PluginBillingExprOption {
|
||||
continue
|
||||
}
|
||||
if !IsModelPricingOption(key) {
|
||||
return fmt.Errorf("unsupported pricing field: %s", key)
|
||||
}
|
||||
@@ -225,19 +409,35 @@ func ValidateModelPricing(name string, values PricingValues) error {
|
||||
if !ok || strings.TrimSpace(expression) == "" {
|
||||
return errors.New("billing expression is required")
|
||||
}
|
||||
generation := jsplugin.DefaultRegistry.Generation()
|
||||
// Even a model expression currently shadowed by every provider must
|
||||
// compile; only its schema-specific smoke tests can be skipped.
|
||||
if _, err := billingexpr.CompileFromCache(expression); err != nil {
|
||||
return fmt.Errorf("model %s: %w", name, err)
|
||||
}
|
||||
var err error
|
||||
if plugin, ok := generation.GetByModel(name); ok {
|
||||
err = billing_setting.SmokeTestTaskExpr(expression, plugin.Meta.UsageSchema)
|
||||
if plugins := generation.PluginsByModel(name); len(plugins) > 0 {
|
||||
for _, plugin := range plugins {
|
||||
if _, overridden := variants[plugin.Meta.Key]; overridden {
|
||||
continue
|
||||
}
|
||||
schema, _ := plugin.Meta.UsageForModel(name)
|
||||
if err = billing_setting.SmokeTestTaskExpr(expression, schema); err != nil {
|
||||
return fmt.Errorf("model %s: plugin %s: %w", name, plugin.Meta.Key, err)
|
||||
}
|
||||
}
|
||||
} else if target, resolved := ResolveTaskModelAlias(generation, name); resolved {
|
||||
if plugin, ok := generation.Get(target.PluginKey); ok {
|
||||
err = billing_setting.SmokeTestTaskExpr(expression, plugin.Meta.UsageSchema)
|
||||
schema, _ := plugin.Meta.UsageForModel(target.Declared)
|
||||
err = billing_setting.SmokeTestTaskExpr(expression, schema)
|
||||
} else {
|
||||
err = billing_setting.SmokeTestExpr(expression)
|
||||
}
|
||||
} else {
|
||||
} else if previous[key] != expression || len(billingexpr.UsedUsageKeys(expression)) == 0 {
|
||||
err = billing_setting.SmokeTestExpr(expression)
|
||||
}
|
||||
// With no remaining plugin, an unchanged stored usage expression has
|
||||
// no schema to test. Preserve it so removing stale overrides or saving
|
||||
// other model prices does not become impossible.
|
||||
if err != nil {
|
||||
return fmt.Errorf("model %s: %w", name, err)
|
||||
}
|
||||
@@ -271,26 +471,22 @@ func UpdateModelPricing(changes []ModelPricingChange) error {
|
||||
if change.ExpectedVersion == "" {
|
||||
return ErrModelPricingConflict
|
||||
}
|
||||
if err := ValidateModelPricing(change.ModelName, change.Pricing); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return mutateModelPricingOptions(func(_ *gorm.DB, values map[string]map[string]any) error {
|
||||
defaults := defaultPricingMaps()
|
||||
for _, change := range changes {
|
||||
if ModelPricingVersion(modelPricingValues(values, change.ModelName)) != change.ExpectedVersion {
|
||||
previous := modelPricingValues(values, change.ModelName)
|
||||
if ModelPricingVersion(previous) != change.ExpectedVersion {
|
||||
return fmt.Errorf("%w: %s", ErrModelPricingConflict, change.ModelName)
|
||||
}
|
||||
pricing := change.Pricing
|
||||
if change.Reset {
|
||||
pricing = modelPricingValues(defaults, change.ModelName)
|
||||
}
|
||||
for _, key := range modelPricingOptionKeys {
|
||||
delete(values[key], change.ModelName)
|
||||
if value, exists := pricing[key]; exists {
|
||||
values[key][change.ModelName] = value
|
||||
}
|
||||
if err := validateModelPricing(change.ModelName, pricing, previous); err != nil {
|
||||
return err
|
||||
}
|
||||
replaceModelPricing(values, change.ModelName, pricing)
|
||||
}
|
||||
return nil
|
||||
})
|
||||
@@ -300,6 +496,7 @@ func UpdateModelPricing(changes []ModelPricingChange) error {
|
||||
// locking, validation and transaction path as the model-level API.
|
||||
func UpdateModelPricingOptions(updates map[string]string) error {
|
||||
return mutateModelPricingOptions(func(_ *gorm.DB, values map[string]map[string]any) error {
|
||||
previous := maps.Clone(values)
|
||||
names := make(map[string]bool)
|
||||
for key, raw := range updates {
|
||||
if !IsModelPricingOption(key) {
|
||||
@@ -312,16 +509,22 @@ func UpdateModelPricingOptions(updates map[string]string) error {
|
||||
if entries == nil {
|
||||
return fmt.Errorf("%s must be a JSON object", key)
|
||||
}
|
||||
for name := range values[key] {
|
||||
names[name] = true
|
||||
for _, entriesForKey := range []map[string]any{values[key], entries} {
|
||||
for name := range entriesForKey {
|
||||
if key == billing_setting.PluginBillingExprOption {
|
||||
_, model, ok := billing_setting.SplitPluginBillingExprKey(name)
|
||||
if !ok {
|
||||
return fmt.Errorf("invalid plugin billing expression key: %s", name)
|
||||
}
|
||||
name = model
|
||||
}
|
||||
names[name] = true
|
||||
}
|
||||
}
|
||||
values[key] = entries
|
||||
for name := range entries {
|
||||
names[name] = true
|
||||
}
|
||||
}
|
||||
for name := range names {
|
||||
if err := ValidateModelPricing(name, modelPricingValues(values, name)); err != nil {
|
||||
if err := validateModelPricing(name, modelPricingValues(values, name), modelPricingValues(previous, name)); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
@@ -334,8 +537,18 @@ func mutateModelPricingOptions(mutate func(*gorm.DB, map[string]map[string]any)
|
||||
defer modelPricingMutationMu.Unlock()
|
||||
var committed map[string]map[string]any
|
||||
err := DB.Transaction(func(tx *gorm.DB) error {
|
||||
values, existing, duplicated, err := readModelPricingMaps(lockForUpdate(tx))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if len(duplicated) > 0 {
|
||||
common.SysError("options table has duplicate pricing keys [" + strings.Join(duplicated, ", ") + "]; the table is missing a primary key")
|
||||
}
|
||||
defaults := defaultPricingMaps()
|
||||
for _, key := range modelPricingOptionKeys {
|
||||
if existing[key] {
|
||||
continue
|
||||
}
|
||||
encoded, err := common.Marshal(defaults[key])
|
||||
if err != nil {
|
||||
return err
|
||||
@@ -345,10 +558,6 @@ func mutateModelPricingOptions(mutate func(*gorm.DB, map[string]map[string]any)
|
||||
return err
|
||||
}
|
||||
}
|
||||
values, err := readModelPricingMaps(lockForUpdate(tx))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := mutate(tx, values); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
@@ -0,0 +1,368 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"maps"
|
||||
"math"
|
||||
"slices"
|
||||
"strings"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/pkg/jsplugin"
|
||||
"github.com/QuantumNous/new-api/setting/billing_setting"
|
||||
"github.com/QuantumNous/new-api/setting/operation_setting"
|
||||
"github.com/QuantumNous/new-api/setting/ratio_setting"
|
||||
hostreasoning "github.com/QuantumNous/new-api/setting/reasoning"
|
||||
"github.com/shopspring/decimal"
|
||||
)
|
||||
|
||||
// ModelPricingConversion is a preview only. Saving still requires a versioned
|
||||
// ModelPricingChange through the existing pricing transaction.
|
||||
type ModelPricingConversion struct {
|
||||
ModelPricingDescription
|
||||
Expression string `json:"expression,omitempty"`
|
||||
UnsupportedReason string `json:"unsupported_reason,omitempty"`
|
||||
}
|
||||
|
||||
// ModelPricingDescription is shared by previews, conversions and snapshot
|
||||
// entries. It describes effective prices without adding persisted settings.
|
||||
type ModelPricingDescription struct {
|
||||
Effective PricingValues `json:"effective,omitempty"`
|
||||
CacheWriteMode CacheWriteMode `json:"cache_write_mode,omitempty"`
|
||||
BillingDetails LegacyBillingDetails `json:"billing_details"`
|
||||
}
|
||||
|
||||
func legacyInputPricePerMillion(ratio float64) (decimal.Decimal, error) {
|
||||
if common.QuotaPerUnit <= 0 || math.IsInf(common.QuotaPerUnit, 0) || math.IsNaN(common.QuotaPerUnit) {
|
||||
return decimal.Zero, errors.New("invalid quota unit")
|
||||
}
|
||||
if ratio < 0 || math.IsInf(ratio, 0) || math.IsNaN(ratio) {
|
||||
return decimal.Zero, errors.New("input ratio must be finite and non-negative")
|
||||
}
|
||||
return decimal.NewFromFloat(ratio).Mul(decimal.NewFromInt(1_000_000)).Div(decimal.NewFromFloat(common.QuotaPerUnit)), nil
|
||||
}
|
||||
|
||||
type LegacyPricingRule struct {
|
||||
Condition string `json:"condition"`
|
||||
Multiplier float64 `json:"multiplier"`
|
||||
}
|
||||
|
||||
// LegacyBillingDetails is display-only. Absolute audio prices cannot always be
|
||||
// represented by a ratio: Gemini audio can be billable with a zero text price.
|
||||
type LegacyBillingDetails struct {
|
||||
AudioInputPrice *float64 `json:"audio_input_price,omitempty"`
|
||||
AudioOutputPrice *float64 `json:"audio_output_price,omitempty"`
|
||||
ImageCount bool `json:"image_count,omitempty"`
|
||||
RequestRules []LegacyPricingRule `json:"request_rules,omitempty"`
|
||||
AudioTextBranches bool `json:"audio_text_branches,omitempty"`
|
||||
InvalidAudioPrice bool `json:"-"`
|
||||
ConflictingAudioPrices bool `json:"-"`
|
||||
}
|
||||
|
||||
func ResolveLegacyBillingDetails(name string, effective, configured PricingValues) LegacyBillingDetails {
|
||||
details := LegacyBillingDetails{}
|
||||
if effective["billing_setting.billing_mode"] == billing_setting.BillingModeTieredExpr {
|
||||
return details
|
||||
}
|
||||
if _, fixed := effective["ModelPrice"]; fixed {
|
||||
details.ImageCount = common.IsImageGenerationModel(name)
|
||||
details.RequestRules = legacyDallePricingRules(name)
|
||||
return details
|
||||
}
|
||||
ratio, priced := effective["ModelRatio"].(float64)
|
||||
if !priced {
|
||||
return details
|
||||
}
|
||||
base, err := legacyInputPricePerMillion(ratio)
|
||||
if err != nil {
|
||||
return details
|
||||
}
|
||||
if price := operation_setting.GetGeminiInputAudioPricePerMillionTokens(name); price > 0 {
|
||||
details.AudioInputPrice = &price
|
||||
// Native Gemini settles with its dedicated input price; compatible
|
||||
// chat enters the audio-ratio path when those ratios are configured.
|
||||
// Do not silently replace one price with the other during migration.
|
||||
audio, hasInput := effective["AudioRatio"].(float64)
|
||||
output, hasOutput := effective["AudioCompletionRatio"].(float64)
|
||||
if hasInput || hasOutput {
|
||||
if !hasInput {
|
||||
audio = 1
|
||||
}
|
||||
if !hasOutput {
|
||||
output = 1
|
||||
}
|
||||
audioPrice := base.Mul(decimal.NewFromFloat(audio))
|
||||
textOutput := base.Mul(decimal.NewFromFloat(effective["CompletionRatio"].(float64)))
|
||||
details.ConflictingAudioPrices = !audioPrice.Equal(decimal.NewFromFloat(price)) || !audioPrice.Mul(decimal.NewFromFloat(output)).Equal(textOutput)
|
||||
}
|
||||
return details
|
||||
}
|
||||
audioRatio, hasAudio := effective["AudioRatio"].(float64)
|
||||
audioCompletionRatio, hasAudioCompletion := effective["AudioCompletionRatio"].(float64)
|
||||
if !hasAudio && !hasAudioCompletion {
|
||||
return details
|
||||
}
|
||||
if !hasAudio {
|
||||
audioRatio = 1
|
||||
}
|
||||
if !hasAudioCompletion {
|
||||
audioCompletionRatio = 1
|
||||
}
|
||||
input := base.Mul(decimal.NewFromFloat(audioRatio))
|
||||
inPrice, outPrice := input.InexactFloat64(), input.Mul(decimal.NewFromFloat(audioCompletionRatio)).InexactFloat64()
|
||||
if math.IsInf(inPrice, 0) || math.IsInf(outPrice, 0) {
|
||||
details.InvalidAudioPrice = true
|
||||
return details
|
||||
}
|
||||
details.AudioInputPrice, details.AudioOutputPrice = &inPrice, &outPrice
|
||||
details.AudioTextBranches = effective["CacheRatio"] != float64(1) || effective["ImageRatio"] != float64(1) || ResolveCacheWriteMode(name, configured) != CacheWriteNone
|
||||
return details
|
||||
}
|
||||
|
||||
type CacheWriteMode string
|
||||
|
||||
const (
|
||||
CacheWriteNone CacheWriteMode = "none"
|
||||
CacheWriteStandard CacheWriteMode = "standard"
|
||||
CacheWriteClaudeTTL CacheWriteMode = "claude_ttl"
|
||||
)
|
||||
|
||||
// ResolveCacheWriteMode describes display and migration only. A generic engine
|
||||
// fallback is not configured cache-write pricing. Only Claude names inherit the
|
||||
// legacy dual-TTL rule; other models retain their configured write price.
|
||||
func ResolveCacheWriteMode(name string, configured PricingValues) CacheWriteMode {
|
||||
if strings.Contains(strings.ToLower(name), "claude") {
|
||||
return CacheWriteClaudeTTL
|
||||
}
|
||||
if _, exists := configured["CreateCacheRatio"]; exists {
|
||||
return CacheWriteStandard
|
||||
}
|
||||
return CacheWriteNone
|
||||
}
|
||||
|
||||
func PreviewModelPricingConversion(name string, draft PricingValues) (*ModelPricingConversion, error) {
|
||||
if draft == nil {
|
||||
return nil, errors.New("pricing draft is required")
|
||||
}
|
||||
if err := ValidateModelPricing(name, draft); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if draft["billing_setting.billing_mode"] == billing_setting.BillingModeTieredExpr {
|
||||
return &ModelPricingConversion{UnsupportedReason: "This model already uses an expression."}, nil
|
||||
}
|
||||
// The draft is the complete editable configuration: omitted fields inherit
|
||||
// relay defaults rather than silently retaining discarded saved values.
|
||||
legacyDraft := make(PricingValues, len(draft)+1)
|
||||
maps.Copy(legacyDraft, draft)
|
||||
legacyDraft["billing_setting.billing_mode"] = billing_setting.BillingModeRatio
|
||||
effective, err := PreviewModelPricing(name, legacyDraft)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := ValidateModelPricing(name, effective); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
_, fixedPrice := effective["ModelPrice"]
|
||||
preview := &ModelPricingConversion{ModelPricingDescription: ModelPricingDescription{Effective: effective, BillingDetails: ResolveLegacyBillingDetails(name, effective, draft)}}
|
||||
if preview.BillingDetails.InvalidAudioPrice {
|
||||
return nil, errors.New("audio prices must be finite")
|
||||
}
|
||||
if preview.BillingDetails.ConflictingAudioPrices {
|
||||
return &ModelPricingConversion{UnsupportedReason: "Gemini and OpenAI audio prices differ for this model. Use separate billing model names to convert them."}, nil
|
||||
}
|
||||
|
||||
generation := jsplugin.DefaultRegistry.Generation()
|
||||
if _, task := generation.GetByModel(name); task {
|
||||
return &ModelPricingConversion{UnsupportedReason: "Task pricing must be converted manually using the task usage schema."}, nil
|
||||
}
|
||||
if _, task := ResolveTaskModelAlias(generation, name); task {
|
||||
return &ModelPricingConversion{UnsupportedReason: "Task pricing must be converted manually using the task usage schema."}, nil
|
||||
}
|
||||
|
||||
// Inspect non-secret routing metadata so aliases cannot disguise a special
|
||||
// settlement path as an ordinary text model.
|
||||
var channels []Channel
|
||||
candidates := DB.Model(&Ability{}).Select("channel_id").Where("model = ?", name)
|
||||
if err := DB.Select("type", "status", "models", "model_mapping").Where("id IN (?)", candidates).Find(&channels).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
names := []string{name}
|
||||
aliPromptExtend, otherImageRoute := false, false
|
||||
for _, channel := range channels {
|
||||
if !slices.Contains(channel.GetModels(), name) {
|
||||
continue
|
||||
}
|
||||
if channel.Type == constant.ChannelTypeTaskPlugin {
|
||||
return &ModelPricingConversion{UnsupportedReason: "Task pricing must be converted manually using the task usage schema."}, nil
|
||||
}
|
||||
if slices.Contains(common.GetEndpointTypesByChannelType(channel.Type, name), constant.EndpointTypeOpenAIVideo) {
|
||||
return &ModelPricingConversion{UnsupportedReason: "Video pricing must be converted manually."}, nil
|
||||
}
|
||||
if channel.Status == common.ChannelStatusEnabled && channel.Type == constant.ChannelTypeOpenRouter && strings.Contains(strings.ToLower(name), "claude") && !fixedPrice {
|
||||
defaultRatio, hasDefault := ratio_setting.GetDefaultModelRatioMap()[name]
|
||||
if hasDefault && effective["ModelRatio"] == defaultRatio && effective["CreateCacheRatio"] != float64(1) {
|
||||
return &ModelPricingConversion{UnsupportedReason: "This OpenRouter Claude price derives cache-write usage from upstream cost and must be converted manually."}, nil
|
||||
}
|
||||
}
|
||||
upstream := name
|
||||
if channel.ModelMapping != nil && *channel.ModelMapping != "" {
|
||||
var mapping map[string]string
|
||||
if err := common.UnmarshalJsonStr(*channel.ModelMapping, &mapping); err != nil {
|
||||
return &ModelPricingConversion{UnsupportedReason: "The model routing configuration could not be verified."}, nil
|
||||
}
|
||||
// Relay also resolves reasoning-suffixed names at any mapping hop.
|
||||
candidates := []string{name}
|
||||
for _, target := range mapping {
|
||||
candidates = append(candidates, target)
|
||||
}
|
||||
for _, candidate := range candidates {
|
||||
if mapping[candidate] == "" {
|
||||
mapping[candidate] = mapping[hostreasoning.BaseModelName(candidate)]
|
||||
}
|
||||
}
|
||||
var cycle bool
|
||||
upstream, cycle = followChannelModelMapping(mapping, name)
|
||||
if cycle {
|
||||
return &ModelPricingConversion{UnsupportedReason: "The model routing configuration could not be verified."}, nil
|
||||
}
|
||||
if upstream != name {
|
||||
names = append(names, upstream)
|
||||
}
|
||||
}
|
||||
if channel.Status == common.ChannelStatusEnabled {
|
||||
if channel.Type == constant.ChannelTypeAli && strings.Contains(upstream, "z-image") {
|
||||
aliPromptExtend = true
|
||||
} else {
|
||||
otherImageRoute = true
|
||||
}
|
||||
}
|
||||
}
|
||||
for _, modelName := range names {
|
||||
lower := strings.ToLower(modelName)
|
||||
if strings.Contains(lower, "realtime") {
|
||||
return &ModelPricingConversion{UnsupportedReason: "Realtime pricing must be converted manually."}, nil
|
||||
}
|
||||
if fixedPrice && ResolveLegacyBillingDetails(modelName, effective, draft).ImageCount {
|
||||
preview.BillingDetails.ImageCount = true
|
||||
}
|
||||
}
|
||||
var metadata []Model
|
||||
if err := DB.Select("model_name", "name_rule", "endpoints").
|
||||
Where("model_name IN ? OR name_rule <> ?", names, NameRuleExact).
|
||||
Where("endpoints <> ?", "").Find(&metadata).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for _, entry := range metadata {
|
||||
if entry.Endpoints == "" || !slices.ContainsFunc(names, entry.MatchesName) {
|
||||
continue
|
||||
}
|
||||
var endpoints map[string]any
|
||||
if err := common.UnmarshalJsonStr(entry.Endpoints, &endpoints); err != nil {
|
||||
return &ModelPricingConversion{UnsupportedReason: "The model routing configuration could not be verified."}, nil
|
||||
}
|
||||
for endpoint := range endpoints {
|
||||
switch constant.EndpointType(endpoint) {
|
||||
case constant.EndpointTypeImageGeneration:
|
||||
if fixedPrice {
|
||||
preview.BillingDetails.ImageCount = true
|
||||
}
|
||||
case constant.EndpointTypeOpenAIVideo:
|
||||
return &ModelPricingConversion{UnsupportedReason: "Video pricing must be converted manually."}, nil
|
||||
case constant.EndpointTypeOpenAI, constant.EndpointTypeOpenAIResponse, constant.EndpointTypeAnthropic, constant.EndpointTypeGemini, constant.EndpointTypeEmbeddings, constant.EndpointTypeJinaRerank:
|
||||
// These endpoints use the ordinary token/fixed-request settlement.
|
||||
}
|
||||
}
|
||||
}
|
||||
var expression string
|
||||
if fixedPrice {
|
||||
price := effective["ModelPrice"].(float64)
|
||||
expression = `tier("request", fixed(` + decimal.NewFromFloat(price).String() + `))`
|
||||
if preview.BillingDetails.ImageCount {
|
||||
if aliPromptExtend && otherImageRoute {
|
||||
return &ModelPricingConversion{UnsupportedReason: "This model has different image request multipliers across channels. Use separate billing model names to convert them."}, nil
|
||||
}
|
||||
expression = `tier("image", fixed(` + decimal.NewFromFloat(price).String() + `)) * image_count`
|
||||
if aliPromptExtend {
|
||||
preview.BillingDetails.RequestRules = append(preview.BillingDetails.RequestRules, LegacyPricingRule{`param("parameters.prompt_extend") == true`, common.ZImagePromptExtendMultiplier})
|
||||
}
|
||||
for _, rule := range preview.BillingDetails.RequestRules {
|
||||
expression += ` * (` + rule.Condition + ` ? ` + decimal.NewFromFloat(rule.Multiplier).String() + ` : 1)`
|
||||
}
|
||||
}
|
||||
} else {
|
||||
ratio, exists := effective["ModelRatio"].(float64)
|
||||
if !exists {
|
||||
return &ModelPricingConversion{UnsupportedReason: "Configure an input price before converting this model."}, nil
|
||||
}
|
||||
preview.CacheWriteMode = ResolveCacheWriteMode(name, draft)
|
||||
base, err := legacyInputPricePerMillion(ratio)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// Cache reads can overlap other separately priced input categories.
|
||||
// Keep cr in that case: clamping the input remainder makes folding it
|
||||
// into p non-equivalent. A zero base price is equivalent either way.
|
||||
ordinaryAudio := preview.BillingDetails.AudioOutputPrice != nil
|
||||
if ordinaryAudio && preview.CacheWriteMode != CacheWriteNone {
|
||||
preview.BillingDetails.AudioTextBranches = true
|
||||
}
|
||||
mergeCacheRead := base.IsZero() || preview.CacheWriteMode == CacheWriteNone && effective["ImageRatio"].(float64) == 1 && (preview.BillingDetails.AudioInputPrice == nil || ordinaryAudio)
|
||||
var expressionBody strings.Builder
|
||||
expressionBody.WriteString(`tier("base", p * `)
|
||||
expressionBody.WriteString(base.String())
|
||||
// Retain explicit zero coefficients for included categories. Only
|
||||
// Claude names inherit the relay's 1h / 5m factor of 6 / 3.75.
|
||||
for _, lane := range []struct {
|
||||
variable, key string
|
||||
multiplier float64
|
||||
}{
|
||||
{"c", "CompletionRatio", 1},
|
||||
{"cr", "CacheRatio", 1},
|
||||
{"cc", "CreateCacheRatio", 1},
|
||||
{"cc1h", "CreateCacheRatio", 6 / 3.75},
|
||||
{"img", "ImageRatio", 1},
|
||||
} {
|
||||
if ordinaryAudio && !preview.BillingDetails.AudioTextBranches && lane.variable != "c" {
|
||||
continue
|
||||
}
|
||||
if lane.variable == "cc" && preview.CacheWriteMode == CacheWriteNone ||
|
||||
lane.variable == "cc1h" && preview.CacheWriteMode != CacheWriteClaudeTTL {
|
||||
continue
|
||||
}
|
||||
multiplier := effective[lane.key].(float64)
|
||||
// Same-price input categories stay in p. Keep distinct prices,
|
||||
// including explicit zero, as independently priced terms.
|
||||
if (lane.variable == "img" || lane.variable == "cr" && mergeCacheRead) && multiplier == 1 {
|
||||
continue
|
||||
}
|
||||
expressionBody.WriteString(" + " + lane.variable + " * ")
|
||||
expressionBody.WriteString(base.Mul(decimal.NewFromFloat(multiplier)).Mul(decimal.NewFromFloat(lane.multiplier)).String())
|
||||
}
|
||||
for _, lane := range []struct {
|
||||
variable string
|
||||
price *float64
|
||||
}{
|
||||
{"ai", preview.BillingDetails.AudioInputPrice},
|
||||
{"ao", preview.BillingDetails.AudioOutputPrice},
|
||||
} {
|
||||
if lane.price != nil && !preview.BillingDetails.AudioTextBranches {
|
||||
expressionBody.WriteString(" + " + lane.variable + " * " + decimal.NewFromFloat(*lane.price).String())
|
||||
}
|
||||
}
|
||||
expressionBody.WriteString(")")
|
||||
expression = expressionBody.String()
|
||||
if ordinaryAudio && preview.BillingDetails.AudioTextBranches {
|
||||
completion := base.Mul(decimal.NewFromFloat(effective["CompletionRatio"].(float64)))
|
||||
// The audio settlement ignores cache/image ratios. Use the full
|
||||
// input total here because variables in the text-only branch still
|
||||
// participate in AST-based token normalization.
|
||||
expression = `(ai > 0 || ao > 0) ? tier("audio", max(len - ai, 0) * ` + base.String() + ` + c * ` + completion.String() + ` + ai * ` + decimal.NewFromFloat(*preview.BillingDetails.AudioInputPrice).String() + ` + ao * ` + decimal.NewFromFloat(*preview.BillingDetails.AudioOutputPrice).String() + `) : ` + expression
|
||||
}
|
||||
}
|
||||
if err := billing_setting.SmokeTestExpr(expression); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
preview.Expression = expression
|
||||
return preview, nil
|
||||
}
|
||||
@@ -195,13 +195,28 @@ func InitOptionMap() {
|
||||
}
|
||||
|
||||
func loadOptionsFromDatabase() {
|
||||
requestPolicyOptionMutex.Lock()
|
||||
defer requestPolicyOptionMutex.Unlock()
|
||||
defer func() {
|
||||
if err := refreshRequestPolicySnapshot(); err != nil {
|
||||
common.SysError("invalid request policy: " + err.Error())
|
||||
}
|
||||
}()
|
||||
passkeyOptionMutex.Lock()
|
||||
defer passkeyOptionMutex.Unlock()
|
||||
options, _ := AllOption()
|
||||
passkeyOptions := make(map[string]string)
|
||||
for _, option := range options {
|
||||
if IsPasskeyDomainOption(option.Key) {
|
||||
passkeyOptions[option.Key] = option.Value
|
||||
continue
|
||||
}
|
||||
err := updateOptionMap(option.Key, option.Value)
|
||||
if err != nil {
|
||||
common.SysLog("failed to update option map: " + err.Error())
|
||||
}
|
||||
}
|
||||
applyPasskeyDomainOptions(passkeyOptions)
|
||||
}
|
||||
|
||||
func SyncOptions(frequency int) {
|
||||
@@ -226,6 +241,13 @@ func validateOptionValue(key string, value string) error {
|
||||
}
|
||||
|
||||
func UpdateOption(key string, value string) error {
|
||||
if IsRequestPolicyOption(key) {
|
||||
return UpdateRequestPolicyOptions(map[string]string{key: value})
|
||||
}
|
||||
if IsPasskeyDomainOption(key) {
|
||||
_, err := UpdatePasskeyDomainOptions(map[string]string{key: value}, false, "")
|
||||
return err
|
||||
}
|
||||
if IsModelPricingOption(key) {
|
||||
return UpdateModelPricingOptions(map[string]string{key: value})
|
||||
}
|
||||
@@ -256,11 +278,37 @@ func UpdateOptionsBulk(values map[string]string) error {
|
||||
if len(values) == 0 {
|
||||
return nil
|
||||
}
|
||||
for key := range values {
|
||||
if IsPasskeyDomainOption(key) {
|
||||
_, err := UpdatePasskeyDomainOptions(values, false, "")
|
||||
return err
|
||||
}
|
||||
}
|
||||
for key, value := range values {
|
||||
if err := validateOptionValue(key, value); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
var policySnapshot *RequestPolicySnapshot
|
||||
for key := range values {
|
||||
if IsRequestPolicyOption(key) {
|
||||
requestPolicyOptionMutex.Lock()
|
||||
defer requestPolicyOptionMutex.Unlock()
|
||||
options := maps.Clone(CurrentRequestPolicy().Options)
|
||||
for key, value := range values {
|
||||
if IsRequestPolicyOption(key) {
|
||||
options[key] = value
|
||||
}
|
||||
}
|
||||
var err error
|
||||
policySnapshot, err = BuildRequestPolicy(options)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
err := DB.Transaction(func(tx *gorm.DB) error {
|
||||
for k, v := range values {
|
||||
option := Option{Key: k}
|
||||
@@ -282,6 +330,9 @@ func UpdateOptionsBulk(values map[string]string) error {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if policySnapshot != nil {
|
||||
requestPolicySnapshot.Store(policySnapshot)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,218 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
const (
|
||||
optionPrimaryKeyTmpTable = "options_pk_tmp"
|
||||
optionPrimaryKeyLockName = "new_api_options_pk"
|
||||
optionPrimaryKeyLockID = 75820193
|
||||
optionLegacyTablePrefix = "options_legacy_"
|
||||
)
|
||||
|
||||
func migrateOptionPrimaryKey(db *gorm.DB) error {
|
||||
if db == nil {
|
||||
return fmt.Errorf("migrate options primary key: database is nil")
|
||||
}
|
||||
if !db.Migrator().HasTable(&Option{}) {
|
||||
return nil
|
||||
}
|
||||
unique, err := optionsKeyIsUnique(db)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if unique {
|
||||
return nil
|
||||
}
|
||||
return withOptionPrimaryKeyLock(db, func(locked *gorm.DB) error {
|
||||
unique, err := optionsKeyIsUnique(locked)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if unique {
|
||||
return nil
|
||||
}
|
||||
return repairOptionPrimaryKey(locked)
|
||||
})
|
||||
}
|
||||
|
||||
func optionsKeyIsUnique(db *gorm.DB) (bool, error) {
|
||||
indexes, err := db.Migrator().GetIndexes(&Option{})
|
||||
if err != nil {
|
||||
return false, fmt.Errorf("inspect options indexes: %w", err)
|
||||
}
|
||||
for _, index := range indexes {
|
||||
if len(index.Columns()) != 1 || index.Columns()[0] != "key" {
|
||||
continue
|
||||
}
|
||||
if unique, ok := index.Unique(); ok && unique {
|
||||
return true, nil
|
||||
}
|
||||
if primary, ok := index.PrimaryKey(); ok && primary {
|
||||
return true, nil
|
||||
}
|
||||
}
|
||||
if db.Dialector.Name() != "postgres" {
|
||||
return false, nil
|
||||
}
|
||||
// PostgreSQL GetIndexes omits constraint-backed primary/unique keys.
|
||||
var count int64
|
||||
if err := db.Raw(`
|
||||
SELECT count(*)
|
||||
FROM pg_catalog.pg_constraint AS constraint_meta
|
||||
JOIN pg_catalog.pg_attribute AS attribute_meta
|
||||
ON attribute_meta.attrelid = constraint_meta.conrelid
|
||||
AND attribute_meta.attnum = constraint_meta.conkey[1]
|
||||
WHERE constraint_meta.conrelid = to_regclass('options')
|
||||
AND constraint_meta.contype IN ('p', 'u')
|
||||
AND cardinality(constraint_meta.conkey) = 1
|
||||
AND attribute_meta.attname = 'key'`).Scan(&count).Error; err != nil {
|
||||
return false, fmt.Errorf("inspect options constraints: %w", err)
|
||||
}
|
||||
return count > 0, nil
|
||||
}
|
||||
|
||||
func withOptionPrimaryKeyLock(db *gorm.DB, fn func(*gorm.DB) error) error {
|
||||
switch db.Dialector.Name() {
|
||||
case "mysql":
|
||||
sqlDB, err := db.DB()
|
||||
if err != nil {
|
||||
return fmt.Errorf("lock options table: %w", err)
|
||||
}
|
||||
ctx := context.Background()
|
||||
conn, err := sqlDB.Conn(ctx)
|
||||
if err != nil {
|
||||
return fmt.Errorf("lock options table: %w", err)
|
||||
}
|
||||
defer conn.Close()
|
||||
var acquired int
|
||||
if err := conn.QueryRowContext(ctx, "SELECT GET_LOCK(?, 60)", optionPrimaryKeyLockName).Scan(&acquired); err != nil {
|
||||
return fmt.Errorf("lock options table: %w", err)
|
||||
}
|
||||
if acquired != 1 {
|
||||
return fmt.Errorf("lock options table: timeout")
|
||||
}
|
||||
defer conn.ExecContext(ctx, "SELECT RELEASE_LOCK(?)", optionPrimaryKeyLockName)
|
||||
return fn(db)
|
||||
case "postgres":
|
||||
return db.Transaction(func(tx *gorm.DB) error {
|
||||
if err := tx.Exec("SELECT pg_advisory_xact_lock(?)", optionPrimaryKeyLockID).Error; err != nil {
|
||||
return fmt.Errorf("lock options table: %w", err)
|
||||
}
|
||||
return fn(tx)
|
||||
})
|
||||
default:
|
||||
return db.Transaction(func(tx *gorm.DB) error {
|
||||
return fn(tx)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func repairOptionPrimaryKey(db *gorm.DB) error {
|
||||
var rows []Option
|
||||
if err := db.Find(&rows).Error; err != nil {
|
||||
return fmt.Errorf("read options rows: %w", err)
|
||||
}
|
||||
deduped, skippedEmpty := dedupeOptionRows(rows)
|
||||
if db.Migrator().HasTable(optionPrimaryKeyTmpTable) {
|
||||
if err := db.Migrator().DropTable(optionPrimaryKeyTmpTable); err != nil {
|
||||
return fmt.Errorf("drop leftover %s: %w", optionPrimaryKeyTmpTable, err)
|
||||
}
|
||||
}
|
||||
if err := db.Table(optionPrimaryKeyTmpTable).Migrator().CreateTable(&Option{}); err != nil {
|
||||
return fmt.Errorf("create %s: %w", optionPrimaryKeyTmpTable, err)
|
||||
}
|
||||
tmpCreated := true
|
||||
defer func() {
|
||||
if tmpCreated && db.Migrator().HasTable(optionPrimaryKeyTmpTable) {
|
||||
_ = db.Migrator().DropTable(optionPrimaryKeyTmpTable)
|
||||
}
|
||||
}()
|
||||
if len(deduped) > 0 {
|
||||
if err := db.Table(optionPrimaryKeyTmpTable).CreateInBatches(deduped, 100).Error; err != nil {
|
||||
return fmt.Errorf("insert rebuilt options: %w", err)
|
||||
}
|
||||
}
|
||||
var written int64
|
||||
if err := db.Table(optionPrimaryKeyTmpTable).Count(&written).Error; err != nil {
|
||||
return fmt.Errorf("count rebuilt options: %w", err)
|
||||
}
|
||||
if int(written) != len(deduped) {
|
||||
return fmt.Errorf("options rebuild wrote %d rows, want %d", written, len(deduped))
|
||||
}
|
||||
backup := fmt.Sprintf("%s%d", optionLegacyTablePrefix, time.Now().UnixNano())
|
||||
if err := swapOptionTables(db, optionPrimaryKeyTmpTable, backup); err != nil {
|
||||
return err
|
||||
}
|
||||
tmpCreated = false
|
||||
unique, err := optionsKeyIsUnique(db)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !unique {
|
||||
return fmt.Errorf("options table still has no unique key after rebuild")
|
||||
}
|
||||
common.SysLog(fmt.Sprintf("rebuilt options table with primary key from %d rows into %d keys; previous rows kept in %s", len(rows)-skippedEmpty, len(deduped), backup))
|
||||
return nil
|
||||
}
|
||||
|
||||
func dedupeOptionRows(rows []Option) ([]Option, int) {
|
||||
seen := make(map[string]int)
|
||||
deduped := make([]Option, 0, len(rows))
|
||||
skippedEmpty := 0
|
||||
for _, row := range rows {
|
||||
if row.Key == "" {
|
||||
skippedEmpty++
|
||||
continue
|
||||
}
|
||||
if i, ok := seen[row.Key]; ok {
|
||||
if deduped[i].Value != row.Value {
|
||||
common.SysError(fmt.Sprintf("options key %q had conflicting values; keeping the last row", row.Key))
|
||||
}
|
||||
deduped[i].Value = row.Value
|
||||
continue
|
||||
}
|
||||
seen[row.Key] = len(deduped)
|
||||
deduped = append(deduped, Option{Key: row.Key, Value: row.Value})
|
||||
}
|
||||
if skippedEmpty > 0 {
|
||||
common.SysError(fmt.Sprintf("skipped %d options rows with empty keys while rebuilding the primary key", skippedEmpty))
|
||||
}
|
||||
return deduped, skippedEmpty
|
||||
}
|
||||
|
||||
func swapOptionTables(db *gorm.DB, tmp, backup string) error {
|
||||
if !optionSafeIdent(tmp) || !optionSafeIdent(backup) {
|
||||
return fmt.Errorf("unsafe options table name")
|
||||
}
|
||||
if db.Dialector.Name() == "mysql" {
|
||||
sql := "RENAME TABLE `options` TO `" + backup + "`, `" + tmp + "` TO `options`"
|
||||
if err := db.Exec(sql).Error; err != nil {
|
||||
return fmt.Errorf("rename options tables: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
if err := db.Migrator().RenameTable("options", backup); err != nil {
|
||||
return fmt.Errorf("rename options to %s: %w", backup, err)
|
||||
}
|
||||
if err := db.Migrator().RenameTable(tmp, "options"); err != nil {
|
||||
return fmt.Errorf("rename %s to options: %w", tmp, err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func optionSafeIdent(name string) bool {
|
||||
if name == "" {
|
||||
return false
|
||||
}
|
||||
return !strings.ContainsFunc(name, func(r rune) bool {
|
||||
return r != '_' && (r < '0' || r > '9') && (r < 'a' || r > 'z')
|
||||
})
|
||||
}
|
||||
+52
-20
@@ -8,6 +8,7 @@ import (
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/setting/system_setting"
|
||||
|
||||
"github.com/go-webauthn/webauthn/protocol"
|
||||
"github.com/go-webauthn/webauthn/webauthn"
|
||||
@@ -22,6 +23,7 @@ var (
|
||||
type PasskeyCredential struct {
|
||||
ID int `json:"id" gorm:"primaryKey"`
|
||||
UserID int `json:"user_id" gorm:"uniqueIndex;not null"`
|
||||
RPID *string `json:"rp_id,omitempty" gorm:"column:rp_id;type:varchar(253)"`
|
||||
CredentialID string `json:"credential_id" gorm:"type:varchar(512);uniqueIndex;not null"` // base64 encoded
|
||||
PublicKey string `json:"public_key" gorm:"type:text;not null"` // base64 encoded
|
||||
AttestationType string `json:"attestation_type" gorm:"type:varchar(255)"`
|
||||
@@ -161,29 +163,49 @@ func GetPasskeyByCredentialID(credentialID []byte) (*PasskeyCredential, error) {
|
||||
// UpdatePasskeyAssertionState persists only fields produced by a successful
|
||||
// assertion. Registration identity (credential ID, public key, AAGUID,
|
||||
// transports and attestation metadata) is immutable on this path.
|
||||
func UpdatePasskeyAssertionState(userID int, credential *webauthn.Credential, lastUsedAt time.Time) error {
|
||||
if userID <= 0 || credential == nil || len(credential.ID) == 0 || lastUsedAt.IsZero() {
|
||||
func UpdatePasskeyAssertionState(userID int, credential *webauthn.Credential, lastUsedAt time.Time, rpID string) error {
|
||||
if userID <= 0 || credential == nil || len(credential.ID) == 0 || lastUsedAt.IsZero() || rpID == "" {
|
||||
return fmt.Errorf("Passkey 保存失败,请重试")
|
||||
}
|
||||
credentialID := base64.StdEncoding.EncodeToString(credential.ID)
|
||||
result := DB.Model(&PasskeyCredential{}).
|
||||
Where("user_id = ? AND credential_id = ?", userID, credentialID).
|
||||
Updates(map[string]any{
|
||||
"sign_count": credential.Authenticator.SignCount,
|
||||
"clone_warning": credential.Authenticator.CloneWarning,
|
||||
"user_present": credential.Flags.UserPresent,
|
||||
"user_verified": credential.Flags.UserVerified,
|
||||
"backup_eligible": credential.Flags.BackupEligible,
|
||||
"backup_state": credential.Flags.BackupState,
|
||||
"last_used_at": lastUsedAt,
|
||||
})
|
||||
if result.Error != nil {
|
||||
return result.Error
|
||||
}
|
||||
if result.RowsAffected != 1 {
|
||||
return ErrPasskeyNotFound
|
||||
}
|
||||
return nil
|
||||
passkeyOptionMutex.Lock()
|
||||
defer passkeyOptionMutex.Unlock()
|
||||
return DB.Transaction(func(tx *gorm.DB) error {
|
||||
if err := validatePasskeyRPIDWithTx(tx, rpID); err != nil {
|
||||
return err
|
||||
}
|
||||
var stored PasskeyCredential
|
||||
if err := lockForUpdate(tx).Where("user_id = ? AND credential_id = ?", userID, credentialID).First(&stored).Error; err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return ErrPasskeyNotFound
|
||||
}
|
||||
return err
|
||||
}
|
||||
// Compare in Go: the database collation may fold historical case.
|
||||
if stored.RPID != nil && *stored.RPID != "" && *stored.RPID != rpID {
|
||||
return system_setting.ErrPasskeyRPIDUnavailable
|
||||
}
|
||||
result := tx.Model(&PasskeyCredential{}).
|
||||
Where("user_id = ? AND credential_id = ?", userID, credentialID).
|
||||
Where("rp_id IS NULL OR rp_id = ? OR rp_id = ?", "", rpID).
|
||||
Updates(map[string]any{
|
||||
"rp_id": rpID,
|
||||
"sign_count": credential.Authenticator.SignCount,
|
||||
"clone_warning": credential.Authenticator.CloneWarning,
|
||||
"user_present": credential.Flags.UserPresent,
|
||||
"user_verified": credential.Flags.UserVerified,
|
||||
"backup_eligible": credential.Flags.BackupEligible,
|
||||
"backup_state": credential.Flags.BackupState,
|
||||
"last_used_at": lastUsedAt,
|
||||
})
|
||||
if result.Error != nil {
|
||||
return result.Error
|
||||
}
|
||||
if result.RowsAffected != 1 {
|
||||
return ErrPasskeyNotFound
|
||||
}
|
||||
return nil
|
||||
})
|
||||
}
|
||||
|
||||
func upsertPasskeyCredentialWithTx(tx *gorm.DB, credential *PasskeyCredential) error {
|
||||
@@ -212,7 +234,17 @@ func upsertPasskeyCredentialWithAuthVersion(credential *PasskeyCredential, ident
|
||||
if credential == nil || credential.UserID <= 0 {
|
||||
return fmt.Errorf("Passkey 保存失败,请重试")
|
||||
}
|
||||
passkeyOptionMutex.Lock()
|
||||
defer passkeyOptionMutex.Unlock()
|
||||
if err := DB.Transaction(func(tx *gorm.DB) error {
|
||||
if identity != nil && (credential.RPID == nil || *credential.RPID == "") {
|
||||
return system_setting.ErrPasskeyRPIDUnavailable
|
||||
}
|
||||
if credential.RPID != nil && *credential.RPID != "" {
|
||||
if err := validatePasskeyRPIDWithTx(tx, *credential.RPID); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if identity != nil {
|
||||
if identity.UserID != credential.UserID {
|
||||
return ErrUserSessionInactive
|
||||
|
||||
@@ -0,0 +1,270 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"crypto/hmac"
|
||||
"errors"
|
||||
"maps"
|
||||
"slices"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/setting/system_setting"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
)
|
||||
|
||||
var passkeyOptionMutex sync.Mutex
|
||||
var errPasskeyDomainPreview = errors.New("passkey domain preview")
|
||||
var ErrPasskeyDomainRemovalConfirmation = errors.New("Review the affected Passkeys and confirm the domain removal.")
|
||||
|
||||
// PasskeyDomainChange is both the preview and the authoritative saved result.
|
||||
// Unknown credentials may belong to a removed domain; they are not a known loss.
|
||||
type PasskeyDomainChange struct {
|
||||
RPID string `json:"rp_id"`
|
||||
LegacyRPIDs string `json:"legacy_rp_ids"`
|
||||
Origins string `json:"origins"`
|
||||
PreviousRPID string `json:"previous_rp_id"`
|
||||
EffectiveRPID string `json:"effective_rp_id"`
|
||||
RemovedRPIDs []string `json:"removed_rp_ids"`
|
||||
AffectedCredentials int64 `json:"affected_credentials"`
|
||||
UnknownCredentials int64 `json:"unknown_credentials"`
|
||||
ConfirmationRequired bool `json:"confirmation_required"`
|
||||
RemovalConfirmation string `json:"removal_confirmation"`
|
||||
}
|
||||
|
||||
type PasskeyDomainRemovalError struct {
|
||||
Change *PasskeyDomainChange
|
||||
}
|
||||
|
||||
func (e *PasskeyDomainRemovalError) Error() string {
|
||||
return ErrPasskeyDomainRemovalConfirmation.Error()
|
||||
}
|
||||
func (e *PasskeyDomainRemovalError) Unwrap() error { return ErrPasskeyDomainRemovalConfirmation }
|
||||
|
||||
// ServerAddress and origins participate because either can change an implicit RP ID.
|
||||
func IsPasskeyDomainOption(key string) bool {
|
||||
return key == "passkey.rp_id" || key == "passkey.legacy_rp_ids" || key == "passkey.origins" || key == "ServerAddress"
|
||||
}
|
||||
|
||||
// lockPasskeyDomainSettings serializes trust changes and credential writes across
|
||||
// nodes. Callers also hold passkeyOptionMutex for SQLite and local cache ordering.
|
||||
// Always acquire these option rows before any user/session/credential row locks.
|
||||
func lockPasskeyDomainSettings(tx *gorm.DB) (system_setting.PasskeySettings, string, error) {
|
||||
common.OptionMapRWMutex.RLock()
|
||||
settings, serverAddress := *system_setting.GetPasskeySettings(), system_setting.ServerAddress
|
||||
common.OptionMapRWMutex.RUnlock()
|
||||
rows := []Option{
|
||||
{Key: "ServerAddress", Value: serverAddress},
|
||||
{Key: "passkey.legacy_rp_ids", Value: settings.LegacyRPIDs},
|
||||
{Key: "passkey.origins", Value: settings.Origins},
|
||||
{Key: "passkey.rp_id", Value: settings.RPID},
|
||||
}
|
||||
// The first write also acquires SQLite's writer lock before any reads.
|
||||
if err := tx.Clauses(clause.OnConflict{DoNothing: true}).Create(&rows).Error; err != nil {
|
||||
return settings, serverAddress, err
|
||||
}
|
||||
for i := range rows {
|
||||
if err := lockForUpdate(tx).Where(&Option{Key: rows[i].Key}).First(&rows[i]).Error; err != nil {
|
||||
return settings, serverAddress, err
|
||||
}
|
||||
}
|
||||
serverAddress = rows[0].Value
|
||||
settings.LegacyRPIDs, settings.Origins, settings.RPID = rows[1].Value, rows[2].Value, rows[3].Value
|
||||
return settings, serverAddress, nil
|
||||
}
|
||||
|
||||
func validatePasskeyRPIDWithTx(tx *gorm.DB, rpID string) error {
|
||||
settings, serverAddress, err := lockPasskeyDomainSettings(tx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
// Preserve the existing request-derived default only when no domain or
|
||||
// origin has ever been configured. Finish still checks the signed origin.
|
||||
if settings.RPID == "" && settings.LegacyRPIDs == "" && settings.Origins == "" && serverAddress == "" {
|
||||
return nil
|
||||
}
|
||||
if !slices.Contains(settings.WithDefaults(serverAddress).RelyingPartyIDs(), rpID) {
|
||||
return system_setting.ErrPasskeyRPIDUnavailable
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// UpdatePasskeyDomainOptions keeps every supplied option in one transaction,
|
||||
// including unrelated keys from UpdateOptionsBulk. Preview rolls back even the
|
||||
// initial default rows and never publishes a local configuration change.
|
||||
func UpdatePasskeyDomainOptions(values map[string]string, preview bool, confirmation string) (*PasskeyDomainChange, error) {
|
||||
passkeyOptionMutex.Lock()
|
||||
defer passkeyOptionMutex.Unlock()
|
||||
values = maps.Clone(values)
|
||||
for key, value := range values {
|
||||
if err := validateOptionValue(key, value); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
var change *PasskeyDomainChange
|
||||
err := DB.Transaction(func(tx *gorm.DB) error {
|
||||
settings, serverAddress, err := lockPasskeyDomainSettings(tx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
previous := settings.WithDefaults(serverAddress)
|
||||
before := [4]string{serverAddress, settings.RPID, settings.LegacyRPIDs, settings.Origins}
|
||||
if value, changed := values["ServerAddress"]; changed {
|
||||
serverAddress = value
|
||||
}
|
||||
if value, changed := values["passkey.origins"]; changed {
|
||||
origins := []string{}
|
||||
for _, origin := range strings.FieldsFunc(value, func(c rune) bool { return c == ',' || c == '\n' || c == '\r' }) {
|
||||
origin = strings.TrimSpace(origin)
|
||||
if origin != "" && !slices.Contains(origins, origin) {
|
||||
origins = append(origins, origin)
|
||||
}
|
||||
}
|
||||
settings.Origins = strings.Join(origins, ",")
|
||||
} else if before[0] != serverAddress && (settings.Origins == "" || settings.Origins == "[]") {
|
||||
// Retain the previously trusted implicit website as well as the new
|
||||
// default, so retaining its RP ID does not leave its keys unusable.
|
||||
origins := []string{}
|
||||
for _, origin := range []string{previous.Origins, settings.WithDefaults(serverAddress).Origins} {
|
||||
if origin != "" && !slices.Contains(origins, origin) {
|
||||
origins = append(origins, origin)
|
||||
}
|
||||
}
|
||||
settings.Origins = strings.Join(origins, ",")
|
||||
}
|
||||
if value, changed := values["passkey.rp_id"]; changed {
|
||||
settings.RPID = strings.TrimSpace(value)
|
||||
if settings.RPID != "" {
|
||||
settings.RPID, err = system_setting.NormalizePasskeyRPID(settings.RPID, settings.WithDefaults(serverAddress).Origins)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
if value, changed := values["passkey.legacy_rp_ids"]; changed {
|
||||
settings.LegacyRPIDs = value
|
||||
}
|
||||
effective := settings.WithDefaults(serverAddress)
|
||||
legacy, err := system_setting.ParsePasskeyRPIDs(settings.LegacyRPIDs, effective.Origins)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
oldRPID, nextRPID := previous.EffectiveRPID(), effective.EffectiveRPID()
|
||||
_, previousDomainError := system_setting.NormalizePasskeyRPID(oldRPID, previous.Origins)
|
||||
if previousDomainError == nil && oldRPID != nextRPID && !slices.Contains(legacy, oldRPID) {
|
||||
legacy = append(legacy, oldRPID)
|
||||
}
|
||||
settings.LegacyRPIDs = strings.Join(legacy, ",")
|
||||
// Validate retained single-label domains against the complete proposed
|
||||
// origin set too. Do not normalize their historical hash inputs.
|
||||
if _, err := system_setting.ParsePasskeyRPIDs(settings.LegacyRPIDs, effective.Origins); err != nil {
|
||||
return err
|
||||
}
|
||||
// ServerAddress also supports non-WebAuthn deployments (for example an
|
||||
// IP during initial setup). Restrict explicit RP IDs, not that URL.
|
||||
if settings.RPID != "" {
|
||||
if _, err := system_setting.NormalizePasskeyRPID(nextRPID, effective.Origins); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
change = &PasskeyDomainChange{
|
||||
RPID: settings.RPID, LegacyRPIDs: settings.LegacyRPIDs, Origins: settings.Origins,
|
||||
PreviousRPID: oldRPID, EffectiveRPID: nextRPID, RemovedRPIDs: []string{},
|
||||
}
|
||||
nextIDs := settings.WithDefaults(serverAddress).RelyingPartyIDs()
|
||||
for _, id := range previous.RelyingPartyIDs() {
|
||||
if !slices.Contains(nextIDs, id) {
|
||||
change.RemovedRPIDs = append(change.RemovedRPIDs, id)
|
||||
}
|
||||
}
|
||||
if len(change.RemovedRPIDs) > 0 {
|
||||
// MySQL may compare/group text case-insensitively. Select only
|
||||
// candidate RP values, then count exact bytes without credential IDs.
|
||||
rows, err := tx.Model(&PasskeyCredential{}).Select("rp_id").
|
||||
Where("rp_id IS NULL OR rp_id = ? OR rp_id IN ?", "", change.RemovedRPIDs).Rows()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer rows.Close()
|
||||
for rows.Next() {
|
||||
var rpID *string
|
||||
if err := rows.Scan(&rpID); err != nil {
|
||||
return err
|
||||
}
|
||||
if rpID == nil || *rpID == "" {
|
||||
change.UnknownCredentials++
|
||||
} else if slices.Contains(change.RemovedRPIDs, *rpID) {
|
||||
change.AffectedCredentials++
|
||||
}
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := rows.Close(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
change.ConfirmationRequired = change.AffectedCredentials > 0 || change.UnknownCredentials > 0
|
||||
// This is an intent/freshness binding, not an authentication credential.
|
||||
// Root authorization remains mandatory. No map iteration affects its hash.
|
||||
payload, err := common.Marshal(struct {
|
||||
Before [4]string
|
||||
After [4]string
|
||||
Removed []string
|
||||
Known, Unknown int64
|
||||
}{before, [4]string{serverAddress, settings.RPID, settings.LegacyRPIDs, settings.Origins}, change.RemovedRPIDs, change.AffectedCredentials, change.UnknownCredentials})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
change.RemovalConfirmation = common.GenerateHMACWithKey([]byte("passkey-domains-v1:"+common.SessionSecret), string(payload))
|
||||
if preview {
|
||||
return errPasskeyDomainPreview
|
||||
}
|
||||
if (change.ConfirmationRequired || confirmation != "") && !hmac.Equal([]byte(confirmation), []byte(change.RemovalConfirmation)) {
|
||||
return &PasskeyDomainRemovalError{Change: change}
|
||||
}
|
||||
values["passkey.rp_id"], values["passkey.legacy_rp_ids"], values["passkey.origins"] = settings.RPID, settings.LegacyRPIDs, settings.Origins
|
||||
values["ServerAddress"] = serverAddress
|
||||
keys := slices.Sorted(maps.Keys(values))
|
||||
for _, key := range keys {
|
||||
if err := tx.Save(&Option{Key: key, Value: values[key]}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
})
|
||||
if errors.Is(err, errPasskeyDomainPreview) {
|
||||
return change, nil
|
||||
}
|
||||
if err != nil {
|
||||
return change, err
|
||||
}
|
||||
applyPasskeyDomainOptions(values)
|
||||
for key, value := range values {
|
||||
if !IsPasskeyDomainOption(key) {
|
||||
if err := updateOptionMap(key, value); err != nil {
|
||||
return change, err
|
||||
}
|
||||
}
|
||||
}
|
||||
return change, nil
|
||||
}
|
||||
|
||||
func applyPasskeyDomainOptions(values map[string]string) {
|
||||
common.OptionMapRWMutex.Lock()
|
||||
defer common.OptionMapRWMutex.Unlock()
|
||||
if common.OptionMap == nil {
|
||||
common.OptionMap = make(map[string]string)
|
||||
}
|
||||
for _, key := range []string{"ServerAddress", "passkey.legacy_rp_ids", "passkey.origins", "passkey.rp_id"} {
|
||||
if value, ok := values[key]; ok {
|
||||
common.OptionMap[key] = value
|
||||
if key == "ServerAddress" {
|
||||
system_setting.ServerAddress = value
|
||||
} else {
|
||||
handleConfigUpdate(key, value)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -8,7 +8,6 @@ import (
|
||||
)
|
||||
|
||||
const prefillGroupNameIndex = "uk_prefill_name"
|
||||
const legacyPrefillGroupNameUnique = "idx_prefill_groups_name"
|
||||
|
||||
type conflictingPrefillGroupUniqueness struct {
|
||||
constraints []string
|
||||
@@ -24,30 +23,6 @@ func (conflicts conflictingPrefillGroupUniqueness) empty() bool {
|
||||
return len(conflicts.constraints) == 0 && len(conflicts.indexes) == 0
|
||||
}
|
||||
|
||||
func (conflicts conflictingPrefillGroupUniqueness) validateAutomaticMigrationScope() error {
|
||||
unexpectedConstraints := make([]string, 0)
|
||||
for _, name := range conflicts.constraints {
|
||||
if name != legacyPrefillGroupNameUnique {
|
||||
unexpectedConstraints = append(unexpectedConstraints, name)
|
||||
}
|
||||
}
|
||||
unexpectedIndexes := make([]string, 0)
|
||||
for _, name := range conflicts.indexes {
|
||||
if name != legacyPrefillGroupNameUnique {
|
||||
unexpectedIndexes = append(unexpectedIndexes, name)
|
||||
}
|
||||
}
|
||||
if len(unexpectedConstraints) == 0 && len(unexpectedIndexes) == 0 {
|
||||
return nil
|
||||
}
|
||||
return fmt.Errorf(
|
||||
"prefill_groups.name has unsupported global unique constraints %q and indexes %q; only legacy object %q can be migrated automatically to partial uniqueness",
|
||||
unexpectedConstraints,
|
||||
unexpectedIndexes,
|
||||
legacyPrefillGroupNameUnique,
|
||||
)
|
||||
}
|
||||
|
||||
func inspectConflictingPrefillGroupUniqueness(db *gorm.DB, tableName string) (conflictingPrefillGroupUniqueness, error) {
|
||||
var conflicts conflictingPrefillGroupUniqueness
|
||||
if err := db.Raw(`
|
||||
@@ -124,9 +99,10 @@ WHERE index_meta.indrelid = to_regclass(?)
|
||||
return prefillGroupNameIndexState{exists: state.Exists, valid: state.Valid}, nil
|
||||
}
|
||||
|
||||
// migratePrefillGroupUniqueness replaces the known global PostgreSQL unique
|
||||
// object left by older GORM versions before AutoMigrate inspects the column.
|
||||
// Unknown conflicting objects are reported without being modified.
|
||||
// migratePrefillGroupUniqueness replaces global PostgreSQL uniqueness on name
|
||||
// before AutoMigrate inspects the column. Match the definition rather than the
|
||||
// object name, which can change across older schemas and database imports.
|
||||
// Composite, expression and partial indexes retain their separate semantics.
|
||||
func migratePrefillGroupUniqueness(db *gorm.DB) error {
|
||||
if db == nil {
|
||||
return fmt.Errorf("migrate prefill group uniqueness: database is nil")
|
||||
@@ -147,10 +123,6 @@ func migratePrefillGroupUniqueness(db *gorm.DB) error {
|
||||
if conflicts.empty() {
|
||||
return nil
|
||||
}
|
||||
if err := conflicts.validateAutomaticMigrationScope(); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return db.Transaction(func(tx *gorm.DB) error {
|
||||
migrator := tx.Migrator()
|
||||
if !migrator.HasTable(&PrefillGroup{}) {
|
||||
@@ -171,16 +143,26 @@ func migratePrefillGroupUniqueness(db *gorm.DB) error {
|
||||
if conflicts.empty() {
|
||||
return nil
|
||||
}
|
||||
if err := conflicts.validateAutomaticMigrationScope(); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if !migrator.HasColumn(&PrefillGroup{}, "DeletedAt") {
|
||||
if err := migrator.AddColumn(&PrefillGroup{}, "DeletedAt"); err != nil {
|
||||
return fmt.Errorf("add prefill groups deleted_at column: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
// A global index or constraint may already use the target index name.
|
||||
// Drop it before creating the replacement under the exclusive table lock;
|
||||
// the transaction restores all old objects if any migration step fails.
|
||||
for _, constraintName := range conflicts.constraints {
|
||||
if err := migrator.DropConstraint(&PrefillGroup{}, constraintName); err != nil {
|
||||
return fmt.Errorf("drop conflicting prefill group constraint %q: %w", constraintName, err)
|
||||
}
|
||||
}
|
||||
for _, indexName := range conflicts.indexes {
|
||||
if err := migrator.DropIndex(&PrefillGroup{}, indexName); err != nil {
|
||||
return fmt.Errorf("drop conflicting prefill group index %q: %w", indexName, err)
|
||||
}
|
||||
}
|
||||
|
||||
targetIndex, err := inspectPrefillGroupNameIndex(tx, tableName)
|
||||
if err != nil {
|
||||
return err
|
||||
@@ -198,17 +180,6 @@ func migratePrefillGroupUniqueness(db *gorm.DB) error {
|
||||
return fmt.Errorf("prefill group index %q has an unexpected definition", prefillGroupNameIndex)
|
||||
}
|
||||
|
||||
for _, constraintName := range conflicts.constraints {
|
||||
if err := migrator.DropConstraint(&PrefillGroup{}, constraintName); err != nil {
|
||||
return fmt.Errorf("drop conflicting prefill group constraint %q: %w", constraintName, err)
|
||||
}
|
||||
}
|
||||
for _, indexName := range conflicts.indexes {
|
||||
if err := migrator.DropIndex(&PrefillGroup{}, indexName); err != nil {
|
||||
return fmt.Errorf("drop conflicting prefill group index %q: %w", indexName, err)
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
})
|
||||
}
|
||||
@@ -3,6 +3,7 @@ package model
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"slices"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
@@ -18,6 +19,13 @@ import (
|
||||
|
||||
func testPrefillGroupMigrationNonPostgreSQL(t *testing.T, db *gorm.DB) {
|
||||
t.Helper()
|
||||
var version string
|
||||
versionQuery := "SELECT version()"
|
||||
if db.Dialector.Name() == "sqlite" {
|
||||
versionQuery = "SELECT sqlite_version()"
|
||||
}
|
||||
require.NoError(t, db.Raw(versionQuery).Scan(&version).Error)
|
||||
t.Logf("%s version: %s", db.Dialector.Name(), version)
|
||||
tableName := fmt.Sprintf("prefill_group_migration_%d", time.Now().UnixNano())
|
||||
t.Cleanup(func() { _ = db.Migrator().DropTable(tableName) })
|
||||
|
||||
@@ -30,15 +38,22 @@ func testPrefillGroupMigrationNonPostgreSQL(t *testing.T, db *gorm.DB) {
|
||||
Description: "preserve me",
|
||||
}).Error)
|
||||
|
||||
for range 2 {
|
||||
require.NoError(t, migratePrefillGroupUniqueness(db))
|
||||
require.NoError(t, tableDB.AutoMigrate(&PrefillGroup{}))
|
||||
recorder := &migrationSQLRecorder{}
|
||||
for pass := range 2 {
|
||||
recorder.reset()
|
||||
require.NoError(t, migratePrefillGroupUniqueness(db.Session(&gorm.Session{Logger: recorder})))
|
||||
require.NoError(t, tableDB.Session(&gorm.Session{Logger: recorder}).AutoMigrate(&PrefillGroup{}))
|
||||
if pass == 1 {
|
||||
assert.Empty(t, recorder.schemaMutations(), "repeated startup must not change the schema")
|
||||
}
|
||||
}
|
||||
|
||||
var preserved PrefillGroup
|
||||
require.NoError(t, tableDB.Where("name = ?", "preserved-name").First(&preserved).Error)
|
||||
assert.Equal(t, "preserve me", preserved.Description)
|
||||
assert.JSONEq(t, `["gpt-test"]`, string(preserved.Items))
|
||||
assert.True(t, tableDB.Migrator().HasIndex(&PrefillGroup{}, prefillGroupNameIndex))
|
||||
require.Error(t, tableDB.Create(&PrefillGroup{Name: preserved.Name, Type: "model"}).Error)
|
||||
}
|
||||
|
||||
func TestMigratePrefillGroupUniquenessSQLite(t *testing.T) {
|
||||
@@ -76,83 +91,64 @@ func TestMigratePrefillGroupUniquenessPostgreSQL(t *testing.T) {
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(func() { require.NoError(t, sqlDB.Close()) })
|
||||
|
||||
var version string
|
||||
require.NoError(t, db.Raw("SELECT version()").Scan(&version).Error)
|
||||
t.Log(version)
|
||||
|
||||
type indexDefinition struct {
|
||||
Name string `gorm:"column:indexname"`
|
||||
Definition string `gorm:"column:indexdef"`
|
||||
}
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
prepareOld func(*testing.T, *gorm.DB)
|
||||
blockedConstraints []string
|
||||
blockedIndexes []string
|
||||
preservedIndexes []string
|
||||
name string
|
||||
constraints []string
|
||||
indexes []string
|
||||
replacePartialIndex bool
|
||||
withoutDeletedAt bool
|
||||
prepareOld func(*testing.T, *gorm.DB)
|
||||
wantError string
|
||||
}{
|
||||
{name: "fresh"},
|
||||
{
|
||||
name: "legacy_constraint",
|
||||
name: "legacy_constraint",
|
||||
constraints: []string{"idx_prefill_groups_name"},
|
||||
},
|
||||
{
|
||||
name: "legacy_standalone_index",
|
||||
indexes: []string{"idx_prefill_groups_name"},
|
||||
replacePartialIndex: true,
|
||||
},
|
||||
{
|
||||
name: "renamed_constraints_and_indexes",
|
||||
constraints: []string{"idx_prefill_groups_name", "prefill_groups_name_key"},
|
||||
indexes: []string{"idx_37606_uk_prefill_name", `custom "name" index`},
|
||||
},
|
||||
{
|
||||
name: "imported_index_without_soft_delete_column",
|
||||
indexes: []string{"idx_37606_uk_prefill_name"},
|
||||
replacePartialIndex: true,
|
||||
withoutDeletedAt: true,
|
||||
},
|
||||
{
|
||||
name: "global_index_uses_target_name",
|
||||
indexes: []string{prefillGroupNameIndex},
|
||||
replacePartialIndex: true,
|
||||
},
|
||||
{
|
||||
name: "global_constraint_uses_target_name",
|
||||
constraints: []string{prefillGroupNameIndex},
|
||||
replacePartialIndex: true,
|
||||
},
|
||||
{
|
||||
name: "non_conflicting_indexes_are_preserved",
|
||||
constraints: []string{"prefill_groups_name_key"},
|
||||
prepareOld: func(t *testing.T, tx *gorm.DB) {
|
||||
t.Helper()
|
||||
require.NoError(t, tx.Exec(
|
||||
"ALTER TABLE ? ADD CONSTRAINT ? UNIQUE (?)",
|
||||
"CREATE INDEX ? ON ? (?)",
|
||||
clause.Column{Name: "keep_prefill_name"},
|
||||
clause.Table{Name: "prefill_groups"},
|
||||
clause.Column{Name: legacyPrefillGroupNameUnique},
|
||||
clause.Column{Name: "name"},
|
||||
).Error)
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "legacy_standalone_index",
|
||||
prepareOld: func(t *testing.T, tx *gorm.DB) {
|
||||
t.Helper()
|
||||
require.NoError(t, tx.Migrator().DropIndex(&PrefillGroup{}, prefillGroupNameIndex))
|
||||
require.NoError(t, tx.Exec(
|
||||
"CREATE UNIQUE INDEX ? ON ? (?)",
|
||||
clause.Column{Name: legacyPrefillGroupNameUnique},
|
||||
clause.Table{Name: "prefill_groups"},
|
||||
clause.Column{Name: "name"},
|
||||
).Error)
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "arbitrary_constraint_name",
|
||||
prepareOld: func(t *testing.T, tx *gorm.DB) {
|
||||
t.Helper()
|
||||
for _, constraintName := range []string{
|
||||
legacyPrefillGroupNameUnique,
|
||||
"prefill_groups_name_key",
|
||||
} {
|
||||
require.NoError(t, tx.Exec(
|
||||
"ALTER TABLE ? ADD CONSTRAINT ? UNIQUE (?)",
|
||||
clause.Table{Name: "prefill_groups"},
|
||||
clause.Column{Name: constraintName},
|
||||
clause.Column{Name: "name"},
|
||||
).Error)
|
||||
}
|
||||
},
|
||||
blockedConstraints: []string{legacyPrefillGroupNameUnique, "prefill_groups_name_key"},
|
||||
},
|
||||
{
|
||||
name: "arbitrary_index_name",
|
||||
prepareOld: func(t *testing.T, tx *gorm.DB) {
|
||||
t.Helper()
|
||||
for _, indexName := range []string{
|
||||
legacyPrefillGroupNameUnique,
|
||||
"prefill_groups_name_key",
|
||||
} {
|
||||
require.NoError(t, tx.Exec(
|
||||
"CREATE UNIQUE INDEX ? ON ? (?)",
|
||||
clause.Column{Name: indexName},
|
||||
clause.Table{Name: "prefill_groups"},
|
||||
clause.Column{Name: "name"},
|
||||
).Error)
|
||||
}
|
||||
},
|
||||
blockedIndexes: []string{legacyPrefillGroupNameUnique, "prefill_groups_name_key"},
|
||||
},
|
||||
{
|
||||
name: "non_conflicting_indexes_are_preserved",
|
||||
prepareOld: func(t *testing.T, tx *gorm.DB) {
|
||||
t.Helper()
|
||||
require.NoError(t, tx.Exec(
|
||||
"ALTER TABLE ? ADD CONSTRAINT ? UNIQUE (?)",
|
||||
clause.Table{Name: "prefill_groups"},
|
||||
clause.Column{Name: legacyPrefillGroupNameUnique},
|
||||
clause.Column{Name: "name"},
|
||||
).Error)
|
||||
require.NoError(t, tx.Exec(
|
||||
@@ -175,11 +171,49 @@ func TestMigratePrefillGroupUniquenessPostgreSQL(t *testing.T) {
|
||||
clause.Column{Name: "name"},
|
||||
).Error)
|
||||
},
|
||||
preservedIndexes: []string{
|
||||
"keep_prefill_name_deleted_at",
|
||||
"keep_prefill_lower_name",
|
||||
"keep_prefill_deleted_name",
|
||||
},
|
||||
{
|
||||
name: "unexpected_target_definition_rolls_back",
|
||||
constraints: []string{"prefill_groups_name_key"},
|
||||
indexes: []string{"idx_37606_uk_prefill_name"},
|
||||
replacePartialIndex: true,
|
||||
prepareOld: func(t *testing.T, tx *gorm.DB) {
|
||||
t.Helper()
|
||||
require.NoError(t, tx.Exec(
|
||||
"CREATE INDEX ? ON ? (?)",
|
||||
clause.Column{Name: prefillGroupNameIndex},
|
||||
clause.Table{Name: "prefill_groups"},
|
||||
clause.Column{Name: "type"},
|
||||
).Error)
|
||||
},
|
||||
wantError: "unexpected definition",
|
||||
},
|
||||
{
|
||||
name: "target_name_on_other_table_rolls_back",
|
||||
indexes: []string{"idx_37606_uk_prefill_name"},
|
||||
replacePartialIndex: true,
|
||||
withoutDeletedAt: true,
|
||||
prepareOld: func(t *testing.T, tx *gorm.DB) {
|
||||
t.Helper()
|
||||
require.NoError(t, tx.Exec("CREATE TABLE other_groups (name varchar(64))").Error)
|
||||
require.NoError(t, tx.Exec(
|
||||
"CREATE INDEX ? ON other_groups (name)",
|
||||
clause.Column{Name: prefillGroupNameIndex},
|
||||
).Error)
|
||||
},
|
||||
wantError: "unexpected definition",
|
||||
},
|
||||
{
|
||||
name: "foreign_key_dependency_rolls_back",
|
||||
constraints: []string{"prefill_groups_name_key"},
|
||||
replacePartialIndex: true,
|
||||
withoutDeletedAt: true,
|
||||
prepareOld: func(t *testing.T, tx *gorm.DB) {
|
||||
t.Helper()
|
||||
require.NoError(t, tx.Exec("CREATE TABLE referenced_groups (name varchar(64) REFERENCES prefill_groups(name))").Error)
|
||||
require.NoError(t, tx.Exec("INSERT INTO referenced_groups (name) VALUES (?)", "shared-name").Error)
|
||||
},
|
||||
wantError: "drop conflicting prefill group constraint",
|
||||
},
|
||||
}
|
||||
|
||||
@@ -208,66 +242,79 @@ func TestMigratePrefillGroupUniquenessPostgreSQL(t *testing.T) {
|
||||
Description: "preserve me",
|
||||
}
|
||||
require.NoError(t, tx.Create(&original).Error)
|
||||
if test.replacePartialIndex {
|
||||
require.NoError(t, tx.Migrator().DropIndex(&PrefillGroup{}, prefillGroupNameIndex))
|
||||
}
|
||||
if test.withoutDeletedAt {
|
||||
require.NoError(t, tx.Migrator().DropColumn(&PrefillGroup{}, "DeletedAt"))
|
||||
}
|
||||
for _, constraintName := range test.constraints {
|
||||
require.NoError(t, tx.Exec(
|
||||
"ALTER TABLE ? ADD CONSTRAINT ? UNIQUE (?)",
|
||||
clause.Table{Name: "prefill_groups"},
|
||||
clause.Column{Name: constraintName},
|
||||
clause.Column{Name: "name"},
|
||||
).Error)
|
||||
}
|
||||
for _, indexName := range test.indexes {
|
||||
require.NoError(t, tx.Exec(
|
||||
"CREATE UNIQUE INDEX ? ON ? (?)",
|
||||
clause.Column{Name: indexName},
|
||||
clause.Table{Name: "prefill_groups"},
|
||||
clause.Column{Name: "name"},
|
||||
).Error)
|
||||
}
|
||||
if test.prepareOld != nil {
|
||||
test.prepareOld(t, tx)
|
||||
}
|
||||
if len(test.blockedConstraints) > 0 || len(test.blockedIndexes) > 0 {
|
||||
var oldIndexes []indexDefinition
|
||||
require.NoError(t, tx.Raw(`
|
||||
SELECT indexname, indexdef FROM pg_catalog.pg_indexes
|
||||
WHERE schemaname = current_schema() AND tablename = 'prefill_groups'
|
||||
ORDER BY indexname`).Scan(&oldIndexes).Error)
|
||||
if test.wantError != "" {
|
||||
err := migratePrefillGroupUniqueness(tx)
|
||||
require.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "prefill_groups_name_key")
|
||||
for _, constraintName := range test.blockedConstraints {
|
||||
require.ErrorContains(t, err, test.wantError)
|
||||
for _, constraintName := range test.constraints {
|
||||
assert.True(t, tx.Migrator().HasConstraint(&PrefillGroup{}, constraintName))
|
||||
}
|
||||
for _, indexName := range test.blockedIndexes {
|
||||
assert.True(t, tx.Migrator().HasIndex(&PrefillGroup{}, indexName))
|
||||
}
|
||||
var restoredIndexes []indexDefinition
|
||||
require.NoError(t, tx.Raw(`
|
||||
SELECT indexname, indexdef FROM pg_catalog.pg_indexes
|
||||
WHERE schemaname = current_schema() AND tablename = 'prefill_groups'
|
||||
ORDER BY indexname`).Scan(&restoredIndexes).Error)
|
||||
assert.Equal(t, oldIndexes, restoredIndexes)
|
||||
assert.Equal(t, !test.withoutDeletedAt, tx.Migrator().HasColumn(&PrefillGroup{}, "DeletedAt"))
|
||||
var preserved PrefillGroup
|
||||
require.NoError(t, tx.Unscoped().First(&preserved, original.Id).Error)
|
||||
assert.Equal(t, original, preserved)
|
||||
return
|
||||
}
|
||||
|
||||
for range 2 {
|
||||
require.NoError(t, migratePrefillGroupUniqueness(tx))
|
||||
require.NoError(t, tx.AutoMigrate(&PrefillGroup{}))
|
||||
recorder := &migrationSQLRecorder{}
|
||||
migrationDB := tx.Session(&gorm.Session{Logger: recorder})
|
||||
for pass := range 2 {
|
||||
recorder.reset()
|
||||
require.NoError(t, migratePrefillGroupUniqueness(migrationDB))
|
||||
require.NoError(t, migrationDB.AutoMigrate(&PrefillGroup{}))
|
||||
if pass == 1 {
|
||||
assert.Empty(t, recorder.schemaMutations(), "repeated startup must not change the schema")
|
||||
}
|
||||
}
|
||||
for _, indexName := range test.preservedIndexes {
|
||||
assert.True(t, tx.Migrator().HasIndex(&PrefillGroup{}, indexName))
|
||||
for _, oldIndex := range oldIndexes {
|
||||
if oldIndex.Name == prefillGroupNameIndex || slices.Contains(test.constraints, oldIndex.Name) || slices.Contains(test.indexes, oldIndex.Name) {
|
||||
continue
|
||||
}
|
||||
var definition string
|
||||
require.NoError(t, tx.Raw(`
|
||||
SELECT indexdef FROM pg_catalog.pg_indexes
|
||||
WHERE schemaname = current_schema() AND indexname = ?`, oldIndex.Name).Scan(&definition).Error)
|
||||
assert.Equal(t, oldIndex.Definition, definition)
|
||||
}
|
||||
|
||||
var preserved PrefillGroup
|
||||
require.NoError(t, tx.First(&preserved, original.Id).Error)
|
||||
assert.Equal(t, original.Name, preserved.Name)
|
||||
assert.Equal(t, original.Description, preserved.Description)
|
||||
|
||||
var globalConstraintCount int64
|
||||
require.NoError(t, tx.Raw(`
|
||||
SELECT count(*)
|
||||
FROM pg_catalog.pg_constraint AS constraint_meta
|
||||
WHERE constraint_meta.conrelid = to_regclass('prefill_groups')
|
||||
AND constraint_meta.contype = 'u'
|
||||
AND cardinality(constraint_meta.conkey) = 1
|
||||
AND EXISTS (
|
||||
SELECT 1
|
||||
FROM pg_catalog.pg_attribute AS attribute_meta
|
||||
WHERE attribute_meta.attrelid = constraint_meta.conrelid
|
||||
AND attribute_meta.attnum = constraint_meta.conkey[1]
|
||||
AND attribute_meta.attname = 'name'
|
||||
)`).Scan(&globalConstraintCount).Error)
|
||||
assert.Zero(t, globalConstraintCount)
|
||||
|
||||
var globalIndexCount int64
|
||||
require.NoError(t, tx.Raw(`
|
||||
SELECT count(*)
|
||||
FROM pg_catalog.pg_index AS index_meta
|
||||
JOIN pg_catalog.pg_attribute AS attribute_meta
|
||||
ON attribute_meta.attrelid = index_meta.indrelid
|
||||
AND attribute_meta.attnum = index_meta.indkey[0]
|
||||
WHERE index_meta.indrelid = to_regclass('prefill_groups')
|
||||
AND index_meta.indisunique
|
||||
AND NOT index_meta.indisprimary
|
||||
AND index_meta.indpred IS NULL
|
||||
AND index_meta.indexprs IS NULL
|
||||
AND index_meta.indnatts = 1
|
||||
AND attribute_meta.attname = 'name'`).Scan(&globalIndexCount).Error)
|
||||
assert.Zero(t, globalIndexCount)
|
||||
assert.Equal(t, original, preserved)
|
||||
|
||||
var targetIndexDefinition string
|
||||
require.NoError(t, tx.Raw(`
|
||||
|
||||
+43
-24
@@ -2,9 +2,7 @@ package model
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"maps"
|
||||
"strings"
|
||||
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
@@ -17,7 +15,18 @@ import (
|
||||
"github.com/QuantumNous/new-api/types"
|
||||
)
|
||||
|
||||
type PricingPluginVariant struct {
|
||||
PluginKey string `json:"plugin_key"`
|
||||
PluginName string `json:"plugin_name"`
|
||||
Icon string `json:"icon,omitempty"`
|
||||
BillingExpr string `json:"billing_expr"`
|
||||
BillingMode string `json:"billing_mode"`
|
||||
BillingUsageSchema map[string]jsplugin.UsageFieldSchema `json:"billing_usage_schema"`
|
||||
BillingUsageExamples []jsplugin.UsageExample `json:"billing_usage_examples,omitempty"`
|
||||
}
|
||||
|
||||
type Pricing struct {
|
||||
BillingPluginVariants []PricingPluginVariant `json:"billing_plugin_variants,omitempty"`
|
||||
ModelName string `json:"model_name"`
|
||||
Description string `json:"description,omitempty"`
|
||||
Icon string `json:"icon,omitempty"`
|
||||
@@ -376,36 +385,46 @@ func updatePricing() {
|
||||
}
|
||||
}
|
||||
}
|
||||
usageModel := model
|
||||
plugin, ok := pluginGeneration.GetByModel(model)
|
||||
if !ok {
|
||||
if target, resolved := ResolveTaskModelAlias(pluginGeneration, model); resolved {
|
||||
plugin, ok = pluginGeneration.Get(target.PluginKey)
|
||||
usageModel = target.Declared
|
||||
}
|
||||
}
|
||||
if ok && plugin != nil && len(plugin.Meta.UsageSchema) > 0 {
|
||||
pricing.BillingUsageSchema = make(map[string]jsplugin.UsageFieldSchema, len(plugin.Meta.UsageSchema))
|
||||
for key, field := range plugin.Meta.UsageSchema {
|
||||
field.Enum = append([]string(nil), field.Enum...)
|
||||
field.Description = maps.Clone(field.Description)
|
||||
if field.EnumLabels != nil {
|
||||
labels := make(map[string]jsplugin.LocalizedText, len(field.EnumLabels))
|
||||
for value, label := range field.EnumLabels {
|
||||
labels[value] = maps.Clone(label)
|
||||
}
|
||||
field.EnumLabels = labels
|
||||
}
|
||||
pricing.BillingUsageSchema[key] = field
|
||||
if ok && plugin != nil {
|
||||
usageSchema, usageExamples := plugin.Meta.UsageForModel(usageModel)
|
||||
pricing.BillingUsageSchema = jsplugin.CloneUsageSchema(usageSchema)
|
||||
pricing.BillingUsageExamples = jsplugin.CloneUsageExamples(usageExamples)
|
||||
}
|
||||
providers := pluginGeneration.PluginsByModel(model)
|
||||
hasProviderOverride := false
|
||||
for _, provider := range providers {
|
||||
if _, configured := billing_setting.GetPluginBillingExpr(provider.Meta.Key, model); configured {
|
||||
hasProviderOverride = true
|
||||
break
|
||||
}
|
||||
if len(plugin.Meta.UsageExamples) > 0 {
|
||||
pricing.BillingUsageExamples = make([]jsplugin.UsageExample, len(plugin.Meta.UsageExamples))
|
||||
for index, example := range plugin.Meta.UsageExamples {
|
||||
facts := make(map[string]any, len(example.Facts))
|
||||
maps.Copy(facts, example.Facts)
|
||||
pricing.BillingUsageExamples[index] = jsplugin.UsageExample{
|
||||
Label: example.Label,
|
||||
Facts: facts,
|
||||
}
|
||||
}
|
||||
if hasProviderOverride || (len(providers) >= 2 && pricing.BillingMode == billing_setting.BillingModeTieredExpr) {
|
||||
for _, provider := range providers {
|
||||
schema, examples := provider.Meta.UsageForModel(model)
|
||||
if schema == nil {
|
||||
schema = map[string]jsplugin.UsageFieldSchema{}
|
||||
}
|
||||
expression, hasExpression := billing_setting.ResolveTaskBillingExpr(provider.Meta.Key, model, "")
|
||||
mode := billing_setting.BillingModeRatio
|
||||
if hasExpression || billing_setting.GetBillingMode(model) == billing_setting.BillingModeTieredExpr {
|
||||
mode = billing_setting.BillingModeTieredExpr
|
||||
}
|
||||
if mode == billing_setting.BillingModeTieredExpr && !billing_setting.TaskExprCompatible(expression, schema) {
|
||||
expression = ""
|
||||
}
|
||||
pricing.BillingPluginVariants = append(pricing.BillingPluginVariants, PricingPluginVariant{
|
||||
PluginKey: provider.Meta.Key, PluginName: provider.Meta.Name, Icon: provider.Meta.Icon,
|
||||
BillingExpr: expression, BillingMode: mode,
|
||||
BillingUsageSchema: jsplugin.CloneUsageSchema(schema), BillingUsageExamples: jsplugin.CloneUsageExamples(examples),
|
||||
})
|
||||
}
|
||||
}
|
||||
pricingMap = append(pricingMap, pricing)
|
||||
|
||||
@@ -2,13 +2,17 @@ package model
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/pkg/jsplugin"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
"github.com/QuantumNous/new-api/setting/billing_setting"
|
||||
"github.com/QuantumNous/new-api/setting/config"
|
||||
"github.com/QuantumNous/new-api/setting/ratio_setting"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
@@ -139,3 +143,157 @@ func pricingByModel(pricings []Pricing) map[string]Pricing {
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
func TestPricingUsageProfilesAndAliases(t *testing.T) {
|
||||
resetPricingEndpointTestTables(t)
|
||||
require.NoError(t, DB.AutoMigrate(&Option{}))
|
||||
const key = "pricing-profiles"
|
||||
const source = `
|
||||
export const meta = {
|
||||
apiVersion:1,key:"pricing-profiles",name:"Pricing Profiles",version:"1.0.0",author:{name:"Test"},
|
||||
models:["profile-image","profile-video"],fetchMode:"per_task",
|
||||
usageSchema:{fallback:{type:"number",unit:"count"}},
|
||||
usageExamples:[{label:"default",facts:{fallback:1}}],
|
||||
usageProfiles:[
|
||||
{models:["profile-image"],schema:{image_count:{type:"number",unit:"count"}}},
|
||||
{models:["profile-video"],schema:{seconds:{type:"number",unit:"second"}},examples:[{label:"5 seconds",facts:{seconds:5}}]}
|
||||
]
|
||||
};
|
||||
export function buildSubmitRequest(){return {};}
|
||||
export function parseSubmitResponse(){return {};}
|
||||
export function buildQueryRequest(){return {};}
|
||||
export function parseTaskResult(){return {};}
|
||||
`
|
||||
_, err := jsplugin.DefaultRegistry.Register(source, jsplugin.Options{})
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(func() { jsplugin.DefaultRegistry.Unregister(key) })
|
||||
for index, mapping := range []string{
|
||||
`{"image-alias":"profile-image","video-alias":"profile-video","ambiguous-alias":"profile-image"}`,
|
||||
`{"ambiguous-alias":"profile-video"}`,
|
||||
} {
|
||||
id := 920 + index
|
||||
require.NoError(t, DB.Create(&Channel{
|
||||
Id: id, Type: constant.ChannelTypeTaskPlugin, Status: 1, Name: fmt.Sprintf("channel-%d", id),
|
||||
Models: "profile-image,profile-video,image-alias,video-alias,ambiguous-alias", ModelMapping: &mapping,
|
||||
}).Error)
|
||||
for _, name := range []string{"profile-image", "profile-video", "image-alias", "video-alias", "ambiguous-alias"} {
|
||||
insertPricingEndpointAbility(t, id, name)
|
||||
}
|
||||
}
|
||||
InitChannelCache()
|
||||
pricing := pricingByModel(GetPricing())
|
||||
names := []string{"profile-image", "profile-video", "image-alias", "video-alias", "ambiguous-alias"}
|
||||
snapshot, err := GetModelPricingSnapshot(names)
|
||||
require.NoError(t, err)
|
||||
entries := make(map[string]ModelPricingEntry, len(snapshot.Entries))
|
||||
for _, entry := range snapshot.Entries {
|
||||
entries[entry.ModelName] = entry
|
||||
}
|
||||
for _, tc := range []struct {
|
||||
name, field, unit string
|
||||
wantExamples int
|
||||
}{
|
||||
{"profile-image", "image_count", "count", 0},
|
||||
{"profile-video", "seconds", "second", 1},
|
||||
{"image-alias", "image_count", "count", 0},
|
||||
{"video-alias", "seconds", "second", 1},
|
||||
{"ambiguous-alias", "fallback", "count", 1},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
expected := map[string]jsplugin.UsageFieldSchema{tc.field: {Type: "number", Unit: tc.unit}}
|
||||
assert.Equal(t, expected, pricing[tc.name].BillingUsageSchema)
|
||||
assert.Equal(t, expected, entries[tc.name].UsageSchema)
|
||||
assert.Len(t, pricing[tc.name].BillingUsageExamples, tc.wantExamples)
|
||||
require.NoError(t, ValidateModelPricing(tc.name, PricingValues{"billing_setting.billing_expr": `u("` + tc.field + `")`}))
|
||||
require.ErrorContains(t, ValidateModelPricing(tc.name, PricingValues{"billing_setting.billing_expr": `u("missing")`}), "not declared")
|
||||
})
|
||||
}
|
||||
require.ErrorContains(t, ValidateModelPricing("profile-image", PricingValues{"billing_setting.billing_expr": `u("seconds")`}), "not declared")
|
||||
require.ErrorContains(t, ValidateModelPricing("video-alias", PricingValues{"billing_setting.billing_expr": `u("image_count")`}), "not declared")
|
||||
|
||||
updated := strings.Replace(source, `version:"1.0.0"`, `version:"1.1.0"`, 1)
|
||||
updated = strings.Replace(updated, `image_count:{type:"number",unit:"count"}`, `images:{type:"number",unit:"count"}`, 1)
|
||||
_, err = jsplugin.DefaultRegistry.Register(updated, jsplugin.Options{})
|
||||
require.NoError(t, err)
|
||||
InvalidatePricingCache()
|
||||
refreshed := pricingByModel(GetPricing())
|
||||
assert.Equal(t, map[string]jsplugin.UsageFieldSchema{"images": {Type: "number", Unit: "count"}}, refreshed["image-alias"].BillingUsageSchema)
|
||||
assert.Equal(t, "second", refreshed["profile-video"].BillingUsageSchema["seconds"].Unit)
|
||||
}
|
||||
|
||||
func TestPricingSharedPluginVariantsUseEachSchemaAndExpression(t *testing.T) {
|
||||
resetPricingEndpointTestTables(t)
|
||||
saved := config.GlobalConfig.ExportAllConfigs()
|
||||
t.Cleanup(func() { require.NoError(t, config.GlobalConfig.LoadFromDB(saved)) })
|
||||
for _, spec := range []struct{ key, field, unit string }{{"pricing-alpha", "seconds", "second"}, {"pricing-beta", "credits", "credit"}} {
|
||||
source := strings.ReplaceAll(pricingUsagePluginSource("1.0.0", `{`+spec.field+`:{type:"number",unit:"`+spec.unit+`"}}`), "pricing-usage-probe", spec.key)
|
||||
_, err := jsplugin.DefaultRegistry.Register(source, jsplugin.Options{})
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(func() { jsplugin.DefaultRegistry.Unregister(spec.key) })
|
||||
}
|
||||
insertPricingEndpointChannel(t, 930, constant.ChannelTypeTaskPlugin, dto.ChannelOtherSettings{})
|
||||
insertPricingEndpointAbility(t, 930, "pricing-usage-model")
|
||||
const expression = `tier("base", u("seconds") * 0.4)`
|
||||
require.NoError(t, config.GlobalConfig.LoadFromDB(map[string]string{
|
||||
"billing_setting.billing_mode": `{"pricing-usage-model":"tiered_expr"}`,
|
||||
"billing_setting.billing_expr": `{"pricing-usage-model":"tier(\"base\", u(\"seconds\") * 0.4)"}`,
|
||||
billing_setting.PluginBillingExprOption: `{}`,
|
||||
}))
|
||||
prices := pricingByModel(GetPricing())
|
||||
variants := prices["pricing-usage-model"].BillingPluginVariants
|
||||
require.Len(t, variants, 2)
|
||||
assert.Equal(t, "pricing-alpha", variants[0].PluginKey)
|
||||
assert.Equal(t, "pricing-beta", variants[1].PluginKey)
|
||||
assert.Equal(t, expression, variants[0].BillingExpr)
|
||||
assert.Empty(t, variants[1].BillingExpr)
|
||||
assert.Equal(t, "credit", variants[1].BillingUsageSchema["credits"].Unit)
|
||||
assert.Equal(t, expression, prices["pricing-usage-model"].BillingExpr)
|
||||
assert.Contains(t, prices["pricing-usage-model"].BillingUsageSchema, "seconds")
|
||||
draft := PricingValues{"billing_setting.billing_expr": expression}
|
||||
require.ErrorContains(t, ValidateModelPricing("pricing-usage-model", draft), "plugin pricing-beta")
|
||||
draft[billing_setting.PluginBillingExprOption] = map[string]any{"pricing-beta": `tier("beta", u("credits") * 2)`}
|
||||
require.NoError(t, ValidateModelPricing("pricing-usage-model", draft))
|
||||
for _, invalid := range []any{map[string]any{"missing": "1"}, map[string]any{"pricing-beta": float64(1)}, map[string]any{"pricing-beta": `u("seconds")`}, map[string]any(nil), []any{}, "invalid"} {
|
||||
draft[billing_setting.PluginBillingExprOption] = invalid
|
||||
assert.Error(t, ValidateModelPricing("pricing-usage-model", draft))
|
||||
}
|
||||
draft = PricingValues{
|
||||
"billing_setting.billing_expr": "invalid expression(",
|
||||
billing_setting.PluginBillingExprOption: map[string]any{"pricing-alpha": expression, "pricing-beta": `u("credits")`},
|
||||
}
|
||||
require.ErrorContains(t, ValidateModelPricing("pricing-usage-model", draft), "compile")
|
||||
require.NoError(t, config.GlobalConfig.LoadFromDB(map[string]string{
|
||||
billing_setting.PluginBillingExprOption: `{"pricing-beta::pricing-usage-model":"tier(\"beta\", u(\"credits\") * 2)"}`,
|
||||
}))
|
||||
InvalidatePricingCache()
|
||||
priced := pricingByModel(GetPricing())["pricing-usage-model"].BillingPluginVariants
|
||||
require.Len(t, priced, 2)
|
||||
assert.Equal(t, `tier("beta", u("credits") * 2)`, priced[1].BillingExpr)
|
||||
previousPrice := ratio_setting.ModelPrice2JSONString()
|
||||
t.Cleanup(func() { require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(previousPrice)) })
|
||||
require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(`{"pricing-usage-model":0.25}`))
|
||||
require.NoError(t, config.GlobalConfig.LoadFromDB(map[string]string{"billing_setting.billing_mode": `{"pricing-usage-model":"ratio"}`}))
|
||||
InvalidatePricingCache()
|
||||
mixed := pricingByModel(GetPricing())["pricing-usage-model"].BillingPluginVariants
|
||||
require.Len(t, mixed, 2)
|
||||
assert.Equal(t, billing_setting.BillingModeRatio, mixed[0].BillingMode)
|
||||
assert.Equal(t, billing_setting.BillingModeTieredExpr, mixed[1].BillingMode)
|
||||
require.NoError(t, config.GlobalConfig.LoadFromDB(map[string]string{billing_setting.PluginBillingExprOption: `{}`}))
|
||||
InvalidatePricingCache()
|
||||
perCall := pricingByModel(GetPricing())["pricing-usage-model"]
|
||||
assert.Empty(t, perCall.BillingPluginVariants)
|
||||
assert.Equal(t, 1, perCall.QuotaType)
|
||||
assert.Equal(t, 0.25, perCall.ModelPrice)
|
||||
require.NoError(t, config.GlobalConfig.LoadFromDB(map[string]string{
|
||||
billing_setting.PluginBillingExprOption: `{"pricing-beta::pricing-usage-model":"tier(\"beta\", u(\"credits\") * 2)"}`,
|
||||
}))
|
||||
require.NoError(t, jsplugin.DefaultRegistry.Unregister("pricing-alpha"))
|
||||
InvalidatePricingCache()
|
||||
single := pricingByModel(GetPricing())["pricing-usage-model"]
|
||||
require.Len(t, single.BillingPluginVariants, 1)
|
||||
assert.Equal(t, "pricing-beta", single.BillingPluginVariants[0].PluginKey)
|
||||
assert.Equal(t, `tier("beta", u("credits") * 2)`, single.BillingPluginVariants[0].BillingExpr)
|
||||
encoded, err := common.Marshal(single)
|
||||
require.NoError(t, err)
|
||||
assert.Contains(t, string(encoded), "billing_plugin_variants")
|
||||
}
|
||||
@@ -0,0 +1,287 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
)
|
||||
|
||||
// RelayAttempt is the durable financial fact and log outbox for one upstream
|
||||
// generation. It contains numeric usage and pricing, never a prompt or key.
|
||||
type RelayAttempt struct {
|
||||
ID int64 `json:"id"`
|
||||
AttemptKey string `json:"attempt_key" gorm:"type:varchar(96);uniqueIndex"`
|
||||
RequestID string `json:"request_id" gorm:"type:varchar(64);index"`
|
||||
Number int `json:"number"`
|
||||
UserID int `json:"-" gorm:"index"`
|
||||
TokenID int `json:"-"`
|
||||
ChannelID int `json:"channel_id"`
|
||||
GroupName string `json:"group"`
|
||||
RequestedModel string `json:"requested_model"`
|
||||
ReturnedModel string `json:"returned_model"`
|
||||
GuardResult string `json:"guard_result"`
|
||||
Funding string `json:"billing_source"`
|
||||
SubscriptionID int `json:"-"`
|
||||
SubscriptionResetTime int64 `json:"-"`
|
||||
Reserved int `json:"-"`
|
||||
Quota int `json:"quota"`
|
||||
ObservedQuota int `json:"-"`
|
||||
EstimateQuota int `json:"-"`
|
||||
Billable bool `json:"-"`
|
||||
UsageSource string `json:"usage_source"`
|
||||
State string `json:"state" gorm:"type:varchar(24);index"`
|
||||
UsageJSON string `json:"-" gorm:"type:text"`
|
||||
PriceJSON string `json:"-" gorm:"type:text"`
|
||||
LogJSON string `json:"-" gorm:"type:text"`
|
||||
LogDelivered bool `json:"-" gorm:"index"`
|
||||
CreatedAt int64 `json:"created_at" gorm:"autoCreateTime:false"`
|
||||
UpdatedAt int64 `json:"-" gorm:"autoUpdateTime:false"`
|
||||
Deadline int64 `json:"-" gorm:"index"`
|
||||
}
|
||||
|
||||
// SQLite cannot upgrade a read transaction while another attempt is writing.
|
||||
// Retry the entire rolled-back transaction, never an individual debit statement.
|
||||
func relayAttemptTransaction(fn func(*gorm.DB) error) error {
|
||||
for retry := 0; ; retry++ {
|
||||
err := DB.Transaction(fn)
|
||||
var coded interface{ Code() int }
|
||||
busy := errors.As(err, &coded) && (coded.Code()&255 == 5 || coded.Code()&255 == 6)
|
||||
if !busy || retry >= 24 {
|
||||
return err
|
||||
}
|
||||
delay := time.Duration(retry+1) * 5 * time.Millisecond
|
||||
if delay > 100*time.Millisecond {
|
||||
delay = 100 * time.Millisecond
|
||||
}
|
||||
time.Sleep(delay)
|
||||
}
|
||||
}
|
||||
|
||||
var ErrAttemptQuota = errors.New("insufficient quota for another upstream attempt")
|
||||
|
||||
func ReserveRelayAttempt(attempt *RelayAttempt, preference string) error {
|
||||
if attempt.Reserved < 0 || attempt.Reserved > common.MaxWalletQuota {
|
||||
return ErrAttemptQuota
|
||||
}
|
||||
attempt.CreatedAt = time.Now().Unix()
|
||||
attempt.UpdatedAt = attempt.CreatedAt
|
||||
attempt.State = "reserved"
|
||||
var tokenKey string
|
||||
original := *attempt
|
||||
err := relayAttemptTransaction(func(tx *gorm.DB) error {
|
||||
*attempt = original
|
||||
var existing RelayAttempt
|
||||
if err := tx.Where("attempt_key = ?", attempt.AttemptKey).First(&existing).Error; err == nil {
|
||||
*attempt = existing
|
||||
return nil
|
||||
} else if !errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return err
|
||||
}
|
||||
var user User
|
||||
if err := lockForUpdate(tx).First(&user, attempt.UserID).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
var token Token
|
||||
if err := lockForUpdate(tx).First(&token, attempt.TokenID).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
tokenKey = token.Key
|
||||
if token.UserId != attempt.UserID || token.Status != common.TokenStatusEnabled || (!token.UnlimitedQuota && token.RemainQuota < attempt.Reserved) {
|
||||
return ErrAttemptQuota
|
||||
}
|
||||
if err := invalidateTokenCacheForMutation(token.Key); err != nil {
|
||||
return err
|
||||
}
|
||||
amount := int64(attempt.Reserved)
|
||||
wallet := func() bool { return user.Quota >= attempt.Reserved }
|
||||
trySub := func() (bool, bool, error) {
|
||||
var subs []UserSubscription
|
||||
if err := lockForUpdate(tx).Where("user_id = ? AND status = ? AND end_time > ?", attempt.UserID, "active", attempt.CreatedAt).Order("end_time asc, id asc").Find(&subs).Error; err != nil {
|
||||
return false, false, err
|
||||
}
|
||||
allow := len(subs) == 0
|
||||
for _, sub := range subs {
|
||||
allow = allow || sub.AllowWalletOverflow
|
||||
plan, err := getSubscriptionPlanByIdTx(tx, sub.PlanId)
|
||||
if err != nil {
|
||||
return false, allow, err
|
||||
}
|
||||
if err := maybeResetUserSubscriptionWithPlanTx(tx, &sub, plan, attempt.CreatedAt); err != nil {
|
||||
return false, allow, err
|
||||
}
|
||||
if sub.AmountTotal > 0 && sub.AmountTotal-sub.AmountUsed < amount {
|
||||
continue
|
||||
}
|
||||
if err := tx.Model(&UserSubscription{}).Where("id = ?", sub.Id).Update("amount_used", gorm.Expr("amount_used + ?", amount)).Error; err != nil {
|
||||
return false, allow, err
|
||||
}
|
||||
attempt.Funding = "subscription"
|
||||
attempt.SubscriptionID = sub.Id
|
||||
attempt.SubscriptionResetTime = sub.LastResetTime
|
||||
return true, allow, nil
|
||||
}
|
||||
return false, allow, nil
|
||||
}
|
||||
useWallet := false
|
||||
switch common.NormalizeBillingPreference(preference) {
|
||||
case "wallet_only":
|
||||
useWallet = wallet()
|
||||
case "wallet_first":
|
||||
useWallet = wallet()
|
||||
if !useWallet {
|
||||
ok, _, err := trySub()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !ok {
|
||||
return ErrAttemptQuota
|
||||
}
|
||||
}
|
||||
case "subscription_only":
|
||||
ok, _, err := trySub()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !ok {
|
||||
return ErrAttemptQuota
|
||||
}
|
||||
default:
|
||||
ok, allow, err := trySub()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !ok {
|
||||
useWallet = allow && wallet()
|
||||
}
|
||||
}
|
||||
if useWallet {
|
||||
attempt.Funding = "wallet"
|
||||
if err := tx.Model(&User{}).Where("id = ?", user.Id).Update("quota", gorm.Expr("quota - ?", attempt.Reserved)).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
} else if attempt.Funding != "subscription" {
|
||||
return ErrAttemptQuota
|
||||
}
|
||||
if err := tx.Model(&Token{}).Where("id = ?", token.Id).Updates(map[string]any{"remain_quota": gorm.Expr("remain_quota - ?", attempt.Reserved), "used_quota": gorm.Expr("used_quota + ?", attempt.Reserved), "accessed_time": attempt.CreatedAt}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Create(attempt).Error
|
||||
})
|
||||
if err == nil {
|
||||
_ = invalidateUserCache(attempt.UserID)
|
||||
_ = invalidateTokenCacheForMutation(tokenKey)
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
func CheckpointRelayAttempt(id int64, quota int, usage, logJSON, returned, guard, source string, billable bool, deadline int64) error {
|
||||
return DB.Model(&RelayAttempt{}).Where("id = ? AND state <> ?", id, "settled").Updates(map[string]any{"observed_quota": quota, "usage_json": usage, "log_json": logJSON, "returned_model": returned, "guard_result": guard, "usage_source": source, "billable": billable, "state": "collecting", "updated_at": time.Now().Unix(), "deadline": deadline}).Error
|
||||
}
|
||||
|
||||
// Settlement and all balances/statistics commit together. Retrying after any
|
||||
// process crash cannot repeat a debit, including the delta-zero case.
|
||||
func SettleRelayAttempt(id int64, quota int, log *Log, details ...RelayAttempt) error {
|
||||
if quota < 0 || quota > common.MaxWalletQuota {
|
||||
return fmt.Errorf("invalid attempt charge")
|
||||
}
|
||||
var attempt RelayAttempt
|
||||
var tokenKey string
|
||||
err := relayAttemptTransaction(func(tx *gorm.DB) error {
|
||||
attempt = RelayAttempt{}
|
||||
if err := lockForUpdate(tx).First(&attempt, id).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if attempt.State == "settled" {
|
||||
return nil
|
||||
}
|
||||
delta := quota - attempt.Reserved
|
||||
var user User
|
||||
if err := lockForUpdate(tx).Unscoped().First(&user, attempt.UserID).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
var token Token
|
||||
if err := lockForUpdate(tx).Unscoped().First(&token, attempt.TokenID).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
tokenKey = token.Key
|
||||
if err := invalidateTokenCacheForMutation(token.Key); err != nil {
|
||||
return err
|
||||
}
|
||||
if attempt.Funding == "subscription" {
|
||||
if err := tx.Model(&UserSubscription{}).Where("id = ? AND last_reset_time = ?", attempt.SubscriptionID, attempt.SubscriptionResetTime).Update("amount_used", gorm.Expr("CASE WHEN amount_used + ? < 0 THEN 0 ELSE amount_used + ? END", delta, delta)).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
} else if err := tx.Unscoped().Model(&User{}).Where("id = ?", attempt.UserID).Update("quota", gorm.Expr("quota - ?", delta)).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tx.Unscoped().Model(&Token{}).Where("id = ?", attempt.TokenID).Updates(map[string]any{"remain_quota": gorm.Expr("remain_quota - ?", delta), "used_quota": gorm.Expr("used_quota + ?", delta)}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if quota > 0 || attempt.Billable || (len(details) > 0 && details[0].Billable) {
|
||||
if err := tx.Unscoped().Model(&User{}).Where("id = ?", attempt.UserID).Updates(map[string]any{"used_quota": gorm.Expr("used_quota + ?", quota), "request_count": gorm.Expr("request_count + 1")}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tx.Model(&Channel{}).Where("id = ?", attempt.ChannelID).Update("used_quota", gorm.Expr("used_quota + ?", quota)).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if len(details) > 0 {
|
||||
d := details[0]
|
||||
attempt.UsageJSON = d.UsageJSON
|
||||
attempt.ReturnedModel = d.ReturnedModel
|
||||
attempt.GuardResult = d.GuardResult
|
||||
attempt.UsageSource = d.UsageSource
|
||||
attempt.Billable = d.Billable
|
||||
attempt.ObservedQuota = quota
|
||||
}
|
||||
attempt.Quota = quota
|
||||
attempt.State = "settled"
|
||||
attempt.UpdatedAt = time.Now().Unix()
|
||||
if log != nil {
|
||||
log.Quota = quota
|
||||
log.AttemptKey = &attempt.AttemptKey
|
||||
if common.DataExportEnabled {
|
||||
if err := tx.Create(&QuotaData{AttemptKey: &attempt.AttemptKey, UserID: attempt.UserID, Username: log.Username, ModelName: log.ModelName, CreatedAt: log.CreatedAt - (log.CreatedAt % 3600), UseGroup: attempt.GroupName, TokenID: attempt.TokenID, ChannelID: attempt.ChannelID, NodeName: common.NodeName, Count: 1, Quota: quota, TokenUsed: log.PromptTokens + log.CompletionTokens}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
body, err := common.Marshal(log)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
attempt.LogJSON = string(body)
|
||||
}
|
||||
return tx.Save(&attempt).Error
|
||||
})
|
||||
if err == nil {
|
||||
_ = invalidateUserCache(attempt.UserID)
|
||||
_ = invalidateTokenCacheForMutation(tokenKey)
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
func DeliverRelayAttemptLogs() error {
|
||||
var attempts []RelayAttempt
|
||||
if err := DB.Where("state = ? AND log_delivered = ? AND log_json <> ?", "settled", false, "").Limit(100).Find(&attempts).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
for _, a := range attempts {
|
||||
var log Log
|
||||
if err := common.UnmarshalJsonStr(a.LogJSON, &log); err != nil {
|
||||
return err
|
||||
}
|
||||
log.Id = 0
|
||||
log.AttemptKey = &a.AttemptKey
|
||||
if err := LOG_DB.Clauses(clause.OnConflict{Columns: []clause.Column{{Name: "attempt_key"}}, DoNothing: true}).Create(&log).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if err := DB.Model(&RelayAttempt{}).Where("id = ?", a.ID).Update("log_delivered", true).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,225 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/glebarez/sqlite"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/driver/mysql"
|
||||
"gorm.io/driver/postgres"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/schema"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestTokenRoutingResolution(t *testing.T) {
|
||||
p := TokenRoutingPolicy{Version: 1, Mode: "custom", Groups: []TokenGroupEntry{{"a", 0}, {"b", 0}, {"c", 2}}, ModelRules: []ModelRoutingRule{{Model: "model-x", Groups: []TokenGroupEntry{{"c", 0}, {"b", 0}}}}}
|
||||
allowed := func(string) bool { return true }
|
||||
require.NoError(t, p.Validate(6, allowed))
|
||||
require.Equal(t, []TokenGroupEntry{{"c", 0}, {"b", 0}, {"a", 1}}, p.Resolve("model-x", nil, allowed))
|
||||
require.Equal(t, []TokenGroupEntry{{"a", 0}, {"b", 0}, {"c", 1}}, p.Resolve("MODEL-X", nil, allowed))
|
||||
}
|
||||
|
||||
func TestTokenRoutingLegacy(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
token Token
|
||||
mode string
|
||||
groups []TokenGroupEntry
|
||||
}{
|
||||
{Token{Group: "auto"}, "inherit", []TokenGroupEntry{}},
|
||||
{Token{Group: ""}, "user", []TokenGroupEntry{}},
|
||||
{Token{Group: "a", AutoGroups: `["b","c"]`, CrossGroupRetry: true}, "custom", []TokenGroupEntry{{"a", 0}, {"b", 1}, {"c", 2}}},
|
||||
{Token{Group: "a", AutoGroups: `["b"]`, CrossGroupRetry: false}, "custom", []TokenGroupEntry{{"a", 0}}},
|
||||
{Token{Group: "a", AutoGroups: `[{"group":"b","priority":0},{"group":"a","priority":0}]`, CrossGroupRetry: true}, "custom", []TokenGroupEntry{{"b", 0}, {"a", 0}}},
|
||||
} {
|
||||
p, err := tc.token.LegacyRoutingPolicy()
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, tc.mode, p.Mode)
|
||||
require.Equal(t, tc.groups, p.Groups)
|
||||
require.Empty(t, p.ModelRules)
|
||||
}
|
||||
p := TokenRoutingPolicy{Version: 1, Mode: "custom", Groups: []TokenGroupEntry{{"a", 0}}, ModelRules: []ModelRoutingRule{{Model: "x", Groups: []TokenGroupEntry{{"b", 0}}}}}
|
||||
require.Error(t, p.Validate(6, func(string) bool { return true }))
|
||||
}
|
||||
|
||||
func TestRelayAttemptDatabaseMatrix(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":
|
||||
dsn := os.Getenv("TEST_MYSQL_DSN")
|
||||
if dsn == "" {
|
||||
t.Skip("TEST_MYSQL_DSN not configured")
|
||||
}
|
||||
driver = mysql.Open(dsn)
|
||||
case "postgres":
|
||||
dsn := os.Getenv("TEST_POSTGRES_DSN")
|
||||
if dsn == "" {
|
||||
t.Skip("TEST_POSTGRES_DSN not configured")
|
||||
}
|
||||
driver = postgres.Open(dsn)
|
||||
}
|
||||
db, err := gorm.Open(driver, &gorm.Config{NamingStrategy: schema.NamingStrategy{TablePrefix: fmt.Sprintf("attempt_%d_", time.Now().UnixNano())}})
|
||||
require.NoError(t, err)
|
||||
sqlDB, err := db.DB()
|
||||
require.NoError(t, err)
|
||||
sqlDB.SetMaxOpenConns(1)
|
||||
oldDB, oldLog, oldType, oldRedis := DB, LOG_DB, common.MainDatabaseType(), common.RedisEnabled
|
||||
DB = db
|
||||
LOG_DB = db
|
||||
common.SetMainDatabaseType(common.DatabaseType(dialect))
|
||||
common.RedisEnabled = false
|
||||
initCol()
|
||||
models := []any{&User{}, &Token{}, &Channel{}, &RelayAttempt{}, &Log{}, &QuotaData{}, &SubscriptionPlan{}, &UserSubscription{}}
|
||||
t.Cleanup(func() {
|
||||
for _, m := range models {
|
||||
require.NoError(t, db.Migrator().DropTable(m))
|
||||
}
|
||||
DB = oldDB
|
||||
LOG_DB = oldLog
|
||||
common.SetMainDatabaseType(oldType)
|
||||
common.RedisEnabled = oldRedis
|
||||
initCol()
|
||||
require.NoError(t, sqlDB.Close())
|
||||
})
|
||||
for range 2 {
|
||||
require.NoError(t, db.AutoMigrate(models...))
|
||||
}
|
||||
user := User{Username: "attempt-owner", Quota: 10000}
|
||||
require.NoError(t, db.Create(&user).Error)
|
||||
token := Token{UserId: user.Id, Key: "attempt-key", Status: common.TokenStatusEnabled, RemainQuota: 10000, Group: "a", CrossGroupRetry: true, AutoGroups: `["b"]`}
|
||||
require.NoError(t, db.Create(&token).Error)
|
||||
channel := Channel{Name: "attempt-channel", Status: common.ChannelStatusEnabled}
|
||||
require.NoError(t, db.Create(&channel).Error)
|
||||
// Upgrade a populated legacy token table, then run startup migrations twice.
|
||||
require.NoError(t, db.Migrator().DropColumn(&Token{}, "RoutingPolicy"))
|
||||
require.NoError(t, db.AutoMigrate(&Token{}))
|
||||
count, err := MigrateTokenRoutingPolicies()
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 1, count)
|
||||
count, err = MigrateTokenRoutingPolicies()
|
||||
require.NoError(t, err)
|
||||
require.Zero(t, count)
|
||||
var migrated Token
|
||||
require.NoError(t, db.First(&migrated, token.Id).Error)
|
||||
p, err := migrated.GetRoutingPolicy()
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []TokenGroupEntry{{"a", 0}, {"b", 1}}, p.Groups)
|
||||
var total int
|
||||
for i, quota := range []int{120, 230, 0} {
|
||||
a := RelayAttempt{AttemptKey: fmt.Sprintf("request:%d", i), RequestID: "request", Number: i + 1, UserID: user.Id, TokenID: token.Id, ChannelID: channel.Id, Reserved: 300}
|
||||
require.NoError(t, ReserveRelayAttempt(&a, "wallet_only"))
|
||||
require.NoError(t, ReserveRelayAttempt(&a, "wallet_only"))
|
||||
require.NoError(t, CheckpointRelayAttempt(a.ID, quota, `{"total_tokens":10}`, "", "luna", "model_mismatch", "upstream", quota > 0, time.Now().Unix()))
|
||||
log := &Log{UserId: user.Id, Type: LogTypeConsume, RequestId: "request"}
|
||||
require.NoError(t, SettleRelayAttempt(a.ID, quota, log))
|
||||
require.NoError(t, SettleRelayAttempt(a.ID, quota+1, log))
|
||||
total += quota
|
||||
}
|
||||
require.NoError(t, DeliverRelayAttemptLogs())
|
||||
require.NoError(t, DeliverRelayAttemptLogs())
|
||||
require.NoError(t, db.First(&user, user.Id).Error)
|
||||
require.Equal(t, 10000-total, user.Quota)
|
||||
require.Equal(t, total, user.UsedQuota)
|
||||
require.Equal(t, 2, user.RequestCount)
|
||||
require.NoError(t, db.First(&token, token.Id).Error)
|
||||
require.Equal(t, 10000-total, token.RemainQuota)
|
||||
require.Equal(t, total, token.UsedQuota)
|
||||
require.NoError(t, db.First(&channel, channel.Id).Error)
|
||||
require.EqualValues(t, total, channel.UsedQuota)
|
||||
var logs int64
|
||||
require.NoError(t, db.Model(&Log{}).Count(&logs).Error)
|
||||
require.EqualValues(t, 3, logs)
|
||||
// A log-store outage never rolls back an already-committed financial fact.
|
||||
require.NoError(t, db.Model(&RelayAttempt{}).Where("request_id = ?", "request").Update("log_delivered", false).Error)
|
||||
require.NoError(t, DeliverRelayAttemptLogs())
|
||||
require.NoError(t, db.Model(&Log{}).Count(&logs).Error)
|
||||
require.EqualValues(t, 3, logs)
|
||||
plan := SubscriptionPlan{Title: "attempt-sub", QuotaResetPeriod: "never"}
|
||||
require.NoError(t, db.Create(&plan).Error)
|
||||
sub := UserSubscription{UserId: user.Id, PlanId: plan.Id, AmountTotal: 1000, Status: "active", EndTime: time.Now().Add(time.Hour).Unix()}
|
||||
require.NoError(t, db.Create(&sub).Error)
|
||||
subAttempt := RelayAttempt{AttemptKey: "sub-request", UserID: user.Id, TokenID: token.Id, ChannelID: channel.Id, Reserved: 300}
|
||||
require.NoError(t, ReserveRelayAttempt(&subAttempt, "subscription_only"))
|
||||
require.Equal(t, "subscription", subAttempt.Funding)
|
||||
require.NoError(t, SettleRelayAttempt(subAttempt.ID, 120, nil))
|
||||
require.NoError(t, SettleRelayAttempt(subAttempt.ID, 999, nil))
|
||||
require.NoError(t, db.First(&sub, sub.Id).Error)
|
||||
require.EqualValues(t, 120, sub.AmountUsed)
|
||||
require.NoError(t, db.First(&user, user.Id).Error)
|
||||
require.Equal(t, 10000-total, user.Quota)
|
||||
rejected := RelayAttempt{AttemptKey: "too-expensive", UserID: user.Id, TokenID: token.Id, Reserved: 20000}
|
||||
require.ErrorIs(t, ReserveRelayAttempt(&rejected, "wallet_only"), ErrAttemptQuota)
|
||||
require.NoError(t, db.First(&user, user.Id).Error)
|
||||
require.Equal(t, 10000-total, user.Quota)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestRelayAttemptSQLiteConcurrentSettlementAndReserve(t *testing.T) {
|
||||
db, err := gorm.Open(sqlite.Open(filepath.Join(t.TempDir(), "attempt.db")), &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
sqlDB, err := db.DB()
|
||||
require.NoError(t, err)
|
||||
sqlDB.SetMaxOpenConns(8)
|
||||
oldDB, oldType, oldRedis := DB, common.MainDatabaseType(), common.RedisEnabled
|
||||
DB = db
|
||||
common.SetMainDatabaseType(common.DatabaseType("sqlite"))
|
||||
common.RedisEnabled = false
|
||||
initCol()
|
||||
t.Cleanup(func() {
|
||||
DB = oldDB
|
||||
common.SetMainDatabaseType(oldType)
|
||||
common.RedisEnabled = oldRedis
|
||||
initCol()
|
||||
_ = sqlDB.Close()
|
||||
})
|
||||
require.NoError(t, db.AutoMigrate(&User{}, &Token{}, &Channel{}, &RelayAttempt{}))
|
||||
user := User{Username: "concurrent", Quota: 10000}
|
||||
require.NoError(t, db.Create(&user).Error)
|
||||
token := Token{UserId: user.Id, Key: "concurrent-token", Status: common.TokenStatusEnabled, RemainQuota: 10000}
|
||||
require.NoError(t, db.Create(&token).Error)
|
||||
channel := Channel{Name: "concurrent"}
|
||||
require.NoError(t, db.Create(&channel).Error)
|
||||
var wg sync.WaitGroup
|
||||
failures := make(chan error, 16)
|
||||
for i := 0; i < 16; i++ {
|
||||
wg.Add(1)
|
||||
go func(i int) {
|
||||
defer wg.Done()
|
||||
a := RelayAttempt{AttemptKey: fmt.Sprintf("concurrent:%d", i), UserID: user.Id, TokenID: token.Id, ChannelID: channel.Id, Reserved: 100}
|
||||
if err := ReserveRelayAttempt(&a, "wallet_only"); err != nil {
|
||||
failures <- err
|
||||
return
|
||||
}
|
||||
if err := SettleRelayAttempt(a.ID, 50, nil); err != nil {
|
||||
failures <- err
|
||||
return
|
||||
}
|
||||
if err := SettleRelayAttempt(a.ID, 50, nil); err != nil {
|
||||
failures <- err
|
||||
}
|
||||
}(i)
|
||||
}
|
||||
wg.Wait()
|
||||
close(failures)
|
||||
for err := range failures {
|
||||
require.NoError(t, err)
|
||||
}
|
||||
require.NoError(t, db.First(&user, user.Id).Error)
|
||||
require.Equal(t, 9200, user.Quota)
|
||||
require.Equal(t, 800, user.UsedQuota)
|
||||
require.NoError(t, db.First(&token, token.Id).Error)
|
||||
require.Equal(t, 9200, token.RemainQuota)
|
||||
require.Equal(t, 800, token.UsedQuota)
|
||||
var count int64
|
||||
require.NoError(t, db.Model(&RelayAttempt{}).Where("state = ?", "settled").Count(&count).Error)
|
||||
require.EqualValues(t, 16, count)
|
||||
}
|
||||
@@ -0,0 +1,258 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"maps"
|
||||
"math"
|
||||
"regexp"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/setting"
|
||||
"github.com/QuantumNous/new-api/setting/config"
|
||||
"github.com/QuantumNous/new-api/setting/operation_setting"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// RequestPolicySnapshot is the validated read model behind the request policy
|
||||
// settings API. It is rebuilt from the option map on every option write, so
|
||||
// the runtime globals updated by the same write never disagree with it.
|
||||
type RequestPolicySnapshot struct {
|
||||
Affinity operation_setting.ChannelAffinitySetting
|
||||
RetryTimes int
|
||||
RetryCodes []operation_setting.StatusCodeRange
|
||||
DisableCodes []operation_setting.StatusCodeRange
|
||||
DisableKeywords []string
|
||||
CheckText bool
|
||||
TextKeywords []string
|
||||
AutoDisable bool
|
||||
Options map[string]string
|
||||
}
|
||||
|
||||
var requestPolicySnapshot atomic.Pointer[RequestPolicySnapshot]
|
||||
var requestPolicyOptionMutex sync.Mutex
|
||||
var requestPolicyDefaults = requestPolicyDefaultOptions()
|
||||
|
||||
func requestPolicyDefaultOptions() map[string]string {
|
||||
defaults := make(map[string]string)
|
||||
for prefix, value := range map[string]any{
|
||||
"channel_affinity_setting.": operation_setting.GetChannelAffinitySetting(),
|
||||
"monitor_setting.": operation_setting.GetMonitorSetting(),
|
||||
} {
|
||||
fields, err := config.ConfigToMap(value)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
for field, value := range fields {
|
||||
defaults[prefix+field] = value
|
||||
}
|
||||
}
|
||||
defaults["RetryTimes"] = strconv.Itoa(common.RetryTimes)
|
||||
defaults["AutomaticRetryStatusCodes"] = operation_setting.AutomaticRetryStatusCodesToString()
|
||||
defaults["AutomaticDisableStatusCodes"] = operation_setting.AutomaticDisableStatusCodesToString()
|
||||
defaults["AutomaticDisableKeywords"] = operation_setting.AutomaticDisableKeywordsToString()
|
||||
defaults["AutomaticDisableChannelEnabled"] = strconv.FormatBool(common.AutomaticDisableChannelEnabled)
|
||||
defaults["CheckSensitiveEnabled"] = strconv.FormatBool(setting.CheckSensitiveEnabled)
|
||||
defaults["CheckSensitiveOnPromptEnabled"] = strconv.FormatBool(setting.CheckSensitiveOnPromptEnabled)
|
||||
defaults["SensitiveWords"] = setting.SensitiveWordsToString()
|
||||
defaults["AutomaticEnableChannelEnabled"] = strconv.FormatBool(common.AutomaticEnableChannelEnabled)
|
||||
defaults["ChannelDisableThreshold"] = strconv.FormatFloat(common.ChannelDisableThreshold, 'f', -1, 64)
|
||||
return defaults
|
||||
}
|
||||
|
||||
func IsRequestPolicyOption(key string) bool {
|
||||
if strings.HasPrefix(key, "channel_affinity_setting.") {
|
||||
return true
|
||||
}
|
||||
switch key {
|
||||
case "CheckSensitiveEnabled", "CheckSensitiveOnPromptEnabled", "SensitiveWords", "AutomaticEnableChannelEnabled", "ChannelDisableThreshold", "monitor_setting.auto_test_channel_enabled", "monitor_setting.auto_test_channel_minutes", "monitor_setting.channel_test_concurrency", "monitor_setting.channel_test_mode", "RetryTimes", "AutomaticRetryStatusCodes", "AutomaticDisableChannelEnabled", "AutomaticDisableStatusCodes", "AutomaticDisableKeywords":
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func CurrentRequestPolicy() *RequestPolicySnapshot {
|
||||
if snapshot := requestPolicySnapshot.Load(); snapshot != nil {
|
||||
return snapshot
|
||||
}
|
||||
// Startup and test fixtures may not have loaded options yet.
|
||||
snapshot, err := BuildRequestPolicy(nil)
|
||||
if err != nil {
|
||||
return &RequestPolicySnapshot{Options: map[string]string{}}
|
||||
}
|
||||
return snapshot
|
||||
}
|
||||
|
||||
func BuildRequestPolicy(options map[string]string) (*RequestPolicySnapshot, error) {
|
||||
defaults := requestPolicyDefaults
|
||||
|
||||
raw := make(map[string]string)
|
||||
for key, value := range defaults {
|
||||
if IsRequestPolicyOption(key) {
|
||||
raw[key] = value
|
||||
}
|
||||
}
|
||||
maps.Copy(raw, options)
|
||||
snapshot := &RequestPolicySnapshot{Options: raw}
|
||||
affinityFields := map[string]string{}
|
||||
for key, value := range raw {
|
||||
if field, ok := strings.CutPrefix(key, "channel_affinity_setting."); ok {
|
||||
switch field {
|
||||
case "enabled", "session_mode", "switch_on_success", "keep_on_channel_disabled", "max_entries", "default_ttl_seconds", "rules":
|
||||
default:
|
||||
return nil, fmt.Errorf("unknown affinity option: %s", key)
|
||||
}
|
||||
affinityFields[field] = value
|
||||
}
|
||||
}
|
||||
// Decode through JSON so malformed scalar/array values cannot be silently ignored.
|
||||
affinityJSON := map[string]json.RawMessage{}
|
||||
for key, value := range affinityFields {
|
||||
if key == "session_mode" {
|
||||
encoded, _ := common.Marshal(value)
|
||||
affinityJSON[key] = encoded
|
||||
} else {
|
||||
affinityJSON[key] = json.RawMessage(value)
|
||||
}
|
||||
}
|
||||
encoded, err := common.Marshal(affinityJSON)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := common.Unmarshal(encoded, &snapshot.Affinity); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
snapshot.RetryTimes, err = strconv.Atoi(raw["RetryTimes"])
|
||||
if err != nil || snapshot.RetryTimes < 0 || snapshot.RetryTimes == math.MaxInt {
|
||||
return nil, fmt.Errorf("retry times must be a non-negative integer with room for the initial attempt")
|
||||
}
|
||||
snapshot.AutoDisable, err = strconv.ParseBool(raw["AutomaticDisableChannelEnabled"])
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
snapshot.RetryCodes, err = operation_setting.ParseHTTPStatusCodeRanges(raw["AutomaticRetryStatusCodes"])
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
snapshot.DisableCodes, err = operation_setting.ParseHTTPStatusCodeRanges(raw["AutomaticDisableStatusCodes"])
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
snapshot.DisableKeywords = strings.Split(raw["AutomaticDisableKeywords"], "\n")
|
||||
for _, key := range []string{"CheckSensitiveEnabled", "CheckSensitiveOnPromptEnabled", "AutomaticEnableChannelEnabled", "monitor_setting.auto_test_channel_enabled"} {
|
||||
if _, err := strconv.ParseBool(raw[key]); err != nil {
|
||||
return nil, fmt.Errorf("invalid boolean: %s", key)
|
||||
}
|
||||
}
|
||||
snapshot.CheckText = raw["CheckSensitiveEnabled"] == "true" && raw["CheckSensitiveOnPromptEnabled"] == "true"
|
||||
for word := range strings.SplitSeq(raw["SensitiveWords"], "\n") {
|
||||
if word = strings.TrimSpace(word); word != "" {
|
||||
snapshot.TextKeywords = append(snapshot.TextKeywords, word)
|
||||
}
|
||||
}
|
||||
if err := operation_setting.ValidateChannelTestConcurrency(raw["monitor_setting.channel_test_concurrency"]); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for _, key := range []string{"ChannelDisableThreshold", "monitor_setting.auto_test_channel_minutes"} {
|
||||
if raw[key] == "" && key == "ChannelDisableThreshold" {
|
||||
continue
|
||||
}
|
||||
value, err := strconv.ParseFloat(raw[key], 64)
|
||||
if err != nil || value < 0 || math.IsNaN(value) || math.IsInf(value, 0) {
|
||||
return nil, fmt.Errorf("invalid numeric value: %s", key)
|
||||
}
|
||||
}
|
||||
switch raw["monitor_setting.channel_test_mode"] {
|
||||
case "scheduled_all", "auto_ban_only", "passive_recovery":
|
||||
default:
|
||||
return nil, fmt.Errorf("invalid channel test mode")
|
||||
}
|
||||
|
||||
if snapshot.Affinity.MaxEntries < 0 || snapshot.Affinity.DefaultTTLSeconds < 0 {
|
||||
return nil, fmt.Errorf("invalid affinity cache limits")
|
||||
}
|
||||
switch snapshot.Affinity.SessionMode {
|
||||
case "", "off", "prefer", "strict":
|
||||
default:
|
||||
return nil, fmt.Errorf("invalid global session mode")
|
||||
}
|
||||
for _, rule := range snapshot.Affinity.Rules {
|
||||
switch rule.SessionMode {
|
||||
case "", "inherit", "off", "prefer", "strict":
|
||||
default:
|
||||
return nil, fmt.Errorf("invalid session mode for rule %q", rule.Name)
|
||||
}
|
||||
patterns := append(append([]string{}, rule.ModelRegex...), rule.PathRegex...)
|
||||
if rule.ValueRegex != "" {
|
||||
patterns = append(patterns, rule.ValueRegex)
|
||||
}
|
||||
for _, pattern := range patterns {
|
||||
if _, err := regexp.Compile(pattern); err != nil {
|
||||
return nil, fmt.Errorf("rule %q: %w", rule.Name, err)
|
||||
}
|
||||
}
|
||||
if rule.TTLSeconds < 0 {
|
||||
return nil, fmt.Errorf("negative affinity TTL")
|
||||
}
|
||||
}
|
||||
return snapshot, nil
|
||||
}
|
||||
|
||||
func refreshRequestPolicySnapshot() error {
|
||||
common.OptionMapRWMutex.RLock()
|
||||
options := make(map[string]string)
|
||||
for key, value := range common.OptionMap {
|
||||
if IsRequestPolicyOption(key) {
|
||||
options[key] = value
|
||||
}
|
||||
}
|
||||
common.OptionMapRWMutex.RUnlock()
|
||||
snapshot, err := BuildRequestPolicy(options)
|
||||
if err == nil {
|
||||
requestPolicySnapshot.Store(snapshot)
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
// UpdateRequestPolicyOptions validates the entire resulting policy before any
|
||||
// write. Unknown JSON rule fields stay in the original option string.
|
||||
func UpdateRequestPolicyOptions(values map[string]string) error {
|
||||
requestPolicyOptionMutex.Lock()
|
||||
defer requestPolicyOptionMutex.Unlock()
|
||||
options := maps.Clone(CurrentRequestPolicy().Options)
|
||||
for key, value := range values {
|
||||
if !IsRequestPolicyOption(key) {
|
||||
return fmt.Errorf("not a request policy option: %s", key)
|
||||
}
|
||||
options[key] = value
|
||||
}
|
||||
snapshot, err := BuildRequestPolicy(options)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := DB.Transaction(func(tx *gorm.DB) error {
|
||||
for key, value := range values {
|
||||
option := Option{Key: key}
|
||||
if err := tx.FirstOrCreate(&option, Option{Key: key}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tx.Model(&option).Update("value", value).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}); err != nil {
|
||||
return err
|
||||
}
|
||||
for key, value := range values {
|
||||
if err := updateOptionMap(key, value); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
requestPolicySnapshot.Store(snapshot)
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,122 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"maps"
|
||||
"math"
|
||||
"os"
|
||||
"strconv"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/setting/operation_setting"
|
||||
"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"
|
||||
)
|
||||
|
||||
func TestRequestPolicyDatabaseMatrix(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":
|
||||
dsn := os.Getenv("TEST_MYSQL_DSN")
|
||||
if dsn == "" {
|
||||
t.Skip("TEST_MYSQL_DSN not configured")
|
||||
}
|
||||
driver = mysql.Open(dsn)
|
||||
case "postgres":
|
||||
dsn := os.Getenv("TEST_POSTGRES_DSN")
|
||||
if dsn == "" {
|
||||
t.Skip("TEST_POSTGRES_DSN not configured")
|
||||
}
|
||||
driver = postgres.Open(dsn)
|
||||
}
|
||||
db, err := gorm.Open(driver, &gorm.Config{NamingStrategy: schema.NamingStrategy{TablePrefix: "policy_test_"}})
|
||||
require.NoError(t, err)
|
||||
sqlDB, err := db.DB()
|
||||
require.NoError(t, err)
|
||||
sqlDB.SetMaxOpenConns(1)
|
||||
previousDB, previousType, previousSnapshot := DB, common.MainDatabaseType(), CurrentRequestPolicy()
|
||||
common.OptionMapRWMutex.Lock()
|
||||
previousOptions := maps.Clone(common.OptionMap)
|
||||
common.OptionMap = maps.Clone(previousSnapshot.Options)
|
||||
common.OptionMapRWMutex.Unlock()
|
||||
DB = db
|
||||
common.SetMainDatabaseType(common.DatabaseType(dialect))
|
||||
initCol()
|
||||
t.Cleanup(func() {
|
||||
require.NoError(t, db.Migrator().DropTable(&Option{}))
|
||||
for k, v := range previousSnapshot.Options {
|
||||
require.NoError(t, updateOptionMap(k, v))
|
||||
}
|
||||
requestPolicySnapshot.Store(previousSnapshot)
|
||||
common.OptionMapRWMutex.Lock()
|
||||
common.OptionMap = previousOptions
|
||||
common.OptionMapRWMutex.Unlock()
|
||||
DB = previousDB
|
||||
common.SetMainDatabaseType(previousType)
|
||||
initCol()
|
||||
require.NoError(t, sqlDB.Close())
|
||||
})
|
||||
require.NoError(t, db.AutoMigrate(&Option{}))
|
||||
var version string
|
||||
query := "SELECT version()"
|
||||
if dialect == "sqlite" {
|
||||
query = "SELECT sqlite_version()"
|
||||
}
|
||||
require.NoError(t, db.Raw(query).Scan(&version).Error)
|
||||
t.Logf("database version: %s", version)
|
||||
rules := `[{"name":"session","model_regex":[".*"],"key_sources":[{"type":"request_header","key":"X-Session"}],"ttl_seconds":0,"skip_retry_on_failure":false,"include_using_group":false,"param_override_template":{"temperature":0},"future_field":{"enabled":false}}]`
|
||||
require.NoError(t, UpdateRequestPolicyOptions(map[string]string{"RetryTimes": "2", "channel_affinity_setting.rules": rules}))
|
||||
assert.Equal(t, 2, CurrentRequestPolicy().RetryTimes)
|
||||
assert.Equal(t, 2, common.RetryTimes, "the runtime global follows the same write")
|
||||
require.NoError(t, UpdateOption("AutomaticRetryStatusCodes", "429,500-503"))
|
||||
assert.Equal(t, "429,500-503", CurrentRequestPolicy().Options["AutomaticRetryStatusCodes"])
|
||||
assert.Equal(t, "429,500-503", operation_setting.AutomaticRetryStatusCodesToString())
|
||||
assert.Error(t, UpdateRequestPolicyOptions(map[string]string{"request_policy_setting.mode": "off"}), "the removed mode switch is not a policy option")
|
||||
assert.False(t, CurrentRequestPolicy().Affinity.Rules[0].SkipRetryOnFailure)
|
||||
assert.Equal(t, rules, CurrentRequestPolicy().Options["channel_affinity_setting.rules"])
|
||||
loadOptionsFromDatabase()
|
||||
loadOptionsFromDatabase()
|
||||
assert.Equal(t, rules, CurrentRequestPolicy().Options["channel_affinity_setting.rules"], "legacy rule JSON survives reloads without dropping extension fields")
|
||||
require.NoError(t, UpdateRequestPolicyOptions(map[string]string{"channel_affinity_setting.session_mode": "strict"}))
|
||||
loadOptionsFromDatabase()
|
||||
loadOptionsFromDatabase()
|
||||
assert.Equal(t, "strict", CurrentRequestPolicy().Affinity.SessionMode)
|
||||
assert.Equal(t, rules, CurrentRequestPolicy().Options["channel_affinity_setting.rules"], "a global mode never rewrites rule fields")
|
||||
snapshot := CurrentRequestPolicy()
|
||||
assert.Error(t, UpdateRequestPolicyOptions(map[string]string{"channel_affinity_setting.session_mode": "unknown"}))
|
||||
assert.Same(t, snapshot, CurrentRequestPolicy())
|
||||
for _, value := range []string{"-1", "1.5", "bad", strconv.Itoa(math.MaxInt)} {
|
||||
assert.Error(t, UpdateRequestPolicyOptions(map[string]string{"RetryTimes": value}))
|
||||
assert.Same(t, snapshot, CurrentRequestPolicy())
|
||||
}
|
||||
assert.Error(t, UpdateRequestPolicyOptions(map[string]string{"channel_affinity_setting.rules": "[", "RetryTimes": "8"}))
|
||||
assert.Same(t, snapshot, CurrentRequestPolicy())
|
||||
var option Option
|
||||
require.NoError(t, db.Where(map[string]any{"key": "RetryTimes"}).First(&option).Error)
|
||||
assert.Equal(t, "2", option.Value)
|
||||
require.NoError(t, db.Callback().Update().Before("gorm:update").Register("policy_fail", func(tx *gorm.DB) {
|
||||
if tx.Statement.Table == "policy_test_options" {
|
||||
tx.AddError(errors.New("save rejected"))
|
||||
}
|
||||
}))
|
||||
assert.Error(t, UpdateRequestPolicyOptions(map[string]string{"RetryTimes": "8", "channel_affinity_setting.session_mode": "off"}))
|
||||
assert.Same(t, snapshot, CurrentRequestPolicy())
|
||||
require.NoError(t, db.Callback().Update().Remove("policy_fail"))
|
||||
option = Option{}
|
||||
require.NoError(t, db.Where(map[string]any{"key": "RetryTimes"}).First(&option).Error)
|
||||
assert.Equal(t, "2", option.Value)
|
||||
loadOptionsFromDatabase()
|
||||
assert.Equal(t, "strict", CurrentRequestPolicy().Affinity.SessionMode, "a failed save keeps the persisted global mode")
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,52 @@
|
||||
package model
|
||||
|
||||
// SelectionTrace is populated by the selector while its candidate snapshot is
|
||||
// still locked. It never retains Channel objects, keys or provider settings.
|
||||
type SelectionTrace struct {
|
||||
Candidates []SelectionCandidate `json:"candidates"`
|
||||
Total int `json:"candidate_total"`
|
||||
SelectedID int `json:"-"`
|
||||
}
|
||||
|
||||
type SelectionCandidate struct {
|
||||
ChannelID int `json:"channel_id"`
|
||||
Name string `json:"name"`
|
||||
Group string `json:"group"`
|
||||
Model string `json:"model"`
|
||||
Priority int64 `json:"priority"`
|
||||
Weight int `json:"weight"`
|
||||
Reason string `json:"reason"`
|
||||
}
|
||||
|
||||
func (t *SelectionTrace) Add(ch *Channel, id int, group, model, reason string) {
|
||||
if t == nil {
|
||||
return
|
||||
}
|
||||
t.Total++
|
||||
if len(t.Candidates) >= 100 {
|
||||
return
|
||||
}
|
||||
c := SelectionCandidate{ChannelID: id, Group: group, Model: model, Reason: reason}
|
||||
if ch != nil {
|
||||
c.Name = ch.Name
|
||||
c.Priority = ch.GetPriority()
|
||||
c.Weight = ch.GetWeight()
|
||||
}
|
||||
t.Candidates = append(t.Candidates, c)
|
||||
}
|
||||
|
||||
func (t *SelectionTrace) Finish(ch *Channel) {
|
||||
if t == nil || ch == nil {
|
||||
return
|
||||
}
|
||||
t.SelectedID = ch.Id
|
||||
for i := range t.Candidates {
|
||||
c := &t.Candidates[i]
|
||||
if c.ChannelID == ch.Id && c.Reason == "eligible" {
|
||||
c.Reason = "selected"
|
||||
}
|
||||
if c.Reason == "eligible" && c.Priority != ch.GetPriority() {
|
||||
c.Reason = "different_priority"
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,90 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/hex"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
)
|
||||
|
||||
const statusPageOptionPrefix = "StatusPageGroup:"
|
||||
|
||||
type StatusPageGroup struct {
|
||||
Alias string `json:"alias"`
|
||||
Hidden bool `json:"hidden"`
|
||||
ResetBefore int64 `json:"reset_before"`
|
||||
}
|
||||
|
||||
func GetStatusPageGroups(ctx context.Context) (map[string]StatusPageGroup, error) {
|
||||
var rows []Option
|
||||
if err := DB.WithContext(ctx).Where(clause.Like{Column: clause.Column{Name: "key"}, Value: statusPageOptionPrefix + "%"}).Find(&rows).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
groups := make(map[string]StatusPageGroup, len(rows))
|
||||
for _, row := range rows {
|
||||
name, err := hex.DecodeString(strings.TrimPrefix(row.Key, statusPageOptionPrefix))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var settings StatusPageGroup
|
||||
if err := common.UnmarshalJsonStr(row.Value, &settings); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
groups[string(name)] = settings
|
||||
}
|
||||
return groups, nil
|
||||
}
|
||||
|
||||
// One persisted cutoff per group preserves shared-channel history and the
|
||||
// detection wallet ledger. Runs started before a reset stay excluded on completion.
|
||||
func SaveStatusPageGroup(ctx context.Context, group string, display *StatusPageGroup) error {
|
||||
key := statusPageOptionPrefix + hex.EncodeToString([]byte(group))
|
||||
return DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
if err := tx.Clauses(clause.OnConflict{DoNothing: true}).Create(&Option{Key: key, Value: "{}"}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
var row Option
|
||||
if err := lockForUpdate(tx).Where(&Option{Key: key}).First(&row).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
var settings StatusPageGroup
|
||||
if err := common.UnmarshalJsonStr(row.Value, &settings); err != nil {
|
||||
return err
|
||||
}
|
||||
if display == nil {
|
||||
settings.ResetBefore = max(settings.ResetBefore, time.Now().Unix())
|
||||
} else {
|
||||
settings.Alias, settings.Hidden = display.Alias, display.Hidden
|
||||
}
|
||||
encoded, err := common.Marshal(settings)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Model(&Option{}).Where(&Option{Key: key}).Update("value", string(encoded)).Error
|
||||
})
|
||||
}
|
||||
|
||||
func GetStatusPageModels(ctx context.Context, channels []StatusChannel) (map[string][]string, error) {
|
||||
ids := make([]int, 0, len(channels))
|
||||
for _, channel := range channels {
|
||||
if channel.Status == common.ChannelStatusEnabled {
|
||||
ids = append(ids, channel.Id)
|
||||
}
|
||||
}
|
||||
result := map[string][]string{}
|
||||
if len(ids) == 0 {
|
||||
return result, nil
|
||||
}
|
||||
var abilities []Ability
|
||||
if err := DB.WithContext(ctx).Select("group", "model").Where("enabled = ? AND channel_id IN ?", true, ids).Distinct().Order("model").Find(&abilities).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for _, ability := range abilities {
|
||||
result[ability.Group] = append(result[ability.Group], ability.Model)
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
+56
-10
@@ -16,11 +16,12 @@ const (
|
||||
SystemTaskStatusSucceeded SystemTaskStatus = "succeeded"
|
||||
SystemTaskStatusFailed SystemTaskStatus = "failed"
|
||||
|
||||
SystemTaskTypeLogCleanup = "log_cleanup"
|
||||
SystemTaskTypeChannelTest = "channel_test"
|
||||
SystemTaskTypeModelUpdate = "model_update"
|
||||
SystemTaskTypeMidjourneyPoll = "midjourney_poll"
|
||||
SystemTaskTypeAsyncTaskPoll = "async_task_poll"
|
||||
SystemTaskTypeLogCleanup = "log_cleanup"
|
||||
SystemTaskTypeChannelTest = "channel_test"
|
||||
SystemTaskTypeChannelAvailability = "channel_availability"
|
||||
SystemTaskTypeModelUpdate = "model_update"
|
||||
SystemTaskTypeMidjourneyPoll = "midjourney_poll"
|
||||
SystemTaskTypeAsyncTaskPoll = "async_task_poll"
|
||||
)
|
||||
|
||||
var ErrSystemTaskLockLost = errors.New("system task lock lost")
|
||||
@@ -175,16 +176,61 @@ func FindEarliestPendingSystemTasks(taskTypes []string) (map[string]*SystemTask,
|
||||
return tasksByType, nil
|
||||
}
|
||||
|
||||
func ListSystemTasks(limit int) ([]*SystemTask, error) {
|
||||
type SystemTaskFilter struct {
|
||||
Scope string `form:"scope" binding:"omitempty,oneof=active history"`
|
||||
Type string `form:"type" binding:"max=64"`
|
||||
Status SystemTaskStatus `form:"status" binding:"omitempty,oneof=pending running succeeded failed"`
|
||||
}
|
||||
|
||||
func (filter SystemTaskFilter) query() *gorm.DB {
|
||||
query := DB.Model(&SystemTask{})
|
||||
switch filter.Scope {
|
||||
case "active":
|
||||
query = query.Where("status IN ?", activeSystemTaskStatuses())
|
||||
case "history":
|
||||
query = query.Where("status IN ?", []SystemTaskStatus{SystemTaskStatusSucceeded, SystemTaskStatusFailed})
|
||||
}
|
||||
if filter.Type != "" {
|
||||
query = query.Where("type = ?", filter.Type)
|
||||
}
|
||||
if filter.Status != "" {
|
||||
query = query.Where("status = ?", filter.Status)
|
||||
}
|
||||
return query
|
||||
}
|
||||
|
||||
func ListSystemTasks(filter SystemTaskFilter, offset, limit int) ([]*SystemTask, int64, error) {
|
||||
if limit <= 0 {
|
||||
limit = 20
|
||||
}
|
||||
if limit > 100 {
|
||||
limit = 100
|
||||
query := filter.query()
|
||||
var total int64
|
||||
if err := query.Count(&total).Error; err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
var tasks []*SystemTask
|
||||
err := DB.Order("id desc").Limit(limit).Find(&tasks).Error
|
||||
return tasks, err
|
||||
err := query.Order("id desc").Offset(max(0, offset)).Limit(min(limit, 100)).Find(&tasks).Error
|
||||
return tasks, total, err
|
||||
}
|
||||
|
||||
// DeleteSystemTaskHistory preserves the latest row of every type because the
|
||||
// scheduler uses it to determine when the next run is due. Active rows and rows
|
||||
// created after this cleanup's snapshot are never removed.
|
||||
func DeleteSystemTaskHistory(filter SystemTaskFilter) (int64, error) {
|
||||
var latestIDs []int64
|
||||
if err := DB.Model(&SystemTask{}).Select("MAX(id)").Group("type").Pluck("MAX(id)", &latestIDs).Error; err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if len(latestIDs) == 0 {
|
||||
return 0, nil
|
||||
}
|
||||
var newestID int64
|
||||
for _, id := range latestIDs {
|
||||
newestID = max(newestID, id)
|
||||
}
|
||||
filter.Scope = "history"
|
||||
result := filter.query().Where("id < ? AND id NOT IN ?", newestID, latestIDs).Delete(&SystemTask{})
|
||||
return result.RowsAffected, result.Error
|
||||
}
|
||||
|
||||
// GetLatestSystemTask returns the most recent task row of the given type
|
||||
|
||||
@@ -1,14 +1,123 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"os"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"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"
|
||||
)
|
||||
|
||||
func TestSystemTaskHistoryDatabaseMatrix(t *testing.T) {
|
||||
for _, dialect := range []string{"sqlite", "mysql", "postgres"} {
|
||||
t.Run(dialect, func(t *testing.T) {
|
||||
var dialector gorm.Dialector
|
||||
switch dialect {
|
||||
case "sqlite":
|
||||
dialector = sqlite.Open(":memory:")
|
||||
case "mysql":
|
||||
dsn := os.Getenv("TEST_MYSQL_DSN")
|
||||
if dsn == "" {
|
||||
t.Skip("TEST_MYSQL_DSN is not configured")
|
||||
}
|
||||
dialector = mysql.Open(dsn)
|
||||
case "postgres":
|
||||
dsn := os.Getenv("TEST_POSTGRES_DSN")
|
||||
if dsn == "" {
|
||||
t.Skip("TEST_POSTGRES_DSN is not configured")
|
||||
}
|
||||
dialector = postgres.Open(dsn)
|
||||
}
|
||||
db, err := gorm.Open(dialector, &gorm.Config{NamingStrategy: schema.NamingStrategy{TablePrefix: "history_test_"}})
|
||||
require.NoError(t, err)
|
||||
sqlDB, err := db.DB()
|
||||
require.NoError(t, err)
|
||||
sqlDB.SetMaxOpenConns(1)
|
||||
previousDB := DB
|
||||
DB = db
|
||||
t.Cleanup(func() {
|
||||
DB = previousDB
|
||||
require.NoError(t, db.Migrator().DropTable(&SystemTask{}))
|
||||
require.NoError(t, sqlDB.Close())
|
||||
})
|
||||
require.NoError(t, db.AutoMigrate(&SystemTask{}))
|
||||
var version string
|
||||
versionQuery := "SELECT version()"
|
||||
if dialect == "sqlite" {
|
||||
versionQuery = "SELECT sqlite_version()"
|
||||
}
|
||||
require.NoError(t, db.Raw(versionQuery).Scan(&version).Error)
|
||||
t.Logf("database version: %s", version)
|
||||
|
||||
tasks := []SystemTask{
|
||||
{TaskID: "model-failed-old", Type: SystemTaskTypeModelUpdate, Status: SystemTaskStatusFailed},
|
||||
{TaskID: "model-success", Type: SystemTaskTypeModelUpdate, Status: SystemTaskStatusSucceeded},
|
||||
{TaskID: "channel-failed", Type: SystemTaskTypeChannelTest, Status: SystemTaskStatusFailed},
|
||||
{TaskID: "model-failed-recent", Type: SystemTaskTypeModelUpdate, Status: SystemTaskStatusFailed},
|
||||
{TaskID: "model-latest", Type: SystemTaskTypeModelUpdate, Status: SystemTaskStatusSucceeded},
|
||||
{TaskID: "channel-pending", Type: SystemTaskTypeChannelTest, Status: SystemTaskStatusPending},
|
||||
{TaskID: "channel-running", Type: SystemTaskTypeChannelTest, Status: SystemTaskStatusRunning},
|
||||
{TaskID: "log-only-history", Type: SystemTaskTypeLogCleanup, Status: SystemTaskStatusSucceeded},
|
||||
}
|
||||
require.NoError(t, db.Create(&tasks).Error)
|
||||
filter := SystemTaskFilter{Scope: "history", Type: SystemTaskTypeModelUpdate, Status: SystemTaskStatusFailed}
|
||||
page, total, err := ListSystemTasks(filter, 1, 1)
|
||||
require.NoError(t, err)
|
||||
assert.EqualValues(t, 2, total)
|
||||
require.Len(t, page, 1)
|
||||
assert.Equal(t, "model-failed-old", page[0].TaskID)
|
||||
page, total, err = ListSystemTasks(filter, 2, 1)
|
||||
require.NoError(t, err)
|
||||
assert.EqualValues(t, 2, total)
|
||||
assert.Empty(t, page)
|
||||
|
||||
active, total, err := ListSystemTasks(SystemTaskFilter{Scope: "active"}, 0, 20)
|
||||
require.NoError(t, err)
|
||||
assert.EqualValues(t, 2, total)
|
||||
require.Len(t, active, 2)
|
||||
assert.Equal(t, "channel-running", active[0].TaskID)
|
||||
assert.Equal(t, "channel-pending", active[1].TaskID)
|
||||
|
||||
deleted, err := DeleteSystemTaskHistory(filter)
|
||||
require.NoError(t, err)
|
||||
assert.EqualValues(t, 2, deleted)
|
||||
_, total, err = ListSystemTasks(SystemTaskFilter{Scope: "history"}, 0, 20)
|
||||
require.NoError(t, err)
|
||||
assert.EqualValues(t, 4, total, "other types and statuses are untouched")
|
||||
|
||||
deleted, err = DeleteSystemTaskHistory(SystemTaskFilter{Scope: "active", Status: SystemTaskStatusRunning})
|
||||
require.NoError(t, err)
|
||||
assert.Zero(t, deleted, "caller cannot override the terminal-status restriction")
|
||||
deleted, err = DeleteSystemTaskHistory(SystemTaskFilter{})
|
||||
require.NoError(t, err)
|
||||
assert.EqualValues(t, 2, deleted)
|
||||
remaining, total, err := ListSystemTasks(SystemTaskFilter{}, 0, 100)
|
||||
require.NoError(t, err)
|
||||
assert.EqualValues(t, 4, total)
|
||||
var remainingIDs []string
|
||||
for _, task := range remaining {
|
||||
remainingIDs = append(remainingIDs, task.TaskID)
|
||||
}
|
||||
assert.ElementsMatch(t, []string{"model-latest", "channel-pending", "channel-running", "log-only-history"}, remainingIDs)
|
||||
latest, err := GetLatestSystemTasks([]string{SystemTaskTypeModelUpdate, SystemTaskTypeChannelTest, SystemTaskTypeLogCleanup})
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "model-latest", latest[SystemTaskTypeModelUpdate].TaskID)
|
||||
assert.Equal(t, "channel-running", latest[SystemTaskTypeChannelTest].TaskID)
|
||||
assert.Equal(t, "log-only-history", latest[SystemTaskTypeLogCleanup].TaskID)
|
||||
deleted, err = DeleteSystemTaskHistory(SystemTaskFilter{})
|
||||
require.NoError(t, err)
|
||||
assert.Zero(t, deleted, "repeated cleanup preserves the scheduler's latest runs")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
type testSystemTaskPayload struct {
|
||||
TargetTimestamp int64 `json:"target_timestamp"`
|
||||
BatchSize int `json:"batch_size"`
|
||||
|
||||
@@ -56,3 +56,27 @@ func identityFilters(key string, channelTypes []int) []dto.ChannelFilter {
|
||||
TaskPluginChannelTypes: channelTypes,
|
||||
}}
|
||||
}
|
||||
|
||||
func TestSharedPluginKeysFilterBothChannelSources(t *testing.T) {
|
||||
truncateTables(t)
|
||||
originalMemoryCache := common.MemoryCacheEnabled
|
||||
common.MemoryCacheEnabled = true
|
||||
t.Cleanup(func() { common.MemoryCacheEnabled = originalMemoryCache; InitChannelCache() })
|
||||
filters := []dto.ChannelFilter{{Kind: dto.FilterTaskPluginIdentity, TaskPluginKey: "alpha", TaskPluginKeys: []string{"alpha", "beta"}}}
|
||||
var abilities []Ability
|
||||
for index, key := range []string{"alpha", "beta", "unrelated"} {
|
||||
setting := `{"task_plugin_key":"` + key + `"}`
|
||||
channel := Channel{Id: 910001 + index, Type: constant.ChannelTypeTaskPlugin, Status: common.ChannelStatusEnabled, Name: key, Models: "shared", Group: "default", Setting: &setting}
|
||||
require.NoError(t, channel.Insert())
|
||||
abilities = append(abilities, Ability{ChannelId: channel.Id, Model: "shared", Group: "default", Enabled: true})
|
||||
matches, _ := ChannelSatisfiesFilters(&channel, "shared", filters)
|
||||
assert.Equal(t, index < 2, matches)
|
||||
}
|
||||
assert.Equal(t, abilities[:2], filterAbilitiesByConstraints(abilities, "shared", filters))
|
||||
InitChannelCache()
|
||||
channelSyncLock.RLock()
|
||||
kept, emptied := filterCandidateIDs([]int{910001, 910002, 910003}, "shared", filters)
|
||||
channelSyncLock.RUnlock()
|
||||
assert.Equal(t, []int{910001, 910002}, kept)
|
||||
assert.Empty(t, emptied)
|
||||
}
|
||||
+80
-22
@@ -1,6 +1,7 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
@@ -12,34 +13,78 @@ import (
|
||||
)
|
||||
|
||||
type Token struct {
|
||||
Id int `json:"id"`
|
||||
UserId int `json:"user_id" gorm:"index"`
|
||||
Key string `json:"key" gorm:"type:varchar(128);uniqueIndex"`
|
||||
Status int `json:"status" gorm:"default:1"`
|
||||
Name string `json:"name" gorm:"index" `
|
||||
CreatedTime int64 `json:"created_time" gorm:"bigint"`
|
||||
AccessedTime int64 `json:"accessed_time" gorm:"bigint"`
|
||||
ExpiredTime int64 `json:"expired_time" gorm:"bigint;default:-1"` // -1 means never expired
|
||||
RemainQuota int `json:"remain_quota" gorm:"default:0"`
|
||||
UnlimitedQuota bool `json:"unlimited_quota"`
|
||||
ModelLimitsEnabled bool `json:"model_limits_enabled"`
|
||||
ModelLimits string `json:"model_limits" gorm:"type:text"`
|
||||
AllowIps *string `json:"allow_ips" gorm:"default:''"`
|
||||
UsedQuota int `json:"used_quota" gorm:"default:0"` // used quota
|
||||
Group string `json:"group" gorm:"default:''"`
|
||||
CrossGroupRetry bool `json:"cross_group_retry"` // 跨分组重试,仅auto分组有效
|
||||
AutoGroups string `json:"-" gorm:"type:text"`
|
||||
DeletedAt gorm.DeletedAt `gorm:"index"`
|
||||
Id int `json:"id"`
|
||||
UserId int `json:"user_id" gorm:"index"`
|
||||
Key string `json:"key" gorm:"type:varchar(128);uniqueIndex"`
|
||||
Status int `json:"status" gorm:"default:1"`
|
||||
Name string `json:"name" gorm:"index" `
|
||||
CreatedTime int64 `json:"created_time" gorm:"bigint"`
|
||||
AccessedTime int64 `json:"accessed_time" gorm:"bigint"`
|
||||
ExpiredTime int64 `json:"expired_time" gorm:"bigint;default:-1"` // -1 means never expired
|
||||
RemainQuota int `json:"remain_quota" gorm:"default:0"`
|
||||
UnlimitedQuota bool `json:"unlimited_quota"`
|
||||
ModelLimitsEnabled bool `json:"model_limits_enabled"`
|
||||
ModelLimits string `json:"model_limits" gorm:"type:text"`
|
||||
AllowIps *string `json:"allow_ips" gorm:"default:''"`
|
||||
UsedQuota int `json:"used_quota" gorm:"default:0"` // used quota
|
||||
Group string `json:"group" gorm:"default:''"`
|
||||
CrossGroupRetry bool `json:"cross_group_retry"` // 跨分组重试,仅auto分组有效
|
||||
AstraDetectionEnabled bool `json:"astra_detection_enabled" gorm:"not null;default:false"`
|
||||
AstraAutoSwitch bool `json:"astra_auto_switch" gorm:"not null;default:false"`
|
||||
AutoGroups string `json:"-" gorm:"type:text"`
|
||||
RoutingPolicy string `json:"-" gorm:"type:text"`
|
||||
DeletedAt gorm.DeletedAt `gorm:"index"`
|
||||
}
|
||||
|
||||
func (token *Token) GetAutoGroups() ([]string, error) {
|
||||
// TokenGroupEntry is one group of a token's routing pool. Lower Priority values
|
||||
// are preferred, and entries sharing a Priority are load-balanced together.
|
||||
type TokenGroupEntry struct {
|
||||
Group string `json:"group"`
|
||||
Priority int `json:"priority"`
|
||||
}
|
||||
|
||||
// GetGroupPool returns the token's routing pool in stored order. Stored entries
|
||||
// may be plain group names, in which case the array position is the priority.
|
||||
func (token *Token) GetGroupPool() ([]TokenGroupEntry, error) {
|
||||
if token.AutoGroups == "" {
|
||||
return nil, nil
|
||||
}
|
||||
var groups []string
|
||||
if err := common.UnmarshalJsonStr(token.AutoGroups, &groups); err != nil {
|
||||
var raw []json.RawMessage
|
||||
if err := common.UnmarshalJsonStr(token.AutoGroups, &raw); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
entries := make([]TokenGroupEntry, 0, len(raw))
|
||||
for index, item := range raw {
|
||||
var name string
|
||||
if err := json.Unmarshal(item, &name); err == nil {
|
||||
if name != "" {
|
||||
entries = append(entries, TokenGroupEntry{Group: name, Priority: index})
|
||||
}
|
||||
continue
|
||||
}
|
||||
var entry TokenGroupEntry
|
||||
if err := json.Unmarshal(item, &entry); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if entry.Group != "" {
|
||||
entries = append(entries, entry)
|
||||
}
|
||||
}
|
||||
return entries, nil
|
||||
}
|
||||
|
||||
func (token *Token) GetAutoGroups() ([]string, error) {
|
||||
entries, err := token.GetGroupPool()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(entries) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
groups := make([]string, 0, len(entries))
|
||||
for _, entry := range entries {
|
||||
groups = append(groups, entry.Group)
|
||||
}
|
||||
return groups, nil
|
||||
}
|
||||
|
||||
@@ -56,6 +101,19 @@ func (token *Token) SetAutoGroups(groups []string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (token *Token) SetGroupPool(entries []TokenGroupEntry) error {
|
||||
if len(entries) == 0 {
|
||||
token.AutoGroups = ""
|
||||
return nil
|
||||
}
|
||||
data, err := common.Marshal(entries)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
token.AutoGroups = string(data)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (token *Token) Clean() {
|
||||
token.Key = ""
|
||||
}
|
||||
@@ -313,7 +371,7 @@ func (token *Token) Update() (err error) {
|
||||
common.SysLog("failed to invalidate token cache before update: " + cacheErr.Error())
|
||||
}
|
||||
return DB.Model(token).Select("name", "status", "expired_time", "remain_quota", "unlimited_quota",
|
||||
"model_limits_enabled", "model_limits", "allow_ips", "group", "cross_group_retry", "auto_groups").Updates(token).Error
|
||||
"model_limits_enabled", "model_limits", "allow_ips", "group", "cross_group_retry", "astra_auto_switch", "astra_detection_enabled", "auto_groups", "routing_policy").Updates(token).Error
|
||||
}
|
||||
|
||||
func (token *Token) SelectUpdate() (err error) {
|
||||
|
||||
@@ -0,0 +1,42 @@
|
||||
package model
|
||||
|
||||
import "gorm.io/gorm"
|
||||
|
||||
const astraAutoSwitchMigrationKey = "migration.token_astra_auto_switch.v1"
|
||||
|
||||
// MigrateTokenAstraAutoSwitch preserves the pre-token-policy behavior for
|
||||
// tokens that existed before the user-level switch was introduced. It is
|
||||
// guarded by an Option marker so newly-created tokens (which default to false)
|
||||
// are never promoted on a later restart.
|
||||
func MigrateTokenAstraAutoSwitch() error {
|
||||
return DB.Transaction(func(tx *gorm.DB) error {
|
||||
var marker Option
|
||||
if err := tx.Where(&Option{Key: astraAutoSwitchMigrationKey}).First(&marker).Error; err == nil {
|
||||
return nil
|
||||
} else if err != gorm.ErrRecordNotFound {
|
||||
return err
|
||||
}
|
||||
if err := tx.Model(&Token{}).Where("deleted_at IS NULL").Update("astra_auto_switch", true).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Create(&Option{Key: astraAutoSwitchMigrationKey, Value: "true"}).Error
|
||||
})
|
||||
}
|
||||
|
||||
// MigrateTokenAstraDetection preserves observation for existing tokens without
|
||||
// resetting an explicit opt-out on subsequent starts.
|
||||
func MigrateTokenAstraDetection() error {
|
||||
return DB.Transaction(func(tx *gorm.DB) error {
|
||||
const key = "migration.token_astra_detection.v1"
|
||||
var marker Option
|
||||
if err := tx.Where(&Option{Key: key}).First(&marker).Error; err == nil {
|
||||
return nil
|
||||
} else if err != gorm.ErrRecordNotFound {
|
||||
return err
|
||||
}
|
||||
if err := tx.Model(&Token{}).Where("deleted_at IS NULL").Update("astra_detection_enabled", true).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Create(&Option{Key: key, Value: "true"}).Error
|
||||
})
|
||||
}
|
||||
@@ -1,6 +1,7 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
@@ -8,6 +9,23 @@ import (
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestTokenAstraCacheUpgradePreservesLiveQuota(t *testing.T) {
|
||||
useUserCacheMiniRedis(t)
|
||||
ctx := context.Background()
|
||||
token := Token{Id: 43, UserId: 7, Key: "legacy-astra-cache", AstraAutoSwitch: true, RemainQuota: 100}
|
||||
require.NoError(t, cacheSetTokenForTest(token))
|
||||
key := getTokenCacheKey(token.Key)
|
||||
require.NoError(t, common.RDB.HDel(ctx, key, "AstraAutoSwitch").Err())
|
||||
require.NoError(t, common.RDB.HSet(ctx, key, "RemainQuota", 70).Err())
|
||||
_, err := cacheGetTokenByKey(token.Key)
|
||||
require.Error(t, err, "legacy cache must fetch the migrated value from the database")
|
||||
require.NoError(t, cacheSetTokenForTest(token))
|
||||
cached, err := cacheGetTokenByKey(token.Key)
|
||||
require.NoError(t, err)
|
||||
require.True(t, cached.AstraAutoSwitch)
|
||||
require.Equal(t, 70, cached.RemainQuota, "metadata hydration must not overwrite concurrent quota consumption")
|
||||
}
|
||||
|
||||
func TestTokenAutoGroupsRoundTripThroughRedisHashCache(t *testing.T) {
|
||||
useUserCacheMiniRedis(t)
|
||||
token := Token{
|
||||
@@ -41,6 +59,7 @@ func TestTokenUpdateSynchronouslyNarrowsPreheatedAutoGroupsCache(t *testing.T) {
|
||||
Group: "auto",
|
||||
CrossGroupRetry: true,
|
||||
AutoGroups: `["default","vip"]`,
|
||||
AstraAutoSwitch: true,
|
||||
}
|
||||
require.NoError(t, token.Insert())
|
||||
require.NoError(t, cacheSetTokenForTest(token))
|
||||
@@ -50,6 +69,7 @@ func TestTokenUpdateSynchronouslyNarrowsPreheatedAutoGroupsCache(t *testing.T) {
|
||||
assert.JSONEq(t, `["default","vip"]`, preheated.AutoGroups)
|
||||
|
||||
require.NoError(t, token.SetAutoGroups([]string{"vip"}))
|
||||
token.AstraAutoSwitch = false
|
||||
require.NoError(t, token.Update())
|
||||
// Update 是限制性变更:写库前删除缓存并设置 fence。缓存不再提供旧的
|
||||
// 宽分组值,下一次读取必须看到收紧后的分组。
|
||||
@@ -58,6 +78,7 @@ func TestTokenUpdateSynchronouslyNarrowsPreheatedAutoGroupsCache(t *testing.T) {
|
||||
reloaded, err := GetTokenByKey(token.Key, false)
|
||||
require.NoError(t, err)
|
||||
assert.JSONEq(t, `["vip"]`, reloaded.AutoGroups)
|
||||
assert.False(t, reloaded.AstraAutoSwitch)
|
||||
}
|
||||
|
||||
// cacheSetTokenForTest 以测试身份写入完整 token 缓存(含额度字段),
|
||||
|
||||
+21
-6
@@ -10,7 +10,7 @@ import (
|
||||
)
|
||||
|
||||
func getTokenCacheKey(key string) string {
|
||||
return fmt.Sprintf("token:%s", common.GenerateHMAC(key))
|
||||
return fmt.Sprintf("token:v2:%s", common.GenerateHMAC(key))
|
||||
}
|
||||
|
||||
func getTokenCacheFenceKey(key string) string {
|
||||
@@ -51,7 +51,7 @@ func invalidateTokenCacheForMutation(key string) error {
|
||||
// active and the hash is cold. An existing hash only gets its TTL refreshed:
|
||||
// its RemainQuota may already be ahead of this snapshot because atomic
|
||||
// pre-consume decrements Redis first, so a snapshot must never overwrite any
|
||||
// field of a live hash.
|
||||
// quota field of a live hash. Only missing policy metadata is hydrated.
|
||||
// 返回值:0=被 fence 拦截,1=完成初始化,2=哈希已存在,仅刷新 TTL。
|
||||
func cacheInitToken(token Token) (int, error) {
|
||||
if !common.RedisEnabled {
|
||||
@@ -66,7 +66,9 @@ if redis.call('EXISTS', KEYS[2]) == 1 then
|
||||
return 0
|
||||
end
|
||||
if redis.call('EXISTS', KEYS[1]) == 1 then
|
||||
redis.call('EXPIRE', KEYS[1], ARGV[17])
|
||||
redis.call('HSETNX', KEYS[1], 'AstraAutoSwitch', ARGV[14])
|
||||
redis.call('HSETNX', KEYS[1], 'AstraDetectionEnabled', ARGV[20])
|
||||
redis.call('EXPIRE', KEYS[1], ARGV[18])
|
||||
return 2
|
||||
end
|
||||
redis.call('HSET', KEYS[1],
|
||||
@@ -74,8 +76,8 @@ redis.call('HSET', KEYS[1],
|
||||
'CreatedTime', ARGV[5], 'AccessedTime', ARGV[6], 'ExpiredTime', ARGV[7],
|
||||
'UnlimitedQuota', ARGV[8], 'ModelLimitsEnabled', ARGV[9], 'ModelLimits', ARGV[10],
|
||||
'AllowIps', ARGV[11], 'Group', ARGV[12], 'CrossGroupRetry', ARGV[13],
|
||||
'AutoGroups', ARGV[14], 'RemainQuota', ARGV[15], 'UsedQuota', ARGV[16])
|
||||
redis.call('EXPIRE', KEYS[1], ARGV[17])
|
||||
'AstraAutoSwitch', ARGV[14], 'AutoGroups', ARGV[15], 'RemainQuota', ARGV[16], 'UsedQuota', ARGV[17], 'RoutingPolicy', ARGV[19], 'AstraDetectionEnabled', ARGV[20])
|
||||
redis.call('EXPIRE', KEYS[1], ARGV[18])
|
||||
return 1`
|
||||
|
||||
return common.RDB.Eval(context.Background(), script, []string{
|
||||
@@ -85,8 +87,9 @@ return 1`
|
||||
token.CreatedTime, token.AccessedTime, token.ExpiredTime,
|
||||
strconv.FormatBool(token.UnlimitedQuota), strconv.FormatBool(token.ModelLimitsEnabled),
|
||||
token.ModelLimits, allowIps, token.Group, strconv.FormatBool(token.CrossGroupRetry),
|
||||
token.AutoGroups, token.RemainQuota, token.UsedQuota,
|
||||
strconv.FormatBool(token.AstraAutoSwitch), token.AutoGroups, token.RemainQuota, token.UsedQuota,
|
||||
tokenCacheTTLSeconds(),
|
||||
token.RoutingPolicy, strconv.FormatBool(token.AstraDetectionEnabled),
|
||||
).Int()
|
||||
}
|
||||
|
||||
@@ -95,6 +98,18 @@ func cacheGetTokenByKey(key string) (*Token, error) {
|
||||
if !common.RedisEnabled {
|
||||
return nil, fmt.Errorf("redis is not enabled")
|
||||
}
|
||||
// A pre-migration cache hash does not contain the policy field. Force a DB
|
||||
// refresh so legacy tokens retain their migrated default instead of being
|
||||
// interpreted as false until the old cache expires.
|
||||
if exists, err := common.RDB.HExists(context.Background(), getTokenCacheKey(key), "AstraAutoSwitch").Result(); err != nil || !exists {
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return nil, fmt.Errorf("token cache is missing AstraAutoSwitch")
|
||||
}
|
||||
if exists, err := common.RDB.HExists(context.Background(), getTokenCacheKey(key), "AstraDetectionEnabled").Result(); err != nil || !exists {
|
||||
return nil, fmt.Errorf("token cache is missing AstraDetectionEnabled")
|
||||
}
|
||||
var token Token
|
||||
if err := common.RedisHGetObj(getTokenCacheKey(key), &token); err != nil {
|
||||
return nil, err
|
||||
|
||||
@@ -0,0 +1,34 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestGetGroupPoolParsesLegacyNamesAndEntries(t *testing.T) {
|
||||
legacy := &Token{AutoGroups: `["vip","pro"]`}
|
||||
pool, err := legacy.GetGroupPool()
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []TokenGroupEntry{{Group: "vip", Priority: 0}, {Group: "pro", Priority: 1}}, pool)
|
||||
|
||||
names, err := legacy.GetAutoGroups()
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []string{"vip", "pro"}, names)
|
||||
|
||||
explicit := &Token{AutoGroups: `[{"group":"vip","priority":2},{"group":"pro","priority":1}]`}
|
||||
pool, err = explicit.GetGroupPool()
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []TokenGroupEntry{{Group: "vip", Priority: 2}, {Group: "pro", Priority: 1}}, pool)
|
||||
}
|
||||
|
||||
func TestSetGroupPoolRoundTrips(t *testing.T) {
|
||||
token := &Token{}
|
||||
require.NoError(t, token.SetGroupPool([]TokenGroupEntry{{Group: "vip", Priority: 1}}))
|
||||
pool, err := token.GetGroupPool()
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []TokenGroupEntry{{Group: "vip", Priority: 1}}, pool)
|
||||
|
||||
require.NoError(t, token.SetGroupPool(nil))
|
||||
require.Empty(t, token.AutoGroups)
|
||||
}
|
||||
@@ -0,0 +1,185 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"slices"
|
||||
"strings"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
)
|
||||
|
||||
const TokenRoutingContextKey = "token_routing_policy"
|
||||
|
||||
type ModelRoutingRule struct {
|
||||
Model string `json:"model"`
|
||||
Groups []TokenGroupEntry `json:"groups"`
|
||||
}
|
||||
|
||||
type TokenRoutingPolicy struct {
|
||||
Version int `json:"version"`
|
||||
Mode string `json:"mode"`
|
||||
Groups []TokenGroupEntry `json:"groups"`
|
||||
ModelRules []ModelRoutingRule `json:"model_rules"`
|
||||
}
|
||||
|
||||
func (p *TokenRoutingPolicy) Validate(maxGroups int, allowed func(string) bool) error {
|
||||
if p.Version != 1 || (p.Mode != "custom" && p.Mode != "inherit" && p.Mode != "user") {
|
||||
return fmt.Errorf("invalid routing policy version or mode")
|
||||
}
|
||||
if p.Mode == "custom" && len(p.Groups) == 0 {
|
||||
return fmt.Errorf("a group pool must contain at least one group")
|
||||
}
|
||||
if p.Mode != "custom" && len(p.Groups) != 0 {
|
||||
return fmt.Errorf("inherited pools cannot contain custom groups")
|
||||
}
|
||||
if len(p.Groups) > maxGroups || len(p.ModelRules) > 128 {
|
||||
return fmt.Errorf("routing policy exceeds configured limits")
|
||||
}
|
||||
pool := map[string]bool{}
|
||||
for _, entry := range p.Groups {
|
||||
if entry.Group == "" || entry.Group == "auto" || entry.Group != strings.TrimSpace(entry.Group) || pool[entry.Group] || entry.Priority < 0 || entry.Priority > 10000 || !allowed(entry.Group) {
|
||||
return fmt.Errorf("invalid or inaccessible pool group: %s", entry.Group)
|
||||
}
|
||||
pool[entry.Group] = true
|
||||
}
|
||||
models := map[string]bool{}
|
||||
for _, rule := range p.ModelRules {
|
||||
if rule.Model == "" || len(rule.Model) > 256 || strings.TrimSpace(rule.Model) != rule.Model || models[rule.Model] || len(rule.Groups) == 0 || len(rule.Groups) > maxGroups {
|
||||
return fmt.Errorf("invalid or duplicate model routing rule: %s", rule.Model)
|
||||
}
|
||||
models[rule.Model] = true
|
||||
seen := map[string]bool{}
|
||||
for _, entry := range rule.Groups {
|
||||
if !pool[entry.Group] || seen[entry.Group] || entry.Priority < 0 || entry.Priority > 10000 {
|
||||
return fmt.Errorf("model rule references an invalid pool group: %s", entry.Group)
|
||||
}
|
||||
seen[entry.Group] = true
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Resolve is shared by routing and the management preview. Model preferences
|
||||
// precede remaining defaults; matching priorities stay in the same tier.
|
||||
func (p *TokenRoutingPolicy) Resolve(modelName string, inherited []string, allowed func(string) bool) []TokenGroupEntry {
|
||||
defaults := slices.Clone(p.Groups)
|
||||
if p.Mode != "custom" {
|
||||
defaults = nil
|
||||
for i, group := range inherited {
|
||||
defaults = append(defaults, TokenGroupEntry{Group: group, Priority: i})
|
||||
}
|
||||
}
|
||||
slices.SortStableFunc(defaults, func(a, b TokenGroupEntry) int { return a.Priority - b.Priority })
|
||||
pool := map[string]bool{}
|
||||
for _, entry := range defaults {
|
||||
if allowed(entry.Group) {
|
||||
pool[entry.Group] = true
|
||||
}
|
||||
}
|
||||
var preferred []TokenGroupEntry
|
||||
for _, rule := range p.ModelRules {
|
||||
if rule.Model == modelName {
|
||||
preferred = slices.Clone(rule.Groups)
|
||||
break
|
||||
}
|
||||
}
|
||||
slices.SortStableFunc(preferred, func(a, b TokenGroupEntry) int { return a.Priority - b.Priority })
|
||||
result := []TokenGroupEntry{}
|
||||
seen := map[string]bool{}
|
||||
next := 0
|
||||
for _, list := range [][]TokenGroupEntry{preferred, defaults} {
|
||||
last, tier := 0, next
|
||||
haveLast := false
|
||||
for _, entry := range list {
|
||||
if !pool[entry.Group] || seen[entry.Group] {
|
||||
continue
|
||||
}
|
||||
if haveLast && last != entry.Priority {
|
||||
tier++
|
||||
}
|
||||
result = append(result, TokenGroupEntry{Group: entry.Group, Priority: tier})
|
||||
seen[entry.Group] = true
|
||||
last = entry.Priority
|
||||
haveLast = true
|
||||
}
|
||||
if haveLast {
|
||||
next = tier + 1
|
||||
}
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
func (token *Token) GetRoutingPolicy() (*TokenRoutingPolicy, error) {
|
||||
if token.RoutingPolicy == "" {
|
||||
return nil, nil
|
||||
}
|
||||
var policy TokenRoutingPolicy
|
||||
if err := common.UnmarshalJsonStr(token.RoutingPolicy, &policy); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if policy.Version != 1 {
|
||||
return nil, fmt.Errorf("unsupported routing policy version")
|
||||
}
|
||||
return &policy, nil
|
||||
}
|
||||
|
||||
// LegacyRoutingPolicy preserves effective order, including an explicit old
|
||||
// pool which already contained the primary group. No model rule is invented.
|
||||
func (token *Token) LegacyRoutingPolicy() (*TokenRoutingPolicy, error) {
|
||||
entries, err := token.GetGroupPool()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
policy := &TokenRoutingPolicy{Version: 1, Mode: "custom", Groups: []TokenGroupEntry{}, ModelRules: []ModelRoutingRule{}}
|
||||
if token.Group == "auto" {
|
||||
if len(entries) == 0 {
|
||||
policy.Mode = "inherit"
|
||||
return policy, nil
|
||||
}
|
||||
policy.Groups = entries
|
||||
return policy, nil
|
||||
}
|
||||
if token.Group == "" {
|
||||
policy.Mode = "user"
|
||||
return policy, nil
|
||||
}
|
||||
if !token.CrossGroupRetry {
|
||||
entries = nil
|
||||
}
|
||||
minPriority := 0
|
||||
hasPrimary := false
|
||||
for _, entry := range entries {
|
||||
minPriority = min(minPriority, entry.Priority)
|
||||
hasPrimary = hasPrimary || entry.Group == token.Group
|
||||
}
|
||||
if !hasPrimary {
|
||||
policy.Groups = append(policy.Groups, TokenGroupEntry{Group: token.Group, Priority: minPriority - 1})
|
||||
}
|
||||
policy.Groups = append(policy.Groups, entries...)
|
||||
policy.Groups = policy.Resolve("", nil, func(string) bool { return true })
|
||||
return policy, nil
|
||||
}
|
||||
|
||||
func MigrateTokenRoutingPolicies() (int, error) {
|
||||
var tokens []Token
|
||||
if err := DB.Where("routing_policy IS NULL OR routing_policy = ?", "").Find(&tokens).Error; err != nil {
|
||||
return 0, err
|
||||
}
|
||||
count := 0
|
||||
for _, token := range tokens {
|
||||
policy, err := token.LegacyRoutingPolicy()
|
||||
if err != nil {
|
||||
return count, fmt.Errorf("token %d: %w", token.Id, err)
|
||||
}
|
||||
body, err := common.Marshal(policy)
|
||||
if err != nil {
|
||||
return count, err
|
||||
}
|
||||
r := DB.Model(&Token{}).Where("id = ? AND (routing_policy IS NULL OR routing_policy = ?)", token.Id, "").Update("routing_policy", string(body))
|
||||
if r.Error != nil {
|
||||
return count, r.Error
|
||||
}
|
||||
count += int(r.RowsAffected)
|
||||
}
|
||||
return count, nil
|
||||
}
|
||||
+15
-14
@@ -11,18 +11,19 @@ import (
|
||||
|
||||
// QuotaData 柱状图数据
|
||||
type QuotaData struct {
|
||||
Id int `json:"id"`
|
||||
UserID int `json:"user_id" gorm:"index"`
|
||||
Username string `json:"username" gorm:"index:idx_qdt_model_user_name,priority:2;size:64;default:''"`
|
||||
ModelName string `json:"model_name" gorm:"index:idx_qdt_model_user_name,priority:1;size:64;default:''"`
|
||||
CreatedAt int64 `json:"created_at" gorm:"bigint;index:idx_qdt_created_at,priority:2"`
|
||||
UseGroup string `json:"use_group" gorm:"index;size:64;default:''"`
|
||||
TokenID int `json:"token_id" gorm:"index;default:0"`
|
||||
ChannelID int `json:"channel_id" gorm:"index;default:0"`
|
||||
NodeName string `json:"node_name" gorm:"index;size:64;default:''"`
|
||||
TokenUsed int `json:"token_used" gorm:"default:0"`
|
||||
Count int `json:"count" gorm:"default:0"`
|
||||
Quota int `json:"quota" gorm:"default:0"`
|
||||
AttemptKey *string `json:"-" gorm:"type:varchar(96);uniqueIndex"`
|
||||
Id int `json:"id"`
|
||||
UserID int `json:"user_id" gorm:"index"`
|
||||
Username string `json:"username" gorm:"index:idx_qdt_model_user_name,priority:2;size:64;default:''"`
|
||||
ModelName string `json:"model_name" gorm:"index:idx_qdt_model_user_name,priority:1;size:64;default:''"`
|
||||
CreatedAt int64 `json:"created_at" gorm:"bigint;index:idx_qdt_created_at,priority:2"`
|
||||
UseGroup string `json:"use_group" gorm:"index;size:64;default:''"`
|
||||
TokenID int `json:"token_id" gorm:"index;default:0"`
|
||||
ChannelID int `json:"channel_id" gorm:"index;default:0"`
|
||||
NodeName string `json:"node_name" gorm:"index;size:64;default:''"`
|
||||
TokenUsed int `json:"token_used" gorm:"default:0"`
|
||||
Count int `json:"count" gorm:"default:0"`
|
||||
Quota int `json:"quota" gorm:"default:0"`
|
||||
}
|
||||
|
||||
type QuotaDataLogParams struct {
|
||||
@@ -107,7 +108,7 @@ func SaveQuotaDataCache() {
|
||||
// 3. 如果没有数据,就插入数据
|
||||
for _, quotaData := range CacheQuotaData {
|
||||
quotaDataDB := &QuotaData{}
|
||||
DB.Table("quota_data").
|
||||
DB.Table("quota_data").Where("attempt_key IS NULL").
|
||||
Where("user_id = ? and username = ? and model_name = ? and created_at = ? and use_group = ? and token_id = ? and channel_id = ? and node_name = ?",
|
||||
quotaData.UserID, quotaData.Username, quotaData.ModelName, quotaData.CreatedAt, quotaData.UseGroup, quotaData.TokenID, quotaData.ChannelID, quotaData.NodeName).
|
||||
First(quotaDataDB)
|
||||
@@ -125,7 +126,7 @@ func SaveQuotaDataCache() {
|
||||
}
|
||||
|
||||
func increaseQuotaData(quotaData *QuotaData) {
|
||||
err := DB.Table("quota_data").
|
||||
err := DB.Table("quota_data").Where("attempt_key IS NULL").
|
||||
Where("user_id = ? and username = ? and model_name = ? and created_at = ? and use_group = ? and token_id = ? and channel_id = ? and node_name = ?",
|
||||
quotaData.UserID, quotaData.Username, quotaData.ModelName, quotaData.CreatedAt, quotaData.UseGroup, quotaData.TokenID, quotaData.ChannelID, quotaData.NodeName).
|
||||
Updates(map[string]any{
|
||||
|
||||
@@ -1106,14 +1106,6 @@ func (user *User) FillUserByGitHubId() error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// UpdateGitHubId updates the user's GitHub ID (used for migration from login to numeric ID)
|
||||
func (user *User) UpdateGitHubId(newGitHubId string) error {
|
||||
if user.Id == 0 {
|
||||
return errors.New("user id is empty")
|
||||
}
|
||||
return DB.Model(user).Update("github_id", newGitHubId).Error
|
||||
}
|
||||
|
||||
func (user *User) FillUserByDiscordId() error {
|
||||
if user.DiscordId == "" {
|
||||
return errors.New("discord id 为空!")
|
||||
|
||||
@@ -10,6 +10,7 @@ import (
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/setting/system_setting"
|
||||
"github.com/go-redis/redis/v8"
|
||||
"github.com/go-webauthn/webauthn/webauthn"
|
||||
"github.com/pquerna/otp/totp"
|
||||
@@ -260,6 +261,22 @@ func TestSecurityFactorMutationsAdvanceUserAuthVersion(t *testing.T) {
|
||||
|
||||
func TestUpdatePasskeyAssertionStateCannotRewriteRegistrationIdentity(t *testing.T) {
|
||||
truncateTables(t)
|
||||
require.NoError(t, DB.AutoMigrate(&Option{}))
|
||||
previousSettings := *system_setting.GetPasskeySettings()
|
||||
domainKeys := map[string]any{"key": []string{"ServerAddress", "passkey.rp_id", "passkey.legacy_rp_ids", "passkey.origins"}}
|
||||
var previousOptions []Option
|
||||
require.NoError(t, DB.Where(domainKeys).Find(&previousOptions).Error)
|
||||
t.Cleanup(func() {
|
||||
*system_setting.GetPasskeySettings() = previousSettings
|
||||
require.NoError(t, DB.Where(domainKeys).Delete(&Option{}).Error)
|
||||
if len(previousOptions) > 0 {
|
||||
require.NoError(t, DB.Create(&previousOptions).Error)
|
||||
}
|
||||
})
|
||||
*system_setting.GetPasskeySettings() = system_setting.PasskeySettings{RPID: "example.com", Origins: "https://example.com"}
|
||||
for _, option := range []Option{{Key: "passkey.rp_id", Value: "example.com"}, {Key: "passkey.legacy_rp_ids", Value: ""}, {Key: "passkey.origins", Value: "https://example.com"}} {
|
||||
require.NoError(t, DB.Save(&option).Error)
|
||||
}
|
||||
|
||||
user := User{Username: "passkey-assertion-state", Password: "password", AuthVersion: 1}
|
||||
require.NoError(t, DB.Create(&user).Error)
|
||||
@@ -292,7 +309,7 @@ func TestUpdatePasskeyAssertionStateCannotRewriteRegistrationIdentity(t *testing
|
||||
CloneWarning: true,
|
||||
},
|
||||
}
|
||||
require.NoError(t, UpdatePasskeyAssertionState(user.Id, validated, usedAt))
|
||||
require.NoError(t, UpdatePasskeyAssertionState(user.Id, validated, usedAt, "example.com"))
|
||||
|
||||
var updated PasskeyCredential
|
||||
require.NoError(t, DB.First(&updated, stored.ID).Error)
|
||||
@@ -312,7 +329,7 @@ func TestUpdatePasskeyAssertionStateCannotRewriteRegistrationIdentity(t *testing
|
||||
assert.Equal(t, usedAt.Unix(), updated.LastUsedAt.Unix())
|
||||
|
||||
validated.ID = []byte("another-credential")
|
||||
assert.ErrorIs(t, UpdatePasskeyAssertionState(user.Id, validated, usedAt), ErrPasskeyNotFound)
|
||||
assert.ErrorIs(t, UpdatePasskeyAssertionState(user.Id, validated, usedAt, "example.com"), ErrPasskeyNotFound)
|
||||
}
|
||||
|
||||
func assertUserAuthVersion(t *testing.T, userID int, expected int64) {
|
||||
|
||||
Reference in new issue
Block a user