easyai-ai-gateway/apps/api/internal/runner/binary_results_test.go

package runner

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"net/http/httptest"
	"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 writeHistoricalLocalBinaryFixture(t *testing.T, service *Service, taskID string, payload []byte) string {
	t.Helper()
	digest := sha256.Sum256(payload)
	taskDir := filepath.Join(service.localBinaryResultRoot(), safeLocalBinaryTaskDir(taskID))
	if err := os.MkdirAll(taskDir, 0o750); err != nil {
		t.Fatalf("create historical result fixture directory: %v", err)
	}
	path := filepath.Join(taskDir, hex.EncodeToString(digest[:])+".bin")
	if err := os.WriteFile(path, payload, 0o640); err != nil {
		t.Fatalf("write historical result fixture: %v", err)
	}
	return path
}

func TestHydrateHistoricalLocalBinaryResult(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.transformLocalBinaryResult(context.Background(), "task-123", input, false, true)
	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))
	}
	writeHistoricalLocalBinaryFixture(t, service, "task-123", payload)

	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 TestHydrateGeneratedResultAssetReference(t *testing.T) {
	payload := []byte("verified generated image bytes")
	digest := sha256.Sum256(payload)
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.Header().Set("Content-Type", "image/png")
		_, _ = w.Write(payload)
	}))
	defer server.Close()

	reference := func(hash string) map[string]any {
		return map[string]any{
			"assetRef": map[string]any{
				"sha256":      hash,
				"contentType": "image/png",
				"size":        len(payload),
				"url":         server.URL + "/generated.png",
			},
			"assetStorage": map[string]any{
				"scene":  store.FileStorageSceneImageResult,
				"source": "b64_json",
			},
		}
	}
	service := &Service{}
	result := map[string]any{"data": []any{map[string]any{"b64_json": reference(hex.EncodeToString(digest[:]))}}}
	hydrated, err := service.HydrateTaskResult(context.Background(), "task-remote", result)
	if err != nil {
		t.Fatalf("hydrate generated result asset: %v", err)
	}
	item := hydrated["data"].([]any)[0].(map[string]any)
	if got, want := item["b64_json"], base64.StdEncoding.EncodeToString(payload); got != want {
		t.Fatalf("hydrated Base64=%v, want %v", got, want)
	}

	result = map[string]any{"data": []any{map[string]any{"b64_json": reference(strings.Repeat("0", 64))}}}
	_, err = service.HydrateTaskResult(context.Background(), "task-corrupted", result)
	assertClientErrorCode(t, err, "binary_result_corrupted")
}

func TestHydrateGeneratedResultRefreshesPrivateObjectURL(t *testing.T) {
	cfg := config.Config{
		MediaOSSDirectEnabled:   true,
		MediaOSSEndpoint:        "https://oss.example.com",
		MediaOSSBucket:          "media-bucket",
		MediaOSSAccessKeyID:     "access-id",
		MediaOSSAccessKeySecret: "access-secret",
		MediaOSSObjectPrefix:    "media",
	}
	service := &Service{cfg: cfg, directOSS: newDirectOSSUploader(cfg)}
	result := map[string]any{
		"data": []any{map[string]any{
			"url":       "https://expired.example.com/result.png",
			"image_url": "https://expired.example.com/result.png",
			"upload": map[string]any{
				"url":         "https://expired.example.com/result.png",
				"objectKey":   "media/image_result/2026/08/04/hash.png",
				"accessScope": "private",
				"storageChannel": map[string]any{
					"channelKey": "environment-direct-oss",
					"provider":   "aliyun_oss",
				},
			},
		}},
	}
	hydrated, err := service.HydrateTaskResult(t.Context(), "task-private-url", result)
	if err != nil {
		t.Fatal(err)
	}
	item := hydrated["data"].([]any)[0].(map[string]any)
	for _, key := range []string{"url", "image_url"} {
		value := stringFromAny(item[key])
		if !strings.Contains(value, "OSSAccessKeyId=access-id") || strings.Contains(value, "expired.example.com") {
			t.Fatalf("%s was not refreshed: %q", key, value)
		}
	}
}

