chore(git): 同步主分支发布与业务变更
# Conflicts: # apps/api/go.mod # apps/api/go.sum
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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/<SHA>.json`,不得修改生产。
|
||||
6. publish 成功后,Agent 必须报告 manifest、组件、digest 和验证结果并停止。只有用户再次明确确认上线,才可单独执行:
|
||||
该命令只能构建、冒烟、推送镜像和生成 `dist/releases/<SHA>.json`,不得修改生产。用户只要求“构建镜像”“推送镜像”或明确限定为 publish 时,执行到此为止。
|
||||
6. 用户明确要求“发布”“上线”或“部署到线上”时,该次授权同时覆盖 publish 和 deploy。publish 成功并核验 manifest、组件、digest 和验证结果后,应直接继续执行下列部署命令,无需再次请求确认:
|
||||
|
||||
```bash
|
||||
./scripts/deploy-production-release.sh dist/releases/<SHA>.json
|
||||
|
||||
+3
-1
@@ -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 && \
|
||||
|
||||
@@ -104,14 +104,14 @@ docker compose -f docker-compose.yml restart web
|
||||
|
||||
本仓库没有 Gitea Actions、Webhook、Tag、`main` Push、轮询或定时发布。`main` 不使用受保护分支限制;commit、push 和 Tag 都不会构建镜像或更新生产。
|
||||
|
||||
生产发布由 Agent 在用户明确指令下分两次执行。第一步在本机构建 `linux/amd64` 镜像、运行临时 PostgreSQL + simulation API 冒烟、推送完整 SHA Tag,并生成带内容完整性校验的 digest-pinned manifest;该命令不会修改生产:
|
||||
生产发布由 Agent 在用户明确指令下按两个阶段连续执行。第一步在本机构建 `linux/amd64` 镜像、运行临时 PostgreSQL + simulation API 冒烟、推送完整 SHA Tag,并生成带内容完整性校验的 digest-pinned manifest;该命令不会修改生产:
|
||||
|
||||
```bash
|
||||
docker login --username=<your-aliyun-account> registry.cn-shanghai.aliyuncs.com
|
||||
./scripts/publish-release-images.sh --components auto
|
||||
```
|
||||
|
||||
Agent 必须报告 `dist/releases/<SHA>.json` 后停止。用户再次明确确认上线后,才执行第二步:
|
||||
用户要求“发布”“上线”或“部署到线上”时,Agent 核验并报告 `dist/releases/<SHA>.json`、组件和 digest 后直接执行第二步,无需再次确认;用户仅要求构建或推送镜像时才在第一步后停止:
|
||||
|
||||
```bash
|
||||
./scripts/deploy-production-release.sh dist/releases/<SHA>.json
|
||||
@@ -124,7 +124,10 @@ Agent 必须报告 `dist/releases/<SHA>.json` 后停止。用户再次明确确
|
||||
./scripts/deploy-production-release.sh --rollback <历史完整 Git SHA>
|
||||
```
|
||||
|
||||
完整安装、验证、失败处理和停用旧自动化步骤见[人工生产发布运行手册](docs/runbooks/production-ci-cd.md),决策背景见 [ADR-003](docs/decisions/003-manual-agent-release.md)。
|
||||
完整安装、验证、失败处理和停用旧自动化步骤见[人工生产发布运行手册](docs/runbooks/production-ci-cd.md);
|
||||
三节点 K3s、CloudNativePG、双 NGINX、备份恢复和跨节点文件验收见
|
||||
[K3s 高可用运行手册](docs/operations/k3s-ha-runbook.md)。决策背景见
|
||||
[ADR-003](docs/decisions/003-manual-agent-release.md)。
|
||||
|
||||
Compose 默认使用独立容器数据库 `postgres:18-alpine`,数据卷会保留在 `postgres_data` 和 `api_data`。为避免本地开发 `.env` 中的 `localhost` 数据库地址污染容器部署,compose 使用 `AI_GATEWAY_COMPOSE_*` 变量作为容器部署专用覆盖,例如:
|
||||
|
||||
@@ -159,6 +162,8 @@ AI_GATEWAY_DATABASE_URL=postgresql://easyai:easyai2025@localhost:5432/easyai_ai_
|
||||
|
||||
如果现有 `easyai-pgvector` 没有把 `5432` 映射到宿主机,就需要补端口映射,或者把 AI Gateway 后端容器化后接入同一个 `easyai` Docker network。
|
||||
|
||||
异步队列 worker 不使用固定并发。服务分别汇总启用平台模型和活跃用户组的有效 `concurrent` 策略,采用两者中更严格的容量,默认每 5 秒在线调整 River 执行容量。策略解析兼容历史 `platformLimits/modelLimits.max_concurrent_requests`,运行时统一转换为 `rules`。`AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT` 默认 `2048`,仅作为单进程资源安全边界;平台模型和用户组的 PostgreSQL concurrency lease 才是业务并发真值。可通过 `AI_GATEWAY_ASYNC_WORKER_REFRESH_INTERVAL_SECONDS` 调整刷新周期。
|
||||
|
||||
## 迁移原则
|
||||
|
||||
1. 新服务先并行运行,不直接删除 `easyai-server-main` 内现有模块。
|
||||
|
||||
@@ -0,0 +1,113 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"flag"
|
||||
"log/slog"
|
||||
"os"
|
||||
"time"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/runner"
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
)
|
||||
|
||||
func main() {
|
||||
apply := flag.Bool("apply", false, "persist compacted results; default is dry-run")
|
||||
batchSize := flag.Int("batch-size", 100, "rows per batch, maximum 100")
|
||||
maxBatches := flag.Int("max-batches", 10, "maximum batches for one invocation")
|
||||
afterID := flag.String("after-id", "", "resume after this task UUID")
|
||||
flag.Parse()
|
||||
|
||||
logger := slog.New(slog.NewJSONHandler(os.Stdout, nil))
|
||||
cfg := config.Load()
|
||||
if err := cfg.Validate(); err != nil {
|
||||
logger.Error("invalid gateway configuration", "error", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
if *batchSize < 1 || *batchSize > 100 || *maxBatches < 1 {
|
||||
logger.Error("invalid backfill bounds", "batchSize", *batchSize, "maxBatches", *maxBatches)
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
ctx := context.Background()
|
||||
db, err := store.Connect(ctx, cfg.DatabaseURL)
|
||||
if err != nil {
|
||||
logger.Error("connect postgres failed", "error", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
defer db.Close()
|
||||
|
||||
service := runner.New(cfg, db, logger)
|
||||
cursor := *afterID
|
||||
scanned := 0
|
||||
matched := 0
|
||||
updated := 0
|
||||
expired := 0
|
||||
for batch := 0; batch < *maxBatches; batch++ {
|
||||
items, err := db.ListTaskBinaryResultBackfillBatch(ctx, cursor, *batchSize)
|
||||
if err != nil {
|
||||
logger.Error("list binary result backfill batch failed", "afterId", cursor, "error", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
if len(items) == 0 {
|
||||
break
|
||||
}
|
||||
for _, item := range items {
|
||||
cursor = item.ID
|
||||
scanned++
|
||||
if !runner.TaskResultHasInlineBinary(item.Result) {
|
||||
continue
|
||||
}
|
||||
matched++
|
||||
if !*apply {
|
||||
continue
|
||||
}
|
||||
isExpired := item.FinishedAt.Before(time.Now().Add(-time.Duration(localResultTTLHours(cfg)) * time.Hour))
|
||||
var persistent map[string]any
|
||||
var changed bool
|
||||
if isExpired {
|
||||
persistent, changed, err = service.CompactExpiredTaskResultForStorage(ctx, item.ID, item.Result)
|
||||
} else {
|
||||
persistent, changed, err = service.MaterializeTaskResultForStorage(ctx, item.ID, item.Result)
|
||||
}
|
||||
if err != nil {
|
||||
logger.Error("materialize historical binary result failed", "taskId", item.ID, "error", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
if !changed {
|
||||
continue
|
||||
}
|
||||
ok, err := db.UpdateTaskBinaryResultBackfill(ctx, item.ID, persistent)
|
||||
if err != nil {
|
||||
logger.Error("update historical binary result failed", "taskId", item.ID, "error", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
updated++
|
||||
if isExpired {
|
||||
expired++
|
||||
}
|
||||
}
|
||||
if len(items) < *batchSize {
|
||||
break
|
||||
}
|
||||
}
|
||||
logger.Info("binary result backfill completed",
|
||||
"apply", *apply,
|
||||
"scanned", scanned,
|
||||
"matched", matched,
|
||||
"updated", updated,
|
||||
"expired", expired,
|
||||
"resumeAfterId", cursor,
|
||||
)
|
||||
}
|
||||
|
||||
func localResultTTLHours(cfg config.Config) int {
|
||||
if cfg.LocalResultTTLHours <= 0 {
|
||||
return 24
|
||||
}
|
||||
return cfg.LocalResultTTLHours
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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 {
|
||||
|
||||
+998
-37
File diff suppressed because it is too large
Load Diff
+693
-37
File diff suppressed because it is too large
Load Diff
@@ -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
|
||||
)
|
||||
|
||||
@@ -73,6 +73,8 @@ golang.org/x/crypto v0.52.0 h1:RMs7fP2rXdep0CftQlK8Uf+kibLm7qkCcradZWYz988=
|
||||
golang.org/x/crypto v0.52.0/go.mod h1:1QgfPxDqh0T2M/elOJtp9RvuR95kVjir0e6/BvEmGbc=
|
||||
golang.org/x/net v0.54.0 h1:2zJIZAxAHV/OHCDTCOHAYehQzLfSXuf/5SoL/Dv6w/w=
|
||||
golang.org/x/net v0.54.0/go.mod h1:Sj4oj8jK6XmHpBZU/zWHw3BV3abl4Kvi+Ut7cQcY+cQ=
|
||||
golang.org/x/image v0.43.0 h1:FLxcP4ec2350nTfOC8ysKtqYSIFbk/QGjw1ZHNP4tsY=
|
||||
golang.org/x/image v0.43.0/go.mod h1:rrpelvGFt+kLPAjPM4HeWPgrl0FtafueU//e5N0qk/Q=
|
||||
golang.org/x/oauth2 v0.36.0 h1:peZ/1z27fi9hUOFCAZaHyrpWG5lwe0RJEEEeH0ThlIs=
|
||||
golang.org/x/oauth2 v0.36.0/go.mod h1:YDBUJMTkDnJS+A4BP4eZBjCqtokkg1hODuPjwiGPO7Q=
|
||||
golang.org/x/sync v0.21.0 h1:HLII4xRRTtCRkxYp4HNFF0Js/Og6q2i++KXbg0gHCwM=
|
||||
|
||||
@@ -78,6 +78,15 @@ func applyAliyunReasoning(body map[string]any, candidate store.RuntimeModelCandi
|
||||
|
||||
body["enable_thinking"] = true
|
||||
model := chatReasoningModelName(body, candidate)
|
||||
if isAliyunQwen38MaxPreview(model) {
|
||||
if budget, ok := positiveIntFromAny(body["thinking_budget_tokens"]); ok {
|
||||
body["thinking_budget"] = budget
|
||||
delete(body, "reasoning_effort")
|
||||
return
|
||||
}
|
||||
body["reasoning_effort"] = qwen38MaxPreviewReasoningEffort(effort)
|
||||
return
|
||||
}
|
||||
if isAliyunHighMaxReasoningModel(model) {
|
||||
body["reasoning_effort"] = highMaxReasoningEffort(effort)
|
||||
return
|
||||
@@ -200,6 +209,10 @@ func isAliyunHighMaxReasoningModel(model string) bool {
|
||||
return strings.Contains(model, "deepseek-v4") || strings.HasPrefix(model, "glm-")
|
||||
}
|
||||
|
||||
func isAliyunQwen38MaxPreview(model string) bool {
|
||||
return strings.Contains(model, "qwen3.8-max-preview")
|
||||
}
|
||||
|
||||
func isAliyunThinkingBudgetModel(model string) bool {
|
||||
return strings.Contains(model, "qwen") || strings.Contains(model, "qwq") || strings.Contains(model, "qvq") || strings.Contains(model, "kimi")
|
||||
}
|
||||
@@ -219,6 +232,17 @@ func highMaxReasoningEffort(effort string) string {
|
||||
return "high"
|
||||
}
|
||||
|
||||
func qwen38MaxPreviewReasoningEffort(effort string) string {
|
||||
switch effort {
|
||||
case "low", "minimal":
|
||||
return "low"
|
||||
case "medium":
|
||||
return "medium"
|
||||
default:
|
||||
return "xhigh"
|
||||
}
|
||||
}
|
||||
|
||||
func zhipuReasoningEffort(effort string) string {
|
||||
if effort == "max" {
|
||||
return "xhigh"
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -326,17 +326,17 @@ func (c MinimaxClient) runSpeech(ctx context.Context, request Request) (Response
|
||||
if err != nil {
|
||||
return Response{}, &ClientError{Code: "invalid_response", Message: "minimax speech audio hex is invalid: " + err.Error(), RequestID: firstNonEmptyString(requestID, requestIDFromResult(result)), ResponseStartedAt: startedAt, ResponseFinishedAt: finishedAt, ResponseDurationMS: responseDurationMS(startedAt, finishedAt), Retryable: false}
|
||||
}
|
||||
normalized := cloneMapAny(result)
|
||||
normalized["status"] = "success"
|
||||
normalized["created"] = time.Now().UnixMilli()
|
||||
normalized["model"] = request.Model
|
||||
normalized["raw_data"] = cloneMapAny(result)
|
||||
normalized["data"] = []any{map[string]any{
|
||||
"type": "audio",
|
||||
"content": "data:audio/mpeg;base64," + base64.StdEncoding.EncodeToString(audioBytes),
|
||||
"mime_type": "audio/mpeg",
|
||||
"uploaded": false,
|
||||
}}
|
||||
normalized := map[string]any{
|
||||
"status": "success",
|
||||
"created": time.Now().UnixMilli(),
|
||||
"model": request.Model,
|
||||
"data": []any{map[string]any{
|
||||
"type": "audio",
|
||||
"content": "data:audio/mpeg;base64," + base64.StdEncoding.EncodeToString(audioBytes),
|
||||
"mime_type": "audio/mpeg",
|
||||
"uploaded": false,
|
||||
}},
|
||||
}
|
||||
return Response{
|
||||
Result: normalized,
|
||||
RequestID: firstNonEmptyString(requestID, requestIDFromResult(result)),
|
||||
@@ -378,12 +378,12 @@ func (c MinimaxClient) runVoiceClone(ctx context.Context, request Request) (Resp
|
||||
if isProviderTaskFailure(providerTaskSpec{Name: "minimax"}, result) {
|
||||
return Response{}, providerTaskFailure(providerTaskSpec{Name: "minimax"}, result, firstNonEmptyString(requestID, uploadRequestID, requestIDFromResult(result)), startedAt)
|
||||
}
|
||||
normalized := cloneMapAny(result)
|
||||
normalized["status"] = "success"
|
||||
normalized["created"] = time.Now().UnixMilli()
|
||||
normalized["model"] = request.Model
|
||||
normalized["voice_id"] = stringFromAny(payload["voice_id"])
|
||||
normalized["raw_data"] = cloneMapAny(result)
|
||||
normalized := map[string]any{
|
||||
"status": "success",
|
||||
"created": time.Now().UnixMilli(),
|
||||
"model": request.Model,
|
||||
"voice_id": stringFromAny(payload["voice_id"]),
|
||||
}
|
||||
if demoAudio := firstNonEmptyString(valueAtPath(result, "demo_audio"), valueAtPath(result, "data.demo_audio")); demoAudio != "" {
|
||||
normalized["demo_audio"] = demoAudio
|
||||
normalized["data"] = []any{map[string]any{"type": "audio", "url": demoAudio}}
|
||||
|
||||
@@ -0,0 +1,232 @@
|
||||
package clients
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
)
|
||||
|
||||
func TestGeminiGenerationConfigDropsDefaultRatioTokensAndNormalizesKSize(t *testing.T) {
|
||||
config := geminiGenerationConfig(map[string]any{
|
||||
"aspect_ratio": "auto",
|
||||
"size": "2k",
|
||||
"generationConfig": map[string]any{
|
||||
"imageConfig": map[string]any{
|
||||
"aspectRatio": "adaptive",
|
||||
"imageSize": "2k",
|
||||
},
|
||||
},
|
||||
}, true)
|
||||
|
||||
imageConfig := mapFromAny(config["imageConfig"])
|
||||
if _, ok := imageConfig["aspectRatio"]; ok {
|
||||
t.Fatalf("non-ratio tokens must not be sent to Gemini: %+v", config)
|
||||
}
|
||||
if imageConfig["imageSize"] != "2K" {
|
||||
t.Fatalf("Gemini imageSize should use canonical K format: %+v", config)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNormalizeOfficialOpenAIImageSizeUsesModelSpecificWireFormat(t *testing.T) {
|
||||
legacy := map[string]any{
|
||||
"model": "gpt-image-1",
|
||||
"size": "2K",
|
||||
"width": 2048,
|
||||
"height": 1152,
|
||||
"aspect_ratio": "16:9",
|
||||
}
|
||||
normalizeOfficialOpenAIImageSize(legacy, map[string]any{
|
||||
"size": "2K",
|
||||
"aspect_ratio": "16:9",
|
||||
})
|
||||
if legacy["size"] != "1536x1024" {
|
||||
t.Fatalf("legacy GPT Image should use an official orientation size: %+v", legacy)
|
||||
}
|
||||
|
||||
flexible := map[string]any{
|
||||
"model": "gpt-image-2",
|
||||
"size": "2K",
|
||||
"resolution": "2k",
|
||||
"aspect_ratio": "16:9",
|
||||
}
|
||||
normalizeOfficialOpenAIImageSize(flexible, map[string]any{
|
||||
"size": "2K",
|
||||
"aspect_ratio": "16:9",
|
||||
})
|
||||
if flexible["size"] != "2048x1152" {
|
||||
t.Fatalf("GPT Image 2 should receive flexible pixel dimensions: %+v", flexible)
|
||||
}
|
||||
|
||||
rounded := map[string]any{
|
||||
"model": "gpt-image-2",
|
||||
"size": "2K",
|
||||
"resolution": "2K",
|
||||
"aspect_ratio": "3:2",
|
||||
}
|
||||
normalizeOfficialOpenAIImageSize(rounded, map[string]any{
|
||||
"size": "2K",
|
||||
"aspect_ratio": "3:2",
|
||||
})
|
||||
if rounded["size"] != "2048x1360" {
|
||||
t.Fatalf("GPT Image 2 dimensions should satisfy the official 16-pixel multiple: %+v", rounded)
|
||||
}
|
||||
|
||||
explicit := map[string]any{
|
||||
"model": "gpt-image-2",
|
||||
"size": "1232x768",
|
||||
"width": 1232,
|
||||
"height": 768,
|
||||
}
|
||||
normalizeOfficialOpenAIImageSize(explicit, map[string]any{
|
||||
"size": "1234x777",
|
||||
})
|
||||
if explicit["size"] != "1234x777" {
|
||||
t.Fatalf("an explicit WxH size must be sent unchanged to OpenAI: %+v", explicit)
|
||||
}
|
||||
|
||||
explicitLegacy := map[string]any{
|
||||
"model": "gpt-image-1",
|
||||
"size": "2048x1152",
|
||||
}
|
||||
normalizeOfficialOpenAIImageSize(explicitLegacy, map[string]any{
|
||||
"size": "2048x1152",
|
||||
})
|
||||
if explicitLegacy["size"] != "2048x1152" {
|
||||
t.Fatalf("explicit WxH must also win for legacy OpenAI image models: %+v", explicitLegacy)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIClientUsesOriginalWxHAndConvertsKResolution(t *testing.T) {
|
||||
received := make([]map[string]any, 0, 3)
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
var body map[string]any
|
||||
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
|
||||
t.Fatalf("decode OpenAI image request: %v", err)
|
||||
}
|
||||
received = append(received, body)
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"data": []any{map[string]any{"b64_json": "aW1hZ2U="}},
|
||||
})
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
candidate := store.RuntimeModelCandidate{
|
||||
BaseURL: server.URL,
|
||||
Provider: "openai",
|
||||
ProviderModelName: "gpt-image-2",
|
||||
ModelType: "image_generate",
|
||||
Credentials: map[string]any{"apiKey": "openai-key"},
|
||||
}
|
||||
requests := []Request{
|
||||
{
|
||||
Kind: "images.generations",
|
||||
ModelType: "image_generate",
|
||||
Body: map[string]any{
|
||||
"prompt": "explicit dimensions",
|
||||
"size": "1232x768",
|
||||
"width": 1232,
|
||||
"height": 768,
|
||||
},
|
||||
OriginalBody: map[string]any{
|
||||
"size": "1234x777",
|
||||
},
|
||||
Candidate: candidate,
|
||||
},
|
||||
{
|
||||
Kind: "images.generations",
|
||||
ModelType: "image_generate",
|
||||
Body: map[string]any{
|
||||
"prompt": "resolution conversion",
|
||||
"size": "2048x1365",
|
||||
"resolution": "2K",
|
||||
"aspect_ratio": "3:2",
|
||||
"width": 2048,
|
||||
"height": 1365,
|
||||
},
|
||||
OriginalBody: map[string]any{
|
||||
"size": "2K",
|
||||
"aspect_ratio": "3:2",
|
||||
},
|
||||
Candidate: candidate,
|
||||
},
|
||||
{
|
||||
Kind: "images.generations",
|
||||
ModelType: "image_generate",
|
||||
Body: map[string]any{
|
||||
"prompt": "resolution field conversion",
|
||||
"size": "1024x1024",
|
||||
"resolution": "1K",
|
||||
"aspect_ratio": "1:1",
|
||||
"width": 1024,
|
||||
"height": 1024,
|
||||
},
|
||||
OriginalBody: map[string]any{
|
||||
"resolution": "1k",
|
||||
"aspect_ratio": "1:1",
|
||||
},
|
||||
Candidate: candidate,
|
||||
},
|
||||
}
|
||||
for _, request := range requests {
|
||||
if _, err := (OpenAIClient{HTTPClient: server.Client()}).Run(context.Background(), request); err != nil {
|
||||
t.Fatalf("run OpenAI image request: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
if len(received) != 3 {
|
||||
t.Fatalf("unexpected request count: %d", len(received))
|
||||
}
|
||||
if received[0]["size"] != "1234x777" {
|
||||
t.Fatalf("explicit WxH should win over preprocessed dimensions: %+v", received[0])
|
||||
}
|
||||
if received[1]["size"] != "2048x1360" {
|
||||
t.Fatalf("2K with 3:2 should use normalized OpenAI dimensions: %+v", received[1])
|
||||
}
|
||||
if received[2]["size"] != "1024x1024" {
|
||||
t.Fatalf("resolution=1K should be converted to OpenAI dimensions: %+v", received[2])
|
||||
}
|
||||
for _, body := range received {
|
||||
for _, key := range []string{"resolution", "aspect_ratio", "width", "height"} {
|
||||
if _, ok := body[key]; ok {
|
||||
t.Fatalf("generic geometry field %q must not reach OpenAI: %+v", key, body)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestVolcesImageBodySelectsConfiguredSizeFormat(t *testing.T) {
|
||||
body := volcesImageBody(Request{
|
||||
Kind: "images.generations",
|
||||
ModelType: "image_generate",
|
||||
Body: map[string]any{
|
||||
"model": "Seedream-5.0-Pro",
|
||||
"resolution": "2k",
|
||||
"size": "2048x1152",
|
||||
"width": 2048,
|
||||
"height": 1152,
|
||||
"aspect_ratio": "16:9",
|
||||
"platformId": "gateway-internal-platform",
|
||||
},
|
||||
Candidate: store.RuntimeModelCandidate{
|
||||
ProviderModelName: "doubao-seedream-5-0-pro-260628",
|
||||
ModelType: "image_generate",
|
||||
Capabilities: map[string]any{
|
||||
"image_generate": map[string]any{
|
||||
"size_param_format": "resolution",
|
||||
},
|
||||
},
|
||||
},
|
||||
})
|
||||
if body["size"] != "2K" {
|
||||
t.Fatalf("resolution-only Volces model should receive canonical K size: %+v", body)
|
||||
}
|
||||
for _, key := range []string{"resolution", "width", "height", "aspect_ratio", "platformId"} {
|
||||
if _, ok := body[key]; ok {
|
||||
t.Fatalf("generic geometry field %q must not reach Volces: %+v", key, body)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,251 @@
|
||||
package clients
|
||||
|
||||
import (
|
||||
"math"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
const (
|
||||
openAIGPTImage2MaxEdge = 3840
|
||||
openAIGPTImage2MinPixels = 655360
|
||||
openAIGPTImage2MaxPixels = 8294400
|
||||
openAIGPTImage2MaxRatio = 3
|
||||
openAIGPTImage2Multiple = 16
|
||||
)
|
||||
|
||||
type providerImageDimensions struct {
|
||||
width int
|
||||
height int
|
||||
}
|
||||
|
||||
func normalizedProviderAspectRatio(value any) string {
|
||||
compact := strings.Map(func(r rune) rune {
|
||||
switch r {
|
||||
case ' ', '\t', '\r', '\n':
|
||||
return -1
|
||||
default:
|
||||
return r
|
||||
}
|
||||
}, strings.TrimSpace(stringFromAny(value)))
|
||||
parts := strings.Split(compact, ":")
|
||||
if len(parts) != 2 {
|
||||
return ""
|
||||
}
|
||||
width, widthErr := strconv.ParseFloat(parts[0], 64)
|
||||
height, heightErr := strconv.ParseFloat(parts[1], 64)
|
||||
if widthErr != nil || heightErr != nil || width <= 0 || height <= 0 {
|
||||
return ""
|
||||
}
|
||||
return parts[0] + ":" + parts[1]
|
||||
}
|
||||
|
||||
func normalizedProviderImageResolution(value any) string {
|
||||
text := strings.ToLower(strings.TrimSpace(stringFromAny(value)))
|
||||
if len(text) < 2 || text[len(text)-1] != 'k' {
|
||||
return ""
|
||||
}
|
||||
magnitude, err := strconv.Atoi(text[:len(text)-1])
|
||||
if err != nil || magnitude <= 0 {
|
||||
return ""
|
||||
}
|
||||
return strconv.Itoa(magnitude) + "K"
|
||||
}
|
||||
|
||||
func normalizeOfficialOpenAIImageSize(body map[string]any, originalBody map[string]any) {
|
||||
if originalBody == nil {
|
||||
originalBody = body
|
||||
}
|
||||
if width, height, ok := providerPixelDimensions(stringFromAny(originalBody["size"])); ok {
|
||||
body["size"] = strconv.Itoa(width) + "x" + strconv.Itoa(height)
|
||||
return
|
||||
}
|
||||
|
||||
model := strings.ToLower(strings.TrimSpace(stringFromAny(body["model"])))
|
||||
resolution := normalizedProviderImageResolution(firstNonEmptyString(
|
||||
originalBody["resolution"],
|
||||
originalBody["size"],
|
||||
body["resolution"],
|
||||
body["size"],
|
||||
))
|
||||
if resolution != "" {
|
||||
width, height, ok := providerBodyPixelDimensions(body)
|
||||
if !ok {
|
||||
width, height, ok = providerResolutionDimensions(
|
||||
resolution,
|
||||
firstNonEmptyString(originalBody["aspect_ratio"], originalBody["aspectRatio"], originalBody["ratio"]),
|
||||
)
|
||||
}
|
||||
if !ok {
|
||||
delete(body, "size")
|
||||
return
|
||||
}
|
||||
switch {
|
||||
case strings.HasPrefix(model, "gpt-image-2"):
|
||||
dimensions := normalizeOpenAIGPTImage2Dimensions(width, height)
|
||||
body["size"] = strconv.Itoa(dimensions.width) + "x" + strconv.Itoa(dimensions.height)
|
||||
case strings.HasPrefix(model, "gpt-image-"):
|
||||
body["size"] = officialOpenAILegacyImageSize(width, height)
|
||||
default:
|
||||
body["size"] = strconv.Itoa(width) + "x" + strconv.Itoa(height)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
if !strings.HasPrefix(model, "gpt-image-") {
|
||||
return
|
||||
}
|
||||
size := strings.ToLower(strings.TrimSpace(stringFromAny(body["size"])))
|
||||
if size == "auto" {
|
||||
body["size"] = "auto"
|
||||
return
|
||||
}
|
||||
if strings.HasPrefix(model, "gpt-image-2") {
|
||||
if width, height, ok := providerPixelDimensions(size); ok {
|
||||
body["size"] = strconv.Itoa(width) + "x" + strconv.Itoa(height)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
switch size {
|
||||
case "1024x1024", "1536x1024", "1024x1536":
|
||||
body["size"] = size
|
||||
return
|
||||
}
|
||||
width, height, ok := requestedProviderDimensions(body)
|
||||
if !ok {
|
||||
body["size"] = "auto"
|
||||
return
|
||||
}
|
||||
body["size"] = officialOpenAILegacyImageSize(width, height)
|
||||
}
|
||||
|
||||
func officialOpenAILegacyImageSize(width int, height int) string {
|
||||
switch {
|
||||
case width > height:
|
||||
return "1536x1024"
|
||||
case height > width:
|
||||
return "1024x1536"
|
||||
default:
|
||||
return "1024x1024"
|
||||
}
|
||||
}
|
||||
|
||||
func providerBodyPixelDimensions(body map[string]any) (int, int, bool) {
|
||||
if width, height, ok := providerPixelDimensions(stringFromAny(body["size"])); ok {
|
||||
return width, height, true
|
||||
}
|
||||
width := intFromAny(body["width"])
|
||||
height := intFromAny(body["height"])
|
||||
return width, height, width > 0 && height > 0
|
||||
}
|
||||
|
||||
func providerResolutionDimensions(resolution string, aspectRatio any) (int, int, bool) {
|
||||
resolution = normalizedProviderImageResolution(resolution)
|
||||
if resolution == "" {
|
||||
return 0, 0, false
|
||||
}
|
||||
magnitude, err := strconv.Atoi(strings.TrimSuffix(resolution, "K"))
|
||||
if err != nil || magnitude <= 0 {
|
||||
return 0, 0, false
|
||||
}
|
||||
baseSide := magnitude * 1024
|
||||
ratio := providerAspectRatioNumber(aspectRatio)
|
||||
width, height := baseSide, baseSide
|
||||
if ratio > 1 {
|
||||
height = max(1, int(math.Round(float64(baseSide)/ratio)))
|
||||
} else if ratio > 0 && ratio < 1 {
|
||||
width = max(1, int(math.Round(float64(baseSide)*ratio)))
|
||||
}
|
||||
return width, height, true
|
||||
}
|
||||
|
||||
func requestedProviderDimensions(body map[string]any) (int, int, bool) {
|
||||
if width, height, ok := providerPixelDimensions(stringFromAny(body["size"])); ok {
|
||||
return width, height, true
|
||||
}
|
||||
width := intFromAny(body["width"])
|
||||
height := intFromAny(body["height"])
|
||||
if width > 0 && height > 0 {
|
||||
return width, height, true
|
||||
}
|
||||
ratio := providerAspectRatioNumber(firstNonEmptyString(body["aspect_ratio"], body["aspectRatio"], body["ratio"]))
|
||||
if ratio > 1 {
|
||||
return 2, 1, true
|
||||
}
|
||||
if ratio > 0 && ratio < 1 {
|
||||
return 1, 2, true
|
||||
}
|
||||
if ratio == 1 {
|
||||
return 1, 1, true
|
||||
}
|
||||
return 0, 0, false
|
||||
}
|
||||
|
||||
func providerPixelDimensions(value string) (int, int, bool) {
|
||||
parts := strings.Split(strings.ToLower(strings.TrimSpace(value)), "x")
|
||||
if len(parts) != 2 {
|
||||
return 0, 0, false
|
||||
}
|
||||
width, widthErr := strconv.Atoi(strings.TrimSpace(parts[0]))
|
||||
height, heightErr := strconv.Atoi(strings.TrimSpace(parts[1]))
|
||||
if widthErr != nil || heightErr != nil || width <= 0 || height <= 0 {
|
||||
return 0, 0, false
|
||||
}
|
||||
return width, height, true
|
||||
}
|
||||
|
||||
func providerAspectRatioNumber(value any) float64 {
|
||||
normalized := normalizedProviderAspectRatio(value)
|
||||
if normalized == "" {
|
||||
return 0
|
||||
}
|
||||
parts := strings.Split(normalized, ":")
|
||||
width, _ := strconv.ParseFloat(parts[0], 64)
|
||||
height, _ := strconv.ParseFloat(parts[1], 64)
|
||||
return width / height
|
||||
}
|
||||
|
||||
func normalizeOpenAIGPTImage2Dimensions(width int, height int) providerImageDimensions {
|
||||
if width <= 0 || height <= 0 {
|
||||
return providerImageDimensions{}
|
||||
}
|
||||
rawRatio := float64(width) / float64(height)
|
||||
ratio := rawRatio
|
||||
if ratio < 1 {
|
||||
ratio = 1 / ratio
|
||||
}
|
||||
ratio = min(ratio, float64(openAIGPTImage2MaxRatio))
|
||||
landscape := rawRatio >= 1
|
||||
longEdge := min(max(width, height), openAIGPTImage2MaxEdge)
|
||||
dimensions := buildOpenAIGPTImage2Dimensions(longEdge, ratio, landscape, false)
|
||||
for dimensions.width*dimensions.height > openAIGPTImage2MaxPixels {
|
||||
longEdge = max(openAIGPTImage2Multiple, max(dimensions.width, dimensions.height)-openAIGPTImage2Multiple)
|
||||
dimensions = buildOpenAIGPTImage2Dimensions(longEdge, ratio, landscape, false)
|
||||
}
|
||||
for dimensions.width*dimensions.height < openAIGPTImage2MinPixels && max(dimensions.width, dimensions.height) < openAIGPTImage2MaxEdge {
|
||||
longEdge = min(openAIGPTImage2MaxEdge, max(dimensions.width, dimensions.height)+openAIGPTImage2Multiple)
|
||||
dimensions = buildOpenAIGPTImage2Dimensions(longEdge, ratio, landscape, true)
|
||||
}
|
||||
return dimensions
|
||||
}
|
||||
|
||||
func buildOpenAIGPTImage2Dimensions(longEdge int, ratio float64, landscape bool, roundUp bool) providerImageDimensions {
|
||||
normalizedLongEdge := roundOpenAIGPTImage2Multiple(float64(longEdge), roundUp)
|
||||
normalizedShortEdge := roundOpenAIGPTImage2Multiple(float64(normalizedLongEdge)/ratio, roundUp)
|
||||
if float64(normalizedLongEdge)/float64(normalizedShortEdge) > openAIGPTImage2MaxRatio {
|
||||
normalizedShortEdge += openAIGPTImage2Multiple
|
||||
}
|
||||
if landscape {
|
||||
return providerImageDimensions{width: normalizedLongEdge, height: normalizedShortEdge}
|
||||
}
|
||||
return providerImageDimensions{width: normalizedShortEdge, height: normalizedLongEdge}
|
||||
}
|
||||
|
||||
func roundOpenAIGPTImage2Multiple(value float64, roundUp bool) int {
|
||||
rounded := math.Round(value / openAIGPTImage2Multiple)
|
||||
if roundUp {
|
||||
rounded = math.Ceil(value / openAIGPTImage2Multiple)
|
||||
}
|
||||
return max(openAIGPTImage2Multiple, int(rounded)*openAIGPTImage2Multiple)
|
||||
}
|
||||
@@ -3,8 +3,14 @@ package clients
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"mime/multipart"
|
||||
"net/http"
|
||||
"net/textproto"
|
||||
"net/url"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
@@ -42,6 +48,11 @@ func (c OpenAIClient) Run(ctx context.Context, request Request) (Response, error
|
||||
}
|
||||
if endpointKind == "chat.completions" {
|
||||
body = NormalizeChatCompletionRequestBody(body)
|
||||
normalizedBody, normalizeErr := normalizeOpenAIChatAudioParts(body, request.Candidate)
|
||||
if normalizeErr != nil {
|
||||
return Response{}, normalizeErr
|
||||
}
|
||||
body = normalizedBody
|
||||
applyOpenAIChatReasoningParams(body, request.Candidate)
|
||||
body = FilterOpenAIChatRequestBody(body)
|
||||
} else if request.Kind == "responses" {
|
||||
@@ -59,15 +70,19 @@ func (c OpenAIClient) Run(ctx context.Context, request Request) (Response, error
|
||||
}
|
||||
}
|
||||
body["model"] = upstreamModelName(request.Candidate)
|
||||
normalizeOpenAIImageRequestBody(endpointKind, body, request.OriginalBody)
|
||||
stream := openAIEndpointSupportsStream(endpointKind) && (request.Stream || boolValue(body, "stream"))
|
||||
ensureOpenAIStreamUsage(body, endpointKind, stream)
|
||||
raw, _ := json.Marshal(body)
|
||||
raw, contentType, err := openAIRequestPayload(endpointKind, body, request.Candidate)
|
||||
if err != nil {
|
||||
return Response{}, err
|
||||
}
|
||||
upstreamEndpoint := joinURL(openAIBaseURL(endpointKind, request.Candidate), endpoint)
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, upstreamEndpoint, bytes.NewReader(raw))
|
||||
if err != nil {
|
||||
return Response{}, err
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Content-Type", contentType)
|
||||
req.Header.Set("Authorization", "Bearer "+apiKey)
|
||||
responseStartedAt := time.Now()
|
||||
if err := notifySubmissionStarted(request); err != nil {
|
||||
@@ -162,6 +177,245 @@ func (c OpenAIClient) Run(ctx context.Context, request Request) (Response, error
|
||||
}, nil
|
||||
}
|
||||
|
||||
func normalizeOpenAIImageRequestBody(endpointKind string, body map[string]any, originalBody map[string]any) {
|
||||
if endpointKind != "images.generations" && endpointKind != "images.edits" {
|
||||
return
|
||||
}
|
||||
normalizeOfficialOpenAIImageSize(body, originalBody)
|
||||
// OpenAI image endpoints express output geometry through size. The Gateway
|
||||
// keeps the generic fields for capability validation and billing, but they
|
||||
// are not valid OpenAI wire parameters.
|
||||
for _, key := range []string{
|
||||
"aspect_ratio",
|
||||
"aspectRatio",
|
||||
"resolution",
|
||||
"width",
|
||||
"height",
|
||||
"platform_id",
|
||||
"platformId",
|
||||
"platform_model_id",
|
||||
"platformModelId",
|
||||
} {
|
||||
delete(body, key)
|
||||
}
|
||||
}
|
||||
|
||||
func openAIRequestPayload(endpointKind string, body map[string]any, candidate store.RuntimeModelCandidate) ([]byte, string, error) {
|
||||
if endpointKind != "images.edits" {
|
||||
raw, err := json.Marshal(body)
|
||||
return raw, "application/json", err
|
||||
}
|
||||
var payload bytes.Buffer
|
||||
writer := multipart.NewWriter(&payload)
|
||||
imageFieldName := openAIImageEditFieldName(candidate)
|
||||
images := openAIImageEditValues(firstPresent(body["images"], body["image"]))
|
||||
if len(images) == 0 {
|
||||
return nil, "", &ClientError{
|
||||
Code: "invalid_parameter",
|
||||
Message: "image is required",
|
||||
Param: "image",
|
||||
StatusCode: http.StatusBadRequest,
|
||||
Retryable: false,
|
||||
}
|
||||
}
|
||||
for index, value := range images {
|
||||
contentType, image, err := openAIImageEditPayload(value)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
if err := writeOpenAIImageEditFile(writer, imageFieldName, fmt.Sprintf("image-%d%s", index+1, openAIImageFileExtension(contentType)), contentType, image); err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
}
|
||||
if mask := firstPresent(body["mask"], body["mask_image"], body["maskImage"]); mask != nil {
|
||||
contentType, image, err := openAIImageEditPayload(mask)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
if err := writeOpenAIImageEditFile(writer, "mask", "mask"+openAIImageFileExtension(contentType), contentType, image); err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
}
|
||||
for key, value := range body {
|
||||
switch key {
|
||||
case "image", "images", "mask", "mask_image", "maskImage":
|
||||
continue
|
||||
}
|
||||
if strings.HasPrefix(key, "_") || value == nil {
|
||||
continue
|
||||
}
|
||||
fieldValue, err := openAIFormFieldValue(value)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
if fieldValue == "" {
|
||||
continue
|
||||
}
|
||||
if err := writer.WriteField(key, fieldValue); err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
}
|
||||
if err := writer.Close(); err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
return payload.Bytes(), writer.FormDataContentType(), nil
|
||||
}
|
||||
|
||||
func openAIImageEditFieldName(candidate store.RuntimeModelCandidate) string {
|
||||
if configured := strings.TrimSpace(stringFromAny(candidate.PlatformConfig["imageEditMultipartFieldName"])); configured == "image" || configured == "image[]" {
|
||||
return configured
|
||||
}
|
||||
parsed, err := url.Parse(strings.TrimSpace(candidate.BaseURL))
|
||||
if err == nil && strings.EqualFold(parsed.Hostname(), "api.openai.com") {
|
||||
return "image[]"
|
||||
}
|
||||
return "image"
|
||||
}
|
||||
|
||||
func openAIImageEditValues(value any) []any {
|
||||
switch typed := value.(type) {
|
||||
case []any:
|
||||
return typed
|
||||
case []string:
|
||||
out := make([]any, 0, len(typed))
|
||||
for _, item := range typed {
|
||||
out = append(out, item)
|
||||
}
|
||||
return out
|
||||
case nil:
|
||||
return nil
|
||||
default:
|
||||
return []any{typed}
|
||||
}
|
||||
}
|
||||
|
||||
func openAIImageEditPayload(value any) (string, []byte, error) {
|
||||
switch typed := value.(type) {
|
||||
case map[string]any:
|
||||
for _, key := range []string{"data", "b64_json", "base64", "url"} {
|
||||
if nested := typed[key]; nested != nil {
|
||||
return openAIImageEditPayload(nested)
|
||||
}
|
||||
}
|
||||
case string:
|
||||
raw := strings.TrimSpace(typed)
|
||||
if raw == "" {
|
||||
break
|
||||
}
|
||||
contentType := ""
|
||||
encoded := raw
|
||||
if strings.HasPrefix(strings.ToLower(raw), "data:") {
|
||||
prefix, payload, ok := strings.Cut(raw, ",")
|
||||
if !ok || !strings.Contains(strings.ToLower(prefix), ";base64") {
|
||||
break
|
||||
}
|
||||
contentType = strings.TrimSpace(strings.Split(strings.TrimPrefix(prefix, "data:"), ";")[0])
|
||||
encoded = payload
|
||||
} else if strings.Contains(raw, "://") {
|
||||
return "", nil, &ClientError{
|
||||
Code: "invalid_parameter",
|
||||
Message: "OpenAI image edit input must be hydrated before multipart submission",
|
||||
Param: "image",
|
||||
StatusCode: http.StatusBadRequest,
|
||||
Retryable: false,
|
||||
}
|
||||
}
|
||||
image, err := decodeOpenAIImageEditBase64(encoded)
|
||||
if err != nil {
|
||||
break
|
||||
}
|
||||
if contentType == "" {
|
||||
contentType = strings.TrimSpace(strings.Split(http.DetectContentType(image), ";")[0])
|
||||
}
|
||||
return contentType, image, nil
|
||||
case []byte:
|
||||
if len(typed) > 0 {
|
||||
return strings.TrimSpace(strings.Split(http.DetectContentType(typed), ";")[0]), typed, nil
|
||||
}
|
||||
}
|
||||
return "", nil, &ClientError{
|
||||
Code: "invalid_parameter",
|
||||
Message: "image must be a base64 or data URL payload",
|
||||
Param: "image",
|
||||
StatusCode: http.StatusBadRequest,
|
||||
Retryable: false,
|
||||
}
|
||||
}
|
||||
|
||||
func decodeOpenAIImageEditBase64(value string) ([]byte, error) {
|
||||
normalized := strings.Map(func(char rune) rune {
|
||||
switch char {
|
||||
case '\n', '\r', '\t', ' ':
|
||||
return -1
|
||||
default:
|
||||
return char
|
||||
}
|
||||
}, value)
|
||||
var lastErr error
|
||||
for _, encoding := range []*base64.Encoding{
|
||||
base64.StdEncoding,
|
||||
base64.RawStdEncoding,
|
||||
base64.URLEncoding,
|
||||
base64.RawURLEncoding,
|
||||
} {
|
||||
payload, err := encoding.DecodeString(normalized)
|
||||
if err == nil && len(payload) > 0 {
|
||||
return payload, nil
|
||||
}
|
||||
lastErr = err
|
||||
}
|
||||
return nil, lastErr
|
||||
}
|
||||
|
||||
func writeOpenAIImageEditFile(writer *multipart.Writer, fieldName string, fileName string, contentType string, payload []byte) error {
|
||||
header := make(textproto.MIMEHeader)
|
||||
header.Set("Content-Disposition", fmt.Sprintf(`form-data; name="%s"; filename="%s"`, fieldName, escapeMultipartFilename(fileName)))
|
||||
if contentType != "" {
|
||||
header.Set("Content-Type", contentType)
|
||||
}
|
||||
part, err := writer.CreatePart(header)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
_, err = io.Copy(part, bytes.NewReader(payload))
|
||||
return err
|
||||
}
|
||||
|
||||
func openAIImageFileExtension(contentType string) string {
|
||||
switch strings.ToLower(strings.TrimSpace(strings.Split(contentType, ";")[0])) {
|
||||
case "image/jpeg":
|
||||
return ".jpg"
|
||||
case "image/webp":
|
||||
return ".webp"
|
||||
case "image/gif":
|
||||
return ".gif"
|
||||
default:
|
||||
return ".png"
|
||||
}
|
||||
}
|
||||
|
||||
func openAIFormFieldValue(value any) (string, error) {
|
||||
switch typed := value.(type) {
|
||||
case string:
|
||||
return typed, nil
|
||||
case bool:
|
||||
return fmt.Sprintf("%t", typed), nil
|
||||
case float64:
|
||||
return fmt.Sprintf("%v", typed), nil
|
||||
case float32:
|
||||
return fmt.Sprintf("%v", typed), nil
|
||||
case int:
|
||||
return fmt.Sprintf("%d", typed), nil
|
||||
case int64:
|
||||
return fmt.Sprintf("%d", typed), nil
|
||||
case json.Number:
|
||||
return typed.String(), nil
|
||||
default:
|
||||
encoded, err := json.Marshal(value)
|
||||
return string(encoded), err
|
||||
}
|
||||
}
|
||||
|
||||
func openAIWireStreamDelta(next StreamDelta, protocol string, response *http.Response) StreamDelta {
|
||||
if next == nil {
|
||||
return nil
|
||||
|
||||
@@ -0,0 +1,169 @@
|
||||
package clients
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
)
|
||||
|
||||
var openAIChatAudioFormatByMIME = map[string]string{
|
||||
"audio/aac": "aac",
|
||||
"audio/flac": "flac",
|
||||
"audio/m4a": "m4a",
|
||||
"audio/mp4": "m4a",
|
||||
"audio/mpeg": "mp3",
|
||||
"audio/ogg": "ogg",
|
||||
"audio/opus": "opus",
|
||||
"audio/wav": "wav",
|
||||
"audio/webm": "webm",
|
||||
"audio/x-m4a": "m4a",
|
||||
"audio/x-wav": "wav",
|
||||
}
|
||||
|
||||
var openAIChatAudioFormatByExtension = map[string]string{
|
||||
"aac": "aac",
|
||||
"flac": "flac",
|
||||
"m4a": "m4a",
|
||||
"mp3": "mp3",
|
||||
"mp4": "m4a",
|
||||
"oga": "ogg",
|
||||
"ogg": "ogg",
|
||||
"opus": "opus",
|
||||
"wav": "wav",
|
||||
"webm": "webm",
|
||||
}
|
||||
|
||||
func normalizeOpenAIChatAudioParts(body map[string]any, candidate store.RuntimeModelCandidate) (map[string]any, error) {
|
||||
if !isAliyunBailianOpenAI(candidate) {
|
||||
return body, nil
|
||||
}
|
||||
messages, ok := body["messages"].([]any)
|
||||
if !ok {
|
||||
return body, nil
|
||||
}
|
||||
|
||||
normalizedMessages := make([]any, 0, len(messages))
|
||||
for _, rawMessage := range messages {
|
||||
message, ok := rawMessage.(map[string]any)
|
||||
if !ok {
|
||||
normalizedMessages = append(normalizedMessages, rawMessage)
|
||||
continue
|
||||
}
|
||||
content, ok := message["content"].([]any)
|
||||
if !ok {
|
||||
normalizedMessages = append(normalizedMessages, rawMessage)
|
||||
continue
|
||||
}
|
||||
|
||||
normalizedContent := make([]any, 0, len(content))
|
||||
for _, rawPart := range content {
|
||||
part, ok := rawPart.(map[string]any)
|
||||
if !ok {
|
||||
normalizedContent = append(normalizedContent, rawPart)
|
||||
continue
|
||||
}
|
||||
normalizedPart, err := normalizeAliyunOpenAIChatAudioPart(part)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
normalizedContent = append(normalizedContent, normalizedPart)
|
||||
}
|
||||
copiedMessage := cloneMapAny(message)
|
||||
copiedMessage["content"] = normalizedContent
|
||||
normalizedMessages = append(normalizedMessages, copiedMessage)
|
||||
}
|
||||
|
||||
out := cloneMapAny(body)
|
||||
out["messages"] = normalizedMessages
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func normalizeAliyunOpenAIChatAudioPart(part map[string]any) (map[string]any, error) {
|
||||
switch strings.TrimSpace(stringFromAny(part["type"])) {
|
||||
case "audio_url":
|
||||
audioURL, _ := part["audio_url"].(map[string]any)
|
||||
data := strings.TrimSpace(stringFromAny(audioURL["url"]))
|
||||
format := resolveOpenAIChatAudioFormat(data, audioURL["format"], audioURL["mime_type"], audioURL["mimeType"])
|
||||
if data == "" || format == "" {
|
||||
return nil, invalidOpenAIChatAudioError()
|
||||
}
|
||||
return map[string]any{
|
||||
"type": "input_audio",
|
||||
"input_audio": map[string]any{
|
||||
"data": data,
|
||||
"format": format,
|
||||
},
|
||||
}, nil
|
||||
case "input_audio":
|
||||
inputAudio, _ := part["input_audio"].(map[string]any)
|
||||
data := firstNonEmptyString(inputAudio["data"], inputAudio["url"])
|
||||
format := resolveOpenAIChatAudioFormat(data, inputAudio["format"], inputAudio["mime_type"], inputAudio["mimeType"])
|
||||
if data == "" || format == "" {
|
||||
return nil, invalidOpenAIChatAudioError()
|
||||
}
|
||||
copiedInput := cloneMapAny(inputAudio)
|
||||
copiedInput["data"] = data
|
||||
copiedInput["format"] = format
|
||||
delete(copiedInput, "url")
|
||||
copiedPart := cloneMapAny(part)
|
||||
copiedPart["input_audio"] = copiedInput
|
||||
return copiedPart, nil
|
||||
default:
|
||||
return part, nil
|
||||
}
|
||||
}
|
||||
|
||||
func resolveOpenAIChatAudioFormat(source string, values ...any) string {
|
||||
if explicit := strings.ToLower(strings.TrimSpace(firstNonEmptyString(values...))); explicit != "" {
|
||||
if normalized := openAIChatAudioFormatByExtension[explicit]; normalized != "" {
|
||||
return normalized
|
||||
}
|
||||
if normalized := openAIChatAudioFormatByMIME[strings.TrimSpace(strings.Split(explicit, ";")[0])]; normalized != "" {
|
||||
return normalized
|
||||
}
|
||||
return explicit
|
||||
}
|
||||
if dataMIME := openAIChatAudioDataMIME(source); dataMIME != "" {
|
||||
if normalized := openAIChatAudioFormatByMIME[dataMIME]; normalized != "" {
|
||||
return normalized
|
||||
}
|
||||
}
|
||||
if extension := openAIChatAudioSourceExtension(source); extension != "" {
|
||||
return openAIChatAudioFormatByExtension[extension]
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func openAIChatAudioDataMIME(source string) string {
|
||||
if !strings.HasPrefix(strings.ToLower(source), "data:") {
|
||||
return ""
|
||||
}
|
||||
header := strings.TrimPrefix(source[:strings.Index(source+",", ",")], "data:")
|
||||
return strings.ToLower(strings.TrimSpace(strings.Split(header, ";")[0]))
|
||||
}
|
||||
|
||||
func openAIChatAudioSourceExtension(source string) string {
|
||||
if source == "" || strings.HasPrefix(strings.ToLower(source), "data:") {
|
||||
return ""
|
||||
}
|
||||
if parsed, err := url.Parse(source); err == nil && parsed.Path != "" {
|
||||
source = parsed.Path
|
||||
}
|
||||
source = strings.TrimSuffix(strings.Split(strings.Split(source, "?")[0], "#")[0], "/")
|
||||
if index := strings.LastIndex(source, "."); index >= 0 && index+1 < len(source) {
|
||||
return strings.ToLower(source[index+1:])
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func invalidOpenAIChatAudioError() error {
|
||||
return &ClientError{
|
||||
Code: "invalid_parameter",
|
||||
Message: "input_audio requires data and a resolvable format",
|
||||
Param: "messages.content.input_audio",
|
||||
StatusCode: http.StatusBadRequest,
|
||||
Retryable: false,
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -357,16 +357,33 @@ func volcesImageBody(request Request) map[string]any {
|
||||
body["seed"] = -1
|
||||
}
|
||||
}
|
||||
if resolution := strings.TrimSpace(stringFromAny(body["resolution"])); resolution != "" {
|
||||
resolution := normalizedProviderImageResolution(body["resolution"])
|
||||
size := widthHeightSize(body)
|
||||
if volcesImageUsesResolutionSize(request) && resolution != "" {
|
||||
body["size"] = resolution
|
||||
} else if size != "" {
|
||||
body["size"] = size
|
||||
} else if resolution != "" {
|
||||
body["size"] = resolution
|
||||
}
|
||||
if size := widthHeightSize(body); size != "" {
|
||||
body["size"] = size
|
||||
for _, key := range []string{"aspect_ratio", "aspectRatio", "ratio", "resolution", "width", "height"} {
|
||||
delete(body, key)
|
||||
}
|
||||
normalizeVolcesSequentialImageGeneration(body, request)
|
||||
return body
|
||||
}
|
||||
|
||||
func volcesImageUsesResolutionSize(request Request) bool {
|
||||
modelType := firstNonEmpty(request.ModelType, request.Candidate.ModelType)
|
||||
if capability, ok := request.Candidate.Capabilities[modelType].(map[string]any); ok {
|
||||
if strings.EqualFold(strings.TrimSpace(stringFromAny(capability["size_param_format"])), "resolution") {
|
||||
return true
|
||||
}
|
||||
}
|
||||
model := strings.ToLower(strings.TrimSpace(upstreamModelName(request.Candidate)))
|
||||
return strings.HasPrefix(model, "doubao-seedream-5-0")
|
||||
}
|
||||
|
||||
func volcesVideoBody(request Request) map[string]any {
|
||||
body := cleanProviderBody(request.Body)
|
||||
body["model"] = upstreamModelName(request.Candidate)
|
||||
@@ -397,6 +414,10 @@ func cleanProviderBody(body map[string]any) map[string]any {
|
||||
"_gateway_target_protocol",
|
||||
"_compat_provider",
|
||||
"_kling_compat_version",
|
||||
"platform_id",
|
||||
"platformId",
|
||||
"platform_model_id",
|
||||
"platformModelId",
|
||||
} {
|
||||
delete(out, key)
|
||||
}
|
||||
@@ -1100,9 +1121,11 @@ func volcesTaskErrorMessage(result map[string]any) string {
|
||||
}
|
||||
|
||||
func volcesVideoSuccessResult(request Request, upstreamTaskID string, raw map[string]any) map[string]any {
|
||||
result := cloneMapAny(raw)
|
||||
if result == nil {
|
||||
result = map[string]any{}
|
||||
result := map[string]any{}
|
||||
for _, key := range []string{"seed", "resolution", "ratio", "duration", "frames", "framespersecond"} {
|
||||
if raw[key] != nil {
|
||||
result[key] = raw[key]
|
||||
}
|
||||
}
|
||||
content, _ := raw["content"].(map[string]any)
|
||||
videoURL := strings.TrimSpace(stringFromAny(content["video_url"]))
|
||||
@@ -1112,18 +1135,18 @@ func volcesVideoSuccessResult(request Request, upstreamTaskID string, raw map[st
|
||||
}
|
||||
data := []any{}
|
||||
if videoURL != "" {
|
||||
data = append(data, map[string]any{"url": videoURL, "type": "video"})
|
||||
item := map[string]any{"url": videoURL, "type": "video"}
|
||||
if lastFrameURL := strings.TrimSpace(stringFromAny(content["last_frame_url"])); lastFrameURL != "" {
|
||||
item["last_frame_url"] = lastFrameURL
|
||||
}
|
||||
data = append(data, item)
|
||||
}
|
||||
result["id"] = firstNonEmpty(stringFromAny(raw["id"]), upstreamTaskID)
|
||||
if strings.TrimSpace(stringFromAny(result["model"])) == "" {
|
||||
result["model"] = upstreamModelName(request.Candidate)
|
||||
}
|
||||
result["model"] = firstNonEmpty(stringFromAny(raw["model"]), upstreamModelName(request.Candidate))
|
||||
result["status"] = "succeeded"
|
||||
result["object"] = "video.generation"
|
||||
result["created"] = created
|
||||
result["upstream_task_id"] = upstreamTaskID
|
||||
result["data"] = data
|
||||
result["raw"] = cloneMapAny(raw)
|
||||
return result
|
||||
}
|
||||
|
||||
|
||||
@@ -79,7 +79,7 @@ func TestVolcesClientSupportsDeyunEnvelope(t *testing.T) {
|
||||
}
|
||||
data, _ := response.Result["data"].([]any)
|
||||
item, _ := data[0].(map[string]any)
|
||||
if response.Result["upstream_task_id"] != "deyun-task-1" || item["url"] != "https://example.com/deyun.mp4" {
|
||||
if response.Result["upstream_task_id"] != nil || item["url"] != "https://example.com/deyun.mp4" {
|
||||
t.Fatalf("unexpected response: %+v", response.Result)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -46,6 +46,46 @@ func TestDecodeHTTPResponseForProtocolCapturesOfficialErrorWire(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestDecodeHTTPResponseForProtocolAllowsLargeImagePayload(t *testing.T) {
|
||||
payload := `{"data":"` + strings.Repeat("a", int(defaultMaxJSONResponseBytes)) + `"}`
|
||||
response := &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Status: "200 OK",
|
||||
Header: http.Header{"Content-Type": {"application/json"}},
|
||||
Body: io.NopCloser(strings.NewReader(payload)),
|
||||
}
|
||||
|
||||
result, wire, err := decodeHTTPResponseForProtocol(response, ProtocolGeminiGenerateContent)
|
||||
if err != nil {
|
||||
t.Fatalf("large Gemini image response failed: %v", err)
|
||||
}
|
||||
if got, _ := result["data"].(string); len(got) != int(defaultMaxJSONResponseBytes) {
|
||||
t.Fatalf("large Gemini image payload length = %d", len(got))
|
||||
}
|
||||
if wire == nil || len(wire.RawJSON) != len(payload) {
|
||||
t.Fatalf("large Gemini wire payload was truncated: %+v", wire)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDecodeHTTPResponseForProtocolRejectsOversizedDefaultPayload(t *testing.T) {
|
||||
payload := `{"data":"` + strings.Repeat("a", int(defaultMaxJSONResponseBytes)) + `"}`
|
||||
response := &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Status: "200 OK",
|
||||
Header: http.Header{"Content-Type": {"application/json"}},
|
||||
Body: io.NopCloser(strings.NewReader(payload)),
|
||||
}
|
||||
|
||||
_, _, err := decodeHTTPResponseForProtocol(response, ProtocolOpenAIResponses)
|
||||
if err == nil {
|
||||
t.Fatal("expected oversized default response to fail")
|
||||
}
|
||||
clientErr, ok := err.(*ClientError)
|
||||
if !ok || clientErr.Code != "response_too_large" {
|
||||
t.Fatalf("unexpected oversized response error: %T %v", err, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGeminiNativeStreamPreservesOfficialEvents(t *testing.T) {
|
||||
var requestedPath string
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
|
||||
@@ -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":
|
||||
|
||||
@@ -26,7 +26,7 @@ func TestLoadIdentitySecretStoreUsesNewEnvironmentNamesAndIgnoresLegacyBusinessV
|
||||
}
|
||||
|
||||
func TestValidateIdentityFileSecretStoreRequiresDirectory(t *testing.T) {
|
||||
cfg := Config{IdentitySecretStore: "file"}
|
||||
cfg := Config{IdentitySecretStore: "file", AsyncWorkerHardLimit: 2048, AsyncWorkerRefreshIntervalSeconds: 5}
|
||||
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "IDENTITY_SECRET_DIR") {
|
||||
t.Fatalf("Validate() error = %v, want missing identity secret directory", err)
|
||||
}
|
||||
@@ -37,7 +37,12 @@ func TestValidateIdentityFileSecretStoreRequiresDirectory(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestValidateIdentityKubernetesSecretStore(t *testing.T) {
|
||||
cfg := Config{IdentitySecretStore: "kubernetes", IdentityKubernetesSecretName: "easyai-gateway-identity"}
|
||||
cfg := Config{
|
||||
IdentitySecretStore: "kubernetes",
|
||||
IdentityKubernetesSecretName: "easyai-gateway-identity",
|
||||
AsyncWorkerHardLimit: 2048,
|
||||
AsyncWorkerRefreshIntervalSeconds: 5,
|
||||
}
|
||||
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "namespace") {
|
||||
t.Fatalf("Validate() error = %v, want missing namespace", err)
|
||||
}
|
||||
@@ -52,8 +57,58 @@ func TestValidateIdentitySecurityEventTiming(t *testing.T) {
|
||||
IdentitySecurityEventHeartbeatIntervalSeconds: 60,
|
||||
IdentitySecurityEventStaleAfterSeconds: 60,
|
||||
IdentitySecurityEventClockSkewSeconds: 60,
|
||||
AsyncWorkerHardLimit: 2048,
|
||||
AsyncWorkerRefreshIntervalSeconds: 5,
|
||||
}
|
||||
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "heartbeat") {
|
||||
t.Fatalf("Validate() error = %v, want invalid stale threshold", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateAsyncWorkerSettings(t *testing.T) {
|
||||
cfg := Config{AsyncWorkerHardLimit: 10001, AsyncWorkerRefreshIntervalSeconds: 5}
|
||||
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "HARD_LIMIT") {
|
||||
t.Fatalf("Validate() error = %v, want invalid hard limit", err)
|
||||
}
|
||||
cfg.AsyncWorkerHardLimit = 2048
|
||||
cfg.AsyncWorkerRefreshIntervalSeconds = 0
|
||||
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "REFRESH_INTERVAL") {
|
||||
t.Fatalf("Validate() error = %v, want invalid refresh interval", err)
|
||||
}
|
||||
t.Setenv("AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT", "not-an-integer")
|
||||
loaded := Load()
|
||||
if err := loaded.Validate(); err == nil || !strings.Contains(err.Error(), "HARD_LIMIT") {
|
||||
t.Fatalf("Validate() error = %v, want invalid non-integer hard limit", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadAsyncQueueWorkerEnabled(t *testing.T) {
|
||||
cfg := Load()
|
||||
if !cfg.AsyncQueueWorkerEnabled {
|
||||
t.Fatal("async queue worker must be enabled by default")
|
||||
}
|
||||
t.Setenv("AI_GATEWAY_ASYNC_QUEUE_WORKER_ENABLED", "false")
|
||||
cfg = Load()
|
||||
if cfg.AsyncQueueWorkerEnabled {
|
||||
t.Fatal("async queue worker remained enabled after explicit disable")
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateTaskHistorySettings(t *testing.T) {
|
||||
cfg := Load()
|
||||
cfg.TaskRetentionDays = 30
|
||||
cfg.TaskAnalysisRetentionDays = 31
|
||||
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "ANALYSIS_RETENTION") {
|
||||
t.Fatalf("Validate() error = %v, want invalid analysis retention", err)
|
||||
}
|
||||
cfg.TaskAnalysisRetentionDays = 7
|
||||
cfg.TaskCleanupIntervalSeconds = 59
|
||||
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "CLEANUP_INTERVAL") {
|
||||
t.Fatalf("Validate() error = %v, want invalid cleanup interval", err)
|
||||
}
|
||||
cfg.TaskCleanupIntervalSeconds = 300
|
||||
cfg.TaskCleanupBatchSize = 5001
|
||||
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "CLEANUP_BATCH") {
|
||||
t.Fatalf("Validate() error = %v, want invalid cleanup batch", err)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,213 @@
|
||||
package httpapi
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
// listAdminTasks godoc
|
||||
// @Summary 管理员列出任务
|
||||
// @Description 跨租户分页查询任务;请求、结果和执行快照中的常见敏感字段默认脱敏。
|
||||
// @Tags admin-tasks
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param q query string false "关键词,匹配任务、用户、租户、模型、平台和 API Key"
|
||||
// @Param tenantId query string false "网关租户 ID"
|
||||
// @Param userId query string false "网关用户 ID"
|
||||
// @Param userGroupId query string false "用户组 ID"
|
||||
// @Param status query string false "任务状态"
|
||||
// @Param platformId query string false "执行平台 ID"
|
||||
// @Param model query string false "调用或实际模型名称"
|
||||
// @Param modelType query string false "模型类型"
|
||||
// @Param runMode query string false "运行模式"
|
||||
// @Param billingStatus query string false "计费状态"
|
||||
// @Param apiKey query string false "API Key ID、名称或前缀"
|
||||
// @Param createdFrom query string false "创建时间起点,支持 RFC3339 或日期格式"
|
||||
// @Param createdTo query string false "创建时间终点,支持 RFC3339 或日期格式"
|
||||
// @Param page query int false "页码" default(1)
|
||||
// @Param pageSize query int false "每页数量" default(50)
|
||||
// @Success 200 {object} AdminTaskListResponse
|
||||
// @Failure 400 {object} ErrorEnvelope
|
||||
// @Failure 401 {object} ErrorEnvelope
|
||||
// @Failure 403 {object} ErrorEnvelope
|
||||
// @Failure 500 {object} ErrorEnvelope
|
||||
// @Router /api/admin/tasks [get]
|
||||
func (s *Server) listAdminTasks(w http.ResponseWriter, r *http.Request) {
|
||||
query := r.URL.Query()
|
||||
page, err := positiveQueryInt(query.Get("page"), 1)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid page")
|
||||
return
|
||||
}
|
||||
pageSize, err := positiveQueryInt(query.Get("pageSize"), 50)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid pageSize")
|
||||
return
|
||||
}
|
||||
createdFrom, err := parseTaskListTime(query.Get("createdFrom"), query.Get("from"))
|
||||
if err != nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid createdFrom")
|
||||
return
|
||||
}
|
||||
createdTo, err := parseTaskListTime(query.Get("createdTo"), query.Get("to"))
|
||||
if err != nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid createdTo")
|
||||
return
|
||||
}
|
||||
for name, value := range map[string]string{
|
||||
"tenantId": query.Get("tenantId"),
|
||||
"userId": query.Get("userId"),
|
||||
"userGroupId": query.Get("userGroupId"),
|
||||
"platformId": query.Get("platformId"),
|
||||
} {
|
||||
if value != "" && !validUUID(value) {
|
||||
writeError(w, http.StatusBadRequest, "invalid "+name)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
result, err := s.store.ListAdminTasks(r.Context(), store.AdminTaskListFilter{
|
||||
Query: firstNonEmpty(query.Get("q"), query.Get("query")),
|
||||
GatewayTenant: query.Get("tenantId"),
|
||||
GatewayUser: query.Get("userId"),
|
||||
UserGroup: query.Get("userGroupId"),
|
||||
Status: query.Get("status"),
|
||||
Platform: query.Get("platformId"),
|
||||
Model: query.Get("model"),
|
||||
ModelType: query.Get("modelType"),
|
||||
RunMode: query.Get("runMode"),
|
||||
BillingStatus: query.Get("billingStatus"),
|
||||
APIKey: query.Get("apiKey"),
|
||||
CreatedFrom: createdFrom,
|
||||
CreatedTo: createdTo,
|
||||
Page: page,
|
||||
PageSize: pageSize,
|
||||
})
|
||||
if err != nil {
|
||||
s.logger.Error("list admin tasks failed", "error", err)
|
||||
writeError(w, http.StatusInternalServerError, "list admin tasks failed")
|
||||
return
|
||||
}
|
||||
for index := range result.Items {
|
||||
result.Items[index] = store.MaskAdminGatewayTask(result.Items[index])
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]any{
|
||||
"items": result.Items,
|
||||
"total": result.Total,
|
||||
"page": result.Page,
|
||||
"pageSize": result.PageSize,
|
||||
})
|
||||
}
|
||||
|
||||
// getAdminTask godoc
|
||||
// @Summary 管理员获取任务详情
|
||||
// @Description 返回任意租户的任务详情;默认脱敏,sensitive=full 时返回完整存储内容。
|
||||
// @Tags admin-tasks
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param taskID path string true "任务 ID"
|
||||
// @Param sensitive query string false "敏感字段模式" Enums(masked,full) default(masked)
|
||||
// @Success 200 {object} store.AdminGatewayTask
|
||||
// @Failure 400 {object} ErrorEnvelope
|
||||
// @Failure 401 {object} ErrorEnvelope
|
||||
// @Failure 403 {object} ErrorEnvelope
|
||||
// @Failure 404 {object} ErrorEnvelope
|
||||
// @Failure 500 {object} ErrorEnvelope
|
||||
// @Router /api/admin/tasks/{taskID} [get]
|
||||
func (s *Server) getAdminTask(w http.ResponseWriter, r *http.Request) {
|
||||
taskID := r.PathValue("taskID")
|
||||
if !validUUID(taskID) {
|
||||
writeError(w, http.StatusBadRequest, "invalid taskID")
|
||||
return
|
||||
}
|
||||
full, ok := parseAdminTaskSensitiveMode(r.URL.Query().Get("sensitive"))
|
||||
if !ok {
|
||||
writeError(w, http.StatusBadRequest, "invalid sensitive mode")
|
||||
return
|
||||
}
|
||||
task, err := s.store.GetAdminTask(r.Context(), taskID)
|
||||
if err != nil {
|
||||
if store.IsNotFound(err) {
|
||||
writeError(w, http.StatusNotFound, "task not found")
|
||||
return
|
||||
}
|
||||
s.logger.Error("get admin task failed", "taskID", taskID, "error", err)
|
||||
writeError(w, http.StatusInternalServerError, "get admin task failed")
|
||||
return
|
||||
}
|
||||
task.GatewayTask, err = s.hydrateTaskResult(r.Context(), task.GatewayTask)
|
||||
if err != nil {
|
||||
writeStoredBinaryResultError(w, err)
|
||||
return
|
||||
}
|
||||
if !full {
|
||||
task = store.MaskAdminGatewayTask(task)
|
||||
}
|
||||
writeJSON(w, http.StatusOK, task)
|
||||
}
|
||||
|
||||
// adminTaskParamPreprocessing godoc
|
||||
// @Summary 管理员获取任务参数预处理日志
|
||||
// @Description 返回任意租户任务的参数改写、校验或模板处理日志;默认脱敏。
|
||||
// @Tags admin-tasks
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param taskID path string true "任务 ID"
|
||||
// @Param sensitive query string false "敏感字段模式" Enums(masked,full) default(masked)
|
||||
// @Success 200 {object} TaskParamPreprocessingLogListResponse
|
||||
// @Failure 400 {object} ErrorEnvelope
|
||||
// @Failure 401 {object} ErrorEnvelope
|
||||
// @Failure 403 {object} ErrorEnvelope
|
||||
// @Failure 404 {object} ErrorEnvelope
|
||||
// @Failure 500 {object} ErrorEnvelope
|
||||
// @Router /api/admin/tasks/{taskID}/param-preprocessing [get]
|
||||
func (s *Server) adminTaskParamPreprocessing(w http.ResponseWriter, r *http.Request) {
|
||||
taskID := r.PathValue("taskID")
|
||||
if !validUUID(taskID) {
|
||||
writeError(w, http.StatusBadRequest, "invalid taskID")
|
||||
return
|
||||
}
|
||||
full, ok := parseAdminTaskSensitiveMode(r.URL.Query().Get("sensitive"))
|
||||
if !ok {
|
||||
writeError(w, http.StatusBadRequest, "invalid sensitive mode")
|
||||
return
|
||||
}
|
||||
if _, err := s.store.GetAdminTask(r.Context(), taskID); err != nil {
|
||||
if store.IsNotFound(err) {
|
||||
writeError(w, http.StatusNotFound, "task not found")
|
||||
return
|
||||
}
|
||||
s.logger.Error("get admin task failed", "taskID", taskID, "error", err)
|
||||
writeError(w, http.StatusInternalServerError, "get admin task failed")
|
||||
return
|
||||
}
|
||||
items, err := s.store.ListTaskParamPreprocessingLogs(r.Context(), taskID)
|
||||
if err != nil {
|
||||
s.logger.Error("list admin task parameter preprocessing logs failed", "taskID", taskID, "error", err)
|
||||
writeError(w, http.StatusInternalServerError, "list admin task parameter preprocessing logs failed")
|
||||
return
|
||||
}
|
||||
if !full {
|
||||
items = store.MaskTaskParamPreprocessingLogs(items)
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]any{"items": items})
|
||||
}
|
||||
|
||||
func parseAdminTaskSensitiveMode(value string) (bool, bool) {
|
||||
switch strings.ToLower(strings.TrimSpace(value)) {
|
||||
case "", "masked":
|
||||
return false, true
|
||||
case "full":
|
||||
return true, true
|
||||
default:
|
||||
return false, false
|
||||
}
|
||||
}
|
||||
|
||||
func validUUID(value string) bool {
|
||||
_, err := uuid.Parse(strings.TrimSpace(value))
|
||||
return err == nil
|
||||
}
|
||||
@@ -0,0 +1,22 @@
|
||||
package httpapi
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestParseAdminTaskSensitiveMode(t *testing.T) {
|
||||
tests := []struct {
|
||||
value string
|
||||
full bool
|
||||
accepted bool
|
||||
}{
|
||||
{value: "", full: false, accepted: true},
|
||||
{value: "masked", full: false, accepted: true},
|
||||
{value: "FULL", full: true, accepted: true},
|
||||
{value: "invalid", full: false, accepted: false},
|
||||
}
|
||||
for _, test := range tests {
|
||||
full, accepted := parseAdminTaskSensitiveMode(test.value)
|
||||
if full != test.full || accepted != test.accepted {
|
||||
t.Fatalf("parseAdminTaskSensitiveMode(%q)=(%v,%v), want (%v,%v)", test.value, full, accepted, test.full, test.accepted)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,321 @@
|
||||
package httpapi
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
"os"
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
)
|
||||
|
||||
func TestAdminTasksCrossTenantQueryAndRedaction(t *testing.T) {
|
||||
databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL"))
|
||||
if databaseURL == "" {
|
||||
t.Skip("set AI_GATEWAY_TEST_DATABASE_URL to run the PostgreSQL integration flow")
|
||||
}
|
||||
|
||||
ctx := context.Background()
|
||||
applyMigration(t, ctx, databaseURL)
|
||||
db, err := store.Connect(ctx, databaseURL)
|
||||
if err != nil {
|
||||
t.Fatalf("connect store: %v", err)
|
||||
}
|
||||
defer db.Close()
|
||||
|
||||
suffix := strconv.FormatInt(time.Now().UnixNano(), 10)
|
||||
createTenant := func(key, name string) store.GatewayTenant {
|
||||
t.Helper()
|
||||
tenant, err := db.CreateTenant(ctx, store.GatewayTenantInput{
|
||||
TenantKey: key + "-" + suffix,
|
||||
Name: name,
|
||||
Source: "gateway",
|
||||
Status: "active",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("create tenant %s: %v", name, err)
|
||||
}
|
||||
return tenant
|
||||
}
|
||||
createUser := func(tenant store.GatewayTenant, prefix string, roles []string) store.GatewayUser {
|
||||
t.Helper()
|
||||
username := prefix + "_" + suffix
|
||||
user, err := db.CreateGatewayUser(ctx, store.GatewayUserInput{
|
||||
UserKey: "gateway:" + username,
|
||||
Username: username,
|
||||
DisplayName: prefix + " display",
|
||||
Email: username + "@example.com",
|
||||
Source: "gateway",
|
||||
GatewayTenantID: tenant.ID,
|
||||
TenantKey: tenant.TenantKey,
|
||||
Roles: roles,
|
||||
Status: "active",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("create gateway user %s: %v", prefix, err)
|
||||
}
|
||||
return user
|
||||
}
|
||||
authUser := func(user store.GatewayUser, apiKeyName string) *auth.User {
|
||||
return &auth.User{
|
||||
ID: user.ID,
|
||||
Username: user.Username,
|
||||
Roles: user.Roles,
|
||||
Source: user.Source,
|
||||
GatewayUserID: user.ID,
|
||||
GatewayTenantID: user.GatewayTenantID,
|
||||
TenantKey: user.TenantKey,
|
||||
APIKeyID: "key-" + apiKeyName + "-" + suffix,
|
||||
APIKeyName: apiKeyName,
|
||||
APIKeyPrefix: "sk-" + apiKeyName,
|
||||
}
|
||||
}
|
||||
createTask := func(user *auth.User, model, secret string) store.GatewayTask {
|
||||
t.Helper()
|
||||
task, err := db.CreateTask(ctx, store.CreateTaskInput{
|
||||
Kind: "chat.completions",
|
||||
Model: model,
|
||||
RunMode: "simulation",
|
||||
Request: map[string]any{
|
||||
"model": model,
|
||||
"diagnostic": map[string]any{
|
||||
"apiKey": secret,
|
||||
},
|
||||
},
|
||||
}, user)
|
||||
if err != nil {
|
||||
t.Fatalf("create task for %s: %v", user.Username, err)
|
||||
}
|
||||
return task
|
||||
}
|
||||
|
||||
tenantA := createTenant("admin-task-a", "Admin Task Tenant A")
|
||||
tenantB := createTenant("admin-task-b", "Admin Task Tenant B")
|
||||
operator := createUser(tenantA, "admin_task_operator", []string{"operator"})
|
||||
userA := createUser(tenantA, "admin_task_user_a", []string{"user"})
|
||||
userB := createUser(tenantB, "admin_task_user_b", []string{"user"})
|
||||
userAAuth := authUser(userA, "alpha-key")
|
||||
userBAuth := authUser(userB, "bravo-key")
|
||||
|
||||
modelA := "admin-task-model-a-" + suffix
|
||||
modelB := "admin-task-model-b-" + suffix
|
||||
taskA := createTask(userAAuth, modelA, "alpha-secret-"+suffix)
|
||||
taskBSecret := "bravo-secret-" + suffix
|
||||
taskB := createTask(userBAuth, modelB, taskBSecret)
|
||||
|
||||
platform, err := db.CreatePlatform(ctx, store.CreatePlatformInput{
|
||||
Provider: "openai",
|
||||
PlatformKey: "admin-task-platform-" + suffix,
|
||||
Name: "Admin Task Platform " + suffix,
|
||||
BaseURL: "https://example.invalid/v1",
|
||||
AuthType: "bearer",
|
||||
Credentials: map[string]any{"apiKey": "platform-secret-" + suffix},
|
||||
Status: "enabled",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("create platform: %v", err)
|
||||
}
|
||||
attemptID, err := db.CreateTaskAttempt(ctx, store.CreateTaskAttemptInput{
|
||||
TaskID: taskB.ID,
|
||||
AttemptNo: 1,
|
||||
PlatformID: platform.ID,
|
||||
QueueKey: "admin-task-integration-" + suffix,
|
||||
Status: "succeeded",
|
||||
Simulated: true,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("create task attempt: %v", err)
|
||||
}
|
||||
if _, err := db.Pool().Exec(ctx, `
|
||||
UPDATE gateway_tasks
|
||||
SET status = 'succeeded',
|
||||
model_type = 'text_generate',
|
||||
requested_model = $2,
|
||||
resolved_model = $2,
|
||||
request_id = $3,
|
||||
billing_status = 'settled',
|
||||
result = $4::jsonb,
|
||||
finished_at = now(),
|
||||
updated_at = now()
|
||||
WHERE id = $1::uuid`,
|
||||
taskB.ID,
|
||||
modelB,
|
||||
"request-admin-task-"+suffix,
|
||||
`{"authorization":"Bearer result-secret","output":"ok"}`,
|
||||
); err != nil {
|
||||
t.Fatalf("enrich task: %v", err)
|
||||
}
|
||||
if _, err := db.Pool().Exec(ctx, `
|
||||
UPDATE gateway_task_attempts
|
||||
SET request_snapshot = $2::jsonb,
|
||||
response_snapshot = $3::jsonb,
|
||||
finished_at = now()
|
||||
WHERE id = $1::uuid`,
|
||||
attemptID,
|
||||
`{"headers":{"authorization":"Bearer attempt-secret"}}`,
|
||||
`{"cookie":"attempt-cookie","status":"ok"}`,
|
||||
); err != nil {
|
||||
t.Fatalf("enrich task attempt: %v", err)
|
||||
}
|
||||
if _, err := db.Pool().Exec(ctx, `
|
||||
INSERT INTO gateway_task_param_preprocessing_logs (
|
||||
task_id, attempt_id, attempt_no, model_type, platform_id, client_id,
|
||||
changed, change_count, actual_input, converted_output, changes, model_snapshot
|
||||
)
|
||||
VALUES (
|
||||
$1::uuid, $2::uuid, 1, 'text_generate', $3::uuid, $4,
|
||||
true, 1, $5::jsonb, $6::jsonb, $7::jsonb, '{}'::jsonb
|
||||
)`,
|
||||
taskB.ID,
|
||||
attemptID,
|
||||
platform.ID,
|
||||
"admin-task-client-"+suffix,
|
||||
`{"password":"input-secret","prompt":"hello"}`,
|
||||
`{"token":"converted-secret","prompt":"hello"}`,
|
||||
`[{"path":"credentials.secret","before":"old-secret","after":"new-secret"}]`,
|
||||
); err != nil {
|
||||
t.Fatalf("create parameter preprocessing log: %v", err)
|
||||
}
|
||||
|
||||
const jwtSecret = "admin-task-integration-jwt-secret"
|
||||
handlerCtx, cancelHandler := context.WithCancel(ctx)
|
||||
defer cancelHandler()
|
||||
server := httptest.NewServer(NewServerWithContext(handlerCtx, config.Config{
|
||||
AppEnv: "test",
|
||||
HTTPAddr: ":0",
|
||||
DatabaseURL: databaseURL,
|
||||
IdentityMode: "hybrid",
|
||||
JWTSecret: jwtSecret,
|
||||
CORSAllowedOrigin: "*",
|
||||
}, db, slog.New(slog.NewTextHandler(io.Discard, nil))))
|
||||
defer server.Close()
|
||||
|
||||
authenticator := auth.New(jwtSecret, "", "")
|
||||
signToken := func(user *auth.User) string {
|
||||
t.Helper()
|
||||
token, err := authenticator.SignJWT(user, time.Hour)
|
||||
if err != nil {
|
||||
t.Fatalf("sign JWT for %s: %v", user.Username, err)
|
||||
}
|
||||
return token
|
||||
}
|
||||
operatorToken := signToken(&auth.User{
|
||||
ID: operator.ID,
|
||||
Username: operator.Username,
|
||||
Roles: operator.Roles,
|
||||
Source: operator.Source,
|
||||
GatewayUserID: operator.ID,
|
||||
GatewayTenantID: operator.GatewayTenantID,
|
||||
TenantKey: operator.TenantKey,
|
||||
})
|
||||
userAToken := signToken(userAAuth)
|
||||
|
||||
doJSON(t, server.URL, http.MethodGet, "/api/admin/tasks", userAToken, nil, http.StatusForbidden, nil)
|
||||
|
||||
query := url.Values{
|
||||
"q": {taskB.ID},
|
||||
"tenantId": {tenantB.ID},
|
||||
"userId": {userB.ID},
|
||||
"status": {"succeeded"},
|
||||
"platformId": {platform.ID},
|
||||
"model": {modelB},
|
||||
"modelType": {"text_generate"},
|
||||
"runMode": {"simulation"},
|
||||
"billingStatus": {"settled"},
|
||||
"apiKey": {"bravo-key"},
|
||||
"page": {"1"},
|
||||
"pageSize": {"10"},
|
||||
}
|
||||
var listResponse struct {
|
||||
Items []store.AdminGatewayTask `json:"items"`
|
||||
Total int `json:"total"`
|
||||
Page int `json:"page"`
|
||||
PageSize int `json:"pageSize"`
|
||||
}
|
||||
doJSON(t, server.URL, http.MethodGet, "/api/admin/tasks?"+query.Encode(), operatorToken, nil, http.StatusOK, &listResponse)
|
||||
if len(listResponse.Items) != 1 || listResponse.Total != 1 || listResponse.Page != 1 || listResponse.PageSize != 10 {
|
||||
t.Fatalf("unexpected filtered admin task list: %+v", listResponse)
|
||||
}
|
||||
listed := listResponse.Items[0]
|
||||
if listed.ID != taskB.ID || listed.AdminContext.User == nil || listed.AdminContext.User.ID != userB.ID ||
|
||||
listed.AdminContext.Tenant == nil || listed.AdminContext.Tenant.ID != tenantB.ID ||
|
||||
listed.AdminContext.LatestPlatform == nil || listed.AdminContext.LatestPlatform.ID != platform.ID {
|
||||
t.Fatalf("admin task list should expose cross-tenant identity and platform summaries: %+v", listed)
|
||||
}
|
||||
if nestedString(listed.Request, "diagnostic", "apiKey") != "***" ||
|
||||
nestedString(listed.Result, "authorization") != "***" ||
|
||||
nestedString(listed.Attempts[0].RequestSnapshot, "headers", "authorization") != "***" {
|
||||
t.Fatalf("admin task list should mask task and attempt secrets: %+v", listed)
|
||||
}
|
||||
|
||||
var maskedDetail store.AdminGatewayTask
|
||||
doJSON(t, server.URL, http.MethodGet, "/api/admin/tasks/"+taskB.ID, operatorToken, nil, http.StatusOK, &maskedDetail)
|
||||
if nestedString(maskedDetail.Request, "diagnostic", "apiKey") != "***" {
|
||||
t.Fatalf("admin task detail should be masked by default: %+v", maskedDetail.Request)
|
||||
}
|
||||
var fullDetail store.AdminGatewayTask
|
||||
doJSON(t, server.URL, http.MethodGet, "/api/admin/tasks/"+taskB.ID+"?sensitive=full", operatorToken, nil, http.StatusOK, &fullDetail)
|
||||
if nestedString(fullDetail.Request, "diagnostic", "apiKey") != taskBSecret {
|
||||
t.Fatalf("explicit full task detail should expose stored content: %+v", fullDetail.Request)
|
||||
}
|
||||
|
||||
var maskedLogs struct {
|
||||
Items []store.TaskParamPreprocessingLog `json:"items"`
|
||||
}
|
||||
doJSON(t, server.URL, http.MethodGet, "/api/admin/tasks/"+taskB.ID+"/param-preprocessing", operatorToken, nil, http.StatusOK, &maskedLogs)
|
||||
if len(maskedLogs.Items) != 1 ||
|
||||
nestedString(maskedLogs.Items[0].ActualInput, "password") != "***" ||
|
||||
nestedString(maskedLogs.Items[0].ConvertedOutput, "token") != "***" {
|
||||
t.Fatalf("admin task preprocessing logs should be masked by default: %+v", maskedLogs.Items)
|
||||
}
|
||||
var fullLogs struct {
|
||||
Items []store.TaskParamPreprocessingLog `json:"items"`
|
||||
}
|
||||
doJSON(t, server.URL, http.MethodGet, "/api/admin/tasks/"+taskB.ID+"/param-preprocessing?sensitive=full", operatorToken, nil, http.StatusOK, &fullLogs)
|
||||
if len(fullLogs.Items) != 1 ||
|
||||
nestedString(fullLogs.Items[0].ActualInput, "password") != "input-secret" ||
|
||||
nestedString(fullLogs.Items[0].ConvertedOutput, "token") != "converted-secret" {
|
||||
t.Fatalf("explicit full preprocessing logs should expose stored content: %+v", fullLogs.Items)
|
||||
}
|
||||
|
||||
var ownTasks struct {
|
||||
Items []store.GatewayTask `json:"items"`
|
||||
}
|
||||
doJSON(t, server.URL, http.MethodGet, "/api/workspace/tasks?pageSize=100", userAToken, nil, http.StatusOK, &ownTasks)
|
||||
if !gatewayTaskListHasID(ownTasks.Items, taskA.ID) || gatewayTaskListHasID(ownTasks.Items, taskB.ID) {
|
||||
t.Fatalf("workspace task list must remain scoped to the current user: %+v", ownTasks.Items)
|
||||
}
|
||||
doJSON(t, server.URL, http.MethodGet, "/api/workspace/tasks/"+taskB.ID, userAToken, nil, http.StatusNotFound, nil)
|
||||
doJSON(t, server.URL, http.MethodGet, "/api/workspace/tasks/"+taskB.ID+"/param-preprocessing", userAToken, nil, http.StatusNotFound, nil)
|
||||
}
|
||||
|
||||
func nestedString(value map[string]any, path ...string) string {
|
||||
var current any = value
|
||||
for _, key := range path {
|
||||
object, ok := current.(map[string]any)
|
||||
if !ok {
|
||||
return ""
|
||||
}
|
||||
current = object[key]
|
||||
}
|
||||
result, _ := current.(string)
|
||||
return result
|
||||
}
|
||||
|
||||
func gatewayTaskListHasID(items []store.GatewayTask, taskID string) bool {
|
||||
for _, item := range items {
|
||||
if item.ID == taskID {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
@@ -11,7 +11,7 @@ import "net/http"
|
||||
// @Security BearerAuth
|
||||
// @Param X-Async header bool false "true 时异步创建任务并返回 202"
|
||||
// @Param input body ImageVectorizeRequest true "图片矢量化请求"
|
||||
// @Success 200 {object} CompatibleResponse
|
||||
// @Success 200 {object} EasyAIGeneratedResponse
|
||||
// @Success 202 {object} TaskAcceptedResponse
|
||||
// @Failure 400 {object} ErrorEnvelope
|
||||
// @Failure 401 {object} ErrorEnvelope
|
||||
@@ -32,7 +32,7 @@ func (s *Server) createImageVectorizeTask() http.Handler {
|
||||
// @Security BearerAuth
|
||||
// @Param X-Async header bool false "true 时异步创建任务并返回 202"
|
||||
// @Param input body VideoUpscaleRequest true "视频超分请求"
|
||||
// @Success 200 {object} CompatibleResponse
|
||||
// @Success 200 {object} EasyAIGeneratedResponse
|
||||
// @Success 202 {object} TaskAcceptedResponse
|
||||
// @Failure 400 {object} ErrorEnvelope
|
||||
// @Failure 401 {object} ErrorEnvelope
|
||||
|
||||
@@ -0,0 +1,479 @@
|
||||
package httpapi
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
)
|
||||
|
||||
func TestAsyncWorkerDynamicConcurrencyAcceptance(t *testing.T) {
|
||||
databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL"))
|
||||
if databaseURL == "" {
|
||||
t.Skip("set AI_GATEWAY_TEST_DATABASE_URL to run dynamic worker acceptance tests")
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 7*time.Minute)
|
||||
defer cancel()
|
||||
applyMigration(t, ctx, databaseURL)
|
||||
|
||||
db, err := store.Connect(ctx, databaseURL)
|
||||
if err != nil {
|
||||
t.Fatalf("connect store: %v", err)
|
||||
}
|
||||
defer db.Close()
|
||||
pool, err := pgxpool.New(ctx, databaseURL)
|
||||
if err != nil {
|
||||
t.Fatalf("connect acceptance pool: %v", err)
|
||||
}
|
||||
defer pool.Close()
|
||||
if _, err := pool.Exec(ctx, `UPDATE integration_platforms SET status = 'disabled' WHERE deleted_at IS NULL`); err != nil {
|
||||
t.Fatalf("isolate acceptance platforms: %v", err)
|
||||
}
|
||||
|
||||
serverCtx, cancelServer := context.WithCancel(ctx)
|
||||
defer cancelServer()
|
||||
server := httptest.NewServer(NewServerWithContext(serverCtx, config.Config{
|
||||
AppEnv: "test",
|
||||
HTTPAddr: ":0",
|
||||
DatabaseURL: databaseURL,
|
||||
IdentityMode: "hybrid",
|
||||
JWTSecret: "test-secret",
|
||||
BillingEngineMode: "observe",
|
||||
CORSAllowedOrigin: "*",
|
||||
AsyncWorkerHardLimit: 256,
|
||||
AsyncWorkerRefreshIntervalSeconds: 1,
|
||||
}, db, slog.New(slog.NewTextHandler(io.Discard, nil))))
|
||||
defer server.Close()
|
||||
|
||||
adminToken := createAsyncAcceptanceAdmin(t, ctx, pool, server.URL)
|
||||
if _, err := pool.Exec(ctx, `
|
||||
UPDATE gateway_user_groups
|
||||
SET rate_limit_policy = '{"rules":[{"metric":"concurrent","limit":256,"leaseTtlSeconds":120}]}'::jsonb
|
||||
WHERE status = 'active'`); err != nil {
|
||||
t.Fatalf("raise acceptance user-group concurrency: %v", err)
|
||||
}
|
||||
suffix := strconv.FormatInt(time.Now().UnixNano(), 10)
|
||||
|
||||
t.Run("超过旧64上限的高并发", func(t *testing.T) {
|
||||
model := "worker-burst-" + suffix
|
||||
platform := createAsyncAcceptancePlatform(t, server.URL, adminToken, "burst-"+suffix, "Burst Simulation", 128)
|
||||
defer updateAsyncAcceptancePlatform(t, server.URL, adminToken, platform, 128, "disabled")
|
||||
modelID := createAsyncAcceptanceModel(t, server.URL, adminToken, platform.ID, model, "video_generate", "override", 96, 120)
|
||||
|
||||
waitForAsyncWorkerMetric(t, server.URL, "easyai_gateway_async_worker_capacity", 96, 15*time.Second)
|
||||
startedAt := time.Now()
|
||||
taskIDs := submitAsyncSimulationTasks(t, server.URL, adminToken, "/api/v1/videos/generations", model, 128, 15*time.Second)
|
||||
peakRunning, peakLeases := waitForAcceptanceTasks(t, ctx, pool, taskIDs, modelID, 60*time.Second)
|
||||
elapsed := time.Since(startedAt)
|
||||
if peakRunning < 80 || peakLeases < 80 {
|
||||
t.Fatalf("running peak=%d lease peak=%d, want both >=80", peakRunning, peakLeases)
|
||||
}
|
||||
if peakLeases > 96 {
|
||||
t.Fatalf("lease peak=%d exceeded model concurrent limit 96", peakLeases)
|
||||
}
|
||||
assertAcceptanceAttempts(t, ctx, pool, taskIDs, len(taskIDs))
|
||||
if elapsed >= 60*time.Second {
|
||||
t.Fatalf("128 tasks completed in %s, want <60s", elapsed)
|
||||
}
|
||||
t.Logf("高并发证据: tasks=128 duration=%s running_peak=%d lease_peak=%d capacity=96", elapsed.Round(time.Millisecond), peakRunning, peakLeases)
|
||||
})
|
||||
|
||||
t.Run("三分钟视频任务在线扩缩容与续租", func(t *testing.T) {
|
||||
model := "worker-video-long-" + suffix
|
||||
platform := createAsyncAcceptancePlatform(t, server.URL, adminToken, "video-long-"+suffix, "Long Video Simulation", 1)
|
||||
modelID := createAsyncAcceptanceModel(t, server.URL, adminToken, platform.ID, model, "video_generate", "inherit", 0, 120)
|
||||
waitForAsyncWorkerMetric(t, server.URL, "easyai_gateway_async_worker_capacity", 1, 15*time.Second)
|
||||
|
||||
longTaskIDs := submitAsyncSimulationTasks(t, server.URL, adminToken, "/api/v1/videos/generations", model, 3, 180*time.Second)
|
||||
waitForTaskCount(t, ctx, pool, longTaskIDs, "running", 1, 10*time.Second)
|
||||
updateStartedAt := time.Now()
|
||||
updateAsyncAcceptancePlatform(t, server.URL, adminToken, platform, 3, "enabled")
|
||||
waitForAsyncWorkerMetric(t, server.URL, "easyai_gateway_async_worker_desired_capacity", 3, 15*time.Second)
|
||||
waitForAsyncWorkerMetric(t, server.URL, "easyai_gateway_async_worker_capacity", 3, 15*time.Second)
|
||||
waitForActiveLeaseCount(t, ctx, pool, modelID, 3, 15*time.Second)
|
||||
waitForTaskCount(t, ctx, pool, longTaskIDs, "running", 3, 5*time.Second)
|
||||
if time.Since(updateStartedAt) > 15*time.Second {
|
||||
t.Fatalf("capacity expansion took %s, want <=15s", time.Since(updateStartedAt))
|
||||
}
|
||||
|
||||
initialExpiry := activeLeaseExpiry(t, ctx, pool, modelID, 3)
|
||||
updateAsyncAcceptancePlatform(t, server.URL, adminToken, platform, 1, "enabled")
|
||||
waitForAsyncWorkerMetric(t, server.URL, "easyai_gateway_async_worker_capacity", 1, 15*time.Second)
|
||||
shortTaskIDs := submitAsyncSimulationTasks(t, server.URL, adminToken, "/api/v1/videos/generations", model, 1, time.Second)
|
||||
time.Sleep(3 * time.Second)
|
||||
if attempts := taskAttemptCount(t, ctx, pool, shortTaskIDs); attempts != 0 {
|
||||
t.Fatalf("short task created %d upstream attempts while three long leases exceeded reduced limit", attempts)
|
||||
}
|
||||
waitForTaskCount(t, ctx, pool, longTaskIDs, "running", 3, 5*time.Second)
|
||||
|
||||
time.Sleep(130 * time.Second)
|
||||
renewedExpiry := activeLeaseExpiry(t, ctx, pool, modelID, 3)
|
||||
if !renewedExpiry.After(initialExpiry.Add(30*time.Second)) || !renewedExpiry.After(time.Now()) {
|
||||
t.Fatalf("lease was not renewed: initial=%s renewed=%s now=%s", initialExpiry, renewedExpiry, time.Now())
|
||||
}
|
||||
|
||||
waitForTaskCount(t, ctx, pool, longTaskIDs, "succeeded", 3, 75*time.Second)
|
||||
shortStartedAt := time.Now()
|
||||
waitForTaskAttempts(t, ctx, pool, shortTaskIDs, 1, 5*time.Second)
|
||||
waitForTaskCount(t, ctx, pool, shortTaskIDs, "succeeded", 1, 10*time.Second)
|
||||
if time.Since(shortStartedAt) > 5*time.Second {
|
||||
t.Fatalf("short task started after %s once long tasks completed, want <=5s", time.Since(shortStartedAt))
|
||||
}
|
||||
allTaskIDs := append(append([]string{}, longTaskIDs...), shortTaskIDs...)
|
||||
assertAcceptanceAttempts(t, ctx, pool, allTaskIDs, 4)
|
||||
assertSimulationOutputsAndEvents(t, ctx, pool, allTaskIDs)
|
||||
t.Logf("长任务证据: initial_capacity=1 expanded_capacity=3 reduced_capacity=1 initial_expiry=%s renewed_expiry=%s tasks=4", initialExpiry.UTC().Format(time.RFC3339Nano), renewedExpiry.UTC().Format(time.RFC3339Nano))
|
||||
})
|
||||
}
|
||||
|
||||
type asyncAcceptancePlatform struct {
|
||||
ID string `json:"id"`
|
||||
Provider string `json:"provider"`
|
||||
PlatformKey string `json:"platformKey"`
|
||||
Name string `json:"name"`
|
||||
BaseURL string `json:"baseUrl"`
|
||||
AuthType string `json:"authType"`
|
||||
Priority int `json:"priority"`
|
||||
}
|
||||
|
||||
func createAsyncAcceptanceAdmin(t *testing.T, ctx context.Context, pool *pgxpool.Pool, baseURL string) string {
|
||||
t.Helper()
|
||||
suffix := strconv.FormatInt(time.Now().UnixNano(), 10)
|
||||
username := "async_acceptance_" + suffix
|
||||
password := "password123"
|
||||
doJSON(t, baseURL, http.MethodPost, "/api/v1/auth/register", "", map[string]any{
|
||||
"username": username,
|
||||
"email": username + "@example.com",
|
||||
"password": password,
|
||||
}, http.StatusCreated, &map[string]any{})
|
||||
if _, err := pool.Exec(ctx, `UPDATE gateway_users SET roles = '["admin"]'::jsonb WHERE username = $1`, username); err != nil {
|
||||
t.Fatalf("promote acceptance user: %v", err)
|
||||
}
|
||||
var login struct {
|
||||
AccessToken string `json:"accessToken"`
|
||||
}
|
||||
doJSON(t, baseURL, http.MethodPost, "/api/v1/auth/login", "", map[string]any{
|
||||
"account": username, "password": password,
|
||||
}, http.StatusOK, &login)
|
||||
if login.AccessToken == "" {
|
||||
t.Fatal("acceptance login returned no access token")
|
||||
}
|
||||
return login.AccessToken
|
||||
}
|
||||
|
||||
func createAsyncAcceptancePlatform(t *testing.T, baseURL, token, key, name string, concurrent int) asyncAcceptancePlatform {
|
||||
t.Helper()
|
||||
var platform asyncAcceptancePlatform
|
||||
doJSON(t, baseURL, http.MethodPost, "/api/admin/platforms", token, asyncAcceptancePlatformPayload(key, name, concurrent, "enabled"), http.StatusCreated, &platform)
|
||||
if platform.ID == "" {
|
||||
t.Fatal("acceptance platform returned no id")
|
||||
}
|
||||
return platform
|
||||
}
|
||||
|
||||
func updateAsyncAcceptancePlatform(t *testing.T, baseURL, token string, platform asyncAcceptancePlatform, concurrent int, status string) {
|
||||
t.Helper()
|
||||
doJSON(t, baseURL, http.MethodPatch, "/api/admin/platforms/"+platform.ID, token,
|
||||
asyncAcceptancePlatformPayload(platform.PlatformKey, platform.Name, concurrent, status), http.StatusOK, &asyncAcceptancePlatform{})
|
||||
}
|
||||
|
||||
func asyncAcceptancePlatformPayload(key, name string, concurrent int, status string) map[string]any {
|
||||
return map[string]any{
|
||||
"provider": "openai",
|
||||
"platformKey": key,
|
||||
"name": name,
|
||||
"baseUrl": "https://api.openai.com/v1",
|
||||
"authType": "bearer",
|
||||
"credentials": map[string]any{"mode": "simulation"},
|
||||
"config": map[string]any{"testMode": true},
|
||||
"priority": 1,
|
||||
"status": status,
|
||||
"rateLimitPolicy": map[string]any{"rules": []any{
|
||||
map[string]any{"metric": "concurrent", "limit": concurrent, "leaseTtlSeconds": 120},
|
||||
}},
|
||||
}
|
||||
}
|
||||
|
||||
func createAsyncAcceptanceModel(t *testing.T, baseURL, token, platformID, model, modelType, mode string, concurrent, ttl int) string {
|
||||
t.Helper()
|
||||
payload := map[string]any{
|
||||
"modelName": model,
|
||||
"providerModelName": model,
|
||||
"modelAlias": model,
|
||||
"modelType": []string{modelType},
|
||||
"displayName": model,
|
||||
"rateLimitPolicyMode": mode,
|
||||
}
|
||||
if mode == "override" {
|
||||
payload["rateLimitPolicy"] = map[string]any{"rules": []any{
|
||||
map[string]any{"metric": "concurrent", "limit": concurrent, "leaseTtlSeconds": ttl},
|
||||
}}
|
||||
}
|
||||
var response struct {
|
||||
ID string `json:"id"`
|
||||
}
|
||||
doJSON(t, baseURL, http.MethodPost, "/api/admin/platforms/"+platformID+"/models", token, payload, http.StatusCreated, &response)
|
||||
if response.ID == "" {
|
||||
t.Fatal("acceptance model returned no id")
|
||||
}
|
||||
return response.ID
|
||||
}
|
||||
|
||||
func submitAsyncSimulationTasks(t *testing.T, baseURL, token, path, model string, count int, duration time.Duration) []string {
|
||||
t.Helper()
|
||||
ids := make([]string, count)
|
||||
errs := make(chan error, count)
|
||||
var wg sync.WaitGroup
|
||||
for index := 0; index < count; index++ {
|
||||
wg.Add(1)
|
||||
go func(index int) {
|
||||
defer wg.Done()
|
||||
payload := map[string]any{
|
||||
"model": model,
|
||||
"runMode": "simulation",
|
||||
"simulation": true,
|
||||
"simulationDurationMs": duration.Milliseconds(),
|
||||
"prompt": fmt.Sprintf("async acceptance %d", index),
|
||||
"input": fmt.Sprintf("async acceptance %d", index),
|
||||
}
|
||||
raw, _ := json.Marshal(payload)
|
||||
req, err := http.NewRequest(http.MethodPost, baseURL+path, bytes.NewReader(raw))
|
||||
if err != nil {
|
||||
errs <- err
|
||||
return
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Authorization", "Bearer "+token)
|
||||
req.Header.Set("X-Async", "true")
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
errs <- err
|
||||
return
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
responseBody, _ := io.ReadAll(resp.Body)
|
||||
if resp.StatusCode != http.StatusAccepted {
|
||||
errs <- fmt.Errorf("submit task %d status=%d body=%s", index, resp.StatusCode, responseBody)
|
||||
return
|
||||
}
|
||||
var response struct {
|
||||
TaskID string `json:"taskId"`
|
||||
}
|
||||
if err := json.Unmarshal(responseBody, &response); err != nil || response.TaskID == "" {
|
||||
errs <- fmt.Errorf("decode task %d: %w body=%s", index, err, responseBody)
|
||||
return
|
||||
}
|
||||
ids[index] = response.TaskID
|
||||
}(index)
|
||||
}
|
||||
wg.Wait()
|
||||
close(errs)
|
||||
for err := range errs {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return ids
|
||||
}
|
||||
|
||||
func waitForAcceptanceTasks(t *testing.T, ctx context.Context, pool *pgxpool.Pool, taskIDs []string, modelID string, timeout time.Duration) (int64, int64) {
|
||||
t.Helper()
|
||||
deadline := time.Now().Add(timeout)
|
||||
var peakRunning, peakLeases int64
|
||||
for time.Now().Before(deadline) {
|
||||
var running, succeeded, failed, leases int64
|
||||
if err := pool.QueryRow(ctx, `
|
||||
SELECT COUNT(*) FILTER (WHERE status = 'running'),
|
||||
COUNT(*) FILTER (WHERE status = 'succeeded'),
|
||||
COUNT(*) FILTER (WHERE status IN ('failed', 'cancelled', 'manual_review'))
|
||||
FROM gateway_tasks
|
||||
WHERE id = ANY($1::uuid[])`, taskIDs).Scan(&running, &succeeded, &failed); err != nil {
|
||||
t.Fatalf("read acceptance tasks: %v", err)
|
||||
}
|
||||
if err := pool.QueryRow(ctx, `
|
||||
SELECT COUNT(*)
|
||||
FROM gateway_concurrency_leases
|
||||
WHERE scope_type = 'platform_model'
|
||||
AND scope_key = $1
|
||||
AND released_at IS NULL
|
||||
AND expires_at > now()`, modelID).Scan(&leases); err != nil {
|
||||
t.Fatalf("read acceptance leases: %v", err)
|
||||
}
|
||||
peakRunning = maxInt64(peakRunning, running)
|
||||
peakLeases = maxInt64(peakLeases, leases)
|
||||
if failed > 0 {
|
||||
t.Fatalf("acceptance tasks entered failed/manual status: %d", failed)
|
||||
}
|
||||
if succeeded == int64(len(taskIDs)) {
|
||||
return peakRunning, peakLeases
|
||||
}
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
}
|
||||
t.Fatalf("tasks did not finish within %s", timeout)
|
||||
return 0, 0
|
||||
}
|
||||
|
||||
func waitForTaskCount(t *testing.T, ctx context.Context, pool *pgxpool.Pool, taskIDs []string, status string, count int, timeout time.Duration) {
|
||||
t.Helper()
|
||||
deadline := time.Now().Add(timeout)
|
||||
for time.Now().Before(deadline) {
|
||||
var got int
|
||||
if err := pool.QueryRow(ctx, `SELECT COUNT(*) FROM gateway_tasks WHERE id = ANY($1::uuid[]) AND status = $2`, taskIDs, status).Scan(&got); err != nil {
|
||||
t.Fatalf("read task count: %v", err)
|
||||
}
|
||||
if got == count {
|
||||
return
|
||||
}
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
}
|
||||
t.Fatalf("task status %s count did not reach %d within %s", status, count, timeout)
|
||||
}
|
||||
|
||||
func activeLeaseExpiry(t *testing.T, ctx context.Context, pool *pgxpool.Pool, modelID string, wantCount int) time.Time {
|
||||
t.Helper()
|
||||
var count int
|
||||
var earliest time.Time
|
||||
if err := pool.QueryRow(ctx, `
|
||||
SELECT COUNT(*), MIN(expires_at)
|
||||
FROM gateway_concurrency_leases
|
||||
WHERE scope_type = 'platform_model'
|
||||
AND scope_key = $1
|
||||
AND released_at IS NULL
|
||||
AND expires_at > now()`, modelID).Scan(&count, &earliest); err != nil {
|
||||
t.Fatalf("read active lease expiry: %v", err)
|
||||
}
|
||||
if count != wantCount {
|
||||
t.Fatalf("active lease count=%d, want=%d", count, wantCount)
|
||||
}
|
||||
return earliest
|
||||
}
|
||||
|
||||
func waitForActiveLeaseCount(t *testing.T, ctx context.Context, pool *pgxpool.Pool, modelID string, count int, timeout time.Duration) {
|
||||
t.Helper()
|
||||
deadline := time.Now().Add(timeout)
|
||||
for time.Now().Before(deadline) {
|
||||
var got int
|
||||
if err := pool.QueryRow(ctx, `
|
||||
SELECT COUNT(*)
|
||||
FROM gateway_concurrency_leases
|
||||
WHERE scope_type = 'platform_model'
|
||||
AND scope_key = $1
|
||||
AND released_at IS NULL
|
||||
AND expires_at > now()`, modelID).Scan(&got); err != nil {
|
||||
t.Fatalf("read active lease count: %v", err)
|
||||
}
|
||||
if got == count {
|
||||
return
|
||||
}
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
}
|
||||
t.Fatalf("active lease count did not reach %d within %s", count, timeout)
|
||||
}
|
||||
|
||||
func taskAttemptCount(t *testing.T, ctx context.Context, pool *pgxpool.Pool, taskIDs []string) int {
|
||||
t.Helper()
|
||||
var count int
|
||||
if err := pool.QueryRow(ctx, `SELECT COUNT(*) FROM gateway_task_attempts WHERE task_id = ANY($1::uuid[])`, taskIDs).Scan(&count); err != nil {
|
||||
t.Fatalf("count task attempts: %v", err)
|
||||
}
|
||||
return count
|
||||
}
|
||||
|
||||
func waitForTaskAttempts(t *testing.T, ctx context.Context, pool *pgxpool.Pool, taskIDs []string, count int, timeout time.Duration) {
|
||||
t.Helper()
|
||||
deadline := time.Now().Add(timeout)
|
||||
for time.Now().Before(deadline) {
|
||||
if taskAttemptCount(t, ctx, pool, taskIDs) == count {
|
||||
return
|
||||
}
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
}
|
||||
t.Fatalf("attempt count did not reach %d within %s", count, timeout)
|
||||
}
|
||||
|
||||
func assertAcceptanceAttempts(t *testing.T, ctx context.Context, pool *pgxpool.Pool, taskIDs []string, want int) {
|
||||
t.Helper()
|
||||
var attempts, duplicateTasks, manualReview int
|
||||
if err := pool.QueryRow(ctx, `
|
||||
SELECT COUNT(*),
|
||||
COUNT(DISTINCT task_id) FILTER (WHERE per_task.attempts > 1),
|
||||
COUNT(*) FILTER (WHERE status = 'manual_review')
|
||||
FROM (
|
||||
SELECT task_id, COUNT(*) AS attempts, MAX(status) AS status
|
||||
FROM gateway_task_attempts
|
||||
WHERE task_id = ANY($1::uuid[])
|
||||
GROUP BY task_id
|
||||
) per_task`, taskIDs).Scan(&attempts, &duplicateTasks, &manualReview); err != nil {
|
||||
t.Fatalf("read attempt acceptance evidence: %v", err)
|
||||
}
|
||||
if attempts != want || duplicateTasks != 0 || manualReview != 0 {
|
||||
t.Fatalf("attempts=%d duplicate_tasks=%d manual_review=%d, want %d/0/0", attempts, duplicateTasks, manualReview, want)
|
||||
}
|
||||
}
|
||||
|
||||
func assertSimulationOutputsAndEvents(t *testing.T, ctx context.Context, pool *pgxpool.Pool, taskIDs []string) {
|
||||
t.Helper()
|
||||
var complete int
|
||||
if err := pool.QueryRow(ctx, `
|
||||
SELECT COUNT(*)
|
||||
FROM gateway_tasks task
|
||||
WHERE task.id = ANY($1::uuid[])
|
||||
AND task.status = 'succeeded'
|
||||
AND task.finished_at IS NOT NULL
|
||||
AND task.result IS NOT NULL
|
||||
AND task.result <> '{}'::jsonb
|
||||
AND task.metrics IS NOT NULL
|
||||
AND jsonb_typeof(task.billings) = 'array'
|
||||
AND jsonb_array_length(task.billings) > 0
|
||||
AND EXISTS (SELECT 1 FROM gateway_task_events event WHERE event.task_id = task.id AND event.event_type = 'task.completed')`,
|
||||
taskIDs).Scan(&complete); err != nil {
|
||||
t.Fatalf("read simulation output evidence: %v", err)
|
||||
}
|
||||
if complete != len(taskIDs) {
|
||||
t.Fatalf("complete simulation outputs/events=%d, want=%d", complete, len(taskIDs))
|
||||
}
|
||||
}
|
||||
|
||||
func waitForAsyncWorkerMetric(t *testing.T, baseURL, metric string, want int64, timeout time.Duration) {
|
||||
t.Helper()
|
||||
deadline := time.Now().Add(timeout)
|
||||
for time.Now().Before(deadline) {
|
||||
resp, err := http.Get(baseURL + "/metrics")
|
||||
if err == nil {
|
||||
raw, _ := io.ReadAll(resp.Body)
|
||||
resp.Body.Close()
|
||||
if resp.StatusCode == http.StatusOK {
|
||||
for _, line := range strings.Split(string(raw), "\n") {
|
||||
fields := strings.Fields(line)
|
||||
if len(fields) == 2 && fields[0] == metric {
|
||||
value, _ := strconv.ParseInt(fields[1], 10, 64)
|
||||
if value == want {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
}
|
||||
t.Fatalf("metric %s did not reach %d within %s", metric, want, timeout)
|
||||
}
|
||||
|
||||
func maxInt64(left, right int64) int64 {
|
||||
if right > left {
|
||||
return right
|
||||
}
|
||||
return left
|
||||
}
|
||||
@@ -0,0 +1,30 @@
|
||||
package httpapi
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/clients"
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
)
|
||||
|
||||
func (s *Server) hydrateTaskResult(ctx context.Context, task store.GatewayTask) (store.GatewayTask, error) {
|
||||
if task.Status != "succeeded" || len(task.Result) == 0 {
|
||||
return task, nil
|
||||
}
|
||||
result, err := s.runner.HydrateTaskResult(ctx, task.ID, task.Result)
|
||||
if err != nil {
|
||||
return store.GatewayTask{}, err
|
||||
}
|
||||
task.Result = result
|
||||
return task, nil
|
||||
}
|
||||
|
||||
func writeStoredBinaryResultError(w http.ResponseWriter, err error) {
|
||||
status := statusFromRunError(err)
|
||||
code := clients.ErrorCode(err)
|
||||
if code == "" {
|
||||
code = "binary_result_corrupted"
|
||||
}
|
||||
writeError(w, status, err.Error(), code)
|
||||
}
|
||||
@@ -182,6 +182,10 @@ func (s *Server) createBaseModel(w http.ResponseWriter, r *http.Request) {
|
||||
}
|
||||
item, err := s.store.CreateBaseModel(r.Context(), input)
|
||||
if err != nil {
|
||||
if errors.Is(err, store.ErrModelAliasConflict) {
|
||||
writeError(w, http.StatusConflict, err.Error(), "model_alias_conflict")
|
||||
return
|
||||
}
|
||||
if store.IsUniqueViolation(err) {
|
||||
writeError(w, http.StatusConflict, "canonical model key already exists")
|
||||
return
|
||||
@@ -230,6 +234,10 @@ func (s *Server) updateBaseModel(w http.ResponseWriter, r *http.Request) {
|
||||
writeError(w, http.StatusConflict, "canonical model key already exists")
|
||||
return
|
||||
}
|
||||
if errors.Is(err, store.ErrModelAliasConflict) {
|
||||
writeError(w, http.StatusConflict, err.Error(), "model_alias_conflict")
|
||||
return
|
||||
}
|
||||
s.logger.Error("update base model failed", "error", err)
|
||||
writeError(w, http.StatusInternalServerError, "update base model failed")
|
||||
return
|
||||
@@ -301,6 +309,7 @@ func (s *Server) resetAllBaseModels(w http.ResponseWriter, r *http.Request) {
|
||||
// @Failure 401 {object} ErrorEnvelope
|
||||
// @Failure 403 {object} ErrorEnvelope
|
||||
// @Failure 404 {object} ErrorEnvelope
|
||||
// @Failure 409 {object} ErrorEnvelope
|
||||
// @Failure 500 {object} ErrorEnvelope
|
||||
// @Router /api/admin/catalog/base-models/{baseModelID} [delete]
|
||||
func (s *Server) deleteBaseModel(w http.ResponseWriter, r *http.Request) {
|
||||
@@ -309,6 +318,10 @@ func (s *Server) deleteBaseModel(w http.ResponseWriter, r *http.Request) {
|
||||
writeError(w, http.StatusNotFound, "base model not found")
|
||||
return
|
||||
}
|
||||
if errors.Is(err, store.ErrBaseModelInUse) {
|
||||
writeError(w, http.StatusConflict, "base model is still referenced by platform models", "base_model_in_use")
|
||||
return
|
||||
}
|
||||
s.logger.Error("delete base model failed", "error", err)
|
||||
writeError(w, http.StatusInternalServerError, "delete base model failed")
|
||||
return
|
||||
|
||||
@@ -123,6 +123,17 @@ func writeOpenAIError(w http.ResponseWriter, status int, message string, details
|
||||
if strings.TrimSpace(code) != "" {
|
||||
payload["code"] = code
|
||||
}
|
||||
if len(details) > 0 {
|
||||
publicDetails := map[string]any{}
|
||||
for key, value := range details {
|
||||
if key != "param" {
|
||||
publicDetails[key] = value
|
||||
}
|
||||
}
|
||||
if len(publicDetails) > 0 {
|
||||
payload["details"] = publicDetails
|
||||
}
|
||||
}
|
||||
writeJSON(w, status, map[string]any{"error": payload})
|
||||
}
|
||||
|
||||
|
||||
@@ -300,23 +300,41 @@ VALUES ($1, 5, '{"purpose":"core-flow"}'::jsonb)`, inviteCode); err != nil {
|
||||
Items []struct {
|
||||
ID string `json:"id"`
|
||||
CanonicalModelKey string `json:"canonicalModelKey"`
|
||||
InvocationName string `json:"invocationName"`
|
||||
ProviderModelName string `json:"providerModelName"`
|
||||
ModelType []string `json:"modelType"`
|
||||
ModelAlias string `json:"modelAlias"`
|
||||
DisplayName string `json:"displayName"`
|
||||
Status string `json:"status"`
|
||||
} `json:"items"`
|
||||
}
|
||||
doJSON(t, server.URL, http.MethodGet, "/api/admin/catalog/base-models", loginResponse.AccessToken, nil, http.StatusOK, &baseModels)
|
||||
if len(baseModels.Items) < 300 {
|
||||
t.Fatalf("server-main seed should include the migrated base model catalog: got %d", len(baseModels.Items))
|
||||
if len(baseModels.Items) == 0 {
|
||||
t.Fatal("migrated base model catalog must not be empty")
|
||||
}
|
||||
requiredLifecycleModels := map[string]string{
|
||||
"gemini:gemini-3.6-flash": "active",
|
||||
"aliyun:qwen3.7-max": "active",
|
||||
}
|
||||
for _, model := range baseModels.Items {
|
||||
if expectedStatus, ok := requiredLifecycleModels[model.CanonicalModelKey]; ok {
|
||||
if model.InvocationName == "" || model.ModelAlias != model.InvocationName || model.Status != expectedStatus {
|
||||
t.Fatalf("base model identity/lifecycle mismatch: %+v", model)
|
||||
}
|
||||
delete(requiredLifecycleModels, model.CanonicalModelKey)
|
||||
}
|
||||
}
|
||||
if len(requiredLifecycleModels) != 0 {
|
||||
t.Fatalf("missing lifecycle-managed base models: %+v", requiredLifecycleModels)
|
||||
}
|
||||
requiredMiniMaxModels := map[string]struct {
|
||||
providerModelName string
|
||||
modelType string
|
||||
alias string
|
||||
}{
|
||||
"minimax:speech-2.8-hd": {providerModelName: "speech-2.8-hd", modelType: "text_to_speech", alias: "MiniMax-Speech-2.8-HD"},
|
||||
"minimax:speech-2.8-turbo": {providerModelName: "speech-2.8-turbo", modelType: "text_to_speech", alias: "MiniMax-Speech-2.8-Turbo"},
|
||||
"minimax:voice-clone": {providerModelName: "voice_clone", modelType: "voice_clone", alias: "MiniMax-Voice-Clone"},
|
||||
"minimax:speech-2.8-hd": {providerModelName: "speech-2.8-hd", modelType: "text_to_speech", alias: "speech-2.8-hd"},
|
||||
"minimax:speech-2.8-turbo": {providerModelName: "speech-2.8-turbo", modelType: "text_to_speech", alias: "speech-2.8-turbo"},
|
||||
"minimax:voice-clone": {providerModelName: "voice_clone", modelType: "voice_clone", alias: "voice_clone"},
|
||||
}
|
||||
seenMiniMaxModels := map[string]bool{}
|
||||
for _, model := range baseModels.Items {
|
||||
@@ -361,9 +379,11 @@ VALUES ($1, 5, '{"purpose":"core-flow"}'::jsonb)`, inviteCode); err != nil {
|
||||
baseModelInput := map[string]any{
|
||||
"providerKey": "openai",
|
||||
"canonicalModelKey": "openai:smoke-base-" + suffixText,
|
||||
"invocationName": "smoke-base-" + suffixText,
|
||||
"providerModelName": "smoke-base-" + suffixText,
|
||||
"modelType": []string{"text_generate"},
|
||||
"modelAlias": "Smoke Base Model",
|
||||
"displayName": "Smoke Base Model",
|
||||
"legacyAliases": []string{"Smoke Base Model Legacy"},
|
||||
"capabilities": map[string]any{"originalTypes": []string{"text_generate"}},
|
||||
"metadata": map[string]any{"source": "test"},
|
||||
}
|
||||
@@ -376,12 +396,13 @@ VALUES ($1, 5, '{"purpose":"core-flow"}'::jsonb)`, inviteCode); err != nil {
|
||||
if createdBaseModel.ID == "" || createdBaseModel.CanonicalModelKey != baseModelInput["canonicalModelKey"] {
|
||||
t.Fatalf("unexpected created base model: %+v", createdBaseModel)
|
||||
}
|
||||
baseModelInput["modelAlias"] = "Smoke Base Model Updated"
|
||||
baseModelInput["displayName"] = "Smoke Base Model Updated"
|
||||
var updatedBaseModel struct {
|
||||
ModelAlias string `json:"modelAlias"`
|
||||
ModelAlias string `json:"modelAlias"`
|
||||
DisplayName string `json:"displayName"`
|
||||
}
|
||||
doJSON(t, server.URL, http.MethodPatch, "/api/admin/catalog/base-models/"+createdBaseModel.ID, loginResponse.AccessToken, baseModelInput, http.StatusOK, &updatedBaseModel)
|
||||
if updatedBaseModel.ModelAlias != "Smoke Base Model Updated" {
|
||||
if updatedBaseModel.ModelAlias != baseModelInput["invocationName"] || updatedBaseModel.DisplayName != "Smoke Base Model Updated" {
|
||||
t.Fatalf("unexpected updated base model: %+v", updatedBaseModel)
|
||||
}
|
||||
doJSON(t, server.URL, http.MethodDelete, "/api/admin/catalog/base-models/"+createdBaseModel.ID, loginResponse.AccessToken, nil, http.StatusNoContent, nil)
|
||||
@@ -466,6 +487,7 @@ VALUES ($1, 5, '{"purpose":"core-flow"}'::jsonb)`, inviteCode); err != nil {
|
||||
cancelReq.Header.Set("Content-Type", "application/json")
|
||||
cancelErrCh := make(chan error, 1)
|
||||
cancelTaskIDCh := make(chan string, 1)
|
||||
cancelRequestStartedAt := time.Now().UTC().Add(-time.Second)
|
||||
go func() {
|
||||
resp, err := http.DefaultClient.Do(cancelReq)
|
||||
if resp != nil {
|
||||
@@ -484,7 +506,7 @@ VALUES ($1, 5, '{"purpose":"core-flow"}'::jsonb)`, inviteCode); err != nil {
|
||||
t.Fatal("cancelled stream did not return response headers")
|
||||
}
|
||||
if cancelTaskID == "" {
|
||||
t.Fatal("cancelled stream response did not expose X-Gateway-Task-Id")
|
||||
cancelTaskID = waitForLatestTaskIDSince(t, ctx, testPool, "chat.completions", cancelRequestStartedAt, 2*time.Second)
|
||||
}
|
||||
cancelRequest()
|
||||
select {
|
||||
@@ -498,11 +520,20 @@ VALUES ($1, 5, '{"purpose":"core-flow"}'::jsonb)`, inviteCode); err != nil {
|
||||
}
|
||||
|
||||
var imageResponse struct {
|
||||
Task struct {
|
||||
Status string `json:"status"`
|
||||
LegacyTaskID string `json:"task_id"`
|
||||
TaskID string `json:"taskId"`
|
||||
QueryURL string `json:"query_url"`
|
||||
Task struct {
|
||||
ID string `json:"id"`
|
||||
Status string `json:"status"`
|
||||
Result map[string]any `json:"result"`
|
||||
} `json:"task"`
|
||||
Next struct {
|
||||
Detail string `json:"detail"`
|
||||
Events string `json:"events"`
|
||||
Result string `json:"result"`
|
||||
} `json:"next"`
|
||||
}
|
||||
doJSONWithHeaders(t, server.URL, http.MethodPost, "/api/v1/images/generations", apiKeyResponse.Secret, map[string]any{
|
||||
"model": defaultImageModel,
|
||||
@@ -513,14 +544,59 @@ VALUES ($1, 5, '{"purpose":"core-flow"}'::jsonb)`, inviteCode); err != nil {
|
||||
"simulation": true,
|
||||
"simulationDurationMs": 5,
|
||||
}, map[string]string{"X-Async": "true"}, http.StatusAccepted, &imageResponse)
|
||||
if imageResponse.Status != "submitted" ||
|
||||
imageResponse.TaskID == "" ||
|
||||
imageResponse.LegacyTaskID != imageResponse.TaskID ||
|
||||
imageResponse.Task.ID != imageResponse.TaskID ||
|
||||
imageResponse.QueryURL != "/api/v1/ai/result/"+imageResponse.TaskID ||
|
||||
imageResponse.Next.Result != imageResponse.QueryURL {
|
||||
t.Fatalf("unexpected EasyAI-compatible image submission: %+v", imageResponse)
|
||||
}
|
||||
waitForTaskStatus(t, server.URL, apiKeyResponse.Secret, imageResponse.Task.ID, []string{"succeeded"}, 10*time.Second)
|
||||
doJSON(t, server.URL, http.MethodGet, "/api/v1/tasks/"+imageResponse.Task.ID, apiKeyResponse.Secret, nil, http.StatusOK, &imageResponse.Task)
|
||||
if imageResponse.Task.Status != "succeeded" || imageResponse.Task.Result["id"] == "" {
|
||||
t.Fatalf("unexpected image generation task: %+v", imageResponse.Task)
|
||||
}
|
||||
var imageEasyAIResult struct {
|
||||
Status string `json:"status"`
|
||||
TaskID string `json:"task_id"`
|
||||
Data []map[string]any `json:"data"`
|
||||
Output []string `json:"output"`
|
||||
OutputContent []map[string]any `json:"output_content"`
|
||||
}
|
||||
doJSON(t, server.URL, http.MethodGet, imageResponse.QueryURL, apiKeyResponse.Secret, nil, http.StatusOK, &imageEasyAIResult)
|
||||
if imageEasyAIResult.Status != "success" ||
|
||||
imageEasyAIResult.TaskID != imageResponse.TaskID ||
|
||||
len(imageEasyAIResult.Data) == 0 ||
|
||||
len(imageEasyAIResult.Output) == 0 ||
|
||||
len(imageEasyAIResult.OutputContent) != len(imageEasyAIResult.Data) ||
|
||||
imageEasyAIResult.Data[0]["type"] != "image" {
|
||||
t.Fatalf("unexpected EasyAI-compatible image result: %+v", imageEasyAIResult)
|
||||
}
|
||||
var ordinaryLoginResponse struct {
|
||||
AccessToken string `json:"accessToken"`
|
||||
}
|
||||
doJSON(t, server.URL, http.MethodPost, "/api/v1/auth/login", "", map[string]any{
|
||||
"account": ordinaryUsername,
|
||||
"password": password,
|
||||
}, http.StatusOK, &ordinaryLoginResponse)
|
||||
var inaccessibleImageResult struct {
|
||||
Status string `json:"status"`
|
||||
Code string `json:"code"`
|
||||
Data []any `json:"data"`
|
||||
}
|
||||
doJSON(t, server.URL, http.MethodGet, imageResponse.QueryURL, ordinaryLoginResponse.AccessToken, nil, http.StatusNotFound, &inaccessibleImageResult)
|
||||
if inaccessibleImageResult.Status != "failed" ||
|
||||
inaccessibleImageResult.Code != "not_found" ||
|
||||
len(inaccessibleImageResult.Data) != 0 {
|
||||
t.Fatalf("EasyAI task result must stay owner-isolated: %+v", inaccessibleImageResult)
|
||||
}
|
||||
|
||||
var imageEditResponse struct {
|
||||
Task struct {
|
||||
Status string `json:"status"`
|
||||
LegacyTaskID string `json:"task_id"`
|
||||
TaskID string `json:"taskId"`
|
||||
Task struct {
|
||||
ID string `json:"id"`
|
||||
Status string `json:"status"`
|
||||
Result map[string]any `json:"result"`
|
||||
@@ -535,6 +611,12 @@ VALUES ($1, 5, '{"purpose":"core-flow"}'::jsonb)`, inviteCode); err != nil {
|
||||
"simulation": true,
|
||||
"simulationDurationMs": 5,
|
||||
}, map[string]string{"X-Async": "true"}, http.StatusAccepted, &imageEditResponse)
|
||||
if imageEditResponse.Status != "submitted" ||
|
||||
imageEditResponse.TaskID == "" ||
|
||||
imageEditResponse.LegacyTaskID != imageEditResponse.TaskID ||
|
||||
imageEditResponse.Task.ID != imageEditResponse.TaskID {
|
||||
t.Fatalf("unexpected EasyAI-compatible image edit submission: %+v", imageEditResponse)
|
||||
}
|
||||
waitForTaskStatus(t, server.URL, apiKeyResponse.Secret, imageEditResponse.Task.ID, []string{"succeeded"}, 10*time.Second)
|
||||
doJSON(t, server.URL, http.MethodGet, "/api/v1/tasks/"+imageEditResponse.Task.ID, apiKeyResponse.Secret, nil, http.StatusOK, &imageEditResponse.Task)
|
||||
if imageEditResponse.Task.Status != "succeeded" || imageEditResponse.Task.Result["id"] == "" {
|
||||
@@ -1613,6 +1695,64 @@ WHERE m.platform_id = $1::uuid
|
||||
if !taskListContains(workspaceTaskList.Items, taskResponse.Task.ID) || !taskListContains(workspaceTaskList.Items, pricingTask.Task.ID) {
|
||||
t.Fatalf("workspace task list should include persisted task records, got %+v", workspaceTaskList.Items)
|
||||
}
|
||||
doJSON(t, server.URL, http.MethodGet, "/api/admin/tasks", ordinaryLoginResponse.AccessToken, nil, http.StatusForbidden, nil)
|
||||
if _, err := testPool.Exec(ctx, `
|
||||
UPDATE gateway_tasks
|
||||
SET request = request || '{"diagnostic":{"apiKey":"private-admin-task-test-key"}}'::jsonb
|
||||
WHERE id = $1::uuid`, taskResponse.Task.ID); err != nil {
|
||||
t.Fatalf("seed admin task redaction payload: %v", err)
|
||||
}
|
||||
var adminTaskList struct {
|
||||
Items []struct {
|
||||
ID string `json:"id"`
|
||||
Request map[string]any `json:"request"`
|
||||
Context struct {
|
||||
User struct {
|
||||
ID string `json:"id"`
|
||||
Username string `json:"username"`
|
||||
} `json:"user"`
|
||||
Tenant struct {
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
} `json:"tenant"`
|
||||
} `json:"adminContext"`
|
||||
} `json:"items"`
|
||||
Total int `json:"total"`
|
||||
}
|
||||
adminTaskQuery := "/api/admin/tasks?q=" + taskResponse.Task.ID +
|
||||
"&userId=" + smokeGatewayUserID +
|
||||
"&status=succeeded" +
|
||||
"&platformId=" + platform.ID +
|
||||
"&model=" + defaultTextModel +
|
||||
"&modelType=text_generate" +
|
||||
"&runMode=simulation" +
|
||||
"&billingStatus=settled" +
|
||||
"&apiKey=smoke" +
|
||||
"&pageSize=10"
|
||||
doJSON(t, server.URL, http.MethodGet, adminTaskQuery, loginResponse.AccessToken, nil, http.StatusOK, &adminTaskList)
|
||||
if len(adminTaskList.Items) != 1 || adminTaskList.Total != 1 || adminTaskList.Items[0].Context.User.Username != username {
|
||||
t.Fatalf("admin task list should expose the matching task and identity summary: %+v", adminTaskList)
|
||||
}
|
||||
maskedDiagnostic, _ := adminTaskList.Items[0].Request["diagnostic"].(map[string]any)
|
||||
if maskedDiagnostic["apiKey"] != "***" {
|
||||
t.Fatalf("admin task list must mask nested secrets: %+v", adminTaskList.Items[0].Request)
|
||||
}
|
||||
var maskedAdminTask struct {
|
||||
Request map[string]any `json:"request"`
|
||||
}
|
||||
doJSON(t, server.URL, http.MethodGet, "/api/admin/tasks/"+taskResponse.Task.ID, loginResponse.AccessToken, nil, http.StatusOK, &maskedAdminTask)
|
||||
maskedDiagnostic, _ = maskedAdminTask.Request["diagnostic"].(map[string]any)
|
||||
if maskedDiagnostic["apiKey"] != "***" {
|
||||
t.Fatalf("admin task detail must be masked by default: %+v", maskedAdminTask.Request)
|
||||
}
|
||||
var fullAdminTask struct {
|
||||
Request map[string]any `json:"request"`
|
||||
}
|
||||
doJSON(t, server.URL, http.MethodGet, "/api/admin/tasks/"+taskResponse.Task.ID+"?sensitive=full", loginResponse.AccessToken, nil, http.StatusOK, &fullAdminTask)
|
||||
fullDiagnostic, _ := fullAdminTask.Request["diagnostic"].(map[string]any)
|
||||
if fullDiagnostic["apiKey"] != "private-admin-task-test-key" {
|
||||
t.Fatalf("explicit full admin task detail should return stored content: %+v", fullAdminTask.Request)
|
||||
}
|
||||
|
||||
req, err := http.NewRequest(http.MethodGet, server.URL+"/api/v1/tasks/"+pricingTask.Task.ID+"/events", nil)
|
||||
if err != nil {
|
||||
@@ -1664,6 +1804,7 @@ WHERE m.platform_id = $1::uuid
|
||||
}
|
||||
|
||||
restartModel := "worker-restart-" + suffixText
|
||||
createSimulationTextBaseModel(t, server.URL, loginResponse.AccessToken, restartModel)
|
||||
createSimulationTextPlatformModel(
|
||||
t,
|
||||
server.URL,
|
||||
@@ -1782,6 +1923,7 @@ func TestCacheAffinityRoutingFlow(t *testing.T) {
|
||||
|
||||
model := "cache-affinity-smoke-" + suffixText
|
||||
cacheKey := "cache-affinity-key-" + suffixText
|
||||
createSimulationTextBaseModel(t, server.URL, loginResponse.AccessToken, model)
|
||||
lowPlatform, lowPlatformModelID := createSimulationTextPlatformModel(t, server.URL, loginResponse.AccessToken, "cache-low-"+suffixText, "Cache Affinity Low Priority", model, 20, map[string]any{
|
||||
"rules": []map[string]any{
|
||||
{"metric": "concurrent", "limit": 1, "leaseTtlSeconds": 120},
|
||||
@@ -1798,8 +1940,31 @@ func TestCacheAffinityRoutingFlow(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
highPlatform, _ := createSimulationTextPlatformModel(t, server.URL, loginResponse.AccessToken, "cache-high-"+suffixText, "Cache Affinity High Priority", model, 1, nil)
|
||||
_ = highPlatform
|
||||
tools := []any{map[string]any{
|
||||
"type": "function",
|
||||
"function": map[string]any{
|
||||
"name": "lookup_weather",
|
||||
"description": "Look up current weather",
|
||||
"parameters": map[string]any{"type": "object", "properties": map[string]any{"city": map[string]any{"type": "string"}}},
|
||||
},
|
||||
}}
|
||||
initialToolMessages := []any{
|
||||
map[string]any{"role": "system", "content": "Use tools when helpful."},
|
||||
map[string]any{"role": "user", "content": "Weather in Hangzhou?"},
|
||||
}
|
||||
for index := 0; index < 3; index++ {
|
||||
detail := runCacheAffinityChatTaskWithBody(t, ctx, testPool, server.URL, apiKeyResponse.Secret, model, "tool-prime-"+strconv.Itoa(index)+"-"+suffixText, map[string]any{
|
||||
"tools": tools,
|
||||
"messages": initialToolMessages,
|
||||
"simulationUsage": map[string]any{"inputTokens": 1000, "cachedInputTokens": 750, "outputTokens": 20},
|
||||
})
|
||||
if detail.Status != "succeeded" || len(detail.Attempts) != 1 || detail.Attempts[0].PlatformName != "Cache Affinity Low Priority" {
|
||||
t.Fatalf("tool-history priming request should use the only platform, got %+v", detail)
|
||||
}
|
||||
}
|
||||
|
||||
peerPlatform, _ := createSimulationTextPlatformModel(t, server.URL, loginResponse.AccessToken, "cache-peer-"+suffixText, "Cache Affinity Peer", model, 20, nil)
|
||||
_ = peerPlatform
|
||||
sameKeyDetail := runCacheAffinityChatTask(t, ctx, testPool, server.URL, apiKeyResponse.Secret, model, cacheKey, "same-key-"+suffixText, map[string]any{
|
||||
"inputTokens": 1000,
|
||||
"cachedInputTokens": 700,
|
||||
@@ -1817,12 +1982,36 @@ func TestCacheAffinityRoutingFlow(t *testing.T) {
|
||||
t.Fatalf("same cache key task detail should expose simulated usage, got %+v", sameKeyDetail.Usage)
|
||||
}
|
||||
|
||||
toolHistoryDetail := runCacheAffinityChatTaskWithBody(t, ctx, testPool, server.URL, apiKeyResponse.Secret, model, "tool-history-"+suffixText, map[string]any{
|
||||
"tools": tools,
|
||||
"messages": append(append([]any{}, initialToolMessages...),
|
||||
map[string]any{"role": "assistant", "tool_calls": []any{map[string]any{
|
||||
"id": "call_weather",
|
||||
"type": "function",
|
||||
"function": map[string]any{
|
||||
"name": "lookup_weather",
|
||||
"arguments": `{"city":"Hangzhou"}`,
|
||||
},
|
||||
}}},
|
||||
map[string]any{"role": "tool", "tool_call_id": "call_weather", "content": `{"temperature":31}`},
|
||||
),
|
||||
"simulationUsage": map[string]any{"inputTokens": 1400, "cachedInputTokens": 900, "outputTokens": 30},
|
||||
})
|
||||
if toolHistoryDetail.Status != "succeeded" || len(toolHistoryDetail.Attempts) != 1 || toolHistoryDetail.Attempts[0].PlatformName != "Cache Affinity Low Priority" {
|
||||
t.Fatalf("tool continuation should retain the cached platform, got %+v", toolHistoryDetail)
|
||||
}
|
||||
if !boolFromTestMap(toolHistoryDetail.Attempts[0].Metrics, "cacheAffinityMatched") ||
|
||||
intFromTestAny(toolHistoryDetail.Attempts[0].Metrics["cacheAffinityMatchedPrefixDepth"]) < 2 ||
|
||||
intFromTestAny(toolHistoryDetail.Attempts[0].Metrics["cacheAffinityCandidateCount"]) < 1 {
|
||||
t.Fatalf("tool continuation should expose prefix and candidate metrics, got %+v", toolHistoryDetail.Attempts[0].Metrics)
|
||||
}
|
||||
|
||||
newKeyDetail := runCacheAffinityChatTask(t, ctx, testPool, server.URL, apiKeyResponse.Secret, model, "cache-affinity-new-"+suffixText, "new-key-"+suffixText, map[string]any{
|
||||
"inputTokens": 1000,
|
||||
"cachedInputTokens": 0,
|
||||
"outputTokens": 20,
|
||||
})
|
||||
if newKeyDetail.Status != "succeeded" || len(newKeyDetail.Attempts) != 1 || newKeyDetail.Attempts[0].PlatformName != "Cache Affinity High Priority" {
|
||||
if newKeyDetail.Status != "succeeded" || len(newKeyDetail.Attempts) != 1 || newKeyDetail.Attempts[0].PlatformName != "Cache Affinity Peer" {
|
||||
t.Fatalf("new cache key should fall back to base priority, got %+v", newKeyDetail)
|
||||
}
|
||||
|
||||
@@ -1832,9 +2021,12 @@ func TestCacheAffinityRoutingFlow(t *testing.T) {
|
||||
"cachedInputTokens": 0,
|
||||
"outputTokens": 20,
|
||||
})
|
||||
if fullAvoidedDetail.Status != "succeeded" || len(fullAvoidedDetail.Attempts) != 1 || fullAvoidedDetail.Attempts[0].PlatformName != "Cache Affinity High Priority" {
|
||||
if fullAvoidedDetail.Status != "succeeded" || len(fullAvoidedDetail.Attempts) != 1 || fullAvoidedDetail.Attempts[0].PlatformName != "Cache Affinity Peer" {
|
||||
t.Fatalf("full cached platform should be avoided before cache affinity boost, got %+v", fullAvoidedDetail)
|
||||
}
|
||||
if fullAvoidedDetail.Attempts[0].Metrics["cacheAffinityOverrideReason"] != "capacity_tier_unavailable" {
|
||||
t.Fatalf("full cached platform override reason should be visible, got %+v", fullAvoidedDetail.Attempts[0].Metrics)
|
||||
}
|
||||
|
||||
var requestCount int
|
||||
var inputTokens int
|
||||
@@ -1844,7 +2036,7 @@ func TestCacheAffinityRoutingFlow(t *testing.T) {
|
||||
SELECT request_count::int, input_tokens::int, cached_input_tokens::int, ema_hit_ratio::float8
|
||||
FROM gateway_cache_affinity_stats
|
||||
WHERE client_id = $1
|
||||
ORDER BY updated_at DESC
|
||||
ORDER BY request_count DESC, cached_input_tokens DESC, updated_at DESC
|
||||
LIMIT 1`, lowPlatform.PlatformKey+":text_generate:"+model).Scan(&requestCount, &inputTokens, &cachedTokens, &emaHitRatio); err != nil {
|
||||
t.Fatalf("read cache affinity stats: %v", err)
|
||||
}
|
||||
@@ -1927,6 +2119,7 @@ func applyMigration(t *testing.T, ctx context.Context, databaseURL string) {
|
||||
t.Fatalf("connect migration db: %v", err)
|
||||
}
|
||||
defer pool.Close()
|
||||
requireDedicatedIntegrationTestDatabase(t, ctx, pool)
|
||||
if _, err := pool.Exec(ctx, `
|
||||
CREATE TABLE IF NOT EXISTS schema_migrations (
|
||||
version text PRIMARY KEY,
|
||||
@@ -1947,6 +2140,22 @@ CREATE TABLE IF NOT EXISTS schema_migrations (
|
||||
if err != nil {
|
||||
t.Fatalf("read migration %s: %v", filepath.Base(migrationPath), err)
|
||||
}
|
||||
migrationSQL := string(migration)
|
||||
const noTransactionMarker = "-- easyai:migration:no-transaction"
|
||||
if strings.HasPrefix(strings.TrimSpace(migrationSQL), noTransactionMarker) {
|
||||
migrationSQL = strings.TrimSpace(strings.TrimPrefix(strings.TrimSpace(migrationSQL), noTransactionMarker))
|
||||
for _, statement := range strings.Split(migrationSQL, "-- easyai:migration:statement") {
|
||||
if statement = strings.TrimSpace(statement); statement != "" {
|
||||
if _, err := pool.Exec(ctx, statement); err != nil {
|
||||
t.Fatalf("apply non-transaction migration %s: %v", filepath.Base(migrationPath), err)
|
||||
}
|
||||
}
|
||||
}
|
||||
if _, err := pool.Exec(ctx, "INSERT INTO schema_migrations(version) VALUES($1)", version); err != nil {
|
||||
t.Fatalf("record non-transaction migration %s: %v", filepath.Base(migrationPath), err)
|
||||
}
|
||||
continue
|
||||
}
|
||||
tx, err := pool.Begin(ctx)
|
||||
if err != nil {
|
||||
t.Fatalf("begin migration %s: %v", filepath.Base(migrationPath), err)
|
||||
@@ -1974,9 +2183,9 @@ SELECT EXISTS (
|
||||
FROM base_model_catalog model
|
||||
WHERE model.provider_key = 'volces'
|
||||
AND model.canonical_model_key = 'volces:doubao-seedream-5-0-pro-260628'
|
||||
AND model.invocation_name = 'doubao-seedream-5-0-pro-260628'
|
||||
AND model.provider_model_name = 'doubao-seedream-5-0-pro-260628'
|
||||
AND model.display_name = 'Seedream-5.0-Pro'
|
||||
AND model.display_name !~ '[[:space:]]'
|
||||
AND model.display_name = 'Seedream 5.0 Pro'
|
||||
AND model.model_type = '["image_edit","image_generate"]'::jsonb
|
||||
AND model.capabilities->'image_edit'->>'input_multiple_images' = 'true'
|
||||
AND model.capabilities->'image_edit'->>'input_max_images_count' = '10'
|
||||
@@ -1991,7 +2200,7 @@ SELECT EXISTS (
|
||||
AND model.capabilities->>'stream' = 'false'
|
||||
AND model.capabilities->>'supportWebSearch' = 'false'
|
||||
AND model.default_rate_limit_policy->'rules'->0->>'limit' = '500'
|
||||
AND model.default_snapshot->>'modelAlias' = 'Seedream-5.0-Pro'
|
||||
AND model.default_snapshot->>'modelAlias' = 'doubao-seedream-5-0-pro-260628'
|
||||
AND model.default_snapshot->>'displayName' = 'Seedream 5.0 Pro'
|
||||
)`).Scan(&valid); err != nil {
|
||||
t.Fatalf("query Seedream 5.0 Pro catalog migration: %v", err)
|
||||
@@ -2048,16 +2257,17 @@ func doJSONWithHeaders(t *testing.T, baseURL string, method string, path string,
|
||||
|
||||
func doAPIV1ChatCompletionAndLoadTask(t *testing.T, ctx context.Context, pool *pgxpool.Pool, baseURL string, token string, payload map[string]any, marker string, expectedStatus int, responseOut any, taskDetailOut any) string {
|
||||
t.Helper()
|
||||
_ = ctx
|
||||
_ = pool
|
||||
_ = marker
|
||||
if responseOut == nil {
|
||||
responseOut = &map[string]any{}
|
||||
}
|
||||
requestStartedAt := time.Now().UTC().Add(-time.Second)
|
||||
responseHeaders := doJSONWithHeaders(t, baseURL, http.MethodPost, "/api/v1/chat/completions", token, payload, nil, expectedStatus, responseOut)
|
||||
taskID := strings.TrimSpace(responseHeaders.Get("X-Gateway-Task-Id"))
|
||||
if taskID == "" {
|
||||
t.Fatal("chat completion response did not expose X-Gateway-Task-Id")
|
||||
taskID = waitForLatestTaskIDSince(t, ctx, pool, "chat.completions", requestStartedAt, 2*time.Second)
|
||||
}
|
||||
if taskID == "" {
|
||||
t.Fatalf("load chat completion task %s failed: headers=%v body=%+v", marker, responseHeaders, responseOut)
|
||||
}
|
||||
if taskDetailOut != nil {
|
||||
doJSON(t, baseURL, http.MethodGet, "/api/v1/tasks/"+taskID, token, nil, http.StatusOK, taskDetailOut)
|
||||
@@ -2065,6 +2275,27 @@ func doAPIV1ChatCompletionAndLoadTask(t *testing.T, ctx context.Context, pool *p
|
||||
return taskID
|
||||
}
|
||||
|
||||
func waitForLatestTaskIDSince(t *testing.T, ctx context.Context, pool *pgxpool.Pool, kind string, since time.Time, timeout time.Duration) string {
|
||||
t.Helper()
|
||||
deadline := time.Now().Add(timeout)
|
||||
for time.Now().Before(deadline) {
|
||||
var taskID string
|
||||
err := pool.QueryRow(ctx, `
|
||||
SELECT id::text
|
||||
FROM gateway_tasks
|
||||
WHERE kind = $1
|
||||
AND created_at >= $2
|
||||
ORDER BY created_at DESC
|
||||
LIMIT 1`, kind, since).Scan(&taskID)
|
||||
if err == nil && taskID != "" {
|
||||
return taskID
|
||||
}
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
}
|
||||
t.Fatalf("latest %s task was not persisted within %s", kind, timeout)
|
||||
return ""
|
||||
}
|
||||
|
||||
type taskWaitDetail struct {
|
||||
ID string `json:"id"`
|
||||
Status string `json:"status"`
|
||||
@@ -2093,6 +2324,20 @@ type cacheAffinityTaskDetail struct {
|
||||
} `json:"attempts"`
|
||||
}
|
||||
|
||||
func createSimulationTextBaseModel(t *testing.T, baseURL string, adminToken string, model string) {
|
||||
t.Helper()
|
||||
doJSON(t, baseURL, http.MethodPost, "/api/admin/catalog/base-models", adminToken, map[string]any{
|
||||
"providerKey": "openai",
|
||||
"canonicalModelKey": "openai:" + model,
|
||||
"invocationName": model,
|
||||
"providerModelName": model,
|
||||
"modelType": []string{"text_generate"},
|
||||
"displayName": model,
|
||||
"capabilities": map[string]any{"originalTypes": []string{"text_generate"}},
|
||||
"metadata": map[string]any{"source": "cache-affinity-integration-test"},
|
||||
}, http.StatusCreated, nil)
|
||||
}
|
||||
|
||||
func createSimulationTextPlatformModel(t *testing.T, baseURL string, adminToken string, platformKey string, platformName string, model string, priority int, rateLimitPolicy map[string]any) (cacheAffinityPlatformFixture, string) {
|
||||
t.Helper()
|
||||
var platform cacheAffinityPlatformFixture
|
||||
@@ -2109,7 +2354,7 @@ func createSimulationTextPlatformModel(t *testing.T, baseURL string, adminToken
|
||||
t.Fatalf("cache affinity platform was not created: %+v", platform)
|
||||
}
|
||||
payload := map[string]any{
|
||||
"canonicalModelKey": "openai:gpt-4o-mini",
|
||||
"canonicalModelKey": "openai:" + model,
|
||||
"modelName": model,
|
||||
"providerModelName": model,
|
||||
"modelAlias": model,
|
||||
@@ -2130,17 +2375,27 @@ func createSimulationTextPlatformModel(t *testing.T, baseURL string, adminToken
|
||||
}
|
||||
|
||||
func runCacheAffinityChatTask(t *testing.T, ctx context.Context, pool *pgxpool.Pool, baseURL string, token string, model string, cacheAffinityKey string, marker string, simulationUsage map[string]any) cacheAffinityTaskDetail {
|
||||
t.Helper()
|
||||
return runCacheAffinityChatTaskWithBody(t, ctx, pool, baseURL, token, model, marker, map[string]any{
|
||||
"cacheAffinityKey": cacheAffinityKey,
|
||||
"simulationUsage": simulationUsage,
|
||||
"messages": []any{map[string]any{"role": "user", "content": "cache affinity route"}},
|
||||
})
|
||||
}
|
||||
|
||||
func runCacheAffinityChatTaskWithBody(t *testing.T, ctx context.Context, pool *pgxpool.Pool, baseURL string, token string, model string, marker string, body map[string]any) cacheAffinityTaskDetail {
|
||||
t.Helper()
|
||||
var detail cacheAffinityTaskDetail
|
||||
doAPIV1ChatCompletionAndLoadTask(t, ctx, pool, baseURL, token, map[string]any{
|
||||
requestBody := map[string]any{
|
||||
"model": model,
|
||||
"runMode": "simulation",
|
||||
"simulation": true,
|
||||
"simulationDurationMs": 5,
|
||||
"cacheAffinityKey": cacheAffinityKey,
|
||||
"simulationUsage": simulationUsage,
|
||||
"messages": []map[string]any{{"role": "user", "content": "cache affinity route"}},
|
||||
}, marker, http.StatusOK, nil, &detail)
|
||||
}
|
||||
for key, value := range body {
|
||||
requestBody[key] = value
|
||||
}
|
||||
doAPIV1ChatCompletionAndLoadTask(t, ctx, pool, baseURL, token, requestBody, marker, http.StatusOK, nil, &detail)
|
||||
return detail
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,489 @@
|
||||
package httpapi
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/runner"
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
)
|
||||
|
||||
func easyAIAsyncMediaRequest(kind string, r *http.Request) bool {
|
||||
if r == nil || !asyncRequest(r) {
|
||||
return false
|
||||
}
|
||||
switch kind {
|
||||
case "images.generations",
|
||||
"images.edits",
|
||||
"images.vectorize",
|
||||
"videos.generations",
|
||||
"videos.upscales",
|
||||
"song.generations",
|
||||
"music.generations",
|
||||
"speech.generations",
|
||||
"voice.clone":
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func easyAITaskAcceptedResponse(task store.GatewayTask) map[string]any {
|
||||
resultPath := "/api/v1/ai/result/" + task.ID
|
||||
return map[string]any{
|
||||
"status": "submitted",
|
||||
"task_id": task.ID,
|
||||
"taskId": task.ID,
|
||||
"query_url": resultPath,
|
||||
"task": task,
|
||||
"next": map[string]string{
|
||||
"events": "/api/v1/tasks/" + task.ID + "/events",
|
||||
"detail": "/api/v1/tasks/" + task.ID,
|
||||
"result": resultPath,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func easyAITaskAcceptedResponseWithExtensions(task store.GatewayTask, extensions map[string]any) map[string]any {
|
||||
response := cloneEasyAIMap(extensions)
|
||||
for key, value := range easyAITaskAcceptedResponse(task) {
|
||||
response[key] = value
|
||||
}
|
||||
return response
|
||||
}
|
||||
|
||||
func writeEasyAIAsyncError(w http.ResponseWriter, status int, message string, details map[string]any, codes ...string) {
|
||||
code := ""
|
||||
if len(codes) > 0 {
|
||||
code = strings.TrimSpace(codes[0])
|
||||
}
|
||||
errorPayload := map[string]any{
|
||||
"message": message,
|
||||
"status": status,
|
||||
}
|
||||
if code != "" {
|
||||
errorPayload["code"] = code
|
||||
}
|
||||
for key, value := range details {
|
||||
errorPayload[key] = value
|
||||
}
|
||||
response := map[string]any{
|
||||
"status": "failed",
|
||||
"message": message,
|
||||
"data": []any{},
|
||||
"error": errorPayload,
|
||||
}
|
||||
if code != "" {
|
||||
response["code"] = code
|
||||
}
|
||||
writeJSON(w, status, response)
|
||||
}
|
||||
|
||||
func easyAISynchronousTaskResponse(task store.GatewayTask, output map[string]any) map[string]any {
|
||||
if task.Kind != "images.vectorize" {
|
||||
return output
|
||||
}
|
||||
response := cloneEasyAIMap(output)
|
||||
response["taskId"] = task.ID
|
||||
response["task_id"] = task.ID
|
||||
response["query_url"] = "/api/v1/ai/result/" + task.ID
|
||||
return response
|
||||
}
|
||||
|
||||
func easyAIFileUploadResponse(upload map[string]any) map[string]any {
|
||||
response := cloneEasyAIMap(upload)
|
||||
response["status"] = "success"
|
||||
response["data"] = cloneEasyAIMap(upload)
|
||||
return response
|
||||
}
|
||||
|
||||
func easyAITaskResultResponse(task store.GatewayTask) map[string]any {
|
||||
sourceResult := cloneEasyAIMap(task.Result)
|
||||
cleanResult := cloneEasyAIMap(sourceResult)
|
||||
delete(cleanResult, "raw")
|
||||
delete(cleanResult, "raw_data")
|
||||
normalizeEasyAIInlineMediaFields(cleanResult)
|
||||
|
||||
data := easyAITaskResultData(task, sourceResult)
|
||||
status := easyAITaskResultStatus(task.Status)
|
||||
if status == "failed" {
|
||||
data = []any{}
|
||||
}
|
||||
cleanResult["data"] = data
|
||||
cleanResult["output_content"] = data
|
||||
|
||||
response := cloneEasyAIMap(cleanResult)
|
||||
response["status"] = status
|
||||
response["task_id"] = task.ID
|
||||
response["taskId"] = task.ID
|
||||
response["query_url"] = "/api/v1/ai/result/" + task.ID
|
||||
response["created"] = task.CreatedAt.UnixMilli()
|
||||
response["data"] = data
|
||||
response["output_content"] = data
|
||||
response["output"] = easyAIOutputURLs(data)
|
||||
response["result"] = cleanResult
|
||||
|
||||
upstreamTaskID := firstNonEmpty(
|
||||
easyAIString(cleanResult["upstream_task_id"]),
|
||||
task.RemoteTaskID,
|
||||
)
|
||||
if upstreamTaskID != "" {
|
||||
response["upstream_task_id"] = upstreamTaskID
|
||||
}
|
||||
if len(task.Usage) > 0 && response["usage"] == nil {
|
||||
response["usage"] = task.Usage
|
||||
}
|
||||
|
||||
cancelState := runner.DescribeTaskCancellation(task)
|
||||
response["cancellable"] = cancelState.Cancellable
|
||||
response["submitted"] = cancelState.Submitted
|
||||
|
||||
message := firstNonEmpty(
|
||||
easyAIString(cleanResult["message"]),
|
||||
task.ErrorMessage,
|
||||
task.Error,
|
||||
task.Message,
|
||||
)
|
||||
if message == "" {
|
||||
switch status {
|
||||
case "submitted":
|
||||
message = "任务已提交,请稍后查看结果"
|
||||
case "process":
|
||||
message = "任务处理中"
|
||||
case "failed":
|
||||
message = "任务执行失败"
|
||||
}
|
||||
}
|
||||
if message != "" {
|
||||
response["message"] = message
|
||||
}
|
||||
if code := firstNonEmpty(task.ErrorCode, easyAIString(cleanResult["code"])); code != "" {
|
||||
response["code"] = code
|
||||
}
|
||||
return response
|
||||
}
|
||||
|
||||
func easyAITaskResultStatus(status string) string {
|
||||
switch strings.ToLower(strings.TrimSpace(status)) {
|
||||
case "succeeded", "success", "completed":
|
||||
return "success"
|
||||
case "failed", "cancelled", "canceled", "rejected":
|
||||
return "failed"
|
||||
case "running", "processing", "in_progress":
|
||||
return "process"
|
||||
default:
|
||||
return "submitted"
|
||||
}
|
||||
}
|
||||
|
||||
func easyAITaskResultData(task store.GatewayTask, result map[string]any) []any {
|
||||
for _, source := range easyAIResultSources(result) {
|
||||
if items := easyAIOutputArray(source); len(items) > 0 {
|
||||
return normalizeEasyAIOutputItems(task, items)
|
||||
}
|
||||
}
|
||||
if item := easyAIOutputMap(result); len(item) > 0 {
|
||||
normalizeEasyAIInlineMediaFields(item)
|
||||
if easyAIOutputURL(item) != "" ||
|
||||
easyAIOutputHasInlineMedia(item) ||
|
||||
(strings.EqualFold(easyAIString(item["type"]), "text") && (item["text"] != nil || item["content"] != nil)) {
|
||||
return normalizeEasyAIOutputItems(task, []any{item})
|
||||
}
|
||||
}
|
||||
return []any{}
|
||||
}
|
||||
|
||||
func easyAIResultSources(result map[string]any) []any {
|
||||
sources := []any{
|
||||
result["data"],
|
||||
result["output_content"],
|
||||
easyAIStructuredOutput(result["content"]),
|
||||
result["images"],
|
||||
result["videos"],
|
||||
result["audios"],
|
||||
result["files"],
|
||||
result["output"],
|
||||
}
|
||||
if raw, ok := result["raw"].(map[string]any); ok {
|
||||
sources = append(sources,
|
||||
raw["data"],
|
||||
raw["output_content"],
|
||||
easyAIStructuredOutput(raw["content"]),
|
||||
raw["images"],
|
||||
raw["videos"],
|
||||
raw["audios"],
|
||||
raw["files"],
|
||||
raw["output"],
|
||||
)
|
||||
}
|
||||
return sources
|
||||
}
|
||||
|
||||
func easyAIStructuredOutput(value any) any {
|
||||
switch value.(type) {
|
||||
case []any, []string, map[string]any:
|
||||
return value
|
||||
default:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
func easyAIOutputArray(value any) []any {
|
||||
switch typed := value.(type) {
|
||||
case []any:
|
||||
return typed
|
||||
case []string:
|
||||
items := make([]any, 0, len(typed))
|
||||
for _, item := range typed {
|
||||
items = append(items, item)
|
||||
}
|
||||
return items
|
||||
case map[string]any, string:
|
||||
return []any{typed}
|
||||
default:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
func normalizeEasyAIOutputItems(task store.GatewayTask, items []any) []any {
|
||||
normalized := make([]any, 0, len(items))
|
||||
for _, item := range items {
|
||||
output := easyAIOutputMap(item)
|
||||
if len(output) == 0 {
|
||||
continue
|
||||
}
|
||||
delete(output, "raw_data")
|
||||
normalizeEasyAIInlineMediaFields(output)
|
||||
|
||||
mediaURL := easyAIOutputURL(output)
|
||||
if mediaURL != "" {
|
||||
output["url"] = mediaURL
|
||||
}
|
||||
if strings.TrimSpace(easyAIString(output["type"])) == "" {
|
||||
if outputType := easyAIOutputType(task, output, mediaURL); outputType != "" {
|
||||
output["type"] = outputType
|
||||
}
|
||||
}
|
||||
normalized = append(normalized, output)
|
||||
}
|
||||
return normalized
|
||||
}
|
||||
|
||||
func easyAIOutputMap(value any) map[string]any {
|
||||
switch typed := value.(type) {
|
||||
case map[string]any:
|
||||
return cloneEasyAIMap(typed)
|
||||
case string:
|
||||
trimmed := strings.TrimSpace(typed)
|
||||
if trimmed == "" {
|
||||
return nil
|
||||
}
|
||||
if strings.HasPrefix(strings.ToLower(trimmed), "data:") {
|
||||
return map[string]any{"url": trimmed}
|
||||
}
|
||||
if strings.HasPrefix(trimmed, "/") ||
|
||||
strings.HasPrefix(strings.ToLower(trimmed), "http://") ||
|
||||
strings.HasPrefix(strings.ToLower(trimmed), "https://") {
|
||||
return map[string]any{"url": trimmed}
|
||||
}
|
||||
return map[string]any{"type": "text", "content": trimmed}
|
||||
default:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
func easyAIOutputURL(item map[string]any) string {
|
||||
for _, key := range []string{
|
||||
"url",
|
||||
"generated_url",
|
||||
"generatedUrl",
|
||||
"output_url",
|
||||
"outputUrl",
|
||||
"download_url",
|
||||
"downloadUrl",
|
||||
"image_url",
|
||||
"video_url",
|
||||
"audio_url",
|
||||
} {
|
||||
if value := easyAINestedURL(item[key]); value != "" {
|
||||
return value
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func easyAINestedURL(value any) string {
|
||||
var candidate string
|
||||
switch typed := value.(type) {
|
||||
case string:
|
||||
candidate = strings.TrimSpace(typed)
|
||||
case map[string]any:
|
||||
candidate = strings.TrimSpace(easyAIString(typed["url"]))
|
||||
default:
|
||||
return ""
|
||||
}
|
||||
if strings.HasPrefix(strings.ToLower(candidate), "data:") {
|
||||
return ""
|
||||
}
|
||||
return candidate
|
||||
}
|
||||
|
||||
func normalizeEasyAIInlineMediaFields(item map[string]any) {
|
||||
for _, key := range []string{
|
||||
"url",
|
||||
"generated_url",
|
||||
"generatedUrl",
|
||||
"output_url",
|
||||
"outputUrl",
|
||||
"download_url",
|
||||
"downloadUrl",
|
||||
"image_url",
|
||||
"video_url",
|
||||
"audio_url",
|
||||
} {
|
||||
value := easyAINestedURLAllowData(item[key])
|
||||
contentType, payload, ok := easyAIBase64DataURL(value)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
if strings.HasPrefix(contentType, "image/") {
|
||||
if easyAIString(item["b64_json"]) == "" {
|
||||
item["b64_json"] = payload
|
||||
}
|
||||
} else if easyAIString(item["content"]) == "" {
|
||||
item["content"] = value
|
||||
}
|
||||
if easyAIString(item["mime_type"]) == "" {
|
||||
item["mime_type"] = contentType
|
||||
}
|
||||
delete(item, key)
|
||||
}
|
||||
}
|
||||
|
||||
func easyAIBase64DataURL(value string) (string, string, bool) {
|
||||
value = strings.TrimSpace(value)
|
||||
if !strings.HasPrefix(strings.ToLower(value), "data:") {
|
||||
return "", "", false
|
||||
}
|
||||
meta, payload, ok := strings.Cut(value, ",")
|
||||
if !ok || strings.TrimSpace(payload) == "" {
|
||||
return "", "", false
|
||||
}
|
||||
parts := strings.Split(strings.TrimSpace(meta[5:]), ";")
|
||||
if len(parts) < 2 || !strings.EqualFold(strings.TrimSpace(parts[len(parts)-1]), "base64") {
|
||||
return "", "", false
|
||||
}
|
||||
contentType := strings.ToLower(strings.TrimSpace(parts[0]))
|
||||
if !strings.HasPrefix(contentType, "image/") &&
|
||||
!strings.HasPrefix(contentType, "audio/") &&
|
||||
!strings.HasPrefix(contentType, "video/") &&
|
||||
!strings.HasPrefix(contentType, "application/") {
|
||||
return "", "", false
|
||||
}
|
||||
return contentType, strings.TrimSpace(payload), true
|
||||
}
|
||||
|
||||
func easyAIOutputHasInlineMedia(item map[string]any) bool {
|
||||
if easyAIString(item["b64_json"]) != "" {
|
||||
return true
|
||||
}
|
||||
content := strings.ToLower(easyAIString(item["content"]))
|
||||
return strings.HasPrefix(content, "data:") && strings.Contains(content, ";base64,")
|
||||
}
|
||||
|
||||
func easyAINestedURLAllowData(value any) string {
|
||||
switch typed := value.(type) {
|
||||
case string:
|
||||
return strings.TrimSpace(typed)
|
||||
case map[string]any:
|
||||
return strings.TrimSpace(easyAIString(typed["url"]))
|
||||
default:
|
||||
return ""
|
||||
}
|
||||
}
|
||||
|
||||
func easyAIOutputType(task store.GatewayTask, item map[string]any, mediaURL string) string {
|
||||
kind := strings.ToLower(strings.TrimSpace(task.Kind))
|
||||
switch {
|
||||
case kind == "images.vectorize":
|
||||
format := strings.ToLower(firstNonEmpty(
|
||||
easyAIString(item["format"]),
|
||||
easyAIString(task.Request["format"]),
|
||||
))
|
||||
if format == "svg" || format == "png" {
|
||||
return "image"
|
||||
}
|
||||
if format != "" {
|
||||
return "file"
|
||||
}
|
||||
case strings.HasPrefix(kind, "images."):
|
||||
return "image"
|
||||
case strings.HasPrefix(kind, "videos.") || kind == "video.upscale":
|
||||
return "video"
|
||||
case kind == "song.generations" || kind == "music.generations" || kind == "speech.generations":
|
||||
return "audio"
|
||||
}
|
||||
|
||||
lowerURL := strings.ToLower(mediaURL)
|
||||
if index := strings.IndexAny(lowerURL, "?#"); index >= 0 {
|
||||
lowerURL = lowerURL[:index]
|
||||
}
|
||||
switch {
|
||||
case strings.HasSuffix(lowerURL, ".svg"),
|
||||
strings.HasSuffix(lowerURL, ".png"),
|
||||
strings.HasSuffix(lowerURL, ".jpg"),
|
||||
strings.HasSuffix(lowerURL, ".jpeg"),
|
||||
strings.HasSuffix(lowerURL, ".webp"),
|
||||
strings.HasSuffix(lowerURL, ".gif"):
|
||||
return "image"
|
||||
case strings.HasSuffix(lowerURL, ".mp4"),
|
||||
strings.HasSuffix(lowerURL, ".mov"),
|
||||
strings.HasSuffix(lowerURL, ".webm"):
|
||||
return "video"
|
||||
case strings.HasSuffix(lowerURL, ".mp3"),
|
||||
strings.HasSuffix(lowerURL, ".wav"),
|
||||
strings.HasSuffix(lowerURL, ".m4a"),
|
||||
strings.HasSuffix(lowerURL, ".aac"),
|
||||
strings.HasSuffix(lowerURL, ".flac"):
|
||||
return "audio"
|
||||
case mediaURL != "":
|
||||
return "file"
|
||||
case item["text"] != nil || item["content"] != nil:
|
||||
return "text"
|
||||
default:
|
||||
return ""
|
||||
}
|
||||
}
|
||||
|
||||
func easyAIOutputURLs(data []any) []string {
|
||||
output := make([]string, 0, len(data))
|
||||
for _, item := range data {
|
||||
if typed, ok := item.(map[string]any); ok {
|
||||
if mediaURL := easyAIOutputURL(typed); mediaURL != "" {
|
||||
output = append(output, mediaURL)
|
||||
}
|
||||
}
|
||||
}
|
||||
return output
|
||||
}
|
||||
|
||||
func easyAIString(value any) string {
|
||||
typed, _ := value.(string)
|
||||
return strings.TrimSpace(typed)
|
||||
}
|
||||
|
||||
func cloneEasyAIMap(source map[string]any) map[string]any {
|
||||
if len(source) == 0 {
|
||||
return map[string]any{}
|
||||
}
|
||||
raw, err := json.Marshal(source)
|
||||
if err != nil {
|
||||
return map[string]any{}
|
||||
}
|
||||
var result map[string]any
|
||||
if err := json.Unmarshal(raw, &result); err != nil {
|
||||
return map[string]any{}
|
||||
}
|
||||
return result
|
||||
}
|
||||
@@ -0,0 +1,311 @@
|
||||
package httpapi
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
)
|
||||
|
||||
func TestEasyAITaskAcceptedResponseKeepsGatewayFieldsAndAddsLegacyFields(t *testing.T) {
|
||||
task := store.GatewayTask{
|
||||
ID: "task-accepted-1",
|
||||
Kind: "speech.generations",
|
||||
Status: "queued",
|
||||
AsyncMode: true,
|
||||
}
|
||||
|
||||
got := easyAITaskAcceptedResponse(task)
|
||||
if got["status"] != "submitted" || got["task_id"] != task.ID || got["taskId"] != task.ID {
|
||||
t.Fatalf("unexpected EasyAI submission identity: %+v", got)
|
||||
}
|
||||
if _, ok := got["task"].(store.GatewayTask); !ok {
|
||||
t.Fatalf("Gateway task envelope was removed: %+v", got)
|
||||
}
|
||||
next, _ := got["next"].(map[string]string)
|
||||
if next["detail"] != "/api/v1/tasks/"+task.ID ||
|
||||
next["events"] != "/api/v1/tasks/"+task.ID+"/events" ||
|
||||
next["result"] != "/api/v1/ai/result/"+task.ID {
|
||||
t.Fatalf("unexpected next links: %+v", next)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEasyAITaskResultStatus(t *testing.T) {
|
||||
for _, item := range []struct {
|
||||
input string
|
||||
want string
|
||||
}{
|
||||
{input: "queued", want: "submitted"},
|
||||
{input: "pending", want: "submitted"},
|
||||
{input: "running", want: "process"},
|
||||
{input: "processing", want: "process"},
|
||||
{input: "succeeded", want: "success"},
|
||||
{input: "completed", want: "success"},
|
||||
{input: "failed", want: "failed"},
|
||||
{input: "cancelled", want: "failed"},
|
||||
} {
|
||||
t.Run(item.input, func(t *testing.T) {
|
||||
if got := easyAITaskResultStatus(item.input); got != item.want {
|
||||
t.Fatalf("easyAITaskResultStatus(%q)=%q, want %q", item.input, got, item.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestEasyAITaskResultResponseNormalizesMediaOutputs(t *testing.T) {
|
||||
createdAt := time.Date(2026, 7, 23, 10, 0, 0, 0, time.UTC)
|
||||
for _, item := range []struct {
|
||||
name string
|
||||
task store.GatewayTask
|
||||
wantType string
|
||||
wantURL string
|
||||
wantCount int
|
||||
}{
|
||||
{
|
||||
name: "image alias",
|
||||
task: store.GatewayTask{
|
||||
ID: "image-1", Kind: "images.generations", Status: "succeeded", CreatedAt: createdAt,
|
||||
Result: map[string]any{"data": []any{map[string]any{
|
||||
"image_url": "https://example.com/output.png",
|
||||
}}},
|
||||
},
|
||||
wantType: "image", wantURL: "https://example.com/output.png", wantCount: 1,
|
||||
},
|
||||
{
|
||||
name: "video raw content",
|
||||
task: store.GatewayTask{
|
||||
ID: "video-1", Kind: "videos.generations", Status: "succeeded", CreatedAt: createdAt,
|
||||
Result: map[string]any{"raw": map[string]any{"content": map[string]any{
|
||||
"video_url": "https://example.com/output.mp4",
|
||||
"duration": 5,
|
||||
}}},
|
||||
},
|
||||
wantType: "video", wantURL: "https://example.com/output.mp4", wantCount: 1,
|
||||
},
|
||||
{
|
||||
name: "audio alias",
|
||||
task: store.GatewayTask{
|
||||
ID: "audio-1", Kind: "speech.generations", Status: "succeeded", CreatedAt: createdAt,
|
||||
Result: map[string]any{"data": []any{map[string]any{
|
||||
"audio_url": "https://example.com/output.wav",
|
||||
}}},
|
||||
},
|
||||
wantType: "audio", wantURL: "https://example.com/output.wav", wantCount: 1,
|
||||
},
|
||||
{
|
||||
name: "voice clone without media",
|
||||
task: store.GatewayTask{
|
||||
ID: "voice-1", Kind: "voice.clone", Status: "succeeded", CreatedAt: createdAt,
|
||||
Result: map[string]any{"voice_id": "voice-test-1", "cloned_voice": map[string]any{"voiceId": "voice-test-1"}},
|
||||
},
|
||||
wantCount: 0,
|
||||
},
|
||||
} {
|
||||
t.Run(item.name, func(t *testing.T) {
|
||||
got := easyAITaskResultResponse(item.task)
|
||||
if got["status"] != "success" || got["task_id"] != item.task.ID || got["created"] != createdAt.UnixMilli() {
|
||||
t.Fatalf("unexpected task result identity: %+v", got)
|
||||
}
|
||||
data, _ := got["data"].([]any)
|
||||
if len(data) != item.wantCount {
|
||||
t.Fatalf("unexpected output count: got=%d response=%+v", len(data), got)
|
||||
}
|
||||
outputContent, _ := got["output_content"].([]any)
|
||||
if len(outputContent) != item.wantCount {
|
||||
t.Fatalf("output_content is not synchronized: %+v", got)
|
||||
}
|
||||
if item.wantCount == 0 {
|
||||
if got["voice_id"] != "voice-test-1" {
|
||||
t.Fatalf("voice clone fields were lost: %+v", got)
|
||||
}
|
||||
return
|
||||
}
|
||||
output, _ := data[0].(map[string]any)
|
||||
if output["type"] != item.wantType || output["url"] != item.wantURL {
|
||||
t.Fatalf("unexpected normalized output: %+v", output)
|
||||
}
|
||||
if _, exposed := output["b64_json"]; exposed {
|
||||
t.Fatalf("base64 output was exposed: %+v", output)
|
||||
}
|
||||
urls, _ := got["output"].([]string)
|
||||
if len(urls) != 1 || urls[0] != item.wantURL {
|
||||
t.Fatalf("unexpected output URL list: %+v", got["output"])
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestEasyAITaskResultResponsePreservesBase64ImageOutput(t *testing.T) {
|
||||
base64Payload := "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAusB9Y9ZlB8AAAAASUVORK5CYII="
|
||||
task := store.GatewayTask{
|
||||
ID: "image-base64-1", Kind: "images.generations", Status: "succeeded",
|
||||
Result: map[string]any{"data": []any{map[string]any{
|
||||
"b64_json": base64Payload,
|
||||
"mime_type": "image/png",
|
||||
}}},
|
||||
}
|
||||
|
||||
got := easyAITaskResultResponse(task)
|
||||
data, _ := got["data"].([]any)
|
||||
if len(data) != 1 {
|
||||
t.Fatalf("unexpected base64 output count: %+v", got)
|
||||
}
|
||||
output, _ := data[0].(map[string]any)
|
||||
if output["type"] != "image" || output["b64_json"] != base64Payload || output["mime_type"] != "image/png" {
|
||||
t.Fatalf("base64 image output was not preserved: %+v", output)
|
||||
}
|
||||
if _, exposedAsURL := output["url"]; exposedAsURL {
|
||||
t.Fatalf("base64 image must not be exposed as a URL: %+v", output)
|
||||
}
|
||||
outputContent, _ := got["output_content"].([]any)
|
||||
content, _ := outputContent[0].(map[string]any)
|
||||
if content["b64_json"] != base64Payload {
|
||||
t.Fatalf("output_content lost base64 image: %+v", outputContent)
|
||||
}
|
||||
urls, _ := got["output"].([]string)
|
||||
if len(urls) != 0 {
|
||||
t.Fatalf("base64-only output must not create URL aliases: %+v", urls)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEasyAITaskResultResponseMovesImageDataURLToB64JSON(t *testing.T) {
|
||||
base64Payload := "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAusB9Y9ZlB8AAAAASUVORK5CYII="
|
||||
task := store.GatewayTask{
|
||||
ID: "image-data-url-1", Kind: "images.generations", Status: "succeeded",
|
||||
Result: map[string]any{"data": []any{map[string]any{
|
||||
"url": "data:image/png;base64," + base64Payload,
|
||||
}}},
|
||||
}
|
||||
|
||||
got := easyAITaskResultResponse(task)
|
||||
data, _ := got["data"].([]any)
|
||||
output, _ := data[0].(map[string]any)
|
||||
if output["b64_json"] != base64Payload || output["mime_type"] != "image/png" || output["type"] != "image" {
|
||||
t.Fatalf("image data URL was not normalized: %+v", output)
|
||||
}
|
||||
if _, exists := output["url"]; exists {
|
||||
t.Fatalf("image data URL should move out of the URL field: %+v", output)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEasyAITaskResultResponsePreservesAudioDataURLAsContent(t *testing.T) {
|
||||
base64Payload := "SUQzBAAAAAAAI1RTU0UAAAAPAAADTGF2ZjYwLjMuMTAwAAAAAAAAAAAAAAD/"
|
||||
task := store.GatewayTask{
|
||||
ID: "audio-data-url-1", Kind: "speech.generations", Status: "succeeded",
|
||||
Result: map[string]any{"data": []any{map[string]any{
|
||||
"audio_url": "data:audio/mpeg;base64," + base64Payload,
|
||||
}}},
|
||||
}
|
||||
|
||||
got := easyAITaskResultResponse(task)
|
||||
data, _ := got["data"].([]any)
|
||||
output, _ := data[0].(map[string]any)
|
||||
if output["content"] != "data:audio/mpeg;base64,"+base64Payload ||
|
||||
output["mime_type"] != "audio/mpeg" ||
|
||||
output["type"] != "audio" {
|
||||
t.Fatalf("audio data URL was not normalized to inline content: %+v", output)
|
||||
}
|
||||
if _, exists := output["audio_url"]; exists {
|
||||
t.Fatalf("audio data URL should move out of the URL field: %+v", output)
|
||||
}
|
||||
urls, _ := got["output"].([]string)
|
||||
if len(urls) != 0 {
|
||||
t.Fatalf("base64-only audio must not create URL aliases: %+v", urls)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEasyAITaskResultResponseMapsVectorFormats(t *testing.T) {
|
||||
for _, item := range []struct {
|
||||
format string
|
||||
wantType string
|
||||
}{
|
||||
{format: "svg", wantType: "image"},
|
||||
{format: "png", wantType: "image"},
|
||||
{format: "pdf", wantType: "file"},
|
||||
{format: "eps", wantType: "file"},
|
||||
} {
|
||||
t.Run(item.format, func(t *testing.T) {
|
||||
task := store.GatewayTask{
|
||||
ID: "vector-" + item.format, Kind: "images.vectorize", Status: "succeeded",
|
||||
Request: map[string]any{"format": item.format},
|
||||
Result: map[string]any{"data": []any{map[string]any{"url": "https://example.com/output." + item.format}}},
|
||||
}
|
||||
got := easyAITaskResultResponse(task)
|
||||
data, _ := got["data"].([]any)
|
||||
output, _ := data[0].(map[string]any)
|
||||
if output["type"] != item.wantType {
|
||||
t.Fatalf("format %s mapped to %v, want %s", item.format, output["type"], item.wantType)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestEasyAITaskResultResponseClearsFailedOutputs(t *testing.T) {
|
||||
task := store.GatewayTask{
|
||||
ID: "failed-1", Kind: "videos.upscales", Status: "cancelled",
|
||||
ErrorCode: "cancelled", ErrorMessage: "任务已取消",
|
||||
Result: map[string]any{"data": []any{map[string]any{
|
||||
"url": "https://example.com/partial.mp4",
|
||||
}}},
|
||||
}
|
||||
got := easyAITaskResultResponse(task)
|
||||
data, _ := got["data"].([]any)
|
||||
outputContent, _ := got["output_content"].([]any)
|
||||
if got["status"] != "failed" || got["code"] != "cancelled" || got["message"] != "任务已取消" ||
|
||||
len(data) != 0 || len(outputContent) != 0 {
|
||||
t.Fatalf("unexpected failed response: %+v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEasyAISynchronousVectorAndFileUploadResponsesAreAdditive(t *testing.T) {
|
||||
vectorTask := store.GatewayTask{ID: "vector-sync-1", Kind: "images.vectorize"}
|
||||
vector := easyAISynchronousTaskResponse(vectorTask, map[string]any{
|
||||
"status": "success",
|
||||
"data": []any{map[string]any{"url": "https://example.com/output.svg"}},
|
||||
})
|
||||
if vector["status"] != "success" || vector["taskId"] != vectorTask.ID || vector["task_id"] != vectorTask.ID {
|
||||
t.Fatalf("unexpected vector response: %+v", vector)
|
||||
}
|
||||
|
||||
upload := easyAIFileUploadResponse(map[string]any{
|
||||
"id": "file-1", "url": "https://example.com/upload.png", "filename": "upload.png",
|
||||
})
|
||||
data, _ := upload["data"].(map[string]any)
|
||||
if upload["status"] != "success" || upload["url"] != data["url"] || upload["id"] != data["id"] {
|
||||
t.Fatalf("unexpected upload response: %+v", upload)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEasyAIAsyncMediaRequestRequiresExplicitHeader(t *testing.T) {
|
||||
syncRequest := httptest.NewRequest(http.MethodPost, "/api/v1/images/generations", nil)
|
||||
if easyAIAsyncMediaRequest("images.generations", syncRequest) {
|
||||
t.Fatal("image request without X-Async must stay synchronous")
|
||||
}
|
||||
asyncRequest := httptest.NewRequest(http.MethodPost, "/api/v1/images/generations", nil)
|
||||
asyncRequest.Header.Set("X-Async", "true")
|
||||
if !easyAIAsyncMediaRequest("images.generations", asyncRequest) {
|
||||
t.Fatal("image request with X-Async must enable EasyAI async compatibility")
|
||||
}
|
||||
if easyAIAsyncMediaRequest("chat.completions", asyncRequest) {
|
||||
t.Fatal("chat completions must not enter EasyAI media async compatibility")
|
||||
}
|
||||
}
|
||||
|
||||
func TestWriteEasyAIAsyncErrorKeepsNestedGatewayError(t *testing.T) {
|
||||
recorder := httptest.NewRecorder()
|
||||
writeEasyAIAsyncError(recorder, http.StatusBadRequest, "invalid request", map[string]any{"param": "model"}, "invalid_parameter")
|
||||
|
||||
var got map[string]any
|
||||
if err := json.NewDecoder(recorder.Body).Decode(&got); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
errorPayload, _ := got["error"].(map[string]any)
|
||||
data, _ := got["data"].([]any)
|
||||
if got["status"] != "failed" || got["code"] != "invalid_parameter" ||
|
||||
errorPayload["message"] != "invalid request" || len(data) != 0 {
|
||||
t.Fatalf("unexpected async error response: %+v", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,911 @@
|
||||
package httpapi
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
)
|
||||
|
||||
type failurePolicyAcceptanceFixture struct {
|
||||
ctx context.Context
|
||||
cancel context.CancelFunc
|
||||
db *store.Store
|
||||
pool *pgxpool.Pool
|
||||
server *httptest.Server
|
||||
adminToken string
|
||||
apiKey string
|
||||
suffix string
|
||||
}
|
||||
|
||||
type controlledProviderStub struct {
|
||||
server *httptest.Server
|
||||
calls atomic.Int64
|
||||
}
|
||||
|
||||
type acceptancePlatform struct {
|
||||
ID string `json:"id"`
|
||||
}
|
||||
|
||||
type acceptancePlatformModel struct {
|
||||
ID string `json:"id"`
|
||||
}
|
||||
|
||||
type acceptanceTaskDetail struct {
|
||||
ID string `json:"id"`
|
||||
Status string `json:"status"`
|
||||
Attempts []struct {
|
||||
AttemptNo int `json:"attemptNo"`
|
||||
PlatformID string `json:"platformId"`
|
||||
PlatformName string `json:"platformName"`
|
||||
PlatformModel string `json:"platformModelId"`
|
||||
Status string `json:"status"`
|
||||
ErrorCode string `json:"errorCode"`
|
||||
DynamicMetrics map[string]any `json:"metrics"`
|
||||
Trace []map[string]any `json:"trace"`
|
||||
} `json:"attempts"`
|
||||
}
|
||||
|
||||
func TestFailurePolicyHTTPAcceptance(t *testing.T) {
|
||||
fixture := newFailurePolicyAcceptanceFixture(t)
|
||||
defer fixture.close()
|
||||
|
||||
originalPolicy, err := fixture.db.GetActiveRunnerPolicy(fixture.ctx)
|
||||
if err != nil {
|
||||
t.Fatalf("read original runner policy: %v", err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
_, _ = fixture.db.UpsertDefaultRunnerPolicy(context.Background(), store.RunnerPolicyInput{
|
||||
PolicyKey: originalPolicy.PolicyKey,
|
||||
Name: originalPolicy.Name,
|
||||
Description: originalPolicy.Description,
|
||||
FailoverPolicy: originalPolicy.FailoverPolicy,
|
||||
HardStopPolicy: originalPolicy.HardStopPolicy,
|
||||
PriorityDemotePolicy: originalPolicy.PriorityDemotePolicy,
|
||||
SingleSourcePolicy: originalPolicy.SingleSourcePolicy,
|
||||
CacheAffinityPolicy: originalPolicy.CacheAffinityPolicy,
|
||||
Metadata: originalPolicy.Metadata,
|
||||
Status: originalPolicy.Status,
|
||||
})
|
||||
})
|
||||
fixture.setRunnerPolicy(t, true)
|
||||
|
||||
t.Run("Gemini 429 后轮转且只冷却模型一次", fixture.testGemini429ImageFailover)
|
||||
t.Run("5xx 同平台重试耗尽后降权并轮转", fixture.testServerErrorRetryThenFailover)
|
||||
t.Run("401 立即禁用平台并轮转", fixture.testAuthFailureDisableThenFailover)
|
||||
t.Run("400 参数错误硬停止", fixture.testBadRequestStops)
|
||||
t.Run("旧 degradePolicy 仅转换一次", fixture.testLegacyDegradeDoesNotDoubleApply)
|
||||
t.Run("并发单源保护至少保留一个来源", fixture.testConcurrentSingleSourceProtection)
|
||||
t.Run("全部冷却同步返回脱敏 429", fixture.testAllCoolingSynchronousContract)
|
||||
t.Run("全部冷却异步零 attempt 重排并恢复", fixture.testAllCoolingAsyncRecovery)
|
||||
t.Run("流式首包边界控制轮转", fixture.testStreamingBoundary)
|
||||
t.Run("管理接口严格校验且不覆盖原策略", fixture.testRunnerPolicyHTTPValidation)
|
||||
}
|
||||
|
||||
func newFailurePolicyAcceptanceFixture(t *testing.T) *failurePolicyAcceptanceFixture {
|
||||
t.Helper()
|
||||
databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL"))
|
||||
if databaseURL == "" {
|
||||
t.Skip("set AI_GATEWAY_TEST_DATABASE_URL to run failure-policy HTTP acceptance")
|
||||
}
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
applyMigration(t, ctx, databaseURL)
|
||||
db, err := store.Connect(ctx, databaseURL)
|
||||
if err != nil {
|
||||
cancel()
|
||||
t.Fatalf("connect acceptance store: %v", err)
|
||||
}
|
||||
pool, err := pgxpool.New(ctx, databaseURL)
|
||||
if err != nil {
|
||||
db.Close()
|
||||
cancel()
|
||||
t.Fatalf("connect acceptance pool: %v", err)
|
||||
}
|
||||
handler := NewServerWithContext(ctx, config.Config{
|
||||
AppEnv: "test",
|
||||
HTTPAddr: ":0",
|
||||
DatabaseURL: databaseURL,
|
||||
IdentityMode: "hybrid",
|
||||
JWTSecret: "failure-policy-http-acceptance-secret",
|
||||
BillingEngineMode: "observe",
|
||||
CORSAllowedOrigin: "*",
|
||||
}, db, slog.New(slog.NewTextHandler(io.Discard, nil)))
|
||||
server := httptest.NewServer(handler)
|
||||
fixture := &failurePolicyAcceptanceFixture{
|
||||
ctx: ctx,
|
||||
cancel: cancel,
|
||||
db: db,
|
||||
pool: pool,
|
||||
server: server,
|
||||
suffix: strconv.FormatInt(time.Now().UnixNano(), 10),
|
||||
}
|
||||
|
||||
username := "failure_acceptance_" + fixture.suffix
|
||||
password := "password123"
|
||||
doJSON(t, server.URL, http.MethodPost, "/api/v1/auth/register", "", map[string]any{
|
||||
"username": username,
|
||||
"email": username + "@example.com",
|
||||
"password": password,
|
||||
}, http.StatusCreated, &map[string]any{})
|
||||
if _, err := pool.Exec(ctx, `UPDATE gateway_users SET roles = '["admin"]'::jsonb WHERE username = $1`, username); err != nil {
|
||||
fixture.close()
|
||||
t.Fatalf("promote acceptance admin: %v", err)
|
||||
}
|
||||
var login struct {
|
||||
AccessToken string `json:"accessToken"`
|
||||
}
|
||||
doJSON(t, server.URL, http.MethodPost, "/api/v1/auth/login", "", map[string]any{
|
||||
"account": username, "password": password,
|
||||
}, http.StatusOK, &login)
|
||||
fixture.adminToken = login.AccessToken
|
||||
var gatewayUserID string
|
||||
if err := pool.QueryRow(ctx, `SELECT id::text FROM gateway_users WHERE username=$1`, username).Scan(&gatewayUserID); err != nil {
|
||||
fixture.close()
|
||||
t.Fatalf("read acceptance user: %v", err)
|
||||
}
|
||||
doJSON(t, server.URL, http.MethodPatch, "/api/admin/users/"+gatewayUserID+"/wallet", fixture.adminToken, map[string]any{
|
||||
"currency": "resource",
|
||||
"balance": 1000000,
|
||||
"reason": "seed isolated failure-policy acceptance wallet",
|
||||
}, http.StatusOK, &map[string]any{})
|
||||
var key struct {
|
||||
Secret string `json:"secret"`
|
||||
}
|
||||
doJSON(t, server.URL, http.MethodPost, "/api/v1/api-keys", fixture.adminToken, map[string]any{
|
||||
"name": "failure policy acceptance",
|
||||
}, http.StatusCreated, &key)
|
||||
fixture.apiKey = key.Secret
|
||||
return fixture
|
||||
}
|
||||
|
||||
func (f *failurePolicyAcceptanceFixture) close() {
|
||||
if f.server != nil {
|
||||
f.server.Close()
|
||||
}
|
||||
if f.cancel != nil {
|
||||
f.cancel()
|
||||
}
|
||||
if f.pool != nil {
|
||||
f.pool.Close()
|
||||
}
|
||||
if f.db != nil {
|
||||
f.db.Close()
|
||||
}
|
||||
}
|
||||
|
||||
func (f *failurePolicyAcceptanceFixture) setRunnerPolicy(t *testing.T, singleSource bool) {
|
||||
t.Helper()
|
||||
doJSON(t, f.server.URL, http.MethodPatch, "/api/admin/runtime/runner-policy", f.adminToken, standardFailureRunnerPolicy(singleSource), http.StatusOK, &map[string]any{})
|
||||
}
|
||||
|
||||
func standardFailureRunnerPolicy(singleSource bool) map[string]any {
|
||||
return map[string]any{
|
||||
"policyKey": "default-runner-v1",
|
||||
"name": "失败策略 HTTP 验收",
|
||||
"status": "active",
|
||||
"description": "真实 provider HTTP client 的失败策略验收",
|
||||
"failoverPolicy": map[string]any{
|
||||
"enabled": true,
|
||||
"maxPlatforms": 99,
|
||||
"maxDurationSeconds": 600,
|
||||
"actionRules": []any{
|
||||
map[string]any{"categories": []any{"rate_limit"}, "statusCodes": []any{429}, "action": "cooldown_and_next", "target": "model", "cooldownSeconds": 30},
|
||||
map[string]any{"categories": []any{"auth_error"}, "statusCodes": []any{401, 403}, "action": "disable_and_next", "target": "platform"},
|
||||
map[string]any{"categories": []any{"timeout", "provider_5xx", "provider_overloaded"}, "statusCodes": []any{408, 500, 502, 503, 504}, "action": "demote_and_next", "target": "platform", "demoteSteps": 1},
|
||||
map[string]any{"categories": []any{"network", "stream_error"}, "action": "next"},
|
||||
map[string]any{"categories": []any{"request_error", "unsupported_model", "user_permission", "insufficient_balance"}, "statusCodes": []any{400, 404, 409, 422}, "action": "stop"},
|
||||
},
|
||||
},
|
||||
"singleSourcePolicy": map[string]any{"enabled": singleSource},
|
||||
"cacheAffinityPolicy": map[string]any{"enabled": false},
|
||||
}
|
||||
}
|
||||
|
||||
func (f *failurePolicyAcceptanceFixture) createBaseModel(t *testing.T, model string, modelType string) string {
|
||||
t.Helper()
|
||||
var created struct {
|
||||
ID string `json:"id"`
|
||||
}
|
||||
doJSON(t, f.server.URL, http.MethodPost, "/api/admin/catalog/base-models", f.adminToken, map[string]any{
|
||||
"providerKey": "openai",
|
||||
"canonicalModelKey": "openai:" + model,
|
||||
"invocationName": model,
|
||||
"providerModelName": model,
|
||||
"modelType": []string{modelType},
|
||||
"displayName": model,
|
||||
"status": "active",
|
||||
"capabilities": map[string]any{},
|
||||
"baseBillingConfig": map[string]any{},
|
||||
}, http.StatusCreated, &created)
|
||||
return created.ID
|
||||
}
|
||||
|
||||
func (f *failurePolicyAcceptanceFixture) createPlatformModel(
|
||||
t *testing.T,
|
||||
name string,
|
||||
provider string,
|
||||
baseURL string,
|
||||
priority int,
|
||||
baseModelID string,
|
||||
providerModelName string,
|
||||
modelType string,
|
||||
retryAttempts int,
|
||||
runtimePolicySetID string,
|
||||
runtimeOverride map[string]any,
|
||||
) (acceptancePlatform, acceptancePlatformModel) {
|
||||
t.Helper()
|
||||
var platform acceptancePlatform
|
||||
doJSON(t, f.server.URL, http.MethodPost, "/api/admin/platforms", f.adminToken, map[string]any{
|
||||
"provider": provider,
|
||||
"platformKey": "acceptance-" + name + "-" + f.suffix,
|
||||
"name": "acceptance-" + name,
|
||||
"baseUrl": strings.TrimRight(baseURL, "/"),
|
||||
"authType": "bearer",
|
||||
"credentials": map[string]any{"apiKey": "controlled-upstream-test-key"},
|
||||
"priority": priority,
|
||||
"status": "enabled",
|
||||
}, http.StatusCreated, &platform)
|
||||
var model acceptancePlatformModel
|
||||
payload := map[string]any{
|
||||
"baseModelId": baseModelID,
|
||||
"providerModelName": providerModelName,
|
||||
"modelType": []string{modelType},
|
||||
"retryPolicy": map[string]any{"enabled": true, "maxAttempts": retryAttempts},
|
||||
}
|
||||
if runtimePolicySetID != "" {
|
||||
payload["runtimePolicySetId"] = runtimePolicySetID
|
||||
}
|
||||
if len(runtimeOverride) > 0 {
|
||||
payload["runtimePolicyOverride"] = runtimeOverride
|
||||
}
|
||||
doJSON(t, f.server.URL, http.MethodPost, "/api/admin/platforms/"+platform.ID+"/models", f.adminToken, payload, http.StatusCreated, &model)
|
||||
return platform, model
|
||||
}
|
||||
|
||||
func newControlledProviderStub(handler func(call int, w http.ResponseWriter, r *http.Request)) *controlledProviderStub {
|
||||
stub := &controlledProviderStub{}
|
||||
stub.server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
call := int(stub.calls.Add(1))
|
||||
handler(call, w, r)
|
||||
}))
|
||||
return stub
|
||||
}
|
||||
|
||||
func (s *controlledProviderStub) close() {
|
||||
if s != nil && s.server != nil {
|
||||
s.server.Close()
|
||||
}
|
||||
}
|
||||
|
||||
func writeStubJSON(w http.ResponseWriter, status int, payload any) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(status)
|
||||
_ = json.NewEncoder(w).Encode(payload)
|
||||
}
|
||||
|
||||
func writeOpenAIChatSuccess(w http.ResponseWriter, content string) {
|
||||
writeStubJSON(w, http.StatusOK, map[string]any{
|
||||
"id": "chatcmpl-controlled", "object": "chat.completion", "created": time.Now().Unix(), "model": "controlled-model",
|
||||
"choices": []any{map[string]any{
|
||||
"index": 0, "message": map[string]any{"role": "assistant", "content": content}, "finish_reason": "stop",
|
||||
}},
|
||||
"usage": map[string]any{"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
|
||||
})
|
||||
}
|
||||
|
||||
func writeOpenAIImageSuccess(w http.ResponseWriter) {
|
||||
writeStubJSON(w, http.StatusOK, map[string]any{
|
||||
"created": time.Now().Unix(),
|
||||
"data": []any{map[string]any{"url": "https://example.invalid/controlled-test.png"}},
|
||||
})
|
||||
}
|
||||
|
||||
func (f *failurePolicyAcceptanceFixture) postJSON(t *testing.T, path string, payload any, headers map[string]string) (int, http.Header, []byte) {
|
||||
t.Helper()
|
||||
raw, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
t.Fatalf("marshal acceptance payload: %v", err)
|
||||
}
|
||||
request, err := http.NewRequest(http.MethodPost, f.server.URL+path, bytes.NewReader(raw))
|
||||
if err != nil {
|
||||
t.Fatalf("build acceptance request: %v", err)
|
||||
}
|
||||
request.Header.Set("Authorization", "Bearer "+f.apiKey)
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
for key, value := range headers {
|
||||
request.Header.Set(key, value)
|
||||
}
|
||||
response, err := http.DefaultClient.Do(request)
|
||||
if err != nil {
|
||||
t.Fatalf("acceptance request %s: %v", path, err)
|
||||
}
|
||||
defer response.Body.Close()
|
||||
body, _ := io.ReadAll(response.Body)
|
||||
return response.StatusCode, response.Header.Clone(), body
|
||||
}
|
||||
|
||||
func (f *failurePolicyAcceptanceFixture) taskIDForModel(t *testing.T, model string, startedAt time.Time) string {
|
||||
t.Helper()
|
||||
var taskID string
|
||||
deadline := time.Now().Add(3 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
err := f.pool.QueryRow(f.ctx, `
|
||||
SELECT id::text
|
||||
FROM gateway_tasks
|
||||
WHERE model = $1
|
||||
AND created_at >= $2
|
||||
ORDER BY created_at DESC
|
||||
LIMIT 1`, model, startedAt.Add(-time.Second)).Scan(&taskID)
|
||||
if err == nil && taskID != "" {
|
||||
return taskID
|
||||
}
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
}
|
||||
t.Fatalf("task for model %q was not persisted", model)
|
||||
return ""
|
||||
}
|
||||
|
||||
func (f *failurePolicyAcceptanceFixture) loadTask(t *testing.T, taskID string) acceptanceTaskDetail {
|
||||
t.Helper()
|
||||
var detail acceptanceTaskDetail
|
||||
doJSON(t, f.server.URL, http.MethodGet, "/api/v1/tasks/"+taskID, f.apiKey, nil, http.StatusOK, &detail)
|
||||
return detail
|
||||
}
|
||||
|
||||
func countFailureEffects(detail acceptanceTaskDetail, effect string) int {
|
||||
count := 0
|
||||
for _, attempt := range detail.Attempts {
|
||||
trace := attempt.Trace
|
||||
for _, raw := range anySlice(attempt.DynamicMetrics["trace"]) {
|
||||
if entry, ok := raw.(map[string]any); ok {
|
||||
trace = append(trace, entry)
|
||||
}
|
||||
}
|
||||
for _, entry := range trace {
|
||||
if entry["event"] == "failure_decision" && entry["effect"] == effect {
|
||||
count++
|
||||
}
|
||||
}
|
||||
}
|
||||
return count
|
||||
}
|
||||
|
||||
func anySlice(value any) []any {
|
||||
items, _ := value.([]any)
|
||||
return items
|
||||
}
|
||||
|
||||
func (f *failurePolicyAcceptanceFixture) testGemini429ImageFailover(t *testing.T) {
|
||||
model := "acceptance-gemini-image-" + f.suffix
|
||||
baseModelID := f.createBaseModel(t, model, "image_generate")
|
||||
gemini := newControlledProviderStub(func(_ int, w http.ResponseWriter, _ *http.Request) {
|
||||
writeStubJSON(w, http.StatusTooManyRequests, map[string]any{
|
||||
"error": map[string]any{"code": 429, "message": "RESOURCE_EXHAUSTED controlled failure", "status": "RESOURCE_EXHAUSTED"},
|
||||
})
|
||||
})
|
||||
defer gemini.close()
|
||||
openai := newControlledProviderStub(func(_ int, w http.ResponseWriter, _ *http.Request) { writeOpenAIImageSuccess(w) })
|
||||
defer openai.close()
|
||||
platformA, modelA := f.createPlatformModel(t, "gemini-image-a", "gemini", gemini.server.URL, 10, baseModelID, "gemini-image-controlled", "image_generate", 3, "", nil)
|
||||
platformB, _ := f.createPlatformModel(t, "openai-image-b", "openai", openai.server.URL+"/v1", 20, baseModelID, "openai-image-controlled", "image_generate", 1, "", nil)
|
||||
|
||||
startedAt := time.Now()
|
||||
status, _, body := f.postJSON(t, "/api/v1/images/generations", map[string]any{"model": model, "prompt": "controlled image"}, nil)
|
||||
if status != http.StatusOK {
|
||||
t.Fatalf("status=%d body=%s", status, body)
|
||||
}
|
||||
if gemini.calls.Load() != 1 || openai.calls.Load() != 1 {
|
||||
t.Fatalf("upstream calls gemini=%d openai=%d, want 1/1", gemini.calls.Load(), openai.calls.Load())
|
||||
}
|
||||
var modelCooling, platformCooling bool
|
||||
if err := f.pool.QueryRow(f.ctx, `
|
||||
SELECT COALESCE(model.cooldown_until > now(), false), COALESCE(platform.cooldown_until > now(), false)
|
||||
FROM platform_models model
|
||||
JOIN integration_platforms platform ON platform.id = model.platform_id
|
||||
WHERE model.id = $1::uuid`, modelA.ID).Scan(&modelCooling, &platformCooling); err != nil {
|
||||
t.Fatalf("read Gemini cooldown state: %v", err)
|
||||
}
|
||||
if !modelCooling || platformCooling {
|
||||
t.Fatalf("cooldown state model=%v platform=%v, want true/false for platform %s", modelCooling, platformCooling, platformA.ID)
|
||||
}
|
||||
detail := f.loadTask(t, f.taskIDForModel(t, model, startedAt))
|
||||
if len(detail.Attempts) != 2 || detail.Attempts[0].PlatformID != platformA.ID || detail.Attempts[1].PlatformID != platformB.ID || countFailureEffects(detail, "cooldown") != 1 {
|
||||
t.Fatalf("unexpected task attempts/trace: %+v", detail)
|
||||
}
|
||||
}
|
||||
|
||||
func (f *failurePolicyAcceptanceFixture) testServerErrorRetryThenFailover(t *testing.T) {
|
||||
model := "acceptance-5xx-" + f.suffix
|
||||
baseModelID := f.createBaseModel(t, model, "text_generate")
|
||||
failed := newControlledProviderStub(func(_ int, w http.ResponseWriter, _ *http.Request) {
|
||||
writeStubJSON(w, http.StatusInternalServerError, map[string]any{"error": map[string]any{"message": "controlled 500"}})
|
||||
})
|
||||
defer failed.close()
|
||||
success := newControlledProviderStub(func(_ int, w http.ResponseWriter, _ *http.Request) { writeOpenAIChatSuccess(w, "fallback ok") })
|
||||
defer success.close()
|
||||
platformA, _ := f.createPlatformModel(t, "5xx-a", "openai", failed.server.URL+"/v1", 10, baseModelID, "controlled-5xx", "text_generate", 2, "", nil)
|
||||
platformB, _ := f.createPlatformModel(t, "5xx-b", "openai", success.server.URL+"/v1", 20, baseModelID, "controlled-success", "text_generate", 1, "", nil)
|
||||
|
||||
startedAt := time.Now()
|
||||
status, _, body := f.postJSON(t, "/api/v1/chat/completions", map[string]any{
|
||||
"model": model, "messages": []any{map[string]any{"role": "user", "content": "retry"}},
|
||||
}, nil)
|
||||
if status != http.StatusOK {
|
||||
t.Fatalf("status=%d body=%s", status, body)
|
||||
}
|
||||
if failed.calls.Load() != 2 || success.calls.Load() != 1 {
|
||||
t.Fatalf("upstream calls failed=%d success=%d, want 2/1", failed.calls.Load(), success.calls.Load())
|
||||
}
|
||||
detail := f.loadTask(t, f.taskIDForModel(t, model, startedAt))
|
||||
if len(detail.Attempts) != 3 ||
|
||||
detail.Attempts[0].PlatformID != platformA.ID ||
|
||||
detail.Attempts[1].PlatformID != platformA.ID ||
|
||||
detail.Attempts[2].PlatformID != platformB.ID ||
|
||||
countFailureEffects(detail, "demote") != 1 {
|
||||
t.Fatalf("unexpected retry/failover chain: %+v", detail)
|
||||
}
|
||||
var dynamicPriority *int
|
||||
if err := f.pool.QueryRow(f.ctx, `SELECT dynamic_priority FROM integration_platforms WHERE id=$1::uuid`, platformA.ID).Scan(&dynamicPriority); err != nil {
|
||||
t.Fatalf("read demoted priority: %v", err)
|
||||
}
|
||||
if dynamicPriority == nil {
|
||||
t.Fatal("failed platform should be demoted exactly once after retry exhaustion")
|
||||
}
|
||||
}
|
||||
|
||||
func (f *failurePolicyAcceptanceFixture) testAuthFailureDisableThenFailover(t *testing.T) {
|
||||
model := "acceptance-auth-" + f.suffix
|
||||
baseModelID := f.createBaseModel(t, model, "text_generate")
|
||||
failed := newControlledProviderStub(func(_ int, w http.ResponseWriter, _ *http.Request) {
|
||||
writeStubJSON(w, http.StatusUnauthorized, map[string]any{"error": map[string]any{"message": "invalid api key"}})
|
||||
})
|
||||
defer failed.close()
|
||||
success := newControlledProviderStub(func(_ int, w http.ResponseWriter, _ *http.Request) { writeOpenAIChatSuccess(w, "auth fallback") })
|
||||
defer success.close()
|
||||
platformA, _ := f.createPlatformModel(t, "auth-a", "openai", failed.server.URL+"/v1", 10, baseModelID, "controlled-auth", "text_generate", 3, "", nil)
|
||||
_, _ = f.createPlatformModel(t, "auth-b", "openai", success.server.URL+"/v1", 20, baseModelID, "controlled-success", "text_generate", 1, "", nil)
|
||||
|
||||
startedAt := time.Now()
|
||||
status, _, body := f.postJSON(t, "/api/v1/chat/completions", map[string]any{
|
||||
"model": model, "messages": []any{map[string]any{"role": "user", "content": "auth"}},
|
||||
}, nil)
|
||||
if status != http.StatusOK {
|
||||
t.Fatalf("status=%d body=%s", status, body)
|
||||
}
|
||||
if failed.calls.Load() != 1 || success.calls.Load() != 1 {
|
||||
t.Fatalf("upstream calls failed=%d success=%d, want 1/1", failed.calls.Load(), success.calls.Load())
|
||||
}
|
||||
var statusValue string
|
||||
if err := f.pool.QueryRow(f.ctx, `SELECT status FROM integration_platforms WHERE id=$1::uuid`, platformA.ID).Scan(&statusValue); err != nil {
|
||||
t.Fatalf("read disabled platform: %v", err)
|
||||
}
|
||||
detail := f.loadTask(t, f.taskIDForModel(t, model, startedAt))
|
||||
if statusValue != "disabled" || countFailureEffects(detail, "disable") != 1 || len(detail.Attempts) != 2 {
|
||||
t.Fatalf("unexpected auth disable result status=%s detail=%+v", statusValue, detail)
|
||||
}
|
||||
}
|
||||
|
||||
func (f *failurePolicyAcceptanceFixture) testBadRequestStops(t *testing.T) {
|
||||
model := "acceptance-400-" + f.suffix
|
||||
baseModelID := f.createBaseModel(t, model, "text_generate")
|
||||
failed := newControlledProviderStub(func(_ int, w http.ResponseWriter, _ *http.Request) {
|
||||
writeStubJSON(w, http.StatusBadRequest, map[string]any{"error": map[string]any{"message": "controlled invalid parameter", "code": "invalid_parameter"}})
|
||||
})
|
||||
defer failed.close()
|
||||
success := newControlledProviderStub(func(_ int, w http.ResponseWriter, _ *http.Request) { writeOpenAIChatSuccess(w, "must not run") })
|
||||
defer success.close()
|
||||
platformA, modelA := f.createPlatformModel(t, "400-a", "openai", failed.server.URL+"/v1", 10, baseModelID, "controlled-bad", "text_generate", 3, "", nil)
|
||||
_, _ = f.createPlatformModel(t, "400-b", "openai", success.server.URL+"/v1", 20, baseModelID, "controlled-success", "text_generate", 1, "", nil)
|
||||
|
||||
status, _, body := f.postJSON(t, "/api/v1/chat/completions", map[string]any{
|
||||
"model": model, "messages": []any{map[string]any{"role": "user", "content": "bad"}},
|
||||
}, nil)
|
||||
if status != http.StatusBadRequest || success.calls.Load() != 0 || failed.calls.Load() != 1 {
|
||||
t.Fatalf("status=%d calls=%d/%d body=%s", status, failed.calls.Load(), success.calls.Load(), body)
|
||||
}
|
||||
var platformStatus string
|
||||
var modelEnabled bool
|
||||
var platformCooldown, modelCooldown *time.Time
|
||||
if err := f.pool.QueryRow(f.ctx, `
|
||||
SELECT platform.status, model.enabled, platform.cooldown_until, model.cooldown_until
|
||||
FROM integration_platforms platform
|
||||
JOIN platform_models model ON model.platform_id=platform.id
|
||||
WHERE platform.id=$1::uuid AND model.id=$2::uuid`, platformA.ID, modelA.ID).Scan(&platformStatus, &modelEnabled, &platformCooldown, &modelCooldown); err != nil {
|
||||
t.Fatalf("read hard-stop state: %v", err)
|
||||
}
|
||||
if platformStatus != "enabled" || !modelEnabled || platformCooldown != nil || modelCooldown != nil {
|
||||
t.Fatalf("hard stop mutated runtime state: status=%s enabled=%v platformCooldown=%v modelCooldown=%v", platformStatus, modelEnabled, platformCooldown, modelCooldown)
|
||||
}
|
||||
}
|
||||
|
||||
func (f *failurePolicyAcceptanceFixture) testLegacyDegradeDoesNotDoubleApply(t *testing.T) {
|
||||
model := "acceptance-legacy-" + f.suffix
|
||||
baseModelID := f.createBaseModel(t, model, "text_generate")
|
||||
var policySet struct {
|
||||
ID string `json:"id"`
|
||||
}
|
||||
doJSON(t, f.server.URL, http.MethodPost, "/api/admin/runtime/policy-sets", f.adminToken, map[string]any{
|
||||
"policyKey": "acceptance-legacy-" + f.suffix,
|
||||
"name": "acceptance legacy degrade",
|
||||
"degradePolicy": map[string]any{
|
||||
"enabled": true, "keywords": []any{"rate_limit"}, "cooldownSeconds": 45,
|
||||
},
|
||||
}, http.StatusCreated, &policySet)
|
||||
failed := newControlledProviderStub(func(_ int, w http.ResponseWriter, _ *http.Request) {
|
||||
writeStubJSON(w, http.StatusTooManyRequests, map[string]any{"error": map[string]any{"message": "rate_limit controlled"}})
|
||||
})
|
||||
defer failed.close()
|
||||
success := newControlledProviderStub(func(_ int, w http.ResponseWriter, _ *http.Request) { writeOpenAIChatSuccess(w, "legacy fallback") })
|
||||
defer success.close()
|
||||
_, modelA := f.createPlatformModel(t, "legacy-a", "openai", failed.server.URL+"/v1", 10, baseModelID, "legacy-failed", "text_generate", 3, policySet.ID, nil)
|
||||
_, _ = f.createPlatformModel(t, "legacy-b", "openai", success.server.URL+"/v1", 20, baseModelID, "legacy-success", "text_generate", 1, "", nil)
|
||||
|
||||
startedAt := time.Now()
|
||||
status, _, body := f.postJSON(t, "/api/v1/chat/completions", map[string]any{
|
||||
"model": model, "messages": []any{map[string]any{"role": "user", "content": "legacy"}},
|
||||
}, nil)
|
||||
if status != http.StatusOK || failed.calls.Load() != 1 || success.calls.Load() != 1 {
|
||||
t.Fatalf("status=%d calls=%d/%d body=%s", status, failed.calls.Load(), success.calls.Load(), body)
|
||||
}
|
||||
detail := f.loadTask(t, f.taskIDForModel(t, model, startedAt))
|
||||
if countFailureEffects(detail, "cooldown") != 1 {
|
||||
t.Fatalf("legacy and global policy must not both execute: %+v", detail)
|
||||
}
|
||||
var remaining float64
|
||||
if err := f.pool.QueryRow(f.ctx, `SELECT EXTRACT(EPOCH FROM cooldown_until-now())::float8 FROM platform_models WHERE id=$1::uuid`, modelA.ID).Scan(&remaining); err != nil {
|
||||
t.Fatalf("read legacy cooldown: %v", err)
|
||||
}
|
||||
if remaining < 35 || remaining > 50 {
|
||||
t.Fatalf("legacy cooldown should be applied once with 45 seconds, remaining=%f", remaining)
|
||||
}
|
||||
}
|
||||
|
||||
func (f *failurePolicyAcceptanceFixture) testConcurrentSingleSourceProtection(t *testing.T) {
|
||||
f.setRunnerPolicy(t, true)
|
||||
model := "acceptance-concurrent-" + f.suffix
|
||||
baseModelID := f.createBaseModel(t, model, "text_generate")
|
||||
stubs := make([]*controlledProviderStub, 0, 3)
|
||||
modelIDs := make([]string, 0, 3)
|
||||
for index := 0; index < 3; index++ {
|
||||
stub := newControlledProviderStub(func(_ int, w http.ResponseWriter, _ *http.Request) {
|
||||
writeStubJSON(w, http.StatusTooManyRequests, map[string]any{"error": map[string]any{"message": "controlled 429"}})
|
||||
})
|
||||
stubs = append(stubs, stub)
|
||||
_, platformModel := f.createPlatformModel(t, fmt.Sprintf("concurrent-%d", index), "openai", stub.server.URL+"/v1", 10+index*10, baseModelID, fmt.Sprintf("concurrent-%d", index), "text_generate", 1, "", nil)
|
||||
modelIDs = append(modelIDs, platformModel.ID)
|
||||
}
|
||||
defer func() {
|
||||
for _, stub := range stubs {
|
||||
stub.close()
|
||||
}
|
||||
}()
|
||||
|
||||
var wait sync.WaitGroup
|
||||
errorsFound := make(chan error, 12)
|
||||
for index := 0; index < 12; index++ {
|
||||
wait.Add(1)
|
||||
go func() {
|
||||
defer wait.Done()
|
||||
status, _, body, err := f.postJSONNoFail("/api/v1/chat/completions", map[string]any{
|
||||
"model": model, "messages": []any{map[string]any{"role": "user", "content": "concurrent"}},
|
||||
})
|
||||
if err != nil {
|
||||
errorsFound <- err
|
||||
return
|
||||
}
|
||||
if status != http.StatusTooManyRequests {
|
||||
errorsFound <- fmt.Errorf("status=%d body=%s", status, body)
|
||||
}
|
||||
}()
|
||||
}
|
||||
wait.Wait()
|
||||
close(errorsFound)
|
||||
for err := range errorsFound {
|
||||
t.Error(err)
|
||||
}
|
||||
|
||||
var available int
|
||||
if err := f.pool.QueryRow(f.ctx, `
|
||||
SELECT count(*)
|
||||
FROM platform_models model
|
||||
JOIN integration_platforms platform ON platform.id=model.platform_id
|
||||
WHERE model.id = ANY($1::uuid[])
|
||||
AND model.enabled=true
|
||||
AND platform.status='enabled'
|
||||
AND (model.cooldown_until IS NULL OR model.cooldown_until <= now())
|
||||
AND (platform.cooldown_until IS NULL OR platform.cooldown_until <= now())`, modelIDs).Scan(&available); err != nil {
|
||||
t.Fatalf("count available sources: %v", err)
|
||||
}
|
||||
if available < 1 {
|
||||
t.Fatal("concurrent cooldown decisions disabled or cooled every source")
|
||||
}
|
||||
before := make([]int64, len(stubs))
|
||||
for index, stub := range stubs {
|
||||
before[index] = stub.calls.Load()
|
||||
}
|
||||
status, _, _, err := f.postJSONNoFail("/api/v1/chat/completions", map[string]any{
|
||||
"model": model, "messages": []any{map[string]any{"role": "user", "content": "post-cooldown"}},
|
||||
})
|
||||
if err != nil || status != http.StatusTooManyRequests {
|
||||
t.Fatalf("post-cooldown request status=%d err=%v", status, err)
|
||||
}
|
||||
for index, modelID := range modelIDs {
|
||||
var cooled bool
|
||||
if err := f.pool.QueryRow(f.ctx, `SELECT COALESCE(cooldown_until > now(), false) FROM platform_models WHERE id=$1::uuid`, modelID).Scan(&cooled); err != nil {
|
||||
t.Fatalf("read source cooldown: %v", err)
|
||||
}
|
||||
if cooled && stubs[index].calls.Load() != before[index] {
|
||||
t.Fatalf("cooled source %d was called again: before=%d after=%d", index, before[index], stubs[index].calls.Load())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (f *failurePolicyAcceptanceFixture) postJSONNoFail(path string, payload any) (int, http.Header, []byte, error) {
|
||||
raw, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
return 0, nil, nil, err
|
||||
}
|
||||
request, err := http.NewRequest(http.MethodPost, f.server.URL+path, bytes.NewReader(raw))
|
||||
if err != nil {
|
||||
return 0, nil, nil, err
|
||||
}
|
||||
request.Header.Set("Authorization", "Bearer "+f.apiKey)
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
response, err := http.DefaultClient.Do(request)
|
||||
if err != nil {
|
||||
return 0, nil, nil, err
|
||||
}
|
||||
defer response.Body.Close()
|
||||
body, err := io.ReadAll(response.Body)
|
||||
return response.StatusCode, response.Header.Clone(), body, err
|
||||
}
|
||||
|
||||
func (f *failurePolicyAcceptanceFixture) testAllCoolingSynchronousContract(t *testing.T) {
|
||||
f.setRunnerPolicy(t, false)
|
||||
defer f.setRunnerPolicy(t, true)
|
||||
model := "acceptance-all-cooling-" + f.suffix
|
||||
baseModelID := f.createBaseModel(t, model, "image_generate")
|
||||
stubA := newControlledProviderStub(func(_ int, w http.ResponseWriter, _ *http.Request) {
|
||||
writeStubJSON(w, http.StatusTooManyRequests, map[string]any{"error": map[string]any{"message": "private-a-raw-error"}})
|
||||
})
|
||||
defer stubA.close()
|
||||
stubB := newControlledProviderStub(func(_ int, w http.ResponseWriter, _ *http.Request) {
|
||||
writeStubJSON(w, http.StatusTooManyRequests, map[string]any{"error": map[string]any{"message": "private-b-raw-error"}})
|
||||
})
|
||||
defer stubB.close()
|
||||
override := func(seconds int) map[string]any {
|
||||
return map[string]any{"failoverPolicy": map[string]any{"actionRules": []any{
|
||||
map[string]any{"statusCodes": []any{429}, "action": "cooldown_and_next", "target": "model", "cooldownSeconds": seconds},
|
||||
}}}
|
||||
}
|
||||
_, modelA := f.createPlatformModel(t, "cooling-a", "openai", stubA.server.URL+"/v1", 10, baseModelID, "cooling-a", "image_generate", 3, "", override(30))
|
||||
_, modelB := f.createPlatformModel(t, "cooling-b", "openai", stubB.server.URL+"/v1", 20, baseModelID, "cooling-b", "image_generate", 3, "", override(60))
|
||||
defer func() {
|
||||
_, _ = f.pool.Exec(context.Background(), `UPDATE platform_models SET cooldown_until=NULL WHERE id IN ($1::uuid,$2::uuid)`, modelA.ID, modelB.ID)
|
||||
}()
|
||||
|
||||
status, _, _ := f.postJSON(t, "/api/v1/images/generations", map[string]any{"model": model, "prompt": "cool both"}, nil)
|
||||
if status != http.StatusTooManyRequests || stubA.calls.Load() != 1 || stubB.calls.Load() != 1 {
|
||||
t.Fatalf("initial cooling request status=%d calls=%d/%d", status, stubA.calls.Load(), stubB.calls.Load())
|
||||
}
|
||||
startedAt := time.Now()
|
||||
status, headers, body := f.postJSON(t, "/api/v1/images/generations", map[string]any{"model": model, "prompt": "all cooling"}, nil)
|
||||
if status != http.StatusTooManyRequests {
|
||||
t.Fatalf("cooldown contract status=%d body=%s", status, body)
|
||||
}
|
||||
retryAfter, err := strconv.Atoi(headers.Get("Retry-After"))
|
||||
if err != nil || retryAfter < 20 || retryAfter > 31 {
|
||||
t.Fatalf("Retry-After=%q, want earliest recovery near 30 seconds", headers.Get("Retry-After"))
|
||||
}
|
||||
bodyText := string(body)
|
||||
for _, forbidden := range []string{"acceptance-cooling-a", "acceptance-cooling-b", stubA.server.URL, stubB.server.URL, "private-a-raw-error", "private-b-raw-error"} {
|
||||
if strings.Contains(bodyText, forbidden) {
|
||||
t.Fatalf("cooldown response leaked %q: %s", forbidden, bodyText)
|
||||
}
|
||||
}
|
||||
var envelope struct {
|
||||
Error struct {
|
||||
Code string `json:"code"`
|
||||
Details map[string]any `json:"details"`
|
||||
} `json:"error"`
|
||||
}
|
||||
if err := json.Unmarshal(body, &envelope); err != nil {
|
||||
t.Fatalf("decode cooldown response: %v body=%s", err, body)
|
||||
}
|
||||
if envelope.Error.Code != "model_cooling_down" || len(envelope.Error.Details) != 2 {
|
||||
t.Fatalf("unexpected cooldown envelope: %+v", envelope)
|
||||
}
|
||||
detail := f.loadTask(t, f.taskIDForModel(t, model, startedAt))
|
||||
if len(detail.Attempts) != 0 {
|
||||
t.Fatalf("candidate selection cooldown must not create provider attempts: %+v", detail.Attempts)
|
||||
}
|
||||
}
|
||||
|
||||
func (f *failurePolicyAcceptanceFixture) testAllCoolingAsyncRecovery(t *testing.T) {
|
||||
f.setRunnerPolicy(t, false)
|
||||
defer f.setRunnerPolicy(t, true)
|
||||
model := "acceptance-async-cooling-" + f.suffix
|
||||
baseModelID := f.createBaseModel(t, model, "image_generate")
|
||||
stubA := newControlledProviderStub(func(call int, w http.ResponseWriter, _ *http.Request) {
|
||||
if call == 1 {
|
||||
writeStubJSON(w, http.StatusTooManyRequests, map[string]any{"error": map[string]any{"message": "controlled first 429"}})
|
||||
return
|
||||
}
|
||||
writeOpenAIImageSuccess(w)
|
||||
})
|
||||
defer stubA.close()
|
||||
stubB := newControlledProviderStub(func(_ int, w http.ResponseWriter, _ *http.Request) {
|
||||
writeStubJSON(w, http.StatusTooManyRequests, map[string]any{"error": map[string]any{"message": "controlled second 429"}})
|
||||
})
|
||||
defer stubB.close()
|
||||
override := func(seconds int) map[string]any {
|
||||
return map[string]any{"failoverPolicy": map[string]any{"actionRules": []any{
|
||||
map[string]any{"statusCodes": []any{429}, "action": "cooldown_and_next", "target": "model", "cooldownSeconds": seconds},
|
||||
}}}
|
||||
}
|
||||
_, modelA := f.createPlatformModel(t, "async-cooling-a", "openai", stubA.server.URL+"/v1", 10, baseModelID, "async-a", "image_generate", 1, "", override(4))
|
||||
_, modelB := f.createPlatformModel(t, "async-cooling-b", "openai", stubB.server.URL+"/v1", 20, baseModelID, "async-b", "image_generate", 1, "", override(8))
|
||||
defer func() {
|
||||
_, _ = f.pool.Exec(context.Background(), `UPDATE platform_models SET cooldown_until=NULL WHERE id IN ($1::uuid,$2::uuid)`, modelA.ID, modelB.ID)
|
||||
}()
|
||||
status, _, _ := f.postJSON(t, "/api/v1/images/generations", map[string]any{"model": model, "prompt": "prime cooldown"}, nil)
|
||||
if status != http.StatusTooManyRequests {
|
||||
t.Fatalf("prime cooldown status=%d", status)
|
||||
}
|
||||
|
||||
startedAt := time.Now()
|
||||
status, _, body := f.postJSON(t, "/api/v1/images/generations", map[string]any{"model": model, "prompt": "async recovery"}, map[string]string{"X-Async": "true"})
|
||||
if status != http.StatusAccepted {
|
||||
t.Fatalf("async submit status=%d body=%s", status, body)
|
||||
}
|
||||
taskID := f.taskIDForModel(t, model, startedAt)
|
||||
deadline := time.Now().Add(2 * time.Second)
|
||||
queuedObserved := false
|
||||
for time.Now().Before(deadline) {
|
||||
detail := f.loadTask(t, taskID)
|
||||
if detail.Status == "queued" && len(detail.Attempts) == 0 {
|
||||
queuedObserved = true
|
||||
break
|
||||
}
|
||||
time.Sleep(25 * time.Millisecond)
|
||||
}
|
||||
if !queuedObserved {
|
||||
t.Fatal("async cooling task was not observed queued with zero provider attempts")
|
||||
}
|
||||
deadline = time.Now().Add(12 * time.Second)
|
||||
var completed acceptanceTaskDetail
|
||||
for time.Now().Before(deadline) {
|
||||
completed = f.loadTask(t, taskID)
|
||||
if completed.Status == "succeeded" {
|
||||
break
|
||||
}
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
}
|
||||
if completed.Status != "succeeded" || len(completed.Attempts) != 1 || stubA.calls.Load() != 2 {
|
||||
t.Fatalf("async task did not recover on first restored candidate: task=%+v callsA=%d callsB=%d", completed, stubA.calls.Load(), stubB.calls.Load())
|
||||
}
|
||||
}
|
||||
|
||||
func (f *failurePolicyAcceptanceFixture) testStreamingBoundary(t *testing.T) {
|
||||
t.Run("首个 delta 前允许轮转", func(t *testing.T) {
|
||||
model := "acceptance-stream-before-" + f.suffix
|
||||
baseModelID := f.createBaseModel(t, model, "text_generate")
|
||||
failed := newControlledProviderStub(func(_ int, w http.ResponseWriter, _ *http.Request) {
|
||||
writeStubJSON(w, http.StatusInternalServerError, map[string]any{"error": map[string]any{"message": "before delta"}})
|
||||
})
|
||||
defer failed.close()
|
||||
success := newControlledProviderStub(func(_ int, w http.ResponseWriter, _ *http.Request) {
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
_, _ = io.WriteString(w, "data: {\"id\":\"chunk-ok\",\"object\":\"chat.completion.chunk\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"content\":\"stream fallback\"},\"finish_reason\":null}]}\n\n")
|
||||
_, _ = io.WriteString(w, "data: {\"id\":\"chunk-ok\",\"object\":\"chat.completion.chunk\",\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n")
|
||||
_, _ = io.WriteString(w, "data: [DONE]\n\n")
|
||||
})
|
||||
defer success.close()
|
||||
_, _ = f.createPlatformModel(t, "stream-before-a", "openai", failed.server.URL+"/v1", 10, baseModelID, "stream-failed", "text_generate", 1, "", nil)
|
||||
_, _ = f.createPlatformModel(t, "stream-before-b", "openai", success.server.URL+"/v1", 20, baseModelID, "stream-success", "text_generate", 1, "", nil)
|
||||
|
||||
status, _, body := f.postJSON(t, "/api/v1/chat/completions", map[string]any{
|
||||
"model": model, "stream": true, "messages": []any{map[string]any{"role": "user", "content": "stream"}},
|
||||
}, nil)
|
||||
if status != http.StatusOK || failed.calls.Load() != 1 || success.calls.Load() != 1 || strings.Count(string(body), "stream fallback") != 1 || !strings.Contains(string(body), "[DONE]") {
|
||||
t.Fatalf("unexpected pre-delta failover status=%d calls=%d/%d body=%s", status, failed.calls.Load(), success.calls.Load(), body)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("已输出 delta 后禁止轮转", func(t *testing.T) {
|
||||
model := "acceptance-stream-after-" + f.suffix
|
||||
baseModelID := f.createBaseModel(t, model, "text_generate")
|
||||
broken := newControlledProviderStub(func(_ int, w http.ResponseWriter, _ *http.Request) {
|
||||
hijacker, ok := w.(http.Hijacker)
|
||||
if !ok {
|
||||
http.Error(w, "hijacking unavailable", http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
connection, buffer, err := hijacker.Hijack()
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
defer connection.Close()
|
||||
_, _ = buffer.WriteString("HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nContent-Length: 4096\r\n\r\n")
|
||||
_, _ = buffer.WriteString("data: {\"id\":\"chunk-broken\",\"object\":\"chat.completion.chunk\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"content\":\"first platform only\"},\"finish_reason\":null}]}\n\n")
|
||||
_ = buffer.Flush()
|
||||
})
|
||||
defer broken.close()
|
||||
fallback := newControlledProviderStub(func(_ int, w http.ResponseWriter, _ *http.Request) { writeOpenAIChatSuccess(w, "must not mix") })
|
||||
defer fallback.close()
|
||||
_, _ = f.createPlatformModel(t, "stream-after-a", "openai", broken.server.URL+"/v1", 10, baseModelID, "stream-broken", "text_generate", 2, "", nil)
|
||||
_, _ = f.createPlatformModel(t, "stream-after-b", "openai", fallback.server.URL+"/v1", 20, baseModelID, "stream-fallback", "text_generate", 1, "", nil)
|
||||
|
||||
status, _, body := f.postJSON(t, "/api/v1/chat/completions", map[string]any{
|
||||
"model": model, "stream": true, "messages": []any{map[string]any{"role": "user", "content": "stream"}},
|
||||
}, nil)
|
||||
if status != http.StatusOK || broken.calls.Load() != 1 || fallback.calls.Load() != 0 {
|
||||
t.Fatalf("post-delta failure rotated unexpectedly status=%d calls=%d/%d body=%s", status, broken.calls.Load(), fallback.calls.Load(), body)
|
||||
}
|
||||
bodyText := string(body)
|
||||
if !strings.Contains(bodyText, "first platform only") || strings.Contains(bodyText, "must not mix") || strings.Count(bodyText, "event: error") != 1 {
|
||||
t.Fatalf("stream response mixed providers or missed terminal error: %s", bodyText)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func (f *failurePolicyAcceptanceFixture) testRunnerPolicyHTTPValidation(t *testing.T) {
|
||||
valid := standardFailureRunnerPolicy(true)
|
||||
var saved map[string]any
|
||||
doJSON(t, f.server.URL, http.MethodPatch, "/api/admin/runtime/runner-policy", f.adminToken, valid, http.StatusOK, &saved)
|
||||
var before map[string]any
|
||||
doJSON(t, f.server.URL, http.MethodGet, "/api/admin/runtime/runner-policy", f.adminToken, nil, http.StatusOK, &before)
|
||||
|
||||
invalidRules := []map[string]any{
|
||||
{"categories": []any{"rate_limit"}, "action": "rotate"},
|
||||
{"categories": []any{"rate_limit"}, "action": "next", "target": "client"},
|
||||
{"action": "next"},
|
||||
{"statusCodes": []any{429}, "action": "cooldown_and_next", "target": "model", "cooldownSeconds": 0},
|
||||
{"statusCodes": []any{500}, "action": "demote_and_next", "target": "platform", "demoteSteps": -1},
|
||||
}
|
||||
for index, rule := range invalidRules {
|
||||
payload := standardFailureRunnerPolicy(true)
|
||||
payload["failoverPolicy"].(map[string]any)["actionRules"] = []any{rule}
|
||||
raw, _ := json.Marshal(payload)
|
||||
request, _ := http.NewRequest(http.MethodPatch, f.server.URL+"/api/admin/runtime/runner-policy", bytes.NewReader(raw))
|
||||
request.Header.Set("Authorization", "Bearer "+f.adminToken)
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
response, err := http.DefaultClient.Do(request)
|
||||
if err != nil {
|
||||
t.Fatalf("invalid policy request %d: %v", index, err)
|
||||
}
|
||||
body, _ := io.ReadAll(response.Body)
|
||||
response.Body.Close()
|
||||
if response.StatusCode != http.StatusBadRequest {
|
||||
t.Fatalf("invalid policy %d status=%d body=%s", index, response.StatusCode, body)
|
||||
}
|
||||
}
|
||||
var after map[string]any
|
||||
doJSON(t, f.server.URL, http.MethodGet, "/api/admin/runtime/runner-policy", f.adminToken, nil, http.StatusOK, &after)
|
||||
beforeFailover, _ := before["failoverPolicy"].(map[string]any)
|
||||
afterFailover, _ := after["failoverPolicy"].(map[string]any)
|
||||
if !jsonObjectsEqual(beforeFailover["actionRules"], afterFailover["actionRules"]) {
|
||||
t.Fatalf("invalid PATCH overwrote runner action rules: before=%+v after=%+v", beforeFailover["actionRules"], afterFailover["actionRules"])
|
||||
}
|
||||
}
|
||||
|
||||
func jsonObjectsEqual(left any, right any) bool {
|
||||
leftRaw, _ := json.Marshal(left)
|
||||
rightRaw, _ := json.Marshal(right)
|
||||
return bytes.Equal(leftRaw, rightRaw)
|
||||
}
|
||||
@@ -62,7 +62,7 @@ func (s *Server) uploadFile(w http.ResponseWriter, r *http.Request) {
|
||||
writeError(w, status, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, upload)
|
||||
writeJSON(w, http.StatusOK, easyAIFileUploadResponse(upload))
|
||||
}
|
||||
|
||||
func firstNonEmptyFormValue(r *http.Request, key string, fallback string) string {
|
||||
|
||||
@@ -170,6 +170,10 @@ func (s *Server) geminiGenerateContent(w http.ResponseWriter, r *http.Request) {
|
||||
writeGeminiTaskError(http.StatusConflict, "Idempotency-Key was reused for a different request", nil, "idempotency_key_reused")
|
||||
return
|
||||
}
|
||||
if errors.Is(err, store.ErrTaskRequestBinaryNotMaterialized) {
|
||||
writeGeminiTaskError(http.StatusBadRequest, err.Error(), nil, clients.ErrorCode(err))
|
||||
return
|
||||
}
|
||||
s.logger.Error("create gemini task failed", "kind", mapping.Kind, "error_category", "task_create_failed")
|
||||
writeGeminiTaskError(http.StatusInternalServerError, "create task failed", nil, "task_create_failed")
|
||||
return
|
||||
@@ -184,6 +188,11 @@ func (s *Server) geminiGenerateContent(w http.ResponseWriter, r *http.Request) {
|
||||
return
|
||||
}
|
||||
if task.Status == "succeeded" {
|
||||
task, err = s.hydrateTaskResult(r.Context(), task)
|
||||
if err != nil {
|
||||
writeGeminiTaskError(statusFromRunError(err), err.Error(), nil, clients.ErrorCode(err))
|
||||
return
|
||||
}
|
||||
if _, native := task.Result["candidates"]; native {
|
||||
writeJSON(w, http.StatusOK, task.Result)
|
||||
return
|
||||
@@ -209,6 +218,7 @@ func (s *Server) geminiGenerateContent(w http.ResponseWriter, r *http.Request) {
|
||||
if !requestStillConnected(r) {
|
||||
return
|
||||
}
|
||||
applyRunErrorHeaders(w, runErr)
|
||||
if wire := clients.ErrorWireResponse(runErr); wireResponseMatches(wire, clients.ProtocolGeminiGenerateContent) {
|
||||
writeWireResponse(w, wire)
|
||||
return
|
||||
@@ -261,6 +271,7 @@ func (s *Server) writeGeminiGenerateContentStream(runCtx context.Context, w http
|
||||
if !requestStillConnected(r) {
|
||||
return
|
||||
}
|
||||
applyRunErrorHeaders(w, runErr)
|
||||
if !nativePassthrough && !convertedFrames {
|
||||
if wire := clients.ErrorWireResponse(runErr); wireResponseMatches(wire, clients.ProtocolGeminiGenerateContent) {
|
||||
writeWireResponse(w, wire)
|
||||
|
||||
@@ -580,12 +580,20 @@ func (s *Server) listModels(w http.ResponseWriter, r *http.Request) {
|
||||
// @Tags playground
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param usage_scene query string false "模型可选场景;不传时保持原有行为" Enums(canvas_model_node,desktop)
|
||||
// @Success 200 {object} PlatformModelListResponse
|
||||
// @Failure 400 {object} ErrorEnvelope
|
||||
// @Failure 401 {object} ErrorEnvelope
|
||||
// @Failure 502 {object} ErrorEnvelope
|
||||
// @Failure 500 {object} ErrorEnvelope
|
||||
// @Router /api/v1/models [get]
|
||||
// @Router /api/v1/playground/models [get]
|
||||
func (s *Server) listPlayableModels(w http.ResponseWriter, r *http.Request) {
|
||||
usageScene, err := parseModelUsageSceneQuery(r.URL.Query())
|
||||
if err != nil {
|
||||
writeError(w, http.StatusBadRequest, err.Error())
|
||||
return
|
||||
}
|
||||
user, _ := auth.UserFromContext(r.Context())
|
||||
models, err := s.store.ListAccessiblePlatformModels(r.Context(), user)
|
||||
if err != nil {
|
||||
@@ -593,6 +601,15 @@ func (s *Server) listPlayableModels(w http.ResponseWriter, r *http.Request) {
|
||||
writeError(w, http.StatusInternalServerError, "list playable models failed")
|
||||
return
|
||||
}
|
||||
if usageScene != "" {
|
||||
availableModelIDs, sceneErr := s.listServerMainUsageSceneModelIDs(r.Context(), usageScene)
|
||||
if sceneErr != nil {
|
||||
s.logger.Error("filter playable models by usage scene failed", "usage_scene", usageScene, "error", sceneErr)
|
||||
writeError(w, http.StatusBadGateway, "model usage scene catalog is unavailable")
|
||||
return
|
||||
}
|
||||
models = filterPlatformModelsByUsageSceneIDs(models, availableModelIDs)
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]any{"items": s.platformModelResponses(r.Context(), models)})
|
||||
}
|
||||
|
||||
@@ -963,6 +980,7 @@ func (s *Server) estimatePricing(w http.ResponseWriter, r *http.Request) {
|
||||
return
|
||||
}
|
||||
if errors.Is(err, store.ErrNoModelCandidate) {
|
||||
applyRunErrorHeaders(w, err)
|
||||
writeErrorWithDetails(w, statusFromRunError(err), runErrorMessage(err), runErrorDetails(err), store.ModelCandidateErrorCode(err))
|
||||
return
|
||||
}
|
||||
@@ -1039,15 +1057,14 @@ func (s *Server) listModelRateLimitStatuses(w http.ResponseWriter, r *http.Reque
|
||||
// @Failure 429 {object} ErrorEnvelope
|
||||
// @Failure 502 {object} ErrorEnvelope
|
||||
// @Router /api/v1/reranks [post]
|
||||
// @Router /api/v1/videos/generations [post]
|
||||
// @Router /api/v1/song/generations [post]
|
||||
// @Router /api/v1/music/generations [post]
|
||||
// @Router /api/v1/speech/generations [post]
|
||||
// @Router /api/v1/voice_clone [post]
|
||||
func (s *Server) createTask(kind string, compatible bool) http.Handler {
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
targetProtocol := targetProtocolForTaskRequest(kind, r)
|
||||
writeTaskError := func(status int, message string, details map[string]any, codes ...string) {
|
||||
if easyAIAsyncMediaRequest(kind, r) {
|
||||
writeEasyAIAsyncError(w, status, message, details, codes...)
|
||||
return
|
||||
}
|
||||
if targetProtocol == "" {
|
||||
writeErrorWithDetails(w, status, message, details, codes...)
|
||||
return
|
||||
@@ -1093,9 +1110,10 @@ func (s *Server) createTask(kind string, compatible bool) http.Handler {
|
||||
return
|
||||
}
|
||||
responsePlan := planTaskResponse(kind, compatible, body, r)
|
||||
if targetProtocol != "" {
|
||||
if targetProtocol != "" && !easyAIAsyncMediaRequest(kind, r) {
|
||||
// OpenAI compatibility endpoints are synchronous APIs. Gateway async
|
||||
// task envelopes belong only to native Gateway routes.
|
||||
// task envelopes belong only to native Gateway routes, except for the
|
||||
// explicit EasyAI media X-Async extension.
|
||||
responsePlan.asyncMode = false
|
||||
}
|
||||
idempotencyKey, hasIdempotencyKey, err := optionalTaskIdempotencyKey(r)
|
||||
@@ -1134,6 +1152,10 @@ func (s *Server) createTask(kind string, compatible bool) http.Handler {
|
||||
writeTaskError(http.StatusConflict, "Idempotency-Key was reused for a different request", nil, "idempotency_key_reused")
|
||||
return
|
||||
}
|
||||
if errors.Is(err, store.ErrTaskRequestBinaryNotMaterialized) {
|
||||
writeTaskError(http.StatusBadRequest, err.Error(), nil, clients.ErrorCode(err))
|
||||
return
|
||||
}
|
||||
s.logger.Error("create task failed", "kind", kind, "error_category", "task_create_failed")
|
||||
writeTaskError(http.StatusInternalServerError, "create task failed", nil, "task_create_failed")
|
||||
return
|
||||
@@ -1150,7 +1172,18 @@ func (s *Server) createTask(kind string, compatible bool) http.Handler {
|
||||
writeTaskError(http.StatusConflict, "streaming idempotent replay is not supported", nil, "idempotency_stream_replay_unsupported")
|
||||
return
|
||||
}
|
||||
if targetProtocol != "" {
|
||||
if task.Status == "succeeded" && !responsePlan.asyncMode && (targetProtocol != "" || responsePlan.compatibleMode) {
|
||||
task, err = s.hydrateTaskResult(r.Context(), task)
|
||||
if err != nil {
|
||||
if targetProtocol != "" {
|
||||
writeProtocolError(w, targetProtocol, statusFromRunError(err), err.Error(), nil, clients.ErrorCode(err))
|
||||
} else {
|
||||
writeStoredBinaryResultError(w, err)
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
if targetProtocol != "" && !responsePlan.asyncMode {
|
||||
if task.Status == "succeeded" {
|
||||
writeJSON(w, http.StatusOK, task.Result)
|
||||
return
|
||||
@@ -1274,7 +1307,7 @@ func openAIEmbeddingsDoc() {}
|
||||
// @Security BearerAuth
|
||||
// @Param input body TaskRequest true "OpenAI Images 请求"
|
||||
// @Param X-Async header bool false "true 时异步创建任务并返回 202"
|
||||
// @Success 200 {object} CompatibleResponse
|
||||
// @Success 200 {object} OpenAIImagesResponse
|
||||
// @Success 202 {object} TaskAcceptedResponse
|
||||
// @Failure 400 {object} OpenAIErrorEnvelope
|
||||
// @Failure 401 {object} OpenAIErrorEnvelope
|
||||
@@ -1284,6 +1317,32 @@ func openAIEmbeddingsDoc() {}
|
||||
// @Router /api/v1/images/edits [post]
|
||||
func openAIImagesDoc() {}
|
||||
|
||||
// easyAIMediaTasksDoc godoc
|
||||
// @Summary 创建 EasyAI 媒体任务
|
||||
// @Description 默认同步返回 EasyAI GeneratedResponse;设置 X-Async=true 时返回同时兼容 Gateway 与 server-main EasyAIClient 的异步提交结构。
|
||||
// @Tags media
|
||||
// @Accept json
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param X-Async header bool false "true 时异步创建任务并返回 202"
|
||||
// @Param Idempotency-Key header string false "可选请求幂等键;同一用户范围内唯一"
|
||||
// @Param input body TaskRequest true "媒体任务请求,字段随任务类型变化"
|
||||
// @Success 200 {object} EasyAIGeneratedResponse
|
||||
// @Success 202 {object} TaskAcceptedResponse
|
||||
// @Failure 400 {object} ErrorEnvelope
|
||||
// @Failure 401 {object} ErrorEnvelope
|
||||
// @Failure 402 {object} ErrorEnvelope
|
||||
// @Failure 403 {object} ErrorEnvelope
|
||||
// @Failure 404 {object} ErrorEnvelope
|
||||
// @Failure 429 {object} ErrorEnvelope
|
||||
// @Failure 502 {object} ErrorEnvelope
|
||||
// @Router /api/v1/videos/generations [post]
|
||||
// @Router /api/v1/song/generations [post]
|
||||
// @Router /api/v1/music/generations [post]
|
||||
// @Router /api/v1/speech/generations [post]
|
||||
// @Router /api/v1/voice_clone [post]
|
||||
func easyAIMediaTasksDoc() {}
|
||||
|
||||
func (s *Server) requestExecutionContext(r *http.Request) (context.Context, context.CancelFunc) {
|
||||
base := context.WithoutCancel(r.Context())
|
||||
if s.ctx == nil {
|
||||
@@ -1363,6 +1422,7 @@ func writeProtocolCompatibleTaskResponse(runCtx context.Context, w http.Response
|
||||
if !requestStillConnected(r) {
|
||||
return
|
||||
}
|
||||
applyRunErrorHeaders(w, runErr)
|
||||
errorPayload := map[string]any{"message": runErrorMessage(runErr), "type": "server_error", "param": nil, "code": runErrorCode(runErr)}
|
||||
if status := statusFromRunError(runErr); status >= 400 && status < 500 {
|
||||
errorPayload["type"] = "invalid_request_error"
|
||||
@@ -1400,6 +1460,7 @@ func writeProtocolCompatibleTaskResponse(runCtx context.Context, w http.Response
|
||||
if !requestStillConnected(r) {
|
||||
return
|
||||
}
|
||||
applyRunErrorHeaders(w, runErr)
|
||||
if wire := clients.ErrorWireResponse(runErr); wireResponseMatches(wire, targetProtocol) {
|
||||
writeWireResponse(w, wire)
|
||||
return
|
||||
@@ -1418,7 +1479,7 @@ func writeProtocolCompatibleTaskResponse(runCtx context.Context, w http.Response
|
||||
writeWireResponse(w, result.Wire)
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, result.Output)
|
||||
writeJSON(w, http.StatusOK, easyAISynchronousTaskResponse(result.Task, result.Output))
|
||||
}
|
||||
|
||||
func streamIncludeUsage(body map[string]any) bool {
|
||||
@@ -1461,14 +1522,7 @@ func streamCompatibleKind(kind string) bool {
|
||||
}
|
||||
|
||||
func writeTaskAccepted(w http.ResponseWriter, task store.GatewayTask) {
|
||||
writeJSON(w, http.StatusAccepted, map[string]any{
|
||||
"taskId": task.ID,
|
||||
"task": task,
|
||||
"next": map[string]string{
|
||||
"events": fmt.Sprintf("/api/v1/tasks/%s/events", task.ID),
|
||||
"detail": fmt.Sprintf("/api/v1/tasks/%s", task.ID),
|
||||
},
|
||||
})
|
||||
writeJSON(w, http.StatusAccepted, easyAITaskAcceptedResponse(task))
|
||||
}
|
||||
|
||||
func apiKeyScopeAllowed(user *auth.User, kind string) bool {
|
||||
@@ -1578,6 +1632,12 @@ func scopeForTaskKind(kind string) string {
|
||||
|
||||
func statusFromRunError(err error) int {
|
||||
switch {
|
||||
case clients.ErrorCode(err) == "binary_result_expired":
|
||||
return http.StatusGone
|
||||
case clients.ErrorCode(err) == "binary_result_corrupted" || clients.ErrorCode(err) == "result_binary_not_materialized":
|
||||
return http.StatusInternalServerError
|
||||
case clients.ErrorCode(err) == "local_result_storage_unavailable":
|
||||
return http.StatusServiceUnavailable
|
||||
case clients.ErrorCode(err) == "billing_hold":
|
||||
return http.StatusServiceUnavailable
|
||||
case runner.IsPricingUnavailable(err):
|
||||
@@ -1637,11 +1697,21 @@ func runErrorDetails(err error) map[string]any {
|
||||
return map[string]any{"rateLimit": detail}
|
||||
}
|
||||
if detail := store.ModelCandidateErrorDetails(err); len(detail) > 0 {
|
||||
return map[string]any{"modelCandidate": detail}
|
||||
return detail
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func applyRunErrorHeaders(w http.ResponseWriter, err error) {
|
||||
if retryAfter := store.ModelCandidateRetryAfter(err); retryAfter > 0 {
|
||||
seconds := int((retryAfter + time.Second - 1) / time.Second)
|
||||
if seconds < 1 {
|
||||
seconds = 1
|
||||
}
|
||||
w.Header().Set("Retry-After", strconv.Itoa(seconds))
|
||||
}
|
||||
}
|
||||
|
||||
func rateLimitErrorSummary(err error) string {
|
||||
var limitErr *store.RateLimitExceededError
|
||||
if !errors.As(err, &limitErr) {
|
||||
@@ -1882,6 +1952,7 @@ func boolValue(body map[string]any, key string) bool {
|
||||
// @Success 200 {object} store.GatewayTask
|
||||
// @Failure 401 {object} ErrorEnvelope
|
||||
// @Failure 404 {object} ErrorEnvelope
|
||||
// @Failure 410 {object} ErrorEnvelope
|
||||
// @Failure 500 {object} ErrorEnvelope
|
||||
// @Router /api/workspace/tasks/{taskID} [get]
|
||||
// @Router /api/v1/tasks/{taskID} [get]
|
||||
@@ -1901,6 +1972,11 @@ func (s *Server) getTask(w http.ResponseWriter, r *http.Request) {
|
||||
task.Cancellable = &cancelState.Cancellable
|
||||
task.Submitted = &cancelState.Submitted
|
||||
task.Message = cancelState.Message
|
||||
task, err = s.hydrateTaskResult(r.Context(), task)
|
||||
if err != nil {
|
||||
writeStoredBinaryResultError(w, err)
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, task)
|
||||
return
|
||||
}
|
||||
|
||||
@@ -0,0 +1,54 @@
|
||||
package httpapi
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
)
|
||||
|
||||
func requireDedicatedIntegrationTestDatabase(t *testing.T, ctx context.Context, pool *pgxpool.Pool) {
|
||||
t.Helper()
|
||||
var databaseName string
|
||||
if err := pool.QueryRow(ctx, `SELECT current_database()`).Scan(&databaseName); err != nil {
|
||||
t.Fatalf("read integration test database name: %v", err)
|
||||
}
|
||||
if !isDedicatedIntegrationTestDatabase(databaseName) {
|
||||
t.Fatalf(
|
||||
"AI_GATEWAY_TEST_DATABASE_URL must reference a dedicated test database; refusing to use %q",
|
||||
databaseName,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
func isDedicatedIntegrationTestDatabase(databaseName string) bool {
|
||||
normalized := strings.ToLower(strings.TrimSpace(databaseName))
|
||||
return normalized == "test" ||
|
||||
strings.HasPrefix(normalized, "test_") ||
|
||||
strings.HasSuffix(normalized, "_test") ||
|
||||
strings.Contains(normalized, "_test_")
|
||||
}
|
||||
|
||||
func TestDedicatedIntegrationTestDatabaseName(t *testing.T) {
|
||||
t.Parallel()
|
||||
for _, testCase := range []struct {
|
||||
name string
|
||||
database string
|
||||
allowed bool
|
||||
}{
|
||||
{name: "documented local test database", database: "easyai_ai_gateway_test_codex", allowed: true},
|
||||
{name: "test prefix", database: "test_gateway", allowed: true},
|
||||
{name: "test suffix", database: "gateway_test", allowed: true},
|
||||
{name: "production database", database: "easyai_ai_gateway", allowed: false},
|
||||
{name: "similar production name", database: "gateway_contest", allowed: false},
|
||||
} {
|
||||
testCase := testCase
|
||||
t.Run(testCase.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
if got := isDedicatedIntegrationTestDatabase(testCase.database); got != testCase.allowed {
|
||||
t.Fatalf("isDedicatedIntegrationTestDatabase(%q) = %v, want %v", testCase.database, got, testCase.allowed)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -177,6 +177,10 @@ func (s *Server) createKelingOmniVideo(w http.ResponseWriter, r *http.Request) {
|
||||
writeKelingCompatError(w, requestID, kelingCompatGatewayError(staged.Err))
|
||||
return
|
||||
}
|
||||
if errors.Is(createErr, store.ErrTaskRequestBinaryNotMaterialized) {
|
||||
writeKelingCompatError(w, requestID, newKelingCompatError(http.StatusBadRequest, 1201, clients.ErrorCode(createErr)))
|
||||
return
|
||||
}
|
||||
s.logger.Error("create Kling-compatible task failed", "error", createErr)
|
||||
writeKelingCompatError(w, requestID, newKelingCompatError(http.StatusInternalServerError, 5000, "create task failed"))
|
||||
return
|
||||
@@ -202,15 +206,11 @@ func (s *Server) createKelingOmniVideo(w http.ResponseWriter, r *http.Request) {
|
||||
writeKelingCompatError(w, requestID, kelingCompatGatewayError(createErr))
|
||||
return
|
||||
}
|
||||
if kelingTaskUsesNativeProtocol(task) && len(task.CompatibilitySubmitBody) > 0 {
|
||||
writeJSON(w, task.CompatibilitySubmitHTTPStatus, task.CompatibilitySubmitBody)
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, KelingCompatibleEnvelope{
|
||||
Code: 0,
|
||||
Message: "SUCCEED",
|
||||
RequestID: requestID,
|
||||
Data: kelingCompatTaskData(task),
|
||||
Data: kelingCompatSubmissionData(task),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -225,6 +225,7 @@ func (s *Server) createKelingOmniVideo(w http.ResponseWriter, r *http.Request) {
|
||||
// @Failure 401 {object} KelingCompatibleEnvelope
|
||||
// @Failure 403 {object} KelingCompatibleEnvelope
|
||||
// @Failure 404 {object} KelingCompatibleEnvelope
|
||||
// @Failure 410 {object} KelingCompatibleEnvelope
|
||||
// @Failure 500 {object} KelingCompatibleEnvelope
|
||||
// @Router /api/v1/videos/omni-video/{taskID} [get]
|
||||
func (s *Server) getKelingOmniVideo(w http.ResponseWriter, r *http.Request) {
|
||||
@@ -251,15 +252,10 @@ func (s *Server) getKelingOmniVideo(w http.ResponseWriter, r *http.Request) {
|
||||
writeKelingCompatError(w, requestID, newKelingCompatError(http.StatusNotFound, 1203, "task not found"))
|
||||
return
|
||||
}
|
||||
if kelingTaskUsesNativeProtocol(task) {
|
||||
if raw, ok := task.Result["raw"].(map[string]any); ok && len(raw) > 0 {
|
||||
writeJSON(w, http.StatusOK, raw)
|
||||
return
|
||||
}
|
||||
if len(task.RemoteTaskPayload) > 0 {
|
||||
writeJSON(w, http.StatusOK, task.RemoteTaskPayload)
|
||||
return
|
||||
}
|
||||
task, err = s.hydrateTaskResult(r.Context(), task)
|
||||
if err != nil {
|
||||
writeKelingCompatError(w, requestID, newKelingCompatError(statusFromRunError(err), 5000, clients.ErrorCode(err)))
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, KelingCompatibleEnvelope{
|
||||
Code: 0,
|
||||
@@ -705,15 +701,22 @@ func kelingCompatTaskData(task store.GatewayTask) map[string]any {
|
||||
return data
|
||||
}
|
||||
|
||||
func kelingCompatPublicID(task store.GatewayTask) string {
|
||||
if publicID := strings.TrimSpace(task.CompatibilityPublicID); publicID != "" {
|
||||
return publicID
|
||||
}
|
||||
return task.ID
|
||||
func kelingCompatSubmissionData(task store.GatewayTask) map[string]any {
|
||||
data := kelingCompatTaskData(task)
|
||||
data["task_status"] = "submitted"
|
||||
delete(data, "task_status_code")
|
||||
delete(data, "task_status_msg")
|
||||
delete(data, "task_result")
|
||||
return data
|
||||
}
|
||||
|
||||
func kelingTaskUsesNativeProtocol(task store.GatewayTask) bool {
|
||||
return task.CompatibilityProtocol == clients.ProtocolKlingV1Omni && task.CompatibilitySourceProtocol == clients.ProtocolKlingV1Omni
|
||||
func kelingCompatPublicID(task store.GatewayTask) string {
|
||||
if task.CompatibilityProtocol == task.CompatibilitySourceProtocol {
|
||||
if remoteTaskID := strings.TrimSpace(task.RemoteTaskID); remoteTaskID != "" {
|
||||
return remoteTaskID
|
||||
}
|
||||
}
|
||||
return task.ID
|
||||
}
|
||||
|
||||
func kelingCompatTaskStatus(status string) string {
|
||||
|
||||
@@ -209,6 +209,10 @@ func TestKelingCompatTaskDataAndOwnership(t *testing.T) {
|
||||
if video["id"] != "video-1" || video["watermark_url"] != "https://example.com/watermarked.mp4" || video["duration"] != "5" {
|
||||
t.Fatalf("unexpected compatible video: %+v", video)
|
||||
}
|
||||
submission := kelingCompatSubmissionData(task)
|
||||
if submission["task_status"] != "submitted" || submission["task_result"] != nil {
|
||||
t.Fatalf("submission response must be stable and minimal: %+v", submission)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRequireKelingAPIKeyWritesOfficialAuthEnvelope(t *testing.T) {
|
||||
|
||||
@@ -144,6 +144,8 @@ func (s *Server) createKlingCompatTask(w http.ResponseWriter, r *http.Request, v
|
||||
writeKlingCompatError(w, http.StatusConflict, err.Error(), "idempotency_key_reused")
|
||||
case errors.Is(err, store.ErrExternalTaskIDReused):
|
||||
writeKlingCompatError(w, http.StatusConflict, err.Error(), "external_task_id_reused")
|
||||
case errors.Is(err, store.ErrTaskRequestBinaryNotMaterialized):
|
||||
writeKlingCompatError(w, http.StatusBadRequest, err.Error(), clients.ErrorCode(err))
|
||||
default:
|
||||
s.logger.Error("create Kling compatibility task failed", "version", version, "model", model, "error", err)
|
||||
writeKlingCompatError(w, http.StatusInternalServerError, "create task failed", "task_create_failed")
|
||||
@@ -177,10 +179,10 @@ func (s *Server) createKlingCompatTask(w http.ResponseWriter, r *http.Request, v
|
||||
w.Header().Set("Idempotent-Replayed", "true")
|
||||
}
|
||||
if version == "v2" {
|
||||
writeJSON(w, http.StatusOK, klingV2Envelope(task))
|
||||
writeJSON(w, http.StatusOK, klingV2SubmissionEnvelope(task))
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, klingV1Envelope(task))
|
||||
writeJSON(w, http.StatusOK, klingV1SubmissionEnvelope(task))
|
||||
}
|
||||
|
||||
func klingCompatTaskBody(version string, model string, native map[string]any) (map[string]any, string, error) {
|
||||
@@ -464,6 +466,7 @@ func validateKlingCompatBody(model string, body map[string]any) error {
|
||||
// @Param taskID path string true "任务 ID"
|
||||
// @Success 200 {object} map[string]interface{}
|
||||
// @Failure 404 {object} map[string]interface{}
|
||||
// @Failure 410 {object} map[string]interface{}
|
||||
// @Router /api/v1/kling/v1/videos/omni-video/{taskID} [get]
|
||||
func (s *Server) klingV1GetOmniVideo(w http.ResponseWriter, r *http.Request) {
|
||||
user, _ := auth.UserFromContext(r.Context())
|
||||
@@ -476,6 +479,11 @@ func (s *Server) klingV1GetOmniVideo(w http.ResponseWriter, r *http.Request) {
|
||||
writeKlingCompatError(w, http.StatusInternalServerError, "get task failed", "task_query_failed")
|
||||
return
|
||||
}
|
||||
task, err = s.hydrateTaskResult(r.Context(), task)
|
||||
if err != nil {
|
||||
writeKlingCompatError(w, statusFromRunError(err), err.Error(), clients.ErrorCode(err))
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, klingV1Envelope(task))
|
||||
}
|
||||
|
||||
@@ -520,6 +528,7 @@ func (s *Server) klingV1ListOmniVideos(w http.ResponseWriter, r *http.Request) {
|
||||
// @Param task_ids query string false "逗号分隔的任务 ID"
|
||||
// @Param external_task_ids query string false "逗号分隔的外部任务 ID"
|
||||
// @Success 200 {object} map[string]interface{}
|
||||
// @Failure 410 {object} map[string]interface{}
|
||||
// @Router /api/v1/kling/v2/tasks [get]
|
||||
func (s *Server) klingV2GetTasks(w http.ResponseWriter, r *http.Request) {
|
||||
user, _ := auth.UserFromContext(r.Context())
|
||||
@@ -543,6 +552,11 @@ func (s *Server) klingV2GetTasks(w http.ResponseWriter, r *http.Request) {
|
||||
writeKlingCompatError(w, http.StatusInternalServerError, "get task failed", "task_query_failed")
|
||||
return
|
||||
}
|
||||
task, err = s.hydrateTaskResult(r.Context(), task)
|
||||
if err != nil {
|
||||
writeKlingCompatError(w, statusFromRunError(err), err.Error(), clients.ErrorCode(err))
|
||||
return
|
||||
}
|
||||
data = append(data, klingV2TaskData(task))
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]any{"code": 0, "message": "success", "request_id": klingRequestIDFromAny(data), "data": data})
|
||||
@@ -612,20 +626,17 @@ func (s *Server) klingV2ListTasks(w http.ResponseWriter, r *http.Request) {
|
||||
}
|
||||
|
||||
func klingV1Envelope(task store.GatewayTask) map[string]any {
|
||||
if klingTaskUsesNativeV1Protocol(task) {
|
||||
if raw := klingMap(task.Result["raw"]); len(raw) > 0 {
|
||||
return raw
|
||||
}
|
||||
if len(task.RemoteTaskPayload) > 0 {
|
||||
return task.RemoteTaskPayload
|
||||
}
|
||||
if raw := klingMap(task.CompatibilitySubmitBody); len(raw) > 0 {
|
||||
return raw
|
||||
}
|
||||
}
|
||||
return map[string]any{"code": 0, "message": "success", "request_id": firstNonEmpty(task.RequestID, task.ID), "data": klingV1TaskData(task)}
|
||||
}
|
||||
|
||||
func klingV1SubmissionEnvelope(task store.GatewayTask) map[string]any {
|
||||
data := klingV1TaskData(task)
|
||||
data["task_status"] = "submitted"
|
||||
delete(data, "task_status_msg")
|
||||
delete(data, "task_result")
|
||||
return map[string]any{"code": 0, "message": "success", "request_id": firstNonEmpty(task.RequestID, task.ID), "data": data}
|
||||
}
|
||||
|
||||
func klingV1TaskData(task store.GatewayTask) map[string]any {
|
||||
data := map[string]any{
|
||||
"task_id": klingCompatibilityPublicID(task), "task_status": klingV1Status(task.Status),
|
||||
@@ -646,20 +657,17 @@ func klingV1TaskData(task store.GatewayTask) map[string]any {
|
||||
}
|
||||
|
||||
func klingV2Envelope(task store.GatewayTask) map[string]any {
|
||||
if task.CompatibilityProtocol == clients.ProtocolKlingV2Omni && task.CompatibilitySourceProtocol == clients.ProtocolKlingV2Omni {
|
||||
if raw := klingMap(task.Result["raw"]); len(raw) > 0 {
|
||||
return raw
|
||||
}
|
||||
if len(task.RemoteTaskPayload) > 0 {
|
||||
return task.RemoteTaskPayload
|
||||
}
|
||||
if len(task.CompatibilitySubmitBody) > 0 {
|
||||
return task.CompatibilitySubmitBody
|
||||
}
|
||||
}
|
||||
return map[string]any{"code": 0, "message": "success", "request_id": firstNonEmpty(task.RequestID, task.ID), "data": klingV2TaskData(task)}
|
||||
}
|
||||
|
||||
func klingV2SubmissionEnvelope(task store.GatewayTask) map[string]any {
|
||||
data := klingV2TaskData(task)
|
||||
data["status"] = "submitted"
|
||||
delete(data, "message")
|
||||
delete(data, "outputs")
|
||||
return map[string]any{"code": 0, "message": "success", "request_id": firstNonEmpty(task.RequestID, task.ID), "data": data}
|
||||
}
|
||||
|
||||
func klingV2TaskData(task store.GatewayTask) map[string]any {
|
||||
data := map[string]any{
|
||||
"id": klingCompatibilityPublicID(task), "status": klingV2Status(task.Status),
|
||||
@@ -676,16 +684,14 @@ func klingV2TaskData(task store.GatewayTask) map[string]any {
|
||||
}
|
||||
|
||||
func klingCompatibilityPublicID(task store.GatewayTask) string {
|
||||
if publicID := strings.TrimSpace(task.CompatibilityPublicID); publicID != "" {
|
||||
return publicID
|
||||
if task.CompatibilityProtocol == task.CompatibilitySourceProtocol {
|
||||
if remoteTaskID := strings.TrimSpace(task.RemoteTaskID); remoteTaskID != "" {
|
||||
return remoteTaskID
|
||||
}
|
||||
}
|
||||
return task.ID
|
||||
}
|
||||
|
||||
func klingTaskUsesNativeV1Protocol(task store.GatewayTask) bool {
|
||||
return task.CompatibilityProtocol == clients.ProtocolKlingV1Omni && task.CompatibilitySourceProtocol == clients.ProtocolKlingV1Omni
|
||||
}
|
||||
|
||||
func klingMap(value any) map[string]any {
|
||||
result, _ := value.(map[string]any)
|
||||
return result
|
||||
|
||||
@@ -6,6 +6,9 @@ import (
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
)
|
||||
|
||||
func TestKlingV1O1CompatibilityBody(t *testing.T) {
|
||||
@@ -86,6 +89,27 @@ func TestKlingV2CompatibilityBody(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestKlingSubmissionEnvelopesDoNotReplayCurrentResult(t *testing.T) {
|
||||
task := store.GatewayTask{
|
||||
ID: "gateway-task", Status: "succeeded", RemoteTaskID: "upstream-task",
|
||||
CompatibilityProtocol: "kling-v1-omni",
|
||||
CompatibilitySourceProtocol: "kling-v1-omni",
|
||||
Result: map[string]any{"data": []any{map[string]any{"url": "https://example.invalid/result.mp4"}}},
|
||||
CreatedAt: time.Unix(100, 0),
|
||||
UpdatedAt: time.Unix(101, 0),
|
||||
}
|
||||
v1 := klingV1SubmissionEnvelope(task)
|
||||
v1Data, _ := v1["data"].(map[string]any)
|
||||
if v1Data["task_status"] != "submitted" || v1Data["task_result"] != nil {
|
||||
t.Fatalf("V1 submission envelope=%#v", v1)
|
||||
}
|
||||
v2 := klingV2SubmissionEnvelope(task)
|
||||
v2Data, _ := v2["data"].(map[string]any)
|
||||
if v2Data["status"] != "submitted" || v2Data["outputs"] != nil {
|
||||
t.Fatalf("V2 submission envelope=%#v", v2)
|
||||
}
|
||||
}
|
||||
|
||||
func TestKlingCompatibilityValidation(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
|
||||
@@ -45,6 +45,58 @@ func (s *Server) cleanupExpiredLocalTempAssets(ctx context.Context, now time.Tim
|
||||
for _, target := range targets {
|
||||
deleted += s.cleanupExpiredLocalTempAssetsInDir(ctx, now, target)
|
||||
}
|
||||
deleted += s.cleanupExpiredLocalBinaryResults(now)
|
||||
return deleted
|
||||
}
|
||||
|
||||
func (s *Server) cleanupExpiredLocalBinaryResults(now time.Time) int {
|
||||
storageDir := strings.TrimSpace(s.cfg.LocalGeneratedStorageDir)
|
||||
if storageDir == "" {
|
||||
storageDir = config.DefaultLocalGeneratedStorageDir
|
||||
}
|
||||
root := filepath.Join(storageDir, "results")
|
||||
ttlHours := s.cfg.LocalResultTTLHours
|
||||
if ttlHours <= 0 {
|
||||
ttlHours = 24
|
||||
}
|
||||
expiredBefore := now.Add(-time.Duration(ttlHours) * time.Hour)
|
||||
deleted := 0
|
||||
var directories []string
|
||||
err := filepath.WalkDir(root, func(path string, entry os.DirEntry, walkErr error) error {
|
||||
if walkErr != nil {
|
||||
if !errors.Is(walkErr, os.ErrNotExist) && s.logger != nil {
|
||||
s.logger.Warn("walk local binary result failed", "path", path, "error", walkErr)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
if entry.IsDir() {
|
||||
if path != root {
|
||||
directories = append(directories, path)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
if entry.Type()&os.ModeSymlink != 0 {
|
||||
return nil
|
||||
}
|
||||
info, infoErr := entry.Info()
|
||||
if infoErr != nil || info.ModTime().After(expiredBefore) {
|
||||
return nil
|
||||
}
|
||||
if removeErr := os.Remove(path); removeErr != nil && !errors.Is(removeErr, os.ErrNotExist) {
|
||||
if s.logger != nil {
|
||||
s.logger.Warn("remove local binary result failed", "path", path, "error", removeErr)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
deleted++
|
||||
return nil
|
||||
})
|
||||
if err != nil && !errors.Is(err, os.ErrNotExist) && s.logger != nil {
|
||||
s.logger.Warn("scan local binary result root failed", "dir", root, "error", err)
|
||||
}
|
||||
for index := len(directories) - 1; index >= 0; index-- {
|
||||
_ = os.Remove(directories[index])
|
||||
}
|
||||
return deleted
|
||||
}
|
||||
|
||||
|
||||
@@ -62,17 +62,19 @@ type ModelCatalogProviderSummary struct {
|
||||
}
|
||||
|
||||
type ModelCatalogSource struct {
|
||||
ID string `json:"id"`
|
||||
PlatformID string `json:"platformId,omitempty"`
|
||||
PlatformName string `json:"platformName,omitempty"`
|
||||
ProviderKey string `json:"providerKey"`
|
||||
ProviderName string `json:"providerName"`
|
||||
ModelName string `json:"modelName"`
|
||||
ModelAlias string `json:"modelAlias,omitempty"`
|
||||
DisplayName string `json:"displayName"`
|
||||
ModelType []string `json:"modelType"`
|
||||
Enabled bool `json:"enabled"`
|
||||
RateLimits ModelCatalogRateLimits `json:"rateLimits"`
|
||||
ID string `json:"id"`
|
||||
PlatformID string `json:"platformId,omitempty"`
|
||||
PlatformName string `json:"platformName,omitempty"`
|
||||
ProviderKey string `json:"providerKey"`
|
||||
ProviderName string `json:"providerName"`
|
||||
ModelName string `json:"modelName"`
|
||||
ProviderModelName string `json:"providerModelName,omitempty"`
|
||||
ModelAlias string `json:"modelAlias,omitempty"`
|
||||
LegacyAliases []string `json:"legacyAliases,omitempty"`
|
||||
DisplayName string `json:"displayName"`
|
||||
ModelType []string `json:"modelType"`
|
||||
Enabled bool `json:"enabled"`
|
||||
RateLimits ModelCatalogRateLimits `json:"rateLimits"`
|
||||
}
|
||||
|
||||
type ModelCatalogText struct {
|
||||
@@ -202,10 +204,10 @@ func buildModelCatalog(
|
||||
sourceCount := 0
|
||||
for _, model := range models {
|
||||
platform := platformMap[model.PlatformID]
|
||||
if !catalogSourceEnabled(model, platform) {
|
||||
baseModel := baseModelMap[model.BaseModelID]
|
||||
if !catalogSourceEnabled(model, platform, baseModel) {
|
||||
continue
|
||||
}
|
||||
baseModel := baseModelMap[model.BaseModelID]
|
||||
providerKey := modelProviderKey(model, baseModel, platform)
|
||||
provider := providerMap[providerKey]
|
||||
providerName := firstNonEmpty(
|
||||
@@ -232,14 +234,14 @@ func buildModelCatalog(
|
||||
rateLimits: effectiveModelRateLimits(model, platform, runtimePolicyMap),
|
||||
}
|
||||
|
||||
key := catalogAliasKey(model)
|
||||
key := catalogIdentityKey(model, baseModel)
|
||||
current := groups[key]
|
||||
if current == nil {
|
||||
current = &catalogGroup{
|
||||
key: key,
|
||||
alias: firstNonEmpty(model.ModelAlias, model.DisplayName, model.ModelName),
|
||||
displayName: firstNonEmpty(model.DisplayName, model.ModelAlias, model.ModelName),
|
||||
modelName: model.ModelName,
|
||||
alias: firstNonEmpty(baseModel.InvocationName, model.ModelName),
|
||||
displayName: firstNonEmpty(baseModel.DisplayName, model.DisplayName, baseModel.InvocationName, model.ModelName),
|
||||
modelName: firstNonEmpty(baseModel.InvocationName, model.ModelName),
|
||||
modelType: cloneStringSlice(model.ModelType),
|
||||
capabilities: cloneObject(model.Capabilities),
|
||||
}
|
||||
@@ -274,8 +276,8 @@ func buildModelCatalog(
|
||||
}
|
||||
}
|
||||
|
||||
func catalogSourceEnabled(model store.PlatformModel, platform store.Platform) bool {
|
||||
return model.Enabled && platform.ID != "" && platform.Status == "enabled"
|
||||
func catalogSourceEnabled(model store.PlatformModel, platform store.Platform, baseModel store.BaseModel) bool {
|
||||
return model.Enabled && platform.ID != "" && platform.Status == "enabled" && (baseModel.ID == "" || baseModel.Status == "" || baseModel.Status == "active")
|
||||
}
|
||||
|
||||
func modelCatalogItem(group *catalogGroup, accessRuleGroups map[string]accessRuleGroup) ModelCatalogItem {
|
||||
@@ -304,17 +306,19 @@ func modelCatalogItem(group *catalogGroup, accessRuleGroups map[string]accessRul
|
||||
itemSources := make([]ModelCatalogSource, 0, len(sources))
|
||||
for _, source := range sources {
|
||||
itemSources = append(itemSources, ModelCatalogSource{
|
||||
ID: source.model.ID,
|
||||
PlatformID: source.model.PlatformID,
|
||||
PlatformName: firstNonEmpty(source.platform.Name, source.model.PlatformName),
|
||||
ProviderKey: source.providerKey,
|
||||
ProviderName: source.providerName,
|
||||
ModelName: source.model.ModelName,
|
||||
ModelAlias: source.model.ModelAlias,
|
||||
DisplayName: firstNonEmpty(source.model.DisplayName, source.model.ModelAlias, source.model.ModelName),
|
||||
ModelType: cloneStringSlice(source.model.ModelType),
|
||||
Enabled: source.model.Enabled && source.platform.Status == "enabled",
|
||||
RateLimits: source.rateLimits,
|
||||
ID: source.model.ID,
|
||||
PlatformID: source.model.PlatformID,
|
||||
PlatformName: firstNonEmpty(source.platform.Name, source.model.PlatformName),
|
||||
ProviderKey: source.providerKey,
|
||||
ProviderName: source.providerName,
|
||||
ModelName: firstNonEmpty(source.baseModel.InvocationName, source.model.ModelName),
|
||||
ProviderModelName: source.model.ProviderModelName,
|
||||
ModelAlias: firstNonEmpty(source.baseModel.InvocationName, source.model.ModelName),
|
||||
LegacyAliases: cloneStringSlice(source.baseModel.LegacyAliases),
|
||||
DisplayName: firstNonEmpty(source.baseModel.DisplayName, source.model.DisplayName, source.baseModel.InvocationName, source.model.ModelName),
|
||||
ModelType: cloneStringSlice(source.model.ModelType),
|
||||
Enabled: source.model.Enabled && source.platform.Status == "enabled",
|
||||
RateLimits: source.rateLimits,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -483,18 +487,20 @@ func discountTitle(label string, discount float64) string {
|
||||
}
|
||||
|
||||
func effectiveModelRateLimits(model store.PlatformModel, platform store.Platform, runtimePolicyMap map[string]store.RuntimePolicySet) ModelCatalogRateLimits {
|
||||
overridePolicyRaw, hasOverridePolicy := model.RuntimePolicyOverride["rateLimitPolicy"]
|
||||
overridePolicy := objectValue(overridePolicyRaw)
|
||||
runtimePolicy := map[string]any(nil)
|
||||
if model.RuntimePolicySetID != "" {
|
||||
runtimePolicy = runtimePolicyMap[model.RuntimePolicySetID].RateLimitPolicy
|
||||
}
|
||||
policies := []rateLimitPolicySource{
|
||||
{policy: overridePolicy, authoritative: hasOverridePolicy},
|
||||
{policy: model.RateLimitPolicy, authoritative: len(model.RateLimitPolicy) > 0},
|
||||
{policy: runtimePolicy, authoritative: strings.TrimSpace(model.RuntimePolicySetID) != ""},
|
||||
{policy: platform.RateLimitPolicy},
|
||||
}
|
||||
policy := store.EffectiveRateLimitPolicy(store.EffectiveRateLimitPolicyInput{
|
||||
BasePolicy: model.BaseRateLimitPolicy,
|
||||
PlatformPolicy: platform.RateLimitPolicy,
|
||||
RuntimePolicy: runtimePolicy,
|
||||
RuntimePolicyExplicit: model.RuntimePolicySetID != "",
|
||||
RuntimePolicyOverride: model.RuntimePolicyOverride,
|
||||
ModelPolicy: model.RateLimitPolicy,
|
||||
ModelPolicyMode: model.RateLimitPolicyMode,
|
||||
})
|
||||
policies := []rateLimitPolicySource{{policy: policy, authoritative: policy != nil}}
|
||||
limits := ModelCatalogRateLimits{
|
||||
RPM: firstRateLimit(policies, "rpm"),
|
||||
TPM: firstRateLimit(policies, "tpm_total"),
|
||||
@@ -1260,8 +1266,8 @@ func modelIconPath(model store.PlatformModel, baseModel store.BaseModel, provide
|
||||
)
|
||||
}
|
||||
|
||||
func catalogAliasKey(model store.PlatformModel) string {
|
||||
key := strings.ToLower(strings.TrimSpace(firstNonEmpty(model.ModelAlias, model.DisplayName, model.ModelName)))
|
||||
func catalogIdentityKey(model store.PlatformModel, baseModel store.BaseModel) string {
|
||||
key := strings.ToLower(strings.TrimSpace(firstNonEmpty(baseModel.InvocationName, baseModel.CanonicalModelKey, model.ModelName)))
|
||||
if key == "" {
|
||||
return model.ID
|
||||
}
|
||||
|
||||
@@ -131,13 +131,24 @@ func TestBuildModelCatalogKeepsDisplayNameSeparateFromCallAlias(t *testing.T) {
|
||||
platforms := []store.Platform{
|
||||
{ID: "platform-volces", Provider: "volces", Name: "火山引擎", Status: "enabled"},
|
||||
}
|
||||
baseModels := []store.BaseModel{
|
||||
{
|
||||
ID: "base-seedream-pro",
|
||||
CanonicalModelKey: "seedream:seedream-5.0-pro",
|
||||
InvocationName: "seedream-5.0-pro",
|
||||
ProviderModelName: "doubao-seedream-5-0-pro-260628",
|
||||
DisplayName: "Seedream 5.0 Pro",
|
||||
Status: "active",
|
||||
},
|
||||
}
|
||||
models[0].BaseModelID = "base-seedream-pro"
|
||||
|
||||
response := buildModelCatalog(models, platforms, nil, nil, nil, nil, nil)
|
||||
response := buildModelCatalog(models, platforms, nil, nil, nil, nil, baseModels)
|
||||
if len(response.Items) != 1 {
|
||||
t.Fatalf("expected one Seedream Pro catalog item, got %+v", response.Items)
|
||||
}
|
||||
item := response.Items[0]
|
||||
if item.Alias != "Seedream-5.0-Pro" || item.DisplayName != "Seedream 5.0 Pro" || item.ModelName != "Seedream-5.0-Pro" {
|
||||
if item.Alias != "seedream-5.0-pro" || item.DisplayName != "Seedream 5.0 Pro" || item.ModelName != "seedream-5.0-pro" {
|
||||
t.Fatalf("catalog should separate the call alias from the display name, got %+v", item)
|
||||
}
|
||||
if strings.Contains(item.Alias, " ") {
|
||||
@@ -145,6 +156,46 @@ func TestBuildModelCatalogKeepsDisplayNameSeparateFromCallAlias(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildModelCatalogAggregatesByInvocationNameNotDisplayName(t *testing.T) {
|
||||
models := []store.PlatformModel{
|
||||
{ID: "official-a", PlatformID: "platform-a", BaseModelID: "base-a", ModelName: "provider/a", ModelType: store.StringList{"text_generate"}, DisplayName: "同名营销标题", Enabled: true},
|
||||
{ID: "official-b", PlatformID: "platform-b", BaseModelID: "base-b", ModelName: "provider/b", ModelType: store.StringList{"text_generate"}, DisplayName: "同名营销标题", Enabled: true},
|
||||
{ID: "official-a-mirror", PlatformID: "platform-b", BaseModelID: "base-a-mirror", ModelName: "mirror/a", ModelType: store.StringList{"text_generate"}, DisplayName: "另一个展示标题", Enabled: true},
|
||||
}
|
||||
platforms := []store.Platform{
|
||||
{ID: "platform-a", Provider: "provider-a", Status: "enabled"},
|
||||
{ID: "platform-b", Provider: "provider-b", Status: "enabled"},
|
||||
}
|
||||
baseModels := []store.BaseModel{
|
||||
{ID: "base-a", CanonicalModelKey: "provider-a:official-a", InvocationName: "official-a", DisplayName: "模型 A", Status: "active"},
|
||||
{ID: "base-a-mirror", CanonicalModelKey: "provider-b:official-a", InvocationName: "official-a", DisplayName: "模型 A 镜像", Status: "active"},
|
||||
{ID: "base-b", CanonicalModelKey: "provider-b:official-b", InvocationName: "official-b", DisplayName: "同名营销标题", Status: "active"},
|
||||
}
|
||||
|
||||
response := buildModelCatalog(models, platforms, nil, nil, nil, nil, baseModels)
|
||||
if response.Summary.ModelCount != 2 || response.Summary.SourceCount != 3 {
|
||||
t.Fatalf("expected two official identities and three sources, got %+v", response.Summary)
|
||||
}
|
||||
for _, item := range response.Items {
|
||||
if item.ModelName == "official-a" && item.SourceCount != 2 {
|
||||
t.Fatalf("official-a mirrors should aggregate into one card: %+v", item)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildModelCatalogExcludesDeprecatedBaseModels(t *testing.T) {
|
||||
models := []store.PlatformModel{
|
||||
{ID: "preview", PlatformID: "platform", BaseModelID: "base-preview", ModelName: "gemini-preview", ModelType: store.StringList{"text_generate"}, Enabled: true},
|
||||
}
|
||||
platforms := []store.Platform{{ID: "platform", Provider: "gemini", Status: "enabled"}}
|
||||
baseModels := []store.BaseModel{{ID: "base-preview", InvocationName: "gemini-preview", Status: "deprecated"}}
|
||||
|
||||
response := buildModelCatalog(models, platforms, nil, nil, nil, nil, baseModels)
|
||||
if response.Summary.ModelCount != 0 || response.Summary.SourceCount != 0 {
|
||||
t.Fatalf("deprecated base models must not create new catalog cards: %+v", response.Summary)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildModelCatalogUsesBaseModelProviderForProviderFilters(t *testing.T) {
|
||||
models := []store.PlatformModel{
|
||||
{
|
||||
|
||||
@@ -0,0 +1,39 @@
|
||||
package httpapi
|
||||
|
||||
import (
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
)
|
||||
|
||||
func TestModelCooldownErrorPublicContract(t *testing.T) {
|
||||
recoveryAt := time.Now().UTC().Add(17 * time.Second).Truncate(time.Second)
|
||||
err := &store.ModelCandidateUnavailableError{
|
||||
Code: "model_cooling_down",
|
||||
Message: "请求的模型暂时不可用,请稍后重试",
|
||||
RetryAfter: 17 * time.Second,
|
||||
RecoveryAt: recoveryAt,
|
||||
Details: map[string]any{
|
||||
"retryAfterSeconds": 17,
|
||||
"recoveryAt": recoveryAt.Format(time.RFC3339),
|
||||
},
|
||||
}
|
||||
recorder := httptest.NewRecorder()
|
||||
applyRunErrorHeaders(recorder, err)
|
||||
|
||||
if got := statusFromRunError(err); got != 429 {
|
||||
t.Fatalf("status = %d, want 429", got)
|
||||
}
|
||||
if got := runErrorCode(err); got != "model_cooling_down" {
|
||||
t.Fatalf("code = %q, want model_cooling_down", got)
|
||||
}
|
||||
if got := recorder.Header().Get("Retry-After"); got != "17" {
|
||||
t.Fatalf("Retry-After = %q, want 17", got)
|
||||
}
|
||||
details := runErrorDetails(err)
|
||||
if len(details) != 2 || details["retryAfterSeconds"] != 17 || details["recoveryAt"] != recoveryAt.Format(time.RFC3339) {
|
||||
t.Fatalf("unexpected public details: %#v", details)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,186 @@
|
||||
package httpapi
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
)
|
||||
|
||||
const maxUsageSceneCatalogResponseBytes = 8 << 20
|
||||
|
||||
var supportedModelUsageScenes = map[string]struct{}{
|
||||
"canvas_model_node": {},
|
||||
"desktop": {},
|
||||
}
|
||||
|
||||
func parseModelUsageSceneQuery(query url.Values) (string, error) {
|
||||
values, present := query["usage_scene"]
|
||||
if !present {
|
||||
return "", nil
|
||||
}
|
||||
if len(values) != 1 {
|
||||
return "", errors.New("usage_scene must be provided exactly once")
|
||||
}
|
||||
scene := strings.TrimSpace(values[0])
|
||||
if _, ok := supportedModelUsageScenes[scene]; !ok {
|
||||
return "", errors.New("usage_scene must be canvas_model_node or desktop")
|
||||
}
|
||||
return scene, nil
|
||||
}
|
||||
|
||||
func (s *Server) listServerMainUsageSceneModelIDs(ctx context.Context, usageScene string) (map[string]struct{}, error) {
|
||||
baseURL := strings.TrimRight(strings.TrimSpace(s.cfg.ServerMainBaseURL), "/")
|
||||
if baseURL == "" {
|
||||
return nil, errors.New("SERVER_MAIN_BASE_URL is not configured")
|
||||
}
|
||||
var requestErrors []error
|
||||
for _, endpoint := range serverMainUsageSceneCatalogEndpoints(baseURL, usageScene) {
|
||||
ids, err := s.requestServerMainUsageSceneModelIDs(ctx, endpoint)
|
||||
if err == nil {
|
||||
return ids, nil
|
||||
}
|
||||
requestErrors = append(requestErrors, err)
|
||||
}
|
||||
return nil, fmt.Errorf("server-main model catalog request failed: %w", errors.Join(requestErrors...))
|
||||
}
|
||||
|
||||
func serverMainUsageSceneCatalogEndpoints(baseURL string, usageScene string) []string {
|
||||
path := "/integration-platform/models/strict/all?usage_scene=" + url.QueryEscape(usageScene)
|
||||
if strings.HasSuffix(baseURL, "/api") {
|
||||
return []string{baseURL + path}
|
||||
}
|
||||
return []string{baseURL + "/api" + path, baseURL + path}
|
||||
}
|
||||
|
||||
func (s *Server) requestServerMainUsageSceneModelIDs(ctx context.Context, endpoint string) (map[string]struct{}, error) {
|
||||
request, err := http.NewRequestWithContext(ctx, http.MethodGet, endpoint, nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("create server-main model catalog request: %w", err)
|
||||
}
|
||||
request.Header.Set("Accept", "application/json")
|
||||
|
||||
client := http.DefaultClient
|
||||
if s.auth != nil && s.auth.HTTPClient != nil {
|
||||
client = s.auth.HTTPClient
|
||||
}
|
||||
response, err := client.Do(request)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("request server-main model catalog: %w", err)
|
||||
}
|
||||
defer response.Body.Close()
|
||||
if response.StatusCode < http.StatusOK || response.StatusCode >= http.StatusMultipleChoices {
|
||||
_, _ = io.Copy(io.Discard, io.LimitReader(response.Body, 4<<10))
|
||||
return nil, fmt.Errorf("server-main model catalog returned HTTP %d", response.StatusCode)
|
||||
}
|
||||
|
||||
var payload any
|
||||
decoder := json.NewDecoder(io.LimitReader(response.Body, maxUsageSceneCatalogResponseBytes))
|
||||
if err := decoder.Decode(&payload); err != nil {
|
||||
return nil, fmt.Errorf("decode server-main model catalog: %w", err)
|
||||
}
|
||||
items, ok := usageSceneCatalogItems(payload)
|
||||
if !ok {
|
||||
return nil, errors.New("server-main model catalog payload is malformed")
|
||||
}
|
||||
return usageSceneModelIDs(items), nil
|
||||
}
|
||||
|
||||
func usageSceneCatalogItems(payload any) ([]any, bool) {
|
||||
if items, ok := payload.([]any); ok {
|
||||
return items, true
|
||||
}
|
||||
record, ok := payload.(map[string]any)
|
||||
if !ok {
|
||||
return nil, false
|
||||
}
|
||||
for _, key := range []string{"modelList", "items", "data"} {
|
||||
value, exists := record[key]
|
||||
if !exists {
|
||||
continue
|
||||
}
|
||||
if items, ok := value.([]any); ok {
|
||||
return items, true
|
||||
}
|
||||
if items, ok := usageSceneCatalogItems(value); ok {
|
||||
return items, true
|
||||
}
|
||||
}
|
||||
return nil, false
|
||||
}
|
||||
|
||||
func usageSceneModelIDs(items []any) map[string]struct{} {
|
||||
ids := make(map[string]struct{}, len(items))
|
||||
for _, item := range items {
|
||||
record, ok := item.(map[string]any)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
for _, key := range []string{
|
||||
"alias",
|
||||
"invocationName",
|
||||
"invocation_name",
|
||||
"modelAlias",
|
||||
"model_alias",
|
||||
"name",
|
||||
"modelName",
|
||||
"model_name",
|
||||
"providerModelName",
|
||||
"provider_model_name",
|
||||
} {
|
||||
addUsageSceneModelID(ids, record[key])
|
||||
}
|
||||
for _, key := range []string{"legacyAliases", "legacy_aliases"} {
|
||||
if aliases, ok := record[key].([]any); ok {
|
||||
for _, alias := range aliases {
|
||||
addUsageSceneModelID(ids, alias)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return ids
|
||||
}
|
||||
|
||||
func filterPlatformModelsByUsageSceneIDs(models []store.PlatformModel, availableModelIDs map[string]struct{}) []store.PlatformModel {
|
||||
filtered := make([]store.PlatformModel, 0, len(models))
|
||||
for _, model := range models {
|
||||
if platformModelMatchesUsageSceneIDs(model, availableModelIDs) {
|
||||
filtered = append(filtered, model)
|
||||
}
|
||||
}
|
||||
return filtered
|
||||
}
|
||||
|
||||
func platformModelMatchesUsageSceneIDs(model store.PlatformModel, availableModelIDs map[string]struct{}) bool {
|
||||
for _, value := range []string{model.ModelName, model.ProviderModelName, model.ModelAlias} {
|
||||
if _, ok := availableModelIDs[normalizeUsageSceneModelID(value)]; ok {
|
||||
return true
|
||||
}
|
||||
}
|
||||
for _, alias := range model.LegacyAliases {
|
||||
if _, ok := availableModelIDs[normalizeUsageSceneModelID(alias)]; ok {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func addUsageSceneModelID(ids map[string]struct{}, value any) {
|
||||
text, ok := value.(string)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
if normalized := normalizeUsageSceneModelID(text); normalized != "" {
|
||||
ids[normalized] = struct{}{}
|
||||
}
|
||||
}
|
||||
|
||||
func normalizeUsageSceneModelID(value string) string {
|
||||
return strings.ToLower(strings.TrimSpace(value))
|
||||
}
|
||||
@@ -0,0 +1,160 @@
|
||||
package httpapi
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
"testing"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
)
|
||||
|
||||
func TestParseModelUsageSceneQuery(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
for _, testCase := range []struct {
|
||||
name string
|
||||
query url.Values
|
||||
want string
|
||||
wantErr bool
|
||||
}{
|
||||
{name: "missing", query: url.Values{}, want: ""},
|
||||
{name: "desktop", query: url.Values{"usage_scene": {"desktop"}}, want: "desktop"},
|
||||
{name: "canvas", query: url.Values{"usage_scene": {"canvas_model_node"}}, want: "canvas_model_node"},
|
||||
{name: "empty", query: url.Values{"usage_scene": {""}}, wantErr: true},
|
||||
{name: "unknown", query: url.Values{"usage_scene": {"mobile"}}, wantErr: true},
|
||||
{name: "duplicate", query: url.Values{"usage_scene": {"desktop", "canvas_model_node"}}, wantErr: true},
|
||||
} {
|
||||
testCase := testCase
|
||||
t.Run(testCase.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
got, err := parseModelUsageSceneQuery(testCase.query)
|
||||
if testCase.wantErr {
|
||||
if err == nil {
|
||||
t.Fatal("expected usage_scene validation error")
|
||||
}
|
||||
return
|
||||
}
|
||||
if err != nil {
|
||||
t.Fatalf("parse usage_scene: %v", err)
|
||||
}
|
||||
if got != testCase.want {
|
||||
t.Fatalf("usage_scene = %q, want %q", got, testCase.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestListServerMainUsageSceneModelIDs(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
serverMain := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path != "/api/integration-platform/models/strict/all" {
|
||||
t.Errorf("path = %q", r.URL.Path)
|
||||
}
|
||||
if got := r.URL.Query().Get("usage_scene"); got != "desktop" {
|
||||
t.Errorf("usage_scene = %q", got)
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = io.WriteString(w, `{"data":[{"alias":"Qwen3.7-Plus"},{"invocationName":"GPT-5"},{"legacyAliases":["gpt-latest"]}]}`)
|
||||
}))
|
||||
defer serverMain.Close()
|
||||
|
||||
authenticator := auth.New("secret", serverMain.URL, "")
|
||||
gateway := &Server{
|
||||
cfg: config.Config{ServerMainBaseURL: serverMain.URL},
|
||||
auth: authenticator,
|
||||
logger: slog.New(slog.NewTextHandler(io.Discard, nil)),
|
||||
}
|
||||
ids, err := gateway.listServerMainUsageSceneModelIDs(context.Background(), "desktop")
|
||||
if err != nil {
|
||||
t.Fatalf("list server-main usage scene models: %v", err)
|
||||
}
|
||||
for _, id := range []string{"qwen3.7-plus", "gpt-5", "gpt-latest"} {
|
||||
if _, ok := ids[id]; !ok {
|
||||
t.Errorf("missing normalized model ID %q in %#v", id, ids)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestFilterPlatformModelsByUsageSceneIDs(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
models := []store.PlatformModel{
|
||||
{ID: "allowed-by-invocation", ModelName: "Qwen3.7-Plus", ProviderModelName: "qwen3.7-plus-2026-07-01"},
|
||||
{ID: "allowed-by-provider-name", ModelName: "GPT-5", ProviderModelName: "gpt-5-2026-07-01"},
|
||||
{ID: "allowed-by-legacy-alias", ModelName: "renamed-model", LegacyAliases: store.StringList{"old-model"}},
|
||||
{ID: "desktop-restricted", ModelName: "Qwen3.6-Plus", ProviderModelName: "qwen3.6-plus-2026-04-02"},
|
||||
}
|
||||
allowed := map[string]struct{}{
|
||||
"qwen3.7-plus": {},
|
||||
"gpt-5-2026-07-01": {},
|
||||
"old-model": {},
|
||||
}
|
||||
|
||||
filtered := filterPlatformModelsByUsageSceneIDs(models, allowed)
|
||||
if len(filtered) != 3 {
|
||||
t.Fatalf("filtered model count = %d, want 3: %#v", len(filtered), filtered)
|
||||
}
|
||||
for _, model := range filtered {
|
||||
if model.ID == "desktop-restricted" {
|
||||
t.Fatalf("desktop-restricted model remained in filtered catalog: %#v", filtered)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestListServerMainUsageSceneModelIDsRejectsMalformedPayload(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
serverMain := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = io.WriteString(w, `{"success":true}`)
|
||||
}))
|
||||
defer serverMain.Close()
|
||||
|
||||
gateway := &Server{
|
||||
cfg: config.Config{ServerMainBaseURL: serverMain.URL},
|
||||
auth: auth.New("secret", serverMain.URL, ""),
|
||||
}
|
||||
if _, err := gateway.listServerMainUsageSceneModelIDs(context.Background(), "desktop"); err == nil {
|
||||
t.Fatal("expected malformed upstream payload error")
|
||||
}
|
||||
}
|
||||
|
||||
func TestListServerMainUsageSceneModelIDsFallsBackToDirectServerMainRoute(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
requestedPaths := make([]string, 0, 2)
|
||||
serverMain := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
requestedPaths = append(requestedPaths, r.URL.Path)
|
||||
if r.URL.Path == "/api/integration-platform/models/strict/all" {
|
||||
http.NotFound(w, r)
|
||||
return
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = io.WriteString(w, `[{"alias":"Qwen3.7-Plus"}]`)
|
||||
}))
|
||||
defer serverMain.Close()
|
||||
|
||||
gateway := &Server{
|
||||
cfg: config.Config{ServerMainBaseURL: serverMain.URL},
|
||||
auth: auth.New("secret", serverMain.URL, ""),
|
||||
}
|
||||
ids, err := gateway.listServerMainUsageSceneModelIDs(context.Background(), "desktop")
|
||||
if err != nil {
|
||||
t.Fatalf("list direct server-main usage scene models: %v", err)
|
||||
}
|
||||
if _, ok := ids["qwen3.7-plus"]; !ok {
|
||||
t.Fatalf("missing direct server-main model ID: %#v", ids)
|
||||
}
|
||||
if len(requestedPaths) != 2 ||
|
||||
requestedPaths[0] != "/api/integration-platform/models/strict/all" ||
|
||||
requestedPaths[1] != "/integration-platform/models/strict/all" {
|
||||
t.Fatalf("unexpected server-main route probes: %#v", requestedPaths)
|
||||
}
|
||||
}
|
||||
@@ -189,6 +189,13 @@ type TaskListResponse struct {
|
||||
PageSize int `json:"pageSize" example:"50"`
|
||||
}
|
||||
|
||||
type AdminTaskListResponse struct {
|
||||
Items []store.AdminGatewayTask `json:"items"`
|
||||
Total int `json:"total" example:"42"`
|
||||
Page int `json:"page" example:"1"`
|
||||
PageSize int `json:"pageSize" example:"50"`
|
||||
}
|
||||
|
||||
type TaskParamPreprocessingLogListResponse struct {
|
||||
Items []store.TaskParamPreprocessingLog `json:"items"`
|
||||
}
|
||||
@@ -208,6 +215,17 @@ type FileUploadResponse struct {
|
||||
ContentType string `json:"contentType,omitempty" example:"image/png"`
|
||||
Size int `json:"size,omitempty" example:"1024"`
|
||||
AssetStorage map[string]interface{} `json:"assetStorage,omitempty"`
|
||||
Status string `json:"status" example:"success"`
|
||||
Data FileUploadData `json:"data"`
|
||||
}
|
||||
|
||||
type FileUploadData struct {
|
||||
ID string `json:"id,omitempty" example:"file_abc123"`
|
||||
URL string `json:"url,omitempty" example:"/static/uploaded/upload-abc123.png"`
|
||||
Filename string `json:"filename,omitempty" example:"image.png"`
|
||||
ContentType string `json:"contentType,omitempty" example:"image/png"`
|
||||
Size int `json:"size,omitempty" example:"1024"`
|
||||
AssetStorage map[string]interface{} `json:"assetStorage,omitempty"`
|
||||
}
|
||||
|
||||
type ReplacePlatformModelsRequest struct {
|
||||
@@ -215,9 +233,12 @@ type ReplacePlatformModelsRequest struct {
|
||||
}
|
||||
|
||||
type TaskAcceptedResponse struct {
|
||||
TaskID string `json:"taskId" example:"9f4d8f3d-5f5f-4bb7-a4be-344a9f930e25"`
|
||||
Task store.GatewayTask `json:"task"`
|
||||
Next TaskNextLinks `json:"next"`
|
||||
Status string `json:"status" example:"submitted"`
|
||||
LegacyTaskID string `json:"task_id" example:"9f4d8f3d-5f5f-4bb7-a4be-344a9f930e25"`
|
||||
TaskID string `json:"taskId" example:"9f4d8f3d-5f5f-4bb7-a4be-344a9f930e25"`
|
||||
QueryURL string `json:"query_url" example:"/api/v1/ai/result/9f4d8f3d-5f5f-4bb7-a4be-344a9f930e25"`
|
||||
Task store.GatewayTask `json:"task"`
|
||||
Next TaskNextLinks `json:"next"`
|
||||
}
|
||||
|
||||
type TaskCancelResponse struct {
|
||||
@@ -231,6 +252,56 @@ type TaskCancelResponse struct {
|
||||
type TaskNextLinks struct {
|
||||
Events string `json:"events" example:"/api/v1/tasks/9f4d8f3d-5f5f-4bb7-a4be-344a9f930e25/events"`
|
||||
Detail string `json:"detail" example:"/api/v1/tasks/9f4d8f3d-5f5f-4bb7-a4be-344a9f930e25"`
|
||||
Result string `json:"result" example:"/api/v1/ai/result/9f4d8f3d-5f5f-4bb7-a4be-344a9f930e25"`
|
||||
}
|
||||
|
||||
type EasyAIGeneratedResponse struct {
|
||||
Status string `json:"status" example:"success" enums:"submitted,process,success,failed"`
|
||||
TaskID string `json:"task_id" example:"9f4d8f3d-5f5f-4bb7-a4be-344a9f930e25"`
|
||||
CamelTaskID string `json:"taskId,omitempty" example:"9f4d8f3d-5f5f-4bb7-a4be-344a9f930e25"`
|
||||
UpstreamTaskID string `json:"upstream_task_id,omitempty" example:"provider-task-123"`
|
||||
QueryURL string `json:"query_url,omitempty" example:"/api/v1/ai/result/9f4d8f3d-5f5f-4bb7-a4be-344a9f930e25"`
|
||||
Created int64 `json:"created" example:"1784772000000"`
|
||||
Message string `json:"message,omitempty"`
|
||||
Code string `json:"code,omitempty"`
|
||||
Data []EasyAIMediaOutput `json:"data"`
|
||||
Output []string `json:"output"`
|
||||
OutputContent []EasyAIMediaOutput `json:"output_content"`
|
||||
Cancellable bool `json:"cancellable"`
|
||||
Submitted bool `json:"submitted"`
|
||||
Result map[string]interface{} `json:"result,omitempty"`
|
||||
Usage map[string]interface{} `json:"usage,omitempty"`
|
||||
}
|
||||
|
||||
type EasyAIMediaOutput struct {
|
||||
Type string `json:"type,omitempty" example:"video" enums:"image,video,audio,file,text"`
|
||||
URL string `json:"url,omitempty" example:"https://cdn.example.com/output.mp4"`
|
||||
// B64JSON is retained when result media transfer is disabled.
|
||||
B64JSON string `json:"b64_json,omitempty"`
|
||||
ImageURL string `json:"image_url,omitempty"`
|
||||
VideoURL string `json:"video_url,omitempty"`
|
||||
AudioURL string `json:"audio_url,omitempty"`
|
||||
Content string `json:"content,omitempty"`
|
||||
Text string `json:"text,omitempty"`
|
||||
MimeType string `json:"mime_type,omitempty"`
|
||||
Format string `json:"format,omitempty"`
|
||||
Duration float64 `json:"duration,omitempty"`
|
||||
Width int `json:"width,omitempty"`
|
||||
Height int `json:"height,omitempty"`
|
||||
SourceResolution string `json:"source_resolution,omitempty"`
|
||||
TargetResolution string `json:"target_resolution,omitempty"`
|
||||
Metadata map[string]interface{} `json:"metadata,omitempty"`
|
||||
}
|
||||
|
||||
type OpenAIImagesResponse struct {
|
||||
Created int64 `json:"created" example:"1710000000"`
|
||||
Data []OpenAIImageData `json:"data"`
|
||||
}
|
||||
|
||||
type OpenAIImageData struct {
|
||||
URL string `json:"url,omitempty" example:"https://cdn.example.com/output.png"`
|
||||
B64JSON string `json:"b64_json,omitempty"`
|
||||
RevisedPrompt string `json:"revised_prompt,omitempty"`
|
||||
}
|
||||
|
||||
type TaskAcceptedEvent struct {
|
||||
|
||||
@@ -11,6 +11,7 @@ import (
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
@@ -235,6 +236,47 @@ func TestCleanupExpiredLocalTempAssetsDeletesExpiredStaticFiles(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestCleanupExpiredLocalTempAssetsDeletesNestedBinaryResultsWithIndependentTTL(t *testing.T) {
|
||||
generatedDir := t.TempDir()
|
||||
taskDir := filepath.Join(generatedDir, "results", "task-123")
|
||||
if err := os.MkdirAll(taskDir, 0o750); err != nil {
|
||||
t.Fatalf("create local result fixture dir: %v", err)
|
||||
}
|
||||
oldResult := filepath.Join(taskDir, strings.Repeat("a", 64)+".bin")
|
||||
freshResult := filepath.Join(taskDir, strings.Repeat("b", 64)+".bin")
|
||||
for _, path := range []string{oldResult, freshResult} {
|
||||
if err := os.WriteFile(path, []byte("asset"), 0o640); err != nil {
|
||||
t.Fatalf("write fixture %s: %v", path, err)
|
||||
}
|
||||
}
|
||||
now := time.Now()
|
||||
if err := os.Chtimes(oldResult, now.Add(-25*time.Hour), now.Add(-25*time.Hour)); err != nil {
|
||||
t.Fatalf("age old binary result: %v", err)
|
||||
}
|
||||
if err := os.Chtimes(freshResult, now.Add(-23*time.Hour), now.Add(-23*time.Hour)); err != nil {
|
||||
t.Fatalf("age fresh binary result: %v", err)
|
||||
}
|
||||
server := &Server{
|
||||
cfg: config.Config{
|
||||
LocalGeneratedStorageDir: generatedDir,
|
||||
LocalResultTTLHours: 24,
|
||||
},
|
||||
logger: slog.New(slog.NewTextHandler(io.Discard, nil)),
|
||||
}
|
||||
|
||||
deleted := server.cleanupExpiredLocalTempAssets(context.Background(), now)
|
||||
|
||||
if deleted != 1 {
|
||||
t.Fatalf("expected one expired binary result delete, got %d", deleted)
|
||||
}
|
||||
if _, err := os.Stat(oldResult); !os.IsNotExist(err) {
|
||||
t.Fatalf("old binary result should be deleted, stat err=%v", err)
|
||||
}
|
||||
if _, err := os.Stat(freshResult); err != nil {
|
||||
t.Fatalf("fresh binary result should remain: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRequestConversationKeyPriority(t *testing.T) {
|
||||
request := httptest.NewRequest(http.MethodPost, "/api/v1/chat/completions", nil)
|
||||
request.Header.Set("X-EasyAI-Conversation-ID", "from-header")
|
||||
|
||||
@@ -3,6 +3,7 @@ package httpapi
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
@@ -71,6 +72,10 @@ func (s *Server) updateRunnerPolicy(w http.ResponseWriter, r *http.Request) {
|
||||
writeError(w, http.StatusBadRequest, "invalid json body")
|
||||
return
|
||||
}
|
||||
if err := validateRunnerPolicyInput(input); err != nil {
|
||||
writeError(w, http.StatusBadRequest, err.Error(), "invalid_parameter")
|
||||
return
|
||||
}
|
||||
item, err := s.store.UpsertDefaultRunnerPolicy(r.Context(), input)
|
||||
if err != nil {
|
||||
s.logger.Error("update runner policy failed", "error", err)
|
||||
@@ -80,6 +85,89 @@ func (s *Server) updateRunnerPolicy(w http.ResponseWriter, r *http.Request) {
|
||||
writeJSON(w, http.StatusOK, item)
|
||||
}
|
||||
|
||||
func validateRunnerPolicyInput(input store.RunnerPolicyInput) error {
|
||||
rawRules, exists := input.FailoverPolicy["actionRules"]
|
||||
if !exists {
|
||||
return nil
|
||||
}
|
||||
rules, ok := rawRules.([]any)
|
||||
if !ok {
|
||||
return errors.New("failoverPolicy.actionRules must be an array")
|
||||
}
|
||||
validActions := map[string]bool{
|
||||
"stop": true, "next": true, "cooldown_and_next": true, "demote_and_next": true, "disable_and_next": true,
|
||||
}
|
||||
for index, rawRule := range rules {
|
||||
rule, ok := rawRule.(map[string]any)
|
||||
if !ok {
|
||||
return fmt.Errorf("failoverPolicy.actionRules[%d] must be an object", index)
|
||||
}
|
||||
action := strings.TrimSpace(stringValue(rule["action"]))
|
||||
if !validActions[action] {
|
||||
return fmt.Errorf("failoverPolicy.actionRules[%d].action is invalid", index)
|
||||
}
|
||||
if target, present := rule["target"]; present {
|
||||
targetValue := strings.TrimSpace(stringValue(target))
|
||||
if targetValue != "model" && targetValue != "platform" {
|
||||
return fmt.Errorf("failoverPolicy.actionRules[%d].target is invalid", index)
|
||||
}
|
||||
}
|
||||
if !runnerActionRuleHasMatch(rule) {
|
||||
return fmt.Errorf("failoverPolicy.actionRules[%d] must include at least one match condition", index)
|
||||
}
|
||||
if action == "cooldown_and_next" && positiveJSONInteger(rule["cooldownSeconds"]) <= 0 {
|
||||
return fmt.Errorf("failoverPolicy.actionRules[%d].cooldownSeconds must be a positive integer", index)
|
||||
}
|
||||
if action == "demote_and_next" && positiveJSONInteger(rule["demoteSteps"]) <= 0 {
|
||||
return fmt.Errorf("failoverPolicy.actionRules[%d].demoteSteps must be a positive integer", index)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func runnerActionRuleHasMatch(rule map[string]any) bool {
|
||||
for _, key := range []string{"categories", "codes", "errorCodes", "statusCodes", "keywords"} {
|
||||
switch values := rule[key].(type) {
|
||||
case []any:
|
||||
for _, value := range values {
|
||||
if strings.TrimSpace(fmt.Sprint(value)) != "" {
|
||||
return true
|
||||
}
|
||||
}
|
||||
case []string:
|
||||
for _, value := range values {
|
||||
if strings.TrimSpace(value) != "" {
|
||||
return true
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func positiveJSONInteger(value any) int {
|
||||
switch typed := value.(type) {
|
||||
case int:
|
||||
if typed > 0 {
|
||||
return typed
|
||||
}
|
||||
case int64:
|
||||
if typed > 0 {
|
||||
return int(typed)
|
||||
}
|
||||
case float64:
|
||||
if typed > 0 && typed == float64(int(typed)) {
|
||||
return int(typed)
|
||||
}
|
||||
case json.Number:
|
||||
parsed, err := typed.Int64()
|
||||
if err == nil && parsed > 0 {
|
||||
return int(parsed)
|
||||
}
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
type updatePlatformDynamicPriorityRequest struct {
|
||||
DynamicPriority *int `json:"dynamicPriority" example:"10"`
|
||||
Reset bool `json:"reset" example:"false"`
|
||||
|
||||
@@ -0,0 +1,57 @@
|
||||
package httpapi
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
)
|
||||
|
||||
func TestValidateRunnerPolicyInputActionRules(t *testing.T) {
|
||||
validRule := func() map[string]any {
|
||||
return map[string]any{
|
||||
"categories": []any{"rate_limit"},
|
||||
"action": "cooldown_and_next",
|
||||
"target": "model",
|
||||
"cooldownSeconds": float64(60),
|
||||
}
|
||||
}
|
||||
tests := []struct {
|
||||
name string
|
||||
mutate func(map[string]any)
|
||||
wantErr string
|
||||
}{
|
||||
{name: "valid"},
|
||||
{name: "invalid action", mutate: func(rule map[string]any) { rule["action"] = "rotate" }, wantErr: ".action is invalid"},
|
||||
{name: "invalid target", mutate: func(rule map[string]any) { rule["target"] = "client" }, wantErr: ".target is invalid"},
|
||||
{name: "empty match", mutate: func(rule map[string]any) { rule["categories"] = []any{} }, wantErr: "at least one match condition"},
|
||||
{name: "missing cooldown", mutate: func(rule map[string]any) { delete(rule, "cooldownSeconds") }, wantErr: ".cooldownSeconds must be a positive integer"},
|
||||
{name: "zero cooldown", mutate: func(rule map[string]any) { rule["cooldownSeconds"] = float64(0) }, wantErr: ".cooldownSeconds must be a positive integer"},
|
||||
{name: "fractional cooldown", mutate: func(rule map[string]any) { rule["cooldownSeconds"] = 1.5 }, wantErr: ".cooldownSeconds must be a positive integer"},
|
||||
{name: "missing demote steps", mutate: func(rule map[string]any) {
|
||||
rule["action"] = "demote_and_next"
|
||||
delete(rule, "cooldownSeconds")
|
||||
}, wantErr: ".demoteSteps must be a positive integer"},
|
||||
}
|
||||
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
rule := validRule()
|
||||
if test.mutate != nil {
|
||||
test.mutate(rule)
|
||||
}
|
||||
err := validateRunnerPolicyInput(store.RunnerPolicyInput{
|
||||
FailoverPolicy: map[string]any{"actionRules": []any{rule}},
|
||||
})
|
||||
if test.wantErr == "" {
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected validation error: %v", err)
|
||||
}
|
||||
return
|
||||
}
|
||||
if err == nil || !strings.Contains(err.Error(), test.wantErr) {
|
||||
t.Fatalf("validation error = %v, want substring %q", err, test.wantErr)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -120,8 +120,13 @@ func NewServerWithContext(ctx context.Context, cfg config.Config, db *store.Stor
|
||||
return server.identityRuntime.LegacyJWTEnabled()
|
||||
}
|
||||
server.auth.LocalAPIKeyVerifier = db.VerifyLocalAPIKey
|
||||
server.runner.StartAsyncQueueWorker(ctx)
|
||||
if cfg.AsyncQueueWorkerEnabled {
|
||||
server.runner.StartAsyncQueueWorker(ctx)
|
||||
} else if logger != nil {
|
||||
logger.Info("asynchronous queue worker disabled for this process")
|
||||
}
|
||||
server.runner.StartBillingSettlementWorker(ctx)
|
||||
server.runner.StartTaskHistoryWorkers(ctx)
|
||||
server.startLocalTempAssetCleanup(ctx)
|
||||
server.startOIDCSessionCleanup(ctx)
|
||||
|
||||
@@ -183,6 +188,9 @@ func NewServerWithContext(ctx context.Context, cfg config.Config, db *store.Stor
|
||||
mux.Handle("POST /api/admin/access-rules/batch", server.requireAdmin(auth.PermissionManager, http.HandlerFunc(server.batchAccessRules)))
|
||||
mux.Handle("PATCH /api/admin/access-rules/{ruleID}", server.requireAdmin(auth.PermissionManager, http.HandlerFunc(server.updateAccessRule)))
|
||||
mux.Handle("DELETE /api/admin/access-rules/{ruleID}", server.requireAdmin(auth.PermissionManager, http.HandlerFunc(server.deleteAccessRule)))
|
||||
mux.Handle("GET /api/admin/tasks", server.requireAdmin(auth.PermissionPower, http.HandlerFunc(server.listAdminTasks)))
|
||||
mux.Handle("GET /api/admin/tasks/{taskID}", server.requireAdmin(auth.PermissionPower, http.HandlerFunc(server.getAdminTask)))
|
||||
mux.Handle("GET /api/admin/tasks/{taskID}/param-preprocessing", server.requireAdmin(auth.PermissionPower, http.HandlerFunc(server.adminTaskParamPreprocessing)))
|
||||
mux.Handle("GET /api/v1/api-keys", server.requireUser(auth.PermissionBasic, http.HandlerFunc(server.listAPIKeys)))
|
||||
mux.Handle("POST /api/v1/api-keys", server.requireUser(auth.PermissionBasic, http.HandlerFunc(server.createAPIKey)))
|
||||
mux.Handle("GET /api/v1/api-keys/access-rules", server.requireUser(auth.PermissionBasic, http.HandlerFunc(server.listAPIKeyAccessRules)))
|
||||
@@ -268,7 +276,7 @@ func NewServerWithContext(ctx context.Context, cfg config.Config, db *store.Stor
|
||||
mux.Handle("POST /api/v1/videos/upscales", server.requireUser(auth.PermissionBasic, server.createVideoUpscaleTask()))
|
||||
mux.Handle("POST /api/v1/video/upscale", server.requireUser(auth.PermissionBasic, server.createVideoUpscaleTask()))
|
||||
mux.Handle("POST /api/v1/video/generations", server.requireUser(auth.PermissionBasic, http.HandlerFunc(server.createLegacyVolcesVideoGeneration)))
|
||||
mux.Handle("GET /api/v1/ai/result/{taskID}", server.requireUser(auth.PermissionBasic, http.HandlerFunc(server.getLegacyVolcesVideoResult)))
|
||||
mux.Handle("GET /api/v1/ai/result/{taskID}", server.requireUser(auth.PermissionBasic, http.HandlerFunc(server.getEasyAITaskResult)))
|
||||
mux.Handle("POST /api/v1/song/generations", server.requireUser(auth.PermissionBasic, server.createTask("song.generations", true)))
|
||||
mux.Handle("POST /api/v1/music/generations", server.requireUser(auth.PermissionBasic, server.createTask("music.generations", true)))
|
||||
mux.Handle("POST /api/v1/speech/generations", server.requireUser(auth.PermissionBasic, server.createTask("speech.generations", true)))
|
||||
|
||||
@@ -1,8 +1,12 @@
|
||||
package httpapi
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
)
|
||||
|
||||
func TestOptionalTaskIdempotencyKeyRejectsMultipleValues(t *testing.T) {
|
||||
@@ -29,3 +33,29 @@ func TestTaskIdempotencyRequestHashIsCanonical(t *testing.T) {
|
||||
t.Fatal("async response semantics must be part of the request hash")
|
||||
}
|
||||
}
|
||||
|
||||
func TestWriteIdempotentAsyncTaskReplayKeepsEasyAICompatibilityFields(t *testing.T) {
|
||||
t.Parallel()
|
||||
task := store.GatewayTask{
|
||||
ID: "async-replay-1",
|
||||
Kind: "images.generations",
|
||||
Status: "queued",
|
||||
AsyncMode: true,
|
||||
}
|
||||
recorder := httptest.NewRecorder()
|
||||
|
||||
writeIdempotentTaskReplay(recorder, task, true)
|
||||
|
||||
if recorder.Code != http.StatusAccepted ||
|
||||
recorder.Header().Get("Idempotent-Replayed") != "true" ||
|
||||
recorder.Header().Get("X-Gateway-Task-Id") != task.ID {
|
||||
t.Fatalf("unexpected replay status/headers: code=%d headers=%v", recorder.Code, recorder.Header())
|
||||
}
|
||||
var got map[string]any
|
||||
if err := json.NewDecoder(recorder.Body).Decode(&got); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got["status"] != "submitted" || got["task_id"] != task.ID || got["taskId"] != task.ID {
|
||||
t.Fatalf("unexpected replay body: %+v", got)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -61,6 +61,7 @@ func (s *Server) deleteClonedVoice(w http.ResponseWriter, r *http.Request) {
|
||||
}
|
||||
result, err := s.runner.DeleteClonedVoice(r.Context(), user, r.PathValue("voiceID"))
|
||||
if err != nil {
|
||||
applyRunErrorHeaders(w, err)
|
||||
writeErrorWithDetails(w, statusFromRunError(err), runErrorMessage(err), runErrorDetails(err), runErrorCode(err))
|
||||
return
|
||||
}
|
||||
|
||||
@@ -56,12 +56,19 @@ func (s *Server) createVolcesContentsGenerationTask(w http.ResponseWriter, r *ht
|
||||
// @Param taskID path string true "任务 ID"
|
||||
// @Success 200 {object} VolcesContentsGenerationTaskResponse
|
||||
// @Failure 404 {object} VolcesErrorEnvelope
|
||||
// @Failure 410 {object} VolcesErrorEnvelope
|
||||
// @Router /api/v1/contents/generations/tasks/{taskID} [get]
|
||||
func (s *Server) getVolcesContentsGenerationTask(w http.ResponseWriter, r *http.Request) {
|
||||
task, ok := s.volcesCompatibleTaskForUser(w, r)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var err error
|
||||
task, err = s.hydrateTaskResult(r.Context(), task)
|
||||
if err != nil {
|
||||
writeVolcesError(w, statusFromRunError(err), err.Error(), clients.ErrorCode(err))
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, volcesCompatibleTask(task))
|
||||
}
|
||||
|
||||
@@ -143,11 +150,11 @@ func (s *Server) deleteVolcesContentsGenerationTask(w http.ResponseWriter, r *ht
|
||||
// createLegacyVolcesVideoGeneration godoc
|
||||
// @Summary 创建 server-main 兼容视频任务
|
||||
// @Description 兼容 server-main 的 /api/v1/video/generations,返回 submitted 和 task_id;额外保留火山任务字段。
|
||||
// @Tags volces-compatible
|
||||
// @Tags media
|
||||
// @Accept json
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Success 200 {object} map[string]any
|
||||
// @Success 200 {object} TaskAcceptedResponse
|
||||
// @Router /api/v1/video/generations [post]
|
||||
func (s *Server) createLegacyVolcesVideoGeneration(w http.ResponseWriter, r *http.Request) {
|
||||
user, ok := auth.UserFromContext(r.Context())
|
||||
@@ -166,35 +173,46 @@ func (s *Server) createLegacyVolcesVideoGeneration(w http.ResponseWriter, r *htt
|
||||
return
|
||||
}
|
||||
response := volcesCompatibleTask(task)
|
||||
response["status"] = "submitted"
|
||||
response["task_id"] = task.ID
|
||||
writeJSON(w, http.StatusOK, response)
|
||||
writeJSON(w, http.StatusOK, easyAITaskAcceptedResponseWithExtensions(task, response))
|
||||
}
|
||||
|
||||
// getLegacyVolcesVideoResult godoc
|
||||
// @Summary 查询 server-main 兼容视频结果
|
||||
// @Tags volces-compatible
|
||||
// getEasyAITaskResult godoc
|
||||
// @Summary 查询 server-main EasyAIClient 兼容任务结果
|
||||
// @Description 查询当前用户的图片、视频、音频、视频高清和矢量化任务,并返回 EasyAI GeneratedResponse 兼容结构。
|
||||
// @Tags tasks
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param taskID path string true "任务 ID"
|
||||
// @Success 200 {object} map[string]any
|
||||
// @Success 200 {object} EasyAIGeneratedResponse
|
||||
// @Failure 404 {object} EasyAIGeneratedResponse
|
||||
// @Failure 410 {object} EasyAIGeneratedResponse
|
||||
// @Router /api/v1/ai/result/{taskID} [get]
|
||||
func (s *Server) getLegacyVolcesVideoResult(w http.ResponseWriter, r *http.Request) {
|
||||
task, ok := s.volcesCompatibleTaskForUser(w, r)
|
||||
if !ok {
|
||||
func (s *Server) getEasyAITaskResult(w http.ResponseWriter, r *http.Request) {
|
||||
user, ok := auth.UserFromContext(r.Context())
|
||||
if !ok || user == nil {
|
||||
writeEasyAIAsyncError(w, http.StatusUnauthorized, "unauthorized", nil, "unauthorized")
|
||||
return
|
||||
}
|
||||
compat := volcesCompatibleTask(task)
|
||||
legacyStatus := "process"
|
||||
switch compat["status"] {
|
||||
case "succeeded":
|
||||
legacyStatus = "success"
|
||||
case "failed", "cancelled":
|
||||
legacyStatus = "failed"
|
||||
task, err := s.store.GetTask(r.Context(), strings.TrimSpace(r.PathValue("taskID")))
|
||||
if err != nil {
|
||||
if store.IsNotFound(err) {
|
||||
writeEasyAIAsyncError(w, http.StatusNotFound, "task not found", nil, "not_found")
|
||||
return
|
||||
}
|
||||
s.logger.Error("get EasyAI-compatible task failed", "error", err)
|
||||
writeError(w, http.StatusInternalServerError, "get task failed", "internal_error")
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]any{
|
||||
"status": legacyStatus, "task_id": task.ID, "data": compat["content"], "result": compat,
|
||||
})
|
||||
if !runner.TaskAccessibleToUser(task, user) {
|
||||
writeEasyAIAsyncError(w, http.StatusNotFound, "task not found", nil, "not_found")
|
||||
return
|
||||
}
|
||||
task, err = s.hydrateTaskResult(r.Context(), task)
|
||||
if err != nil {
|
||||
writeEasyAIAsyncError(w, statusFromRunError(err), err.Error(), nil, clients.ErrorCode(err))
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, easyAITaskResultResponse(task))
|
||||
}
|
||||
|
||||
func (s *Server) createVolcesCompatibleTask(r *http.Request, user *auth.User, body map[string]any) (store.GatewayTask, error) {
|
||||
@@ -218,7 +236,6 @@ func (s *Server) createVolcesCompatibleTask(r *http.Request, user *auth.User, bo
|
||||
return store.GatewayTask{}, err
|
||||
}
|
||||
task.CompatibilityProtocol = clients.ProtocolVolcesContents
|
||||
task.CompatibilityPublicID = task.ID
|
||||
if err := s.runner.EnqueueAsyncTask(r.Context(), task); err != nil {
|
||||
return store.GatewayTask{}, &clients.ClientError{Code: "enqueue_failed", Message: err.Error(), StatusCode: http.StatusServiceUnavailable, Retryable: true}
|
||||
}
|
||||
@@ -253,7 +270,6 @@ func (s *Server) waitForCompatibilitySubmission(r *http.Request, task store.Gate
|
||||
Message: firstNonEmpty(current.ErrorMessage, current.Error, current.Message),
|
||||
StatusCode: status,
|
||||
Retryable: false,
|
||||
Wire: compatibilitySubmissionWire(current),
|
||||
}
|
||||
}
|
||||
select {
|
||||
@@ -266,34 +282,6 @@ func (s *Server) waitForCompatibilitySubmission(r *http.Request, task store.Gate
|
||||
}
|
||||
}
|
||||
|
||||
func compatibilitySubmissionWire(task store.GatewayTask) *clients.WireResponse {
|
||||
if task.CompatibilitySubmitHTTPStatus == 0 || strings.TrimSpace(task.CompatibilitySourceProtocol) == "" {
|
||||
return nil
|
||||
}
|
||||
headers := map[string][]string{}
|
||||
for name, value := range task.CompatibilitySubmitHeaders {
|
||||
switch typed := value.(type) {
|
||||
case []any:
|
||||
for _, item := range typed {
|
||||
if text := strings.TrimSpace(volcesCompatString(item)); text != "" {
|
||||
headers[name] = append(headers[name], text)
|
||||
}
|
||||
}
|
||||
case []string:
|
||||
headers[name] = append([]string(nil), typed...)
|
||||
case string:
|
||||
headers[name] = []string{typed}
|
||||
}
|
||||
}
|
||||
return &clients.WireResponse{
|
||||
Protocol: task.CompatibilitySourceProtocol,
|
||||
StatusCode: task.CompatibilitySubmitHTTPStatus,
|
||||
Headers: headers,
|
||||
Body: cloneVolcesCompatibleMap(task.CompatibilitySubmitBody),
|
||||
Converted: task.CompatibilitySourceProtocol != task.CompatibilityProtocol,
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Server) volcesCompatibleTaskForUser(w http.ResponseWriter, r *http.Request) (store.GatewayTask, bool) {
|
||||
user, ok := auth.UserFromContext(r.Context())
|
||||
if !ok || user == nil {
|
||||
@@ -322,15 +310,30 @@ func isVolcesCompatibleTask(task store.GatewayTask) bool {
|
||||
}
|
||||
|
||||
func volcesCompatibleTask(task store.GatewayTask) map[string]any {
|
||||
if volcesTaskUsesNativeProtocol(task) {
|
||||
if raw := cloneVolcesCompatibleMap(mapFromVolcesCompat(task.Result["raw"])); len(raw) > 0 {
|
||||
return raw
|
||||
response := map[string]any{}
|
||||
for _, key := range []string{"model", "seed", "resolution", "ratio", "duration", "frames", "framespersecond"} {
|
||||
if task.Result[key] != nil {
|
||||
response[key] = task.Result[key]
|
||||
}
|
||||
if raw := cloneVolcesCompatibleMap(task.RemoteTaskPayload); len(raw) > 0 {
|
||||
return raw
|
||||
}
|
||||
if content, ok := task.Result["content"].(map[string]any); ok && len(content) > 0 {
|
||||
// Transitional read support for historical canonical results. New
|
||||
// results only persist data[] and reconstruct this protocol field.
|
||||
response["content"] = content
|
||||
} else if data, ok := task.Result["data"].([]any); ok && len(data) > 0 {
|
||||
if item, ok := data[0].(map[string]any); ok {
|
||||
content := map[string]any{}
|
||||
if videoURL := firstNonEmpty(volcesCompatString(item["url"]), volcesCompatString(item["video_url"])); videoURL != "" {
|
||||
content["video_url"] = videoURL
|
||||
}
|
||||
if lastFrameURL := volcesCompatString(item["last_frame_url"]); lastFrameURL != "" {
|
||||
content["last_frame_url"] = lastFrameURL
|
||||
}
|
||||
if len(content) > 0 {
|
||||
response["content"] = content
|
||||
}
|
||||
}
|
||||
}
|
||||
response := map[string]any{}
|
||||
response["id"] = volcesCompatiblePublicID(task)
|
||||
response["model"] = firstNonEmpty(volcesCompatString(response["model"]), task.Model)
|
||||
response["status"] = volcesCompatibleTaskStatus(task.Status)
|
||||
@@ -341,8 +344,10 @@ func volcesCompatibleTask(task store.GatewayTask) map[string]any {
|
||||
response[key] = task.Request[key]
|
||||
}
|
||||
}
|
||||
if len(task.Usage) > 0 && response["usage"] == nil {
|
||||
response["usage"] = task.Usage
|
||||
if usage := volcesCompatibleUsage(task.Usage); len(usage) > 0 {
|
||||
response["usage"] = usage
|
||||
} else if legacyUsage, ok := task.Result["usage"].(map[string]any); ok && len(legacyUsage) > 0 {
|
||||
response["usage"] = legacyUsage
|
||||
}
|
||||
if task.Status == "failed" || task.Status == "cancelled" {
|
||||
response["error"] = map[string]any{"code": firstNonEmpty(task.ErrorCode, strings.ToUpper(task.Status)), "message": firstNonEmpty(task.ErrorMessage, task.Error, task.Message)}
|
||||
@@ -350,19 +355,28 @@ func volcesCompatibleTask(task store.GatewayTask) map[string]any {
|
||||
return response
|
||||
}
|
||||
|
||||
func volcesCompatibleCreateResponse(task store.GatewayTask) map[string]any {
|
||||
if volcesTaskUsesNativeProtocol(task) {
|
||||
if raw := cloneVolcesCompatibleMap(task.RemoteTaskPayload); len(raw) > 0 {
|
||||
return raw
|
||||
func volcesCompatibleUsage(usage map[string]any) map[string]any {
|
||||
out := map[string]any{}
|
||||
for outputKey, inputKeys := range map[string][]string{
|
||||
"prompt_tokens": {"inputTokens", "promptTokens"},
|
||||
"completion_tokens": {"outputTokens", "completionTokens"},
|
||||
"total_tokens": {"totalTokens"},
|
||||
} {
|
||||
for _, inputKey := range inputKeys {
|
||||
if value := usage[inputKey]; value != nil {
|
||||
out[outputKey] = value
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func volcesCompatibleCreateResponse(task store.GatewayTask) map[string]any {
|
||||
return map[string]any{"id": volcesCompatiblePublicID(task)}
|
||||
}
|
||||
|
||||
func volcesCompatiblePublicID(task store.GatewayTask) string {
|
||||
if strings.TrimSpace(task.CompatibilityPublicID) != "" {
|
||||
return strings.TrimSpace(task.CompatibilityPublicID)
|
||||
}
|
||||
if volcesTaskUsesNativeProtocol(task) && strings.TrimSpace(task.RemoteTaskID) != "" {
|
||||
return strings.TrimSpace(task.RemoteTaskID)
|
||||
}
|
||||
@@ -381,11 +395,6 @@ func volcesTaskUsesNativeProtocol(task store.GatewayTask) bool {
|
||||
return false
|
||||
}
|
||||
|
||||
func mapFromVolcesCompat(value any) map[string]any {
|
||||
result, _ := value.(map[string]any)
|
||||
return result
|
||||
}
|
||||
|
||||
func volcesCompatibleTaskStatus(status string) string {
|
||||
switch strings.ToLower(strings.TrimSpace(status)) {
|
||||
case "succeeded", "success", "completed":
|
||||
@@ -401,21 +410,6 @@ func volcesCompatibleTaskStatus(status string) string {
|
||||
}
|
||||
}
|
||||
|
||||
func cloneVolcesCompatibleMap(source map[string]any) map[string]any {
|
||||
if len(source) == 0 {
|
||||
return nil
|
||||
}
|
||||
raw, err := json.Marshal(source)
|
||||
if err != nil {
|
||||
return map[string]any{}
|
||||
}
|
||||
var out map[string]any
|
||||
if err := json.Unmarshal(raw, &out); err != nil {
|
||||
return map[string]any{}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func writeVolcesCompatibleTaskError(w http.ResponseWriter, err error) {
|
||||
status := http.StatusInternalServerError
|
||||
var staged *gatewayTaskCreationError
|
||||
@@ -427,6 +421,8 @@ func writeVolcesCompatibleTaskError(w http.ResponseWriter, err error) {
|
||||
status = clientErr.StatusCode
|
||||
} else if errors.As(err, &clientErr) {
|
||||
status = http.StatusBadRequest
|
||||
} else if errors.Is(err, store.ErrTaskRequestBinaryNotMaterialized) {
|
||||
status = http.StatusBadRequest
|
||||
}
|
||||
if wire := clients.ErrorWireResponse(err); wireResponseMatches(wire, clients.ProtocolVolcesContents) {
|
||||
writeWireResponse(w, wire)
|
||||
|
||||
@@ -7,23 +7,22 @@ import (
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
)
|
||||
|
||||
func TestVolcesCompatibleTaskPreservesNativeOfficialFieldsWithoutGatewayExtensions(t *testing.T) {
|
||||
func TestVolcesCompatibleTaskBuildsProtocolResponseFromCanonicalResult(t *testing.T) {
|
||||
now := time.Date(2026, 7, 18, 8, 0, 0, 0, time.UTC)
|
||||
task := store.GatewayTask{
|
||||
ID: "gateway-task-1", Kind: "videos.generations", Status: "succeeded", Model: "doubao-seedance-2-0-mini-260615",
|
||||
RemoteTaskID: "cgt-upstream-1", CreatedAt: now, UpdatedAt: now.Add(time.Second),
|
||||
Result: map[string]any{
|
||||
"raw": map[string]any{
|
||||
"id": "cgt-upstream-1", "model": "doubao-seedance-2-0-mini-260615", "status": "succeeded",
|
||||
"content": map[string]any{"video_url": "https://example.com/out.mp4"}, "usage": map[string]any{"total_tokens": 9},
|
||||
"future_official_field": "preserved",
|
||||
},
|
||||
"id": "cgt-upstream-1", "model": "doubao-seedance-2-0-mini-260615", "status": "succeeded",
|
||||
"data": []any{map[string]any{"url": "https://example.com/out.mp4", "type": "video"}},
|
||||
"future_official_field": "not-whitelisted",
|
||||
},
|
||||
Usage: map[string]any{"totalTokens": 9},
|
||||
Billings: []any{map[string]any{"amount": 3}}, BillingSummary: map[string]any{"currency": "resource"}, FinalChargeAmount: 3,
|
||||
Attempts: []store.TaskAttempt{{Provider: "volces"}},
|
||||
}
|
||||
got := volcesCompatibleTask(task)
|
||||
if got["id"] != task.RemoteTaskID || got["status"] != "succeeded" || got["future_official_field"] != "preserved" {
|
||||
if got["id"] != task.RemoteTaskID || got["status"] != "succeeded" || got["future_official_field"] != nil {
|
||||
t.Fatalf("unexpected compatibility identity/status: %+v", got)
|
||||
}
|
||||
content, _ := got["content"].(map[string]any)
|
||||
|
||||
@@ -0,0 +1,721 @@
|
||||
package runner
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/base64"
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
"syscall"
|
||||
"time"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/clients"
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
|
||||
)
|
||||
|
||||
const (
|
||||
localBinaryResultDirName = "results"
|
||||
localBinaryPlaceholderPrefix = "[GatewayBinary:v1;"
|
||||
localBinaryGenericBase64MinLength = 4096
|
||||
localBinaryMaxDepth = 64
|
||||
defaultLocalResultTTLHours = 24
|
||||
defaultLocalResultMinFreeBytes = int64(10 * 1024 * 1024 * 1024)
|
||||
defaultLocalResultMaxBytes = int64(256 * 1024 * 1024)
|
||||
defaultLocalResultMaxTaskBytes = int64(512 * 1024 * 1024)
|
||||
)
|
||||
|
||||
type localBinaryDescriptor struct {
|
||||
Prefix string
|
||||
SHA256 string
|
||||
Size int64
|
||||
ContentType string
|
||||
Encoding string
|
||||
}
|
||||
|
||||
type localBinaryMaterializer struct {
|
||||
service *Service
|
||||
taskDir string
|
||||
writeFiles bool
|
||||
enforceLimits bool
|
||||
totalBytes int64
|
||||
seen map[string]struct{}
|
||||
createdFiles []string
|
||||
}
|
||||
|
||||
// materializeLocalBinaryResult replaces every inline binary value with a
|
||||
// bounded placeholder after atomically writing and verifying the bytes locally.
|
||||
func (s *Service) materializeLocalBinaryResult(ctx context.Context, taskID string, result map[string]any) (map[string]any, bool, error) {
|
||||
if _, _, err := s.transformLocalBinaryResult(ctx, taskID, result, false, true); err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
return s.transformLocalBinaryResult(ctx, taskID, result, true, true)
|
||||
}
|
||||
|
||||
func (s *Service) transformLocalBinaryResult(ctx context.Context, taskID string, result map[string]any, writeFiles bool, enforceLimits bool) (map[string]any, bool, error) {
|
||||
root := s.localBinaryResultRoot()
|
||||
taskDir := filepath.Join(root, safeLocalBinaryTaskDir(taskID))
|
||||
materializer := &localBinaryMaterializer{
|
||||
service: s,
|
||||
taskDir: taskDir,
|
||||
writeFiles: writeFiles,
|
||||
enforceLimits: enforceLimits,
|
||||
seen: map[string]struct{}{},
|
||||
}
|
||||
next, changed, err := materializer.materializeValue(ctx, result, "", nil, 0)
|
||||
if err != nil {
|
||||
materializer.rollbackCreatedFiles()
|
||||
return nil, false, err
|
||||
}
|
||||
mapped, ok := next.(map[string]any)
|
||||
if !ok {
|
||||
materializer.rollbackCreatedFiles()
|
||||
return nil, false, &clients.ClientError{
|
||||
Code: "result_binary_not_materialized",
|
||||
Message: "generated result is not a JSON object",
|
||||
StatusCode: 500,
|
||||
Retryable: false,
|
||||
}
|
||||
}
|
||||
if changed && enforceLimits && !writeFiles {
|
||||
if err := os.MkdirAll(root, 0o750); err != nil {
|
||||
return nil, false, localBinaryStorageError(err)
|
||||
}
|
||||
if err := ensureLocalBinaryDiskHeadroom(root, materializer.totalBytes, s.localResultMinFreeBytes()); err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
}
|
||||
return mapped, changed, nil
|
||||
}
|
||||
|
||||
// MaterializeTaskResultForStorage exposes the same verified materialization
|
||||
// path to the explicit historical maintenance command.
|
||||
func (s *Service) MaterializeTaskResultForStorage(ctx context.Context, taskID string, result map[string]any) (map[string]any, bool, error) {
|
||||
return s.materializeLocalBinaryResult(ctx, taskID, result)
|
||||
}
|
||||
|
||||
// CompactExpiredTaskResultForStorage computes the same deterministic
|
||||
// placeholders without creating files that are already outside the recovery
|
||||
// window.
|
||||
func (s *Service) CompactExpiredTaskResultForStorage(ctx context.Context, taskID string, result map[string]any) (map[string]any, bool, error) {
|
||||
return s.transformLocalBinaryResult(ctx, taskID, result, false, false)
|
||||
}
|
||||
|
||||
func TaskResultHasInlineBinary(result map[string]any) bool {
|
||||
return localBinaryValueHasPayload(result, "", nil, 0)
|
||||
}
|
||||
|
||||
func localBinaryValueHasPayload(value any, key string, siblings map[string]any, depth int) bool {
|
||||
if depth >= localBinaryMaxDepth {
|
||||
return false
|
||||
}
|
||||
switch typed := value.(type) {
|
||||
case map[string]any:
|
||||
if _, _, ok := localBufferObjectBytes(typed); ok {
|
||||
return true
|
||||
}
|
||||
for childKey, child := range typed {
|
||||
if localBinaryValueHasPayload(child, childKey, typed, depth+1) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
case []any:
|
||||
if localBinaryKey(key) {
|
||||
if _, ok := bytesFromNumberArray(typed); ok {
|
||||
return true
|
||||
}
|
||||
}
|
||||
for _, child := range typed {
|
||||
if localBinaryValueHasPayload(child, key, siblings, depth+1) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
case []byte:
|
||||
return len(typed) > 0
|
||||
case string:
|
||||
_, _, _, ok := localBinaryStringBytes(key, typed, siblings)
|
||||
return ok
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func (m *localBinaryMaterializer) materializeValue(ctx context.Context, value any, key string, siblings map[string]any, depth int) (any, bool, error) {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
if depth >= localBinaryMaxDepth {
|
||||
return nil, false, &clients.ClientError{
|
||||
Code: "result_binary_not_materialized",
|
||||
Message: "generated result exceeds the maximum JSON depth",
|
||||
StatusCode: 500,
|
||||
Retryable: false,
|
||||
}
|
||||
}
|
||||
switch typed := value.(type) {
|
||||
case map[string]any:
|
||||
if payload, contentType, ok := localBufferObjectBytes(typed); ok {
|
||||
return m.persistBinary(ctx, payload, contentType, "buffer")
|
||||
}
|
||||
next := make(map[string]any, len(typed))
|
||||
changed := false
|
||||
for childKey, childValue := range typed {
|
||||
child, childChanged, err := m.materializeValue(ctx, childValue, childKey, typed, depth+1)
|
||||
if err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
next[childKey] = child
|
||||
changed = changed || childChanged
|
||||
}
|
||||
if !changed {
|
||||
return value, false, nil
|
||||
}
|
||||
return next, true, nil
|
||||
case []any:
|
||||
if localBinaryKey(key) {
|
||||
if payload, ok := bytesFromNumberArray(typed); ok {
|
||||
return m.persistBinary(ctx, payload, mediaContentTypeFromItem(siblings), "buffer")
|
||||
}
|
||||
}
|
||||
next := make([]any, len(typed))
|
||||
changed := false
|
||||
for index, item := range typed {
|
||||
child, childChanged, err := m.materializeValue(ctx, item, key, siblings, depth+1)
|
||||
if err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
next[index] = child
|
||||
changed = changed || childChanged
|
||||
}
|
||||
if !changed {
|
||||
return value, false, nil
|
||||
}
|
||||
return next, true, nil
|
||||
case []byte:
|
||||
if len(typed) == 0 {
|
||||
return value, false, nil
|
||||
}
|
||||
return m.persistBinary(ctx, append([]byte(nil), typed...), mediaContentTypeFromItem(siblings), "buffer")
|
||||
case string:
|
||||
payload, contentType, encoding, ok := localBinaryStringBytes(key, typed, siblings)
|
||||
if !ok {
|
||||
return value, false, nil
|
||||
}
|
||||
return m.persistBinary(ctx, payload, contentType, encoding)
|
||||
default:
|
||||
return value, false, nil
|
||||
}
|
||||
}
|
||||
|
||||
func (m *localBinaryMaterializer) persistBinary(ctx context.Context, payload []byte, contentType string, encoding string) (any, bool, error) {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
if len(payload) == 0 {
|
||||
return nil, false, nil
|
||||
}
|
||||
size := int64(len(payload))
|
||||
if m.enforceLimits && size > m.service.localResultMaxBytes() {
|
||||
return nil, false, &clients.ClientError{
|
||||
Code: "binary_result_too_large",
|
||||
Message: "one generated binary result exceeds the local storage limit",
|
||||
StatusCode: 502,
|
||||
Retryable: false,
|
||||
}
|
||||
}
|
||||
digest := sha256.Sum256(payload)
|
||||
digestHex := hex.EncodeToString(digest[:])
|
||||
if _, exists := m.seen[digestHex]; !exists {
|
||||
if m.enforceLimits && m.totalBytes+size > m.service.localResultMaxTaskBytes() {
|
||||
return nil, false, &clients.ClientError{
|
||||
Code: "binary_result_too_large",
|
||||
Message: "generated binary results exceed the per-task local storage limit",
|
||||
StatusCode: 502,
|
||||
Retryable: false,
|
||||
}
|
||||
}
|
||||
if m.writeFiles {
|
||||
created, err := m.service.writeLocalBinaryResult(m.taskDir, digestHex, payload)
|
||||
if err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
if created {
|
||||
m.createdFiles = append(m.createdFiles, filepath.Join(m.taskDir, digestHex+".bin"))
|
||||
}
|
||||
}
|
||||
m.seen[digestHex] = struct{}{}
|
||||
m.totalBytes += size
|
||||
}
|
||||
descriptor := localBinaryDescriptor{
|
||||
Prefix: localBinaryPrefix(payload),
|
||||
SHA256: digestHex,
|
||||
Size: size,
|
||||
ContentType: normalizedLocalBinaryContentType(contentType),
|
||||
Encoding: normalizedLocalBinaryEncoding(encoding),
|
||||
}
|
||||
return localBinaryPlaceholder(descriptor), true, nil
|
||||
}
|
||||
|
||||
func (m *localBinaryMaterializer) rollbackCreatedFiles() {
|
||||
for _, path := range m.createdFiles {
|
||||
_ = os.Remove(path)
|
||||
}
|
||||
_ = os.Remove(m.taskDir)
|
||||
m.createdFiles = nil
|
||||
}
|
||||
|
||||
func (s *Service) writeLocalBinaryResult(taskDir string, digestHex string, payload []byte) (bool, error) {
|
||||
if err := os.MkdirAll(taskDir, 0o750); err != nil {
|
||||
return false, localBinaryStorageError(err)
|
||||
}
|
||||
targetPath := filepath.Join(taskDir, digestHex+".bin")
|
||||
if info, err := os.Stat(targetPath); err == nil {
|
||||
if !info.IsDir() && info.Size() == int64(len(payload)) {
|
||||
if err := verifyLocalBinaryFile(targetPath, digestHex, int64(len(payload))); err == nil {
|
||||
now := time.Now()
|
||||
if err := os.Chtimes(targetPath, now, now); err != nil {
|
||||
return false, localBinaryStorageError(err)
|
||||
}
|
||||
return false, nil
|
||||
}
|
||||
}
|
||||
return false, &clients.ClientError{
|
||||
Code: "binary_result_corrupted",
|
||||
Message: "existing local result file does not match its content hash",
|
||||
StatusCode: 500,
|
||||
Retryable: false,
|
||||
}
|
||||
} else if !errors.Is(err, os.ErrNotExist) {
|
||||
return false, localBinaryStorageError(err)
|
||||
}
|
||||
if err := ensureLocalBinaryDiskHeadroom(taskDir, int64(len(payload)), s.localResultMinFreeBytes()); err != nil {
|
||||
return false, err
|
||||
}
|
||||
tempFile, err := os.CreateTemp(taskDir, ".gateway-result-*")
|
||||
if err != nil {
|
||||
return false, localBinaryStorageError(err)
|
||||
}
|
||||
tempPath := tempFile.Name()
|
||||
cleanup := func() {
|
||||
_ = tempFile.Close()
|
||||
_ = os.Remove(tempPath)
|
||||
}
|
||||
if err := tempFile.Chmod(0o640); err != nil {
|
||||
cleanup()
|
||||
return false, localBinaryStorageError(err)
|
||||
}
|
||||
if _, err := tempFile.Write(payload); err != nil {
|
||||
cleanup()
|
||||
return false, localBinaryStorageError(err)
|
||||
}
|
||||
if err := tempFile.Sync(); err != nil {
|
||||
cleanup()
|
||||
return false, localBinaryStorageError(err)
|
||||
}
|
||||
if err := tempFile.Close(); err != nil {
|
||||
_ = os.Remove(tempPath)
|
||||
return false, localBinaryStorageError(err)
|
||||
}
|
||||
if err := os.Rename(tempPath, targetPath); err != nil {
|
||||
_ = os.Remove(tempPath)
|
||||
return false, localBinaryStorageError(err)
|
||||
}
|
||||
if err := verifyLocalBinaryFile(targetPath, digestHex, int64(len(payload))); err != nil {
|
||||
_ = os.Remove(targetPath)
|
||||
return false, err
|
||||
}
|
||||
return true, nil
|
||||
}
|
||||
|
||||
func ensureLocalBinaryDiskHeadroom(path string, incomingBytes int64, minFreeBytes int64) error {
|
||||
var stat syscall.Statfs_t
|
||||
if err := syscall.Statfs(path, &stat); err != nil {
|
||||
return localBinaryStorageError(err)
|
||||
}
|
||||
freeBytes := int64(stat.Bavail) * int64(stat.Bsize)
|
||||
if freeBytes-incomingBytes < minFreeBytes {
|
||||
return &clients.ClientError{
|
||||
Code: "local_result_storage_unavailable",
|
||||
Message: "local result storage does not have enough free space",
|
||||
StatusCode: 503,
|
||||
Retryable: false,
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func localBinaryStorageError(err error) error {
|
||||
return &clients.ClientError{
|
||||
Code: "local_result_storage_unavailable",
|
||||
Message: "local result storage failed: " + err.Error(),
|
||||
StatusCode: 503,
|
||||
Retryable: false,
|
||||
}
|
||||
}
|
||||
|
||||
func verifyLocalBinaryFile(path string, expectedHash string, expectedSize int64) error {
|
||||
file, err := os.Open(path)
|
||||
if err != nil {
|
||||
return localBinaryStorageError(err)
|
||||
}
|
||||
defer file.Close()
|
||||
hasher := sha256.New()
|
||||
size, err := io.Copy(hasher, file)
|
||||
if err != nil {
|
||||
return localBinaryStorageError(err)
|
||||
}
|
||||
if size != expectedSize || hex.EncodeToString(hasher.Sum(nil)) != expectedHash {
|
||||
return &clients.ClientError{
|
||||
Code: "binary_result_corrupted",
|
||||
Message: "local result file failed size or hash verification",
|
||||
StatusCode: 500,
|
||||
Retryable: false,
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// HydrateTaskResult restores placeholders from verified local files. It is only
|
||||
// used by result/detail/replay endpoints, never by task lists or callbacks.
|
||||
func (s *Service) HydrateTaskResult(ctx context.Context, taskID string, result map[string]any) (map[string]any, error) {
|
||||
next, changed, err := s.hydrateLocalBinaryValue(ctx, taskID, result, 0)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if !changed {
|
||||
return result, nil
|
||||
}
|
||||
mapped, ok := next.(map[string]any)
|
||||
if !ok {
|
||||
return nil, &clients.ClientError{Code: "binary_result_corrupted", Message: "stored result is not a JSON object", StatusCode: 500}
|
||||
}
|
||||
return mapped, nil
|
||||
}
|
||||
|
||||
func (s *Service) hydrateLocalBinaryValue(ctx context.Context, taskID string, value any, depth int) (any, bool, error) {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
if depth >= localBinaryMaxDepth {
|
||||
return nil, false, &clients.ClientError{Code: "binary_result_corrupted", Message: "stored result exceeds the maximum JSON depth", StatusCode: 500}
|
||||
}
|
||||
switch typed := value.(type) {
|
||||
case map[string]any:
|
||||
next := make(map[string]any, len(typed))
|
||||
changed := false
|
||||
for key, childValue := range typed {
|
||||
child, childChanged, err := s.hydrateLocalBinaryValue(ctx, taskID, childValue, depth+1)
|
||||
if err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
next[key] = child
|
||||
changed = changed || childChanged
|
||||
}
|
||||
if !changed {
|
||||
return value, false, nil
|
||||
}
|
||||
return next, true, nil
|
||||
case []any:
|
||||
next := make([]any, len(typed))
|
||||
changed := false
|
||||
for index, childValue := range typed {
|
||||
child, childChanged, err := s.hydrateLocalBinaryValue(ctx, taskID, childValue, depth+1)
|
||||
if err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
next[index] = child
|
||||
changed = changed || childChanged
|
||||
}
|
||||
if !changed {
|
||||
return value, false, nil
|
||||
}
|
||||
return next, true, nil
|
||||
case string:
|
||||
descriptor, ok := parseLocalBinaryPlaceholder(typed)
|
||||
if !ok {
|
||||
return value, false, nil
|
||||
}
|
||||
payload, err := s.readLocalBinaryResult(taskID, descriptor)
|
||||
if err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
encoded := base64.StdEncoding.EncodeToString(payload)
|
||||
if descriptor.Encoding == "data-uri" {
|
||||
return "data:" + descriptor.ContentType + ";base64," + encoded, true, nil
|
||||
}
|
||||
return encoded, true, nil
|
||||
default:
|
||||
return value, false, nil
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Service) readLocalBinaryResult(taskID string, descriptor localBinaryDescriptor) ([]byte, error) {
|
||||
path := filepath.Join(s.localBinaryResultRoot(), safeLocalBinaryTaskDir(taskID), descriptor.SHA256+".bin")
|
||||
info, err := os.Stat(path)
|
||||
if err != nil {
|
||||
if errors.Is(err, os.ErrNotExist) {
|
||||
return nil, &clients.ClientError{Code: "binary_result_expired", Message: "local binary result has expired", StatusCode: 410, Retryable: false}
|
||||
}
|
||||
return nil, localBinaryStorageError(err)
|
||||
}
|
||||
if info.IsDir() || info.Size() != descriptor.Size {
|
||||
return nil, &clients.ClientError{Code: "binary_result_corrupted", Message: "local binary result has an invalid size", StatusCode: 500, Retryable: false}
|
||||
}
|
||||
if info.ModTime().Before(time.Now().Add(-time.Duration(s.localResultTTLHours()) * time.Hour)) {
|
||||
return nil, &clients.ClientError{Code: "binary_result_expired", Message: "local binary result has expired", StatusCode: 410, Retryable: false}
|
||||
}
|
||||
payload, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
return nil, localBinaryStorageError(err)
|
||||
}
|
||||
digest := sha256.Sum256(payload)
|
||||
if int64(len(payload)) != descriptor.Size || hex.EncodeToString(digest[:]) != descriptor.SHA256 {
|
||||
return nil, &clients.ClientError{Code: "binary_result_corrupted", Message: "local binary result failed size or hash verification", StatusCode: 500, Retryable: false}
|
||||
}
|
||||
return payload, nil
|
||||
}
|
||||
|
||||
func localBinaryStringBytes(key string, value string, siblings map[string]any) ([]byte, string, string, bool) {
|
||||
raw := strings.TrimSpace(value)
|
||||
if raw == "" || strings.HasPrefix(raw, localBinaryPlaceholderPrefix) {
|
||||
return nil, "", "", false
|
||||
}
|
||||
if strings.HasPrefix(strings.ToLower(raw), "data:") {
|
||||
contentType, encoded, ok, err := parseBase64DataURL(raw)
|
||||
if err == nil && ok {
|
||||
payload, decodeErr := decodeBase64Payload(encoded)
|
||||
if decodeErr == nil && len(payload) > 0 {
|
||||
return payload, contentType, "data-uri", true
|
||||
}
|
||||
}
|
||||
return nil, "", "", false
|
||||
}
|
||||
strict := localBinaryKey(key)
|
||||
if !strict && len(raw) < localBinaryGenericBase64MinLength {
|
||||
return nil, "", "", false
|
||||
}
|
||||
payload, err := decodeBase64Payload(raw)
|
||||
if err != nil || len(payload) == 0 {
|
||||
return nil, "", "", false
|
||||
}
|
||||
return payload, firstNonEmptyString(mediaContentTypeFromItem(siblings), defaultContentTypeForRawMediaKey(key)), "raw", true
|
||||
}
|
||||
|
||||
func localBufferObjectBytes(value map[string]any) ([]byte, string, bool) {
|
||||
if normalizeLocalBinaryKey(stringFromAny(value["type"])) != "buffer" {
|
||||
return nil, "", false
|
||||
}
|
||||
contentType := firstNonEmptyString(
|
||||
stringFromAny(value["mime_type"]),
|
||||
stringFromAny(value["mimeType"]),
|
||||
stringFromAny(value["contentType"]),
|
||||
)
|
||||
switch data := value["data"].(type) {
|
||||
case []byte:
|
||||
if len(data) == 0 {
|
||||
return nil, "", false
|
||||
}
|
||||
return append([]byte(nil), data...), contentType, true
|
||||
case []any:
|
||||
payload, ok := bytesFromNumberArray(data)
|
||||
return payload, contentType, ok
|
||||
default:
|
||||
return nil, "", false
|
||||
}
|
||||
}
|
||||
|
||||
func localBinaryKey(key string) bool {
|
||||
normalized := normalizeLocalBinaryKey(key)
|
||||
return normalized == "b64" ||
|
||||
normalized == "b64json" ||
|
||||
normalized == "base64" ||
|
||||
normalized == "buffer" ||
|
||||
normalized == "bytes" ||
|
||||
strings.Contains(normalized, "base64") ||
|
||||
strings.Contains(normalized, "buffer") ||
|
||||
strings.Contains(normalized, "binary") ||
|
||||
strings.HasSuffix(normalized, "b64") ||
|
||||
strings.HasSuffix(normalized, "bytes")
|
||||
}
|
||||
|
||||
func normalizeLocalBinaryKey(value string) string {
|
||||
return strings.Map(func(char rune) rune {
|
||||
switch {
|
||||
case char >= 'a' && char <= 'z':
|
||||
return char
|
||||
case char >= 'A' && char <= 'Z':
|
||||
return char + ('a' - 'A')
|
||||
case char >= '0' && char <= '9':
|
||||
return char
|
||||
default:
|
||||
return -1
|
||||
}
|
||||
}, value)
|
||||
}
|
||||
|
||||
func localBinaryPrefix(payload []byte) string {
|
||||
prefix := base64.StdEncoding.EncodeToString(payload)
|
||||
if len(prefix) > 16 {
|
||||
prefix = prefix[:16]
|
||||
}
|
||||
return prefix
|
||||
}
|
||||
|
||||
func normalizedLocalBinaryContentType(value string) string {
|
||||
value = normalizeGeneratedContentType(value)
|
||||
if value == "" || len(value) > 32 || !localBinaryContentTypeSafe(value) {
|
||||
return "application/octet-stream"
|
||||
}
|
||||
return value
|
||||
}
|
||||
|
||||
func localBinaryContentTypeSafe(value string) bool {
|
||||
for _, char := range value {
|
||||
switch {
|
||||
case char >= 'a' && char <= 'z':
|
||||
case char >= 'A' && char <= 'Z':
|
||||
case char >= '0' && char <= '9':
|
||||
case char == '/', char == '.', char == '+', char == '-':
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func normalizedLocalBinaryEncoding(value string) string {
|
||||
switch strings.ToLower(strings.TrimSpace(value)) {
|
||||
case "data-uri":
|
||||
return "data-uri"
|
||||
case "buffer":
|
||||
return "buffer"
|
||||
default:
|
||||
return "raw"
|
||||
}
|
||||
}
|
||||
|
||||
func localBinaryPlaceholder(descriptor localBinaryDescriptor) string {
|
||||
return fmt.Sprintf(
|
||||
"[GatewayBinary:v1;prefix=%s;sha256=%s;bytes=%d;mime=%s;encoding=%s]",
|
||||
descriptor.Prefix,
|
||||
descriptor.SHA256,
|
||||
descriptor.Size,
|
||||
descriptor.ContentType,
|
||||
descriptor.Encoding,
|
||||
)
|
||||
}
|
||||
|
||||
func parseLocalBinaryPlaceholder(value string) (localBinaryDescriptor, bool) {
|
||||
if !strings.HasPrefix(value, localBinaryPlaceholderPrefix) || !strings.HasSuffix(value, "]") {
|
||||
return localBinaryDescriptor{}, false
|
||||
}
|
||||
content := strings.TrimSuffix(strings.TrimPrefix(value, localBinaryPlaceholderPrefix), "]")
|
||||
fields := map[string]string{}
|
||||
for _, item := range strings.Split(content, ";") {
|
||||
key, fieldValue, ok := strings.Cut(item, "=")
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
fields[key] = fieldValue
|
||||
}
|
||||
size, err := strconv.ParseInt(fields["bytes"], 10, 64)
|
||||
if err != nil || size <= 0 || len(fields["sha256"]) != sha256.Size*2 {
|
||||
return localBinaryDescriptor{}, false
|
||||
}
|
||||
if _, err := hex.DecodeString(fields["sha256"]); err != nil {
|
||||
return localBinaryDescriptor{}, false
|
||||
}
|
||||
descriptor := localBinaryDescriptor{
|
||||
Prefix: fields["prefix"],
|
||||
SHA256: strings.ToLower(fields["sha256"]),
|
||||
Size: size,
|
||||
ContentType: normalizedLocalBinaryContentType(fields["mime"]),
|
||||
Encoding: normalizedLocalBinaryEncoding(fields["encoding"]),
|
||||
}
|
||||
if len(descriptor.Prefix) > 16 {
|
||||
return localBinaryDescriptor{}, false
|
||||
}
|
||||
return descriptor, true
|
||||
}
|
||||
|
||||
func localBinaryResultHasPlaceholders(value any) bool {
|
||||
switch typed := value.(type) {
|
||||
case map[string]any:
|
||||
for _, child := range typed {
|
||||
if localBinaryResultHasPlaceholders(child) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
case []any:
|
||||
for _, child := range typed {
|
||||
if localBinaryResultHasPlaceholders(child) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
case string:
|
||||
return strings.HasPrefix(typed, localBinaryPlaceholderPrefix)
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func safeLocalBinaryTaskDir(taskID string) string {
|
||||
taskID = strings.TrimSpace(taskID)
|
||||
if taskID != "" {
|
||||
safe := true
|
||||
for _, char := range taskID {
|
||||
if (char >= 'a' && char <= 'z') ||
|
||||
(char >= 'A' && char <= 'Z') ||
|
||||
(char >= '0' && char <= '9') ||
|
||||
char == '-' || char == '_' {
|
||||
continue
|
||||
}
|
||||
safe = false
|
||||
break
|
||||
}
|
||||
if safe && taskID != "." && taskID != ".." && len(taskID) <= 128 {
|
||||
return taskID
|
||||
}
|
||||
}
|
||||
digest := sha256.Sum256([]byte(taskID))
|
||||
return "task-" + hex.EncodeToString(digest[:16])
|
||||
}
|
||||
|
||||
func (s *Service) localBinaryResultRoot() string {
|
||||
root := strings.TrimSpace(s.cfg.LocalGeneratedStorageDir)
|
||||
if root == "" {
|
||||
root = config.DefaultLocalGeneratedStorageDir
|
||||
}
|
||||
return filepath.Join(root, localBinaryResultDirName)
|
||||
}
|
||||
|
||||
func (s *Service) localResultTTLHours() int {
|
||||
if s.cfg.LocalResultTTLHours <= 0 {
|
||||
return defaultLocalResultTTLHours
|
||||
}
|
||||
return s.cfg.LocalResultTTLHours
|
||||
}
|
||||
|
||||
func (s *Service) localResultMinFreeBytes() int64 {
|
||||
if s.cfg.LocalResultMinFreeBytes <= 0 {
|
||||
return defaultLocalResultMinFreeBytes
|
||||
}
|
||||
return s.cfg.LocalResultMinFreeBytes
|
||||
}
|
||||
|
||||
func (s *Service) localResultMaxBytes() int64 {
|
||||
if s.cfg.LocalResultMaxBytes <= 0 {
|
||||
return defaultLocalResultMaxBytes
|
||||
}
|
||||
return s.cfg.LocalResultMaxBytes
|
||||
}
|
||||
|
||||
func (s *Service) localResultMaxTaskBytes() int64 {
|
||||
if s.cfg.LocalResultMaxTaskBytes <= 0 {
|
||||
return defaultLocalResultMaxTaskBytes
|
||||
}
|
||||
return s.cfg.LocalResultMaxTaskBytes
|
||||
}
|
||||
@@ -0,0 +1,233 @@
|
||||
package runner
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/clients"
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
)
|
||||
|
||||
func newLocalBinaryTestService(t *testing.T) *Service {
|
||||
t.Helper()
|
||||
return &Service{cfg: config.Config{
|
||||
LocalGeneratedStorageDir: t.TempDir(),
|
||||
LocalResultTTLHours: 24,
|
||||
LocalResultMinFreeBytes: 1,
|
||||
LocalResultMaxBytes: 1024 * 1024,
|
||||
LocalResultMaxTaskBytes: 2 * 1024 * 1024,
|
||||
}}
|
||||
}
|
||||
|
||||
func TestMaterializeAndHydrateLocalBinaryResult(t *testing.T) {
|
||||
service := newLocalBinaryTestService(t)
|
||||
payload := []byte("one binary result shared across representations")
|
||||
encoded := base64.StdEncoding.EncodeToString(payload)
|
||||
input := map[string]any{
|
||||
"data": []any{
|
||||
map[string]any{
|
||||
"b64_json": encoded,
|
||||
"data_uri": "data:image/png;base64," + encoded,
|
||||
"buffer": map[string]any{
|
||||
"type": "Buffer",
|
||||
"data": bytesToAny(payload),
|
||||
"mimeType": "image/png",
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
persistent, changed, err := service.materializeLocalBinaryResult(context.Background(), "task-123", input)
|
||||
if err != nil {
|
||||
t.Fatalf("materialize local binary result: %v", err)
|
||||
}
|
||||
if !changed {
|
||||
t.Fatal("expected binary result to be materialized")
|
||||
}
|
||||
item := persistent["data"].([]any)[0].(map[string]any)
|
||||
for _, key := range []string{"b64_json", "data_uri", "buffer"} {
|
||||
placeholder, ok := item[key].(string)
|
||||
if !ok || !strings.HasPrefix(placeholder, localBinaryPlaceholderPrefix) {
|
||||
t.Fatalf("%s was not replaced with a placeholder: %+v", key, item[key])
|
||||
}
|
||||
if len(placeholder) > 200 {
|
||||
t.Fatalf("%s placeholder exceeds 200 bytes: %d", key, len(placeholder))
|
||||
}
|
||||
}
|
||||
if input["data"].([]any)[0].(map[string]any)["b64_json"] != encoded {
|
||||
t.Fatal("materializer mutated the provider result")
|
||||
}
|
||||
persistentJSON, err := json.Marshal(persistent)
|
||||
if err != nil {
|
||||
t.Fatalf("marshal persistent result: %v", err)
|
||||
}
|
||||
if len(persistentJSON) > 32*1024 {
|
||||
t.Fatalf("binary-only persistent result exceeds 32KiB: %d", len(persistentJSON))
|
||||
}
|
||||
|
||||
taskDir := filepath.Join(service.localBinaryResultRoot(), "task-123")
|
||||
entries, err := os.ReadDir(taskDir)
|
||||
if err != nil {
|
||||
t.Fatalf("read local result dir: %v", err)
|
||||
}
|
||||
if len(entries) != 1 {
|
||||
t.Fatalf("same payload should reuse one local file, got %d", len(entries))
|
||||
}
|
||||
taskInfo, err := os.Stat(taskDir)
|
||||
if err != nil {
|
||||
t.Fatalf("stat task directory: %v", err)
|
||||
}
|
||||
if taskInfo.Mode().Perm() != 0o750 {
|
||||
t.Fatalf("task directory mode = %v, want 0750", taskInfo.Mode().Perm())
|
||||
}
|
||||
fileInfo, err := entries[0].Info()
|
||||
if err != nil {
|
||||
t.Fatalf("stat result file: %v", err)
|
||||
}
|
||||
if fileInfo.Mode().Perm() != 0o640 {
|
||||
t.Fatalf("result file mode = %v, want 0640", fileInfo.Mode().Perm())
|
||||
}
|
||||
|
||||
wire, err := service.HydrateTaskResult(context.Background(), "task-123", persistent)
|
||||
if err != nil {
|
||||
t.Fatalf("hydrate local binary result: %v", err)
|
||||
}
|
||||
wireItem := wire["data"].([]any)[0].(map[string]any)
|
||||
if wireItem["b64_json"] != encoded {
|
||||
t.Fatalf("raw Base64 mismatch: got %v", wireItem["b64_json"])
|
||||
}
|
||||
if wireItem["data_uri"] != "data:image/png;base64,"+encoded {
|
||||
t.Fatalf("data URI mismatch: got %v", wireItem["data_uri"])
|
||||
}
|
||||
if wireItem["buffer"] != encoded {
|
||||
t.Fatalf("Buffer should hydrate as Base64: got %v", wireItem["buffer"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestHydrateLocalBinaryResultReturnsExpiredAndCorruptedErrors(t *testing.T) {
|
||||
service := newLocalBinaryTestService(t)
|
||||
service.cfg.LocalResultTTLHours = 1
|
||||
encoded := base64.StdEncoding.EncodeToString([]byte("expiring result"))
|
||||
persistent, _, err := service.materializeLocalBinaryResult(context.Background(), "task-expired", map[string]any{"b64_json": encoded})
|
||||
if err != nil {
|
||||
t.Fatalf("materialize fixture: %v", err)
|
||||
}
|
||||
entries, err := os.ReadDir(filepath.Join(service.localBinaryResultRoot(), "task-expired"))
|
||||
if err != nil || len(entries) != 1 {
|
||||
t.Fatalf("read fixture result: entries=%v err=%v", entries, err)
|
||||
}
|
||||
path := filepath.Join(service.localBinaryResultRoot(), "task-expired", entries[0].Name())
|
||||
old := time.Now().Add(-2 * time.Hour)
|
||||
if err := os.Chtimes(path, old, old); err != nil {
|
||||
t.Fatalf("age fixture: %v", err)
|
||||
}
|
||||
_, err = service.HydrateTaskResult(context.Background(), "task-expired", persistent)
|
||||
assertClientErrorCode(t, err, "binary_result_expired")
|
||||
|
||||
now := time.Now()
|
||||
if err := os.Chtimes(path, now, now); err != nil {
|
||||
t.Fatalf("refresh fixture: %v", err)
|
||||
}
|
||||
if err := os.WriteFile(path, []byte("tampered"), 0o640); err != nil {
|
||||
t.Fatalf("tamper fixture: %v", err)
|
||||
}
|
||||
_, err = service.HydrateTaskResult(context.Background(), "task-expired", persistent)
|
||||
assertClientErrorCode(t, err, "binary_result_corrupted")
|
||||
}
|
||||
|
||||
func TestMaterializeLocalBinaryResultEnforcesLimitsAndKeepsText(t *testing.T) {
|
||||
service := newLocalBinaryTestService(t)
|
||||
service.cfg.LocalResultMaxBytes = 4
|
||||
_, _, err := service.materializeLocalBinaryResult(context.Background(), "task-large", map[string]any{
|
||||
"b64_json": base64.StdEncoding.EncodeToString([]byte("too large")),
|
||||
})
|
||||
assertClientErrorCode(t, err, "binary_result_too_large")
|
||||
|
||||
persistent, changed, err := service.materializeLocalBinaryResult(context.Background(), "task-text", map[string]any{
|
||||
"message": "YWJjZA==",
|
||||
})
|
||||
if err != nil || changed || persistent["message"] != "YWJjZA==" {
|
||||
t.Fatalf("short Base64-like text should remain unchanged: result=%+v changed=%v err=%v", persistent, changed, err)
|
||||
}
|
||||
|
||||
service = newLocalBinaryTestService(t)
|
||||
service.cfg.LocalResultMaxBytes = 1024
|
||||
service.cfg.LocalResultMaxTaskBytes = 8
|
||||
_, _, err = service.materializeLocalBinaryResult(context.Background(), "task-total-large", map[string]any{
|
||||
"first_base64": base64.StdEncoding.EncodeToString([]byte("12345")),
|
||||
"second_base64": base64.StdEncoding.EncodeToString([]byte("67890")),
|
||||
})
|
||||
assertClientErrorCode(t, err, "binary_result_too_large")
|
||||
if _, statErr := os.Stat(filepath.Join(service.localBinaryResultRoot(), "task-total-large")); !errors.Is(statErr, os.ErrNotExist) {
|
||||
t.Fatalf("preflight task limit must not create partial files: %v", statErr)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCompactExpiredTaskResultDoesNotWriteFiles(t *testing.T) {
|
||||
service := newLocalBinaryTestService(t)
|
||||
service.cfg.LocalResultMaxBytes = 1
|
||||
encoded := base64.StdEncoding.EncodeToString([]byte("already expired result"))
|
||||
|
||||
persistent, changed, err := service.CompactExpiredTaskResultForStorage(context.Background(), "task-old", map[string]any{
|
||||
"b64_json": encoded,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("compact expired result: %v", err)
|
||||
}
|
||||
if !changed || !localBinaryResultHasPlaceholders(persistent) {
|
||||
t.Fatalf("expired result was not compacted: %+v", persistent)
|
||||
}
|
||||
if _, err := os.Stat(filepath.Join(service.localBinaryResultRoot(), "task-old")); !errors.Is(err, os.ErrNotExist) {
|
||||
t.Fatalf("expired compaction must not create a task directory: %v", err)
|
||||
}
|
||||
_, err = service.HydrateTaskResult(context.Background(), "task-old", persistent)
|
||||
assertClientErrorCode(t, err, "binary_result_expired")
|
||||
}
|
||||
|
||||
func TestLocalBinaryStorageUnavailableDoesNotRetryProvider(t *testing.T) {
|
||||
err := localBinaryStorageError(errors.New("write failed"))
|
||||
assertClientErrorCode(t, err, "local_result_storage_unavailable")
|
||||
if clients.IsRetryable(err) {
|
||||
t.Fatal("local result storage failure must not call the provider again")
|
||||
}
|
||||
if retryDecisionForCandidate(store.RuntimeModelCandidate{}, err).Retry {
|
||||
t.Fatal("local result storage failure must not retry the same provider client")
|
||||
}
|
||||
if failoverDecisionForCandidate(store.RunnerPolicy{}, store.RuntimeModelCandidate{}, err).Retry {
|
||||
t.Fatal("local result storage failure must not fail over to another provider")
|
||||
}
|
||||
|
||||
service := newLocalBinaryTestService(t)
|
||||
service.cfg.LocalResultMinFreeBytes = 1 << 62
|
||||
_, _, err = service.materializeLocalBinaryResult(context.Background(), "task-low-disk", map[string]any{
|
||||
"b64_json": base64.StdEncoding.EncodeToString([]byte("disk preflight")),
|
||||
})
|
||||
assertClientErrorCode(t, err, "local_result_storage_unavailable")
|
||||
if _, statErr := os.Stat(filepath.Join(service.localBinaryResultRoot(), "task-low-disk")); !errors.Is(statErr, os.ErrNotExist) {
|
||||
t.Fatalf("disk preflight must not create partial task files: %v", statErr)
|
||||
}
|
||||
}
|
||||
|
||||
func bytesToAny(payload []byte) []any {
|
||||
result := make([]any, len(payload))
|
||||
for index, value := range payload {
|
||||
result[index] = float64(value)
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
func assertClientErrorCode(t *testing.T, err error, code string) {
|
||||
t.Helper()
|
||||
var clientErr *clients.ClientError
|
||||
if !errors.As(err, &clientErr) || clientErr.Code != code {
|
||||
t.Fatalf("error = %v, want client error code %s", err, code)
|
||||
}
|
||||
}
|
||||
@@ -13,6 +13,8 @@ type cacheAffinityKeys struct {
|
||||
Record []string
|
||||
}
|
||||
|
||||
const cacheAffinityMinimumStableMessages = 2
|
||||
|
||||
func buildCacheAffinityKey(kind string, modelType string, body map[string]any) string {
|
||||
return buildCacheAffinityKeys(kind, modelType, body).Primary
|
||||
}
|
||||
@@ -37,48 +39,28 @@ func buildCacheAffinityKeys(kind string, modelType string, body map[string]any)
|
||||
return cacheAffinityKeys{}
|
||||
}
|
||||
basePayload := map[string]any{
|
||||
"kind": kind,
|
||||
"modelType": modelType,
|
||||
"version": "v2",
|
||||
}
|
||||
if tools, ok := body["tools"]; ok {
|
||||
basePayload["tools"] = tools
|
||||
}
|
||||
if instructions, ok := body["instructions"]; ok {
|
||||
basePayload["instructions"] = instructions
|
||||
}
|
||||
messageHashes := cacheAffinityMessageHashes(body)
|
||||
if len(messageHashes) > 0 {
|
||||
if len(messageHashes) >= cacheAffinityMinimumStableMessages {
|
||||
baseHash := stableJSONHash(basePayload)
|
||||
keys := make([]string, 0, len(messageHashes))
|
||||
for index := range messageHashes {
|
||||
keyPayload := map[string]any{
|
||||
"base": baseHash,
|
||||
"messageHashes": messageHashes[:index+1],
|
||||
"prefixLength": index + 1,
|
||||
}
|
||||
keys = append(keys, "prompt_lcp:"+stableJSONHash(keyPayload))
|
||||
}
|
||||
keys := cacheAffinityPrefixKeys("prompt_lcp_v2:", baseHash, messageHashes)
|
||||
if len(keys) == 0 {
|
||||
return cacheAffinityKeys{}
|
||||
}
|
||||
lookup := make([]string, 0, len(keys))
|
||||
for index := len(keys) - 1; index >= 0; index-- {
|
||||
lookup = append(lookup, keys[index])
|
||||
}
|
||||
lookup := reverseCacheAffinityKeys(keys)
|
||||
legacyKeys := legacyCacheAffinityPrefixKeys(kind, modelType, body)
|
||||
lookup = normalizedCacheAffinityKeyList(append(lookup, reverseCacheAffinityKeys(legacyKeys)...))
|
||||
return cacheAffinityKeys{
|
||||
Primary: keys[len(keys)-1],
|
||||
Lookup: lookup,
|
||||
Record: keys,
|
||||
}
|
||||
}
|
||||
if input, ok := body["input"]; ok {
|
||||
basePayload["input"] = input
|
||||
}
|
||||
if len(basePayload) <= 2 {
|
||||
return cacheAffinityKeys{}
|
||||
}
|
||||
key := "prompt:" + stableJSONHash(basePayload)
|
||||
return cacheAffinityKeys{Primary: key, Lookup: []string{key}, Record: []string{key}}
|
||||
return cacheAffinityKeys{}
|
||||
}
|
||||
|
||||
func cacheAffinityPromptHashSupported(kind string, modelType string) bool {
|
||||
@@ -95,6 +77,146 @@ func cacheAffinityPromptHashSupported(kind string, modelType string) bool {
|
||||
}
|
||||
|
||||
func cacheAffinityMessageHashes(body map[string]any) []string {
|
||||
if hashes := stringListFromAny(body["messageHashes"]); len(hashes) > 0 {
|
||||
return hashes
|
||||
}
|
||||
if hashes := stringListFromAny(body["message_hashes"]); len(hashes) > 0 {
|
||||
return hashes
|
||||
}
|
||||
messages := cacheAffinityStableMessages(body)
|
||||
if len(messages) == 0 {
|
||||
return nil
|
||||
}
|
||||
hashes := make([]string, 0, len(messages))
|
||||
for _, message := range messages {
|
||||
hash := stableJSONHash(message)
|
||||
if hash == "" {
|
||||
continue
|
||||
}
|
||||
hashes = append(hashes, hash)
|
||||
}
|
||||
return hashes
|
||||
}
|
||||
|
||||
func cacheAffinityStableMessages(body map[string]any) []any {
|
||||
messages := make([]any, 0)
|
||||
if instructions := strings.TrimSpace(stringFromAny(body["instructions"])); instructions != "" {
|
||||
messages = append(messages, map[string]any{"role": "system", "content": instructions})
|
||||
}
|
||||
if rawMessages, ok := body["messages"].([]any); ok {
|
||||
for _, message := range rawMessages {
|
||||
if normalized := normalizeCacheAffinityMessage(message); normalized != nil {
|
||||
messages = append(messages, normalized)
|
||||
}
|
||||
}
|
||||
return messages
|
||||
}
|
||||
switch input := body["input"].(type) {
|
||||
case string:
|
||||
if strings.TrimSpace(input) != "" {
|
||||
messages = append(messages, map[string]any{"role": "user", "content": input})
|
||||
}
|
||||
case []any:
|
||||
for _, message := range input {
|
||||
if normalized := normalizeCacheAffinityMessage(message); normalized != nil {
|
||||
messages = append(messages, normalized)
|
||||
}
|
||||
}
|
||||
}
|
||||
return messages
|
||||
}
|
||||
|
||||
func normalizeCacheAffinityMessage(value any) any {
|
||||
message, ok := value.(map[string]any)
|
||||
if !ok {
|
||||
if text := strings.TrimSpace(stringFromAny(value)); text != "" {
|
||||
return map[string]any{"role": "user", "content": text}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
messageType := strings.TrimSpace(stringFromAny(message["type"]))
|
||||
switch messageType {
|
||||
case "function_call":
|
||||
return map[string]any{
|
||||
"role": "assistant",
|
||||
"tool_calls": []any{map[string]any{
|
||||
"id": firstNonEmptyCacheAffinityString(message["call_id"], message["id"]),
|
||||
"type": "function",
|
||||
"function": map[string]any{
|
||||
"name": message["name"],
|
||||
"arguments": message["arguments"],
|
||||
},
|
||||
}},
|
||||
}
|
||||
case "function_call_output":
|
||||
return map[string]any{
|
||||
"role": "tool",
|
||||
"tool_call_id": firstNonEmptyCacheAffinityString(message["call_id"], message["id"]),
|
||||
"content": message["output"],
|
||||
}
|
||||
}
|
||||
normalized := make(map[string]any, len(message))
|
||||
for key, item := range message {
|
||||
switch key {
|
||||
case "id", "status":
|
||||
continue
|
||||
case "type":
|
||||
if messageType == "message" {
|
||||
continue
|
||||
}
|
||||
}
|
||||
normalized[key] = item
|
||||
}
|
||||
if role := strings.TrimSpace(stringFromAny(normalized["role"])); role != "" {
|
||||
normalized["role"] = role
|
||||
}
|
||||
return normalized
|
||||
}
|
||||
|
||||
func firstNonEmptyCacheAffinityString(values ...any) string {
|
||||
for _, value := range values {
|
||||
if text := strings.TrimSpace(stringFromAny(value)); text != "" {
|
||||
return text
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func cacheAffinityPrefixKeys(prefix string, baseHash string, messageHashes []string) []string {
|
||||
if len(messageHashes) < cacheAffinityMinimumStableMessages {
|
||||
return nil
|
||||
}
|
||||
keys := make([]string, 0, len(messageHashes)-cacheAffinityMinimumStableMessages+1)
|
||||
for prefixLength := cacheAffinityMinimumStableMessages; prefixLength <= len(messageHashes); prefixLength++ {
|
||||
keyPayload := map[string]any{
|
||||
"base": baseHash,
|
||||
"messageHashes": messageHashes[:prefixLength],
|
||||
"prefixLength": prefixLength,
|
||||
}
|
||||
keys = append(keys, prefix+stableJSONHash(keyPayload))
|
||||
}
|
||||
return keys
|
||||
}
|
||||
|
||||
func legacyCacheAffinityPrefixKeys(kind string, modelType string, body map[string]any) []string {
|
||||
messageHashes := legacyCacheAffinityMessageHashes(body)
|
||||
if len(messageHashes) < cacheAffinityMinimumStableMessages {
|
||||
return nil
|
||||
}
|
||||
basePayload := map[string]any{
|
||||
"kind": kind,
|
||||
"modelType": modelType,
|
||||
}
|
||||
if tools, ok := body["tools"]; ok {
|
||||
basePayload["tools"] = tools
|
||||
}
|
||||
if instructions, ok := body["instructions"]; ok {
|
||||
basePayload["instructions"] = instructions
|
||||
}
|
||||
return cacheAffinityPrefixKeys("prompt_lcp:", stableJSONHash(basePayload), messageHashes)
|
||||
}
|
||||
|
||||
func legacyCacheAffinityMessageHashes(body map[string]any) []string {
|
||||
if hashes := stringListFromAny(body["messageHashes"]); len(hashes) > 0 {
|
||||
return hashes
|
||||
}
|
||||
@@ -102,20 +224,41 @@ func cacheAffinityMessageHashes(body map[string]any) []string {
|
||||
return hashes
|
||||
}
|
||||
messages, ok := body["messages"].([]any)
|
||||
if !ok || len(messages) == 0 {
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
hashes := make([]string, 0, len(messages))
|
||||
for _, message := range messages {
|
||||
raw, err := json.Marshal(message)
|
||||
if err != nil || len(raw) == 0 {
|
||||
continue
|
||||
if err == nil && len(raw) > 0 {
|
||||
hashes = append(hashes, sha256Text(string(raw)))
|
||||
}
|
||||
hashes = append(hashes, sha256Text(string(raw)))
|
||||
}
|
||||
return hashes
|
||||
}
|
||||
|
||||
func reverseCacheAffinityKeys(keys []string) []string {
|
||||
reversed := make([]string, 0, len(keys))
|
||||
for index := len(keys) - 1; index >= 0; index-- {
|
||||
reversed = append(reversed, keys[index])
|
||||
}
|
||||
return reversed
|
||||
}
|
||||
|
||||
func normalizedCacheAffinityKeyList(keys []string) []string {
|
||||
seen := map[string]bool{}
|
||||
out := make([]string, 0, len(keys))
|
||||
for _, key := range keys {
|
||||
key = strings.TrimSpace(key)
|
||||
if key == "" || seen[key] {
|
||||
continue
|
||||
}
|
||||
seen[key] = true
|
||||
out = append(out, key)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func stableJSONHash(value any) string {
|
||||
raw, err := json.Marshal(value)
|
||||
if err != nil || len(raw) == 0 {
|
||||
|
||||
@@ -39,7 +39,7 @@ func TestBuildCacheAffinityKeyFallsBackToPromptHash(t *testing.T) {
|
||||
|
||||
first := buildCacheAffinityKey("chat.completions", "text_generate", body)
|
||||
second := buildCacheAffinityKey("chat.completions", "text_generate", body)
|
||||
if first == "" || first != second || !strings.HasPrefix(first, "prompt_lcp:") {
|
||||
if first == "" || first != second || !strings.HasPrefix(first, "prompt_lcp_v2:") {
|
||||
t.Fatalf("expected stable prompt LCP cache affinity key, first=%q second=%q", first, second)
|
||||
}
|
||||
}
|
||||
@@ -54,10 +54,10 @@ func TestBuildCacheAffinityKeysUsesMessagePrefixLCPKeys(t *testing.T) {
|
||||
"tools": []any{map[string]any{"type": "function", "function": map[string]any{"name": "lookup"}}},
|
||||
})
|
||||
|
||||
if !strings.HasPrefix(first.Primary, "prompt_lcp:") {
|
||||
if !strings.HasPrefix(first.Primary, "prompt_lcp_v2:") {
|
||||
t.Fatalf("expected LCP prompt key, got %+v", first)
|
||||
}
|
||||
if len(first.Record) != 2 || len(next.Lookup) != 4 {
|
||||
if len(first.Record) != 1 || len(next.Record) != 3 || len(next.Lookup) != 6 {
|
||||
t.Fatalf("expected per-message prefix keys, first=%+v next=%+v", first, next)
|
||||
}
|
||||
found := false
|
||||
@@ -71,3 +71,81 @@ func TestBuildCacheAffinityKeysUsesMessagePrefixLCPKeys(t *testing.T) {
|
||||
t.Fatalf("next lookup keys should include first request prefix key, first=%+v next=%+v", first, next)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildCacheAffinityKeysRequiresTwoStableMessages(t *testing.T) {
|
||||
systemOnly := buildCacheAffinityKeys("chat.completions", "text_generate", map[string]any{
|
||||
"messages": []any{
|
||||
map[string]any{"role": "system", "content": "shared system prompt"},
|
||||
},
|
||||
})
|
||||
firstConversation := buildCacheAffinityKeys("chat.completions", "text_generate", map[string]any{
|
||||
"messages": []any{
|
||||
map[string]any{"role": "system", "content": "shared system prompt"},
|
||||
map[string]any{"role": "user", "content": "conversation A"},
|
||||
},
|
||||
})
|
||||
secondConversation := buildCacheAffinityKeys("chat.completions", "text_generate", map[string]any{
|
||||
"messages": []any{
|
||||
map[string]any{"role": "system", "content": "shared system prompt"},
|
||||
map[string]any{"role": "user", "content": "conversation B"},
|
||||
},
|
||||
})
|
||||
|
||||
if systemOnly.Primary != "" || len(systemOnly.Lookup) != 0 || len(systemOnly.Record) != 0 {
|
||||
t.Fatalf("shared system prompt alone must not create affinity, got %+v", systemOnly)
|
||||
}
|
||||
if firstConversation.Primary == "" || secondConversation.Primary == "" {
|
||||
t.Fatalf("two stable messages should create affinity, first=%+v second=%+v", firstConversation, secondConversation)
|
||||
}
|
||||
for _, firstKey := range firstConversation.Lookup {
|
||||
for _, secondKey := range secondConversation.Lookup {
|
||||
if firstKey == secondKey {
|
||||
t.Fatalf("different conversations sharing only a system prompt must not intersect: key=%s", firstKey)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildCacheAffinityKeysNormalizesResponsesToolHistory(t *testing.T) {
|
||||
tools := []any{map[string]any{
|
||||
"type": "function",
|
||||
"function": map[string]any{
|
||||
"name": "lookup_weather",
|
||||
},
|
||||
}}
|
||||
chat := buildCacheAffinityKeys("chat.completions", "text_generate", map[string]any{
|
||||
"tools": tools,
|
||||
"messages": []any{
|
||||
map[string]any{"role": "system", "content": "Use tools when helpful."},
|
||||
map[string]any{"role": "user", "content": "Weather in Hangzhou?"},
|
||||
},
|
||||
})
|
||||
responses := buildCacheAffinityKeys("responses", "text_generate", map[string]any{
|
||||
"tools": tools,
|
||||
"instructions": "Use tools when helpful.",
|
||||
"input": "Weather in Hangzhou?",
|
||||
})
|
||||
continued := buildCacheAffinityKeys("responses", "text_generate", map[string]any{
|
||||
"tools": tools,
|
||||
"instructions": "Use tools when helpful.",
|
||||
"input": []any{
|
||||
map[string]any{"role": "user", "content": "Weather in Hangzhou?"},
|
||||
map[string]any{"type": "function_call", "call_id": "call_weather", "name": "lookup_weather", "arguments": `{"city":"Hangzhou"}`},
|
||||
map[string]any{"type": "function_call_output", "call_id": "call_weather", "output": `{"temperature":31}`},
|
||||
},
|
||||
})
|
||||
|
||||
if chat.Primary == "" || chat.Primary != responses.Primary {
|
||||
t.Fatalf("chat and responses should share the V2 protocol-independent prefix, chat=%+v responses=%+v", chat, responses)
|
||||
}
|
||||
found := false
|
||||
for _, key := range continued.Lookup {
|
||||
if key == responses.Primary {
|
||||
found = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Fatalf("responses tool continuation should retain its stable prefix, initial=%+v continued=%+v", responses, continued)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -8,6 +8,7 @@ import (
|
||||
)
|
||||
|
||||
const unsupportedRequestResolutionCode = "unsupported_request_resolution"
|
||||
const unsupportedRequestAspectRatioCode = "unsupported_request_aspect_ratio"
|
||||
|
||||
type requestResolutionRequirement struct {
|
||||
Kind string
|
||||
@@ -29,6 +30,15 @@ type videoResolutionReferenceStats struct {
|
||||
}
|
||||
|
||||
func filterRuntimeCandidatesByRequest(kind string, requestedModel string, modelType string, body map[string]any, candidates []store.RuntimeModelCandidate) ([]store.RuntimeModelCandidate, map[string]any, error) {
|
||||
filtered, resolutionSummary, err := filterRuntimeCandidatesByRequestResolution(kind, requestedModel, modelType, body, candidates)
|
||||
if err != nil {
|
||||
return filtered, resolutionSummary, err
|
||||
}
|
||||
filtered, aspectRatioSummary, err := filterRuntimeCandidatesByRequestAspectRatio(kind, requestedModel, modelType, body, filtered)
|
||||
return filtered, mergeCandidateFilterSummaries(resolutionSummary, aspectRatioSummary), err
|
||||
}
|
||||
|
||||
func filterRuntimeCandidatesByRequestResolution(kind string, requestedModel string, modelType string, body map[string]any, candidates []store.RuntimeModelCandidate) ([]store.RuntimeModelCandidate, map[string]any, error) {
|
||||
requirement, ok := requestResolutionRequirementFor(kind, requestedModel, modelType, body)
|
||||
if !ok || len(candidates) == 0 {
|
||||
return candidates, nil, nil
|
||||
@@ -82,13 +92,13 @@ func requestResolutionRequirementFor(kind string, requestedModel string, modelTy
|
||||
}
|
||||
|
||||
func requestResolutionValue(body map[string]any, modelType string) (string, string) {
|
||||
if value := normalizedRequestResolution(stringFromAny(body["resolution"])); value != "" {
|
||||
if value := normalizedRequestResolution(stringFromAny(body["resolution"]), modelType); value != "" {
|
||||
if _, _, ok := parsePixelSizeString(value); ok {
|
||||
return "", ""
|
||||
}
|
||||
return value, "resolution"
|
||||
}
|
||||
size := normalizedRequestResolution(stringFromAny(body["size"]))
|
||||
size := normalizedRequestResolution(stringFromAny(body["size"]), modelType)
|
||||
if size == "" {
|
||||
return "", ""
|
||||
}
|
||||
@@ -98,7 +108,7 @@ func requestResolutionValue(body map[string]any, modelType string) (string, stri
|
||||
return "", ""
|
||||
}
|
||||
|
||||
func normalizedRequestResolution(value string) string {
|
||||
func normalizedRequestResolution(value string, modelType string) string {
|
||||
value = strings.TrimSpace(value)
|
||||
if value == "" || isEmptyParamString(value) {
|
||||
return ""
|
||||
@@ -106,9 +116,11 @@ func normalizedRequestResolution(value string) string {
|
||||
switch strings.ToLower(value) {
|
||||
case "auto", "automatic", "adaptive", "default":
|
||||
return ""
|
||||
default:
|
||||
}
|
||||
if _, _, ok := parsePixelSizeString(value); ok {
|
||||
return value
|
||||
}
|
||||
return normalizeMediaResolution(modelType, value)
|
||||
}
|
||||
|
||||
func isResolutionFilteredModelType(modelType string) bool {
|
||||
@@ -120,10 +132,9 @@ func candidateSupportsRequestResolution(candidate store.RuntimeModelCandidate, r
|
||||
capability := capabilityForType(effectiveModelCapability(candidate), modelType)
|
||||
detail := candidateResolutionDetail(candidate, requirement, modelType)
|
||||
if capability == nil {
|
||||
detail["reason"] = "capability_missing"
|
||||
detail["message"] = "候选平台模型未配置对应模型类型能力。"
|
||||
detail["capabilityPath"] = capabilityPath(modelType, "output_resolutions")
|
||||
return false, detail
|
||||
detail["reason"] = "capability_unrestricted"
|
||||
detail["message"] = "候选平台模型未配置分辨率白名单,按未限制处理。"
|
||||
return true, detail
|
||||
}
|
||||
|
||||
allowed, configured := outputResolutionAllowedValues(capability["output_resolutions"], requirement.Scopes)
|
||||
@@ -131,9 +142,19 @@ func candidateSupportsRequestResolution(candidate store.RuntimeModelCandidate, r
|
||||
detail["capabilityPath"] = capabilityPath(modelType, "output_resolutions")
|
||||
detail["capabilityValue"] = cloneAny(capability["output_resolutions"])
|
||||
if !configured {
|
||||
detail["reason"] = "output_resolutions_missing"
|
||||
detail["message"] = "候选平台模型未声明 output_resolutions。"
|
||||
return false, detail
|
||||
detail["reason"] = "output_resolutions_unrestricted"
|
||||
detail["message"] = "候选平台模型未声明 output_resolutions,按未限制处理。"
|
||||
return true, detail
|
||||
}
|
||||
if len(allowed) == 0 {
|
||||
if _, scoped := capability["output_resolutions"].(map[string]any); scoped && len(requirement.Scopes) > 0 && !outputResolutionScopeConfigured(capability["output_resolutions"], requirement.Scopes) {
|
||||
detail["reason"] = "resolution_scope_not_allowed"
|
||||
detail["message"] = "候选平台模型未声明当前请求场景的 output_resolutions。"
|
||||
return false, detail
|
||||
}
|
||||
detail["reason"] = "output_resolutions_empty"
|
||||
detail["message"] = "候选平台模型声明了空分辨率白名单,按未限制处理。"
|
||||
return true, detail
|
||||
}
|
||||
if containsResolution(allowed, requirement.Resolution) {
|
||||
detail["reason"] = "supported"
|
||||
@@ -162,7 +183,7 @@ func outputResolutionAllowedValues(value any, scopes []string) ([]string, bool)
|
||||
for _, raw := range typed {
|
||||
values = append(values, stringListFromAny(raw)...)
|
||||
}
|
||||
return uniqueStringList(values), len(values) > 0
|
||||
return uniqueStringList(values), true
|
||||
}
|
||||
return nil, true
|
||||
default:
|
||||
@@ -170,6 +191,148 @@ func outputResolutionAllowedValues(value any, scopes []string) ([]string, bool)
|
||||
}
|
||||
}
|
||||
|
||||
func outputResolutionScopeConfigured(value any, scopes []string) bool {
|
||||
typed, ok := value.(map[string]any)
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
for _, scope := range append(append([]string(nil), scopes...), "default", "*", "all") {
|
||||
if scope == "" {
|
||||
continue
|
||||
}
|
||||
if _, exists := typed[scope]; exists {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
type requestAspectRatioRequirement struct {
|
||||
Kind string
|
||||
RequestedModel string
|
||||
ModelType string
|
||||
AspectRatio string
|
||||
Source string
|
||||
Resolution string
|
||||
}
|
||||
|
||||
func filterRuntimeCandidatesByRequestAspectRatio(kind string, requestedModel string, modelType string, body map[string]any, candidates []store.RuntimeModelCandidate) ([]store.RuntimeModelCandidate, map[string]any, error) {
|
||||
requirement, ok := requestAspectRatioRequirementFor(kind, requestedModel, modelType, body)
|
||||
if !ok || len(candidates) == 0 {
|
||||
return candidates, nil, nil
|
||||
}
|
||||
|
||||
filtered := make([]store.RuntimeModelCandidate, 0, len(candidates))
|
||||
rejected := make([]map[string]any, 0)
|
||||
supportedRatios := make([]string, 0)
|
||||
for _, candidate := range candidates {
|
||||
modelType := firstNonEmptyString(candidate.ModelType, requirement.ModelType)
|
||||
capability := capabilityForType(effectiveModelCapability(candidate), modelType)
|
||||
detail := candidateAspectRatioDetail(candidate, requirement, modelType)
|
||||
if capability == nil {
|
||||
detail["reason"] = "capability_unrestricted"
|
||||
filtered = append(filtered, candidate)
|
||||
continue
|
||||
}
|
||||
allowed, configured := outputResolutionAllowedValues(capability["aspect_ratio_allowed"], []string{requirement.Resolution})
|
||||
detail["allowedAspectRatios"] = allowed
|
||||
detail["capabilityPath"] = capabilityPath(modelType, "aspect_ratio_allowed")
|
||||
detail["capabilityValue"] = cloneAny(capability["aspect_ratio_allowed"])
|
||||
for _, value := range allowed {
|
||||
appendUniqueString(&supportedRatios, value)
|
||||
}
|
||||
if !configured || len(allowed) == 0 || containsResolution(allowed, requirement.AspectRatio) {
|
||||
detail["reason"] = "supported"
|
||||
filtered = append(filtered, candidate)
|
||||
continue
|
||||
}
|
||||
detail["reason"] = "aspect_ratio_not_allowed"
|
||||
detail["message"] = "请求比例不在候选平台模型 aspect_ratio_allowed 中。"
|
||||
rejected = append(rejected, detail)
|
||||
}
|
||||
|
||||
summary := map[string]any{
|
||||
"code": unsupportedRequestAspectRatioCode,
|
||||
"filter": "request_aspect_ratio",
|
||||
"kind": requirement.Kind,
|
||||
"requestedModel": requirement.RequestedModel,
|
||||
"modelType": requirement.ModelType,
|
||||
"requestedAspectRatio": requirement.AspectRatio,
|
||||
"aspectRatioSource": requirement.Source,
|
||||
"resolution": requirement.Resolution,
|
||||
"capabilityPath": capabilityPath(requirement.ModelType, "aspect_ratio_allowed"),
|
||||
"candidateCount": len(candidates),
|
||||
"supportedCandidateCount": len(filtered),
|
||||
"filteredCandidateCount": len(rejected),
|
||||
"supportedAspectRatios": uniqueStringList(supportedRatios),
|
||||
"rejectedCandidates": rejected,
|
||||
}
|
||||
if len(filtered) == 0 {
|
||||
return nil, summary, &store.ModelCandidateUnavailableError{
|
||||
Code: unsupportedRequestAspectRatioCode,
|
||||
Message: fmt.Sprintf("请求的%s比例 %s 没有可用平台模型支持,已过滤 %d 个候选平台模型", mediaResourceName(requirement.ModelType), requirement.AspectRatio, len(rejected)),
|
||||
Details: summary,
|
||||
}
|
||||
}
|
||||
return filtered, summary, nil
|
||||
}
|
||||
|
||||
func requestAspectRatioRequirementFor(kind string, requestedModel string, modelType string, body map[string]any) (requestAspectRatioRequirement, bool) {
|
||||
if !isResolutionFilteredModelType(modelType) {
|
||||
return requestAspectRatioRequirement{}, false
|
||||
}
|
||||
for _, key := range []string{"aspect_ratio", "aspectRatio", "ratio"} {
|
||||
raw, exists := body[key]
|
||||
if !exists {
|
||||
continue
|
||||
}
|
||||
aspectRatio, ok := normalizedAspectRatio(stringFromAny(raw))
|
||||
if !ok {
|
||||
return requestAspectRatioRequirement{}, false
|
||||
}
|
||||
resolution, _ := requestResolutionValue(body, modelType)
|
||||
return requestAspectRatioRequirement{
|
||||
Kind: kind,
|
||||
RequestedModel: requestedModel,
|
||||
ModelType: modelType,
|
||||
AspectRatio: aspectRatio,
|
||||
Source: key,
|
||||
Resolution: resolution,
|
||||
}, true
|
||||
}
|
||||
return requestAspectRatioRequirement{}, false
|
||||
}
|
||||
|
||||
func candidateAspectRatioDetail(candidate store.RuntimeModelCandidate, requirement requestAspectRatioRequirement, modelType string) map[string]any {
|
||||
return map[string]any{
|
||||
"platformId": candidate.PlatformID,
|
||||
"platformKey": candidate.PlatformKey,
|
||||
"platformName": candidate.PlatformName,
|
||||
"provider": candidate.Provider,
|
||||
"platformModelId": candidate.PlatformModelID,
|
||||
"modelName": candidate.ModelName,
|
||||
"modelAlias": candidate.ModelAlias,
|
||||
"displayName": candidate.DisplayName,
|
||||
"providerModelName": candidate.ProviderModelName,
|
||||
"modelType": modelType,
|
||||
"requested": map[string]any{
|
||||
"aspectRatio": requirement.AspectRatio,
|
||||
"source": requirement.Source,
|
||||
"resolution": requirement.Resolution,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func mediaResourceName(modelType string) string {
|
||||
if modelType == "image_generate" || modelType == "image_edit" {
|
||||
return "图像"
|
||||
}
|
||||
if isVideoModelType(modelType) {
|
||||
return "视频"
|
||||
}
|
||||
return "媒体"
|
||||
}
|
||||
|
||||
func containsResolution(values []string, target string) bool {
|
||||
for _, value := range values {
|
||||
if strings.EqualFold(strings.TrimSpace(value), strings.TrimSpace(target)) {
|
||||
|
||||
@@ -0,0 +1,49 @@
|
||||
package runner
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
)
|
||||
|
||||
func filterCandidatesByRequestedPlatform(
|
||||
candidates []store.RuntimeModelCandidate,
|
||||
body map[string]any,
|
||||
) ([]store.RuntimeModelCandidate, error) {
|
||||
platformID := firstNonEmptyString(
|
||||
stringFromAny(body["platform_id"]),
|
||||
stringFromAny(body["platformId"]),
|
||||
)
|
||||
platformModelID := firstNonEmptyString(
|
||||
stringFromAny(body["platform_model_id"]),
|
||||
stringFromAny(body["platformModelId"]),
|
||||
)
|
||||
if platformID == "" && platformModelID == "" {
|
||||
return candidates, nil
|
||||
}
|
||||
|
||||
filtered := make([]store.RuntimeModelCandidate, 0, len(candidates))
|
||||
for _, candidate := range candidates {
|
||||
if platformID != "" && !strings.EqualFold(candidate.PlatformID, platformID) {
|
||||
continue
|
||||
}
|
||||
if platformModelID != "" && !strings.EqualFold(candidate.PlatformModelID, platformModelID) {
|
||||
continue
|
||||
}
|
||||
filtered = append(filtered, candidate)
|
||||
}
|
||||
if len(filtered) > 0 {
|
||||
return filtered, nil
|
||||
}
|
||||
|
||||
requested := platformID
|
||||
if platformModelID != "" {
|
||||
requested = platformModelID
|
||||
}
|
||||
return nil, fmt.Errorf(
|
||||
"%w: requested platform %q is unavailable for this model",
|
||||
store.ErrNoModelCandidate,
|
||||
requested,
|
||||
)
|
||||
}
|
||||
@@ -0,0 +1,49 @@
|
||||
package runner
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"testing"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
)
|
||||
|
||||
func TestFilterCandidatesByRequestedPlatform(t *testing.T) {
|
||||
candidates := []store.RuntimeModelCandidate{
|
||||
{PlatformID: "platform-a", PlatformModelID: "model-a"},
|
||||
{PlatformID: "platform-b", PlatformModelID: "model-b"},
|
||||
}
|
||||
|
||||
t.Run("without selection", func(t *testing.T) {
|
||||
filtered, err := filterCandidatesByRequestedPlatform(candidates, map[string]any{})
|
||||
if err != nil || len(filtered) != len(candidates) {
|
||||
t.Fatalf("unexpected unfiltered candidates: %+v err=%v", filtered, err)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("snake case platform id", func(t *testing.T) {
|
||||
filtered, err := filterCandidatesByRequestedPlatform(candidates, map[string]any{
|
||||
"platform_id": "platform-b",
|
||||
})
|
||||
if err != nil || len(filtered) != 1 || filtered[0].PlatformID != "platform-b" {
|
||||
t.Fatalf("unexpected platform filter result: %+v err=%v", filtered, err)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("camel case platform model id", func(t *testing.T) {
|
||||
filtered, err := filterCandidatesByRequestedPlatform(candidates, map[string]any{
|
||||
"platformModelId": "model-a",
|
||||
})
|
||||
if err != nil || len(filtered) != 1 || filtered[0].PlatformModelID != "model-a" {
|
||||
t.Fatalf("unexpected platform model filter result: %+v err=%v", filtered, err)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("selection must be accessible and enabled", func(t *testing.T) {
|
||||
filtered, err := filterCandidatesByRequestedPlatform(candidates, map[string]any{
|
||||
"platformId": "platform-missing",
|
||||
})
|
||||
if len(filtered) != 0 || !errors.Is(err, store.ErrNoModelCandidate) {
|
||||
t.Fatalf("unexpected unavailable platform result: %+v err=%v", filtered, err)
|
||||
}
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,381 @@
|
||||
package runner
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"fmt"
|
||||
"image"
|
||||
imagedraw "image/draw"
|
||||
"image/jpeg"
|
||||
_ "image/png"
|
||||
"io"
|
||||
"math"
|
||||
"net/http"
|
||||
"os"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/clients"
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
_ "golang.org/x/image/bmp"
|
||||
"golang.org/x/image/draw"
|
||||
_ "golang.org/x/image/tiff"
|
||||
_ "golang.org/x/image/webp"
|
||||
)
|
||||
|
||||
const (
|
||||
maxInputImageConversionBytes = 32 << 20
|
||||
maxInputImageConversionPixels = 100_000_000
|
||||
)
|
||||
|
||||
type inputImageResolutionRange struct {
|
||||
MinLong int
|
||||
MinShort int
|
||||
MaxLong int
|
||||
MaxShort int
|
||||
}
|
||||
|
||||
type inputImageConstraints struct {
|
||||
Resolution inputImageResolutionRange
|
||||
MinAspect float64
|
||||
MaxAspect float64
|
||||
}
|
||||
|
||||
func (s *Service) normalizeVideoInputImages(ctx context.Context, body map[string]any, candidate store.RuntimeModelCandidate) (map[string]any, error) {
|
||||
constraints, ok := candidateInputImageConstraints(candidate)
|
||||
if !ok {
|
||||
return body, nil
|
||||
}
|
||||
value, err := s.normalizeVideoInputImageValue(ctx, body, nil, constraints)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out, _ := value.(map[string]any)
|
||||
if out == nil {
|
||||
return map[string]any{}, nil
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (s *Service) normalizeVideoInputImageValue(ctx context.Context, value any, path []string, constraints inputImageConstraints) (any, error) {
|
||||
switch typed := value.(type) {
|
||||
case map[string]any:
|
||||
next := make(map[string]any, len(typed))
|
||||
for key, item := range typed {
|
||||
normalized, err := s.normalizeVideoInputImageValue(ctx, item, append(path, key), constraints)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
next[key] = normalized
|
||||
}
|
||||
return next, nil
|
||||
case []any:
|
||||
next := make([]any, 0, len(typed))
|
||||
for index, item := range typed {
|
||||
normalized, err := s.normalizeVideoInputImageValue(ctx, item, append(path, fmt.Sprintf("[%d]", index)), constraints)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
next = append(next, normalized)
|
||||
}
|
||||
return next, nil
|
||||
case string:
|
||||
if !imageInputFieldNeedsHydration(path) {
|
||||
return value, nil
|
||||
}
|
||||
normalized, converted, original, target, err := s.normalizeVideoInputImageSource(ctx, typed, constraints)
|
||||
if err != nil {
|
||||
param := requestInputImageParam(path)
|
||||
return nil, &clients.ClientError{
|
||||
Code: "invalid_parameter",
|
||||
Message: fmt.Sprintf("输入图片 %s 自动转换失败:%s", param, err.Error()),
|
||||
Param: param,
|
||||
StatusCode: http.StatusBadRequest,
|
||||
Retryable: false,
|
||||
}
|
||||
}
|
||||
if converted && s.logger != nil {
|
||||
s.logger.Info(
|
||||
"video input image auto-converted",
|
||||
"param", requestInputImageParam(path),
|
||||
"original", fmt.Sprintf("%dx%d", original.X, original.Y),
|
||||
"converted", fmt.Sprintf("%dx%d", target.X, target.Y),
|
||||
)
|
||||
}
|
||||
return normalized, nil
|
||||
default:
|
||||
return value, nil
|
||||
}
|
||||
}
|
||||
|
||||
func candidateInputImageConstraints(candidate store.RuntimeModelCandidate) (inputImageConstraints, bool) {
|
||||
modelTypes := []string{candidate.ModelType, "image_to_video", "omni_video"}
|
||||
seen := map[string]bool{}
|
||||
for _, modelType := range modelTypes {
|
||||
modelType = strings.TrimSpace(modelType)
|
||||
if modelType == "" || seen[modelType] {
|
||||
continue
|
||||
}
|
||||
seen[modelType] = true
|
||||
values, _ := candidate.Capabilities[modelType].(map[string]any)
|
||||
if constraints, ok := parseInputImageConstraints(values); ok {
|
||||
return constraints, true
|
||||
}
|
||||
}
|
||||
return inputImageConstraints{}, false
|
||||
}
|
||||
|
||||
func parseInputImageConstraints(values map[string]any) (inputImageConstraints, bool) {
|
||||
if values == nil {
|
||||
return inputImageConstraints{}, false
|
||||
}
|
||||
resolution, _ := values["input_image_resolution_range"].(map[string]any)
|
||||
minimum, _ := resolution["min"].(map[string]any)
|
||||
maximum, _ := resolution["max"].(map[string]any)
|
||||
aspect, ok := inputImageNumberPair(values["input_image_aspect_ratio_range"])
|
||||
constraints := inputImageConstraints{
|
||||
Resolution: inputImageResolutionRange{
|
||||
MinLong: int(math.Round(floatFromAny(minimum["long_edge"]))),
|
||||
MinShort: int(math.Round(floatFromAny(minimum["short_edge"]))),
|
||||
MaxLong: int(math.Round(floatFromAny(maximum["long_edge"]))),
|
||||
MaxShort: int(math.Round(floatFromAny(maximum["short_edge"]))),
|
||||
},
|
||||
MinAspect: aspect[0],
|
||||
MaxAspect: aspect[1],
|
||||
}
|
||||
if !ok ||
|
||||
constraints.Resolution.MinLong <= 0 ||
|
||||
constraints.Resolution.MinShort <= 0 ||
|
||||
constraints.Resolution.MaxLong < constraints.Resolution.MinLong ||
|
||||
constraints.Resolution.MaxShort < constraints.Resolution.MinShort ||
|
||||
constraints.MinAspect <= 0 ||
|
||||
constraints.MaxAspect < constraints.MinAspect {
|
||||
return inputImageConstraints{}, false
|
||||
}
|
||||
return constraints, true
|
||||
}
|
||||
|
||||
func inputImageNumberPair(value any) ([2]float64, bool) {
|
||||
switch typed := value.(type) {
|
||||
case []any:
|
||||
if len(typed) != 2 {
|
||||
return [2]float64{}, false
|
||||
}
|
||||
pair := [2]float64{floatFromAny(typed[0]), floatFromAny(typed[1])}
|
||||
return pair, pair[0] > 0 && pair[1] > 0
|
||||
case []float64:
|
||||
if len(typed) != 2 {
|
||||
return [2]float64{}, false
|
||||
}
|
||||
return [2]float64{typed[0], typed[1]}, typed[0] > 0 && typed[1] > 0
|
||||
default:
|
||||
return [2]float64{}, false
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Service) normalizeVideoInputImageSource(ctx context.Context, source string, constraints inputImageConstraints) (string, bool, image.Point, image.Point, error) {
|
||||
source = strings.TrimSpace(source)
|
||||
if source == "" || strings.HasPrefix(strings.ToLower(source), "asset://") {
|
||||
return source, false, image.Point{}, image.Point{}, nil
|
||||
}
|
||||
payload, err := s.readVideoInputImageBytes(ctx, source)
|
||||
if err != nil {
|
||||
return "", false, image.Point{}, image.Point{}, err
|
||||
}
|
||||
config, _, err := image.DecodeConfig(bytes.NewReader(payload))
|
||||
if err != nil || config.Width <= 0 || config.Height <= 0 {
|
||||
return "", false, image.Point{}, image.Point{}, fmt.Errorf("图片无法解码")
|
||||
}
|
||||
if config.Width > 50_000 || config.Height > 50_000 || int64(config.Width)*int64(config.Height) > maxInputImageConversionPixels {
|
||||
return "", false, image.Point{}, image.Point{}, fmt.Errorf("图片像素数量超过转换上限")
|
||||
}
|
||||
original := image.Pt(config.Width, config.Height)
|
||||
if inputImageWithinConstraints(original, constraints) {
|
||||
return source, false, original, original, nil
|
||||
}
|
||||
|
||||
decoded, _, err := image.Decode(bytes.NewReader(payload))
|
||||
if err != nil {
|
||||
return "", false, original, image.Point{}, fmt.Errorf("图片无法解码")
|
||||
}
|
||||
target := resolveInputImageTarget(original, constraints)
|
||||
canvas := image.NewNRGBA(image.Rect(0, 0, target.X, target.Y))
|
||||
imagedraw.Draw(canvas, canvas.Bounds(), &image.Uniform{C: averageImageColor(decoded)}, image.Point{}, imagedraw.Src)
|
||||
|
||||
sourceBounds := decoded.Bounds()
|
||||
scale := math.Min(float64(target.X)/float64(sourceBounds.Dx()), float64(target.Y)/float64(sourceBounds.Dy()))
|
||||
contentWidth := max(1, int(math.Round(float64(sourceBounds.Dx())*scale)))
|
||||
contentHeight := max(1, int(math.Round(float64(sourceBounds.Dy())*scale)))
|
||||
left := (target.X - contentWidth) / 2
|
||||
top := (target.Y - contentHeight) / 2
|
||||
draw.CatmullRom.Scale(
|
||||
canvas,
|
||||
image.Rect(left, top, left+contentWidth, top+contentHeight),
|
||||
decoded,
|
||||
sourceBounds,
|
||||
draw.Over,
|
||||
nil,
|
||||
)
|
||||
|
||||
var output bytes.Buffer
|
||||
if err := jpeg.Encode(&output, canvas, &jpeg.Options{Quality: 90}); err != nil {
|
||||
return "", false, original, target, fmt.Errorf("图片编码失败")
|
||||
}
|
||||
if !inputImageWithinConstraints(target, constraints) {
|
||||
return "", false, original, target, fmt.Errorf("转换结果仍不符合平台限制")
|
||||
}
|
||||
return "data:image/jpeg;base64," + base64.StdEncoding.EncodeToString(output.Bytes()), true, original, target, nil
|
||||
}
|
||||
|
||||
func (s *Service) readVideoInputImageBytes(ctx context.Context, source string) ([]byte, error) {
|
||||
lower := strings.ToLower(source)
|
||||
if strings.HasPrefix(lower, "data:") || (!strings.Contains(source, "://") && !strings.HasPrefix(source, "/")) {
|
||||
payload, err := decodeBase64Payload(source)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("图片数据不是有效 Base64")
|
||||
}
|
||||
if len(payload) > maxInputImageConversionBytes {
|
||||
return nil, fmt.Errorf("图片超过 32 MiB 转换上限")
|
||||
}
|
||||
return payload, nil
|
||||
}
|
||||
if localPath := s.localPathFromRequestAssetURL(source); localPath != "" {
|
||||
payload, err := os.ReadFile(localPath)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("托管图片读取失败")
|
||||
}
|
||||
if len(payload) > maxInputImageConversionBytes {
|
||||
return nil, fmt.Errorf("图片超过 32 MiB 转换上限")
|
||||
}
|
||||
return payload, nil
|
||||
}
|
||||
if !strings.HasPrefix(lower, "http://") && !strings.HasPrefix(lower, "https://") {
|
||||
return nil, fmt.Errorf("图片地址格式不受支持")
|
||||
}
|
||||
|
||||
requestCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
|
||||
defer cancel()
|
||||
request, err := http.NewRequestWithContext(requestCtx, http.MethodGet, source, nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("图片地址无效")
|
||||
}
|
||||
httpClient := generatedAssetHTTPClient(false)
|
||||
httpClient.Timeout = 10 * time.Second
|
||||
httpClient.CheckRedirect = func(request *http.Request, via []*http.Request) error {
|
||||
if len(via) >= 3 {
|
||||
return fmt.Errorf("图片地址重定向次数过多")
|
||||
}
|
||||
if !requestAssetURLIsPublic("", request.URL.String()) {
|
||||
return fmt.Errorf("图片地址重定向到受限网络")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
response, err := httpClient.Do(request)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("图片读取失败")
|
||||
}
|
||||
defer response.Body.Close()
|
||||
if response.StatusCode < 200 || response.StatusCode >= 300 {
|
||||
return nil, fmt.Errorf("图片读取失败,HTTP %d", response.StatusCode)
|
||||
}
|
||||
payload, err := io.ReadAll(io.LimitReader(response.Body, maxInputImageConversionBytes+1))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("图片读取失败")
|
||||
}
|
||||
if len(payload) > maxInputImageConversionBytes {
|
||||
return nil, fmt.Errorf("图片超过 32 MiB 转换上限")
|
||||
}
|
||||
return payload, nil
|
||||
}
|
||||
|
||||
func inputImageWithinConstraints(size image.Point, constraints inputImageConstraints) bool {
|
||||
longEdge := max(size.X, size.Y)
|
||||
shortEdge := min(size.X, size.Y)
|
||||
ratio := float64(size.X) / float64(size.Y)
|
||||
return longEdge >= constraints.Resolution.MinLong &&
|
||||
longEdge <= constraints.Resolution.MaxLong &&
|
||||
shortEdge >= constraints.Resolution.MinShort &&
|
||||
shortEdge <= constraints.Resolution.MaxShort &&
|
||||
ratio >= constraints.MinAspect &&
|
||||
ratio <= constraints.MaxAspect
|
||||
}
|
||||
|
||||
func resolveInputImageTarget(source image.Point, constraints inputImageConstraints) image.Point {
|
||||
sourceRatio := float64(source.X) / float64(source.Y)
|
||||
targetRatio := math.Min(constraints.MaxAspect, math.Max(constraints.MinAspect, sourceRatio))
|
||||
landscape := targetRatio >= 1
|
||||
normalizedRatio := targetRatio
|
||||
var sourceShort float64
|
||||
if landscape {
|
||||
sourceShort = math.Max(float64(source.Y), float64(source.X)/targetRatio)
|
||||
} else {
|
||||
normalizedRatio = 1 / targetRatio
|
||||
sourceShort = math.Max(float64(source.X), float64(source.Y)*targetRatio)
|
||||
}
|
||||
|
||||
shortEdge := int(math.Round(sourceShort))
|
||||
shortEdge = min(constraints.Resolution.MaxShort, max(constraints.Resolution.MinShort, shortEdge))
|
||||
longEdge := int(math.Floor(float64(shortEdge) * normalizedRatio))
|
||||
if longEdge > constraints.Resolution.MaxLong {
|
||||
longEdge = constraints.Resolution.MaxLong
|
||||
shortEdge = int(math.Ceil(float64(longEdge) / normalizedRatio))
|
||||
}
|
||||
if longEdge < constraints.Resolution.MinLong {
|
||||
longEdge = constraints.Resolution.MinLong
|
||||
shortEdge = int(math.Ceil(float64(longEdge) / normalizedRatio))
|
||||
}
|
||||
shortEdge = min(constraints.Resolution.MaxShort, max(constraints.Resolution.MinShort, shortEdge))
|
||||
longEdge = int(math.Floor(float64(shortEdge) * normalizedRatio))
|
||||
longEdge = min(constraints.Resolution.MaxLong, max(constraints.Resolution.MinLong, longEdge))
|
||||
if landscape {
|
||||
return image.Pt(longEdge, shortEdge)
|
||||
}
|
||||
return image.Pt(shortEdge, longEdge)
|
||||
}
|
||||
|
||||
func averageImageColor(source image.Image) imageColor {
|
||||
bounds := source.Bounds()
|
||||
stepX := max(1, bounds.Dx()/32)
|
||||
stepY := max(1, bounds.Dy()/32)
|
||||
var red, green, blue, count uint64
|
||||
for y := bounds.Min.Y; y < bounds.Max.Y; y += stepY {
|
||||
for x := bounds.Min.X; x < bounds.Max.X; x += stepX {
|
||||
r, g, b, _ := source.At(x, y).RGBA()
|
||||
red += uint64(r >> 8)
|
||||
green += uint64(g >> 8)
|
||||
blue += uint64(b >> 8)
|
||||
count++
|
||||
}
|
||||
}
|
||||
if count == 0 {
|
||||
return imageColor{R: 255, G: 255, B: 255, A: 255}
|
||||
}
|
||||
return imageColor{R: uint8(red / count), G: uint8(green / count), B: uint8(blue / count), A: 255}
|
||||
}
|
||||
|
||||
type imageColor struct {
|
||||
R, G, B, A uint8
|
||||
}
|
||||
|
||||
func (c imageColor) RGBA() (r, g, b, a uint32) {
|
||||
return uint32(c.R) * 0x101, uint32(c.G) * 0x101, uint32(c.B) * 0x101, uint32(c.A) * 0x101
|
||||
}
|
||||
|
||||
func requestInputImageParam(path []string) string {
|
||||
var builder strings.Builder
|
||||
for _, segment := range path {
|
||||
if strings.HasPrefix(segment, "[") {
|
||||
builder.WriteString(segment)
|
||||
continue
|
||||
}
|
||||
if builder.Len() > 0 {
|
||||
builder.WriteByte('.')
|
||||
}
|
||||
builder.WriteString(segment)
|
||||
}
|
||||
if builder.Len() == 0 {
|
||||
return "image"
|
||||
}
|
||||
return builder.String()
|
||||
}
|
||||
@@ -0,0 +1,140 @@
|
||||
package runner
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"image"
|
||||
"image/color"
|
||||
"image/jpeg"
|
||||
"image/png"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/clients"
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
)
|
||||
|
||||
func seedanceImageConstraintCandidate() store.RuntimeModelCandidate {
|
||||
return store.RuntimeModelCandidate{
|
||||
ModelType: "image_to_video",
|
||||
Capabilities: map[string]any{
|
||||
"image_to_video": map[string]any{
|
||||
"input_image_resolution_range": map[string]any{
|
||||
"min": map[string]any{"long_edge": float64(360), "short_edge": float64(360)},
|
||||
"max": map[string]any{"long_edge": float64(1920), "short_edge": float64(1080)},
|
||||
},
|
||||
"input_image_aspect_ratio_range": []any{float64(0.39), float64(2.5)},
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func pngImageDataURL(t *testing.T, width int, height int) string {
|
||||
t.Helper()
|
||||
source := image.NewRGBA(image.Rect(0, 0, width, height))
|
||||
for y := 0; y < height; y++ {
|
||||
for x := 0; x < width; x++ {
|
||||
source.Set(x, y, color.RGBA{R: 80, G: 120, B: 160, A: 255})
|
||||
}
|
||||
}
|
||||
var payload bytes.Buffer
|
||||
if err := png.Encode(&payload, source); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return "data:image/png;base64," + base64.StdEncoding.EncodeToString(payload.Bytes())
|
||||
}
|
||||
|
||||
func decodedImageSize(t *testing.T, source string) image.Point {
|
||||
t.Helper()
|
||||
if !strings.HasPrefix(source, "data:image/jpeg;base64,") {
|
||||
t.Fatalf("expected converted JPEG data URL, got prefix %.32q", source)
|
||||
}
|
||||
payload, err := base64.StdEncoding.DecodeString(strings.TrimPrefix(source, "data:image/jpeg;base64,"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
config, err := jpeg.DecodeConfig(bytes.NewReader(payload))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return image.Pt(config.Width, config.Height)
|
||||
}
|
||||
|
||||
func TestNormalizeVideoInputImagesResizesOversizedSeedanceImage(t *testing.T) {
|
||||
service := &Service{}
|
||||
body := map[string]any{
|
||||
"content": []any{
|
||||
map[string]any{
|
||||
"type": "image_url",
|
||||
"role": "first_frame",
|
||||
"image_url": map[string]any{"url": pngImageDataURL(t, 3840, 2160)},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
normalized, err := service.normalizeVideoInputImages(context.Background(), body, seedanceImageConstraintCandidate())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
content := normalized["content"].([]any)
|
||||
item := content[0].(map[string]any)
|
||||
imageURL := item["image_url"].(map[string]any)["url"].(string)
|
||||
if got := decodedImageSize(t, imageURL); got != image.Pt(1920, 1080) {
|
||||
t.Fatalf("unexpected converted size: %v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNormalizeVideoInputImagesPadsAspectRatioViolation(t *testing.T) {
|
||||
service := &Service{}
|
||||
source := pngImageDataURL(t, 1530, 500)
|
||||
body := map[string]any{"first_frame": source}
|
||||
|
||||
normalized, err := service.normalizeVideoInputImages(context.Background(), body, seedanceImageConstraintCandidate())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got := decodedImageSize(t, normalized["first_frame"].(string)); got != image.Pt(1530, 612) {
|
||||
t.Fatalf("unexpected converted size: %v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNormalizeVideoInputImagesUpscalesBelowMinimumResolution(t *testing.T) {
|
||||
service := &Service{}
|
||||
body := map[string]any{"first_frame": pngImageDataURL(t, 200, 100)}
|
||||
|
||||
normalized, err := service.normalizeVideoInputImages(context.Background(), body, seedanceImageConstraintCandidate())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got := decodedImageSize(t, normalized["first_frame"].(string)); got != image.Pt(720, 360) {
|
||||
t.Fatalf("unexpected converted size: %v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNormalizeVideoInputImagesKeepsCompliantImage(t *testing.T) {
|
||||
service := &Service{}
|
||||
source := pngImageDataURL(t, 1280, 720)
|
||||
body := map[string]any{"first_frame": source}
|
||||
|
||||
normalized, err := service.normalizeVideoInputImages(context.Background(), body, seedanceImageConstraintCandidate())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if normalized["first_frame"] != source {
|
||||
t.Fatal("compliant image should remain unchanged")
|
||||
}
|
||||
}
|
||||
|
||||
func TestNormalizeVideoInputImagesReturnsActionableDecodeError(t *testing.T) {
|
||||
service := &Service{}
|
||||
body := map[string]any{"first_frame": "data:image/png;base64,bm90LWltYWdl"}
|
||||
|
||||
_, err := service.normalizeVideoInputImages(context.Background(), body, seedanceImageConstraintCandidate())
|
||||
if err == nil {
|
||||
t.Fatal("expected decode error")
|
||||
}
|
||||
if clients.ErrorCode(err) != "invalid_parameter" || clients.ErrorParam(err) != "first_frame" {
|
||||
t.Fatalf("unexpected error contract: code=%s param=%s err=%v", clients.ErrorCode(err), clients.ErrorParam(err), err)
|
||||
}
|
||||
}
|
||||
@@ -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 {
|
||||
|
||||
@@ -102,8 +102,10 @@ func TestEffectiveRateLimitPolicyTreatsEmptyRuntimePolicyAsUnlimited(t *testing.
|
||||
PlatformRateLimitPolicy: map[string]any{"rules": []any{
|
||||
map[string]any{"metric": "rpm", "limit": 500},
|
||||
}},
|
||||
RuntimePolicySetID: "runtime-policy-1",
|
||||
RuntimeRateLimitPolicy: map[string]any{"rules": []any{}},
|
||||
RuntimePolicySetID: "runtime-policy-1",
|
||||
RuntimePolicyExplicit: true,
|
||||
RuntimeRateLimitPolicy: map[string]any{"rules": []any{}},
|
||||
ModelRateLimitPolicyMode: "inherit",
|
||||
})
|
||||
|
||||
if hasRules(policy) {
|
||||
|
||||
@@ -0,0 +1,173 @@
|
||||
package runner
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
)
|
||||
|
||||
func TestParamProcessorTreatsNonRatioValuesAsUnspecifiedForImagesAndVideos(t *testing.T) {
|
||||
for _, modelType := range []string{"image_generate", "video_generate"} {
|
||||
t.Run(modelType, func(t *testing.T) {
|
||||
result := preprocessRequestWithLog(mediaKindForTest(modelType), map[string]any{
|
||||
"aspect_ratio": "auto",
|
||||
"aspectRatio": "adaptive",
|
||||
"ratio": "keep_ratio",
|
||||
}, store.RuntimeModelCandidate{
|
||||
ModelType: modelType,
|
||||
Capabilities: map[string]any{
|
||||
modelType: map[string]any{
|
||||
"aspect_ratio_allowed": []any{"adaptive"},
|
||||
},
|
||||
},
|
||||
})
|
||||
for _, key := range []string{"aspect_ratio", "aspectRatio", "ratio"} {
|
||||
if _, ok := result.Body[key]; ok {
|
||||
t.Fatalf("%s should be removed for non-ratio input: %+v", key, result.Body)
|
||||
}
|
||||
}
|
||||
if result.Err != nil {
|
||||
t.Fatalf("non-ratio input should use provider default instead of failing: %v", result.Err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestParamProcessorNormalizesImageKSizeAndPreservesProviderResolutionChoice(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
capability map[string]any
|
||||
wantSize string
|
||||
}{
|
||||
{
|
||||
name: "dimension upstream",
|
||||
capability: map[string]any{
|
||||
"output_resolutions": []any{"1K", "2K", "4K"},
|
||||
"aspect_ratio_allowed": []any{"16:9"},
|
||||
},
|
||||
wantSize: "2048x1152",
|
||||
},
|
||||
{
|
||||
name: "resolution upstream",
|
||||
capability: map[string]any{
|
||||
"output_resolutions": []any{"1K", "2K"},
|
||||
"aspect_ratio_allowed": []any{"16:9"},
|
||||
"size_param_format": "resolution",
|
||||
},
|
||||
wantSize: "2K",
|
||||
},
|
||||
}
|
||||
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
result := preprocessRequestWithLog("images.generations", map[string]any{
|
||||
"size": "2k",
|
||||
"aspect_ratio": "16 : 9",
|
||||
}, store.RuntimeModelCandidate{
|
||||
ModelType: "image_generate",
|
||||
Capabilities: map[string]any{
|
||||
"image_generate": test.capability,
|
||||
},
|
||||
})
|
||||
if result.Err != nil {
|
||||
t.Fatalf("preprocess image size: %v", result.Err)
|
||||
}
|
||||
if result.Body["resolution"] != "2K" || result.Body["size"] != test.wantSize {
|
||||
t.Fatalf("unexpected normalized size: %+v", result.Body)
|
||||
}
|
||||
if result.Body["width"] != 2048 || result.Body["height"] != 1152 || result.Body["aspect_ratio"] != "16:9" {
|
||||
t.Fatalf("unexpected normalized dimensions: %+v", result.Body)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestParamProcessorNormalizesVideo4KSize(t *testing.T) {
|
||||
result := preprocessRequestWithLog("videos.generations", map[string]any{
|
||||
"size": "4K",
|
||||
"aspect_ratio": "auto",
|
||||
}, store.RuntimeModelCandidate{
|
||||
ModelType: "video_generate",
|
||||
Capabilities: map[string]any{
|
||||
"video_generate": map[string]any{
|
||||
"output_resolutions": []any{"2160p"},
|
||||
"aspect_ratio_allowed": []any{"16:9"},
|
||||
},
|
||||
},
|
||||
})
|
||||
if result.Body["resolution"] != "2160p" {
|
||||
t.Fatalf("video 4K size should normalize to 2160p: %+v", result.Body)
|
||||
}
|
||||
if _, ok := result.Body["aspect_ratio"]; ok {
|
||||
t.Fatalf("video auto ratio should stay unspecified: %+v", result.Body)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCandidateFilterIgnoresNonRatioValuesAndFiltersRealRatios(t *testing.T) {
|
||||
candidates := []store.RuntimeModelCandidate{
|
||||
candidateWithAspectRatios("square", "image_generate", "1:1"),
|
||||
candidateWithAspectRatios("landscape", "image_generate", "16:9"),
|
||||
}
|
||||
|
||||
filtered, summary, err := filterRuntimeCandidatesByRequest("images.generations", "demo-image", "image_generate", map[string]any{
|
||||
"aspect_ratio": "auto",
|
||||
}, candidates)
|
||||
if err != nil || len(filtered) != 2 || summary != nil {
|
||||
t.Fatalf("auto ratio should not participate in candidate matching: filtered=%+v summary=%+v err=%v", filtered, summary, err)
|
||||
}
|
||||
|
||||
filtered, summary, err = filterRuntimeCandidatesByRequest("images.generations", "demo-image", "image_generate", map[string]any{
|
||||
"aspect_ratio": "16 : 9",
|
||||
}, candidates)
|
||||
if err != nil || len(filtered) != 1 || filtered[0].PlatformKey != "landscape" {
|
||||
t.Fatalf("real ratio should filter candidates: filtered=%+v summary=%+v err=%v", filtered, summary, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCandidateFilterNormalizesKSizeAndAllowsUnconfiguredResolution(t *testing.T) {
|
||||
candidates := []store.RuntimeModelCandidate{
|
||||
candidateWithImageResolutions("one-k", "1K"),
|
||||
{
|
||||
PlatformID: "platform-unrestricted",
|
||||
PlatformKey: "unrestricted",
|
||||
PlatformModelID: "model-unrestricted",
|
||||
ModelName: "demo-image",
|
||||
ModelType: "image_generate",
|
||||
Capabilities: map[string]any{
|
||||
"image_generate": map[string]any{},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
filtered, _, err := filterRuntimeCandidatesByRequest("images.generations", "demo-image", "image_generate", map[string]any{
|
||||
"size": "2k",
|
||||
}, candidates)
|
||||
if err != nil {
|
||||
t.Fatalf("2k size should be accepted and normalized for matching: %v", err)
|
||||
}
|
||||
if len(filtered) != 1 || filtered[0].PlatformKey != "unrestricted" {
|
||||
t.Fatalf("unexpected candidates for normalized 2K size: %+v", filtered)
|
||||
}
|
||||
}
|
||||
|
||||
func candidateWithAspectRatios(platformKey string, modelType string, ratios ...string) store.RuntimeModelCandidate {
|
||||
return store.RuntimeModelCandidate{
|
||||
PlatformID: "platform-" + platformKey,
|
||||
PlatformKey: platformKey,
|
||||
PlatformModelID: "model-" + platformKey,
|
||||
ModelName: "demo-media",
|
||||
ModelType: modelType,
|
||||
Capabilities: map[string]any{
|
||||
modelType: map[string]any{
|
||||
"aspect_ratio_allowed": stringsToAny(ratios),
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func mediaKindForTest(modelType string) string {
|
||||
if modelType == "image_generate" {
|
||||
return "images.generations"
|
||||
}
|
||||
return "videos.generations"
|
||||
}
|
||||
@@ -10,7 +10,6 @@ import (
|
||||
|
||||
const (
|
||||
volcesOutputTokenThreshold = 10240
|
||||
volcesOutputCapabilityErrorCode = "model_capability_configuration_error"
|
||||
volcesOutputInvalidParameterCode = "invalid_parameter"
|
||||
)
|
||||
|
||||
@@ -25,20 +24,25 @@ func (outputTokenLimitProcessor) ShouldProcess(_ map[string]any, modelType strin
|
||||
func (outputTokenLimitProcessor) Process(params map[string]any, modelType string, context *paramProcessContext) bool {
|
||||
modelMax, sourceType, capabilityValue, ok := candidateMaxOutputTokens(context.candidate, modelType)
|
||||
path := capabilityPath(sourceType, "max_output_tokens")
|
||||
if parameter, value, explicit := explicitOutputTokenParameter(context.kind, params); explicit {
|
||||
if _, valid := positiveInteger(value); !valid {
|
||||
return context.reject(
|
||||
"OutputTokenLimitProcessor",
|
||||
parameter,
|
||||
value,
|
||||
fmt.Sprintf("%s must be a positive integer", parameter),
|
||||
path,
|
||||
capabilityValue,
|
||||
)
|
||||
}
|
||||
}
|
||||
if !ok {
|
||||
return context.reject(
|
||||
"OutputTokenLimitProcessor",
|
||||
outputTokenParameter(context.kind),
|
||||
nil,
|
||||
"火山引擎文本候选未配置正整数 max_output_tokens,已禁止执行该候选。",
|
||||
path,
|
||||
capabilityValue,
|
||||
)
|
||||
return true
|
||||
}
|
||||
|
||||
if context.kind == "chat.completions" {
|
||||
if value, explicit := nonNullParameter(params, "max_tokens"); explicit {
|
||||
if parsed, valid := positiveInteger(value); !valid || parsed > modelMax {
|
||||
if parsed, _ := positiveInteger(value); parsed > modelMax {
|
||||
return context.reject("OutputTokenLimitProcessor", "max_tokens", value, fmt.Sprintf("max_tokens must be a positive integer no greater than the selected Volcengine model limit (%d)", modelMax), path, capabilityValue)
|
||||
}
|
||||
return true
|
||||
@@ -57,7 +61,7 @@ func (outputTokenLimitProcessor) Process(params map[string]any, modelType string
|
||||
}
|
||||
|
||||
if value, explicit := nonNullParameter(params, "max_output_tokens"); explicit {
|
||||
if parsed, valid := positiveInteger(value); !valid || parsed > modelMax {
|
||||
if parsed, _ := positiveInteger(value); parsed > modelMax {
|
||||
return context.reject("OutputTokenLimitProcessor", "max_output_tokens", value, fmt.Sprintf("max_output_tokens must be a positive integer no greater than the selected Volcengine model limit (%d)", modelMax), path, capabilityValue)
|
||||
}
|
||||
return true
|
||||
@@ -76,8 +80,35 @@ func filterRuntimeCandidatesByOutputTokens(kind string, requestedModel string, m
|
||||
if !isOpenAITextGenerationKind(kind) || !isTextOutputModelType(modelType) || len(candidates) == 0 {
|
||||
return candidates, nil, nil
|
||||
}
|
||||
hasVolcesCandidate := false
|
||||
for _, candidate := range candidates {
|
||||
if isVolcesCandidate(candidate) {
|
||||
hasVolcesCandidate = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !hasVolcesCandidate {
|
||||
return candidates, nil, nil
|
||||
}
|
||||
explicitParameter, explicitValue, hasExplicitParameter := explicitOutputTokenParameter(kind, body)
|
||||
explicitTokens := 0
|
||||
if hasExplicitParameter {
|
||||
var valid bool
|
||||
explicitTokens, valid = positiveInteger(explicitValue)
|
||||
if !valid {
|
||||
summary := map[string]any{
|
||||
"filter": "volces_output_token_limit", "kind": kind, "requestedModel": requestedModel, "modelType": modelType,
|
||||
"candidateCount": len(candidates), "supportedCandidateCount": 0, "filteredCandidateCount": len(candidates),
|
||||
"code": volcesOutputInvalidParameterCode, "reason": "invalid_output_token_parameter",
|
||||
"parameter": explicitParameter, "requestedValue": explicitValue,
|
||||
}
|
||||
message := fmt.Sprintf("%s must be a positive integer", explicitParameter)
|
||||
return nil, summary, &store.ModelCandidateUnavailableError{Code: volcesOutputInvalidParameterCode, Message: message, Details: summary}
|
||||
}
|
||||
}
|
||||
filtered := make([]store.RuntimeModelCandidate, 0, len(candidates))
|
||||
rejected := make([]map[string]any, 0)
|
||||
missingCapability := make([]map[string]any, 0)
|
||||
invalidExplicit := false
|
||||
for _, candidate := range candidates {
|
||||
if !isVolcesCandidate(candidate) {
|
||||
@@ -92,15 +123,15 @@ func filterRuntimeCandidatesByOutputTokens(kind string, requestedModel string, m
|
||||
}
|
||||
if !configured {
|
||||
detail["reason"] = "max_output_tokens_missing"
|
||||
rejected = append(rejected, detail)
|
||||
missingCapability = append(missingCapability, detail)
|
||||
filtered = append(filtered, candidate)
|
||||
continue
|
||||
}
|
||||
if parameter, value, explicit := explicitOutputTokenParameter(kind, body); explicit {
|
||||
parsed, valid := positiveInteger(value)
|
||||
if !valid || (parameter != "max_completion_tokens" && parsed > modelMax) {
|
||||
if hasExplicitParameter {
|
||||
if explicitParameter != "max_completion_tokens" && explicitTokens > modelMax {
|
||||
detail["reason"] = "requested_output_tokens_exceed_capability"
|
||||
detail["parameter"] = parameter
|
||||
detail["requestedValue"] = value
|
||||
detail["parameter"] = explicitParameter
|
||||
detail["requestedValue"] = explicitValue
|
||||
invalidExplicit = true
|
||||
rejected = append(rejected, detail)
|
||||
continue
|
||||
@@ -108,26 +139,25 @@ func filterRuntimeCandidatesByOutputTokens(kind string, requestedModel string, m
|
||||
}
|
||||
filtered = append(filtered, candidate)
|
||||
}
|
||||
if len(rejected) == 0 {
|
||||
if len(rejected) == 0 && len(missingCapability) == 0 {
|
||||
return filtered, nil, nil
|
||||
}
|
||||
summary := map[string]any{
|
||||
"filter": "volces_output_token_limit", "kind": kind, "requestedModel": requestedModel, "modelType": modelType,
|
||||
"candidateCount": len(candidates), "supportedCandidateCount": len(filtered), "filteredCandidateCount": len(rejected),
|
||||
"missingCapabilityCandidateCount": len(missingCapability), "passthroughMissingCapability": len(missingCapability) > 0,
|
||||
"threshold": volcesOutputTokenThreshold, "formula": "third=floor(modelMaxOutputTokens/3); default=third>=10240?third:modelMaxOutputTokens",
|
||||
"rejectedCandidates": rejected,
|
||||
"rejectedCandidates": rejected, "missingCapabilityCandidates": missingCapability,
|
||||
}
|
||||
if len(filtered) > 0 {
|
||||
return filtered, summary, nil
|
||||
}
|
||||
code := volcesOutputCapabilityErrorCode
|
||||
message := "所有火山引擎文本候选都缺少有效的 max_output_tokens 能力配置"
|
||||
if invalidExplicit {
|
||||
code = volcesOutputInvalidParameterCode
|
||||
message = "请求的输出 token 上限超过所有可用火山引擎候选的模型能力"
|
||||
message := "请求的输出 token 上限超过所有可用火山引擎候选的模型能力"
|
||||
summary["code"] = volcesOutputInvalidParameterCode
|
||||
return nil, summary, &store.ModelCandidateUnavailableError{Code: volcesOutputInvalidParameterCode, Message: message, Details: summary}
|
||||
}
|
||||
summary["code"] = code
|
||||
return nil, summary, &store.ModelCandidateUnavailableError{Code: code, Message: message, Details: summary}
|
||||
return filtered, summary, nil
|
||||
}
|
||||
|
||||
func defaultVolcesOutputTokens(modelMax int) int {
|
||||
@@ -195,13 +225,6 @@ func explicitOutputTokenParameter(kind string, body map[string]any) (string, any
|
||||
return "max_completion_tokens", value, ok
|
||||
}
|
||||
|
||||
func outputTokenParameter(kind string) string {
|
||||
if kind == "responses" {
|
||||
return "max_output_tokens"
|
||||
}
|
||||
return "max_tokens"
|
||||
}
|
||||
|
||||
func isOpenAITextGenerationKind(kind string) bool {
|
||||
return kind == "chat.completions" || kind == "responses"
|
||||
}
|
||||
|
||||
@@ -72,32 +72,90 @@ func TestVolcesOutputProcessorPreservesExplicitLimits(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestVolcesOutputCandidateFilterSupportsFailover(t *testing.T) {
|
||||
func TestVolcesOutputCandidateFilterPreservesMissingCapabilityOrder(t *testing.T) {
|
||||
missing := volcTextCandidate(nil)
|
||||
missing.PlatformID = "missing"
|
||||
valid := volcTextCandidate(32768)
|
||||
valid.PlatformID = "valid"
|
||||
filtered, summary, err := filterRuntimeCandidatesByOutputTokens("chat.completions", "demo", "text_generate", map[string]any{"messages": []any{}}, []store.RuntimeModelCandidate{missing, valid})
|
||||
if err != nil || len(filtered) != 1 || filtered[0].PlatformID != "valid" {
|
||||
t.Fatalf("expected missing capability candidate to be skipped: filtered=%+v summary=%+v err=%v", filtered, summary, err)
|
||||
if err != nil || len(filtered) != 2 || filtered[0].PlatformID != "missing" || filtered[1].PlatformID != "valid" {
|
||||
t.Fatalf("expected missing capability candidate order to be preserved: filtered=%+v summary=%+v err=%v", filtered, summary, err)
|
||||
}
|
||||
result := preprocessRequestWithLog("chat.completions", map[string]any{"messages": []any{}}, filtered[0])
|
||||
if summary["missingCapabilityCandidateCount"] != 1 || summary["passthroughMissingCapability"] != true {
|
||||
t.Fatalf("expected missing capability audit summary, got %+v", summary)
|
||||
}
|
||||
missingResult := preprocessRequestWithLog("chat.completions", map[string]any{"messages": []any{}}, filtered[0])
|
||||
if missingResult.Err != nil || len(missingResult.Log.Changes) != 0 {
|
||||
t.Fatalf("missing capability candidate must pass through unchanged: %+v", missingResult)
|
||||
}
|
||||
if _, ok := missingResult.Body["max_tokens"]; ok {
|
||||
t.Fatalf("missing capability candidate must not inject max_tokens: %+v", missingResult.Body)
|
||||
}
|
||||
result := preprocessRequestWithLog("chat.completions", map[string]any{"messages": []any{}}, filtered[1])
|
||||
if result.Body["max_tokens"] != 10922 {
|
||||
t.Fatalf("failover candidate must recalculate its own default, got %+v", result.Body)
|
||||
t.Fatalf("configured candidate must calculate its own default, got %+v", result.Body)
|
||||
}
|
||||
}
|
||||
|
||||
func TestVolcesOutputCandidateFilterRejectsMissingAndExceededCapabilities(t *testing.T) {
|
||||
_, _, err := filterRuntimeCandidatesByOutputTokens("responses", "demo", "text_generate", map[string]any{"input": "hello"}, []store.RuntimeModelCandidate{volcTextCandidate(nil)})
|
||||
if store.ModelCandidateErrorCode(err) != volcesOutputCapabilityErrorCode {
|
||||
t.Fatalf("expected capability configuration error, got %v", err)
|
||||
func TestVolcesOutputCandidateFilterPassesThroughMissingCapability(t *testing.T) {
|
||||
filtered, summary, err := filterRuntimeCandidatesByOutputTokens("responses", "demo", "text_generate", map[string]any{"input": "hello"}, []store.RuntimeModelCandidate{volcTextCandidate(nil)})
|
||||
if err != nil || len(filtered) != 1 {
|
||||
t.Fatalf("expected missing capability candidate to remain available: filtered=%+v summary=%+v err=%v", filtered, summary, err)
|
||||
}
|
||||
_, _, err = filterRuntimeCandidatesByOutputTokens("chat.completions", "demo", "text_generate", map[string]any{"messages": []any{}, "max_tokens": 32769}, []store.RuntimeModelCandidate{volcTextCandidate(32768)})
|
||||
result := preprocessRequestWithLog("responses", map[string]any{"input": "hello"}, filtered[0])
|
||||
if result.Err != nil {
|
||||
t.Fatalf("missing capability preprocessing must pass through: %v", result.Err)
|
||||
}
|
||||
if _, ok := result.Body["max_output_tokens"]; ok {
|
||||
t.Fatalf("missing capability preprocessing must not inject max_output_tokens: %+v", result.Body)
|
||||
}
|
||||
}
|
||||
|
||||
func TestVolcesOutputCandidateFilterRejectsExceededKnownCapability(t *testing.T) {
|
||||
_, _, err := filterRuntimeCandidatesByOutputTokens("chat.completions", "demo", "text_generate", map[string]any{"messages": []any{}, "max_tokens": 32769}, []store.RuntimeModelCandidate{volcTextCandidate(32768)})
|
||||
if store.ModelCandidateErrorCode(err) != volcesOutputInvalidParameterCode {
|
||||
t.Fatalf("expected invalid_parameter for explicit limit, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestVolcesOutputCandidateFilterAllowsUnknownCapabilityForExplicitLimit(t *testing.T) {
|
||||
missing := volcTextCandidate(nil)
|
||||
missing.PlatformID = "missing"
|
||||
configured := volcTextCandidate(32768)
|
||||
configured.PlatformID = "configured"
|
||||
filtered, summary, err := filterRuntimeCandidatesByOutputTokens(
|
||||
"chat.completions",
|
||||
"demo",
|
||||
"text_generate",
|
||||
map[string]any{"messages": []any{}, "max_tokens": 65536},
|
||||
[]store.RuntimeModelCandidate{missing, configured},
|
||||
)
|
||||
if err != nil || len(filtered) != 1 || filtered[0].PlatformID != "missing" {
|
||||
t.Fatalf("expected unknown capability candidate to remain available: filtered=%+v summary=%+v err=%v", filtered, summary, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestVolcesOutputProcessorValidatesExplicitLimitWithoutCapability(t *testing.T) {
|
||||
valid := preprocessRequestWithLog("chat.completions", map[string]any{"messages": []any{}, "max_tokens": 2048}, volcTextCandidate(nil))
|
||||
if valid.Err != nil || valid.Body["max_tokens"] != 2048 {
|
||||
t.Fatalf("valid explicit limit must pass through: %+v", valid)
|
||||
}
|
||||
invalid := preprocessRequestWithLog("chat.completions", map[string]any{"messages": []any{}, "max_tokens": 0}, volcTextCandidate(nil))
|
||||
if invalid.Err == nil {
|
||||
t.Fatalf("invalid explicit limit must still be rejected: %+v", invalid)
|
||||
}
|
||||
_, _, err := filterRuntimeCandidatesByOutputTokens(
|
||||
"chat.completions",
|
||||
"demo",
|
||||
"text_generate",
|
||||
map[string]any{"messages": []any{}, "max_tokens": 0},
|
||||
[]store.RuntimeModelCandidate{volcTextCandidate(nil)},
|
||||
)
|
||||
if store.ModelCandidateErrorCode(err) != volcesOutputInvalidParameterCode {
|
||||
t.Fatalf("candidate filter must reject an invalid explicit limit, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNonVolcesOutputProcessorDoesNotInject(t *testing.T) {
|
||||
candidate := volcTextCandidate(131072)
|
||||
candidate.Provider = "openai"
|
||||
|
||||
@@ -3,6 +3,7 @@ package runner
|
||||
import (
|
||||
"fmt"
|
||||
"math"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
@@ -13,28 +14,43 @@ type resolutionNormalizeProcessor struct{}
|
||||
func (resolutionNormalizeProcessor) Name() string { return "ResolutionNormalizeProcessor" }
|
||||
|
||||
func (resolutionNormalizeProcessor) ShouldProcess(params map[string]any, modelType string, context *paramProcessContext) bool {
|
||||
if stringFromAny(params["resolution"]) != "" {
|
||||
return false
|
||||
if normalizeMediaResolution(modelType, stringFromAny(params["resolution"])) != "" {
|
||||
return true
|
||||
}
|
||||
size := stringFromAny(params["size"])
|
||||
if size == "" {
|
||||
return false
|
||||
}
|
||||
return isImageResolution(modelType, size) || isVideoResolution(modelType, size)
|
||||
return normalizeMediaResolution(modelType, stringFromAny(params["size"])) != ""
|
||||
}
|
||||
|
||||
func (resolutionNormalizeProcessor) Process(params map[string]any, modelType string, context *paramProcessContext) bool {
|
||||
if resolution := normalizeMediaResolution(modelType, stringFromAny(params["resolution"])); resolution != "" {
|
||||
before := params["resolution"]
|
||||
params["resolution"] = resolution
|
||||
context.resolution = resolution
|
||||
if stringFromAny(before) != resolution {
|
||||
_, capabilityValue := capabilityEvidence(context.modelCapability, modelType, "output_resolutions")
|
||||
context.recordChange(
|
||||
"ResolutionNormalizeProcessor",
|
||||
"adjust",
|
||||
"resolution",
|
||||
before,
|
||||
resolution,
|
||||
"resolution 使用了可识别的媒体分辨率格式,已归一为标准大小写和值。",
|
||||
capabilityPath(modelType, "output_resolutions"),
|
||||
capabilityValue,
|
||||
)
|
||||
}
|
||||
return true
|
||||
}
|
||||
size := stringFromAny(params["size"])
|
||||
if stringFromAny(params["resolution"]) == "" && (isImageResolution(modelType, size) || isVideoResolution(modelType, size)) {
|
||||
if resolution := normalizeMediaResolution(modelType, size); resolution != "" {
|
||||
_, capabilityValue := capabilityEvidence(context.modelCapability, modelType, "output_resolutions")
|
||||
params["resolution"] = size
|
||||
context.resolution = size
|
||||
params["resolution"] = resolution
|
||||
context.resolution = resolution
|
||||
context.recordChange(
|
||||
"ResolutionNormalizeProcessor",
|
||||
"set",
|
||||
"resolution",
|
||||
nil,
|
||||
size,
|
||||
resolution,
|
||||
"size 使用分辨率格式,归一到 resolution 供后续能力校验和计费使用。",
|
||||
capabilityPath(modelType, "output_resolutions"),
|
||||
capabilityValue,
|
||||
@@ -48,30 +64,69 @@ type aspectRatioProcessor struct{}
|
||||
func (aspectRatioProcessor) Name() string { return "AspectRatioProcessor" }
|
||||
|
||||
func (aspectRatioProcessor) ShouldProcess(params map[string]any, modelType string, context *paramProcessContext) bool {
|
||||
return modelType != "text_generate" && (stringFromAny(params["aspect_ratio"]) != "" || stringFromAny(params["size"]) != "")
|
||||
if modelType == "text_generate" {
|
||||
return false
|
||||
}
|
||||
for _, key := range []string{"aspect_ratio", "aspectRatio", "ratio"} {
|
||||
if _, ok := params[key]; ok {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return stringFromAny(params["size"]) != ""
|
||||
}
|
||||
|
||||
func (aspectRatioProcessor) Process(params map[string]any, modelType string, context *paramProcessContext) bool {
|
||||
capability := capabilityForType(context.modelCapability, modelType)
|
||||
if capability == nil {
|
||||
aspectRatio, ratioSource, ratioPresent := requestAspectRatio(params)
|
||||
if ratioPresent {
|
||||
originalAspectRatio := aspectRatio
|
||||
normalized, valid := normalizedAspectRatio(aspectRatio)
|
||||
if !valid {
|
||||
before := map[string]any{}
|
||||
for _, key := range []string{"aspect_ratio", "aspectRatio", "ratio"} {
|
||||
if value, ok := params[key]; ok {
|
||||
before[key] = cloneAny(value)
|
||||
delete(params, key)
|
||||
}
|
||||
}
|
||||
aspectRatio = ""
|
||||
context.aspectRatio = ""
|
||||
context.recordChange(
|
||||
"AspectRatioProcessor",
|
||||
"remove",
|
||||
ratioSource,
|
||||
before,
|
||||
nil,
|
||||
"比例不是有效的正数 x:x 格式,按未指定比例处理。",
|
||||
"",
|
||||
nil,
|
||||
)
|
||||
return true
|
||||
} else {
|
||||
aspectRatio = normalized
|
||||
params["aspect_ratio"] = normalized
|
||||
for _, key := range []string{"aspectRatio", "ratio"} {
|
||||
delete(params, key)
|
||||
}
|
||||
if ratioSource != "aspect_ratio" || normalized != originalAspectRatio {
|
||||
context.recordChange(
|
||||
"AspectRatioProcessor",
|
||||
"adjust",
|
||||
"aspect_ratio",
|
||||
map[string]any{ratioSource: originalAspectRatio},
|
||||
normalized,
|
||||
"比例已归一为标准 x:x 格式。",
|
||||
"",
|
||||
nil,
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
if !ratioPresent {
|
||||
return true
|
||||
}
|
||||
|
||||
aspectRatio := stringFromAny(params["aspect_ratio"])
|
||||
if isEmptyParamString(aspectRatio) {
|
||||
before := params["aspect_ratio"]
|
||||
delete(params, "aspect_ratio")
|
||||
context.aspectRatio = ""
|
||||
context.recordChange(
|
||||
"AspectRatioProcessor",
|
||||
"remove",
|
||||
"aspect_ratio",
|
||||
before,
|
||||
nil,
|
||||
"aspect_ratio 是空值字符串,不能作为有效比例传给上游。",
|
||||
"",
|
||||
nil,
|
||||
)
|
||||
capability := capabilityForType(context.modelCapability, modelType)
|
||||
if capability == nil {
|
||||
return true
|
||||
}
|
||||
|
||||
@@ -174,6 +229,15 @@ func (aspectRatioProcessor) Process(params map[string]any, modelType string, con
|
||||
return true
|
||||
}
|
||||
|
||||
func requestAspectRatio(params map[string]any) (string, string, bool) {
|
||||
for _, key := range []string{"aspect_ratio", "aspectRatio", "ratio"} {
|
||||
if value, ok := params[key]; ok {
|
||||
return stringFromAny(value), key, true
|
||||
}
|
||||
}
|
||||
return "", "aspect_ratio", false
|
||||
}
|
||||
|
||||
func validateExplicitAspectRatio(aspectRatio string, capability map[string]any, allowed []string, modelType string) (string, bool) {
|
||||
if isVideoModelType(modelType) && len(allowed) > 0 && !containsString(allowed, aspectRatio) {
|
||||
if aspectRatio == "adaptive" || aspectRatio == "keep_ratio" {
|
||||
@@ -194,17 +258,14 @@ func (imageSizeProcessor) ShouldProcess(params map[string]any, modelType string,
|
||||
if modelType != "image_generate" && modelType != "image_edit" {
|
||||
return false
|
||||
}
|
||||
if _, _, ok := imageDimensionsFromParams(params); !ok {
|
||||
return false
|
||||
}
|
||||
capability := capabilityForType(context.modelCapability, modelType)
|
||||
return capability != nil && imageSizeCapabilityConfigured(capability)
|
||||
_, _, ok := imageDimensionsFromParams(params)
|
||||
return ok
|
||||
}
|
||||
|
||||
func (imageSizeProcessor) Process(params map[string]any, modelType string, context *paramProcessContext) bool {
|
||||
capability := capabilityForType(context.modelCapability, modelType)
|
||||
if capability == nil {
|
||||
return true
|
||||
capability = map[string]any{}
|
||||
}
|
||||
width, height, ok := imageDimensionsFromParams(params)
|
||||
if !ok {
|
||||
@@ -222,7 +283,11 @@ func (imageSizeProcessor) Process(params map[string]any, modelType string, conte
|
||||
params["width"] = width
|
||||
params["height"] = height
|
||||
_, _, sizeHasPixelDimensions := parsePixelSizeString(stringFromAny(params["size"]))
|
||||
if stringFromAny(params["aspect_ratio"]) == "" || sizeHasPixelDimensions {
|
||||
_, _, resolutionHasPixelDimensions := parsePixelSizeString(stringFromAny(params["resolution"]))
|
||||
explicitWidthHeight := positiveIntegerFromAny(params["width"]) > 0 && positiveIntegerFromAny(params["height"]) > 0 &&
|
||||
normalizeImageResolution(stringFromAny(params["size"])) == "" &&
|
||||
normalizeImageResolution(stringFromAny(params["resolution"])) == ""
|
||||
if sizeHasPixelDimensions || resolutionHasPixelDimensions || (stringFromAny(params["aspect_ratio"]) == "" && explicitWidthHeight) {
|
||||
aspectRatio := aspectRatioFromDimensions(width, height)
|
||||
allowed := aspectRatioAllowed(capability["aspect_ratio_allowed"], firstNonEmptyString(stringFromAny(params["resolution"]), context.resolution))
|
||||
if processed, ok := validateAndAdjustAspectRatio(aspectRatio, capability, allowed); ok && processed != "" {
|
||||
@@ -274,7 +339,37 @@ func imageDimensionsFromParams(params map[string]any) (int, int, bool) {
|
||||
if width > 0 && height > 0 {
|
||||
return width, height, true
|
||||
}
|
||||
return parsePixelSizeString(stringFromAny(params["resolution"]))
|
||||
if width, height, ok := parsePixelSizeString(stringFromAny(params["resolution"])); ok {
|
||||
return width, height, true
|
||||
}
|
||||
resolution := firstNonEmptyString(
|
||||
normalizeImageResolution(stringFromAny(params["resolution"])),
|
||||
normalizeImageResolution(stringFromAny(params["size"])),
|
||||
)
|
||||
if resolution == "" {
|
||||
return 0, 0, false
|
||||
}
|
||||
return imageDimensionsForResolution(resolution, stringFromAny(params["aspect_ratio"]))
|
||||
}
|
||||
|
||||
func imageDimensionsForResolution(resolution string, aspectRatio string) (int, int, bool) {
|
||||
resolution = normalizeImageResolution(resolution)
|
||||
if resolution == "" {
|
||||
return 0, 0, false
|
||||
}
|
||||
magnitude, err := strconv.Atoi(strings.TrimSuffix(resolution, "K"))
|
||||
if err != nil || magnitude <= 0 {
|
||||
return 0, 0, false
|
||||
}
|
||||
baseSide := magnitude * 1024
|
||||
ratio, validRatio := aspectRatioNumber(aspectRatio)
|
||||
if !validRatio {
|
||||
return baseSide, baseSide, true
|
||||
}
|
||||
if ratio >= 1 {
|
||||
return baseSide, max(1, int(math.Round(float64(baseSide)/ratio))), true
|
||||
}
|
||||
return max(1, int(math.Round(float64(baseSide)*ratio))), baseSide, true
|
||||
}
|
||||
|
||||
func imageSizeCapabilityConfigured(capability map[string]any) bool {
|
||||
|
||||
@@ -393,23 +393,39 @@ func numberPair(value any) ([2]float64, bool) {
|
||||
}
|
||||
|
||||
func validAspectRatio(value string) bool {
|
||||
if value == "adaptive" || value == "keep_ratio" {
|
||||
return true
|
||||
}
|
||||
_, ok := aspectRatioNumber(value)
|
||||
_, ok := normalizedAspectRatio(value)
|
||||
return ok
|
||||
}
|
||||
|
||||
func aspectRatioNumber(value string) (float64, bool) {
|
||||
parts := strings.Split(value, ":")
|
||||
func normalizedAspectRatio(value string) (string, bool) {
|
||||
compact := strings.Map(func(r rune) rune {
|
||||
switch r {
|
||||
case ' ', '\t', '\r', '\n':
|
||||
return -1
|
||||
default:
|
||||
return r
|
||||
}
|
||||
}, strings.TrimSpace(value))
|
||||
parts := strings.Split(compact, ":")
|
||||
if len(parts) != 2 {
|
||||
return 0, false
|
||||
return "", false
|
||||
}
|
||||
width := parsePositiveFloat(parts[0])
|
||||
height := parsePositiveFloat(parts[1])
|
||||
if width <= 0 || height <= 0 {
|
||||
return "", false
|
||||
}
|
||||
return parts[0] + ":" + parts[1], true
|
||||
}
|
||||
|
||||
func aspectRatioNumber(value string) (float64, bool) {
|
||||
normalized, ok := normalizedAspectRatio(value)
|
||||
if !ok {
|
||||
return 0, false
|
||||
}
|
||||
parts := strings.Split(normalized, ":")
|
||||
width := parsePositiveFloat(parts[0])
|
||||
height := parsePositiveFloat(parts[1])
|
||||
return width / height, true
|
||||
}
|
||||
|
||||
@@ -537,11 +553,53 @@ func isEmptyParamString(value string) bool {
|
||||
}
|
||||
|
||||
func isImageResolution(modelType string, value string) bool {
|
||||
return (modelType == "image_generate" || modelType == "image_edit") && containsString([]string{"1K", "2K", "3K", "4K", "8K"}, value)
|
||||
if modelType != "image_generate" && modelType != "image_edit" {
|
||||
return false
|
||||
}
|
||||
return normalizeImageResolution(value) != ""
|
||||
}
|
||||
|
||||
func isVideoResolution(modelType string, value string) bool {
|
||||
return isVideoModelType(modelType) && containsString([]string{"480p", "720p", "1080p", "1440p", "2160p"}, value)
|
||||
return isVideoModelType(modelType) && normalizeVideoResolution(value) != ""
|
||||
}
|
||||
|
||||
func normalizeMediaResolution(modelType string, value string) string {
|
||||
if modelType == "image_generate" || modelType == "image_edit" {
|
||||
return normalizeImageResolution(value)
|
||||
}
|
||||
if isVideoModelType(modelType) {
|
||||
return normalizeVideoResolution(value)
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func normalizeImageResolution(value string) string {
|
||||
value = strings.ToLower(strings.TrimSpace(value))
|
||||
if len(value) < 2 || value[len(value)-1] != 'k' {
|
||||
return ""
|
||||
}
|
||||
magnitude, err := strconv.Atoi(value[:len(value)-1])
|
||||
if err != nil || magnitude <= 0 {
|
||||
return ""
|
||||
}
|
||||
return strconv.Itoa(magnitude) + "K"
|
||||
}
|
||||
|
||||
func normalizeVideoResolution(value string) string {
|
||||
switch strings.ToLower(strings.TrimSpace(value)) {
|
||||
case "480p":
|
||||
return "480p"
|
||||
case "720p":
|
||||
return "720p"
|
||||
case "1080p":
|
||||
return "1080p"
|
||||
case "1440p":
|
||||
return "1440p"
|
||||
case "2160p", "4k":
|
||||
return "2160p"
|
||||
default:
|
||||
return ""
|
||||
}
|
||||
}
|
||||
|
||||
func isVideoModelType(modelType string) bool {
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -11,6 +11,7 @@ import (
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
"github.com/google/uuid"
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/riverqueue/river"
|
||||
"github.com/riverqueue/river/riverdriver/riverpgxv5"
|
||||
"github.com/riverqueue/river/rivermigrate"
|
||||
@@ -23,6 +24,12 @@ type asyncTaskArgs struct {
|
||||
TaskID string `json:"task_id" river:"unique"`
|
||||
}
|
||||
|
||||
type asyncExecutionClient interface {
|
||||
Start(context.Context) error
|
||||
Stop(context.Context) error
|
||||
StopAndCancel(context.Context) error
|
||||
}
|
||||
|
||||
func (asyncTaskArgs) Kind() string { return "gateway_task_run" }
|
||||
|
||||
type asyncTaskWorker struct {
|
||||
@@ -87,17 +94,70 @@ func (s *Service) startRiverQueue(ctx context.Context) error {
|
||||
return err
|
||||
}
|
||||
|
||||
workers := river.NewWorkers()
|
||||
if err := river.AddWorkerSafely(workers, &asyncTaskWorker{service: s}); err != nil {
|
||||
controlClient, err := river.NewClient(driver, &river.Config{
|
||||
ID: asyncWorkerID() + "-control",
|
||||
Logger: s.logger,
|
||||
TestOnly: s.cfg.AppEnv == "test",
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
riverClient, err := river.NewClient(driver, &river.Config{
|
||||
ID: asyncWorkerID(),
|
||||
snapshot, err := s.loadAsyncWorkerCapacity(ctx)
|
||||
if err != nil {
|
||||
return fmt.Errorf("calculate initial async worker capacity: %w", err)
|
||||
}
|
||||
executionClient, err := s.makeAsyncExecutionClient(snapshot.Capacity)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := executionClient.Start(ctx); err != nil {
|
||||
cleanupCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
_ = executionClient.StopAndCancel(cleanupCtx)
|
||||
cancel()
|
||||
return err
|
||||
}
|
||||
s.riverMu.Lock()
|
||||
s.riverControlClient = controlClient
|
||||
s.riverExecutionClient = executionClient
|
||||
s.riverWorkerCapacity = snapshot.Capacity
|
||||
s.riverDrainingClients = make(map[asyncExecutionClient]struct{})
|
||||
s.riverMu.Unlock()
|
||||
s.observeAsyncWorkerCapacity(snapshot)
|
||||
s.logger.Info("async worker capacity initialized",
|
||||
"capacity", snapshot.Capacity,
|
||||
"desiredCapacity", snapshot.Desired,
|
||||
"hardLimit", snapshot.HardLimit,
|
||||
"enabledModels", snapshot.EnabledModels,
|
||||
"unlimitedModels", snapshot.UnlimitedModels,
|
||||
"modelDesired", snapshot.ModelDesired,
|
||||
"enabledGroups", snapshot.EnabledGroups,
|
||||
"unlimitedGroups", snapshot.UnlimitedGroups,
|
||||
"groupDesired", snapshot.GroupDesired,
|
||||
"capped", snapshot.Capped,
|
||||
)
|
||||
if err := s.recoverAsyncRiverJobs(ctx); err != nil {
|
||||
cleanupCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
_ = executionClient.StopAndCancel(cleanupCtx)
|
||||
cancel()
|
||||
return err
|
||||
}
|
||||
go s.refreshAsyncWorkerCapacity(ctx)
|
||||
go s.stopAsyncWorkersOnShutdown(ctx)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Service) newRiverAsyncExecutionClient(capacity int) (*river.Client[pgx.Tx], error) {
|
||||
workers := river.NewWorkers()
|
||||
if err := river.AddWorkerSafely(workers, &asyncTaskWorker{service: s}); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return river.NewClient(riverpgxv5.New(s.store.Pool()), &river.Config{
|
||||
ID: fmt.Sprintf("%s-exec-%d-%d", asyncWorkerID(), capacity, time.Now().UnixNano()),
|
||||
JobTimeout: -1,
|
||||
Logger: s.logger,
|
||||
CompletedJobRetentionPeriod: 24 * time.Hour,
|
||||
Queues: map[string]river.QueueConfig{
|
||||
asyncTaskQueueName: {MaxWorkers: 32},
|
||||
asyncTaskQueueName: {MaxWorkers: capacity},
|
||||
},
|
||||
// Provider-backed media jobs commonly poll for 10-20 minutes. River may
|
||||
// execute a still-running job again once this window elapses, so keep the
|
||||
@@ -106,32 +166,160 @@ func (s *Service) startRiverQueue(ctx context.Context) error {
|
||||
TestOnly: s.cfg.AppEnv == "test",
|
||||
Workers: workers,
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Service) makeAsyncExecutionClient(capacity int) (asyncExecutionClient, error) {
|
||||
if s.asyncClientFactory != nil {
|
||||
return s.asyncClientFactory(capacity)
|
||||
}
|
||||
return s.newRiverAsyncExecutionClient(capacity)
|
||||
}
|
||||
|
||||
func (s *Service) loadAsyncWorkerCapacity(ctx context.Context) (store.AsyncWorkerCapacitySnapshot, error) {
|
||||
if s.asyncCapacityLoader != nil {
|
||||
return s.asyncCapacityLoader(ctx, s.cfg.AsyncWorkerHardLimit)
|
||||
}
|
||||
return s.store.AsyncWorkerCapacity(ctx, s.cfg.AsyncWorkerHardLimit)
|
||||
}
|
||||
|
||||
func (s *Service) refreshAsyncWorkerCapacity(ctx context.Context) {
|
||||
ticker := time.NewTicker(time.Duration(s.cfg.AsyncWorkerRefreshIntervalSeconds) * time.Second)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
s.resizeAsyncWorkerCapacity(ctx)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Service) resizeAsyncWorkerCapacity(ctx context.Context) {
|
||||
snapshot, err := s.loadAsyncWorkerCapacity(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
s.observeAsyncWorkerResize("refresh_failed")
|
||||
s.logger.Warn("refresh async worker capacity failed; keeping current client", "error", err)
|
||||
return
|
||||
}
|
||||
s.riverClient = riverClient
|
||||
if err := riverClient.Start(ctx); err != nil {
|
||||
return err
|
||||
s.riverMu.RLock()
|
||||
currentCapacity := s.riverWorkerCapacity
|
||||
s.riverMu.RUnlock()
|
||||
s.observeAsyncWorkerCapacity(snapshot)
|
||||
if snapshot.Capacity == currentCapacity {
|
||||
return
|
||||
}
|
||||
if err := s.recoverAsyncRiverJobs(ctx); err != nil {
|
||||
return err
|
||||
newClient, err := s.makeAsyncExecutionClient(snapshot.Capacity)
|
||||
if err != nil {
|
||||
s.observeAsyncWorkerResize("create_failed")
|
||||
s.logger.Warn("create replacement async worker client failed; keeping current client",
|
||||
"error", err, "currentCapacity", currentCapacity, "desiredCapacity", snapshot.Capacity)
|
||||
return
|
||||
}
|
||||
go func() {
|
||||
<-ctx.Done()
|
||||
if err := newClient.Start(ctx); err != nil {
|
||||
cleanupCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
_ = newClient.StopAndCancel(cleanupCtx)
|
||||
cancel()
|
||||
s.observeAsyncWorkerResize("start_failed")
|
||||
s.logger.Warn("start replacement async worker client failed; keeping current client",
|
||||
"error", err, "currentCapacity", currentCapacity, "desiredCapacity", snapshot.Capacity)
|
||||
return
|
||||
}
|
||||
if ctx.Err() != nil {
|
||||
cleanupCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
_ = newClient.StopAndCancel(cleanupCtx)
|
||||
cancel()
|
||||
return
|
||||
}
|
||||
s.riverMu.Lock()
|
||||
oldClient := s.riverExecutionClient
|
||||
s.riverExecutionClient = newClient
|
||||
s.riverWorkerCapacity = snapshot.Capacity
|
||||
if oldClient != nil {
|
||||
if s.riverDrainingClients == nil {
|
||||
s.riverDrainingClients = make(map[asyncExecutionClient]struct{})
|
||||
}
|
||||
s.riverDrainingClients[oldClient] = struct{}{}
|
||||
}
|
||||
s.riverMu.Unlock()
|
||||
s.observeAsyncWorkerResize("success")
|
||||
s.logger.Info("async worker capacity resized",
|
||||
"previousCapacity", currentCapacity,
|
||||
"capacity", snapshot.Capacity,
|
||||
"desiredCapacity", snapshot.Desired,
|
||||
"hardLimit", snapshot.HardLimit,
|
||||
"modelDesired", snapshot.ModelDesired,
|
||||
"groupDesired", snapshot.GroupDesired,
|
||||
"capped", snapshot.Capped,
|
||||
)
|
||||
if oldClient != nil {
|
||||
go s.drainAsyncWorkerClient(oldClient)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Service) drainAsyncWorkerClient(client asyncExecutionClient) {
|
||||
stopCtx, cancel := context.WithTimeout(context.Background(), time.Hour)
|
||||
defer cancel()
|
||||
if err := client.Stop(stopCtx); err != nil {
|
||||
s.logger.Warn("gracefully drain previous async worker client failed", "error", err)
|
||||
}
|
||||
s.riverMu.Lock()
|
||||
delete(s.riverDrainingClients, client)
|
||||
s.riverMu.Unlock()
|
||||
}
|
||||
|
||||
func (s *Service) stopAsyncWorkersOnShutdown(ctx context.Context) {
|
||||
<-ctx.Done()
|
||||
s.riverMu.Lock()
|
||||
clients := make([]asyncExecutionClient, 0, 1+len(s.riverDrainingClients))
|
||||
if s.riverExecutionClient != nil {
|
||||
clients = append(clients, s.riverExecutionClient)
|
||||
}
|
||||
for client := range s.riverDrainingClients {
|
||||
clients = append(clients, client)
|
||||
}
|
||||
s.riverExecutionClient = nil
|
||||
s.riverDrainingClients = nil
|
||||
s.riverMu.Unlock()
|
||||
for _, client := range clients {
|
||||
stopCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
if err := riverClient.StopAndCancel(stopCtx); err != nil {
|
||||
if err := client.StopAndCancel(stopCtx); err != nil {
|
||||
s.logger.Warn("stop river async queue failed", "error", err)
|
||||
}
|
||||
}()
|
||||
return nil
|
||||
cancel()
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Service) asyncControlClient() *river.Client[pgx.Tx] {
|
||||
s.riverMu.RLock()
|
||||
defer s.riverMu.RUnlock()
|
||||
return s.riverControlClient
|
||||
}
|
||||
|
||||
func (s *Service) observeAsyncWorkerCapacity(snapshot store.AsyncWorkerCapacitySnapshot) {
|
||||
observer, ok := s.billingMetrics.(interface {
|
||||
SetAsyncWorkerCapacity(current, desired, hardLimit int, capped bool)
|
||||
})
|
||||
if ok {
|
||||
observer.SetAsyncWorkerCapacity(snapshot.Capacity, snapshot.Desired, snapshot.HardLimit, snapshot.Capped)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Service) observeAsyncWorkerResize(outcome string) {
|
||||
observer, ok := s.billingMetrics.(interface {
|
||||
ObserveAsyncWorkerResize(string)
|
||||
})
|
||||
if ok {
|
||||
observer.ObserveAsyncWorkerResize(outcome)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Service) EnqueueAsyncTask(ctx context.Context, task store.GatewayTask) error {
|
||||
if s.riverClient == nil {
|
||||
riverClient := s.asyncControlClient()
|
||||
if riverClient == nil {
|
||||
return errors.New("river async queue is not started")
|
||||
}
|
||||
result, err := s.riverClient.Insert(ctx, asyncTaskArgs{TaskID: task.ID}, asyncTaskInsertOpts(task))
|
||||
result, err := riverClient.Insert(ctx, asyncTaskArgs{TaskID: task.ID}, asyncTaskInsertOpts(task))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
@@ -151,12 +339,16 @@ func (s *Service) RunAsyncTask(ctx context.Context, task store.GatewayTask, user
|
||||
}
|
||||
|
||||
func (s *Service) recoverAsyncRiverJobs(ctx context.Context) error {
|
||||
riverClient := s.asyncControlClient()
|
||||
if riverClient == nil {
|
||||
return errors.New("river async queue is not started")
|
||||
}
|
||||
items, err := s.store.ListRecoverableAsyncTasks(ctx, 1000)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
for _, item := range items {
|
||||
result, err := s.riverClient.Insert(ctx, asyncTaskArgs{TaskID: item.ID}, asyncTaskRecoveryInsertOpts(item, time.Now()))
|
||||
result, err := riverClient.Insert(ctx, asyncTaskArgs{TaskID: item.ID}, asyncTaskRecoveryInsertOpts(item, time.Now()))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
@@ -0,0 +1,189 @@
|
||||
package runner
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"log/slog"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
)
|
||||
|
||||
type fakeAsyncExecutionClient struct {
|
||||
startErr error
|
||||
stopGate <-chan struct{}
|
||||
started atomic.Int64
|
||||
stopped atomic.Int64
|
||||
stopCancelled atomic.Int64
|
||||
}
|
||||
|
||||
func (c *fakeAsyncExecutionClient) Start(context.Context) error {
|
||||
c.started.Add(1)
|
||||
return c.startErr
|
||||
}
|
||||
|
||||
func (c *fakeAsyncExecutionClient) Stop(ctx context.Context) error {
|
||||
c.stopped.Add(1)
|
||||
if c.stopGate == nil {
|
||||
return nil
|
||||
}
|
||||
select {
|
||||
case <-c.stopGate:
|
||||
return nil
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
}
|
||||
}
|
||||
|
||||
func (c *fakeAsyncExecutionClient) StopAndCancel(context.Context) error {
|
||||
c.stopCancelled.Add(1)
|
||||
return nil
|
||||
}
|
||||
|
||||
func TestResizeAsyncWorkerCapacityStartsReplacementBeforeGracefulDrain(t *testing.T) {
|
||||
oldClient := &fakeAsyncExecutionClient{}
|
||||
drainGate := make(chan struct{})
|
||||
oldClient.stopGate = drainGate
|
||||
newClient := &fakeAsyncExecutionClient{}
|
||||
service := asyncWorkerManagerTestService(1)
|
||||
service.riverExecutionClient = oldClient
|
||||
service.asyncCapacityLoader = fixedAsyncCapacity(3)
|
||||
service.asyncClientFactory = func(capacity int) (asyncExecutionClient, error) {
|
||||
if capacity != 3 {
|
||||
t.Fatalf("factory capacity=%d, want=3", capacity)
|
||||
}
|
||||
return newClient, nil
|
||||
}
|
||||
|
||||
service.resizeAsyncWorkerCapacity(context.Background())
|
||||
if newClient.started.Load() != 1 {
|
||||
t.Fatalf("replacement start count=%d, want=1", newClient.started.Load())
|
||||
}
|
||||
if service.riverExecutionClient != newClient || service.riverWorkerCapacity != 3 {
|
||||
t.Fatalf("replacement was not installed: capacity=%d client=%T", service.riverWorkerCapacity, service.riverExecutionClient)
|
||||
}
|
||||
waitForAtomicValue(t, &oldClient.stopped, 1)
|
||||
if oldClient.stopCancelled.Load() != 0 {
|
||||
t.Fatal("graceful drain cancelled an already running old task")
|
||||
}
|
||||
close(drainGate)
|
||||
waitForDrainingClients(t, service, 0)
|
||||
}
|
||||
|
||||
func TestResizeAsyncWorkerCapacityFailuresKeepCurrentClient(t *testing.T) {
|
||||
t.Run("refresh failure", func(t *testing.T) {
|
||||
oldClient := &fakeAsyncExecutionClient{}
|
||||
service := asyncWorkerManagerTestService(4)
|
||||
service.riverExecutionClient = oldClient
|
||||
service.asyncCapacityLoader = func(context.Context, int) (store.AsyncWorkerCapacitySnapshot, error) {
|
||||
return store.AsyncWorkerCapacitySnapshot{}, errors.New("database unavailable")
|
||||
}
|
||||
service.asyncClientFactory = func(int) (asyncExecutionClient, error) {
|
||||
t.Fatal("factory must not run after refresh failure")
|
||||
return nil, nil
|
||||
}
|
||||
service.resizeAsyncWorkerCapacity(context.Background())
|
||||
if service.riverExecutionClient != oldClient || service.riverWorkerCapacity != 4 {
|
||||
t.Fatal("refresh failure replaced the current client")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("create failure", func(t *testing.T) {
|
||||
oldClient := &fakeAsyncExecutionClient{}
|
||||
service := asyncWorkerManagerTestService(4)
|
||||
service.riverExecutionClient = oldClient
|
||||
service.asyncCapacityLoader = fixedAsyncCapacity(8)
|
||||
service.asyncClientFactory = func(int) (asyncExecutionClient, error) {
|
||||
return nil, errors.New("factory failed")
|
||||
}
|
||||
service.resizeAsyncWorkerCapacity(context.Background())
|
||||
if service.riverExecutionClient != oldClient || service.riverWorkerCapacity != 4 {
|
||||
t.Fatal("create failure replaced the current client")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("start failure", func(t *testing.T) {
|
||||
oldClient := &fakeAsyncExecutionClient{}
|
||||
failedClient := &fakeAsyncExecutionClient{startErr: errors.New("start failed")}
|
||||
service := asyncWorkerManagerTestService(4)
|
||||
service.riverExecutionClient = oldClient
|
||||
service.asyncCapacityLoader = fixedAsyncCapacity(8)
|
||||
service.asyncClientFactory = func(int) (asyncExecutionClient, error) {
|
||||
return failedClient, nil
|
||||
}
|
||||
service.resizeAsyncWorkerCapacity(context.Background())
|
||||
if service.riverExecutionClient != oldClient || service.riverWorkerCapacity != 4 {
|
||||
t.Fatal("start failure replaced the current client")
|
||||
}
|
||||
if failedClient.stopCancelled.Load() != 1 {
|
||||
t.Fatal("failed replacement client was not cleaned up")
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestStopAsyncWorkersIncludesCurrentAndDrainingClients(t *testing.T) {
|
||||
current := &fakeAsyncExecutionClient{}
|
||||
draining := &fakeAsyncExecutionClient{}
|
||||
service := asyncWorkerManagerTestService(2)
|
||||
service.riverExecutionClient = current
|
||||
service.riverDrainingClients[draining] = struct{}{}
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
service.stopAsyncWorkersOnShutdown(ctx)
|
||||
if current.stopCancelled.Load() != 1 || draining.stopCancelled.Load() != 1 {
|
||||
t.Fatalf("shutdown cancellations current=%d draining=%d, want 1/1", current.stopCancelled.Load(), draining.stopCancelled.Load())
|
||||
}
|
||||
}
|
||||
|
||||
func asyncWorkerManagerTestService(capacity int) *Service {
|
||||
return &Service{
|
||||
cfg: config.Config{
|
||||
AsyncWorkerHardLimit: 2048,
|
||||
AsyncWorkerRefreshIntervalSeconds: 5,
|
||||
},
|
||||
logger: slog.New(slog.NewTextHandler(io.Discard, nil)),
|
||||
riverDrainingClients: make(map[asyncExecutionClient]struct{}),
|
||||
riverWorkerCapacity: capacity,
|
||||
riverExecutionClient: nil,
|
||||
}
|
||||
}
|
||||
|
||||
func fixedAsyncCapacity(capacity int) func(context.Context, int) (store.AsyncWorkerCapacitySnapshot, error) {
|
||||
return func(context.Context, int) (store.AsyncWorkerCapacitySnapshot, error) {
|
||||
return store.AsyncWorkerCapacitySnapshot{Capacity: capacity, Desired: capacity, HardLimit: 2048}, nil
|
||||
}
|
||||
}
|
||||
|
||||
func waitForAtomicValue(t *testing.T, value *atomic.Int64, want int64) {
|
||||
t.Helper()
|
||||
deadline := time.Now().Add(time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
if value.Load() == want {
|
||||
return
|
||||
}
|
||||
time.Sleep(time.Millisecond)
|
||||
}
|
||||
t.Fatalf("atomic value=%d, want=%d", value.Load(), want)
|
||||
}
|
||||
|
||||
func waitForDrainingClients(t *testing.T, service *Service, want int) {
|
||||
t.Helper()
|
||||
deadline := time.Now().Add(time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
service.riverMu.RLock()
|
||||
count := len(service.riverDrainingClients)
|
||||
service.riverMu.RUnlock()
|
||||
if count == want {
|
||||
return
|
||||
}
|
||||
time.Sleep(time.Millisecond)
|
||||
}
|
||||
service.riverMu.RLock()
|
||||
count := len(service.riverDrainingClients)
|
||||
service.riverMu.RUnlock()
|
||||
t.Fatalf("draining client count=%d, want=%d", count, want)
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -159,7 +159,12 @@ func (s *Service) hydrateProviderRequestAssetString(ctx context.Context, value s
|
||||
if raw == "" || !imageInputFieldNeedsHydration(path) {
|
||||
return value, nil
|
||||
}
|
||||
style, ok := requestAssetCapabilityHydrationForMedia("image", candidate, raw, "")
|
||||
style, ok := requestAssetHydrateUnsupported, false
|
||||
if openAIImageEditRequiresMultipartBytes(candidate) {
|
||||
style, ok = requestAssetHydrateDataURL, true
|
||||
} else {
|
||||
style, ok = requestAssetCapabilityHydrationForMedia("image", candidate, raw, "")
|
||||
}
|
||||
if !ok {
|
||||
return value, nil
|
||||
}
|
||||
@@ -317,6 +322,9 @@ func requestAssetHydrationForField(path []string, asset store.RequestAsset, cand
|
||||
return requestAssetHydrateDataURL
|
||||
}
|
||||
if requestAssetMediaKindForHydration(path, asset) == "image" {
|
||||
if openAIImageEditRequiresMultipartBytes(candidate) {
|
||||
return requestAssetHydrateDataURL
|
||||
}
|
||||
if style, ok := requestAssetCapabilityHydrationForMedia("image", candidate, asset.URL, asset.StorageProvider); ok {
|
||||
return style
|
||||
}
|
||||
@@ -332,6 +340,11 @@ func requestAssetHydrationForField(path []string, asset store.RequestAsset, cand
|
||||
return requestAssetHydrateURL
|
||||
}
|
||||
|
||||
func openAIImageEditRequiresMultipartBytes(candidate store.RuntimeModelCandidate) bool {
|
||||
return normalizeProviderKey(candidate.Provider) == "openai" &&
|
||||
strings.TrimSpace(candidate.ModelType) == "image_edit"
|
||||
}
|
||||
|
||||
func requestAssetMediaKindForHydration(path []string, asset store.RequestAsset) string {
|
||||
if mediaURLFieldNeedsHydration(path) {
|
||||
return requestAssetMediaURLKind(path)
|
||||
|
||||
@@ -45,6 +45,56 @@ func TestHydrateProviderRequestAssetsConvertsStrictBase64Field(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestHydrateProviderRequestAssetsConvertsInputAudioDataToRawBase64(t *testing.T) {
|
||||
storageDir := t.TempDir()
|
||||
fileName := "gateway-request-asset-audio.mp3"
|
||||
payload := []byte("audio bytes")
|
||||
if err := os.WriteFile(filepath.Join(storageDir, fileName), payload, 0o644); err != nil {
|
||||
t.Fatalf("write request asset: %v", err)
|
||||
}
|
||||
service := &Service{cfg: config.Config{LocalUploadedStorageDir: storageDir}}
|
||||
body := map[string]any{
|
||||
"messages": []any{
|
||||
map[string]any{
|
||||
"role": "user",
|
||||
"content": []any{
|
||||
map[string]any{
|
||||
"type": "input_audio",
|
||||
"input_audio": map[string]any{
|
||||
"data": map[string]any{
|
||||
"assetRef": map[string]any{
|
||||
"sha256": "sha-input-audio",
|
||||
"contentType": "audio/mpeg",
|
||||
"url": "/static/uploaded/" + fileName,
|
||||
"storageProvider": "local_static",
|
||||
},
|
||||
"url": "/static/uploaded/" + fileName,
|
||||
},
|
||||
"format": "mp3",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
hydrated, err := service.hydrateProviderRequestAssets(context.Background(), body, store.RuntimeModelCandidate{})
|
||||
if err != nil {
|
||||
t.Fatalf("hydrate request assets: %v", err)
|
||||
}
|
||||
messages := hydrated["messages"].([]any)
|
||||
message := messages[0].(map[string]any)
|
||||
content := message["content"].([]any)
|
||||
audioPart := content[0].(map[string]any)
|
||||
inputAudio := audioPart["input_audio"].(map[string]any)
|
||||
if got, want := stringFromAny(inputAudio["data"]), base64.StdEncoding.EncodeToString(payload); got != want {
|
||||
t.Fatalf("unexpected hydrated input audio: got %q want %q", got, want)
|
||||
}
|
||||
if inputAudio["format"] != "mp3" {
|
||||
t.Fatalf("input audio format was not preserved: %+v", inputAudio)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHydrateProviderRequestAssetsConvertsBase64ArrayField(t *testing.T) {
|
||||
storageDir := t.TempDir()
|
||||
fileName := "gateway-request-asset-array.png"
|
||||
@@ -205,6 +255,41 @@ func TestHydrateProviderRequestAssetsUsesImageCapabilityBase64ForTopLevelImageAs
|
||||
}
|
||||
}
|
||||
|
||||
func TestHydrateProviderRequestAssetsConvertsOpenAIEditImagesForMultipart(t *testing.T) {
|
||||
payload := []byte{0x89, 'P', 'N', 'G', '\r', '\n', 0x1a, '\n', 1, 2, 3, 4}
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "image/png")
|
||||
_, _ = w.Write(payload)
|
||||
}))
|
||||
defer server.Close()
|
||||
service := &Service{}
|
||||
body := map[string]any{
|
||||
"images": []any{
|
||||
server.URL + "/first.png",
|
||||
server.URL + "/second.png",
|
||||
},
|
||||
}
|
||||
|
||||
hydrated, err := service.hydrateProviderRequestAssets(context.Background(), body, store.RuntimeModelCandidate{
|
||||
Provider: "openai",
|
||||
ModelType: "image_edit",
|
||||
Capabilities: map[string]any{
|
||||
"image_edit": map[string]any{
|
||||
"support_url_input": true,
|
||||
"support_base64_input": false,
|
||||
},
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("hydrate OpenAI edit images: %v", err)
|
||||
}
|
||||
images := hydrated["images"].([]any)
|
||||
want := "data:image/png;base64," + base64.StdEncoding.EncodeToString(payload)
|
||||
if len(images) != 2 || stringFromAny(images[0]) != want || stringFromAny(images[1]) != want {
|
||||
t.Fatalf("OpenAI edit images should be hydrated for multipart: %+v", images)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHydrateProviderRequestAssetsImageCapabilityOverridesProviderDataURLDefault(t *testing.T) {
|
||||
service := &Service{}
|
||||
body := map[string]any{
|
||||
|
||||
@@ -42,6 +42,29 @@ type failoverDecision struct {
|
||||
Info failureInfo
|
||||
}
|
||||
|
||||
type failureDecision struct {
|
||||
Route string
|
||||
Effect string
|
||||
Target string
|
||||
CooldownSeconds int
|
||||
DemoteSteps int
|
||||
Reason string
|
||||
Match policyRuleMatch
|
||||
Info failureInfo
|
||||
}
|
||||
|
||||
type resolveCandidateFailureInput struct {
|
||||
RunnerPolicy store.RunnerPolicy
|
||||
Candidate store.RuntimeModelCandidate
|
||||
Err error
|
||||
ClientAttempt int
|
||||
MaxClientAttempts int
|
||||
HasNextCandidate bool
|
||||
FailoverExpired bool
|
||||
Async bool
|
||||
DownstreamStarted bool
|
||||
}
|
||||
|
||||
type priorityDemoteDecision struct {
|
||||
Demote bool
|
||||
Reason string
|
||||
@@ -57,6 +80,9 @@ func shouldRetrySameClient(candidate store.RuntimeModelCandidate, err error) boo
|
||||
func retryDecisionForCandidate(candidate store.RuntimeModelCandidate, err error) retryDecision {
|
||||
policy := effectiveRetryPolicy(candidate)
|
||||
info := failureInfoFromError(err)
|
||||
if isResultPersistenceFailure(err) {
|
||||
return retryDecision{Retry: false, Reason: "result_persistence_failed", Match: policyRuleMatch{Source: "gateway_result_storage", Policy: "persistence", Rule: "providerRetry", Value: "disabled"}, Info: info}
|
||||
}
|
||||
if errors.Is(err, store.ErrRateLimited) {
|
||||
return retryDecision{Retry: false, Reason: "local_rate_limit_wait_queue", Match: policyRuleMatch{Source: "gateway_rate_limits", Policy: "rateLimitPolicy", Rule: "localCapacity", Value: "exceeded"}, Info: info}
|
||||
}
|
||||
@@ -77,6 +103,9 @@ func retryDecisionForCandidate(candidate store.RuntimeModelCandidate, err error)
|
||||
|
||||
func failoverDecisionForCandidate(runnerPolicy store.RunnerPolicy, candidate store.RuntimeModelCandidate, err error) failoverDecision {
|
||||
info := failureInfoFromError(err)
|
||||
if isResultPersistenceFailure(err) {
|
||||
return failoverDecision{Retry: false, Action: "stop", Reason: "result_persistence_failed", Match: policyRuleMatch{Source: "gateway_result_storage", Policy: "persistence", Rule: "providerFailover", Value: "disabled"}, Info: info}
|
||||
}
|
||||
if strings.TrimSpace(runnerPolicy.Status) != "" && runnerPolicy.Status != "active" {
|
||||
return failoverDecision{Retry: false, Action: "stop", Reason: "runner_policy_disabled", Match: policyRuleMatch{Source: "gateway_runner_policies", Policy: "runnerPolicy", Rule: "status", Value: runnerPolicy.Status}, Info: info}
|
||||
}
|
||||
@@ -98,7 +127,7 @@ func failoverDecisionForCandidate(runnerPolicy store.RunnerPolicy, candidate sto
|
||||
if errors.Is(err, store.ErrRateLimited) {
|
||||
return failoverDecision{Retry: false, Action: "queue", Reason: "local_rate_limit_wait_queue", Match: policyRuleMatch{Source: "gateway_rate_limits", Policy: "rateLimitPolicy", Rule: "localCapacity", Value: "exceeded"}, Info: info}
|
||||
}
|
||||
if ruleDecision, ok := failoverActionRuleDecisionWithSources(runnerPolicy.FailoverPolicy, overridePolicy, info); ok {
|
||||
if ruleDecision, ok := candidateFailoverActionRuleDecision(runnerPolicy.FailoverPolicy, overridePolicy, candidate, info); ok {
|
||||
return ruleDecision
|
||||
}
|
||||
if match, ok := failoverDenyMatchWithSources(runnerPolicy.FailoverPolicy, overridePolicy, info); ok {
|
||||
@@ -119,6 +148,122 @@ func failoverDecisionForCandidate(runnerPolicy store.RunnerPolicy, candidate sto
|
||||
return failoverDecision{Retry: false, Action: "stop", Reason: "client_non_retryable", Match: policyRuleMatch{Source: "provider_client", Policy: "ClientError", Rule: "Retryable", Value: "false"}, Info: info}
|
||||
}
|
||||
|
||||
func resolveCandidateFailure(input resolveCandidateFailureInput) failureDecision {
|
||||
info := failureInfoFromError(input.Err)
|
||||
if isResultPersistenceFailure(input.Err) {
|
||||
return failureDecision{
|
||||
Route: "stop",
|
||||
Effect: "none",
|
||||
Reason: "result_persistence_failed",
|
||||
Match: policyRuleMatch{Source: "gateway_result_storage", Policy: "persistence", Rule: "providerFailover", Value: "disabled"},
|
||||
Info: info,
|
||||
}
|
||||
}
|
||||
if errors.Is(input.Err, store.ErrRateLimited) {
|
||||
route := "stop"
|
||||
if input.Async && store.RateLimitRetryable(input.Err) {
|
||||
route = "requeue"
|
||||
}
|
||||
return failureDecision{
|
||||
Route: route,
|
||||
Effect: "none",
|
||||
Reason: "local_rate_limit_wait_queue",
|
||||
Match: policyRuleMatch{Source: "gateway_rate_limits", Policy: "rateLimitPolicy", Rule: "localCapacity", Value: "exceeded"},
|
||||
Info: info,
|
||||
}
|
||||
}
|
||||
|
||||
failover := failoverDecisionForCandidate(input.RunnerPolicy, input.Candidate, input.Err)
|
||||
effect := failoverEffect(failover.Action)
|
||||
if input.DownstreamStarted {
|
||||
return failureDecision{
|
||||
Route: "stop",
|
||||
Effect: effect,
|
||||
Target: failover.Target,
|
||||
CooldownSeconds: failover.CooldownSeconds,
|
||||
DemoteSteps: failover.DemoteSteps,
|
||||
Reason: "downstream_response_started",
|
||||
Match: policyRuleMatch{Source: "gateway_stream", Policy: "streamSafety", Rule: "downstreamStarted", Value: "true"},
|
||||
Info: info,
|
||||
}
|
||||
}
|
||||
explicitStop := failover.Action == "stop" && (failover.Reason == "failover_action_rule" ||
|
||||
failover.Reason == "hard_stop_policy" ||
|
||||
failover.Reason == "failover_deny_policy" ||
|
||||
failover.Reason == "failover_disabled" ||
|
||||
failover.Reason == "runner_policy_disabled")
|
||||
if explicitStop {
|
||||
return failureDecision{
|
||||
Route: "stop",
|
||||
Effect: effect,
|
||||
Target: failover.Target,
|
||||
CooldownSeconds: failover.CooldownSeconds,
|
||||
DemoteSteps: failover.DemoteSteps,
|
||||
Reason: failover.Reason,
|
||||
Match: failover.Match,
|
||||
Info: failover.Info,
|
||||
}
|
||||
}
|
||||
|
||||
immediateCandidateInvalidation := failover.Action == "cooldown_and_next" || failover.Action == "disable_and_next"
|
||||
retry := retryDecisionForCandidate(input.Candidate, input.Err)
|
||||
if !immediateCandidateInvalidation && retry.Retry && input.ClientAttempt < input.MaxClientAttempts {
|
||||
return failureDecision{
|
||||
Route: "retry_same",
|
||||
Effect: "none",
|
||||
Reason: retry.Reason,
|
||||
Match: retry.Match,
|
||||
Info: retry.Info,
|
||||
}
|
||||
}
|
||||
|
||||
route := "next"
|
||||
reason := failover.Reason
|
||||
match := failover.Match
|
||||
if input.FailoverExpired {
|
||||
route = "stop"
|
||||
reason = "failover_time_budget_exceeded"
|
||||
match = policyRuleMatch{
|
||||
Source: "gateway_runner_policies.failover_policy",
|
||||
Policy: "failoverPolicy",
|
||||
Rule: "maxDurationSeconds",
|
||||
Value: "exceeded",
|
||||
}
|
||||
} else if !input.HasNextCandidate {
|
||||
route = "stop"
|
||||
reason = "no_next_platform"
|
||||
match = policyRuleMatch{
|
||||
Source: "runner_candidates",
|
||||
Policy: "candidateSelection",
|
||||
Rule: "candidateCount",
|
||||
Value: "exhausted",
|
||||
}
|
||||
}
|
||||
return failureDecision{
|
||||
Route: route,
|
||||
Effect: effect,
|
||||
Target: failover.Target,
|
||||
CooldownSeconds: failover.CooldownSeconds,
|
||||
DemoteSteps: failover.DemoteSteps,
|
||||
Reason: reason,
|
||||
Match: match,
|
||||
Info: failover.Info,
|
||||
}
|
||||
}
|
||||
|
||||
func failoverEffect(action string) string {
|
||||
switch action {
|
||||
case "cooldown_and_next":
|
||||
return "cooldown"
|
||||
case "demote_and_next":
|
||||
return "demote"
|
||||
case "disable_and_next":
|
||||
return "disable"
|
||||
default:
|
||||
return "none"
|
||||
}
|
||||
}
|
||||
|
||||
func shouldDemoteCandidatePriority(runnerPolicy store.RunnerPolicy, err error) bool {
|
||||
decision := priorityDemoteDecisionForCandidate(runnerPolicy, err)
|
||||
return decision.Demote
|
||||
@@ -126,6 +271,9 @@ func shouldDemoteCandidatePriority(runnerPolicy store.RunnerPolicy, err error) b
|
||||
|
||||
func priorityDemoteDecisionForCandidate(runnerPolicy store.RunnerPolicy, err error) priorityDemoteDecision {
|
||||
info := failureInfoFromError(err)
|
||||
if isResultPersistenceFailure(err) {
|
||||
return priorityDemoteDecision{Demote: false, Reason: "result_persistence_failed", Info: info}
|
||||
}
|
||||
if strings.TrimSpace(runnerPolicy.Status) != "" && runnerPolicy.Status != "active" {
|
||||
return priorityDemoteDecision{Demote: false, Reason: "runner_policy_disabled", Info: info}
|
||||
}
|
||||
@@ -145,6 +293,19 @@ func priorityDemoteDecisionForCandidate(runnerPolicy store.RunnerPolicy, err err
|
||||
return priorityDemoteDecision{Demote: false, Reason: "priority_demote_no_match", Info: info}
|
||||
}
|
||||
|
||||
func isResultPersistenceFailure(err error) bool {
|
||||
switch strings.ToLower(strings.TrimSpace(clients.ErrorCode(err))) {
|
||||
case "local_result_storage_unavailable",
|
||||
"binary_result_too_large",
|
||||
"binary_result_corrupted",
|
||||
"binary_result_expired",
|
||||
"result_binary_not_materialized":
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func effectiveFailoverPolicy(base map[string]any, override map[string]any) map[string]any {
|
||||
policy := base
|
||||
if nested := failoverOverridePolicy(override); len(nested) > 0 {
|
||||
@@ -324,6 +485,46 @@ func failoverActionRuleDecisionWithSources(base map[string]any, override map[str
|
||||
return failoverActionRuleDecision(actionRulesFromPolicy(base), "gateway_runner_policies.failover_policy", info)
|
||||
}
|
||||
|
||||
func candidateFailoverActionRuleDecision(base map[string]any, override map[string]any, candidate store.RuntimeModelCandidate, info failureInfo) (failoverDecision, bool) {
|
||||
if rules := actionRulesFromPolicy(override); len(rules) > 0 {
|
||||
if decision, ok := failoverActionRuleDecision(rules, "runtime_policy_override.failoverPolicy", info); ok {
|
||||
return decision, true
|
||||
}
|
||||
}
|
||||
autoDisablePolicy := effectiveRuntimePolicy(candidate.AutoDisablePolicy, candidate.RuntimePolicyOverride, "autoDisablePolicy")
|
||||
if boolFromPolicy(autoDisablePolicy, "enabled", false) && intFromPolicy(autoDisablePolicy, "threshold") <= 1 {
|
||||
if value, ok := matchingKeywordValue(autoDisablePolicy, "keywords", info.Target); ok {
|
||||
return failoverDecision{
|
||||
Retry: true,
|
||||
Action: "disable_and_next",
|
||||
Target: "platform",
|
||||
Reason: "legacy_auto_disable_policy",
|
||||
Match: policyRuleMatch{Source: "model_runtime_policy_sets.auto_disable_policy", Policy: "autoDisablePolicy", Rule: "keywords", Value: value},
|
||||
Info: info,
|
||||
}, true
|
||||
}
|
||||
}
|
||||
degradePolicy := effectiveRuntimePolicy(candidate.DegradePolicy, candidate.RuntimePolicyOverride, "degradePolicy")
|
||||
if boolFromPolicy(degradePolicy, "enabled", false) {
|
||||
if value, ok := matchingKeywordValue(degradePolicy, "keywords", info.Target); ok {
|
||||
cooldownSeconds := intFromPolicy(degradePolicy, "cooldownSeconds")
|
||||
if cooldownSeconds <= 0 {
|
||||
cooldownSeconds = 300
|
||||
}
|
||||
return failoverDecision{
|
||||
Retry: true,
|
||||
Action: "cooldown_and_next",
|
||||
Target: "model",
|
||||
Reason: "legacy_degrade_policy",
|
||||
CooldownSeconds: cooldownSeconds,
|
||||
Match: policyRuleMatch{Source: "model_runtime_policy_sets.degrade_policy", Policy: "degradePolicy", Rule: "keywords", Value: value},
|
||||
Info: info,
|
||||
}, true
|
||||
}
|
||||
}
|
||||
return failoverActionRuleDecision(actionRulesFromPolicy(base), "gateway_runner_policies.failover_policy", info)
|
||||
}
|
||||
|
||||
func failoverActionRuleDecision(rules []map[string]any, source string, info failureInfo) (failoverDecision, bool) {
|
||||
for index, rule := range rules {
|
||||
if !boolFromPolicy(rule, "enabled", true) {
|
||||
|
||||
@@ -334,3 +334,151 @@ func TestPriorityDemotePolicyIsSkippedWhenFailoverActionRulesExist(t *testing.T)
|
||||
t.Fatalf("legacy priority demotion should be skipped when actionRules are active, got %+v", decision)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveCandidateFailureCooldownSkipsSameClientRetry(t *testing.T) {
|
||||
decision := resolveCandidateFailure(resolveCandidateFailureInput{
|
||||
RunnerPolicy: store.RunnerPolicy{
|
||||
Status: "active",
|
||||
FailoverPolicy: map[string]any{
|
||||
"enabled": true,
|
||||
"actionRules": []any{
|
||||
map[string]any{
|
||||
"statusCodes": []any{429},
|
||||
"action": "cooldown_and_next",
|
||||
"target": "model",
|
||||
"cooldownSeconds": 60,
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
Candidate: store.RuntimeModelCandidate{
|
||||
ModelRetryPolicy: map[string]any{"enabled": true, "maxAttempts": 3},
|
||||
},
|
||||
Err: &clients.ClientError{Code: "resource_exhausted", StatusCode: 429, Retryable: true},
|
||||
ClientAttempt: 1,
|
||||
MaxClientAttempts: 3,
|
||||
HasNextCandidate: true,
|
||||
})
|
||||
|
||||
if decision.Route != "next" || decision.Effect != "cooldown" || decision.Target != "model" || decision.CooldownSeconds != 60 {
|
||||
t.Fatalf("429 cooldown rule should invalidate the candidate immediately, got %+v", decision)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveCandidateFailureDemoteRetriesBeforeNext(t *testing.T) {
|
||||
input := resolveCandidateFailureInput{
|
||||
RunnerPolicy: store.RunnerPolicy{
|
||||
Status: "active",
|
||||
FailoverPolicy: map[string]any{
|
||||
"enabled": true,
|
||||
"actionRules": []any{
|
||||
map[string]any{
|
||||
"statusCodes": []any{500},
|
||||
"action": "demote_and_next",
|
||||
"target": "platform",
|
||||
"demoteSteps": 2,
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
Candidate: store.RuntimeModelCandidate{
|
||||
ModelRetryPolicy: map[string]any{"enabled": true, "maxAttempts": 2},
|
||||
},
|
||||
Err: &clients.ClientError{Code: "server_error", StatusCode: 500, Retryable: true},
|
||||
ClientAttempt: 1,
|
||||
MaxClientAttempts: 2,
|
||||
HasNextCandidate: true,
|
||||
}
|
||||
|
||||
first := resolveCandidateFailure(input)
|
||||
if first.Route != "retry_same" || first.Effect != "none" {
|
||||
t.Fatalf("demotion must wait until same-client retries are exhausted, got %+v", first)
|
||||
}
|
||||
|
||||
input.ClientAttempt = 2
|
||||
second := resolveCandidateFailure(input)
|
||||
if second.Route != "next" || second.Effect != "demote" || second.DemoteSteps != 2 {
|
||||
t.Fatalf("exhausted retry should demote once and rotate, got %+v", second)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveCandidateFailureStopsAfterDownstreamStarts(t *testing.T) {
|
||||
decision := resolveCandidateFailure(resolveCandidateFailureInput{
|
||||
RunnerPolicy: store.RunnerPolicy{
|
||||
Status: "active",
|
||||
FailoverPolicy: map[string]any{
|
||||
"enabled": true,
|
||||
"actionRules": []any{
|
||||
map[string]any{
|
||||
"categories": []any{"stream_error"},
|
||||
"action": "next",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
Candidate: store.RuntimeModelCandidate{ModelRetryPolicy: map[string]any{"enabled": true, "maxAttempts": 3}},
|
||||
Err: &clients.ClientError{Code: "stream_read_error", Retryable: true},
|
||||
ClientAttempt: 1,
|
||||
MaxClientAttempts: 3,
|
||||
HasNextCandidate: true,
|
||||
DownstreamStarted: true,
|
||||
})
|
||||
|
||||
if decision.Route != "stop" || decision.Reason != "downstream_response_started" {
|
||||
t.Fatalf("stream output must prevent retries and failover, got %+v", decision)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveCandidateFailureLegacyPolicyPrecedence(t *testing.T) {
|
||||
candidate := store.RuntimeModelCandidate{
|
||||
DegradePolicy: map[string]any{
|
||||
"enabled": true,
|
||||
"keywords": []any{"resource_exhausted"},
|
||||
"cooldownSeconds": 45,
|
||||
},
|
||||
}
|
||||
decision := resolveCandidateFailure(resolveCandidateFailureInput{
|
||||
RunnerPolicy: store.RunnerPolicy{
|
||||
Status: "active",
|
||||
FailoverPolicy: map[string]any{
|
||||
"enabled": true,
|
||||
"actionRules": []any{
|
||||
map[string]any{
|
||||
"statusCodes": []any{429},
|
||||
"action": "cooldown_and_next",
|
||||
"target": "platform",
|
||||
"cooldownSeconds": 90,
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
Candidate: candidate,
|
||||
Err: &clients.ClientError{Code: "resource_exhausted", Message: "resource_exhausted", StatusCode: 429, Retryable: true},
|
||||
ClientAttempt: 1,
|
||||
MaxClientAttempts: 3,
|
||||
HasNextCandidate: true,
|
||||
})
|
||||
|
||||
if decision.Match.Policy != "degradePolicy" || decision.Target != "model" || decision.CooldownSeconds != 45 {
|
||||
t.Fatalf("legacy compatibility rule should precede global action rules, got %+v", decision)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveCandidateFailureUnmatchedRetryableUsesSameThenNext(t *testing.T) {
|
||||
input := resolveCandidateFailureInput{
|
||||
RunnerPolicy: store.RunnerPolicy{Status: "active", FailoverPolicy: map[string]any{"enabled": true, "actionRules": []any{}}},
|
||||
Candidate: store.RuntimeModelCandidate{ModelRetryPolicy: map[string]any{"enabled": true, "maxAttempts": 2}},
|
||||
Err: &clients.ClientError{Code: "unknown_transient", Retryable: true},
|
||||
ClientAttempt: 1,
|
||||
MaxClientAttempts: 2,
|
||||
HasNextCandidate: true,
|
||||
}
|
||||
|
||||
if decision := resolveCandidateFailure(input); decision.Route != "retry_same" {
|
||||
t.Fatalf("unmatched retryable error should retry the same client first, got %+v", decision)
|
||||
}
|
||||
input.ClientAttempt = 2
|
||||
if decision := resolveCandidateFailure(input); decision.Route != "next" || decision.Effect != "none" {
|
||||
t.Fatalf("unmatched exhausted error should rotate without side effects, got %+v", decision)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2,153 +2,59 @@ package runner
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"strings"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/clients"
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
)
|
||||
|
||||
func (s *Service) applyCandidateFailurePolicies(ctx context.Context, taskID string, candidate store.RuntimeModelCandidate, cause error, simulated bool, singleSourceProtected bool) {
|
||||
code := clients.ErrorCode(cause)
|
||||
message := ""
|
||||
if cause != nil {
|
||||
message = cause.Error()
|
||||
func (s *Service) applyFailureDecision(ctx context.Context, taskID string, requestedModel string, candidate store.RuntimeModelCandidate, decision failureDecision, singleSourceProtection bool, simulated bool) (store.CandidateFailureEffectResult, error) {
|
||||
result, err := s.store.ApplyCandidateFailureEffect(ctx, store.CandidateFailureEffectInput{
|
||||
Effect: decision.Effect,
|
||||
Target: decision.Target,
|
||||
PlatformID: candidate.PlatformID,
|
||||
PlatformModelID: candidate.PlatformModelID,
|
||||
RequestedModel: requestedModel,
|
||||
ModelType: candidate.ModelType,
|
||||
CooldownSeconds: decision.CooldownSeconds,
|
||||
DemoteSteps: decision.DemoteSteps,
|
||||
SingleSourceProtection: singleSourceProtection,
|
||||
})
|
||||
if err != nil {
|
||||
return result, err
|
||||
}
|
||||
|
||||
autoDisablePolicy := effectiveRuntimePolicy(candidate.AutoDisablePolicy, candidate.RuntimePolicyOverride, "autoDisablePolicy")
|
||||
if failurePolicyMatches(autoDisablePolicy, code, message) && intFromPolicy(autoDisablePolicy, "threshold") <= 1 {
|
||||
if singleSourceProtected {
|
||||
s.emitSingleSourceProtected(ctx, taskID, candidate, "auto_disable", code, simulated)
|
||||
} else if err := s.store.DisableCandidatePlatform(ctx, candidate.PlatformID); err == nil {
|
||||
_ = s.emit(ctx, taskID, "task.policy.auto_disabled", "running", "auto_disable", 0.48, "candidate platform disabled by failure policy", map[string]any{
|
||||
"platformId": candidate.PlatformID,
|
||||
"platformModelId": candidate.PlatformModelID,
|
||||
"code": code,
|
||||
}, simulated)
|
||||
}
|
||||
}
|
||||
|
||||
degradePolicy := effectiveRuntimePolicy(candidate.DegradePolicy, candidate.RuntimePolicyOverride, "degradePolicy")
|
||||
if failurePolicyMatches(degradePolicy, code, message) {
|
||||
cooldownSeconds := intFromPolicy(degradePolicy, "cooldownSeconds")
|
||||
if singleSourceProtected {
|
||||
s.emitSingleSourceProtected(ctx, taskID, candidate, "degrade_cooldown", code, simulated)
|
||||
} else if err := s.store.CooldownCandidatePlatformModel(ctx, candidate.PlatformModelID, cooldownSeconds); err == nil {
|
||||
_ = s.emit(ctx, taskID, "task.policy.degraded", "running", "degrade", 0.5, "candidate model cooled down by failure policy", map[string]any{
|
||||
"platformId": candidate.PlatformID,
|
||||
"platformModelId": candidate.PlatformModelID,
|
||||
"cooldownSeconds": cooldownSeconds,
|
||||
"code": code,
|
||||
}, simulated)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Service) applyFailoverAction(ctx context.Context, taskID string, candidate store.RuntimeModelCandidate, requestedModel string, decision failoverDecision, simulated bool, singleSourceProtected bool) {
|
||||
switch decision.Action {
|
||||
case "disable_and_next":
|
||||
if singleSourceProtected {
|
||||
s.emitSingleSourceProtected(ctx, taskID, candidate, decision.Action, decision.Info.Code, simulated)
|
||||
return
|
||||
}
|
||||
target := decision.Target
|
||||
if target == "" {
|
||||
target = "platform"
|
||||
}
|
||||
var err error
|
||||
if target == "model" {
|
||||
err = s.store.DisableCandidatePlatformModel(ctx, candidate.PlatformModelID)
|
||||
} else {
|
||||
err = s.store.DisableCandidatePlatform(ctx, candidate.PlatformID)
|
||||
}
|
||||
if err == nil {
|
||||
_ = s.emit(ctx, taskID, "task.policy.failover_disabled", "running", "failover", 0.51, "candidate platform disabled by runner failover policy", addPolicyTracePayload(map[string]any{
|
||||
"platformId": candidate.PlatformID,
|
||||
"platformModelId": candidate.PlatformModelID,
|
||||
"action": decision.Action,
|
||||
"target": target,
|
||||
"reason": decision.Reason,
|
||||
}, decision.Match, decision.Info), simulated)
|
||||
}
|
||||
case "cooldown_and_next":
|
||||
if singleSourceProtected {
|
||||
s.emitSingleSourceProtected(ctx, taskID, candidate, decision.Action, decision.Info.Code, simulated)
|
||||
return
|
||||
}
|
||||
target := decision.Target
|
||||
if target == "" {
|
||||
target = "model"
|
||||
}
|
||||
var err error
|
||||
if target == "platform" {
|
||||
err = s.store.CooldownCandidatePlatform(ctx, candidate.PlatformID, decision.CooldownSeconds)
|
||||
} else {
|
||||
err = s.store.CooldownCandidatePlatformModel(ctx, candidate.PlatformModelID, decision.CooldownSeconds)
|
||||
}
|
||||
if err == nil {
|
||||
_ = s.emit(ctx, taskID, "task.policy.failover_cooled_down", "running", "failover", 0.51, "candidate model cooled down by runner failover policy", addPolicyTracePayload(map[string]any{
|
||||
"platformId": candidate.PlatformID,
|
||||
"platformModelId": candidate.PlatformModelID,
|
||||
"cooldownSeconds": decision.CooldownSeconds,
|
||||
"action": decision.Action,
|
||||
"target": target,
|
||||
"reason": decision.Reason,
|
||||
}, decision.Match, decision.Info), simulated)
|
||||
}
|
||||
case "demote_and_next":
|
||||
demoteSteps := decision.DemoteSteps
|
||||
if demoteSteps <= 0 {
|
||||
demoteSteps = 1
|
||||
}
|
||||
if dynamicPriority, err := s.store.DemoteCandidatePlatformPriorityBySteps(ctx, candidate.PlatformID, candidate.PlatformModelID, requestedModel, candidate.ModelType, demoteSteps); err == nil {
|
||||
_ = s.emit(ctx, taskID, "task.policy.priority_demoted", "running", "priority_demote", 0.52, "candidate platform priority demoted by runner failover policy", addPolicyTracePayload(map[string]any{
|
||||
"platformId": candidate.PlatformID,
|
||||
"platformModelId": candidate.PlatformModelID,
|
||||
"dynamicPriority": dynamicPriority,
|
||||
"demoteSteps": demoteSteps,
|
||||
"action": decision.Action,
|
||||
"target": "platform",
|
||||
"reason": decision.Reason,
|
||||
"errorMessage": decision.Info.Message,
|
||||
}, decision.Match, decision.Info), simulated)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Service) emitSingleSourceProtected(ctx context.Context, taskID string, candidate store.RuntimeModelCandidate, action string, code string, simulated bool) {
|
||||
_ = s.emit(ctx, taskID, "task.policy.single_source_protected", "running", "policy_guard", 0.5, "single source protected from disable or cooldown", map[string]any{
|
||||
"platformId": candidate.PlatformID,
|
||||
"platformModelId": candidate.PlatformModelID,
|
||||
"action": action,
|
||||
"code": code,
|
||||
}, simulated)
|
||||
}
|
||||
|
||||
func (s *Service) applyPriorityDemotePolicy(ctx context.Context, taskID string, attemptNo int, runnerPolicy store.RunnerPolicy, candidate store.RuntimeModelCandidate, requestedModel string, cause error, simulated bool) {
|
||||
if errors.Is(cause, store.ErrRateLimited) {
|
||||
return
|
||||
}
|
||||
decision := priorityDemoteDecisionForCandidate(runnerPolicy, cause)
|
||||
if !decision.Demote {
|
||||
return
|
||||
}
|
||||
demoteSteps := decision.DemoteSteps
|
||||
if demoteSteps <= 0 {
|
||||
demoteSteps = 99
|
||||
}
|
||||
if dynamicPriority, err := s.store.DemoteCandidatePlatformPriorityBySteps(ctx, candidate.PlatformID, candidate.PlatformModelID, requestedModel, candidate.ModelType, demoteSteps); err == nil {
|
||||
s.recordAttemptTrace(ctx, taskID, attemptNo, priorityDemoteTraceEntry(decision, candidate.PlatformID, candidate.PlatformModelID, dynamicPriority))
|
||||
_ = s.emit(ctx, taskID, "task.policy.priority_demoted", "running", "priority_demote", 0.52, "candidate platform priority demoted by runner policy", addPolicyTracePayload(map[string]any{
|
||||
if result.GuardReason == "single_source" {
|
||||
_ = s.emit(ctx, taskID, "task.policy.single_source_protected", "running", "policy_guard", 0.5, "single source protected from disable or cooldown", addPolicyTracePayload(map[string]any{
|
||||
"platformId": candidate.PlatformID,
|
||||
"platformModelId": candidate.PlatformModelID,
|
||||
"dynamicPriority": dynamicPriority,
|
||||
"demoteSteps": demoteSteps,
|
||||
"code": clients.ErrorCode(cause),
|
||||
"reason": decision.Reason,
|
||||
"errorMessage": messageFromError(cause),
|
||||
"effect": decision.Effect,
|
||||
"target": decision.Target,
|
||||
"guardReason": result.GuardReason,
|
||||
}, decision.Match, decision.Info), simulated)
|
||||
return result, nil
|
||||
}
|
||||
if !result.Applied {
|
||||
return result, nil
|
||||
}
|
||||
|
||||
payload := addPolicyTracePayload(map[string]any{
|
||||
"platformId": candidate.PlatformID,
|
||||
"platformModelId": candidate.PlatformModelID,
|
||||
"effect": decision.Effect,
|
||||
"target": decision.Target,
|
||||
"reason": decision.Reason,
|
||||
}, decision.Match, decision.Info)
|
||||
switch decision.Effect {
|
||||
case "disable":
|
||||
_ = s.emit(ctx, taskID, "task.policy.failover_disabled", "running", "failover", 0.51, "candidate disabled by runner failover policy", payload, simulated)
|
||||
case "cooldown":
|
||||
payload["cooldownSeconds"] = decision.CooldownSeconds
|
||||
_ = s.emit(ctx, taskID, "task.policy.failover_cooled_down", "running", "failover", 0.51, "candidate cooled down by runner failover policy", payload, simulated)
|
||||
case "demote":
|
||||
payload["demoteSteps"] = decision.DemoteSteps
|
||||
payload["dynamicPriority"] = result.DynamicPriority
|
||||
_ = s.emit(ctx, taskID, "task.policy.priority_demoted", "running", "priority_demote", 0.52, "candidate platform priority demoted by runner failover policy", payload, simulated)
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func messageFromError(err error) string {
|
||||
|
||||
+433
-138
@@ -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)
|
||||
}
|
||||
|
||||
@@ -58,10 +58,11 @@ func (s *Service) CancelTask(ctx context.Context, taskID string, user *auth.User
|
||||
return taskCancelUnavailable(task, "任务已开始执行,当前阶段不可取消,请继续查询结果"), nil
|
||||
}
|
||||
if task.RiverJobID > 0 {
|
||||
if s.riverClient == nil {
|
||||
riverClient := s.asyncControlClient()
|
||||
if riverClient == nil {
|
||||
return taskCancelUnavailable(task, "任务取消队列未就绪,请继续查询结果"), nil
|
||||
}
|
||||
job, err := s.riverClient.JobGet(ctx, task.RiverJobID)
|
||||
job, err := riverClient.JobGet(ctx, task.RiverJobID)
|
||||
if errors.Is(err, rivertype.ErrNotFound) {
|
||||
return taskCancelUnavailable(task, "任务已不在本地排队队列,可能已提交上游,当前不可取消,请继续查询结果"), nil
|
||||
}
|
||||
@@ -71,7 +72,7 @@ func (s *Service) CancelTask(ctx context.Context, taskID string, user *auth.User
|
||||
if job == nil || !riverJobStateCancellable(job.State) {
|
||||
return taskCancelUnavailable(task, "任务已不在可取消队列状态,请继续查询结果"), nil
|
||||
}
|
||||
if _, err := s.riverClient.JobDelete(ctx, task.RiverJobID); err != nil {
|
||||
if _, err := riverClient.JobDelete(ctx, task.RiverJobID); err != nil {
|
||||
if errors.Is(err, rivertype.ErrJobRunning) || errors.Is(err, rivertype.ErrNotFound) {
|
||||
return taskCancelUnavailable(task, "任务已被工作进程领取,当前不可取消,请继续查询结果"), nil
|
||||
}
|
||||
|
||||
@@ -0,0 +1,181 @@
|
||||
package runner
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"hash/fnv"
|
||||
"io"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
func (s *Service) StartTaskHistoryWorkers(ctx context.Context) {
|
||||
if s.cfg.TaskProgressCallbackEnabled && s.cfg.TaskProgressCallbackURL != "" {
|
||||
go s.runTaskCallbackWorker(ctx, "callback-"+uuid.NewString())
|
||||
}
|
||||
if s.cfg.TaskCleanupEnabled {
|
||||
go s.runTaskHistoryCleanupWorker(ctx)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Service) runTaskCallbackWorker(ctx context.Context, workerID string) {
|
||||
ticker := time.NewTicker(time.Second)
|
||||
defer ticker.Stop()
|
||||
client := &http.Client{Timeout: time.Duration(s.cfg.TaskProgressCallbackTimeoutMS) * time.Millisecond}
|
||||
for {
|
||||
s.processTaskCallbackBatch(ctx, client, workerID)
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Service) processTaskCallbackBatch(ctx context.Context, client *http.Client, workerID string) {
|
||||
items, err := s.store.ClaimTaskCallbacks(ctx, workerID, store.TaskCallbackBatchSize, store.TaskCallbackLockTTL)
|
||||
if err != nil {
|
||||
if ctx.Err() == nil {
|
||||
s.logger.Error("claim task callbacks failed", "error_category", "task_callback_claim_failed")
|
||||
}
|
||||
return
|
||||
}
|
||||
for _, item := range items {
|
||||
if ctx.Err() != nil {
|
||||
return
|
||||
}
|
||||
statusCode, deliveryErr := deliverTaskCallback(ctx, client, item, s.cfg.ServerMainInternalToken)
|
||||
if deliveryErr == nil && statusCode >= 200 && statusCode < 300 {
|
||||
if err := s.store.MarkTaskCallbackDelivered(context.WithoutCancel(ctx), item); err != nil {
|
||||
s.logger.Error("mark task callback delivered failed", "taskID", item.TaskID, "seq", item.Seq, "error_category", "task_callback_state_failed")
|
||||
}
|
||||
continue
|
||||
}
|
||||
retryable := deliveryErr != nil || statusCode == http.StatusRequestTimeout || statusCode == http.StatusTooManyRequests || statusCode >= 500
|
||||
retry := retryable && item.Attempts < s.cfg.TaskProgressCallbackMaxAttempts
|
||||
message := "callback_network_error"
|
||||
if deliveryErr == nil {
|
||||
message = fmt.Sprintf("callback_http_%d", statusCode)
|
||||
}
|
||||
nextAttemptAt := time.Now().Add(taskCallbackRetryDelay(item.ID, item.Attempts))
|
||||
if err := s.store.MarkTaskCallbackFailed(context.WithoutCancel(ctx), item, retry, nextAttemptAt, message); err != nil {
|
||||
s.logger.Error("mark task callback failed", "taskID", item.TaskID, "seq", item.Seq, "error_category", "task_callback_state_failed")
|
||||
continue
|
||||
}
|
||||
if !retry {
|
||||
s.logger.Warn("task callback moved to dead letter", "taskID", item.TaskID, "seq", item.Seq, "statusCode", statusCode, "error_category", "task_callback_failed")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func deliverTaskCallback(ctx context.Context, client *http.Client, item store.TaskCallbackDelivery, bearerToken string) (int, error) {
|
||||
body, err := json.Marshal(map[string]any{
|
||||
"taskId": item.TaskID,
|
||||
"seq": item.Seq,
|
||||
"eventType": item.EventType,
|
||||
"status": item.TaskStatus,
|
||||
"createdAt": item.CreatedAt.UTC().Format(time.RFC3339Nano),
|
||||
})
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
request, err := http.NewRequestWithContext(ctx, http.MethodPost, item.CallbackURL, bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
request.Header.Set("Idempotency-Key", fmt.Sprintf("%s:%d", item.TaskID, item.Seq))
|
||||
request.Header.Set("X-EasyAI-Event-Type", item.EventType)
|
||||
if bearerToken != "" {
|
||||
request.Header.Set("Authorization", "Bearer "+bearerToken)
|
||||
}
|
||||
response, err := client.Do(request)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
defer response.Body.Close()
|
||||
_, _ = io.Copy(io.Discard, io.LimitReader(response.Body, 4096))
|
||||
return response.StatusCode, nil
|
||||
}
|
||||
|
||||
func taskCallbackRetryDelay(deliveryID string, attempt int) time.Duration {
|
||||
delays := []time.Duration{
|
||||
5 * time.Second,
|
||||
30 * time.Second,
|
||||
2 * time.Minute,
|
||||
10 * time.Minute,
|
||||
30 * time.Minute,
|
||||
}
|
||||
index := attempt - 1
|
||||
if index < 0 {
|
||||
index = 0
|
||||
}
|
||||
if index >= len(delays) {
|
||||
index = len(delays) - 1
|
||||
}
|
||||
base := delays[index]
|
||||
hash := fnv.New32a()
|
||||
_, _ = hash.Write([]byte(deliveryID))
|
||||
jitter := time.Duration(hash.Sum32()%21) * base / 100
|
||||
return base + jitter
|
||||
}
|
||||
|
||||
func (s *Service) runTaskHistoryCleanupWorker(ctx context.Context) {
|
||||
interval := time.Duration(s.cfg.TaskCleanupIntervalSeconds) * time.Second
|
||||
ticker := time.NewTicker(interval)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
s.processTaskHistoryCleanup(ctx)
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Service) processTaskHistoryCleanup(ctx context.Context) {
|
||||
for batch := 0; batch < 10; batch++ {
|
||||
now := time.Now().UTC()
|
||||
result, err := s.store.CleanupTaskHistory(
|
||||
ctx,
|
||||
now.AddDate(0, 0, -s.cfg.TaskAnalysisRetentionDays),
|
||||
now.AddDate(0, 0, -s.cfg.TaskRetentionDays),
|
||||
s.cfg.TaskCleanupBatchSize,
|
||||
)
|
||||
if err != nil {
|
||||
if ctx.Err() == nil {
|
||||
s.logger.Error("cleanup task history failed", "error_category", "task_history_cleanup_failed")
|
||||
}
|
||||
return
|
||||
}
|
||||
if result.Total() == 0 {
|
||||
return
|
||||
}
|
||||
s.logger.Info("task history cleanup batch completed",
|
||||
"compactedTasks", result.CompactedTasks,
|
||||
"compactedCheckpoints", result.CompactedCheckpoints,
|
||||
"compactedAttempts", result.CompactedAttempts,
|
||||
"compactedEvents", result.CompactedEvents,
|
||||
"compactedCallbacks", result.CompactedCallbacks,
|
||||
"retiredLegacyCallbacks", result.RetiredCallbacks,
|
||||
"compactedParamLogs", result.CompactedParamLogs,
|
||||
"compactedClonedVoices", result.CompactedClonedVoices,
|
||||
"deletedCallbacks", result.DeletedCallbacks,
|
||||
"deletedParamLogs", result.DeletedParamLogs,
|
||||
"deletedEvents", result.DeletedEvents,
|
||||
"deletedAttempts", result.DeletedAttempts,
|
||||
"deletedTasks", result.DeletedTasks,
|
||||
)
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-time.After(100 * time.Millisecond):
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,118 @@
|
||||
package runner
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
)
|
||||
|
||||
func TestMinimalRemoteTaskCheckpointUsesAllowlist(t *testing.T) {
|
||||
checkpoint := minimalRemoteTaskCheckpoint("keling", "", map[string]any{
|
||||
"endpoint": "/videos/task",
|
||||
"mode": "omni_video",
|
||||
"cleanupElementIds": []any{"element-1"},
|
||||
"phase": "uploaded",
|
||||
"submit": map[string]any{"large": "provider response"},
|
||||
"payload": map[string]any{"prompt": "duplicate request"},
|
||||
"raw": map[string]any{"data": "duplicate result"},
|
||||
"status": "processing",
|
||||
})
|
||||
if checkpoint["endpoint"] != "/videos/task" || checkpoint["mode"] != "omni_video" || checkpoint["cleanupElementIds"] == nil {
|
||||
t.Fatalf("checkpoint lost required fields: %#v", checkpoint)
|
||||
}
|
||||
for _, forbidden := range []string{"phase", "submit", "payload", "raw", "status"} {
|
||||
if _, ok := checkpoint[forbidden]; ok {
|
||||
t.Fatalf("checkpoint contains forbidden field %q: %#v", forbidden, checkpoint)
|
||||
}
|
||||
}
|
||||
if unknown := minimalRemoteTaskCheckpoint("unknown-provider", "", map[string]any{"endpoint": "/forbidden"}); len(unknown) != 0 {
|
||||
t.Fatalf("unknown provider checkpoint=%#v, want empty", unknown)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCanonicalTaskResultDropsProviderRawCopy(t *testing.T) {
|
||||
result := canonicalTaskResult(map[string]any{
|
||||
"id": "task",
|
||||
"data": []any{map[string]any{"url": "https://example.com/output"}},
|
||||
"raw": map[string]any{"provider": "duplicate"},
|
||||
"raw_data": map[string]any{"provider": "duplicate"},
|
||||
"provider_response": map[string]any{"provider": "duplicate"},
|
||||
"upstream_task_id": "remote-task",
|
||||
"compatibility_data": "kept",
|
||||
})
|
||||
if result["raw"] != nil || result["raw_data"] != nil || result["provider_response"] != nil ||
|
||||
result["upstream_task_id"] != nil || result["id"] != "task" || result["data"] == nil ||
|
||||
result["compatibility_data"] != "kept" {
|
||||
t.Fatalf("canonical result=%#v", result)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMinimalParameterPreprocessingChangesDropsValuesAndCapsSize(t *testing.T) {
|
||||
changes := make([]parameterPreprocessChange, 0, 100)
|
||||
for index := 0; index < 100; index++ {
|
||||
changes = append(changes, parameterPreprocessChange{
|
||||
Path: "input.image",
|
||||
Action: "converted",
|
||||
Before: map[string]any{"base64": "forbidden"},
|
||||
After: map[string]any{"url": "forbidden"},
|
||||
})
|
||||
}
|
||||
minimal := minimalParameterPreprocessingChanges(changes)
|
||||
encoded, err := json.Marshal(minimal)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(encoded) > 1024 {
|
||||
t.Fatalf("minimal changes size=%d, want <=1024", len(encoded))
|
||||
}
|
||||
for _, raw := range minimal {
|
||||
item := raw.(map[string]any)
|
||||
if len(item) != 2 || item["path"] != "input.image" || item["action"] != "converted" {
|
||||
t.Fatalf("unexpected minimal change: %#v", item)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestDeliverTaskCallbackSendsMinimalAuthenticatedBody(t *testing.T) {
|
||||
createdAt := time.Date(2026, time.July, 24, 10, 0, 0, 0, time.UTC)
|
||||
var received map[string]any
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) {
|
||||
if request.Header.Get("Authorization") != "Bearer internal-token" {
|
||||
t.Errorf("Authorization=%q", request.Header.Get("Authorization"))
|
||||
}
|
||||
if request.Header.Get("Idempotency-Key") != "task-id:3" {
|
||||
t.Errorf("Idempotency-Key=%q", request.Header.Get("Idempotency-Key"))
|
||||
}
|
||||
if request.Header.Get("X-EasyAI-Event-Type") != "task.completed" {
|
||||
t.Errorf("X-EasyAI-Event-Type=%q", request.Header.Get("X-EasyAI-Event-Type"))
|
||||
}
|
||||
if err := json.NewDecoder(request.Body).Decode(&received); err != nil {
|
||||
t.Error(err)
|
||||
}
|
||||
w.WriteHeader(http.StatusNoContent)
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
status, err := deliverTaskCallback(context.Background(), server.Client(), store.TaskCallbackDelivery{
|
||||
TaskID: "task-id",
|
||||
Seq: 3,
|
||||
CallbackURL: server.URL,
|
||||
EventType: "task.completed",
|
||||
TaskStatus: "succeeded",
|
||||
CreatedAt: createdAt,
|
||||
}, "internal-token")
|
||||
if err != nil || status != http.StatusNoContent {
|
||||
t.Fatalf("status=%d err=%v", status, err)
|
||||
}
|
||||
if len(received) != 5 {
|
||||
t.Fatalf("callback fields=%#v, want exactly five fields", received)
|
||||
}
|
||||
if received["taskId"] != "task-id" || received["eventType"] != "task.completed" || received["status"] != "succeeded" {
|
||||
t.Fatalf("callback body=%#v", received)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -167,16 +167,7 @@ func (s *Service) persistVoiceCloneResult(ctx context.Context, task store.Gatewa
|
||||
DemoAudioURL: demoAudioURL,
|
||||
Status: "active",
|
||||
ExpiresAt: &expiresAt,
|
||||
Metadata: map[string]any{
|
||||
"request": map[string]any{
|
||||
"textValidation": body["text_validation"],
|
||||
"languageBoost": body["language_boost"],
|
||||
"needNoiseReduction": body["need_noise_reduction"],
|
||||
"needVolumeNormalization": body["need_volume_normalization"],
|
||||
"aigcWatermark": body["aigc_watermark"],
|
||||
},
|
||||
"rawData": result["raw_data"],
|
||||
},
|
||||
Metadata: map[string]any{},
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
@@ -44,6 +44,20 @@ type Metrics struct {
|
||||
billingEstimateFailed atomic.Uint64
|
||||
billingIdempotentReplay atomic.Uint64
|
||||
billingPricingUnavailable atomic.Uint64
|
||||
asyncWorkerCapacity atomic.Int64
|
||||
asyncWorkerDesiredCapacity atomic.Int64
|
||||
asyncWorkerHardLimit atomic.Int64
|
||||
asyncWorkerCapacityCapped atomic.Int64
|
||||
asyncWorkerResizeSuccess atomic.Uint64
|
||||
asyncWorkerRefreshFailed atomic.Uint64
|
||||
asyncWorkerCreateFailed atomic.Uint64
|
||||
asyncWorkerStartFailed atomic.Uint64
|
||||
leaseRenewalSuccess atomic.Uint64
|
||||
leaseRenewalFailure atomic.Uint64
|
||||
leaseRenewalLost atomic.Uint64
|
||||
taskEventDuplicate atomic.Uint64
|
||||
taskEventUnknownType atomic.Uint64
|
||||
taskEventBudgetExceeded atomic.Uint64
|
||||
}
|
||||
|
||||
var processingDurationBounds = [...]time.Duration{
|
||||
@@ -123,6 +137,52 @@ func (m *Metrics) ObserveBillingEvent(event string) {
|
||||
}
|
||||
}
|
||||
|
||||
func (m *Metrics) SetAsyncWorkerCapacity(current, desired, hardLimit int, capped bool) {
|
||||
m.asyncWorkerCapacity.Store(int64(current))
|
||||
m.asyncWorkerDesiredCapacity.Store(int64(desired))
|
||||
m.asyncWorkerHardLimit.Store(int64(hardLimit))
|
||||
cappedValue := int64(0)
|
||||
if capped {
|
||||
cappedValue = 1
|
||||
}
|
||||
m.asyncWorkerCapacityCapped.Store(cappedValue)
|
||||
}
|
||||
|
||||
func (m *Metrics) ObserveAsyncWorkerResize(outcome string) {
|
||||
switch outcome {
|
||||
case "success":
|
||||
m.asyncWorkerResizeSuccess.Add(1)
|
||||
case "refresh_failed":
|
||||
m.asyncWorkerRefreshFailed.Add(1)
|
||||
case "create_failed":
|
||||
m.asyncWorkerCreateFailed.Add(1)
|
||||
case "start_failed":
|
||||
m.asyncWorkerStartFailed.Add(1)
|
||||
}
|
||||
}
|
||||
|
||||
func (m *Metrics) ObserveConcurrencyLeaseRenewal(outcome string) {
|
||||
switch outcome {
|
||||
case "success":
|
||||
m.leaseRenewalSuccess.Add(1)
|
||||
case "lost":
|
||||
m.leaseRenewalLost.Add(1)
|
||||
default:
|
||||
m.leaseRenewalFailure.Add(1)
|
||||
}
|
||||
}
|
||||
|
||||
func (m *Metrics) ObserveTaskEventSkip(reason string) {
|
||||
switch reason {
|
||||
case "duplicate":
|
||||
m.taskEventDuplicate.Add(1)
|
||||
case "budget_exceeded":
|
||||
m.taskEventBudgetExceeded.Add(1)
|
||||
default:
|
||||
m.taskEventUnknownType.Add(1)
|
||||
}
|
||||
}
|
||||
|
||||
func (m *Metrics) Handler(provider MetricsSnapshotProvider, issuer, audience string, enabled bool) http.Handler {
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
mode := "disabled"
|
||||
@@ -212,6 +272,26 @@ func (m *Metrics) Handler(provider MetricsSnapshotProvider, issuer, audience str
|
||||
plainCounter(w, "easyai_gateway_billing_estimate_failures_total", "Pricing estimate requests that failed.", m.billingEstimateFailed.Load())
|
||||
plainCounter(w, "easyai_gateway_billing_idempotent_replays_total", "Generation requests replayed idempotently.", m.billingIdempotentReplay.Load())
|
||||
plainCounter(w, "easyai_gateway_billing_pricing_unavailable_total", "Pricing requests rejected because no effective price was available.", m.billingPricingUnavailable.Load())
|
||||
plainGauge(w, "easyai_gateway_async_worker_capacity", "Current River asynchronous worker capacity.", m.asyncWorkerCapacity.Load())
|
||||
plainGauge(w, "easyai_gateway_async_worker_desired_capacity", "Uncapped asynchronous worker capacity derived from model and user-group policies.", m.asyncWorkerDesiredCapacity.Load())
|
||||
plainGauge(w, "easyai_gateway_async_worker_hard_limit", "Per-process asynchronous worker safety limit.", m.asyncWorkerHardLimit.Load())
|
||||
plainGauge(w, "easyai_gateway_async_worker_capacity_capped", "Whether desired asynchronous worker capacity is capped by the hard limit.", m.asyncWorkerCapacityCapped.Load())
|
||||
outcomeCounters(w, "easyai_gateway_async_worker_resizes_total", "Asynchronous worker resize attempts by bounded outcome.", []outcomeValue{
|
||||
{"success", m.asyncWorkerResizeSuccess.Load()},
|
||||
{"refresh_failed", m.asyncWorkerRefreshFailed.Load()},
|
||||
{"create_failed", m.asyncWorkerCreateFailed.Load()},
|
||||
{"start_failed", m.asyncWorkerStartFailed.Load()},
|
||||
})
|
||||
outcomeCounters(w, "easyai_gateway_concurrency_lease_renewals_total", "Concurrency lease renewals by bounded outcome.", []outcomeValue{
|
||||
{"success", m.leaseRenewalSuccess.Load()},
|
||||
{"failure", m.leaseRenewalFailure.Load()},
|
||||
{"lost", m.leaseRenewalLost.Load()},
|
||||
})
|
||||
outcomeCounters(w, "easyai_gateway_task_events_skipped_total", "Task events skipped before persistence by bounded reason.", []outcomeValue{
|
||||
{"duplicate", m.taskEventDuplicate.Load()},
|
||||
{"unknown_type", m.taskEventUnknownType.Load()},
|
||||
{"budget_exceeded", m.taskEventBudgetExceeded.Load()},
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
@@ -249,3 +329,7 @@ func outcomeCounters(w http.ResponseWriter, name, help string, values []outcomeV
|
||||
func plainCounter(w http.ResponseWriter, name, help string, value uint64) {
|
||||
fmt.Fprintf(w, "# HELP %s %s\n# TYPE %s counter\n%s %d\n", name, help, name, name, value)
|
||||
}
|
||||
|
||||
func plainGauge(w http.ResponseWriter, name, help string, value int64) {
|
||||
fmt.Fprintf(w, "# HELP %s %s\n# TYPE %s gauge\n%s %d\n", name, help, name, name, value)
|
||||
}
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user