feat(gateway): 补齐桌面端高级媒体直连接口
ci / verify (pull_request) Successful in 15m34s

新增图片矢量化、视频超分、每日用量、计价与任务隔离能力,并通过环境变量解析平台凭据。

已通过 Go 全量门禁、迁移检查、镜像构建以及 Vectorizer 五格式和 Topaz 3 秒视频真实 DEV 验收。
This commit is contained in:
2026-07-22 14:02:53 +08:00
parent cbebfd7baa
commit 3056cf8fca
37 changed files with 3336 additions and 29 deletions
@@ -0,0 +1,32 @@
package runner
import (
"strings"
"testing"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
func TestCandidateWithEnvironmentCredentialsResolvesNamesAtRuntime(t *testing.T) {
t.Setenv("TEST_GATEWAY_PROVIDER_SECRET", "secret-value")
candidate, err := candidateWithEnvironmentCredentials(store.RuntimeModelCandidate{
Credentials: map[string]any{"safe": "existing"},
PlatformConfig: map[string]any{"credentialEnv": map[string]any{
"apiKey": "TEST_GATEWAY_PROVIDER_SECRET",
}},
})
if err != nil {
t.Fatal(err)
}
if candidate.Credentials["apiKey"] != "secret-value" || candidate.Credentials["safe"] != "existing" {
t.Fatalf("credentials not resolved: %+v", candidate.Credentials)
}
}
func TestCandidateWithEnvironmentCredentialsDoesNotLeakMissingValue(t *testing.T) {
candidate := store.RuntimeModelCandidate{PlatformConfig: map[string]any{"credentialEnv": map[string]any{"apiKey": "MISSING_GATEWAY_PROVIDER_SECRET"}}}
_, err := candidateWithEnvironmentCredentials(candidate)
if err == nil || strings.Contains(err.Error(), "MISSING_GATEWAY_PROVIDER_SECRET") {
t.Fatalf("missing credential error should be generic: %v", err)
}
}
@@ -39,6 +39,20 @@ func TestVideoModelTypeInferenceReadsContentArray(t *testing.T) {
}
}
func TestAdvancedMediaModelTypeInference(t *testing.T) {
if got := modelTypeFromKind("images.vectorize", nil); got != "image_vectorize" {
t.Fatalf("images.vectorize model type = %s", got)
}
if got := modelTypeFromKind("videos.upscales", nil); got != "video_enhance" {
t.Fatalf("videos.upscales model type = %s", got)
}
for _, alias := range []string{"vectorize", "image-vectorizer", "image_vectorize"} {
if got := canonicalModelType(alias); got != "image_vectorize" {
t.Fatalf("canonical image vectorize alias %q = %q", alias, got)
}
}
}
func TestVideoContentTextContributesToTokenEstimate(t *testing.T) {
tokens := estimateRequestTokens(map[string]any{
"model": "demo-video",
+169
View File
@@ -3,6 +3,7 @@ package runner
import (
"context"
"math"
"strconv"
"strings"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
@@ -131,6 +132,20 @@ func (s *Service) billings(ctx context.Context, user *auth.User, kind string, bo
resource = "image_edit"
baseKey = "editBase"
}
if kind == "images.vectorize" {
resource = "image_vectorize"
unit = "conversion"
baseKey = "vectorizeBase"
}
if kind == "videos.upscales" {
resource = "video_enhance"
unit = "5s_video"
baseKey = "videoEnhanceBase"
inputs := resolveVideoEnhancePricingInputs(body, response)
amount, details := videoEnhanceAmount(config, inputs, float64(count), resourcePrice(config, resource, baseKey, "basePrice"))
amount = math.Ceil(amount*1e8) / 1e8 * discount
return []any{billingLineWithDetails(candidate, resource, unit, float64(count)*inputs.DurationUnits, roundPrice(amount), discount, simulated, details)}
}
if kind == "videos.generations" {
resource = "video"
unit = "5s_video"
@@ -193,6 +208,160 @@ func (s *Service) billings(ctx context.Context, user *auth.User, kind string, bo
return []any{billingLine(candidate, resource, unit, count, roundPrice(amount), discount, simulated)}
}
type videoEnhancePricingInputs struct {
DurationSeconds, DurationUnits, SourceFrameRate, TargetFrameRate, SlowMotionRate float64
SourceResolution, TargetResolution, Model string
}
var videoEnhanceResolutionPixels = map[string]float64{
"480p": 854 * 480, "720p": 1280 * 720, "1080p": 1920 * 1080,
"1440p": 2560 * 1440, "2160p": 3840 * 2160,
}
var videoEnhanceResolutionBaseline = map[string]float64{"480p": 1, "720p": 1.4, "1080p": 1.8, "1440p": 2.1, "2160p": 2.6}
func resolveVideoEnhancePricingInputs(body map[string]any, response clients.Response) videoEnhancePricingInputs {
merged := cloneMap(body)
if data, ok := response.Result["data"].([]any); ok && len(data) > 0 {
if item, ok := data[0].(map[string]any); ok {
for _, key := range []string{"duration", "source_resolution", "target_resolution", "source_frame_rate", "target_frame_rate", "slow_motion_rate"} {
if merged[key] == nil && item[key] != nil {
merged[key] = item[key]
}
}
}
}
duration := floatFromAny(merged["duration"])
if duration <= 0 {
duration = 5
}
targetValue := firstNonEmptyString(stringFromMap(merged, "target_resolution"), stringFromMap(merged, "output_resolution"), stringFromMap(merged, "resolution"))
if dimensions := videoEnhanceDimensionsValue(merged, "output_width", "output_height"); dimensions != "" {
targetValue = dimensions
}
target := normalizeVideoEnhanceResolution(targetValue, "1080p")
sourceValue := firstNonEmptyString(stringFromMap(merged, "source_resolution"), stringFromMap(merged, "resolution"))
if dimensions := videoEnhanceDimensionsValue(merged, "source_width", "source_height"); dimensions != "" {
sourceValue = dimensions
}
source := normalizeVideoEnhanceResolution(sourceValue, target)
sourceFPS := floatFromAny(firstPresentValue(merged, "source_frame_rate", "frame_rate", "frameRate"))
if sourceFPS <= 0 {
sourceFPS = 24
}
targetFPS := floatFromAny(firstPresentValue(merged, "target_frame_rate", "output_frame_rate"))
if targetFPS <= 0 {
targetFPS = sourceFPS
}
slowMotion := floatFromAny(merged["slow_motion_rate"])
if slowMotion <= 0 {
slowMotion = 1
}
return videoEnhancePricingInputs{
DurationSeconds: duration, DurationUnits: math.Max(1, math.Round(duration)/5),
SourceResolution: source, TargetResolution: target,
SourceFrameRate: sourceFPS, TargetFrameRate: targetFPS, SlowMotionRate: slowMotion,
Model: firstNonEmptyString(stringFromMap(merged, "model"), stringFromMap(merged, "enhancement_model")),
}
}
func videoEnhanceDimensionsValue(values map[string]any, widthKey, heightKey string) string {
width := int(math.Round(floatFromAny(values[widthKey])))
height := int(math.Round(floatFromAny(values[heightKey])))
if width <= 0 || height <= 0 {
return ""
}
return strconv.Itoa(width) + "x" + strconv.Itoa(height)
}
func normalizeVideoEnhanceResolution(value, fallback string) string {
normalized := strings.TrimSuffix(strings.ToLower(strings.TrimSpace(value)), "_upscale")
if normalized == "2k" {
normalized = "1440p"
}
if normalized == "4k" {
normalized = "2160p"
}
if width, height, ok := parseVideoEnhanceDimensions(normalized); ok {
pixels := float64(width * height)
for _, bucket := range []string{"480p", "720p", "1080p", "1440p", "2160p"} {
if pixels <= videoEnhanceResolutionPixels[bucket] {
return bucket
}
}
return "2160p"
}
if _, ok := videoEnhanceResolutionPixels[normalized]; ok {
return normalized
}
if _, ok := videoEnhanceResolutionPixels[fallback]; ok {
return fallback
}
return "1080p"
}
func parseVideoEnhanceDimensions(value string) (int, int, bool) {
parts := strings.Split(strings.ReplaceAll(value, " ", ""), "x")
if len(parts) != 2 {
return 0, 0, false
}
width, widthErr := strconv.Atoi(parts[0])
height, heightErr := strconv.Atoi(parts[1])
return width, height, widthErr == nil && heightErr == nil && width > 0 && height > 0
}
func videoEnhanceAmount(config map[string]any, input videoEnhancePricingInputs, count, basePrice float64) (float64, map[string]any) {
transitionKey := input.SourceResolution + "->" + input.TargetResolution
minimumResolutionWeight := videoEnhanceResolutionBaseline[input.TargetResolution]
resolutionWeight, configuredResolution := pricingDynamicWeight(config, "video_enhance", "resolutionTransitions", transitionKey)
if !configuredResolution {
resolutionWeight = math.Max(minimumResolutionWeight, videoEnhanceResolutionPixels[input.TargetResolution]/videoEnhanceResolutionPixels[input.SourceResolution])
} else {
resolutionWeight = math.Max(minimumResolutionWeight, resolutionWeight)
}
sourceFPS := math.Max(1, math.Round(input.SourceFrameRate))
targetFPS := math.Max(1, math.Round(input.TargetFrameRate))
frameKey := strconv.Itoa(int(sourceFPS)) + "->" + strconv.Itoa(int(targetFPS))
minimumFrameWeight := math.Max(1, targetFPS/24)
frameWeight, configuredFrame := pricingDynamicWeight(config, "video_enhance", "frameRateTransitions", frameKey)
if !configuredFrame {
frameWeight = minimumFrameWeight
} else {
frameWeight = math.Max(minimumFrameWeight, frameWeight)
}
modelWeight, configuredModel := pricingDynamicWeight(config, "video_enhance", "modelWeights", input.Model)
if !configuredModel {
modelWeight = 1
}
markup, foundMarkup := pricingDynamicWeight(config, "video_enhance", "", "markup")
if !foundMarkup {
markup = 1
}
amount := count * input.DurationUnits * basePrice * resolutionWeight * frameWeight * input.SlowMotionRate * modelWeight * markup
return amount, map[string]any{
"count": count, "durationSeconds": input.DurationSeconds, "durationUnit": "5s", "durationUnitCount": input.DurationUnits,
"sourceResolution": input.SourceResolution, "targetResolution": input.TargetResolution,
"resolutionTransitionKey": transitionKey, "resolutionTransitionWeight": resolutionWeight, "resolutionTransitionFallback": !configuredResolution,
"sourceFrameRate": input.SourceFrameRate, "targetFrameRate": input.TargetFrameRate,
"frameRateTransitionKey": frameKey, "frameRateTransitionWeight": frameWeight, "frameRateTransitionFallback": !configuredFrame,
"slowMotionRate": input.SlowMotionRate, "model": input.Model, "modelWeight": modelWeight, "markup": markup,
}
}
func pricingDynamicWeight(config map[string]any, resource, group, key string) (float64, bool) {
resourceConfig, _ := config[resource].(map[string]any)
dynamic, _ := resourceConfig["dynamicWeight"].(map[string]any)
if group == "" {
value, ok := dynamic[key]
weight := floatFromAny(value)
return weight, ok && weight > 0
}
values, _ := dynamic[group].(map[string]any)
value, ok := values[key]
weight := floatFromAny(value)
return weight, ok && weight > 0
}
func (s *Service) effectiveBillingConfig(ctx context.Context, candidate store.RuntimeModelCandidate) map[string]any {
var inheritedRuleSetConfig map[string]any
if ruleSetID := firstNonEmptyString(candidate.BasePricingRuleSetID, candidate.PlatformPricingRuleSetID); ruleSetID != "" && s.store != nil {
+58
View File
@@ -93,6 +93,64 @@ func TestVideoBillingEstimateProratesFiveSecondUnitsAndDynamicWeights(t *testing
}
}
func TestAdvancedMediaBillingUsesConversionAndTransitionMatrix(t *testing.T) {
service := &Service{}
vectorCandidate := store.RuntimeModelCandidate{ModelName: "easy-image-vectorizer-1", BaseBillingConfig: map[string]any{
"image_vectorize": map[string]any{"basePrice": 200},
}}
vectorLine := firstBillingLine(t, service.billings(context.Background(), nil, "images.vectorize", map[string]any{"count": 1}, vectorCandidate, clients.Response{}, true))
if vectorLine["resourceType"] != "image_vectorize" || vectorLine["unit"] != "conversion" || floatFromAny(vectorLine["amount"]) != 200 {
t.Fatalf("unexpected vectorizer billing: %+v", vectorLine)
}
topazCandidate := store.RuntimeModelCandidate{ModelName: "easy-proteus-standard-4", BaseBillingConfig: map[string]any{
"video_enhance": map[string]any{
"basePrice": 100,
"dynamicWeight": map[string]any{
"resolutionTransitions": map[string]any{"720p->1080p": 1.8},
"frameRateTransitions": map[string]any{"24->60": 2.5},
"markup": 1,
},
},
}}
topazLine := firstBillingLine(t, service.billings(context.Background(), nil, "videos.upscales", map[string]any{
"count": 1, "duration": 5, "source_resolution": "720p", "target_resolution": "1080p",
"source_frame_rate": 24, "target_frame_rate": 60, "slow_motion_rate": 2,
"model": "easy-proteus-standard-4",
}, topazCandidate, clients.Response{}, true))
if got, want := floatFromAny(topazLine["amount"]), 900.0; got != want {
t.Fatalf("video enhance amount=%v, want=%v, line=%+v", got, want, topazLine)
}
if topazLine["resolutionTransitionKey"] != "720p->1080p" || topazLine["frameRateTransitionKey"] != "24->60" {
t.Fatalf("missing transition evidence: %+v", topazLine)
}
}
func TestVideoEnhancePricingFallbackUsesPixelsAndBaseline(t *testing.T) {
amount, details := videoEnhanceAmount(map[string]any{"video_enhance": map[string]any{"dynamicWeight": map[string]any{}}}, videoEnhancePricingInputs{
DurationSeconds: 5, DurationUnits: 1, SourceResolution: "720p", TargetResolution: "2160p",
SourceFrameRate: 30, TargetFrameRate: 60, SlowMotionRate: 1,
}, 1, 100)
if amount != 2250 {
t.Fatalf("fallback amount=%v, want=2250 details=%+v", amount, details)
}
}
func TestVideoEnhancePricingPrefersProbedDimensions(t *testing.T) {
inputs := resolveVideoEnhancePricingInputs(map[string]any{
"duration": 3,
"source_resolution": "180p",
"source_width": 320,
"source_height": 180,
"target_resolution": "720p",
"output_width": 1280,
"output_height": 720,
}, clients.Response{})
if inputs.SourceResolution != "480p" || inputs.TargetResolution != "720p" {
t.Fatalf("dimension buckets = %s->%s, want 480p->720p", inputs.SourceResolution, inputs.TargetResolution)
}
}
func TestMusicBillingUsesSongResourceAndOutputCount(t *testing.T) {
service := &Service{}
candidate := store.RuntimeModelCandidate{
+22
View File
@@ -659,6 +659,28 @@ func (s *Service) billingsWithResolvedPricingV2(
resource = "image_edit"
baseKey = "editBase"
}
if kind == "images.vectorize" {
resource = "image_vectorize"
unit = "conversion"
baseKey = "vectorizeBase"
}
if kind == "videos.upscales" {
resource = "video_enhance"
unit = "5s_video"
baseKey = "videoEnhanceBase"
inputs := resolveVideoEnhancePricingInputs(body, response)
price, priceErr := pricing.requiredPrice(resource, baseKey, "basePrice")
if priceErr != nil {
return nil, 0, resolvedPricing{}, priceErr
}
rawAmount, details := videoEnhanceAmount(pricing.Config, inputs, float64(count), price.Float64())
rawAmount = math.Ceil(rawAmount*1e8) / 1e8
amount, amountErr := fixedAmountFromAny(rawAmount * discount.Float64())
if amountErr != nil {
return nil, 0, resolvedPricing{}, pricing.calculationError(resource, amountErr)
}
return []any{buildLine(resource, unit, float64(count)*inputs.DurationUnits, amount, details)}, amount, pricing, nil
}
if kind == "videos.generations" {
resource = "video"
unit = "5s_video"
+103 -1
View File
@@ -5,6 +5,8 @@ import (
"errors"
"fmt"
"log/slog"
"os"
"regexp"
"strconv"
"strings"
"time"
@@ -90,6 +92,8 @@ func New(cfg config.Config, db *store.Store, logger *slog.Logger, observers ...b
"volces": clients.VolcesClient{HTTPClient: httpClients.none},
"keling": clients.KelingClient{HTTPClient: httpClients.none},
"kling": clients.KelingClient{HTTPClient: httpClients.none},
"vectorizer": clients.VectorizerClient{HTTPClient: httpClients.none},
"topaz": clients.TopazClient{HTTPClient: httpClients.none},
"universal": clients.UniversalClient{HTTPClient: httpClients.none, ScriptExecutor: scriptExecutor},
"simulation": clients.SimulationClient{},
},
@@ -857,6 +861,11 @@ func billingItemsFixedTotal(items []any) fixedAmount {
}
func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user *auth.User, body map[string]any, preprocessing parameterPreprocessingLog, candidate store.RuntimeModelCandidate, pricing resolvedPricing, attemptNo int, onDelta clients.StreamDelta, responseExecution responseExecutionContext, singleSourceProtected bool, cacheAffinityPolicy map[string]any, cacheAffinityRecordKeys []string) (clients.Response, error) {
var err error
candidate, err = candidateWithEnvironmentCredentials(candidate)
if err != nil {
return clients.Response{}, err
}
simulated := isSimulation(task, candidate)
baseAttemptMetrics := mergeMetrics(attemptMetrics(candidate, attemptNo, simulated), parameterPreprocessingMetrics(preprocessing))
reservations := s.rateLimitReservations(ctx, user, candidate, body)
@@ -949,6 +958,15 @@ func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user
})
return clients.Response{}, err
}
providerBody, err = s.preparePrivateProviderRequest(ctx, user, task.Kind, providerBody)
if err != nil {
_ = s.store.FinishTaskAttempt(ctx, store.FinishTaskAttemptInput{
AttemptID: attemptID, Status: "failed", Retryable: false,
Metrics: mergeMetrics(baseAttemptMetrics, map[string]any{"error": err.Error(), "retryable": false}),
ErrorCode: clients.ErrorCode(err), ErrorMessage: err.Error(),
})
return clients.Response{}, err
}
providerBody, err = s.hydrateProviderRequestAssets(ctx, providerBody, candidate)
if err != nil {
_ = s.store.FinishTaskAttempt(ctx, store.FinishTaskAttemptInput{
@@ -1421,6 +1439,8 @@ func modelTypeFromKind(kind string, body map[string]any) string {
return "image_edit"
}
return "image_generate"
case "images.vectorize":
return "image_vectorize"
case "videos.generations":
if videoRequestHasVideoOrAudioReference(body) {
return "omni_video"
@@ -1429,6 +1449,8 @@ func modelTypeFromKind(kind string, body map[string]any) string {
return "image_to_video"
}
return "video_generate"
case "videos.upscales":
return "video_enhance"
case "song.generations", "music.generations":
return "audio_generate"
case "speech.generations":
@@ -1463,6 +1485,10 @@ func canonicalModelType(value string) string {
return "text_to_speech"
case "voice", "voice_clone", "voiceclone", "voice.cloning":
return "voice_clone"
case "vectorize", "image_vectorizer", "image_vectorize":
return "image_vectorize"
case "video_upscale", "video_enhance", "upscale":
return "video_enhance"
default:
return normalized
}
@@ -1470,7 +1496,7 @@ func canonicalModelType(value string) string {
func isKnownModelType(value string) bool {
switch value {
case "text_generate", "text_embedding", "text_rerank", "image_generate", "image_edit", "video_generate", "image_to_video", "text_to_video", "video_edit", "video_reference", "video_first_last_frame", "omni_video", "omni", "audio_generate", "text_to_speech", "voice_clone":
case "text_generate", "text_embedding", "text_rerank", "image_generate", "image_edit", "image_vectorize", "video_generate", "video_enhance", "image_to_video", "text_to_video", "video_edit", "video_reference", "video_first_last_frame", "omni_video", "omni", "audio_generate", "text_to_speech", "voice_clone":
return true
default:
return false
@@ -1703,6 +1729,14 @@ func validateRequest(kind string, body map[string]any) error {
if strings.TrimSpace(stringFromMap(body, "prompt")) == "" {
return errors.New("prompt is required")
}
case "images.vectorize":
if vectorizerSourceURL(body) == "" && vectorizerSourceTaskID(body) == "" {
return &clients.ClientError{Code: "invalid_parameter", Message: "source.url or source.vectorizerTaskId is required", Param: "source", StatusCode: 400}
}
case "videos.upscales":
if firstNonEmptyString(stringFromMap(body, "video_url"), stringFromMap(body, "videoUrl")) == "" {
return &clients.ClientError{Code: "invalid_parameter", Message: "video_url is required", Param: "video_url", StatusCode: 400}
}
case "song.generations", "music.generations":
if strings.TrimSpace(stringFromMap(body, "prompt")) == "" {
return errors.New("prompt is required")
@@ -1722,6 +1756,74 @@ func validateRequest(kind string, body map[string]any) error {
return nil
}
var credentialEnvironmentName = regexp.MustCompile(`^[A-Z_][A-Z0-9_]*$`)
func candidateWithEnvironmentCredentials(candidate store.RuntimeModelCandidate) (store.RuntimeModelCandidate, error) {
references, _ := candidate.PlatformConfig["credentialEnv"].(map[string]any)
if len(references) == 0 {
references, _ = candidate.PlatformConfig["credential_env"].(map[string]any)
}
if len(references) == 0 {
return candidate, nil
}
credentials := cloneMap(candidate.Credentials)
for credentialName, rawEnvironmentName := range references {
environmentName := strings.TrimSpace(fmt.Sprint(rawEnvironmentName))
if !credentialEnvironmentName.MatchString(environmentName) {
return candidate, &clients.ClientError{Code: "missing_credentials", Message: "platform credential environment reference is invalid", Retryable: false}
}
value, ok := os.LookupEnv(environmentName)
if !ok || strings.TrimSpace(value) == "" {
return candidate, &clients.ClientError{Code: "missing_credentials", Message: "platform credential environment variable is not configured", Retryable: false}
}
credentials[credentialName] = value
}
candidate.Credentials = credentials
return candidate, nil
}
func (s *Service) preparePrivateProviderRequest(ctx context.Context, user *auth.User, kind string, body map[string]any) (map[string]any, error) {
if kind != "images.vectorize" {
return body, nil
}
taskID := vectorizerSourceTaskID(body)
if taskID == "" {
return body, nil
}
sourceTask, err := s.store.GetTask(ctx, taskID)
if err != nil || !taskAccessibleToUser(sourceTask, user) || sourceTask.Kind != "images.vectorize" {
return nil, &clients.ClientError{Code: "not_found", Message: "source vectorizer task not found", StatusCode: 404, Retryable: false}
}
imageToken := strings.TrimSpace(fmt.Sprint(sourceTask.RemoteTaskPayload["imageToken"]))
if imageToken == "" {
return nil, &clients.ClientError{Code: "invalid_parameter", Message: "source vectorizer task cannot be reused", Param: "source.vectorizerTaskId", StatusCode: 400, Retryable: false}
}
privateBody := cloneMap(body)
privateBody["_vectorizer_image_token"] = imageToken
if receipt := strings.TrimSpace(fmt.Sprint(sourceTask.RemoteTaskPayload["receipt"])); receipt != "" {
privateBody["_vectorizer_receipt"] = receipt
}
return privateBody, nil
}
func vectorizerSourceURL(body map[string]any) string {
if source, ok := body["source"].(map[string]any); ok {
return strings.TrimSpace(firstNonEmptyString(stringFromMap(source, "url"), stringFromMap(source, "image_url"), stringFromMap(source, "imageUrl")))
}
return strings.TrimSpace(firstNonEmptyString(stringFromMap(body, "image_url"), stringFromMap(body, "imageUrl")))
}
func vectorizerSourceTaskID(body map[string]any) string {
for _, key := range []string{"source", "sourceContext"} {
if source, ok := body[key].(map[string]any); ok {
if taskID := strings.TrimSpace(firstNonEmptyString(stringFromMap(source, "vectorizerTaskId"), stringFromMap(source, "vectorizer_task_id"), stringFromMap(source, "taskId"), stringFromMap(source, "task_id"))); taskID != "" {
return taskID
}
}
}
return strings.TrimSpace(firstNonEmptyString(stringFromMap(body, "vectorizerTaskId"), stringFromMap(body, "vectorizer_task_id")))
}
func hasRerankDocuments(value any) bool {
switch typed := value.(type) {
case []any:
+7
View File
@@ -211,6 +211,13 @@ func taskAccessibleToUser(task store.GatewayTask, user *auth.User) bool {
return userID != "" && strings.TrimSpace(task.UserID) == userID
}
// TaskAccessibleToUser applies the same ownership boundary to every public
// task surface. API-key requests are deliberately scoped to the exact key,
// while JWT requests may see the current gateway user's tasks.
func TaskAccessibleToUser(task store.GatewayTask, user *auth.User) bool {
return taskAccessibleToUser(task, user)
}
func gatewayUserIDForAuth(user *auth.User) string {
if user == nil {
return ""
+80 -3
View File
@@ -8,10 +8,12 @@ import (
"encoding/base64"
"encoding/hex"
"encoding/json"
"errors"
"fmt"
"io"
"mime"
"mime/multipart"
"net"
"net/http"
"net/textproto"
"net/url"
@@ -85,6 +87,11 @@ 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}
}
// 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" {
policy.UploadURLMedia = true
}
decisions := make([]generatedAssetDecision, len(data))
needsUpload := false
changed := false
@@ -723,7 +730,8 @@ func (s *Service) readGeneratedURLAsset(ctx context.Context, asset *generatedURL
if err != nil {
return nil, "", err
}
resp, err := http.DefaultClient.Do(req)
allowLocal := strings.HasPrefix(strings.TrimSpace(asset.URL), "/")
resp, err := generatedAssetHTTPClient(allowLocal).Do(req)
if err != nil {
return nil, "", &clients.ClientError{Code: "upload_source_fetch_failed", Message: err.Error(), Retryable: true}
}
@@ -752,6 +760,47 @@ func (s *Service) readGeneratedURLAsset(ctx context.Context, asset *generatedURL
return payload, strings.TrimSpace(strings.Split(contentType, ";")[0]), nil
}
func generatedAssetHTTPClient(allowLocal bool) *http.Client {
transport := http.DefaultTransport.(*http.Transport).Clone()
transport.Proxy = nil
transport.DialContext = func(ctx context.Context, network, address string) (net.Conn, error) {
host, port, err := net.SplitHostPort(address)
if err != nil {
return nil, err
}
addresses, err := net.DefaultResolver.LookupIPAddr(ctx, host)
if err != nil || len(addresses) == 0 {
return nil, errors.New("generated media DNS resolution failed")
}
for _, address := range addresses {
if generatedAssetBlockedAddress(address.IP) && !(allowLocal && address.IP.IsLoopback()) {
return nil, errors.New("generated media resolved to a blocked network")
}
}
dialer := &net.Dialer{Timeout: 10 * time.Second}
var attempts []error
for _, resolved := range addresses {
connection, dialErr := dialer.DialContext(ctx, network, net.JoinHostPort(resolved.IP.String(), port))
if dialErr == nil {
return connection, nil
}
attempts = append(attempts, dialErr)
}
return nil, errors.Join(attempts...)
}
return &http.Client{
Transport: transport,
Timeout: 10 * time.Minute,
CheckRedirect: func(*http.Request, []*http.Request) error {
return http.ErrUseLastResponse
},
}
}
func generatedAssetBlockedAddress(ip net.IP) bool {
return ip == nil || ip.IsLoopback() || ip.IsPrivate() || ip.IsLinkLocalUnicast() || ip.IsLinkLocalMulticast() || ip.IsUnspecified() || ip.IsMulticast()
}
func (s *Service) generatedAssetFetchURL(raw string) (string, error) {
value := strings.TrimSpace(raw)
if value == "" {
@@ -1304,6 +1353,9 @@ func mediaContentTypeFromItem(item map[string]any) string {
func mediaKindForAsset(taskKind string, item map[string]any, sourceKey string, contentType string) string {
contentType = strings.ToLower(strings.TrimSpace(contentType))
if generatedContentTypeIsDocument(contentType) {
return "file"
}
if strings.HasPrefix(contentType, "image/") {
return "image"
}
@@ -1314,6 +1366,9 @@ func mediaKindForAsset(taskKind string, item map[string]any, sourceKey string, c
return "audio"
}
itemType := strings.ToLower(strings.TrimSpace(stringFromAny(item["type"])))
if itemType == "file" || strings.Contains(itemType, "document") {
return "file"
}
if strings.Contains(itemType, "video") {
return "video"
}
@@ -1346,6 +1401,8 @@ func defaultContentTypeForGeneratedAsset(kind string) string {
return "video/mp4"
case "audio":
return "audio/mpeg"
case "file":
return "application/octet-stream"
default:
return "image/png"
}
@@ -1354,10 +1411,10 @@ func defaultContentTypeForGeneratedAsset(kind string) string {
func resolvedGeneratedAssetContentType(declared string, kind string, payload []byte) string {
declared = normalizeGeneratedContentType(declared)
detected := detectGeneratedAssetContentType(payload)
if generatedContentTypeIsMedia(detected) {
if generatedContentTypeIsMedia(detected) || generatedContentTypeIsDocument(detected) {
return detected
}
if generatedContentTypeIsMedia(declared) {
if generatedContentTypeIsMedia(declared) || generatedContentTypeIsDocument(declared) {
return declared
}
return defaultContentTypeForGeneratedAsset(kind)
@@ -1380,6 +1437,15 @@ func generatedContentTypeIsMedia(contentType string) bool {
strings.HasPrefix(contentType, "audio/")
}
func generatedContentTypeIsDocument(contentType string) bool {
switch normalizeGeneratedContentType(contentType) {
case "application/pdf", "application/postscript", "application/eps", "application/dxf", "image/vnd.dxf":
return true
default:
return false
}
}
func generatedAssetKindFromContentType(fallback string, contentType string) string {
contentType = normalizeGeneratedContentType(contentType)
if strings.HasPrefix(contentType, "image/") {
@@ -1391,6 +1457,9 @@ func generatedAssetKindFromContentType(fallback string, contentType string) stri
if strings.HasPrefix(contentType, "audio/") {
return "audio"
}
if generatedContentTypeIsDocument(contentType) {
return "file"
}
fallback = strings.ToLower(strings.TrimSpace(fallback))
if fallback != "" {
return fallback
@@ -1434,6 +1503,14 @@ func randomHexSuffix(byteCount int) string {
func fileExtensionForContentType(contentType string, kind string) string {
normalized := strings.ToLower(strings.TrimSpace(strings.Split(contentType, ";")[0]))
switch normalized {
case "application/pdf":
return ".pdf"
case "application/postscript", "application/eps":
return ".eps"
case "application/dxf", "image/vnd.dxf":
return ".dxf"
}
switch normalized {
case "image/jpeg", "image/jpg":
return ".jpg"
case "image/webp":
+19
View File
@@ -52,6 +52,25 @@ func TestGeneratedAssetDecisionUploadsInlineImageBase64(t *testing.T) {
}
}
func TestGeneratedAssetDecisionUploadsVectorDocumentWithoutChangingType(t *testing.T) {
item := map[string]any{
"type": "file",
"b64_json": base64.StdEncoding.EncodeToString([]byte("%PDF-1.7\n")),
"mime_type": "application/pdf",
}
decision, err := generatedAssetDecisionForItem("images.vectorize", item, defaultGeneratedAssetUploadPolicy())
if err != nil {
t.Fatal(err)
}
if decision.Inline == nil || decision.Inline.Kind != "file" || decision.Inline.ContentType != "application/pdf" {
t.Fatalf("unexpected vector document decision: %+v", decision)
}
contentType := resolvedGeneratedAssetContentType(decision.Inline.ContentType, decision.Inline.Kind, decision.Inline.Bytes)
if contentType != "application/pdf" || fileExtensionForContentType(contentType, "file") != ".pdf" {
t.Fatalf("vector document type changed: contentType=%s", contentType)
}
}
func TestGeneratedAssetDecisionUploadsInlineVideoBuffer(t *testing.T) {
item := map[string]any{
"type": "video",
+4
View File
@@ -215,6 +215,10 @@ func (s *Service) DeleteClonedVoice(ctx context.Context, user *auth.User, rawID
if !ok {
return DeletedClonedVoiceResult{}, &clients.ClientError{Code: "cloned_voice_platform_unavailable", Message: "cloned voice platform binding is unavailable", StatusCode: 400, Retryable: false}
}
candidate, err = candidateWithEnvironmentCredentials(candidate)
if err != nil {
return DeletedClonedVoiceResult{}, err
}
requestHTTPClient, err := s.httpClientForCandidate(candidate, false)
if err != nil {
return DeletedClonedVoiceResult{}, err