Co-Authored-By: Codex <noreply@anthropic.com>master
| @@ -4,6 +4,7 @@ import ( | |||||
| "context" | "context" | ||||
| "errors" | "errors" | ||||
| "fmt" | "fmt" | ||||
| "strings" | |||||
| "time" | "time" | ||||
| "github.com/QuantumNous/new-api/common" | "github.com/QuantumNous/new-api/common" | ||||
| @@ -42,6 +43,69 @@ type Log struct { | |||||
| Other string `json:"other"` | 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 | // don't use iota, avoid change log type value | ||||
| const ( | const ( | ||||
| LogTypeUnknown = 0 | LogTypeUnknown = 0 | ||||
| @@ -301,6 +365,10 @@ func GetAllLogs(logType int, startTimestamp int64, endTimestamp int64, modelName | |||||
| if err != nil { | if err != nil { | ||||
| return nil, 0, err | return nil, 0, err | ||||
| } | } | ||||
| logs, err = enrichTaskBillingLogs(logs) | |||||
| if err != nil { | |||||
| return nil, 0, err | |||||
| } | |||||
| channelIds := types.NewSet[int]() | channelIds := types.NewSet[int]() | ||||
| for _, log := range logs { | for _, log := range logs { | ||||
| @@ -394,6 +462,11 @@ func GetUserLogs(userId int, logType int, startTimestamp int64, endTimestamp int | |||||
| return nil, 0, errors.New("查询日志失败") | 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) | formatUserLogs(logs, startIdx) | ||||
| return logs, total, err | return logs, total, err | ||||
| } | } | ||||
| @@ -152,6 +152,56 @@ func TestGetUserLogsFiltersByChatIDAndUpstreamID(t *testing.T) { | |||||
| require.Equal(t, "up-owner-a", logs[0].UpstreamId) | 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) { | func TestRecordConsumeLogPersistsContextChatIDAndUpstreamID(t *testing.T) { | ||||
| db := setupLogIdentityDB(t) | db := setupLogIdentityDB(t) | ||||
| @@ -274,6 +274,9 @@ func RecalculateTaskQuota(ctx context.Context, task *model.Task, actualQuota int | |||||
| taskAdjustTokenQuota(ctx, task, quotaDelta) | taskAdjustTokenQuota(ctx, task, quotaDelta) | ||||
| task.Quota = actualQuota | 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 logType int | ||||
| var logQuota int | var logQuota int | ||||
| @@ -380,6 +380,7 @@ func TestRecalculate_PositiveDelta(t *testing.T) { | |||||
| seedChannel(t, channelID) | seedChannel(t, channelID) | ||||
| task := makeTask(userID, channelID, preConsumed, tokenID, BillingSourceWallet, 0) | task := makeTask(userID, channelID, preConsumed, tokenID, BillingSourceWallet, 0) | ||||
| require.NoError(t, model.DB.Create(task).Error) | |||||
| RecalculateTaskQuota(ctx, task, actualQuota, "adaptor adjustment") | RecalculateTaskQuota(ctx, task, actualQuota, "adaptor adjustment") | ||||
| @@ -391,6 +392,9 @@ func TestRecalculate_PositiveDelta(t *testing.T) { | |||||
| // task.Quota should be updated to actualQuota | // task.Quota should be updated to actualQuota | ||||
| assert.Equal(t, actualQuota, task.Quota) | 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 type should be Consume (additional charge) | ||||
| log := getLastLog(t) | log := getLastLog(t) | ||||