From fedd394e79eb75835413121b85538c2733930194 Mon Sep 17 00:00:00 2001 From: fengsilin Date: Mon, 20 Jul 2026 10:22:45 +0800 Subject: [PATCH] fix(task-billing): persist settlement and merge paged logs Co-Authored-By: Codex --- model/log.go | 73 ++++++++++++++++++++++++++++++++++++ model/log_identity_test.go | 50 ++++++++++++++++++++++++ service/task_billing.go | 3 ++ service/task_billing_test.go | 4 ++ 4 files changed, 130 insertions(+) diff --git a/model/log.go b/model/log.go index 3dd6d1b..9eb7368 100644 --- a/model/log.go +++ b/model/log.go @@ -4,6 +4,7 @@ import ( "context" "errors" "fmt" + "strings" "time" "github.com/QuantumNous/new-api/common" @@ -42,6 +43,69 @@ type Log struct { Other string `json:"other"` } +type taskBillingLogOther struct { + BillingMode string `json:"billing_mode"` + BillingPhase string `json:"billing_phase"` + TaskID string `json:"task_id"` +} + +// enrichTaskBillingLogs adds terminal billing records for Matrix preconsume logs +// that landed on a different raw log page. +func enrichTaskBillingLogs(logs []*Log) ([]*Log, error) { + taskIDs := make(map[string]struct{}) + userIDs := make(map[int]struct{}) + existingIDs := make(map[int]struct{}, len(logs)) + for _, log := range logs { + existingIDs[log.Id] = struct{}{} + var other taskBillingLogOther + if err := common.UnmarshalJsonStr(log.Other, &other); err != nil { + continue + } + if other.BillingMode == "matrix" && other.BillingPhase == "preconsume" && other.TaskID != "" { + taskIDs[other.TaskID] = struct{}{} + userIDs[log.UserId] = struct{}{} + } + } + if len(taskIDs) == 0 { + return logs, nil + } + + patterns := make([]string, 0, len(taskIDs)) + for taskID := range taskIDs { + patterns = append(patterns, "%\"task_id\":\""+taskID+"\"%") + } + users := make([]int, 0, len(userIDs)) + for userID := range userIDs { + users = append(users, userID) + } + conditions := make([]string, len(patterns)) + args := make([]any, len(patterns)) + for i, pattern := range patterns { + conditions[i] = "other LIKE ?" + args[i] = pattern + } + tx := LOG_DB.Where("type IN ? AND user_id IN ?", []int{LogTypeConsume, LogTypeRefund}, users). + Where("("+strings.Join(conditions, " OR ")+")", args...) + var candidates []*Log + if err := tx.Find(&candidates).Error; err != nil { + return nil, err + } + for _, log := range candidates { + if _, exists := existingIDs[log.Id]; exists { + continue + } + var other taskBillingLogOther + if err := common.UnmarshalJsonStr(log.Other, &other); err != nil { + continue + } + if _, wanted := taskIDs[other.TaskID]; wanted && other.BillingMode == "matrix" && + (other.BillingPhase == "settlement" || other.BillingPhase == "refund") { + logs = append(logs, log) + } + } + return logs, nil +} + // don't use iota, avoid change log type value const ( LogTypeUnknown = 0 @@ -301,6 +365,10 @@ func GetAllLogs(logType int, startTimestamp int64, endTimestamp int64, modelName if err != nil { return nil, 0, err } + logs, err = enrichTaskBillingLogs(logs) + if err != nil { + return nil, 0, err + } channelIds := types.NewSet[int]() for _, log := range logs { @@ -394,6 +462,11 @@ func GetUserLogs(userId int, logType int, startTimestamp int64, endTimestamp int return nil, 0, errors.New("查询日志失败") } + logs, err = enrichTaskBillingLogs(logs) + if err != nil { + common.SysError("failed to enrich task billing logs: " + err.Error()) + return nil, 0, err + } formatUserLogs(logs, startIdx) return logs, total, err } diff --git a/model/log_identity_test.go b/model/log_identity_test.go index afccdef..048ed59 100644 --- a/model/log_identity_test.go +++ b/model/log_identity_test.go @@ -152,6 +152,56 @@ func TestGetUserLogsFiltersByChatIDAndUpstreamID(t *testing.T) { require.Equal(t, "up-owner-a", logs[0].UpstreamId) } +func TestGetUserLogsIncludesTaskSettlementOutsidePage(t *testing.T) { + db := setupLogIdentityDB(t) + preconsume := &Log{ + UserId: 1, CreatedAt: 100, Type: LogTypeConsume, ModelName: "seedance", + Other: `{"billing_mode":"matrix","billing_phase":"preconsume","task_id":"task_cross_page"}`, + } + settlement := &Log{ + UserId: 1, CreatedAt: 200, Type: LogTypeConsume, ModelName: "seedance", + Other: `{"billing_mode":"matrix","billing_phase":"settlement","task_id":"task_cross_page"}`, + } + require.NoError(t, db.Create(preconsume).Error) + require.NoError(t, db.Create(settlement).Error) + + logs, total, err := GetUserLogs(1, LogTypeConsume, 0, 0, "", "", 1, 1, "", "", "", "") + require.NoError(t, err) + require.Equal(t, int64(2), total) + require.Len(t, logs, 2) + require.Equal(t, "preconsume", logBillingPhase(t, logs[0])) + require.Equal(t, "settlement", logBillingPhase(t, logs[1])) +} + +func TestGetAllLogsIncludesTaskSettlementOutsidePage(t *testing.T) { + db := setupLogIdentityDB(t) + preconsume := &Log{ + UserId: 1, CreatedAt: 100, Type: LogTypeConsume, ModelName: "seedance", + Other: `{"billing_mode":"matrix","billing_phase":"preconsume","task_id":"task_cross_page_admin"}`, + } + settlement := &Log{ + UserId: 1, CreatedAt: 200, Type: LogTypeConsume, ModelName: "seedance", + Other: `{"billing_mode":"matrix","billing_phase":"settlement","task_id":"task_cross_page_admin"}`, + } + require.NoError(t, db.Create(preconsume).Error) + require.NoError(t, db.Create(settlement).Error) + + logs, total, err := GetAllLogs(LogTypeConsume, 0, 0, "", "", "", 1, 1, 0, "", "", "", "") + require.NoError(t, err) + require.Equal(t, int64(2), total) + require.Len(t, logs, 2) + require.Equal(t, "preconsume", logBillingPhase(t, logs[0])) + require.Equal(t, "settlement", logBillingPhase(t, logs[1])) +} + +func logBillingPhase(t *testing.T, log *Log) string { + t.Helper() + other := map[string]any{} + require.NoError(t, common.UnmarshalJsonStr(log.Other, &other)) + phase, _ := other["billing_phase"].(string) + return phase +} + func TestRecordConsumeLogPersistsContextChatIDAndUpstreamID(t *testing.T) { db := setupLogIdentityDB(t) diff --git a/service/task_billing.go b/service/task_billing.go index 18fc066..54b1050 100644 --- a/service/task_billing.go +++ b/service/task_billing.go @@ -274,6 +274,9 @@ func RecalculateTaskQuota(ctx context.Context, task *model.Task, actualQuota int taskAdjustTokenQuota(ctx, task, quotaDelta) task.Quota = actualQuota + if err := task.Update(); err != nil { + logger.LogError(ctx, fmt.Sprintf("更新任务 quota 失败 task %s: %s", task.TaskID, err.Error())) + } var logType int var logQuota int diff --git a/service/task_billing_test.go b/service/task_billing_test.go index 98a3376..ee702eb 100644 --- a/service/task_billing_test.go +++ b/service/task_billing_test.go @@ -380,6 +380,7 @@ func TestRecalculate_PositiveDelta(t *testing.T) { seedChannel(t, channelID) task := makeTask(userID, channelID, preConsumed, tokenID, BillingSourceWallet, 0) + require.NoError(t, model.DB.Create(task).Error) RecalculateTaskQuota(ctx, task, actualQuota, "adaptor adjustment") @@ -391,6 +392,9 @@ func TestRecalculate_PositiveDelta(t *testing.T) { // task.Quota should be updated to actualQuota assert.Equal(t, actualQuota, task.Quota) + var reloaded model.Task + require.NoError(t, model.DB.First(&reloaded, task.ID).Error) + assert.Equal(t, actualQuota, reloaded.Quota) // Log type should be Consume (additional charge) log := getLastLog(t)