wangbo
810dcfeee6
停止持久化 provider 原始响应、兼容响应快照、attempt/event/outbox 重复 JSON，并由标准任务结果动态生成 Kling/Keling/Volces 兼容响应。 增加事件去重与预算、极简 callback 投递、7/30 天分批清理、安全删除条件、并发迁移索引及可实际恢复的任务域排除备份。历史清理默认关闭，待兼容协议和异步恢复在线验证后单独启用。 验证：Go 全量测试与 go vet、PostgreSQL 18 集成与实际备份恢复、迁移安全测试、bash -n、ShellCheck、Compose 配置和人工发布脚本测试均通过。
395 lines
14 KiB
Go
395 lines
14 KiB
Go
package store
|
|
|
|
import (
|
|
"context"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
|
|
"github.com/google/uuid"
|
|
)
|
|
|
|
func TestTaskEventAllowlistRejectsSyntheticProgress(t *testing.T) {
|
|
if taskEventAllowed("task.progress") || taskEventAllowed("task.attempt.started") || taskEventAllowed("unknown") {
|
|
t.Fatal("synthetic or unknown event type was allowed")
|
|
}
|
|
for _, eventType := range []string{"task.accepted", "task.running", "task.completed", "task.attempt.failed", "task.policy.priority_demoted"} {
|
|
if !taskEventAllowed(eventType) {
|
|
t.Fatalf("required event type %q was rejected", eventType)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestMinimalTaskResultDropsProviderRawCopy(t *testing.T) {
|
|
result := minimalTaskResult(map[string]any{
|
|
"data": []any{map[string]any{"url": "https://example.invalid/result"}},
|
|
"raw": map[string]any{"provider": "duplicate"},
|
|
"raw_data": map[string]any{"provider": "duplicate"},
|
|
"providerResponse": map[string]any{"provider": "duplicate"},
|
|
"upstream_task_id": "remote-task",
|
|
"provider_specific": "kept",
|
|
})
|
|
if result["raw"] != nil || result["raw_data"] != nil || result["providerResponse"] != nil ||
|
|
result["upstream_task_id"] != nil || result["data"] == nil || result["provider_specific"] != "kept" {
|
|
t.Fatalf("minimal task result=%#v", result)
|
|
}
|
|
}
|
|
|
|
func TestTaskEventsAreMinimalAndConsecutiveDuplicatesAreSkipped(t *testing.T) {
|
|
db := billingV2IntegrationStore(t)
|
|
ctx := context.Background()
|
|
user := &auth.User{ID: "event-minimal-" + uuid.NewString()}
|
|
task, err := db.CreateTask(ctx, CreateTaskInput{
|
|
Kind: "images.generations", Model: "event-minimal", RunMode: "simulation",
|
|
Request: map[string]any{"prompt": "minimal"},
|
|
}, user)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
t.Cleanup(func() {
|
|
_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_tasks WHERE id=$1::uuid`, task.ID)
|
|
})
|
|
|
|
running, err := db.AddTaskEvent(ctx, task.ID, "task.running", "running", "starting", 0.1, "ignored", map[string]any{"large": "ignored"}, true)
|
|
if err != nil || running.ID == "" || len(running.Payload) != 0 || running.Phase != "" || running.Progress != 0 || running.Message != "" {
|
|
t.Fatalf("running=%+v err=%v", running, err)
|
|
}
|
|
duplicate, err := db.AddTaskEvent(ctx, task.ID, "task.running", "running", "polling", 0.8, "ignored again", map[string]any{"other": "ignored"}, true)
|
|
if err != nil || duplicate.ID != "" || duplicate.SkippedReason != "duplicate" {
|
|
t.Fatalf("duplicate=%+v err=%v", duplicate, err)
|
|
}
|
|
synthetic, err := db.AddTaskEvent(ctx, task.ID, "task.progress", "running", "polling", 0.9, "ignored", nil, true)
|
|
if err != nil || synthetic.ID != "" || synthetic.SkippedReason != "unknown_type" {
|
|
t.Fatalf("synthetic=%+v err=%v", synthetic, err)
|
|
}
|
|
completed, err := db.AddTaskEvent(ctx, task.ID, "task.completed", "succeeded", "completed", 1, "ignored", map[string]any{"result": map[string]any{"duplicate": true}}, true)
|
|
if err != nil || completed.ID == "" || len(completed.Payload) != 0 {
|
|
t.Fatalf("completed=%+v err=%v", completed, err)
|
|
}
|
|
if err := db.QueueTaskCallback(ctx, completed, "https://callback.invalid/task"); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
var events int
|
|
var nonEmptyPayloads int
|
|
if err := db.pool.QueryRow(ctx, `
|
|
SELECT count(*), count(*) FILTER (WHERE payload <> '{}'::jsonb)
|
|
FROM gateway_task_events
|
|
WHERE task_id=$1::uuid`, task.ID).Scan(&events, &nonEmptyPayloads); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if events != 3 || nonEmptyPayloads != 0 {
|
|
t.Fatalf("events=%d nonEmptyPayloads=%d, want 3/0", events, nonEmptyPayloads)
|
|
}
|
|
var callbackPayloadBytes int
|
|
if err := db.pool.QueryRow(ctx, `
|
|
SELECT pg_column_size(payload)
|
|
FROM gateway_task_callback_outbox
|
|
WHERE task_id=$1::uuid`, task.ID).Scan(&callbackPayloadBytes); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if callbackPayloadBytes > 16 {
|
|
t.Fatalf("callback payload storage=%d, want an empty json object", callbackPayloadBytes)
|
|
}
|
|
budgetExceeded := false
|
|
for index := 0; index < 20; index++ {
|
|
eventType := "task.running"
|
|
status := "running"
|
|
if index%2 == 1 {
|
|
eventType = "task.queued"
|
|
status = "queued"
|
|
}
|
|
event, err := db.AddTaskEvent(ctx, task.ID, eventType, status, "", 0, "", nil, true)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if event.SkippedReason == "budget_exceeded" {
|
|
budgetExceeded = true
|
|
break
|
|
}
|
|
}
|
|
if !budgetExceeded {
|
|
t.Fatal("non-terminal event budget was not enforced")
|
|
}
|
|
var nonTerminalEvents int
|
|
if err := db.pool.QueryRow(ctx, `
|
|
SELECT count(*)
|
|
FROM gateway_task_events
|
|
WHERE task_id=$1::uuid
|
|
AND event_type NOT IN (
|
|
'task.completed', 'task.failed', 'task.cancelled',
|
|
'task.billing.settled', 'task.billing.released', 'task.billing.review'
|
|
)`, task.ID).Scan(&nonTerminalEvents); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if nonTerminalEvents != 16 {
|
|
t.Fatalf("non-terminal events=%d, want 16", nonTerminalEvents)
|
|
}
|
|
}
|
|
|
|
func TestTaskAttemptDropsDuplicateJSONSnapshots(t *testing.T) {
|
|
db := billingV2IntegrationStore(t)
|
|
ctx := context.Background()
|
|
user := &auth.User{ID: "attempt-minimal-" + uuid.NewString()}
|
|
task, err := db.CreateTask(ctx, CreateTaskInput{
|
|
Kind: "images.generations", Model: "attempt-minimal", RunMode: "simulation",
|
|
Request: map[string]any{"prompt": "canonical"},
|
|
}, user)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
t.Cleanup(func() {
|
|
_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_tasks WHERE id=$1::uuid`, task.ID)
|
|
})
|
|
attemptID, err := db.CreateTaskAttempt(ctx, CreateTaskAttemptInput{
|
|
TaskID: task.ID, AttemptNo: 1, QueueKey: "test", Status: "running",
|
|
RequestSnapshot: map[string]any{"prompt": "duplicate"},
|
|
Metrics: map[string]any{"trace": []any{map[string]any{"large": "duplicate"}}},
|
|
PricingSnapshot: map[string]any{"price": 1},
|
|
RequestFingerprint: "duplicate-fingerprint",
|
|
})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := db.FinishTaskAttempt(ctx, FinishTaskAttemptInput{
|
|
AttemptID: attemptID,
|
|
Status: "failed",
|
|
Retryable: true,
|
|
StatusCode: 503,
|
|
Usage: map[string]any{"output_tokens": 100},
|
|
Metrics: map[string]any{"trace": []any{map[string]any{"large": "duplicate"}}},
|
|
ResponseSnapshot: map[string]any{"result": "duplicate"},
|
|
ErrorCode: "upstream_error",
|
|
ErrorMessage: strings.Repeat("错误", 4096),
|
|
}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
attempts, err := db.ListTaskAttempts(ctx, task.ID)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if len(attempts) != 1 {
|
|
t.Fatalf("attempts=%d", len(attempts))
|
|
}
|
|
attempt := attempts[0]
|
|
if len(attempt.RequestSnapshot) != 0 || len(attempt.ResponseSnapshot) != 0 ||
|
|
len(attempt.Usage) != 0 || len(attempt.Metrics) != 0 || len(attempt.PricingSnapshot) != 0 {
|
|
t.Fatalf("attempt contains duplicate JSON: %+v", attempt)
|
|
}
|
|
if attempt.StatusCode != 503 || len(attempt.ErrorMessage) > 2048 {
|
|
t.Fatalf("attempt status/error not preserved safely: %+v", attempt)
|
|
}
|
|
}
|
|
|
|
func TestTaskRunningDoesNotPersistNormalizedRequestCopy(t *testing.T) {
|
|
db := billingV2IntegrationStore(t)
|
|
ctx := context.Background()
|
|
user := &auth.User{ID: "request-minimal-" + uuid.NewString()}
|
|
task, err := db.CreateTask(ctx, CreateTaskInput{
|
|
Kind: "images.generations", Model: "request-minimal", RunMode: "simulation",
|
|
Request: map[string]any{"prompt": "canonical"},
|
|
}, user)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
t.Cleanup(func() {
|
|
_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_tasks WHERE id=$1::uuid`, task.ID)
|
|
})
|
|
executionToken := uuid.NewString()
|
|
if _, err := db.ClaimTaskExecution(ctx, task.ID, executionToken, time.Minute); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := db.MarkTaskRunning(ctx, task.ID, executionToken, "image_generate", map[string]any{
|
|
"prompt": "duplicate normalized request",
|
|
}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
var normalizedBytes int
|
|
if err := db.pool.QueryRow(ctx, `
|
|
SELECT pg_column_size(normalized_request)
|
|
FROM gateway_tasks
|
|
WHERE id=$1::uuid`, task.ID).Scan(&normalizedBytes); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if normalizedBytes > 16 {
|
|
t.Fatalf("normalized request storage=%d, want an empty json object", normalizedBytes)
|
|
}
|
|
}
|
|
|
|
func TestTaskHistoryCleanupRemovesHistoricalProviderCopies(t *testing.T) {
|
|
db := billingV2IntegrationStore(t)
|
|
ctx := context.Background()
|
|
user := &auth.User{ID: "history-compaction-" + uuid.NewString()}
|
|
task, err := db.CreateTask(ctx, CreateTaskInput{
|
|
Kind: "images.generations", Model: "history-compaction", RunMode: "simulation",
|
|
Request: map[string]any{"prompt": "canonical"},
|
|
}, user)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
voiceID := uuid.NewString()
|
|
t.Cleanup(func() {
|
|
_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_cloned_voices WHERE id=$1::uuid`, voiceID)
|
|
_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_tasks WHERE id=$1::uuid`, task.ID)
|
|
})
|
|
if _, err := db.pool.Exec(ctx, `
|
|
UPDATE gateway_tasks
|
|
SET result='{"data":[{"url":"https://example.invalid/result"}],"raw_data":{"provider":"duplicate"},"upstream_task_id":"remote"}'::jsonb
|
|
WHERE id=$1::uuid`,
|
|
task.ID); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, err := db.pool.Exec(ctx, `
|
|
INSERT INTO gateway_cloned_voices (id, user_id, provider, voice_id, metadata)
|
|
VALUES ($1::uuid, 'history-compaction', 'test', $2, '{"request":{"duplicate":true},"rawData":{"provider":"duplicate"},"keep":"value"}'::jsonb)`,
|
|
voiceID, "voice-"+voiceID); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
result, err := db.CleanupTaskHistory(ctx, time.Now().AddDate(0, 0, -7), time.Now().AddDate(0, 0, -30), 100)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if result.CompactedTasks < 1 || result.CompactedClonedVoices < 1 {
|
|
t.Fatalf("cleanup result=%+v", result)
|
|
}
|
|
var taskResult map[string]any
|
|
var voiceMetadata map[string]any
|
|
if err := db.pool.QueryRow(ctx, `SELECT result FROM gateway_tasks WHERE id=$1::uuid`, task.ID).Scan(&taskResult); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := db.pool.QueryRow(ctx, `SELECT metadata FROM gateway_cloned_voices WHERE id=$1::uuid`, voiceID).Scan(&voiceMetadata); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if taskResult["raw_data"] != nil || taskResult["upstream_task_id"] != nil || taskResult["data"] == nil {
|
|
t.Fatalf("compacted task result=%+v", taskResult)
|
|
}
|
|
if voiceMetadata["request"] != nil || voiceMetadata["rawData"] != nil || voiceMetadata["keep"] != "value" {
|
|
t.Fatalf("compacted voice metadata=%+v", voiceMetadata)
|
|
}
|
|
}
|
|
|
|
func TestTaskHistoryCleanupKeepsTaskUntilCallbackRetentionExpires(t *testing.T) {
|
|
db := billingV2IntegrationStore(t)
|
|
ctx := context.Background()
|
|
user := &auth.User{ID: "cleanup-" + uuid.NewString()}
|
|
createOldTask := func(label string) GatewayTask {
|
|
t.Helper()
|
|
task, err := db.CreateTask(ctx, CreateTaskInput{
|
|
Kind: "images.generations", Model: label, RunMode: "simulation",
|
|
Request: map[string]any{"prompt": label},
|
|
}, user)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, err := db.pool.Exec(ctx, `
|
|
UPDATE gateway_tasks
|
|
SET status='succeeded', billing_status='not_required',
|
|
finished_at=now()-interval '40 days', updated_at=now()-interval '40 days'
|
|
WHERE id=$1::uuid`, task.ID); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, err := db.pool.Exec(ctx, `
|
|
UPDATE gateway_task_events
|
|
SET created_at=now()-interval '10 days'
|
|
WHERE task_id=$1::uuid`, task.ID); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
return task
|
|
}
|
|
safe := createOldTask("cleanup-safe")
|
|
protected := createOldTask("cleanup-protected")
|
|
t.Cleanup(func() {
|
|
_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_tasks WHERE id IN ($1::uuid, $2::uuid)`, safe.ID, protected.ID)
|
|
})
|
|
completed, err := db.AddTaskEvent(ctx, protected.ID, "task.completed", "succeeded", "", 0, "", nil, true)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := db.QueueTaskCallback(ctx, completed, "https://callback.invalid/task"); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, err := db.pool.Exec(ctx, `
|
|
UPDATE gateway_task_callback_outbox
|
|
SET status='delivered', delivered_at=now(), updated_at=now()
|
|
WHERE task_id=$1::uuid`, protected.ID); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
result, err := db.CleanupTaskHistory(ctx, time.Now().AddDate(0, 0, -7), time.Now().AddDate(0, 0, -30), 100)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if result.DeletedTasks < 1 {
|
|
t.Fatalf("deletedTasks=%d, want at least 1", result.DeletedTasks)
|
|
}
|
|
var safeExists bool
|
|
var protectedExists bool
|
|
if err := db.pool.QueryRow(ctx, `SELECT EXISTS (SELECT 1 FROM gateway_tasks WHERE id=$1::uuid)`, safe.ID).Scan(&safeExists); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := db.pool.QueryRow(ctx, `SELECT EXISTS (SELECT 1 FROM gateway_tasks WHERE id=$1::uuid)`, protected.ID).Scan(&protectedExists); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if safeExists || !protectedExists {
|
|
t.Fatalf("safeExists=%t protectedExists=%t", safeExists, protectedExists)
|
|
}
|
|
}
|
|
|
|
func TestLegacyCallbacksAreNotReplayed(t *testing.T) {
|
|
db := billingV2IntegrationStore(t)
|
|
ctx := context.Background()
|
|
user := &auth.User{ID: "legacy-callback-" + uuid.NewString()}
|
|
task, err := db.CreateTask(ctx, CreateTaskInput{
|
|
Kind: "images.generations", Model: "legacy-callback", RunMode: "simulation",
|
|
Request: map[string]any{"prompt": "legacy"},
|
|
}, user)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
t.Cleanup(func() {
|
|
_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_tasks WHERE id=$1::uuid`, task.ID)
|
|
})
|
|
event, err := db.AddTaskEvent(ctx, task.ID, "task.running", "running", "", 0, "", nil, true)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := db.QueueTaskCallback(ctx, event, "https://callback.invalid/legacy"); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, err := db.pool.Exec(ctx, `
|
|
UPDATE gateway_task_callback_outbox
|
|
SET created_at=(
|
|
SELECT applied_at - interval '1 second'
|
|
FROM schema_migrations
|
|
WHERE version='0083_task_history_minimal_storage'
|
|
)
|
|
WHERE task_id=$1::uuid`, task.ID); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
claimed, err := db.ClaimTaskCallbacks(ctx, "legacy-test", 10, time.Minute)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if len(claimed) != 0 {
|
|
t.Fatalf("claimed legacy callbacks=%d, want 0", len(claimed))
|
|
}
|
|
result, err := db.CleanupTaskHistory(ctx, time.Now().AddDate(0, 0, -7), time.Now().AddDate(0, 0, -30), 100)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if result.RetiredCallbacks < 1 {
|
|
t.Fatalf("retired legacy callbacks=%d, want at least 1", result.RetiredCallbacks)
|
|
}
|
|
var status string
|
|
if err := db.pool.QueryRow(ctx, `
|
|
SELECT status
|
|
FROM gateway_task_callback_outbox
|
|
WHERE task_id=$1::uuid`, task.ID).Scan(&status); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if status != "failed" {
|
|
t.Fatalf("legacy callback status=%q, want failed", status)
|
|
}
|
|
}
|