Co-Authored-By: Codex <noreply@anthropic.com>master
| @@ -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 | |||
| } | |||
| @@ -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) | |||
| @@ -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 | |||
| @@ -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) | |||