diff --git a/.env.example b/.env.example index f551733..8c010a6 100644 --- a/.env.example +++ b/.env.example @@ -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 diff --git a/AGENTS.md b/AGENTS.md index a04d1fe..5bc4846 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -52,14 +52,14 @@ docker compose -f docker-compose.yml config --quiet 2. `main` 不设置受保护分支或 required status；Agent 在获得用户明确授权后，可以合并已验证的 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/.json`，不得修改生产。 -6. publish 成功后，Agent 必须报告 manifest、组件、digest 和验证结果并停止。只有用户再次明确确认上线，才可单独执行： + 该命令只能构建、冒烟、推送镜像和生成 `dist/releases/.json`，不得修改生产。用户只要求“构建镜像”“推送镜像”或明确限定为 publish 时，执行到此为止。 +6. 用户明确要求“发布”“上线”或“部署到线上”时，该次授权同时覆盖 publish 和 deploy。publish 成功并核验 manifest、组件、digest 和验证结果后，应直接继续执行下列部署命令，无需再次请求确认： ```bash ./scripts/deploy-production-release.sh dist/releases/.json diff --git a/Dockerfile b/Dockerfile index a5248b0..032afa9 100644 --- a/Dockerfile +++ b/Dockerfile @@ -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 && \ diff --git a/README.md b/README.md index 7a6ac21..2ba3421 100644 --- a/README.md +++ b/README.md @@ -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= registry.cn-shanghai.aliyuncs.com ./scripts/publish-release-images.sh --components auto ``` -Agent 必须报告 `dist/releases/.json` 后停止。用户再次明确确认上线后，才执行第二步： +用户要求“发布”“上线”或“部署到线上”时，Agent 核验并报告 `dist/releases/.json`、组件和 digest 后直接执行第二步，无需再次确认；用户仅要求构建或推送镜像时才在第一步后停止： ```bash ./scripts/deploy-production-release.sh dist/releases/.json @@ -124,7 +124,10 @@ Agent 必须报告 `dist/releases/.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` 内现有模块。 diff --git a/apps/api/cmd/backfill-binary-results/main.go b/apps/api/cmd/backfill-binary-results/main.go new file mode 100644 index 0000000..cffadec --- /dev/null +++ b/apps/api/cmd/backfill-binary-results/main.go @@ -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 +} diff --git a/apps/api/cmd/migrate/main.go b/apps/api/cmd/migrate/main.go index fe5a9f8..77a7455 100644 --- a/apps/api/cmd/migrate/main.go +++ b/apps/api/cmd/migrate/main.go @@ -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 +} diff --git a/apps/api/cmd/migrate/main_test.go b/apps/api/cmd/migrate/main_test.go index e886f51..63249b3 100644 --- a/apps/api/cmd/migrate/main_test.go +++ b/apps/api/cmd/migrate/main_test.go @@ -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 { diff --git a/apps/api/docs/swagger.json b/apps/api/docs/swagger.json index a8684a0..e4f5bc9 100644 --- a/apps/api/docs/swagger.json +++ b/apps/api/docs/swagger.json @@ -584,6 +584,12 @@ "$ref": "#/definitions/httpapi.ErrorEnvelope" } }, + "409": { + "description": "Conflict", + "schema": { + "$ref": "#/definitions/httpapi.ErrorEnvelope" + } + }, "500": { "description": "Internal Server Error", "schema": { @@ -3760,6 +3766,299 @@ } } }, + "/api/admin/tasks": { + "get": { + "security": [ + { + "BearerAuth": [] + } + ], + "description": "跨租户分页查询任务；请求、结果和执行快照中的常见敏感字段默认脱敏。", + "produces": [ + "application/json" + ], + "tags": [ + "admin-tasks" + ], + "summary": "管理员列出任务", + "parameters": [ + { + "type": "string", + "description": "关键词，匹配任务、用户、租户、模型、平台和 API Key", + "name": "q", + "in": "query" + }, + { + "type": "string", + "description": "网关租户 ID", + "name": "tenantId", + "in": "query" + }, + { + "type": "string", + "description": "网关用户 ID", + "name": "userId", + "in": "query" + }, + { + "type": "string", + "description": "用户组 ID", + "name": "userGroupId", + "in": "query" + }, + { + "type": "string", + "description": "任务状态", + "name": "status", + "in": "query" + }, + { + "type": "string", + "description": "执行平台 ID", + "name": "platformId", + "in": "query" + }, + { + "type": "string", + "description": "调用或实际模型名称", + "name": "model", + "in": "query" + }, + { + "type": "string", + "description": "模型类型", + "name": "modelType", + "in": "query" + }, + { + "type": "string", + "description": "运行模式", + "name": "runMode", + "in": "query" + }, + { + "type": "string", + "description": "计费状态", + "name": "billingStatus", + "in": "query" + }, + { + "type": "string", + "description": "API Key ID、名称或前缀", + "name": "apiKey", + "in": "query" + }, + { + "type": "string", + "description": "创建时间起点，支持 RFC3339 或日期格式", + "name": "createdFrom", + "in": "query" + }, + { + "type": "string", + "description": "创建时间终点，支持 RFC3339 或日期格式", + "name": "createdTo", + "in": "query" + }, + { + "type": "integer", + "default": 1, + "description": "页码", + "name": "page", + "in": "query" + }, + { + "type": "integer", + "default": 50, + "description": "每页数量", + "name": "pageSize", + "in": "query" + } + ], + "responses": { + "200": { + "description": "OK", + "schema": { + "$ref": "#/definitions/httpapi.AdminTaskListResponse" + } + }, + "400": { + "description": "Bad Request", + "schema": { + "$ref": "#/definitions/httpapi.ErrorEnvelope" + } + }, + "401": { + "description": "Unauthorized", + "schema": { + "$ref": "#/definitions/httpapi.ErrorEnvelope" + } + }, + "403": { + "description": "Forbidden", + "schema": { + "$ref": "#/definitions/httpapi.ErrorEnvelope" + } + }, + "500": { + "description": "Internal Server Error", + "schema": { + "$ref": "#/definitions/httpapi.ErrorEnvelope" + } + } + } + } + }, + "/api/admin/tasks/{taskID}": { + "get": { + "security": [ + { + "BearerAuth": [] + } + ], + "description": "返回任意租户的任务详情；默认脱敏，sensitive=full 时返回完整存储内容。", + "produces": [ + "application/json" + ], + "tags": [ + "admin-tasks" + ], + "summary": "管理员获取任务详情", + "parameters": [ + { + "type": "string", + "description": "任务 ID", + "name": "taskID", + "in": "path", + "required": true + }, + { + "enum": [ + "masked", + "full" + ], + "type": "string", + "default": "masked", + "description": "敏感字段模式", + "name": "sensitive", + "in": "query" + } + ], + "responses": { + "200": { + "description": "OK", + "schema": { + "$ref": "#/definitions/store.AdminGatewayTask" + } + }, + "400": { + "description": "Bad Request", + "schema": { + "$ref": "#/definitions/httpapi.ErrorEnvelope" + } + }, + "401": { + "description": "Unauthorized", + "schema": { + "$ref": "#/definitions/httpapi.ErrorEnvelope" + } + }, + "403": { + "description": "Forbidden", + "schema": { + "$ref": "#/definitions/httpapi.ErrorEnvelope" + } + }, + "404": { + "description": "Not Found", + "schema": { + "$ref": "#/definitions/httpapi.ErrorEnvelope" + } + }, + "500": { + "description": "Internal Server Error", + "schema": { + "$ref": "#/definitions/httpapi.ErrorEnvelope" + } + } + } + } + }, + "/api/admin/tasks/{taskID}/param-preprocessing": { + "get": { + "security": [ + { + "BearerAuth": [] + } + ], + "description": "返回任意租户任务的参数改写、校验或模板处理日志；默认脱敏。", + "produces": [ + "application/json" + ], + "tags": [ + "admin-tasks" + ], + "summary": "管理员获取任务参数预处理日志", + "parameters": [ + { + "type": "string", + "description": "任务 ID", + "name": "taskID", + "in": "path", + "required": true + }, + { + "enum": [ + "masked", + "full" + ], + "type": "string", + "default": "masked", + "description": "敏感字段模式", + "name": "sensitive", + "in": "query" + } + ], + "responses": { + "200": { + "description": "OK", + "schema": { + "$ref": "#/definitions/httpapi.TaskParamPreprocessingLogListResponse" + } + }, + "400": { + "description": "Bad Request", + "schema": { + "$ref": "#/definitions/httpapi.ErrorEnvelope" + } + }, + "401": { + "description": "Unauthorized", + "schema": { + "$ref": "#/definitions/httpapi.ErrorEnvelope" + } + }, + "403": { + "description": "Forbidden", + "schema": { + "$ref": "#/definitions/httpapi.ErrorEnvelope" + } + }, + "404": { + "description": "Not Found", + "schema": { + "$ref": "#/definitions/httpapi.ErrorEnvelope" + } + }, + "500": { + "description": "Internal Server Error", + "schema": { + "$ref": "#/definitions/httpapi.ErrorEnvelope" + } + } + } + } + }, "/api/admin/tenants": { "get": { "security": [ @@ -4703,13 +5002,14 @@ "BearerAuth": [] } ], + "description": "查询当前用户的图片、视频、音频、视频高清和矢量化任务，并返回 EasyAI GeneratedResponse 兼容结构。", "produces": [ "application/json" ], "tags": [ - "volces-compatible" + "tasks" ], - "summary": "查询 server-main 兼容视频结果", + "summary": "查询 server-main EasyAIClient 兼容任务结果", "parameters": [ { "type": "string", @@ -4723,8 +5023,19 @@ "200": { "description": "OK", "schema": { - "type": "object", - "additionalProperties": true + "$ref": "#/definitions/httpapi.EasyAIGeneratedResponse" + } + }, + "404": { + "description": "Not Found", + "schema": { + "$ref": "#/definitions/httpapi.EasyAIGeneratedResponse" + } + }, + "410": { + "description": "Gone", + "schema": { + "$ref": "#/definitions/httpapi.EasyAIGeneratedResponse" } } } @@ -5633,6 +5944,12 @@ "schema": { "$ref": "#/definitions/httpapi.VolcesErrorEnvelope" } + }, + "410": { + "description": "Gone", + "schema": { + "$ref": "#/definitions/httpapi.VolcesErrorEnvelope" + } } } }, @@ -5970,7 +6287,7 @@ "200": { "description": "OK", "schema": { - "$ref": "#/definitions/httpapi.CompatibleResponse" + "$ref": "#/definitions/httpapi.OpenAIImagesResponse" } }, "202": { @@ -6045,7 +6362,7 @@ "200": { "description": "OK", "schema": { - "$ref": "#/definitions/httpapi.CompatibleResponse" + "$ref": "#/definitions/httpapi.OpenAIImagesResponse" } }, "202": { @@ -6120,7 +6437,7 @@ "200": { "description": "OK", "schema": { - "$ref": "#/definitions/httpapi.CompatibleResponse" + "$ref": "#/definitions/httpapi.EasyAIGeneratedResponse" } }, "202": { @@ -6292,6 +6609,13 @@ "type": "object", "additionalProperties": true } + }, + "410": { + "description": "Gone", + "schema": { + "type": "object", + "additionalProperties": true + } } } } @@ -6392,6 +6716,13 @@ "type": "object", "additionalProperties": true } + }, + "410": { + "description": "Gone", + "schema": { + "type": "object", + "additionalProperties": true + } } } }, @@ -6529,6 +6860,18 @@ "playground" ], "summary": "列出可调用模型", + "parameters": [ + { + "enum": [ + "canvas_model_node", + "desktop" + ], + "type": "string", + "description": "模型可选场景；不传时保持原有行为", + "name": "usage_scene", + "in": "query" + } + ], "responses": { "200": { "description": "OK", @@ -6536,6 +6879,12 @@ "$ref": "#/definitions/httpapi.PlatformModelListResponse" } }, + "400": { + "description": "Bad Request", + "schema": { + "$ref": "#/definitions/httpapi.ErrorEnvelope" + } + }, "401": { "description": "Unauthorized", "schema": { @@ -6547,6 +6896,12 @@ "schema": { "$ref": "#/definitions/httpapi.ErrorEnvelope" } + }, + "502": { + "description": "Bad Gateway", + "schema": { + "$ref": "#/definitions/httpapi.ErrorEnvelope" + } } } } @@ -6704,7 +7059,7 @@ "BearerAuth": [] } ], - "description": "统一公开入口按 model 选择平台模型并默认同步返回兼容响应；设置 X-Async=true 时异步创建任务并返回 202。", + "description": "默认同步返回 EasyAI GeneratedResponse；设置 X-Async=true 时返回同时兼容 Gateway 与 server-main EasyAIClient 的异步提交结构。", "consumes": [ "application/json" ], @@ -6712,9 +7067,9 @@ "application/json" ], "tags": [ - "tasks" + "media" ], - "summary": "创建或执行 AI 任务", + "summary": "创建 EasyAI 媒体任务", "parameters": [ { "type": "boolean", @@ -6729,7 +7084,7 @@ "in": "header" }, { - "description": "AI 任务请求，字段随任务类型变化", + "description": "媒体任务请求，字段随任务类型变化", "name": "input", "in": "body", "required": true, @@ -6742,7 +7097,7 @@ "200": { "description": "OK", "schema": { - "$ref": "#/definitions/httpapi.CompatibleResponse" + "$ref": "#/definitions/httpapi.EasyAIGeneratedResponse" } }, "202": { @@ -6889,6 +7244,18 @@ "playground" ], "summary": "列出可调用模型", + "parameters": [ + { + "enum": [ + "canvas_model_node", + "desktop" + ], + "type": "string", + "description": "模型可选场景；不传时保持原有行为", + "name": "usage_scene", + "in": "query" + } + ], "responses": { "200": { "description": "OK", @@ -6896,6 +7263,12 @@ "$ref": "#/definitions/httpapi.PlatformModelListResponse" } }, + "400": { + "description": "Bad Request", + "schema": { + "$ref": "#/definitions/httpapi.ErrorEnvelope" + } + }, "401": { "description": "Unauthorized", "schema": { @@ -6907,6 +7280,12 @@ "schema": { "$ref": "#/definitions/httpapi.ErrorEnvelope" } + }, + "502": { + "description": "Bad Gateway", + "schema": { + "$ref": "#/definitions/httpapi.ErrorEnvelope" + } } } } @@ -7542,7 +7921,7 @@ "BearerAuth": [] } ], - "description": "统一公开入口按 model 选择平台模型并默认同步返回兼容响应；设置 X-Async=true 时异步创建任务并返回 202。", + "description": "默认同步返回 EasyAI GeneratedResponse；设置 X-Async=true 时返回同时兼容 Gateway 与 server-main EasyAIClient 的异步提交结构。", "consumes": [ "application/json" ], @@ -7550,9 +7929,9 @@ "application/json" ], "tags": [ - "tasks" + "media" ], - "summary": "创建或执行 AI 任务", + "summary": "创建 EasyAI 媒体任务", "parameters": [ { "type": "boolean", @@ -7567,7 +7946,7 @@ "in": "header" }, { - "description": "AI 任务请求，字段随任务类型变化", + "description": "媒体任务请求，字段随任务类型变化", "name": "input", "in": "body", "required": true, @@ -7580,7 +7959,7 @@ "200": { "description": "OK", "schema": { - "$ref": "#/definitions/httpapi.CompatibleResponse" + "$ref": "#/definitions/httpapi.EasyAIGeneratedResponse" } }, "202": { @@ -7641,7 +8020,7 @@ "BearerAuth": [] } ], - "description": "统一公开入口按 model 选择平台模型并默认同步返回兼容响应；设置 X-Async=true 时异步创建任务并返回 202。", + "description": "默认同步返回 EasyAI GeneratedResponse；设置 X-Async=true 时返回同时兼容 Gateway 与 server-main EasyAIClient 的异步提交结构。", "consumes": [ "application/json" ], @@ -7649,9 +8028,9 @@ "application/json" ], "tags": [ - "tasks" + "media" ], - "summary": "创建或执行 AI 任务", + "summary": "创建 EasyAI 媒体任务", "parameters": [ { "type": "boolean", @@ -7666,7 +8045,7 @@ "in": "header" }, { - "description": "AI 任务请求，字段随任务类型变化", + "description": "媒体任务请求，字段随任务类型变化", "name": "input", "in": "body", "required": true, @@ -7679,7 +8058,7 @@ "200": { "description": "OK", "schema": { - "$ref": "#/definitions/httpapi.CompatibleResponse" + "$ref": "#/definitions/httpapi.EasyAIGeneratedResponse" } }, "202": { @@ -7871,6 +8250,12 @@ "$ref": "#/definitions/httpapi.ErrorEnvelope" } }, + "410": { + "description": "Gone", + "schema": { + "$ref": "#/definitions/httpapi.ErrorEnvelope" + } + }, "500": { "description": "Internal Server Error", "schema": { @@ -8057,15 +8442,14 @@ "application/json" ], "tags": [ - "volces-compatible" + "media" ], "summary": "创建 server-main 兼容视频任务", "responses": { "200": { "description": "OK", "schema": { - "type": "object", - "additionalProperties": true + "$ref": "#/definitions/httpapi.TaskAcceptedResponse" } } } @@ -8078,7 +8462,7 @@ "BearerAuth": [] } ], - "description": "统一公开入口按 model 选择平台模型并默认同步返回兼容响应；设置 X-Async=true 时异步创建任务并返回 202。", + "description": "默认同步返回 EasyAI GeneratedResponse；设置 X-Async=true 时返回同时兼容 Gateway 与 server-main EasyAIClient 的异步提交结构。", "consumes": [ "application/json" ], @@ -8086,9 +8470,9 @@ "application/json" ], "tags": [ - "tasks" + "media" ], - "summary": "创建或执行 AI 任务", + "summary": "创建 EasyAI 媒体任务", "parameters": [ { "type": "boolean", @@ -8103,7 +8487,7 @@ "in": "header" }, { - "description": "AI 任务请求，字段随任务类型变化", + "description": "媒体任务请求，字段随任务类型变化", "name": "input", "in": "body", "required": true, @@ -8116,7 +8500,7 @@ "200": { "description": "OK", "schema": { - "$ref": "#/definitions/httpapi.CompatibleResponse" + "$ref": "#/definitions/httpapi.EasyAIGeneratedResponse" } }, "202": { @@ -8294,6 +8678,12 @@ "$ref": "#/definitions/httpapi.KelingCompatibleEnvelope" } }, + "410": { + "description": "Gone", + "schema": { + "$ref": "#/definitions/httpapi.KelingCompatibleEnvelope" + } + }, "500": { "description": "Internal Server Error", "schema": { @@ -8342,7 +8732,7 @@ "200": { "description": "OK", "schema": { - "$ref": "#/definitions/httpapi.CompatibleResponse" + "$ref": "#/definitions/httpapi.EasyAIGeneratedResponse" } }, "202": { @@ -8391,7 +8781,7 @@ "BearerAuth": [] } ], - "description": "统一公开入口按 model 选择平台模型并默认同步返回兼容响应；设置 X-Async=true 时异步创建任务并返回 202。", + "description": "默认同步返回 EasyAI GeneratedResponse；设置 X-Async=true 时返回同时兼容 Gateway 与 server-main EasyAIClient 的异步提交结构。", "consumes": [ "application/json" ], @@ -8399,9 +8789,9 @@ "application/json" ], "tags": [ - "tasks" + "media" ], - "summary": "创建或执行 AI 任务", + "summary": "创建 EasyAI 媒体任务", "parameters": [ { "type": "boolean", @@ -8416,7 +8806,7 @@ "in": "header" }, { - "description": "AI 任务请求，字段随任务类型变化", + "description": "媒体任务请求，字段随任务类型变化", "name": "input", "in": "body", "required": true, @@ -8429,7 +8819,7 @@ "200": { "description": "OK", "schema": { - "$ref": "#/definitions/httpapi.CompatibleResponse" + "$ref": "#/definitions/httpapi.EasyAIGeneratedResponse" } }, "202": { @@ -8755,6 +9145,12 @@ "$ref": "#/definitions/httpapi.ErrorEnvelope" } }, + "410": { + "description": "Gone", + "schema": { + "$ref": "#/definitions/httpapi.ErrorEnvelope" + } + }, "500": { "description": "Internal Server Error", "schema": { @@ -9396,6 +9792,29 @@ } } }, + "httpapi.AdminTaskListResponse": { + "type": "object", + "properties": { + "items": { + "type": "array", + "items": { + "$ref": "#/definitions/store.AdminGatewayTask" + } + }, + "page": { + "type": "integer", + "example": 1 + }, + "pageSize": { + "type": "integer", + "example": 50 + }, + "total": { + "type": "integer", + "example": 42 + } + } + }, "httpapi.AuditLogListResponse": { "type": "object", "properties": { @@ -9807,6 +10226,143 @@ } } }, + "httpapi.EasyAIGeneratedResponse": { + "type": "object", + "properties": { + "cancellable": { + "type": "boolean" + }, + "code": { + "type": "string" + }, + "created": { + "type": "integer", + "example": 1784772000000 + }, + "data": { + "type": "array", + "items": { + "$ref": "#/definitions/httpapi.EasyAIMediaOutput" + } + }, + "message": { + "type": "string" + }, + "output": { + "type": "array", + "items": { + "type": "string" + } + }, + "output_content": { + "type": "array", + "items": { + "$ref": "#/definitions/httpapi.EasyAIMediaOutput" + } + }, + "query_url": { + "type": "string", + "example": "/api/v1/ai/result/9f4d8f3d-5f5f-4bb7-a4be-344a9f930e25" + }, + "result": { + "type": "object", + "additionalProperties": true + }, + "status": { + "type": "string", + "enum": [ + "submitted", + "process", + "success", + "failed" + ], + "example": "success" + }, + "submitted": { + "type": "boolean" + }, + "taskId": { + "type": "string", + "example": "9f4d8f3d-5f5f-4bb7-a4be-344a9f930e25" + }, + "task_id": { + "type": "string", + "example": "9f4d8f3d-5f5f-4bb7-a4be-344a9f930e25" + }, + "upstream_task_id": { + "type": "string", + "example": "provider-task-123" + }, + "usage": { + "type": "object", + "additionalProperties": true + } + } + }, + "httpapi.EasyAIMediaOutput": { + "type": "object", + "properties": { + "audio_url": { + "type": "string" + }, + "b64_json": { + "description": "B64JSON is retained when result media transfer is disabled.", + "type": "string" + }, + "content": { + "type": "string" + }, + "duration": { + "type": "number" + }, + "format": { + "type": "string" + }, + "height": { + "type": "integer" + }, + "image_url": { + "type": "string" + }, + "metadata": { + "type": "object", + "additionalProperties": true + }, + "mime_type": { + "type": "string" + }, + "source_resolution": { + "type": "string" + }, + "target_resolution": { + "type": "string" + }, + "text": { + "type": "string" + }, + "type": { + "type": "string", + "enum": [ + "image", + "video", + "audio", + "file", + "text" + ], + "example": "video" + }, + "url": { + "type": "string", + "example": "https://cdn.example.com/output.mp4" + }, + "video_url": { + "type": "string" + }, + "width": { + "type": "integer" + } + } + }, "httpapi.ErrorEnvelope": { "type": "object", "properties": { @@ -9848,7 +10404,7 @@ } } }, - "httpapi.FileUploadResponse": { + "httpapi.FileUploadData": { "type": "object", "properties": { "assetStorage": { @@ -9877,6 +10433,42 @@ } } }, + "httpapi.FileUploadResponse": { + "type": "object", + "properties": { + "assetStorage": { + "type": "object", + "additionalProperties": true + }, + "contentType": { + "type": "string", + "example": "image/png" + }, + "data": { + "$ref": "#/definitions/httpapi.FileUploadData" + }, + "filename": { + "type": "string", + "example": "image.png" + }, + "id": { + "type": "string", + "example": "file_abc123" + }, + "size": { + "type": "integer", + "example": 1024 + }, + "status": { + "type": "string", + "example": "success" + }, + "url": { + "type": "string", + "example": "/static/uploaded/upload-abc123.png" + } + } + }, "httpapi.GeminiErrorEnvelope": { "type": "object", "properties": { @@ -10391,6 +10983,12 @@ "id": { "type": "string" }, + "legacyAliases": { + "type": "array", + "items": { + "type": "string" + } + }, "modelAlias": { "type": "string" }, @@ -10412,6 +11010,9 @@ "providerKey": { "type": "string" }, + "providerModelName": { + "type": "string" + }, "providerName": { "type": "string" }, @@ -10499,6 +11100,36 @@ } } }, + "httpapi.OpenAIImageData": { + "type": "object", + "properties": { + "b64_json": { + "type": "string" + }, + "revised_prompt": { + "type": "string" + }, + "url": { + "type": "string", + "example": "https://cdn.example.com/output.png" + } + } + }, + "httpapi.OpenAIImagesResponse": { + "type": "object", + "properties": { + "created": { + "type": "integer", + "example": 1710000000 + }, + "data": { + "type": "array", + "items": { + "$ref": "#/definitions/httpapi.OpenAIImageData" + } + } + } + }, "httpapi.PlatformListResponse": { "type": "object", "properties": { @@ -10879,12 +11510,24 @@ "next": { "$ref": "#/definitions/httpapi.TaskNextLinks" }, + "query_url": { + "type": "string", + "example": "/api/v1/ai/result/9f4d8f3d-5f5f-4bb7-a4be-344a9f930e25" + }, + "status": { + "type": "string", + "example": "submitted" + }, "task": { "$ref": "#/definitions/store.GatewayTask" }, "taskId": { "type": "string", "example": "9f4d8f3d-5f5f-4bb7-a4be-344a9f930e25" + }, + "task_id": { + "type": "string", + "example": "9f4d8f3d-5f5f-4bb7-a4be-344a9f930e25" } } }, @@ -10946,6 +11589,10 @@ "events": { "type": "string", "example": "/api/v1/tasks/9f4d8f3d-5f5f-4bb7-a4be-344a9f930e25/events" + }, + "result": { + "type": "string", + "example": "/api/v1/ai/result/9f4d8f3d-5f5f-4bb7-a4be-344a9f930e25" } } }, @@ -12051,6 +12698,274 @@ } } }, + "store.AdminGatewayTask": { + "type": "object", + "properties": { + "adminContext": { + "$ref": "#/definitions/store.AdminTaskContext" + }, + "apiKeyId": { + "type": "string" + }, + "apiKeyName": { + "type": "string" + }, + "apiKeyPrefix": { + "type": "string" + }, + "asyncMode": { + "type": "boolean" + }, + "attemptCount": { + "type": "integer" + }, + "attempts": { + "type": "array", + "items": { + "$ref": "#/definitions/store.TaskAttempt" + } + }, + "billingCurrency": { + "type": "string" + }, + "billingSettledAt": { + "type": "string" + }, + "billingStatus": { + "type": "string" + }, + "billingSummary": { + "type": "object", + "additionalProperties": {} + }, + "billingUpdatedAt": { + "type": "string" + }, + "billingVersion": { + "type": "string" + }, + "billings": { + "type": "array", + "items": {} + }, + "cancellable": { + "type": "boolean" + }, + "compatibilityProtocol": { + "type": "string" + }, + "compatibilityPublicId": { + "type": "string" + }, + "compatibilitySourceProtocol": { + "type": "string" + }, + "compatibilitySubmitBody": { + "type": "object", + "additionalProperties": {} + }, + "compatibilitySubmitHeaders": { + "type": "object", + "additionalProperties": {} + }, + "compatibilitySubmitHttpStatus": { + "type": "integer" + }, + "conversationId": { + "type": "string" + }, + "createdAt": { + "type": "string" + }, + "error": { + "type": "string" + }, + "errorCode": { + "type": "string" + }, + "errorMessage": { + "type": "string" + }, + "executionLeaseExpiresAt": { + "type": "string" + }, + "externalTaskId": { + "type": "string" + }, + "finalChargeAmount": { + "type": "number" + }, + "finishedAt": { + "type": "string" + }, + "gatewayTenantId": { + "type": "string" + }, + "gatewayUserId": { + "type": "string" + }, + "id": { + "type": "string" + }, + "kind": { + "type": "string" + }, + "message": { + "type": "string" + }, + "metrics": { + "type": "object", + "additionalProperties": {} + }, + "model": { + "type": "string" + }, + "modelType": { + "type": "string" + }, + "newMessageCount": { + "type": "integer" + }, + "pricingSnapshot": { + "type": "object", + "additionalProperties": {} + }, + "remoteTaskId": { + "type": "string" + }, + "request": { + "type": "object", + "additionalProperties": {} + }, + "requestFingerprint": { + "type": "string" + }, + "requestId": { + "type": "string" + }, + "requestedModel": { + "type": "string" + }, + "reservationAmount": { + "type": "number" + }, + "resolvedModel": { + "type": "string" + }, + "responseDurationMs": { + "type": "integer" + }, + "responseFinishedAt": { + "type": "string" + }, + "responseStartedAt": { + "type": "string" + }, + "result": { + "type": "object", + "additionalProperties": {} + }, + "riverJobId": { + "type": "integer" + }, + "runMode": { + "type": "string" + }, + "status": { + "type": "string" + }, + "submitted": { + "type": "boolean" + }, + "tenantId": { + "type": "string" + }, + "tenantKey": { + "type": "string" + }, + "updatedAt": { + "type": "string" + }, + "usage": { + "type": "object", + "additionalProperties": {} + }, + "userGroupId": { + "type": "string" + }, + "userGroupKey": { + "type": "string" + }, + "userId": { + "type": "string" + }, + "userSource": { + "type": "string" + } + } + }, + "store.AdminTaskContext": { + "type": "object", + "properties": { + "latestPlatform": { + "$ref": "#/definitions/store.AdminTaskPlatformSummary" + }, + "tenant": { + "$ref": "#/definitions/store.AdminTaskTenantSummary" + }, + "user": { + "$ref": "#/definitions/store.AdminTaskUserSummary" + } + } + }, + "store.AdminTaskPlatformSummary": { + "type": "object", + "properties": { + "id": { + "type": "string" + }, + "name": { + "type": "string" + }, + "provider": { + "type": "string" + } + } + }, + "store.AdminTaskTenantSummary": { + "type": "object", + "properties": { + "id": { + "type": "string" + }, + "name": { + "type": "string" + }, + "tenantKey": { + "type": "string" + } + } + }, + "store.AdminTaskUserSummary": { + "type": "object", + "properties": { + "displayName": { + "type": "string" + }, + "email": { + "type": "string" + }, + "id": { + "type": "string" + }, + "source": { + "type": "string" + }, + "username": { + "type": "string" + } + } + }, "store.AuditLog": { "type": "object", "properties": { @@ -12147,14 +13062,27 @@ "type": "object", "additionalProperties": {} }, + "displayName": { + "type": "string" + }, "id": { "type": "string" }, + "invocationName": { + "type": "string" + }, + "legacyAliases": { + "type": "array", + "items": { + "type": "string" + } + }, "metadata": { "type": "object", "additionalProperties": {} }, "modelAlias": { + "description": "Deprecated: use InvocationName.", "type": "string" }, "modelType": { @@ -12175,6 +13103,9 @@ "providerModelName": { "type": "string" }, + "referenceCount": { + "type": "integer" + }, "runtimePolicyOverride": { "type": "object", "additionalProperties": {} @@ -12218,6 +13149,15 @@ "displayName": { "type": "string" }, + "invocationName": { + "type": "string" + }, + "legacyAliases": { + "type": "array", + "items": { + "type": "string" + } + }, "metadata": { "type": "object", "additionalProperties": {} @@ -12584,6 +13524,13 @@ "type": "object", "additionalProperties": {} }, + "rateLimitPolicyMode": { + "type": "string", + "enum": [ + "inherit", + "override" + ] + }, "retryPolicy": { "type": "object", "additionalProperties": {} @@ -13529,7 +14476,14 @@ "id": { "type": "string" }, + "legacyAliases": { + "type": "array", + "items": { + "type": "string" + } + }, "modelAlias": { + "description": "Deprecated: use ModelName. This field mirrors the canonical invocation name during the compatibility window.", "type": "string" }, "modelName": { @@ -13567,6 +14521,13 @@ "type": "object", "additionalProperties": {} }, + "rateLimitPolicyMode": { + "type": "string", + "enum": [ + "inherit", + "override" + ] + }, "retryPolicy": { "type": "object", "additionalProperties": {} diff --git a/apps/api/docs/swagger.yaml b/apps/api/docs/swagger.yaml index f48e4d3..514e8cb 100644 --- a/apps/api/docs/swagger.yaml +++ b/apps/api/docs/swagger.yaml @@ -57,6 +57,22 @@ definitions: $ref: '#/definitions/store.AccessRule' type: array type: object + httpapi.AdminTaskListResponse: + properties: + items: + items: + $ref: '#/definitions/store.AdminGatewayTask' + type: array + page: + example: 1 + type: integer + pageSize: + example: 50 + type: integer + total: + example: 42 + type: integer + type: object httpapi.AuditLogListResponse: properties: items: @@ -346,6 +362,103 @@ definitions: $ref: '#/definitions/store.DailyTokenUsage' type: array type: object + httpapi.EasyAIGeneratedResponse: + properties: + cancellable: + type: boolean + code: + type: string + created: + example: 1784772000000 + type: integer + data: + items: + $ref: '#/definitions/httpapi.EasyAIMediaOutput' + type: array + message: + type: string + output: + items: + type: string + type: array + output_content: + items: + $ref: '#/definitions/httpapi.EasyAIMediaOutput' + type: array + query_url: + example: /api/v1/ai/result/9f4d8f3d-5f5f-4bb7-a4be-344a9f930e25 + type: string + result: + additionalProperties: true + type: object + status: + enum: + - submitted + - process + - success + - failed + example: success + type: string + submitted: + type: boolean + task_id: + example: 9f4d8f3d-5f5f-4bb7-a4be-344a9f930e25 + type: string + taskId: + example: 9f4d8f3d-5f5f-4bb7-a4be-344a9f930e25 + type: string + upstream_task_id: + example: provider-task-123 + type: string + usage: + additionalProperties: true + type: object + type: object + httpapi.EasyAIMediaOutput: + properties: + audio_url: + type: string + b64_json: + description: B64JSON is retained when result media transfer is disabled. + type: string + content: + type: string + duration: + type: number + format: + type: string + height: + type: integer + image_url: + type: string + metadata: + additionalProperties: true + type: object + mime_type: + type: string + source_resolution: + type: string + target_resolution: + type: string + text: + type: string + type: + enum: + - image + - video + - audio + - file + - text + example: video + type: string + url: + example: https://cdn.example.com/output.mp4 + type: string + video_url: + type: string + width: + type: integer + type: object httpapi.ErrorEnvelope: properties: error: @@ -374,7 +487,7 @@ definitions: $ref: '#/definitions/store.FileStorageChannel' type: array type: object - httpapi.FileUploadResponse: + httpapi.FileUploadData: properties: assetStorage: additionalProperties: true @@ -395,6 +508,32 @@ definitions: example: /static/uploaded/upload-abc123.png type: string type: object + httpapi.FileUploadResponse: + properties: + assetStorage: + additionalProperties: true + type: object + contentType: + example: image/png + type: string + data: + $ref: '#/definitions/httpapi.FileUploadData' + filename: + example: image.png + type: string + id: + example: file_abc123 + type: string + size: + example: 1024 + type: integer + status: + example: success + type: string + url: + example: /static/uploaded/upload-abc123.png + type: string + type: object httpapi.GeminiErrorEnvelope: properties: error: @@ -748,6 +887,10 @@ definitions: type: boolean id: type: string + legacyAliases: + items: + type: string + type: array modelAlias: type: string modelName: @@ -762,6 +905,8 @@ definitions: type: string providerKey: type: string + providerModelName: + type: string providerName: type: string rateLimits: @@ -820,6 +965,26 @@ definitions: example: invalid_request_error type: string type: object + httpapi.OpenAIImageData: + properties: + b64_json: + type: string + revised_prompt: + type: string + url: + example: https://cdn.example.com/output.png + type: string + type: object + httpapi.OpenAIImagesResponse: + properties: + created: + example: 1710000000 + type: integer + data: + items: + $ref: '#/definitions/httpapi.OpenAIImageData' + type: array + type: object httpapi.PlatformListResponse: properties: items: @@ -1086,8 +1251,17 @@ definitions: properties: next: $ref: '#/definitions/httpapi.TaskNextLinks' + query_url: + example: /api/v1/ai/result/9f4d8f3d-5f5f-4bb7-a4be-344a9f930e25 + type: string + status: + example: submitted + type: string task: $ref: '#/definitions/store.GatewayTask' + task_id: + example: 9f4d8f3d-5f5f-4bb7-a4be-344a9f930e25 + type: string taskId: example: 9f4d8f3d-5f5f-4bb7-a4be-344a9f930e25 type: string @@ -1134,6 +1308,9 @@ definitions: events: example: /api/v1/tasks/9f4d8f3d-5f5f-4bb7-a4be-344a9f930e25/events type: string + result: + example: /api/v1/ai/result/9f4d8f3d-5f5f-4bb7-a4be-344a9f930e25 + type: string type: object httpapi.TaskParamPreprocessingLogListResponse: properties: @@ -1898,6 +2075,186 @@ definitions: status: type: string type: object + store.AdminGatewayTask: + properties: + adminContext: + $ref: '#/definitions/store.AdminTaskContext' + apiKeyId: + type: string + apiKeyName: + type: string + apiKeyPrefix: + type: string + asyncMode: + type: boolean + attemptCount: + type: integer + attempts: + items: + $ref: '#/definitions/store.TaskAttempt' + type: array + billingCurrency: + type: string + billingSettledAt: + type: string + billingStatus: + type: string + billingSummary: + additionalProperties: {} + type: object + billingUpdatedAt: + type: string + billingVersion: + type: string + billings: + items: {} + type: array + cancellable: + type: boolean + compatibilityProtocol: + type: string + compatibilityPublicId: + type: string + compatibilitySourceProtocol: + type: string + compatibilitySubmitBody: + additionalProperties: {} + type: object + compatibilitySubmitHeaders: + additionalProperties: {} + type: object + compatibilitySubmitHttpStatus: + type: integer + conversationId: + type: string + createdAt: + type: string + error: + type: string + errorCode: + type: string + errorMessage: + type: string + executionLeaseExpiresAt: + type: string + externalTaskId: + type: string + finalChargeAmount: + type: number + finishedAt: + type: string + gatewayTenantId: + type: string + gatewayUserId: + type: string + id: + type: string + kind: + type: string + message: + type: string + metrics: + additionalProperties: {} + type: object + model: + type: string + modelType: + type: string + newMessageCount: + type: integer + pricingSnapshot: + additionalProperties: {} + type: object + remoteTaskId: + type: string + request: + additionalProperties: {} + type: object + requestFingerprint: + type: string + requestId: + type: string + requestedModel: + type: string + reservationAmount: + type: number + resolvedModel: + type: string + responseDurationMs: + type: integer + responseFinishedAt: + type: string + responseStartedAt: + type: string + result: + additionalProperties: {} + type: object + riverJobId: + type: integer + runMode: + type: string + status: + type: string + submitted: + type: boolean + tenantId: + type: string + tenantKey: + type: string + updatedAt: + type: string + usage: + additionalProperties: {} + type: object + userGroupId: + type: string + userGroupKey: + type: string + userId: + type: string + userSource: + type: string + type: object + store.AdminTaskContext: + properties: + latestPlatform: + $ref: '#/definitions/store.AdminTaskPlatformSummary' + tenant: + $ref: '#/definitions/store.AdminTaskTenantSummary' + user: + $ref: '#/definitions/store.AdminTaskUserSummary' + type: object + store.AdminTaskPlatformSummary: + properties: + id: + type: string + name: + type: string + provider: + type: string + type: object + store.AdminTaskTenantSummary: + properties: + id: + type: string + name: + type: string + tenantKey: + type: string + type: object + store.AdminTaskUserSummary: + properties: + displayName: + type: string + email: + type: string + id: + type: string + source: + type: string + username: + type: string + type: object store.AuditLog: properties: action: @@ -1964,12 +2321,21 @@ definitions: defaultSnapshot: additionalProperties: {} type: object + displayName: + type: string id: type: string + invocationName: + type: string + legacyAliases: + items: + type: string + type: array metadata: additionalProperties: {} type: object modelAlias: + description: 'Deprecated: use InvocationName.' type: string modelType: items: @@ -1983,6 +2349,8 @@ definitions: type: string providerModelName: type: string + referenceCount: + type: integer runtimePolicyOverride: additionalProperties: {} type: object @@ -2013,6 +2381,12 @@ definitions: type: object displayName: type: string + invocationName: + type: string + legacyAliases: + items: + type: string + type: array metadata: additionalProperties: {} type: object @@ -2261,6 +2635,11 @@ definitions: rateLimitPolicy: additionalProperties: {} type: object + rateLimitPolicyMode: + enum: + - inherit + - override + type: string retryPolicy: additionalProperties: {} type: object @@ -2898,7 +3277,13 @@ definitions: type: boolean id: type: string + legacyAliases: + items: + type: string + type: array modelAlias: + description: 'Deprecated: use ModelName. This field mirrors the canonical + invocation name during the compatibility window.' type: string modelName: type: string @@ -2924,6 +3309,11 @@ definitions: rateLimitPolicy: additionalProperties: {} type: object + rateLimitPolicyMode: + enum: + - inherit + - override + type: string retryPolicy: additionalProperties: {} type: object @@ -3931,6 +4321,10 @@ paths: description: Not Found schema: $ref: '#/definitions/httpapi.ErrorEnvelope' + "409": + description: Conflict + schema: + $ref: '#/definitions/httpapi.ErrorEnvelope' "500": description: Internal Server Error schema: @@ -6003,6 +6397,198 @@ paths: summary: 验证统一认证 Revision tags: - identity + /api/admin/tasks: + get: + description: 跨租户分页查询任务；请求、结果和执行快照中的常见敏感字段默认脱敏。 + parameters: + - description: 关键词，匹配任务、用户、租户、模型、平台和 API Key + in: query + name: q + type: string + - description: 网关租户 ID + in: query + name: tenantId + type: string + - description: 网关用户 ID + in: query + name: userId + type: string + - description: 用户组 ID + in: query + name: userGroupId + type: string + - description: 任务状态 + in: query + name: status + type: string + - description: 执行平台 ID + in: query + name: platformId + type: string + - description: 调用或实际模型名称 + in: query + name: model + type: string + - description: 模型类型 + in: query + name: modelType + type: string + - description: 运行模式 + in: query + name: runMode + type: string + - description: 计费状态 + in: query + name: billingStatus + type: string + - description: API Key ID、名称或前缀 + in: query + name: apiKey + type: string + - description: 创建时间起点，支持 RFC3339 或日期格式 + in: query + name: createdFrom + type: string + - description: 创建时间终点，支持 RFC3339 或日期格式 + in: query + name: createdTo + type: string + - default: 1 + description: 页码 + in: query + name: page + type: integer + - default: 50 + description: 每页数量 + in: query + name: pageSize + type: integer + produces: + - application/json + responses: + "200": + description: OK + schema: + $ref: '#/definitions/httpapi.AdminTaskListResponse' + "400": + description: Bad Request + schema: + $ref: '#/definitions/httpapi.ErrorEnvelope' + "401": + description: Unauthorized + schema: + $ref: '#/definitions/httpapi.ErrorEnvelope' + "403": + description: Forbidden + schema: + $ref: '#/definitions/httpapi.ErrorEnvelope' + "500": + description: Internal Server Error + schema: + $ref: '#/definitions/httpapi.ErrorEnvelope' + security: + - BearerAuth: [] + summary: 管理员列出任务 + tags: + - admin-tasks + /api/admin/tasks/{taskID}: + get: + description: 返回任意租户的任务详情；默认脱敏，sensitive=full 时返回完整存储内容。 + parameters: + - description: 任务 ID + in: path + name: taskID + required: true + type: string + - default: masked + description: 敏感字段模式 + enum: + - masked + - full + in: query + name: sensitive + type: string + produces: + - application/json + responses: + "200": + description: OK + schema: + $ref: '#/definitions/store.AdminGatewayTask' + "400": + description: Bad Request + schema: + $ref: '#/definitions/httpapi.ErrorEnvelope' + "401": + description: Unauthorized + schema: + $ref: '#/definitions/httpapi.ErrorEnvelope' + "403": + description: Forbidden + schema: + $ref: '#/definitions/httpapi.ErrorEnvelope' + "404": + description: Not Found + schema: + $ref: '#/definitions/httpapi.ErrorEnvelope' + "500": + description: Internal Server Error + schema: + $ref: '#/definitions/httpapi.ErrorEnvelope' + security: + - BearerAuth: [] + summary: 管理员获取任务详情 + tags: + - admin-tasks + /api/admin/tasks/{taskID}/param-preprocessing: + get: + description: 返回任意租户任务的参数改写、校验或模板处理日志；默认脱敏。 + parameters: + - description: 任务 ID + in: path + name: taskID + required: true + type: string + - default: masked + description: 敏感字段模式 + enum: + - masked + - full + in: query + name: sensitive + type: string + produces: + - application/json + responses: + "200": + description: OK + schema: + $ref: '#/definitions/httpapi.TaskParamPreprocessingLogListResponse' + "400": + description: Bad Request + schema: + $ref: '#/definitions/httpapi.ErrorEnvelope' + "401": + description: Unauthorized + schema: + $ref: '#/definitions/httpapi.ErrorEnvelope' + "403": + description: Forbidden + schema: + $ref: '#/definitions/httpapi.ErrorEnvelope' + "404": + description: Not Found + schema: + $ref: '#/definitions/httpapi.ErrorEnvelope' + "500": + description: Internal Server Error + schema: + $ref: '#/definitions/httpapi.ErrorEnvelope' + security: + - BearerAuth: [] + summary: 管理员获取任务参数预处理日志 + tags: + - admin-tasks /api/admin/tenants: get: description: 管理端返回网关租户列表。 @@ -6605,6 +7191,7 @@ paths: - playground /api/v1/ai/result/{taskID}: get: + description: 查询当前用户的图片、视频、音频、视频高清和矢量化任务，并返回 EasyAI GeneratedResponse 兼容结构。 parameters: - description: 任务 ID in: path @@ -6617,13 +7204,20 @@ paths: "200": description: OK schema: - additionalProperties: true - type: object + $ref: '#/definitions/httpapi.EasyAIGeneratedResponse' + "404": + description: Not Found + schema: + $ref: '#/definitions/httpapi.EasyAIGeneratedResponse' + "410": + description: Gone + schema: + $ref: '#/definitions/httpapi.EasyAIGeneratedResponse' security: - BearerAuth: [] - summary: 查询 server-main 兼容视频结果 + summary: 查询 server-main EasyAIClient 兼容任务结果 tags: - - volces-compatible + - tasks /api/v1/api-keys: get: description: 返回当前用户创建的 API Key 元数据，secret 只在创建时返回。 @@ -7227,6 +7821,10 @@ paths: description: Not Found schema: $ref: '#/definitions/httpapi.VolcesErrorEnvelope' + "410": + description: Gone + schema: + $ref: '#/definitions/httpapi.VolcesErrorEnvelope' security: - BearerAuth: [] summary: 查询火山内容生成任务 @@ -7418,7 +8016,7 @@ paths: "200": description: OK schema: - $ref: '#/definitions/httpapi.CompatibleResponse' + $ref: '#/definitions/httpapi.OpenAIImagesResponse' "202": description: Accepted schema: @@ -7466,7 +8064,7 @@ paths: "200": description: OK schema: - $ref: '#/definitions/httpapi.CompatibleResponse' + $ref: '#/definitions/httpapi.OpenAIImagesResponse' "202": description: Accepted schema: @@ -7515,7 +8113,7 @@ paths: "200": description: OK schema: - $ref: '#/definitions/httpapi.CompatibleResponse' + $ref: '#/definitions/httpapi.EasyAIGeneratedResponse' "202": description: Accepted schema: @@ -7627,6 +8225,11 @@ paths: schema: additionalProperties: true type: object + "410": + description: Gone + schema: + additionalProperties: true + type: object security: - BearerAuth: [] summary: 查询可灵 V1 Omni 视频任务 @@ -7691,6 +8294,11 @@ paths: schema: additionalProperties: true type: object + "410": + description: Gone + schema: + additionalProperties: true + type: object security: - BearerAuth: [] summary: 按 ID 查询可灵 API 2.0 任务 @@ -7773,6 +8381,14 @@ paths: /api/v1/models: get: description: 按当前用户权限返回可用于 Playground 或 API 调用的模型列表。 + parameters: + - description: 模型可选场景；不传时保持原有行为 + enum: + - canvas_model_node + - desktop + in: query + name: usage_scene + type: string produces: - application/json responses: @@ -7780,6 +8396,10 @@ paths: description: OK schema: $ref: '#/definitions/httpapi.PlatformModelListResponse' + "400": + description: Bad Request + schema: + $ref: '#/definitions/httpapi.ErrorEnvelope' "401": description: Unauthorized schema: @@ -7788,6 +8408,10 @@ paths: description: Internal Server Error schema: $ref: '#/definitions/httpapi.ErrorEnvelope' + "502": + description: Bad Gateway + schema: + $ref: '#/definitions/httpapi.ErrorEnvelope' security: - BearerAuth: [] summary: 列出可调用模型 @@ -7893,7 +8517,8 @@ paths: post: consumes: - application/json - description: 统一公开入口按 model 选择平台模型并默认同步返回兼容响应；设置 X-Async=true 时异步创建任务并返回 202。 + description: 默认同步返回 EasyAI GeneratedResponse；设置 X-Async=true 时返回同时兼容 Gateway + 与 server-main EasyAIClient 的异步提交结构。 parameters: - description: true 时异步创建任务并返回 202 in: header @@ -7903,7 +8528,7 @@ paths: in: header name: Idempotency-Key type: string - - description: AI 任务请求，字段随任务类型变化 + - description: 媒体任务请求，字段随任务类型变化 in: body name: input required: true @@ -7915,7 +8540,7 @@ paths: "200": description: OK schema: - $ref: '#/definitions/httpapi.CompatibleResponse' + $ref: '#/definitions/httpapi.EasyAIGeneratedResponse' "202": description: Accepted schema: @@ -7950,9 +8575,9 @@ paths: $ref: '#/definitions/httpapi.ErrorEnvelope' security: - BearerAuth: [] - summary: 创建或执行 AI 任务 + summary: 创建 EasyAI 媒体任务 tags: - - tasks + - media /api/v1/openapi.json: get: description: 返回当前构建内嵌的完整机器可读 Swagger JSON，供 Agent 在 SKILL references 未覆盖接口时查询。 @@ -8006,6 +8631,14 @@ paths: /api/v1/playground/models: get: description: 按当前用户权限返回可用于 Playground 或 API 调用的模型列表。 + parameters: + - description: 模型可选场景；不传时保持原有行为 + enum: + - canvas_model_node + - desktop + in: query + name: usage_scene + type: string produces: - application/json responses: @@ -8013,6 +8646,10 @@ paths: description: OK schema: $ref: '#/definitions/httpapi.PlatformModelListResponse' + "400": + description: Bad Request + schema: + $ref: '#/definitions/httpapi.ErrorEnvelope' "401": description: Unauthorized schema: @@ -8021,6 +8658,10 @@ paths: description: Internal Server Error schema: $ref: '#/definitions/httpapi.ErrorEnvelope' + "502": + description: Bad Gateway + schema: + $ref: '#/definitions/httpapi.ErrorEnvelope' security: - BearerAuth: [] summary: 列出可调用模型 @@ -8437,7 +9078,8 @@ paths: post: consumes: - application/json - description: 统一公开入口按 model 选择平台模型并默认同步返回兼容响应；设置 X-Async=true 时异步创建任务并返回 202。 + description: 默认同步返回 EasyAI GeneratedResponse；设置 X-Async=true 时返回同时兼容 Gateway + 与 server-main EasyAIClient 的异步提交结构。 parameters: - description: true 时异步创建任务并返回 202 in: header @@ -8447,7 +9089,7 @@ paths: in: header name: Idempotency-Key type: string - - description: AI 任务请求，字段随任务类型变化 + - description: 媒体任务请求，字段随任务类型变化 in: body name: input required: true @@ -8459,7 +9101,7 @@ paths: "200": description: OK schema: - $ref: '#/definitions/httpapi.CompatibleResponse' + $ref: '#/definitions/httpapi.EasyAIGeneratedResponse' "202": description: Accepted schema: @@ -8494,14 +9136,15 @@ paths: $ref: '#/definitions/httpapi.ErrorEnvelope' security: - BearerAuth: [] - summary: 创建或执行 AI 任务 + summary: 创建 EasyAI 媒体任务 tags: - - tasks + - media /api/v1/speech/generations: post: consumes: - application/json - description: 统一公开入口按 model 选择平台模型并默认同步返回兼容响应；设置 X-Async=true 时异步创建任务并返回 202。 + description: 默认同步返回 EasyAI GeneratedResponse；设置 X-Async=true 时返回同时兼容 Gateway + 与 server-main EasyAIClient 的异步提交结构。 parameters: - description: true 时异步创建任务并返回 202 in: header @@ -8511,7 +9154,7 @@ paths: in: header name: Idempotency-Key type: string - - description: AI 任务请求，字段随任务类型变化 + - description: 媒体任务请求，字段随任务类型变化 in: body name: input required: true @@ -8523,7 +9166,7 @@ paths: "200": description: OK schema: - $ref: '#/definitions/httpapi.CompatibleResponse' + $ref: '#/definitions/httpapi.EasyAIGeneratedResponse' "202": description: Accepted schema: @@ -8558,9 +9201,9 @@ paths: $ref: '#/definitions/httpapi.ErrorEnvelope' security: - BearerAuth: [] - summary: 创建或执行 AI 任务 + summary: 创建 EasyAI 媒体任务 tags: - - tasks + - media /api/v1/tasks: get: description: 按当前用户列出任务，支持关键字、模型类型、时间范围和分页过滤。 @@ -8647,6 +9290,10 @@ paths: description: Not Found schema: $ref: '#/definitions/httpapi.ErrorEnvelope' + "410": + description: Gone + schema: + $ref: '#/definitions/httpapi.ErrorEnvelope' "500": description: Internal Server Error schema: @@ -8770,18 +9417,18 @@ paths: "200": description: OK schema: - additionalProperties: true - type: object + $ref: '#/definitions/httpapi.TaskAcceptedResponse' security: - BearerAuth: [] summary: 创建 server-main 兼容视频任务 tags: - - volces-compatible + - media /api/v1/videos/generations: post: consumes: - application/json - description: 统一公开入口按 model 选择平台模型并默认同步返回兼容响应；设置 X-Async=true 时异步创建任务并返回 202。 + description: 默认同步返回 EasyAI GeneratedResponse；设置 X-Async=true 时返回同时兼容 Gateway + 与 server-main EasyAIClient 的异步提交结构。 parameters: - description: true 时异步创建任务并返回 202 in: header @@ -8791,7 +9438,7 @@ paths: in: header name: Idempotency-Key type: string - - description: AI 任务请求，字段随任务类型变化 + - description: 媒体任务请求，字段随任务类型变化 in: body name: input required: true @@ -8803,7 +9450,7 @@ paths: "200": description: OK schema: - $ref: '#/definitions/httpapi.CompatibleResponse' + $ref: '#/definitions/httpapi.EasyAIGeneratedResponse' "202": description: Accepted schema: @@ -8838,9 +9485,9 @@ paths: $ref: '#/definitions/httpapi.ErrorEnvelope' security: - BearerAuth: [] - summary: 创建或执行 AI 任务 + summary: 创建 EasyAI 媒体任务 tags: - - tasks + - media /api/v1/videos/omni-video: post: consumes: @@ -8918,6 +9565,10 @@ paths: description: Not Found schema: $ref: '#/definitions/httpapi.KelingCompatibleEnvelope' + "410": + description: Gone + schema: + $ref: '#/definitions/httpapi.KelingCompatibleEnvelope' "500": description: Internal Server Error schema: @@ -8950,7 +9601,7 @@ paths: "200": description: OK schema: - $ref: '#/definitions/httpapi.CompatibleResponse' + $ref: '#/definitions/httpapi.EasyAIGeneratedResponse' "202": description: Accepted schema: @@ -8984,7 +9635,8 @@ paths: post: consumes: - application/json - description: 统一公开入口按 model 选择平台模型并默认同步返回兼容响应；设置 X-Async=true 时异步创建任务并返回 202。 + description: 默认同步返回 EasyAI GeneratedResponse；设置 X-Async=true 时返回同时兼容 Gateway + 与 server-main EasyAIClient 的异步提交结构。 parameters: - description: true 时异步创建任务并返回 202 in: header @@ -8994,7 +9646,7 @@ paths: in: header name: Idempotency-Key type: string - - description: AI 任务请求，字段随任务类型变化 + - description: 媒体任务请求，字段随任务类型变化 in: body name: input required: true @@ -9006,7 +9658,7 @@ paths: "200": description: OK schema: - $ref: '#/definitions/httpapi.CompatibleResponse' + $ref: '#/definitions/httpapi.EasyAIGeneratedResponse' "202": description: Accepted schema: @@ -9041,9 +9693,9 @@ paths: $ref: '#/definitions/httpapi.ErrorEnvelope' security: - BearerAuth: [] - summary: 创建或执行 AI 任务 + summary: 创建 EasyAI 媒体任务 tags: - - tasks + - media /api/v1/voice_clone/voices: get: description: 返回当前用户在网关中维护的克隆音色，以及克隆时绑定的平台与平台模型。 @@ -9215,6 +9867,10 @@ paths: description: Not Found schema: $ref: '#/definitions/httpapi.ErrorEnvelope' + "410": + description: Gone + schema: + $ref: '#/definitions/httpapi.ErrorEnvelope' "500": description: Internal Server Error schema: diff --git a/apps/api/go.mod b/apps/api/go.mod index cd2eb53..f4338e5 100644 --- a/apps/api/go.mod +++ b/apps/api/go.mod @@ -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 ) diff --git a/apps/api/go.sum b/apps/api/go.sum index 2f9a8e9..2a02291 100644 --- a/apps/api/go.sum +++ b/apps/api/go.sum @@ -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= diff --git a/apps/api/internal/clients/chat_reasoning.go b/apps/api/internal/clients/chat_reasoning.go index 5bb21bc..bf8ff4d 100644 --- a/apps/api/internal/clients/chat_reasoning.go +++ b/apps/api/internal/clients/chat_reasoning.go @@ -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" diff --git a/apps/api/internal/clients/clients_test.go b/apps/api/internal/clients/clients_test.go index 9f80763..1b201f3 100644 --- a/apps/api/internal/clients/clients_test.go +++ b/apps/api/internal/clients/clients_test.go @@ -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) } } diff --git a/apps/api/internal/clients/gemini.go b/apps/api/internal/clients/gemini.go index c1dd487..b3a0e25 100644 --- a/apps/api/internal/clients/gemini.go +++ b/apps/api/internal/clients/gemini.go @@ -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, } } diff --git a/apps/api/internal/clients/helpers.go b/apps/api/internal/clients/helpers.go index 00e2255..c7dfd12 100644 --- a/apps/api/internal/clients/helpers.go +++ b/apps/api/internal/clients/helpers.go @@ -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 diff --git a/apps/api/internal/clients/keling.go b/apps/api/internal/clients/keling.go index 69ec78d..a353162 100644 --- a/apps/api/internal/clients/keling.go +++ b/apps/api/internal/clients/keling.go @@ -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, } } diff --git a/apps/api/internal/clients/media_clients.go b/apps/api/internal/clients/media_clients.go index 8b7c297..148935c 100644 --- a/apps/api/internal/clients/media_clients.go +++ b/apps/api/internal/clients/media_clients.go @@ -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}} diff --git a/apps/api/internal/clients/media_parameter_normalization_test.go b/apps/api/internal/clients/media_parameter_normalization_test.go new file mode 100644 index 0000000..68d6887 --- /dev/null +++ b/apps/api/internal/clients/media_parameter_normalization_test.go @@ -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) + } + } +} diff --git a/apps/api/internal/clients/media_parameters.go b/apps/api/internal/clients/media_parameters.go new file mode 100644 index 0000000..39f9a3f --- /dev/null +++ b/apps/api/internal/clients/media_parameters.go @@ -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) +} diff --git a/apps/api/internal/clients/openai.go b/apps/api/internal/clients/openai.go index 5e4404e..2f96131 100644 --- a/apps/api/internal/clients/openai.go +++ b/apps/api/internal/clients/openai.go @@ -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 diff --git a/apps/api/internal/clients/openai_chat_audio.go b/apps/api/internal/clients/openai_chat_audio.go new file mode 100644 index 0000000..c0e9d8b --- /dev/null +++ b/apps/api/internal/clients/openai_chat_audio.go @@ -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, + } +} diff --git a/apps/api/internal/clients/provider_task.go b/apps/api/internal/clients/provider_task.go index 2fb8c71..7c8aed8 100644 --- a/apps/api/internal/clients/provider_task.go +++ b/apps/api/internal/clients/provider_task.go @@ -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 } diff --git a/apps/api/internal/clients/types.go b/apps/api/internal/clients/types.go index 599deef..78fdbad 100644 --- a/apps/api/internal/clients/types.go +++ b/apps/api/internal/clients/types.go @@ -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 diff --git a/apps/api/internal/clients/volces.go b/apps/api/internal/clients/volces.go index fc727ac..4b770d8 100644 --- a/apps/api/internal/clients/volces.go +++ b/apps/api/internal/clients/volces.go @@ -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 } diff --git a/apps/api/internal/clients/volces_deyun_test.go b/apps/api/internal/clients/volces_deyun_test.go index 04485a1..788b2fa 100644 --- a/apps/api/internal/clients/volces_deyun_test.go +++ b/apps/api/internal/clients/volces_deyun_test.go @@ -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) } } diff --git a/apps/api/internal/clients/wire_response_test.go b/apps/api/internal/clients/wire_response_test.go index 685fa37..d78d4e7 100644 --- a/apps/api/internal/clients/wire_response_test.go +++ b/apps/api/internal/clients/wire_response_test.go @@ -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) { diff --git a/apps/api/internal/config/config.go b/apps/api/internal/config/config.go index 728aa62..89a7e39 100644 --- a/apps/api/internal/config/config.go +++ b/apps/api/internal/config/config.go @@ -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": diff --git a/apps/api/internal/config/config_test.go b/apps/api/internal/config/config_test.go index 1e3f7b9..16b6920 100644 --- a/apps/api/internal/config/config_test.go +++ b/apps/api/internal/config/config_test.go @@ -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) + } +} diff --git a/apps/api/internal/httpapi/admin_task_handlers.go b/apps/api/internal/httpapi/admin_task_handlers.go new file mode 100644 index 0000000..bc89687 --- /dev/null +++ b/apps/api/internal/httpapi/admin_task_handlers.go @@ -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 +} diff --git a/apps/api/internal/httpapi/admin_task_handlers_test.go b/apps/api/internal/httpapi/admin_task_handlers_test.go new file mode 100644 index 0000000..fb2b362 --- /dev/null +++ b/apps/api/internal/httpapi/admin_task_handlers_test.go @@ -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) + } + } +} diff --git a/apps/api/internal/httpapi/admin_tasks_integration_test.go b/apps/api/internal/httpapi/admin_tasks_integration_test.go new file mode 100644 index 0000000..61693ee --- /dev/null +++ b/apps/api/internal/httpapi/admin_tasks_integration_test.go @@ -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 +} diff --git a/apps/api/internal/httpapi/advanced_task_handlers.go b/apps/api/internal/httpapi/advanced_task_handlers.go index b4dc5db..d3e7c92 100644 --- a/apps/api/internal/httpapi/advanced_task_handlers.go +++ b/apps/api/internal/httpapi/advanced_task_handlers.go @@ -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 diff --git a/apps/api/internal/httpapi/async_worker_acceptance_integration_test.go b/apps/api/internal/httpapi/async_worker_acceptance_integration_test.go new file mode 100644 index 0000000..3210ab7 --- /dev/null +++ b/apps/api/internal/httpapi/async_worker_acceptance_integration_test.go @@ -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 +} diff --git a/apps/api/internal/httpapi/binary_results.go b/apps/api/internal/httpapi/binary_results.go new file mode 100644 index 0000000..2455c3d --- /dev/null +++ b/apps/api/internal/httpapi/binary_results.go @@ -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) +} diff --git a/apps/api/internal/httpapi/catalog_handlers.go b/apps/api/internal/httpapi/catalog_handlers.go index 3affb97..e31e2bf 100644 --- a/apps/api/internal/httpapi/catalog_handlers.go +++ b/apps/api/internal/httpapi/catalog_handlers.go @@ -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 diff --git a/apps/api/internal/httpapi/compat_protocol.go b/apps/api/internal/httpapi/compat_protocol.go index 603afbc..eb82700 100644 --- a/apps/api/internal/httpapi/compat_protocol.go +++ b/apps/api/internal/httpapi/compat_protocol.go @@ -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}) } diff --git a/apps/api/internal/httpapi/core_flow_integration_test.go b/apps/api/internal/httpapi/core_flow_integration_test.go index d9f4b0c..73e5800 100644 --- a/apps/api/internal/httpapi/core_flow_integration_test.go +++ b/apps/api/internal/httpapi/core_flow_integration_test.go @@ -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 } diff --git a/apps/api/internal/httpapi/easyai_compat.go b/apps/api/internal/httpapi/easyai_compat.go new file mode 100644 index 0000000..2a25a0e --- /dev/null +++ b/apps/api/internal/httpapi/easyai_compat.go @@ -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 +} diff --git a/apps/api/internal/httpapi/easyai_compat_test.go b/apps/api/internal/httpapi/easyai_compat_test.go new file mode 100644 index 0000000..5a64eee --- /dev/null +++ b/apps/api/internal/httpapi/easyai_compat_test.go @@ -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) + } +} diff --git a/apps/api/internal/httpapi/failure_policy_http_acceptance_integration_test.go b/apps/api/internal/httpapi/failure_policy_http_acceptance_integration_test.go new file mode 100644 index 0000000..c6bf2d7 --- /dev/null +++ b/apps/api/internal/httpapi/failure_policy_http_acceptance_integration_test.go @@ -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) +} diff --git a/apps/api/internal/httpapi/file_upload_handlers.go b/apps/api/internal/httpapi/file_upload_handlers.go index 07f49ef..0a02b14 100644 --- a/apps/api/internal/httpapi/file_upload_handlers.go +++ b/apps/api/internal/httpapi/file_upload_handlers.go @@ -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 { diff --git a/apps/api/internal/httpapi/gemini_compat.go b/apps/api/internal/httpapi/gemini_compat.go index 4fa8c59..98ce2b0 100644 --- a/apps/api/internal/httpapi/gemini_compat.go +++ b/apps/api/internal/httpapi/gemini_compat.go @@ -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) diff --git a/apps/api/internal/httpapi/handlers.go b/apps/api/internal/httpapi/handlers.go index 2a69032..14fd893 100644 --- a/apps/api/internal/httpapi/handlers.go +++ b/apps/api/internal/httpapi/handlers.go @@ -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 } diff --git a/apps/api/internal/httpapi/integration_database_safety_test.go b/apps/api/internal/httpapi/integration_database_safety_test.go new file mode 100644 index 0000000..cf35662 --- /dev/null +++ b/apps/api/internal/httpapi/integration_database_safety_test.go @@ -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) + } + }) + } +} diff --git a/apps/api/internal/httpapi/keling_compat_handlers.go b/apps/api/internal/httpapi/keling_compat_handlers.go index cba14fd..002720c 100644 --- a/apps/api/internal/httpapi/keling_compat_handlers.go +++ b/apps/api/internal/httpapi/keling_compat_handlers.go @@ -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 { diff --git a/apps/api/internal/httpapi/keling_compat_handlers_test.go b/apps/api/internal/httpapi/keling_compat_handlers_test.go index 6d5b41c..c35d3c8 100644 --- a/apps/api/internal/httpapi/keling_compat_handlers_test.go +++ b/apps/api/internal/httpapi/keling_compat_handlers_test.go @@ -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) { diff --git a/apps/api/internal/httpapi/kling_compat.go b/apps/api/internal/httpapi/kling_compat.go index 39e0a5a..8507519 100644 --- a/apps/api/internal/httpapi/kling_compat.go +++ b/apps/api/internal/httpapi/kling_compat.go @@ -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 diff --git a/apps/api/internal/httpapi/kling_compat_test.go b/apps/api/internal/httpapi/kling_compat_test.go index 07766de..36f750f 100644 --- a/apps/api/internal/httpapi/kling_compat_test.go +++ b/apps/api/internal/httpapi/kling_compat_test.go @@ -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 diff --git a/apps/api/internal/httpapi/local_temp_assets.go b/apps/api/internal/httpapi/local_temp_assets.go index 1a31ce9..f0a66b1 100644 --- a/apps/api/internal/httpapi/local_temp_assets.go +++ b/apps/api/internal/httpapi/local_temp_assets.go @@ -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 } diff --git a/apps/api/internal/httpapi/model_catalog.go b/apps/api/internal/httpapi/model_catalog.go index 965326c..1a120fd 100644 --- a/apps/api/internal/httpapi/model_catalog.go +++ b/apps/api/internal/httpapi/model_catalog.go @@ -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 } diff --git a/apps/api/internal/httpapi/model_catalog_test.go b/apps/api/internal/httpapi/model_catalog_test.go index 22a7f29..ba8f98b 100644 --- a/apps/api/internal/httpapi/model_catalog_test.go +++ b/apps/api/internal/httpapi/model_catalog_test.go @@ -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{ { diff --git a/apps/api/internal/httpapi/model_cooldown_error_test.go b/apps/api/internal/httpapi/model_cooldown_error_test.go new file mode 100644 index 0000000..e0d72bd --- /dev/null +++ b/apps/api/internal/httpapi/model_cooldown_error_test.go @@ -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) + } +} diff --git a/apps/api/internal/httpapi/model_usage_scene.go b/apps/api/internal/httpapi/model_usage_scene.go new file mode 100644 index 0000000..07bbdbe --- /dev/null +++ b/apps/api/internal/httpapi/model_usage_scene.go @@ -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)) +} diff --git a/apps/api/internal/httpapi/model_usage_scene_test.go b/apps/api/internal/httpapi/model_usage_scene_test.go new file mode 100644 index 0000000..b3ba35b --- /dev/null +++ b/apps/api/internal/httpapi/model_usage_scene_test.go @@ -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) + } +} diff --git a/apps/api/internal/httpapi/openapi_models.go b/apps/api/internal/httpapi/openapi_models.go index 24701de..aa70b9b 100644 --- a/apps/api/internal/httpapi/openapi_models.go +++ b/apps/api/internal/httpapi/openapi_models.go @@ -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 { diff --git a/apps/api/internal/httpapi/request_preparation_test.go b/apps/api/internal/httpapi/request_preparation_test.go index ff15157..25fa05e 100644 --- a/apps/api/internal/httpapi/request_preparation_test.go +++ b/apps/api/internal/httpapi/request_preparation_test.go @@ -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") diff --git a/apps/api/internal/httpapi/runtime_policy_handlers.go b/apps/api/internal/httpapi/runtime_policy_handlers.go index 50f37bb..6f0c2e4 100644 --- a/apps/api/internal/httpapi/runtime_policy_handlers.go +++ b/apps/api/internal/httpapi/runtime_policy_handlers.go @@ -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"` diff --git a/apps/api/internal/httpapi/runtime_policy_validation_test.go b/apps/api/internal/httpapi/runtime_policy_validation_test.go new file mode 100644 index 0000000..3fd4a9f --- /dev/null +++ b/apps/api/internal/httpapi/runtime_policy_validation_test.go @@ -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) + } + }) + } +} diff --git a/apps/api/internal/httpapi/server.go b/apps/api/internal/httpapi/server.go index 13a0734..e0682ee 100644 --- a/apps/api/internal/httpapi/server.go +++ b/apps/api/internal/httpapi/server.go @@ -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))) diff --git a/apps/api/internal/httpapi/task_idempotency_test.go b/apps/api/internal/httpapi/task_idempotency_test.go index 4d49609..ddf8bcb 100644 --- a/apps/api/internal/httpapi/task_idempotency_test.go +++ b/apps/api/internal/httpapi/task_idempotency_test.go @@ -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) + } +} diff --git a/apps/api/internal/httpapi/voice_clone_handlers.go b/apps/api/internal/httpapi/voice_clone_handlers.go index 33a6ca1..7f570df 100644 --- a/apps/api/internal/httpapi/voice_clone_handlers.go +++ b/apps/api/internal/httpapi/voice_clone_handlers.go @@ -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 } diff --git a/apps/api/internal/httpapi/volces_compat_handlers.go b/apps/api/internal/httpapi/volces_compat_handlers.go index 96be443..236fae6 100644 --- a/apps/api/internal/httpapi/volces_compat_handlers.go +++ b/apps/api/internal/httpapi/volces_compat_handlers.go @@ -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) diff --git a/apps/api/internal/httpapi/volces_compat_handlers_test.go b/apps/api/internal/httpapi/volces_compat_handlers_test.go index 12e7725..0900aab 100644 --- a/apps/api/internal/httpapi/volces_compat_handlers_test.go +++ b/apps/api/internal/httpapi/volces_compat_handlers_test.go @@ -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) diff --git a/apps/api/internal/runner/binary_results.go b/apps/api/internal/runner/binary_results.go new file mode 100644 index 0000000..c83b48a --- /dev/null +++ b/apps/api/internal/runner/binary_results.go @@ -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 +} diff --git a/apps/api/internal/runner/binary_results_test.go b/apps/api/internal/runner/binary_results_test.go new file mode 100644 index 0000000..0b8d699 --- /dev/null +++ b/apps/api/internal/runner/binary_results_test.go @@ -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) + } +} diff --git a/apps/api/internal/runner/cache_affinity.go b/apps/api/internal/runner/cache_affinity.go index 90a2b6c..90a96b5 100644 --- a/apps/api/internal/runner/cache_affinity.go +++ b/apps/api/internal/runner/cache_affinity.go @@ -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 { diff --git a/apps/api/internal/runner/cache_affinity_test.go b/apps/api/internal/runner/cache_affinity_test.go index e00280b..7292f8f 100644 --- a/apps/api/internal/runner/cache_affinity_test.go +++ b/apps/api/internal/runner/cache_affinity_test.go @@ -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) + } +} diff --git a/apps/api/internal/runner/candidate_filter.go b/apps/api/internal/runner/candidate_filter.go index c8549f0..1c9b8ac 100644 --- a/apps/api/internal/runner/candidate_filter.go +++ b/apps/api/internal/runner/candidate_filter.go @@ -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)) { diff --git a/apps/api/internal/runner/candidate_platform_filter.go b/apps/api/internal/runner/candidate_platform_filter.go new file mode 100644 index 0000000..a67e73a --- /dev/null +++ b/apps/api/internal/runner/candidate_platform_filter.go @@ -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, + ) +} diff --git a/apps/api/internal/runner/candidate_platform_filter_test.go b/apps/api/internal/runner/candidate_platform_filter_test.go new file mode 100644 index 0000000..2c2405f --- /dev/null +++ b/apps/api/internal/runner/candidate_platform_filter_test.go @@ -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) + } + }) +} diff --git a/apps/api/internal/runner/input_image_normalizer.go b/apps/api/internal/runner/input_image_normalizer.go new file mode 100644 index 0000000..7d5db55 --- /dev/null +++ b/apps/api/internal/runner/input_image_normalizer.go @@ -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() +} diff --git a/apps/api/internal/runner/input_image_normalizer_test.go b/apps/api/internal/runner/input_image_normalizer_test.go new file mode 100644 index 0000000..275b318 --- /dev/null +++ b/apps/api/internal/runner/input_image_normalizer_test.go @@ -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) + } +} diff --git a/apps/api/internal/runner/limits.go b/apps/api/internal/runner/limits.go index c533c93..636b680 100644 --- a/apps/api/internal/runner/limits.go +++ b/apps/api/internal/runner/limits.go @@ -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 { diff --git a/apps/api/internal/runner/limits_test.go b/apps/api/internal/runner/limits_test.go index 31490b3..0169ad7 100644 --- a/apps/api/internal/runner/limits_test.go +++ b/apps/api/internal/runner/limits_test.go @@ -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) { diff --git a/apps/api/internal/runner/media_parameter_normalization_test.go b/apps/api/internal/runner/media_parameter_normalization_test.go new file mode 100644 index 0000000..8c8640f --- /dev/null +++ b/apps/api/internal/runner/media_parameter_normalization_test.go @@ -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" +} diff --git a/apps/api/internal/runner/output_token_limit.go b/apps/api/internal/runner/output_token_limit.go index 329e546..3825fab 100644 --- a/apps/api/internal/runner/output_token_limit.go +++ b/apps/api/internal/runner/output_token_limit.go @@ -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" } diff --git a/apps/api/internal/runner/output_token_limit_test.go b/apps/api/internal/runner/output_token_limit_test.go index 37262d4..b8fcd49 100644 --- a/apps/api/internal/runner/output_token_limit_test.go +++ b/apps/api/internal/runner/output_token_limit_test.go @@ -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" diff --git a/apps/api/internal/runner/param_processor_media.go b/apps/api/internal/runner/param_processor_media.go index bf0ffda..4b905ba 100644 --- a/apps/api/internal/runner/param_processor_media.go +++ b/apps/api/internal/runner/param_processor_media.go @@ -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 { diff --git a/apps/api/internal/runner/param_processor_utils.go b/apps/api/internal/runner/param_processor_utils.go index fe5f951..d271957 100644 --- a/apps/api/internal/runner/param_processor_utils.go +++ b/apps/api/internal/runner/param_processor_utils.go @@ -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 { diff --git a/apps/api/internal/runner/pricing.go b/apps/api/internal/runner/pricing.go index 724acec..6bb6ef3 100644 --- a/apps/api/internal/runner/pricing.go +++ b/apps/api/internal/runner/pricing.go @@ -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 diff --git a/apps/api/internal/runner/proxy.go b/apps/api/internal/runner/proxy.go index f650ada..9f74f10 100644 --- a/apps/api/internal/runner/proxy.go +++ b/apps/api/internal/runner/proxy.go @@ -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, } } diff --git a/apps/api/internal/runner/proxy_test.go b/apps/api/internal/runner/proxy_test.go index ebb8aca..cd6c385 100644 --- a/apps/api/internal/runner/proxy_test.go +++ b/apps/api/internal/runner/proxy_test.go @@ -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) { diff --git a/apps/api/internal/runner/queue_worker.go b/apps/api/internal/runner/queue_worker.go index 6809d7e..72e0b9b 100644 --- a/apps/api/internal/runner/queue_worker.go +++ b/apps/api/internal/runner/queue_worker.go @@ -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 } diff --git a/apps/api/internal/runner/queue_worker_manager_test.go b/apps/api/internal/runner/queue_worker_manager_test.go new file mode 100644 index 0000000..b01e993 --- /dev/null +++ b/apps/api/internal/runner/queue_worker_manager_test.go @@ -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) +} diff --git a/apps/api/internal/runner/recording.go b/apps/api/internal/runner/recording.go index 1022b5c..0754b0d 100644 --- a/apps/api/internal/runner/recording.go +++ b/apps/api/internal/runner/recording.go @@ -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 diff --git a/apps/api/internal/runner/recording_test.go b/apps/api/internal/runner/recording_test.go index a024a52..1fd6f0d 100644 --- a/apps/api/internal/runner/recording_test.go +++ b/apps/api/internal/runner/recording_test.go @@ -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) + } } diff --git a/apps/api/internal/runner/request_assets.go b/apps/api/internal/runner/request_assets.go index b3efcdc..85670e8 100644 --- a/apps/api/internal/runner/request_assets.go +++ b/apps/api/internal/runner/request_assets.go @@ -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) diff --git a/apps/api/internal/runner/request_assets_test.go b/apps/api/internal/runner/request_assets_test.go index 2a8659f..1d05812 100644 --- a/apps/api/internal/runner/request_assets_test.go +++ b/apps/api/internal/runner/request_assets_test.go @@ -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{ diff --git a/apps/api/internal/runner/retry_decision.go b/apps/api/internal/runner/retry_decision.go index 30f47ee..c854dd7 100644 --- a/apps/api/internal/runner/retry_decision.go +++ b/apps/api/internal/runner/retry_decision.go @@ -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) { diff --git a/apps/api/internal/runner/retry_decision_test.go b/apps/api/internal/runner/retry_decision_test.go index aebc1f1..0788a02 100644 --- a/apps/api/internal/runner/retry_decision_test.go +++ b/apps/api/internal/runner/retry_decision_test.go @@ -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) + } +} diff --git a/apps/api/internal/runner/runtime_policy.go b/apps/api/internal/runner/runtime_policy.go index 131a90e..41c41b0 100644 --- a/apps/api/internal/runner/runtime_policy.go +++ b/apps/api/internal/runner/runtime_policy.go @@ -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 { diff --git a/apps/api/internal/runner/service.go b/apps/api/internal/runner/service.go index 1a6d975..cfdd3c1 100644 --- a/apps/api/internal/runner/service.go +++ b/apps/api/internal/runner/service.go @@ -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) } diff --git a/apps/api/internal/runner/task_cancel.go b/apps/api/internal/runner/task_cancel.go index 29e9491..5a304dd 100644 --- a/apps/api/internal/runner/task_cancel.go +++ b/apps/api/internal/runner/task_cancel.go @@ -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 } diff --git a/apps/api/internal/runner/task_history_workers.go b/apps/api/internal/runner/task_history_workers.go new file mode 100644 index 0000000..78f62dd --- /dev/null +++ b/apps/api/internal/runner/task_history_workers.go @@ -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): + } + } +} diff --git a/apps/api/internal/runner/task_history_workers_test.go b/apps/api/internal/runner/task_history_workers_test.go new file mode 100644 index 0000000..795a645 --- /dev/null +++ b/apps/api/internal/runner/task_history_workers_test.go @@ -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) + } +} diff --git a/apps/api/internal/runner/trace.go b/apps/api/internal/runner/trace.go index d585bcf..7a8fa40 100644 --- a/apps/api/internal/runner/trace.go +++ b/apps/api/internal/runner/trace.go @@ -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 diff --git a/apps/api/internal/runner/upload.go b/apps/api/internal/runner/upload.go index a3fc312..526748d 100644 --- a/apps/api/internal/runner/upload.go +++ b/apps/api/internal/runner/upload.go @@ -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) diff --git a/apps/api/internal/runner/upload_test.go b/apps/api/internal/runner/upload_test.go index 97ee6fc..ee97fd4 100644 --- a/apps/api/internal/runner/upload_test.go +++ b/apps/api/internal/runner/upload_test.go @@ -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) } diff --git a/apps/api/internal/runner/voice_clone.go b/apps/api/internal/runner/voice_clone.go index 97b594c..cb59c96 100644 --- a/apps/api/internal/runner/voice_clone.go +++ b/apps/api/internal/runner/voice_clone.go @@ -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{}, }) } diff --git a/apps/api/internal/securityevents/metrics.go b/apps/api/internal/securityevents/metrics.go index 6bc24b3..556705f 100644 --- a/apps/api/internal/securityevents/metrics.go +++ b/apps/api/internal/securityevents/metrics.go @@ -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) +} diff --git a/apps/api/internal/securityevents/metrics_test.go b/apps/api/internal/securityevents/metrics_test.go index 1ba6d30..ff76210 100644 --- a/apps/api/internal/securityevents/metrics_test.go +++ b/apps/api/internal/securityevents/metrics_test.go @@ -26,6 +26,12 @@ func TestMetricsExposeBoundedOutcomesAndState(t *testing.T) { metrics.ObserveHeartbeat("accepted") metrics.ObserveIntrospection("failed") metrics.ObserveJWKSRefreshFailure("ssf") + metrics.SetAsyncWorkerCapacity(96, 128, 96, true) + metrics.ObserveAsyncWorkerResize("success") + metrics.ObserveConcurrencyLeaseRenewal("success") + metrics.ObserveConcurrencyLeaseRenewal("lost") + metrics.ObserveTaskEventSkip("duplicate") + metrics.ObserveTaskEventSkip("budget_exceeded") recorder := httptest.NewRecorder() metrics.Handler(metricsSnapshot{mode: "push_healthy", last: time.Now().Add(-time.Minute)}, "issuer", "audience", true). @@ -42,6 +48,14 @@ func TestMetricsExposeBoundedOutcomesAndState(t *testing.T) { `easyai_gateway_ssf_mode{mode="push_healthy"} 1`, `easyai_gateway_oidc_introspection_total{outcome="failed"} 1`, `easyai_gateway_jwks_refresh_failures_total{outcome="ssf"} 1`, + `easyai_gateway_async_worker_capacity 96`, + `easyai_gateway_async_worker_desired_capacity 128`, + `easyai_gateway_async_worker_capacity_capped 1`, + `easyai_gateway_async_worker_resizes_total{outcome="success"} 1`, + `easyai_gateway_concurrency_lease_renewals_total{outcome="success"} 1`, + `easyai_gateway_concurrency_lease_renewals_total{outcome="lost"} 1`, + `easyai_gateway_task_events_skipped_total{outcome="duplicate"} 1`, + `easyai_gateway_task_events_skipped_total{outcome="budget_exceeded"} 1`, } { if !strings.Contains(recorder.Body.String(), expected) { t.Fatalf("missing metric %q in:\n%s", expected, recorder.Body.String()) diff --git a/apps/api/internal/store/admin_task_redaction.go b/apps/api/internal/store/admin_task_redaction.go new file mode 100644 index 0000000..a6e7e95 --- /dev/null +++ b/apps/api/internal/store/admin_task_redaction.go @@ -0,0 +1,94 @@ +package store + +import "strings" + +const maskedAdminTaskValue = "***" + +func MaskAdminGatewayTask(task AdminGatewayTask) AdminGatewayTask { + task.Request = maskSensitiveObject(task.Request) + task.Result = maskSensitiveObject(task.Result) + task.Billings = maskSensitiveArray(task.Billings) + task.Usage = maskSensitiveObject(task.Usage) + task.Metrics = maskSensitiveObject(task.Metrics) + task.BillingSummary = maskSensitiveObject(task.BillingSummary) + task.PricingSnapshot = maskSensitiveObject(task.PricingSnapshot) + task.CompatibilitySubmitHeaders = maskSensitiveObject(task.CompatibilitySubmitHeaders) + task.CompatibilitySubmitBody = maskSensitiveObject(task.CompatibilitySubmitBody) + task.Attempts = append([]TaskAttempt(nil), task.Attempts...) + for index := range task.Attempts { + task.Attempts[index].Usage = maskSensitiveObject(task.Attempts[index].Usage) + task.Attempts[index].Metrics = maskSensitiveObject(task.Attempts[index].Metrics) + task.Attempts[index].RequestSnapshot = maskSensitiveObject(task.Attempts[index].RequestSnapshot) + task.Attempts[index].ResponseSnapshot = maskSensitiveObject(task.Attempts[index].ResponseSnapshot) + task.Attempts[index].PricingSnapshot = maskSensitiveObject(task.Attempts[index].PricingSnapshot) + } + return task +} + +func MaskTaskParamPreprocessingLogs(items []TaskParamPreprocessingLog) []TaskParamPreprocessingLog { + masked := make([]TaskParamPreprocessingLog, len(items)) + for index, item := range items { + item.ActualInput = maskSensitiveObject(item.ActualInput) + item.ConvertedOutput = maskSensitiveObject(item.ConvertedOutput) + item.Changes = maskSensitiveArray(item.Changes) + item.ModelSnapshot = maskSensitiveObject(item.ModelSnapshot) + masked[index] = item + } + return masked +} + +func maskSensitiveObject(value map[string]any) map[string]any { + if value == nil { + return nil + } + masked, _ := maskSensitiveValue(value).(map[string]any) + return masked +} + +func maskSensitiveArray(value []any) []any { + if value == nil { + return nil + } + masked, _ := maskSensitiveValue(value).([]any) + return masked +} + +func maskSensitiveValue(value any) any { + switch typed := value.(type) { + case map[string]any: + masked := make(map[string]any, len(typed)) + for key, item := range typed { + if adminTaskSensitiveKey(key) { + masked[key] = maskedAdminTaskValue + continue + } + masked[key] = maskSensitiveValue(item) + } + return masked + case []any: + masked := make([]any, len(typed)) + for index, item := range typed { + masked[index] = maskSensitiveValue(item) + } + return masked + default: + return value + } +} + +func adminTaskSensitiveKey(key string) bool { + normalized := strings.NewReplacer("_", "", "-", "", ".", "", " ", "").Replace(strings.ToLower(strings.TrimSpace(key))) + switch normalized { + case "authorization", "proxyauthorization", "apikey", "xapikey", "accesstoken", "refreshtoken", + "idtoken", "token", "secret", "clientsecret", "password", "passwd", "cookie", "setcookie", + "credentials", "credential", "privatekey": + return true + default: + return strings.HasSuffix(normalized, "apikey") || + strings.HasSuffix(normalized, "accesstoken") || + strings.HasSuffix(normalized, "refreshtoken") || + strings.HasSuffix(normalized, "clientsecret") || + strings.HasSuffix(normalized, "password") || + strings.HasSuffix(normalized, "privatekey") + } +} diff --git a/apps/api/internal/store/admin_task_redaction_test.go b/apps/api/internal/store/admin_task_redaction_test.go new file mode 100644 index 0000000..21bc9b9 --- /dev/null +++ b/apps/api/internal/store/admin_task_redaction_test.go @@ -0,0 +1,87 @@ +package store + +import ( + "strings" + "testing" +) + +func TestMaskAdminGatewayTaskRecursivelyMasksSecretsWithoutMutatingSource(t *testing.T) { + source := AdminGatewayTask{ + GatewayTask: GatewayTask{ + Request: map[string]any{ + "model": "example", + "headers": map[string]any{ + "Authorization": "Bearer private", + "X-Api-Key": "secret-key", + }, + "usage": map[string]any{"input_tokens": float64(12)}, + }, + Result: map[string]any{ + "nested": []any{map[string]any{"password": "private-password"}}, + }, + Attempts: []TaskAttempt{{ + RequestSnapshot: map[string]any{"client_secret": "private-client-secret"}, + }}, + }, + } + + masked := MaskAdminGatewayTask(source) + headers := masked.Request["headers"].(map[string]any) + if headers["Authorization"] != maskedAdminTaskValue || headers["X-Api-Key"] != maskedAdminTaskValue { + t.Fatalf("sensitive headers were not masked: %#v", headers) + } + if masked.Request["usage"].(map[string]any)["input_tokens"] != float64(12) { + t.Fatalf("token usage must remain visible: %#v", masked.Request) + } + nested := masked.Result["nested"].([]any)[0].(map[string]any) + if nested["password"] != maskedAdminTaskValue { + t.Fatalf("nested password was not masked: %#v", nested) + } + if masked.Attempts[0].RequestSnapshot["client_secret"] != maskedAdminTaskValue { + t.Fatalf("attempt snapshot secret was not masked: %#v", masked.Attempts[0]) + } + if source.Request["headers"].(map[string]any)["Authorization"] != "Bearer private" { + t.Fatal("masking mutated the source request") + } + if source.Attempts[0].RequestSnapshot["client_secret"] != "private-client-secret" { + t.Fatal("masking mutated the source attempts") + } +} + +func TestBuildAdminTaskWhereUsesSharedPlaceholdersAndAllFilters(t *testing.T) { + where, args := buildAdminTaskWhere(AdminTaskListFilter{ + Query: "needle", + GatewayTenant: "00000000-0000-0000-0000-000000000001", + GatewayUser: "00000000-0000-0000-0000-000000000002", + UserGroup: "00000000-0000-0000-0000-000000000003", + Status: "failed", + Platform: "00000000-0000-0000-0000-000000000004", + Model: "model-a", + ModelType: "image_generate", + RunMode: "production", + BillingStatus: "settled", + APIKey: "ops", + }) + sql := strings.Join(where, "\n") + if strings.Contains(sql, "%!") || strings.Contains(sql, "$%d") { + t.Fatalf("SQL contains an unresolved placeholder: %s", sql) + } + if len(args) != 11 { + t.Fatalf("argument count=%d, want 11", len(args)) + } + if strings.Count(sql, "$1") < 10 { + t.Fatalf("keyword search should reuse one placeholder, got: %s", sql) + } + for _, fragment := range []string{ + "t.gateway_tenant_id", + "t.gateway_user_id", + "t.user_group_id", + "platform_attempt.platform_id", + "t.billing_status", + "t.api_key_prefix", + } { + if !strings.Contains(sql, fragment) { + t.Fatalf("SQL is missing %q: %s", fragment, sql) + } + } +} diff --git a/apps/api/internal/store/admin_tasks.go b/apps/api/internal/store/admin_tasks.go new file mode 100644 index 0000000..01c12a1 --- /dev/null +++ b/apps/api/internal/store/admin_tasks.go @@ -0,0 +1,385 @@ +package store + +import ( + "context" + "fmt" + "strings" + "time" +) + +type AdminTaskListFilter struct { + Query string + GatewayTenant string + GatewayUser string + UserGroup string + Status string + Platform string + Model string + ModelType string + RunMode string + BillingStatus string + APIKey string + CreatedFrom *time.Time + CreatedTo *time.Time + Page int + PageSize int +} + +type AdminTaskUserSummary struct { + ID string `json:"id"` + Username string `json:"username,omitempty"` + DisplayName string `json:"displayName,omitempty"` + Email string `json:"email,omitempty"` + Source string `json:"source,omitempty"` +} + +type AdminTaskTenantSummary struct { + ID string `json:"id"` + TenantKey string `json:"tenantKey,omitempty"` + Name string `json:"name,omitempty"` +} + +type AdminTaskPlatformSummary struct { + ID string `json:"id"` + Name string `json:"name,omitempty"` + Provider string `json:"provider,omitempty"` +} + +type AdminTaskContext struct { + User *AdminTaskUserSummary `json:"user,omitempty"` + Tenant *AdminTaskTenantSummary `json:"tenant,omitempty"` + LatestPlatform *AdminTaskPlatformSummary `json:"latestPlatform,omitempty"` +} + +type AdminGatewayTask struct { + GatewayTask + AdminContext AdminTaskContext `json:"adminContext"` +} + +type AdminTaskListResult struct { + Items []AdminGatewayTask + Total int + Page int + PageSize int +} + +func (s *Store) ListAdminTasks(ctx context.Context, filter AdminTaskListFilter) (AdminTaskListResult, error) { + page := filter.Page + if page <= 0 { + page = 1 + } + pageSize := filter.PageSize + if pageSize <= 0 { + pageSize = 50 + } + if pageSize > 100 { + pageSize = 100 + } + + where, args := buildAdminTaskWhere(filter) + var total int + if err := s.pool.QueryRow(ctx, ` +SELECT count(*) +FROM gateway_tasks t +WHERE `+strings.Join(where, "\n AND "), args...).Scan(&total); err != nil { + return AdminTaskListResult{}, err + } + + queryArgs := append([]any{}, args...) + queryArgs = append(queryArgs, pageSize, (page-1)*pageSize) + rows, err := s.pool.Query(ctx, ` +SELECT t.id::text +FROM gateway_tasks t +WHERE `+strings.Join(where, "\n AND ")+` +ORDER BY t.created_at DESC, t.id DESC +LIMIT $`+fmt.Sprint(len(args)+1)+` OFFSET $`+fmt.Sprint(len(args)+2), queryArgs...) + if err != nil { + return AdminTaskListResult{}, err + } + defer rows.Close() + + taskIDs := make([]string, 0, pageSize) + for rows.Next() { + var taskID string + if err := rows.Scan(&taskID); err != nil { + return AdminTaskListResult{}, err + } + taskIDs = append(taskIDs, taskID) + } + if err := rows.Err(); err != nil { + return AdminTaskListResult{}, err + } + + items, err := s.loadAdminTasksByIDs(ctx, taskIDs) + if err != nil { + return AdminTaskListResult{}, err + } + return AdminTaskListResult{ + Items: items, + Total: total, + Page: page, + PageSize: pageSize, + }, nil +} + +func (s *Store) GetAdminTask(ctx context.Context, taskID string) (AdminGatewayTask, error) { + task, err := s.GetTask(ctx, taskID) + if err != nil { + return AdminGatewayTask{}, err + } + contexts, err := s.loadAdminTaskContexts(ctx, []string{task.ID}) + if err != nil { + return AdminGatewayTask{}, err + } + return adminTaskWithContext(task, contexts[task.ID]), nil +} + +func buildAdminTaskWhere(filter AdminTaskListFilter) ([]string, []any) { + where := []string{"TRUE"} + args := make([]any, 0, 14) + add := func(clause string, value any) { + args = append(args, value) + placeholder := fmt.Sprintf("$%d", len(args)) + where = append(where, strings.ReplaceAll(clause, "$%d", placeholder)) + } + + if query := strings.TrimSpace(filter.Query); query != "" { + add(`( + t.id::text ILIKE $%d + OR COALESCE(t.request_id, '') ILIKE $%d + OR t.kind ILIKE $%d + OR t.model ILIKE $%d + OR COALESCE(t.requested_model, '') ILIKE $%d + OR COALESCE(t.resolved_model, '') ILIKE $%d + OR COALESCE(t.api_key_id, '') ILIKE $%d + OR COALESCE(t.api_key_name, '') ILIKE $%d + OR COALESCE(t.api_key_prefix, '') ILIKE $%d + OR COALESCE(t.model_type, '') ILIKE $%d + OR EXISTS ( + SELECT 1 + FROM gateway_users admin_user + WHERE admin_user.id = t.gateway_user_id + AND ( + admin_user.username ILIKE $%d + OR COALESCE(admin_user.display_name, '') ILIKE $%d + OR COALESCE(admin_user.email, '') ILIKE $%d + OR admin_user.user_key ILIKE $%d + ) + ) + OR EXISTS ( + SELECT 1 + FROM gateway_tenants admin_tenant + WHERE admin_tenant.id = t.gateway_tenant_id + AND ( + admin_tenant.name ILIKE $%d + OR admin_tenant.tenant_key ILIKE $%d + ) + ) + OR EXISTS ( + SELECT 1 + FROM gateway_task_attempts admin_attempt + LEFT JOIN integration_platforms admin_platform ON admin_platform.id = admin_attempt.platform_id + LEFT JOIN platform_models admin_model ON admin_model.id = admin_attempt.platform_model_id + WHERE admin_attempt.task_id = t.id + AND ( + COALESCE(admin_platform.name, '') ILIKE $%d + OR COALESCE(admin_platform.internal_name, '') ILIKE $%d + OR COALESCE(admin_platform.provider, '') ILIKE $%d + OR COALESCE(admin_model.model_name, '') ILIKE $%d + OR COALESCE(admin_model.provider_model_name, '') ILIKE $%d + OR COALESCE(admin_model.model_alias, '') ILIKE $%d + ) + ) + )`, "%"+query+"%") + } + if value := strings.TrimSpace(filter.GatewayTenant); value != "" { + add("t.gateway_tenant_id = $%d::uuid", value) + } + if value := strings.TrimSpace(filter.GatewayUser); value != "" { + add("t.gateway_user_id = $%d::uuid", value) + } + if value := strings.TrimSpace(filter.UserGroup); value != "" { + add("t.user_group_id = $%d::uuid", value) + } + if value := strings.TrimSpace(filter.Status); value != "" { + add("t.status = $%d", value) + } + if value := strings.TrimSpace(filter.Platform); value != "" { + add(`EXISTS ( + SELECT 1 + FROM gateway_task_attempts platform_attempt + WHERE platform_attempt.task_id = t.id + AND platform_attempt.platform_id = $%d::uuid + )`, value) + } + if value := strings.TrimSpace(filter.Model); value != "" { + add(`( + t.model = $%d + OR COALESCE(t.requested_model, '') = $%d + OR COALESCE(t.resolved_model, '') = $%d + OR EXISTS ( + SELECT 1 + FROM gateway_task_attempts model_attempt + LEFT JOIN platform_models selected_model ON selected_model.id = model_attempt.platform_model_id + WHERE model_attempt.task_id = t.id + AND ( + COALESCE(selected_model.model_name, '') = $%d + OR COALESCE(selected_model.provider_model_name, '') = $%d + OR COALESCE(selected_model.model_alias, '') = $%d + ) + ) + )`, value) + } + if value := strings.TrimSpace(filter.ModelType); value != "" { + add("COALESCE(t.model_type, '') = $%d", value) + } + if value := strings.TrimSpace(filter.RunMode); value != "" { + add("t.run_mode = $%d", value) + } + if value := strings.TrimSpace(filter.BillingStatus); value != "" { + add("t.billing_status = $%d", value) + } + if value := strings.TrimSpace(filter.APIKey); value != "" { + add(`( + COALESCE(t.api_key_id, '') ILIKE $%d + OR COALESCE(t.api_key_name, '') ILIKE $%d + OR COALESCE(t.api_key_prefix, '') ILIKE $%d + )`, "%"+value+"%") + } + if filter.CreatedFrom != nil { + add("t.created_at >= $%d::timestamptz", *filter.CreatedFrom) + } + if filter.CreatedTo != nil { + add("t.created_at <= $%d::timestamptz", *filter.CreatedTo) + } + return where, args +} + +func (s *Store) loadAdminTasksByIDs(ctx context.Context, taskIDs []string) ([]AdminGatewayTask, error) { + if len(taskIDs) == 0 { + return []AdminGatewayTask{}, nil + } + rows, err := s.pool.Query(ctx, ` +SELECT `+gatewayTaskColumns+` +FROM gateway_tasks +WHERE id::text = ANY($1)`, taskIDs) + if err != nil { + return nil, err + } + defer rows.Close() + + tasksByID := make(map[string]GatewayTask, len(taskIDs)) + tasks := make([]GatewayTask, 0, len(taskIDs)) + for rows.Next() { + task, err := scanGatewayTask(rows) + if err != nil { + return nil, err + } + tasks = append(tasks, task) + } + if err := rows.Err(); err != nil { + return nil, err + } + tasks, err = s.attachTaskAttempts(ctx, tasks) + if err != nil { + return nil, err + } + for _, task := range tasks { + tasksByID[task.ID] = task + } + contexts, err := s.loadAdminTaskContexts(ctx, taskIDs) + if err != nil { + return nil, err + } + + items := make([]AdminGatewayTask, 0, len(taskIDs)) + for _, taskID := range taskIDs { + task, ok := tasksByID[taskID] + if !ok { + continue + } + items = append(items, adminTaskWithContext(task, contexts[taskID])) + } + return items, nil +} + +func (s *Store) loadAdminTaskContexts(ctx context.Context, taskIDs []string) (map[string]AdminTaskContext, error) { + contexts := make(map[string]AdminTaskContext, len(taskIDs)) + if len(taskIDs) == 0 { + return contexts, nil + } + rows, err := s.pool.Query(ctx, ` +SELECT t.id::text, + COALESCE(u.id::text, ''), COALESCE(u.username, ''), COALESCE(u.display_name, ''), + COALESCE(u.email, ''), COALESCE(u.source, ''), + COALESCE(tenant.id::text, ''), COALESCE(tenant.tenant_key, ''), COALESCE(tenant.name, '') +FROM gateway_tasks t +LEFT JOIN gateway_users u ON u.id = t.gateway_user_id +LEFT JOIN gateway_tenants tenant ON tenant.id = t.gateway_tenant_id +WHERE t.id::text = ANY($1)`, taskIDs) + if err != nil { + return nil, err + } + defer rows.Close() + for rows.Next() { + var taskID string + var user AdminTaskUserSummary + var tenant AdminTaskTenantSummary + if err := rows.Scan( + &taskID, + &user.ID, + &user.Username, + &user.DisplayName, + &user.Email, + &user.Source, + &tenant.ID, + &tenant.TenantKey, + &tenant.Name, + ); err != nil { + return nil, err + } + context := AdminTaskContext{} + if user.ID != "" { + context.User = &user + } + if tenant.ID != "" { + context.Tenant = &tenant + } + contexts[taskID] = context + } + return contexts, rows.Err() +} + +func adminTaskWithContext(task GatewayTask, context AdminTaskContext) AdminGatewayTask { + if context.User == nil { + userID := strings.TrimSpace(task.GatewayUserID) + if userID == "" { + userID = strings.TrimSpace(task.UserID) + } + if userID != "" { + context.User = &AdminTaskUserSummary{ID: userID, Source: task.UserSource} + } + } + if context.Tenant == nil { + tenantID := strings.TrimSpace(task.GatewayTenantID) + if tenantID == "" { + tenantID = strings.TrimSpace(task.TenantID) + } + if tenantID != "" || strings.TrimSpace(task.TenantKey) != "" { + context.Tenant = &AdminTaskTenantSummary{ID: tenantID, TenantKey: task.TenantKey} + } + } + for index := len(task.Attempts) - 1; index >= 0; index-- { + attempt := task.Attempts[index] + if strings.TrimSpace(attempt.PlatformID) == "" { + continue + } + context.LatestPlatform = &AdminTaskPlatformSummary{ + ID: attempt.PlatformID, + Name: attempt.PlatformName, + Provider: attempt.Provider, + } + break + } + return AdminGatewayTask{GatewayTask: task, AdminContext: context} +} diff --git a/apps/api/internal/store/async_worker_capacity.go b/apps/api/internal/store/async_worker_capacity.go new file mode 100644 index 0000000..07a5d76 --- /dev/null +++ b/apps/api/internal/store/async_worker_capacity.go @@ -0,0 +1,152 @@ +package store + +import ( + "context" + "fmt" +) + +type AsyncWorkerCapacitySnapshot struct { + Capacity int + Desired int + HardLimit int + Capped bool + EnabledModels int + UnlimitedModels int + EnabledGroups int + UnlimitedGroups int + ModelDesired int + GroupDesired int +} + +func (s *Store) AsyncWorkerCapacity(ctx context.Context, hardLimit int) (AsyncWorkerCapacitySnapshot, error) { + if hardLimit < 1 { + return AsyncWorkerCapacitySnapshot{}, fmt.Errorf("async worker hard limit must be positive") + } + rows, err := s.pool.Query(ctx, ` +SELECT COALESCE(b.default_rate_limit_policy, '{}'::jsonb), + p.rate_limit_policy, + COALESCE(rp.rate_limit_policy, '{}'::jsonb), + (m.runtime_policy_set_id IS NOT NULL), + COALESCE(m.runtime_policy_override, '{}'::jsonb), + m.rate_limit_policy, + m.rate_limit_policy_mode +FROM platform_models m +JOIN integration_platforms p ON p.id = m.platform_id +LEFT JOIN base_model_catalog b ON b.id = m.base_model_id +LEFT JOIN model_runtime_policy_sets rp ON rp.id = COALESCE(m.runtime_policy_set_id, b.runtime_policy_set_id) +WHERE p.status = 'enabled' + AND p.deleted_at IS NULL + AND m.enabled = true`) + if err != nil { + return AsyncWorkerCapacitySnapshot{}, err + } + defer rows.Close() + + modelPolicies := make([]map[string]any, 0) + for rows.Next() { + var basePolicyBytes, platformPolicyBytes, runtimePolicyBytes []byte + var runtimeOverrideBytes, modelPolicyBytes []byte + var runtimeExplicit bool + var modelPolicyMode string + if err := rows.Scan( + &basePolicyBytes, + &platformPolicyBytes, + &runtimePolicyBytes, + &runtimeExplicit, + &runtimeOverrideBytes, + &modelPolicyBytes, + &modelPolicyMode, + ); err != nil { + return AsyncWorkerCapacitySnapshot{}, err + } + modelPolicies = append(modelPolicies, EffectiveRateLimitPolicy(EffectiveRateLimitPolicyInput{ + BasePolicy: decodeObject(basePolicyBytes), + PlatformPolicy: decodeObject(platformPolicyBytes), + RuntimePolicy: decodeObject(runtimePolicyBytes), + RuntimePolicyExplicit: runtimeExplicit, + RuntimePolicyOverride: decodeObject(runtimeOverrideBytes), + ModelPolicy: decodeObject(modelPolicyBytes), + ModelPolicyMode: modelPolicyMode, + })) + } + if err := rows.Err(); err != nil { + return AsyncWorkerCapacitySnapshot{}, err + } + groupRows, err := s.pool.Query(ctx, ` +SELECT rate_limit_policy +FROM gateway_user_groups +WHERE status = 'active'`) + if err != nil { + return AsyncWorkerCapacitySnapshot{}, err + } + defer groupRows.Close() + groupPolicies := make([]map[string]any, 0) + for groupRows.Next() { + var policyBytes []byte + if err := groupRows.Scan(&policyBytes); err != nil { + return AsyncWorkerCapacitySnapshot{}, err + } + groupPolicies = append(groupPolicies, NormalizeRateLimitPolicy(decodeObject(policyBytes))) + } + if err := groupRows.Err(); err != nil { + return AsyncWorkerCapacitySnapshot{}, err + } + return asyncWorkerCapacityFromPolicySets(modelPolicies, groupPolicies, hardLimit), nil +} + +func asyncWorkerCapacityFromPolicies(policies []map[string]any, hardLimit int) AsyncWorkerCapacitySnapshot { + return asyncWorkerCapacityFromPolicySets(policies, nil, hardLimit) +} + +func asyncWorkerCapacityFromPolicySets(modelPolicies []map[string]any, groupPolicies []map[string]any, hardLimit int) AsyncWorkerCapacitySnapshot { + snapshot := AsyncWorkerCapacitySnapshot{ + HardLimit: hardLimit, + EnabledModels: len(modelPolicies), + EnabledGroups: len(groupPolicies), + } + modelDesired, modelFinite, unlimitedModels := concurrentPolicySetCapacity(modelPolicies) + groupDesired, groupFinite, unlimitedGroups := concurrentPolicySetCapacity(groupPolicies) + snapshot.ModelDesired = modelDesired + snapshot.GroupDesired = groupDesired + snapshot.UnlimitedModels = unlimitedModels + snapshot.UnlimitedGroups = unlimitedGroups + + desired := hardLimit + switch { + case snapshot.EnabledModels == 0: + desired = 1 + case modelFinite && groupFinite: + desired = min(modelDesired, groupDesired) + case modelFinite: + desired = modelDesired + case groupFinite: + desired = groupDesired + } + if desired < 1 { + desired = 1 + } + snapshot.Desired = desired + snapshot.Capacity = desired + if snapshot.Capacity > hardLimit { + snapshot.Capacity = hardLimit + snapshot.Capped = true + } + return snapshot +} + +func concurrentPolicySetCapacity(policies []map[string]any) (total int, finite bool, unlimited int) { + if len(policies) == 0 { + return 0, false, 0 + } + finite = true + for _, policy := range policies { + capacity, policyFinite := ConcurrentPolicyCapacity(policy) + if !policyFinite { + unlimited++ + finite = false + continue + } + total += capacity + } + return total, finite, unlimited +} diff --git a/apps/api/internal/store/base_models.go b/apps/api/internal/store/base_models.go index e21bc30..3b38452 100644 --- a/apps/api/internal/store/base_models.go +++ b/apps/api/internal/store/base_models.go @@ -3,25 +3,36 @@ package store import ( "context" "encoding/json" + "fmt" "strings" "github.com/jackc/pgx/v5" ) const baseModelColumns = ` -id::text, provider_key, canonical_model_key, provider_model_name, model_type, display_name, +id::text, provider_key, canonical_model_key, invocation_name, provider_model_name, model_type, display_name, capabilities, base_billing_config, default_rate_limit_policy, COALESCE(pricing_rule_set_id::text, ''), COALESCE(runtime_policy_set_id::text, ''), runtime_policy_override, metadata, catalog_type, COALESCE(default_snapshot, '{}'::jsonb), COALESCE(customized_at::text, ''), -pricing_version, status, created_at, updated_at` +pricing_version, status, created_at, updated_at, +COALESCE(( + SELECT jsonb_agg(DISTINCT compatibility_alias.alias ORDER BY compatibility_alias.alias) + FROM model_compatibility_aliases compatibility_alias + WHERE compatibility_alias.base_model_id = base_model_catalog.id + AND compatibility_alias.active = true + AND (compatibility_alias.expires_at IS NULL OR compatibility_alias.expires_at > now()) +), '[]'::jsonb), +(SELECT count(*)::int FROM platform_models platform_model WHERE platform_model.base_model_id = base_model_catalog.id)` type BaseModelInput struct { ProviderKey string `json:"providerKey"` CanonicalModelKey string `json:"canonicalModelKey"` + InvocationName string `json:"invocationName"` ProviderModelName string `json:"providerModelName"` ModelType StringList `json:"modelType"` ModelAlias string `json:"modelAlias"` DisplayName string `json:"displayName"` + LegacyAliases StringList `json:"legacyAliases"` Capabilities map[string]any `json:"capabilities"` BaseBillingConfig map[string]any `json:"baseBillingConfig"` DefaultRateLimitPolicy map[string]any `json:"defaultRateLimitPolicy"` @@ -86,25 +97,32 @@ func (s *Store) CreateBaseModel(ctx context.Context, input BaseModelInput) (Base defaultSnapshot, _ := json.Marshal(emptyObjectIfNil(input.DefaultSnapshot)) modelType, _ := json.Marshal(input.ModelType) - return scanBaseModel(s.pool.QueryRow(ctx, ` + tx, err := s.pool.Begin(ctx) + if err != nil { + return BaseModel{}, err + } + defer tx.Rollback(ctx) + + item, err := scanBaseModel(tx.QueryRow(ctx, ` INSERT INTO base_model_catalog ( - provider_id, provider_key, canonical_model_key, provider_model_name, model_type, display_name, + provider_id, provider_key, canonical_model_key, invocation_name, provider_model_name, model_type, display_name, capabilities, base_billing_config, default_rate_limit_policy, pricing_rule_set_id, runtime_policy_set_id, runtime_policy_override, metadata, catalog_type, default_snapshot, pricing_version, status ) VALUES ( (SELECT id FROM model_catalog_providers WHERE provider_key = $1 OR provider_code = $1 LIMIT 1), - $1, $2, $3, $4::jsonb, $5, $6, $7, $8, - COALESCE(NULLIF($9, '')::uuid, (SELECT id FROM model_pricing_rule_sets WHERE rule_set_key = 'default-multimodal-v1' LIMIT 1)), - COALESCE(NULLIF($10, '')::uuid, (SELECT id FROM model_runtime_policy_sets WHERE policy_key = 'default-runtime-v1' LIMIT 1)), - $11, $12, NULLIF($13, ''), NULLIF($14::jsonb, '{}'::jsonb), $15, $16 + $1, $2, $3, $4, $5::jsonb, $6, $7, $8, $9, + COALESCE(NULLIF($10, '')::uuid, (SELECT id FROM model_pricing_rule_sets WHERE rule_set_key = 'default-multimodal-v1' LIMIT 1)), + COALESCE(NULLIF($11, '')::uuid, (SELECT id FROM model_runtime_policy_sets WHERE policy_key = 'default-runtime-v1' LIMIT 1)), + $12, $13, NULLIF($14, ''), NULLIF($15::jsonb, '{}'::jsonb), $16, $17 ) RETURNING `+baseModelColumns, input.ProviderKey, input.CanonicalModelKey, + input.InvocationName, input.ProviderModelName, string(modelType), - input.ModelAlias, + input.DisplayName, capabilities, billingConfig, rateLimitPolicy, @@ -117,6 +135,17 @@ RETURNING `+baseModelColumns, input.PricingVersion, input.Status, )) + if err != nil { + return BaseModel{}, err + } + if err := replaceBaseModelAliases(ctx, tx, item.ID, input); err != nil { + return BaseModel{}, err + } + if err := tx.Commit(ctx); err != nil { + return BaseModel{}, err + } + item.LegacyAliases = input.LegacyAliases + return item, nil } func (s *Store) UpdateBaseModel(ctx context.Context, id string, input BaseModelInput) (BaseModel, error) { @@ -129,35 +158,43 @@ func (s *Store) UpdateBaseModel(ctx context.Context, id string, input BaseModelI defaultSnapshot, _ := json.Marshal(emptyObjectIfNil(input.DefaultSnapshot)) modelType, _ := json.Marshal(input.ModelType) - return scanBaseModel(s.pool.QueryRow(ctx, ` + tx, err := s.pool.Begin(ctx) + if err != nil { + return BaseModel{}, err + } + defer tx.Rollback(ctx) + + item, err := scanBaseModel(tx.QueryRow(ctx, ` UPDATE base_model_catalog SET provider_id = (SELECT id FROM model_catalog_providers WHERE provider_key = $2 OR provider_code = $2 LIMIT 1), provider_key = $2, canonical_model_key = $3, - provider_model_name = $4, - model_type = $5::jsonb, - display_name = $6, - capabilities = $7, - base_billing_config = $8, - default_rate_limit_policy = $9, - pricing_rule_set_id = COALESCE(NULLIF($10, '')::uuid, (SELECT id FROM model_pricing_rule_sets WHERE rule_set_key = 'default-multimodal-v1' LIMIT 1)), - runtime_policy_set_id = COALESCE(NULLIF($11, '')::uuid, (SELECT id FROM model_runtime_policy_sets WHERE policy_key = 'default-runtime-v1' LIMIT 1)), - runtime_policy_override = $12, - metadata = $13, - catalog_type = NULLIF($14, ''), - default_snapshot = COALESCE(NULLIF($15::jsonb, '{}'::jsonb), default_snapshot), - customized_at = CASE WHEN NULLIF($14, '') = 'system' THEN now() ELSE NULL END, - pricing_version = $16, - status = $17, + invocation_name = $4, + provider_model_name = $5, + model_type = $6::jsonb, + display_name = $7, + capabilities = $8, + base_billing_config = $9, + default_rate_limit_policy = $10, + pricing_rule_set_id = COALESCE(NULLIF($11, '')::uuid, (SELECT id FROM model_pricing_rule_sets WHERE rule_set_key = 'default-multimodal-v1' LIMIT 1)), + runtime_policy_set_id = COALESCE(NULLIF($12, '')::uuid, (SELECT id FROM model_runtime_policy_sets WHERE policy_key = 'default-runtime-v1' LIMIT 1)), + runtime_policy_override = $13, + metadata = $14, + catalog_type = NULLIF($15, ''), + default_snapshot = COALESCE(NULLIF($16::jsonb, '{}'::jsonb), default_snapshot), + customized_at = CASE WHEN NULLIF($15, '') = 'system' THEN now() ELSE NULL END, + pricing_version = $17, + status = $18, updated_at = now() WHERE id = $1::uuid RETURNING `+baseModelColumns, id, input.ProviderKey, input.CanonicalModelKey, + input.InvocationName, input.ProviderModelName, string(modelType), - input.ModelAlias, + input.DisplayName, capabilities, billingConfig, rateLimitPolicy, @@ -170,6 +207,17 @@ RETURNING `+baseModelColumns, input.PricingVersion, input.Status, )) + if err != nil { + return BaseModel{}, err + } + if err := replaceBaseModelAliases(ctx, tx, item.ID, input); err != nil { + return BaseModel{}, err + } + if err := tx.Commit(ctx); err != nil { + return BaseModel{}, err + } + item.LegacyAliases = input.LegacyAliases + return item, nil } func (s *Store) ResetBaseModelToDefault(ctx context.Context, id string) (BaseModel, error) { @@ -194,18 +242,19 @@ UPDATE base_model_catalog SET provider_id = (SELECT id FROM model_catalog_providers WHERE provider_key = COALESCE($2::text, provider_key) OR provider_code = COALESCE($2::text, provider_key) LIMIT 1), provider_key = COALESCE($2::text, provider_key), canonical_model_key = COALESCE($3::text, canonical_model_key), - provider_model_name = COALESCE($4::text, provider_model_name), - model_type = COALESCE($5::jsonb, model_type), - display_name = COALESCE($6::text, display_name), - capabilities = COALESCE($7::jsonb, capabilities), - base_billing_config = COALESCE($8::jsonb, base_billing_config), - default_rate_limit_policy = COALESCE($9::jsonb, default_rate_limit_policy), - pricing_rule_set_id = COALESCE(NULLIF($10::text, '')::uuid, pricing_rule_set_id), - runtime_policy_set_id = COALESCE(NULLIF($11::text, '')::uuid, runtime_policy_set_id), - runtime_policy_override = COALESCE($12::jsonb, runtime_policy_override), - metadata = COALESCE($13::jsonb, metadata), - pricing_version = COALESCE($14::integer, pricing_version), - status = COALESCE($15::text, status), + invocation_name = COALESCE($4::text, invocation_name), + provider_model_name = COALESCE($5::text, provider_model_name), + model_type = COALESCE($6::jsonb, model_type), + display_name = COALESCE($7::text, display_name), + capabilities = COALESCE($8::jsonb, capabilities), + base_billing_config = COALESCE($9::jsonb, base_billing_config), + default_rate_limit_policy = COALESCE($10::jsonb, default_rate_limit_policy), + pricing_rule_set_id = COALESCE(NULLIF($11::text, '')::uuid, pricing_rule_set_id), + runtime_policy_set_id = COALESCE(NULLIF($12::text, '')::uuid, runtime_policy_set_id), + runtime_policy_override = COALESCE($13::jsonb, runtime_policy_override), + metadata = COALESCE($14::jsonb, metadata), + pricing_version = COALESCE($15::integer, pricing_version), + status = COALESCE($16::text, status), customized_at = NULL, updated_at = now() WHERE id = $1::uuid @@ -213,9 +262,10 @@ RETURNING `+baseModelColumns, id, stringFromSnapshot(snapshot, "providerKey"), stringFromSnapshot(snapshot, "canonicalModelKey"), + stringFromSnapshot(snapshot, "invocationName", "modelAlias", "providerModelName"), stringFromSnapshot(snapshot, "providerModelName"), jsonStringListFromSnapshot(snapshot, "modelType"), - stringFromSnapshot(snapshot, "modelAlias", "displayName"), + stringFromSnapshot(snapshot, "displayName", "modelAlias", "providerModelName"), jsonFromSnapshot(snapshot, "capabilities"), jsonFromSnapshot(snapshot, "baseBillingConfig"), jsonFromSnapshot(snapshot, "defaultRateLimitPolicy"), @@ -240,13 +290,14 @@ SET provider_id = ( ), provider_key = COALESCE(NULLIF(default_snapshot->>'providerKey', ''), provider_key), canonical_model_key = COALESCE(NULLIF(default_snapshot->>'canonicalModelKey', ''), canonical_model_key), + invocation_name = COALESCE(NULLIF(default_snapshot->>'invocationName', ''), NULLIF(default_snapshot->>'modelAlias', ''), invocation_name), provider_model_name = COALESCE(NULLIF(default_snapshot->>'providerModelName', ''), provider_model_name), model_type = COALESCE(NULLIF(CASE WHEN jsonb_typeof(default_snapshot->'modelType') = 'array' THEN default_snapshot->'modelType' WHEN COALESCE(default_snapshot->>'modelType', '') <> '' THEN jsonb_build_array(default_snapshot->>'modelType') ELSE NULL END, '[]'::jsonb), model_type), - display_name = COALESCE(NULLIF(COALESCE(default_snapshot->>'modelAlias', default_snapshot->>'displayName'), ''), display_name), + display_name = COALESCE(NULLIF(COALESCE(default_snapshot->>'displayName', default_snapshot->>'modelAlias'), ''), display_name), capabilities = COALESCE(default_snapshot->'capabilities', capabilities), base_billing_config = COALESCE(default_snapshot->'baseBillingConfig', base_billing_config), default_rate_limit_policy = COALESCE(default_snapshot->'defaultRateLimitPolicy', default_rate_limit_policy), @@ -269,11 +320,23 @@ RETURNING `+baseModelColumns) } func (s *Store) DeleteBaseModel(ctx context.Context, id string) error { - result, err := s.pool.Exec(ctx, `DELETE FROM base_model_catalog WHERE id = $1::uuid`, id) + result, err := s.pool.Exec(ctx, ` +DELETE FROM base_model_catalog base_model +WHERE base_model.id = $1::uuid + AND NOT EXISTS ( + SELECT 1 FROM platform_models platform_model WHERE platform_model.base_model_id = base_model.id + )`, id) if err != nil { return err } if result.RowsAffected() == 0 { + var exists bool + if err := s.pool.QueryRow(ctx, `SELECT EXISTS (SELECT 1 FROM base_model_catalog WHERE id = $1::uuid)`, id).Scan(&exists); err != nil { + return err + } + if exists { + return ErrBaseModelInUse + } return pgx.ErrNoRows } return nil @@ -294,7 +357,7 @@ func scanBaseModelRows(rows pgx.Rows) ([]BaseModel, error) { func scanBaseModel(scanner baseModelScanner) (BaseModel, error) { var item BaseModel var modelType []byte - var modelAlias string + var legacyAliases []byte var capabilities []byte var billingConfig []byte var rateLimitPolicy []byte @@ -305,9 +368,10 @@ func scanBaseModel(scanner baseModelScanner) (BaseModel, error) { &item.ID, &item.ProviderKey, &item.CanonicalModelKey, + &item.InvocationName, &item.ProviderModelName, &modelType, - &modelAlias, + &item.DisplayName, &capabilities, &billingConfig, &rateLimitPolicy, @@ -322,6 +386,8 @@ func scanBaseModel(scanner baseModelScanner) (BaseModel, error) { &item.Status, &item.CreatedAt, &item.UpdatedAt, + &legacyAliases, + &item.ReferenceCount, ); err != nil { return BaseModel{}, err } @@ -332,18 +398,20 @@ func scanBaseModel(scanner baseModelScanner) (BaseModel, error) { item.Metadata = decodeObject(metadata) item.DefaultSnapshot = decodeObject(defaultSnapshot) item.ModelType = decodeStringArray(modelType) - item.ModelAlias = modelAlias - item.DisplayName = modelAlias + item.LegacyAliases = decodeStringArray(legacyAliases) + item.ModelAlias = item.InvocationName return item, nil } func normalizeBaseModelInput(input BaseModelInput) BaseModelInput { input.ProviderKey = strings.TrimSpace(input.ProviderKey) input.CanonicalModelKey = strings.TrimSpace(input.CanonicalModelKey) + input.InvocationName = strings.TrimSpace(input.InvocationName) input.ProviderModelName = strings.TrimSpace(input.ProviderModelName) input.ModelType = uniqueStringList(input.ModelType) input.ModelAlias = strings.TrimSpace(input.ModelAlias) input.DisplayName = strings.TrimSpace(input.DisplayName) + input.LegacyAliases = normalizeCompatibilityAliases(input.LegacyAliases) input.PricingRuleSetID = strings.TrimSpace(input.PricingRuleSetID) input.RuntimePolicySetID = strings.TrimSpace(input.RuntimePolicySetID) input.CatalogType = strings.TrimSpace(input.CatalogType) @@ -351,12 +419,17 @@ func normalizeBaseModelInput(input BaseModelInput) BaseModelInput { if input.CanonicalModelKey == "" && input.ProviderKey != "" && input.ProviderModelName != "" { input.CanonicalModelKey = input.ProviderKey + ":" + input.ProviderModelName } - if input.ModelAlias == "" { - input.ModelAlias = input.DisplayName + if input.InvocationName == "" { + input.InvocationName = input.ModelAlias } - if input.ModelAlias == "" { - input.ModelAlias = input.ProviderModelName + if input.InvocationName == "" { + input.InvocationName = input.ProviderModelName } + if input.DisplayName == "" { + input.DisplayName = input.InvocationName + } + input.ModelAlias = input.InvocationName + input.LegacyAliases = withoutStrings(input.LegacyAliases, input.InvocationName) if len(input.ModelType) == 0 { input.ModelType = StringList{"text_generate"} } @@ -372,6 +445,78 @@ func normalizeBaseModelInput(input BaseModelInput) BaseModelInput { return input } +func replaceBaseModelAliases(ctx context.Context, tx pgx.Tx, baseModelID string, input BaseModelInput) error { + for _, alias := range input.LegacyAliases { + for _, modelType := range input.ModelType { + var conflictingCanonicalKey string + if err := tx.QueryRow(ctx, ` +SELECT COALESCE(( + SELECT other.canonical_model_key + FROM base_model_catalog other + WHERE other.id <> $1::uuid + AND other.invocation_name <> $4::text + AND other.model_type ? $3::text + AND ( + other.invocation_name = $2::text + OR EXISTS ( + SELECT 1 + FROM model_compatibility_aliases other_alias + WHERE other_alias.base_model_id = other.id + AND other_alias.alias = $2::text + AND other_alias.model_type = $3::text + AND other_alias.active = true + AND (other_alias.expires_at IS NULL OR other_alias.expires_at > now()) + ) + ) + LIMIT 1 +), '')`, baseModelID, alias, modelType, input.InvocationName).Scan(&conflictingCanonicalKey); err != nil { + return err + } + if conflictingCanonicalKey != "" { + return fmt.Errorf("%w: %q (%s) conflicts with %s", ErrModelAliasConflict, alias, modelType, conflictingCanonicalKey) + } + } + } + + if _, err := tx.Exec(ctx, `DELETE FROM model_compatibility_aliases WHERE base_model_id = $1::uuid`, baseModelID); err != nil { + return err + } + for _, alias := range input.LegacyAliases { + for _, modelType := range input.ModelType { + if _, err := tx.Exec(ctx, ` +INSERT INTO model_compatibility_aliases (base_model_id, alias, model_type, expires_at) +VALUES ($1::uuid, $2, $3, now() + interval '14 days')`, baseModelID, alias, modelType); err != nil { + return err + } + } + } + return nil +} + +func normalizeCompatibilityAliases(values []string) StringList { + out := make(StringList, 0, len(values)) + seen := map[string]bool{} + for _, value := range values { + value = strings.TrimSpace(value) + if value == "" || seen[value] { + continue + } + seen[value] = true + out = append(out, value) + } + return out +} + +func withoutStrings(values []string, excluded string) StringList { + out := make(StringList, 0, len(values)) + for _, value := range values { + if value != excluded { + out = append(out, value) + } + } + return out +} + func stringFromSnapshot(snapshot map[string]any, keys ...string) any { for _, key := range keys { value, ok := snapshot[key] diff --git a/apps/api/internal/store/billing_settlements.go b/apps/api/internal/store/billing_settlements.go index 338e6bf..567df9b 100644 --- a/apps/api/internal/store/billing_settlements.go +++ b/apps/api/internal/store/billing_settlements.go @@ -432,13 +432,13 @@ SELECT task.id, COALESCE((SELECT MAX(event.seq) + 1 FROM gateway_task_events event WHERE event.task_id = task.id), 1), CASE WHEN $2 = 'settled' THEN 'task.billing.settled' ELSE 'task.billing.released' END, task.status, - 'billing', - 1, - CASE WHEN $2 = 'settled' THEN 'task billing settled' ELSE 'task billing reservation released' END, - jsonb_build_object('settlementId', $3::text, 'billingStatus', $2::text), + NULL, + 0, + NULL, + '{}'::jsonb, task.run_mode = 'simulation' FROM gateway_tasks task -WHERE task.id = $1::uuid`, taskID, billingStatus, settlementID); err != nil { +WHERE task.id = $1::uuid`, taskID, billingStatus); err != nil { return err } tag, err := tx.Exec(ctx, ` diff --git a/apps/api/internal/store/binary_result_backfill.go b/apps/api/internal/store/binary_result_backfill.go new file mode 100644 index 0000000..5df9528 --- /dev/null +++ b/apps/api/internal/store/binary_result_backfill.go @@ -0,0 +1,87 @@ +package store + +import ( + "context" + "encoding/json" + "time" + + "github.com/jackc/pgx/v5" +) + +type TaskBinaryResultBackfillItem struct { + ID string + Result map[string]any + FinishedAt time.Time +} + +func (s *Store) ListTaskBinaryResultBackfillBatch(ctx context.Context, afterID string, batchSize int) ([]TaskBinaryResultBackfillItem, error) { + if batchSize < 1 || batchSize > 100 { + batchSize = 100 + } + var items []TaskBinaryResultBackfillItem + err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error { + if _, err := tx.Exec(ctx, `SET LOCAL lock_timeout = '5s'; SET LOCAL statement_timeout = '5s'`); err != nil { + return err + } + rows, err := tx.Query(ctx, ` +SELECT id::text, result, COALESCE(finished_at, updated_at) +FROM gateway_tasks +WHERE status = 'succeeded' + AND result <> '{}'::jsonb + AND (NULLIF($1::text, '') IS NULL OR id > NULLIF($1::text, '')::uuid) +ORDER BY id +LIMIT $2`, afterID, batchSize) + if err != nil { + return err + } + defer rows.Close() + items = make([]TaskBinaryResultBackfillItem, 0, batchSize) + for rows.Next() { + var item TaskBinaryResultBackfillItem + var resultJSON []byte + if err := rows.Scan(&item.ID, &resultJSON, &item.FinishedAt); err != nil { + return err + } + item.Result = decodeObject(resultJSON) + items = append(items, item) + } + return rows.Err() + }) + return items, err +} + +func (s *Store) UpdateTaskBinaryResultBackfill(ctx context.Context, taskID string, result map[string]any) (bool, error) { + report := sanitizeJSONForStorageWithReport(minimalTaskResult(result)) + if report.BinaryCount > 0 { + return false, &taskPayloadBinaryError{ + target: ErrTaskResultBinaryNotMaterialized, + code: "result_binary_not_materialized", + count: report.BinaryCount, + } + } + resultJSON, err := json.Marshal(report.Value) + if err != nil { + return false, err + } + updated := false + err = pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error { + if _, err := tx.Exec(ctx, `SET LOCAL lock_timeout = '5s'; SET LOCAL statement_timeout = '5s'`); err != nil { + return err + } + tag, err := tx.Exec(ctx, ` +UPDATE gateway_tasks +SET result = $2::jsonb, + updated_at = now() +WHERE id = $1::uuid + AND status = 'succeeded'`, + taskID, + string(resultJSON), + ) + if err != nil { + return err + } + updated = tag.RowsAffected() == 1 + return nil + }) + return updated, err +} diff --git a/apps/api/internal/store/candidates.go b/apps/api/internal/store/candidates.go index 40d6cc8..d532cfb 100644 --- a/apps/api/internal/store/candidates.go +++ b/apps/api/internal/store/candidates.go @@ -2,13 +2,15 @@ package store import ( "context" + "errors" "fmt" "math" "sort" "strings" - "unicode" + "time" "github.com/easyai/easyai-ai-gateway/apps/api/internal/auth" + "github.com/jackc/pgx/v5" ) type ListModelCandidatesOptions struct { @@ -19,7 +21,6 @@ type ListModelCandidatesOptions struct { func (s *Store) ListModelCandidates(ctx context.Context, model string, modelType string, user *auth.User, options ...ListModelCandidatesOptions) ([]RuntimeModelCandidate, error) { exactModel := strings.TrimSpace(model) - modelMatchKey := normalizeModelMatchKey(exactModel) listOptions := normalizeListModelCandidatesOptions(modelType, options...) rows, err := s.pool.Query(ctx, ` SELECT p.id::text, p.platform_key, p.name, p.provider, @@ -27,27 +28,40 @@ func (s *Store) ListModelCandidates(ctx context.Context, model string, modelType COALESCE(p.base_url, ''), p.auth_type, p.credentials, p.config, p.default_pricing_mode, p.default_discount_factor::float8, COALESCE(p.pricing_rule_set_id::text, ''), - p.retry_policy, p.rate_limit_policy, + p.retry_policy, p.rate_limit_policy, COALESCE(b.default_rate_limit_policy, '{}'::jsonb), COALESCE(p.dynamic_priority, p.priority) AS effective_priority, m.id::text, COALESCE(m.base_model_id::text, ''), COALESCE(b.canonical_model_key, ''), COALESCE(NULLIF(m.provider_model_name, ''), m.model_name), m.model_name, COALESCE(m.model_alias, ''), $2::text AS requested_model_type, m.display_name, + EXISTS ( + SELECT 1 + FROM model_compatibility_aliases compatibility_alias + JOIN base_model_catalog alias_base ON alias_base.id = compatibility_alias.base_model_id + WHERE compatibility_alias.alias = $1::text + AND compatibility_alias.model_type = $2::text + AND compatibility_alias.active = true + AND (compatibility_alias.expires_at IS NULL OR compatibility_alias.expires_at > now()) + AND (compatibility_alias.base_model_id = b.id OR alias_base.invocation_name = b.invocation_name) + ) AS legacy_alias_used, CASE WHEN b.capabilities #> '{text_generate,supportedApiProtocols}' IS NOT NULL THEN jsonb_set( - COALESCE(m.capabilities, '{}'::jsonb), + COALESCE(b.capabilities, '{}'::jsonb) || COALESCE(m.capabilities, '{}'::jsonb), '{text_generate,supportedApiProtocols}', b.capabilities #> '{text_generate,supportedApiProtocols}', true ) - ELSE m.capabilities + ELSE COALESCE(b.capabilities, '{}'::jsonb) || COALESCE(m.capabilities, '{}'::jsonb) END AS effective_capabilities, m.capability_override, COALESCE(b.base_billing_config, '{}'::jsonb), m.billing_config, m.billing_config_override, m.pricing_mode, COALESCE(m.discount_factor, 0)::float8, COALESCE(m.pricing_rule_set_id::text, ''), COALESCE(b.pricing_rule_set_id::text, ''), - m.permission_config, m.retry_policy, m.rate_limit_policy, COALESCE(m.runtime_policy_set_id::text, COALESCE(b.runtime_policy_set_id::text, '')), + m.permission_config, m.retry_policy, m.rate_limit_policy, m.rate_limit_policy_mode, + COALESCE(m.runtime_policy_set_id::text, COALESCE(b.runtime_policy_set_id::text, '')), + (m.runtime_policy_set_id IS NOT NULL), COALESCE(NULLIF(m.runtime_policy_override, '{}'::jsonb), b.runtime_policy_override, '{}'::jsonb), + COALESCE(m.runtime_policy_override, '{}'::jsonb), COALESCE(rp.retry_policy, '{}'::jsonb), COALESCE(rp.rate_limit_policy, '{}'::jsonb), COALESCE(rp.auto_disable_policy, '{}'::jsonb), COALESCE(rp.degrade_policy, '{}'::jsonb), COALESCE(con.active, 0)::float8, @@ -74,10 +88,10 @@ func (s *Store) ListModelCandidates(ctx context.Context, model string, modelType ON s.client_id = p.platform_key || ':' || $2::text || ':' || COALESCE(NULLIF(m.provider_model_name, ''), m.model_name) LEFT JOIN LATERAL ( SELECT ca.cache_affinity_key, ca.request_count, ca.input_tokens, ca.cached_input_tokens, ca.ema_hit_ratio, ca.last_hit_ratio, ca.last_observed_at - FROM unnest($4::text[]) WITH ORDINALITY AS affinity_keys(cache_affinity_key, affinity_rank) + FROM unnest($3::text[]) WITH ORDINALITY AS affinity_keys(cache_affinity_key, affinity_rank) JOIN gateway_cache_affinity_stats ca ON ca.cache_affinity_key = affinity_keys.cache_affinity_key WHERE ca.client_id = p.platform_key || ':' || $2::text || ':' || COALESCE(NULLIF(m.provider_model_name, ''), m.model_name) - AND ($5::int <= 0 OR ca.last_observed_at >= now() - ($5::int * interval '1 second')) + AND ($4::int <= 0 OR ca.last_observed_at >= now() - ($4::int * interval '1 second')) ORDER BY affinity_keys.affinity_rank ASC, ca.cached_input_tokens DESC, ca.ema_hit_ratio DESC, ca.last_observed_at DESC LIMIT 1 ) ca ON TRUE @@ -138,61 +152,25 @@ WHERE p.status = 'enabled' AND m.model_type @> jsonb_build_array($2::text) AND (p.cooldown_until IS NULL OR p.cooldown_until <= now()) AND (m.cooldown_until IS NULL OR m.cooldown_until <= now()) - AND ( - ( - $2::text IN ('audio_generate', 'text_to_speech', 'voice_clone') - AND ( - m.model_alias = $1::text - OR m.model_name = $1::text - OR b.canonical_model_key = $1::text - OR b.provider_model_name = $1::text - OR ( - NULLIF($3::text, '') IS NOT NULL - AND ( - regexp_replace(COALESCE(m.model_alias, ''), '[[:space:]]+', '', 'g') = $3::text - OR regexp_replace(COALESCE(m.model_name, ''), '[[:space:]]+', '', 'g') = $3::text - OR regexp_replace(COALESCE(b.canonical_model_key, ''), '[[:space:]]+', '', 'g') = $3::text - OR regexp_replace(COALESCE(b.provider_model_name, ''), '[[:space:]]+', '', 'g') = $3::text - ) - ) - ) - ) - OR ( - $2::text NOT IN ('audio_generate', 'text_to_speech', 'voice_clone') - AND ( - ( - COALESCE(m.model_alias, '') <> '' - AND ( - m.model_alias = $1::text - OR ( - NULLIF($3::text, '') IS NOT NULL - AND regexp_replace(COALESCE(m.model_alias, ''), '[[:space:]]+', '', 'g') = $3::text - ) - ) - ) - OR ( - m.model_name = $1::text - OR COALESCE(NULLIF(m.provider_model_name, ''), m.model_name) = $1::text - OR b.canonical_model_key = $1::text - OR b.provider_model_name = $1::text - OR ( - NULLIF($3::text, '') IS NOT NULL - AND ( - regexp_replace(COALESCE(m.model_name, ''), '[[:space:]]+', '', 'g') = $3::text - OR regexp_replace(COALESCE(NULLIF(m.provider_model_name, ''), m.model_name), '[[:space:]]+', '', 'g') = $3::text - OR regexp_replace(COALESCE(b.canonical_model_key, ''), '[[:space:]]+', '', 'g') = $3::text - OR regexp_replace(COALESCE(b.provider_model_name, ''), '[[:space:]]+', '', 'g') = $3::text - ) - ) - ) - ) - ) - ) + AND ( + b.invocation_name = $1::text + OR (b.id IS NULL AND m.model_name = $1::text) + OR EXISTS ( + SELECT 1 + FROM model_compatibility_aliases compatibility_alias + JOIN base_model_catalog alias_base ON alias_base.id = compatibility_alias.base_model_id + WHERE compatibility_alias.alias = $1::text + AND compatibility_alias.model_type = $2::text + AND compatibility_alias.active = true + AND (compatibility_alias.expires_at IS NULL OR compatibility_alias.expires_at > now()) + AND (compatibility_alias.base_model_id = b.id OR alias_base.invocation_name = b.invocation_name) + ) + ) ORDER BY effective_priority ASC, COALESCE(s.running_count, 0) ASC, COALESCE(s.waiting_count, 0) ASC, COALESCE(s.last_assigned_at, to_timestamp(0)) ASC, - m.created_at ASC`, exactModel, modelType, modelMatchKey, listOptions.CacheAffinityKeys, cacheAffinityStaleAfterSeconds(listOptions.CacheAffinityPolicy)) + m.created_at ASC`, exactModel, modelType, listOptions.CacheAffinityKeys, cacheAffinityStaleAfterSeconds(listOptions.CacheAffinityPolicy)) if err != nil { return nil, err } @@ -205,6 +183,7 @@ WHERE p.status = 'enabled' var platformConfig []byte var platformRetryPolicy []byte var platformRateLimitPolicy []byte + var baseRateLimitPolicy []byte var capabilities []byte var capabilityOverride []byte var baseBilling []byte @@ -214,6 +193,7 @@ WHERE p.status = 'enabled' var modelRetryPolicy []byte var modelRateLimitPolicy []byte var runtimePolicyOverride []byte + var rateLimitRuntimeOverride []byte var runtimeRetryPolicy []byte var runtimeRateLimitPolicy []byte var autoDisablePolicy []byte @@ -250,6 +230,7 @@ WHERE p.status = 'enabled' &item.PlatformPricingRuleSetID, &platformRetryPolicy, &platformRateLimitPolicy, + &baseRateLimitPolicy, &item.PlatformPriority, &item.PlatformModelID, &item.BaseModelID, @@ -259,6 +240,7 @@ WHERE p.status = 'enabled' &item.ModelAlias, &item.ModelType, &item.DisplayName, + &item.LegacyAliasUsed, &capabilities, &capabilityOverride, &baseBilling, @@ -271,8 +253,11 @@ WHERE p.status = 'enabled' &permissionConfig, &modelRetryPolicy, &modelRateLimitPolicy, + &item.ModelRateLimitPolicyMode, &item.RuntimePolicySetID, + &item.RuntimePolicyExplicit, &runtimePolicyOverride, + &rateLimitRuntimeOverride, &runtimeRetryPolicy, &runtimeRateLimitPolicy, &autoDisablePolicy, @@ -301,6 +286,7 @@ WHERE p.status = 'enabled' item.PlatformConfig = decodeObject(platformConfig) item.PlatformRetryPolicy = decodeObject(platformRetryPolicy) item.PlatformRateLimitPolicy = decodeObject(platformRateLimitPolicy) + item.BaseRateLimitPolicy = decodeObject(baseRateLimitPolicy) item.Capabilities = decodeObject(capabilities) item.CapabilityOverride = decodeObject(capabilityOverride) item.BaseBillingConfig = decodeObject(baseBilling) @@ -310,6 +296,7 @@ WHERE p.status = 'enabled' item.ModelRetryPolicy = decodeObject(modelRetryPolicy) item.ModelRateLimitPolicy = decodeObject(modelRateLimitPolicy) item.RuntimePolicyOverride = decodeObject(runtimePolicyOverride) + item.RateLimitRuntimeOverride = decodeObject(rateLimitRuntimeOverride) item.RuntimeRetryPolicy = decodeObject(runtimeRetryPolicy) item.RuntimeRateLimitPolicy = decodeObject(runtimeRateLimitPolicy) item.AutoDisablePolicy = decodeObject(autoDisablePolicy) @@ -330,7 +317,15 @@ WHERE p.status = 'enabled' LastObservedUnix: cacheLastObservedUnix, }) applyRuntimeCandidateLoad(&item, runtimeCandidateLoadInput{ - Policy: effectiveModelRateLimitPolicy(item.PlatformRateLimitPolicy, item.RuntimeRateLimitPolicy, item.RuntimePolicySetID, item.RuntimePolicyOverride, item.ModelRateLimitPolicy), + Policy: EffectiveRateLimitPolicy(EffectiveRateLimitPolicyInput{ + BasePolicy: item.BaseRateLimitPolicy, + PlatformPolicy: item.PlatformRateLimitPolicy, + RuntimePolicy: item.RuntimeRateLimitPolicy, + RuntimePolicyExplicit: item.RuntimePolicyExplicit, + RuntimePolicyOverride: item.RateLimitRuntimeOverride, + ModelPolicy: item.ModelRateLimitPolicy, + ModelPolicyMode: item.ModelRateLimitPolicyMode, + }), ConcurrentActive: concurrentActive, QueuedWaiting: queuedWaiting, RPMUsed: rpmUsed, @@ -361,10 +356,31 @@ WHERE p.status = 'enabled' if len(items) == 0 { return nil, ErrNoModelCandidate } + s.recordLegacyAliasUsage(ctx, exactModel, items) sortRuntimeModelCandidates(items) return items, nil } +func (s *Store) recordLegacyAliasUsage(ctx context.Context, alias string, items []RuntimeModelCandidate) { + seen := map[string]bool{} + for _, item := range items { + if !item.LegacyAliasUsed || item.CanonicalModelKey == "" { + continue + } + key := item.CanonicalModelKey + "\x00" + item.ModelType + if seen[key] { + continue + } + seen[key] = true + _, _ = s.pool.Exec(ctx, ` +INSERT INTO model_alias_usage_metrics (alias, canonical_model_key, model_type, hit_count) +VALUES ($1, $2, $3, 1) +ON CONFLICT (alias, canonical_model_key, model_type) DO UPDATE +SET hit_count = model_alias_usage_metrics.hit_count + 1, + last_used_at = now()`, alias, item.CanonicalModelKey, item.ModelType) + } +} + func (s *Store) GetRuntimeModelCandidateForVoiceCloneDeletion(ctx context.Context, platformModelID string, platformID string) (RuntimeModelCandidate, bool, error) { platformModelID = strings.TrimSpace(platformModelID) platformID = strings.TrimSpace(platformID) @@ -487,14 +503,15 @@ func applyRuntimeCandidateCacheAffinity(candidate *RuntimeModelCandidate, option key = options.CacheAffinityKey } affinity := RuntimeCandidateCacheAffinity{ - Key: key, - RequestCount: input.RequestCount, - InputTokens: input.InputTokens, - CachedInputTokens: input.CachedInputTokens, - EMAHitRatio: boundedRatio(input.EMAHitRatio), - LastHitRatio: boundedRatio(input.LastHitRatio), - LastObservedUnix: input.LastObservedUnix, - AdjustedPriority: float64(candidate.PlatformPriority), + Key: key, + RequestCount: input.RequestCount, + InputTokens: input.InputTokens, + CachedInputTokens: input.CachedInputTokens, + EMAHitRatio: boundedRatio(input.EMAHitRatio), + LastHitRatio: boundedRatio(input.LastHitRatio), + LastObservedUnix: input.LastObservedUnix, + AdjustedPriority: float64(candidate.PlatformPriority), + MatchedPrefixDepth: cacheAffinityMatchedPrefixDepth(key, options.CacheAffinityKeys), } minSamples := cacheAffinityMinSamples(options.CacheAffinityPolicy) hasObservedCachedHit := affinity.CachedInputTokens > 0 || affinity.LastHitRatio > 0 @@ -525,6 +542,32 @@ func applyRuntimeCandidateCacheAffinity(candidate *RuntimeModelCandidate, option candidate.CacheAffinity = affinity } +func cacheAffinityMatchedPrefixDepth(key string, lookup []string) int { + prefix := "" + switch { + case strings.HasPrefix(key, "prompt_lcp_v2:"): + prefix = "prompt_lcp_v2:" + case strings.HasPrefix(key, "prompt_lcp:"): + prefix = "prompt_lcp:" + default: + return 0 + } + matchingKeys := make([]string, 0, len(lookup)) + for _, candidateKey := range lookup { + if strings.HasPrefix(candidateKey, prefix) { + matchingKeys = append(matchingKeys, candidateKey) + } + } + for index, candidateKey := range matchingKeys { + if candidateKey == key { + return len(matchingKeys) - index + cacheAffinityMinimumPrefixDepth - 1 + } + } + return 0 +} + +const cacheAffinityMinimumPrefixDepth = 2 + func cacheAffinityPolicyEnabled(policy map[string]any, modelType string) bool { if enabled, ok := policy["enabled"].(bool); ok && !enabled { return false @@ -699,6 +742,9 @@ func sortRuntimeModelCandidates(items []RuntimeModelCandidate) { if aFull != bFull { return !aFull } + if items[i].PlatformPriority != items[j].PlatformPriority { + return items[i].PlatformPriority < items[j].PlatformPriority + } if items[i].CacheAffinity.Applied != items[j].CacheAffinity.Applied { return items[i].CacheAffinity.Applied } @@ -733,6 +779,32 @@ func sortRuntimeModelCandidates(items []RuntimeModelCandidate) { } return false }) + appliedCount := 0 + for index := range items { + if items[index].CacheAffinity.Applied { + appliedCount++ + } + } + for index := range items { + items[index].CacheAffinity.CandidateCount = appliedCount + } + if len(items) == 0 || items[0].CacheAffinity.Applied { + return + } + for index := 1; index < len(items); index++ { + if !items[index].CacheAffinity.Applied { + continue + } + switch { + case runtimeCandidateFull(items[index]) && !runtimeCandidateFull(items[0]): + items[0].CacheAffinity.OverrideReason = "capacity_tier_unavailable" + case items[index].PlatformPriority != items[0].PlatformPriority: + items[0].CacheAffinity.OverrideReason = "different_effective_priority_tier" + } + if items[0].CacheAffinity.OverrideReason != "" { + break + } + } } func runtimeCandidateFull(candidate RuntimeModelCandidate) bool { @@ -741,14 +813,20 @@ func runtimeCandidateFull(candidate RuntimeModelCandidate) bool { func (s *Store) modelCandidateCooldownError(ctx context.Context, model string, modelType string) (error, error) { exactModel := strings.TrimSpace(model) - modelMatchKey := normalizeModelMatchKey(exactModel) - rows, err := s.pool.Query(ctx, ` -SELECT p.name, - COALESCE(NULLIF(m.display_name, ''), NULLIF(m.model_alias, ''), m.model_name), - COALESCE(to_char(p.cooldown_until AT TIME ZONE 'UTC', 'YYYY-MM-DD"T"HH24:MI:SS.MS"Z"'), ''), - GREATEST(COALESCE(EXTRACT(EPOCH FROM p.cooldown_until - now()), 0), 0)::float8, - COALESCE(to_char(m.cooldown_until AT TIME ZONE 'UTC', 'YYYY-MM-DD"T"HH24:MI:SS.MS"Z"'), ''), - GREATEST(COALESCE(EXTRACT(EPOCH FROM m.cooldown_until - now()), 0), 0)::float8 + var code string + var recoveryAt time.Time + var remainingSeconds float64 + err := s.pool.QueryRow(ctx, ` +SELECT CASE + WHEN COALESCE(m.cooldown_until, to_timestamp(0)) >= COALESCE(p.cooldown_until, to_timestamp(0)) + THEN 'model_cooling_down' + ELSE 'platform_cooling_down' + END, + GREATEST(COALESCE(p.cooldown_until, to_timestamp(0)), COALESCE(m.cooldown_until, to_timestamp(0))) AS recovery_at, + GREATEST( + EXTRACT(EPOCH FROM GREATEST(COALESCE(p.cooldown_until, to_timestamp(0)), COALESCE(m.cooldown_until, to_timestamp(0))) - now()), + 0 + )::float8 FROM platform_models m JOIN integration_platforms p ON p.id = m.platform_id LEFT JOIN base_model_catalog b ON b.id = m.base_model_id @@ -756,129 +834,51 @@ WHERE p.status = 'enabled' AND p.deleted_at IS NULL AND m.enabled = true AND m.model_type @> jsonb_build_array($2::text) + AND ( + b.invocation_name = $1::text + OR (b.id IS NULL AND m.model_name = $1::text) + OR EXISTS ( + SELECT 1 + FROM model_compatibility_aliases compatibility_alias + JOIN base_model_catalog alias_base ON alias_base.id = compatibility_alias.base_model_id + WHERE compatibility_alias.alias = $1::text + AND compatibility_alias.model_type = $2::text + AND compatibility_alias.active = true + AND (compatibility_alias.expires_at IS NULL OR compatibility_alias.expires_at > now()) + AND (compatibility_alias.base_model_id = b.id OR alias_base.invocation_name = b.invocation_name) + ) + ) AND ( - ( - $2::text IN ('audio_generate', 'text_to_speech', 'voice_clone') - AND ( - m.model_alias = $1::text - OR m.model_name = $1::text - OR b.canonical_model_key = $1::text - OR b.provider_model_name = $1::text - OR ( - NULLIF($3::text, '') IS NOT NULL - AND ( - regexp_replace(COALESCE(m.model_alias, ''), '[[:space:]]+', '', 'g') = $3::text - OR regexp_replace(COALESCE(m.model_name, ''), '[[:space:]]+', '', 'g') = $3::text - OR regexp_replace(COALESCE(b.canonical_model_key, ''), '[[:space:]]+', '', 'g') = $3::text - OR regexp_replace(COALESCE(b.provider_model_name, ''), '[[:space:]]+', '', 'g') = $3::text - ) - ) - ) - ) - OR ( - $2::text NOT IN ('audio_generate', 'text_to_speech', 'voice_clone') - AND ( - ( - COALESCE(m.model_alias, '') <> '' - AND ( - m.model_alias = $1::text - OR ( - NULLIF($3::text, '') IS NOT NULL - AND regexp_replace(COALESCE(m.model_alias, ''), '[[:space:]]+', '', 'g') = $3::text - ) - ) - ) - OR ( - COALESCE(m.model_alias, '') = '' - AND ( - m.model_name = $1::text - OR b.canonical_model_key = $1::text - OR b.provider_model_name = $1::text - OR ( - NULLIF($3::text, '') IS NOT NULL - AND ( - regexp_replace(COALESCE(m.model_name, ''), '[[:space:]]+', '', 'g') = $3::text - OR regexp_replace(COALESCE(b.canonical_model_key, ''), '[[:space:]]+', '', 'g') = $3::text - OR regexp_replace(COALESCE(b.provider_model_name, ''), '[[:space:]]+', '', 'g') = $3::text - ) - ) - ) - ) - ) - ) + p.cooldown_until > now() + OR m.cooldown_until > now() ) -ORDER BY GREATEST(COALESCE(p.cooldown_until, to_timestamp(0)), COALESCE(m.cooldown_until, to_timestamp(0))) DESC, - p.priority ASC, - m.created_at ASC`, exactModel, modelType, modelMatchKey) +ORDER BY GREATEST(COALESCE(p.cooldown_until, to_timestamp(0)), COALESCE(m.cooldown_until, to_timestamp(0))) ASC, + p.priority ASC, + m.created_at ASC +LIMIT 1`, exactModel, modelType).Scan(&code, &recoveryAt, &remainingSeconds) + if errors.Is(err, pgx.ErrNoRows) { + return nil, nil + } if err != nil { return nil, err } - defer rows.Close() - - for rows.Next() { - var platformName string - var displayName string - var platformCooldownUntil string - var platformRemainingSeconds float64 - var modelCooldownUntil string - var modelRemainingSeconds float64 - if err := rows.Scan( - &platformName, - &displayName, - &platformCooldownUntil, - &platformRemainingSeconds, - &modelCooldownUntil, - &modelRemainingSeconds, - ); err != nil { - return nil, err - } - if modelRemainingSeconds > 0 { - return &ModelCandidateUnavailableError{ - Code: "model_cooling_down", - Message: cooldownErrorMessage("模型", displayName, modelRemainingSeconds, modelCooldownUntil), - }, nil - } - if platformRemainingSeconds > 0 { - return &ModelCandidateUnavailableError{ - Code: "platform_cooling_down", - Message: cooldownErrorMessage("平台", platformName, platformRemainingSeconds, platformCooldownUntil), - }, nil - } + retryAfterSeconds := int(math.Ceil(remainingSeconds)) + if retryAfterSeconds < 1 { + retryAfterSeconds = 1 } - if err := rows.Err(); err != nil { - return nil, err + message := "请求的模型暂时不可用，请稍后重试" + if code == "platform_cooling_down" { + message = "可用平台暂时不可用，请稍后重试" } - return nil, nil -} - -func cooldownErrorMessage(scope string, name string, remainingSeconds float64, cooldownUntil string) string { - name = strings.TrimSpace(name) - if name == "" { - name = "候选" - } - remainingMinutes := remainingSeconds / 60 - if remainingMinutes < 0.1 { - remainingMinutes = 0.1 - } - message := fmt.Sprintf("%s %s 冷却中，剩余 %.1f 分钟", scope, name, remainingMinutes) - if strings.TrimSpace(cooldownUntil) != "" { - message += "，预计恢复时间 " + cooldownUntil - } - return message -} - -func normalizeModelMatchKey(value string) string { - value = strings.TrimSpace(value) - if value == "" { - return "" - } - var builder strings.Builder - builder.Grow(len(value)) - for _, char := range value { - if unicode.IsSpace(char) { - continue - } - builder.WriteRune(char) - } - return builder.String() + recoveryAt = recoveryAt.UTC() + return &ModelCandidateUnavailableError{ + Code: code, + Message: message, + RetryAfter: time.Duration(retryAfterSeconds) * time.Second, + RecoveryAt: recoveryAt, + Details: map[string]any{ + "retryAfterSeconds": retryAfterSeconds, + "recoveryAt": recoveryAt.Format(time.RFC3339), + }, + }, nil } diff --git a/apps/api/internal/store/candidates_test.go b/apps/api/internal/store/candidates_test.go index b66c857..c11689f 100644 --- a/apps/api/internal/store/candidates_test.go +++ b/apps/api/internal/store/candidates_test.go @@ -2,13 +2,6 @@ package store import "testing" -func TestNormalizeModelMatchKeyRemovesWhitespace(t *testing.T) { - got := normalizeModelMatchKey(" doubao-5.0 图像\t编辑\n") - if got != "doubao-5.0图像编辑" { - t.Fatalf("expected whitespace-insensitive model key, got %q", got) - } -} - func TestNormalizeModelTypeListExpandsOmniVideoBaseCapabilities(t *testing.T) { got := normalizeModelTypeList([]string{"omni_video"}) want := StringList{"video_generate", "image_to_video", "omni_video"} @@ -127,7 +120,7 @@ func TestRuntimeCandidateSortingPrefersCacheAffinityWithinAvailableCandidates(t }, { PlatformID: "cache-affinity", - PlatformPriority: 20, + PlatformPriority: 10, }, } applyRuntimeCandidateCacheAffinity(&candidates[0], ListModelCandidatesOptions{ @@ -155,7 +148,7 @@ func TestRuntimeCandidateSortingPrefersCacheAffinityWithinAvailableCandidates(t } } -func TestRuntimeCandidateSortingPromotesObservedCacheHitAbovePriority(t *testing.T) { +func TestRuntimeCandidateSortingKeepsEffectivePriorityAheadOfCacheAffinity(t *testing.T) { candidates := []RuntimeModelCandidate{ { PlatformID: "volces-priority", @@ -184,14 +177,17 @@ func TestRuntimeCandidateSortingPromotesObservedCacheHitAbovePriority(t *testing sortRuntimeModelCandidates(candidates) - if candidates[0].PlatformID != "deepseek-cache-hit" || !candidates[0].CacheAffinity.Applied { - t.Fatalf("expected observed cache hit to outrank platform priority, got %+v", candidates) + if candidates[0].PlatformID != "volces-priority" || candidates[0].CacheAffinity.Applied { + t.Fatalf("expected effective priority tier to remain first, got %+v", candidates) } - if candidates[0].CacheAffinity.Score <= 0 || candidates[0].CacheAffinity.Boost <= 0 { - t.Fatalf("expected observed cache hit to carry score and boost, got %+v", candidates[0].CacheAffinity) + if candidates[0].CacheAffinity.OverrideReason != "different_effective_priority_tier" { + t.Fatalf("expected priority override reason on selected candidate, got %+v", candidates[0].CacheAffinity) } - if candidates[1].PlatformID != "volces-priority" || candidates[1].CacheAffinity.Applied { - t.Fatalf("expected priority-only candidate second, got %+v", candidates) + if candidates[1].PlatformID != "deepseek-cache-hit" || !candidates[1].CacheAffinity.Applied { + t.Fatalf("expected cache affinity candidate to remain visible as fallback, got %+v", candidates) + } + if candidates[1].CacheAffinity.Score <= 0 || candidates[1].CacheAffinity.Boost <= 0 { + t.Fatalf("expected observed cache hit to retain score and boost diagnostics, got %+v", candidates[1].CacheAffinity) } } @@ -203,7 +199,7 @@ func TestRuntimeCandidateSortingOrdersMultipleCacheAffinityCandidatesByScore(t * }, { PlatformID: "higher-cache-score", - PlatformPriority: 50, + PlatformPriority: 1, }, } policy := map[string]any{"enabled": true, "minSamples": 1, "maxPriorityBoost": 20, "modelTypes": []any{"text_generate"}} @@ -267,6 +263,29 @@ func TestRuntimeCandidateSortingKeepsFullCacheAffinityCandidateAvoided(t *testin if candidates[1].PlatformID != "cache-affinity-full" || !candidates[1].LoadAvoided { t.Fatalf("expected full cache-affinity candidate to remain avoided fallback, got %+v", candidates) } + if candidates[0].CacheAffinity.OverrideReason != "capacity_tier_unavailable" { + t.Fatalf("expected capacity override reason on selected candidate, got %+v", candidates[0].CacheAffinity) + } +} + +func TestCacheAffinityMatchedPrefixDepthIgnoresLegacyCompatibilityKeys(t *testing.T) { + lookup := []string{ + "prompt_lcp_v2:depth-four", + "prompt_lcp_v2:depth-three", + "prompt_lcp_v2:depth-two", + "prompt_lcp:depth-four", + "prompt_lcp:depth-three", + "prompt_lcp:depth-two", + } + if depth := cacheAffinityMatchedPrefixDepth("prompt_lcp_v2:depth-three", lookup); depth != 3 { + t.Fatalf("V2 matched prefix depth=%d want=3", depth) + } + if depth := cacheAffinityMatchedPrefixDepth("prompt_lcp:depth-two", lookup); depth != 2 { + t.Fatalf("legacy matched prefix depth=%d want=2", depth) + } + if depth := cacheAffinityMatchedPrefixDepth("explicit:key", lookup); depth != 0 { + t.Fatalf("explicit key prefix depth=%d want=0", depth) + } } func TestDefaultRunnerPriorityDemotePolicyUsesAutoMode(t *testing.T) { diff --git a/apps/api/internal/store/compatibility_tasks.go b/apps/api/internal/store/compatibility_tasks.go index a59d476..b83593d 100644 --- a/apps/api/internal/store/compatibility_tasks.go +++ b/apps/api/internal/store/compatibility_tasks.go @@ -2,7 +2,6 @@ package store import ( "context" - "encoding/json" "strings" ) @@ -16,25 +15,19 @@ type CompatibilitySubmission struct { } func (s *Store) SetTaskCompatibilitySubmission(ctx context.Context, taskID string, submission CompatibilitySubmission) error { - headers, _ := json.Marshal(emptyObjectIfNil(submission.Headers)) - body, _ := json.Marshal(emptyObjectIfNil(submission.Body)) _, err := s.pool.Exec(ctx, ` UPDATE gateway_tasks -SET compatibility_protocol = NULLIF($2, ''), - compatibility_public_id = NULLIF($3, ''), - compatibility_source_protocol = NULLIF($4, ''), - compatibility_submit_http_status = NULLIF($5, 0), - compatibility_submit_headers = $6::jsonb, - compatibility_submit_body = $7::jsonb, +SET compatibility_protocol = COALESCE(NULLIF($2, ''), compatibility_protocol), + compatibility_public_id = NULL, + compatibility_source_protocol = COALESCE(NULLIF($3, ''), compatibility_source_protocol), + compatibility_submit_http_status = NULL, + compatibility_submit_headers = '{}'::jsonb, + compatibility_submit_body = '{}'::jsonb, updated_at = now() WHERE id = $1::uuid`, taskID, strings.TrimSpace(submission.TargetProtocol), - strings.TrimSpace(submission.PublicID), strings.TrimSpace(submission.SourceProtocol), - submission.HTTPStatus, - string(headers), - string(body), ) return err } @@ -44,7 +37,11 @@ func (s *Store) GetCompatibilityTask(ctx context.Context, protocol string, publi SELECT `+gatewayTaskColumns+` FROM gateway_tasks WHERE compatibility_protocol = $1 - AND compatibility_public_id = $2 + AND ( + id::text = $2 + OR remote_task_id = $2 + OR compatibility_public_id = $2 + ) LIMIT 1`, strings.TrimSpace(protocol), strings.TrimSpace(publicID))) if err != nil { return GatewayTask{}, err diff --git a/apps/api/internal/store/conversations.go b/apps/api/internal/store/conversations.go index e686c64..1bf0292 100644 --- a/apps/api/internal/store/conversations.go +++ b/apps/api/internal/store/conversations.go @@ -41,7 +41,7 @@ func (s *Store) EnsureConversation(ctx context.Context, user *auth.User, convers if userID == "" { userID = "anonymous" } - metadataJSON, _ := json.Marshal(emptyObjectIfNil(metadata)) + metadataJSON, _ := json.Marshal(sanitizeJSONForStorage(emptyObjectIfNil(metadata))) var conversationID string err := s.pool.QueryRow(ctx, ` INSERT INTO gateway_conversations (user_id, gateway_user_id, conversation_key, metadata) @@ -67,7 +67,15 @@ func (s *Store) UpsertConversationMessages(ctx context.Context, conversationID s refs := make([]TaskMessageRefInput, 0, len(messages)) newCount := 0 for index, message := range messages { - snapshotJSON, _ := json.Marshal(emptyObjectIfNil(message.Snapshot)) + snapshotReport := sanitizeJSONForStorageWithReport(emptyObjectIfNil(message.Snapshot)) + if snapshotReport.BinaryCount > 0 { + return nil, 0, &taskPayloadBinaryError{ + target: ErrTaskRequestBinaryNotMaterialized, + code: "request_binary_not_materialized", + count: snapshotReport.BinaryCount, + } + } + snapshotJSON, _ := json.Marshal(snapshotReport.Value) var messageID string var inserted bool if err := tx.QueryRow(ctx, ` diff --git a/apps/api/internal/store/file_storage_channels.go b/apps/api/internal/store/file_storage_channels.go index 479a4f5..e8c3987 100644 --- a/apps/api/internal/store/file_storage_channels.go +++ b/apps/api/internal/store/file_storage_channels.go @@ -578,7 +578,7 @@ func normalizeFileStorageScene(scene string) string { } func defaultFileStorageScenes() []string { - return []string{FileStorageSceneUpload, FileStorageSceneImageResult} + return []string{FileStorageSceneUpload, FileStorageSceneImageResult, FileStorageSceneRequestAsset} } func defaultFileStorageRetryPolicyIfEmpty(policy map[string]any) map[string]any { diff --git a/apps/api/internal/store/file_storage_channels_test.go b/apps/api/internal/store/file_storage_channels_test.go new file mode 100644 index 0000000..61dbafb --- /dev/null +++ b/apps/api/internal/store/file_storage_channels_test.go @@ -0,0 +1,17 @@ +package store + +import ( + "reflect" + "testing" +) + +func TestDefaultFileStorageScenesIncludeCrossNodeRequestAssets(t *testing.T) { + want := []string{ + FileStorageSceneUpload, + FileStorageSceneImageResult, + FileStorageSceneRequestAsset, + } + if got := defaultFileStorageScenes(); !reflect.DeepEqual(got, want) { + t.Fatalf("default file storage scenes = %#v, want %#v", got, want) + } +} diff --git a/apps/api/internal/store/json_storage_guard.go b/apps/api/internal/store/json_storage_guard.go new file mode 100644 index 0000000..a8486f0 --- /dev/null +++ b/apps/api/internal/store/json_storage_guard.go @@ -0,0 +1,344 @@ +package store + +import ( + "crypto/sha256" + "encoding/base64" + "encoding/json" + "errors" + "fmt" + "strconv" + "strings" +) + +const ( + storageBinaryPrefixChars = 16 + storageGenericBase64MinLength = 4096 + storageInvalidBinaryMinLength = 512 + storageJSONSanitizerMaxDepth = 64 + storagePlaceholderPrefix = "[GatewayBinary:v1;" + storageBufferObjectType = "buffer" + storageDataURLPrefix = "data:" + storageDataURLBase64Marker = ";base64" + storageDataURLMaxContentType = 64 +) + +var ( + ErrTaskRequestBinaryNotMaterialized = errors.New("task request binary payload was not materialized") + ErrTaskResultBinaryNotMaterialized = errors.New("task result binary payload was not materialized") +) + +type taskPayloadBinaryError struct { + target error + code string + count int +} + +func (e *taskPayloadBinaryError) Error() string { + return fmt.Sprintf("%s: detected %d inline binary value(s)", e.code, e.count) +} + +func (e *taskPayloadBinaryError) ErrorCode() string { + return e.code +} + +func (e *taskPayloadBinaryError) Is(target error) bool { + return target == e.target +} + +type storageSanitizeReport struct { + Value any + BinaryCount int +} + +// sanitizeJSONForStorage is the final task-domain persistence guard. It always +// returns a detached JSON-compatible value and replaces inline binary payloads +// with a bounded, deterministic placeholder. +func sanitizeJSONForStorage(value any) any { + return sanitizeJSONForStorageWithReport(value).Value +} + +func sanitizeJSONForStorageWithReport(value any) storageSanitizeReport { + next, count := sanitizeJSONStorageValue(value, nil, 0) + return storageSanitizeReport{Value: next, BinaryCount: count} +} + +func sanitizeJSONStorageValue(value any, path []string, depth int) (any, int) { + if depth >= storageJSONSanitizerMaxDepth { + return "[JSON,max-depth]", 0 + } + switch typed := value.(type) { + case map[string]any: + if payload, contentType, ok := storageBufferObjectBytes(typed); ok { + return storageBinaryPlaceholder(payload, contentType, "buffer"), 1 + } + next := make(map[string]any, len(typed)) + count := 0 + for key, child := range typed { + sanitized, childCount := sanitizeJSONStorageValue(child, appendStoragePath(path, key), depth+1) + next[key] = sanitized + count += childCount + } + return next, count + case []any: + if storagePathIsBinary(path) { + if payload, ok := storageNumberArrayBytes(typed); ok { + return storageBinaryPlaceholder(payload, "", "buffer"), 1 + } + } + next := make([]any, len(typed)) + count := 0 + for index, child := range typed { + sanitized, childCount := sanitizeJSONStorageValue(child, path, depth+1) + next[index] = sanitized + count += childCount + } + return next, count + case []byte: + return storageBinaryPlaceholder(typed, "", "buffer"), 1 + case string: + if placeholder, ok := storageStringPlaceholder(typed, path); ok { + return placeholder, 1 + } + return typed, 0 + default: + return value, 0 + } +} + +func appendStoragePath(path []string, key string) []string { + next := make([]string, len(path)+1) + copy(next, path) + next[len(path)] = key + return next +} + +func storageStringPlaceholder(value string, path []string) (string, bool) { + raw := strings.TrimSpace(value) + if raw == "" || strings.HasPrefix(raw, storagePlaceholderPrefix) { + return "", false + } + encoded, contentType, dataURL := storageBase64StringParts(raw) + strict := dataURL || storagePathIsBinary(path) + if !strict && len(encoded) < storageGenericBase64MinLength { + return "", false + } + payload, ok := storageDecodeBase64(encoded) + if ok { + encoding := "raw" + if dataURL { + encoding = "data-uri" + } + return storageBinaryPlaceholder(payload, contentType, encoding), true + } + if strict && len(raw) >= storageInvalidBinaryMinLength { + return storageBinaryPlaceholder([]byte(raw), contentType, "raw"), true + } + return "", false +} + +func storageBase64StringParts(value string) (encoded string, contentType string, dataURL bool) { + if !strings.HasPrefix(strings.ToLower(value), storageDataURLPrefix) { + return value, "", false + } + prefix, payload, ok := strings.Cut(value, ",") + if !ok || !strings.Contains(strings.ToLower(prefix), storageDataURLBase64Marker) { + return value, "", false + } + mediaType := strings.TrimSpace(prefix[len(storageDataURLPrefix):]) + if before, _, found := strings.Cut(mediaType, ";"); found { + mediaType = before + } + if len(mediaType) > storageDataURLMaxContentType || !storageSafeContentType(mediaType) { + mediaType = "" + } + return payload, strings.ToLower(mediaType), true +} + +func storageSafeContentType(value string) bool { + if value == "" { + return true + } + for _, char := range value { + switch { + case char >= 'a' && char <= 'z': + case char >= 'A' && char <= 'Z': + case char >= '0' && char <= '9': + case char == '/', char == '.', char == '+', char == '-': + default: + return false + } + } + return true +} + +func storageDecodeBase64(value string) ([]byte, bool) { + normalized := removeStorageASCIIWhitespace(value) + if normalized == "" { + return nil, false + } + for _, encoding := range []*base64.Encoding{ + base64.StdEncoding, + base64.RawStdEncoding, + base64.URLEncoding, + base64.RawURLEncoding, + } { + payload, err := encoding.DecodeString(normalized) + if err != nil || len(payload) == 0 { + continue + } + canonical := encoding.EncodeToString(payload) + if strings.TrimRight(normalized, "=") == strings.TrimRight(canonical, "=") { + return payload, true + } + } + return nil, false +} + +func removeStorageASCIIWhitespace(value string) string { + return strings.Map(func(char rune) rune { + switch char { + case ' ', '\n', '\r', '\t': + return -1 + default: + return char + } + }, value) +} + +func storagePathIsBinary(path []string) bool { + if len(path) == 0 { + return false + } + key := normalizeStorageBinaryKey(path[len(path)-1]) + if storageBinaryKey(key) { + return true + } + if len(path) < 2 { + return false + } + parent := normalizeStorageBinaryKey(path[len(path)-2]) + return (parent == "inlinedata" || parent == "binary" || parent == "media") && + (key == "data" || key == "content") +} + +func normalizeStorageBinaryKey(value string) string { + return strings.Map(func(char rune) rune { + switch { + case char >= 'a' && char <= 'z': + return char + case char >= 'A' && char <= 'Z': + return char + ('a' - 'A') + case char >= '0' && char <= '9': + return char + default: + return -1 + } + }, value) +} + +func storageBinaryKey(key string) bool { + return key == "b64" || + key == "b64json" || + key == "base64" || + key == "buffer" || + key == "bytes" || + strings.Contains(key, "base64") || + strings.Contains(key, "buffer") || + strings.Contains(key, "binary") || + strings.HasSuffix(key, "b64") || + strings.HasSuffix(key, "bytes") +} + +func storageBufferObjectBytes(value map[string]any) ([]byte, string, bool) { + kind, _ := value["type"].(string) + if normalizeStorageBinaryKey(kind) != storageBufferObjectType { + return nil, "", false + } + contentType := firstNonEmpty( + stringFromAny(value["mime_type"]), + stringFromAny(value["mimeType"]), + stringFromAny(value["contentType"]), + ) + switch data := value["data"].(type) { + case []byte: + if len(data) == 0 { + return nil, "", false + } + return append([]byte(nil), data...), contentType, true + case []any: + payload, ok := storageNumberArrayBytes(data) + return payload, contentType, ok + default: + return nil, "", false + } +} + +func storageNumberArrayBytes(values []any) ([]byte, bool) { + if len(values) == 0 { + return nil, false + } + payload := make([]byte, len(values)) + for index, value := range values { + next, ok := storageByteFromAny(value) + if !ok { + return nil, false + } + payload[index] = next + } + return payload, true +} + +func storageByteFromAny(value any) (byte, bool) { + switch typed := value.(type) { + case byte: + return typed, true + case int: + if typed >= 0 && typed <= 255 { + return byte(typed), true + } + case int32: + if typed >= 0 && typed <= 255 { + return byte(typed), true + } + case int64: + if typed >= 0 && typed <= 255 { + return byte(typed), true + } + case float64: + asInt := int(typed) + if typed == float64(asInt) && asInt >= 0 && asInt <= 255 { + return byte(asInt), true + } + case json.Number: + asInt, err := strconv.ParseInt(string(typed), 10, 16) + if err == nil && asInt >= 0 && asInt <= 255 { + return byte(asInt), true + } + } + return 0, false +} + +func storageBinaryPlaceholder(payload []byte, contentType string, encoding string) string { + digest := sha256.Sum256(payload) + prefix := base64.StdEncoding.EncodeToString(payload) + if len(prefix) > storageBinaryPrefixChars { + prefix = prefix[:storageBinaryPrefixChars] + } + contentType = strings.ToLower(strings.TrimSpace(contentType)) + if contentType == "" || len(contentType) > 32 || !storageSafeContentType(contentType) { + contentType = "application/octet-stream" + } + switch encoding { + case "data-uri", "buffer": + default: + encoding = "raw" + } + return fmt.Sprintf( + "[GatewayBinary:v1;prefix=%s;sha256=%x;bytes=%d;mime=%s;encoding=%s]", + prefix, + digest, + len(payload), + contentType, + encoding, + ) +} diff --git a/apps/api/internal/store/json_storage_guard_test.go b/apps/api/internal/store/json_storage_guard_test.go new file mode 100644 index 0000000..4bec55d --- /dev/null +++ b/apps/api/internal/store/json_storage_guard_test.go @@ -0,0 +1,103 @@ +package store + +import ( + "encoding/base64" + "encoding/json" + "strings" + "testing" +) + +func TestSanitizeJSONForStorageReplacesBinaryWithoutMutatingInput(t *testing.T) { + payload := []byte("shared binary payload") + encoded := base64.StdEncoding.EncodeToString(payload) + input := map[string]any{ + "data": []any{ + map[string]any{"b64_json": encoded}, + }, + "buffer": map[string]any{ + "type": "Buffer", + "data": []any{float64(1), float64(2), float64(3)}, + }, + "text": "ordinary text", + } + + report := sanitizeJSONForStorageWithReport(input) + if report.BinaryCount != 2 { + t.Fatalf("binary count = %d, want 2", report.BinaryCount) + } + next := report.Value.(map[string]any) + data := next["data"].([]any) + placeholder := data[0].(map[string]any)["b64_json"].(string) + if !strings.HasPrefix(placeholder, storagePlaceholderPrefix) { + t.Fatalf("unexpected placeholder: %q", placeholder) + } + if len(placeholder) > 200 { + t.Fatalf("placeholder exceeds 200 bytes: %d", len(placeholder)) + } + if !strings.Contains(placeholder, ";prefix="+encoded[:16]+";") { + t.Fatalf("placeholder should retain the bounded Base64 prefix: %q", placeholder) + } + if got := input["data"].([]any)[0].(map[string]any)["b64_json"]; got != encoded { + t.Fatalf("sanitizer mutated input: %v", got) + } + if next["text"] != "ordinary text" { + t.Fatalf("ordinary text changed: %+v", next) + } +} + +func TestSanitizeJSONForStorageUsesDecodedBytesForEquivalentRepresentations(t *testing.T) { + payload := []byte("equivalent payload") + encoded := base64.StdEncoding.EncodeToString(payload) + report := sanitizeJSONForStorageWithReport(map[string]any{ + "rawBase64": encoded, + "dataURL": "data:image/png;base64," + encoded, + "bytes": []any{float64('e'), float64('q'), float64('u'), float64('i'), float64('v'), float64('a'), float64('l'), float64('e'), float64('n'), float64('t'), float64(' '), float64('p'), float64('a'), float64('y'), float64('l'), float64('o'), float64('a'), float64('d')}, + }) + if report.BinaryCount != 3 { + t.Fatalf("binary count = %d, want 3", report.BinaryCount) + } + next := report.Value.(map[string]any) + hashes := map[string]struct{}{} + for _, key := range []string{"rawBase64", "dataURL", "bytes"} { + value := next[key].(string) + hashStart := strings.Index(value, ";sha256=") + hashEnd := strings.Index(value[hashStart+1:], ";bytes=") + if hashStart < 0 || hashEnd < 0 { + t.Fatalf("missing hash in %s placeholder: %q", key, value) + } + hashes[value[hashStart+8:hashStart+1+hashEnd]] = struct{}{} + } + if len(hashes) != 1 { + t.Fatalf("equivalent binary values produced different hashes: %+v", next) + } +} + +func TestSanitizeJSONForStorageAvoidsShortGenericTextFalsePositive(t *testing.T) { + input := map[string]any{"message": "YWJjZA==", "count": json.Number("12")} + report := sanitizeJSONForStorageWithReport(input) + if report.BinaryCount != 0 { + t.Fatalf("short generic Base64-like text should not be sanitized: %+v", report) + } +} + +func TestSanitizeJSONForStorageRecognizesExistingPlaceholder(t *testing.T) { + value := storageBinaryPlaceholder([]byte("payload"), "image/png", "raw") + report := sanitizeJSONForStorageWithReport(map[string]any{"b64_json": value}) + if report.BinaryCount != 0 { + t.Fatalf("existing placeholder should pass the final guard: %+v", report) + } +} + +func TestStorageBinaryPlaceholderIsAlwaysBounded(t *testing.T) { + value := storageBinaryPlaceholder( + []byte("payload"), + "application/vnd.a-very-long-provider-specific-generated-binary-result+json", + "data-uri", + ) + if len(value) > 200 { + t.Fatalf("placeholder exceeds 200 bytes: %d %q", len(value), value) + } + if !strings.Contains(value, ";mime=application/octet-stream;") { + t.Fatalf("oversized MIME should use the bounded fallback: %q", value) + } +} diff --git a/apps/api/internal/store/model_identity_integration_test.go b/apps/api/internal/store/model_identity_integration_test.go new file mode 100644 index 0000000..ce47c0b --- /dev/null +++ b/apps/api/internal/store/model_identity_integration_test.go @@ -0,0 +1,119 @@ +package store + +import ( + "context" + "errors" + "os" + "strings" + "testing" + "time" +) + +func TestModelIdentityRoutingAndDeleteProtection(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 model identity PostgreSQL integration tests") + } + + ctx := context.Background() + db, err := Connect(ctx, databaseURL) + if err != nil { + t.Fatalf("connect test database: %v", err) + } + defer db.Close() + + suffix := time.Now().UTC().Format("20060102150405.000000000") + invocationName := "official-model-" + suffix + legacyAlias := "Legacy Model " + suffix + displayName := "Model Card " + suffix + baseModel, err := db.CreateBaseModel(ctx, BaseModelInput{ + ProviderKey: "openai", + CanonicalModelKey: "openai:" + invocationName, + InvocationName: invocationName, + ProviderModelName: "provider/" + invocationName, + ModelType: StringList{"text_generate"}, + DisplayName: displayName, + LegacyAliases: StringList{legacyAlias}, + Status: "active", + }) + if err != nil { + t.Fatalf("create base model: %v", err) + } + defer func() { + _, _ = db.pool.Exec(ctx, `DELETE FROM platform_models WHERE base_model_id = $1::uuid`, baseModel.ID) + _, _ = db.pool.Exec(ctx, `DELETE FROM base_model_catalog WHERE id = $1::uuid`, baseModel.ID) + }() + + var platformID string + if err := db.pool.QueryRow(ctx, ` +SELECT id::text +FROM integration_platforms +WHERE status = 'enabled' AND deleted_at IS NULL +ORDER BY created_at ASC +LIMIT 1`).Scan(&platformID); err != nil { + t.Fatalf("find enabled platform: %v", err) + } + var platformModelID string + if err := db.pool.QueryRow(ctx, ` +INSERT INTO platform_models ( + platform_id, base_model_id, model_name, provider_model_name, model_alias, + model_type, display_name, enabled +) +VALUES ($1::uuid, $2::uuid, $3, $4, $3, '["text_generate"]'::jsonb, $5, true) +RETURNING id::text`, platformID, baseModel.ID, invocationName, "provider/"+invocationName, displayName).Scan(&platformModelID); err != nil { + t.Fatalf("bind platform model: %v", err) + } + + if err := db.DeleteBaseModel(ctx, baseModel.ID); !errors.Is(err, ErrBaseModelInUse) { + t.Fatalf("bound base model deletion should be protected: %v", err) + } + + officialCandidates, err := db.ListModelCandidates(ctx, invocationName, "text_generate", nil) + if err != nil || len(officialCandidates) != 1 || officialCandidates[0].PlatformModelID != platformModelID { + t.Fatalf("official invocation should resolve the bound source: candidates=%+v err=%v", officialCandidates, err) + } + legacyCandidates, err := db.ListModelCandidates(ctx, legacyAlias, "text_generate", nil) + if err != nil || len(legacyCandidates) != 1 || !legacyCandidates[0].LegacyAliasUsed { + t.Fatalf("registered legacy alias should resolve and be marked: candidates=%+v err=%v", legacyCandidates, err) + } + if _, err := db.ListModelCandidates(ctx, displayName, "text_generate", nil); !errors.Is(err, ErrNoModelCandidate) { + t.Fatalf("display name must not become an invocation alias: %v", err) + } + + var legacyHits int64 + if err := db.pool.QueryRow(ctx, ` +SELECT hit_count +FROM model_alias_usage_metrics +WHERE alias = $1 AND canonical_model_key = $2 AND model_type = 'text_generate'`, legacyAlias, baseModel.CanonicalModelKey).Scan(&legacyHits); err != nil { + t.Fatalf("read legacy alias metric: %v", err) + } + if legacyHits != 1 { + t.Fatalf("expected one legacy alias hit, got %d", legacyHits) + } + + deprecatedInvocation := "deprecated-model-" + suffix + deprecatedBase, err := db.CreateBaseModel(ctx, BaseModelInput{ + ProviderKey: "openai", + CanonicalModelKey: "openai:" + deprecatedInvocation, + InvocationName: deprecatedInvocation, + ProviderModelName: deprecatedInvocation, + ModelType: StringList{"text_generate"}, + DisplayName: "Deprecated Model", + Status: "deprecated", + }) + if err != nil { + t.Fatalf("create deprecated base model: %v", err) + } + defer func() { + _, _ = db.pool.Exec(ctx, `DELETE FROM base_model_catalog WHERE id = $1::uuid`, deprecatedBase.ID) + }() + if _, err := db.CreatePlatformModel(ctx, CreatePlatformModelInput{ + PlatformID: platformID, + BaseModelID: deprecatedBase.ID, + ModelName: deprecatedInvocation, + ProviderModelName: deprecatedInvocation, + ModelType: StringList{"text_generate"}, + }); !errors.Is(err, ErrInvalidPlatformModelConfiguration) { + t.Fatalf("deprecated base model must reject new bindings: %v", err) + } +} diff --git a/apps/api/internal/store/oidc_users_integration_test.go b/apps/api/internal/store/oidc_users_integration_test.go index 3a3aff7..9fd0e09 100644 --- a/apps/api/internal/store/oidc_users_integration_test.go +++ b/apps/api/internal/store/oidc_users_integration_test.go @@ -307,6 +307,22 @@ func applyOIDCJITTestMigrations(t *testing.T, ctx context.Context, databaseURL s if err != nil { t.Fatalf("read migration %s: %v", version, 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", version, err) + } + } + } + if _, err := pool.Exec(ctx, `INSERT INTO schema_migrations(version) VALUES($1)`, version); err != nil { + t.Fatalf("record non-transaction migration %s: %v", version, err) + } + continue + } tx, err := pool.Begin(ctx) if err != nil { t.Fatalf("begin migration %s: %v", version, err) diff --git a/apps/api/internal/store/platform_models.go b/apps/api/internal/store/platform_models.go index 9eb70ca..7183ad2 100644 --- a/apps/api/internal/store/platform_models.go +++ b/apps/api/internal/store/platform_models.go @@ -18,6 +18,7 @@ type modelCatalogSnapshot struct { ID string ProviderKey string CanonicalModelKey string + InvocationName string ProviderModelName string ModelType StringList DisplayName string @@ -27,6 +28,7 @@ type modelCatalogSnapshot struct { DefaultRateLimitPolicy map[string]any RuntimePolicySetID string RuntimePolicyOverride map[string]any + Status string } func (s *Store) CreatePlatformModel(ctx context.Context, input CreatePlatformModelInput) (PlatformModel, error) { @@ -95,6 +97,18 @@ func (s *Store) createPlatformModel(ctx context.Context, q platformModelQuerier, if err != nil && !IsNotFound(err) { return PlatformModel{}, err } + if base.ID != "" && base.Status != "" && base.Status != "active" { + var existingBinding bool + if err := q.QueryRow(ctx, ` +SELECT EXISTS ( + SELECT 1 FROM platform_models WHERE platform_id = $1::uuid AND base_model_id = $2::uuid +)`, input.PlatformID, base.ID).Scan(&existingBinding); err != nil { + return PlatformModel{}, err + } + if !existingBinding { + return PlatformModel{}, fmt.Errorf("%w: base model %q is %s and cannot be newly bound", ErrInvalidPlatformModelConfiguration, base.InvocationName, base.Status) + } + } if len(input.ModelType) == 0 { input.ModelType = base.ModelType } @@ -102,7 +116,9 @@ func (s *Store) createPlatformModel(ctx context.Context, q platformModelQuerier, if len(input.ModelType) == 0 { input.ModelType = StringList{"text_generate"} } - if input.ModelName == "" { + if base.InvocationName != "" { + input.ModelName = base.InvocationName + } else if input.ModelName == "" { input.ModelName = base.ProviderModelName } if input.ProviderModelName == "" { @@ -127,18 +143,13 @@ func (s *Store) createPlatformModel(ctx context.Context, q platformModelQuerier, // soon as the base pricing rule changes and can mask the authoritative rule. billingConfig := input.BillingConfig explicitRuntimePolicySetID := strings.TrimSpace(input.RuntimePolicySetID) + rateLimitPolicyMode := NormalizeRateLimitPolicyMode(input.RateLimitPolicyMode, input.RateLimitPolicy != nil) rateLimitPolicy := input.RateLimitPolicy - if len(rateLimitPolicy) == 0 && explicitRuntimePolicySetID == "" { - rateLimitPolicy = base.DefaultRateLimitPolicy + if rateLimitPolicyMode == RateLimitPolicyModeInherit { + rateLimitPolicy = nil } runtimePolicySetID := explicitRuntimePolicySetID - if runtimePolicySetID == "" { - runtimePolicySetID = base.RuntimePolicySetID - } runtimePolicyOverride := input.RuntimePolicyOverride - if len(runtimePolicyOverride) == 0 { - runtimePolicyOverride = base.RuntimePolicyOverride - } capabilityOverrideJSON, _ := json.Marshal(emptyObjectIfNil(input.CapabilityOverride)) capabilitiesJSON, _ := json.Marshal(emptyObjectIfNil(capabilities)) @@ -173,14 +184,14 @@ func (s *Store) createPlatformModel(ctx context.Context, q platformModelQuerier, INSERT INTO platform_models ( platform_id, base_model_id, model_name, provider_model_name, model_alias, model_type, display_name, capability_override, capabilities, pricing_mode, discount_factor, - pricing_rule_set_id, billing_config_override, billing_config, permission_config, retry_policy, rate_limit_policy, + pricing_rule_set_id, billing_config_override, billing_config, permission_config, retry_policy, rate_limit_policy, rate_limit_policy_mode, runtime_policy_set_id, runtime_policy_override, enabled ) VALUES ( $1::uuid, $2::uuid, $3, NULLIF($4, ''), NULLIF($5, ''), $6::jsonb, $7, $8::jsonb, $9::jsonb, $10, $11::numeric, - NULLIF($12, '')::uuid, $13::jsonb, $14::jsonb, $15::jsonb, $16::jsonb, $17::jsonb, - NULLIF($18, '')::uuid, $19::jsonb, true + NULLIF($12, '')::uuid, $13::jsonb, $14::jsonb, $15::jsonb, $16::jsonb, $17::jsonb, $18, + NULLIF($19, '')::uuid, $20::jsonb, true ) ON CONFLICT (platform_id, model_name) DO UPDATE SET base_model_id = EXCLUDED.base_model_id, @@ -197,6 +208,7 @@ SET base_model_id = EXCLUDED.base_model_id, permission_config = EXCLUDED.permission_config, retry_policy = EXCLUDED.retry_policy, rate_limit_policy = EXCLUDED.rate_limit_policy, + rate_limit_policy_mode = EXCLUDED.rate_limit_policy_mode, runtime_policy_set_id = EXCLUDED.runtime_policy_set_id, runtime_policy_override = EXCLUDED.runtime_policy_override, enabled = true, @@ -205,7 +217,7 @@ RETURNING id::text, platform_id::text, COALESCE(base_model_id::text, ''), model_ COALESCE(NULLIF(provider_model_name, ''), model_name), COALESCE(model_alias, ''), model_type, display_name, capability_override, capabilities, pricing_mode, COALESCE(discount_factor, 0)::float8, COALESCE(pricing_rule_set_id::text, ''), billing_config_override, billing_config, - permission_config, retry_policy, rate_limit_policy, COALESCE(runtime_policy_set_id::text, ''), runtime_policy_override, + permission_config, retry_policy, rate_limit_policy, rate_limit_policy_mode, COALESCE(runtime_policy_set_id::text, ''), runtime_policy_override, COALESCE(to_char(cooldown_until AT TIME ZONE 'UTC', 'YYYY-MM-DD"T"HH24:MI:SS.MS"Z"'), ''), enabled, created_at, updated_at`, input.PlatformID, @@ -225,6 +237,7 @@ RETURNING id::text, platform_id::text, COALESCE(base_model_id::text, ''), model_ string(permissionJSON), string(retryJSON), string(rateLimitJSON), + rateLimitPolicyMode, runtimePolicySetID, string(runtimePolicyOverrideJSON), ).Scan( @@ -246,6 +259,7 @@ RETURNING id::text, platform_id::text, COALESCE(base_model_id::text, ''), model_ &permissionBytes, &retryPolicyBytes, &rateLimitPolicyBytes, + &model.RateLimitPolicyMode, &model.RuntimePolicySetID, &runtimePolicyOverrideBytes, &model.CooldownUntil, @@ -370,9 +384,9 @@ func (s *Store) lookupBaseModel(ctx context.Context, q platformModelQuerier, id var runtimePolicyOverride []byte var modelTypeBytes []byte err := q.QueryRow(ctx, ` -SELECT id::text, provider_key, canonical_model_key, provider_model_name, model_type, display_name, +SELECT id::text, provider_key, canonical_model_key, invocation_name, provider_model_name, model_type, display_name, capabilities, base_billing_config, COALESCE(pricing_rule_set_id::text, ''), default_rate_limit_policy, - COALESCE(runtime_policy_set_id::text, ''), runtime_policy_override + COALESCE(runtime_policy_set_id::text, ''), runtime_policy_override, status FROM base_model_catalog WHERE ($1 <> '' AND id = NULLIF($1, '')::uuid) OR ($2 <> '' AND canonical_model_key = $2) @@ -382,6 +396,7 @@ LIMIT 1`, strings.TrimSpace(id), strings.TrimSpace(canonicalKey), strings.TrimSp &item.ID, &item.ProviderKey, &item.CanonicalModelKey, + &item.InvocationName, &item.ProviderModelName, &modelTypeBytes, &item.DisplayName, @@ -391,6 +406,7 @@ LIMIT 1`, strings.TrimSpace(id), strings.TrimSpace(canonicalKey), strings.TrimSp &rateLimitPolicy, &item.RuntimePolicySetID, &runtimePolicyOverride, + &item.Status, ) if err != nil { if err == pgx.ErrNoRows { @@ -407,15 +423,8 @@ LIMIT 1`, strings.TrimSpace(id), strings.TrimSpace(canonicalKey), strings.TrimSp } func normalizePlatformModelAlias(alias string, base modelCatalogSnapshot) string { - alias = strings.TrimSpace(alias) - if alias == "" { - alias = firstNonEmpty(base.ProviderModelName, base.DisplayName, base.CanonicalModelKey) - } - if base.ProviderKey != "" { - alias = strings.TrimPrefix(alias, base.ProviderKey+":") - } - if alias == base.CanonicalModelKey { - alias = stripAliasPrefix(alias) + if base.InvocationName != "" { + return base.InvocationName } return strings.TrimSpace(alias) } diff --git a/apps/api/internal/store/postgres.go b/apps/api/internal/store/postgres.go index a587f7d..63ae9e1 100644 --- a/apps/api/internal/store/postgres.go +++ b/apps/api/internal/store/postgres.go @@ -69,6 +69,8 @@ var ( ErrTaskExecutionFinished = errors.New("task execution already finished") ErrTaskExecutionManualReview = errors.New("task execution requires manual review") ErrProtectedDefault = errors.New("protected default resource cannot be deleted") + ErrBaseModelInUse = errors.New("base model is still referenced by platform models") + ErrModelAliasConflict = errors.New("model compatibility alias conflicts with another model identity") ErrUserAlreadyExists = errors.New("user already exists") ErrWeakPassword = errors.New("password must be at least 8 characters") ) @@ -217,20 +219,24 @@ type CreatedAPIKey struct { } type PlatformModel struct { - ID string `json:"id"` - PlatformID string `json:"platformId"` - BaseModelID string `json:"baseModelId,omitempty"` - Provider string `json:"provider,omitempty"` - PlatformName string `json:"platformName,omitempty"` - ModelName string `json:"modelName"` - ProviderModelName string `json:"providerModelName,omitempty"` + ID string `json:"id"` + PlatformID string `json:"platformId"` + BaseModelID string `json:"baseModelId,omitempty"` + Provider string `json:"provider,omitempty"` + PlatformName string `json:"platformName,omitempty"` + ModelName string `json:"modelName"` + ProviderModelName string `json:"providerModelName,omitempty"` + // Deprecated: use ModelName. This field mirrors the canonical invocation name during the compatibility window. ModelAlias string `json:"modelAlias,omitempty"` + LegacyAliases StringList `json:"legacyAliases,omitempty"` ModelType StringList `json:"modelType"` DisplayName string `json:"displayName"` CapabilityOverride map[string]any `json:"capabilityOverride,omitempty"` Capabilities map[string]any `json:"capabilities,omitempty"` BaseCapabilities map[string]any `json:"-"` BaseBillingConfig map[string]any `json:"-"` + BaseRateLimitPolicy map[string]any `json:"-"` + BaseRuntimePolicySetID string `json:"-"` BasePricingRuleSetID string `json:"-"` PlatformPricingRuleSetID string `json:"-"` PricingMode string `json:"pricingMode"` @@ -241,6 +247,7 @@ type PlatformModel struct { PermissionConfig map[string]any `json:"permissionConfig,omitempty"` RetryPolicy map[string]any `json:"retryPolicy,omitempty"` RateLimitPolicy map[string]any `json:"rateLimitPolicy,omitempty"` + RateLimitPolicyMode string `json:"rateLimitPolicyMode" enums:"inherit,override"` RuntimePolicySetID string `json:"runtimePolicySetId,omitempty"` RuntimePolicyOverride map[string]any `json:"runtimePolicyOverride,omitempty"` CooldownUntil string `json:"cooldownUntil,omitempty"` @@ -284,13 +291,17 @@ type CatalogProvider struct { } type BaseModel struct { - ID string `json:"id"` - ProviderKey string `json:"providerKey"` - CanonicalModelKey string `json:"canonicalModelKey"` - ProviderModelName string `json:"providerModelName"` - ModelType StringList `json:"modelType"` + ID string `json:"id"` + ProviderKey string `json:"providerKey"` + CanonicalModelKey string `json:"canonicalModelKey"` + InvocationName string `json:"invocationName"` + ProviderModelName string `json:"providerModelName"` + ModelType StringList `json:"modelType"` + // Deprecated: use InvocationName. ModelAlias string `json:"modelAlias"` - DisplayName string `json:"-"` + DisplayName string `json:"displayName"` + LegacyAliases StringList `json:"legacyAliases,omitempty"` + ReferenceCount int `json:"referenceCount"` Capabilities map[string]any `json:"capabilities,omitempty"` BaseBillingConfig map[string]any `json:"baseBillingConfig,omitempty"` DefaultRateLimitPolicy map[string]any `json:"defaultRateLimitPolicy,omitempty"` @@ -576,17 +587,19 @@ COALESCE(compatibility_submit_headers, '{}'::jsonb), COALESCE(compatibility_subm created_at, updated_at, COALESCE(finished_at::text, '')` type TaskEvent struct { - ID string `json:"id"` - TaskID string `json:"taskId"` - Seq int64 `json:"seq"` - EventType string `json:"eventType"` - Status string `json:"status,omitempty"` - Phase string `json:"phase,omitempty"` - Progress float64 `json:"progress,omitempty"` - Message string `json:"message,omitempty"` - Payload map[string]any `json:"payload,omitempty"` - Simulated bool `json:"simulated"` - CreatedAt time.Time `json:"createdAt"` + ID string `json:"id"` + TaskID string `json:"taskId"` + PlatformID string `json:"-"` + SkippedReason string `json:"-"` + Seq int64 `json:"seq"` + EventType string `json:"eventType"` + Status string `json:"status,omitempty"` + Phase string `json:"phase,omitempty"` + Progress float64 `json:"progress,omitempty"` + Message string `json:"message,omitempty"` + Payload map[string]any `json:"payload,omitempty"` + Simulated bool `json:"simulated"` + CreatedAt time.Time `json:"createdAt"` } type TaskAttempt struct { @@ -938,18 +951,34 @@ func (s *Store) listModels(ctx context.Context, platformID string) ([]PlatformMo rows, err := s.pool.Query(ctx, ` SELECT m.id::text, m.platform_id::text, COALESCE(m.base_model_id::text, ''), p.provider, p.name, - m.model_name, COALESCE(NULLIF(m.provider_model_name, ''), m.model_name), COALESCE(m.model_alias, ''), m.model_type, m.display_name, + COALESCE(NULLIF(b.invocation_name, ''), m.model_name), + COALESCE(NULLIF(m.provider_model_name, ''), m.model_name), + COALESCE(NULLIF(b.invocation_name, ''), NULLIF(m.model_alias, ''), m.model_name), + COALESCE(b.model_type, m.model_type), + COALESCE(NULLIF(b.display_name, ''), NULLIF(m.display_name, ''), COALESCE(NULLIF(b.invocation_name, ''), m.model_name)), + COALESCE(b.legacy_aliases, '[]'::jsonb), m.capability_override, m.capabilities, COALESCE(b.capabilities, '{}'::jsonb), COALESCE(b.base_billing_config, '{}'::jsonb), COALESCE(b.pricing_rule_set_id::text, ''), COALESCE(p.pricing_rule_set_id::text, ''), m.pricing_mode, COALESCE(m.discount_factor, 0)::float8, COALESCE(m.pricing_rule_set_id::text, ''), m.billing_config_override, m.billing_config, - m.permission_config, m.retry_policy, m.rate_limit_policy, COALESCE(m.runtime_policy_set_id::text, ''), m.runtime_policy_override, + m.permission_config, m.retry_policy, m.rate_limit_policy, m.rate_limit_policy_mode, + COALESCE(m.runtime_policy_set_id::text, ''), m.runtime_policy_override, + COALESCE(b.default_rate_limit_policy, '{}'::jsonb), COALESCE(b.runtime_policy_set_id::text, ''), COALESCE(to_char(m.cooldown_until AT TIME ZONE 'UTC', 'YYYY-MM-DD"T"HH24:MI:SS.MS"Z"'), ''), m.enabled, m.created_at, m.updated_at FROM platform_models m JOIN integration_platforms p ON p.id = m.platform_id LEFT JOIN LATERAL ( - SELECT catalog.capabilities, catalog.base_billing_config, catalog.pricing_rule_set_id + SELECT catalog.invocation_name, catalog.display_name, catalog.model_type, + catalog.capabilities, catalog.base_billing_config, catalog.pricing_rule_set_id, + catalog.default_rate_limit_policy, catalog.runtime_policy_set_id, + COALESCE(( + SELECT jsonb_agg(DISTINCT compatibility_alias.alias ORDER BY compatibility_alias.alias) + FROM model_compatibility_aliases compatibility_alias + WHERE compatibility_alias.base_model_id = catalog.id + AND compatibility_alias.active = true + AND (compatibility_alias.expires_at IS NULL OR compatibility_alias.expires_at > now()) + ), '[]'::jsonb) AS legacy_aliases FROM base_model_catalog catalog WHERE (m.base_model_id IS NOT NULL AND catalog.id = m.base_model_id) OR ( @@ -983,7 +1012,9 @@ ORDER BY m.model_type ASC, m.model_name ASC`, args...) var retryPolicy []byte var rateLimitPolicy []byte var runtimePolicyOverride []byte + var baseRateLimitPolicy []byte var modelTypeBytes []byte + var legacyAliasesBytes []byte if err := rows.Scan( &model.ID, &model.PlatformID, @@ -995,6 +1026,7 @@ ORDER BY m.model_type ASC, m.model_name ASC`, args...) &model.ModelAlias, &modelTypeBytes, &model.DisplayName, + &legacyAliasesBytes, &capabilityOverride, &capabilities, &baseCapabilities, @@ -1009,8 +1041,11 @@ ORDER BY m.model_type ASC, m.model_name ASC`, args...) &permissionConfig, &retryPolicy, &rateLimitPolicy, + &model.RateLimitPolicyMode, &model.RuntimePolicySetID, &runtimePolicyOverride, + &baseRateLimitPolicy, + &model.BaseRuntimePolicySetID, &model.CooldownUntil, &model.Enabled, &model.CreatedAt, @@ -1023,12 +1058,14 @@ ORDER BY m.model_type ASC, m.model_name ASC`, args...) model.BaseCapabilities = decodeObject(baseCapabilities) model.BaseBillingConfig = decodeObject(baseBillingConfig) model.ModelType = decodeStringArray(modelTypeBytes) + model.LegacyAliases = decodeStringArray(legacyAliasesBytes) model.BillingConfigOverride = decodeObject(billingConfigOverride) model.BillingConfig = decodeObject(billingConfig) model.PermissionConfig = decodeObject(permissionConfig) model.RetryPolicy = decodeObject(retryPolicy) model.RateLimitPolicy = decodeObject(rateLimitPolicy) model.RuntimePolicyOverride = decodeObject(runtimePolicyOverride) + model.BaseRateLimitPolicy = decodeObject(baseRateLimitPolicy) models = append(models, model) } return models, rows.Err() @@ -1956,7 +1993,15 @@ func (s *Store) CreateTask(ctx context.Context, input CreateTaskInput, user *aut } func (s *Store) CreateTaskIdempotent(ctx context.Context, input CreateTaskInput, user *auth.User) (CreateTaskResult, error) { - requestBody, _ := json.Marshal(input.Request) + requestReport := sanitizeJSONForStorageWithReport(input.Request) + if requestReport.BinaryCount > 0 { + return CreateTaskResult{}, &taskPayloadBinaryError{ + target: ErrTaskRequestBinaryNotMaterialized, + code: "request_binary_not_materialized", + count: requestReport.BinaryCount, + } + } + requestBody, _ := json.Marshal(requestReport.Value) runMode := normalizeRunMode(input.RunMode, input.Request) status := "queued" resultBody, _ := json.Marshal(map[string]any(nil)) @@ -2011,7 +2056,7 @@ WHERE user_id = $1 AND idempotency_key_hash = $2`, user.ID, strings.TrimSpace(in } events := taskEventsForCreate(task.ID, runMode, status, nil) for _, event := range events { - payload, _ := json.Marshal(event.Payload) + payload, _ := json.Marshal(sanitizeJSONForStorage(event.Payload)) if _, err := tx.Exec(ctx, ` INSERT INTO gateway_task_events (task_id, seq, event_type, status, phase, progress, message, payload, simulated) VALUES ($1::uuid, $2, $3::text, NULLIF($4::text, ''), NULLIF($5::text, ''), $6, NULLIF($7::text, ''), $8::jsonb, $9)`, @@ -2146,7 +2191,8 @@ func scanGatewayTask(scanner taskScanner) (GatewayTask, error) { func (s *Store) ListTaskEvents(ctx context.Context, taskID string) ([]TaskEvent, error) { rows, err := s.pool.Query(ctx, ` SELECT id::text, task_id::text, seq, event_type, COALESCE(status, ''), COALESCE(phase, ''), - COALESCE(progress, 0)::float8, COALESCE(message, ''), payload, simulated, created_at + COALESCE(progress, 0)::float8, COALESCE(message, ''), payload, simulated, created_at, + COALESCE(platform_id::text, '') FROM gateway_task_events WHERE task_id = $1::uuid ORDER BY seq ASC`, taskID) @@ -2171,6 +2217,7 @@ ORDER BY seq ASC`, taskID) &payload, &item.Simulated, &item.CreatedAt, + &item.PlatformID, ); err != nil { return nil, err } @@ -2386,10 +2433,7 @@ func taskEventsForCreate(taskID string, runMode string, status string, result ma Seq: 1, EventType: "task.accepted", Status: "queued", - Phase: "queued", - Progress: 0, - Message: "task accepted", - Payload: map[string]any{"taskId": taskID}, + Payload: map[string]any{}, Simulated: runMode == "simulation", }} } diff --git a/apps/api/internal/store/rate_limit_policy.go b/apps/api/internal/store/rate_limit_policy.go new file mode 100644 index 0000000..d717909 --- /dev/null +++ b/apps/api/internal/store/rate_limit_policy.go @@ -0,0 +1,167 @@ +package store + +import ( + "math" + "strings" +) + +const ( + RateLimitPolicyModeInherit = "inherit" + RateLimitPolicyModeOverride = "override" +) + +type EffectiveRateLimitPolicyInput struct { + BasePolicy map[string]any + PlatformPolicy map[string]any + RuntimePolicy map[string]any + RuntimePolicyExplicit bool + RuntimePolicyOverride map[string]any + ModelPolicy map[string]any + ModelPolicyMode string +} + +// EffectiveRateLimitPolicy resolves one authoritative policy. Platform +// policies are defaults for every bound model; explicit model/runtime settings +// replace the complete policy instead of merging individual metrics. +func EffectiveRateLimitPolicy(input EffectiveRateLimitPolicyInput) map[string]any { + policy := input.BasePolicy + if policySpecified(input.PlatformPolicy) { + policy = input.PlatformPolicy + } + if input.RuntimePolicyExplicit { + policy = input.RuntimePolicy + } + if raw, ok := input.RuntimePolicyOverride["rateLimitPolicy"]; ok { + policy, _ = raw.(map[string]any) + } + mode := NormalizeRateLimitPolicyMode(input.ModelPolicyMode, input.ModelPolicy != nil) + if mode == RateLimitPolicyModeOverride { + policy = input.ModelPolicy + } + return NormalizeRateLimitPolicy(policy) +} + +func NormalizeRateLimitPolicyMode(mode string, policyProvided bool) string { + switch strings.ToLower(strings.TrimSpace(mode)) { + case RateLimitPolicyModeOverride: + return RateLimitPolicyModeOverride + case RateLimitPolicyModeInherit: + return RateLimitPolicyModeInherit + default: + if policyProvided { + return RateLimitPolicyModeOverride + } + return RateLimitPolicyModeInherit + } +} + +func RateLimitPolicyMetric(policy map[string]any, metric string) (float64, bool) { + rules, _ := NormalizeRateLimitPolicy(policy)["rules"].([]any) + for _, rawRule := range rules { + rule, _ := rawRule.(map[string]any) + if strings.TrimSpace(stringValue(rule["metric"])) != metric { + continue + } + value := floatValue(rule["limit"]) + return value, true + } + return 0, false +} + +// NormalizeRateLimitPolicy keeps the canonical rules contract while accepting +// policies imported from server-main before that contract existed. Runtime +// enforcement and worker sizing must agree on these legacy limits; otherwise a +// configured max_concurrent_requests silently becomes unlimited. +func NormalizeRateLimitPolicy(policy map[string]any) map[string]any { + if policy == nil { + return nil + } + out := clonePolicy(policy) + rules, _ := out["rules"].([]any) + if len(rules) > 0 { + return out + } + + legacyScopes := []map[string]any{policy} + for _, key := range []string{"platformLimits", "modelLimits", "platform_limits", "model_limits"} { + if nested, ok := policy[key].(map[string]any); ok { + legacyScopes = append(legacyScopes, nested) + } + } + if limit, ok := lowestPositiveLegacyLimit(legacyScopes, + "max_concurrent_requests", "maxConcurrentRequests", "concurrent"); ok { + rules = append(rules, map[string]any{ + "metric": "concurrent", + "limit": limit, + "leaseTtlSeconds": 120, + }) + } + if limit, ok := lowestPositiveLegacyLimit(legacyScopes, + "max_request_per_minute", "maxRequestsPerMinute", "rpm"); ok { + rules = append(rules, map[string]any{ + "metric": "rpm", + "limit": limit, + "windowSeconds": 60, + }) + } + if limit, ok := lowestPositiveLegacyLimit(legacyScopes, + "max_tokens_per_minute", "maxTokensPerMinute", "tpm_total"); ok { + rules = append(rules, map[string]any{ + "metric": "tpm_total", + "limit": limit, + "windowSeconds": 60, + }) + } + if len(rules) > 0 { + out["rules"] = rules + } + return out +} + +func ConcurrentPolicyCapacity(policy map[string]any) (int, bool) { + limit, ok := RateLimitPolicyMetric(policy, "concurrent") + if !ok || limit <= 0 { + return 0, false + } + if limit < 1 { + return 1, true + } + if limit >= float64(math.MaxInt) { + return math.MaxInt, true + } + return int(math.Floor(limit)), true +} + +func policySpecified(policy map[string]any) bool { + rules, _ := NormalizeRateLimitPolicy(policy)["rules"].([]any) + return len(rules) > 0 +} + +func lowestPositiveLegacyLimit(scopes []map[string]any, keys ...string) (float64, bool) { + limit := 0.0 + found := false + for _, scope := range scopes { + for _, key := range keys { + value := floatValue(scope[key]) + if value <= 0 { + continue + } + if !found || value < limit { + limit = value + found = true + } + } + } + return limit, found +} + +func clonePolicy(policy map[string]any) map[string]any { + if policy == nil { + return nil + } + out := make(map[string]any, len(policy)) + for key, value := range policy { + out[key] = value + } + return out +} diff --git a/apps/api/internal/store/rate_limit_policy_migration_integration_test.go b/apps/api/internal/store/rate_limit_policy_migration_integration_test.go new file mode 100644 index 0000000..4d8464f --- /dev/null +++ b/apps/api/internal/store/rate_limit_policy_migration_integration_test.go @@ -0,0 +1,240 @@ +package store + +import ( + "context" + "fmt" + "net/url" + "os" + "path/filepath" + "runtime" + "strings" + "testing" + "time" + + "github.com/jackc/pgx/v5/pgxpool" +) + +func TestRateLimitPolicyModeMigrationClassifiesExistingRows(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 rate-limit migration PostgreSQL integration tests") + } + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second) + defer cancel() + schema := fmt.Sprintf("rate_limit_mode_%d", time.Now().UnixNano()) + adminPool, err := pgxpool.New(ctx, databaseURL) + if err != nil { + t.Fatalf("connect migration test database: %v", err) + } + defer adminPool.Close() + if _, err := adminPool.Exec(ctx, `CREATE SCHEMA `+schema); err != nil { + t.Fatalf("create migration test schema: %v", err) + } + defer adminPool.Exec(context.Background(), `DROP SCHEMA IF EXISTS `+schema+` CASCADE`) + + schemaURL, err := url.Parse(databaseURL) + if err != nil { + t.Fatalf("parse test database url: %v", err) + } + query := schemaURL.Query() + query.Set("search_path", schema) + schemaURL.RawQuery = query.Encode() + pool, err := pgxpool.New(ctx, schemaURL.String()) + if err != nil { + t.Fatalf("connect migration test schema: %v", err) + } + defer pool.Close() + if _, err := pool.Exec(ctx, ` +CREATE TABLE base_model_catalog ( + id uuid PRIMARY KEY, + default_rate_limit_policy jsonb NOT NULL DEFAULT '{}'::jsonb, + runtime_policy_set_id uuid, + runtime_policy_override jsonb NOT NULL DEFAULT '{}'::jsonb +); +CREATE TABLE integration_platforms ( + id uuid PRIMARY KEY, + rate_limit_policy jsonb NOT NULL DEFAULT '{}'::jsonb, + updated_at timestamptz NOT NULL DEFAULT now() +); +CREATE TABLE platform_models ( + id uuid PRIMARY KEY, + base_model_id uuid, + rate_limit_policy jsonb NOT NULL DEFAULT '{}'::jsonb, + runtime_policy_set_id uuid, + runtime_policy_override jsonb NOT NULL DEFAULT '{}'::jsonb, + updated_at timestamptz NOT NULL DEFAULT now() +); +INSERT INTO integration_platforms (id, rate_limit_policy) +VALUES ('20000000-0000-0000-0000-000000000001', '{"rules":[]}'); +INSERT INTO base_model_catalog (id, default_rate_limit_policy) +VALUES ('00000000-0000-0000-0000-000000000001', '{"rules":[{"metric":"concurrent","limit":4}]}'); +INSERT INTO platform_models (id, base_model_id, rate_limit_policy, runtime_policy_override) VALUES + ('10000000-0000-0000-0000-000000000001', '00000000-0000-0000-0000-000000000001', '{"rules":[{"metric":"concurrent","limit":4}]}', '{}'), + ('10000000-0000-0000-0000-000000000002', '00000000-0000-0000-0000-000000000001', '{"rules":[{"metric":"concurrent","limit":8}]}', '{}'), + ('10000000-0000-0000-0000-000000000003', '00000000-0000-0000-0000-000000000001', '{"rules":[{"metric":"concurrent","limit":4}]}', '{"rateLimitPolicy":{"rules":[]}}'), + ('10000000-0000-0000-0000-000000000004', NULL, '{}', '{}');`); err != nil { + t.Fatalf("seed pre-migration rows: %v", err) + } + _, currentFile, _, _ := runtime.Caller(0) + migrationPath := filepath.Join(filepath.Dir(currentFile), "..", "..", "migrations", "0080_platform_model_rate_limit_policy_mode.sql") + migration, err := os.ReadFile(migrationPath) + if err != nil { + t.Fatalf("read migration: %v", err) + } + if _, err := pool.Exec(ctx, string(migration)); err != nil { + t.Fatalf("apply migration: %v", err) + } + rows, err := pool.Query(ctx, `SELECT id::text, rate_limit_policy_mode FROM platform_models ORDER BY id`) + if err != nil { + t.Fatalf("read classified rows: %v", err) + } + defer rows.Close() + want := []string{"inherit", "override", "override", "inherit"} + index := 0 + for rows.Next() { + var id, mode string + if err := rows.Scan(&id, &mode); err != nil { + t.Fatalf("scan classified row: %v", err) + } + if index >= len(want) || mode != want[index] { + t.Fatalf("row %s mode=%s, want=%s", id, mode, want[index]) + } + index++ + } + if err := rows.Err(); err != nil { + t.Fatalf("read classified rows: %v", err) + } + if index != len(want) { + t.Fatalf("classified %d rows, want %d", index, len(want)) + } + var normalizedPlatformPolicy string + if err := pool.QueryRow(ctx, ` +SELECT rate_limit_policy::text +FROM integration_platforms +WHERE id = '20000000-0000-0000-0000-000000000001'`).Scan(&normalizedPlatformPolicy); err != nil { + t.Fatalf("read normalized platform policy: %v", err) + } + if normalizedPlatformPolicy != "{}" { + t.Fatalf("normalized platform policy=%s, want={}", normalizedPlatformPolicy) + } +} + +func TestRemoveUnconfiguredGPTImageConcurrencyLimitMigration(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 rate-limit migration PostgreSQL integration tests") + } + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second) + defer cancel() + schema := fmt.Sprintf("remove_gpt_image_limit_%d", time.Now().UnixNano()) + adminPool, err := pgxpool.New(ctx, databaseURL) + if err != nil { + t.Fatalf("connect migration test database: %v", err) + } + defer adminPool.Close() + if _, err := adminPool.Exec(ctx, `CREATE SCHEMA `+schema); err != nil { + t.Fatalf("create migration test schema: %v", err) + } + defer adminPool.Exec(context.Background(), `DROP SCHEMA IF EXISTS `+schema+` CASCADE`) + + schemaURL, err := url.Parse(databaseURL) + if err != nil { + t.Fatalf("parse test database url: %v", err) + } + query := schemaURL.Query() + query.Set("search_path", schema) + schemaURL.RawQuery = query.Encode() + pool, err := pgxpool.New(ctx, schemaURL.String()) + if err != nil { + t.Fatalf("connect migration test schema: %v", err) + } + defer pool.Close() + if _, err := pool.Exec(ctx, ` +CREATE TABLE base_model_catalog ( + id uuid PRIMARY KEY, + invocation_name text NOT NULL, + default_rate_limit_policy jsonb NOT NULL DEFAULT '{}'::jsonb, + metadata jsonb NOT NULL DEFAULT '{}'::jsonb, + updated_at timestamptz NOT NULL DEFAULT now() +); +INSERT INTO base_model_catalog (id, invocation_name, default_rate_limit_policy, metadata) VALUES + ( + '00000000-0000-0000-0000-000000000001', + 'gpt-image-2', + '{"rules":[{"metric":"concurrent","limit":5,"leaseTtlSeconds":300}]}', + '{"rateLimitSource":"image-gateway-migration","rateLimitUpdatedAt":"2026-07-24"}' + ), + ( + '00000000-0000-0000-0000-000000000002', + 'gpt-image-2', + '{"rules":[{"metric":"concurrent","limit":6,"leaseTtlSeconds":300}]}', + '{"rateLimitSource":"image-gateway-migration"}' + ), + ( + '00000000-0000-0000-0000-000000000003', + 'gpt-image-2', + '{"rules":[{"metric":"concurrent","limit":5,"leaseTtlSeconds":300}]}', + '{"rateLimitSource":"operator"}' + ), + ( + '00000000-0000-0000-0000-000000000004', + 'gemini-3-pro-image', + '{"rules":[{"metric":"concurrent","limit":10,"leaseTtlSeconds":600}]}', + '{"rateLimitSource":"image-gateway-migration"}' + );`); err != nil { + t.Fatalf("seed pre-migration rows: %v", err) + } + + _, currentFile, _, _ := runtime.Caller(0) + migrationPath := filepath.Join(filepath.Dir(currentFile), "..", "..", "migrations", "0082_remove_unconfigured_gpt_image_concurrency_limit.sql") + migration, err := os.ReadFile(migrationPath) + if err != nil { + t.Fatalf("read migration: %v", err) + } + if _, err := pool.Exec(ctx, string(migration)); err != nil { + t.Fatalf("apply migration: %v", err) + } + + rows, err := pool.Query(ctx, ` +SELECT id::text, default_rate_limit_policy::text, metadata->>'rateLimitSource', + metadata->>'rateLimitRemovalReason' +FROM base_model_catalog +ORDER BY id`) + if err != nil { + t.Fatalf("read migrated rows: %v", err) + } + defer rows.Close() + type result struct { + policy string + source *string + removalReason *string + } + got := make([]result, 0, 4) + for rows.Next() { + var id string + var item result + if err := rows.Scan(&id, &item.policy, &item.source, &item.removalReason); err != nil { + t.Fatalf("scan migrated row: %v", err) + } + got = append(got, item) + } + if err := rows.Err(); err != nil { + t.Fatalf("read migrated rows: %v", err) + } + if len(got) != 4 { + t.Fatalf("migrated %d rows, want 4", len(got)) + } + if got[0].policy != "{}" || got[0].source != nil || + got[0].removalReason == nil || *got[0].removalReason != "upstream-limit-unconfigured" { + t.Fatalf("migration-owned GPT policy not removed: %+v", got[0]) + } + if got[1].policy == "{}" || got[1].source == nil || *got[1].source != "image-gateway-migration" { + t.Fatalf("custom GPT policy was changed: %+v", got[1]) + } + if got[2].policy == "{}" || got[2].source == nil || *got[2].source != "operator" { + t.Fatalf("operator GPT policy was changed: %+v", got[2]) + } + if got[3].policy == "{}" || got[3].source == nil || *got[3].source != "image-gateway-migration" { + t.Fatalf("Gemini policy was changed: %+v", got[3]) + } +} diff --git a/apps/api/internal/store/rate_limit_policy_test.go b/apps/api/internal/store/rate_limit_policy_test.go new file mode 100644 index 0000000..24fae8c --- /dev/null +++ b/apps/api/internal/store/rate_limit_policy_test.go @@ -0,0 +1,169 @@ +package store + +import "testing" + +func TestEffectiveRateLimitPolicyPrecedence(t *testing.T) { + policy := func(limit float64) map[string]any { + return map[string]any{"rules": []any{map[string]any{"metric": "concurrent", "limit": limit}}} + } + tests := []struct { + name string + input EffectiveRateLimitPolicyInput + want float64 + ok bool + }{ + {name: "base", input: EffectiveRateLimitPolicyInput{BasePolicy: policy(2)}, want: 2, ok: true}, + {name: "platform", input: EffectiveRateLimitPolicyInput{BasePolicy: policy(2), PlatformPolicy: policy(4)}, want: 4, ok: true}, + {name: "empty platform inherits base", input: EffectiveRateLimitPolicyInput{BasePolicy: policy(2), PlatformPolicy: map[string]any{"rules": []any{}}}, want: 2, ok: true}, + {name: "runtime", input: EffectiveRateLimitPolicyInput{PlatformPolicy: policy(4), RuntimePolicy: policy(8), RuntimePolicyExplicit: true}, want: 8, ok: true}, + {name: "runtime override", input: EffectiveRateLimitPolicyInput{PlatformPolicy: policy(4), RuntimePolicyOverride: map[string]any{"rateLimitPolicy": policy(16)}}, want: 16, ok: true}, + {name: "model override", input: EffectiveRateLimitPolicyInput{PlatformPolicy: policy(4), ModelPolicy: policy(32), ModelPolicyMode: "override"}, want: 32, ok: true}, + {name: "model explicit unlimited", input: EffectiveRateLimitPolicyInput{PlatformPolicy: policy(4), ModelPolicy: map[string]any{}, ModelPolicyMode: "override"}, ok: false}, + {name: "model inherit", input: EffectiveRateLimitPolicyInput{PlatformPolicy: policy(4), ModelPolicy: policy(32), ModelPolicyMode: "inherit"}, want: 4, ok: true}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, ok := ConcurrentPolicyCapacity(EffectiveRateLimitPolicy(tt.input)) + if ok != tt.ok || (ok && got != int(tt.want)) { + t.Fatalf("capacity = (%d, %v), want (%d, %v)", got, ok, int(tt.want), tt.ok) + } + }) + } +} + +func TestNormalizeRateLimitPolicyLegacyShapes(t *testing.T) { + tests := []struct { + name string + policy map[string]any + metric string + want float64 + }{ + { + name: "platform concurrent", + policy: map[string]any{"platformLimits": map[string]any{"max_concurrent_requests": 5.0}}, + metric: "concurrent", + want: 5, + }, + { + name: "model concurrent camel case", + policy: map[string]any{"modelLimits": map[string]any{"maxConcurrentRequests": 10.0}}, + metric: "concurrent", + want: 10, + }, + { + name: "stricter duplicate wins", + policy: map[string]any{ + "platformLimits": map[string]any{"max_concurrent_requests": 8.0}, + "modelLimits": map[string]any{"max_concurrent_requests": 3.0}, + }, + metric: "concurrent", + want: 3, + }, + { + name: "requests per minute", + policy: map[string]any{"model_limits": map[string]any{"max_request_per_minute": 60.0}}, + metric: "rpm", + want: 60, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, ok := RateLimitPolicyMetric(tt.policy, tt.metric) + if !ok || got != tt.want { + t.Fatalf("metric %s = (%v, %v), want (%v, true)", tt.metric, got, ok, tt.want) + } + }) + } +} + +func TestConcurrentPolicyCapacityRoundsDown(t *testing.T) { + policy := map[string]any{"rules": []any{map[string]any{"metric": "concurrent", "limit": 96.9}}} + got, ok := ConcurrentPolicyCapacity(policy) + if !ok || got != 96 { + t.Fatalf("capacity = (%d, %v), want (96, true)", got, ok) + } +} + +func TestAsyncWorkerCapacityAggregation(t *testing.T) { + policy := func(limit float64) map[string]any { + return map[string]any{"rules": []any{map[string]any{"metric": "concurrent", "limit": limit}}} + } + tests := []struct { + name string + policies []map[string]any + hardLimit int + wantCapacity int + wantDesired int + wantCapped bool + }{ + {name: "no enabled models", hardLimit: 2048, wantCapacity: 1, wantDesired: 1}, + {name: "finite sum", policies: []map[string]any{policy(64), policy(32)}, hardLimit: 2048, wantCapacity: 96, wantDesired: 96}, + {name: "unlimited model", policies: []map[string]any{policy(64), {}}, hardLimit: 2048, wantCapacity: 2048, wantDesired: 2048}, + {name: "hard limit cap", policies: []map[string]any{policy(80), policy(80)}, hardLimit: 96, wantCapacity: 96, wantDesired: 160, wantCapped: true}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := asyncWorkerCapacityFromPolicies(tt.policies, tt.hardLimit) + if got.Capacity != tt.wantCapacity || got.Desired != tt.wantDesired || got.Capped != tt.wantCapped { + t.Fatalf("snapshot=%+v, want capacity=%d desired=%d capped=%v", got, tt.wantCapacity, tt.wantDesired, tt.wantCapped) + } + }) + } +} + +func TestAsyncWorkerCapacityRespectsUserGroupCeiling(t *testing.T) { + policy := func(limit float64) map[string]any { + return map[string]any{"rules": []any{map[string]any{"metric": "concurrent", "limit": limit}}} + } + tests := []struct { + name string + models []map[string]any + groups []map[string]any + hardLimit int + wantCapacity int + wantDesired int + wantCapped bool + }{ + { + name: "group ceiling prevents worker oversubscription", + models: []map[string]any{{}}, + groups: []map[string]any{policy(3), policy(10)}, + hardLimit: 2048, + wantCapacity: 13, + wantDesired: 13, + }, + { + name: "model ceiling is stricter", + models: []map[string]any{policy(5), policy(7)}, + groups: []map[string]any{policy(300)}, + hardLimit: 2048, + wantCapacity: 12, + wantDesired: 12, + }, + { + name: "both policy sets unlimited use hard limit", + models: []map[string]any{{}}, + groups: []map[string]any{{}}, + hardLimit: 256, + wantCapacity: 256, + wantDesired: 256, + }, + { + name: "finite group desired still reports hard cap", + models: []map[string]any{{}}, + groups: []map[string]any{policy(500)}, + hardLimit: 256, + wantCapacity: 256, + wantDesired: 500, + wantCapped: true, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := asyncWorkerCapacityFromPolicySets(tt.models, tt.groups, tt.hardLimit) + if got.Capacity != tt.wantCapacity || got.Desired != tt.wantDesired || got.Capped != tt.wantCapped { + t.Fatalf("snapshot=%+v, want capacity=%d desired=%d capped=%v", got, tt.wantCapacity, tt.wantDesired, tt.wantCapped) + } + }) + } +} diff --git a/apps/api/internal/store/rate_limit_retry_test.go b/apps/api/internal/store/rate_limit_retry_test.go new file mode 100644 index 0000000..dea2f2c --- /dev/null +++ b/apps/api/internal/store/rate_limit_retry_test.go @@ -0,0 +1,19 @@ +package store + +import ( + "testing" + "time" +) + +func TestConcurrencyRetryAfterBoundsPolling(t *testing.T) { + if got := concurrencyRetryAfter(time.Time{}); got != 5*time.Second { + t.Fatalf("zero expiry retry=%s, want=5s", got) + } + if got := concurrencyRetryAfter(time.Now().Add(30 * time.Second)); got != 5*time.Second { + t.Fatalf("distant expiry retry=%s, want=5s", got) + } + got := concurrencyRetryAfter(time.Now().Add(2500 * time.Millisecond)) + if got < 2*time.Second || got > 2500*time.Millisecond { + t.Fatalf("near expiry retry=%s, want within [2s,2.5s]", got) + } +} diff --git a/apps/api/internal/store/rate_limit_status.go b/apps/api/internal/store/rate_limit_status.go index 7c4cde6..cb16963 100644 --- a/apps/api/internal/store/rate_limit_status.go +++ b/apps/api/internal/store/rate_limit_status.go @@ -143,8 +143,10 @@ func (s *Store) ListModelRateLimitStatuses(ctx context.Context) ([]ModelRateLimi p.priority, p.dynamic_priority, COALESCE(p.dynamic_priority, p.priority), m.model_name, COALESCE(NULLIF(m.provider_model_name, ''), m.model_name), COALESCE(m.model_alias, ''), m.model_type, m.display_name, m.enabled, - p.rate_limit_policy, COALESCE(rp.rate_limit_policy, '{}'::jsonb), COALESCE(m.runtime_policy_set_id::text, b.runtime_policy_set_id::text, ''), - COALESCE(NULLIF(m.runtime_policy_override, '{}'::jsonb), b.runtime_policy_override, '{}'::jsonb), m.rate_limit_policy, + COALESCE(b.default_rate_limit_policy, '{}'::jsonb), p.rate_limit_policy, + COALESCE(rp.rate_limit_policy, '{}'::jsonb), + (m.runtime_policy_set_id IS NOT NULL), + COALESCE(m.runtime_policy_override, '{}'::jsonb), m.rate_limit_policy, m.rate_limit_policy_mode, COALESCE(to_char(p.cooldown_until AT TIME ZONE 'UTC', 'YYYY-MM-DD"T"HH24:MI:SS.MS"Z"'), ''), COALESCE(to_char(m.cooldown_until AT TIME ZONE 'UTC', 'YYYY-MM-DD"T"HH24:MI:SS.MS"Z"'), ''), COALESCE(con.active, 0)::float8, @@ -217,11 +219,13 @@ ORDER BY p.priority ASC, m.model_name ASC`) for rows.Next() { var item ModelRateLimitStatus var modelTypeBytes []byte + var basePolicyBytes []byte var platformPolicyBytes []byte var runtimePolicyBytes []byte - var runtimePolicySetID string + var runtimePolicyExplicit bool var runtimeOverrideBytes []byte var modelPolicyBytes []byte + var modelPolicyMode string var platformDynamicPriority sql.NullInt64 var platformCooldownUntil string var modelCooldownUntil string @@ -248,11 +252,13 @@ ORDER BY p.priority ASC, m.model_name ASC`) &modelTypeBytes, &item.DisplayName, &item.Enabled, + &basePolicyBytes, &platformPolicyBytes, &runtimePolicyBytes, - &runtimePolicySetID, + &runtimePolicyExplicit, &runtimeOverrideBytes, &modelPolicyBytes, + &modelPolicyMode, &platformCooldownUntil, &modelCooldownUntil, &concurrentCurrent, @@ -268,13 +274,15 @@ ORDER BY p.priority ASC, m.model_name ASC`) } item.PlatformDynamicPriority = intPointerFromNull(platformDynamicPriority) item.ModelType = decodeStringArray(modelTypeBytes) - policy := effectiveModelRateLimitPolicy( - decodeObject(platformPolicyBytes), - decodeObject(runtimePolicyBytes), - runtimePolicySetID, - decodeObject(runtimeOverrideBytes), - decodeObject(modelPolicyBytes), - ) + policy := EffectiveRateLimitPolicy(EffectiveRateLimitPolicyInput{ + BasePolicy: decodeObject(basePolicyBytes), + PlatformPolicy: decodeObject(platformPolicyBytes), + RuntimePolicy: decodeObject(runtimePolicyBytes), + RuntimePolicyExplicit: runtimePolicyExplicit, + RuntimePolicyOverride: decodeObject(runtimeOverrideBytes), + ModelPolicy: decodeObject(modelPolicyBytes), + ModelPolicyMode: modelPolicyMode, + }) item.PlatformCooldownUntil = platformCooldownUntil item.ModelCooldownUntil = modelCooldownUntil item.RateLimitPolicy = policy @@ -330,19 +338,20 @@ func (s *Store) listRecentPriorityDemotionsByPlatform(ctx context.Context, statu return out, nil } rows, err := s.pool.Query(ctx, ` - SELECT id::text, task_id::text, COALESCE(message, ''), payload, created_at + SELECT id::text, task_id::text, COALESCE(message, ''), payload, created_at, event_platform_id FROM ( SELECT e.*, + COALESCE(e.platform_id::text, e.payload->>'platformId') AS event_platform_id, row_number() OVER ( - PARTITION BY e.payload->>'platformId' + PARTITION BY COALESCE(e.platform_id::text, e.payload->>'platformId') ORDER BY e.created_at DESC, e.seq DESC ) AS demotion_rank FROM gateway_task_events e WHERE e.event_type = 'task.policy.priority_demoted' - AND e.payload->>'platformId' = ANY($1::text[]) + AND COALESCE(e.platform_id::text, e.payload->>'platformId') = ANY($1::text[]) ) ranked WHERE demotion_rank <= $2 - ORDER BY payload->>'platformId' ASC, created_at DESC`, platformIDs, limit) + ORDER BY event_platform_id ASC, created_at DESC`, platformIDs, limit) if err != nil { return nil, err } @@ -353,10 +362,16 @@ func (s *Store) listRecentPriorityDemotionsByPlatform(ctx context.Context, statu var message string var payloadBytes []byte var createdAt time.Time - if err := rows.Scan(&id, &taskID, &message, &payloadBytes, &createdAt); err != nil { + var platformID string + if err := rows.Scan(&id, &taskID, &message, &payloadBytes, &createdAt, &platformID); err != nil { return nil, err } - record := priorityDemotionRecordFromEventPayload(id, taskID, message, decodeObject(payloadBytes), createdAt) + payload := decodeObject(payloadBytes) + if payload == nil { + payload = map[string]any{} + } + payload["platformId"] = platformID + record := priorityDemotionRecordFromEventPayload(id, taskID, message, payload, createdAt) if record.PlatformID == "" { continue } @@ -412,12 +427,13 @@ func (s *Store) listLatestPlatformDisabledReasons(ctx context.Context, statuses return out, nil } rows, err := s.pool.Query(ctx, ` -SELECT id::text, task_id::text, event_type, COALESCE(message, ''), payload, COALESCE(attempt_error_message, ''), created_at +SELECT id::text, task_id::text, event_type, COALESCE(message, ''), payload, COALESCE(attempt_error_message, ''), created_at, event_platform_id FROM ( SELECT e.*, + COALESCE(e.platform_id::text, e.payload->>'platformId') AS event_platform_id, a.error_message AS attempt_error_message, row_number() OVER ( - PARTITION BY e.payload->>'platformId' + PARTITION BY COALESCE(e.platform_id::text, e.payload->>'platformId') ORDER BY e.created_at DESC, e.seq DESC ) AS disabled_rank FROM gateway_task_events e @@ -425,12 +441,12 @@ FROM ( SELECT error_message FROM gateway_task_attempts attempt WHERE attempt.task_id = e.task_id - AND attempt.platform_id::text = e.payload->>'platformId' + AND attempt.platform_id::text = COALESCE(e.platform_id::text, e.payload->>'platformId') ORDER BY attempt.attempt_no DESC, attempt.started_at DESC LIMIT 1 ) a ON TRUE WHERE e.event_type IN ('task.policy.failover_disabled', 'task.policy.auto_disabled') - AND e.payload->>'platformId' = ANY($1::text[]) + AND COALESCE(e.platform_id::text, e.payload->>'platformId') = ANY($1::text[]) ) ranked WHERE disabled_rank = 1`, platformIDs) if err != nil { @@ -445,10 +461,16 @@ WHERE disabled_rank = 1`, platformIDs) var payloadBytes []byte var attemptErrorMessage string var createdAt time.Time - if err := rows.Scan(&id, &taskID, &eventType, &message, &payloadBytes, &attemptErrorMessage, &createdAt); err != nil { + var platformID string + if err := rows.Scan(&id, &taskID, &eventType, &message, &payloadBytes, &attemptErrorMessage, &createdAt, &platformID); err != nil { return nil, err } - record := platformPolicyEventFromPayload(id, taskID, eventType, message, attemptErrorMessage, decodeObject(payloadBytes), createdAt) + payload := decodeObject(payloadBytes) + if payload == nil { + payload = map[string]any{} + } + payload["platformId"] = platformID + record := platformPolicyEventFromPayload(id, taskID, eventType, message, attemptErrorMessage, payload, createdAt) if record.PlatformID == "" { continue } @@ -491,46 +513,6 @@ func platformPolicyEventFromPayload(id string, taskID string, eventType string, } } -func effectiveModelRateLimitPolicy(platformPolicy map[string]any, runtimePolicy map[string]any, runtimePolicySetID string, runtimeOverride map[string]any, modelPolicy map[string]any) map[string]any { - policy := platformPolicy - if strings.TrimSpace(runtimePolicySetID) != "" { - policy = runtimePolicy - } else if hasRateLimitRules(runtimePolicy) { - policy = shallowMergeMap(policy, runtimePolicy) - } - if _, hasOverride := runtimeOverride["rateLimitPolicy"]; hasOverride { - nested, _ := runtimeOverride["rateLimitPolicy"].(map[string]any) - if len(nested) == 0 { - policy = nil - } else { - policy = shallowMergeMap(policy, nested) - } - } - if hasRateLimitRules(modelPolicy) { - policy = shallowMergeMap(policy, modelPolicy) - } - if hasRateLimitRules(policy) { - return policy - } - return nil -} - -func hasRateLimitRules(policy map[string]any) bool { - rules, _ := policy["rules"].([]any) - return len(rules) > 0 -} - -func shallowMergeMap(base map[string]any, override map[string]any) map[string]any { - out := map[string]any{} - for key, value := range base { - out[key] = value - } - for key, value := range override { - out[key] = value - } - return out -} - func rateLimitForMetric(policy map[string]any, metric string) float64 { rules, _ := policy["rules"].([]any) for _, rawRule := range rules { diff --git a/apps/api/internal/store/rate_limit_status_test.go b/apps/api/internal/store/rate_limit_status_test.go index 84f0ee0..6f71ff9 100644 --- a/apps/api/internal/store/rate_limit_status_test.go +++ b/apps/api/internal/store/rate_limit_status_test.go @@ -6,23 +6,23 @@ import ( ) func TestEffectiveModelRateLimitPolicyTreatsModelRulesAsAuthoritative(t *testing.T) { - policy := effectiveModelRateLimitPolicy( - map[string]any{"rules": []any{ + policy := EffectiveRateLimitPolicy(EffectiveRateLimitPolicyInput{ + PlatformPolicy: map[string]any{"rules": []any{ map[string]any{"metric": "rpm", "limit": 500}, map[string]any{"metric": "tpm_total", "limit": 100000}, }}, - map[string]any{"rules": []any{ + RuntimePolicy: map[string]any{"rules": []any{ map[string]any{"metric": "rpm", "limit": 120}, map[string]any{"metric": "tpm_total", "limit": 240000}, map[string]any{"metric": "concurrent", "limit": 6}, }}, - "runtime-policy-1", - map[string]any{}, - map[string]any{"rules": []any{ + RuntimePolicyExplicit: true, + ModelPolicy: map[string]any{"rules": []any{ map[string]any{"metric": "rpm", "limit": 30}, map[string]any{"metric": "concurrent", "limit": 2}, }}, - ) + ModelPolicyMode: "override", + }) if got := rateLimitForMetric(policy, "rpm"); got != 30 { t.Fatalf("expected model rpm limit to win, got %v", got) @@ -36,16 +36,15 @@ func TestEffectiveModelRateLimitPolicyTreatsModelRulesAsAuthoritative(t *testing } func TestEffectiveModelRateLimitPolicyTreatsEmptyRuntimePolicyAsUnlimited(t *testing.T) { - policy := effectiveModelRateLimitPolicy( - map[string]any{"rules": []any{ + policy := EffectiveRateLimitPolicy(EffectiveRateLimitPolicyInput{ + PlatformPolicy: map[string]any{"rules": []any{ map[string]any{"metric": "rpm", "limit": 500}, map[string]any{"metric": "tpm_total", "limit": 100000}, }}, - map[string]any{"rules": []any{}}, - "runtime-policy-1", - map[string]any{}, - map[string]any{}, - ) + RuntimePolicy: map[string]any{"rules": []any{}}, + RuntimePolicyExplicit: true, + ModelPolicyMode: "inherit", + }) if got := rateLimitForMetric(policy, "rpm"); got != 0 { t.Fatalf("expected empty runtime policy rpm to mean unlimited, got %v", got) @@ -56,17 +55,16 @@ func TestEffectiveModelRateLimitPolicyTreatsEmptyRuntimePolicyAsUnlimited(t *tes } func TestEffectiveModelRateLimitPolicyTreatsNegativeLimitAsUnlimited(t *testing.T) { - policy := effectiveModelRateLimitPolicy( - map[string]any{"rules": []any{ + policy := EffectiveRateLimitPolicy(EffectiveRateLimitPolicyInput{ + PlatformPolicy: map[string]any{"rules": []any{ map[string]any{"metric": "rpm", "limit": 500}, }}, - map[string]any{"rules": []any{ + RuntimePolicy: map[string]any{"rules": []any{ map[string]any{"metric": "rpm", "limit": -1}, }}, - "runtime-policy-1", - map[string]any{}, - map[string]any{}, - ) + RuntimePolicyExplicit: true, + ModelPolicyMode: "inherit", + }) if got := rateLimitForMetric(policy, "rpm"); got != -1 { t.Fatalf("expected negative runtime rpm marker to be preserved, got %v", got) diff --git a/apps/api/internal/store/rate_limits.go b/apps/api/internal/store/rate_limits.go index af61f2b..f9e6a8f 100644 --- a/apps/api/internal/store/rate_limits.go +++ b/apps/api/internal/store/rate_limits.go @@ -4,6 +4,7 @@ import ( "context" "errors" "fmt" + "sort" "time" "github.com/jackc/pgx/v5" @@ -17,6 +18,8 @@ type RuntimeRecoveryResult struct { RequeuedAsyncTasks int64 `json:"requeuedAsyncTasks"` } +var ErrConcurrencyLeaseLost = errors.New("concurrency lease lost") + func (s *Store) ReserveRateLimits(ctx context.Context, taskID string, attemptID string, reservations []RateLimitReservation) (RateLimitResult, error) { tx, err := s.pool.Begin(ctx) if err != nil { @@ -24,6 +27,26 @@ func (s *Store) ReserveRateLimits(ctx context.Context, taskID string, attemptID } defer tx.Rollback(ctx) + lockKeys := make([]string, 0) + lockKeySet := make(map[string]struct{}) + for _, reservation := range reservations { + if reservation.Metric != "concurrent" || reservation.Limit <= 0 || reservation.Amount <= 0 { + continue + } + key := fmt.Sprintf("%d:%s%d:%s", len(reservation.ScopeType), reservation.ScopeType, len(reservation.ScopeKey), reservation.ScopeKey) + if _, exists := lockKeySet[key]; exists { + continue + } + lockKeySet[key] = struct{}{} + lockKeys = append(lockKeys, key) + } + sort.Strings(lockKeys) + for _, key := range lockKeys { + if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended($1, 0))`, key); err != nil { + return RateLimitResult{}, err + } + } + result := RateLimitResult{} for _, reservation := range reservations { if reservation.Limit <= 0 || reservation.Amount <= 0 { @@ -49,11 +72,11 @@ func (s *Store) ReserveRateLimits(ctx context.Context, taskID string, attemptID reservation.WindowSeconds = 60 } if reservation.Metric == "concurrent" { - leaseID, err := reserveConcurrencyLease(ctx, tx, taskID, attemptID, reservation) + lease, err := reserveConcurrencyLease(ctx, tx, taskID, attemptID, reservation) if err != nil { return RateLimitResult{}, err } - result.LeaseIDs = append(result.LeaseIDs, leaseID) + result.Leases = append(result.Leases, lease) continue } normalized, err := reserveCounterWindow(ctx, tx, taskID, attemptID, reservation) @@ -65,7 +88,7 @@ func (s *Store) ReserveRateLimits(ctx context.Context, taskID string, attemptID return result, tx.Commit(ctx) } -func reserveConcurrencyLease(ctx context.Context, tx pgx.Tx, taskID string, attemptID string, reservation RateLimitReservation) (string, error) { +func reserveConcurrencyLease(ctx context.Context, tx pgx.Tx, taskID string, attemptID string, reservation RateLimitReservation) (ConcurrencyLease, error) { if reservation.LeaseTTLSeconds <= 0 { reservation.LeaseTTLSeconds = 120 } @@ -83,10 +106,10 @@ WHERE scope_type = $1 reservation.ScopeKey, reservation.LeaseTTLSeconds, ).Scan(&active, &nextAvailableAt); err != nil { - return "", err + return ConcurrencyLease{}, err } if active+reservation.Amount > reservation.Limit { - return "", &RateLimitExceededError{ + return ConcurrencyLease{}, &RateLimitExceededError{ ScopeType: reservation.ScopeType, ScopeKey: reservation.ScopeKey, ScopeName: reservation.ScopeName, @@ -117,9 +140,9 @@ RETURNING id::text`, reservation.Amount, reservation.LeaseTTLSeconds, ).Scan(&leaseID); err != nil { - return "", err + return ConcurrencyLease{}, err } - return leaseID, nil + return ConcurrencyLease{ID: leaseID, TTL: time.Duration(reservation.LeaseTTLSeconds) * time.Second}, nil } func reserveCounterWindow(ctx context.Context, tx pgx.Tx, taskID string, attemptID string, reservation RateLimitReservation) (RateLimitReservation, error) { @@ -232,13 +255,16 @@ func retryAfterUntil(when time.Time) time.Duration { func concurrencyRetryAfter(leaseExpiresAt time.Time) time.Duration { if leaseExpiresAt.IsZero() { - return time.Second + return 5 * time.Second } duration := time.Until(leaseExpiresAt) if duration <= time.Second { return time.Second } - return time.Second + if duration > 5*time.Second { + return 5 * time.Second + } + return duration } func (s *Store) CommitRateLimitReservations(ctx context.Context, reservations []RateLimitReservation, actualByMetric map[string]float64) error { @@ -249,26 +275,57 @@ func (s *Store) ReleaseRateLimitReservations(ctx context.Context, reservations [ return s.finishRateLimitReservations(ctx, reservations, nil, "released", reason) } -func (s *Store) ReleaseConcurrencyLeases(ctx context.Context, leaseIDs []string) error { +func (s *Store) ReleaseConcurrencyLeases(ctx context.Context, leases []ConcurrencyLease) error { + leaseIDs := concurrencyLeaseIDs(leases) if len(leaseIDs) == 0 { return nil } - for _, leaseID := range leaseIDs { - if leaseID == "" { - continue - } - if _, err := s.pool.Exec(ctx, ` + _, err := s.pool.Exec(ctx, ` UPDATE gateway_concurrency_leases SET released_at = now() -WHERE id = $1::uuid AND released_at IS NULL`, leaseID); err != nil && !errors.Is(err, ErrRateLimited) { - return err +WHERE id = ANY($1::uuid[]) AND released_at IS NULL`, leaseIDs) + return err +} + +func (s *Store) RenewConcurrencyLeases(ctx context.Context, leases []ConcurrencyLease) error { + leaseIDs := make([]string, 0, len(leases)) + ttlSeconds := make([]int32, 0, len(leases)) + for _, lease := range leases { + if lease.ID == "" { + continue } + ttl := lease.TTL + if ttl <= 0 { + ttl = 120 * time.Second + } + seconds := int32(ttl / time.Second) + if seconds < 1 { + seconds = 1 + } + leaseIDs = append(leaseIDs, lease.ID) + ttlSeconds = append(ttlSeconds, seconds) + } + if len(leaseIDs) == 0 { + return nil + } + tag, err := s.pool.Exec(ctx, ` +UPDATE gateway_concurrency_leases lease +SET expires_at = now() + (renewal.ttl_seconds * interval '1 second') +FROM unnest($1::uuid[], $2::int[]) AS renewal(id, ttl_seconds) +WHERE lease.id = renewal.id + AND lease.released_at IS NULL + AND lease.expires_at > now()`, leaseIDs, ttlSeconds) + if err != nil { + return err + } + if tag.RowsAffected() != int64(len(leaseIDs)) { + return fmt.Errorf("%w: renewed %d of %d leases", ErrConcurrencyLeaseLost, tag.RowsAffected(), len(leaseIDs)) } return nil } func (s *Store) AttachRateLimitResultToAttempt(ctx context.Context, attemptID string, result RateLimitResult) error { - if attemptID == "" || (len(result.Reservations) == 0 && len(result.LeaseIDs) == 0) { + if attemptID == "" || (len(result.Reservations) == 0 && len(result.Leases) == 0) { return nil } tx, err := s.pool.Begin(ctx) @@ -289,7 +346,8 @@ WHERE id = $1::uuid`, reservation.ReservationID, attemptID); err != nil { return err } } - for _, leaseID := range result.LeaseIDs { + for _, lease := range result.Leases { + leaseID := lease.ID if leaseID == "" { continue } @@ -303,6 +361,16 @@ WHERE id = $1::uuid`, leaseID, attemptID); err != nil { return tx.Commit(ctx) } +func concurrencyLeaseIDs(leases []ConcurrencyLease) []string { + ids := make([]string, 0, len(leases)) + for _, lease := range leases { + if lease.ID != "" { + ids = append(ids, lease.ID) + } + } + return ids +} + func (s *Store) RecoverInterruptedRuntimeState(ctx context.Context) (RuntimeRecoveryResult, error) { tx, err := s.pool.Begin(ctx) if err != nil { @@ -401,17 +469,37 @@ RETURNING id::text`) for _, taskID := range asyncTaskIDs { if _, err := tx.Exec(ctx, ` INSERT INTO gateway_task_events (task_id, seq, event_type, status, phase, progress, message, payload, simulated) -VALUES ( - $1::uuid, +SELECT + task.id, COALESCE((SELECT MAX(seq) + 1 FROM gateway_task_events WHERE task_id = $1::uuid), 1), - 'task.recovered', + 'task.queued', 'queued', - 'recovered', - 0.2, - 'async task recovered after service restart', - '{"code":"server_restarted"}'::jsonb, + NULL, + 0, + NULL, + '{}'::jsonb, false -)`, taskID); err != nil { +FROM gateway_tasks task +WHERE task.id = $1::uuid + AND ( + SELECT count(*) + FROM gateway_task_events event + WHERE event.task_id = task.id + AND event.event_type NOT IN ( + 'task.completed', 'task.failed', 'task.cancelled', + 'task.billing.settled', 'task.billing.released', 'task.billing.review' + ) + ) < 16 + AND NOT EXISTS ( + SELECT 1 + FROM gateway_task_events event + WHERE event.task_id = task.id + AND event.seq = (SELECT MAX(last_event.seq) FROM gateway_task_events last_event WHERE last_event.task_id = task.id) + AND event.event_type = 'task.queued' + AND COALESCE(event.status, '') = 'queued' + AND event.platform_id IS NULL + AND event.simulated = false + )`, taskID); err != nil { return RuntimeRecoveryResult{}, err } } @@ -420,9 +508,10 @@ VALUES ( taskRows, err := tx.Query(ctx, ` UPDATE gateway_tasks SET status = 'failed', - error = 'task interrupted by service restart', + error = NULL, error_code = 'server_restarted', error_message = 'task interrupted by service restart', + remote_task_payload = '{}'::jsonb, finished_at = now(), updated_at = now() WHERE async_mode = false @@ -452,12 +541,12 @@ INSERT INTO gateway_task_events (task_id, seq, event_type, status, phase, progre VALUES ( $1::uuid, COALESCE((SELECT MAX(seq) + 1 FROM gateway_task_events WHERE task_id = $1::uuid), 1), - 'task.recovered', + 'task.failed', 'failed', - 'recovered', - 1, - 'task interrupted by service restart', - '{"code":"server_restarted"}'::jsonb, + NULL, + 0, + NULL, + '{}'::jsonb, false )`, taskID); err != nil { return RuntimeRecoveryResult{}, err diff --git a/apps/api/internal/store/rate_limits_integration_test.go b/apps/api/internal/store/rate_limits_integration_test.go new file mode 100644 index 0000000..572d591 --- /dev/null +++ b/apps/api/internal/store/rate_limits_integration_test.go @@ -0,0 +1,198 @@ +package store + +import ( + "context" + "errors" + "os" + "strings" + "sync" + "sync/atomic" + "testing" + "time" +) + +func TestConcurrencyLeaseReservationIsAtomicAcrossPools(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 concurrency lease PostgreSQL integration tests") + } + ctx, cancel := context.WithTimeout(context.Background(), 45*time.Second) + defer cancel() + first, err := Connect(ctx, databaseURL) + if err != nil { + t.Fatalf("connect first store: %v", err) + } + defer first.Close() + second, err := Connect(ctx, databaseURL) + if err != nil { + t.Fatalf("connect second store: %v", err) + } + defer second.Close() + + scopeKey := "atomic-" + time.Now().UTC().Format("20060102150405.000000000") + taskIDs := createLeaseTestTasks(t, ctx, first, 256, scopeKey) + defer deleteLeaseTestTasks(t, first, taskIDs) + + var successes atomic.Int64 + var peak atomic.Int64 + monitorCtx, stopMonitor := context.WithCancel(ctx) + var monitorWG sync.WaitGroup + monitorWG.Add(1) + go func() { + defer monitorWG.Done() + ticker := time.NewTicker(2 * time.Millisecond) + defer ticker.Stop() + for { + select { + case <-monitorCtx.Done(): + return + case <-ticker.C: + var active int64 + if err := first.Pool().QueryRow(monitorCtx, ` +SELECT COUNT(*) +FROM gateway_concurrency_leases +WHERE scope_type = 'platform_model' + AND scope_key = $1 + AND released_at IS NULL + AND expires_at > now()`, scopeKey).Scan(&active); err == nil { + for active > peak.Load() && !peak.CompareAndSwap(peak.Load(), active) { + } + } + } + } + }() + + var wg sync.WaitGroup + errs := make(chan error, len(taskIDs)) + for index, taskID := range taskIDs { + wg.Add(1) + go func(index int, taskID string) { + defer wg.Done() + target := first + if index%2 == 1 { + target = second + } + _, err := target.ReserveRateLimits(ctx, taskID, "", []RateLimitReservation{{ + ScopeType: "platform_model", + ScopeKey: scopeKey, + Metric: "concurrent", + Limit: 64, + Amount: 1, + LeaseTTLSeconds: 30, + }}) + if err == nil { + successes.Add(1) + return + } + if !errors.Is(err, ErrRateLimited) { + errs <- err + } + }(index, taskID) + } + wg.Wait() + stopMonitor() + monitorWG.Wait() + close(errs) + for err := range errs { + t.Fatalf("unexpected reservation error: %v", err) + } + + var active int64 + if err := first.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()`, scopeKey).Scan(&active); err != nil { + t.Fatalf("count active leases: %v", err) + } + if successes.Load() != 64 || active != 64 { + t.Fatalf("successful reservations=%d active leases=%d, want exactly 64", successes.Load(), active) + } + if peak.Load() > 64 { + t.Fatalf("active lease peak=%d, want <=64", peak.Load()) + } +} + +func TestConcurrencyLeaseRenewalExtendsAndReleases(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 concurrency lease PostgreSQL integration tests") + } + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + db, err := Connect(ctx, databaseURL) + if err != nil { + t.Fatalf("connect store: %v", err) + } + defer db.Close() + scopeKey := "renew-" + time.Now().UTC().Format("20060102150405.000000000") + taskIDs := createLeaseTestTasks(t, ctx, db, 1, scopeKey) + defer deleteLeaseTestTasks(t, db, taskIDs) + + result, err := db.ReserveRateLimits(ctx, taskIDs[0], "", []RateLimitReservation{{ + ScopeType: "platform_model", + ScopeKey: scopeKey, + Metric: "concurrent", + Limit: 1, + Amount: 1, + LeaseTTLSeconds: 2, + }}) + if err != nil { + t.Fatalf("reserve short lease: %v", err) + } + time.Sleep(time.Second) + if err := db.RenewConcurrencyLeases(ctx, result.Leases); err != nil { + t.Fatalf("renew short lease: %v", err) + } + time.Sleep(1500 * time.Millisecond) + var active bool + if err := db.Pool().QueryRow(ctx, ` +SELECT EXISTS ( + SELECT 1 FROM gateway_concurrency_leases + WHERE id = $1::uuid AND released_at IS NULL AND expires_at > now() +)`, result.Leases[0].ID).Scan(&active); err != nil { + t.Fatalf("read renewed lease: %v", err) + } + if !active { + t.Fatal("renewed lease expired at its original TTL") + } + if err := db.ReleaseConcurrencyLeases(ctx, result.Leases); err != nil { + t.Fatalf("release renewed lease: %v", err) + } +} + +func createLeaseTestTasks(t *testing.T, ctx context.Context, db *Store, count int, marker string) []string { + t.Helper() + rows, err := db.Pool().Query(ctx, ` +INSERT INTO gateway_tasks (kind, run_mode, user_id, model, model_type, request, status, queue_key) +SELECT 'lease-test', 'simulation', $2, 'lease-test', 'text_generate', '{}'::jsonb, 'queued', $2 +FROM generate_series(1, $1) +RETURNING id::text`, count, marker) + if err != nil { + t.Fatalf("create lease test tasks: %v", err) + } + defer rows.Close() + ids := make([]string, 0, count) + for rows.Next() { + var id string + if err := rows.Scan(&id); err != nil { + t.Fatalf("scan lease test task: %v", err) + } + ids = append(ids, id) + } + if err := rows.Err(); err != nil { + t.Fatalf("create lease test tasks: %v", err) + } + return ids +} + +func deleteLeaseTestTasks(t *testing.T, db *Store, taskIDs []string) { + t.Helper() + cleanupCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + if _, err := db.Pool().Exec(cleanupCtx, `DELETE FROM gateway_tasks WHERE id = ANY($1::uuid[])`, taskIDs); err != nil { + t.Errorf("delete lease test tasks: %v", err) + } +} diff --git a/apps/api/internal/store/response_chains.go b/apps/api/internal/store/response_chains.go index 45840de..8831012 100644 --- a/apps/api/internal/store/response_chains.go +++ b/apps/api/internal/store/response_chains.go @@ -45,15 +45,15 @@ type CreateResponseChainInput struct { } func (s *Store) CreateResponseChain(ctx context.Context, input CreateResponseChainInput) error { - requestJSON, err := json.Marshal(input.RequestSnapshot) + requestJSON, err := json.Marshal(sanitizeJSONForStorage(input.RequestSnapshot)) if err != nil { return err } - responseJSON, err := json.Marshal(input.ResponseSnapshot) + responseJSON, err := json.Marshal(sanitizeJSONForStorage(input.ResponseSnapshot)) if err != nil { return err } - internalJSON, err := json.Marshal(input.InternalSnapshot) + internalJSON, err := json.Marshal(sanitizeJSONForStorage(input.InternalSnapshot)) if err != nil { return err } diff --git a/apps/api/internal/store/runtime_policies.go b/apps/api/internal/store/runtime_policies.go index 3b04730..a2dfe11 100644 --- a/apps/api/internal/store/runtime_policies.go +++ b/apps/api/internal/store/runtime_policies.go @@ -8,6 +8,7 @@ import ( "time" "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgconn" ) const runtimePolicyColumns = ` @@ -38,6 +39,24 @@ type PlatformDynamicPriorityState struct { UpdatedAt time.Time `json:"updatedAt"` } +type CandidateFailureEffectInput struct { + Effect string + Target string + PlatformID string + PlatformModelID string + RequestedModel string + ModelType string + CooldownSeconds int + DemoteSteps int + SingleSourceProtection bool +} + +type CandidateFailureEffectResult struct { + Applied bool + GuardReason string + DynamicPriority int +} + func (s *Store) ListRuntimePolicySets(ctx context.Context) ([]RuntimePolicySet, error) { rows, err := s.pool.Query(ctx, `SELECT `+runtimePolicyColumns+` FROM model_runtime_policy_sets ORDER BY policy_key ASC`) if err != nil { @@ -171,6 +190,164 @@ WHERE id = $1::uuid`, platformModelID, cooldownSeconds) return err } +func (s *Store) RuntimeCandidateAvailable(ctx context.Context, platformID string, platformModelID string) (bool, error) { + if strings.TrimSpace(platformID) == "" || strings.TrimSpace(platformModelID) == "" { + return false, nil + } + var available bool + err := s.pool.QueryRow(ctx, ` +SELECT EXISTS ( + SELECT 1 + FROM platform_models model + JOIN integration_platforms platform ON platform.id = model.platform_id + WHERE platform.id = $1::uuid + AND model.id = $2::uuid + AND platform.status = 'enabled' + AND platform.deleted_at IS NULL + AND (platform.cooldown_until IS NULL OR platform.cooldown_until <= now()) + AND model.enabled = true + AND (model.cooldown_until IS NULL OR model.cooldown_until <= now()) +)`, platformID, platformModelID).Scan(&available) + return available, err +} + +func (s *Store) ApplyCandidateFailureEffect(ctx context.Context, input CandidateFailureEffectInput) (CandidateFailureEffectResult, error) { + if input.Effect == "" || input.Effect == "none" { + return CandidateFailureEffectResult{}, nil + } + if input.Effect == "demote" { + steps := input.DemoteSteps + if steps <= 0 { + steps = 1 + } + priority, err := s.DemoteCandidatePlatformPriorityBySteps( + ctx, + input.PlatformID, + input.PlatformModelID, + input.RequestedModel, + input.ModelType, + steps, + ) + return CandidateFailureEffectResult{Applied: err == nil, DynamicPriority: priority}, err + } + if input.Effect != "cooldown" && input.Effect != "disable" { + return CandidateFailureEffectResult{}, nil + } + + tx, err := s.pool.Begin(ctx) + if err != nil { + return CandidateFailureEffectResult{}, err + } + defer func() { + _ = tx.Rollback(ctx) + }() + + lockKey := strings.TrimSpace(input.RequestedModel) + "\x1f" + strings.TrimSpace(input.ModelType) + if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended($1::text, 0))`, lockKey); err != nil { + return CandidateFailureEffectResult{}, err + } + if input.SingleSourceProtection { + alternativeAvailable, err := candidateAlternativeAvailable(ctx, tx, input) + if err != nil { + return CandidateFailureEffectResult{}, err + } + if !alternativeAvailable { + if err := tx.Commit(ctx); err != nil { + return CandidateFailureEffectResult{}, err + } + return CandidateFailureEffectResult{GuardReason: "single_source"}, nil + } + } + + var commandTag pgconn.CommandTag + switch { + case input.Effect == "disable" && input.Target == "model": + commandTag, err = tx.Exec(ctx, ` +UPDATE platform_models +SET enabled = false, + updated_at = now() +WHERE id = $1::uuid + AND enabled = true`, input.PlatformModelID) + case input.Effect == "disable": + commandTag, err = tx.Exec(ctx, ` +UPDATE integration_platforms +SET status = 'disabled', + updated_at = now() +WHERE id = $1::uuid + AND status = 'enabled' + AND deleted_at IS NULL`, input.PlatformID) + case input.Effect == "cooldown" && input.Target == "platform": + commandTag, err = tx.Exec(ctx, ` +UPDATE integration_platforms +SET cooldown_until = GREATEST( + COALESCE(cooldown_until, to_timestamp(0)), + now() + ($2::int * interval '1 second') + ), + updated_at = now() +WHERE id = $1::uuid + AND status = 'enabled' + AND deleted_at IS NULL`, input.PlatformID, positiveCooldownSeconds(input.CooldownSeconds)) + default: + commandTag, err = tx.Exec(ctx, ` +UPDATE platform_models +SET cooldown_until = GREATEST( + COALESCE(cooldown_until, to_timestamp(0)), + now() + ($2::int * interval '1 second') + ), + updated_at = now() +WHERE id = $1::uuid + AND enabled = true`, input.PlatformModelID, positiveCooldownSeconds(input.CooldownSeconds)) + } + if err != nil { + return CandidateFailureEffectResult{}, err + } + applied := commandTag.RowsAffected() > 0 + if err := tx.Commit(ctx); err != nil { + return CandidateFailureEffectResult{}, err + } + return CandidateFailureEffectResult{Applied: applied}, nil +} + +func candidateAlternativeAvailable(ctx context.Context, tx pgx.Tx, input CandidateFailureEffectInput) (bool, error) { + var available bool + err := tx.QueryRow(ctx, ` +SELECT EXISTS ( + SELECT 1 + FROM platform_models model + JOIN integration_platforms platform ON platform.id = model.platform_id + LEFT JOIN base_model_catalog base ON base.id = model.base_model_id + WHERE platform.id <> $1::uuid + AND platform.status = 'enabled' + AND platform.deleted_at IS NULL + AND (platform.cooldown_until IS NULL OR platform.cooldown_until <= now()) + AND model.enabled = true + AND (model.cooldown_until IS NULL OR model.cooldown_until <= now()) + AND (NULLIF($3::text, '') IS NULL OR model.model_type @> jsonb_build_array($3::text)) + AND ( + base.invocation_name = $2::text + OR (base.id IS NULL AND model.model_name = $2::text) + OR EXISTS ( + SELECT 1 + FROM model_compatibility_aliases compatibility_alias + JOIN base_model_catalog alias_base ON alias_base.id = compatibility_alias.base_model_id + WHERE compatibility_alias.alias = $2::text + AND compatibility_alias.model_type = $3::text + AND compatibility_alias.active = true + AND (compatibility_alias.expires_at IS NULL OR compatibility_alias.expires_at > now()) + AND (compatibility_alias.base_model_id = base.id OR alias_base.invocation_name = base.invocation_name) + ) + ) +)`, input.PlatformID, strings.TrimSpace(input.RequestedModel), strings.TrimSpace(input.ModelType)).Scan(&available) + return available, err +} + +func positiveCooldownSeconds(value int) int { + if value <= 0 { + return 300 + } + return value +} + func (s *Store) DemoteCandidatePlatformPriority(ctx context.Context, platformID string, platformModelID string, requestedModel string, modelType string) (int, error) { return s.DemoteCandidatePlatformPriorityBySteps(ctx, platformID, platformModelID, requestedModel, modelType, 99) } diff --git a/apps/api/internal/store/runtime_types.go b/apps/api/internal/store/runtime_types.go index cd2e7ed..cdf7640 100644 --- a/apps/api/internal/store/runtime_types.go +++ b/apps/api/internal/store/runtime_types.go @@ -13,9 +13,11 @@ var ( ) type ModelCandidateUnavailableError struct { - Code string - Message string - Details map[string]any + Code string + Message string + Details map[string]any + RetryAfter time.Duration + RecoveryAt time.Time } func (e *ModelCandidateUnavailableError) Error() string { @@ -42,6 +44,22 @@ func ModelCandidateErrorDetails(err error) map[string]any { return nil } +func ModelCandidateRetryAfter(err error) time.Duration { + var candidateErr *ModelCandidateUnavailableError + if errors.As(err, &candidateErr) && candidateErr.RetryAfter > 0 { + return candidateErr.RetryAfter + } + return 0 +} + +func ModelCandidateRecoveryAt(err error) time.Time { + var candidateErr *ModelCandidateUnavailableError + if errors.As(err, &candidateErr) { + return candidateErr.RecoveryAt + } + return time.Time{} +} + type RateLimitExceededError struct { ScopeType string ScopeKey string @@ -111,6 +129,7 @@ type CreatePlatformModelInput struct { PermissionConfig map[string]any `json:"permissionConfig"` RetryPolicy map[string]any `json:"retryPolicy"` RateLimitPolicy map[string]any `json:"rateLimitPolicy"` + RateLimitPolicyMode string `json:"rateLimitPolicyMode" enums:"inherit,override"` RuntimePolicySetID string `json:"runtimePolicySetId"` RuntimePolicyOverride map[string]any `json:"runtimePolicyOverride"` Enabled bool `json:"enabled"` @@ -130,6 +149,7 @@ type RuntimeModelCandidate struct { DefaultDiscountFactor float64 PlatformRetryPolicy map[string]any PlatformRateLimitPolicy map[string]any + BaseRateLimitPolicy map[string]any PlatformPriority int PlatformModelID string BaseModelID string @@ -139,6 +159,7 @@ type RuntimeModelCandidate struct { ModelAlias string ModelType string DisplayName string + LegacyAliasUsed bool Capabilities map[string]any CapabilityOverride map[string]any BaseBillingConfig map[string]any @@ -152,8 +173,11 @@ type RuntimeModelCandidate struct { ModelPricingRuleSetID string ModelRetryPolicy map[string]any ModelRateLimitPolicy map[string]any + ModelRateLimitPolicyMode string RuntimePolicySetID string + RuntimePolicyExplicit bool RuntimePolicyOverride map[string]any + RateLimitRuntimeOverride map[string]any RuntimeRetryPolicy map[string]any RuntimeRateLimitPolicy map[string]any AutoDisablePolicy map[string]any @@ -172,18 +196,21 @@ type RuntimeModelCandidate struct { } type RuntimeCandidateCacheAffinity struct { - Key string - RequestCount int - InputTokens int - CachedInputTokens int - EMAHitRatio float64 - LastHitRatio float64 - LastObservedUnix float64 - Confidence float64 - Score float64 - Boost float64 - AdjustedPriority float64 - Applied bool + Key string + RequestCount int + InputTokens int + CachedInputTokens int + EMAHitRatio float64 + LastHitRatio float64 + LastObservedUnix float64 + Confidence float64 + Score float64 + Boost float64 + AdjustedPriority float64 + MatchedPrefixDepth int + CandidateCount int + OverrideReason string + Applied bool } type RuntimeCandidateLoadMetrics struct { @@ -218,10 +245,15 @@ type RateLimitReservation struct { } type RateLimitResult struct { - LeaseIDs []string + Leases []ConcurrencyLease Reservations []RateLimitReservation } +type ConcurrencyLease struct { + ID string + TTL time.Duration +} + type CreateTaskAttemptInput struct { TaskID string AttemptNo int @@ -248,6 +280,7 @@ type FinishTaskAttemptInput struct { Status string Retryable bool RequestID string + StatusCode int Usage map[string]any Metrics map[string]any ResponseSnapshot map[string]any diff --git a/apps/api/internal/store/task_history.go b/apps/api/internal/store/task_history.go new file mode 100644 index 0000000..022e154 --- /dev/null +++ b/apps/api/internal/store/task_history.go @@ -0,0 +1,496 @@ +package store + +import ( + "context" + "time" + + "github.com/jackc/pgx/v5" +) + +const ( + TaskCallbackBatchSize = 100 + TaskCallbackLockTTL = 60 * time.Second +) + +type TaskCallbackDelivery struct { + ID string + TaskID string + EventID string + Seq int64 + CallbackURL string + Status string + Attempts int + LockToken string + EventType string + TaskStatus string + CreatedAt time.Time +} + +type TaskHistoryCleanupResult struct { + CompactedTasks int64 + CompactedCheckpoints int64 + CompactedAttempts int64 + CompactedEvents int64 + CompactedCallbacks int64 + RetiredCallbacks int64 + CompactedParamLogs int64 + CompactedClonedVoices int64 + DeletedCallbacks int64 + DeletedParamLogs int64 + DeletedEvents int64 + DeletedAttempts int64 + DeletedTasks int64 +} + +func (result TaskHistoryCleanupResult) Total() int64 { + return result.CompactedTasks + + result.CompactedCheckpoints + + result.CompactedAttempts + + result.CompactedEvents + + result.CompactedCallbacks + + result.RetiredCallbacks + + result.CompactedParamLogs + + result.CompactedClonedVoices + + result.DeletedCallbacks + + result.DeletedParamLogs + + result.DeletedEvents + + result.DeletedAttempts + + result.DeletedTasks +} + +func (s *Store) ClaimTaskCallbacks(ctx context.Context, workerID string, limit int, staleAfter time.Duration) ([]TaskCallbackDelivery, error) { + if limit <= 0 || limit > TaskCallbackBatchSize { + limit = TaskCallbackBatchSize + } + if staleAfter <= 0 { + staleAfter = TaskCallbackLockTTL + } + rows, err := s.pool.Query(ctx, ` +WITH picked AS ( + SELECT outbox.id, event.event_type, COALESCE(event.status, '') AS task_status, event.created_at + FROM gateway_task_callback_outbox outbox + JOIN gateway_task_events event ON event.id = outbox.event_id + WHERE outbox.created_at >= COALESCE(( + SELECT applied_at + FROM schema_migrations + WHERE version = '0083_task_history_minimal_storage' + ), now()) + AND ( + ( + outbox.status = 'pending' + AND outbox.next_attempt_at <= now() + ) + OR ( + outbox.status = 'processing' + AND outbox.locked_at < now() - ($3::int * interval '1 second') + ) + ) + ORDER BY outbox.next_attempt_at, outbox.created_at + LIMIT $2 + FOR UPDATE OF outbox SKIP LOCKED +) +UPDATE gateway_task_callback_outbox outbox +SET status = 'processing', + attempts = outbox.attempts + 1, + locked_by = $1, + lock_token = gen_random_uuid(), + locked_at = now(), + updated_at = now() +FROM picked +WHERE outbox.id = picked.id +RETURNING outbox.id::text, outbox.task_id::text, COALESCE(outbox.event_id::text, ''), + outbox.seq, outbox.callback_url, outbox.status, outbox.attempts, + COALESCE(outbox.lock_token::text, ''), picked.event_type, picked.task_status, + picked.created_at`, + workerID, limit, int(staleAfter/time.Second)) + if err != nil { + return nil, err + } + defer rows.Close() + items := make([]TaskCallbackDelivery, 0) + for rows.Next() { + var item TaskCallbackDelivery + if err := rows.Scan( + &item.ID, + &item.TaskID, + &item.EventID, + &item.Seq, + &item.CallbackURL, + &item.Status, + &item.Attempts, + &item.LockToken, + &item.EventType, + &item.TaskStatus, + &item.CreatedAt, + ); err != nil { + return nil, err + } + items = append(items, item) + } + return items, rows.Err() +} + +func (s *Store) MarkTaskCallbackDelivered(ctx context.Context, item TaskCallbackDelivery) error { + _, err := s.pool.Exec(ctx, ` +UPDATE gateway_task_callback_outbox +SET status = 'delivered', + delivered_at = now(), + failed_at = NULL, + last_error = NULL, + locked_by = NULL, + lock_token = NULL, + locked_at = NULL, + updated_at = now() +WHERE id = $1::uuid + AND status = 'processing' + AND lock_token = $2::uuid`, item.ID, item.LockToken) + return err +} + +func (s *Store) MarkTaskCallbackFailed(ctx context.Context, item TaskCallbackDelivery, retry bool, nextAttemptAt time.Time, message string) error { + status := "failed" + if retry { + status = "pending" + } + _, err := s.pool.Exec(ctx, ` +UPDATE gateway_task_callback_outbox +SET status = $3, + next_attempt_at = CASE WHEN $3 = 'pending' THEN $4::timestamptz ELSE next_attempt_at END, + failed_at = CASE WHEN $3 = 'failed' THEN now() ELSE NULL END, + last_error = NULLIF(left($5, 2048), ''), + locked_by = NULL, + lock_token = NULL, + locked_at = NULL, + updated_at = now() +WHERE id = $1::uuid + AND status = 'processing' + AND lock_token = $2::uuid`, + item.ID, item.LockToken, status, nextAttemptAt, message) + return err +} + +func (s *Store) CleanupTaskHistory(ctx context.Context, analysisCutoff time.Time, taskCutoff time.Time, batchSize int) (TaskHistoryCleanupResult, error) { + if batchSize < 1 { + batchSize = 1000 + } + var result TaskHistoryCleanupResult + err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error { + if _, err := tx.Exec(ctx, `SET LOCAL lock_timeout = '5s'; SET LOCAL statement_timeout = '5s'`); err != nil { + return err + } + steps := []struct { + target *int64 + sql string + args []any + }{ + {&result.CompactedTasks, ` +WITH picked AS ( + SELECT id + FROM gateway_tasks + WHERE COALESCE(normalized_request, '{}'::jsonb) <> '{}'::jsonb + OR compatibility_public_id IS NOT NULL + OR compatibility_submit_http_status IS NOT NULL + OR COALESCE(compatibility_submit_headers, '{}'::jsonb) <> '{}'::jsonb + OR COALESCE(compatibility_submit_body, '{}'::jsonb) <> '{}'::jsonb + OR result ?| ARRAY[ + 'raw', 'raw_data', 'rawData', 'provider_response', 'providerResponse', + 'submit', 'file_retrieve', 'fileRetrieve', 'upstream_task_id', 'remote_task_id' + ] + OR error IS NOT NULL + OR octet_length(COALESCE(error_message, '')) > 2048 + ORDER BY updated_at + LIMIT $1 + FOR UPDATE SKIP LOCKED +) +UPDATE gateway_tasks task +SET normalized_request = '{}'::jsonb, + compatibility_public_id = NULL, + compatibility_submit_http_status = NULL, + compatibility_submit_headers = '{}'::jsonb, + compatibility_submit_body = '{}'::jsonb, + result = COALESCE(task.result, '{}'::jsonb) + - 'raw' - 'raw_data' - 'rawData' + - 'provider_response' - 'providerResponse' + - 'submit' - 'file_retrieve' - 'fileRetrieve' + - 'upstream_task_id' - 'remote_task_id', + error = NULL, + error_message = CASE + WHEN octet_length(COALESCE(task.error_message, '')) > 2048 THEN left(task.error_message, 512) + ELSE task.error_message + END, + remote_task_payload = CASE + WHEN task.status IN ('succeeded', 'failed', 'cancelled') THEN '{}'::jsonb + ELSE ( + SELECT COALESCE(jsonb_object_agg(entry.key, entry.value), '{}'::jsonb) + FROM jsonb_each(COALESCE(task.remote_task_payload, '{}'::jsonb)) entry + WHERE entry.key IN ( + 'endpoint', 'pollEndpoint', 'phase', 'mode', 'taskType', + 'targetResolution', 'imageToken', 'receipt', 'cleanupElementIds' + ) + ) + END +FROM picked +WHERE task.id = picked.id`, []any{batchSize}}, + {&result.CompactedCheckpoints, ` +WITH picked AS ( + SELECT id + FROM gateway_tasks + WHERE COALESCE(remote_task_payload, '{}'::jsonb) <> '{}'::jsonb + AND ( + status IN ('succeeded', 'failed', 'cancelled') + OR EXISTS ( + SELECT 1 + FROM jsonb_object_keys(COALESCE(remote_task_payload, '{}'::jsonb)) AS key + WHERE key NOT IN ( + 'endpoint', 'pollEndpoint', 'phase', 'mode', 'taskType', + 'targetResolution', 'imageToken', 'receipt', 'cleanupElementIds' + ) + ) + ) + ORDER BY updated_at + LIMIT $1 + FOR UPDATE SKIP LOCKED +) +UPDATE gateway_tasks task +SET remote_task_payload = CASE + WHEN task.status IN ('succeeded', 'failed', 'cancelled') THEN '{}'::jsonb + ELSE ( + SELECT COALESCE(jsonb_object_agg(entry.key, entry.value), '{}'::jsonb) + FROM jsonb_each(COALESCE(task.remote_task_payload, '{}'::jsonb)) entry + WHERE entry.key IN ( + 'endpoint', 'pollEndpoint', 'phase', 'mode', 'taskType', + 'targetResolution', 'imageToken', 'receipt', 'cleanupElementIds' + ) + ) + END +FROM picked +WHERE task.id = picked.id`, []any{batchSize}}, + {&result.CompactedAttempts, ` +WITH picked AS ( + SELECT id + FROM gateway_task_attempts + WHERE request_snapshot <> '{}'::jsonb + OR COALESCE(response_snapshot, '{}'::jsonb) <> '{}'::jsonb + OR COALESCE(usage, '{}'::jsonb) <> '{}'::jsonb + OR COALESCE(metrics, '{}'::jsonb) <> '{}'::jsonb + OR COALESCE(pricing_snapshot, '{}'::jsonb) <> '{}'::jsonb + OR request_fingerprint IS NOT NULL + ORDER BY started_at + LIMIT $1 + FOR UPDATE SKIP LOCKED +) +UPDATE gateway_task_attempts attempt +SET request_snapshot = '{}'::jsonb, + response_snapshot = '{}'::jsonb, + usage = '{}'::jsonb, + metrics = '{}'::jsonb, + pricing_snapshot = '{}'::jsonb, + request_fingerprint = NULL, + error_message = NULLIF(left(error_message, 2048), '') +FROM picked +WHERE attempt.id = picked.id`, []any{batchSize}}, + {&result.CompactedEvents, ` +WITH picked AS ( + SELECT id + FROM gateway_task_events + WHERE payload <> '{}'::jsonb + OR phase IS NOT NULL + OR COALESCE(progress, 0) <> 0 + OR message IS NOT NULL + ORDER BY created_at + LIMIT $1 + FOR UPDATE SKIP LOCKED +) +UPDATE gateway_task_events event +SET platform_id = CASE + WHEN event.platform_id IS NULL + AND event.payload->>'platformId' ~* '^[0-9a-f]{8}-[0-9a-f]{4}-[1-5][0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$' + THEN (event.payload->>'platformId')::uuid + ELSE event.platform_id + END, + phase = NULL, + progress = 0, + message = NULL, + payload = '{}'::jsonb +FROM picked +WHERE event.id = picked.id`, []any{batchSize}}, + {&result.CompactedCallbacks, ` +WITH picked AS ( + SELECT id + FROM gateway_task_callback_outbox + WHERE payload <> '{}'::jsonb + ORDER BY created_at + LIMIT $1 + FOR UPDATE SKIP LOCKED +) +UPDATE gateway_task_callback_outbox outbox +SET payload = '{}'::jsonb +FROM picked +WHERE outbox.id = picked.id`, []any{batchSize}}, + {&result.RetiredCallbacks, ` +WITH migration_boundary AS ( + SELECT applied_at + FROM schema_migrations + WHERE version = '0083_task_history_minimal_storage' +), picked AS ( + SELECT outbox.id + FROM gateway_task_callback_outbox outbox + CROSS JOIN migration_boundary boundary + WHERE outbox.created_at < boundary.applied_at + AND outbox.status IN ('pending', 'processing') + ORDER BY outbox.created_at + LIMIT $1 + FOR UPDATE OF outbox SKIP LOCKED +) +UPDATE gateway_task_callback_outbox outbox +SET status = 'failed', + failed_at = outbox.created_at, + last_error = 'legacy_callback_not_replayed', + locked_by = NULL, + lock_token = NULL, + locked_at = NULL, + updated_at = now() +FROM picked +WHERE outbox.id = picked.id`, []any{batchSize}}, + {&result.CompactedParamLogs, ` +WITH picked AS ( + SELECT id + FROM gateway_task_param_preprocessing_logs + WHERE actual_input <> '{}'::jsonb + OR converted_output <> '{}'::jsonb + OR model_snapshot <> '{}'::jsonb + OR octet_length(changes::text) > 1024 + ORDER BY created_at + LIMIT $1 + FOR UPDATE SKIP LOCKED +) +UPDATE gateway_task_param_preprocessing_logs log +SET actual_input = '{}'::jsonb, + converted_output = '{}'::jsonb, + model_snapshot = '{}'::jsonb, + changes = CASE WHEN octet_length(log.changes::text) > 1024 THEN '[]'::jsonb ELSE log.changes END +FROM picked +WHERE log.id = picked.id`, []any{batchSize}}, + {&result.CompactedClonedVoices, ` +WITH picked AS ( + SELECT id + FROM gateway_cloned_voices + WHERE metadata ?| ARRAY['raw', 'raw_data', 'rawData', 'provider_response', 'providerResponse', 'request'] + ORDER BY updated_at + LIMIT $1 + FOR UPDATE SKIP LOCKED +) +UPDATE gateway_cloned_voices voice +SET metadata = COALESCE(voice.metadata, '{}'::jsonb) + - 'raw' - 'raw_data' - 'rawData' + - 'provider_response' - 'providerResponse' - 'request', + updated_at = now() +FROM picked +WHERE voice.id = picked.id`, []any{batchSize}}, + {&result.DeletedCallbacks, ` +WITH picked AS ( + SELECT id + FROM gateway_task_callback_outbox + WHERE (status = 'delivered' AND delivered_at < now() - interval '24 hours') + OR (status = 'failed' AND COALESCE(failed_at, updated_at) < $1::timestamptz) + ORDER BY updated_at + LIMIT $2 + FOR UPDATE SKIP LOCKED +) +DELETE FROM gateway_task_callback_outbox outbox +USING picked +WHERE outbox.id = picked.id`, []any{analysisCutoff, batchSize}}, + {&result.DeletedParamLogs, ` +WITH picked AS ( + SELECT id + FROM gateway_task_param_preprocessing_logs + WHERE created_at < $1::timestamptz + ORDER BY created_at + LIMIT $2 + FOR UPDATE SKIP LOCKED +) +DELETE FROM gateway_task_param_preprocessing_logs log +USING picked +WHERE log.id = picked.id`, []any{analysisCutoff, batchSize}}, + {&result.DeletedEvents, ` +WITH picked AS ( + SELECT event.id + FROM gateway_task_events event + WHERE event.created_at < $1::timestamptz + AND NOT EXISTS ( + SELECT 1 + FROM gateway_task_callback_outbox outbox + WHERE outbox.event_id = event.id + AND outbox.status IN ('pending', 'processing') + ) + ORDER BY event.created_at + LIMIT $2 + FOR UPDATE OF event SKIP LOCKED +) +DELETE FROM gateway_task_events event +USING picked +WHERE event.id = picked.id`, []any{analysisCutoff, batchSize}}, + {&result.DeletedAttempts, ` +WITH picked AS ( + SELECT id + FROM gateway_task_attempts + WHERE COALESCE(finished_at, started_at) < $1::timestamptz + AND status <> 'running' + ORDER BY COALESCE(finished_at, started_at) + LIMIT $2 + FOR UPDATE SKIP LOCKED +) +DELETE FROM gateway_task_attempts attempt +USING picked +WHERE attempt.id = picked.id`, []any{analysisCutoff, batchSize}}, + {&result.DeletedTasks, ` +WITH picked AS ( + SELECT task.id + FROM gateway_tasks task + WHERE task.finished_at < $1::timestamptz + AND task.status IN ('succeeded', 'failed', 'cancelled') + AND task.billing_status IN ('settled', 'released', 'not_required') + AND (task.execution_lease_expires_at IS NULL OR task.execution_lease_expires_at <= now()) + AND NOT EXISTS ( + SELECT 1 + FROM settlement_outbox settlement + WHERE settlement.task_id = task.id + AND settlement.status IN ('pending', 'processing', 'retryable_failed', 'manual_review') + ) + AND NOT EXISTS ( + SELECT 1 + FROM gateway_task_callback_outbox callback + WHERE callback.task_id = task.id + ) + AND NOT EXISTS ( + SELECT 1 + FROM gateway_concurrency_leases lease + WHERE lease.task_id = task.id + AND lease.released_at IS NULL + AND lease.expires_at > now() + ) + AND NOT EXISTS ( + SELECT 1 + FROM gateway_rate_limit_reservations reservation + WHERE reservation.task_id = task.id + AND reservation.status = 'reserved' + ) + ORDER BY task.finished_at + LIMIT $2 + FOR UPDATE OF task SKIP LOCKED +) +DELETE FROM gateway_tasks task +USING picked +WHERE task.id = picked.id`, []any{taskCutoff, batchSize}}, + } + for _, step := range steps { + tag, err := tx.Exec(ctx, step.sql, step.args...) + if err != nil { + return err + } + *step.target = tag.RowsAffected() + } + return nil + }) + return result, err +} diff --git a/apps/api/internal/store/task_history_test.go b/apps/api/internal/store/task_history_test.go new file mode 100644 index 0000000..a954f6c --- /dev/null +++ b/apps/api/internal/store/task_history_test.go @@ -0,0 +1,401 @@ +package store + +import ( + "context" + "strings" + "testing" + "time" + + "github.com/easyai/easyai-ai-gateway/apps/api/internal/auth" + "github.com/google/uuid" +) + +func TestTaskEventAllowlistRejectsSyntheticProgress(t *testing.T) { + if taskEventAllowed("task.progress") || taskEventAllowed("task.attempt.started") || taskEventAllowed("unknown") { + t.Fatal("synthetic or unknown event type was allowed") + } + for _, eventType := range []string{"task.accepted", "task.running", "task.completed", "task.attempt.failed", "task.policy.priority_demoted"} { + if !taskEventAllowed(eventType) { + t.Fatalf("required event type %q was rejected", eventType) + } + } +} + +func TestMinimalTaskResultDropsProviderRawCopy(t *testing.T) { + result := minimalTaskResult(map[string]any{ + "data": []any{map[string]any{"url": "https://example.invalid/result"}}, + "raw": map[string]any{"provider": "duplicate"}, + "raw_data": map[string]any{"provider": "duplicate"}, + "providerResponse": map[string]any{"provider": "duplicate"}, + "upstream_task_id": "remote-task", + "provider_specific": "kept", + }) + if result["raw"] != nil || result["raw_data"] != nil || result["providerResponse"] != nil || + result["upstream_task_id"] != nil || result["data"] == nil || result["provider_specific"] != "kept" { + t.Fatalf("minimal task result=%#v", result) + } +} + +func TestTaskEventsAreMinimalAndConsecutiveDuplicatesAreSkipped(t *testing.T) { + db := billingV2IntegrationStore(t) + ctx := context.Background() + user := &auth.User{ID: "event-minimal-" + uuid.NewString()} + task, err := db.CreateTask(ctx, CreateTaskInput{ + Kind: "images.generations", Model: "event-minimal", RunMode: "simulation", + Request: map[string]any{"prompt": "minimal"}, + }, user) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + _, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_tasks WHERE id=$1::uuid`, task.ID) + }) + + running, err := db.AddTaskEvent(ctx, task.ID, "task.running", "running", "starting", 0.1, "ignored", map[string]any{"large": "ignored"}, true) + if err != nil || running.ID == "" || len(running.Payload) != 0 || running.Phase != "" || running.Progress != 0 || running.Message != "" { + t.Fatalf("running=%+v err=%v", running, err) + } + duplicate, err := db.AddTaskEvent(ctx, task.ID, "task.running", "running", "polling", 0.8, "ignored again", map[string]any{"other": "ignored"}, true) + if err != nil || duplicate.ID != "" || duplicate.SkippedReason != "duplicate" { + t.Fatalf("duplicate=%+v err=%v", duplicate, err) + } + synthetic, err := db.AddTaskEvent(ctx, task.ID, "task.progress", "running", "polling", 0.9, "ignored", nil, true) + if err != nil || synthetic.ID != "" || synthetic.SkippedReason != "unknown_type" { + t.Fatalf("synthetic=%+v err=%v", synthetic, err) + } + completed, err := db.AddTaskEvent(ctx, task.ID, "task.completed", "succeeded", "completed", 1, "ignored", map[string]any{"result": map[string]any{"duplicate": true}}, true) + if err != nil || completed.ID == "" || len(completed.Payload) != 0 { + t.Fatalf("completed=%+v err=%v", completed, err) + } + if err := db.QueueTaskCallback(ctx, completed, "https://callback.invalid/task"); err != nil { + t.Fatal(err) + } + + var events int + var nonEmptyPayloads int + if err := db.pool.QueryRow(ctx, ` +SELECT count(*), count(*) FILTER (WHERE payload <> '{}'::jsonb) +FROM gateway_task_events +WHERE task_id=$1::uuid`, task.ID).Scan(&events, &nonEmptyPayloads); err != nil { + t.Fatal(err) + } + if events != 3 || nonEmptyPayloads != 0 { + t.Fatalf("events=%d nonEmptyPayloads=%d, want 3/0", events, nonEmptyPayloads) + } + var callbackPayloadBytes int + if err := db.pool.QueryRow(ctx, ` +SELECT pg_column_size(payload) +FROM gateway_task_callback_outbox +WHERE task_id=$1::uuid`, task.ID).Scan(&callbackPayloadBytes); err != nil { + t.Fatal(err) + } + if callbackPayloadBytes > 16 { + t.Fatalf("callback payload storage=%d, want an empty json object", callbackPayloadBytes) + } + budgetExceeded := false + for index := 0; index < 20; index++ { + eventType := "task.running" + status := "running" + if index%2 == 1 { + eventType = "task.queued" + status = "queued" + } + event, err := db.AddTaskEvent(ctx, task.ID, eventType, status, "", 0, "", nil, true) + if err != nil { + t.Fatal(err) + } + if event.SkippedReason == "budget_exceeded" { + budgetExceeded = true + break + } + } + if !budgetExceeded { + t.Fatal("non-terminal event budget was not enforced") + } + var nonTerminalEvents int + if err := db.pool.QueryRow(ctx, ` +SELECT count(*) +FROM gateway_task_events +WHERE task_id=$1::uuid + AND event_type NOT IN ( + 'task.completed', 'task.failed', 'task.cancelled', + 'task.billing.settled', 'task.billing.released', 'task.billing.review' + )`, task.ID).Scan(&nonTerminalEvents); err != nil { + t.Fatal(err) + } + if nonTerminalEvents != 16 { + t.Fatalf("non-terminal events=%d, want 16", nonTerminalEvents) + } +} + +func TestTaskAttemptKeepsOnlyCompactDiagnostics(t *testing.T) { + db := billingV2IntegrationStore(t) + ctx := context.Background() + user := &auth.User{ID: "attempt-minimal-" + uuid.NewString()} + task, err := db.CreateTask(ctx, CreateTaskInput{ + Kind: "images.generations", Model: "attempt-minimal", RunMode: "simulation", + Request: map[string]any{"prompt": "canonical"}, + }, user) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + _, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_tasks WHERE id=$1::uuid`, task.ID) + }) + attemptID, err := db.CreateTaskAttempt(ctx, CreateTaskAttemptInput{ + TaskID: task.ID, AttemptNo: 1, QueueKey: "test", Status: "running", + RequestSnapshot: map[string]any{"prompt": "duplicate"}, + Metrics: map[string]any{"trace": []any{map[string]any{"large": "duplicate"}}}, + PricingSnapshot: map[string]any{"price": 1}, + RequestFingerprint: "duplicate-fingerprint", + }) + if err != nil { + t.Fatal(err) + } + if err := db.FinishTaskAttempt(ctx, FinishTaskAttemptInput{ + AttemptID: attemptID, + Status: "failed", + Retryable: true, + StatusCode: 503, + Usage: map[string]any{"output_tokens": 100, "raw": strings.Repeat("duplicate", 100)}, + Metrics: map[string]any{"trace": []any{map[string]any{"large": "duplicate"}}, "cacheAffinityMatched": true, "cacheAffinityMatchedPrefixDepth": 3}, + ResponseSnapshot: map[string]any{"result": "duplicate"}, + ErrorCode: "upstream_error", + ErrorMessage: strings.Repeat("错误", 4096), + }); err != nil { + t.Fatal(err) + } + attempts, err := db.ListTaskAttempts(ctx, task.ID) + if err != nil { + t.Fatal(err) + } + if len(attempts) != 1 { + t.Fatalf("attempts=%d", len(attempts)) + } + attempt := attempts[0] + if len(attempt.RequestSnapshot) != 0 || len(attempt.ResponseSnapshot) != 0 || len(attempt.PricingSnapshot) != 0 { + t.Fatalf("attempt contains duplicate snapshots: %+v", attempt) + } + if len(attempt.Usage) != 1 || taskAttemptMetricInt(attempt.Usage, "output_tokens") != 100 { + t.Fatalf("attempt compact usage mismatch: %+v", attempt.Usage) + } + if len(attempt.Metrics) != 2 || + attempt.Metrics["cacheAffinityMatched"] != true || + taskAttemptMetricInt(attempt.Metrics, "cacheAffinityMatchedPrefixDepth") != 3 { + t.Fatalf("attempt compact metrics mismatch: %+v", attempt.Metrics) + } + if attempt.StatusCode != 503 || len(attempt.ErrorMessage) > 2048 { + t.Fatalf("attempt status/error not preserved safely: %+v", attempt) + } +} + +func TestTaskRunningDoesNotPersistNormalizedRequestCopy(t *testing.T) { + db := billingV2IntegrationStore(t) + ctx := context.Background() + user := &auth.User{ID: "request-minimal-" + uuid.NewString()} + task, err := db.CreateTask(ctx, CreateTaskInput{ + Kind: "images.generations", Model: "request-minimal", RunMode: "simulation", + Request: map[string]any{"prompt": "canonical"}, + }, user) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + _, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_tasks WHERE id=$1::uuid`, task.ID) + }) + executionToken := uuid.NewString() + if _, err := db.ClaimTaskExecution(ctx, task.ID, executionToken, time.Minute); err != nil { + t.Fatal(err) + } + if err := db.MarkTaskRunning(ctx, task.ID, executionToken, "image_generate", map[string]any{ + "prompt": "duplicate normalized request", + }); err != nil { + t.Fatal(err) + } + var normalizedBytes int + if err := db.pool.QueryRow(ctx, ` +SELECT pg_column_size(normalized_request) +FROM gateway_tasks +WHERE id=$1::uuid`, task.ID).Scan(&normalizedBytes); err != nil { + t.Fatal(err) + } + if normalizedBytes > 16 { + t.Fatalf("normalized request storage=%d, want an empty json object", normalizedBytes) + } +} + +func TestTaskHistoryCleanupRemovesHistoricalProviderCopies(t *testing.T) { + db := billingV2IntegrationStore(t) + ctx := context.Background() + user := &auth.User{ID: "history-compaction-" + uuid.NewString()} + task, err := db.CreateTask(ctx, CreateTaskInput{ + Kind: "images.generations", Model: "history-compaction", RunMode: "simulation", + Request: map[string]any{"prompt": "canonical"}, + }, user) + if err != nil { + t.Fatal(err) + } + voiceID := uuid.NewString() + t.Cleanup(func() { + _, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_cloned_voices WHERE id=$1::uuid`, voiceID) + _, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_tasks WHERE id=$1::uuid`, task.ID) + }) + if _, err := db.pool.Exec(ctx, ` +UPDATE gateway_tasks +SET result='{"data":[{"url":"https://example.invalid/result"}],"raw_data":{"provider":"duplicate"},"upstream_task_id":"remote"}'::jsonb +WHERE id=$1::uuid`, + task.ID); err != nil { + t.Fatal(err) + } + if _, err := db.pool.Exec(ctx, ` +INSERT INTO gateway_cloned_voices (id, user_id, provider, voice_id, metadata) +VALUES ($1::uuid, 'history-compaction', 'test', $2, '{"request":{"duplicate":true},"rawData":{"provider":"duplicate"},"keep":"value"}'::jsonb)`, + voiceID, "voice-"+voiceID); err != nil { + t.Fatal(err) + } + result, err := db.CleanupTaskHistory(ctx, time.Now().AddDate(0, 0, -7), time.Now().AddDate(0, 0, -30), 100) + if err != nil { + t.Fatal(err) + } + if result.CompactedTasks < 1 || result.CompactedClonedVoices < 1 { + t.Fatalf("cleanup result=%+v", result) + } + var taskResult map[string]any + var voiceMetadata map[string]any + if err := db.pool.QueryRow(ctx, `SELECT result FROM gateway_tasks WHERE id=$1::uuid`, task.ID).Scan(&taskResult); err != nil { + t.Fatal(err) + } + if err := db.pool.QueryRow(ctx, `SELECT metadata FROM gateway_cloned_voices WHERE id=$1::uuid`, voiceID).Scan(&voiceMetadata); err != nil { + t.Fatal(err) + } + if taskResult["raw_data"] != nil || taskResult["upstream_task_id"] != nil || taskResult["data"] == nil { + t.Fatalf("compacted task result=%+v", taskResult) + } + if voiceMetadata["request"] != nil || voiceMetadata["rawData"] != nil || voiceMetadata["keep"] != "value" { + t.Fatalf("compacted voice metadata=%+v", voiceMetadata) + } +} + +func TestTaskHistoryCleanupKeepsTaskUntilCallbackRetentionExpires(t *testing.T) { + db := billingV2IntegrationStore(t) + ctx := context.Background() + user := &auth.User{ID: "cleanup-" + uuid.NewString()} + createOldTask := func(label string) GatewayTask { + t.Helper() + task, err := db.CreateTask(ctx, CreateTaskInput{ + Kind: "images.generations", Model: label, RunMode: "simulation", + Request: map[string]any{"prompt": label}, + }, user) + if err != nil { + t.Fatal(err) + } + if _, err := db.pool.Exec(ctx, ` +UPDATE gateway_tasks +SET status='succeeded', billing_status='not_required', + finished_at=now()-interval '40 days', updated_at=now()-interval '40 days' +WHERE id=$1::uuid`, task.ID); err != nil { + t.Fatal(err) + } + if _, err := db.pool.Exec(ctx, ` +UPDATE gateway_task_events +SET created_at=now()-interval '10 days' +WHERE task_id=$1::uuid`, task.ID); err != nil { + t.Fatal(err) + } + return task + } + safe := createOldTask("cleanup-safe") + protected := createOldTask("cleanup-protected") + t.Cleanup(func() { + _, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_tasks WHERE id IN ($1::uuid, $2::uuid)`, safe.ID, protected.ID) + }) + completed, err := db.AddTaskEvent(ctx, protected.ID, "task.completed", "succeeded", "", 0, "", nil, true) + if err != nil { + t.Fatal(err) + } + if err := db.QueueTaskCallback(ctx, completed, "https://callback.invalid/task"); err != nil { + t.Fatal(err) + } + if _, err := db.pool.Exec(ctx, ` +UPDATE gateway_task_callback_outbox +SET status='delivered', delivered_at=now(), updated_at=now() +WHERE task_id=$1::uuid`, protected.ID); err != nil { + t.Fatal(err) + } + + result, err := db.CleanupTaskHistory(ctx, time.Now().AddDate(0, 0, -7), time.Now().AddDate(0, 0, -30), 100) + if err != nil { + t.Fatal(err) + } + if result.DeletedTasks < 1 { + t.Fatalf("deletedTasks=%d, want at least 1", result.DeletedTasks) + } + var safeExists bool + var protectedExists bool + if err := db.pool.QueryRow(ctx, `SELECT EXISTS (SELECT 1 FROM gateway_tasks WHERE id=$1::uuid)`, safe.ID).Scan(&safeExists); err != nil { + t.Fatal(err) + } + if err := db.pool.QueryRow(ctx, `SELECT EXISTS (SELECT 1 FROM gateway_tasks WHERE id=$1::uuid)`, protected.ID).Scan(&protectedExists); err != nil { + t.Fatal(err) + } + if safeExists || !protectedExists { + t.Fatalf("safeExists=%t protectedExists=%t", safeExists, protectedExists) + } +} + +func TestLegacyCallbacksAreNotReplayed(t *testing.T) { + db := billingV2IntegrationStore(t) + ctx := context.Background() + user := &auth.User{ID: "legacy-callback-" + uuid.NewString()} + task, err := db.CreateTask(ctx, CreateTaskInput{ + Kind: "images.generations", Model: "legacy-callback", RunMode: "simulation", + Request: map[string]any{"prompt": "legacy"}, + }, user) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + _, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_tasks WHERE id=$1::uuid`, task.ID) + }) + event, err := db.AddTaskEvent(ctx, task.ID, "task.running", "running", "", 0, "", nil, true) + if err != nil { + t.Fatal(err) + } + if err := db.QueueTaskCallback(ctx, event, "https://callback.invalid/legacy"); err != nil { + t.Fatal(err) + } + if _, err := db.pool.Exec(ctx, ` +UPDATE gateway_task_callback_outbox +SET created_at=( + SELECT applied_at - interval '1 second' + FROM schema_migrations + WHERE version='0083_task_history_minimal_storage' +) +WHERE task_id=$1::uuid`, task.ID); err != nil { + t.Fatal(err) + } + claimed, err := db.ClaimTaskCallbacks(ctx, "legacy-test", 10, time.Minute) + if err != nil { + t.Fatal(err) + } + if len(claimed) != 0 { + t.Fatalf("claimed legacy callbacks=%d, want 0", len(claimed)) + } + result, err := db.CleanupTaskHistory(ctx, time.Now().AddDate(0, 0, -7), time.Now().AddDate(0, 0, -30), 100) + if err != nil { + t.Fatal(err) + } + if result.RetiredCallbacks < 1 { + t.Fatalf("retired legacy callbacks=%d, want at least 1", result.RetiredCallbacks) + } + var status string + if err := db.pool.QueryRow(ctx, ` +SELECT status +FROM gateway_task_callback_outbox +WHERE task_id=$1::uuid`, task.ID).Scan(&status); err != nil { + t.Fatal(err) + } + if status != "failed" { + t.Fatalf("legacy callback status=%q, want failed", status) + } +} diff --git a/apps/api/internal/store/tasks_runtime.go b/apps/api/internal/store/tasks_runtime.go index a23f811..03f9139 100644 --- a/apps/api/internal/store/tasks_runtime.go +++ b/apps/api/internal/store/tasks_runtime.go @@ -8,6 +8,7 @@ import ( "strconv" "strings" "time" + "unicode/utf8" "github.com/easyai/easyai-ai-gateway/apps/api/internal/auth" "github.com/jackc/pgx/v5" @@ -324,7 +325,7 @@ UPDATE gateway_tasks SET status = 'failed', billing_status = CASE WHEN $2 THEN 'manual_review' ELSE 'not_required' END, billing_updated_at = now(), - error = 'upstream submission result is unknown', + error = NULL, error_code = 'upstream_submission_unknown', error_message = 'upstream submission result is unknown', locked_by = NULL, @@ -332,6 +333,7 @@ SET status = 'failed', heartbeat_at = NULL, execution_token = NULL, execution_lease_expires_at = NULL, + remote_task_payload = '{}'::jsonb, finished_at = now(), updated_at = now() WHERE id = $1::uuid`, taskID, hasGatewayUser); err != nil { @@ -416,19 +418,18 @@ SELECT COALESCE((SELECT status = 'running' FROM gateway_tasks WHERE id = $1::uui return nil } -func (s *Store) MarkTaskRunning(ctx context.Context, taskID string, executionToken string, modelType string, normalizedRequest map[string]any) error { - normalizedJSON, _ := json.Marshal(emptyObjectIfNil(normalizedRequest)) +func (s *Store) MarkTaskRunning(ctx context.Context, taskID string, executionToken string, modelType string, _ map[string]any) error { tag, err := s.pool.Exec(ctx, ` UPDATE gateway_tasks SET status = 'running', model_type = NULLIF($3::text, ''), - normalized_request = $4::jsonb, + normalized_request = '{}'::jsonb, locked_at = now(), heartbeat_at = now(), updated_at = now() WHERE id = $1::uuid AND status = 'running' - AND execution_token = $2::uuid`, taskID, executionToken, modelType, string(normalizedJSON)) + AND execution_token = $2::uuid`, taskID, executionToken, modelType) if err != nil { return err } @@ -469,6 +470,13 @@ func (s *Store) RequeueTask(ctx context.Context, taskID string, executionToken s delay = 10 * time.Minute } nextRunAt := time.Now().Add(delay) + return s.RequeueTaskUntil(ctx, taskID, executionToken, nextRunAt, queueKey) +} + +func (s *Store) RequeueTaskUntil(ctx context.Context, taskID string, executionToken string, nextRunAt time.Time, queueKey string) (GatewayTask, error) { + if nextRunAt.Before(time.Now().Add(time.Second)) { + nextRunAt = time.Now().Add(time.Second) + } return scanGatewayTask(s.pool.QueryRow(ctx, ` UPDATE gateway_tasks SET status = 'queued', @@ -502,7 +510,7 @@ WHERE id = $1::uuid`, taskID, riverJobID) } func (s *Store) SetTaskRemoteTask(ctx context.Context, taskID string, executionToken string, attemptID string, remoteTaskID string, payload map[string]any) error { - payloadJSON, _ := json.Marshal(emptyObjectIfNil(payload)) + payloadJSON, _ := json.Marshal(sanitizeJSONForStorage(emptyObjectIfNil(payload))) return pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error { tag, err := tx.Exec(ctx, ` UPDATE gateway_tasks @@ -557,12 +565,14 @@ func (s *Store) CancelQueuedTask(ctx context.Context, taskID string, message str if message == "" { message = "任务已取消" } + message = truncateUTF8Bytes(message, 2048) tag, err := s.pool.Exec(ctx, ` UPDATE gateway_tasks SET status = 'cancelled', - error = NULLIF($2::text, ''), + error = NULL, error_code = 'task_cancelled', error_message = NULLIF($2::text, ''), + remote_task_payload = '{}'::jsonb, locked_by = NULL, locked_at = NULL, heartbeat_at = NULL, @@ -592,6 +602,7 @@ func (s *Store) CancelSubmittedTask(ctx context.Context, taskID string, executio if message == "" { message = "任务已由上游取消" } + message = truncateUTF8Bytes(message, 2048) var task GatewayTask changed := false err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error { @@ -599,7 +610,7 @@ func (s *Store) CancelSubmittedTask(ctx context.Context, taskID string, executio task, err = scanGatewayTask(tx.QueryRow(ctx, ` UPDATE gateway_tasks SET status = 'cancelled', - error = NULLIF($2, ''), + error = NULL, error_code = 'task_cancelled', error_message = NULLIF($2, ''), billing_status = CASE @@ -613,6 +624,7 @@ SET status = 'cancelled', heartbeat_at = NULL, execution_token = NULL, execution_lease_expires_at = NULL, + remote_task_payload = '{}'::jsonb, finished_at = now(), updated_at = now() WHERE id = $1::uuid @@ -683,9 +695,6 @@ LIMIT $1`, limit) } func (s *Store) CreateTaskAttempt(ctx context.Context, input CreateTaskAttemptInput) (string, error) { - requestJSON, _ := json.Marshal(emptyObjectIfNil(input.RequestSnapshot)) - metricsJSON, _ := json.Marshal(emptyObjectIfNil(input.Metrics)) - pricingJSON, _ := json.Marshal(emptyObjectIfNil(input.PricingSnapshot)) tx, err := s.pool.Begin(ctx) if err != nil { return "", err @@ -701,7 +710,7 @@ INSERT INTO gateway_task_attempts ( ) VALUES ( $1::uuid, $2::int, NULLIF($3::text, '')::uuid, NULLIF($4::text, '')::uuid, NULLIF($5::text, ''), $6, - $7, $8, $9::jsonb, $10::jsonb, $11::jsonb, NULLIF($12, ''), + $7, $8, '{}'::jsonb, '{}'::jsonb, '{}'::jsonb, NULL, 'not_submitted', now() ) RETURNING id::text`, @@ -713,10 +722,6 @@ RETURNING id::text`, input.QueueKey, firstNonEmpty(input.Status, "running"), input.Simulated, - string(requestJSON), - string(metricsJSON), - string(pricingJSON), - input.RequestFingerprint, ).Scan(&attemptID) if err != nil { return "", err @@ -731,13 +736,16 @@ WHERE id = $1::uuid`, input.TaskID, input.AttemptNo); err != nil { } func (s *Store) CreateTaskParamPreprocessingLog(ctx context.Context, input CreateTaskParamPreprocessingLogInput) (string, error) { - actualInputJSON, _ := json.Marshal(emptyObjectIfNil(input.ActualInput)) - convertedOutputJSON, _ := json.Marshal(emptyObjectIfNil(input.ConvertedOutput)) - changesJSON, _ := json.Marshal(input.Changes) + if !input.Changed { + return "", nil + } + changesJSON, _ := json.Marshal(sanitizeJSONForStorage(input.Changes)) if input.Changes == nil { changesJSON = []byte("[]") } - modelSnapshotJSON, _ := json.Marshal(emptyObjectIfNil(input.ModelSnapshot)) + if len(changesJSON) > 1024 { + changesJSON = []byte("[]") + } var attemptNo any if input.AttemptNo > 0 { attemptNo = input.AttemptNo @@ -751,7 +759,7 @@ INSERT INTO gateway_task_param_preprocessing_logs ( VALUES ( $1::uuid, NULLIF($2::text, '')::uuid, $3::int, NULLIF($4::text, ''), NULLIF($5::text, '')::uuid, NULLIF($6::text, '')::uuid, NULLIF($7::text, ''), - $8, $9::int, $10::jsonb, $11::jsonb, $12::jsonb, $13::jsonb + true, $8::int, '{}'::jsonb, '{}'::jsonb, $9::jsonb, '{}'::jsonb ) RETURNING id::text`, input.TaskID, @@ -761,12 +769,8 @@ RETURNING id::text`, input.PlatformID, input.PlatformModelID, input.ClientID, - input.Changed, input.ChangeCount, - string(actualInputJSON), - string(convertedOutputJSON), string(changesJSON), - string(modelSnapshotJSON), ).Scan(&id) return id, err } @@ -823,24 +827,30 @@ ORDER BY COALESCE(attempt_no, 0), created_at`, taskID) } func (s *Store) AppendTaskAttemptTrace(ctx context.Context, taskID string, attemptNo int, entry map[string]any) error { - entryJSON, _ := json.Marshal(emptyObjectIfNil(entry)) - _, err := s.pool.Exec(ctx, ` + if strings.TrimSpace(taskID) == "" || attemptNo <= 0 || len(entry) == 0 { + return nil + } + encoded, err := json.Marshal(entry) + if err != nil { + return err + } + result, err := s.pool.Exec(ctx, ` UPDATE gateway_task_attempts SET metrics = jsonb_set( COALESCE(metrics, '{}'::jsonb), '{trace}', - ( - CASE - WHEN jsonb_typeof(COALESCE(metrics->'trace', '[]'::jsonb)) = 'array' - THEN COALESCE(metrics->'trace', '[]'::jsonb) - ELSE '[]'::jsonb - END - ) || jsonb_build_array($3::jsonb), + COALESCE(metrics->'trace', '[]'::jsonb) || jsonb_build_array($3::jsonb), true ) WHERE task_id = $1::uuid - AND attempt_no = $2::int`, taskID, attemptNo, string(entryJSON)) - return err + AND attempt_no = $2::int`, taskID, attemptNo, encoded) + if err != nil { + return err + } + if result.RowsAffected() == 0 { + return pgx.ErrNoRows + } + return nil } func (s *Store) listTaskAttemptsByTaskIDs(ctx context.Context, taskIDs []string) (map[string][]TaskAttempt, error) { @@ -855,7 +865,8 @@ SELECT a.id::text, a.task_id::text, a.attempt_no, COALESCE(NULLIF(pm.provider_model_name, ''), pm.model_name, ''), COALESCE(pm.model_alias, ''), COALESCE(a.client_id, ''), a.queue_key, a.status, a.retryable, a.simulated, - COALESCE(a.request_id, ''), COALESCE(a.usage, '{}'::jsonb), COALESCE(a.metrics, '{}'::jsonb), + COALESCE(a.request_id, ''), COALESCE(a.status_code, 0), + COALESCE(a.usage, '{}'::jsonb), COALESCE(a.metrics, '{}'::jsonb), a.request_snapshot, COALESCE(a.response_snapshot, '{}'::jsonb), COALESCE(a.response_started_at::text, ''), COALESCE(a.response_finished_at::text, ''), COALESCE(a.response_duration_ms, 0), COALESCE(a.error_code, ''), COALESCE(a.error_message, ''), @@ -908,6 +919,7 @@ func scanTaskAttempt(scanner taskScanner) (TaskAttempt, error) { &item.Retryable, &item.Simulated, &item.RequestID, + &item.StatusCode, &usageBytes, &metricsBytes, &requestBytes, @@ -980,7 +992,9 @@ func enrichTaskAttemptFromMetrics(item *TaskAttempt) { item.ModelAlias = firstNonEmpty(item.ModelAlias, taskAttemptMetricString(item.Metrics, "modelAlias")) item.ModelType = firstNonEmpty(item.ModelType, taskAttemptMetricString(item.Metrics, "modelType")) item.ClientID = firstNonEmpty(item.ClientID, taskAttemptMetricString(item.Metrics, "clientId")) - item.StatusCode = taskAttemptMetricInt(item.Metrics, "statusCode") + if item.StatusCode == 0 { + item.StatusCode = taskAttemptMetricInt(item.Metrics, "statusCode") + } } func taskAttemptMetricString(metrics map[string]any, key string) string { @@ -1008,47 +1022,96 @@ func taskAttemptMetricInt(metrics map[string]any, key string) int { } func (s *Store) FinishTaskAttempt(ctx context.Context, input FinishTaskAttemptInput) error { - responseJSON, _ := json.Marshal(emptyObjectIfNil(input.ResponseSnapshot)) - usageJSON, _ := json.Marshal(emptyObjectIfNil(input.Usage)) - metricsJSON, _ := json.Marshal(emptyObjectIfNil(input.Metrics)) + statusCode := input.StatusCode + if statusCode == 0 { + statusCode = taskAttemptMetricInt(input.Metrics, "statusCode") + } + usageJSON, _ := json.Marshal(sanitizeJSONForStorage(minimalTaskAttemptUsage(input.Usage))) + metricsJSON, _ := json.Marshal(sanitizeJSONForStorage(minimalTaskAttemptMetrics(input.Metrics))) _, err := s.pool.Exec(ctx, ` UPDATE gateway_task_attempts SET status = $2::text, retryable = $3, request_id = NULLIF($4::text, ''), - usage = $5::jsonb, - metrics = $6::jsonb, - response_snapshot = $7::jsonb, - response_started_at = $8::timestamptz, - response_finished_at = $9::timestamptz, - response_duration_ms = $10, - error_code = NULLIF($11::text, ''), - error_message = NULLIF($12::text, ''), + status_code = NULLIF($5::int, 0), + usage = $11::jsonb, + metrics = $12::jsonb, + response_snapshot = '{}'::jsonb, + pricing_snapshot = '{}'::jsonb, + request_fingerprint = NULL, + response_started_at = $6::timestamptz, + response_finished_at = $7::timestamptz, + response_duration_ms = $8, + error_code = NULLIF($9::text, ''), + error_message = NULLIF(left($10::text, 2048), ''), finished_at = now() WHERE id = $1::uuid`, input.AttemptID, input.Status, input.Retryable, input.RequestID, - string(usageJSON), - string(metricsJSON), - string(responseJSON), + statusCode, nullableTime(input.ResponseStartedAt), nullableTime(input.ResponseFinishedAt), input.ResponseDurationMS, input.ErrorCode, - input.ErrorMessage, + truncateUTF8Bytes(input.ErrorMessage, 2048), + usageJSON, + metricsJSON, ) return err } +func minimalTaskAttemptUsage(usage map[string]any) map[string]any { + return whitelistedTaskAttemptMap(usage, []string{ + "inputTokens", "promptTokens", "input_tokens", "prompt_tokens", + "cachedInputTokens", "cachedPromptTokens", "cached_input_tokens", "cached_tokens", + "cachedInputTokensKnown", + "outputTokens", "completionTokens", "output_tokens", "completion_tokens", + "totalTokens", "total_tokens", + }) +} + +func minimalTaskAttemptMetrics(metrics map[string]any) map[string]any { + return whitelistedTaskAttemptMap(metrics, []string{ + "platformPriority", "currentPriority", "loadRatio", "loadAvoided", + "cacheAffinityKey", "cacheAdjustedPriority", "cacheAffinitySamples", + "cacheAffinityScore", "cacheAffinityConfidence", "cacheAffinityHitRatio", + "cacheAffinityEMAHitRatio", "cacheAffinityLastHitRatio", + "cacheAffinityCachedInputTokens", "cacheAffinityBoost", + "cacheAffinityApplied", "cacheAffinityMatched", + "cacheAffinityCandidateCount", "cacheAffinityMatchedPrefixDepth", + "cacheAffinityOverrideReason", + }) +} + +func whitelistedTaskAttemptMap(input map[string]any, keys []string) map[string]any { + out := map[string]any{} + for _, key := range keys { + if value, ok := input[key]; ok { + out[key] = value + } + } + return out +} + func (s *Store) FinishTaskSuccess(ctx context.Context, input FinishTaskSuccessInput) (GatewayTask, error) { - resultJSON, _ := json.Marshal(emptyObjectIfNil(input.Result)) - billingsJSON, _ := json.Marshal(input.Billings) - usageJSON, _ := json.Marshal(emptyObjectIfNil(input.Usage)) - metricsJSON, _ := json.Marshal(emptyObjectIfNil(input.Metrics)) - billingSummaryJSON, _ := json.Marshal(emptyObjectIfNil(input.BillingSummary)) - pricingSnapshotJSON, _ := json.Marshal(emptyObjectIfNil(input.PricingSnapshot)) + resultReport := sanitizeJSONForStorageWithReport(minimalTaskResult(input.Result)) + if resultReport.BinaryCount > 0 { + return GatewayTask{}, &taskPayloadBinaryError{ + target: ErrTaskResultBinaryNotMaterialized, + code: "result_binary_not_materialized", + count: resultReport.BinaryCount, + } + } + resultJSON, _ := json.Marshal(resultReport.Value) + billingsJSON, _ := json.Marshal(sanitizeJSONForStorage(input.Billings)) + usageJSON, _ := json.Marshal(sanitizeJSONForStorage(emptyObjectIfNil(input.Usage))) + metricsJSON, _ := json.Marshal(sanitizeJSONForStorage(emptyObjectIfNil(input.Metrics))) + attemptUsageJSON, _ := json.Marshal(sanitizeJSONForStorage(minimalTaskAttemptUsage(input.Usage))) + attemptMetricsJSON, _ := json.Marshal(sanitizeJSONForStorage(minimalTaskAttemptMetrics(input.Metrics))) + billingSummaryJSON, _ := json.Marshal(sanitizeJSONForStorage(emptyObjectIfNil(input.BillingSummary))) + pricingSnapshotJSON, _ := json.Marshal(sanitizeJSONForStorage(emptyObjectIfNil(input.PricingSnapshot))) finalChargeAmount := strings.TrimSpace(input.FinalChargeAmountText) if finalChargeAmount == "" { finalChargeAmount = strconv.FormatFloat(input.FinalChargeAmount, 'f', 9, 64) @@ -1061,23 +1124,24 @@ UPDATE gateway_task_attempts SET status = 'succeeded', retryable = false, request_id = NULLIF($2, ''), - usage = $3::jsonb, - metrics = $4::jsonb, - response_snapshot = $5::jsonb, - pricing_snapshot = $6::jsonb, - request_fingerprint = NULLIF($7, ''), + status_code = NULL, + usage = $6::jsonb, + metrics = $7::jsonb, + response_snapshot = '{}'::jsonb, + pricing_snapshot = '{}'::jsonb, + request_fingerprint = NULL, upstream_submission_status = 'response_received', upstream_submission_updated_at = now(), - response_started_at = $8::timestamptz, - response_finished_at = $9::timestamptz, - response_duration_ms = $10, + response_started_at = $3::timestamptz, + response_finished_at = $4::timestamptz, + response_duration_ms = $5, error_code = NULL, error_message = NULL, finished_at = now() WHERE id = $1::uuid`, - input.AttemptID, input.RequestID, string(usageJSON), string(metricsJSON), - string(resultJSON), string(pricingSnapshotJSON), input.RequestFingerprint, + input.AttemptID, input.RequestID, nullableTime(input.ResponseStartedAt), nullableTime(input.ResponseFinishedAt), input.ResponseDurationMS, + attemptUsageJSON, attemptMetricsJSON, ); err != nil { return err } @@ -1113,6 +1177,7 @@ SET status = 'succeeded', heartbeat_at = NULL, execution_token = NULL, execution_lease_expires_at = NULL, + remote_task_payload = '{}'::jsonb, finished_at = now(), updated_at = now() WHERE id = $1::uuid @@ -1166,8 +1231,21 @@ func (s *Store) FinishTaskManualReview(ctx context.Context, input FinishTaskManu if status != "succeeded" && status != "failed" { return GatewayTask{}, fmt.Errorf("manual review task status must be succeeded or failed") } - resultJSON, _ := json.Marshal(emptyObjectIfNil(input.Result)) - pricingSnapshotJSON, _ := json.Marshal(emptyObjectIfNil(input.PricingSnapshot)) + message := truncateUTF8Bytes(input.Message, 2048) + result := input.Result + if status == "failed" { + result = nil + } + resultReport := sanitizeJSONForStorageWithReport(minimalTaskResult(result)) + if resultReport.BinaryCount > 0 { + return GatewayTask{}, &taskPayloadBinaryError{ + target: ErrTaskResultBinaryNotMaterialized, + code: "result_binary_not_materialized", + count: resultReport.BinaryCount, + } + } + resultJSON, _ := json.Marshal(resultReport.Value) + pricingSnapshotJSON, _ := json.Marshal(sanitizeJSONForStorage(emptyObjectIfNil(input.PricingSnapshot))) err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error { if strings.TrimSpace(input.AttemptID) != "" { attemptStatus := "failed" @@ -1186,7 +1264,7 @@ SET status = $2, response_finished_at = $7::timestamptz, response_duration_ms = $8, finished_at = now() -WHERE id = $1::uuid`, input.AttemptID, attemptStatus, input.RequestID, input.Code, input.Message, +WHERE id = $1::uuid`, input.AttemptID, attemptStatus, input.RequestID, input.Code, message, nullableTime(input.ResponseStartedAt), nullableTime(input.ResponseFinishedAt), input.ResponseDurationMS); err != nil { return err } @@ -1196,7 +1274,7 @@ UPDATE gateway_tasks SET status = $2, result = $3::jsonb, request_id = NULLIF($4, ''), - error = CASE WHEN $2 = 'failed' THEN NULLIF($6, '') ELSE NULL END, + error = NULL, error_code = NULLIF($5, ''), error_message = NULLIF($6, ''), billing_status = 'manual_review', @@ -1211,11 +1289,12 @@ SET status = $2, heartbeat_at = NULL, execution_token = NULL, execution_lease_expires_at = NULL, + remote_task_payload = '{}'::jsonb, finished_at = now(), updated_at = now() WHERE id = $1::uuid AND status = 'running' - AND execution_token = $12::uuid`, input.TaskID, status, string(resultJSON), input.RequestID, input.Code, input.Message, + AND execution_token = $12::uuid`, input.TaskID, status, string(resultJSON), input.RequestID, input.Code, message, string(pricingSnapshotJSON), input.RequestFingerprint, nullableTime(input.ResponseStartedAt), nullableTime(input.ResponseFinishedAt), input.ResponseDurationMS, input.ExecutionToken) if err != nil { @@ -1268,7 +1347,7 @@ func (s *Store) SettleTaskBilling(ctx context.Context, task GatewayTask) error { "billings": task.Billings, "billingSummary": task.BillingSummary, } - metadata, _ := json.Marshal(metadataMap) + metadata, _ := json.Marshal(sanitizeJSONForStorage(metadataMap)) return pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error { if _, err := tx.Exec(ctx, ` INSERT INTO gateway_wallet_accounts ( @@ -1360,7 +1439,7 @@ ON CONFLICT (account_id, idempotency_key) WHERE idempotency_key IS NOT NULL DO N "frozenBefore": roundMoney(frozenBefore), "frozenAfter": frozenAfter, }) - metadata, _ = json.Marshal(billingMetadata) + metadata, _ = json.Marshal(sanitizeJSONForStorage(billingMetadata)) if _, err := tx.Exec(ctx, ` INSERT INTO gateway_wallet_transactions ( account_id, gateway_tenant_id, gateway_user_id, direction, transaction_type, @@ -1396,13 +1475,14 @@ func taskBillingString(value any) string { } func (s *Store) FinishTaskFailure(ctx context.Context, input FinishTaskFailureInput) (GatewayTask, error) { - metricsJSON, _ := json.Marshal(emptyObjectIfNil(input.Metrics)) - resultJSON, _ := json.Marshal(emptyObjectIfNil(input.Result)) + metricsJSON, _ := json.Marshal(sanitizeJSONForStorage(emptyObjectIfNil(input.Metrics))) + resultJSON, _ := json.Marshal(minimalTaskResult(nil)) + message := truncateUTF8Bytes(input.Message, 2048) err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error { tag, err := tx.Exec(ctx, ` UPDATE gateway_tasks SET status = 'failed', - error = NULLIF($2::text, ''), + error = NULL, error_code = NULLIF($3::text, ''), error_message = NULLIF($2::text, ''), request_id = NULLIF($4::text, ''), @@ -1422,13 +1502,14 @@ func (s *Store) FinishTaskFailure(ctx context.Context, input FinishTaskFailureIn heartbeat_at = NULL, execution_token = NULL, execution_lease_expires_at = NULL, + remote_task_payload = '{}'::jsonb, finished_at = now(), updated_at = now() WHERE id = $1::uuid AND status = 'running' AND execution_token = $10::uuid`, input.TaskID, - input.Message, + message, input.Code, input.RequestID, string(metricsJSON), @@ -1472,29 +1553,152 @@ func nullableTime(value time.Time) any { return value } +func truncateUTF8Bytes(value string, maxBytes int) string { + if maxBytes <= 0 || len(value) <= maxBytes { + return value + } + end := maxBytes + for end > 0 && !utf8.ValidString(value[:end]) { + end-- + } + return value[:end] +} + +func minimalTaskResult(input map[string]any) map[string]any { + result := make(map[string]any, len(input)) + for key, value := range input { + switch key { + case "raw", + "raw_data", + "rawData", + "provider_response", + "providerResponse", + "submit", + "file_retrieve", + "fileRetrieve", + "upstream_task_id", + "remote_task_id": + continue + } + result[key] = value + } + return result +} + +var allowedTaskEventTypes = map[string]struct{}{ + "task.accepted": {}, + "task.running": {}, + "task.queued": {}, + "task.completed": {}, + "task.failed": {}, + "task.cancelled": {}, + "task.attempt.failed": {}, + "task.billing.settled": {}, + "task.billing.released": {}, + "task.billing.review": {}, + "task.policy.auto_disabled": {}, + "task.policy.degraded": {}, + "task.policy.failover_disabled": {}, + "task.policy.failover_cooled_down": {}, + "task.policy.priority_demoted": {}, +} + +func taskEventAllowed(eventType string) bool { + _, ok := allowedTaskEventTypes[strings.TrimSpace(eventType)] + return ok +} + +func taskEventTerminal(eventType string) bool { + switch strings.TrimSpace(eventType) { + case "task.completed", "task.failed", "task.cancelled", + "task.billing.settled", "task.billing.released", "task.billing.review": + return true + default: + return false + } +} + +func taskEventPlatformID(payload map[string]any) string { + platformID, _ := payload["platformId"].(string) + return strings.TrimSpace(platformID) +} + func (s *Store) AddTaskEvent(ctx context.Context, taskID string, eventType string, status string, phase string, progress float64, message string, payload map[string]any, simulated bool) (TaskEvent, error) { - payloadJSON, _ := json.Marshal(emptyObjectIfNil(payload)) + eventType = strings.TrimSpace(eventType) + if !taskEventAllowed(eventType) { + return TaskEvent{SkippedReason: "unknown_type"}, nil + } + platformID := taskEventPlatformID(payload) + tx, err := s.pool.Begin(ctx) + if err != nil { + return TaskEvent{}, err + } + defer tx.Rollback(ctx) + if _, err := tx.Exec(ctx, `SELECT 1 FROM gateway_tasks WHERE id = $1::uuid FOR UPDATE`, taskID); err != nil { + return TaskEvent{}, err + } + var previousType string + var previousStatus string + var previousPlatformID string + var previousSimulated bool + err = tx.QueryRow(ctx, ` +SELECT event_type, COALESCE(status, ''), COALESCE(platform_id::text, ''), simulated +FROM gateway_task_events +WHERE task_id = $1::uuid +ORDER BY seq DESC +LIMIT 1`, taskID).Scan(&previousType, &previousStatus, &previousPlatformID, &previousSimulated) + if err != nil && !errors.Is(err, pgx.ErrNoRows) { + return TaskEvent{}, err + } + if err == nil && + previousType == eventType && + previousStatus == strings.TrimSpace(status) && + previousPlatformID == platformID && + previousSimulated == simulated { + if err := tx.Commit(ctx); err != nil { + return TaskEvent{}, err + } + return TaskEvent{SkippedReason: "duplicate"}, nil + } + if !taskEventTerminal(eventType) { + var nonTerminalCount int + if err := tx.QueryRow(ctx, ` +SELECT count(*) +FROM gateway_task_events +WHERE task_id = $1::uuid + AND event_type NOT IN ( + 'task.completed', 'task.failed', 'task.cancelled', + 'task.billing.settled', 'task.billing.released', 'task.billing.review' + )`, taskID).Scan(&nonTerminalCount); err != nil { + return TaskEvent{}, err + } + if nonTerminalCount >= 16 { + if err := tx.Commit(ctx); err != nil { + return TaskEvent{}, err + } + return TaskEvent{SkippedReason: "budget_exceeded"}, nil + } + } var event TaskEvent var payloadBytes []byte - err := s.pool.QueryRow(ctx, ` + err = tx.QueryRow(ctx, ` WITH next_seq AS ( SELECT COALESCE(MAX(seq), 0) + 1 AS seq FROM gateway_task_events WHERE task_id = $1::uuid ) -INSERT INTO gateway_task_events (task_id, seq, event_type, status, phase, progress, message, payload, simulated) -SELECT $1::uuid, next_seq.seq, $2::text, NULLIF($3::text, ''), NULLIF($4::text, ''), $5, NULLIF($6::text, ''), $7::jsonb, $8 +INSERT INTO gateway_task_events (task_id, seq, event_type, status, phase, progress, message, payload, simulated, platform_id) +SELECT $1::uuid, next_seq.seq, $2::text, NULLIF($3::text, ''), NULL, 0, NULL, '{}'::jsonb, $4, + NULLIF($5::text, '')::uuid FROM next_seq RETURNING id::text, task_id::text, seq, event_type, COALESCE(status, ''), COALESCE(phase, ''), - COALESCE(progress, 0)::float8, COALESCE(message, ''), payload, simulated, created_at`, + COALESCE(progress, 0)::float8, COALESCE(message, ''), payload, simulated, created_at, + COALESCE(platform_id::text, '')`, taskID, eventType, status, - phase, - progress, - message, - string(payloadJSON), simulated, + platformID, ).Scan( &event.ID, &event.TaskID, @@ -1507,39 +1711,30 @@ RETURNING id::text, task_id::text, seq, event_type, COALESCE(status, ''), COALES &payloadBytes, &event.Simulated, &event.CreatedAt, + &event.PlatformID, ) if err != nil { return TaskEvent{}, err } event.Payload = decodeObject(payloadBytes) + if err := tx.Commit(ctx); err != nil { + return TaskEvent{}, err + } return event, nil } func (s *Store) QueueTaskCallback(ctx context.Context, event TaskEvent, callbackURL string) error { - if callbackURL == "" { + if callbackURL == "" || event.ID == "" { return nil } - payloadJSON, _ := json.Marshal(map[string]any{ - "taskId": event.TaskID, - "seq": event.Seq, - "eventType": event.EventType, - "status": event.Status, - "phase": event.Phase, - "progress": event.Progress, - "message": event.Message, - "payload": event.Payload, - "simulated": event.Simulated, - "createdAt": event.CreatedAt, - }) _, err := s.pool.Exec(ctx, ` INSERT INTO gateway_task_callback_outbox (task_id, event_id, seq, callback_url, payload) -VALUES ($1::uuid, $2::uuid, $3, $4, $5::jsonb) +VALUES ($1::uuid, $2::uuid, $3, $4, '{}'::jsonb) ON CONFLICT (task_id, seq, callback_url) DO NOTHING`, event.TaskID, event.ID, event.Seq, callbackURL, - string(payloadJSON), ) return err } diff --git a/apps/api/internal/store/wallet.go b/apps/api/internal/store/wallet.go index dfc996f..a3191ea 100644 --- a/apps/api/internal/store/wallet.go +++ b/apps/api/internal/store/wallet.go @@ -161,7 +161,7 @@ func (s *Store) ReserveTaskBilling(ctx context.Context, task GatewayTask, user * } reservations := make([]WalletBillingReservation, 0, len(amounts)) - pricingSnapshotJSON, _ := json.Marshal(pricingSnapshot) + pricingSnapshotJSON, _ := json.Marshal(sanitizeJSONForStorage(pricingSnapshot)) requestFingerprint := walletString(pricingSnapshot["requestFingerprint"]) err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error { for currency, rawAmount := range amounts { @@ -306,7 +306,7 @@ func (s *Store) reserveTaskBillingExact(ctx context.Context, task GatewayTask, g if currency != "resource" { return nil, fmt.Errorf("unsupported billing currency %q", currency) } - pricingSnapshotJSON, _ := json.Marshal(pricingSnapshot) + pricingSnapshotJSON, _ := json.Marshal(sanitizeJSONForStorage(pricingSnapshot)) requestFingerprint := walletString(pricingSnapshot["requestFingerprint"]) var reservations []WalletBillingReservation err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error { diff --git a/apps/api/migrations/0076_model_identity_contract.sql b/apps/api/migrations/0076_model_identity_contract.sql new file mode 100644 index 0000000..6f37934 --- /dev/null +++ b/apps/api/migrations/0076_model_identity_contract.sql @@ -0,0 +1,173 @@ +ALTER TABLE base_model_catalog + ADD COLUMN IF NOT EXISTS invocation_name text; + +UPDATE base_model_catalog +SET invocation_name = COALESCE( + NULLIF(metadata->>'officialModelId', ''), + NULLIF(provider_model_name, ''), + canonical_model_key + ) +WHERE invocation_name IS NULL OR trim(invocation_name) = ''; + +UPDATE base_model_catalog +SET display_name = COALESCE( + NULLIF(metadata->>'officialDisplayName', ''), + NULLIF(default_snapshot->>'displayName', ''), + NULLIF(display_name, ''), + invocation_name + ) +WHERE catalog_type = 'system'; + +ALTER TABLE base_model_catalog + ALTER COLUMN invocation_name SET NOT NULL; + +CREATE INDEX IF NOT EXISTS idx_base_model_catalog_invocation_name + ON base_model_catalog(invocation_name); + +CREATE TABLE IF NOT EXISTS model_compatibility_aliases ( + id uuid PRIMARY KEY DEFAULT gen_random_uuid(), + base_model_id uuid NOT NULL REFERENCES base_model_catalog(id) ON DELETE CASCADE, + alias text NOT NULL, + model_type text NOT NULL, + active boolean NOT NULL DEFAULT true, + expires_at timestamptz, + created_at timestamptz NOT NULL DEFAULT now(), + updated_at timestamptz NOT NULL DEFAULT now(), + CONSTRAINT model_compatibility_aliases_nonempty CHECK (trim(alias) <> ''), + CONSTRAINT model_compatibility_aliases_unique UNIQUE (base_model_id, alias, model_type) +); + +CREATE INDEX IF NOT EXISTS idx_model_compatibility_aliases_lookup + ON model_compatibility_aliases(alias, model_type) + WHERE active = true; + +CREATE TABLE IF NOT EXISTS model_alias_usage_metrics ( + alias text NOT NULL, + canonical_model_key text NOT NULL, + model_type text NOT NULL, + hit_count bigint NOT NULL DEFAULT 0, + first_used_at timestamptz NOT NULL DEFAULT now(), + last_used_at timestamptz NOT NULL DEFAULT now(), + PRIMARY KEY (alias, canonical_model_key, model_type) +); + +-- Preserve unambiguous historical card/call names for the 14-day compatibility +-- window. Ambiguous names are deliberately not inferred: they must be resolved +-- to one canonical identity by an administrator before they can be registered. +WITH alias_candidates AS ( + SELECT b.id AS base_model_id, + old_name.alias, + model_type.value AS model_type, + b.invocation_name + FROM base_model_catalog b + CROSS JOIN LATERAL jsonb_array_elements_text(b.model_type) AS model_type(value) + CROSS JOIN LATERAL ( + SELECT b.display_name AS alias + UNION + SELECT b.canonical_model_key AS alias + ) old_name + WHERE trim(COALESCE(old_name.alias, '')) <> '' + AND old_name.alias <> b.invocation_name +), unambiguous AS ( + SELECT alias, model_type, min(base_model_id::text)::uuid AS base_model_id + FROM alias_candidates + GROUP BY alias, model_type + HAVING count(DISTINCT invocation_name) = 1 +) +INSERT INTO model_compatibility_aliases (base_model_id, alias, model_type, expires_at) +SELECT base_model_id, alias, model_type, now() + interval '14 days' +FROM unambiguous +ON CONFLICT (base_model_id, alias, model_type) DO NOTHING; + +-- Keep the old platform model_name/model_alias values as explicit compatibility +-- aliases before canonical invocation names are written to platform_models. +WITH platform_alias_candidates AS ( + SELECT b.id AS base_model_id, + old_name.alias, + model_type.value AS model_type, + b.invocation_name + FROM platform_models pm + JOIN base_model_catalog b ON b.id = pm.base_model_id + CROSS JOIN LATERAL jsonb_array_elements_text(pm.model_type) AS model_type(value) + CROSS JOIN LATERAL ( + SELECT pm.model_name AS alias + UNION + SELECT pm.model_alias AS alias + ) old_name + WHERE trim(COALESCE(old_name.alias, '')) <> '' + AND old_name.alias <> b.invocation_name +), unambiguous AS ( + SELECT alias, model_type, min(base_model_id::text)::uuid AS base_model_id + FROM platform_alias_candidates + GROUP BY alias, model_type + HAVING count(DISTINCT invocation_name) = 1 +) +INSERT INTO model_compatibility_aliases (base_model_id, alias, model_type, expires_at) +SELECT base_model_id, alias, model_type, now() + interval '14 days' +FROM unambiguous +ON CONFLICT (base_model_id, alias, model_type) DO NOTHING; + +WITH ranked_targets AS ( + SELECT pm.id, + b.invocation_name, + COALESCE(NULLIF(b.display_name, ''), b.invocation_name) AS display_name, + row_number() OVER ( + PARTITION BY pm.platform_id, b.invocation_name + ORDER BY CASE WHEN pm.model_name = b.invocation_name THEN 0 ELSE 1 END, + pm.created_at ASC, + pm.id ASC + ) AS identity_rank + FROM platform_models pm + JOIN base_model_catalog b ON b.id = pm.base_model_id +) +UPDATE platform_models pm +SET model_name = target.invocation_name, + model_alias = target.invocation_name, + display_name = target.display_name, + updated_at = now() +FROM ranked_targets target +WHERE target.id = pm.id + AND target.identity_rank = 1 + AND NOT EXISTS ( + SELECT 1 + FROM platform_models conflict + WHERE conflict.platform_id = pm.platform_id + AND conflict.id <> pm.id + AND conflict.model_name = target.invocation_name + ); + +UPDATE base_model_catalog +SET default_snapshot = default_snapshot || jsonb_build_object( + 'invocationName', invocation_name, + 'modelAlias', invocation_name, + 'displayName', display_name + ) +WHERE catalog_type = 'system' + AND COALESCE(default_snapshot, '{}'::jsonb) <> '{}'::jsonb; + +CREATE OR REPLACE FUNCTION fill_system_base_model_default_snapshot() +RETURNS trigger AS $$ +BEGIN + IF NEW.catalog_type = 'system' AND NEW.default_snapshot IS NULL THEN + NEW.default_snapshot = jsonb_build_object( + 'providerKey', NEW.provider_key, + 'canonicalModelKey', NEW.canonical_model_key, + 'invocationName', NEW.invocation_name, + 'providerModelName', NEW.provider_model_name, + 'modelType', NEW.model_type, + 'modelAlias', NEW.invocation_name, + 'displayName', NEW.display_name, + 'capabilities', NEW.capabilities, + 'baseBillingConfig', NEW.base_billing_config, + 'defaultRateLimitPolicy', NEW.default_rate_limit_policy, + 'pricingRuleSetId', COALESCE(NEW.pricing_rule_set_id::text, ''), + 'runtimePolicySetId', COALESCE(NEW.runtime_policy_set_id::text, ''), + 'runtimePolicyOverride', NEW.runtime_policy_override, + 'metadata', NEW.metadata, + 'pricingVersion', NEW.pricing_version, + 'status', NEW.status + ); + END IF; + RETURN NEW; +END; +$$ LANGUAGE plpgsql; diff --git a/apps/api/migrations/0077_model_lifecycle_cleanup_batch_one.sql b/apps/api/migrations/0077_model_lifecycle_cleanup_batch_one.sql new file mode 100644 index 0000000..e658ca3 --- /dev/null +++ b/apps/api/migrations/0077_model_lifecycle_cleanup_batch_one.sql @@ -0,0 +1,334 @@ +-- Gemini models that Google has retired or superseded stay callable only through +-- the compatibility window; they are no longer eligible for new catalog cards. +UPDATE base_model_catalog +SET status = 'deprecated', + metadata = metadata || jsonb_build_object( + 'officialModelId', invocation_name, + 'officialDisplayName', COALESCE(NULLIF(display_name, ''), invocation_name), + 'lifecycleStage', 'deprecated', + 'replacementModel', CASE + WHEN invocation_name = 'gemini-2.0-flash-exp-image-generation' THEN 'gemini-3.1-flash-image' + WHEN invocation_name = 'gemini-3-pro-preview' THEN 'gemini-3.1-pro-preview' + WHEN invocation_name = 'gemini-3.1-flash-lite-preview' THEN 'gemini-3.5-flash-lite' + WHEN invocation_name = 'gemini-3-pro-image-preview' THEN 'gemini-3-pro-image' + WHEN invocation_name = 'gemini-3.1-flash-image-preview' THEN 'gemini-3.1-flash-image' + WHEN invocation_name = 'gemini-3-flash-preview' THEN 'gemini-3.6-flash' + ELSE metadata->>'replacementModel' + END, + 'retireAt', '2026-07-22', + 'sourceUrl', 'https://ai.google.dev/gemini-api/docs/deprecations', + 'verifiedAt', '2026-07-22' + ), + updated_at = now() +WHERE catalog_type = 'system' + AND customized_at IS NULL + AND invocation_name IN ( + 'gemini-2.0-flash-exp-image-generation', + 'gemini-3-pro-preview', + 'gemini-3.1-flash-lite-preview', + 'gemini-3-pro-image-preview', + 'gemini-3.1-flash-image-preview', + 'gemini-3-flash-preview' + ); + +WITH targets(model_id, display_name, template_model, fallback_type, lifecycle_stage) AS ( + VALUES + ('gemini-3.1-pro-preview', 'Gemini 3.1 Pro Preview', 'gemini-3-pro-preview', 'text_generate', 'preview'), + ('gemini-3.5-flash-lite', 'Gemini 3.5 Flash-Lite', 'gemini-3.1-flash-lite-preview', 'text_generate', 'stable'), + ('gemini-3-pro-image', 'Gemini 3 Pro Image', 'gemini-3-pro-image-preview', 'image_generate', 'stable'), + ('gemini-3.1-flash-image', 'Gemini 3.1 Flash Image', 'gemini-3.1-flash-image-preview', 'image_generate', 'stable'), + ('gemini-3.6-flash', 'Gemini 3.6 Flash', 'gemini-3-flash-preview', 'text_generate', 'stable') +) +INSERT INTO base_model_catalog ( + provider_id, provider_key, canonical_model_key, invocation_name, provider_model_name, + model_type, display_name, capabilities, base_billing_config, default_rate_limit_policy, + pricing_rule_set_id, runtime_policy_set_id, runtime_policy_override, metadata, + catalog_type, pricing_version, status +) +SELECT provider.id, + 'gemini', + 'gemini:' || target.model_id, + target.model_id, + target.model_id, + COALESCE(template.model_type, jsonb_build_array(target.fallback_type)), + target.display_name, + COALESCE(template.capabilities, '{}'::jsonb), + COALESCE(template.base_billing_config, '{}'::jsonb), + COALESCE(template.default_rate_limit_policy, '{}'::jsonb), + template.pricing_rule_set_id, + template.runtime_policy_set_id, + COALESCE(template.runtime_policy_override, '{}'::jsonb), + jsonb_build_object( + 'officialModelId', target.model_id, + 'officialDisplayName', target.display_name, + 'lifecycleStage', target.lifecycle_stage, + 'sourceUrl', 'https://ai.google.dev/gemini-api/docs/deprecations', + 'verifiedAt', '2026-07-22', + 'bindingStatus', 'unverified' + ), + 'system', + COALESCE(template.pricing_version, 1), + 'active' +FROM targets target +LEFT JOIN model_catalog_providers provider ON provider.provider_key = 'gemini' OR provider.provider_code = 'gemini' +LEFT JOIN LATERAL ( + SELECT source.model_type, source.capabilities, source.base_billing_config, + source.default_rate_limit_policy, source.pricing_rule_set_id, + source.runtime_policy_set_id, source.runtime_policy_override, source.pricing_version + FROM base_model_catalog source + WHERE source.invocation_name = target.template_model + ORDER BY CASE WHEN source.provider_key = 'gemini' THEN 0 ELSE 1 END, source.updated_at DESC + LIMIT 1 +) template ON true +ON CONFLICT (canonical_model_key) DO UPDATE +SET invocation_name = CASE WHEN base_model_catalog.customized_at IS NULL THEN EXCLUDED.invocation_name ELSE base_model_catalog.invocation_name END, + provider_model_name = CASE WHEN base_model_catalog.customized_at IS NULL THEN EXCLUDED.provider_model_name ELSE base_model_catalog.provider_model_name END, + display_name = CASE WHEN base_model_catalog.customized_at IS NULL THEN EXCLUDED.display_name ELSE base_model_catalog.display_name END, + metadata = CASE WHEN base_model_catalog.customized_at IS NULL THEN base_model_catalog.metadata || EXCLUDED.metadata ELSE base_model_catalog.metadata END, + status = CASE WHEN base_model_catalog.customized_at IS NULL THEN 'active' ELSE base_model_catalog.status END, + updated_at = now(); + +-- Alibaba Cloud announced the following Qwen IDs for retirement by 2026-10-10. +UPDATE base_model_catalog +SET status = 'deprecated', + metadata = metadata || jsonb_build_object( + 'officialModelId', invocation_name, + 'officialDisplayName', COALESCE(NULLIF(display_name, ''), invocation_name), + 'lifecycleStage', 'deprecated', + 'replacementModel', CASE + WHEN invocation_name = 'qwen3-max' THEN 'qwen3.7-max' + WHEN invocation_name = 'qwen-turbo' THEN 'qwen3.7-plus' + WHEN invocation_name IN ('qwen-vl-max', 'qwen-vl-plus') THEN 'qwen3.6-flash' + WHEN invocation_name = 'qwen3-235b-a22b' THEN 'qwen3.7-plus' + ELSE metadata->>'replacementModel' + END, + 'retireAt', '2026-10-10', + 'sourceUrl', 'https://help.aliyun.com/en/model-studio/model-depreciation', + 'verifiedAt', '2026-07-22' + ), + updated_at = now() +WHERE catalog_type = 'system' + AND customized_at IS NULL + AND invocation_name IN ('qwen3-max', 'qwen-turbo', 'qwen-vl-max', 'qwen-vl-plus', 'qwen3-235b-a22b'); + +WITH targets(model_id, display_name, template_model, fallback_type) AS ( + VALUES + ('qwen3.7-max', 'Qwen3.7-Max', 'qwen3-max', 'text_generate'), + ('qwen3.7-plus', 'Qwen3.7-Plus', 'qwen3-235b-a22b', 'text_generate'), + ('qwen3.6-flash', 'Qwen3.6-Flash', 'qwen-vl-plus', 'text_generate') +) +INSERT INTO base_model_catalog ( + provider_id, provider_key, canonical_model_key, invocation_name, provider_model_name, + model_type, display_name, capabilities, base_billing_config, default_rate_limit_policy, + pricing_rule_set_id, runtime_policy_set_id, runtime_policy_override, metadata, + catalog_type, pricing_version, status +) +SELECT provider.id, + 'aliyun', + 'aliyun:' || target.model_id, + target.model_id, + target.model_id, + COALESCE(template.model_type, jsonb_build_array(target.fallback_type)), + target.display_name, + COALESCE(template.capabilities, '{}'::jsonb), + COALESCE(template.base_billing_config, '{}'::jsonb), + COALESCE(template.default_rate_limit_policy, '{}'::jsonb), + template.pricing_rule_set_id, + template.runtime_policy_set_id, + COALESCE(template.runtime_policy_override, '{}'::jsonb), + jsonb_build_object( + 'officialModelId', target.model_id, + 'officialDisplayName', target.display_name, + 'lifecycleStage', 'stable', + 'sourceUrl', 'https://help.aliyun.com/en/model-studio/model-depreciation', + 'verifiedAt', '2026-07-22', + 'bindingStatus', 'unverified' + ), + 'system', + COALESCE(template.pricing_version, 1), + 'active' +FROM targets target +LEFT JOIN model_catalog_providers provider ON provider.provider_key = 'aliyun' OR provider.provider_code = 'aliyun' +LEFT JOIN LATERAL ( + SELECT source.model_type, source.capabilities, source.base_billing_config, + source.default_rate_limit_policy, source.pricing_rule_set_id, + source.runtime_policy_set_id, source.runtime_policy_override, source.pricing_version + FROM base_model_catalog source + WHERE source.invocation_name = target.template_model + ORDER BY CASE WHEN source.provider_key = 'aliyun' THEN 0 ELSE 1 END, source.updated_at DESC + LIMIT 1 +) template ON true +ON CONFLICT (canonical_model_key) DO UPDATE +SET invocation_name = CASE WHEN base_model_catalog.customized_at IS NULL THEN EXCLUDED.invocation_name ELSE base_model_catalog.invocation_name END, + provider_model_name = CASE WHEN base_model_catalog.customized_at IS NULL THEN EXCLUDED.provider_model_name ELSE base_model_catalog.provider_model_name END, + display_name = CASE WHEN base_model_catalog.customized_at IS NULL THEN EXCLUDED.display_name ELSE base_model_catalog.display_name END, + metadata = CASE WHEN base_model_catalog.customized_at IS NULL THEN base_model_catalog.metadata || EXCLUDED.metadata ELSE base_model_catalog.metadata END, + status = CASE WHEN base_model_catalog.customized_at IS NULL THEN 'active' ELSE base_model_catalog.status END, + updated_at = now(); + +-- DeepSeek direct API has two official public IDs. Aggregator-only V4 names are +-- kept but explicitly classified as provider identities. +UPDATE base_model_catalog +SET metadata = metadata || jsonb_build_object( + 'officialModelId', invocation_name, + 'officialDisplayName', COALESCE(NULLIF(display_name, ''), invocation_name), + 'lifecycleStage', 'stable', + 'sourceUrl', 'https://api-docs.deepseek.com/quick_start/pricing-details-usd', + 'verifiedAt', '2026-07-22', + 'identityAuthority', 'official' + ), + updated_at = now() +WHERE invocation_name IN ('deepseek-chat', 'deepseek-reasoner'); + +UPDATE base_model_catalog +SET metadata = metadata || jsonb_build_object( + 'lifecycleStage', 'provider', + 'verifiedAt', '2026-07-22', + 'identityAuthority', 'provider' + ), + updated_at = now() +WHERE lower(invocation_name) LIKE 'deepseek-v4-%'; + +-- OpenRouter keeps its namespaced upstream ID, while callers use Anthropic's +-- official model ID. +WITH claude_mapping(provider_model_name, invocation_name, display_name) AS ( + VALUES + ('anthropic/claude-opus-4.5', 'claude-opus-4-5', 'Claude Opus 4.5'), + ('anthropic/claude-opus-4.6', 'claude-opus-4-6', 'Claude Opus 4.6'), + ('anthropic/claude-sonnet-4.5', 'claude-sonnet-4-5', 'Claude Sonnet 4.5'), + ('anthropic/claude-sonnet-4.6', 'claude-sonnet-4-6', 'Claude Sonnet 4.6'), + ('anthropic/claude-haiku-4.5', 'claude-haiku-4-5', 'Claude Haiku 4.5') +), affected AS ( + SELECT base.id, base.invocation_name AS old_invocation_name, mapping.invocation_name, mapping.display_name + FROM base_model_catalog base + JOIN claude_mapping mapping ON mapping.provider_model_name = base.provider_model_name + WHERE base.catalog_type = 'system' AND base.customized_at IS NULL +), old_aliases AS ( + INSERT INTO model_compatibility_aliases (base_model_id, alias, model_type, expires_at) + SELECT affected.id, affected.old_invocation_name, model_type.value, now() + interval '14 days' + FROM affected + JOIN base_model_catalog base ON base.id = affected.id + CROSS JOIN LATERAL jsonb_array_elements_text(base.model_type) AS model_type(value) + WHERE affected.old_invocation_name <> affected.invocation_name + ON CONFLICT (base_model_id, alias, model_type) DO NOTHING +) +UPDATE base_model_catalog base +SET invocation_name = affected.invocation_name, + display_name = affected.display_name, + metadata = base.metadata || jsonb_build_object( + 'officialModelId', affected.invocation_name, + 'officialDisplayName', affected.display_name, + 'lifecycleStage', 'stable', + 'sourceUrl', 'https://platform.claude.com/docs/en/about-claude/models/model-ids-and-versions', + 'verifiedAt', '2026-07-22', + 'identityAuthority', 'official' + ), + updated_at = now() +FROM affected +WHERE base.id = affected.id; + +WITH ranked_targets AS ( + SELECT platform_model.id, + base.invocation_name, + COALESCE(NULLIF(base.display_name, ''), base.invocation_name) AS display_name, + row_number() OVER ( + PARTITION BY platform_model.platform_id, base.invocation_name + ORDER BY CASE WHEN platform_model.model_name = base.invocation_name THEN 0 ELSE 1 END, + platform_model.created_at ASC, + platform_model.id ASC + ) AS identity_rank + FROM platform_models platform_model + JOIN base_model_catalog base ON base.id = platform_model.base_model_id +) +UPDATE platform_models platform_model +SET model_name = target.invocation_name, + model_alias = target.invocation_name, + display_name = target.display_name, + updated_at = now() +FROM ranked_targets target +WHERE target.id = platform_model.id + AND target.identity_rank = 1 + AND platform_model.model_name <> target.invocation_name + AND NOT EXISTS ( + SELECT 1 + FROM platform_models conflict + WHERE conflict.platform_id = platform_model.platform_id + AND conflict.id <> platform_model.id + AND conflict.model_name = target.invocation_name + ); + +-- MiniMax keeps current main and fallback families; older text/speech previews +-- remain present only for compatibility and reference-safe migration. +UPDATE base_model_catalog +SET status = 'deprecated', + metadata = metadata || jsonb_build_object( + 'lifecycleStage', 'deprecated', + 'replacementModel', CASE + WHEN lower(invocation_name) LIKE '%speech-2.5%' THEN 'speech-2.8' + ELSE 'minimax-m2.5' + END, + 'sourceUrl', 'https://platform.minimaxi.com/docs/api-reference/api-overview', + 'verifiedAt', '2026-07-22' + ), + updated_at = now() +WHERE catalog_type = 'system' + AND customized_at IS NULL + AND ( + lower(invocation_name) IN ('minimax-m2', 'minimax-m2.1', 'minimax-m2-highspeed', 'minimax-m2.1-highspeed') + OR lower(invocation_name) LIKE '%speech-2.5%preview%' + ); + +-- Consolidate duplicate qwen-max rows conservatively. A row is hidden only +-- after every safe platform binding has moved to the selected keeper. +WITH ranked AS ( + SELECT base.id, + first_value(base.id) OVER ( + PARTITION BY base.provider_key, base.invocation_name + ORDER BY (SELECT count(*) FROM platform_models model WHERE model.base_model_id = base.id) DESC, + CASE WHEN base.customized_at IS NULL THEN 0 ELSE 1 END, + base.created_at ASC, + base.id ASC + ) AS keeper_id + FROM base_model_catalog base + WHERE lower(base.invocation_name) = 'qwen-max' +), rebound AS ( + UPDATE platform_models model + SET base_model_id = ranked.keeper_id, + updated_at = now() + FROM ranked + WHERE model.base_model_id = ranked.id + AND ranked.id <> ranked.keeper_id + AND NOT EXISTS ( + SELECT 1 FROM platform_models keeper_model + WHERE keeper_model.platform_id = model.platform_id + AND keeper_model.base_model_id = ranked.keeper_id + ) + RETURNING model.id +) +UPDATE base_model_catalog base +SET status = 'hidden', + metadata = metadata || jsonb_build_object( + 'lifecycleStage', 'duplicate', + 'cleanupCandidate', true, + 'verifiedAt', '2026-07-22' + ), + updated_at = now() +FROM ranked +WHERE base.id = ranked.id + AND ranked.id <> ranked.keeper_id + AND base.catalog_type = 'system' + AND base.customized_at IS NULL + AND NOT EXISTS (SELECT 1 FROM platform_models model WHERE model.base_model_id = base.id); + +-- Release B may physically delete only candidates which remain unreferenced +-- after the observation window. Custom and manually modified rows are excluded. +UPDATE base_model_catalog base +SET metadata = metadata || jsonb_build_object( + 'cleanupCandidate', true, + 'cleanupEligibleAfter', '2026-08-05' + ), + updated_at = now() +WHERE base.status IN ('deprecated', 'hidden') + AND base.catalog_type = 'system' + AND base.customized_at IS NULL + AND NOT EXISTS (SELECT 1 FROM platform_models model WHERE model.base_model_id = base.id); diff --git a/apps/api/migrations/0078_seedance_input_image_constraints.sql b/apps/api/migrations/0078_seedance_input_image_constraints.sql new file mode 100644 index 0000000..6e7fcf7 --- /dev/null +++ b/apps/api/migrations/0078_seedance_input_image_constraints.sql @@ -0,0 +1,98 @@ +-- Tencent Seedance SE accepts input images only within the observed provider +-- limits. The runtime uses these capability fields to auto-convert images +-- before forwarding the request. +WITH input_constraints AS ( + SELECT jsonb_build_object( + 'input_image_resolution_range', + '{"min":{"long_edge":360,"short_edge":360},"max":{"long_edge":1920,"short_edge":1080}}'::jsonb, + 'input_image_aspect_ratio_range', + '[0.39,2.5]'::jsonb + ) AS value +), +target_base_models AS ( + SELECT + base_model.id, + COALESCE(base_model.capabilities, '{}'::jsonb) || jsonb_build_object( + 'image_to_video', + COALESCE(base_model.capabilities->'image_to_video', '{}'::jsonb) || input_constraints.value, + 'omni_video', + COALESCE(base_model.capabilities->'omni_video', '{}'::jsonb) || input_constraints.value + ) AS image_capabilities + FROM base_model_catalog base_model + CROSS JOIN input_constraints + WHERE base_model.canonical_model_key IN ( + 'easyai:豆包Seedance-2.0', + 'easyai:豆包Seedance-2.0-fast' + ) + OR ( + base_model.provider_key = 'easyai' + AND base_model.provider_model_name IN ('豆包Seedance-2.0', '豆包Seedance-2.0-fast') + ) + OR ( + base_model.provider_key = 'tencent-hunyuan-video' + AND lower(base_model.provider_model_name) IN ('2.0', '2.0-fast') + ) +), +updated_base_models AS ( + UPDATE base_model_catalog base_model + SET capabilities = target.image_capabilities, + metadata = jsonb_set( + COALESCE(base_model.metadata, '{}'::jsonb) || jsonb_build_object( + 'rawModel', + COALESCE(base_model.metadata->'rawModel', '{}'::jsonb) + ), + '{rawModel,capabilities}', + target.image_capabilities - 'originalTypes', + true + ), + default_snapshot = CASE + WHEN COALESCE(base_model.default_snapshot, '{}'::jsonb) = '{}'::jsonb + THEN base_model.default_snapshot + ELSE jsonb_set( + jsonb_set( + base_model.default_snapshot || jsonb_build_object( + 'metadata', + COALESCE(base_model.default_snapshot->'metadata', '{}'::jsonb) || jsonb_build_object( + 'rawModel', + COALESCE(base_model.default_snapshot#>'{metadata,rawModel}', '{}'::jsonb) + ) + ), + '{capabilities}', + target.image_capabilities, + true + ), + '{metadata,rawModel,capabilities}', + target.image_capabilities - 'originalTypes', + true + ) + END, + updated_at = now() + FROM target_base_models target + WHERE base_model.id = target.id + RETURNING base_model.id +) +UPDATE platform_models platform_model +SET capabilities = COALESCE(platform_model.capabilities, '{}'::jsonb) || jsonb_build_object( + 'image_to_video', + COALESCE(platform_model.capabilities->'image_to_video', '{}'::jsonb) || input_constraints.value, + 'omni_video', + COALESCE(platform_model.capabilities->'omni_video', '{}'::jsonb) || input_constraints.value + ), + updated_at = now() +FROM integration_platforms platform +CROSS JOIN input_constraints +WHERE platform_model.platform_id = platform.id + AND platform.deleted_at IS NULL + AND ( + platform_model.base_model_id IN (SELECT id FROM updated_base_models) + OR ( + platform.provider = 'easyai' + AND COALESCE(NULLIF(platform_model.provider_model_name, ''), platform_model.model_name) + IN ('豆包Seedance-2.0', '豆包Seedance-2.0-fast') + ) + OR ( + platform.provider = 'tencent-hunyuan-video' + AND lower(COALESCE(NULLIF(platform_model.provider_model_name, ''), platform_model.model_name)) + IN ('2.0', '2.0-fast') + ) + ); diff --git a/apps/api/migrations/0079_image_model_migration_contract.sql b/apps/api/migrations/0079_image_model_migration_contract.sql new file mode 100644 index 0000000..a6d4e98 --- /dev/null +++ b/apps/api/migrations/0079_image_model_migration_contract.sql @@ -0,0 +1,181 @@ +-- Finalize the public/runtime contract for the three image-model migration. +-- Callers use stable official Gemini IDs while provider bindings retain their +-- upstream preview IDs. Capabilities and pricing are copied from the source +-- EasyAI runtime contract instead of falling back to generic defaults. +WITH model_contract( + invocation_name, + preview_invocation_name, + capabilities, + billing_config, + pricing_rule_set_key +) AS ( + VALUES + ( + 'gemini-3-pro-image', + 'gemini-3-pro-image-preview', + '{ + "image_generate": { + "output_multiple_images": true, + "output_max_images_count": 4, + "output_resolutions": ["1K", "2K", "4K"], + "aspect_ratio_allowed": ["1:1", "16:9", "4:3", "3:4", "5:4", "4:5", "9:16", "21:9", "3:2", "2:3", "adaptive"] + }, + "image_edit": { + "output_multiple_images": true, + "input_multiple_images": true, + "input_max_images_count": 14, + "input_max_file_size_bytes": 20971520, + "output_resolutions": ["1K", "2K", "4K"], + "aspect_ratio_allowed": ["1:1", "16:9", "4:3", "3:4", "5:4", "4:5", "9:16", "21:9", "3:2", "2:3", "adaptive"] + }, + "originalTypes": ["image_generate", "image_edit"] + }'::jsonb, + '{ + "image": { + "basePrice": 1, + "baseWeight": 1, + "dynamicWeight": { + "1K": 1, + "2K": 1, + "3K": 1, + "4K": 1.8, + "8K": 4, + "qualityLow": 0.5, + "qualityMedium": 1, + "qualityHigh": 1.5 + } + } + }'::jsonb, + 'migration-pricing-f721f8076252bd77' + ), + ( + 'gemini-3.1-flash-image', + 'gemini-3.1-flash-image-preview', + '{ + "image_generate": { + "output_multiple_images": true, + "output_max_images_count": 4, + "output_resolutions": ["1K", "2K", "4K"], + "aspect_ratio_allowed": ["1:1", "16:9", "4:3", "3:4", "5:4", "4:5", "9:16", "21:9", "3:2", "2:3", "1:4", "4:1", "1:8", "8:1", "adaptive"] + }, + "image_edit": { + "output_multiple_images": true, + "input_multiple_images": true, + "input_max_images_count": 14, + "input_max_file_size_bytes": 20971520, + "output_resolutions": ["1K", "2K", "4K"], + "aspect_ratio_allowed": ["1:1", "16:9", "4:3", "3:4", "5:4", "4:5", "9:16", "21:9", "3:2", "2:3", "1:4", "4:1", "1:8", "8:1", "adaptive"] + }, + "originalTypes": ["image_generate", "image_edit"] + }'::jsonb, + '{ + "image": { + "basePrice": 1, + "baseWeight": 1, + "dynamicWeight": { + "1K": 1, + "2K": 1, + "3K": 1, + "4K": 1.8, + "8K": 4, + "qualityLow": 0.5, + "qualityMedium": 1, + "qualityHigh": 1.5 + } + } + }'::jsonb, + 'migration-pricing-f721f8076252bd77' + ) +), +updated_base_models AS ( + UPDATE base_model_catalog base_model + SET capabilities = contract.capabilities, + base_billing_config = contract.billing_config, + pricing_rule_set_id = COALESCE( + ( + SELECT rule_set.id + FROM model_pricing_rule_sets rule_set + WHERE rule_set.rule_set_key = contract.pricing_rule_set_key + AND rule_set.status = 'active' + LIMIT 1 + ), + base_model.pricing_rule_set_id + ), + metadata = COALESCE(base_model.metadata, '{}'::jsonb) || jsonb_build_object( + 'bindingStatus', 'validating', + 'capabilitySource', 'server-main-model-runtime', + 'pricingSource', contract.pricing_rule_set_key + ), + status = 'active', + updated_at = now() + FROM model_contract contract + WHERE base_model.provider_key = 'gemini' + AND base_model.invocation_name = contract.invocation_name + RETURNING base_model.id, base_model.invocation_name, base_model.capabilities +) +UPDATE platform_models platform_model +SET capabilities = updated.capabilities, + updated_at = now() +FROM updated_base_models updated +WHERE platform_model.base_model_id = updated.id; + +-- The preview catalog rows originated as customized import records, so the +-- generic lifecycle migration deliberately did not alter them. This migration +-- scopes the lifecycle change to the two migrated image identities. +WITH replacements(preview_invocation_name, invocation_name) AS ( + VALUES + ('gemini-3-pro-image-preview', 'gemini-3-pro-image'), + ('gemini-3.1-flash-image-preview', 'gemini-3.1-flash-image') +), +preview_models AS ( + UPDATE base_model_catalog preview + SET status = 'deprecated', + metadata = COALESCE(preview.metadata, '{}'::jsonb) || jsonb_build_object( + 'lifecycleStage', 'deprecated', + 'replacementModel', replacements.invocation_name, + 'retireAt', '2026-07-22', + 'sourceUrl', 'https://ai.google.dev/gemini-api/docs/deprecations', + 'verifiedAt', '2026-07-24' + ), + updated_at = now() + FROM replacements + WHERE preview.provider_key = 'gemini' + AND preview.invocation_name = replacements.preview_invocation_name + RETURNING preview.id, preview.invocation_name +), +stable_models AS ( + SELECT stable.id, stable.invocation_name, replacements.preview_invocation_name + FROM replacements + JOIN base_model_catalog stable + ON stable.provider_key = 'gemini' + AND stable.invocation_name = replacements.invocation_name +), +disabled_conflicting_aliases AS ( + UPDATE model_compatibility_aliases compatibility_alias + SET active = false, + expires_at = COALESCE(compatibility_alias.expires_at, now()), + updated_at = now() + FROM preview_models preview, stable_models stable + WHERE compatibility_alias.base_model_id = preview.id + AND stable.preview_invocation_name = preview.invocation_name + AND compatibility_alias.alias IN ( + stable.invocation_name, + 'gemini:' || preview.invocation_name + ) + RETURNING compatibility_alias.id +) +INSERT INTO model_compatibility_aliases ( + base_model_id, + alias, + model_type, + active +) +SELECT stable.id, stable.preview_invocation_name, model_type.value, true +FROM stable_models stable +CROSS JOIN LATERAL ( + VALUES ('image_generate'), ('image_edit') +) AS model_type(value) +ON CONFLICT (base_model_id, alias, model_type) DO UPDATE +SET active = true, + expires_at = NULL, + updated_at = now(); diff --git a/apps/api/migrations/0080_platform_model_rate_limit_policy_mode.sql b/apps/api/migrations/0080_platform_model_rate_limit_policy_mode.sql new file mode 100644 index 0000000..3b0c120 --- /dev/null +++ b/apps/api/migrations/0080_platform_model_rate_limit_policy_mode.sql @@ -0,0 +1,58 @@ +ALTER TABLE platform_models + ADD COLUMN IF NOT EXISTS rate_limit_policy_mode text; + +-- Historical server-main imports wrote {"rules": []} as a placeholder. It +-- means "no override", not an explicit unlimited policy. +UPDATE integration_platforms +SET rate_limit_policy = '{}'::jsonb, + updated_at = now() +WHERE rate_limit_policy = '{"rules":[]}'::jsonb; + +UPDATE platform_models +SET rate_limit_policy = '{}'::jsonb, + updated_at = now() +WHERE rate_limit_policy = '{"rules":[]}'::jsonb; + +UPDATE platform_models model +SET runtime_policy_set_id = NULL, + runtime_policy_override = '{}'::jsonb +FROM base_model_catalog base +WHERE model.base_model_id = base.id + AND model.runtime_policy_set_id IS NOT DISTINCT FROM base.runtime_policy_set_id + AND model.runtime_policy_override IS NOT DISTINCT FROM base.runtime_policy_override; + +UPDATE platform_models model +SET rate_limit_policy_mode = CASE + WHEN model.runtime_policy_override ? 'rateLimitPolicy' THEN 'override' + WHEN model.rate_limit_policy <> '{}'::jsonb + AND ( + base.id IS NULL + OR model.rate_limit_policy IS DISTINCT FROM base.default_rate_limit_policy + ) + THEN 'override' + ELSE 'inherit' +END +FROM ( + SELECT platform_model.id AS platform_model_id, catalog.id, catalog.default_rate_limit_policy + FROM platform_models platform_model + LEFT JOIN base_model_catalog catalog ON catalog.id = platform_model.base_model_id +) base +WHERE model.id = base.platform_model_id + AND model.rate_limit_policy_mode IS NULL; + +UPDATE platform_models +SET rate_limit_policy_mode = 'inherit' +WHERE rate_limit_policy_mode IS NULL; + +ALTER TABLE platform_models + ALTER COLUMN rate_limit_policy_mode SET DEFAULT 'inherit'; + +ALTER TABLE platform_models + ADD CONSTRAINT platform_models_rate_limit_policy_mode_check + CHECK ( + COALESCE(rate_limit_policy_mode, '') IN ('inherit', 'override') + ) + NOT VALID; + +ALTER TABLE platform_models + VALIDATE CONSTRAINT platform_models_rate_limit_policy_mode_check; diff --git a/apps/api/migrations/0081_image_gateway_rate_limits.sql b/apps/api/migrations/0081_image_gateway_rate_limits.sql new file mode 100644 index 0000000..1a46a87 --- /dev/null +++ b/apps/api/migrations/0081_image_gateway_rate_limits.sql @@ -0,0 +1,44 @@ +-- The three migrated image identities previously relied on a global +-- user-group limit. Give every upstream binding its own canonical concurrency +-- lease so raising the service group no longer removes provider protection. +WITH limits(invocation_name, policy) AS ( + VALUES + ( + 'gemini-3-pro-image', + '{"rules":[{"metric":"concurrent","limit":10,"leaseTtlSeconds":600}]}'::jsonb + ), + ( + 'gemini-3.1-flash-image', + '{"rules":[{"metric":"concurrent","limit":10,"leaseTtlSeconds":600}]}'::jsonb + ), + ( + 'gpt-image-2', + '{"rules":[{"metric":"concurrent","limit":5,"leaseTtlSeconds":300}]}'::jsonb + ) +), +updated_base_models AS ( + UPDATE base_model_catalog base_model + SET default_rate_limit_policy = limits.policy, + metadata = COALESCE(base_model.metadata, '{}'::jsonb) || jsonb_build_object( + 'rateLimitSource', 'image-gateway-migration', + 'rateLimitUpdatedAt', '2026-07-24' + ), + updated_at = now() + FROM limits + WHERE base_model.invocation_name = limits.invocation_name + RETURNING base_model.id +) +UPDATE platform_models platform_model +SET rate_limit_policy = '{}'::jsonb, + rate_limit_policy_mode = 'inherit', + updated_at = now() +FROM updated_base_models base_model +WHERE platform_model.base_model_id = base_model.id + AND ( + platform_model.rate_limit_policy = '{}'::jsonb + OR platform_model.rate_limit_policy = '{"rules":[]}'::jsonb + OR platform_model.rate_limit_policy ? 'platformLimits' + OR platform_model.rate_limit_policy ? 'modelLimits' + OR platform_model.rate_limit_policy ? 'platform_limits' + OR platform_model.rate_limit_policy ? 'model_limits' + ); diff --git a/apps/api/migrations/0082_remove_unconfigured_gpt_image_concurrency_limit.sql b/apps/api/migrations/0082_remove_unconfigured_gpt_image_concurrency_limit.sql new file mode 100644 index 0000000..b0c0ea3 --- /dev/null +++ b/apps/api/migrations/0082_remove_unconfigured_gpt_image_concurrency_limit.sql @@ -0,0 +1,18 @@ +-- gpt-image-2 platforms do not declare an upstream concurrency limit. The +-- migration-only base default added in 0081 therefore turned an unspecified +-- limit into an artificial per-platform-model cap of five. +-- +-- Remove only the exact policy written by 0081. Explicit platform, runtime, or +-- platform-model policies continue to take precedence and remain untouched. +UPDATE base_model_catalog +SET default_rate_limit_policy = '{}'::jsonb, + metadata = (COALESCE(metadata, '{}'::jsonb) - 'rateLimitSource' - 'rateLimitUpdatedAt') + || jsonb_build_object( + 'rateLimitRemovedAt', '2026-07-24', + 'rateLimitRemovalReason', 'upstream-limit-unconfigured' + ), + updated_at = now() +WHERE invocation_name = 'gpt-image-2' + AND metadata->>'rateLimitSource' = 'image-gateway-migration' + AND default_rate_limit_policy = + '{"rules":[{"metric":"concurrent","limit":5,"leaseTtlSeconds":300}]}'::jsonb; diff --git a/apps/api/migrations/0083_task_history_minimal_storage.sql b/apps/api/migrations/0083_task_history_minimal_storage.sql new file mode 100644 index 0000000..1a0a615 --- /dev/null +++ b/apps/api/migrations/0083_task_history_minimal_storage.sql @@ -0,0 +1,11 @@ +ALTER TABLE gateway_task_events + ADD COLUMN IF NOT EXISTS platform_id uuid; + +ALTER TABLE gateway_task_attempts + ADD COLUMN IF NOT EXISTS status_code integer; + +ALTER TABLE gateway_task_callback_outbox + ADD COLUMN IF NOT EXISTS locked_by text, + ADD COLUMN IF NOT EXISTS lock_token uuid, + ADD COLUMN IF NOT EXISTS locked_at timestamptz, + ADD COLUMN IF NOT EXISTS failed_at timestamptz; diff --git a/apps/api/migrations/0084_task_history_retention_index.sql b/apps/api/migrations/0084_task_history_retention_index.sql new file mode 100644 index 0000000..d5a0f4a --- /dev/null +++ b/apps/api/migrations/0084_task_history_retention_index.sql @@ -0,0 +1,68 @@ +-- easyai:migration:no-transaction +CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_gateway_tasks_retention + ON gateway_tasks(finished_at, id) + WHERE status IN ('succeeded', 'failed', 'cancelled'); + +-- easyai:migration:statement +CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_gateway_task_attempts_retention + ON gateway_task_attempts((COALESCE(finished_at, started_at)), id) + WHERE status <> 'running'; + +-- easyai:migration:statement +CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_gateway_task_events_retention + ON gateway_task_events(created_at, id); + +-- easyai:migration:statement +CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_gateway_task_param_logs_retention + ON gateway_task_param_preprocessing_logs(created_at, id); + +-- easyai:migration:statement +CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_gateway_tasks_minimal_compaction + ON gateway_tasks(updated_at, id) + WHERE COALESCE(normalized_request, '{}'::jsonb) <> '{}'::jsonb + OR compatibility_public_id IS NOT NULL + OR compatibility_submit_http_status IS NOT NULL + OR COALESCE(compatibility_submit_headers, '{}'::jsonb) <> '{}'::jsonb + OR COALESCE(compatibility_submit_body, '{}'::jsonb) <> '{}'::jsonb + OR result ?| ARRAY[ + 'raw', 'raw_data', 'rawData', 'provider_response', 'providerResponse', + 'submit', 'file_retrieve', 'fileRetrieve', 'upstream_task_id', 'remote_task_id' + ] + OR error IS NOT NULL + OR octet_length(COALESCE(error_message, '')) > 2048; + +-- easyai:migration:statement +CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_gateway_tasks_checkpoint_compaction + ON gateway_tasks(updated_at, id) + WHERE COALESCE(remote_task_payload, '{}'::jsonb) <> '{}'::jsonb; + +-- easyai:migration:statement +CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_gateway_task_attempts_minimal_compaction + ON gateway_task_attempts(started_at, id) + WHERE request_snapshot <> '{}'::jsonb + OR COALESCE(response_snapshot, '{}'::jsonb) <> '{}'::jsonb + OR COALESCE(usage, '{}'::jsonb) <> '{}'::jsonb + OR COALESCE(metrics, '{}'::jsonb) <> '{}'::jsonb + OR COALESCE(pricing_snapshot, '{}'::jsonb) <> '{}'::jsonb + OR request_fingerprint IS NOT NULL; + +-- easyai:migration:statement +CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_gateway_task_events_minimal_compaction + ON gateway_task_events(created_at, id) + WHERE payload <> '{}'::jsonb + OR phase IS NOT NULL + OR COALESCE(progress, 0) <> 0 + OR message IS NOT NULL; + +-- easyai:migration:statement +CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_gateway_task_param_logs_minimal_compaction + ON gateway_task_param_preprocessing_logs(created_at, id) + WHERE actual_input <> '{}'::jsonb + OR converted_output <> '{}'::jsonb + OR model_snapshot <> '{}'::jsonb + OR octet_length(changes::text) > 1024; + +-- easyai:migration:statement +CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_gateway_cloned_voices_minimal_compaction + ON gateway_cloned_voices(updated_at, id) + WHERE metadata ?| ARRAY['raw', 'raw_data', 'rawData', 'provider_response', 'providerResponse', 'request']; diff --git a/apps/api/migrations/0085_task_event_platform_index.sql b/apps/api/migrations/0085_task_event_platform_index.sql new file mode 100644 index 0000000..e3ea0bc --- /dev/null +++ b/apps/api/migrations/0085_task_event_platform_index.sql @@ -0,0 +1,4 @@ +-- easyai:migration:no-transaction +CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_gateway_task_events_platform_created + ON gateway_task_events(platform_id, created_at DESC) + WHERE platform_id IS NOT NULL; diff --git a/apps/api/migrations/0086_task_callback_lock_index.sql b/apps/api/migrations/0086_task_callback_lock_index.sql new file mode 100644 index 0000000..ffdd389 --- /dev/null +++ b/apps/api/migrations/0086_task_callback_lock_index.sql @@ -0,0 +1,24 @@ +-- easyai:migration:no-transaction +CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_gateway_task_callback_outbox_stale + ON gateway_task_callback_outbox(status, locked_at) + WHERE status = 'processing'; + +-- easyai:migration:statement +CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_gateway_task_callback_outbox_delivered_retention + ON gateway_task_callback_outbox(delivered_at, id) + WHERE status = 'delivered'; + +-- easyai:migration:statement +CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_gateway_task_callback_outbox_failed_retention + ON gateway_task_callback_outbox(failed_at, id) + WHERE status = 'failed'; + +-- easyai:migration:statement +CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_gateway_task_callback_outbox_payload_compaction + ON gateway_task_callback_outbox(created_at, id) + WHERE payload <> '{}'::jsonb; + +-- easyai:migration:statement +CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_gateway_task_callback_outbox_legacy + ON gateway_task_callback_outbox(created_at, id) + WHERE status IN ('pending', 'processing'); diff --git a/apps/api/migrations/0087_admin_task_query_indexes.sql b/apps/api/migrations/0087_admin_task_query_indexes.sql new file mode 100644 index 0000000..55b288a --- /dev/null +++ b/apps/api/migrations/0087_admin_task_query_indexes.sql @@ -0,0 +1,22 @@ +-- easyai:migration:no-transaction +CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_gateway_tasks_admin_created + ON gateway_tasks(created_at DESC, id DESC); + +-- easyai:migration:statement +CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_gateway_tasks_admin_status_created + ON gateway_tasks(status, created_at DESC, id DESC); + +-- easyai:migration:statement +CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_gateway_tasks_admin_user_created + ON gateway_tasks(gateway_user_id, created_at DESC, id DESC) + WHERE gateway_user_id IS NOT NULL; + +-- easyai:migration:statement +CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_gateway_tasks_admin_tenant_created + ON gateway_tasks(gateway_tenant_id, created_at DESC, id DESC) + WHERE gateway_tenant_id IS NOT NULL; + +-- easyai:migration:statement +CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_gateway_task_attempts_admin_platform_task + ON gateway_task_attempts(platform_id, task_id) + WHERE platform_id IS NOT NULL; diff --git a/apps/api/migrations/0088_kimi_k3_qwen38_base_models.sql b/apps/api/migrations/0088_kimi_k3_qwen38_base_models.sql new file mode 100644 index 0000000..b5c33e6 --- /dev/null +++ b/apps/api/migrations/0088_kimi_k3_qwen38_base_models.sql @@ -0,0 +1,424 @@ +-- Register the latest Kimi and Qwen foundation models in both the EasyAI +-- catalog and their real upstream providers. +-- +-- Official references: +-- https://www.volcengine.com/docs/82379/2578669 +-- https://help.aliyun.com/zh/model-studio/token-plan-personal-overview +-- https://help.aliyun.com/zh/model-studio/qwen-function-calling +-- https://help.aliyun.com/zh/model-studio/qwen-api-via-openai-responses + +WITH model_defs AS ( + SELECT * + FROM ( + VALUES + ( + 'easyai', + 'easyai', + 'EasyAI', + 'easyai', + 'easyai:Kimi-K3', + 'Kimi-K3', + 'Kimi-K3', + 'Kimi-K3', + 'Kimi-K3', + 'Kimi 最新旗舰基础模型，原生支持视觉理解，面向 Coding 与 Agent 场景，统一模型 ID 持续升级，支持 100 万上下文。', + 'https://static.51easyai.com/kimi-logo-124.png', + '["text_generate","image_analysis","tools_call"]'::jsonb, + '{ + "text_generate": { + "supportedApiProtocols": ["openai_chat_completions", "openai_responses"], + "max_context_tokens": 1000000, + "supportTool": true, + "supportThinking": true + }, + "image_analysis": { + "max_context_tokens": 1000000, + "supportThinking": true + }, + "tools_call": { + "supportedApiProtocols": ["openai_chat_completions", "openai_responses"], + "max_context_tokens": 1000000, + "supportTool": true, + "supportThinking": true + }, + "originalTypes": ["text_generate", "image_analysis", "tools_call"] + }'::jsonb, + 'easyai:Kimi-K2.7-Code', + 'https://www.volcengine.com/docs/82379/2578669', + false + ), + ( + 'volces-openai', + 'volces-openai', + '火山引擎(OpenAI兼容)', + 'openai', + 'volces-openai:kimi-k3', + 'kimi-k3', + 'kimi-k3', + 'Kimi-K3', + 'Kimi-K3', + 'Kimi 最新旗舰基础模型，原生支持视觉理解，面向 Coding 与 Agent 场景，统一模型 ID 持续升级，支持 100 万上下文。', + 'https://static.51easyai.com/kimi-logo-124.png', + '["text_generate","image_analysis","tools_call"]'::jsonb, + '{ + "text_generate": { + "supportedApiProtocols": ["openai_chat_completions", "openai_responses"], + "max_context_tokens": 1000000, + "supportTool": true, + "supportThinking": true + }, + "image_analysis": { + "max_context_tokens": 1000000, + "supportThinking": true + }, + "tools_call": { + "supportedApiProtocols": ["openai_chat_completions", "openai_responses"], + "max_context_tokens": 1000000, + "supportTool": true, + "supportThinking": true + }, + "originalTypes": ["text_generate", "image_analysis", "tools_call"] + }'::jsonb, + 'volces-openai:doubao-seed-2-0-pro-260215', + 'https://www.volcengine.com/docs/82379/2578669', + false + ), + ( + 'easyai', + 'easyai', + 'EasyAI', + 'easyai', + 'easyai:Qwen3.8-Max-Preview', + 'Qwen3.8-Max-Preview', + 'Qwen3.8-Max-Preview', + 'Qwen3.8-Max-Preview', + 'Qwen3.8-Max-Preview', + '通义千问最新一代 2.4T 旗舰基础模型预览版，支持推理、视觉理解、文本生成与工具调用；仅面向阿里云百炼 Token Plan，需使用订阅专属 Base URL 与 API Key。', + 'https://static.51easyai.com/qwen-color.webp', + '["text_generate","image_analysis","tools_call"]'::jsonb, + '{ + "text_generate": { + "supportedApiProtocols": ["openai_chat_completions", "openai_responses"], + "supportTool": true, + "supportThinking": true, + "supportThinkingModeSwitch": true, + "thinkingEffortLevels": ["none", "minimal", "low", "medium", "high", "xhigh", "max"], + "thinkingBudgetTokensRange": [4096, 262144], + "supportWebSearch": true + }, + "image_analysis": { + "supportThinking": true, + "supportThinkingModeSwitch": true, + "thinkingEffortLevels": ["none", "minimal", "low", "medium", "high", "xhigh", "max"], + "thinkingBudgetTokensRange": [4096, 262144] + }, + "tools_call": { + "supportedApiProtocols": ["openai_chat_completions", "openai_responses"], + "supportTool": true, + "supportThinking": true, + "supportThinkingModeSwitch": true, + "thinkingEffortLevels": ["none", "minimal", "low", "medium", "high", "xhigh", "max"], + "thinkingBudgetTokensRange": [4096, 262144], + "supportWebSearch": true + }, + "originalTypes": ["text_generate", "image_analysis", "tools_call"] + }'::jsonb, + 'easyai:Qwen3.7-Max', + 'https://help.aliyun.com/zh/model-studio/token-plan-personal-overview', + true + ), + ( + 'aliyun-bailian-openai', + 'aliyun-bailian-openai', + '阿里云百炼(OpenAI兼容）', + 'openai', + 'aliyun-bailian-openai:qwen3.8-max-preview', + 'qwen3.8-max-preview', + 'qwen3.8-max-preview', + 'Qwen3.8-Max-Preview', + 'Qwen3.8-Max-Preview', + '通义千问最新一代 2.4T 旗舰基础模型预览版，支持推理、视觉理解、文本生成与工具调用；仅面向阿里云百炼 Token Plan，需使用订阅专属 Base URL 与 API Key。', + 'https://static.51easyai.com/qwen-color.webp', + '["text_generate","image_analysis","tools_call"]'::jsonb, + '{ + "text_generate": { + "supportedApiProtocols": ["openai_chat_completions", "openai_responses"], + "supportTool": true, + "supportThinking": true, + "supportThinkingModeSwitch": true, + "thinkingEffortLevels": ["none", "minimal", "low", "medium", "high", "xhigh", "max"], + "thinkingBudgetTokensRange": [4096, 262144], + "supportWebSearch": true + }, + "image_analysis": { + "supportThinking": true, + "supportThinkingModeSwitch": true, + "thinkingEffortLevels": ["none", "minimal", "low", "medium", "high", "xhigh", "max"], + "thinkingBudgetTokensRange": [4096, 262144] + }, + "tools_call": { + "supportedApiProtocols": ["openai_chat_completions", "openai_responses"], + "supportTool": true, + "supportThinking": true, + "supportThinkingModeSwitch": true, + "thinkingEffortLevels": ["none", "minimal", "low", "medium", "high", "xhigh", "max"], + "thinkingBudgetTokensRange": [4096, 262144], + "supportWebSearch": true + }, + "originalTypes": ["text_generate", "image_analysis", "tools_call"] + }'::jsonb, + 'aliyun-bailian-openai:qwen3.7-max', + 'https://help.aliyun.com/zh/model-studio/token-plan-personal-overview', + true + ) + ) AS defs( + provider_key, + source_provider_code, + source_provider_name, + source_spec_type, + canonical_model_key, + provider_model_name, + raw_model_name, + display_name, + model_alias, + description, + icon_path, + model_type, + capabilities, + template_key, + reference_document, + token_plan_only + ) +), +source_rows AS ( + SELECT + providers.id AS provider_id, + model_defs.*, + COALESCE( + template.base_billing_config, + '{"text":{"basePrice":0.01,"baseWeight":1},"image":{"basePrice":10,"baseWeight":1},"currency":"resource"}'::jsonb + ) AS base_billing_config, + COALESCE( + template.default_rate_limit_policy, + jsonb_build_object( + 'platformLimits', + jsonb_build_object('max_concurrent_requests', 5) + ) + ) AS default_rate_limit_policy, + COALESCE( + template.pricing_rule_set_id, + (SELECT id FROM model_pricing_rule_sets WHERE rule_set_key = 'default-multimodal-v1' LIMIT 1) + ) AS pricing_rule_set_id, + COALESCE( + template.runtime_policy_set_id, + (SELECT id FROM model_runtime_policy_sets WHERE policy_key = 'default-runtime-v1' LIMIT 1) + ) AS runtime_policy_set_id, + COALESCE(template.runtime_policy_override, '{}'::jsonb) AS runtime_policy_override, + COALESCE(template.pricing_version, 1) AS pricing_version + FROM model_defs + LEFT JOIN model_catalog_providers providers + ON providers.provider_key = model_defs.provider_key + OR providers.provider_code = model_defs.provider_key + LEFT JOIN base_model_catalog template + ON template.canonical_model_key = model_defs.template_key +), +payload AS ( + SELECT + source_rows.*, + jsonb_build_object( + 'source', 'server-main.integration-platform', + 'sourceProviderCode', source_provider_code, + 'sourceProviderName', source_provider_name, + 'sourceSpecType', source_spec_type, + 'officialModelId', provider_model_name, + 'officialDisplayName', display_name, + 'originalTypes', model_type, + 'alias', model_alias, + 'description', description, + 'iconPath', icon_path, + 'billingType', 'external-api', + 'billingMode', '', + 'referenceModel', '', + 'modelWeight', NULL, + 'selectable', true, + 'referenceDocument', reference_document, + 'tokenPlanOnly', token_plan_only, + 'rawModel', jsonb_build_object( + 'name', raw_model_name, + 'types', model_type, + 'icon_path', icon_path, + 'alias', model_alias, + 'description', description, + 'capabilities', capabilities - 'originalTypes' + ) + ) AS metadata + FROM source_rows +), +snapshot AS ( + SELECT + payload.*, + jsonb_build_object( + 'providerKey', provider_key, + 'canonicalModelKey', canonical_model_key, + 'invocationName', provider_model_name, + 'providerModelName', provider_model_name, + 'modelType', model_type, + 'modelAlias', provider_model_name, + 'displayName', display_name, + 'capabilities', capabilities, + 'baseBillingConfig', base_billing_config, + 'defaultRateLimitPolicy', default_rate_limit_policy, + 'pricingRuleSetId', COALESCE(pricing_rule_set_id::text, ''), + 'runtimePolicySetId', COALESCE(runtime_policy_set_id::text, ''), + 'runtimePolicyOverride', runtime_policy_override, + 'metadata', metadata, + 'pricingVersion', pricing_version, + 'status', 'active' + ) AS default_snapshot + FROM payload +) +INSERT INTO base_model_catalog ( + provider_id, + provider_key, + canonical_model_key, + invocation_name, + provider_model_name, + model_type, + display_name, + capabilities, + base_billing_config, + default_rate_limit_policy, + pricing_rule_set_id, + runtime_policy_set_id, + runtime_policy_override, + metadata, + catalog_type, + default_snapshot, + pricing_version, + status +) +SELECT + provider_id, + provider_key, + canonical_model_key, + provider_model_name, + provider_model_name, + model_type, + display_name, + capabilities, + base_billing_config, + default_rate_limit_policy, + pricing_rule_set_id, + runtime_policy_set_id, + runtime_policy_override, + metadata, + 'system', + default_snapshot, + pricing_version, + 'active' +FROM snapshot +ON CONFLICT (canonical_model_key) DO UPDATE +SET provider_id = EXCLUDED.provider_id, + provider_key = EXCLUDED.provider_key, + invocation_name = CASE + WHEN base_model_catalog.customized_at IS NULL THEN EXCLUDED.invocation_name + ELSE base_model_catalog.invocation_name + END, + provider_model_name = CASE + WHEN base_model_catalog.customized_at IS NULL THEN EXCLUDED.provider_model_name + ELSE base_model_catalog.provider_model_name + END, + model_type = CASE + WHEN base_model_catalog.customized_at IS NULL THEN EXCLUDED.model_type + ELSE base_model_catalog.model_type + END, + display_name = CASE + WHEN base_model_catalog.customized_at IS NULL THEN EXCLUDED.display_name + ELSE base_model_catalog.display_name + END, + capabilities = CASE + WHEN base_model_catalog.customized_at IS NULL THEN EXCLUDED.capabilities + ELSE base_model_catalog.capabilities + END, + base_billing_config = CASE + WHEN base_model_catalog.customized_at IS NULL THEN EXCLUDED.base_billing_config + ELSE base_model_catalog.base_billing_config + END, + default_rate_limit_policy = CASE + WHEN base_model_catalog.customized_at IS NULL THEN EXCLUDED.default_rate_limit_policy + ELSE base_model_catalog.default_rate_limit_policy + END, + pricing_rule_set_id = COALESCE(base_model_catalog.pricing_rule_set_id, EXCLUDED.pricing_rule_set_id), + runtime_policy_set_id = COALESCE(base_model_catalog.runtime_policy_set_id, EXCLUDED.runtime_policy_set_id), + runtime_policy_override = CASE + WHEN base_model_catalog.customized_at IS NULL THEN EXCLUDED.runtime_policy_override + ELSE base_model_catalog.runtime_policy_override + END, + metadata = CASE + WHEN base_model_catalog.customized_at IS NULL THEN EXCLUDED.metadata + ELSE base_model_catalog.metadata + END, + catalog_type = 'system', + default_snapshot = EXCLUDED.default_snapshot, + pricing_version = CASE + WHEN base_model_catalog.customized_at IS NULL THEN EXCLUDED.pricing_version + ELSE base_model_catalog.pricing_version + END, + status = CASE + WHEN base_model_catalog.customized_at IS NULL THEN 'active' + ELSE base_model_catalog.status + END, + updated_at = now(); + +INSERT INTO platform_models ( + platform_id, + base_model_id, + model_name, + provider_model_name, + model_alias, + model_type, + display_name, + capabilities, + pricing_mode, + billing_config, + retry_policy, + rate_limit_policy, + enabled +) +SELECT + platform.id, + base_model.id, + base_model.invocation_name, + base_model.provider_model_name, + base_model.invocation_name, + base_model.model_type, + base_model.display_name, + base_model.capabilities, + 'inherit_discount', + base_model.base_billing_config, + '{"enabled":true,"maxAttempts":1}'::jsonb, + base_model.default_rate_limit_policy, + true +FROM integration_platforms platform +JOIN base_model_catalog base_model + ON base_model.provider_key = platform.provider +WHERE base_model.canonical_model_key IN ( + 'easyai:Kimi-K3', + 'volces-openai:kimi-k3', + 'easyai:Qwen3.8-Max-Preview', + 'aliyun-bailian-openai:qwen3.8-max-preview' + ) + AND platform.deleted_at IS NULL +ON CONFLICT (platform_id, model_name) DO UPDATE +SET base_model_id = EXCLUDED.base_model_id, + provider_model_name = EXCLUDED.provider_model_name, + model_alias = EXCLUDED.model_alias, + display_name = EXCLUDED.display_name, + model_type = EXCLUDED.model_type, + capabilities = EXCLUDED.capabilities, + pricing_mode = EXCLUDED.pricing_mode, + billing_config = EXCLUDED.billing_config, + retry_policy = EXCLUDED.retry_policy, + rate_limit_policy = EXCLUDED.rate_limit_policy, + enabled = EXCLUDED.enabled, + updated_at = now(); diff --git a/apps/api/migrations/0089_file_storage_request_asset_scene.sql b/apps/api/migrations/0089_file_storage_request_asset_scene.sql new file mode 100644 index 0000000..6b1ad51 --- /dev/null +++ b/apps/api/migrations/0089_file_storage_request_asset_scene.sql @@ -0,0 +1,19 @@ +UPDATE file_storage_channels +SET + config = jsonb_set( + config, + '{scenes}', + CASE + WHEN jsonb_typeof(config->'scenes') = 'array' + THEN (config->'scenes') || '["request_asset"]'::jsonb + ELSE '["upload", "image_result", "request_asset"]'::jsonb + END, + true + ), + updated_at = now() +WHERE deleted_at IS NULL + AND provider = 'server_main_openapi' + AND ( + jsonb_typeof(config->'scenes') IS DISTINCT FROM 'array' + OR NOT (COALESCE(config->'scenes', '[]'::jsonb) ? 'request_asset') + ); diff --git a/apps/web/src/App.tsx b/apps/web/src/App.tsx index 9740241..9b9f99d 100644 --- a/apps/web/src/App.tsx +++ b/apps/web/src/App.tsx @@ -1,5 +1,6 @@ import { useEffect, useMemo, useRef, useState, type FormEvent } from 'react'; import type { + AdminGatewayTask, BaseModelCatalogItem, CatalogProvider, ClientCustomizationSettings, @@ -67,6 +68,7 @@ import { getSecurityEventConnection, getWalletSummary, listAccessRules, + listAdminTasks, listAuditLogs, listApiKeyAccessRules, listApiKeyAssignableModels, @@ -141,7 +143,9 @@ import { ModelsPage } from './pages/ModelsPage'; import { PlaygroundPage } from './pages/PlaygroundPage'; import { WorkspacePage } from './pages/WorkspacePage'; import { + adminTaskQueryKey, parseAppRoute, + pathForAdminTaskQuery, pathForAdminSection, pathForApiDocSection, pathForPage, @@ -153,6 +157,7 @@ import { } from './routing'; import type { AdminSection, + AdminTaskQuery, ApiDocSection, ApiKeyForm, AuthMode, @@ -197,6 +202,7 @@ type DataKey = | 'tenants' | 'users' | 'userGroups' + | 'adminTasks' | 'tasks' | 'wallet' | 'walletTransactions' @@ -208,6 +214,7 @@ export function App() { const initialRoute = parseAppRoute(); const [activePage, setActivePage] = useState(initialRoute.activePage); const [adminSection, setAdminSection] = useState(initialRoute.adminSection); + const [adminTaskQuery, setAdminTaskQuery] = useState(initialRoute.adminTaskQuery); const [workspaceSection, setWorkspaceSection] = useState(initialRoute.workspaceSection); const [workspaceTaskQuery, setWorkspaceTaskQuery] = useState(initialRoute.workspaceTaskQuery); const [workspaceTransactionQuery, setWorkspaceTransactionQuery] = useState(() => defaultWorkspaceTransactionQuery()); @@ -258,6 +265,10 @@ export function App() { const [taskResult, setTaskResult] = useState(null); const [tasks, setTasks] = useState([]); const [taskTotal, setTaskTotal] = useState(0); + const [adminTasks, setAdminTasks] = useState([]); + const [adminTaskTotal, setAdminTaskTotal] = useState(0); + const [adminTasksUpdatedAt, setAdminTasksUpdatedAt] = useState(null); + const [adminTaskState, setAdminTaskState] = useState('idle'); const [walletAccounts, setWalletAccounts] = useState([]); const [walletTransactions, setWalletTransactions] = useState([]); const [walletTransactionTotal, setWalletTransactionTotal] = useState(0); @@ -279,6 +290,8 @@ export function App() { const loadingDataKeysRef = useRef(new Set()); const loadedTaskQueryKeyRef = useRef(''); const currentTaskQueryKeyRef = useRef(''); + const loadedAdminTaskQueryKeyRef = useRef(''); + const currentAdminTaskQueryKeyRef = useRef(''); const loadedTransactionQueryKeyRef = useRef(''); const currentTransactionQueryKeyRef = useRef(''); const { removeBaseModel, removeProvider, resetAllBaseModelsToDefault, resetBaseModelToDefault, saveBaseModel, saveProvider } = useCatalogOperations({ @@ -302,8 +315,10 @@ export function App() { token, }); const taskListRequestKey = workspaceTaskQueryKey(workspaceTaskQuery); + const adminTaskListRequestKey = adminTaskQueryKey(adminTaskQuery); const transactionListRequestKey = workspaceTransactionQueryKey(workspaceTransactionQuery); currentTaskQueryKeyRef.current = taskListRequestKey; + currentAdminTaskQueryKeyRef.current = adminTaskListRequestKey; currentTransactionQueryKeyRef.current = transactionListRequestKey; useEffect(() => { @@ -360,7 +375,7 @@ export function App() { }, []); useEffect(() => { void ensureRouteData(token); - }, [activePage, adminSection, taskListRequestKey, transactionListRequestKey, workspaceSection, token]); + }, [activePage, adminSection, adminTaskListRequestKey, taskListRequestKey, transactionListRequestKey, workspaceSection, token]); useEffect(() => { if (!token || activePage !== 'admin' || adminSection !== 'realtimeLoad') return undefined; const timer = window.setInterval(() => { @@ -409,6 +424,7 @@ export function App() { const data = useMemo(() => ({ accessRules, + adminTasks, auditLogs, apiKeys, baseModels, @@ -437,7 +453,7 @@ export function App() { users, walletAccounts, walletTransactions, - }), [accessRules, apiKeys, auditLogs, baseModels, clientCustomizationSettings, currentUser, currentUserGroups, fileStorageChannels, fileStorageSettings, modelCatalog, modelRateLimits, modelRateLimitsUpdatedAt, models, networkProxyConfig, platforms, pricingRuleSets, pricingRules, providers, rateLimitWindows, runnerPolicy, runtimePolicySets, securityEventConnection, taskResult, tasks, tenants, userGroups, users, walletAccounts, walletTransactions]); + }), [accessRules, adminTasks, apiKeys, auditLogs, baseModels, clientCustomizationSettings, currentUser, currentUserGroups, fileStorageChannels, fileStorageSettings, modelCatalog, modelRateLimits, modelRateLimitsUpdatedAt, models, networkProxyConfig, platforms, pricingRuleSets, pricingRules, providers, rateLimitWindows, runnerPolicy, runtimePolicySets, securityEventConnection, taskResult, tasks, tenants, userGroups, users, walletAccounts, walletTransactions]); async function refresh(nextToken = token) { await ensureRouteData(nextToken, true); @@ -452,6 +468,10 @@ export function App() { loadedDataKeysRef.current.delete('tasks'); loadingDataKeysRef.current.delete('tasks'); } + if (activePage === 'admin' && adminSection === 'tasks' && loadedAdminTaskQueryKeyRef.current !== adminTaskListRequestKey) { + loadedDataKeysRef.current.delete('adminTasks'); + loadingDataKeysRef.current.delete('adminTasks'); + } if (activePage === 'workspace' && workspaceSection === 'transactions' && loadedTransactionQueryKeyRef.current !== transactionListRequestKey) { loadedDataKeysRef.current.delete('walletTransactions'); loadingDataKeysRef.current.delete('walletTransactions'); @@ -590,6 +610,26 @@ export function App() { loadedTaskQueryKeyRef.current = requestKey; } return; + case 'adminTasks': + { + const requestKey = adminTaskListRequestKey; + setAdminTaskState('loading'); + try { + const response = await listAdminTasks(nextToken, adminTaskQuery); + if (requestKey !== currentAdminTaskQueryKeyRef.current) return; + setAdminTasks(response.items); + setAdminTaskTotal(response.total ?? response.items.length); + setAdminTasksUpdatedAt(Date.now()); + loadedAdminTaskQueryKeyRef.current = requestKey; + setAdminTaskState('ready'); + } catch (err) { + if (requestKey === currentAdminTaskQueryKeyRef.current) { + setAdminTaskState('error'); + } + throw err; + } + } + return; case 'wallet': { const response = await getWalletSummary(nextToken); setWalletAccounts(response.accounts); @@ -1219,6 +1259,10 @@ export function App() { setTaskResult(null); setTasks([]); setTaskTotal(0); + setAdminTasks([]); + setAdminTaskTotal(0); + setAdminTasksUpdatedAt(null); + setAdminTaskState('idle'); setWalletAccounts([]); setWalletTransactions([]); setWalletTransactionTotal(0); @@ -1270,12 +1314,13 @@ export function App() { } function currentRouteState(): AppRouteState { - return { activePage, adminSection, apiDocSection, playgroundMode, workspaceSection, workspaceTaskQuery }; + return { activePage, adminSection, adminTaskQuery, apiDocSection, playgroundMode, workspaceSection, workspaceTaskQuery }; } function applyRoute(route: AppRouteState) { setActivePage(route.activePage); setAdminSection(route.adminSection); + setAdminTaskQuery(route.adminTaskQuery); setApiDocSection(route.apiDocSection); setPlaygroundMode(route.playgroundMode); setWorkspaceSection(route.workspaceSection); @@ -1298,6 +1343,15 @@ export function App() { navigatePath(pathForAdminSection(section)); } + function navigateAdminTaskQuery(query: AdminTaskQuery) { + navigatePath(pathForAdminTaskQuery(query)); + } + + async function refreshAdminTasks() { + loadedDataKeysRef.current.delete('adminTasks'); + await ensureData(['adminTasks'], token, true); + } + function navigateWorkspaceSection(section: WorkspaceSection) { navigatePath(pathForWorkspaceSection(section)); } @@ -1405,7 +1459,11 @@ export function App() { {activePage === 'admin' && ( isAuthenticated ? ( setCoreMessage('')} + onAdminTaskQueryChange={navigateAdminTaskQuery} + onRefreshAdminTasks={refreshAdminTasks} onSectionChange={navigateAdminSection} /> ) : ( @@ -1513,6 +1573,7 @@ function mergeExistingPlatformModelInput(input: PlatformModelBindingInput, curre discountFactor: (input.discountFactor ?? existing.discountFactor) || undefined, pricingRuleSetId: input.pricingRuleSetId ?? existing.pricingRuleSetId, rateLimitPolicy: input.rateLimitPolicy ?? existing.rateLimitPolicy, + rateLimitPolicyMode: input.rateLimitPolicyMode ?? existing.rateLimitPolicyMode, retryPolicy: input.retryPolicy ?? existing.retryPolicy, runtimePolicyOverride: input.runtimePolicyOverride ?? (existing.runtimePolicyOverride as Record | undefined), runtimePolicySetId: input.runtimePolicySetId ?? existing.runtimePolicySetId, @@ -1646,22 +1707,24 @@ function dataKeysForRoute( return ['pricingRuleSets']; case 'runtime': return ['runtimePolicySets', 'runnerPolicy']; + case 'billingSettlements': + return []; case 'baseModels': return ['baseModels', 'providers', 'pricingRuleSets', 'runtimePolicySets']; case 'platforms': return ['platforms', 'models', 'providers', 'baseModels', 'pricingRuleSets', 'networkProxyConfig']; case 'realtimeLoad': return ['platforms', 'modelRateLimits']; + case 'tasks': + return ['adminTasks', 'tenants', 'users', 'userGroups', 'platforms', 'models']; case 'tenants': return ['tenants', 'userGroups']; case 'users': return ['users', 'tenants', 'userGroups']; case 'userGroups': - return ['userGroups']; + return ['userGroups', 'accessRules', 'platforms', 'models']; case 'auditLogs': return ['auditLogs']; - case 'accessRules': - return ['accessRules', 'userGroups', 'platforms', 'models']; case 'systemSettings': return ['clientCustomizationSettings', 'fileStorageSettings', 'fileStorageChannels', 'securityEventConnection']; default: diff --git a/apps/web/src/api.ts b/apps/web/src/api.ts index 30ba323..4f96374 100644 --- a/apps/web/src/api.ts +++ b/apps/web/src/api.ts @@ -1,4 +1,5 @@ import type { + AdminGatewayTask, AuthResponse, AuthUser, BaseModelCatalogItem, @@ -59,7 +60,7 @@ import type { WalletRechargeRequest, WalletSummaryResponse, } from '@easyai-ai-gateway/contracts'; -import type { PlatformCreateInput, PlatformModelBindingInput, WorkspaceTaskQuery } from './types'; +import type { AdminTaskQuery, PlatformCreateInput, PlatformModelBindingInput, WorkspaceTaskQuery } from './types'; const API_BASE = import.meta.env.VITE_GATEWAY_API_BASE_URL ?? 'http://localhost:8088'; export const OIDC_BROWSER_SESSION_CREDENTIAL = '__easyai_gateway_oidc_browser_session__'; @@ -967,6 +968,42 @@ export async function listTasks(token: string, query: WorkspaceTaskQuery): Promi return request>(`/api/workspace/tasks?${search.toString()}`, { token }); } +export async function listAdminTasks(token: string, query: AdminTaskQuery): Promise> { + const search = new URLSearchParams({ + page: String(query.page), + pageSize: String(query.pageSize), + }); + if (query.query) search.set('q', query.query); + if (query.tenantId) search.set('tenantId', query.tenantId); + if (query.userId) search.set('userId', query.userId); + if (query.userGroupId) search.set('userGroupId', query.userGroupId); + if (query.status) search.set('status', query.status); + if (query.platformId) search.set('platformId', query.platformId); + if (query.model) search.set('model', query.model); + if (query.modelType) search.set('modelType', query.modelType); + if (query.runMode) search.set('runMode', query.runMode); + if (query.billingStatus) search.set('billingStatus', query.billingStatus); + if (query.apiKey) search.set('apiKey', query.apiKey); + if (query.createdFrom) search.set('createdFrom', query.createdFrom); + if (query.createdTo) search.set('createdTo', query.createdTo); + return request>(`/api/admin/tasks?${search.toString()}`, { token }); +} + +export async function getAdminTask(token: string, taskId: string, sensitive: 'masked' | 'full' = 'masked'): Promise { + return request(`/api/admin/tasks/${taskId}?sensitive=${sensitive}`, { token }); +} + +export async function listAdminTaskParamPreprocessing( + token: string, + taskId: string, + sensitive: 'masked' | 'full' = 'masked', +): Promise> { + return request>( + `/api/admin/tasks/${taskId}/param-preprocessing?sensitive=${sensitive}`, + { token }, + ); +} + export async function getWalletSummary(token: string): Promise { return request('/api/workspace/wallet', { token }); } diff --git a/apps/web/src/app-state.ts b/apps/web/src/app-state.ts index 9ef928f..77301f7 100644 --- a/apps/web/src/app-state.ts +++ b/apps/web/src/app-state.ts @@ -1,4 +1,5 @@ import type { + AdminGatewayTask, AuthUser, BaseModelCatalogItem, CatalogProvider, @@ -29,6 +30,7 @@ import type { export interface ConsoleData { accessRules: GatewayAccessRule[]; + adminTasks: AdminGatewayTask[]; auditLogs: GatewayAuditLog[]; apiKeys: GatewayApiKey[]; baseModels: BaseModelCatalogItem[]; diff --git a/apps/web/src/components/AuthPanel.tsx b/apps/web/src/components/AuthPanel.tsx index 1440d99..34bc373 100644 --- a/apps/web/src/components/AuthPanel.tsx +++ b/apps/web/src/components/AuthPanel.tsx @@ -1,6 +1,6 @@ import type { FormEvent } from 'react'; import { LogIn, UserPlus } from 'lucide-react'; -import { Button, Card, CardContent, CardHeader, CardTitle, Input, Label, Tabs } from './ui'; +import { Button, Card, CardContent, CardHeader, CardTitle, Input, Label, PasswordInput, Tabs } from './ui'; import type { AuthMode, LoadState, LoginForm, RegisterForm } from '../types'; const tabs = [ @@ -70,9 +70,8 @@ function LoginFormView(props: {