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:
@@ -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, `
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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",
|
||||
}}
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user