func TestHydrateCanonicalUploadedResultForExplicitSynchronousBase64(t *testing.T) {
	payload := []byte("canonical uploaded image")
	digest := sha256.Sum256(payload)
	getCount := 0
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Method == http.MethodGet {
			getCount++
			_, _ = w.Write(payload)
			return
		}
		w.WriteHeader(http.StatusOK)
	}))
	defer server.Close()
	cfg := config.Config{
		MediaOSSDirectEnabled: true, MediaOSSEndpoint: server.URL, MediaOSSBucket: "media-bucket",
		MediaOSSAccessKeyID: "access-id", MediaOSSAccessKeySecret: "access-secret", MediaOSSObjectPrefix: "media",
	}
	service := &Service{cfg: cfg, directOSS: newDirectOSSUploader(cfg)}
	result := map[string]any{"thinking_bytes": strings.Repeat("opaque", 1024), "thought_signature": "signature", "data": []any{map[string]any{
		"type": "image", "url": "https://expired.example/result.png", "mime_type": "image/png",
		"assetStorage": map[string]any{"scene": store.FileStorageSceneImageResult, "source": "b64_json"},
		"upload": map[string]any{
			"url": "https://expired.example/result.png", "objectKey": "media/image_result/hash.png", "accessScope": "private",
			"sha256": hex.EncodeToString(digest[:]), "size": len(payload), "contentType": "image/png",
			"storageChannel": map[string]any{"channelKey": "environment-direct-oss", "provider": "aliyun_oss"},
		},
	}}}
	projected, err := service.ProjectTaskResultURLs(t.Context(), "task-sync-b64", result)
	if err != nil {
		t.Fatal(err)
	}
	projectedItem := projected["data"].([]any)[0].(map[string]any)
	if stringFromAny(projectedItem["url"]) == "" || getCount != 0 {
		t.Fatalf("URL projection read object: item=%#v getCount=%d", projectedItem, getCount)
	}
	if projected["thinking_bytes"] != nil || projected["thought_signature"] != nil {
		t.Fatalf("provider metadata leaked: %#v", projected)
	}

	hydrated, err := service.HydrateTaskResult(t.Context(), "task-sync-b64", result)
	if err != nil {
		t.Fatal(err)
	}
	item := hydrated["data"].([]any)[0].(map[string]any)
	if got := stringFromAny(item["b64_json"]); got != base64.StdEncoding.EncodeToString(payload) {
		t.Fatalf("Base64=%q", got)
	}
	if item["url"] != nil || item["upload"] != nil || getCount != 1 {
		t.Fatalf("unexpected hydrated item=%#v getCount=%d", item, getCount)
	}
	if hydrated["thinking_bytes"] != nil || hydrated["thought_signature"] != nil {
		t.Fatalf("provider metadata leaked into synchronous response: %#v", hydrated)
	}
}

func TestSynchronousInlineResultBytesUsesMetadataBeforeObjectRead(t *testing.T) {
	result := map[string]any{"data": []any{
		map[string]any{
			"assetStorage": map[string]any{"scene": store.FileStorageSceneImageResult, "source": "b64_json"},
			"upload": map[string]any{
				"objectKey": "media/image_result/large.png", "sha256": strings.Repeat("a", 64),
				"size": MaxSynchronousInlineResponseBytes + 1, "contentType": "image/png",
				"storageChannel": map[string]any{"channelKey": "environment-direct-oss"},
			},
		},
	}}

	if got := SynchronousInlineResultBytes(result); got != MaxSynchronousInlineResponseBytes+1 {
		t.Fatalf("stored bytes=%d", got)
	}
}

