BCAI/easyai-ai-gateway
3
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

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

Browse Source 
停止持久化 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:
wangbo
2026-07-24 18:23:22 +08:00
parent 09375bfae7
commit 810dcfeee6
46 changed files with 2246 additions and 458 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/api/cmd/migrate/main.go
+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
}
apps/api/cmd/migrate/main_test.go
+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 {
apps/api/internal/clients/clients_test.go
+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)
}
}
apps/api/internal/clients/gemini.go
-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,
}
}
apps/api/internal/clients/keling.go
-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,
}
}
apps/api/internal/clients/media_clients.go
+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}}
apps/api/internal/clients/provider_task.go
+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
}
apps/api/internal/clients/volces.go
+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
}
apps/api/internal/clients/volces_deyun_test.go
+1 -1
View File
@@ -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)
}
}
apps/api/internal/config/config.go
+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":
apps/api/internal/config/config_test.go
+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)
}
}
apps/api/internal/httpapi/core_flow_integration_test.go
+16
View File
@@ -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)
apps/api/internal/httpapi/keling_compat_handlers.go
+15 -22
View File
@@ -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 {
apps/api/internal/httpapi/keling_compat_handlers_test.go
+4
View File
@@ -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) {
apps/api/internal/httpapi/kling_compat.go
+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
apps/api/internal/httpapi/kling_compat_test.go
+24
View File
@@ -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
apps/api/internal/httpapi/server.go
+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)
apps/api/internal/httpapi/volces_compat_handlers.go
+41 -65
View File
@@ -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
apps/api/internal/httpapi/volces_compat_handlers_test.go
+6 -7
View File
@@ -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)
apps/api/internal/runner/service.go
+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)
}
apps/api/internal/runner/task_history_workers.go
+181
View File
@@ -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):
}
}
}
apps/api/internal/runner/task_history_workers_test.go
+118
View File
@@ -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)
}
}
apps/api/internal/runner/voice_clone.go
+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{},
})
}
apps/api/internal/securityevents/metrics.go
+19
View File
@@ -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()},
})
})
}
apps/api/internal/securityevents/metrics_test.go
+4
View File
@@ -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())
apps/api/internal/store/billing_settlements.go
+5 -5
View File
@@ -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, `
apps/api/internal/store/compatibility_tasks.go
+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
apps/api/internal/store/oidc_users_integration_test.go
+16
View File
@@ -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)
apps/api/internal/store/postgres.go
+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",
}}
}
apps/api/internal/store/rate_limit_status.go
+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
}
apps/api/internal/store/rate_limits.go
+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
apps/api/internal/store/runtime_types.go
+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
apps/api/internal/store/task_history.go
+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
}
apps/api/internal/store/task_history_test.go
+394
View File
@@ -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)
}
}
apps/api/internal/store/tasks_runtime.go
+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
}
apps/api/migrations/0083_task_history_minimal_storage.sql
+11
View File
@@ -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;
apps/api/migrations/0084_task_history_retention_index.sql
+68
View File
@@ -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'];
apps/api/migrations/0085_task_event_platform_index.sql
+4
View File
@@ -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;
apps/api/migrations/0086_task_callback_lock_index.sql
+24
View File
@@ -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');