perf(storage): 拦截任务二进制并本地暂存结果

原因:任务标准结果中的 Base64、Data URI 和 Buffer 会进入 PostgreSQL JSON,导致 TOAST 与备份体积快速增长。

影响:新增统一 JSON 持久化关口;upload_none 将二进制原子写入本地结果目录,数据库仅保存带 SHA-256 的有界占位符;任务详情、同步响应、异步查询和兼容协议按需校验恢复。补充 24 小时清理、容量上限、历史小批量治理命令及管理端说明。

风险:本地结果超过 TTL、丢失或损坏时分别返回明确的 410/500;空间不足时返回 503 且不重试上游。未自动执行历史治理。

验证:三种真实图片模型同步/异步与幂等重放通过;Go vet/全量测试、前端 111 测试、lint/typecheck/build、OpenAPI、迁移安全、govulncheck、依赖审计、手工发布测试及 Linux amd64 构建通过。
This commit is contained in:
2026-07-24 21:13:09 +08:00
parent 2457de6a56
commit 4f163ea6d7
31 changed files with 2067 additions and 30 deletions
+4
View File
@@ -80,6 +80,10 @@ AI_GATEWAY_TASK_RETENTION_DAYS=30
AI_GATEWAY_TASK_ANALYSIS_RETENTION_DAYS=7 AI_GATEWAY_TASK_ANALYSIS_RETENTION_DAYS=7
AI_GATEWAY_TASK_CLEANUP_INTERVAL_SECONDS=300 AI_GATEWAY_TASK_CLEANUP_INTERVAL_SECONDS=300
AI_GATEWAY_TASK_CLEANUP_BATCH_SIZE=1000 AI_GATEWAY_TASK_CLEANUP_BATCH_SIZE=1000
AI_GATEWAY_LOCAL_RESULT_TTL_HOURS=24
AI_GATEWAY_LOCAL_RESULT_MIN_FREE_BYTES=10737418240
AI_GATEWAY_LOCAL_RESULT_MAX_BYTES=268435456
AI_GATEWAY_LOCAL_RESULT_MAX_TASK_BYTES=536870912
CORS_ALLOWED_ORIGIN=http://localhost:5178,http://127.0.0.1:5178 CORS_ALLOWED_ORIGIN=http://localhost:5178,http://127.0.0.1:5178
VITE_GATEWAY_API_BASE_URL=http://localhost:8088 VITE_GATEWAY_API_BASE_URL=http://localhost:8088
+3 -1
View File
@@ -32,7 +32,8 @@ RUN --mount=type=cache,target=/go/pkg/mod \
--mount=type=cache,target=/root/.cache/go-build \ --mount=type=cache,target=/root/.cache/go-build \
cd apps/api && \ cd apps/api && \
CGO_ENABLED=0 GOOS=$TARGETOS GOARCH=$TARGETARCH go build -trimpath -ldflags="-s -w" -o /out/easyai-ai-gateway ./cmd/gateway && \ CGO_ENABLED=0 GOOS=$TARGETOS GOARCH=$TARGETARCH go build -trimpath -ldflags="-s -w" -o /out/easyai-ai-gateway ./cmd/gateway && \
CGO_ENABLED=0 GOOS=$TARGETOS GOARCH=$TARGETARCH go build -trimpath -ldflags="-s -w" -o /out/easyai-ai-gateway-migrate ./cmd/migrate CGO_ENABLED=0 GOOS=$TARGETOS GOARCH=$TARGETARCH go build -trimpath -ldflags="-s -w" -o /out/easyai-ai-gateway-migrate ./cmd/migrate && \
CGO_ENABLED=0 GOOS=$TARGETOS GOARCH=$TARGETARCH go build -trimpath -ldflags="-s -w" -o /out/easyai-ai-gateway-backfill-binary-results ./cmd/backfill-binary-results
FROM ${API_RUNTIME_IMAGE} AS api FROM ${API_RUNTIME_IMAGE} AS api
@@ -42,6 +43,7 @@ WORKDIR /app
COPY --from=api-builder /out/easyai-ai-gateway /app/easyai-ai-gateway COPY --from=api-builder /out/easyai-ai-gateway /app/easyai-ai-gateway
COPY --from=api-builder /out/easyai-ai-gateway-migrate /app/easyai-ai-gateway-migrate COPY --from=api-builder /out/easyai-ai-gateway-migrate /app/easyai-ai-gateway-migrate
COPY --from=api-builder /out/easyai-ai-gateway-backfill-binary-results /app/easyai-ai-gateway-backfill-binary-results
COPY apps/api/migrations /app/migrations COPY apps/api/migrations /app/migrations
RUN mkdir -p /app/data/static/generated /app/data/static/uploaded && \ RUN mkdir -p /app/data/static/generated /app/data/static/uploaded && \
@@ -0,0 +1,113 @@
package main
import (
"context"
"flag"
"log/slog"
"os"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/runner"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
func main() {
apply := flag.Bool("apply", false, "persist compacted results; default is dry-run")
batchSize := flag.Int("batch-size", 100, "rows per batch, maximum 100")
maxBatches := flag.Int("max-batches", 10, "maximum batches for one invocation")
afterID := flag.String("after-id", "", "resume after this task UUID")
flag.Parse()
logger := slog.New(slog.NewJSONHandler(os.Stdout, nil))
cfg := config.Load()
if err := cfg.Validate(); err != nil {
logger.Error("invalid gateway configuration", "error", err)
os.Exit(1)
}
if *batchSize < 1 || *batchSize > 100 || *maxBatches < 1 {
logger.Error("invalid backfill bounds", "batchSize", *batchSize, "maxBatches", *maxBatches)
os.Exit(1)
}
ctx := context.Background()
db, err := store.Connect(ctx, cfg.DatabaseURL)
if err != nil {
logger.Error("connect postgres failed", "error", err)
os.Exit(1)
}
defer db.Close()
service := runner.New(cfg, db, logger)
cursor := *afterID
scanned := 0
matched := 0
updated := 0
expired := 0
for batch := 0; batch < *maxBatches; batch++ {
items, err := db.ListTaskBinaryResultBackfillBatch(ctx, cursor, *batchSize)
if err != nil {
logger.Error("list binary result backfill batch failed", "afterId", cursor, "error", err)
os.Exit(1)
}
if len(items) == 0 {
break
}
for _, item := range items {
cursor = item.ID
scanned++
if !runner.TaskResultHasInlineBinary(item.Result) {
continue
}
matched++
if !*apply {
continue
}
isExpired := item.FinishedAt.Before(time.Now().Add(-time.Duration(localResultTTLHours(cfg)) * time.Hour))
var persistent map[string]any
var changed bool
if isExpired {
persistent, changed, err = service.CompactExpiredTaskResultForStorage(ctx, item.ID, item.Result)
} else {
persistent, changed, err = service.MaterializeTaskResultForStorage(ctx, item.ID, item.Result)
}
if err != nil {
logger.Error("materialize historical binary result failed", "taskId", item.ID, "error", err)
os.Exit(1)
}
if !changed {
continue
}
ok, err := db.UpdateTaskBinaryResultBackfill(ctx, item.ID, persistent)
if err != nil {
logger.Error("update historical binary result failed", "taskId", item.ID, "error", err)
os.Exit(1)
}
if !ok {
continue
}
updated++
if isExpired {
expired++
}
}
if len(items) < *batchSize {
break
}
}
logger.Info("binary result backfill completed",
"apply", *apply,
"scanned", scanned,
"matched", matched,
"updated", updated,
"expired", expired,
"resumeAfterId", cursor,
)
}
func localResultTTLHours(cfg config.Config) int {
if cfg.LocalResultTTLHours <= 0 {
return 24
}
return cfg.LocalResultTTLHours
}
+44
View File
@@ -4738,6 +4738,12 @@
"schema": { "schema": {
"$ref": "#/definitions/httpapi.EasyAIGeneratedResponse" "$ref": "#/definitions/httpapi.EasyAIGeneratedResponse"
} }
},
"410": {
"description": "Gone",
"schema": {
"$ref": "#/definitions/httpapi.EasyAIGeneratedResponse"
}
} }
} }
} }
@@ -5645,6 +5651,12 @@
"schema": { "schema": {
"$ref": "#/definitions/httpapi.VolcesErrorEnvelope" "$ref": "#/definitions/httpapi.VolcesErrorEnvelope"
} }
},
"410": {
"description": "Gone",
"schema": {
"$ref": "#/definitions/httpapi.VolcesErrorEnvelope"
}
} }
} }
}, },
@@ -6304,6 +6316,13 @@
"type": "object", "type": "object",
"additionalProperties": true "additionalProperties": true
} }
},
"410": {
"description": "Gone",
"schema": {
"type": "object",
"additionalProperties": true
}
} }
} }
} }
@@ -6404,6 +6423,13 @@
"type": "object", "type": "object",
"additionalProperties": true "additionalProperties": true
} }
},
"410": {
"description": "Gone",
"schema": {
"type": "object",
"additionalProperties": true
}
} }
} }
}, },
@@ -7883,6 +7909,12 @@
"$ref": "#/definitions/httpapi.ErrorEnvelope" "$ref": "#/definitions/httpapi.ErrorEnvelope"
} }
}, },
"410": {
"description": "Gone",
"schema": {
"$ref": "#/definitions/httpapi.ErrorEnvelope"
}
},
"500": { "500": {
"description": "Internal Server Error", "description": "Internal Server Error",
"schema": { "schema": {
@@ -8305,6 +8337,12 @@
"$ref": "#/definitions/httpapi.KelingCompatibleEnvelope" "$ref": "#/definitions/httpapi.KelingCompatibleEnvelope"
} }
}, },
"410": {
"description": "Gone",
"schema": {
"$ref": "#/definitions/httpapi.KelingCompatibleEnvelope"
}
},
"500": { "500": {
"description": "Internal Server Error", "description": "Internal Server Error",
"schema": { "schema": {
@@ -8766,6 +8804,12 @@
"$ref": "#/definitions/httpapi.ErrorEnvelope" "$ref": "#/definitions/httpapi.ErrorEnvelope"
} }
}, },
"410": {
"description": "Gone",
"schema": {
"$ref": "#/definitions/httpapi.ErrorEnvelope"
}
},
"500": { "500": {
"description": "Internal Server Error", "description": "Internal Server Error",
"schema": { "schema": {
+30
View File
@@ -6821,6 +6821,10 @@ paths:
description: Not Found description: Not Found
schema: schema:
$ref: '#/definitions/httpapi.EasyAIGeneratedResponse' $ref: '#/definitions/httpapi.EasyAIGeneratedResponse'
"410":
description: Gone
schema:
$ref: '#/definitions/httpapi.EasyAIGeneratedResponse'
security: security:
- BearerAuth: [] - BearerAuth: []
summary: 查询 server-main EasyAIClient 兼容任务结果 summary: 查询 server-main EasyAIClient 兼容任务结果
@@ -7429,6 +7433,10 @@ paths:
description: Not Found description: Not Found
schema: schema:
$ref: '#/definitions/httpapi.VolcesErrorEnvelope' $ref: '#/definitions/httpapi.VolcesErrorEnvelope'
"410":
description: Gone
schema:
$ref: '#/definitions/httpapi.VolcesErrorEnvelope'
security: security:
- BearerAuth: [] - BearerAuth: []
summary: 查询火山内容生成任务 summary: 查询火山内容生成任务
@@ -7829,6 +7837,11 @@ paths:
schema: schema:
additionalProperties: true additionalProperties: true
type: object type: object
"410":
description: Gone
schema:
additionalProperties: true
type: object
security: security:
- BearerAuth: [] - BearerAuth: []
summary: 查询可灵 V1 Omni 视频任务 summary: 查询可灵 V1 Omni 视频任务
@@ -7893,6 +7906,11 @@ paths:
schema: schema:
additionalProperties: true additionalProperties: true
type: object type: object
"410":
description: Gone
schema:
additionalProperties: true
type: object
security: security:
- BearerAuth: [] - BearerAuth: []
summary: 按 ID 查询可灵 API 2.0 任务 summary: 按 ID 查询可灵 API 2.0 任务
@@ -8852,6 +8870,10 @@ paths:
description: Not Found description: Not Found
schema: schema:
$ref: '#/definitions/httpapi.ErrorEnvelope' $ref: '#/definitions/httpapi.ErrorEnvelope'
"410":
description: Gone
schema:
$ref: '#/definitions/httpapi.ErrorEnvelope'
"500": "500":
description: Internal Server Error description: Internal Server Error
schema: schema:
@@ -9123,6 +9145,10 @@ paths:
description: Not Found description: Not Found
schema: schema:
$ref: '#/definitions/httpapi.KelingCompatibleEnvelope' $ref: '#/definitions/httpapi.KelingCompatibleEnvelope'
"410":
description: Gone
schema:
$ref: '#/definitions/httpapi.KelingCompatibleEnvelope'
"500": "500":
description: Internal Server Error description: Internal Server Error
schema: schema:
@@ -9421,6 +9447,10 @@ paths:
description: Not Found description: Not Found
schema: schema:
$ref: '#/definitions/httpapi.ErrorEnvelope' $ref: '#/definitions/httpapi.ErrorEnvelope'
"410":
description: Gone
schema:
$ref: '#/definitions/httpapi.ErrorEnvelope'
"500": "500":
description: Internal Server Error description: Internal Server Error
schema: schema:
+32
View File
@@ -39,6 +39,10 @@ type Config struct {
LocalGeneratedStorageDir string LocalGeneratedStorageDir string
LocalUploadedStorageDir string LocalUploadedStorageDir string
LocalTempAssetTTLHours int LocalTempAssetTTLHours int
LocalResultTTLHours int
LocalResultMinFreeBytes int64
LocalResultMaxBytes int64
LocalResultMaxTaskBytes int64
TaskProgressCallbackEnabled bool TaskProgressCallbackEnabled bool
TaskProgressCallbackURL string TaskProgressCallbackURL string
TaskProgressCallbackTimeoutMS int TaskProgressCallbackTimeoutMS int
@@ -88,6 +92,10 @@ func Load() Config {
LocalGeneratedStorageDir: env("AI_GATEWAY_GENERATED_STORAGE_DIR", env("LOCAL_GENERATED_STORAGE_DIR", env("AI_GATEWAY_STATIC_STORAGE_DIR", DefaultLocalGeneratedStorageDir))), LocalGeneratedStorageDir: env("AI_GATEWAY_GENERATED_STORAGE_DIR", env("LOCAL_GENERATED_STORAGE_DIR", env("AI_GATEWAY_STATIC_STORAGE_DIR", DefaultLocalGeneratedStorageDir))),
LocalUploadedStorageDir: env("AI_GATEWAY_UPLOADED_STORAGE_DIR", env("LOCAL_UPLOADED_STORAGE_DIR", DefaultLocalUploadedStorageDir)), LocalUploadedStorageDir: env("AI_GATEWAY_UPLOADED_STORAGE_DIR", env("LOCAL_UPLOADED_STORAGE_DIR", DefaultLocalUploadedStorageDir)),
LocalTempAssetTTLHours: envInt("AI_GATEWAY_LOCAL_TEMP_ASSET_TTL_HOURS", 24), LocalTempAssetTTLHours: envInt("AI_GATEWAY_LOCAL_TEMP_ASSET_TTL_HOURS", 24),
LocalResultTTLHours: envIntValidated("AI_GATEWAY_LOCAL_RESULT_TTL_HOURS", 24),
LocalResultMinFreeBytes: envInt64Validated("AI_GATEWAY_LOCAL_RESULT_MIN_FREE_BYTES", 10*1024*1024*1024),
LocalResultMaxBytes: envInt64Validated("AI_GATEWAY_LOCAL_RESULT_MAX_BYTES", 256*1024*1024),
LocalResultMaxTaskBytes: envInt64Validated("AI_GATEWAY_LOCAL_RESULT_MAX_TASK_BYTES", 512*1024*1024),
TaskProgressCallbackEnabled: env("TASK_PROGRESS_CALLBACK_ENABLED", "true") == "true", TaskProgressCallbackEnabled: env("TASK_PROGRESS_CALLBACK_ENABLED", "true") == "true",
TaskProgressCallbackURL: env("TASK_PROGRESS_CALLBACK_URL", TaskProgressCallbackURL: env("TASK_PROGRESS_CALLBACK_URL",
strings.TrimRight(env("SERVER_MAIN_BASE_URL", "http://localhost:3000"), "/")+"/internal/platform/task-progress-callbacks", strings.TrimRight(env("SERVER_MAIN_BASE_URL", "http://localhost:3000"), "/")+"/internal/platform/task-progress-callbacks",
@@ -139,6 +147,18 @@ func (c Config) Validate() error {
if c.TaskCleanupBatchSize != 0 && (c.TaskCleanupBatchSize < 100 || c.TaskCleanupBatchSize > 5000) { 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") return errors.New("AI_GATEWAY_TASK_CLEANUP_BATCH_SIZE must be between 100 and 5000")
} }
if c.LocalResultTTLHours != 0 && (c.LocalResultTTLHours < 1 || c.LocalResultTTLHours > 24*30) {
return errors.New("AI_GATEWAY_LOCAL_RESULT_TTL_HOURS must be between 1 and 720")
}
if c.LocalResultMinFreeBytes < 0 {
return errors.New("AI_GATEWAY_LOCAL_RESULT_MIN_FREE_BYTES must not be negative")
}
if c.LocalResultMaxBytes != 0 && c.LocalResultMaxBytes < 1024 {
return errors.New("AI_GATEWAY_LOCAL_RESULT_MAX_BYTES must be at least 1024")
}
if c.LocalResultMaxTaskBytes != 0 && c.LocalResultMaxTaskBytes < c.LocalResultMaxBytes {
return errors.New("AI_GATEWAY_LOCAL_RESULT_MAX_TASK_BYTES must be at least AI_GATEWAY_LOCAL_RESULT_MAX_BYTES")
}
switch strings.ToLower(strings.TrimSpace(c.IdentitySecretStore)) { switch strings.ToLower(strings.TrimSpace(c.IdentitySecretStore)) {
case "": case "":
case "file": case "file":
@@ -260,6 +280,18 @@ func envIntValidated(key string, fallback int) int {
return parsed return parsed
} }
func envInt64Validated(key string, fallback int64) int64 {
value := envValue(key)
if value == "" {
return fallback
}
parsed, err := strconv.ParseInt(value, 10, 64)
if err != nil {
return 0
}
return parsed
}
func logLevel(value string) slog.Level { func logLevel(value string) slog.Level {
switch strings.ToLower(value) { switch strings.ToLower(value) {
case "debug": case "debug":
@@ -0,0 +1,30 @@
package httpapi
import (
"context"
"net/http"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/clients"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
func (s *Server) hydrateTaskResult(ctx context.Context, task store.GatewayTask) (store.GatewayTask, error) {
if task.Status != "succeeded" || len(task.Result) == 0 {
return task, nil
}
result, err := s.runner.HydrateTaskResult(ctx, task.ID, task.Result)
if err != nil {
return store.GatewayTask{}, err
}
task.Result = result
return task, nil
}
func writeStoredBinaryResultError(w http.ResponseWriter, err error) {
status := statusFromRunError(err)
code := clients.ErrorCode(err)
if code == "" {
code = "binary_result_corrupted"
}
writeError(w, status, err.Error(), code)
}
@@ -170,6 +170,10 @@ func (s *Server) geminiGenerateContent(w http.ResponseWriter, r *http.Request) {
writeGeminiTaskError(http.StatusConflict, "Idempotency-Key was reused for a different request", nil, "idempotency_key_reused") writeGeminiTaskError(http.StatusConflict, "Idempotency-Key was reused for a different request", nil, "idempotency_key_reused")
return return
} }
if errors.Is(err, store.ErrTaskRequestBinaryNotMaterialized) {
writeGeminiTaskError(http.StatusBadRequest, err.Error(), nil, clients.ErrorCode(err))
return
}
s.logger.Error("create gemini task failed", "kind", mapping.Kind, "error_category", "task_create_failed") s.logger.Error("create gemini task failed", "kind", mapping.Kind, "error_category", "task_create_failed")
writeGeminiTaskError(http.StatusInternalServerError, "create task failed", nil, "task_create_failed") writeGeminiTaskError(http.StatusInternalServerError, "create task failed", nil, "task_create_failed")
return return
@@ -184,6 +188,11 @@ func (s *Server) geminiGenerateContent(w http.ResponseWriter, r *http.Request) {
return return
} }
if task.Status == "succeeded" { if task.Status == "succeeded" {
task, err = s.hydrateTaskResult(r.Context(), task)
if err != nil {
writeGeminiTaskError(statusFromRunError(err), err.Error(), nil, clients.ErrorCode(err))
return
}
if _, native := task.Result["candidates"]; native { if _, native := task.Result["candidates"]; native {
writeJSON(w, http.StatusOK, task.Result) writeJSON(w, http.StatusOK, task.Result)
return return
+27
View File
@@ -1134,6 +1134,10 @@ func (s *Server) createTask(kind string, compatible bool) http.Handler {
writeTaskError(http.StatusConflict, "Idempotency-Key was reused for a different request", nil, "idempotency_key_reused") writeTaskError(http.StatusConflict, "Idempotency-Key was reused for a different request", nil, "idempotency_key_reused")
return return
} }
if errors.Is(err, store.ErrTaskRequestBinaryNotMaterialized) {
writeTaskError(http.StatusBadRequest, err.Error(), nil, clients.ErrorCode(err))
return
}
s.logger.Error("create task failed", "kind", kind, "error_category", "task_create_failed") s.logger.Error("create task failed", "kind", kind, "error_category", "task_create_failed")
writeTaskError(http.StatusInternalServerError, "create task failed", nil, "task_create_failed") writeTaskError(http.StatusInternalServerError, "create task failed", nil, "task_create_failed")
return return
@@ -1150,6 +1154,17 @@ func (s *Server) createTask(kind string, compatible bool) http.Handler {
writeTaskError(http.StatusConflict, "streaming idempotent replay is not supported", nil, "idempotency_stream_replay_unsupported") writeTaskError(http.StatusConflict, "streaming idempotent replay is not supported", nil, "idempotency_stream_replay_unsupported")
return return
} }
if task.Status == "succeeded" && !responsePlan.asyncMode && (targetProtocol != "" || responsePlan.compatibleMode) {
task, err = s.hydrateTaskResult(r.Context(), task)
if err != nil {
if targetProtocol != "" {
writeProtocolError(w, targetProtocol, statusFromRunError(err), err.Error(), nil, clients.ErrorCode(err))
} else {
writeStoredBinaryResultError(w, err)
}
return
}
}
if targetProtocol != "" && !responsePlan.asyncMode { if targetProtocol != "" && !responsePlan.asyncMode {
if task.Status == "succeeded" { if task.Status == "succeeded" {
writeJSON(w, http.StatusOK, task.Result) writeJSON(w, http.StatusOK, task.Result)
@@ -1597,6 +1612,12 @@ func scopeForTaskKind(kind string) string {
func statusFromRunError(err error) int { func statusFromRunError(err error) int {
switch { switch {
case clients.ErrorCode(err) == "binary_result_expired":
return http.StatusGone
case clients.ErrorCode(err) == "binary_result_corrupted" || clients.ErrorCode(err) == "result_binary_not_materialized":
return http.StatusInternalServerError
case clients.ErrorCode(err) == "local_result_storage_unavailable":
return http.StatusServiceUnavailable
case clients.ErrorCode(err) == "billing_hold": case clients.ErrorCode(err) == "billing_hold":
return http.StatusServiceUnavailable return http.StatusServiceUnavailable
case runner.IsPricingUnavailable(err): case runner.IsPricingUnavailable(err):
@@ -1901,6 +1922,7 @@ func boolValue(body map[string]any, key string) bool {
// @Success 200 {object} store.GatewayTask // @Success 200 {object} store.GatewayTask
// @Failure 401 {object} ErrorEnvelope // @Failure 401 {object} ErrorEnvelope
// @Failure 404 {object} ErrorEnvelope // @Failure 404 {object} ErrorEnvelope
// @Failure 410 {object} ErrorEnvelope
// @Failure 500 {object} ErrorEnvelope // @Failure 500 {object} ErrorEnvelope
// @Router /api/workspace/tasks/{taskID} [get] // @Router /api/workspace/tasks/{taskID} [get]
// @Router /api/v1/tasks/{taskID} [get] // @Router /api/v1/tasks/{taskID} [get]
@@ -1920,6 +1942,11 @@ func (s *Server) getTask(w http.ResponseWriter, r *http.Request) {
task.Cancellable = &cancelState.Cancellable task.Cancellable = &cancelState.Cancellable
task.Submitted = &cancelState.Submitted task.Submitted = &cancelState.Submitted
task.Message = cancelState.Message task.Message = cancelState.Message
task, err = s.hydrateTaskResult(r.Context(), task)
if err != nil {
writeStoredBinaryResultError(w, err)
return
}
writeJSON(w, http.StatusOK, task) writeJSON(w, http.StatusOK, task)
return return
} }
@@ -177,6 +177,10 @@ func (s *Server) createKelingOmniVideo(w http.ResponseWriter, r *http.Request) {
writeKelingCompatError(w, requestID, kelingCompatGatewayError(staged.Err)) writeKelingCompatError(w, requestID, kelingCompatGatewayError(staged.Err))
return return
} }
if errors.Is(createErr, store.ErrTaskRequestBinaryNotMaterialized) {
writeKelingCompatError(w, requestID, newKelingCompatError(http.StatusBadRequest, 1201, clients.ErrorCode(createErr)))
return
}
s.logger.Error("create Kling-compatible task failed", "error", createErr) s.logger.Error("create Kling-compatible task failed", "error", createErr)
writeKelingCompatError(w, requestID, newKelingCompatError(http.StatusInternalServerError, 5000, "create task failed")) writeKelingCompatError(w, requestID, newKelingCompatError(http.StatusInternalServerError, 5000, "create task failed"))
return return
@@ -221,6 +225,7 @@ func (s *Server) createKelingOmniVideo(w http.ResponseWriter, r *http.Request) {
// @Failure 401 {object} KelingCompatibleEnvelope // @Failure 401 {object} KelingCompatibleEnvelope
// @Failure 403 {object} KelingCompatibleEnvelope // @Failure 403 {object} KelingCompatibleEnvelope
// @Failure 404 {object} KelingCompatibleEnvelope // @Failure 404 {object} KelingCompatibleEnvelope
// @Failure 410 {object} KelingCompatibleEnvelope
// @Failure 500 {object} KelingCompatibleEnvelope // @Failure 500 {object} KelingCompatibleEnvelope
// @Router /api/v1/videos/omni-video/{taskID} [get] // @Router /api/v1/videos/omni-video/{taskID} [get]
func (s *Server) getKelingOmniVideo(w http.ResponseWriter, r *http.Request) { func (s *Server) getKelingOmniVideo(w http.ResponseWriter, r *http.Request) {
@@ -247,6 +252,11 @@ func (s *Server) getKelingOmniVideo(w http.ResponseWriter, r *http.Request) {
writeKelingCompatError(w, requestID, newKelingCompatError(http.StatusNotFound, 1203, "task not found")) writeKelingCompatError(w, requestID, newKelingCompatError(http.StatusNotFound, 1203, "task not found"))
return return
} }
task, err = s.hydrateTaskResult(r.Context(), task)
if err != nil {
writeKelingCompatError(w, requestID, newKelingCompatError(statusFromRunError(err), 5000, clients.ErrorCode(err)))
return
}
writeJSON(w, http.StatusOK, KelingCompatibleEnvelope{ writeJSON(w, http.StatusOK, KelingCompatibleEnvelope{
Code: 0, Code: 0,
Message: "SUCCEED", Message: "SUCCEED",
+14
View File
@@ -144,6 +144,8 @@ func (s *Server) createKlingCompatTask(w http.ResponseWriter, r *http.Request, v
writeKlingCompatError(w, http.StatusConflict, err.Error(), "idempotency_key_reused") writeKlingCompatError(w, http.StatusConflict, err.Error(), "idempotency_key_reused")
case errors.Is(err, store.ErrExternalTaskIDReused): case errors.Is(err, store.ErrExternalTaskIDReused):
writeKlingCompatError(w, http.StatusConflict, err.Error(), "external_task_id_reused") writeKlingCompatError(w, http.StatusConflict, err.Error(), "external_task_id_reused")
case errors.Is(err, store.ErrTaskRequestBinaryNotMaterialized):
writeKlingCompatError(w, http.StatusBadRequest, err.Error(), clients.ErrorCode(err))
default: default:
s.logger.Error("create Kling compatibility task failed", "version", version, "model", model, "error", err) s.logger.Error("create Kling compatibility task failed", "version", version, "model", model, "error", err)
writeKlingCompatError(w, http.StatusInternalServerError, "create task failed", "task_create_failed") writeKlingCompatError(w, http.StatusInternalServerError, "create task failed", "task_create_failed")
@@ -464,6 +466,7 @@ func validateKlingCompatBody(model string, body map[string]any) error {
// @Param taskID path string true "任务 ID" // @Param taskID path string true "任务 ID"
// @Success 200 {object} map[string]interface{} // @Success 200 {object} map[string]interface{}
// @Failure 404 {object} map[string]interface{} // @Failure 404 {object} map[string]interface{}
// @Failure 410 {object} map[string]interface{}
// @Router /api/v1/kling/v1/videos/omni-video/{taskID} [get] // @Router /api/v1/kling/v1/videos/omni-video/{taskID} [get]
func (s *Server) klingV1GetOmniVideo(w http.ResponseWriter, r *http.Request) { func (s *Server) klingV1GetOmniVideo(w http.ResponseWriter, r *http.Request) {
user, _ := auth.UserFromContext(r.Context()) user, _ := auth.UserFromContext(r.Context())
@@ -476,6 +479,11 @@ func (s *Server) klingV1GetOmniVideo(w http.ResponseWriter, r *http.Request) {
writeKlingCompatError(w, http.StatusInternalServerError, "get task failed", "task_query_failed") writeKlingCompatError(w, http.StatusInternalServerError, "get task failed", "task_query_failed")
return return
} }
task, err = s.hydrateTaskResult(r.Context(), task)
if err != nil {
writeKlingCompatError(w, statusFromRunError(err), err.Error(), clients.ErrorCode(err))
return
}
writeJSON(w, http.StatusOK, klingV1Envelope(task)) writeJSON(w, http.StatusOK, klingV1Envelope(task))
} }
@@ -520,6 +528,7 @@ func (s *Server) klingV1ListOmniVideos(w http.ResponseWriter, r *http.Request) {
// @Param task_ids query string false "逗号分隔的任务 ID" // @Param task_ids query string false "逗号分隔的任务 ID"
// @Param external_task_ids query string false "逗号分隔的外部任务 ID" // @Param external_task_ids query string false "逗号分隔的外部任务 ID"
// @Success 200 {object} map[string]interface{} // @Success 200 {object} map[string]interface{}
// @Failure 410 {object} map[string]interface{}
// @Router /api/v1/kling/v2/tasks [get] // @Router /api/v1/kling/v2/tasks [get]
func (s *Server) klingV2GetTasks(w http.ResponseWriter, r *http.Request) { func (s *Server) klingV2GetTasks(w http.ResponseWriter, r *http.Request) {
user, _ := auth.UserFromContext(r.Context()) user, _ := auth.UserFromContext(r.Context())
@@ -543,6 +552,11 @@ func (s *Server) klingV2GetTasks(w http.ResponseWriter, r *http.Request) {
writeKlingCompatError(w, http.StatusInternalServerError, "get task failed", "task_query_failed") writeKlingCompatError(w, http.StatusInternalServerError, "get task failed", "task_query_failed")
return return
} }
task, err = s.hydrateTaskResult(r.Context(), task)
if err != nil {
writeKlingCompatError(w, statusFromRunError(err), err.Error(), clients.ErrorCode(err))
return
}
data = append(data, klingV2TaskData(task)) data = append(data, klingV2TaskData(task))
} }
writeJSON(w, http.StatusOK, map[string]any{"code": 0, "message": "success", "request_id": klingRequestIDFromAny(data), "data": data}) writeJSON(w, http.StatusOK, map[string]any{"code": 0, "message": "success", "request_id": klingRequestIDFromAny(data), "data": data})
@@ -45,6 +45,58 @@ func (s *Server) cleanupExpiredLocalTempAssets(ctx context.Context, now time.Tim
for _, target := range targets { for _, target := range targets {
deleted += s.cleanupExpiredLocalTempAssetsInDir(ctx, now, target) deleted += s.cleanupExpiredLocalTempAssetsInDir(ctx, now, target)
} }
deleted += s.cleanupExpiredLocalBinaryResults(now)
return deleted
}
func (s *Server) cleanupExpiredLocalBinaryResults(now time.Time) int {
storageDir := strings.TrimSpace(s.cfg.LocalGeneratedStorageDir)
if storageDir == "" {
storageDir = config.DefaultLocalGeneratedStorageDir
}
root := filepath.Join(storageDir, "results")
ttlHours := s.cfg.LocalResultTTLHours
if ttlHours <= 0 {
ttlHours = 24
}
expiredBefore := now.Add(-time.Duration(ttlHours) * time.Hour)
deleted := 0
var directories []string
err := filepath.WalkDir(root, func(path string, entry os.DirEntry, walkErr error) error {
if walkErr != nil {
if !errors.Is(walkErr, os.ErrNotExist) && s.logger != nil {
s.logger.Warn("walk local binary result failed", "path", path, "error", walkErr)
}
return nil
}
if entry.IsDir() {
if path != root {
directories = append(directories, path)
}
return nil
}
if entry.Type()&os.ModeSymlink != 0 {
return nil
}
info, infoErr := entry.Info()
if infoErr != nil || info.ModTime().After(expiredBefore) {
return nil
}
if removeErr := os.Remove(path); removeErr != nil && !errors.Is(removeErr, os.ErrNotExist) {
if s.logger != nil {
s.logger.Warn("remove local binary result failed", "path", path, "error", removeErr)
}
return nil
}
deleted++
return nil
})
if err != nil && !errors.Is(err, os.ErrNotExist) && s.logger != nil {
s.logger.Warn("scan local binary result root failed", "dir", root, "error", err)
}
for index := len(directories) - 1; index >= 0; index-- {
_ = os.Remove(directories[index])
}
return deleted return deleted
} }
@@ -11,6 +11,7 @@ import (
"net/http/httptest" "net/http/httptest"
"os" "os"
"path/filepath" "path/filepath"
"strings"
"testing" "testing"
"time" "time"
@@ -235,6 +236,47 @@ func TestCleanupExpiredLocalTempAssetsDeletesExpiredStaticFiles(t *testing.T) {
} }
} }
func TestCleanupExpiredLocalTempAssetsDeletesNestedBinaryResultsWithIndependentTTL(t *testing.T) {
generatedDir := t.TempDir()
taskDir := filepath.Join(generatedDir, "results", "task-123")
if err := os.MkdirAll(taskDir, 0o750); err != nil {
t.Fatalf("create local result fixture dir: %v", err)
}
oldResult := filepath.Join(taskDir, strings.Repeat("a", 64)+".bin")
freshResult := filepath.Join(taskDir, strings.Repeat("b", 64)+".bin")
for _, path := range []string{oldResult, freshResult} {
if err := os.WriteFile(path, []byte("asset"), 0o640); err != nil {
t.Fatalf("write fixture %s: %v", path, err)
}
}
now := time.Now()
if err := os.Chtimes(oldResult, now.Add(-25*time.Hour), now.Add(-25*time.Hour)); err != nil {
t.Fatalf("age old binary result: %v", err)
}
if err := os.Chtimes(freshResult, now.Add(-23*time.Hour), now.Add(-23*time.Hour)); err != nil {
t.Fatalf("age fresh binary result: %v", err)
}
server := &Server{
cfg: config.Config{
LocalGeneratedStorageDir: generatedDir,
LocalResultTTLHours: 24,
},
logger: slog.New(slog.NewTextHandler(io.Discard, nil)),
}
deleted := server.cleanupExpiredLocalTempAssets(context.Background(), now)
if deleted != 1 {
t.Fatalf("expected one expired binary result delete, got %d", deleted)
}
if _, err := os.Stat(oldResult); !os.IsNotExist(err) {
t.Fatalf("old binary result should be deleted, stat err=%v", err)
}
if _, err := os.Stat(freshResult); err != nil {
t.Fatalf("fresh binary result should remain: %v", err)
}
}
func TestRequestConversationKeyPriority(t *testing.T) { func TestRequestConversationKeyPriority(t *testing.T) {
request := httptest.NewRequest(http.MethodPost, "/api/v1/chat/completions", nil) request := httptest.NewRequest(http.MethodPost, "/api/v1/chat/completions", nil)
request.Header.Set("X-EasyAI-Conversation-ID", "from-header") request.Header.Set("X-EasyAI-Conversation-ID", "from-header")
@@ -56,12 +56,19 @@ func (s *Server) createVolcesContentsGenerationTask(w http.ResponseWriter, r *ht
// @Param taskID path string true "任务 ID" // @Param taskID path string true "任务 ID"
// @Success 200 {object} VolcesContentsGenerationTaskResponse // @Success 200 {object} VolcesContentsGenerationTaskResponse
// @Failure 404 {object} VolcesErrorEnvelope // @Failure 404 {object} VolcesErrorEnvelope
// @Failure 410 {object} VolcesErrorEnvelope
// @Router /api/v1/contents/generations/tasks/{taskID} [get] // @Router /api/v1/contents/generations/tasks/{taskID} [get]
func (s *Server) getVolcesContentsGenerationTask(w http.ResponseWriter, r *http.Request) { func (s *Server) getVolcesContentsGenerationTask(w http.ResponseWriter, r *http.Request) {
task, ok := s.volcesCompatibleTaskForUser(w, r) task, ok := s.volcesCompatibleTaskForUser(w, r)
if !ok { if !ok {
return return
} }
var err error
task, err = s.hydrateTaskResult(r.Context(), task)
if err != nil {
writeVolcesError(w, statusFromRunError(err), err.Error(), clients.ErrorCode(err))
return
}
writeJSON(w, http.StatusOK, volcesCompatibleTask(task)) writeJSON(w, http.StatusOK, volcesCompatibleTask(task))
} }
@@ -178,6 +185,7 @@ func (s *Server) createLegacyVolcesVideoGeneration(w http.ResponseWriter, r *htt
// @Param taskID path string true "任务 ID" // @Param taskID path string true "任务 ID"
// @Success 200 {object} EasyAIGeneratedResponse // @Success 200 {object} EasyAIGeneratedResponse
// @Failure 404 {object} EasyAIGeneratedResponse // @Failure 404 {object} EasyAIGeneratedResponse
// @Failure 410 {object} EasyAIGeneratedResponse
// @Router /api/v1/ai/result/{taskID} [get] // @Router /api/v1/ai/result/{taskID} [get]
func (s *Server) getEasyAITaskResult(w http.ResponseWriter, r *http.Request) { func (s *Server) getEasyAITaskResult(w http.ResponseWriter, r *http.Request) {
user, ok := auth.UserFromContext(r.Context()) user, ok := auth.UserFromContext(r.Context())
@@ -199,6 +207,11 @@ func (s *Server) getEasyAITaskResult(w http.ResponseWriter, r *http.Request) {
writeEasyAIAsyncError(w, http.StatusNotFound, "task not found", nil, "not_found") writeEasyAIAsyncError(w, http.StatusNotFound, "task not found", nil, "not_found")
return return
} }
task, err = s.hydrateTaskResult(r.Context(), task)
if err != nil {
writeEasyAIAsyncError(w, statusFromRunError(err), err.Error(), nil, clients.ErrorCode(err))
return
}
writeJSON(w, http.StatusOK, easyAITaskResultResponse(task)) writeJSON(w, http.StatusOK, easyAITaskResultResponse(task))
} }
@@ -408,6 +421,8 @@ func writeVolcesCompatibleTaskError(w http.ResponseWriter, err error) {
status = clientErr.StatusCode status = clientErr.StatusCode
} else if errors.As(err, &clientErr) { } else if errors.As(err, &clientErr) {
status = http.StatusBadRequest status = http.StatusBadRequest
} else if errors.Is(err, store.ErrTaskRequestBinaryNotMaterialized) {
status = http.StatusBadRequest
} }
if wire := clients.ErrorWireResponse(err); wireResponseMatches(wire, clients.ProtocolVolcesContents) { if wire := clients.ErrorWireResponse(err); wireResponseMatches(wire, clients.ProtocolVolcesContents) {
writeWireResponse(w, wire) writeWireResponse(w, wire)
+721
View File
@@ -0,0 +1,721 @@
package runner
import (
"context"
"crypto/sha256"
"encoding/base64"
"encoding/hex"
"errors"
"fmt"
"io"
"os"
"path/filepath"
"strconv"
"strings"
"syscall"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/clients"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
)
const (
localBinaryResultDirName = "results"
localBinaryPlaceholderPrefix = "[GatewayBinary:v1;"
localBinaryGenericBase64MinLength = 4096
localBinaryMaxDepth = 64
defaultLocalResultTTLHours = 24
defaultLocalResultMinFreeBytes = int64(10 * 1024 * 1024 * 1024)
defaultLocalResultMaxBytes = int64(256 * 1024 * 1024)
defaultLocalResultMaxTaskBytes = int64(512 * 1024 * 1024)
)
type localBinaryDescriptor struct {
Prefix string
SHA256 string
Size int64
ContentType string
Encoding string
}
type localBinaryMaterializer struct {
service *Service
taskDir string
writeFiles bool
enforceLimits bool
totalBytes int64
seen map[string]struct{}
createdFiles []string
}
// materializeLocalBinaryResult replaces every inline binary value with a
// bounded placeholder after atomically writing and verifying the bytes locally.
func (s *Service) materializeLocalBinaryResult(ctx context.Context, taskID string, result map[string]any) (map[string]any, bool, error) {
if _, _, err := s.transformLocalBinaryResult(ctx, taskID, result, false, true); err != nil {
return nil, false, err
}
return s.transformLocalBinaryResult(ctx, taskID, result, true, true)
}
func (s *Service) transformLocalBinaryResult(ctx context.Context, taskID string, result map[string]any, writeFiles bool, enforceLimits bool) (map[string]any, bool, error) {
root := s.localBinaryResultRoot()
taskDir := filepath.Join(root, safeLocalBinaryTaskDir(taskID))
materializer := &localBinaryMaterializer{
service: s,
taskDir: taskDir,
writeFiles: writeFiles,
enforceLimits: enforceLimits,
seen: map[string]struct{}{},
}
next, changed, err := materializer.materializeValue(ctx, result, "", nil, 0)
if err != nil {
materializer.rollbackCreatedFiles()
return nil, false, err
}
mapped, ok := next.(map[string]any)
if !ok {
materializer.rollbackCreatedFiles()
return nil, false, &clients.ClientError{
Code: "result_binary_not_materialized",
Message: "generated result is not a JSON object",
StatusCode: 500,
Retryable: false,
}
}
if changed && enforceLimits && !writeFiles {
if err := os.MkdirAll(root, 0o750); err != nil {
return nil, false, localBinaryStorageError(err)
}
if err := ensureLocalBinaryDiskHeadroom(root, materializer.totalBytes, s.localResultMinFreeBytes()); err != nil {
return nil, false, err
}
}
return mapped, changed, nil
}
// MaterializeTaskResultForStorage exposes the same verified materialization
// path to the explicit historical maintenance command.
func (s *Service) MaterializeTaskResultForStorage(ctx context.Context, taskID string, result map[string]any) (map[string]any, bool, error) {
return s.materializeLocalBinaryResult(ctx, taskID, result)
}
// CompactExpiredTaskResultForStorage computes the same deterministic
// placeholders without creating files that are already outside the recovery
// window.
func (s *Service) CompactExpiredTaskResultForStorage(ctx context.Context, taskID string, result map[string]any) (map[string]any, bool, error) {
return s.transformLocalBinaryResult(ctx, taskID, result, false, false)
}
func TaskResultHasInlineBinary(result map[string]any) bool {
return localBinaryValueHasPayload(result, "", nil, 0)
}
func localBinaryValueHasPayload(value any, key string, siblings map[string]any, depth int) bool {
if depth >= localBinaryMaxDepth {
return false
}
switch typed := value.(type) {
case map[string]any:
if _, _, ok := localBufferObjectBytes(typed); ok {
return true
}
for childKey, child := range typed {
if localBinaryValueHasPayload(child, childKey, typed, depth+1) {
return true
}
}
case []any:
if localBinaryKey(key) {
if _, ok := bytesFromNumberArray(typed); ok {
return true
}
}
for _, child := range typed {
if localBinaryValueHasPayload(child, key, siblings, depth+1) {
return true
}
}
case []byte:
return len(typed) > 0
case string:
_, _, _, ok := localBinaryStringBytes(key, typed, siblings)
return ok
}
return false
}
func (m *localBinaryMaterializer) materializeValue(ctx context.Context, value any, key string, siblings map[string]any, depth int) (any, bool, error) {
if err := ctx.Err(); err != nil {
return nil, false, err
}
if depth >= localBinaryMaxDepth {
return nil, false, &clients.ClientError{
Code: "result_binary_not_materialized",
Message: "generated result exceeds the maximum JSON depth",
StatusCode: 500,
Retryable: false,
}
}
switch typed := value.(type) {
case map[string]any:
if payload, contentType, ok := localBufferObjectBytes(typed); ok {
return m.persistBinary(ctx, payload, contentType, "buffer")
}
next := make(map[string]any, len(typed))
changed := false
for childKey, childValue := range typed {
child, childChanged, err := m.materializeValue(ctx, childValue, childKey, typed, depth+1)
if err != nil {
return nil, false, err
}
next[childKey] = child
changed = changed || childChanged
}
if !changed {
return value, false, nil
}
return next, true, nil
case []any:
if localBinaryKey(key) {
if payload, ok := bytesFromNumberArray(typed); ok {
return m.persistBinary(ctx, payload, mediaContentTypeFromItem(siblings), "buffer")
}
}
next := make([]any, len(typed))
changed := false
for index, item := range typed {
child, childChanged, err := m.materializeValue(ctx, item, key, siblings, depth+1)
if err != nil {
return nil, false, err
}
next[index] = child
changed = changed || childChanged
}
if !changed {
return value, false, nil
}
return next, true, nil
case []byte:
if len(typed) == 0 {
return value, false, nil
}
return m.persistBinary(ctx, append([]byte(nil), typed...), mediaContentTypeFromItem(siblings), "buffer")
case string:
payload, contentType, encoding, ok := localBinaryStringBytes(key, typed, siblings)
if !ok {
return value, false, nil
}
return m.persistBinary(ctx, payload, contentType, encoding)
default:
return value, false, nil
}
}
func (m *localBinaryMaterializer) persistBinary(ctx context.Context, payload []byte, contentType string, encoding string) (any, bool, error) {
if err := ctx.Err(); err != nil {
return nil, false, err
}
if len(payload) == 0 {
return nil, false, nil
}
size := int64(len(payload))
if m.enforceLimits && size > m.service.localResultMaxBytes() {
return nil, false, &clients.ClientError{
Code: "binary_result_too_large",
Message: "one generated binary result exceeds the local storage limit",
StatusCode: 502,
Retryable: false,
}
}
digest := sha256.Sum256(payload)
digestHex := hex.EncodeToString(digest[:])
if _, exists := m.seen[digestHex]; !exists {
if m.enforceLimits && m.totalBytes+size > m.service.localResultMaxTaskBytes() {
return nil, false, &clients.ClientError{
Code: "binary_result_too_large",
Message: "generated binary results exceed the per-task local storage limit",
StatusCode: 502,
Retryable: false,
}
}
if m.writeFiles {
created, err := m.service.writeLocalBinaryResult(m.taskDir, digestHex, payload)
if err != nil {
return nil, false, err
}
if created {
m.createdFiles = append(m.createdFiles, filepath.Join(m.taskDir, digestHex+".bin"))
}
}
m.seen[digestHex] = struct{}{}
m.totalBytes += size
}
descriptor := localBinaryDescriptor{
Prefix: localBinaryPrefix(payload),
SHA256: digestHex,
Size: size,
ContentType: normalizedLocalBinaryContentType(contentType),
Encoding: normalizedLocalBinaryEncoding(encoding),
}
return localBinaryPlaceholder(descriptor), true, nil
}
func (m *localBinaryMaterializer) rollbackCreatedFiles() {
for _, path := range m.createdFiles {
_ = os.Remove(path)
}
_ = os.Remove(m.taskDir)
m.createdFiles = nil
}
func (s *Service) writeLocalBinaryResult(taskDir string, digestHex string, payload []byte) (bool, error) {
if err := os.MkdirAll(taskDir, 0o750); err != nil {
return false, localBinaryStorageError(err)
}
targetPath := filepath.Join(taskDir, digestHex+".bin")
if info, err := os.Stat(targetPath); err == nil {
if !info.IsDir() && info.Size() == int64(len(payload)) {
if err := verifyLocalBinaryFile(targetPath, digestHex, int64(len(payload))); err == nil {
now := time.Now()
if err := os.Chtimes(targetPath, now, now); err != nil {
return false, localBinaryStorageError(err)
}
return false, nil
}
}
return false, &clients.ClientError{
Code: "binary_result_corrupted",
Message: "existing local result file does not match its content hash",
StatusCode: 500,
Retryable: false,
}
} else if !errors.Is(err, os.ErrNotExist) {
return false, localBinaryStorageError(err)
}
if err := ensureLocalBinaryDiskHeadroom(taskDir, int64(len(payload)), s.localResultMinFreeBytes()); err != nil {
return false, err
}
tempFile, err := os.CreateTemp(taskDir, ".gateway-result-*")
if err != nil {
return false, localBinaryStorageError(err)
}
tempPath := tempFile.Name()
cleanup := func() {
_ = tempFile.Close()
_ = os.Remove(tempPath)
}
if err := tempFile.Chmod(0o640); err != nil {
cleanup()
return false, localBinaryStorageError(err)
}
if _, err := tempFile.Write(payload); err != nil {
cleanup()
return false, localBinaryStorageError(err)
}
if err := tempFile.Sync(); err != nil {
cleanup()
return false, localBinaryStorageError(err)
}
if err := tempFile.Close(); err != nil {
_ = os.Remove(tempPath)
return false, localBinaryStorageError(err)
}
if err := os.Rename(tempPath, targetPath); err != nil {
_ = os.Remove(tempPath)
return false, localBinaryStorageError(err)
}
if err := verifyLocalBinaryFile(targetPath, digestHex, int64(len(payload))); err != nil {
_ = os.Remove(targetPath)
return false, err
}
return true, nil
}
func ensureLocalBinaryDiskHeadroom(path string, incomingBytes int64, minFreeBytes int64) error {
var stat syscall.Statfs_t
if err := syscall.Statfs(path, &stat); err != nil {
return localBinaryStorageError(err)
}
freeBytes := int64(stat.Bavail) * int64(stat.Bsize)
if freeBytes-incomingBytes < minFreeBytes {
return &clients.ClientError{
Code: "local_result_storage_unavailable",
Message: "local result storage does not have enough free space",
StatusCode: 503,
Retryable: false,
}
}
return nil
}
func localBinaryStorageError(err error) error {
return &clients.ClientError{
Code: "local_result_storage_unavailable",
Message: "local result storage failed: " + err.Error(),
StatusCode: 503,
Retryable: false,
}
}
func verifyLocalBinaryFile(path string, expectedHash string, expectedSize int64) error {
file, err := os.Open(path)
if err != nil {
return localBinaryStorageError(err)
}
defer file.Close()
hasher := sha256.New()
size, err := io.Copy(hasher, file)
if err != nil {
return localBinaryStorageError(err)
}
if size != expectedSize || hex.EncodeToString(hasher.Sum(nil)) != expectedHash {
return &clients.ClientError{
Code: "binary_result_corrupted",
Message: "local result file failed size or hash verification",
StatusCode: 500,
Retryable: false,
}
}
return nil
}
// HydrateTaskResult restores placeholders from verified local files. It is only
// used by result/detail/replay endpoints, never by task lists or callbacks.
func (s *Service) HydrateTaskResult(ctx context.Context, taskID string, result map[string]any) (map[string]any, error) {
next, changed, err := s.hydrateLocalBinaryValue(ctx, taskID, result, 0)
if err != nil {
return nil, err
}
if !changed {
return result, nil
}
mapped, ok := next.(map[string]any)
if !ok {
return nil, &clients.ClientError{Code: "binary_result_corrupted", Message: "stored result is not a JSON object", StatusCode: 500}
}
return mapped, nil
}
func (s *Service) hydrateLocalBinaryValue(ctx context.Context, taskID string, value any, depth int) (any, bool, error) {
if err := ctx.Err(); err != nil {
return nil, false, err
}
if depth >= localBinaryMaxDepth {
return nil, false, &clients.ClientError{Code: "binary_result_corrupted", Message: "stored result exceeds the maximum JSON depth", StatusCode: 500}
}
switch typed := value.(type) {
case map[string]any:
next := make(map[string]any, len(typed))
changed := false
for key, childValue := range typed {
child, childChanged, err := s.hydrateLocalBinaryValue(ctx, taskID, childValue, depth+1)
if err != nil {
return nil, false, err
}
next[key] = child
changed = changed || childChanged
}
if !changed {
return value, false, nil
}
return next, true, nil
case []any:
next := make([]any, len(typed))
changed := false
for index, childValue := range typed {
child, childChanged, err := s.hydrateLocalBinaryValue(ctx, taskID, childValue, depth+1)
if err != nil {
return nil, false, err
}
next[index] = child
changed = changed || childChanged
}
if !changed {
return value, false, nil
}
return next, true, nil
case string:
descriptor, ok := parseLocalBinaryPlaceholder(typed)
if !ok {
return value, false, nil
}
payload, err := s.readLocalBinaryResult(taskID, descriptor)
if err != nil {
return nil, false, err
}
encoded := base64.StdEncoding.EncodeToString(payload)
if descriptor.Encoding == "data-uri" {
return "data:" + descriptor.ContentType + ";base64," + encoded, true, nil
}
return encoded, true, nil
default:
return value, false, nil
}
}
func (s *Service) readLocalBinaryResult(taskID string, descriptor localBinaryDescriptor) ([]byte, error) {
path := filepath.Join(s.localBinaryResultRoot(), safeLocalBinaryTaskDir(taskID), descriptor.SHA256+".bin")
info, err := os.Stat(path)
if err != nil {
if errors.Is(err, os.ErrNotExist) {
return nil, &clients.ClientError{Code: "binary_result_expired", Message: "local binary result has expired", StatusCode: 410, Retryable: false}
}
return nil, localBinaryStorageError(err)
}
if info.IsDir() || info.Size() != descriptor.Size {
return nil, &clients.ClientError{Code: "binary_result_corrupted", Message: "local binary result has an invalid size", StatusCode: 500, Retryable: false}
}
if info.ModTime().Before(time.Now().Add(-time.Duration(s.localResultTTLHours()) * time.Hour)) {
return nil, &clients.ClientError{Code: "binary_result_expired", Message: "local binary result has expired", StatusCode: 410, Retryable: false}
}
payload, err := os.ReadFile(path)
if err != nil {
return nil, localBinaryStorageError(err)
}
digest := sha256.Sum256(payload)
if int64(len(payload)) != descriptor.Size || hex.EncodeToString(digest[:]) != descriptor.SHA256 {
return nil, &clients.ClientError{Code: "binary_result_corrupted", Message: "local binary result failed size or hash verification", StatusCode: 500, Retryable: false}
}
return payload, nil
}
func localBinaryStringBytes(key string, value string, siblings map[string]any) ([]byte, string, string, bool) {
raw := strings.TrimSpace(value)
if raw == "" || strings.HasPrefix(raw, localBinaryPlaceholderPrefix) {
return nil, "", "", false
}
if strings.HasPrefix(strings.ToLower(raw), "data:") {
contentType, encoded, ok, err := parseBase64DataURL(raw)
if err == nil && ok {
payload, decodeErr := decodeBase64Payload(encoded)
if decodeErr == nil && len(payload) > 0 {
return payload, contentType, "data-uri", true
}
}
return nil, "", "", false
}
strict := localBinaryKey(key)
if !strict && len(raw) < localBinaryGenericBase64MinLength {
return nil, "", "", false
}
payload, err := decodeBase64Payload(raw)
if err != nil || len(payload) == 0 {
return nil, "", "", false
}
return payload, firstNonEmptyString(mediaContentTypeFromItem(siblings), defaultContentTypeForRawMediaKey(key)), "raw", true
}
func localBufferObjectBytes(value map[string]any) ([]byte, string, bool) {
if normalizeLocalBinaryKey(stringFromAny(value["type"])) != "buffer" {
return nil, "", false
}
contentType := firstNonEmptyString(
stringFromAny(value["mime_type"]),
stringFromAny(value["mimeType"]),
stringFromAny(value["contentType"]),
)
switch data := value["data"].(type) {
case []byte:
if len(data) == 0 {
return nil, "", false
}
return append([]byte(nil), data...), contentType, true
case []any:
payload, ok := bytesFromNumberArray(data)
return payload, contentType, ok
default:
return nil, "", false
}
}
func localBinaryKey(key string) bool {
normalized := normalizeLocalBinaryKey(key)
return normalized == "b64" ||
normalized == "b64json" ||
normalized == "base64" ||
normalized == "buffer" ||
normalized == "bytes" ||
strings.Contains(normalized, "base64") ||
strings.Contains(normalized, "buffer") ||
strings.Contains(normalized, "binary") ||
strings.HasSuffix(normalized, "b64") ||
strings.HasSuffix(normalized, "bytes")
}
func normalizeLocalBinaryKey(value string) string {
return strings.Map(func(char rune) rune {
switch {
case char >= 'a' && char <= 'z':
return char
case char >= 'A' && char <= 'Z':
return char + ('a' - 'A')
case char >= '0' && char <= '9':
return char
default:
return -1
}
}, value)
}
func localBinaryPrefix(payload []byte) string {
prefix := base64.StdEncoding.EncodeToString(payload)
if len(prefix) > 16 {
prefix = prefix[:16]
}
return prefix
}
func normalizedLocalBinaryContentType(value string) string {
value = normalizeGeneratedContentType(value)
if value == "" || len(value) > 32 || !localBinaryContentTypeSafe(value) {
return "application/octet-stream"
}
return value
}
func localBinaryContentTypeSafe(value string) bool {
for _, char := range value {
switch {
case char >= 'a' && char <= 'z':
case char >= 'A' && char <= 'Z':
case char >= '0' && char <= '9':
case char == '/', char == '.', char == '+', char == '-':
default:
return false
}
}
return true
}
func normalizedLocalBinaryEncoding(value string) string {
switch strings.ToLower(strings.TrimSpace(value)) {
case "data-uri":
return "data-uri"
case "buffer":
return "buffer"
default:
return "raw"
}
}
func localBinaryPlaceholder(descriptor localBinaryDescriptor) string {
return fmt.Sprintf(
"[GatewayBinary:v1;prefix=%s;sha256=%s;bytes=%d;mime=%s;encoding=%s]",
descriptor.Prefix,
descriptor.SHA256,
descriptor.Size,
descriptor.ContentType,
descriptor.Encoding,
)
}
func parseLocalBinaryPlaceholder(value string) (localBinaryDescriptor, bool) {
if !strings.HasPrefix(value, localBinaryPlaceholderPrefix) || !strings.HasSuffix(value, "]") {
return localBinaryDescriptor{}, false
}
content := strings.TrimSuffix(strings.TrimPrefix(value, localBinaryPlaceholderPrefix), "]")
fields := map[string]string{}
for _, item := range strings.Split(content, ";") {
key, fieldValue, ok := strings.Cut(item, "=")
if !ok {
continue
}
fields[key] = fieldValue
}
size, err := strconv.ParseInt(fields["bytes"], 10, 64)
if err != nil || size <= 0 || len(fields["sha256"]) != sha256.Size*2 {
return localBinaryDescriptor{}, false
}
if _, err := hex.DecodeString(fields["sha256"]); err != nil {
return localBinaryDescriptor{}, false
}
descriptor := localBinaryDescriptor{
Prefix: fields["prefix"],
SHA256: strings.ToLower(fields["sha256"]),
Size: size,
ContentType: normalizedLocalBinaryContentType(fields["mime"]),
Encoding: normalizedLocalBinaryEncoding(fields["encoding"]),
}
if len(descriptor.Prefix) > 16 {
return localBinaryDescriptor{}, false
}
return descriptor, true
}
func localBinaryResultHasPlaceholders(value any) bool {
switch typed := value.(type) {
case map[string]any:
for _, child := range typed {
if localBinaryResultHasPlaceholders(child) {
return true
}
}
case []any:
for _, child := range typed {
if localBinaryResultHasPlaceholders(child) {
return true
}
}
case string:
return strings.HasPrefix(typed, localBinaryPlaceholderPrefix)
}
return false
}
func safeLocalBinaryTaskDir(taskID string) string {
taskID = strings.TrimSpace(taskID)
if taskID != "" {
safe := true
for _, char := range taskID {
if (char >= 'a' && char <= 'z') ||
(char >= 'A' && char <= 'Z') ||
(char >= '0' && char <= '9') ||
char == '-' || char == '_' {
continue
}
safe = false
break
}
if safe && taskID != "." && taskID != ".." && len(taskID) <= 128 {
return taskID
}
}
digest := sha256.Sum256([]byte(taskID))
return "task-" + hex.EncodeToString(digest[:16])
}
func (s *Service) localBinaryResultRoot() string {
root := strings.TrimSpace(s.cfg.LocalGeneratedStorageDir)
if root == "" {
root = config.DefaultLocalGeneratedStorageDir
}
return filepath.Join(root, localBinaryResultDirName)
}
func (s *Service) localResultTTLHours() int {
if s.cfg.LocalResultTTLHours <= 0 {
return defaultLocalResultTTLHours
}
return s.cfg.LocalResultTTLHours
}
func (s *Service) localResultMinFreeBytes() int64 {
if s.cfg.LocalResultMinFreeBytes <= 0 {
return defaultLocalResultMinFreeBytes
}
return s.cfg.LocalResultMinFreeBytes
}
func (s *Service) localResultMaxBytes() int64 {
if s.cfg.LocalResultMaxBytes <= 0 {
return defaultLocalResultMaxBytes
}
return s.cfg.LocalResultMaxBytes
}
func (s *Service) localResultMaxTaskBytes() int64 {
if s.cfg.LocalResultMaxTaskBytes <= 0 {
return defaultLocalResultMaxTaskBytes
}
return s.cfg.LocalResultMaxTaskBytes
}
@@ -0,0 +1,233 @@
package runner
import (
"context"
"encoding/base64"
"encoding/json"
"errors"
"os"
"path/filepath"
"strings"
"testing"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/clients"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
func newLocalBinaryTestService(t *testing.T) *Service {
t.Helper()
return &Service{cfg: config.Config{
LocalGeneratedStorageDir: t.TempDir(),
LocalResultTTLHours: 24,
LocalResultMinFreeBytes: 1,
LocalResultMaxBytes: 1024 * 1024,
LocalResultMaxTaskBytes: 2 * 1024 * 1024,
}}
}
func TestMaterializeAndHydrateLocalBinaryResult(t *testing.T) {
service := newLocalBinaryTestService(t)
payload := []byte("one binary result shared across representations")
encoded := base64.StdEncoding.EncodeToString(payload)
input := map[string]any{
"data": []any{
map[string]any{
"b64_json": encoded,
"data_uri": "data:image/png;base64," + encoded,
"buffer": map[string]any{
"type": "Buffer",
"data": bytesToAny(payload),
"mimeType": "image/png",
},
},
},
}
persistent, changed, err := service.materializeLocalBinaryResult(context.Background(), "task-123", input)
if err != nil {
t.Fatalf("materialize local binary result: %v", err)
}
if !changed {
t.Fatal("expected binary result to be materialized")
}
item := persistent["data"].([]any)[0].(map[string]any)
for _, key := range []string{"b64_json", "data_uri", "buffer"} {
placeholder, ok := item[key].(string)
if !ok || !strings.HasPrefix(placeholder, localBinaryPlaceholderPrefix) {
t.Fatalf("%s was not replaced with a placeholder: %+v", key, item[key])
}
if len(placeholder) > 200 {
t.Fatalf("%s placeholder exceeds 200 bytes: %d", key, len(placeholder))
}
}
if input["data"].([]any)[0].(map[string]any)["b64_json"] != encoded {
t.Fatal("materializer mutated the provider result")
}
persistentJSON, err := json.Marshal(persistent)
if err != nil {
t.Fatalf("marshal persistent result: %v", err)
}
if len(persistentJSON) > 32*1024 {
t.Fatalf("binary-only persistent result exceeds 32KiB: %d", len(persistentJSON))
}
taskDir := filepath.Join(service.localBinaryResultRoot(), "task-123")
entries, err := os.ReadDir(taskDir)
if err != nil {
t.Fatalf("read local result dir: %v", err)
}
if len(entries) != 1 {
t.Fatalf("same payload should reuse one local file, got %d", len(entries))
}
taskInfo, err := os.Stat(taskDir)
if err != nil {
t.Fatalf("stat task directory: %v", err)
}
if taskInfo.Mode().Perm() != 0o750 {
t.Fatalf("task directory mode = %v, want 0750", taskInfo.Mode().Perm())
}
fileInfo, err := entries[0].Info()
if err != nil {
t.Fatalf("stat result file: %v", err)
}
if fileInfo.Mode().Perm() != 0o640 {
t.Fatalf("result file mode = %v, want 0640", fileInfo.Mode().Perm())
}
wire, err := service.HydrateTaskResult(context.Background(), "task-123", persistent)
if err != nil {
t.Fatalf("hydrate local binary result: %v", err)
}
wireItem := wire["data"].([]any)[0].(map[string]any)
if wireItem["b64_json"] != encoded {
t.Fatalf("raw Base64 mismatch: got %v", wireItem["b64_json"])
}
if wireItem["data_uri"] != "data:image/png;base64,"+encoded {
t.Fatalf("data URI mismatch: got %v", wireItem["data_uri"])
}
if wireItem["buffer"] != encoded {
t.Fatalf("Buffer should hydrate as Base64: got %v", wireItem["buffer"])
}
}
func TestHydrateLocalBinaryResultReturnsExpiredAndCorruptedErrors(t *testing.T) {
service := newLocalBinaryTestService(t)
service.cfg.LocalResultTTLHours = 1
encoded := base64.StdEncoding.EncodeToString([]byte("expiring result"))
persistent, _, err := service.materializeLocalBinaryResult(context.Background(), "task-expired", map[string]any{"b64_json": encoded})
if err != nil {
t.Fatalf("materialize fixture: %v", err)
}
entries, err := os.ReadDir(filepath.Join(service.localBinaryResultRoot(), "task-expired"))
if err != nil || len(entries) != 1 {
t.Fatalf("read fixture result: entries=%v err=%v", entries, err)
}
path := filepath.Join(service.localBinaryResultRoot(), "task-expired", entries[0].Name())
old := time.Now().Add(-2 * time.Hour)
if err := os.Chtimes(path, old, old); err != nil {
t.Fatalf("age fixture: %v", err)
}
_, err = service.HydrateTaskResult(context.Background(), "task-expired", persistent)
assertClientErrorCode(t, err, "binary_result_expired")
now := time.Now()
if err := os.Chtimes(path, now, now); err != nil {
t.Fatalf("refresh fixture: %v", err)
}
if err := os.WriteFile(path, []byte("tampered"), 0o640); err != nil {
t.Fatalf("tamper fixture: %v", err)
}
_, err = service.HydrateTaskResult(context.Background(), "task-expired", persistent)
assertClientErrorCode(t, err, "binary_result_corrupted")
}
func TestMaterializeLocalBinaryResultEnforcesLimitsAndKeepsText(t *testing.T) {
service := newLocalBinaryTestService(t)
service.cfg.LocalResultMaxBytes = 4
_, _, err := service.materializeLocalBinaryResult(context.Background(), "task-large", map[string]any{
"b64_json": base64.StdEncoding.EncodeToString([]byte("too large")),
})
assertClientErrorCode(t, err, "binary_result_too_large")
persistent, changed, err := service.materializeLocalBinaryResult(context.Background(), "task-text", map[string]any{
"message": "YWJjZA==",
})
if err != nil || changed || persistent["message"] != "YWJjZA==" {
t.Fatalf("short Base64-like text should remain unchanged: result=%+v changed=%v err=%v", persistent, changed, err)
}
service = newLocalBinaryTestService(t)
service.cfg.LocalResultMaxBytes = 1024
service.cfg.LocalResultMaxTaskBytes = 8
_, _, err = service.materializeLocalBinaryResult(context.Background(), "task-total-large", map[string]any{
"first_base64": base64.StdEncoding.EncodeToString([]byte("12345")),
"second_base64": base64.StdEncoding.EncodeToString([]byte("67890")),
})
assertClientErrorCode(t, err, "binary_result_too_large")
if _, statErr := os.Stat(filepath.Join(service.localBinaryResultRoot(), "task-total-large")); !errors.Is(statErr, os.ErrNotExist) {
t.Fatalf("preflight task limit must not create partial files: %v", statErr)
}
}
func TestCompactExpiredTaskResultDoesNotWriteFiles(t *testing.T) {
service := newLocalBinaryTestService(t)
service.cfg.LocalResultMaxBytes = 1
encoded := base64.StdEncoding.EncodeToString([]byte("already expired result"))
persistent, changed, err := service.CompactExpiredTaskResultForStorage(context.Background(), "task-old", map[string]any{
"b64_json": encoded,
})
if err != nil {
t.Fatalf("compact expired result: %v", err)
}
if !changed || !localBinaryResultHasPlaceholders(persistent) {
t.Fatalf("expired result was not compacted: %+v", persistent)
}
if _, err := os.Stat(filepath.Join(service.localBinaryResultRoot(), "task-old")); !errors.Is(err, os.ErrNotExist) {
t.Fatalf("expired compaction must not create a task directory: %v", err)
}
_, err = service.HydrateTaskResult(context.Background(), "task-old", persistent)
assertClientErrorCode(t, err, "binary_result_expired")
}
func TestLocalBinaryStorageUnavailableDoesNotRetryProvider(t *testing.T) {
err := localBinaryStorageError(errors.New("write failed"))
assertClientErrorCode(t, err, "local_result_storage_unavailable")
if clients.IsRetryable(err) {
t.Fatal("local result storage failure must not call the provider again")
}
if retryDecisionForCandidate(store.RuntimeModelCandidate{}, err).Retry {
t.Fatal("local result storage failure must not retry the same provider client")
}
if failoverDecisionForCandidate(store.RunnerPolicy{}, store.RuntimeModelCandidate{}, err).Retry {
t.Fatal("local result storage failure must not fail over to another provider")
}
service := newLocalBinaryTestService(t)
service.cfg.LocalResultMinFreeBytes = 1 << 62
_, _, err = service.materializeLocalBinaryResult(context.Background(), "task-low-disk", map[string]any{
"b64_json": base64.StdEncoding.EncodeToString([]byte("disk preflight")),
})
assertClientErrorCode(t, err, "local_result_storage_unavailable")
if _, statErr := os.Stat(filepath.Join(service.localBinaryResultRoot(), "task-low-disk")); !errors.Is(statErr, os.ErrNotExist) {
t.Fatalf("disk preflight must not create partial task files: %v", statErr)
}
}
func bytesToAny(payload []byte) []any {
result := make([]any, len(payload))
for index, value := range payload {
result[index] = float64(value)
}
return result
}
func assertClientErrorCode(t *testing.T, err error, code string) {
t.Helper()
var clientErr *clients.ClientError
if !errors.As(err, &clientErr) || clientErr.Code != code {
t.Fatalf("error = %v, want client error code %s", err, code)
}
}
@@ -57,6 +57,9 @@ func shouldRetrySameClient(candidate store.RuntimeModelCandidate, err error) boo
func retryDecisionForCandidate(candidate store.RuntimeModelCandidate, err error) retryDecision { func retryDecisionForCandidate(candidate store.RuntimeModelCandidate, err error) retryDecision {
policy := effectiveRetryPolicy(candidate) policy := effectiveRetryPolicy(candidate)
info := failureInfoFromError(err) info := failureInfoFromError(err)
if isResultPersistenceFailure(err) {
return retryDecision{Retry: false, Reason: "result_persistence_failed", Match: policyRuleMatch{Source: "gateway_result_storage", Policy: "persistence", Rule: "providerRetry", Value: "disabled"}, Info: info}
}
if errors.Is(err, store.ErrRateLimited) { if errors.Is(err, store.ErrRateLimited) {
return retryDecision{Retry: false, Reason: "local_rate_limit_wait_queue", Match: policyRuleMatch{Source: "gateway_rate_limits", Policy: "rateLimitPolicy", Rule: "localCapacity", Value: "exceeded"}, Info: info} return retryDecision{Retry: false, Reason: "local_rate_limit_wait_queue", Match: policyRuleMatch{Source: "gateway_rate_limits", Policy: "rateLimitPolicy", Rule: "localCapacity", Value: "exceeded"}, Info: info}
} }
@@ -77,6 +80,9 @@ func retryDecisionForCandidate(candidate store.RuntimeModelCandidate, err error)
func failoverDecisionForCandidate(runnerPolicy store.RunnerPolicy, candidate store.RuntimeModelCandidate, err error) failoverDecision { func failoverDecisionForCandidate(runnerPolicy store.RunnerPolicy, candidate store.RuntimeModelCandidate, err error) failoverDecision {
info := failureInfoFromError(err) info := failureInfoFromError(err)
if isResultPersistenceFailure(err) {
return failoverDecision{Retry: false, Action: "stop", Reason: "result_persistence_failed", Match: policyRuleMatch{Source: "gateway_result_storage", Policy: "persistence", Rule: "providerFailover", Value: "disabled"}, Info: info}
}
if strings.TrimSpace(runnerPolicy.Status) != "" && runnerPolicy.Status != "active" { if strings.TrimSpace(runnerPolicy.Status) != "" && runnerPolicy.Status != "active" {
return failoverDecision{Retry: false, Action: "stop", Reason: "runner_policy_disabled", Match: policyRuleMatch{Source: "gateway_runner_policies", Policy: "runnerPolicy", Rule: "status", Value: runnerPolicy.Status}, Info: info} return failoverDecision{Retry: false, Action: "stop", Reason: "runner_policy_disabled", Match: policyRuleMatch{Source: "gateway_runner_policies", Policy: "runnerPolicy", Rule: "status", Value: runnerPolicy.Status}, Info: info}
} }
@@ -126,6 +132,9 @@ func shouldDemoteCandidatePriority(runnerPolicy store.RunnerPolicy, err error) b
func priorityDemoteDecisionForCandidate(runnerPolicy store.RunnerPolicy, err error) priorityDemoteDecision { func priorityDemoteDecisionForCandidate(runnerPolicy store.RunnerPolicy, err error) priorityDemoteDecision {
info := failureInfoFromError(err) info := failureInfoFromError(err)
if isResultPersistenceFailure(err) {
return priorityDemoteDecision{Demote: false, Reason: "result_persistence_failed", Info: info}
}
if strings.TrimSpace(runnerPolicy.Status) != "" && runnerPolicy.Status != "active" { if strings.TrimSpace(runnerPolicy.Status) != "" && runnerPolicy.Status != "active" {
return priorityDemoteDecision{Demote: false, Reason: "runner_policy_disabled", Info: info} return priorityDemoteDecision{Demote: false, Reason: "runner_policy_disabled", Info: info}
} }
@@ -145,6 +154,19 @@ func priorityDemoteDecisionForCandidate(runnerPolicy store.RunnerPolicy, err err
return priorityDemoteDecision{Demote: false, Reason: "priority_demote_no_match", Info: info} return priorityDemoteDecision{Demote: false, Reason: "priority_demote_no_match", Info: info}
} }
func isResultPersistenceFailure(err error) bool {
switch strings.ToLower(strings.TrimSpace(clients.ErrorCode(err))) {
case "local_result_storage_unavailable",
"binary_result_too_large",
"binary_result_corrupted",
"binary_result_expired",
"result_binary_not_materialized":
return true
default:
return false
}
}
func effectiveFailoverPolicy(base map[string]any, override map[string]any) map[string]any { func effectiveFailoverPolicy(base map[string]any, override map[string]any) map[string]any {
policy := base policy := base
if nested := failoverOverridePolicy(override); len(nested) > 0 { if nested := failoverOverridePolicy(override); len(nested) > 0 {
+49 -3
View File
@@ -630,7 +630,11 @@ candidatesLoop:
} }
walletReservationFinalized = true walletReservationFinalized = true
s.logger.Warn("task succeeded but billing requires manual review", "taskID", task.ID, "error_category", "billing_calculation_failed") s.logger.Warn("task succeeded but billing requires manual review", "taskID", task.ID, "error_category", "billing_calculation_failed")
return Result{Task: review, Output: response.Result, Wire: response.Wire}, nil output, hydrateErr := s.HydrateTaskResult(ctx, task.ID, response.Result)
if hydrateErr != nil {
return Result{Task: review}, hydrateErr
}
return Result{Task: review, Output: output, Wire: response.Wire}, nil
} }
finalAmountText = finalAmount.String() finalAmountText = finalAmount.String()
} else { } else {
@@ -661,6 +665,38 @@ candidatesLoop:
ResponseDurationMS: record.ResponseDurationMS, ResponseDurationMS: record.ResponseDurationMS,
}) })
if finishErr != nil { if finishErr != nil {
if errors.Is(finishErr, store.ErrTaskResultBinaryNotMaterialized) {
s.logger.Error("task result rejected by binary persistence guard",
"taskID", task.ID,
"error_category", "result_binary_not_materialized",
)
failureCtx := context.WithoutCancel(ctx)
_ = s.store.FinishTaskAttempt(failureCtx, store.FinishTaskAttemptInput{
AttemptID: response.AttemptID,
Status: "failed",
Retryable: false,
RequestID: response.RequestID,
ResponseStartedAt: response.ResponseStartedAt,
ResponseFinishedAt: response.ResponseFinishedAt,
ResponseDurationMS: response.ResponseDurationMS,
ErrorCode: clients.ErrorCode(finishErr),
ErrorMessage: finishErr.Error(),
})
failed, failErr := s.failTask(
failureCtx,
task.ID,
task.ExecutionToken,
clients.ErrorCode(finishErr),
finishErr.Error(),
isSimulation(task, candidate),
finishErr,
)
if failErr != nil {
return Result{}, failErr
}
walletReservationFinalized = true
return Result{Task: failed, Output: failed.Result}, finishErr
}
if errors.Is(finishErr, store.ErrTaskExecutionLeaseLost) { if errors.Is(finishErr, store.ErrTaskExecutionLeaseLost) {
latest, latestErr := s.store.GetTask(ctx, task.ID) latest, latestErr := s.store.GetTask(ctx, task.ID)
if latestErr == nil && latest.Status == "cancelled" { if latestErr == nil && latest.Status == "cancelled" {
@@ -688,7 +724,11 @@ candidatesLoop:
}, isSimulation(task, candidate)); err != nil { }, isSimulation(task, candidate)); err != nil {
return Result{}, err return Result{}, err
} }
return Result{Task: finished, Output: response.Result, Wire: response.Wire}, nil output, hydrateErr := s.HydrateTaskResult(ctx, task.ID, response.Result)
if hydrateErr != nil {
return Result{Task: finished}, hydrateErr
}
return Result{Task: finished, Output: output, Wire: response.Wire}, nil
} }
var submissionUnknown *upstreamSubmissionUnknownError var submissionUnknown *upstreamSubmissionUnknownError
if errors.As(err, &submissionUnknown) { if errors.As(err, &submissionUnknown) {
@@ -811,7 +851,7 @@ candidatesLoop:
break break
} }
decision := failoverDecisionForCandidate(runnerPolicy, candidate, candidateErr) decision := failoverDecisionForCandidate(runnerPolicy, candidate, candidateErr)
if !decision.Retry && hasLoadAvoidanceFallback(candidates, index, maxPlatforms) { if !decision.Retry && !isResultPersistenceFailure(candidateErr) && hasLoadAvoidanceFallback(candidates, index, maxPlatforms) {
decision = loadAvoidanceFallbackDecision(candidateErr) decision = loadAvoidanceFallbackDecision(candidateErr)
} }
s.recordAttemptTrace(ctx, task.ID, attemptNo, failoverTraceEntry(decision, candidate)) s.recordAttemptTrace(ctx, task.ID, attemptNo, failoverTraceEntry(decision, candidate))
@@ -1216,6 +1256,12 @@ func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user
return clients.Response{}, err return clients.Response{}, err
} }
response.Result = uploadedResult response.Result = uploadedResult
if localBinaryResultHasPlaceholders(response.Result) {
// A provider wire response can still contain its original Base64 body.
// Force protocol responses to be rebuilt from the verified standard
// result instead of bypassing local result materialization.
response.Wire = nil
}
if task.Kind == "responses" { if task.Kind == "responses" {
response.UpstreamProtocol = candidate.ResponseProtocol response.UpstreamProtocol = candidate.ResponseProtocol
response.ParentResponseID = responseExecution.PublicPreviousResponseID response.ParentResponseID = responseExecution.PublicPreviousResponseID
+16 -2
View File
@@ -47,6 +47,7 @@ type generatedAssetUploadPolicy struct {
UploadInlineMedia bool UploadInlineMedia bool
UploadURLMedia bool UploadURLMedia bool
PreserveInlineMedia bool PreserveInlineMedia bool
LocalizeInlineMedia bool
} }
type generatedAssetDecision struct { type generatedAssetDecision struct {
@@ -79,7 +80,8 @@ func defaultGeneratedAssetUploadPolicy() generatedAssetUploadPolicy {
func (s *Service) uploadGeneratedAssets(ctx context.Context, taskID string, taskKind string, result map[string]any) (map[string]any, error) { func (s *Service) uploadGeneratedAssets(ctx context.Context, taskID string, taskKind string, result map[string]any) (map[string]any, error) {
data, _ := result["data"].([]any) data, _ := result["data"].([]any)
rawNeedsUpload := generatedRawValueHasInlineMedia(result["raw"], "", nil) rawNeedsUpload := generatedRawValueHasInlineMedia(result["raw"], "", nil)
if len(data) == 0 && !rawNeedsUpload { hasInlineBinary := TaskResultHasInlineBinary(result)
if len(data) == 0 && !rawNeedsUpload && !hasInlineBinary {
redactGeneratedResultRawData(result) redactGeneratedResultRawData(result)
return result, nil return result, nil
} }
@@ -87,6 +89,18 @@ func (s *Service) uploadGeneratedAssets(ctx context.Context, taskID string, task
if err != nil { if err != nil {
return nil, &clients.ClientError{Code: "upload_config_failed", Message: err.Error(), Retryable: true} return nil, &clients.ClientError{Code: "upload_config_failed", Message: err.Error(), Retryable: true}
} }
if policy.LocalizeInlineMedia {
next, _, err := s.materializeLocalBinaryResult(ctx, taskID, result)
if err != nil {
return nil, err
}
redactGeneratedResultRawData(next)
return next, nil
}
if len(data) == 0 && !rawNeedsUpload {
redactGeneratedResultRawData(result)
return result, nil
}
// Topaz download URLs are short-lived. Persist them before the task can be // Topaz download URLs are short-lived. Persist them before the task can be
// marked succeeded even when the global policy normally keeps URL media. // marked succeeded even when the global policy normally keeps URL media.
if taskKind == "videos.upscales" { if taskKind == "videos.upscales" {
@@ -542,7 +556,7 @@ func generatedAssetUploadPolicyFromName(policyName string) generatedAssetUploadP
case store.FileStorageResultUploadPolicyUploadAll: case store.FileStorageResultUploadPolicyUploadAll:
return generatedAssetUploadPolicy{UploadInlineMedia: true, UploadURLMedia: true} return generatedAssetUploadPolicy{UploadInlineMedia: true, UploadURLMedia: true}
case store.FileStorageResultUploadPolicyUploadNone: case store.FileStorageResultUploadPolicyUploadNone:
return generatedAssetUploadPolicy{UploadInlineMedia: false, UploadURLMedia: false, PreserveInlineMedia: true} return generatedAssetUploadPolicy{UploadInlineMedia: false, UploadURLMedia: false, PreserveInlineMedia: true, LocalizeInlineMedia: true}
default: default:
return defaultGeneratedAssetUploadPolicy() return defaultGeneratedAssetUploadPolicy()
} }
+1 -1
View File
@@ -222,7 +222,7 @@ func TestGeneratedAssetUploadPolicyFromName(t *testing.T) {
{ {
name: "upload none", name: "upload none",
policyName: store.FileStorageResultUploadPolicyUploadNone, policyName: store.FileStorageResultUploadPolicyUploadNone,
want: generatedAssetUploadPolicy{UploadInlineMedia: false, UploadURLMedia: false, PreserveInlineMedia: true}, want: generatedAssetUploadPolicy{UploadInlineMedia: false, UploadURLMedia: false, PreserveInlineMedia: true, LocalizeInlineMedia: true},
}, },
} }
@@ -0,0 +1,87 @@
package store
import (
"context"
"encoding/json"
"time"
"github.com/jackc/pgx/v5"
)
type TaskBinaryResultBackfillItem struct {
ID string
Result map[string]any
FinishedAt time.Time
}
func (s *Store) ListTaskBinaryResultBackfillBatch(ctx context.Context, afterID string, batchSize int) ([]TaskBinaryResultBackfillItem, error) {
if batchSize < 1 || batchSize > 100 {
batchSize = 100
}
var items []TaskBinaryResultBackfillItem
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
}
rows, err := tx.Query(ctx, `
SELECT id::text, result, COALESCE(finished_at, updated_at)
FROM gateway_tasks
WHERE status = 'succeeded'
AND result <> '{}'::jsonb
AND (NULLIF($1::text, '') IS NULL OR id > NULLIF($1::text, '')::uuid)
ORDER BY id
LIMIT $2`, afterID, batchSize)
if err != nil {
return err
}
defer rows.Close()
items = make([]TaskBinaryResultBackfillItem, 0, batchSize)
for rows.Next() {
var item TaskBinaryResultBackfillItem
var resultJSON []byte
if err := rows.Scan(&item.ID, &resultJSON, &item.FinishedAt); err != nil {
return err
}
item.Result = decodeObject(resultJSON)
items = append(items, item)
}
return rows.Err()
})
return items, err
}
func (s *Store) UpdateTaskBinaryResultBackfill(ctx context.Context, taskID string, result map[string]any) (bool, error) {
report := sanitizeJSONForStorageWithReport(minimalTaskResult(result))
if report.BinaryCount > 0 {
return false, &taskPayloadBinaryError{
target: ErrTaskResultBinaryNotMaterialized,
code: "result_binary_not_materialized",
count: report.BinaryCount,
}
}
resultJSON, err := json.Marshal(report.Value)
if err != nil {
return false, err
}
updated := false
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
}
tag, err := tx.Exec(ctx, `
UPDATE gateway_tasks
SET result = $2::jsonb,
updated_at = now()
WHERE id = $1::uuid
AND status = 'succeeded'`,
taskID,
string(resultJSON),
)
if err != nil {
return err
}
updated = tag.RowsAffected() == 1
return nil
})
return updated, err
}
+10 -2
View File
@@ -41,7 +41,7 @@ func (s *Store) EnsureConversation(ctx context.Context, user *auth.User, convers
if userID == "" { if userID == "" {
userID = "anonymous" userID = "anonymous"
} }
metadataJSON, _ := json.Marshal(emptyObjectIfNil(metadata)) metadataJSON, _ := json.Marshal(sanitizeJSONForStorage(emptyObjectIfNil(metadata)))
var conversationID string var conversationID string
err := s.pool.QueryRow(ctx, ` err := s.pool.QueryRow(ctx, `
INSERT INTO gateway_conversations (user_id, gateway_user_id, conversation_key, metadata) INSERT INTO gateway_conversations (user_id, gateway_user_id, conversation_key, metadata)
@@ -67,7 +67,15 @@ func (s *Store) UpsertConversationMessages(ctx context.Context, conversationID s
refs := make([]TaskMessageRefInput, 0, len(messages)) refs := make([]TaskMessageRefInput, 0, len(messages))
newCount := 0 newCount := 0
for index, message := range messages { for index, message := range messages {
snapshotJSON, _ := json.Marshal(emptyObjectIfNil(message.Snapshot)) snapshotReport := sanitizeJSONForStorageWithReport(emptyObjectIfNil(message.Snapshot))
if snapshotReport.BinaryCount > 0 {
return nil, 0, &taskPayloadBinaryError{
target: ErrTaskRequestBinaryNotMaterialized,
code: "request_binary_not_materialized",
count: snapshotReport.BinaryCount,
}
}
snapshotJSON, _ := json.Marshal(snapshotReport.Value)
var messageID string var messageID string
var inserted bool var inserted bool
if err := tx.QueryRow(ctx, ` if err := tx.QueryRow(ctx, `
@@ -0,0 +1,344 @@
package store
import (
"crypto/sha256"
"encoding/base64"
"encoding/json"
"errors"
"fmt"
"strconv"
"strings"
)
const (
storageBinaryPrefixChars = 16
storageGenericBase64MinLength = 4096
storageInvalidBinaryMinLength = 512
storageJSONSanitizerMaxDepth = 64
storagePlaceholderPrefix = "[GatewayBinary:v1;"
storageBufferObjectType = "buffer"
storageDataURLPrefix = "data:"
storageDataURLBase64Marker = ";base64"
storageDataURLMaxContentType = 64
)
var (
ErrTaskRequestBinaryNotMaterialized = errors.New("task request binary payload was not materialized")
ErrTaskResultBinaryNotMaterialized = errors.New("task result binary payload was not materialized")
)
type taskPayloadBinaryError struct {
target error
code string
count int
}
func (e *taskPayloadBinaryError) Error() string {
return fmt.Sprintf("%s: detected %d inline binary value(s)", e.code, e.count)
}
func (e *taskPayloadBinaryError) ErrorCode() string {
return e.code
}
func (e *taskPayloadBinaryError) Is(target error) bool {
return target == e.target
}
type storageSanitizeReport struct {
Value any
BinaryCount int
}
// sanitizeJSONForStorage is the final task-domain persistence guard. It always
// returns a detached JSON-compatible value and replaces inline binary payloads
// with a bounded, deterministic placeholder.
func sanitizeJSONForStorage(value any) any {
return sanitizeJSONForStorageWithReport(value).Value
}
func sanitizeJSONForStorageWithReport(value any) storageSanitizeReport {
next, count := sanitizeJSONStorageValue(value, nil, 0)
return storageSanitizeReport{Value: next, BinaryCount: count}
}
func sanitizeJSONStorageValue(value any, path []string, depth int) (any, int) {
if depth >= storageJSONSanitizerMaxDepth {
return "[JSON,max-depth]", 0
}
switch typed := value.(type) {
case map[string]any:
if payload, contentType, ok := storageBufferObjectBytes(typed); ok {
return storageBinaryPlaceholder(payload, contentType, "buffer"), 1
}
next := make(map[string]any, len(typed))
count := 0
for key, child := range typed {
sanitized, childCount := sanitizeJSONStorageValue(child, appendStoragePath(path, key), depth+1)
next[key] = sanitized
count += childCount
}
return next, count
case []any:
if storagePathIsBinary(path) {
if payload, ok := storageNumberArrayBytes(typed); ok {
return storageBinaryPlaceholder(payload, "", "buffer"), 1
}
}
next := make([]any, len(typed))
count := 0
for index, child := range typed {
sanitized, childCount := sanitizeJSONStorageValue(child, path, depth+1)
next[index] = sanitized
count += childCount
}
return next, count
case []byte:
return storageBinaryPlaceholder(typed, "", "buffer"), 1
case string:
if placeholder, ok := storageStringPlaceholder(typed, path); ok {
return placeholder, 1
}
return typed, 0
default:
return value, 0
}
}
func appendStoragePath(path []string, key string) []string {
next := make([]string, len(path)+1)
copy(next, path)
next[len(path)] = key
return next
}
func storageStringPlaceholder(value string, path []string) (string, bool) {
raw := strings.TrimSpace(value)
if raw == "" || strings.HasPrefix(raw, storagePlaceholderPrefix) {
return "", false
}
encoded, contentType, dataURL := storageBase64StringParts(raw)
strict := dataURL || storagePathIsBinary(path)
if !strict && len(encoded) < storageGenericBase64MinLength {
return "", false
}
payload, ok := storageDecodeBase64(encoded)
if ok {
encoding := "raw"
if dataURL {
encoding = "data-uri"
}
return storageBinaryPlaceholder(payload, contentType, encoding), true
}
if strict && len(raw) >= storageInvalidBinaryMinLength {
return storageBinaryPlaceholder([]byte(raw), contentType, "raw"), true
}
return "", false
}
func storageBase64StringParts(value string) (encoded string, contentType string, dataURL bool) {
if !strings.HasPrefix(strings.ToLower(value), storageDataURLPrefix) {
return value, "", false
}
prefix, payload, ok := strings.Cut(value, ",")
if !ok || !strings.Contains(strings.ToLower(prefix), storageDataURLBase64Marker) {
return value, "", false
}
mediaType := strings.TrimSpace(prefix[len(storageDataURLPrefix):])
if before, _, found := strings.Cut(mediaType, ";"); found {
mediaType = before
}
if len(mediaType) > storageDataURLMaxContentType || !storageSafeContentType(mediaType) {
mediaType = ""
}
return payload, strings.ToLower(mediaType), true
}
func storageSafeContentType(value string) bool {
if value == "" {
return true
}
for _, char := range value {
switch {
case char >= 'a' && char <= 'z':
case char >= 'A' && char <= 'Z':
case char >= '0' && char <= '9':
case char == '/', char == '.', char == '+', char == '-':
default:
return false
}
}
return true
}
func storageDecodeBase64(value string) ([]byte, bool) {
normalized := removeStorageASCIIWhitespace(value)
if normalized == "" {
return nil, false
}
for _, encoding := range []*base64.Encoding{
base64.StdEncoding,
base64.RawStdEncoding,
base64.URLEncoding,
base64.RawURLEncoding,
} {
payload, err := encoding.DecodeString(normalized)
if err != nil || len(payload) == 0 {
continue
}
canonical := encoding.EncodeToString(payload)
if strings.TrimRight(normalized, "=") == strings.TrimRight(canonical, "=") {
return payload, true
}
}
return nil, false
}
func removeStorageASCIIWhitespace(value string) string {
return strings.Map(func(char rune) rune {
switch char {
case ' ', '\n', '\r', '\t':
return -1
default:
return char
}
}, value)
}
func storagePathIsBinary(path []string) bool {
if len(path) == 0 {
return false
}
key := normalizeStorageBinaryKey(path[len(path)-1])
if storageBinaryKey(key) {
return true
}
if len(path) < 2 {
return false
}
parent := normalizeStorageBinaryKey(path[len(path)-2])
return (parent == "inlinedata" || parent == "binary" || parent == "media") &&
(key == "data" || key == "content")
}
func normalizeStorageBinaryKey(value string) string {
return strings.Map(func(char rune) rune {
switch {
case char >= 'a' && char <= 'z':
return char
case char >= 'A' && char <= 'Z':
return char + ('a' - 'A')
case char >= '0' && char <= '9':
return char
default:
return -1
}
}, value)
}
func storageBinaryKey(key string) bool {
return key == "b64" ||
key == "b64json" ||
key == "base64" ||
key == "buffer" ||
key == "bytes" ||
strings.Contains(key, "base64") ||
strings.Contains(key, "buffer") ||
strings.Contains(key, "binary") ||
strings.HasSuffix(key, "b64") ||
strings.HasSuffix(key, "bytes")
}
func storageBufferObjectBytes(value map[string]any) ([]byte, string, bool) {
kind, _ := value["type"].(string)
if normalizeStorageBinaryKey(kind) != storageBufferObjectType {
return nil, "", false
}
contentType := firstNonEmpty(
stringFromAny(value["mime_type"]),
stringFromAny(value["mimeType"]),
stringFromAny(value["contentType"]),
)
switch data := value["data"].(type) {
case []byte:
if len(data) == 0 {
return nil, "", false
}
return append([]byte(nil), data...), contentType, true
case []any:
payload, ok := storageNumberArrayBytes(data)
return payload, contentType, ok
default:
return nil, "", false
}
}
func storageNumberArrayBytes(values []any) ([]byte, bool) {
if len(values) == 0 {
return nil, false
}
payload := make([]byte, len(values))
for index, value := range values {
next, ok := storageByteFromAny(value)
if !ok {
return nil, false
}
payload[index] = next
}
return payload, true
}
func storageByteFromAny(value any) (byte, bool) {
switch typed := value.(type) {
case byte:
return typed, true
case int:
if typed >= 0 && typed <= 255 {
return byte(typed), true
}
case int32:
if typed >= 0 && typed <= 255 {
return byte(typed), true
}
case int64:
if typed >= 0 && typed <= 255 {
return byte(typed), true
}
case float64:
asInt := int(typed)
if typed == float64(asInt) && asInt >= 0 && asInt <= 255 {
return byte(asInt), true
}
case json.Number:
asInt, err := strconv.ParseInt(string(typed), 10, 16)
if err == nil && asInt >= 0 && asInt <= 255 {
return byte(asInt), true
}
}
return 0, false
}
func storageBinaryPlaceholder(payload []byte, contentType string, encoding string) string {
digest := sha256.Sum256(payload)
prefix := base64.StdEncoding.EncodeToString(payload)
if len(prefix) > storageBinaryPrefixChars {
prefix = prefix[:storageBinaryPrefixChars]
}
contentType = strings.ToLower(strings.TrimSpace(contentType))
if contentType == "" || len(contentType) > 32 || !storageSafeContentType(contentType) {
contentType = "application/octet-stream"
}
switch encoding {
case "data-uri", "buffer":
default:
encoding = "raw"
}
return fmt.Sprintf(
"[GatewayBinary:v1;prefix=%s;sha256=%x;bytes=%d;mime=%s;encoding=%s]",
prefix,
digest,
len(payload),
contentType,
encoding,
)
}
@@ -0,0 +1,103 @@
package store
import (
"encoding/base64"
"encoding/json"
"strings"
"testing"
)
func TestSanitizeJSONForStorageReplacesBinaryWithoutMutatingInput(t *testing.T) {
payload := []byte("shared binary payload")
encoded := base64.StdEncoding.EncodeToString(payload)
input := map[string]any{
"data": []any{
map[string]any{"b64_json": encoded},
},
"buffer": map[string]any{
"type": "Buffer",
"data": []any{float64(1), float64(2), float64(3)},
},
"text": "ordinary text",
}
report := sanitizeJSONForStorageWithReport(input)
if report.BinaryCount != 2 {
t.Fatalf("binary count = %d, want 2", report.BinaryCount)
}
next := report.Value.(map[string]any)
data := next["data"].([]any)
placeholder := data[0].(map[string]any)["b64_json"].(string)
if !strings.HasPrefix(placeholder, storagePlaceholderPrefix) {
t.Fatalf("unexpected placeholder: %q", placeholder)
}
if len(placeholder) > 200 {
t.Fatalf("placeholder exceeds 200 bytes: %d", len(placeholder))
}
if !strings.Contains(placeholder, ";prefix="+encoded[:16]+";") {
t.Fatalf("placeholder should retain the bounded Base64 prefix: %q", placeholder)
}
if got := input["data"].([]any)[0].(map[string]any)["b64_json"]; got != encoded {
t.Fatalf("sanitizer mutated input: %v", got)
}
if next["text"] != "ordinary text" {
t.Fatalf("ordinary text changed: %+v", next)
}
}
func TestSanitizeJSONForStorageUsesDecodedBytesForEquivalentRepresentations(t *testing.T) {
payload := []byte("equivalent payload")
encoded := base64.StdEncoding.EncodeToString(payload)
report := sanitizeJSONForStorageWithReport(map[string]any{
"rawBase64": encoded,
"dataURL": "data:image/png;base64," + encoded,
"bytes": []any{float64('e'), float64('q'), float64('u'), float64('i'), float64('v'), float64('a'), float64('l'), float64('e'), float64('n'), float64('t'), float64(' '), float64('p'), float64('a'), float64('y'), float64('l'), float64('o'), float64('a'), float64('d')},
})
if report.BinaryCount != 3 {
t.Fatalf("binary count = %d, want 3", report.BinaryCount)
}
next := report.Value.(map[string]any)
hashes := map[string]struct{}{}
for _, key := range []string{"rawBase64", "dataURL", "bytes"} {
value := next[key].(string)
hashStart := strings.Index(value, ";sha256=")
hashEnd := strings.Index(value[hashStart+1:], ";bytes=")
if hashStart < 0 || hashEnd < 0 {
t.Fatalf("missing hash in %s placeholder: %q", key, value)
}
hashes[value[hashStart+8:hashStart+1+hashEnd]] = struct{}{}
}
if len(hashes) != 1 {
t.Fatalf("equivalent binary values produced different hashes: %+v", next)
}
}
func TestSanitizeJSONForStorageAvoidsShortGenericTextFalsePositive(t *testing.T) {
input := map[string]any{"message": "YWJjZA==", "count": json.Number("12")}
report := sanitizeJSONForStorageWithReport(input)
if report.BinaryCount != 0 {
t.Fatalf("short generic Base64-like text should not be sanitized: %+v", report)
}
}
func TestSanitizeJSONForStorageRecognizesExistingPlaceholder(t *testing.T) {
value := storageBinaryPlaceholder([]byte("payload"), "image/png", "raw")
report := sanitizeJSONForStorageWithReport(map[string]any{"b64_json": value})
if report.BinaryCount != 0 {
t.Fatalf("existing placeholder should pass the final guard: %+v", report)
}
}
func TestStorageBinaryPlaceholderIsAlwaysBounded(t *testing.T) {
value := storageBinaryPlaceholder(
[]byte("payload"),
"application/vnd.a-very-long-provider-specific-generated-binary-result+json",
"data-uri",
)
if len(value) > 200 {
t.Fatalf("placeholder exceeds 200 bytes: %d %q", len(value), value)
}
if !strings.Contains(value, ";mime=application/octet-stream;") {
t.Fatalf("oversized MIME should use the bounded fallback: %q", value)
}
}
+10 -2
View File
@@ -1993,7 +1993,15 @@ func (s *Store) CreateTask(ctx context.Context, input CreateTaskInput, user *aut
} }
func (s *Store) CreateTaskIdempotent(ctx context.Context, input CreateTaskInput, user *auth.User) (CreateTaskResult, error) { func (s *Store) CreateTaskIdempotent(ctx context.Context, input CreateTaskInput, user *auth.User) (CreateTaskResult, error) {
requestBody, _ := json.Marshal(input.Request) requestReport := sanitizeJSONForStorageWithReport(input.Request)
if requestReport.BinaryCount > 0 {
return CreateTaskResult{}, &taskPayloadBinaryError{
target: ErrTaskRequestBinaryNotMaterialized,
code: "request_binary_not_materialized",
count: requestReport.BinaryCount,
}
}
requestBody, _ := json.Marshal(requestReport.Value)
runMode := normalizeRunMode(input.RunMode, input.Request) runMode := normalizeRunMode(input.RunMode, input.Request)
status := "queued" status := "queued"
resultBody, _ := json.Marshal(map[string]any(nil)) resultBody, _ := json.Marshal(map[string]any(nil))
@@ -2048,7 +2056,7 @@ WHERE user_id = $1 AND idempotency_key_hash = $2`, user.ID, strings.TrimSpace(in
} }
events := taskEventsForCreate(task.ID, runMode, status, nil) events := taskEventsForCreate(task.ID, runMode, status, nil)
for _, event := range events { for _, event := range events {
payload, _ := json.Marshal(event.Payload) payload, _ := json.Marshal(sanitizeJSONForStorage(event.Payload))
if _, err := tx.Exec(ctx, ` if _, err := tx.Exec(ctx, `
INSERT INTO gateway_task_events (task_id, seq, event_type, status, phase, progress, message, payload, simulated) INSERT INTO gateway_task_events (task_id, seq, event_type, status, phase, progress, message, payload, simulated)
VALUES ($1::uuid, $2, $3::text, NULLIF($4::text, ''), NULLIF($5::text, ''), $6, NULLIF($7::text, ''), $8::jsonb, $9)`, VALUES ($1::uuid, $2, $3::text, NULLIF($4::text, ''), NULLIF($5::text, ''), $6, NULLIF($7::text, ''), $8::jsonb, $9)`,
+3 -3
View File
@@ -45,15 +45,15 @@ type CreateResponseChainInput struct {
} }
func (s *Store) CreateResponseChain(ctx context.Context, input CreateResponseChainInput) error { func (s *Store) CreateResponseChain(ctx context.Context, input CreateResponseChainInput) error {
requestJSON, err := json.Marshal(input.RequestSnapshot) requestJSON, err := json.Marshal(sanitizeJSONForStorage(input.RequestSnapshot))
if err != nil { if err != nil {
return err return err
} }
responseJSON, err := json.Marshal(input.ResponseSnapshot) responseJSON, err := json.Marshal(sanitizeJSONForStorage(input.ResponseSnapshot))
if err != nil { if err != nil {
return err return err
} }
internalJSON, err := json.Marshal(input.InternalSnapshot) internalJSON, err := json.Marshal(sanitizeJSONForStorage(input.InternalSnapshot))
if err != nil { if err != nil {
return err return err
} }
+29 -13
View File
@@ -503,7 +503,7 @@ WHERE id = $1::uuid`, taskID, riverJobID)
} }
func (s *Store) SetTaskRemoteTask(ctx context.Context, taskID string, executionToken string, attemptID string, remoteTaskID string, payload map[string]any) error { func (s *Store) SetTaskRemoteTask(ctx context.Context, taskID string, executionToken string, attemptID string, remoteTaskID string, payload map[string]any) error {
payloadJSON, _ := json.Marshal(emptyObjectIfNil(payload)) payloadJSON, _ := json.Marshal(sanitizeJSONForStorage(emptyObjectIfNil(payload)))
return pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error { return pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
tag, err := tx.Exec(ctx, ` tag, err := tx.Exec(ctx, `
UPDATE gateway_tasks UPDATE gateway_tasks
@@ -732,7 +732,7 @@ func (s *Store) CreateTaskParamPreprocessingLog(ctx context.Context, input Creat
if !input.Changed { if !input.Changed {
return "", nil return "", nil
} }
changesJSON, _ := json.Marshal(input.Changes) changesJSON, _ := json.Marshal(sanitizeJSONForStorage(input.Changes))
if input.Changes == nil { if input.Changes == nil {
changesJSON = []byte("[]") changesJSON = []byte("[]")
} }
@@ -1029,12 +1029,20 @@ WHERE id = $1::uuid`,
} }
func (s *Store) FinishTaskSuccess(ctx context.Context, input FinishTaskSuccessInput) (GatewayTask, error) { func (s *Store) FinishTaskSuccess(ctx context.Context, input FinishTaskSuccessInput) (GatewayTask, error) {
resultJSON, _ := json.Marshal(minimalTaskResult(input.Result)) resultReport := sanitizeJSONForStorageWithReport(minimalTaskResult(input.Result))
billingsJSON, _ := json.Marshal(input.Billings) if resultReport.BinaryCount > 0 {
usageJSON, _ := json.Marshal(emptyObjectIfNil(input.Usage)) return GatewayTask{}, &taskPayloadBinaryError{
metricsJSON, _ := json.Marshal(emptyObjectIfNil(input.Metrics)) target: ErrTaskResultBinaryNotMaterialized,
billingSummaryJSON, _ := json.Marshal(emptyObjectIfNil(input.BillingSummary)) code: "result_binary_not_materialized",
pricingSnapshotJSON, _ := json.Marshal(emptyObjectIfNil(input.PricingSnapshot)) count: resultReport.BinaryCount,
}
}
resultJSON, _ := json.Marshal(resultReport.Value)
billingsJSON, _ := json.Marshal(sanitizeJSONForStorage(input.Billings))
usageJSON, _ := json.Marshal(sanitizeJSONForStorage(emptyObjectIfNil(input.Usage)))
metricsJSON, _ := json.Marshal(sanitizeJSONForStorage(emptyObjectIfNil(input.Metrics)))
billingSummaryJSON, _ := json.Marshal(sanitizeJSONForStorage(emptyObjectIfNil(input.BillingSummary)))
pricingSnapshotJSON, _ := json.Marshal(sanitizeJSONForStorage(emptyObjectIfNil(input.PricingSnapshot)))
finalChargeAmount := strings.TrimSpace(input.FinalChargeAmountText) finalChargeAmount := strings.TrimSpace(input.FinalChargeAmountText)
if finalChargeAmount == "" { if finalChargeAmount == "" {
finalChargeAmount = strconv.FormatFloat(input.FinalChargeAmount, 'f', 9, 64) finalChargeAmount = strconv.FormatFloat(input.FinalChargeAmount, 'f', 9, 64)
@@ -1158,8 +1166,16 @@ func (s *Store) FinishTaskManualReview(ctx context.Context, input FinishTaskManu
if status == "failed" { if status == "failed" {
result = nil result = nil
} }
resultJSON, _ := json.Marshal(minimalTaskResult(result)) resultReport := sanitizeJSONForStorageWithReport(minimalTaskResult(result))
pricingSnapshotJSON, _ := json.Marshal(emptyObjectIfNil(input.PricingSnapshot)) if resultReport.BinaryCount > 0 {
return GatewayTask{}, &taskPayloadBinaryError{
target: ErrTaskResultBinaryNotMaterialized,
code: "result_binary_not_materialized",
count: resultReport.BinaryCount,
}
}
resultJSON, _ := json.Marshal(resultReport.Value)
pricingSnapshotJSON, _ := json.Marshal(sanitizeJSONForStorage(emptyObjectIfNil(input.PricingSnapshot)))
err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error { err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
if strings.TrimSpace(input.AttemptID) != "" { if strings.TrimSpace(input.AttemptID) != "" {
attemptStatus := "failed" attemptStatus := "failed"
@@ -1261,7 +1277,7 @@ func (s *Store) SettleTaskBilling(ctx context.Context, task GatewayTask) error {
"billings": task.Billings, "billings": task.Billings,
"billingSummary": task.BillingSummary, "billingSummary": task.BillingSummary,
} }
metadata, _ := json.Marshal(metadataMap) metadata, _ := json.Marshal(sanitizeJSONForStorage(metadataMap))
return pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error { return pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
if _, err := tx.Exec(ctx, ` if _, err := tx.Exec(ctx, `
INSERT INTO gateway_wallet_accounts ( INSERT INTO gateway_wallet_accounts (
@@ -1353,7 +1369,7 @@ ON CONFLICT (account_id, idempotency_key) WHERE idempotency_key IS NOT NULL DO N
"frozenBefore": roundMoney(frozenBefore), "frozenBefore": roundMoney(frozenBefore),
"frozenAfter": frozenAfter, "frozenAfter": frozenAfter,
}) })
metadata, _ = json.Marshal(billingMetadata) metadata, _ = json.Marshal(sanitizeJSONForStorage(billingMetadata))
if _, err := tx.Exec(ctx, ` if _, err := tx.Exec(ctx, `
INSERT INTO gateway_wallet_transactions ( INSERT INTO gateway_wallet_transactions (
account_id, gateway_tenant_id, gateway_user_id, direction, transaction_type, account_id, gateway_tenant_id, gateway_user_id, direction, transaction_type,
@@ -1389,7 +1405,7 @@ func taskBillingString(value any) string {
} }
func (s *Store) FinishTaskFailure(ctx context.Context, input FinishTaskFailureInput) (GatewayTask, error) { func (s *Store) FinishTaskFailure(ctx context.Context, input FinishTaskFailureInput) (GatewayTask, error) {
metricsJSON, _ := json.Marshal(emptyObjectIfNil(input.Metrics)) metricsJSON, _ := json.Marshal(sanitizeJSONForStorage(emptyObjectIfNil(input.Metrics)))
resultJSON, _ := json.Marshal(minimalTaskResult(nil)) resultJSON, _ := json.Marshal(minimalTaskResult(nil))
message := truncateUTF8Bytes(input.Message, 2048) message := truncateUTF8Bytes(input.Message, 2048)
err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error { err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
+2 -2
View File
@@ -161,7 +161,7 @@ func (s *Store) ReserveTaskBilling(ctx context.Context, task GatewayTask, user *
} }
reservations := make([]WalletBillingReservation, 0, len(amounts)) reservations := make([]WalletBillingReservation, 0, len(amounts))
pricingSnapshotJSON, _ := json.Marshal(pricingSnapshot) pricingSnapshotJSON, _ := json.Marshal(sanitizeJSONForStorage(pricingSnapshot))
requestFingerprint := walletString(pricingSnapshot["requestFingerprint"]) requestFingerprint := walletString(pricingSnapshot["requestFingerprint"])
err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error { err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
for currency, rawAmount := range amounts { for currency, rawAmount := range amounts {
@@ -306,7 +306,7 @@ func (s *Store) reserveTaskBillingExact(ctx context.Context, task GatewayTask, g
if currency != "resource" { if currency != "resource" {
return nil, fmt.Errorf("unsupported billing currency %q", currency) return nil, fmt.Errorf("unsupported billing currency %q", currency)
} }
pricingSnapshotJSON, _ := json.Marshal(pricingSnapshot) pricingSnapshotJSON, _ := json.Marshal(sanitizeJSONForStorage(pricingSnapshot))
requestFingerprint := walletString(pricingSnapshot["requestFingerprint"]) requestFingerprint := walletString(pricingSnapshot["requestFingerprint"])
var reservations []WalletBillingReservation var reservations []WalletBillingReservation
err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error { err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
@@ -51,7 +51,7 @@ const sceneOptions = [
const resultUploadPolicyOptions = [ const resultUploadPolicyOptions = [
{ value: 'default', label: '默认:仅非链接资源转存', description: 'URL 结果直接保存;base64 / buffer 等结果转存后保存 URL' }, { value: 'default', label: '默认:仅非链接资源转存', description: 'URL 结果直接保存;base64 / buffer 等结果转存后保存 URL' },
{ value: 'upload_all', label: '全部转存', description: 'URL、base64、buffer 等生成媒体结果都会转存到当前文件渠道' }, { value: 'upload_all', label: '全部转存', description: 'URL、base64、buffer 等生成媒体结果都会转存到当前文件渠道' },
{ value: 'upload_none', label: '全部不转存', description: '链接结果直接保存;base64 / buffer 结果保留在对应响应字段中' }, { value: 'upload_none', label: '不做外部转存', description: 'base64 / buffer 临时写入本地静态文件;数据库仅保存占位符,默认 24 小时内按需恢复' },
]; ];
export function SystemSettingsPanel(props: { export function SystemSettingsPanel(props: {
+4
View File
@@ -29,6 +29,10 @@ x-api-environment: &api-environment
AI_GATEWAY_WEB_BASE_URL: ${AI_GATEWAY_COMPOSE_WEB_BASE_URL:-http://localhost:5178} AI_GATEWAY_WEB_BASE_URL: ${AI_GATEWAY_COMPOSE_WEB_BASE_URL:-http://localhost:5178}
AI_GATEWAY_GENERATED_STORAGE_DIR: /app/data/static/generated AI_GATEWAY_GENERATED_STORAGE_DIR: /app/data/static/generated
AI_GATEWAY_UPLOADED_STORAGE_DIR: /app/data/static/uploaded AI_GATEWAY_UPLOADED_STORAGE_DIR: /app/data/static/uploaded
AI_GATEWAY_LOCAL_RESULT_TTL_HOURS: ${AI_GATEWAY_LOCAL_RESULT_TTL_HOURS:-24}
AI_GATEWAY_LOCAL_RESULT_MIN_FREE_BYTES: ${AI_GATEWAY_LOCAL_RESULT_MIN_FREE_BYTES:-10737418240}
AI_GATEWAY_LOCAL_RESULT_MAX_BYTES: ${AI_GATEWAY_LOCAL_RESULT_MAX_BYTES:-268435456}
AI_GATEWAY_LOCAL_RESULT_MAX_TASK_BYTES: ${AI_GATEWAY_LOCAL_RESULT_MAX_TASK_BYTES:-536870912}
services: services:
postgres: postgres:
+7
View File
@@ -1823,11 +1823,18 @@ AI_GATEWAY_TASK_RETENTION_DAYS=30
AI_GATEWAY_TASK_ANALYSIS_RETENTION_DAYS=7 AI_GATEWAY_TASK_ANALYSIS_RETENTION_DAYS=7
AI_GATEWAY_TASK_CLEANUP_INTERVAL_SECONDS=300 AI_GATEWAY_TASK_CLEANUP_INTERVAL_SECONDS=300
AI_GATEWAY_TASK_CLEANUP_BATCH_SIZE=1000 AI_GATEWAY_TASK_CLEANUP_BATCH_SIZE=1000
AI_GATEWAY_LOCAL_RESULT_TTL_HOURS=24
AI_GATEWAY_LOCAL_RESULT_MIN_FREE_BYTES=10737418240
AI_GATEWAY_LOCAL_RESULT_MAX_BYTES=268435456
AI_GATEWAY_LOCAL_RESULT_MAX_TASK_BYTES=536870912
``` ```
运行时要求: 运行时要求:
- 首次上线先保持清理关闭;兼容协议和异步恢复验收通过后,显式改为 `AI_GATEWAY_TASK_CLEANUP_ENABLED=true`,再启动历史压缩和 7/30 天删除。 - 首次上线先保持清理关闭;兼容协议和异步恢复验收通过后,显式改为 `AI_GATEWAY_TASK_CLEANUP_ENABLED=true`,再启动历史压缩和 7/30 天删除。
- 结果策略为 `upload_none` 时,Base64、Data URI 和 Buffer 只写入 `data/static/generated/results/<task-id>/<sha256>.bin`;任务 JSON 保存固定占位符,查询结果时校验大小与 SHA-256 后恢复 Base64。
- 本地二进制结果默认保留 24 小时;文件过期返回 `410 binary_result_expired`,文件损坏返回 `500 binary_result_corrupted`。写入前保留至少 10GiB 空间,单文件和单任务上限分别为 256MiB、512MiB。
- 历史二进制治理使用镜像内 `/app/easyai-ai-gateway-backfill-binary-results` 先 dry-run;确认范围后带 `--apply`,每批最多 100 行,使用输出的 `resumeAfterId` 分段继续。该命令不会随服务启动自动执行。
- runtime 仅在任务真实状态变化时写 `gateway_task_events`;事件 payload 固定为空,轮询不会生成事件。 - runtime 仅在任务真实状态变化时写 `gateway_task_events`;事件 payload 固定为空,轮询不会生成事件。
- callback worker 以 `0083_task_history_minimal_storage` 的应用时间为边界,只投递新 outbox;旧 pending/processing callback 在启用治理后分批标记为 `failed`,不会向业务端补发形成历史回调风暴。 - callback worker 以 `0083_task_history_minimal_storage` 的应用时间为边界,只投递新 outbox;旧 pending/processing callback 在启用治理后分批标记为 `failed`,不会向业务端补发形成历史回调风暴。
- callback worker 使用 `SERVER_MAIN_INTERNAL_TOKEN` 调用 `TASK_PROGRESS_CALLBACK_URL` - callback worker 使用 `SERVER_MAIN_INTERNAL_TOKEN` 调用 `TASK_PROGRESS_CALLBACK_URL`