diff --git a/common/init.go b/common/init.go index aa41e44..434149b 100644 --- a/common/init.go +++ b/common/init.go @@ -146,30 +146,30 @@ func InitEnv() { } func initConstantEnv() { - constant.StreamingTimeout = GetEnvOrDefault("STREAMING_TIMEOUT", 300) - constant.DifyDebug = GetEnvOrDefaultBool("DIFY_DEBUG", true) - constant.MaxFileDownloadMB = GetEnvOrDefault("MAX_FILE_DOWNLOAD_MB", 64) - constant.StreamScannerMaxBufferMB = GetEnvOrDefault("STREAM_SCANNER_MAX_BUFFER_MB", 128) + constant.StreamingTimeout = GetEnvOrDefault("STREAMING_TIMEOUT", constant.DefaultStreamingTimeout) + constant.DifyDebug = GetEnvOrDefaultBool("DIFY_DEBUG", constant.DefaultDifyDebug) + constant.MaxFileDownloadMB = GetEnvOrDefault("MAX_FILE_DOWNLOAD_MB", constant.DefaultMaxFileDownloadMB) + constant.StreamScannerMaxBufferMB = GetEnvOrDefault("STREAM_SCANNER_MAX_BUFFER_MB", constant.DefaultStreamScannerMaxBufferMB) // MaxRequestBodyMB 请求体最大大小（解压后），用于防止超大请求/zip bomb导致内存暴涨 - constant.MaxRequestBodyMB = GetEnvOrDefault("MAX_REQUEST_BODY_MB", 128) - constant.AnonymousRequestBodyLimitKB = GetEnvOrDefault("ANONYMOUS_REQUEST_BODY_LIMIT_KB", 512) + constant.MaxRequestBodyMB = GetEnvOrDefault("MAX_REQUEST_BODY_MB", constant.DefaultMaxRequestBodyMB) + constant.AnonymousRequestBodyLimitKB = GetEnvOrDefault("ANONYMOUS_REQUEST_BODY_LIMIT_KB", constant.DefaultAnonymousRequestBodyLimitKB) // ForceStreamOption 覆盖请求参数，强制返回usage信息 - constant.ForceStreamOption = GetEnvOrDefaultBool("FORCE_STREAM_OPTION", true) - constant.CountToken = GetEnvOrDefaultBool("CountToken", true) - constant.GetMediaToken = GetEnvOrDefaultBool("GET_MEDIA_TOKEN", true) - constant.GetMediaTokenNotStream = GetEnvOrDefaultBool("GET_MEDIA_TOKEN_NOT_STREAM", false) - constant.UpdateTask = GetEnvOrDefaultBool("UPDATE_TASK", true) - constant.AzureDefaultAPIVersion = GetEnvOrDefaultString("AZURE_DEFAULT_API_VERSION", "2025-04-01-preview") - constant.NotifyLimitCount = GetEnvOrDefault("NOTIFY_LIMIT_COUNT", 2) - constant.NotificationLimitDurationMinute = GetEnvOrDefault("NOTIFICATION_LIMIT_DURATION_MINUTE", 10) + constant.ForceStreamOption = GetEnvOrDefaultBool("FORCE_STREAM_OPTION", constant.DefaultForceStreamOption) + constant.CountToken = GetEnvOrDefaultBool("CountToken", constant.DefaultCountToken) + constant.GetMediaToken = GetEnvOrDefaultBool("GET_MEDIA_TOKEN", constant.DefaultGetMediaToken) + constant.GetMediaTokenNotStream = GetEnvOrDefaultBool("GET_MEDIA_TOKEN_NOT_STREAM", constant.DefaultGetMediaTokenNotStream) + constant.UpdateTask = GetEnvOrDefaultBool("UPDATE_TASK", constant.DefaultUpdateTask) + constant.AzureDefaultAPIVersion = GetEnvOrDefaultString("AZURE_DEFAULT_API_VERSION", constant.DefaultAzureDefaultAPIVersion) + constant.NotifyLimitCount = GetEnvOrDefault("NOTIFY_LIMIT_COUNT", constant.DefaultNotifyLimitCount) + constant.NotificationLimitDurationMinute = GetEnvOrDefault("NOTIFICATION_LIMIT_DURATION_MINUTE", constant.DefaultNotificationLimitDurationMinute) // GenerateDefaultToken 是否生成初始令牌，默认关闭。 - constant.GenerateDefaultToken = GetEnvOrDefaultBool("GENERATE_DEFAULT_TOKEN", false) + constant.GenerateDefaultToken = GetEnvOrDefaultBool("GENERATE_DEFAULT_TOKEN", constant.DefaultGenerateDefaultToken) // 是否启用错误日志 - constant.ErrorLogEnabled = GetEnvOrDefaultBool("ERROR_LOG_ENABLED", false) + constant.ErrorLogEnabled = GetEnvOrDefaultBool("ERROR_LOG_ENABLED", constant.DefaultErrorLogEnabled) // 任务轮询时查询的最大数量 - constant.TaskQueryLimit = GetEnvOrDefault("TASK_QUERY_LIMIT", 1000) + constant.TaskQueryLimit = GetEnvOrDefault("TASK_QUERY_LIMIT", constant.DefaultTaskQueryLimit) // 异步任务超时时间（分钟），超过此时间未完成的任务将被标记为失败并退款。0 表示禁用。 - constant.TaskTimeoutMinutes = GetEnvOrDefault("TASK_TIMEOUT_MINUTES", 1440) + constant.TaskTimeoutMinutes = GetEnvOrDefault("TASK_TIMEOUT_MINUTES", constant.DefaultTaskTimeoutMinutes) soraPatchStr := GetEnvOrDefaultString("TASK_PRICE_PATCH", "") if soraPatchStr != "" { diff --git a/constant/env.go b/constant/env.go index 512bfc3..446c514 100644 --- a/constant/env.go +++ b/constant/env.go @@ -1,23 +1,44 @@ package constant -var StreamingTimeout int -var DifyDebug bool -var MaxFileDownloadMB int -var StreamScannerMaxBufferMB int -var ForceStreamOption bool -var CountToken bool -var GetMediaToken bool -var GetMediaTokenNotStream bool -var UpdateTask bool -var MaxRequestBodyMB int -var AnonymousRequestBodyLimitKB int -var AzureDefaultAPIVersion string -var NotifyLimitCount int -var NotificationLimitDurationMinute int -var GenerateDefaultToken bool -var ErrorLogEnabled bool -var TaskQueryLimit int -var TaskTimeoutMinutes int +const ( + DefaultStreamingTimeout = 300 + DefaultMaxFileDownloadMB = 64 + DefaultStreamScannerMaxBufferMB = 128 + DefaultMaxRequestBodyMB = 128 + DefaultAnonymousRequestBodyLimitKB = 512 + DefaultAzureDefaultAPIVersion = "2025-04-01-preview" + DefaultNotifyLimitCount = 2 + DefaultNotificationLimitDurationMinute = 10 + DefaultTaskQueryLimit = 1000 + DefaultTaskTimeoutMinutes = 1440 + DefaultDifyDebug = true + DefaultForceStreamOption = true + DefaultCountToken = true + DefaultGetMediaToken = true + DefaultGetMediaTokenNotStream = false + DefaultUpdateTask = true + DefaultGenerateDefaultToken = false + DefaultErrorLogEnabled = false +) + +var StreamingTimeout = DefaultStreamingTimeout +var DifyDebug = DefaultDifyDebug +var MaxFileDownloadMB = DefaultMaxFileDownloadMB +var StreamScannerMaxBufferMB = DefaultStreamScannerMaxBufferMB +var ForceStreamOption = DefaultForceStreamOption +var CountToken = DefaultCountToken +var GetMediaToken = DefaultGetMediaToken +var GetMediaTokenNotStream = DefaultGetMediaTokenNotStream +var UpdateTask = DefaultUpdateTask +var MaxRequestBodyMB = DefaultMaxRequestBodyMB +var AnonymousRequestBodyLimitKB = DefaultAnonymousRequestBodyLimitKB +var AzureDefaultAPIVersion = DefaultAzureDefaultAPIVersion +var NotifyLimitCount = DefaultNotifyLimitCount +var NotificationLimitDurationMinute = DefaultNotificationLimitDurationMinute +var GenerateDefaultToken = DefaultGenerateDefaultToken +var ErrorLogEnabled = DefaultErrorLogEnabled +var TaskQueryLimit = DefaultTaskQueryLimit +var TaskTimeoutMinutes = DefaultTaskTimeoutMinutes // temporary variable for sora patch, will be removed in future var TaskPricePatches []string diff --git a/constant/env_test.go b/constant/env_test.go new file mode 100644 index 0000000..b54a96c --- /dev/null +++ b/constant/env_test.go @@ -0,0 +1,31 @@ +package constant + +import "testing" + +func TestEnvironmentBackedDefaultsAreInitialized(t *testing.T) { + checks := map[string]bool{ + "StreamingTimeout": StreamingTimeout == DefaultStreamingTimeout, + "DifyDebug": DifyDebug == DefaultDifyDebug, + "MaxFileDownloadMB": MaxFileDownloadMB == DefaultMaxFileDownloadMB, + "StreamScannerMaxBufferMB": StreamScannerMaxBufferMB == DefaultStreamScannerMaxBufferMB, + "ForceStreamOption": ForceStreamOption == DefaultForceStreamOption, + "CountToken": CountToken == DefaultCountToken, + "GetMediaToken": GetMediaToken == DefaultGetMediaToken, + "GetMediaTokenNotStream": GetMediaTokenNotStream == DefaultGetMediaTokenNotStream, + "UpdateTask": UpdateTask == DefaultUpdateTask, + "MaxRequestBodyMB": MaxRequestBodyMB == DefaultMaxRequestBodyMB, + "AnonymousRequestBodyLimitKB": AnonymousRequestBodyLimitKB == DefaultAnonymousRequestBodyLimitKB, + "AzureDefaultAPIVersion": AzureDefaultAPIVersion == DefaultAzureDefaultAPIVersion, + "NotifyLimitCount": NotifyLimitCount == DefaultNotifyLimitCount, + "NotificationLimitDurationMinute": NotificationLimitDurationMinute == DefaultNotificationLimitDurationMinute, + "GenerateDefaultToken": GenerateDefaultToken == DefaultGenerateDefaultToken, + "ErrorLogEnabled": ErrorLogEnabled == DefaultErrorLogEnabled, + "TaskQueryLimit": TaskQueryLimit == DefaultTaskQueryLimit, + "TaskTimeoutMinutes": TaskTimeoutMinutes == DefaultTaskTimeoutMinutes, + } + for name, ok := range checks { + if !ok { + t.Fatalf("%s package default does not match expected init default", name) + } + } +} diff --git a/model/log.go b/model/log.go index 8c09db8..3db1ee1 100644 --- a/model/log.go +++ b/model/log.go @@ -95,16 +95,35 @@ func applyLogTypeFilter(tx *gorm.DB, logType int) *gorm.DB { return tx.Where("logs.type = ?", logType) } -func applyRetryLogFilter(tx *gorm.DB) *gorm.DB { - return tx.Where("logs.is_retry = ?", true) +var ensureLogRetryMarkerBackfillCompletedForRead = ensureLogRetryMarkerBackfillCompletedForReadDefault + +func applyRetryLogFilter(tx *gorm.DB) (*gorm.DB, error) { + if err := ensureLogRetryMarkerBackfillCompletedForRead(); err != nil { + return nil, err + } + return tx.Where("logs.is_retry = ?", true), nil } -func applyLogFilter(tx *gorm.DB, filter string) *gorm.DB { +func ensureLogRetryMarkerBackfillCompletedForReadDefault() error { + completed, err := isLogRetryMarkerBackfillCompleted() + if err != nil { + return fmt.Errorf("failed to check log retry marker backfill status before retry log read: %w", err) + } + if completed { + return nil + } + if err := backfillLogRetryMarker(); err != nil { + return fmt.Errorf("failed to backfill log retry markers before retry log read: %w", err) + } + return nil +} + +func applyLogFilter(tx *gorm.DB, filter string) (*gorm.DB, error) { switch filter { case LogFilterRetry: return applyRetryLogFilter(tx) default: - return tx + return tx, nil } } @@ -522,7 +541,10 @@ func RecordTaskBillingLog(params RecordTaskBillingLogParams) { } func GetAllLogs(logType int, logFilter string, startTimestamp int64, endTimestamp int64, modelName string, username string, tokenName string, startIdx int, num int, channel int, group string, requestId string, upstreamRequestId string) (logs []*Log, total int64, err error) { - tx := applyLogFilter(applyLogTypeFilter(LOG_DB, logType), logFilter) + tx, err := applyLogFilter(applyLogTypeFilter(LOG_DB, logType), logFilter) + if err != nil { + return nil, 0, err + } if tx, err = applyExplicitLogTextFilter(tx, "logs.model_name", modelName); err != nil { return nil, 0, err @@ -607,7 +629,10 @@ func GetAllLogs(logType int, logFilter string, startTimestamp int64, endTimestam const logSearchCountLimit = 10000 func GetUserLogs(userId int, logType int, logFilter string, startTimestamp int64, endTimestamp int64, modelName string, tokenName string, startIdx int, num int, group string, requestId string, upstreamRequestId string) (logs []*Log, total int64, err error) { - tx := applyLogFilter(applyLogTypeFilter(LOG_DB.Where("logs.user_id = ?", userId), logType), logFilter) + tx, err := applyLogFilter(applyLogTypeFilter(LOG_DB.Where("logs.user_id = ?", userId), logType), logFilter) + if err != nil { + return nil, 0, err + } if tx, err = applyExplicitLogTextFilter(tx, "logs.model_name", modelName); err != nil { return nil, 0, err @@ -688,8 +713,14 @@ func SumUsedQuota(logType int, logFilter string, startTimestamp int64, endTimest // 为rpm和tpm创建单独的查询 rpmTpmQuery := LOG_DB.Table("logs").Select("count(*) rpm, sum(prompt_tokens) + sum(completion_tokens) tpm") - tx = applyLogFilter(tx, logFilter) - rpmTpmQuery = applyLogFilter(rpmTpmQuery, logFilter) + tx, err = applyLogFilter(tx, logFilter) + if err != nil { + return stat, err + } + rpmTpmQuery, err = applyLogFilter(rpmTpmQuery, logFilter) + if err != nil { + return stat, err + } if tx, err = applyExplicitLogTextFilter(tx, "username", username); err != nil { return stat, err @@ -703,9 +734,11 @@ func SumUsedQuota(logType int, logFilter string, startTimestamp int64, endTimest } if startTimestamp != 0 { tx = tx.Where("created_at >= ?", startTimestamp) + rpmTpmQuery = rpmTpmQuery.Where("created_at >= ?", startTimestamp) } if endTimestamp != 0 { tx = tx.Where("created_at <= ?", endTimestamp) + rpmTpmQuery = rpmTpmQuery.Where("created_at <= ?", endTimestamp) } if tx, err = applyExplicitLogTextFilter(tx, "model_name", modelName); err != nil { return stat, err @@ -722,8 +755,13 @@ func SumUsedQuota(logType int, logFilter string, startTimestamp int64, endTimest rpmTpmQuery = rpmTpmQuery.Where(logGroupCol+" = ?", group) } - tx = tx.Where("type = ?", LogTypeConsume) - rpmTpmQuery = rpmTpmQuery.Where("type = ?", LogTypeConsume) + if logType != LogTypeUnknown { + tx = tx.Where("logs.type = ?", logType) + rpmTpmQuery = rpmTpmQuery.Where("logs.type = ?", logType) + } else if logFilter != LogFilterRetry { + tx = tx.Where("logs.type = ?", LogTypeConsume) + rpmTpmQuery = rpmTpmQuery.Where("logs.type = ?", LogTypeConsume) + } // 只统计最近60秒的rpm和tpm rpmTpmQuery = rpmTpmQuery.Where("created_at >= ?", time.Now().Add(-60*time.Second).Unix()) diff --git a/model/log_test.go b/model/log_test.go index 26e9675..114df37 100644 --- a/model/log_test.go +++ b/model/log_test.go @@ -1,6 +1,7 @@ package model import ( + "errors" "os" "testing" "time" @@ -33,9 +34,9 @@ func TestGetAllLogsRetryFilter(t *testing.T) { got, total, err := GetAllLogs(LogTypeUnknown, LogFilterRetry, 0, 0, "", "", "", 0, 10, 0, "", "", "") require.NoError(t, err) - require.EqualValues(t, 3, total) - require.Len(t, got, 3) - require.ElementsMatch(t, []int{logs[0].Id, logs[1].Id, logs[4].Id}, []int{got[0].Id, got[1].Id, got[2].Id}) + require.EqualValues(t, 4, total) + require.Len(t, got, 4) + require.ElementsMatch(t, []int{logs[0].Id, logs[1].Id, logs[4].Id, logs[5].Id}, []int{got[0].Id, got[1].Id, got[2].Id, got[3].Id}) } func TestGetUserLogsRetryFilter(t *testing.T) { @@ -48,10 +49,10 @@ func TestGetUserLogsRetryFilter(t *testing.T) { got, total, err := GetUserLogs(1, LogTypeUnknown, LogFilterRetry, 0, 0, "", "", 0, 10, "", "", "") require.NoError(t, err) - require.EqualValues(t, 2, total) - require.Len(t, got, 2) - gotLogIds := []int{got[0].LogId, got[1].LogId} - require.ElementsMatch(t, []int{logs[0].Id, logs[1].Id}, gotLogIds) + require.EqualValues(t, 3, total) + require.Len(t, got, 3) + gotLogIds := []int{got[0].LogId, got[1].LogId, got[2].LogId} + require.ElementsMatch(t, []int{logs[0].Id, logs[1].Id, logs[5].Id}, gotLogIds) require.NotContains(t, gotLogIds, logs[4].Id) for _, log := range got { require.Equal(t, 1, log.UserId) @@ -69,10 +70,25 @@ func TestSumUsedQuotaRetryFilter(t *testing.T) { stat, err := SumUsedQuota(LogTypeUnknown, LogFilterRetry, 0, 0, "", "", "", 0, "") require.NoError(t, err) require.Equal(t, 1200, stat.Quota) - require.Equal(t, 3, stat.Rpm) + require.Equal(t, 4, stat.Rpm) require.Equal(t, 71, stat.Tpm) } +func TestSumUsedQuotaAppliesExplicitLogType(t *testing.T) { + require.NoError(t, LOG_DB.Where("1 = 1").Delete(&Log{}).Error) + t.Cleanup(func() { + require.NoError(t, LOG_DB.Where("1 = 1").Delete(&Log{}).Error) + }) + + createRetryFilterLogs(t) + + stat, err := SumUsedQuota(LogTypeError, "", 0, 0, "", "", "", 0, "") + require.NoError(t, err) + require.Equal(t, 0, stat.Quota) + require.Equal(t, 1, stat.Rpm) + require.Equal(t, 0, stat.Tpm) +} + func TestRetryFilterIgnoresNestedRetryMarker(t *testing.T) { require.NoError(t, LOG_DB.Where("1 = 1").Delete(&Log{}).Error) t.Cleanup(func() { @@ -121,6 +137,114 @@ func TestRetryFilterIgnoresNestedRetryMarker(t *testing.T) { require.Equal(t, 16, stat.Tpm) } +func TestRetryFilterBackfillsLegacyMarkersBeforeCompletion(t *testing.T) { + markerKey := logRetryMarkerBackfillCompletionKey() + require.NoError(t, LOG_DB.Where("1 = 1").Delete(&Log{}).Error) + require.NoError(t, DB.Where(commonKeyCol+" = ?", markerKey).Delete(&Option{}).Error) + t.Cleanup(func() { + require.NoError(t, LOG_DB.Where("1 = 1").Delete(&Log{}).Error) + require.NoError(t, DB.Where(commonKeyCol+" = ?", markerKey).Delete(&Option{}).Error) + }) + + logs := []Log{ + { + UserId: 1, + CreatedAt: time.Now().Unix() - 10, + Type: LogTypeConsume, + Quota: 100, + PromptTokens: 3, + CompletionTokens: 5, + Other: common.MapToJsonStr(map[string]interface{}{ + "retry_log": true, + }), + }, + { + UserId: 1, + CreatedAt: time.Now().Unix() - 20, + Type: LogTypeConsume, + Quota: 200, + PromptTokens: 7, + CompletionTokens: 11, + Other: common.MapToJsonStr(map[string]interface{}{ + "admin_info": map[string]interface{}{ + "retry_log": true, + }, + }), + }, + } + require.NoError(t, LOG_DB.Create(&logs).Error) + require.NoError(t, LOG_DB.Model(&Log{}).Where("1 = 1").UpdateColumn("is_retry", false).Error) + + got, total, err := GetAllLogs(LogTypeUnknown, LogFilterRetry, 0, 0, "", "", "", 0, 10, 0, "", "", "") + require.NoError(t, err) + require.EqualValues(t, 1, total) + require.Len(t, got, 1) + require.Equal(t, logs[0].Id, got[0].Id) + + var reloaded []Log + require.NoError(t, LOG_DB.Order("id asc").Find(&reloaded).Error) + require.Len(t, reloaded, 2) + require.True(t, reloaded[0].IsRetry) + require.False(t, reloaded[1].IsRetry) +} + +func TestRetryFilterUsesIsRetryAfterBackfillCompletion(t *testing.T) { + markerKey := logRetryMarkerBackfillCompletionKey() + require.NoError(t, LOG_DB.Where("1 = 1").Delete(&Log{}).Error) + require.NoError(t, DB.Where(commonKeyCol+" = ?", markerKey).Delete(&Option{}).Error) + t.Cleanup(func() { + require.NoError(t, LOG_DB.Where("1 = 1").Delete(&Log{}).Error) + require.NoError(t, DB.Where(commonKeyCol+" = ?", markerKey).Delete(&Option{}).Error) + }) + + log := Log{ + UserId: 1, + CreatedAt: time.Now().Unix() - 10, + Type: LogTypeConsume, + Quota: 100, + Other: common.MapToJsonStr(map[string]interface{}{ + "retry_log": true, + }), + } + require.NoError(t, LOG_DB.Create(&log).Error) + require.NoError(t, LOG_DB.Model(&Log{}).Where("id = ?", log.Id).UpdateColumn("is_retry", false).Error) + require.NoError(t, markLogRetryMarkerBackfillCompleted()) + + got, total, err := GetAllLogs(LogTypeUnknown, LogFilterRetry, 0, 0, "", "", "", 0, 10, 0, "", "", "") + require.NoError(t, err) + require.EqualValues(t, 0, total) + require.Empty(t, got) + + var reloaded Log + require.NoError(t, LOG_DB.First(&reloaded, log.Id).Error) + require.False(t, reloaded.IsRetry) +} + +func TestRetryFilterReadPathsReturnReadinessError(t *testing.T) { + originalEnsure := ensureLogRetryMarkerBackfillCompletedForRead + expectedErr := errors.New("readiness unavailable") + ensureLogRetryMarkerBackfillCompletedForRead = func() error { + return expectedErr + } + t.Cleanup(func() { + ensureLogRetryMarkerBackfillCompletedForRead = originalEnsure + }) + + got, total, err := GetAllLogs(LogTypeUnknown, LogFilterRetry, 0, 0, "", "", "", 0, 10, 0, "", "", "") + require.ErrorIs(t, err, expectedErr) + require.Nil(t, got) + require.Zero(t, total) + + got, total, err = GetUserLogs(1, LogTypeUnknown, LogFilterRetry, 0, 0, "", "", 0, 10, "", "", "") + require.ErrorIs(t, err, expectedErr) + require.Nil(t, got) + require.Zero(t, total) + + stat, err := SumUsedQuota(LogTypeUnknown, LogFilterRetry, 0, 0, "", "", "", 0, "") + require.ErrorIs(t, err, expectedErr) + require.Zero(t, stat) +} + func TestBackfillLogRetryMarkerUsesTopLevelMarkersOnly(t *testing.T) { markerKey := logRetryMarkerBackfillCompletionKey() require.NoError(t, LOG_DB.Where("1 = 1").Delete(&Log{}).Error) @@ -248,6 +372,89 @@ func TestBackfillLogRetryMarkerCompletionIsScopedToLogDBIdentity(t *testing.T) { require.True(t, reloaded.IsRetry) } +func TestMigrateLOGDBSchedulesRetryMarkerBackfillWithoutRunningInline(t *testing.T) { + originalDB := DB + originalLOGDB := LOG_DB + originalRunner := logRetryMarkerBackfillAsyncRunner + originalLogSQLType := common.LogSqlType + t.Cleanup(func() { + DB = originalDB + LOG_DB = originalLOGDB + logRetryMarkerBackfillAsyncRunner = originalRunner + common.LogSqlType = originalLogSQLType + initCol() + }) + + mainDB := newRetryBackfillTestDB(t, &Option{}) + logDB := newRetryBackfillTestDB(t) + + DB = mainDB + LOG_DB = logDB + common.LogSqlType = common.DatabaseTypeSQLite + t.Setenv("LOG_SQL_DSN", "sqlite://async-backfill") + initCol() + + var scheduled []func() + logRetryMarkerBackfillAsyncRunner = func(fn func()) { + scheduled = append(scheduled, fn) + } + + require.NoError(t, migrateLOGDB()) + require.Len(t, scheduled, 1) + + log := Log{ + UserId: 1, + Type: LogTypeConsume, + Other: common.MapToJsonStr(map[string]interface{}{ + "retry_log": true, + }), + } + require.NoError(t, LOG_DB.Create(&log).Error) + require.NoError(t, LOG_DB.Model(&Log{}).Where("1 = 1").UpdateColumn("is_retry", false).Error) + + var reloaded Log + require.NoError(t, LOG_DB.First(&reloaded, log.Id).Error) + require.False(t, reloaded.IsRetry) + + scheduled[0]() + + reloaded = Log{} + require.NoError(t, LOG_DB.First(&reloaded, log.Id).Error) + require.True(t, reloaded.IsRetry) +} + +func TestScheduleLogRetryMarkerBackfillSkipsCompletedMarker(t *testing.T) { + originalDB := DB + originalLOGDB := LOG_DB + originalRunner := logRetryMarkerBackfillAsyncRunner + originalLogSQLType := common.LogSqlType + t.Cleanup(func() { + DB = originalDB + LOG_DB = originalLOGDB + logRetryMarkerBackfillAsyncRunner = originalRunner + common.LogSqlType = originalLogSQLType + initCol() + }) + + mainDB := newRetryBackfillTestDB(t, &Option{}) + logDB := newRetryBackfillTestDB(t, &Log{}) + + DB = mainDB + LOG_DB = logDB + common.LogSqlType = common.DatabaseTypeSQLite + t.Setenv("LOG_SQL_DSN", "sqlite://completed-backfill") + initCol() + require.NoError(t, markLogRetryMarkerBackfillCompleted()) + + called := false + logRetryMarkerBackfillAsyncRunner = func(fn func()) { + called = true + } + + scheduleLogRetryMarkerBackfill() + require.False(t, called) +} + func TestLogRetryMarkerCompletionKeyUsesHashedLogDBIdentity(t *testing.T) { originalLogSQLType := common.LogSqlType t.Cleanup(func() { @@ -385,6 +592,20 @@ func createRetryFilterLogs(t *testing.T) []Log { }, }), }, + { + UserId: 1, + CreatedAt: time.Now().Unix() - 5, + Type: LogTypeError, + Quota: 0, + PromptTokens: 0, + CompletionTokens: 0, + Other: common.MapToJsonStr(map[string]interface{}{ + "retry_log": true, + "admin_info": map[string]interface{}{ + "use_channel": []string{"8"}, + }, + }), + }, } require.NoError(t, LOG_DB.Create(&logs).Error) diff --git a/model/main.go b/model/main.go index c8c5a0c..b432ff6 100644 --- a/model/main.go +++ b/model/main.go @@ -12,6 +12,7 @@ import ( "github.com/MAX-API-Next/MAX-API/common" "github.com/MAX-API-Next/MAX-API/constant" + "github.com/bytedance/gopkg/util/gopool" "github.com/glebarez/sqlite" "gorm.io/driver/mysql" "gorm.io/driver/postgres" @@ -28,6 +29,12 @@ var logGroupCol string const logRetryMarkerBackfillOptionKey = "LogRetryMarkerBackfillCompleted" +var logRetryMarkerBackfillMu sync.Mutex + +var logRetryMarkerBackfillAsyncRunner = func(fn func()) { + gopool.Go(fn) +} + func logRetryMarkerBackfillCompletionKey() string { logDSN := strings.TrimSpace(os.Getenv("LOG_SQL_DSN")) if logDSN == "" { @@ -315,9 +322,7 @@ func migrateDB() error { } if os.Getenv("LOG_SQL_DSN") == "" { LOG_DB = DB - if err := backfillLogRetryMarker(); err != nil { - return err - } + scheduleLogRetryMarkerBackfill() } return nil } @@ -399,13 +404,34 @@ func migrateLOGDB() error { if err = LOG_DB.AutoMigrate(&Log{}); err != nil { return err } - if err = backfillLogRetryMarker(); err != nil { - return err - } + scheduleLogRetryMarkerBackfill() return nil } +func scheduleLogRetryMarkerBackfill() { + if LOG_DB == nil || !LOG_DB.Migrator().HasColumn(&Log{}, "is_retry") { + return + } + if completed, err := isLogRetryMarkerBackfillCompleted(); err == nil && completed { + return + } else if err != nil { + common.SysLog("failed to check log retry marker backfill status: " + err.Error()) + } + + logRetryMarkerBackfillAsyncRunner(func() { + common.SysLog("log retry marker backfill started") + if err := backfillLogRetryMarker(); err != nil { + common.SysLog("log retry marker backfill failed: " + err.Error()) + return + } + common.SysLog("log retry marker backfill finished") + }) +} + func backfillLogRetryMarker() error { + logRetryMarkerBackfillMu.Lock() + defer logRetryMarkerBackfillMu.Unlock() + if LOG_DB == nil || !LOG_DB.Migrator().HasColumn(&Log{}, "is_retry") { return nil } diff --git a/relay/channel/gemini/relay_gemini_usage_test.go b/relay/channel/gemini/relay_gemini_usage_test.go index 13e479d..2c95a3d 100644 --- a/relay/channel/gemini/relay_gemini_usage_test.go +++ b/relay/channel/gemini/relay_gemini_usage_test.go @@ -8,7 +8,6 @@ import ( "testing" "github.com/MAX-API-Next/MAX-API/common" - "github.com/MAX-API-Next/MAX-API/constant" "github.com/MAX-API-Next/MAX-API/dto" relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" "github.com/MAX-API-Next/MAX-API/types" @@ -72,12 +71,6 @@ func TestGeminiStreamHandlerCompletionTokensExcludeToolUsePromptTokens(t *testin c, _ := gin.CreateTestContext(httptest.NewRecorder()) c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) - oldStreamingTimeout := constant.StreamingTimeout - constant.StreamingTimeout = 300 - t.Cleanup(func() { - constant.StreamingTimeout = oldStreamingTimeout - }) - info := &relaycommon.RelayInfo{ OriginModelName: "gemini-3-flash-preview", ChannelMeta: &relaycommon.ChannelMeta{ @@ -230,12 +223,6 @@ func TestGeminiStreamHandlerUsesEstimatedPromptTokensWhenUsagePromptMissing(t *t c, _ := gin.CreateTestContext(httptest.NewRecorder()) c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) - oldStreamingTimeout := constant.StreamingTimeout - constant.StreamingTimeout = 300 - t.Cleanup(func() { - constant.StreamingTimeout = oldStreamingTimeout - }) - info := &relaycommon.RelayInfo{ OriginModelName: "gemini-3-flash-preview", ChannelMeta: &relaycommon.ChannelMeta{ diff --git a/relay/channel/openai/empty_completion_test.go b/relay/channel/openai/empty_completion_test.go index 2230a3e..5baf78b 100644 --- a/relay/channel/openai/empty_completion_test.go +++ b/relay/channel/openai/empty_completion_test.go @@ -8,7 +8,6 @@ import ( "testing" "github.com/MAX-API-Next/MAX-API/common" - "github.com/MAX-API-Next/MAX-API/constant" "github.com/MAX-API-Next/MAX-API/dto" relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" "github.com/MAX-API-Next/MAX-API/types" @@ -143,14 +142,11 @@ func TestOaiResponsesStreamHandlerRetriesEmptyBeforeWriting(t *testing.T) { oldRetryTimes := common.RetryTimes oldEmptyCompletionRetryEnabled := common.EmptyCompletionRetryEnabled - oldStreamingTimeout := constant.StreamingTimeout common.RetryTimes = 1 common.EmptyCompletionRetryEnabled = true - constant.StreamingTimeout = 300 defer func() { common.RetryTimes = oldRetryTimes common.EmptyCompletionRetryEnabled = oldEmptyCompletionRetryEnabled - constant.StreamingTimeout = oldStreamingTimeout }() resp := &http.Response{ @@ -194,14 +190,11 @@ func TestOaiResponsesStreamHandlerRetriesEmptyOnEOFBeforeWriting(t *testing.T) { oldRetryTimes := common.RetryTimes oldEmptyCompletionRetryEnabled := common.EmptyCompletionRetryEnabled - oldStreamingTimeout := constant.StreamingTimeout common.RetryTimes = 1 common.EmptyCompletionRetryEnabled = true - constant.StreamingTimeout = 300 defer func() { common.RetryTimes = oldRetryTimes common.EmptyCompletionRetryEnabled = oldEmptyCompletionRetryEnabled - constant.StreamingTimeout = oldStreamingTimeout }() resp := &http.Response{ @@ -243,14 +236,11 @@ func TestOaiResponsesStreamHandlerFlushesImmediatelyWhenEmptyRetryDisabled(t *te oldRetryTimes := common.RetryTimes oldEmptyCompletionRetryEnabled := common.EmptyCompletionRetryEnabled - oldStreamingTimeout := constant.StreamingTimeout common.RetryTimes = 1 common.EmptyCompletionRetryEnabled = false - constant.StreamingTimeout = 300 defer func() { common.RetryTimes = oldRetryTimes common.EmptyCompletionRetryEnabled = oldEmptyCompletionRetryEnabled - constant.StreamingTimeout = oldStreamingTimeout }() resp := &http.Response{ @@ -285,14 +275,11 @@ func TestOaiResponsesStreamHandlerCapsPendingEvents(t *testing.T) { oldRetryTimes := common.RetryTimes oldEmptyCompletionRetryEnabled := common.EmptyCompletionRetryEnabled - oldStreamingTimeout := constant.StreamingTimeout common.RetryTimes = 1 common.EmptyCompletionRetryEnabled = true - constant.StreamingTimeout = 300 defer func() { common.RetryTimes = oldRetryTimes common.EmptyCompletionRetryEnabled = oldEmptyCompletionRetryEnabled - constant.StreamingTimeout = oldStreamingTimeout }() events := make([]string, 0, maxPendingResponsesStreamEvents+2) @@ -329,14 +316,11 @@ func TestOaiResponsesToChatStreamHandlerRetriesEmptyOnEOFBeforeWriting(t *testin oldRetryTimes := common.RetryTimes oldEmptyCompletionRetryEnabled := common.EmptyCompletionRetryEnabled - oldStreamingTimeout := constant.StreamingTimeout common.RetryTimes = 1 common.EmptyCompletionRetryEnabled = true - constant.StreamingTimeout = 300 defer func() { common.RetryTimes = oldRetryTimes common.EmptyCompletionRetryEnabled = oldEmptyCompletionRetryEnabled - constant.StreamingTimeout = oldStreamingTimeout }() resp := &http.Response{ diff --git a/relay/channel/openai/image_stream_test.go b/relay/channel/openai/image_stream_test.go index 2a034a0..f009878 100644 --- a/relay/channel/openai/image_stream_test.go +++ b/relay/channel/openai/image_stream_test.go @@ -7,7 +7,6 @@ import ( "strings" "testing" - "github.com/MAX-API-Next/MAX-API/constant" relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" "github.com/gin-gonic/gin" "github.com/stretchr/testify/require" @@ -37,10 +36,6 @@ func TestOpenaiImageStreamHandlerForwardsSSEAndUsage(t *testing.T) { gin.SetMode(gin.TestMode) t.Cleanup(func() { gin.SetMode(oldMode) }) - oldTimeout := constant.StreamingTimeout - constant.StreamingTimeout = 30 - t.Cleanup(func() { constant.StreamingTimeout = oldTimeout }) - body := strings.Join([]string{ `event: image_generation.partial_image`, `data: {"type":"image_generation.partial_image","b64_json":"partial"}`, @@ -130,10 +125,6 @@ func TestOpenaiImageStreamHandlerRecordsUpstreamErrorEvent(t *testing.T) { gin.SetMode(gin.TestMode) t.Cleanup(func() { gin.SetMode(oldMode) }) - oldTimeout := constant.StreamingTimeout - constant.StreamingTimeout = 30 - t.Cleanup(func() { constant.StreamingTimeout = oldTimeout }) - body := strings.Join([]string{ `event: image_generation.partial_image`, `data: {"type":"image_generation.partial_image","b64_json":"partial"}`, diff --git a/relay/helper/stream_scanner.go b/relay/helper/stream_scanner.go index c37af73..3c79eac 100644 --- a/relay/helper/stream_scanner.go +++ b/relay/helper/stream_scanner.go @@ -40,6 +40,14 @@ func NewStreamScanner(reader io.Reader) *bufio.Scanner { return scanner } +func newStreamingTimeoutTicker(timeoutSeconds int) (*time.Ticker, <-chan time.Time) { + if timeoutSeconds <= 0 { + return nil, nil + } + ticker := time.NewTicker(time.Duration(timeoutSeconds) * time.Second) + return ticker, ticker.C +} + func StreamScannerHandler(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo, dataHandler func(data string, sr *StreamResult)) { if resp == nil || dataHandler == nil { @@ -58,12 +66,13 @@ func StreamScannerHandler(c *gin.Context, resp *http.Response, info *relaycommon } }() - streamingTimeout := time.Duration(constant.StreamingTimeout) * time.Second + streamingTimeoutSeconds := constant.StreamingTimeout + streamingTimeout := time.Duration(streamingTimeoutSeconds) * time.Second + ticker, timeoutChan := newStreamingTimeoutTicker(streamingTimeoutSeconds) var ( stopChan = make(chan bool, 3) // 增加缓冲区避免阻塞 scanner = NewStreamScanner(resp.Body) - ticker = time.NewTicker(streamingTimeout) pingTicker *time.Ticker writeMutex sync.Mutex // Mutex to protect concurrent writes wg sync.WaitGroup // 用于等待所有 goroutine 退出 @@ -91,7 +100,9 @@ func StreamScannerHandler(c *gin.Context, resp *http.Response, info *relaycommon // 通知所有 goroutine 停止 common.SafeSendBool(stopChan, true) - ticker.Stop() + if ticker != nil { + ticker.Stop() + } if pingTicker != nil { pingTicker.Stop() } @@ -233,7 +244,9 @@ func StreamScannerHandler(c *gin.Context, resp *http.Response, info *relaycommon default: } - ticker.Reset(streamingTimeout) + if ticker != nil { + ticker.Reset(streamingTimeout) + } data := scanner.Text() logger.LogDebug(c, "stream scanner data: %s", data) @@ -277,7 +290,7 @@ func StreamScannerHandler(c *gin.Context, resp *http.Response, info *relaycommon // 主循环等待完成或超时 select { - case <-ticker.C: + case <-timeoutChan: info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonTimeout, nil) case <-stopChan: // EndReason already set by the goroutine that triggered stopChan diff --git a/relay/helper/stream_scanner_test.go b/relay/helper/stream_scanner_test.go index d49cc5e..1ecd645 100644 --- a/relay/helper/stream_scanner_test.go +++ b/relay/helper/stream_scanner_test.go @@ -27,12 +27,6 @@ func init() { func setupStreamTest(t *testing.T, body io.Reader) (*gin.Context, *http.Response, *relaycommon.RelayInfo) { t.Helper() - oldTimeout := constant.StreamingTimeout - constant.StreamingTimeout = 30 - t.Cleanup(func() { - constant.StreamingTimeout = oldTimeout - }) - recorder := httptest.NewRecorder() c, _ := gin.CreateTestContext(recorder) c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) @@ -98,6 +92,13 @@ func TestNewStreamScanner_AllowsLargeStreamLine(t *testing.T) { require.NoError(t, scanner.Err()) } +func TestNewStreamingTimeoutTicker_DisabledWhenZero(t *testing.T) { + ticker, timeoutChan := newStreamingTimeoutTicker(0) + + assert.Nil(t, ticker) + assert.Nil(t, timeoutChan) +} + func TestStreamScannerHandler_EmptyBody(t *testing.T) { t.Parallel() @@ -269,10 +270,6 @@ func TestStreamScannerHandler_ScannerDecoupledFromSlowHandler(t *testing.T) { c, _ := gin.CreateTestContext(recorder) c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) - oldTimeout := constant.StreamingTimeout - constant.StreamingTimeout = 30 - t.Cleanup(func() { constant.StreamingTimeout = oldTimeout }) - resp := &http.Response{Body: pr} info := &relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{}} @@ -366,12 +363,6 @@ func TestStreamScannerHandler_PingSentDuringSlowUpstream(t *testing.T) { c, _ := gin.CreateTestContext(recorder) c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) - oldTimeout := constant.StreamingTimeout - constant.StreamingTimeout = 30 - t.Cleanup(func() { - constant.StreamingTimeout = oldTimeout - }) - resp := &http.Response{Body: pr} info := &relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{}} @@ -426,12 +417,6 @@ func TestStreamScannerHandler_PingDisabledByRelayInfo(t *testing.T) { c, _ := gin.CreateTestContext(recorder) c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) - oldTimeout := constant.StreamingTimeout - constant.StreamingTimeout = 30 - t.Cleanup(func() { - constant.StreamingTimeout = oldTimeout - }) - resp := &http.Response{Body: pr} info := &relaycommon.RelayInfo{ DisablePing: true, @@ -638,6 +623,36 @@ func TestStreamScannerHandler_StreamStatus_InitializedIfNil(t *testing.T) { assert.NotNil(t, info.StreamStatus) } +func TestStreamScannerHandler_AllowsDisabledStreamingTimeout(t *testing.T) { + oldTimeout := constant.StreamingTimeout + constant.StreamingTimeout = 0 + t.Cleanup(func() { + constant.StreamingTimeout = oldTimeout + }) + + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) + + resp := &http.Response{ + Body: io.NopCloser(strings.NewReader(buildSSEBody(1))), + } + info := &relaycommon.RelayInfo{ + ChannelMeta: &relaycommon.ChannelMeta{}, + } + + var count atomic.Int64 + require.NotPanics(t, func() { + StreamScannerHandler(c, resp, info, func(data string, sr *StreamResult) { + count.Add(1) + }) + }) + + assert.Equal(t, int64(1), count.Load()) + require.NotNil(t, info.StreamStatus) + assert.Equal(t, relaycommon.StreamEndReasonDone, info.StreamStatus.EndReason) +} + func TestStreamScannerHandler_StreamStatus_PreInitialized(t *testing.T) { t.Parallel() @@ -680,12 +695,6 @@ func TestStreamScannerHandler_PingInterleavesWithSlowUpstream(t *testing.T) { c, _ := gin.CreateTestContext(recorder) c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) - oldTimeout := constant.StreamingTimeout - constant.StreamingTimeout = 30 - t.Cleanup(func() { - constant.StreamingTimeout = oldTimeout - }) - resp := &http.Response{Body: pr} info := &relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{}} diff --git a/service/quota.go b/service/quota.go index 09fe4f6..0e7cb2e 100644 --- a/service/quota.go +++ b/service/quota.go @@ -214,6 +214,7 @@ func PostWssConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, mod } // record all the consume log even if quota is 0 + shouldUpdateUsageStats := totalTokens != 0 if totalTokens == 0 { // in this case, must be some error happened // we cannot just return, because we may have to return the pre-consumed quota @@ -221,7 +222,14 @@ func PostWssConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, mod logContent += "（可能是上游超时）" logger.LogError(ctx, fmt.Sprintf("total tokens is 0, cannot consume quota, userId %d, channelId %d, "+ "tokenId %d, model %s， pre-consumed quota %d", relayInfo.UserId, relayInfo.ChannelId, relayInfo.TokenId, modelName, relayInfo.FinalPreConsumedQuota)) - } else { + } + if fallbackQuota, ok := streamFallbackQuota(relayInfo, quota); ok { + quota = fallbackQuota + shouldUpdateUsageStats = true + logContent += fmt.Sprintf(", stream ended abnormally (%s), billed pre-consumed quota %d", relayInfo.StreamStatus.EndReason, fallbackQuota) + logger.LogError(ctx, fmt.Sprintf("stream ended abnormally without billable usage, billing fallback quota, userId %d, channelId %d, tokenId %d, model %s, reason %s, fallback quota %d", relayInfo.UserId, relayInfo.ChannelId, relayInfo.TokenId, modelName, relayInfo.StreamStatus.EndReason, fallbackQuota)) + } + if shouldUpdateUsageStats { model.UpdateUserUsedQuotaAndRequestCount(relayInfo.UserId, quota) model.UpdateChannelUsedQuota(relayInfo.ChannelId, quota) } @@ -335,6 +343,7 @@ func PostAudioConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, u } // record all the consume log even if quota is 0 + shouldUpdateUsageStats := totalTokens != 0 if totalTokens == 0 { // in this case, must be some error happened // we cannot just return, because we may have to return the pre-consumed quota @@ -342,7 +351,14 @@ func PostAudioConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, u logContent += "（可能是上游超时）" logger.LogError(ctx, fmt.Sprintf("total tokens is 0, cannot consume quota, userId %d, channelId %d, "+ "tokenId %d, model %s， pre-consumed quota %d", relayInfo.UserId, relayInfo.ChannelId, relayInfo.TokenId, relayInfo.OriginModelName, relayInfo.FinalPreConsumedQuota)) - } else { + } + if fallbackQuota, ok := streamFallbackQuota(relayInfo, quota); ok { + quota = fallbackQuota + shouldUpdateUsageStats = true + logContent += fmt.Sprintf(", stream ended abnormally (%s), billed pre-consumed quota %d", relayInfo.StreamStatus.EndReason, fallbackQuota) + logger.LogError(ctx, fmt.Sprintf("stream ended abnormally without billable usage, billing fallback quota, userId %d, channelId %d, tokenId %d, model %s, reason %s, fallback quota %d", relayInfo.UserId, relayInfo.ChannelId, relayInfo.TokenId, relayInfo.OriginModelName, relayInfo.StreamStatus.EndReason, fallbackQuota)) + } + if shouldUpdateUsageStats { model.UpdateUserUsedQuotaAndRequestCount(relayInfo.UserId, quota) model.UpdateChannelUsedQuota(relayInfo.ChannelId, quota) } diff --git a/service/text_quota.go b/service/text_quota.go index 7af2781..273bb2f 100644 --- a/service/text_quota.go +++ b/service/text_quota.go @@ -319,6 +319,29 @@ func usageSemanticFromUsage(relayInfo *relaycommon.RelayInfo, usage *dto.Usage) return "openai" } +func streamFallbackQuota(relayInfo *relaycommon.RelayInfo, quota int) (int, bool) { + if quota != 0 || relayInfo == nil || !relayInfo.IsStream || relayInfo.StreamStatus == nil { + return quota, false + } + switch relayInfo.StreamStatus.EndReason { + case relaycommon.StreamEndReasonTimeout, + relaycommon.StreamEndReasonClientGone, + relaycommon.StreamEndReasonScannerErr, + relaycommon.StreamEndReasonPanic, + relaycommon.StreamEndReasonPingFail: + default: + return quota, false + } + fallbackQuota := relayInfo.FinalPreConsumedQuota + if fallbackQuota == 0 { + fallbackQuota = relayInfo.PriceData.QuotaToPreConsume + } + if fallbackQuota <= 0 { + return quota, false + } + return fallbackQuota, true +} + func PostTextConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usage *dto.Usage, extraContent []string) { originUsage := usage if usage == nil { @@ -362,10 +385,22 @@ func PostTextConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, us extraContent = append(extraContent, fmt.Sprintf("Image Generation Call 花费 %s", decimal.NewFromFloat(summary.ImageGenerationCallPrice).Mul(decimal.NewFromFloat(summary.GroupRatio)).Mul(decimal.NewFromFloat(common.QuotaPerUnit)).String())) } + shouldUpdateUsageStats := summary.TotalTokens != 0 if summary.TotalTokens == 0 { extraContent = append(extraContent, "上游没有返回计费信息，无法扣费（可能是上游超时）") logger.LogError(ctx, fmt.Sprintf("total tokens is 0, cannot consume quota, userId %d, channelId %d, tokenId %d, model %s， pre-consumed quota %d", relayInfo.UserId, relayInfo.ChannelId, relayInfo.TokenId, summary.ModelName, relayInfo.FinalPreConsumedQuota)) - } else { + } + fallbackBaseQuota := summary.Quota + if originUsage == nil { + fallbackBaseQuota = 0 + } + if fallbackQuota, ok := streamFallbackQuota(relayInfo, fallbackBaseQuota); ok { + summary.Quota = fallbackQuota + shouldUpdateUsageStats = true + extraContent = append(extraContent, fmt.Sprintf("stream ended abnormally (%s), billed pre-consumed quota %d", relayInfo.StreamStatus.EndReason, fallbackQuota)) + logger.LogError(ctx, fmt.Sprintf("stream ended abnormally without billable usage, billing fallback quota, userId %d, channelId %d, tokenId %d, model %s, reason %s, fallback quota %d", relayInfo.UserId, relayInfo.ChannelId, relayInfo.TokenId, summary.ModelName, relayInfo.StreamStatus.EndReason, fallbackQuota)) + } + if shouldUpdateUsageStats { model.UpdateUserUsedQuotaAndRequestCount(relayInfo.UserId, summary.Quota) model.UpdateChannelUsedQuota(relayInfo.ChannelId, summary.Quota) } diff --git a/service/text_quota_test.go b/service/text_quota_test.go index 2d4becd..d1c3964 100644 --- a/service/text_quota_test.go +++ b/service/text_quota_test.go @@ -7,6 +7,7 @@ import ( "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/pkg/billingexpr" relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" "github.com/MAX-API-Next/MAX-API/types" @@ -68,6 +69,166 @@ func TestCalculateTextQuotaSummaryUnifiedForClaudeSemantic(t *testing.T) { require.Equal(t, 1488, chatSummary.Quota) } +func TestStreamFallbackQuotaUsesPreConsumedQuotaOnAbnormalStreamEnd(t *testing.T) { + relayInfo := &relaycommon.RelayInfo{ + IsStream: true, + FinalPreConsumedQuota: 300, + PriceData: types.PriceData{ + QuotaToPreConsume: 500, + }, + StreamStatus: relaycommon.NewStreamStatus(), + } + relayInfo.StreamStatus.SetEndReason(relaycommon.StreamEndReasonTimeout, nil) + + quota, ok := streamFallbackQuota(relayInfo, 0) + + require.True(t, ok) + require.Equal(t, 300, quota) +} + +func TestStreamFallbackQuotaFallsBackToPriceDataPreConsume(t *testing.T) { + relayInfo := &relaycommon.RelayInfo{ + IsStream: true, + PriceData: types.PriceData{ + QuotaToPreConsume: 500, + }, + StreamStatus: relaycommon.NewStreamStatus(), + } + relayInfo.StreamStatus.SetEndReason(relaycommon.StreamEndReasonClientGone, nil) + + quota, ok := streamFallbackQuota(relayInfo, 0) + + require.True(t, ok) + require.Equal(t, 500, quota) +} + +func TestStreamFallbackQuotaSkipsNormalEndAndExistingQuota(t *testing.T) { + relayInfo := &relaycommon.RelayInfo{ + IsStream: true, + FinalPreConsumedQuota: 300, + PriceData: types.PriceData{ + QuotaToPreConsume: 500, + }, + StreamStatus: relaycommon.NewStreamStatus(), + } + relayInfo.StreamStatus.SetEndReason(relaycommon.StreamEndReasonDone, nil) + + quota, ok := streamFallbackQuota(relayInfo, 0) + require.False(t, ok) + require.Equal(t, 0, quota) + + relayInfo.StreamStatus = relaycommon.NewStreamStatus() + relayInfo.StreamStatus.SetEndReason(relaycommon.StreamEndReasonTimeout, nil) + quota, ok = streamFallbackQuota(relayInfo, 200) + require.False(t, ok) + require.Equal(t, 200, quota) +} + +func TestPostTextConsumeQuotaUpdatesUsageStatsForStreamFallback(t *testing.T) { + truncate(t) + gin.SetMode(gin.TestMode) + w := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(w) + + const userID = 101 + const channelID = 201 + const fallbackQuota = 300 + + seedUser(t, userID, 10000) + seedChannel(t, channelID) + + relayInfo := &relaycommon.RelayInfo{ + UserId: userID, + OriginModelName: "test-model", + StartTime: time.Now(), + IsStream: true, + IsPlayground: true, + FinalPreConsumedQuota: fallbackQuota, + StreamStatus: relaycommon.NewStreamStatus(), + UsingGroup: "default", + FinalRequestRelayFormat: types.RelayFormatOpenAI, + ChannelMeta: &relaycommon.ChannelMeta{ + ChannelId: channelID, + }, + PriceData: types.PriceData{ + ModelRatio: 1, + GroupRatioInfo: types.GroupRatioInfo{ + GroupRatio: 1, + }, + }, + } + relayInfo.StreamStatus.SetEndReason(relaycommon.StreamEndReasonTimeout, nil) + + PostTextConsumeQuota(ctx, relayInfo, &dto.Usage{}, nil) + + var user model.User + require.NoError(t, model.DB.Select("used_quota", "request_count").Where("id = ?", userID).First(&user).Error) + require.Equal(t, fallbackQuota, user.UsedQuota) + require.Equal(t, 1, user.RequestCount) + + var channel model.Channel + require.NoError(t, model.DB.Select("used_quota").Where("id = ?", channelID).First(&channel).Error) + require.Equal(t, int64(fallbackQuota), channel.UsedQuota) + + log := getLastLog(t) + require.NotNil(t, log) + require.Equal(t, model.LogTypeConsume, log.Type) + require.Equal(t, fallbackQuota, log.Quota) +} + +func TestPostTextConsumeQuotaUsesFallbackForNilUsageWithEstimate(t *testing.T) { + truncate(t) + gin.SetMode(gin.TestMode) + w := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(w) + + const userID = 102 + const channelID = 202 + const fallbackQuota = 300 + + seedUser(t, userID, 10000) + seedChannel(t, channelID) + + relayInfo := &relaycommon.RelayInfo{ + UserId: userID, + OriginModelName: "test-model", + StartTime: time.Now(), + IsStream: true, + IsPlayground: true, + FinalPreConsumedQuota: fallbackQuota, + StreamStatus: relaycommon.NewStreamStatus(), + UsingGroup: "default", + FinalRequestRelayFormat: types.RelayFormatOpenAI, + ChannelMeta: &relaycommon.ChannelMeta{ + ChannelId: channelID, + }, + PriceData: types.PriceData{ + ModelRatio: 2, + GroupRatioInfo: types.GroupRatioInfo{ + GroupRatio: 1, + }, + }, + } + relayInfo.SetEstimatePromptTokens(20) + relayInfo.StreamStatus.SetEndReason(relaycommon.StreamEndReasonTimeout, nil) + + PostTextConsumeQuota(ctx, relayInfo, nil, nil) + + var user model.User + require.NoError(t, model.DB.Select("used_quota", "request_count").Where("id = ?", userID).First(&user).Error) + require.Equal(t, fallbackQuota, user.UsedQuota) + require.Equal(t, 1, user.RequestCount) + + var channel model.Channel + require.NoError(t, model.DB.Select("used_quota").Where("id = ?", channelID).First(&channel).Error) + require.Equal(t, int64(fallbackQuota), channel.UsedQuota) + + log := getLastLog(t) + require.NotNil(t, log) + require.Equal(t, fallbackQuota, log.Quota) + require.Equal(t, 20, log.PromptTokens) +} + func TestCalculateTextQuotaSummaryUsesSplitClaudeCacheCreationRatios(t *testing.T) { gin.SetMode(gin.TestMode) w := httptest.NewRecorder()