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
+7
View File
@@ -73,6 +73,13 @@ TASK_PROGRESS_CALLBACK_ENABLED=true
TASK_PROGRESS_CALLBACK_URL=http://localhost:3000/internal/platform/task-progress-callbacks
TASK_PROGRESS_CALLBACK_TIMEOUT_MS=5000
TASK_PROGRESS_CALLBACK_MAX_ATTEMPTS=10
# First deploy with cleanup disabled; enable only after compatibility and
# asynchronous-resume verification has passed.
AI_GATEWAY_TASK_CLEANUP_ENABLED=false
AI_GATEWAY_TASK_RETENTION_DAYS=30
AI_GATEWAY_TASK_ANALYSIS_RETENTION_DAYS=7
AI_GATEWAY_TASK_CLEANUP_INTERVAL_SECONDS=300
AI_GATEWAY_TASK_CLEANUP_BATCH_SIZE=1000
CORS_ALLOWED_ORIGIN=http://localhost:5178,http://127.0.0.1:5178
VITE_GATEWAY_API_BASE_URL=http://localhost:8088
+37
View File
@@ -13,6 +13,11 @@ import (
"github.com/jackc/pgx/v5"
)
const (
noTransactionMigrationMarker = "-- easyai:migration:no-transaction"
migrationStatementSeparator = "-- easyai:migration:statement"
)
func main() {
cfg := config.Load()
logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
@@ -63,6 +68,22 @@ CREATE TABLE IF NOT EXISTS schema_migrations (
os.Exit(1)
}
noTransaction, statements := migrationStatements(string(sqlBytes))
if noTransaction {
for _, statement := range statements {
if _, err := conn.Exec(ctx, statement); err != nil {
logger.Error("execute non-transaction migration failed", "version", version, "error", err)
os.Exit(1)
}
}
if _, err := conn.Exec(ctx, "INSERT INTO schema_migrations(version) VALUES($1)", version); err != nil {
logger.Error("record non-transaction migration failed", "version", version, "error", err)
os.Exit(1)
}
logger.Info("migration applied", "version", version, "transactional", false)
continue
}
tx, err := conn.Begin(ctx)
if err != nil {
logger.Error("begin migration failed", "version", version, "error", err)
@@ -95,3 +116,19 @@ CREATE TABLE IF NOT EXISTS schema_migrations (
fmt.Println("migrations complete")
}
func migrationStatements(sql string) (bool, []string) {
trimmed := strings.TrimSpace(sql)
if !strings.HasPrefix(trimmed, noTransactionMigrationMarker) {
return false, []string{sql}
}
trimmed = strings.TrimSpace(strings.TrimPrefix(trimmed, noTransactionMigrationMarker))
parts := strings.Split(trimmed, migrationStatementSeparator)
statements := make([]string, 0, len(parts))
for _, part := range parts {
if statement := strings.TrimSpace(part); statement != "" {
statements = append(statements, statement)
}
}
return true, statements
}
+18
View File
@@ -14,6 +14,24 @@ import (
"github.com/jackc/pgx/v5/pgxpool"
)
func TestMigrationStatementsSupportsConcurrentIndexes(t *testing.T) {
payload := `-- easyai:migration:no-transaction
CREATE INDEX CONCURRENTLY IF NOT EXISTS first_index ON first_table(id);
-- easyai:migration:statement
CREATE INDEX CONCURRENTLY IF NOT EXISTS second_index ON second_table(id);
`
noTransaction, statements := migrationStatements(payload)
if !noTransaction {
t.Fatal("expected a non-transaction migration")
}
if len(statements) != 2 {
t.Fatalf("statements=%d, want 2", len(statements))
}
if !strings.Contains(statements[0], "first_index") || !strings.Contains(statements[1], "second_index") {
t.Fatalf("unexpected statements: %#v", statements)
}
}
func TestSeedanceInputImageConstraintMigrationKeepsCatalogSnapshotsInSync(t *testing.T) {
payload, err := os.ReadFile("../../migrations/0078_seedance_input_image_constraints.sql")
if err != nil {
+5 -4
View File
@@ -2101,7 +2101,7 @@ func TestVolcesClientVideoSubmitsAndPollsTask(t *testing.T) {
}
data, _ := response.Result["data"].([]any)
item, _ := data[0].(map[string]any)
if item["url"] != "https://example.com/out.mp4" || response.Usage.TotalTokens != 9 {
if item["url"] != "https://example.com/out.mp4" || response.Result["content"] != nil || response.Result["usage"] != nil || response.Usage.TotalTokens != 9 {
t.Fatalf("unexpected response: %+v usage=%+v", response.Result, response.Usage)
}
}
@@ -2145,8 +2145,9 @@ func TestVolcesClientVideoRetriesTransientPollAndKeepsOfficialResult(t *testing.
if polls != 2 || len(persisted) != 1 || persisted[0] != "cgt-retry:succeeded" {
t.Fatalf("unexpected poll state polls=%d persisted=%+v", polls, persisted)
}
if response.Result["updated_at"] != float64(124) || response.Result["seed"] != float64(7) || response.Result["raw"] == nil {
t.Fatalf("official result fields lost: %+v", response.Result)
if response.Result["seed"] != float64(7) || response.Result["raw"] != nil || response.Result["updated_at"] != nil ||
response.Result["content"] != nil || response.Result["usage"] != nil || response.Result["upstream_task_id"] != nil {
t.Fatalf("provider response should be reduced to the canonical result: %+v", response.Result)
}
}
@@ -2424,7 +2425,7 @@ func TestVolcesClientVideoResumePollsExistingTaskID(t *testing.T) {
}
data, _ := response.Result["data"].([]any)
item, _ := data[0].(map[string]any)
if response.Result["upstream_task_id"] != "cgt-existing" || item["url"] != "https://example.com/resumed.mp4" {
if response.Result["upstream_task_id"] != nil || item["url"] != "https://example.com/resumed.mp4" {
t.Fatalf("unexpected resumed response: %+v", response.Result)
}
}
-2
View File
@@ -627,7 +627,6 @@ func geminiResult(request Request, raw map[string]any) map[string]any {
"created": nowUnix(),
"model": request.Model,
"data": data,
"raw": raw,
}
}
message, finishReason := geminiChatMessage(raw)
@@ -642,7 +641,6 @@ func geminiResult(request Request, raw map[string]any) map[string]any {
"message": message,
}},
"usage": geminiUsageMap(raw),
"raw": raw,
}
}
-2
View File
@@ -1306,7 +1306,6 @@ func kelingVideoSuccessResult(request Request, upstreamTaskID string, raw map[st
"status": "succeeded",
"upstream_task_id": upstreamTaskID,
"data": items,
"raw": raw,
}
}
@@ -1342,7 +1341,6 @@ func kelingTaskAPIVideoSuccessResult(request Request, upstreamTaskID string, tas
"status": "succeeded",
"upstream_task_id": upstreamTaskID,
"data": items,
"raw": raw,
}
}
+17 -17
View File
@@ -326,17 +326,17 @@ func (c MinimaxClient) runSpeech(ctx context.Context, request Request) (Response
if err != nil {
return Response{}, &ClientError{Code: "invalid_response", Message: "minimax speech audio hex is invalid: " + err.Error(), RequestID: firstNonEmptyString(requestID, requestIDFromResult(result)), ResponseStartedAt: startedAt, ResponseFinishedAt: finishedAt, ResponseDurationMS: responseDurationMS(startedAt, finishedAt), Retryable: false}
}
normalized := cloneMapAny(result)
normalized["status"] = "success"
normalized["created"] = time.Now().UnixMilli()
normalized["model"] = request.Model
normalized["raw_data"] = cloneMapAny(result)
normalized["data"] = []any{map[string]any{
"type": "audio",
"content": "data:audio/mpeg;base64," + base64.StdEncoding.EncodeToString(audioBytes),
"mime_type": "audio/mpeg",
"uploaded": false,
}}
normalized := map[string]any{
"status": "success",
"created": time.Now().UnixMilli(),
"model": request.Model,
"data": []any{map[string]any{
"type": "audio",
"content": "data:audio/mpeg;base64," + base64.StdEncoding.EncodeToString(audioBytes),
"mime_type": "audio/mpeg",
"uploaded": false,
}},
}
return Response{
Result: normalized,
RequestID: firstNonEmptyString(requestID, requestIDFromResult(result)),
@@ -378,12 +378,12 @@ func (c MinimaxClient) runVoiceClone(ctx context.Context, request Request) (Resp
if isProviderTaskFailure(providerTaskSpec{Name: "minimax"}, result) {
return Response{}, providerTaskFailure(providerTaskSpec{Name: "minimax"}, result, firstNonEmptyString(requestID, uploadRequestID, requestIDFromResult(result)), startedAt)
}
normalized := cloneMapAny(result)
normalized["status"] = "success"
normalized["created"] = time.Now().UnixMilli()
normalized["model"] = request.Model
normalized["voice_id"] = stringFromAny(payload["voice_id"])
normalized["raw_data"] = cloneMapAny(result)
normalized := map[string]any{
"status": "success",
"created": time.Now().UnixMilli(),
"model": request.Model,
"voice_id": stringFromAny(payload["voice_id"]),
}
if demoAudio := firstNonEmptyString(valueAtPath(result, "demo_audio"), valueAtPath(result, "data.demo_audio")); demoAudio != "" {
normalized["demo_audio"] = demoAudio
normalized["data"] = []any{map[string]any{"type": "audio", "url": demoAudio}}
+11 -15
View File
@@ -340,24 +340,20 @@ func hasProviderTaskResult(result map[string]any) bool {
return result["data"] != nil || valueAtPath(result, "data.result") != nil || valueAtPath(result, "data.audio") != nil || valueAtPath(result, "output.image_urls") != nil || valueAtPath(result, "output.video_url") != nil || valueAtPath(result, "Response.ResultVideoUrl") != nil || valueAtPath(result, "Response.ResultImages") != nil || result["audio_url"] != nil || result["urls"] != nil
}
func normalizeProviderTaskResult(request Request, spec providerTaskSpec, result map[string]any, upstreamTaskID string) map[string]any {
out := cloneMapAny(result)
out["status"] = "success"
func normalizeProviderTaskResult(request Request, _ providerTaskSpec, result map[string]any, upstreamTaskID string) map[string]any {
data, ok := result["data"].([]any)
if !ok {
data = providerTaskData(request, result)
}
out := map[string]any{
"status": "success",
"created": time.Now().UnixMilli(),
"model": request.Model,
"data": data,
}
if upstreamTaskID != "" {
out["upstream_task_id"] = upstreamTaskID
}
if out["created"] == nil {
out["created"] = time.Now().UnixMilli()
}
if out["model"] == nil {
out["model"] = request.Model
}
if _, ok := out["data"].([]any); !ok {
if out["data"] != nil {
out["raw_data"] = out["data"]
}
out["data"] = providerTaskData(request, result)
}
return out
}
+11 -9
View File
@@ -1100,9 +1100,11 @@ func volcesTaskErrorMessage(result map[string]any) string {
}
func volcesVideoSuccessResult(request Request, upstreamTaskID string, raw map[string]any) map[string]any {
result := cloneMapAny(raw)
if result == nil {
result = map[string]any{}
result := map[string]any{}
for _, key := range []string{"seed", "resolution", "ratio", "duration", "frames", "framespersecond"} {
if raw[key] != nil {
result[key] = raw[key]
}
}
content, _ := raw["content"].(map[string]any)
videoURL := strings.TrimSpace(stringFromAny(content["video_url"]))
@@ -1112,18 +1114,18 @@ func volcesVideoSuccessResult(request Request, upstreamTaskID string, raw map[st
}
data := []any{}
if videoURL != "" {
data = append(data, map[string]any{"url": videoURL, "type": "video"})
item := map[string]any{"url": videoURL, "type": "video"}
if lastFrameURL := strings.TrimSpace(stringFromAny(content["last_frame_url"])); lastFrameURL != "" {
item["last_frame_url"] = lastFrameURL
}
data = append(data, item)
}
result["id"] = firstNonEmpty(stringFromAny(raw["id"]), upstreamTaskID)
if strings.TrimSpace(stringFromAny(result["model"])) == "" {
result["model"] = upstreamModelName(request.Candidate)
}
result["model"] = firstNonEmpty(stringFromAny(raw["model"]), upstreamModelName(request.Candidate))
result["status"] = "succeeded"
result["object"] = "video.generation"
result["created"] = created
result["upstream_task_id"] = upstreamTaskID
result["data"] = data
result["raw"] = cloneMapAny(raw)
return result
}
@@ -79,7 +79,7 @@ func TestVolcesClientSupportsDeyunEnvelope(t *testing.T) {
}
data, _ := response.Result["data"].([]any)
item, _ := data[0].(map[string]any)
if response.Result["upstream_task_id"] != "deyun-task-1" || item["url"] != "https://example.com/deyun.mp4" {
if response.Result["upstream_task_id"] != nil || item["url"] != "https://example.com/deyun.mp4" {
t.Fatalf("unexpected response: %+v", response.Result)
}
}
+32 -4
View File
@@ -41,8 +41,13 @@ type Config struct {
LocalTempAssetTTLHours int
TaskProgressCallbackEnabled bool
TaskProgressCallbackURL string
TaskProgressCallbackTimeoutMS string
TaskProgressCallbackMaxAttempts string
TaskProgressCallbackTimeoutMS int
TaskProgressCallbackMaxAttempts int
TaskCleanupEnabled bool
TaskRetentionDays int
TaskAnalysisRetentionDays int
TaskCleanupIntervalSeconds int
TaskCleanupBatchSize int
CORSAllowedOrigin string
GlobalHTTPProxy string
GlobalHTTPProxySource string
@@ -87,8 +92,13 @@ func Load() Config {
TaskProgressCallbackURL: env("TASK_PROGRESS_CALLBACK_URL",
strings.TrimRight(env("SERVER_MAIN_BASE_URL", "http://localhost:3000"), "/")+"/internal/platform/task-progress-callbacks",
),
TaskProgressCallbackTimeoutMS: env("TASK_PROGRESS_CALLBACK_TIMEOUT_MS", "5000"),
TaskProgressCallbackMaxAttempts: env("TASK_PROGRESS_CALLBACK_MAX_ATTEMPTS", "10"),
TaskProgressCallbackTimeoutMS: envIntValidated("TASK_PROGRESS_CALLBACK_TIMEOUT_MS", 5000),
TaskProgressCallbackMaxAttempts: envIntValidated("TASK_PROGRESS_CALLBACK_MAX_ATTEMPTS", 10),
TaskCleanupEnabled: env("AI_GATEWAY_TASK_CLEANUP_ENABLED", "false") == "true",
TaskRetentionDays: envIntValidated("AI_GATEWAY_TASK_RETENTION_DAYS", 30),
TaskAnalysisRetentionDays: envIntValidated("AI_GATEWAY_TASK_ANALYSIS_RETENTION_DAYS", 7),
TaskCleanupIntervalSeconds: envIntValidated("AI_GATEWAY_TASK_CLEANUP_INTERVAL_SECONDS", 300),
TaskCleanupBatchSize: envIntValidated("AI_GATEWAY_TASK_CLEANUP_BATCH_SIZE", 1000),
CORSAllowedOrigin: env("CORS_ALLOWED_ORIGIN", "http://localhost:5178,http://127.0.0.1:5178"),
GlobalHTTPProxy: globalProxy.HTTPProxy,
GlobalHTTPProxySource: globalProxy.Source,
@@ -111,6 +121,24 @@ func (c Config) Validate() error {
if c.AsyncWorkerRefreshIntervalSeconds < 1 {
return errors.New("AI_GATEWAY_ASYNC_WORKER_REFRESH_INTERVAL_SECONDS must be positive")
}
if c.TaskProgressCallbackTimeoutMS != 0 && (c.TaskProgressCallbackTimeoutMS < 100 || c.TaskProgressCallbackTimeoutMS > 60000) {
return errors.New("TASK_PROGRESS_CALLBACK_TIMEOUT_MS must be between 100 and 60000")
}
if c.TaskProgressCallbackMaxAttempts != 0 && (c.TaskProgressCallbackMaxAttempts < 1 || c.TaskProgressCallbackMaxAttempts > 100) {
return errors.New("TASK_PROGRESS_CALLBACK_MAX_ATTEMPTS must be between 1 and 100")
}
if c.TaskRetentionDays != 0 && (c.TaskRetentionDays < 1 || c.TaskRetentionDays > 3650) {
return errors.New("AI_GATEWAY_TASK_RETENTION_DAYS must be between 1 and 3650")
}
if c.TaskAnalysisRetentionDays != 0 && (c.TaskAnalysisRetentionDays < 1 || (c.TaskRetentionDays != 0 && c.TaskAnalysisRetentionDays > c.TaskRetentionDays)) {
return errors.New("AI_GATEWAY_TASK_ANALYSIS_RETENTION_DAYS must be between 1 and AI_GATEWAY_TASK_RETENTION_DAYS")
}
if c.TaskCleanupIntervalSeconds != 0 && c.TaskCleanupIntervalSeconds < 60 {
return errors.New("AI_GATEWAY_TASK_CLEANUP_INTERVAL_SECONDS must be at least 60")
}
if c.TaskCleanupBatchSize != 0 && (c.TaskCleanupBatchSize < 100 || c.TaskCleanupBatchSize > 5000) {
return errors.New("AI_GATEWAY_TASK_CLEANUP_BATCH_SIZE must be between 100 and 5000")
}
switch strings.ToLower(strings.TrimSpace(c.IdentitySecretStore)) {
case "":
case "file":
+19
View File
@@ -81,3 +81,22 @@ func TestValidateAsyncWorkerSettings(t *testing.T) {
t.Fatalf("Validate() error = %v, want invalid non-integer hard limit", err)
}
}
func TestValidateTaskHistorySettings(t *testing.T) {
cfg := Load()
cfg.TaskRetentionDays = 30
cfg.TaskAnalysisRetentionDays = 31
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "ANALYSIS_RETENTION") {
t.Fatalf("Validate() error = %v, want invalid analysis retention", err)
}
cfg.TaskAnalysisRetentionDays = 7
cfg.TaskCleanupIntervalSeconds = 59
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "CLEANUP_INTERVAL") {
t.Fatalf("Validate() error = %v, want invalid cleanup interval", err)
}
cfg.TaskCleanupIntervalSeconds = 300
cfg.TaskCleanupBatchSize = 5001
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "CLEANUP_BATCH") {
t.Fatalf("Validate() error = %v, want invalid cleanup batch", err)
}
}
@@ -2029,6 +2029,22 @@ CREATE TABLE IF NOT EXISTS schema_migrations (
if err != nil {
t.Fatalf("read migration %s: %v", filepath.Base(migrationPath), 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", filepath.Base(migrationPath), err)
}
}
}
if _, err := pool.Exec(ctx, "INSERT INTO schema_migrations(version) VALUES($1)", version); err != nil {
t.Fatalf("record non-transaction migration %s: %v", filepath.Base(migrationPath), err)
}
continue
}
tx, err := pool.Begin(ctx)
if err != nil {
t.Fatalf("begin migration %s: %v", filepath.Base(migrationPath), err)
@@ -202,15 +202,11 @@ func (s *Server) createKelingOmniVideo(w http.ResponseWriter, r *http.Request) {
writeKelingCompatError(w, requestID, kelingCompatGatewayError(createErr))
return
}
if kelingTaskUsesNativeProtocol(task) && len(task.CompatibilitySubmitBody) > 0 {
writeJSON(w, task.CompatibilitySubmitHTTPStatus, task.CompatibilitySubmitBody)
return
}
writeJSON(w, http.StatusOK, KelingCompatibleEnvelope{
Code: 0,
Message: "SUCCEED",
RequestID: requestID,
Data: kelingCompatTaskData(task),
Data: kelingCompatSubmissionData(task),
})
}
@@ -251,16 +247,6 @@ func (s *Server) getKelingOmniVideo(w http.ResponseWriter, r *http.Request) {
writeKelingCompatError(w, requestID, newKelingCompatError(http.StatusNotFound, 1203, "task not found"))
return
}
if kelingTaskUsesNativeProtocol(task) {
if raw, ok := task.Result["raw"].(map[string]any); ok && len(raw) > 0 {
writeJSON(w, http.StatusOK, raw)
return
}
if len(task.RemoteTaskPayload) > 0 {
writeJSON(w, http.StatusOK, task.RemoteTaskPayload)
return
}
}
writeJSON(w, http.StatusOK, KelingCompatibleEnvelope{
Code: 0,
Message: "SUCCEED",
@@ -705,15 +691,22 @@ func kelingCompatTaskData(task store.GatewayTask) map[string]any {
return data
}
func kelingCompatPublicID(task store.GatewayTask) string {
if publicID := strings.TrimSpace(task.CompatibilityPublicID); publicID != "" {
return publicID
}
return task.ID
func kelingCompatSubmissionData(task store.GatewayTask) map[string]any {
data := kelingCompatTaskData(task)
data["task_status"] = "submitted"
delete(data, "task_status_code")
delete(data, "task_status_msg")
delete(data, "task_result")
return data
}
func kelingTaskUsesNativeProtocol(task store.GatewayTask) bool {
return task.CompatibilityProtocol == clients.ProtocolKlingV1Omni && task.CompatibilitySourceProtocol == clients.ProtocolKlingV1Omni
func kelingCompatPublicID(task store.GatewayTask) string {
if task.CompatibilityProtocol == task.CompatibilitySourceProtocol {
if remoteTaskID := strings.TrimSpace(task.RemoteTaskID); remoteTaskID != "" {
return remoteTaskID
}
}
return task.ID
}
func kelingCompatTaskStatus(status string) string {
@@ -209,6 +209,10 @@ func TestKelingCompatTaskDataAndOwnership(t *testing.T) {
if video["id"] != "video-1" || video["watermark_url"] != "https://example.com/watermarked.mp4" || video["duration"] != "5" {
t.Fatalf("unexpected compatible video: %+v", video)
}
submission := kelingCompatSubmissionData(task)
if submission["task_status"] != "submitted" || submission["task_result"] != nil {
t.Fatalf("submission response must be stable and minimal: %+v", submission)
}
}
func TestRequireKelingAPIKeyWritesOfficialAuthEnvelope(t *testing.T) {
+22 -30
View File
@@ -177,10 +177,10 @@ func (s *Server) createKlingCompatTask(w http.ResponseWriter, r *http.Request, v
w.Header().Set("Idempotent-Replayed", "true")
}
if version == "v2" {
writeJSON(w, http.StatusOK, klingV2Envelope(task))
writeJSON(w, http.StatusOK, klingV2SubmissionEnvelope(task))
return
}
writeJSON(w, http.StatusOK, klingV1Envelope(task))
writeJSON(w, http.StatusOK, klingV1SubmissionEnvelope(task))
}
func klingCompatTaskBody(version string, model string, native map[string]any) (map[string]any, string, error) {
@@ -612,20 +612,17 @@ func (s *Server) klingV2ListTasks(w http.ResponseWriter, r *http.Request) {
}
func klingV1Envelope(task store.GatewayTask) map[string]any {
if klingTaskUsesNativeV1Protocol(task) {
if raw := klingMap(task.Result["raw"]); len(raw) > 0 {
return raw
}
if len(task.RemoteTaskPayload) > 0 {
return task.RemoteTaskPayload
}
if raw := klingMap(task.CompatibilitySubmitBody); len(raw) > 0 {
return raw
}
}
return map[string]any{"code": 0, "message": "success", "request_id": firstNonEmpty(task.RequestID, task.ID), "data": klingV1TaskData(task)}
}
func klingV1SubmissionEnvelope(task store.GatewayTask) map[string]any {
data := klingV1TaskData(task)
data["task_status"] = "submitted"
delete(data, "task_status_msg")
delete(data, "task_result")
return map[string]any{"code": 0, "message": "success", "request_id": firstNonEmpty(task.RequestID, task.ID), "data": data}
}
func klingV1TaskData(task store.GatewayTask) map[string]any {
data := map[string]any{
"task_id": klingCompatibilityPublicID(task), "task_status": klingV1Status(task.Status),
@@ -646,20 +643,17 @@ func klingV1TaskData(task store.GatewayTask) map[string]any {
}
func klingV2Envelope(task store.GatewayTask) map[string]any {
if task.CompatibilityProtocol == clients.ProtocolKlingV2Omni && task.CompatibilitySourceProtocol == clients.ProtocolKlingV2Omni {
if raw := klingMap(task.Result["raw"]); len(raw) > 0 {
return raw
}
if len(task.RemoteTaskPayload) > 0 {
return task.RemoteTaskPayload
}
if len(task.CompatibilitySubmitBody) > 0 {
return task.CompatibilitySubmitBody
}
}
return map[string]any{"code": 0, "message": "success", "request_id": firstNonEmpty(task.RequestID, task.ID), "data": klingV2TaskData(task)}
}
func klingV2SubmissionEnvelope(task store.GatewayTask) map[string]any {
data := klingV2TaskData(task)
data["status"] = "submitted"
delete(data, "message")
delete(data, "outputs")
return map[string]any{"code": 0, "message": "success", "request_id": firstNonEmpty(task.RequestID, task.ID), "data": data}
}
func klingV2TaskData(task store.GatewayTask) map[string]any {
data := map[string]any{
"id": klingCompatibilityPublicID(task), "status": klingV2Status(task.Status),
@@ -676,16 +670,14 @@ func klingV2TaskData(task store.GatewayTask) map[string]any {
}
func klingCompatibilityPublicID(task store.GatewayTask) string {
if publicID := strings.TrimSpace(task.CompatibilityPublicID); publicID != "" {
return publicID
if task.CompatibilityProtocol == task.CompatibilitySourceProtocol {
if remoteTaskID := strings.TrimSpace(task.RemoteTaskID); remoteTaskID != "" {
return remoteTaskID
}
}
return task.ID
}
func klingTaskUsesNativeV1Protocol(task store.GatewayTask) bool {
return task.CompatibilityProtocol == clients.ProtocolKlingV1Omni && task.CompatibilitySourceProtocol == clients.ProtocolKlingV1Omni
}
func klingMap(value any) map[string]any {
result, _ := value.(map[string]any)
return result
@@ -6,6 +6,9 @@ import (
"net/http/httptest"
"strings"
"testing"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
func TestKlingV1O1CompatibilityBody(t *testing.T) {
@@ -86,6 +89,27 @@ func TestKlingV2CompatibilityBody(t *testing.T) {
}
}
func TestKlingSubmissionEnvelopesDoNotReplayCurrentResult(t *testing.T) {
task := store.GatewayTask{
ID: "gateway-task", Status: "succeeded", RemoteTaskID: "upstream-task",
CompatibilityProtocol: "kling-v1-omni",
CompatibilitySourceProtocol: "kling-v1-omni",
Result: map[string]any{"data": []any{map[string]any{"url": "https://example.invalid/result.mp4"}}},
CreatedAt: time.Unix(100, 0),
UpdatedAt: time.Unix(101, 0),
}
v1 := klingV1SubmissionEnvelope(task)
v1Data, _ := v1["data"].(map[string]any)
if v1Data["task_status"] != "submitted" || v1Data["task_result"] != nil {
t.Fatalf("V1 submission envelope=%#v", v1)
}
v2 := klingV2SubmissionEnvelope(task)
v2Data, _ := v2["data"].(map[string]any)
if v2Data["status"] != "submitted" || v2Data["outputs"] != nil {
t.Fatalf("V2 submission envelope=%#v", v2)
}
}
func TestKlingCompatibilityValidation(t *testing.T) {
tests := []struct {
name string
+1
View File
@@ -122,6 +122,7 @@ func NewServerWithContext(ctx context.Context, cfg config.Config, db *store.Stor
server.auth.LocalAPIKeyVerifier = db.VerifyLocalAPIKey
server.runner.StartAsyncQueueWorker(ctx)
server.runner.StartBillingSettlementWorker(ctx)
server.runner.StartTaskHistoryWorkers(ctx)
server.startLocalTempAssetCleanup(ctx)
server.startOIDCSessionCleanup(ctx)
@@ -223,7 +223,6 @@ func (s *Server) createVolcesCompatibleTask(r *http.Request, user *auth.User, bo
return store.GatewayTask{}, err
}
task.CompatibilityProtocol = clients.ProtocolVolcesContents
task.CompatibilityPublicID = task.ID
if err := s.runner.EnqueueAsyncTask(r.Context(), task); err != nil {
return store.GatewayTask{}, &clients.ClientError{Code: "enqueue_failed", Message: err.Error(), StatusCode: http.StatusServiceUnavailable, Retryable: true}
}
@@ -258,7 +257,6 @@ func (s *Server) waitForCompatibilitySubmission(r *http.Request, task store.Gate
Message: firstNonEmpty(current.ErrorMessage, current.Error, current.Message),
StatusCode: status,
Retryable: false,
Wire: compatibilitySubmissionWire(current),
}
}
select {
@@ -271,34 +269,6 @@ func (s *Server) waitForCompatibilitySubmission(r *http.Request, task store.Gate
}
}
func compatibilitySubmissionWire(task store.GatewayTask) *clients.WireResponse {
if task.CompatibilitySubmitHTTPStatus == 0 || strings.TrimSpace(task.CompatibilitySourceProtocol) == "" {
return nil
}
headers := map[string][]string{}
for name, value := range task.CompatibilitySubmitHeaders {
switch typed := value.(type) {
case []any:
for _, item := range typed {
if text := strings.TrimSpace(volcesCompatString(item)); text != "" {
headers[name] = append(headers[name], text)
}
}
case []string:
headers[name] = append([]string(nil), typed...)
case string:
headers[name] = []string{typed}
}
}
return &clients.WireResponse{
Protocol: task.CompatibilitySourceProtocol,
StatusCode: task.CompatibilitySubmitHTTPStatus,
Headers: headers,
Body: cloneVolcesCompatibleMap(task.CompatibilitySubmitBody),
Converted: task.CompatibilitySourceProtocol != task.CompatibilityProtocol,
}
}
func (s *Server) volcesCompatibleTaskForUser(w http.ResponseWriter, r *http.Request) (store.GatewayTask, bool) {
user, ok := auth.UserFromContext(r.Context())
if !ok || user == nil {
@@ -327,15 +297,30 @@ func isVolcesCompatibleTask(task store.GatewayTask) bool {
}
func volcesCompatibleTask(task store.GatewayTask) map[string]any {
if volcesTaskUsesNativeProtocol(task) {
if raw := cloneVolcesCompatibleMap(mapFromVolcesCompat(task.Result["raw"])); len(raw) > 0 {
return raw
response := map[string]any{}
for _, key := range []string{"model", "seed", "resolution", "ratio", "duration", "frames", "framespersecond"} {
if task.Result[key] != nil {
response[key] = task.Result[key]
}
if raw := cloneVolcesCompatibleMap(task.RemoteTaskPayload); len(raw) > 0 {
return raw
}
if content, ok := task.Result["content"].(map[string]any); ok && len(content) > 0 {
// Transitional read support for historical canonical results. New
// results only persist data[] and reconstruct this protocol field.
response["content"] = content
} else if data, ok := task.Result["data"].([]any); ok && len(data) > 0 {
if item, ok := data[0].(map[string]any); ok {
content := map[string]any{}
if videoURL := firstNonEmpty(volcesCompatString(item["url"]), volcesCompatString(item["video_url"])); videoURL != "" {
content["video_url"] = videoURL
}
if lastFrameURL := volcesCompatString(item["last_frame_url"]); lastFrameURL != "" {
content["last_frame_url"] = lastFrameURL
}
if len(content) > 0 {
response["content"] = content
}
}
}
response := map[string]any{}
response["id"] = volcesCompatiblePublicID(task)
response["model"] = firstNonEmpty(volcesCompatString(response["model"]), task.Model)
response["status"] = volcesCompatibleTaskStatus(task.Status)
@@ -346,8 +331,10 @@ func volcesCompatibleTask(task store.GatewayTask) map[string]any {
response[key] = task.Request[key]
}
}
if len(task.Usage) > 0 && response["usage"] == nil {
response["usage"] = task.Usage
if usage := volcesCompatibleUsage(task.Usage); len(usage) > 0 {
response["usage"] = usage
} else if legacyUsage, ok := task.Result["usage"].(map[string]any); ok && len(legacyUsage) > 0 {
response["usage"] = legacyUsage
}
if task.Status == "failed" || task.Status == "cancelled" {
response["error"] = map[string]any{"code": firstNonEmpty(task.ErrorCode, strings.ToUpper(task.Status)), "message": firstNonEmpty(task.ErrorMessage, task.Error, task.Message)}
@@ -355,19 +342,28 @@ func volcesCompatibleTask(task store.GatewayTask) map[string]any {
return response
}
func volcesCompatibleCreateResponse(task store.GatewayTask) map[string]any {
if volcesTaskUsesNativeProtocol(task) {
if raw := cloneVolcesCompatibleMap(task.RemoteTaskPayload); len(raw) > 0 {
return raw
func volcesCompatibleUsage(usage map[string]any) map[string]any {
out := map[string]any{}
for outputKey, inputKeys := range map[string][]string{
"prompt_tokens": {"inputTokens", "promptTokens"},
"completion_tokens": {"outputTokens", "completionTokens"},
"total_tokens": {"totalTokens"},
} {
for _, inputKey := range inputKeys {
if value := usage[inputKey]; value != nil {
out[outputKey] = value
break
}
}
}
return out
}
func volcesCompatibleCreateResponse(task store.GatewayTask) map[string]any {
return map[string]any{"id": volcesCompatiblePublicID(task)}
}
func volcesCompatiblePublicID(task store.GatewayTask) string {
if strings.TrimSpace(task.CompatibilityPublicID) != "" {
return strings.TrimSpace(task.CompatibilityPublicID)
}
if volcesTaskUsesNativeProtocol(task) && strings.TrimSpace(task.RemoteTaskID) != "" {
return strings.TrimSpace(task.RemoteTaskID)
}
@@ -386,11 +382,6 @@ func volcesTaskUsesNativeProtocol(task store.GatewayTask) bool {
return false
}
func mapFromVolcesCompat(value any) map[string]any {
result, _ := value.(map[string]any)
return result
}
func volcesCompatibleTaskStatus(status string) string {
switch strings.ToLower(strings.TrimSpace(status)) {
case "succeeded", "success", "completed":
@@ -406,21 +397,6 @@ func volcesCompatibleTaskStatus(status string) string {
}
}
func cloneVolcesCompatibleMap(source map[string]any) map[string]any {
if len(source) == 0 {
return nil
}
raw, err := json.Marshal(source)
if err != nil {
return map[string]any{}
}
var out map[string]any
if err := json.Unmarshal(raw, &out); err != nil {
return map[string]any{}
}
return out
}
func writeVolcesCompatibleTaskError(w http.ResponseWriter, err error) {
status := http.StatusInternalServerError
var staged *gatewayTaskCreationError
@@ -7,23 +7,22 @@ import (
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
func TestVolcesCompatibleTaskPreservesNativeOfficialFieldsWithoutGatewayExtensions(t *testing.T) {
func TestVolcesCompatibleTaskBuildsProtocolResponseFromCanonicalResult(t *testing.T) {
now := time.Date(2026, 7, 18, 8, 0, 0, 0, time.UTC)
task := store.GatewayTask{
ID: "gateway-task-1", Kind: "videos.generations", Status: "succeeded", Model: "doubao-seedance-2-0-mini-260615",
RemoteTaskID: "cgt-upstream-1", CreatedAt: now, UpdatedAt: now.Add(time.Second),
Result: map[string]any{
"raw": map[string]any{
"id": "cgt-upstream-1", "model": "doubao-seedance-2-0-mini-260615", "status": "succeeded",
"content": map[string]any{"video_url": "https://example.com/out.mp4"}, "usage": map[string]any{"total_tokens": 9},
"future_official_field": "preserved",
},
"id": "cgt-upstream-1", "model": "doubao-seedance-2-0-mini-260615", "status": "succeeded",
"data": []any{map[string]any{"url": "https://example.com/out.mp4", "type": "video"}},
"future_official_field": "not-whitelisted",
},
Usage: map[string]any{"totalTokens": 9},
Billings: []any{map[string]any{"amount": 3}}, BillingSummary: map[string]any{"currency": "resource"}, FinalChargeAmount: 3,
Attempts: []store.TaskAttempt{{Provider: "volces"}},
}
got := volcesCompatibleTask(task)
if got["id"] != task.RemoteTaskID || got["status"] != "succeeded" || got["future_official_field"] != "preserved" {
if got["id"] != task.RemoteTaskID || got["status"] != "succeeded" || got["future_official_field"] != nil {
t.Fatalf("unexpected compatibility identity/status: %+v", got)
}
content, _ := got["content"].(map[string]any)
+123 -46
View File
@@ -2,10 +2,10 @@ package runner
import (
"context"
"encoding/json"
"errors"
"fmt"
"log/slog"
"net/http"
"os"
"regexp"
"strconv"
@@ -88,6 +88,24 @@ func New(cfg config.Config, db *store.Store, logger *slog.Logger, observers ...b
if cfg.AsyncWorkerRefreshIntervalSeconds == 0 {
cfg.AsyncWorkerRefreshIntervalSeconds = 5
}
if cfg.TaskProgressCallbackTimeoutMS == 0 {
cfg.TaskProgressCallbackTimeoutMS = 5000
}
if cfg.TaskProgressCallbackMaxAttempts == 0 {
cfg.TaskProgressCallbackMaxAttempts = 10
}
if cfg.TaskRetentionDays == 0 {
cfg.TaskRetentionDays = 30
}
if cfg.TaskAnalysisRetentionDays == 0 {
cfg.TaskAnalysisRetentionDays = 7
}
if cfg.TaskCleanupIntervalSeconds == 0 {
cfg.TaskCleanupIntervalSeconds = 300
}
if cfg.TaskCleanupBatchSize == 0 {
cfg.TaskCleanupBatchSize = 1000
}
httpClients := newHTTPClientCache()
scriptExecutor := &scriptengine.Executor{Logger: logger}
service := &Service{
@@ -130,6 +148,15 @@ func (s *Service) observeBillingEvent(event string) {
}
}
func (s *Service) observeTaskEventSkip(reason string) {
observer, ok := s.billingMetrics.(interface {
ObserveTaskEventSkip(string)
})
if ok {
observer.ObserveTaskEventSkip(reason)
}
}
func (s *Service) Execute(ctx context.Context, task store.GatewayTask, user *auth.User) (Result, error) {
return s.execute(ctx, task, user, nil)
}
@@ -569,6 +596,7 @@ candidatesLoop:
}
if err == nil {
attemptNo = nextAttemptNo
response.Result = canonicalTaskResult(response.Result)
var billings []any
finalAmount := fixedAmount(0)
finalAmountText := ""
@@ -1058,12 +1086,13 @@ func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user
if strings.TrimSpace(remoteTaskID) == "" {
return nil
}
if err := s.store.SetTaskRemoteTask(context.WithoutCancel(ctx), task.ID, task.ExecutionToken, attemptID, remoteTaskID, payload); err != nil {
checkpoint := minimalRemoteTaskCheckpoint(candidate.Provider, candidate.SpecType, payload)
if err := s.store.SetTaskRemoteTask(context.WithoutCancel(ctx), task.ID, task.ExecutionToken, attemptID, remoteTaskID, checkpoint); err != nil {
return err
}
task.RemoteTaskID = remoteTaskID
task.RemoteTaskPayload = payload
if err := s.persistCompatibilitySubmission(context.WithoutCancel(ctx), task, candidate, remoteTaskID, payload, submissionWire); err != nil {
task.RemoteTaskPayload = checkpoint
if err := s.persistCompatibilitySubmission(context.WithoutCancel(ctx), task, candidate, remoteTaskID, checkpoint, submissionWire); err != nil {
return err
}
return setSubmissionStatus("response_received")
@@ -1072,11 +1101,12 @@ func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user
if strings.TrimSpace(remoteTaskID) == "" {
return nil
}
if err := s.store.SetTaskRemoteTask(context.WithoutCancel(ctx), task.ID, task.ExecutionToken, attemptID, remoteTaskID, payload); err != nil {
checkpoint := minimalRemoteTaskCheckpoint(candidate.Provider, candidate.SpecType, payload)
if err := s.store.SetTaskRemoteTask(context.WithoutCancel(ctx), task.ID, task.ExecutionToken, attemptID, remoteTaskID, checkpoint); err != nil {
return err
}
task.RemoteTaskID = remoteTaskID
task.RemoteTaskPayload = payload
task.RemoteTaskPayload = checkpoint
return setSubmissionStatus("response_received")
},
OnUpstreamSubmissionStarted: func() error {
@@ -1258,6 +1288,59 @@ func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user
return response, nil
}
func minimalRemoteTaskCheckpoint(provider string, specType string, payload map[string]any) map[string]any {
const maxBytes = 8192
provider = strings.ToLower(strings.TrimSpace(provider))
specType = strings.ToLower(strings.TrimSpace(specType))
allowedByProvider := map[string][]string{
"keling": {"endpoint", "pollEndpoint", "mode", "taskType", "cleanupElementIds"},
"kling": {"endpoint", "pollEndpoint", "mode", "taskType", "cleanupElementIds"},
"topaz": {"phase", "targetResolution"},
"vectorizer": {"imageToken", "receipt"},
}
allowedKeys := allowedByProvider[provider]
if len(allowedKeys) == 0 {
allowedKeys = allowedByProvider[specType]
}
checkpoint := map[string]any{}
for _, key := range allowedKeys {
value, ok := payload[key]
if !ok || value == nil {
continue
}
if text, ok := value.(string); ok {
value = truncateText(text, 4096)
}
next := cloneMap(checkpoint)
next[key] = value
encoded, err := json.Marshal(next)
if err != nil || len(encoded) > maxBytes {
continue
}
checkpoint = next
}
return checkpoint
}
func canonicalTaskResult(result map[string]any) map[string]any {
canonical := cloneMap(result)
for _, key := range []string{
"raw",
"raw_data",
"rawData",
"provider_response",
"providerResponse",
"submit",
"file_retrieve",
"fileRetrieve",
"upstream_task_id",
"remote_task_id",
} {
delete(canonical, key)
}
return canonical
}
func (s *Service) persistCompatibilitySubmission(ctx context.Context, task store.GatewayTask, candidate store.RuntimeModelCandidate, remoteTaskID string, payload map[string]any, wire *clients.WireResponse) error {
targetProtocol := strings.TrimSpace(stringFromMap(task.Request, "_gateway_target_protocol"))
if targetProtocol == "" {
@@ -1267,32 +1350,9 @@ func (s *Service) persistCompatibilitySubmission(ctx context.Context, task store
if wire != nil && strings.TrimSpace(wire.Protocol) != "" {
sourceProtocol = strings.TrimSpace(wire.Protocol)
}
publicID := task.ID
if sourceProtocol == targetProtocol && strings.TrimSpace(remoteTaskID) != "" {
publicID = strings.TrimSpace(remoteTaskID)
}
submitBody := payload
if nested, ok := payload["submit"].(map[string]any); ok && len(nested) > 0 {
submitBody = nested
}
httpStatus := http.StatusOK
headers := map[string]any{}
if wire != nil {
httpStatus = wire.StatusCode
if len(wire.Body) > 0 {
submitBody = wire.Body
}
for name, values := range wire.Headers {
headers[name] = append([]string(nil), values...)
}
}
return s.store.SetTaskCompatibilitySubmission(ctx, task.ID, store.CompatibilitySubmission{
TargetProtocol: targetProtocol,
PublicID: publicID,
SourceProtocol: sourceProtocol,
HTTPStatus: httpStatus,
Headers: headers,
Body: submitBody,
})
}
@@ -1327,9 +1387,10 @@ func compatibilitySourceProtocol(kind string, candidate store.RuntimeModelCandid
}
func (s *Service) recordTaskParameterPreprocessing(ctx context.Context, taskID string, attemptID string, attemptNo int, candidate store.RuntimeModelCandidate, log parameterPreprocessingLog) error {
if skipTaskParameterPreprocessingLog(log.ModelType) && !log.Changed {
if !log.Changed {
return nil
}
changes := minimalParameterPreprocessingChanges(log.Changes)
_, err := s.store.CreateTaskParamPreprocessingLog(ctx, store.CreateTaskParamPreprocessingLogInput{
TaskID: taskID,
AttemptID: attemptID,
@@ -1338,16 +1399,31 @@ func (s *Service) recordTaskParameterPreprocessing(ctx context.Context, taskID s
PlatformID: candidate.PlatformID,
PlatformModelID: candidate.PlatformModelID,
ClientID: candidate.ClientID,
Changed: log.Changed,
Changed: true,
ChangeCount: len(log.Changes),
ActualInput: log.Input,
ConvertedOutput: log.Output,
Changes: log.Changes,
ModelSnapshot: log.Model,
Changes: changes,
})
return err
}
func minimalParameterPreprocessingChanges(changes []parameterPreprocessChange) []any {
const maxBytes = 1024
out := make([]any, 0, len(changes))
for _, change := range changes {
item := map[string]any{
"path": truncateText(change.Path, 256),
"action": truncateText(change.Action, 64),
}
next := append(append([]any(nil), out...), item)
encoded, err := json.Marshal(next)
if err != nil || len(encoded) > maxBytes {
break
}
out = next
}
return out
}
func skipTaskParameterPreprocessingLog(modelType string) bool {
switch strings.TrimSpace(modelType) {
case "text_generate", "text_embedding", "text_rerank", "chat", "responses", "text":
@@ -1357,6 +1433,15 @@ func skipTaskParameterPreprocessingLog(modelType string) bool {
}
}
func truncateText(value string, maxRunes int) string {
value = strings.TrimSpace(value)
runes := []rune(value)
if len(runes) <= maxRunes {
return value
}
return string(runes[:maxRunes])
}
func (s *Service) clientFor(candidate store.RuntimeModelCandidate, simulated bool) clients.Client {
if simulated {
return s.clients["simulation"]
@@ -1628,17 +1713,6 @@ func (s *Service) requeueInterruptedAsyncTask(ctx context.Context, task store.Ga
}
func (s *Service) withAttemptHistory(ctx context.Context, taskID string, metrics map[string]any) map[string]any {
attempts, err := s.store.ListTaskAttempts(ctx, taskID)
if err != nil {
s.logger.Warn("list task attempts for metrics failed", "taskID", taskID, "error", err)
return metrics
}
if len(attempts) == 0 {
return metrics
}
metrics = mergeMetrics(metrics)
metrics["attemptCount"] = len(attempts)
metrics["attempts"] = summarizeAttempts(attempts)
return metrics
}
@@ -1647,6 +1721,9 @@ func (s *Service) emit(ctx context.Context, taskID string, eventType string, sta
if err != nil {
return err
}
if event.ID == "" {
s.observeTaskEventSkip(event.SkippedReason)
}
if s.cfg.TaskProgressCallbackEnabled {
return s.store.QueueTaskCallback(ctx, event, s.cfg.TaskProgressCallbackURL)
}
@@ -0,0 +1,181 @@
package runner
import (
"bytes"
"context"
"encoding/json"
"fmt"
"hash/fnv"
"io"
"net/http"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
"github.com/google/uuid"
)
func (s *Service) StartTaskHistoryWorkers(ctx context.Context) {
if s.cfg.TaskProgressCallbackEnabled && s.cfg.TaskProgressCallbackURL != "" {
go s.runTaskCallbackWorker(ctx, "callback-"+uuid.NewString())
}
if s.cfg.TaskCleanupEnabled {
go s.runTaskHistoryCleanupWorker(ctx)
}
}
func (s *Service) runTaskCallbackWorker(ctx context.Context, workerID string) {
ticker := time.NewTicker(time.Second)
defer ticker.Stop()
client := &http.Client{Timeout: time.Duration(s.cfg.TaskProgressCallbackTimeoutMS) * time.Millisecond}
for {
s.processTaskCallbackBatch(ctx, client, workerID)
select {
case <-ctx.Done():
return
case <-ticker.C:
}
}
}
func (s *Service) processTaskCallbackBatch(ctx context.Context, client *http.Client, workerID string) {
items, err := s.store.ClaimTaskCallbacks(ctx, workerID, store.TaskCallbackBatchSize, store.TaskCallbackLockTTL)
if err != nil {
if ctx.Err() == nil {
s.logger.Error("claim task callbacks failed", "error_category", "task_callback_claim_failed")
}
return
}
for _, item := range items {
if ctx.Err() != nil {
return
}
statusCode, deliveryErr := deliverTaskCallback(ctx, client, item, s.cfg.ServerMainInternalToken)
if deliveryErr == nil && statusCode >= 200 && statusCode < 300 {
if err := s.store.MarkTaskCallbackDelivered(context.WithoutCancel(ctx), item); err != nil {
s.logger.Error("mark task callback delivered failed", "taskID", item.TaskID, "seq", item.Seq, "error_category", "task_callback_state_failed")
}
continue
}
retryable := deliveryErr != nil || statusCode == http.StatusRequestTimeout || statusCode == http.StatusTooManyRequests || statusCode >= 500
retry := retryable && item.Attempts < s.cfg.TaskProgressCallbackMaxAttempts
message := "callback_network_error"
if deliveryErr == nil {
message = fmt.Sprintf("callback_http_%d", statusCode)
}
nextAttemptAt := time.Now().Add(taskCallbackRetryDelay(item.ID, item.Attempts))
if err := s.store.MarkTaskCallbackFailed(context.WithoutCancel(ctx), item, retry, nextAttemptAt, message); err != nil {
s.logger.Error("mark task callback failed", "taskID", item.TaskID, "seq", item.Seq, "error_category", "task_callback_state_failed")
continue
}
if !retry {
s.logger.Warn("task callback moved to dead letter", "taskID", item.TaskID, "seq", item.Seq, "statusCode", statusCode, "error_category", "task_callback_failed")
}
}
}
func deliverTaskCallback(ctx context.Context, client *http.Client, item store.TaskCallbackDelivery, bearerToken string) (int, error) {
body, err := json.Marshal(map[string]any{
"taskId": item.TaskID,
"seq": item.Seq,
"eventType": item.EventType,
"status": item.TaskStatus,
"createdAt": item.CreatedAt.UTC().Format(time.RFC3339Nano),
})
if err != nil {
return 0, err
}
request, err := http.NewRequestWithContext(ctx, http.MethodPost, item.CallbackURL, bytes.NewReader(body))
if err != nil {
return 0, err
}
request.Header.Set("Content-Type", "application/json")
request.Header.Set("Idempotency-Key", fmt.Sprintf("%s:%d", item.TaskID, item.Seq))
request.Header.Set("X-EasyAI-Event-Type", item.EventType)
if bearerToken != "" {
request.Header.Set("Authorization", "Bearer "+bearerToken)
}
response, err := client.Do(request)
if err != nil {
return 0, err
}
defer response.Body.Close()
_, _ = io.Copy(io.Discard, io.LimitReader(response.Body, 4096))
return response.StatusCode, nil
}
func taskCallbackRetryDelay(deliveryID string, attempt int) time.Duration {
delays := []time.Duration{
5 * time.Second,
30 * time.Second,
2 * time.Minute,
10 * time.Minute,
30 * time.Minute,
}
index := attempt - 1
if index < 0 {
index = 0
}
if index >= len(delays) {
index = len(delays) - 1
}
base := delays[index]
hash := fnv.New32a()
_, _ = hash.Write([]byte(deliveryID))
jitter := time.Duration(hash.Sum32()%21) * base / 100
return base + jitter
}
func (s *Service) runTaskHistoryCleanupWorker(ctx context.Context) {
interval := time.Duration(s.cfg.TaskCleanupIntervalSeconds) * time.Second
ticker := time.NewTicker(interval)
defer ticker.Stop()
for {
s.processTaskHistoryCleanup(ctx)
select {
case <-ctx.Done():
return
case <-ticker.C:
}
}
}
func (s *Service) processTaskHistoryCleanup(ctx context.Context) {
for batch := 0; batch < 10; batch++ {
now := time.Now().UTC()
result, err := s.store.CleanupTaskHistory(
ctx,
now.AddDate(0, 0, -s.cfg.TaskAnalysisRetentionDays),
now.AddDate(0, 0, -s.cfg.TaskRetentionDays),
s.cfg.TaskCleanupBatchSize,
)
if err != nil {
if ctx.Err() == nil {
s.logger.Error("cleanup task history failed", "error_category", "task_history_cleanup_failed")
}
return
}
if result.Total() == 0 {
return
}
s.logger.Info("task history cleanup batch completed",
"compactedTasks", result.CompactedTasks,
"compactedCheckpoints", result.CompactedCheckpoints,
"compactedAttempts", result.CompactedAttempts,
"compactedEvents", result.CompactedEvents,
"compactedCallbacks", result.CompactedCallbacks,
"retiredLegacyCallbacks", result.RetiredCallbacks,
"compactedParamLogs", result.CompactedParamLogs,
"compactedClonedVoices", result.CompactedClonedVoices,
"deletedCallbacks", result.DeletedCallbacks,
"deletedParamLogs", result.DeletedParamLogs,
"deletedEvents", result.DeletedEvents,
"deletedAttempts", result.DeletedAttempts,
"deletedTasks", result.DeletedTasks,
)
select {
case <-ctx.Done():
return
case <-time.After(100 * time.Millisecond):
}
}
}
@@ -0,0 +1,118 @@
package runner
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"testing"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
func TestMinimalRemoteTaskCheckpointUsesAllowlist(t *testing.T) {
checkpoint := minimalRemoteTaskCheckpoint("keling", "", map[string]any{
"endpoint": "/videos/task",
"mode": "omni_video",
"cleanupElementIds": []any{"element-1"},
"phase": "uploaded",
"submit": map[string]any{"large": "provider response"},
"payload": map[string]any{"prompt": "duplicate request"},
"raw": map[string]any{"data": "duplicate result"},
"status": "processing",
})
if checkpoint["endpoint"] != "/videos/task" || checkpoint["mode"] != "omni_video" || checkpoint["cleanupElementIds"] == nil {
t.Fatalf("checkpoint lost required fields: %#v", checkpoint)
}
for _, forbidden := range []string{"phase", "submit", "payload", "raw", "status"} {
if _, ok := checkpoint[forbidden]; ok {
t.Fatalf("checkpoint contains forbidden field %q: %#v", forbidden, checkpoint)
}
}
if unknown := minimalRemoteTaskCheckpoint("unknown-provider", "", map[string]any{"endpoint": "/forbidden"}); len(unknown) != 0 {
t.Fatalf("unknown provider checkpoint=%#v, want empty", unknown)
}
}
func TestCanonicalTaskResultDropsProviderRawCopy(t *testing.T) {
result := canonicalTaskResult(map[string]any{
"id": "task",
"data": []any{map[string]any{"url": "https://example.com/output"}},
"raw": map[string]any{"provider": "duplicate"},
"raw_data": map[string]any{"provider": "duplicate"},
"provider_response": map[string]any{"provider": "duplicate"},
"upstream_task_id": "remote-task",
"compatibility_data": "kept",
})
if result["raw"] != nil || result["raw_data"] != nil || result["provider_response"] != nil ||
result["upstream_task_id"] != nil || result["id"] != "task" || result["data"] == nil ||
result["compatibility_data"] != "kept" {
t.Fatalf("canonical result=%#v", result)
}
}
func TestMinimalParameterPreprocessingChangesDropsValuesAndCapsSize(t *testing.T) {
changes := make([]parameterPreprocessChange, 0, 100)
for index := 0; index < 100; index++ {
changes = append(changes, parameterPreprocessChange{
Path: "input.image",
Action: "converted",
Before: map[string]any{"base64": "forbidden"},
After: map[string]any{"url": "forbidden"},
})
}
minimal := minimalParameterPreprocessingChanges(changes)
encoded, err := json.Marshal(minimal)
if err != nil {
t.Fatal(err)
}
if len(encoded) > 1024 {
t.Fatalf("minimal changes size=%d, want <=1024", len(encoded))
}
for _, raw := range minimal {
item := raw.(map[string]any)
if len(item) != 2 || item["path"] != "input.image" || item["action"] != "converted" {
t.Fatalf("unexpected minimal change: %#v", item)
}
}
}
func TestDeliverTaskCallbackSendsMinimalAuthenticatedBody(t *testing.T) {
createdAt := time.Date(2026, time.July, 24, 10, 0, 0, 0, time.UTC)
var received map[string]any
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) {
if request.Header.Get("Authorization") != "Bearer internal-token" {
t.Errorf("Authorization=%q", request.Header.Get("Authorization"))
}
if request.Header.Get("Idempotency-Key") != "task-id:3" {
t.Errorf("Idempotency-Key=%q", request.Header.Get("Idempotency-Key"))
}
if request.Header.Get("X-EasyAI-Event-Type") != "task.completed" {
t.Errorf("X-EasyAI-Event-Type=%q", request.Header.Get("X-EasyAI-Event-Type"))
}
if err := json.NewDecoder(request.Body).Decode(&received); err != nil {
t.Error(err)
}
w.WriteHeader(http.StatusNoContent)
}))
defer server.Close()
status, err := deliverTaskCallback(context.Background(), server.Client(), store.TaskCallbackDelivery{
TaskID: "task-id",
Seq: 3,
CallbackURL: server.URL,
EventType: "task.completed",
TaskStatus: "succeeded",
CreatedAt: createdAt,
}, "internal-token")
if err != nil || status != http.StatusNoContent {
t.Fatalf("status=%d err=%v", status, err)
}
if len(received) != 5 {
t.Fatalf("callback fields=%#v, want exactly five fields", received)
}
if received["taskId"] != "task-id" || received["eventType"] != "task.completed" || received["status"] != "succeeded" {
t.Fatalf("callback body=%#v", received)
}
}
+1 -10
View File
@@ -167,16 +167,7 @@ func (s *Service) persistVoiceCloneResult(ctx context.Context, task store.Gatewa
DemoAudioURL: demoAudioURL,
Status: "active",
ExpiresAt: &expiresAt,
Metadata: map[string]any{
"request": map[string]any{
"textValidation": body["text_validation"],
"languageBoost": body["language_boost"],
"needNoiseReduction": body["need_noise_reduction"],
"needVolumeNormalization": body["need_volume_normalization"],
"aigcWatermark": body["aigc_watermark"],
},
"rawData": result["raw_data"],
},
Metadata: map[string]any{},
})
}
@@ -55,6 +55,9 @@ type Metrics struct {
leaseRenewalSuccess atomic.Uint64
leaseRenewalFailure atomic.Uint64
leaseRenewalLost atomic.Uint64
taskEventDuplicate atomic.Uint64
taskEventUnknownType atomic.Uint64
taskEventBudgetExceeded atomic.Uint64
}
var processingDurationBounds = [...]time.Duration{
@@ -169,6 +172,17 @@ func (m *Metrics) ObserveConcurrencyLeaseRenewal(outcome string) {
}
}
func (m *Metrics) ObserveTaskEventSkip(reason string) {
switch reason {
case "duplicate":
m.taskEventDuplicate.Add(1)
case "budget_exceeded":
m.taskEventBudgetExceeded.Add(1)
default:
m.taskEventUnknownType.Add(1)
}
}
func (m *Metrics) Handler(provider MetricsSnapshotProvider, issuer, audience string, enabled bool) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
mode := "disabled"
@@ -273,6 +287,11 @@ func (m *Metrics) Handler(provider MetricsSnapshotProvider, issuer, audience str
{"failure", m.leaseRenewalFailure.Load()},
{"lost", m.leaseRenewalLost.Load()},
})
outcomeCounters(w, "easyai_gateway_task_events_skipped_total", "Task events skipped before persistence by bounded reason.", []outcomeValue{
{"duplicate", m.taskEventDuplicate.Load()},
{"unknown_type", m.taskEventUnknownType.Load()},
{"budget_exceeded", m.taskEventBudgetExceeded.Load()},
})
})
}
@@ -30,6 +30,8 @@ func TestMetricsExposeBoundedOutcomesAndState(t *testing.T) {
metrics.ObserveAsyncWorkerResize("success")
metrics.ObserveConcurrencyLeaseRenewal("success")
metrics.ObserveConcurrencyLeaseRenewal("lost")
metrics.ObserveTaskEventSkip("duplicate")
metrics.ObserveTaskEventSkip("budget_exceeded")
recorder := httptest.NewRecorder()
metrics.Handler(metricsSnapshot{mode: "push_healthy", last: time.Now().Add(-time.Minute)}, "issuer", "audience", true).
@@ -52,6 +54,8 @@ func TestMetricsExposeBoundedOutcomesAndState(t *testing.T) {
`easyai_gateway_async_worker_resizes_total{outcome="success"} 1`,
`easyai_gateway_concurrency_lease_renewals_total{outcome="success"} 1`,
`easyai_gateway_concurrency_lease_renewals_total{outcome="lost"} 1`,
`easyai_gateway_task_events_skipped_total{outcome="duplicate"} 1`,
`easyai_gateway_task_events_skipped_total{outcome="budget_exceeded"} 1`,
} {
if !strings.Contains(recorder.Body.String(), expected) {
t.Fatalf("missing metric %q in:\n%s", expected, recorder.Body.String())
@@ -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
}
@@ -0,0 +1,11 @@
ALTER TABLE gateway_task_events
ADD COLUMN IF NOT EXISTS platform_id uuid;
ALTER TABLE gateway_task_attempts
ADD COLUMN IF NOT EXISTS status_code integer;
ALTER TABLE gateway_task_callback_outbox
ADD COLUMN IF NOT EXISTS locked_by text,
ADD COLUMN IF NOT EXISTS lock_token uuid,
ADD COLUMN IF NOT EXISTS locked_at timestamptz,
ADD COLUMN IF NOT EXISTS failed_at timestamptz;
@@ -0,0 +1,68 @@
-- easyai:migration:no-transaction
CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_gateway_tasks_retention
ON gateway_tasks(finished_at, id)
WHERE status IN ('succeeded', 'failed', 'cancelled');
-- easyai:migration:statement
CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_gateway_task_attempts_retention
ON gateway_task_attempts((COALESCE(finished_at, started_at)), id)
WHERE status <> 'running';
-- easyai:migration:statement
CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_gateway_task_events_retention
ON gateway_task_events(created_at, id);
-- easyai:migration:statement
CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_gateway_task_param_logs_retention
ON gateway_task_param_preprocessing_logs(created_at, id);
-- easyai:migration:statement
CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_gateway_tasks_minimal_compaction
ON gateway_tasks(updated_at, id)
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;
-- easyai:migration:statement
CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_gateway_tasks_checkpoint_compaction
ON gateway_tasks(updated_at, id)
WHERE COALESCE(remote_task_payload, '{}'::jsonb) <> '{}'::jsonb;
-- easyai:migration:statement
CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_gateway_task_attempts_minimal_compaction
ON gateway_task_attempts(started_at, id)
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;
-- easyai:migration:statement
CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_gateway_task_events_minimal_compaction
ON gateway_task_events(created_at, id)
WHERE payload <> '{}'::jsonb
OR phase IS NOT NULL
OR COALESCE(progress, 0) <> 0
OR message IS NOT NULL;
-- easyai:migration:statement
CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_gateway_task_param_logs_minimal_compaction
ON gateway_task_param_preprocessing_logs(created_at, id)
WHERE actual_input <> '{}'::jsonb
OR converted_output <> '{}'::jsonb
OR model_snapshot <> '{}'::jsonb
OR octet_length(changes::text) > 1024;
-- easyai:migration:statement
CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_gateway_cloned_voices_minimal_compaction
ON gateway_cloned_voices(updated_at, id)
WHERE metadata ?| ARRAY['raw', 'raw_data', 'rawData', 'provider_response', 'providerResponse', 'request'];
@@ -0,0 +1,4 @@
-- easyai:migration:no-transaction
CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_gateway_task_events_platform_created
ON gateway_task_events(platform_id, created_at DESC)
WHERE platform_id IS NOT NULL;
@@ -0,0 +1,24 @@
-- easyai:migration:no-transaction
CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_gateway_task_callback_outbox_stale
ON gateway_task_callback_outbox(status, locked_at)
WHERE status = 'processing';
-- easyai:migration:statement
CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_gateway_task_callback_outbox_delivered_retention
ON gateway_task_callback_outbox(delivered_at, id)
WHERE status = 'delivered';
-- easyai:migration:statement
CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_gateway_task_callback_outbox_failed_retention
ON gateway_task_callback_outbox(failed_at, id)
WHERE status = 'failed';
-- easyai:migration:statement
CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_gateway_task_callback_outbox_payload_compaction
ON gateway_task_callback_outbox(created_at, id)
WHERE payload <> '{}'::jsonb;
-- easyai:migration:statement
CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_gateway_task_callback_outbox_legacy
ON gateway_task_callback_outbox(created_at, id)
WHERE status IN ('pending', 'processing');
+100 -4
View File
@@ -30,7 +30,7 @@ source "$config_file"
[[ $PUBLIC_BASE_URL =~ ^https://[A-Za-z0-9.-]+$ ]] || { echo 'invalid PUBLIC_BASE_URL' >&2; exit 1; }
[[ -d $DEPLOY_DIR && -f $COMPOSE_FILE ]] || { echo 'deployment directory is unavailable' >&2; exit 1; }
for command in docker curl node flock install; do
for command in docker curl node flock install sha256sum date wc; do
command -v "$command" >/dev/null 2>&1 || { echo "$command is required" >&2; exit 1; }
done
docker info >/dev/null 2>&1 || { echo 'Docker Engine is unavailable' >&2; exit 1; }
@@ -116,18 +116,114 @@ backup_database() {
local source_sha=$1
local temp_backup=$BACKUP_DIR/.${source_sha}.dump.tmp
local final_backup=$BACKUP_DIR/${source_sha}.dump
local checksum_file=$BACKUP_DIR/${source_sha}.dump.sha256
local exclusions_file=$BACKUP_DIR/${source_sha}.excluded-tables
local restore_database=easyai_restore_${source_sha:0:12}_$$
local source_counts restored_counts
local backup_started_at backup_seconds backup_bytes restore_started_at restore_seconds
local restore_failed=false
local count_sql="SELECT (SELECT count(*) FROM gateway_wallet_transactions)::text || ':' || (SELECT count(*) FROM gateway_upload_assets)::text || ':' || (SELECT count(*) FROM gateway_cloned_voices)::text;"
local task_count_sql='SELECT (SELECT count(*) FROM gateway_tasks) + (SELECT count(*) FROM gateway_task_attempts) + (SELECT count(*) FROM gateway_task_events);'
local -a excluded_tables=(
gateway_tasks
gateway_task_attempts
gateway_task_events
gateway_task_callback_outbox
gateway_task_param_preprocessing_logs
gateway_task_message_refs
gateway_concurrency_leases
gateway_rate_limit_reservations
gateway_response_chains
settlement_outbox
river_job
)
local -a exclude_args=()
local table
for table in "${excluded_tables[@]}"; do
exclude_args+=("--exclude-table-data=$table")
done
umask 077
backup_started_at=$(date +%s)
"${compose[@]}" exec -T "$POSTGRES_SERVICE" \
pg_dump -U "$POSTGRES_USER" -d "$POSTGRES_DATABASE" -Fc >"$temp_backup"
pg_dump -U "$POSTGRES_USER" -d "$POSTGRES_DATABASE" -Fc "${exclude_args[@]}" >"$temp_backup"
backup_seconds=$(( $(date +%s) - backup_started_at ))
backup_bytes=$(wc -c <"$temp_backup")
[[ -s $temp_backup ]] || { echo 'database backup is empty' >&2; rm -f "$temp_backup"; return 1; }
if (( backup_bytes > 1073741824 || backup_seconds > 300 )); then
echo "database backup control target exceeded: bytes=$backup_bytes seconds=$backup_seconds" >&2
rm -f "$temp_backup"
return 1
fi
"${compose[@]}" exec -T "$POSTGRES_SERVICE" pg_restore -l <"$temp_backup" >/dev/null
restore_started_at=$(date +%s)
"${compose[@]}" exec -T "$POSTGRES_SERVICE" \
dropdb -U "$POSTGRES_USER" --if-exists "$restore_database"
"${compose[@]}" exec -T "$POSTGRES_SERVICE" \
createdb -U "$POSTGRES_USER" "$restore_database"
if ! "${compose[@]}" exec -T "$POSTGRES_SERVICE" \
pg_restore -U "$POSTGRES_USER" -d "$restore_database" --section=pre-data <"$temp_backup"; then
restore_failed=true
fi
if [[ $restore_failed == false ]] && ! "${compose[@]}" exec -T "$POSTGRES_SERVICE" \
pg_restore -U "$POSTGRES_USER" -d "$restore_database" --section=data <"$temp_backup"; then
restore_failed=true
fi
if [[ $restore_failed == false ]]; then
if ! "${compose[@]}" exec -T "$POSTGRES_SERVICE" \
psql -U "$POSTGRES_USER" -d "$restore_database" -v ON_ERROR_STOP=1 -c \
'UPDATE gateway_upload_assets SET task_id = NULL; UPDATE gateway_cloned_voices SET source_task_id = NULL, source_attempt_id = NULL;'; then
restore_failed=true
fi
fi
if [[ $restore_failed == false ]] && ! "${compose[@]}" exec -T "$POSTGRES_SERVICE" \
pg_restore -U "$POSTGRES_USER" -d "$restore_database" --section=post-data <"$temp_backup"; then
restore_failed=true
fi
if [[ $restore_failed == false ]]; then
if ! source_counts=$("${compose[@]}" exec -T "$POSTGRES_SERVICE" \
psql -U "$POSTGRES_USER" -d "$POSTGRES_DATABASE" -At -v ON_ERROR_STOP=1 -c \
"$count_sql"); then
restore_failed=true
fi
if ! restored_counts=$("${compose[@]}" exec -T "$POSTGRES_SERVICE" \
psql -U "$POSTGRES_USER" -d "$restore_database" -At -v ON_ERROR_STOP=1 -c \
"$count_sql"); then
restore_failed=true
fi
[[ $source_counts == "$restored_counts" ]] || restore_failed=true
fi
if [[ $restore_failed == false ]]; then
if ! restored_counts=$("${compose[@]}" exec -T "$POSTGRES_SERVICE" \
psql -U "$POSTGRES_USER" -d "$restore_database" -At -v ON_ERROR_STOP=1 -c \
"$task_count_sql"); then
restore_failed=true
fi
[[ $restored_counts == 0 ]] || restore_failed=true
fi
"${compose[@]}" exec -T "$POSTGRES_SERVICE" \
dropdb -U "$POSTGRES_USER" --if-exists "$restore_database"
restore_seconds=$(( $(date +%s) - restore_started_at ))
if (( restore_seconds > 300 )); then
restore_failed=true
fi
if [[ $restore_failed == true ]]; then
echo "database backup restore verification failed: seconds=$restore_seconds" >&2
rm -f "$temp_backup"
return 1
fi
mv "$temp_backup" "$final_backup"
sha256sum "$final_backup" >"$checksum_file"
printf '%s\n' "${excluded_tables[@]}" >"$exclusions_file"
mapfile -t expired < <(find "$BACKUP_DIR" -maxdepth 1 -type f -name '*.dump' -printf '%T@ %p\n' | sort -rn | awk -v keep="$BACKUP_RETENTION" 'NR > keep { sub(/^[^ ]+ /, ""); print }')
if [[ ${#expired[@]} -gt 0 ]]; then
rm -f -- "${expired[@]}"
for table in "${expired[@]}"; do
rm -f -- "$table" "$table.sha256" "${table%.dump}.excluded-tables"
done
fi
echo "[release] verified database backup: $final_backup"
echo "[release] verified restorable task-history-free database backup: $final_backup bytes=$backup_bytes backup_seconds=$backup_seconds restore_seconds=$restore_seconds"
}
record_current() {
+5
View File
@@ -19,6 +19,11 @@ x-api-environment: &api-environment
TASK_PROGRESS_CALLBACK_URL: ${AI_GATEWAY_COMPOSE_TASK_PROGRESS_CALLBACK_URL:-http://host.docker.internal:3000/internal/platform/task-progress-callbacks}
TASK_PROGRESS_CALLBACK_TIMEOUT_MS: ${TASK_PROGRESS_CALLBACK_TIMEOUT_MS:-5000}
TASK_PROGRESS_CALLBACK_MAX_ATTEMPTS: ${TASK_PROGRESS_CALLBACK_MAX_ATTEMPTS:-10}
AI_GATEWAY_TASK_CLEANUP_ENABLED: ${AI_GATEWAY_TASK_CLEANUP_ENABLED:-false}
AI_GATEWAY_TASK_RETENTION_DAYS: ${AI_GATEWAY_TASK_RETENTION_DAYS:-30}
AI_GATEWAY_TASK_ANALYSIS_RETENTION_DAYS: ${AI_GATEWAY_TASK_ANALYSIS_RETENTION_DAYS:-7}
AI_GATEWAY_TASK_CLEANUP_INTERVAL_SECONDS: ${AI_GATEWAY_TASK_CLEANUP_INTERVAL_SECONDS:-300}
AI_GATEWAY_TASK_CLEANUP_BATCH_SIZE: ${AI_GATEWAY_TASK_CLEANUP_BATCH_SIZE:-1000}
CORS_ALLOWED_ORIGIN: ${AI_GATEWAY_COMPOSE_CORS_ALLOWED_ORIGIN:-http://localhost:5178,http://127.0.0.1:5178}
AI_GATEWAY_PUBLIC_BASE_URL: ${AI_GATEWAY_COMPOSE_PUBLIC_BASE_URL:-http://localhost:8088}
AI_GATEWAY_WEB_BASE_URL: ${AI_GATEWAY_COMPOSE_WEB_BASE_URL:-http://localhost:5178}
+43 -40
View File
@@ -454,13 +454,13 @@ Gateway Web Console 默认使用 HttpOnly OIDC Cookie 会话解决标签页隔
- `user_group_id` / `user_group_key`
- `user_group_policy_snapshot`:任务创建时解析出的用户组策略快照,用于审计和重试稳定性。
- `model`
- `request`
- `request`:唯一完整请求;`normalized_request` 为兼容列,新数据固定为 `{}`
- `status``queued``running``succeeded``failed``cancelled`
- `queue_key`:限流队列 key,例如 `${platformKey}-${model}``${provider}-${methodName}`
- `priority`
- `idempotency_key`
- `remote_task_id`
- `remote_task_payload`
- `remote_task_payload`:仅保存 provider 显式白名单内、跨进程恢复不可重算的 checkpoint;禁止保存提交响应、最终响应和完整请求副本。
- `run_mode``production``simulation`
- `simulation_profile`
- `simulation_seed`
@@ -472,9 +472,11 @@ Gateway Web Console 默认使用 HttpOnly OIDC Cookie 会话解决标签页隔
- `billings`
- `error`
Provider 原始响应只在内存中参与标准结果转换,不进入 `result.raw`。兼容协议响应由任务状态、标准 `result`、错误、时间和 `remote_task_id` 动态生成;兼容响应 body、Header、HTTP status 和 public ID 快照均不再写入。终态任务默认保留 30 天。
### 7.4 `gateway_task_attempts`
保存每一次客户端尝试,支持“上一个客户端失败,下一个客户端重试”的完整审计
保存每一次客户端尝试的标量分析信息,支持还原重试顺序和错误
- `task_id`
- `attempt_no`
@@ -487,25 +489,27 @@ Gateway Web Console 默认使用 HttpOnly OIDC Cookie 会话解决标签页隔
- `error_message`
- `remote_task_id`
- `simulated`
- `request_snapshot`
- `response_snapshot`
- `request_id` / `status_code`
- `upstream_submission_status`
- `response_started_at` / `response_finished_at` / `response_duration_ms`
- `started_at`
- `finished_at`
`request_snapshot``response_snapshot``usage``metrics``pricing_snapshot``request_fingerprint` 为兼容字段,新数据固定为空。错误信息最多 2KB,attempt 默认保留 7 天。
### 7.5 `gateway_task_events`
保存任务执行过程中的事件,用于控制台 SSE / WebSocket 重放、进度回调 outbox 投递与服务重启后的进度恢复
保存真实状态变化、真实失败 attempt、计费结果和真实策略动作,用于控制台 SSE / WebSocket 重放与 callback outbox 投递
- `task_id`
- `seq`:单任务内递增序号
- `event_type``queue_status``node_status``progress``partial_result``completed``failed`
- `event_type`代码白名单内的 accepted、running、终态、billing、attempt failed 和 policy 事件。
- `status`
- `progress`
- `message`
- `payload`
- `platform_id`
- `simulated`
- `created_at`
客户端断线重连时可通过 `Last-Event-ID` `afterSeq` 回放事件
`phase``progress``message` `payload` 为兼容字段,新数据固定为空。轮询和 synthetic progress 不写事件;同一 task 行锁内的连续重复状态被跳过且不递增 `seq`。普通成功任务最多 4 条事件,每个任务最多 16 条非终态事件,终态和计费事件不受此限制。事件默认保留 7 天
### 7.5.1 `gateway_task_callback_outbox`
@@ -514,13 +518,15 @@ Gateway Web Console 默认使用 HttpOnly OIDC Cookie 会话解决标签页隔
- `task_id`
- `event_id` / `seq`:与 `gateway_task_events` 对应,用于幂等与顺序控制。
- `callback_url`:当前使用的回调地址快照,避免配置变化影响已排队事件。
- `payload`回调给 server-main 的标准事件载荷
- `status``pending``delivering``delivered``failed``skipped`
- `payload`兼容列,新数据固定为 `{}`worker 根据 event 构造极简回调
- `status``pending``processing``delivered``failed`
- `attempts`
- `next_attempt_at`
- `last_error`
- `delivered_at`
请求体固定只包含 `taskId``seq``eventType``status``createdAt`,幂等键为 `<taskId>:<seq>`。2xx 成功;网络错误、408、429 和 5xx 重试,其他 4xx 直接失败,最多投递 10 次。`delivered` 保留 24 小时,`failed` 保留 7 天;任一 callback 仍处于保留期时不得删除对应 task。
### 7.6 `runtime_client_states`
保存平台客户端的运行时状态:
@@ -1044,8 +1050,8 @@ Gateway Web Console 默认使用 HttpOnly OIDC Cookie 会话解决标签页隔
| `retryable` | `boolean` | not null, default `false` | 是否可重试 |
| `simulated` | `boolean` | not null, default `false` | 是否模拟 |
| `remote_task_id` | `text` | nullable | 供应商任务 ID |
| `request_snapshot` | `jsonb` | not null, default `{}` | 本次请求快照 |
| `response_snapshot` | `jsonb` | nullable | 本次响应快照 |
| `request_snapshot` | `jsonb` | not null, default `{}` | 兼容列,新数据固定为 `{}` |
| `response_snapshot` | `jsonb` | nullable | 兼容列,新数据固定为 `{}` |
| `error_code` | `text` | nullable | 错误码 |
| `error_message` | `text` | nullable | 错误信息 |
| `started_at` | `timestamptz` | not null | 开始时间 |
@@ -1064,12 +1070,13 @@ Gateway Web Console 默认使用 HttpOnly OIDC Cookie 会话解决标签页隔
| `id` | `uuid` | PK | 事件 ID |
| `task_id` | `uuid` | FK | 所属任务 |
| `seq` | `bigint` | not null | 单任务递增序号 |
| `event_type` | `text` | not null | `queue_status``progress``partial_result``completed``failed` |
| `event_type` | `text` | not null | 代码白名单内的真实状态、失败、计费和策略事件 |
| `status` | `text` | nullable | 任务状态 |
| `phase` | `text` | nullable | `routing``queued``submit``polling``upload` |
| `progress` | `numeric` | nullable | 0 到 1 |
| `message` | `text` | nullable | 展示消息 |
| `payload` | `jsonb` | not null, default `{}` | 事件载荷 |
| `platform_id` | `uuid` | nullable | 策略事件平台 ID |
| `phase` | `text` | nullable | 兼容列,新数据为空 |
| `progress` | `numeric` | nullable | 兼容列,新数据为 0 |
| `message` | `text` | nullable | 兼容列,新数据为空 |
| `payload` | `jsonb` | not null, default `{}` | 兼容列,新数据固定为 `{}` |
| `simulated` | `boolean` | not null, default `false` | 是否模拟 |
| `created_at` | `timestamptz` | not null | 创建时间 |
@@ -1087,8 +1094,8 @@ Gateway Web Console 默认使用 HttpOnly OIDC Cookie 会话解决标签页隔
| `event_id` | `uuid` | FK nullable | 对应 `gateway_task_events.id` |
| `seq` | `bigint` | not null | 单任务事件序号 |
| `callback_url` | `text` | not null | 投递地址快照,来自 `TASK_PROGRESS_CALLBACK_URL` |
| `payload` | `jsonb` | not null | 回调载荷 |
| `status` | `text` | not null | `pending``delivering``delivered``failed``skipped` |
| `payload` | `jsonb` | not null | 兼容列,新数据固定为 `{}` |
| `status` | `text` | not null | `pending``processing``delivered``failed` |
| `attempts` | `int` | not null, default `0` | 已投递次数 |
| `next_attempt_at` | `timestamptz` | not null | 下次投递时间 |
| `last_error` | `text` | nullable | 最近错误 |
@@ -1615,26 +1622,15 @@ POST ${TASK_PROGRESS_CALLBACK_URL}
Authorization: Bearer ${SERVER_MAIN_INTERNAL_TOKEN}
Content-Type: application/json
Idempotency-Key: ${taskId}:${seq}
X-EasyAI-Event-Type: task.progress
X-EasyAI-Event-Type: task.completed
```
```json
{
"eventId": "uuid",
"taskId": "uuid",
"externalTaskId": "server-main-task-id",
"userId": "user-id",
"tenantId": "tenant-id",
"apiKeyId": "optional",
"kind": "images.generations",
"model": "gpt-image-1",
"seq": 12,
"event": "progress",
"status": "running",
"phase": "polling",
"progress": 0.42,
"message": "Generating video frames",
"payload": {},
"seq": 3,
"eventType": "task.completed",
"status": "succeeded",
"createdAt": "2026-05-09T12:00:00Z"
}
```
@@ -1642,7 +1638,7 @@ X-EasyAI-Event-Type: task.progress
`server-main` 收到后不重新执行任务,只做三件事:
1. 按 `Idempotency-Key``taskId + seq` 幂等落库/去重。
2. 根据 `externalTaskId``taskId``userId``tenantId` 找到原业务会话或任务频道
2. 根据 `taskId` 找到原业务会话或任务频道;需要完整结果时调用 task detail
3. 通过原 WebSocket 网关推送给业务前端,保持前端订阅协议不大改。
### 12.3 进度持久化要求
@@ -1822,15 +1818,22 @@ TASK_PROGRESS_CALLBACK_URL=http://easyai-server-main:3000/internal/platform/task
TASK_PROGRESS_CALLBACK_ENABLED=true
TASK_PROGRESS_CALLBACK_TIMEOUT_MS=5000
TASK_PROGRESS_CALLBACK_MAX_ATTEMPTS=10
AI_GATEWAY_TASK_CLEANUP_ENABLED=false
AI_GATEWAY_TASK_RETENTION_DAYS=30
AI_GATEWAY_TASK_ANALYSIS_RETENTION_DAYS=7
AI_GATEWAY_TASK_CLEANUP_INTERVAL_SECONDS=300
AI_GATEWAY_TASK_CLEANUP_BATCH_SIZE=1000
```
运行时要求:
- `BaseClient` / runtime 每产生一个进度事件,先写 `gateway_task_events`,再写 `gateway_task_callback_outbox`
- 首次上线先保持清理关闭;兼容协议和异步恢复验收通过后,显式改为 `AI_GATEWAY_TASK_CLEANUP_ENABLED=true`,再启动历史压缩和 7/30 天删除
- runtime 仅在任务真实状态变化时写 `gateway_task_events`;事件 payload 固定为空,轮询不会生成事件。
- callback worker 以 `0083_task_history_minimal_storage` 的应用时间为边界,只投递新 outbox;旧 pending/processing callback 在启用治理后分批标记为 `failed`,不会向业务端补发形成历史回调风暴。
- callback worker 使用 `SERVER_MAIN_INTERNAL_TOKEN` 调用 `TASK_PROGRESS_CALLBACK_URL`
- `server-main` 收到事件后进入主服务内部推送流程,再由原 WebSocket 网关推送给业务前端。
- Gateway 控制台仍可通过 SSE 读 `gateway_task_events`,用于运维诊断,不影响业务前端推送通道。
- 如果 `server-main` 短暂不可用,Gateway 按 outbox 重试;超过最大次数后标记 `failed`,控制台可手动 replay
- 如果 `server-main` 短暂不可用,Gateway 按 outbox 重试;超过最大次数后标记 `failed` 并保留 7 天供诊断
### 14.3 中期:前端直接打 Gateway
+36 -1
View File
@@ -110,13 +110,48 @@ dist/releases/<40 位 Git SHA>.json
- API/Web 均为完整 Registry digest
- 镜像为 `linux/amd64`
- 本次变化组件的 OCI revision 等于 release SHA
- 迁移变化时备份可被 `pg_restore -l` 读取
- 迁移变化时生成排除任务历史数据的完整 schema/业务状态备份,并实际恢复到一次性数据库验证长期数据、外键和空任务历史表
- 内部 API health/readiness/OpenAPI
- Web 页面和 Web API 反代;
- 公网 health/readiness/OpenAPI。
失败时自动恢复上一组应用 digest 并重复探活。已经执行的 expand-only 数据库迁移不会自动恢复。
## 任务历史治理
首次上线极简写入时保持 `AI_GATEWAY_TASK_CLEANUP_ENABLED=false`。先验证 Kling V1/V2、Volces 的提交、查询、失败响应和异步恢复均不依赖兼容快照或 provider 原始响应,再单独开启 cleanup。不要把首次清理和应用切换放在同一个维护动作中。
开启后每 5 分钟确认 7 天分析数据和 30 天主任务的过期积压;任一积压最长时间超过 24 小时、callback 待投递增长、结算待办增长或事件跳过指标异常时暂停治理并排查。以下查询只用于只读验收:
```sql
SELECT relname, pg_size_pretty(pg_total_relation_size(relid)) AS total_size
FROM pg_catalog.pg_statio_user_tables
WHERE relname IN (
'gateway_tasks',
'gateway_task_attempts',
'gateway_task_events',
'gateway_task_callback_outbox',
'gateway_task_param_preprocessing_logs'
)
ORDER BY pg_total_relation_size(relid) DESC;
SELECT
(SELECT count(*) FROM gateway_task_attempts WHERE created_at < now() - interval '7 days') AS expired_attempts,
(SELECT count(*) FROM gateway_task_events WHERE created_at < now() - interval '7 days') AS expired_events,
(SELECT count(*) FROM gateway_task_param_preprocessing_logs WHERE created_at < now() - interval '7 days') AS expired_parameter_logs,
(SELECT count(*) FROM gateway_tasks WHERE finished_at < now() - interval '30 days') AS candidate_tasks;
```
逻辑删除稳定后执行 `VACUUM (ANALYZE)` 更新统计信息;它不会把空间归还给操作系统。物理回收只在独立维护窗口中逐表执行 `pg_repack`,执行前确认扩展、额外磁盘空间、主从延迟和锁等待均满足要求。优先处理 `gateway_task_events``gateway_task_attempts``gateway_tasks`,不得与 deploy、迁移或大批量清理并行。
治理控制目标:
- 新任务的兼容 body/Header、event/outbox payload、attempt 大 JSON 均为空;
- 普通成功任务不超过 4 条事件,每任务非终态事件不超过 16 条;
- 分析数据过期积压低于 24 小时,稳态分析表总量不超过主任务表的 20%;
- 首轮治理后任务域低于治理前 10.9GB 的 40%;
- 发布备份不超过 1GB,备份和实际恢复各不超过 5 分钟。
## 状态与回滚
```bash
+6 -17
View File
@@ -169,35 +169,24 @@ Idempotency-Key: ${groupKey}:${version}
### 1.6 任务进度回调到 server-main
AI Gateway 不直接替换原业务前端 WebSocket 通道。Gateway 配置任务进度回调地址,所有任务中间状态先写入 Gateway 本地事件表和 callback outbox,再回调给 `server-main`,由 `server-main` 内部推送流程复用原 WebSocket 网关推送给业务前端。
AI Gateway 不直接替换原业务前端 WebSocket 通道。Gateway 仅在真实任务状态变化时写入本地事件表和 callback outbox,再回调给 `server-main`,由 `server-main` 内部推送流程复用原 WebSocket 网关推送给业务前端。
```http
POST /internal/platform/task-progress-callbacks
Authorization: Bearer ${SERVER_MAIN_INTERNAL_TOKEN}
Content-Type: application/json
Idempotency-Key: ${taskId}:${seq}
X-EasyAI-Event-Type: task.progress
X-EasyAI-Event-Type: task.completed
```
请求体:
```json
{
"eventId": "uuid",
"taskId": "gateway-task-id",
"externalTaskId": "server-main-task-id",
"userId": "user-id",
"tenantId": "tenant-id",
"apiKeyId": "optional",
"kind": "images.generations",
"model": "gpt-image-1",
"seq": 12,
"event": "progress",
"status": "running",
"phase": "polling",
"progress": 0.42,
"message": "Generating video frames",
"payload": {},
"seq": 3,
"eventType": "task.completed",
"status": "succeeded",
"createdAt": "2026-05-09T12:00:00Z"
}
```
@@ -205,7 +194,7 @@ X-EasyAI-Event-Type: task.progress
`server-main` 处理要求:
- 使用 `Idempotency-Key``taskId + seq` 幂等去重。
- 根据 `externalTaskId` / `taskId` / `userId` / `tenantId` 定位原业务频道
- 根据 `taskId` 定位原业务频道;需要完整结果时调用 task detail
- 复用现有 WebSocket 网关事件格式推给前端,尽量不改业务前端订阅协议。
- 只负责推送与必要状态同步,不重新执行任务、不重新计算计费。
+6 -1
View File
@@ -126,7 +126,7 @@ node "$root/scripts/release-manifest.mjs" validate "$manifest" >/dev/null
[[ $(node "$root/scripts/release-manifest.mjs" get "$manifest" images.api) == "$api_image" ]]
bootstrap_manifest=$tmp/bootstrap.json
RELEASE_SOURCE_SHA=$source_sha \
RELEASE_BASE_SHA= \
RELEASE_BASE_SHA='' \
RELEASE_COMPONENTS=api \
RELEASE_MIGRATIONS_CHANGED=true \
RELEASE_API_IMAGE=$api_image \
@@ -270,6 +270,11 @@ if find "$root/.gitea/workflows" -type f -print -quit 2>/dev/null | grep -q .; t
exit 1
fi
grep -Fq 'stale release manifest' "$root/deploy/manual/easyai-ai-gateway-release"
grep -Fq -- "--exclude-table-data=\$table" "$root/deploy/manual/easyai-ai-gateway-release"
grep -Fq -- "pg_restore -U \"\$POSTGRES_USER\" -d \"\$restore_database\" --section=post-data" "$root/deploy/manual/easyai-ai-gateway-release"
grep -Fq 'gateway_tasks' "$root/deploy/manual/easyai-ai-gateway-release"
grep -Fq 'gateway_task_attempts' "$root/deploy/manual/easyai-ai-gateway-release"
grep -Fq 'gateway_task_events' "$root/deploy/manual/easyai-ai-gateway-release"
grep -Fq 'production_changed=false' "$root/scripts/publish-release-images.sh"
echo 'manual_release_tests=PASS'