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
+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 {
policy := effectiveRetryPolicy(candidate)
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) {
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 {
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" {
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 {
info := failureInfoFromError(err)
if isResultPersistenceFailure(err) {
return priorityDemoteDecision{Demote: false, Reason: "result_persistence_failed", Info: info}
}
if strings.TrimSpace(runnerPolicy.Status) != "" && runnerPolicy.Status != "active" {
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}
}
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 {
policy := base
if nested := failoverOverridePolicy(override); len(nested) > 0 {
+49 -3
View File
@@ -630,7 +630,11 @@ candidatesLoop:
}
walletReservationFinalized = true
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()
} else {
@@ -661,6 +665,38 @@ candidatesLoop:
ResponseDurationMS: record.ResponseDurationMS,
})
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) {
latest, latestErr := s.store.GetTask(ctx, task.ID)
if latestErr == nil && latest.Status == "cancelled" {
@@ -688,7 +724,11 @@ candidatesLoop:
}, isSimulation(task, candidate)); err != nil {
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
if errors.As(err, &submissionUnknown) {
@@ -811,7 +851,7 @@ candidatesLoop:
break
}
decision := failoverDecisionForCandidate(runnerPolicy, candidate, candidateErr)
if !decision.Retry && hasLoadAvoidanceFallback(candidates, index, maxPlatforms) {
if !decision.Retry && !isResultPersistenceFailure(candidateErr) && hasLoadAvoidanceFallback(candidates, index, maxPlatforms) {
decision = loadAvoidanceFallbackDecision(candidateErr)
}
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
}
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" {
response.UpstreamProtocol = candidate.ResponseProtocol
response.ParentResponseID = responseExecution.PublicPreviousResponseID
+16 -2
View File
@@ -47,6 +47,7 @@ type generatedAssetUploadPolicy struct {
UploadInlineMedia bool
UploadURLMedia bool
PreserveInlineMedia bool
LocalizeInlineMedia bool
}
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) {
data, _ := result["data"].([]any)
rawNeedsUpload := generatedRawValueHasInlineMedia(result["raw"], "", nil)
if len(data) == 0 && !rawNeedsUpload {
hasInlineBinary := TaskResultHasInlineBinary(result)
if len(data) == 0 && !rawNeedsUpload && !hasInlineBinary {
redactGeneratedResultRawData(result)
return result, nil
}
@@ -87,6 +89,18 @@ func (s *Service) uploadGeneratedAssets(ctx context.Context, taskID string, task
if err != nil {
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
// marked succeeded even when the global policy normally keeps URL media.
if taskKind == "videos.upscales" {
@@ -542,7 +556,7 @@ func generatedAssetUploadPolicyFromName(policyName string) generatedAssetUploadP
case store.FileStorageResultUploadPolicyUploadAll:
return generatedAssetUploadPolicy{UploadInlineMedia: true, UploadURLMedia: true}
case store.FileStorageResultUploadPolicyUploadNone:
return generatedAssetUploadPolicy{UploadInlineMedia: false, UploadURLMedia: false, PreserveInlineMedia: true}
return generatedAssetUploadPolicy{UploadInlineMedia: false, UploadURLMedia: false, PreserveInlineMedia: true, LocalizeInlineMedia: true}
default:
return defaultGeneratedAssetUploadPolicy()
}
+1 -1
View File
@@ -222,7 +222,7 @@ func TestGeneratedAssetUploadPolicyFromName(t *testing.T) {
{
name: "upload none",
policyName: store.FileStorageResultUploadPolicyUploadNone,
want: generatedAssetUploadPolicy{UploadInlineMedia: false, UploadURLMedia: false, PreserveInlineMedia: true},
want: generatedAssetUploadPolicy{UploadInlineMedia: false, UploadURLMedia: false, PreserveInlineMedia: true, LocalizeInlineMedia: true},
},
}