chore(git): 同步主分支发布与业务变更

# Conflicts:
#	apps/api/go.mod
#	apps/api/go.sum
This commit is contained in:
2026-07-28 10:47:41 +08:00
238 changed files with 25370 additions and 2079 deletions
+17
View File
@@ -29,6 +29,11 @@ IDENTITY_MODE=hybrid
# - hold: reject new production generation before any upstream request; existing settlements continue.
BILLING_ENGINE_MODE=observe
# River 执行 worker 会按平台模型和活跃用户组中更严格的 concurrent 总量每 5 秒重算。
# hard limit 是单进程安全边界;平台模型/用户组 concurrency lease 才是业务并发真值。
AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT=2048
AI_GATEWAY_ASYNC_WORKER_REFRESH_INTERVAL_SECONDS=5
# Unified identity business settings are managed in System Settings > Unified
# Identity. Deployment only supplies the SecretStore and infrastructure timing.
AI_GATEWAY_PUBLIC_BASE_URL=http://localhost:8088
@@ -68,6 +73,18 @@ TASK_PROGRESS_CALLBACK_ENABLED=true
TASK_PROGRESS_CALLBACK_URL=http://localhost:3000/internal/platform/task-progress-callbacks
TASK_PROGRESS_CALLBACK_TIMEOUT_MS=5000
TASK_PROGRESS_CALLBACK_MAX_ATTEMPTS=10
# First deploy with cleanup disabled; enable only after compatibility and
# asynchronous-resume verification has passed.
AI_GATEWAY_TASK_CLEANUP_ENABLED=false
AI_GATEWAY_TASK_RETENTION_DAYS=30
AI_GATEWAY_TASK_ANALYSIS_RETENTION_DAYS=7
AI_GATEWAY_TASK_CLEANUP_INTERVAL_SECONDS=300
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_ASYNC_QUEUE_WORKER_ENABLED=true
CORS_ALLOWED_ORIGIN=http://localhost:5178,http://127.0.0.1:5178
VITE_GATEWAY_API_BASE_URL=http://localhost:8088
+3 -3
View File
@@ -52,14 +52,14 @@ docker compose -f docker-compose.yml config --quiet
2. `main` 不设置受保护分支或 required statusAgent 在获得用户明确授权后,可以合并已验证的 PR,也可以直接提交并推送 `main`。操作前必须获取最新 `origin/main`、确认工作区边界、检查提交差异和敏感信息,并执行与风险匹配的本地验证。
3. 不得等待、伪造或手工补写已经取消的 CI status,也不得为了满足旧流程擅自恢复 Gitea Actions。Git 提交、合并、Push 和 Tag 授权只覆盖源码历史变更,永远不等于 publish 或 deploy 授权。
4. 发布只能使用工作区干净、已经提交并属于 `origin/main` 历史的完整 SHA。镜像 Tag 使用完整 SHA,线上只能使用 Registry digest,禁止使用 `latest`
5. 本地镜像发布必须由用户明确要求后单独执行:
5. 本地镜像发布必须由用户明确要求后执行:
```bash
./scripts/publish-release-images.sh --components auto
```
该命令只能构建、冒烟、推送镜像和生成 `dist/releases/<SHA>.json`,不得修改生产。
6. publish 成功后,Agent 必须报告 manifest、组件、digest 和验证结果并停止。只有用户再次明确确认上线,才可单独执行
该命令只能构建、冒烟、推送镜像和生成 `dist/releases/<SHA>.json`,不得修改生产。用户只要求“构建镜像”“推送镜像”或明确限定为 publish 时,执行到此为止。
6. 用户明确要求“发布”“上线”或“部署到线上”时,该次授权同时覆盖 publish 和 deploy。publish 成功并核验 manifest、组件、digest 和验证结果后,应直接继续执行下列部署命令,无需再次请求确认
```bash
./scripts/deploy-production-release.sh dist/releases/<SHA>.json
+3 -1
View File
@@ -32,7 +32,8 @@ RUN --mount=type=cache,target=/go/pkg/mod \
--mount=type=cache,target=/root/.cache/go-build \
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-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
@@ -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-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
RUN mkdir -p /app/data/static/generated /app/data/static/uploaded && \
+8 -3
View File
@@ -104,14 +104,14 @@ docker compose -f docker-compose.yml restart web
本仓库没有 Gitea Actions、Webhook、Tag、`main` Push、轮询或定时发布。`main` 不使用受保护分支限制;commit、push 和 Tag 都不会构建镜像或更新生产。
生产发布由 Agent 在用户明确指令下分两次执行。第一步在本机构建 `linux/amd64` 镜像、运行临时 PostgreSQL + simulation API 冒烟、推送完整 SHA Tag,并生成带内容完整性校验的 digest-pinned manifest;该命令不会修改生产:
生产发布由 Agent 在用户明确指令下按两个阶段连续执行。第一步在本机构建 `linux/amd64` 镜像、运行临时 PostgreSQL + simulation API 冒烟、推送完整 SHA Tag,并生成带内容完整性校验的 digest-pinned manifest;该命令不会修改生产:
```bash
docker login --username=<your-aliyun-account> registry.cn-shanghai.aliyuncs.com
./scripts/publish-release-images.sh --components auto
```
Agent 必须报告 `dist/releases/<SHA>.json` 后停止。用户再次明确确认上线后,才执行第二步
用户要求“发布”“上线”或“部署到线上”时,Agent 核验并报告 `dist/releases/<SHA>.json`、组件和 digest 后直接执行第二步,无需再次确认;用户仅要求构建或推送镜像时才在第一步后停止
```bash
./scripts/deploy-production-release.sh dist/releases/<SHA>.json
@@ -124,7 +124,10 @@ Agent 必须报告 `dist/releases/<SHA>.json` 后停止。用户再次明确确
./scripts/deploy-production-release.sh --rollback <历史完整 Git SHA>
```
完整安装、验证、失败处理和停用旧自动化步骤见[人工生产发布运行手册](docs/runbooks/production-ci-cd.md),决策背景见 [ADR-003](docs/decisions/003-manual-agent-release.md)。
完整安装、验证、失败处理和停用旧自动化步骤见[人工生产发布运行手册](docs/runbooks/production-ci-cd.md)
三节点 K3s、CloudNativePG、双 NGINX、备份恢复和跨节点文件验收见
[K3s 高可用运行手册](docs/operations/k3s-ha-runbook.md)。决策背景见
[ADR-003](docs/decisions/003-manual-agent-release.md)。
Compose 默认使用独立容器数据库 `postgres:18-alpine`,数据卷会保留在 `postgres_data``api_data`。为避免本地开发 `.env` 中的 `localhost` 数据库地址污染容器部署,compose 使用 `AI_GATEWAY_COMPOSE_*` 变量作为容器部署专用覆盖,例如:
@@ -159,6 +162,8 @@ AI_GATEWAY_DATABASE_URL=postgresql://easyai:easyai2025@localhost:5432/easyai_ai_
如果现有 `easyai-pgvector` 没有把 `5432` 映射到宿主机,就需要补端口映射,或者把 AI Gateway 后端容器化后接入同一个 `easyai` Docker network。
异步队列 worker 不使用固定并发。服务分别汇总启用平台模型和活跃用户组的有效 `concurrent` 策略,采用两者中更严格的容量,默认每 5 秒在线调整 River 执行容量。策略解析兼容历史 `platformLimits/modelLimits.max_concurrent_requests`,运行时统一转换为 `rules``AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT` 默认 `2048`,仅作为单进程资源安全边界;平台模型和用户组的 PostgreSQL concurrency lease 才是业务并发真值。可通过 `AI_GATEWAY_ASYNC_WORKER_REFRESH_INTERVAL_SECONDS` 调整刷新周期。
## 迁移原则
1. 新服务先并行运行,不直接删除 `easyai-server-main` 内现有模块。
@@ -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
}
+37
View File
@@ -13,6 +13,11 @@ import (
"github.com/jackc/pgx/v5"
)
const (
noTransactionMigrationMarker = "-- easyai:migration:no-transaction"
migrationStatementSeparator = "-- easyai:migration:statement"
)
func main() {
cfg := config.Load()
logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
@@ -63,6 +68,22 @@ CREATE TABLE IF NOT EXISTS schema_migrations (
os.Exit(1)
}
noTransaction, statements := migrationStatements(string(sqlBytes))
if noTransaction {
for _, statement := range statements {
if _, err := conn.Exec(ctx, statement); err != nil {
logger.Error("execute non-transaction migration failed", "version", version, "error", err)
os.Exit(1)
}
}
if _, err := conn.Exec(ctx, "INSERT INTO schema_migrations(version) VALUES($1)", version); err != nil {
logger.Error("record non-transaction migration failed", "version", version, "error", err)
os.Exit(1)
}
logger.Info("migration applied", "version", version, "transactional", false)
continue
}
tx, err := conn.Begin(ctx)
if err != nil {
logger.Error("begin migration failed", "version", version, "error", err)
@@ -95,3 +116,19 @@ CREATE TABLE IF NOT EXISTS schema_migrations (
fmt.Println("migrations complete")
}
func migrationStatements(sql string) (bool, []string) {
trimmed := strings.TrimSpace(sql)
if !strings.HasPrefix(trimmed, noTransactionMigrationMarker) {
return false, []string{sql}
}
trimmed = strings.TrimSpace(strings.TrimPrefix(trimmed, noTransactionMigrationMarker))
parts := strings.Split(trimmed, migrationStatementSeparator)
statements := make([]string, 0, len(parts))
for _, part := range parts {
if statement := strings.TrimSpace(part); statement != "" {
statements = append(statements, statement)
}
}
return true, statements
}
+38
View File
@@ -14,6 +14,44 @@ import (
"github.com/jackc/pgx/v5/pgxpool"
)
func TestMigrationStatementsSupportsConcurrentIndexes(t *testing.T) {
payload := `-- easyai:migration:no-transaction
CREATE INDEX CONCURRENTLY IF NOT EXISTS first_index ON first_table(id);
-- easyai:migration:statement
CREATE INDEX CONCURRENTLY IF NOT EXISTS second_index ON second_table(id);
`
noTransaction, statements := migrationStatements(payload)
if !noTransaction {
t.Fatal("expected a non-transaction migration")
}
if len(statements) != 2 {
t.Fatalf("statements=%d, want 2", len(statements))
}
if !strings.Contains(statements[0], "first_index") || !strings.Contains(statements[1], "second_index") {
t.Fatalf("unexpected statements: %#v", statements)
}
}
func TestSeedanceInputImageConstraintMigrationKeepsCatalogSnapshotsInSync(t *testing.T) {
payload, err := os.ReadFile("../../migrations/0078_seedance_input_image_constraints.sql")
if err != nil {
t.Fatal(err)
}
content := string(payload)
for _, required := range []string{
"easyai:豆包Seedance-2.0",
"easyai:豆包Seedance-2.0-fast",
"input_image_resolution_range",
"input_image_aspect_ratio_range",
"'{metadata,rawModel,capabilities}'",
"UPDATE platform_models",
} {
if !strings.Contains(content, required) {
t.Fatalf("Seedance input image migration is missing %q", required)
}
}
}
func TestSecurityEventSchemaMigrationsDefineCurrentLifecycle(t *testing.T) {
streamPayload, err := os.ReadFile("../../migrations/0063_oidc_security_events.sql")
if err != nil {
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
+1
View File
@@ -15,6 +15,7 @@ require (
github.com/riverqueue/river/riverdriver/riverpgxv5 v0.24.0
github.com/riverqueue/river/rivertype v0.24.0
golang.org/x/crypto v0.52.0
golang.org/x/image v0.43.0
golang.org/x/net v0.54.0
golang.org/x/oauth2 v0.36.0
)
+2
View File
@@ -73,6 +73,8 @@ golang.org/x/crypto v0.52.0 h1:RMs7fP2rXdep0CftQlK8Uf+kibLm7qkCcradZWYz988=
golang.org/x/crypto v0.52.0/go.mod h1:1QgfPxDqh0T2M/elOJtp9RvuR95kVjir0e6/BvEmGbc=
golang.org/x/net v0.54.0 h1:2zJIZAxAHV/OHCDTCOHAYehQzLfSXuf/5SoL/Dv6w/w=
golang.org/x/net v0.54.0/go.mod h1:Sj4oj8jK6XmHpBZU/zWHw3BV3abl4Kvi+Ut7cQcY+cQ=
golang.org/x/image v0.43.0 h1:FLxcP4ec2350nTfOC8ysKtqYSIFbk/QGjw1ZHNP4tsY=
golang.org/x/image v0.43.0/go.mod h1:rrpelvGFt+kLPAjPM4HeWPgrl0FtafueU//e5N0qk/Q=
golang.org/x/oauth2 v0.36.0 h1:peZ/1z27fi9hUOFCAZaHyrpWG5lwe0RJEEEeH0ThlIs=
golang.org/x/oauth2 v0.36.0/go.mod h1:YDBUJMTkDnJS+A4BP4eZBjCqtokkg1hODuPjwiGPO7Q=
golang.org/x/sync v0.21.0 h1:HLII4xRRTtCRkxYp4HNFF0Js/Og6q2i++KXbg0gHCwM=
@@ -78,6 +78,15 @@ func applyAliyunReasoning(body map[string]any, candidate store.RuntimeModelCandi
body["enable_thinking"] = true
model := chatReasoningModelName(body, candidate)
if isAliyunQwen38MaxPreview(model) {
if budget, ok := positiveIntFromAny(body["thinking_budget_tokens"]); ok {
body["thinking_budget"] = budget
delete(body, "reasoning_effort")
return
}
body["reasoning_effort"] = qwen38MaxPreviewReasoningEffort(effort)
return
}
if isAliyunHighMaxReasoningModel(model) {
body["reasoning_effort"] = highMaxReasoningEffort(effort)
return
@@ -200,6 +209,10 @@ func isAliyunHighMaxReasoningModel(model string) bool {
return strings.Contains(model, "deepseek-v4") || strings.HasPrefix(model, "glm-")
}
func isAliyunQwen38MaxPreview(model string) bool {
return strings.Contains(model, "qwen3.8-max-preview")
}
func isAliyunThinkingBudgetModel(model string) bool {
return strings.Contains(model, "qwen") || strings.Contains(model, "qwq") || strings.Contains(model, "qvq") || strings.Contains(model, "kimi")
}
@@ -219,6 +232,17 @@ func highMaxReasoningEffort(effort string) string {
return "high"
}
func qwen38MaxPreviewReasoningEffort(effort string) string {
switch effort {
case "low", "minimal":
return "low"
case "medium":
return "medium"
default:
return "xhigh"
}
}
func zhipuReasoningEffort(effort string) string {
if effort == "max" {
return "xhigh"
+232 -4
View File
@@ -344,6 +344,203 @@ func TestOpenAIClientChatContract(t *testing.T) {
}
}
func TestOpenAIClientAliyunChatSendsStandardInputAudioAndNormalizesLegacyAudioURL(t *testing.T) {
var captured map[string]any
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if err := json.NewDecoder(r.Body).Decode(&captured); err != nil {
t.Fatalf("decode request: %v", err)
}
_ = json.NewEncoder(w).Encode(map[string]any{
"id": "chatcmpl-audio",
"object": "chat.completion",
"model": captured["model"],
"choices": []any{map[string]any{
"message": map[string]any{"role": "assistant", "content": "ok"},
}},
})
}))
defer server.Close()
_, err := (OpenAIClient{HTTPClient: server.Client()}).Run(context.Background(), Request{
Kind: "chat.completions",
Model: "qwen-omni",
Body: map[string]any{
"model": "qwen-omni",
"messages": []any{
map[string]any{
"role": "user",
"content": []any{
map[string]any{
"type": "input_audio",
"input_audio": map[string]any{
"data": "https://cdn.example.com/audio/standard.wav",
"format": "wav",
},
},
map[string]any{
"type": "audio_url",
"audio_url": map[string]any{"url": "https://cdn.example.com/audio/legacy.m4a"},
},
},
},
},
},
Candidate: store.RuntimeModelCandidate{
Provider: "aliyun-bailian-openai",
BaseURL: server.URL,
ProviderModelName: "qwen-omni",
Credentials: map[string]any{"apiKey": "test-key"},
},
})
if err != nil {
t.Fatalf("run aliyun openai-compatible client: %v", err)
}
messages, _ := captured["messages"].([]any)
message, _ := messages[0].(map[string]any)
content, _ := message["content"].([]any)
standard, _ := content[0].(map[string]any)
if standard["type"] != "input_audio" {
t.Fatalf("standard input_audio was not preserved: %+v", standard)
}
standardAudio, _ := standard["input_audio"].(map[string]any)
if standardAudio["data"] != "https://cdn.example.com/audio/standard.wav" || standardAudio["format"] != "wav" {
t.Fatalf("unexpected standard input_audio: %+v", standardAudio)
}
legacy, _ := content[1].(map[string]any)
if legacy["type"] != "input_audio" {
t.Fatalf("legacy audio_url was not normalized: %+v", legacy)
}
legacyAudio, _ := legacy["input_audio"].(map[string]any)
if legacyAudio["data"] != "https://cdn.example.com/audio/legacy.m4a" || legacyAudio["format"] != "m4a" {
t.Fatalf("unexpected normalized legacy audio: %+v", legacyAudio)
}
}
func TestOpenAIClientImageEditUsesMultipartWithMultipleImages(t *testing.T) {
var contentType string
var receivedFields map[string][]string
var receivedImageCount int
var receivedImageFields []string
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
contentType = r.Header.Get("Content-Type")
if err := r.ParseMultipartForm(1 << 20); err != nil {
t.Fatalf("parse multipart image edit: %v", err)
}
receivedFields = r.MultipartForm.Value
for field, files := range r.MultipartForm.File {
receivedImageFields = append(receivedImageFields, field)
receivedImageCount += len(files)
}
_ = json.NewEncoder(w).Encode(map[string]any{
"data": []any{map[string]any{"b64_json": "aW1hZ2U="}},
})
}))
defer server.Close()
_, err := (OpenAIClient{HTTPClient: server.Client()}).Run(context.Background(), Request{
Kind: "images.edits",
ModelType: "image_edit",
Model: "gpt-image-2",
Body: map[string]any{
"model": "gpt-image-2",
"prompt": "combine the references",
"size": "1024x1536",
"quality": "medium",
"aspect_ratio": "2:3",
"resolution": "2K",
"width": 1024,
"height": 1536,
"images": []any{
"data:image/png;base64," + base64.StdEncoding.EncodeToString([]byte("image one")),
"data:image/jpeg;base64," + base64.StdEncoding.EncodeToString([]byte("image two")),
},
"_metadata": map[string]any{"private": true},
},
Candidate: store.RuntimeModelCandidate{
BaseURL: server.URL,
Provider: "openai",
ProviderModelName: "gpt-image-2-vip",
ModelType: "image_edit",
Credentials: map[string]any{"apiKey": "openai-key"},
},
})
if err != nil {
t.Fatalf("run OpenAI image edit: %v", err)
}
if !strings.HasPrefix(contentType, "multipart/form-data; boundary=") {
t.Fatalf("image edit should use multipart/form-data, got %q", contentType)
}
if receivedImageCount != 2 || len(receivedImageFields) != 1 || receivedImageFields[0] != "image" {
t.Fatalf("unexpected multipart image files: fields=%v count=%d", receivedImageFields, receivedImageCount)
}
if receivedFields["model"][0] != "gpt-image-2-vip" ||
receivedFields["prompt"][0] != "combine the references" ||
receivedFields["size"][0] != "1024x1536" ||
receivedFields["quality"][0] != "medium" {
t.Fatalf("unexpected multipart fields: %+v", receivedFields)
}
if _, ok := receivedFields["_metadata"]; ok {
t.Fatalf("internal metadata must not be forwarded: %+v", receivedFields)
}
for _, field := range []string{"aspect_ratio", "resolution", "width", "height"} {
if _, ok := receivedFields[field]; ok {
t.Fatalf("generic image geometry field %q must not be forwarded: %+v", field, receivedFields)
}
}
}
func TestOpenAIClientImageGenerationUsesSizeWithoutGenericGeometryFields(t *testing.T) {
var received map[string]any
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if err := json.NewDecoder(r.Body).Decode(&received); err != nil {
t.Fatalf("decode image generation body: %v", err)
}
_ = json.NewEncoder(w).Encode(map[string]any{
"data": []any{map[string]any{"b64_json": "aW1hZ2U="}},
})
}))
defer server.Close()
_, err := (OpenAIClient{HTTPClient: server.Client()}).Run(context.Background(), Request{
Kind: "images.generations",
ModelType: "image_generate",
Model: "gpt-image-2",
Body: map[string]any{
"model": "gpt-image-2",
"prompt": "draw",
"size": "2048x1152",
"quality": "high",
"aspect_ratio": "16:9",
"aspectRatio": "16:9",
"resolution": "2K",
"width": 2048,
"height": 1152,
},
Candidate: store.RuntimeModelCandidate{
BaseURL: server.URL,
Provider: "openai",
ProviderModelName: "gpt-image-2-vip",
ModelType: "image_generate",
Credentials: map[string]any{"apiKey": "openai-key"},
},
})
if err != nil {
t.Fatalf("run OpenAI image generation: %v", err)
}
if received["model"] != "gpt-image-2-vip" ||
received["prompt"] != "draw" ||
received["size"] != "2048x1152" ||
received["quality"] != "high" {
t.Fatalf("unexpected image generation fields: %+v", received)
}
for _, field := range []string{"aspect_ratio", "aspectRatio", "resolution", "width", "height"} {
if _, ok := received[field]; ok {
t.Fatalf("generic image geometry field %q must not be forwarded: %+v", field, received)
}
}
}
func TestOpenAIClientChatReasoningParamsByProvider(t *testing.T) {
cases := []struct {
name string
@@ -405,6 +602,36 @@ func TestOpenAIClientChatReasoningParamsByProvider(t *testing.T) {
}
},
},
{
name: "aliyun qwen 3.8 preview preserves supported effort",
provider: "aliyun-bailian-openai",
model: "qwen3.8-max-preview",
body: map[string]any{
"reasoning_effort": "high",
},
assertion: func(t *testing.T, body map[string]any) {
if body["enable_thinking"] != true || body["reasoning_effort"] != "xhigh" {
t.Fatalf("aliyun qwen 3.8 should use xhigh reasoning effort, got %+v", body)
}
},
},
{
name: "aliyun qwen 3.8 preview uses explicit budget exclusively",
provider: "aliyun-bailian-openai",
model: "qwen3.8-max-preview",
body: map[string]any{
"reasoning_effort": "medium",
"thinking_budget_tokens": 16384,
},
assertion: func(t *testing.T, body map[string]any) {
if body["enable_thinking"] != true || body["thinking_budget"] != float64(16384) {
t.Fatalf("aliyun qwen 3.8 should use the explicit thinking budget, got %+v", body)
}
if _, ok := body["reasoning_effort"]; ok {
t.Fatalf("aliyun qwen 3.8 must not combine reasoning effort with thinking budget: %+v", body)
}
},
},
{
name: "same deepseek model on aliyun uses aliyun protocol",
provider: "aliyun-bailian-openai",
@@ -1977,7 +2204,7 @@ func TestVolcesClientVideoSubmitsAndPollsTask(t *testing.T) {
}
data, _ := response.Result["data"].([]any)
item, _ := data[0].(map[string]any)
if item["url"] != "https://example.com/out.mp4" || response.Usage.TotalTokens != 9 {
if item["url"] != "https://example.com/out.mp4" || response.Result["content"] != nil || response.Result["usage"] != nil || response.Usage.TotalTokens != 9 {
t.Fatalf("unexpected response: %+v usage=%+v", response.Result, response.Usage)
}
}
@@ -2021,8 +2248,9 @@ func TestVolcesClientVideoRetriesTransientPollAndKeepsOfficialResult(t *testing.
if polls != 2 || len(persisted) != 1 || persisted[0] != "cgt-retry:succeeded" {
t.Fatalf("unexpected poll state polls=%d persisted=%+v", polls, persisted)
}
if response.Result["updated_at"] != float64(124) || response.Result["seed"] != float64(7) || response.Result["raw"] == nil {
t.Fatalf("official result fields lost: %+v", response.Result)
if response.Result["seed"] != float64(7) || response.Result["raw"] != nil || response.Result["updated_at"] != nil ||
response.Result["content"] != nil || response.Result["usage"] != nil || response.Result["upstream_task_id"] != nil {
t.Fatalf("provider response should be reduced to the canonical result: %+v", response.Result)
}
}
@@ -2300,7 +2528,7 @@ func TestVolcesClientVideoResumePollsExistingTaskID(t *testing.T) {
}
data, _ := response.Result["data"].([]any)
item, _ := data[0].(map[string]any)
if response.Result["upstream_task_id"] != "cgt-existing" || item["url"] != "https://example.com/resumed.mp4" {
if response.Result["upstream_task_id"] != nil || item["url"] != "https://example.com/resumed.mp4" {
t.Fatalf("unexpected resumed response: %+v", response.Result)
}
}
+23 -10
View File
@@ -272,24 +272,39 @@ func geminiGenerationConfig(body map[string]any, imageResponse bool) map[string]
if out == nil {
out = map[string]any{}
}
if aspectRatio := firstNonEmptyString(body["aspect_ratio"], body["aspectRatio"]); aspectRatio != "" {
imageConfig := cloneMapAny(mapFromAny(firstPresent(out["imageConfig"], out["image_config"])))
imageConfig := cloneMapAny(mapFromAny(firstPresent(out["imageConfig"], out["image_config"])))
if imageConfig != nil {
aspectRatio := normalizedProviderAspectRatio(firstNonEmptyString(imageConfig["aspectRatio"], imageConfig["aspect_ratio"]))
imageSize := normalizedProviderImageResolution(firstNonEmptyString(imageConfig["imageSize"], imageConfig["image_size"]))
delete(imageConfig, "aspectRatio")
delete(imageConfig, "aspect_ratio")
delete(imageConfig, "imageSize")
delete(imageConfig, "image_size")
if aspectRatio != "" {
imageConfig["aspectRatio"] = aspectRatio
}
if imageSize != "" {
imageConfig["imageSize"] = imageSize
}
}
if aspectRatio := normalizedProviderAspectRatio(firstNonEmptyString(body["aspect_ratio"], body["aspectRatio"], body["ratio"])); aspectRatio != "" {
if imageConfig == nil {
imageConfig = map[string]any{}
}
imageConfig["aspectRatio"] = aspectRatio
out["imageConfig"] = imageConfig
delete(out, "image_config")
}
if imageSize := firstNonEmptyString(body["resolution"], body["imageSize"], body["image_size"], body["size"]); imageSize != "" {
imageConfig := cloneMapAny(mapFromAny(firstPresent(out["imageConfig"], out["image_config"])))
if imageSize := normalizedProviderImageResolution(firstNonEmptyString(body["resolution"], body["imageSize"], body["image_size"], body["size"])); imageSize != "" {
if imageConfig == nil {
imageConfig = map[string]any{}
}
imageConfig["imageSize"] = imageSize
out["imageConfig"] = imageConfig
delete(out, "image_config")
}
if len(imageConfig) > 0 {
out["imageConfig"] = imageConfig
} else {
delete(out, "imageConfig")
}
delete(out, "image_config")
if count := intFromAny(firstPresent(body["n"], body["candidateCount"], body["candidate_count"])); count > 0 {
out["candidateCount"] = count
}
@@ -627,7 +642,6 @@ func geminiResult(request Request, raw map[string]any) map[string]any {
"created": nowUnix(),
"model": request.Model,
"data": data,
"raw": raw,
}
}
message, finishReason := geminiChatMessage(raw)
@@ -642,7 +656,6 @@ func geminiResult(request Request, raw map[string]any) map[string]any {
"message": message,
}},
"usage": geminiUsageMap(raw),
"raw": raw,
}
}
+40 -1
View File
@@ -13,6 +13,11 @@ import (
"time"
)
const (
defaultMaxJSONResponseBytes int64 = 16 << 20
mediaMaxJSONResponseBytes int64 = 128 << 20
)
func credential(candidate map[string]any, keys ...string) string {
for _, key := range keys {
if value, ok := candidate[key].(string); ok && strings.TrimSpace(value) != "" {
@@ -50,7 +55,12 @@ func decodeHTTPResponse(resp *http.Response) (map[string]any, error) {
func decodeHTTPResponseForProtocol(resp *http.Response, protocol string) (map[string]any, *WireResponse, error) {
defer resp.Body.Close()
raw, _ := io.ReadAll(io.LimitReader(resp.Body, 16*1024*1024))
maxBytes := maxJSONResponseBytesForProtocol(protocol)
raw, readErr := io.ReadAll(io.LimitReader(resp.Body, maxBytes+1))
tooLarge := int64(len(raw)) > maxBytes
if tooLarge {
raw = raw[:maxBytes]
}
wire := &WireResponse{
Protocol: strings.TrimSpace(protocol),
StatusCode: resp.StatusCode,
@@ -60,6 +70,26 @@ func decodeHTTPResponseForProtocol(resp *http.Response, protocol string) (map[st
if len(raw) > 0 {
_ = json.Unmarshal(raw, &wire.Body)
}
if readErr != nil {
return nil, wire, &ClientError{
Code: "response_read_error",
Message: readErr.Error(),
StatusCode: resp.StatusCode,
RequestID: requestIDFromHTTPResponse(resp),
Retryable: true,
Wire: wire,
}
}
if tooLarge {
return nil, wire, &ClientError{
Code: "response_too_large",
Message: fmt.Sprintf("upstream JSON response exceeds %d bytes", maxBytes),
StatusCode: resp.StatusCode,
RequestID: requestIDFromHTTPResponse(resp),
Retryable: false,
Wire: wire,
}
}
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
return nil, wire, &ClientError{
Code: statusCodeName(resp.StatusCode),
@@ -82,6 +112,15 @@ func decodeHTTPResponseForProtocol(resp *http.Response, protocol string) (map[st
return out, wire, nil
}
func maxJSONResponseBytesForProtocol(protocol string) int64 {
switch strings.TrimSpace(protocol) {
case ProtocolGeminiGenerateContent, ProtocolOpenAIImages:
return mediaMaxJSONResponseBytes
default:
return defaultMaxJSONResponseBytes
}
}
func compatibleResponseHeaders(source http.Header) map[string][]string {
if len(source) == 0 {
return nil
-2
View File
@@ -1306,7 +1306,6 @@ func kelingVideoSuccessResult(request Request, upstreamTaskID string, raw map[st
"status": "succeeded",
"upstream_task_id": upstreamTaskID,
"data": items,
"raw": raw,
}
}
@@ -1342,7 +1341,6 @@ func kelingTaskAPIVideoSuccessResult(request Request, upstreamTaskID string, tas
"status": "succeeded",
"upstream_task_id": upstreamTaskID,
"data": items,
"raw": raw,
}
}
+17 -17
View File
@@ -326,17 +326,17 @@ func (c MinimaxClient) runSpeech(ctx context.Context, request Request) (Response
if err != nil {
return Response{}, &ClientError{Code: "invalid_response", Message: "minimax speech audio hex is invalid: " + err.Error(), RequestID: firstNonEmptyString(requestID, requestIDFromResult(result)), ResponseStartedAt: startedAt, ResponseFinishedAt: finishedAt, ResponseDurationMS: responseDurationMS(startedAt, finishedAt), Retryable: false}
}
normalized := cloneMapAny(result)
normalized["status"] = "success"
normalized["created"] = time.Now().UnixMilli()
normalized["model"] = request.Model
normalized["raw_data"] = cloneMapAny(result)
normalized["data"] = []any{map[string]any{
"type": "audio",
"content": "data:audio/mpeg;base64," + base64.StdEncoding.EncodeToString(audioBytes),
"mime_type": "audio/mpeg",
"uploaded": false,
}}
normalized := map[string]any{
"status": "success",
"created": time.Now().UnixMilli(),
"model": request.Model,
"data": []any{map[string]any{
"type": "audio",
"content": "data:audio/mpeg;base64," + base64.StdEncoding.EncodeToString(audioBytes),
"mime_type": "audio/mpeg",
"uploaded": false,
}},
}
return Response{
Result: normalized,
RequestID: firstNonEmptyString(requestID, requestIDFromResult(result)),
@@ -378,12 +378,12 @@ func (c MinimaxClient) runVoiceClone(ctx context.Context, request Request) (Resp
if isProviderTaskFailure(providerTaskSpec{Name: "minimax"}, result) {
return Response{}, providerTaskFailure(providerTaskSpec{Name: "minimax"}, result, firstNonEmptyString(requestID, uploadRequestID, requestIDFromResult(result)), startedAt)
}
normalized := cloneMapAny(result)
normalized["status"] = "success"
normalized["created"] = time.Now().UnixMilli()
normalized["model"] = request.Model
normalized["voice_id"] = stringFromAny(payload["voice_id"])
normalized["raw_data"] = cloneMapAny(result)
normalized := map[string]any{
"status": "success",
"created": time.Now().UnixMilli(),
"model": request.Model,
"voice_id": stringFromAny(payload["voice_id"]),
}
if demoAudio := firstNonEmptyString(valueAtPath(result, "demo_audio"), valueAtPath(result, "data.demo_audio")); demoAudio != "" {
normalized["demo_audio"] = demoAudio
normalized["data"] = []any{map[string]any{"type": "audio", "url": demoAudio}}
@@ -0,0 +1,232 @@
package clients
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"testing"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
func TestGeminiGenerationConfigDropsDefaultRatioTokensAndNormalizesKSize(t *testing.T) {
config := geminiGenerationConfig(map[string]any{
"aspect_ratio": "auto",
"size": "2k",
"generationConfig": map[string]any{
"imageConfig": map[string]any{
"aspectRatio": "adaptive",
"imageSize": "2k",
},
},
}, true)
imageConfig := mapFromAny(config["imageConfig"])
if _, ok := imageConfig["aspectRatio"]; ok {
t.Fatalf("non-ratio tokens must not be sent to Gemini: %+v", config)
}
if imageConfig["imageSize"] != "2K" {
t.Fatalf("Gemini imageSize should use canonical K format: %+v", config)
}
}
func TestNormalizeOfficialOpenAIImageSizeUsesModelSpecificWireFormat(t *testing.T) {
legacy := map[string]any{
"model": "gpt-image-1",
"size": "2K",
"width": 2048,
"height": 1152,
"aspect_ratio": "16:9",
}
normalizeOfficialOpenAIImageSize(legacy, map[string]any{
"size": "2K",
"aspect_ratio": "16:9",
})
if legacy["size"] != "1536x1024" {
t.Fatalf("legacy GPT Image should use an official orientation size: %+v", legacy)
}
flexible := map[string]any{
"model": "gpt-image-2",
"size": "2K",
"resolution": "2k",
"aspect_ratio": "16:9",
}
normalizeOfficialOpenAIImageSize(flexible, map[string]any{
"size": "2K",
"aspect_ratio": "16:9",
})
if flexible["size"] != "2048x1152" {
t.Fatalf("GPT Image 2 should receive flexible pixel dimensions: %+v", flexible)
}
rounded := map[string]any{
"model": "gpt-image-2",
"size": "2K",
"resolution": "2K",
"aspect_ratio": "3:2",
}
normalizeOfficialOpenAIImageSize(rounded, map[string]any{
"size": "2K",
"aspect_ratio": "3:2",
})
if rounded["size"] != "2048x1360" {
t.Fatalf("GPT Image 2 dimensions should satisfy the official 16-pixel multiple: %+v", rounded)
}
explicit := map[string]any{
"model": "gpt-image-2",
"size": "1232x768",
"width": 1232,
"height": 768,
}
normalizeOfficialOpenAIImageSize(explicit, map[string]any{
"size": "1234x777",
})
if explicit["size"] != "1234x777" {
t.Fatalf("an explicit WxH size must be sent unchanged to OpenAI: %+v", explicit)
}
explicitLegacy := map[string]any{
"model": "gpt-image-1",
"size": "2048x1152",
}
normalizeOfficialOpenAIImageSize(explicitLegacy, map[string]any{
"size": "2048x1152",
})
if explicitLegacy["size"] != "2048x1152" {
t.Fatalf("explicit WxH must also win for legacy OpenAI image models: %+v", explicitLegacy)
}
}
func TestOpenAIClientUsesOriginalWxHAndConvertsKResolution(t *testing.T) {
received := make([]map[string]any, 0, 3)
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
var body map[string]any
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
t.Fatalf("decode OpenAI image request: %v", err)
}
received = append(received, body)
_ = json.NewEncoder(w).Encode(map[string]any{
"data": []any{map[string]any{"b64_json": "aW1hZ2U="}},
})
}))
defer server.Close()
candidate := store.RuntimeModelCandidate{
BaseURL: server.URL,
Provider: "openai",
ProviderModelName: "gpt-image-2",
ModelType: "image_generate",
Credentials: map[string]any{"apiKey": "openai-key"},
}
requests := []Request{
{
Kind: "images.generations",
ModelType: "image_generate",
Body: map[string]any{
"prompt": "explicit dimensions",
"size": "1232x768",
"width": 1232,
"height": 768,
},
OriginalBody: map[string]any{
"size": "1234x777",
},
Candidate: candidate,
},
{
Kind: "images.generations",
ModelType: "image_generate",
Body: map[string]any{
"prompt": "resolution conversion",
"size": "2048x1365",
"resolution": "2K",
"aspect_ratio": "3:2",
"width": 2048,
"height": 1365,
},
OriginalBody: map[string]any{
"size": "2K",
"aspect_ratio": "3:2",
},
Candidate: candidate,
},
{
Kind: "images.generations",
ModelType: "image_generate",
Body: map[string]any{
"prompt": "resolution field conversion",
"size": "1024x1024",
"resolution": "1K",
"aspect_ratio": "1:1",
"width": 1024,
"height": 1024,
},
OriginalBody: map[string]any{
"resolution": "1k",
"aspect_ratio": "1:1",
},
Candidate: candidate,
},
}
for _, request := range requests {
if _, err := (OpenAIClient{HTTPClient: server.Client()}).Run(context.Background(), request); err != nil {
t.Fatalf("run OpenAI image request: %v", err)
}
}
if len(received) != 3 {
t.Fatalf("unexpected request count: %d", len(received))
}
if received[0]["size"] != "1234x777" {
t.Fatalf("explicit WxH should win over preprocessed dimensions: %+v", received[0])
}
if received[1]["size"] != "2048x1360" {
t.Fatalf("2K with 3:2 should use normalized OpenAI dimensions: %+v", received[1])
}
if received[2]["size"] != "1024x1024" {
t.Fatalf("resolution=1K should be converted to OpenAI dimensions: %+v", received[2])
}
for _, body := range received {
for _, key := range []string{"resolution", "aspect_ratio", "width", "height"} {
if _, ok := body[key]; ok {
t.Fatalf("generic geometry field %q must not reach OpenAI: %+v", key, body)
}
}
}
}
func TestVolcesImageBodySelectsConfiguredSizeFormat(t *testing.T) {
body := volcesImageBody(Request{
Kind: "images.generations",
ModelType: "image_generate",
Body: map[string]any{
"model": "Seedream-5.0-Pro",
"resolution": "2k",
"size": "2048x1152",
"width": 2048,
"height": 1152,
"aspect_ratio": "16:9",
"platformId": "gateway-internal-platform",
},
Candidate: store.RuntimeModelCandidate{
ProviderModelName: "doubao-seedream-5-0-pro-260628",
ModelType: "image_generate",
Capabilities: map[string]any{
"image_generate": map[string]any{
"size_param_format": "resolution",
},
},
},
})
if body["size"] != "2K" {
t.Fatalf("resolution-only Volces model should receive canonical K size: %+v", body)
}
for _, key := range []string{"resolution", "width", "height", "aspect_ratio", "platformId"} {
if _, ok := body[key]; ok {
t.Fatalf("generic geometry field %q must not reach Volces: %+v", key, body)
}
}
}
@@ -0,0 +1,251 @@
package clients
import (
"math"
"strconv"
"strings"
)
const (
openAIGPTImage2MaxEdge = 3840
openAIGPTImage2MinPixels = 655360
openAIGPTImage2MaxPixels = 8294400
openAIGPTImage2MaxRatio = 3
openAIGPTImage2Multiple = 16
)
type providerImageDimensions struct {
width int
height int
}
func normalizedProviderAspectRatio(value any) string {
compact := strings.Map(func(r rune) rune {
switch r {
case ' ', '\t', '\r', '\n':
return -1
default:
return r
}
}, strings.TrimSpace(stringFromAny(value)))
parts := strings.Split(compact, ":")
if len(parts) != 2 {
return ""
}
width, widthErr := strconv.ParseFloat(parts[0], 64)
height, heightErr := strconv.ParseFloat(parts[1], 64)
if widthErr != nil || heightErr != nil || width <= 0 || height <= 0 {
return ""
}
return parts[0] + ":" + parts[1]
}
func normalizedProviderImageResolution(value any) string {
text := strings.ToLower(strings.TrimSpace(stringFromAny(value)))
if len(text) < 2 || text[len(text)-1] != 'k' {
return ""
}
magnitude, err := strconv.Atoi(text[:len(text)-1])
if err != nil || magnitude <= 0 {
return ""
}
return strconv.Itoa(magnitude) + "K"
}
func normalizeOfficialOpenAIImageSize(body map[string]any, originalBody map[string]any) {
if originalBody == nil {
originalBody = body
}
if width, height, ok := providerPixelDimensions(stringFromAny(originalBody["size"])); ok {
body["size"] = strconv.Itoa(width) + "x" + strconv.Itoa(height)
return
}
model := strings.ToLower(strings.TrimSpace(stringFromAny(body["model"])))
resolution := normalizedProviderImageResolution(firstNonEmptyString(
originalBody["resolution"],
originalBody["size"],
body["resolution"],
body["size"],
))
if resolution != "" {
width, height, ok := providerBodyPixelDimensions(body)
if !ok {
width, height, ok = providerResolutionDimensions(
resolution,
firstNonEmptyString(originalBody["aspect_ratio"], originalBody["aspectRatio"], originalBody["ratio"]),
)
}
if !ok {
delete(body, "size")
return
}
switch {
case strings.HasPrefix(model, "gpt-image-2"):
dimensions := normalizeOpenAIGPTImage2Dimensions(width, height)
body["size"] = strconv.Itoa(dimensions.width) + "x" + strconv.Itoa(dimensions.height)
case strings.HasPrefix(model, "gpt-image-"):
body["size"] = officialOpenAILegacyImageSize(width, height)
default:
body["size"] = strconv.Itoa(width) + "x" + strconv.Itoa(height)
}
return
}
if !strings.HasPrefix(model, "gpt-image-") {
return
}
size := strings.ToLower(strings.TrimSpace(stringFromAny(body["size"])))
if size == "auto" {
body["size"] = "auto"
return
}
if strings.HasPrefix(model, "gpt-image-2") {
if width, height, ok := providerPixelDimensions(size); ok {
body["size"] = strconv.Itoa(width) + "x" + strconv.Itoa(height)
}
return
}
switch size {
case "1024x1024", "1536x1024", "1024x1536":
body["size"] = size
return
}
width, height, ok := requestedProviderDimensions(body)
if !ok {
body["size"] = "auto"
return
}
body["size"] = officialOpenAILegacyImageSize(width, height)
}
func officialOpenAILegacyImageSize(width int, height int) string {
switch {
case width > height:
return "1536x1024"
case height > width:
return "1024x1536"
default:
return "1024x1024"
}
}
func providerBodyPixelDimensions(body map[string]any) (int, int, bool) {
if width, height, ok := providerPixelDimensions(stringFromAny(body["size"])); ok {
return width, height, true
}
width := intFromAny(body["width"])
height := intFromAny(body["height"])
return width, height, width > 0 && height > 0
}
func providerResolutionDimensions(resolution string, aspectRatio any) (int, int, bool) {
resolution = normalizedProviderImageResolution(resolution)
if resolution == "" {
return 0, 0, false
}
magnitude, err := strconv.Atoi(strings.TrimSuffix(resolution, "K"))
if err != nil || magnitude <= 0 {
return 0, 0, false
}
baseSide := magnitude * 1024
ratio := providerAspectRatioNumber(aspectRatio)
width, height := baseSide, baseSide
if ratio > 1 {
height = max(1, int(math.Round(float64(baseSide)/ratio)))
} else if ratio > 0 && ratio < 1 {
width = max(1, int(math.Round(float64(baseSide)*ratio)))
}
return width, height, true
}
func requestedProviderDimensions(body map[string]any) (int, int, bool) {
if width, height, ok := providerPixelDimensions(stringFromAny(body["size"])); ok {
return width, height, true
}
width := intFromAny(body["width"])
height := intFromAny(body["height"])
if width > 0 && height > 0 {
return width, height, true
}
ratio := providerAspectRatioNumber(firstNonEmptyString(body["aspect_ratio"], body["aspectRatio"], body["ratio"]))
if ratio > 1 {
return 2, 1, true
}
if ratio > 0 && ratio < 1 {
return 1, 2, true
}
if ratio == 1 {
return 1, 1, true
}
return 0, 0, false
}
func providerPixelDimensions(value string) (int, int, bool) {
parts := strings.Split(strings.ToLower(strings.TrimSpace(value)), "x")
if len(parts) != 2 {
return 0, 0, false
}
width, widthErr := strconv.Atoi(strings.TrimSpace(parts[0]))
height, heightErr := strconv.Atoi(strings.TrimSpace(parts[1]))
if widthErr != nil || heightErr != nil || width <= 0 || height <= 0 {
return 0, 0, false
}
return width, height, true
}
func providerAspectRatioNumber(value any) float64 {
normalized := normalizedProviderAspectRatio(value)
if normalized == "" {
return 0
}
parts := strings.Split(normalized, ":")
width, _ := strconv.ParseFloat(parts[0], 64)
height, _ := strconv.ParseFloat(parts[1], 64)
return width / height
}
func normalizeOpenAIGPTImage2Dimensions(width int, height int) providerImageDimensions {
if width <= 0 || height <= 0 {
return providerImageDimensions{}
}
rawRatio := float64(width) / float64(height)
ratio := rawRatio
if ratio < 1 {
ratio = 1 / ratio
}
ratio = min(ratio, float64(openAIGPTImage2MaxRatio))
landscape := rawRatio >= 1
longEdge := min(max(width, height), openAIGPTImage2MaxEdge)
dimensions := buildOpenAIGPTImage2Dimensions(longEdge, ratio, landscape, false)
for dimensions.width*dimensions.height > openAIGPTImage2MaxPixels {
longEdge = max(openAIGPTImage2Multiple, max(dimensions.width, dimensions.height)-openAIGPTImage2Multiple)
dimensions = buildOpenAIGPTImage2Dimensions(longEdge, ratio, landscape, false)
}
for dimensions.width*dimensions.height < openAIGPTImage2MinPixels && max(dimensions.width, dimensions.height) < openAIGPTImage2MaxEdge {
longEdge = min(openAIGPTImage2MaxEdge, max(dimensions.width, dimensions.height)+openAIGPTImage2Multiple)
dimensions = buildOpenAIGPTImage2Dimensions(longEdge, ratio, landscape, true)
}
return dimensions
}
func buildOpenAIGPTImage2Dimensions(longEdge int, ratio float64, landscape bool, roundUp bool) providerImageDimensions {
normalizedLongEdge := roundOpenAIGPTImage2Multiple(float64(longEdge), roundUp)
normalizedShortEdge := roundOpenAIGPTImage2Multiple(float64(normalizedLongEdge)/ratio, roundUp)
if float64(normalizedLongEdge)/float64(normalizedShortEdge) > openAIGPTImage2MaxRatio {
normalizedShortEdge += openAIGPTImage2Multiple
}
if landscape {
return providerImageDimensions{width: normalizedLongEdge, height: normalizedShortEdge}
}
return providerImageDimensions{width: normalizedShortEdge, height: normalizedLongEdge}
}
func roundOpenAIGPTImage2Multiple(value float64, roundUp bool) int {
rounded := math.Round(value / openAIGPTImage2Multiple)
if roundUp {
rounded = math.Ceil(value / openAIGPTImage2Multiple)
}
return max(openAIGPTImage2Multiple, int(rounded)*openAIGPTImage2Multiple)
}
+256 -2
View File
@@ -3,8 +3,14 @@ package clients
import (
"bytes"
"context"
"encoding/base64"
"encoding/json"
"fmt"
"io"
"mime/multipart"
"net/http"
"net/textproto"
"net/url"
"strings"
"time"
@@ -42,6 +48,11 @@ func (c OpenAIClient) Run(ctx context.Context, request Request) (Response, error
}
if endpointKind == "chat.completions" {
body = NormalizeChatCompletionRequestBody(body)
normalizedBody, normalizeErr := normalizeOpenAIChatAudioParts(body, request.Candidate)
if normalizeErr != nil {
return Response{}, normalizeErr
}
body = normalizedBody
applyOpenAIChatReasoningParams(body, request.Candidate)
body = FilterOpenAIChatRequestBody(body)
} else if request.Kind == "responses" {
@@ -59,15 +70,19 @@ func (c OpenAIClient) Run(ctx context.Context, request Request) (Response, error
}
}
body["model"] = upstreamModelName(request.Candidate)
normalizeOpenAIImageRequestBody(endpointKind, body, request.OriginalBody)
stream := openAIEndpointSupportsStream(endpointKind) && (request.Stream || boolValue(body, "stream"))
ensureOpenAIStreamUsage(body, endpointKind, stream)
raw, _ := json.Marshal(body)
raw, contentType, err := openAIRequestPayload(endpointKind, body, request.Candidate)
if err != nil {
return Response{}, err
}
upstreamEndpoint := joinURL(openAIBaseURL(endpointKind, request.Candidate), endpoint)
req, err := http.NewRequestWithContext(ctx, http.MethodPost, upstreamEndpoint, bytes.NewReader(raw))
if err != nil {
return Response{}, err
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Content-Type", contentType)
req.Header.Set("Authorization", "Bearer "+apiKey)
responseStartedAt := time.Now()
if err := notifySubmissionStarted(request); err != nil {
@@ -162,6 +177,245 @@ func (c OpenAIClient) Run(ctx context.Context, request Request) (Response, error
}, nil
}
func normalizeOpenAIImageRequestBody(endpointKind string, body map[string]any, originalBody map[string]any) {
if endpointKind != "images.generations" && endpointKind != "images.edits" {
return
}
normalizeOfficialOpenAIImageSize(body, originalBody)
// OpenAI image endpoints express output geometry through size. The Gateway
// keeps the generic fields for capability validation and billing, but they
// are not valid OpenAI wire parameters.
for _, key := range []string{
"aspect_ratio",
"aspectRatio",
"resolution",
"width",
"height",
"platform_id",
"platformId",
"platform_model_id",
"platformModelId",
} {
delete(body, key)
}
}
func openAIRequestPayload(endpointKind string, body map[string]any, candidate store.RuntimeModelCandidate) ([]byte, string, error) {
if endpointKind != "images.edits" {
raw, err := json.Marshal(body)
return raw, "application/json", err
}
var payload bytes.Buffer
writer := multipart.NewWriter(&payload)
imageFieldName := openAIImageEditFieldName(candidate)
images := openAIImageEditValues(firstPresent(body["images"], body["image"]))
if len(images) == 0 {
return nil, "", &ClientError{
Code: "invalid_parameter",
Message: "image is required",
Param: "image",
StatusCode: http.StatusBadRequest,
Retryable: false,
}
}
for index, value := range images {
contentType, image, err := openAIImageEditPayload(value)
if err != nil {
return nil, "", err
}
if err := writeOpenAIImageEditFile(writer, imageFieldName, fmt.Sprintf("image-%d%s", index+1, openAIImageFileExtension(contentType)), contentType, image); err != nil {
return nil, "", err
}
}
if mask := firstPresent(body["mask"], body["mask_image"], body["maskImage"]); mask != nil {
contentType, image, err := openAIImageEditPayload(mask)
if err != nil {
return nil, "", err
}
if err := writeOpenAIImageEditFile(writer, "mask", "mask"+openAIImageFileExtension(contentType), contentType, image); err != nil {
return nil, "", err
}
}
for key, value := range body {
switch key {
case "image", "images", "mask", "mask_image", "maskImage":
continue
}
if strings.HasPrefix(key, "_") || value == nil {
continue
}
fieldValue, err := openAIFormFieldValue(value)
if err != nil {
return nil, "", err
}
if fieldValue == "" {
continue
}
if err := writer.WriteField(key, fieldValue); err != nil {
return nil, "", err
}
}
if err := writer.Close(); err != nil {
return nil, "", err
}
return payload.Bytes(), writer.FormDataContentType(), nil
}
func openAIImageEditFieldName(candidate store.RuntimeModelCandidate) string {
if configured := strings.TrimSpace(stringFromAny(candidate.PlatformConfig["imageEditMultipartFieldName"])); configured == "image" || configured == "image[]" {
return configured
}
parsed, err := url.Parse(strings.TrimSpace(candidate.BaseURL))
if err == nil && strings.EqualFold(parsed.Hostname(), "api.openai.com") {
return "image[]"
}
return "image"
}
func openAIImageEditValues(value any) []any {
switch typed := value.(type) {
case []any:
return typed
case []string:
out := make([]any, 0, len(typed))
for _, item := range typed {
out = append(out, item)
}
return out
case nil:
return nil
default:
return []any{typed}
}
}
func openAIImageEditPayload(value any) (string, []byte, error) {
switch typed := value.(type) {
case map[string]any:
for _, key := range []string{"data", "b64_json", "base64", "url"} {
if nested := typed[key]; nested != nil {
return openAIImageEditPayload(nested)
}
}
case string:
raw := strings.TrimSpace(typed)
if raw == "" {
break
}
contentType := ""
encoded := raw
if strings.HasPrefix(strings.ToLower(raw), "data:") {
prefix, payload, ok := strings.Cut(raw, ",")
if !ok || !strings.Contains(strings.ToLower(prefix), ";base64") {
break
}
contentType = strings.TrimSpace(strings.Split(strings.TrimPrefix(prefix, "data:"), ";")[0])
encoded = payload
} else if strings.Contains(raw, "://") {
return "", nil, &ClientError{
Code: "invalid_parameter",
Message: "OpenAI image edit input must be hydrated before multipart submission",
Param: "image",
StatusCode: http.StatusBadRequest,
Retryable: false,
}
}
image, err := decodeOpenAIImageEditBase64(encoded)
if err != nil {
break
}
if contentType == "" {
contentType = strings.TrimSpace(strings.Split(http.DetectContentType(image), ";")[0])
}
return contentType, image, nil
case []byte:
if len(typed) > 0 {
return strings.TrimSpace(strings.Split(http.DetectContentType(typed), ";")[0]), typed, nil
}
}
return "", nil, &ClientError{
Code: "invalid_parameter",
Message: "image must be a base64 or data URL payload",
Param: "image",
StatusCode: http.StatusBadRequest,
Retryable: false,
}
}
func decodeOpenAIImageEditBase64(value string) ([]byte, error) {
normalized := strings.Map(func(char rune) rune {
switch char {
case '\n', '\r', '\t', ' ':
return -1
default:
return char
}
}, value)
var lastErr error
for _, encoding := range []*base64.Encoding{
base64.StdEncoding,
base64.RawStdEncoding,
base64.URLEncoding,
base64.RawURLEncoding,
} {
payload, err := encoding.DecodeString(normalized)
if err == nil && len(payload) > 0 {
return payload, nil
}
lastErr = err
}
return nil, lastErr
}
func writeOpenAIImageEditFile(writer *multipart.Writer, fieldName string, fileName string, contentType string, payload []byte) error {
header := make(textproto.MIMEHeader)
header.Set("Content-Disposition", fmt.Sprintf(`form-data; name="%s"; filename="%s"`, fieldName, escapeMultipartFilename(fileName)))
if contentType != "" {
header.Set("Content-Type", contentType)
}
part, err := writer.CreatePart(header)
if err != nil {
return err
}
_, err = io.Copy(part, bytes.NewReader(payload))
return err
}
func openAIImageFileExtension(contentType string) string {
switch strings.ToLower(strings.TrimSpace(strings.Split(contentType, ";")[0])) {
case "image/jpeg":
return ".jpg"
case "image/webp":
return ".webp"
case "image/gif":
return ".gif"
default:
return ".png"
}
}
func openAIFormFieldValue(value any) (string, error) {
switch typed := value.(type) {
case string:
return typed, nil
case bool:
return fmt.Sprintf("%t", typed), nil
case float64:
return fmt.Sprintf("%v", typed), nil
case float32:
return fmt.Sprintf("%v", typed), nil
case int:
return fmt.Sprintf("%d", typed), nil
case int64:
return fmt.Sprintf("%d", typed), nil
case json.Number:
return typed.String(), nil
default:
encoded, err := json.Marshal(value)
return string(encoded), err
}
}
func openAIWireStreamDelta(next StreamDelta, protocol string, response *http.Response) StreamDelta {
if next == nil {
return nil
@@ -0,0 +1,169 @@
package clients
import (
"net/http"
"net/url"
"strings"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
var openAIChatAudioFormatByMIME = map[string]string{
"audio/aac": "aac",
"audio/flac": "flac",
"audio/m4a": "m4a",
"audio/mp4": "m4a",
"audio/mpeg": "mp3",
"audio/ogg": "ogg",
"audio/opus": "opus",
"audio/wav": "wav",
"audio/webm": "webm",
"audio/x-m4a": "m4a",
"audio/x-wav": "wav",
}
var openAIChatAudioFormatByExtension = map[string]string{
"aac": "aac",
"flac": "flac",
"m4a": "m4a",
"mp3": "mp3",
"mp4": "m4a",
"oga": "ogg",
"ogg": "ogg",
"opus": "opus",
"wav": "wav",
"webm": "webm",
}
func normalizeOpenAIChatAudioParts(body map[string]any, candidate store.RuntimeModelCandidate) (map[string]any, error) {
if !isAliyunBailianOpenAI(candidate) {
return body, nil
}
messages, ok := body["messages"].([]any)
if !ok {
return body, nil
}
normalizedMessages := make([]any, 0, len(messages))
for _, rawMessage := range messages {
message, ok := rawMessage.(map[string]any)
if !ok {
normalizedMessages = append(normalizedMessages, rawMessage)
continue
}
content, ok := message["content"].([]any)
if !ok {
normalizedMessages = append(normalizedMessages, rawMessage)
continue
}
normalizedContent := make([]any, 0, len(content))
for _, rawPart := range content {
part, ok := rawPart.(map[string]any)
if !ok {
normalizedContent = append(normalizedContent, rawPart)
continue
}
normalizedPart, err := normalizeAliyunOpenAIChatAudioPart(part)
if err != nil {
return nil, err
}
normalizedContent = append(normalizedContent, normalizedPart)
}
copiedMessage := cloneMapAny(message)
copiedMessage["content"] = normalizedContent
normalizedMessages = append(normalizedMessages, copiedMessage)
}
out := cloneMapAny(body)
out["messages"] = normalizedMessages
return out, nil
}
func normalizeAliyunOpenAIChatAudioPart(part map[string]any) (map[string]any, error) {
switch strings.TrimSpace(stringFromAny(part["type"])) {
case "audio_url":
audioURL, _ := part["audio_url"].(map[string]any)
data := strings.TrimSpace(stringFromAny(audioURL["url"]))
format := resolveOpenAIChatAudioFormat(data, audioURL["format"], audioURL["mime_type"], audioURL["mimeType"])
if data == "" || format == "" {
return nil, invalidOpenAIChatAudioError()
}
return map[string]any{
"type": "input_audio",
"input_audio": map[string]any{
"data": data,
"format": format,
},
}, nil
case "input_audio":
inputAudio, _ := part["input_audio"].(map[string]any)
data := firstNonEmptyString(inputAudio["data"], inputAudio["url"])
format := resolveOpenAIChatAudioFormat(data, inputAudio["format"], inputAudio["mime_type"], inputAudio["mimeType"])
if data == "" || format == "" {
return nil, invalidOpenAIChatAudioError()
}
copiedInput := cloneMapAny(inputAudio)
copiedInput["data"] = data
copiedInput["format"] = format
delete(copiedInput, "url")
copiedPart := cloneMapAny(part)
copiedPart["input_audio"] = copiedInput
return copiedPart, nil
default:
return part, nil
}
}
func resolveOpenAIChatAudioFormat(source string, values ...any) string {
if explicit := strings.ToLower(strings.TrimSpace(firstNonEmptyString(values...))); explicit != "" {
if normalized := openAIChatAudioFormatByExtension[explicit]; normalized != "" {
return normalized
}
if normalized := openAIChatAudioFormatByMIME[strings.TrimSpace(strings.Split(explicit, ";")[0])]; normalized != "" {
return normalized
}
return explicit
}
if dataMIME := openAIChatAudioDataMIME(source); dataMIME != "" {
if normalized := openAIChatAudioFormatByMIME[dataMIME]; normalized != "" {
return normalized
}
}
if extension := openAIChatAudioSourceExtension(source); extension != "" {
return openAIChatAudioFormatByExtension[extension]
}
return ""
}
func openAIChatAudioDataMIME(source string) string {
if !strings.HasPrefix(strings.ToLower(source), "data:") {
return ""
}
header := strings.TrimPrefix(source[:strings.Index(source+",", ",")], "data:")
return strings.ToLower(strings.TrimSpace(strings.Split(header, ";")[0]))
}
func openAIChatAudioSourceExtension(source string) string {
if source == "" || strings.HasPrefix(strings.ToLower(source), "data:") {
return ""
}
if parsed, err := url.Parse(source); err == nil && parsed.Path != "" {
source = parsed.Path
}
source = strings.TrimSuffix(strings.Split(strings.Split(source, "?")[0], "#")[0], "/")
if index := strings.LastIndex(source, "."); index >= 0 && index+1 < len(source) {
return strings.ToLower(source[index+1:])
}
return ""
}
func invalidOpenAIChatAudioError() error {
return &ClientError{
Code: "invalid_parameter",
Message: "input_audio requires data and a resolvable format",
Param: "messages.content.input_audio",
StatusCode: http.StatusBadRequest,
Retryable: false,
}
}
+11 -15
View File
@@ -340,24 +340,20 @@ func hasProviderTaskResult(result map[string]any) bool {
return result["data"] != nil || valueAtPath(result, "data.result") != nil || valueAtPath(result, "data.audio") != nil || valueAtPath(result, "output.image_urls") != nil || valueAtPath(result, "output.video_url") != nil || valueAtPath(result, "Response.ResultVideoUrl") != nil || valueAtPath(result, "Response.ResultImages") != nil || result["audio_url"] != nil || result["urls"] != nil
}
func normalizeProviderTaskResult(request Request, spec providerTaskSpec, result map[string]any, upstreamTaskID string) map[string]any {
out := cloneMapAny(result)
out["status"] = "success"
func normalizeProviderTaskResult(request Request, _ providerTaskSpec, result map[string]any, upstreamTaskID string) map[string]any {
data, ok := result["data"].([]any)
if !ok {
data = providerTaskData(request, result)
}
out := map[string]any{
"status": "success",
"created": time.Now().UnixMilli(),
"model": request.Model,
"data": data,
}
if upstreamTaskID != "" {
out["upstream_task_id"] = upstreamTaskID
}
if out["created"] == nil {
out["created"] = time.Now().UnixMilli()
}
if out["model"] == nil {
out["model"] = request.Model
}
if _, ok := out["data"].([]any); !ok {
if out["data"] != nil {
out["raw_data"] = out["data"]
}
out["data"] = providerTaskData(request, result)
}
return out
}
+1
View File
@@ -15,6 +15,7 @@ type Request struct {
ModelType string
Model string
Body map[string]any
OriginalBody map[string]any
Candidate store.RuntimeModelCandidate
HTTPClient *http.Client
RemoteTaskID string
+35 -12
View File
@@ -357,16 +357,33 @@ func volcesImageBody(request Request) map[string]any {
body["seed"] = -1
}
}
if resolution := strings.TrimSpace(stringFromAny(body["resolution"])); resolution != "" {
resolution := normalizedProviderImageResolution(body["resolution"])
size := widthHeightSize(body)
if volcesImageUsesResolutionSize(request) && resolution != "" {
body["size"] = resolution
} else if size != "" {
body["size"] = size
} else if resolution != "" {
body["size"] = resolution
}
if size := widthHeightSize(body); size != "" {
body["size"] = size
for _, key := range []string{"aspect_ratio", "aspectRatio", "ratio", "resolution", "width", "height"} {
delete(body, key)
}
normalizeVolcesSequentialImageGeneration(body, request)
return body
}
func volcesImageUsesResolutionSize(request Request) bool {
modelType := firstNonEmpty(request.ModelType, request.Candidate.ModelType)
if capability, ok := request.Candidate.Capabilities[modelType].(map[string]any); ok {
if strings.EqualFold(strings.TrimSpace(stringFromAny(capability["size_param_format"])), "resolution") {
return true
}
}
model := strings.ToLower(strings.TrimSpace(upstreamModelName(request.Candidate)))
return strings.HasPrefix(model, "doubao-seedream-5-0")
}
func volcesVideoBody(request Request) map[string]any {
body := cleanProviderBody(request.Body)
body["model"] = upstreamModelName(request.Candidate)
@@ -397,6 +414,10 @@ func cleanProviderBody(body map[string]any) map[string]any {
"_gateway_target_protocol",
"_compat_provider",
"_kling_compat_version",
"platform_id",
"platformId",
"platform_model_id",
"platformModelId",
} {
delete(out, key)
}
@@ -1100,9 +1121,11 @@ func volcesTaskErrorMessage(result map[string]any) string {
}
func volcesVideoSuccessResult(request Request, upstreamTaskID string, raw map[string]any) map[string]any {
result := cloneMapAny(raw)
if result == nil {
result = map[string]any{}
result := map[string]any{}
for _, key := range []string{"seed", "resolution", "ratio", "duration", "frames", "framespersecond"} {
if raw[key] != nil {
result[key] = raw[key]
}
}
content, _ := raw["content"].(map[string]any)
videoURL := strings.TrimSpace(stringFromAny(content["video_url"]))
@@ -1112,18 +1135,18 @@ func volcesVideoSuccessResult(request Request, upstreamTaskID string, raw map[st
}
data := []any{}
if videoURL != "" {
data = append(data, map[string]any{"url": videoURL, "type": "video"})
item := map[string]any{"url": videoURL, "type": "video"}
if lastFrameURL := strings.TrimSpace(stringFromAny(content["last_frame_url"])); lastFrameURL != "" {
item["last_frame_url"] = lastFrameURL
}
data = append(data, item)
}
result["id"] = firstNonEmpty(stringFromAny(raw["id"]), upstreamTaskID)
if strings.TrimSpace(stringFromAny(result["model"])) == "" {
result["model"] = upstreamModelName(request.Candidate)
}
result["model"] = firstNonEmpty(stringFromAny(raw["model"]), upstreamModelName(request.Candidate))
result["status"] = "succeeded"
result["object"] = "video.generation"
result["created"] = created
result["upstream_task_id"] = upstreamTaskID
result["data"] = data
result["raw"] = cloneMapAny(raw)
return result
}
@@ -79,7 +79,7 @@ func TestVolcesClientSupportsDeyunEnvelope(t *testing.T) {
}
data, _ := response.Result["data"].([]any)
item, _ := data[0].(map[string]any)
if response.Result["upstream_task_id"] != "deyun-task-1" || item["url"] != "https://example.com/deyun.mp4" {
if response.Result["upstream_task_id"] != nil || item["url"] != "https://example.com/deyun.mp4" {
t.Fatalf("unexpected response: %+v", response.Result)
}
}
@@ -46,6 +46,46 @@ func TestDecodeHTTPResponseForProtocolCapturesOfficialErrorWire(t *testing.T) {
}
}
func TestDecodeHTTPResponseForProtocolAllowsLargeImagePayload(t *testing.T) {
payload := `{"data":"` + strings.Repeat("a", int(defaultMaxJSONResponseBytes)) + `"}`
response := &http.Response{
StatusCode: http.StatusOK,
Status: "200 OK",
Header: http.Header{"Content-Type": {"application/json"}},
Body: io.NopCloser(strings.NewReader(payload)),
}
result, wire, err := decodeHTTPResponseForProtocol(response, ProtocolGeminiGenerateContent)
if err != nil {
t.Fatalf("large Gemini image response failed: %v", err)
}
if got, _ := result["data"].(string); len(got) != int(defaultMaxJSONResponseBytes) {
t.Fatalf("large Gemini image payload length = %d", len(got))
}
if wire == nil || len(wire.RawJSON) != len(payload) {
t.Fatalf("large Gemini wire payload was truncated: %+v", wire)
}
}
func TestDecodeHTTPResponseForProtocolRejectsOversizedDefaultPayload(t *testing.T) {
payload := `{"data":"` + strings.Repeat("a", int(defaultMaxJSONResponseBytes)) + `"}`
response := &http.Response{
StatusCode: http.StatusOK,
Status: "200 OK",
Header: http.Header{"Content-Type": {"application/json"}},
Body: io.NopCloser(strings.NewReader(payload)),
}
_, _, err := decodeHTTPResponseForProtocol(response, ProtocolOpenAIResponses)
if err == nil {
t.Fatal("expected oversized default response to fail")
}
clientErr, ok := err.(*ClientError)
if !ok || clientErr.Code != "response_too_large" {
t.Fatalf("unexpected oversized response error: %T %v", err, err)
}
}
func TestGeminiNativeStreamPreservesOfficialEvents(t *testing.T) {
var requestedPath string
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+93 -9
View File
@@ -39,15 +39,27 @@ type Config struct {
LocalGeneratedStorageDir string
LocalUploadedStorageDir string
LocalTempAssetTTLHours int
LocalResultTTLHours int
LocalResultMinFreeBytes int64
LocalResultMaxBytes int64
LocalResultMaxTaskBytes int64
TaskProgressCallbackEnabled bool
TaskProgressCallbackURL string
TaskProgressCallbackTimeoutMS string
TaskProgressCallbackMaxAttempts string
TaskProgressCallbackTimeoutMS int
TaskProgressCallbackMaxAttempts int
TaskCleanupEnabled bool
TaskRetentionDays int
TaskAnalysisRetentionDays int
TaskCleanupIntervalSeconds int
TaskCleanupBatchSize int
CORSAllowedOrigin string
GlobalHTTPProxy string
GlobalHTTPProxySource string
LogLevel slog.Level
BillingEngineMode string
AsyncQueueWorkerEnabled bool
AsyncWorkerHardLimit int
AsyncWorkerRefreshIntervalSeconds int
}
func Load() Config {
@@ -81,17 +93,29 @@ func Load() Config {
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)),
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",
TaskProgressCallbackURL: env("TASK_PROGRESS_CALLBACK_URL",
strings.TrimRight(env("SERVER_MAIN_BASE_URL", "http://localhost:3000"), "/")+"/internal/platform/task-progress-callbacks",
),
TaskProgressCallbackTimeoutMS: env("TASK_PROGRESS_CALLBACK_TIMEOUT_MS", "5000"),
TaskProgressCallbackMaxAttempts: env("TASK_PROGRESS_CALLBACK_MAX_ATTEMPTS", "10"),
CORSAllowedOrigin: env("CORS_ALLOWED_ORIGIN", "http://localhost:5178,http://127.0.0.1:5178"),
GlobalHTTPProxy: globalProxy.HTTPProxy,
GlobalHTTPProxySource: globalProxy.Source,
LogLevel: logLevel(env("LOG_LEVEL", "info")),
BillingEngineMode: strings.ToLower(env("BILLING_ENGINE_MODE", "observe")),
TaskProgressCallbackTimeoutMS: envIntValidated("TASK_PROGRESS_CALLBACK_TIMEOUT_MS", 5000),
TaskProgressCallbackMaxAttempts: envIntValidated("TASK_PROGRESS_CALLBACK_MAX_ATTEMPTS", 10),
TaskCleanupEnabled: env("AI_GATEWAY_TASK_CLEANUP_ENABLED", "false") == "true",
TaskRetentionDays: envIntValidated("AI_GATEWAY_TASK_RETENTION_DAYS", 30),
TaskAnalysisRetentionDays: envIntValidated("AI_GATEWAY_TASK_ANALYSIS_RETENTION_DAYS", 7),
TaskCleanupIntervalSeconds: envIntValidated("AI_GATEWAY_TASK_CLEANUP_INTERVAL_SECONDS", 300),
TaskCleanupBatchSize: envIntValidated("AI_GATEWAY_TASK_CLEANUP_BATCH_SIZE", 1000),
CORSAllowedOrigin: env("CORS_ALLOWED_ORIGIN", "http://localhost:5178,http://127.0.0.1:5178"),
GlobalHTTPProxy: globalProxy.HTTPProxy,
GlobalHTTPProxySource: globalProxy.Source,
LogLevel: logLevel(env("LOG_LEVEL", "info")),
BillingEngineMode: strings.ToLower(env("BILLING_ENGINE_MODE", "observe")),
AsyncQueueWorkerEnabled: env("AI_GATEWAY_ASYNC_QUEUE_WORKER_ENABLED", "true") == "true",
AsyncWorkerHardLimit: envIntValidated("AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT", 2048),
AsyncWorkerRefreshIntervalSeconds: envIntValidated("AI_GATEWAY_ASYNC_WORKER_REFRESH_INTERVAL_SECONDS", 5),
}
}
@@ -101,6 +125,42 @@ func (c Config) Validate() error {
default:
return errors.New("BILLING_ENGINE_MODE must be observe, enforce, or hold")
}
if c.AsyncWorkerHardLimit < 1 || c.AsyncWorkerHardLimit > 10000 {
return errors.New("AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT must be between 1 and 10000")
}
if c.AsyncWorkerRefreshIntervalSeconds < 1 {
return errors.New("AI_GATEWAY_ASYNC_WORKER_REFRESH_INTERVAL_SECONDS must be positive")
}
if c.TaskProgressCallbackTimeoutMS != 0 && (c.TaskProgressCallbackTimeoutMS < 100 || c.TaskProgressCallbackTimeoutMS > 60000) {
return errors.New("TASK_PROGRESS_CALLBACK_TIMEOUT_MS must be between 100 and 60000")
}
if c.TaskProgressCallbackMaxAttempts != 0 && (c.TaskProgressCallbackMaxAttempts < 1 || c.TaskProgressCallbackMaxAttempts > 100) {
return errors.New("TASK_PROGRESS_CALLBACK_MAX_ATTEMPTS must be between 1 and 100")
}
if c.TaskRetentionDays != 0 && (c.TaskRetentionDays < 1 || c.TaskRetentionDays > 3650) {
return errors.New("AI_GATEWAY_TASK_RETENTION_DAYS must be between 1 and 3650")
}
if c.TaskAnalysisRetentionDays != 0 && (c.TaskAnalysisRetentionDays < 1 || (c.TaskRetentionDays != 0 && c.TaskAnalysisRetentionDays > c.TaskRetentionDays)) {
return errors.New("AI_GATEWAY_TASK_ANALYSIS_RETENTION_DAYS must be between 1 and AI_GATEWAY_TASK_RETENTION_DAYS")
}
if c.TaskCleanupIntervalSeconds != 0 && c.TaskCleanupIntervalSeconds < 60 {
return errors.New("AI_GATEWAY_TASK_CLEANUP_INTERVAL_SECONDS must be at least 60")
}
if c.TaskCleanupBatchSize != 0 && (c.TaskCleanupBatchSize < 100 || c.TaskCleanupBatchSize > 5000) {
return errors.New("AI_GATEWAY_TASK_CLEANUP_BATCH_SIZE must be between 100 and 5000")
}
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)) {
case "":
case "file":
@@ -210,6 +270,30 @@ func envInt(key string, fallback int) int {
return parsed
}
func envIntValidated(key string, fallback int) int {
value := envValue(key)
if value == "" {
return fallback
}
parsed, err := strconv.Atoi(value)
if err != nil {
return 0
}
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 {
switch strings.ToLower(value) {
case "debug":
+57 -2
View File
@@ -26,7 +26,7 @@ func TestLoadIdentitySecretStoreUsesNewEnvironmentNamesAndIgnoresLegacyBusinessV
}
func TestValidateIdentityFileSecretStoreRequiresDirectory(t *testing.T) {
cfg := Config{IdentitySecretStore: "file"}
cfg := Config{IdentitySecretStore: "file", AsyncWorkerHardLimit: 2048, AsyncWorkerRefreshIntervalSeconds: 5}
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "IDENTITY_SECRET_DIR") {
t.Fatalf("Validate() error = %v, want missing identity secret directory", err)
}
@@ -37,7 +37,12 @@ func TestValidateIdentityFileSecretStoreRequiresDirectory(t *testing.T) {
}
func TestValidateIdentityKubernetesSecretStore(t *testing.T) {
cfg := Config{IdentitySecretStore: "kubernetes", IdentityKubernetesSecretName: "easyai-gateway-identity"}
cfg := Config{
IdentitySecretStore: "kubernetes",
IdentityKubernetesSecretName: "easyai-gateway-identity",
AsyncWorkerHardLimit: 2048,
AsyncWorkerRefreshIntervalSeconds: 5,
}
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "namespace") {
t.Fatalf("Validate() error = %v, want missing namespace", err)
}
@@ -52,8 +57,58 @@ func TestValidateIdentitySecurityEventTiming(t *testing.T) {
IdentitySecurityEventHeartbeatIntervalSeconds: 60,
IdentitySecurityEventStaleAfterSeconds: 60,
IdentitySecurityEventClockSkewSeconds: 60,
AsyncWorkerHardLimit: 2048,
AsyncWorkerRefreshIntervalSeconds: 5,
}
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "heartbeat") {
t.Fatalf("Validate() error = %v, want invalid stale threshold", err)
}
}
func TestValidateAsyncWorkerSettings(t *testing.T) {
cfg := Config{AsyncWorkerHardLimit: 10001, AsyncWorkerRefreshIntervalSeconds: 5}
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "HARD_LIMIT") {
t.Fatalf("Validate() error = %v, want invalid hard limit", err)
}
cfg.AsyncWorkerHardLimit = 2048
cfg.AsyncWorkerRefreshIntervalSeconds = 0
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "REFRESH_INTERVAL") {
t.Fatalf("Validate() error = %v, want invalid refresh interval", err)
}
t.Setenv("AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT", "not-an-integer")
loaded := Load()
if err := loaded.Validate(); err == nil || !strings.Contains(err.Error(), "HARD_LIMIT") {
t.Fatalf("Validate() error = %v, want invalid non-integer hard limit", err)
}
}
func TestLoadAsyncQueueWorkerEnabled(t *testing.T) {
cfg := Load()
if !cfg.AsyncQueueWorkerEnabled {
t.Fatal("async queue worker must be enabled by default")
}
t.Setenv("AI_GATEWAY_ASYNC_QUEUE_WORKER_ENABLED", "false")
cfg = Load()
if cfg.AsyncQueueWorkerEnabled {
t.Fatal("async queue worker remained enabled after explicit disable")
}
}
func TestValidateTaskHistorySettings(t *testing.T) {
cfg := Load()
cfg.TaskRetentionDays = 30
cfg.TaskAnalysisRetentionDays = 31
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "ANALYSIS_RETENTION") {
t.Fatalf("Validate() error = %v, want invalid analysis retention", err)
}
cfg.TaskAnalysisRetentionDays = 7
cfg.TaskCleanupIntervalSeconds = 59
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "CLEANUP_INTERVAL") {
t.Fatalf("Validate() error = %v, want invalid cleanup interval", err)
}
cfg.TaskCleanupIntervalSeconds = 300
cfg.TaskCleanupBatchSize = 5001
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "CLEANUP_BATCH") {
t.Fatalf("Validate() error = %v, want invalid cleanup batch", err)
}
}
@@ -0,0 +1,213 @@
package httpapi
import (
"net/http"
"strings"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
"github.com/google/uuid"
)
// listAdminTasks godoc
// @Summary 管理员列出任务
// @Description 跨租户分页查询任务;请求、结果和执行快照中的常见敏感字段默认脱敏。
// @Tags admin-tasks
// @Produce json
// @Security BearerAuth
// @Param q query string false "关键词,匹配任务、用户、租户、模型、平台和 API Key"
// @Param tenantId query string false "网关租户 ID"
// @Param userId query string false "网关用户 ID"
// @Param userGroupId query string false "用户组 ID"
// @Param status query string false "任务状态"
// @Param platformId query string false "执行平台 ID"
// @Param model query string false "调用或实际模型名称"
// @Param modelType query string false "模型类型"
// @Param runMode query string false "运行模式"
// @Param billingStatus query string false "计费状态"
// @Param apiKey query string false "API Key ID、名称或前缀"
// @Param createdFrom query string false "创建时间起点,支持 RFC3339 或日期格式"
// @Param createdTo query string false "创建时间终点,支持 RFC3339 或日期格式"
// @Param page query int false "页码" default(1)
// @Param pageSize query int false "每页数量" default(50)
// @Success 200 {object} AdminTaskListResponse
// @Failure 400 {object} ErrorEnvelope
// @Failure 401 {object} ErrorEnvelope
// @Failure 403 {object} ErrorEnvelope
// @Failure 500 {object} ErrorEnvelope
// @Router /api/admin/tasks [get]
func (s *Server) listAdminTasks(w http.ResponseWriter, r *http.Request) {
query := r.URL.Query()
page, err := positiveQueryInt(query.Get("page"), 1)
if err != nil {
writeError(w, http.StatusBadRequest, "invalid page")
return
}
pageSize, err := positiveQueryInt(query.Get("pageSize"), 50)
if err != nil {
writeError(w, http.StatusBadRequest, "invalid pageSize")
return
}
createdFrom, err := parseTaskListTime(query.Get("createdFrom"), query.Get("from"))
if err != nil {
writeError(w, http.StatusBadRequest, "invalid createdFrom")
return
}
createdTo, err := parseTaskListTime(query.Get("createdTo"), query.Get("to"))
if err != nil {
writeError(w, http.StatusBadRequest, "invalid createdTo")
return
}
for name, value := range map[string]string{
"tenantId": query.Get("tenantId"),
"userId": query.Get("userId"),
"userGroupId": query.Get("userGroupId"),
"platformId": query.Get("platformId"),
} {
if value != "" && !validUUID(value) {
writeError(w, http.StatusBadRequest, "invalid "+name)
return
}
}
result, err := s.store.ListAdminTasks(r.Context(), store.AdminTaskListFilter{
Query: firstNonEmpty(query.Get("q"), query.Get("query")),
GatewayTenant: query.Get("tenantId"),
GatewayUser: query.Get("userId"),
UserGroup: query.Get("userGroupId"),
Status: query.Get("status"),
Platform: query.Get("platformId"),
Model: query.Get("model"),
ModelType: query.Get("modelType"),
RunMode: query.Get("runMode"),
BillingStatus: query.Get("billingStatus"),
APIKey: query.Get("apiKey"),
CreatedFrom: createdFrom,
CreatedTo: createdTo,
Page: page,
PageSize: pageSize,
})
if err != nil {
s.logger.Error("list admin tasks failed", "error", err)
writeError(w, http.StatusInternalServerError, "list admin tasks failed")
return
}
for index := range result.Items {
result.Items[index] = store.MaskAdminGatewayTask(result.Items[index])
}
writeJSON(w, http.StatusOK, map[string]any{
"items": result.Items,
"total": result.Total,
"page": result.Page,
"pageSize": result.PageSize,
})
}
// getAdminTask godoc
// @Summary 管理员获取任务详情
// @Description 返回任意租户的任务详情;默认脱敏,sensitive=full 时返回完整存储内容。
// @Tags admin-tasks
// @Produce json
// @Security BearerAuth
// @Param taskID path string true "任务 ID"
// @Param sensitive query string false "敏感字段模式" Enums(masked,full) default(masked)
// @Success 200 {object} store.AdminGatewayTask
// @Failure 400 {object} ErrorEnvelope
// @Failure 401 {object} ErrorEnvelope
// @Failure 403 {object} ErrorEnvelope
// @Failure 404 {object} ErrorEnvelope
// @Failure 500 {object} ErrorEnvelope
// @Router /api/admin/tasks/{taskID} [get]
func (s *Server) getAdminTask(w http.ResponseWriter, r *http.Request) {
taskID := r.PathValue("taskID")
if !validUUID(taskID) {
writeError(w, http.StatusBadRequest, "invalid taskID")
return
}
full, ok := parseAdminTaskSensitiveMode(r.URL.Query().Get("sensitive"))
if !ok {
writeError(w, http.StatusBadRequest, "invalid sensitive mode")
return
}
task, err := s.store.GetAdminTask(r.Context(), taskID)
if err != nil {
if store.IsNotFound(err) {
writeError(w, http.StatusNotFound, "task not found")
return
}
s.logger.Error("get admin task failed", "taskID", taskID, "error", err)
writeError(w, http.StatusInternalServerError, "get admin task failed")
return
}
task.GatewayTask, err = s.hydrateTaskResult(r.Context(), task.GatewayTask)
if err != nil {
writeStoredBinaryResultError(w, err)
return
}
if !full {
task = store.MaskAdminGatewayTask(task)
}
writeJSON(w, http.StatusOK, task)
}
// adminTaskParamPreprocessing godoc
// @Summary 管理员获取任务参数预处理日志
// @Description 返回任意租户任务的参数改写、校验或模板处理日志;默认脱敏。
// @Tags admin-tasks
// @Produce json
// @Security BearerAuth
// @Param taskID path string true "任务 ID"
// @Param sensitive query string false "敏感字段模式" Enums(masked,full) default(masked)
// @Success 200 {object} TaskParamPreprocessingLogListResponse
// @Failure 400 {object} ErrorEnvelope
// @Failure 401 {object} ErrorEnvelope
// @Failure 403 {object} ErrorEnvelope
// @Failure 404 {object} ErrorEnvelope
// @Failure 500 {object} ErrorEnvelope
// @Router /api/admin/tasks/{taskID}/param-preprocessing [get]
func (s *Server) adminTaskParamPreprocessing(w http.ResponseWriter, r *http.Request) {
taskID := r.PathValue("taskID")
if !validUUID(taskID) {
writeError(w, http.StatusBadRequest, "invalid taskID")
return
}
full, ok := parseAdminTaskSensitiveMode(r.URL.Query().Get("sensitive"))
if !ok {
writeError(w, http.StatusBadRequest, "invalid sensitive mode")
return
}
if _, err := s.store.GetAdminTask(r.Context(), taskID); err != nil {
if store.IsNotFound(err) {
writeError(w, http.StatusNotFound, "task not found")
return
}
s.logger.Error("get admin task failed", "taskID", taskID, "error", err)
writeError(w, http.StatusInternalServerError, "get admin task failed")
return
}
items, err := s.store.ListTaskParamPreprocessingLogs(r.Context(), taskID)
if err != nil {
s.logger.Error("list admin task parameter preprocessing logs failed", "taskID", taskID, "error", err)
writeError(w, http.StatusInternalServerError, "list admin task parameter preprocessing logs failed")
return
}
if !full {
items = store.MaskTaskParamPreprocessingLogs(items)
}
writeJSON(w, http.StatusOK, map[string]any{"items": items})
}
func parseAdminTaskSensitiveMode(value string) (bool, bool) {
switch strings.ToLower(strings.TrimSpace(value)) {
case "", "masked":
return false, true
case "full":
return true, true
default:
return false, false
}
}
func validUUID(value string) bool {
_, err := uuid.Parse(strings.TrimSpace(value))
return err == nil
}
@@ -0,0 +1,22 @@
package httpapi
import "testing"
func TestParseAdminTaskSensitiveMode(t *testing.T) {
tests := []struct {
value string
full bool
accepted bool
}{
{value: "", full: false, accepted: true},
{value: "masked", full: false, accepted: true},
{value: "FULL", full: true, accepted: true},
{value: "invalid", full: false, accepted: false},
}
for _, test := range tests {
full, accepted := parseAdminTaskSensitiveMode(test.value)
if full != test.full || accepted != test.accepted {
t.Fatalf("parseAdminTaskSensitiveMode(%q)=(%v,%v), want (%v,%v)", test.value, full, accepted, test.full, test.accepted)
}
}
}
@@ -0,0 +1,321 @@
package httpapi
import (
"context"
"io"
"log/slog"
"net/http"
"net/http/httptest"
"net/url"
"os"
"strconv"
"strings"
"testing"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
func TestAdminTasksCrossTenantQueryAndRedaction(t *testing.T) {
databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL"))
if databaseURL == "" {
t.Skip("set AI_GATEWAY_TEST_DATABASE_URL to run the PostgreSQL integration flow")
}
ctx := context.Background()
applyMigration(t, ctx, databaseURL)
db, err := store.Connect(ctx, databaseURL)
if err != nil {
t.Fatalf("connect store: %v", err)
}
defer db.Close()
suffix := strconv.FormatInt(time.Now().UnixNano(), 10)
createTenant := func(key, name string) store.GatewayTenant {
t.Helper()
tenant, err := db.CreateTenant(ctx, store.GatewayTenantInput{
TenantKey: key + "-" + suffix,
Name: name,
Source: "gateway",
Status: "active",
})
if err != nil {
t.Fatalf("create tenant %s: %v", name, err)
}
return tenant
}
createUser := func(tenant store.GatewayTenant, prefix string, roles []string) store.GatewayUser {
t.Helper()
username := prefix + "_" + suffix
user, err := db.CreateGatewayUser(ctx, store.GatewayUserInput{
UserKey: "gateway:" + username,
Username: username,
DisplayName: prefix + " display",
Email: username + "@example.com",
Source: "gateway",
GatewayTenantID: tenant.ID,
TenantKey: tenant.TenantKey,
Roles: roles,
Status: "active",
})
if err != nil {
t.Fatalf("create gateway user %s: %v", prefix, err)
}
return user
}
authUser := func(user store.GatewayUser, apiKeyName string) *auth.User {
return &auth.User{
ID: user.ID,
Username: user.Username,
Roles: user.Roles,
Source: user.Source,
GatewayUserID: user.ID,
GatewayTenantID: user.GatewayTenantID,
TenantKey: user.TenantKey,
APIKeyID: "key-" + apiKeyName + "-" + suffix,
APIKeyName: apiKeyName,
APIKeyPrefix: "sk-" + apiKeyName,
}
}
createTask := func(user *auth.User, model, secret string) store.GatewayTask {
t.Helper()
task, err := db.CreateTask(ctx, store.CreateTaskInput{
Kind: "chat.completions",
Model: model,
RunMode: "simulation",
Request: map[string]any{
"model": model,
"diagnostic": map[string]any{
"apiKey": secret,
},
},
}, user)
if err != nil {
t.Fatalf("create task for %s: %v", user.Username, err)
}
return task
}
tenantA := createTenant("admin-task-a", "Admin Task Tenant A")
tenantB := createTenant("admin-task-b", "Admin Task Tenant B")
operator := createUser(tenantA, "admin_task_operator", []string{"operator"})
userA := createUser(tenantA, "admin_task_user_a", []string{"user"})
userB := createUser(tenantB, "admin_task_user_b", []string{"user"})
userAAuth := authUser(userA, "alpha-key")
userBAuth := authUser(userB, "bravo-key")
modelA := "admin-task-model-a-" + suffix
modelB := "admin-task-model-b-" + suffix
taskA := createTask(userAAuth, modelA, "alpha-secret-"+suffix)
taskBSecret := "bravo-secret-" + suffix
taskB := createTask(userBAuth, modelB, taskBSecret)
platform, err := db.CreatePlatform(ctx, store.CreatePlatformInput{
Provider: "openai",
PlatformKey: "admin-task-platform-" + suffix,
Name: "Admin Task Platform " + suffix,
BaseURL: "https://example.invalid/v1",
AuthType: "bearer",
Credentials: map[string]any{"apiKey": "platform-secret-" + suffix},
Status: "enabled",
})
if err != nil {
t.Fatalf("create platform: %v", err)
}
attemptID, err := db.CreateTaskAttempt(ctx, store.CreateTaskAttemptInput{
TaskID: taskB.ID,
AttemptNo: 1,
PlatformID: platform.ID,
QueueKey: "admin-task-integration-" + suffix,
Status: "succeeded",
Simulated: true,
})
if err != nil {
t.Fatalf("create task attempt: %v", err)
}
if _, err := db.Pool().Exec(ctx, `
UPDATE gateway_tasks
SET status = 'succeeded',
model_type = 'text_generate',
requested_model = $2,
resolved_model = $2,
request_id = $3,
billing_status = 'settled',
result = $4::jsonb,
finished_at = now(),
updated_at = now()
WHERE id = $1::uuid`,
taskB.ID,
modelB,
"request-admin-task-"+suffix,
`{"authorization":"Bearer result-secret","output":"ok"}`,
); err != nil {
t.Fatalf("enrich task: %v", err)
}
if _, err := db.Pool().Exec(ctx, `
UPDATE gateway_task_attempts
SET request_snapshot = $2::jsonb,
response_snapshot = $3::jsonb,
finished_at = now()
WHERE id = $1::uuid`,
attemptID,
`{"headers":{"authorization":"Bearer attempt-secret"}}`,
`{"cookie":"attempt-cookie","status":"ok"}`,
); err != nil {
t.Fatalf("enrich task attempt: %v", err)
}
if _, err := db.Pool().Exec(ctx, `
INSERT INTO gateway_task_param_preprocessing_logs (
task_id, attempt_id, attempt_no, model_type, platform_id, client_id,
changed, change_count, actual_input, converted_output, changes, model_snapshot
)
VALUES (
$1::uuid, $2::uuid, 1, 'text_generate', $3::uuid, $4,
true, 1, $5::jsonb, $6::jsonb, $7::jsonb, '{}'::jsonb
)`,
taskB.ID,
attemptID,
platform.ID,
"admin-task-client-"+suffix,
`{"password":"input-secret","prompt":"hello"}`,
`{"token":"converted-secret","prompt":"hello"}`,
`[{"path":"credentials.secret","before":"old-secret","after":"new-secret"}]`,
); err != nil {
t.Fatalf("create parameter preprocessing log: %v", err)
}
const jwtSecret = "admin-task-integration-jwt-secret"
handlerCtx, cancelHandler := context.WithCancel(ctx)
defer cancelHandler()
server := httptest.NewServer(NewServerWithContext(handlerCtx, config.Config{
AppEnv: "test",
HTTPAddr: ":0",
DatabaseURL: databaseURL,
IdentityMode: "hybrid",
JWTSecret: jwtSecret,
CORSAllowedOrigin: "*",
}, db, slog.New(slog.NewTextHandler(io.Discard, nil))))
defer server.Close()
authenticator := auth.New(jwtSecret, "", "")
signToken := func(user *auth.User) string {
t.Helper()
token, err := authenticator.SignJWT(user, time.Hour)
if err != nil {
t.Fatalf("sign JWT for %s: %v", user.Username, err)
}
return token
}
operatorToken := signToken(&auth.User{
ID: operator.ID,
Username: operator.Username,
Roles: operator.Roles,
Source: operator.Source,
GatewayUserID: operator.ID,
GatewayTenantID: operator.GatewayTenantID,
TenantKey: operator.TenantKey,
})
userAToken := signToken(userAAuth)
doJSON(t, server.URL, http.MethodGet, "/api/admin/tasks", userAToken, nil, http.StatusForbidden, nil)
query := url.Values{
"q": {taskB.ID},
"tenantId": {tenantB.ID},
"userId": {userB.ID},
"status": {"succeeded"},
"platformId": {platform.ID},
"model": {modelB},
"modelType": {"text_generate"},
"runMode": {"simulation"},
"billingStatus": {"settled"},
"apiKey": {"bravo-key"},
"page": {"1"},
"pageSize": {"10"},
}
var listResponse struct {
Items []store.AdminGatewayTask `json:"items"`
Total int `json:"total"`
Page int `json:"page"`
PageSize int `json:"pageSize"`
}
doJSON(t, server.URL, http.MethodGet, "/api/admin/tasks?"+query.Encode(), operatorToken, nil, http.StatusOK, &listResponse)
if len(listResponse.Items) != 1 || listResponse.Total != 1 || listResponse.Page != 1 || listResponse.PageSize != 10 {
t.Fatalf("unexpected filtered admin task list: %+v", listResponse)
}
listed := listResponse.Items[0]
if listed.ID != taskB.ID || listed.AdminContext.User == nil || listed.AdminContext.User.ID != userB.ID ||
listed.AdminContext.Tenant == nil || listed.AdminContext.Tenant.ID != tenantB.ID ||
listed.AdminContext.LatestPlatform == nil || listed.AdminContext.LatestPlatform.ID != platform.ID {
t.Fatalf("admin task list should expose cross-tenant identity and platform summaries: %+v", listed)
}
if nestedString(listed.Request, "diagnostic", "apiKey") != "***" ||
nestedString(listed.Result, "authorization") != "***" ||
nestedString(listed.Attempts[0].RequestSnapshot, "headers", "authorization") != "***" {
t.Fatalf("admin task list should mask task and attempt secrets: %+v", listed)
}
var maskedDetail store.AdminGatewayTask
doJSON(t, server.URL, http.MethodGet, "/api/admin/tasks/"+taskB.ID, operatorToken, nil, http.StatusOK, &maskedDetail)
if nestedString(maskedDetail.Request, "diagnostic", "apiKey") != "***" {
t.Fatalf("admin task detail should be masked by default: %+v", maskedDetail.Request)
}
var fullDetail store.AdminGatewayTask
doJSON(t, server.URL, http.MethodGet, "/api/admin/tasks/"+taskB.ID+"?sensitive=full", operatorToken, nil, http.StatusOK, &fullDetail)
if nestedString(fullDetail.Request, "diagnostic", "apiKey") != taskBSecret {
t.Fatalf("explicit full task detail should expose stored content: %+v", fullDetail.Request)
}
var maskedLogs struct {
Items []store.TaskParamPreprocessingLog `json:"items"`
}
doJSON(t, server.URL, http.MethodGet, "/api/admin/tasks/"+taskB.ID+"/param-preprocessing", operatorToken, nil, http.StatusOK, &maskedLogs)
if len(maskedLogs.Items) != 1 ||
nestedString(maskedLogs.Items[0].ActualInput, "password") != "***" ||
nestedString(maskedLogs.Items[0].ConvertedOutput, "token") != "***" {
t.Fatalf("admin task preprocessing logs should be masked by default: %+v", maskedLogs.Items)
}
var fullLogs struct {
Items []store.TaskParamPreprocessingLog `json:"items"`
}
doJSON(t, server.URL, http.MethodGet, "/api/admin/tasks/"+taskB.ID+"/param-preprocessing?sensitive=full", operatorToken, nil, http.StatusOK, &fullLogs)
if len(fullLogs.Items) != 1 ||
nestedString(fullLogs.Items[0].ActualInput, "password") != "input-secret" ||
nestedString(fullLogs.Items[0].ConvertedOutput, "token") != "converted-secret" {
t.Fatalf("explicit full preprocessing logs should expose stored content: %+v", fullLogs.Items)
}
var ownTasks struct {
Items []store.GatewayTask `json:"items"`
}
doJSON(t, server.URL, http.MethodGet, "/api/workspace/tasks?pageSize=100", userAToken, nil, http.StatusOK, &ownTasks)
if !gatewayTaskListHasID(ownTasks.Items, taskA.ID) || gatewayTaskListHasID(ownTasks.Items, taskB.ID) {
t.Fatalf("workspace task list must remain scoped to the current user: %+v", ownTasks.Items)
}
doJSON(t, server.URL, http.MethodGet, "/api/workspace/tasks/"+taskB.ID, userAToken, nil, http.StatusNotFound, nil)
doJSON(t, server.URL, http.MethodGet, "/api/workspace/tasks/"+taskB.ID+"/param-preprocessing", userAToken, nil, http.StatusNotFound, nil)
}
func nestedString(value map[string]any, path ...string) string {
var current any = value
for _, key := range path {
object, ok := current.(map[string]any)
if !ok {
return ""
}
current = object[key]
}
result, _ := current.(string)
return result
}
func gatewayTaskListHasID(items []store.GatewayTask, taskID string) bool {
for _, item := range items {
if item.ID == taskID {
return true
}
}
return false
}
@@ -11,7 +11,7 @@ import "net/http"
// @Security BearerAuth
// @Param X-Async header bool false "true 时异步创建任务并返回 202"
// @Param input body ImageVectorizeRequest true "图片矢量化请求"
// @Success 200 {object} CompatibleResponse
// @Success 200 {object} EasyAIGeneratedResponse
// @Success 202 {object} TaskAcceptedResponse
// @Failure 400 {object} ErrorEnvelope
// @Failure 401 {object} ErrorEnvelope
@@ -32,7 +32,7 @@ func (s *Server) createImageVectorizeTask() http.Handler {
// @Security BearerAuth
// @Param X-Async header bool false "true 时异步创建任务并返回 202"
// @Param input body VideoUpscaleRequest true "视频超分请求"
// @Success 200 {object} CompatibleResponse
// @Success 200 {object} EasyAIGeneratedResponse
// @Success 202 {object} TaskAcceptedResponse
// @Failure 400 {object} ErrorEnvelope
// @Failure 401 {object} ErrorEnvelope
@@ -0,0 +1,479 @@
package httpapi
import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"log/slog"
"net/http"
"net/http/httptest"
"os"
"strconv"
"strings"
"sync"
"testing"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
"github.com/jackc/pgx/v5/pgxpool"
)
func TestAsyncWorkerDynamicConcurrencyAcceptance(t *testing.T) {
databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL"))
if databaseURL == "" {
t.Skip("set AI_GATEWAY_TEST_DATABASE_URL to run dynamic worker acceptance tests")
}
ctx, cancel := context.WithTimeout(context.Background(), 7*time.Minute)
defer cancel()
applyMigration(t, ctx, databaseURL)
db, err := store.Connect(ctx, databaseURL)
if err != nil {
t.Fatalf("connect store: %v", err)
}
defer db.Close()
pool, err := pgxpool.New(ctx, databaseURL)
if err != nil {
t.Fatalf("connect acceptance pool: %v", err)
}
defer pool.Close()
if _, err := pool.Exec(ctx, `UPDATE integration_platforms SET status = 'disabled' WHERE deleted_at IS NULL`); err != nil {
t.Fatalf("isolate acceptance platforms: %v", err)
}
serverCtx, cancelServer := context.WithCancel(ctx)
defer cancelServer()
server := httptest.NewServer(NewServerWithContext(serverCtx, config.Config{
AppEnv: "test",
HTTPAddr: ":0",
DatabaseURL: databaseURL,
IdentityMode: "hybrid",
JWTSecret: "test-secret",
BillingEngineMode: "observe",
CORSAllowedOrigin: "*",
AsyncWorkerHardLimit: 256,
AsyncWorkerRefreshIntervalSeconds: 1,
}, db, slog.New(slog.NewTextHandler(io.Discard, nil))))
defer server.Close()
adminToken := createAsyncAcceptanceAdmin(t, ctx, pool, server.URL)
if _, err := pool.Exec(ctx, `
UPDATE gateway_user_groups
SET rate_limit_policy = '{"rules":[{"metric":"concurrent","limit":256,"leaseTtlSeconds":120}]}'::jsonb
WHERE status = 'active'`); err != nil {
t.Fatalf("raise acceptance user-group concurrency: %v", err)
}
suffix := strconv.FormatInt(time.Now().UnixNano(), 10)
t.Run("超过旧64上限的高并发", func(t *testing.T) {
model := "worker-burst-" + suffix
platform := createAsyncAcceptancePlatform(t, server.URL, adminToken, "burst-"+suffix, "Burst Simulation", 128)
defer updateAsyncAcceptancePlatform(t, server.URL, adminToken, platform, 128, "disabled")
modelID := createAsyncAcceptanceModel(t, server.URL, adminToken, platform.ID, model, "video_generate", "override", 96, 120)
waitForAsyncWorkerMetric(t, server.URL, "easyai_gateway_async_worker_capacity", 96, 15*time.Second)
startedAt := time.Now()
taskIDs := submitAsyncSimulationTasks(t, server.URL, adminToken, "/api/v1/videos/generations", model, 128, 15*time.Second)
peakRunning, peakLeases := waitForAcceptanceTasks(t, ctx, pool, taskIDs, modelID, 60*time.Second)
elapsed := time.Since(startedAt)
if peakRunning < 80 || peakLeases < 80 {
t.Fatalf("running peak=%d lease peak=%d, want both >=80", peakRunning, peakLeases)
}
if peakLeases > 96 {
t.Fatalf("lease peak=%d exceeded model concurrent limit 96", peakLeases)
}
assertAcceptanceAttempts(t, ctx, pool, taskIDs, len(taskIDs))
if elapsed >= 60*time.Second {
t.Fatalf("128 tasks completed in %s, want <60s", elapsed)
}
t.Logf("高并发证据: tasks=128 duration=%s running_peak=%d lease_peak=%d capacity=96", elapsed.Round(time.Millisecond), peakRunning, peakLeases)
})
t.Run("三分钟视频任务在线扩缩容与续租", func(t *testing.T) {
model := "worker-video-long-" + suffix
platform := createAsyncAcceptancePlatform(t, server.URL, adminToken, "video-long-"+suffix, "Long Video Simulation", 1)
modelID := createAsyncAcceptanceModel(t, server.URL, adminToken, platform.ID, model, "video_generate", "inherit", 0, 120)
waitForAsyncWorkerMetric(t, server.URL, "easyai_gateway_async_worker_capacity", 1, 15*time.Second)
longTaskIDs := submitAsyncSimulationTasks(t, server.URL, adminToken, "/api/v1/videos/generations", model, 3, 180*time.Second)
waitForTaskCount(t, ctx, pool, longTaskIDs, "running", 1, 10*time.Second)
updateStartedAt := time.Now()
updateAsyncAcceptancePlatform(t, server.URL, adminToken, platform, 3, "enabled")
waitForAsyncWorkerMetric(t, server.URL, "easyai_gateway_async_worker_desired_capacity", 3, 15*time.Second)
waitForAsyncWorkerMetric(t, server.URL, "easyai_gateway_async_worker_capacity", 3, 15*time.Second)
waitForActiveLeaseCount(t, ctx, pool, modelID, 3, 15*time.Second)
waitForTaskCount(t, ctx, pool, longTaskIDs, "running", 3, 5*time.Second)
if time.Since(updateStartedAt) > 15*time.Second {
t.Fatalf("capacity expansion took %s, want <=15s", time.Since(updateStartedAt))
}
initialExpiry := activeLeaseExpiry(t, ctx, pool, modelID, 3)
updateAsyncAcceptancePlatform(t, server.URL, adminToken, platform, 1, "enabled")
waitForAsyncWorkerMetric(t, server.URL, "easyai_gateway_async_worker_capacity", 1, 15*time.Second)
shortTaskIDs := submitAsyncSimulationTasks(t, server.URL, adminToken, "/api/v1/videos/generations", model, 1, time.Second)
time.Sleep(3 * time.Second)
if attempts := taskAttemptCount(t, ctx, pool, shortTaskIDs); attempts != 0 {
t.Fatalf("short task created %d upstream attempts while three long leases exceeded reduced limit", attempts)
}
waitForTaskCount(t, ctx, pool, longTaskIDs, "running", 3, 5*time.Second)
time.Sleep(130 * time.Second)
renewedExpiry := activeLeaseExpiry(t, ctx, pool, modelID, 3)
if !renewedExpiry.After(initialExpiry.Add(30*time.Second)) || !renewedExpiry.After(time.Now()) {
t.Fatalf("lease was not renewed: initial=%s renewed=%s now=%s", initialExpiry, renewedExpiry, time.Now())
}
waitForTaskCount(t, ctx, pool, longTaskIDs, "succeeded", 3, 75*time.Second)
shortStartedAt := time.Now()
waitForTaskAttempts(t, ctx, pool, shortTaskIDs, 1, 5*time.Second)
waitForTaskCount(t, ctx, pool, shortTaskIDs, "succeeded", 1, 10*time.Second)
if time.Since(shortStartedAt) > 5*time.Second {
t.Fatalf("short task started after %s once long tasks completed, want <=5s", time.Since(shortStartedAt))
}
allTaskIDs := append(append([]string{}, longTaskIDs...), shortTaskIDs...)
assertAcceptanceAttempts(t, ctx, pool, allTaskIDs, 4)
assertSimulationOutputsAndEvents(t, ctx, pool, allTaskIDs)
t.Logf("长任务证据: initial_capacity=1 expanded_capacity=3 reduced_capacity=1 initial_expiry=%s renewed_expiry=%s tasks=4", initialExpiry.UTC().Format(time.RFC3339Nano), renewedExpiry.UTC().Format(time.RFC3339Nano))
})
}
type asyncAcceptancePlatform struct {
ID string `json:"id"`
Provider string `json:"provider"`
PlatformKey string `json:"platformKey"`
Name string `json:"name"`
BaseURL string `json:"baseUrl"`
AuthType string `json:"authType"`
Priority int `json:"priority"`
}
func createAsyncAcceptanceAdmin(t *testing.T, ctx context.Context, pool *pgxpool.Pool, baseURL string) string {
t.Helper()
suffix := strconv.FormatInt(time.Now().UnixNano(), 10)
username := "async_acceptance_" + suffix
password := "password123"
doJSON(t, baseURL, http.MethodPost, "/api/v1/auth/register", "", map[string]any{
"username": username,
"email": username + "@example.com",
"password": password,
}, http.StatusCreated, &map[string]any{})
if _, err := pool.Exec(ctx, `UPDATE gateway_users SET roles = '["admin"]'::jsonb WHERE username = $1`, username); err != nil {
t.Fatalf("promote acceptance user: %v", err)
}
var login struct {
AccessToken string `json:"accessToken"`
}
doJSON(t, baseURL, http.MethodPost, "/api/v1/auth/login", "", map[string]any{
"account": username, "password": password,
}, http.StatusOK, &login)
if login.AccessToken == "" {
t.Fatal("acceptance login returned no access token")
}
return login.AccessToken
}
func createAsyncAcceptancePlatform(t *testing.T, baseURL, token, key, name string, concurrent int) asyncAcceptancePlatform {
t.Helper()
var platform asyncAcceptancePlatform
doJSON(t, baseURL, http.MethodPost, "/api/admin/platforms", token, asyncAcceptancePlatformPayload(key, name, concurrent, "enabled"), http.StatusCreated, &platform)
if platform.ID == "" {
t.Fatal("acceptance platform returned no id")
}
return platform
}
func updateAsyncAcceptancePlatform(t *testing.T, baseURL, token string, platform asyncAcceptancePlatform, concurrent int, status string) {
t.Helper()
doJSON(t, baseURL, http.MethodPatch, "/api/admin/platforms/"+platform.ID, token,
asyncAcceptancePlatformPayload(platform.PlatformKey, platform.Name, concurrent, status), http.StatusOK, &asyncAcceptancePlatform{})
}
func asyncAcceptancePlatformPayload(key, name string, concurrent int, status string) map[string]any {
return map[string]any{
"provider": "openai",
"platformKey": key,
"name": name,
"baseUrl": "https://api.openai.com/v1",
"authType": "bearer",
"credentials": map[string]any{"mode": "simulation"},
"config": map[string]any{"testMode": true},
"priority": 1,
"status": status,
"rateLimitPolicy": map[string]any{"rules": []any{
map[string]any{"metric": "concurrent", "limit": concurrent, "leaseTtlSeconds": 120},
}},
}
}
func createAsyncAcceptanceModel(t *testing.T, baseURL, token, platformID, model, modelType, mode string, concurrent, ttl int) string {
t.Helper()
payload := map[string]any{
"modelName": model,
"providerModelName": model,
"modelAlias": model,
"modelType": []string{modelType},
"displayName": model,
"rateLimitPolicyMode": mode,
}
if mode == "override" {
payload["rateLimitPolicy"] = map[string]any{"rules": []any{
map[string]any{"metric": "concurrent", "limit": concurrent, "leaseTtlSeconds": ttl},
}}
}
var response struct {
ID string `json:"id"`
}
doJSON(t, baseURL, http.MethodPost, "/api/admin/platforms/"+platformID+"/models", token, payload, http.StatusCreated, &response)
if response.ID == "" {
t.Fatal("acceptance model returned no id")
}
return response.ID
}
func submitAsyncSimulationTasks(t *testing.T, baseURL, token, path, model string, count int, duration time.Duration) []string {
t.Helper()
ids := make([]string, count)
errs := make(chan error, count)
var wg sync.WaitGroup
for index := 0; index < count; index++ {
wg.Add(1)
go func(index int) {
defer wg.Done()
payload := map[string]any{
"model": model,
"runMode": "simulation",
"simulation": true,
"simulationDurationMs": duration.Milliseconds(),
"prompt": fmt.Sprintf("async acceptance %d", index),
"input": fmt.Sprintf("async acceptance %d", index),
}
raw, _ := json.Marshal(payload)
req, err := http.NewRequest(http.MethodPost, baseURL+path, bytes.NewReader(raw))
if err != nil {
errs <- err
return
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", "Bearer "+token)
req.Header.Set("X-Async", "true")
resp, err := http.DefaultClient.Do(req)
if err != nil {
errs <- err
return
}
defer resp.Body.Close()
responseBody, _ := io.ReadAll(resp.Body)
if resp.StatusCode != http.StatusAccepted {
errs <- fmt.Errorf("submit task %d status=%d body=%s", index, resp.StatusCode, responseBody)
return
}
var response struct {
TaskID string `json:"taskId"`
}
if err := json.Unmarshal(responseBody, &response); err != nil || response.TaskID == "" {
errs <- fmt.Errorf("decode task %d: %w body=%s", index, err, responseBody)
return
}
ids[index] = response.TaskID
}(index)
}
wg.Wait()
close(errs)
for err := range errs {
t.Fatal(err)
}
return ids
}
func waitForAcceptanceTasks(t *testing.T, ctx context.Context, pool *pgxpool.Pool, taskIDs []string, modelID string, timeout time.Duration) (int64, int64) {
t.Helper()
deadline := time.Now().Add(timeout)
var peakRunning, peakLeases int64
for time.Now().Before(deadline) {
var running, succeeded, failed, leases int64
if err := pool.QueryRow(ctx, `
SELECT COUNT(*) FILTER (WHERE status = 'running'),
COUNT(*) FILTER (WHERE status = 'succeeded'),
COUNT(*) FILTER (WHERE status IN ('failed', 'cancelled', 'manual_review'))
FROM gateway_tasks
WHERE id = ANY($1::uuid[])`, taskIDs).Scan(&running, &succeeded, &failed); err != nil {
t.Fatalf("read acceptance tasks: %v", err)
}
if err := pool.QueryRow(ctx, `
SELECT COUNT(*)
FROM gateway_concurrency_leases
WHERE scope_type = 'platform_model'
AND scope_key = $1
AND released_at IS NULL
AND expires_at > now()`, modelID).Scan(&leases); err != nil {
t.Fatalf("read acceptance leases: %v", err)
}
peakRunning = maxInt64(peakRunning, running)
peakLeases = maxInt64(peakLeases, leases)
if failed > 0 {
t.Fatalf("acceptance tasks entered failed/manual status: %d", failed)
}
if succeeded == int64(len(taskIDs)) {
return peakRunning, peakLeases
}
time.Sleep(50 * time.Millisecond)
}
t.Fatalf("tasks did not finish within %s", timeout)
return 0, 0
}
func waitForTaskCount(t *testing.T, ctx context.Context, pool *pgxpool.Pool, taskIDs []string, status string, count int, timeout time.Duration) {
t.Helper()
deadline := time.Now().Add(timeout)
for time.Now().Before(deadline) {
var got int
if err := pool.QueryRow(ctx, `SELECT COUNT(*) FROM gateway_tasks WHERE id = ANY($1::uuid[]) AND status = $2`, taskIDs, status).Scan(&got); err != nil {
t.Fatalf("read task count: %v", err)
}
if got == count {
return
}
time.Sleep(100 * time.Millisecond)
}
t.Fatalf("task status %s count did not reach %d within %s", status, count, timeout)
}
func activeLeaseExpiry(t *testing.T, ctx context.Context, pool *pgxpool.Pool, modelID string, wantCount int) time.Time {
t.Helper()
var count int
var earliest time.Time
if err := pool.QueryRow(ctx, `
SELECT COUNT(*), MIN(expires_at)
FROM gateway_concurrency_leases
WHERE scope_type = 'platform_model'
AND scope_key = $1
AND released_at IS NULL
AND expires_at > now()`, modelID).Scan(&count, &earliest); err != nil {
t.Fatalf("read active lease expiry: %v", err)
}
if count != wantCount {
t.Fatalf("active lease count=%d, want=%d", count, wantCount)
}
return earliest
}
func waitForActiveLeaseCount(t *testing.T, ctx context.Context, pool *pgxpool.Pool, modelID string, count int, timeout time.Duration) {
t.Helper()
deadline := time.Now().Add(timeout)
for time.Now().Before(deadline) {
var got int
if err := pool.QueryRow(ctx, `
SELECT COUNT(*)
FROM gateway_concurrency_leases
WHERE scope_type = 'platform_model'
AND scope_key = $1
AND released_at IS NULL
AND expires_at > now()`, modelID).Scan(&got); err != nil {
t.Fatalf("read active lease count: %v", err)
}
if got == count {
return
}
time.Sleep(100 * time.Millisecond)
}
t.Fatalf("active lease count did not reach %d within %s", count, timeout)
}
func taskAttemptCount(t *testing.T, ctx context.Context, pool *pgxpool.Pool, taskIDs []string) int {
t.Helper()
var count int
if err := pool.QueryRow(ctx, `SELECT COUNT(*) FROM gateway_task_attempts WHERE task_id = ANY($1::uuid[])`, taskIDs).Scan(&count); err != nil {
t.Fatalf("count task attempts: %v", err)
}
return count
}
func waitForTaskAttempts(t *testing.T, ctx context.Context, pool *pgxpool.Pool, taskIDs []string, count int, timeout time.Duration) {
t.Helper()
deadline := time.Now().Add(timeout)
for time.Now().Before(deadline) {
if taskAttemptCount(t, ctx, pool, taskIDs) == count {
return
}
time.Sleep(50 * time.Millisecond)
}
t.Fatalf("attempt count did not reach %d within %s", count, timeout)
}
func assertAcceptanceAttempts(t *testing.T, ctx context.Context, pool *pgxpool.Pool, taskIDs []string, want int) {
t.Helper()
var attempts, duplicateTasks, manualReview int
if err := pool.QueryRow(ctx, `
SELECT COUNT(*),
COUNT(DISTINCT task_id) FILTER (WHERE per_task.attempts > 1),
COUNT(*) FILTER (WHERE status = 'manual_review')
FROM (
SELECT task_id, COUNT(*) AS attempts, MAX(status) AS status
FROM gateway_task_attempts
WHERE task_id = ANY($1::uuid[])
GROUP BY task_id
) per_task`, taskIDs).Scan(&attempts, &duplicateTasks, &manualReview); err != nil {
t.Fatalf("read attempt acceptance evidence: %v", err)
}
if attempts != want || duplicateTasks != 0 || manualReview != 0 {
t.Fatalf("attempts=%d duplicate_tasks=%d manual_review=%d, want %d/0/0", attempts, duplicateTasks, manualReview, want)
}
}
func assertSimulationOutputsAndEvents(t *testing.T, ctx context.Context, pool *pgxpool.Pool, taskIDs []string) {
t.Helper()
var complete int
if err := pool.QueryRow(ctx, `
SELECT COUNT(*)
FROM gateway_tasks task
WHERE task.id = ANY($1::uuid[])
AND task.status = 'succeeded'
AND task.finished_at IS NOT NULL
AND task.result IS NOT NULL
AND task.result <> '{}'::jsonb
AND task.metrics IS NOT NULL
AND jsonb_typeof(task.billings) = 'array'
AND jsonb_array_length(task.billings) > 0
AND EXISTS (SELECT 1 FROM gateway_task_events event WHERE event.task_id = task.id AND event.event_type = 'task.completed')`,
taskIDs).Scan(&complete); err != nil {
t.Fatalf("read simulation output evidence: %v", err)
}
if complete != len(taskIDs) {
t.Fatalf("complete simulation outputs/events=%d, want=%d", complete, len(taskIDs))
}
}
func waitForAsyncWorkerMetric(t *testing.T, baseURL, metric string, want int64, timeout time.Duration) {
t.Helper()
deadline := time.Now().Add(timeout)
for time.Now().Before(deadline) {
resp, err := http.Get(baseURL + "/metrics")
if err == nil {
raw, _ := io.ReadAll(resp.Body)
resp.Body.Close()
if resp.StatusCode == http.StatusOK {
for _, line := range strings.Split(string(raw), "\n") {
fields := strings.Fields(line)
if len(fields) == 2 && fields[0] == metric {
value, _ := strconv.ParseInt(fields[1], 10, 64)
if value == want {
return
}
}
}
}
}
time.Sleep(100 * time.Millisecond)
}
t.Fatalf("metric %s did not reach %d within %s", metric, want, timeout)
}
func maxInt64(left, right int64) int64 {
if right > left {
return right
}
return left
}
@@ -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)
}
@@ -182,6 +182,10 @@ func (s *Server) createBaseModel(w http.ResponseWriter, r *http.Request) {
}
item, err := s.store.CreateBaseModel(r.Context(), input)
if err != nil {
if errors.Is(err, store.ErrModelAliasConflict) {
writeError(w, http.StatusConflict, err.Error(), "model_alias_conflict")
return
}
if store.IsUniqueViolation(err) {
writeError(w, http.StatusConflict, "canonical model key already exists")
return
@@ -230,6 +234,10 @@ func (s *Server) updateBaseModel(w http.ResponseWriter, r *http.Request) {
writeError(w, http.StatusConflict, "canonical model key already exists")
return
}
if errors.Is(err, store.ErrModelAliasConflict) {
writeError(w, http.StatusConflict, err.Error(), "model_alias_conflict")
return
}
s.logger.Error("update base model failed", "error", err)
writeError(w, http.StatusInternalServerError, "update base model failed")
return
@@ -301,6 +309,7 @@ func (s *Server) resetAllBaseModels(w http.ResponseWriter, r *http.Request) {
// @Failure 401 {object} ErrorEnvelope
// @Failure 403 {object} ErrorEnvelope
// @Failure 404 {object} ErrorEnvelope
// @Failure 409 {object} ErrorEnvelope
// @Failure 500 {object} ErrorEnvelope
// @Router /api/admin/catalog/base-models/{baseModelID} [delete]
func (s *Server) deleteBaseModel(w http.ResponseWriter, r *http.Request) {
@@ -309,6 +318,10 @@ func (s *Server) deleteBaseModel(w http.ResponseWriter, r *http.Request) {
writeError(w, http.StatusNotFound, "base model not found")
return
}
if errors.Is(err, store.ErrBaseModelInUse) {
writeError(w, http.StatusConflict, "base model is still referenced by platform models", "base_model_in_use")
return
}
s.logger.Error("delete base model failed", "error", err)
writeError(w, http.StatusInternalServerError, "delete base model failed")
return
@@ -123,6 +123,17 @@ func writeOpenAIError(w http.ResponseWriter, status int, message string, details
if strings.TrimSpace(code) != "" {
payload["code"] = code
}
if len(details) > 0 {
publicDetails := map[string]any{}
for key, value := range details {
if key != "param" {
publicDetails[key] = value
}
}
if len(publicDetails) > 0 {
payload["details"] = publicDetails
}
}
writeJSON(w, status, map[string]any{"error": payload})
}
@@ -300,23 +300,41 @@ VALUES ($1, 5, '{"purpose":"core-flow"}'::jsonb)`, inviteCode); err != nil {
Items []struct {
ID string `json:"id"`
CanonicalModelKey string `json:"canonicalModelKey"`
InvocationName string `json:"invocationName"`
ProviderModelName string `json:"providerModelName"`
ModelType []string `json:"modelType"`
ModelAlias string `json:"modelAlias"`
DisplayName string `json:"displayName"`
Status string `json:"status"`
} `json:"items"`
}
doJSON(t, server.URL, http.MethodGet, "/api/admin/catalog/base-models", loginResponse.AccessToken, nil, http.StatusOK, &baseModels)
if len(baseModels.Items) < 300 {
t.Fatalf("server-main seed should include the migrated base model catalog: got %d", len(baseModels.Items))
if len(baseModels.Items) == 0 {
t.Fatal("migrated base model catalog must not be empty")
}
requiredLifecycleModels := map[string]string{
"gemini:gemini-3.6-flash": "active",
"aliyun:qwen3.7-max": "active",
}
for _, model := range baseModels.Items {
if expectedStatus, ok := requiredLifecycleModels[model.CanonicalModelKey]; ok {
if model.InvocationName == "" || model.ModelAlias != model.InvocationName || model.Status != expectedStatus {
t.Fatalf("base model identity/lifecycle mismatch: %+v", model)
}
delete(requiredLifecycleModels, model.CanonicalModelKey)
}
}
if len(requiredLifecycleModels) != 0 {
t.Fatalf("missing lifecycle-managed base models: %+v", requiredLifecycleModels)
}
requiredMiniMaxModels := map[string]struct {
providerModelName string
modelType string
alias string
}{
"minimax:speech-2.8-hd": {providerModelName: "speech-2.8-hd", modelType: "text_to_speech", alias: "MiniMax-Speech-2.8-HD"},
"minimax:speech-2.8-turbo": {providerModelName: "speech-2.8-turbo", modelType: "text_to_speech", alias: "MiniMax-Speech-2.8-Turbo"},
"minimax:voice-clone": {providerModelName: "voice_clone", modelType: "voice_clone", alias: "MiniMax-Voice-Clone"},
"minimax:speech-2.8-hd": {providerModelName: "speech-2.8-hd", modelType: "text_to_speech", alias: "speech-2.8-hd"},
"minimax:speech-2.8-turbo": {providerModelName: "speech-2.8-turbo", modelType: "text_to_speech", alias: "speech-2.8-turbo"},
"minimax:voice-clone": {providerModelName: "voice_clone", modelType: "voice_clone", alias: "voice_clone"},
}
seenMiniMaxModels := map[string]bool{}
for _, model := range baseModels.Items {
@@ -361,9 +379,11 @@ VALUES ($1, 5, '{"purpose":"core-flow"}'::jsonb)`, inviteCode); err != nil {
baseModelInput := map[string]any{
"providerKey": "openai",
"canonicalModelKey": "openai:smoke-base-" + suffixText,
"invocationName": "smoke-base-" + suffixText,
"providerModelName": "smoke-base-" + suffixText,
"modelType": []string{"text_generate"},
"modelAlias": "Smoke Base Model",
"displayName": "Smoke Base Model",
"legacyAliases": []string{"Smoke Base Model Legacy"},
"capabilities": map[string]any{"originalTypes": []string{"text_generate"}},
"metadata": map[string]any{"source": "test"},
}
@@ -376,12 +396,13 @@ VALUES ($1, 5, '{"purpose":"core-flow"}'::jsonb)`, inviteCode); err != nil {
if createdBaseModel.ID == "" || createdBaseModel.CanonicalModelKey != baseModelInput["canonicalModelKey"] {
t.Fatalf("unexpected created base model: %+v", createdBaseModel)
}
baseModelInput["modelAlias"] = "Smoke Base Model Updated"
baseModelInput["displayName"] = "Smoke Base Model Updated"
var updatedBaseModel struct {
ModelAlias string `json:"modelAlias"`
ModelAlias string `json:"modelAlias"`
DisplayName string `json:"displayName"`
}
doJSON(t, server.URL, http.MethodPatch, "/api/admin/catalog/base-models/"+createdBaseModel.ID, loginResponse.AccessToken, baseModelInput, http.StatusOK, &updatedBaseModel)
if updatedBaseModel.ModelAlias != "Smoke Base Model Updated" {
if updatedBaseModel.ModelAlias != baseModelInput["invocationName"] || updatedBaseModel.DisplayName != "Smoke Base Model Updated" {
t.Fatalf("unexpected updated base model: %+v", updatedBaseModel)
}
doJSON(t, server.URL, http.MethodDelete, "/api/admin/catalog/base-models/"+createdBaseModel.ID, loginResponse.AccessToken, nil, http.StatusNoContent, nil)
@@ -466,6 +487,7 @@ VALUES ($1, 5, '{"purpose":"core-flow"}'::jsonb)`, inviteCode); err != nil {
cancelReq.Header.Set("Content-Type", "application/json")
cancelErrCh := make(chan error, 1)
cancelTaskIDCh := make(chan string, 1)
cancelRequestStartedAt := time.Now().UTC().Add(-time.Second)
go func() {
resp, err := http.DefaultClient.Do(cancelReq)
if resp != nil {
@@ -484,7 +506,7 @@ VALUES ($1, 5, '{"purpose":"core-flow"}'::jsonb)`, inviteCode); err != nil {
t.Fatal("cancelled stream did not return response headers")
}
if cancelTaskID == "" {
t.Fatal("cancelled stream response did not expose X-Gateway-Task-Id")
cancelTaskID = waitForLatestTaskIDSince(t, ctx, testPool, "chat.completions", cancelRequestStartedAt, 2*time.Second)
}
cancelRequest()
select {
@@ -498,11 +520,20 @@ VALUES ($1, 5, '{"purpose":"core-flow"}'::jsonb)`, inviteCode); err != nil {
}
var imageResponse struct {
Task struct {
Status string `json:"status"`
LegacyTaskID string `json:"task_id"`
TaskID string `json:"taskId"`
QueryURL string `json:"query_url"`
Task struct {
ID string `json:"id"`
Status string `json:"status"`
Result map[string]any `json:"result"`
} `json:"task"`
Next struct {
Detail string `json:"detail"`
Events string `json:"events"`
Result string `json:"result"`
} `json:"next"`
}
doJSONWithHeaders(t, server.URL, http.MethodPost, "/api/v1/images/generations", apiKeyResponse.Secret, map[string]any{
"model": defaultImageModel,
@@ -513,14 +544,59 @@ VALUES ($1, 5, '{"purpose":"core-flow"}'::jsonb)`, inviteCode); err != nil {
"simulation": true,
"simulationDurationMs": 5,
}, map[string]string{"X-Async": "true"}, http.StatusAccepted, &imageResponse)
if imageResponse.Status != "submitted" ||
imageResponse.TaskID == "" ||
imageResponse.LegacyTaskID != imageResponse.TaskID ||
imageResponse.Task.ID != imageResponse.TaskID ||
imageResponse.QueryURL != "/api/v1/ai/result/"+imageResponse.TaskID ||
imageResponse.Next.Result != imageResponse.QueryURL {
t.Fatalf("unexpected EasyAI-compatible image submission: %+v", imageResponse)
}
waitForTaskStatus(t, server.URL, apiKeyResponse.Secret, imageResponse.Task.ID, []string{"succeeded"}, 10*time.Second)
doJSON(t, server.URL, http.MethodGet, "/api/v1/tasks/"+imageResponse.Task.ID, apiKeyResponse.Secret, nil, http.StatusOK, &imageResponse.Task)
if imageResponse.Task.Status != "succeeded" || imageResponse.Task.Result["id"] == "" {
t.Fatalf("unexpected image generation task: %+v", imageResponse.Task)
}
var imageEasyAIResult struct {
Status string `json:"status"`
TaskID string `json:"task_id"`
Data []map[string]any `json:"data"`
Output []string `json:"output"`
OutputContent []map[string]any `json:"output_content"`
}
doJSON(t, server.URL, http.MethodGet, imageResponse.QueryURL, apiKeyResponse.Secret, nil, http.StatusOK, &imageEasyAIResult)
if imageEasyAIResult.Status != "success" ||
imageEasyAIResult.TaskID != imageResponse.TaskID ||
len(imageEasyAIResult.Data) == 0 ||
len(imageEasyAIResult.Output) == 0 ||
len(imageEasyAIResult.OutputContent) != len(imageEasyAIResult.Data) ||
imageEasyAIResult.Data[0]["type"] != "image" {
t.Fatalf("unexpected EasyAI-compatible image result: %+v", imageEasyAIResult)
}
var ordinaryLoginResponse struct {
AccessToken string `json:"accessToken"`
}
doJSON(t, server.URL, http.MethodPost, "/api/v1/auth/login", "", map[string]any{
"account": ordinaryUsername,
"password": password,
}, http.StatusOK, &ordinaryLoginResponse)
var inaccessibleImageResult struct {
Status string `json:"status"`
Code string `json:"code"`
Data []any `json:"data"`
}
doJSON(t, server.URL, http.MethodGet, imageResponse.QueryURL, ordinaryLoginResponse.AccessToken, nil, http.StatusNotFound, &inaccessibleImageResult)
if inaccessibleImageResult.Status != "failed" ||
inaccessibleImageResult.Code != "not_found" ||
len(inaccessibleImageResult.Data) != 0 {
t.Fatalf("EasyAI task result must stay owner-isolated: %+v", inaccessibleImageResult)
}
var imageEditResponse struct {
Task struct {
Status string `json:"status"`
LegacyTaskID string `json:"task_id"`
TaskID string `json:"taskId"`
Task struct {
ID string `json:"id"`
Status string `json:"status"`
Result map[string]any `json:"result"`
@@ -535,6 +611,12 @@ VALUES ($1, 5, '{"purpose":"core-flow"}'::jsonb)`, inviteCode); err != nil {
"simulation": true,
"simulationDurationMs": 5,
}, map[string]string{"X-Async": "true"}, http.StatusAccepted, &imageEditResponse)
if imageEditResponse.Status != "submitted" ||
imageEditResponse.TaskID == "" ||
imageEditResponse.LegacyTaskID != imageEditResponse.TaskID ||
imageEditResponse.Task.ID != imageEditResponse.TaskID {
t.Fatalf("unexpected EasyAI-compatible image edit submission: %+v", imageEditResponse)
}
waitForTaskStatus(t, server.URL, apiKeyResponse.Secret, imageEditResponse.Task.ID, []string{"succeeded"}, 10*time.Second)
doJSON(t, server.URL, http.MethodGet, "/api/v1/tasks/"+imageEditResponse.Task.ID, apiKeyResponse.Secret, nil, http.StatusOK, &imageEditResponse.Task)
if imageEditResponse.Task.Status != "succeeded" || imageEditResponse.Task.Result["id"] == "" {
@@ -1613,6 +1695,64 @@ WHERE m.platform_id = $1::uuid
if !taskListContains(workspaceTaskList.Items, taskResponse.Task.ID) || !taskListContains(workspaceTaskList.Items, pricingTask.Task.ID) {
t.Fatalf("workspace task list should include persisted task records, got %+v", workspaceTaskList.Items)
}
doJSON(t, server.URL, http.MethodGet, "/api/admin/tasks", ordinaryLoginResponse.AccessToken, nil, http.StatusForbidden, nil)
if _, err := testPool.Exec(ctx, `
UPDATE gateway_tasks
SET request = request || '{"diagnostic":{"apiKey":"private-admin-task-test-key"}}'::jsonb
WHERE id = $1::uuid`, taskResponse.Task.ID); err != nil {
t.Fatalf("seed admin task redaction payload: %v", err)
}
var adminTaskList struct {
Items []struct {
ID string `json:"id"`
Request map[string]any `json:"request"`
Context struct {
User struct {
ID string `json:"id"`
Username string `json:"username"`
} `json:"user"`
Tenant struct {
ID string `json:"id"`
Name string `json:"name"`
} `json:"tenant"`
} `json:"adminContext"`
} `json:"items"`
Total int `json:"total"`
}
adminTaskQuery := "/api/admin/tasks?q=" + taskResponse.Task.ID +
"&userId=" + smokeGatewayUserID +
"&status=succeeded" +
"&platformId=" + platform.ID +
"&model=" + defaultTextModel +
"&modelType=text_generate" +
"&runMode=simulation" +
"&billingStatus=settled" +
"&apiKey=smoke" +
"&pageSize=10"
doJSON(t, server.URL, http.MethodGet, adminTaskQuery, loginResponse.AccessToken, nil, http.StatusOK, &adminTaskList)
if len(adminTaskList.Items) != 1 || adminTaskList.Total != 1 || adminTaskList.Items[0].Context.User.Username != username {
t.Fatalf("admin task list should expose the matching task and identity summary: %+v", adminTaskList)
}
maskedDiagnostic, _ := adminTaskList.Items[0].Request["diagnostic"].(map[string]any)
if maskedDiagnostic["apiKey"] != "***" {
t.Fatalf("admin task list must mask nested secrets: %+v", adminTaskList.Items[0].Request)
}
var maskedAdminTask struct {
Request map[string]any `json:"request"`
}
doJSON(t, server.URL, http.MethodGet, "/api/admin/tasks/"+taskResponse.Task.ID, loginResponse.AccessToken, nil, http.StatusOK, &maskedAdminTask)
maskedDiagnostic, _ = maskedAdminTask.Request["diagnostic"].(map[string]any)
if maskedDiagnostic["apiKey"] != "***" {
t.Fatalf("admin task detail must be masked by default: %+v", maskedAdminTask.Request)
}
var fullAdminTask struct {
Request map[string]any `json:"request"`
}
doJSON(t, server.URL, http.MethodGet, "/api/admin/tasks/"+taskResponse.Task.ID+"?sensitive=full", loginResponse.AccessToken, nil, http.StatusOK, &fullAdminTask)
fullDiagnostic, _ := fullAdminTask.Request["diagnostic"].(map[string]any)
if fullDiagnostic["apiKey"] != "private-admin-task-test-key" {
t.Fatalf("explicit full admin task detail should return stored content: %+v", fullAdminTask.Request)
}
req, err := http.NewRequest(http.MethodGet, server.URL+"/api/v1/tasks/"+pricingTask.Task.ID+"/events", nil)
if err != nil {
@@ -1664,6 +1804,7 @@ WHERE m.platform_id = $1::uuid
}
restartModel := "worker-restart-" + suffixText
createSimulationTextBaseModel(t, server.URL, loginResponse.AccessToken, restartModel)
createSimulationTextPlatformModel(
t,
server.URL,
@@ -1782,6 +1923,7 @@ func TestCacheAffinityRoutingFlow(t *testing.T) {
model := "cache-affinity-smoke-" + suffixText
cacheKey := "cache-affinity-key-" + suffixText
createSimulationTextBaseModel(t, server.URL, loginResponse.AccessToken, model)
lowPlatform, lowPlatformModelID := createSimulationTextPlatformModel(t, server.URL, loginResponse.AccessToken, "cache-low-"+suffixText, "Cache Affinity Low Priority", model, 20, map[string]any{
"rules": []map[string]any{
{"metric": "concurrent", "limit": 1, "leaseTtlSeconds": 120},
@@ -1798,8 +1940,31 @@ func TestCacheAffinityRoutingFlow(t *testing.T) {
}
}
highPlatform, _ := createSimulationTextPlatformModel(t, server.URL, loginResponse.AccessToken, "cache-high-"+suffixText, "Cache Affinity High Priority", model, 1, nil)
_ = highPlatform
tools := []any{map[string]any{
"type": "function",
"function": map[string]any{
"name": "lookup_weather",
"description": "Look up current weather",
"parameters": map[string]any{"type": "object", "properties": map[string]any{"city": map[string]any{"type": "string"}}},
},
}}
initialToolMessages := []any{
map[string]any{"role": "system", "content": "Use tools when helpful."},
map[string]any{"role": "user", "content": "Weather in Hangzhou?"},
}
for index := 0; index < 3; index++ {
detail := runCacheAffinityChatTaskWithBody(t, ctx, testPool, server.URL, apiKeyResponse.Secret, model, "tool-prime-"+strconv.Itoa(index)+"-"+suffixText, map[string]any{
"tools": tools,
"messages": initialToolMessages,
"simulationUsage": map[string]any{"inputTokens": 1000, "cachedInputTokens": 750, "outputTokens": 20},
})
if detail.Status != "succeeded" || len(detail.Attempts) != 1 || detail.Attempts[0].PlatformName != "Cache Affinity Low Priority" {
t.Fatalf("tool-history priming request should use the only platform, got %+v", detail)
}
}
peerPlatform, _ := createSimulationTextPlatformModel(t, server.URL, loginResponse.AccessToken, "cache-peer-"+suffixText, "Cache Affinity Peer", model, 20, nil)
_ = peerPlatform
sameKeyDetail := runCacheAffinityChatTask(t, ctx, testPool, server.URL, apiKeyResponse.Secret, model, cacheKey, "same-key-"+suffixText, map[string]any{
"inputTokens": 1000,
"cachedInputTokens": 700,
@@ -1817,12 +1982,36 @@ func TestCacheAffinityRoutingFlow(t *testing.T) {
t.Fatalf("same cache key task detail should expose simulated usage, got %+v", sameKeyDetail.Usage)
}
toolHistoryDetail := runCacheAffinityChatTaskWithBody(t, ctx, testPool, server.URL, apiKeyResponse.Secret, model, "tool-history-"+suffixText, map[string]any{
"tools": tools,
"messages": append(append([]any{}, initialToolMessages...),
map[string]any{"role": "assistant", "tool_calls": []any{map[string]any{
"id": "call_weather",
"type": "function",
"function": map[string]any{
"name": "lookup_weather",
"arguments": `{"city":"Hangzhou"}`,
},
}}},
map[string]any{"role": "tool", "tool_call_id": "call_weather", "content": `{"temperature":31}`},
),
"simulationUsage": map[string]any{"inputTokens": 1400, "cachedInputTokens": 900, "outputTokens": 30},
})
if toolHistoryDetail.Status != "succeeded" || len(toolHistoryDetail.Attempts) != 1 || toolHistoryDetail.Attempts[0].PlatformName != "Cache Affinity Low Priority" {
t.Fatalf("tool continuation should retain the cached platform, got %+v", toolHistoryDetail)
}
if !boolFromTestMap(toolHistoryDetail.Attempts[0].Metrics, "cacheAffinityMatched") ||
intFromTestAny(toolHistoryDetail.Attempts[0].Metrics["cacheAffinityMatchedPrefixDepth"]) < 2 ||
intFromTestAny(toolHistoryDetail.Attempts[0].Metrics["cacheAffinityCandidateCount"]) < 1 {
t.Fatalf("tool continuation should expose prefix and candidate metrics, got %+v", toolHistoryDetail.Attempts[0].Metrics)
}
newKeyDetail := runCacheAffinityChatTask(t, ctx, testPool, server.URL, apiKeyResponse.Secret, model, "cache-affinity-new-"+suffixText, "new-key-"+suffixText, map[string]any{
"inputTokens": 1000,
"cachedInputTokens": 0,
"outputTokens": 20,
})
if newKeyDetail.Status != "succeeded" || len(newKeyDetail.Attempts) != 1 || newKeyDetail.Attempts[0].PlatformName != "Cache Affinity High Priority" {
if newKeyDetail.Status != "succeeded" || len(newKeyDetail.Attempts) != 1 || newKeyDetail.Attempts[0].PlatformName != "Cache Affinity Peer" {
t.Fatalf("new cache key should fall back to base priority, got %+v", newKeyDetail)
}
@@ -1832,9 +2021,12 @@ func TestCacheAffinityRoutingFlow(t *testing.T) {
"cachedInputTokens": 0,
"outputTokens": 20,
})
if fullAvoidedDetail.Status != "succeeded" || len(fullAvoidedDetail.Attempts) != 1 || fullAvoidedDetail.Attempts[0].PlatformName != "Cache Affinity High Priority" {
if fullAvoidedDetail.Status != "succeeded" || len(fullAvoidedDetail.Attempts) != 1 || fullAvoidedDetail.Attempts[0].PlatformName != "Cache Affinity Peer" {
t.Fatalf("full cached platform should be avoided before cache affinity boost, got %+v", fullAvoidedDetail)
}
if fullAvoidedDetail.Attempts[0].Metrics["cacheAffinityOverrideReason"] != "capacity_tier_unavailable" {
t.Fatalf("full cached platform override reason should be visible, got %+v", fullAvoidedDetail.Attempts[0].Metrics)
}
var requestCount int
var inputTokens int
@@ -1844,7 +2036,7 @@ func TestCacheAffinityRoutingFlow(t *testing.T) {
SELECT request_count::int, input_tokens::int, cached_input_tokens::int, ema_hit_ratio::float8
FROM gateway_cache_affinity_stats
WHERE client_id = $1
ORDER BY updated_at DESC
ORDER BY request_count DESC, cached_input_tokens DESC, updated_at DESC
LIMIT 1`, lowPlatform.PlatformKey+":text_generate:"+model).Scan(&requestCount, &inputTokens, &cachedTokens, &emaHitRatio); err != nil {
t.Fatalf("read cache affinity stats: %v", err)
}
@@ -1927,6 +2119,7 @@ func applyMigration(t *testing.T, ctx context.Context, databaseURL string) {
t.Fatalf("connect migration db: %v", err)
}
defer pool.Close()
requireDedicatedIntegrationTestDatabase(t, ctx, pool)
if _, err := pool.Exec(ctx, `
CREATE TABLE IF NOT EXISTS schema_migrations (
version text PRIMARY KEY,
@@ -1947,6 +2140,22 @@ CREATE TABLE IF NOT EXISTS schema_migrations (
if err != nil {
t.Fatalf("read migration %s: %v", filepath.Base(migrationPath), err)
}
migrationSQL := string(migration)
const noTransactionMarker = "-- easyai:migration:no-transaction"
if strings.HasPrefix(strings.TrimSpace(migrationSQL), noTransactionMarker) {
migrationSQL = strings.TrimSpace(strings.TrimPrefix(strings.TrimSpace(migrationSQL), noTransactionMarker))
for _, statement := range strings.Split(migrationSQL, "-- easyai:migration:statement") {
if statement = strings.TrimSpace(statement); statement != "" {
if _, err := pool.Exec(ctx, statement); err != nil {
t.Fatalf("apply non-transaction migration %s: %v", filepath.Base(migrationPath), err)
}
}
}
if _, err := pool.Exec(ctx, "INSERT INTO schema_migrations(version) VALUES($1)", version); err != nil {
t.Fatalf("record non-transaction migration %s: %v", filepath.Base(migrationPath), err)
}
continue
}
tx, err := pool.Begin(ctx)
if err != nil {
t.Fatalf("begin migration %s: %v", filepath.Base(migrationPath), err)
@@ -1974,9 +2183,9 @@ SELECT EXISTS (
FROM base_model_catalog model
WHERE model.provider_key = 'volces'
AND model.canonical_model_key = 'volces:doubao-seedream-5-0-pro-260628'
AND model.invocation_name = 'doubao-seedream-5-0-pro-260628'
AND model.provider_model_name = 'doubao-seedream-5-0-pro-260628'
AND model.display_name = 'Seedream-5.0-Pro'
AND model.display_name !~ '[[:space:]]'
AND model.display_name = 'Seedream 5.0 Pro'
AND model.model_type = '["image_edit","image_generate"]'::jsonb
AND model.capabilities->'image_edit'->>'input_multiple_images' = 'true'
AND model.capabilities->'image_edit'->>'input_max_images_count' = '10'
@@ -1991,7 +2200,7 @@ SELECT EXISTS (
AND model.capabilities->>'stream' = 'false'
AND model.capabilities->>'supportWebSearch' = 'false'
AND model.default_rate_limit_policy->'rules'->0->>'limit' = '500'
AND model.default_snapshot->>'modelAlias' = 'Seedream-5.0-Pro'
AND model.default_snapshot->>'modelAlias' = 'doubao-seedream-5-0-pro-260628'
AND model.default_snapshot->>'displayName' = 'Seedream 5.0 Pro'
)`).Scan(&valid); err != nil {
t.Fatalf("query Seedream 5.0 Pro catalog migration: %v", err)
@@ -2048,16 +2257,17 @@ func doJSONWithHeaders(t *testing.T, baseURL string, method string, path string,
func doAPIV1ChatCompletionAndLoadTask(t *testing.T, ctx context.Context, pool *pgxpool.Pool, baseURL string, token string, payload map[string]any, marker string, expectedStatus int, responseOut any, taskDetailOut any) string {
t.Helper()
_ = ctx
_ = pool
_ = marker
if responseOut == nil {
responseOut = &map[string]any{}
}
requestStartedAt := time.Now().UTC().Add(-time.Second)
responseHeaders := doJSONWithHeaders(t, baseURL, http.MethodPost, "/api/v1/chat/completions", token, payload, nil, expectedStatus, responseOut)
taskID := strings.TrimSpace(responseHeaders.Get("X-Gateway-Task-Id"))
if taskID == "" {
t.Fatal("chat completion response did not expose X-Gateway-Task-Id")
taskID = waitForLatestTaskIDSince(t, ctx, pool, "chat.completions", requestStartedAt, 2*time.Second)
}
if taskID == "" {
t.Fatalf("load chat completion task %s failed: headers=%v body=%+v", marker, responseHeaders, responseOut)
}
if taskDetailOut != nil {
doJSON(t, baseURL, http.MethodGet, "/api/v1/tasks/"+taskID, token, nil, http.StatusOK, taskDetailOut)
@@ -2065,6 +2275,27 @@ func doAPIV1ChatCompletionAndLoadTask(t *testing.T, ctx context.Context, pool *p
return taskID
}
func waitForLatestTaskIDSince(t *testing.T, ctx context.Context, pool *pgxpool.Pool, kind string, since time.Time, timeout time.Duration) string {
t.Helper()
deadline := time.Now().Add(timeout)
for time.Now().Before(deadline) {
var taskID string
err := pool.QueryRow(ctx, `
SELECT id::text
FROM gateway_tasks
WHERE kind = $1
AND created_at >= $2
ORDER BY created_at DESC
LIMIT 1`, kind, since).Scan(&taskID)
if err == nil && taskID != "" {
return taskID
}
time.Sleep(10 * time.Millisecond)
}
t.Fatalf("latest %s task was not persisted within %s", kind, timeout)
return ""
}
type taskWaitDetail struct {
ID string `json:"id"`
Status string `json:"status"`
@@ -2093,6 +2324,20 @@ type cacheAffinityTaskDetail struct {
} `json:"attempts"`
}
func createSimulationTextBaseModel(t *testing.T, baseURL string, adminToken string, model string) {
t.Helper()
doJSON(t, baseURL, http.MethodPost, "/api/admin/catalog/base-models", adminToken, map[string]any{
"providerKey": "openai",
"canonicalModelKey": "openai:" + model,
"invocationName": model,
"providerModelName": model,
"modelType": []string{"text_generate"},
"displayName": model,
"capabilities": map[string]any{"originalTypes": []string{"text_generate"}},
"metadata": map[string]any{"source": "cache-affinity-integration-test"},
}, http.StatusCreated, nil)
}
func createSimulationTextPlatformModel(t *testing.T, baseURL string, adminToken string, platformKey string, platformName string, model string, priority int, rateLimitPolicy map[string]any) (cacheAffinityPlatformFixture, string) {
t.Helper()
var platform cacheAffinityPlatformFixture
@@ -2109,7 +2354,7 @@ func createSimulationTextPlatformModel(t *testing.T, baseURL string, adminToken
t.Fatalf("cache affinity platform was not created: %+v", platform)
}
payload := map[string]any{
"canonicalModelKey": "openai:gpt-4o-mini",
"canonicalModelKey": "openai:" + model,
"modelName": model,
"providerModelName": model,
"modelAlias": model,
@@ -2130,17 +2375,27 @@ func createSimulationTextPlatformModel(t *testing.T, baseURL string, adminToken
}
func runCacheAffinityChatTask(t *testing.T, ctx context.Context, pool *pgxpool.Pool, baseURL string, token string, model string, cacheAffinityKey string, marker string, simulationUsage map[string]any) cacheAffinityTaskDetail {
t.Helper()
return runCacheAffinityChatTaskWithBody(t, ctx, pool, baseURL, token, model, marker, map[string]any{
"cacheAffinityKey": cacheAffinityKey,
"simulationUsage": simulationUsage,
"messages": []any{map[string]any{"role": "user", "content": "cache affinity route"}},
})
}
func runCacheAffinityChatTaskWithBody(t *testing.T, ctx context.Context, pool *pgxpool.Pool, baseURL string, token string, model string, marker string, body map[string]any) cacheAffinityTaskDetail {
t.Helper()
var detail cacheAffinityTaskDetail
doAPIV1ChatCompletionAndLoadTask(t, ctx, pool, baseURL, token, map[string]any{
requestBody := map[string]any{
"model": model,
"runMode": "simulation",
"simulation": true,
"simulationDurationMs": 5,
"cacheAffinityKey": cacheAffinityKey,
"simulationUsage": simulationUsage,
"messages": []map[string]any{{"role": "user", "content": "cache affinity route"}},
}, marker, http.StatusOK, nil, &detail)
}
for key, value := range body {
requestBody[key] = value
}
doAPIV1ChatCompletionAndLoadTask(t, ctx, pool, baseURL, token, requestBody, marker, http.StatusOK, nil, &detail)
return detail
}
+489
View File
@@ -0,0 +1,489 @@
package httpapi
import (
"encoding/json"
"net/http"
"strings"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/runner"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
func easyAIAsyncMediaRequest(kind string, r *http.Request) bool {
if r == nil || !asyncRequest(r) {
return false
}
switch kind {
case "images.generations",
"images.edits",
"images.vectorize",
"videos.generations",
"videos.upscales",
"song.generations",
"music.generations",
"speech.generations",
"voice.clone":
return true
default:
return false
}
}
func easyAITaskAcceptedResponse(task store.GatewayTask) map[string]any {
resultPath := "/api/v1/ai/result/" + task.ID
return map[string]any{
"status": "submitted",
"task_id": task.ID,
"taskId": task.ID,
"query_url": resultPath,
"task": task,
"next": map[string]string{
"events": "/api/v1/tasks/" + task.ID + "/events",
"detail": "/api/v1/tasks/" + task.ID,
"result": resultPath,
},
}
}
func easyAITaskAcceptedResponseWithExtensions(task store.GatewayTask, extensions map[string]any) map[string]any {
response := cloneEasyAIMap(extensions)
for key, value := range easyAITaskAcceptedResponse(task) {
response[key] = value
}
return response
}
func writeEasyAIAsyncError(w http.ResponseWriter, status int, message string, details map[string]any, codes ...string) {
code := ""
if len(codes) > 0 {
code = strings.TrimSpace(codes[0])
}
errorPayload := map[string]any{
"message": message,
"status": status,
}
if code != "" {
errorPayload["code"] = code
}
for key, value := range details {
errorPayload[key] = value
}
response := map[string]any{
"status": "failed",
"message": message,
"data": []any{},
"error": errorPayload,
}
if code != "" {
response["code"] = code
}
writeJSON(w, status, response)
}
func easyAISynchronousTaskResponse(task store.GatewayTask, output map[string]any) map[string]any {
if task.Kind != "images.vectorize" {
return output
}
response := cloneEasyAIMap(output)
response["taskId"] = task.ID
response["task_id"] = task.ID
response["query_url"] = "/api/v1/ai/result/" + task.ID
return response
}
func easyAIFileUploadResponse(upload map[string]any) map[string]any {
response := cloneEasyAIMap(upload)
response["status"] = "success"
response["data"] = cloneEasyAIMap(upload)
return response
}
func easyAITaskResultResponse(task store.GatewayTask) map[string]any {
sourceResult := cloneEasyAIMap(task.Result)
cleanResult := cloneEasyAIMap(sourceResult)
delete(cleanResult, "raw")
delete(cleanResult, "raw_data")
normalizeEasyAIInlineMediaFields(cleanResult)
data := easyAITaskResultData(task, sourceResult)
status := easyAITaskResultStatus(task.Status)
if status == "failed" {
data = []any{}
}
cleanResult["data"] = data
cleanResult["output_content"] = data
response := cloneEasyAIMap(cleanResult)
response["status"] = status
response["task_id"] = task.ID
response["taskId"] = task.ID
response["query_url"] = "/api/v1/ai/result/" + task.ID
response["created"] = task.CreatedAt.UnixMilli()
response["data"] = data
response["output_content"] = data
response["output"] = easyAIOutputURLs(data)
response["result"] = cleanResult
upstreamTaskID := firstNonEmpty(
easyAIString(cleanResult["upstream_task_id"]),
task.RemoteTaskID,
)
if upstreamTaskID != "" {
response["upstream_task_id"] = upstreamTaskID
}
if len(task.Usage) > 0 && response["usage"] == nil {
response["usage"] = task.Usage
}
cancelState := runner.DescribeTaskCancellation(task)
response["cancellable"] = cancelState.Cancellable
response["submitted"] = cancelState.Submitted
message := firstNonEmpty(
easyAIString(cleanResult["message"]),
task.ErrorMessage,
task.Error,
task.Message,
)
if message == "" {
switch status {
case "submitted":
message = "任务已提交,请稍后查看结果"
case "process":
message = "任务处理中"
case "failed":
message = "任务执行失败"
}
}
if message != "" {
response["message"] = message
}
if code := firstNonEmpty(task.ErrorCode, easyAIString(cleanResult["code"])); code != "" {
response["code"] = code
}
return response
}
func easyAITaskResultStatus(status string) string {
switch strings.ToLower(strings.TrimSpace(status)) {
case "succeeded", "success", "completed":
return "success"
case "failed", "cancelled", "canceled", "rejected":
return "failed"
case "running", "processing", "in_progress":
return "process"
default:
return "submitted"
}
}
func easyAITaskResultData(task store.GatewayTask, result map[string]any) []any {
for _, source := range easyAIResultSources(result) {
if items := easyAIOutputArray(source); len(items) > 0 {
return normalizeEasyAIOutputItems(task, items)
}
}
if item := easyAIOutputMap(result); len(item) > 0 {
normalizeEasyAIInlineMediaFields(item)
if easyAIOutputURL(item) != "" ||
easyAIOutputHasInlineMedia(item) ||
(strings.EqualFold(easyAIString(item["type"]), "text") && (item["text"] != nil || item["content"] != nil)) {
return normalizeEasyAIOutputItems(task, []any{item})
}
}
return []any{}
}
func easyAIResultSources(result map[string]any) []any {
sources := []any{
result["data"],
result["output_content"],
easyAIStructuredOutput(result["content"]),
result["images"],
result["videos"],
result["audios"],
result["files"],
result["output"],
}
if raw, ok := result["raw"].(map[string]any); ok {
sources = append(sources,
raw["data"],
raw["output_content"],
easyAIStructuredOutput(raw["content"]),
raw["images"],
raw["videos"],
raw["audios"],
raw["files"],
raw["output"],
)
}
return sources
}
func easyAIStructuredOutput(value any) any {
switch value.(type) {
case []any, []string, map[string]any:
return value
default:
return nil
}
}
func easyAIOutputArray(value any) []any {
switch typed := value.(type) {
case []any:
return typed
case []string:
items := make([]any, 0, len(typed))
for _, item := range typed {
items = append(items, item)
}
return items
case map[string]any, string:
return []any{typed}
default:
return nil
}
}
func normalizeEasyAIOutputItems(task store.GatewayTask, items []any) []any {
normalized := make([]any, 0, len(items))
for _, item := range items {
output := easyAIOutputMap(item)
if len(output) == 0 {
continue
}
delete(output, "raw_data")
normalizeEasyAIInlineMediaFields(output)
mediaURL := easyAIOutputURL(output)
if mediaURL != "" {
output["url"] = mediaURL
}
if strings.TrimSpace(easyAIString(output["type"])) == "" {
if outputType := easyAIOutputType(task, output, mediaURL); outputType != "" {
output["type"] = outputType
}
}
normalized = append(normalized, output)
}
return normalized
}
func easyAIOutputMap(value any) map[string]any {
switch typed := value.(type) {
case map[string]any:
return cloneEasyAIMap(typed)
case string:
trimmed := strings.TrimSpace(typed)
if trimmed == "" {
return nil
}
if strings.HasPrefix(strings.ToLower(trimmed), "data:") {
return map[string]any{"url": trimmed}
}
if strings.HasPrefix(trimmed, "/") ||
strings.HasPrefix(strings.ToLower(trimmed), "http://") ||
strings.HasPrefix(strings.ToLower(trimmed), "https://") {
return map[string]any{"url": trimmed}
}
return map[string]any{"type": "text", "content": trimmed}
default:
return nil
}
}
func easyAIOutputURL(item map[string]any) string {
for _, key := range []string{
"url",
"generated_url",
"generatedUrl",
"output_url",
"outputUrl",
"download_url",
"downloadUrl",
"image_url",
"video_url",
"audio_url",
} {
if value := easyAINestedURL(item[key]); value != "" {
return value
}
}
return ""
}
func easyAINestedURL(value any) string {
var candidate string
switch typed := value.(type) {
case string:
candidate = strings.TrimSpace(typed)
case map[string]any:
candidate = strings.TrimSpace(easyAIString(typed["url"]))
default:
return ""
}
if strings.HasPrefix(strings.ToLower(candidate), "data:") {
return ""
}
return candidate
}
func normalizeEasyAIInlineMediaFields(item map[string]any) {
for _, key := range []string{
"url",
"generated_url",
"generatedUrl",
"output_url",
"outputUrl",
"download_url",
"downloadUrl",
"image_url",
"video_url",
"audio_url",
} {
value := easyAINestedURLAllowData(item[key])
contentType, payload, ok := easyAIBase64DataURL(value)
if !ok {
continue
}
if strings.HasPrefix(contentType, "image/") {
if easyAIString(item["b64_json"]) == "" {
item["b64_json"] = payload
}
} else if easyAIString(item["content"]) == "" {
item["content"] = value
}
if easyAIString(item["mime_type"]) == "" {
item["mime_type"] = contentType
}
delete(item, key)
}
}
func easyAIBase64DataURL(value string) (string, string, bool) {
value = strings.TrimSpace(value)
if !strings.HasPrefix(strings.ToLower(value), "data:") {
return "", "", false
}
meta, payload, ok := strings.Cut(value, ",")
if !ok || strings.TrimSpace(payload) == "" {
return "", "", false
}
parts := strings.Split(strings.TrimSpace(meta[5:]), ";")
if len(parts) < 2 || !strings.EqualFold(strings.TrimSpace(parts[len(parts)-1]), "base64") {
return "", "", false
}
contentType := strings.ToLower(strings.TrimSpace(parts[0]))
if !strings.HasPrefix(contentType, "image/") &&
!strings.HasPrefix(contentType, "audio/") &&
!strings.HasPrefix(contentType, "video/") &&
!strings.HasPrefix(contentType, "application/") {
return "", "", false
}
return contentType, strings.TrimSpace(payload), true
}
func easyAIOutputHasInlineMedia(item map[string]any) bool {
if easyAIString(item["b64_json"]) != "" {
return true
}
content := strings.ToLower(easyAIString(item["content"]))
return strings.HasPrefix(content, "data:") && strings.Contains(content, ";base64,")
}
func easyAINestedURLAllowData(value any) string {
switch typed := value.(type) {
case string:
return strings.TrimSpace(typed)
case map[string]any:
return strings.TrimSpace(easyAIString(typed["url"]))
default:
return ""
}
}
func easyAIOutputType(task store.GatewayTask, item map[string]any, mediaURL string) string {
kind := strings.ToLower(strings.TrimSpace(task.Kind))
switch {
case kind == "images.vectorize":
format := strings.ToLower(firstNonEmpty(
easyAIString(item["format"]),
easyAIString(task.Request["format"]),
))
if format == "svg" || format == "png" {
return "image"
}
if format != "" {
return "file"
}
case strings.HasPrefix(kind, "images."):
return "image"
case strings.HasPrefix(kind, "videos.") || kind == "video.upscale":
return "video"
case kind == "song.generations" || kind == "music.generations" || kind == "speech.generations":
return "audio"
}
lowerURL := strings.ToLower(mediaURL)
if index := strings.IndexAny(lowerURL, "?#"); index >= 0 {
lowerURL = lowerURL[:index]
}
switch {
case strings.HasSuffix(lowerURL, ".svg"),
strings.HasSuffix(lowerURL, ".png"),
strings.HasSuffix(lowerURL, ".jpg"),
strings.HasSuffix(lowerURL, ".jpeg"),
strings.HasSuffix(lowerURL, ".webp"),
strings.HasSuffix(lowerURL, ".gif"):
return "image"
case strings.HasSuffix(lowerURL, ".mp4"),
strings.HasSuffix(lowerURL, ".mov"),
strings.HasSuffix(lowerURL, ".webm"):
return "video"
case strings.HasSuffix(lowerURL, ".mp3"),
strings.HasSuffix(lowerURL, ".wav"),
strings.HasSuffix(lowerURL, ".m4a"),
strings.HasSuffix(lowerURL, ".aac"),
strings.HasSuffix(lowerURL, ".flac"):
return "audio"
case mediaURL != "":
return "file"
case item["text"] != nil || item["content"] != nil:
return "text"
default:
return ""
}
}
func easyAIOutputURLs(data []any) []string {
output := make([]string, 0, len(data))
for _, item := range data {
if typed, ok := item.(map[string]any); ok {
if mediaURL := easyAIOutputURL(typed); mediaURL != "" {
output = append(output, mediaURL)
}
}
}
return output
}
func easyAIString(value any) string {
typed, _ := value.(string)
return strings.TrimSpace(typed)
}
func cloneEasyAIMap(source map[string]any) map[string]any {
if len(source) == 0 {
return map[string]any{}
}
raw, err := json.Marshal(source)
if err != nil {
return map[string]any{}
}
var result map[string]any
if err := json.Unmarshal(raw, &result); err != nil {
return map[string]any{}
}
return result
}
@@ -0,0 +1,311 @@
package httpapi
import (
"encoding/json"
"net/http"
"net/http/httptest"
"testing"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
func TestEasyAITaskAcceptedResponseKeepsGatewayFieldsAndAddsLegacyFields(t *testing.T) {
task := store.GatewayTask{
ID: "task-accepted-1",
Kind: "speech.generations",
Status: "queued",
AsyncMode: true,
}
got := easyAITaskAcceptedResponse(task)
if got["status"] != "submitted" || got["task_id"] != task.ID || got["taskId"] != task.ID {
t.Fatalf("unexpected EasyAI submission identity: %+v", got)
}
if _, ok := got["task"].(store.GatewayTask); !ok {
t.Fatalf("Gateway task envelope was removed: %+v", got)
}
next, _ := got["next"].(map[string]string)
if next["detail"] != "/api/v1/tasks/"+task.ID ||
next["events"] != "/api/v1/tasks/"+task.ID+"/events" ||
next["result"] != "/api/v1/ai/result/"+task.ID {
t.Fatalf("unexpected next links: %+v", next)
}
}
func TestEasyAITaskResultStatus(t *testing.T) {
for _, item := range []struct {
input string
want string
}{
{input: "queued", want: "submitted"},
{input: "pending", want: "submitted"},
{input: "running", want: "process"},
{input: "processing", want: "process"},
{input: "succeeded", want: "success"},
{input: "completed", want: "success"},
{input: "failed", want: "failed"},
{input: "cancelled", want: "failed"},
} {
t.Run(item.input, func(t *testing.T) {
if got := easyAITaskResultStatus(item.input); got != item.want {
t.Fatalf("easyAITaskResultStatus(%q)=%q, want %q", item.input, got, item.want)
}
})
}
}
func TestEasyAITaskResultResponseNormalizesMediaOutputs(t *testing.T) {
createdAt := time.Date(2026, 7, 23, 10, 0, 0, 0, time.UTC)
for _, item := range []struct {
name string
task store.GatewayTask
wantType string
wantURL string
wantCount int
}{
{
name: "image alias",
task: store.GatewayTask{
ID: "image-1", Kind: "images.generations", Status: "succeeded", CreatedAt: createdAt,
Result: map[string]any{"data": []any{map[string]any{
"image_url": "https://example.com/output.png",
}}},
},
wantType: "image", wantURL: "https://example.com/output.png", wantCount: 1,
},
{
name: "video raw content",
task: store.GatewayTask{
ID: "video-1", Kind: "videos.generations", Status: "succeeded", CreatedAt: createdAt,
Result: map[string]any{"raw": map[string]any{"content": map[string]any{
"video_url": "https://example.com/output.mp4",
"duration": 5,
}}},
},
wantType: "video", wantURL: "https://example.com/output.mp4", wantCount: 1,
},
{
name: "audio alias",
task: store.GatewayTask{
ID: "audio-1", Kind: "speech.generations", Status: "succeeded", CreatedAt: createdAt,
Result: map[string]any{"data": []any{map[string]any{
"audio_url": "https://example.com/output.wav",
}}},
},
wantType: "audio", wantURL: "https://example.com/output.wav", wantCount: 1,
},
{
name: "voice clone without media",
task: store.GatewayTask{
ID: "voice-1", Kind: "voice.clone", Status: "succeeded", CreatedAt: createdAt,
Result: map[string]any{"voice_id": "voice-test-1", "cloned_voice": map[string]any{"voiceId": "voice-test-1"}},
},
wantCount: 0,
},
} {
t.Run(item.name, func(t *testing.T) {
got := easyAITaskResultResponse(item.task)
if got["status"] != "success" || got["task_id"] != item.task.ID || got["created"] != createdAt.UnixMilli() {
t.Fatalf("unexpected task result identity: %+v", got)
}
data, _ := got["data"].([]any)
if len(data) != item.wantCount {
t.Fatalf("unexpected output count: got=%d response=%+v", len(data), got)
}
outputContent, _ := got["output_content"].([]any)
if len(outputContent) != item.wantCount {
t.Fatalf("output_content is not synchronized: %+v", got)
}
if item.wantCount == 0 {
if got["voice_id"] != "voice-test-1" {
t.Fatalf("voice clone fields were lost: %+v", got)
}
return
}
output, _ := data[0].(map[string]any)
if output["type"] != item.wantType || output["url"] != item.wantURL {
t.Fatalf("unexpected normalized output: %+v", output)
}
if _, exposed := output["b64_json"]; exposed {
t.Fatalf("base64 output was exposed: %+v", output)
}
urls, _ := got["output"].([]string)
if len(urls) != 1 || urls[0] != item.wantURL {
t.Fatalf("unexpected output URL list: %+v", got["output"])
}
})
}
}
func TestEasyAITaskResultResponsePreservesBase64ImageOutput(t *testing.T) {
base64Payload := "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAusB9Y9ZlB8AAAAASUVORK5CYII="
task := store.GatewayTask{
ID: "image-base64-1", Kind: "images.generations", Status: "succeeded",
Result: map[string]any{"data": []any{map[string]any{
"b64_json": base64Payload,
"mime_type": "image/png",
}}},
}
got := easyAITaskResultResponse(task)
data, _ := got["data"].([]any)
if len(data) != 1 {
t.Fatalf("unexpected base64 output count: %+v", got)
}
output, _ := data[0].(map[string]any)
if output["type"] != "image" || output["b64_json"] != base64Payload || output["mime_type"] != "image/png" {
t.Fatalf("base64 image output was not preserved: %+v", output)
}
if _, exposedAsURL := output["url"]; exposedAsURL {
t.Fatalf("base64 image must not be exposed as a URL: %+v", output)
}
outputContent, _ := got["output_content"].([]any)
content, _ := outputContent[0].(map[string]any)
if content["b64_json"] != base64Payload {
t.Fatalf("output_content lost base64 image: %+v", outputContent)
}
urls, _ := got["output"].([]string)
if len(urls) != 0 {
t.Fatalf("base64-only output must not create URL aliases: %+v", urls)
}
}
func TestEasyAITaskResultResponseMovesImageDataURLToB64JSON(t *testing.T) {
base64Payload := "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAusB9Y9ZlB8AAAAASUVORK5CYII="
task := store.GatewayTask{
ID: "image-data-url-1", Kind: "images.generations", Status: "succeeded",
Result: map[string]any{"data": []any{map[string]any{
"url": "data:image/png;base64," + base64Payload,
}}},
}
got := easyAITaskResultResponse(task)
data, _ := got["data"].([]any)
output, _ := data[0].(map[string]any)
if output["b64_json"] != base64Payload || output["mime_type"] != "image/png" || output["type"] != "image" {
t.Fatalf("image data URL was not normalized: %+v", output)
}
if _, exists := output["url"]; exists {
t.Fatalf("image data URL should move out of the URL field: %+v", output)
}
}
func TestEasyAITaskResultResponsePreservesAudioDataURLAsContent(t *testing.T) {
base64Payload := "SUQzBAAAAAAAI1RTU0UAAAAPAAADTGF2ZjYwLjMuMTAwAAAAAAAAAAAAAAD/"
task := store.GatewayTask{
ID: "audio-data-url-1", Kind: "speech.generations", Status: "succeeded",
Result: map[string]any{"data": []any{map[string]any{
"audio_url": "data:audio/mpeg;base64," + base64Payload,
}}},
}
got := easyAITaskResultResponse(task)
data, _ := got["data"].([]any)
output, _ := data[0].(map[string]any)
if output["content"] != "data:audio/mpeg;base64,"+base64Payload ||
output["mime_type"] != "audio/mpeg" ||
output["type"] != "audio" {
t.Fatalf("audio data URL was not normalized to inline content: %+v", output)
}
if _, exists := output["audio_url"]; exists {
t.Fatalf("audio data URL should move out of the URL field: %+v", output)
}
urls, _ := got["output"].([]string)
if len(urls) != 0 {
t.Fatalf("base64-only audio must not create URL aliases: %+v", urls)
}
}
func TestEasyAITaskResultResponseMapsVectorFormats(t *testing.T) {
for _, item := range []struct {
format string
wantType string
}{
{format: "svg", wantType: "image"},
{format: "png", wantType: "image"},
{format: "pdf", wantType: "file"},
{format: "eps", wantType: "file"},
} {
t.Run(item.format, func(t *testing.T) {
task := store.GatewayTask{
ID: "vector-" + item.format, Kind: "images.vectorize", Status: "succeeded",
Request: map[string]any{"format": item.format},
Result: map[string]any{"data": []any{map[string]any{"url": "https://example.com/output." + item.format}}},
}
got := easyAITaskResultResponse(task)
data, _ := got["data"].([]any)
output, _ := data[0].(map[string]any)
if output["type"] != item.wantType {
t.Fatalf("format %s mapped to %v, want %s", item.format, output["type"], item.wantType)
}
})
}
}
func TestEasyAITaskResultResponseClearsFailedOutputs(t *testing.T) {
task := store.GatewayTask{
ID: "failed-1", Kind: "videos.upscales", Status: "cancelled",
ErrorCode: "cancelled", ErrorMessage: "任务已取消",
Result: map[string]any{"data": []any{map[string]any{
"url": "https://example.com/partial.mp4",
}}},
}
got := easyAITaskResultResponse(task)
data, _ := got["data"].([]any)
outputContent, _ := got["output_content"].([]any)
if got["status"] != "failed" || got["code"] != "cancelled" || got["message"] != "任务已取消" ||
len(data) != 0 || len(outputContent) != 0 {
t.Fatalf("unexpected failed response: %+v", got)
}
}
func TestEasyAISynchronousVectorAndFileUploadResponsesAreAdditive(t *testing.T) {
vectorTask := store.GatewayTask{ID: "vector-sync-1", Kind: "images.vectorize"}
vector := easyAISynchronousTaskResponse(vectorTask, map[string]any{
"status": "success",
"data": []any{map[string]any{"url": "https://example.com/output.svg"}},
})
if vector["status"] != "success" || vector["taskId"] != vectorTask.ID || vector["task_id"] != vectorTask.ID {
t.Fatalf("unexpected vector response: %+v", vector)
}
upload := easyAIFileUploadResponse(map[string]any{
"id": "file-1", "url": "https://example.com/upload.png", "filename": "upload.png",
})
data, _ := upload["data"].(map[string]any)
if upload["status"] != "success" || upload["url"] != data["url"] || upload["id"] != data["id"] {
t.Fatalf("unexpected upload response: %+v", upload)
}
}
func TestEasyAIAsyncMediaRequestRequiresExplicitHeader(t *testing.T) {
syncRequest := httptest.NewRequest(http.MethodPost, "/api/v1/images/generations", nil)
if easyAIAsyncMediaRequest("images.generations", syncRequest) {
t.Fatal("image request without X-Async must stay synchronous")
}
asyncRequest := httptest.NewRequest(http.MethodPost, "/api/v1/images/generations", nil)
asyncRequest.Header.Set("X-Async", "true")
if !easyAIAsyncMediaRequest("images.generations", asyncRequest) {
t.Fatal("image request with X-Async must enable EasyAI async compatibility")
}
if easyAIAsyncMediaRequest("chat.completions", asyncRequest) {
t.Fatal("chat completions must not enter EasyAI media async compatibility")
}
}
func TestWriteEasyAIAsyncErrorKeepsNestedGatewayError(t *testing.T) {
recorder := httptest.NewRecorder()
writeEasyAIAsyncError(recorder, http.StatusBadRequest, "invalid request", map[string]any{"param": "model"}, "invalid_parameter")
var got map[string]any
if err := json.NewDecoder(recorder.Body).Decode(&got); err != nil {
t.Fatal(err)
}
errorPayload, _ := got["error"].(map[string]any)
data, _ := got["data"].([]any)
if got["status"] != "failed" || got["code"] != "invalid_parameter" ||
errorPayload["message"] != "invalid request" || len(data) != 0 {
t.Fatalf("unexpected async error response: %+v", got)
}
}
@@ -0,0 +1,911 @@
package httpapi
import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"log/slog"
"net/http"
"net/http/httptest"
"os"
"strconv"
"strings"
"sync"
"sync/atomic"
"testing"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
"github.com/jackc/pgx/v5/pgxpool"
)
type failurePolicyAcceptanceFixture struct {
ctx context.Context
cancel context.CancelFunc
db *store.Store
pool *pgxpool.Pool
server *httptest.Server
adminToken string
apiKey string
suffix string
}
type controlledProviderStub struct {
server *httptest.Server
calls atomic.Int64
}
type acceptancePlatform struct {
ID string `json:"id"`
}
type acceptancePlatformModel struct {
ID string `json:"id"`
}
type acceptanceTaskDetail struct {
ID string `json:"id"`
Status string `json:"status"`
Attempts []struct {
AttemptNo int `json:"attemptNo"`
PlatformID string `json:"platformId"`
PlatformName string `json:"platformName"`
PlatformModel string `json:"platformModelId"`
Status string `json:"status"`
ErrorCode string `json:"errorCode"`
DynamicMetrics map[string]any `json:"metrics"`
Trace []map[string]any `json:"trace"`
} `json:"attempts"`
}
func TestFailurePolicyHTTPAcceptance(t *testing.T) {
fixture := newFailurePolicyAcceptanceFixture(t)
defer fixture.close()
originalPolicy, err := fixture.db.GetActiveRunnerPolicy(fixture.ctx)
if err != nil {
t.Fatalf("read original runner policy: %v", err)
}
t.Cleanup(func() {
_, _ = fixture.db.UpsertDefaultRunnerPolicy(context.Background(), store.RunnerPolicyInput{
PolicyKey: originalPolicy.PolicyKey,
Name: originalPolicy.Name,
Description: originalPolicy.Description,
FailoverPolicy: originalPolicy.FailoverPolicy,
HardStopPolicy: originalPolicy.HardStopPolicy,
PriorityDemotePolicy: originalPolicy.PriorityDemotePolicy,
SingleSourcePolicy: originalPolicy.SingleSourcePolicy,
CacheAffinityPolicy: originalPolicy.CacheAffinityPolicy,
Metadata: originalPolicy.Metadata,
Status: originalPolicy.Status,
})
})
fixture.setRunnerPolicy(t, true)
t.Run("Gemini 429 后轮转且只冷却模型一次", fixture.testGemini429ImageFailover)
t.Run("5xx 同平台重试耗尽后降权并轮转", fixture.testServerErrorRetryThenFailover)
t.Run("401 立即禁用平台并轮转", fixture.testAuthFailureDisableThenFailover)
t.Run("400 参数错误硬停止", fixture.testBadRequestStops)
t.Run("旧 degradePolicy 仅转换一次", fixture.testLegacyDegradeDoesNotDoubleApply)
t.Run("并发单源保护至少保留一个来源", fixture.testConcurrentSingleSourceProtection)
t.Run("全部冷却同步返回脱敏 429", fixture.testAllCoolingSynchronousContract)
t.Run("全部冷却异步零 attempt 重排并恢复", fixture.testAllCoolingAsyncRecovery)
t.Run("流式首包边界控制轮转", fixture.testStreamingBoundary)
t.Run("管理接口严格校验且不覆盖原策略", fixture.testRunnerPolicyHTTPValidation)
}
func newFailurePolicyAcceptanceFixture(t *testing.T) *failurePolicyAcceptanceFixture {
t.Helper()
databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL"))
if databaseURL == "" {
t.Skip("set AI_GATEWAY_TEST_DATABASE_URL to run failure-policy HTTP acceptance")
}
ctx, cancel := context.WithCancel(context.Background())
applyMigration(t, ctx, databaseURL)
db, err := store.Connect(ctx, databaseURL)
if err != nil {
cancel()
t.Fatalf("connect acceptance store: %v", err)
}
pool, err := pgxpool.New(ctx, databaseURL)
if err != nil {
db.Close()
cancel()
t.Fatalf("connect acceptance pool: %v", err)
}
handler := NewServerWithContext(ctx, config.Config{
AppEnv: "test",
HTTPAddr: ":0",
DatabaseURL: databaseURL,
IdentityMode: "hybrid",
JWTSecret: "failure-policy-http-acceptance-secret",
BillingEngineMode: "observe",
CORSAllowedOrigin: "*",
}, db, slog.New(slog.NewTextHandler(io.Discard, nil)))
server := httptest.NewServer(handler)
fixture := &failurePolicyAcceptanceFixture{
ctx: ctx,
cancel: cancel,
db: db,
pool: pool,
server: server,
suffix: strconv.FormatInt(time.Now().UnixNano(), 10),
}
username := "failure_acceptance_" + fixture.suffix
password := "password123"
doJSON(t, server.URL, http.MethodPost, "/api/v1/auth/register", "", map[string]any{
"username": username,
"email": username + "@example.com",
"password": password,
}, http.StatusCreated, &map[string]any{})
if _, err := pool.Exec(ctx, `UPDATE gateway_users SET roles = '["admin"]'::jsonb WHERE username = $1`, username); err != nil {
fixture.close()
t.Fatalf("promote acceptance admin: %v", err)
}
var login struct {
AccessToken string `json:"accessToken"`
}
doJSON(t, server.URL, http.MethodPost, "/api/v1/auth/login", "", map[string]any{
"account": username, "password": password,
}, http.StatusOK, &login)
fixture.adminToken = login.AccessToken
var gatewayUserID string
if err := pool.QueryRow(ctx, `SELECT id::text FROM gateway_users WHERE username=$1`, username).Scan(&gatewayUserID); err != nil {
fixture.close()
t.Fatalf("read acceptance user: %v", err)
}
doJSON(t, server.URL, http.MethodPatch, "/api/admin/users/"+gatewayUserID+"/wallet", fixture.adminToken, map[string]any{
"currency": "resource",
"balance": 1000000,
"reason": "seed isolated failure-policy acceptance wallet",
}, http.StatusOK, &map[string]any{})
var key struct {
Secret string `json:"secret"`
}
doJSON(t, server.URL, http.MethodPost, "/api/v1/api-keys", fixture.adminToken, map[string]any{
"name": "failure policy acceptance",
}, http.StatusCreated, &key)
fixture.apiKey = key.Secret
return fixture
}
func (f *failurePolicyAcceptanceFixture) close() {
if f.server != nil {
f.server.Close()
}
if f.cancel != nil {
f.cancel()
}
if f.pool != nil {
f.pool.Close()
}
if f.db != nil {
f.db.Close()
}
}
func (f *failurePolicyAcceptanceFixture) setRunnerPolicy(t *testing.T, singleSource bool) {
t.Helper()
doJSON(t, f.server.URL, http.MethodPatch, "/api/admin/runtime/runner-policy", f.adminToken, standardFailureRunnerPolicy(singleSource), http.StatusOK, &map[string]any{})
}
func standardFailureRunnerPolicy(singleSource bool) map[string]any {
return map[string]any{
"policyKey": "default-runner-v1",
"name": "失败策略 HTTP 验收",
"status": "active",
"description": "真实 provider HTTP client 的失败策略验收",
"failoverPolicy": map[string]any{
"enabled": true,
"maxPlatforms": 99,
"maxDurationSeconds": 600,
"actionRules": []any{
map[string]any{"categories": []any{"rate_limit"}, "statusCodes": []any{429}, "action": "cooldown_and_next", "target": "model", "cooldownSeconds": 30},
map[string]any{"categories": []any{"auth_error"}, "statusCodes": []any{401, 403}, "action": "disable_and_next", "target": "platform"},
map[string]any{"categories": []any{"timeout", "provider_5xx", "provider_overloaded"}, "statusCodes": []any{408, 500, 502, 503, 504}, "action": "demote_and_next", "target": "platform", "demoteSteps": 1},
map[string]any{"categories": []any{"network", "stream_error"}, "action": "next"},
map[string]any{"categories": []any{"request_error", "unsupported_model", "user_permission", "insufficient_balance"}, "statusCodes": []any{400, 404, 409, 422}, "action": "stop"},
},
},
"singleSourcePolicy": map[string]any{"enabled": singleSource},
"cacheAffinityPolicy": map[string]any{"enabled": false},
}
}
func (f *failurePolicyAcceptanceFixture) createBaseModel(t *testing.T, model string, modelType string) string {
t.Helper()
var created struct {
ID string `json:"id"`
}
doJSON(t, f.server.URL, http.MethodPost, "/api/admin/catalog/base-models", f.adminToken, map[string]any{
"providerKey": "openai",
"canonicalModelKey": "openai:" + model,
"invocationName": model,
"providerModelName": model,
"modelType": []string{modelType},
"displayName": model,
"status": "active",
"capabilities": map[string]any{},
"baseBillingConfig": map[string]any{},
}, http.StatusCreated, &created)
return created.ID
}
func (f *failurePolicyAcceptanceFixture) createPlatformModel(
t *testing.T,
name string,
provider string,
baseURL string,
priority int,
baseModelID string,
providerModelName string,
modelType string,
retryAttempts int,
runtimePolicySetID string,
runtimeOverride map[string]any,
) (acceptancePlatform, acceptancePlatformModel) {
t.Helper()
var platform acceptancePlatform
doJSON(t, f.server.URL, http.MethodPost, "/api/admin/platforms", f.adminToken, map[string]any{
"provider": provider,
"platformKey": "acceptance-" + name + "-" + f.suffix,
"name": "acceptance-" + name,
"baseUrl": strings.TrimRight(baseURL, "/"),
"authType": "bearer",
"credentials": map[string]any{"apiKey": "controlled-upstream-test-key"},
"priority": priority,
"status": "enabled",
}, http.StatusCreated, &platform)
var model acceptancePlatformModel
payload := map[string]any{
"baseModelId": baseModelID,
"providerModelName": providerModelName,
"modelType": []string{modelType},
"retryPolicy": map[string]any{"enabled": true, "maxAttempts": retryAttempts},
}
if runtimePolicySetID != "" {
payload["runtimePolicySetId"] = runtimePolicySetID
}
if len(runtimeOverride) > 0 {
payload["runtimePolicyOverride"] = runtimeOverride
}
doJSON(t, f.server.URL, http.MethodPost, "/api/admin/platforms/"+platform.ID+"/models", f.adminToken, payload, http.StatusCreated, &model)
return platform, model
}
func newControlledProviderStub(handler func(call int, w http.ResponseWriter, r *http.Request)) *controlledProviderStub {
stub := &controlledProviderStub{}
stub.server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
call := int(stub.calls.Add(1))
handler(call, w, r)
}))
return stub
}
func (s *controlledProviderStub) close() {
if s != nil && s.server != nil {
s.server.Close()
}
}
func writeStubJSON(w http.ResponseWriter, status int, payload any) {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(status)
_ = json.NewEncoder(w).Encode(payload)
}
func writeOpenAIChatSuccess(w http.ResponseWriter, content string) {
writeStubJSON(w, http.StatusOK, map[string]any{
"id": "chatcmpl-controlled", "object": "chat.completion", "created": time.Now().Unix(), "model": "controlled-model",
"choices": []any{map[string]any{
"index": 0, "message": map[string]any{"role": "assistant", "content": content}, "finish_reason": "stop",
}},
"usage": map[string]any{"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
})
}
func writeOpenAIImageSuccess(w http.ResponseWriter) {
writeStubJSON(w, http.StatusOK, map[string]any{
"created": time.Now().Unix(),
"data": []any{map[string]any{"url": "https://example.invalid/controlled-test.png"}},
})
}
func (f *failurePolicyAcceptanceFixture) postJSON(t *testing.T, path string, payload any, headers map[string]string) (int, http.Header, []byte) {
t.Helper()
raw, err := json.Marshal(payload)
if err != nil {
t.Fatalf("marshal acceptance payload: %v", err)
}
request, err := http.NewRequest(http.MethodPost, f.server.URL+path, bytes.NewReader(raw))
if err != nil {
t.Fatalf("build acceptance request: %v", err)
}
request.Header.Set("Authorization", "Bearer "+f.apiKey)
request.Header.Set("Content-Type", "application/json")
for key, value := range headers {
request.Header.Set(key, value)
}
response, err := http.DefaultClient.Do(request)
if err != nil {
t.Fatalf("acceptance request %s: %v", path, err)
}
defer response.Body.Close()
body, _ := io.ReadAll(response.Body)
return response.StatusCode, response.Header.Clone(), body
}
func (f *failurePolicyAcceptanceFixture) taskIDForModel(t *testing.T, model string, startedAt time.Time) string {
t.Helper()
var taskID string
deadline := time.Now().Add(3 * time.Second)
for time.Now().Before(deadline) {
err := f.pool.QueryRow(f.ctx, `
SELECT id::text
FROM gateway_tasks
WHERE model = $1
AND created_at >= $2
ORDER BY created_at DESC
LIMIT 1`, model, startedAt.Add(-time.Second)).Scan(&taskID)
if err == nil && taskID != "" {
return taskID
}
time.Sleep(10 * time.Millisecond)
}
t.Fatalf("task for model %q was not persisted", model)
return ""
}
func (f *failurePolicyAcceptanceFixture) loadTask(t *testing.T, taskID string) acceptanceTaskDetail {
t.Helper()
var detail acceptanceTaskDetail
doJSON(t, f.server.URL, http.MethodGet, "/api/v1/tasks/"+taskID, f.apiKey, nil, http.StatusOK, &detail)
return detail
}
func countFailureEffects(detail acceptanceTaskDetail, effect string) int {
count := 0
for _, attempt := range detail.Attempts {
trace := attempt.Trace
for _, raw := range anySlice(attempt.DynamicMetrics["trace"]) {
if entry, ok := raw.(map[string]any); ok {
trace = append(trace, entry)
}
}
for _, entry := range trace {
if entry["event"] == "failure_decision" && entry["effect"] == effect {
count++
}
}
}
return count
}
func anySlice(value any) []any {
items, _ := value.([]any)
return items
}
func (f *failurePolicyAcceptanceFixture) testGemini429ImageFailover(t *testing.T) {
model := "acceptance-gemini-image-" + f.suffix
baseModelID := f.createBaseModel(t, model, "image_generate")
gemini := newControlledProviderStub(func(_ int, w http.ResponseWriter, _ *http.Request) {
writeStubJSON(w, http.StatusTooManyRequests, map[string]any{
"error": map[string]any{"code": 429, "message": "RESOURCE_EXHAUSTED controlled failure", "status": "RESOURCE_EXHAUSTED"},
})
})
defer gemini.close()
openai := newControlledProviderStub(func(_ int, w http.ResponseWriter, _ *http.Request) { writeOpenAIImageSuccess(w) })
defer openai.close()
platformA, modelA := f.createPlatformModel(t, "gemini-image-a", "gemini", gemini.server.URL, 10, baseModelID, "gemini-image-controlled", "image_generate", 3, "", nil)
platformB, _ := f.createPlatformModel(t, "openai-image-b", "openai", openai.server.URL+"/v1", 20, baseModelID, "openai-image-controlled", "image_generate", 1, "", nil)
startedAt := time.Now()
status, _, body := f.postJSON(t, "/api/v1/images/generations", map[string]any{"model": model, "prompt": "controlled image"}, nil)
if status != http.StatusOK {
t.Fatalf("status=%d body=%s", status, body)
}
if gemini.calls.Load() != 1 || openai.calls.Load() != 1 {
t.Fatalf("upstream calls gemini=%d openai=%d, want 1/1", gemini.calls.Load(), openai.calls.Load())
}
var modelCooling, platformCooling bool
if err := f.pool.QueryRow(f.ctx, `
SELECT COALESCE(model.cooldown_until > now(), false), COALESCE(platform.cooldown_until > now(), false)
FROM platform_models model
JOIN integration_platforms platform ON platform.id = model.platform_id
WHERE model.id = $1::uuid`, modelA.ID).Scan(&modelCooling, &platformCooling); err != nil {
t.Fatalf("read Gemini cooldown state: %v", err)
}
if !modelCooling || platformCooling {
t.Fatalf("cooldown state model=%v platform=%v, want true/false for platform %s", modelCooling, platformCooling, platformA.ID)
}
detail := f.loadTask(t, f.taskIDForModel(t, model, startedAt))
if len(detail.Attempts) != 2 || detail.Attempts[0].PlatformID != platformA.ID || detail.Attempts[1].PlatformID != platformB.ID || countFailureEffects(detail, "cooldown") != 1 {
t.Fatalf("unexpected task attempts/trace: %+v", detail)
}
}
func (f *failurePolicyAcceptanceFixture) testServerErrorRetryThenFailover(t *testing.T) {
model := "acceptance-5xx-" + f.suffix
baseModelID := f.createBaseModel(t, model, "text_generate")
failed := newControlledProviderStub(func(_ int, w http.ResponseWriter, _ *http.Request) {
writeStubJSON(w, http.StatusInternalServerError, map[string]any{"error": map[string]any{"message": "controlled 500"}})
})
defer failed.close()
success := newControlledProviderStub(func(_ int, w http.ResponseWriter, _ *http.Request) { writeOpenAIChatSuccess(w, "fallback ok") })
defer success.close()
platformA, _ := f.createPlatformModel(t, "5xx-a", "openai", failed.server.URL+"/v1", 10, baseModelID, "controlled-5xx", "text_generate", 2, "", nil)
platformB, _ := f.createPlatformModel(t, "5xx-b", "openai", success.server.URL+"/v1", 20, baseModelID, "controlled-success", "text_generate", 1, "", nil)
startedAt := time.Now()
status, _, body := f.postJSON(t, "/api/v1/chat/completions", map[string]any{
"model": model, "messages": []any{map[string]any{"role": "user", "content": "retry"}},
}, nil)
if status != http.StatusOK {
t.Fatalf("status=%d body=%s", status, body)
}
if failed.calls.Load() != 2 || success.calls.Load() != 1 {
t.Fatalf("upstream calls failed=%d success=%d, want 2/1", failed.calls.Load(), success.calls.Load())
}
detail := f.loadTask(t, f.taskIDForModel(t, model, startedAt))
if len(detail.Attempts) != 3 ||
detail.Attempts[0].PlatformID != platformA.ID ||
detail.Attempts[1].PlatformID != platformA.ID ||
detail.Attempts[2].PlatformID != platformB.ID ||
countFailureEffects(detail, "demote") != 1 {
t.Fatalf("unexpected retry/failover chain: %+v", detail)
}
var dynamicPriority *int
if err := f.pool.QueryRow(f.ctx, `SELECT dynamic_priority FROM integration_platforms WHERE id=$1::uuid`, platformA.ID).Scan(&dynamicPriority); err != nil {
t.Fatalf("read demoted priority: %v", err)
}
if dynamicPriority == nil {
t.Fatal("failed platform should be demoted exactly once after retry exhaustion")
}
}
func (f *failurePolicyAcceptanceFixture) testAuthFailureDisableThenFailover(t *testing.T) {
model := "acceptance-auth-" + f.suffix
baseModelID := f.createBaseModel(t, model, "text_generate")
failed := newControlledProviderStub(func(_ int, w http.ResponseWriter, _ *http.Request) {
writeStubJSON(w, http.StatusUnauthorized, map[string]any{"error": map[string]any{"message": "invalid api key"}})
})
defer failed.close()
success := newControlledProviderStub(func(_ int, w http.ResponseWriter, _ *http.Request) { writeOpenAIChatSuccess(w, "auth fallback") })
defer success.close()
platformA, _ := f.createPlatformModel(t, "auth-a", "openai", failed.server.URL+"/v1", 10, baseModelID, "controlled-auth", "text_generate", 3, "", nil)
_, _ = f.createPlatformModel(t, "auth-b", "openai", success.server.URL+"/v1", 20, baseModelID, "controlled-success", "text_generate", 1, "", nil)
startedAt := time.Now()
status, _, body := f.postJSON(t, "/api/v1/chat/completions", map[string]any{
"model": model, "messages": []any{map[string]any{"role": "user", "content": "auth"}},
}, nil)
if status != http.StatusOK {
t.Fatalf("status=%d body=%s", status, body)
}
if failed.calls.Load() != 1 || success.calls.Load() != 1 {
t.Fatalf("upstream calls failed=%d success=%d, want 1/1", failed.calls.Load(), success.calls.Load())
}
var statusValue string
if err := f.pool.QueryRow(f.ctx, `SELECT status FROM integration_platforms WHERE id=$1::uuid`, platformA.ID).Scan(&statusValue); err != nil {
t.Fatalf("read disabled platform: %v", err)
}
detail := f.loadTask(t, f.taskIDForModel(t, model, startedAt))
if statusValue != "disabled" || countFailureEffects(detail, "disable") != 1 || len(detail.Attempts) != 2 {
t.Fatalf("unexpected auth disable result status=%s detail=%+v", statusValue, detail)
}
}
func (f *failurePolicyAcceptanceFixture) testBadRequestStops(t *testing.T) {
model := "acceptance-400-" + f.suffix
baseModelID := f.createBaseModel(t, model, "text_generate")
failed := newControlledProviderStub(func(_ int, w http.ResponseWriter, _ *http.Request) {
writeStubJSON(w, http.StatusBadRequest, map[string]any{"error": map[string]any{"message": "controlled invalid parameter", "code": "invalid_parameter"}})
})
defer failed.close()
success := newControlledProviderStub(func(_ int, w http.ResponseWriter, _ *http.Request) { writeOpenAIChatSuccess(w, "must not run") })
defer success.close()
platformA, modelA := f.createPlatformModel(t, "400-a", "openai", failed.server.URL+"/v1", 10, baseModelID, "controlled-bad", "text_generate", 3, "", nil)
_, _ = f.createPlatformModel(t, "400-b", "openai", success.server.URL+"/v1", 20, baseModelID, "controlled-success", "text_generate", 1, "", nil)
status, _, body := f.postJSON(t, "/api/v1/chat/completions", map[string]any{
"model": model, "messages": []any{map[string]any{"role": "user", "content": "bad"}},
}, nil)
if status != http.StatusBadRequest || success.calls.Load() != 0 || failed.calls.Load() != 1 {
t.Fatalf("status=%d calls=%d/%d body=%s", status, failed.calls.Load(), success.calls.Load(), body)
}
var platformStatus string
var modelEnabled bool
var platformCooldown, modelCooldown *time.Time
if err := f.pool.QueryRow(f.ctx, `
SELECT platform.status, model.enabled, platform.cooldown_until, model.cooldown_until
FROM integration_platforms platform
JOIN platform_models model ON model.platform_id=platform.id
WHERE platform.id=$1::uuid AND model.id=$2::uuid`, platformA.ID, modelA.ID).Scan(&platformStatus, &modelEnabled, &platformCooldown, &modelCooldown); err != nil {
t.Fatalf("read hard-stop state: %v", err)
}
if platformStatus != "enabled" || !modelEnabled || platformCooldown != nil || modelCooldown != nil {
t.Fatalf("hard stop mutated runtime state: status=%s enabled=%v platformCooldown=%v modelCooldown=%v", platformStatus, modelEnabled, platformCooldown, modelCooldown)
}
}
func (f *failurePolicyAcceptanceFixture) testLegacyDegradeDoesNotDoubleApply(t *testing.T) {
model := "acceptance-legacy-" + f.suffix
baseModelID := f.createBaseModel(t, model, "text_generate")
var policySet struct {
ID string `json:"id"`
}
doJSON(t, f.server.URL, http.MethodPost, "/api/admin/runtime/policy-sets", f.adminToken, map[string]any{
"policyKey": "acceptance-legacy-" + f.suffix,
"name": "acceptance legacy degrade",
"degradePolicy": map[string]any{
"enabled": true, "keywords": []any{"rate_limit"}, "cooldownSeconds": 45,
},
}, http.StatusCreated, &policySet)
failed := newControlledProviderStub(func(_ int, w http.ResponseWriter, _ *http.Request) {
writeStubJSON(w, http.StatusTooManyRequests, map[string]any{"error": map[string]any{"message": "rate_limit controlled"}})
})
defer failed.close()
success := newControlledProviderStub(func(_ int, w http.ResponseWriter, _ *http.Request) { writeOpenAIChatSuccess(w, "legacy fallback") })
defer success.close()
_, modelA := f.createPlatformModel(t, "legacy-a", "openai", failed.server.URL+"/v1", 10, baseModelID, "legacy-failed", "text_generate", 3, policySet.ID, nil)
_, _ = f.createPlatformModel(t, "legacy-b", "openai", success.server.URL+"/v1", 20, baseModelID, "legacy-success", "text_generate", 1, "", nil)
startedAt := time.Now()
status, _, body := f.postJSON(t, "/api/v1/chat/completions", map[string]any{
"model": model, "messages": []any{map[string]any{"role": "user", "content": "legacy"}},
}, nil)
if status != http.StatusOK || failed.calls.Load() != 1 || success.calls.Load() != 1 {
t.Fatalf("status=%d calls=%d/%d body=%s", status, failed.calls.Load(), success.calls.Load(), body)
}
detail := f.loadTask(t, f.taskIDForModel(t, model, startedAt))
if countFailureEffects(detail, "cooldown") != 1 {
t.Fatalf("legacy and global policy must not both execute: %+v", detail)
}
var remaining float64
if err := f.pool.QueryRow(f.ctx, `SELECT EXTRACT(EPOCH FROM cooldown_until-now())::float8 FROM platform_models WHERE id=$1::uuid`, modelA.ID).Scan(&remaining); err != nil {
t.Fatalf("read legacy cooldown: %v", err)
}
if remaining < 35 || remaining > 50 {
t.Fatalf("legacy cooldown should be applied once with 45 seconds, remaining=%f", remaining)
}
}
func (f *failurePolicyAcceptanceFixture) testConcurrentSingleSourceProtection(t *testing.T) {
f.setRunnerPolicy(t, true)
model := "acceptance-concurrent-" + f.suffix
baseModelID := f.createBaseModel(t, model, "text_generate")
stubs := make([]*controlledProviderStub, 0, 3)
modelIDs := make([]string, 0, 3)
for index := 0; index < 3; index++ {
stub := newControlledProviderStub(func(_ int, w http.ResponseWriter, _ *http.Request) {
writeStubJSON(w, http.StatusTooManyRequests, map[string]any{"error": map[string]any{"message": "controlled 429"}})
})
stubs = append(stubs, stub)
_, platformModel := f.createPlatformModel(t, fmt.Sprintf("concurrent-%d", index), "openai", stub.server.URL+"/v1", 10+index*10, baseModelID, fmt.Sprintf("concurrent-%d", index), "text_generate", 1, "", nil)
modelIDs = append(modelIDs, platformModel.ID)
}
defer func() {
for _, stub := range stubs {
stub.close()
}
}()
var wait sync.WaitGroup
errorsFound := make(chan error, 12)
for index := 0; index < 12; index++ {
wait.Add(1)
go func() {
defer wait.Done()
status, _, body, err := f.postJSONNoFail("/api/v1/chat/completions", map[string]any{
"model": model, "messages": []any{map[string]any{"role": "user", "content": "concurrent"}},
})
if err != nil {
errorsFound <- err
return
}
if status != http.StatusTooManyRequests {
errorsFound <- fmt.Errorf("status=%d body=%s", status, body)
}
}()
}
wait.Wait()
close(errorsFound)
for err := range errorsFound {
t.Error(err)
}
var available int
if err := f.pool.QueryRow(f.ctx, `
SELECT count(*)
FROM platform_models model
JOIN integration_platforms platform ON platform.id=model.platform_id
WHERE model.id = ANY($1::uuid[])
AND model.enabled=true
AND platform.status='enabled'
AND (model.cooldown_until IS NULL OR model.cooldown_until <= now())
AND (platform.cooldown_until IS NULL OR platform.cooldown_until <= now())`, modelIDs).Scan(&available); err != nil {
t.Fatalf("count available sources: %v", err)
}
if available < 1 {
t.Fatal("concurrent cooldown decisions disabled or cooled every source")
}
before := make([]int64, len(stubs))
for index, stub := range stubs {
before[index] = stub.calls.Load()
}
status, _, _, err := f.postJSONNoFail("/api/v1/chat/completions", map[string]any{
"model": model, "messages": []any{map[string]any{"role": "user", "content": "post-cooldown"}},
})
if err != nil || status != http.StatusTooManyRequests {
t.Fatalf("post-cooldown request status=%d err=%v", status, err)
}
for index, modelID := range modelIDs {
var cooled bool
if err := f.pool.QueryRow(f.ctx, `SELECT COALESCE(cooldown_until > now(), false) FROM platform_models WHERE id=$1::uuid`, modelID).Scan(&cooled); err != nil {
t.Fatalf("read source cooldown: %v", err)
}
if cooled && stubs[index].calls.Load() != before[index] {
t.Fatalf("cooled source %d was called again: before=%d after=%d", index, before[index], stubs[index].calls.Load())
}
}
}
func (f *failurePolicyAcceptanceFixture) postJSONNoFail(path string, payload any) (int, http.Header, []byte, error) {
raw, err := json.Marshal(payload)
if err != nil {
return 0, nil, nil, err
}
request, err := http.NewRequest(http.MethodPost, f.server.URL+path, bytes.NewReader(raw))
if err != nil {
return 0, nil, nil, err
}
request.Header.Set("Authorization", "Bearer "+f.apiKey)
request.Header.Set("Content-Type", "application/json")
response, err := http.DefaultClient.Do(request)
if err != nil {
return 0, nil, nil, err
}
defer response.Body.Close()
body, err := io.ReadAll(response.Body)
return response.StatusCode, response.Header.Clone(), body, err
}
func (f *failurePolicyAcceptanceFixture) testAllCoolingSynchronousContract(t *testing.T) {
f.setRunnerPolicy(t, false)
defer f.setRunnerPolicy(t, true)
model := "acceptance-all-cooling-" + f.suffix
baseModelID := f.createBaseModel(t, model, "image_generate")
stubA := newControlledProviderStub(func(_ int, w http.ResponseWriter, _ *http.Request) {
writeStubJSON(w, http.StatusTooManyRequests, map[string]any{"error": map[string]any{"message": "private-a-raw-error"}})
})
defer stubA.close()
stubB := newControlledProviderStub(func(_ int, w http.ResponseWriter, _ *http.Request) {
writeStubJSON(w, http.StatusTooManyRequests, map[string]any{"error": map[string]any{"message": "private-b-raw-error"}})
})
defer stubB.close()
override := func(seconds int) map[string]any {
return map[string]any{"failoverPolicy": map[string]any{"actionRules": []any{
map[string]any{"statusCodes": []any{429}, "action": "cooldown_and_next", "target": "model", "cooldownSeconds": seconds},
}}}
}
_, modelA := f.createPlatformModel(t, "cooling-a", "openai", stubA.server.URL+"/v1", 10, baseModelID, "cooling-a", "image_generate", 3, "", override(30))
_, modelB := f.createPlatformModel(t, "cooling-b", "openai", stubB.server.URL+"/v1", 20, baseModelID, "cooling-b", "image_generate", 3, "", override(60))
defer func() {
_, _ = f.pool.Exec(context.Background(), `UPDATE platform_models SET cooldown_until=NULL WHERE id IN ($1::uuid,$2::uuid)`, modelA.ID, modelB.ID)
}()
status, _, _ := f.postJSON(t, "/api/v1/images/generations", map[string]any{"model": model, "prompt": "cool both"}, nil)
if status != http.StatusTooManyRequests || stubA.calls.Load() != 1 || stubB.calls.Load() != 1 {
t.Fatalf("initial cooling request status=%d calls=%d/%d", status, stubA.calls.Load(), stubB.calls.Load())
}
startedAt := time.Now()
status, headers, body := f.postJSON(t, "/api/v1/images/generations", map[string]any{"model": model, "prompt": "all cooling"}, nil)
if status != http.StatusTooManyRequests {
t.Fatalf("cooldown contract status=%d body=%s", status, body)
}
retryAfter, err := strconv.Atoi(headers.Get("Retry-After"))
if err != nil || retryAfter < 20 || retryAfter > 31 {
t.Fatalf("Retry-After=%q, want earliest recovery near 30 seconds", headers.Get("Retry-After"))
}
bodyText := string(body)
for _, forbidden := range []string{"acceptance-cooling-a", "acceptance-cooling-b", stubA.server.URL, stubB.server.URL, "private-a-raw-error", "private-b-raw-error"} {
if strings.Contains(bodyText, forbidden) {
t.Fatalf("cooldown response leaked %q: %s", forbidden, bodyText)
}
}
var envelope struct {
Error struct {
Code string `json:"code"`
Details map[string]any `json:"details"`
} `json:"error"`
}
if err := json.Unmarshal(body, &envelope); err != nil {
t.Fatalf("decode cooldown response: %v body=%s", err, body)
}
if envelope.Error.Code != "model_cooling_down" || len(envelope.Error.Details) != 2 {
t.Fatalf("unexpected cooldown envelope: %+v", envelope)
}
detail := f.loadTask(t, f.taskIDForModel(t, model, startedAt))
if len(detail.Attempts) != 0 {
t.Fatalf("candidate selection cooldown must not create provider attempts: %+v", detail.Attempts)
}
}
func (f *failurePolicyAcceptanceFixture) testAllCoolingAsyncRecovery(t *testing.T) {
f.setRunnerPolicy(t, false)
defer f.setRunnerPolicy(t, true)
model := "acceptance-async-cooling-" + f.suffix
baseModelID := f.createBaseModel(t, model, "image_generate")
stubA := newControlledProviderStub(func(call int, w http.ResponseWriter, _ *http.Request) {
if call == 1 {
writeStubJSON(w, http.StatusTooManyRequests, map[string]any{"error": map[string]any{"message": "controlled first 429"}})
return
}
writeOpenAIImageSuccess(w)
})
defer stubA.close()
stubB := newControlledProviderStub(func(_ int, w http.ResponseWriter, _ *http.Request) {
writeStubJSON(w, http.StatusTooManyRequests, map[string]any{"error": map[string]any{"message": "controlled second 429"}})
})
defer stubB.close()
override := func(seconds int) map[string]any {
return map[string]any{"failoverPolicy": map[string]any{"actionRules": []any{
map[string]any{"statusCodes": []any{429}, "action": "cooldown_and_next", "target": "model", "cooldownSeconds": seconds},
}}}
}
_, modelA := f.createPlatformModel(t, "async-cooling-a", "openai", stubA.server.URL+"/v1", 10, baseModelID, "async-a", "image_generate", 1, "", override(4))
_, modelB := f.createPlatformModel(t, "async-cooling-b", "openai", stubB.server.URL+"/v1", 20, baseModelID, "async-b", "image_generate", 1, "", override(8))
defer func() {
_, _ = f.pool.Exec(context.Background(), `UPDATE platform_models SET cooldown_until=NULL WHERE id IN ($1::uuid,$2::uuid)`, modelA.ID, modelB.ID)
}()
status, _, _ := f.postJSON(t, "/api/v1/images/generations", map[string]any{"model": model, "prompt": "prime cooldown"}, nil)
if status != http.StatusTooManyRequests {
t.Fatalf("prime cooldown status=%d", status)
}
startedAt := time.Now()
status, _, body := f.postJSON(t, "/api/v1/images/generations", map[string]any{"model": model, "prompt": "async recovery"}, map[string]string{"X-Async": "true"})
if status != http.StatusAccepted {
t.Fatalf("async submit status=%d body=%s", status, body)
}
taskID := f.taskIDForModel(t, model, startedAt)
deadline := time.Now().Add(2 * time.Second)
queuedObserved := false
for time.Now().Before(deadline) {
detail := f.loadTask(t, taskID)
if detail.Status == "queued" && len(detail.Attempts) == 0 {
queuedObserved = true
break
}
time.Sleep(25 * time.Millisecond)
}
if !queuedObserved {
t.Fatal("async cooling task was not observed queued with zero provider attempts")
}
deadline = time.Now().Add(12 * time.Second)
var completed acceptanceTaskDetail
for time.Now().Before(deadline) {
completed = f.loadTask(t, taskID)
if completed.Status == "succeeded" {
break
}
time.Sleep(100 * time.Millisecond)
}
if completed.Status != "succeeded" || len(completed.Attempts) != 1 || stubA.calls.Load() != 2 {
t.Fatalf("async task did not recover on first restored candidate: task=%+v callsA=%d callsB=%d", completed, stubA.calls.Load(), stubB.calls.Load())
}
}
func (f *failurePolicyAcceptanceFixture) testStreamingBoundary(t *testing.T) {
t.Run("首个 delta 前允许轮转", func(t *testing.T) {
model := "acceptance-stream-before-" + f.suffix
baseModelID := f.createBaseModel(t, model, "text_generate")
failed := newControlledProviderStub(func(_ int, w http.ResponseWriter, _ *http.Request) {
writeStubJSON(w, http.StatusInternalServerError, map[string]any{"error": map[string]any{"message": "before delta"}})
})
defer failed.close()
success := newControlledProviderStub(func(_ int, w http.ResponseWriter, _ *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
_, _ = io.WriteString(w, "data: {\"id\":\"chunk-ok\",\"object\":\"chat.completion.chunk\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"content\":\"stream fallback\"},\"finish_reason\":null}]}\n\n")
_, _ = io.WriteString(w, "data: {\"id\":\"chunk-ok\",\"object\":\"chat.completion.chunk\",\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n")
_, _ = io.WriteString(w, "data: [DONE]\n\n")
})
defer success.close()
_, _ = f.createPlatformModel(t, "stream-before-a", "openai", failed.server.URL+"/v1", 10, baseModelID, "stream-failed", "text_generate", 1, "", nil)
_, _ = f.createPlatformModel(t, "stream-before-b", "openai", success.server.URL+"/v1", 20, baseModelID, "stream-success", "text_generate", 1, "", nil)
status, _, body := f.postJSON(t, "/api/v1/chat/completions", map[string]any{
"model": model, "stream": true, "messages": []any{map[string]any{"role": "user", "content": "stream"}},
}, nil)
if status != http.StatusOK || failed.calls.Load() != 1 || success.calls.Load() != 1 || strings.Count(string(body), "stream fallback") != 1 || !strings.Contains(string(body), "[DONE]") {
t.Fatalf("unexpected pre-delta failover status=%d calls=%d/%d body=%s", status, failed.calls.Load(), success.calls.Load(), body)
}
})
t.Run("已输出 delta 后禁止轮转", func(t *testing.T) {
model := "acceptance-stream-after-" + f.suffix
baseModelID := f.createBaseModel(t, model, "text_generate")
broken := newControlledProviderStub(func(_ int, w http.ResponseWriter, _ *http.Request) {
hijacker, ok := w.(http.Hijacker)
if !ok {
http.Error(w, "hijacking unavailable", http.StatusInternalServerError)
return
}
connection, buffer, err := hijacker.Hijack()
if err != nil {
return
}
defer connection.Close()
_, _ = buffer.WriteString("HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nContent-Length: 4096\r\n\r\n")
_, _ = buffer.WriteString("data: {\"id\":\"chunk-broken\",\"object\":\"chat.completion.chunk\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"content\":\"first platform only\"},\"finish_reason\":null}]}\n\n")
_ = buffer.Flush()
})
defer broken.close()
fallback := newControlledProviderStub(func(_ int, w http.ResponseWriter, _ *http.Request) { writeOpenAIChatSuccess(w, "must not mix") })
defer fallback.close()
_, _ = f.createPlatformModel(t, "stream-after-a", "openai", broken.server.URL+"/v1", 10, baseModelID, "stream-broken", "text_generate", 2, "", nil)
_, _ = f.createPlatformModel(t, "stream-after-b", "openai", fallback.server.URL+"/v1", 20, baseModelID, "stream-fallback", "text_generate", 1, "", nil)
status, _, body := f.postJSON(t, "/api/v1/chat/completions", map[string]any{
"model": model, "stream": true, "messages": []any{map[string]any{"role": "user", "content": "stream"}},
}, nil)
if status != http.StatusOK || broken.calls.Load() != 1 || fallback.calls.Load() != 0 {
t.Fatalf("post-delta failure rotated unexpectedly status=%d calls=%d/%d body=%s", status, broken.calls.Load(), fallback.calls.Load(), body)
}
bodyText := string(body)
if !strings.Contains(bodyText, "first platform only") || strings.Contains(bodyText, "must not mix") || strings.Count(bodyText, "event: error") != 1 {
t.Fatalf("stream response mixed providers or missed terminal error: %s", bodyText)
}
})
}
func (f *failurePolicyAcceptanceFixture) testRunnerPolicyHTTPValidation(t *testing.T) {
valid := standardFailureRunnerPolicy(true)
var saved map[string]any
doJSON(t, f.server.URL, http.MethodPatch, "/api/admin/runtime/runner-policy", f.adminToken, valid, http.StatusOK, &saved)
var before map[string]any
doJSON(t, f.server.URL, http.MethodGet, "/api/admin/runtime/runner-policy", f.adminToken, nil, http.StatusOK, &before)
invalidRules := []map[string]any{
{"categories": []any{"rate_limit"}, "action": "rotate"},
{"categories": []any{"rate_limit"}, "action": "next", "target": "client"},
{"action": "next"},
{"statusCodes": []any{429}, "action": "cooldown_and_next", "target": "model", "cooldownSeconds": 0},
{"statusCodes": []any{500}, "action": "demote_and_next", "target": "platform", "demoteSteps": -1},
}
for index, rule := range invalidRules {
payload := standardFailureRunnerPolicy(true)
payload["failoverPolicy"].(map[string]any)["actionRules"] = []any{rule}
raw, _ := json.Marshal(payload)
request, _ := http.NewRequest(http.MethodPatch, f.server.URL+"/api/admin/runtime/runner-policy", bytes.NewReader(raw))
request.Header.Set("Authorization", "Bearer "+f.adminToken)
request.Header.Set("Content-Type", "application/json")
response, err := http.DefaultClient.Do(request)
if err != nil {
t.Fatalf("invalid policy request %d: %v", index, err)
}
body, _ := io.ReadAll(response.Body)
response.Body.Close()
if response.StatusCode != http.StatusBadRequest {
t.Fatalf("invalid policy %d status=%d body=%s", index, response.StatusCode, body)
}
}
var after map[string]any
doJSON(t, f.server.URL, http.MethodGet, "/api/admin/runtime/runner-policy", f.adminToken, nil, http.StatusOK, &after)
beforeFailover, _ := before["failoverPolicy"].(map[string]any)
afterFailover, _ := after["failoverPolicy"].(map[string]any)
if !jsonObjectsEqual(beforeFailover["actionRules"], afterFailover["actionRules"]) {
t.Fatalf("invalid PATCH overwrote runner action rules: before=%+v after=%+v", beforeFailover["actionRules"], afterFailover["actionRules"])
}
}
func jsonObjectsEqual(left any, right any) bool {
leftRaw, _ := json.Marshal(left)
rightRaw, _ := json.Marshal(right)
return bytes.Equal(leftRaw, rightRaw)
}
@@ -62,7 +62,7 @@ func (s *Server) uploadFile(w http.ResponseWriter, r *http.Request) {
writeError(w, status, err.Error())
return
}
writeJSON(w, http.StatusOK, upload)
writeJSON(w, http.StatusOK, easyAIFileUploadResponse(upload))
}
func firstNonEmptyFormValue(r *http.Request, key string, fallback string) string {
@@ -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")
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")
writeGeminiTaskError(http.StatusInternalServerError, "create task failed", nil, "task_create_failed")
return
@@ -184,6 +188,11 @@ func (s *Server) geminiGenerateContent(w http.ResponseWriter, r *http.Request) {
return
}
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 {
writeJSON(w, http.StatusOK, task.Result)
return
@@ -209,6 +218,7 @@ func (s *Server) geminiGenerateContent(w http.ResponseWriter, r *http.Request) {
if !requestStillConnected(r) {
return
}
applyRunErrorHeaders(w, runErr)
if wire := clients.ErrorWireResponse(runErr); wireResponseMatches(wire, clients.ProtocolGeminiGenerateContent) {
writeWireResponse(w, wire)
return
@@ -261,6 +271,7 @@ func (s *Server) writeGeminiGenerateContentStream(runCtx context.Context, w http
if !requestStillConnected(r) {
return
}
applyRunErrorHeaders(w, runErr)
if !nativePassthrough && !convertedFrames {
if wire := clients.ErrorWireResponse(runErr); wireResponseMatches(wire, clients.ProtocolGeminiGenerateContent) {
writeWireResponse(w, wire)
+95 -19
View File
@@ -580,12 +580,20 @@ func (s *Server) listModels(w http.ResponseWriter, r *http.Request) {
// @Tags playground
// @Produce json
// @Security BearerAuth
// @Param usage_scene query string false "模型可选场景;不传时保持原有行为" Enums(canvas_model_node,desktop)
// @Success 200 {object} PlatformModelListResponse
// @Failure 400 {object} ErrorEnvelope
// @Failure 401 {object} ErrorEnvelope
// @Failure 502 {object} ErrorEnvelope
// @Failure 500 {object} ErrorEnvelope
// @Router /api/v1/models [get]
// @Router /api/v1/playground/models [get]
func (s *Server) listPlayableModels(w http.ResponseWriter, r *http.Request) {
usageScene, err := parseModelUsageSceneQuery(r.URL.Query())
if err != nil {
writeError(w, http.StatusBadRequest, err.Error())
return
}
user, _ := auth.UserFromContext(r.Context())
models, err := s.store.ListAccessiblePlatformModels(r.Context(), user)
if err != nil {
@@ -593,6 +601,15 @@ func (s *Server) listPlayableModels(w http.ResponseWriter, r *http.Request) {
writeError(w, http.StatusInternalServerError, "list playable models failed")
return
}
if usageScene != "" {
availableModelIDs, sceneErr := s.listServerMainUsageSceneModelIDs(r.Context(), usageScene)
if sceneErr != nil {
s.logger.Error("filter playable models by usage scene failed", "usage_scene", usageScene, "error", sceneErr)
writeError(w, http.StatusBadGateway, "model usage scene catalog is unavailable")
return
}
models = filterPlatformModelsByUsageSceneIDs(models, availableModelIDs)
}
writeJSON(w, http.StatusOK, map[string]any{"items": s.platformModelResponses(r.Context(), models)})
}
@@ -963,6 +980,7 @@ func (s *Server) estimatePricing(w http.ResponseWriter, r *http.Request) {
return
}
if errors.Is(err, store.ErrNoModelCandidate) {
applyRunErrorHeaders(w, err)
writeErrorWithDetails(w, statusFromRunError(err), runErrorMessage(err), runErrorDetails(err), store.ModelCandidateErrorCode(err))
return
}
@@ -1039,15 +1057,14 @@ func (s *Server) listModelRateLimitStatuses(w http.ResponseWriter, r *http.Reque
// @Failure 429 {object} ErrorEnvelope
// @Failure 502 {object} ErrorEnvelope
// @Router /api/v1/reranks [post]
// @Router /api/v1/videos/generations [post]
// @Router /api/v1/song/generations [post]
// @Router /api/v1/music/generations [post]
// @Router /api/v1/speech/generations [post]
// @Router /api/v1/voice_clone [post]
func (s *Server) createTask(kind string, compatible bool) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
targetProtocol := targetProtocolForTaskRequest(kind, r)
writeTaskError := func(status int, message string, details map[string]any, codes ...string) {
if easyAIAsyncMediaRequest(kind, r) {
writeEasyAIAsyncError(w, status, message, details, codes...)
return
}
if targetProtocol == "" {
writeErrorWithDetails(w, status, message, details, codes...)
return
@@ -1093,9 +1110,10 @@ func (s *Server) createTask(kind string, compatible bool) http.Handler {
return
}
responsePlan := planTaskResponse(kind, compatible, body, r)
if targetProtocol != "" {
if targetProtocol != "" && !easyAIAsyncMediaRequest(kind, r) {
// OpenAI compatibility endpoints are synchronous APIs. Gateway async
// task envelopes belong only to native Gateway routes.
// task envelopes belong only to native Gateway routes, except for the
// explicit EasyAI media X-Async extension.
responsePlan.asyncMode = false
}
idempotencyKey, hasIdempotencyKey, err := optionalTaskIdempotencyKey(r)
@@ -1134,6 +1152,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")
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")
writeTaskError(http.StatusInternalServerError, "create task failed", nil, "task_create_failed")
return
@@ -1150,7 +1172,18 @@ func (s *Server) createTask(kind string, compatible bool) http.Handler {
writeTaskError(http.StatusConflict, "streaming idempotent replay is not supported", nil, "idempotency_stream_replay_unsupported")
return
}
if targetProtocol != "" {
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 task.Status == "succeeded" {
writeJSON(w, http.StatusOK, task.Result)
return
@@ -1274,7 +1307,7 @@ func openAIEmbeddingsDoc() {}
// @Security BearerAuth
// @Param input body TaskRequest true "OpenAI Images 请求"
// @Param X-Async header bool false "true 时异步创建任务并返回 202"
// @Success 200 {object} CompatibleResponse
// @Success 200 {object} OpenAIImagesResponse
// @Success 202 {object} TaskAcceptedResponse
// @Failure 400 {object} OpenAIErrorEnvelope
// @Failure 401 {object} OpenAIErrorEnvelope
@@ -1284,6 +1317,32 @@ func openAIEmbeddingsDoc() {}
// @Router /api/v1/images/edits [post]
func openAIImagesDoc() {}
// easyAIMediaTasksDoc godoc
// @Summary 创建 EasyAI 媒体任务
// @Description 默认同步返回 EasyAI GeneratedResponse;设置 X-Async=true 时返回同时兼容 Gateway 与 server-main EasyAIClient 的异步提交结构。
// @Tags media
// @Accept json
// @Produce json
// @Security BearerAuth
// @Param X-Async header bool false "true 时异步创建任务并返回 202"
// @Param Idempotency-Key header string false "可选请求幂等键;同一用户范围内唯一"
// @Param input body TaskRequest true "媒体任务请求,字段随任务类型变化"
// @Success 200 {object} EasyAIGeneratedResponse
// @Success 202 {object} TaskAcceptedResponse
// @Failure 400 {object} ErrorEnvelope
// @Failure 401 {object} ErrorEnvelope
// @Failure 402 {object} ErrorEnvelope
// @Failure 403 {object} ErrorEnvelope
// @Failure 404 {object} ErrorEnvelope
// @Failure 429 {object} ErrorEnvelope
// @Failure 502 {object} ErrorEnvelope
// @Router /api/v1/videos/generations [post]
// @Router /api/v1/song/generations [post]
// @Router /api/v1/music/generations [post]
// @Router /api/v1/speech/generations [post]
// @Router /api/v1/voice_clone [post]
func easyAIMediaTasksDoc() {}
func (s *Server) requestExecutionContext(r *http.Request) (context.Context, context.CancelFunc) {
base := context.WithoutCancel(r.Context())
if s.ctx == nil {
@@ -1363,6 +1422,7 @@ func writeProtocolCompatibleTaskResponse(runCtx context.Context, w http.Response
if !requestStillConnected(r) {
return
}
applyRunErrorHeaders(w, runErr)
errorPayload := map[string]any{"message": runErrorMessage(runErr), "type": "server_error", "param": nil, "code": runErrorCode(runErr)}
if status := statusFromRunError(runErr); status >= 400 && status < 500 {
errorPayload["type"] = "invalid_request_error"
@@ -1400,6 +1460,7 @@ func writeProtocolCompatibleTaskResponse(runCtx context.Context, w http.Response
if !requestStillConnected(r) {
return
}
applyRunErrorHeaders(w, runErr)
if wire := clients.ErrorWireResponse(runErr); wireResponseMatches(wire, targetProtocol) {
writeWireResponse(w, wire)
return
@@ -1418,7 +1479,7 @@ func writeProtocolCompatibleTaskResponse(runCtx context.Context, w http.Response
writeWireResponse(w, result.Wire)
return
}
writeJSON(w, http.StatusOK, result.Output)
writeJSON(w, http.StatusOK, easyAISynchronousTaskResponse(result.Task, result.Output))
}
func streamIncludeUsage(body map[string]any) bool {
@@ -1461,14 +1522,7 @@ func streamCompatibleKind(kind string) bool {
}
func writeTaskAccepted(w http.ResponseWriter, task store.GatewayTask) {
writeJSON(w, http.StatusAccepted, map[string]any{
"taskId": task.ID,
"task": task,
"next": map[string]string{
"events": fmt.Sprintf("/api/v1/tasks/%s/events", task.ID),
"detail": fmt.Sprintf("/api/v1/tasks/%s", task.ID),
},
})
writeJSON(w, http.StatusAccepted, easyAITaskAcceptedResponse(task))
}
func apiKeyScopeAllowed(user *auth.User, kind string) bool {
@@ -1578,6 +1632,12 @@ func scopeForTaskKind(kind string) string {
func statusFromRunError(err error) int {
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":
return http.StatusServiceUnavailable
case runner.IsPricingUnavailable(err):
@@ -1637,11 +1697,21 @@ func runErrorDetails(err error) map[string]any {
return map[string]any{"rateLimit": detail}
}
if detail := store.ModelCandidateErrorDetails(err); len(detail) > 0 {
return map[string]any{"modelCandidate": detail}
return detail
}
return nil
}
func applyRunErrorHeaders(w http.ResponseWriter, err error) {
if retryAfter := store.ModelCandidateRetryAfter(err); retryAfter > 0 {
seconds := int((retryAfter + time.Second - 1) / time.Second)
if seconds < 1 {
seconds = 1
}
w.Header().Set("Retry-After", strconv.Itoa(seconds))
}
}
func rateLimitErrorSummary(err error) string {
var limitErr *store.RateLimitExceededError
if !errors.As(err, &limitErr) {
@@ -1882,6 +1952,7 @@ func boolValue(body map[string]any, key string) bool {
// @Success 200 {object} store.GatewayTask
// @Failure 401 {object} ErrorEnvelope
// @Failure 404 {object} ErrorEnvelope
// @Failure 410 {object} ErrorEnvelope
// @Failure 500 {object} ErrorEnvelope
// @Router /api/workspace/tasks/{taskID} [get]
// @Router /api/v1/tasks/{taskID} [get]
@@ -1901,6 +1972,11 @@ func (s *Server) getTask(w http.ResponseWriter, r *http.Request) {
task.Cancellable = &cancelState.Cancellable
task.Submitted = &cancelState.Submitted
task.Message = cancelState.Message
task, err = s.hydrateTaskResult(r.Context(), task)
if err != nil {
writeStoredBinaryResultError(w, err)
return
}
writeJSON(w, http.StatusOK, task)
return
}
@@ -0,0 +1,54 @@
package httpapi
import (
"context"
"strings"
"testing"
"github.com/jackc/pgx/v5/pgxpool"
)
func requireDedicatedIntegrationTestDatabase(t *testing.T, ctx context.Context, pool *pgxpool.Pool) {
t.Helper()
var databaseName string
if err := pool.QueryRow(ctx, `SELECT current_database()`).Scan(&databaseName); err != nil {
t.Fatalf("read integration test database name: %v", err)
}
if !isDedicatedIntegrationTestDatabase(databaseName) {
t.Fatalf(
"AI_GATEWAY_TEST_DATABASE_URL must reference a dedicated test database; refusing to use %q",
databaseName,
)
}
}
func isDedicatedIntegrationTestDatabase(databaseName string) bool {
normalized := strings.ToLower(strings.TrimSpace(databaseName))
return normalized == "test" ||
strings.HasPrefix(normalized, "test_") ||
strings.HasSuffix(normalized, "_test") ||
strings.Contains(normalized, "_test_")
}
func TestDedicatedIntegrationTestDatabaseName(t *testing.T) {
t.Parallel()
for _, testCase := range []struct {
name string
database string
allowed bool
}{
{name: "documented local test database", database: "easyai_ai_gateway_test_codex", allowed: true},
{name: "test prefix", database: "test_gateway", allowed: true},
{name: "test suffix", database: "gateway_test", allowed: true},
{name: "production database", database: "easyai_ai_gateway", allowed: false},
{name: "similar production name", database: "gateway_contest", allowed: false},
} {
testCase := testCase
t.Run(testCase.name, func(t *testing.T) {
t.Parallel()
if got := isDedicatedIntegrationTestDatabase(testCase.database); got != testCase.allowed {
t.Fatalf("isDedicatedIntegrationTestDatabase(%q) = %v, want %v", testCase.database, got, testCase.allowed)
}
})
}
}
@@ -177,6 +177,10 @@ func (s *Server) createKelingOmniVideo(w http.ResponseWriter, r *http.Request) {
writeKelingCompatError(w, requestID, kelingCompatGatewayError(staged.Err))
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)
writeKelingCompatError(w, requestID, newKelingCompatError(http.StatusInternalServerError, 5000, "create task failed"))
return
@@ -202,15 +206,11 @@ func (s *Server) createKelingOmniVideo(w http.ResponseWriter, r *http.Request) {
writeKelingCompatError(w, requestID, kelingCompatGatewayError(createErr))
return
}
if kelingTaskUsesNativeProtocol(task) && len(task.CompatibilitySubmitBody) > 0 {
writeJSON(w, task.CompatibilitySubmitHTTPStatus, task.CompatibilitySubmitBody)
return
}
writeJSON(w, http.StatusOK, KelingCompatibleEnvelope{
Code: 0,
Message: "SUCCEED",
RequestID: requestID,
Data: kelingCompatTaskData(task),
Data: kelingCompatSubmissionData(task),
})
}
@@ -225,6 +225,7 @@ func (s *Server) createKelingOmniVideo(w http.ResponseWriter, r *http.Request) {
// @Failure 401 {object} KelingCompatibleEnvelope
// @Failure 403 {object} KelingCompatibleEnvelope
// @Failure 404 {object} KelingCompatibleEnvelope
// @Failure 410 {object} KelingCompatibleEnvelope
// @Failure 500 {object} KelingCompatibleEnvelope
// @Router /api/v1/videos/omni-video/{taskID} [get]
func (s *Server) getKelingOmniVideo(w http.ResponseWriter, r *http.Request) {
@@ -251,15 +252,10 @@ func (s *Server) getKelingOmniVideo(w http.ResponseWriter, r *http.Request) {
writeKelingCompatError(w, requestID, newKelingCompatError(http.StatusNotFound, 1203, "task not found"))
return
}
if kelingTaskUsesNativeProtocol(task) {
if raw, ok := task.Result["raw"].(map[string]any); ok && len(raw) > 0 {
writeJSON(w, http.StatusOK, raw)
return
}
if len(task.RemoteTaskPayload) > 0 {
writeJSON(w, http.StatusOK, task.RemoteTaskPayload)
return
}
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{
Code: 0,
@@ -705,15 +701,22 @@ func kelingCompatTaskData(task store.GatewayTask) map[string]any {
return data
}
func kelingCompatPublicID(task store.GatewayTask) string {
if publicID := strings.TrimSpace(task.CompatibilityPublicID); publicID != "" {
return publicID
}
return task.ID
func kelingCompatSubmissionData(task store.GatewayTask) map[string]any {
data := kelingCompatTaskData(task)
data["task_status"] = "submitted"
delete(data, "task_status_code")
delete(data, "task_status_msg")
delete(data, "task_result")
return data
}
func kelingTaskUsesNativeProtocol(task store.GatewayTask) bool {
return task.CompatibilityProtocol == clients.ProtocolKlingV1Omni && task.CompatibilitySourceProtocol == clients.ProtocolKlingV1Omni
func kelingCompatPublicID(task store.GatewayTask) string {
if task.CompatibilityProtocol == task.CompatibilitySourceProtocol {
if remoteTaskID := strings.TrimSpace(task.RemoteTaskID); remoteTaskID != "" {
return remoteTaskID
}
}
return task.ID
}
func kelingCompatTaskStatus(status string) string {
@@ -209,6 +209,10 @@ func TestKelingCompatTaskDataAndOwnership(t *testing.T) {
if video["id"] != "video-1" || video["watermark_url"] != "https://example.com/watermarked.mp4" || video["duration"] != "5" {
t.Fatalf("unexpected compatible video: %+v", video)
}
submission := kelingCompatSubmissionData(task)
if submission["task_status"] != "submitted" || submission["task_result"] != nil {
t.Fatalf("submission response must be stable and minimal: %+v", submission)
}
}
func TestRequireKelingAPIKeyWritesOfficialAuthEnvelope(t *testing.T) {
+36 -30
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")
case errors.Is(err, store.ErrExternalTaskIDReused):
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:
s.logger.Error("create Kling compatibility task failed", "version", version, "model", model, "error", err)
writeKlingCompatError(w, http.StatusInternalServerError, "create task failed", "task_create_failed")
@@ -177,10 +179,10 @@ func (s *Server) createKlingCompatTask(w http.ResponseWriter, r *http.Request, v
w.Header().Set("Idempotent-Replayed", "true")
}
if version == "v2" {
writeJSON(w, http.StatusOK, klingV2Envelope(task))
writeJSON(w, http.StatusOK, klingV2SubmissionEnvelope(task))
return
}
writeJSON(w, http.StatusOK, klingV1Envelope(task))
writeJSON(w, http.StatusOK, klingV1SubmissionEnvelope(task))
}
func klingCompatTaskBody(version string, model string, native map[string]any) (map[string]any, string, error) {
@@ -464,6 +466,7 @@ func validateKlingCompatBody(model string, body map[string]any) error {
// @Param taskID path string true "任务 ID"
// @Success 200 {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]
func (s *Server) klingV1GetOmniVideo(w http.ResponseWriter, r *http.Request) {
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")
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))
}
@@ -520,6 +528,7 @@ func (s *Server) klingV1ListOmniVideos(w http.ResponseWriter, r *http.Request) {
// @Param task_ids query string false "逗号分隔的任务 ID"
// @Param external_task_ids query string false "逗号分隔的外部任务 ID"
// @Success 200 {object} map[string]interface{}
// @Failure 410 {object} map[string]interface{}
// @Router /api/v1/kling/v2/tasks [get]
func (s *Server) klingV2GetTasks(w http.ResponseWriter, r *http.Request) {
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")
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))
}
writeJSON(w, http.StatusOK, map[string]any{"code": 0, "message": "success", "request_id": klingRequestIDFromAny(data), "data": data})
@@ -612,20 +626,17 @@ func (s *Server) klingV2ListTasks(w http.ResponseWriter, r *http.Request) {
}
func klingV1Envelope(task store.GatewayTask) map[string]any {
if klingTaskUsesNativeV1Protocol(task) {
if raw := klingMap(task.Result["raw"]); len(raw) > 0 {
return raw
}
if len(task.RemoteTaskPayload) > 0 {
return task.RemoteTaskPayload
}
if raw := klingMap(task.CompatibilitySubmitBody); len(raw) > 0 {
return raw
}
}
return map[string]any{"code": 0, "message": "success", "request_id": firstNonEmpty(task.RequestID, task.ID), "data": klingV1TaskData(task)}
}
func klingV1SubmissionEnvelope(task store.GatewayTask) map[string]any {
data := klingV1TaskData(task)
data["task_status"] = "submitted"
delete(data, "task_status_msg")
delete(data, "task_result")
return map[string]any{"code": 0, "message": "success", "request_id": firstNonEmpty(task.RequestID, task.ID), "data": data}
}
func klingV1TaskData(task store.GatewayTask) map[string]any {
data := map[string]any{
"task_id": klingCompatibilityPublicID(task), "task_status": klingV1Status(task.Status),
@@ -646,20 +657,17 @@ func klingV1TaskData(task store.GatewayTask) map[string]any {
}
func klingV2Envelope(task store.GatewayTask) map[string]any {
if task.CompatibilityProtocol == clients.ProtocolKlingV2Omni && task.CompatibilitySourceProtocol == clients.ProtocolKlingV2Omni {
if raw := klingMap(task.Result["raw"]); len(raw) > 0 {
return raw
}
if len(task.RemoteTaskPayload) > 0 {
return task.RemoteTaskPayload
}
if len(task.CompatibilitySubmitBody) > 0 {
return task.CompatibilitySubmitBody
}
}
return map[string]any{"code": 0, "message": "success", "request_id": firstNonEmpty(task.RequestID, task.ID), "data": klingV2TaskData(task)}
}
func klingV2SubmissionEnvelope(task store.GatewayTask) map[string]any {
data := klingV2TaskData(task)
data["status"] = "submitted"
delete(data, "message")
delete(data, "outputs")
return map[string]any{"code": 0, "message": "success", "request_id": firstNonEmpty(task.RequestID, task.ID), "data": data}
}
func klingV2TaskData(task store.GatewayTask) map[string]any {
data := map[string]any{
"id": klingCompatibilityPublicID(task), "status": klingV2Status(task.Status),
@@ -676,16 +684,14 @@ func klingV2TaskData(task store.GatewayTask) map[string]any {
}
func klingCompatibilityPublicID(task store.GatewayTask) string {
if publicID := strings.TrimSpace(task.CompatibilityPublicID); publicID != "" {
return publicID
if task.CompatibilityProtocol == task.CompatibilitySourceProtocol {
if remoteTaskID := strings.TrimSpace(task.RemoteTaskID); remoteTaskID != "" {
return remoteTaskID
}
}
return task.ID
}
func klingTaskUsesNativeV1Protocol(task store.GatewayTask) bool {
return task.CompatibilityProtocol == clients.ProtocolKlingV1Omni && task.CompatibilitySourceProtocol == clients.ProtocolKlingV1Omni
}
func klingMap(value any) map[string]any {
result, _ := value.(map[string]any)
return result
@@ -6,6 +6,9 @@ import (
"net/http/httptest"
"strings"
"testing"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
func TestKlingV1O1CompatibilityBody(t *testing.T) {
@@ -86,6 +89,27 @@ func TestKlingV2CompatibilityBody(t *testing.T) {
}
}
func TestKlingSubmissionEnvelopesDoNotReplayCurrentResult(t *testing.T) {
task := store.GatewayTask{
ID: "gateway-task", Status: "succeeded", RemoteTaskID: "upstream-task",
CompatibilityProtocol: "kling-v1-omni",
CompatibilitySourceProtocol: "kling-v1-omni",
Result: map[string]any{"data": []any{map[string]any{"url": "https://example.invalid/result.mp4"}}},
CreatedAt: time.Unix(100, 0),
UpdatedAt: time.Unix(101, 0),
}
v1 := klingV1SubmissionEnvelope(task)
v1Data, _ := v1["data"].(map[string]any)
if v1Data["task_status"] != "submitted" || v1Data["task_result"] != nil {
t.Fatalf("V1 submission envelope=%#v", v1)
}
v2 := klingV2SubmissionEnvelope(task)
v2Data, _ := v2["data"].(map[string]any)
if v2Data["status"] != "submitted" || v2Data["outputs"] != nil {
t.Fatalf("V2 submission envelope=%#v", v2)
}
}
func TestKlingCompatibilityValidation(t *testing.T) {
tests := []struct {
name string
@@ -45,6 +45,58 @@ func (s *Server) cleanupExpiredLocalTempAssets(ctx context.Context, now time.Tim
for _, target := range targets {
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
}
+46 -40
View File
@@ -62,17 +62,19 @@ type ModelCatalogProviderSummary struct {
}
type ModelCatalogSource struct {
ID string `json:"id"`
PlatformID string `json:"platformId,omitempty"`
PlatformName string `json:"platformName,omitempty"`
ProviderKey string `json:"providerKey"`
ProviderName string `json:"providerName"`
ModelName string `json:"modelName"`
ModelAlias string `json:"modelAlias,omitempty"`
DisplayName string `json:"displayName"`
ModelType []string `json:"modelType"`
Enabled bool `json:"enabled"`
RateLimits ModelCatalogRateLimits `json:"rateLimits"`
ID string `json:"id"`
PlatformID string `json:"platformId,omitempty"`
PlatformName string `json:"platformName,omitempty"`
ProviderKey string `json:"providerKey"`
ProviderName string `json:"providerName"`
ModelName string `json:"modelName"`
ProviderModelName string `json:"providerModelName,omitempty"`
ModelAlias string `json:"modelAlias,omitempty"`
LegacyAliases []string `json:"legacyAliases,omitempty"`
DisplayName string `json:"displayName"`
ModelType []string `json:"modelType"`
Enabled bool `json:"enabled"`
RateLimits ModelCatalogRateLimits `json:"rateLimits"`
}
type ModelCatalogText struct {
@@ -202,10 +204,10 @@ func buildModelCatalog(
sourceCount := 0
for _, model := range models {
platform := platformMap[model.PlatformID]
if !catalogSourceEnabled(model, platform) {
baseModel := baseModelMap[model.BaseModelID]
if !catalogSourceEnabled(model, platform, baseModel) {
continue
}
baseModel := baseModelMap[model.BaseModelID]
providerKey := modelProviderKey(model, baseModel, platform)
provider := providerMap[providerKey]
providerName := firstNonEmpty(
@@ -232,14 +234,14 @@ func buildModelCatalog(
rateLimits: effectiveModelRateLimits(model, platform, runtimePolicyMap),
}
key := catalogAliasKey(model)
key := catalogIdentityKey(model, baseModel)
current := groups[key]
if current == nil {
current = &catalogGroup{
key: key,
alias: firstNonEmpty(model.ModelAlias, model.DisplayName, model.ModelName),
displayName: firstNonEmpty(model.DisplayName, model.ModelAlias, model.ModelName),
modelName: model.ModelName,
alias: firstNonEmpty(baseModel.InvocationName, model.ModelName),
displayName: firstNonEmpty(baseModel.DisplayName, model.DisplayName, baseModel.InvocationName, model.ModelName),
modelName: firstNonEmpty(baseModel.InvocationName, model.ModelName),
modelType: cloneStringSlice(model.ModelType),
capabilities: cloneObject(model.Capabilities),
}
@@ -274,8 +276,8 @@ func buildModelCatalog(
}
}
func catalogSourceEnabled(model store.PlatformModel, platform store.Platform) bool {
return model.Enabled && platform.ID != "" && platform.Status == "enabled"
func catalogSourceEnabled(model store.PlatformModel, platform store.Platform, baseModel store.BaseModel) bool {
return model.Enabled && platform.ID != "" && platform.Status == "enabled" && (baseModel.ID == "" || baseModel.Status == "" || baseModel.Status == "active")
}
func modelCatalogItem(group *catalogGroup, accessRuleGroups map[string]accessRuleGroup) ModelCatalogItem {
@@ -304,17 +306,19 @@ func modelCatalogItem(group *catalogGroup, accessRuleGroups map[string]accessRul
itemSources := make([]ModelCatalogSource, 0, len(sources))
for _, source := range sources {
itemSources = append(itemSources, ModelCatalogSource{
ID: source.model.ID,
PlatformID: source.model.PlatformID,
PlatformName: firstNonEmpty(source.platform.Name, source.model.PlatformName),
ProviderKey: source.providerKey,
ProviderName: source.providerName,
ModelName: source.model.ModelName,
ModelAlias: source.model.ModelAlias,
DisplayName: firstNonEmpty(source.model.DisplayName, source.model.ModelAlias, source.model.ModelName),
ModelType: cloneStringSlice(source.model.ModelType),
Enabled: source.model.Enabled && source.platform.Status == "enabled",
RateLimits: source.rateLimits,
ID: source.model.ID,
PlatformID: source.model.PlatformID,
PlatformName: firstNonEmpty(source.platform.Name, source.model.PlatformName),
ProviderKey: source.providerKey,
ProviderName: source.providerName,
ModelName: firstNonEmpty(source.baseModel.InvocationName, source.model.ModelName),
ProviderModelName: source.model.ProviderModelName,
ModelAlias: firstNonEmpty(source.baseModel.InvocationName, source.model.ModelName),
LegacyAliases: cloneStringSlice(source.baseModel.LegacyAliases),
DisplayName: firstNonEmpty(source.baseModel.DisplayName, source.model.DisplayName, source.baseModel.InvocationName, source.model.ModelName),
ModelType: cloneStringSlice(source.model.ModelType),
Enabled: source.model.Enabled && source.platform.Status == "enabled",
RateLimits: source.rateLimits,
})
}
@@ -483,18 +487,20 @@ func discountTitle(label string, discount float64) string {
}
func effectiveModelRateLimits(model store.PlatformModel, platform store.Platform, runtimePolicyMap map[string]store.RuntimePolicySet) ModelCatalogRateLimits {
overridePolicyRaw, hasOverridePolicy := model.RuntimePolicyOverride["rateLimitPolicy"]
overridePolicy := objectValue(overridePolicyRaw)
runtimePolicy := map[string]any(nil)
if model.RuntimePolicySetID != "" {
runtimePolicy = runtimePolicyMap[model.RuntimePolicySetID].RateLimitPolicy
}
policies := []rateLimitPolicySource{
{policy: overridePolicy, authoritative: hasOverridePolicy},
{policy: model.RateLimitPolicy, authoritative: len(model.RateLimitPolicy) > 0},
{policy: runtimePolicy, authoritative: strings.TrimSpace(model.RuntimePolicySetID) != ""},
{policy: platform.RateLimitPolicy},
}
policy := store.EffectiveRateLimitPolicy(store.EffectiveRateLimitPolicyInput{
BasePolicy: model.BaseRateLimitPolicy,
PlatformPolicy: platform.RateLimitPolicy,
RuntimePolicy: runtimePolicy,
RuntimePolicyExplicit: model.RuntimePolicySetID != "",
RuntimePolicyOverride: model.RuntimePolicyOverride,
ModelPolicy: model.RateLimitPolicy,
ModelPolicyMode: model.RateLimitPolicyMode,
})
policies := []rateLimitPolicySource{{policy: policy, authoritative: policy != nil}}
limits := ModelCatalogRateLimits{
RPM: firstRateLimit(policies, "rpm"),
TPM: firstRateLimit(policies, "tpm_total"),
@@ -1260,8 +1266,8 @@ func modelIconPath(model store.PlatformModel, baseModel store.BaseModel, provide
)
}
func catalogAliasKey(model store.PlatformModel) string {
key := strings.ToLower(strings.TrimSpace(firstNonEmpty(model.ModelAlias, model.DisplayName, model.ModelName)))
func catalogIdentityKey(model store.PlatformModel, baseModel store.BaseModel) string {
key := strings.ToLower(strings.TrimSpace(firstNonEmpty(baseModel.InvocationName, baseModel.CanonicalModelKey, model.ModelName)))
if key == "" {
return model.ID
}
@@ -131,13 +131,24 @@ func TestBuildModelCatalogKeepsDisplayNameSeparateFromCallAlias(t *testing.T) {
platforms := []store.Platform{
{ID: "platform-volces", Provider: "volces", Name: "火山引擎", Status: "enabled"},
}
baseModels := []store.BaseModel{
{
ID: "base-seedream-pro",
CanonicalModelKey: "seedream:seedream-5.0-pro",
InvocationName: "seedream-5.0-pro",
ProviderModelName: "doubao-seedream-5-0-pro-260628",
DisplayName: "Seedream 5.0 Pro",
Status: "active",
},
}
models[0].BaseModelID = "base-seedream-pro"
response := buildModelCatalog(models, platforms, nil, nil, nil, nil, nil)
response := buildModelCatalog(models, platforms, nil, nil, nil, nil, baseModels)
if len(response.Items) != 1 {
t.Fatalf("expected one Seedream Pro catalog item, got %+v", response.Items)
}
item := response.Items[0]
if item.Alias != "Seedream-5.0-Pro" || item.DisplayName != "Seedream 5.0 Pro" || item.ModelName != "Seedream-5.0-Pro" {
if item.Alias != "seedream-5.0-pro" || item.DisplayName != "Seedream 5.0 Pro" || item.ModelName != "seedream-5.0-pro" {
t.Fatalf("catalog should separate the call alias from the display name, got %+v", item)
}
if strings.Contains(item.Alias, " ") {
@@ -145,6 +156,46 @@ func TestBuildModelCatalogKeepsDisplayNameSeparateFromCallAlias(t *testing.T) {
}
}
func TestBuildModelCatalogAggregatesByInvocationNameNotDisplayName(t *testing.T) {
models := []store.PlatformModel{
{ID: "official-a", PlatformID: "platform-a", BaseModelID: "base-a", ModelName: "provider/a", ModelType: store.StringList{"text_generate"}, DisplayName: "同名营销标题", Enabled: true},
{ID: "official-b", PlatformID: "platform-b", BaseModelID: "base-b", ModelName: "provider/b", ModelType: store.StringList{"text_generate"}, DisplayName: "同名营销标题", Enabled: true},
{ID: "official-a-mirror", PlatformID: "platform-b", BaseModelID: "base-a-mirror", ModelName: "mirror/a", ModelType: store.StringList{"text_generate"}, DisplayName: "另一个展示标题", Enabled: true},
}
platforms := []store.Platform{
{ID: "platform-a", Provider: "provider-a", Status: "enabled"},
{ID: "platform-b", Provider: "provider-b", Status: "enabled"},
}
baseModels := []store.BaseModel{
{ID: "base-a", CanonicalModelKey: "provider-a:official-a", InvocationName: "official-a", DisplayName: "模型 A", Status: "active"},
{ID: "base-a-mirror", CanonicalModelKey: "provider-b:official-a", InvocationName: "official-a", DisplayName: "模型 A 镜像", Status: "active"},
{ID: "base-b", CanonicalModelKey: "provider-b:official-b", InvocationName: "official-b", DisplayName: "同名营销标题", Status: "active"},
}
response := buildModelCatalog(models, platforms, nil, nil, nil, nil, baseModels)
if response.Summary.ModelCount != 2 || response.Summary.SourceCount != 3 {
t.Fatalf("expected two official identities and three sources, got %+v", response.Summary)
}
for _, item := range response.Items {
if item.ModelName == "official-a" && item.SourceCount != 2 {
t.Fatalf("official-a mirrors should aggregate into one card: %+v", item)
}
}
}
func TestBuildModelCatalogExcludesDeprecatedBaseModels(t *testing.T) {
models := []store.PlatformModel{
{ID: "preview", PlatformID: "platform", BaseModelID: "base-preview", ModelName: "gemini-preview", ModelType: store.StringList{"text_generate"}, Enabled: true},
}
platforms := []store.Platform{{ID: "platform", Provider: "gemini", Status: "enabled"}}
baseModels := []store.BaseModel{{ID: "base-preview", InvocationName: "gemini-preview", Status: "deprecated"}}
response := buildModelCatalog(models, platforms, nil, nil, nil, nil, baseModels)
if response.Summary.ModelCount != 0 || response.Summary.SourceCount != 0 {
t.Fatalf("deprecated base models must not create new catalog cards: %+v", response.Summary)
}
}
func TestBuildModelCatalogUsesBaseModelProviderForProviderFilters(t *testing.T) {
models := []store.PlatformModel{
{
@@ -0,0 +1,39 @@
package httpapi
import (
"net/http/httptest"
"testing"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
func TestModelCooldownErrorPublicContract(t *testing.T) {
recoveryAt := time.Now().UTC().Add(17 * time.Second).Truncate(time.Second)
err := &store.ModelCandidateUnavailableError{
Code: "model_cooling_down",
Message: "请求的模型暂时不可用,请稍后重试",
RetryAfter: 17 * time.Second,
RecoveryAt: recoveryAt,
Details: map[string]any{
"retryAfterSeconds": 17,
"recoveryAt": recoveryAt.Format(time.RFC3339),
},
}
recorder := httptest.NewRecorder()
applyRunErrorHeaders(recorder, err)
if got := statusFromRunError(err); got != 429 {
t.Fatalf("status = %d, want 429", got)
}
if got := runErrorCode(err); got != "model_cooling_down" {
t.Fatalf("code = %q, want model_cooling_down", got)
}
if got := recorder.Header().Get("Retry-After"); got != "17" {
t.Fatalf("Retry-After = %q, want 17", got)
}
details := runErrorDetails(err)
if len(details) != 2 || details["retryAfterSeconds"] != 17 || details["recoveryAt"] != recoveryAt.Format(time.RFC3339) {
t.Fatalf("unexpected public details: %#v", details)
}
}
@@ -0,0 +1,186 @@
package httpapi
import (
"context"
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
"net/url"
"strings"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
const maxUsageSceneCatalogResponseBytes = 8 << 20
var supportedModelUsageScenes = map[string]struct{}{
"canvas_model_node": {},
"desktop": {},
}
func parseModelUsageSceneQuery(query url.Values) (string, error) {
values, present := query["usage_scene"]
if !present {
return "", nil
}
if len(values) != 1 {
return "", errors.New("usage_scene must be provided exactly once")
}
scene := strings.TrimSpace(values[0])
if _, ok := supportedModelUsageScenes[scene]; !ok {
return "", errors.New("usage_scene must be canvas_model_node or desktop")
}
return scene, nil
}
func (s *Server) listServerMainUsageSceneModelIDs(ctx context.Context, usageScene string) (map[string]struct{}, error) {
baseURL := strings.TrimRight(strings.TrimSpace(s.cfg.ServerMainBaseURL), "/")
if baseURL == "" {
return nil, errors.New("SERVER_MAIN_BASE_URL is not configured")
}
var requestErrors []error
for _, endpoint := range serverMainUsageSceneCatalogEndpoints(baseURL, usageScene) {
ids, err := s.requestServerMainUsageSceneModelIDs(ctx, endpoint)
if err == nil {
return ids, nil
}
requestErrors = append(requestErrors, err)
}
return nil, fmt.Errorf("server-main model catalog request failed: %w", errors.Join(requestErrors...))
}
func serverMainUsageSceneCatalogEndpoints(baseURL string, usageScene string) []string {
path := "/integration-platform/models/strict/all?usage_scene=" + url.QueryEscape(usageScene)
if strings.HasSuffix(baseURL, "/api") {
return []string{baseURL + path}
}
return []string{baseURL + "/api" + path, baseURL + path}
}
func (s *Server) requestServerMainUsageSceneModelIDs(ctx context.Context, endpoint string) (map[string]struct{}, error) {
request, err := http.NewRequestWithContext(ctx, http.MethodGet, endpoint, nil)
if err != nil {
return nil, fmt.Errorf("create server-main model catalog request: %w", err)
}
request.Header.Set("Accept", "application/json")
client := http.DefaultClient
if s.auth != nil && s.auth.HTTPClient != nil {
client = s.auth.HTTPClient
}
response, err := client.Do(request)
if err != nil {
return nil, fmt.Errorf("request server-main model catalog: %w", err)
}
defer response.Body.Close()
if response.StatusCode < http.StatusOK || response.StatusCode >= http.StatusMultipleChoices {
_, _ = io.Copy(io.Discard, io.LimitReader(response.Body, 4<<10))
return nil, fmt.Errorf("server-main model catalog returned HTTP %d", response.StatusCode)
}
var payload any
decoder := json.NewDecoder(io.LimitReader(response.Body, maxUsageSceneCatalogResponseBytes))
if err := decoder.Decode(&payload); err != nil {
return nil, fmt.Errorf("decode server-main model catalog: %w", err)
}
items, ok := usageSceneCatalogItems(payload)
if !ok {
return nil, errors.New("server-main model catalog payload is malformed")
}
return usageSceneModelIDs(items), nil
}
func usageSceneCatalogItems(payload any) ([]any, bool) {
if items, ok := payload.([]any); ok {
return items, true
}
record, ok := payload.(map[string]any)
if !ok {
return nil, false
}
for _, key := range []string{"modelList", "items", "data"} {
value, exists := record[key]
if !exists {
continue
}
if items, ok := value.([]any); ok {
return items, true
}
if items, ok := usageSceneCatalogItems(value); ok {
return items, true
}
}
return nil, false
}
func usageSceneModelIDs(items []any) map[string]struct{} {
ids := make(map[string]struct{}, len(items))
for _, item := range items {
record, ok := item.(map[string]any)
if !ok {
continue
}
for _, key := range []string{
"alias",
"invocationName",
"invocation_name",
"modelAlias",
"model_alias",
"name",
"modelName",
"model_name",
"providerModelName",
"provider_model_name",
} {
addUsageSceneModelID(ids, record[key])
}
for _, key := range []string{"legacyAliases", "legacy_aliases"} {
if aliases, ok := record[key].([]any); ok {
for _, alias := range aliases {
addUsageSceneModelID(ids, alias)
}
}
}
}
return ids
}
func filterPlatformModelsByUsageSceneIDs(models []store.PlatformModel, availableModelIDs map[string]struct{}) []store.PlatformModel {
filtered := make([]store.PlatformModel, 0, len(models))
for _, model := range models {
if platformModelMatchesUsageSceneIDs(model, availableModelIDs) {
filtered = append(filtered, model)
}
}
return filtered
}
func platformModelMatchesUsageSceneIDs(model store.PlatformModel, availableModelIDs map[string]struct{}) bool {
for _, value := range []string{model.ModelName, model.ProviderModelName, model.ModelAlias} {
if _, ok := availableModelIDs[normalizeUsageSceneModelID(value)]; ok {
return true
}
}
for _, alias := range model.LegacyAliases {
if _, ok := availableModelIDs[normalizeUsageSceneModelID(alias)]; ok {
return true
}
}
return false
}
func addUsageSceneModelID(ids map[string]struct{}, value any) {
text, ok := value.(string)
if !ok {
return
}
if normalized := normalizeUsageSceneModelID(text); normalized != "" {
ids[normalized] = struct{}{}
}
}
func normalizeUsageSceneModelID(value string) string {
return strings.ToLower(strings.TrimSpace(value))
}
@@ -0,0 +1,160 @@
package httpapi
import (
"context"
"io"
"log/slog"
"net/http"
"net/http/httptest"
"net/url"
"testing"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
func TestParseModelUsageSceneQuery(t *testing.T) {
t.Parallel()
for _, testCase := range []struct {
name string
query url.Values
want string
wantErr bool
}{
{name: "missing", query: url.Values{}, want: ""},
{name: "desktop", query: url.Values{"usage_scene": {"desktop"}}, want: "desktop"},
{name: "canvas", query: url.Values{"usage_scene": {"canvas_model_node"}}, want: "canvas_model_node"},
{name: "empty", query: url.Values{"usage_scene": {""}}, wantErr: true},
{name: "unknown", query: url.Values{"usage_scene": {"mobile"}}, wantErr: true},
{name: "duplicate", query: url.Values{"usage_scene": {"desktop", "canvas_model_node"}}, wantErr: true},
} {
testCase := testCase
t.Run(testCase.name, func(t *testing.T) {
t.Parallel()
got, err := parseModelUsageSceneQuery(testCase.query)
if testCase.wantErr {
if err == nil {
t.Fatal("expected usage_scene validation error")
}
return
}
if err != nil {
t.Fatalf("parse usage_scene: %v", err)
}
if got != testCase.want {
t.Fatalf("usage_scene = %q, want %q", got, testCase.want)
}
})
}
}
func TestListServerMainUsageSceneModelIDs(t *testing.T) {
t.Parallel()
serverMain := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/api/integration-platform/models/strict/all" {
t.Errorf("path = %q", r.URL.Path)
}
if got := r.URL.Query().Get("usage_scene"); got != "desktop" {
t.Errorf("usage_scene = %q", got)
}
w.Header().Set("Content-Type", "application/json")
_, _ = io.WriteString(w, `{"data":[{"alias":"Qwen3.7-Plus"},{"invocationName":"GPT-5"},{"legacyAliases":["gpt-latest"]}]}`)
}))
defer serverMain.Close()
authenticator := auth.New("secret", serverMain.URL, "")
gateway := &Server{
cfg: config.Config{ServerMainBaseURL: serverMain.URL},
auth: authenticator,
logger: slog.New(slog.NewTextHandler(io.Discard, nil)),
}
ids, err := gateway.listServerMainUsageSceneModelIDs(context.Background(), "desktop")
if err != nil {
t.Fatalf("list server-main usage scene models: %v", err)
}
for _, id := range []string{"qwen3.7-plus", "gpt-5", "gpt-latest"} {
if _, ok := ids[id]; !ok {
t.Errorf("missing normalized model ID %q in %#v", id, ids)
}
}
}
func TestFilterPlatformModelsByUsageSceneIDs(t *testing.T) {
t.Parallel()
models := []store.PlatformModel{
{ID: "allowed-by-invocation", ModelName: "Qwen3.7-Plus", ProviderModelName: "qwen3.7-plus-2026-07-01"},
{ID: "allowed-by-provider-name", ModelName: "GPT-5", ProviderModelName: "gpt-5-2026-07-01"},
{ID: "allowed-by-legacy-alias", ModelName: "renamed-model", LegacyAliases: store.StringList{"old-model"}},
{ID: "desktop-restricted", ModelName: "Qwen3.6-Plus", ProviderModelName: "qwen3.6-plus-2026-04-02"},
}
allowed := map[string]struct{}{
"qwen3.7-plus": {},
"gpt-5-2026-07-01": {},
"old-model": {},
}
filtered := filterPlatformModelsByUsageSceneIDs(models, allowed)
if len(filtered) != 3 {
t.Fatalf("filtered model count = %d, want 3: %#v", len(filtered), filtered)
}
for _, model := range filtered {
if model.ID == "desktop-restricted" {
t.Fatalf("desktop-restricted model remained in filtered catalog: %#v", filtered)
}
}
}
func TestListServerMainUsageSceneModelIDsRejectsMalformedPayload(t *testing.T) {
t.Parallel()
serverMain := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
w.Header().Set("Content-Type", "application/json")
_, _ = io.WriteString(w, `{"success":true}`)
}))
defer serverMain.Close()
gateway := &Server{
cfg: config.Config{ServerMainBaseURL: serverMain.URL},
auth: auth.New("secret", serverMain.URL, ""),
}
if _, err := gateway.listServerMainUsageSceneModelIDs(context.Background(), "desktop"); err == nil {
t.Fatal("expected malformed upstream payload error")
}
}
func TestListServerMainUsageSceneModelIDsFallsBackToDirectServerMainRoute(t *testing.T) {
t.Parallel()
requestedPaths := make([]string, 0, 2)
serverMain := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
requestedPaths = append(requestedPaths, r.URL.Path)
if r.URL.Path == "/api/integration-platform/models/strict/all" {
http.NotFound(w, r)
return
}
w.Header().Set("Content-Type", "application/json")
_, _ = io.WriteString(w, `[{"alias":"Qwen3.7-Plus"}]`)
}))
defer serverMain.Close()
gateway := &Server{
cfg: config.Config{ServerMainBaseURL: serverMain.URL},
auth: auth.New("secret", serverMain.URL, ""),
}
ids, err := gateway.listServerMainUsageSceneModelIDs(context.Background(), "desktop")
if err != nil {
t.Fatalf("list direct server-main usage scene models: %v", err)
}
if _, ok := ids["qwen3.7-plus"]; !ok {
t.Fatalf("missing direct server-main model ID: %#v", ids)
}
if len(requestedPaths) != 2 ||
requestedPaths[0] != "/api/integration-platform/models/strict/all" ||
requestedPaths[1] != "/integration-platform/models/strict/all" {
t.Fatalf("unexpected server-main route probes: %#v", requestedPaths)
}
}
+74 -3
View File
@@ -189,6 +189,13 @@ type TaskListResponse struct {
PageSize int `json:"pageSize" example:"50"`
}
type AdminTaskListResponse struct {
Items []store.AdminGatewayTask `json:"items"`
Total int `json:"total" example:"42"`
Page int `json:"page" example:"1"`
PageSize int `json:"pageSize" example:"50"`
}
type TaskParamPreprocessingLogListResponse struct {
Items []store.TaskParamPreprocessingLog `json:"items"`
}
@@ -208,6 +215,17 @@ type FileUploadResponse struct {
ContentType string `json:"contentType,omitempty" example:"image/png"`
Size int `json:"size,omitempty" example:"1024"`
AssetStorage map[string]interface{} `json:"assetStorage,omitempty"`
Status string `json:"status" example:"success"`
Data FileUploadData `json:"data"`
}
type FileUploadData struct {
ID string `json:"id,omitempty" example:"file_abc123"`
URL string `json:"url,omitempty" example:"/static/uploaded/upload-abc123.png"`
Filename string `json:"filename,omitempty" example:"image.png"`
ContentType string `json:"contentType,omitempty" example:"image/png"`
Size int `json:"size,omitempty" example:"1024"`
AssetStorage map[string]interface{} `json:"assetStorage,omitempty"`
}
type ReplacePlatformModelsRequest struct {
@@ -215,9 +233,12 @@ type ReplacePlatformModelsRequest struct {
}
type TaskAcceptedResponse struct {
TaskID string `json:"taskId" example:"9f4d8f3d-5f5f-4bb7-a4be-344a9f930e25"`
Task store.GatewayTask `json:"task"`
Next TaskNextLinks `json:"next"`
Status string `json:"status" example:"submitted"`
LegacyTaskID string `json:"task_id" example:"9f4d8f3d-5f5f-4bb7-a4be-344a9f930e25"`
TaskID string `json:"taskId" example:"9f4d8f3d-5f5f-4bb7-a4be-344a9f930e25"`
QueryURL string `json:"query_url" example:"/api/v1/ai/result/9f4d8f3d-5f5f-4bb7-a4be-344a9f930e25"`
Task store.GatewayTask `json:"task"`
Next TaskNextLinks `json:"next"`
}
type TaskCancelResponse struct {
@@ -231,6 +252,56 @@ type TaskCancelResponse struct {
type TaskNextLinks struct {
Events string `json:"events" example:"/api/v1/tasks/9f4d8f3d-5f5f-4bb7-a4be-344a9f930e25/events"`
Detail string `json:"detail" example:"/api/v1/tasks/9f4d8f3d-5f5f-4bb7-a4be-344a9f930e25"`
Result string `json:"result" example:"/api/v1/ai/result/9f4d8f3d-5f5f-4bb7-a4be-344a9f930e25"`
}
type EasyAIGeneratedResponse struct {
Status string `json:"status" example:"success" enums:"submitted,process,success,failed"`
TaskID string `json:"task_id" example:"9f4d8f3d-5f5f-4bb7-a4be-344a9f930e25"`
CamelTaskID string `json:"taskId,omitempty" example:"9f4d8f3d-5f5f-4bb7-a4be-344a9f930e25"`
UpstreamTaskID string `json:"upstream_task_id,omitempty" example:"provider-task-123"`
QueryURL string `json:"query_url,omitempty" example:"/api/v1/ai/result/9f4d8f3d-5f5f-4bb7-a4be-344a9f930e25"`
Created int64 `json:"created" example:"1784772000000"`
Message string `json:"message,omitempty"`
Code string `json:"code,omitempty"`
Data []EasyAIMediaOutput `json:"data"`
Output []string `json:"output"`
OutputContent []EasyAIMediaOutput `json:"output_content"`
Cancellable bool `json:"cancellable"`
Submitted bool `json:"submitted"`
Result map[string]interface{} `json:"result,omitempty"`
Usage map[string]interface{} `json:"usage,omitempty"`
}
type EasyAIMediaOutput struct {
Type string `json:"type,omitempty" example:"video" enums:"image,video,audio,file,text"`
URL string `json:"url,omitempty" example:"https://cdn.example.com/output.mp4"`
// B64JSON is retained when result media transfer is disabled.
B64JSON string `json:"b64_json,omitempty"`
ImageURL string `json:"image_url,omitempty"`
VideoURL string `json:"video_url,omitempty"`
AudioURL string `json:"audio_url,omitempty"`
Content string `json:"content,omitempty"`
Text string `json:"text,omitempty"`
MimeType string `json:"mime_type,omitempty"`
Format string `json:"format,omitempty"`
Duration float64 `json:"duration,omitempty"`
Width int `json:"width,omitempty"`
Height int `json:"height,omitempty"`
SourceResolution string `json:"source_resolution,omitempty"`
TargetResolution string `json:"target_resolution,omitempty"`
Metadata map[string]interface{} `json:"metadata,omitempty"`
}
type OpenAIImagesResponse struct {
Created int64 `json:"created" example:"1710000000"`
Data []OpenAIImageData `json:"data"`
}
type OpenAIImageData struct {
URL string `json:"url,omitempty" example:"https://cdn.example.com/output.png"`
B64JSON string `json:"b64_json,omitempty"`
RevisedPrompt string `json:"revised_prompt,omitempty"`
}
type TaskAcceptedEvent struct {
@@ -11,6 +11,7 @@ import (
"net/http/httptest"
"os"
"path/filepath"
"strings"
"testing"
"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) {
request := httptest.NewRequest(http.MethodPost, "/api/v1/chat/completions", nil)
request.Header.Set("X-EasyAI-Conversation-ID", "from-header")
@@ -3,6 +3,7 @@ package httpapi
import (
"encoding/json"
"errors"
"fmt"
"net/http"
"strings"
@@ -71,6 +72,10 @@ func (s *Server) updateRunnerPolicy(w http.ResponseWriter, r *http.Request) {
writeError(w, http.StatusBadRequest, "invalid json body")
return
}
if err := validateRunnerPolicyInput(input); err != nil {
writeError(w, http.StatusBadRequest, err.Error(), "invalid_parameter")
return
}
item, err := s.store.UpsertDefaultRunnerPolicy(r.Context(), input)
if err != nil {
s.logger.Error("update runner policy failed", "error", err)
@@ -80,6 +85,89 @@ func (s *Server) updateRunnerPolicy(w http.ResponseWriter, r *http.Request) {
writeJSON(w, http.StatusOK, item)
}
func validateRunnerPolicyInput(input store.RunnerPolicyInput) error {
rawRules, exists := input.FailoverPolicy["actionRules"]
if !exists {
return nil
}
rules, ok := rawRules.([]any)
if !ok {
return errors.New("failoverPolicy.actionRules must be an array")
}
validActions := map[string]bool{
"stop": true, "next": true, "cooldown_and_next": true, "demote_and_next": true, "disable_and_next": true,
}
for index, rawRule := range rules {
rule, ok := rawRule.(map[string]any)
if !ok {
return fmt.Errorf("failoverPolicy.actionRules[%d] must be an object", index)
}
action := strings.TrimSpace(stringValue(rule["action"]))
if !validActions[action] {
return fmt.Errorf("failoverPolicy.actionRules[%d].action is invalid", index)
}
if target, present := rule["target"]; present {
targetValue := strings.TrimSpace(stringValue(target))
if targetValue != "model" && targetValue != "platform" {
return fmt.Errorf("failoverPolicy.actionRules[%d].target is invalid", index)
}
}
if !runnerActionRuleHasMatch(rule) {
return fmt.Errorf("failoverPolicy.actionRules[%d] must include at least one match condition", index)
}
if action == "cooldown_and_next" && positiveJSONInteger(rule["cooldownSeconds"]) <= 0 {
return fmt.Errorf("failoverPolicy.actionRules[%d].cooldownSeconds must be a positive integer", index)
}
if action == "demote_and_next" && positiveJSONInteger(rule["demoteSteps"]) <= 0 {
return fmt.Errorf("failoverPolicy.actionRules[%d].demoteSteps must be a positive integer", index)
}
}
return nil
}
func runnerActionRuleHasMatch(rule map[string]any) bool {
for _, key := range []string{"categories", "codes", "errorCodes", "statusCodes", "keywords"} {
switch values := rule[key].(type) {
case []any:
for _, value := range values {
if strings.TrimSpace(fmt.Sprint(value)) != "" {
return true
}
}
case []string:
for _, value := range values {
if strings.TrimSpace(value) != "" {
return true
}
}
}
}
return false
}
func positiveJSONInteger(value any) int {
switch typed := value.(type) {
case int:
if typed > 0 {
return typed
}
case int64:
if typed > 0 {
return int(typed)
}
case float64:
if typed > 0 && typed == float64(int(typed)) {
return int(typed)
}
case json.Number:
parsed, err := typed.Int64()
if err == nil && parsed > 0 {
return int(parsed)
}
}
return 0
}
type updatePlatformDynamicPriorityRequest struct {
DynamicPriority *int `json:"dynamicPriority" example:"10"`
Reset bool `json:"reset" example:"false"`
@@ -0,0 +1,57 @@
package httpapi
import (
"strings"
"testing"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
func TestValidateRunnerPolicyInputActionRules(t *testing.T) {
validRule := func() map[string]any {
return map[string]any{
"categories": []any{"rate_limit"},
"action": "cooldown_and_next",
"target": "model",
"cooldownSeconds": float64(60),
}
}
tests := []struct {
name string
mutate func(map[string]any)
wantErr string
}{
{name: "valid"},
{name: "invalid action", mutate: func(rule map[string]any) { rule["action"] = "rotate" }, wantErr: ".action is invalid"},
{name: "invalid target", mutate: func(rule map[string]any) { rule["target"] = "client" }, wantErr: ".target is invalid"},
{name: "empty match", mutate: func(rule map[string]any) { rule["categories"] = []any{} }, wantErr: "at least one match condition"},
{name: "missing cooldown", mutate: func(rule map[string]any) { delete(rule, "cooldownSeconds") }, wantErr: ".cooldownSeconds must be a positive integer"},
{name: "zero cooldown", mutate: func(rule map[string]any) { rule["cooldownSeconds"] = float64(0) }, wantErr: ".cooldownSeconds must be a positive integer"},
{name: "fractional cooldown", mutate: func(rule map[string]any) { rule["cooldownSeconds"] = 1.5 }, wantErr: ".cooldownSeconds must be a positive integer"},
{name: "missing demote steps", mutate: func(rule map[string]any) {
rule["action"] = "demote_and_next"
delete(rule, "cooldownSeconds")
}, wantErr: ".demoteSteps must be a positive integer"},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
rule := validRule()
if test.mutate != nil {
test.mutate(rule)
}
err := validateRunnerPolicyInput(store.RunnerPolicyInput{
FailoverPolicy: map[string]any{"actionRules": []any{rule}},
})
if test.wantErr == "" {
if err != nil {
t.Fatalf("unexpected validation error: %v", err)
}
return
}
if err == nil || !strings.Contains(err.Error(), test.wantErr) {
t.Fatalf("validation error = %v, want substring %q", err, test.wantErr)
}
})
}
}
+10 -2
View File
@@ -120,8 +120,13 @@ func NewServerWithContext(ctx context.Context, cfg config.Config, db *store.Stor
return server.identityRuntime.LegacyJWTEnabled()
}
server.auth.LocalAPIKeyVerifier = db.VerifyLocalAPIKey
server.runner.StartAsyncQueueWorker(ctx)
if cfg.AsyncQueueWorkerEnabled {
server.runner.StartAsyncQueueWorker(ctx)
} else if logger != nil {
logger.Info("asynchronous queue worker disabled for this process")
}
server.runner.StartBillingSettlementWorker(ctx)
server.runner.StartTaskHistoryWorkers(ctx)
server.startLocalTempAssetCleanup(ctx)
server.startOIDCSessionCleanup(ctx)
@@ -183,6 +188,9 @@ func NewServerWithContext(ctx context.Context, cfg config.Config, db *store.Stor
mux.Handle("POST /api/admin/access-rules/batch", server.requireAdmin(auth.PermissionManager, http.HandlerFunc(server.batchAccessRules)))
mux.Handle("PATCH /api/admin/access-rules/{ruleID}", server.requireAdmin(auth.PermissionManager, http.HandlerFunc(server.updateAccessRule)))
mux.Handle("DELETE /api/admin/access-rules/{ruleID}", server.requireAdmin(auth.PermissionManager, http.HandlerFunc(server.deleteAccessRule)))
mux.Handle("GET /api/admin/tasks", server.requireAdmin(auth.PermissionPower, http.HandlerFunc(server.listAdminTasks)))
mux.Handle("GET /api/admin/tasks/{taskID}", server.requireAdmin(auth.PermissionPower, http.HandlerFunc(server.getAdminTask)))
mux.Handle("GET /api/admin/tasks/{taskID}/param-preprocessing", server.requireAdmin(auth.PermissionPower, http.HandlerFunc(server.adminTaskParamPreprocessing)))
mux.Handle("GET /api/v1/api-keys", server.requireUser(auth.PermissionBasic, http.HandlerFunc(server.listAPIKeys)))
mux.Handle("POST /api/v1/api-keys", server.requireUser(auth.PermissionBasic, http.HandlerFunc(server.createAPIKey)))
mux.Handle("GET /api/v1/api-keys/access-rules", server.requireUser(auth.PermissionBasic, http.HandlerFunc(server.listAPIKeyAccessRules)))
@@ -268,7 +276,7 @@ func NewServerWithContext(ctx context.Context, cfg config.Config, db *store.Stor
mux.Handle("POST /api/v1/videos/upscales", server.requireUser(auth.PermissionBasic, server.createVideoUpscaleTask()))
mux.Handle("POST /api/v1/video/upscale", server.requireUser(auth.PermissionBasic, server.createVideoUpscaleTask()))
mux.Handle("POST /api/v1/video/generations", server.requireUser(auth.PermissionBasic, http.HandlerFunc(server.createLegacyVolcesVideoGeneration)))
mux.Handle("GET /api/v1/ai/result/{taskID}", server.requireUser(auth.PermissionBasic, http.HandlerFunc(server.getLegacyVolcesVideoResult)))
mux.Handle("GET /api/v1/ai/result/{taskID}", server.requireUser(auth.PermissionBasic, http.HandlerFunc(server.getEasyAITaskResult)))
mux.Handle("POST /api/v1/song/generations", server.requireUser(auth.PermissionBasic, server.createTask("song.generations", true)))
mux.Handle("POST /api/v1/music/generations", server.requireUser(auth.PermissionBasic, server.createTask("music.generations", true)))
mux.Handle("POST /api/v1/speech/generations", server.requireUser(auth.PermissionBasic, server.createTask("speech.generations", true)))
@@ -1,8 +1,12 @@
package httpapi
import (
"encoding/json"
"net/http"
"net/http/httptest"
"testing"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
func TestOptionalTaskIdempotencyKeyRejectsMultipleValues(t *testing.T) {
@@ -29,3 +33,29 @@ func TestTaskIdempotencyRequestHashIsCanonical(t *testing.T) {
t.Fatal("async response semantics must be part of the request hash")
}
}
func TestWriteIdempotentAsyncTaskReplayKeepsEasyAICompatibilityFields(t *testing.T) {
t.Parallel()
task := store.GatewayTask{
ID: "async-replay-1",
Kind: "images.generations",
Status: "queued",
AsyncMode: true,
}
recorder := httptest.NewRecorder()
writeIdempotentTaskReplay(recorder, task, true)
if recorder.Code != http.StatusAccepted ||
recorder.Header().Get("Idempotent-Replayed") != "true" ||
recorder.Header().Get("X-Gateway-Task-Id") != task.ID {
t.Fatalf("unexpected replay status/headers: code=%d headers=%v", recorder.Code, recorder.Header())
}
var got map[string]any
if err := json.NewDecoder(recorder.Body).Decode(&got); err != nil {
t.Fatal(err)
}
if got["status"] != "submitted" || got["task_id"] != task.ID || got["taskId"] != task.ID {
t.Fatalf("unexpected replay body: %+v", got)
}
}
@@ -61,6 +61,7 @@ func (s *Server) deleteClonedVoice(w http.ResponseWriter, r *http.Request) {
}
result, err := s.runner.DeleteClonedVoice(r.Context(), user, r.PathValue("voiceID"))
if err != nil {
applyRunErrorHeaders(w, err)
writeErrorWithDetails(w, statusFromRunError(err), runErrorMessage(err), runErrorDetails(err), runErrorCode(err))
return
}
@@ -56,12 +56,19 @@ func (s *Server) createVolcesContentsGenerationTask(w http.ResponseWriter, r *ht
// @Param taskID path string true "任务 ID"
// @Success 200 {object} VolcesContentsGenerationTaskResponse
// @Failure 404 {object} VolcesErrorEnvelope
// @Failure 410 {object} VolcesErrorEnvelope
// @Router /api/v1/contents/generations/tasks/{taskID} [get]
func (s *Server) getVolcesContentsGenerationTask(w http.ResponseWriter, r *http.Request) {
task, ok := s.volcesCompatibleTaskForUser(w, r)
if !ok {
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))
}
@@ -143,11 +150,11 @@ func (s *Server) deleteVolcesContentsGenerationTask(w http.ResponseWriter, r *ht
// createLegacyVolcesVideoGeneration godoc
// @Summary 创建 server-main 兼容视频任务
// @Description 兼容 server-main 的 /api/v1/video/generations,返回 submitted 和 task_id;额外保留火山任务字段。
// @Tags volces-compatible
// @Tags media
// @Accept json
// @Produce json
// @Security BearerAuth
// @Success 200 {object} map[string]any
// @Success 200 {object} TaskAcceptedResponse
// @Router /api/v1/video/generations [post]
func (s *Server) createLegacyVolcesVideoGeneration(w http.ResponseWriter, r *http.Request) {
user, ok := auth.UserFromContext(r.Context())
@@ -166,35 +173,46 @@ func (s *Server) createLegacyVolcesVideoGeneration(w http.ResponseWriter, r *htt
return
}
response := volcesCompatibleTask(task)
response["status"] = "submitted"
response["task_id"] = task.ID
writeJSON(w, http.StatusOK, response)
writeJSON(w, http.StatusOK, easyAITaskAcceptedResponseWithExtensions(task, response))
}
// getLegacyVolcesVideoResult godoc
// @Summary 查询 server-main 兼容视频结果
// @Tags volces-compatible
// getEasyAITaskResult godoc
// @Summary 查询 server-main EasyAIClient 兼容任务结果
// @Description 查询当前用户的图片、视频、音频、视频高清和矢量化任务,并返回 EasyAI GeneratedResponse 兼容结构。
// @Tags tasks
// @Produce json
// @Security BearerAuth
// @Param taskID path string true "任务 ID"
// @Success 200 {object} map[string]any
// @Success 200 {object} EasyAIGeneratedResponse
// @Failure 404 {object} EasyAIGeneratedResponse
// @Failure 410 {object} EasyAIGeneratedResponse
// @Router /api/v1/ai/result/{taskID} [get]
func (s *Server) getLegacyVolcesVideoResult(w http.ResponseWriter, r *http.Request) {
task, ok := s.volcesCompatibleTaskForUser(w, r)
if !ok {
func (s *Server) getEasyAITaskResult(w http.ResponseWriter, r *http.Request) {
user, ok := auth.UserFromContext(r.Context())
if !ok || user == nil {
writeEasyAIAsyncError(w, http.StatusUnauthorized, "unauthorized", nil, "unauthorized")
return
}
compat := volcesCompatibleTask(task)
legacyStatus := "process"
switch compat["status"] {
case "succeeded":
legacyStatus = "success"
case "failed", "cancelled":
legacyStatus = "failed"
task, err := s.store.GetTask(r.Context(), strings.TrimSpace(r.PathValue("taskID")))
if err != nil {
if store.IsNotFound(err) {
writeEasyAIAsyncError(w, http.StatusNotFound, "task not found", nil, "not_found")
return
}
s.logger.Error("get EasyAI-compatible task failed", "error", err)
writeError(w, http.StatusInternalServerError, "get task failed", "internal_error")
return
}
writeJSON(w, http.StatusOK, map[string]any{
"status": legacyStatus, "task_id": task.ID, "data": compat["content"], "result": compat,
})
if !runner.TaskAccessibleToUser(task, user) {
writeEasyAIAsyncError(w, http.StatusNotFound, "task not found", nil, "not_found")
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))
}
func (s *Server) createVolcesCompatibleTask(r *http.Request, user *auth.User, body map[string]any) (store.GatewayTask, error) {
@@ -218,7 +236,6 @@ func (s *Server) createVolcesCompatibleTask(r *http.Request, user *auth.User, bo
return store.GatewayTask{}, err
}
task.CompatibilityProtocol = clients.ProtocolVolcesContents
task.CompatibilityPublicID = task.ID
if err := s.runner.EnqueueAsyncTask(r.Context(), task); err != nil {
return store.GatewayTask{}, &clients.ClientError{Code: "enqueue_failed", Message: err.Error(), StatusCode: http.StatusServiceUnavailable, Retryable: true}
}
@@ -253,7 +270,6 @@ func (s *Server) waitForCompatibilitySubmission(r *http.Request, task store.Gate
Message: firstNonEmpty(current.ErrorMessage, current.Error, current.Message),
StatusCode: status,
Retryable: false,
Wire: compatibilitySubmissionWire(current),
}
}
select {
@@ -266,34 +282,6 @@ func (s *Server) waitForCompatibilitySubmission(r *http.Request, task store.Gate
}
}
func compatibilitySubmissionWire(task store.GatewayTask) *clients.WireResponse {
if task.CompatibilitySubmitHTTPStatus == 0 || strings.TrimSpace(task.CompatibilitySourceProtocol) == "" {
return nil
}
headers := map[string][]string{}
for name, value := range task.CompatibilitySubmitHeaders {
switch typed := value.(type) {
case []any:
for _, item := range typed {
if text := strings.TrimSpace(volcesCompatString(item)); text != "" {
headers[name] = append(headers[name], text)
}
}
case []string:
headers[name] = append([]string(nil), typed...)
case string:
headers[name] = []string{typed}
}
}
return &clients.WireResponse{
Protocol: task.CompatibilitySourceProtocol,
StatusCode: task.CompatibilitySubmitHTTPStatus,
Headers: headers,
Body: cloneVolcesCompatibleMap(task.CompatibilitySubmitBody),
Converted: task.CompatibilitySourceProtocol != task.CompatibilityProtocol,
}
}
func (s *Server) volcesCompatibleTaskForUser(w http.ResponseWriter, r *http.Request) (store.GatewayTask, bool) {
user, ok := auth.UserFromContext(r.Context())
if !ok || user == nil {
@@ -322,15 +310,30 @@ func isVolcesCompatibleTask(task store.GatewayTask) bool {
}
func volcesCompatibleTask(task store.GatewayTask) map[string]any {
if volcesTaskUsesNativeProtocol(task) {
if raw := cloneVolcesCompatibleMap(mapFromVolcesCompat(task.Result["raw"])); len(raw) > 0 {
return raw
response := map[string]any{}
for _, key := range []string{"model", "seed", "resolution", "ratio", "duration", "frames", "framespersecond"} {
if task.Result[key] != nil {
response[key] = task.Result[key]
}
if raw := cloneVolcesCompatibleMap(task.RemoteTaskPayload); len(raw) > 0 {
return raw
}
if content, ok := task.Result["content"].(map[string]any); ok && len(content) > 0 {
// Transitional read support for historical canonical results. New
// results only persist data[] and reconstruct this protocol field.
response["content"] = content
} else if data, ok := task.Result["data"].([]any); ok && len(data) > 0 {
if item, ok := data[0].(map[string]any); ok {
content := map[string]any{}
if videoURL := firstNonEmpty(volcesCompatString(item["url"]), volcesCompatString(item["video_url"])); videoURL != "" {
content["video_url"] = videoURL
}
if lastFrameURL := volcesCompatString(item["last_frame_url"]); lastFrameURL != "" {
content["last_frame_url"] = lastFrameURL
}
if len(content) > 0 {
response["content"] = content
}
}
}
response := map[string]any{}
response["id"] = volcesCompatiblePublicID(task)
response["model"] = firstNonEmpty(volcesCompatString(response["model"]), task.Model)
response["status"] = volcesCompatibleTaskStatus(task.Status)
@@ -341,8 +344,10 @@ func volcesCompatibleTask(task store.GatewayTask) map[string]any {
response[key] = task.Request[key]
}
}
if len(task.Usage) > 0 && response["usage"] == nil {
response["usage"] = task.Usage
if usage := volcesCompatibleUsage(task.Usage); len(usage) > 0 {
response["usage"] = usage
} else if legacyUsage, ok := task.Result["usage"].(map[string]any); ok && len(legacyUsage) > 0 {
response["usage"] = legacyUsage
}
if task.Status == "failed" || task.Status == "cancelled" {
response["error"] = map[string]any{"code": firstNonEmpty(task.ErrorCode, strings.ToUpper(task.Status)), "message": firstNonEmpty(task.ErrorMessage, task.Error, task.Message)}
@@ -350,19 +355,28 @@ func volcesCompatibleTask(task store.GatewayTask) map[string]any {
return response
}
func volcesCompatibleCreateResponse(task store.GatewayTask) map[string]any {
if volcesTaskUsesNativeProtocol(task) {
if raw := cloneVolcesCompatibleMap(task.RemoteTaskPayload); len(raw) > 0 {
return raw
func volcesCompatibleUsage(usage map[string]any) map[string]any {
out := map[string]any{}
for outputKey, inputKeys := range map[string][]string{
"prompt_tokens": {"inputTokens", "promptTokens"},
"completion_tokens": {"outputTokens", "completionTokens"},
"total_tokens": {"totalTokens"},
} {
for _, inputKey := range inputKeys {
if value := usage[inputKey]; value != nil {
out[outputKey] = value
break
}
}
}
return out
}
func volcesCompatibleCreateResponse(task store.GatewayTask) map[string]any {
return map[string]any{"id": volcesCompatiblePublicID(task)}
}
func volcesCompatiblePublicID(task store.GatewayTask) string {
if strings.TrimSpace(task.CompatibilityPublicID) != "" {
return strings.TrimSpace(task.CompatibilityPublicID)
}
if volcesTaskUsesNativeProtocol(task) && strings.TrimSpace(task.RemoteTaskID) != "" {
return strings.TrimSpace(task.RemoteTaskID)
}
@@ -381,11 +395,6 @@ func volcesTaskUsesNativeProtocol(task store.GatewayTask) bool {
return false
}
func mapFromVolcesCompat(value any) map[string]any {
result, _ := value.(map[string]any)
return result
}
func volcesCompatibleTaskStatus(status string) string {
switch strings.ToLower(strings.TrimSpace(status)) {
case "succeeded", "success", "completed":
@@ -401,21 +410,6 @@ func volcesCompatibleTaskStatus(status string) string {
}
}
func cloneVolcesCompatibleMap(source map[string]any) map[string]any {
if len(source) == 0 {
return nil
}
raw, err := json.Marshal(source)
if err != nil {
return map[string]any{}
}
var out map[string]any
if err := json.Unmarshal(raw, &out); err != nil {
return map[string]any{}
}
return out
}
func writeVolcesCompatibleTaskError(w http.ResponseWriter, err error) {
status := http.StatusInternalServerError
var staged *gatewayTaskCreationError
@@ -427,6 +421,8 @@ func writeVolcesCompatibleTaskError(w http.ResponseWriter, err error) {
status = clientErr.StatusCode
} else if errors.As(err, &clientErr) {
status = http.StatusBadRequest
} else if errors.Is(err, store.ErrTaskRequestBinaryNotMaterialized) {
status = http.StatusBadRequest
}
if wire := clients.ErrorWireResponse(err); wireResponseMatches(wire, clients.ProtocolVolcesContents) {
writeWireResponse(w, wire)
@@ -7,23 +7,22 @@ import (
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
func TestVolcesCompatibleTaskPreservesNativeOfficialFieldsWithoutGatewayExtensions(t *testing.T) {
func TestVolcesCompatibleTaskBuildsProtocolResponseFromCanonicalResult(t *testing.T) {
now := time.Date(2026, 7, 18, 8, 0, 0, 0, time.UTC)
task := store.GatewayTask{
ID: "gateway-task-1", Kind: "videos.generations", Status: "succeeded", Model: "doubao-seedance-2-0-mini-260615",
RemoteTaskID: "cgt-upstream-1", CreatedAt: now, UpdatedAt: now.Add(time.Second),
Result: map[string]any{
"raw": map[string]any{
"id": "cgt-upstream-1", "model": "doubao-seedance-2-0-mini-260615", "status": "succeeded",
"content": map[string]any{"video_url": "https://example.com/out.mp4"}, "usage": map[string]any{"total_tokens": 9},
"future_official_field": "preserved",
},
"id": "cgt-upstream-1", "model": "doubao-seedance-2-0-mini-260615", "status": "succeeded",
"data": []any{map[string]any{"url": "https://example.com/out.mp4", "type": "video"}},
"future_official_field": "not-whitelisted",
},
Usage: map[string]any{"totalTokens": 9},
Billings: []any{map[string]any{"amount": 3}}, BillingSummary: map[string]any{"currency": "resource"}, FinalChargeAmount: 3,
Attempts: []store.TaskAttempt{{Provider: "volces"}},
}
got := volcesCompatibleTask(task)
if got["id"] != task.RemoteTaskID || got["status"] != "succeeded" || got["future_official_field"] != "preserved" {
if got["id"] != task.RemoteTaskID || got["status"] != "succeeded" || got["future_official_field"] != nil {
t.Fatalf("unexpected compatibility identity/status: %+v", got)
}
content, _ := got["content"].(map[string]any)
+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)
}
}
+174 -31
View File
@@ -13,6 +13,8 @@ type cacheAffinityKeys struct {
Record []string
}
const cacheAffinityMinimumStableMessages = 2
func buildCacheAffinityKey(kind string, modelType string, body map[string]any) string {
return buildCacheAffinityKeys(kind, modelType, body).Primary
}
@@ -37,48 +39,28 @@ func buildCacheAffinityKeys(kind string, modelType string, body map[string]any)
return cacheAffinityKeys{}
}
basePayload := map[string]any{
"kind": kind,
"modelType": modelType,
"version": "v2",
}
if tools, ok := body["tools"]; ok {
basePayload["tools"] = tools
}
if instructions, ok := body["instructions"]; ok {
basePayload["instructions"] = instructions
}
messageHashes := cacheAffinityMessageHashes(body)
if len(messageHashes) > 0 {
if len(messageHashes) >= cacheAffinityMinimumStableMessages {
baseHash := stableJSONHash(basePayload)
keys := make([]string, 0, len(messageHashes))
for index := range messageHashes {
keyPayload := map[string]any{
"base": baseHash,
"messageHashes": messageHashes[:index+1],
"prefixLength": index + 1,
}
keys = append(keys, "prompt_lcp:"+stableJSONHash(keyPayload))
}
keys := cacheAffinityPrefixKeys("prompt_lcp_v2:", baseHash, messageHashes)
if len(keys) == 0 {
return cacheAffinityKeys{}
}
lookup := make([]string, 0, len(keys))
for index := len(keys) - 1; index >= 0; index-- {
lookup = append(lookup, keys[index])
}
lookup := reverseCacheAffinityKeys(keys)
legacyKeys := legacyCacheAffinityPrefixKeys(kind, modelType, body)
lookup = normalizedCacheAffinityKeyList(append(lookup, reverseCacheAffinityKeys(legacyKeys)...))
return cacheAffinityKeys{
Primary: keys[len(keys)-1],
Lookup: lookup,
Record: keys,
}
}
if input, ok := body["input"]; ok {
basePayload["input"] = input
}
if len(basePayload) <= 2 {
return cacheAffinityKeys{}
}
key := "prompt:" + stableJSONHash(basePayload)
return cacheAffinityKeys{Primary: key, Lookup: []string{key}, Record: []string{key}}
return cacheAffinityKeys{}
}
func cacheAffinityPromptHashSupported(kind string, modelType string) bool {
@@ -95,6 +77,146 @@ func cacheAffinityPromptHashSupported(kind string, modelType string) bool {
}
func cacheAffinityMessageHashes(body map[string]any) []string {
if hashes := stringListFromAny(body["messageHashes"]); len(hashes) > 0 {
return hashes
}
if hashes := stringListFromAny(body["message_hashes"]); len(hashes) > 0 {
return hashes
}
messages := cacheAffinityStableMessages(body)
if len(messages) == 0 {
return nil
}
hashes := make([]string, 0, len(messages))
for _, message := range messages {
hash := stableJSONHash(message)
if hash == "" {
continue
}
hashes = append(hashes, hash)
}
return hashes
}
func cacheAffinityStableMessages(body map[string]any) []any {
messages := make([]any, 0)
if instructions := strings.TrimSpace(stringFromAny(body["instructions"])); instructions != "" {
messages = append(messages, map[string]any{"role": "system", "content": instructions})
}
if rawMessages, ok := body["messages"].([]any); ok {
for _, message := range rawMessages {
if normalized := normalizeCacheAffinityMessage(message); normalized != nil {
messages = append(messages, normalized)
}
}
return messages
}
switch input := body["input"].(type) {
case string:
if strings.TrimSpace(input) != "" {
messages = append(messages, map[string]any{"role": "user", "content": input})
}
case []any:
for _, message := range input {
if normalized := normalizeCacheAffinityMessage(message); normalized != nil {
messages = append(messages, normalized)
}
}
}
return messages
}
func normalizeCacheAffinityMessage(value any) any {
message, ok := value.(map[string]any)
if !ok {
if text := strings.TrimSpace(stringFromAny(value)); text != "" {
return map[string]any{"role": "user", "content": text}
}
return nil
}
messageType := strings.TrimSpace(stringFromAny(message["type"]))
switch messageType {
case "function_call":
return map[string]any{
"role": "assistant",
"tool_calls": []any{map[string]any{
"id": firstNonEmptyCacheAffinityString(message["call_id"], message["id"]),
"type": "function",
"function": map[string]any{
"name": message["name"],
"arguments": message["arguments"],
},
}},
}
case "function_call_output":
return map[string]any{
"role": "tool",
"tool_call_id": firstNonEmptyCacheAffinityString(message["call_id"], message["id"]),
"content": message["output"],
}
}
normalized := make(map[string]any, len(message))
for key, item := range message {
switch key {
case "id", "status":
continue
case "type":
if messageType == "message" {
continue
}
}
normalized[key] = item
}
if role := strings.TrimSpace(stringFromAny(normalized["role"])); role != "" {
normalized["role"] = role
}
return normalized
}
func firstNonEmptyCacheAffinityString(values ...any) string {
for _, value := range values {
if text := strings.TrimSpace(stringFromAny(value)); text != "" {
return text
}
}
return ""
}
func cacheAffinityPrefixKeys(prefix string, baseHash string, messageHashes []string) []string {
if len(messageHashes) < cacheAffinityMinimumStableMessages {
return nil
}
keys := make([]string, 0, len(messageHashes)-cacheAffinityMinimumStableMessages+1)
for prefixLength := cacheAffinityMinimumStableMessages; prefixLength <= len(messageHashes); prefixLength++ {
keyPayload := map[string]any{
"base": baseHash,
"messageHashes": messageHashes[:prefixLength],
"prefixLength": prefixLength,
}
keys = append(keys, prefix+stableJSONHash(keyPayload))
}
return keys
}
func legacyCacheAffinityPrefixKeys(kind string, modelType string, body map[string]any) []string {
messageHashes := legacyCacheAffinityMessageHashes(body)
if len(messageHashes) < cacheAffinityMinimumStableMessages {
return nil
}
basePayload := map[string]any{
"kind": kind,
"modelType": modelType,
}
if tools, ok := body["tools"]; ok {
basePayload["tools"] = tools
}
if instructions, ok := body["instructions"]; ok {
basePayload["instructions"] = instructions
}
return cacheAffinityPrefixKeys("prompt_lcp:", stableJSONHash(basePayload), messageHashes)
}
func legacyCacheAffinityMessageHashes(body map[string]any) []string {
if hashes := stringListFromAny(body["messageHashes"]); len(hashes) > 0 {
return hashes
}
@@ -102,20 +224,41 @@ func cacheAffinityMessageHashes(body map[string]any) []string {
return hashes
}
messages, ok := body["messages"].([]any)
if !ok || len(messages) == 0 {
if !ok {
return nil
}
hashes := make([]string, 0, len(messages))
for _, message := range messages {
raw, err := json.Marshal(message)
if err != nil || len(raw) == 0 {
continue
if err == nil && len(raw) > 0 {
hashes = append(hashes, sha256Text(string(raw)))
}
hashes = append(hashes, sha256Text(string(raw)))
}
return hashes
}
func reverseCacheAffinityKeys(keys []string) []string {
reversed := make([]string, 0, len(keys))
for index := len(keys) - 1; index >= 0; index-- {
reversed = append(reversed, keys[index])
}
return reversed
}
func normalizedCacheAffinityKeyList(keys []string) []string {
seen := map[string]bool{}
out := make([]string, 0, len(keys))
for _, key := range keys {
key = strings.TrimSpace(key)
if key == "" || seen[key] {
continue
}
seen[key] = true
out = append(out, key)
}
return out
}
func stableJSONHash(value any) string {
raw, err := json.Marshal(value)
if err != nil || len(raw) == 0 {
@@ -39,7 +39,7 @@ func TestBuildCacheAffinityKeyFallsBackToPromptHash(t *testing.T) {
first := buildCacheAffinityKey("chat.completions", "text_generate", body)
second := buildCacheAffinityKey("chat.completions", "text_generate", body)
if first == "" || first != second || !strings.HasPrefix(first, "prompt_lcp:") {
if first == "" || first != second || !strings.HasPrefix(first, "prompt_lcp_v2:") {
t.Fatalf("expected stable prompt LCP cache affinity key, first=%q second=%q", first, second)
}
}
@@ -54,10 +54,10 @@ func TestBuildCacheAffinityKeysUsesMessagePrefixLCPKeys(t *testing.T) {
"tools": []any{map[string]any{"type": "function", "function": map[string]any{"name": "lookup"}}},
})
if !strings.HasPrefix(first.Primary, "prompt_lcp:") {
if !strings.HasPrefix(first.Primary, "prompt_lcp_v2:") {
t.Fatalf("expected LCP prompt key, got %+v", first)
}
if len(first.Record) != 2 || len(next.Lookup) != 4 {
if len(first.Record) != 1 || len(next.Record) != 3 || len(next.Lookup) != 6 {
t.Fatalf("expected per-message prefix keys, first=%+v next=%+v", first, next)
}
found := false
@@ -71,3 +71,81 @@ func TestBuildCacheAffinityKeysUsesMessagePrefixLCPKeys(t *testing.T) {
t.Fatalf("next lookup keys should include first request prefix key, first=%+v next=%+v", first, next)
}
}
func TestBuildCacheAffinityKeysRequiresTwoStableMessages(t *testing.T) {
systemOnly := buildCacheAffinityKeys("chat.completions", "text_generate", map[string]any{
"messages": []any{
map[string]any{"role": "system", "content": "shared system prompt"},
},
})
firstConversation := buildCacheAffinityKeys("chat.completions", "text_generate", map[string]any{
"messages": []any{
map[string]any{"role": "system", "content": "shared system prompt"},
map[string]any{"role": "user", "content": "conversation A"},
},
})
secondConversation := buildCacheAffinityKeys("chat.completions", "text_generate", map[string]any{
"messages": []any{
map[string]any{"role": "system", "content": "shared system prompt"},
map[string]any{"role": "user", "content": "conversation B"},
},
})
if systemOnly.Primary != "" || len(systemOnly.Lookup) != 0 || len(systemOnly.Record) != 0 {
t.Fatalf("shared system prompt alone must not create affinity, got %+v", systemOnly)
}
if firstConversation.Primary == "" || secondConversation.Primary == "" {
t.Fatalf("two stable messages should create affinity, first=%+v second=%+v", firstConversation, secondConversation)
}
for _, firstKey := range firstConversation.Lookup {
for _, secondKey := range secondConversation.Lookup {
if firstKey == secondKey {
t.Fatalf("different conversations sharing only a system prompt must not intersect: key=%s", firstKey)
}
}
}
}
func TestBuildCacheAffinityKeysNormalizesResponsesToolHistory(t *testing.T) {
tools := []any{map[string]any{
"type": "function",
"function": map[string]any{
"name": "lookup_weather",
},
}}
chat := buildCacheAffinityKeys("chat.completions", "text_generate", map[string]any{
"tools": tools,
"messages": []any{
map[string]any{"role": "system", "content": "Use tools when helpful."},
map[string]any{"role": "user", "content": "Weather in Hangzhou?"},
},
})
responses := buildCacheAffinityKeys("responses", "text_generate", map[string]any{
"tools": tools,
"instructions": "Use tools when helpful.",
"input": "Weather in Hangzhou?",
})
continued := buildCacheAffinityKeys("responses", "text_generate", map[string]any{
"tools": tools,
"instructions": "Use tools when helpful.",
"input": []any{
map[string]any{"role": "user", "content": "Weather in Hangzhou?"},
map[string]any{"type": "function_call", "call_id": "call_weather", "name": "lookup_weather", "arguments": `{"city":"Hangzhou"}`},
map[string]any{"type": "function_call_output", "call_id": "call_weather", "output": `{"temperature":31}`},
},
})
if chat.Primary == "" || chat.Primary != responses.Primary {
t.Fatalf("chat and responses should share the V2 protocol-independent prefix, chat=%+v responses=%+v", chat, responses)
}
found := false
for _, key := range continued.Lookup {
if key == responses.Primary {
found = true
break
}
}
if !found {
t.Fatalf("responses tool continuation should retain its stable prefix, initial=%+v continued=%+v", responses, continued)
}
}
+175 -12
View File
@@ -8,6 +8,7 @@ import (
)
const unsupportedRequestResolutionCode = "unsupported_request_resolution"
const unsupportedRequestAspectRatioCode = "unsupported_request_aspect_ratio"
type requestResolutionRequirement struct {
Kind string
@@ -29,6 +30,15 @@ type videoResolutionReferenceStats struct {
}
func filterRuntimeCandidatesByRequest(kind string, requestedModel string, modelType string, body map[string]any, candidates []store.RuntimeModelCandidate) ([]store.RuntimeModelCandidate, map[string]any, error) {
filtered, resolutionSummary, err := filterRuntimeCandidatesByRequestResolution(kind, requestedModel, modelType, body, candidates)
if err != nil {
return filtered, resolutionSummary, err
}
filtered, aspectRatioSummary, err := filterRuntimeCandidatesByRequestAspectRatio(kind, requestedModel, modelType, body, filtered)
return filtered, mergeCandidateFilterSummaries(resolutionSummary, aspectRatioSummary), err
}
func filterRuntimeCandidatesByRequestResolution(kind string, requestedModel string, modelType string, body map[string]any, candidates []store.RuntimeModelCandidate) ([]store.RuntimeModelCandidate, map[string]any, error) {
requirement, ok := requestResolutionRequirementFor(kind, requestedModel, modelType, body)
if !ok || len(candidates) == 0 {
return candidates, nil, nil
@@ -82,13 +92,13 @@ func requestResolutionRequirementFor(kind string, requestedModel string, modelTy
}
func requestResolutionValue(body map[string]any, modelType string) (string, string) {
if value := normalizedRequestResolution(stringFromAny(body["resolution"])); value != "" {
if value := normalizedRequestResolution(stringFromAny(body["resolution"]), modelType); value != "" {
if _, _, ok := parsePixelSizeString(value); ok {
return "", ""
}
return value, "resolution"
}
size := normalizedRequestResolution(stringFromAny(body["size"]))
size := normalizedRequestResolution(stringFromAny(body["size"]), modelType)
if size == "" {
return "", ""
}
@@ -98,7 +108,7 @@ func requestResolutionValue(body map[string]any, modelType string) (string, stri
return "", ""
}
func normalizedRequestResolution(value string) string {
func normalizedRequestResolution(value string, modelType string) string {
value = strings.TrimSpace(value)
if value == "" || isEmptyParamString(value) {
return ""
@@ -106,9 +116,11 @@ func normalizedRequestResolution(value string) string {
switch strings.ToLower(value) {
case "auto", "automatic", "adaptive", "default":
return ""
default:
}
if _, _, ok := parsePixelSizeString(value); ok {
return value
}
return normalizeMediaResolution(modelType, value)
}
func isResolutionFilteredModelType(modelType string) bool {
@@ -120,10 +132,9 @@ func candidateSupportsRequestResolution(candidate store.RuntimeModelCandidate, r
capability := capabilityForType(effectiveModelCapability(candidate), modelType)
detail := candidateResolutionDetail(candidate, requirement, modelType)
if capability == nil {
detail["reason"] = "capability_missing"
detail["message"] = "候选平台模型未配置对应模型类型能力。"
detail["capabilityPath"] = capabilityPath(modelType, "output_resolutions")
return false, detail
detail["reason"] = "capability_unrestricted"
detail["message"] = "候选平台模型未配置分辨率白名单,按未限制处理。"
return true, detail
}
allowed, configured := outputResolutionAllowedValues(capability["output_resolutions"], requirement.Scopes)
@@ -131,9 +142,19 @@ func candidateSupportsRequestResolution(candidate store.RuntimeModelCandidate, r
detail["capabilityPath"] = capabilityPath(modelType, "output_resolutions")
detail["capabilityValue"] = cloneAny(capability["output_resolutions"])
if !configured {
detail["reason"] = "output_resolutions_missing"
detail["message"] = "候选平台模型未声明 output_resolutions。"
return false, detail
detail["reason"] = "output_resolutions_unrestricted"
detail["message"] = "候选平台模型未声明 output_resolutions,按未限制处理。"
return true, detail
}
if len(allowed) == 0 {
if _, scoped := capability["output_resolutions"].(map[string]any); scoped && len(requirement.Scopes) > 0 && !outputResolutionScopeConfigured(capability["output_resolutions"], requirement.Scopes) {
detail["reason"] = "resolution_scope_not_allowed"
detail["message"] = "候选平台模型未声明当前请求场景的 output_resolutions。"
return false, detail
}
detail["reason"] = "output_resolutions_empty"
detail["message"] = "候选平台模型声明了空分辨率白名单,按未限制处理。"
return true, detail
}
if containsResolution(allowed, requirement.Resolution) {
detail["reason"] = "supported"
@@ -162,7 +183,7 @@ func outputResolutionAllowedValues(value any, scopes []string) ([]string, bool)
for _, raw := range typed {
values = append(values, stringListFromAny(raw)...)
}
return uniqueStringList(values), len(values) > 0
return uniqueStringList(values), true
}
return nil, true
default:
@@ -170,6 +191,148 @@ func outputResolutionAllowedValues(value any, scopes []string) ([]string, bool)
}
}
func outputResolutionScopeConfigured(value any, scopes []string) bool {
typed, ok := value.(map[string]any)
if !ok {
return false
}
for _, scope := range append(append([]string(nil), scopes...), "default", "*", "all") {
if scope == "" {
continue
}
if _, exists := typed[scope]; exists {
return true
}
}
return false
}
type requestAspectRatioRequirement struct {
Kind string
RequestedModel string
ModelType string
AspectRatio string
Source string
Resolution string
}
func filterRuntimeCandidatesByRequestAspectRatio(kind string, requestedModel string, modelType string, body map[string]any, candidates []store.RuntimeModelCandidate) ([]store.RuntimeModelCandidate, map[string]any, error) {
requirement, ok := requestAspectRatioRequirementFor(kind, requestedModel, modelType, body)
if !ok || len(candidates) == 0 {
return candidates, nil, nil
}
filtered := make([]store.RuntimeModelCandidate, 0, len(candidates))
rejected := make([]map[string]any, 0)
supportedRatios := make([]string, 0)
for _, candidate := range candidates {
modelType := firstNonEmptyString(candidate.ModelType, requirement.ModelType)
capability := capabilityForType(effectiveModelCapability(candidate), modelType)
detail := candidateAspectRatioDetail(candidate, requirement, modelType)
if capability == nil {
detail["reason"] = "capability_unrestricted"
filtered = append(filtered, candidate)
continue
}
allowed, configured := outputResolutionAllowedValues(capability["aspect_ratio_allowed"], []string{requirement.Resolution})
detail["allowedAspectRatios"] = allowed
detail["capabilityPath"] = capabilityPath(modelType, "aspect_ratio_allowed")
detail["capabilityValue"] = cloneAny(capability["aspect_ratio_allowed"])
for _, value := range allowed {
appendUniqueString(&supportedRatios, value)
}
if !configured || len(allowed) == 0 || containsResolution(allowed, requirement.AspectRatio) {
detail["reason"] = "supported"
filtered = append(filtered, candidate)
continue
}
detail["reason"] = "aspect_ratio_not_allowed"
detail["message"] = "请求比例不在候选平台模型 aspect_ratio_allowed 中。"
rejected = append(rejected, detail)
}
summary := map[string]any{
"code": unsupportedRequestAspectRatioCode,
"filter": "request_aspect_ratio",
"kind": requirement.Kind,
"requestedModel": requirement.RequestedModel,
"modelType": requirement.ModelType,
"requestedAspectRatio": requirement.AspectRatio,
"aspectRatioSource": requirement.Source,
"resolution": requirement.Resolution,
"capabilityPath": capabilityPath(requirement.ModelType, "aspect_ratio_allowed"),
"candidateCount": len(candidates),
"supportedCandidateCount": len(filtered),
"filteredCandidateCount": len(rejected),
"supportedAspectRatios": uniqueStringList(supportedRatios),
"rejectedCandidates": rejected,
}
if len(filtered) == 0 {
return nil, summary, &store.ModelCandidateUnavailableError{
Code: unsupportedRequestAspectRatioCode,
Message: fmt.Sprintf("请求的%s比例 %s 没有可用平台模型支持,已过滤 %d 个候选平台模型", mediaResourceName(requirement.ModelType), requirement.AspectRatio, len(rejected)),
Details: summary,
}
}
return filtered, summary, nil
}
func requestAspectRatioRequirementFor(kind string, requestedModel string, modelType string, body map[string]any) (requestAspectRatioRequirement, bool) {
if !isResolutionFilteredModelType(modelType) {
return requestAspectRatioRequirement{}, false
}
for _, key := range []string{"aspect_ratio", "aspectRatio", "ratio"} {
raw, exists := body[key]
if !exists {
continue
}
aspectRatio, ok := normalizedAspectRatio(stringFromAny(raw))
if !ok {
return requestAspectRatioRequirement{}, false
}
resolution, _ := requestResolutionValue(body, modelType)
return requestAspectRatioRequirement{
Kind: kind,
RequestedModel: requestedModel,
ModelType: modelType,
AspectRatio: aspectRatio,
Source: key,
Resolution: resolution,
}, true
}
return requestAspectRatioRequirement{}, false
}
func candidateAspectRatioDetail(candidate store.RuntimeModelCandidate, requirement requestAspectRatioRequirement, modelType string) map[string]any {
return map[string]any{
"platformId": candidate.PlatformID,
"platformKey": candidate.PlatformKey,
"platformName": candidate.PlatformName,
"provider": candidate.Provider,
"platformModelId": candidate.PlatformModelID,
"modelName": candidate.ModelName,
"modelAlias": candidate.ModelAlias,
"displayName": candidate.DisplayName,
"providerModelName": candidate.ProviderModelName,
"modelType": modelType,
"requested": map[string]any{
"aspectRatio": requirement.AspectRatio,
"source": requirement.Source,
"resolution": requirement.Resolution,
},
}
}
func mediaResourceName(modelType string) string {
if modelType == "image_generate" || modelType == "image_edit" {
return "图像"
}
if isVideoModelType(modelType) {
return "视频"
}
return "媒体"
}
func containsResolution(values []string, target string) bool {
for _, value := range values {
if strings.EqualFold(strings.TrimSpace(value), strings.TrimSpace(target)) {
@@ -0,0 +1,49 @@
package runner
import (
"fmt"
"strings"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
func filterCandidatesByRequestedPlatform(
candidates []store.RuntimeModelCandidate,
body map[string]any,
) ([]store.RuntimeModelCandidate, error) {
platformID := firstNonEmptyString(
stringFromAny(body["platform_id"]),
stringFromAny(body["platformId"]),
)
platformModelID := firstNonEmptyString(
stringFromAny(body["platform_model_id"]),
stringFromAny(body["platformModelId"]),
)
if platformID == "" && platformModelID == "" {
return candidates, nil
}
filtered := make([]store.RuntimeModelCandidate, 0, len(candidates))
for _, candidate := range candidates {
if platformID != "" && !strings.EqualFold(candidate.PlatformID, platformID) {
continue
}
if platformModelID != "" && !strings.EqualFold(candidate.PlatformModelID, platformModelID) {
continue
}
filtered = append(filtered, candidate)
}
if len(filtered) > 0 {
return filtered, nil
}
requested := platformID
if platformModelID != "" {
requested = platformModelID
}
return nil, fmt.Errorf(
"%w: requested platform %q is unavailable for this model",
store.ErrNoModelCandidate,
requested,
)
}
@@ -0,0 +1,49 @@
package runner
import (
"errors"
"testing"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
func TestFilterCandidatesByRequestedPlatform(t *testing.T) {
candidates := []store.RuntimeModelCandidate{
{PlatformID: "platform-a", PlatformModelID: "model-a"},
{PlatformID: "platform-b", PlatformModelID: "model-b"},
}
t.Run("without selection", func(t *testing.T) {
filtered, err := filterCandidatesByRequestedPlatform(candidates, map[string]any{})
if err != nil || len(filtered) != len(candidates) {
t.Fatalf("unexpected unfiltered candidates: %+v err=%v", filtered, err)
}
})
t.Run("snake case platform id", func(t *testing.T) {
filtered, err := filterCandidatesByRequestedPlatform(candidates, map[string]any{
"platform_id": "platform-b",
})
if err != nil || len(filtered) != 1 || filtered[0].PlatformID != "platform-b" {
t.Fatalf("unexpected platform filter result: %+v err=%v", filtered, err)
}
})
t.Run("camel case platform model id", func(t *testing.T) {
filtered, err := filterCandidatesByRequestedPlatform(candidates, map[string]any{
"platformModelId": "model-a",
})
if err != nil || len(filtered) != 1 || filtered[0].PlatformModelID != "model-a" {
t.Fatalf("unexpected platform model filter result: %+v err=%v", filtered, err)
}
})
t.Run("selection must be accessible and enabled", func(t *testing.T) {
filtered, err := filterCandidatesByRequestedPlatform(candidates, map[string]any{
"platformId": "platform-missing",
})
if len(filtered) != 0 || !errors.Is(err, store.ErrNoModelCandidate) {
t.Fatalf("unexpected unavailable platform result: %+v err=%v", filtered, err)
}
})
}
@@ -0,0 +1,381 @@
package runner
import (
"bytes"
"context"
"encoding/base64"
"fmt"
"image"
imagedraw "image/draw"
"image/jpeg"
_ "image/png"
"io"
"math"
"net/http"
"os"
"strings"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/clients"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
_ "golang.org/x/image/bmp"
"golang.org/x/image/draw"
_ "golang.org/x/image/tiff"
_ "golang.org/x/image/webp"
)
const (
maxInputImageConversionBytes = 32 << 20
maxInputImageConversionPixels = 100_000_000
)
type inputImageResolutionRange struct {
MinLong int
MinShort int
MaxLong int
MaxShort int
}
type inputImageConstraints struct {
Resolution inputImageResolutionRange
MinAspect float64
MaxAspect float64
}
func (s *Service) normalizeVideoInputImages(ctx context.Context, body map[string]any, candidate store.RuntimeModelCandidate) (map[string]any, error) {
constraints, ok := candidateInputImageConstraints(candidate)
if !ok {
return body, nil
}
value, err := s.normalizeVideoInputImageValue(ctx, body, nil, constraints)
if err != nil {
return nil, err
}
out, _ := value.(map[string]any)
if out == nil {
return map[string]any{}, nil
}
return out, nil
}
func (s *Service) normalizeVideoInputImageValue(ctx context.Context, value any, path []string, constraints inputImageConstraints) (any, error) {
switch typed := value.(type) {
case map[string]any:
next := make(map[string]any, len(typed))
for key, item := range typed {
normalized, err := s.normalizeVideoInputImageValue(ctx, item, append(path, key), constraints)
if err != nil {
return nil, err
}
next[key] = normalized
}
return next, nil
case []any:
next := make([]any, 0, len(typed))
for index, item := range typed {
normalized, err := s.normalizeVideoInputImageValue(ctx, item, append(path, fmt.Sprintf("[%d]", index)), constraints)
if err != nil {
return nil, err
}
next = append(next, normalized)
}
return next, nil
case string:
if !imageInputFieldNeedsHydration(path) {
return value, nil
}
normalized, converted, original, target, err := s.normalizeVideoInputImageSource(ctx, typed, constraints)
if err != nil {
param := requestInputImageParam(path)
return nil, &clients.ClientError{
Code: "invalid_parameter",
Message: fmt.Sprintf("输入图片 %s 自动转换失败:%s", param, err.Error()),
Param: param,
StatusCode: http.StatusBadRequest,
Retryable: false,
}
}
if converted && s.logger != nil {
s.logger.Info(
"video input image auto-converted",
"param", requestInputImageParam(path),
"original", fmt.Sprintf("%dx%d", original.X, original.Y),
"converted", fmt.Sprintf("%dx%d", target.X, target.Y),
)
}
return normalized, nil
default:
return value, nil
}
}
func candidateInputImageConstraints(candidate store.RuntimeModelCandidate) (inputImageConstraints, bool) {
modelTypes := []string{candidate.ModelType, "image_to_video", "omni_video"}
seen := map[string]bool{}
for _, modelType := range modelTypes {
modelType = strings.TrimSpace(modelType)
if modelType == "" || seen[modelType] {
continue
}
seen[modelType] = true
values, _ := candidate.Capabilities[modelType].(map[string]any)
if constraints, ok := parseInputImageConstraints(values); ok {
return constraints, true
}
}
return inputImageConstraints{}, false
}
func parseInputImageConstraints(values map[string]any) (inputImageConstraints, bool) {
if values == nil {
return inputImageConstraints{}, false
}
resolution, _ := values["input_image_resolution_range"].(map[string]any)
minimum, _ := resolution["min"].(map[string]any)
maximum, _ := resolution["max"].(map[string]any)
aspect, ok := inputImageNumberPair(values["input_image_aspect_ratio_range"])
constraints := inputImageConstraints{
Resolution: inputImageResolutionRange{
MinLong: int(math.Round(floatFromAny(minimum["long_edge"]))),
MinShort: int(math.Round(floatFromAny(minimum["short_edge"]))),
MaxLong: int(math.Round(floatFromAny(maximum["long_edge"]))),
MaxShort: int(math.Round(floatFromAny(maximum["short_edge"]))),
},
MinAspect: aspect[0],
MaxAspect: aspect[1],
}
if !ok ||
constraints.Resolution.MinLong <= 0 ||
constraints.Resolution.MinShort <= 0 ||
constraints.Resolution.MaxLong < constraints.Resolution.MinLong ||
constraints.Resolution.MaxShort < constraints.Resolution.MinShort ||
constraints.MinAspect <= 0 ||
constraints.MaxAspect < constraints.MinAspect {
return inputImageConstraints{}, false
}
return constraints, true
}
func inputImageNumberPair(value any) ([2]float64, bool) {
switch typed := value.(type) {
case []any:
if len(typed) != 2 {
return [2]float64{}, false
}
pair := [2]float64{floatFromAny(typed[0]), floatFromAny(typed[1])}
return pair, pair[0] > 0 && pair[1] > 0
case []float64:
if len(typed) != 2 {
return [2]float64{}, false
}
return [2]float64{typed[0], typed[1]}, typed[0] > 0 && typed[1] > 0
default:
return [2]float64{}, false
}
}
func (s *Service) normalizeVideoInputImageSource(ctx context.Context, source string, constraints inputImageConstraints) (string, bool, image.Point, image.Point, error) {
source = strings.TrimSpace(source)
if source == "" || strings.HasPrefix(strings.ToLower(source), "asset://") {
return source, false, image.Point{}, image.Point{}, nil
}
payload, err := s.readVideoInputImageBytes(ctx, source)
if err != nil {
return "", false, image.Point{}, image.Point{}, err
}
config, _, err := image.DecodeConfig(bytes.NewReader(payload))
if err != nil || config.Width <= 0 || config.Height <= 0 {
return "", false, image.Point{}, image.Point{}, fmt.Errorf("图片无法解码")
}
if config.Width > 50_000 || config.Height > 50_000 || int64(config.Width)*int64(config.Height) > maxInputImageConversionPixels {
return "", false, image.Point{}, image.Point{}, fmt.Errorf("图片像素数量超过转换上限")
}
original := image.Pt(config.Width, config.Height)
if inputImageWithinConstraints(original, constraints) {
return source, false, original, original, nil
}
decoded, _, err := image.Decode(bytes.NewReader(payload))
if err != nil {
return "", false, original, image.Point{}, fmt.Errorf("图片无法解码")
}
target := resolveInputImageTarget(original, constraints)
canvas := image.NewNRGBA(image.Rect(0, 0, target.X, target.Y))
imagedraw.Draw(canvas, canvas.Bounds(), &image.Uniform{C: averageImageColor(decoded)}, image.Point{}, imagedraw.Src)
sourceBounds := decoded.Bounds()
scale := math.Min(float64(target.X)/float64(sourceBounds.Dx()), float64(target.Y)/float64(sourceBounds.Dy()))
contentWidth := max(1, int(math.Round(float64(sourceBounds.Dx())*scale)))
contentHeight := max(1, int(math.Round(float64(sourceBounds.Dy())*scale)))
left := (target.X - contentWidth) / 2
top := (target.Y - contentHeight) / 2
draw.CatmullRom.Scale(
canvas,
image.Rect(left, top, left+contentWidth, top+contentHeight),
decoded,
sourceBounds,
draw.Over,
nil,
)
var output bytes.Buffer
if err := jpeg.Encode(&output, canvas, &jpeg.Options{Quality: 90}); err != nil {
return "", false, original, target, fmt.Errorf("图片编码失败")
}
if !inputImageWithinConstraints(target, constraints) {
return "", false, original, target, fmt.Errorf("转换结果仍不符合平台限制")
}
return "data:image/jpeg;base64," + base64.StdEncoding.EncodeToString(output.Bytes()), true, original, target, nil
}
func (s *Service) readVideoInputImageBytes(ctx context.Context, source string) ([]byte, error) {
lower := strings.ToLower(source)
if strings.HasPrefix(lower, "data:") || (!strings.Contains(source, "://") && !strings.HasPrefix(source, "/")) {
payload, err := decodeBase64Payload(source)
if err != nil {
return nil, fmt.Errorf("图片数据不是有效 Base64")
}
if len(payload) > maxInputImageConversionBytes {
return nil, fmt.Errorf("图片超过 32 MiB 转换上限")
}
return payload, nil
}
if localPath := s.localPathFromRequestAssetURL(source); localPath != "" {
payload, err := os.ReadFile(localPath)
if err != nil {
return nil, fmt.Errorf("托管图片读取失败")
}
if len(payload) > maxInputImageConversionBytes {
return nil, fmt.Errorf("图片超过 32 MiB 转换上限")
}
return payload, nil
}
if !strings.HasPrefix(lower, "http://") && !strings.HasPrefix(lower, "https://") {
return nil, fmt.Errorf("图片地址格式不受支持")
}
requestCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
defer cancel()
request, err := http.NewRequestWithContext(requestCtx, http.MethodGet, source, nil)
if err != nil {
return nil, fmt.Errorf("图片地址无效")
}
httpClient := generatedAssetHTTPClient(false)
httpClient.Timeout = 10 * time.Second
httpClient.CheckRedirect = func(request *http.Request, via []*http.Request) error {
if len(via) >= 3 {
return fmt.Errorf("图片地址重定向次数过多")
}
if !requestAssetURLIsPublic("", request.URL.String()) {
return fmt.Errorf("图片地址重定向到受限网络")
}
return nil
}
response, err := httpClient.Do(request)
if err != nil {
return nil, fmt.Errorf("图片读取失败")
}
defer response.Body.Close()
if response.StatusCode < 200 || response.StatusCode >= 300 {
return nil, fmt.Errorf("图片读取失败,HTTP %d", response.StatusCode)
}
payload, err := io.ReadAll(io.LimitReader(response.Body, maxInputImageConversionBytes+1))
if err != nil {
return nil, fmt.Errorf("图片读取失败")
}
if len(payload) > maxInputImageConversionBytes {
return nil, fmt.Errorf("图片超过 32 MiB 转换上限")
}
return payload, nil
}
func inputImageWithinConstraints(size image.Point, constraints inputImageConstraints) bool {
longEdge := max(size.X, size.Y)
shortEdge := min(size.X, size.Y)
ratio := float64(size.X) / float64(size.Y)
return longEdge >= constraints.Resolution.MinLong &&
longEdge <= constraints.Resolution.MaxLong &&
shortEdge >= constraints.Resolution.MinShort &&
shortEdge <= constraints.Resolution.MaxShort &&
ratio >= constraints.MinAspect &&
ratio <= constraints.MaxAspect
}
func resolveInputImageTarget(source image.Point, constraints inputImageConstraints) image.Point {
sourceRatio := float64(source.X) / float64(source.Y)
targetRatio := math.Min(constraints.MaxAspect, math.Max(constraints.MinAspect, sourceRatio))
landscape := targetRatio >= 1
normalizedRatio := targetRatio
var sourceShort float64
if landscape {
sourceShort = math.Max(float64(source.Y), float64(source.X)/targetRatio)
} else {
normalizedRatio = 1 / targetRatio
sourceShort = math.Max(float64(source.X), float64(source.Y)*targetRatio)
}
shortEdge := int(math.Round(sourceShort))
shortEdge = min(constraints.Resolution.MaxShort, max(constraints.Resolution.MinShort, shortEdge))
longEdge := int(math.Floor(float64(shortEdge) * normalizedRatio))
if longEdge > constraints.Resolution.MaxLong {
longEdge = constraints.Resolution.MaxLong
shortEdge = int(math.Ceil(float64(longEdge) / normalizedRatio))
}
if longEdge < constraints.Resolution.MinLong {
longEdge = constraints.Resolution.MinLong
shortEdge = int(math.Ceil(float64(longEdge) / normalizedRatio))
}
shortEdge = min(constraints.Resolution.MaxShort, max(constraints.Resolution.MinShort, shortEdge))
longEdge = int(math.Floor(float64(shortEdge) * normalizedRatio))
longEdge = min(constraints.Resolution.MaxLong, max(constraints.Resolution.MinLong, longEdge))
if landscape {
return image.Pt(longEdge, shortEdge)
}
return image.Pt(shortEdge, longEdge)
}
func averageImageColor(source image.Image) imageColor {
bounds := source.Bounds()
stepX := max(1, bounds.Dx()/32)
stepY := max(1, bounds.Dy()/32)
var red, green, blue, count uint64
for y := bounds.Min.Y; y < bounds.Max.Y; y += stepY {
for x := bounds.Min.X; x < bounds.Max.X; x += stepX {
r, g, b, _ := source.At(x, y).RGBA()
red += uint64(r >> 8)
green += uint64(g >> 8)
blue += uint64(b >> 8)
count++
}
}
if count == 0 {
return imageColor{R: 255, G: 255, B: 255, A: 255}
}
return imageColor{R: uint8(red / count), G: uint8(green / count), B: uint8(blue / count), A: 255}
}
type imageColor struct {
R, G, B, A uint8
}
func (c imageColor) RGBA() (r, g, b, a uint32) {
return uint32(c.R) * 0x101, uint32(c.G) * 0x101, uint32(c.B) * 0x101, uint32(c.A) * 0x101
}
func requestInputImageParam(path []string) string {
var builder strings.Builder
for _, segment := range path {
if strings.HasPrefix(segment, "[") {
builder.WriteString(segment)
continue
}
if builder.Len() > 0 {
builder.WriteByte('.')
}
builder.WriteString(segment)
}
if builder.Len() == 0 {
return "image"
}
return builder.String()
}
@@ -0,0 +1,140 @@
package runner
import (
"bytes"
"context"
"encoding/base64"
"image"
"image/color"
"image/jpeg"
"image/png"
"strings"
"testing"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/clients"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
func seedanceImageConstraintCandidate() store.RuntimeModelCandidate {
return store.RuntimeModelCandidate{
ModelType: "image_to_video",
Capabilities: map[string]any{
"image_to_video": map[string]any{
"input_image_resolution_range": map[string]any{
"min": map[string]any{"long_edge": float64(360), "short_edge": float64(360)},
"max": map[string]any{"long_edge": float64(1920), "short_edge": float64(1080)},
},
"input_image_aspect_ratio_range": []any{float64(0.39), float64(2.5)},
},
},
}
}
func pngImageDataURL(t *testing.T, width int, height int) string {
t.Helper()
source := image.NewRGBA(image.Rect(0, 0, width, height))
for y := 0; y < height; y++ {
for x := 0; x < width; x++ {
source.Set(x, y, color.RGBA{R: 80, G: 120, B: 160, A: 255})
}
}
var payload bytes.Buffer
if err := png.Encode(&payload, source); err != nil {
t.Fatal(err)
}
return "data:image/png;base64," + base64.StdEncoding.EncodeToString(payload.Bytes())
}
func decodedImageSize(t *testing.T, source string) image.Point {
t.Helper()
if !strings.HasPrefix(source, "data:image/jpeg;base64,") {
t.Fatalf("expected converted JPEG data URL, got prefix %.32q", source)
}
payload, err := base64.StdEncoding.DecodeString(strings.TrimPrefix(source, "data:image/jpeg;base64,"))
if err != nil {
t.Fatal(err)
}
config, err := jpeg.DecodeConfig(bytes.NewReader(payload))
if err != nil {
t.Fatal(err)
}
return image.Pt(config.Width, config.Height)
}
func TestNormalizeVideoInputImagesResizesOversizedSeedanceImage(t *testing.T) {
service := &Service{}
body := map[string]any{
"content": []any{
map[string]any{
"type": "image_url",
"role": "first_frame",
"image_url": map[string]any{"url": pngImageDataURL(t, 3840, 2160)},
},
},
}
normalized, err := service.normalizeVideoInputImages(context.Background(), body, seedanceImageConstraintCandidate())
if err != nil {
t.Fatal(err)
}
content := normalized["content"].([]any)
item := content[0].(map[string]any)
imageURL := item["image_url"].(map[string]any)["url"].(string)
if got := decodedImageSize(t, imageURL); got != image.Pt(1920, 1080) {
t.Fatalf("unexpected converted size: %v", got)
}
}
func TestNormalizeVideoInputImagesPadsAspectRatioViolation(t *testing.T) {
service := &Service{}
source := pngImageDataURL(t, 1530, 500)
body := map[string]any{"first_frame": source}
normalized, err := service.normalizeVideoInputImages(context.Background(), body, seedanceImageConstraintCandidate())
if err != nil {
t.Fatal(err)
}
if got := decodedImageSize(t, normalized["first_frame"].(string)); got != image.Pt(1530, 612) {
t.Fatalf("unexpected converted size: %v", got)
}
}
func TestNormalizeVideoInputImagesUpscalesBelowMinimumResolution(t *testing.T) {
service := &Service{}
body := map[string]any{"first_frame": pngImageDataURL(t, 200, 100)}
normalized, err := service.normalizeVideoInputImages(context.Background(), body, seedanceImageConstraintCandidate())
if err != nil {
t.Fatal(err)
}
if got := decodedImageSize(t, normalized["first_frame"].(string)); got != image.Pt(720, 360) {
t.Fatalf("unexpected converted size: %v", got)
}
}
func TestNormalizeVideoInputImagesKeepsCompliantImage(t *testing.T) {
service := &Service{}
source := pngImageDataURL(t, 1280, 720)
body := map[string]any{"first_frame": source}
normalized, err := service.normalizeVideoInputImages(context.Background(), body, seedanceImageConstraintCandidate())
if err != nil {
t.Fatal(err)
}
if normalized["first_frame"] != source {
t.Fatal("compliant image should remain unchanged")
}
}
func TestNormalizeVideoInputImagesReturnsActionableDecodeError(t *testing.T) {
service := &Service{}
body := map[string]any{"first_frame": "data:image/png;base64,bm90LWltYWdl"}
_, err := service.normalizeVideoInputImages(context.Background(), body, seedanceImageConstraintCandidate())
if err == nil {
t.Fatal("expected decode error")
}
if clients.ErrorCode(err) != "invalid_parameter" || clients.ErrorParam(err) != "first_frame" {
t.Fatalf("unexpected error contract: code=%s param=%s err=%v", clients.ErrorCode(err), clients.ErrorParam(err), err)
}
}
+10 -22
View File
@@ -74,7 +74,7 @@ func (s *Service) rateLimitReservations(ctx context.Context, user *auth.User, ca
"groupKey": group.GroupKey,
"name": group.Name,
},
group.RateLimitPolicy,
store.NormalizeRateLimitPolicy(group.RateLimitPolicy),
body,
)...)
}
@@ -82,27 +82,15 @@ func (s *Service) rateLimitReservations(ctx context.Context, user *auth.User, ca
}
func effectiveRateLimitPolicy(candidate store.RuntimeModelCandidate) map[string]any {
policy := candidate.PlatformRateLimitPolicy
if strings.TrimSpace(candidate.RuntimePolicySetID) != "" {
policy = candidate.RuntimeRateLimitPolicy
} else if hasRules(candidate.RuntimeRateLimitPolicy) {
policy = mergeMap(policy, candidate.RuntimeRateLimitPolicy)
}
if _, hasOverride := candidate.RuntimePolicyOverride["rateLimitPolicy"]; hasOverride {
nested, _ := candidate.RuntimePolicyOverride["rateLimitPolicy"].(map[string]any)
if len(nested) == 0 {
policy = nil
} else {
policy = mergeMap(policy, nested)
}
}
if hasRules(candidate.ModelRateLimitPolicy) {
policy = mergeMap(policy, candidate.ModelRateLimitPolicy)
}
if hasRules(policy) {
return policy
}
return nil
return store.EffectiveRateLimitPolicy(store.EffectiveRateLimitPolicyInput{
BasePolicy: candidate.BaseRateLimitPolicy,
PlatformPolicy: candidate.PlatformRateLimitPolicy,
RuntimePolicy: candidate.RuntimeRateLimitPolicy,
RuntimePolicyExplicit: candidate.RuntimePolicyExplicit,
RuntimePolicyOverride: candidate.RateLimitRuntimeOverride,
ModelPolicy: candidate.ModelRateLimitPolicy,
ModelPolicyMode: candidate.ModelRateLimitPolicyMode,
})
}
func effectiveRetryPolicy(candidate store.RuntimeModelCandidate) map[string]any {
+4 -2
View File
@@ -102,8 +102,10 @@ func TestEffectiveRateLimitPolicyTreatsEmptyRuntimePolicyAsUnlimited(t *testing.
PlatformRateLimitPolicy: map[string]any{"rules": []any{
map[string]any{"metric": "rpm", "limit": 500},
}},
RuntimePolicySetID: "runtime-policy-1",
RuntimeRateLimitPolicy: map[string]any{"rules": []any{}},
RuntimePolicySetID: "runtime-policy-1",
RuntimePolicyExplicit: true,
RuntimeRateLimitPolicy: map[string]any{"rules": []any{}},
ModelRateLimitPolicyMode: "inherit",
})
if hasRules(policy) {
@@ -0,0 +1,173 @@
package runner
import (
"testing"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
func TestParamProcessorTreatsNonRatioValuesAsUnspecifiedForImagesAndVideos(t *testing.T) {
for _, modelType := range []string{"image_generate", "video_generate"} {
t.Run(modelType, func(t *testing.T) {
result := preprocessRequestWithLog(mediaKindForTest(modelType), map[string]any{
"aspect_ratio": "auto",
"aspectRatio": "adaptive",
"ratio": "keep_ratio",
}, store.RuntimeModelCandidate{
ModelType: modelType,
Capabilities: map[string]any{
modelType: map[string]any{
"aspect_ratio_allowed": []any{"adaptive"},
},
},
})
for _, key := range []string{"aspect_ratio", "aspectRatio", "ratio"} {
if _, ok := result.Body[key]; ok {
t.Fatalf("%s should be removed for non-ratio input: %+v", key, result.Body)
}
}
if result.Err != nil {
t.Fatalf("non-ratio input should use provider default instead of failing: %v", result.Err)
}
})
}
}
func TestParamProcessorNormalizesImageKSizeAndPreservesProviderResolutionChoice(t *testing.T) {
tests := []struct {
name string
capability map[string]any
wantSize string
}{
{
name: "dimension upstream",
capability: map[string]any{
"output_resolutions": []any{"1K", "2K", "4K"},
"aspect_ratio_allowed": []any{"16:9"},
},
wantSize: "2048x1152",
},
{
name: "resolution upstream",
capability: map[string]any{
"output_resolutions": []any{"1K", "2K"},
"aspect_ratio_allowed": []any{"16:9"},
"size_param_format": "resolution",
},
wantSize: "2K",
},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
result := preprocessRequestWithLog("images.generations", map[string]any{
"size": "2k",
"aspect_ratio": "16 : 9",
}, store.RuntimeModelCandidate{
ModelType: "image_generate",
Capabilities: map[string]any{
"image_generate": test.capability,
},
})
if result.Err != nil {
t.Fatalf("preprocess image size: %v", result.Err)
}
if result.Body["resolution"] != "2K" || result.Body["size"] != test.wantSize {
t.Fatalf("unexpected normalized size: %+v", result.Body)
}
if result.Body["width"] != 2048 || result.Body["height"] != 1152 || result.Body["aspect_ratio"] != "16:9" {
t.Fatalf("unexpected normalized dimensions: %+v", result.Body)
}
})
}
}
func TestParamProcessorNormalizesVideo4KSize(t *testing.T) {
result := preprocessRequestWithLog("videos.generations", map[string]any{
"size": "4K",
"aspect_ratio": "auto",
}, store.RuntimeModelCandidate{
ModelType: "video_generate",
Capabilities: map[string]any{
"video_generate": map[string]any{
"output_resolutions": []any{"2160p"},
"aspect_ratio_allowed": []any{"16:9"},
},
},
})
if result.Body["resolution"] != "2160p" {
t.Fatalf("video 4K size should normalize to 2160p: %+v", result.Body)
}
if _, ok := result.Body["aspect_ratio"]; ok {
t.Fatalf("video auto ratio should stay unspecified: %+v", result.Body)
}
}
func TestCandidateFilterIgnoresNonRatioValuesAndFiltersRealRatios(t *testing.T) {
candidates := []store.RuntimeModelCandidate{
candidateWithAspectRatios("square", "image_generate", "1:1"),
candidateWithAspectRatios("landscape", "image_generate", "16:9"),
}
filtered, summary, err := filterRuntimeCandidatesByRequest("images.generations", "demo-image", "image_generate", map[string]any{
"aspect_ratio": "auto",
}, candidates)
if err != nil || len(filtered) != 2 || summary != nil {
t.Fatalf("auto ratio should not participate in candidate matching: filtered=%+v summary=%+v err=%v", filtered, summary, err)
}
filtered, summary, err = filterRuntimeCandidatesByRequest("images.generations", "demo-image", "image_generate", map[string]any{
"aspect_ratio": "16 : 9",
}, candidates)
if err != nil || len(filtered) != 1 || filtered[0].PlatformKey != "landscape" {
t.Fatalf("real ratio should filter candidates: filtered=%+v summary=%+v err=%v", filtered, summary, err)
}
}
func TestCandidateFilterNormalizesKSizeAndAllowsUnconfiguredResolution(t *testing.T) {
candidates := []store.RuntimeModelCandidate{
candidateWithImageResolutions("one-k", "1K"),
{
PlatformID: "platform-unrestricted",
PlatformKey: "unrestricted",
PlatformModelID: "model-unrestricted",
ModelName: "demo-image",
ModelType: "image_generate",
Capabilities: map[string]any{
"image_generate": map[string]any{},
},
},
}
filtered, _, err := filterRuntimeCandidatesByRequest("images.generations", "demo-image", "image_generate", map[string]any{
"size": "2k",
}, candidates)
if err != nil {
t.Fatalf("2k size should be accepted and normalized for matching: %v", err)
}
if len(filtered) != 1 || filtered[0].PlatformKey != "unrestricted" {
t.Fatalf("unexpected candidates for normalized 2K size: %+v", filtered)
}
}
func candidateWithAspectRatios(platformKey string, modelType string, ratios ...string) store.RuntimeModelCandidate {
return store.RuntimeModelCandidate{
PlatformID: "platform-" + platformKey,
PlatformKey: platformKey,
PlatformModelID: "model-" + platformKey,
ModelName: "demo-media",
ModelType: modelType,
Capabilities: map[string]any{
modelType: map[string]any{
"aspect_ratio_allowed": stringsToAny(ratios),
},
},
}
}
func mediaKindForTest(modelType string) string {
if modelType == "image_generate" {
return "images.generations"
}
return "videos.generations"
}
+55 -32
View File
@@ -10,7 +10,6 @@ import (
const (
volcesOutputTokenThreshold = 10240
volcesOutputCapabilityErrorCode = "model_capability_configuration_error"
volcesOutputInvalidParameterCode = "invalid_parameter"
)
@@ -25,20 +24,25 @@ func (outputTokenLimitProcessor) ShouldProcess(_ map[string]any, modelType strin
func (outputTokenLimitProcessor) Process(params map[string]any, modelType string, context *paramProcessContext) bool {
modelMax, sourceType, capabilityValue, ok := candidateMaxOutputTokens(context.candidate, modelType)
path := capabilityPath(sourceType, "max_output_tokens")
if parameter, value, explicit := explicitOutputTokenParameter(context.kind, params); explicit {
if _, valid := positiveInteger(value); !valid {
return context.reject(
"OutputTokenLimitProcessor",
parameter,
value,
fmt.Sprintf("%s must be a positive integer", parameter),
path,
capabilityValue,
)
}
}
if !ok {
return context.reject(
"OutputTokenLimitProcessor",
outputTokenParameter(context.kind),
nil,
"火山引擎文本候选未配置正整数 max_output_tokens,已禁止执行该候选。",
path,
capabilityValue,
)
return true
}
if context.kind == "chat.completions" {
if value, explicit := nonNullParameter(params, "max_tokens"); explicit {
if parsed, valid := positiveInteger(value); !valid || parsed > modelMax {
if parsed, _ := positiveInteger(value); parsed > modelMax {
return context.reject("OutputTokenLimitProcessor", "max_tokens", value, fmt.Sprintf("max_tokens must be a positive integer no greater than the selected Volcengine model limit (%d)", modelMax), path, capabilityValue)
}
return true
@@ -57,7 +61,7 @@ func (outputTokenLimitProcessor) Process(params map[string]any, modelType string
}
if value, explicit := nonNullParameter(params, "max_output_tokens"); explicit {
if parsed, valid := positiveInteger(value); !valid || parsed > modelMax {
if parsed, _ := positiveInteger(value); parsed > modelMax {
return context.reject("OutputTokenLimitProcessor", "max_output_tokens", value, fmt.Sprintf("max_output_tokens must be a positive integer no greater than the selected Volcengine model limit (%d)", modelMax), path, capabilityValue)
}
return true
@@ -76,8 +80,35 @@ func filterRuntimeCandidatesByOutputTokens(kind string, requestedModel string, m
if !isOpenAITextGenerationKind(kind) || !isTextOutputModelType(modelType) || len(candidates) == 0 {
return candidates, nil, nil
}
hasVolcesCandidate := false
for _, candidate := range candidates {
if isVolcesCandidate(candidate) {
hasVolcesCandidate = true
break
}
}
if !hasVolcesCandidate {
return candidates, nil, nil
}
explicitParameter, explicitValue, hasExplicitParameter := explicitOutputTokenParameter(kind, body)
explicitTokens := 0
if hasExplicitParameter {
var valid bool
explicitTokens, valid = positiveInteger(explicitValue)
if !valid {
summary := map[string]any{
"filter": "volces_output_token_limit", "kind": kind, "requestedModel": requestedModel, "modelType": modelType,
"candidateCount": len(candidates), "supportedCandidateCount": 0, "filteredCandidateCount": len(candidates),
"code": volcesOutputInvalidParameterCode, "reason": "invalid_output_token_parameter",
"parameter": explicitParameter, "requestedValue": explicitValue,
}
message := fmt.Sprintf("%s must be a positive integer", explicitParameter)
return nil, summary, &store.ModelCandidateUnavailableError{Code: volcesOutputInvalidParameterCode, Message: message, Details: summary}
}
}
filtered := make([]store.RuntimeModelCandidate, 0, len(candidates))
rejected := make([]map[string]any, 0)
missingCapability := make([]map[string]any, 0)
invalidExplicit := false
for _, candidate := range candidates {
if !isVolcesCandidate(candidate) {
@@ -92,15 +123,15 @@ func filterRuntimeCandidatesByOutputTokens(kind string, requestedModel string, m
}
if !configured {
detail["reason"] = "max_output_tokens_missing"
rejected = append(rejected, detail)
missingCapability = append(missingCapability, detail)
filtered = append(filtered, candidate)
continue
}
if parameter, value, explicit := explicitOutputTokenParameter(kind, body); explicit {
parsed, valid := positiveInteger(value)
if !valid || (parameter != "max_completion_tokens" && parsed > modelMax) {
if hasExplicitParameter {
if explicitParameter != "max_completion_tokens" && explicitTokens > modelMax {
detail["reason"] = "requested_output_tokens_exceed_capability"
detail["parameter"] = parameter
detail["requestedValue"] = value
detail["parameter"] = explicitParameter
detail["requestedValue"] = explicitValue
invalidExplicit = true
rejected = append(rejected, detail)
continue
@@ -108,26 +139,25 @@ func filterRuntimeCandidatesByOutputTokens(kind string, requestedModel string, m
}
filtered = append(filtered, candidate)
}
if len(rejected) == 0 {
if len(rejected) == 0 && len(missingCapability) == 0 {
return filtered, nil, nil
}
summary := map[string]any{
"filter": "volces_output_token_limit", "kind": kind, "requestedModel": requestedModel, "modelType": modelType,
"candidateCount": len(candidates), "supportedCandidateCount": len(filtered), "filteredCandidateCount": len(rejected),
"missingCapabilityCandidateCount": len(missingCapability), "passthroughMissingCapability": len(missingCapability) > 0,
"threshold": volcesOutputTokenThreshold, "formula": "third=floor(modelMaxOutputTokens/3); default=third>=10240?third:modelMaxOutputTokens",
"rejectedCandidates": rejected,
"rejectedCandidates": rejected, "missingCapabilityCandidates": missingCapability,
}
if len(filtered) > 0 {
return filtered, summary, nil
}
code := volcesOutputCapabilityErrorCode
message := "所有火山引擎文本候选都缺少有效的 max_output_tokens 能力配置"
if invalidExplicit {
code = volcesOutputInvalidParameterCode
message = "请求的输出 token 上限超过所有可用火山引擎候选的模型能力"
message := "请求的输出 token 上限超过所有可用火山引擎候选的模型能力"
summary["code"] = volcesOutputInvalidParameterCode
return nil, summary, &store.ModelCandidateUnavailableError{Code: volcesOutputInvalidParameterCode, Message: message, Details: summary}
}
summary["code"] = code
return nil, summary, &store.ModelCandidateUnavailableError{Code: code, Message: message, Details: summary}
return filtered, summary, nil
}
func defaultVolcesOutputTokens(modelMax int) int {
@@ -195,13 +225,6 @@ func explicitOutputTokenParameter(kind string, body map[string]any) (string, any
return "max_completion_tokens", value, ok
}
func outputTokenParameter(kind string) string {
if kind == "responses" {
return "max_output_tokens"
}
return "max_tokens"
}
func isOpenAITextGenerationKind(kind string) bool {
return kind == "chat.completions" || kind == "responses"
}
@@ -72,32 +72,90 @@ func TestVolcesOutputProcessorPreservesExplicitLimits(t *testing.T) {
}
}
func TestVolcesOutputCandidateFilterSupportsFailover(t *testing.T) {
func TestVolcesOutputCandidateFilterPreservesMissingCapabilityOrder(t *testing.T) {
missing := volcTextCandidate(nil)
missing.PlatformID = "missing"
valid := volcTextCandidate(32768)
valid.PlatformID = "valid"
filtered, summary, err := filterRuntimeCandidatesByOutputTokens("chat.completions", "demo", "text_generate", map[string]any{"messages": []any{}}, []store.RuntimeModelCandidate{missing, valid})
if err != nil || len(filtered) != 1 || filtered[0].PlatformID != "valid" {
t.Fatalf("expected missing capability candidate to be skipped: filtered=%+v summary=%+v err=%v", filtered, summary, err)
if err != nil || len(filtered) != 2 || filtered[0].PlatformID != "missing" || filtered[1].PlatformID != "valid" {
t.Fatalf("expected missing capability candidate order to be preserved: filtered=%+v summary=%+v err=%v", filtered, summary, err)
}
result := preprocessRequestWithLog("chat.completions", map[string]any{"messages": []any{}}, filtered[0])
if summary["missingCapabilityCandidateCount"] != 1 || summary["passthroughMissingCapability"] != true {
t.Fatalf("expected missing capability audit summary, got %+v", summary)
}
missingResult := preprocessRequestWithLog("chat.completions", map[string]any{"messages": []any{}}, filtered[0])
if missingResult.Err != nil || len(missingResult.Log.Changes) != 0 {
t.Fatalf("missing capability candidate must pass through unchanged: %+v", missingResult)
}
if _, ok := missingResult.Body["max_tokens"]; ok {
t.Fatalf("missing capability candidate must not inject max_tokens: %+v", missingResult.Body)
}
result := preprocessRequestWithLog("chat.completions", map[string]any{"messages": []any{}}, filtered[1])
if result.Body["max_tokens"] != 10922 {
t.Fatalf("failover candidate must recalculate its own default, got %+v", result.Body)
t.Fatalf("configured candidate must calculate its own default, got %+v", result.Body)
}
}
func TestVolcesOutputCandidateFilterRejectsMissingAndExceededCapabilities(t *testing.T) {
_, _, err := filterRuntimeCandidatesByOutputTokens("responses", "demo", "text_generate", map[string]any{"input": "hello"}, []store.RuntimeModelCandidate{volcTextCandidate(nil)})
if store.ModelCandidateErrorCode(err) != volcesOutputCapabilityErrorCode {
t.Fatalf("expected capability configuration error, got %v", err)
func TestVolcesOutputCandidateFilterPassesThroughMissingCapability(t *testing.T) {
filtered, summary, err := filterRuntimeCandidatesByOutputTokens("responses", "demo", "text_generate", map[string]any{"input": "hello"}, []store.RuntimeModelCandidate{volcTextCandidate(nil)})
if err != nil || len(filtered) != 1 {
t.Fatalf("expected missing capability candidate to remain available: filtered=%+v summary=%+v err=%v", filtered, summary, err)
}
_, _, err = filterRuntimeCandidatesByOutputTokens("chat.completions", "demo", "text_generate", map[string]any{"messages": []any{}, "max_tokens": 32769}, []store.RuntimeModelCandidate{volcTextCandidate(32768)})
result := preprocessRequestWithLog("responses", map[string]any{"input": "hello"}, filtered[0])
if result.Err != nil {
t.Fatalf("missing capability preprocessing must pass through: %v", result.Err)
}
if _, ok := result.Body["max_output_tokens"]; ok {
t.Fatalf("missing capability preprocessing must not inject max_output_tokens: %+v", result.Body)
}
}
func TestVolcesOutputCandidateFilterRejectsExceededKnownCapability(t *testing.T) {
_, _, err := filterRuntimeCandidatesByOutputTokens("chat.completions", "demo", "text_generate", map[string]any{"messages": []any{}, "max_tokens": 32769}, []store.RuntimeModelCandidate{volcTextCandidate(32768)})
if store.ModelCandidateErrorCode(err) != volcesOutputInvalidParameterCode {
t.Fatalf("expected invalid_parameter for explicit limit, got %v", err)
}
}
func TestVolcesOutputCandidateFilterAllowsUnknownCapabilityForExplicitLimit(t *testing.T) {
missing := volcTextCandidate(nil)
missing.PlatformID = "missing"
configured := volcTextCandidate(32768)
configured.PlatformID = "configured"
filtered, summary, err := filterRuntimeCandidatesByOutputTokens(
"chat.completions",
"demo",
"text_generate",
map[string]any{"messages": []any{}, "max_tokens": 65536},
[]store.RuntimeModelCandidate{missing, configured},
)
if err != nil || len(filtered) != 1 || filtered[0].PlatformID != "missing" {
t.Fatalf("expected unknown capability candidate to remain available: filtered=%+v summary=%+v err=%v", filtered, summary, err)
}
}
func TestVolcesOutputProcessorValidatesExplicitLimitWithoutCapability(t *testing.T) {
valid := preprocessRequestWithLog("chat.completions", map[string]any{"messages": []any{}, "max_tokens": 2048}, volcTextCandidate(nil))
if valid.Err != nil || valid.Body["max_tokens"] != 2048 {
t.Fatalf("valid explicit limit must pass through: %+v", valid)
}
invalid := preprocessRequestWithLog("chat.completions", map[string]any{"messages": []any{}, "max_tokens": 0}, volcTextCandidate(nil))
if invalid.Err == nil {
t.Fatalf("invalid explicit limit must still be rejected: %+v", invalid)
}
_, _, err := filterRuntimeCandidatesByOutputTokens(
"chat.completions",
"demo",
"text_generate",
map[string]any{"messages": []any{}, "max_tokens": 0},
[]store.RuntimeModelCandidate{volcTextCandidate(nil)},
)
if store.ModelCandidateErrorCode(err) != volcesOutputInvalidParameterCode {
t.Fatalf("candidate filter must reject an invalid explicit limit, got %v", err)
}
}
func TestNonVolcesOutputProcessorDoesNotInject(t *testing.T) {
candidate := volcTextCandidate(131072)
candidate.Provider = "openai"
+132 -37
View File
@@ -3,6 +3,7 @@ package runner
import (
"fmt"
"math"
"strconv"
"strings"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
@@ -13,28 +14,43 @@ type resolutionNormalizeProcessor struct{}
func (resolutionNormalizeProcessor) Name() string { return "ResolutionNormalizeProcessor" }
func (resolutionNormalizeProcessor) ShouldProcess(params map[string]any, modelType string, context *paramProcessContext) bool {
if stringFromAny(params["resolution"]) != "" {
return false
if normalizeMediaResolution(modelType, stringFromAny(params["resolution"])) != "" {
return true
}
size := stringFromAny(params["size"])
if size == "" {
return false
}
return isImageResolution(modelType, size) || isVideoResolution(modelType, size)
return normalizeMediaResolution(modelType, stringFromAny(params["size"])) != ""
}
func (resolutionNormalizeProcessor) Process(params map[string]any, modelType string, context *paramProcessContext) bool {
if resolution := normalizeMediaResolution(modelType, stringFromAny(params["resolution"])); resolution != "" {
before := params["resolution"]
params["resolution"] = resolution
context.resolution = resolution
if stringFromAny(before) != resolution {
_, capabilityValue := capabilityEvidence(context.modelCapability, modelType, "output_resolutions")
context.recordChange(
"ResolutionNormalizeProcessor",
"adjust",
"resolution",
before,
resolution,
"resolution 使用了可识别的媒体分辨率格式,已归一为标准大小写和值。",
capabilityPath(modelType, "output_resolutions"),
capabilityValue,
)
}
return true
}
size := stringFromAny(params["size"])
if stringFromAny(params["resolution"]) == "" && (isImageResolution(modelType, size) || isVideoResolution(modelType, size)) {
if resolution := normalizeMediaResolution(modelType, size); resolution != "" {
_, capabilityValue := capabilityEvidence(context.modelCapability, modelType, "output_resolutions")
params["resolution"] = size
context.resolution = size
params["resolution"] = resolution
context.resolution = resolution
context.recordChange(
"ResolutionNormalizeProcessor",
"set",
"resolution",
nil,
size,
resolution,
"size 使用分辨率格式,归一到 resolution 供后续能力校验和计费使用。",
capabilityPath(modelType, "output_resolutions"),
capabilityValue,
@@ -48,30 +64,69 @@ type aspectRatioProcessor struct{}
func (aspectRatioProcessor) Name() string { return "AspectRatioProcessor" }
func (aspectRatioProcessor) ShouldProcess(params map[string]any, modelType string, context *paramProcessContext) bool {
return modelType != "text_generate" && (stringFromAny(params["aspect_ratio"]) != "" || stringFromAny(params["size"]) != "")
if modelType == "text_generate" {
return false
}
for _, key := range []string{"aspect_ratio", "aspectRatio", "ratio"} {
if _, ok := params[key]; ok {
return true
}
}
return stringFromAny(params["size"]) != ""
}
func (aspectRatioProcessor) Process(params map[string]any, modelType string, context *paramProcessContext) bool {
capability := capabilityForType(context.modelCapability, modelType)
if capability == nil {
aspectRatio, ratioSource, ratioPresent := requestAspectRatio(params)
if ratioPresent {
originalAspectRatio := aspectRatio
normalized, valid := normalizedAspectRatio(aspectRatio)
if !valid {
before := map[string]any{}
for _, key := range []string{"aspect_ratio", "aspectRatio", "ratio"} {
if value, ok := params[key]; ok {
before[key] = cloneAny(value)
delete(params, key)
}
}
aspectRatio = ""
context.aspectRatio = ""
context.recordChange(
"AspectRatioProcessor",
"remove",
ratioSource,
before,
nil,
"比例不是有效的正数 x:x 格式,按未指定比例处理。",
"",
nil,
)
return true
} else {
aspectRatio = normalized
params["aspect_ratio"] = normalized
for _, key := range []string{"aspectRatio", "ratio"} {
delete(params, key)
}
if ratioSource != "aspect_ratio" || normalized != originalAspectRatio {
context.recordChange(
"AspectRatioProcessor",
"adjust",
"aspect_ratio",
map[string]any{ratioSource: originalAspectRatio},
normalized,
"比例已归一为标准 x:x 格式。",
"",
nil,
)
}
}
}
if !ratioPresent {
return true
}
aspectRatio := stringFromAny(params["aspect_ratio"])
if isEmptyParamString(aspectRatio) {
before := params["aspect_ratio"]
delete(params, "aspect_ratio")
context.aspectRatio = ""
context.recordChange(
"AspectRatioProcessor",
"remove",
"aspect_ratio",
before,
nil,
"aspect_ratio 是空值字符串,不能作为有效比例传给上游。",
"",
nil,
)
capability := capabilityForType(context.modelCapability, modelType)
if capability == nil {
return true
}
@@ -174,6 +229,15 @@ func (aspectRatioProcessor) Process(params map[string]any, modelType string, con
return true
}
func requestAspectRatio(params map[string]any) (string, string, bool) {
for _, key := range []string{"aspect_ratio", "aspectRatio", "ratio"} {
if value, ok := params[key]; ok {
return stringFromAny(value), key, true
}
}
return "", "aspect_ratio", false
}
func validateExplicitAspectRatio(aspectRatio string, capability map[string]any, allowed []string, modelType string) (string, bool) {
if isVideoModelType(modelType) && len(allowed) > 0 && !containsString(allowed, aspectRatio) {
if aspectRatio == "adaptive" || aspectRatio == "keep_ratio" {
@@ -194,17 +258,14 @@ func (imageSizeProcessor) ShouldProcess(params map[string]any, modelType string,
if modelType != "image_generate" && modelType != "image_edit" {
return false
}
if _, _, ok := imageDimensionsFromParams(params); !ok {
return false
}
capability := capabilityForType(context.modelCapability, modelType)
return capability != nil && imageSizeCapabilityConfigured(capability)
_, _, ok := imageDimensionsFromParams(params)
return ok
}
func (imageSizeProcessor) Process(params map[string]any, modelType string, context *paramProcessContext) bool {
capability := capabilityForType(context.modelCapability, modelType)
if capability == nil {
return true
capability = map[string]any{}
}
width, height, ok := imageDimensionsFromParams(params)
if !ok {
@@ -222,7 +283,11 @@ func (imageSizeProcessor) Process(params map[string]any, modelType string, conte
params["width"] = width
params["height"] = height
_, _, sizeHasPixelDimensions := parsePixelSizeString(stringFromAny(params["size"]))
if stringFromAny(params["aspect_ratio"]) == "" || sizeHasPixelDimensions {
_, _, resolutionHasPixelDimensions := parsePixelSizeString(stringFromAny(params["resolution"]))
explicitWidthHeight := positiveIntegerFromAny(params["width"]) > 0 && positiveIntegerFromAny(params["height"]) > 0 &&
normalizeImageResolution(stringFromAny(params["size"])) == "" &&
normalizeImageResolution(stringFromAny(params["resolution"])) == ""
if sizeHasPixelDimensions || resolutionHasPixelDimensions || (stringFromAny(params["aspect_ratio"]) == "" && explicitWidthHeight) {
aspectRatio := aspectRatioFromDimensions(width, height)
allowed := aspectRatioAllowed(capability["aspect_ratio_allowed"], firstNonEmptyString(stringFromAny(params["resolution"]), context.resolution))
if processed, ok := validateAndAdjustAspectRatio(aspectRatio, capability, allowed); ok && processed != "" {
@@ -274,7 +339,37 @@ func imageDimensionsFromParams(params map[string]any) (int, int, bool) {
if width > 0 && height > 0 {
return width, height, true
}
return parsePixelSizeString(stringFromAny(params["resolution"]))
if width, height, ok := parsePixelSizeString(stringFromAny(params["resolution"])); ok {
return width, height, true
}
resolution := firstNonEmptyString(
normalizeImageResolution(stringFromAny(params["resolution"])),
normalizeImageResolution(stringFromAny(params["size"])),
)
if resolution == "" {
return 0, 0, false
}
return imageDimensionsForResolution(resolution, stringFromAny(params["aspect_ratio"]))
}
func imageDimensionsForResolution(resolution string, aspectRatio string) (int, int, bool) {
resolution = normalizeImageResolution(resolution)
if resolution == "" {
return 0, 0, false
}
magnitude, err := strconv.Atoi(strings.TrimSuffix(resolution, "K"))
if err != nil || magnitude <= 0 {
return 0, 0, false
}
baseSide := magnitude * 1024
ratio, validRatio := aspectRatioNumber(aspectRatio)
if !validRatio {
return baseSide, baseSide, true
}
if ratio >= 1 {
return baseSide, max(1, int(math.Round(float64(baseSide)/ratio))), true
}
return max(1, int(math.Round(float64(baseSide)*ratio))), baseSide, true
}
func imageSizeCapabilityConfigured(capability map[string]any) bool {
@@ -393,23 +393,39 @@ func numberPair(value any) ([2]float64, bool) {
}
func validAspectRatio(value string) bool {
if value == "adaptive" || value == "keep_ratio" {
return true
}
_, ok := aspectRatioNumber(value)
_, ok := normalizedAspectRatio(value)
return ok
}
func aspectRatioNumber(value string) (float64, bool) {
parts := strings.Split(value, ":")
func normalizedAspectRatio(value string) (string, bool) {
compact := strings.Map(func(r rune) rune {
switch r {
case ' ', '\t', '\r', '\n':
return -1
default:
return r
}
}, strings.TrimSpace(value))
parts := strings.Split(compact, ":")
if len(parts) != 2 {
return 0, false
return "", false
}
width := parsePositiveFloat(parts[0])
height := parsePositiveFloat(parts[1])
if width <= 0 || height <= 0 {
return "", false
}
return parts[0] + ":" + parts[1], true
}
func aspectRatioNumber(value string) (float64, bool) {
normalized, ok := normalizedAspectRatio(value)
if !ok {
return 0, false
}
parts := strings.Split(normalized, ":")
width := parsePositiveFloat(parts[0])
height := parsePositiveFloat(parts[1])
return width / height, true
}
@@ -537,11 +553,53 @@ func isEmptyParamString(value string) bool {
}
func isImageResolution(modelType string, value string) bool {
return (modelType == "image_generate" || modelType == "image_edit") && containsString([]string{"1K", "2K", "3K", "4K", "8K"}, value)
if modelType != "image_generate" && modelType != "image_edit" {
return false
}
return normalizeImageResolution(value) != ""
}
func isVideoResolution(modelType string, value string) bool {
return isVideoModelType(modelType) && containsString([]string{"480p", "720p", "1080p", "1440p", "2160p"}, value)
return isVideoModelType(modelType) && normalizeVideoResolution(value) != ""
}
func normalizeMediaResolution(modelType string, value string) string {
if modelType == "image_generate" || modelType == "image_edit" {
return normalizeImageResolution(value)
}
if isVideoModelType(modelType) {
return normalizeVideoResolution(value)
}
return ""
}
func normalizeImageResolution(value string) string {
value = strings.ToLower(strings.TrimSpace(value))
if len(value) < 2 || value[len(value)-1] != 'k' {
return ""
}
magnitude, err := strconv.Atoi(value[:len(value)-1])
if err != nil || magnitude <= 0 {
return ""
}
return strconv.Itoa(magnitude) + "K"
}
func normalizeVideoResolution(value string) string {
switch strings.ToLower(strings.TrimSpace(value)) {
case "480p":
return "480p"
case "720p":
return "720p"
case "1080p":
return "1080p"
case "1440p":
return "1440p"
case "2160p", "4k":
return "2160p"
default:
return ""
}
}
func isVideoModelType(modelType string) bool {
+4
View File
@@ -31,6 +31,10 @@ func (s *Service) Estimate(ctx context.Context, kind string, model string, body
if err != nil {
return EstimateResult{}, err
}
candidates, err = filterCandidatesByRequestedPlatform(candidates, body)
if err != nil {
return EstimateResult{}, err
}
candidates, _, err = filterRuntimeCandidatesByRequest(kind, model, modelType, body, candidates)
if err != nil {
return EstimateResult{}, err
+3 -1
View File
@@ -19,6 +19,8 @@ type httpClientCache struct {
custom map[string]*http.Client
}
const providerHTTPClientTimeout = 10 * time.Minute
func newHTTPClientCache() *httpClientCache {
return &httpClientCache{
none: newHTTPClient(nil),
@@ -67,7 +69,7 @@ func newHTTPClient(proxy func(*http.Request) (*url.URL, error)) *http.Client {
transport := http.DefaultTransport.(*http.Transport).Clone()
transport.Proxy = proxy
return &http.Client{
Timeout: 120 * time.Second,
Timeout: providerHTTPClientTimeout,
Transport: transport,
}
}
+8
View File
@@ -5,11 +5,19 @@ import (
"net/http"
"net/http/httptest"
"testing"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
func TestProviderHTTPClientTimeoutAllowsLongRunningMediaRequests(t *testing.T) {
client := newHTTPClient(nil)
if client.Timeout != 10*time.Minute {
t.Fatalf("unexpected provider HTTP timeout: got %s want %s", client.Timeout, 10*time.Minute)
}
}
func TestPlatformProxyModeNoneIgnoresEnvironmentProxy(t *testing.T) {
var proxyHits int
proxy := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+212 -20
View File
@@ -11,6 +11,7 @@ import (
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
"github.com/google/uuid"
"github.com/jackc/pgx/v5"
"github.com/riverqueue/river"
"github.com/riverqueue/river/riverdriver/riverpgxv5"
"github.com/riverqueue/river/rivermigrate"
@@ -23,6 +24,12 @@ type asyncTaskArgs struct {
TaskID string `json:"task_id" river:"unique"`
}
type asyncExecutionClient interface {
Start(context.Context) error
Stop(context.Context) error
StopAndCancel(context.Context) error
}
func (asyncTaskArgs) Kind() string { return "gateway_task_run" }
type asyncTaskWorker struct {
@@ -87,17 +94,70 @@ func (s *Service) startRiverQueue(ctx context.Context) error {
return err
}
workers := river.NewWorkers()
if err := river.AddWorkerSafely(workers, &asyncTaskWorker{service: s}); err != nil {
controlClient, err := river.NewClient(driver, &river.Config{
ID: asyncWorkerID() + "-control",
Logger: s.logger,
TestOnly: s.cfg.AppEnv == "test",
})
if err != nil {
return err
}
riverClient, err := river.NewClient(driver, &river.Config{
ID: asyncWorkerID(),
snapshot, err := s.loadAsyncWorkerCapacity(ctx)
if err != nil {
return fmt.Errorf("calculate initial async worker capacity: %w", err)
}
executionClient, err := s.makeAsyncExecutionClient(snapshot.Capacity)
if err != nil {
return err
}
if err := executionClient.Start(ctx); err != nil {
cleanupCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
_ = executionClient.StopAndCancel(cleanupCtx)
cancel()
return err
}
s.riverMu.Lock()
s.riverControlClient = controlClient
s.riverExecutionClient = executionClient
s.riverWorkerCapacity = snapshot.Capacity
s.riverDrainingClients = make(map[asyncExecutionClient]struct{})
s.riverMu.Unlock()
s.observeAsyncWorkerCapacity(snapshot)
s.logger.Info("async worker capacity initialized",
"capacity", snapshot.Capacity,
"desiredCapacity", snapshot.Desired,
"hardLimit", snapshot.HardLimit,
"enabledModels", snapshot.EnabledModels,
"unlimitedModels", snapshot.UnlimitedModels,
"modelDesired", snapshot.ModelDesired,
"enabledGroups", snapshot.EnabledGroups,
"unlimitedGroups", snapshot.UnlimitedGroups,
"groupDesired", snapshot.GroupDesired,
"capped", snapshot.Capped,
)
if err := s.recoverAsyncRiverJobs(ctx); err != nil {
cleanupCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
_ = executionClient.StopAndCancel(cleanupCtx)
cancel()
return err
}
go s.refreshAsyncWorkerCapacity(ctx)
go s.stopAsyncWorkersOnShutdown(ctx)
return nil
}
func (s *Service) newRiverAsyncExecutionClient(capacity int) (*river.Client[pgx.Tx], error) {
workers := river.NewWorkers()
if err := river.AddWorkerSafely(workers, &asyncTaskWorker{service: s}); err != nil {
return nil, err
}
return river.NewClient(riverpgxv5.New(s.store.Pool()), &river.Config{
ID: fmt.Sprintf("%s-exec-%d-%d", asyncWorkerID(), capacity, time.Now().UnixNano()),
JobTimeout: -1,
Logger: s.logger,
CompletedJobRetentionPeriod: 24 * time.Hour,
Queues: map[string]river.QueueConfig{
asyncTaskQueueName: {MaxWorkers: 32},
asyncTaskQueueName: {MaxWorkers: capacity},
},
// Provider-backed media jobs commonly poll for 10-20 minutes. River may
// execute a still-running job again once this window elapses, so keep the
@@ -106,32 +166,160 @@ func (s *Service) startRiverQueue(ctx context.Context) error {
TestOnly: s.cfg.AppEnv == "test",
Workers: workers,
})
}
func (s *Service) makeAsyncExecutionClient(capacity int) (asyncExecutionClient, error) {
if s.asyncClientFactory != nil {
return s.asyncClientFactory(capacity)
}
return s.newRiverAsyncExecutionClient(capacity)
}
func (s *Service) loadAsyncWorkerCapacity(ctx context.Context) (store.AsyncWorkerCapacitySnapshot, error) {
if s.asyncCapacityLoader != nil {
return s.asyncCapacityLoader(ctx, s.cfg.AsyncWorkerHardLimit)
}
return s.store.AsyncWorkerCapacity(ctx, s.cfg.AsyncWorkerHardLimit)
}
func (s *Service) refreshAsyncWorkerCapacity(ctx context.Context) {
ticker := time.NewTicker(time.Duration(s.cfg.AsyncWorkerRefreshIntervalSeconds) * time.Second)
defer ticker.Stop()
for {
select {
case <-ctx.Done():
return
case <-ticker.C:
s.resizeAsyncWorkerCapacity(ctx)
}
}
}
func (s *Service) resizeAsyncWorkerCapacity(ctx context.Context) {
snapshot, err := s.loadAsyncWorkerCapacity(ctx)
if err != nil {
return err
s.observeAsyncWorkerResize("refresh_failed")
s.logger.Warn("refresh async worker capacity failed; keeping current client", "error", err)
return
}
s.riverClient = riverClient
if err := riverClient.Start(ctx); err != nil {
return err
s.riverMu.RLock()
currentCapacity := s.riverWorkerCapacity
s.riverMu.RUnlock()
s.observeAsyncWorkerCapacity(snapshot)
if snapshot.Capacity == currentCapacity {
return
}
if err := s.recoverAsyncRiverJobs(ctx); err != nil {
return err
newClient, err := s.makeAsyncExecutionClient(snapshot.Capacity)
if err != nil {
s.observeAsyncWorkerResize("create_failed")
s.logger.Warn("create replacement async worker client failed; keeping current client",
"error", err, "currentCapacity", currentCapacity, "desiredCapacity", snapshot.Capacity)
return
}
go func() {
<-ctx.Done()
if err := newClient.Start(ctx); err != nil {
cleanupCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
_ = newClient.StopAndCancel(cleanupCtx)
cancel()
s.observeAsyncWorkerResize("start_failed")
s.logger.Warn("start replacement async worker client failed; keeping current client",
"error", err, "currentCapacity", currentCapacity, "desiredCapacity", snapshot.Capacity)
return
}
if ctx.Err() != nil {
cleanupCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
_ = newClient.StopAndCancel(cleanupCtx)
cancel()
return
}
s.riverMu.Lock()
oldClient := s.riverExecutionClient
s.riverExecutionClient = newClient
s.riverWorkerCapacity = snapshot.Capacity
if oldClient != nil {
if s.riverDrainingClients == nil {
s.riverDrainingClients = make(map[asyncExecutionClient]struct{})
}
s.riverDrainingClients[oldClient] = struct{}{}
}
s.riverMu.Unlock()
s.observeAsyncWorkerResize("success")
s.logger.Info("async worker capacity resized",
"previousCapacity", currentCapacity,
"capacity", snapshot.Capacity,
"desiredCapacity", snapshot.Desired,
"hardLimit", snapshot.HardLimit,
"modelDesired", snapshot.ModelDesired,
"groupDesired", snapshot.GroupDesired,
"capped", snapshot.Capped,
)
if oldClient != nil {
go s.drainAsyncWorkerClient(oldClient)
}
}
func (s *Service) drainAsyncWorkerClient(client asyncExecutionClient) {
stopCtx, cancel := context.WithTimeout(context.Background(), time.Hour)
defer cancel()
if err := client.Stop(stopCtx); err != nil {
s.logger.Warn("gracefully drain previous async worker client failed", "error", err)
}
s.riverMu.Lock()
delete(s.riverDrainingClients, client)
s.riverMu.Unlock()
}
func (s *Service) stopAsyncWorkersOnShutdown(ctx context.Context) {
<-ctx.Done()
s.riverMu.Lock()
clients := make([]asyncExecutionClient, 0, 1+len(s.riverDrainingClients))
if s.riverExecutionClient != nil {
clients = append(clients, s.riverExecutionClient)
}
for client := range s.riverDrainingClients {
clients = append(clients, client)
}
s.riverExecutionClient = nil
s.riverDrainingClients = nil
s.riverMu.Unlock()
for _, client := range clients {
stopCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
if err := riverClient.StopAndCancel(stopCtx); err != nil {
if err := client.StopAndCancel(stopCtx); err != nil {
s.logger.Warn("stop river async queue failed", "error", err)
}
}()
return nil
cancel()
}
}
func (s *Service) asyncControlClient() *river.Client[pgx.Tx] {
s.riverMu.RLock()
defer s.riverMu.RUnlock()
return s.riverControlClient
}
func (s *Service) observeAsyncWorkerCapacity(snapshot store.AsyncWorkerCapacitySnapshot) {
observer, ok := s.billingMetrics.(interface {
SetAsyncWorkerCapacity(current, desired, hardLimit int, capped bool)
})
if ok {
observer.SetAsyncWorkerCapacity(snapshot.Capacity, snapshot.Desired, snapshot.HardLimit, snapshot.Capped)
}
}
func (s *Service) observeAsyncWorkerResize(outcome string) {
observer, ok := s.billingMetrics.(interface {
ObserveAsyncWorkerResize(string)
})
if ok {
observer.ObserveAsyncWorkerResize(outcome)
}
}
func (s *Service) EnqueueAsyncTask(ctx context.Context, task store.GatewayTask) error {
if s.riverClient == nil {
riverClient := s.asyncControlClient()
if riverClient == nil {
return errors.New("river async queue is not started")
}
result, err := s.riverClient.Insert(ctx, asyncTaskArgs{TaskID: task.ID}, asyncTaskInsertOpts(task))
result, err := riverClient.Insert(ctx, asyncTaskArgs{TaskID: task.ID}, asyncTaskInsertOpts(task))
if err != nil {
return err
}
@@ -151,12 +339,16 @@ func (s *Service) RunAsyncTask(ctx context.Context, task store.GatewayTask, user
}
func (s *Service) recoverAsyncRiverJobs(ctx context.Context) error {
riverClient := s.asyncControlClient()
if riverClient == nil {
return errors.New("river async queue is not started")
}
items, err := s.store.ListRecoverableAsyncTasks(ctx, 1000)
if err != nil {
return err
}
for _, item := range items {
result, err := s.riverClient.Insert(ctx, asyncTaskArgs{TaskID: item.ID}, asyncTaskRecoveryInsertOpts(item, time.Now()))
result, err := riverClient.Insert(ctx, asyncTaskArgs{TaskID: item.ID}, asyncTaskRecoveryInsertOpts(item, time.Now()))
if err != nil {
return err
}
@@ -0,0 +1,189 @@
package runner
import (
"context"
"errors"
"io"
"log/slog"
"sync/atomic"
"testing"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
type fakeAsyncExecutionClient struct {
startErr error
stopGate <-chan struct{}
started atomic.Int64
stopped atomic.Int64
stopCancelled atomic.Int64
}
func (c *fakeAsyncExecutionClient) Start(context.Context) error {
c.started.Add(1)
return c.startErr
}
func (c *fakeAsyncExecutionClient) Stop(ctx context.Context) error {
c.stopped.Add(1)
if c.stopGate == nil {
return nil
}
select {
case <-c.stopGate:
return nil
case <-ctx.Done():
return ctx.Err()
}
}
func (c *fakeAsyncExecutionClient) StopAndCancel(context.Context) error {
c.stopCancelled.Add(1)
return nil
}
func TestResizeAsyncWorkerCapacityStartsReplacementBeforeGracefulDrain(t *testing.T) {
oldClient := &fakeAsyncExecutionClient{}
drainGate := make(chan struct{})
oldClient.stopGate = drainGate
newClient := &fakeAsyncExecutionClient{}
service := asyncWorkerManagerTestService(1)
service.riverExecutionClient = oldClient
service.asyncCapacityLoader = fixedAsyncCapacity(3)
service.asyncClientFactory = func(capacity int) (asyncExecutionClient, error) {
if capacity != 3 {
t.Fatalf("factory capacity=%d, want=3", capacity)
}
return newClient, nil
}
service.resizeAsyncWorkerCapacity(context.Background())
if newClient.started.Load() != 1 {
t.Fatalf("replacement start count=%d, want=1", newClient.started.Load())
}
if service.riverExecutionClient != newClient || service.riverWorkerCapacity != 3 {
t.Fatalf("replacement was not installed: capacity=%d client=%T", service.riverWorkerCapacity, service.riverExecutionClient)
}
waitForAtomicValue(t, &oldClient.stopped, 1)
if oldClient.stopCancelled.Load() != 0 {
t.Fatal("graceful drain cancelled an already running old task")
}
close(drainGate)
waitForDrainingClients(t, service, 0)
}
func TestResizeAsyncWorkerCapacityFailuresKeepCurrentClient(t *testing.T) {
t.Run("refresh failure", func(t *testing.T) {
oldClient := &fakeAsyncExecutionClient{}
service := asyncWorkerManagerTestService(4)
service.riverExecutionClient = oldClient
service.asyncCapacityLoader = func(context.Context, int) (store.AsyncWorkerCapacitySnapshot, error) {
return store.AsyncWorkerCapacitySnapshot{}, errors.New("database unavailable")
}
service.asyncClientFactory = func(int) (asyncExecutionClient, error) {
t.Fatal("factory must not run after refresh failure")
return nil, nil
}
service.resizeAsyncWorkerCapacity(context.Background())
if service.riverExecutionClient != oldClient || service.riverWorkerCapacity != 4 {
t.Fatal("refresh failure replaced the current client")
}
})
t.Run("create failure", func(t *testing.T) {
oldClient := &fakeAsyncExecutionClient{}
service := asyncWorkerManagerTestService(4)
service.riverExecutionClient = oldClient
service.asyncCapacityLoader = fixedAsyncCapacity(8)
service.asyncClientFactory = func(int) (asyncExecutionClient, error) {
return nil, errors.New("factory failed")
}
service.resizeAsyncWorkerCapacity(context.Background())
if service.riverExecutionClient != oldClient || service.riverWorkerCapacity != 4 {
t.Fatal("create failure replaced the current client")
}
})
t.Run("start failure", func(t *testing.T) {
oldClient := &fakeAsyncExecutionClient{}
failedClient := &fakeAsyncExecutionClient{startErr: errors.New("start failed")}
service := asyncWorkerManagerTestService(4)
service.riverExecutionClient = oldClient
service.asyncCapacityLoader = fixedAsyncCapacity(8)
service.asyncClientFactory = func(int) (asyncExecutionClient, error) {
return failedClient, nil
}
service.resizeAsyncWorkerCapacity(context.Background())
if service.riverExecutionClient != oldClient || service.riverWorkerCapacity != 4 {
t.Fatal("start failure replaced the current client")
}
if failedClient.stopCancelled.Load() != 1 {
t.Fatal("failed replacement client was not cleaned up")
}
})
}
func TestStopAsyncWorkersIncludesCurrentAndDrainingClients(t *testing.T) {
current := &fakeAsyncExecutionClient{}
draining := &fakeAsyncExecutionClient{}
service := asyncWorkerManagerTestService(2)
service.riverExecutionClient = current
service.riverDrainingClients[draining] = struct{}{}
ctx, cancel := context.WithCancel(context.Background())
cancel()
service.stopAsyncWorkersOnShutdown(ctx)
if current.stopCancelled.Load() != 1 || draining.stopCancelled.Load() != 1 {
t.Fatalf("shutdown cancellations current=%d draining=%d, want 1/1", current.stopCancelled.Load(), draining.stopCancelled.Load())
}
}
func asyncWorkerManagerTestService(capacity int) *Service {
return &Service{
cfg: config.Config{
AsyncWorkerHardLimit: 2048,
AsyncWorkerRefreshIntervalSeconds: 5,
},
logger: slog.New(slog.NewTextHandler(io.Discard, nil)),
riverDrainingClients: make(map[asyncExecutionClient]struct{}),
riverWorkerCapacity: capacity,
riverExecutionClient: nil,
}
}
func fixedAsyncCapacity(capacity int) func(context.Context, int) (store.AsyncWorkerCapacitySnapshot, error) {
return func(context.Context, int) (store.AsyncWorkerCapacitySnapshot, error) {
return store.AsyncWorkerCapacitySnapshot{Capacity: capacity, Desired: capacity, HardLimit: 2048}, nil
}
}
func waitForAtomicValue(t *testing.T, value *atomic.Int64, want int64) {
t.Helper()
deadline := time.Now().Add(time.Second)
for time.Now().Before(deadline) {
if value.Load() == want {
return
}
time.Sleep(time.Millisecond)
}
t.Fatalf("atomic value=%d, want=%d", value.Load(), want)
}
func waitForDrainingClients(t *testing.T, service *Service, want int) {
t.Helper()
deadline := time.Now().Add(time.Second)
for time.Now().Before(deadline) {
service.riverMu.RLock()
count := len(service.riverDrainingClients)
service.riverMu.RUnlock()
if count == want {
return
}
time.Sleep(time.Millisecond)
}
service.riverMu.RLock()
count := len(service.riverDrainingClients)
service.riverMu.RUnlock()
t.Fatalf("draining client count=%d, want=%d", count, want)
}
+8
View File
@@ -170,6 +170,14 @@ func attemptMetrics(candidate store.RuntimeModelCandidate, attemptNo int, simula
metrics["cacheAffinityCachedInputTokens"] = candidate.CacheAffinity.CachedInputTokens
metrics["cacheAffinityBoost"] = candidate.CacheAffinity.Boost
metrics["cacheAffinityApplied"] = candidate.CacheAffinity.Applied
metrics["cacheAffinityMatched"] = candidate.CacheAffinity.Applied
metrics["cacheAffinityCandidateCount"] = candidate.CacheAffinity.CandidateCount
if candidate.CacheAffinity.MatchedPrefixDepth > 0 {
metrics["cacheAffinityMatchedPrefixDepth"] = candidate.CacheAffinity.MatchedPrefixDepth
}
if candidate.CacheAffinity.OverrideReason != "" {
metrics["cacheAffinityOverrideReason"] = candidate.CacheAffinity.OverrideReason
}
}
if attemptNo > 0 {
metrics["attempt"] = attemptNo
+19 -10
View File
@@ -13,16 +13,19 @@ func TestAttemptMetricsIncludesCacheAffinityCoefficient(t *testing.T) {
PlatformID: "deepseek-platform",
PlatformPriority: 100,
CacheAffinity: store.RuntimeCandidateCacheAffinity{
Key: "conversation:test",
RequestCount: 2,
CachedInputTokens: 7296,
EMAHitRatio: 0.91,
LastHitRatio: 0.92,
Confidence: 1,
Score: 1.74,
Boost: 20,
AdjustedPriority: 80,
Applied: true,
Key: "prompt_lcp_v2:test",
RequestCount: 2,
CachedInputTokens: 7296,
EMAHitRatio: 0.91,
LastHitRatio: 0.92,
Confidence: 1,
Score: 1.74,
Boost: 20,
AdjustedPriority: 80,
MatchedPrefixDepth: 4,
CandidateCount: 2,
OverrideReason: "different_effective_priority_tier",
Applied: true,
},
}, 1, false)
@@ -35,4 +38,10 @@ func TestAttemptMetricsIncludesCacheAffinityCoefficient(t *testing.T) {
if metrics["cacheAffinityHitRatio"] != 0.91 || metrics["cacheAffinityLastHitRatio"] != 0.92 {
t.Fatalf("expected hit ratios in attempt metrics, got %+v", metrics)
}
if metrics["cacheAffinityMatched"] != true ||
metrics["cacheAffinityMatchedPrefixDepth"] != 4 ||
metrics["cacheAffinityCandidateCount"] != 2 ||
metrics["cacheAffinityOverrideReason"] != "different_effective_priority_tier" {
t.Fatalf("expected cache affinity routing diagnostics in attempt metrics, got %+v", metrics)
}
}
+14 -1
View File
@@ -159,7 +159,12 @@ func (s *Service) hydrateProviderRequestAssetString(ctx context.Context, value s
if raw == "" || !imageInputFieldNeedsHydration(path) {
return value, nil
}
style, ok := requestAssetCapabilityHydrationForMedia("image", candidate, raw, "")
style, ok := requestAssetHydrateUnsupported, false
if openAIImageEditRequiresMultipartBytes(candidate) {
style, ok = requestAssetHydrateDataURL, true
} else {
style, ok = requestAssetCapabilityHydrationForMedia("image", candidate, raw, "")
}
if !ok {
return value, nil
}
@@ -317,6 +322,9 @@ func requestAssetHydrationForField(path []string, asset store.RequestAsset, cand
return requestAssetHydrateDataURL
}
if requestAssetMediaKindForHydration(path, asset) == "image" {
if openAIImageEditRequiresMultipartBytes(candidate) {
return requestAssetHydrateDataURL
}
if style, ok := requestAssetCapabilityHydrationForMedia("image", candidate, asset.URL, asset.StorageProvider); ok {
return style
}
@@ -332,6 +340,11 @@ func requestAssetHydrationForField(path []string, asset store.RequestAsset, cand
return requestAssetHydrateURL
}
func openAIImageEditRequiresMultipartBytes(candidate store.RuntimeModelCandidate) bool {
return normalizeProviderKey(candidate.Provider) == "openai" &&
strings.TrimSpace(candidate.ModelType) == "image_edit"
}
func requestAssetMediaKindForHydration(path []string, asset store.RequestAsset) string {
if mediaURLFieldNeedsHydration(path) {
return requestAssetMediaURLKind(path)
@@ -45,6 +45,56 @@ func TestHydrateProviderRequestAssetsConvertsStrictBase64Field(t *testing.T) {
}
}
func TestHydrateProviderRequestAssetsConvertsInputAudioDataToRawBase64(t *testing.T) {
storageDir := t.TempDir()
fileName := "gateway-request-asset-audio.mp3"
payload := []byte("audio bytes")
if err := os.WriteFile(filepath.Join(storageDir, fileName), payload, 0o644); err != nil {
t.Fatalf("write request asset: %v", err)
}
service := &Service{cfg: config.Config{LocalUploadedStorageDir: storageDir}}
body := map[string]any{
"messages": []any{
map[string]any{
"role": "user",
"content": []any{
map[string]any{
"type": "input_audio",
"input_audio": map[string]any{
"data": map[string]any{
"assetRef": map[string]any{
"sha256": "sha-input-audio",
"contentType": "audio/mpeg",
"url": "/static/uploaded/" + fileName,
"storageProvider": "local_static",
},
"url": "/static/uploaded/" + fileName,
},
"format": "mp3",
},
},
},
},
},
}
hydrated, err := service.hydrateProviderRequestAssets(context.Background(), body, store.RuntimeModelCandidate{})
if err != nil {
t.Fatalf("hydrate request assets: %v", err)
}
messages := hydrated["messages"].([]any)
message := messages[0].(map[string]any)
content := message["content"].([]any)
audioPart := content[0].(map[string]any)
inputAudio := audioPart["input_audio"].(map[string]any)
if got, want := stringFromAny(inputAudio["data"]), base64.StdEncoding.EncodeToString(payload); got != want {
t.Fatalf("unexpected hydrated input audio: got %q want %q", got, want)
}
if inputAudio["format"] != "mp3" {
t.Fatalf("input audio format was not preserved: %+v", inputAudio)
}
}
func TestHydrateProviderRequestAssetsConvertsBase64ArrayField(t *testing.T) {
storageDir := t.TempDir()
fileName := "gateway-request-asset-array.png"
@@ -205,6 +255,41 @@ func TestHydrateProviderRequestAssetsUsesImageCapabilityBase64ForTopLevelImageAs
}
}
func TestHydrateProviderRequestAssetsConvertsOpenAIEditImagesForMultipart(t *testing.T) {
payload := []byte{0x89, 'P', 'N', 'G', '\r', '\n', 0x1a, '\n', 1, 2, 3, 4}
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "image/png")
_, _ = w.Write(payload)
}))
defer server.Close()
service := &Service{}
body := map[string]any{
"images": []any{
server.URL + "/first.png",
server.URL + "/second.png",
},
}
hydrated, err := service.hydrateProviderRequestAssets(context.Background(), body, store.RuntimeModelCandidate{
Provider: "openai",
ModelType: "image_edit",
Capabilities: map[string]any{
"image_edit": map[string]any{
"support_url_input": true,
"support_base64_input": false,
},
},
})
if err != nil {
t.Fatalf("hydrate OpenAI edit images: %v", err)
}
images := hydrated["images"].([]any)
want := "data:image/png;base64," + base64.StdEncoding.EncodeToString(payload)
if len(images) != 2 || stringFromAny(images[0]) != want || stringFromAny(images[1]) != want {
t.Fatalf("OpenAI edit images should be hydrated for multipart: %+v", images)
}
}
func TestHydrateProviderRequestAssetsImageCapabilityOverridesProviderDataURLDefault(t *testing.T) {
service := &Service{}
body := map[string]any{
+202 -1
View File
@@ -42,6 +42,29 @@ type failoverDecision struct {
Info failureInfo
}
type failureDecision struct {
Route string
Effect string
Target string
CooldownSeconds int
DemoteSteps int
Reason string
Match policyRuleMatch
Info failureInfo
}
type resolveCandidateFailureInput struct {
RunnerPolicy store.RunnerPolicy
Candidate store.RuntimeModelCandidate
Err error
ClientAttempt int
MaxClientAttempts int
HasNextCandidate bool
FailoverExpired bool
Async bool
DownstreamStarted bool
}
type priorityDemoteDecision struct {
Demote bool
Reason string
@@ -57,6 +80,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 +103,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}
}
@@ -98,7 +127,7 @@ func failoverDecisionForCandidate(runnerPolicy store.RunnerPolicy, candidate sto
if errors.Is(err, store.ErrRateLimited) {
return failoverDecision{Retry: false, Action: "queue", Reason: "local_rate_limit_wait_queue", Match: policyRuleMatch{Source: "gateway_rate_limits", Policy: "rateLimitPolicy", Rule: "localCapacity", Value: "exceeded"}, Info: info}
}
if ruleDecision, ok := failoverActionRuleDecisionWithSources(runnerPolicy.FailoverPolicy, overridePolicy, info); ok {
if ruleDecision, ok := candidateFailoverActionRuleDecision(runnerPolicy.FailoverPolicy, overridePolicy, candidate, info); ok {
return ruleDecision
}
if match, ok := failoverDenyMatchWithSources(runnerPolicy.FailoverPolicy, overridePolicy, info); ok {
@@ -119,6 +148,122 @@ func failoverDecisionForCandidate(runnerPolicy store.RunnerPolicy, candidate sto
return failoverDecision{Retry: false, Action: "stop", Reason: "client_non_retryable", Match: policyRuleMatch{Source: "provider_client", Policy: "ClientError", Rule: "Retryable", Value: "false"}, Info: info}
}
func resolveCandidateFailure(input resolveCandidateFailureInput) failureDecision {
info := failureInfoFromError(input.Err)
if isResultPersistenceFailure(input.Err) {
return failureDecision{
Route: "stop",
Effect: "none",
Reason: "result_persistence_failed",
Match: policyRuleMatch{Source: "gateway_result_storage", Policy: "persistence", Rule: "providerFailover", Value: "disabled"},
Info: info,
}
}
if errors.Is(input.Err, store.ErrRateLimited) {
route := "stop"
if input.Async && store.RateLimitRetryable(input.Err) {
route = "requeue"
}
return failureDecision{
Route: route,
Effect: "none",
Reason: "local_rate_limit_wait_queue",
Match: policyRuleMatch{Source: "gateway_rate_limits", Policy: "rateLimitPolicy", Rule: "localCapacity", Value: "exceeded"},
Info: info,
}
}
failover := failoverDecisionForCandidate(input.RunnerPolicy, input.Candidate, input.Err)
effect := failoverEffect(failover.Action)
if input.DownstreamStarted {
return failureDecision{
Route: "stop",
Effect: effect,
Target: failover.Target,
CooldownSeconds: failover.CooldownSeconds,
DemoteSteps: failover.DemoteSteps,
Reason: "downstream_response_started",
Match: policyRuleMatch{Source: "gateway_stream", Policy: "streamSafety", Rule: "downstreamStarted", Value: "true"},
Info: info,
}
}
explicitStop := failover.Action == "stop" && (failover.Reason == "failover_action_rule" ||
failover.Reason == "hard_stop_policy" ||
failover.Reason == "failover_deny_policy" ||
failover.Reason == "failover_disabled" ||
failover.Reason == "runner_policy_disabled")
if explicitStop {
return failureDecision{
Route: "stop",
Effect: effect,
Target: failover.Target,
CooldownSeconds: failover.CooldownSeconds,
DemoteSteps: failover.DemoteSteps,
Reason: failover.Reason,
Match: failover.Match,
Info: failover.Info,
}
}
immediateCandidateInvalidation := failover.Action == "cooldown_and_next" || failover.Action == "disable_and_next"
retry := retryDecisionForCandidate(input.Candidate, input.Err)
if !immediateCandidateInvalidation && retry.Retry && input.ClientAttempt < input.MaxClientAttempts {
return failureDecision{
Route: "retry_same",
Effect: "none",
Reason: retry.Reason,
Match: retry.Match,
Info: retry.Info,
}
}
route := "next"
reason := failover.Reason
match := failover.Match
if input.FailoverExpired {
route = "stop"
reason = "failover_time_budget_exceeded"
match = policyRuleMatch{
Source: "gateway_runner_policies.failover_policy",
Policy: "failoverPolicy",
Rule: "maxDurationSeconds",
Value: "exceeded",
}
} else if !input.HasNextCandidate {
route = "stop"
reason = "no_next_platform"
match = policyRuleMatch{
Source: "runner_candidates",
Policy: "candidateSelection",
Rule: "candidateCount",
Value: "exhausted",
}
}
return failureDecision{
Route: route,
Effect: effect,
Target: failover.Target,
CooldownSeconds: failover.CooldownSeconds,
DemoteSteps: failover.DemoteSteps,
Reason: reason,
Match: match,
Info: failover.Info,
}
}
func failoverEffect(action string) string {
switch action {
case "cooldown_and_next":
return "cooldown"
case "demote_and_next":
return "demote"
case "disable_and_next":
return "disable"
default:
return "none"
}
}
func shouldDemoteCandidatePriority(runnerPolicy store.RunnerPolicy, err error) bool {
decision := priorityDemoteDecisionForCandidate(runnerPolicy, err)
return decision.Demote
@@ -126,6 +271,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 +293,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 {
@@ -324,6 +485,46 @@ func failoverActionRuleDecisionWithSources(base map[string]any, override map[str
return failoverActionRuleDecision(actionRulesFromPolicy(base), "gateway_runner_policies.failover_policy", info)
}
func candidateFailoverActionRuleDecision(base map[string]any, override map[string]any, candidate store.RuntimeModelCandidate, info failureInfo) (failoverDecision, bool) {
if rules := actionRulesFromPolicy(override); len(rules) > 0 {
if decision, ok := failoverActionRuleDecision(rules, "runtime_policy_override.failoverPolicy", info); ok {
return decision, true
}
}
autoDisablePolicy := effectiveRuntimePolicy(candidate.AutoDisablePolicy, candidate.RuntimePolicyOverride, "autoDisablePolicy")
if boolFromPolicy(autoDisablePolicy, "enabled", false) && intFromPolicy(autoDisablePolicy, "threshold") <= 1 {
if value, ok := matchingKeywordValue(autoDisablePolicy, "keywords", info.Target); ok {
return failoverDecision{
Retry: true,
Action: "disable_and_next",
Target: "platform",
Reason: "legacy_auto_disable_policy",
Match: policyRuleMatch{Source: "model_runtime_policy_sets.auto_disable_policy", Policy: "autoDisablePolicy", Rule: "keywords", Value: value},
Info: info,
}, true
}
}
degradePolicy := effectiveRuntimePolicy(candidate.DegradePolicy, candidate.RuntimePolicyOverride, "degradePolicy")
if boolFromPolicy(degradePolicy, "enabled", false) {
if value, ok := matchingKeywordValue(degradePolicy, "keywords", info.Target); ok {
cooldownSeconds := intFromPolicy(degradePolicy, "cooldownSeconds")
if cooldownSeconds <= 0 {
cooldownSeconds = 300
}
return failoverDecision{
Retry: true,
Action: "cooldown_and_next",
Target: "model",
Reason: "legacy_degrade_policy",
CooldownSeconds: cooldownSeconds,
Match: policyRuleMatch{Source: "model_runtime_policy_sets.degrade_policy", Policy: "degradePolicy", Rule: "keywords", Value: value},
Info: info,
}, true
}
}
return failoverActionRuleDecision(actionRulesFromPolicy(base), "gateway_runner_policies.failover_policy", info)
}
func failoverActionRuleDecision(rules []map[string]any, source string, info failureInfo) (failoverDecision, bool) {
for index, rule := range rules {
if !boolFromPolicy(rule, "enabled", true) {
@@ -334,3 +334,151 @@ func TestPriorityDemotePolicyIsSkippedWhenFailoverActionRulesExist(t *testing.T)
t.Fatalf("legacy priority demotion should be skipped when actionRules are active, got %+v", decision)
}
}
func TestResolveCandidateFailureCooldownSkipsSameClientRetry(t *testing.T) {
decision := resolveCandidateFailure(resolveCandidateFailureInput{
RunnerPolicy: store.RunnerPolicy{
Status: "active",
FailoverPolicy: map[string]any{
"enabled": true,
"actionRules": []any{
map[string]any{
"statusCodes": []any{429},
"action": "cooldown_and_next",
"target": "model",
"cooldownSeconds": 60,
},
},
},
},
Candidate: store.RuntimeModelCandidate{
ModelRetryPolicy: map[string]any{"enabled": true, "maxAttempts": 3},
},
Err: &clients.ClientError{Code: "resource_exhausted", StatusCode: 429, Retryable: true},
ClientAttempt: 1,
MaxClientAttempts: 3,
HasNextCandidate: true,
})
if decision.Route != "next" || decision.Effect != "cooldown" || decision.Target != "model" || decision.CooldownSeconds != 60 {
t.Fatalf("429 cooldown rule should invalidate the candidate immediately, got %+v", decision)
}
}
func TestResolveCandidateFailureDemoteRetriesBeforeNext(t *testing.T) {
input := resolveCandidateFailureInput{
RunnerPolicy: store.RunnerPolicy{
Status: "active",
FailoverPolicy: map[string]any{
"enabled": true,
"actionRules": []any{
map[string]any{
"statusCodes": []any{500},
"action": "demote_and_next",
"target": "platform",
"demoteSteps": 2,
},
},
},
},
Candidate: store.RuntimeModelCandidate{
ModelRetryPolicy: map[string]any{"enabled": true, "maxAttempts": 2},
},
Err: &clients.ClientError{Code: "server_error", StatusCode: 500, Retryable: true},
ClientAttempt: 1,
MaxClientAttempts: 2,
HasNextCandidate: true,
}
first := resolveCandidateFailure(input)
if first.Route != "retry_same" || first.Effect != "none" {
t.Fatalf("demotion must wait until same-client retries are exhausted, got %+v", first)
}
input.ClientAttempt = 2
second := resolveCandidateFailure(input)
if second.Route != "next" || second.Effect != "demote" || second.DemoteSteps != 2 {
t.Fatalf("exhausted retry should demote once and rotate, got %+v", second)
}
}
func TestResolveCandidateFailureStopsAfterDownstreamStarts(t *testing.T) {
decision := resolveCandidateFailure(resolveCandidateFailureInput{
RunnerPolicy: store.RunnerPolicy{
Status: "active",
FailoverPolicy: map[string]any{
"enabled": true,
"actionRules": []any{
map[string]any{
"categories": []any{"stream_error"},
"action": "next",
},
},
},
},
Candidate: store.RuntimeModelCandidate{ModelRetryPolicy: map[string]any{"enabled": true, "maxAttempts": 3}},
Err: &clients.ClientError{Code: "stream_read_error", Retryable: true},
ClientAttempt: 1,
MaxClientAttempts: 3,
HasNextCandidate: true,
DownstreamStarted: true,
})
if decision.Route != "stop" || decision.Reason != "downstream_response_started" {
t.Fatalf("stream output must prevent retries and failover, got %+v", decision)
}
}
func TestResolveCandidateFailureLegacyPolicyPrecedence(t *testing.T) {
candidate := store.RuntimeModelCandidate{
DegradePolicy: map[string]any{
"enabled": true,
"keywords": []any{"resource_exhausted"},
"cooldownSeconds": 45,
},
}
decision := resolveCandidateFailure(resolveCandidateFailureInput{
RunnerPolicy: store.RunnerPolicy{
Status: "active",
FailoverPolicy: map[string]any{
"enabled": true,
"actionRules": []any{
map[string]any{
"statusCodes": []any{429},
"action": "cooldown_and_next",
"target": "platform",
"cooldownSeconds": 90,
},
},
},
},
Candidate: candidate,
Err: &clients.ClientError{Code: "resource_exhausted", Message: "resource_exhausted", StatusCode: 429, Retryable: true},
ClientAttempt: 1,
MaxClientAttempts: 3,
HasNextCandidate: true,
})
if decision.Match.Policy != "degradePolicy" || decision.Target != "model" || decision.CooldownSeconds != 45 {
t.Fatalf("legacy compatibility rule should precede global action rules, got %+v", decision)
}
}
func TestResolveCandidateFailureUnmatchedRetryableUsesSameThenNext(t *testing.T) {
input := resolveCandidateFailureInput{
RunnerPolicy: store.RunnerPolicy{Status: "active", FailoverPolicy: map[string]any{"enabled": true, "actionRules": []any{}}},
Candidate: store.RuntimeModelCandidate{ModelRetryPolicy: map[string]any{"enabled": true, "maxAttempts": 2}},
Err: &clients.ClientError{Code: "unknown_transient", Retryable: true},
ClientAttempt: 1,
MaxClientAttempts: 2,
HasNextCandidate: true,
}
if decision := resolveCandidateFailure(input); decision.Route != "retry_same" {
t.Fatalf("unmatched retryable error should retry the same client first, got %+v", decision)
}
input.ClientAttempt = 2
if decision := resolveCandidateFailure(input); decision.Route != "next" || decision.Effect != "none" {
t.Fatalf("unmatched exhausted error should rotate without side effects, got %+v", decision)
}
}
+43 -137
View File
@@ -2,153 +2,59 @@ package runner
import (
"context"
"errors"
"strings"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/clients"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
func (s *Service) applyCandidateFailurePolicies(ctx context.Context, taskID string, candidate store.RuntimeModelCandidate, cause error, simulated bool, singleSourceProtected bool) {
code := clients.ErrorCode(cause)
message := ""
if cause != nil {
message = cause.Error()
func (s *Service) applyFailureDecision(ctx context.Context, taskID string, requestedModel string, candidate store.RuntimeModelCandidate, decision failureDecision, singleSourceProtection bool, simulated bool) (store.CandidateFailureEffectResult, error) {
result, err := s.store.ApplyCandidateFailureEffect(ctx, store.CandidateFailureEffectInput{
Effect: decision.Effect,
Target: decision.Target,
PlatformID: candidate.PlatformID,
PlatformModelID: candidate.PlatformModelID,
RequestedModel: requestedModel,
ModelType: candidate.ModelType,
CooldownSeconds: decision.CooldownSeconds,
DemoteSteps: decision.DemoteSteps,
SingleSourceProtection: singleSourceProtection,
})
if err != nil {
return result, err
}
autoDisablePolicy := effectiveRuntimePolicy(candidate.AutoDisablePolicy, candidate.RuntimePolicyOverride, "autoDisablePolicy")
if failurePolicyMatches(autoDisablePolicy, code, message) && intFromPolicy(autoDisablePolicy, "threshold") <= 1 {
if singleSourceProtected {
s.emitSingleSourceProtected(ctx, taskID, candidate, "auto_disable", code, simulated)
} else if err := s.store.DisableCandidatePlatform(ctx, candidate.PlatformID); err == nil {
_ = s.emit(ctx, taskID, "task.policy.auto_disabled", "running", "auto_disable", 0.48, "candidate platform disabled by failure policy", map[string]any{
"platformId": candidate.PlatformID,
"platformModelId": candidate.PlatformModelID,
"code": code,
}, simulated)
}
}
degradePolicy := effectiveRuntimePolicy(candidate.DegradePolicy, candidate.RuntimePolicyOverride, "degradePolicy")
if failurePolicyMatches(degradePolicy, code, message) {
cooldownSeconds := intFromPolicy(degradePolicy, "cooldownSeconds")
if singleSourceProtected {
s.emitSingleSourceProtected(ctx, taskID, candidate, "degrade_cooldown", code, simulated)
} else if err := s.store.CooldownCandidatePlatformModel(ctx, candidate.PlatformModelID, cooldownSeconds); err == nil {
_ = s.emit(ctx, taskID, "task.policy.degraded", "running", "degrade", 0.5, "candidate model cooled down by failure policy", map[string]any{
"platformId": candidate.PlatformID,
"platformModelId": candidate.PlatformModelID,
"cooldownSeconds": cooldownSeconds,
"code": code,
}, simulated)
}
}
}
func (s *Service) applyFailoverAction(ctx context.Context, taskID string, candidate store.RuntimeModelCandidate, requestedModel string, decision failoverDecision, simulated bool, singleSourceProtected bool) {
switch decision.Action {
case "disable_and_next":
if singleSourceProtected {
s.emitSingleSourceProtected(ctx, taskID, candidate, decision.Action, decision.Info.Code, simulated)
return
}
target := decision.Target
if target == "" {
target = "platform"
}
var err error
if target == "model" {
err = s.store.DisableCandidatePlatformModel(ctx, candidate.PlatformModelID)
} else {
err = s.store.DisableCandidatePlatform(ctx, candidate.PlatformID)
}
if err == nil {
_ = s.emit(ctx, taskID, "task.policy.failover_disabled", "running", "failover", 0.51, "candidate platform disabled by runner failover policy", addPolicyTracePayload(map[string]any{
"platformId": candidate.PlatformID,
"platformModelId": candidate.PlatformModelID,
"action": decision.Action,
"target": target,
"reason": decision.Reason,
}, decision.Match, decision.Info), simulated)
}
case "cooldown_and_next":
if singleSourceProtected {
s.emitSingleSourceProtected(ctx, taskID, candidate, decision.Action, decision.Info.Code, simulated)
return
}
target := decision.Target
if target == "" {
target = "model"
}
var err error
if target == "platform" {
err = s.store.CooldownCandidatePlatform(ctx, candidate.PlatformID, decision.CooldownSeconds)
} else {
err = s.store.CooldownCandidatePlatformModel(ctx, candidate.PlatformModelID, decision.CooldownSeconds)
}
if err == nil {
_ = s.emit(ctx, taskID, "task.policy.failover_cooled_down", "running", "failover", 0.51, "candidate model cooled down by runner failover policy", addPolicyTracePayload(map[string]any{
"platformId": candidate.PlatformID,
"platformModelId": candidate.PlatformModelID,
"cooldownSeconds": decision.CooldownSeconds,
"action": decision.Action,
"target": target,
"reason": decision.Reason,
}, decision.Match, decision.Info), simulated)
}
case "demote_and_next":
demoteSteps := decision.DemoteSteps
if demoteSteps <= 0 {
demoteSteps = 1
}
if dynamicPriority, err := s.store.DemoteCandidatePlatformPriorityBySteps(ctx, candidate.PlatformID, candidate.PlatformModelID, requestedModel, candidate.ModelType, demoteSteps); err == nil {
_ = s.emit(ctx, taskID, "task.policy.priority_demoted", "running", "priority_demote", 0.52, "candidate platform priority demoted by runner failover policy", addPolicyTracePayload(map[string]any{
"platformId": candidate.PlatformID,
"platformModelId": candidate.PlatformModelID,
"dynamicPriority": dynamicPriority,
"demoteSteps": demoteSteps,
"action": decision.Action,
"target": "platform",
"reason": decision.Reason,
"errorMessage": decision.Info.Message,
}, decision.Match, decision.Info), simulated)
}
}
}
func (s *Service) emitSingleSourceProtected(ctx context.Context, taskID string, candidate store.RuntimeModelCandidate, action string, code string, simulated bool) {
_ = s.emit(ctx, taskID, "task.policy.single_source_protected", "running", "policy_guard", 0.5, "single source protected from disable or cooldown", map[string]any{
"platformId": candidate.PlatformID,
"platformModelId": candidate.PlatformModelID,
"action": action,
"code": code,
}, simulated)
}
func (s *Service) applyPriorityDemotePolicy(ctx context.Context, taskID string, attemptNo int, runnerPolicy store.RunnerPolicy, candidate store.RuntimeModelCandidate, requestedModel string, cause error, simulated bool) {
if errors.Is(cause, store.ErrRateLimited) {
return
}
decision := priorityDemoteDecisionForCandidate(runnerPolicy, cause)
if !decision.Demote {
return
}
demoteSteps := decision.DemoteSteps
if demoteSteps <= 0 {
demoteSteps = 99
}
if dynamicPriority, err := s.store.DemoteCandidatePlatformPriorityBySteps(ctx, candidate.PlatformID, candidate.PlatformModelID, requestedModel, candidate.ModelType, demoteSteps); err == nil {
s.recordAttemptTrace(ctx, taskID, attemptNo, priorityDemoteTraceEntry(decision, candidate.PlatformID, candidate.PlatformModelID, dynamicPriority))
_ = s.emit(ctx, taskID, "task.policy.priority_demoted", "running", "priority_demote", 0.52, "candidate platform priority demoted by runner policy", addPolicyTracePayload(map[string]any{
if result.GuardReason == "single_source" {
_ = s.emit(ctx, taskID, "task.policy.single_source_protected", "running", "policy_guard", 0.5, "single source protected from disable or cooldown", addPolicyTracePayload(map[string]any{
"platformId": candidate.PlatformID,
"platformModelId": candidate.PlatformModelID,
"dynamicPriority": dynamicPriority,
"demoteSteps": demoteSteps,
"code": clients.ErrorCode(cause),
"reason": decision.Reason,
"errorMessage": messageFromError(cause),
"effect": decision.Effect,
"target": decision.Target,
"guardReason": result.GuardReason,
}, decision.Match, decision.Info), simulated)
return result, nil
}
if !result.Applied {
return result, nil
}
payload := addPolicyTracePayload(map[string]any{
"platformId": candidate.PlatformID,
"platformModelId": candidate.PlatformModelID,
"effect": decision.Effect,
"target": decision.Target,
"reason": decision.Reason,
}, decision.Match, decision.Info)
switch decision.Effect {
case "disable":
_ = s.emit(ctx, taskID, "task.policy.failover_disabled", "running", "failover", 0.51, "candidate disabled by runner failover policy", payload, simulated)
case "cooldown":
payload["cooldownSeconds"] = decision.CooldownSeconds
_ = s.emit(ctx, taskID, "task.policy.failover_cooled_down", "running", "failover", 0.51, "candidate cooled down by runner failover policy", payload, simulated)
case "demote":
payload["demoteSteps"] = decision.DemoteSteps
payload["dynamicPriority"] = result.DynamicPriority
_ = s.emit(ctx, taskID, "task.policy.priority_demoted", "running", "priority_demote", 0.52, "candidate platform priority demoted by runner failover policy", payload, simulated)
}
return result, nil
}
func messageFromError(err error) string {
+433 -138
View File
@@ -2,14 +2,16 @@ package runner
import (
"context"
"encoding/json"
"errors"
"fmt"
"log/slog"
"net/http"
"os"
"regexp"
"strconv"
"strings"
"sync"
"sync/atomic"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
@@ -23,14 +25,20 @@ import (
)
type Service struct {
cfg config.Config
store *store.Store
logger *slog.Logger
clients map[string]clients.Client
scriptExecutor *scriptengine.Executor
httpClients *httpClientCache
riverClient *river.Client[pgx.Tx]
billingMetrics billingMetricsObserver
cfg config.Config
store *store.Store
logger *slog.Logger
clients map[string]clients.Client
scriptExecutor *scriptengine.Executor
httpClients *httpClientCache
riverMu sync.RWMutex
riverControlClient *river.Client[pgx.Tx]
riverExecutionClient asyncExecutionClient
riverDrainingClients map[asyncExecutionClient]struct{}
riverWorkerCapacity int
asyncCapacityLoader func(context.Context, int) (store.AsyncWorkerCapacitySnapshot, error)
asyncClientFactory func(int) (asyncExecutionClient, error)
billingMetrics billingMetricsObserver
}
type billingMetricsObserver interface {
@@ -75,6 +83,30 @@ func (e *TaskQueuedError) Is(target error) bool {
}
func New(cfg config.Config, db *store.Store, logger *slog.Logger, observers ...billingMetricsObserver) *Service {
if cfg.AsyncWorkerHardLimit == 0 {
cfg.AsyncWorkerHardLimit = 2048
}
if cfg.AsyncWorkerRefreshIntervalSeconds == 0 {
cfg.AsyncWorkerRefreshIntervalSeconds = 5
}
if cfg.TaskProgressCallbackTimeoutMS == 0 {
cfg.TaskProgressCallbackTimeoutMS = 5000
}
if cfg.TaskProgressCallbackMaxAttempts == 0 {
cfg.TaskProgressCallbackMaxAttempts = 10
}
if cfg.TaskRetentionDays == 0 {
cfg.TaskRetentionDays = 30
}
if cfg.TaskAnalysisRetentionDays == 0 {
cfg.TaskAnalysisRetentionDays = 7
}
if cfg.TaskCleanupIntervalSeconds == 0 {
cfg.TaskCleanupIntervalSeconds = 300
}
if cfg.TaskCleanupBatchSize == 0 {
cfg.TaskCleanupBatchSize = 1000
}
httpClients := newHTTPClientCache()
scriptExecutor := &scriptengine.Executor{Logger: logger}
service := &Service{
@@ -117,6 +149,15 @@ func (s *Service) observeBillingEvent(event string) {
}
}
func (s *Service) observeTaskEventSkip(reason string) {
observer, ok := s.billingMetrics.(interface {
ObserveTaskEventSkip(string)
})
if ok {
observer.ObserveTaskEventSkip(reason)
}
}
func (s *Service) Execute(ctx context.Context, task store.GatewayTask, user *auth.User) (Result, error) {
return s.execute(ctx, task, user, nil)
}
@@ -222,10 +263,27 @@ func (s *Service) executeWithToken(ctx context.Context, task store.GatewayTask,
CacheAffinityKeys: cacheAffinityKeys.Lookup,
CacheAffinityPolicy: runnerPolicy.CacheAffinityPolicy,
})
if err == nil {
candidates, err = filterCandidatesByRequestedPlatform(candidates, body)
}
if err != nil {
if task.Kind == "responses" && responseExecution.PreviousChain != nil {
err = responseChainUnavailableError()
}
if task.AsyncMode && store.ModelCandidateRetryAfter(err) > 0 {
queued, delay, queueErr := s.requeueModelCoolingTask(ctx, task, err)
if queueErr != nil {
return Result{}, queueErr
}
return Result{Task: queued, Output: queued.Result}, &TaskQueuedError{Delay: delay}
}
if store.ModelCandidateRetryAfter(err) > 0 {
failed, finishErr := s.failTask(ctx, task.ID, task.ExecutionToken, store.ModelCandidateErrorCode(err), err.Error(), task.RunMode == "simulation", err)
if finishErr != nil {
return Result{}, finishErr
}
return Result{Task: failed, Output: failed.Result}, err
}
s.recordFailedAttempt(ctx, failedAttemptRecord{
Task: task,
Body: body,
@@ -496,18 +554,38 @@ func (s *Service) executeWithToken(ctx context.Context, task store.GatewayTask,
maxPlatforms := maxPlatformsForCandidates(candidates, runnerPolicy)
maxFailoverDuration := maxFailoverDurationForCandidates(candidates, runnerPolicy)
singleSourceProtected := singleSourceProtectionActive(candidates, maxPlatforms, runnerPolicy)
singleSourceProtection := runnerSingleSourceProtectionEnabled(runnerPolicy)
var downstreamStarted atomic.Bool
candidateDelta := onDelta
if onDelta != nil {
candidateDelta = func(event clients.StreamDeltaEvent) error {
downstreamStarted.Store(true)
return onDelta(event)
}
}
var lastErr error
var lastCandidate store.RuntimeModelCandidate
var lastPreprocessing *parameterPreprocessingLog
platformsVisited := 0
candidatesLoop:
for index, candidate := range candidates {
if index >= maxPlatforms {
if platformsVisited >= maxPlatforms {
break
}
available, availabilityErr := s.store.RuntimeCandidateAvailable(ctx, candidate.PlatformID, candidate.PlatformModelID)
if availabilityErr != nil {
return Result{}, availabilityErr
}
if !available {
continue
}
platformsVisited++
lastCandidate = candidate
clientAttempts := clientAttemptsForCandidate(candidate)
var candidateErr error
var candidateDecision failureDecision
var candidateDecisionAttempt int
var candidateClientAttempt int
for clientAttempt := 1; clientAttempt <= clientAttempts; clientAttempt++ {
nextAttemptNo := attemptNo + 1
preprocessing, ok := preprocessingByCandidate[pricingCandidateKey(candidate)]
@@ -536,7 +614,7 @@ candidatesLoop:
}
candidateBody := preprocessing.Body
candidatePricing := pricingByCandidate[pricingCandidateKey(candidate)]
response, err := s.runCandidate(ctx, task, user, candidateBody, preprocessing.Log, candidate, candidatePricing, nextAttemptNo, onDelta, responseExecution, singleSourceProtected, runnerPolicy.CacheAffinityPolicy, cacheAffinityKeys.Record)
response, err := s.runCandidate(ctx, task, user, candidateBody, preprocessing.Log, candidate, candidatePricing, nextAttemptNo, candidateDelta, responseExecution, runnerPolicy.CacheAffinityPolicy, cacheAffinityKeys.Record)
if err != nil && isVolcesRemoteTaskCancellation(candidate, err) {
cancelled, changed, cancelErr := s.store.CancelSubmittedTask(ctx, task.ID, task.ExecutionToken, "任务已由火山引擎取消")
if cancelErr != nil {
@@ -553,6 +631,7 @@ candidatesLoop:
}
if err == nil {
attemptNo = nextAttemptNo
response.Result = canonicalTaskResult(response.Result)
var billings []any
finalAmount := fixedAmount(0)
finalAmountText := ""
@@ -586,7 +665,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 {
@@ -617,6 +700,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" {
@@ -644,7 +759,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) {
@@ -664,7 +783,18 @@ candidatesLoop:
if isLocalRateLimitError(err) {
lastErr = err
candidateErr = err
if task.AsyncMode && store.RateLimitRetryable(err) {
candidateDecision = resolveCandidateFailure(resolveCandidateFailureInput{
RunnerPolicy: runnerPolicy,
Candidate: candidate,
Err: err,
ClientAttempt: clientAttempt,
MaxClientAttempts: clientAttempts,
HasNextCandidate: platformsVisited < maxPlatforms && index+1 < len(candidates),
FailoverExpired: failoverTimeBudgetExceeded(executeStartedAt, maxFailoverDuration),
Async: task.AsyncMode,
DownstreamStarted: downstreamStarted.Load(),
})
if candidateDecision.Route == "requeue" {
queued, delay, queueErr := s.requeueRateLimitedTask(ctx, task, err, candidate)
if queueErr != nil {
return Result{}, queueErr
@@ -684,104 +814,88 @@ candidatesLoop:
ExtraMetrics: []map[string]any{parameterPreprocessingMetrics(preprocessing.Log)},
ModelType: candidate.ModelType,
})
candidateDecisionAttempt = attemptNo
candidateClientAttempt = clientAttempt
s.recordAttemptTrace(ctx, task.ID, attemptNo, failureDecisionTraceEntry(candidateDecision, candidate, clientAttempt, clientAttempts, false, ""))
break candidatesLoop
}
attemptNo = nextAttemptNo
lastErr = err
candidateErr = err
retryDecision := retryDecisionForCandidate(candidate, err)
retryAction := "retry"
if !retryDecision.Retry {
retryAction = "stop"
}
if clientAttempt >= clientAttempts {
retryDecision.Retry = false
retryDecision.Reason = "same_client_max_attempts"
retryDecision.Match = policyRuleMatch{
Source: "model_runtime_policy_sets.retry_policy",
Policy: "retryPolicy",
Rule: "maxAttempts",
Value: strconv.Itoa(clientAttempts),
}
retryDecision.Info = failureInfoFromError(err)
retryAction = "stop"
}
if failoverTimeBudgetExceeded(executeStartedAt, maxFailoverDuration) {
retryDecision.Retry = false
retryDecision.Reason = "failover_time_budget_exceeded"
retryDecision.Match = policyRuleMatch{
Source: "gateway_runner_policies.failover_policy",
Policy: "failoverPolicy",
Rule: "maxDurationSeconds",
Value: strconv.Itoa(int(maxFailoverDuration.Seconds())),
}
retryDecision.Info = failureInfoFromError(err)
retryAction = "stop"
}
s.recordAttemptTrace(ctx, task.ID, attemptNo, retryTraceEntry(candidate, retryDecision, retryAction, clientAttempt, clientAttempts))
if !retryDecision.Retry {
candidateDecision = resolveCandidateFailure(resolveCandidateFailureInput{
RunnerPolicy: runnerPolicy,
Candidate: candidate,
Err: err,
ClientAttempt: clientAttempt,
MaxClientAttempts: clientAttempts,
HasNextCandidate: platformsVisited < maxPlatforms && index+1 < len(candidates),
FailoverExpired: failoverTimeBudgetExceeded(executeStartedAt, maxFailoverDuration),
Async: task.AsyncMode,
DownstreamStarted: downstreamStarted.Load(),
})
candidateDecisionAttempt = attemptNo
candidateClientAttempt = clientAttempt
if candidateDecision.Route != "retry_same" {
break
}
s.recordAttemptTrace(ctx, task.ID, attemptNo, failureDecisionTraceEntry(candidateDecision, candidate, clientAttempt, clientAttempts, false, ""))
if err := s.emit(ctx, task.ID, "task.retrying", "running", "retry", 0.45, "retrying same client", addPolicyTracePayload(map[string]any{
"attempt": attemptNo,
"clientAttempt": clientAttempt,
"clientId": candidate.ClientID,
"error": err.Error(),
"reason": retryDecision.Reason,
"reason": candidateDecision.Reason,
"scope": "same_client",
}, retryDecision.Match, retryDecision.Info), isSimulation(task, candidate)); err != nil {
}, candidateDecision.Match, candidateDecision.Info), isSimulation(task, candidate)); err != nil {
return Result{}, err
}
}
if candidateErr == nil || index+1 >= len(candidates) || index+1 >= maxPlatforms {
if candidateErr != nil {
s.applyPriorityDemotePolicy(ctx, task.ID, attemptNo, runnerPolicy, candidate, task.Model, candidateErr, isSimulation(task, candidate))
decision := failoverDecisionForCandidate(runnerPolicy, candidate, candidateErr)
if decision.Retry {
decision.Retry = false
decision.Action = "stop"
decision.Reason = "no_next_platform"
decision.Match = policyRuleMatch{Source: "runner_candidates", Policy: "candidateSelection", Rule: "candidateCount", Value: strconv.Itoa(len(candidates))}
if index+1 >= maxPlatforms {
decision.Reason = "max_platforms_reached"
decision.Match = policyRuleMatch{Source: "gateway_runner_policies.failover_policy", Policy: "failoverPolicy", Rule: "maxPlatforms", Value: strconv.Itoa(maxPlatforms)}
}
if candidateErr == nil {
break
}
effectResult, effectErr := s.applyFailureDecision(
context.WithoutCancel(ctx),
task.ID,
task.Model,
candidate,
candidateDecision,
singleSourceProtection,
isSimulation(task, candidate),
)
if effectErr != nil {
s.logger.Warn("apply candidate failure effect failed", "taskID", task.ID, "effect", candidateDecision.Effect, "error", effectErr)
effectResult.GuardReason = "effect_error"
}
s.recordAttemptTrace(ctx, task.ID, candidateDecisionAttempt, failureDecisionTraceEntry(
candidateDecision,
candidate,
candidateClientAttempt,
clientAttempts,
effectResult.Applied,
effectResult.GuardReason,
))
if candidateDecision.Route != "next" {
if candidateDecision.Reason == "failover_time_budget_exceeded" {
elapsedSeconds := int(time.Since(executeStartedAt).Seconds())
if err := s.emit(ctx, task.ID, "task.failover.stopped", "running", "retry", 0.55, "failover time budget exceeded", map[string]any{
"elapsedSeconds": elapsedSeconds,
"maxDurationSeconds": int(maxFailoverDuration.Seconds()),
"scope": "next_platform",
"statusCode": clients.ErrorResponseMetadata(candidateErr).StatusCode,
}, isSimulation(task, candidate)); err != nil {
return Result{}, err
}
s.recordAttemptTrace(ctx, task.ID, attemptNo, failoverTraceEntry(decision, candidate))
}
break
}
s.applyPriorityDemotePolicy(ctx, task.ID, attemptNo, runnerPolicy, candidate, task.Model, candidateErr, isSimulation(task, candidate))
if failoverTimeBudgetExceeded(executeStartedAt, maxFailoverDuration) {
elapsedSeconds := int(time.Since(executeStartedAt).Seconds())
maxDurationSeconds := int(maxFailoverDuration.Seconds())
s.recordAttemptTrace(ctx, task.ID, attemptNo, failoverTimeBudgetTraceEntry(elapsedSeconds, maxDurationSeconds, failureInfoFromError(candidateErr), candidate))
if err := s.emit(ctx, task.ID, "task.failover.stopped", "running", "retry", 0.55, "failover time budget exceeded", map[string]any{
"elapsedSeconds": elapsedSeconds,
"maxDurationSeconds": maxDurationSeconds,
"scope": "next_platform",
"statusCode": clients.ErrorResponseMetadata(candidateErr).StatusCode,
}, isSimulation(task, candidate)); err != nil {
return Result{}, err
}
break
}
decision := failoverDecisionForCandidate(runnerPolicy, candidate, candidateErr)
if !decision.Retry && hasLoadAvoidanceFallback(candidates, index, maxPlatforms) {
decision = loadAvoidanceFallbackDecision(candidateErr)
}
s.recordAttemptTrace(ctx, task.ID, attemptNo, failoverTraceEntry(decision, candidate))
if !decision.Retry {
break
}
s.applyFailoverAction(ctx, task.ID, candidate, task.Model, decision, isSimulation(task, candidate), singleSourceProtected)
if err := s.emit(ctx, task.ID, "task.retrying", "running", "retry", 0.55, "retrying next client", addPolicyTracePayload(map[string]any{
"attempt": attemptNo,
"action": decision.Action,
"route": candidateDecision.Route,
"effect": candidateDecision.Effect,
"error": candidateErr.Error(),
"reason": decision.Reason,
"reason": candidateDecision.Reason,
"scope": "next_platform",
}, decision.Match, decision.Info), isSimulation(task, candidate)); err != nil {
}, candidateDecision.Match, candidateDecision.Info), isSimulation(task, candidate)); err != nil {
return Result{}, err
}
}
@@ -867,7 +981,7 @@ func billingItemsFixedTotal(items []any) fixedAmount {
return total
}
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) {
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, cacheAffinityPolicy map[string]any, cacheAffinityRecordKeys []string) (clients.Response, error) {
var err error
candidate, err = candidateWithEnvironmentCredentials(candidate)
if err != nil {
@@ -888,7 +1002,7 @@ func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user
_ = s.store.ReleaseRateLimitReservations(context.WithoutCancel(ctx), limitResult.Reservations, "attempt_failed")
}
}()
defer s.store.ReleaseConcurrencyLeases(context.WithoutCancel(ctx), limitResult.LeaseIDs)
defer s.store.ReleaseConcurrencyLeases(context.WithoutCancel(ctx), limitResult.Leases)
attemptID, err := s.store.CreateTaskAttempt(ctx, store.CreateTaskAttemptInput{
TaskID: task.ID,
@@ -986,6 +1100,20 @@ func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user
})
return clients.Response{}, err
}
if strings.TrimSpace(task.RemoteTaskID) == "" {
providerBody, err = s.normalizeVideoInputImages(ctx, providerBody, candidate)
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, "trace": []any{failureTraceEntry(err, false)}}),
ErrorCode: clients.ErrorCode(err),
ErrorMessage: err.Error(),
})
return clients.Response{}, err
}
}
callStartedAt := time.Now()
publicResponseID := ""
publicPreviousResponseID := ""
@@ -1014,11 +1142,13 @@ func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user
return nil
}
var submissionWire *clients.WireResponse
response, err := client.Run(ctx, clients.Request{
runCtx, stopLeaseRenewal := s.startConcurrencyLeaseRenewal(ctx, task.ID, limitResult.Leases)
response, err := client.Run(runCtx, clients.Request{
Kind: task.Kind,
ModelType: candidate.ModelType,
Model: task.Model,
Body: providerBody,
OriginalBody: preprocessing.Input,
Candidate: candidate,
HTTPClient: requestHTTPClient,
RemoteTaskID: task.RemoteTaskID,
@@ -1027,12 +1157,13 @@ func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user
if strings.TrimSpace(remoteTaskID) == "" {
return nil
}
if err := s.store.SetTaskRemoteTask(context.WithoutCancel(ctx), task.ID, task.ExecutionToken, attemptID, remoteTaskID, payload); err != nil {
checkpoint := minimalRemoteTaskCheckpoint(candidate.Provider, candidate.SpecType, payload)
if err := s.store.SetTaskRemoteTask(context.WithoutCancel(ctx), task.ID, task.ExecutionToken, attemptID, remoteTaskID, checkpoint); err != nil {
return err
}
task.RemoteTaskID = remoteTaskID
task.RemoteTaskPayload = payload
if err := s.persistCompatibilitySubmission(context.WithoutCancel(ctx), task, candidate, remoteTaskID, payload, submissionWire); err != nil {
task.RemoteTaskPayload = checkpoint
if err := s.persistCompatibilitySubmission(context.WithoutCancel(ctx), task, candidate, remoteTaskID, checkpoint, submissionWire); err != nil {
return err
}
return setSubmissionStatus("response_received")
@@ -1041,11 +1172,12 @@ func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user
if strings.TrimSpace(remoteTaskID) == "" {
return nil
}
if err := s.store.SetTaskRemoteTask(context.WithoutCancel(ctx), task.ID, task.ExecutionToken, attemptID, remoteTaskID, payload); err != nil {
checkpoint := minimalRemoteTaskCheckpoint(candidate.Provider, candidate.SpecType, payload)
if err := s.store.SetTaskRemoteTask(context.WithoutCancel(ctx), task.ID, task.ExecutionToken, attemptID, remoteTaskID, checkpoint); err != nil {
return err
}
task.RemoteTaskID = remoteTaskID
task.RemoteTaskPayload = payload
task.RemoteTaskPayload = checkpoint
return setSubmissionStatus("response_received")
},
OnUpstreamSubmissionStarted: func() error {
@@ -1066,6 +1198,13 @@ func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user
UpstreamPreviousResponseID: responseExecution.UpstreamPreviousResponseID,
PreviousResponseTurns: responseExecution.PreviousTurns,
})
if leaseErr := stopLeaseRenewal(); leaseErr != nil {
err = &clients.ClientError{
Code: "concurrency_lease_lost",
Message: leaseErr.Error(),
Retryable: true,
}
}
callFinishedAt := time.Now()
if err == nil {
if markErr := setSubmissionStatus("response_received"); markErr != nil {
@@ -1121,7 +1260,6 @@ func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user
if !simulated && submissionStatus == "submitting" {
return clients.Response{}, &upstreamSubmissionUnknownError{AttemptID: attemptID, Cause: err}
}
s.applyCandidateFailurePolicies(ctx, task.ID, candidate, err, simulated, singleSourceProtected)
return clients.Response{}, err
}
uploadedResult, err := s.uploadGeneratedAssets(ctx, task.ID, task.Kind, response.Result)
@@ -1148,6 +1286,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
@@ -1220,6 +1364,59 @@ func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user
return response, nil
}
func minimalRemoteTaskCheckpoint(provider string, specType string, payload map[string]any) map[string]any {
const maxBytes = 8192
provider = strings.ToLower(strings.TrimSpace(provider))
specType = strings.ToLower(strings.TrimSpace(specType))
allowedByProvider := map[string][]string{
"keling": {"endpoint", "pollEndpoint", "mode", "taskType", "cleanupElementIds"},
"kling": {"endpoint", "pollEndpoint", "mode", "taskType", "cleanupElementIds"},
"topaz": {"phase", "targetResolution"},
"vectorizer": {"imageToken", "receipt"},
}
allowedKeys := allowedByProvider[provider]
if len(allowedKeys) == 0 {
allowedKeys = allowedByProvider[specType]
}
checkpoint := map[string]any{}
for _, key := range allowedKeys {
value, ok := payload[key]
if !ok || value == nil {
continue
}
if text, ok := value.(string); ok {
value = truncateText(text, 4096)
}
next := cloneMap(checkpoint)
next[key] = value
encoded, err := json.Marshal(next)
if err != nil || len(encoded) > maxBytes {
continue
}
checkpoint = next
}
return checkpoint
}
func canonicalTaskResult(result map[string]any) map[string]any {
canonical := cloneMap(result)
for _, key := range []string{
"raw",
"raw_data",
"rawData",
"provider_response",
"providerResponse",
"submit",
"file_retrieve",
"fileRetrieve",
"upstream_task_id",
"remote_task_id",
} {
delete(canonical, key)
}
return canonical
}
func (s *Service) persistCompatibilitySubmission(ctx context.Context, task store.GatewayTask, candidate store.RuntimeModelCandidate, remoteTaskID string, payload map[string]any, wire *clients.WireResponse) error {
targetProtocol := strings.TrimSpace(stringFromMap(task.Request, "_gateway_target_protocol"))
if targetProtocol == "" {
@@ -1229,32 +1426,9 @@ func (s *Service) persistCompatibilitySubmission(ctx context.Context, task store
if wire != nil && strings.TrimSpace(wire.Protocol) != "" {
sourceProtocol = strings.TrimSpace(wire.Protocol)
}
publicID := task.ID
if sourceProtocol == targetProtocol && strings.TrimSpace(remoteTaskID) != "" {
publicID = strings.TrimSpace(remoteTaskID)
}
submitBody := payload
if nested, ok := payload["submit"].(map[string]any); ok && len(nested) > 0 {
submitBody = nested
}
httpStatus := http.StatusOK
headers := map[string]any{}
if wire != nil {
httpStatus = wire.StatusCode
if len(wire.Body) > 0 {
submitBody = wire.Body
}
for name, values := range wire.Headers {
headers[name] = append([]string(nil), values...)
}
}
return s.store.SetTaskCompatibilitySubmission(ctx, task.ID, store.CompatibilitySubmission{
TargetProtocol: targetProtocol,
PublicID: publicID,
SourceProtocol: sourceProtocol,
HTTPStatus: httpStatus,
Headers: headers,
Body: submitBody,
})
}
@@ -1289,9 +1463,10 @@ func compatibilitySourceProtocol(kind string, candidate store.RuntimeModelCandid
}
func (s *Service) recordTaskParameterPreprocessing(ctx context.Context, taskID string, attemptID string, attemptNo int, candidate store.RuntimeModelCandidate, log parameterPreprocessingLog) error {
if skipTaskParameterPreprocessingLog(log.ModelType) && !log.Changed {
if !log.Changed {
return nil
}
changes := minimalParameterPreprocessingChanges(log.Changes)
_, err := s.store.CreateTaskParamPreprocessingLog(ctx, store.CreateTaskParamPreprocessingLogInput{
TaskID: taskID,
AttemptID: attemptID,
@@ -1300,16 +1475,31 @@ func (s *Service) recordTaskParameterPreprocessing(ctx context.Context, taskID s
PlatformID: candidate.PlatformID,
PlatformModelID: candidate.PlatformModelID,
ClientID: candidate.ClientID,
Changed: log.Changed,
Changed: true,
ChangeCount: len(log.Changes),
ActualInput: log.Input,
ConvertedOutput: log.Output,
Changes: log.Changes,
ModelSnapshot: log.Model,
Changes: changes,
})
return err
}
func minimalParameterPreprocessingChanges(changes []parameterPreprocessChange) []any {
const maxBytes = 1024
out := make([]any, 0, len(changes))
for _, change := range changes {
item := map[string]any{
"path": truncateText(change.Path, 256),
"action": truncateText(change.Action, 64),
}
next := append(append([]any(nil), out...), item)
encoded, err := json.Marshal(next)
if err != nil || len(encoded) > maxBytes {
break
}
out = next
}
return out
}
func skipTaskParameterPreprocessingLog(modelType string) bool {
switch strings.TrimSpace(modelType) {
case "text_generate", "text_embedding", "text_rerank", "chat", "responses", "text":
@@ -1319,6 +1509,15 @@ func skipTaskParameterPreprocessingLog(modelType string) bool {
}
}
func truncateText(value string, maxRunes int) string {
value = strings.TrimSpace(value)
runes := []rune(value)
if len(runes) <= maxRunes {
return value
}
return string(runes[:maxRunes])
}
func (s *Service) clientFor(candidate store.RuntimeModelCandidate, simulated bool) clients.Client {
if simulated {
return s.clients["simulation"]
@@ -1483,6 +1682,7 @@ func (s *Service) requeueRateLimitedTask(ctx context.Context, task store.Gateway
if delay <= 0 {
delay = 5 * time.Second
}
delay += taskRetryJitter(task.ID)
queued, err := s.store.RequeueTask(ctx, task.ID, task.ExecutionToken, delay, candidate.QueueKey)
if err != nil {
return store.GatewayTask{}, 0, err
@@ -1498,6 +1698,109 @@ func (s *Service) requeueRateLimitedTask(ctx context.Context, task store.Gateway
return queued, delay, nil
}
func (s *Service) requeueModelCoolingTask(ctx context.Context, task store.GatewayTask, cause error) (store.GatewayTask, time.Duration, error) {
delay := store.ModelCandidateRetryAfter(cause)
recoveryAt := store.ModelCandidateRecoveryAt(cause)
if recoveryAt.IsZero() {
recoveryAt = time.Now().Add(delay)
}
if delay <= 0 {
delay = time.Until(recoveryAt)
}
if delay < time.Second {
delay = time.Second
recoveryAt = time.Now().Add(delay)
}
queued, err := s.store.RequeueTaskUntil(ctx, task.ID, task.ExecutionToken, recoveryAt, "")
if err != nil {
return store.GatewayTask{}, 0, err
}
payload := map[string]any{
"code": store.ModelCandidateErrorCode(cause),
"retryAfterSeconds": int((delay + time.Second - 1) / time.Second),
"recoveryAt": recoveryAt.UTC().Format(time.RFC3339),
}
if eventErr := s.emit(ctx, task.ID, "task.queued", "queued", "model_cooling_down", 0.2, "task queued until a model candidate recovers", payload, task.RunMode == "simulation"); eventErr != nil {
return store.GatewayTask{}, 0, eventErr
}
return queued, delay, nil
}
func taskRetryJitter(taskID string) time.Duration {
var sum uint32
for _, value := range []byte(taskID) {
sum = sum*33 + uint32(value)
}
return time.Duration(sum%251) * time.Millisecond
}
func (s *Service) startConcurrencyLeaseRenewal(ctx context.Context, taskID string, leases []store.ConcurrencyLease) (context.Context, func() error) {
if len(leases) == 0 {
return ctx, func() error { return nil }
}
interval := 30 * time.Second
for _, lease := range leases {
ttl := lease.TTL
if ttl <= 0 {
ttl = 120 * time.Second
}
candidate := ttl / 3
if candidate < time.Second {
candidate = time.Second
}
if candidate < interval {
interval = candidate
}
}
runCtx, cancelRun := context.WithCancel(ctx)
renewCtx, cancelRenew := context.WithCancel(ctx)
done := make(chan error, 1)
go func() {
ticker := time.NewTicker(interval)
defer ticker.Stop()
for {
select {
case <-renewCtx.Done():
done <- nil
return
case <-ticker.C:
if err := s.store.RenewConcurrencyLeases(renewCtx, leases); err != nil {
if renewCtx.Err() != nil {
done <- nil
return
}
outcome := "failure"
if errors.Is(err, store.ErrConcurrencyLeaseLost) {
outcome = "lost"
}
s.observeConcurrencyLeaseRenewal(outcome)
s.logger.Error("concurrency lease renewal failed; cancelling upstream execution",
"taskID", taskID, "leaseCount", len(leases), "outcome", outcome, "error", err)
done <- err
cancelRun()
return
}
s.observeConcurrencyLeaseRenewal("success")
}
}
}()
return runCtx, func() error {
cancelRenew()
renewalErr := <-done
cancelRun()
return renewalErr
}
}
func (s *Service) observeConcurrencyLeaseRenewal(outcome string) {
observer, ok := s.billingMetrics.(interface {
ObserveConcurrencyLeaseRenewal(string)
})
if ok {
observer.ObserveConcurrencyLeaseRenewal(outcome)
}
}
func (s *Service) requeueInterruptedAsyncTask(ctx context.Context, task store.GatewayTask) (store.GatewayTask, error) {
queued, err := s.store.RequeueTask(ctx, task.ID, task.ExecutionToken, 0, "")
if err != nil {
@@ -1514,17 +1817,6 @@ func (s *Service) requeueInterruptedAsyncTask(ctx context.Context, task store.Ga
}
func (s *Service) withAttemptHistory(ctx context.Context, taskID string, metrics map[string]any) map[string]any {
attempts, err := s.store.ListTaskAttempts(ctx, taskID)
if err != nil {
s.logger.Warn("list task attempts for metrics failed", "taskID", taskID, "error", err)
return metrics
}
if len(attempts) == 0 {
return metrics
}
metrics = mergeMetrics(metrics)
metrics["attemptCount"] = len(attempts)
metrics["attempts"] = summarizeAttempts(attempts)
return metrics
}
@@ -1533,6 +1825,9 @@ func (s *Service) emit(ctx context.Context, taskID string, eventType string, sta
if err != nil {
return err
}
if event.ID == "" {
s.observeTaskEventSkip(event.SkippedReason)
}
if s.cfg.TaskProgressCallbackEnabled {
return s.store.QueueTaskCallback(ctx, event, s.cfg.TaskProgressCallbackURL)
}
+4 -3
View File
@@ -58,10 +58,11 @@ func (s *Service) CancelTask(ctx context.Context, taskID string, user *auth.User
return taskCancelUnavailable(task, "任务已开始执行,当前阶段不可取消,请继续查询结果"), nil
}
if task.RiverJobID > 0 {
if s.riverClient == nil {
riverClient := s.asyncControlClient()
if riverClient == nil {
return taskCancelUnavailable(task, "任务取消队列未就绪,请继续查询结果"), nil
}
job, err := s.riverClient.JobGet(ctx, task.RiverJobID)
job, err := riverClient.JobGet(ctx, task.RiverJobID)
if errors.Is(err, rivertype.ErrNotFound) {
return taskCancelUnavailable(task, "任务已不在本地排队队列,可能已提交上游,当前不可取消,请继续查询结果"), nil
}
@@ -71,7 +72,7 @@ func (s *Service) CancelTask(ctx context.Context, taskID string, user *auth.User
if job == nil || !riverJobStateCancellable(job.State) {
return taskCancelUnavailable(task, "任务已不在可取消队列状态,请继续查询结果"), nil
}
if _, err := s.riverClient.JobDelete(ctx, task.RiverJobID); err != nil {
if _, err := riverClient.JobDelete(ctx, task.RiverJobID); err != nil {
if errors.Is(err, rivertype.ErrJobRunning) || errors.Is(err, rivertype.ErrNotFound) {
return taskCancelUnavailable(task, "任务已被工作进程领取,当前不可取消,请继续查询结果"), nil
}
@@ -0,0 +1,181 @@
package runner
import (
"bytes"
"context"
"encoding/json"
"fmt"
"hash/fnv"
"io"
"net/http"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
"github.com/google/uuid"
)
func (s *Service) StartTaskHistoryWorkers(ctx context.Context) {
if s.cfg.TaskProgressCallbackEnabled && s.cfg.TaskProgressCallbackURL != "" {
go s.runTaskCallbackWorker(ctx, "callback-"+uuid.NewString())
}
if s.cfg.TaskCleanupEnabled {
go s.runTaskHistoryCleanupWorker(ctx)
}
}
func (s *Service) runTaskCallbackWorker(ctx context.Context, workerID string) {
ticker := time.NewTicker(time.Second)
defer ticker.Stop()
client := &http.Client{Timeout: time.Duration(s.cfg.TaskProgressCallbackTimeoutMS) * time.Millisecond}
for {
s.processTaskCallbackBatch(ctx, client, workerID)
select {
case <-ctx.Done():
return
case <-ticker.C:
}
}
}
func (s *Service) processTaskCallbackBatch(ctx context.Context, client *http.Client, workerID string) {
items, err := s.store.ClaimTaskCallbacks(ctx, workerID, store.TaskCallbackBatchSize, store.TaskCallbackLockTTL)
if err != nil {
if ctx.Err() == nil {
s.logger.Error("claim task callbacks failed", "error_category", "task_callback_claim_failed")
}
return
}
for _, item := range items {
if ctx.Err() != nil {
return
}
statusCode, deliveryErr := deliverTaskCallback(ctx, client, item, s.cfg.ServerMainInternalToken)
if deliveryErr == nil && statusCode >= 200 && statusCode < 300 {
if err := s.store.MarkTaskCallbackDelivered(context.WithoutCancel(ctx), item); err != nil {
s.logger.Error("mark task callback delivered failed", "taskID", item.TaskID, "seq", item.Seq, "error_category", "task_callback_state_failed")
}
continue
}
retryable := deliveryErr != nil || statusCode == http.StatusRequestTimeout || statusCode == http.StatusTooManyRequests || statusCode >= 500
retry := retryable && item.Attempts < s.cfg.TaskProgressCallbackMaxAttempts
message := "callback_network_error"
if deliveryErr == nil {
message = fmt.Sprintf("callback_http_%d", statusCode)
}
nextAttemptAt := time.Now().Add(taskCallbackRetryDelay(item.ID, item.Attempts))
if err := s.store.MarkTaskCallbackFailed(context.WithoutCancel(ctx), item, retry, nextAttemptAt, message); err != nil {
s.logger.Error("mark task callback failed", "taskID", item.TaskID, "seq", item.Seq, "error_category", "task_callback_state_failed")
continue
}
if !retry {
s.logger.Warn("task callback moved to dead letter", "taskID", item.TaskID, "seq", item.Seq, "statusCode", statusCode, "error_category", "task_callback_failed")
}
}
}
func deliverTaskCallback(ctx context.Context, client *http.Client, item store.TaskCallbackDelivery, bearerToken string) (int, error) {
body, err := json.Marshal(map[string]any{
"taskId": item.TaskID,
"seq": item.Seq,
"eventType": item.EventType,
"status": item.TaskStatus,
"createdAt": item.CreatedAt.UTC().Format(time.RFC3339Nano),
})
if err != nil {
return 0, err
}
request, err := http.NewRequestWithContext(ctx, http.MethodPost, item.CallbackURL, bytes.NewReader(body))
if err != nil {
return 0, err
}
request.Header.Set("Content-Type", "application/json")
request.Header.Set("Idempotency-Key", fmt.Sprintf("%s:%d", item.TaskID, item.Seq))
request.Header.Set("X-EasyAI-Event-Type", item.EventType)
if bearerToken != "" {
request.Header.Set("Authorization", "Bearer "+bearerToken)
}
response, err := client.Do(request)
if err != nil {
return 0, err
}
defer response.Body.Close()
_, _ = io.Copy(io.Discard, io.LimitReader(response.Body, 4096))
return response.StatusCode, nil
}
func taskCallbackRetryDelay(deliveryID string, attempt int) time.Duration {
delays := []time.Duration{
5 * time.Second,
30 * time.Second,
2 * time.Minute,
10 * time.Minute,
30 * time.Minute,
}
index := attempt - 1
if index < 0 {
index = 0
}
if index >= len(delays) {
index = len(delays) - 1
}
base := delays[index]
hash := fnv.New32a()
_, _ = hash.Write([]byte(deliveryID))
jitter := time.Duration(hash.Sum32()%21) * base / 100
return base + jitter
}
func (s *Service) runTaskHistoryCleanupWorker(ctx context.Context) {
interval := time.Duration(s.cfg.TaskCleanupIntervalSeconds) * time.Second
ticker := time.NewTicker(interval)
defer ticker.Stop()
for {
s.processTaskHistoryCleanup(ctx)
select {
case <-ctx.Done():
return
case <-ticker.C:
}
}
}
func (s *Service) processTaskHistoryCleanup(ctx context.Context) {
for batch := 0; batch < 10; batch++ {
now := time.Now().UTC()
result, err := s.store.CleanupTaskHistory(
ctx,
now.AddDate(0, 0, -s.cfg.TaskAnalysisRetentionDays),
now.AddDate(0, 0, -s.cfg.TaskRetentionDays),
s.cfg.TaskCleanupBatchSize,
)
if err != nil {
if ctx.Err() == nil {
s.logger.Error("cleanup task history failed", "error_category", "task_history_cleanup_failed")
}
return
}
if result.Total() == 0 {
return
}
s.logger.Info("task history cleanup batch completed",
"compactedTasks", result.CompactedTasks,
"compactedCheckpoints", result.CompactedCheckpoints,
"compactedAttempts", result.CompactedAttempts,
"compactedEvents", result.CompactedEvents,
"compactedCallbacks", result.CompactedCallbacks,
"retiredLegacyCallbacks", result.RetiredCallbacks,
"compactedParamLogs", result.CompactedParamLogs,
"compactedClonedVoices", result.CompactedClonedVoices,
"deletedCallbacks", result.DeletedCallbacks,
"deletedParamLogs", result.DeletedParamLogs,
"deletedEvents", result.DeletedEvents,
"deletedAttempts", result.DeletedAttempts,
"deletedTasks", result.DeletedTasks,
)
select {
case <-ctx.Done():
return
case <-time.After(100 * time.Millisecond):
}
}
}
@@ -0,0 +1,118 @@
package runner
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"testing"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
func TestMinimalRemoteTaskCheckpointUsesAllowlist(t *testing.T) {
checkpoint := minimalRemoteTaskCheckpoint("keling", "", map[string]any{
"endpoint": "/videos/task",
"mode": "omni_video",
"cleanupElementIds": []any{"element-1"},
"phase": "uploaded",
"submit": map[string]any{"large": "provider response"},
"payload": map[string]any{"prompt": "duplicate request"},
"raw": map[string]any{"data": "duplicate result"},
"status": "processing",
})
if checkpoint["endpoint"] != "/videos/task" || checkpoint["mode"] != "omni_video" || checkpoint["cleanupElementIds"] == nil {
t.Fatalf("checkpoint lost required fields: %#v", checkpoint)
}
for _, forbidden := range []string{"phase", "submit", "payload", "raw", "status"} {
if _, ok := checkpoint[forbidden]; ok {
t.Fatalf("checkpoint contains forbidden field %q: %#v", forbidden, checkpoint)
}
}
if unknown := minimalRemoteTaskCheckpoint("unknown-provider", "", map[string]any{"endpoint": "/forbidden"}); len(unknown) != 0 {
t.Fatalf("unknown provider checkpoint=%#v, want empty", unknown)
}
}
func TestCanonicalTaskResultDropsProviderRawCopy(t *testing.T) {
result := canonicalTaskResult(map[string]any{
"id": "task",
"data": []any{map[string]any{"url": "https://example.com/output"}},
"raw": map[string]any{"provider": "duplicate"},
"raw_data": map[string]any{"provider": "duplicate"},
"provider_response": map[string]any{"provider": "duplicate"},
"upstream_task_id": "remote-task",
"compatibility_data": "kept",
})
if result["raw"] != nil || result["raw_data"] != nil || result["provider_response"] != nil ||
result["upstream_task_id"] != nil || result["id"] != "task" || result["data"] == nil ||
result["compatibility_data"] != "kept" {
t.Fatalf("canonical result=%#v", result)
}
}
func TestMinimalParameterPreprocessingChangesDropsValuesAndCapsSize(t *testing.T) {
changes := make([]parameterPreprocessChange, 0, 100)
for index := 0; index < 100; index++ {
changes = append(changes, parameterPreprocessChange{
Path: "input.image",
Action: "converted",
Before: map[string]any{"base64": "forbidden"},
After: map[string]any{"url": "forbidden"},
})
}
minimal := minimalParameterPreprocessingChanges(changes)
encoded, err := json.Marshal(minimal)
if err != nil {
t.Fatal(err)
}
if len(encoded) > 1024 {
t.Fatalf("minimal changes size=%d, want <=1024", len(encoded))
}
for _, raw := range minimal {
item := raw.(map[string]any)
if len(item) != 2 || item["path"] != "input.image" || item["action"] != "converted" {
t.Fatalf("unexpected minimal change: %#v", item)
}
}
}
func TestDeliverTaskCallbackSendsMinimalAuthenticatedBody(t *testing.T) {
createdAt := time.Date(2026, time.July, 24, 10, 0, 0, 0, time.UTC)
var received map[string]any
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) {
if request.Header.Get("Authorization") != "Bearer internal-token" {
t.Errorf("Authorization=%q", request.Header.Get("Authorization"))
}
if request.Header.Get("Idempotency-Key") != "task-id:3" {
t.Errorf("Idempotency-Key=%q", request.Header.Get("Idempotency-Key"))
}
if request.Header.Get("X-EasyAI-Event-Type") != "task.completed" {
t.Errorf("X-EasyAI-Event-Type=%q", request.Header.Get("X-EasyAI-Event-Type"))
}
if err := json.NewDecoder(request.Body).Decode(&received); err != nil {
t.Error(err)
}
w.WriteHeader(http.StatusNoContent)
}))
defer server.Close()
status, err := deliverTaskCallback(context.Background(), server.Client(), store.TaskCallbackDelivery{
TaskID: "task-id",
Seq: 3,
CallbackURL: server.URL,
EventType: "task.completed",
TaskStatus: "succeeded",
CreatedAt: createdAt,
}, "internal-token")
if err != nil || status != http.StatusNoContent {
t.Fatalf("status=%d err=%v", status, err)
}
if len(received) != 5 {
t.Fatalf("callback fields=%#v, want exactly five fields", received)
}
if received["taskId"] != "task-id" || received["eventType"] != "task.completed" || received["status"] != "succeeded" {
t.Fatalf("callback body=%#v", received)
}
}
+23
View File
@@ -48,6 +48,29 @@ func failoverTraceEntry(decision failoverDecision, candidate store.RuntimeModelC
return entry
}
func failureDecisionTraceEntry(decision failureDecision, candidate store.RuntimeModelCandidate, clientAttempt int, maxAttempts int, applied bool, guardReason string) map[string]any {
entry := policyTraceEntry("failure_decision", "candidate_failure", decision.Route, decision.Reason, decision.Match, decision.Info)
entry["route"] = decision.Route
entry["effect"] = decision.Effect
entry["applied"] = applied
entry["guardReason"] = guardReason
entry["clientAttempt"] = clientAttempt
entry["maxAttempts"] = maxAttempts
entry["rule"] = decision.Match.Rule
entry["value"] = decision.Match.Value
if decision.Target != "" {
entry["target"] = decision.Target
}
if decision.CooldownSeconds > 0 {
entry["cooldownSeconds"] = decision.CooldownSeconds
}
if decision.DemoteSteps > 0 {
entry["demoteSteps"] = decision.DemoteSteps
}
addCandidatePriorityTraceFields(entry, candidate)
return entry
}
func priorityDemoteTraceEntry(decision priorityDemoteDecision, platformID string, platformModelID string, dynamicPriority int) map[string]any {
entry := policyTraceEntry("priority_demoted", "priority_demote", "demote", decision.Reason, decision.Match, decision.Info)
entry["demote"] = decision.Demote
+82 -27
View File
@@ -44,9 +44,10 @@ type FileUploadPayload struct {
}
type generatedAssetUploadPolicy struct {
UploadInlineMedia bool
UploadURLMedia bool
StoreInlineMediaLocally bool
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,17 @@ 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 {
return s.finalizeGeneratedAssets(ctx, taskID, taskKind, result, policy, nil, false, 0)
}
// 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" {
@@ -114,15 +127,16 @@ func (s *Service) uploadGeneratedAssets(ctx context.Context, taskID string, task
}
rawNeedsUpload = rawNeedsUpload && policy.UploadInlineMedia
if !needsUpload && !changed && !rawNeedsUpload {
redactGeneratedResultRawData(result)
return result, nil
return s.finalizeGeneratedAssets(ctx, taskID, taskKind, result, policy, nil, false, 0)
}
var channels []store.FileStorageChannel
if (needsUpload && generatedAssetNeedsChannelLookup(policy, decisions)) || (rawNeedsUpload && !policy.StoreInlineMediaLocally) {
channelsLoaded := false
if needsUpload || rawNeedsUpload {
channels, err = s.activeFileStorageChannels(ctx, store.FileStorageSceneImageResult)
if err != nil {
return nil, &clients.ClientError{Code: "upload_config_failed", Message: err.Error(), Retryable: true}
}
channelsLoaded = true
}
next := map[string]any{}
for key, value := range result {
@@ -151,7 +165,7 @@ func (s *Service) uploadGeneratedAssets(ctx context.Context, taskID string, task
var contentType string
var err error
if decision.Inline != nil {
upload, contentType, kind, strategy, err = s.uploadGeneratedAsset(ctx, taskID, decision.Inline, index, channels, policy.StoreInlineMediaLocally)
upload, contentType, kind, strategy, err = s.uploadGeneratedAsset(ctx, taskID, decision.Inline, index, channels)
sourceKey = decision.Inline.SourceKey
} else {
upload, contentType, kind, strategy, err = s.uploadGeneratedURLAsset(ctx, taskID, decision.URL, index, channels)
@@ -203,8 +217,59 @@ func (s *Service) uploadGeneratedAssets(ctx context.Context, taskID string, task
next["raw"] = raw
}
}
redactGeneratedResultRawData(next)
return next, nil
return s.finalizeGeneratedAssets(ctx, taskID, taskKind, next, policy, channels, channelsLoaded, len(nextData))
}
func (s *Service) finalizeGeneratedAssets(
ctx context.Context,
taskID string,
taskKind string,
result map[string]any,
policy generatedAssetUploadPolicy,
channels []store.FileStorageChannel,
channelsLoaded bool,
startIndex int,
) (map[string]any, error) {
redactGeneratedResultRawData(result)
if !TaskResultHasInlineBinary(result) {
return result, nil
}
next := result
if policy.UploadInlineMedia {
var err error
if !channelsLoaded {
channels, err = s.activeFileStorageChannels(ctx, store.FileStorageSceneImageResult)
if err != nil {
return nil, &clients.ClientError{Code: "upload_config_failed", Message: err.Error(), Retryable: true}
}
}
index := startIndex
uploaded, changed, err := s.uploadGeneratedRawMediaValue(ctx, taskID, taskKind, next, "", nil, policy, channels, &index)
if err != nil {
return nil, err
}
if changed {
mapped, ok := uploaded.(map[string]any)
if !ok {
return nil, &clients.ClientError{
Code: "result_binary_not_materialized",
Message: "generated result is not a JSON object",
StatusCode: 500,
Retryable: false,
}
}
next = mapped
}
}
if !TaskResultHasInlineBinary(next) {
return next, nil
}
persistent, _, err := s.materializeLocalBinaryResult(ctx, taskID, next)
if err != nil {
return nil, err
}
redactGeneratedResultRawData(persistent)
return persistent, nil
}
func generatedRawValueHasInlineMedia(value any, key string, siblings map[string]any) bool {
@@ -269,7 +334,7 @@ func (s *Service) uploadGeneratedRawMediaValue(ctx context.Context, taskID strin
if !ok {
return value, false, nil
}
upload, contentType, kind, strategy, err := s.uploadGeneratedAsset(ctx, taskID, asset, *index, channels, policy.StoreInlineMediaLocally)
upload, contentType, kind, strategy, err := s.uploadGeneratedAsset(ctx, taskID, asset, *index, channels)
if err != nil {
return nil, false, err
}
@@ -542,25 +607,13 @@ func generatedAssetUploadPolicyFromName(policyName string) generatedAssetUploadP
case store.FileStorageResultUploadPolicyUploadAll:
return generatedAssetUploadPolicy{UploadInlineMedia: true, UploadURLMedia: true}
case store.FileStorageResultUploadPolicyUploadNone:
return generatedAssetUploadPolicy{UploadInlineMedia: true, UploadURLMedia: false, StoreInlineMediaLocally: true}
return generatedAssetUploadPolicy{UploadInlineMedia: false, UploadURLMedia: false, PreserveInlineMedia: true, LocalizeInlineMedia: true}
default:
return defaultGeneratedAssetUploadPolicy()
}
}
func generatedAssetNeedsChannelLookup(policy generatedAssetUploadPolicy, decisions []generatedAssetDecision) bool {
for _, decision := range decisions {
if decision.URL != nil {
return true
}
if decision.Inline != nil && !policy.StoreInlineMediaLocally {
return true
}
}
return false
}
func (s *Service) uploadGeneratedAsset(ctx context.Context, taskID string, asset *generatedInlineAsset, index int, channels []store.FileStorageChannel, forceLocal bool) (map[string]any, string, string, string, error) {
func (s *Service) uploadGeneratedAsset(ctx context.Context, taskID string, asset *generatedInlineAsset, index int, channels []store.FileStorageChannel) (map[string]any, string, string, string, error) {
contentType := resolvedGeneratedAssetContentType(asset.ContentType, asset.Kind, asset.Bytes)
kind := generatedAssetKindFromContentType(asset.Kind, contentType)
payload := FileUploadPayload{
@@ -570,7 +623,7 @@ func (s *Service) uploadGeneratedAsset(ctx context.Context, taskID string, asset
Scene: store.FileStorageSceneImageResult,
Source: "ai-gateway",
}
if forceLocal || len(channels) == 0 {
if len(channels) == 0 {
upload, err := s.storeFileLocally(payload, s.cfg.LocalGeneratedStorageDir, config.DefaultLocalGeneratedStorageDir, localStaticGeneratedPathPrefix)
return upload, contentType, kind, "local_static_inline_media", err
}
@@ -1005,7 +1058,9 @@ func generatedAssetDecisionForItem(taskKind string, item map[string]any, policy
urlKey, mediaURL := mediaURLSourceFromItem(item)
if mediaURL != "" {
if !policy.UploadURLMedia {
decision.StripKeys = inlineMediaKeys(item)
if !policy.PreserveInlineMedia {
decision.StripKeys = inlineMediaKeys(item)
}
return decision, nil
}
contentType := mediaContentTypeFromItem(item)
+78 -11
View File
@@ -8,6 +8,7 @@ import (
"path/filepath"
"strings"
"testing"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
@@ -170,7 +171,7 @@ func TestGeneratedAssetDecisionUploadsURLWhenPolicyUploadAll(t *testing.T) {
}
}
func TestGeneratedAssetDecisionStoresInlineLocallyWhenPolicyUploadNone(t *testing.T) {
func TestGeneratedAssetDecisionPreservesInlineWhenPolicyUploadNone(t *testing.T) {
item := map[string]any{
"b64_json": base64.StdEncoding.EncodeToString([]byte("inline image")),
}
@@ -179,11 +180,26 @@ func TestGeneratedAssetDecisionStoresInlineLocallyWhenPolicyUploadNone(t *testin
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if decision.Inline == nil || decision.URL != nil {
t.Fatalf("upload_none should still turn inline payloads into static URLs: %+v", decision)
if decision.Inline != nil || decision.URL != nil {
t.Fatalf("upload_none should not transfer inline payloads: %+v", decision)
}
if !containsString(decision.StripKeys, "b64_json") {
t.Fatalf("inline payload should be stripped before persistence: %+v", decision.StripKeys)
if len(decision.StripKeys) != 0 {
t.Fatalf("upload_none should preserve b64_json for the caller: %+v", decision.StripKeys)
}
}
func TestGeneratedAssetDecisionPreservesInlineAlongsideURLWhenPolicyUploadNone(t *testing.T) {
item := map[string]any{
"url": "https://cdn.example.com/generated.png",
"b64_json": base64.StdEncoding.EncodeToString([]byte("inline image")),
}
decision, err := generatedAssetDecisionForItem("images.generations", item, generatedAssetUploadPolicyFromName(store.FileStorageResultUploadPolicyUploadNone))
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if decision.Inline != nil || decision.URL != nil || len(decision.StripKeys) != 0 {
t.Fatalf("upload_none should preserve the upstream URL and base64 fields unchanged: %+v", decision)
}
}
@@ -196,17 +212,17 @@ func TestGeneratedAssetUploadPolicyFromName(t *testing.T) {
{
name: "default",
policyName: store.FileStorageResultUploadPolicyDefault,
want: generatedAssetUploadPolicy{UploadInlineMedia: true, UploadURLMedia: false, StoreInlineMediaLocally: false},
want: generatedAssetUploadPolicy{UploadInlineMedia: true, UploadURLMedia: false, PreserveInlineMedia: false},
},
{
name: "upload all",
policyName: store.FileStorageResultUploadPolicyUploadAll,
want: generatedAssetUploadPolicy{UploadInlineMedia: true, UploadURLMedia: true, StoreInlineMediaLocally: false},
want: generatedAssetUploadPolicy{UploadInlineMedia: true, UploadURLMedia: true, PreserveInlineMedia: false},
},
{
name: "upload none",
policyName: store.FileStorageResultUploadPolicyUploadNone,
want: generatedAssetUploadPolicy{UploadInlineMedia: true, UploadURLMedia: false, StoreInlineMediaLocally: true},
want: generatedAssetUploadPolicy{UploadInlineMedia: false, UploadURLMedia: false, PreserveInlineMedia: true, LocalizeInlineMedia: true},
},
}
@@ -220,6 +236,52 @@ func TestGeneratedAssetUploadPolicyFromName(t *testing.T) {
}
}
func TestFinalizeGeneratedAssetsUploadsNestedInlineBinaryUnderDefaultPolicy(t *testing.T) {
storageDir := t.TempDir()
service := &Service{cfg: config.Config{LocalGeneratedStorageDir: storageDir}}
payload := append([]byte{0x89, 'P', 'N', 'G', 0x0d, 0x0a, 0x1a, 0x0a}, bytes.Repeat([]byte{0}, 160)...)
encoded := base64.StdEncoding.EncodeToString(payload)
result := map[string]any{
"id": "image-edit",
"provider_result": map[string]any{
"binary_data_base64": encoded,
"mime_type": "image/png",
},
}
finalized, err := service.finalizeGeneratedAssets(
context.Background(),
"task-nested-binary",
"images.edits",
result,
defaultGeneratedAssetUploadPolicy(),
nil,
true,
0,
)
if err != nil {
t.Fatalf("finalize generated assets: %v", err)
}
if TaskResultHasInlineBinary(finalized) {
t.Fatal("finalized result still contains inline binary")
}
nested := finalized["provider_result"].(map[string]any)
reference, ok := nested["binary_data_base64"].(map[string]any)
if !ok {
t.Fatalf("nested binary should be replaced by an asset reference: %+v", nested)
}
if urlValue := stringFromAny(reference["url"]); !strings.HasPrefix(urlValue, "/static/generated/gateway-result-task-nested-binary-01-") {
t.Fatalf("unexpected local static URL: %s", urlValue)
}
entries, err := os.ReadDir(storageDir)
if err != nil {
t.Fatalf("read generated storage: %v", err)
}
if len(entries) != 1 {
t.Fatalf("expected one localized result file, got %d", len(entries))
}
}
func TestResolvedGeneratedAssetContentTypePrefersDetectedMedia(t *testing.T) {
pngPayload := []byte{0x89, 'P', 'N', 'G', 0x0d, 0x0a, 0x1a, 0x0a, 0, 0, 0, 0}
@@ -261,7 +323,7 @@ func TestUploadGeneratedAssetStoresLocalWhenNoChannels(t *testing.T) {
SourceKey: "b64_json",
}
upload, contentType, kind, strategy, err := service.uploadGeneratedAsset(context.Background(), "task-123", asset, 0, nil, false)
upload, contentType, kind, strategy, err := service.uploadGeneratedAsset(context.Background(), "task-123", asset, 0, nil)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
@@ -272,9 +334,14 @@ func TestUploadGeneratedAssetStoresLocalWhenNoChannels(t *testing.T) {
if !strings.HasPrefix(urlValue, "/static/generated/gateway-result-task-123-01-") || !strings.HasSuffix(urlValue, ".png") {
t.Fatalf("unexpected local static URL: %s", urlValue)
}
if stringFromAny(upload["expiresAt"]) == "" {
expiresAt, err := time.Parse(time.RFC3339, stringFromAny(upload["expiresAt"]))
if err != nil {
t.Fatalf("local static upload should expose expiresAt: %+v", upload)
}
remaining := time.Until(expiresAt)
if remaining < 23*time.Hour+59*time.Minute || remaining > 24*time.Hour+time.Minute {
t.Fatalf("local static upload should expire after one day, remaining=%s", remaining)
}
entries, err := os.ReadDir(storageDir)
if err != nil {
t.Fatalf("failed to read local static dir: %v", err)
@@ -301,7 +368,7 @@ func TestUploadGeneratedAssetStoresAudioLocalWhenNoChannels(t *testing.T) {
SourceKey: "content",
}
upload, contentType, kind, strategy, err := service.uploadGeneratedAsset(context.Background(), "task-tts", asset, 0, nil, false)
upload, contentType, kind, strategy, err := service.uploadGeneratedAsset(context.Background(), "task-tts", asset, 0, nil)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
+1 -10
View File
@@ -167,16 +167,7 @@ func (s *Service) persistVoiceCloneResult(ctx context.Context, task store.Gatewa
DemoAudioURL: demoAudioURL,
Status: "active",
ExpiresAt: &expiresAt,
Metadata: map[string]any{
"request": map[string]any{
"textValidation": body["text_validation"],
"languageBoost": body["language_boost"],
"needNoiseReduction": body["need_noise_reduction"],
"needVolumeNormalization": body["need_volume_normalization"],
"aigcWatermark": body["aigc_watermark"],
},
"rawData": result["raw_data"],
},
Metadata: map[string]any{},
})
}
@@ -44,6 +44,20 @@ type Metrics struct {
billingEstimateFailed atomic.Uint64
billingIdempotentReplay atomic.Uint64
billingPricingUnavailable atomic.Uint64
asyncWorkerCapacity atomic.Int64
asyncWorkerDesiredCapacity atomic.Int64
asyncWorkerHardLimit atomic.Int64
asyncWorkerCapacityCapped atomic.Int64
asyncWorkerResizeSuccess atomic.Uint64
asyncWorkerRefreshFailed atomic.Uint64
asyncWorkerCreateFailed atomic.Uint64
asyncWorkerStartFailed atomic.Uint64
leaseRenewalSuccess atomic.Uint64
leaseRenewalFailure atomic.Uint64
leaseRenewalLost atomic.Uint64
taskEventDuplicate atomic.Uint64
taskEventUnknownType atomic.Uint64
taskEventBudgetExceeded atomic.Uint64
}
var processingDurationBounds = [...]time.Duration{
@@ -123,6 +137,52 @@ func (m *Metrics) ObserveBillingEvent(event string) {
}
}
func (m *Metrics) SetAsyncWorkerCapacity(current, desired, hardLimit int, capped bool) {
m.asyncWorkerCapacity.Store(int64(current))
m.asyncWorkerDesiredCapacity.Store(int64(desired))
m.asyncWorkerHardLimit.Store(int64(hardLimit))
cappedValue := int64(0)
if capped {
cappedValue = 1
}
m.asyncWorkerCapacityCapped.Store(cappedValue)
}
func (m *Metrics) ObserveAsyncWorkerResize(outcome string) {
switch outcome {
case "success":
m.asyncWorkerResizeSuccess.Add(1)
case "refresh_failed":
m.asyncWorkerRefreshFailed.Add(1)
case "create_failed":
m.asyncWorkerCreateFailed.Add(1)
case "start_failed":
m.asyncWorkerStartFailed.Add(1)
}
}
func (m *Metrics) ObserveConcurrencyLeaseRenewal(outcome string) {
switch outcome {
case "success":
m.leaseRenewalSuccess.Add(1)
case "lost":
m.leaseRenewalLost.Add(1)
default:
m.leaseRenewalFailure.Add(1)
}
}
func (m *Metrics) ObserveTaskEventSkip(reason string) {
switch reason {
case "duplicate":
m.taskEventDuplicate.Add(1)
case "budget_exceeded":
m.taskEventBudgetExceeded.Add(1)
default:
m.taskEventUnknownType.Add(1)
}
}
func (m *Metrics) Handler(provider MetricsSnapshotProvider, issuer, audience string, enabled bool) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
mode := "disabled"
@@ -212,6 +272,26 @@ func (m *Metrics) Handler(provider MetricsSnapshotProvider, issuer, audience str
plainCounter(w, "easyai_gateway_billing_estimate_failures_total", "Pricing estimate requests that failed.", m.billingEstimateFailed.Load())
plainCounter(w, "easyai_gateway_billing_idempotent_replays_total", "Generation requests replayed idempotently.", m.billingIdempotentReplay.Load())
plainCounter(w, "easyai_gateway_billing_pricing_unavailable_total", "Pricing requests rejected because no effective price was available.", m.billingPricingUnavailable.Load())
plainGauge(w, "easyai_gateway_async_worker_capacity", "Current River asynchronous worker capacity.", m.asyncWorkerCapacity.Load())
plainGauge(w, "easyai_gateway_async_worker_desired_capacity", "Uncapped asynchronous worker capacity derived from model and user-group policies.", m.asyncWorkerDesiredCapacity.Load())
plainGauge(w, "easyai_gateway_async_worker_hard_limit", "Per-process asynchronous worker safety limit.", m.asyncWorkerHardLimit.Load())
plainGauge(w, "easyai_gateway_async_worker_capacity_capped", "Whether desired asynchronous worker capacity is capped by the hard limit.", m.asyncWorkerCapacityCapped.Load())
outcomeCounters(w, "easyai_gateway_async_worker_resizes_total", "Asynchronous worker resize attempts by bounded outcome.", []outcomeValue{
{"success", m.asyncWorkerResizeSuccess.Load()},
{"refresh_failed", m.asyncWorkerRefreshFailed.Load()},
{"create_failed", m.asyncWorkerCreateFailed.Load()},
{"start_failed", m.asyncWorkerStartFailed.Load()},
})
outcomeCounters(w, "easyai_gateway_concurrency_lease_renewals_total", "Concurrency lease renewals by bounded outcome.", []outcomeValue{
{"success", m.leaseRenewalSuccess.Load()},
{"failure", m.leaseRenewalFailure.Load()},
{"lost", m.leaseRenewalLost.Load()},
})
outcomeCounters(w, "easyai_gateway_task_events_skipped_total", "Task events skipped before persistence by bounded reason.", []outcomeValue{
{"duplicate", m.taskEventDuplicate.Load()},
{"unknown_type", m.taskEventUnknownType.Load()},
{"budget_exceeded", m.taskEventBudgetExceeded.Load()},
})
})
}
@@ -249,3 +329,7 @@ func outcomeCounters(w http.ResponseWriter, name, help string, values []outcomeV
func plainCounter(w http.ResponseWriter, name, help string, value uint64) {
fmt.Fprintf(w, "# HELP %s %s\n# TYPE %s counter\n%s %d\n", name, help, name, name, value)
}
func plainGauge(w http.ResponseWriter, name, help string, value int64) {
fmt.Fprintf(w, "# HELP %s %s\n# TYPE %s gauge\n%s %d\n", name, help, name, name, value)
}

Some files were not shown because too many files have changed in this diff Show More