perf(storage): 极简化任务历史并增加保留治理

停止持久化 provider 原始响应、兼容响应快照、attempt/event/outbox 重复 JSON,并由标准任务结果动态生成 Kling/Keling/Volces 兼容响应。

增加事件去重与预算、极简 callback 投递、7/30 天分批清理、安全删除条件、并发迁移索引及可实际恢复的任务域排除备份。历史清理默认关闭,待兼容协议和异步恢复在线验证后单独启用。

验证:Go 全量测试与 go vet、PostgreSQL 18 集成与实际备份恢复、迁移安全测试、bash -n、ShellCheck、Compose 配置和人工发布脚本测试均通过。
This commit is contained in:
2026-07-24 18:23:22 +08:00
parent 09375bfae7
commit 810dcfeee6
46 changed files with 2246 additions and 458 deletions
@@ -432,13 +432,13 @@ SELECT task.id,
COALESCE((SELECT MAX(event.seq) + 1 FROM gateway_task_events event WHERE event.task_id = task.id), 1),
CASE WHEN $2 = 'settled' THEN 'task.billing.settled' ELSE 'task.billing.released' END,
task.status,
'billing',
1,
CASE WHEN $2 = 'settled' THEN 'task billing settled' ELSE 'task billing reservation released' END,
jsonb_build_object('settlementId', $3::text, 'billingStatus', $2::text),
NULL,
0,
NULL,
'{}'::jsonb,
task.run_mode = 'simulation'
FROM gateway_tasks task
WHERE task.id = $1::uuid`, taskID, billingStatus, settlementID); err != nil {
WHERE task.id = $1::uuid`, taskID, billingStatus); err != nil {
return err
}
tag, err := tx.Exec(ctx, `
+11 -14
View File
@@ -2,7 +2,6 @@ package store
import (
"context"
"encoding/json"
"strings"
)
@@ -16,25 +15,19 @@ type CompatibilitySubmission struct {
}
func (s *Store) SetTaskCompatibilitySubmission(ctx context.Context, taskID string, submission CompatibilitySubmission) error {
headers, _ := json.Marshal(emptyObjectIfNil(submission.Headers))
body, _ := json.Marshal(emptyObjectIfNil(submission.Body))
_, err := s.pool.Exec(ctx, `
UPDATE gateway_tasks
SET compatibility_protocol = NULLIF($2, ''),
compatibility_public_id = NULLIF($3, ''),
compatibility_source_protocol = NULLIF($4, ''),
compatibility_submit_http_status = NULLIF($5, 0),
compatibility_submit_headers = $6::jsonb,
compatibility_submit_body = $7::jsonb,
SET compatibility_protocol = COALESCE(NULLIF($2, ''), compatibility_protocol),
compatibility_public_id = NULL,
compatibility_source_protocol = COALESCE(NULLIF($3, ''), compatibility_source_protocol),
compatibility_submit_http_status = NULL,
compatibility_submit_headers = '{}'::jsonb,
compatibility_submit_body = '{}'::jsonb,
updated_at = now()
WHERE id = $1::uuid`,
taskID,
strings.TrimSpace(submission.TargetProtocol),
strings.TrimSpace(submission.PublicID),
strings.TrimSpace(submission.SourceProtocol),
submission.HTTPStatus,
string(headers),
string(body),
)
return err
}
@@ -44,7 +37,11 @@ func (s *Store) GetCompatibilityTask(ctx context.Context, protocol string, publi
SELECT `+gatewayTaskColumns+`
FROM gateway_tasks
WHERE compatibility_protocol = $1
AND compatibility_public_id = $2
AND (
id::text = $2
OR remote_task_id = $2
OR compatibility_public_id = $2
)
LIMIT 1`, strings.TrimSpace(protocol), strings.TrimSpace(publicID)))
if err != nil {
return GatewayTask{}, err
@@ -307,6 +307,22 @@ func applyOIDCJITTestMigrations(t *testing.T, ctx context.Context, databaseURL s
if err != nil {
t.Fatalf("read migration %s: %v", version, err)
}
migrationSQL := string(migration)
const noTransactionMarker = "-- easyai:migration:no-transaction"
if strings.HasPrefix(strings.TrimSpace(migrationSQL), noTransactionMarker) {
migrationSQL = strings.TrimSpace(strings.TrimPrefix(strings.TrimSpace(migrationSQL), noTransactionMarker))
for _, statement := range strings.Split(migrationSQL, "-- easyai:migration:statement") {
if statement = strings.TrimSpace(statement); statement != "" {
if _, err := pool.Exec(ctx, statement); err != nil {
t.Fatalf("apply non-transaction migration %s: %v", version, err)
}
}
}
if _, err := pool.Exec(ctx, `INSERT INTO schema_migrations(version) VALUES($1)`, version); err != nil {
t.Fatalf("record non-transaction migration %s: %v", version, err)
}
continue
}
tx, err := pool.Begin(ctx)
if err != nil {
t.Fatalf("begin migration %s: %v", version, err)
+17 -16
View File
@@ -587,17 +587,19 @@ COALESCE(compatibility_submit_headers, '{}'::jsonb), COALESCE(compatibility_subm
created_at, updated_at, COALESCE(finished_at::text, '')`
type TaskEvent struct {
ID string `json:"id"`
TaskID string `json:"taskId"`
Seq int64 `json:"seq"`
EventType string `json:"eventType"`
Status string `json:"status,omitempty"`
Phase string `json:"phase,omitempty"`
Progress float64 `json:"progress,omitempty"`
Message string `json:"message,omitempty"`
Payload map[string]any `json:"payload,omitempty"`
Simulated bool `json:"simulated"`
CreatedAt time.Time `json:"createdAt"`
ID string `json:"id"`
TaskID string `json:"taskId"`
PlatformID string `json:"-"`
SkippedReason string `json:"-"`
Seq int64 `json:"seq"`
EventType string `json:"eventType"`
Status string `json:"status,omitempty"`
Phase string `json:"phase,omitempty"`
Progress float64 `json:"progress,omitempty"`
Message string `json:"message,omitempty"`
Payload map[string]any `json:"payload,omitempty"`
Simulated bool `json:"simulated"`
CreatedAt time.Time `json:"createdAt"`
}
type TaskAttempt struct {
@@ -2181,7 +2183,8 @@ func scanGatewayTask(scanner taskScanner) (GatewayTask, error) {
func (s *Store) ListTaskEvents(ctx context.Context, taskID string) ([]TaskEvent, error) {
rows, err := s.pool.Query(ctx, `
SELECT id::text, task_id::text, seq, event_type, COALESCE(status, ''), COALESCE(phase, ''),
COALESCE(progress, 0)::float8, COALESCE(message, ''), payload, simulated, created_at
COALESCE(progress, 0)::float8, COALESCE(message, ''), payload, simulated, created_at,
COALESCE(platform_id::text, '')
FROM gateway_task_events
WHERE task_id = $1::uuid
ORDER BY seq ASC`, taskID)
@@ -2206,6 +2209,7 @@ ORDER BY seq ASC`, taskID)
&payload,
&item.Simulated,
&item.CreatedAt,
&item.PlatformID,
); err != nil {
return nil, err
}
@@ -2421,10 +2425,7 @@ func taskEventsForCreate(taskID string, runMode string, status string, result ma
Seq: 1,
EventType: "task.accepted",
Status: "queued",
Phase: "queued",
Progress: 0,
Message: "task accepted",
Payload: map[string]any{"taskId": taskID},
Payload: map[string]any{},
Simulated: runMode == "simulation",
}}
}
+26 -12
View File
@@ -338,19 +338,20 @@ func (s *Store) listRecentPriorityDemotionsByPlatform(ctx context.Context, statu
return out, nil
}
rows, err := s.pool.Query(ctx, `
SELECT id::text, task_id::text, COALESCE(message, ''), payload, created_at
SELECT id::text, task_id::text, COALESCE(message, ''), payload, created_at, event_platform_id
FROM (
SELECT e.*,
COALESCE(e.platform_id::text, e.payload->>'platformId') AS event_platform_id,
row_number() OVER (
PARTITION BY e.payload->>'platformId'
PARTITION BY COALESCE(e.platform_id::text, e.payload->>'platformId')
ORDER BY e.created_at DESC, e.seq DESC
) AS demotion_rank
FROM gateway_task_events e
WHERE e.event_type = 'task.policy.priority_demoted'
AND e.payload->>'platformId' = ANY($1::text[])
AND COALESCE(e.platform_id::text, e.payload->>'platformId') = ANY($1::text[])
) ranked
WHERE demotion_rank <= $2
ORDER BY payload->>'platformId' ASC, created_at DESC`, platformIDs, limit)
ORDER BY event_platform_id ASC, created_at DESC`, platformIDs, limit)
if err != nil {
return nil, err
}
@@ -361,10 +362,16 @@ func (s *Store) listRecentPriorityDemotionsByPlatform(ctx context.Context, statu
var message string
var payloadBytes []byte
var createdAt time.Time
if err := rows.Scan(&id, &taskID, &message, &payloadBytes, &createdAt); err != nil {
var platformID string
if err := rows.Scan(&id, &taskID, &message, &payloadBytes, &createdAt, &platformID); err != nil {
return nil, err
}
record := priorityDemotionRecordFromEventPayload(id, taskID, message, decodeObject(payloadBytes), createdAt)
payload := decodeObject(payloadBytes)
if payload == nil {
payload = map[string]any{}
}
payload["platformId"] = platformID
record := priorityDemotionRecordFromEventPayload(id, taskID, message, payload, createdAt)
if record.PlatformID == "" {
continue
}
@@ -420,12 +427,13 @@ func (s *Store) listLatestPlatformDisabledReasons(ctx context.Context, statuses
return out, nil
}
rows, err := s.pool.Query(ctx, `
SELECT id::text, task_id::text, event_type, COALESCE(message, ''), payload, COALESCE(attempt_error_message, ''), created_at
SELECT id::text, task_id::text, event_type, COALESCE(message, ''), payload, COALESCE(attempt_error_message, ''), created_at, event_platform_id
FROM (
SELECT e.*,
COALESCE(e.platform_id::text, e.payload->>'platformId') AS event_platform_id,
a.error_message AS attempt_error_message,
row_number() OVER (
PARTITION BY e.payload->>'platformId'
PARTITION BY COALESCE(e.platform_id::text, e.payload->>'platformId')
ORDER BY e.created_at DESC, e.seq DESC
) AS disabled_rank
FROM gateway_task_events e
@@ -433,12 +441,12 @@ FROM (
SELECT error_message
FROM gateway_task_attempts attempt
WHERE attempt.task_id = e.task_id
AND attempt.platform_id::text = e.payload->>'platformId'
AND attempt.platform_id::text = COALESCE(e.platform_id::text, e.payload->>'platformId')
ORDER BY attempt.attempt_no DESC, attempt.started_at DESC
LIMIT 1
) a ON TRUE
WHERE e.event_type IN ('task.policy.failover_disabled', 'task.policy.auto_disabled')
AND e.payload->>'platformId' = ANY($1::text[])
AND COALESCE(e.platform_id::text, e.payload->>'platformId') = ANY($1::text[])
) ranked
WHERE disabled_rank = 1`, platformIDs)
if err != nil {
@@ -453,10 +461,16 @@ WHERE disabled_rank = 1`, platformIDs)
var payloadBytes []byte
var attemptErrorMessage string
var createdAt time.Time
if err := rows.Scan(&id, &taskID, &eventType, &message, &payloadBytes, &attemptErrorMessage, &createdAt); err != nil {
var platformID string
if err := rows.Scan(&id, &taskID, &eventType, &message, &payloadBytes, &attemptErrorMessage, &createdAt, &platformID); err != nil {
return nil, err
}
record := platformPolicyEventFromPayload(id, taskID, eventType, message, attemptErrorMessage, decodeObject(payloadBytes), createdAt)
payload := decodeObject(payloadBytes)
if payload == nil {
payload = map[string]any{}
}
payload["platformId"] = platformID
record := platformPolicyEventFromPayload(id, taskID, eventType, message, attemptErrorMessage, payload, createdAt)
if record.PlatformID == "" {
continue
}
+35 -14
View File
@@ -469,17 +469,37 @@ RETURNING id::text`)
for _, taskID := range asyncTaskIDs {
if _, err := tx.Exec(ctx, `
INSERT INTO gateway_task_events (task_id, seq, event_type, status, phase, progress, message, payload, simulated)
VALUES (
$1::uuid,
SELECT
task.id,
COALESCE((SELECT MAX(seq) + 1 FROM gateway_task_events WHERE task_id = $1::uuid), 1),
'task.recovered',
'task.queued',
'queued',
'recovered',
0.2,
'async task recovered after service restart',
'{"code":"server_restarted"}'::jsonb,
NULL,
0,
NULL,
'{}'::jsonb,
false
)`, taskID); err != nil {
FROM gateway_tasks task
WHERE task.id = $1::uuid
AND (
SELECT count(*)
FROM gateway_task_events event
WHERE event.task_id = task.id
AND event.event_type NOT IN (
'task.completed', 'task.failed', 'task.cancelled',
'task.billing.settled', 'task.billing.released', 'task.billing.review'
)
) < 16
AND NOT EXISTS (
SELECT 1
FROM gateway_task_events event
WHERE event.task_id = task.id
AND event.seq = (SELECT MAX(last_event.seq) FROM gateway_task_events last_event WHERE last_event.task_id = task.id)
AND event.event_type = 'task.queued'
AND COALESCE(event.status, '') = 'queued'
AND event.platform_id IS NULL
AND event.simulated = false
)`, taskID); err != nil {
return RuntimeRecoveryResult{}, err
}
}
@@ -488,9 +508,10 @@ VALUES (
taskRows, err := tx.Query(ctx, `
UPDATE gateway_tasks
SET status = 'failed',
error = 'task interrupted by service restart',
error = NULL,
error_code = 'server_restarted',
error_message = 'task interrupted by service restart',
remote_task_payload = '{}'::jsonb,
finished_at = now(),
updated_at = now()
WHERE async_mode = false
@@ -520,12 +541,12 @@ INSERT INTO gateway_task_events (task_id, seq, event_type, status, phase, progre
VALUES (
$1::uuid,
COALESCE((SELECT MAX(seq) + 1 FROM gateway_task_events WHERE task_id = $1::uuid), 1),
'task.recovered',
'task.failed',
'failed',
'recovered',
1,
'task interrupted by service restart',
'{"code":"server_restarted"}'::jsonb,
NULL,
0,
NULL,
'{}'::jsonb,
false
)`, taskID); err != nil {
return RuntimeRecoveryResult{}, err
+1
View File
@@ -259,6 +259,7 @@ type FinishTaskAttemptInput struct {
Status string
Retryable bool
RequestID string
StatusCode int
Usage map[string]any
Metrics map[string]any
ResponseSnapshot map[string]any
+496
View File
@@ -0,0 +1,496 @@
package store
import (
"context"
"time"
"github.com/jackc/pgx/v5"
)
const (
TaskCallbackBatchSize = 100
TaskCallbackLockTTL = 60 * time.Second
)
type TaskCallbackDelivery struct {
ID string
TaskID string
EventID string
Seq int64
CallbackURL string
Status string
Attempts int
LockToken string
EventType string
TaskStatus string
CreatedAt time.Time
}
type TaskHistoryCleanupResult struct {
CompactedTasks int64
CompactedCheckpoints int64
CompactedAttempts int64
CompactedEvents int64
CompactedCallbacks int64
RetiredCallbacks int64
CompactedParamLogs int64
CompactedClonedVoices int64
DeletedCallbacks int64
DeletedParamLogs int64
DeletedEvents int64
DeletedAttempts int64
DeletedTasks int64
}
func (result TaskHistoryCleanupResult) Total() int64 {
return result.CompactedTasks +
result.CompactedCheckpoints +
result.CompactedAttempts +
result.CompactedEvents +
result.CompactedCallbacks +
result.RetiredCallbacks +
result.CompactedParamLogs +
result.CompactedClonedVoices +
result.DeletedCallbacks +
result.DeletedParamLogs +
result.DeletedEvents +
result.DeletedAttempts +
result.DeletedTasks
}
func (s *Store) ClaimTaskCallbacks(ctx context.Context, workerID string, limit int, staleAfter time.Duration) ([]TaskCallbackDelivery, error) {
if limit <= 0 || limit > TaskCallbackBatchSize {
limit = TaskCallbackBatchSize
}
if staleAfter <= 0 {
staleAfter = TaskCallbackLockTTL
}
rows, err := s.pool.Query(ctx, `
WITH picked AS (
SELECT outbox.id, event.event_type, COALESCE(event.status, '') AS task_status, event.created_at
FROM gateway_task_callback_outbox outbox
JOIN gateway_task_events event ON event.id = outbox.event_id
WHERE outbox.created_at >= COALESCE((
SELECT applied_at
FROM schema_migrations
WHERE version = '0083_task_history_minimal_storage'
), now())
AND (
(
outbox.status = 'pending'
AND outbox.next_attempt_at <= now()
)
OR (
outbox.status = 'processing'
AND outbox.locked_at < now() - ($3::int * interval '1 second')
)
)
ORDER BY outbox.next_attempt_at, outbox.created_at
LIMIT $2
FOR UPDATE OF outbox SKIP LOCKED
)
UPDATE gateway_task_callback_outbox outbox
SET status = 'processing',
attempts = outbox.attempts + 1,
locked_by = $1,
lock_token = gen_random_uuid(),
locked_at = now(),
updated_at = now()
FROM picked
WHERE outbox.id = picked.id
RETURNING outbox.id::text, outbox.task_id::text, COALESCE(outbox.event_id::text, ''),
outbox.seq, outbox.callback_url, outbox.status, outbox.attempts,
COALESCE(outbox.lock_token::text, ''), picked.event_type, picked.task_status,
picked.created_at`,
workerID, limit, int(staleAfter/time.Second))
if err != nil {
return nil, err
}
defer rows.Close()
items := make([]TaskCallbackDelivery, 0)
for rows.Next() {
var item TaskCallbackDelivery
if err := rows.Scan(
&item.ID,
&item.TaskID,
&item.EventID,
&item.Seq,
&item.CallbackURL,
&item.Status,
&item.Attempts,
&item.LockToken,
&item.EventType,
&item.TaskStatus,
&item.CreatedAt,
); err != nil {
return nil, err
}
items = append(items, item)
}
return items, rows.Err()
}
func (s *Store) MarkTaskCallbackDelivered(ctx context.Context, item TaskCallbackDelivery) error {
_, err := s.pool.Exec(ctx, `
UPDATE gateway_task_callback_outbox
SET status = 'delivered',
delivered_at = now(),
failed_at = NULL,
last_error = NULL,
locked_by = NULL,
lock_token = NULL,
locked_at = NULL,
updated_at = now()
WHERE id = $1::uuid
AND status = 'processing'
AND lock_token = $2::uuid`, item.ID, item.LockToken)
return err
}
func (s *Store) MarkTaskCallbackFailed(ctx context.Context, item TaskCallbackDelivery, retry bool, nextAttemptAt time.Time, message string) error {
status := "failed"
if retry {
status = "pending"
}
_, err := s.pool.Exec(ctx, `
UPDATE gateway_task_callback_outbox
SET status = $3,
next_attempt_at = CASE WHEN $3 = 'pending' THEN $4::timestamptz ELSE next_attempt_at END,
failed_at = CASE WHEN $3 = 'failed' THEN now() ELSE NULL END,
last_error = NULLIF(left($5, 2048), ''),
locked_by = NULL,
lock_token = NULL,
locked_at = NULL,
updated_at = now()
WHERE id = $1::uuid
AND status = 'processing'
AND lock_token = $2::uuid`,
item.ID, item.LockToken, status, nextAttemptAt, message)
return err
}
func (s *Store) CleanupTaskHistory(ctx context.Context, analysisCutoff time.Time, taskCutoff time.Time, batchSize int) (TaskHistoryCleanupResult, error) {
if batchSize < 1 {
batchSize = 1000
}
var result TaskHistoryCleanupResult
err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
if _, err := tx.Exec(ctx, `SET LOCAL lock_timeout = '5s'; SET LOCAL statement_timeout = '5s'`); err != nil {
return err
}
steps := []struct {
target *int64
sql string
args []any
}{
{&result.CompactedTasks, `
WITH picked AS (
SELECT id
FROM gateway_tasks
WHERE COALESCE(normalized_request, '{}'::jsonb) <> '{}'::jsonb
OR compatibility_public_id IS NOT NULL
OR compatibility_submit_http_status IS NOT NULL
OR COALESCE(compatibility_submit_headers, '{}'::jsonb) <> '{}'::jsonb
OR COALESCE(compatibility_submit_body, '{}'::jsonb) <> '{}'::jsonb
OR result ?| ARRAY[
'raw', 'raw_data', 'rawData', 'provider_response', 'providerResponse',
'submit', 'file_retrieve', 'fileRetrieve', 'upstream_task_id', 'remote_task_id'
]
OR error IS NOT NULL
OR octet_length(COALESCE(error_message, '')) > 2048
ORDER BY updated_at
LIMIT $1
FOR UPDATE SKIP LOCKED
)
UPDATE gateway_tasks task
SET normalized_request = '{}'::jsonb,
compatibility_public_id = NULL,
compatibility_submit_http_status = NULL,
compatibility_submit_headers = '{}'::jsonb,
compatibility_submit_body = '{}'::jsonb,
result = COALESCE(task.result, '{}'::jsonb)
- 'raw' - 'raw_data' - 'rawData'
- 'provider_response' - 'providerResponse'
- 'submit' - 'file_retrieve' - 'fileRetrieve'
- 'upstream_task_id' - 'remote_task_id',
error = NULL,
error_message = CASE
WHEN octet_length(COALESCE(task.error_message, '')) > 2048 THEN left(task.error_message, 512)
ELSE task.error_message
END,
remote_task_payload = CASE
WHEN task.status IN ('succeeded', 'failed', 'cancelled') THEN '{}'::jsonb
ELSE (
SELECT COALESCE(jsonb_object_agg(entry.key, entry.value), '{}'::jsonb)
FROM jsonb_each(COALESCE(task.remote_task_payload, '{}'::jsonb)) entry
WHERE entry.key IN (
'endpoint', 'pollEndpoint', 'phase', 'mode', 'taskType',
'targetResolution', 'imageToken', 'receipt', 'cleanupElementIds'
)
)
END
FROM picked
WHERE task.id = picked.id`, []any{batchSize}},
{&result.CompactedCheckpoints, `
WITH picked AS (
SELECT id
FROM gateway_tasks
WHERE COALESCE(remote_task_payload, '{}'::jsonb) <> '{}'::jsonb
AND (
status IN ('succeeded', 'failed', 'cancelled')
OR EXISTS (
SELECT 1
FROM jsonb_object_keys(COALESCE(remote_task_payload, '{}'::jsonb)) AS key
WHERE key NOT IN (
'endpoint', 'pollEndpoint', 'phase', 'mode', 'taskType',
'targetResolution', 'imageToken', 'receipt', 'cleanupElementIds'
)
)
)
ORDER BY updated_at
LIMIT $1
FOR UPDATE SKIP LOCKED
)
UPDATE gateway_tasks task
SET remote_task_payload = CASE
WHEN task.status IN ('succeeded', 'failed', 'cancelled') THEN '{}'::jsonb
ELSE (
SELECT COALESCE(jsonb_object_agg(entry.key, entry.value), '{}'::jsonb)
FROM jsonb_each(COALESCE(task.remote_task_payload, '{}'::jsonb)) entry
WHERE entry.key IN (
'endpoint', 'pollEndpoint', 'phase', 'mode', 'taskType',
'targetResolution', 'imageToken', 'receipt', 'cleanupElementIds'
)
)
END
FROM picked
WHERE task.id = picked.id`, []any{batchSize}},
{&result.CompactedAttempts, `
WITH picked AS (
SELECT id
FROM gateway_task_attempts
WHERE request_snapshot <> '{}'::jsonb
OR COALESCE(response_snapshot, '{}'::jsonb) <> '{}'::jsonb
OR COALESCE(usage, '{}'::jsonb) <> '{}'::jsonb
OR COALESCE(metrics, '{}'::jsonb) <> '{}'::jsonb
OR COALESCE(pricing_snapshot, '{}'::jsonb) <> '{}'::jsonb
OR request_fingerprint IS NOT NULL
ORDER BY started_at
LIMIT $1
FOR UPDATE SKIP LOCKED
)
UPDATE gateway_task_attempts attempt
SET request_snapshot = '{}'::jsonb,
response_snapshot = '{}'::jsonb,
usage = '{}'::jsonb,
metrics = '{}'::jsonb,
pricing_snapshot = '{}'::jsonb,
request_fingerprint = NULL,
error_message = NULLIF(left(error_message, 2048), '')
FROM picked
WHERE attempt.id = picked.id`, []any{batchSize}},
{&result.CompactedEvents, `
WITH picked AS (
SELECT id
FROM gateway_task_events
WHERE payload <> '{}'::jsonb
OR phase IS NOT NULL
OR COALESCE(progress, 0) <> 0
OR message IS NOT NULL
ORDER BY created_at
LIMIT $1
FOR UPDATE SKIP LOCKED
)
UPDATE gateway_task_events event
SET platform_id = CASE
WHEN event.platform_id IS NULL
AND event.payload->>'platformId' ~* '^[0-9a-f]{8}-[0-9a-f]{4}-[1-5][0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$'
THEN (event.payload->>'platformId')::uuid
ELSE event.platform_id
END,
phase = NULL,
progress = 0,
message = NULL,
payload = '{}'::jsonb
FROM picked
WHERE event.id = picked.id`, []any{batchSize}},
{&result.CompactedCallbacks, `
WITH picked AS (
SELECT id
FROM gateway_task_callback_outbox
WHERE payload <> '{}'::jsonb
ORDER BY created_at
LIMIT $1
FOR UPDATE SKIP LOCKED
)
UPDATE gateway_task_callback_outbox outbox
SET payload = '{}'::jsonb
FROM picked
WHERE outbox.id = picked.id`, []any{batchSize}},
{&result.RetiredCallbacks, `
WITH migration_boundary AS (
SELECT applied_at
FROM schema_migrations
WHERE version = '0083_task_history_minimal_storage'
), picked AS (
SELECT outbox.id
FROM gateway_task_callback_outbox outbox
CROSS JOIN migration_boundary boundary
WHERE outbox.created_at < boundary.applied_at
AND outbox.status IN ('pending', 'processing')
ORDER BY outbox.created_at
LIMIT $1
FOR UPDATE OF outbox SKIP LOCKED
)
UPDATE gateway_task_callback_outbox outbox
SET status = 'failed',
failed_at = outbox.created_at,
last_error = 'legacy_callback_not_replayed',
locked_by = NULL,
lock_token = NULL,
locked_at = NULL,
updated_at = now()
FROM picked
WHERE outbox.id = picked.id`, []any{batchSize}},
{&result.CompactedParamLogs, `
WITH picked AS (
SELECT id
FROM gateway_task_param_preprocessing_logs
WHERE actual_input <> '{}'::jsonb
OR converted_output <> '{}'::jsonb
OR model_snapshot <> '{}'::jsonb
OR octet_length(changes::text) > 1024
ORDER BY created_at
LIMIT $1
FOR UPDATE SKIP LOCKED
)
UPDATE gateway_task_param_preprocessing_logs log
SET actual_input = '{}'::jsonb,
converted_output = '{}'::jsonb,
model_snapshot = '{}'::jsonb,
changes = CASE WHEN octet_length(log.changes::text) > 1024 THEN '[]'::jsonb ELSE log.changes END
FROM picked
WHERE log.id = picked.id`, []any{batchSize}},
{&result.CompactedClonedVoices, `
WITH picked AS (
SELECT id
FROM gateway_cloned_voices
WHERE metadata ?| ARRAY['raw', 'raw_data', 'rawData', 'provider_response', 'providerResponse', 'request']
ORDER BY updated_at
LIMIT $1
FOR UPDATE SKIP LOCKED
)
UPDATE gateway_cloned_voices voice
SET metadata = COALESCE(voice.metadata, '{}'::jsonb)
- 'raw' - 'raw_data' - 'rawData'
- 'provider_response' - 'providerResponse' - 'request',
updated_at = now()
FROM picked
WHERE voice.id = picked.id`, []any{batchSize}},
{&result.DeletedCallbacks, `
WITH picked AS (
SELECT id
FROM gateway_task_callback_outbox
WHERE (status = 'delivered' AND delivered_at < now() - interval '24 hours')
OR (status = 'failed' AND COALESCE(failed_at, updated_at) < $1::timestamptz)
ORDER BY updated_at
LIMIT $2
FOR UPDATE SKIP LOCKED
)
DELETE FROM gateway_task_callback_outbox outbox
USING picked
WHERE outbox.id = picked.id`, []any{analysisCutoff, batchSize}},
{&result.DeletedParamLogs, `
WITH picked AS (
SELECT id
FROM gateway_task_param_preprocessing_logs
WHERE created_at < $1::timestamptz
ORDER BY created_at
LIMIT $2
FOR UPDATE SKIP LOCKED
)
DELETE FROM gateway_task_param_preprocessing_logs log
USING picked
WHERE log.id = picked.id`, []any{analysisCutoff, batchSize}},
{&result.DeletedEvents, `
WITH picked AS (
SELECT event.id
FROM gateway_task_events event
WHERE event.created_at < $1::timestamptz
AND NOT EXISTS (
SELECT 1
FROM gateway_task_callback_outbox outbox
WHERE outbox.event_id = event.id
AND outbox.status IN ('pending', 'processing')
)
ORDER BY event.created_at
LIMIT $2
FOR UPDATE OF event SKIP LOCKED
)
DELETE FROM gateway_task_events event
USING picked
WHERE event.id = picked.id`, []any{analysisCutoff, batchSize}},
{&result.DeletedAttempts, `
WITH picked AS (
SELECT id
FROM gateway_task_attempts
WHERE COALESCE(finished_at, started_at) < $1::timestamptz
AND status <> 'running'
ORDER BY COALESCE(finished_at, started_at)
LIMIT $2
FOR UPDATE SKIP LOCKED
)
DELETE FROM gateway_task_attempts attempt
USING picked
WHERE attempt.id = picked.id`, []any{analysisCutoff, batchSize}},
{&result.DeletedTasks, `
WITH picked AS (
SELECT task.id
FROM gateway_tasks task
WHERE task.finished_at < $1::timestamptz
AND task.status IN ('succeeded', 'failed', 'cancelled')
AND task.billing_status IN ('settled', 'released', 'not_required')
AND (task.execution_lease_expires_at IS NULL OR task.execution_lease_expires_at <= now())
AND NOT EXISTS (
SELECT 1
FROM settlement_outbox settlement
WHERE settlement.task_id = task.id
AND settlement.status IN ('pending', 'processing', 'retryable_failed', 'manual_review')
)
AND NOT EXISTS (
SELECT 1
FROM gateway_task_callback_outbox callback
WHERE callback.task_id = task.id
)
AND NOT EXISTS (
SELECT 1
FROM gateway_concurrency_leases lease
WHERE lease.task_id = task.id
AND lease.released_at IS NULL
AND lease.expires_at > now()
)
AND NOT EXISTS (
SELECT 1
FROM gateway_rate_limit_reservations reservation
WHERE reservation.task_id = task.id
AND reservation.status = 'reserved'
)
ORDER BY task.finished_at
LIMIT $2
FOR UPDATE OF task SKIP LOCKED
)
DELETE FROM gateway_tasks task
USING picked
WHERE task.id = picked.id`, []any{taskCutoff, batchSize}},
}
for _, step := range steps {
tag, err := tx.Exec(ctx, step.sql, step.args...)
if err != nil {
return err
}
*step.target = tag.RowsAffected()
}
return nil
})
return result, err
}
@@ -0,0 +1,394 @@
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)
}
}
+209 -100
View File
@@ -8,6 +8,7 @@ import (
"strconv"
"strings"
"time"
"unicode/utf8"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
"github.com/jackc/pgx/v5"
@@ -324,7 +325,7 @@ UPDATE gateway_tasks
SET status = 'failed',
billing_status = CASE WHEN $2 THEN 'manual_review' ELSE 'not_required' END,
billing_updated_at = now(),
error = 'upstream submission result is unknown',
error = NULL,
error_code = 'upstream_submission_unknown',
error_message = 'upstream submission result is unknown',
locked_by = NULL,
@@ -332,6 +333,7 @@ SET status = 'failed',
heartbeat_at = NULL,
execution_token = NULL,
execution_lease_expires_at = NULL,
remote_task_payload = '{}'::jsonb,
finished_at = now(),
updated_at = now()
WHERE id = $1::uuid`, taskID, hasGatewayUser); err != nil {
@@ -416,19 +418,18 @@ SELECT COALESCE((SELECT status = 'running' FROM gateway_tasks WHERE id = $1::uui
return nil
}
func (s *Store) MarkTaskRunning(ctx context.Context, taskID string, executionToken string, modelType string, normalizedRequest map[string]any) error {
normalizedJSON, _ := json.Marshal(emptyObjectIfNil(normalizedRequest))
func (s *Store) MarkTaskRunning(ctx context.Context, taskID string, executionToken string, modelType string, _ map[string]any) error {
tag, err := s.pool.Exec(ctx, `
UPDATE gateway_tasks
SET status = 'running',
model_type = NULLIF($3::text, ''),
normalized_request = $4::jsonb,
normalized_request = '{}'::jsonb,
locked_at = now(),
heartbeat_at = now(),
updated_at = now()
WHERE id = $1::uuid
AND status = 'running'
AND execution_token = $2::uuid`, taskID, executionToken, modelType, string(normalizedJSON))
AND execution_token = $2::uuid`, taskID, executionToken, modelType)
if err != nil {
return err
}
@@ -557,12 +558,14 @@ func (s *Store) CancelQueuedTask(ctx context.Context, taskID string, message str
if message == "" {
message = "任务已取消"
}
message = truncateUTF8Bytes(message, 2048)
tag, err := s.pool.Exec(ctx, `
UPDATE gateway_tasks
SET status = 'cancelled',
error = NULLIF($2::text, ''),
error = NULL,
error_code = 'task_cancelled',
error_message = NULLIF($2::text, ''),
remote_task_payload = '{}'::jsonb,
locked_by = NULL,
locked_at = NULL,
heartbeat_at = NULL,
@@ -592,6 +595,7 @@ func (s *Store) CancelSubmittedTask(ctx context.Context, taskID string, executio
if message == "" {
message = "任务已由上游取消"
}
message = truncateUTF8Bytes(message, 2048)
var task GatewayTask
changed := false
err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
@@ -599,7 +603,7 @@ func (s *Store) CancelSubmittedTask(ctx context.Context, taskID string, executio
task, err = scanGatewayTask(tx.QueryRow(ctx, `
UPDATE gateway_tasks
SET status = 'cancelled',
error = NULLIF($2, ''),
error = NULL,
error_code = 'task_cancelled',
error_message = NULLIF($2, ''),
billing_status = CASE
@@ -613,6 +617,7 @@ SET status = 'cancelled',
heartbeat_at = NULL,
execution_token = NULL,
execution_lease_expires_at = NULL,
remote_task_payload = '{}'::jsonb,
finished_at = now(),
updated_at = now()
WHERE id = $1::uuid
@@ -683,9 +688,6 @@ LIMIT $1`, limit)
}
func (s *Store) CreateTaskAttempt(ctx context.Context, input CreateTaskAttemptInput) (string, error) {
requestJSON, _ := json.Marshal(emptyObjectIfNil(input.RequestSnapshot))
metricsJSON, _ := json.Marshal(emptyObjectIfNil(input.Metrics))
pricingJSON, _ := json.Marshal(emptyObjectIfNil(input.PricingSnapshot))
tx, err := s.pool.Begin(ctx)
if err != nil {
return "", err
@@ -701,7 +703,7 @@ INSERT INTO gateway_task_attempts (
)
VALUES (
$1::uuid, $2::int, NULLIF($3::text, '')::uuid, NULLIF($4::text, '')::uuid, NULLIF($5::text, ''), $6,
$7, $8, $9::jsonb, $10::jsonb, $11::jsonb, NULLIF($12, ''),
$7, $8, '{}'::jsonb, '{}'::jsonb, '{}'::jsonb, NULL,
'not_submitted', now()
)
RETURNING id::text`,
@@ -713,10 +715,6 @@ RETURNING id::text`,
input.QueueKey,
firstNonEmpty(input.Status, "running"),
input.Simulated,
string(requestJSON),
string(metricsJSON),
string(pricingJSON),
input.RequestFingerprint,
).Scan(&attemptID)
if err != nil {
return "", err
@@ -731,13 +729,16 @@ WHERE id = $1::uuid`, input.TaskID, input.AttemptNo); err != nil {
}
func (s *Store) CreateTaskParamPreprocessingLog(ctx context.Context, input CreateTaskParamPreprocessingLogInput) (string, error) {
actualInputJSON, _ := json.Marshal(emptyObjectIfNil(input.ActualInput))
convertedOutputJSON, _ := json.Marshal(emptyObjectIfNil(input.ConvertedOutput))
if !input.Changed {
return "", nil
}
changesJSON, _ := json.Marshal(input.Changes)
if input.Changes == nil {
changesJSON = []byte("[]")
}
modelSnapshotJSON, _ := json.Marshal(emptyObjectIfNil(input.ModelSnapshot))
if len(changesJSON) > 1024 {
changesJSON = []byte("[]")
}
var attemptNo any
if input.AttemptNo > 0 {
attemptNo = input.AttemptNo
@@ -751,7 +752,7 @@ INSERT INTO gateway_task_param_preprocessing_logs (
VALUES (
$1::uuid, NULLIF($2::text, '')::uuid, $3::int, NULLIF($4::text, ''),
NULLIF($5::text, '')::uuid, NULLIF($6::text, '')::uuid, NULLIF($7::text, ''),
$8, $9::int, $10::jsonb, $11::jsonb, $12::jsonb, $13::jsonb
true, $8::int, '{}'::jsonb, '{}'::jsonb, $9::jsonb, '{}'::jsonb
)
RETURNING id::text`,
input.TaskID,
@@ -761,12 +762,8 @@ RETURNING id::text`,
input.PlatformID,
input.PlatformModelID,
input.ClientID,
input.Changed,
input.ChangeCount,
string(actualInputJSON),
string(convertedOutputJSON),
string(changesJSON),
string(modelSnapshotJSON),
).Scan(&id)
return id, err
}
@@ -823,24 +820,7 @@ ORDER BY COALESCE(attempt_no, 0), created_at`, taskID)
}
func (s *Store) AppendTaskAttemptTrace(ctx context.Context, taskID string, attemptNo int, entry map[string]any) error {
entryJSON, _ := json.Marshal(emptyObjectIfNil(entry))
_, err := s.pool.Exec(ctx, `
UPDATE gateway_task_attempts
SET metrics = jsonb_set(
COALESCE(metrics, '{}'::jsonb),
'{trace}',
(
CASE
WHEN jsonb_typeof(COALESCE(metrics->'trace', '[]'::jsonb)) = 'array'
THEN COALESCE(metrics->'trace', '[]'::jsonb)
ELSE '[]'::jsonb
END
) || jsonb_build_array($3::jsonb),
true
)
WHERE task_id = $1::uuid
AND attempt_no = $2::int`, taskID, attemptNo, string(entryJSON))
return err
return nil
}
func (s *Store) listTaskAttemptsByTaskIDs(ctx context.Context, taskIDs []string) (map[string][]TaskAttempt, error) {
@@ -855,7 +835,8 @@ SELECT a.id::text, a.task_id::text, a.attempt_no,
COALESCE(NULLIF(pm.provider_model_name, ''), pm.model_name, ''),
COALESCE(pm.model_alias, ''),
COALESCE(a.client_id, ''), a.queue_key, a.status, a.retryable, a.simulated,
COALESCE(a.request_id, ''), COALESCE(a.usage, '{}'::jsonb), COALESCE(a.metrics, '{}'::jsonb),
COALESCE(a.request_id, ''), COALESCE(a.status_code, 0),
COALESCE(a.usage, '{}'::jsonb), COALESCE(a.metrics, '{}'::jsonb),
a.request_snapshot, COALESCE(a.response_snapshot, '{}'::jsonb),
COALESCE(a.response_started_at::text, ''), COALESCE(a.response_finished_at::text, ''),
COALESCE(a.response_duration_ms, 0), COALESCE(a.error_code, ''), COALESCE(a.error_message, ''),
@@ -908,6 +889,7 @@ func scanTaskAttempt(scanner taskScanner) (TaskAttempt, error) {
&item.Retryable,
&item.Simulated,
&item.RequestID,
&item.StatusCode,
&usageBytes,
&metricsBytes,
&requestBytes,
@@ -980,7 +962,9 @@ func enrichTaskAttemptFromMetrics(item *TaskAttempt) {
item.ModelAlias = firstNonEmpty(item.ModelAlias, taskAttemptMetricString(item.Metrics, "modelAlias"))
item.ModelType = firstNonEmpty(item.ModelType, taskAttemptMetricString(item.Metrics, "modelType"))
item.ClientID = firstNonEmpty(item.ClientID, taskAttemptMetricString(item.Metrics, "clientId"))
item.StatusCode = taskAttemptMetricInt(item.Metrics, "statusCode")
if item.StatusCode == 0 {
item.StatusCode = taskAttemptMetricInt(item.Metrics, "statusCode")
}
}
func taskAttemptMetricString(metrics map[string]any, key string) string {
@@ -1008,42 +992,44 @@ func taskAttemptMetricInt(metrics map[string]any, key string) int {
}
func (s *Store) FinishTaskAttempt(ctx context.Context, input FinishTaskAttemptInput) error {
responseJSON, _ := json.Marshal(emptyObjectIfNil(input.ResponseSnapshot))
usageJSON, _ := json.Marshal(emptyObjectIfNil(input.Usage))
metricsJSON, _ := json.Marshal(emptyObjectIfNil(input.Metrics))
statusCode := input.StatusCode
if statusCode == 0 {
statusCode = taskAttemptMetricInt(input.Metrics, "statusCode")
}
_, err := s.pool.Exec(ctx, `
UPDATE gateway_task_attempts
SET status = $2::text,
retryable = $3,
request_id = NULLIF($4::text, ''),
usage = $5::jsonb,
metrics = $6::jsonb,
response_snapshot = $7::jsonb,
response_started_at = $8::timestamptz,
response_finished_at = $9::timestamptz,
response_duration_ms = $10,
error_code = NULLIF($11::text, ''),
error_message = NULLIF($12::text, ''),
status_code = NULLIF($5::int, 0),
usage = '{}'::jsonb,
metrics = '{}'::jsonb,
response_snapshot = '{}'::jsonb,
pricing_snapshot = '{}'::jsonb,
request_fingerprint = NULL,
response_started_at = $6::timestamptz,
response_finished_at = $7::timestamptz,
response_duration_ms = $8,
error_code = NULLIF($9::text, ''),
error_message = NULLIF(left($10::text, 2048), ''),
finished_at = now()
WHERE id = $1::uuid`,
input.AttemptID,
input.Status,
input.Retryable,
input.RequestID,
string(usageJSON),
string(metricsJSON),
string(responseJSON),
statusCode,
nullableTime(input.ResponseStartedAt),
nullableTime(input.ResponseFinishedAt),
input.ResponseDurationMS,
input.ErrorCode,
input.ErrorMessage,
truncateUTF8Bytes(input.ErrorMessage, 2048),
)
return err
}
func (s *Store) FinishTaskSuccess(ctx context.Context, input FinishTaskSuccessInput) (GatewayTask, error) {
resultJSON, _ := json.Marshal(emptyObjectIfNil(input.Result))
resultJSON, _ := json.Marshal(minimalTaskResult(input.Result))
billingsJSON, _ := json.Marshal(input.Billings)
usageJSON, _ := json.Marshal(emptyObjectIfNil(input.Usage))
metricsJSON, _ := json.Marshal(emptyObjectIfNil(input.Metrics))
@@ -1061,22 +1047,22 @@ UPDATE gateway_task_attempts
SET status = 'succeeded',
retryable = false,
request_id = NULLIF($2, ''),
usage = $3::jsonb,
metrics = $4::jsonb,
response_snapshot = $5::jsonb,
pricing_snapshot = $6::jsonb,
request_fingerprint = NULLIF($7, ''),
status_code = NULL,
usage = '{}'::jsonb,
metrics = '{}'::jsonb,
response_snapshot = '{}'::jsonb,
pricing_snapshot = '{}'::jsonb,
request_fingerprint = NULL,
upstream_submission_status = 'response_received',
upstream_submission_updated_at = now(),
response_started_at = $8::timestamptz,
response_finished_at = $9::timestamptz,
response_duration_ms = $10,
response_started_at = $3::timestamptz,
response_finished_at = $4::timestamptz,
response_duration_ms = $5,
error_code = NULL,
error_message = NULL,
finished_at = now()
WHERE id = $1::uuid`,
input.AttemptID, input.RequestID, string(usageJSON), string(metricsJSON),
string(resultJSON), string(pricingSnapshotJSON), input.RequestFingerprint,
input.AttemptID, input.RequestID,
nullableTime(input.ResponseStartedAt), nullableTime(input.ResponseFinishedAt), input.ResponseDurationMS,
); err != nil {
return err
@@ -1113,6 +1099,7 @@ SET status = 'succeeded',
heartbeat_at = NULL,
execution_token = NULL,
execution_lease_expires_at = NULL,
remote_task_payload = '{}'::jsonb,
finished_at = now(),
updated_at = now()
WHERE id = $1::uuid
@@ -1166,7 +1153,12 @@ func (s *Store) FinishTaskManualReview(ctx context.Context, input FinishTaskManu
if status != "succeeded" && status != "failed" {
return GatewayTask{}, fmt.Errorf("manual review task status must be succeeded or failed")
}
resultJSON, _ := json.Marshal(emptyObjectIfNil(input.Result))
message := truncateUTF8Bytes(input.Message, 2048)
result := input.Result
if status == "failed" {
result = nil
}
resultJSON, _ := json.Marshal(minimalTaskResult(result))
pricingSnapshotJSON, _ := json.Marshal(emptyObjectIfNil(input.PricingSnapshot))
err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
if strings.TrimSpace(input.AttemptID) != "" {
@@ -1186,7 +1178,7 @@ SET status = $2,
response_finished_at = $7::timestamptz,
response_duration_ms = $8,
finished_at = now()
WHERE id = $1::uuid`, input.AttemptID, attemptStatus, input.RequestID, input.Code, input.Message,
WHERE id = $1::uuid`, input.AttemptID, attemptStatus, input.RequestID, input.Code, message,
nullableTime(input.ResponseStartedAt), nullableTime(input.ResponseFinishedAt), input.ResponseDurationMS); err != nil {
return err
}
@@ -1196,7 +1188,7 @@ UPDATE gateway_tasks
SET status = $2,
result = $3::jsonb,
request_id = NULLIF($4, ''),
error = CASE WHEN $2 = 'failed' THEN NULLIF($6, '') ELSE NULL END,
error = NULL,
error_code = NULLIF($5, ''),
error_message = NULLIF($6, ''),
billing_status = 'manual_review',
@@ -1211,11 +1203,12 @@ SET status = $2,
heartbeat_at = NULL,
execution_token = NULL,
execution_lease_expires_at = NULL,
remote_task_payload = '{}'::jsonb,
finished_at = now(),
updated_at = now()
WHERE id = $1::uuid
AND status = 'running'
AND execution_token = $12::uuid`, input.TaskID, status, string(resultJSON), input.RequestID, input.Code, input.Message,
AND execution_token = $12::uuid`, input.TaskID, status, string(resultJSON), input.RequestID, input.Code, message,
string(pricingSnapshotJSON), input.RequestFingerprint, nullableTime(input.ResponseStartedAt),
nullableTime(input.ResponseFinishedAt), input.ResponseDurationMS, input.ExecutionToken)
if err != nil {
@@ -1397,12 +1390,13 @@ func taskBillingString(value any) string {
func (s *Store) FinishTaskFailure(ctx context.Context, input FinishTaskFailureInput) (GatewayTask, error) {
metricsJSON, _ := json.Marshal(emptyObjectIfNil(input.Metrics))
resultJSON, _ := json.Marshal(emptyObjectIfNil(input.Result))
resultJSON, _ := json.Marshal(minimalTaskResult(nil))
message := truncateUTF8Bytes(input.Message, 2048)
err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
tag, err := tx.Exec(ctx, `
UPDATE gateway_tasks
SET status = 'failed',
error = NULLIF($2::text, ''),
error = NULL,
error_code = NULLIF($3::text, ''),
error_message = NULLIF($2::text, ''),
request_id = NULLIF($4::text, ''),
@@ -1422,13 +1416,14 @@ func (s *Store) FinishTaskFailure(ctx context.Context, input FinishTaskFailureIn
heartbeat_at = NULL,
execution_token = NULL,
execution_lease_expires_at = NULL,
remote_task_payload = '{}'::jsonb,
finished_at = now(),
updated_at = now()
WHERE id = $1::uuid
AND status = 'running'
AND execution_token = $10::uuid`,
input.TaskID,
input.Message,
message,
input.Code,
input.RequestID,
string(metricsJSON),
@@ -1472,29 +1467,152 @@ func nullableTime(value time.Time) any {
return value
}
func truncateUTF8Bytes(value string, maxBytes int) string {
if maxBytes <= 0 || len(value) <= maxBytes {
return value
}
end := maxBytes
for end > 0 && !utf8.ValidString(value[:end]) {
end--
}
return value[:end]
}
func minimalTaskResult(input map[string]any) map[string]any {
result := make(map[string]any, len(input))
for key, value := range input {
switch key {
case "raw",
"raw_data",
"rawData",
"provider_response",
"providerResponse",
"submit",
"file_retrieve",
"fileRetrieve",
"upstream_task_id",
"remote_task_id":
continue
}
result[key] = value
}
return result
}
var allowedTaskEventTypes = map[string]struct{}{
"task.accepted": {},
"task.running": {},
"task.queued": {},
"task.completed": {},
"task.failed": {},
"task.cancelled": {},
"task.attempt.failed": {},
"task.billing.settled": {},
"task.billing.released": {},
"task.billing.review": {},
"task.policy.auto_disabled": {},
"task.policy.degraded": {},
"task.policy.failover_disabled": {},
"task.policy.failover_cooled_down": {},
"task.policy.priority_demoted": {},
}
func taskEventAllowed(eventType string) bool {
_, ok := allowedTaskEventTypes[strings.TrimSpace(eventType)]
return ok
}
func taskEventTerminal(eventType string) bool {
switch strings.TrimSpace(eventType) {
case "task.completed", "task.failed", "task.cancelled",
"task.billing.settled", "task.billing.released", "task.billing.review":
return true
default:
return false
}
}
func taskEventPlatformID(payload map[string]any) string {
platformID, _ := payload["platformId"].(string)
return strings.TrimSpace(platformID)
}
func (s *Store) AddTaskEvent(ctx context.Context, taskID string, eventType string, status string, phase string, progress float64, message string, payload map[string]any, simulated bool) (TaskEvent, error) {
payloadJSON, _ := json.Marshal(emptyObjectIfNil(payload))
eventType = strings.TrimSpace(eventType)
if !taskEventAllowed(eventType) {
return TaskEvent{SkippedReason: "unknown_type"}, nil
}
platformID := taskEventPlatformID(payload)
tx, err := s.pool.Begin(ctx)
if err != nil {
return TaskEvent{}, err
}
defer tx.Rollback(ctx)
if _, err := tx.Exec(ctx, `SELECT 1 FROM gateway_tasks WHERE id = $1::uuid FOR UPDATE`, taskID); err != nil {
return TaskEvent{}, err
}
var previousType string
var previousStatus string
var previousPlatformID string
var previousSimulated bool
err = tx.QueryRow(ctx, `
SELECT event_type, COALESCE(status, ''), COALESCE(platform_id::text, ''), simulated
FROM gateway_task_events
WHERE task_id = $1::uuid
ORDER BY seq DESC
LIMIT 1`, taskID).Scan(&previousType, &previousStatus, &previousPlatformID, &previousSimulated)
if err != nil && !errors.Is(err, pgx.ErrNoRows) {
return TaskEvent{}, err
}
if err == nil &&
previousType == eventType &&
previousStatus == strings.TrimSpace(status) &&
previousPlatformID == platformID &&
previousSimulated == simulated {
if err := tx.Commit(ctx); err != nil {
return TaskEvent{}, err
}
return TaskEvent{SkippedReason: "duplicate"}, nil
}
if !taskEventTerminal(eventType) {
var nonTerminalCount int
if err := tx.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'
)`, taskID).Scan(&nonTerminalCount); err != nil {
return TaskEvent{}, err
}
if nonTerminalCount >= 16 {
if err := tx.Commit(ctx); err != nil {
return TaskEvent{}, err
}
return TaskEvent{SkippedReason: "budget_exceeded"}, nil
}
}
var event TaskEvent
var payloadBytes []byte
err := s.pool.QueryRow(ctx, `
err = tx.QueryRow(ctx, `
WITH next_seq AS (
SELECT COALESCE(MAX(seq), 0) + 1 AS seq
FROM gateway_task_events
WHERE task_id = $1::uuid
)
INSERT INTO gateway_task_events (task_id, seq, event_type, status, phase, progress, message, payload, simulated)
SELECT $1::uuid, next_seq.seq, $2::text, NULLIF($3::text, ''), NULLIF($4::text, ''), $5, NULLIF($6::text, ''), $7::jsonb, $8
INSERT INTO gateway_task_events (task_id, seq, event_type, status, phase, progress, message, payload, simulated, platform_id)
SELECT $1::uuid, next_seq.seq, $2::text, NULLIF($3::text, ''), NULL, 0, NULL, '{}'::jsonb, $4,
NULLIF($5::text, '')::uuid
FROM next_seq
RETURNING id::text, task_id::text, seq, event_type, COALESCE(status, ''), COALESCE(phase, ''),
COALESCE(progress, 0)::float8, COALESCE(message, ''), payload, simulated, created_at`,
COALESCE(progress, 0)::float8, COALESCE(message, ''), payload, simulated, created_at,
COALESCE(platform_id::text, '')`,
taskID,
eventType,
status,
phase,
progress,
message,
string(payloadJSON),
simulated,
platformID,
).Scan(
&event.ID,
&event.TaskID,
@@ -1507,39 +1625,30 @@ RETURNING id::text, task_id::text, seq, event_type, COALESCE(status, ''), COALES
&payloadBytes,
&event.Simulated,
&event.CreatedAt,
&event.PlatformID,
)
if err != nil {
return TaskEvent{}, err
}
event.Payload = decodeObject(payloadBytes)
if err := tx.Commit(ctx); err != nil {
return TaskEvent{}, err
}
return event, nil
}
func (s *Store) QueueTaskCallback(ctx context.Context, event TaskEvent, callbackURL string) error {
if callbackURL == "" {
if callbackURL == "" || event.ID == "" {
return nil
}
payloadJSON, _ := json.Marshal(map[string]any{
"taskId": event.TaskID,
"seq": event.Seq,
"eventType": event.EventType,
"status": event.Status,
"phase": event.Phase,
"progress": event.Progress,
"message": event.Message,
"payload": event.Payload,
"simulated": event.Simulated,
"createdAt": event.CreatedAt,
})
_, err := s.pool.Exec(ctx, `
INSERT INTO gateway_task_callback_outbox (task_id, event_id, seq, callback_url, payload)
VALUES ($1::uuid, $2::uuid, $3, $4, $5::jsonb)
VALUES ($1::uuid, $2::uuid, $3, $4, '{}'::jsonb)
ON CONFLICT (task_id, seq, callback_url) DO NOTHING`,
event.TaskID,
event.ID,
event.Seq,
callbackURL,
string(payloadJSON),
)
return err
}