fix(api): 文本回放 - 增加增量配额与生命周期清理 (#144)

Closes #144
This commit is contained in:
ddCat committed 2026-08-05 16:54:47 +08:00
1 parent 38c4f63e2e
commit 4721c1e97f
11 files changed
+389 -3

No files matched your search

@@ -58,6 +58,9 @@ func (r *Repository) BulkDisableUsers(actorID string, userIDs []string, events [
if err := tx.Delete(&model.AuthSession{}, "user_id IN ?", userIDs).Error; err != nil {
return err
}
if err := tx.Delete(&model.TaskTextDelta{}, "user_id IN ?", userIDs).Error; err != nil {
return err
}
if err := tx.Model(&model.User{}).Where("id IN ?", userIDs).Updates(map[string]any{"status": model.UserStatusDisabled, "updated_at": now}).Error; err != nil {
return err
}
+4 -1
View File
@@ -222,7 +222,7 @@ func (r *Repository) RetryTaskWithBilling(userID string, taskID string, order *m
}
updates := map[string]any{
"status": model.TaskStatusQueued, "stage": "等待队列调度", "progress": 5, "error": "", "result_json": "",
"started_at": nil, "completed_at": nil, "updated_at": time.Now(),
"text_draft": "", "started_at": nil, "completed_at": nil, "updated_at": time.Now(),
}
if order != nil {
updates["billing_order_id"] = order.ID
@@ -239,6 +239,9 @@ func (r *Repository) RetryTaskWithBilling(userID string, taskID string, order *m
if updated.RowsAffected != 1 {
return ErrTaskNotRetryable
}
if err := tx.Delete(&model.TaskTextDelta{}, "user_id = ? AND task_id = ?", userID, taskID).Error; err != nil {
return err
}
return tx.First(&task, "id = ? AND user_id = ?", taskID, userID).Error
})
return &task, err
+16 -2
View File
@@ -16,6 +16,10 @@ var ErrDailyUploadLimitExceeded = errors.New("daily upload limit exceeded")
var ErrTaskProviderRecoveryConflict = errors.New("task provider recovery is already running")
var ErrTextReplayQuotaExceeded = errors.New("text replay quota exceeded")
var ErrTextReplayClosed = errors.New("text replay task is closed")
type Repository struct {
db *gorm.DB
}
@@ -52,9 +56,10 @@ func (r *Repository) UserStorageUsage(userID string) (UserStorageUsage, error) {
(SELECT COALESCE(SUM(length(CAST(COALESCE(prompt, '') AS BLOB)) + length(CAST(COALESCE(canvas_snapshot_json, '') AS BLOB)) + length(CAST(COALESCE(canvas_ops_json, '') AS BLOB))), 0) FROM sessions WHERE user_id = ?)
+ (SELECT COALESCE(SUM(length(CAST(COALESCE(content, '') AS BLOB)) + length(CAST(COALESCE(payload, '') AS BLOB))), 0) FROM messages WHERE user_id = ?) AS session_bytes,
(SELECT COUNT(*) FROM tasks WHERE user_id = ?) AS task_count,
(SELECT COALESCE(SUM(length(CAST(COALESCE(prompt, '') AS BLOB)) + length(CAST(COALESCE(input_json, '') AS BLOB)) + length(CAST(COALESCE(result_json, '') AS BLOB)) + length(CAST(COALESCE(error, '') AS BLOB))), 0) FROM tasks WHERE user_id = ?)
(SELECT COALESCE(SUM(length(CAST(COALESCE(prompt, '') AS BLOB)) + length(CAST(COALESCE(input_json, '') AS BLOB)) + length(CAST(COALESCE(result_json, '') AS BLOB)) + length(CAST(COALESCE(text_draft, '') AS BLOB)) + length(CAST(COALESCE(error, '') AS BLOB))), 0) FROM tasks WHERE user_id = ?)
+ (SELECT COALESCE(SUM(length(CAST(COALESCE(message, '') AS BLOB)) + length(CAST(COALESCE(payload, '') AS BLOB))), 0) FROM task_logs WHERE user_id = ?)
+ (SELECT COALESCE(SUM(length(CAST(COALESCE(url, '') AS BLOB)) + length(CAST(COALESCE(payload, '') AS BLOB))), 0) FROM results WHERE user_id = ?)
+ (SELECT COALESCE(SUM(byte_count), 0) FROM task_text_deltas WHERE user_id = ?)
+ (SELECT COALESCE(SUM(length(CAST(COALESCE(path, '') AS BLOB)) + length(CAST(COALESCE(model, '') AS BLOB)) + length(CAST(COALESCE(provider_request_id, '') AS BLOB)) + length(CAST(COALESCE(error_code, '') AS BLOB)) + length(CAST(COALESCE(error, '') AS BLOB)) + length(CAST(COALESCE(upstream_url, '') AS BLOB)) + length(CAST(COALESCE(request_body, '') AS BLOB)) + length(CAST(COALESCE(response_body, '') AS BLOB))), 0) FROM api_call_logs WHERE user_id = ?) AS task_bytes,
(SELECT COUNT(*) FROM api_call_logs WHERE user_id = ?) AS api_call_count
`
@@ -62,7 +67,7 @@ func (r *Repository) UserStorageUsage(userID string) (UserStorageUsage, error) {
query = strings.ReplaceAll(query, "length(CAST(COALESCE(", "octet_length(COALESCE(")
query = strings.ReplaceAll(query, ", '') AS BLOB))", ", ''))")
}
err := r.db.Raw(query, userID, userID, userID, userID, userID, userID, userID, userID, userID, userID, userID, userID, userID).Scan(&usage).Error
err := r.db.Raw(query, userID, userID, userID, userID, userID, userID, userID, userID, userID, userID, userID, userID, userID, userID).Scan(&usage).Error
return usage, err
}
@@ -450,6 +455,15 @@ func (r *Repository) SessionForUser(userID string, id string) (*model.Session, e
func (r *Repository) DeleteSessionDraft(userID string, id string) error {
return r.db.Transaction(func(tx *gorm.DB) error {
var taskIDs []string
if err := tx.Model(&model.Task{}).Where("user_id = ? AND session_id = ?", userID, id).Pluck("id", &taskIDs).Error; err != nil {
return err
}
if len(taskIDs) > 0 {
if err := tx.Delete(&model.TaskTextDelta{}, "user_id = ? AND task_id IN ?", userID, taskIDs).Error; err != nil {
return err
}
}
if err := tx.Delete(&model.Message{}, "user_id = ? AND session_id = ?", userID, id).Error; err != nil {
return err
}
+135
View File
@@ -0,0 +1,135 @@
package repository
import (
"errors"
"time"
"infinite-canvas/backend/internal/model"
"gorm.io/gorm"
"gorm.io/gorm/clause"
)
type TextReplayLimits struct {
MaxTaskBytes int64
MaxUserBytes int64
MaxTaskEvents int64
}
type TextReplayStats struct {
EventCount int64 `json:"eventCount"`
TaskCount int64 `json:"taskCount"`
ByteCount int64 `json:"byteCount"`
OldestAt *time.Time `json:"oldestAt,omitempty"`
}
func (r *Repository) AppendTaskTextDelta(userID string, taskID string, content string, expiresAt time.Time, limits TextReplayLimits) (*model.TaskTextDelta, error) {
var created model.TaskTextDelta
err := r.db.Transaction(func(tx *gorm.DB) error {
var task model.Task
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Select("id", "user_id", "status").First(&task, "id = ? AND user_id = ?", taskID, userID).Error; err != nil {
return err
}
if task.Status == model.TaskStatusSucceeded || task.Status == model.TaskStatusFailed || task.Status == model.TaskStatusCancelled {
return ErrTextReplayClosed
}
var taskUsage struct {
Count int64
Bytes int64
Max int64
}
if err := tx.Model(&model.TaskTextDelta{}).Select("count(*) AS count, COALESCE(sum(byte_count), 0) AS bytes, COALESCE(max(sequence), 0) AS max").Where("task_id = ?", taskID).Scan(&taskUsage).Error; err != nil {
return err
}
var userBytes int64
if err := tx.Model(&model.TaskTextDelta{}).Select("COALESCE(sum(byte_count), 0)").Where("user_id = ?", userID).Scan(&userBytes).Error; err != nil {
return err
}
incoming := int64(len([]byte(content)))
if taskUsage.Count+1 > limits.MaxTaskEvents || taskUsage.Bytes+incoming > limits.MaxTaskBytes || userBytes+incoming > limits.MaxUserBytes {
return ErrTextReplayQuotaExceeded
}
created = model.TaskTextDelta{ID: newRepositoryID(), UserID: userID, TaskID: taskID, Sequence: taskUsage.Max + 1, Content: content, ByteCount: incoming, CreatedAt: time.Now(), ExpiresAt: expiresAt}
return tx.Create(&created).Error
})
return &created, err
}
func (r *Repository) TaskTextDeltas(userID string, taskID string, after int64, limit int) ([]model.TaskTextDelta, error) {
var items []model.TaskTextDelta
if limit <= 0 || limit > 1000 {
limit = 1000
}
err := r.db.Where("user_id = ? AND task_id = ? AND sequence > ?", userID, taskID, after).Order("sequence asc").Limit(limit).Find(&items).Error
return items, err
}
func (r *Repository) CompactTaskTextDeltas(taskID string, expiresAt time.Time, keepDraft bool) error {
return r.db.Transaction(func(tx *gorm.DB) error {
var task model.Task
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).First(&task, "id = ?", taskID).Error; err != nil {
return err
}
if keepDraft {
var items []model.TaskTextDelta
if err := tx.Select("content").Where("task_id = ?", taskID).Order("sequence asc").Find(&items).Error; err != nil {
return err
}
if task.TextDraft == "" {
for _, item := range items {
task.TextDraft += item.Content
}
if err := tx.Model(&task).Update("text_draft", task.TextDraft).Error; err != nil {
return err
}
}
}
return tx.Model(&model.TaskTextDelta{}).Where("task_id = ?", taskID).Update("expires_at", expiresAt).Error
})
}
func (r *Repository) CleanupTaskTextDeltas(now time.Time) (int64, error) {
var deleted int64
err := r.db.Transaction(func(tx *gorm.DB) error {
var taskIDs []string
if err := tx.Model(&model.TaskTextDelta{}).Where("expires_at <= ?", now).Distinct().Pluck("task_id", &taskIDs).Error; err != nil {
return err
}
for _, taskID := range taskIDs {
var task model.Task
if err := tx.Select("id", "status", "text_draft").First(&task, "id = ?", taskID).Error; err != nil {
if errors.Is(err, gorm.ErrRecordNotFound) {
continue
}
return err
}
if task.TextDraft != "" || (task.Status != model.TaskStatusFailed && task.Status != model.TaskStatusCancelled) {
continue
}
var items []model.TaskTextDelta
if err := tx.Select("content").Where("task_id = ?", taskID).Order("sequence asc").Find(&items).Error; err != nil {
return err
}
for _, item := range items {
task.TextDraft += item.Content
}
if err := tx.Model(&task).Update("text_draft", task.TextDraft).Error; err != nil {
return err
}
}
result := tx.Where("expires_at <= ? OR NOT EXISTS (SELECT 1 FROM tasks WHERE tasks.id = task_text_deltas.task_id)", now).Delete(&model.TaskTextDelta{})
deleted = result.RowsAffected
return result.Error
})
return deleted, err
}
func (r *Repository) DeleteUserTaskTextDeltas(userID string) error {
return r.db.Delete(&model.TaskTextDelta{}, "user_id = ?", userID).Error
}
func (r *Repository) TextReplayStats() (TextReplayStats, error) {
var stats TextReplayStats
err := r.db.Model(&model.TaskTextDelta{}).Select("count(*) AS event_count, count(DISTINCT task_id) AS task_count, COALESCE(sum(byte_count), 0) AS byte_count, min(created_at) AS oldest_at").Scan(&stats).Error
return stats, err
}