diff --git a/common/init.go b/common/init.go
index ff3229aa..2b28a566 100644
--- a/common/init.go
+++ b/common/init.go
@@ -5,6 +5,7 @@ import (
"fmt"
"log"
"net/http"
+ "net/url"
"os"
"path/filepath"
"strconv"
@@ -14,6 +15,11 @@ import (
"github.com/MAX-API-Next/MAX-API/constant"
)
+func SessionCookieSecureForServerAddress(serverAddress string) bool {
+ parsed, err := url.Parse(strings.TrimSpace(serverAddress))
+ return err == nil && strings.EqualFold(parsed.Scheme, "https") && parsed.Host != ""
+}
+
var (
Port = flag.Int("port", 3000, "the listening port")
PrintVersion = flag.Bool("version", false, "print version and exit")
@@ -174,7 +180,8 @@ func initConstantEnv() {
constant.ErrorLogEnabled = GetEnvOrDefaultBool("ERROR_LOG_ENABLED", constant.DefaultErrorLogEnabled)
// 任务轮询时查询的最大数量
constant.TaskQueryLimit = GetEnvOrDefault("TASK_QUERY_LIMIT", constant.DefaultTaskQueryLimit)
- // 异步任务超时时间（分钟），超过此时间未完成的任务将被标记为失败并退款。0 表示禁用。
+ // 异步任务超时时间（分钟）。已确认上游 ID 的任务失败后退款；无法确认
+ // 上游是否接受的提交占位只标记人工核对，避免错误退款。0 表示禁用。
constant.TaskTimeoutMinutes = GetEnvOrDefault("TASK_TIMEOUT_MINUTES", constant.DefaultTaskTimeoutMinutes)
soraPatchStr := GetEnvOrDefaultString("TASK_PRICE_PATCH", "")
diff --git a/common/session_cookie_test.go b/common/session_cookie_test.go
new file mode 100644
index 00000000..6ee555e9
--- /dev/null
+++ b/common/session_cookie_test.go
@@ -0,0 +1,15 @@
+package common
+
+import (
+ "testing"
+
+ "github.com/stretchr/testify/require"
+)
+
+func TestSessionCookieSecureForServerAddress(t *testing.T) {
+ require.True(t, SessionCookieSecureForServerAddress("https://max-api.example"))
+ require.True(t, SessionCookieSecureForServerAddress(" HTTPS://max-api.example/path "))
+ require.False(t, SessionCookieSecureForServerAddress("http://max-api.example"))
+ require.False(t, SessionCookieSecureForServerAddress("localhost:3000"))
+ require.False(t, SessionCookieSecureForServerAddress(""))
+}
diff --git a/controller/channel.go b/controller/channel.go
index ae770a93..345df4b8 100644
--- a/controller/channel.go
+++ b/controller/channel.go
@@ -462,6 +462,9 @@ func validateChannel(channel *model.Channel, isAdd bool) error {
if channel == nil {
return fmt.Errorf("channel cannot be empty")
}
+ if err := service.ValidateStatusCodeMapping(channel.GetStatusCodeMapping()); err != nil {
+ return fmt.Errorf("渠道状态码映射[status code mapping]格式错误：%s", err.Error())
+ }
// 校验 channel settings
if err := channel.ValidateSettings(); err != nil {
return fmt.Errorf("渠道额外设置[channel setting] 格式错误：%s", err.Error())
@@ -1381,6 +1384,21 @@ func ManageMultiKeys(c *gin.Context) {
return
}
+ // Avoid allocating a permanent per-channel lock for arbitrary nonexistent
+ // channel IDs. The mutation still reloads after taking the lock because
+ // multi-key operations rewrite the complete ChannelInfo value.
+ if _, err := model.GetChannelById(request.ChannelId, false); err != nil {
+ c.JSON(http.StatusOK, gin.H{
+ "success": false,
+ "message": "渠道不存在",
+ })
+ return
+ }
+
+ lock := model.GetChannelPollingLock(request.ChannelId)
+ lock.Lock()
+ defer lock.Unlock()
+
channel, err := model.GetChannelById(request.ChannelId, true)
if err != nil {
c.JSON(http.StatusOK, gin.H{
@@ -1408,10 +1426,6 @@ func ManageMultiKeys(c *gin.Context) {
})
}
- lock := model.GetChannelPollingLock(channel.Id)
- lock.Lock()
- defer lock.Unlock()
-
switch request.Action {
case "get_key_status":
keys := channel.GetKeys()
diff --git a/controller/channel_add_test.go b/controller/channel_add_test.go
index 54a1a77c..aea9a244 100644
--- a/controller/channel_add_test.go
+++ b/controller/channel_add_test.go
@@ -26,3 +26,16 @@ func TestBuildChannelsForInsertUsesIndependentNamesWithKeyPrefixes(t *testing.T)
assert.Equal(t, "bravo-key-2", channels[1].Key)
assert.Equal(t, []string{"OpenAI Prod alpha-ke", "OpenAI Prod bravo-ke"}, names)
}
+
+func TestValidateChannelRejectsInvalidStatusCodeMapping(t *testing.T) {
+ invalidMapping := `{"429":999}`
+ channel := &model.Channel{
+ Name: "OpenAI Prod",
+ Type: 1,
+ Key: "key",
+ Models: "gpt-test",
+ StatusCodeMapping: &invalidMapping,
+ }
+
+ require.ErrorContains(t, validateChannel(channel, true), "status code mapping")
+}
diff --git a/controller/channel_multi_key_test.go b/controller/channel_multi_key_test.go
new file mode 100644
index 00000000..81e8ab17
--- /dev/null
+++ b/controller/channel_multi_key_test.go
@@ -0,0 +1,84 @@
+package controller
+
+import (
+ "net/http"
+ "sync/atomic"
+ "testing"
+ "time"
+
+ "github.com/MAX-API-Next/MAX-API/common"
+ "github.com/MAX-API-Next/MAX-API/model"
+ "github.com/stretchr/testify/require"
+ "gorm.io/gorm"
+)
+
+func TestManageMultiKeysValidatesThenReloadsChannelUnderPollingLock(t *testing.T) {
+ db := openTokenControllerTestDB(t)
+ require.NoError(t, db.AutoMigrate(&model.Channel{}, &model.Ability{}, &model.Log{}, &model.User{}))
+
+ previousMemoryCacheEnabled := common.MemoryCacheEnabled
+ common.MemoryCacheEnabled = false
+ t.Cleanup(func() {
+ common.MemoryCacheEnabled = previousMemoryCacheEnabled
+ })
+
+ channel := model.Channel{
+ Name: "multi-key-lock-test",
+ Key: "first-key\nsecond-key",
+ Status: common.ChannelStatusEnabled,
+ Models: "gpt-4o",
+ Group: "default",
+ ChannelInfo: model.ChannelInfo{
+ IsMultiKey: true,
+ MultiKeySize: 2,
+ },
+ }
+ require.NoError(t, db.Create(&channel).Error)
+
+ var channelReads atomic.Int32
+ bothReads := make(chan struct{})
+ require.NoError(t, db.Callback().Query().Before("gorm:query").Register("test:observe-multi-key-channel-load", func(tx *gorm.DB) {
+ if tx.Statement.Table == "channels" && channelReads.Add(1) == 2 {
+ close(bothReads)
+ }
+ }))
+
+ lock := model.GetChannelPollingLock(channel.Id)
+ lock.Lock()
+ results := make(chan tokenAPIResponse, 2)
+ for _, keyIndex := range []int{0, 1} {
+ keyIndex := keyIndex
+ go func() {
+ ctx, recorder := newAuthenticatedContext(t, http.MethodPost, "/api/channel/multi-key", MultiKeyManageRequest{
+ ChannelId: channel.Id,
+ Action: "disable_key",
+ KeyIndex: &keyIndex,
+ }, 1)
+ ManageMultiKeys(ctx)
+ results <- decodeAPIResponse(t, recorder)
+ }()
+ }
+
+ validatedBeforeLock := false
+ select {
+ case <-bothReads:
+ validatedBeforeLock = true
+ case <-time.After(250 * time.Millisecond):
+ }
+ lock.Unlock()
+ require.True(t, validatedBeforeLock, "multi-key handlers must validate channel existence before allocating a polling lock")
+
+ for range 2 {
+ select {
+ case result := <-results:
+ require.True(t, result.Success, "multi-key mutation failed: %s", result.Message)
+ case <-time.After(2 * time.Second):
+ t.Fatal("multi-key mutation did not finish")
+ }
+ }
+
+ var stored model.Channel
+ require.NoError(t, db.First(&stored, channel.Id).Error)
+ require.Equal(t, 2, stored.ChannelInfo.MultiKeyStatusList[0])
+ require.Equal(t, 2, stored.ChannelInfo.MultiKeyStatusList[1])
+}
diff --git a/controller/midjourney.go b/controller/midjourney.go
index 7bdabfdd..3a8e9fb3 100644
--- a/controller/midjourney.go
+++ b/controller/midjourney.go
@@ -271,7 +271,7 @@ func GetAllMidjourney(c *gin.Context) {
if setting.MjForwardUrlEnabled {
for i, midjourney := range items {
- midjourney.ImageUrl = system_setting.ServerAddress + "/mj/image/" + midjourney.MjId
+ midjourney.ImageUrl = service.BuildMidjourneyImageURL(system_setting.ServerAddress, midjourney.MjId, midjourney.UserId, time.Now())
items[i] = midjourney
}
}
@@ -297,7 +297,7 @@ func GetUserMidjourney(c *gin.Context) {
if setting.MjForwardUrlEnabled {
for i, midjourney := range items {
- midjourney.ImageUrl = system_setting.ServerAddress + "/mj/image/" + midjourney.MjId
+ midjourney.ImageUrl = service.BuildMidjourneyImageURL(system_setting.ServerAddress, midjourney.MjId, midjourney.UserId, time.Now())
items[i] = midjourney
}
}
diff --git a/controller/model.go b/controller/model.go
index f709b989..16966bb6 100644
--- a/controller/model.go
+++ b/controller/model.go
@@ -312,12 +312,15 @@ func ListModels(c *gin.Context, modelType int) {
Type: "model",
}
}
- c.JSON(200, gin.H{
+ response := gin.H{
"data": useranthropicModels,
- "first_id": useranthropicModels[0].ID,
"has_more": false,
- "last_id": useranthropicModels[len(useranthropicModels)-1].ID,
- })
+ }
+ if len(useranthropicModels) > 0 {
+ response["first_id"] = useranthropicModels[0].ID
+ response["last_id"] = useranthropicModels[len(useranthropicModels)-1].ID
+ }
+ c.JSON(200, response)
case constant.ChannelTypeGemini:
userGeminiModels := make([]dto.GeminiModel, len(userOpenAiModels))
for i, model := range userOpenAiModels {
diff --git a/controller/model_list_test.go b/controller/model_list_test.go
index f4224d86..916a07d5 100644
--- a/controller/model_list_test.go
+++ b/controller/model_list_test.go
@@ -244,6 +244,27 @@ func TestListModelsTokenLimitIncludesTieredBillingModel(t *testing.T) {
require.NotContains(t, ids, "zz-token-unpriced-model")
}
+func TestListModelsAnthropicAllowsEmptyModelList(t *testing.T) {
+ withSelfUseModeDisabled(t)
+
+ recorder := httptest.NewRecorder()
+ ctx, _ := gin.CreateTestContext(recorder)
+ ctx.Request = httptest.NewRequest(http.MethodGet, "/v1/models", nil)
+ common.SetContextKey(ctx, constant.ContextKeyTokenModelLimitEnabled, true)
+ common.SetContextKey(ctx, constant.ContextKeyTokenModelLimit, map[string]bool{})
+
+ ListModels(ctx, constant.ChannelTypeAnthropic)
+
+ require.Equal(t, http.StatusOK, recorder.Code)
+ var payload struct {
+ Data []dto.AnthropicModel `json:"data"`
+ HasMore bool `json:"has_more"`
+ }
+ require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &payload))
+ require.Empty(t, payload.Data)
+ require.False(t, payload.HasMore)
+}
+
func TestCheckUpdatePasswordRequiresCurrentPassword(t *testing.T) {
db := setupModelListControllerTestDB(t)
hashedPassword, err := common.Password2Hash("CurrentPassword123")
diff --git a/controller/model_meta.go b/controller/model_meta.go
index d18a33b0..93232600 100644
--- a/controller/model_meta.go
+++ b/controller/model_meta.go
@@ -152,7 +152,12 @@ func DeleteModelMeta(c *gin.Context) {
common.ApiError(c, err)
return
}
- if err := model.DB.Delete(&model.Model{}, id).Error; err != nil {
+ var m model.Model
+ if err := model.DB.First(&m, id).Error; err != nil {
+ common.ApiError(c, err)
+ return
+ }
+ if err := m.Delete(); err != nil {
common.ApiError(c, err)
return
}
diff --git a/controller/model_sync.go b/controller/model_sync.go
index 37dcad8e..1a8ac625 100644
--- a/controller/model_sync.go
+++ b/controller/model_sync.go
@@ -61,7 +61,7 @@ type upstreamModel struct {
Icon string `json:"icon"`
ModelName string `json:"model_name"`
NameRule int `json:"name_rule"`
- Status int `json:"status"`
+ Status *int `json:"status"`
Tags string `json:"tags"`
VendorName string `json:"vendor_name"`
}
@@ -70,7 +70,7 @@ type upstreamVendor struct {
Description string `json:"description"`
Icon string `json:"icon"`
Name string `json:"name"`
- Status int `json:"status"`
+ Status *int `json:"status"`
}
var (
@@ -262,6 +262,45 @@ func ensureVendorID(vendorName string, vendorByName map[string]upstreamVendor, v
return 0
}
+func fetchUpstreamMetadata(ctx context.Context, modelsURL, vendorsURL string) (upstreamEnvelope[upstreamVendor], upstreamEnvelope[upstreamModel], error) {
+ var vendorsEnv upstreamEnvelope[upstreamVendor]
+ var modelsEnv upstreamEnvelope[upstreamModel]
+ var vendorsErr error
+ var modelsErr error
+ var wg sync.WaitGroup
+ wg.Add(2)
+ go func() {
+ defer wg.Done()
+ vendorsErr = fetchJSON(ctx, vendorsURL, &vendorsEnv)
+ }()
+ go func() {
+ defer wg.Done()
+ modelsErr = fetchJSON(ctx, modelsURL, &modelsEnv)
+ }()
+ wg.Wait()
+ if vendorsErr != nil {
+ return vendorsEnv, modelsEnv, fmt.Errorf("fetch upstream vendors: %w", vendorsErr)
+ }
+ if !vendorsEnv.Success {
+ message := vendorsEnv.Message
+ if message == "" {
+ message = "upstream response reported failure"
+ }
+ return vendorsEnv, modelsEnv, fmt.Errorf("fetch upstream vendors: %s", message)
+ }
+ if modelsErr != nil {
+ return vendorsEnv, modelsEnv, fmt.Errorf("fetch upstream models: %w", modelsErr)
+ }
+ if !modelsEnv.Success {
+ message := modelsEnv.Message
+ if message == "" {
+ message = "upstream response reported failure"
+ }
+ return vendorsEnv, modelsEnv, fmt.Errorf("fetch upstream models: %s", message)
+ }
+ return vendorsEnv, modelsEnv, nil
+}
+
// SyncUpstreamModels 同步上游模型与供应商：
// - 默认仅创建「未配置模型」
// - 可通过 overwrite 选择性覆盖更新本地已有模型的字段（前提：sync_official <> 0）
@@ -305,25 +344,9 @@ func SyncUpstreamModels(c *gin.Context) {
defer cancel()
modelsURL, vendorsURL := getUpstreamURLs(req.Locale)
- var vendorsEnv upstreamEnvelope[upstreamVendor]
- var modelsEnv upstreamEnvelope[upstreamModel]
- var fetchErr error
- var wg sync.WaitGroup
- wg.Add(2)
- go func() {
- defer wg.Done()
- // vendor 失败不拦截
- _ = fetchJSON(ctx, vendorsURL, &vendorsEnv)
- }()
- go func() {
- defer wg.Done()
- if err := fetchJSON(ctx, modelsURL, &modelsEnv); err != nil {
- fetchErr = err
- }
- }()
- wg.Wait()
+ vendorsEnv, modelsEnv, fetchErr := fetchUpstreamMetadata(ctx, modelsURL, vendorsURL)
if fetchErr != nil {
- c.JSON(http.StatusOK, gin.H{"success": false, "message": "获取上游模型失败: " + fetchErr.Error(), "locale": req.Locale, "source_urls": gin.H{"models_url": modelsURL, "vendors_url": vendorsURL}})
+ c.JSON(http.StatusOK, gin.H{"success": false, "message": "获取上游元数据失败: " + fetchErr.Error(), "locale": req.Locale, "source_urls": gin.H{"models_url": modelsURL, "vendors_url": vendorsURL}})
return
}
@@ -485,14 +508,11 @@ func coalesce(a, b string) string {
return b
}
-func chooseStatus(primary, fallback int) int {
- if primary == 0 && fallback != 0 {
- return fallback
+func chooseStatus(primary *int, fallback int) int {
+ if primary != nil {
+ return *primary
}
- if primary != 0 {
- return primary
- }
- return 1
+ return fallback
}
// SyncUpstreamPreview 预览上游与本地的差异（仅用于弹窗选择）
@@ -505,24 +525,9 @@ func SyncUpstreamPreview(c *gin.Context) {
locale := c.Query("locale")
modelsURL, vendorsURL := getUpstreamURLs(locale)
- var vendorsEnv upstreamEnvelope[upstreamVendor]
- var modelsEnv upstreamEnvelope[upstreamModel]
- var fetchErr error
- var wg sync.WaitGroup
- wg.Add(2)
- go func() {
- defer wg.Done()
- _ = fetchJSON(ctx, vendorsURL, &vendorsEnv)
- }()
- go func() {
- defer wg.Done()
- if err := fetchJSON(ctx, modelsURL, &modelsEnv); err != nil {
- fetchErr = err
- }
- }()
- wg.Wait()
+ vendorsEnv, modelsEnv, fetchErr := fetchUpstreamMetadata(ctx, modelsURL, vendorsURL)
if fetchErr != nil {
- c.JSON(http.StatusOK, gin.H{"success": false, "message": "获取上游模型失败: " + fetchErr.Error(), "locale": locale, "source_urls": gin.H{"models_url": modelsURL, "vendors_url": vendorsURL}})
+ c.JSON(http.StatusOK, gin.H{"success": false, "message": "获取上游元数据失败: " + fetchErr.Error(), "locale": locale, "source_urls": gin.H{"models_url": modelsURL, "vendors_url": vendorsURL}})
return
}
diff --git a/controller/model_sync_test.go b/controller/model_sync_test.go
new file mode 100644
index 00000000..449dd2ce
--- /dev/null
+++ b/controller/model_sync_test.go
@@ -0,0 +1,101 @@
+package controller
+
+import (
+ "context"
+ "net/http"
+ "net/http/httptest"
+ "testing"
+
+ "github.com/stretchr/testify/require"
+)
+
+func TestFetchUpstreamMetadataFailsWhenVendorsAreUnavailable(t *testing.T) {
+ t.Setenv("SYNC_HTTP_RETRY", "1")
+ server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ switch r.URL.Path {
+ case "/models":
+ w.Header().Set("Content-Type", "application/json")
+ _, _ = w.Write([]byte(`[]`))
+ case "/vendors":
+ w.Header().Set("Content-Type", "application/json")
+ _, _ = w.Write([]byte(`{"success":false,"message":"temporarily unavailable"}`))
+ default:
+ http.NotFound(w, r)
+ }
+ }))
+ defer server.Close()
+
+ _, _, err := fetchUpstreamMetadata(context.Background(), server.URL+"/models", server.URL+"/vendors")
+ require.Error(t, err)
+ require.ErrorContains(t, err, "fetch upstream vendors")
+}
+
+func TestFetchUpstreamMetadataFailsWhenModelsAreUnavailable(t *testing.T) {
+ t.Setenv("SYNC_HTTP_RETRY", "1")
+ server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ w.Header().Set("Content-Type", "application/json")
+ switch r.URL.Path {
+ case "/models":
+ _, _ = w.Write([]byte(`{"success":false,"message":"models unavailable"}`))
+ case "/vendors":
+ _, _ = w.Write([]byte(`{"success":true,"data":[{"name":"Example Vendor"}]}`))
+ default:
+ http.NotFound(w, r)
+ }
+ }))
+ defer server.Close()
+
+ _, _, err := fetchUpstreamMetadata(context.Background(), server.URL+"/models", server.URL+"/vendors")
+ require.Error(t, err)
+ require.ErrorContains(t, err, "fetch upstream models")
+ require.ErrorContains(t, err, "models unavailable")
+}
+
+func TestFetchUpstreamMetadataFailsWhenModelsTransportFails(t *testing.T) {
+ t.Setenv("SYNC_HTTP_RETRY", "1")
+ server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ w.Header().Set("Content-Type", "application/json")
+ if r.URL.Path == "/vendors" {
+ _, _ = w.Write([]byte(`{"success":true,"data":[{"name":"Example Vendor"}]}`))
+ return
+ }
+ http.Error(w, "not found", http.StatusNotFound)
+ }))
+ defer server.Close()
+
+ _, _, err := fetchUpstreamMetadata(context.Background(), server.URL+"/models", server.URL+"/vendors")
+ require.Error(t, err)
+ require.ErrorContains(t, err, "fetch upstream models")
+ require.ErrorContains(t, err, "404 Not Found")
+}
+
+func TestFetchUpstreamMetadataAcceptsArrayResponses(t *testing.T) {
+ t.Setenv("SYNC_HTTP_RETRY", "1")
+ server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ w.Header().Set("Content-Type", "application/json")
+ switch r.URL.Path {
+ case "/models":
+ _, _ = w.Write([]byte(`[{"model_name":"example-model","vendor_name":"Example Vendor"}]`))
+ case "/vendors":
+ _, _ = w.Write([]byte(`[{"name":"Example Vendor"}]`))
+ default:
+ http.NotFound(w, r)
+ }
+ }))
+ defer server.Close()
+
+ vendors, models, err := fetchUpstreamMetadata(context.Background(), server.URL+"/models", server.URL+"/vendors")
+ require.NoError(t, err)
+ require.True(t, vendors.Success)
+ require.True(t, models.Success)
+ require.Equal(t, "Example Vendor", vendors.Data[0].Name)
+ require.Equal(t, "example-model", models.Data[0].ModelName)
+}
+
+func TestChooseStatusPreservesExplicitDisabledValue(t *testing.T) {
+ disabled := 0
+ enabled := 1
+
+ require.Equal(t, 0, chooseStatus(&disabled, enabled))
+ require.Equal(t, enabled, chooseStatus(nil, enabled))
+}
diff --git a/controller/oauth.go b/controller/oauth.go
index 17e63dbc..0d81e1ea 100644
--- a/controller/oauth.go
+++ b/controller/oauth.go
@@ -291,7 +291,7 @@ func findOrCreateOAuthUser(c *gin.Context, provider oauth.Provider, oauthUser *o
if errors.Is(err, model.ErrUserDeleted) {
return nil, &OAuthUserDeletedError{}
}
- return nil, err
+ return nil, &oauthIdentityLookupError{provider: provider.GetName(), err: err}
}
// Check if user has been deleted
if user.Id == 0 {
@@ -312,7 +312,7 @@ func findOrCreateOAuthUser(c *gin.Context, provider oauth.Provider, oauthUser *o
if errors.Is(err, model.ErrUserDeleted) {
return nil, &OAuthUserDeletedError{}
}
- return nil, err
+ return nil, &oauthIdentityLookupError{provider: provider.GetName(), err: err}
}
if user.Id != 0 {
// Found user with legacy ID, migrate to new ID
diff --git a/controller/oauth_test.go b/controller/oauth_test.go
index 6e11377e..26fd1502 100644
--- a/controller/oauth_test.go
+++ b/controller/oauth_test.go
@@ -96,6 +96,29 @@ type publicLookupFailingOAuthProvider struct {
deletedUserOAuthProvider
}
+type publicFillFailingOAuthProvider struct {
+ publicLookupFailingOAuthProvider
+}
+
+func (*publicFillFailingOAuthProvider) IsUserIDTaken(string) (bool, error) { return true, nil }
+func (*publicFillFailingOAuthProvider) FillUserByProviderID(*model.User, string) error {
+ return errors.New("internal database host db.internal:5432")
+}
+
+func TestFindOrCreateOAuthUserWrapsProviderFillDatabaseError(t *testing.T) {
+ _, err := findOrCreateOAuthUser(nil, &publicFillFailingOAuthProvider{}, &oauth.OAuthUser{
+ ProviderUserID: "existing-user-id",
+ Extra: map[string]any{},
+ }, nil)
+ var lookupErr *oauthIdentityLookupError
+ if !errors.As(err, &lookupErr) {
+ t.Fatalf("expected masked identity lookup error, got %v", err)
+ }
+ if strings.Contains(err.Error(), "db.internal") {
+ t.Fatalf("public error leaked database details: %q", err.Error())
+ }
+}
+
func (*publicLookupFailingOAuthProvider) ExchangeToken(context.Context, string, *gin.Context) (*oauth.OAuthToken, error) {
return &oauth.OAuthToken{AccessToken: "test-token"}, nil
}
diff --git a/controller/option.go b/controller/option.go
index ddc2f582..dd2bee96 100644
--- a/controller/option.go
+++ b/controller/option.go
@@ -128,6 +128,13 @@ func UpdateOption(c *gin.Context) {
})
return
}
+ if !model.IsRegisteredOptionKey(option.Key) {
+ c.JSON(http.StatusBadRequest, gin.H{
+ "success": false,
+ "message": "不支持的设置项",
+ })
+ return
+ }
switch option.Value.(type) {
case bool:
option.Value = common.Interface2String(option.Value.(bool))
diff --git a/controller/option_test.go b/controller/option_test.go
new file mode 100644
index 00000000..cf198a49
--- /dev/null
+++ b/controller/option_test.go
@@ -0,0 +1,34 @@
+package controller
+
+import (
+ "net/http"
+ "net/http/httptest"
+ "strings"
+ "testing"
+
+ "github.com/MAX-API-Next/MAX-API/common"
+ "github.com/gin-gonic/gin"
+ "github.com/stretchr/testify/require"
+)
+
+func TestUpdateOptionRejectsUnregisteredKey(t *testing.T) {
+ gin.SetMode(gin.TestMode)
+ common.OptionMapRWMutex.Lock()
+ originalOptionMap := common.OptionMap
+ common.OptionMap = map[string]string{"SystemName": "MAX API"}
+ common.OptionMapRWMutex.Unlock()
+ t.Cleanup(func() {
+ common.OptionMapRWMutex.Lock()
+ common.OptionMap = originalOptionMap
+ common.OptionMapRWMutex.Unlock()
+ })
+
+ recorder := httptest.NewRecorder()
+ ctx, _ := gin.CreateTestContext(recorder)
+ ctx.Request = httptest.NewRequest(http.MethodPut, "/api/option/", strings.NewReader(`{"key":"unexpected.option","value":"x"}`))
+ ctx.Request.Header.Set("Content-Type", "application/json")
+
+ UpdateOption(ctx)
+
+ require.Equal(t, http.StatusBadRequest, recorder.Code)
+}
diff --git a/controller/relay.go b/controller/relay.go
index 390f2e52..b6eb5735 100644
--- a/controller/relay.go
+++ b/controller/relay.go
@@ -171,10 +171,7 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) {
// Only return quota if downstream failed and quota was actually pre-consumed
if maxAPIError != nil {
maxAPIError = service.NormalizeViolationFeeError(maxAPIError)
- if relayInfo.Billing != nil {
- relayInfo.Billing.Refund(c)
- }
- service.ChargeViolationFeeIfNeeded(c, relayInfo, maxAPIError)
+ service.HandleFailedBilling(c, relayInfo, maxAPIError)
}
}()
@@ -234,6 +231,7 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) {
if !shouldRetry(c, maxAPIError, common.RetryTimes-retryParam.GetRetry()) {
break
}
+ retryParam.ExcludeChannel(channel.Id)
}
useChannel := c.GetStringSlice("use_channel")
@@ -505,9 +503,37 @@ func RelayTask(c *gin.Context) {
var result *relay.TaskSubmitResult
var taskErr *dto.TaskError
+ var settlementEffectOperationKey string
defer func() {
- if taskErr != nil && relayInfo.Billing != nil {
- relayInfo.Billing.Refund(c)
+ if taskErr != nil {
+ upstreamAmbiguous := relayInfo.UpstreamTaskResponseReceived || relayInfo.UpstreamTaskOutcomeUnknown
+ upstreamPersisted := result != nil && result.Task != nil
+ if relayInfo.Billing != nil && !upstreamAmbiguous {
+ relayInfo.Billing.Refund(c)
+ }
+ if relayInfo.PersistedTaskID > 0 {
+ reason := taskErr.Message
+ if reason == "" && taskErr.Error != nil {
+ reason = taskErr.Error.Error()
+ }
+ if upstreamPersisted {
+ reason += "（上游已接受任务，本地结算记录未完成；已保留预扣费，请人工核对）"
+ if markErr := model.MarkTaskSubmitNeedsReview(result.Task, reason); markErr != nil {
+ common.SysError("mark task submit for manual review error: " + markErr.Error())
+ }
+ } else if upstreamAmbiguous {
+ if relayInfo.UpstreamTaskOutcomeUnknown {
+ reason += "（上游请求已发送但响应结果未知；已保留预扣费，请人工核对）"
+ } else {
+ reason += "（已收到上游成功响应但无法确认任务 ID；已保留预扣费，请人工核对）"
+ }
+ if markErr := model.MarkTaskSubmitAmbiguous(relayInfo.PersistedTaskID, reason); markErr != nil {
+ common.SysError("mark ambiguous task submit error: " + markErr.Error())
+ }
+ } else if markErr := model.MarkTaskSubmitFailed(relayInfo.PersistedTaskID, reason); markErr != nil {
+ common.SysError("mark failed task submit error: " + markErr.Error())
+ }
+ }
}
}()
@@ -556,6 +582,9 @@ func RelayTask(c *gin.Context) {
if taskErr == nil {
break
}
+ if relayInfo.UpstreamTaskResponseReceived || relayInfo.UpstreamTaskOutcomeUnknown {
+ break
+ }
if !taskErr.LocalError {
processChannelError(c,
@@ -567,6 +596,9 @@ func RelayTask(c *gin.Context) {
if !shouldRetryTaskRelay(c, channel.Id, taskErr, common.RetryTimes-retryParam.GetRetry()) {
break
}
+ if relayInfo.LockedChannel == nil {
+ retryParam.ExcludeChannel(channel.Id)
+ }
}
useChannel := c.GetStringSlice("use_channel")
@@ -575,38 +607,44 @@ func RelayTask(c *gin.Context) {
logger.LogInfo(c, retryLogStr)
}
- // ── 成功：结算 + 日志 + 插入任务 ──
+ // ── 成功：原子持久化任务 + 结算 intent，再释放缓冲响应 ──
if taskErr == nil {
- if settleErr := service.SettleBilling(c, relayInfo, result.Quota); settleErr != nil {
- common.SysError("settle task billing error: " + settleErr.Error())
- }
- service.LogTaskConsumption(c, relayInfo)
-
- deltaSettlementDisabled := false
- if relayInfo.ChannelMeta != nil && relayInfo.ChannelType == constant.ChannelTypeDoubaoVideo {
- deltaSettlementDisabled = relayInfo.ChannelOtherSettings.DisableTaskDeltaSettlement
- }
- task := model.InitTask(result.Platform, relayInfo)
- task.PrivateData.UpstreamTaskID = result.UpstreamTaskID
- task.PrivateData.BillingSource = relayInfo.BillingSource
- task.PrivateData.SubscriptionId = relayInfo.SubscriptionId
- task.PrivateData.TokenId = relayInfo.TokenId
- task.PrivateData.NodeName = common.NodeName
- task.PrivateData.BillingContext = &model.TaskBillingContext{
- ModelPrice: relayInfo.PriceData.ModelPrice,
- GroupRatio: relayInfo.PriceData.GroupRatioInfo.GroupRatio,
- ModelRatio: relayInfo.PriceData.ModelRatio,
- OtherRatios: relayInfo.PriceData.OtherRatios,
- TaskBilling: relayInfo.TaskBilling,
- OriginModelName: relayInfo.OriginModelName,
- PerCallBilling: common.StringsContains(constant.TaskPricePatches, relayInfo.OriginModelName) || relayInfo.PriceData.UsePrice || relayInfo.TaskBilling != nil,
- DeltaSettlementDisabled: common.GetPointer(deltaSettlementDisabled),
+ if result == nil || result.Task == nil {
+ taskErr = service.TaskErrorWrapperLocal(errors.New("task submit result is incomplete"), "persist_task_failed", http.StatusInternalServerError)
+ } else {
+ var settlementIntent *model.BillingSettlementInput
+ if relayInfo.Billing != nil {
+ preparer, ok := relayInfo.Billing.(service.SettlementPreparer)
+ if !ok {
+ taskErr = service.TaskErrorWrapperLocal(errors.New("task billing session cannot prepare a durable settlement"), "persist_billing_intent_failed", http.StatusInternalServerError)
+ } else if settlementIntent, err = preparer.PrepareSettlement(result.Quota); err != nil {
+ taskErr = service.TaskErrorWrapperLocal(err, "persist_billing_intent_failed", http.StatusInternalServerError)
+ }
+ }
+ if taskErr == nil {
+ if settlementIntent != nil {
+ settlementIntent.Effect = service.BuildTaskSubmissionSettlementEffect(c, relayInfo, result.Quota)
+ settlementEffectOperationKey = settlementIntent.OperationKey
+ }
+ if err = result.Task.UpdateWithSettlementIntent(settlementIntent); err != nil {
+ taskErr = relay.TaskPersistenceError(err, "persist_task_result_failed", "failed to persist task result")
+ }
+ }
}
- task.Quota = result.Quota
- task.Data = result.TaskData
- task.Action = relayInfo.Action
- if insertErr := task.Insert(); insertErr != nil {
- common.SysError("insert task error: " + insertErr.Error())
+ if taskErr == nil {
+ if settleErr := service.SettleBilling(c, relayInfo, result.Quota); settleErr != nil {
+ common.SysError("settle task billing remains pending/manual: " + settleErr.Error())
+ }
+ if settlementEffectOperationKey != "" {
+ if effectErr := model.ProcessBillingSettlementEffect(settlementEffectOperationKey); effectErr != nil {
+ common.SysError("task billing effect remains pending/manual: " + effectErr.Error())
+ }
+ } else {
+ service.LogTaskConsumption(c, relayInfo)
+ }
+ if writeErr := result.WriteResponse(c); writeErr != nil {
+ common.SysError("write task response error: " + writeErr.Error())
+ }
}
}
@@ -639,27 +677,34 @@ func shouldRetryTaskRelay(c *gin.Context, channelId int, taskErr *dto.TaskError,
if taskErr.LocalError {
return false
}
- if taskErr.StatusCode == http.StatusTooManyRequests {
+ if taskErr.UpstreamStatusCode != 0 && !shouldRetryTaskStatusCode(taskErr.UpstreamStatusCode) {
+ return false
+ }
+ return shouldRetryTaskStatusCode(taskErr.StatusCode)
+}
+
+func shouldRetryTaskStatusCode(statusCode int) bool {
+ if statusCode == http.StatusTooManyRequests {
return true
}
- if taskErr.StatusCode == 307 {
+ if statusCode == 307 {
return true
}
- if taskErr.StatusCode/100 == 5 {
+ if statusCode/100 == 5 {
// 超时不重试
- if operation_setting.IsAlwaysSkipRetryStatusCode(taskErr.StatusCode) {
+ if operation_setting.IsAlwaysSkipRetryStatusCode(statusCode) {
return false
}
return true
}
- if taskErr.StatusCode == http.StatusBadRequest {
+ if statusCode == http.StatusBadRequest {
return false
}
- if taskErr.StatusCode == 408 {
+ if statusCode == 408 {
// azure处理超时不重试
return false
}
- if taskErr.StatusCode/100 == 2 {
+ if statusCode/100 == 2 {
return false
}
return true
diff --git a/controller/relay_retry_test.go b/controller/relay_retry_test.go
index 5e0eb901..94e8c39d 100644
--- a/controller/relay_retry_test.go
+++ b/controller/relay_retry_test.go
@@ -96,3 +96,23 @@ func TestShouldRetryTaskRelaySkipsLocalServerError(t *testing.T) {
LocalError: true,
}, 1))
}
+
+func TestShouldRetryTaskRelayDoesNotEnableRetryThroughStatusMapping(t *testing.T) {
+ gin.SetMode(gin.TestMode)
+ c, _ := gin.CreateTestContext(httptest.NewRecorder())
+
+ require.False(t, shouldRetryTaskRelay(c, 1, &dto.TaskError{
+ StatusCode: http.StatusServiceUnavailable,
+ UpstreamStatusCode: http.StatusBadRequest,
+ }, 1))
+}
+
+func TestShouldRetryTaskRelayAllowsStatusMappingToSuppressRetry(t *testing.T) {
+ gin.SetMode(gin.TestMode)
+ c, _ := gin.CreateTestContext(httptest.NewRecorder())
+
+ require.False(t, shouldRetryTaskRelay(c, 1, &dto.TaskError{
+ StatusCode: http.StatusBadRequest,
+ UpstreamStatusCode: http.StatusServiceUnavailable,
+ }, 1))
+}
diff --git a/controller/secure_verification.go b/controller/secure_verification.go
index 64e1e612..ffbb2952 100644
--- a/controller/secure_verification.go
+++ b/controller/secure_verification.go
@@ -14,14 +14,15 @@ import (
const (
// SecureVerificationSessionKey means the user has fully passed secure verification.
- SecureVerificationSessionKey = "secure_verified_at"
- secureVerificationMethodSessionKey = "secure_verified_method"
- secureVerificationUserSessionKey = "secure_verified_user_id"
- secureVerificationScopeSessionKey = "secure_verified_scope"
- secureVerificationMethod2FA = "2fa"
- secureVerificationMethodPasskey = "passkey"
- secureVerificationMethodPassword = "password"
- secureVerificationScopeAccessToken = "access_token"
+ SecureVerificationSessionKey = "secure_verified_at"
+ secureVerificationMethodSessionKey = "secure_verified_method"
+ secureVerificationUserSessionKey = "secure_verified_user_id"
+ secureVerificationScopeSessionKey = "secure_verified_scope"
+ secureVerificationMethod2FA = "2fa"
+ secureVerificationMethodPasskey = "passkey"
+ secureVerificationMethodPassword = "password"
+ secureVerificationScopeAccessToken = "access_token"
+ secureVerificationScopeAccountDelete = "account_delete"
// PasskeyReadySessionKey means WebAuthn finished and /api/verify can finalize step-up verification.
PasskeyReadySessionKey = "secure_passkey_ready_at"
// SecureVerificationTimeout 验证有效期（秒）
@@ -94,7 +95,12 @@ func GetVerificationMethods(c *gin.Context) {
return
}
- methods, err := loadSecureVerificationMethods(user, c.Query("scope") == secureVerificationScopeAccessToken)
+ scope := c.Query("scope")
+ if !isSupportedSecureVerificationScope(scope) {
+ common.ApiError(c, fmt.Errorf("不支持的验证作用域: %s", scope))
+ return
+ }
+ methods, err := loadSecureVerificationMethods(user, passwordVerificationAllowed(scope))
if err != nil {
common.ApiError(c, err)
return
@@ -140,7 +146,11 @@ func UniversalVerify(c *gin.Context) {
return
}
- methods, err := loadSecureVerificationMethods(user, req.Scope == secureVerificationScopeAccessToken)
+ if !isSupportedSecureVerificationScope(req.Scope) {
+ common.ApiError(c, fmt.Errorf("不支持的验证作用域: %s", req.Scope))
+ return
+ }
+ methods, err := loadSecureVerificationMethods(user, passwordVerificationAllowed(req.Scope))
if err != nil {
common.ApiError(c, err)
return
@@ -226,6 +236,19 @@ func UniversalVerify(c *gin.Context) {
})
}
+func isSupportedSecureVerificationScope(scope string) bool {
+ switch scope {
+ case "", secureVerificationScopeAccessToken, secureVerificationScopeAccountDelete:
+ return true
+ default:
+ return false
+ }
+}
+
+func passwordVerificationAllowed(scope string) bool {
+ return scope == secureVerificationScopeAccessToken || scope == secureVerificationScopeAccountDelete
+}
+
func setSecureVerificationSession(c *gin.Context, userId int, method string, scope string) (int64, error) {
session := sessions.Default(c)
session.Delete(PasskeyReadySessionKey)
diff --git a/controller/secure_verification_test.go b/controller/secure_verification_test.go
index 1ad98d32..c9446839 100644
--- a/controller/secure_verification_test.go
+++ b/controller/secure_verification_test.go
@@ -112,6 +112,15 @@ func TestUniversalVerifyPasswordRequiresAccessTokenScope(t *testing.T) {
require.Contains(t, recorder.Body.String(), "密码验证不可用")
}
+func TestSecureVerificationScopeAllowlist(t *testing.T) {
+ require.True(t, isSupportedSecureVerificationScope(""))
+ require.True(t, isSupportedSecureVerificationScope(secureVerificationScopeAccessToken))
+ require.True(t, isSupportedSecureVerificationScope(secureVerificationScopeAccountDelete))
+ require.False(t, isSupportedSecureVerificationScope("admin_delete"))
+ require.True(t, passwordVerificationAllowed(secureVerificationScopeAccountDelete))
+ require.False(t, passwordVerificationAllowed(""))
+}
+
func TestVerificationMethodsUsePasswordOnlyAsFallback(t *testing.T) {
db := setupUserSettingControllerTestDB(t)
require.NoError(t, db.AutoMigrate(
diff --git a/controller/token.go b/controller/token.go
index 04708978..215f7413 100644
--- a/controller/token.go
+++ b/controller/token.go
@@ -260,6 +260,11 @@ func UpdateToken(c *gin.Context) {
common.ApiError(c, err)
return
}
+ // Full updates omit status with its zero value. Status-only updates must persist a valid status.
+ if (statusOnly != "" || token.Status != 0) && !isValidTokenStatus(token.Status) {
+ common.ApiErrorI18n(c, i18n.MsgInvalidParams)
+ return
+ }
if len(token.Name) > 50 {
common.ApiErrorI18n(c, i18n.MsgTokenNameTooLong)
return
@@ -281,11 +286,19 @@ func UpdateToken(c *gin.Context) {
return
}
if token.Status == common.TokenStatusEnabled {
- if cleanToken.Status == common.TokenStatusExpired && cleanToken.ExpiredTime <= common.GetTimestamp() && cleanToken.ExpiredTime != -1 {
+ expiredTime := cleanToken.ExpiredTime
+ remainQuota := cleanToken.RemainQuota
+ unlimitedQuota := cleanToken.UnlimitedQuota
+ if statusOnly == "" {
+ expiredTime = token.ExpiredTime
+ remainQuota = token.RemainQuota
+ unlimitedQuota = token.UnlimitedQuota
+ }
+ if cleanToken.Status == common.TokenStatusExpired && expiredTime <= common.GetTimestamp() && expiredTime != -1 {
common.ApiErrorI18n(c, i18n.MsgTokenExpiredCannotEnable)
return
}
- if cleanToken.Status == common.TokenStatusExhausted && cleanToken.RemainQuota <= 0 && !cleanToken.UnlimitedQuota {
+ if cleanToken.Status == common.TokenStatusExhausted && remainQuota <= 0 && !unlimitedQuota {
common.ApiErrorI18n(c, i18n.MsgTokenExhaustedCannotEable)
return
}
@@ -307,6 +320,9 @@ func UpdateToken(c *gin.Context) {
cleanToken.AllowIps = token.AllowIps
cleanToken.Group = token.Group
cleanToken.CrossGroupRetry = token.CrossGroupRetry
+ if token.Status != 0 {
+ cleanToken.Status = token.Status
+ }
}
err = cleanToken.Update()
if err != nil {
@@ -320,6 +336,15 @@ func UpdateToken(c *gin.Context) {
})
}
+func isValidTokenStatus(status int) bool {
+ switch status {
+ case common.TokenStatusEnabled, common.TokenStatusDisabled, common.TokenStatusExpired, common.TokenStatusExhausted:
+ return true
+ default:
+ return false
+ }
+}
+
type TokenBatch struct {
Ids []int `json:"ids"`
}
diff --git a/controller/token_test.go b/controller/token_test.go
index c4314acd..fc45b7f2 100644
--- a/controller/token_test.go
+++ b/controller/token_test.go
@@ -637,6 +637,101 @@ func TestUpdateTokenMasksKeyInResponse(t *testing.T) {
}
}
+func TestUpdateTokenCanRestoreExpiredTokenWithNewValues(t *testing.T) {
+ db := setupTokenControllerTestDB(t)
+ token := seedToken(t, db, 1, "expired-token", "expired1234token5678")
+ token.Status = common.TokenStatusExpired
+ token.ExpiredTime = common.GetTimestamp() - 1
+ token.RemainQuota = 0
+ token.UnlimitedQuota = false
+ if err := db.Save(token).Error; err != nil {
+ t.Fatalf("failed to expire token: %v", err)
+ }
+
+ newExpiry := common.GetTimestamp() + 3600
+ body := map[string]any{
+ "id": token.Id,
+ "name": "restored-token",
+ "status": common.TokenStatusEnabled,
+ "expired_time": newExpiry,
+ "remain_quota": 100,
+ "unlimited_quota": false,
+ "model_limits_enabled": false,
+ "model_limits": "",
+ "group": "default",
+ "cross_group_retry": false,
+ }
+
+ ctx, recorder := newAuthenticatedContext(t, http.MethodPut, "/api/token/", body, 1)
+ UpdateToken(ctx)
+
+ response := decodeAPIResponse(t, recorder)
+ if !response.Success {
+ t.Fatalf("expected one-step token restore to succeed, got message: %s", response.Message)
+ }
+
+ var stored model.Token
+ if err := db.First(&stored, token.Id).Error; err != nil {
+ t.Fatalf("failed to reload restored token: %v", err)
+ }
+ if stored.Status != common.TokenStatusEnabled {
+ t.Fatalf("expected restored token status %d, got %d", common.TokenStatusEnabled, stored.Status)
+ }
+ if stored.ExpiredTime != newExpiry || stored.RemainQuota != 100 {
+ t.Fatalf("expected restored values expiry=%d quota=100, got expiry=%d quota=%d", newExpiry, stored.ExpiredTime, stored.RemainQuota)
+ }
+}
+
+func TestUpdateTokenRejectsInvalidStatus(t *testing.T) {
+ tests := []struct {
+ name string
+ target string
+ status int
+ statusOnly bool
+ }{
+ {name: "full update unsupported status", target: "/api/token/", status: 99},
+ {name: "status only zero status", target: "/api/token/?status_only=true", status: 0, statusOnly: true},
+ {name: "status only unsupported status", target: "/api/token/?status_only=true", status: 99, statusOnly: true},
+ }
+
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ db := setupTokenControllerTestDB(t)
+ token := seedToken(t, db, 1, "editable-token", "invalidstatus123456")
+ body := map[string]any{
+ "id": token.Id,
+ "status": tt.status,
+ }
+ if !tt.statusOnly {
+ body["name"] = "updated-token"
+ body["expired_time"] = -1
+ body["remain_quota"] = 100
+ body["unlimited_quota"] = true
+ body["model_limits_enabled"] = false
+ body["model_limits"] = ""
+ body["group"] = "default"
+ body["cross_group_retry"] = false
+ }
+
+ ctx, recorder := newAuthenticatedContext(t, http.MethodPut, tt.target, body, 1)
+ UpdateToken(ctx)
+
+ response := decodeAPIResponse(t, recorder)
+ if response.Success {
+ t.Fatal("expected invalid token status update to fail")
+ }
+
+ var stored model.Token
+ if err := db.First(&stored, token.Id).Error; err != nil {
+ t.Fatalf("failed to reload token: %v", err)
+ }
+ if stored.Status != common.TokenStatusEnabled {
+ t.Fatalf("expected status to remain %d, got %d", common.TokenStatusEnabled, stored.Status)
+ }
+ })
+ }
+}
+
func TestUpdateTokenAllowsRuntimeAutoRoute(t *testing.T) {
db := setupTokenControllerTestDB(t)
setupHiddenAutoRouteForTokenTest(t)
diff --git a/controller/topup_stripe.go b/controller/topup_stripe.go
index f8dd23e8..2c66a371 100644
--- a/controller/topup_stripe.go
+++ b/controller/topup_stripe.go
@@ -173,48 +173,54 @@ func StripeWebhook(c *gin.Context) {
callerIp := c.ClientIP()
logger.LogInfo(ctx, fmt.Sprintf("Stripe webhook 验签成功 event_type=%s client_ip=%s path=%q", string(event.Type), callerIp, c.Request.RequestURI))
+ var handleErr error
switch event.Type {
case stripe.EventTypeCheckoutSessionCompleted:
- sessionCompleted(ctx, event, callerIp)
+ handleErr = sessionCompleted(ctx, event, callerIp)
case stripe.EventTypeCheckoutSessionExpired:
sessionExpired(ctx, event)
case stripe.EventTypeCheckoutSessionAsyncPaymentSucceeded:
- sessionAsyncPaymentSucceeded(ctx, event, callerIp)
+ handleErr = sessionAsyncPaymentSucceeded(ctx, event, callerIp)
case stripe.EventTypeCheckoutSessionAsyncPaymentFailed:
sessionAsyncPaymentFailed(ctx, event, callerIp)
default:
logger.LogInfo(ctx, fmt.Sprintf("Stripe webhook 忽略事件 event_type=%s client_ip=%s", string(event.Type), callerIp))
}
+ if handleErr != nil {
+ logger.LogError(ctx, fmt.Sprintf("Stripe webhook 处理失败，等待重试 event_type=%s client_ip=%s error=%q", string(event.Type), callerIp, handleErr.Error()))
+ c.AbortWithStatus(http.StatusInternalServerError)
+ return
+ }
c.Status(http.StatusOK)
}
-func sessionCompleted(ctx context.Context, event stripe.Event, callerIp string) {
+func sessionCompleted(ctx context.Context, event stripe.Event, callerIp string) error {
customerId := event.GetObjectValue("customer")
referenceId := event.GetObjectValue("client_reference_id")
status := event.GetObjectValue("status")
if "complete" != status {
logger.LogWarn(ctx, fmt.Sprintf("Stripe checkout.completed 状态异常，忽略处理 trade_no=%s status=%s client_ip=%s", referenceId, status, callerIp))
- return
+ return nil
}
paymentStatus := event.GetObjectValue("payment_status")
if paymentStatus != "paid" {
logger.LogInfo(ctx, fmt.Sprintf("Stripe Checkout 支付未完成，等待异步结果 trade_no=%s payment_status=%s client_ip=%s", referenceId, paymentStatus, callerIp))
- return
+ return nil
}
- fulfillOrder(ctx, event, referenceId, customerId, callerIp)
+ return fulfillOrder(ctx, event, referenceId, customerId, callerIp)
}
// sessionAsyncPaymentSucceeded handles delayed payment methods (bank transfer, SEPA, etc.)
// that confirm payment after the checkout session completes.
-func sessionAsyncPaymentSucceeded(ctx context.Context, event stripe.Event, callerIp string) {
+func sessionAsyncPaymentSucceeded(ctx context.Context, event stripe.Event, callerIp string) error {
customerId := event.GetObjectValue("customer")
referenceId := event.GetObjectValue("client_reference_id")
logger.LogInfo(ctx, fmt.Sprintf("Stripe 异步支付成功 trade_no=%s client_ip=%s", referenceId, callerIp))
- fulfillOrder(ctx, event, referenceId, customerId, callerIp)
+ return fulfillOrder(ctx, event, referenceId, customerId, callerIp)
}
// sessionAsyncPaymentFailed marks orders as failed when delayed payment methods
@@ -255,11 +261,12 @@ func sessionAsyncPaymentFailed(ctx context.Context, event stripe.Event, callerIp
logger.LogInfo(ctx, fmt.Sprintf("Stripe 充值订单已标记为失败 trade_no=%s client_ip=%s", referenceId, callerIp))
}
-// fulfillOrder is the shared logic for crediting quota after payment is confirmed.
-func fulfillOrder(ctx context.Context, event stripe.Event, referenceId string, customerId string, callerIp string) {
+// fulfillOrder returns only retryable processing errors. Replayed events and
+// permanent local-order mismatches must be acknowledged without a second credit.
+func fulfillOrder(ctx context.Context, event stripe.Event, referenceId string, customerId string, callerIp string) error {
if len(referenceId) == 0 {
logger.LogWarn(ctx, fmt.Sprintf("Stripe 完成订单时缺少订单号 client_ip=%s", callerIp))
- return
+ return nil
}
LockOrder(referenceId)
@@ -272,21 +279,49 @@ func fulfillOrder(ctx context.Context, event stripe.Event, referenceId string, c
}
if err := model.CompleteSubscriptionOrder(referenceId, common.GetJsonString(payload), model.PaymentProviderStripe, ""); err == nil {
logger.LogInfo(ctx, fmt.Sprintf("Stripe 订阅订单处理成功 trade_no=%s event_type=%s client_ip=%s", referenceId, string(event.Type), callerIp))
- return
- } else if err != nil && !errors.Is(err, model.ErrSubscriptionOrderNotFound) {
- logger.LogError(ctx, fmt.Sprintf("Stripe 订阅订单处理失败 trade_no=%s event_type=%s client_ip=%s error=%q", referenceId, string(event.Type), callerIp, err.Error()))
- return
+ return nil
+ } else if !errors.Is(err, model.ErrSubscriptionOrderNotFound) {
+ if errors.Is(err, model.ErrPaymentMethodMismatch) || errors.Is(err, model.ErrSubscriptionOrderStatusInvalid) {
+ logger.LogWarn(ctx, fmt.Sprintf("Stripe 订阅订单处于终态，忽略重放 trade_no=%s event_type=%s client_ip=%s error=%q", referenceId, string(event.Type), callerIp, err.Error()))
+ return nil
+ }
+ return fmt.Errorf("complete subscription order: %w", err)
}
- err := model.Recharge(referenceId, customerId, callerIp)
+ topUp, err := model.GetTopUpByTradeNoWithError(referenceId)
+ if errors.Is(err, model.ErrTopUpNotFound) {
+ logger.LogWarn(ctx, fmt.Sprintf("Stripe 充值订单不存在，忽略回调 trade_no=%s event_type=%s client_ip=%s", referenceId, string(event.Type), callerIp))
+ return nil
+ }
if err != nil {
- logger.LogError(ctx, fmt.Sprintf("Stripe 充值处理失败 trade_no=%s event_type=%s client_ip=%s error=%q", referenceId, string(event.Type), callerIp, err.Error()))
- return
+ return fmt.Errorf("load top-up order: %w", err)
+ }
+ if topUp.PaymentProvider != model.PaymentProviderStripe {
+ logger.LogWarn(ctx, fmt.Sprintf("Stripe 充值订单支付网关不匹配，忽略回调 trade_no=%s payment_provider=%s client_ip=%s", referenceId, topUp.PaymentProvider, callerIp))
+ return nil
+ }
+ if topUp.Status == common.TopUpStatusSuccess {
+ logger.LogInfo(ctx, fmt.Sprintf("Stripe 充值订单已处理，确认重放 trade_no=%s client_ip=%s", referenceId, callerIp))
+ return nil
+ }
+ if topUp.Status != common.TopUpStatusPending {
+ logger.LogWarn(ctx, fmt.Sprintf("Stripe 充值订单处于终态，忽略回调 trade_no=%s status=%s client_ip=%s", referenceId, topUp.Status, callerIp))
+ return nil
+ }
+
+ if err := model.Recharge(referenceId, customerId, callerIp); err != nil {
+ latest, latestErr := model.GetTopUpByTradeNoWithError(referenceId)
+ if latestErr == nil && latest.Status == common.TopUpStatusSuccess {
+ logger.LogInfo(ctx, fmt.Sprintf("Stripe 充值订单已由并发回调完成 trade_no=%s client_ip=%s", referenceId, callerIp))
+ return nil
+ }
+ return fmt.Errorf("recharge top-up: %w", err)
}
total, _ := strconv.ParseFloat(event.GetObjectValue("amount_total"), 64)
currency := strings.ToUpper(event.GetObjectValue("currency"))
logger.LogInfo(ctx, fmt.Sprintf("Stripe 充值成功 trade_no=%s amount_total=%.2f currency=%s event_type=%s client_ip=%s", referenceId, total/100, currency, string(event.Type), callerIp))
+ return nil
}
func sessionExpired(ctx context.Context, event stripe.Event) {
diff --git a/controller/topup_stripe_test.go b/controller/topup_stripe_test.go
new file mode 100644
index 00000000..0a56b8a4
--- /dev/null
+++ b/controller/topup_stripe_test.go
@@ -0,0 +1,106 @@
+package controller
+
+import (
+ "context"
+ "fmt"
+ "strings"
+ "testing"
+
+ "github.com/MAX-API-Next/MAX-API/common"
+ "github.com/MAX-API-Next/MAX-API/model"
+ "github.com/glebarez/sqlite"
+ "github.com/stretchr/testify/require"
+ "github.com/stripe/stripe-go/v81"
+ "gorm.io/gorm"
+)
+
+func setupStripeWebhookTestDB(t *testing.T) *gorm.DB {
+ t.Helper()
+ oldDB := model.DB
+ oldLogDB := model.LOG_DB
+ oldRedisEnabled := common.RedisEnabled
+ oldUsingSQLite := common.UsingSQLite
+ oldUsingMySQL := common.UsingMySQL
+ oldUsingPostgreSQL := common.UsingPostgreSQL
+
+ common.RedisEnabled = false
+ common.UsingSQLite = true
+ common.UsingMySQL = false
+ common.UsingPostgreSQL = false
+ dsn := fmt.Sprintf("file:%s?mode=memory&cache=shared", strings.ReplaceAll(t.Name(), "/", "_"))
+ db, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{})
+ require.NoError(t, err)
+ model.DB = db
+ model.LOG_DB = db
+ require.NoError(t, db.AutoMigrate(&model.User{}, &model.TopUp{}, &model.SubscriptionOrder{}))
+
+ t.Cleanup(func() {
+ sqlDB, err := db.DB()
+ if err == nil {
+ _ = sqlDB.Close()
+ }
+ model.DB = oldDB
+ model.LOG_DB = oldLogDB
+ common.RedisEnabled = oldRedisEnabled
+ common.UsingSQLite = oldUsingSQLite
+ common.UsingMySQL = oldUsingMySQL
+ common.UsingPostgreSQL = oldUsingPostgreSQL
+ })
+ return db
+}
+
+func completedStripeEvent(tradeNo string) stripe.Event {
+ return stripe.Event{
+ Type: stripe.EventTypeCheckoutSessionCompleted,
+ Data: &stripe.EventData{Object: map[string]interface{}{
+ "client_reference_id": tradeNo,
+ "customer": "cus_test",
+ "status": "complete",
+ "payment_status": "paid",
+ }},
+ }
+}
+
+func TestStripeSessionCompletedReturnsRetryableErrorWhenRechargeFails(t *testing.T) {
+ db := setupStripeWebhookTestDB(t)
+ topUp := model.TopUp{
+ Id: 8101,
+ UserId: 9999,
+ Money: 1,
+ TradeNo: "stripe-retryable-failure",
+ PaymentProvider: model.PaymentProviderStripe,
+ Status: common.TopUpStatusPending,
+ }
+ require.NoError(t, db.Create(&topUp).Error)
+
+ err := sessionCompleted(context.Background(), completedStripeEvent(topUp.TradeNo), "127.0.0.1")
+ require.Error(t, err)
+ var stored model.TopUp
+ require.NoError(t, db.First(&stored, topUp.Id).Error)
+ require.Equal(t, common.TopUpStatusPending, stored.Status)
+}
+
+func TestStripeSessionCompletedAcknowledgesCompletedTopUpReplay(t *testing.T) {
+ db := setupStripeWebhookTestDB(t)
+ topUp := model.TopUp{
+ Id: 8102,
+ UserId: 9999,
+ Money: 1,
+ TradeNo: "stripe-completed-replay",
+ PaymentProvider: model.PaymentProviderStripe,
+ Status: common.TopUpStatusSuccess,
+ }
+ require.NoError(t, db.Create(&topUp).Error)
+
+ require.NoError(t, sessionCompleted(context.Background(), completedStripeEvent(topUp.TradeNo), "127.0.0.1"))
+}
+
+func TestStripeSessionCompletedReturnsRetryableErrorWhenOrderLookupFails(t *testing.T) {
+ db := setupStripeWebhookTestDB(t)
+ sqlDB, err := db.DB()
+ require.NoError(t, err)
+ require.NoError(t, sqlDB.Close())
+
+ err = sessionCompleted(context.Background(), completedStripeEvent("stripe-database-unavailable"), "127.0.0.1")
+ require.Error(t, err)
+}
diff --git a/controller/user.go b/controller/user.go
index 131d84a5..9c666962 100644
--- a/controller/user.go
+++ b/controller/user.go
@@ -915,9 +915,6 @@ func DeleteUser(c *gin.Context) {
common.ApiError(c, err)
return
}
- if err := model.InvalidateUserTokensCache(id); err != nil {
- common.SysLog(fmt.Sprintf("failed to invalidate tokens cache for user %d: %s", id, err.Error()))
- }
recordManageAuditFor(c, originUser.Id, "user.delete", map[string]interface{}{
"username": originUser.Username,
"id": originUser.Id,
@@ -943,9 +940,6 @@ func DeleteSelf(c *gin.Context) {
common.ApiError(c, err)
return
}
- if err := model.InvalidateUserTokensCache(id); err != nil {
- common.SysLog(fmt.Sprintf("failed to invalidate tokens cache for user %d: %s", id, err.Error()))
- }
c.JSON(http.StatusOK, gin.H{
"success": true,
"message": "",
@@ -1169,14 +1163,9 @@ func ManageUser(c *gin.Context) {
common.ApiError(c, err)
return
}
- // 禁用 / 角色调整后，强制失效用户缓存与其全部令牌缓存，
- // 避免在 Redis TTL 过期前仍使用旧状态（尤其是禁用后仍可发起请求的问题）。
- // InvalidateUserCache 会让下一次 GetUserCache 从数据库重新加载，
- // InvalidateUserTokensCache 则确保令牌侧的缓存也同步刷新。
- if req.Action == "disable" || req.Action == "promote" || req.Action == "demote" {
- if err := model.InvalidateUserCache(user.Id); err != nil {
- common.SysLog(fmt.Sprintf("failed to invalidate user cache for user %d: %s", user.Id, err.Error()))
- }
+ // UpdateFields 已在同一事务中登记用户缓存失效；状态或角色变化时
+ // 额外刷新该用户的 Token 缓存，避免令牌侧继续使用旧快照。
+ if req.Action == "enable" || req.Action == "disable" || req.Action == "promote" || req.Action == "demote" {
if err := model.InvalidateUserTokensCache(user.Id); err != nil {
common.SysLog(fmt.Sprintf("failed to invalidate tokens cache for user %d: %s", user.Id, err.Error()))
}
diff --git a/controller/vendor_meta.go b/controller/vendor_meta.go
index 4820bf97..efd0d0ce 100644
--- a/controller/vendor_meta.go
+++ b/controller/vendor_meta.go
@@ -1,6 +1,7 @@
package controller
import (
+ "errors"
"strconv"
"github.com/MAX-API-Next/MAX-API/common"
@@ -116,7 +117,16 @@ func DeleteVendorMeta(c *gin.Context) {
common.ApiError(c, err)
return
}
- if err := model.DB.Delete(&model.Vendor{}, id).Error; err != nil {
+ v, err := model.GetVendorByID(id)
+ if err != nil {
+ common.ApiError(c, err)
+ return
+ }
+ if err := v.Delete(); err != nil {
+ if errors.Is(err, model.ErrVendorHasModels) {
+ common.ApiErrorMsg(c, "供应商仍有关联模型，请先重新分配或删除这些模型")
+ return
+ }
common.ApiError(c, err)
return
}
diff --git a/docs/openapi/api.json b/docs/openapi/api.json
index 01f39ce7..4527a57e 100644
--- a/docs/openapi/api.json
+++ b/docs/openapi/api.json
@@ -2320,13 +2320,14 @@
"password": {
"type": "string",
"format": "password",
- "description": "当前账户密码，仅method=password且scope=access_token时使用"
+ "description": "当前账户密码，仅method=password且scope=access_token或account_delete时使用"
},
"scope": {
"type": "string",
- "description": "验证作用域；密码验证仅支持access_token",
+ "description": "验证作用域；密码验证仅支持access_token或account_delete",
"enum": [
- "access_token"
+ "access_token",
+ "account_delete"
]
}
}
@@ -2363,11 +2364,12 @@
"name": "scope",
"in": "query",
"required": false,
- "description": "验证作用域；传access_token时，无2FA或Passkey的用户可使用密码验证",
+ "description": "验证作用域；传access_token或account_delete时，无2FA或Passkey的用户可使用密码验证",
"schema": {
"type": "string",
"enum": [
- "access_token"
+ "access_token",
+ "account_delete"
]
}
}
diff --git a/dto/openai_request.go b/dto/openai_request.go
index 789e7683..16cf20f9 100644
--- a/dto/openai_request.go
+++ b/dto/openai_request.go
@@ -257,12 +257,18 @@ type StreamOptions struct {
IncludeObfuscation bool `json:"include_obfuscation,omitempty"`
}
-func (r *GeneralOpenAIRequest) GetMaxTokens() uint {
- maxCompletionTokens := lo.FromPtrOr(r.MaxCompletionTokens, uint(0))
- if maxCompletionTokens != 0 {
+// ResolveMaxTokens prefers a non-zero max_completion_tokens value. An explicit
+// zero max_completion_tokens is retained when max_tokens is absent so callers
+// can preserve the distinction between zero and an omitted limit.
+func ResolveMaxTokens(maxTokens, maxCompletionTokens *uint) *uint {
+ if maxCompletionTokens != nil && (*maxCompletionTokens != 0 || maxTokens == nil) {
return maxCompletionTokens
}
- return lo.FromPtrOr(r.MaxTokens, uint(0))
+ return maxTokens
+}
+
+func (r *GeneralOpenAIRequest) GetMaxTokens() uint {
+ return lo.FromPtrOr(ResolveMaxTokens(r.MaxTokens, r.MaxCompletionTokens), uint(0))
}
func (r *GeneralOpenAIRequest) ParseInput() []string {
diff --git a/dto/openai_request_zero_value_test.go b/dto/openai_request_zero_value_test.go
index e46e2023..2ffbb01b 100644
--- a/dto/openai_request_zero_value_test.go
+++ b/dto/openai_request_zero_value_test.go
@@ -50,6 +50,50 @@ func TestGeneralOpenAIRequestPreserveExplicitZeroValues(t *testing.T) {
require.True(t, gjson.GetBytes(encoded, "return_related_questions").Exists())
}
+func TestResolveMaxTokens(t *testing.T) {
+ legacy := uint(64)
+ completion := uint(128)
+ zero := uint(0)
+ tests := []struct {
+ name string
+ maxTokens *uint
+ maxCompletionTokens *uint
+ want *uint
+ }{
+ {
+ name: "prefers non-zero max completion tokens",
+ maxTokens: &legacy,
+ maxCompletionTokens: &completion,
+ want: &completion,
+ },
+ {
+ name: "keeps legacy max tokens when completion is zero",
+ maxTokens: &legacy,
+ maxCompletionTokens: &zero,
+ want: &legacy,
+ },
+ {
+ name: "preserves zero max completion tokens without legacy value",
+ maxCompletionTokens: &zero,
+ want: &zero,
+ },
+ {
+ name: "keeps legacy max tokens",
+ maxTokens: &legacy,
+ want: &legacy,
+ },
+ {
+ name: "returns nil when both values are absent",
+ },
+ }
+
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ require.Same(t, tt.want, ResolveMaxTokens(tt.maxTokens, tt.maxCompletionTokens))
+ })
+ }
+}
+
func TestOpenAIResponsesRequestPreserveExplicitZeroValues(t *testing.T) {
raw := []byte(`{
"model":"gpt-4.1",
diff --git a/dto/task.go b/dto/task.go
index 4a9a8e2e..78b61676 100644
--- a/dto/task.go
+++ b/dto/task.go
@@ -9,8 +9,10 @@ type TaskError struct {
Message string `json:"message"`
Data any `json:"data"`
StatusCode int `json:"-"`
- LocalError bool `json:"-"`
- Error error `json:"-"`
+ // UpstreamStatusCode preserves retry classification when a channel mapping changes the downstream status.
+ UpstreamStatusCode int `json:"-"`
+ LocalError bool `json:"-"`
+ Error error `json:"-"`
}
type TaskData interface {
diff --git a/main.go b/main.go
index 52e46e71..53036daf 100644
--- a/main.go
+++ b/main.go
@@ -30,6 +30,7 @@ import (
"github.com/MAX-API-Next/MAX-API/service"
_ "github.com/MAX-API-Next/MAX-API/setting/performance_setting"
"github.com/MAX-API-Next/MAX-API/setting/ratio_setting"
+ "github.com/MAX-API-Next/MAX-API/setting/system_setting"
"github.com/bytedance/gopkg/util/gopool"
"github.com/gin-contrib/sessions"
@@ -70,9 +71,13 @@ func main() {
}
}()
+ // Recover durable balance settlements after process restarts or transient DB failures.
+ model.StartBillingSettlementTaskRunner()
+
if common.RedisEnabled {
// for compatibility with old versions
common.MemoryCacheEnabled = true
+ model.StartCacheInvalidationTaskRunner()
}
if common.MemoryCacheEnabled {
common.SysLog("memory cache enabled")
@@ -184,6 +189,9 @@ func main() {
server.Use(middleware.I18n())
middleware.SetUpLogger(server)
// Initialize session store
+ if err := configureSessionCookieSecure(system_setting.ServerAddress); err != nil {
+ common.FatalLog(err.Error())
+ }
store := cookie.NewStore([]byte(common.SessionSecret))
store.Options(sessions.Options{
Path: "/",
@@ -234,6 +242,16 @@ func main() {
ctx, cancel := context.WithTimeout(context.Background(), shutdownTimeout)
defer cancel()
shutdownHTTPServer(ctx, srv)
+ stopBackgroundRunnerCtx, stopBackgroundRunnerCancel := context.WithTimeout(context.Background(), shutdownTimeout)
+ if err := model.StopBillingSettlementTaskRunner(stopBackgroundRunnerCtx); err != nil {
+ common.SysError(fmt.Sprintf("timed out stopping billing settlement runner: %v", err))
+ }
+ if common.RedisEnabled {
+ if err := model.StopCacheInvalidationTaskRunner(stopBackgroundRunnerCtx); err != nil {
+ common.SysError(fmt.Sprintf("timed out stopping cache invalidation runner: %v", err))
+ }
+ }
+ stopBackgroundRunnerCancel()
if common.DataExportEnabled {
saveTimeout := time.Duration(common.GetEnvOrDefault("QUOTA_DATA_CACHE_SAVE_TIMEOUT_SECONDS", 30)) * time.Second
if !runWithTimeout(saveTimeout, model.WaitPendingLogQuotaData) {
@@ -253,6 +271,25 @@ func main() {
common.SysLog("server exited")
}
+func configureSessionCookieSecure(serverAddress string) error {
+ value, explicitlyConfigured := os.LookupEnv("SESSION_COOKIE_SECURE")
+ value = strings.TrimSpace(value)
+ if explicitlyConfigured && value != "" {
+ secure, err := strconv.ParseBool(value)
+ if err != nil {
+ return fmt.Errorf("invalid SESSION_COOKIE_SECURE value %q: %w", value, err)
+ }
+ common.SessionCookieSecure = secure
+ return nil
+ }
+
+ common.SessionCookieSecure = common.SessionCookieSecureForServerAddress(serverAddress)
+ if !common.SessionCookieSecure {
+ common.SysLog("SESSION_COOKIE_SECURE is not set and ServerAddress is not HTTPS; session cookies may traverse plain HTTP. Set SESSION_COOKIE_SECURE=true for HTTPS deployments")
+ }
+ return nil
+}
+
func shutdownHTTPServer(ctx context.Context, srv *http.Server) {
if err := srv.Shutdown(ctx); err != nil {
common.SysError(fmt.Sprintf("server forced to shutdown: %v", err))
diff --git a/main_test.go b/main_test.go
index c15b7d46..8a7753fc 100644
--- a/main_test.go
+++ b/main_test.go
@@ -8,6 +8,7 @@ import (
"testing"
"time"
+ "github.com/MAX-API-Next/MAX-API/common"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
@@ -100,3 +101,29 @@ func TestShutdownHTTPServerClosesActiveHandlersAfterTimeout(t *testing.T) {
t.Fatal("client request did not finish after forced close")
}
}
+
+func TestConfigureSessionCookieSecureTreatsBlankAsUnset(t *testing.T) {
+ oldSecure := common.SessionCookieSecure
+ t.Cleanup(func() { common.SessionCookieSecure = oldSecure })
+ t.Setenv("SESSION_COOKIE_SECURE", "")
+
+ require.NoError(t, configureSessionCookieSecure("https://example.com"))
+
+ assert.True(t, common.SessionCookieSecure)
+}
+
+func TestConfigureSessionCookieSecurePreservesExplicitFalse(t *testing.T) {
+ oldSecure := common.SessionCookieSecure
+ t.Cleanup(func() { common.SessionCookieSecure = oldSecure })
+ t.Setenv("SESSION_COOKIE_SECURE", "false")
+
+ require.NoError(t, configureSessionCookieSecure("https://example.com"))
+
+ assert.False(t, common.SessionCookieSecure)
+}
+
+func TestConfigureSessionCookieSecureRejectsInvalidValue(t *testing.T) {
+ t.Setenv("SESSION_COOKIE_SECURE", "sometimes")
+
+ require.Error(t, configureSessionCookieSecure("https://example.com"))
+}
diff --git a/middleware/rate-limit.go b/middleware/rate-limit.go
index 6a67b37b..6250845b 100644
--- a/middleware/rate-limit.go
+++ b/middleware/rate-limit.go
@@ -82,11 +82,14 @@ func applyRateLimit(c *gin.Context, maxRequestNum int, duration int64, key strin
if maxRequestNum <= 0 {
return
}
- if common.RedisEnabled {
- redisRateLimiter(c, maxRequestNum, duration, key, policy)
+ if common.RedisEnabled && redisRateLimiter(c, maxRequestNum, duration, key, policy) {
return
}
+ applyInMemoryRateLimit(c, maxRequestNum, duration, key, policy)
+}
+
+func applyInMemoryRateLimit(c *gin.Context, maxRequestNum int, duration int64, key string, policy string) {
inMemoryRateLimiter.Init(common.RateLimitKeyExpirationDuration)
allowed, retryAfter := inMemoryRateLimiter.RequestWithRetry(key, maxRequestNum, duration)
if !allowed {
@@ -94,41 +97,48 @@ func applyRateLimit(c *gin.Context, maxRequestNum int, duration int64, key strin
}
}
-func redisRateLimiter(c *gin.Context, maxRequestNum int, duration int64, key string, policy string) {
+// redisRateLimiter returns false when Redis cannot produce a valid decision so
+// callers can preserve a local rate limit without converting an outage to 500.
+func redisRateLimiter(c *gin.Context, maxRequestNum int, duration int64, key string, policy string) bool {
ctx, cancel := context.WithTimeout(c.Request.Context(), rateLimitRedisTimeout)
defer cancel()
result, err := common.RDB.Eval(ctx, rollingWindowRateLimitScript, []string{key}, duration*1000, maxRequestNum).Result()
if err != nil {
logRedisRateLimitFailure(c, policy, "evaluation", err)
- return
+ return false
}
values, ok := result.([]interface{})
if !ok || len(values) != 2 {
logRedisRateLimitFailure(c, policy, "response", fmt.Errorf("unexpected result type %T", result))
- return
+ return false
}
allowed, err := redisRateLimitInt(values[0])
if err != nil {
logRedisRateLimitFailure(c, policy, "allowed value", err)
- return
+ return false
}
if allowed == 1 {
- return
+ return true
+ }
+ if allowed != 0 {
+ logRedisRateLimitFailure(c, policy, "allowed value", fmt.Errorf("unexpected decision %d", allowed))
+ return false
}
retryAfter, err := redisRateLimitInt(values[1])
if err != nil {
logRedisRateLimitFailure(c, policy, "retry value", err)
- return
+ return false
}
if retryAfter < 1 {
retryAfter = 1
}
abortRateLimited(c, maxRequestNum, time.Duration(retryAfter)*time.Second, policy)
+ return true
}
func logRedisRateLimitFailure(c *gin.Context, policy string, stage string, err error) {
- logger.LogError(c.Request.Context(), fmt.Sprintf("rate limit Redis %s failed for policy %s; request allowed: %v", stage, policy, err))
+ logger.LogError(c.Request.Context(), fmt.Sprintf("rate limit Redis %s failed for policy %s; falling back to local limiter: %v", stage, policy, err))
}
func redisRateLimitInt(value interface{}) (int64, error) {
diff --git a/middleware/rate_limit_test.go b/middleware/rate_limit_test.go
index 95656a8f..dd9035ca 100644
--- a/middleware/rate_limit_test.go
+++ b/middleware/rate_limit_test.go
@@ -283,8 +283,9 @@ func TestRedisRateLimitIsAtomicUnderConcurrentRequests(t *testing.T) {
require.Equal(t, int32(5), atomic.LoadInt32(&allowed))
}
-func TestRedisRateLimitFailureFailsOpen(t *testing.T) {
+func TestRedisRateLimitFailureFallsBackToInMemory(t *testing.T) {
gin.SetMode(gin.TestMode)
+ testRun := atomic.AddUint64(&rateLimitTestRun, 1)
server, err := miniredis.Run()
require.NoError(t, err)
address := server.Addr()
@@ -305,17 +306,27 @@ func TestRedisRateLimitFailureFailsOpen(t *testing.T) {
})
engine := gin.New()
- engine.GET("/", rateLimitFactory(1, 60, "redis-failure-test"), func(c *gin.Context) {
+ mark := fmt.Sprintf("redis-failure-fallback-test-%d", testRun)
+ engine.GET("/", rateLimitFactory(1, 60, mark), func(c *gin.Context) {
c.Status(http.StatusNoContent)
})
- recorder := httptest.NewRecorder()
- engine.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, "/", nil))
+ request := func() *httptest.ResponseRecorder {
+ recorder := httptest.NewRecorder()
+ req := httptest.NewRequest(http.MethodGet, "/", nil)
+ req.RemoteAddr = fmt.Sprintf("[2001:db8:%x::8]:1234", testRun)
+ engine.ServeHTTP(recorder, req)
+ return recorder
+ }
- require.Equal(t, http.StatusNoContent, recorder.Code)
+ require.Equal(t, http.StatusNoContent, request().Code)
+ limited := request()
+ require.Equal(t, http.StatusTooManyRequests, limited.Code)
+ require.Equal(t, mark, limited.Header().Get("X-RateLimit-Policy"))
}
-func TestRedisRateLimitUsesBoundedTimeoutAndFailsOpen(t *testing.T) {
+func TestRedisRateLimitUsesBoundedTimeoutAndFallsBackToInMemory(t *testing.T) {
gin.SetMode(gin.TestMode)
+ testRun := atomic.AddUint64(&rateLimitTestRun, 1)
listener, err := net.Listen("tcp", "127.0.0.1:0")
require.NoError(t, err)
accepted := make(chan net.Conn, 1)
@@ -351,13 +362,23 @@ func TestRedisRateLimitUsesBoundedTimeoutAndFailsOpen(t *testing.T) {
})
engine := gin.New()
- engine.GET("/", rateLimitFactory(1, 60, "redis-timeout-test"), func(c *gin.Context) {
+ mark := fmt.Sprintf("redis-timeout-fallback-test-%d", testRun)
+ engine.GET("/", rateLimitFactory(1, 60, mark), func(c *gin.Context) {
c.Status(http.StatusNoContent)
})
- recorder := httptest.NewRecorder()
+ request := func() *httptest.ResponseRecorder {
+ recorder := httptest.NewRecorder()
+ req := httptest.NewRequest(http.MethodGet, "/", nil)
+ req.RemoteAddr = fmt.Sprintf("[2001:db8:%x::9]:1234", testRun)
+ engine.ServeHTTP(recorder, req)
+ return recorder
+ }
startedAt := time.Now()
- engine.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, "/", nil))
+ recorder := request()
require.Less(t, time.Since(startedAt), time.Second)
require.Equal(t, http.StatusNoContent, recorder.Code)
+ limited := request()
+ require.Equal(t, http.StatusTooManyRequests, limited.Code)
+ require.Equal(t, mark, limited.Header().Get("X-RateLimit-Policy"))
}
diff --git a/middleware/secure_verification.go b/middleware/secure_verification.go
index 6b9181e4..03ba00e2 100644
--- a/middleware/secure_verification.go
+++ b/middleware/secure_verification.go
@@ -98,6 +98,20 @@ func SecureVerificationRequired(requiredScopes ...string) gin.HandlerFunc {
c.Abort()
return
}
+ requiredScope := ""
+ if len(requiredScopes) > 0 {
+ requiredScope = requiredScopes[0]
+ }
+ verifiedScope, _ := session.Get(secureVerificationScopeSessionKey).(string)
+ if requiredScope != "" && verifiedScope != requiredScope {
+ c.JSON(http.StatusForbidden, gin.H{
+ "success": false,
+ "message": "需要对应操作的安全验证",
+ "code": "VERIFICATION_REQUIRED",
+ })
+ c.Abort()
+ return
+ }
if verifiedMethod == secureVerificationMethodPassword {
requiredScope := ""
if len(requiredScopes) > 0 {
diff --git a/middleware/secure_verification_test.go b/middleware/secure_verification_test.go
index 96c5daf5..59414bcd 100644
--- a/middleware/secure_verification_test.go
+++ b/middleware/secure_verification_test.go
@@ -71,14 +71,17 @@ func TestPasswordVerificationIsRestrictedToMatchingScope(t *testing.T) {
router.GET("/other-sensitive", setUser, SecureVerificationRequired(), func(c *gin.Context) {
c.Status(http.StatusNoContent)
})
+ router.DELETE("/account", setUser, SecureVerificationRequired("account_delete"), func(c *gin.Context) {
+ c.Status(http.StatusNoContent)
+ })
seedRecorder := httptest.NewRecorder()
router.ServeHTTP(seedRecorder, httptest.NewRequest(http.MethodGet, "/seed", nil))
require.Equal(t, http.StatusNoContent, seedRecorder.Code)
- perform := func(path string) *httptest.ResponseRecorder {
+ perform := func(method, path string) *httptest.ResponseRecorder {
recorder := httptest.NewRecorder()
- request := httptest.NewRequest(http.MethodGet, path, nil)
+ request := httptest.NewRequest(method, path, nil)
for _, sessionCookie := range seedRecorder.Result().Cookies() {
request.AddCookie(sessionCookie)
}
@@ -86,8 +89,54 @@ func TestPasswordVerificationIsRestrictedToMatchingScope(t *testing.T) {
return recorder
}
- require.Equal(t, http.StatusNoContent, perform("/token").Code)
- otherRecorder := perform("/other-sensitive")
+ require.Equal(t, http.StatusNoContent, perform(http.MethodGet, "/token").Code)
+ otherRecorder := perform(http.MethodGet, "/other-sensitive")
require.Equal(t, http.StatusForbidden, otherRecorder.Code)
require.Contains(t, otherRecorder.Body.String(), `"code":"VERIFICATION_REQUIRED"`)
+ deleteRecorder := perform(http.MethodDelete, "/account")
+ require.Equal(t, http.StatusForbidden, deleteRecorder.Code)
+ require.Contains(t, deleteRecorder.Body.String(), `"code":"VERIFICATION_REQUIRED"`)
+}
+
+func TestNonPasswordVerificationIsRestrictedToMatchingScope(t *testing.T) {
+ router := gin.New()
+ router.Use(sessions.Sessions("session", cookie.NewStore([]byte("secure-verification-passkey-scope-test"))))
+ router.GET("/seed", func(c *gin.Context) {
+ session := sessions.Default(c)
+ session.Set(SecureVerificationSessionKey, time.Now().Unix())
+ session.Set(secureVerificationUserSessionKey, 1001)
+ session.Set(secureVerificationMethodSessionKey, "passkey")
+ session.Set(secureVerificationScopeSessionKey, "access_token")
+ require.NoError(t, session.Save())
+ c.Status(http.StatusNoContent)
+ })
+ setUser := func(c *gin.Context) {
+ c.Set("id", 1001)
+ c.Next()
+ }
+ router.GET("/token", setUser, SecureVerificationRequired("access_token"), func(c *gin.Context) {
+ c.Status(http.StatusNoContent)
+ })
+ router.DELETE("/account", setUser, SecureVerificationRequired("account_delete"), func(c *gin.Context) {
+ c.Status(http.StatusNoContent)
+ })
+
+ seedRecorder := httptest.NewRecorder()
+ router.ServeHTTP(seedRecorder, httptest.NewRequest(http.MethodGet, "/seed", nil))
+ require.Equal(t, http.StatusNoContent, seedRecorder.Code)
+
+ perform := func(method, path string) *httptest.ResponseRecorder {
+ recorder := httptest.NewRecorder()
+ request := httptest.NewRequest(method, path, nil)
+ for _, sessionCookie := range seedRecorder.Result().Cookies() {
+ request.AddCookie(sessionCookie)
+ }
+ router.ServeHTTP(recorder, request)
+ return recorder
+ }
+
+ require.Equal(t, http.StatusNoContent, perform(http.MethodGet, "/token").Code)
+ deleteRecorder := perform(http.MethodDelete, "/account")
+ require.Equal(t, http.StatusForbidden, deleteRecorder.Code)
+ require.Contains(t, deleteRecorder.Body.String(), `"code":"VERIFICATION_REQUIRED"`)
}
diff --git a/model/ability.go b/model/ability.go
index 289f99cb..c5c54e63 100644
--- a/model/ability.go
+++ b/model/ability.go
@@ -66,6 +66,12 @@ func GetAllEnableAbilities() []Ability {
}
func GetChannel(group string, model string, retry int, requestPath string) (*Channel, error) {
+ return GetChannelExcluding(group, model, retry, requestPath, nil)
+}
+
+// GetChannelExcluding selects a channel while excluding channels already tried
+// during the current request's retry sequence.
+func GetChannelExcluding(group string, model string, retry int, requestPath string, excludedChannelIDs map[int]struct{}) (*Channel, error) {
channelQuery, err := getChannelQuery(group, model, retry)
if err != nil {
return nil, err
@@ -88,7 +94,7 @@ func GetChannel(group string, model string, retry int, requestPath string) (*Cha
return nil, nil
}
- channelId, ok := selectChannelIdFromAbilities(abilities, retry)
+ channelId, ok := selectChannelIdFromAbilities(abilities, retry, excludedChannelIDs)
if !ok {
return nil, nil
}
@@ -160,13 +166,16 @@ func abilityCol(name string) string {
return "`abilities`.`" + name + "`"
}
-func selectChannelIdFromAbilities(abilities []Ability, retry int) (int, bool) {
+func selectChannelIdFromAbilities(abilities []Ability, retry int, excludedChannelIDs map[int]struct{}) (int, bool) {
if len(abilities) == 0 {
return 0, false
}
priorities := make([]int64, 0, len(abilities))
seenPriorities := make(map[int64]struct{}, len(abilities))
for _, ability := range abilities {
+ if _, excluded := excludedChannelIDs[ability.ChannelId]; excluded {
+ continue
+ }
priority := abilityPriority(ability)
if _, ok := seenPriorities[priority]; ok {
continue
@@ -184,7 +193,7 @@ func selectChannelIdFromAbilities(abilities []Ability, retry int) (int, bool) {
targetAbilities := make([]Ability, 0, len(abilities))
for _, ability := range abilities {
- if abilityPriority(ability) == targetPriority {
+ if _, excluded := excludedChannelIDs[ability.ChannelId]; !excluded && abilityPriority(ability) == targetPriority {
targetAbilities = append(targetAbilities, ability)
}
}
diff --git a/model/billing_settlement.go b/model/billing_settlement.go
new file mode 100644
index 00000000..4063ce2b
--- /dev/null
+++ b/model/billing_settlement.go
@@ -0,0 +1,934 @@
+package model
+
+import (
+ "context"
+ "errors"
+ "fmt"
+ "sync"
+ "time"
+
+ "github.com/MAX-API-Next/MAX-API/common"
+ "github.com/bytedance/gopkg/util/gopool"
+ "gorm.io/gorm"
+ "gorm.io/gorm/clause"
+)
+
+const (
+ BillingSettlementSourceWallet = "wallet"
+ BillingSettlementSourceSubscription = "subscription"
+ BillingSettlementStatusPending = "pending"
+ BillingSettlementStatusApplied = "applied"
+ BillingSettlementStatusManual = "manual"
+ BillingSettlementEffectPending = "pending"
+ BillingSettlementEffectApplying = "applying"
+ BillingSettlementEffectApplied = "applied"
+)
+
+var (
+ ErrBillingSettlementManualReview = errors.New("billing settlement requires manual review")
+ ErrBillingSettlementTaskConflict = errors.New("billing settlement task quota conflict")
+ ErrBillingSettlementOperationConflict = errors.New("billing settlement operation conflict")
+ ErrSubscriptionSettlementUnbound = errors.New("subscription settlement is not bound to its pre-consume request")
+ ErrSubscriptionSettlementPeriodChanged = errors.New("subscription settlement crossed a quota reset period")
+)
+
+type permanentBillingSettlementError struct {
+ err error
+}
+
+func (e *permanentBillingSettlementError) Error() string { return e.err.Error() }
+func (e *permanentBillingSettlementError) Unwrap() error { return e.err }
+
+func permanentBillingSettlement(err error) error {
+ if err == nil {
+ return nil
+ }
+ var permanentErr *permanentBillingSettlementError
+ if errors.As(err, &permanentErr) {
+ return err
+ }
+ return &permanentBillingSettlementError{err: err}
+}
+
+func isPermanentBillingSettlementError(err error) bool {
+ var permanentErr *permanentBillingSettlementError
+ return errors.As(err, &permanentErr)
+}
+
+// BillingSettlement is the durable idempotency record for a balance mutation.
+// The record and every related DB mutation are committed in one transaction.
+type BillingSettlement struct {
+ ID int64 `gorm:"primaryKey"`
+ OperationKey string `gorm:"type:varchar(191);uniqueIndex;not null"`
+ Source string `gorm:"type:varchar(32);not null"`
+ UserID int `gorm:"not null;default:0"`
+ SubscriptionID int `gorm:"not null;default:0"`
+ TokenID int `gorm:"not null;default:0"`
+ FundingDelta int64 `gorm:"not null;default:0"`
+ AppliedFundingDelta int64 `gorm:"not null;default:0"`
+ TokenDelta int64 `gorm:"not null;default:0"`
+ AppliedTokenDelta int64 `gorm:"not null;default:0"`
+ TaskID int64 `gorm:"not null;default:0"`
+ TaskQuota int64 `gorm:"not null;default:0"`
+ TaskQuotaTarget int64 `gorm:"not null;default:0"`
+ Status string `gorm:"type:varchar(16);index;not null;default:'applied'"`
+ Attempts int `gorm:"not null;default:0"`
+ LastError string `gorm:"type:text"`
+ NextAttempt int64 `gorm:"index;not null;default:0"`
+ Revision int64 `gorm:"not null;default:1"`
+ SubscriptionPreConsumeRequestID string `gorm:"type:varchar(64);index"`
+ FinalizeSubscriptionPreConsume bool `gorm:"not null;default:false"`
+ AllowMissingToken bool `gorm:"not null;default:false"`
+ ManualOnFailure bool `gorm:"not null;default:false"`
+ PreConsumeRequestID string `gorm:"type:varchar(64);index"`
+ PreConsumeModelName string `gorm:"type:varchar(191);not null;default:''"`
+ PreConsumeRequestedQuota int64 `gorm:"not null;default:0"`
+ PreConsumeEffectiveQuota int64 `gorm:"not null;default:0"`
+ EffectPayload string `gorm:"type:text"`
+ EffectStatus string `gorm:"type:varchar(16);index"`
+ CreatedAt int64 `gorm:"not null"`
+ UpdatedAt int64 `gorm:"index;not null"`
+}
+
+// BillingPreConsumeSelection is the immutable funding-source decision for one
+// request. Wallet and subscription mutations claim it inside their own balance
+// transaction so concurrent attempts cannot commit against different sources.
+type BillingPreConsumeSelection struct {
+ ID int64 `gorm:"primaryKey"`
+ RequestID string `gorm:"type:varchar(64);uniqueIndex;not null"`
+ Source string `gorm:"type:varchar(32);not null"`
+ UserID int `gorm:"not null"`
+ TokenID int `gorm:"not null;default:0"`
+ ModelName string `gorm:"type:varchar(191);not null;default:''"`
+ RequestedQuota int64 `gorm:"not null;default:0"`
+ EffectiveQuota int64 `gorm:"not null;default:0"`
+ CreatedAt int64 `gorm:"not null"`
+}
+
+type BillingSettlementEffect struct {
+ LogType int `json:"log_type"`
+ Content string `json:"content"`
+ ChannelID int `json:"channel_id"`
+ ModelName string `json:"model_name"`
+ TokenID int `json:"token_id"`
+ Group string `json:"group"`
+ Other map[string]interface{} `json:"other"`
+ NodeName string `json:"node_name"`
+ UpdateUsage bool `json:"update_usage"`
+ Quota int64 `json:"quota,omitempty"`
+}
+
+type BillingSettlementInput struct {
+ OperationKey string
+ Source string
+ UserID int
+ SubscriptionID int
+ TokenID int
+ TokenKey string
+ FundingDelta int64
+ TokenDelta int64
+ TaskID int64
+ TaskQuota int64
+ TaskQuotaTarget int64
+ SubscriptionPreConsumeRequestID string
+ FinalizeSubscriptionPreConsume bool
+ // AllowMissingToken is only for terminal task refunds. A token may be
+ // deleted while its asynchronous task is still pending; that must not
+ // prevent the user funding balance from being refunded.
+ AllowMissingToken bool
+ ManualOnFailure bool
+ Effect *BillingSettlementEffect
+ PreConsumeRequestID string
+ PreConsumeModelName string
+ PreConsumeRequestedQuota int64
+ PreConsumeEffectiveQuota int64
+}
+
+var billingSettlementRunnerOnce sync.Once
+var billingSettlementRunnerState struct {
+ sync.Mutex
+ cancel context.CancelFunc
+ done chan struct{}
+}
+
+// StartBillingSettlementTaskRunner retries durable settlement intents after a
+// process restart or a transient database failure. All balance mutations are
+// still protected by ApplyBillingSettlementOnce's operation key.
+func StartBillingSettlementTaskRunner() {
+ if DB == nil {
+ return
+ }
+ billingSettlementRunnerOnce.Do(func() {
+ ctx, cancel := context.WithCancel(context.Background())
+ done := make(chan struct{})
+ billingSettlementRunnerState.Lock()
+ billingSettlementRunnerState.cancel = cancel
+ billingSettlementRunnerState.done = done
+ billingSettlementRunnerState.Unlock()
+ gopool.Go(func() {
+ defer close(done)
+ ProcessPendingBillingSettlementsOnce()
+ ticker := time.NewTicker(time.Second)
+ defer ticker.Stop()
+ for {
+ select {
+ case <-ctx.Done():
+ return
+ case <-ticker.C:
+ ProcessPendingBillingSettlementsOnce()
+ }
+ }
+ })
+ })
+}
+
+func StopBillingSettlementTaskRunner(ctx context.Context) error {
+ billingSettlementRunnerState.Lock()
+ cancel := billingSettlementRunnerState.cancel
+ done := billingSettlementRunnerState.done
+ billingSettlementRunnerState.Unlock()
+ if cancel == nil || done == nil {
+ return nil
+ }
+ cancel()
+ select {
+ case <-done:
+ return nil
+ case <-ctx.Done():
+ return ctx.Err()
+ }
+}
+
+func ProcessPendingBillingSettlementsOnce() {
+ processPendingBillingSettlements()
+}
+
+func ApplyBillingSettlementOnce(input BillingSettlementInput) (appliedFundingDelta int64, alreadyApplied bool, err error) {
+ if input.OperationKey == "" {
+ return 0, false, errors.New("billing settlement operation key is required")
+ }
+
+ record, alreadyApplied, err := ensureBillingSettlementRecord(input)
+ if err != nil {
+ return 0, false, err
+ }
+ resolvedTokenKey := input.TokenKey
+ effectiveTokenID := input.TokenID
+ effectiveTokenDelta := input.TokenDelta
+ appliedTokenDelta := int64(0)
+ var cacheTasks []CacheInvalidationTask
+ if alreadyApplied {
+ if input.PreConsumeRequestID != "" {
+ if claimErr := DB.Transaction(func(tx *gorm.DB) error {
+ return claimBillingPreConsumeSourceTx(tx, billingPreConsumeSelectionFromSettlement(input))
+ }); claimErr != nil {
+ return 0, false, claimErr
+ }
+ }
+ if resolvedTokenKey == "" {
+ resolvedTokenKey = lookupBillingTokenKey(record.TokenID)
+ }
+ invalidateBillingSettlementCaches(record, resolvedTokenKey)
+ return record.AppliedFundingDelta, true, nil
+ }
+
+ err = DB.Transaction(func(tx *gorm.DB) error {
+ var current BillingSettlement
+ if err := withRowLock(tx).Where("id = ?", record.ID).First(¤t).Error; err != nil {
+ return err
+ }
+ if err := validateBillingSettlement(current, input); err != nil {
+ return err
+ }
+ if current.Status == "" || current.Status == BillingSettlementStatusApplied {
+ appliedFundingDelta = current.AppliedFundingDelta
+ alreadyApplied = true
+ return nil
+ }
+ if current.Status == BillingSettlementStatusManual {
+ return permanentBillingSettlement(fmt.Errorf("%w: %s", ErrBillingSettlementManualReview, current.LastError))
+ }
+ if current.Status != BillingSettlementStatusPending {
+ return permanentBillingSettlement(fmt.Errorf("unknown billing settlement status %q", current.Status))
+ }
+ if err := validateBillingSettlementInput(input); err != nil {
+ return err
+ }
+ if input.PreConsumeRequestID != "" {
+ if err := claimBillingPreConsumeSourceTx(tx, billingPreConsumeSelectionFromSettlement(input)); err != nil {
+ return err
+ }
+ }
+
+ if effectiveTokenID > 0 && effectiveTokenDelta != 0 {
+ var token Token
+ if lookupErr := tx.Select("key", "user_id").First(&token, effectiveTokenID).Error; lookupErr != nil {
+ if input.AllowMissingToken && errors.Is(lookupErr, gorm.ErrRecordNotFound) {
+ effectiveTokenID = 0
+ effectiveTokenDelta = 0
+ } else {
+ return lookupErr
+ }
+ } else {
+ if token.UserId != input.UserID {
+ return permanentBillingSettlement(fmt.Errorf("billing settlement token user identity mismatch: token_id=%d", effectiveTokenID))
+ }
+ resolvedTokenKey = token.Key
+ }
+ }
+
+ var applyErr error
+ appliedFundingDelta, applyErr = applyFundingDeltaTx(tx, input)
+ if applyErr != nil {
+ return applyErr
+ }
+ if input.Source == BillingSettlementSourceSubscription && input.TokenDelta == input.FundingDelta {
+ effectiveTokenDelta = appliedFundingDelta
+ }
+ if err := applyTokenDeltaTx(tx, input.UserID, effectiveTokenID, effectiveTokenDelta, input.AllowMissingToken); err != nil {
+ return err
+ }
+ if effectiveTokenID > 0 {
+ appliedTokenDelta = effectiveTokenDelta
+ }
+
+ if input.TaskID > 0 {
+ appliedTaskQuotaTarget := input.TaskQuota + appliedFundingDelta
+ result := tx.Model(&Task{}).
+ Where("id = ? AND quota = ?", input.TaskID, input.TaskQuota).
+ Update("quota", appliedTaskQuotaTarget)
+ if result.Error != nil {
+ return result.Error
+ }
+ if result.RowsAffected != 1 {
+ return permanentBillingSettlement(fmt.Errorf("%w: task_id=%d expected_quota=%d", ErrBillingSettlementTaskConflict, input.TaskID, input.TaskQuota))
+ }
+ }
+
+ result := tx.Model(&BillingSettlement{}).
+ Where("id = ? AND status = ?", record.ID, BillingSettlementStatusPending).
+ Updates(map[string]interface{}{
+ "applied_funding_delta": appliedFundingDelta,
+ "applied_token_delta": appliedTokenDelta,
+ "status": BillingSettlementStatusApplied,
+ "effect_status": effectStatusAfterSettlement(current.EffectPayload),
+ "last_error": "",
+ "next_attempt": 0,
+ "updated_at": time.Now().Unix(),
+ "revision": gorm.Expr("revision + ?", 1),
+ })
+ if result.Error != nil {
+ return result.Error
+ }
+ if result.RowsAffected != 1 {
+ return permanentBillingSettlement(fmt.Errorf("%w: %s", ErrBillingSettlementOperationConflict, input.OperationKey))
+ }
+ if input.UserID > 0 {
+ userTask, err := stageUserCacheInvalidationTx(tx, input.UserID, false)
+ if err != nil {
+ return err
+ }
+ cacheTasks = append(cacheTasks, userTask)
+ }
+ if effectiveTokenID > 0 && effectiveTokenDelta != 0 && resolvedTokenKey != "" {
+ tokenTask, err := stageTokenCacheInvalidationTx(tx, resolvedTokenKey, false)
+ if err != nil {
+ return err
+ }
+ cacheTasks = append(cacheTasks, tokenTask)
+ }
+ return nil
+ })
+ if err != nil {
+ markBillingSettlementFailure(input.OperationKey, err)
+ return 0, false, err
+ }
+
+ dispatchStagedCacheInvalidations(cacheTasks)
+ return appliedFundingDelta, alreadyApplied, nil
+}
+
+// ResolveBillingPreConsumeSource returns the funding source already selected by
+// a request's durable pre-consume record. A request must never move between a
+// wallet settlement and a subscription pre-consume on replay.
+func ResolveBillingPreConsumeSource(requestID string) (string, bool, error) {
+ if requestID == "" {
+ return "", false, nil
+ }
+ if DB == nil {
+ return "", false, errors.New("database is not initialized")
+ }
+
+ var selection BillingPreConsumeSelection
+ selectionQuery := DB.Select("source").
+ Where("request_id = ?", requestID).
+ Limit(1).
+ Find(&selection)
+ if selectionQuery.Error != nil {
+ return "", false, selectionQuery.Error
+ }
+ if selectionQuery.RowsAffected > 0 {
+ if selection.Source != BillingSettlementSourceWallet && selection.Source != BillingSettlementSourceSubscription {
+ return "", false, fmt.Errorf("%w: request %s has invalid pre-consume source %q", ErrBillingSettlementOperationConflict, requestID, selection.Source)
+ }
+ return selection.Source, true, nil
+ }
+
+ // Backward-compatible discovery for operations created before the selection
+ // table existed. The next replay backfills the immutable selection in the
+ // same transaction as its idempotent balance operation.
+ var walletRecord BillingSettlement
+ walletQuery := DB.Select("source").
+ Where("operation_key = ?", "request:"+requestID+":pre-consume").
+ Limit(1).
+ Find(&walletRecord)
+ if walletQuery.Error != nil {
+ return "", false, walletQuery.Error
+ }
+
+ var subscriptionRecord SubscriptionPreConsumeRecord
+ subscriptionQuery := DB.Select("id").
+ Where("request_id = ?", requestID).
+ Limit(1).
+ Find(&subscriptionRecord)
+ if subscriptionQuery.Error != nil {
+ return "", false, subscriptionQuery.Error
+ }
+
+ hasWallet := walletQuery.RowsAffected > 0
+ hasSubscription := subscriptionQuery.RowsAffected > 0
+ if hasWallet && hasSubscription {
+ return "", false, fmt.Errorf("%w: request %s has multiple pre-consume funding sources", ErrBillingSettlementOperationConflict, requestID)
+ }
+ if hasWallet {
+ if walletRecord.Source != BillingSettlementSourceWallet {
+ return "", false, fmt.Errorf("%w: request %s has invalid wallet pre-consume source %q", ErrBillingSettlementOperationConflict, requestID, walletRecord.Source)
+ }
+ return BillingSettlementSourceWallet, true, nil
+ }
+ if hasSubscription {
+ return BillingSettlementSourceSubscription, true, nil
+ }
+ return "", false, nil
+}
+
+func billingPreConsumeSelectionFromSettlement(input BillingSettlementInput) BillingPreConsumeSelection {
+ return BillingPreConsumeSelection{
+ RequestID: input.PreConsumeRequestID,
+ Source: input.Source,
+ UserID: input.UserID,
+ TokenID: input.TokenID,
+ ModelName: input.PreConsumeModelName,
+ RequestedQuota: input.PreConsumeRequestedQuota,
+ EffectiveQuota: input.PreConsumeEffectiveQuota,
+ }
+}
+
+func claimBillingPreConsumeSourceTx(tx *gorm.DB, selection BillingPreConsumeSelection) error {
+ if tx == nil {
+ return errors.New("billing pre-consume selection transaction is required")
+ }
+ if selection.RequestID == "" || selection.UserID <= 0 || selection.RequestedQuota < 0 || selection.EffectiveQuota < 0 {
+ return fmt.Errorf("%w: invalid billing pre-consume selection", ErrBillingSettlementOperationConflict)
+ }
+ if selection.Source != BillingSettlementSourceWallet && selection.Source != BillingSettlementSourceSubscription {
+ return fmt.Errorf("%w: invalid billing pre-consume source %q", ErrBillingSettlementOperationConflict, selection.Source)
+ }
+ selection.CreatedAt = time.Now().Unix()
+ if err := tx.Clauses(clause.OnConflict{
+ Columns: []clause.Column{{Name: "request_id"}},
+ DoNothing: true,
+ }).Create(&selection).Error; err != nil {
+ return err
+ }
+
+ var existing BillingPreConsumeSelection
+ if err := tx.Where("request_id = ?", selection.RequestID).First(&existing).Error; err != nil {
+ return err
+ }
+ if existing.Source != selection.Source ||
+ existing.UserID != selection.UserID ||
+ existing.TokenID != selection.TokenID ||
+ existing.ModelName != selection.ModelName ||
+ existing.RequestedQuota != selection.RequestedQuota ||
+ existing.EffectiveQuota != selection.EffectiveQuota {
+ return fmt.Errorf("%w: request %s reused with different pre-consume source or parameters", ErrBillingSettlementOperationConflict, selection.RequestID)
+ }
+ return nil
+}
+
+func ensureBillingSettlementRecord(input BillingSettlementInput) (BillingSettlement, bool, error) {
+ return ensureBillingSettlementRecordDB(DB, input)
+}
+
+func ensureBillingSettlementRecordDB(db *gorm.DB, input BillingSettlementInput) (BillingSettlement, bool, error) {
+ if db == nil {
+ return BillingSettlement{}, false, errors.New("database is not initialized")
+ }
+ effectPayload, err := billingSettlementEffectPayload(input.Effect)
+ if err != nil {
+ return BillingSettlement{}, false, err
+ }
+ now := time.Now().Unix()
+ record := BillingSettlement{
+ OperationKey: input.OperationKey, Source: input.Source, UserID: input.UserID,
+ SubscriptionID: input.SubscriptionID, TokenID: input.TokenID,
+ FundingDelta: input.FundingDelta, TokenDelta: input.TokenDelta,
+ TaskID: input.TaskID, TaskQuota: input.TaskQuota, TaskQuotaTarget: input.TaskQuotaTarget,
+ SubscriptionPreConsumeRequestID: input.SubscriptionPreConsumeRequestID,
+ AllowMissingToken: input.AllowMissingToken,
+ FinalizeSubscriptionPreConsume: input.FinalizeSubscriptionPreConsume,
+ ManualOnFailure: input.ManualOnFailure,
+ PreConsumeRequestID: input.PreConsumeRequestID,
+ PreConsumeModelName: input.PreConsumeModelName,
+ PreConsumeRequestedQuota: input.PreConsumeRequestedQuota,
+ PreConsumeEffectiveQuota: input.PreConsumeEffectiveQuota,
+ EffectPayload: effectPayload,
+ Status: BillingSettlementStatusPending, NextAttempt: now, CreatedAt: now, UpdatedAt: now, Revision: 1,
+ }
+ result := db.Clauses(clause.OnConflict{
+ Columns: []clause.Column{{Name: "operation_key"}},
+ DoNothing: true,
+ }).Create(&record)
+ if result.Error != nil {
+ return BillingSettlement{}, false, result.Error
+ }
+ if err := db.Where("operation_key = ?", input.OperationKey).First(&record).Error; err != nil {
+ return BillingSettlement{}, false, err
+ }
+ if err := validateBillingSettlement(record, input); err != nil {
+ return BillingSettlement{}, false, err
+ }
+ if record.Status == BillingSettlementStatusManual {
+ return BillingSettlement{}, false, fmt.Errorf("%w: %s", ErrBillingSettlementManualReview, record.LastError)
+ }
+ return record, record.Status == "" || record.Status == BillingSettlementStatusApplied, nil
+}
+
+func billingSettlementEffectPayload(effect *BillingSettlementEffect) (string, error) {
+ if effect == nil {
+ return "", nil
+ }
+ data, err := common.Marshal(effect)
+ if err != nil {
+ return "", err
+ }
+ return string(data), nil
+}
+
+func decodeBillingSettlementEffect(payload string) (*BillingSettlementEffect, error) {
+ if payload == "" {
+ return nil, nil
+ }
+ var effect BillingSettlementEffect
+ if err := common.UnmarshalJsonStr(payload, &effect); err != nil {
+ return nil, err
+ }
+ return &effect, nil
+}
+
+func effectStatusAfterSettlement(payload string) string {
+ if payload == "" {
+ return ""
+ }
+ return BillingSettlementEffectPending
+}
+
+func lookupBillingTokenKey(tokenID int) string {
+ if tokenID <= 0 || DB == nil {
+ return ""
+ }
+ var token Token
+ if err := DB.Select("key").First(&token, tokenID).Error; err != nil {
+ return ""
+ }
+ return token.Key
+}
+
+func invalidateBillingSettlementCaches(record BillingSettlement, tokenKey string) {
+ if record.UserID > 0 {
+ if cacheErr := invalidateUserQuotaCache(record.UserID); cacheErr != nil {
+ enqueueUserCacheInvalidationRetry(record.UserID, cacheErr)
+ }
+ }
+ if record.TokenID > 0 && record.AppliedTokenDelta != 0 && tokenKey != "" {
+ invalidateTokenQuotaCache(tokenKey)
+ }
+}
+
+func markBillingSettlementFailure(operationKey string, cause error) {
+ if DB == nil || operationKey == "" {
+ return
+ }
+ var record BillingSettlement
+ if err := DB.Where("operation_key = ?", operationKey).First(&record).Error; err != nil {
+ common.SysLog("failed to load billing settlement retry record: " + err.Error())
+ return
+ }
+ attempts := record.Attempts + 1
+ delay := 1 << min(attempts, 6)
+ status := BillingSettlementStatusPending
+ nextAttempt := time.Now().Add(time.Duration(delay) * time.Second).Unix()
+ if record.ManualOnFailure || isPermanentBillingSettlementError(cause) {
+ status = BillingSettlementStatusManual
+ nextAttempt = 0
+ }
+ lastError := ""
+ if cause != nil {
+ lastError = cause.Error()
+ }
+ if err := DB.Model(&BillingSettlement{}).Where("id = ? AND status = ?", record.ID, BillingSettlementStatusPending).Updates(map[string]interface{}{
+ "attempts": attempts, "last_error": lastError,
+ "status": status, "next_attempt": nextAttempt,
+ "updated_at": time.Now().Unix(), "revision": gorm.Expr("revision + ?", 1),
+ }).Error; err != nil {
+ common.SysLog("failed to reschedule billing settlement: " + err.Error())
+ }
+}
+
+func processPendingBillingSettlements() {
+ if DB == nil {
+ return
+ }
+ var records []BillingSettlement
+ if err := DB.Where("status = ? AND next_attempt <= ?", BillingSettlementStatusPending, time.Now().Unix()).Order("id").Limit(100).Find(&records).Error; err != nil {
+ common.SysLog("failed to load billing settlement retries: " + err.Error())
+ return
+ }
+ for _, record := range records {
+ effect, decodeErr := decodeBillingSettlementEffect(record.EffectPayload)
+ if decodeErr != nil {
+ markBillingSettlementFailure(record.OperationKey, permanentBillingSettlement(decodeErr))
+ continue
+ }
+ _, _, err := ApplyBillingSettlementOnce(BillingSettlementInput{
+ OperationKey: record.OperationKey, Source: record.Source, UserID: record.UserID,
+ SubscriptionID: record.SubscriptionID, TokenID: record.TokenID,
+ FundingDelta: record.FundingDelta, TokenDelta: record.TokenDelta,
+ TaskID: record.TaskID, TaskQuota: record.TaskQuota, TaskQuotaTarget: record.TaskQuotaTarget,
+ SubscriptionPreConsumeRequestID: record.SubscriptionPreConsumeRequestID,
+ FinalizeSubscriptionPreConsume: record.FinalizeSubscriptionPreConsume,
+ AllowMissingToken: record.AllowMissingToken,
+ ManualOnFailure: record.ManualOnFailure,
+ PreConsumeRequestID: record.PreConsumeRequestID,
+ PreConsumeModelName: record.PreConsumeModelName,
+ PreConsumeRequestedQuota: record.PreConsumeRequestedQuota,
+ PreConsumeEffectiveQuota: record.PreConsumeEffectiveQuota,
+ Effect: effect,
+ })
+ if err != nil {
+ common.SysLog("billing settlement retry remains pending: " + err.Error())
+ continue
+ }
+ if effectErr := ProcessBillingSettlementEffect(record.OperationKey); effectErr != nil {
+ common.SysLog("billing settlement effect remains pending: " + effectErr.Error())
+ }
+ }
+ processPendingBillingSettlementEffects()
+}
+
+func validateBillingSettlement(existing BillingSettlement, input BillingSettlementInput) error {
+ // Effect payload is non-financial metadata. The first persisted payload wins;
+ // replays must not overwrite it or turn harmless description drift into a balance conflict.
+ if existing.Source != input.Source ||
+ existing.UserID != input.UserID ||
+ existing.SubscriptionID != input.SubscriptionID ||
+ existing.TokenID != input.TokenID ||
+ existing.FundingDelta != input.FundingDelta ||
+ existing.TokenDelta != input.TokenDelta ||
+ existing.TaskID != input.TaskID ||
+ existing.TaskQuota != input.TaskQuota ||
+ existing.TaskQuotaTarget != input.TaskQuotaTarget ||
+ existing.SubscriptionPreConsumeRequestID != input.SubscriptionPreConsumeRequestID ||
+ existing.FinalizeSubscriptionPreConsume != input.FinalizeSubscriptionPreConsume ||
+ existing.AllowMissingToken != input.AllowMissingToken ||
+ existing.ManualOnFailure != input.ManualOnFailure {
+ return permanentBillingSettlement(fmt.Errorf("%w: operation key reused with different parameters: %s", ErrBillingSettlementOperationConflict, input.OperationKey))
+ }
+ if billingSettlementHasPreConsumeFields(existing) &&
+ (existing.PreConsumeRequestID != input.PreConsumeRequestID ||
+ existing.PreConsumeModelName != input.PreConsumeModelName ||
+ existing.PreConsumeRequestedQuota != input.PreConsumeRequestedQuota ||
+ existing.PreConsumeEffectiveQuota != input.PreConsumeEffectiveQuota) {
+ return permanentBillingSettlement(fmt.Errorf("%w: operation key reused with different pre-consume parameters: %s", ErrBillingSettlementOperationConflict, input.OperationKey))
+ }
+ return nil
+}
+
+func billingSettlementHasPreConsumeFields(record BillingSettlement) bool {
+ return record.PreConsumeRequestID != "" ||
+ record.PreConsumeModelName != "" ||
+ record.PreConsumeRequestedQuota != 0 ||
+ record.PreConsumeEffectiveQuota != 0
+}
+
+func processPendingBillingSettlementEffects() {
+ if DB == nil {
+ return
+ }
+ var records []BillingSettlement
+ if err := DB.Where("status = ? AND effect_status = ?", BillingSettlementStatusApplied, BillingSettlementEffectPending).
+ Order("id").Limit(100).Find(&records).Error; err != nil {
+ common.SysLog("failed to load billing settlement effects: " + err.Error())
+ return
+ }
+ for _, record := range records {
+ if err := ProcessBillingSettlementEffect(record.OperationKey); err != nil {
+ common.SysLog("billing settlement effect remains pending: " + err.Error())
+ }
+ }
+}
+
+func ProcessBillingSettlementEffect(operationKey string) error {
+ if DB == nil || operationKey == "" {
+ return errors.New("billing settlement effect operation key is required")
+ }
+ var record BillingSettlement
+ if err := DB.Where("operation_key = ?", operationKey).First(&record).Error; err != nil {
+ return err
+ }
+ if record.EffectPayload == "" || record.EffectStatus == BillingSettlementEffectApplied {
+ return nil
+ }
+ if record.Status != BillingSettlementStatusApplied || record.EffectStatus != BillingSettlementEffectPending {
+ return fmt.Errorf("billing settlement effect is not ready: operation=%s status=%s effect_status=%s", operationKey, record.Status, record.EffectStatus)
+ }
+ effect, err := decodeBillingSettlementEffect(record.EffectPayload)
+ if err != nil {
+ return err
+ }
+ if effect == nil {
+ return nil
+ }
+
+ appliedDelta := record.AppliedFundingDelta
+ logQuota := appliedDelta
+ if effect.Quota > 0 {
+ logQuota = effect.Quota
+ }
+ if logQuota != 0 {
+ other := effect.Other
+ if other == nil {
+ other = make(map[string]interface{})
+ }
+ if record.TaskID > 0 {
+ other["actual_quota"] = record.TaskQuota + appliedDelta
+ }
+ quota := logQuota
+ if quota < 0 {
+ quota = -quota
+ }
+ if err := RecordTaskBillingLogOnce(operationKey, RecordTaskBillingLogParams{
+ UserId: record.UserID, LogType: effect.LogType, Content: effect.Content,
+ ChannelId: effect.ChannelID, ModelName: effect.ModelName, Quota: int(quota),
+ TokenId: effect.TokenID, Group: effect.Group, Other: other, NodeName: effect.NodeName,
+ }); err != nil {
+ return err
+ }
+ }
+
+ return DB.Transaction(func(tx *gorm.DB) error {
+ claim := tx.Model(&BillingSettlement{}).
+ Where("id = ? AND status = ? AND effect_status = ?", record.ID, BillingSettlementStatusApplied, BillingSettlementEffectPending).
+ Update("effect_status", BillingSettlementEffectApplying)
+ if claim.Error != nil {
+ return claim.Error
+ }
+ if claim.RowsAffected == 0 {
+ return nil
+ }
+ usageQuota := appliedDelta
+ if effect.Quota > 0 {
+ usageQuota = effect.Quota
+ }
+ if effect.UpdateUsage && usageQuota > 0 {
+ userResult := tx.Model(&User{}).Where("id = ?", record.UserID).Updates(map[string]interface{}{
+ "used_quota": gorm.Expr("used_quota + ?", usageQuota),
+ "request_count": gorm.Expr("request_count + ?", 1),
+ })
+ if userResult.Error != nil {
+ return userResult.Error
+ }
+ if userResult.RowsAffected != 1 {
+ return fmt.Errorf("billing settlement effect user not found: user_id=%d", record.UserID)
+ }
+ if effect.ChannelID > 0 {
+ if err := tx.Model(&Channel{}).Where("id = ?", effect.ChannelID).
+ Update("used_quota", gorm.Expr("used_quota + ?", usageQuota)).Error; err != nil {
+ return err
+ }
+ }
+ }
+ result := tx.Model(&BillingSettlement{}).
+ Where("id = ? AND effect_status = ?", record.ID, BillingSettlementEffectApplying).
+ Update("effect_status", BillingSettlementEffectApplied)
+ if result.Error != nil {
+ return result.Error
+ }
+ if result.RowsAffected != 1 {
+ return fmt.Errorf("billing settlement effect state conflict: operation=%s", operationKey)
+ }
+ return nil
+ })
+}
+
+func validateBillingSettlementInput(input BillingSettlementInput) error {
+ if input.Source != BillingSettlementSourceWallet && input.Source != BillingSettlementSourceSubscription {
+ return permanentBillingSettlement(fmt.Errorf("unsupported billing settlement source: %s", input.Source))
+ }
+ if input.TaskID > 0 && input.TaskQuotaTarget != input.TaskQuota+input.FundingDelta {
+ return permanentBillingSettlement(fmt.Errorf("task quota target mismatch: task_id=%d current=%d target=%d funding_delta=%d", input.TaskID, input.TaskQuota, input.TaskQuotaTarget, input.FundingDelta))
+ }
+ if input.Source == BillingSettlementSourceSubscription && input.SubscriptionPreConsumeRequestID == "" {
+ return permanentBillingSettlement(ErrSubscriptionSettlementUnbound)
+ }
+ if input.FinalizeSubscriptionPreConsume && (input.Source != BillingSettlementSourceSubscription || input.FundingDelta >= 0) {
+ return permanentBillingSettlement(errors.New("subscription pre-consume finalization requires a negative subscription delta"))
+ }
+ return nil
+}
+
+func applyFundingDeltaTx(tx *gorm.DB, input BillingSettlementInput) (int64, error) {
+ switch input.Source {
+ case BillingSettlementSourceWallet:
+ if input.UserID <= 0 {
+ return 0, permanentBillingSettlement(errors.New("wallet settlement user id is required"))
+ }
+ if input.FundingDelta > 0 {
+ result := tx.Model(&User{}).
+ Where("id = ? AND quota >= ?", input.UserID, input.FundingDelta).
+ Update("quota", gorm.Expr("quota - ?", input.FundingDelta))
+ if result.Error != nil {
+ return 0, result.Error
+ }
+ if result.RowsAffected != 1 {
+ return 0, permanentBillingSettlement(fmt.Errorf("%w: id=%d, need=%d", ErrUserQuotaInsufficient, input.UserID, input.FundingDelta))
+ }
+ } else if input.FundingDelta < 0 {
+ result := tx.Model(&User{}).
+ Where("id = ?", input.UserID).
+ Update("quota", gorm.Expr("quota + ?", -input.FundingDelta))
+ if result.Error != nil {
+ return 0, result.Error
+ }
+ if result.RowsAffected != 1 {
+ return 0, permanentBillingSettlement(fmt.Errorf("%w: id=%d", ErrUserNotFound, input.UserID))
+ }
+ }
+ return input.FundingDelta, nil
+ case BillingSettlementSourceSubscription:
+ return applySubscriptionDeltaTx(tx, input)
+ default:
+ return 0, permanentBillingSettlement(fmt.Errorf("unsupported billing settlement source: %s", input.Source))
+ }
+}
+
+func applyTokenDeltaTx(tx *gorm.DB, userID int, tokenID int, delta int64, allowMissing bool) error {
+ if tokenID <= 0 || delta == 0 {
+ return nil
+ }
+ if delta > 0 {
+ result := tx.Model(&Token{}).
+ Where("id = ? AND user_id = ? AND (unlimited_quota = ? OR remain_quota >= ?)", tokenID, userID, true, delta).
+ Updates(map[string]interface{}{
+ "remain_quota": gorm.Expr("remain_quota - ?", delta),
+ "used_quota": gorm.Expr("used_quota + ?", delta),
+ "accessed_time": common.GetTimestamp(),
+ })
+ if result.Error != nil {
+ return result.Error
+ }
+ if result.RowsAffected != 1 {
+ var token Token
+ if err := tx.Select("id").First(&token, tokenID).Error; errors.Is(err, gorm.ErrRecordNotFound) {
+ return err
+ }
+ return permanentBillingSettlement(fmt.Errorf("%w: id=%d, need=%d", ErrTokenQuotaInsufficient, tokenID, delta))
+ }
+ return nil
+ }
+
+ result := tx.Model(&Token{}).
+ Where("id = ? AND user_id = ? AND used_quota >= ?", tokenID, userID, -delta).
+ Updates(map[string]interface{}{
+ "remain_quota": gorm.Expr("remain_quota + ?", -delta),
+ "used_quota": gorm.Expr("used_quota - ?", -delta),
+ "accessed_time": common.GetTimestamp(),
+ })
+ if result.Error != nil {
+ return result.Error
+ }
+ if result.RowsAffected != 1 {
+ if allowMissing && delta < 0 {
+ var token Token
+ if err := tx.Select("id").First(&token, tokenID).Error; errors.Is(err, gorm.ErrRecordNotFound) {
+ return nil
+ }
+ }
+ return permanentBillingSettlement(fmt.Errorf("%w: id=%d", ErrTokenQuotaInsufficient, tokenID))
+ }
+ return nil
+}
+
+func applySubscriptionDeltaTx(tx *gorm.DB, input BillingSettlementInput) (int64, error) {
+ requestID := input.SubscriptionPreConsumeRequestID
+ if requestID == "" {
+ return 0, permanentBillingSettlement(ErrSubscriptionSettlementUnbound)
+ }
+ var record SubscriptionPreConsumeRecord
+ if err := withRowLock(tx).Where("request_id = ?", requestID).First(&record).Error; err != nil {
+ if errors.Is(err, gorm.ErrRecordNotFound) {
+ return 0, permanentBillingSettlement(fmt.Errorf("%w: request=%s", ErrSubscriptionSettlementUnbound, requestID))
+ }
+ return 0, err
+ }
+ if input.SubscriptionID <= 0 || record.UserSubscriptionId != input.SubscriptionID {
+ return 0, permanentBillingSettlement(fmt.Errorf("subscription settlement identity mismatch: request=%s", requestID))
+ }
+ if record.UserId != input.UserID || (input.TokenDelta != 0 && record.TokenId != input.TokenID) {
+ return 0, permanentBillingSettlement(fmt.Errorf("subscription settlement owner identity mismatch: request=%s", requestID))
+ }
+ if record.Status == "refunded" {
+ return 0, permanentBillingSettlement(fmt.Errorf("subscription pre-consume already refunded: request=%s", requestID))
+ }
+ var sub UserSubscription
+ if err := withRowLock(tx).Where("id = ?", record.UserSubscriptionId).First(&sub).Error; err != nil {
+ if errors.Is(err, gorm.ErrRecordNotFound) {
+ return 0, permanentBillingSettlement(fmt.Errorf("subscription settlement target no longer exists: request=%s", requestID))
+ }
+ return 0, err
+ }
+ if sub.UserId != input.UserID {
+ return 0, permanentBillingSettlement(fmt.Errorf("subscription settlement user identity mismatch: request=%s", requestID))
+ }
+ if err := validateSubscriptionPreConsumePeriod(record, sub); err != nil {
+ return 0, permanentBillingSettlement(err)
+ }
+ if input.FinalizeSubscriptionPreConsume {
+ if input.FundingDelta >= 0 || -input.FundingDelta < record.PreConsumed {
+ return 0, permanentBillingSettlement(fmt.Errorf("subscription refund does not cover pre-consumed amount: request=%s", requestID))
+ }
+ }
+ applied, err := postConsumeUserSubscriptionDeltaTx(tx, record.UserSubscriptionId, input.FundingDelta)
+ if err != nil {
+ if errors.Is(err, ErrSubscriptionQuotaInsufficient) {
+ return 0, permanentBillingSettlement(err)
+ }
+ return 0, err
+ }
+ if input.FinalizeSubscriptionPreConsume {
+ if applied != input.FundingDelta {
+ return 0, permanentBillingSettlement(fmt.Errorf("subscription refund was clamped: request=%s requested=%d applied=%d", requestID, input.FundingDelta, applied))
+ }
+ result := tx.Model(&SubscriptionPreConsumeRecord{}).
+ Where("id = ? AND status = ?", record.Id, "consumed").
+ Updates(map[string]interface{}{"status": "refunded", "updated_at": common.GetTimestamp()})
+ if result.Error != nil {
+ return 0, result.Error
+ }
+ if result.RowsAffected != 1 {
+ return 0, permanentBillingSettlement(errors.New("subscription refund claim lost"))
+ }
+ }
+ return applied, nil
+}
diff --git a/model/billing_settlement_regression_test.go b/model/billing_settlement_regression_test.go
new file mode 100644
index 00000000..51397d16
--- /dev/null
+++ b/model/billing_settlement_regression_test.go
@@ -0,0 +1,621 @@
+package model
+
+import (
+ "errors"
+ "testing"
+ "time"
+
+ "github.com/MAX-API-Next/MAX-API/common"
+ "github.com/stretchr/testify/assert"
+ "github.com/stretchr/testify/require"
+ "gorm.io/gorm"
+)
+
+func TestConcurrentBillingPreConsumeSelectionAllowsOnlyOneFundingSource(t *testing.T) {
+ setupUserUpdateTestState(t)
+ const requestID = "regression:concurrent-funding-source"
+
+ start := make(chan struct{})
+ results := make(chan error, 2)
+ for _, source := range []string{BillingSettlementSourceWallet, BillingSettlementSourceSubscription} {
+ source := source
+ go func() {
+ <-start
+ results <- DB.Transaction(func(tx *gorm.DB) error {
+ return claimBillingPreConsumeSourceTx(tx, BillingPreConsumeSelection{
+ RequestID: requestID, Source: source, UserID: 940, TokenID: 941,
+ ModelName: "concurrent-model", RequestedQuota: 10, EffectiveQuota: 10,
+ })
+ })
+ }()
+ }
+ close(start)
+
+ var succeeded, conflicted int
+ for range 2 {
+ err := <-results
+ switch {
+ case err == nil:
+ succeeded++
+ case errors.Is(err, ErrBillingSettlementOperationConflict):
+ conflicted++
+ default:
+ require.NoError(t, err)
+ }
+ }
+ assert.Equal(t, 1, succeeded)
+ assert.Equal(t, 1, conflicted)
+
+ source, found, err := ResolveBillingPreConsumeSource(requestID)
+ require.NoError(t, err)
+ require.True(t, found)
+ assert.Contains(t, []string{BillingSettlementSourceWallet, BillingSettlementSourceSubscription}, source)
+}
+
+func TestSubscriptionSettlementUsesActualAppliedRefundForToken(t *testing.T) {
+ setupUserUpdateTestState(t)
+
+ user := User{Id: 901, Username: "settlement-refund-user", Quota: 0, Status: common.UserStatusEnabled}
+ token := Token{Id: 902, UserId: user.Id, Key: "settlement-refund-token", RemainQuota: 100, UsedQuota: 10, Status: common.TokenStatusEnabled}
+ subscription := UserSubscription{
+ Id: 903, UserId: user.Id, AmountTotal: 100, AmountUsed: 0,
+ StartTime: time.Now().Add(-time.Hour).Unix(), EndTime: time.Now().Add(time.Hour).Unix(), Status: "active",
+ }
+ preConsume := SubscriptionPreConsumeRecord{
+ RequestId: "regression:subscription-partial-refund-request", UserId: user.Id,
+ TokenId: token.Id, UserSubscriptionId: subscription.Id, PreConsumed: 10, Status: "consumed",
+ CreatedAt: time.Now().Unix(), UpdatedAt: time.Now().Unix(),
+ }
+ require.NoError(t, DB.Create(&user).Error)
+ require.NoError(t, DB.Create(&token).Error)
+ require.NoError(t, DB.Create(&subscription).Error)
+ require.NoError(t, DB.Create(&preConsume).Error)
+
+ _, _, err := ApplyBillingSettlementOnce(BillingSettlementInput{
+ OperationKey: "regression:subscription-partial-refund",
+ Source: BillingSettlementSourceSubscription,
+ UserID: user.Id,
+ SubscriptionID: subscription.Id,
+ TokenID: token.Id,
+ FundingDelta: -10,
+ TokenDelta: -10,
+ SubscriptionPreConsumeRequestID: preConsume.RequestId,
+ })
+ require.NoError(t, err)
+
+ var gotToken Token
+ require.NoError(t, DB.First(&gotToken, token.Id).Error)
+ require.EqualValues(t, 100, gotToken.RemainQuota)
+ require.EqualValues(t, 10, gotToken.UsedQuota)
+}
+
+func TestFailedSettlementLeavesDurablePendingOperation(t *testing.T) {
+ setupUserUpdateTestState(t)
+ user := User{Id: 904, Username: "pending-settlement-user", Quota: 100, Status: common.UserStatusEnabled}
+ require.NoError(t, DB.Create(&user).Error)
+
+ input := BillingSettlementInput{
+ OperationKey: "regression:pending-settlement",
+ Source: BillingSettlementSourceWallet,
+ UserID: user.Id,
+ TokenID: 905,
+ FundingDelta: 10,
+ TokenDelta: 10,
+ }
+ _, _, err := ApplyBillingSettlementOnce(input)
+ require.Error(t, err)
+
+ var pending BillingSettlement
+ require.NoError(t, DB.Where("operation_key = ?", input.OperationKey).First(&pending).Error)
+ require.Equal(t, "pending", pending.Status)
+
+ token := Token{Id: 905, UserId: user.Id, Key: "pending-settlement-token", RemainQuota: 100, Status: common.TokenStatusEnabled}
+ require.NoError(t, DB.Create(&token).Error)
+ _, _, err = ApplyBillingSettlementOnce(input)
+ require.NoError(t, err)
+ require.EqualValues(t, 90, getRegressionUserQuota(t, user.Id))
+ require.EqualValues(t, 90, getRegressionTokenRemainQuota(t, token.Id))
+}
+
+func TestBillingSettlementPendingInvalidationBypassesStaleUserCache(t *testing.T) {
+ setupUserUpdateTestState(t)
+ user := User{Id: 977, Username: "settlement-pending-cache-user", Quota: 100, Status: common.UserStatusEnabled}
+ require.NoError(t, DB.Create(&user).Error)
+ client := useCacheMutationRedis(t)
+ cacheUserForRetryTest(t, client, user)
+ hook := newBlockingCacheInvalidationHook()
+ client.AddHook(hook)
+ t.Cleanup(hook.unblock)
+ done := make(chan error, 1)
+
+ go func() {
+ _, _, err := ApplyBillingSettlementOnce(BillingSettlementInput{
+ OperationKey: "regression:settlement-pending-cache",
+ Source: BillingSettlementSourceWallet,
+ UserID: user.Id,
+ FundingDelta: 10,
+ })
+ done <- err
+ }()
+ waitForCacheHook(t, hook.started, "settlement invalidation start")
+
+ quota, err := GetUserQuota(user.Id, false)
+ require.NoError(t, err)
+ assert.EqualValues(t, 90, quota)
+
+ hook.unblock()
+ require.NoError(t, <-done)
+}
+
+func TestPendingPreConsumeSettlementRetryClaimsPersistedSelection(t *testing.T) {
+ setupUserUpdateTestState(t)
+ user := User{Id: 966, Username: "pending-pre-consume-user", Quota: 100, Status: common.UserStatusEnabled}
+ token := Token{Id: 967, UserId: user.Id, Key: "pending-pre-consume-token", RemainQuota: 100, Status: common.TokenStatusEnabled}
+ require.NoError(t, DB.Create(&user).Error)
+ require.NoError(t, DB.Create(&token).Error)
+ input := BillingSettlementInput{
+ OperationKey: "request:pending-pre-consume-selection:pre-consume",
+ Source: BillingSettlementSourceWallet,
+ UserID: user.Id,
+ TokenID: token.Id,
+ TokenKey: token.Key,
+ FundingDelta: 10,
+ TokenDelta: 10,
+ ManualOnFailure: true,
+ PreConsumeRequestID: "pending-pre-consume-selection",
+ PreConsumeModelName: "selection-model",
+ PreConsumeRequestedQuota: 12,
+ PreConsumeEffectiveQuota: 10,
+ }
+ _, alreadyApplied, err := ensureBillingSettlementRecord(input)
+ require.NoError(t, err)
+ require.False(t, alreadyApplied)
+
+ ProcessPendingBillingSettlementsOnce()
+
+ var selection BillingPreConsumeSelection
+ require.NoError(t, DB.Where("request_id = ?", input.PreConsumeRequestID).First(&selection).Error)
+ assert.Equal(t, BillingSettlementSourceWallet, selection.Source)
+ assert.EqualValues(t, user.Id, selection.UserID)
+ assert.EqualValues(t, token.Id, selection.TokenID)
+ assert.Equal(t, "selection-model", selection.ModelName)
+ assert.EqualValues(t, 12, selection.RequestedQuota)
+ assert.EqualValues(t, 10, selection.EffectiveQuota)
+ assert.EqualValues(t, 90, getRegressionUserQuota(t, user.Id))
+ assert.EqualValues(t, 90, getRegressionTokenRemainQuota(t, token.Id))
+}
+
+func TestWalletSettlementRollsBackWhenCacheOutboxCannotPersist(t *testing.T) {
+ setupUserUpdateTestState(t)
+ user := User{Id: 960, Username: "wallet-outbox-rollback", Quota: 100, Status: common.UserStatusEnabled}
+ token := Token{Id: 961, UserId: user.Id, Key: "wallet-outbox-token", RemainQuota: 100, Status: common.TokenStatusEnabled}
+ require.NoError(t, DB.Create(&user).Error)
+ require.NoError(t, DB.Create(&token).Error)
+ failCacheInvalidationOutboxInserts(t)
+
+ _, _, err := ApplyBillingSettlementOnce(BillingSettlementInput{
+ OperationKey: "regression:wallet-cache-outbox",
+ Source: BillingSettlementSourceWallet,
+ UserID: user.Id,
+ TokenID: token.Id,
+ FundingDelta: 10,
+ TokenDelta: 10,
+ })
+ require.Error(t, err)
+ assert.EqualValues(t, 100, getRegressionUserQuota(t, user.Id))
+ assert.EqualValues(t, 100, getRegressionTokenRemainQuota(t, token.Id))
+}
+
+func TestSubscriptionPreConsumeRollsBackWhenCacheOutboxCannotPersist(t *testing.T) {
+ setupUserUpdateTestState(t)
+ now := time.Now()
+ user := User{Id: 962, Username: "subscription-outbox-rollback", Status: common.UserStatusEnabled}
+ plan := SubscriptionPlan{Id: 963, Title: "outbox-plan", Enabled: true, TotalAmount: 100}
+ subscription := UserSubscription{
+ Id: 964, UserId: user.Id, PlanId: plan.Id, AmountTotal: 100,
+ StartTime: now.Add(-time.Hour).Unix(), EndTime: now.Add(time.Hour).Unix(), Status: "active",
+ }
+ token := Token{Id: 965, UserId: user.Id, Key: "subscription-outbox-token", RemainQuota: 100, Status: common.TokenStatusEnabled}
+ require.NoError(t, DB.Create(&user).Error)
+ require.NoError(t, DB.Create(&plan).Error)
+ require.NoError(t, DB.Create(&subscription).Error)
+ require.NoError(t, DB.Create(&token).Error)
+ failCacheInvalidationOutboxInserts(t)
+
+ _, err := PreConsumeTokenAndUserSubscription(
+ "regression:subscription-cache-outbox",
+ user.Id,
+ token.Id,
+ token.Key,
+ "outbox-model",
+ 0,
+ 10,
+ )
+ require.Error(t, err)
+ assert.EqualValues(t, 100, getRegressionTokenRemainQuota(t, token.Id))
+ var storedSubscription UserSubscription
+ require.NoError(t, DB.First(&storedSubscription, subscription.Id).Error)
+ assert.Zero(t, storedSubscription.AmountUsed)
+ var preConsumeCount int64
+ require.NoError(t, DB.Model(&SubscriptionPreConsumeRecord{}).
+ Where("request_id = ?", "regression:subscription-cache-outbox").
+ Count(&preConsumeCount).Error)
+ assert.Zero(t, preConsumeCount)
+}
+
+func failCacheInvalidationOutboxInserts(t *testing.T) {
+ t.Helper()
+ oldRedisEnabled := common.RedisEnabled
+ common.RedisEnabled = true
+ t.Cleanup(func() { common.RedisEnabled = oldRedisEnabled })
+ failCacheOutboxInserts(t)
+}
+
+func TestFailedRealtimeReserveIsNotRetriedAfterRequestAborts(t *testing.T) {
+ setupUserUpdateTestState(t)
+ user := User{Id: 906, Username: "failed-reserve-user", Quota: 100, Status: common.UserStatusEnabled}
+ require.NoError(t, DB.Create(&user).Error)
+ input := BillingSettlementInput{
+ OperationKey: "request:failed-reserve:reserve:1", Source: BillingSettlementSourceWallet,
+ UserID: user.Id, TokenID: 907, FundingDelta: 10, TokenDelta: 10, ManualOnFailure: true,
+ }
+ _, _, err := ApplyBillingSettlementOnce(input)
+ require.Error(t, err)
+
+ var settlement BillingSettlement
+ require.NoError(t, DB.Where("operation_key = ?", input.OperationKey).First(&settlement).Error)
+ require.Equal(t, BillingSettlementStatusManual, settlement.Status)
+ token := Token{Id: 907, UserId: user.Id, Key: "failed-reserve-token", RemainQuota: 100, Status: common.TokenStatusEnabled}
+ require.NoError(t, DB.Create(&token).Error)
+ processPendingBillingSettlements()
+ require.EqualValues(t, 100, getRegressionUserQuota(t, user.Id))
+ require.EqualValues(t, 100, getRegressionTokenRemainQuota(t, token.Id))
+}
+
+func TestBillingSettlementRejectsOperationKeyReuseWithoutSecondMutation(t *testing.T) {
+ setupUserUpdateTestState(t)
+ user := User{Id: 915, Username: "settlement-operation-reuse-user", Quota: 100, Status: common.UserStatusEnabled}
+ token := Token{Id: 916, UserId: user.Id, Key: "settlement-operation-reuse-token", RemainQuota: 100, Status: common.TokenStatusEnabled}
+ require.NoError(t, DB.Create(&user).Error)
+ require.NoError(t, DB.Create(&token).Error)
+
+ input := BillingSettlementInput{
+ OperationKey: "regression:operation-reuse",
+ Source: BillingSettlementSourceWallet,
+ UserID: user.Id,
+ TokenID: token.Id,
+ FundingDelta: 10,
+ TokenDelta: 10,
+ }
+ _, _, err := ApplyBillingSettlementOnce(input)
+ require.NoError(t, err)
+
+ input.FundingDelta = 20
+ input.TokenDelta = 20
+ _, _, err = ApplyBillingSettlementOnce(input)
+ require.Error(t, err)
+ assert.ErrorIs(t, err, ErrBillingSettlementOperationConflict)
+ assert.EqualValues(t, 90, getRegressionUserQuota(t, user.Id))
+ assert.EqualValues(t, 90, getRegressionTokenRemainQuota(t, token.Id))
+
+ var settlement BillingSettlement
+ require.NoError(t, DB.Where("operation_key = ?", "regression:operation-reuse").First(&settlement).Error)
+ assert.Equal(t, BillingSettlementStatusApplied, settlement.Status)
+ assert.EqualValues(t, 10, settlement.FundingDelta)
+ assert.EqualValues(t, 10, settlement.AppliedFundingDelta)
+}
+
+func TestBillingSettlementDoesNotTrustDuplicateInsertRowsAffected(t *testing.T) {
+ setupUserUpdateTestState(t)
+ user := User{Id: 917, Username: "settlement-found-rows-user", Quota: 100, Status: common.UserStatusEnabled}
+ token := Token{Id: 918, UserId: user.Id, Key: "settlement-found-rows-token", RemainQuota: 100, Status: common.TokenStatusEnabled}
+ require.NoError(t, DB.Create(&user).Error)
+ require.NoError(t, DB.Create(&token).Error)
+ input := BillingSettlementInput{
+ OperationKey: "regression:duplicate-found-rows", Source: BillingSettlementSourceWallet,
+ UserID: user.Id, TokenID: token.Id, FundingDelta: 10, TokenDelta: 10,
+ }
+ _, _, err := ApplyBillingSettlementOnce(input)
+ require.NoError(t, err)
+
+ callbackName := "test:billing-settlement-duplicate-found-rows"
+ require.NoError(t, DB.Callback().Create().After("gorm:create").Register(callbackName, func(tx *gorm.DB) {
+ if tx.Statement.Schema != nil && tx.Statement.Schema.Table == "billing_settlements" {
+ tx.RowsAffected = 1
+ }
+ }))
+ t.Cleanup(func() { DB.Callback().Create().Remove(callbackName) })
+
+ applied, alreadyApplied, err := ApplyBillingSettlementOnce(input)
+ require.NoError(t, err)
+ assert.True(t, alreadyApplied)
+ assert.EqualValues(t, 10, applied)
+ assert.EqualValues(t, 90, getRegressionUserQuota(t, user.Id))
+ assert.EqualValues(t, 90, getRegressionTokenRemainQuota(t, token.Id))
+}
+
+func TestBillingSettlementUpdatesTaskQuotaAtomically(t *testing.T) {
+ setupUserUpdateTestState(t)
+ user := User{Id: 906, Username: "task-settlement-user", Quota: 100, Status: common.UserStatusEnabled}
+ token := Token{Id: 907, UserId: user.Id, Key: "task-settlement-token", RemainQuota: 100, Status: common.TokenStatusEnabled}
+ task := Task{ID: 908, TaskID: "task-settlement", UserId: user.Id, Quota: 10, Status: TaskStatusSuccess}
+ require.NoError(t, DB.Create(&user).Error)
+ require.NoError(t, DB.Create(&token).Error)
+ require.NoError(t, DB.Create(&task).Error)
+
+ _, _, err := ApplyBillingSettlementOnce(BillingSettlementInput{
+ OperationKey: "regression:task-settlement",
+ Source: BillingSettlementSourceWallet,
+ UserID: user.Id,
+ TokenID: token.Id,
+ FundingDelta: 5,
+ TokenDelta: 5,
+ TaskID: task.ID,
+ TaskQuota: 10,
+ TaskQuotaTarget: 15,
+ })
+ require.NoError(t, err)
+ require.EqualValues(t, 15, getRegressionTaskQuota(t, task.ID))
+ require.EqualValues(t, 95, getRegressionUserQuota(t, user.Id))
+}
+
+func TestBillingSettlementTaskConflictRollsBackAllLedgerLegs(t *testing.T) {
+ setupUserUpdateTestState(t)
+ user := User{Id: 909, Username: "task-conflict-user", Quota: 100, Status: common.UserStatusEnabled}
+ token := Token{Id: 910, UserId: user.Id, Key: "task-conflict-token", RemainQuota: 100, UsedQuota: 10, Status: common.TokenStatusEnabled}
+ task := Task{ID: 911, TaskID: "task-conflict", UserId: user.Id, Quota: 20, Status: TaskStatusSuccess}
+ require.NoError(t, DB.Create(&user).Error)
+ require.NoError(t, DB.Create(&token).Error)
+ require.NoError(t, DB.Create(&task).Error)
+
+ _, _, err := ApplyBillingSettlementOnce(BillingSettlementInput{
+ OperationKey: "regression:task-conflict", Source: BillingSettlementSourceWallet,
+ UserID: user.Id, TokenID: token.Id, FundingDelta: 5, TokenDelta: 5,
+ TaskID: task.ID, TaskQuota: 10, TaskQuotaTarget: 15,
+ })
+ require.Error(t, err)
+ assert.EqualValues(t, 100, getRegressionUserQuota(t, user.Id))
+ assert.EqualValues(t, 100, getRegressionTokenRemainQuota(t, token.Id))
+ assert.EqualValues(t, 20, getRegressionTaskQuota(t, task.ID))
+ var settlement BillingSettlement
+ require.NoError(t, DB.Where("operation_key = ?", "regression:task-conflict").First(&settlement).Error)
+ assert.Equal(t, BillingSettlementStatusManual, settlement.Status)
+}
+
+func TestSubscriptionSettlementRefusesCrossResetPeriodReplay(t *testing.T) {
+ setupUserUpdateTestState(t)
+ user := User{Id: 912, Username: "subscription-period-user", Quota: 0, Status: common.UserStatusEnabled}
+ token := Token{Id: 913, UserId: user.Id, Key: "subscription-period-token", RemainQuota: 100, UsedQuota: 10, Status: common.TokenStatusEnabled}
+ subscription := UserSubscription{
+ Id: 914, UserId: user.Id, AmountTotal: 100, AmountUsed: 10,
+ LastResetTime: time.Now().Add(time.Hour).Unix(), StartTime: time.Now().Add(-time.Hour).Unix(),
+ EndTime: time.Now().Add(time.Hour).Unix(), Status: "active",
+ }
+ preConsume := SubscriptionPreConsumeRecord{
+ RequestId: "regression:period-request", UserId: user.Id,
+ TokenId: token.Id, UserSubscriptionId: subscription.Id, PreConsumed: 10, Status: "consumed",
+ CreatedAt: time.Now().Add(-time.Hour).Unix(), UpdatedAt: time.Now().Add(-time.Hour).Unix(),
+ }
+ require.NoError(t, DB.Create(&user).Error)
+ require.NoError(t, DB.Create(&token).Error)
+ require.NoError(t, DB.Create(&subscription).Error)
+ require.NoError(t, DB.Create(&preConsume).Error)
+
+ _, _, err := ApplyBillingSettlementOnce(BillingSettlementInput{
+ OperationKey: "regression:period-replay", Source: BillingSettlementSourceSubscription,
+ UserID: user.Id, SubscriptionID: subscription.Id, TokenID: token.Id,
+ FundingDelta: -10, TokenDelta: -10,
+ SubscriptionPreConsumeRequestID: preConsume.RequestId,
+ })
+ require.Error(t, err)
+ assert.EqualValues(t, 100, getRegressionTokenRemainQuota(t, token.Id))
+ var gotSub UserSubscription
+ require.NoError(t, DB.First(&gotSub, subscription.Id).Error)
+ assert.EqualValues(t, 10, gotSub.AmountUsed)
+ var settlement BillingSettlement
+ require.NoError(t, DB.Where("operation_key = ?", "regression:period-replay").First(&settlement).Error)
+ assert.Equal(t, BillingSettlementStatusManual, settlement.Status)
+}
+
+func TestSubscriptionSettlementRefusesSameSecondResetReplay(t *testing.T) {
+ setupUserUpdateTestState(t)
+ user := User{Id: 919, Username: "subscription-same-second-user", Quota: 0, Status: common.UserStatusEnabled}
+ token := Token{Id: 920, UserId: user.Id, Key: "subscription-same-second-token", RemainQuota: 100, UsedQuota: 10, Status: common.TokenStatusEnabled}
+ subscription := UserSubscription{
+ Id: 921, UserId: user.Id, AmountTotal: 100, AmountUsed: 10,
+ StartTime: time.Now().Add(-time.Hour).Unix(), EndTime: time.Now().Add(time.Hour).Unix(), Status: "active",
+ }
+ preConsume := SubscriptionPreConsumeRecord{
+ RequestId: "regression:same-second-period-request", UserId: user.Id,
+ TokenId: token.Id, UserSubscriptionId: subscription.Id, PreConsumed: 10, Status: "consumed",
+ }
+ require.NoError(t, DB.Create(&user).Error)
+ require.NoError(t, DB.Create(&token).Error)
+ require.NoError(t, DB.Create(&subscription).Error)
+ require.NoError(t, DB.Create(&preConsume).Error)
+ periodBoundary := time.Now().Unix()
+ require.NoError(t, DB.Model(&UserSubscription{}).Where("id = ?", subscription.Id).Update("last_reset_time", periodBoundary).Error)
+ require.NoError(t, DB.Model(&SubscriptionPreConsumeRecord{}).Where("id = ?", preConsume.Id).Update("created_at", periodBoundary).Error)
+
+ _, _, err := ApplyBillingSettlementOnce(BillingSettlementInput{
+ OperationKey: "regression:same-second-period-replay", Source: BillingSettlementSourceSubscription,
+ UserID: user.Id, SubscriptionID: subscription.Id, TokenID: token.Id,
+ FundingDelta: -10, TokenDelta: -10,
+ SubscriptionPreConsumeRequestID: preConsume.RequestId,
+ })
+ require.Error(t, err)
+ assert.ErrorIs(t, err, ErrSubscriptionSettlementPeriodChanged)
+ assert.EqualValues(t, 100, getRegressionTokenRemainQuota(t, token.Id))
+ var gotSub UserSubscription
+ require.NoError(t, DB.First(&gotSub, subscription.Id).Error)
+ assert.EqualValues(t, 10, gotSub.AmountUsed)
+}
+
+func TestSubscriptionPreConsumeReplayRejectsChangedUserAndAmount(t *testing.T) {
+ setupUserUpdateTestState(t)
+ now := time.Now()
+ subscription := UserSubscription{
+ Id: 921, UserId: 922, AmountTotal: 100, AmountUsed: 10,
+ StartTime: now.Add(-time.Hour).Unix(), EndTime: now.Add(time.Hour).Unix(), Status: "active",
+ }
+ record := SubscriptionPreConsumeRecord{
+ RequestId: "regression:pre-consume-identity", UserId: subscription.UserId,
+ UserSubscriptionId: subscription.Id, PreConsumed: 10, Status: "consumed",
+ CreatedAt: now.Unix(), UpdatedAt: now.Unix(),
+ }
+ require.NoError(t, DB.Create(&subscription).Error)
+ require.NoError(t, DB.Create(&record).Error)
+
+ _, err := PreConsumeUserSubscription(record.RequestId, subscription.UserId, "model-a", 0, 20)
+ require.Error(t, err)
+ _, err = PreConsumeUserSubscription(record.RequestId, subscription.UserId+1, "model-a", 0, 10)
+ require.Error(t, err)
+}
+
+func TestSubscriptionPreConsumeReplayRejectsChangedToken(t *testing.T) {
+ setupUserUpdateTestState(t)
+ now := time.Now()
+ user := User{Id: 923, Username: "subscription-replay-user", Status: common.UserStatusEnabled}
+ plan := SubscriptionPlan{Id: 924, Title: "replay-plan", Enabled: true, TotalAmount: 100}
+ subscription := UserSubscription{
+ Id: 925, UserId: user.Id, PlanId: plan.Id, AmountTotal: 100,
+ StartTime: now.Add(-time.Hour).Unix(), EndTime: now.Add(time.Hour).Unix(), Status: "active",
+ }
+ firstToken := Token{Id: 926, UserId: user.Id, Key: "subscription-first-token", RemainQuota: 100, Status: common.TokenStatusEnabled}
+ secondToken := Token{Id: 927, UserId: user.Id, Key: "subscription-second-token", RemainQuota: 100, Status: common.TokenStatusEnabled}
+ require.NoError(t, DB.Create(&user).Error)
+ require.NoError(t, DB.Create(&plan).Error)
+ require.NoError(t, DB.Create(&subscription).Error)
+ require.NoError(t, DB.Create(&firstToken).Error)
+ require.NoError(t, DB.Create(&secondToken).Error)
+
+ const requestID = "regression:pre-consume-token"
+ _, err := PreConsumeTokenAndUserSubscription(requestID, user.Id, firstToken.Id, firstToken.Key, "model-a", 0, 10)
+ require.NoError(t, err)
+ _, err = PreConsumeTokenAndUserSubscription(requestID, user.Id, secondToken.Id, secondToken.Key, "model-a", 0, 10)
+ require.Error(t, err)
+ assert.EqualValues(t, 100, getRegressionTokenRemainQuota(t, secondToken.Id))
+}
+
+func TestBillingSettlementRejectsTokenOwnedByAnotherUser(t *testing.T) {
+ setupUserUpdateTestState(t)
+ owner := User{Id: 928, Username: "settlement-owner", AffCode: "settlement-owner", Quota: 100, Status: common.UserStatusEnabled}
+ other := User{Id: 929, Username: "settlement-other", AffCode: "settlement-other", Quota: 100, Status: common.UserStatusEnabled}
+ otherToken := Token{Id: 930, UserId: other.Id, Key: "settlement-other-token", RemainQuota: 100, Status: common.TokenStatusEnabled}
+ require.NoError(t, DB.Create(&owner).Error)
+ require.NoError(t, DB.Create(&other).Error)
+ require.NoError(t, DB.Create(&otherToken).Error)
+
+ _, _, err := ApplyBillingSettlementOnce(BillingSettlementInput{
+ OperationKey: "regression:cross-user-token", Source: BillingSettlementSourceWallet,
+ UserID: owner.Id, TokenID: otherToken.Id, FundingDelta: 10, TokenDelta: 10,
+ })
+ require.Error(t, err)
+ assert.EqualValues(t, 100, getRegressionUserQuota(t, owner.Id))
+ assert.EqualValues(t, 100, getRegressionTokenRemainQuota(t, otherToken.Id))
+}
+
+func TestSubscriptionSettlementRejectsPreConsumeOwnedByAnotherUser(t *testing.T) {
+ setupUserUpdateTestState(t)
+ now := time.Now()
+ owner := User{Id: 931, Username: "subscription-owner", AffCode: "subscription-owner", Status: common.UserStatusEnabled}
+ other := User{Id: 932, Username: "subscription-other", AffCode: "subscription-other", Status: common.UserStatusEnabled}
+ subscription := UserSubscription{
+ Id: 933, UserId: owner.Id, AmountTotal: 100, AmountUsed: 10,
+ StartTime: now.Add(-time.Hour).Unix(), EndTime: now.Add(time.Hour).Unix(), Status: "active",
+ }
+ record := SubscriptionPreConsumeRecord{
+ RequestId: "regression:cross-user-subscription", UserId: owner.Id,
+ UserSubscriptionId: subscription.Id, PreConsumed: 10, Status: "consumed",
+ CreatedAt: now.Unix(), UpdatedAt: now.Unix(),
+ }
+ token := Token{Id: 934, UserId: other.Id, Key: "subscription-other-token", RemainQuota: 100, Status: common.TokenStatusEnabled}
+ require.NoError(t, DB.Create(&owner).Error)
+ require.NoError(t, DB.Create(&other).Error)
+ require.NoError(t, DB.Create(&subscription).Error)
+ require.NoError(t, DB.Create(&record).Error)
+ require.NoError(t, DB.Create(&token).Error)
+
+ _, _, err := ApplyBillingSettlementOnce(BillingSettlementInput{
+ OperationKey: "regression:cross-user-subscription-settlement", Source: BillingSettlementSourceSubscription,
+ UserID: other.Id, SubscriptionID: subscription.Id, TokenID: token.Id,
+ FundingDelta: 10, TokenDelta: 10, SubscriptionPreConsumeRequestID: record.RequestId,
+ })
+ require.Error(t, err)
+ assert.EqualValues(t, 100, getRegressionTokenRemainQuota(t, token.Id))
+ var got UserSubscription
+ require.NoError(t, DB.First(&got, subscription.Id).Error)
+ assert.EqualValues(t, 10, got.AmountUsed)
+}
+
+func TestSubscriptionPreConsumeCleanupRetainsIdempotencyTombstones(t *testing.T) {
+ setupUserUpdateTestState(t)
+ record := SubscriptionPreConsumeRecord{
+ RequestId: "regression:retained-pre-consume", UserId: 935, UserSubscriptionId: 936,
+ PreConsumed: 10, Status: "consumed",
+ CreatedAt: time.Now().Add(-30 * 24 * time.Hour).Unix(), UpdatedAt: time.Now().Add(-30 * 24 * time.Hour).Unix(),
+ }
+ require.NoError(t, DB.Create(&record).Error)
+
+ deleted, err := CleanupSubscriptionPreConsumeRecords(7 * 24 * 3600)
+ require.NoError(t, err)
+ assert.Zero(t, deleted)
+ var count int64
+ require.NoError(t, DB.Model(&SubscriptionPreConsumeRecord{}).Where("request_id = ?", record.RequestId).Count(&count).Error)
+ assert.EqualValues(t, 1, count)
+}
+
+func TestSubscriptionPreConsumeCleanupDeletesOldTerminalTombstones(t *testing.T) {
+ setupUserUpdateTestState(t)
+ record := SubscriptionPreConsumeRecord{
+ RequestId: "regression:cleanup-terminal-pre-consume", UserId: 937, UserSubscriptionId: 938,
+ PreConsumed: 10, Status: "consumed",
+ }
+ require.NoError(t, DB.Create(&record).Error)
+ oldTimestamp := time.Now().Add(-30 * 24 * time.Hour).Unix()
+ require.NoError(t, DB.Model(&SubscriptionPreConsumeRecord{}).
+ Where("request_id = ?", record.RequestId).
+ UpdateColumns(map[string]interface{}{"created_at": oldTimestamp, "updated_at": oldTimestamp}).Error)
+ now := time.Now().Unix()
+ require.NoError(t, DB.Create(&BillingSettlement{
+ OperationKey: "regression:cleanup-terminal-settlement",
+ Source: BillingSettlementSourceSubscription,
+ UserID: record.UserId,
+ SubscriptionID: record.UserSubscriptionId,
+ FundingDelta: -10,
+ AppliedFundingDelta: -10,
+ Status: BillingSettlementStatusApplied,
+ SubscriptionPreConsumeRequestID: record.RequestId,
+ CreatedAt: now,
+ UpdatedAt: now,
+ Revision: 1,
+ }).Error)
+
+ deleted, err := CleanupSubscriptionPreConsumeRecords(7 * 24 * 3600)
+ require.NoError(t, err)
+ assert.EqualValues(t, 1, deleted)
+ var count int64
+ require.NoError(t, DB.Model(&SubscriptionPreConsumeRecord{}).Where("request_id = ?", record.RequestId).Count(&count).Error)
+ assert.Zero(t, count)
+}
+
+func getRegressionUserQuota(t *testing.T, userID int) int64 {
+ t.Helper()
+ var user User
+ require.NoError(t, DB.First(&user, userID).Error)
+ return user.Quota
+}
+
+func getRegressionTokenRemainQuota(t *testing.T, tokenID int) int64 {
+ t.Helper()
+ var token Token
+ require.NoError(t, DB.First(&token, tokenID).Error)
+ return token.RemainQuota
+}
+
+func getRegressionTaskQuota(t *testing.T, taskID int64) int {
+ t.Helper()
+ var task Task
+ require.NoError(t, DB.First(&task, taskID).Error)
+ return task.Quota
+}
diff --git a/model/cache_invalidation_task.go b/model/cache_invalidation_task.go
new file mode 100644
index 00000000..d8e0909f
--- /dev/null
+++ b/model/cache_invalidation_task.go
@@ -0,0 +1,246 @@
+package model
+
+import (
+ "context"
+ "errors"
+ "fmt"
+ "strconv"
+ "sync"
+ "time"
+
+ "github.com/MAX-API-Next/MAX-API/common"
+ "github.com/bytedance/gopkg/util/gopool"
+ "gorm.io/gorm"
+ "gorm.io/gorm/clause"
+)
+
+const (
+ cacheInvalidationKindUser = "user"
+ cacheInvalidationKindToken = "token"
+)
+
+var errCacheInvalidationPermanent = errors.New("permanent cache invalidation failure")
+
+// CacheInvalidationTask persists Redis invalidations so process restarts do not
+// lose work queued after a committed database mutation.
+type CacheInvalidationTask struct {
+ ID int64 `gorm:"primaryKey"`
+ Kind string `gorm:"type:varchar(16);uniqueIndex:idx_cache_invalidation_target,priority:1;not null"`
+ EntityKey string `gorm:"type:varchar(96);uniqueIndex:idx_cache_invalidation_target,priority:2;not null"`
+ VersionKey string `gorm:"type:varchar(96);not null"`
+ DeleteEntry bool `gorm:"uniqueIndex:idx_cache_invalidation_target,priority:3;not null;default:false"`
+ Attempts int `gorm:"not null;default:0"`
+ LastError string `gorm:"type:text"`
+ NextAttempt int64 `gorm:"index;not null;default:0"`
+ CreatedAt int64 `gorm:"not null"`
+ UpdatedAt int64 `gorm:"not null"`
+ Revision int64 `gorm:"not null;default:1"`
+}
+
+func persistCacheInvalidationTask(kind, entityKey, versionKey string, deleteEntry bool, cause error) {
+ if _, err := upsertCacheInvalidationTask(DB, kind, entityKey, versionKey, deleteEntry, cause); err != nil {
+ common.SysLog("failed to persist cache invalidation task: " + err.Error())
+ }
+}
+
+func stageCacheInvalidationTaskTx(tx *gorm.DB, kind, entityKey, versionKey string, deleteEntry bool) (CacheInvalidationTask, error) {
+ if !common.RedisEnabled {
+ return CacheInvalidationTask{}, nil
+ }
+ return upsertCacheInvalidationTask(tx, kind, entityKey, versionKey, deleteEntry, nil)
+}
+
+func upsertCacheInvalidationTask(db *gorm.DB, kind, entityKey, versionKey string, deleteEntry bool, cause error) (CacheInvalidationTask, error) {
+ if db == nil {
+ return CacheInvalidationTask{}, errors.New("database is not initialized")
+ }
+ if entityKey == "" {
+ return CacheInvalidationTask{}, errors.New("cache invalidation entity key is required")
+ }
+ now := time.Now().Unix()
+ lastError := ""
+ if cause != nil {
+ lastError = cause.Error()
+ }
+ task := CacheInvalidationTask{
+ Kind: kind,
+ EntityKey: entityKey,
+ VersionKey: versionKey,
+ DeleteEntry: deleteEntry,
+ LastError: lastError,
+ NextAttempt: now,
+ CreatedAt: now,
+ UpdatedAt: now,
+ Revision: 1,
+ }
+ if err := db.Clauses(clause.OnConflict{
+ Columns: []clause.Column{{Name: "kind"}, {Name: "entity_key"}, {Name: "delete_entry"}},
+ DoUpdates: clause.Assignments(map[string]interface{}{
+ "version_key": versionKey,
+ "updated_at": now,
+ "revision": gorm.Expr("revision + ?", 1),
+ }),
+ }).Create(&task).Error; err != nil {
+ return CacheInvalidationTask{}, err
+ }
+ var stored CacheInvalidationTask
+ if err := db.Where("kind = ? AND entity_key = ? AND delete_entry = ?", kind, entityKey, deleteEntry).Take(&stored).Error; err != nil {
+ return CacheInvalidationTask{}, err
+ }
+ return stored, nil
+}
+
+// cacheInvalidationTaskPending makes the durable outbox a read-side cache fence.
+// A committed DB mutation may be visible before its Redis invalidation is
+// delivered; readers must not trust or refill that entity cache while its task
+// is still pending.
+func cacheInvalidationTaskPending(kind, entityKey string) (bool, error) {
+ if !common.RedisEnabled {
+ return false, nil
+ }
+ if DB == nil {
+ return false, errors.New("database is not initialized")
+ }
+ if kind == "" || entityKey == "" {
+ return false, errors.New("cache invalidation target is required")
+ }
+ var task CacheInvalidationTask
+ result := DB.Select("id").
+ Where("kind = ? AND entity_key = ?", kind, entityKey).
+ Limit(1).
+ Find(&task)
+ if result.Error != nil {
+ return false, result.Error
+ }
+ return result.RowsAffected > 0, nil
+}
+
+func dispatchStagedCacheInvalidation(task CacheInvalidationTask) {
+ if task.ID <= 0 || !common.RedisEnabled {
+ return
+ }
+ if err := executeCacheInvalidationTask(task); err != nil {
+ common.SysLog(fmt.Sprintf("cache invalidation remains pending after commit: kind=%s entity=%s error=%v", task.Kind, task.EntityKey, err))
+ return
+ }
+ if err := DB.Where("id = ? AND revision = ?", task.ID, task.Revision).Delete(&CacheInvalidationTask{}).Error; err != nil {
+ common.SysLog("failed to complete immediate cache invalidation task: " + err.Error())
+ }
+}
+
+func dispatchStagedCacheInvalidations(tasks []CacheInvalidationTask) {
+ for _, task := range tasks {
+ dispatchStagedCacheInvalidation(task)
+ }
+}
+
+var cacheInvalidationRunnerOnce sync.Once
+var cacheInvalidationRunnerState struct {
+ sync.Mutex
+ cancel context.CancelFunc
+ done chan struct{}
+}
+
+func StartCacheInvalidationTaskRunner() {
+ if !common.RedisEnabled || DB == nil {
+ return
+ }
+ cacheInvalidationRunnerOnce.Do(func() {
+ ctx, cancel := context.WithCancel(context.Background())
+ done := make(chan struct{})
+ cacheInvalidationRunnerState.Lock()
+ cacheInvalidationRunnerState.cancel = cancel
+ cacheInvalidationRunnerState.done = done
+ cacheInvalidationRunnerState.Unlock()
+ gopool.Go(func() {
+ defer close(done)
+ processCacheInvalidationTasks()
+ ticker := time.NewTicker(time.Second)
+ defer ticker.Stop()
+ for {
+ select {
+ case <-ctx.Done():
+ return
+ case <-ticker.C:
+ processCacheInvalidationTasks()
+ }
+ }
+ })
+ })
+}
+
+func StopCacheInvalidationTaskRunner(ctx context.Context) error {
+ cacheInvalidationRunnerState.Lock()
+ cancel := cacheInvalidationRunnerState.cancel
+ done := cacheInvalidationRunnerState.done
+ cacheInvalidationRunnerState.Unlock()
+ if cancel == nil || done == nil {
+ return nil
+ }
+ cancel()
+ select {
+ case <-done:
+ return nil
+ case <-ctx.Done():
+ return ctx.Err()
+ }
+}
+
+func processCacheInvalidationTasks() {
+ now := time.Now().Unix()
+ var tasks []CacheInvalidationTask
+ if err := DB.Where("next_attempt <= ?", now).Order("id").Limit(100).Find(&tasks).Error; err != nil {
+ common.SysLog("failed to load cache invalidation tasks: " + err.Error())
+ return
+ }
+ for _, task := range tasks {
+ if err := executeCacheInvalidationTask(task); err != nil {
+ if errors.Is(err, errCacheInvalidationPermanent) {
+ if deleteErr := DB.Where("id = ? AND revision = ?", task.ID, task.Revision).Delete(&CacheInvalidationTask{}).Error; deleteErr != nil {
+ common.SysLog("failed to drop permanent cache invalidation task: " + deleteErr.Error())
+ }
+ common.SysLog(fmt.Sprintf("dropped permanent cache invalidation task: kind=%s entity=%s error=%v", task.Kind, task.EntityKey, err))
+ continue
+ }
+ attempts := task.Attempts + 1
+ delay := 1 << min(attempts, 6)
+ if updateErr := DB.Model(&CacheInvalidationTask{}).Where("id = ? AND revision = ?", task.ID, task.Revision).Updates(map[string]interface{}{
+ "attempts": attempts,
+ "last_error": err.Error(),
+ "next_attempt": time.Now().Add(time.Duration(delay) * time.Second).Unix(),
+ "updated_at": time.Now().Unix(),
+ "revision": gorm.Expr("revision + ?", 1),
+ }).Error; updateErr != nil {
+ common.SysLog("failed to reschedule cache invalidation task: " + updateErr.Error())
+ }
+ continue
+ }
+ if err := DB.Where("id = ? AND revision = ?", task.ID, task.Revision).Delete(&CacheInvalidationTask{}).Error; err != nil {
+ common.SysLog("failed to complete cache invalidation task: " + err.Error())
+ }
+ }
+}
+
+func executeCacheInvalidationTask(task CacheInvalidationTask) error {
+ switch task.Kind {
+ case cacheInvalidationKindUser:
+ userID, err := strconv.Atoi(task.EntityKey)
+ if err != nil || userID <= 0 {
+ return fmt.Errorf("%w: invalid persisted user cache key: %s", errCacheInvalidationPermanent, task.EntityKey)
+ }
+ if task.DeleteEntry {
+ return common.RedisDeleteVersionedHash(getUserCacheKey(userID), getUserCacheVersionKey(userID))
+ }
+ return common.RedisInvalidateVersionedHash(getUserCacheKey(userID), getUserCacheVersionKey(userID))
+ case cacheInvalidationKindToken:
+ if task.EntityKey == "" || task.VersionKey == "" {
+ return fmt.Errorf("%w: invalid persisted token cache key", errCacheInvalidationPermanent)
+ }
+ if task.DeleteEntry {
+ return common.RedisDeleteVersionedHash(task.EntityKey, task.VersionKey)
+ }
+ return common.RedisInvalidateVersionedHash(task.EntityKey, task.VersionKey)
+ default:
+ return fmt.Errorf("%w: unknown cache invalidation kind: %s", errCacheInvalidationPermanent, task.Kind)
+ }
+}
diff --git a/model/cache_invalidation_task_test.go b/model/cache_invalidation_task_test.go
new file mode 100644
index 00000000..83fe09a6
--- /dev/null
+++ b/model/cache_invalidation_task_test.go
@@ -0,0 +1,216 @@
+package model
+
+import (
+ "context"
+ "errors"
+ "sync"
+ "testing"
+ "time"
+
+ "github.com/MAX-API-Next/MAX-API/common"
+ miniredis "github.com/alicebob/miniredis/v2"
+ "github.com/go-redis/redis/v8"
+ "github.com/stretchr/testify/assert"
+ "github.com/stretchr/testify/require"
+ "gorm.io/gorm/schema"
+)
+
+func TestCacheInvalidationTargetIndexFitsMySQL57CompactUtf8mb4Limit(t *testing.T) {
+ parsed, err := schema.Parse(&CacheInvalidationTask{}, &sync.Map{}, schema.NamingStrategy{})
+ require.NoError(t, err)
+
+ kind := parsed.LookUpField("Kind")
+ entityKey := parsed.LookUpField("EntityKey")
+ require.NotNil(t, kind)
+ require.NotNil(t, entityKey)
+
+ // MySQL 5.7.8 with utf8mb4 and COMPACT row format limits index keys to
+ // 767 bytes. DeleteEntry is a one-byte boolean column.
+ indexBytes := (kind.Size+entityKey.Size)*4 + 1
+ require.LessOrEqual(t, indexBytes, 767)
+}
+
+func useCacheMutationRedis(t *testing.T) *redis.Client {
+ t.Helper()
+ server := miniredis.RunT(t)
+ client := redis.NewClient(&redis.Options{Addr: server.Addr()})
+
+ oldRDB := common.RDB
+ oldRedisEnabled := common.RedisEnabled
+ common.RDB = client
+ common.RedisEnabled = true
+ t.Cleanup(func() {
+ _ = client.Close()
+ common.RDB = oldRDB
+ common.RedisEnabled = oldRedisEnabled
+ })
+ return client
+}
+
+func useRecoverableCacheMutationRedis(t *testing.T) (*redis.Client, *recoverableCacheMutationHook) {
+ t.Helper()
+ client := useCacheMutationRedis(t)
+ require.NoError(t, common.RedisInvalidateVersionedHash("persistent-retry-warmup", "persistent-retry-warmup-version"))
+ hook := &recoverableCacheMutationHook{}
+ client.AddHook(hook)
+ return client, hook
+}
+
+func failCacheOutboxInserts(t *testing.T) {
+ t.Helper()
+ if !common.UsingSQLite {
+ t.Skip("cache outbox trigger failure injection uses SQLite trigger syntax")
+ }
+ require.NoError(t, DB.Exec("DROP TRIGGER IF EXISTS cache_outbox_insert_failure").Error)
+ require.NoError(t, DB.Exec("CREATE TRIGGER cache_outbox_insert_failure BEFORE INSERT ON cache_invalidation_tasks BEGIN SELECT RAISE(FAIL, 'cache outbox unavailable'); END").Error)
+ t.Cleanup(func() {
+ require.NoError(t, DB.Exec("DROP TRIGGER IF EXISTS cache_outbox_insert_failure").Error)
+ })
+}
+
+func clearTokenCacheRetryMemory(t *testing.T) {
+ t.Helper()
+ tokenCacheRetries.Lock()
+ tokenCacheRetries.pending = make(map[string]*tokenCacheRetryState)
+ tokenCacheRetries.Unlock()
+
+ require.Eventually(t, func() bool {
+ tokenCacheRetries.Lock()
+ defer tokenCacheRetries.Unlock()
+ return !tokenCacheRetries.running
+ }, time.Second, 10*time.Millisecond)
+}
+
+func TestApplyBillingSettlementOnceResolvesTokenKeyForCacheInvalidation(t *testing.T) {
+ setupUserUpdateTestState(t)
+ user := User{Id: 301, Username: "settlement-cache-user", Quota: 100, Status: common.UserStatusEnabled}
+ token := Token{Id: 302, UserId: user.Id, Key: "settlement-cache-token", RemainQuota: 100, Status: common.TokenStatusEnabled}
+ require.NoError(t, DB.Create(&user).Error)
+ require.NoError(t, DB.Create(&token).Error)
+
+ client := useCacheMutationRedis(t)
+ cacheTokenForRetryTest(t, client, token)
+
+ _, _, err := ApplyBillingSettlementOnce(BillingSettlementInput{
+ OperationKey: "request:resolve-token-key:settle",
+ Source: BillingSettlementSourceWallet,
+ UserID: user.Id,
+ TokenID: token.Id,
+ FundingDelta: 10,
+ TokenDelta: 10,
+ })
+ require.NoError(t, err)
+
+ exists, err := client.Exists(context.Background(), getTokenCacheKey(token.Key)).Result()
+ require.NoError(t, err)
+ assert.Zero(t, exists)
+}
+
+func TestPersistedTokenCacheInvalidationSurvivesRetryWorkerRestart(t *testing.T) {
+ setupUserUpdateTestState(t)
+ token := Token{Id: 303, UserId: 304, Key: "persistent-token-cache-retry", RemainQuota: 10, Status: common.TokenStatusEnabled}
+ require.NoError(t, DB.Create(&token).Error)
+ client, hook := useRecoverableCacheMutationRedis(t)
+ cacheTokenForRetryTest(t, client, token)
+
+ enqueueTokenCacheRetry(token.Key, false, errors.New("injected cache outage"))
+ require.Eventually(t, func() bool { return hook.seen.Load() > 0 }, time.Second, 10*time.Millisecond)
+ clearTokenCacheRetryMemory(t)
+
+ var task CacheInvalidationTask
+ require.NoError(t, DB.First(&task).Error)
+ assert.Equal(t, cacheInvalidationKindToken, task.Kind)
+ assert.Equal(t, getTokenCacheKey(token.Key), task.EntityKey)
+ assert.NotContains(t, task.EntityKey, token.Key)
+
+ hook.available.Store(true)
+ processCacheInvalidationTasks()
+
+ var count int64
+ require.NoError(t, DB.Model(&CacheInvalidationTask{}).Count(&count).Error)
+ assert.Zero(t, count)
+ exists, err := client.Exists(context.Background(), getTokenCacheKey(token.Key)).Result()
+ require.NoError(t, err)
+ assert.Zero(t, exists)
+}
+
+func TestPersistedCacheInvalidationRevisionProtectsNewerTask(t *testing.T) {
+ setupUserUpdateTestState(t)
+ token := Token{Id: 305, UserId: 306, Key: "persistent-revision-cache-retry", RemainQuota: 10, Status: common.TokenStatusEnabled}
+ require.NoError(t, DB.Create(&token).Error)
+ client, hook := useRecoverableCacheMutationRedis(t)
+ cacheTokenForRetryTest(t, client, token)
+
+ enqueueTokenCacheRetry(token.Key, false, errors.New("first cache outage"))
+ require.Eventually(t, func() bool { return hook.seen.Load() > 0 }, time.Second, 10*time.Millisecond)
+ clearTokenCacheRetryMemory(t)
+
+ var stale CacheInvalidationTask
+ require.NoError(t, DB.First(&stale).Error)
+ staleRevision := stale.Revision
+ persistCacheInvalidationTask(cacheInvalidationKindToken, stale.EntityKey, stale.VersionKey, stale.DeleteEntry, errors.New("newer cache outage"))
+
+ result := DB.Where("id = ? AND revision = ?", stale.ID, staleRevision).Delete(&CacheInvalidationTask{})
+ require.NoError(t, result.Error)
+ assert.Zero(t, result.RowsAffected)
+
+ hook.available.Store(true)
+ processCacheInvalidationTasks()
+ var count int64
+ require.NoError(t, DB.Model(&CacheInvalidationTask{}).Count(&count).Error)
+ assert.Zero(t, count)
+}
+
+func TestCacheInvalidationRestagePreservesDurableBackoff(t *testing.T) {
+ setupUserUpdateTestState(t)
+ oldRedisEnabled := common.RedisEnabled
+ common.RedisEnabled = true
+ t.Cleanup(func() { common.RedisEnabled = oldRedisEnabled })
+
+ task, err := upsertCacheInvalidationTask(DB, cacheInvalidationKindUser, "301", getUserCacheVersionKey(301), false, errors.New("first outage"))
+ require.NoError(t, err)
+ nextAttempt := time.Now().Add(time.Minute).Unix()
+ require.NoError(t, DB.Model(&CacheInvalidationTask{}).Where("id = ?", task.ID).Updates(map[string]interface{}{
+ "attempts": 4,
+ "next_attempt": nextAttempt,
+ "last_error": "first outage",
+ }).Error)
+
+ restaged, err := upsertCacheInvalidationTask(DB, cacheInvalidationKindUser, "301", "new-version", false, errors.New("second outage"))
+ require.NoError(t, err)
+
+ assert.EqualValues(t, 4, restaged.Attempts)
+ assert.EqualValues(t, nextAttempt, restaged.NextAttempt)
+ assert.Equal(t, "first outage", restaged.LastError)
+ assert.Equal(t, "new-version", restaged.VersionKey)
+ assert.Greater(t, restaged.Revision, task.Revision)
+}
+
+func TestProcessCacheInvalidationTasksDropsPermanentFailures(t *testing.T) {
+ setupUserUpdateTestState(t)
+ now := time.Now().Unix()
+ require.NoError(t, DB.Create(&CacheInvalidationTask{
+ Kind: "unknown",
+ EntityKey: "entity",
+ VersionKey: "version",
+ NextAttempt: now,
+ CreatedAt: now,
+ UpdatedAt: now,
+ Revision: 1,
+ }).Error)
+ require.NoError(t, DB.Create(&CacheInvalidationTask{
+ Kind: cacheInvalidationKindUser,
+ EntityKey: "not-a-user-id",
+ VersionKey: "version",
+ NextAttempt: now,
+ CreatedAt: now,
+ UpdatedAt: now,
+ Revision: 1,
+ }).Error)
+
+ processCacheInvalidationTasks()
+
+ var count int64
+ require.NoError(t, DB.Model(&CacheInvalidationTask{}).Count(&count).Error)
+ assert.Zero(t, count)
+}
diff --git a/model/channel.go b/model/channel.go
index c0a29337..b1f57056 100644
--- a/model/channel.go
+++ b/model/channel.go
@@ -197,10 +197,26 @@ func (channel *Channel) GetKeys() []string {
}
func (channel *Channel) GetNextEnabledKey() (string, int, *types.MaxAPIError) {
+ channelID := channel.Id
+ foundCachedChannel := false
+ if common.MemoryCacheEnabled {
+ // Keep the cache read lock while taking the per-channel lock. Cache sync
+ // takes the same locks in this order before replacing the channel pointer.
+ channelSyncLock.RLock()
+ defer channelSyncLock.RUnlock()
+ if cached, ok := channelsIDM[channelID]; ok {
+ channel = cached
+ foundCachedChannel = true
+ }
+ }
+
// If not in multi-key mode, return the original key string directly.
if !channel.ChannelInfo.IsMultiKey {
return channel.Key, 0, nil
}
+ if common.MemoryCacheEnabled && !foundCachedChannel {
+ return "", 0, types.NewError(fmt.Errorf("渠道# %d，已不存在", channelID), types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry())
+ }
// Obtain all keys (split by \n)
keys := channel.GetKeys()
@@ -246,10 +262,13 @@ func (channel *Channel) GetNextEnabledKey() (string, int, *types.MaxAPIError) {
return keys[selectedIdx], selectedIdx, nil
case constant.MultiKeyModePolling:
// Use channel-specific lock to ensure thread-safe polling
-
- channelInfo, err := CacheGetChannelInfo(channel.Id)
- if err != nil {
- return "", 0, types.NewError(err, types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry())
+ channelInfo := &channel.ChannelInfo
+ if !common.MemoryCacheEnabled {
+ var err error
+ channelInfo, err = CacheGetChannelInfo(channel.Id)
+ if err != nil {
+ return "", 0, types.NewError(err, types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry())
+ }
}
defer func() {
if common.DebugEnabled {
diff --git a/model/channel_cache.go b/model/channel_cache.go
index b7d5a6f4..8896d1ad 100644
--- a/model/channel_cache.go
+++ b/model/channel_cache.go
@@ -88,12 +88,15 @@ func InitChannelCache() {
if channel.ChannelInfo.IsMultiKey {
channel.Keys = channel.GetKeys()
if channel.ChannelInfo.MultiKeyMode == constant.MultiKeyModePolling {
+ pollingLock := GetChannelPollingLock(i)
+ pollingLock.Lock()
if oldChannel, ok := channelsIDM[i]; ok {
// 存在旧的渠道，如果是多key且轮询，保留轮询索引信息
if oldChannel.ChannelInfo.IsMultiKey && oldChannel.ChannelInfo.MultiKeyMode == constant.MultiKeyModePolling {
channel.ChannelInfo.MultiKeyPollingIndex = oldChannel.ChannelInfo.MultiKeyPollingIndex
}
}
+ pollingLock.Unlock()
}
}
}
@@ -113,9 +116,15 @@ func SyncChannelCache(frequency int) {
}
func GetRandomSatisfiedChannel(group string, model string, retry int, requestPath string) (*Channel, error) {
+ return GetRandomSatisfiedChannelExcluding(group, model, retry, requestPath, nil)
+}
+
+// GetRandomSatisfiedChannelExcluding selects from the cache while excluding
+// channels already attempted by the current relay request.
+func GetRandomSatisfiedChannelExcluding(group string, model string, retry int, requestPath string, excludedChannelIDs map[int]struct{}) (*Channel, error) {
// if memory cache is disabled, get channel directly from database
if !common.MemoryCacheEnabled {
- return GetChannel(group, model, retry, requestPath)
+ return GetChannelExcluding(group, model, retry, requestPath, excludedChannelIDs)
}
channelSyncLock.RLock()
@@ -136,6 +145,18 @@ func GetRandomSatisfiedChannel(group string, model string, retry int, requestPat
}
}
+ if len(channels) == 0 {
+ return nil, nil
+ }
+ if len(excludedChannelIDs) > 0 {
+ availableChannels := make([]int, 0, len(channels))
+ for _, channelID := range channels {
+ if _, excluded := excludedChannelIDs[channelID]; !excluded {
+ availableChannels = append(availableChannels, channelID)
+ }
+ }
+ channels = availableChannels
+ }
if len(channels) == 0 {
return nil, nil
}
diff --git a/model/channel_cache_concurrency_test.go b/model/channel_cache_concurrency_test.go
new file mode 100644
index 00000000..5a2133ec
--- /dev/null
+++ b/model/channel_cache_concurrency_test.go
@@ -0,0 +1,98 @@
+package model
+
+import (
+ "fmt"
+ "sync"
+ "testing"
+
+ "github.com/MAX-API-Next/MAX-API/common"
+ "github.com/MAX-API-Next/MAX-API/constant"
+ "github.com/stretchr/testify/require"
+)
+
+func TestMultiKeyPollingRemainsSynchronizedDuringCacheRefresh(t *testing.T) {
+ clearPreferredOwnerTables(t)
+ const channelID = 9811
+
+ previousMemoryCacheEnabled := common.MemoryCacheEnabled
+ channelSyncLock.Lock()
+ previousGroup2Model2Channels := group2model2channels
+ previousChannelsIDM := channelsIDM
+ previousAdvancedConfig := channel2advancedCustomConfig
+ previousAdvancedConfigError := channel2advancedCustomConfigError
+ group2model2channels = nil
+ channelsIDM = nil
+ channel2advancedCustomConfig = nil
+ channel2advancedCustomConfigError = nil
+ channelSyncLock.Unlock()
+ channelPollingLocks.Delete(channelID)
+ t.Cleanup(func() {
+ clearPreferredOwnerTables(t)
+ channelPollingLocks.Delete(channelID)
+ channelSyncLock.Lock()
+ group2model2channels = previousGroup2Model2Channels
+ channelsIDM = previousChannelsIDM
+ channel2advancedCustomConfig = previousAdvancedConfig
+ channel2advancedCustomConfigError = previousAdvancedConfigError
+ channelSyncLock.Unlock()
+ common.MemoryCacheEnabled = previousMemoryCacheEnabled
+ })
+
+ common.MemoryCacheEnabled = true
+ key, index, keyErr := (&Channel{Id: channelID + 1, Key: "single-key"}).GetNextEnabledKey()
+ require.Nil(t, keyErr)
+ require.Equal(t, "single-key", key)
+ require.Zero(t, index)
+
+ channel := &Channel{
+ Id: channelID,
+ Type: constant.ChannelTypeOpenAI,
+ Key: "first-key\nsecond-key",
+ Status: common.ChannelStatusEnabled,
+ Name: "cache-refresh-polling-channel",
+ Models: "gpt-cache-refresh",
+ Group: "default",
+ ChannelInfo: ChannelInfo{
+ IsMultiKey: true,
+ MultiKeyMode: constant.MultiKeyModePolling,
+ },
+ }
+ priority := int64(0)
+ require.NoError(t, DB.Create(channel).Error)
+ require.NoError(t, DB.Create(&Ability{
+ Group: "default", Model: "gpt-cache-refresh", ChannelId: channelID,
+ Enabled: true, Priority: &priority, Weight: 100,
+ }).Error)
+ InitChannelCache()
+
+ start := make(chan struct{})
+ errs := make(chan error, 4)
+ var waitGroup sync.WaitGroup
+ for range 3 {
+ waitGroup.Add(1)
+ go func() {
+ defer waitGroup.Done()
+ <-start
+ for i := 0; i < 100; i++ {
+ if _, _, err := channel.GetNextEnabledKey(); err != nil {
+ errs <- err
+ return
+ }
+ }
+ }()
+ }
+ waitGroup.Add(1)
+ go func() {
+ defer waitGroup.Done()
+ <-start
+ for i := 0; i < 40; i++ {
+ InitChannelCache()
+ }
+ }()
+ close(start)
+ waitGroup.Wait()
+ close(errs)
+ for err := range errs {
+ require.NoError(t, err, fmt.Sprintf("polling during cache refresh failed: %v", err))
+ }
+}
diff --git a/model/channel_selection_test.go b/model/channel_selection_test.go
index b2c98008..389865b2 100644
--- a/model/channel_selection_test.go
+++ b/model/channel_selection_test.go
@@ -158,3 +158,42 @@ func TestMemoryCacheSelectionSkipsInvalidAdvancedCustomConfig(t *testing.T) {
require.NotNil(t, channel)
assert.Equal(t, 2, channel.Id)
}
+
+func TestChannelSelectionExcludesPreviouslyTriedChannel(t *testing.T) {
+ clearPreferredOwnerTables(t)
+ t.Cleanup(func() {
+ clearPreferredOwnerTables(t)
+ channelSyncLock.Lock()
+ group2model2channels = nil
+ channelsIDM = nil
+ channel2advancedCustomConfig = nil
+ channel2advancedCustomConfigError = nil
+ channelSyncLock.Unlock()
+ })
+
+ const (
+ groupName = "default"
+ modelName = "gpt-exclude-failed-channel"
+ )
+ insertChannelSelectionCandidate(t, 1, constant.ChannelTypeOpenAI, "", modelName, groupName, 10)
+ insertChannelSelectionCandidate(t, 2, constant.ChannelTypeOpenAI, "", modelName, groupName, 10)
+ excluded := map[int]struct{}{1: {}}
+
+ previousMemoryCacheEnabled := common.MemoryCacheEnabled
+ t.Cleanup(func() {
+ common.MemoryCacheEnabled = previousMemoryCacheEnabled
+ })
+
+ common.MemoryCacheEnabled = false
+ channel, err := GetChannelExcluding(groupName, modelName, 0, "/v1/chat/completions", excluded)
+ require.NoError(t, err)
+ require.NotNil(t, channel)
+ assert.Equal(t, 2, channel.Id)
+
+ common.MemoryCacheEnabled = true
+ InitChannelCache()
+ channel, err = GetRandomSatisfiedChannelExcluding(groupName, modelName, 0, "/v1/chat/completions", excluded)
+ require.NoError(t, err)
+ require.NotNil(t, channel)
+ assert.Equal(t, 2, channel.Id)
+}
diff --git a/model/checkin.go b/model/checkin.go
index d9032133..f240bd38 100644
--- a/model/checkin.go
+++ b/model/checkin.go
@@ -2,7 +2,6 @@ package model
import (
"errors"
- "fmt"
"math/rand"
"time"
@@ -94,6 +93,7 @@ func UserCheckin(userId int) (*Checkin, error) {
// userCheckinWithTransaction 使用事务执行签到（适用于 MySQL 和 PostgreSQL）
func userCheckinWithTransaction(checkin *Checkin, userId int, quotaAwarded int) (*Checkin, error) {
+ var cacheTask CacheInvalidationTask
err := DB.Transaction(func(tx *gorm.DB) error {
// 步骤1: 创建签到记录
// 数据库有唯一约束 (user_id, checkin_date)，可以防止并发重复签到
@@ -107,19 +107,16 @@ func userCheckinWithTransaction(checkin *Checkin, userId int, quotaAwarded int)
return errors.New("签到失败：更新额度出错")
}
- return nil
+ var err error
+ cacheTask, err = stageUserCacheInvalidationTx(tx, userId, false)
+ return err
})
if err != nil {
return nil, err
}
- // 事务成功后同步失效缓存，下一次读取以数据库结果回填。
- if cacheErr := invalidateUserQuotaCache(userId); cacheErr != nil {
- common.SysLog(fmt.Sprintf("failed to invalidate check-in quota cache (user_id=%d): %v", userId, cacheErr))
- enqueueUserCacheInvalidationRetry(userId, cacheErr)
- }
-
+ dispatchStagedCacheInvalidation(cacheTask)
return checkin, nil
}
diff --git a/model/log.go b/model/log.go
index 46f03be8..71f980af 100644
--- a/model/log.go
+++ b/model/log.go
@@ -14,8 +14,10 @@ import (
"github.com/bytedance/gopkg/util/gopool"
"github.com/gin-gonic/gin"
+ "github.com/google/uuid"
"gorm.io/gorm"
+ "gorm.io/gorm/clause"
)
func applyExplicitLogTextFilter(tx *gorm.DB, column string, value string) (*gorm.DB, error) {
@@ -60,6 +62,15 @@ type Log struct {
LogId int `json:"log_id,omitempty" gorm:"-"`
}
+// BillingLogReceipt keeps task-billing log idempotency off the high-volume
+// logs table. The receipt and its log row are committed in one LOG_DB transaction.
+type BillingLogReceipt struct {
+ ID int64 `gorm:"primaryKey"`
+ OperationKey string `gorm:"type:varchar(191);uniqueIndex;not null"`
+ ClaimToken string `gorm:"type:varchar(36);not null;default:''"`
+ CreatedAt int64 `gorm:"not null"`
+}
+
func (log *Log) BeforeSave(*gorm.DB) error {
log.syncRetryMarker()
return nil
@@ -642,8 +653,24 @@ type RecordTaskBillingLogParams struct {
}
func RecordTaskBillingLog(params RecordTaskBillingLogParams) {
+ if err := recordTaskBillingLog("", params); err != nil {
+ common.SysLog("failed to record task billing log: " + err.Error())
+ }
+}
+
+func RecordTaskBillingLogOnce(operationKey string, params RecordTaskBillingLogParams) error {
+ if operationKey == "" {
+ return errors.New("task billing log operation key is required")
+ }
+ return recordTaskBillingLog(operationKey, params)
+}
+
+func recordTaskBillingLog(operationKey string, params RecordTaskBillingLogParams) error {
if params.LogType == LogTypeConsume && !common.LogConsumeEnabled {
- return
+ return nil
+ }
+ if LOG_DB == nil {
+ return errors.New("log database is not initialized")
}
username, _ := GetUsernameById(params.UserId, false)
tokenName := ""
@@ -667,9 +694,45 @@ func RecordTaskBillingLog(params RecordTaskBillingLogParams) {
Group: params.Group,
Other: common.MapToJsonStr(params.Other),
}
- err := LOG_DB.Create(log).Error
- if err != nil {
- common.SysLog("failed to record task billing log: " + err.Error())
+ inserted := true
+ if operationKey == "" {
+ if err := LOG_DB.Create(log).Error; err != nil {
+ return err
+ }
+ } else {
+ err := LOG_DB.Transaction(func(tx *gorm.DB) error {
+ claimToken := uuid.NewString()
+ receipt := BillingLogReceipt{
+ OperationKey: operationKey,
+ ClaimToken: claimToken,
+ CreatedAt: createdAt,
+ }
+ result := tx.Clauses(clause.OnConflict{
+ Columns: []clause.Column{{Name: "operation_key"}},
+ DoNothing: true,
+ }).Create(&receipt)
+ if result.Error != nil {
+ return result.Error
+ }
+ // MySQL implements DoNothing as a no-op UPDATE. Its RowsAffected can
+ // report 1 when clientFoundRows is enabled, so verify ownership using
+ // a per-attempt claim instead of relying on driver-specific row counts.
+ var storedReceipt BillingLogReceipt
+ if err := tx.Select("claim_token").Where("operation_key = ?", operationKey).Take(&storedReceipt).Error; err != nil {
+ return err
+ }
+ if storedReceipt.ClaimToken != claimToken {
+ inserted = false
+ return nil
+ }
+ return tx.Create(log).Error
+ })
+ if err != nil {
+ return err
+ }
+ }
+ if !inserted {
+ return nil
}
if params.LogType == LogTypeConsume && common.DataExportEnabled {
nodeName := params.NodeName
@@ -688,6 +751,7 @@ func RecordTaskBillingLog(params RecordTaskBillingLogParams) {
NodeName: nodeName,
})
}
+ return nil
}
func GetAllLogs(params LogQueryParams) (logs []*Log, total int64, err error) {
@@ -931,8 +995,10 @@ func SumUsedQuota(params LogQueryParams) (stat Stat, err error) {
return stat, nil
}
+const sumUsedTokenSelect = "COALESCE(SUM(prompt_tokens), 0) + COALESCE(SUM(completion_tokens), 0)"
+
func SumUsedToken(logType int, startTimestamp int64, endTimestamp int64, modelName string, username string, tokenName string) (token int) {
- tx := LOG_DB.Table("logs").Select("ifnull(sum(prompt_tokens),0) + ifnull(sum(completion_tokens),0)")
+ tx := LOG_DB.Table("logs").Select(sumUsedTokenSelect)
if username != "" {
tx = tx.Where("username = ?", username)
}
@@ -979,17 +1045,30 @@ func DeleteOldLogBatch(ctx context.Context, targetTimestamp int64, limit int) (i
return 0, err
}
if len(ids) == 0 {
+ _, err := DeleteOldBillingLogReceiptsBatch(ctx, targetTimestamp, limit)
+ if err != nil {
+ return 0, err
+ }
return 0, nil
}
if err := ctx.Err(); err != nil {
return 0, err
}
- result := LOG_DB.WithContext(ctx).Where("id IN ?", ids).Delete(&Log{})
- if nil != result.Error {
- return 0, result.Error
+ var rowsAffected int64
+ err = LOG_DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
+ result := tx.Where("id IN ?", ids).Delete(&Log{})
+ if nil != result.Error {
+ return result.Error
+ }
+ rowsAffected = result.RowsAffected
+ _, err := deleteOldBillingLogReceiptsBatchTx(ctx, tx, targetTimestamp, limit)
+ return err
+ })
+ if err != nil {
+ return 0, err
}
- return result.RowsAffected, nil
+ return rowsAffected, nil
}
func DeleteOldLog(ctx context.Context, targetTimestamp int64, limit int) (int64, error) {
@@ -1009,7 +1088,67 @@ func DeleteOldLog(ctx context.Context, targetTimestamp int64, limit int) (int64,
}
total += rowsAffected
if rowsAffected < int64(limit) {
+ if _, err := DeleteOldBillingLogReceipts(ctx, targetTimestamp, limit); err != nil {
+ return total, err
+ }
return total, nil
}
}
}
+
+func DeleteOldBillingLogReceipts(ctx context.Context, targetTimestamp int64, limit int) (int64, error) {
+ if limit <= 0 {
+ limit = 100
+ }
+
+ var total int64
+ for {
+ if err := ctx.Err(); err != nil {
+ return total, err
+ }
+ rowsAffected, err := DeleteOldBillingLogReceiptsBatch(ctx, targetTimestamp, limit)
+ if err != nil {
+ return total, err
+ }
+ total += rowsAffected
+ if rowsAffected < int64(limit) {
+ return total, nil
+ }
+ }
+}
+
+func DeleteOldBillingLogReceiptsBatch(ctx context.Context, targetTimestamp int64, limit int) (int64, error) {
+ if limit <= 0 {
+ limit = 100
+ }
+ if err := ctx.Err(); err != nil {
+ return 0, err
+ }
+ return deleteOldBillingLogReceiptsBatchTx(ctx, LOG_DB, targetTimestamp, limit)
+}
+
+func deleteOldBillingLogReceiptsBatchTx(ctx context.Context, db *gorm.DB, targetTimestamp int64, limit int) (int64, error) {
+ if db == nil {
+ return 0, errors.New("log database is not initialized")
+ }
+ ids := make([]int64, 0, limit)
+ if err := db.WithContext(ctx).
+ Model(&BillingLogReceipt{}).
+ Where("created_at < ?", targetTimestamp).
+ Order("created_at ASC, id ASC").
+ Limit(limit).
+ Pluck("id", &ids).Error; err != nil {
+ return 0, err
+ }
+ if len(ids) == 0 {
+ return 0, nil
+ }
+ if err := ctx.Err(); err != nil {
+ return 0, err
+ }
+ result := db.WithContext(ctx).Where("id IN ?", ids).Delete(&BillingLogReceipt{})
+ if result.Error != nil {
+ return 0, result.Error
+ }
+ return result.RowsAffected, nil
+}
diff --git a/model/log_cleanup_test.go b/model/log_cleanup_test.go
index 92794a4e..9f4a81fe 100644
--- a/model/log_cleanup_test.go
+++ b/model/log_cleanup_test.go
@@ -104,3 +104,66 @@ func TestDeleteOldLogDeletesAllMatchingRows(t *testing.T) {
require.NoError(t, LOG_DB.Model(&Log{}).Count(&total).Error)
assert.EqualValues(t, 1, total)
}
+
+func TestDeleteOldLogBatchDeletesOldBillingLogReceipts(t *testing.T) {
+ truncateTables(t)
+
+ logs := []Log{
+ {UserId: 1, CreatedAt: 10, Type: LogTypeConsume},
+ {UserId: 1, CreatedAt: 1000, Type: LogTypeConsume},
+ }
+ require.NoError(t, LOG_DB.Create(&logs).Error)
+ receipts := []BillingLogReceipt{
+ {OperationKey: "old-receipt", ClaimToken: "old", CreatedAt: 10},
+ {OperationKey: "new-receipt", ClaimToken: "new", CreatedAt: 1000},
+ }
+ require.NoError(t, LOG_DB.Create(&receipts).Error)
+
+ rows, err := DeleteOldLogBatch(context.Background(), 100, 10)
+ require.NoError(t, err)
+ assert.EqualValues(t, 1, rows)
+
+ var oldReceiptCount int64
+ require.NoError(t, LOG_DB.Model(&BillingLogReceipt{}).Where("operation_key = ?", "old-receipt").Count(&oldReceiptCount).Error)
+ assert.Zero(t, oldReceiptCount)
+ var newReceiptCount int64
+ require.NoError(t, LOG_DB.Model(&BillingLogReceipt{}).Where("operation_key = ?", "new-receipt").Count(&newReceiptCount).Error)
+ assert.EqualValues(t, 1, newReceiptCount)
+}
+
+func TestDeleteOldLogBatchDeletesOldBillingLogReceiptsWithoutOldLogs(t *testing.T) {
+ truncateTables(t)
+
+ require.NoError(t, LOG_DB.Create(&BillingLogReceipt{OperationKey: "old-receipt-only", ClaimToken: "old", CreatedAt: 10}).Error)
+
+ rows, err := DeleteOldLogBatch(context.Background(), 100, 10)
+ require.NoError(t, err)
+ assert.Zero(t, rows)
+
+ var count int64
+ require.NoError(t, LOG_DB.Model(&BillingLogReceipt{}).Where("operation_key = ?", "old-receipt-only").Count(&count).Error)
+ assert.Zero(t, count)
+}
+
+func TestDeleteOldLogDeletesAllOldBillingLogReceipts(t *testing.T) {
+ truncateTables(t)
+
+ receipts := []BillingLogReceipt{
+ {OperationKey: "old-receipt-1", ClaimToken: "old-1", CreatedAt: 10},
+ {OperationKey: "old-receipt-2", ClaimToken: "old-2", CreatedAt: 11},
+ {OperationKey: "old-receipt-3", ClaimToken: "old-3", CreatedAt: 12},
+ {OperationKey: "new-receipt", ClaimToken: "new", CreatedAt: 1000},
+ }
+ require.NoError(t, LOG_DB.Create(&receipts).Error)
+
+ rows, err := DeleteOldLog(context.Background(), 100, 2)
+ require.NoError(t, err)
+ assert.Zero(t, rows)
+
+ var oldReceiptCount int64
+ require.NoError(t, LOG_DB.Model(&BillingLogReceipt{}).Where("created_at < ?", 100).Count(&oldReceiptCount).Error)
+ assert.Zero(t, oldReceiptCount)
+ var newReceiptCount int64
+ require.NoError(t, LOG_DB.Model(&BillingLogReceipt{}).Where("operation_key = ?", "new-receipt").Count(&newReceiptCount).Error)
+ assert.EqualValues(t, 1, newReceiptCount)
+}
diff --git a/model/log_test.go b/model/log_test.go
index 3c44f29b..06679274 100644
--- a/model/log_test.go
+++ b/model/log_test.go
@@ -27,6 +27,40 @@ func withLogAuditSettings(t *testing.T, requestEnabled bool, responseEnabled boo
})
}
+func TestRecordTaskBillingLogOnceUsesReceiptClaimForIdempotency(t *testing.T) {
+ require.NoError(t, LOG_DB.Where("1 = 1").Delete(&BillingLogReceipt{}).Error)
+ require.NoError(t, LOG_DB.Where("1 = 1").Delete(&Log{}).Error)
+ t.Cleanup(func() {
+ require.NoError(t, LOG_DB.Where("1 = 1").Delete(&BillingLogReceipt{}).Error)
+ require.NoError(t, LOG_DB.Where("1 = 1").Delete(&Log{}).Error)
+ })
+
+ params := RecordTaskBillingLogParams{
+ UserId: 1,
+ LogType: LogTypeSystem,
+ Content: "idempotent task billing effect",
+ }
+ require.NoError(t, RecordTaskBillingLogOnce("task:test:effect", params))
+ require.NoError(t, RecordTaskBillingLogOnce("task:test:effect", params))
+
+ var logCount int64
+ require.NoError(t, LOG_DB.Model(&Log{}).Count(&logCount).Error)
+ require.EqualValues(t, 1, logCount)
+
+ var receipt BillingLogReceipt
+ require.NoError(t, LOG_DB.Where("operation_key = ?", "task:test:effect").Take(&receipt).Error)
+ require.NotEmpty(t, receipt.ClaimToken)
+
+ require.NoError(t, LOG_DB.Create(&BillingLogReceipt{
+ OperationKey: "task:legacy:effect",
+ ClaimToken: "",
+ CreatedAt: time.Now().Unix(),
+ }).Error)
+ require.NoError(t, RecordTaskBillingLogOnce("task:legacy:effect", params))
+ require.NoError(t, LOG_DB.Model(&Log{}).Count(&logCount).Error)
+ require.EqualValues(t, 1, logCount)
+}
+
func markRetryLogBackfillCompletedForTest(t *testing.T) {
t.Helper()
markerKey := logRetryMarkerBackfillCompletionKey()
@@ -76,6 +110,12 @@ func TestLogQuotaFilterIndexes(t *testing.T) {
require.True(t, db.Migrator().HasIndex(&Log{}, "idx_logs_type_quota_created_at"))
}
+func TestLogSchemaDoesNotAddBillingSettlementIdempotencyColumn(t *testing.T) {
+ db := newRetryBackfillTestDB(t, &Log{})
+
+ require.False(t, db.Migrator().HasColumn(&Log{}, "billing_settlement_operation_key"))
+}
+
func TestGetAllLogsRetryFilter(t *testing.T) {
require.NoError(t, LOG_DB.Where("1 = 1").Delete(&Log{}).Error)
t.Cleanup(func() {
@@ -1151,3 +1191,7 @@ func TestFormatUserLogsHidesAuditContentWhenDisabled(t *testing.T) {
require.NotContains(t, other, "admin_info")
require.NotContains(t, other, "audit_info")
}
+
+func TestSumUsedTokenUsesPortableAggregate(t *testing.T) {
+ require.Equal(t, "COALESCE(SUM(prompt_tokens), 0) + COALESCE(SUM(completion_tokens), 0)", sumUsedTokenSelect)
+}
diff --git a/model/main.go b/model/main.go
index 496685f2..131bb2f1 100644
--- a/model/main.go
+++ b/model/main.go
@@ -215,7 +215,11 @@ func chooseDB(envName string, isLog bool) (*gorm.DB, error) {
}
// Use SQLite
common.SysLog("SQL_DSN not set, using SQLite as database")
- common.UsingSQLite = true
+ if !isLog {
+ common.UsingSQLite = true
+ } else {
+ common.LogSqlType = common.DatabaseTypeSQLite
+ }
return gorm.Open(sqlite.Open(common.SQLitePath), &gorm.Config{
PrepareStmt: true, // precompile SQL
})
@@ -295,8 +299,6 @@ func InitLogDB() (err error) {
}
func migrateDB() error {
- // Migrate price_amount column from float/double to decimal for existing tables
- migrateSubscriptionPlanPriceAmount()
// Migrate model_limits column from varchar to text for existing tables
if err := migrateTokenModelLimitsToText(); err != nil {
return err
@@ -317,6 +319,7 @@ func migrateDB() error {
&Redemption{},
&Ability{},
&Log{},
+ &BillingLogReceipt{},
&Midjourney{},
&TopUp{},
&QuotaData{},
@@ -333,6 +336,9 @@ func migrateDB() error {
&SubscriptionOrder{},
&UserSubscription{},
&SubscriptionPreConsumeRecord{},
+ &BillingSettlement{},
+ &BillingPreConsumeSelection{},
+ &CacheInvalidationTask{},
&CustomOAuthProvider{},
&UserOAuthBinding{},
&PerfMetric{},
@@ -342,98 +348,13 @@ func migrateDB() error {
if err != nil {
return err
}
- if err := migrateQuotaDataAggregateKeysOnStartup(); err != nil {
- return err
- }
- if err := migrateUserOAuthIdentityConstraints(); err != nil {
- return err
- }
- if err := backfillUserNormalizedEmails(); err != nil {
- return err
- }
- if common.UsingSQLite {
- if err := ensureSubscriptionPlanTableSQLite(); err != nil {
- return err
- }
- } else {
- if err := DB.AutoMigrate(&SubscriptionPlan{}); err != nil {
- return err
- }
- }
- if os.Getenv("LOG_SQL_DSN") == "" {
- LOG_DB = DB
- scheduleLogRetryMarkerBackfill()
- }
- return nil
-}
-
-func migrateDBFast() error {
- if err := migrateUserQuotaColumnsToBigInt(); err != nil {
+ if err := migrateMetadataNameKeys(); err != nil {
return err
}
- if err := migrateTokenQuotaColumnsToBigInt(); err != nil {
+ // Run the legacy subscription price migration only after current core schemas exist.
+ if err := migrateSubscriptionPlanPriceAmount(); err != nil {
return err
}
-
- var wg sync.WaitGroup
-
- migrations := []struct {
- model interface{}
- name string
- }{
- {&Channel{}, "Channel"},
- {&Token{}, "Token"},
- {&User{}, "User"},
- {&PasskeyCredential{}, "PasskeyCredential"},
- {&Option{}, "Option"},
- {&Redemption{}, "Redemption"},
- {&Ability{}, "Ability"},
- {&Log{}, "Log"},
- {&Midjourney{}, "Midjourney"},
- {&TopUp{}, "TopUp"},
- {&QuotaData{}, "QuotaData"},
- {&QuotaDataSnapshot{}, "QuotaDataSnapshot"},
- {&QuotaDataSnapshotRetry{}, "QuotaDataSnapshotRetry"},
- {&Task{}, "Task"},
- {&Model{}, "Model"},
- {&Vendor{}, "Vendor"},
- {&PrefillGroup{}, "PrefillGroup"},
- {&Setup{}, "Setup"},
- {&TwoFA{}, "TwoFA"},
- {&TwoFABackupCode{}, "TwoFABackupCode"},
- {&Checkin{}, "Checkin"},
- {&SubscriptionOrder{}, "SubscriptionOrder"},
- {&UserSubscription{}, "UserSubscription"},
- {&SubscriptionPreConsumeRecord{}, "SubscriptionPreConsumeRecord"},
- {&CustomOAuthProvider{}, "CustomOAuthProvider"},
- {&UserOAuthBinding{}, "UserOAuthBinding"},
- {&PerfMetric{}, "PerfMetric"},
- {&SystemTask{}, "SystemTask"},
- {&SystemInstance{}, "SystemInstance"},
- }
- // 动态计算migration数量，确保errChan缓冲区足够大
- errChan := make(chan error, len(migrations))
-
- for _, m := range migrations {
- wg.Add(1)
- go func(model interface{}, name string) {
- defer wg.Done()
- if err := DB.AutoMigrate(model); err != nil {
- errChan <- fmt.Errorf("failed to migrate %s: %v", name, err)
- }
- }(m.model, m.name)
- }
-
- // Wait for all migrations to complete
- wg.Wait()
- close(errChan)
-
- // Check for any errors
- for err := range errChan {
- if err != nil {
- return err
- }
- }
if err := migrateQuotaDataAggregateKeysOnStartup(); err != nil {
return err
}
@@ -452,7 +373,10 @@ func migrateDBFast() error {
return err
}
}
- common.SysLog("database migrated")
+ if os.Getenv("LOG_SQL_DSN") == "" {
+ LOG_DB = DB
+ scheduleLogRetryMarkerBackfill()
+ }
return nil
}
@@ -683,7 +607,7 @@ func userOAuthIdentityGeneratedColumnSQL(identity userOAuthIdentityMigration) st
func migrateLOGDB() error {
var err error
- if err = LOG_DB.AutoMigrate(&Log{}); err != nil {
+ if err = LOG_DB.AutoMigrate(&Log{}, &BillingLogReceipt{}); err != nil {
return err
}
scheduleLogRetryMarkerBackfill()
@@ -1022,11 +946,11 @@ func isZeroColumnDefault(defaultValue sql.NullString) bool {
// migrateSubscriptionPlanPriceAmount migrates price_amount column from float/double to decimal(10,6)
// This is safe to run multiple times - it checks the column type first
-func migrateSubscriptionPlanPriceAmount() {
+func migrateSubscriptionPlanPriceAmount() error {
// SQLite doesn't support ALTER COLUMN, and its type affinity handles this automatically
// Skip early to avoid GORM parsing the existing table DDL which may cause issues
if common.UsingSQLite {
- return
+ return nil
}
tableName := "subscription_plans"
@@ -1034,12 +958,12 @@ func migrateSubscriptionPlanPriceAmount() {
// Check if table exists first
if !DB.Migrator().HasTable(tableName) {
- return
+ return nil
}
// Check if column exists
if !DB.Migrator().HasColumn(&SubscriptionPlan{}, columnName) {
- return
+ return nil
}
var alterSQL string
@@ -1049,9 +973,9 @@ func migrateSubscriptionPlanPriceAmount() {
if err := DB.Raw(`SELECT data_type FROM information_schema.columns
WHERE table_schema = current_schema() AND table_name = ? AND column_name = ?`,
tableName, columnName).Scan(&dataType).Error; err != nil {
- common.SysLog(fmt.Sprintf("Warning: failed to query metadata for %s.%s: %v", tableName, columnName, err))
+ return fmt.Errorf("failed to query metadata for %s.%s: %w", tableName, columnName, err)
} else if dataType == "numeric" {
- return // Already decimal/numeric
+ return nil // Already decimal/numeric
}
alterSQL = fmt.Sprintf(`ALTER TABLE %s ALTER COLUMN %s TYPE decimal(10,6) USING %s::decimal(10,6)`,
tableName, columnName, columnName)
@@ -1061,23 +985,23 @@ func migrateSubscriptionPlanPriceAmount() {
if err := DB.Raw(`SELECT COLUMN_TYPE FROM information_schema.columns
WHERE table_schema = DATABASE() AND table_name = ? AND column_name = ?`,
tableName, columnName).Scan(&columnType).Error; err != nil {
- common.SysLog(fmt.Sprintf("Warning: failed to query metadata for %s.%s: %v", tableName, columnName, err))
+ return fmt.Errorf("failed to query metadata for %s.%s: %w", tableName, columnName, err)
} else if strings.HasPrefix(strings.ToLower(columnType), "decimal") {
- return // Already decimal
+ return nil // Already decimal
}
alterSQL = fmt.Sprintf("ALTER TABLE %s MODIFY COLUMN %s decimal(10,6) NOT NULL DEFAULT 0",
tableName, columnName)
} else {
- return
+ return nil
}
if alterSQL != "" {
if err := DB.Exec(alterSQL).Error; err != nil {
- common.SysLog(fmt.Sprintf("Warning: failed to migrate %s.%s to decimal: %v", tableName, columnName, err))
- } else {
- common.SysLog(fmt.Sprintf("Successfully migrated %s.%s to decimal(10,6)", tableName, columnName))
+ return fmt.Errorf("failed to migrate %s.%s to decimal: %w", tableName, columnName, err)
}
+ common.SysLog(fmt.Sprintf("Successfully migrated %s.%s to decimal(10,6)", tableName, columnName))
}
+ return nil
}
func closeDB(db *gorm.DB) error {
diff --git a/model/main_migration_test.go b/model/main_migration_test.go
index d9a51b6a..dec3ccaa 100644
--- a/model/main_migration_test.go
+++ b/model/main_migration_test.go
@@ -6,6 +6,10 @@ import (
"reflect"
"strings"
"testing"
+
+ "github.com/MAX-API-Next/MAX-API/common"
+ "github.com/glebarez/sqlite"
+ "gorm.io/gorm"
)
func TestTokenQuotaFieldsUseInt64(t *testing.T) {
@@ -61,3 +65,39 @@ func TestStartupMigrationSchemaChecksUseModelValues(t *testing.T) {
}
}
}
+
+func TestMigrateDBRunsCoreAutoMigrateBeforeSubscriptionPlanPriceMigrationFailure(t *testing.T) {
+ previousDB := DB
+ previousLogDB := LOG_DB
+ previousSQLite := common.UsingSQLite
+ previousMySQL := common.UsingMySQL
+ previousPostgreSQL := common.UsingPostgreSQL
+ t.Cleanup(func() {
+ DB = previousDB
+ LOG_DB = previousLogDB
+ common.UsingSQLite = previousSQLite
+ common.UsingMySQL = previousMySQL
+ common.UsingPostgreSQL = previousPostgreSQL
+ })
+
+ db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
+ if err != nil {
+ t.Fatalf("failed to open test database: %v", err)
+ }
+ DB = db
+ LOG_DB = db
+ common.UsingSQLite = false
+ common.UsingMySQL = true
+ common.UsingPostgreSQL = false
+
+ if err := db.AutoMigrate(&SubscriptionPlan{}); err != nil {
+ t.Fatalf("failed to create subscription plans table: %v", err)
+ }
+
+ if err := migrateDB(); err == nil {
+ t.Fatal("expected subscription price migration failure to be returned")
+ }
+ if !db.Migrator().HasTable(&Channel{}) {
+ t.Fatal("expected core schema auto-migration to run before subscription price migration failure")
+ }
+}
diff --git a/model/metadata_name_key.go b/model/metadata_name_key.go
new file mode 100644
index 00000000..73f49da6
--- /dev/null
+++ b/model/metadata_name_key.go
@@ -0,0 +1,197 @@
+package model
+
+import (
+ "fmt"
+ "sort"
+
+ "github.com/MAX-API-Next/MAX-API/common"
+
+ "gorm.io/gorm"
+)
+
+const (
+ vendorNameKeyIndex = "uk_vendors_active_name_key"
+ modelNameKeyIndex = "uk_models_active_name_key"
+ legacyVendorNameKeyIndex = "uk_vendor_name_delete_at"
+ legacyModelNameKeyIndex = "uk_model_name_delete_at"
+)
+
+func activeMetadataNameKey(name string) string {
+ return name
+}
+
+func retiredMetadataNameKey(kind string, id int) string {
+ return fmt.Sprintf("deleted:%s:%d", kind, id)
+}
+
+func softDeleteMetadataWithNameKey(tx *gorm.DB, entity interface{}, id int, nameKey string) error {
+ result := tx.Model(entity).Where("id = ?", id).UpdateColumn("name_key", nameKey)
+ if result.Error != nil {
+ return result.Error
+ }
+ if result.RowsAffected == 0 {
+ return gorm.ErrRecordNotFound
+ }
+ return tx.Delete(entity, id).Error
+}
+
+// migrateMetadataNameKeys replaces the nullable deleted_at composite indexes
+// with non-null active-name keys. MySQL and PostgreSQL treat NULLs as distinct
+// in unique indexes, so the old index cannot prevent concurrent active names.
+func migrateMetadataNameKeys() error {
+ if DB == nil {
+ return nil
+ }
+ if common.UsingMySQL {
+ if err := backfillMetadataNameKeys(DB); err != nil {
+ return err
+ }
+ if err := resolveDuplicateMetadataNameKeys(DB); err != nil {
+ return err
+ }
+ if err := ensureMetadataNameKeyIndexes(DB); err != nil {
+ return err
+ }
+ return dropLegacyMetadataNameIndexes(DB)
+ }
+ return DB.Transaction(func(tx *gorm.DB) error {
+ if err := backfillMetadataNameKeys(tx); err != nil {
+ return err
+ }
+ if err := resolveDuplicateMetadataNameKeys(tx); err != nil {
+ return err
+ }
+ if err := ensureMetadataNameKeyIndexes(tx); err != nil {
+ return err
+ }
+ return dropLegacyMetadataNameIndexes(tx)
+ })
+}
+
+// resolveDuplicateMetadataNameKeys preserves the earliest active record for
+// each name and retires later duplicates before the unique indexes are built.
+// Legacy databases may already contain duplicate active names because the old
+// (name, deleted_at) index treats NULL deleted_at values as distinct.
+func resolveDuplicateMetadataNameKeys(tx *gorm.DB) error {
+ var vendors []Vendor
+ if err := tx.Unscoped().Where("deleted_at IS NULL").Order("id ASC").Find(&vendors).Error; err != nil {
+ return fmt.Errorf("load active vendors for duplicate name keys: %w", err)
+ }
+ if err := resolveDuplicateVendorNameKeys(tx, vendors); err != nil {
+ return err
+ }
+
+ var models []Model
+ if err := tx.Unscoped().Where("deleted_at IS NULL").Order("id ASC").Find(&models).Error; err != nil {
+ return fmt.Errorf("load active models for duplicate name keys: %w", err)
+ }
+ return resolveDuplicateModelNameKeys(tx, models)
+}
+
+func resolveDuplicateVendorNameKeys(tx *gorm.DB, vendors []Vendor) error {
+ seen := make(map[string]struct{}, len(vendors))
+ // Keep the ordering deterministic even if a database ignores the requested
+ // order for an unusual legacy query plan.
+ sort.SliceStable(vendors, func(i, j int) bool { return vendors[i].Id < vendors[j].Id })
+ for _, vendor := range vendors {
+ if _, exists := seen[vendor.NameKey]; !exists {
+ seen[vendor.NameKey] = struct{}{}
+ continue
+ }
+ nameKey := retiredMetadataNameKey("vendor", vendor.Id)
+ if err := tx.Unscoped().Model(&Vendor{}).Where("id = ?", vendor.Id).UpdateColumn("name_key", nameKey).Error; err != nil {
+ return fmt.Errorf("retire duplicate vendor %d name key: %w", vendor.Id, err)
+ }
+ }
+ return nil
+}
+
+func resolveDuplicateModelNameKeys(tx *gorm.DB, models []Model) error {
+ seen := make(map[string]struct{}, len(models))
+ sort.SliceStable(models, func(i, j int) bool { return models[i].Id < models[j].Id })
+ for _, model := range models {
+ if _, exists := seen[model.NameKey]; !exists {
+ seen[model.NameKey] = struct{}{}
+ continue
+ }
+ nameKey := retiredMetadataNameKey("model", model.Id)
+ if err := tx.Unscoped().Model(&Model{}).Where("id = ?", model.Id).UpdateColumn("name_key", nameKey).Error; err != nil {
+ return fmt.Errorf("retire duplicate model %d name key: %w", model.Id, err)
+ }
+ }
+ return nil
+}
+
+func backfillMetadataNameKeys(tx *gorm.DB) error {
+ var vendors []Vendor
+ if err := tx.Unscoped().Find(&vendors).Error; err != nil {
+ return fmt.Errorf("load vendors for name-key migration: %w", err)
+ }
+ for _, vendor := range vendors {
+ nameKey := activeMetadataNameKey(vendor.Name)
+ if vendor.DeletedAt.Valid {
+ nameKey = retiredMetadataNameKey("vendor", vendor.Id)
+ }
+ if vendor.NameKey == nameKey {
+ continue
+ }
+ if err := tx.Unscoped().Model(&Vendor{}).Where("id = ?", vendor.Id).UpdateColumn("name_key", nameKey).Error; err != nil {
+ return fmt.Errorf("backfill vendor %d name key: %w", vendor.Id, err)
+ }
+ }
+
+ var models []Model
+ if err := tx.Unscoped().Find(&models).Error; err != nil {
+ return fmt.Errorf("load models for name-key migration: %w", err)
+ }
+ for _, model := range models {
+ nameKey := activeMetadataNameKey(model.ModelName)
+ if model.DeletedAt.Valid {
+ nameKey = retiredMetadataNameKey("model", model.Id)
+ }
+ if model.NameKey == nameKey {
+ continue
+ }
+ if err := tx.Unscoped().Model(&Model{}).Where("id = ?", model.Id).UpdateColumn("name_key", nameKey).Error; err != nil {
+ return fmt.Errorf("backfill model %d name key: %w", model.Id, err)
+ }
+ }
+ return nil
+}
+
+func ensureMetadataNameKeyIndexes(db *gorm.DB) error {
+ if !db.Migrator().HasIndex(&Vendor{}, vendorNameKeyIndex) {
+ if err := createUniqueIndex(db, "vendors", vendorNameKeyIndex, "name_key"); err != nil {
+ return fmt.Errorf("create vendor active-name unique index: %w", err)
+ }
+ }
+ if !db.Migrator().HasIndex(&Model{}, modelNameKeyIndex) {
+ if err := createUniqueIndex(db, "models", modelNameKeyIndex, "name_key"); err != nil {
+ return fmt.Errorf("create model active-name unique index: %w", err)
+ }
+ }
+ return nil
+}
+
+func dropLegacyMetadataNameIndexes(db *gorm.DB) error {
+ if db.Migrator().HasIndex(&Vendor{}, legacyVendorNameKeyIndex) {
+ if err := db.Migrator().DropIndex(&Vendor{}, legacyVendorNameKeyIndex); err != nil {
+ return fmt.Errorf("drop legacy vendor name unique index: %w", err)
+ }
+ }
+ if db.Migrator().HasIndex(&Model{}, legacyModelNameKeyIndex) {
+ if err := db.Migrator().DropIndex(&Model{}, legacyModelNameKeyIndex); err != nil {
+ return fmt.Errorf("drop legacy model name unique index: %w", err)
+ }
+ }
+ return nil
+}
+
+func createUniqueIndex(db *gorm.DB, table, index, column string) error {
+ return db.Exec(fmt.Sprintf(
+ "CREATE UNIQUE INDEX %s ON %s (%s)",
+ quoteDBIdentifier(index),
+ quoteDBIdentifier(table),
+ quoteDBIdentifier(column),
+ )).Error
+}
diff --git a/model/metadata_name_key_test.go b/model/metadata_name_key_test.go
new file mode 100644
index 00000000..85bdab66
--- /dev/null
+++ b/model/metadata_name_key_test.go
@@ -0,0 +1,125 @@
+package model
+
+import (
+ "errors"
+ "testing"
+
+ "github.com/MAX-API-Next/MAX-API/common"
+ "github.com/glebarez/sqlite"
+ "github.com/stretchr/testify/require"
+ "gorm.io/gorm"
+)
+
+func setupMetadataNameKeyTestDB(t *testing.T) *gorm.DB {
+ t.Helper()
+
+ previousDB := DB
+ previousSQLite := common.UsingSQLite
+ previousMySQL := common.UsingMySQL
+ previousPostgreSQL := common.UsingPostgreSQL
+ t.Cleanup(func() {
+ DB = previousDB
+ common.UsingSQLite = previousSQLite
+ common.UsingMySQL = previousMySQL
+ common.UsingPostgreSQL = previousPostgreSQL
+ })
+
+ db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
+ require.NoError(t, err)
+ DB = db
+ common.UsingSQLite = true
+ common.UsingMySQL = false
+ common.UsingPostgreSQL = false
+ initCol()
+ require.NoError(t, db.AutoMigrate(&Vendor{}, &Model{}))
+ return db
+}
+
+func TestMetadataNameKeyMigrationEnforcesActiveUniquenessAndAllowsReuseAfterDelete(t *testing.T) {
+ db := setupMetadataNameKeyTestDB(t)
+
+ vendor := &Vendor{Name: "OpenAI", Status: 1}
+ require.NoError(t, vendor.Insert())
+ model := &Model{ModelName: "gpt-test", VendorID: vendor.Id, Status: 1, SyncOfficial: 1}
+ require.NoError(t, model.Insert())
+ require.NoError(t, db.Model(&Vendor{}).Where("id = ?", vendor.Id).UpdateColumn("name_key", "").Error)
+ require.NoError(t, db.Model(&Model{}).Where("id = ?", model.Id).UpdateColumn("name_key", "").Error)
+ require.NoError(t, db.Exec("CREATE UNIQUE INDEX "+legacyVendorNameKeyIndex+" ON vendors (name, deleted_at)").Error)
+ require.NoError(t, db.Exec("CREATE UNIQUE INDEX "+legacyModelNameKeyIndex+" ON models (model_name, deleted_at)").Error)
+
+ require.NoError(t, migrateMetadataNameKeys())
+ require.True(t, db.Migrator().HasIndex(&Vendor{}, vendorNameKeyIndex))
+ require.True(t, db.Migrator().HasIndex(&Model{}, modelNameKeyIndex))
+ require.False(t, db.Migrator().HasIndex(&Vendor{}, legacyVendorNameKeyIndex))
+ require.False(t, db.Migrator().HasIndex(&Model{}, legacyModelNameKeyIndex))
+
+ duplicateVendor := &Vendor{Name: "OpenAI", Status: 1}
+ require.Error(t, duplicateVendor.Insert())
+ duplicateModel := &Model{ModelName: "gpt-test", Status: 1, SyncOfficial: 1}
+ require.Error(t, duplicateModel.Insert())
+
+ require.NoError(t, model.Delete())
+ secondModel := &Model{ModelName: "gpt-test", Status: 1, SyncOfficial: 1}
+ require.NoError(t, secondModel.Insert())
+
+ require.NoError(t, secondModel.Delete())
+ require.NoError(t, vendor.Delete())
+ secondVendor := &Vendor{Name: "OpenAI", Status: 1}
+ require.NoError(t, secondVendor.Insert())
+}
+
+func TestMetadataNameKeyMigrationRetiresDuplicateActiveKeysDeterministically(t *testing.T) {
+ db := setupMetadataNameKeyTestDB(t)
+
+ firstVendor := &Vendor{Name: "Duplicate Vendor", Status: 1}
+ secondVendor := &Vendor{Name: "Duplicate Vendor", Status: 1}
+ require.NoError(t, db.Create(firstVendor).Error)
+ require.NoError(t, db.Create(secondVendor).Error)
+ firstModel := &Model{ModelName: "duplicate-model", Status: 1, SyncOfficial: 1}
+ secondModel := &Model{ModelName: "duplicate-model", Status: 1, SyncOfficial: 1}
+ require.NoError(t, db.Create(firstModel).Error)
+ require.NoError(t, db.Create(secondModel).Error)
+
+ require.NoError(t, migrateMetadataNameKeys())
+
+ var vendors []Vendor
+ require.NoError(t, db.Order("id ASC").Find(&vendors).Error)
+ require.Len(t, vendors, 2)
+ require.Equal(t, "Duplicate Vendor", vendors[0].NameKey)
+ require.Equal(t, retiredMetadataNameKey("vendor", vendors[1].Id), vendors[1].NameKey)
+
+ var models []Model
+ require.NoError(t, db.Order("id ASC").Find(&models).Error)
+ require.Len(t, models, 2)
+ require.Equal(t, "duplicate-model", models[0].NameKey)
+ require.Equal(t, retiredMetadataNameKey("model", models[1].Id), models[1].NameKey)
+}
+
+func TestModelDeleteRollsBackNameKeyChangeWhenDeleteFails(t *testing.T) {
+ db := setupMetadataNameKeyTestDB(t)
+ model := &Model{ModelName: "rollback-model", Status: 1, SyncOfficial: 1}
+ require.NoError(t, model.Insert())
+ require.NoError(t, migrateMetadataNameKeys())
+ require.NoError(t, db.Exec(`CREATE TRIGGER fail_model_soft_delete BEFORE UPDATE OF deleted_at ON models WHEN NEW.deleted_at IS NOT NULL BEGIN SELECT RAISE(ABORT, 'delete blocked'); END`).Error)
+
+ require.Error(t, model.Delete())
+ var stored Model
+ require.NoError(t, db.Unscoped().First(&stored, model.Id).Error)
+ require.Equal(t, "rollback-model", stored.NameKey)
+ require.False(t, stored.DeletedAt.Valid)
+}
+
+func TestVendorDeleteRejectsReferencedModels(t *testing.T) {
+ setupMetadataNameKeyTestDB(t)
+ vendor := &Vendor{Name: "Referenced Vendor", Status: 1}
+ require.NoError(t, vendor.Insert())
+ model := &Model{ModelName: "referenced-model", VendorID: vendor.Id, Status: 1, SyncOfficial: 1}
+ require.NoError(t, model.Insert())
+
+ err := vendor.Delete()
+ require.True(t, errors.Is(err, ErrVendorHasModels))
+
+ var count int64
+ require.NoError(t, DB.Model(&Vendor{}).Where("id = ?", vendor.Id).Count(&count).Error)
+ require.EqualValues(t, 1, count)
+}
diff --git a/model/model_meta.go b/model/model_meta.go
index 45f21a88..c1c457c2 100644
--- a/model/model_meta.go
+++ b/model/model_meta.go
@@ -23,7 +23,8 @@ type BoundChannel struct {
type Model struct {
Id int `json:"id"`
- ModelName string `json:"model_name" gorm:"size:128;not null;uniqueIndex:uk_model_name_delete_at,priority:1"`
+ ModelName string `json:"model_name" gorm:"size:128;not null"`
+ NameKey string `json:"-" gorm:"size:191;not null;default:''"`
Description string `json:"description,omitempty" gorm:"type:text"`
Icon string `json:"icon,omitempty" gorm:"type:varchar(128)"`
Tags string `json:"tags,omitempty" gorm:"type:varchar(255)"`
@@ -33,7 +34,7 @@ type Model struct {
SyncOfficial int `json:"sync_official" gorm:"default:1"`
CreatedTime int64 `json:"created_time" gorm:"bigint"`
UpdatedTime int64 `json:"updated_time" gorm:"bigint"`
- DeletedAt gorm.DeletedAt `json:"-" gorm:"index;uniqueIndex:uk_model_name_delete_at,priority:2"`
+ DeletedAt gorm.DeletedAt `json:"-" gorm:"index"`
BoundChannels []BoundChannel `json:"bound_channels,omitempty" gorm:"-"`
EnableGroups []string `json:"enable_groups,omitempty" gorm:"-"`
@@ -48,6 +49,7 @@ func (mi *Model) Insert() error {
now := common.GetTimestamp()
mi.CreatedTime = now
mi.UpdatedTime = now
+ mi.NameKey = activeMetadataNameKey(mi.ModelName)
// 保存原始值（因为 Create 后可能被 GORM 的 default 标签覆盖为 1）
originalStatus := mi.Status
@@ -76,14 +78,22 @@ func IsModelNameDuplicated(id int, name string) (bool, error) {
func (mi *Model) Update() error {
mi.UpdatedTime = common.GetTimestamp()
+ mi.NameKey = activeMetadataNameKey(mi.ModelName)
// 使用 Select 强制更新所有字段，包括零值
return DB.Model(&Model{}).Where("id = ?", mi.Id).
- Select("model_name", "description", "icon", "tags", "vendor_id", "endpoints", "status", "sync_official", "name_rule", "updated_time").
+ Select("model_name", "name_key", "description", "icon", "tags", "vendor_id", "endpoints", "status", "sync_official", "name_rule", "updated_time").
Updates(mi).Error
}
func (mi *Model) Delete() error {
- return DB.Delete(mi).Error
+ return DB.Transaction(func(tx *gorm.DB) error {
+ return softDeleteMetadataWithNameKey(tx, &Model{}, mi.Id, retiredMetadataNameKey("model", mi.Id))
+ })
+}
+
+func (mi *Model) BeforeCreate(_ *gorm.DB) error {
+ mi.NameKey = activeMetadataNameKey(mi.ModelName)
+ return nil
}
func GetVendorModelCounts() (map[int64]int64, error) {
diff --git a/model/option.go b/model/option.go
index e3fa0c18..609f2541 100644
--- a/model/option.go
+++ b/model/option.go
@@ -1,6 +1,8 @@
package model
import (
+ "errors"
+ "fmt"
"strconv"
"strings"
"time"
@@ -73,6 +75,10 @@ func InitOptionMap() {
common.OptionMap["Notice"] = ""
common.OptionMap["About"] = ""
common.OptionMap["HomePageContent"] = ""
+ common.OptionMap["Midjourney"] = ""
+ common.OptionMap["HeaderNavModules"] = ""
+ common.OptionMap["RankingsModule"] = ""
+ common.OptionMap["SidebarModulesAdmin"] = ""
common.OptionMap["Footer"] = common.Footer
common.OptionMap["SystemName"] = common.SystemName
common.OptionMap["Logo"] = common.Logo
@@ -195,6 +201,9 @@ func InitOptionMap() {
func loadOptionsFromDatabase() {
options, _ := AllOption()
for _, option := range options {
+ if !IsRegisteredOptionKey(option.Key) {
+ continue
+ }
err := updateOptionMap(option.Key, option.Value)
if err != nil {
common.SysLog("failed to update option map: " + err.Error())
@@ -202,6 +211,16 @@ func loadOptionsFromDatabase() {
}
}
+// IsRegisteredOptionKey reports whether a key is part of the initialized
+// system/configuration option registry. It intentionally excludes arbitrary
+// rows that may exist in the options table from older deployments.
+func IsRegisteredOptionKey(key string) bool {
+ common.OptionMapRWMutex.RLock()
+ defer common.OptionMapRWMutex.RUnlock()
+ _, ok := common.OptionMap[key]
+ return ok
+}
+
func SyncOptions(frequency int) {
for {
time.Sleep(time.Duration(frequency) * time.Second)
@@ -211,6 +230,9 @@ func SyncOptions(frequency int) {
}
func UpdateOption(key string, value string) error {
+ if err := requireRegisteredOptionKey(key); err != nil {
+ return err
+ }
var err error
value, err = normalizeOptionUpdateValue(key, value)
if err != nil {
@@ -244,6 +266,11 @@ func UpdateOptionsBulk(values map[string]string) error {
if len(values) == 0 {
return nil
}
+ for key := range values {
+ if err := requireRegisteredOptionKey(key); err != nil {
+ return err
+ }
+ }
normalizedValues := make(map[string]string, len(values))
for k, v := range values {
normalized, err := normalizeOptionUpdateValue(k, v)
@@ -281,16 +308,40 @@ func UpdateOptionsBulk(values map[string]string) error {
return nil
}
+func requireRegisteredOptionKey(key string) error {
+ if IsRegisteredOptionKey(key) {
+ return nil
+ }
+ return fmt.Errorf("unsupported option key: %s", key)
+}
+
func normalizeOptionUpdateValue(key string, value string) (string, error) {
switch key {
case "GroupRatio", "group_ratio_setting.group_ratio":
return ratio_setting.NormalizeGroupRatioJSONString(value)
+ case "DataExportInterval":
+ interval, err := parseDataExportInterval(value)
+ if err != nil {
+ return "", err
+ }
+ return strconv.Itoa(interval), nil
default:
return value, nil
}
}
+func parseDataExportInterval(value string) (int, error) {
+ interval, err := strconv.Atoi(strings.TrimSpace(value))
+ if err != nil || interval < 1 || interval > 1440 {
+ return 0, errors.New("DataExportInterval must be an integer between 1 and 1440 minutes")
+ }
+ return interval, nil
+}
+
func validateOptionUpdate(key string, value string) error {
+ if err := validateRegisteredConfigOption(key, value); err != nil {
+ return err
+ }
switch key {
case "Chats":
return validateJSONOption[[]map[string]string](value)
@@ -304,7 +355,7 @@ func validateOptionUpdate(key string, value string) error {
case "ModelRequestRateLimitGroup":
return setting.CheckModelRequestRateLimitGroup(value)
case "ModelRatio", "ModelPrice", "CacheRatio", "CreateCacheRatio", "CompletionRatio", "ImageRatio", "AudioRatio", "AudioCompletionRatio":
- return validateJSONOption[map[string]float64](value)
+ return ratio_setting.ValidatePricingMapJSONString(value)
case "GroupRatio", "group_ratio_setting.group_ratio":
return ratio_setting.CheckGroupRatio(value)
case "GroupGroupRatio":
@@ -323,6 +374,18 @@ func validateOptionUpdate(key string, value string) error {
}
}
+func validateRegisteredConfigOption(key string, value string) error {
+ parts := strings.SplitN(key, ".", 2)
+ if len(parts) != 2 {
+ return nil
+ }
+ cfg := config.GlobalConfig.Get(parts[0])
+ if cfg == nil {
+ return nil
+ }
+ return config.ValidateConfigPointerNulls(cfg, map[string]string{parts[1]: value})
+}
+
func validateJSONOption[T any](value string) error {
var target T
return common.UnmarshalJsonStr(value, &target)
diff --git a/model/option_test.go b/model/option_test.go
index 6df69326..8c9ce17b 100644
--- a/model/option_test.go
+++ b/model/option_test.go
@@ -26,11 +26,37 @@ func setupOptionMapTestState(t *testing.T) {
})
}
-func optionMapContainsForTest(key string) bool {
- common.OptionMapRWMutex.RLock()
- defer common.OptionMapRWMutex.RUnlock()
- _, ok := common.OptionMap[key]
- return ok
+func registerOptionKeysForTest(keys ...string) {
+ common.OptionMapRWMutex.Lock()
+ defer common.OptionMapRWMutex.Unlock()
+ for _, key := range keys {
+ common.OptionMap[key] = ""
+ }
+}
+
+func TestUpdateOptionRejectsUnregisteredKeyBeforePersistence(t *testing.T) {
+ setupOptionMapTestState(t)
+ deleteOptionsForTest(t, "unregistered.option")
+ t.Cleanup(func() { deleteOptionsForTest(t, "unregistered.option") })
+
+ err := UpdateOption("unregistered.option", "value")
+ require.Error(t, err)
+ require.False(t, optionExistsForTest(t, "unregistered.option"))
+}
+
+func TestUpdateOptionsBulkRejectsUnregisteredKeyBeforePersistence(t *testing.T) {
+ setupOptionMapTestState(t)
+ registerOptionKeysForTest("SystemName")
+ deleteOptionsForTest(t, "SystemName", "unregistered.option")
+ t.Cleanup(func() { deleteOptionsForTest(t, "SystemName", "unregistered.option") })
+
+ err := UpdateOptionsBulk(map[string]string{
+ "SystemName": "should-not-persist",
+ "unregistered.option": "value",
+ })
+ require.Error(t, err)
+ require.False(t, optionExistsForTest(t, "SystemName"))
+ require.False(t, optionExistsForTest(t, "unregistered.option"))
}
func deleteOptionsForTest(t *testing.T, keys ...string) {
@@ -60,6 +86,7 @@ func optionMapValueForTest(key string) string {
func TestUpdateOptionFiltersAutoRouteGroupRatioNamesBeforePersistence(t *testing.T) {
setupOptionMapTestState(t)
+ registerOptionKeysForTest("GroupRatio", "group_ratio_setting.group_ratio")
deleteOptionsForTest(t, "GroupRatio", "group_ratio_setting.group_ratio")
t.Cleanup(func() {
deleteOptionsForTest(t, "GroupRatio", "group_ratio_setting.group_ratio")
@@ -154,8 +181,79 @@ func TestValidateOptionUpdateRejectsRuntimeConfigParseErrors(t *testing.T) {
}
}
+func TestValidateOptionUpdateRejectsUnsafePricingValues(t *testing.T) {
+ pricingKeys := []string{
+ "ModelRatio",
+ "ModelPrice",
+ "CacheRatio",
+ "CreateCacheRatio",
+ "CompletionRatio",
+ "ImageRatio",
+ "AudioRatio",
+ "AudioCompletionRatio",
+ }
+
+ for _, key := range pricingKeys {
+ t.Run(key+" rejects null", func(t *testing.T) {
+ require.Error(t, validateOptionUpdate(key, `{"unsafe-model":null}`))
+ })
+ t.Run(key+" rejects negative values", func(t *testing.T) {
+ require.Error(t, validateOptionUpdate(key, `{"unsafe-model":-0.01}`))
+ })
+ t.Run(key+" allows zero", func(t *testing.T) {
+ require.NoError(t, validateOptionUpdate(key, `{"free-model":0}`))
+ })
+ }
+}
+
+func TestValidateOptionUpdateRejectsNullRWMapConfigs(t *testing.T) {
+ for _, key := range []string{
+ "billing_setting.billing_mode",
+ "billing_setting.billing_expr",
+ "task_billing_setting.rate_cards",
+ } {
+ t.Run(key, func(t *testing.T) {
+ require.Error(t, validateOptionUpdate(key, "null"))
+ })
+ }
+}
+
+func TestNormalizeDataExportIntervalRejectsUnsafeValues(t *testing.T) {
+ for _, value := range []string{"0", "-1", "1441", "not-a-number"} {
+ t.Run(value, func(t *testing.T) {
+ _, err := normalizeOptionUpdateValue("DataExportInterval", value)
+ require.Error(t, err)
+ })
+ }
+
+ normalized, err := normalizeOptionUpdateValue("DataExportInterval", " 60 ")
+ require.NoError(t, err)
+ require.Equal(t, "60", normalized)
+}
+
+func TestUpdateOptionsBulkRejectsNullRWMapConfigBeforePersistence(t *testing.T) {
+ setupOptionMapTestState(t)
+ registerOptionKeysForTest("SystemName", "billing_setting.billing_mode")
+ deleteOptionsForTest(t, "SystemName", "billing_setting.billing_mode")
+ t.Cleanup(func() {
+ deleteOptionsForTest(t, "SystemName", "billing_setting.billing_mode")
+ })
+
+ err := UpdateOptionsBulk(map[string]string{
+ "SystemName": "should-not-persist",
+ "billing_setting.billing_mode": "null",
+ })
+
+ require.Error(t, err)
+ require.False(t, optionExistsForTest(t, "SystemName"))
+ require.False(t, optionExistsForTest(t, "billing_setting.billing_mode"))
+ require.Equal(t, "", optionMapValueForTest("SystemName"))
+ require.Equal(t, "", optionMapValueForTest("billing_setting.billing_mode"))
+}
+
func TestUpdateOptionsBulkRejectsRuntimeConfigErrorsBeforePersistence(t *testing.T) {
setupOptionMapTestState(t)
+ registerOptionKeysForTest("SystemName", "ModelRatio")
deleteOptionsForTest(t, "SystemName", "ModelRatio")
t.Cleanup(func() {
deleteOptionsForTest(t, "SystemName", "ModelRatio")
@@ -168,6 +266,6 @@ func TestUpdateOptionsBulkRejectsRuntimeConfigErrorsBeforePersistence(t *testing
require.Error(t, err)
require.False(t, optionExistsForTest(t, "SystemName"))
require.False(t, optionExistsForTest(t, "ModelRatio"))
- require.False(t, optionMapContainsForTest("SystemName"))
- require.False(t, optionMapContainsForTest("ModelRatio"))
+ require.Equal(t, "", optionMapValueForTest("SystemName"))
+ require.Equal(t, "", optionMapValueForTest("ModelRatio"))
}
diff --git a/model/redemption.go b/model/redemption.go
index 76d518f2..279088a9 100644
--- a/model/redemption.go
+++ b/model/redemption.go
@@ -152,6 +152,7 @@ func Redeem(key string, userId int) (quota int, err error) {
keyCol = `"key"`
}
common.RandomSleep()
+ var cacheTask CacheInvalidationTask
err = DB.Transaction(func(tx *gorm.DB) error {
err := withRowLock(tx).Where(keyCol+" = ?", key).First(redemption).Error
if err != nil {
@@ -185,6 +186,10 @@ func Redeem(key string, userId int) (quota int, err error) {
if err := ensureUserUpdateMatchedTx(tx, userResult, userId, errors.New("无效的 user id")); err != nil {
return err
}
+ cacheTask, err = stageUserCacheInvalidationTx(tx, userId, false)
+ if err != nil {
+ return err
+ }
redemption.RedeemedTime = now
redemption.Status = common.RedemptionCodeStatusUsed
redemption.UsedUserId = userId
@@ -195,6 +200,7 @@ func Redeem(key string, userId int) (quota int, err error) {
common.SysError("redemption failed: " + err.Error())
return 0, ErrRedeemFailed
}
+ dispatchStagedCacheInvalidation(cacheTask)
RecordLog(userId, LogTypeTopup, fmt.Sprintf("通过兑换码充值 %s，兑换码ID %d", logger.LogQuota(redemption.Quota), redemption.Id))
return redemption.Quota, nil
}
diff --git a/model/subscription.go b/model/subscription.go
index d0813ade..77dfd88f 100644
--- a/model/subscription.go
+++ b/model/subscription.go
@@ -36,6 +36,7 @@ const (
var (
ErrSubscriptionOrderNotFound = errors.New("subscription order not found")
ErrSubscriptionOrderStatusInvalid = errors.New("subscription order status invalid")
+ ErrSubscriptionQuotaInsufficient = errors.New("subscription quota insufficient")
)
func invalidateSubscriptionUserCache(userId int, operation string) {
@@ -508,13 +509,15 @@ func downgradeUserGroupForSubscriptionTx(tx *gorm.DB, sub *UserSubscription, now
if currentGroup != upgradeGroup {
return "", nil
}
- var activeSub UserSubscription
activeQuery := tx.Where("user_id = ? AND status = ? AND end_time > ? AND id <> ? AND upgrade_group <> ''",
sub.UserId, "active", now, sub.Id).
- Order("end_time desc, id desc").
+ Where("TRIM(upgrade_group) = ?", currentGroup).
Limit(1).
- Find(&activeSub)
- if activeQuery.Error == nil && activeQuery.RowsAffected > 0 {
+ Find(&UserSubscription{})
+ if activeQuery.Error != nil {
+ return "", activeQuery.Error
+ }
+ if activeQuery.RowsAffected > 0 {
return "", nil
}
prevGroup := strings.TrimSpace(sub.PrevUserGroup)
@@ -617,7 +620,10 @@ func CompleteSubscriptionOrder(tradeNo string, providerPayload string, expectedP
err := DB.Transaction(func(tx *gorm.DB) error {
var order SubscriptionOrder
if err := withRowLock(tx).Where(refCol+" = ?", tradeNo).First(&order).Error; err != nil {
- return ErrSubscriptionOrderNotFound
+ if errors.Is(err, gorm.ErrRecordNotFound) {
+ return ErrSubscriptionOrderNotFound
+ }
+ return err
}
if expectedPaymentProvider != "" && order.PaymentProvider != expectedPaymentProvider {
return ErrPaymentMethodMismatch
@@ -1164,18 +1170,8 @@ func ExpireDueSubscriptions(limit int) (int, error) {
}
expiredCount += int(res.RowsAffected)
- // If there's an active upgraded subscription, keep current group.
- var activeSub UserSubscription
- activeQuery := tx.Where("user_id = ? AND status = ? AND end_time > ? AND upgrade_group <> ''",
- userId, "active", now).
- Order("end_time desc, id desc").
- Limit(1).
- Find(&activeSub)
- if activeQuery.Error == nil && activeQuery.RowsAffected > 0 {
- return nil
- }
-
- // No active upgraded subscription, downgrade to previous group if needed.
+ // Select the most recent expired upgrade and restore its previous group
+ // unless an active subscription still grants the current group.
var lastExpired UserSubscription
expiredQuery := tx.Where("user_id = ? AND status = ? AND upgrade_group <> ''",
userId, "expired").
@@ -1197,6 +1193,20 @@ func ExpireDueSubscriptions(limit int) (int, error) {
if currentGroup != upgradeGroup || currentGroup == prevGroup {
return nil
}
+ // Keep the current group only when another active subscription grants
+ // that exact group. Any other active upgrade must not retain a higher
+ // group from the subscription that just expired.
+ activeQuery := tx.Where("user_id = ? AND status = ? AND end_time > ? AND upgrade_group <> ''",
+ userId, "active", now).
+ Where("TRIM(upgrade_group) = ?", currentGroup).
+ Limit(1).
+ Find(&UserSubscription{})
+ if activeQuery.Error != nil {
+ return activeQuery.Error
+ }
+ if activeQuery.RowsAffected > 0 {
+ return nil
+ }
if err := tx.Model(&User{}).Where("id = ?", userId).
Update("group", prevGroup).Error; err != nil {
return err
@@ -1216,14 +1226,18 @@ func ExpireDueSubscriptions(limit int) (int, error) {
// SubscriptionPreConsumeRecord stores idempotent pre-consume operations per request.
type SubscriptionPreConsumeRecord struct {
- Id int `json:"id"`
- RequestId string `json:"request_id" gorm:"type:varchar(64);uniqueIndex"`
- UserId int `json:"user_id" gorm:"index"`
- UserSubscriptionId int `json:"user_subscription_id" gorm:"index"`
- PreConsumed int64 `json:"pre_consumed" gorm:"type:bigint;not null;default:0"`
- Status string `json:"status" gorm:"type:varchar(32);index"` // consumed/refunded
- CreatedAt int64 `json:"created_at" gorm:"bigint"`
- UpdatedAt int64 `json:"updated_at" gorm:"bigint;index"`
+ Id int `json:"id"`
+ RequestId string `json:"request_id" gorm:"type:varchar(64);uniqueIndex"`
+ UserId int `json:"user_id" gorm:"index"`
+ TokenId int `json:"token_id" gorm:"index"`
+ UserSubscriptionId int `json:"user_subscription_id" gorm:"index"`
+ SubscriptionLastResetTime int64 `json:"subscription_last_reset_time" gorm:"bigint;not null;default:0"`
+ ModelName string `json:"model_name" gorm:"type:varchar(191)"`
+ QuotaType int `json:"quota_type"`
+ PreConsumed int64 `json:"pre_consumed" gorm:"type:bigint;not null;default:0"`
+ Status string `json:"status" gorm:"type:varchar(32);index"` // consumed/refunded
+ CreatedAt int64 `json:"created_at" gorm:"bigint"`
+ UpdatedAt int64 `json:"updated_at" gorm:"bigint;index"`
}
func (r *SubscriptionPreConsumeRecord) BeforeCreate(tx *gorm.DB) error {
@@ -1287,115 +1301,237 @@ func PreConsumeUserSubscription(requestId string, userId int, modelName string,
}
now := GetDBTimestamp()
- returnValue := &SubscriptionPreConsumeResult{}
+ var returnValue SubscriptionPreConsumeResult
+ err := DB.Transaction(func(tx *gorm.DB) error {
+ if err := claimBillingPreConsumeSourceTx(tx, BillingPreConsumeSelection{
+ RequestID: requestId, Source: BillingSettlementSourceSubscription,
+ UserID: userId, ModelName: modelName,
+ RequestedQuota: amount, EffectiveQuota: amount,
+ }); err != nil {
+ return err
+ }
+ var err error
+ returnValue, _, err = preConsumeUserSubscriptionTx(tx, requestId, userId, 0, modelName, quotaType, now, amount)
+ return err
+ })
+ if err != nil {
+ return nil, err
+ }
+ return &returnValue, nil
+}
+// PreConsumeTokenAndUserSubscription atomically reserves subscription and
+// token quota. Existing request IDs are returned without applying the token
+// mutation again.
+func PreConsumeTokenAndUserSubscription(requestId string, userId int, tokenId int, tokenKey string, modelName string, quotaType int, amount int64) (*SubscriptionPreConsumeResult, error) {
+ if tokenId <= 0 {
+ return nil, errors.New("invalid tokenId")
+ }
+ if userId <= 0 {
+ return nil, errors.New("invalid userId")
+ }
+ if strings.TrimSpace(requestId) == "" {
+ return nil, errors.New("requestId is empty")
+ }
+ if amount <= 0 {
+ return nil, errors.New("amount must be > 0")
+ }
+ now := GetDBTimestamp()
+ var returnValue SubscriptionPreConsumeResult
+ var cacheTask CacheInvalidationTask
+ resolvedTokenKey := tokenKey
+ alreadyConsumed := false
err := DB.Transaction(func(tx *gorm.DB) error {
- var existing SubscriptionPreConsumeRecord
- query := tx.Where("request_id = ?", requestId).Limit(1).Find(&existing)
- if query.Error != nil {
- return query.Error
- }
- if query.RowsAffected > 0 {
- if existing.Status == "refunded" {
- return errors.New("subscription pre-consume already refunded")
- }
- var sub UserSubscription
- if err := tx.Where("id = ?", existing.UserSubscriptionId).First(&sub).Error; err != nil {
- return err
+ if err := claimBillingPreConsumeSourceTx(tx, BillingPreConsumeSelection{
+ RequestID: requestId, Source: BillingSettlementSourceSubscription,
+ UserID: userId, TokenID: tokenId, ModelName: modelName,
+ RequestedQuota: amount, EffectiveQuota: amount,
+ }); err != nil {
+ return err
+ }
+ var err error
+ returnValue, alreadyConsumed, err = preConsumeUserSubscriptionTx(tx, requestId, userId, tokenId, modelName, quotaType, now, amount)
+ if err != nil || alreadyConsumed {
+ if alreadyConsumed {
+ var token Token
+ if tokenErr := tx.Unscoped().Select("key").First(&token, tokenId).Error; tokenErr == nil {
+ resolvedTokenKey = token.Key
+ }
}
- returnValue.UserSubscriptionId = sub.Id
- returnValue.PreConsumed = existing.PreConsumed
- returnValue.AmountTotal = sub.AmountTotal
- returnValue.AmountUsedBefore = sub.AmountUsed
- returnValue.AmountUsedAfter = sub.AmountUsed
- return nil
+ return err
}
+ result := tx.Model(&Token{}).
+ Where("id = ? AND user_id = ? AND (unlimited_quota = ? OR remain_quota >= ?)", tokenId, userId, true, amount).
+ Updates(map[string]interface{}{
+ "remain_quota": gorm.Expr("remain_quota - ?", amount),
+ "used_quota": gorm.Expr("used_quota + ?", amount),
+ "accessed_time": common.GetTimestamp(),
+ })
+ if result.Error != nil {
+ return result.Error
+ }
+ if result.RowsAffected != 1 {
+ return fmt.Errorf("%w: id=%d, need=%d", ErrTokenQuotaInsufficient, tokenId, amount)
+ }
+ var token Token
+ if err := tx.First(&token, tokenId).Error; err != nil {
+ return err
+ }
+ resolvedTokenKey = token.Key
+ cacheTask, err = stageTokenCacheInvalidationTx(tx, resolvedTokenKey, false)
+ return err
+ })
+ if err != nil {
+ return nil, err
+ }
+ if alreadyConsumed {
+ invalidateTokenQuotaCache(resolvedTokenKey)
+ } else {
+ dispatchStagedCacheInvalidation(cacheTask)
+ }
+ return &returnValue, nil
+}
- var subs []UserSubscription
- if err := withRowLock(tx).
- Where("user_id = ? AND status = ? AND end_time > ?", userId, "active", now).
- Order("end_time asc, id asc").
- Find(&subs).Error; err != nil {
- return errors.New("no active subscription")
+func preConsumeUserSubscriptionTx(tx *gorm.DB, requestId string, userId int, tokenId int, modelName string, quotaType int, now int64, amount int64) (SubscriptionPreConsumeResult, bool, error) {
+ returnValue := SubscriptionPreConsumeResult{}
+ var existing SubscriptionPreConsumeRecord
+ query := tx.Where("request_id = ?", requestId).Limit(1).Find(&existing)
+ if query.Error != nil {
+ return returnValue, false, query.Error
+ }
+ if query.RowsAffected > 0 {
+ if err := validateSubscriptionPreConsumeReplay(existing, userId, tokenId, modelName, quotaType, amount); err != nil {
+ return returnValue, true, err
}
- if len(subs) == 0 {
- return errors.New("no active subscription")
+ if existing.Status == "refunded" {
+ return returnValue, true, errors.New("subscription pre-consume already refunded")
}
- for _, candidate := range subs {
- sub := candidate
- plan, err := getSubscriptionPlanByIdTx(tx, sub.PlanId)
- if err != nil {
- return err
- }
- if err := maybeResetUserSubscriptionWithPlanTx(tx, &sub, plan, now); err != nil {
- return err
- }
- usedBefore := sub.AmountUsed
- if sub.AmountTotal > 0 {
- remain := sub.AmountTotal - usedBefore
- if remain < amount {
- continue
+ var sub UserSubscription
+ if err := tx.Where("id = ?", existing.UserSubscriptionId).First(&sub).Error; err != nil {
+ return returnValue, true, err
+ }
+ if sub.UserId != userId {
+ return returnValue, true, errors.New("subscription pre-consume user identity mismatch")
+ }
+ if err := validateSubscriptionPreConsumePeriod(existing, sub); err != nil {
+ return returnValue, true, err
+ }
+ returnValue.UserSubscriptionId = sub.Id
+ returnValue.PreConsumed = existing.PreConsumed
+ returnValue.AmountTotal = sub.AmountTotal
+ returnValue.AmountUsedBefore = sub.AmountUsed
+ returnValue.AmountUsedAfter = sub.AmountUsed
+ return returnValue, true, nil
+ }
+
+ var subs []UserSubscription
+ if err := withRowLock(tx).
+ Where("user_id = ? AND status = ? AND end_time > ?", userId, "active", now).
+ Order("end_time asc, id asc").
+ Find(&subs).Error; err != nil {
+ return returnValue, false, err
+ }
+ if len(subs) == 0 {
+ return returnValue, false, errors.New("no active subscription")
+ }
+ for _, candidate := range subs {
+ sub := candidate
+ plan, err := getSubscriptionPlanByIdTx(tx, sub.PlanId)
+ if err != nil {
+ return returnValue, false, err
+ }
+ if err := maybeResetUserSubscriptionWithPlanTx(tx, &sub, plan, now); err != nil {
+ return returnValue, false, err
+ }
+ usedBefore := sub.AmountUsed
+ if sub.AmountTotal > 0 && sub.AmountTotal-usedBefore < amount {
+ continue
+ }
+ record := &SubscriptionPreConsumeRecord{
+ RequestId: requestId,
+ UserId: userId,
+ TokenId: tokenId,
+ UserSubscriptionId: sub.Id,
+ SubscriptionLastResetTime: sub.LastResetTime,
+ ModelName: modelName,
+ QuotaType: quotaType,
+ PreConsumed: amount,
+ Status: "consumed",
+ }
+ if err := tx.Create(record).Error; err != nil {
+ var duplicate SubscriptionPreConsumeRecord
+ if duplicateErr := tx.Where("request_id = ?", requestId).First(&duplicate).Error; duplicateErr == nil {
+ if replayErr := validateSubscriptionPreConsumeReplay(duplicate, userId, tokenId, modelName, quotaType, amount); replayErr != nil {
+ return returnValue, true, replayErr
}
- }
- record := &SubscriptionPreConsumeRecord{
- RequestId: requestId,
- UserId: userId,
- UserSubscriptionId: sub.Id,
- PreConsumed: amount,
- Status: "consumed",
- }
- if err := tx.Create(record).Error; err != nil {
- var dup SubscriptionPreConsumeRecord
- if err2 := tx.Where("request_id = ?", requestId).First(&dup).Error; err2 == nil {
- if dup.Status == "refunded" {
- return errors.New("subscription pre-consume already refunded")
- }
- var dupSub UserSubscription
- if err3 := tx.Where("id = ?", dup.UserSubscriptionId).First(&dupSub).Error; err3 != nil {
- return err3
- }
- returnValue.UserSubscriptionId = dupSub.Id
- returnValue.PreConsumed = dup.PreConsumed
- returnValue.AmountTotal = dupSub.AmountTotal
- returnValue.AmountUsedBefore = dupSub.AmountUsed
- returnValue.AmountUsedAfter = dupSub.AmountUsed
- return nil
+ if duplicate.Status == "refunded" {
+ return returnValue, true, errors.New("subscription pre-consume already refunded")
}
- return err
- }
- updateQuery := tx.Model(&UserSubscription{}).
- Where("id = ? AND status = ? AND end_time > ?", sub.Id, "active", now).
- Where("amount_total <= 0 OR amount_used + ? <= amount_total", amount)
- result := updateQuery.Updates(map[string]interface{}{
- "amount_used": gorm.Expr("amount_used + ?", amount),
- "updated_at": common.GetTimestamp(),
- })
- if result.Error != nil {
- return result.Error
- }
- if result.RowsAffected != 1 {
- if err := tx.Delete(record).Error; err != nil {
- return err
+ var duplicateSub UserSubscription
+ if err := tx.Where("id = ?", duplicate.UserSubscriptionId).First(&duplicateSub).Error; err != nil {
+ return returnValue, true, err
+ }
+ if duplicateSub.UserId != userId {
+ return returnValue, true, errors.New("subscription pre-consume user identity mismatch")
}
- continue
+ if err := validateSubscriptionPreConsumePeriod(duplicate, duplicateSub); err != nil {
+ return returnValue, true, err
+ }
+ returnValue.UserSubscriptionId = duplicateSub.Id
+ returnValue.PreConsumed = duplicate.PreConsumed
+ returnValue.AmountTotal = duplicateSub.AmountTotal
+ returnValue.AmountUsedBefore = duplicateSub.AmountUsed
+ returnValue.AmountUsedAfter = duplicateSub.AmountUsed
+ return returnValue, true, nil
}
- var updatedSub UserSubscription
- if err := tx.Where("id = ?", sub.Id).First(&updatedSub).Error; err != nil {
- return err
+ return returnValue, false, err
+ }
+ updateQuery := tx.Model(&UserSubscription{}).
+ Where("id = ? AND status = ? AND end_time > ?", sub.Id, "active", now).
+ Where("amount_total <= 0 OR amount_used + ? <= amount_total", amount)
+ result := updateQuery.Updates(map[string]interface{}{
+ "amount_used": gorm.Expr("amount_used + ?", amount),
+ "updated_at": common.GetTimestamp(),
+ })
+ if result.Error != nil {
+ return returnValue, false, result.Error
+ }
+ if result.RowsAffected != 1 {
+ if err := tx.Delete(record).Error; err != nil {
+ return returnValue, false, err
}
- returnValue.UserSubscriptionId = sub.Id
- returnValue.PreConsumed = amount
- returnValue.AmountTotal = updatedSub.AmountTotal
- returnValue.AmountUsedBefore = usedBefore
- returnValue.AmountUsedAfter = updatedSub.AmountUsed
- return nil
+ continue
}
- return fmt.Errorf("subscription quota insufficient, need=%d", amount)
- })
- if err != nil {
- return nil, err
+ var updatedSub UserSubscription
+ if err := tx.Where("id = ?", sub.Id).First(&updatedSub).Error; err != nil {
+ return returnValue, false, err
+ }
+ returnValue.UserSubscriptionId = sub.Id
+ returnValue.PreConsumed = amount
+ returnValue.AmountTotal = updatedSub.AmountTotal
+ returnValue.AmountUsedBefore = usedBefore
+ returnValue.AmountUsedAfter = updatedSub.AmountUsed
+ return returnValue, false, nil
}
- return returnValue, nil
+ return returnValue, false, fmt.Errorf("%w: need=%d", ErrSubscriptionQuotaInsufficient, amount)
+}
+
+func validateSubscriptionPreConsumeReplay(existing SubscriptionPreConsumeRecord, userId int, tokenId int, modelName string, quotaType int, amount int64) error {
+ if existing.UserId != userId ||
+ existing.TokenId != tokenId ||
+ existing.ModelName != modelName ||
+ existing.QuotaType != quotaType ||
+ existing.PreConsumed != amount {
+ return fmt.Errorf("subscription pre-consume request reused with different identity or parameters: request=%s", existing.RequestId)
+ }
+ return nil
+}
+
+func validateSubscriptionPreConsumePeriod(record SubscriptionPreConsumeRecord, sub UserSubscription) error {
+ if record.SubscriptionLastResetTime != sub.LastResetTime {
+ return fmt.Errorf("%w: request=%s", ErrSubscriptionSettlementPeriodChanged, record.RequestId)
+ }
+ return nil
}
// RefundSubscriptionPreConsume is idempotent and refunds pre-consumed subscription quota by requestId.
@@ -1412,6 +1548,13 @@ func RefundSubscriptionPreConsume(requestId string) error {
if record.Status == "refunded" {
return nil
}
+ var sub UserSubscription
+ if err := withRowLock(tx).Where("id = ?", record.UserSubscriptionId).First(&sub).Error; err != nil {
+ return err
+ }
+ if err := validateSubscriptionPreConsumePeriod(record, sub); err != nil {
+ return err
+ }
result := tx.Model(&SubscriptionPreConsumeRecord{}).
Where("id = ? AND status <> ?", record.Id, "refunded").
Updates(map[string]interface{}{
@@ -1427,9 +1570,13 @@ func RefundSubscriptionPreConsume(requestId string) error {
if record.PreConsumed <= 0 {
return nil
}
- if _, err := postConsumeUserSubscriptionDeltaTx(tx, record.UserSubscriptionId, -record.PreConsumed); err != nil {
+ applied, err := postConsumeUserSubscriptionDeltaTx(tx, record.UserSubscriptionId, -record.PreConsumed)
+ if err != nil {
return err
}
+ if applied != -record.PreConsumed {
+ return fmt.Errorf("subscription refund was clamped: requested=%d applied=%d", -record.PreConsumed, applied)
+ }
return nil
})
}
@@ -1477,14 +1624,33 @@ func ResetDueSubscriptions(limit int) (int, error) {
return resetCount, nil
}
-// CleanupSubscriptionPreConsumeRecords removes old idempotency records to keep table small.
+// CleanupSubscriptionPreConsumeRecords removes only old replay tombstones whose
+// linked billing settlements have reached terminal applied/manual states.
+// olderThanSeconds is retained for the scheduled cleanup API; non-positive
+// values disable deletion, and pending settlement rows are always preserved.
func CleanupSubscriptionPreConsumeRecords(olderThanSeconds int64) (int64, error) {
if olderThanSeconds <= 0 {
- olderThanSeconds = 7 * 24 * 3600
+ return 0, nil
+ }
+ cutoff := common.GetTimestamp() - olderThanSeconds
+ if cutoff <= 0 {
+ return 0, nil
}
- cutoff := GetDBTimestamp() - olderThanSeconds
- res := DB.Where("updated_at < ?", cutoff).Delete(&SubscriptionPreConsumeRecord{})
- return res.RowsAffected, res.Error
+ terminalStatuses := []string{BillingSettlementStatusApplied, BillingSettlementStatusManual}
+ nonterminalSettlements := DB.Model(&BillingSettlement{}).
+ Select("1").
+ Where("billing_settlements.subscription_pre_consume_request_id = subscription_pre_consume_records.request_id").
+ Where("billing_settlements.status NOT IN ?", terminalStatuses)
+ terminalSettlements := DB.Model(&BillingSettlement{}).
+ Select("1").
+ Where("billing_settlements.subscription_pre_consume_request_id = subscription_pre_consume_records.request_id").
+ Where("billing_settlements.status IN ?", terminalStatuses)
+ result := DB.
+ Where("created_at < ?", cutoff).
+ Where("EXISTS (?)", terminalSettlements).
+ Where("NOT EXISTS (?)", nonterminalSettlements).
+ Delete(&SubscriptionPreConsumeRecord{})
+ return result.RowsAffected, result.Error
}
type SubscriptionPlanInfo struct {
@@ -1547,7 +1713,7 @@ func postConsumeUserSubscriptionDeltaTx(tx *gorm.DB, userSubscriptionId int, del
if delta > 0 {
newUsed := sub.AmountUsed + delta
if sub.AmountTotal > 0 && newUsed > sub.AmountTotal {
- return 0, fmt.Errorf("subscription used exceeds total, used=%d total=%d", newUsed, sub.AmountTotal)
+ return 0, fmt.Errorf("%w: used=%d total=%d", ErrSubscriptionQuotaInsufficient, newUsed, sub.AmountTotal)
}
updateQuery := tx.Model(&UserSubscription{}).Where("id = ?", userSubscriptionId)
if sub.AmountTotal > 0 {
@@ -1561,7 +1727,7 @@ func postConsumeUserSubscriptionDeltaTx(tx *gorm.DB, userSubscriptionId int, del
return 0, result.Error
}
if result.RowsAffected != 1 {
- return 0, fmt.Errorf("subscription used exceeds total, used=%d total=%d", newUsed, sub.AmountTotal)
+ return 0, fmt.Errorf("%w: used=%d total=%d", ErrSubscriptionQuotaInsufficient, newUsed, sub.AmountTotal)
}
return delta, nil
}
diff --git a/model/subscription_reset_test.go b/model/subscription_reset_test.go
index 338f18a7..0a690f1b 100644
--- a/model/subscription_reset_test.go
+++ b/model/subscription_reset_test.go
@@ -5,8 +5,10 @@ import (
"testing"
"time"
+ "github.com/MAX-API-Next/MAX-API/common"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
+ "gorm.io/gorm"
)
func seedSubscriptionResetPlan(t *testing.T, plan *SubscriptionPlan) {
@@ -199,3 +201,118 @@ func TestAdminResetPlanSubscriptionsNoMatchSucceeds(t *testing.T) {
assert.Zero(t, result.UserCount)
assert.Empty(t, result.AffectedUserIds)
}
+
+func TestDowngradeUserGroupDoesNotRetainHigherGroupForDifferentActiveSubscription(t *testing.T) {
+ truncateTables(t)
+ now := GetDBTimestamp()
+ user := User{Id: 9701, Username: "subscription-downgrade-user", Group: "svip", Status: common.UserStatusEnabled}
+ replaced := UserSubscription{
+ Id: 9702, UserId: user.Id, Status: "cancelled", EndTime: now + 3600,
+ UpgradeGroup: "svip", PrevUserGroup: "vip",
+ }
+ remaining := UserSubscription{
+ Id: 9703, UserId: user.Id, Status: "active", EndTime: now + 7200,
+ UpgradeGroup: "vip", PrevUserGroup: "default",
+ }
+ require.NoError(t, DB.Create(&user).Error)
+ require.NoError(t, DB.Create(&replaced).Error)
+ require.NoError(t, DB.Create(&remaining).Error)
+
+ var group string
+ require.NoError(t, DB.Transaction(func(tx *gorm.DB) error {
+ var err error
+ group, err = downgradeUserGroupForSubscriptionTx(tx, &replaced, now)
+ return err
+ }))
+ assert.Equal(t, "vip", group)
+ var stored User
+ require.NoError(t, DB.First(&stored, user.Id).Error)
+ assert.Equal(t, "vip", stored.Group)
+}
+
+func TestExpireDueSubscriptionsDowngradesToRemainingActiveUpgradeGroup(t *testing.T) {
+ truncateTables(t)
+ now := GetDBTimestamp()
+ user := User{Id: 9711, Username: "subscription-expiry-user", Group: "svip", Status: common.UserStatusEnabled}
+ expired := UserSubscription{
+ Id: 9712, UserId: user.Id, Status: "active", EndTime: now - 1,
+ UpgradeGroup: "svip", PrevUserGroup: "vip",
+ }
+ remaining := UserSubscription{
+ Id: 9713, UserId: user.Id, Status: "active", EndTime: now + 7200,
+ UpgradeGroup: "vip", PrevUserGroup: "default",
+ }
+ require.NoError(t, DB.Create(&user).Error)
+ require.NoError(t, DB.Create(&expired).Error)
+ require.NoError(t, DB.Create(&remaining).Error)
+
+ count, err := ExpireDueSubscriptions(10)
+ require.NoError(t, err)
+ assert.Equal(t, 1, count)
+ var stored User
+ require.NoError(t, DB.First(&stored, user.Id).Error)
+ assert.Equal(t, "vip", stored.Group)
+ assert.Equal(t, "expired", getSubscriptionResetSub(t, expired.Id).Status)
+}
+
+func TestDowngradeUserGroupRetainsCurrentGroupWhenAnyActiveSubscriptionGrantsIt(t *testing.T) {
+ truncateTables(t)
+ now := GetDBTimestamp()
+ user := User{Id: 9721, Username: "subscription-retain-current-group", Group: "svip", Status: common.UserStatusEnabled}
+ replaced := UserSubscription{
+ Id: 9722, UserId: user.Id, Status: "cancelled", EndTime: now + 3600,
+ UpgradeGroup: "svip", PrevUserGroup: "vip",
+ }
+ matching := UserSubscription{
+ Id: 9723, UserId: user.Id, Status: "active", EndTime: now + 7200,
+ UpgradeGroup: "svip", PrevUserGroup: "vip",
+ }
+ lower := UserSubscription{
+ Id: 9724, UserId: user.Id, Status: "active", EndTime: now + 10800,
+ UpgradeGroup: "vip", PrevUserGroup: "default",
+ }
+ require.NoError(t, DB.Create(&user).Error)
+ require.NoError(t, DB.Create(&replaced).Error)
+ require.NoError(t, DB.Create(&matching).Error)
+ require.NoError(t, DB.Create(&lower).Error)
+
+ var group string
+ require.NoError(t, DB.Transaction(func(tx *gorm.DB) error {
+ var err error
+ group, err = downgradeUserGroupForSubscriptionTx(tx, &replaced, now)
+ return err
+ }))
+ assert.Empty(t, group)
+ var stored User
+ require.NoError(t, DB.First(&stored, user.Id).Error)
+ assert.Equal(t, "svip", stored.Group)
+}
+
+func TestExpireDueSubscriptionsRetainsCurrentGroupWhenAnyActiveSubscriptionGrantsIt(t *testing.T) {
+ truncateTables(t)
+ now := GetDBTimestamp()
+ user := User{Id: 9731, Username: "subscription-expiry-retain-current", Group: "svip", Status: common.UserStatusEnabled}
+ expired := UserSubscription{
+ Id: 9732, UserId: user.Id, Status: "active", EndTime: now - 1,
+ UpgradeGroup: "svip", PrevUserGroup: "vip",
+ }
+ matching := UserSubscription{
+ Id: 9733, UserId: user.Id, Status: "active", EndTime: now + 7200,
+ UpgradeGroup: "svip", PrevUserGroup: "vip",
+ }
+ lower := UserSubscription{
+ Id: 9734, UserId: user.Id, Status: "active", EndTime: now + 10800,
+ UpgradeGroup: "vip", PrevUserGroup: "default",
+ }
+ require.NoError(t, DB.Create(&user).Error)
+ require.NoError(t, DB.Create(&expired).Error)
+ require.NoError(t, DB.Create(&matching).Error)
+ require.NoError(t, DB.Create(&lower).Error)
+
+ count, err := ExpireDueSubscriptions(10)
+ require.NoError(t, err)
+ assert.Equal(t, 1, count)
+ var stored User
+ require.NoError(t, DB.First(&stored, user.Id).Error)
+ assert.Equal(t, "svip", stored.Group)
+}
diff --git a/model/task.go b/model/task.go
index 905632fd..1fe3624d 100644
--- a/model/task.go
+++ b/model/task.go
@@ -4,6 +4,8 @@ import (
"bytes"
"database/sql/driver"
"encoding/json"
+ "errors"
+ "fmt"
"time"
"github.com/MAX-API-Next/MAX-API/common"
@@ -11,6 +13,7 @@ import (
"github.com/MAX-API-Next/MAX-API/dto"
commonRelay "github.com/MAX-API-Next/MAX-API/relay/common"
"github.com/MAX-API-Next/MAX-API/types"
+ "gorm.io/gorm"
)
type TaskStatus string
@@ -64,11 +67,29 @@ type Task struct {
// 禁止返回给用户，内部可能包含key等隐私信息
PrivateData TaskPrivateData `json:"-" gorm:"column:private_data;type:json"`
Data json.RawMessage `json:"data" gorm:"type:json"`
+
+ includeDataInUpdate bool `gorm:"-"`
+ includePrivateDataInUpdate bool `gorm:"-"`
}
func (t *Task) SetData(data any) {
b, _ := common.Marshal(data)
t.Data = json.RawMessage(b)
+ t.includeDataInUpdate = true
+}
+
+func (t *Task) ClearDataForUpdate() {
+ t.Data = nil
+ t.includeDataInUpdate = true
+}
+
+func (t *Task) ClearPrivateDataForUpdate() {
+ t.PrivateData = TaskPrivateData{}
+ t.includePrivateDataInUpdate = true
+}
+
+func (t *Task) IncludePrivateDataInUpdate() {
+ t.includePrivateDataInUpdate = true
}
func (t *Task) GetData(v any) error {
@@ -100,13 +121,17 @@ func (m Properties) Value() (driver.Value, error) {
type TaskPrivateData struct {
Key string `json:"key,omitempty"`
UpstreamTaskID string `json:"upstream_task_id,omitempty"` // 上游真实 task ID
- ResultURL string `json:"result_url,omitempty"` // 任务成功后的结果 URL（视频地址等）
+ // AwaitingUpstreamID distinguishes a newly persisted placeholder from
+ // legacy rows where TaskID itself is the provider task identifier.
+ AwaitingUpstreamID bool `json:"awaiting_upstream_id,omitempty"`
+ ResultURL string `json:"result_url,omitempty"` // 任务成功后的结果 URL（视频地址等）
// 计费上下文：用于异步退款/差额结算（轮询阶段读取）
- BillingSource string `json:"billing_source,omitempty"` // "wallet" 或 "subscription"
- SubscriptionId int `json:"subscription_id,omitempty"` // 订阅 ID，用于订阅退款
- TokenId int `json:"token_id,omitempty"` // 令牌 ID，用于令牌额度退款
- NodeName string `json:"node_name,omitempty"` // 发起任务的节点名，轮询结算阶段据此归属日志而非最后查询节点
- BillingContext *TaskBillingContext `json:"billing_context,omitempty"` // 计费参数快照（用于轮询阶段重新计算）
+ BillingSource string `json:"billing_source,omitempty"` // "wallet" 或 "subscription"
+ BillingRequestId string `json:"billing_request_id,omitempty"` // 原始预扣 request id，用于跨周期安全结算
+ SubscriptionId int `json:"subscription_id,omitempty"` // 订阅 ID，用于订阅退款
+ TokenId int `json:"token_id,omitempty"` // 令牌 ID，用于令牌额度退款
+ NodeName string `json:"node_name,omitempty"` // 发起任务的节点名，轮询结算阶段据此归属日志而非最后查询节点
+ BillingContext *TaskBillingContext `json:"billing_context,omitempty"` // 计费参数快照（用于轮询阶段重新计算）
}
// TaskBillingContext 记录任务提交时的计费参数，以便轮询阶段可以重新计算额度。
@@ -144,6 +169,9 @@ func (t *Task) GetUpstreamTaskID() string {
if t.PrivateData.UpstreamTaskID != "" {
return t.PrivateData.UpstreamTaskID
}
+ if t.PrivateData.AwaitingUpstreamID {
+ return ""
+ }
return t.TaskID
}
@@ -327,9 +355,35 @@ func GetTimedOutUnfinishedTasks(cutoffUnix int64, limit int) []*Task {
}
func GetAllUnFinishSyncTasks(limit int) []*Task {
+ if limit <= 0 {
+ limit = 100
+ }
var tasks []*Task
- var err error
- err = DB.Where("status NOT IN ?", []string{TaskStatusFailure, TaskStatusSuccess}).Limit(limit).Order("id").Find(&tasks).Error
+ err := DB.Transaction(func(tx *gorm.DB) error {
+ if err := withRowLock(tx).Where("status NOT IN ?", []string{TaskStatusFailure, TaskStatusSuccess}).
+ Limit(limit).
+ Order("updated_at ASC, id ASC").
+ Find(&tasks).Error; err != nil {
+ return err
+ }
+ if len(tasks) == 0 {
+ return nil
+ }
+ ids := make([]int64, 0, len(tasks))
+ claimedAt := time.Now().Unix()
+ for _, task := range tasks {
+ ids = append(ids, task.ID)
+ if task.UpdatedAt >= claimedAt {
+ claimedAt = task.UpdatedAt + 1
+ }
+ }
+ for _, task := range tasks {
+ task.UpdatedAt = claimedAt
+ }
+ return tx.Model(&Task{}).
+ Where("id IN ? AND status NOT IN ?", ids, []string{TaskStatusFailure, TaskStatusSuccess}).
+ Update("updated_at", claimedAt).Error
+ })
if err != nil {
return nil
}
@@ -417,6 +471,46 @@ func (t *Task) Snapshot() taskSnapshot {
}
}
+func (t *Task) statusUpdateValues() map[string]interface{} {
+ updatedAt := t.UpdatedAt
+ if updatedAt <= 0 {
+ updatedAt = time.Now().Unix()
+ }
+ values := map[string]interface{}{
+ "status": t.Status,
+ "progress": t.Progress,
+ "start_time": t.StartTime,
+ "finish_time": t.FinishTime,
+ "fail_reason": t.FailReason,
+ "updated_at": updatedAt,
+ }
+ if t.Data != nil || t.includeDataInUpdate {
+ values["data"] = t.Data
+ }
+ if t.PrivateData != (TaskPrivateData{}) || t.includePrivateDataInUpdate {
+ values["private_data"] = t.PrivateData
+ }
+ return values
+}
+
+func (t *Task) submitResultUpdateValues() map[string]interface{} {
+ updatedAt := t.UpdatedAt
+ if updatedAt <= 0 {
+ updatedAt = time.Now().Unix()
+ }
+ values := map[string]interface{}{
+ "quota": t.Quota,
+ "updated_at": updatedAt,
+ }
+ if t.Data != nil || t.includeDataInUpdate {
+ values["data"] = t.Data
+ }
+ if t.PrivateData != (TaskPrivateData{}) || t.includePrivateDataInUpdate {
+ values["private_data"] = t.PrivateData
+ }
+ return values
+}
+
func (Task *Task) Update() error {
var err error
err = DB.Save(Task).Error
@@ -431,17 +525,147 @@ func (t *Task) UpdateQuota() error {
// Returns (true, nil) if this caller won the update, (false, nil) if
// another process already moved the task out of fromStatus.
//
-// Uses Model().Select("*").Updates() instead of Save() because GORM's Save
-// falls back to INSERT ON CONFLICT when the WHERE-guarded UPDATE matches
-// zero rows, which silently bypasses the CAS guard.
+// Uses a column map instead of Save() because GORM's Save falls back to
+// INSERT ON CONFLICT when the WHERE-guarded UPDATE matches zero rows, which
+// silently bypasses the CAS guard.
func (t *Task) UpdateWithStatus(fromStatus TaskStatus) (bool, error) {
- result := DB.Model(t).Where("status = ?", fromStatus).Select("*").Updates(t)
+ result := DB.Model(t).Where("status = ?", fromStatus).Updates(t.statusUpdateValues())
if result.Error != nil {
return false, result.Error
}
return result.RowsAffected > 0, nil
}
+var errTaskStatusCASLost = errors.New("task status compare-and-swap lost")
+
+// UpdateWithStatusAndSettlement commits a terminal task transition and its
+// durable settlement intent atomically. Balance mutations are applied after
+// this transaction by ApplyBillingSettlementOnce or its retry runner.
+func (t *Task) UpdateWithStatusAndSettlement(fromStatus TaskStatus, input BillingSettlementInput) (bool, error) {
+ if t == nil || t.ID <= 0 {
+ return false, errors.New("persisted task is required")
+ }
+ if input.TaskID != t.ID {
+ return false, fmt.Errorf("billing settlement task identity mismatch: task=%d input=%d", t.ID, input.TaskID)
+ }
+ if err := validateBillingSettlementInput(input); err != nil {
+ return false, err
+ }
+
+ err := DB.Transaction(func(tx *gorm.DB) error {
+ if _, _, err := ensureBillingSettlementRecordDB(tx, input); err != nil {
+ return err
+ }
+ result := tx.Model(t).Where("status = ?", fromStatus).Updates(t.statusUpdateValues())
+ if result.Error != nil {
+ return result.Error
+ }
+ if result.RowsAffected == 0 {
+ return errTaskStatusCASLost
+ }
+ return nil
+ })
+ if errors.Is(err, errTaskStatusCASLost) {
+ return false, nil
+ }
+ return err == nil, err
+}
+
+// UpdateWithSettlementIntent persists upstream task identity and the request's
+// final billing intent in one transaction before the HTTP success response is released.
+func (t *Task) UpdateWithSettlementIntent(input *BillingSettlementInput) error {
+ if t == nil || t.ID <= 0 {
+ return errors.New("persisted task is required")
+ }
+ return DB.Transaction(func(tx *gorm.DB) error {
+ if input != nil {
+ if err := validateBillingSettlementInput(*input); err != nil {
+ return err
+ }
+ if _, _, err := ensureBillingSettlementRecordDB(tx, *input); err != nil {
+ return err
+ }
+ }
+ result := tx.Model(&Task{}).
+ Where("id = ? AND status NOT IN ?", t.ID, []string{TaskStatusFailure, TaskStatusSuccess}).
+ Updates(t.submitResultUpdateValues())
+ if result.Error != nil {
+ return result.Error
+ }
+ if result.RowsAffected != 1 {
+ return taskSubmitResultUpdateMissError(tx, t.ID)
+ }
+ return nil
+ })
+}
+
+func taskSubmitResultUpdateMissError(tx *gorm.DB, taskID int64) error {
+ var existing Task
+ if err := tx.Select("id", "status").First(&existing, taskID).Error; err != nil {
+ if errors.Is(err, gorm.ErrRecordNotFound) {
+ return fmt.Errorf("persisted task not found: id=%d", taskID)
+ }
+ return err
+ }
+ if existing.Status == TaskStatusFailure || existing.Status == TaskStatusSuccess {
+ return fmt.Errorf("task already terminal (status=%s): id=%d", existing.Status, taskID)
+ }
+ return fmt.Errorf("persisted task not updated: id=%d status=%s", taskID, existing.Status)
+}
+
+func MarkTaskSubmitFailed(taskID int64, reason string) error {
+ if taskID <= 0 {
+ return nil
+ }
+ return DB.Model(&Task{}).
+ Where("id = ? AND status NOT IN ?", taskID, []string{TaskStatusFailure, TaskStatusSuccess}).
+ Updates(map[string]interface{}{
+ "status": TaskStatusFailure,
+ "progress": "100%",
+ "finish_time": time.Now().Unix(),
+ "fail_reason": reason,
+ "quota": 0,
+ "updated_at": time.Now().Unix(),
+ }).Error
+}
+
+// MarkTaskSubmitNeedsReview preserves the pre-consumed quota and any upstream
+// identity after the provider accepted a task but local durable finalization
+// failed. Refunding this state would create an unbilled upstream task.
+func MarkTaskSubmitNeedsReview(task *Task, reason string) error {
+ if task == nil || task.ID <= 0 {
+ return nil
+ }
+ task.PrivateData.AwaitingUpstreamID = false
+ return DB.Model(&Task{}).
+ Where("id = ? AND status <> ?", task.ID, TaskStatusSuccess).
+ Updates(map[string]interface{}{
+ "status": TaskStatusFailure,
+ "progress": "100%",
+ "finish_time": time.Now().Unix(),
+ "fail_reason": reason,
+ "quota": task.Quota,
+ "private_data": task.PrivateData,
+ "data": task.Data,
+ "updated_at": time.Now().Unix(),
+ }).Error
+}
+
+func MarkTaskSubmitAmbiguous(taskID int64, reason string) error {
+ if taskID <= 0 {
+ return nil
+ }
+ return DB.Model(&Task{}).
+ Where("id = ? AND status NOT IN ?", taskID, []string{TaskStatusFailure, TaskStatusSuccess}).
+ Updates(map[string]interface{}{
+ "status": TaskStatusFailure,
+ "progress": "100%",
+ "finish_time": time.Now().Unix(),
+ "fail_reason": reason,
+ "updated_at": time.Now().Unix(),
+ }).Error
+}
+
// TaskBulkUpdate performs an unconditional bulk UPDATE by upstream task_id strings.
// Same caveats as TaskBulkUpdateByID — no CAS guard.
func TaskBulkUpdate(taskIds []string, params map[string]any) error {
diff --git a/model/task_cas_test.go b/model/task_cas_test.go
index 4a977d0f..1ee7dcc5 100644
--- a/model/task_cas_test.go
+++ b/model/task_cas_test.go
@@ -2,6 +2,7 @@ package model
import (
"encoding/json"
+ "fmt"
"os"
"sync"
"testing"
@@ -39,6 +40,7 @@ func TestMain(m *testing.M) {
&User{},
&Token{},
&Log{},
+ &BillingLogReceipt{},
&Option{},
&Channel{},
&Ability{},
@@ -48,6 +50,9 @@ func TestMain(m *testing.M) {
&SubscriptionOrder{},
&UserSubscription{},
&SubscriptionPreConsumeRecord{},
+ &BillingSettlement{},
+ &BillingPreConsumeSelection{},
+ &CacheInvalidationTask{},
&PerfMetric{},
&SystemTask{},
&SystemInstance{},
@@ -65,6 +70,7 @@ func truncateTables(t *testing.T) {
DB.Exec("DELETE FROM users")
DB.Exec("DELETE FROM tokens")
DB.Exec("DELETE FROM logs")
+ DB.Exec("DELETE FROM billing_log_receipts")
DB.Exec("DELETE FROM channels")
DB.Exec("DELETE FROM abilities")
DB.Exec("DELETE FROM redemptions")
@@ -73,6 +79,9 @@ func truncateTables(t *testing.T) {
DB.Exec("DELETE FROM subscription_plans")
DB.Exec("DELETE FROM user_subscriptions")
DB.Exec("DELETE FROM subscription_pre_consume_records")
+ DB.Exec("DELETE FROM billing_settlements")
+ DB.Exec("DELETE FROM billing_pre_consume_selections")
+ DB.Exec("DELETE FROM cache_invalidation_tasks")
DB.Exec("DELETE FROM perf_metrics")
DB.Exec("DELETE FROM system_tasks")
DB.Exec("DELETE FROM system_instances")
@@ -177,6 +186,209 @@ func TestUpdateWithStatus_Win(t *testing.T) {
assert.Equal(t, "100%", reloaded.Progress)
}
+func TestUpdateWithStatusCanExplicitlyClearJSONColumns(t *testing.T) {
+ truncateTables(t)
+
+ task := &Task{
+ TaskID: "task_clear_json_columns",
+ Status: TaskStatusInProgress,
+ Progress: "50%",
+ Data: json.RawMessage(`{"id":"upstream"}`),
+ PrivateData: TaskPrivateData{ResultURL: "https://example.com/result.mp4"},
+ }
+ insertTask(t, task)
+
+ task.Status = TaskStatusQueued
+ task.Progress = "75%"
+ task.ClearDataForUpdate()
+ task.ClearPrivateDataForUpdate()
+ won, err := task.UpdateWithStatus(TaskStatusInProgress)
+ require.NoError(t, err)
+ require.True(t, won)
+
+ var reloaded Task
+ require.NoError(t, DB.First(&reloaded, task.ID).Error)
+ assert.EqualValues(t, TaskStatusQueued, reloaded.Status)
+ assert.Empty(t, reloaded.Data)
+ assert.Equal(t, TaskPrivateData{}, reloaded.PrivateData)
+}
+
+func TestUpdateWithStatusAndSettlementCommitsIntentAtomically(t *testing.T) {
+ truncateTables(t)
+ task := &Task{TaskID: "atomic-settlement", Status: TaskStatusInProgress, Quota: 10}
+ insertTask(t, task)
+ task.Status = TaskStatusFailure
+
+ won, err := task.UpdateWithStatusAndSettlement(TaskStatusInProgress, BillingSettlementInput{
+ OperationKey: fmt.Sprintf("task:%d:refund", task.ID),
+ Source: BillingSettlementSourceWallet,
+ UserID: 1,
+ FundingDelta: -10,
+ TaskID: task.ID,
+ TaskQuota: 10,
+ TaskQuotaTarget: 0,
+ })
+ require.NoError(t, err)
+ require.True(t, won)
+
+ var settlement BillingSettlement
+ require.NoError(t, DB.Where("operation_key = ?", fmt.Sprintf("task:%d:refund", task.ID)).First(&settlement).Error)
+ require.Equal(t, BillingSettlementStatusPending, settlement.Status)
+ var reloaded Task
+ require.NoError(t, DB.First(&reloaded, task.ID).Error)
+ require.EqualValues(t, TaskStatusFailure, reloaded.Status)
+ require.Equal(t, 10, reloaded.Quota)
+}
+
+func TestUpdateWithStatusAndSettlementCASLossLeavesNoIntent(t *testing.T) {
+ truncateTables(t)
+ task := &Task{TaskID: "lost-settlement", Status: TaskStatusSuccess, Quota: 10}
+ insertTask(t, task)
+ task.Status = TaskStatusFailure
+ operationKey := fmt.Sprintf("task:%d:refund", task.ID)
+
+ won, err := task.UpdateWithStatusAndSettlement(TaskStatusInProgress, BillingSettlementInput{
+ OperationKey: operationKey, Source: BillingSettlementSourceWallet,
+ UserID: 1, FundingDelta: -10, TaskID: task.ID, TaskQuota: 10, TaskQuotaTarget: 0,
+ })
+ require.NoError(t, err)
+ require.False(t, won)
+
+ var count int64
+ require.NoError(t, DB.Model(&BillingSettlement{}).Where("operation_key = ?", operationKey).Count(&count).Error)
+ require.Zero(t, count)
+}
+
+func TestUpdateWithSettlementIntentCommitsTaskAndIntentTogether(t *testing.T) {
+ truncateTables(t)
+ task := &Task{
+ TaskID: "submit-intent-commit", Status: TaskStatusNotStart, Quota: 10,
+ PrivateData: TaskPrivateData{AwaitingUpstreamID: true},
+ }
+ insertTask(t, task)
+ task.Quota = 15
+ task.Data = json.RawMessage(`{"id":"upstream-1"}`)
+ task.PrivateData.UpstreamTaskID = "upstream-1"
+ task.PrivateData.AwaitingUpstreamID = false
+ operationKey := "request:submit-intent-commit:finalize"
+
+ err := task.UpdateWithSettlementIntent(&BillingSettlementInput{
+ OperationKey: operationKey,
+ Source: BillingSettlementSourceWallet,
+ UserID: 1,
+ FundingDelta: 5,
+ TokenDelta: 5,
+ })
+ require.NoError(t, err)
+
+ var reloaded Task
+ require.NoError(t, DB.First(&reloaded, task.ID).Error)
+ require.Equal(t, 15, reloaded.Quota)
+ require.Equal(t, "upstream-1", reloaded.PrivateData.UpstreamTaskID)
+ require.False(t, reloaded.PrivateData.AwaitingUpstreamID)
+ var settlement BillingSettlement
+ require.NoError(t, DB.Where("operation_key = ?", operationKey).First(&settlement).Error)
+ require.Equal(t, int64(5), settlement.FundingDelta)
+}
+
+func TestUpdateWithSettlementIntentRollsBackIntentWhenTaskIsMissing(t *testing.T) {
+ truncateTables(t)
+ operationKey := "request:submit-intent-missing:finalize"
+ task := &Task{ID: 999999, TaskID: "missing", Status: TaskStatusNotStart}
+
+ err := task.UpdateWithSettlementIntent(&BillingSettlementInput{
+ OperationKey: operationKey,
+ Source: BillingSettlementSourceWallet,
+ UserID: 1,
+ FundingDelta: 5,
+ TokenDelta: 5,
+ })
+ require.Error(t, err)
+
+ var count int64
+ require.NoError(t, DB.Model(&BillingSettlement{}).Where("operation_key = ?", operationKey).Count(&count).Error)
+ require.Zero(t, count)
+}
+
+func TestUpdateWithSettlementIntentDoesNotResurrectTerminalTask(t *testing.T) {
+ truncateTables(t)
+ task := &Task{TaskID: "terminal-submit", Status: TaskStatusFailure, Quota: 10}
+ insertTask(t, task)
+ task.Status = TaskStatusNotStart
+ task.Quota = 15
+ operationKey := "request:terminal-submit:finalize"
+
+ err := task.UpdateWithSettlementIntent(&BillingSettlementInput{
+ OperationKey: operationKey,
+ Source: BillingSettlementSourceWallet,
+ UserID: 1,
+ FundingDelta: 5,
+ TokenDelta: 5,
+ })
+ require.Error(t, err)
+ require.ErrorContains(t, err, "task already terminal")
+ require.NotContains(t, err.Error(), "persisted task not found")
+
+ var reloaded Task
+ require.NoError(t, DB.First(&reloaded, task.ID).Error)
+ require.EqualValues(t, TaskStatusFailure, reloaded.Status)
+ require.Equal(t, 10, reloaded.Quota)
+ var count int64
+ require.NoError(t, DB.Model(&BillingSettlement{}).Where("operation_key = ?", operationKey).Count(&count).Error)
+ require.Zero(t, count)
+}
+
+func TestGetUpstreamTaskIDDistinguishesPlaceholderFromLegacyTask(t *testing.T) {
+ placeholder := &Task{
+ TaskID: "task_public", PrivateData: TaskPrivateData{AwaitingUpstreamID: true},
+ }
+ require.Empty(t, placeholder.GetUpstreamTaskID())
+
+ placeholder.PrivateData.UpstreamTaskID = "provider-task"
+ require.Equal(t, "provider-task", placeholder.GetUpstreamTaskID())
+
+ legacy := &Task{TaskID: "legacy-provider-task"}
+ require.Equal(t, "legacy-provider-task", legacy.GetUpstreamTaskID())
+}
+
+func TestMarkTaskSubmitNeedsReviewPreservesQuotaAndUpstreamIdentity(t *testing.T) {
+ truncateTables(t)
+ task := &Task{
+ TaskID: "manual-review", Status: TaskStatusNotStart, Quota: 20,
+ PrivateData: TaskPrivateData{AwaitingUpstreamID: true},
+ }
+ insertTask(t, task)
+ task.Quota = 25
+ task.PrivateData.UpstreamTaskID = "provider-manual-review"
+ task.Data = json.RawMessage(`{"id":"provider-manual-review"}`)
+
+ require.NoError(t, MarkTaskSubmitNeedsReview(task, "local finalize failed"))
+
+ var reloaded Task
+ require.NoError(t, DB.First(&reloaded, task.ID).Error)
+ require.EqualValues(t, TaskStatusFailure, reloaded.Status)
+ require.Equal(t, 25, reloaded.Quota)
+ require.Equal(t, "provider-manual-review", reloaded.PrivateData.UpstreamTaskID)
+ require.False(t, reloaded.PrivateData.AwaitingUpstreamID)
+ require.Contains(t, reloaded.FailReason, "local finalize failed")
+}
+
+func TestGetAllUnFinishSyncTasksRotatesPollingBatch(t *testing.T) {
+ truncateTables(t)
+ for i := 1; i <= 3; i++ {
+ insertTask(t, &Task{
+ TaskID: fmt.Sprintf("poll-fair-%d", i), Status: TaskStatusInProgress,
+ })
+ }
+
+ first := GetAllUnFinishSyncTasks(2)
+ require.Len(t, first, 2)
+ require.Equal(t, []string{"poll-fair-1", "poll-fair-2"}, []string{first[0].TaskID, first[1].TaskID})
+ second := GetAllUnFinishSyncTasks(2)
+ require.Len(t, second, 2)
+ require.Contains(t, []string{second[0].TaskID, second[1].TaskID}, "poll-fair-3")
+}
+
func TestGetAllUnFinishSyncTasksIncludesNonTerminalHundredProgress(t *testing.T) {
truncateTables(t)
diff --git a/model/token.go b/model/token.go
index 3c98c71b..2d23a30f 100644
--- a/model/token.go
+++ b/model/token.go
@@ -308,38 +308,56 @@ func (token *Token) Insert() error {
// Update Make sure your token's fields is completed, because this will update non-zero values
func (token *Token) Update() (err error) {
- err = DB.Model(token).Select("name", "status", "expired_time", "remain_quota", "unlimited_quota",
- "model_limits_enabled", "model_limits", "allow_ips", "group", "cross_group_retry").Updates(token).Error
- if err == nil && common.RedisEnabled {
- if cacheErr := invalidateTokenCache(token.Key); cacheErr != nil {
- common.SysLog("failed to invalidate token cache: " + cacheErr.Error())
- enqueueTokenCacheRetry(token.Key, false, cacheErr)
+ var cacheTask CacheInvalidationTask
+ err = DB.Transaction(func(tx *gorm.DB) error {
+ if err := tx.Model(token).Select("name", "status", "expired_time", "remain_quota", "unlimited_quota",
+ "model_limits_enabled", "model_limits", "allow_ips", "group", "cross_group_retry").Updates(token).Error; err != nil {
+ return err
}
+ var err error
+ cacheTask, err = stageTokenCacheInvalidationTx(tx, token.Key, false)
+ return err
+ })
+ if err != nil {
+ return err
}
- return err
+ dispatchStagedCacheInvalidation(cacheTask)
+ return nil
}
func (token *Token) SelectUpdate() (err error) {
// This can update zero values
- err = DB.Model(token).Select("accessed_time", "status").Updates(token).Error
- if err == nil && common.RedisEnabled {
- if cacheErr := invalidateTokenCache(token.Key); cacheErr != nil {
- common.SysLog("failed to invalidate token cache: " + cacheErr.Error())
- enqueueTokenCacheRetry(token.Key, false, cacheErr)
+ var cacheTask CacheInvalidationTask
+ err = DB.Transaction(func(tx *gorm.DB) error {
+ if err := tx.Model(token).Select("accessed_time", "status").Updates(token).Error; err != nil {
+ return err
}
+ var err error
+ cacheTask, err = stageTokenCacheInvalidationTx(tx, token.Key, false)
+ return err
+ })
+ if err != nil {
+ return err
}
- return err
+ dispatchStagedCacheInvalidation(cacheTask)
+ return nil
}
func (token *Token) Delete() (err error) {
- err = DB.Delete(token).Error
- if err == nil && common.RedisEnabled {
- if cacheErr := deleteTokenCache(token.Key); cacheErr != nil {
- common.SysLog("failed to invalidate token cache: " + cacheErr.Error())
- enqueueTokenCacheRetry(token.Key, true, cacheErr)
+ var cacheTask CacheInvalidationTask
+ err = DB.Transaction(func(tx *gorm.DB) error {
+ if err := tx.Delete(token).Error; err != nil {
+ return err
}
+ var err error
+ cacheTask, err = stageTokenCacheInvalidationTx(tx, token.Key, true)
+ return err
+ })
+ if err != nil {
+ return err
}
- return err
+ dispatchStagedCacheInvalidation(cacheTask)
+ return nil
}
func (token *Token) IsModelLimitsEnabled() bool {
@@ -433,53 +451,6 @@ func DecreaseTokenQuota(id int, key string, quota int) (err error) {
return nil
}
-// DecreaseTokenAndUserQuota atomically consumes both balances. This is used by
-// legacy pre-consume callers so a failed user update cannot strand a token
-// deduction that would otherwise require a best-effort rollback.
-func DecreaseTokenAndUserQuota(tokenId, userId int, key string, quota int) error {
- if quota < 0 {
- return errors.New("quota 不能为负数！")
- }
- if quota == 0 {
- return nil
- }
- err := DB.Transaction(func(tx *gorm.DB) error {
- tokenResult := tx.Model(&Token{}).
- Where("id = ? AND user_id = ? AND (unlimited_quota = ? OR remain_quota >= ?)", tokenId, userId, true, quota).
- Updates(map[string]interface{}{
- "remain_quota": gorm.Expr("remain_quota - ?", quota),
- "used_quota": gorm.Expr("used_quota + ?", quota),
- "accessed_time": common.GetTimestamp(),
- })
- if tokenResult.Error != nil {
- return tokenResult.Error
- }
- if tokenResult.RowsAffected != 1 {
- return fmt.Errorf("%w: id=%d, need=%d", ErrTokenQuotaInsufficient, tokenId, quota)
- }
-
- userResult := tx.Model(&User{}).
- Where("id = ? AND quota >= ?", userId, quota).
- Update("quota", gorm.Expr("quota - ?", quota))
- if userResult.Error != nil {
- return userResult.Error
- }
- if userResult.RowsAffected != 1 {
- return fmt.Errorf("%w: id=%d, need=%d", ErrUserQuotaInsufficient, userId, quota)
- }
- return nil
- })
- if err != nil {
- return err
- }
- invalidateTokenQuotaCache(key)
- if cacheErr := invalidateUserQuotaCache(userId); cacheErr != nil {
- common.SysLog(fmt.Sprintf("failed to invalidate user quota cache after atomic pre-consume (user_id=%d): %v", userId, cacheErr))
- enqueueUserCacheInvalidationRetry(userId, cacheErr)
- }
- return nil
-}
-
func decreaseTokenQuota(id int, quota int64) (err error) {
if quota == 0 {
return nil
@@ -563,30 +534,17 @@ func GetTokenKeysByIds(ids []int, userId int) ([]Token, error) {
// 配合 InvalidateUserCache 使用，可在用户被禁用/删除时立即阻断其令牌的请求。
// 下一次请求将从数据库重新加载令牌及用户状态，从而立即识别出被禁用的用户。
func InvalidateUserTokensCache(userId int) error {
- if !common.RedisEnabled {
- return nil
- }
if userId <= 0 {
return errors.New("userId 无效")
}
- var tokens []Token
- if err := DB.Unscoped().
- Select("id", commonKeyCol).
- Where("user_id = ?", userId).
- Find(&tokens).Error; err != nil {
+ var tasks []CacheInvalidationTask
+ if err := DB.Transaction(func(tx *gorm.DB) error {
+ var err error
+ tasks, err = stageUserTokenCacheInvalidationsTx(tx, userId, false)
+ return err
+ }); err != nil {
return err
}
- var firstErr error
- for _, t := range tokens {
- if t.Key == "" {
- continue
- }
- if err := invalidateTokenCache(t.Key); err != nil {
- enqueueTokenCacheRetry(t.Key, false, err)
- if firstErr == nil {
- firstErr = err
- }
- }
- }
- return firstErr
+ dispatchStagedCacheInvalidations(tasks)
+ return nil
}
diff --git a/model/token_cache.go b/model/token_cache.go
index 055e52ad..167f5cd2 100644
--- a/model/token_cache.go
+++ b/model/token_cache.go
@@ -1,12 +1,14 @@
package model
import (
+ "errors"
"fmt"
"sync"
"time"
"github.com/MAX-API-Next/MAX-API/common"
"github.com/bytedance/gopkg/util/gopool"
+ "gorm.io/gorm"
)
const (
@@ -55,9 +57,14 @@ func getTokenCacheVersionKey(key string) string {
func cacheSetTokenIfVersion(token Token, version int64) error {
key := token.Key
+ cacheKey := getTokenCacheKey(key)
+ pending, err := cacheInvalidationTaskPending(cacheInvalidationKindToken, cacheKey)
+ if err != nil || pending {
+ return err
+ }
token.Clean()
- _, err := common.RedisHSetObjIfVersion(
- getTokenCacheKey(key),
+ _, err = common.RedisHSetObjIfVersion(
+ cacheKey,
getTokenCacheVersionKey(key),
version,
&token,
@@ -80,6 +87,44 @@ func deleteTokenCache(key string) error {
return common.RedisDeleteVersionedHash(getTokenCacheKey(key), getTokenCacheVersionKey(key))
}
+func stageTokenCacheInvalidationTx(tx *gorm.DB, key string, deleteEntry bool) (CacheInvalidationTask, error) {
+ if key == "" {
+ return CacheInvalidationTask{}, errors.New("token key is required for cache invalidation")
+ }
+ return stageCacheInvalidationTaskTx(
+ tx,
+ cacheInvalidationKindToken,
+ getTokenCacheKey(key),
+ getTokenCacheVersionKey(key),
+ deleteEntry,
+ )
+}
+
+func stageUserTokenCacheInvalidationsTx(tx *gorm.DB, userID int, deleteEntry bool) ([]CacheInvalidationTask, error) {
+ if userID <= 0 {
+ return nil, errors.New("user id is required for token cache invalidation")
+ }
+ if !common.RedisEnabled {
+ return nil, nil
+ }
+ var tokens []Token
+ if err := tx.Unscoped().Select("key").Where("user_id = ?", userID).Find(&tokens).Error; err != nil {
+ return nil, err
+ }
+ tasks := make([]CacheInvalidationTask, 0, len(tokens))
+ for _, token := range tokens {
+ if token.Key == "" {
+ continue
+ }
+ task, err := stageTokenCacheInvalidationTx(tx, token.Key, deleteEntry)
+ if err != nil {
+ return nil, err
+ }
+ tasks = append(tasks, task)
+ }
+ return tasks, nil
+}
+
func tokenCacheRetryWindow() time.Duration {
window := time.Duration(common.RedisKeyCacheSeconds()) * time.Second
if window <= 0 {
@@ -92,13 +137,22 @@ func enqueueTokenCacheRetry(key string, deleteEntry bool, cause error) {
if key == "" {
return
}
- now := time.Now()
cacheKey := getTokenCacheKey(key)
versionKey := getTokenCacheVersionKey(key)
+ enqueueTokenCacheRetryByTarget(cacheKey, versionKey, deleteEntry, cause)
+}
+
+func enqueueTokenCacheRetryByTarget(cacheKey string, versionKey string, deleteEntry bool, cause error) {
+ if cacheKey == "" || versionKey == "" {
+ return
+ }
+ now := time.Now()
+ shouldPersist := false
tokenCacheRetries.Lock()
state, exists := tokenCacheRetries.pending[cacheKey]
if !exists {
+ shouldPersist = true
state = &tokenCacheRetryState{
cacheKey: cacheKey,
versionKey: versionKey,
@@ -119,6 +173,7 @@ func enqueueTokenCacheRetry(key string, deleteEntry bool, cause error) {
state.cause = cause
}
if deleteEntry && !state.deleteEntry {
+ shouldPersist = true
state.deleteEntry = true
state.deadline = now.Add(tokenCacheRetryWindow())
state.delay = tokenCacheRetryInitialDelay
@@ -131,6 +186,10 @@ func enqueueTokenCacheRetry(key string, deleteEntry bool, cause error) {
}
tokenCacheRetries.Unlock()
+ if shouldPersist {
+ persistCacheInvalidationTask(cacheInvalidationKindToken, cacheKey, versionKey, deleteEntry, cause)
+ }
+
select {
case tokenCacheRetries.wake <- struct{}{}:
default:
@@ -277,10 +336,18 @@ func cacheGetTokenByKey(key string) (*Token, error) {
return nil, fmt.Errorf("redis is not enabled")
}
var token Token
- err := common.RedisHGetObj(getTokenCacheKey(key), &token)
+ cacheKey := getTokenCacheKey(key)
+ err := common.RedisHGetObj(cacheKey, &token)
if err != nil {
return nil, err
}
+ pending, err := cacheInvalidationTaskPending(cacheInvalidationKindToken, cacheKey)
+ if err != nil {
+ return nil, err
+ }
+ if pending {
+ return nil, fmt.Errorf("token cache invalidation pending")
+ }
token.Key = key
return &token, nil
}
diff --git a/model/topup.go b/model/topup.go
index 763d9ca2..0486a609 100644
--- a/model/topup.go
+++ b/model/topup.go
@@ -109,15 +109,23 @@ func GetTopUpById(id int) *TopUp {
}
func GetTopUpByTradeNo(tradeNo string) *TopUp {
- var topUp *TopUp
- var err error
- err = DB.Where("trade_no = ?", tradeNo).First(&topUp).Error
- if err != nil {
- return nil
- }
+ topUp, _ := GetTopUpByTradeNoWithError(tradeNo)
return topUp
}
+// GetTopUpByTradeNoWithError keeps database failures distinct from a missing
+// order for payment callbacks, which must retry when the database is unhealthy.
+func GetTopUpByTradeNoWithError(tradeNo string) (*TopUp, error) {
+ topUp := &TopUp{}
+ if err := DB.Where("trade_no = ?", tradeNo).First(topUp).Error; err != nil {
+ if errors.Is(err, gorm.ErrRecordNotFound) {
+ return nil, ErrTopUpNotFound
+ }
+ return nil, err
+ }
+ return topUp, nil
+}
+
func UpdatePendingTopUpStatus(tradeNo string, expectedPaymentProvider string, targetStatus string) error {
if tradeNo == "" {
return errors.New("未提供支付单号")
diff --git a/model/user.go b/model/user.go
index 40d7c3fc..d00d2265 100644
--- a/model/user.go
+++ b/model/user.go
@@ -155,14 +155,20 @@ func UpdateUserSetting(userId int, setting dto.UserSetting) error {
return err
}
settingValue := string(settingBytes)
- result := DB.Model(&User{}).Where("id = ?", userId).Update("setting", settingValue)
- if err = ensureUserUpdateMatchedTx(DB, result, userId, errors.New("用户不存在")); err != nil {
+ var cacheTask CacheInvalidationTask
+ err = DB.Transaction(func(tx *gorm.DB) error {
+ result := tx.Model(&User{}).Where("id = ?", userId).Update("setting", settingValue)
+ if err := ensureUserUpdateMatchedTx(tx, result, userId, errors.New("用户不存在")); err != nil {
+ return err
+ }
+ var err error
+ cacheTask, err = stageUserCacheInvalidationTx(tx, userId, false)
+ return err
+ })
+ if err != nil {
return err
}
- if err = invalidateUserCache(userId); err != nil {
- common.SysLog(fmt.Sprintf("failed to update user setting cache: user_id=%d, error=%v", userId, err))
- enqueueUserCacheInvalidationRetry(userId, err)
- }
+ dispatchStagedCacheInvalidation(cacheTask)
return nil
}
@@ -789,6 +795,7 @@ func ensureEmailAvailableWithTx(tx *gorm.DB, email string, excludeUserID int) er
// user, preventing concurrent binds from sharing the same normalized address.
func BindEmailToUser(user *User, email string) error {
email = NormalizeEmail(email)
+ var cacheTask CacheInvalidationTask
if err := WithNormalizedEmailWriteTx(email, func(tx *gorm.DB) error {
if err := ensureEmailAvailableWithTx(tx, email, user.Id); err != nil {
return err
@@ -801,11 +808,17 @@ func BindEmailToUser(user *User, email string) error {
}
user.Email = email
user.NormalizedEmail = email
- return tx.First(user, user.Id).Error
+ if err := tx.First(user, user.Id).Error; err != nil {
+ return err
+ }
+ var err error
+ cacheTask, err = stageUserCacheInvalidationTx(tx, user.Id, false)
+ return err
}); err != nil {
return err
}
- return updateUserCache(*user)
+ dispatchStagedCacheInvalidation(cacheTask)
+ return nil
}
func (user *User) Insert(inviterId int) error {
@@ -915,20 +928,34 @@ func (user *User) FinalizeOAuthUserCreation(inviterId int) {
}
func (user *User) Update(updatePassword bool) error {
+ var cacheTask CacheInvalidationTask
if err := withUserOAuthIdentityMutationLock(DB, func(db *gorm.DB) error {
- return user.updateWithTx(db, updatePassword, nil)
+ return db.Transaction(func(tx *gorm.DB) error {
+ if err := user.updateWithTx(tx, updatePassword, nil); err != nil {
+ return err
+ }
+ var err error
+ cacheTask, err = stageUserCacheInvalidationTx(tx, user.Id, false)
+ return err
+ })
}); err != nil {
return err
}
- if err := updateUserCache(*user); err != nil {
- common.SysLog(fmt.Sprintf("failed to update user cache: user_id=%d, error=%v", user.Id, err))
- }
+ dispatchStagedCacheInvalidation(cacheTask)
return nil
}
func (user *User) UpdateFields(updatePassword bool, fields ...UserUpdateField) error {
+ var cacheTask CacheInvalidationTask
update := func(db *gorm.DB) error {
- return user.updateWithTx(db, updatePassword, fields)
+ return db.Transaction(func(tx *gorm.DB) error {
+ if err := user.updateWithTx(tx, updatePassword, fields); err != nil {
+ return err
+ }
+ var err error
+ cacheTask, err = stageUserCacheInvalidationTx(tx, user.Id, false)
+ return err
+ })
}
var err error
if isAccessTokenOnlyUserUpdate(updatePassword, fields) {
@@ -939,9 +966,7 @@ func (user *User) UpdateFields(updatePassword bool, fields ...UserUpdateField) e
if err != nil {
return err
}
- if err := updateUserCache(*user); err != nil {
- common.SysLog(fmt.Sprintf("failed to update user cache: user_id=%d, error=%v", user.Id, err))
- }
+ dispatchStagedCacheInvalidation(cacheTask)
return nil
}
@@ -1171,31 +1196,13 @@ func copyUnspecifiedUserUpdateValues(updates map[string]interface{}, current Use
}
func (user *User) Edit(updatePassword bool) error {
- var err error
- if updatePassword {
- user.Password, err = common.Password2Hash(user.Password)
- if err != nil {
- return err
- }
- }
-
- newUser := *user
- updates := map[string]interface{}{
- "username": newUser.Username,
- "display_name": newUser.DisplayName,
- "group": newUser.Group,
- "remark": newUser.Remark,
- }
- if updatePassword {
- updates["password"] = newUser.Password
- }
-
- DB.First(&user, user.Id)
- if err = DB.Model(user).Updates(updates).Error; err != nil {
- return err
- }
-
- return invalidateUserCache(user.Id)
+ return user.UpdateFields(
+ updatePassword,
+ UserUpdateFieldUsername,
+ UserUpdateFieldDisplayName,
+ UserUpdateFieldGroup,
+ UserUpdateFieldRemark,
+ )
}
func (user *User) ClearBinding(bindingType string) error {
@@ -1203,75 +1210,78 @@ func (user *User) ClearBinding(bindingType string) error {
return errors.New("user id is empty")
}
- bindingColumnMap := map[string]string{
- "email": "email",
- "github": "github_id",
- "discord": "discord_id",
- "oidc": "oidc_id",
- "wechat": "wechat_id",
- "telegram": "telegram_id",
- "linuxdo": "linux_do_id",
- }
-
- column, ok := bindingColumnMap[bindingType]
- if !ok {
+ var field UserUpdateField
+ switch bindingType {
+ case "email":
+ user.Email = ""
+ field = UserUpdateFieldEmail
+ case "github":
+ user.GitHubId = ""
+ field = UserUpdateFieldGitHubId
+ case "discord":
+ user.DiscordId = ""
+ field = UserUpdateFieldDiscordId
+ case "oidc":
+ user.OidcId = ""
+ field = UserUpdateFieldOidcId
+ case "wechat":
+ user.WeChatId = ""
+ field = UserUpdateFieldWeChatId
+ case "telegram":
+ user.TelegramId = ""
+ field = UserUpdateFieldTelegramId
+ case "linuxdo":
+ user.LinuxDOId = ""
+ field = UserUpdateFieldLinuxDOId
+ default:
return errors.New("invalid binding type")
}
-
- updates := map[string]interface{}{column: ""}
- if bindingType == "email" {
- updates["normalized_email"] = ""
- }
- update := func(db *gorm.DB) error {
- if err := db.Model(&User{}).Where("id = ?", user.Id).Updates(updates).Error; err != nil {
- return err
- }
- return db.Where("id = ?", user.Id).First(user).Error
- }
- var err error
- if bindingType == "email" {
- err = update(DB)
- } else {
- err = withUserOAuthIdentityMutationLock(DB, update)
- }
- if err != nil {
- return err
- }
-
- return updateUserCache(*user)
+ return user.UpdateFields(false, field)
}
func (user *User) Delete() error {
- if user.Id == 0 {
- return errors.New("id 为空！")
- }
- if err := user.ensureCanDelete(); err != nil {
- return err
- }
- if err := DB.Delete(user).Error; err != nil {
- return err
- }
-
- // 清除缓存
- deleteUserCacheAfterCommittedDelete(user.Id)
- return nil
+ return user.deleteWithCacheInvalidation(false)
}
func (user *User) HardDelete() error {
+ return user.deleteWithCacheInvalidation(true)
+}
+
+func (user *User) deleteWithCacheInvalidation(hardDelete bool) error {
if user.Id == 0 {
return errors.New("id 为空！")
}
- if err := user.ensureCanDelete(); err != nil {
- return err
- }
- if err := DB.Unscoped().Delete(user).Error; err != nil {
+ var tasks []CacheInvalidationTask
+ err := DB.Transaction(func(tx *gorm.DB) error {
+ if err := user.ensureCanDeleteTx(tx); err != nil {
+ return err
+ }
+ userTask, err := stageUserCacheInvalidationTx(tx, user.Id, true)
+ if err != nil {
+ return err
+ }
+ tokenTasks, err := stageUserTokenCacheInvalidationsTx(tx, user.Id, hardDelete)
+ if err != nil {
+ return err
+ }
+ tasks = append(tasks, userTask)
+ tasks = append(tasks, tokenTasks...)
+ if hardDelete {
+ if err := tx.Unscoped().Where("user_id = ?", user.Id).Delete(&Token{}).Error; err != nil {
+ return err
+ }
+ return tx.Unscoped().Where("id = ?", user.Id).Delete(&User{}).Error
+ }
+ return tx.Delete(user).Error
+ })
+ if err != nil {
return err
}
- deleteUserCacheAfterCommittedDelete(user.Id)
+ dispatchStagedCacheInvalidations(tasks)
return nil
}
-func (user *User) ensureCanDelete() error {
+func (user *User) ensureCanDeleteTx(tx *gorm.DB) error {
if user.Role == common.RoleRootUser {
return errors.New("cannot delete root user")
}
@@ -1279,7 +1289,7 @@ func (user *User) ensureCanDelete() error {
return nil
}
var existing User
- if err := DB.Unscoped().Select("role").First(&existing, "id = ?", user.Id).Error; err != nil {
+ if err := tx.Unscoped().Select("role").First(&existing, "id = ?", user.Id).Error; err != nil {
return err
}
if existing.Role == common.RoleRootUser {
@@ -1346,9 +1356,8 @@ func (user *User) UpdateGitHubId(newGitHubId string) error {
if err := validateOAuthIdentityLength("github_id", newGitHubId); err != nil {
return err
}
- return withUserOAuthIdentityMutationLock(DB, func(db *gorm.DB) error {
- return db.Model(user).Update("github_id", newGitHubId).Error
- })
+ user.GitHubId = newGitHubId
+ return user.UpdateFields(false, UserUpdateFieldGitHubId)
}
func (user *User) FillUserByDiscordId() error {
@@ -1636,21 +1645,27 @@ func IncreaseUserQuota[T quotaDeltaInteger](id int, quota T, _ bool) (err error)
if delta == 0 {
return nil
}
- if err := increaseUserQuota(id, delta); err != nil {
+ var cacheTask CacheInvalidationTask
+ err = DB.Transaction(func(tx *gorm.DB) error {
+ if err := increaseUserQuotaTx(tx, id, delta); err != nil {
+ return err
+ }
+ var err error
+ cacheTask, err = stageUserCacheInvalidationTx(tx, id, false)
+ return err
+ })
+ if err != nil {
return err
}
- if cacheErr := invalidateUserQuotaCache(id); cacheErr != nil {
- common.SysLog("failed to invalidate user quota cache: " + cacheErr.Error())
- enqueueUserCacheInvalidationRetry(id, cacheErr)
- }
+ dispatchStagedCacheInvalidation(cacheTask)
return nil
}
-func increaseUserQuota(id int, quota int64) (err error) {
+func increaseUserQuotaTx(tx *gorm.DB, id int, quota int64) (err error) {
if quota == 0 {
return nil
}
- result := DB.Model(&User{}).Where("id = ?", id).Update("quota", gorm.Expr("quota + ?", quota))
+ result := tx.Model(&User{}).Where("id = ?", id).Update("quota", gorm.Expr("quota + ?", quota))
if result.Error != nil {
return result.Error
}
@@ -1669,21 +1684,27 @@ func DecreaseUserQuota[T quotaDeltaInteger](id int, quota T, _ bool) (err error)
if delta == 0 {
return nil
}
- if err := decreaseUserQuota(id, delta); err != nil {
+ var cacheTask CacheInvalidationTask
+ err = DB.Transaction(func(tx *gorm.DB) error {
+ if err := decreaseUserQuotaTx(tx, id, delta); err != nil {
+ return err
+ }
+ var err error
+ cacheTask, err = stageUserCacheInvalidationTx(tx, id, false)
+ return err
+ })
+ if err != nil {
return err
}
- if cacheErr := invalidateUserQuotaCache(id); cacheErr != nil {
- common.SysLog("failed to invalidate user quota cache: " + cacheErr.Error())
- enqueueUserCacheInvalidationRetry(id, cacheErr)
- }
+ dispatchStagedCacheInvalidation(cacheTask)
return nil
}
-func decreaseUserQuota(id int, quota int64) (err error) {
+func decreaseUserQuotaTx(tx *gorm.DB, id int, quota int64) (err error) {
if quota == 0 {
return nil
}
- result := DB.Model(&User{}).Where("id = ? AND quota >= ?", id, quota).
+ result := tx.Model(&User{}).Where("id = ? AND quota >= ?", id, quota).
Update("quota", gorm.Expr("quota - ?", quota))
if result.Error != nil {
return result.Error
diff --git a/model/user_cache.go b/model/user_cache.go
index 30034b9d..9b9273d2 100644
--- a/model/user_cache.go
+++ b/model/user_cache.go
@@ -1,6 +1,7 @@
package model
import (
+ "errors"
"fmt"
"sync"
"time"
@@ -12,6 +13,7 @@ import (
"github.com/gin-gonic/gin"
"github.com/bytedance/gopkg/util/gopool"
+ "gorm.io/gorm"
)
const (
@@ -101,10 +103,6 @@ func enqueueUserCacheInvalidationRetry(userId int, cause error) {
enqueueUserCacheRetry(userId, false, cause)
}
-func enqueueUserCacheDeletionRetry(userId int, cause error) {
- enqueueUserCacheRetry(userId, true, cause)
-}
-
func deleteUserCache(userId int) error {
if !common.RedisEnabled {
return nil
@@ -112,11 +110,17 @@ func deleteUserCache(userId int) error {
return common.RedisDeleteVersionedHash(getUserCacheKey(userId), getUserCacheVersionKey(userId))
}
-func deleteUserCacheAfterCommittedDelete(userId int) {
- if err := deleteUserCache(userId); err != nil {
- common.SysLog(fmt.Sprintf("failed to delete user cache after user deletion: user_id=%d, error=%v", userId, err))
- enqueueUserCacheDeletionRetry(userId, err)
+func stageUserCacheInvalidationTx(tx *gorm.DB, userID int, deleteEntry bool) (CacheInvalidationTask, error) {
+ if userID <= 0 {
+ return CacheInvalidationTask{}, errors.New("user id is required for cache invalidation")
}
+ return stageCacheInvalidationTaskTx(
+ tx,
+ cacheInvalidationKindUser,
+ fmt.Sprintf("%d", userID),
+ getUserCacheVersionKey(userID),
+ deleteEntry,
+ )
}
func userCacheRetryWindow() time.Duration {
@@ -132,9 +136,11 @@ func enqueueUserCacheRetry(userId int, deleteEntry bool, cause error) {
return
}
now := time.Now()
+ shouldPersist := false
userCacheRetries.Lock()
state, exists := userCacheRetries.pending[userId]
if !exists {
+ shouldPersist = true
state = &userCacheRetryState{
userId: userId,
revision: 1,
@@ -154,6 +160,7 @@ func enqueueUserCacheRetry(userId int, deleteEntry bool, cause error) {
state.cause = cause
}
if deleteEntry && !state.deleteEntry {
+ shouldPersist = true
state.deleteEntry = true
state.deadline = now.Add(userCacheRetryWindow())
state.delay = userCacheRetryInitialDelay
@@ -166,6 +173,10 @@ func enqueueUserCacheRetry(userId int, deleteEntry bool, cause error) {
}
userCacheRetries.Unlock()
+ if shouldPersist {
+ persistCacheInvalidationTask(cacheInvalidationKindUser, fmt.Sprintf("%d", userId), getUserCacheVersionKey(userId), deleteEntry, cause)
+ }
+
select {
case userCacheRetries.wake <- struct{}{}:
default:
@@ -309,14 +320,22 @@ func completeUserCacheRetryAttempt(attempt userCacheRetryAttempt, err error) {
// InvalidateUserCache is the exported version of invalidateUserCache.
// 供 controller 等上层包在用户状态变更（如禁用、删除、角色变更）后主动清理缓存。
func InvalidateUserCache(userId int) error {
- return invalidateUserCache(userId)
+ if err := invalidateUserCache(userId); err != nil {
+ enqueueUserCacheInvalidationRetry(userId, err)
+ return err
+ }
+ return nil
}
func populateUserCacheIfVersion(user User, version int64) error {
if !common.RedisEnabled {
return nil
}
- _, err := common.RedisHSetObjIfVersion(
+ pending, err := cacheInvalidationTaskPending(cacheInvalidationKindUser, fmt.Sprintf("%d", user.Id))
+ if err != nil || pending {
+ return err
+ }
+ _, err = common.RedisHSetObjIfVersion(
getUserCacheKey(user.Id),
getUserCacheVersionKey(user.Id),
version,
@@ -335,7 +354,11 @@ func updateUserCache(user User) error {
}
// User mutations must advance the generation before any later DB refill can
// write a stale snapshot. The next read repopulates the complete hash.
- return invalidateUserCache(user.Id)
+ if err := invalidateUserCache(user.Id); err != nil {
+ enqueueUserCacheInvalidationRetry(user.Id, err)
+ return err
+ }
+ return nil
}
// GetUserCache gets complete user cache from hash
@@ -397,6 +420,13 @@ func cacheGetUserBase(userId int) (*UserBase, error) {
if userCache.Role < common.RoleCommonUser || userCache.Status == 0 {
return nil, fmt.Errorf("stale user cache")
}
+ pending, err := cacheInvalidationTaskPending(cacheInvalidationKindUser, fmt.Sprintf("%d", userId))
+ if err != nil {
+ return nil, err
+ }
+ if pending {
+ return nil, fmt.Errorf("user cache invalidation pending")
+ }
return &userCache, nil
}
@@ -445,7 +475,11 @@ func updateUserQuotaCacheIfVersion(userId int, quota int64, version int64) error
if !common.RedisEnabled {
return nil
}
- _, err := common.RedisHSetFieldIfVersion(
+ pending, err := cacheInvalidationTaskPending(cacheInvalidationKindUser, fmt.Sprintf("%d", userId))
+ if err != nil || pending {
+ return err
+ }
+ _, err = common.RedisHSetFieldIfVersion(
getUserCacheKey(userId),
getUserCacheVersionKey(userId),
version,
@@ -457,7 +491,11 @@ func updateUserQuotaCacheIfVersion(userId int, quota int64, version int64) error
}
func updateUserCacheFieldIfVersion(userId int, field string, value interface{}, version int64) error {
- _, err := common.RedisHSetFieldIfVersion(
+ pending, err := cacheInvalidationTaskPending(cacheInvalidationKindUser, fmt.Sprintf("%d", userId))
+ if err != nil || pending {
+ return err
+ }
+ _, err = common.RedisHSetFieldIfVersion(
getUserCacheKey(userId),
getUserCacheVersionKey(userId),
version,
diff --git a/model/user_update_test.go b/model/user_update_test.go
index 0f945f9c..664080d8 100644
--- a/model/user_update_test.go
+++ b/model/user_update_test.go
@@ -95,24 +95,6 @@ func TestUnlimitedTokenQuotaDeductionRemainsUnbounded(t *testing.T) {
assert.EqualValues(t, 7, got.UsedQuota)
}
-func TestAtomicTokenAndUserPreConsumeRollsBackTokenOnUserFailure(t *testing.T) {
- setupUserUpdateTestState(t)
- user := User{Id: 38, Username: "atomic-preconsume-user", Quota: 5, Status: common.UserStatusEnabled}
- token := Token{Id: 39, UserId: user.Id, Key: "atomic-preconsume-token", RemainQuota: 10}
- require.NoError(t, DB.Create(&user).Error)
- require.NoError(t, DB.Create(&token).Error)
-
- require.ErrorIs(t, DecreaseTokenAndUserQuota(token.Id, user.Id, token.Key, 7), ErrUserQuotaInsufficient)
-
- var storedToken Token
- require.NoError(t, DB.First(&storedToken, token.Id).Error)
- assert.EqualValues(t, 10, storedToken.RemainQuota)
- assert.Zero(t, storedToken.UsedQuota)
- var storedUser User
- require.NoError(t, DB.First(&storedUser, user.Id).Error)
- assert.EqualValues(t, 5, storedUser.Quota)
-}
-
func TestQuotaMutationDoesNotTouchCacheAfterDatabaseFailure(t *testing.T) {
setupUserUpdateTestState(t)
@@ -267,6 +249,13 @@ type blockingCacheFillHook struct {
afterOnce sync.Once
}
+type blockingCacheInvalidationHook struct {
+ started chan struct{}
+ release chan struct{}
+ releaseOnce sync.Once
+ seen atomic.Int32
+}
+
type failCacheMutationHook struct {
seen atomic.Int32
failAt int32
@@ -318,6 +307,43 @@ func (*overlappingCacheMutationHook) AfterProcessPipeline(context.Context, []red
return nil
}
+func newBlockingCacheInvalidationHook() *blockingCacheInvalidationHook {
+ return &blockingCacheInvalidationHook{
+ started: make(chan struct{}),
+ release: make(chan struct{}),
+ }
+}
+
+func (h *blockingCacheInvalidationHook) BeforeProcess(ctx context.Context, cmd redis.Cmder) (context.Context, error) {
+ name := strings.ToLower(cmd.Name())
+ if name != "eval" && name != "evalsha" {
+ return ctx, nil
+ }
+ if h.seen.Add(1) == 1 {
+ close(h.started)
+ <-h.release
+ }
+ return ctx, nil
+}
+
+func (*blockingCacheInvalidationHook) AfterProcess(context.Context, redis.Cmder) error {
+ return nil
+}
+
+func (*blockingCacheInvalidationHook) BeforeProcessPipeline(ctx context.Context, _ []redis.Cmder) (context.Context, error) {
+ return ctx, nil
+}
+
+func (*blockingCacheInvalidationHook) AfterProcessPipeline(context.Context, []redis.Cmder) error {
+ return nil
+}
+
+func (h *blockingCacheInvalidationHook) unblock() {
+ h.releaseOnce.Do(func() {
+ close(h.release)
+ })
+}
+
func (h *recoverableCacheMutationHook) BeforeProcess(ctx context.Context, cmd redis.Cmder) (context.Context, error) {
name := strings.ToLower(cmd.Name())
if name != "eval" && name != "evalsha" {
@@ -431,26 +457,125 @@ func TestUserQuotaMutationsRetryCacheInvalidation(t *testing.T) {
cacheUserForRetryTest(t, client, user)
require.NoError(t, tt.mutate(user.Id))
+ processCacheInvalidationTasks()
requireCacheKeyDeletedEventually(t, client, getUserCacheKey(user.Id))
})
}
}
-func TestAtomicPreConsumeRetriesBothCacheInvalidations(t *testing.T) {
+func TestUserAuthFieldMutationsRetryCacheInvalidation(t *testing.T) {
setupUserUpdateTestState(t)
- user := User{Id: 83, Username: "atomic-cache-retry-user", Quota: 10, Status: common.UserStatusEnabled}
- token := Token{Id: 84, UserId: user.Id, Key: "atomic-cache-retry-token", RemainQuota: 10, Status: common.TokenStatusEnabled}
+ user := User{Id: 82, Username: "auth-cache-retry-user", Quota: 10, Status: common.UserStatusEnabled, Role: common.RoleCommonUser}
require.NoError(t, DB.Create(&user).Error)
- require.NoError(t, DB.Create(&token).Error)
- client, _ := useFailingCacheMutationRedis(t, 2)
+ client, _ := useFailingCacheMutationRedis(t, 1)
cacheUserForRetryTest(t, client, user)
- cacheTokenForRetryTest(t, client, token)
- require.NoError(t, DecreaseTokenAndUserQuota(token.Id, user.Id, token.Key, 1))
- requireCacheKeyDeletedEventually(t, client, getTokenCacheKey(token.Key))
+ user.Status = common.UserStatusDisabled
+ require.NoError(t, user.UpdateFields(false, UserUpdateFieldStatus))
+ processCacheInvalidationTasks()
requireCacheKeyDeletedEventually(t, client, getUserCacheKey(user.Id))
}
+func TestUserMutationRollsBackWhenCacheOutboxCannotPersist(t *testing.T) {
+ setupUserUpdateTestState(t)
+ user := User{Id: 93, Username: "cache-outbox-user", Quota: 10, Status: common.UserStatusEnabled}
+ require.NoError(t, DB.Create(&user).Error)
+ useFailingUserUpdateRedis(t)
+ failCacheOutboxInserts(t)
+
+ user.Status = common.UserStatusDisabled
+ err := user.UpdateFields(false, UserUpdateFieldStatus)
+
+ require.Error(t, err)
+ var stored User
+ require.NoError(t, DB.First(&stored, user.Id).Error)
+ assert.Equal(t, common.UserStatusEnabled, stored.Status)
+}
+
+func TestDirectUserMutationsRollBackWhenCacheOutboxCannotPersist(t *testing.T) {
+ tests := []struct {
+ name string
+ mutate func(*User) error
+ check func(*testing.T, User)
+ }{
+ {
+ name: "edit",
+ mutate: func(user *User) error {
+ user.DisplayName = "after"
+ return user.Edit(false)
+ },
+ check: func(t *testing.T, stored User) { assert.Equal(t, "before", stored.DisplayName) },
+ },
+ {
+ name: "clear binding",
+ mutate: func(user *User) error {
+ return user.ClearBinding("github")
+ },
+ check: func(t *testing.T, stored User) { assert.Equal(t, "github-before", stored.GitHubId) },
+ },
+ {
+ name: "bind email",
+ mutate: func(user *User) error {
+ return BindEmailToUser(user, "after@example.com")
+ },
+ check: func(t *testing.T, stored User) { assert.Equal(t, "before@example.com", stored.Email) },
+ },
+ {
+ name: "github migration",
+ mutate: func(user *User) error {
+ return user.UpdateGitHubId("github-after")
+ },
+ check: func(t *testing.T, stored User) { assert.Equal(t, "github-before", stored.GitHubId) },
+ },
+ {
+ name: "setting",
+ mutate: func(user *User) error {
+ return UpdateUserSetting(user.Id, dto.UserSetting{Language: "zh"})
+ },
+ check: func(t *testing.T, stored User) { assert.Equal(t, `{"language":"en"}`, stored.Setting) },
+ },
+ }
+
+ for index, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ setupUserUpdateTestState(t)
+ user := User{
+ Id: 320 + index,
+ Username: "direct-outbox-" + tt.name,
+ DisplayName: "before",
+ Email: "before@example.com",
+ GitHubId: "github-before",
+ Setting: `{"language":"en"}`,
+ Status: common.UserStatusEnabled,
+ }
+ require.NoError(t, DB.Create(&user).Error)
+ useFailingUserUpdateRedis(t)
+ failCacheOutboxInserts(t)
+
+ require.Error(t, tt.mutate(&user))
+ var stored User
+ require.NoError(t, DB.First(&stored, user.Id).Error)
+ tt.check(t, stored)
+ })
+ }
+}
+
+func TestTokenMutationRollsBackWhenCacheOutboxCannotPersist(t *testing.T) {
+ setupUserUpdateTestState(t)
+ token := Token{Id: 94, UserId: 95, Key: "cache-outbox-token", Status: common.TokenStatusEnabled, RemainQuota: 10}
+ require.NoError(t, DB.Create(&token).Error)
+ useFailingUserUpdateRedis(t)
+ failCacheOutboxInserts(t)
+
+ token.Status = common.TokenStatusDisabled
+ err := token.Update()
+
+ require.Error(t, err)
+ var stored Token
+ require.NoError(t, DB.First(&stored, token.Id).Error)
+ assert.Equal(t, common.TokenStatusEnabled, stored.Status)
+}
+
func TestTokenQuotaMutationRetriesCacheInvalidation(t *testing.T) {
setupUserUpdateTestState(t)
token := Token{Id: 85, UserId: 86, Key: "token-quota-cache-retry", RemainQuota: 10, Status: common.TokenStatusEnabled}
@@ -470,6 +595,7 @@ func TestTokenDeleteRetriesCacheDeletion(t *testing.T) {
cacheTokenForRetryTest(t, client, token)
require.NoError(t, token.Delete())
+ processCacheInvalidationTasks()
requireCacheKeyDeletedEventually(t, client, getTokenCacheKey(token.Key))
}
@@ -481,6 +607,7 @@ func TestUpdateUserSettingRetriesCacheInvalidation(t *testing.T) {
cacheUserForRetryTest(t, client, user)
require.NoError(t, UpdateUserSetting(user.Id, dto.UserSetting{Language: "zh"}))
+ processCacheInvalidationTasks()
requireCacheKeyDeletedEventually(t, client, getUserCacheKey(user.Id))
}
@@ -489,19 +616,22 @@ func TestUserDeletesRetryCacheDeletionAfterDatabaseCommit(t *testing.T) {
name string
userID int
deleteUser func(*User) error
- wantRows int64
+ wantUsers int64
+ wantTokens int64
}{
{
name: "soft delete",
userID: 93,
deleteUser: (*User).Delete,
- wantRows: 1,
+ wantUsers: 1,
+ wantTokens: 1,
},
{
name: "hard delete",
userID: 94,
deleteUser: (*User).HardDelete,
- wantRows: 0,
+ wantUsers: 0,
+ wantTokens: 0,
},
}
@@ -510,19 +640,39 @@ func TestUserDeletesRetryCacheDeletionAfterDatabaseCommit(t *testing.T) {
setupUserUpdateTestState(t)
user := User{Id: tt.userID, Username: "delete-cache-retry-user-" + tt.name, Status: common.UserStatusEnabled}
require.NoError(t, DB.Create(&user).Error)
+ token := Token{Id: tt.userID + 1000, UserId: user.Id, Key: "delete-cache-token-" + tt.name, Status: common.TokenStatusEnabled}
+ require.NoError(t, DB.Create(&token).Error)
client, _ := useFailingCacheMutationRedis(t, 1)
cacheUserForRetryTest(t, client, user)
+ cacheTokenForRetryTest(t, client, token)
require.NoError(t, tt.deleteUser(&user))
+ processCacheInvalidationTasks()
var count int64
require.NoError(t, DB.Unscoped().Model(&User{}).Where("id = ?", user.Id).Count(&count).Error)
- assert.Equal(t, tt.wantRows, count)
+ assert.Equal(t, tt.wantUsers, count)
+ require.NoError(t, DB.Unscoped().Model(&Token{}).Where("id = ?", token.Id).Count(&count).Error)
+ assert.Equal(t, tt.wantTokens, count)
requireCacheKeyDeletedEventually(t, client, getUserCacheKey(user.Id))
+ requireCacheKeyDeletedEventually(t, client, getTokenCacheKey(token.Key))
})
}
}
-func TestTokenDeleteRetriesUntilCacheRecovers(t *testing.T) {
+func TestUserDeleteRollsBackWhenCacheOutboxCannotPersist(t *testing.T) {
+ setupUserUpdateTestState(t)
+ user := User{Id: 395, Username: "delete-outbox-rollback", Status: common.UserStatusEnabled}
+ require.NoError(t, DB.Create(&user).Error)
+ useFailingUserUpdateRedis(t)
+ failCacheOutboxInserts(t)
+
+ require.Error(t, user.Delete())
+ var count int64
+ require.NoError(t, DB.Unscoped().Model(&User{}).Where("id = ?", user.Id).Count(&count).Error)
+ assert.EqualValues(t, 1, count)
+}
+
+func TestTokenDeletePersistsUntilCacheRecovers(t *testing.T) {
setupUserUpdateTestState(t)
token := Token{Id: 91, UserId: 92, Key: "token-delete-cache-recovery", RemainQuota: 10, Status: common.TokenStatusEnabled}
require.NoError(t, DB.Create(&token).Error)
@@ -544,12 +694,15 @@ func TestTokenDeleteRetriesUntilCacheRecovers(t *testing.T) {
cacheTokenForRetryTest(t, client, token)
require.NoError(t, token.Delete())
- require.Eventually(t, func() bool { return hook.seen.Load() >= 4 }, 2*time.Second, 10*time.Millisecond)
exists, err := client.Exists(context.Background(), getTokenCacheKey(token.Key)).Result()
require.NoError(t, err)
require.EqualValues(t, 1, exists)
+ var pending int64
+ require.NoError(t, DB.Model(&CacheInvalidationTask{}).Count(&pending).Error)
+ require.EqualValues(t, 1, pending)
hook.available.Store(true)
+ processCacheInvalidationTasks()
requireCacheKeyDeletedEventually(t, client, getTokenCacheKey(token.Key))
}
@@ -687,6 +840,30 @@ func TestUserQuotaInvalidationRejectsOlderAsyncRefill(t *testing.T) {
}, 3*time.Second, 10*time.Millisecond)
}
+func TestUserQuotaMutationPendingInvalidationBypassesStaleCache(t *testing.T) {
+ setupUserUpdateTestState(t)
+ user := User{Id: 39, Username: "quota-pending-invalidation-user", Quota: 100, Status: common.UserStatusEnabled}
+ require.NoError(t, DB.Create(&user).Error)
+ client := useCacheMutationRedis(t)
+ cacheUserForRetryTest(t, client, user)
+ hook := newBlockingCacheInvalidationHook()
+ client.AddHook(hook)
+ done := make(chan error, 1)
+
+ go func() {
+ done <- IncreaseUserQuota(user.Id, 10, false)
+ }()
+ waitForCacheHook(t, hook.started, "invalidation start")
+ t.Cleanup(hook.unblock)
+
+ quota, err := GetUserQuota(user.Id, false)
+ require.NoError(t, err)
+ assert.EqualValues(t, 110, quota)
+
+ hook.unblock()
+ require.NoError(t, <-done)
+}
+
func TestUserGroupInvalidationRejectsOlderNarrowRefill(t *testing.T) {
setupUserUpdateTestState(t)
hook, _ := useBlockingCacheFillRedis(t)
@@ -732,6 +909,29 @@ func TestTokenQuotaInvalidationRejectsOlderAsyncRefill(t *testing.T) {
}, 3*time.Second, 10*time.Millisecond)
}
+func TestTokenPendingInvalidationBypassesStaleCache(t *testing.T) {
+ setupUserUpdateTestState(t)
+ client := useCacheMutationRedis(t)
+ token := Token{Id: 38, UserId: 36, Key: "token-pending-cache", RemainQuota: 90, Status: common.TokenStatusEnabled}
+ require.NoError(t, DB.Create(&token).Error)
+ staleToken := token
+ staleToken.RemainQuota = 100
+ cacheTokenForRetryTest(t, client, staleToken)
+ _, err := upsertCacheInvalidationTask(
+ DB,
+ cacheInvalidationKindToken,
+ getTokenCacheKey(token.Key),
+ getTokenCacheVersionKey(token.Key),
+ false,
+ nil,
+ )
+ require.NoError(t, err)
+
+ got, err := GetTokenByKey(token.Key, false)
+ require.NoError(t, err)
+ assert.EqualValues(t, 90, got.RemainQuota)
+}
+
func setupUserUpdateTestState(t *testing.T) {
t.Helper()
truncateTables(t)
diff --git a/model/vendor_meta.go b/model/vendor_meta.go
index 5443207c..fa6ac4fe 100644
--- a/model/vendor_meta.go
+++ b/model/vendor_meta.go
@@ -1,11 +1,15 @@
package model
import (
+ "errors"
+
"github.com/MAX-API-Next/MAX-API/common"
"gorm.io/gorm"
)
+var ErrVendorHasModels = errors.New("vendor is still referenced by models")
+
// Vendor 用于存储供应商信息，供模型引用
// Name 唯一，用于在模型中关联
// Icon 采用 @lobehub/icons 的图标名，前端可直接渲染
@@ -14,13 +18,14 @@ import (
type Vendor struct {
Id int `json:"id"`
- Name string `json:"name" gorm:"size:128;not null;uniqueIndex:uk_vendor_name_delete_at,priority:1"`
+ Name string `json:"name" gorm:"size:128;not null"`
+ NameKey string `json:"-" gorm:"size:191;not null;default:''"`
Description string `json:"description,omitempty" gorm:"type:text"`
Icon string `json:"icon,omitempty" gorm:"type:varchar(128)"`
Status int `json:"status" gorm:"default:1"`
CreatedTime int64 `json:"created_time" gorm:"bigint"`
UpdatedTime int64 `json:"updated_time" gorm:"bigint"`
- DeletedAt gorm.DeletedAt `json:"-" gorm:"index;uniqueIndex:uk_vendor_name_delete_at,priority:2"`
+ DeletedAt gorm.DeletedAt `json:"-" gorm:"index"`
}
// Insert 创建新的供应商记录
@@ -28,6 +33,7 @@ func (v *Vendor) Insert() error {
now := common.GetTimestamp()
v.CreatedTime = now
v.UpdatedTime = now
+ v.NameKey = activeMetadataNameKey(v.Name)
return DB.Create(v).Error
}
@@ -44,12 +50,27 @@ func IsVendorNameDuplicated(id int, name string) (bool, error) {
// Update 更新供应商记录
func (v *Vendor) Update() error {
v.UpdatedTime = common.GetTimestamp()
+ v.NameKey = activeMetadataNameKey(v.Name)
return DB.Save(v).Error
}
// Delete 软删除供应商
func (v *Vendor) Delete() error {
- return DB.Delete(v).Error
+ return DB.Transaction(func(tx *gorm.DB) error {
+ var modelCount int64
+ if err := tx.Model(&Model{}).Where("vendor_id = ?", v.Id).Count(&modelCount).Error; err != nil {
+ return err
+ }
+ if modelCount > 0 {
+ return ErrVendorHasModels
+ }
+ return softDeleteMetadataWithNameKey(tx, &Vendor{}, v.Id, retiredMetadataNameKey("vendor", v.Id))
+ })
+}
+
+func (v *Vendor) BeforeCreate(_ *gorm.DB) error {
+ v.NameKey = activeMetadataNameKey(v.Name)
+ return nil
}
// GetVendorByID 根据 ID 获取供应商
diff --git a/relay/channel/api_request.go b/relay/channel/api_request.go
index 53aac675..30f28a8f 100644
--- a/relay/channel/api_request.go
+++ b/relay/channel/api_request.go
@@ -6,9 +6,11 @@ import (
"fmt"
"io"
"net/http"
+ "net/http/httptrace"
"regexp"
"strings"
"sync"
+ "sync/atomic"
"time"
common2 "github.com/MAX-API-Next/MAX-API/common"
@@ -607,6 +609,14 @@ func DoTaskApiRequest(a TaskAdaptor, c *gin.Context, info *common.RelayInfo, req
if err != nil {
return nil, fmt.Errorf("new request failed: %w", err)
}
+ var requestWritten atomic.Bool
+ req = req.WithContext(httptrace.WithClientTrace(req.Context(), &httptrace.ClientTrace{
+ WroteRequest: func(traceInfo httptrace.WroteRequestInfo) {
+ if traceInfo.Err == nil {
+ requestWritten.Store(true)
+ }
+ },
+ }))
err = a.BuildRequestHeader(c, req, info)
if err != nil {
@@ -619,6 +629,9 @@ func DoTaskApiRequest(a TaskAdaptor, c *gin.Context, info *common.RelayInfo, req
applyHeaderOverrideToRequest(req, headerOverride)
resp, err := doRequest(c, req, info)
if err != nil {
+ if info != nil && info.TaskRelayInfo != nil && requestWritten.Load() {
+ info.UpstreamTaskOutcomeUnknown = true
+ }
return nil, fmt.Errorf("do request failed: %w", err)
}
return resp, nil
diff --git a/relay/channel/api_request_test.go b/relay/channel/api_request_test.go
index e5646a7f..a34a0608 100644
--- a/relay/channel/api_request_test.go
+++ b/relay/channel/api_request_test.go
@@ -400,6 +400,41 @@ func TestDoTaskApiRequestAppliesRuntimeHeaderOverride(t *testing.T) {
require.Equal(t, "enabled", gotRuntime)
}
+func TestDoTaskApiRequestMarksWrittenTransportFailureAsOutcomeUnknown(t *testing.T) {
+ gin.SetMode(gin.TestMode)
+ service.InitHttpClient()
+ requestReceived := make(chan struct{})
+ server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ _, _ = io.Copy(io.Discard, r.Body)
+ close(requestReceived)
+ hijacker, ok := w.(http.Hijacker)
+ require.True(t, ok)
+ conn, _, err := hijacker.Hijack()
+ require.NoError(t, err)
+ _ = conn.Close()
+ }))
+ t.Cleanup(server.Close)
+
+ recorder := httptest.NewRecorder()
+ ctx, _ := gin.CreateTestContext(recorder)
+ ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/videos", strings.NewReader(`{"prompt":"test"}`))
+ info := &relaycommon.RelayInfo{
+ ChannelMeta: &relaycommon.ChannelMeta{},
+ TaskRelayInfo: &relaycommon.TaskRelayInfo{},
+ }
+
+ resp, err := DoTaskApiRequest(taskHeaderAdaptor{url: server.URL}, ctx, info, strings.NewReader(`{"prompt":"test"}`))
+
+ require.Error(t, err)
+ require.Nil(t, resp)
+ require.True(t, info.UpstreamTaskOutcomeUnknown)
+ select {
+ case <-requestReceived:
+ default:
+ t.Fatal("upstream did not receive the task request")
+ }
+}
+
func TestSendPingDataReturnsWhenWriterBlocks(t *testing.T) {
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
diff --git a/relay/channel/aws/dto.go b/relay/channel/aws/dto.go
index d66b4483..4d1c754b 100644
--- a/relay/channel/aws/dto.go
+++ b/relay/channel/aws/dto.go
@@ -72,10 +72,10 @@ type NovaRequest struct {
}
type NovaInferenceConfig struct {
- MaxTokens int `json:"maxTokens,omitempty"` // 最大生成的 token 数
- Temperature float64 `json:"temperature,omitempty"` // 随机性 (默认 0.7, 范围 0-1)
- TopP float64 `json:"topP,omitempty"` // nucleus sampling (默认 0.9, 范围 0-1)
- TopK int `json:"topK,omitempty"` // 限制候选 token 数 (默认 50, 范围 0-128)
+ MaxTokens *int `json:"maxTokens,omitempty"` // 最大生成的 token 数
+ Temperature *float64 `json:"temperature,omitempty"` // 随机性 (默认 0.7, 范围 0-1)
+ TopP *float64 `json:"topP,omitempty"` // nucleus sampling (默认 0.9, 范围 0-1)
+ TopK *int `json:"topK,omitempty"` // 限制候选 token 数 (默认 50, 范围 0-128)
StopSequences []string `json:"stopSequences,omitempty"` // 停止生成的序列
}
@@ -95,19 +95,21 @@ func convertToNovaRequest(req *dto.GeneralOpenAIRequest) *NovaRequest {
}
// 设置推理配置
- if (req.MaxTokens != nil && *req.MaxTokens != 0) || (req.Temperature != nil && *req.Temperature != 0) || (req.TopP != nil && *req.TopP != 0) || (req.TopK != nil && *req.TopK != 0) || req.Stop != nil {
+ maxTokens := dto.ResolveMaxTokens(req.MaxTokens, req.MaxCompletionTokens)
+ if maxTokens != nil || req.Temperature != nil || req.TopP != nil || req.TopK != nil || req.Stop != nil {
novaReq.InferenceConfig = &NovaInferenceConfig{}
- if req.MaxTokens != nil && *req.MaxTokens != 0 {
- novaReq.InferenceConfig.MaxTokens = int(*req.MaxTokens)
+ if maxTokens != nil {
+ value := int(*maxTokens)
+ novaReq.InferenceConfig.MaxTokens = &value
}
- if req.Temperature != nil && *req.Temperature != 0 {
- novaReq.InferenceConfig.Temperature = *req.Temperature
+ if req.Temperature != nil {
+ novaReq.InferenceConfig.Temperature = req.Temperature
}
- if req.TopP != nil && *req.TopP != 0 {
- novaReq.InferenceConfig.TopP = *req.TopP
+ if req.TopP != nil {
+ novaReq.InferenceConfig.TopP = req.TopP
}
- if req.TopK != nil && *req.TopK != 0 {
- novaReq.InferenceConfig.TopK = *req.TopK
+ if req.TopK != nil {
+ novaReq.InferenceConfig.TopK = req.TopK
}
if req.Stop != nil {
if stopSequences := parseStopSequences(req.Stop); len(stopSequences) > 0 {
diff --git a/relay/channel/aws/relay-aws.go b/relay/channel/aws/relay-aws.go
index bf631773..1c7de080 100644
--- a/relay/channel/aws/relay-aws.go
+++ b/relay/channel/aws/relay-aws.go
@@ -2,7 +2,6 @@ package aws
import (
"context"
- "encoding/json"
"fmt"
"io"
"net/http"
@@ -40,11 +39,15 @@ func getAwsErrorStatusCode(err error) int {
return http.StatusInternalServerError
}
-func newAwsInvokeContext() (context.Context, context.CancelFunc) {
+func newAwsInvokeContext(c *gin.Context) (context.Context, context.CancelFunc) {
+ base := context.Background()
+ if c != nil && c.Request != nil {
+ base = c.Request.Context()
+ }
if common.RelayTimeout <= 0 {
- return context.Background(), func() {}
+ return context.WithCancel(base)
}
- return context.WithTimeout(context.Background(), time.Duration(common.RelayTimeout)*time.Second)
+ return context.WithTimeout(base, time.Duration(common.RelayTimeout)*time.Second)
}
func newAwsClient(c *gin.Context, info *relaycommon.RelayInfo) (*bedrockruntime.Client, error) {
@@ -223,7 +226,7 @@ func getAwsModelID(requestModel string) string {
func awsHandler(c *gin.Context, info *relaycommon.RelayInfo, a *Adaptor) (*types.MaxAPIError, *dto.Usage) {
- ctx, cancel := newAwsInvokeContext()
+ ctx, cancel := newAwsInvokeContext(c)
defer cancel()
awsResp, err := a.AwsClient.InvokeModel(ctx, a.AwsReq.(*bedrockruntime.InvokeModelInput))
@@ -253,7 +256,7 @@ func awsHandler(c *gin.Context, info *relaycommon.RelayInfo, a *Adaptor) (*types
}
func awsStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, a *Adaptor) (*types.MaxAPIError, *dto.Usage) {
- ctx, cancel := newAwsInvokeContext()
+ ctx, cancel := newAwsInvokeContext(c)
defer cancel()
awsResp, err := a.AwsClient.InvokeModelWithResponseStream(ctx, a.AwsReq.(*bedrockruntime.InvokeModelWithResponseStreamInput))
@@ -296,7 +299,7 @@ func awsStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, a *Adaptor) (
// Nova模型处理函数
func handleNovaRequest(c *gin.Context, info *relaycommon.RelayInfo, a *Adaptor) (*types.MaxAPIError, *dto.Usage) {
- ctx, cancel := newAwsInvokeContext()
+ ctx, cancel := newAwsInvokeContext(c)
defer cancel()
awsResp, err := a.AwsClient.InvokeModel(ctx, a.AwsReq.(*bedrockruntime.InvokeModelInput))
@@ -305,24 +308,9 @@ func handleNovaRequest(c *gin.Context, info *relaycommon.RelayInfo, a *Adaptor)
return types.NewOpenAIError(errors.Wrap(err, "InvokeModel"), types.ErrorCodeAwsInvokeError, statusCode), nil
}
- // 解析Nova响应
- var novaResp struct {
- Output struct {
- Message struct {
- Content []struct {
- Text string `json:"text"`
- } `json:"content"`
- } `json:"message"`
- } `json:"output"`
- Usage struct {
- InputTokens int `json:"inputTokens"`
- OutputTokens int `json:"outputTokens"`
- TotalTokens int `json:"totalTokens"`
- } `json:"usage"`
- }
-
- if err := json.Unmarshal(awsResp.Body, &novaResp); err != nil {
- return types.NewError(errors.Wrap(err, "unmarshal nova response"), types.ErrorCodeBadResponseBody), nil
+ content, usage, err := parseNovaResponse(awsResp.Body)
+ if err != nil {
+ return types.NewError(err, types.ErrorCodeBadResponseBody), nil
}
// 构造OpenAI格式响应
@@ -335,17 +323,43 @@ func handleNovaRequest(c *gin.Context, info *relaycommon.RelayInfo, a *Adaptor)
Index: 0,
Message: dto.Message{
Role: "assistant",
- Content: novaResp.Output.Message.Content[0].Text,
+ Content: content,
},
FinishReason: "stop",
}},
- Usage: dto.Usage{
- PromptTokens: novaResp.Usage.InputTokens,
- CompletionTokens: novaResp.Usage.OutputTokens,
- TotalTokens: novaResp.Usage.TotalTokens,
- },
+ Usage: usage,
}
c.JSON(http.StatusOK, response)
return nil, &response.Usage
}
+
+func parseNovaResponse(body []byte) (string, dto.Usage, error) {
+ var novaResp struct {
+ Output struct {
+ Message struct {
+ Content []struct {
+ Text string `json:"text"`
+ } `json:"content"`
+ } `json:"message"`
+ } `json:"output"`
+ Usage struct {
+ InputTokens int `json:"inputTokens"`
+ OutputTokens int `json:"outputTokens"`
+ TotalTokens int `json:"totalTokens"`
+ } `json:"usage"`
+ }
+
+ if err := common.Unmarshal(body, &novaResp); err != nil {
+ return "", dto.Usage{}, errors.Wrap(err, "unmarshal nova response")
+ }
+ if len(novaResp.Output.Message.Content) == 0 {
+ return "", dto.Usage{}, errors.New("nova response content is empty")
+ }
+
+ return novaResp.Output.Message.Content[0].Text, dto.Usage{
+ PromptTokens: novaResp.Usage.InputTokens,
+ CompletionTokens: novaResp.Usage.OutputTokens,
+ TotalTokens: novaResp.Usage.TotalTokens,
+ }, nil
+}
diff --git a/relay/channel/aws/relay_aws_test.go b/relay/channel/aws/relay_aws_test.go
index 66e0ede8..b03009d7 100644
--- a/relay/channel/aws/relay_aws_test.go
+++ b/relay/channel/aws/relay_aws_test.go
@@ -2,17 +2,39 @@ package aws
import (
"bytes"
+ "context"
"net/http"
"net/http/httptest"
"testing"
+ "time"
"github.com/MAX-API-Next/MAX-API/common"
+ "github.com/MAX-API-Next/MAX-API/dto"
relaycommon "github.com/MAX-API-Next/MAX-API/relay/common"
"github.com/aws/aws-sdk-go-v2/service/bedrockruntime"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
+func TestNewAwsInvokeContextFollowsRequestCancellation(t *testing.T) {
+ gin.SetMode(gin.TestMode)
+ recorder := httptest.NewRecorder()
+ c, _ := gin.CreateTestContext(recorder)
+ requestContext, cancelRequest := context.WithCancel(context.Background())
+ c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil).WithContext(requestContext)
+
+ ctx, cancel := newAwsInvokeContext(c)
+ defer cancel()
+ cancelRequest()
+
+ select {
+ case <-ctx.Done():
+ require.ErrorIs(t, ctx.Err(), context.Canceled)
+ case <-time.After(time.Second):
+ t.Fatal("AWS invoke context did not follow client cancellation")
+ }
+}
+
func TestDoAwsClientRequest_AppliesRuntimeHeaderOverrideToAnthropicBeta(t *testing.T) {
t.Parallel()
@@ -53,3 +75,29 @@ func TestDoAwsClientRequest_AppliesRuntimeHeaderOverrideToAnthropicBeta(t *testi
require.True(t, ok)
require.Equal(t, []any{"computer-use-2025-01-24"}, values)
}
+
+func TestParseNovaResponseRejectsEmptyContent(t *testing.T) {
+ content, usage, err := parseNovaResponse([]byte(
+ `{"output":{"message":{"content":[]}},"usage":{"inputTokens":1,"outputTokens":0,"totalTokens":1}}`,
+ ))
+
+ require.ErrorContains(t, err, "nova response content is empty")
+ require.Empty(t, content)
+ require.Zero(t, usage)
+}
+
+func TestConvertToNovaRequestPreservesExplicitZeroValues(t *testing.T) {
+ zeroFloat := 0.0
+ zeroInt := 0
+ zeroUint := uint(0)
+ novaRequest := convertToNovaRequest(&dto.GeneralOpenAIRequest{
+ MaxTokens: &zeroUint,
+ Temperature: &zeroFloat,
+ TopP: &zeroFloat,
+ TopK: &zeroInt,
+ })
+
+ payload, err := common.Marshal(novaRequest)
+ require.NoError(t, err)
+ require.JSONEq(t, `{"schemaVersion":"messages-v1","messages":[],"inferenceConfig":{"maxTokens":0,"temperature":0,"topP":0,"topK":0}}`, string(payload))
+}
diff --git a/relay/channel/baidu/dto.go b/relay/channel/baidu/dto.go
index 3e9d6244..bf37efd8 100644
--- a/relay/channel/baidu/dto.go
+++ b/relay/channel/baidu/dto.go
@@ -15,9 +15,9 @@ type BaiduMessage struct {
type BaiduChatRequest struct {
Messages []BaiduMessage `json:"messages"`
Temperature *float64 `json:"temperature,omitempty"`
- TopP float64 `json:"top_p,omitempty"`
- PenaltyScore float64 `json:"penalty_score,omitempty"`
- Stream bool `json:"stream,omitempty"`
+ TopP *float64 `json:"top_p,omitempty"`
+ PenaltyScore *float64 `json:"penalty_score,omitempty"`
+ Stream *bool `json:"stream,omitempty"`
System string `json:"system,omitempty"`
DisableSearch bool `json:"disable_search,omitempty"`
EnableCitation bool `json:"enable_citation,omitempty"`
diff --git a/relay/channel/baidu/relay-baidu.go b/relay/channel/baidu/relay-baidu.go
index 26e947cf..1f27fac0 100644
--- a/relay/channel/baidu/relay-baidu.go
+++ b/relay/channel/baidu/relay-baidu.go
@@ -17,8 +17,6 @@ import (
"github.com/MAX-API-Next/MAX-API/relay/helper"
"github.com/MAX-API-Next/MAX-API/service"
"github.com/MAX-API-Next/MAX-API/types"
- "github.com/samber/lo"
-
"github.com/gin-gonic/gin"
)
@@ -29,9 +27,9 @@ var baiduTokenStore sync.Map
func requestOpenAI2Baidu(request dto.GeneralOpenAIRequest) *BaiduChatRequest {
baiduRequest := BaiduChatRequest{
Temperature: request.Temperature,
- TopP: lo.FromPtrOr(request.TopP, 0),
- PenaltyScore: lo.FromPtrOr(request.FrequencyPenalty, 0),
- Stream: lo.FromPtrOr(request.Stream, false),
+ TopP: request.TopP,
+ PenaltyScore: request.FrequencyPenalty,
+ Stream: request.Stream,
DisableSearch: false,
EnableCitation: false,
UserId: request.User,
@@ -160,6 +158,9 @@ func baiduHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Respon
c.Writer.Header().Set("Content-Type", "application/json")
c.Writer.WriteHeader(resp.StatusCode)
_, err = c.Writer.Write(jsonResponse)
+ if err != nil {
+ common.SysLog("failed to write Baidu response: " + err.Error())
+ }
return nil, &fullTextResponse.Usage
}
@@ -185,6 +186,9 @@ func baiduEmbeddingHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *ht
c.Writer.Header().Set("Content-Type", "application/json")
c.Writer.WriteHeader(resp.StatusCode)
_, err = c.Writer.Write(jsonResponse)
+ if err != nil {
+ common.SysLog("failed to write Baidu embedding response: " + err.Error())
+ }
return nil, &fullTextResponse.Usage
}
@@ -229,8 +233,25 @@ func getBaiduAccessTokenHelper(apiKey string) (*BaiduAccessToken, error) {
}
defer res.Body.Close()
+ accessToken, err := parseBaiduAccessTokenResponse(res)
+ if err != nil {
+ return nil, err
+ }
+ accessToken.ExpiresAt = time.Now().Add(time.Duration(accessToken.ExpiresIn) * time.Second)
+ baiduTokenStore.Store(apiKey, *accessToken)
+ return accessToken, nil
+}
+
+func parseBaiduAccessTokenResponse(res *http.Response) (*BaiduAccessToken, error) {
+ if res == nil || res.Body == nil {
+ return nil, errors.New("baidu access token response is empty")
+ }
+ if res.StatusCode < http.StatusOK || res.StatusCode >= http.StatusMultipleChoices {
+ return nil, fmt.Errorf("baidu access token endpoint returned HTTP %d", res.StatusCode)
+ }
+
var accessToken BaiduAccessToken
- err = json.NewDecoder(res.Body).Decode(&accessToken)
+ err := common.DecodeJson(res.Body, &accessToken)
if err != nil {
return nil, err
}
@@ -240,7 +261,5 @@ func getBaiduAccessTokenHelper(apiKey string) (*BaiduAccessToken, error) {
if accessToken.AccessToken == "" {
return nil, errors.New("getBaiduAccessTokenHelper get empty access token")
}
- accessToken.ExpiresAt = time.Now().Add(time.Duration(accessToken.ExpiresIn) * time.Second)
- baiduTokenStore.Store(apiKey, accessToken)
return &accessToken, nil
}
diff --git a/relay/channel/baidu/relay_baidu_test.go b/relay/channel/baidu/relay_baidu_test.go
new file mode 100644
index 00000000..4a6391bf
--- /dev/null
+++ b/relay/channel/baidu/relay_baidu_test.go
@@ -0,0 +1,47 @@
+package baidu
+
+import (
+ "io"
+ "net/http"
+ "strings"
+ "testing"
+
+ "github.com/MAX-API-Next/MAX-API/common"
+ "github.com/MAX-API-Next/MAX-API/dto"
+ "github.com/stretchr/testify/require"
+)
+
+func TestParseBaiduAccessTokenResponseRejectsNonSuccessStatus(t *testing.T) {
+ accessToken, err := parseBaiduAccessTokenResponse(&http.Response{
+ StatusCode: http.StatusUnauthorized,
+ Body: io.NopCloser(strings.NewReader(`{"error":"invalid_client"}`)),
+ })
+
+ require.Nil(t, accessToken)
+ require.EqualError(t, err, "baidu access token endpoint returned HTTP 401")
+}
+
+func TestParseBaiduAccessTokenResponseAcceptsSuccessStatus(t *testing.T) {
+ accessToken, err := parseBaiduAccessTokenResponse(&http.Response{
+ StatusCode: http.StatusOK,
+ Body: io.NopCloser(strings.NewReader(`{"access_token":"token","expires_in":3600}`)),
+ })
+
+ require.NoError(t, err)
+ require.Equal(t, "token", accessToken.AccessToken)
+ require.Equal(t, int64(3600), accessToken.ExpiresIn)
+}
+
+func TestRequestOpenAI2BaiduPreservesExplicitZeroValues(t *testing.T) {
+ zero := 0.0
+ stream := false
+ request := requestOpenAI2Baidu(dto.GeneralOpenAIRequest{
+ TopP: &zero,
+ FrequencyPenalty: &zero,
+ Stream: &stream,
+ })
+
+ payload, err := common.Marshal(request)
+ require.NoError(t, err)
+ require.JSONEq(t, `{"messages":null,"top_p":0,"penalty_score":0,"stream":false}`, string(payload))
+}
diff --git a/relay/channel/cohere/relay-cohere.go b/relay/channel/cohere/relay-cohere.go
index e7b8e95b..63701019 100644
--- a/relay/channel/cohere/relay-cohere.go
+++ b/relay/channel/cohere/relay-cohere.go
@@ -99,16 +99,33 @@ func cohereStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http
return 0, nil, nil
})
dataChan := make(chan string)
- stopChan := make(chan bool)
+ stopChan := make(chan struct{}, 1)
+ done := make(chan struct{})
+ producerDone := make(chan struct{})
go func() {
+ defer close(producerDone)
+ defer func() {
+ select {
+ case stopChan <- struct{}{}:
+ case <-done:
+ }
+ }()
for scanner.Scan() {
data := scanner.Text()
- dataChan <- data
+ select {
+ case dataChan <- data:
+ case <-done:
+ return
+ }
}
if err := scanner.Err(); err != nil {
+ select {
+ case <-done:
+ return
+ default:
+ }
common.SysLog("error reading stream: " + err.Error())
}
- stopChan <- true
}()
helper.SetEventStreamHeaders(c)
isFirst := true
@@ -166,8 +183,13 @@ func cohereStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http
case <-stopChan:
c.Render(-1, common.CustomEvent{Data: "data: [DONE]"})
return false
+ case <-c.Request.Context().Done():
+ return false
}
})
+ close(done)
+ service.CloseResponseBodyGracefully(resp)
+ <-producerDone
if usage.PromptTokens == 0 {
usage = service.ResponseText2Usage(c, responseText, info.UpstreamModelName, info.GetEstimatePromptTokens())
}
@@ -212,7 +234,9 @@ func cohereHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Respo
}
c.Writer.Header().Set("Content-Type", "application/json")
c.Writer.WriteHeader(resp.StatusCode)
- _, _ = c.Writer.Write(jsonResponse)
+ if _, writeErr := c.Writer.Write(jsonResponse); writeErr != nil {
+ common.SysLog("failed to write Cohere response: " + writeErr.Error())
+ }
return &usage, nil
}
@@ -249,5 +273,8 @@ func cohereRerankHandler(c *gin.Context, resp *http.Response, info *relaycommon.
c.Writer.Header().Set("Content-Type", "application/json")
c.Writer.WriteHeader(resp.StatusCode)
_, err = c.Writer.Write(jsonResponse)
+ if err != nil {
+ common.SysLog("failed to write Cohere rerank response: " + err.Error())
+ }
return &usage, nil
}
diff --git a/relay/channel/cohere/stream_cancel_test.go b/relay/channel/cohere/stream_cancel_test.go
new file mode 100644
index 00000000..11aad76a
--- /dev/null
+++ b/relay/channel/cohere/stream_cancel_test.go
@@ -0,0 +1,63 @@
+package cohere
+
+import (
+ "context"
+ "io"
+ "net/http"
+ "net/http/httptest"
+ "testing"
+ "time"
+
+ relaycommon "github.com/MAX-API-Next/MAX-API/relay/common"
+ "github.com/gin-gonic/gin"
+ "github.com/stretchr/testify/require"
+)
+
+type closeNotifyRecorder struct {
+ *httptest.ResponseRecorder
+ closeNotify chan bool
+}
+
+func (r *closeNotifyRecorder) CloseNotify() <-chan bool {
+ return r.closeNotify
+}
+
+func TestCohereStreamHandlerStopsAfterRequestCancellation(t *testing.T) {
+ gin.SetMode(gin.TestMode)
+ c, _ := gin.CreateTestContext(&closeNotifyRecorder{
+ ResponseRecorder: httptest.NewRecorder(),
+ closeNotify: make(chan bool),
+ })
+ request := httptest.NewRequest(http.MethodGet, "/", nil)
+ ctx, cancel := context.WithCancel(request.Context())
+ c.Request = request.WithContext(ctx)
+ reader, writer := io.Pipe()
+ t.Cleanup(func() { _ = writer.Close() })
+
+ returned := make(chan struct{})
+ go func() {
+ _, _ = cohereStreamHandler(c, &relaycommon.RelayInfo{
+ ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "cohere-test"},
+ }, &http.Response{Body: reader})
+ close(returned)
+ }()
+ cancel()
+
+ select {
+ case <-returned:
+ case <-time.After(time.Second):
+ t.Fatal("cohere stream handler did not stop after client cancellation")
+ }
+
+ writeResult := make(chan error, 1)
+ go func() {
+ _, err := writer.Write([]byte("event\n"))
+ writeResult <- err
+ }()
+ select {
+ case err := <-writeResult:
+ require.ErrorIs(t, err, io.ErrClosedPipe)
+ case <-time.After(time.Second):
+ t.Fatal("cohere stream response body was not closed")
+ }
+}
diff --git a/relay/channel/gemini/relay-gemini.go b/relay/channel/gemini/relay-gemini.go
index 0bde747b..c47656bb 100644
--- a/relay/channel/gemini/relay-gemini.go
+++ b/relay/channel/gemini/relay-gemini.go
@@ -1643,11 +1643,11 @@ func GeminiChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *
}
func GeminiChatHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.MaxAPIError) {
+ defer service.CloseResponseBodyGracefully(resp)
responseBody, err := io.ReadAll(resp.Body)
if err != nil {
return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
}
- service.CloseResponseBodyGracefully(resp)
logger.LogDebug(c, "Gemini response body: %s", responseBody)
var geminiResponse dto.GeminiChatResponse
err = common.Unmarshal(responseBody, &geminiResponse)
@@ -1766,12 +1766,11 @@ func GeminiEmbeddingHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *h
}
func GeminiImageHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.MaxAPIError) {
+ defer service.CloseResponseBodyGracefully(resp)
responseBody, readErr := io.ReadAll(resp.Body)
if readErr != nil {
return nil, types.NewOpenAIError(readErr, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
}
- _ = resp.Body.Close()
-
var geminiResponse dto.GeminiImageResponse
if jsonErr := common.Unmarshal(responseBody, &geminiResponse); jsonErr != nil {
return nil, types.NewOpenAIError(jsonErr, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
diff --git a/relay/channel/ollama/dto.go b/relay/channel/ollama/dto.go
index 07aeb17a..24d2d006 100644
--- a/relay/channel/ollama/dto.go
+++ b/relay/channel/ollama/dto.go
@@ -36,7 +36,7 @@ type OllamaChatRequest struct {
Messages []OllamaChatMessage `json:"messages"`
Tools interface{} `json:"tools,omitempty"`
Format interface{} `json:"format,omitempty"`
- Stream bool `json:"stream,omitempty"`
+ Stream *bool `json:"stream,omitempty"`
Options map[string]any `json:"options,omitempty"`
KeepAlive interface{} `json:"keep_alive,omitempty"`
Think json.RawMessage `json:"think,omitempty"`
@@ -48,7 +48,7 @@ type OllamaGenerateRequest struct {
Suffix string `json:"suffix,omitempty"`
Images []string `json:"images,omitempty"`
Format interface{} `json:"format,omitempty"`
- Stream bool `json:"stream,omitempty"`
+ Stream *bool `json:"stream,omitempty"`
Options map[string]any `json:"options,omitempty"`
KeepAlive interface{} `json:"keep_alive,omitempty"`
Think json.RawMessage `json:"think,omitempty"`
diff --git a/relay/channel/ollama/relay-ollama.go b/relay/channel/ollama/relay-ollama.go
index 51f8dd44..8722b05b 100644
--- a/relay/channel/ollama/relay-ollama.go
+++ b/relay/channel/ollama/relay-ollama.go
@@ -22,7 +22,7 @@ import (
func openAIChatToOllamaChat(c *gin.Context, r *dto.GeneralOpenAIRequest) (*OllamaChatRequest, error) {
chatReq := &OllamaChatRequest{
Model: r.Model,
- Stream: lo.FromPtrOr(r.Stream, false),
+ Stream: r.Stream,
Options: map[string]any{},
Think: r.Think,
}
@@ -149,7 +149,7 @@ func openAIChatToOllamaChat(c *gin.Context, r *dto.GeneralOpenAIRequest) (*Ollam
func openAIToGenerate(c *gin.Context, r *dto.GeneralOpenAIRequest) (*OllamaGenerateRequest, error) {
gen := &OllamaGenerateRequest{
Model: r.Model,
- Stream: lo.FromPtrOr(r.Stream, false),
+ Stream: r.Stream,
Options: map[string]any{},
Think: r.Think,
}
diff --git a/relay/channel/ollama/stream_test.go b/relay/channel/ollama/stream_test.go
index 449eb332..42cfc4b6 100644
--- a/relay/channel/ollama/stream_test.go
+++ b/relay/channel/ollama/stream_test.go
@@ -94,6 +94,43 @@ func TestOllamaChatHandlerNonStreamToolCalls(t *testing.T) {
}
}
+func TestOllamaRequestConvertersPreserveStreamFalse(t *testing.T) {
+ gin.SetMode(gin.TestMode)
+ c, _ := gin.CreateTestContext(httptest.NewRecorder())
+ stream := false
+ request := &dto.GeneralOpenAIRequest{Model: "llama-test", Stream: &stream}
+
+ chatRequest, err := openAIChatToOllamaChat(c, request)
+ require.NoError(t, err)
+ chatPayload, err := common.Marshal(chatRequest)
+ require.NoError(t, err)
+ require.JSONEq(t, `{"model":"llama-test","messages":[],"stream":false}`, string(chatPayload))
+
+ generateRequest, err := openAIToGenerate(c, request)
+ require.NoError(t, err)
+ generatePayload, err := common.Marshal(generateRequest)
+ require.NoError(t, err)
+ require.JSONEq(t, `{"model":"llama-test","stream":false}`, string(generatePayload))
+}
+
+func TestOllamaRequestConvertersOmitUnsetStream(t *testing.T) {
+ gin.SetMode(gin.TestMode)
+ c, _ := gin.CreateTestContext(httptest.NewRecorder())
+ request := &dto.GeneralOpenAIRequest{Model: "llama-test"}
+
+ chatRequest, err := openAIChatToOllamaChat(c, request)
+ require.NoError(t, err)
+ chatPayload, err := common.Marshal(chatRequest)
+ require.NoError(t, err)
+ require.JSONEq(t, `{"model":"llama-test","messages":[]}`, string(chatPayload))
+
+ generateRequest, err := openAIToGenerate(c, request)
+ require.NoError(t, err)
+ generatePayload, err := common.Marshal(generateRequest)
+ require.NoError(t, err)
+ require.JSONEq(t, `{"model":"llama-test"}`, string(generatePayload))
+}
+
func TestOllamaStreamHandlerToolCalls(t *testing.T) {
gin.SetMode(gin.TestMode)
diff --git a/relay/channel/openai/adaptor.go b/relay/channel/openai/adaptor.go
index b63bbf70..bc9cf272 100644
--- a/relay/channel/openai/adaptor.go
+++ b/relay/channel/openai/adaptor.go
@@ -8,9 +8,7 @@ import (
"io"
"mime/multipart"
"net/http"
- "net/textproto"
"net/url"
- "path/filepath"
"strings"
"github.com/MAX-API-Next/MAX-API/common"
@@ -27,6 +25,7 @@ import (
relaycommon "github.com/MAX-API-Next/MAX-API/relay/common"
"github.com/MAX-API-Next/MAX-API/relay/common_handler"
relayconstant "github.com/MAX-API-Next/MAX-API/relay/constant"
+ "github.com/MAX-API-Next/MAX-API/relay/helper"
"github.com/MAX-API-Next/MAX-API/service"
"github.com/MAX-API-Next/MAX-API/setting/model_setting"
"github.com/MAX-API-Next/MAX-API/setting/reasoning"
@@ -194,19 +193,24 @@ func (a *Adaptor) SetupRequestHeader(c *gin.Context, header *http.Header, info *
}
}
if info.RelayMode == relayconstant.RelayModeRealtime {
+ legacyRealtimeBeta := strings.Contains(info.UpstreamModelName, "-realtime-preview")
swp := c.Request.Header.Get("Sec-WebSocket-Protocol")
if swp != "" {
items := []string{
"realtime",
"openai-insecure-api-key." + info.ApiKey,
- "openai-beta.realtime-v1",
+ }
+ if legacyRealtimeBeta {
+ items = append(items, "openai-beta.realtime-v1")
}
header.Set("Sec-WebSocket-Protocol", strings.Join(items, ","))
//req.Header.Set("Sec-WebSocket-Key", c.Request.Header.Get("Sec-WebSocket-Key"))
//req.Header.Set("Sec-Websocket-Extensions", c.Request.Header.Get("Sec-Websocket-Extensions"))
//req.Header.Set("Sec-Websocket-Version", c.Request.Header.Get("Sec-Websocket-Version"))
} else {
- header.Set("openai-beta", "realtime=v1")
+ if legacyRealtimeBeta {
+ header.Set("openai-beta", "realtime=v1")
+ }
if !hasAuthOverride {
header.Set("Authorization", "Bearer "+info.ApiKey)
}
@@ -488,9 +492,9 @@ func (a *Adaptor) ConvertImageRequest(c *gin.Context, info *relaycommon.RelayInf
// Process all image files
for i, fileHeader := range imageFiles {
- file, err := fileHeader.Open()
+ file, mimeType, err := openValidatedImageFile(fileHeader)
if err != nil {
- return nil, fmt.Errorf("failed to open image file %d: %w", i, err)
+ return nil, fmt.Errorf("invalid image file %d: %w", i, err)
}
// If multiple images, use image[] as the field name
@@ -499,20 +503,14 @@ func (a *Adaptor) ConvertImageRequest(c *gin.Context, info *relaycommon.RelayInf
fieldName = "image[]"
}
- // Determine MIME type based on file extension
- mimeType := detectImageMimeType(fileHeader.Filename)
-
- // Create a form file with the appropriate content type
- h := make(textproto.MIMEHeader)
- h.Set("Content-Disposition", fmt.Sprintf(`form-data; name="%s"; filename="%s"`, fieldName, fileHeader.Filename))
- h.Set("Content-Type", mimeType)
-
- part, err := writer.CreatePart(h)
+ part, err := helper.CreateFormFileWithContentType(writer, fieldName, fileHeader.Filename, mimeType)
if err != nil {
+ _ = file.Close()
return nil, fmt.Errorf("create form part failed for image %d: %w", i, err)
}
if _, err := io.Copy(part, file); err != nil {
+ _ = file.Close()
return nil, fmt.Errorf("copy file failed for image %d: %w", i, err)
}
@@ -522,27 +520,21 @@ func (a *Adaptor) ConvertImageRequest(c *gin.Context, info *relaycommon.RelayInf
// Handle mask file if present
if maskFiles, exists := mf.File["mask"]; exists && len(maskFiles) > 0 {
- maskFile, err := maskFiles[0].Open()
+ maskFile, mimeType, err := openValidatedImageFile(maskFiles[0])
if err != nil {
- return nil, errors.New("failed to open mask file")
+ return nil, fmt.Errorf("invalid mask file: %w", err)
}
// 复制完立即关闭，避免在循环内使用 defer 占用资源
- // Determine MIME type for mask file
- mimeType := detectImageMimeType(maskFiles[0].Filename)
-
- // Create a form file with the appropriate content type
- h := make(textproto.MIMEHeader)
- h.Set("Content-Disposition", fmt.Sprintf(`form-data; name="mask"; filename="%s"`, maskFiles[0].Filename))
- h.Set("Content-Type", mimeType)
-
- maskPart, err := writer.CreatePart(h)
+ maskPart, err := helper.CreateFormFileWithContentType(writer, "mask", maskFiles[0].Filename, mimeType)
if err != nil {
- return nil, errors.New("create form file failed for mask")
+ _ = maskFile.Close()
+ return nil, fmt.Errorf("create form part failed for mask: %w", err)
}
if _, err := io.Copy(maskPart, maskFile); err != nil {
- return nil, errors.New("copy mask file failed")
+ _ = maskFile.Close()
+ return nil, fmt.Errorf("copy mask file failed: %w", err)
}
_ = maskFile.Close()
}
@@ -567,24 +559,34 @@ func isJSONRequest(c *gin.Context) bool {
return strings.HasPrefix(c.Request.Header.Get("Content-Type"), "application/json")
}
-// detectImageMimeType determines the MIME type based on the file extension
-func detectImageMimeType(filename string) string {
- ext := strings.ToLower(filepath.Ext(filename))
- switch ext {
- case ".jpg", ".jpeg":
- return "image/jpeg"
- case ".png":
- return "image/png"
- case ".webp":
- return "image/webp"
- default:
- // Try to detect from extension if possible
- if strings.HasPrefix(ext, ".jp") {
- return "image/jpeg"
- }
- // Default to png as a fallback
- return "image/png"
+func openValidatedImageFile(fileHeader *multipart.FileHeader) (multipart.File, string, error) {
+ if fileHeader == nil {
+ return nil, "", errors.New("image file is nil")
+ }
+ file, err := fileHeader.Open()
+ if err != nil {
+ return nil, "", err
+ }
+ header := make([]byte, 512)
+ n, readErr := io.ReadFull(file, header)
+ if readErr != nil && readErr != io.ErrUnexpectedEOF && readErr != io.EOF {
+ _ = file.Close()
+ return nil, "", fmt.Errorf("read image header: %w", readErr)
+ }
+ if n == 0 {
+ _ = file.Close()
+ return nil, "", errors.New("image file is empty")
+ }
+ mimeType := http.DetectContentType(header[:n])
+ if mimeType != "image/jpeg" && mimeType != "image/png" && mimeType != "image/webp" {
+ _ = file.Close()
+ return nil, "", fmt.Errorf("unsupported image content type %s", mimeType)
+ }
+ if _, err := file.Seek(0, io.SeekStart); err != nil {
+ _ = file.Close()
+ return nil, "", fmt.Errorf("rewind image file: %w", err)
}
+ return file, mimeType, nil
}
func (a *Adaptor) ConvertOpenAIResponsesRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.OpenAIResponsesRequest) (any, error) {
diff --git a/relay/channel/openai/adaptor_realtime_test.go b/relay/channel/openai/adaptor_realtime_test.go
new file mode 100644
index 00000000..75fe5e08
--- /dev/null
+++ b/relay/channel/openai/adaptor_realtime_test.go
@@ -0,0 +1,56 @@
+package openai
+
+import (
+ "net/http"
+ "net/http/httptest"
+ "strings"
+ "testing"
+
+ "github.com/MAX-API-Next/MAX-API/constant"
+ relaycommon "github.com/MAX-API-Next/MAX-API/relay/common"
+ relayconstant "github.com/MAX-API-Next/MAX-API/relay/constant"
+ "github.com/gin-gonic/gin"
+ "github.com/stretchr/testify/assert"
+ "github.com/stretchr/testify/require"
+)
+
+func TestRealtimeBetaHeaderOnlyForPreviewModels(t *testing.T) {
+ gin.SetMode(gin.TestMode)
+ tests := []struct {
+ name string
+ model string
+ wantBeta bool
+ websocket bool
+ }{
+ {name: "ga http", model: "gpt-realtime", wantBeta: false},
+ {name: "preview http", model: "gpt-4o-realtime-preview", wantBeta: true},
+ {name: "ga websocket", model: "gpt-realtime", wantBeta: false, websocket: true},
+ {name: "preview websocket", model: "gpt-4o-realtime-preview", wantBeta: true, websocket: true},
+ }
+
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ c, _ := gin.CreateTestContext(httptest.NewRecorder())
+ c.Request = httptest.NewRequest(http.MethodPost, "/v1/realtime", nil)
+ if tt.websocket {
+ c.Request.Header.Set("Sec-WebSocket-Protocol", "realtime")
+ }
+ header := http.Header{}
+ info := &relaycommon.RelayInfo{
+ RelayMode: relayconstant.RelayModeRealtime,
+ ChannelMeta: &relaycommon.ChannelMeta{
+ ChannelType: constant.ChannelTypeOpenAI,
+ ApiKey: "test-key",
+ UpstreamModelName: tt.model,
+ },
+ }
+ require.NoError(t, (&Adaptor{}).SetupRequestHeader(c, &header, info))
+ if tt.websocket {
+ hasBeta := strings.Contains(header.Get("Sec-WebSocket-Protocol"), "openai-beta.realtime-v1")
+ assert.Equal(t, tt.wantBeta, hasBeta)
+ } else {
+ assert.Equal(t, tt.wantBeta, header.Get("openai-beta") != "")
+ }
+ })
+ }
+}
diff --git a/relay/channel/openai/image_edit_test.go b/relay/channel/openai/image_edit_test.go
index 73d7a5a1..049b8e66 100644
--- a/relay/channel/openai/image_edit_test.go
+++ b/relay/channel/openai/image_edit_test.go
@@ -18,6 +18,10 @@ import (
func TestConvertImageEditRequestMultipart(t *testing.T) {
gin.SetMode(gin.TestMode)
+ pngContent := []byte{
+ 0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a,
+ 0x00, 0x00, 0x00, 0x0d, 0x49, 0x48, 0x44, 0x52,
+ }
newMultipartContext := func(t *testing.T, prompt string) *gin.Context {
var body bytes.Buffer
@@ -28,7 +32,7 @@ func TestConvertImageEditRequestMultipart(t *testing.T) {
require.NoError(t, writer.WriteField("partial_images", "3"))
part, err := writer.CreateFormFile("image", "input.png")
require.NoError(t, err)
- _, err = part.Write([]byte("fake image"))
+ _, err = part.Write(pngContent)
require.NoError(t, err)
require.NoError(t, writer.Close())
@@ -68,7 +72,7 @@ func TestConvertImageEditRequestMultipart(t *testing.T) {
defer file.Close()
fileBytes, err := io.ReadAll(file)
require.NoError(t, err)
- require.Equal(t, []byte("fake image"), fileBytes)
+ require.Equal(t, pngContent, fileBytes)
}
t.Run("with pre-parsed form", func(t *testing.T) {
@@ -92,3 +96,22 @@ func TestConvertImageEditRequestMultipart(t *testing.T) {
convertAndReplay(t, c, prompt)
})
}
+
+func TestConvertImageEditRequestRejectsNonImageContent(t *testing.T) {
+ gin.SetMode(gin.TestMode)
+ var body bytes.Buffer
+ writer := multipart.NewWriter(&body)
+ require.NoError(t, writer.WriteField("model", "gpt-image-1"))
+ require.NoError(t, writer.WriteField("prompt", "edit this"))
+ part, err := writer.CreateFormFile("image", "claimed-image.png")
+ require.NoError(t, err)
+ _, err = part.Write([]byte("not an image"))
+ require.NoError(t, err)
+ require.NoError(t, writer.Close())
+
+ c, _ := gin.CreateTestContext(httptest.NewRecorder())
+ c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/edits", &body)
+ c.Request.Header.Set("Content-Type", writer.FormDataContentType())
+ _, err = (&Adaptor{}).ConvertImageRequest(c, &relaycommon.RelayInfo{RelayMode: relayconstant.RelayModeImagesEdits}, dto.ImageRequest{Model: "gpt-image-1", Prompt: "edit this"})
+ require.ErrorContains(t, err, "unsupported image content type")
+}
diff --git a/relay/channel/openai/realtime_connection.go b/relay/channel/openai/realtime_connection.go
new file mode 100644
index 00000000..f8e6392c
--- /dev/null
+++ b/relay/channel/openai/realtime_connection.go
@@ -0,0 +1,41 @@
+package openai
+
+import (
+ "fmt"
+ "time"
+
+ "github.com/gorilla/websocket"
+)
+
+const (
+ realtimePongWait = 90 * time.Second
+ realtimePingPeriod = 75 * time.Second
+ realtimeWriteWait = 10 * time.Second
+)
+
+func configureRealtimeConnection(conn *websocket.Conn, pongWait time.Duration) error {
+ if err := conn.SetReadDeadline(time.Now().Add(pongWait)); err != nil {
+ return err
+ }
+ conn.SetPongHandler(func(string) error {
+ return conn.SetReadDeadline(time.Now().Add(pongWait))
+ })
+ return nil
+}
+
+func runRealtimePingLoop(conn *websocket.Conn, done <-chan struct{}, pingPeriod, writeWait time.Duration, stop func(error)) {
+ ticker := time.NewTicker(pingPeriod)
+ defer ticker.Stop()
+
+ for {
+ select {
+ case <-done:
+ return
+ case <-ticker.C:
+ if err := conn.WriteControl(websocket.PingMessage, nil, time.Now().Add(writeWait)); err != nil {
+ stop(fmt.Errorf("websocket ping failed: %w", err))
+ return
+ }
+ }
+ }
+}
diff --git a/relay/channel/openai/realtime_connection_test.go b/relay/channel/openai/realtime_connection_test.go
new file mode 100644
index 00000000..1730c600
--- /dev/null
+++ b/relay/channel/openai/realtime_connection_test.go
@@ -0,0 +1,49 @@
+package openai
+
+import (
+ "net/http"
+ "net/http/httptest"
+ "strings"
+ "testing"
+ "time"
+
+ "github.com/gorilla/websocket"
+ "github.com/stretchr/testify/require"
+)
+
+func TestConfigureRealtimeConnectionTimesOutSilentPeer(t *testing.T) {
+ readErr := make(chan error, 1)
+ upgrader := websocket.Upgrader{}
+ server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ conn, err := upgrader.Upgrade(w, r, nil)
+ if err != nil {
+ readErr <- err
+ return
+ }
+ defer conn.Close()
+ if err := configureRealtimeConnection(conn, 25*time.Millisecond); err != nil {
+ readErr <- err
+ return
+ }
+ _, _, err = conn.ReadMessage()
+ readErr <- err
+ }))
+ t.Cleanup(server.Close)
+
+ conn, _, err := websocket.DefaultDialer.Dial("ws"+strings.TrimPrefix(server.URL, "http"), nil)
+ require.NoError(t, err)
+ t.Cleanup(func() { _ = conn.Close() })
+
+ select {
+ case err := <-readErr:
+ require.Error(t, err)
+ require.True(t, websocket.IsCloseError(err, websocket.CloseNormalClosure) || isTimeoutError(err), "expected timeout, got %v", err)
+ case <-time.After(time.Second):
+ t.Fatal("silent websocket peer did not time out")
+ }
+}
+
+func isTimeoutError(err error) bool {
+ timeout, ok := err.(interface{ Timeout() bool })
+ return ok && timeout.Timeout()
+}
diff --git a/relay/channel/openai/relay-openai.go b/relay/channel/openai/relay-openai.go
index 4ba3c139..0f19a4c1 100644
--- a/relay/channel/openai/relay-openai.go
+++ b/relay/channel/openai/relay-openai.go
@@ -5,6 +5,7 @@ import (
"io"
"net/http"
"strings"
+ "sync"
"github.com/MAX-API-Next/MAX-API/common"
"github.com/MAX-API-Next/MAX-API/constant"
@@ -340,18 +341,27 @@ func OpenaiRealtimeHandler(c *gin.Context, info *relaycommon.RelayInfo) (*types.
info.IsStream = true
clientConn := info.ClientWs
targetConn := info.TargetWs
+ if err := configureRealtimeConnection(clientConn, realtimePongWait); err != nil {
+ return types.NewError(fmt.Errorf("configure client websocket: %w", err), types.ErrorCodeBadResponse), nil
+ }
+ if err := configureRealtimeConnection(targetConn, realtimePongWait); err != nil {
+ return types.NewError(fmt.Errorf("configure target websocket: %w", err), types.ErrorCodeBadResponse), nil
+ }
clientClosed := make(chan struct{})
targetClosed := make(chan struct{})
- sendChan := make(chan []byte, 100)
- receiveChan := make(chan []byte, 100)
- errChan := make(chan error, 2)
+ errChan := make(chan error, 4)
+ done := make(chan struct{})
+ var workers sync.WaitGroup
+ var stateMu sync.Mutex
usage := &dto.RealtimeUsage{}
localUsage := &dto.RealtimeUsage{}
sumUsage := &dto.RealtimeUsage{}
+ workers.Add(4)
gopool.Go(func() {
+ defer workers.Done()
defer func() {
if r := recover(); r != nil {
errChan <- fmt.Errorf("panic in client reader: %v", r)
@@ -378,6 +388,7 @@ func OpenaiRealtimeHandler(c *gin.Context, info *relaycommon.RelayInfo) (*types.
return
}
+ stateMu.Lock()
if realtimeEvent.Type == dto.RealtimeEventTypeSessionUpdate {
if realtimeEvent.Session != nil {
if realtimeEvent.Session.Tools != nil {
@@ -388,6 +399,7 @@ func OpenaiRealtimeHandler(c *gin.Context, info *relaycommon.RelayInfo) (*types.
textToken, audioToken, err := service.CountTokenRealtime(info, *realtimeEvent, info.UpstreamModelName)
if err != nil {
+ stateMu.Unlock()
errChan <- fmt.Errorf("error counting text token: %v", err)
return
}
@@ -396,6 +408,7 @@ func OpenaiRealtimeHandler(c *gin.Context, info *relaycommon.RelayInfo) (*types.
localUsage.InputTokens += textToken + audioToken
localUsage.InputTokenDetails.TextTokens += textToken
localUsage.InputTokenDetails.AudioTokens += audioToken
+ stateMu.Unlock()
err = helper.WssString(c, targetConn, string(message))
if err != nil {
@@ -403,15 +416,12 @@ func OpenaiRealtimeHandler(c *gin.Context, info *relaycommon.RelayInfo) (*types.
return
}
- select {
- case sendChan <- message:
- default:
- }
}
}
})
gopool.Go(func() {
+ defer workers.Done()
defer func() {
if r := recover(); r != nil {
errChan <- fmt.Errorf("panic in target reader: %v", r)
@@ -438,6 +448,7 @@ func OpenaiRealtimeHandler(c *gin.Context, info *relaycommon.RelayInfo) (*types.
return
}
+ stateMu.Lock()
if realtimeEvent.Type == dto.RealtimeEventTypeResponseDone {
realtimeUsage := realtimeEvent.Response.Usage
if realtimeUsage != nil {
@@ -451,6 +462,9 @@ func OpenaiRealtimeHandler(c *gin.Context, info *relaycommon.RelayInfo) (*types.
usage.OutputTokenDetails.TextTokens += realtimeUsage.OutputTokenDetails.TextTokens
err := preConsumeUsage(c, info, usage, sumUsage)
if err != nil {
+ usage = &dto.RealtimeUsage{}
+ localUsage = &dto.RealtimeUsage{}
+ stateMu.Unlock()
errChan <- fmt.Errorf("error consume usage: %v", err)
return
}
@@ -461,6 +475,7 @@ func OpenaiRealtimeHandler(c *gin.Context, info *relaycommon.RelayInfo) (*types.
} else {
textToken, audioToken, err := service.CountTokenRealtime(info, *realtimeEvent, info.UpstreamModelName)
if err != nil {
+ stateMu.Unlock()
errChan <- fmt.Errorf("error counting text token: %v", err)
return
}
@@ -472,6 +487,8 @@ func OpenaiRealtimeHandler(c *gin.Context, info *relaycommon.RelayInfo) (*types.
localUsage.InputTokenDetails.AudioTokens += audioToken
err = preConsumeUsage(c, info, localUsage, sumUsage)
if err != nil {
+ localUsage = &dto.RealtimeUsage{}
+ stateMu.Unlock()
errChan <- fmt.Errorf("error consume usage: %v", err)
return
}
@@ -493,6 +510,7 @@ func OpenaiRealtimeHandler(c *gin.Context, info *relaycommon.RelayInfo) (*types.
} else {
textToken, audioToken, err := service.CountTokenRealtime(info, *realtimeEvent, info.UpstreamModelName)
if err != nil {
+ stateMu.Unlock()
errChan <- fmt.Errorf("error counting text token: %v", err)
return
}
@@ -502,6 +520,7 @@ func OpenaiRealtimeHandler(c *gin.Context, info *relaycommon.RelayInfo) (*types.
localUsage.OutputTokenDetails.TextTokens += textToken
localUsage.OutputTokenDetails.AudioTokens += audioToken
}
+ stateMu.Unlock()
err = helper.WssString(c, clientConn, string(message))
if err != nil {
@@ -509,14 +528,25 @@ func OpenaiRealtimeHandler(c *gin.Context, info *relaycommon.RelayInfo) (*types.
return
}
- select {
- case receiveChan <- message:
- default:
- }
}
}
})
+ reportPingError := func(err error) {
+ select {
+ case errChan <- err:
+ default:
+ }
+ }
+ gopool.Go(func() {
+ defer workers.Done()
+ runRealtimePingLoop(clientConn, done, realtimePingPeriod, realtimeWriteWait, reportPingError)
+ })
+ gopool.Go(func() {
+ defer workers.Done()
+ runRealtimePingLoop(targetConn, done, realtimePingPeriod, realtimeWriteWait, reportPingError)
+ })
+
select {
case <-clientClosed:
case <-targetClosed:
@@ -525,7 +555,13 @@ func OpenaiRealtimeHandler(c *gin.Context, info *relaycommon.RelayInfo) (*types.
logger.LogError(c, "realtime error: "+err.Error())
case <-c.Done():
}
+ close(done)
+ _ = clientConn.Close()
+ _ = targetConn.Close()
+ workers.Wait()
+ stateMu.Lock()
+ defer stateMu.Unlock()
if usage.TotalTokens != 0 {
_ = preConsumeUsage(c, info, usage, sumUsage)
}
@@ -544,17 +580,17 @@ func preConsumeUsage(ctx *gin.Context, info *relaycommon.RelayInfo, usage *dto.R
return fmt.Errorf("invalid usage pointer")
}
- totalUsage.TotalTokens += usage.TotalTokens
- totalUsage.InputTokens += usage.InputTokens
- totalUsage.OutputTokens += usage.OutputTokens
- totalUsage.InputTokenDetails.CachedTokens += usage.InputTokenDetails.CachedTokens
- totalUsage.InputTokenDetails.TextTokens += usage.InputTokenDetails.TextTokens
- totalUsage.InputTokenDetails.AudioTokens += usage.InputTokenDetails.AudioTokens
- totalUsage.OutputTokenDetails.TextTokens += usage.OutputTokenDetails.TextTokens
- totalUsage.OutputTokenDetails.AudioTokens += usage.OutputTokenDetails.AudioTokens
- // clear usage
- err := service.PreWssConsumeQuota(ctx, info, usage)
- return err
+ nextUsage := *totalUsage
+ nextUsage.TotalTokens += usage.TotalTokens
+ nextUsage.InputTokens += usage.InputTokens
+ nextUsage.OutputTokens += usage.OutputTokens
+ nextUsage.InputTokenDetails.CachedTokens += usage.InputTokenDetails.CachedTokens
+ nextUsage.InputTokenDetails.TextTokens += usage.InputTokenDetails.TextTokens
+ nextUsage.InputTokenDetails.AudioTokens += usage.InputTokenDetails.AudioTokens
+ nextUsage.OutputTokenDetails.TextTokens += usage.OutputTokenDetails.TextTokens
+ nextUsage.OutputTokenDetails.AudioTokens += usage.OutputTokenDetails.AudioTokens
+ *totalUsage = nextUsage
+ return service.PreWssConsumeQuota(ctx, info, totalUsage)
}
func OpenaiHandlerWithUsage(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.MaxAPIError) {
diff --git a/relay/channel/palm/relay-palm.go b/relay/channel/palm/relay-palm.go
index c765d827..7a327f3e 100644
--- a/relay/channel/palm/relay-palm.go
+++ b/relay/channel/palm/relay-palm.go
@@ -55,20 +55,31 @@ func palmStreamHandler(c *gin.Context, resp *http.Response) (*types.MaxAPIError,
responseId := helper.GetResponseID(c)
createdTime := common.GetTimestamp()
dataChan := make(chan string)
- stopChan := make(chan bool)
+ stopChan := make(chan struct{}, 1)
+ done := make(chan struct{})
+ producerDone := make(chan struct{})
go func() {
+ defer close(producerDone)
+ defer func() {
+ select {
+ case stopChan <- struct{}{}:
+ case <-done:
+ }
+ }()
responseBody, err := io.ReadAll(resp.Body)
if err != nil {
+ select {
+ case <-done:
+ return
+ default:
+ }
common.SysLog("error reading stream response: " + err.Error())
- stopChan <- true
return
}
- service.CloseResponseBodyGracefully(resp)
var palmResponse PaLMChatResponse
err = json.Unmarshal(responseBody, &palmResponse)
if err != nil {
common.SysLog("error unmarshalling stream response: " + err.Error())
- stopChan <- true
return
}
fullTextResponse := streamResponsePaLM2OpenAI(&palmResponse)
@@ -80,11 +91,13 @@ func palmStreamHandler(c *gin.Context, resp *http.Response) (*types.MaxAPIError,
jsonResponse, err := json.Marshal(fullTextResponse)
if err != nil {
common.SysLog("error marshalling stream response: " + err.Error())
- stopChan <- true
return
}
- dataChan <- string(jsonResponse)
- stopChan <- true
+ select {
+ case dataChan <- string(jsonResponse):
+ case <-done:
+ return
+ }
}()
helper.SetEventStreamHeaders(c)
c.Stream(func(w io.Writer) bool {
@@ -95,9 +108,13 @@ func palmStreamHandler(c *gin.Context, resp *http.Response) (*types.MaxAPIError,
case <-stopChan:
c.Render(-1, common.CustomEvent{Data: "data: [DONE]"})
return false
+ case <-c.Request.Context().Done():
+ return false
}
})
+ close(done)
service.CloseResponseBodyGracefully(resp)
+ <-producerDone
return nil, responseText
}
diff --git a/relay/channel/palm/stream_cancel_test.go b/relay/channel/palm/stream_cancel_test.go
new file mode 100644
index 00000000..33d7401c
--- /dev/null
+++ b/relay/channel/palm/stream_cancel_test.go
@@ -0,0 +1,60 @@
+package palm
+
+import (
+ "context"
+ "io"
+ "net/http"
+ "net/http/httptest"
+ "testing"
+ "time"
+
+ "github.com/gin-gonic/gin"
+ "github.com/stretchr/testify/require"
+)
+
+type closeNotifyRecorder struct {
+ *httptest.ResponseRecorder
+ closeNotify chan bool
+}
+
+func (r *closeNotifyRecorder) CloseNotify() <-chan bool {
+ return r.closeNotify
+}
+
+func TestPaLMStreamHandlerStopsAfterRequestCancellation(t *testing.T) {
+ gin.SetMode(gin.TestMode)
+ c, _ := gin.CreateTestContext(&closeNotifyRecorder{
+ ResponseRecorder: httptest.NewRecorder(),
+ closeNotify: make(chan bool),
+ })
+ request := httptest.NewRequest(http.MethodGet, "/", nil)
+ ctx, cancel := context.WithCancel(request.Context())
+ c.Request = request.WithContext(ctx)
+ reader, writer := io.Pipe()
+ t.Cleanup(func() { _ = writer.Close() })
+
+ returned := make(chan struct{})
+ go func() {
+ _, _ = palmStreamHandler(c, &http.Response{Body: reader})
+ close(returned)
+ }()
+ cancel()
+
+ select {
+ case <-returned:
+ case <-time.After(time.Second):
+ t.Fatal("palm stream handler did not stop after client cancellation")
+ }
+
+ writeResult := make(chan error, 1)
+ go func() {
+ _, err := writer.Write([]byte("{}"))
+ writeResult <- err
+ }()
+ select {
+ case err := <-writeResult:
+ require.ErrorIs(t, err, io.ErrClosedPipe)
+ case <-time.After(time.Second):
+ t.Fatal("palm stream response body was not closed")
+ }
+}
diff --git a/relay/channel/replicate/adaptor.go b/relay/channel/replicate/adaptor.go
index 70bb251b..da1ac070 100644
--- a/relay/channel/replicate/adaptor.go
+++ b/relay/channel/replicate/adaptor.go
@@ -8,7 +8,6 @@ import (
"io"
"mime/multipart"
"net/http"
- "net/textproto"
"strconv"
"strings"
@@ -18,6 +17,7 @@ import (
"github.com/MAX-API-Next/MAX-API/relay/channel"
relaycommon "github.com/MAX-API-Next/MAX-API/relay/common"
relayconstant "github.com/MAX-API-Next/MAX-API/relay/constant"
+ "github.com/MAX-API-Next/MAX-API/relay/helper"
"github.com/MAX-API-Next/MAX-API/service"
"github.com/MAX-API-Next/MAX-API/types"
@@ -445,15 +445,8 @@ func uploadFileFromForm(c *gin.Context, info *relaycommon.RelayInfo, fieldCandid
var body bytes.Buffer
writer := multipart.NewWriter(&body)
- hdr := make(textproto.MIMEHeader)
- hdr.Set("Content-Disposition", fmt.Sprintf("form-data; name=\"content\"; filename=\"%s\"", fileHeader.Filename))
contentType := fileHeader.Header.Get("Content-Type")
- if contentType == "" {
- contentType = "application/octet-stream"
- }
- hdr.Set("Content-Type", contentType)
-
- part, err := writer.CreatePart(hdr)
+ part, err := helper.CreateFormFileWithContentType(writer, "content", fileHeader.Filename, contentType)
if err != nil {
writer.Close()
return "", fmt.Errorf("replicate adaptor: create upload form failed: %w", err)
diff --git a/relay/channel/task/ali/adaptor.go b/relay/channel/task/ali/adaptor.go
index 15438ca3..8db5f56d 100644
--- a/relay/channel/task/ali/adaptor.go
+++ b/relay/channel/task/ali/adaptor.go
@@ -59,10 +59,10 @@ type AliVideoParameters struct {
Resolution string `json:"resolution,omitempty"` // 分辨率: 480P/720P/1080P（图生视频、首尾帧生视频）
Size string `json:"size,omitempty"` // 尺寸: 如 "832*480"（文生视频）
Duration int `json:"duration,omitempty"` // 时长: 3-10秒
- PromptExtend bool `json:"prompt_extend,omitempty"` // 是否开启prompt智能改写
- Watermark bool `json:"watermark,omitempty"` // 是否添加水印
+ PromptExtend *bool `json:"prompt_extend,omitempty"` // 是否开启prompt智能改写
+ Watermark *bool `json:"watermark,omitempty"` // 是否添加水印
Audio *bool `json:"audio,omitempty"` // 是否添加音频（wan2.5）
- Seed int `json:"seed,omitempty"` // 随机数种子
+ Seed *int `json:"seed,omitempty"` // 随机数种子
}
// AliVideoResponse 阿里通义万相响应
@@ -464,8 +464,8 @@ func (a *TaskAdaptor) convertToAliRequest(info *relaycommon.RelayInfo, req relay
ImgURL: firstTaskImage(req),
},
Parameters: &AliVideoParameters{
- PromptExtend: true, // 默认开启智能改写
- Watermark: false,
+ PromptExtend: lo.ToPtr(true), // 默认开启智能改写
+ Watermark: lo.ToPtr(false),
},
}
@@ -544,7 +544,7 @@ func (a *TaskAdaptor) convertToAliKlingRequest(upstreamModel string, req relayco
AspectRatio: sizeToKlingAspectRatio(req.Size),
Duration: 5,
Audio: &audio,
- Watermark: true,
+ Watermark: lo.ToPtr(true),
},
}
@@ -601,7 +601,7 @@ func applyKlingOfficialMetadata(metadata map[string]interface{}, aliReq *AliVide
aliReq.Parameters.Mode = mode
}
if watermark, ok := metadataBool(metadata, "watermark"); ok {
- aliReq.Parameters.Watermark = watermark
+ aliReq.Parameters.Watermark = &watermark
}
if sound := strings.ToLower(metadataString(metadata, "sound")); sound != "" {
audio := sound == "on" || sound == "true" || sound == "1"
diff --git a/relay/channel/task/ali/adaptor_wan27_test.go b/relay/channel/task/ali/adaptor_wan27_test.go
index a0c441c2..38cb9668 100644
--- a/relay/channel/task/ali/adaptor_wan27_test.go
+++ b/relay/channel/task/ali/adaptor_wan27_test.go
@@ -6,6 +6,7 @@ import (
"github.com/MAX-API-Next/MAX-API/common"
relaycommon "github.com/MAX-API-Next/MAX-API/relay/common"
+ "github.com/samber/lo"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
@@ -92,6 +93,27 @@ func TestConvertToAliRequestResolvesDurationFromTaskSubmitReq(t *testing.T) {
}
}
+func TestApplyAliMetadataPreservesExplicitFalseAndZeroParameters(t *testing.T) {
+ aliReq := &AliVideoRequest{
+ Parameters: &AliVideoParameters{
+ PromptExtend: lo.ToPtr(true),
+ Watermark: lo.ToPtr(true),
+ },
+ }
+ metadata := map[string]interface{}{
+ "parameters": map[string]interface{}{
+ "prompt_extend": false,
+ "watermark": false,
+ "seed": 0,
+ },
+ }
+
+ require.NoError(t, applyAliMetadata(metadata, aliReq))
+ payload, err := common.Marshal(aliReq)
+ require.NoError(t, err)
+ require.JSONEq(t, `{"model":"","input":{},"parameters":{"prompt_extend":false,"watermark":false,"seed":0}}`, string(payload))
+}
+
func TestConvertToAliRequestWan27I2VPrefersInputReferenceOverImage(t *testing.T) {
adaptor := &TaskAdaptor{}
req := relaycommon.TaskSubmitReq{
diff --git a/relay/channel/task/hailuo/adaptor.go b/relay/channel/task/hailuo/adaptor.go
index 5adad1cf..c0dbebf8 100644
--- a/relay/channel/task/hailuo/adaptor.go
+++ b/relay/channel/task/hailuo/adaptor.go
@@ -212,9 +212,17 @@ func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, e
taskResult.Progress = "50%"
}
case TaskStatusSuccess:
- taskResult.Status = model.TaskStatusSuccess
- taskResult.Progress = "100%"
- taskResult.Url = a.buildVideoURL(resTask.TaskID, resTask.FileID)
+ // A completed Hailuo task still needs a retrievable file URL. Reporting
+ // success without one settles billing while the video proxy has nothing
+ // to serve. Keep polling so a transient file-retrieval failure can recover.
+ if videoURL := a.buildVideoURL(resTask.TaskID, resTask.FileID); videoURL != "" {
+ taskResult.Status = model.TaskStatusSuccess
+ taskResult.Progress = "100%"
+ taskResult.Url = videoURL
+ } else {
+ taskResult.Status = model.TaskStatusInProgress
+ taskResult.Progress = "90%"
+ }
case TaskStatusFailed:
taskResult.Status = model.TaskStatusFailure
taskResult.Progress = "100%"
diff --git a/relay/channel/task/hailuo/adaptor_test.go b/relay/channel/task/hailuo/adaptor_test.go
new file mode 100644
index 00000000..1bdb51ed
--- /dev/null
+++ b/relay/channel/task/hailuo/adaptor_test.go
@@ -0,0 +1,43 @@
+package hailuo
+
+import (
+ "net/http"
+ "net/http/httptest"
+ "testing"
+ "time"
+
+ "github.com/MAX-API-Next/MAX-API/model"
+ "github.com/MAX-API-Next/MAX-API/service"
+ "github.com/stretchr/testify/require"
+)
+
+func TestParseTaskResultWaitsForRetrievableVideoURL(t *testing.T) {
+ service.InitHttpClient()
+
+ requestPath := make(chan string, 1)
+ server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ requestPath <- r.URL.Path
+ w.WriteHeader(http.StatusServiceUnavailable)
+ }))
+ t.Cleanup(server.Close)
+
+ adaptor := &TaskAdaptor{apiKey: "test-key", baseURL: server.URL}
+ result, err := adaptor.ParseTaskResult([]byte(`{
+ "task_id":"task-1",
+ "status":"Success",
+ "file_id":"file-1",
+ "base_resp":{"status_code":0}
+ }`))
+
+ require.NoError(t, err)
+ require.Equal(t, string(model.TaskStatusInProgress), result.Status)
+ require.Equal(t, "90%", result.Progress)
+ require.Empty(t, result.Url)
+
+ select {
+ case path := <-requestPath:
+ require.Equal(t, "/v1/files/retrieve", path)
+ case <-time.After(time.Second):
+ t.Fatal("timed out waiting for the Hailuo file retrieval request")
+ }
+}
diff --git a/relay/channel/task/jimeng/adaptor.go b/relay/channel/task/jimeng/adaptor.go
index f1083fc9..c0e1d6eb 100644
--- a/relay/channel/task/jimeng/adaptor.go
+++ b/relay/channel/task/jimeng/adaptor.go
@@ -444,6 +444,7 @@ func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, e
taskResult.Reason = resTask.Message
taskResult.Status = model.TaskStatusFailure
taskResult.Progress = "100%"
+ return &taskResult, nil
}
switch resTask.Data.Status {
case "in_queue":
diff --git a/relay/channel/task/jimeng/adaptor_test.go b/relay/channel/task/jimeng/adaptor_test.go
new file mode 100644
index 00000000..dfc060aa
--- /dev/null
+++ b/relay/channel/task/jimeng/adaptor_test.go
@@ -0,0 +1,21 @@
+package jimeng
+
+import (
+ "testing"
+
+ "github.com/MAX-API-Next/MAX-API/model"
+ "github.com/stretchr/testify/require"
+)
+
+func TestParseTaskResultKeepsErrorResponseAsFailure(t *testing.T) {
+ result, err := (&TaskAdaptor{}).ParseTaskResult([]byte(`{
+ "code":50001,
+ "message":"upstream failed",
+ "data":{"status":"done","video_url":"https://example.com/video.mp4"}
+ }`))
+
+ require.NoError(t, err)
+ require.Equal(t, string(model.TaskStatusFailure), result.Status)
+ require.Equal(t, "100%", result.Progress)
+ require.Equal(t, "upstream failed", result.Reason)
+}
diff --git a/relay/channel/task/sora/adaptor.go b/relay/channel/task/sora/adaptor.go
index 75abdc33..93cc0311 100644
--- a/relay/channel/task/sora/adaptor.go
+++ b/relay/channel/task/sora/adaptor.go
@@ -6,7 +6,6 @@ import (
"io"
"mime/multipart"
"net/http"
- "net/textproto"
"strconv"
"strings"
@@ -17,6 +16,7 @@ import (
"github.com/MAX-API-Next/MAX-API/relay/channel"
taskcommon "github.com/MAX-API-Next/MAX-API/relay/channel/task/taskcommon"
relaycommon "github.com/MAX-API-Next/MAX-API/relay/common"
+ "github.com/MAX-API-Next/MAX-API/relay/helper"
"github.com/MAX-API-Next/MAX-API/service"
"github.com/gin-gonic/gin"
@@ -202,10 +202,7 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn
continue
}
}
- h := make(textproto.MIMEHeader)
- h.Set("Content-Disposition", fmt.Sprintf(`form-data; name="%s"; filename="%s"`, fieldName, fh.Filename))
- h.Set("Content-Type", ct)
- part, err := writer.CreatePart(h)
+ part, err := helper.CreateFormFileWithContentType(writer, fieldName, fh.Filename, ct)
if err != nil {
f.Close()
continue
diff --git a/relay/channel/vertex/relay-vertex.go b/relay/channel/vertex/relay-vertex.go
index 39e59391..eb6b35de 100644
--- a/relay/channel/vertex/relay-vertex.go
+++ b/relay/channel/vertex/relay-vertex.go
@@ -2,6 +2,11 @@ package vertex
import "github.com/MAX-API-Next/MAX-API/common"
+func configuredRegion(value any) (string, bool) {
+ region, ok := value.(string)
+ return region, ok && region != ""
+}
+
func GetModelRegion(other string, localModelName string) string {
// if other is json string
if common.IsJsonObject(other) {
@@ -9,14 +14,13 @@ func GetModelRegion(other string, localModelName string) string {
if err != nil {
return other // return original if parsing fails
}
- if m[localModelName] != nil {
- return m[localModelName].(string)
- } else {
- if v, ok := m["default"]; ok {
- return v.(string)
- }
- return "global"
+ if region, ok := configuredRegion(m[localModelName]); ok {
+ return region
+ }
+ if region, ok := configuredRegion(m["default"]); ok {
+ return region
}
+ return "global"
}
return other
}
diff --git a/relay/channel/vertex/relay_vertex_test.go b/relay/channel/vertex/relay_vertex_test.go
new file mode 100644
index 00000000..762f448a
--- /dev/null
+++ b/relay/channel/vertex/relay_vertex_test.go
@@ -0,0 +1,12 @@
+package vertex
+
+import (
+ "testing"
+
+ "github.com/stretchr/testify/require"
+)
+
+func TestGetModelRegionIgnoresNonStringConfiguration(t *testing.T) {
+ require.Equal(t, "global", GetModelRegion(`{"model-a":123,"default":false}`, "model-a"))
+ require.Equal(t, "us-central1", GetModelRegion(`{"model-a":123,"default":"us-central1"}`, "model-a"))
+}
diff --git a/relay/channel/xunfei/dto.go b/relay/channel/xunfei/dto.go
index a13ed71d..2f8166eb 100644
--- a/relay/channel/xunfei/dto.go
+++ b/relay/channel/xunfei/dto.go
@@ -15,8 +15,8 @@ type XunfeiChatRequest struct {
Chat struct {
Domain string `json:"domain,omitempty"`
Temperature *float64 `json:"temperature,omitempty"`
- TopK int `json:"top_k,omitempty"`
- MaxTokens uint `json:"max_tokens,omitempty"`
+ TopK *int `json:"top_k,omitempty"`
+ MaxTokens *uint `json:"max_tokens,omitempty"`
Auditing bool `json:"auditing,omitempty"`
} `json:"chat"`
} `json:"parameter"`
diff --git a/relay/channel/xunfei/relay-xunfei.go b/relay/channel/xunfei/relay-xunfei.go
index 5578ee6f..927c78ee 100644
--- a/relay/channel/xunfei/relay-xunfei.go
+++ b/relay/channel/xunfei/relay-xunfei.go
@@ -1,12 +1,13 @@
package xunfei
import (
+ "context"
"crypto/hmac"
"crypto/sha256"
"encoding/base64"
- "encoding/json"
"fmt"
"io"
+ "net/http"
"net/url"
"strings"
"time"
@@ -16,8 +17,6 @@ import (
"github.com/MAX-API-Next/MAX-API/dto"
"github.com/MAX-API-Next/MAX-API/relay/helper"
"github.com/MAX-API-Next/MAX-API/types"
- "github.com/samber/lo"
-
"github.com/gin-gonic/gin"
"github.com/gorilla/websocket"
)
@@ -49,8 +48,8 @@ func requestOpenAI2Xunfei(request dto.GeneralOpenAIRequest, xunfeiAppId string,
xunfeiRequest.Header.AppId = xunfeiAppId
xunfeiRequest.Parameter.Chat.Domain = domain
xunfeiRequest.Parameter.Chat.Temperature = request.Temperature
- xunfeiRequest.Parameter.Chat.TopK = lo.FromPtrOr(request.N, 0)
- xunfeiRequest.Parameter.Chat.MaxTokens = request.GetMaxTokens()
+ xunfeiRequest.Parameter.Chat.TopK = request.TopK
+ xunfeiRequest.Parameter.Chat.MaxTokens = dto.ResolveMaxTokens(request.MaxTokens, request.MaxCompletionTokens)
xunfeiRequest.Payload.Message.Text = messages
return &xunfeiRequest
}
@@ -130,47 +129,68 @@ func buildXunfeiAuthUrl(hostUrl string, apiKey, apiSecret string) string {
func xunfeiStreamHandler(c *gin.Context, textRequest dto.GeneralOpenAIRequest, appId string, apiSecret string, apiKey string) (*dto.Usage, *types.MaxAPIError) {
domain, authUrl := getXunfeiAuthUrl(c, apiKey, apiSecret, textRequest.Model)
- dataChan, stopChan, err := xunfeiMakeRequest(textRequest, domain, authUrl, appId)
+ events, err := xunfeiMakeRequest(c.Request.Context(), textRequest, domain, authUrl, appId)
if err != nil {
return nil, types.NewError(err, types.ErrorCodeDoRequestFailed)
}
helper.SetEventStreamHeaders(c)
var usage dto.Usage
+ var streamErr error
+ wrote := false
c.Stream(func(w io.Writer) bool {
- select {
- case xunfeiResponse := <-dataChan:
+ event, ok := <-events
+ if !ok {
+ c.Render(-1, common.CustomEvent{Data: "data: [DONE]"})
+ return false
+ }
+ if event.Err != nil {
+ streamErr = event.Err
+ if wrote {
+ c.Render(-1, common.CustomEvent{Data: "data: [DONE]"})
+ }
+ return false
+ }
+ if event.Response != nil {
+ xunfeiResponse := *event.Response
usage.PromptTokens += xunfeiResponse.Payload.Usage.Text.PromptTokens
usage.CompletionTokens += xunfeiResponse.Payload.Usage.Text.CompletionTokens
usage.TotalTokens += xunfeiResponse.Payload.Usage.Text.TotalTokens
response := streamResponseXunfei2OpenAI(&xunfeiResponse)
- jsonResponse, err := json.Marshal(response)
+ jsonResponse, err := common.Marshal(response)
if err != nil {
common.SysLog("error marshalling stream response: " + err.Error())
return true
}
c.Render(-1, common.CustomEvent{Data: "data: " + string(jsonResponse)})
+ wrote = true
return true
- case <-stopChan:
- c.Render(-1, common.CustomEvent{Data: "data: [DONE]"})
- return false
}
+ return true
})
+ if streamErr != nil && !wrote {
+ return nil, types.NewError(streamErr, types.ErrorCodeBadResponse)
+ }
+ if streamErr != nil {
+ common.SysLog("xunfei stream terminated early: " + streamErr.Error())
+ }
return &usage, nil
}
func xunfeiHandler(c *gin.Context, textRequest dto.GeneralOpenAIRequest, appId string, apiSecret string, apiKey string) (*dto.Usage, *types.MaxAPIError) {
domain, authUrl := getXunfeiAuthUrl(c, apiKey, apiSecret, textRequest.Model)
- dataChan, stopChan, err := xunfeiMakeRequest(textRequest, domain, authUrl, appId)
+ events, err := xunfeiMakeRequest(c.Request.Context(), textRequest, domain, authUrl, appId)
if err != nil {
return nil, types.NewError(err, types.ErrorCodeDoRequestFailed)
}
var usage dto.Usage
var content string
var xunfeiResponse XunfeiChatResponse
- stop := false
- for !stop {
- select {
- case xunfeiResponse = <-dataChan:
+ for event := range events {
+ if event.Err != nil {
+ return nil, types.NewError(event.Err, types.ErrorCodeBadResponse)
+ }
+ if event.Response != nil {
+ xunfeiResponse = *event.Response
if len(xunfeiResponse.Payload.Choices.Text) == 0 {
continue
}
@@ -178,7 +198,6 @@ func xunfeiHandler(c *gin.Context, textRequest dto.GeneralOpenAIRequest, appId s
usage.PromptTokens += xunfeiResponse.Payload.Usage.Text.PromptTokens
usage.CompletionTokens += xunfeiResponse.Payload.Usage.Text.CompletionTokens
usage.TotalTokens += xunfeiResponse.Payload.Usage.Text.TotalTokens
- case stop = <-stopChan:
}
}
if len(xunfeiResponse.Payload.Choices.Text) == 0 {
@@ -191,7 +210,7 @@ func xunfeiHandler(c *gin.Context, textRequest dto.GeneralOpenAIRequest, appId s
xunfeiResponse.Payload.Choices.Text[0].Content = content
response := responseXunfei2OpenAI(&xunfeiResponse)
- jsonResponse, err := json.Marshal(response)
+ jsonResponse, err := common.Marshal(response)
if err != nil {
return nil, types.NewError(err, types.ErrorCodeBadResponseBody)
}
@@ -200,51 +219,107 @@ func xunfeiHandler(c *gin.Context, textRequest dto.GeneralOpenAIRequest, appId s
return &usage, nil
}
-func xunfeiMakeRequest(textRequest dto.GeneralOpenAIRequest, domain, authUrl, appId string) (chan XunfeiChatResponse, chan bool, error) {
+type xunfeiEvent struct {
+ Response *XunfeiChatResponse
+ Err error
+}
+
+func emitXunfeiEvent(ctx context.Context, events chan<- xunfeiEvent, event xunfeiEvent) bool {
+ select {
+ case events <- event:
+ return true
+ case <-ctx.Done():
+ return false
+ }
+}
+
+func xunfeiMakeRequest(ctx context.Context, textRequest dto.GeneralOpenAIRequest, domain, authUrl, appId string) (<-chan xunfeiEvent, error) {
d := websocket.Dialer{
HandshakeTimeout: 5 * time.Second,
}
- conn, resp, err := d.Dial(authUrl, nil)
- if err != nil || resp.StatusCode != 101 {
- return nil, nil, err
+ conn, resp, err := d.DialContext(ctx, authUrl, nil)
+ if err != nil {
+ status := 0
+ if resp != nil {
+ status = resp.StatusCode
+ if resp.Body != nil {
+ _ = resp.Body.Close()
+ }
+ }
+ return nil, fmt.Errorf("xunfei websocket handshake failed (status=%d): %w", status, err)
+ }
+ if resp == nil || resp.StatusCode != http.StatusSwitchingProtocols {
+ if conn != nil {
+ _ = conn.Close()
+ }
+ status := 0
+ if resp != nil {
+ status = resp.StatusCode
+ }
+ return nil, fmt.Errorf("xunfei websocket handshake failed: unexpected status %d", status)
}
data := requestOpenAI2Xunfei(textRequest, appId, domain)
+ _ = conn.SetWriteDeadline(xunfeiOperationDeadline(ctx, 10*time.Second))
err = conn.WriteJSON(data)
if err != nil {
- return nil, nil, err
+ _ = conn.Close()
+ return nil, err
}
- dataChan := make(chan XunfeiChatResponse)
- stopChan := make(chan bool)
+ events := make(chan xunfeiEvent)
+ done := make(chan struct{})
+ go func() {
+ select {
+ case <-ctx.Done():
+ _ = conn.Close()
+ case <-done:
+ }
+ }()
go func() {
defer func() {
- conn.Close()
+ close(done)
+ close(events)
+ _ = conn.Close()
}()
for {
+ _ = conn.SetReadDeadline(xunfeiOperationDeadline(ctx, 60*time.Second))
_, msg, err := conn.ReadMessage()
if err != nil {
- common.SysLog("error reading stream response: " + err.Error())
- break
+ if ctx.Err() != nil {
+ err = ctx.Err()
+ }
+ emitXunfeiEvent(ctx, events, xunfeiEvent{Err: fmt.Errorf("read xunfei websocket response: %w", err)})
+ return
}
var response XunfeiChatResponse
- err = json.Unmarshal(msg, &response)
+ err = common.Unmarshal(msg, &response)
if err != nil {
- common.SysLog("error unmarshalling stream response: " + err.Error())
- break
+ emitXunfeiEvent(ctx, events, xunfeiEvent{Err: fmt.Errorf("decode xunfei websocket response: %w", err)})
+ return
+ }
+ if response.Header.Code != 0 {
+ emitXunfeiEvent(ctx, events, xunfeiEvent{Err: fmt.Errorf("xunfei upstream error %d: %s", response.Header.Code, response.Header.Message)})
+ return
+ }
+ if !emitXunfeiEvent(ctx, events, xunfeiEvent{Response: &response}) {
+ return
}
- dataChan <- response
if response.Payload.Choices.Status == 2 {
- if err != nil {
- common.SysLog("error closing websocket connection: " + err.Error())
- }
- break
+ return
}
}
- stopChan <- true
}()
- return dataChan, stopChan, nil
+ return events, nil
+}
+
+func xunfeiOperationDeadline(ctx context.Context, fallback time.Duration) time.Time {
+ deadline := time.Now().Add(fallback)
+ if ctxDeadline, ok := ctx.Deadline(); ok && ctxDeadline.Before(deadline) {
+ return ctxDeadline
+ }
+ return deadline
}
func apiVersion2domain(apiVersion string) string {
diff --git a/relay/channel/xunfei/relay_xunfei_test.go b/relay/channel/xunfei/relay_xunfei_test.go
new file mode 100644
index 00000000..f7b44c37
--- /dev/null
+++ b/relay/channel/xunfei/relay_xunfei_test.go
@@ -0,0 +1,93 @@
+package xunfei
+
+import (
+ "context"
+ "net/http"
+ "net/http/httptest"
+ "strings"
+ "testing"
+ "time"
+
+ "github.com/MAX-API-Next/MAX-API/common"
+ "github.com/MAX-API-Next/MAX-API/dto"
+ "github.com/gorilla/websocket"
+ "github.com/stretchr/testify/require"
+)
+
+func TestEmitXunfeiEventStopsWhenRequestContextIsCanceled(t *testing.T) {
+ ctx, cancel := context.WithCancel(context.Background())
+ cancel()
+
+ done := make(chan bool, 1)
+ go func() {
+ done <- emitXunfeiEvent(ctx, make(chan xunfeiEvent), xunfeiEvent{})
+ }()
+
+ select {
+ case emitted := <-done:
+ require.False(t, emitted)
+ case <-time.After(time.Second):
+ t.Fatal("xunfei event delivery blocked after request cancellation")
+ }
+}
+
+func TestXunfeiMakeRequestRejectsNonSwitchingProtocolsResponse(t *testing.T) {
+ server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+ w.WriteHeader(http.StatusUnauthorized)
+ }))
+ defer server.Close()
+
+ authURL := "ws" + strings.TrimPrefix(server.URL, "http")
+ events, err := xunfeiMakeRequest(context.Background(), dto.GeneralOpenAIRequest{}, "generalv3", authURL, "app-id")
+ require.Error(t, err)
+ require.Nil(t, events)
+ require.Contains(t, err.Error(), "handshake failed")
+}
+
+func TestXunfeiMakeRequestSurfacesUpstreamHeaderError(t *testing.T) {
+ upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }}
+ server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ conn, err := upgrader.Upgrade(w, r, nil)
+ require.NoError(t, err)
+ defer conn.Close()
+
+ _, _, err = conn.ReadMessage()
+ require.NoError(t, err)
+ require.NoError(t, conn.WriteJSON(map[string]any{
+ "header": map[string]any{
+ "code": 10013,
+ "message": "invalid api key",
+ },
+ }))
+ }))
+ defer server.Close()
+
+ authURL := "ws" + strings.TrimPrefix(server.URL, "http")
+ events, err := xunfeiMakeRequest(context.Background(), dto.GeneralOpenAIRequest{}, "generalv3", authURL, "app-id")
+ require.NoError(t, err)
+
+ select {
+ case event, ok := <-events:
+ require.True(t, ok)
+ require.Nil(t, event.Response)
+ require.Error(t, event.Err)
+ require.Contains(t, event.Err.Error(), "xunfei upstream error 10013: invalid api key")
+ case <-time.After(time.Second):
+ t.Fatal("xunfei upstream error was not delivered")
+ }
+}
+
+func TestRequestOpenAI2XunfeiPreservesExplicitZeroValues(t *testing.T) {
+ zeroInt := 0
+ zeroUint := uint(0)
+ completionCount := 2
+ request := requestOpenAI2Xunfei(dto.GeneralOpenAIRequest{
+ TopK: &zeroInt,
+ N: &completionCount,
+ MaxTokens: &zeroUint,
+ }, "app-id", "generalv3")
+
+ payload, err := common.Marshal(request)
+ require.NoError(t, err)
+ require.JSONEq(t, `{"header":{"app_id":"app-id"},"parameter":{"chat":{"domain":"generalv3","top_k":0,"max_tokens":0}},"payload":{"message":{"text":[]}}}`, string(payload))
+}
diff --git a/relay/channel/zhipu/dto.go b/relay/channel/zhipu/dto.go
index 7ecd198d..52b33c89 100644
--- a/relay/channel/zhipu/dto.go
+++ b/relay/channel/zhipu/dto.go
@@ -14,7 +14,7 @@ type ZhipuMessage struct {
type ZhipuRequest struct {
Prompt []ZhipuMessage `json:"prompt"`
Temperature *float64 `json:"temperature,omitempty"`
- TopP float64 `json:"top_p,omitempty"`
+ TopP *float64 `json:"top_p,omitempty"`
RequestId string `json:"request_id,omitempty"`
Incremental bool `json:"incremental,omitempty"`
}
diff --git a/relay/channel/zhipu/relay-zhipu.go b/relay/channel/zhipu/relay-zhipu.go
index 233d3d96..f160d676 100644
--- a/relay/channel/zhipu/relay-zhipu.go
+++ b/relay/channel/zhipu/relay-zhipu.go
@@ -16,8 +16,6 @@ import (
"github.com/MAX-API-Next/MAX-API/relay/helper"
"github.com/MAX-API-Next/MAX-API/service"
"github.com/MAX-API-Next/MAX-API/types"
- "github.com/samber/lo"
-
"github.com/gin-gonic/gin"
"github.com/golang-jwt/jwt/v5"
)
@@ -99,7 +97,7 @@ func requestOpenAI2Zhipu(request dto.GeneralOpenAIRequest) *ZhipuRequest {
return &ZhipuRequest{
Prompt: messages,
Temperature: request.Temperature,
- TopP: lo.FromPtrOr(request.TopP, 0),
+ TopP: request.TopP,
Incremental: false,
}
}
@@ -161,8 +159,17 @@ func zhipuStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.
scanner.Split(bufio.ScanLines)
dataChan := make(chan string)
metaChan := make(chan string)
- stopChan := make(chan bool)
+ stopChan := make(chan struct{}, 1)
+ done := make(chan struct{})
+ producerDone := make(chan struct{})
go func() {
+ defer close(producerDone)
+ defer func() {
+ select {
+ case stopChan <- struct{}{}:
+ case <-done:
+ }
+ }()
for scanner.Scan() {
data := scanner.Text()
lines := strings.Split(data, "\n")
@@ -171,16 +178,35 @@ func zhipuStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.
continue
}
if line[:5] == "data:" {
- dataChan <- line[5:]
+ select {
+ case dataChan <- line[5:]:
+ case <-done:
+ return
+ }
if i != len(lines)-1 {
- dataChan <- "\n"
+ select {
+ case dataChan <- "\n":
+ case <-done:
+ return
+ }
}
} else if line[:5] == "meta:" {
- metaChan <- line[5:]
+ select {
+ case metaChan <- line[5:]:
+ case <-done:
+ return
+ }
}
}
}
- stopChan <- true
+ if err := scanner.Err(); err != nil {
+ select {
+ case <-done:
+ return
+ default:
+ }
+ common.SysLog("error reading stream: " + err.Error())
+ }
}()
helper.SetEventStreamHeaders(c)
c.Stream(func(w io.Writer) bool {
@@ -213,9 +239,13 @@ func zhipuStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.
case <-stopChan:
c.Render(-1, common.CustomEvent{Data: "data: [DONE]"})
return false
+ case <-c.Request.Context().Done():
+ return false
}
})
+ close(done)
service.CloseResponseBodyGracefully(resp)
+ <-producerDone
return usage, nil
}
@@ -244,5 +274,8 @@ func zhipuHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Respon
c.Writer.Header().Set("Content-Type", "application/json")
c.Writer.WriteHeader(resp.StatusCode)
_, err = c.Writer.Write(jsonResponse)
+ if err != nil {
+ common.SysLog("failed to write Zhipu response: " + err.Error())
+ }
return &fullTextResponse.Usage, nil
}
diff --git a/relay/channel/zhipu/stream_cancel_test.go b/relay/channel/zhipu/stream_cancel_test.go
new file mode 100644
index 00000000..a264462b
--- /dev/null
+++ b/relay/channel/zhipu/stream_cancel_test.go
@@ -0,0 +1,72 @@
+package zhipu
+
+import (
+ "context"
+ "io"
+ "net/http"
+ "net/http/httptest"
+ "testing"
+ "time"
+
+ "github.com/MAX-API-Next/MAX-API/common"
+ "github.com/MAX-API-Next/MAX-API/dto"
+ relaycommon "github.com/MAX-API-Next/MAX-API/relay/common"
+ "github.com/gin-gonic/gin"
+ "github.com/stretchr/testify/require"
+)
+
+type closeNotifyRecorder struct {
+ *httptest.ResponseRecorder
+ closeNotify chan bool
+}
+
+func (r *closeNotifyRecorder) CloseNotify() <-chan bool {
+ return r.closeNotify
+}
+
+func TestZhipuStreamHandlerStopsAfterRequestCancellation(t *testing.T) {
+ gin.SetMode(gin.TestMode)
+ c, _ := gin.CreateTestContext(&closeNotifyRecorder{
+ ResponseRecorder: httptest.NewRecorder(),
+ closeNotify: make(chan bool),
+ })
+ request := httptest.NewRequest(http.MethodGet, "/", nil)
+ ctx, cancel := context.WithCancel(request.Context())
+ c.Request = request.WithContext(ctx)
+ reader, writer := io.Pipe()
+ t.Cleanup(func() { _ = writer.Close() })
+
+ returned := make(chan struct{})
+ go func() {
+ _, _ = zhipuStreamHandler(c, &relaycommon.RelayInfo{}, &http.Response{Body: reader})
+ close(returned)
+ }()
+ cancel()
+
+ select {
+ case <-returned:
+ case <-time.After(time.Second):
+ t.Fatal("zhipu stream handler did not stop after client cancellation")
+ }
+
+ writeResult := make(chan error, 1)
+ go func() {
+ _, err := writer.Write([]byte("data: event\n"))
+ writeResult <- err
+ }()
+ select {
+ case err := <-writeResult:
+ require.ErrorIs(t, err, io.ErrClosedPipe)
+ case <-time.After(time.Second):
+ t.Fatal("zhipu stream response body was not closed")
+ }
+}
+
+func TestRequestOpenAI2ZhipuPreservesExplicitTopPZero(t *testing.T) {
+ zero := 0.0
+ request := requestOpenAI2Zhipu(dto.GeneralOpenAIRequest{TopP: &zero})
+
+ payload, err := common.Marshal(request)
+ require.NoError(t, err)
+ require.JSONEq(t, `{"prompt":[],"top_p":0}`, string(payload))
+}
diff --git a/relay/common/relay_info.go b/relay/common/relay_info.go
index 77602847..aa14a5c6 100644
--- a/relay/common/relay_info.go
+++ b/relay/common/relay_info.go
@@ -680,6 +680,17 @@ type TaskRelayInfo struct {
// PublicTaskID 是提交时预生成的 task_xxxx 格式公开 ID，
// 供 DoResponse 在返回给客户端时使用（避免暴露上游真实 ID）。
PublicTaskID string
+ // PersistedTaskID is the local placeholder row created before the upstream
+ // request. It is reused across channel retries for the same client request.
+ PersistedTaskID int64
+ // UpstreamTaskResponseReceived is set after an upstream 2xx response. Until
+ // parsing proves otherwise, retrying or refunding that state could create a
+ // duplicate or unbilled provider task.
+ UpstreamTaskResponseReceived bool
+ // UpstreamTaskOutcomeUnknown is set when the complete request was written but
+ // the transport failed before a response arrived. Retrying or refunding that
+ // state could duplicate provider work or lose provider spend.
+ UpstreamTaskOutcomeUnknown bool
ConsumeQuota bool
diff --git a/relay/helper/multipart.go b/relay/helper/multipart.go
new file mode 100644
index 00000000..8a8d11db
--- /dev/null
+++ b/relay/helper/multipart.go
@@ -0,0 +1,37 @@
+package helper
+
+import (
+ "errors"
+ "io"
+ "mime"
+ "mime/multipart"
+ "net/textproto"
+ "strings"
+)
+
+// CreateFormFileWithContentType creates a multipart file part without copying
+// client-controlled names directly into MIME headers.
+func CreateFormFileWithContentType(writer *multipart.Writer, fieldName, fileName, contentType string) (io.Writer, error) {
+ if writer == nil {
+ return nil, errors.New("multipart writer is nil")
+ }
+ disposition := mime.FormatMediaType("form-data", map[string]string{
+ "name": fieldName,
+ "filename": fileName,
+ })
+ if disposition == "" {
+ return nil, errors.New("invalid multipart file name")
+ }
+ header := make(textproto.MIMEHeader)
+ header.Set("Content-Disposition", disposition)
+ header.Set("Content-Type", safeMultipartContentType(contentType))
+ return writer.CreatePart(header)
+}
+
+func safeMultipartContentType(contentType string) string {
+ mediaType, _, err := mime.ParseMediaType(strings.TrimSpace(contentType))
+ if err != nil || mediaType == "" {
+ return "application/octet-stream"
+ }
+ return mediaType
+}
diff --git a/relay/helper/multipart_test.go b/relay/helper/multipart_test.go
new file mode 100644
index 00000000..5e0e7aa1
--- /dev/null
+++ b/relay/helper/multipart_test.go
@@ -0,0 +1,44 @@
+package helper
+
+import (
+ "bytes"
+ "mime/multipart"
+ "strings"
+ "testing"
+
+ "github.com/stretchr/testify/assert"
+ "github.com/stretchr/testify/require"
+)
+
+func TestCreateFormFileWithContentTypeEscapesFilename(t *testing.T) {
+ var body bytes.Buffer
+ writer := multipart.NewWriter(&body)
+ part, err := CreateFormFileWithContentType(writer, "file", `quote".png`, "image/png")
+ require.NoError(t, err)
+ _, err = part.Write([]byte("image"))
+ require.NoError(t, err)
+ require.NoError(t, writer.Close())
+
+ reader := multipart.NewReader(bytes.NewReader(body.Bytes()), writer.Boundary())
+ form, err := reader.ReadForm(1024)
+ require.NoError(t, err)
+ defer form.RemoveAll()
+ require.Len(t, form.File["file"], 1)
+ assert.Equal(t, `quote".png`, form.File["file"][0].Filename)
+}
+
+func TestCreateFormFileWithContentTypeEncodesHeaderInjectionCharacters(t *testing.T) {
+ var body bytes.Buffer
+ writer := multipart.NewWriter(&body)
+ part, err := CreateFormFileWithContentType(writer, "file", "image.png\r\nX-Injected: true", "image/png")
+ require.NoError(t, err)
+ _, err = part.Write([]byte("image"))
+ require.NoError(t, err)
+ require.NoError(t, writer.Close())
+
+ reader := multipart.NewReader(bytes.NewReader(body.Bytes()), writer.Boundary())
+ parsedPart, err := reader.NextPart()
+ require.NoError(t, err)
+ assert.Empty(t, parsedPart.Header.Get("X-Injected"))
+ assert.False(t, strings.Contains(parsedPart.Header.Get("Content-Disposition"), "\r\n"))
+}
diff --git a/relay/mjproxy_handler.go b/relay/mjproxy_handler.go
index 2d2504e6..0a5c2279 100644
--- a/relay/mjproxy_handler.go
+++ b/relay/mjproxy_handler.go
@@ -26,15 +26,29 @@ import (
"github.com/gin-gonic/gin"
)
+func writeMidjourneyStatusCode(c *gin.Context, upstreamStatusCode int) {
+ c.Writer.WriteHeader(service.MapStatusCode(upstreamStatusCode, c.GetString("status_code_mapping")))
+}
+
func RelayMidjourneyImage(c *gin.Context) {
taskId := c.Param("id")
- midjourneyTask := model.GetByOnlyMJId(taskId)
+ userID, userErr := strconv.Atoi(c.Query("uid"))
+ expiresAt, expiresErr := strconv.ParseInt(c.Query("expires"), 10, 64)
+ if userErr != nil || expiresErr != nil ||
+ !service.ValidateMidjourneyImageURL(taskId, userID, expiresAt, c.Query("signature"), time.Now()) {
+ c.JSON(http.StatusForbidden, gin.H{
+ "error": "midjourney_image_authorization_failed",
+ })
+ return
+ }
+ midjourneyTask := model.GetByMJId(userID, taskId)
if midjourneyTask == nil {
- c.JSON(400, gin.H{
- "error": "midjourney_task_not_found",
+ c.JSON(http.StatusForbidden, gin.H{
+ "error": "midjourney_image_authorization_failed",
})
return
}
+ c.Header("Cache-Control", "private, max-age=300")
var httpClient *http.Client
if channel, err := model.CacheGetChannel(midjourneyTask.ChannelId); err == nil {
proxy := channel.GetSetting().Proxy
@@ -141,9 +155,9 @@ func coverMidjourneyTaskDto(c *gin.Context, originTask *model.Midjourney) (midjo
midjourneyTask.FinishTime = originTask.FinishTime
midjourneyTask.ImageUrl = ""
if originTask.ImageUrl != "" && setting.MjForwardUrlEnabled {
- midjourneyTask.ImageUrl = system_setting.ServerAddress + "/mj/image/" + originTask.MjId
+ midjourneyTask.ImageUrl = service.BuildMidjourneyImageURL(system_setting.ServerAddress, originTask.MjId, originTask.UserId, time.Now())
if originTask.Status != "SUCCESS" {
- midjourneyTask.ImageUrl += "?rand=" + strconv.FormatInt(time.Now().UnixNano(), 10)
+ midjourneyTask.ImageUrl += "&rand=" + strconv.FormatInt(time.Now().UnixNano(), 10)
}
} else {
midjourneyTask.ImageUrl = originTask.ImageUrl
@@ -225,7 +239,7 @@ func RelaySwapFace(c *gin.Context, info *relaycommon.RelayInfo) *dto.MidjourneyR
}
defer func() {
if mjResp.StatusCode == 200 && mjResp.Response.Code == 1 {
- err := service.PostConsumeQuota(info, priceData.Quota, 0, true)
+ err := service.PostConsumeQuotaOnce(info, "midjourney-success", priceData.Quota, 0, true)
if err != nil {
common.SysLog("error consuming token remain quota: " + err.Error())
}
@@ -271,7 +285,7 @@ func RelaySwapFace(c *gin.Context, info *relaycommon.RelayInfo) *dto.MidjourneyR
if err != nil {
return service.MidjourneyErrorWrapper(constant.MjRequestError, "insert_midjourney_task_failed")
}
- c.Writer.WriteHeader(mjResp.StatusCode)
+ writeMidjourneyStatusCode(c, mjResp.StatusCode)
respBody, err := json.Marshal(midjResponse)
if err != nil {
return service.MidjourneyErrorWrapper(constant.MjRequestError, "unmarshal_response_body_failed")
@@ -307,7 +321,7 @@ func RelayMidjourneyTaskImageSeed(c *gin.Context) *dto.MidjourneyResponse {
return &midjResponseWithStatus.Response
}
midjResponse := &midjResponseWithStatus.Response
- c.Writer.WriteHeader(midjResponseWithStatus.StatusCode)
+ writeMidjourneyStatusCode(c, midjResponseWithStatus.StatusCode)
respBody, err := json.Marshal(midjResponse)
if err != nil {
return service.MidjourneyErrorWrapper(constant.MjRequestError, "unmarshal_response_body_failed")
@@ -532,7 +546,7 @@ func RelayMidjourneySubmit(c *gin.Context, relayInfo *relaycommon.RelayInfo) *dt
defer func() {
if consumeQuota && midjResponseWithStatus.StatusCode == 200 {
- err := service.PostConsumeQuota(relayInfo, priceData.Quota, 0, true)
+ err := service.PostConsumeQuotaOnce(relayInfo, "midjourney-success", priceData.Quota, 0, true)
if err != nil {
common.SysLog("error consuming token remain quota: " + err.Error())
}
@@ -642,7 +656,7 @@ func RelayMidjourneySubmit(c *gin.Context, relayInfo *relaycommon.RelayInfo) *dt
//for k, v := range resp.Header {
// c.Writer.Header().Set(k, v[0])
//}
- c.Writer.WriteHeader(midjResponseWithStatus.StatusCode)
+ writeMidjourneyStatusCode(c, midjResponseWithStatus.StatusCode)
_, err = io.Copy(c.Writer, bodyReader)
if err != nil {
diff --git a/relay/mjproxy_handler_test.go b/relay/mjproxy_handler_test.go
new file mode 100644
index 00000000..bb95259e
--- /dev/null
+++ b/relay/mjproxy_handler_test.go
@@ -0,0 +1,67 @@
+package relay
+
+import (
+ "net/http"
+ "net/http/httptest"
+ "strconv"
+ "testing"
+ "time"
+
+ "github.com/MAX-API-Next/MAX-API/common"
+ "github.com/MAX-API-Next/MAX-API/model"
+ "github.com/MAX-API-Next/MAX-API/service"
+ "github.com/gin-gonic/gin"
+ "github.com/glebarez/sqlite"
+ "github.com/stretchr/testify/require"
+ "gorm.io/gorm"
+)
+
+func TestWriteMidjourneyStatusCodeUsesChannelMapping(t *testing.T) {
+ gin.SetMode(gin.TestMode)
+ recorder := httptest.NewRecorder()
+ c, _ := gin.CreateTestContext(recorder)
+ c.Set("status_code_mapping", `{"429":"503"}`)
+
+ writeMidjourneyStatusCode(c, http.StatusTooManyRequests)
+
+ require.Equal(t, http.StatusServiceUnavailable, c.Writer.Status())
+}
+
+func TestRelayMidjourneyImageRejectsUnsignedAndWrongUserURLs(t *testing.T) {
+ gin.SetMode(gin.TestMode)
+ oldDB := model.DB
+ oldSecret := common.CryptoSecret
+ common.CryptoSecret = "midjourney-handler-test-secret"
+ db, err := gorm.Open(sqlite.Open("file:mjproxy_auth?mode=memory&cache=shared"), &gorm.Config{})
+ require.NoError(t, err)
+ model.DB = db
+ require.NoError(t, db.AutoMigrate(&model.Midjourney{}))
+ require.NoError(t, db.Create(&model.Midjourney{
+ Id: 1, UserId: 42, MjId: "mj-private-task", ImageUrl: "https://example.com/private.png",
+ }).Error)
+ t.Cleanup(func() {
+ model.DB = oldDB
+ common.CryptoSecret = oldSecret
+ if sqlDB, dbErr := db.DB(); dbErr == nil {
+ _ = sqlDB.Close()
+ }
+ })
+
+ router := gin.New()
+ router.GET("/mj/image/:id", RelayMidjourneyImage)
+
+ unsigned := httptest.NewRecorder()
+ router.ServeHTTP(unsigned, httptest.NewRequest(http.MethodGet, "/mj/image/mj-private-task", nil))
+ require.Equal(t, http.StatusForbidden, unsigned.Code)
+ require.Contains(t, unsigned.Body.String(), "midjourney_image_authorization_failed")
+
+ now := time.Now()
+ expiresAt := now.Add(service.MidjourneyImageURLTTL).Unix()
+ wrongUserURL := "/mj/image/mj-private-task?uid=43&expires=" +
+ strconv.FormatInt(expiresAt, 10) + "&signature=" +
+ service.SignMidjourneyImageURL("mj-private-task", 43, expiresAt)
+ wrongUser := httptest.NewRecorder()
+ router.ServeHTTP(wrongUser, httptest.NewRequest(http.MethodGet, wrongUserURL, nil))
+ require.Equal(t, http.StatusForbidden, wrongUser.Code)
+ require.Contains(t, wrongUser.Body.String(), "midjourney_image_authorization_failed")
+}
diff --git a/relay/relay_task.go b/relay/relay_task.go
index 34e015b2..5a7b6e43 100644
--- a/relay/relay_task.go
+++ b/relay/relay_task.go
@@ -28,9 +28,84 @@ type TaskSubmitResult struct {
TaskData []byte
Platform constant.TaskPlatform
Quota int
+ Task *model.Task
+ response *taskResponseSnapshot
//PerCallPrice types.PriceData
}
+func (r *TaskSubmitResult) WriteResponse(c *gin.Context) error {
+ if r == nil {
+ return errors.New("task submit result is nil")
+ }
+ return r.response.writeTo(c)
+}
+
+func populateTaskBillingMetadata(task *model.Task, info *relaycommon.RelayInfo) {
+ if task == nil || info == nil {
+ return
+ }
+ deltaSettlementDisabled := false
+ if info.ChannelMeta != nil && info.ChannelType == constant.ChannelTypeDoubaoVideo {
+ deltaSettlementDisabled = info.ChannelOtherSettings.DisableTaskDeltaSettlement
+ }
+ task.PrivateData.BillingSource = info.BillingSource
+ task.PrivateData.BillingRequestId = info.RequestId
+ task.PrivateData.SubscriptionId = info.SubscriptionId
+ task.PrivateData.TokenId = info.TokenId
+ task.PrivateData.NodeName = common.NodeName
+ task.PrivateData.BillingContext = &model.TaskBillingContext{
+ ModelPrice: info.PriceData.ModelPrice,
+ GroupRatio: info.PriceData.GroupRatioInfo.GroupRatio,
+ ModelRatio: info.PriceData.ModelRatio,
+ OtherRatios: info.PriceData.OtherRatios,
+ TaskBilling: info.TaskBilling,
+ OriginModelName: info.OriginModelName,
+ PerCallBilling: common.StringsContains(constant.TaskPricePatches, info.OriginModelName) || info.PriceData.UsePrice || info.TaskBilling != nil,
+ DeltaSettlementDisabled: common.GetPointer(deltaSettlementDisabled),
+ }
+ task.Action = info.Action
+}
+
+func ensureTaskPlaceholder(platform constant.TaskPlatform, info *relaycommon.RelayInfo) (*model.Task, *dto.TaskError) {
+ if info == nil || info.TaskRelayInfo == nil {
+ return nil, service.TaskErrorWrapperLocal(errors.New("task relay metadata is unavailable"), "persist_task_failed", http.StatusInternalServerError)
+ }
+ if info.PersistedTaskID > 0 {
+ var task model.Task
+ if err := model.DB.First(&task, info.PersistedTaskID).Error; err != nil {
+ return nil, TaskPersistenceError(err, "load_persisted_task_failed", "failed to load persisted task")
+ }
+ fresh := model.InitTask(platform, info)
+ task.Platform = platform
+ task.ChannelId = fresh.ChannelId
+ task.Group = fresh.Group
+ task.Properties = fresh.Properties
+ task.PrivateData.Key = fresh.PrivateData.Key
+ task.PrivateData.AwaitingUpstreamID = true
+ task.Quota = info.PriceData.Quota
+ populateTaskBillingMetadata(&task, info)
+ if err := task.Update(); err != nil {
+ return nil, TaskPersistenceError(err, "update_persisted_task_failed", "failed to update persisted task")
+ }
+ return &task, nil
+ }
+
+ task := model.InitTask(platform, info)
+ task.Quota = info.PriceData.Quota
+ task.PrivateData.AwaitingUpstreamID = true
+ populateTaskBillingMetadata(task, info)
+ if err := task.Insert(); err != nil {
+ return nil, TaskPersistenceError(err, "persist_task_failed", "failed to persist task")
+ }
+ info.PersistedTaskID = task.ID
+ return task, nil
+}
+
+func TaskPersistenceError(err error, code string, safeMessage string) *dto.TaskError {
+ common.SysLog(fmt.Sprintf("%s: %s", code, err.Error()))
+ return service.TaskErrorWrapperLocal(errors.New(safeMessage), code, http.StatusInternalServerError)
+}
+
type taskBillingEstimator interface {
EstimateTaskBilling(c *gin.Context, info *relaycommon.RelayInfo) (*types.TaskBillingResult, error)
}
@@ -235,6 +310,10 @@ func RelayTaskSubmit(c *gin.Context, info *relaycommon.RelayInfo) (*TaskSubmitRe
return nil, service.TaskErrorFromAPIError(apiErr)
}
}
+ task, taskErr := ensureTaskPlaceholder(platform, info)
+ if taskErr != nil {
+ return nil, taskErr
+ }
// 9. 发送请求
resp, err := adaptor.DoRequest(c, info, requestBody)
@@ -245,23 +324,39 @@ func RelayTaskSubmit(c *gin.Context, info *relaycommon.RelayInfo) (*TaskSubmitRe
return nil, service.TaskErrorWrapper(err, "do_request_failed", http.StatusInternalServerError)
}
if resp != nil && resp.StatusCode != http.StatusOK {
+ if resp.Body != nil {
+ defer resp.Body.Close()
+ }
responseBody, _ := io.ReadAll(resp.Body)
- return nil, service.TaskErrorWrapper(fmt.Errorf("%s", string(responseBody)), "fail_to_fetch_task", resp.StatusCode)
- }
-
- // 10. 返回 OtherRatios 给下游（header 必须在 DoResponse 写 body 之前设置）
- otherRatios := info.PriceData.OtherRatios
- if otherRatios == nil {
- otherRatios = map[string]float64{}
- }
- ratiosJSON, _ := common.Marshal(otherRatios)
- c.Header("X-Max-Api-Other-Ratios", string(ratiosJSON))
- c.Header("X-New-Api-Other-Ratios", string(ratiosJSON))
-
- // 11. 解析响应
- upstreamTaskID, taskData, taskErr := adaptor.DoResponse(c, resp, info)
+ taskErr := service.TaskErrorWrapper(fmt.Errorf("%s", string(responseBody)), "fail_to_fetch_task", resp.StatusCode)
+ return nil, mapUpstreamTaskError(c, taskErr)
+ }
+ if resp == nil {
+ return nil, service.TaskErrorWrapperLocal(errors.New("task upstream response is nil"), "empty_upstream_response", http.StatusBadGateway)
+ }
+ info.UpstreamTaskResponseReceived = true
+
+ // 10-11. Buffer the adapter response until the task row and settlement
+ // intent are durable. Adapters may write the body while parsing.
+ var upstreamTaskID string
+ var taskData []byte
+ var responseSnapshot *taskResponseSnapshot
+ func() {
+ originalWriter := c.Writer
+ bufferedWriter := newTaskResponseBuffer(originalWriter)
+ c.Writer = bufferedWriter
+ defer func() { c.Writer = originalWriter }()
+
+ setTaskOtherRatioHeaders(bufferedWriter.Header(), info.PriceData.OtherRatios)
+
+ upstreamTaskID, taskData, taskErr = adaptor.DoResponse(c, resp, info)
+ responseSnapshot = bufferedWriter.snapshot()
+ }()
if taskErr != nil {
- return nil, taskErr
+ return nil, mapUpstreamTaskError(c, taskErr)
+ }
+ if strings.TrimSpace(upstreamTaskID) == "" {
+ return nil, service.TaskErrorWrapperLocal(errors.New("task upstream response did not contain a task id"), "missing_upstream_task_id", http.StatusBadGateway)
}
// 11. 提交后计费调整：让适配器根据上游实际返回调整 OtherRatios
@@ -274,15 +369,49 @@ func RelayTaskSubmit(c *gin.Context, info *relaycommon.RelayInfo) (*TaskSubmitRe
info.PriceData.Quota = finalQuota
}
}
+ if responseSnapshot != nil {
+ setTaskOtherRatioHeaders(responseSnapshot.header, info.PriceData.OtherRatios)
+ }
+ task.PrivateData.UpstreamTaskID = upstreamTaskID
+ task.PrivateData.AwaitingUpstreamID = false
+ task.Quota = finalQuota
+ task.Data = taskData
+ populateTaskBillingMetadata(task, info)
return &TaskSubmitResult{
UpstreamTaskID: upstreamTaskID,
TaskData: taskData,
Platform: platform,
Quota: finalQuota,
+ Task: task,
+ response: responseSnapshot,
}, nil
}
+func setTaskOtherRatioHeaders(header http.Header, otherRatios map[string]float64) {
+ if header == nil {
+ return
+ }
+ if otherRatios == nil {
+ otherRatios = map[string]float64{}
+ }
+ ratiosJSON, err := common.Marshal(otherRatios)
+ if err != nil {
+ common.SysLog("failed to marshal task ratio headers: " + err.Error())
+ ratiosJSON = []byte("{}")
+ }
+ header.Set("X-Max-Api-Other-Ratios", string(ratiosJSON))
+ header.Set("X-New-Api-Other-Ratios", string(ratiosJSON))
+}
+
+// mapUpstreamTaskError is only called after an upstream response is available.
+func mapUpstreamTaskError(c *gin.Context, taskErr *dto.TaskError) *dto.TaskError {
+ if c != nil {
+ service.ResetTaskStatusCode(taskErr, c.GetString("status_code_mapping"))
+ }
+ return taskErr
+}
+
func estimateTaskBilling(c *gin.Context, info *relaycommon.RelayInfo, adaptor channel.TaskAdaptor, platform constant.TaskPlatform) (*types.TaskBillingResult, error) {
if estimator, ok := adaptor.(taskBillingEstimator); ok {
taskBilling, err := estimator.EstimateTaskBilling(c, info)
@@ -619,7 +748,10 @@ func tryRealtimeFetch(task *model.Task, isOpenAIVideoAPI bool) []byte {
task.PrivateData.ResultURL = taskcommon.BuildProxyURL(task.TaskID)
}
- if !snap.Equal(task.Snapshot()) {
+ // Terminal transitions must go through the polling settlement path. The
+ // realtime fetch may report terminal data to this caller, but must not make
+ // that state durable without the matching refund/final-settlement intent.
+ if !isTerminalTaskStatus(task.Status) && !snap.Equal(task.Snapshot()) {
_, _ = task.UpdateWithStatus(snap.Status)
}
@@ -645,6 +777,10 @@ func tryRealtimeFetch(task *model.Task, isOpenAIVideoAPI bool) []byte {
return respBody
}
+func isTerminalTaskStatus(status model.TaskStatus) bool {
+ return status == model.TaskStatusSuccess || status == model.TaskStatusFailure
+}
+
// detectVideoFormat 从 Gemini/Vertex 原始响应中探测视频格式
func detectVideoFormat(rawBody []byte) string {
var raw map[string]any
diff --git a/relay/relay_task_test.go b/relay/relay_task_test.go
index 7dbed281..958e672c 100644
--- a/relay/relay_task_test.go
+++ b/relay/relay_task_test.go
@@ -1,6 +1,7 @@
package relay
import (
+ "errors"
"io"
"net/http"
"net/http/httptest"
@@ -9,6 +10,8 @@ import (
"github.com/MAX-API-Next/MAX-API/common"
"github.com/MAX-API-Next/MAX-API/constant"
+ "github.com/MAX-API-Next/MAX-API/dto"
+ "github.com/MAX-API-Next/MAX-API/model"
"github.com/MAX-API-Next/MAX-API/relay/channel/task/doubao"
relaycommon "github.com/MAX-API-Next/MAX-API/relay/common"
"github.com/MAX-API-Next/MAX-API/setting/config"
@@ -230,3 +233,93 @@ func TestPrepareTaskSubmitRequestBodyParamOverrideReturnErrorIsLocal(t *testing.
assert.Equal(t, "forced_bad_request", taskErr.Code)
assert.Equal(t, "forced bad request by param override", taskErr.Message)
}
+
+func TestMapUpstreamTaskErrorAppliesMappingWithoutChangingLocalClassification(t *testing.T) {
+ gin.SetMode(gin.TestMode)
+ c, _ := gin.CreateTestContext(httptest.NewRecorder())
+ c.Set("status_code_mapping", `{"400":"503"}`)
+ taskErr := &dto.TaskError{
+ StatusCode: http.StatusBadRequest,
+ LocalError: true,
+ }
+
+ got := mapUpstreamTaskError(c, taskErr)
+
+ require.Same(t, taskErr, got)
+ assert.Equal(t, http.StatusServiceUnavailable, got.StatusCode)
+ assert.Equal(t, http.StatusBadRequest, got.UpstreamStatusCode)
+ assert.True(t, got.LocalError)
+}
+
+func TestTaskResponseBufferDefersHeadersStatusAndBody(t *testing.T) {
+ gin.SetMode(gin.TestMode)
+ recorder := httptest.NewRecorder()
+ c, _ := gin.CreateTestContext(recorder)
+ originalWriter := c.Writer
+ bufferedWriter := newTaskResponseBuffer(originalWriter)
+ c.Writer = bufferedWriter
+
+ c.Header("X-Task-Result", "buffered")
+ c.Status(http.StatusAccepted)
+ _, err := c.Writer.WriteString(`{"id":"task_public"}`)
+ require.NoError(t, err)
+
+ require.False(t, originalWriter.Written())
+ require.Empty(t, recorder.Body.String())
+ require.Empty(t, recorder.Header().Get("X-Task-Result"))
+
+ snapshot := bufferedWriter.snapshot()
+ c.Writer = originalWriter
+ require.NoError(t, snapshot.writeTo(c))
+ require.Equal(t, http.StatusAccepted, recorder.Code)
+ require.Equal(t, "buffered", recorder.Header().Get("X-Task-Result"))
+ require.JSONEq(t, `{"id":"task_public"}`, recorder.Body.String())
+}
+
+func TestTaskResponseSnapshotClearsHeadersRemovedFromBuffer(t *testing.T) {
+ gin.SetMode(gin.TestMode)
+ recorder := httptest.NewRecorder()
+ c, _ := gin.CreateTestContext(recorder)
+ c.Header("X-Stale", "stale")
+ originalWriter := c.Writer
+ bufferedWriter := newTaskResponseBuffer(originalWriter)
+ c.Writer = bufferedWriter
+
+ bufferedWriter.Header().Del("X-Stale")
+ bufferedWriter.Header().Set("X-Fresh", "fresh")
+ _, err := c.Writer.WriteString(`{"ok":true}`)
+ require.NoError(t, err)
+
+ snapshot := bufferedWriter.snapshot()
+ c.Writer = originalWriter
+ require.NoError(t, snapshot.writeTo(c))
+ require.Empty(t, recorder.Header().Get("X-Stale"))
+ require.Equal(t, "fresh", recorder.Header().Get("X-Fresh"))
+}
+
+func TestSetTaskOtherRatioHeadersUsesFinalRatios(t *testing.T) {
+ header := http.Header{}
+
+ setTaskOtherRatioHeaders(header, map[string]float64{"duration": 1.25})
+ setTaskOtherRatioHeaders(header, map[string]float64{"duration": 2})
+
+ require.JSONEq(t, `{"duration":2}`, header.Get("X-Max-Api-Other-Ratios"))
+ require.JSONEq(t, `{"duration":2}`, header.Get("X-New-Api-Other-Ratios"))
+}
+
+func TestTaskPersistenceErrorReturnsSafeClientMessage(t *testing.T) {
+ taskErr := TaskPersistenceError(errors.New("sql: failed near secret_table"), "persist_task_failed", "failed to persist task")
+
+ require.NotNil(t, taskErr)
+ assert.Equal(t, "persist_task_failed", taskErr.Code)
+ assert.Equal(t, "failed to persist task", taskErr.Message)
+ assert.NotContains(t, taskErr.Message, "secret_table")
+ assert.NotContains(t, taskErr.Error.Error(), "secret_table")
+}
+
+func TestTerminalTaskStatusClassification(t *testing.T) {
+ require.True(t, isTerminalTaskStatus(model.TaskStatusSuccess))
+ require.True(t, isTerminalTaskStatus(model.TaskStatusFailure))
+ require.False(t, isTerminalTaskStatus(model.TaskStatusInProgress))
+ require.False(t, isTerminalTaskStatus(model.TaskStatusNotStart))
+}
diff --git a/relay/task_response_buffer.go b/relay/task_response_buffer.go
new file mode 100644
index 00000000..866dd456
--- /dev/null
+++ b/relay/task_response_buffer.go
@@ -0,0 +1,100 @@
+package relay
+
+import (
+ "bufio"
+ "bytes"
+ "errors"
+ "net"
+ "net/http"
+
+ "github.com/gin-gonic/gin"
+)
+
+type taskResponseSnapshot struct {
+ header http.Header
+ status int
+ body []byte
+}
+
+func (s *taskResponseSnapshot) writeTo(c *gin.Context) error {
+ if s == nil || c == nil {
+ return errors.New("task response snapshot is unavailable")
+ }
+ for key := range c.Writer.Header() {
+ delete(c.Writer.Header(), key)
+ }
+ for key, values := range s.header {
+ c.Writer.Header()[key] = append([]string(nil), values...)
+ }
+ status := s.status
+ if status == 0 {
+ status = http.StatusOK
+ }
+ c.Writer.WriteHeader(status)
+ _, err := c.Writer.Write(s.body)
+ return err
+}
+
+type taskResponseBuffer struct {
+ original gin.ResponseWriter
+ header http.Header
+ body bytes.Buffer
+ status int
+ written bool
+}
+
+func newTaskResponseBuffer(original gin.ResponseWriter) *taskResponseBuffer {
+ header := make(http.Header, len(original.Header()))
+ for key, values := range original.Header() {
+ header[key] = append([]string(nil), values...)
+ }
+ return &taskResponseBuffer{original: original, header: header, status: http.StatusOK}
+}
+
+func (w *taskResponseBuffer) Header() http.Header { return w.header }
+
+func (w *taskResponseBuffer) Write(data []byte) (int, error) {
+ if !w.written {
+ w.WriteHeaderNow()
+ }
+ return w.body.Write(data)
+}
+
+func (w *taskResponseBuffer) WriteString(data string) (int, error) {
+ if !w.written {
+ w.WriteHeaderNow()
+ }
+ return w.body.WriteString(data)
+}
+
+func (w *taskResponseBuffer) WriteHeader(status int) {
+ if w.written {
+ return
+ }
+ w.status = status
+}
+
+func (w *taskResponseBuffer) WriteHeaderNow() { w.written = true }
+func (w *taskResponseBuffer) Status() int { return w.status }
+func (w *taskResponseBuffer) Size() int { return w.body.Len() }
+func (w *taskResponseBuffer) Written() bool { return w.written }
+func (w *taskResponseBuffer) Flush() {}
+func (w *taskResponseBuffer) Pusher() http.Pusher {
+ return w.original.Pusher()
+}
+func (w *taskResponseBuffer) CloseNotify() <-chan bool { return w.original.CloseNotify() }
+func (w *taskResponseBuffer) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+ return nil, nil, errors.New("task submission responses cannot hijack the connection")
+}
+
+func (w *taskResponseBuffer) snapshot() *taskResponseSnapshot {
+ header := make(http.Header, len(w.header))
+ for key, values := range w.header {
+ header[key] = append([]string(nil), values...)
+ }
+ return &taskResponseSnapshot{
+ header: header,
+ status: w.status,
+ body: append([]byte(nil), w.body.Bytes()...),
+ }
+}
diff --git a/router/api-router.go b/router/api-router.go
index 3f0e4a8b..55caf4d4 100644
--- a/router/api-router.go
+++ b/router/api-router.go
@@ -85,7 +85,7 @@ func SetApiRouter(router *gin.Engine) {
selfRoute.GET("/self", controller.GetSelf)
selfRoute.GET("/models", controller.GetUserModels)
selfRoute.PUT("/self", middleware.CriticalRateLimit(), controller.UpdateSelf)
- selfRoute.DELETE("/self", controller.DeleteSelf)
+ selfRoute.DELETE("/self", middleware.CriticalRateLimit(), middleware.DisableCache(), middleware.SecureVerificationRequired("account_delete"), controller.DeleteSelf)
selfRoute.POST("/token", middleware.CriticalRateLimit(), middleware.DisableCache(), middleware.SecureVerificationRequired("access_token"), controller.GenerateAccessToken)
selfRoute.GET("/passkey", controller.PasskeyStatus)
selfRoute.POST("/passkey/register/begin", controller.PasskeyRegisterBegin)
diff --git a/router/api_router_test.go b/router/api_router_test.go
index 10ac9115..10c08c90 100644
--- a/router/api_router_test.go
+++ b/router/api_router_test.go
@@ -190,10 +190,11 @@ func TestSecureVerificationOpenAPIIncludesScopeAndRequestBody(t *testing.T) {
}
schema, schemaOK := value["schema"].(map[string]any)
values, valuesOK := schema["enum"].([]any)
- return schemaOK && valuesOK && len(values) == 1 && values[0] == "access_token"
+ return schemaOK && valuesOK && len(values) == 2 &&
+ values[0] == "access_token" && values[1] == "account_delete"
}
return false
- }, "expected scope query parameter for GET /api/verify/methods")
+ }, "expected supported scope query parameter for GET /api/verify/methods")
verifyOperation := document.Paths["/api/verify"]["post"]
requestBody, ok := verifyOperation["requestBody"].(map[string]any)
diff --git a/service/billing.go b/service/billing.go
index 6c62de58..cd3ac711 100644
--- a/service/billing.go
+++ b/service/billing.go
@@ -96,7 +96,7 @@ func SettleBilling(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, actualQuo
// 回退：无 BillingSession 时使用旧路径
quotaDelta := actualQuota - relayInfo.FinalPreConsumedQuota
if quotaDelta != 0 {
- return PostConsumeQuota(relayInfo, quotaDelta, relayInfo.FinalPreConsumedQuota, true)
+ return PostConsumeQuotaOnce(relayInfo, "finalize", quotaDelta, relayInfo.FinalPreConsumedQuota, true)
}
return nil
}
diff --git a/service/billing_session.go b/service/billing_session.go
index 2faf7fed..c0869fbb 100644
--- a/service/billing_session.go
+++ b/service/billing_session.go
@@ -1,6 +1,7 @@
package service
import (
+ "errors"
"fmt"
"net/http"
"strings"
@@ -13,7 +14,6 @@ import (
relaycommon "github.com/MAX-API-Next/MAX-API/relay/common"
"github.com/MAX-API-Next/MAX-API/types"
- "github.com/bytedance/gopkg/util/gopool"
"github.com/gin-gonic/gin"
)
@@ -34,22 +34,41 @@ type BillingSession struct {
appliedFundingDelta int64 // 已提交、尚未完成令牌对账的资金差额
compensationFailed bool // 资金补偿返回错误且结果不确定；后续仅重试令牌
fundingReconcilePending bool // 已知部分补偿已生效；重试令牌前先补齐资金差额
+ reserveRevision int64 // 单次请求内累计预留的单调修订号
settled bool // Settle 全部完成（资金 + 令牌）
refunded bool // Refund 已调用
+ refundInFlight bool
+ settleInFlight bool
mu sync.Mutex
}
+type SettlementPreparer interface {
+ PrepareSettlement(actualQuota int) (*model.BillingSettlementInput, error)
+}
+
+var _ SettlementPreparer = (*BillingSession)(nil)
+
+type billingSettleIntent struct {
+ input model.BillingSettlementInput
+}
+
// Settle 根据实际消耗额度进行结算。
// 资金来源和令牌额度分两步提交。若令牌调整失败，会优先反向补偿已提交的
// 资金差额。已确认的部分补偿会记录剩余差额并在重试令牌前补齐；补偿
// 返回错误时结果不明确，不会自动重复该非幂等资金操作，后续仅重试令牌。
func (s *BillingSession) Settle(actualQuota int) error {
- s.mu.Lock()
- defer s.mu.Unlock()
+ return s.settleWithEffect(actualQuota, nil)
+}
+
+func (s *BillingSession) SettleWithEffect(actualQuota int, effect *model.BillingSettlementEffect) error {
+ return s.settleWithEffect(actualQuota, effect)
+}
+
+func (s *BillingSession) settleWithEffect(actualQuota int, effect *model.BillingSettlementEffect) error {
var lastErr error
const maxAttempts = 3
for attempt := 0; attempt < maxAttempts; attempt++ {
- if err := s.settleLocked(actualQuota); err == nil {
+ if err := s.settleAttempt(actualQuota, effect); err == nil {
return nil
} else {
lastErr = err
@@ -61,15 +80,72 @@ func (s *BillingSession) Settle(actualQuota int) error {
return fmt.Errorf("billing settlement failed after %d attempts: %w", maxAttempts, lastErr)
}
-func (s *BillingSession) settleLocked(actualQuota int) error {
- if s.settled {
- return nil
+func (s *BillingSession) settleAttempt(actualQuota int, effect *model.BillingSettlementEffect) error {
+ intent, ok, err := s.beginSettleAttempt(actualQuota, effect)
+ if err != nil || !ok {
+ return err
+ }
+ return s.applyDurableSettleIntent(intent)
+}
+
+func (s *BillingSession) beginSettleAttempt(actualQuota int, effect *model.BillingSettlementEffect) (*billingSettleIntent, bool, error) {
+ s.mu.Lock()
+ defer s.mu.Unlock()
+ if s.settled || s.refunded {
+ return nil, false, nil
+ }
+ if s.refundInFlight {
+ return nil, false, errors.New("billing refund is already in progress")
+ }
+ if s.settleInFlight {
+ return nil, false, errors.New("billing settlement is already in progress")
}
delta := actualQuota - s.preConsumedQuota
- if delta == 0 {
+ if delta == 0 && effect == nil {
s.settled = true
- return nil
+ return nil, false, nil
+ }
+ if input, ok := s.prepareDurableSettlementLocked(actualQuota, effect); ok {
+ s.settleInFlight = true
+ return &billingSettleIntent{input: *input}, true, nil
+ }
+ if effect != nil {
+ return nil, false, errors.New("billing settlement effects require a durable funding source and request id")
+ }
+ return nil, false, s.settleNonDurableLocked(delta)
+}
+
+func (s *BillingSession) applyDurableSettleIntent(intent *billingSettleIntent) error {
+ applied, _, err := model.ApplyBillingSettlementOnce(intent.input)
+ if err != nil {
+ s.finishDurableSettleIntent(false, 0)
+ return err
+ }
+
+ var effectErr error
+ if intent.input.Effect != nil {
+ effectErr = model.ProcessBillingSettlementEffect(intent.input.OperationKey)
+ }
+ s.finishDurableSettleIntent(true, applied)
+ return effectErr
+}
+
+func (s *BillingSession) finishDurableSettleIntent(applied bool, appliedDelta int64) {
+ s.mu.Lock()
+ defer s.mu.Unlock()
+ s.settleInFlight = false
+ if !applied {
+ return
}
+ s.appliedFundingDelta = appliedDelta
+ s.fundingSettled = true
+ if s.funding.Source() == BillingSourceSubscription {
+ s.relayInfo.SubscriptionPostDelta += appliedDelta
+ }
+ s.settled = true
+}
+
+func (s *BillingSession) settleNonDurableLocked(delta int) error {
// 1) 调整资金来源（仅在尚未提交时执行，防止重复调用）
if !s.fundingSettled {
applied, err := s.funding.Settle(delta)
@@ -155,48 +231,171 @@ func (s *BillingSession) settleLocked(actualQuota int) error {
return nil
}
-// Refund 退还所有预扣费，幂等安全，异步执行。
-func (s *BillingSession) Refund(c *gin.Context) {
+// PrepareSettlement returns the exact durable request-finalize intent without
+// applying it. Task submission uses this to commit upstream identity and the
+// billing intent atomically before releasing the success response.
+func (s *BillingSession) PrepareSettlement(actualQuota int) (*model.BillingSettlementInput, error) {
+ s.mu.Lock()
+ defer s.mu.Unlock()
+ if s.settled || s.refunded {
+ return nil, nil
+ }
+ if actualQuota == s.preConsumedQuota {
+ input, ok := s.prepareDurableSettlementLocked(actualQuota, nil)
+ if !ok {
+ return nil, nil
+ }
+ return input, nil
+ }
+ input, ok := s.prepareDurableSettlementLocked(actualQuota, nil)
+ if !ok {
+ return nil, errors.New("billing settlement intent requires a durable funding source and request id")
+ }
+ return input, nil
+}
+
+func (s *BillingSession) prepareDurableSettlementLocked(actualQuota int, effect *model.BillingSettlementEffect) (*model.BillingSettlementInput, bool) {
+ if s.relayInfo == nil || s.relayInfo.RequestId == "" {
+ return nil, false
+ }
+ source, userID, subscriptionID, ok := durableFundingIdentity(s.funding)
+ if !ok {
+ return nil, false
+ }
+ delta := actualQuota - s.preConsumedQuota
+ tokenDelta := int64(delta)
+ if s.relayInfo.IsPlayground {
+ tokenDelta = 0
+ }
+ return &model.BillingSettlementInput{
+ OperationKey: "request:" + s.relayInfo.RequestId + ":finalize",
+ Source: source,
+ UserID: userID,
+ SubscriptionID: subscriptionID,
+ TokenID: s.relayInfo.TokenId,
+ TokenKey: s.relayInfo.TokenKey,
+ FundingDelta: int64(delta),
+ TokenDelta: tokenDelta,
+ SubscriptionPreConsumeRequestID: subscriptionPreConsumeRequestID(s.funding),
+ Effect: effect,
+ }, true
+}
+
+func durableFundingIdentity(funding FundingSource) (source string, userID int, subscriptionID int, ok bool) {
+ switch value := funding.(type) {
+ case *WalletFunding:
+ return model.BillingSettlementSourceWallet, value.userId, 0, true
+ case *SubscriptionFunding:
+ return model.BillingSettlementSourceSubscription, value.userId, value.subscriptionId, true
+ default:
+ return "", 0, 0, false
+ }
+}
+
+type billingRefundIntent struct {
+ input model.BillingSettlementInput
+ userID int
+ requestID string
+ tokenConsumed int
+ fundingSource string
+}
+
+func (s *BillingSession) prepareRefundIntent() (*billingRefundIntent, bool) {
s.mu.Lock()
- if s.settled || s.refunded || !s.needsRefundLocked() {
- s.mu.Unlock()
+ defer s.mu.Unlock()
+ if s.settled || s.refunded || s.refundInFlight || s.settleInFlight || !s.needsRefundLocked() {
+ return nil, false
+ }
+ refundFunding := s.refundFundingAmountLocked()
+ refundToken := int64(s.tokenConsumed)
+ if s.relayInfo.IsPlayground {
+ refundToken = 0
+ }
+ requestID := s.relayInfo.RequestId
+ funding := s.funding
+ if requestID == "" || refundFunding <= 0 {
+ common.SysLog(fmt.Sprintf("billing refund requires manual review (userId=%d, requestId=%q, funding=%d)", s.relayInfo.UserId, requestID, refundFunding))
+ return nil, false
+ }
+ source, userID, subscriptionID, ok := durableFundingIdentity(funding)
+ if !ok {
+ common.SysLog(fmt.Sprintf("billing refund skipped for non-durable funding source (userId=%d, requestId=%s)", s.relayInfo.UserId, requestID))
+ return nil, false
+ }
+
+ s.refundInFlight = true
+ return &billingRefundIntent{
+ input: model.BillingSettlementInput{
+ OperationKey: "request:" + requestID + ":finalize",
+ Source: source,
+ UserID: userID,
+ SubscriptionID: subscriptionID,
+ TokenID: s.relayInfo.TokenId,
+ TokenKey: s.relayInfo.TokenKey,
+ FundingDelta: -refundFunding,
+ TokenDelta: -refundToken,
+ SubscriptionPreConsumeRequestID: subscriptionPreConsumeRequestID(funding),
+ FinalizeSubscriptionPreConsume: source == model.BillingSettlementSourceSubscription,
+ AllowMissingToken: true,
+ },
+ userID: s.relayInfo.UserId,
+ requestID: requestID,
+ tokenConsumed: s.tokenConsumed,
+ fundingSource: funding.Source(),
+ }, true
+}
+
+func (s *BillingSession) finishRefundIntent(applied bool) {
+ s.mu.Lock()
+ defer s.mu.Unlock()
+ s.refundInFlight = false
+ if !applied {
return
}
s.refunded = true
- s.mu.Unlock()
+ s.tokenConsumed = 0
+ s.extraReserved = 0
+ s.preConsumedQuota = 0
+}
+
+// Refund 退还所有预扣费。普通钱包/订阅会话必须通过持久化结算记录执行；
+// 没有稳定 request id 的自定义资金来源不会被自动修改，避免无法证明幂等时重复退款。
+func (s *BillingSession) Refund(c *gin.Context) {
+ intent, ok := s.prepareRefundIntent()
+ if !ok {
+ return
+ }
logger.LogInfo(c, fmt.Sprintf("用户 %d 请求失败, 返还预扣费（token_quota=%s, funding=%s）",
- s.relayInfo.UserId,
- logger.FormatQuota(s.tokenConsumed),
- s.funding.Source(),
+ intent.userID,
+ logger.FormatQuota(intent.tokenConsumed),
+ intent.fundingSource,
))
+ _, _, err := model.ApplyBillingSettlementOnce(intent.input)
+ if err != nil {
+ common.SysLog(fmt.Sprintf("billing refund remains pending/manual (userId=%d, requestId=%s): %s", intent.userID, intent.requestID, err.Error()))
+ s.finishRefundIntent(false)
+ return
+ }
+ s.finishRefundIntent(true)
+}
- // 复制需要的值到闭包中
- tokenId := s.relayInfo.TokenId
- tokenKey := s.relayInfo.TokenKey
- isPlayground := s.relayInfo.IsPlayground
- tokenConsumed := s.tokenConsumed
- extraReserved := s.extraReserved
- subscriptionId := s.relayInfo.SubscriptionId
- funding := s.funding
+func (s *BillingSession) refundFundingAmountLocked() int64 {
+ switch funding := s.funding.(type) {
+ case *WalletFunding:
+ return int64(funding.consumed)
+ case *SubscriptionFunding:
+ return funding.preConsumed + int64(s.extraReserved)
+ default:
+ return int64(s.preConsumedQuota)
+ }
+}
- gopool.Go(func() {
- // 1) 退还资金来源
- if err := funding.Refund(); err != nil {
- common.SysLog("error refunding billing source: " + err.Error())
- }
- if extraReserved > 0 && funding.Source() == BillingSourceSubscription && subscriptionId > 0 {
- if _, err := model.PostConsumeUserSubscriptionDelta(subscriptionId, -int64(extraReserved)); err != nil {
- common.SysLog("error refunding subscription extra reserved quota: " + err.Error())
- }
- }
- // 2) 退还令牌额度
- if tokenConsumed > 0 && !isPlayground {
- if err := model.IncreaseTokenQuota(tokenId, tokenKey, tokenConsumed); err != nil {
- common.SysLog("error refunding token quota: " + err.Error())
- }
- }
- })
+func subscriptionPreConsumeRequestID(funding FundingSource) string {
+ if subscription, ok := funding.(*SubscriptionFunding); ok {
+ return subscription.requestId
+ }
+ return ""
}
// NeedsRefund 返回是否存在需要退还的预扣状态。
@@ -207,7 +406,7 @@ func (s *BillingSession) NeedsRefund() bool {
}
func (s *BillingSession) needsRefundLocked() bool {
- if s.settled || s.refunded || s.fundingSettled {
+ if s.settled || s.refunded || s.fundingSettled || s.settleInFlight {
// fundingSettled 时资金来源已提交结算，不能再退预扣费
return false
}
@@ -239,12 +438,43 @@ func (s *BillingSession) Reserve(targetQuota int) error {
return nil
}
- if err := s.reserveFunding(delta); err != nil {
- return err
+ atomicReserved := false
+ if s.relayInfo.RequestId != "" {
+ if source, userID, subscriptionID, ok := durableFundingIdentity(s.funding); ok {
+ tokenDelta := int64(delta)
+ if s.relayInfo.IsPlayground {
+ tokenDelta = 0
+ }
+ nextRevision := s.reserveRevision + 1
+ if _, _, err := model.ApplyBillingSettlementOnce(model.BillingSettlementInput{
+ OperationKey: fmt.Sprintf("request:%s:reserve:%d", s.relayInfo.RequestId, nextRevision),
+ Source: source,
+ UserID: userID,
+ SubscriptionID: subscriptionID,
+ TokenID: s.relayInfo.TokenId,
+ TokenKey: s.relayInfo.TokenKey,
+ FundingDelta: int64(delta),
+ TokenDelta: tokenDelta,
+ SubscriptionPreConsumeRequestID: subscriptionPreConsumeRequestID(s.funding),
+ ManualOnFailure: true,
+ }); err != nil {
+ return err
+ }
+ if wallet, ok := s.funding.(*WalletFunding); ok {
+ wallet.consumed += delta
+ }
+ s.reserveRevision = nextRevision
+ atomicReserved = true
+ }
}
- if err := s.reserveToken(delta); err != nil {
- s.rollbackFundingReserve(delta)
- return err
+ if !atomicReserved {
+ if err := s.reserveFunding(delta); err != nil {
+ return err
+ }
+ if err := s.reserveToken(delta); err != nil {
+ s.rollbackFundingReserve(delta)
+ return err
+ }
}
s.preConsumedQuota += delta
@@ -272,30 +502,84 @@ func (s *BillingSession) preConsume(c *gin.Context, quota int) *types.MaxAPIErro
logger.LogInfo(c, fmt.Sprintf("用户 %d 需要预扣费 %s (funding=%s)", s.relayInfo.UserId, logger.FormatQuota(effectiveQuota), s.funding.Source()))
}
- // ---- 1) 预扣令牌额度 ----
- if effectiveQuota > 0 {
+ atomicPreConsumed := false
+ if !s.relayInfo.IsPlayground {
+ switch funding := s.funding.(type) {
+ case *WalletFunding:
+ if s.relayInfo.RequestId == "" {
+ if effectiveQuota > 0 {
+ return mapAtomicPreConsumeError(errors.New("wallet pre-consume requires a stable request id"))
+ }
+ break
+ }
+ applied, _, err := model.ApplyBillingSettlementOnce(model.BillingSettlementInput{
+ OperationKey: "request:" + s.relayInfo.RequestId + ":pre-consume",
+ Source: model.BillingSettlementSourceWallet,
+ UserID: s.relayInfo.UserId,
+ TokenID: s.relayInfo.TokenId,
+ TokenKey: s.relayInfo.TokenKey,
+ FundingDelta: int64(effectiveQuota),
+ TokenDelta: int64(effectiveQuota),
+ ManualOnFailure: true,
+ PreConsumeRequestID: s.relayInfo.RequestId,
+ PreConsumeModelName: s.relayInfo.OriginModelName,
+ PreConsumeRequestedQuota: int64(quota),
+ PreConsumeEffectiveQuota: int64(effectiveQuota),
+ })
+ if err != nil {
+ return mapAtomicPreConsumeError(err)
+ }
+ if applied != int64(effectiveQuota) {
+ return mapAtomicPreConsumeError(fmt.Errorf("wallet pre-consume applied unexpected amount: requested=%d applied=%d", effectiveQuota, applied))
+ }
+ funding.consumed = int(applied)
+ s.tokenConsumed = int(applied)
+ atomicPreConsumed = true
+ case *SubscriptionFunding:
+ if effectiveQuota <= 0 {
+ break
+ }
+ res, err := model.PreConsumeTokenAndUserSubscription(
+ s.relayInfo.RequestId,
+ s.relayInfo.UserId,
+ s.relayInfo.TokenId,
+ s.relayInfo.TokenKey,
+ funding.modelName,
+ 0,
+ int64(effectiveQuota),
+ )
+ if err != nil {
+ return mapAtomicPreConsumeError(err)
+ }
+ funding.applyPreConsumeResult(res)
+ s.tokenConsumed = effectiveQuota
+ atomicPreConsumed = true
+ }
+ }
+
+ // Custom funding sources and playground requests retain their existing
+ // hooks; normal wallet/subscription requests use the atomic branches above.
+ if !atomicPreConsumed && effectiveQuota > 0 {
if err := PreConsumeTokenQuota(s.relayInfo, effectiveQuota); err != nil {
return types.NewErrorWithStatusCode(err, types.ErrorCodePreConsumeTokenQuotaFailed, http.StatusForbidden, types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog())
}
- s.tokenConsumed = effectiveQuota
+ if !s.relayInfo.IsPlayground {
+ s.tokenConsumed = effectiveQuota
+ }
}
- // ---- 2) 预扣资金来源 ----
- if err := s.funding.PreConsume(effectiveQuota); err != nil {
- // 预扣费失败，回滚令牌额度
- if s.tokenConsumed > 0 && !s.relayInfo.IsPlayground {
- if rollbackErr := model.IncreaseTokenQuota(s.relayInfo.TokenId, s.relayInfo.TokenKey, s.tokenConsumed); rollbackErr != nil {
- common.SysLog(fmt.Sprintf("error rolling back token quota (userId=%d, tokenId=%d, amount=%d, fundingErr=%s): %s",
- s.relayInfo.UserId, s.relayInfo.TokenId, s.tokenConsumed, err.Error(), rollbackErr.Error()))
+ if !atomicPreConsumed {
+ if err := s.funding.PreConsume(effectiveQuota); err != nil {
+ // 预扣费失败，回滚令牌额度
+ if s.tokenConsumed > 0 && !s.relayInfo.IsPlayground {
+ if rollbackErr := model.IncreaseTokenQuota(s.relayInfo.TokenId, s.relayInfo.TokenKey, s.tokenConsumed); rollbackErr != nil {
+ common.SysLog(fmt.Sprintf("error rolling back token quota (userId=%d, tokenId=%d, amount=%d, fundingErr=%s): %s",
+ s.relayInfo.UserId, s.relayInfo.TokenId, s.tokenConsumed, err.Error(), rollbackErr.Error()))
+ }
+ s.tokenConsumed = 0
}
- s.tokenConsumed = 0
- }
- // TODO: model 层应定义哨兵错误（如 ErrNoActiveSubscription），用 errors.Is 替代字符串匹配
- errMsg := err.Error()
- if strings.Contains(errMsg, "no active subscription") || strings.Contains(errMsg, "subscription quota insufficient") {
- return types.NewErrorWithStatusCode(fmt.Errorf("订阅额度不足或未配置订阅: %s", errMsg), types.ErrorCodeInsufficientUserQuota, http.StatusForbidden, types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog())
+ return mapAtomicPreConsumeError(err)
}
- return types.NewError(err, types.ErrorCodeUpdateDataError, types.ErrOptionWithSkipRetry())
}
s.preConsumedQuota = effectiveQuota
@@ -306,6 +590,23 @@ func (s *BillingSession) preConsume(c *gin.Context, quota int) *types.MaxAPIErro
return nil
}
+func mapAtomicPreConsumeError(err error) *types.MaxAPIError {
+ if errors.Is(err, model.ErrTokenQuotaInsufficient) {
+ return types.NewErrorWithStatusCode(err, types.ErrorCodePreConsumeTokenQuotaFailed, http.StatusForbidden, types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog())
+ }
+ if errors.Is(err, model.ErrUserQuotaInsufficient) {
+ return types.NewErrorWithStatusCode(err, types.ErrorCodeInsufficientUserQuota, http.StatusForbidden, types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog())
+ }
+ if errors.Is(err, model.ErrSubscriptionQuotaInsufficient) {
+ return types.NewErrorWithStatusCode(fmt.Errorf("订阅额度不足或未配置订阅: %s", err.Error()), types.ErrorCodeInsufficientUserQuota, http.StatusForbidden, types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog())
+ }
+ errMsg := err.Error()
+ if strings.Contains(errMsg, "no active subscription") || strings.Contains(errMsg, "subscription quota insufficient") {
+ return types.NewErrorWithStatusCode(fmt.Errorf("订阅额度不足或未配置订阅: %s", errMsg), types.ErrorCodeInsufficientUserQuota, http.StatusForbidden, types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog())
+ }
+ return types.NewError(err, types.ErrorCodeUpdateDataError, types.ErrOptionWithSkipRetry())
+}
+
func (s *BillingSession) reserveFunding(delta int) error {
switch funding := s.funding.(type) {
case *WalletFunding:
@@ -316,13 +617,7 @@ func (s *BillingSession) reserveFunding(delta int) error {
return nil
case *SubscriptionFunding:
if _, err := model.PostConsumeUserSubscriptionDelta(funding.subscriptionId, int64(delta)); err != nil {
- return types.NewErrorWithStatusCode(
- fmt.Errorf("订阅额度不足或未配置订阅: %s", err.Error()),
- types.ErrorCodeInsufficientUserQuota,
- http.StatusForbidden,
- types.ErrOptionWithSkipRetry(),
- types.ErrOptionWithNoRecordErrorLog(),
- )
+ return mapAtomicPreConsumeError(err)
}
return nil
default:
@@ -422,19 +717,20 @@ func NewBillingSession(c *gin.Context, relayInfo *relaycommon.RelayInfo, preCons
pref := common.NormalizeBillingPreference(relayInfo.UserSetting.BillingPreference)
- // 钱包路径需要先检查用户额度
- tryWallet := func() (*BillingSession, *types.MaxAPIError) {
+ // 新钱包请求先做友好的额度检查；已持久化的钱包重放必须进入
+ // operation-key 校验，不能因扣款后的当前余额变化而切换资金源。
+ tryWallet := func(replay bool) (*BillingSession, *types.MaxAPIError) {
userQuota, err := model.GetUserQuota(relayInfo.UserId, false)
if err != nil {
return nil, types.NewError(err, types.ErrorCodeQueryDataError, types.ErrOptionWithSkipRetry())
}
- if userQuota <= 0 {
+ if !replay && userQuota <= 0 {
return nil, types.NewErrorWithStatusCode(
fmt.Errorf("用户额度不足, 剩余额度: %s", logger.FormatQuota(userQuota)),
types.ErrorCodeInsufficientUserQuota, http.StatusForbidden,
types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog())
}
- if userQuota-int64(preConsumedQuota) < 0 {
+ if !replay && userQuota-int64(preConsumedQuota) < 0 {
return nil, types.NewErrorWithStatusCode(
fmt.Errorf("预扣费额度失败, 用户剩余额度: %s, 需要预扣费额度: %s", logger.FormatQuota(userQuota), logger.FormatQuota(preConsumedQuota)),
types.ErrorCodeInsufficientUserQuota, http.StatusForbidden,
@@ -472,13 +768,32 @@ func NewBillingSession(c *gin.Context, relayInfo *relaycommon.RelayInfo, preCons
return session, nil
}
+ persistedSource, hasPersistedSource, sourceErr := model.ResolveBillingPreConsumeSource(relayInfo.RequestId)
+ if sourceErr != nil {
+ return nil, types.NewError(sourceErr, types.ErrorCodeUpdateDataError, types.ErrOptionWithSkipRetry())
+ }
+ if hasPersistedSource {
+ switch persistedSource {
+ case model.BillingSettlementSourceWallet:
+ return tryWallet(true)
+ case model.BillingSettlementSourceSubscription:
+ return trySubscription()
+ default:
+ return nil, types.NewError(
+ fmt.Errorf("unsupported persisted billing source: %s", persistedSource),
+ types.ErrorCodeUpdateDataError,
+ types.ErrOptionWithSkipRetry(),
+ )
+ }
+ }
+
switch pref {
case "subscription_only":
return trySubscription()
case "wallet_only":
- return tryWallet()
+ return tryWallet(false)
case "wallet_first":
- session, err := tryWallet()
+ session, err := tryWallet(false)
if err != nil {
if err.GetErrorCode() == types.ErrorCodeInsufficientUserQuota {
return trySubscription()
@@ -494,12 +809,12 @@ func NewBillingSession(c *gin.Context, relayInfo *relaycommon.RelayInfo, preCons
return nil, types.NewError(subCheckErr, types.ErrorCodeQueryDataError, types.ErrOptionWithSkipRetry())
}
if !hasSub {
- return tryWallet()
+ return tryWallet(false)
}
session, apiErr := trySubscription()
if apiErr != nil {
if apiErr.GetErrorCode() == types.ErrorCodeInsufficientUserQuota {
- return tryWallet()
+ return tryWallet(false)
}
return nil, apiErr
}
diff --git a/service/billing_session_test.go b/service/billing_session_test.go
index 7f855b86..0c2da86d 100644
--- a/service/billing_session_test.go
+++ b/service/billing_session_test.go
@@ -6,8 +6,10 @@ import (
"fmt"
"net/http"
"testing"
+ "time"
"github.com/MAX-API-Next/MAX-API/common"
+ "github.com/MAX-API-Next/MAX-API/dto"
"github.com/MAX-API-Next/MAX-API/model"
relaycommon "github.com/MAX-API-Next/MAX-API/relay/common"
"github.com/MAX-API-Next/MAX-API/types"
@@ -30,6 +32,43 @@ func (f *recordingFundingSource) Settle(delta int) (int64, error) {
return int64(delta), nil
}
+type recordingBillingSettler struct {
+ preConsumed int
+ reserves []int
+}
+
+func (s *recordingBillingSettler) Settle(int) error { return nil }
+func (s *recordingBillingSettler) Refund(*gin.Context) {}
+func (s *recordingBillingSettler) NeedsRefund() bool { return false }
+func (s *recordingBillingSettler) GetPreConsumedQuota() int { return s.preConsumed }
+func (s *recordingBillingSettler) Reserve(target int) error {
+ s.reserves = append(s.reserves, target)
+ s.preConsumed = target
+ return nil
+}
+
+func TestBillingSessionPrepareSettlementPersistsZeroDeltaTaskEffect(t *testing.T) {
+ session := &BillingSession{
+ relayInfo: &relaycommon.RelayInfo{
+ RequestId: "task-fixed-price-request",
+ UserId: 41,
+ TokenId: 42,
+ TokenKey: "fixed-price-token",
+ },
+ funding: &WalletFunding{userId: 41},
+ preConsumedQuota: 10,
+ tokenConsumed: 10,
+ }
+
+ input, err := session.PrepareSettlement(10)
+
+ require.NoError(t, err)
+ require.NotNil(t, input)
+ assert.Equal(t, "request:task-fixed-price-request:finalize", input.OperationKey)
+ assert.Zero(t, input.FundingDelta)
+ assert.Zero(t, input.TokenDelta)
+}
+
func TestBillingSessionCompensatesFundingWhenTokenSettlementFails(t *testing.T) {
truncate(t)
funding := &recordingFundingSource{}
@@ -63,6 +102,326 @@ func TestBillingSessionCompensatesFundingWhenTokenSettlementFails(t *testing.T)
assert.EqualValues(t, 5, token.UsedQuota)
}
+func TestBillingSessionCumulativeReservationSettlesExactlyOnce(t *testing.T) {
+ truncate(t)
+ const userID, tokenID = 801, 802
+ const initiallyReserved = 5
+
+ require.NoError(t, model.DB.Create(&model.User{
+ Id: userID, Username: "realtime-once-user", Quota: 995, Status: common.UserStatusEnabled,
+ }).Error)
+ require.NoError(t, model.DB.Create(&model.Token{
+ Id: tokenID, UserId: userID, Key: "realtime-once-token", Name: "realtime-once-token",
+ Status: common.TokenStatusEnabled, RemainQuota: 995, UsedQuota: initiallyReserved,
+ }).Error)
+
+ ctx, _ := gin.CreateTestContext(nil)
+ info := &relaycommon.RelayInfo{
+ RequestId: "realtime-once-request",
+ UserId: userID,
+ TokenId: tokenID,
+ TokenKey: "realtime-once-token",
+ OriginModelName: "gpt-4o-realtime-preview",
+ UsingGroup: "default",
+ }
+ session := &BillingSession{
+ relayInfo: info,
+ funding: &WalletFunding{
+ userId: userID, consumed: initiallyReserved,
+ },
+ preConsumedQuota: initiallyReserved,
+ tokenConsumed: initiallyReserved,
+ }
+ info.Billing = session
+ info.FinalPreConsumedQuota = initiallyReserved
+ require.NoError(t, session.Reserve(10))
+ require.NoError(t, SettleBilling(ctx, info, 10))
+
+ assert.EqualValues(t, 990, getUserQuota(t, userID))
+ assert.EqualValues(t, 990, getTokenRemainQuota(t, tokenID))
+ assert.Equal(t, 10, session.GetPreConsumedQuota())
+}
+
+func TestBillingSessionWalletPreConsumeUsesDurableRequestIdentity(t *testing.T) {
+ truncate(t)
+ const userID, tokenID = 803, 804
+ seedUser(t, userID, 100)
+ seedToken(t, tokenID, userID, "durable-pre-consume-token", 100)
+ ctx, _ := gin.CreateTestContext(nil)
+ newSession := func() *BillingSession {
+ return &BillingSession{
+ relayInfo: &relaycommon.RelayInfo{
+ RequestId: "durable-pre-consume-request", UserId: userID,
+ TokenId: tokenID, TokenKey: "durable-pre-consume-token",
+ },
+ funding: &WalletFunding{userId: userID},
+ }
+ }
+
+ first := newSession()
+ require.Nil(t, first.preConsume(ctx, 10))
+ second := newSession()
+ require.Nil(t, second.preConsume(ctx, 10))
+
+ assert.EqualValues(t, 90, getUserQuota(t, userID))
+ assert.EqualValues(t, 90, getTokenRemainQuota(t, tokenID))
+ var settlement model.BillingSettlement
+ require.NoError(t, model.DB.Where("operation_key = ?", "request:durable-pre-consume-request:pre-consume").First(&settlement).Error)
+ assert.Equal(t, model.BillingSettlementStatusApplied, settlement.Status)
+}
+
+func TestBillingSessionWalletFirstReplayCannotSwitchFromWalletToSubscription(t *testing.T) {
+ truncate(t)
+ const userID, tokenID, planID, subscriptionID = 805, 806, 807, 808
+ seedUser(t, userID, 10)
+ seedToken(t, tokenID, userID, "wallet-replay-token", 100)
+ seedBillingSubscription(t, planID, subscriptionID, userID, 100)
+
+ ctx, _ := gin.CreateTestContext(nil)
+ newInfo := func() *relaycommon.RelayInfo {
+ return &relaycommon.RelayInfo{
+ RequestId: "wallet-replay-request", UserId: userID,
+ TokenId: tokenID, TokenKey: "wallet-replay-token",
+ OriginModelName: "wallet-replay-model",
+ UserSetting: dto.UserSetting{BillingPreference: "wallet_first"},
+ }
+ }
+
+ first, apiErr := NewBillingSession(ctx, newInfo(), 10)
+ require.Nil(t, apiErr)
+ require.Equal(t, BillingSourceWallet, first.funding.Source())
+
+ second, apiErr := NewBillingSession(ctx, newInfo(), 10)
+ require.Nil(t, apiErr)
+ require.Equal(t, BillingSourceWallet, second.funding.Source())
+ assert.EqualValues(t, 0, getUserQuota(t, userID))
+ assert.EqualValues(t, 90, getTokenRemainQuota(t, tokenID))
+ assert.EqualValues(t, 0, getSubscriptionAmountUsed(t, subscriptionID))
+ assert.EqualValues(t, 0, countSubscriptionPreConsumeRecords(t, "wallet-replay-request"))
+}
+
+func TestBillingSessionWalletFirstReplayCannotSwitchFromSubscriptionToWallet(t *testing.T) {
+ truncate(t)
+ const userID, tokenID, planID, subscriptionID = 809, 810, 813, 814
+ seedUser(t, userID, 0)
+ seedToken(t, tokenID, userID, "subscription-replay-token", 100)
+ seedBillingSubscription(t, planID, subscriptionID, userID, 100)
+
+ ctx, _ := gin.CreateTestContext(nil)
+ newInfo := func() *relaycommon.RelayInfo {
+ return &relaycommon.RelayInfo{
+ RequestId: "subscription-replay-request", UserId: userID,
+ TokenId: tokenID, TokenKey: "subscription-replay-token",
+ OriginModelName: "subscription-replay-model",
+ UserSetting: dto.UserSetting{BillingPreference: "wallet_first"},
+ }
+ }
+
+ first, apiErr := NewBillingSession(ctx, newInfo(), 10)
+ require.Nil(t, apiErr)
+ require.Equal(t, BillingSourceSubscription, first.funding.Source())
+ require.NoError(t, model.DB.Model(&model.User{}).Where("id = ?", userID).Update("quota", 100).Error)
+
+ second, apiErr := NewBillingSession(ctx, newInfo(), 10)
+ require.Nil(t, apiErr)
+ require.Equal(t, BillingSourceSubscription, second.funding.Source())
+ assert.EqualValues(t, 100, getUserQuota(t, userID))
+ assert.EqualValues(t, 90, getTokenRemainQuota(t, tokenID))
+ assert.EqualValues(t, 10, getSubscriptionAmountUsed(t, subscriptionID))
+ assert.EqualValues(t, 1, countSubscriptionPreConsumeRecords(t, "subscription-replay-request"))
+}
+
+func TestBillingSessionTrustedWalletReplayFailsClosedInsteadOfSwitchingFunding(t *testing.T) {
+ truncate(t)
+ const userID, tokenID, planID, subscriptionID = 815, 816, 817, 818
+ trustQuota := common.GetTrustQuota()
+ require.Positive(t, trustQuota, "trust quota must be enabled for this scenario")
+ seedUser(t, userID, trustQuota+100)
+ seedToken(t, tokenID, userID, "trusted-wallet-replay-token", trustQuota+100)
+ seedBillingSubscription(t, planID, subscriptionID, userID, 100)
+
+ ctx, _ := gin.CreateTestContext(nil)
+ ctx.Set("token_quota", int64(trustQuota+100))
+ newInfo := func() *relaycommon.RelayInfo {
+ return &relaycommon.RelayInfo{
+ RequestId: "trusted-wallet-replay-request", UserId: userID,
+ TokenId: tokenID, TokenKey: "trusted-wallet-replay-token",
+ OriginModelName: "trusted-wallet-replay-model",
+ UserSetting: dto.UserSetting{BillingPreference: "wallet_first"},
+ }
+ }
+
+ first, apiErr := NewBillingSession(ctx, newInfo(), 10)
+ require.Nil(t, apiErr)
+ require.True(t, first.trusted)
+ var tombstone model.BillingSettlement
+ require.NoError(t, model.DB.Where("operation_key = ?", "request:trusted-wallet-replay-request:pre-consume").First(&tombstone).Error)
+ assert.Zero(t, tombstone.FundingDelta)
+ assert.Zero(t, tombstone.TokenDelta)
+
+ require.NoError(t, model.DB.Model(&model.User{}).Where("id = ?", userID).Update("quota", 0).Error)
+ ctx.Set("token_quota", int64(0))
+ second, apiErr := NewBillingSession(ctx, newInfo(), 10)
+ require.Nil(t, second)
+ require.NotNil(t, apiErr)
+ assert.ErrorIs(t, apiErr, model.ErrBillingSettlementOperationConflict)
+ assert.EqualValues(t, trustQuota+100, getTokenRemainQuota(t, tokenID))
+ assert.EqualValues(t, 0, getSubscriptionAmountUsed(t, subscriptionID))
+ assert.EqualValues(t, 0, countSubscriptionPreConsumeRecords(t, "trusted-wallet-replay-request"))
+}
+
+func seedBillingSubscription(t *testing.T, planID int, subscriptionID int, userID int, amountTotal int64) {
+ t.Helper()
+ require.NoError(t, model.DB.Create(&model.SubscriptionPlan{
+ Id: planID, Title: fmt.Sprintf("billing-plan-%d", planID), Enabled: true,
+ TotalAmount: amountTotal, QuotaResetPeriod: model.SubscriptionResetNever,
+ }).Error)
+ now := time.Now()
+ require.NoError(t, model.DB.Create(&model.UserSubscription{
+ Id: subscriptionID, UserId: userID, PlanId: planID,
+ AmountTotal: amountTotal, Status: "active",
+ StartTime: now.Add(-time.Hour).Unix(), EndTime: now.Add(time.Hour).Unix(),
+ }).Error)
+}
+
+func getSubscriptionAmountUsed(t *testing.T, subscriptionID int) int64 {
+ t.Helper()
+ var subscription model.UserSubscription
+ require.NoError(t, model.DB.First(&subscription, subscriptionID).Error)
+ return subscription.AmountUsed
+}
+
+func countSubscriptionPreConsumeRecords(t *testing.T, requestID string) int64 {
+ t.Helper()
+ var count int64
+ require.NoError(t, model.DB.Model(&model.SubscriptionPreConsumeRecord{}).Where("request_id = ?", requestID).Count(&count).Error)
+ return count
+}
+
+func TestPreWssConsumeQuotaUsesCumulativeReservationInsteadOfIndependentCharge(t *testing.T) {
+ server := miniredis.RunT(t)
+ client := redis.NewClient(&redis.Options{Addr: server.Addr()})
+ oldRDB := common.RDB
+ oldRedisEnabled := common.RedisEnabled
+ common.RDB = client
+ common.RedisEnabled = true
+ t.Cleanup(func() {
+ _ = client.Close()
+ common.RDB = oldRDB
+ common.RedisEnabled = oldRedisEnabled
+ })
+ require.NoError(t, client.HSet(context.Background(), "user:811", map[string]interface{}{
+ "Id": 811, "Quota": 100000, "Role": common.RoleCommonUser, "Status": common.UserStatusEnabled,
+ }).Err())
+ require.NoError(t, client.HSet(context.Background(), fmt.Sprintf("token:%s", common.GenerateHMAC("realtime-reserve-token")), map[string]interface{}{
+ "Id": 812, "UserId": 811, "Status": common.TokenStatusEnabled, "RemainQuota": 100000,
+ }).Err())
+
+ ctx, _ := gin.CreateTestContext(nil)
+ billing := &recordingBillingSettler{}
+ info := &relaycommon.RelayInfo{
+ UserId: 811, TokenId: 812, TokenKey: "realtime-reserve-token",
+ OriginModelName: "gpt-4o-realtime-preview", UsingGroup: "default", Billing: billing,
+ }
+ usage := &dto.RealtimeUsage{
+ TotalTokens: 4, InputTokens: 4,
+ InputTokenDetails: dto.InputTokenDetails{TextTokens: 4},
+ }
+ require.NoError(t, PreWssConsumeQuota(ctx, info, usage))
+ require.Len(t, billing.reserves, 1)
+ assert.Greater(t, billing.reserves[0], 0)
+
+ firstTarget := billing.reserves[0]
+ require.NoError(t, client.HSet(context.Background(), "user:811", "Quota", firstTarget).Err())
+ require.NoError(t, client.HSet(context.Background(), fmt.Sprintf("token:%s", common.GenerateHMAC("realtime-reserve-token")), "RemainQuota", firstTarget).Err())
+ usage.TotalTokens = 8
+ usage.InputTokens = 8
+ usage.InputTokenDetails.TextTokens = 8
+ require.NoError(t, PreWssConsumeQuota(ctx, info, usage))
+ require.Len(t, billing.reserves, 2)
+ assert.Greater(t, billing.reserves[1], firstTarget)
+ assert.LessOrEqual(t, billing.reserves[1]-firstTarget, firstTarget)
+}
+
+func TestPostConsumeQuotaRejectsAmountChangedReplay(t *testing.T) {
+ truncate(t)
+ const userID, tokenID = 821, 822
+ seedUser(t, userID, 100)
+ seedToken(t, tokenID, userID, "stable-post-token", 100)
+ info := &relaycommon.RelayInfo{
+ RequestId: "stable-post-request", TaskRelayInfo: &relaycommon.TaskRelayInfo{Action: "violation-fee"},
+ UserId: userID, TokenId: tokenID, TokenKey: "stable-post-token",
+ }
+
+ require.NoError(t, PostConsumeQuota(info, 10, 0, false))
+ err := PostConsumeQuota(info, 20, 0, false)
+ require.Error(t, err)
+ assert.ErrorIs(t, err, model.ErrBillingSettlementOperationConflict)
+ assert.EqualValues(t, 90, getUserQuota(t, userID))
+ assert.EqualValues(t, 90, getTokenRemainQuota(t, tokenID))
+}
+
+func TestViolationFeeSettlesOriginalSubscriptionPreConsume(t *testing.T) {
+ truncate(t)
+ const userID, tokenID, planID, subscriptionID = 831, 832, 833, 834
+ const preConsumed int64 = 10
+ seedUser(t, userID, 0)
+ seedToken(t, tokenID, userID, "violation-subscription-token", 100)
+ seedChannel(t, 1)
+ seedBillingSubscription(t, planID, subscriptionID, userID, 100)
+
+ result, err := model.PreConsumeTokenAndUserSubscription(
+ "violation-subscription-request", userID, tokenID,
+ "violation-subscription-token", "grok-test", 0, preConsumed,
+ )
+ require.NoError(t, err)
+ require.Equal(t, subscriptionID, result.UserSubscriptionId)
+
+ oldQuotaPerUnit := common.QuotaPerUnit
+ common.QuotaPerUnit = 100
+ t.Cleanup(func() { common.QuotaPerUnit = oldQuotaPerUnit })
+
+ ctx, _ := gin.CreateTestContext(nil)
+ info := &relaycommon.RelayInfo{
+ RequestId: "violation-subscription-request", UserId: userID,
+ TokenId: tokenID, TokenKey: "violation-subscription-token",
+ ChannelMeta: &relaycommon.ChannelMeta{ChannelId: 1},
+ OriginModelName: "grok-test", UsingGroup: "default",
+ BillingSource: BillingSourceSubscription, SubscriptionId: subscriptionID,
+ StartTime: time.Now(),
+ }
+ info.PriceData.GroupRatioInfo.GroupRatio = 1
+ session := &BillingSession{
+ relayInfo: info,
+ funding: &SubscriptionFunding{
+ requestId: "violation-subscription-request", userId: userID,
+ subscriptionId: subscriptionID, preConsumed: preConsumed,
+ },
+ preConsumedQuota: int(preConsumed), tokenConsumed: int(preConsumed),
+ }
+ info.Billing = session
+
+ apiErr := types.NewErrorWithStatusCode(
+ errors.New(CSAMViolationMarker), types.ErrorCodeBadResponseStatusCode,
+ http.StatusBadRequest,
+ )
+ HandleFailedBilling(ctx, info, NormalizeViolationFeeError(apiErr))
+
+ feeQuota := calcViolationFeeQuota(0.05, 1)
+ require.EqualValues(t, feeQuota, getSubscriptionAmountUsed(t, subscriptionID))
+ require.EqualValues(t, 100-feeQuota, getTokenRemainQuota(t, tokenID))
+ require.Equal(t, int64(1), countLogs(t))
+ require.Equal(t, feeQuota, getLastLog(t).Quota)
+
+ var finalize model.BillingSettlement
+ require.NoError(t, model.DB.Where("operation_key = ?", "request:violation-subscription-request:finalize").First(&finalize).Error)
+ require.Equal(t, model.BillingSettlementStatusApplied, finalize.Status)
+ var legacyCount int64
+ require.NoError(t, model.DB.Model(&model.BillingSettlement{}).
+ Where("operation_key LIKE ?", "request:violation-subscription-request:legacy-post:%").Count(&legacyCount).Error)
+ require.Zero(t, legacyCount)
+}
+
type compensationFailingFundingSource struct {
deltas []int
}
@@ -167,76 +526,6 @@ func TestBillingSessionReconcilesPartialFundingCompensationBeforeRetry(t *testin
assert.Equal(t, int64(2), int64(funding.deltas[len(funding.deltas)-1]))
}
-type exhaustUserAfterTokenCacheReadHook struct {
- userID int
- tokenKey string
- exhausted bool
- err error
-}
-
-func (h *exhaustUserAfterTokenCacheReadHook) BeforeProcess(ctx context.Context, _ redis.Cmder) (context.Context, error) {
- return ctx, nil
-}
-
-func (h *exhaustUserAfterTokenCacheReadHook) AfterProcess(_ context.Context, cmd redis.Cmder) error {
- args := cmd.Args()
- if h.exhausted || cmd.Name() != "hgetall" || len(args) < 2 || fmt.Sprint(args[1]) != h.tokenKey {
- return nil
- }
- h.exhausted = true
- h.err = model.DB.Model(&model.User{}).Where("id = ?", h.userID).Update("quota", 0).Error
- return nil
-}
-
-func (*exhaustUserAfterTokenCacheReadHook) BeforeProcessPipeline(ctx context.Context, _ []redis.Cmder) (context.Context, error) {
- return ctx, nil
-}
-
-func (*exhaustUserAfterTokenCacheReadHook) AfterProcessPipeline(context.Context, []redis.Cmder) error {
- return nil
-}
-
-func TestPreConsumeQuotaMapsConcurrentUserExhaustion(t *testing.T) {
- truncate(t)
- seedUser(t, 73, 10)
- seedToken(t, 74, 73, "concurrent-user-exhaustion", 10)
-
- server := miniredis.RunT(t)
- client := redis.NewClient(&redis.Options{Addr: server.Addr()})
- oldRDB := common.RDB
- oldRedisEnabled := common.RedisEnabled
- common.RDB = client
- common.RedisEnabled = true
- t.Cleanup(func() {
- _ = client.Close()
- common.RDB = oldRDB
- common.RedisEnabled = oldRedisEnabled
- })
- userCacheKey := "user:73"
- tokenCacheKey := fmt.Sprintf("token:%s", common.GenerateHMAC("concurrent-user-exhaustion"))
- require.NoError(t, client.HSet(context.Background(), userCacheKey, map[string]interface{}{
- "Id": 73, "Quota": 10, "Role": common.RoleCommonUser, "Status": common.UserStatusEnabled,
- }).Err())
- require.NoError(t, client.HSet(context.Background(), tokenCacheKey, map[string]interface{}{
- "Id": 74, "UserId": 73, "Status": common.TokenStatusEnabled, "RemainQuota": 10,
- }).Err())
- hook := &exhaustUserAfterTokenCacheReadHook{userID: 73, tokenKey: tokenCacheKey}
- client.AddHook(hook)
-
- ctx, _ := gin.CreateTestContext(nil)
- apiErr := PreConsumeQuota(ctx, 7, &relaycommon.RelayInfo{
- UserId: 73, TokenId: 74, TokenKey: "concurrent-user-exhaustion",
- })
- require.NotNil(t, apiErr)
- require.NoError(t, hook.err)
- assert.True(t, hook.exhausted)
- assert.ErrorIs(t, apiErr, model.ErrUserQuotaInsufficient)
- assert.Equal(t, types.ErrorCodeInsufficientUserQuota, apiErr.GetErrorCode())
- assert.Equal(t, http.StatusForbidden, apiErr.StatusCode)
- assert.True(t, types.IsSkipRetryError(apiErr))
- assert.False(t, types.IsRecordErrorLog(apiErr))
-}
-
type tokenRecoveringFundingSource struct {
t *testing.T
deltas []int
@@ -285,25 +574,78 @@ func TestBillingSessionTrustsFiniteInt64TokenQuota(t *testing.T) {
assert.True(t, session.shouldTrust(ctx))
}
-func TestPreConsumeQuotaTrustsFiniteInt64TokenQuota(t *testing.T) {
+func TestBillingSessionRefundUsesDurableOperation(t *testing.T) {
truncate(t)
- trustQuota := common.GetTrustQuota()
- seedUser(t, 71, trustQuota+100)
- seedToken(t, 72, 71, "finite-int64-token", trustQuota+100)
+ const userID, tokenID = 81, 82
+ seedUser(t, userID, 100)
+ seedToken(t, tokenID, userID, "durable-refund-token", 100)
+ require.NoError(t, model.DB.Model(&model.Token{}).Where("id = ?", tokenID).Update("used_quota", 10).Error)
ctx, _ := gin.CreateTestContext(nil)
- ctx.Set("token_quota", int64(trustQuota+100))
- relayInfo := &relaycommon.RelayInfo{
- UserId: 71, TokenId: 72, TokenKey: "finite-int64-token",
+ session := &BillingSession{
+ relayInfo: &relaycommon.RelayInfo{
+ RequestId: "durable-refund-request",
+ UserId: userID,
+ TokenId: tokenID,
+ TokenKey: "durable-refund-token",
+ },
+ funding: &WalletFunding{userId: userID, consumed: 10},
+ preConsumedQuota: 10,
+ tokenConsumed: 10,
}
- require.Nil(t, PreConsumeQuota(ctx, 10, relayInfo))
- assert.Zero(t, relayInfo.FinalPreConsumedQuota)
+ session.Refund(ctx)
+ require.EqualValues(t, 110, getUserQuota(t, userID))
+ require.EqualValues(t, 110, getTokenRemainQuota(t, tokenID))
- var user model.User
- require.NoError(t, model.DB.First(&user, 71).Error)
- assert.EqualValues(t, trustQuota+100, user.Quota)
- var token model.Token
- require.NoError(t, model.DB.First(&token, 72).Error)
- assert.EqualValues(t, trustQuota+100, token.RemainQuota)
+ var settlement model.BillingSettlement
+ require.NoError(t, model.DB.Where("operation_key = ?", "request:durable-refund-request:finalize").First(&settlement).Error)
+ require.Equal(t, "applied", settlement.Status)
+}
+
+func TestBillingSessionDurableSettleEffectConflictDoesNotRemainRefundable(t *testing.T) {
+ truncate(t)
+ const userID, tokenID = 83, 84
+ const requestID = "durable-effect-conflict-request"
+ seedUser(t, userID, 100)
+ seedToken(t, tokenID, userID, "durable-effect-conflict-token", 100)
+
+ effect := &model.BillingSettlementEffect{
+ LogType: model.LogTypeConsume, Content: "settlement effect conflict",
+ ModelName: "test-model", TokenID: tokenID, Group: "default",
+ }
+ operationKey := "request:" + requestID + ":finalize"
+ applied, _, err := model.ApplyBillingSettlementOnce(model.BillingSettlementInput{
+ OperationKey: operationKey,
+ Source: model.BillingSettlementSourceWallet,
+ UserID: userID,
+ TokenID: tokenID,
+ TokenKey: "durable-effect-conflict-token",
+ FundingDelta: 5,
+ TokenDelta: 5,
+ Effect: effect,
+ })
+ require.NoError(t, err)
+ require.EqualValues(t, 5, applied)
+ require.NoError(t, model.DB.Model(&model.BillingSettlement{}).
+ Where("operation_key = ?", operationKey).
+ Update("effect_status", model.BillingSettlementEffectApplying).Error)
+
+ session := &BillingSession{
+ relayInfo: &relaycommon.RelayInfo{
+ RequestId: requestID, UserId: userID,
+ TokenId: tokenID, TokenKey: "durable-effect-conflict-token",
+ },
+ funding: &WalletFunding{userId: userID, consumed: 10},
+ preConsumedQuota: 10,
+ tokenConsumed: 10,
+ }
+
+ require.NoError(t, session.SettleWithEffect(15, effect))
+ assert.True(t, session.settled)
+ assert.True(t, session.fundingSettled)
+ assert.EqualValues(t, 5, session.appliedFundingDelta)
+ assert.False(t, session.NeedsRefund())
+ assert.EqualValues(t, 95, getUserQuota(t, userID))
+ assert.EqualValues(t, 95, getTokenRemainQuota(t, tokenID))
}
diff --git a/service/channel_select.go b/service/channel_select.go
index e53f8504..f8071ad4 100644
--- a/service/channel_select.go
+++ b/service/channel_select.go
@@ -12,12 +12,13 @@ import (
)
type RetryParam struct {
- Ctx *gin.Context
- TokenGroup string
- ModelName string
- RequestPath string
- Retry *int
- resetNextTry bool
+ Ctx *gin.Context
+ TokenGroup string
+ ModelName string
+ RequestPath string
+ Retry *int
+ resetNextTry bool
+ excludedChannelIDs map[int]struct{}
}
func (p *RetryParam) GetRetry() int {
@@ -46,6 +47,18 @@ func (p *RetryParam) ResetRetryNextTry() {
p.resetNextTry = true
}
+// ExcludeChannel prevents a retried request from selecting the same failed
+// channel again while leaving each new request independent.
+func (p *RetryParam) ExcludeChannel(channelID int) {
+ if channelID <= 0 {
+ return
+ }
+ if p.excludedChannelIDs == nil {
+ p.excludedChannelIDs = make(map[int]struct{})
+ }
+ p.excludedChannelIDs[channelID] = struct{}{}
+}
+
// CacheGetRandomSatisfiedChannel tries to get a random channel that satisfies the requirements.
// 尝试获取一个满足要求的随机渠道。
//
@@ -116,7 +129,7 @@ func CacheGetRandomSatisfiedChannel(param *RetryParam) (*model.Channel, string,
}
logger.LogDebug(param.Ctx, "Auto selecting group: %s, priorityRetry: %d", autoGroup, priorityRetry)
- channel, err = model.GetRandomSatisfiedChannel(autoGroup, param.ModelName, priorityRetry, param.RequestPath)
+ channel, err = model.GetRandomSatisfiedChannelExcluding(autoGroup, param.ModelName, priorityRetry, param.RequestPath, param.excludedChannelIDs)
if err != nil {
return nil, autoGroup, err
}
@@ -157,7 +170,7 @@ func CacheGetRandomSatisfiedChannel(param *RetryParam) (*model.Channel, string,
break
}
} else {
- channel, err = model.GetRandomSatisfiedChannel(param.TokenGroup, param.ModelName, param.GetRetry(), param.RequestPath)
+ channel, err = model.GetRandomSatisfiedChannelExcluding(param.TokenGroup, param.ModelName, param.GetRetry(), param.RequestPath, param.excludedChannelIDs)
if err != nil {
return nil, param.TokenGroup, err
}
diff --git a/service/error.go b/service/error.go
index 57e2374c..de9e3e8a 100644
--- a/service/error.go
+++ b/service/error.go
@@ -150,25 +150,86 @@ func ResetStatusCode(maxApiErr *types.MaxAPIError, statusCodeMappingStr string)
if maxApiErr == nil {
return
}
- if statusCodeMappingStr == "" || statusCodeMappingStr == "{}" {
+ maxApiErr.StatusCode = MapStatusCode(maxApiErr.StatusCode, statusCodeMappingStr)
+}
+
+func ResetTaskStatusCode(taskErr *dto.TaskError, statusCodeMappingStr string) {
+ if taskErr == nil {
return
}
- statusCodeMapping := make(map[string]any)
- err := common.Unmarshal([]byte(statusCodeMappingStr), &statusCodeMapping)
- if err != nil {
+ mappedStatusCode := MapStatusCode(taskErr.StatusCode, statusCodeMappingStr)
+ if mappedStatusCode == taskErr.StatusCode {
return
}
- if maxApiErr.StatusCode == http.StatusOK {
- return
+ if taskErr.UpstreamStatusCode == 0 {
+ taskErr.UpstreamStatusCode = taskErr.StatusCode
}
- codeStr := strconv.Itoa(maxApiErr.StatusCode)
- if value, ok := statusCodeMapping[codeStr]; ok {
- intCode, ok := parseStatusCodeMappingValue(value)
- if !ok {
- return
+ taskErr.StatusCode = mappedStatusCode
+}
+
+func MapStatusCode(statusCode int, statusCodeMappingStr string) int {
+ if statusCode == http.StatusOK {
+ return statusCode
+ }
+ rawMapping, err := decodeStatusCodeMapping(statusCodeMappingStr)
+ if err != nil {
+ return statusCode
+ }
+ value, ok := rawMapping[strconv.Itoa(statusCode)]
+ if !ok {
+ return statusCode
+ }
+ mappedStatusCode, ok := parseStatusCodeMappingValue(value)
+ if !ok || mappedStatusCode < http.StatusContinue || mappedStatusCode > 599 {
+ return statusCode
+ }
+ return mappedStatusCode
+}
+
+func ValidateStatusCodeMapping(statusCodeMappingStr string) error {
+ _, err := parseStatusCodeMapping(statusCodeMappingStr)
+ return err
+}
+
+func parseStatusCodeMapping(statusCodeMappingStr string) (map[int]int, error) {
+ rawMapping, err := decodeStatusCodeMapping(statusCodeMappingStr)
+ if err != nil {
+ return nil, err
+ }
+
+ statusCodeMapping := make(map[int]int, len(rawMapping))
+ for source, value := range rawMapping {
+ normalizedSource := strings.TrimSpace(source)
+ sourceCode, err := strconv.Atoi(normalizedSource)
+ if err != nil || sourceCode < http.StatusContinue || sourceCode > 599 || strconv.Itoa(sourceCode) != normalizedSource {
+ return nil, fmt.Errorf("invalid source status code %q", source)
+ }
+ targetCode, ok := parseStatusCodeMappingValue(value)
+ if !ok || targetCode < http.StatusContinue || targetCode > 599 {
+ return nil, fmt.Errorf("invalid target status code for %q", source)
+ }
+ if _, exists := statusCodeMapping[sourceCode]; exists {
+ return nil, fmt.Errorf("duplicate source status code %d", sourceCode)
}
- maxApiErr.StatusCode = intCode
+ statusCodeMapping[sourceCode] = targetCode
+ }
+ return statusCodeMapping, nil
+}
+
+func decodeStatusCodeMapping(statusCodeMappingStr string) (map[string]any, error) {
+ trimmed := strings.TrimSpace(statusCodeMappingStr)
+ if trimmed == "" || trimmed == "{}" {
+ return map[string]any{}, nil
+ }
+
+ rawMapping := make(map[string]any)
+ if err := common.Unmarshal([]byte(trimmed), &rawMapping); err != nil {
+ return nil, fmt.Errorf("invalid JSON object: %w", err)
+ }
+ if rawMapping == nil {
+ return nil, fmt.Errorf("mapping must be a JSON object")
}
+ return rawMapping, nil
}
func parseStatusCodeMappingValue(value any) (int, bool) {
diff --git a/service/error_test.go b/service/error_test.go
index 8eb50256..dc3467cd 100644
--- a/service/error_test.go
+++ b/service/error_test.go
@@ -10,6 +10,7 @@ import (
"testing"
"github.com/MAX-API-Next/MAX-API/common"
+ "github.com/MAX-API-Next/MAX-API/dto"
"github.com/MAX-API-Next/MAX-API/types"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
@@ -48,6 +49,18 @@ func TestResetStatusCode(t *testing.T) {
statusCodeConfig: `{"200":503}`,
expectedCode: 200,
},
+ {
+ name: "skip out of range target",
+ statusCode: 429,
+ statusCodeConfig: `{"429":999}`,
+ expectedCode: 429,
+ },
+ {
+ name: "ignore unrelated invalid entry",
+ statusCode: 429,
+ statusCodeConfig: `{"429":503,"500":999}`,
+ expectedCode: 503,
+ },
}
for _, tc := range testCases {
@@ -64,6 +77,25 @@ func TestResetStatusCode(t *testing.T) {
}
}
+func TestResetTaskStatusCodePreservesOriginalUpstreamStatus(t *testing.T) {
+ t.Parallel()
+
+ upstreamErr := &dto.TaskError{StatusCode: http.StatusTooManyRequests}
+ ResetTaskStatusCode(upstreamErr, `{"429":"503"}`)
+ require.Equal(t, http.StatusServiceUnavailable, upstreamErr.StatusCode)
+ require.Equal(t, http.StatusTooManyRequests, upstreamErr.UpstreamStatusCode)
+}
+
+func TestValidateStatusCodeMapping(t *testing.T) {
+ t.Parallel()
+
+ require.NoError(t, ValidateStatusCodeMapping(`{"429":"503","500":502}`))
+ require.Error(t, ValidateStatusCodeMapping(`null`))
+ require.Error(t, ValidateStatusCodeMapping(`{"bad":503}`))
+ require.Error(t, ValidateStatusCodeMapping(`{"429":999}`))
+ require.Error(t, ValidateStatusCodeMapping(`{"429":503,"0429":502}`))
+}
+
func TestRelayErrorHandlerTruncatesInvalidJSONBodyInLog(t *testing.T) {
withDebugEnabled(t, false)
diff --git a/service/funding_source.go b/service/funding_source.go
index e6228d92..2d55fa20 100644
--- a/service/funding_source.go
+++ b/service/funding_source.go
@@ -90,6 +90,14 @@ func (s *SubscriptionFunding) PreConsume(amount int) error {
if err != nil {
return err
}
+ s.applyPreConsumeResult(res)
+ return nil
+}
+
+func (s *SubscriptionFunding) applyPreConsumeResult(res *model.SubscriptionPreConsumeResult) {
+ if res == nil {
+ return
+ }
s.subscriptionId = res.UserSubscriptionId
s.preConsumed = res.PreConsumed
s.AmountTotal = res.AmountTotal
@@ -99,7 +107,6 @@ func (s *SubscriptionFunding) PreConsume(amount int) error {
s.PlanId = planInfo.PlanId
s.PlanTitle = planInfo.PlanTitle
}
- return nil
}
func (s *SubscriptionFunding) Settle(delta int) (int64, error) {
diff --git a/service/midjourney_image_url.go b/service/midjourney_image_url.go
new file mode 100644
index 00000000..29b58c06
--- /dev/null
+++ b/service/midjourney_image_url.go
@@ -0,0 +1,39 @@
+package service
+
+import (
+ "crypto/hmac"
+ "fmt"
+ "net/url"
+ "strconv"
+ "strings"
+ "time"
+
+ "github.com/MAX-API-Next/MAX-API/common"
+)
+
+const MidjourneyImageURLTTL = 15 * time.Minute
+
+func midjourneyImageSignaturePayload(mjID string, userID int, expiresAt int64) string {
+ return fmt.Sprintf("midjourney-image\n%s\n%d\n%d", mjID, userID, expiresAt)
+}
+
+func SignMidjourneyImageURL(mjID string, userID int, expiresAt int64) string {
+ return common.GenerateHMAC(midjourneyImageSignaturePayload(mjID, userID, expiresAt))
+}
+
+func ValidateMidjourneyImageURL(mjID string, userID int, expiresAt int64, signature string, now time.Time) bool {
+ if strings.TrimSpace(mjID) == "" || userID <= 0 || expiresAt <= now.Unix() || signature == "" {
+ return false
+ }
+ expected := SignMidjourneyImageURL(mjID, userID, expiresAt)
+ return hmac.Equal([]byte(expected), []byte(signature))
+}
+
+func BuildMidjourneyImageURL(serverAddress, mjID string, userID int, now time.Time) string {
+ expiresAt := now.Add(MidjourneyImageURLTTL).Unix()
+ query := url.Values{}
+ query.Set("uid", strconv.Itoa(userID))
+ query.Set("expires", strconv.FormatInt(expiresAt, 10))
+ query.Set("signature", SignMidjourneyImageURL(mjID, userID, expiresAt))
+ return strings.TrimRight(serverAddress, "/") + "/mj/image/" + url.PathEscape(mjID) + "?" + query.Encode()
+}
diff --git a/service/midjourney_image_url_test.go b/service/midjourney_image_url_test.go
new file mode 100644
index 00000000..afce6dee
--- /dev/null
+++ b/service/midjourney_image_url_test.go
@@ -0,0 +1,31 @@
+package service
+
+import (
+ "net/url"
+ "strconv"
+ "testing"
+ "time"
+
+ "github.com/MAX-API-Next/MAX-API/common"
+ "github.com/stretchr/testify/require"
+)
+
+func TestMidjourneyImageURLBindsTaskUserAndExpiry(t *testing.T) {
+ oldSecret := common.CryptoSecret
+ common.CryptoSecret = "midjourney-image-test-secret"
+ t.Cleanup(func() { common.CryptoSecret = oldSecret })
+
+ now := time.Unix(1_700_000_000, 0)
+ rawURL := BuildMidjourneyImageURL("https://max-api.example/", "mj-task/one", 42, now)
+ parsed, err := url.Parse(rawURL)
+ require.NoError(t, err)
+ require.Equal(t, "/mj/image/mj-task/one", parsed.Path)
+
+ expiresAt, err := strconv.ParseInt(parsed.Query().Get("expires"), 10, 64)
+ require.NoError(t, err)
+ signature := parsed.Query().Get("signature")
+ require.True(t, ValidateMidjourneyImageURL("mj-task/one", 42, expiresAt, signature, now))
+ require.False(t, ValidateMidjourneyImageURL("mj-task/one", 43, expiresAt, signature, now))
+ require.False(t, ValidateMidjourneyImageURL("other-task", 42, expiresAt, signature, now))
+ require.False(t, ValidateMidjourneyImageURL("mj-task/one", 42, expiresAt, signature, time.Unix(expiresAt, 0)))
+}
diff --git a/service/pre_consume_quota.go b/service/pre_consume_quota.go
deleted file mode 100644
index e4e8e03b..00000000
--- a/service/pre_consume_quota.go
+++ /dev/null
@@ -1,94 +0,0 @@
-package service
-
-import (
- "errors"
- "fmt"
- "net/http"
-
- "github.com/MAX-API-Next/MAX-API/common"
- "github.com/MAX-API-Next/MAX-API/logger"
- "github.com/MAX-API-Next/MAX-API/model"
- relaycommon "github.com/MAX-API-Next/MAX-API/relay/common"
- "github.com/MAX-API-Next/MAX-API/types"
-
- "github.com/bytedance/gopkg/util/gopool"
- "github.com/gin-gonic/gin"
-)
-
-func ReturnPreConsumedQuota(c *gin.Context, relayInfo *relaycommon.RelayInfo) {
- if relayInfo.FinalPreConsumedQuota != 0 {
- logger.LogInfo(c, fmt.Sprintf("用户 %d 请求失败, 返还预扣费额度 %s", relayInfo.UserId, logger.FormatQuota(relayInfo.FinalPreConsumedQuota)))
- gopool.Go(func() {
- relayInfoCopy := *relayInfo
-
- err := PostConsumeQuota(&relayInfoCopy, -relayInfoCopy.FinalPreConsumedQuota, 0, false)
- if err != nil {
- common.SysLog("error return pre-consumed quota: " + err.Error())
- }
- })
- }
-}
-
-// PreConsumeQuota checks if the user has enough quota to pre-consume.
-// It returns the pre-consumed quota if successful, or an error if not.
-func PreConsumeQuota(c *gin.Context, preConsumedQuota int, relayInfo *relaycommon.RelayInfo) *types.MaxAPIError {
- if apiErr := validatePreConsumedQuota(preConsumedQuota, relayInfo); apiErr != nil {
- return apiErr
- }
- userQuota, err := model.GetUserQuota(relayInfo.UserId, false)
- if err != nil {
- return types.NewError(err, types.ErrorCodeQueryDataError, types.ErrOptionWithSkipRetry())
- }
- if userQuota <= 0 {
- return types.NewErrorWithStatusCode(fmt.Errorf("用户额度不足, 剩余额度: %s", logger.FormatQuota(userQuota)), types.ErrorCodeInsufficientUserQuota, http.StatusForbidden, types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog())
- }
- if userQuota-int64(preConsumedQuota) < 0 {
- return types.NewErrorWithStatusCode(fmt.Errorf("预扣费额度失败, 用户剩余额度: %s, 需要预扣费额度: %s", logger.FormatQuota(userQuota), logger.FormatQuota(preConsumedQuota)), types.ErrorCodeInsufficientUserQuota, http.StatusForbidden, types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog())
- }
-
- trustQuota := common.GetTrustQuota()
-
- relayInfo.UserQuota = userQuota
- if userQuota > int64(trustQuota) {
- // 用户额度充足，判断令牌额度是否充足
- if !relayInfo.TokenUnlimited {
- // 非无限令牌，判断令牌额度是否充足
- tokenQuota := c.GetInt64("token_quota")
- if tokenQuota > int64(trustQuota) {
- // 令牌额度充足，信任令牌
- preConsumedQuota = 0
- logger.LogInfo(c, fmt.Sprintf("用户 %d 剩余额度 %s 且令牌 %d 额度 %d 充足, 信任且不需要预扣费", relayInfo.UserId, logger.FormatQuota(userQuota), relayInfo.TokenId, tokenQuota))
- }
- } else {
- // in this case, we do not pre-consume quota
- // because the user has enough quota
- preConsumedQuota = 0
- logger.LogInfo(c, fmt.Sprintf("用户 %d 额度充足且为无限额度令牌, 信任且不需要预扣费", relayInfo.UserId))
- }
- }
-
- if preConsumedQuota > 0 {
- token, err := model.GetTokenByKey(relayInfo.TokenKey, false)
- if err != nil {
- return types.NewErrorWithStatusCode(err, types.ErrorCodePreConsumeTokenQuotaFailed, http.StatusForbidden, types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog())
- }
- if !relayInfo.TokenUnlimited && token.RemainQuota < int64(preConsumedQuota) {
- return types.NewErrorWithStatusCode(
- fmt.Errorf("token quota is not enough, token remain quota: %s, need quota: %s", logger.FormatQuota(token.RemainQuota), logger.FormatQuota(preConsumedQuota)),
- types.ErrorCodePreConsumeTokenQuotaFailed, http.StatusForbidden, types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog())
- }
- err = model.DecreaseTokenAndUserQuota(relayInfo.TokenId, relayInfo.UserId, relayInfo.TokenKey, preConsumedQuota)
- if errors.Is(err, model.ErrTokenQuotaInsufficient) {
- return types.NewErrorWithStatusCode(err, types.ErrorCodePreConsumeTokenQuotaFailed, http.StatusForbidden, types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog())
- }
- if errors.Is(err, model.ErrUserQuotaInsufficient) {
- return types.NewErrorWithStatusCode(err, types.ErrorCodeInsufficientUserQuota, http.StatusForbidden, types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog())
- }
- if err != nil {
- return types.NewError(err, types.ErrorCodeUpdateDataError, types.ErrOptionWithSkipRetry())
- }
- logger.LogInfo(c, fmt.Sprintf("用户 %d 预扣费 %s, 预扣费后剩余额度: %s", relayInfo.UserId, logger.FormatQuota(preConsumedQuota), logger.FormatQuota(userQuota-int64(preConsumedQuota))))
- }
- relayInfo.FinalPreConsumedQuota = preConsumedQuota
- return nil
-}
diff --git a/service/pre_consume_validation_test.go b/service/pre_consume_validation_test.go
index 082bc8b1..43e2ca35 100644
--- a/service/pre_consume_validation_test.go
+++ b/service/pre_consume_validation_test.go
@@ -11,51 +11,31 @@ import (
"github.com/stretchr/testify/require"
)
-func TestPreConsumeValidationSharedByBillingAndQuota(t *testing.T) {
- entries := []struct {
- name string
- call func(int, *relaycommon.RelayInfo) *types.MaxAPIError
- }{
- {
- name: "billing session",
- call: func(preConsumedQuota int, relayInfo *relaycommon.RelayInfo) *types.MaxAPIError {
- return PreConsumeBilling(nil, preConsumedQuota, relayInfo)
- },
- },
- {
- name: "legacy quota",
- call: func(preConsumedQuota int, relayInfo *relaycommon.RelayInfo) *types.MaxAPIError {
- return PreConsumeQuota(nil, preConsumedQuota, relayInfo)
- },
- },
- }
-
- for _, entry := range entries {
- t.Run(entry.name+"/quota_clamp", func(t *testing.T) {
- clamp := &common.QuotaClamp{
- Op: "test",
- Kind: common.QuotaClampOverflow,
- Original: float64(common.MaxQuota) + 1,
- Clamped: common.MaxQuota,
- }
-
- apiErr := entry.call(1, &relaycommon.RelayInfo{QuotaClamp: clamp})
-
- require.NotNil(t, apiErr)
- require.Equal(t, clamp, apiErr.Err)
- require.Equal(t, types.ErrorCodeModelPriceError, apiErr.GetErrorCode())
- require.Equal(t, http.StatusBadRequest, apiErr.StatusCode)
- require.True(t, types.IsSkipRetryError(apiErr))
- })
-
- t.Run(entry.name+"/negative_quota", func(t *testing.T) {
- apiErr := entry.call(-1, &relaycommon.RelayInfo{})
-
- require.NotNil(t, apiErr)
- require.ErrorContains(t, apiErr.Err, "pre-consume quota cannot be negative: -1")
- require.Equal(t, types.ErrorCodeModelPriceError, apiErr.GetErrorCode())
- require.Equal(t, http.StatusBadRequest, apiErr.StatusCode)
- require.True(t, types.IsSkipRetryError(apiErr))
- })
- }
+func TestPreConsumeBillingValidation(t *testing.T) {
+ t.Run("quota_clamp", func(t *testing.T) {
+ clamp := &common.QuotaClamp{
+ Op: "test",
+ Kind: common.QuotaClampOverflow,
+ Original: float64(common.MaxQuota) + 1,
+ Clamped: common.MaxQuota,
+ }
+
+ apiErr := PreConsumeBilling(nil, 1, &relaycommon.RelayInfo{QuotaClamp: clamp})
+
+ require.NotNil(t, apiErr)
+ require.Equal(t, clamp, apiErr.Err)
+ require.Equal(t, types.ErrorCodeModelPriceError, apiErr.GetErrorCode())
+ require.Equal(t, http.StatusBadRequest, apiErr.StatusCode)
+ require.True(t, types.IsSkipRetryError(apiErr))
+ })
+
+ t.Run("negative_quota", func(t *testing.T) {
+ apiErr := PreConsumeBilling(nil, -1, &relaycommon.RelayInfo{})
+
+ require.NotNil(t, apiErr)
+ require.ErrorContains(t, apiErr.Err, "pre-consume quota cannot be negative: -1")
+ require.Equal(t, types.ErrorCodeModelPriceError, apiErr.GetErrorCode())
+ require.Equal(t, http.StatusBadRequest, apiErr.StatusCode)
+ require.True(t, types.IsSkipRetryError(apiErr))
+ })
}
diff --git a/service/quota.go b/service/quota.go
index 50855077..91728f1f 100644
--- a/service/quota.go
+++ b/service/quota.go
@@ -139,19 +139,26 @@ func PreWssConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usag
quota, clamp := calculateAudioQuota(quotaInfo)
noteQuotaClamp(relayInfo, clamp)
- if userQuota < int64(quota) {
- return fmt.Errorf("user quota is not enough, user quota: %s, need quota: %s", logger.FormatQuota(userQuota), logger.FormatQuota(quota))
+ if quota == 0 {
+ return nil
}
-
- if !token.UnlimitedQuota && token.RemainQuota < int64(quota) {
- return fmt.Errorf("token quota is not enough, token remain quota: %s, need quota: %s", logger.FormatQuota(token.RemainQuota), logger.FormatQuota(quota))
+ if relayInfo.Billing == nil {
+ return errors.New("realtime cumulative reservation requires a billing session")
}
-
- err = PostConsumeQuota(relayInfo, quota, 0, false)
- if err != nil {
+ neededQuota := quota - relayInfo.Billing.GetPreConsumedQuota()
+ if neededQuota <= 0 {
+ return nil
+ }
+ if userQuota < int64(neededQuota) {
+ return fmt.Errorf("user quota is not enough, user quota: %s, need quota: %s", logger.FormatQuota(userQuota), logger.FormatQuota(neededQuota))
+ }
+ if !token.UnlimitedQuota && token.RemainQuota < int64(neededQuota) {
+ return fmt.Errorf("token quota is not enough, token remain quota: %s, need quota: %s", logger.FormatQuota(token.RemainQuota), logger.FormatQuota(neededQuota))
+ }
+ if err := relayInfo.Billing.Reserve(quota); err != nil {
return err
}
- logger.LogInfo(ctx, "realtime streaming consume quota success, quota: "+fmt.Sprintf("%d", quota))
+ logger.LogInfo(ctx, "realtime streaming cumulative quota reserved, target: "+fmt.Sprintf("%d", quota))
return nil
}
@@ -425,50 +432,55 @@ func PreConsumeTokenQuota(relayInfo *relaycommon.RelayInfo, quota int) error {
}
func PostConsumeQuota(relayInfo *relaycommon.RelayInfo, quota int, preConsumedQuota int, sendEmail bool) (err error) {
+ return PostConsumeQuotaOnce(relayInfo, "post-consume", quota, preConsumedQuota, sendEmail)
+}
- // 1) Consume from wallet quota OR subscription item
- if relayInfo != nil && relayInfo.BillingSource == BillingSourceSubscription {
- if relayInfo.SubscriptionId == 0 {
- return errors.New("subscription id is missing")
- }
- delta := int64(quota)
- if delta != 0 {
- appliedDelta, err := model.PostConsumeUserSubscriptionDelta(relayInfo.SubscriptionId, delta)
- if err != nil {
- return err
- }
- relayInfo.SubscriptionPostDelta += appliedDelta
- }
- } else {
- // Wallet
- if quota > 0 {
- err = model.DecreaseUserQuota(relayInfo.UserId, quota, false)
- } else {
- err = model.IncreaseUserQuota(relayInfo.UserId, -quota, false)
- }
- if err != nil {
- return err
- }
+func PostConsumeQuotaOnce(relayInfo *relaycommon.RelayInfo, operationKind string, quota int, preConsumedQuota int, sendEmail bool) (err error) {
+ if relayInfo == nil {
+ return errors.New("relayInfo is nil")
}
-
- if !relayInfo.IsPlayground {
- if quota > 0 {
- err = model.DecreaseTokenQuota(relayInfo.TokenId, relayInfo.TokenKey, quota)
- } else {
- err = model.IncreaseTokenQuota(relayInfo.TokenId, relayInfo.TokenKey, -quota)
+ if operationKind == "" {
+ return errors.New("post-consume operation kind is required")
+ }
+ if quota == 0 {
+ return nil
+ }
+ if relayInfo.RequestId != "" {
+ source := model.BillingSettlementSourceWallet
+ if relayInfo.BillingSource == BillingSourceSubscription {
+ source = model.BillingSettlementSourceSubscription
}
- if err != nil {
- return err
+ tokenDelta := int64(quota)
+ if relayInfo.IsPlayground {
+ tokenDelta = 0
}
- }
-
- if sendEmail {
- if (quota + preConsumedQuota) != 0 {
- checkAndSendQuotaNotify(relayInfo, quota, preConsumedQuota)
+ operationKey := fmt.Sprintf("request:%s:legacy-post:%s", relayInfo.RequestId, operationKind)
+ _, _, durableErr := model.ApplyBillingSettlementOnce(model.BillingSettlementInput{
+ OperationKey: operationKey,
+ Source: source,
+ UserID: relayInfo.UserId,
+ SubscriptionID: relayInfo.SubscriptionId,
+ TokenID: relayInfo.TokenId,
+ TokenKey: relayInfo.TokenKey,
+ FundingDelta: int64(quota),
+ TokenDelta: tokenDelta,
+ SubscriptionPreConsumeRequestID: relayInfo.RequestId,
+ AllowMissingToken: quota < 0,
+ })
+ if durableErr != nil {
+ return durableErr
}
+ if sendEmail && (quota+preConsumedQuota) != 0 {
+ if relayInfo.BillingSource == BillingSourceSubscription {
+ checkAndSendSubscriptionQuotaNotify(relayInfo)
+ } else {
+ checkAndSendQuotaNotify(relayInfo, quota, preConsumedQuota)
+ }
+ }
+ return nil
}
- return nil
+ return errors.New("balance mutation requires a stable request id")
}
func checkAndSendQuotaNotify(relayInfo *relaycommon.RelayInfo, quota int, preConsumedQuota int) {
diff --git a/service/system_task.go b/service/system_task.go
index b178125a..e1482a7b 100644
--- a/service/system_task.go
+++ b/service/system_task.go
@@ -30,6 +30,8 @@ var (
systemTaskHandlers = map[string]SystemTaskHandler{}
)
+var updateSystemTaskState = model.UpdateSystemTaskState
+
func RegisterSystemTaskHandler(h SystemTaskHandler) {
if h == nil {
return
@@ -218,7 +220,7 @@ func runLogCleanupTask(ctx context.Context, task *model.SystemTask, runnerID str
return
}
syncLogCleanupStateFromRemaining(&state, remaining)
- if err := model.UpdateSystemTaskState(task.TaskID, runnerID, state, systemTaskLockUntil()); err != nil {
+ if err := updateSystemTaskState(task.TaskID, runnerID, state, systemTaskLockUntil()); err != nil {
logSystemTaskLockError(ctx, task, err)
return
}
@@ -249,7 +251,7 @@ func runLogCleanupTask(ctx context.Context, task *model.SystemTask, runnerID str
}
state.Progress = logCleanupProgress(state.Processed, state.Total)
- if err := model.UpdateSystemTaskState(task.TaskID, runnerID, state, systemTaskLockUntil()); err != nil {
+ if err := updateSystemTaskState(task.TaskID, runnerID, state, systemTaskLockUntil()); err != nil {
logSystemTaskLockError(ctx, task, err)
return
}
@@ -266,9 +268,19 @@ func runLogCleanupTask(ctx context.Context, task *model.SystemTask, runnerID str
if state.Total < state.Processed {
state.Total = state.Processed
}
- if err := model.UpdateSystemTaskState(task.TaskID, runnerID, state, systemTaskLockUntil()); err != nil {
- logSystemTaskLockError(ctx, task, err)
- return
+ for {
+ rowsAffected, err := model.DeleteOldBillingLogReceiptsBatch(ctx, payload.TargetTimestamp, payload.BatchSize)
+ if err != nil {
+ failSystemTask(task, runnerID, err)
+ return
+ }
+ if err := updateSystemTaskState(task.TaskID, runnerID, state, systemTaskLockUntil()); err != nil {
+ logSystemTaskLockError(ctx, task, err)
+ return
+ }
+ if rowsAffected < int64(payload.BatchSize) {
+ break
+ }
}
result := LogCleanupResult{DeletedCount: state.Processed}
@@ -346,7 +358,7 @@ func NewSystemTaskProgressReporter(task *model.SystemTask, runnerID string) func
lastWriteAt = time.Now()
state := SystemTaskProgress{Total: total, Processed: processed, Progress: progress}
- _ = model.UpdateSystemTaskState(task.TaskID, runnerID, state, systemTaskLockUntil())
+ _ = updateSystemTaskState(task.TaskID, runnerID, state, systemTaskLockUntil())
}
}
diff --git a/service/system_task_test.go b/service/system_task_test.go
new file mode 100644
index 00000000..82f0c14a
--- /dev/null
+++ b/service/system_task_test.go
@@ -0,0 +1,61 @@
+package service
+
+import (
+ "context"
+ "sync/atomic"
+ "testing"
+
+ "github.com/MAX-API-Next/MAX-API/model"
+ "github.com/stretchr/testify/assert"
+ "github.com/stretchr/testify/require"
+)
+
+func TestLogCleanupTaskRenewsLockBetweenBillingReceiptBatches(t *testing.T) {
+ truncate(t)
+ require.NoError(t, model.DB.AutoMigrate(&model.SystemTask{}))
+ require.NoError(t, model.DB.Exec("DELETE FROM system_tasks").Error)
+
+ task, err := model.CreateSystemTask(
+ model.SystemTaskTypeLogCleanup,
+ model.SystemTaskTypeLogCleanup,
+ LogCleanupPayload{TargetTimestamp: 100, BatchSize: 1},
+ LogCleanupState{},
+ )
+ require.NoError(t, err)
+ const runnerID = "receipt-cleanup-runner"
+ claimedTask, claimed, err := model.ClaimSystemTask(task.ID, model.SystemTaskTypeLogCleanup, runnerID, systemTaskLockUntil())
+ require.NoError(t, err)
+ require.True(t, claimed)
+
+ var renewals atomic.Int64
+ oldUpdateSystemTaskState := updateSystemTaskState
+ updateSystemTaskState = func(taskID string, lockedBy string, state any, lockUntil int64) error {
+ err := oldUpdateSystemTaskState(taskID, lockedBy, state, lockUntil)
+ if err == nil {
+ renewals.Add(1)
+ }
+ return err
+ }
+ t.Cleanup(func() {
+ updateSystemTaskState = oldUpdateSystemTaskState
+ _ = model.DB.Exec("DELETE FROM system_tasks").Error
+ })
+
+ require.NoError(t, model.DB.Create(&[]model.BillingLogReceipt{
+ {OperationKey: "old-receipt-1", ClaimToken: "1", CreatedAt: 10},
+ {OperationKey: "old-receipt-2", ClaimToken: "2", CreatedAt: 20},
+ {OperationKey: "old-receipt-3", ClaimToken: "3", CreatedAt: 30},
+ }).Error)
+
+ runLogCleanupTask(context.Background(), claimedTask, runnerID)
+
+ var remaining int64
+ require.NoError(t, model.DB.Model(&model.BillingLogReceipt{}).Where("created_at < ?", 100).Count(&remaining).Error)
+ assert.Zero(t, remaining)
+
+ assert.EqualValues(t, 5, renewals.Load())
+
+ var reloaded model.SystemTask
+ require.NoError(t, model.DB.First(&reloaded, claimedTask.ID).Error)
+ assert.Equal(t, model.SystemTaskStatusSucceeded, reloaded.Status)
+}
diff --git a/service/task_billing.go b/service/task_billing.go
index c5462bd3..1f739fce 100644
--- a/service/task_billing.go
+++ b/service/task_billing.go
@@ -14,10 +14,7 @@ import (
"github.com/gin-gonic/gin"
)
-// LogTaskConsumption 记录任务消费日志和统计信息（仅记录，不涉及实际扣费）。
-// 实际扣费已由 BillingSession（PreConsumeBilling + SettleBilling）完成。
-func LogTaskConsumption(c *gin.Context, info *relaycommon.RelayInfo) {
- tokenName := c.GetString("token_name")
+func buildTaskConsumptionLogData(c *gin.Context, info *relaycommon.RelayInfo) (string, map[string]interface{}) {
logContent := fmt.Sprintf("操作 %s", info.Action)
// 支持任务仅按次计费
if info.TaskBilling != nil {
@@ -39,7 +36,9 @@ func LogTaskConsumption(c *gin.Context, info *relaycommon.RelayInfo) {
}
other := make(map[string]interface{})
other["is_task"] = true
- other["request_path"] = c.Request.URL.Path
+ if c != nil && c.Request != nil && c.Request.URL != nil {
+ other["request_path"] = c.Request.URL.Path
+ }
other["model_price"] = info.PriceData.ModelPrice
if info.PriceData.ModelRatio > 0 {
other["model_ratio"] = info.PriceData.ModelRatio
@@ -56,6 +55,39 @@ func LogTaskConsumption(c *gin.Context, info *relaycommon.RelayInfo) {
other["upstream_model_name"] = info.UpstreamModelName
}
attachQuotaSaturation(c, info, other)
+ return logContent, other
+}
+
+// BuildTaskSubmissionSettlementEffect creates the durable non-financial effect
+// for a successfully submitted task. The enclosing settlement operation keeps
+// the log, usage counters, and dashboard update replayable after a restart.
+func BuildTaskSubmissionSettlementEffect(c *gin.Context, info *relaycommon.RelayInfo, quota int) *model.BillingSettlementEffect {
+ if info == nil {
+ return nil
+ }
+ logContent, other := buildTaskConsumptionLogData(c, info)
+ return &model.BillingSettlementEffect{
+ LogType: model.LogTypeConsume,
+ Content: logContent,
+ ChannelID: info.ChannelId,
+ ModelName: info.OriginModelName,
+ TokenID: info.TokenId,
+ Group: info.UsingGroup,
+ Other: other,
+ NodeName: common.NodeName,
+ UpdateUsage: true,
+ Quota: int64(quota),
+ }
+}
+
+// LogTaskConsumption records task usage for legacy/custom funding paths that
+// do not have a durable BillingSettlement operation to carry the effect.
+func LogTaskConsumption(c *gin.Context, info *relaycommon.RelayInfo) {
+ tokenName := ""
+ if c != nil {
+ tokenName = c.GetString("token_name")
+ }
+ logContent, other := buildTaskConsumptionLogData(c, info)
model.RecordConsumeLog(c, info.UserId, model.RecordConsumeLogParams{
ChannelId: info.ChannelId,
ModelName: info.OriginModelName,
@@ -74,55 +106,11 @@ func LogTaskConsumption(c *gin.Context, info *relaycommon.RelayInfo) {
// 异步任务计费辅助函数
// ---------------------------------------------------------------------------
-// resolveTokenKey 通过 TokenId 运行时获取令牌 Key（用于 Redis 缓存操作）。
-// 如果令牌已被删除或查询失败，返回空字符串。
-func resolveTokenKey(ctx context.Context, tokenId int, taskID string) string {
- token, err := model.GetTokenById(tokenId)
- if err != nil {
- logger.LogWarn(ctx, fmt.Sprintf("获取令牌 key 失败 (tokenId=%d, task=%s): %s", tokenId, taskID, err.Error()))
- return ""
- }
- return token.Key
-}
-
// taskIsSubscription 判断任务是否通过订阅计费。
func taskIsSubscription(task *model.Task) bool {
return task.PrivateData.BillingSource == BillingSourceSubscription && task.PrivateData.SubscriptionId > 0
}
-// taskAdjustFunding 调整任务的资金来源（钱包或订阅），delta > 0 表示扣费，delta < 0 表示退还。
-func taskAdjustFunding(task *model.Task, delta int) error {
- if taskIsSubscription(task) {
- _, err := model.PostConsumeUserSubscriptionDelta(task.PrivateData.SubscriptionId, int64(delta))
- return err
- }
- if delta > 0 {
- return model.DecreaseUserQuota(task.UserId, delta, false)
- }
- return model.IncreaseUserQuota(task.UserId, -delta, false)
-}
-
-// taskAdjustTokenQuota 调整任务的令牌额度，delta > 0 表示扣费，delta < 0 表示退还。
-// 需要通过 resolveTokenKey 运行时获取 key（不从 PrivateData 中读取）。
-func taskAdjustTokenQuota(ctx context.Context, task *model.Task, delta int) {
- if task.PrivateData.TokenId <= 0 || delta == 0 {
- return
- }
- tokenKey := resolveTokenKey(ctx, task.PrivateData.TokenId, task.TaskID)
- if tokenKey == "" {
- return
- }
- var err error
- if delta > 0 {
- err = model.DecreaseTokenQuota(task.PrivateData.TokenId, tokenKey, delta)
- } else {
- err = model.IncreaseTokenQuota(task.PrivateData.TokenId, tokenKey, -delta)
- }
- if err != nil {
- logger.LogWarn(ctx, fmt.Sprintf("调整令牌额度失败 (delta=%d, task=%s): %s", delta, task.TaskID, err.Error()))
- }
-}
-
// taskBillingOther 从 task 的 BillingContext 构建日志 Other 字段。
func taskBillingOther(task *model.Task) map[string]interface{} {
other := make(map[string]interface{})
@@ -157,39 +145,119 @@ func taskModelName(task *model.Task) string {
return task.Properties.OriginModelName
}
-// RefundTaskQuota 统一的任务失败退款逻辑。
-// 当异步任务失败时，将预扣的 quota 退还给用户（支持钱包和订阅），并退还令牌额度。
-func RefundTaskQuota(ctx context.Context, task *model.Task, reason string) {
- quota := task.Quota
- if quota == 0 {
- return
+func buildTaskRefundSettlementInput(task *model.Task, reason string) *model.BillingSettlementInput {
+ if task == nil || task.Quota == 0 || task.ID <= 0 {
+ return nil
}
-
- // 1. 退还资金来源（钱包或订阅）
- if err := taskAdjustFunding(task, -quota); err != nil {
- logger.LogWarn(ctx, fmt.Sprintf("退还资金来源失败 task %s: %s", task.TaskID, err.Error()))
- return
+ source := model.BillingSettlementSourceWallet
+ if taskIsSubscription(task) {
+ source = model.BillingSettlementSourceSubscription
}
-
- // 2. 退还令牌额度
- taskAdjustTokenQuota(ctx, task, -quota)
-
- // 3. 记录日志
other := taskBillingOther(task)
other["task_id"] = task.TaskID
other["reason"] = reason
- model.RecordTaskBillingLog(model.RecordTaskBillingLogParams{
- UserId: task.UserId,
- LogType: model.LogTypeRefund,
- Content: "",
- ChannelId: task.ChannelId,
- ModelName: taskModelName(task),
- Quota: quota,
- TokenId: task.PrivateData.TokenId,
- Group: task.Group,
- Other: other,
- NodeName: task.PrivateData.NodeName,
- })
+ return &model.BillingSettlementInput{
+ OperationKey: fmt.Sprintf("task:%d:refund", task.ID),
+ Source: source,
+ UserID: task.UserId,
+ SubscriptionID: task.PrivateData.SubscriptionId,
+ TokenID: task.PrivateData.TokenId,
+ FundingDelta: -int64(task.Quota),
+ TokenDelta: -int64(task.Quota),
+ TaskID: task.ID,
+ TaskQuota: int64(task.Quota),
+ TaskQuotaTarget: 0,
+ SubscriptionPreConsumeRequestID: task.PrivateData.BillingRequestId,
+ FinalizeSubscriptionPreConsume: source == model.BillingSettlementSourceSubscription,
+ AllowMissingToken: true,
+ Effect: &model.BillingSettlementEffect{
+ LogType: model.LogTypeRefund, Content: reason, ChannelID: task.ChannelId,
+ ModelName: taskModelName(task), TokenID: task.PrivateData.TokenId,
+ Group: task.Group, Other: other, NodeName: task.PrivateData.NodeName,
+ },
+ }
+}
+
+func buildTaskFinalSettlementInput(task *model.Task, actualQuota int, reason string, clamps ...*common.QuotaClamp) *model.BillingSettlementInput {
+ if task == nil || task.ID <= 0 || actualQuota <= 0 || actualQuota == task.Quota {
+ return nil
+ }
+ quotaDelta := actualQuota - task.Quota
+ source := model.BillingSettlementSourceWallet
+ if taskIsSubscription(task) {
+ source = model.BillingSettlementSourceSubscription
+ }
+ tokenDelta := int64(quotaDelta)
+ if task.PrivateData.TokenId <= 0 {
+ tokenDelta = 0
+ }
+ logType := model.LogTypeRefund
+ if quotaDelta > 0 {
+ logType = model.LogTypeConsume
+ }
+ other := taskBillingOther(task)
+ other["task_id"] = task.TaskID
+ other["pre_consumed_quota"] = task.Quota
+ other["actual_quota"] = actualQuota
+ for _, clamp := range clamps {
+ attachQuotaSaturationToOther(other, clamp)
+ }
+ return &model.BillingSettlementInput{
+ OperationKey: fmt.Sprintf("task:%d:finalize", task.ID),
+ Source: source,
+ UserID: task.UserId,
+ SubscriptionID: task.PrivateData.SubscriptionId,
+ TokenID: task.PrivateData.TokenId,
+ FundingDelta: int64(quotaDelta),
+ TokenDelta: tokenDelta,
+ TaskID: task.ID,
+ TaskQuota: int64(task.Quota),
+ TaskQuotaTarget: int64(actualQuota),
+ SubscriptionPreConsumeRequestID: task.PrivateData.BillingRequestId,
+ AllowMissingToken: quotaDelta < 0,
+ Effect: &model.BillingSettlementEffect{
+ LogType: logType, Content: reason, ChannelID: task.ChannelId,
+ ModelName: taskModelName(task), TokenID: task.PrivateData.TokenId,
+ Group: task.Group, Other: other, NodeName: task.PrivateData.NodeName,
+ UpdateUsage: quotaDelta > 0,
+ },
+ }
+}
+
+func applyTaskBillingSettlement(ctx context.Context, task *model.Task, input *model.BillingSettlementInput) bool {
+ if input == nil {
+ return true
+ }
+ appliedFundingDelta, _, err := model.ApplyBillingSettlementOnce(*input)
+ if err != nil {
+ logger.LogWarn(ctx, fmt.Sprintf("任务结算仍待恢复 task %s operation=%s: %s", task.TaskID, input.OperationKey, err.Error()))
+ return false
+ }
+ task.Quota = int(input.TaskQuota + appliedFundingDelta)
+ if effectErr := model.ProcessBillingSettlementEffect(input.OperationKey); effectErr != nil {
+ logger.LogWarn(ctx, fmt.Sprintf("任务结算日志/统计仍待恢复 task %s: %s", task.TaskID, effectErr.Error()))
+ }
+ if appliedFundingDelta != input.FundingDelta {
+ logger.LogWarn(ctx, fmt.Sprintf("任务 %s 结算资金实际应用=%d，目标=%d，账面额度=%d", task.TaskID, appliedFundingDelta, input.FundingDelta, task.Quota))
+ }
+ return true
+}
+
+// RefundTaskQuota 统一的任务失败退款逻辑。
+// 当异步任务失败时，将预扣的 quota 退还给用户（支持钱包和订阅），并退还令牌额度。
+func RefundTaskQuota(ctx context.Context, task *model.Task, reason string) bool {
+ if task == nil {
+ return false
+ }
+ if task.Quota == 0 {
+ return true
+ }
+ if task.ID <= 0 {
+ logger.LogWarn(ctx, fmt.Sprintf("拒绝退款未持久化任务 task %s", task.TaskID))
+ return false
+ }
+
+ return applyTaskBillingSettlement(ctx, task, buildTaskRefundSettlementInput(task, reason))
}
// RecalculateTaskQuota 通用的异步差额结算。
@@ -199,6 +267,10 @@ func RecalculateTaskQuota(ctx context.Context, task *model.Task, actualQuota int
if actualQuota <= 0 {
return
}
+ if task == nil || task.ID <= 0 {
+ logger.LogWarn(ctx, "拒绝对未持久化任务执行差额结算")
+ return
+ }
preConsumedQuota := task.Quota
quotaDelta := actualQuota - preConsumedQuota
@@ -216,59 +288,24 @@ func RecalculateTaskQuota(ctx context.Context, task *model.Task, actualQuota int
reason,
))
- // 调整资金来源
- if err := taskAdjustFunding(task, quotaDelta); err != nil {
- logger.LogError(ctx, fmt.Sprintf("差额结算资金调整失败 task %s: %s", task.TaskID, err.Error()))
- return
- }
-
- // 调整令牌额度
- taskAdjustTokenQuota(ctx, task, quotaDelta)
-
- task.Quota = actualQuota
- if err := task.UpdateQuota(); err != nil {
- logger.LogError(ctx, fmt.Sprintf("差额结算回写 quota 失败 task %s: %s", task.TaskID, err.Error()))
- }
-
- var logType int
- var logQuota int
- if quotaDelta > 0 {
- logType = model.LogTypeConsume
- logQuota = quotaDelta
- model.UpdateUserUsedQuotaAndRequestCount(task.UserId, quotaDelta)
- model.UpdateChannelUsedQuota(task.ChannelId, quotaDelta)
- } else {
- logType = model.LogTypeRefund
- logQuota = -quotaDelta
- }
- other := taskBillingOther(task)
- other["task_id"] = task.TaskID
- other["pre_consumed_quota"] = preConsumedQuota
- other["actual_quota"] = actualQuota
- for _, clamp := range clamps {
- attachQuotaSaturationToOther(other, clamp)
- }
- model.RecordTaskBillingLog(model.RecordTaskBillingLogParams{
- UserId: task.UserId,
- LogType: logType,
- Content: reason,
- ChannelId: task.ChannelId,
- ModelName: taskModelName(task),
- Quota: logQuota,
- TokenId: task.PrivateData.TokenId,
- Group: task.Group,
- Other: other,
- NodeName: task.PrivateData.NodeName,
- })
+ applyTaskBillingSettlement(ctx, task, buildTaskFinalSettlementInput(task, actualQuota, reason, clamps...))
}
// RecalculateTaskQuotaByTokens 根据实际 token 消耗重新计费（异步差额结算）。
// 当任务成功且返回了 totalTokens 时，根据模型倍率和分组倍率重新计算实际扣费额度，
// 与预扣费的差额进行补扣或退还。支持钱包和订阅计费来源。
func RecalculateTaskQuotaByTokens(ctx context.Context, task *model.Task, totalTokens int) {
- if totalTokens <= 0 {
+ actualQuota, reason, clamp, ok := calculateTaskQuotaByTokens(task, totalTokens)
+ if !ok {
return
}
+ RecalculateTaskQuota(ctx, task, actualQuota, reason, clamp)
+}
+
+func calculateTaskQuotaByTokens(task *model.Task, totalTokens int) (int, string, *common.QuotaClamp, bool) {
+ if task == nil || totalTokens <= 0 {
+ return 0, "", nil, false
+ }
modelName := taskModelName(task)
@@ -276,7 +313,7 @@ func RecalculateTaskQuotaByTokens(ctx context.Context, task *model.Task, totalTo
modelRatio, hasRatioSetting, _ := ratio_setting.GetModelRatio(modelName)
// 只有配置了倍率(非固定价格)时才按 token 重新计费
if !hasRatioSetting || modelRatio <= 0 {
- return
+ return 0, "", nil, false
}
// 获取用户和组的倍率信息
@@ -288,7 +325,7 @@ func RecalculateTaskQuotaByTokens(ctx context.Context, task *model.Task, totalTo
}
}
if group == "" {
- return
+ return 0, "", nil, false
}
groupRatio := ratio_setting.GetGroupRatio(group)
@@ -315,5 +352,5 @@ func RecalculateTaskQuotaByTokens(ctx context.Context, task *model.Task, totalTo
actualQuota, clamp := common.QuotaFromFloatChecked(float64(totalTokens) * modelRatio * finalGroupRatio * otherMultiplier)
reason := fmt.Sprintf("token重算：tokens=%d, modelRatio=%.2f, groupRatio=%.2f, otherMultiplier=%.4f", totalTokens, modelRatio, finalGroupRatio, otherMultiplier)
- RecalculateTaskQuota(ctx, task, actualQuota, reason, clamp)
+ return actualQuota, reason, clamp, true
}
diff --git a/service/task_billing_test.go b/service/task_billing_test.go
index 9d309999..b8c27528 100644
--- a/service/task_billing_test.go
+++ b/service/task_billing_test.go
@@ -3,8 +3,10 @@ package service
import (
"context"
"encoding/json"
+ "fmt"
"net/http"
"os"
+ "sync"
"testing"
"time"
@@ -43,9 +45,15 @@ func TestMain(m *testing.M) {
&model.User{},
&model.Token{},
&model.Log{},
+ &model.BillingLogReceipt{},
&model.Channel{},
&model.TopUp{},
+ &model.SubscriptionPlan{},
&model.UserSubscription{},
+ &model.SubscriptionPreConsumeRecord{},
+ &model.BillingSettlement{},
+ &model.BillingPreConsumeSelection{},
+ &model.CacheInvalidationTask{},
); err != nil {
panic("failed to migrate: " + err.Error())
}
@@ -68,6 +76,37 @@ func TestShouldApplyTaskResultProgressSkipsNonTerminalCompleteProgress(t *testin
}))
}
+func TestSweepTimedOutUnconfirmedSubmitRequiresReviewWithoutRefund(t *testing.T) {
+ truncate(t)
+ originalTimeout := constant.TaskTimeoutMinutes
+ constant.TaskTimeoutMinutes = 1
+ t.Cleanup(func() { constant.TaskTimeoutMinutes = originalTimeout })
+
+ task := &model.Task{
+ TaskID: "task_unconfirmed_submit",
+ Status: model.TaskStatusNotStart,
+ Quota: 30,
+ SubmitTime: time.Now().Add(-2 * time.Minute).Unix(),
+ PrivateData: model.TaskPrivateData{
+ AwaitingUpstreamID: true,
+ BillingRequestId: "unconfirmed-submit-request",
+ BillingSource: model.BillingSettlementSourceWallet,
+ },
+ }
+ require.NoError(t, model.DB.Create(task).Error)
+
+ sweepTimedOutTasks(context.Background())
+
+ var reloaded model.Task
+ require.NoError(t, model.DB.First(&reloaded, task.ID).Error)
+ require.EqualValues(t, model.TaskStatusFailure, reloaded.Status)
+ require.Equal(t, 30, reloaded.Quota)
+ require.Contains(t, reloaded.FailReason, "未自动退款")
+ var settlementCount int64
+ require.NoError(t, model.DB.Model(&model.BillingSettlement{}).Count(&settlementCount).Error)
+ require.Zero(t, settlementCount)
+}
+
// ---------------------------------------------------------------------------
// Seed helpers
// ---------------------------------------------------------------------------
@@ -79,9 +118,15 @@ func truncate(t *testing.T) {
model.DB.Exec("DELETE FROM users")
model.DB.Exec("DELETE FROM tokens")
model.DB.Exec("DELETE FROM logs")
+ model.DB.Exec("DELETE FROM billing_log_receipts")
model.DB.Exec("DELETE FROM channels")
model.DB.Exec("DELETE FROM top_ups")
+ model.DB.Exec("DELETE FROM subscription_plans")
model.DB.Exec("DELETE FROM user_subscriptions")
+ model.DB.Exec("DELETE FROM subscription_pre_consume_records")
+ model.DB.Exec("DELETE FROM billing_settlements")
+ model.DB.Exec("DELETE FROM billing_pre_consume_selections")
+ model.DB.Exec("DELETE FROM cache_invalidation_tasks")
})
}
@@ -100,7 +145,8 @@ func seedToken(t *testing.T, id int, userId int, key string, remainQuota int) {
Name: "test_token",
Status: common.TokenStatusEnabled,
RemainQuota: int64(remainQuota),
- UsedQuota: 0,
+ // Task fixtures model a token that already paid the pre-consumed quota.
+ UsedQuota: int64(remainQuota),
}
require.NoError(t, model.DB.Create(token).Error)
}
@@ -152,6 +198,21 @@ func makeTask(userId, channelId, quota, tokenId int, billingSource string, subsc
}
}
+func persistTask(t *testing.T, task *model.Task) {
+ t.Helper()
+ if task.PrivateData.BillingSource == BillingSourceSubscription && task.PrivateData.BillingRequestId == "" {
+ task.PrivateData.BillingRequestId = "task-preconsume:" + task.TaskID
+ var sub model.UserSubscription
+ require.NoError(t, model.DB.Where("id = ?", task.PrivateData.SubscriptionId).First(&sub).Error)
+ require.NoError(t, model.DB.Create(&model.SubscriptionPreConsumeRecord{
+ RequestId: task.PrivateData.BillingRequestId, UserId: task.UserId,
+ TokenId: task.PrivateData.TokenId, UserSubscriptionId: sub.Id, PreConsumed: int64(task.Quota), Status: "consumed",
+ CreatedAt: time.Now().Unix(), UpdatedAt: time.Now().Unix(),
+ }).Error)
+ }
+ require.NoError(t, model.DB.Create(task).Error)
+}
+
// ---------------------------------------------------------------------------
// Read-back helpers
// ---------------------------------------------------------------------------
@@ -225,6 +286,7 @@ func TestRefundTaskQuota_Wallet(t *testing.T) {
seedChannel(t, channelID)
task := makeTask(userID, channelID, preConsumed, tokenID, BillingSourceWallet, 0)
+ persistTask(t, task)
RefundTaskQuota(ctx, task, "task failed: upstream error")
@@ -233,7 +295,7 @@ func TestRefundTaskQuota_Wallet(t *testing.T) {
// Token remain_quota should increase, used_quota should decrease
assert.Equal(t, tokenRemain+preConsumed, getTokenRemainQuota(t, tokenID))
- assert.Equal(t, -preConsumed, getTokenUsedQuota(t, tokenID))
+ assert.Equal(t, tokenRemain-preConsumed, getTokenUsedQuota(t, tokenID))
// A refund log should be created
log := getLastLog(t)
@@ -241,6 +303,9 @@ func TestRefundTaskQuota_Wallet(t *testing.T) {
assert.Equal(t, model.LogTypeRefund, log.Type)
assert.Equal(t, preConsumed, log.Quota)
assert.Equal(t, "test-model", log.ModelName)
+ var reloaded model.Task
+ require.NoError(t, model.DB.First(&reloaded, task.ID).Error)
+ assert.Zero(t, reloaded.Quota)
}
func TestRefundTaskQuota_Subscription(t *testing.T) {
@@ -258,6 +323,7 @@ func TestRefundTaskQuota_Subscription(t *testing.T) {
seedSubscription(t, subID, userID, subTotal, subUsed)
task := makeTask(userID, channelID, preConsumed, tokenID, BillingSourceSubscription, subID)
+ persistTask(t, task)
RefundTaskQuota(ctx, task, "subscription task failed")
@@ -280,6 +346,7 @@ func TestRefundTaskQuota_ZeroQuota(t *testing.T) {
seedUser(t, userID, 5000)
task := makeTask(userID, 0, 0, 0, BillingSourceWallet, 0)
+ persistTask(t, task)
RefundTaskQuota(ctx, task, "zero quota task")
@@ -301,6 +368,7 @@ func TestRefundTaskQuota_NoToken(t *testing.T) {
seedChannel(t, channelID)
task := makeTask(userID, channelID, preConsumed, 0, BillingSourceWallet, 0) // TokenId=0
+ persistTask(t, task)
RefundTaskQuota(ctx, task, "no token task failed")
@@ -313,6 +381,188 @@ func TestRefundTaskQuota_NoToken(t *testing.T) {
assert.Equal(t, model.LogTypeRefund, log.Type)
}
+func TestRefundTaskQuota_IsDurablyIdempotent(t *testing.T) {
+ truncate(t)
+ ctx := context.Background()
+
+ const userID, tokenID, channelID = 5, 5, 5
+ const initQuota, preConsumed, tokenRemain = 10000, 2500, 4000
+ seedUser(t, userID, initQuota)
+ seedToken(t, tokenID, userID, "sk-durable-refund", tokenRemain)
+ seedChannel(t, channelID)
+ task := makeTask(userID, channelID, preConsumed, tokenID, BillingSourceWallet, 0)
+ persistTask(t, task)
+
+ assert.True(t, RefundTaskQuota(ctx, task, "first failure"))
+ task.Quota = preConsumed // Simulate a stale copy held by another poller.
+ assert.True(t, RefundTaskQuota(ctx, task, "duplicate failure"))
+
+ assert.EqualValues(t, initQuota+preConsumed, getUserQuota(t, userID))
+ assert.Equal(t, tokenRemain+preConsumed, getTokenRemainQuota(t, tokenID))
+ assert.Equal(t, int64(1), countLogs(t))
+ var logOther map[string]interface{}
+ require.NoError(t, common.UnmarshalJsonStr(getLastLog(t).Other, &logOther))
+ assert.Equal(t, "first failure", logOther["reason"])
+ var reloaded model.Task
+ require.NoError(t, model.DB.First(&reloaded, task.ID).Error)
+ assert.Zero(t, reloaded.Quota)
+}
+
+func TestRefundTaskQuotaRefundsFundingWhenTokenWasDeleted(t *testing.T) {
+ truncate(t)
+ ctx := context.Background()
+
+ const userID, tokenID, channelID = 8, 8, 8
+ const initialUserQuota, preConsumed = 1000, 250
+ seedUser(t, userID, initialUserQuota)
+ seedToken(t, tokenID, userID, "deleted-before-refund", 2000)
+ seedChannel(t, channelID)
+ task := makeTask(userID, channelID, preConsumed, tokenID, BillingSourceWallet, 0)
+ persistTask(t, task)
+ require.NoError(t, model.DB.Delete(&model.Token{}, tokenID).Error)
+
+ assert.True(t, RefundTaskQuota(ctx, task, "token deleted while task was running"))
+ assert.EqualValues(t, initialUserQuota+preConsumed, getUserQuota(t, userID))
+ var reloaded model.Task
+ require.NoError(t, model.DB.First(&reloaded, task.ID).Error)
+ assert.Zero(t, reloaded.Quota)
+}
+
+func TestBackgroundTaskSettlementRecoveryRestoresLogAndUsageExactlyOnce(t *testing.T) {
+ truncate(t)
+ ctx := context.Background()
+ const userID, tokenID, channelID = 41, 42, 43
+ seedUser(t, userID, 1000)
+ seedChannel(t, channelID)
+ task := makeTask(userID, channelID, 10, tokenID, BillingSourceWallet, 0)
+ task.Status = model.TaskStatusSuccess
+ persistTask(t, task)
+
+ RecalculateTaskQuota(ctx, task, 15, "background recovery")
+ var pending model.BillingSettlement
+ require.NoError(t, model.DB.Where("operation_key = ?", "task:"+fmt.Sprint(task.ID)+":finalize").First(&pending).Error)
+ require.Equal(t, model.BillingSettlementStatusPending, pending.Status)
+
+ seedToken(t, tokenID, userID, "background-recovery-token", 100)
+ require.NoError(t, model.DB.Model(&model.BillingSettlement{}).Where("id = ?", pending.ID).Update("next_attempt", 0).Error)
+ model.ProcessPendingBillingSettlementsOnce()
+
+ var reloadedTask model.Task
+ require.NoError(t, model.DB.First(&reloadedTask, task.ID).Error)
+ assert.EqualValues(t, 15, reloadedTask.Quota)
+ var user model.User
+ require.NoError(t, model.DB.First(&user, userID).Error)
+ assert.EqualValues(t, 995, user.Quota)
+ assert.EqualValues(t, 5, user.UsedQuota)
+ assert.EqualValues(t, 1, user.RequestCount)
+ var token model.Token
+ require.NoError(t, model.DB.First(&token, tokenID).Error)
+ assert.EqualValues(t, 95, token.RemainQuota)
+ var channel model.Channel
+ require.NoError(t, model.DB.First(&channel, channelID).Error)
+ assert.EqualValues(t, 5, channel.UsedQuota)
+ assert.Equal(t, int64(1), countLogs(t))
+}
+
+func TestConcurrentTaskSettlementEffectProcessingIsIdempotent(t *testing.T) {
+ truncate(t)
+ const userID, tokenID, channelID = 51, 52, 53
+ seedUser(t, userID, 1000)
+ seedToken(t, tokenID, userID, "concurrent-effect-token", 100)
+ seedChannel(t, channelID)
+ task := makeTask(userID, channelID, 10, tokenID, BillingSourceWallet, 0)
+ task.Status = model.TaskStatusSuccess
+ persistTask(t, task)
+
+ operationKey := "task:" + fmt.Sprint(task.ID) + ":concurrent-effect"
+ _, _, err := model.ApplyBillingSettlementOnce(model.BillingSettlementInput{
+ OperationKey: operationKey, Source: model.BillingSettlementSourceWallet,
+ UserID: userID, TokenID: tokenID, FundingDelta: 5, TokenDelta: 5,
+ TaskID: task.ID, TaskQuota: 10, TaskQuotaTarget: 15,
+ Effect: &model.BillingSettlementEffect{
+ LogType: model.LogTypeConsume, Content: "concurrent effect",
+ ChannelID: channelID, ModelName: "test-model", TokenID: tokenID,
+ Group: "default", Other: map[string]interface{}{"task_id": task.TaskID},
+ UpdateUsage: true,
+ },
+ })
+ require.NoError(t, err)
+
+ var wg sync.WaitGroup
+ errs := make(chan error, 8)
+ for i := 0; i < 8; i++ {
+ wg.Add(1)
+ go func() {
+ defer wg.Done()
+ errs <- model.ProcessBillingSettlementEffect(operationKey)
+ }()
+ }
+ wg.Wait()
+ close(errs)
+ for effectErr := range errs {
+ require.NoError(t, effectErr)
+ }
+
+ assert.Equal(t, int64(1), countLogs(t))
+ var user model.User
+ var channel model.Channel
+ require.NoError(t, model.DB.First(&user, userID).Error)
+ require.NoError(t, model.DB.First(&channel, channelID).Error)
+ assert.EqualValues(t, 5, user.UsedQuota)
+ assert.EqualValues(t, 1, user.RequestCount)
+ assert.EqualValues(t, 5, channel.UsedQuota)
+}
+
+func TestTaskFinalizationRejectsChangedTargetWithoutSecondMutation(t *testing.T) {
+ truncate(t)
+ ctx := context.Background()
+ const userID, tokenID, channelID = 61, 62, 63
+ seedUser(t, userID, 1000)
+ seedToken(t, tokenID, userID, "stable-task-finalize-token", 100)
+ seedChannel(t, channelID)
+ task := makeTask(userID, channelID, 10, tokenID, BillingSourceWallet, 0)
+ task.Status = model.TaskStatusSuccess
+ persistTask(t, task)
+
+ RecalculateTaskQuota(ctx, task, 15, "first final result")
+ RecalculateTaskQuota(ctx, task, 20, "changed replay result")
+
+ assert.EqualValues(t, 995, getUserQuota(t, userID))
+ assert.Equal(t, 95, getTokenRemainQuota(t, tokenID))
+ assert.Equal(t, int64(1), countLogs(t))
+ var reloaded model.Task
+ require.NoError(t, model.DB.First(&reloaded, task.ID).Error)
+ assert.Equal(t, 15, reloaded.Quota)
+}
+
+func TestProcessSunoTaskResponsePersistsRefundIntentBeforeApplying(t *testing.T) {
+ truncate(t)
+ ctx := context.Background()
+
+ const userID, tokenID, channelID = 6, 6, 6
+ const pendingRefund, userQuota, tokenRemain = 250, 1000, 2000
+ seedUser(t, userID, userQuota)
+ seedToken(t, tokenID, userID, "sk-suno-pending-refund", tokenRemain)
+ seedChannel(t, channelID)
+ task := makeTask(userID, channelID, pendingRefund, tokenID, BillingSourceWallet, 0)
+ task.Status = model.TaskStatusInProgress
+ persistTask(t, task)
+
+ processSunoTaskResponse(ctx, task, dto.SunoDataResponse{
+ TaskID: task.TaskID, Status: string(model.TaskStatusFailure),
+ FailReason: "upstream failed", FinishTime: time.Now().Unix(),
+ })
+
+ assert.EqualValues(t, userQuota+pendingRefund, getUserQuota(t, userID))
+ assert.Equal(t, tokenRemain+pendingRefund, getTokenRemainQuota(t, tokenID))
+ var reloaded model.Task
+ require.NoError(t, model.DB.First(&reloaded, task.ID).Error)
+ assert.Zero(t, reloaded.Quota)
+ var settlement model.BillingSettlement
+ require.NoError(t, model.DB.Where("operation_key = ?", fmt.Sprintf("task:%d:refund", task.ID)).First(&settlement).Error)
+ assert.Equal(t, model.BillingSettlementStatusApplied, settlement.Status)
+}
+
// ===========================================================================
// RecalculateTaskQuota tests
// ===========================================================================
@@ -331,6 +581,7 @@ func TestRecalculate_PositiveDelta(t *testing.T) {
seedChannel(t, channelID)
task := makeTask(userID, channelID, preConsumed, tokenID, BillingSourceWallet, 0)
+ persistTask(t, task)
RecalculateTaskQuota(ctx, task, actualQuota, "adaptor adjustment")
@@ -376,6 +627,7 @@ func TestRecalculate_PositiveDeltaAttributesQuotaDataToSubmitNode(t *testing.T)
task := makeTask(userID, channelID, preConsumed, tokenID, BillingSourceWallet, 0)
task.PrivateData.NodeName = "submit-node"
+ persistTask(t, task)
RecalculateTaskQuota(ctx, task, actualQuota, "node attribution")
@@ -410,6 +662,7 @@ func TestRecalculate_NegativeDelta(t *testing.T) {
seedChannel(t, channelID)
task := makeTask(userID, channelID, preConsumed, tokenID, BillingSourceWallet, 0)
+ persistTask(t, task)
RecalculateTaskQuota(ctx, task, actualQuota, "adaptor adjustment")
@@ -483,6 +736,7 @@ func TestRecalculate_Subscription_NegativeDelta(t *testing.T) {
seedSubscription(t, subID, userID, subTotal, subUsed)
task := makeTask(userID, channelID, preConsumed, tokenID, BillingSourceSubscription, subID)
+ persistTask(t, task)
RecalculateTaskQuota(ctx, task, actualQuota, "subscription over-charge")
@@ -714,6 +968,7 @@ func TestSettle_PerCallBilling_SkipsAdaptorAdjust(t *testing.T) {
task := makeTask(userID, channelID, preConsumed, tokenID, BillingSourceWallet, 0)
task.PrivateData.BillingContext.PerCallBilling = true
+ persistTask(t, task)
adaptor := &mockAdaptor{adjustReturn: 2000}
taskResult := &relaycommon.TaskInfo{Status: model.TaskStatusSuccess}
@@ -741,6 +996,7 @@ func TestSettle_PerCallBilling_SkipsTotalTokens(t *testing.T) {
task := makeTask(userID, channelID, preConsumed, tokenID, BillingSourceWallet, 0)
task.PrivateData.BillingContext.PerCallBilling = true
+ persistTask(t, task)
adaptor := &mockAdaptor{adjustReturn: 0}
taskResult := &relaycommon.TaskInfo{Status: model.TaskStatusSuccess, TotalTokens: 9999}
@@ -768,6 +1024,7 @@ func TestSettle_DeltaSettlementDisabledSnapshotSkipsAdjustments(t *testing.T) {
task := makeTask(userID, channelID, preConsumed, tokenID, BillingSourceWallet, 0)
task.PrivateData.BillingContext.DeltaSettlementDisabled = common.GetPointer(true)
+ persistTask(t, task)
adaptor := &mockAdaptor{adjustReturn: 3000}
taskResult := &relaycommon.TaskInfo{Status: model.TaskStatusSuccess, TotalTokens: 9999}
@@ -794,6 +1051,7 @@ func TestSettle_DeltaSettlementDisabledChannelFallbackForLegacyTask(t *testing.T
task := makeTask(userID, channelID, preConsumed, tokenID, BillingSourceWallet, 0)
task.PrivateData.BillingContext.DeltaSettlementDisabled = nil
+ persistTask(t, task)
adaptor := &mockAdaptor{adjustReturn: 3000}
taskResult := &relaycommon.TaskInfo{Status: model.TaskStatusSuccess, TotalTokens: 9999}
@@ -823,6 +1081,7 @@ func TestSettle_DeltaSettlementSnapshotOverridesCurrentChannelSetting(t *testing
task := makeTask(userID, channelID, preConsumed, tokenID, BillingSourceWallet, 0)
task.PrivateData.BillingContext.DeltaSettlementDisabled = common.GetPointer(false)
+ persistTask(t, task)
adaptor := &mockAdaptor{adjustReturn: adaptorQuota}
taskResult := &relaycommon.TaskInfo{Status: model.TaskStatusSuccess}
@@ -852,6 +1111,7 @@ func TestSettle_DeltaSettlementDisabledIgnoredForNonDoubaoVideo(t *testing.T) {
task := makeTask(userID, channelID, preConsumed, tokenID, BillingSourceWallet, 0)
task.PrivateData.BillingContext.DeltaSettlementDisabled = common.GetPointer(true)
+ persistTask(t, task)
adaptor := &mockAdaptor{adjustReturn: adaptorQuota}
taskResult := &relaycommon.TaskInfo{Status: model.TaskStatusSuccess}
@@ -881,6 +1141,7 @@ func TestSettle_NonPerCall_AdaptorAdjustWorks(t *testing.T) {
task := makeTask(userID, channelID, preConsumed, tokenID, BillingSourceWallet, 0)
// PerCallBilling defaults to false
+ persistTask(t, task)
adaptor := &mockAdaptor{adjustReturn: adaptorQuota}
taskResult := &relaycommon.TaskInfo{Status: model.TaskStatusSuccess}
diff --git a/service/task_polling.go b/service/task_polling.go
index 89a83546..8981e158 100644
--- a/service/task_polling.go
+++ b/service/task_polling.go
@@ -35,6 +35,8 @@ type TaskPollingAdaptor interface {
// 打破 service -> relay -> relay/channel -> service 的循环依赖。
var GetTaskAdaptorFunc func(platform constant.TaskPlatform) TaskPollingAdaptor
+const legacyTaskRefundCutoff int64 = 1740182400 // 2025-02-22 00:00:00 UTC
+
// sweepTimedOutTasks 在主轮询之前独立清理超时任务。
// 每次最多处理 100 条，剩余的下个周期继续处理。
// 使用 per-task CAS (UpdateWithStatus) 防止覆盖被正常轮询已推进的任务。
@@ -48,14 +50,14 @@ func sweepTimedOutTasks(ctx context.Context) {
return
}
- const legacyTaskCutoff int64 = 1740182400 // 2026-02-22 00:00:00 UTC
reason := fmt.Sprintf("任务超时（%d分钟）", constant.TaskTimeoutMinutes)
legacyReason := "任务超时（旧系统遗留任务，不进行退款，请联系管理员）"
now := time.Now().Unix()
timedOutCount := 0
for _, task := range tasks {
- isLegacy := task.SubmitTime > 0 && task.SubmitTime < legacyTaskCutoff
+ isLegacy := task.SubmitTime > 0 && task.SubmitTime < legacyTaskRefundCutoff
+ isUnconfirmedSubmit := task.PrivateData.AwaitingUpstreamID
oldStatus := task.Status
task.Status = model.TaskStatusFailure
@@ -63,11 +65,23 @@ func sweepTimedOutTasks(ctx context.Context) {
task.FinishTime = now
if isLegacy {
task.FailReason = legacyReason
+ } else if isUnconfirmedSubmit {
+ task.FailReason = "任务提交超时（无法确认上游是否已创建任务，未自动退款，请人工核对）"
} else {
task.FailReason = reason
}
- won, err := task.UpdateWithStatus(oldStatus)
+ var settlement *model.BillingSettlementInput
+ if !isLegacy && !isUnconfirmedSubmit {
+ settlement = buildTaskRefundSettlementInput(task, reason)
+ }
+ var won bool
+ var err error
+ if settlement != nil {
+ won, err = task.UpdateWithStatusAndSettlement(oldStatus, *settlement)
+ } else {
+ won, err = task.UpdateWithStatus(oldStatus)
+ }
if err != nil {
logger.LogError(ctx, fmt.Sprintf("sweepTimedOutTasks CAS update error for task %s: %v", task.TaskID, err))
continue
@@ -77,8 +91,8 @@ func sweepTimedOutTasks(ctx context.Context) {
continue
}
timedOutCount++
- if !isLegacy && task.Quota != 0 {
- RefundTaskQuota(ctx, task, reason)
+ if settlement != nil {
+ applyTaskBillingSettlement(ctx, task, settlement)
}
}
@@ -105,28 +119,16 @@ func TaskPollingLoop() {
}
taskChannelM := make(map[int][]string)
taskM := make(map[string]*model.Task)
- nullTaskIds := make([]int64, 0)
for _, task := range tasks {
upstreamID := task.GetUpstreamTaskID()
if upstreamID == "" {
- // 统计失败的未完成任务
- nullTaskIds = append(nullTaskIds, task.ID)
+ // Submission placeholders remain non-terminal until the timeout
+ // sweeper can classify them for manual review.
continue
}
taskM[upstreamID] = task
taskChannelM[task.ChannelId] = append(taskChannelM[task.ChannelId], upstreamID)
}
- if len(nullTaskIds) > 0 {
- err := model.TaskBulkUpdateByID(nullTaskIds, map[string]any{
- "status": "FAILURE",
- "progress": "100%",
- })
- if err != nil {
- logger.LogError(ctx, fmt.Sprintf("Fix null task_id task error: %v", err))
- } else {
- logger.LogInfo(ctx, fmt.Sprintf("Fix null task_id task success: %v", nullTaskIds))
- }
- }
if len(taskChannelM) == 0 {
continue
}
@@ -222,31 +224,55 @@ func updateSunoTasks(ctx context.Context, channelId int, taskIds []string, taskM
for _, responseItem := range responseItems.Data {
task := taskM[responseItem.TaskID]
- if !taskNeedsUpdate(task, responseItem) {
+ if task == nil {
+ logger.LogWarn(ctx, fmt.Sprintf("Suno task response ignored: unknown task_id=%s", responseItem.TaskID))
continue
}
+ processSunoTaskResponse(ctx, task, responseItem)
+ }
+ return nil
+}
- task.Status = lo.If(model.TaskStatus(responseItem.Status) != "", model.TaskStatus(responseItem.Status)).Else(task.Status)
- task.FailReason = lo.If(responseItem.FailReason != "", responseItem.FailReason).Else(task.FailReason)
- task.SubmitTime = lo.If(responseItem.SubmitTime != 0, responseItem.SubmitTime).Else(task.SubmitTime)
- task.StartTime = lo.If(responseItem.StartTime != 0, responseItem.StartTime).Else(task.StartTime)
- task.FinishTime = lo.If(responseItem.FinishTime != 0, responseItem.FinishTime).Else(task.FinishTime)
- if responseItem.FailReason != "" || task.Status == model.TaskStatusFailure {
- logger.LogInfo(ctx, task.TaskID+" 构建失败，"+task.FailReason)
- task.Progress = "100%"
- RefundTaskQuota(ctx, task, task.FailReason)
- }
- if responseItem.Status == model.TaskStatusSuccess {
- task.Progress = "100%"
- }
- task.Data = responseItem.Data
+func processSunoTaskResponse(ctx context.Context, task *model.Task, responseItem dto.SunoDataResponse) {
+ if !taskNeedsUpdate(task, responseItem) {
+ return
+ }
- err = task.Update()
- if err != nil {
- common.SysLog("UpdateSunoTask task error: " + err.Error())
- }
+ previousStatus := task.Status
+ task.Status = lo.If(model.TaskStatus(responseItem.Status) != "", model.TaskStatus(responseItem.Status)).Else(task.Status)
+ task.FailReason = lo.If(responseItem.FailReason != "", responseItem.FailReason).Else(task.FailReason)
+ task.SubmitTime = lo.If(responseItem.SubmitTime != 0, responseItem.SubmitTime).Else(task.SubmitTime)
+ task.StartTime = lo.If(responseItem.StartTime != 0, responseItem.StartTime).Else(task.StartTime)
+ task.FinishTime = lo.If(responseItem.FinishTime != 0, responseItem.FinishTime).Else(task.FinishTime)
+ isFailure := responseItem.FailReason != "" || task.Status == model.TaskStatusFailure
+ if isFailure {
+ logger.LogInfo(ctx, task.TaskID+" 构建失败，"+task.FailReason)
+ task.Status = model.TaskStatusFailure
+ task.Progress = "100%"
+ }
+ if responseItem.Status == model.TaskStatusSuccess {
+ task.Progress = "100%"
+ }
+ task.Data = responseItem.Data
+
+ var settlement *model.BillingSettlementInput
+ if isFailure {
+ settlement = buildTaskRefundSettlementInput(task, task.FailReason)
+ }
+ var won bool
+ var err error
+ if settlement != nil {
+ won, err = task.UpdateWithStatusAndSettlement(previousStatus, *settlement)
+ } else {
+ won, err = task.UpdateWithStatus(previousStatus)
+ }
+ if err != nil {
+ common.SysLog("UpdateSunoTask task error: " + err.Error())
+ } else if !won {
+ logger.LogWarn(ctx, fmt.Sprintf("Task %s CAS lost or no-op update, skip billing", task.TaskID))
+ } else if settlement != nil {
+ applyTaskBillingSettlement(ctx, task, settlement)
}
- return nil
}
// taskNeedsUpdate 检查 Suno 任务是否需要更新
@@ -480,8 +506,21 @@ func updateVideoSingleTask(ctx context.Context, adaptor TaskPollingAdaptor, ch *
}
isDone := task.Status == model.TaskStatusSuccess || task.Status == model.TaskStatusFailure
+ var settlement *model.BillingSettlementInput
+ transitionWon := false
if isDone && snap.Status != task.Status {
- won, err := task.UpdateWithStatus(snap.Status)
+ if shouldSettle {
+ settlement = prepareTaskCompletionSettlement(ctx, adaptor, task, taskResult, ch.Type, ch.GetOtherSettings())
+ } else if shouldRefund {
+ settlement = buildTaskRefundSettlementInput(task, task.FailReason)
+ }
+ var won bool
+ var err error
+ if settlement != nil {
+ won, err = task.UpdateWithStatusAndSettlement(snap.Status, *settlement)
+ } else {
+ won, err = task.UpdateWithStatus(snap.Status)
+ }
if err != nil {
logger.LogError(ctx, fmt.Sprintf("UpdateWithStatus failed for task %s: %s", task.TaskID, err.Error()))
shouldRefund = false
@@ -490,6 +529,8 @@ func updateVideoSingleTask(ctx context.Context, adaptor TaskPollingAdaptor, ch *
logger.LogWarn(ctx, fmt.Sprintf("Task %s already transitioned by another process, skip billing", task.TaskID))
shouldRefund = false
shouldSettle = false
+ } else {
+ transitionWon = true
}
} else if !snap.Equal(task.Snapshot()) {
if _, err := task.UpdateWithStatus(snap.Status); err != nil {
@@ -500,11 +541,8 @@ func updateVideoSingleTask(ctx context.Context, adaptor TaskPollingAdaptor, ch *
logger.LogDebug(ctx, "No update needed for task %s", task.TaskID)
}
- if shouldSettle {
- settleTaskBillingOnComplete(ctx, adaptor, task, taskResult, ch.Type, ch.GetOtherSettings())
- }
- if shouldRefund {
- RefundTaskQuota(ctx, task, task.FailReason)
+ if transitionWon && settlement != nil {
+ applyTaskBillingSettlement(ctx, task, settlement)
}
return nil
@@ -564,24 +602,30 @@ func truncateBase64(s string) string {
// 2. taskResult.TotalTokens > 0 → 按 token 重算
// 3. 都不满足 → 保持预扣额度不变
func settleTaskBillingOnComplete(ctx context.Context, adaptor TaskPollingAdaptor, task *model.Task, taskResult *relaycommon.TaskInfo, channelType int, channelSettings ...dto.ChannelOtherSettings) {
+ applyTaskBillingSettlement(ctx, task, prepareTaskCompletionSettlement(ctx, adaptor, task, taskResult, channelType, channelSettings...))
+}
+
+func prepareTaskCompletionSettlement(ctx context.Context, adaptor TaskPollingAdaptor, task *model.Task, taskResult *relaycommon.TaskInfo, channelType int, channelSettings ...dto.ChannelOtherSettings) *model.BillingSettlementInput {
// 0. 按次计费的任务不做差额结算
if bc := task.PrivateData.BillingContext; bc != nil && bc.PerCallBilling {
logger.LogInfo(ctx, fmt.Sprintf("任务 %s 按次计费，跳过差额结算", task.TaskID))
- return
+ return nil
}
if task.DeltaSettlementDisabledForChannel(channelType, channelSettings...) {
logger.LogInfo(ctx, fmt.Sprintf("任务 %s 已按渠道设置跳过完成态差额结算", task.TaskID))
- return
+ return nil
}
// 1. 优先让 adaptor 决定最终额度
if actualQuota := adaptor.AdjustBillingOnComplete(task, taskResult); actualQuota > 0 {
- RecalculateTaskQuota(ctx, task, actualQuota, "adaptor计费调整")
- return
+ return buildTaskFinalSettlementInput(task, actualQuota, "adaptor计费调整")
}
// 2. 回退到 token 重算
if taskResult.TotalTokens > 0 {
- RecalculateTaskQuotaByTokens(ctx, task, taskResult.TotalTokens)
- return
+ actualQuota, reason, clamp, ok := calculateTaskQuotaByTokens(task, taskResult.TotalTokens)
+ if ok {
+ return buildTaskFinalSettlementInput(task, actualQuota, reason, clamp)
+ }
}
// 3. 无调整，保持预扣额度
+ return nil
}
diff --git a/service/violation_fee.go b/service/violation_fee.go
index bf77175a..caa232e1 100644
--- a/service/violation_fee.go
+++ b/service/violation_fee.go
@@ -99,38 +99,47 @@ func calcViolationFeeQuota(amount, groupRatio float64) int {
return int(quota)
}
-// ChargeViolationFeeIfNeeded charges an additional fee after the normal flow finishes (including refund).
-// It uses Grok fee settings as the fee policy.
+type violationFeeChargeResult struct {
+ applicable bool
+ settled bool
+}
+
+// HandleFailedBilling finalizes a failed request as either a violation fee or
+// a normal refund. Once a violation settlement intent exists, it must not be
+// replaced by a refund using the same request-finalize operation key.
+func HandleFailedBilling(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, apiErr *types.MaxAPIError) {
+ result := chargeViolationFeeIfNeeded(ctx, relayInfo, apiErr)
+ if !result.applicable && relayInfo != nil && relayInfo.Billing != nil {
+ relayInfo.Billing.Refund(ctx)
+ }
+}
+
func ChargeViolationFeeIfNeeded(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, apiErr *types.MaxAPIError) bool {
+ return chargeViolationFeeIfNeeded(ctx, relayInfo, apiErr).settled
+}
+
+func chargeViolationFeeIfNeeded(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, apiErr *types.MaxAPIError) violationFeeChargeResult {
if ctx == nil || relayInfo == nil || apiErr == nil {
- return false
+ return violationFeeChargeResult{}
}
//if relayInfo.IsPlayground {
// return false
//}
if !shouldChargeViolationFee(apiErr) {
- return false
+ return violationFeeChargeResult{}
}
settings := model_setting.GetGrokSettings()
if settings == nil || !settings.ViolationDeductionEnabled {
- return false
+ return violationFeeChargeResult{}
}
groupRatio := relayInfo.PriceData.GroupRatioInfo.GroupRatio
feeQuota := calcViolationFeeQuota(settings.ViolationDeductionAmount, groupRatio)
if feeQuota <= 0 {
- return false
- }
-
- if err := PostConsumeQuota(relayInfo, feeQuota, 0, true); err != nil {
- logger.LogError(ctx, fmt.Sprintf("failed to charge violation fee: %s", err.Error()))
- return false
+ return violationFeeChargeResult{}
}
- model.UpdateUserUsedQuotaAndRequestCount(relayInfo.UserId, feeQuota)
- model.UpdateChannelUsedQuota(relayInfo.ChannelId, feeQuota)
-
useTimeSeconds := time.Now().Unix() - relayInfo.StartTime.Unix()
tokenName := ctx.GetString("token_name")
oai := apiErr.ToOpenAIError()
@@ -146,6 +155,29 @@ func ChargeViolationFeeIfNeeded(ctx *gin.Context, relayInfo *relaycommon.RelayIn
"upstream_error_code": fmt.Sprintf("%v", oai.Code),
"violation_fee_marker": CSAMViolationMarker,
}
+ effect := &model.BillingSettlementEffect{
+ LogType: model.LogTypeConsume, Content: "Violation fee charged",
+ ChannelID: relayInfo.ChannelId, ModelName: relayInfo.OriginModelName,
+ TokenID: relayInfo.TokenId, Group: relayInfo.UsingGroup, Other: other,
+ UpdateUsage: true, Quota: int64(feeQuota),
+ }
+
+ if settler, ok := relayInfo.Billing.(interface {
+ SettleWithEffect(int, *model.BillingSettlementEffect) error
+ }); ok {
+ if err := settler.SettleWithEffect(feeQuota, effect); err != nil {
+ logger.LogError(ctx, fmt.Sprintf("violation fee settlement remains pending/manual: %s", err.Error()))
+ return violationFeeChargeResult{applicable: true}
+ }
+ return violationFeeChargeResult{applicable: true, settled: true}
+ }
+
+ if err := SettleBilling(ctx, relayInfo, feeQuota); err != nil {
+ logger.LogError(ctx, fmt.Sprintf("violation fee settlement remains pending/manual: %s", err.Error()))
+ return violationFeeChargeResult{applicable: true}
+ }
+ model.UpdateUserUsedQuotaAndRequestCount(relayInfo.UserId, feeQuota)
+ model.UpdateChannelUsedQuota(relayInfo.ChannelId, feeQuota)
model.RecordConsumeLog(ctx, relayInfo.UserId, model.RecordConsumeLogParams{
ChannelId: relayInfo.ChannelId,
@@ -160,5 +192,5 @@ func ChargeViolationFeeIfNeeded(ctx *gin.Context, relayInfo *relaycommon.RelayIn
Other: other,
})
- return true
+ return violationFeeChargeResult{applicable: true, settled: true}
}
diff --git a/setting/billing_setting/tiered_billing.go b/setting/billing_setting/tiered_billing.go
index bde2d7bc..c9f87b02 100644
--- a/setting/billing_setting/tiered_billing.go
+++ b/setting/billing_setting/tiered_billing.go
@@ -5,6 +5,7 @@ import (
"github.com/MAX-API-Next/MAX-API/pkg/billingexpr"
"github.com/MAX-API-Next/MAX-API/setting/config"
+ "github.com/MAX-API-Next/MAX-API/types"
"github.com/samber/lo"
)
@@ -18,13 +19,13 @@ const (
// BillingSetting is managed by config.GlobalConfig.Register.
// DB keys: billing_setting.billing_mode, billing_setting.billing_expr
type BillingSetting struct {
- BillingMode map[string]string `json:"billing_mode"`
- BillingExpr map[string]string `json:"billing_expr"`
+ BillingMode *types.RWMap[string, string] `json:"billing_mode"`
+ BillingExpr *types.RWMap[string, string] `json:"billing_expr"`
}
var billingSetting = BillingSetting{
- BillingMode: make(map[string]string),
- BillingExpr: make(map[string]string),
+ BillingMode: types.NewRWMap[string, string](),
+ BillingExpr: types.NewRWMap[string, string](),
}
func init() {
@@ -36,23 +37,22 @@ func init() {
// ---------------------------------------------------------------------------
func GetBillingMode(model string) string {
- if mode, ok := billingSetting.BillingMode[model]; ok {
+ if mode, ok := billingSetting.BillingMode.Get(model); ok {
return mode
}
return BillingModeRatio
}
func GetBillingExpr(model string) (string, bool) {
- expr, ok := billingSetting.BillingExpr[model]
- return expr, ok
+ return billingSetting.BillingExpr.Get(model)
}
func GetBillingModeCopy() map[string]string {
- return lo.Assign(billingSetting.BillingMode)
+ return billingSetting.BillingMode.ReadAll()
}
func GetBillingExprCopy() map[string]string {
- return lo.Assign(billingSetting.BillingExpr)
+ return billingSetting.BillingExpr.ReadAll()
}
func GetPricingSyncData(base map[string]any) map[string]any {
diff --git a/setting/billing_setting/tiered_billing_test.go b/setting/billing_setting/tiered_billing_test.go
new file mode 100644
index 00000000..393c651f
--- /dev/null
+++ b/setting/billing_setting/tiered_billing_test.go
@@ -0,0 +1,44 @@
+package billing_setting
+
+import (
+ "sync"
+ "testing"
+
+ "github.com/MAX-API-Next/MAX-API/setting/config"
+ "github.com/MAX-API-Next/MAX-API/types"
+ "github.com/stretchr/testify/require"
+)
+
+func TestBillingSettingMapsSupportConcurrentReloads(t *testing.T) {
+ setting := BillingSetting{
+ BillingMode: types.NewRWMap[string, string](),
+ BillingExpr: types.NewRWMap[string, string](),
+ }
+
+ var readers sync.WaitGroup
+ for range 4 {
+ readers.Add(1)
+ go func() {
+ defer readers.Done()
+ for range 200 {
+ _, _ = setting.BillingMode.Get("model")
+ _, _ = setting.BillingExpr.Get("model")
+ }
+ }()
+ }
+
+ for i := 0; i < 200; i++ {
+ require.NoError(t, config.UpdateConfigFromMap(&setting, map[string]string{
+ BillingModeField: `{"model":"tiered_expr"}`,
+ BillingExprField: `{"model":"p + c"}`,
+ }))
+ }
+ readers.Wait()
+
+ mode, ok := setting.BillingMode.Get("model")
+ require.True(t, ok)
+ require.Equal(t, BillingModeTieredExpr, mode)
+ expr, ok := setting.BillingExpr.Get("model")
+ require.True(t, ok)
+ require.Equal(t, "p + c", expr)
+}
diff --git a/setting/config/config.go b/setting/config/config.go
index 27007b3d..25f96ae8 100644
--- a/setting/config/config.go
+++ b/setting/config/config.go
@@ -18,6 +18,12 @@ type ConfigManager struct {
var GlobalConfig = NewConfigManager()
+type jsonUnmarshaler interface {
+ UnmarshalJSON([]byte) error
+}
+
+var jsonUnmarshalerType = reflect.TypeOf((*jsonUnmarshaler)(nil)).Elem()
+
func NewConfigManager() *ConfigManager {
return &ConfigManager{
configs: make(map[string]interface{}),
@@ -172,6 +178,9 @@ func updateConfigFromMap(config interface{}, configMap map[string]string) error
if val.Kind() != reflect.Struct {
return fmt.Errorf("config must point to a struct, got %s", val.Kind())
}
+ if err := validateConfigPointerNulls(val, configMap); err != nil {
+ return err
+ }
typ := val.Type()
for i := 0; i < val.NumField(); i++ {
@@ -281,6 +290,30 @@ func updateConfigFromMap(config interface{}, configMap map[string]string) error
return nil
}
+func validateConfigPointerNulls(val reflect.Value, configMap map[string]string) error {
+ typ := val.Type()
+ for i := 0; i < val.NumField(); i++ {
+ field := val.Field(i)
+ fieldType := typ.Field(i)
+ if !fieldType.IsExported() || !field.CanSet() {
+ continue
+ }
+
+ key := fieldType.Tag.Get("json")
+ if key == "" || key == "-" {
+ key = fieldType.Name
+ }
+ strValue, ok := configMap[key]
+ if !ok || strings.TrimSpace(strValue) != "null" {
+ continue
+ }
+ if field.Kind() == reflect.Ptr && field.Type().Implements(jsonUnmarshalerType) {
+ return fmt.Errorf("invalid value for %s: null is not allowed", key)
+ }
+ }
+ return nil
+}
+
// ConfigToMap 将配置对象转换为map（导出函数）
func ConfigToMap(config interface{}) (map[string]string, error) {
return configToMap(config)
@@ -291,6 +324,19 @@ func UpdateConfigFromMap(config interface{}, configMap map[string]string) error
return updateConfigFromMap(config, configMap)
}
+// ValidateConfigPointerNulls rejects null for pointer-backed JSON config fields without mutating the config.
+func ValidateConfigPointerNulls(config interface{}, configMap map[string]string) error {
+ val := reflect.ValueOf(config)
+ if val.Kind() != reflect.Ptr {
+ return fmt.Errorf("config must be a pointer, got %s", val.Kind())
+ }
+ val = val.Elem()
+ if val.Kind() != reflect.Struct {
+ return fmt.Errorf("config must point to a struct, got %s", val.Kind())
+ }
+ return validateConfigPointerNulls(val, configMap)
+}
+
// ExportAllConfigs 导出所有已注册的配置为扁平结构
func (cm *ConfigManager) ExportAllConfigs() map[string]string {
cm.mutex.RLock()
diff --git a/setting/config/config_test.go b/setting/config/config_test.go
index 5c0c496f..df134d14 100644
--- a/setting/config/config_test.go
+++ b/setting/config/config_test.go
@@ -117,6 +117,18 @@ func TestUpdateConfigFromMap_PointerUnmarshalerUpdatesInPlace(t *testing.T) {
assert.Equal(t, map[string]float64{"new": 2}, cfg.Ratios.ReadAll())
}
+func TestUpdateConfigFromMapRejectsNullRWMap(t *testing.T) {
+ ratios := types.NewRWMap[string, float64]()
+ ratios.Set("existing", 1)
+ cfg := &testConfigWithRWMap{Ratios: ratios}
+
+ err := UpdateConfigFromMap(cfg, map[string]string{"ratios": " null "})
+
+ require.Error(t, err)
+ require.Same(t, ratios, cfg.Ratios)
+ assert.Equal(t, map[string]float64{"existing": 1}, cfg.Ratios.ReadAll())
+}
+
func TestLoadFromDBReturnsInvalidSubConfigError(t *testing.T) {
manager := NewConfigManager()
cfg := &testConfigWithMap{Modes: map[string]string{"existing": "tiered_expr"}}
diff --git a/setting/ratio_setting/cache_ratio.go b/setting/ratio_setting/cache_ratio.go
index 52496823..3399aeb9 100644
--- a/setting/ratio_setting/cache_ratio.go
+++ b/setting/ratio_setting/cache_ratio.go
@@ -144,12 +144,12 @@ func CreateCacheRatio2JSONString() string {
// UpdateCacheRatioByJSONString updates the cache ratio map from a JSON string
func UpdateCacheRatioByJSONString(jsonStr string) error {
- return types.LoadFromJsonStringWithCallback(cacheRatioMap, jsonStr, InvalidateExposedDataCache)
+ return loadPricingMap(cacheRatioMap, jsonStr)
}
// UpdateCreateCacheRatioByJSONString updates the create cache ratio map from a JSON string
func UpdateCreateCacheRatioByJSONString(jsonStr string) error {
- return types.LoadFromJsonStringWithCallback(createCacheRatioMap, jsonStr, InvalidateExposedDataCache)
+ return loadPricingMap(createCacheRatioMap, jsonStr)
}
// GetCacheRatio returns the cache ratio for a model
diff --git a/setting/ratio_setting/model_ratio.go b/setting/ratio_setting/model_ratio.go
index 5d95164e..07011079 100644
--- a/setting/ratio_setting/model_ratio.go
+++ b/setting/ratio_setting/model_ratio.go
@@ -366,7 +366,7 @@ func ModelPrice2JSONString() string {
}
func UpdateModelPriceByJSONString(jsonStr string) error {
- return types.LoadFromJsonStringWithCallback(modelPriceMap, jsonStr, InvalidateExposedDataCache)
+ return loadPricingMap(modelPriceMap, jsonStr)
}
// GetModelPrice 返回模型的价格，如果模型不存在则返回-1，false
@@ -395,7 +395,7 @@ func GetModelPrice(name string, printErr bool) (float64, bool) {
}
func UpdateModelRatioByJSONString(jsonStr string) error {
- return types.LoadFromJsonStringWithCallback(modelRatioMap, jsonStr, InvalidateExposedDataCache)
+ return loadPricingMap(modelRatioMap, jsonStr)
}
// 处理带有思考预算的模型名称，方便统一定价
@@ -443,7 +443,7 @@ func CompletionRatio2JSONString() string {
}
func UpdateCompletionRatioByJSONString(jsonStr string) error {
- return types.LoadFromJsonStringWithCallback(completionRatioMap, jsonStr, InvalidateExposedDataCache)
+ return loadPricingMap(completionRatioMap, jsonStr)
}
func GetCompletionRatio(name string) float64 {
@@ -676,7 +676,7 @@ func ImageRatio2JSONString() string {
}
func UpdateImageRatioByJSONString(jsonStr string) error {
- return types.LoadFromJsonString(imageRatioMap, jsonStr)
+ return loadPricingMap(imageRatioMap, jsonStr)
}
func GetImageRatio(name string) (float64, bool) {
@@ -692,7 +692,7 @@ func AudioRatio2JSONString() string {
}
func UpdateAudioRatioByJSONString(jsonStr string) error {
- return types.LoadFromJsonStringWithCallback(audioRatioMap, jsonStr, InvalidateExposedDataCache)
+ return loadPricingMap(audioRatioMap, jsonStr)
}
func AudioCompletionRatio2JSONString() string {
@@ -700,7 +700,7 @@ func AudioCompletionRatio2JSONString() string {
}
func UpdateAudioCompletionRatioByJSONString(jsonStr string) error {
- return types.LoadFromJsonStringWithCallback(audioCompletionRatioMap, jsonStr, InvalidateExposedDataCache)
+ return loadPricingMap(audioCompletionRatioMap, jsonStr)
}
func GetModelRatioCopy() map[string]float64 {
diff --git a/setting/ratio_setting/pricing_validation.go b/setting/ratio_setting/pricing_validation.go
new file mode 100644
index 00000000..2cfc10c3
--- /dev/null
+++ b/setting/ratio_setting/pricing_validation.go
@@ -0,0 +1,35 @@
+package ratio_setting
+
+import (
+ "fmt"
+ "math"
+
+ "github.com/MAX-API-Next/MAX-API/common"
+ "github.com/MAX-API-Next/MAX-API/types"
+)
+
+func ValidatePricingMapJSONString(jsonStr string) error {
+ var values map[string]*float64
+ if err := common.UnmarshalJsonStr(jsonStr, &values); err != nil {
+ return err
+ }
+ if values == nil {
+ return fmt.Errorf("pricing configuration must be a JSON object")
+ }
+ for modelName, value := range values {
+ if value == nil {
+ return fmt.Errorf("pricing value for model %q must not be null", modelName)
+ }
+ if *value < 0 || math.IsNaN(*value) || math.IsInf(*value, 0) {
+ return fmt.Errorf("pricing value for model %q must be finite and non-negative", modelName)
+ }
+ }
+ return nil
+}
+
+func loadPricingMap(m *types.RWMap[string, float64], jsonStr string) error {
+ if err := ValidatePricingMapJSONString(jsonStr); err != nil {
+ return err
+ }
+ return types.LoadFromJsonStringWithCallback(m, jsonStr, InvalidateExposedDataCache)
+}
diff --git a/setting/task_billing_setting/rate_card.go b/setting/task_billing_setting/rate_card.go
index 2820ee85..966c98f3 100644
--- a/setting/task_billing_setting/rate_card.go
+++ b/setting/task_billing_setting/rate_card.go
@@ -30,11 +30,11 @@ type RateCard struct {
}
type TaskBillingSetting struct {
- RateCards map[string]RateCard `json:"rate_cards"`
+ RateCards *types.RWMap[string, RateCard] `json:"rate_cards"`
}
var taskBillingSetting = TaskBillingSetting{
- RateCards: defaultRateCards(),
+ RateCards: newRateCardMap(defaultRateCards()),
}
func init() {
@@ -42,7 +42,7 @@ func init() {
}
func GetRateCardsCopy() map[string]RateCard {
- return cloneRateCards(taskBillingSetting.RateCards)
+ return cloneRateCards(taskBillingSetting.RateCards.ReadAll())
}
func GetRateCardCopy(models ...string) (*RateCard, string) {
@@ -137,7 +137,7 @@ func validQuantity(value float64) bool {
}
func findRateCard(models ...string) (*RateCard, string) {
- rateCards := taskBillingSetting.RateCards
+ rateCards := taskBillingSetting.RateCards.ReadAll()
for _, model := range models {
model = strings.TrimSpace(model)
if model == "" {
@@ -170,6 +170,12 @@ func findRateCard(models ...string) (*RateCard, string) {
return nil, ""
}
+func newRateCardMap(rateCards map[string]RateCard) *types.RWMap[string, RateCard] {
+ m := types.NewRWMap[string, RateCard]()
+ m.AddAll(rateCards)
+ return m
+}
+
func matchPattern(pattern, value string) bool {
matched, err := path.Match(pattern, value)
if err == nil {
diff --git a/setting/task_billing_setting/rate_card_test.go b/setting/task_billing_setting/rate_card_test.go
index f578511c..3d4f27aa 100644
--- a/setting/task_billing_setting/rate_card_test.go
+++ b/setting/task_billing_setting/rate_card_test.go
@@ -2,10 +2,13 @@ package task_billing_setting
import (
"strings"
+ "sync"
"testing"
"github.com/MAX-API-Next/MAX-API/common"
+ "github.com/MAX-API-Next/MAX-API/setting/config"
"github.com/MAX-API-Next/MAX-API/types"
+ "github.com/stretchr/testify/require"
)
func withQuotaPerUnit(t *testing.T, value float64) {
@@ -75,3 +78,28 @@ func TestCalculateUnknownModelFallsBack(t *testing.T) {
t.Fatalf("Calculate returned %+v, want nil", got)
}
}
+
+func TestRateCardsSupportConcurrentReloads(t *testing.T) {
+ setting := TaskBillingSetting{RateCards: newRateCardMap(nil)}
+ raw := `{"model":{"rows":[{"match":{"quality":"std"},"unit_price":1}]}}`
+
+ var readers sync.WaitGroup
+ for range 4 {
+ readers.Add(1)
+ go func() {
+ defer readers.Done()
+ for range 200 {
+ _, _ = setting.RateCards.Get("model")
+ }
+ }()
+ }
+
+ for range 200 {
+ require.NoError(t, config.UpdateConfigFromMap(&setting, map[string]string{"rate_cards": raw}))
+ }
+ readers.Wait()
+
+ card, ok := setting.RateCards.Get("model")
+ require.True(t, ok)
+ require.Len(t, card.Rows, 1)
+}
diff --git a/tools/jsonwrapcheck/allowlist.txt b/tools/jsonwrapcheck/allowlist.txt
index c214a600..b1473e14 100644
--- a/tools/jsonwrapcheck/allowlist.txt
+++ b/tools/jsonwrapcheck/allowlist.txt
@@ -84,12 +84,10 @@ oauth/oidc.go|*OIDCProvider.GetUserInfo|NewDecoder|64ee82d1a886a8c84d5771d6a7733
relay/channel/ali/rerank.go|RerankHandler|Unmarshal|66eb02820ea009d9d095c26550faf57ec59dba488d801f0cf3a611d01b3cfcf7
relay/channel/ali/rerank.go|RerankHandler|Marshal|ea6233db65ed7e55f6f9ae4ca1f4d9b9ca04f9e508ab5eef3dc68f139231720d
relay/channel/aws/dto.go|formatRequest|Marshal|0811da817c96b432dc92949bfc702d029e13adb47ab075abd27b6b82eae30478
-relay/channel/aws/relay-aws.go|handleNovaRequest|Unmarshal|d42307b92b274aacb69d483fb90f0053be946f9f9e1ec208c3f6d526cb27d9d0
relay/channel/baidu/relay-baidu.go|baiduHandler|Unmarshal|6c88b37cda6b69cb880852f57074d792626b82bd771acbab3a46754f80631684
relay/channel/baidu/relay-baidu.go|baiduHandler|Marshal|92981973316a59184c1f9d4fdca6fe3a1045b5da6a30f89abb52b516f1849571
relay/channel/baidu/relay-baidu.go|baiduEmbeddingHandler|Unmarshal|6c88b37cda6b69cb880852f57074d792626b82bd771acbab3a46754f80631684
relay/channel/baidu/relay-baidu.go|baiduEmbeddingHandler|Marshal|92981973316a59184c1f9d4fdca6fe3a1045b5da6a30f89abb52b516f1849571
-relay/channel/baidu/relay-baidu.go|getBaiduAccessTokenHelper|NewDecoder|64ee82d1a886a8c84d5771d6a7733571f892d8b735ff5b5ffa0167dc5f1ecb39
relay/channel/claude/relay-claude.go|ResponseClaude2OpenAI|Marshal|9f23cc0eecff96adc0410bf3360508a44815f9bc18567a76278a14aaa6e46e0a
relay/channel/claude/relay-claude.go|HandleClaudeResponseData|Marshal|0e01a2a781af1a773195f919e8bc2c53345aed7dd99776d98cfe550576ea831a
relay/channel/cloudflare/relay_cloudflare.go|cfStreamHandler|Unmarshal|5a9c4fba44410fd4d5ee918ce49cdd49a713d756c491be89ccc2d9f3de44cdc5
@@ -150,9 +148,6 @@ relay/channel/volcengine/adaptor.go|*Adaptor.ConvertAudioRequest|Unmarshal|21d62
relay/channel/volcengine/adaptor.go|*Adaptor.ConvertAudioRequest|Marshal|6df4bc98beb165214f3179900a7dc69f4c1aaa89816f2699cf59ee2bd83ef3d2
relay/channel/volcengine/tts.go|handleTTSResponse|Unmarshal|1c19ee75effdb9440195e572a78088232532eacc7938c6bf68a28030b8265b5f
relay/channel/volcengine/tts.go|handleTTSWebSocketResponse|Marshal|6df4bc98beb165214f3179900a7dc69f4c1aaa89816f2699cf59ee2bd83ef3d2
-relay/channel/xunfei/relay-xunfei.go|xunfeiStreamHandler|Marshal|168eb97ce25654c31be502868a95131c40bf62f506fa17dca8eec1cd5263b090
-relay/channel/xunfei/relay-xunfei.go|xunfeiHandler|Marshal|168eb97ce25654c31be502868a95131c40bf62f506fa17dca8eec1cd5263b090
-relay/channel/xunfei/relay-xunfei.go|xunfeiMakeRequest|Unmarshal|d2d86b3177f2d5720b8dc29e1ae601a20d37184045670a165bd9bbec01237187
relay/channel/zhipu/relay-zhipu.go|zhipuStreamHandler|Marshal|168eb97ce25654c31be502868a95131c40bf62f506fa17dca8eec1cd5263b090
relay/channel/zhipu/relay-zhipu.go|zhipuStreamHandler|Unmarshal|d6c2892f24d52e3e062e3eb924a5a5e290199f1238c21310625defc58a4d1cdf
relay/channel/zhipu/relay-zhipu.go|zhipuStreamHandler|Marshal|168eb97ce25654c31be502868a95131c40bf62f506fa17dca8eec1cd5263b090
diff --git a/web/default/src/components/json-editor.tsx b/web/default/src/components/json-editor.tsx
index 2129cb8e..af22547f 100644
--- a/web/default/src/components/json-editor.tsx
+++ b/web/default/src/components/json-editor.tsx
@@ -16,7 +16,7 @@ along with this program. If not, see .
For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API/issues
*/
-import { useState, useEffect } from 'react'
+import { useEffect, useRef, useState } from 'react'
import { Code, Table, Plus, Trash2 } from 'lucide-react'
import { useTranslation } from 'react-i18next'
import { cn } from '@/lib/utils'
@@ -43,6 +43,23 @@ type EditorRow = {
value: string
}
+let nextEditorRowId = 0
+
+function parseJsonRows(json: string): EditorRow[] | null {
+ try {
+ if (!json.trim()) return []
+ const parsed = JSON.parse(json)
+ if (typeof parsed !== 'object' || parsed === null) return null
+ return Object.entries(parsed).map(([key, val]) => ({
+ id: `json-editor-row-${nextEditorRowId++}`,
+ key,
+ value: typeof val === 'object' ? JSON.stringify(val) : String(val),
+ }))
+ } catch {
+ return null
+ }
+}
+
export function JsonEditor({
value,
onChange,
@@ -63,38 +80,35 @@ export function JsonEditor({
const resolvedKeyLabel = keyLabel ?? t('Key')
const resolvedValueLabel = valueLabel ?? t('Value')
const [mode, setMode] = useState<'visual' | 'json'>('visual')
- const [rows, setRows] = useState([])
+ const [rows, setRows] = useState(
+ () => parseJsonRows(value) ?? []
+ )
const [jsonValue, setJsonValue] = useState(value)
+ const handledValueRef = useRef(value)
- const parseJsonToRows = (json: string) => {
- try {
- if (!json.trim()) {
- setRows([])
- return
- }
- const parsed = JSON.parse(json)
- const newRows: EditorRow[] = Object.entries(parsed).map(
- ([key, val], index) => ({
- id: `${Date.now()}-${index}`,
- key,
- value: typeof val === 'object' ? JSON.stringify(val) : String(val),
- })
- )
- setRows(newRows)
- } catch (_error) {
- // Invalid JSON, keep current rows
- }
+ const syncRowsFromJson = (json: string) => {
+ const parsedRows = parseJsonRows(json)
+ if (parsedRows !== null) setRows(parsedRows)
+ }
+
+ const emitChange = (newValue: string) => {
+ handledValueRef.current = newValue
+ setJsonValue(newValue)
+ onChange(newValue)
}
// Parse JSON to rows when value changes externally
useEffect(() => {
- if (value !== jsonValue) {
+ if (value !== handledValueRef.current) {
+ handledValueRef.current = value
// External updates replace the local visual/JSON draft together.
- // eslint-disable-next-line react-hooks/set-state-in-effect
setJsonValue(value)
- parseJsonToRows(value)
+ const parsedRows = parseJsonRows(value)
+ if (parsedRows !== null) {
+ // eslint-disable-next-line react-hooks/set-state-in-effect
+ setRows(parsedRows)
+ }
}
- // eslint-disable-next-line react-hooks/exhaustive-deps
}, [value])
const convertRowsToJson = (updatedRows: EditorRow[]): string => {
@@ -139,8 +153,7 @@ export function JsonEditor({
const updatedRows = rows.filter((row) => row.id !== id)
setRows(updatedRows)
const json = convertRowsToJson(updatedRows)
- setJsonValue(json)
- onChange(json)
+ emitChange(json)
}
const handleRowChange = (
@@ -153,36 +166,24 @@ export function JsonEditor({
)
setRows(updatedRows)
const json = convertRowsToJson(updatedRows)
- setJsonValue(json)
- onChange(json)
+ emitChange(json)
}
const handleJsonChange = (newJson: string) => {
- setJsonValue(newJson)
- onChange(newJson)
- parseJsonToRows(newJson)
+ emitChange(newJson)
+ syncRowsFromJson(newJson)
}
const handleFillTemplate = () => {
if (!template) return
const templateJson = JSON.stringify(template, null, 2)
- setJsonValue(templateJson)
- onChange(templateJson)
- parseJsonToRows(templateJson)
+ emitChange(templateJson)
+ syncRowsFromJson(templateJson)
}
const toggleMode = () => {
- if (mode === 'visual') {
- // Switching to JSON mode: sync rows to JSON
- const json = convertRowsToJson(rows)
- setJsonValue(json)
- onChange(json)
- setMode('json')
- } else {
- // Switching to visual mode: sync JSON to rows
- parseJsonToRows(jsonValue)
- setMode('visual')
- }
+ if (mode === 'json') syncRowsFromJson(jsonValue)
+ setMode(mode === 'visual' ? 'json' : 'visual')
}
return (
diff --git a/web/default/src/features/auth/lib/app-session-verifier.ts b/web/default/src/features/auth/lib/app-session-verifier.ts
new file mode 100644
index 00000000..d95b1430
--- /dev/null
+++ b/web/default/src/features/auth/lib/app-session-verifier.ts
@@ -0,0 +1,29 @@
+/*
+Copyright (C) 2023-2026 MAX-API-Next
+
+This program is free software: you can redistribute it and/or modify
+it under the terms of the GNU Affero General Public License as
+published by the Free Software Foundation, either version 3 of the
+License, or (at your option) any later version.
+
+This program is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+GNU Affero General Public License for more details.
+
+You should have received a copy of the GNU Affero General Public License
+along with this program. If not, see .
+
+For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API/issues
+*/
+import { useAuthStore } from '@/stores/auth-store'
+import { getSelf } from '@/lib/api'
+import { createSessionVerifier } from './session-verifier'
+
+export const appSessionVerifier = createSessionVerifier({
+ getAuthState: () => useAuthStore.getState().auth,
+ getSelf,
+ getUserIdentity: (user) => user.id,
+ isUnauthorized: (error: unknown) =>
+ (error as { response?: { status?: number } })?.response?.status === 401,
+})
diff --git a/web/default/src/features/auth/lib/session-verifier.ts b/web/default/src/features/auth/lib/session-verifier.ts
new file mode 100644
index 00000000..fe9088cf
--- /dev/null
+++ b/web/default/src/features/auth/lib/session-verifier.ts
@@ -0,0 +1,139 @@
+/*
+Copyright (C) 2023-2026 MAX-API-Next
+
+This program is free software: you can redistribute it and/or modify
+it under the terms of the GNU Affero General Public License as
+published by the Free Software Foundation, either version 3 of the
+License, or (at your option) any later version.
+
+This program is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+GNU Affero General Public License for more details.
+
+You should have received a copy of the GNU Affero General Public License
+along with this program. If not, see .
+
+For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API/issues
+*/
+type SessionAuthState = {
+ user: TUser | null
+ setUser: (user: TUser | null) => void
+ reset: () => void
+}
+
+type SessionResponse = {
+ success?: boolean
+ data?: TUser
+}
+
+type SessionVerifierDependencies = {
+ getAuthState: () => SessionAuthState
+ getSelf: () => Promise>
+ getUserIdentity: (user: TUser) => string | number
+ isUnauthorized: (error: unknown) => boolean
+ now?: () => number
+ revalidateAfterMs?: number
+}
+
+export const AUTH_SESSION_REVALIDATE_MS = 30_000
+
+export function createSessionVerifier(
+ dependencies: SessionVerifierDependencies
+) {
+ const now = dependencies.now ?? Date.now
+ const revalidateAfterMs =
+ dependencies.revalidateAfterMs ?? AUTH_SESSION_REVALIDATE_MS
+ let verifiedAt = 0
+ let verifiedUserIdentity: string | number | null = null
+ let inFlight: {
+ userIdentity: string | number | null
+ promise: Promise
+ } | null = null
+
+ const getCurrentUserIdentity = () => {
+ const user = dependencies.getAuthState().user
+ return user ? dependencies.getUserIdentity(user) : null
+ }
+
+ const verifyRemote = async (
+ startingUserIdentity: string | number | null
+ ): Promise => {
+ let response: SessionResponse
+ try {
+ response = await dependencies.getSelf()
+ } catch (error: unknown) {
+ if (!dependencies.isUnauthorized(error)) throw error
+ if (getCurrentUserIdentity() === startingUserIdentity) {
+ dependencies.getAuthState().reset()
+ verifiedAt = 0
+ verifiedUserIdentity = null
+ }
+ return false
+ }
+
+ if (getCurrentUserIdentity() !== startingUserIdentity) {
+ return false
+ }
+
+ const currentAuth = dependencies.getAuthState()
+ if (response.success && response.data) {
+ currentAuth.setUser(response.data)
+ verifiedUserIdentity = dependencies.getUserIdentity(response.data)
+ verifiedAt = now()
+ return true
+ }
+
+ currentAuth.reset()
+ verifiedAt = 0
+ verifiedUserIdentity = null
+ return false
+ }
+
+ const verify = async (): Promise => {
+ while (true) {
+ const currentUserIdentity = getCurrentUserIdentity()
+ if (
+ currentUserIdentity !== null &&
+ currentUserIdentity === verifiedUserIdentity &&
+ verifiedAt > 0 &&
+ now() - verifiedAt < revalidateAfterMs
+ ) {
+ return true
+ }
+
+ const active = inFlight
+ if (active) {
+ if (active.userIdentity === currentUserIdentity) {
+ try {
+ const result = await active.promise
+ if (getCurrentUserIdentity() === currentUserIdentity) return result
+ } catch (error: unknown) {
+ if (getCurrentUserIdentity() === currentUserIdentity) throw error
+ }
+ } else {
+ try {
+ await active.promise
+ } catch {
+ // A request for another identity must not decide this verification.
+ }
+ }
+ continue
+ }
+
+ const promise = verifyRemote(currentUserIdentity)
+ const entry = { userIdentity: currentUserIdentity, promise }
+ inFlight = entry
+ try {
+ const result = await promise
+ if (getCurrentUserIdentity() === currentUserIdentity) return result
+ } catch (error: unknown) {
+ if (getCurrentUserIdentity() === currentUserIdentity) throw error
+ } finally {
+ if (inFlight === entry) inFlight = null
+ }
+ }
+ }
+
+ return { verify }
+}
diff --git a/web/default/src/features/auth/secure-verification/types.ts b/web/default/src/features/auth/secure-verification/types.ts
index 60c3cae6..c3b7e99a 100644
--- a/web/default/src/features/auth/secure-verification/types.ts
+++ b/web/default/src/features/auth/secure-verification/types.ts
@@ -17,7 +17,7 @@ along with this program. If not, see .
For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API/issues
*/
export type VerificationMethod = '2fa' | 'passkey' | 'password'
-export type VerificationScope = 'access_token'
+export type VerificationScope = 'access_token' | 'account_delete'
export interface VerificationMethods {
has2FA: boolean
diff --git a/web/default/src/features/channels/components/drawers/channel-editor-state.ts b/web/default/src/features/channels/components/drawers/channel-editor-state.ts
index 448ba5c0..a464beb5 100644
--- a/web/default/src/features/channels/components/drawers/channel-editor-state.ts
+++ b/web/default/src/features/channels/components/drawers/channel-editor-state.ts
@@ -166,6 +166,8 @@ export const ADVANCED_ERROR_FIELDS = [
'video_task_task_id_path',
'video_task_status_path',
'video_task_progress_path',
+ 'video_task_created_at_path',
+ 'video_task_updated_at_path',
'video_task_result_url_paths',
'video_task_error_message_path',
'video_task_status_submitted',
diff --git a/web/default/src/features/channels/components/drawers/channel-mutate-drawer.tsx b/web/default/src/features/channels/components/drawers/channel-mutate-drawer.tsx
index 055273db..3d093319 100644
--- a/web/default/src/features/channels/components/drawers/channel-mutate-drawer.tsx
+++ b/web/default/src/features/channels/components/drawers/channel-mutate-drawer.tsx
@@ -3148,6 +3148,54 @@ export function ChannelMutateDrawer({
)}
/>
+
+ (
+
+
+ {t('Created Time Path')}
+
+
+
+
+
+ {t(
+ 'Response path for task creation timestamp'
+ )}
+
+
+
+ )}
+ />
+
+ (
+
+
+ {t('Updated Time Path')}
+
+
+
+
+
+ {t(
+ 'Response path for task update timestamp'
+ )}
+
+
+
+ )}
+ />
{
- test('caps selected chips while preserving the complete model count', () => {
+ test('shows every selected model chip and preserves the complete count', () => {
const markup = renderToStaticMarkup(
@@ -76,9 +76,10 @@ describe('ChannelModelsSection', () => {
)
assert.match(markup, /Selected 8/)
- assert.match(markup, /\+2 more/)
- assert.doesNotMatch(markup, />long-vendor-model-7long-vendor-model-8${model}<`))
+ }
})
test('keeps long mapping warnings wrap-safe and form descriptions linked', () => {
diff --git a/web/default/src/features/channels/components/drawers/sections/channel-models-section.tsx b/web/default/src/features/channels/components/drawers/sections/channel-models-section.tsx
index 88a0085f..c6aa427f 100644
--- a/web/default/src/features/channels/components/drawers/sections/channel-models-section.tsx
+++ b/web/default/src/features/channels/components/drawers/sections/channel-models-section.tsx
@@ -188,7 +188,6 @@ export function ChannelModelsSection(props: ChannelModelsSectionProps) {
placeholder={t('Select models or add custom ones')}
allowCreate
createLabel='Add custom model "{{value}}"'
- maxVisibleChips={6}
/>
{modelMappingGuardrail.exposedTargetModels.length > 0 && (
diff --git a/web/default/src/features/channels/lib/channel-form.test.ts b/web/default/src/features/channels/lib/channel-form.test.ts
new file mode 100644
index 00000000..58fd638d
--- /dev/null
+++ b/web/default/src/features/channels/lib/channel-form.test.ts
@@ -0,0 +1,93 @@
+/*
+Copyright (C) 2023-2026 MAX-API-Next
+
+This program is free software: you can redistribute it and/or modify
+it under the terms of the GNU Affero General Public License as
+published by the Free Software Foundation, either version 3 of the
+License, or (at your option) any later version.
+
+This program is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+GNU Affero General Public License for more details.
+
+You should have received a copy of the GNU Affero General Public License
+along with this program. If not, see .
+
+For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API/issues
+*/
+import assert from 'node:assert/strict'
+import { describe, test } from 'node:test'
+import { CHANNEL_TYPE_DOUBAO_VIDEO } from '../constants'
+import { channelSchema } from '../types'
+import {
+ CHANNEL_FORM_DEFAULT_VALUES,
+ channelFormSchema,
+ transformChannelToFormDefaults,
+ transformFormDataToUpdatePayload,
+} from './channel-form'
+
+describe('channel form settings round trip', () => {
+ test('preserves status mapping and all configured video response paths', () => {
+ const channel = channelSchema.parse({
+ id: 7,
+ type: CHANNEL_TYPE_DOUBAO_VIDEO,
+ key: 'not-returned-to-form',
+ status: 1,
+ name: 'video channel',
+ created_time: 1,
+ test_time: 1,
+ response_time: 1,
+ balance_updated_time: 1,
+ models: 'video-model',
+ status_code_mapping: '{"429":"503"}',
+ settings: JSON.stringify({
+ task_protocol: 'generic_video_task',
+ task_protocol_config: {
+ submit_path: '/submit',
+ query_path: '/tasks/{task_id}',
+ task_id_path: 'data.id',
+ status_path: 'data.status',
+ progress_path: 'data.progress',
+ result_url_paths: ['data.url'],
+ error_message_path: 'data.error.message',
+ created_at_path: 'data.created_at',
+ updated_at_path: 'data.updated_at',
+ status_map: { completed: 'SUCCESS' },
+ vendor_extension: { version: 2 },
+ },
+ }),
+ })
+
+ const defaults = transformChannelToFormDefaults(channel)
+ assert.equal(defaults.status_code_mapping, '{"429":"503"}')
+ assert.equal(defaults.video_task_created_at_path, 'data.created_at')
+ assert.equal(defaults.video_task_updated_at_path, 'data.updated_at')
+
+ const payload = transformFormDataToUpdatePayload(defaults, channel.id)
+ const settings = JSON.parse(String(payload.settings))
+ assert.equal(payload.status_code_mapping, '{"429":"503"}')
+ assert.equal(
+ settings.task_protocol_config.created_at_path,
+ 'data.created_at'
+ )
+ assert.equal(
+ settings.task_protocol_config.updated_at_path,
+ 'data.updated_at'
+ )
+ assert.deepEqual(settings.task_protocol_config.vendor_extension, {
+ version: 2,
+ })
+ })
+
+ test('rejects aliased numeric status-code source keys', () => {
+ const result = channelFormSchema.safeParse({
+ ...CHANNEL_FORM_DEFAULT_VALUES,
+ name: 'channel',
+ models: 'model',
+ status_code_mapping: '{"429":503," 429":502}',
+ })
+
+ assert.equal(result.success, false)
+ })
+})
diff --git a/web/default/src/features/channels/lib/channel-form.ts b/web/default/src/features/channels/lib/channel-form.ts
index 71a9030c..2d91c7e9 100644
--- a/web/default/src/features/channels/lib/channel-form.ts
+++ b/web/default/src/features/channels/lib/channel-form.ts
@@ -48,6 +48,8 @@ const DEFAULT_VIDEO_TASK_QUERY_PATH = '/v1/videos/{task_id}'
const DEFAULT_VIDEO_TASK_TASK_ID_PATH = 'task_id'
const DEFAULT_VIDEO_TASK_STATUS_PATH = 'status'
const DEFAULT_VIDEO_TASK_PROGRESS_PATH = 'progress'
+const DEFAULT_VIDEO_TASK_CREATED_AT_PATH = 'created_at'
+const DEFAULT_VIDEO_TASK_UPDATED_AT_PATH = 'updated_at'
const DEFAULT_VIDEO_TASK_RESULT_URL_PATHS =
'result.primary_url,result.urls.0,result.url,result.video_url,result.output_url,data.result.primary_url,data.result.urls.0,data.result.url,data.result.video_url,data.result.output_url,url,video_url,output_url,file_url,download_url,result'
const DEFAULT_VIDEO_TASK_ERROR_MESSAGE_PATH = 'error_message'
@@ -96,17 +98,22 @@ function isOptionalStatusCodeMapping(value: string | undefined): boolean {
const parsed = parseOptionalJson(value)
if (parsed === undefined) return true
if (!isJsonObjectValue(parsed)) return false
+ const seenSourceCodes = new Set()
return Object.entries(parsed).every(([from, to]) => {
+ const normalizedFrom = from.trim()
const fromCode = Number(from)
const toCode = Number(to)
- return (
+ const valid =
+ /^[1-5]\d{2}$/.test(normalizedFrom) &&
Number.isInteger(fromCode) &&
Number.isInteger(toCode) &&
fromCode >= 100 &&
fromCode <= 599 &&
toCode >= 100 &&
- toCode <= 599
- )
+ toCode <= 599 &&
+ !seenSourceCodes.has(fromCode)
+ seenSourceCodes.add(fromCode)
+ return valid
})
} catch {
return false
@@ -253,6 +260,8 @@ export const channelFormSchema = z
video_task_task_id_path: z.string().optional(),
video_task_status_path: z.string().optional(),
video_task_progress_path: z.string().optional(),
+ video_task_created_at_path: z.string().optional(),
+ video_task_updated_at_path: z.string().optional(),
video_task_result_url_paths: z.string().optional(),
video_task_error_message_path: z.string().optional(),
video_task_status_submitted: z.string().optional(),
@@ -444,6 +453,8 @@ export const CHANNEL_FORM_DEFAULT_VALUES: ChannelFormValues = {
video_task_task_id_path: DEFAULT_VIDEO_TASK_TASK_ID_PATH,
video_task_status_path: DEFAULT_VIDEO_TASK_STATUS_PATH,
video_task_progress_path: DEFAULT_VIDEO_TASK_PROGRESS_PATH,
+ video_task_created_at_path: DEFAULT_VIDEO_TASK_CREATED_AT_PATH,
+ video_task_updated_at_path: DEFAULT_VIDEO_TASK_UPDATED_AT_PATH,
video_task_result_url_paths: DEFAULT_VIDEO_TASK_RESULT_URL_PATHS,
video_task_error_message_path: DEFAULT_VIDEO_TASK_ERROR_MESSAGE_PATH,
video_task_status_submitted: DEFAULT_VIDEO_TASK_STATUS_SUBMITTED,
@@ -595,6 +606,8 @@ export function transformChannelToFormDefaults(
let videoTaskIDPath = DEFAULT_VIDEO_TASK_TASK_ID_PATH
let videoTaskStatusPath = DEFAULT_VIDEO_TASK_STATUS_PATH
let videoTaskProgressPath = DEFAULT_VIDEO_TASK_PROGRESS_PATH
+ let videoTaskCreatedAtPath = DEFAULT_VIDEO_TASK_CREATED_AT_PATH
+ let videoTaskUpdatedAtPath = DEFAULT_VIDEO_TASK_UPDATED_AT_PATH
let videoTaskResultURLPaths = DEFAULT_VIDEO_TASK_RESULT_URL_PATHS
let videoTaskErrorMessagePath = DEFAULT_VIDEO_TASK_ERROR_MESSAGE_PATH
let videoTaskStatusSubmitted = DEFAULT_VIDEO_TASK_STATUS_SUBMITTED
@@ -663,6 +676,12 @@ export function transformChannelToFormDefaults(
videoTaskProgressPath =
String(taskProtocolConfig.progress_path || '').trim() ||
DEFAULT_VIDEO_TASK_PROGRESS_PATH
+ videoTaskCreatedAtPath =
+ String(taskProtocolConfig.created_at_path || '').trim() ||
+ DEFAULT_VIDEO_TASK_CREATED_AT_PATH
+ videoTaskUpdatedAtPath =
+ String(taskProtocolConfig.updated_at_path || '').trim() ||
+ DEFAULT_VIDEO_TASK_UPDATED_AT_PATH
videoTaskResultURLPaths = joinCSV(
taskProtocolConfig.result_url_paths,
DEFAULT_VIDEO_TASK_RESULT_URL_PATHS
@@ -754,6 +773,8 @@ export function transformChannelToFormDefaults(
video_task_task_id_path: videoTaskIDPath,
video_task_status_path: videoTaskStatusPath,
video_task_progress_path: videoTaskProgressPath,
+ video_task_created_at_path: videoTaskCreatedAtPath,
+ video_task_updated_at_path: videoTaskUpdatedAtPath,
video_task_result_url_paths: videoTaskResultURLPaths,
video_task_error_message_path: videoTaskErrorMessagePath,
video_task_status_submitted: videoTaskStatusSubmitted,
@@ -904,6 +925,11 @@ function buildSettingsJSON(formData: ChannelFormValues): string {
enableVideoTaskProtocol)
if (enableVideoTaskProtocol) {
+ const existingTaskProtocolConfig = isJsonObjectValue(
+ settingsObj.task_protocol_config
+ )
+ ? settingsObj.task_protocol_config
+ : {}
const statusMap: Record = {}
setVideoTaskStatusMap(
statusMap,
@@ -933,6 +959,7 @@ function buildSettingsJSON(formData: ChannelFormValues): string {
settingsObj.task_protocol = VIDEO_TASK_PROTOCOL
settingsObj.task_protocol_config = {
+ ...existingTaskProtocolConfig,
submit_path:
formData.video_task_submit_path?.trim() ||
DEFAULT_VIDEO_TASK_SUBMIT_PATH,
@@ -947,6 +974,12 @@ function buildSettingsJSON(formData: ChannelFormValues): string {
progress_path:
formData.video_task_progress_path?.trim() ||
DEFAULT_VIDEO_TASK_PROGRESS_PATH,
+ created_at_path:
+ formData.video_task_created_at_path?.trim() ||
+ DEFAULT_VIDEO_TASK_CREATED_AT_PATH,
+ updated_at_path:
+ formData.video_task_updated_at_path?.trim() ||
+ DEFAULT_VIDEO_TASK_UPDATED_AT_PATH,
result_url_paths: splitCSV(
formData.video_task_result_url_paths,
DEFAULT_VIDEO_TASK_RESULT_URL_PATHS
diff --git a/web/default/src/features/models/components/drawers/model-mutate-drawer.test.ts b/web/default/src/features/models/components/drawers/model-mutate-drawer.test.ts
new file mode 100644
index 00000000..a806750a
--- /dev/null
+++ b/web/default/src/features/models/components/drawers/model-mutate-drawer.test.ts
@@ -0,0 +1,90 @@
+/*
+Copyright (C) 2023-2026 MAX-API-Next
+
+This program is free software: you can redistribute it and/or modify
+it under the terms of the GNU Affero General Public License as
+published by the Free Software Foundation, either version 3 of the
+License, or (at your option) any later version.
+
+This program is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+GNU Affero General Public License for more details.
+
+You should have received a copy of the GNU Affero General Public License
+along with this program. If not, see .
+
+For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API/issues
+*/
+import assert from 'node:assert/strict'
+import { describe, test } from 'node:test'
+import type { ModelSettings } from '@/features/system-settings/types'
+import {
+ parseNonNegativeFiniteNumber,
+ readPricingConfig,
+ shouldReplacePricingConfig,
+} from './model-pricing-config'
+
+function createPricingSettings(
+ overrides: Partial
+): ModelSettings {
+ return {
+ ModelPrice: '{}',
+ ModelRatio: '{}',
+ CacheRatio: '{}',
+ CompletionRatio: '{}',
+ ImageRatio: '{}',
+ AudioRatio: '{}',
+ AudioCompletionRatio: '{}',
+ ...overrides,
+ } as ModelSettings
+}
+
+describe('model drawer pricing', () => {
+ test('strictly parses only finite non-negative pricing values', () => {
+ assert.equal(parseNonNegativeFiniteNumber('0'), 0)
+ assert.equal(parseNonNegativeFiniteNumber('1.25e-3'), 0.00125)
+ assert.equal(parseNonNegativeFiniteNumber(''), undefined)
+ assert.equal(parseNonNegativeFiniteNumber('-1'), undefined)
+ assert.equal(parseNonNegativeFiniteNumber('Infinity'), undefined)
+ assert.equal(parseNonNegativeFiniteNumber('1abc'), undefined)
+ })
+
+ test('loads existing pricing for a same-name create flow', () => {
+ const pricing = readPricingConfig(
+ createPricingSettings({
+ ModelRatio: '{"existing-model":1.5}',
+ CompletionRatio: '{"existing-model":2}',
+ CacheRatio: '{"existing-model":0}',
+ }),
+ 'existing-model'
+ )
+
+ assert.equal(pricing.mode, 'per-token')
+ assert.equal(pricing.fields.ratio, '1.5')
+ assert.equal(pricing.promptPrice, '3')
+ assert.equal(pricing.completionPrice, '6')
+ assert.equal(pricing.fields.cacheRatio, '0')
+ assert.equal(pricing.advancedOpen, true)
+ })
+
+ test('keeps fixed price authoritative over stale ratio entries', () => {
+ const pricing = readPricingConfig(
+ createPricingSettings({
+ ModelPrice: '{"priced-model":0.25}',
+ ModelRatio: '{"priced-model":9}',
+ }),
+ 'priced-model'
+ )
+
+ assert.equal(pricing.mode, 'per-request')
+ assert.equal(pricing.fields.price, '0.25')
+ assert.equal(pricing.fields.ratio, '')
+ })
+
+ test('only replaces pricing the form loaded or explicitly supplied', () => {
+ assert.equal(shouldReplacePricingConfig('loaded', 'loaded', false), true)
+ assert.equal(shouldReplacePricingConfig('renamed', 'loaded', true), true)
+ assert.equal(shouldReplacePricingConfig('foreign', 'loaded', false), false)
+ })
+})
diff --git a/web/default/src/features/models/components/drawers/model-mutate-drawer.tsx b/web/default/src/features/models/components/drawers/model-mutate-drawer.tsx
index c95e7f09..fbd3aa0c 100644
--- a/web/default/src/features/models/components/drawers/model-mutate-drawer.tsx
+++ b/web/default/src/features/models/components/drawers/model-mutate-drawer.tsx
@@ -16,7 +16,7 @@ along with this program. If not, see .
For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API/issues
*/
-import { useEffect, useState, useCallback, useMemo } from 'react'
+import { useEffect, useState, useCallback, useMemo, useRef } from 'react'
import * as z from 'zod'
import { useForm } from 'react-hook-form'
import { zodResolver } from '@hookform/resolvers/zod'
@@ -78,11 +78,22 @@ import {
import { useUpdateOption } from '@/features/system-settings/hooks/use-update-option'
import { normalizeJsonString } from '@/features/system-settings/models/utils'
import type { ModelSettings } from '@/features/system-settings/types'
-import { safeJsonParse } from '@/features/system-settings/utils/json-parser'
import { createModel, updateModel, getModel, getVendors } from '../../api'
import { getNameRuleOptions, ENDPOINT_TEMPLATES } from '../../constants'
import { modelsQueryKeys, vendorsQueryKeys, parseModelTags } from '../../lib'
import type { Model } from '../../types'
+import {
+ parseNonNegativeFiniteNumber,
+ parsePricingMap,
+ readPricingConfig,
+ shouldReplacePricingConfig,
+ type PricingMode,
+} from './model-pricing-config'
+
+const isOptionalPricingNumber = (value: string | undefined) =>
+ value === undefined ||
+ value.trim() === '' ||
+ parseNonNegativeFiniteNumber(value) !== undefined
// Extended schema for ratio configuration (internal form state only)
const extendedModelFormSchema = z.object({
@@ -96,18 +107,17 @@ const extendedModelFormSchema = z.object({
name_rule: z.number(),
status: z.boolean(),
sync_official: z.boolean(),
- price: z.string().optional(),
- ratio: z.string().optional(),
- cacheRatio: z.string().optional(),
- completionRatio: z.string().optional(),
- imageRatio: z.string().optional(),
- audioRatio: z.string().optional(),
- audioCompletionRatio: z.string().optional(),
+ price: z.string().optional().refine(isOptionalPricingNumber),
+ ratio: z.string().optional().refine(isOptionalPricingNumber),
+ cacheRatio: z.string().optional().refine(isOptionalPricingNumber),
+ completionRatio: z.string().optional().refine(isOptionalPricingNumber),
+ imageRatio: z.string().optional().refine(isOptionalPricingNumber),
+ audioRatio: z.string().optional().refine(isOptionalPricingNumber),
+ audioCompletionRatio: z.string().optional().refine(isOptionalPricingNumber),
})
type ExtendedModelFormValues = z.infer
-type PricingMode = 'per-token' | 'per-request'
type PricingSubMode = 'ratio' | 'price'
type ModelMutateDrawerProps = {
@@ -131,6 +141,8 @@ export function ModelMutateDrawer({
const [promptPrice, setPromptPrice] = useState('')
const [completionPrice, setCompletionPrice] = useState('')
const [oldModelName, setOldModelName] = useState('')
+ const [loadedPricingName, setLoadedPricingName] = useState('')
+ const modelSettingsRef = useRef(null)
// Fetch vendors for dropdown
const { data: vendorsData } = useQuery({
@@ -208,6 +220,11 @@ export function ModelMutateDrawer({
return getOptionValue(systemOptionsData.data, defaultModelSettings)
}, [systemOptionsData])
+ const hasModelSettings = modelSettings !== null
+ useEffect(() => {
+ modelSettingsRef.current = modelSettings
+ })
+
const form = useForm({
resolver: zodResolver(extendedModelFormSchema),
defaultValues: {
@@ -232,13 +249,14 @@ export function ModelMutateDrawer({
const validateNumber = (value: string) => {
if (value === '') return true
- return !isNaN(parseFloat(value))
+ return parseNonNegativeFiniteNumber(value) !== undefined
}
const handlePromptPriceChange = (value: string) => {
setPromptPrice(value)
- if (value && !isNaN(parseFloat(value))) {
- const ratio = parseFloat(value) / 2
+ const parsedValue = parseNonNegativeFiniteNumber(value)
+ if (parsedValue !== undefined) {
+ const ratio = parsedValue / 2
form.setValue('ratio', ratio.toString())
} else {
form.setValue('ratio', '')
@@ -247,14 +265,14 @@ export function ModelMutateDrawer({
const handleCompletionPriceChange = (value: string) => {
setCompletionPrice(value)
+ const parsedValue = parseNonNegativeFiniteNumber(value)
+ const parsedPromptPrice = parseNonNegativeFiniteNumber(promptPrice)
if (
- value &&
- !isNaN(parseFloat(value)) &&
- promptPrice &&
- !isNaN(parseFloat(promptPrice)) &&
- parseFloat(promptPrice) > 0
+ parsedValue !== undefined &&
+ parsedPromptPrice !== undefined &&
+ parsedPromptPrice > 0
) {
- const completionRatio = parseFloat(value) / parseFloat(promptPrice)
+ const completionRatio = parsedValue / parsedPromptPrice
form.setValue('completionRatio', completionRatio.toString())
} else {
form.setValue('completionRatio', '')
@@ -262,13 +280,22 @@ export function ModelMutateDrawer({
}
// Load model data for editing and ratio configuration
+ /* eslint-disable react-hooks/set-state-in-effect -- async query snapshots initialize the drawer form */
useEffect(() => {
if (open && isEditing && modelData?.data) {
const model = modelData.data
setOldModelName(model.model_name)
- // Base model data reset
- const baseModelData = {
+ const pricing = readPricingConfig(
+ modelSettingsRef.current,
+ model.model_name
+ )
+ setLoadedPricingName(model.model_name)
+ setPricingMode(pricing.mode)
+ setPromptPrice(pricing.promptPrice)
+ setCompletionPrice(pricing.completionPrice)
+ setAdvancedOpen(pricing.advancedOpen)
+ form.reset({
id: model.id,
model_name: model.model_name,
description: model.description || '',
@@ -279,102 +306,20 @@ export function ModelMutateDrawer({
name_rule: model.name_rule || 0,
status: model.status === 1,
sync_official: model.sync_official === 1,
- price: '',
- ratio: '',
- cacheRatio: '',
- completionRatio: '',
- imageRatio: '',
- audioRatio: '',
- audioCompletionRatio: '',
- }
-
- // Parse ratio configurations from system settings if available
- if (modelSettings) {
- const priceMap = safeJsonParse>(
- modelSettings.ModelPrice,
- { fallback: {}, silent: true }
- )
- const ratioMap = safeJsonParse>(
- modelSettings.ModelRatio,
- { fallback: {}, silent: true }
- )
- const cacheMap = safeJsonParse>(
- modelSettings.CacheRatio,
- { fallback: {}, silent: true }
- )
- const completionMap = safeJsonParse>(
- modelSettings.CompletionRatio,
- { fallback: {}, silent: true }
- )
- const imageMap = safeJsonParse>(
- modelSettings.ImageRatio,
- { fallback: {}, silent: true }
- )
- const audioMap = safeJsonParse>(
- modelSettings.AudioRatio,
- { fallback: {}, silent: true }
- )
- const audioCompletionMap = safeJsonParse>(
- modelSettings.AudioCompletionRatio,
- { fallback: {}, silent: true }
- )
-
- // Extract ratio config for this model
- const modelName = model.model_name
- const price = priceMap[modelName]
- const ratio = ratioMap[modelName]
- const cacheRatio = cacheMap[modelName]
- const completionRatio = completionMap[modelName]
- const imageRatio = imageMap[modelName]
- const audioRatio = audioMap[modelName]
- const audioCompletionRatio = audioCompletionMap[modelName]
-
- // Determine pricing mode
- if (price !== undefined && price !== null) {
- setPricingMode('per-request')
- form.reset({
- ...baseModelData,
- price: price.toString(),
- })
- } else {
- setPricingMode('per-token')
- if (ratio !== undefined && ratio !== null) {
- const tokenPrice = ratio * 2
- setPromptPrice(tokenPrice.toString())
- if (completionRatio !== undefined && completionRatio !== null) {
- const compPrice = tokenPrice * completionRatio
- setCompletionPrice(compPrice.toString())
- }
- }
- form.reset({
- ...baseModelData,
- ratio: ratio?.toString() || '',
- cacheRatio: cacheRatio?.toString() || '',
- completionRatio: completionRatio?.toString() || '',
- imageRatio: imageRatio?.toString() || '',
- audioRatio: audioRatio?.toString() || '',
- audioCompletionRatio: audioCompletionRatio?.toString() || '',
- })
- setAdvancedOpen(
- !!(cacheRatio || imageRatio || audioRatio || audioCompletionRatio)
- )
- }
- } else {
- // If system settings not loaded yet, just load base model data
- setPricingMode('per-token')
- form.reset(baseModelData)
- setAdvancedOpen(false)
- }
+ ...pricing.fields,
+ })
} else if (open && !isEditing) {
- // Pre-fill model name if passed from missing models
+ const modelName = currentRow?.model_name || ''
+ const pricing = readPricingConfig(modelSettingsRef.current, modelName)
setOldModelName('')
- setPricingMode('per-token')
+ setLoadedPricingName(modelName)
setPricingSubMode('ratio')
- setPromptPrice('')
- setCompletionPrice('')
- setAdvancedOpen(false)
+ setPricingMode(pricing.mode)
+ setPromptPrice(pricing.promptPrice)
+ setCompletionPrice(pricing.completionPrice)
+ setAdvancedOpen(pricing.advancedOpen)
form.reset({
- model_name: currentRow?.model_name || '',
+ model_name: modelName,
description: '',
icon: '',
tags: [],
@@ -383,16 +328,11 @@ export function ModelMutateDrawer({
name_rule: 0,
status: true,
sync_official: true,
- price: '',
- ratio: '',
- cacheRatio: '',
- completionRatio: '',
- imageRatio: '',
- audioRatio: '',
- audioCompletionRatio: '',
+ ...pricing.fields,
})
}
- }, [open, isEditing, modelData, currentRow, form, modelSettings])
+ }, [open, isEditing, modelData, currentRow, form, hasModelSettings])
+ /* eslint-enable react-hooks/set-state-in-effect */
const onSubmit = useCallback(
async (values: ExtendedModelFormValues): Promise => {
@@ -425,7 +365,7 @@ export function ModelMutateDrawer({
if (response.success) {
// Handle ratio configuration updates in system settings
const finalModelName = values.model_name
- const hasRatioConfig =
+ const hasRatioConfig = Boolean(
(pricingMode === 'per-request' &&
values.price &&
values.price !== '') ||
@@ -436,38 +376,20 @@ export function ModelMutateDrawer({
values.imageRatio ||
values.audioRatio ||
values.audioCompletionRatio))
+ )
// Always process system settings updates if we have modelSettings
// This ensures we can remove stale entries even when clearing all pricing fields
if (modelSettings) {
// Read existing configurations
- const priceMap = safeJsonParse>(
- modelSettings.ModelPrice,
- { fallback: {}, silent: true }
- )
- const ratioMap = safeJsonParse>(
- modelSettings.ModelRatio,
- { fallback: {}, silent: true }
- )
- const cacheMap = safeJsonParse>(
- modelSettings.CacheRatio,
- { fallback: {}, silent: true }
- )
- const completionMap = safeJsonParse>(
- modelSettings.CompletionRatio,
- { fallback: {}, silent: true }
- )
- const imageMap = safeJsonParse>(
- modelSettings.ImageRatio,
- { fallback: {}, silent: true }
- )
- const audioMap = safeJsonParse>(
- modelSettings.AudioRatio,
- { fallback: {}, silent: true }
- )
- const audioCompletionMap = safeJsonParse>(
- modelSettings.AudioCompletionRatio,
- { fallback: {}, silent: true }
+ const priceMap = parsePricingMap(modelSettings.ModelPrice)
+ const ratioMap = parsePricingMap(modelSettings.ModelRatio)
+ const cacheMap = parsePricingMap(modelSettings.CacheRatio)
+ const completionMap = parsePricingMap(modelSettings.CompletionRatio)
+ const imageMap = parsePricingMap(modelSettings.ImageRatio)
+ const audioMap = parsePricingMap(modelSettings.AudioRatio)
+ const audioCompletionMap = parsePricingMap(
+ modelSettings.AudioCompletionRatio
)
// Remove old model name entries if model name changed (always, even if no new config)
@@ -481,15 +403,21 @@ export function ModelMutateDrawer({
delete audioCompletionMap[oldModelName]
}
- // Remove current model name from all maps first (always, to handle mode switches or clearing)
- // This ensures stale entries are removed even when user clears all fields
- delete priceMap[finalModelName]
- delete ratioMap[finalModelName]
- delete cacheMap[finalModelName]
- delete completionMap[finalModelName]
- delete imageMap[finalModelName]
- delete audioMap[finalModelName]
- delete audioCompletionMap[finalModelName]
+ if (
+ shouldReplacePricingConfig(
+ finalModelName,
+ loadedPricingName,
+ hasRatioConfig
+ )
+ ) {
+ delete priceMap[finalModelName]
+ delete ratioMap[finalModelName]
+ delete cacheMap[finalModelName]
+ delete completionMap[finalModelName]
+ delete imageMap[finalModelName]
+ delete audioMap[finalModelName]
+ delete audioCompletionMap[finalModelName]
+ }
// Only add new entries if user provided new configuration
if (hasRatioConfig) {
@@ -498,32 +426,41 @@ export function ModelMutateDrawer({
values.price &&
values.price !== ''
) {
- priceMap[finalModelName] = parseFloat(values.price)
+ priceMap[finalModelName] = parseNonNegativeFiniteNumber(
+ values.price
+ )!
} else if (pricingMode === 'per-token') {
if (values.ratio && values.ratio !== '') {
- ratioMap[finalModelName] = parseFloat(values.ratio)
+ ratioMap[finalModelName] = parseNonNegativeFiniteNumber(
+ values.ratio
+ )!
}
if (values.cacheRatio && values.cacheRatio !== '') {
- cacheMap[finalModelName] = parseFloat(values.cacheRatio)
+ cacheMap[finalModelName] = parseNonNegativeFiniteNumber(
+ values.cacheRatio
+ )!
}
if (values.completionRatio && values.completionRatio !== '') {
- completionMap[finalModelName] = parseFloat(
+ completionMap[finalModelName] = parseNonNegativeFiniteNumber(
values.completionRatio
- )
+ )!
}
if (values.imageRatio && values.imageRatio !== '') {
- imageMap[finalModelName] = parseFloat(values.imageRatio)
+ imageMap[finalModelName] = parseNonNegativeFiniteNumber(
+ values.imageRatio
+ )!
}
if (values.audioRatio && values.audioRatio !== '') {
- audioMap[finalModelName] = parseFloat(values.audioRatio)
+ audioMap[finalModelName] = parseNonNegativeFiniteNumber(
+ values.audioRatio
+ )!
}
if (
values.audioCompletionRatio &&
values.audioCompletionRatio !== ''
) {
- audioCompletionMap[finalModelName] = parseFloat(
- values.audioCompletionRatio
- )
+ audioCompletionMap[finalModelName] =
+ parseNonNegativeFiniteNumber(values.audioCompletionRatio)!
}
}
}
@@ -622,6 +559,7 @@ export function ModelMutateDrawer({
onOpenChange,
pricingMode,
oldModelName,
+ loadedPricingName,
modelSettings,
updateOption,
]
diff --git a/web/default/src/features/models/components/drawers/model-pricing-config.ts b/web/default/src/features/models/components/drawers/model-pricing-config.ts
new file mode 100644
index 00000000..a90ee561
--- /dev/null
+++ b/web/default/src/features/models/components/drawers/model-pricing-config.ts
@@ -0,0 +1,159 @@
+/*
+Copyright (C) 2023-2026 MAX-API-Next
+
+This program is free software: you can redistribute it and/or modify
+it under the terms of the GNU Affero General Public License as
+published by the Free Software Foundation, either version 3 of the
+License, or (at your option) any later version.
+
+This program is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+GNU Affero General Public License for more details.
+
+You should have received a copy of the GNU Affero General Public License
+along with this program. If not, see .
+
+For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API/issues
+*/
+import type { ModelSettings } from '@/features/system-settings/types'
+import { safeJsonParse } from '@/features/system-settings/utils/json-parser'
+
+export type PricingMode = 'per-token' | 'per-request'
+
+const NON_NEGATIVE_DECIMAL_PATTERN =
+ /^(?:\d+(?:\.\d*)?|\.\d+)(?:[eE][+-]?\d+)?$/
+
+export function parseNonNegativeFiniteNumber(
+ value: string
+): number | undefined {
+ const normalized = value.trim()
+ if (!normalized || !NON_NEGATIVE_DECIMAL_PATTERN.test(normalized)) {
+ return undefined
+ }
+ const parsed = Number(normalized)
+ return Number.isFinite(parsed) && parsed >= 0 ? parsed : undefined
+}
+
+export function parsePricingMap(rawMap: string): Record {
+ const parsed = safeJsonParse>(rawMap, {
+ fallback: {},
+ silent: true,
+ })
+ return Object.fromEntries(
+ Object.entries(parsed).filter(
+ (entry): entry is [string, number] =>
+ typeof entry[1] === 'number' &&
+ Number.isFinite(entry[1]) &&
+ entry[1] >= 0
+ )
+ )
+}
+
+type PricingFields = {
+ price: string
+ ratio: string
+ cacheRatio: string
+ completionRatio: string
+ imageRatio: string
+ audioRatio: string
+ audioCompletionRatio: string
+}
+
+type PricingConfig = {
+ mode: PricingMode
+ fields: PricingFields
+ promptPrice: string
+ completionPrice: string
+ advancedOpen: boolean
+}
+
+const EMPTY_PRICING_FIELDS: PricingFields = {
+ price: '',
+ ratio: '',
+ cacheRatio: '',
+ completionRatio: '',
+ imageRatio: '',
+ audioRatio: '',
+ audioCompletionRatio: '',
+}
+
+const EMPTY_PRICING_CONFIG: PricingConfig = {
+ mode: 'per-token',
+ fields: EMPTY_PRICING_FIELDS,
+ promptPrice: '',
+ completionPrice: '',
+ advancedOpen: false,
+}
+
+function lookupModelRatio(
+ rawMap: string,
+ modelName: string
+): number | undefined {
+ return parsePricingMap(rawMap)[modelName]
+}
+
+export function readPricingConfig(
+ settings: ModelSettings | null,
+ modelName: string
+): PricingConfig {
+ if (!settings || !modelName) return EMPTY_PRICING_CONFIG
+
+ const price = lookupModelRatio(settings.ModelPrice, modelName)
+ const ratio = lookupModelRatio(settings.ModelRatio, modelName)
+ const cacheRatio = lookupModelRatio(settings.CacheRatio, modelName)
+ const completionRatio = lookupModelRatio(settings.CompletionRatio, modelName)
+ const imageRatio = lookupModelRatio(settings.ImageRatio, modelName)
+ const audioRatio = lookupModelRatio(settings.AudioRatio, modelName)
+ const audioCompletionRatio = lookupModelRatio(
+ settings.AudioCompletionRatio,
+ modelName
+ )
+
+ if (price !== undefined && price !== null) {
+ return {
+ ...EMPTY_PRICING_CONFIG,
+ mode: 'per-request',
+ fields: { ...EMPTY_PRICING_FIELDS, price: price.toString() },
+ }
+ }
+
+ let promptPrice = ''
+ let completionPrice = ''
+ if (ratio !== undefined && ratio !== null) {
+ const tokenPrice = ratio * 2
+ promptPrice = tokenPrice.toString()
+ if (completionRatio !== undefined && completionRatio !== null) {
+ completionPrice = (tokenPrice * completionRatio).toString()
+ }
+ }
+
+ return {
+ mode: 'per-token',
+ fields: {
+ price: '',
+ ratio: ratio?.toString() || '',
+ cacheRatio: cacheRatio?.toString() || '',
+ completionRatio: completionRatio?.toString() || '',
+ imageRatio: imageRatio?.toString() || '',
+ audioRatio: audioRatio?.toString() || '',
+ audioCompletionRatio: audioCompletionRatio?.toString() || '',
+ },
+ promptPrice,
+ completionPrice,
+ advancedOpen: [
+ cacheRatio,
+ imageRatio,
+ audioRatio,
+ audioCompletionRatio,
+ ].some((value) => value !== undefined && value !== null),
+ }
+}
+
+export function shouldReplacePricingConfig(
+ finalModelName: string,
+ loadedPricingName: string,
+ hasPricingConfig: boolean
+): boolean {
+ return hasPricingConfig || finalModelName === loadedPricingName
+}
diff --git a/web/default/src/features/profile/api.ts b/web/default/src/features/profile/api.ts
index 2fcd9a97..aa4a4d2b 100644
--- a/web/default/src/features/profile/api.ts
+++ b/web/default/src/features/profile/api.ts
@@ -81,7 +81,10 @@ export async function updateUserLanguage(
export async function deleteUserAccount(
data?: DeleteAccountRequest
): Promise {
- const res = await api.delete('/api/user/self', { data })
+ const res = await api.delete('/api/user/self', {
+ ...sensitiveActionConfig,
+ data,
+ })
return res.data
}
diff --git a/web/default/src/features/profile/components/dialogs/delete-account-dialog.tsx b/web/default/src/features/profile/components/dialogs/delete-account-dialog.tsx
index feb22da1..46fdf0e8 100644
--- a/web/default/src/features/profile/components/dialogs/delete-account-dialog.tsx
+++ b/web/default/src/features/profile/components/dialogs/delete-account-dialog.tsx
@@ -23,6 +23,7 @@ import { useTranslation } from 'react-i18next'
import { toast } from 'sonner'
import { useAuthStore } from '@/stores/auth-store'
import { api } from '@/lib/api'
+import { handleServerError } from '@/lib/handle-server-error'
import { Alert, AlertDescription } from '@/components/ui/alert'
import { Button } from '@/components/ui/button'
import {
@@ -35,6 +36,10 @@ import {
} from '@/components/ui/dialog'
import { Input } from '@/components/ui/input'
import { Label } from '@/components/ui/label'
+import {
+ SecureVerificationDialog,
+ useSecureVerification,
+} from '@/features/auth/secure-verification'
import { deleteUserAccount } from '../../api'
// ============================================================================
@@ -57,6 +62,41 @@ export function DeleteAccountDialog({
const { reset } = useAuthStore((state) => state.auth)
const [loading, setLoading] = useState(false)
const [confirmation, setConfirmation] = useState('')
+ const {
+ open: verificationOpen,
+ methods: verificationMethods,
+ state: verificationState,
+ executeVerification,
+ withVerification,
+ cancel: cancelVerification,
+ setCode: setVerificationCode,
+ switchMethod: switchVerificationMethod,
+ } = useSecureVerification()
+
+ const performDelete = async () => {
+ try {
+ setLoading(true)
+ const response = await deleteUserAccount()
+
+ if (!response.success) {
+ throw new Error(response.message || t('Failed to delete account'))
+ }
+
+ toast.success(t('Account deleted successfully'))
+
+ try {
+ await api.get('/api/user/logout')
+ } catch {
+ // The account is already deleted; local auth state still needs clearing.
+ }
+
+ reset()
+ navigate({ to: '/sign-in' })
+ return response
+ } finally {
+ setLoading(false)
+ }
+ }
const handleDelete = async () => {
if (confirmation !== username) {
@@ -65,29 +105,16 @@ export function DeleteAccountDialog({
}
try {
- setLoading(true)
- const response = await deleteUserAccount()
-
- if (response.success) {
- toast.success(t('Account deleted successfully'))
-
- // Logout and redirect
- try {
- await api.get('/api/user/logout')
- } catch {
- // Ignore logout errors
- }
-
- reset()
- localStorage.removeItem('user')
- navigate({ to: '/sign-in' })
- } else {
- toast.error(response.message || t('Failed to delete account'))
- }
- } catch (_error) {
- toast.error(t('Failed to delete account'))
- } finally {
- setLoading(false)
+ await withVerification(performDelete, {
+ preferredMethod: 'passkey',
+ scope: 'account_delete',
+ title: t('Security verification'),
+ description: t(
+ 'Verify your identity before permanently deleting your account.'
+ ),
+ })
+ } catch (error) {
+ handleServerError(error, { fallback: t('Failed to delete account') })
}
}
@@ -101,7 +128,8 @@ export function DeleteAccountDialog({
}
return (
-
+
+ {
+ if (!nextOpen) cancelVerification()
+ }}
+ methods={verificationMethods}
+ state={verificationState}
+ onVerify={async (method, code) => {
+ try {
+ await executeVerification(method, code)
+ } catch {
+ // executeVerification already reports verification failures.
+ }
+ }}
+ onCancel={cancelVerification}
+ onCodeChange={setVerificationCode}
+ onMethodChange={switchVerificationMethod}
+ />
+
)
}
diff --git a/web/default/src/features/system-settings/models/model-pricing-sheet.tsx b/web/default/src/features/system-settings/models/model-pricing-sheet.tsx
index f15ced41..3cf94889 100644
--- a/web/default/src/features/system-settings/models/model-pricing-sheet.tsx
+++ b/web/default/src/features/system-settings/models/model-pricing-sheet.tsx
@@ -443,7 +443,22 @@ export function ModelPricingSheet({
}
export function ModelPricingEditorPanel(props: ModelPricingEditorPanelProps) {
- const identity = props.editData?.name ?? '__new-model-pricing__'
+ const identity = props.editData
+ ? JSON.stringify([
+ props.editData.name,
+ props.editData.billingMode,
+ props.editData.billingExpr,
+ props.editData.requestRuleExpr,
+ props.editData.price,
+ props.editData.ratio,
+ props.editData.cacheRatio,
+ props.editData.createCacheRatio,
+ props.editData.completionRatio,
+ props.editData.imageRatio,
+ props.editData.audioRatio,
+ props.editData.audioCompletionRatio,
+ ])
+ : '__new-model-pricing__'
return
}
diff --git a/web/default/src/features/system-settings/models/tiered-pricing-editor.runtime.test.tsx b/web/default/src/features/system-settings/models/tiered-pricing-editor.runtime.test.tsx
index 7b82f923..267fb8b9 100644
--- a/web/default/src/features/system-settings/models/tiered-pricing-editor.runtime.test.tsx
+++ b/web/default/src/features/system-settings/models/tiered-pricing-editor.runtime.test.tsx
@@ -27,6 +27,7 @@ import {
ModelPricingEditorPanel,
type ModelRatioData,
} from './model-pricing-sheet'
+import { ModelRatioVisualEditor } from './model-ratio-visual-editor'
import { TieredPricingEditor } from './tiered-pricing-editor'
const dom = new JSDOM('', {
@@ -48,6 +49,7 @@ const globalKeys = [
'requestAnimationFrame',
'cancelAnimationFrame',
'ResizeObserver',
+ 'localStorage',
'IS_REACT_ACT_ENVIRONMENT',
] as const
const previousDescriptors = new Map<
@@ -92,6 +94,20 @@ async function unmount(root: Root, container: HTMLElement) {
container.remove()
}
+function getInputValueByLabel(container: HTMLElement, labelText: string) {
+ const label = Array.from(container.querySelectorAll('label')).find(
+ (candidate) => candidate.textContent === labelText
+ )
+ assert.ok(label, `missing ${labelText} label`)
+
+ const inputId = label.getAttribute('for')
+ assert.ok(inputId, `missing input id for ${labelText}`)
+
+ const input = dom.window.document.getElementById(inputId) as HTMLInputElement
+ assert.ok(input, `missing input for ${labelText}`)
+ return input.value
+}
+
const modelA: ModelRatioData = {
name: 'model-a',
billingMode: 'tiered_expr',
@@ -144,6 +160,7 @@ before(async () => {
)
setGlobal('cancelAnimationFrame', (handle: number) => clearTimeout(handle))
setGlobal('ResizeObserver', ResizeObserverStub)
+ setGlobal('localStorage', window.localStorage)
setGlobal('IS_REACT_ACT_ENVIRONMENT', true)
})
@@ -228,6 +245,8 @@ describe('TieredPricingEditor runtime behavior', () => {
try {
await act(async () => root.render(renderPanel(modelA)))
assert.match(container.textContent || '', /tier\("a", p \* 1 \+ c \* 2\)/)
+ assert.equal(getInputValueByLabel(container, 'Input price'), '1')
+ assert.equal(getInputValueByLabel(container, 'Output price'), '2')
await act(async () => root.render(renderPanel(modelB)))
@@ -235,6 +254,103 @@ describe('TieredPricingEditor runtime behavior', () => {
assert.match(text, /tier\("b", p \* 3 \+ c \* 4\)/)
assert.doesNotMatch(text, /tier\("a", p \* 1 \+ c \* 2\)/)
assert.doesNotMatch(text, /tier\("base", p \* 0 \+ c \* 0\)/)
+ assert.equal(getInputValueByLabel(container, 'Input price'), '3')
+ assert.equal(getInputValueByLabel(container, 'Output price'), '4')
+ } finally {
+ await unmount(root, container)
+ }
+ })
+
+ test('replaces expression state when the selected model snapshot changes', async () => {
+ const container = createContainer()
+ const root = createRoot(container)
+ const updatedModelA: ModelRatioData = {
+ ...modelA,
+ billingExpr: 'tier("a-updated", p * 5 + c * 6)',
+ }
+ const renderPanel = (editData: ModelRatioData) => (
+
+ undefined}
+ onCancel={() => undefined}
+ />
+
+ )
+
+ try {
+ await act(async () => root.render(renderPanel(modelA)))
+ assert.equal(getInputValueByLabel(container, 'Input price'), '1')
+ assert.equal(getInputValueByLabel(container, 'Output price'), '2')
+
+ await act(async () => root.render(renderPanel(updatedModelA)))
+ assert.equal(getInputValueByLabel(container, 'Input price'), '5')
+ assert.equal(getInputValueByLabel(container, 'Output price'), '6')
+ } finally {
+ await unmount(root, container)
+ }
+ })
+
+ test('updates expression prices when selecting another model row', async () => {
+ const container = createContainer()
+ const root = createRoot(container)
+ const modeMap = JSON.stringify({
+ [modelA.name]: 'tiered_expr',
+ [modelB.name]: 'tiered_expr',
+ })
+ const exprMap = JSON.stringify({
+ [modelA.name]: modelA.billingExpr,
+ [modelB.name]: modelB.billingExpr,
+ })
+ const renderEditor = () => (
+
+ undefined}
+ />
+
+ )
+ const selectRow = async (modelName: string) => {
+ const cell = Array.from(container.querySelectorAll('td')).find(
+ (candidate) => candidate.textContent?.trim().startsWith(modelName)
+ )
+ assert.ok(cell, `missing ${modelName} table cell`)
+ const row = cell.closest('tr')
+ assert.ok(row, `missing ${modelName} table row`)
+ await act(async () =>
+ row.dispatchEvent(new MouseEvent('click', { bubbles: true }))
+ )
+ }
+
+ try {
+ await act(async () => root.render(renderEditor()))
+
+ await selectRow(modelA.name)
+ assert.equal(getInputValueByLabel(container, 'Input price'), '1')
+ assert.equal(getInputValueByLabel(container, 'Output price'), '2')
+
+ await selectRow(modelB.name)
+ assert.equal(getInputValueByLabel(container, 'Input price'), '3')
+ assert.equal(getInputValueByLabel(container, 'Output price'), '4')
} finally {
await unmount(root, container)
}
diff --git a/web/default/src/i18n/locales/en.json b/web/default/src/i18n/locales/en.json
index 3cafee7d..8d3b7d36 100644
--- a/web/default/src/i18n/locales/en.json
+++ b/web/default/src/i18n/locales/en.json
@@ -1243,6 +1243,7 @@
"Create, revoke, and audit API tokens.": "Create, revoke, and audit API tokens.",
"Created": "Created",
"Created At": "Created At",
+ "Created Time Path": "Created Time Path",
"Creates a Pancake product in the saved store using this plan’s title and price. Requires Waffo Pancake to be fully configured in Payment settings first.": "Creates a Pancake product in the saved store using this plan’s title and price. Requires Waffo Pancake to be fully configured in Payment settings first.",
"Creating...": "Creating...",
"Creation failed": "Creation failed",
@@ -3935,8 +3936,10 @@
"Response Content": "Response Content",
"Response Mapping": "Response Mapping",
"Response path for failure reason": "Response path for failure reason",
+ "Response path for task creation timestamp": "Response path for task creation timestamp",
"Response path for task progress": "Response path for task progress",
"Response path for task status": "Response path for task status",
+ "Response path for task update timestamp": "Response path for task update timestamp",
"Response path for upstream task ID": "Response path for upstream task ID",
"Response Time": "Response Time",
"Response time: {{duration}}": "Response time: {{duration}}",
@@ -4140,6 +4143,7 @@
"Security boundary": "Security boundary",
"Security Check": "Security Check",
"Security verification": "Security verification",
+ "Verify your identity before permanently deleting your account.": "Verify your identity before permanently deleting your account.",
"Seed": "Seed",
"Seedance": "Seedance",
"Select": "Select",
@@ -4945,6 +4949,7 @@
"Updated": "Updated",
"Updated daily": "Updated daily",
"Updated successfully": "Updated successfully",
+ "Updated Time Path": "Updated Time Path",
"Updating all channel balances. This may take a while. Please refresh to see results.": "Updating all channel balances. This may take a while. Please refresh to see results.",
"Updating...": "Updating...",
"Upgrade Group": "Upgrade Group",
diff --git a/web/default/src/i18n/locales/fr.json b/web/default/src/i18n/locales/fr.json
index c2cc84db..38b37915 100644
--- a/web/default/src/i18n/locales/fr.json
+++ b/web/default/src/i18n/locales/fr.json
@@ -1243,6 +1243,7 @@
"Create, revoke, and audit API tokens.": "Créer, révoquer et auditer les jetons API.",
"Created": "Créé",
"Created At": "Créé le",
+ "Created Time Path": "Chemin de l’heure de création",
"Creates a Pancake product in the saved store using this plan’s title and price. Requires Waffo Pancake to be fully configured in Payment settings first.": "Crée un produit Pancake dans la boutique enregistrée avec le titre et le prix de ce forfait. Waffo Pancake doit d’abord être entièrement configuré dans les paramètres de paiement.",
"Creating...": "Création...",
"Creation failed": "Échec de la création",
@@ -3935,8 +3936,10 @@
"Response Content": "Contenu de la réponse",
"Response Mapping": "Mappage de réponse",
"Response path for failure reason": "Chemin de réponse pour la raison d’échec",
+ "Response path for task creation timestamp": "Chemin de réponse de l’horodatage de création de la tâche",
"Response path for task progress": "Chemin de réponse pour la progression",
"Response path for task status": "Chemin de réponse pour le statut",
+ "Response path for task update timestamp": "Chemin de réponse de l’horodatage de mise à jour de la tâche",
"Response path for upstream task ID": "Chemin de réponse pour l’ID de tâche amont",
"Response Time": "Temps de réponse",
"Response time: {{duration}}": "Temps de réponse : {{duration}}",
@@ -4140,6 +4143,7 @@
"Security boundary": "Frontière de sécurité",
"Security Check": "Vérification de sécurité",
"Security verification": "Vérification de sécurité",
+ "Verify your identity before permanently deleting your account.": "Vérifiez votre identité avant de supprimer définitivement votre compte.",
"Seed": "Graine",
"Seedance": "Seedance",
"Select": "Sélectionner",
@@ -4945,6 +4949,7 @@
"Updated": "Mis à jour",
"Updated daily": "Mis à jour quotidiennement",
"Updated successfully": "Mise à jour réussie",
+ "Updated Time Path": "Chemin de l’heure de mise à jour",
"Updating all channel balances. This may take a while. Please refresh to see results.": "Mise à jour de tous les soldes des canaux. Cela peut prendre un certain temps. Veuillez actualiser pour voir les résultats.",
"Updating...": "Mise à jour...",
"Upgrade Group": "Groupe de mise à niveau",
diff --git a/web/default/src/i18n/locales/ja.json b/web/default/src/i18n/locales/ja.json
index 87bc3973..db54eac4 100644
--- a/web/default/src/i18n/locales/ja.json
+++ b/web/default/src/i18n/locales/ja.json
@@ -1243,6 +1243,7 @@
"Create, revoke, and audit API tokens.": "APIトークンを作成、取り消し、監査。",
"Created": "作成済み",
"Created At": "作成日時",
+ "Created Time Path": "作成日時パス",
"Creates a Pancake product in the saved store using this plan’s title and price. Requires Waffo Pancake to be fully configured in Payment settings first.": "保存済みストアに、このプランのタイトルと価格を使って Pancake 商品を作成します。事前に支払い設定で Waffo Pancake を完全に設定する必要があります。",
"Creating...": "作成中...",
"Creation failed": "作成に失敗しました",
@@ -3935,8 +3936,10 @@
"Response Content": "レスポンス内容",
"Response Mapping": "レスポンスマッピング",
"Response path for failure reason": "失敗理由のレスポンスパス",
+ "Response path for task creation timestamp": "タスク作成タイムスタンプのレスポンスパス",
"Response path for task progress": "タスク進行状況のレスポンスパス",
"Response path for task status": "タスクステータスのレスポンスパス",
+ "Response path for task update timestamp": "タスク更新タイムスタンプのレスポンスパス",
"Response path for upstream task ID": "上流タスク ID のレスポンスパス",
"Response Time": "応答時間",
"Response time: {{duration}}": "応答時間: {{duration}}",
@@ -4140,6 +4143,7 @@
"Security boundary": "セキュリティ境界",
"Security Check": "セキュリティチェック",
"Security verification": "セキュリティ確認",
+ "Verify your identity before permanently deleting your account.": "アカウントを完全に削除する前に本人確認を行ってください。",
"Seed": "シード",
"Seedance": "Seedance（動画生成）",
"Select": "選択",
@@ -4945,6 +4949,7 @@
"Updated": "更新されました",
"Updated daily": "毎日更新",
"Updated successfully": "正常に更新されました",
+ "Updated Time Path": "更新日時パス",
"Updating all channel balances. This may take a while. Please refresh to see results.": "すべてのチャネル残高を更新中です。これには少し時間がかかる場合があります。結果を確認するには更新してください。",
"Updating...": "更新中...",
"Upgrade Group": "グループをアップグレード",
diff --git a/web/default/src/i18n/locales/ru.json b/web/default/src/i18n/locales/ru.json
index d51719a9..42632cd3 100644
--- a/web/default/src/i18n/locales/ru.json
+++ b/web/default/src/i18n/locales/ru.json
@@ -1243,6 +1243,7 @@
"Create, revoke, and audit API tokens.": "Создать, отозвать и аудитировать токены API.",
"Created": "Создано",
"Created At": "Дата создания",
+ "Created Time Path": "Путь времени создания",
"Creates a Pancake product in the saved store using this plan’s title and price. Requires Waffo Pancake to be fully configured in Payment settings first.": "Создает продукт Pancake в сохраненном магазине с названием и ценой этого плана. Сначала необходимо полностью настроить Waffo Pancake в настройках платежей.",
"Creating...": "Создание...",
"Creation failed": "Создание не удалось",
@@ -3935,8 +3936,10 @@
"Response Content": "Содержимое ответа",
"Response Mapping": "Сопоставление ответа",
"Response path for failure reason": "Путь ответа для причины ошибки",
+ "Response path for task creation timestamp": "Путь в ответе к метке времени создания задачи",
"Response path for task progress": "Путь ответа для прогресса задачи",
"Response path for task status": "Путь ответа для статуса задачи",
+ "Response path for task update timestamp": "Путь в ответе к метке времени обновления задачи",
"Response path for upstream task ID": "Путь ответа для ID задачи апстрима",
"Response Time": "Время ответа",
"Response time: {{duration}}": "Время ответа: {{duration}}",
@@ -4140,6 +4143,7 @@
"Security boundary": "Граница безопасности",
"Security Check": "Проверка безопасности",
"Security verification": "Подтверждение безопасности",
+ "Verify your identity before permanently deleting your account.": "Подтвердите свою личность перед окончательным удалением аккаунта.",
"Seed": "Сид",
"Seedance": "Seedance (генерация видео)",
"Select": "Выбрать",
@@ -4945,6 +4949,7 @@
"Updated": "Обновлено",
"Updated daily": "Обновляется ежедневно",
"Updated successfully": "Обновлено успешно",
+ "Updated Time Path": "Путь времени обновления",
"Updating all channel balances. This may take a while. Please refresh to see results.": "Обновление балансов всех каналов. Это может занять некоторое время. Пожалуйста, обновите страницу, чтобы увидеть результаты.",
"Updating...": "Обновление...",
"Upgrade Group": "Повысить группу",
diff --git a/web/default/src/i18n/locales/vi.json b/web/default/src/i18n/locales/vi.json
index 4010ef4a..1dfd1b80 100644
--- a/web/default/src/i18n/locales/vi.json
+++ b/web/default/src/i18n/locales/vi.json
@@ -1243,6 +1243,7 @@
"Create, revoke, and audit API tokens.": "Tạo, thu hồi và kiểm toán token API.",
"Created": "Đã tạo",
"Created At": "Ngày tạo",
+ "Created Time Path": "Đường dẫn thời gian tạo",
"Creates a Pancake product in the saved store using this plan’s title and price. Requires Waffo Pancake to be fully configured in Payment settings first.": "Tạo một sản phẩm Pancake trong cửa hàng đã lưu bằng tiêu đề và giá của gói này. Trước tiên cần cấu hình đầy đủ Waffo Pancake trong cài đặt Thanh toán.",
"Creating...": "Đang tạo...",
"Creation failed": "Tạo thất bại",
@@ -3935,8 +3936,10 @@
"Response Content": "Nội dung phản hồi",
"Response Mapping": "Ánh xạ phản hồi",
"Response path for failure reason": "Đường dẫn phản hồi cho lý do thất bại",
+ "Response path for task creation timestamp": "Đường dẫn phản hồi cho dấu thời gian tạo tác vụ",
"Response path for task progress": "Đường dẫn phản hồi cho tiến độ tác vụ",
"Response path for task status": "Đường dẫn phản hồi cho trạng thái tác vụ",
+ "Response path for task update timestamp": "Đường dẫn phản hồi cho dấu thời gian cập nhật tác vụ",
"Response path for upstream task ID": "Đường dẫn phản hồi cho ID tác vụ upstream",
"Response Time": "Thời gian phản hồi",
"Response time: {{duration}}": "Thời gian phản hồi: {{duration}}",
@@ -4140,6 +4143,7 @@
"Security boundary": "Ranh giới bảo mật",
"Security Check": "Kiểm tra bảo mật",
"Security verification": "Xác minh bảo mật",
+ "Verify your identity before permanently deleting your account.": "Xác minh danh tính trước khi xóa vĩnh viễn tài khoản của bạn.",
"Seed": "Seed",
"Seedance": "Seedance",
"Select": "Chọn",
@@ -4945,6 +4949,7 @@
"Updated": "Đã cập nhật",
"Updated daily": "Cập nhật hàng ngày",
"Updated successfully": "Cập nhật thành công",
+ "Updated Time Path": "Đường dẫn thời gian cập nhật",
"Updating all channel balances. This may take a while. Please refresh to see results.": "Đang cập nhật tất cả số dư kênh. Quá trình này có thể mất một chút thời gian. Vui lòng làm mới để xem kết quả.",
"Updating...": "Đang cập nhật...",
"Upgrade Group": "Nhóm nâng cấp",
diff --git a/web/default/src/i18n/locales/zh.json b/web/default/src/i18n/locales/zh.json
index 1a72aa02..2888408b 100644
--- a/web/default/src/i18n/locales/zh.json
+++ b/web/default/src/i18n/locales/zh.json
@@ -1243,6 +1243,7 @@
"Create, revoke, and audit API tokens.": "创建、撤销和审计 API 令牌。",
"Created": "创建时间",
"Created At": "创建时间",
+ "Created Time Path": "创建时间路径",
"Creates a Pancake product in the saved store using this plan’s title and price. Requires Waffo Pancake to be fully configured in Payment settings first.": "使用此套餐的标题和价格，在已保存的店铺中创建 Pancake 产品。需要先在支付设置中完整配置 Waffo Pancake。",
"Creating...": "创建中...",
"Creation failed": "创建失败",
@@ -3935,8 +3936,10 @@
"Response Content": "响应内容",
"Response Mapping": "响应映射",
"Response path for failure reason": "失败原因的响应路径",
+ "Response path for task creation timestamp": "任务创建时间戳的响应路径",
"Response path for task progress": "任务进度的响应路径",
"Response path for task status": "任务状态的响应路径",
+ "Response path for task update timestamp": "任务更新时间戳的响应路径",
"Response path for upstream task ID": "上游任务 ID 的响应路径",
"Response Time": "响应时间",
"Response time: {{duration}}": "响应时间：{{duration}}",
@@ -4140,6 +4143,7 @@
"Security boundary": "安全边界",
"Security Check": "安全验证",
"Security verification": "安全验证",
+ "Verify your identity before permanently deleting your account.": "永久删除账号前，请先验证您的身份。",
"Seed": "随机种子",
"Seedance": "Seedance / 豆包视频",
"Select": "选择",
@@ -4945,6 +4949,7 @@
"Updated": "更新时间",
"Updated daily": "每日更新",
"Updated successfully": "更新成功",
+ "Updated Time Path": "更新时间路径",
"Updating all channel balances. This may take a while. Please refresh to see results.": "正在更新所有渠道余额。这可能需要一段时间。请刷新以查看结果。",
"Updating...": "正在更新...",
"Upgrade Group": "升级分组",
diff --git a/web/default/src/lib/handle-server-error.ts b/web/default/src/lib/handle-server-error.ts
index af923e3c..59a6cd34 100644
--- a/web/default/src/lib/handle-server-error.ts
+++ b/web/default/src/lib/handle-server-error.ts
@@ -27,6 +27,10 @@ interface SafeErrorDebugInfo {
status?: number
}
+interface HandleServerErrorOptions {
+ fallback?: string
+}
+
export function getSafeErrorDebugInfo(error: unknown): SafeErrorDebugInfo {
if (error instanceof AxiosError) {
return {
@@ -42,13 +46,16 @@ export function getSafeErrorDebugInfo(error: unknown): SafeErrorDebugInfo {
return { name: 'UnknownError', message: 'Unknown error' }
}
-export function handleServerError(error: unknown) {
+export function handleServerError(
+ error: unknown,
+ options: HandleServerErrorOptions = {}
+): void {
if (import.meta.env.DEV) {
// eslint-disable-next-line no-console
console.error('[handleServerError]', getSafeErrorDebugInfo(error))
}
- let errMsg = i18next.t('Something went wrong!')
+ let errMsg = options.fallback || i18next.t('Something went wrong!')
if (
error &&
diff --git a/web/default/src/routes/__root.tsx b/web/default/src/routes/__root.tsx
index d1f88639..eba21b94 100644
--- a/web/default/src/routes/__root.tsx
+++ b/web/default/src/routes/__root.tsx
@@ -29,6 +29,7 @@ import { ThemeCustomizationProvider } from '@/context/theme-customization-provid
import { useSystemConfig } from '@/hooks/use-system-config'
import { Toaster } from '@/components/ui/sonner'
import { NavigationProgress } from '@/components/navigation-progress'
+import { appSessionVerifier } from '@/features/auth/lib/app-session-verifier'
import { saveAffiliateCode } from '@/features/auth/lib/storage'
import { GeneralError } from '@/features/errors/general-error'
import { NotFoundError } from '@/features/errors/not-found-error'
@@ -73,9 +74,6 @@ export const Route = createRootRouteWithContext<{
const needsSetupCheck =
!setupStatusChecked && !pathname.startsWith('/setup')
- // 根路由只负责 setup 检查；受保护路由会在 _authenticated 的
- // beforeLoad 中调用 getSelf() 校验服务端会话并刷新本地用户信息。
-
// 只检查 setup 状态（如果需要）
if (needsSetupCheck) {
const status = await getSetupStatus().catch((error) => {
@@ -91,6 +89,17 @@ export const Route = createRootRouteWithContext<{
}
setupStatusChecked = true
}
+
+ // Restore a valid cookie session before public pages render. Public routes
+ // ignore non-authentication failures and never redirect on a 401; the
+ // protected route still owns access control.
+ await appSessionVerifier.verify().catch((error: unknown) => {
+ if (import.meta.env.DEV) {
+ // eslint-disable-next-line no-console
+ console.warn('[root.beforeLoad] session check failed', error)
+ }
+ return true
+ })
},
component: RootComponent,
notFoundComponent: NotFoundError,
diff --git a/web/default/src/routes/_authenticated/route.tsx b/web/default/src/routes/_authenticated/route.tsx
index 4c9ff631..7ca06333 100644
--- a/web/default/src/routes/_authenticated/route.tsx
+++ b/web/default/src/routes/_authenticated/route.tsx
@@ -17,48 +17,17 @@ along with this program. If not, see .
For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API/issues
*/
import { createFileRoute, redirect } from '@tanstack/react-router'
-import { useAuthStore } from '@/stores/auth-store'
-import { getSelf } from '@/lib/api'
import { AuthenticatedLayout } from '@/components/layout'
-
-// 内存中的验证标记，避免同一会话中重复验证
-let sessionVerified = false
+import { appSessionVerifier } from '@/features/auth/lib/app-session-verifier'
export const Route = createFileRoute('/_authenticated')({
beforeLoad: async ({ location }) => {
- const { auth } = useAuthStore.getState()
-
- // 如果本地没有用户信息，直接跳转登录页
- if (!auth.user) {
+ if (!(await appSessionVerifier.verify())) {
throw redirect({
to: '/sign-in',
search: { redirect: location.href },
})
}
-
- // 本地有用户信息，但需要验证 session 是否有效（每个会话只验证一次）
- if (!sessionVerified) {
- const res = await getSelf().catch((err: unknown) => {
- if (
- (err as { response?: { status?: number } })?.response?.status === 401
- ) {
- return { success: false }
- }
- throw err
- })
- if (res?.success && res.data) {
- // 验证成功，更新用户信息（可能有变化）
- auth.setUser(res.data)
- sessionVerified = true
- } else if (res) {
- // 验证失败，清除本地缓存并跳转登录页
- auth.reset()
- throw redirect({
- to: '/sign-in',
- search: { redirect: location.href },
- })
- }
- }
},
component: AuthenticatedLayout,
})
diff --git a/web/default/src/stores/auth-store.ts b/web/default/src/stores/auth-store.ts
index ecc07cfc..84764bab 100644
--- a/web/default/src/stores/auth-store.ts
+++ b/web/default/src/stores/auth-store.ts
@@ -58,48 +58,15 @@ interface AuthState {
}
}
-export const useAuthStore = create()((set) => {
- // Restore user info from localStorage
- const initUser = (() => {
- try {
- if (typeof window !== 'undefined') {
- const saved = window.localStorage.getItem('user')
- return saved ? JSON.parse(saved) : null
- }
- } catch {
- // Clear dirty data when parsing fails
- if (typeof window !== 'undefined') {
- window.localStorage.removeItem('user')
- }
- }
- return null
- })()
-
- return {
- auth: {
- user: initUser,
- setUser: (user) =>
- set((state) => {
- // Persist user to localStorage
- if (typeof window !== 'undefined') {
- if (user) {
- window.localStorage.setItem('user', JSON.stringify(user))
- } else {
- window.localStorage.removeItem('user')
- }
- }
- return { ...state, auth: { ...state.auth, user } }
- }),
- reset: () =>
- set((state) => {
- if (typeof window !== 'undefined') {
- window.localStorage.removeItem('user')
- }
- return {
- ...state,
- auth: { ...state.auth, user: null },
- }
- }),
- },
- }
-})
+export const useAuthStore = create()((set) => ({
+ auth: {
+ user: null,
+ setUser: (user) =>
+ set((state) => ({ ...state, auth: { ...state.auth, user } })),
+ reset: () =>
+ set((state) => ({
+ ...state,
+ auth: { ...state.auth, user: null },
+ })),
+ },
+}))
diff --git a/web/default/src/stores/notification-store.test.ts b/web/default/src/stores/notification-store.test.ts
deleted file mode 100644
index f46f4039..00000000
--- a/web/default/src/stores/notification-store.test.ts
+++ /dev/null
@@ -1,151 +0,0 @@
-/*
-Copyright (C) 2023-2026 MAX-API-Next
-
-This program is free software: you can redistribute it and/or modify
-it under the terms of the GNU Affero General Public License as
-published by the Free Software Foundation, either version 3 of the
-License, or (at your option) any later version.
-
-This program is distributed in the hope that it will be useful,
-but WITHOUT ANY WARRANTY; without even the implied warranty of
-MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
-GNU Affero General Public License for more details.
-
-You should have received a copy of the GNU Affero General Public License
-along with this program. If not, see .
-
-For commercial licensing, please contact https://github.com/MAX-API-Next/MAX-API/issues
-*/
-import assert from 'node:assert/strict'
-import { after, beforeEach, describe, test } from 'node:test'
-
-const storage = new Map()
-const testLocalStorage = {
- getItem: (key: string) => storage.get(key) ?? null,
- removeItem: (key: string) => {
- storage.delete(key)
- },
- setItem: (key: string, value: string) => {
- storage.set(key, value)
- },
-}
-const previousLocalStorageDescriptor = Object.getOwnPropertyDescriptor(
- globalThis,
- 'localStorage'
-)
-const previousWindowDescriptor = Object.getOwnPropertyDescriptor(
- globalThis,
- 'window'
-)
-
-Object.defineProperty(globalThis, 'localStorage', {
- configurable: true,
- value: testLocalStorage,
-})
-Object.defineProperty(globalThis, 'window', {
- configurable: true,
- value: {
- localStorage: testLocalStorage,
- },
-})
-
-const { useNotificationStore } = (await import(
- './notification-store'
-)) as typeof import('./notification-store')
-/* eslint-disable no-console */
-const originalWarn = console.warn
-
-console.warn = (...args: unknown[]) => {
- if (
- typeof args[0] === 'string' &&
- args[0].includes("Unable to update item 'notification-storage'")
- ) {
- return
- }
-
- originalWarn(...args)
-}
-
-after(() => {
- console.warn = originalWarn
-
- if (previousLocalStorageDescriptor) {
- Object.defineProperty(
- globalThis,
- 'localStorage',
- previousLocalStorageDescriptor
- )
- } else {
- delete (globalThis as Partial).localStorage
- }
-
- if (previousWindowDescriptor) {
- Object.defineProperty(globalThis, 'window', previousWindowDescriptor)
- } else {
- delete (globalThis as Partial).window
- }
-})
-/* eslint-enable no-console */
-
-describe('notification store', () => {
- beforeEach(() => {
- storage.clear()
- useNotificationStore.setState({
- byUser: {},
- autoOpenedSignatures: [],
- })
- })
-
- test('keeps read announcement state scoped by user', () => {
- const store = useNotificationStore.getState()
-
- store.markAnnouncementsRead('user:1', ['id:10'])
-
- assert.equal(store.isAnnouncementRead('user:1', 'id:10'), true)
- assert.equal(store.isAnnouncementRead('user:2', 'id:10'), false)
- })
-
- test('keeps close-for-today state scoped by user', () => {
- const today = new Date().toDateString()
- const store = useNotificationStore.getState()
-
- store.setClosedUntilDate('user:1', today)
-
- assert.equal(store.isNoticeClosed('user:1'), true)
- assert.equal(store.isNoticeClosed('user:2'), false)
- })
-
- test('migrates legacy persisted notification state to anonymous user scope', async () => {
- storage.set(
- 'notification-storage',
- JSON.stringify({
- state: {
- lastReadNotice: 'legacy notice',
- readAnnouncementKeys: ['id:10'],
- closedUntilDate: 'Mon Jun 29 2026',
- },
- version: 0,
- })
- )
-
- await useNotificationStore.persist.rehydrate()
-
- const anonymousState =
- useNotificationStore.getState().getUserState('anonymous')
- assert.deepEqual(anonymousState, {
- lastReadNotice: 'legacy notice',
- readAnnouncementKeys: ['id:10'],
- closedUntilDate: 'Mon Jun 29 2026',
- })
- })
-
- test('remembers every scoped auto-open signature seen in this page session', () => {
- const store = useNotificationStore.getState()
- const userOneSignature = 'user:1:new-notice'
- const userTwoSignature = 'user:2:new-notice'
-
- assert.equal(store.rememberAutoOpenedSignature(userOneSignature), true)
- assert.equal(store.rememberAutoOpenedSignature(userTwoSignature), true)
- assert.equal(store.rememberAutoOpenedSignature(userOneSignature), false)
- })
-})
diff --git a/web/default/tests/channel-form.test.ts b/web/default/tests/channel-form.test.ts
index 12d6c3d0..5babf0fd 100644
--- a/web/default/tests/channel-form.test.ts
+++ b/web/default/tests/channel-form.test.ts
@@ -78,26 +78,6 @@ function legacyChannel(): Channel {
}
}
-function genericVideoTaskChannel(): Channel {
- return {
- ...legacyChannel(),
- settings: JSON.stringify({
- task_protocol: 'generic_video_task',
- task_protocol_config: {
- submit_path: '/v1/videos',
- query_path: '/v1/videos/{task_id}',
- request_body_mode: 'media_generation',
- request_body_mapping: {
- model: 'model',
- },
- task_id_path: 'data.id',
- status_path: 'data.status',
- },
- keep_existing: true,
- }),
- }
-}
-
test('loads legacy video task protocol as path override only', () => {
const defaults = transformChannelToFormDefaults(legacyChannel())
@@ -119,22 +99,6 @@ test('clears legacy video task protocol when saving path override', () => {
})
})
-test('drops stale request body config when saving generic video task channel', () => {
- const defaults = transformChannelToFormDefaults(genericVideoTaskChannel())
- const savedSettings = buildSettings(defaults)
- const taskProtocolConfig = savedSettings.task_protocol_config as Record<
- string,
- unknown
- >
-
- assert.equal(defaults.video_task_protocol_enabled, true)
- assert.equal(taskProtocolConfig.request_body_mode, undefined)
- assert.equal(taskProtocolConfig.request_body_mapping, undefined)
- assert.equal(taskProtocolConfig.request_body_defaults, undefined)
- assert.equal(taskProtocolConfig.task_id_path, 'data.id')
- assert.equal(taskProtocolConfig.status_path, 'data.status')
-})
-
test('channel type switch clears scoped base url and other values', () => {
assert.deepEqual(getChannelTypeScopedFieldDefaults(49), {
base_url: '',