func TestProjectTaskResultURLsConvertsLegacyAssetWithoutReadingObject(t *testing.T) {
	payload := []byte("must never be downloaded")
	digest := sha256.Sum256(payload)
	getCount := 0
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Method == http.MethodGet {
			getCount++
		}
		http.Error(w, "object read is forbidden", http.StatusInternalServerError)
	}))
	defer server.Close()

	service := &Service{}
	result := map[string]any{"data": []any{map[string]any{
		"b64_json": map[string]any{
			"assetRef": map[string]any{
				"sha256": hex.EncodeToString(digest[:]), "contentType": "image/png", "size": len(payload), "url": server.URL + "/result.png",
			},
			"assetStorage": map[string]any{"scene": store.FileStorageSceneImageResult, "source": "b64_json"},
		},
	}}}

	projected, err := service.ProjectTaskResultURLs(t.Context(), "task-url-only", result)
	if err != nil {
		t.Fatal(err)
	}
	item := projected["data"].([]any)[0].(map[string]any)
	if got := stringFromAny(item["url"]); got != server.URL+"/result.png" {
		t.Fatalf("projected URL=%q", got)
	}
	if item["b64_json"] != nil || item["assetRef"] != nil || item["assetStorage"] != nil || item["upload"] != nil {
		t.Fatalf("internal or inline fields leaked: %#v", item)
	}
	if getCount != 0 {
		t.Fatalf("projector downloaded object %d time(s)", getCount)
	}
}

func TestProjectTaskResultURLsRejectsHistoricalInlinePayload(t *testing.T) {
	service := &Service{}
	_, err := service.ProjectTaskResultURLs(t.Context(), "task-inline", map[string]any{
		"data": []any{map[string]any{"b64_json": base64.StdEncoding.EncodeToString([]byte("inline"))}},
	})
	assertClientErrorCode(t, err, "result_materialization_required")
}

func TestMigrateTaskResultToURLsRewritesLegacyAssetReference(t *testing.T) {
	payload := []byte("legacy")
	digest := sha256.Sum256(payload)
	service := &Service{}
	result := map[string]any{"data": []any{map[string]any{
		"b64_json": map[string]any{
			"assetRef": map[string]any{
				"sha256": hex.EncodeToString(digest[:]), "contentType": "image/png", "size": len(payload), "url": "https://cdn.example/result.png",
			},
			"assetStorage": map[string]any{"scene": store.FileStorageSceneImageResult, "source": "b64_json"},
		},
	}}}

	migrated, changed, err := service.MigrateTaskResultToURLs(t.Context(), "task-legacy", result)
	if err != nil {
		t.Fatal(err)
	}
	if !changed || TaskResultNeedsURLMigration(migrated) {
		t.Fatalf("legacy result was not fully migrated: %#v", migrated)
	}
	item := migrated["data"].([]any)[0].(map[string]any)
	if stringFromAny(item["url"]) != "https://cdn.example/result.png" || item["upload"] == nil || item["b64_json"] != nil {
		t.Fatalf("unexpected migrated result: %#v", item)
	}
}

func TestMigrateTaskResultToURLsUploadsActiveLocalPlaceholder(t *testing.T) {
	payload := []byte("historical local image")
	digest := sha256.Sum256(payload)
	putCount := 0
	storageServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Method == http.MethodPut {
			putCount++
			w.WriteHeader(http.StatusOK)
			return
		}
		http.Error(w, "unexpected method", http.StatusMethodNotAllowed)
	}))
	defer storageServer.Close()

	service := newLocalBinaryTestService(t)
	service.directOSS = &directOSSUploader{
		endpoint: storageServer.URL, bucket: "bucket", accessKeyID: "access-id", accessKeySecret: "access-secret", objectPrefix: "media",
	}
	writeHistoricalLocalBinaryFixture(t, service, "task-local-migrate", payload)
	placeholder := fmt.Sprintf("%ssha256=%s;bytes=%d;mime=image/png;encoding=base64]", localBinaryPlaceholderPrefix, hex.EncodeToString(digest[:]), len(payload))
	result := map[string]any{"data": []any{map[string]any{"b64_json": placeholder}}}

	migrated, changed, err := service.MigrateTaskResultToURLs(t.Context(), "task-local-migrate", result)
	if err != nil {
		t.Fatal(err)
	}
	if !changed || TaskResultNeedsURLMigration(migrated) || putCount != 1 {
		t.Fatalf("placeholder migration changed=%t putCount=%d result=%#v", changed, putCount, migrated)
	}
}

func TestHydrateLocalBinaryResultReturnsExpiredAndCorruptedErrors(t *testing.T) {
	service := newLocalBinaryTestService(t)
	service.cfg.LocalResultTTLHours = 1
	encoded := base64.StdEncoding.EncodeToString([]byte("expiring result"))
	persistent, _, err := service.transformLocalBinaryResult(context.Background(), "task-expired", map[string]any{"b64_json": encoded}, false, true)
	if err != nil {
		t.Fatalf("materialize fixture: %v", err)
	}
	path := writeHistoricalLocalBinaryFixture(t, service, "task-expired", []byte("expiring result"))
	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 TestHistoricalLocalBinaryFixtureEnforcesLimitsAndKeepsText(t *testing.T) {
	service := newLocalBinaryTestService(t)
	service.cfg.LocalResultMaxBytes = 4
	_, _, err := service.transformLocalBinaryResult(context.Background(), "task-large", map[string]any{
		"b64_json": base64.StdEncoding.EncodeToString([]byte("too large")),
	}, false, true)
	assertClientErrorCode(t, err, "binary_result_too_large")

	persistent, changed, err := service.transformLocalBinaryResult(context.Background(), "task-text", map[string]any{
		"message": "YWJjZA==",
	}, false, true)
	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.transformLocalBinaryResult(context.Background(), "task-total-large", map[string]any{
		"first_base64":  base64.StdEncoding.EncodeToString([]byte("12345")),
		"second_base64": base64.StdEncoding.EncodeToString([]byte("67890")),
	}, false, true)
	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 TestCompactExpiredTaskResultUsesObjectStorageAndDoesNotWriteFiles(t *testing.T) {
	var uploads int
	storageServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodPut {
			t.Fatalf("method=%s, want PUT", r.Method)
		}
		uploads++
		w.WriteHeader(http.StatusOK)
	}))
	defer storageServer.Close()
	service := newLocalBinaryTestService(t)
	service.directOSS = &directOSSUploader{
		endpoint: storageServer.URL, bucket: "bucket", accessKeyID: "access-id", accessKeySecret: "access-secret", objectPrefix: "media",
	}
	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) || TaskResultHasInlineBinary(persistent) {
		t.Fatalf("expired result was not objectified: %+v", persistent)
	}
	if uploads != 1 {
		t.Fatalf("object storage uploads=%d, want 1", uploads)
	}
	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)
	}
}

func TestHistoricalLocalBinaryStorageUnavailableDoesNotRetryProvider(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")
	}
}

func TestTaskResultIgnoresLongOpaqueBase64Metadata(t *testing.T) {
	opaque := base64.StdEncoding.EncodeToString([]byte(strings.Repeat("opaque metadata ", 400)))
	result := map[string]any{
		"thought_signature": opaque,
		"thinking_bytes":    opaque,
	}

	if TaskResultHasInlineBinary(result) {
		t.Fatal("opaque Base64 metadata, including bytes/buffer-style keys, must not be mistaken for generated media")
	}
}

func TestTaskResultDetectsLongBase64MediaWithoutSemanticKey(t *testing.T) {
	payload := append([]byte{0x89, 'P', 'N', 'G', 0x0d, 0x0a, 0x1a, 0x0a}, bytes.Repeat([]byte{0}, 4096)...)
	result := map[string]any{"provider_payload": base64.StdEncoding.EncodeToString(payload)}

	if !TaskResultHasInlineBinary(result) {
		t.Fatal("signature-detected generated media must still be materialized")
	}
	diagnostics := taskResultInlineBinaryDiagnostics(result)
	if len(diagnostics) != 1 || !strings.Contains(diagnostics[0], "$.provider_payload") || !strings.Contains(diagnostics[0], "contentType=image/png") {
		t.Fatalf("unexpected safe inline binary diagnostics: %+v", diagnostics)
	}
}

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)
	}
}